blob: 8f06f093c94c38d396192d8ace9f8e2b4277322f [file] [log] [blame]
Shikha Panwar55957112022-08-22 13:54:33 +00001// 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 Panwar049c28b2023-01-10 05:25:05 +000015//! Attach/Detach bootconfigs to initrd image
16use anyhow::{bail, Result};
Victor Hsiehb5bcfab2022-09-12 13:06:26 -070017use clap::Parser;
Shikha Panwar049c28b2023-01-10 05:25:05 +000018use std::cmp::min;
Shikha Panwar55957112022-08-22 13:54:33 +000019use std::fs::File;
Shikha Panwar049c28b2023-01-10 05:25:05 +000020use std::io::{Read, Seek, SeekFrom, Write};
21use std::mem::size_of;
Shikha Panwar55957112022-08-22 13:54:33 +000022use std::path::PathBuf;
Shikha Panwar55957112022-08-22 13:54:33 +000023
24const FOOTER_ALIGNMENT: usize = 4;
25const ZEROS: [u8; 4] = [0u8; 4_usize];
Shikha Panwar049c28b2023-01-10 05:25:05 +000026const BOOTCONFIG_MAGIC: &str = "#BOOTCONFIG\n";
27// Footer includes [size(le32)][checksum(le32)][#BOOTCONFIG\n] at the end of bootconfigs.
28const INITRD_FOOTER_LEN: usize = 2 * std::mem::size_of::<u32>() + BOOTCONFIG_MAGIC.len();
Shikha Panwar55957112022-08-22 13:54:33 +000029
Victor Hsiehb5bcfab2022-09-12 13:06:26 -070030#[derive(Parser, Debug)]
Shikha Panwar049c28b2023-01-10 05:25:05 +000031enum 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 Panwar55957112022-08-22 13:54:33 +000053}
54
55fn get_checksum(file_path: &PathBuf) -> Result<u32> {
56 File::open(file_path)?.bytes().map(|x| Ok(x? as u32)).sum()
57}
58
Shikha Panwar049c28b2023-01-10 05:25:05 +000059// 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.
61fn 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
74fn detach_bootconfig(initrd_bc: PathBuf, initrd: PathBuf, bootconfig: PathBuf) -> Result<()> {
Charisee804381c2023-01-30 23:38:54 +000075 let mut initrd_bc = File::open(initrd_bc)?;
76 let mut bootconfig = File::create(bootconfig)?;
77 let mut initrd = File::create(initrd)?;
Shikha Panwar049c28b2023-01-10 05:25:05 +000078 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 Wailes738be142023-02-14 16:06:32 -080093 initrd_bc.rewind()?;
Shikha Panwar049c28b2023-01-10 05:25:05 +000094 copyfile2file(&mut initrd_bc, &mut initrd, initrd_size)?;
95 copyfile2file(&mut initrd_bc, &mut bootconfig, bc_size)?;
96 Ok(())
97}
98
Shikha Panwar55957112022-08-22 13:54:33 +000099// Bootconfig is attached to the initrd in the following way:
100// [initrd][bootconfig][padding][size(le32)][checksum(le32)][#BOOTCONFIG\n]
101fn attach_bootconfig(initrd: PathBuf, bootconfigs: Vec<PathBuf>, output: PathBuf) -> Result<()> {
Charisee96113f32023-01-26 09:00:42 +0000102 let mut output_file = File::create(output)?;
103 let mut initrd_file = File::open(initrd)?;
Shikha Panwar55957112022-08-22 13:54:33 +0000104 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 Panwar58fdf032023-01-10 05:36:20 +0000116 let padding_size: usize =
117 (FOOTER_ALIGNMENT - (initrd_size + bootconfig_size) % FOOTER_ALIGNMENT) % FOOTER_ALIGNMENT;
Shikha Panwar55957112022-08-22 13:54:33 +0000118 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 Panwar049c28b2023-01-10 05:25:05 +0000121 output_file.write_all(BOOTCONFIG_MAGIC.as_bytes())?;
Shikha Panwar55957112022-08-22 13:54:33 +0000122 output_file.flush()?;
123 Ok(())
124}
125
126fn try_main() -> Result<()> {
Shikha Panwar049c28b2023-01-10 05:25:05 +0000127 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 Panwar55957112022-08-22 13:54:33 +0000136 Ok(())
137}
138
139fn main() {
140 try_main().unwrap()
141}
Andrew Walbranda8786d2022-12-01 14:54:27 +0000142
143#[cfg(test)]
144mod 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 Panwar049c28b2023-01-10 05:25:05 +0000151 Opt::command().debug_assert();
Andrew Walbranda8786d2022-12-01 14:54:27 +0000152 }
153}