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