1 /*
2 * Copyright (c) 2026 Netflix, Inc. Written by Warner Losh
3 *
4 * SPDX-License-Identifier: BSD-2-Clause
5 */
6
7 #include "stand.h"
8 #include "loader_efi.h"
9 #include <efilib.h>
10 #include "decompress.h"
11
12 #include <zlib.h>
13 #include <bzlib.h>
14 #include <xz.h>
15 #ifdef LOADER_ZFS_SUPPORT /* ZSTD and lzma only available with ZFS */
16 #include <zstd.h>
17 #endif
18 #include <sys/_param.h>
19
20 #define ULL(x) ((unsigned long long)(x))
21
22 /*
23 * While memdisk_uefi.efi used EfiReservedMemoryType, we need to use
24 * EfiLoaderData instead, the standard memory type the loader uses to pass
25 * memory to the kernel. This is also safe if we're chain booting to another efi
26 * progarm since the memory will be available to that program. The other type
27 * isn't necessarily in the pmap, so the md driver fails.
28 */
29 static EFI_MEMORY_TYPE mem_type = EfiLoaderData;
30
31 struct decomp_state
32 {
33 enum compression type;
34 size_t size; /* Best guess at final size */
35 size_t alloc_size; /* Current size of `buf` */
36 size_t pages; /* Alloc_size in pages */
37 uint8_t *buf_cur; /* Current output buffer */
38 uint8_t *buf_end; /* Current end of allocated buffer */
39 EFI_PHYSICAL_ADDRESS buf; /* Decompression buffer */
40 union {
41 z_stream zstrm;
42 bz_stream bzstrm;
43 struct xz_dec *xzstrm;
44 #ifdef LOADER_ZFS_SUPPORT
45 ZSTD_DStream *zstdstrm;
46 #endif
47 };
48 EFI_STATUS (*init)(decomp_state *dctx, uint8_t *first_buf, size_t buflen,
49 size_t size_hint);
50 enum step_return (*step)(decomp_state *dctx, uint8_t *buf, size_t len, size_t offset);
51 void (*fini)(decomp_state *dctx, bool flush);
52 };
53
54 static enum compression
what_compressed(uint8_t * buf,size_t len)55 what_compressed(uint8_t *buf, size_t len)
56 {
57 /*
58 * Failsafe
59 */
60 if (len < 32)
61 return (none);
62 if (memcmp(buf, "\x1f\x8b", 2) == 0) {
63 printf("GZIP\n");
64 return (zlib);
65 }
66 if (memcmp(buf, "BZh", 3) == 0) {
67 printf("BZIP2\n");
68 return (bzip2);
69 }
70 if (memcmp(buf, "\x28\xb5\x2f\xfd", 4) == 0) {
71 printf("zstd\n");
72 return (zstd);
73 }
74 if (memcmp(buf, "\xfd""7zXZ\x00", 6) == 0) {
75 printf("xz\n");
76 return (xz);
77 }
78 printf("Not compressed\n");
79 return (none);
80 }
81
82 static EFI_STATUS
alloc_buffer(decomp_state * dctx,size_t size)83 alloc_buffer(decomp_state *dctx, size_t size)
84 {
85 dctx->alloc_size = roundup2(size, EFI_PAGE_SIZE);
86 dctx->pages = dctx->alloc_size / EFI_PAGE_SIZE;
87 EFI_STATUS status = BS->AllocatePages(AllocateAnyPages, mem_type, dctx->pages, &dctx->buf);
88 if (EFI_ERROR(status)) {
89 printf("Failed to allocate memory for %llu bytes\n", ULL(dctx->alloc_size));
90 return (status);
91 }
92 BS->SetMem((void *)(uintptr_t)dctx->buf, dctx->alloc_size, 0);
93 dctx->buf_cur = (uint8_t *)(uintptr_t)dctx->buf;
94 dctx->buf_end = (uint8_t *)(uintptr_t)dctx->buf + dctx->alloc_size;
95 return (EFI_SUCCESS);
96 }
97
98 static EFI_STATUS
grow_buffer(decomp_state * dctx)99 grow_buffer(decomp_state *dctx)
100 {
101 /*
102 * a 1.5 exp growth trades a few more copies for a little less waste.
103 */
104 size_t newsz = roundup2(dctx->alloc_size * 3 / 2, EFI_PAGE_SIZE);
105 size_t newpages = newsz / EFI_PAGE_SIZE;
106 EFI_PHYSICAL_ADDRESS newbuf;
107 EFI_STATUS status = BS->AllocatePages(AllocateAnyPages, mem_type, newpages, &newbuf);
108 if (EFI_ERROR(status)) {
109 printf("Failed to allocate memory for %llu bytes\n", ULL(newsz));
110 return (status);
111 }
112 memcpy((void *)(uintptr_t)newbuf, (void *)(uintptr_t)dctx->buf, dctx->alloc_size);
113 BS->FreePages(dctx->buf, dctx->pages);
114 dctx->buf = newbuf;
115 dctx->pages = newpages;
116 dctx->buf_cur = (uint8_t *)(uintptr_t)dctx->buf + dctx->alloc_size;
117 dctx->buf_end = (uint8_t *)(uintptr_t)dctx->buf + newsz;
118 BS->SetMem(dctx->buf_cur, newsz - dctx->alloc_size, 0);
119 dctx->alloc_size = newsz;
120 return (EFI_SUCCESS);
121 }
122
123 static void
free_buffer(decomp_state * dctx)124 free_buffer(decomp_state *dctx)
125 {
126 if (dctx->buf)
127 BS->FreePages(dctx->buf, dctx->pages);
128 dctx->buf = 0;
129 }
130
131 /*
132 * zlib supprot
133 */
134 static EFI_STATUS
zlib_init(decomp_state * dctx,uint8_t * first_buf,size_t buflen,size_t size_hint)135 zlib_init(decomp_state *dctx, uint8_t *first_buf, size_t buflen, size_t size_hint)
136 {
137 z_stream *strm = &dctx->zstrm;
138
139 /*
140 * Assume 4x compression, but start at 64MB
141 */
142 dctx->size = max(size_hint * 4, M(64));
143 EFI_STATUS status = alloc_buffer(dctx, dctx->size);
144 if (EFI_ERROR(status))
145 return (status);
146 memset(strm, 0, sizeof(*strm));
147 strm->next_in = first_buf;
148 strm->avail_in = buflen;
149 return (inflateInit2(strm, 15 + 16) == Z_OK ? EFI_SUCCESS : EFI_VOLUME_CORRUPTED);
150 }
151
152 static enum step_return
zlib_step(decomp_state * dctx,uint8_t * buf,size_t len,size_t offset)153 zlib_step(decomp_state *dctx, uint8_t *buf, size_t len, size_t offset)
154 {
155 z_stream *strm = &dctx->zstrm;
156 size_t outlen = dctx->buf_end - dctx->buf_cur;
157
158 strm->next_in = buf;
159 strm->avail_in = len;
160 strm->next_out = dctx->buf_cur;
161 strm->avail_out = outlen;
162
163 int ret = inflate(strm, Z_NO_FLUSH);
164 dctx->buf_cur += outlen - strm->avail_out;
165
166 if (ret == Z_STREAM_END)
167 return (done);
168 if (ret != Z_OK)
169 return (err);
170 if (dctx->buf_cur < dctx->buf_end) /* Have output space */
171 return (ok);
172
173 /*
174 * We're out of space, grow the buffer and try again if there's buffer
175 * space. We try again recursively since we know that will usually go
176 * only 1 deep.
177 */
178 if (EFI_ERROR(grow_buffer(dctx)))
179 return (err);
180 if (strm->avail_in == 0)
181 return (ok);
182 size_t consumed = len - strm->avail_in;
183 return (zlib_step(dctx, buf + consumed, strm->avail_in, offset + consumed));
184 }
185
186 static void
zlib_fini(decomp_state * dctx,bool flush)187 zlib_fini(decomp_state *dctx, bool flush)
188 {
189 inflateEnd(&dctx->zstrm);
190 if (!flush)
191 return;
192 free_buffer(dctx);
193 }
194
195 /*
196 * Bzip2 supprot
197 */
198 static EFI_STATUS
bzip2_init(decomp_state * dctx,uint8_t * first_buf,size_t buflen,size_t size_hint)199 bzip2_init(decomp_state *dctx, uint8_t *first_buf, size_t buflen, size_t size_hint)
200 {
201 bz_stream *strm = &dctx->bzstrm;
202
203 /*
204 * Assume 4x compression, but start at 64MB
205 */
206 dctx->size = max(size_hint * 4, M(64));
207 EFI_STATUS status = alloc_buffer(dctx, dctx->size);
208 if (EFI_ERROR(status))
209 return (status);
210 memset(strm, 0, sizeof(*strm));
211 strm->next_in = first_buf;
212 strm->avail_in = buflen;
213 return (BZ2_bzDecompressInit(strm, 0, 0) == BZ_OK ? EFI_SUCCESS : EFI_VOLUME_CORRUPTED);
214 }
215
216 static enum step_return
bzip2_step(decomp_state * dctx,uint8_t * buf,size_t len,size_t offset)217 bzip2_step(decomp_state *dctx, uint8_t *buf, size_t len, size_t offset)
218 {
219 bz_stream *strm = &dctx->bzstrm;
220 size_t outlen = dctx->buf_end - dctx->buf_cur;
221
222 strm->next_in = buf;
223 strm->avail_in = len;
224 strm->next_out = dctx->buf_cur;
225 strm->avail_out = outlen;
226
227 int ret = BZ2_bzDecompress(strm);
228 dctx->buf_cur += outlen - strm->avail_out;
229
230 if (ret == BZ_STREAM_END)
231 return (done);
232 if (ret != BZ_OK)
233 return (err);
234 if (dctx->buf_cur < dctx->buf_end) /* Have output space */
235 return (ok);
236
237 /*
238 * We're out of space, grow the buffer and try again if there's buffer
239 * space. We try again recursively since we know that will usually go
240 * only 1 deep.
241 */
242 if (EFI_ERROR(grow_buffer(dctx)))
243 return (err);
244 if (strm->avail_in == 0)
245 return (ok);
246 size_t consumed = len - strm->avail_in;
247 return (bzip2_step(dctx, buf + consumed, strm->avail_in, offset + consumed));
248 }
249
250 static void
bzip2_fini(decomp_state * dctx,bool flush)251 bzip2_fini(decomp_state *dctx, bool flush)
252 {
253 BZ2_bzDecompressEnd(&dctx->bzstrm);
254 if (!flush)
255 return;
256 free_buffer(dctx);
257 }
258
259 /*
260 * XZ support
261 */
262 static EFI_STATUS
xz_init(decomp_state * dctx,uint8_t * first_buf,size_t buflen,size_t size_hint)263 xz_init(decomp_state *dctx, uint8_t *first_buf, size_t buflen, size_t size_hint)
264 {
265 /*
266 * Assume 4x compression, but start at 64MB
267 */
268 dctx->size = max(size_hint * 4, M(64));
269 EFI_STATUS status = alloc_buffer(dctx, dctx->size);
270 if (EFI_ERROR(status))
271 return (status);
272 xz_crc32_init();
273 xz_crc64_init();
274 dctx->xzstrm = xz_dec_init(XZ_DYNALLOC, (uint32_t)-1);
275 return (dctx->xzstrm != NULL ? EFI_SUCCESS : EFI_VOLUME_CORRUPTED);
276 }
277
278
279 static enum step_return
xz_step(decomp_state * dctx,uint8_t * buf,size_t len,size_t offset)280 xz_step(decomp_state *dctx, uint8_t *buf, size_t len, size_t offset)
281 {
282 struct xz_dec *strm = dctx->xzstrm;
283 size_t outlen = dctx->buf_end - dctx->buf_cur;
284 struct xz_buf b = { .in = buf, .in_size = len, .in_pos = 0,
285 .out = dctx->buf_cur, .out_size = outlen, .out_pos = 0 };
286 int ret;
287
288 ret = xz_dec_run(strm, &b);
289 dctx->buf_cur += b.out_pos;
290
291 if (ret == XZ_STREAM_END)
292 return (done);
293 if (ret != XZ_OK) {
294 switch(ret) {
295 case XZ_MEM_ERROR:
296 printf("xz no memory ");
297 break;
298 case XZ_DATA_ERROR:
299 printf("xz file corrupted ");
300 break;
301 case XZ_FORMAT_ERROR:
302 printf("xz format not found ");
303 break;
304 case XZ_OPTIONS_ERROR:
305 printf("unsupported xz option ");
306 break;
307 case XZ_MEMLIMIT_ERROR:
308 printf("xz dictionary too small ");
309 break;
310 default:
311 printf("xz step error %d ", ret);
312 break;
313 }
314 printf(" len %d offset %d\n", (int)len, (int)offset);
315 return (err);
316 }
317 if (dctx->buf_cur < dctx->buf_end) /* Have output space */
318 return (ok);
319
320 /*
321 * We're out of space, grow the buffer and try again if there's buffer
322 * space. We try again recursively since we know that will usually go
323 * only 1 deep.
324 */
325 if (EFI_ERROR(grow_buffer(dctx)))
326 return (err);
327 if (b.in_pos == b.in_size)
328 return (ok);
329 size_t consumed = b.in_pos;
330 return (xz_step(dctx, buf + consumed, len - consumed, offset + consumed));
331 }
332
333 static void
xz_fini(decomp_state * dctx,bool flush)334 xz_fini(decomp_state *dctx, bool flush)
335 {
336 xz_dec_end(dctx->xzstrm);
337 dctx->xzstrm = NULL;
338 if (!flush)
339 return;
340 free_buffer(dctx);
341 }
342
343 /*
344 * ZSTD supprot
345 */
346 #ifdef LOADER_ZFS_SUPPORT
347 static EFI_STATUS
zstd_init(decomp_state * dctx,uint8_t * first_buf,size_t buflen,size_t size_hint)348 zstd_init(decomp_state *dctx, uint8_t *first_buf, size_t buflen, size_t size_hint)
349 {
350 unsigned long long size = ZSTD_getFrameContentSize(first_buf, buflen);
351 if (size == ZSTD_CONTENTSIZE_ERROR)
352 return (EFI_VOLUME_CORRUPTED);
353 if (size == ZSTD_CONTENTSIZE_UNKNOWN)
354 dctx->size = max(size_hint * 4, M(64)); /* Guess 4x compression or 64M */
355 else
356 dctx->size = size; /* We know the size */
357 EFI_STATUS status = alloc_buffer(dctx, dctx->size);
358 if (EFI_ERROR(status))
359 return (status);
360
361 dctx->zstdstrm = ZSTD_createDStream();
362 if (dctx->zstdstrm == NULL)
363 return (EFI_OUT_OF_RESOURCES);
364 if (ZSTD_isError(ZSTD_initDStream(dctx->zstdstrm))) {
365 ZSTD_freeDStream(dctx->zstdstrm);
366 dctx->zstdstrm = NULL;
367 return (EFI_OUT_OF_RESOURCES);
368 }
369 return (EFI_SUCCESS);
370 }
371
372 static enum step_return
zstd_step(decomp_state * dctx,uint8_t * buf,size_t len,size_t offset)373 zstd_step(decomp_state *dctx, uint8_t *buf, size_t len, size_t offset)
374 {
375 size_t outlen = dctx->buf_end - dctx->buf_cur;
376 ZSTD_inBuffer inbuf = { buf, len, 0 };
377 ZSTD_outBuffer outbuf = { dctx->buf_cur, outlen, 0 };
378 size_t ret;
379
380 ret = ZSTD_decompressStream(dctx->zstdstrm, &outbuf, &inbuf);
381 dctx->buf_cur += outbuf.pos;
382
383 if (ZSTD_isError(ret))
384 return (err);
385 if (ret == 0)
386 return (done);
387 if (dctx->buf_cur < dctx->buf_end) /* Have output space */
388 return (ok);
389
390 /*
391 * We're out of space, grow the buffer and try again if there's buffer
392 * space. We try again recursively since we know that will usually go
393 * only 1 deep.
394 */
395 if (EFI_ERROR(grow_buffer(dctx)))
396 return (err);
397 if (inbuf.size == inbuf.pos)
398 return (ok);
399 return (zstd_step(dctx, buf + inbuf.pos, inbuf.size - inbuf.pos, offset + inbuf.pos));
400 }
401
402 static void
zstd_fini(decomp_state * dctx,bool flush)403 zstd_fini(decomp_state *dctx, bool flush)
404 {
405 ZSTD_freeDStream(dctx->zstdstrm);
406 if (!flush)
407 return;
408 free_buffer(dctx);
409 }
410 #endif
411
412 /*
413 * No / Unknown decompression fallback
414 */
415 static EFI_STATUS
null_init(decomp_state * dctx,uint8_t * first_buf,size_t buflen,size_t size_hint)416 null_init(decomp_state *dctx, uint8_t *first_buf, size_t buflen, size_t size_hint)
417 {
418 dctx->size = size_hint;
419 return (alloc_buffer(dctx, size_hint));
420 }
421
422 static enum step_return
null_step(decomp_state * dctx,uint8_t * buf,size_t len,size_t offset)423 null_step(decomp_state *dctx, uint8_t *buf, size_t len, size_t offset)
424 {
425 size_t end = offset + len;
426
427 if (end > dctx->size) {
428 printf("Too much data recieved!");
429 return (err);
430 }
431 if ((uintptr_t)dctx->buf_cur - (uintptr_t)dctx->buf != offset) {
432 printf("OH NO! The offset is %llu but I expected %llu\n", ULL(offset),
433 ULL((uintptr_t)dctx->buf_cur - (uintptr_t)dctx->buf));
434 return (err);
435 }
436
437 memcpy(dctx->buf_cur, buf, len);
438 dctx->buf_cur += len;
439 return (end == dctx->size ? done : ok);
440 }
441
442 static void
null_fini(decomp_state * dctx,bool flush)443 null_fini(decomp_state *dctx, bool flush)
444 {
445 if (!flush)
446 return;
447 free_buffer(dctx);
448 }
449
450 decomp_state *
decomp_init(uint8_t * buf,size_t buflen,size_t size_hint)451 decomp_init(uint8_t *buf, size_t buflen, size_t size_hint)
452 {
453 decomp_state *dctx;
454
455 dctx = malloc(sizeof(*dctx));
456 memset(dctx, 0, sizeof(*dctx));
457 dctx->type = what_compressed(buf, buflen);
458 switch (dctx->type) {
459 case zlib:
460 dctx->init = zlib_init;
461 dctx->step = zlib_step;
462 dctx->fini = zlib_fini;
463 break;
464 case bzip2:
465 dctx->init = bzip2_init;
466 dctx->step = bzip2_step;
467 dctx->fini = bzip2_fini;
468 break;
469 case xz:
470 dctx->init = xz_init;
471 dctx->step = xz_step;
472 dctx->fini = xz_fini;
473 break;
474 #ifdef LOADER_ZFS_SUPPORT
475 case zstd:
476 dctx->init = zstd_init;
477 dctx->step = zstd_step;
478 dctx->fini = zstd_fini;
479 break;
480 #endif
481 case none:
482 dctx->init = null_init;
483 dctx->step = null_step;
484 dctx->fini = null_fini;
485 break;
486 default:
487 return (NULL);
488 }
489
490 if (EFI_ERROR(dctx->init(dctx, buf, buflen, size_hint))) {
491 free(dctx);
492 dctx = NULL;
493 }
494 return (dctx);
495 }
496
497 enum step_return
decomp_step(decomp_state * dctx,uint8_t * buf,size_t len,size_t offset)498 decomp_step(decomp_state *dctx, uint8_t *buf, size_t len, size_t offset)
499 {
500 return (dctx->step(dctx, buf, len, offset));
501 }
502
503 void
decomp_fini(decomp_state * dctx,bool flush)504 decomp_fini(decomp_state *dctx, bool flush)
505 {
506 return (dctx->fini(dctx, flush));
507 }
508
509 EFI_PHYSICAL_ADDRESS
decomp_buffer(decomp_state * dctx)510 decomp_buffer(decomp_state *dctx)
511 {
512 if (dctx == NULL)
513 return (0);
514 return (dctx->buf);
515 }
516
517 size_t
decomp_buffer_length(decomp_state * dctx)518 decomp_buffer_length(decomp_state *dctx)
519 {
520 if (dctx == NULL)
521 return (0);
522 return ((uintptr_t)dctx->buf_cur - (uintptr_t)dctx->buf);
523 }
524