1#![cfg_attr(not(feature = "std"), no_std)]
2
3extern crate alloc;
4
5use core::fmt;
6
7use alloc::vec::Vec;
8#[cfg(feature = "std")]
9use std::io::Read;
10
11const ARM64_IMAGE_MIN_SIZE: usize = 64;
12const ARM64_IMAGE_SIZE_OFFSET: usize = 16;
13const ARM64_IMAGE_MAGIC_OFFSET: usize = 56;
14const ARM64_IMAGE_MAGIC: &[u8; 4] = b"ARM\x64";
15const PACKAGE_ALIGN: u64 = 0x1000;
16const PACKAGE_MAGIC: &[u8; 8] = b"ABLXPKG1";
17const PACKAGE_HEADER_LEN: usize = 48;
18const PACKAGE_COMPRESSION_LZ4: u32 = 2;
19const RAMDISK_ALIGN: u64 = 0x1000;
20const RAMDISK_MAGIC: &[u8; 8] = b"ABLXRD1\0";
21const RAMDISK_HEADER_LEN: usize = 72;
22
23#[cfg(feature = "std")]
24const GZIP_MAGIC: &[u8; 2] = b"\x1f\x8b";
25#[cfg(feature = "std")]
26const ZSTD_MAGIC: &[u8; 4] = b"\x28\xb5\x2f\xfd";
27#[cfg(feature = "std")]
28const LINUX_ZBOOT_MIN_SIZE: usize = 28;
29#[cfg(feature = "std")]
30const LINUX_ZBOOT_IMAGE_TYPE_OFFSET: usize = 4;
31#[cfg(feature = "std")]
32const LINUX_ZBOOT_PAYLOAD_OFFSET_OFFSET: usize = 8;
33#[cfg(feature = "std")]
34const LINUX_ZBOOT_PAYLOAD_SIZE_OFFSET: usize = 12;
35#[cfg(feature = "std")]
36const LINUX_ZBOOT_COMP_TYPE_OFFSET: usize = 24;
37#[cfg(feature = "std")]
38const LINUX_ZBOOT_COMP_TYPE_MAX_LEN: usize = 32;
39
40#[derive(Clone, Copy, Debug, Eq, PartialEq)]
41pub enum ImageKind {
42 Kernel,
43 Shim,
44}
45
46impl fmt::Display for ImageKind {
47 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
48 match self {
49 Self::Kernel => f.write_str("kernel"),
50 Self::Shim => f.write_str("abl-exorcist"),
51 }
52 }
53}
54
55#[derive(Clone, Debug, Eq, PartialEq)]
56pub enum AssembleError {
57 ImageTooSmall { image: ImageKind, size: usize },
58 NotArm64Image { image: ImageKind },
59 ZeroImageSize { image: ImageKind },
60 ShimTooLarge { len: u64, payload_offset: u64 },
61 KernelLargerThanImageSize { len: u64, image_size: u64 },
62 EmptyInitrd,
63 InvalidAlignment(u64),
64 SizeOverflow(&'static str),
65}
66
67impl fmt::Display for AssembleError {
68 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
69 match self {
70 Self::ImageTooSmall { image, size } => write!(
71 f,
72 "{image} is smaller than an arm64 Image header: {size} < {ARM64_IMAGE_MIN_SIZE}"
73 ),
74 Self::NotArm64Image { image } => write!(f, "{image} is not a raw arm64 Image"),
75 Self::ZeroImageSize { image } => {
76 write!(f, "{image} arm64 Image header has zero image_size")
77 }
78 Self::ShimTooLarge {
79 len,
80 payload_offset,
81 } => write!(
82 f,
83 "abl-exorcist image length 0x{len:x} exceeds payload offset 0x{payload_offset:x}"
84 ),
85 Self::KernelLargerThanImageSize { len, image_size } => write!(
86 f,
87 "kernel file length 0x{len:x} exceeds arm64 image_size 0x{image_size:x}"
88 ),
89 Self::EmptyInitrd => f.write_str("initrd must not be empty"),
90 Self::InvalidAlignment(alignment) => write!(
91 f,
92 "alignment must be a non-zero power of two: 0x{alignment:x}"
93 ),
94 Self::SizeOverflow(description) => write!(f, "{description} overflows u64/usize"),
95 }
96 }
97}
98
99impl core::error::Error for AssembleError {}
100
101#[cfg(feature = "std")]
102#[derive(Debug)]
103pub enum KernelImageError {
104 InvalidImage(AssembleError),
105 InvalidGzip(std::io::Error),
106 InvalidZstd(std::io::Error),
107 ZbootTooSmall {
108 size: usize,
109 },
110 ZbootPayloadOutOfBounds {
111 payload_offset: usize,
112 payload_size: usize,
113 image_size: usize,
114 },
115 UnsupportedZbootCompression(String),
116 SizeOverflow(&'static str),
117}
118
119#[cfg(feature = "std")]
120impl fmt::Display for KernelImageError {
121 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
122 match self {
123 Self::InvalidImage(err) => write!(f, "{err}"),
124 Self::InvalidGzip(err) => write!(f, "decompress gzip kernel image: {err}"),
125 Self::InvalidZstd(err) => write!(f, "decompress zstd kernel image: {err}"),
126 Self::ZbootTooSmall { size } => {
127 write!(
128 f,
129 "Linux EFI zboot image is too small: {size} < {LINUX_ZBOOT_MIN_SIZE}"
130 )
131 }
132 Self::ZbootPayloadOutOfBounds {
133 payload_offset,
134 payload_size,
135 image_size,
136 } => write!(
137 f,
138 "Linux EFI zboot payload extends beyond image: offset 0x{payload_offset:x} + size 0x{payload_size:x} > 0x{image_size:x}"
139 ),
140 Self::UnsupportedZbootCompression(compression) => write!(
141 f,
142 "unsupported Linux EFI zboot compression type: {compression}"
143 ),
144 Self::SizeOverflow(description) => write!(f, "{description} overflows usize"),
145 }
146 }
147}
148
149#[cfg(feature = "std")]
150impl std::error::Error for KernelImageError {
151 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
152 match self {
153 Self::InvalidImage(err) => Some(err),
154 Self::InvalidGzip(err) | Self::InvalidZstd(err) => Some(err),
155 _ => None,
156 }
157 }
158}
159
160#[cfg(feature = "std")]
161pub fn canonicalize_kernel(kernel: &[u8]) -> Result<Vec<u8>, KernelImageError> {
162 if is_arm64_image(kernel) {
163 verify_kernel(kernel)?;
164 return Ok(kernel.to_vec());
165 }
166
167 if is_linux_zboot(kernel) {
168 return canonicalize_linux_zboot(kernel);
169 }
170
171 if kernel.starts_with(GZIP_MAGIC) {
172 return decompress_gzip_kernel(kernel);
173 }
174
175 if kernel.starts_with(ZSTD_MAGIC) {
176 return decompress_zstd_kernel(kernel);
177 }
178
179 Err(KernelImageError::InvalidImage(
180 arm64_image_size(kernel, ImageKind::Kernel).unwrap_err(),
181 ))
182}
183
184pub fn assemble(kernel: &[u8], shim: &[u8]) -> Result<Vec<u8>, AssembleError> {
188 let kernel_size = arm64_image_size(kernel, ImageKind::Kernel)?;
189 let shim_size = arm64_image_size(shim, ImageKind::Shim)?;
190 let package_offset = package_source_offset(shim_size)?;
191 let shim_file_len =
192 u64::try_from(shim.len()).map_err(|_| AssembleError::SizeOverflow("shim file length"))?;
193
194 if shim_file_len > package_offset {
195 return Err(AssembleError::ShimTooLarge {
196 len: shim_file_len,
197 payload_offset: package_offset,
198 });
199 }
200
201 let kernel_file_len = u64::try_from(kernel.len())
202 .map_err(|_| AssembleError::SizeOverflow("kernel file length"))?;
203 if kernel_file_len > kernel_size {
204 return Err(AssembleError::KernelLargerThanImageSize {
205 len: kernel_file_len,
206 image_size: kernel_size,
207 });
208 }
209
210 let compressed_kernel = lz4_flex::block::compress(kernel);
211 let compressed_kernel_len = u64::try_from(compressed_kernel.len())
212 .map_err(|_| AssembleError::SizeOverflow("compressed kernel length"))?;
213 let package_len = (PACKAGE_HEADER_LEN as u64)
214 .checked_add(compressed_kernel_len)
215 .ok_or(AssembleError::SizeOverflow("kernel package length"))?;
216 let output_len = package_offset
217 .checked_add(package_len)
218 .ok_or(AssembleError::SizeOverflow("assembled image length"))?;
219 let package_offset = usize::try_from(package_offset)
220 .map_err(|_| AssembleError::SizeOverflow("package offset"))?;
221 let output_len = usize::try_from(output_len)
222 .map_err(|_| AssembleError::SizeOverflow("assembled image length"))?;
223
224 let mut out = Vec::with_capacity(output_len);
225 out.extend_from_slice(shim);
226 out.resize(package_offset, 0);
227 write_package_header(
228 &mut out,
229 compressed_kernel_len,
230 kernel_file_len,
231 kernel_size,
232 );
233 out.extend_from_slice(&compressed_kernel);
234 Ok(out)
235}
236
237pub fn assemble_ramdisk(kernel: &[u8], initrd: &[u8]) -> Result<Vec<u8>, AssembleError> {
243 if initrd.is_empty() {
244 return Err(AssembleError::EmptyInitrd);
245 }
246 let kernel_size = arm64_image_size(kernel, ImageKind::Kernel)?;
247 let kernel_file_len = u64::try_from(kernel.len())
248 .map_err(|_| AssembleError::SizeOverflow("kernel file length"))?;
249 if kernel_file_len > kernel_size {
250 return Err(AssembleError::KernelLargerThanImageSize {
251 len: kernel_file_len,
252 image_size: kernel_size,
253 });
254 }
255
256 let compressed_kernel = lz4_flex::block::compress(kernel);
257 let compressed_kernel_len = u64::try_from(compressed_kernel.len())
258 .map_err(|_| AssembleError::SizeOverflow("compressed kernel length"))?;
259 let initrd_len =
260 u64::try_from(initrd.len()).map_err(|_| AssembleError::SizeOverflow("initrd length"))?;
261 let kernel_offset = align_up(RAMDISK_HEADER_LEN as u64, RAMDISK_ALIGN)?;
262 let initrd_offset = align_up(
263 kernel_offset
264 .checked_add(compressed_kernel_len)
265 .ok_or(AssembleError::SizeOverflow("ramdisk kernel bounds"))?,
266 RAMDISK_ALIGN,
267 )?;
268 let output_len = initrd_offset
269 .checked_add(initrd_len)
270 .ok_or(AssembleError::SizeOverflow("ramdisk container length"))?;
271 let kernel_offset = usize::try_from(kernel_offset)
272 .map_err(|_| AssembleError::SizeOverflow("ramdisk kernel offset"))?;
273 let initrd_offset = usize::try_from(initrd_offset)
274 .map_err(|_| AssembleError::SizeOverflow("ramdisk initrd offset"))?;
275 let output_len = usize::try_from(output_len)
276 .map_err(|_| AssembleError::SizeOverflow("ramdisk container length"))?;
277
278 let mut out = Vec::with_capacity(output_len);
279 write_ramdisk_header(
280 &mut out,
281 kernel_offset as u64,
282 compressed_kernel_len,
283 kernel_file_len,
284 kernel_size,
285 initrd_offset as u64,
286 initrd_len,
287 );
288 out.resize(kernel_offset, 0);
289 out.extend_from_slice(&compressed_kernel);
290 out.resize(initrd_offset, 0);
291 out.extend_from_slice(initrd);
292 Ok(out)
293}
294
295fn write_package_header(
296 out: &mut Vec<u8>,
297 compressed_size: u64,
298 uncompressed_size: u64,
299 image_size: u64,
300) {
301 let start_len = out.len();
302 out.extend_from_slice(PACKAGE_MAGIC);
303 out.extend_from_slice(&(PACKAGE_HEADER_LEN as u32).to_le_bytes());
304 out.extend_from_slice(&PACKAGE_COMPRESSION_LZ4.to_le_bytes());
305 out.extend_from_slice(&compressed_size.to_le_bytes());
306 out.extend_from_slice(&uncompressed_size.to_le_bytes());
307 out.extend_from_slice(&image_size.to_le_bytes());
308 out.extend_from_slice(&0u64.to_le_bytes());
309 debug_assert_eq!(out.len() - start_len, PACKAGE_HEADER_LEN);
310}
311
312fn write_ramdisk_header(
313 out: &mut Vec<u8>,
314 kernel_offset: u64,
315 compressed_size: u64,
316 uncompressed_size: u64,
317 image_size: u64,
318 initrd_offset: u64,
319 initrd_size: u64,
320) {
321 let start_len = out.len();
322 out.extend_from_slice(RAMDISK_MAGIC);
323 out.extend_from_slice(&(RAMDISK_HEADER_LEN as u32).to_le_bytes());
324 out.extend_from_slice(&PACKAGE_COMPRESSION_LZ4.to_le_bytes());
325 out.extend_from_slice(&kernel_offset.to_le_bytes());
326 out.extend_from_slice(&compressed_size.to_le_bytes());
327 out.extend_from_slice(&uncompressed_size.to_le_bytes());
328 out.extend_from_slice(&image_size.to_le_bytes());
329 out.extend_from_slice(&initrd_offset.to_le_bytes());
330 out.extend_from_slice(&initrd_size.to_le_bytes());
331 out.extend_from_slice(&0u64.to_le_bytes());
332 debug_assert_eq!(out.len() - start_len, RAMDISK_HEADER_LEN);
333}
334
335#[cfg(feature = "std")]
336fn canonicalize_linux_zboot(kernel: &[u8]) -> Result<Vec<u8>, KernelImageError> {
337 if kernel.len() < LINUX_ZBOOT_MIN_SIZE {
338 return Err(KernelImageError::ZbootTooSmall { size: kernel.len() });
339 }
340
341 let payload_offset = read_le_u32(kernel, LINUX_ZBOOT_PAYLOAD_OFFSET_OFFSET)?;
342 let payload_size = read_le_u32(kernel, LINUX_ZBOOT_PAYLOAD_SIZE_OFFSET)?;
343 let payload_end =
344 payload_offset
345 .checked_add(payload_size)
346 .ok_or(KernelImageError::SizeOverflow(
347 "Linux EFI zboot payload bounds",
348 ))?;
349 if payload_end > kernel.len() {
350 return Err(KernelImageError::ZbootPayloadOutOfBounds {
351 payload_offset,
352 payload_size,
353 image_size: kernel.len(),
354 });
355 }
356
357 let payload = &kernel[payload_offset..payload_end];
358 let compression = linux_zboot_compression(kernel);
359 if compression.starts_with(b"gzip") {
360 decompress_gzip_kernel(payload)
361 } else if compression.starts_with(b"zstd") {
362 decompress_zstd_kernel(payload)
363 } else {
364 Err(KernelImageError::UnsupportedZbootCompression(
365 String::from_utf8_lossy(compression).into_owned(),
366 ))
367 }
368}
369
370#[cfg(feature = "std")]
371fn decompress_gzip_kernel(kernel: &[u8]) -> Result<Vec<u8>, KernelImageError> {
372 let mut out = Vec::new();
373 flate2::read::GzDecoder::new(kernel)
374 .read_to_end(&mut out)
375 .map_err(KernelImageError::InvalidGzip)?;
376 verify_kernel(&out)?;
377 Ok(out)
378}
379
380#[cfg(feature = "std")]
381fn decompress_zstd_kernel(kernel: &[u8]) -> Result<Vec<u8>, KernelImageError> {
382 let out = zstd::stream::decode_all(kernel).map_err(KernelImageError::InvalidZstd)?;
383 verify_kernel(&out)?;
384 Ok(out)
385}
386
387#[cfg(feature = "std")]
388fn verify_kernel(kernel: &[u8]) -> Result<(), KernelImageError> {
389 arm64_image_size(kernel, ImageKind::Kernel)
390 .map(|_| ())
391 .map_err(KernelImageError::InvalidImage)
392}
393
394#[cfg(feature = "std")]
395fn is_arm64_image(image: &[u8]) -> bool {
396 image.len() >= ARM64_IMAGE_MIN_SIZE
397 && &image[ARM64_IMAGE_MAGIC_OFFSET..ARM64_IMAGE_MAGIC_OFFSET + ARM64_IMAGE_MAGIC.len()]
398 == ARM64_IMAGE_MAGIC
399}
400
401#[cfg(feature = "std")]
402fn is_linux_zboot(image: &[u8]) -> bool {
403 image.len() >= LINUX_ZBOOT_MIN_SIZE
404 && image.starts_with(b"MZ")
405 && &image[LINUX_ZBOOT_IMAGE_TYPE_OFFSET..LINUX_ZBOOT_IMAGE_TYPE_OFFSET + 4] == b"zimg"
406}
407
408#[cfg(feature = "std")]
409fn linux_zboot_compression(image: &[u8]) -> &[u8] {
410 let end = image
411 .len()
412 .min(LINUX_ZBOOT_COMP_TYPE_OFFSET + LINUX_ZBOOT_COMP_TYPE_MAX_LEN);
413 let compression = &image[LINUX_ZBOOT_COMP_TYPE_OFFSET..end];
414 let nul = compression
415 .iter()
416 .position(|byte| *byte == 0)
417 .unwrap_or(compression.len());
418 &compression[..nul]
419}
420
421#[cfg(feature = "std")]
422fn read_le_u32(image: &[u8], offset: usize) -> Result<usize, KernelImageError> {
423 let end = offset
424 .checked_add(4)
425 .ok_or(KernelImageError::SizeOverflow("u32 field bounds"))?;
426 if end > image.len() {
427 return Err(KernelImageError::ZbootTooSmall { size: image.len() });
428 }
429 Ok(u32::from_le_bytes(image[offset..end].try_into().unwrap()) as usize)
430}
431
432pub fn arm64_image_size(image: &[u8], kind: ImageKind) -> Result<u64, AssembleError> {
433 if image.len() < ARM64_IMAGE_MIN_SIZE {
434 return Err(AssembleError::ImageTooSmall {
435 image: kind,
436 size: image.len(),
437 });
438 }
439 if &image[ARM64_IMAGE_MAGIC_OFFSET..ARM64_IMAGE_MAGIC_OFFSET + ARM64_IMAGE_MAGIC.len()]
440 != ARM64_IMAGE_MAGIC
441 {
442 return Err(AssembleError::NotArm64Image { image: kind });
443 }
444
445 let image_size = u64::from_le_bytes(
446 image[ARM64_IMAGE_SIZE_OFFSET..ARM64_IMAGE_SIZE_OFFSET + 8]
447 .try_into()
448 .unwrap(),
449 );
450 if image_size == 0 {
451 return Err(AssembleError::ZeroImageSize { image: kind });
452 }
453 Ok(image_size)
454}
455
456fn package_source_offset(shim_size: u64) -> Result<u64, AssembleError> {
457 align_up(shim_size, PACKAGE_ALIGN)
458}
459
460fn align_up(value: u64, alignment: u64) -> Result<u64, AssembleError> {
461 if alignment == 0 || !alignment.is_power_of_two() {
462 return Err(AssembleError::InvalidAlignment(alignment));
463 }
464
465 let mask = alignment - 1;
466 value
467 .checked_add(mask)
468 .map(|value| value & !mask)
469 .ok_or(AssembleError::SizeOverflow("aligned value"))
470}
471
472#[cfg(test)]
473mod tests {
474 use super::*;
475 use alloc::vec;
476
477 #[cfg(feature = "std")]
478 use std::io::Write;
479
480 #[test]
481 fn assembles_shim_payload_and_padding() {
482 let shim = image(0x1000, 128);
483 let kernel = image(0x2000, 256);
484
485 let assembled = assemble(&kernel, &shim).unwrap();
486
487 assert_eq!(&assembled[..shim.len()], shim.as_slice());
488 assert!(
489 assembled[shim.len()..PACKAGE_ALIGN as usize]
490 .iter()
491 .all(|byte| *byte == 0)
492 );
493
494 assert_eq!(
495 &assembled[PACKAGE_ALIGN as usize..PACKAGE_ALIGN as usize + PACKAGE_MAGIC.len()],
496 PACKAGE_MAGIC
497 );
498 assert_eq!(read_package_u32(&assembled, 0), PACKAGE_HEADER_LEN as u32);
499 assert_eq!(read_package_u32(&assembled, 4), PACKAGE_COMPRESSION_LZ4);
500 assert_eq!(read_package_u64(&assembled, 16), kernel.len() as u64);
501 assert_eq!(read_package_u64(&assembled, 24), 0x2000);
502
503 let compressed_size = read_package_u64(&assembled, 8) as usize;
504 let compressed = &assembled[PACKAGE_ALIGN as usize + PACKAGE_HEADER_LEN
505 ..PACKAGE_ALIGN as usize + PACKAGE_HEADER_LEN + compressed_size];
506 let mut decompressed = vec![0; kernel.len()];
507 let written = lz4_flex::block::decompress_into(compressed, &mut decompressed).unwrap();
508 assert_eq!(written, kernel.len());
509 assert_eq!(decompressed, kernel);
510 }
511
512 #[test]
513 fn assembles_ramdisk_container() {
514 let kernel = image(0x2000, 256);
515 let initrd = b"real initrd";
516
517 let assembled = assemble_ramdisk(&kernel, initrd).unwrap();
518
519 assert_eq!(&assembled[..RAMDISK_MAGIC.len()], RAMDISK_MAGIC);
520 assert_eq!(read_u32(&assembled, 8), RAMDISK_HEADER_LEN as u32);
521 assert_eq!(read_u32(&assembled, 12), PACKAGE_COMPRESSION_LZ4);
522 let kernel_offset = read_u64(&assembled, 16) as usize;
523 let compressed_size = read_u64(&assembled, 24) as usize;
524 let uncompressed_size = read_u64(&assembled, 32);
525 let image_size = read_u64(&assembled, 40);
526 let initrd_offset = read_u64(&assembled, 48) as usize;
527 let initrd_size = read_u64(&assembled, 56) as usize;
528
529 assert_eq!(uncompressed_size, kernel.len() as u64);
530 assert_eq!(image_size, 0x2000);
531 assert_eq!(initrd_size, initrd.len());
532 assert_eq!(
533 &assembled[initrd_offset..initrd_offset + initrd_size],
534 initrd
535 );
536
537 let compressed = &assembled[kernel_offset..kernel_offset + compressed_size];
538 let mut decompressed = vec![0; kernel.len()];
539 let written = lz4_flex::block::decompress_into(compressed, &mut decompressed).unwrap();
540 assert_eq!(written, kernel.len());
541 assert_eq!(decompressed, kernel);
542 }
543
544 #[test]
545 fn ramdisk_rejects_empty_initrd() {
546 let kernel = image(0x2000, 256);
547 assert_eq!(
548 assemble_ramdisk(&kernel, &[]),
549 Err(AssembleError::EmptyInitrd)
550 );
551 }
552
553 #[test]
554 fn package_offset_is_4k_aligned_after_shim_image() {
555 assert_eq!(package_source_offset(1).unwrap(), PACKAGE_ALIGN);
556 assert_eq!(package_source_offset(PACKAGE_ALIGN).unwrap(), PACKAGE_ALIGN);
557 assert_eq!(
558 package_source_offset(PACKAGE_ALIGN + 1).unwrap(),
559 PACKAGE_ALIGN * 2
560 );
561 }
562
563 #[test]
564 fn reads_arm64_image_size() {
565 let image = image(1234, ARM64_IMAGE_MIN_SIZE);
566
567 assert_eq!(arm64_image_size(&image, ImageKind::Kernel), Ok(1234));
568 }
569
570 #[test]
571 fn assembly_rejects_invalid_kernel_headers() {
572 let shim = image(0x1000, 128);
573 let cases = [
574 (
575 vec![0; 63],
576 AssembleError::ImageTooSmall {
577 image: ImageKind::Kernel,
578 size: 63,
579 },
580 ),
581 (
582 vec![0; 64],
583 AssembleError::NotArm64Image {
584 image: ImageKind::Kernel,
585 },
586 ),
587 (
588 image(0, 64),
589 AssembleError::ZeroImageSize {
590 image: ImageKind::Kernel,
591 },
592 ),
593 (
594 image(64, 128),
595 AssembleError::KernelLargerThanImageSize {
596 len: 128,
597 image_size: 64,
598 },
599 ),
600 ];
601 for (kernel, expected) in cases {
602 assert_eq!(assemble(&kernel, &shim), Err(expected.clone()));
603 assert_eq!(assemble_ramdisk(&kernel, b"initrd"), Err(expected));
604 }
605 }
606
607 #[test]
608 fn assembly_rejects_invalid_shim_layout() {
609 let kernel = image(0x2000, 256);
610 assert_eq!(
611 assemble(&kernel, &[0; 64]),
612 Err(AssembleError::NotArm64Image {
613 image: ImageKind::Shim
614 })
615 );
616 assert_eq!(
617 assemble(&kernel, &image(1, 4097)),
618 Err(AssembleError::ShimTooLarge {
619 len: 4097,
620 payload_offset: 4096
621 })
622 );
623 assert_eq!(
624 assemble(&kernel, &image(u64::MAX, 64)),
625 Err(AssembleError::SizeOverflow("aligned value"))
626 );
627 }
628
629 #[cfg(feature = "std")]
630 #[test]
631 fn canonicalizes_raw_arm64_image() {
632 let image = image(0x2000, 256);
633
634 assert_eq!(canonicalize_kernel(&image).unwrap(), image);
635 }
636
637 #[cfg(feature = "std")]
638 #[test]
639 fn canonicalizes_gzip_arm64_image() {
640 let image = image(0x2000, 256);
641 let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
642 encoder.write_all(&image).unwrap();
643 let compressed = encoder.finish().unwrap();
644
645 assert_eq!(canonicalize_kernel(&compressed).unwrap(), image);
646 }
647
648 #[cfg(feature = "std")]
649 #[test]
650 fn canonicalizes_zstd_arm64_image() {
651 let image = image(0x2000, 256);
652 let compressed = zstd::stream::encode_all(image.as_slice(), 0).unwrap();
653
654 assert_eq!(canonicalize_kernel(&compressed).unwrap(), image);
655 }
656
657 #[cfg(feature = "std")]
658 #[test]
659 fn canonicalizes_linux_zboot_gzip_image() {
660 let image = image(0x2000, 256);
661 let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
662 encoder.write_all(&image).unwrap();
663 let compressed = encoder.finish().unwrap();
664 let zboot = linux_zboot(b"gzip", &compressed);
665
666 assert_eq!(canonicalize_kernel(&zboot).unwrap(), image);
667 }
668
669 #[cfg(feature = "std")]
670 #[test]
671 fn canonicalizes_linux_zboot_zstd_image() {
672 let image = image(0x2000, 256);
673 let compressed = zstd::stream::encode_all(image.as_slice(), 0).unwrap();
674 let zboot = linux_zboot(b"zstd", &compressed);
675
676 assert_eq!(canonicalize_kernel(&zboot).unwrap(), image);
677 }
678
679 fn image(image_size: u64, file_len: usize) -> Vec<u8> {
680 let mut image = vec![0; file_len.max(ARM64_IMAGE_MIN_SIZE)];
681 image[ARM64_IMAGE_SIZE_OFFSET..ARM64_IMAGE_SIZE_OFFSET + 8]
682 .copy_from_slice(&image_size.to_le_bytes());
683 image[ARM64_IMAGE_MAGIC_OFFSET..ARM64_IMAGE_MAGIC_OFFSET + ARM64_IMAGE_MAGIC.len()]
684 .copy_from_slice(ARM64_IMAGE_MAGIC);
685 image
686 }
687
688 #[cfg(feature = "std")]
689 fn linux_zboot(compression: &[u8], payload: &[u8]) -> Vec<u8> {
690 let payload_offset = 0x100usize;
691 let mut image = vec![0; payload_offset];
692 image[..2].copy_from_slice(b"MZ");
693 image[LINUX_ZBOOT_IMAGE_TYPE_OFFSET..LINUX_ZBOOT_IMAGE_TYPE_OFFSET + 4]
694 .copy_from_slice(b"zimg");
695 image[LINUX_ZBOOT_PAYLOAD_OFFSET_OFFSET..LINUX_ZBOOT_PAYLOAD_OFFSET_OFFSET + 4]
696 .copy_from_slice(&(payload_offset as u32).to_le_bytes());
697 image[LINUX_ZBOOT_PAYLOAD_SIZE_OFFSET..LINUX_ZBOOT_PAYLOAD_SIZE_OFFSET + 4]
698 .copy_from_slice(&(payload.len() as u32).to_le_bytes());
699 image[LINUX_ZBOOT_COMP_TYPE_OFFSET..LINUX_ZBOOT_COMP_TYPE_OFFSET + compression.len()]
700 .copy_from_slice(compression);
701 image.extend_from_slice(payload);
702 image
703 }
704
705 fn read_package_u32(image: &[u8], offset: usize) -> u32 {
706 let offset = PACKAGE_ALIGN as usize + PACKAGE_MAGIC.len() + offset;
707 u32::from_le_bytes(image[offset..offset + 4].try_into().unwrap())
708 }
709
710 fn read_package_u64(image: &[u8], offset: usize) -> u64 {
711 let offset = PACKAGE_ALIGN as usize + PACKAGE_MAGIC.len() + offset;
712 u64::from_le_bytes(image[offset..offset + 8].try_into().unwrap())
713 }
714
715 fn read_u32(image: &[u8], offset: usize) -> u32 {
716 u32::from_le_bytes(image[offset..offset + 4].try_into().unwrap())
717 }
718
719 fn read_u64(image: &[u8], offset: usize) -> u64 {
720 u64::from_le_bytes(image[offset..offset + 8].try_into().unwrap())
721 }
722}