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