Shikha Panwar | 5595711 | 2022-08-22 13:54:33 +0000 | [diff] [blame] | 1 | // Copyright 2022, The Android Open Source Project |
| 2 | // |
| 3 | // Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | // you may not use this file except in compliance with the License. |
| 5 | // You may obtain a copy of the License at |
| 6 | // |
| 7 | // http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | // |
| 9 | // Unless required by applicable law or agreed to in writing, software |
| 10 | // distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | // See the License for the specific language governing permissions and |
| 13 | // limitations under the License. |
| 14 | |
Shikha Panwar | 049c28b | 2023-01-10 05:25:05 +0000 | [diff] [blame] | 15 | //! Attach/Detach bootconfigs to initrd image |
| 16 | use anyhow::{bail, Result}; |
Victor Hsieh | b5bcfab | 2022-09-12 13:06:26 -0700 | [diff] [blame] | 17 | use clap::Parser; |
Shikha Panwar | 049c28b | 2023-01-10 05:25:05 +0000 | [diff] [blame] | 18 | use std::cmp::min; |
Shikha Panwar | 5595711 | 2022-08-22 13:54:33 +0000 | [diff] [blame] | 19 | use std::fs::File; |
Shikha Panwar | 049c28b | 2023-01-10 05:25:05 +0000 | [diff] [blame] | 20 | use std::io::{Read, Seek, SeekFrom, Write}; |
| 21 | use std::mem::size_of; |
Shikha Panwar | 5595711 | 2022-08-22 13:54:33 +0000 | [diff] [blame] | 22 | use std::path::PathBuf; |
Shikha Panwar | 5595711 | 2022-08-22 13:54:33 +0000 | [diff] [blame] | 23 | |
| 24 | const FOOTER_ALIGNMENT: usize = 4; |
| 25 | const ZEROS: [u8; 4] = [0u8; 4_usize]; |
Shikha Panwar | 049c28b | 2023-01-10 05:25:05 +0000 | [diff] [blame] | 26 | const BOOTCONFIG_MAGIC: &str = "#BOOTCONFIG\n"; |
| 27 | // Footer includes [size(le32)][checksum(le32)][#BOOTCONFIG\n] at the end of bootconfigs. |
| 28 | const INITRD_FOOTER_LEN: usize = 2 * std::mem::size_of::<u32>() + BOOTCONFIG_MAGIC.len(); |
Shikha Panwar | 5595711 | 2022-08-22 13:54:33 +0000 | [diff] [blame] | 29 | |
Victor Hsieh | b5bcfab | 2022-09-12 13:06:26 -0700 | [diff] [blame] | 30 | #[derive(Parser, Debug)] |
Shikha Panwar | 049c28b | 2023-01-10 05:25:05 +0000 | [diff] [blame] | 31 | enum Opt { |
| 32 | /// Append bootconfig(s) to initrd image |
| 33 | Attach { |
| 34 | /// Initrd (without bootconfigs) <- Input |
| 35 | initrd: PathBuf, |
| 36 | /// Bootconfigs <- Input |
| 37 | bootconfigs: Vec<PathBuf>, |
| 38 | /// Initrd (with bootconfigs) <- Output |
| 39 | #[clap(long = "output")] |
| 40 | output: PathBuf, |
| 41 | }, |
| 42 | |
| 43 | /// Detach the initrd & bootconfigs - this is required for cases when we update |
| 44 | /// bootconfigs in sign_virt_apex |
| 45 | Detach { |
| 46 | /// Initrd (with bootconfigs) <- Input |
| 47 | initrd_with_bootconfigs: PathBuf, |
| 48 | /// Initrd (without bootconfigs) <- Output |
| 49 | initrd: PathBuf, |
| 50 | /// Bootconfigs <- Output |
| 51 | bootconfigs: PathBuf, |
| 52 | }, |
Shikha Panwar | 5595711 | 2022-08-22 13:54:33 +0000 | [diff] [blame] | 53 | } |
| 54 | |
| 55 | fn get_checksum(file_path: &PathBuf) -> Result<u32> { |
| 56 | File::open(file_path)?.bytes().map(|x| Ok(x? as u32)).sum() |
| 57 | } |
| 58 | |
Shikha Panwar | 049c28b | 2023-01-10 05:25:05 +0000 | [diff] [blame] | 59 | // Copy n bytes of file_in to file_out. Note: copying starts from the current cursors of files. |
| 60 | // On successful return, the files' cursors would have moved forward by k bytes. |
| 61 | fn copyfile2file(file_in: &mut File, file_out: &mut File, n: usize) -> Result<()> { |
| 62 | let mut buf = vec![0; 1024]; |
| 63 | let mut copied: usize = 0; |
| 64 | while copied < n { |
| 65 | let k = min(n - copied, buf.len()); |
| 66 | file_in.read_exact(&mut buf[..k])?; |
| 67 | file_out.write_all(&buf[..k])?; |
| 68 | copied += k; |
| 69 | } |
| 70 | Ok(()) |
| 71 | } |
| 72 | |
| 73 | // Note: attaching & then detaching bootconfigs can lead to extra padding in bootconfigs |
| 74 | fn detach_bootconfig(initrd_bc: PathBuf, initrd: PathBuf, bootconfig: PathBuf) -> Result<()> { |
Charisee | 804381c | 2023-01-30 23:38:54 +0000 | [diff] [blame] | 75 | let mut initrd_bc = File::open(initrd_bc)?; |
| 76 | let mut bootconfig = File::create(bootconfig)?; |
| 77 | let mut initrd = File::create(initrd)?; |
Shikha Panwar | 049c28b | 2023-01-10 05:25:05 +0000 | [diff] [blame] | 78 | let initrd_bc_size: usize = initrd_bc.metadata()?.len().try_into()?; |
| 79 | |
| 80 | initrd_bc.seek(SeekFrom::End(-(BOOTCONFIG_MAGIC.len() as i64)))?; |
| 81 | let mut magic_buf = [0; BOOTCONFIG_MAGIC.len()]; |
| 82 | initrd_bc.read_exact(&mut magic_buf)?; |
| 83 | if magic_buf != BOOTCONFIG_MAGIC.as_bytes() { |
| 84 | bail!("BOOTCONFIG_MAGIC not found in initrd. Bootconfigs might not be attached correctly"); |
| 85 | } |
| 86 | let mut size_buf = [0; size_of::<u32>()]; |
| 87 | initrd_bc.seek(SeekFrom::End(-(INITRD_FOOTER_LEN as i64)))?; |
| 88 | initrd_bc.read_exact(&mut size_buf)?; |
| 89 | let bc_size: usize = u32::from_le_bytes(size_buf) as usize; |
| 90 | |
| 91 | let initrd_size: usize = initrd_bc_size - bc_size - INITRD_FOOTER_LEN; |
| 92 | |
Chris Wailes | 738be14 | 2023-02-14 16:06:32 -0800 | [diff] [blame] | 93 | initrd_bc.rewind()?; |
Shikha Panwar | 049c28b | 2023-01-10 05:25:05 +0000 | [diff] [blame] | 94 | copyfile2file(&mut initrd_bc, &mut initrd, initrd_size)?; |
| 95 | copyfile2file(&mut initrd_bc, &mut bootconfig, bc_size)?; |
| 96 | Ok(()) |
| 97 | } |
| 98 | |
Shikha Panwar | 5595711 | 2022-08-22 13:54:33 +0000 | [diff] [blame] | 99 | // Bootconfig is attached to the initrd in the following way: |
| 100 | // [initrd][bootconfig][padding][size(le32)][checksum(le32)][#BOOTCONFIG\n] |
| 101 | fn attach_bootconfig(initrd: PathBuf, bootconfigs: Vec<PathBuf>, output: PathBuf) -> Result<()> { |
Charisee | 96113f3 | 2023-01-26 09:00:42 +0000 | [diff] [blame] | 102 | let mut output_file = File::create(output)?; |
| 103 | let mut initrd_file = File::open(initrd)?; |
Shikha Panwar | 5595711 | 2022-08-22 13:54:33 +0000 | [diff] [blame] | 104 | let initrd_size: usize = initrd_file.metadata()?.len().try_into()?; |
| 105 | let mut bootconfig_size: usize = 0; |
| 106 | let mut checksum: u32 = 0; |
| 107 | |
| 108 | std::io::copy(&mut initrd_file, &mut output_file)?; |
| 109 | for bootconfig in bootconfigs { |
| 110 | let mut bootconfig_file = File::open(&bootconfig)?; |
| 111 | std::io::copy(&mut bootconfig_file, &mut output_file)?; |
| 112 | bootconfig_size += bootconfig_file.metadata()?.len() as usize; |
| 113 | checksum += get_checksum(&bootconfig)?; |
| 114 | } |
| 115 | |
Shikha Panwar | 58fdf03 | 2023-01-10 05:36:20 +0000 | [diff] [blame] | 116 | let padding_size: usize = |
| 117 | (FOOTER_ALIGNMENT - (initrd_size + bootconfig_size) % FOOTER_ALIGNMENT) % FOOTER_ALIGNMENT; |
Shikha Panwar | 5595711 | 2022-08-22 13:54:33 +0000 | [diff] [blame] | 118 | output_file.write_all(&ZEROS[..padding_size])?; |
| 119 | output_file.write_all(&((padding_size + bootconfig_size) as u32).to_le_bytes())?; |
| 120 | output_file.write_all(&checksum.to_le_bytes())?; |
Shikha Panwar | 049c28b | 2023-01-10 05:25:05 +0000 | [diff] [blame] | 121 | output_file.write_all(BOOTCONFIG_MAGIC.as_bytes())?; |
Shikha Panwar | 5595711 | 2022-08-22 13:54:33 +0000 | [diff] [blame] | 122 | output_file.flush()?; |
| 123 | Ok(()) |
| 124 | } |
| 125 | |
| 126 | fn try_main() -> Result<()> { |
Shikha Panwar | 049c28b | 2023-01-10 05:25:05 +0000 | [diff] [blame] | 127 | let args = Opt::parse(); |
| 128 | match args { |
| 129 | Opt::Attach { initrd, bootconfigs, output } => { |
| 130 | attach_bootconfig(initrd, bootconfigs, output)? |
| 131 | } |
| 132 | Opt::Detach { initrd_with_bootconfigs, initrd, bootconfigs } => { |
| 133 | detach_bootconfig(initrd_with_bootconfigs, initrd, bootconfigs)? |
| 134 | } |
| 135 | }; |
Shikha Panwar | 5595711 | 2022-08-22 13:54:33 +0000 | [diff] [blame] | 136 | Ok(()) |
| 137 | } |
| 138 | |
| 139 | fn main() { |
| 140 | try_main().unwrap() |
| 141 | } |
Andrew Walbran | da8786d | 2022-12-01 14:54:27 +0000 | [diff] [blame] | 142 | |
| 143 | #[cfg(test)] |
| 144 | mod tests { |
| 145 | use super::*; |
| 146 | use clap::CommandFactory; |
| 147 | |
| 148 | #[test] |
| 149 | fn verify_args() { |
| 150 | // Check that the command parsing has been configured in a valid way. |
Shikha Panwar | 049c28b | 2023-01-10 05:25:05 +0000 | [diff] [blame] | 151 | Opt::command().debug_assert(); |
Andrew Walbran | da8786d | 2022-12-01 14:54:27 +0000 | [diff] [blame] | 152 | } |
| 153 | } |