use anyhow::{anyhow, bail, Context, Result};
use byte_unit::Byte;
use nix::unistd::isatty;
use reqwest::Url;
use std::fs::{remove_file, File, OpenOptions};
use std::io::{self, copy, stderr, BufReader, BufWriter, Cursor, Read, Seek, SeekFrom, Write};
use std::num::{NonZeroU32, NonZeroU64};
use std::os::unix::io::AsRawFd;
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};
use crate::blockdev::{detect_formatted_sector_size, get_gpt_size, SavedPartitions};
use crate::cmdline::*;
use crate::io::*;
use crate::source::*;
pub fn download(config: DownloadConfig) -> Result<()> {
let location: Box<dyn ImageLocation> = if let Some(image_url) = &config.image_url {
Box::new(UrlLocation::new(image_url, config.fetch_retries))
} else {
Box::new(StreamLocation::new(
&config.stream,
config.architecture.as_str(),
&config.platform,
&config.format,
config.stream_base_url.as_ref(),
config.fetch_retries,
)?)
};
eprintln!("{}", location);
let mut sources = location.sources()?;
if sources.is_empty() {
bail!("no artifacts found");
}
for source in sources.iter_mut() {
if source.signature.is_none() {
if config.insecure {
eprintln!("Signature not found; skipping verification as requested");
} else {
bail!("--insecure not specified and signature not found");
}
}
let (decompress, filename) = should_decompress(config.decompress, &source.filename);
let mut path = PathBuf::new();
path.push(&config.directory);
path.push(filename);
let sig_path = path.with_file_name(format!("{}.sig", filename));
if !decompress
&& check_image_and_sig(source, &path, &sig_path, VerifyKeys::Production).is_ok()
{
println!("{}", path.display());
continue;
}
if let Err(err) = write_image_and_sig(
source,
&path,
&sig_path,
decompress,
!config.decompress,
VerifyKeys::Production,
) {
let _ = remove_file(&path);
let _ = remove_file(&sig_path);
return Err(err);
}
println!("{}", path.display());
}
Ok(())
}
fn should_decompress(enabled: bool, filename: &str) -> (bool, &str) {
#[allow(clippy::if_same_then_else)] if !enabled {
(false, filename)
} else if filename.ends_with(".tar.gz") || filename.ends_with(".tar.xz") {
(false, filename)
} else if filename.ends_with(".gz") {
(true, filename.trim_end_matches(".gz"))
} else if filename.ends_with(".xz") {
(true, filename.trim_end_matches(".xz"))
} else {
(false, filename)
}
}
fn check_image_and_sig(
source: &ImageSource,
path: &Path,
sig_path: &Path,
keys: VerifyKeys,
) -> Result<()> {
if source.signature.is_none() {
bail!("no signature available; can't check existing file");
}
let signature = source.signature.as_ref().unwrap();
let mut sig_file = OpenOptions::new()
.read(true)
.open(sig_path)
.with_context(|| format!("opening {}", sig_path.display()))?;
let mut buf = Vec::new();
sig_file
.read_to_end(&mut buf)
.with_context(|| format!("reading {}", sig_path.display()))?;
if &buf != signature {
bail!("signature file doesn't match source");
}
let mut file = OpenOptions::new()
.read(true)
.open(path)
.with_context(|| format!("opening {}", path.display()))?;
let mut reader = VerifyReader::new(
BufReader::with_capacity(BUFFER_SIZE, &mut file),
Some(signature),
keys,
)?;
copy(&mut reader, &mut io::sink())?;
reader.verify_without_logging_failure()?;
Ok(())
}
fn write_image_and_sig(
source: &mut ImageSource,
path: &Path,
sig_path: &Path,
decompress: bool,
save_sig: bool,
keys: VerifyKeys,
) -> Result<()> {
let mut dest = OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(path)
.with_context(|| format!("opening {}", path.display()))?;
write_image(
source,
&mut dest,
path,
image_copy_default,
decompress,
None,
None,
keys,
)?;
if let (true, Some(signature)) = (save_sig, source.signature.as_ref()) {
let mut sig_dest = OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(sig_path)
.with_context(|| format!("opening {}", sig_path.display()))?;
sig_dest
.write_all(signature)
.context("writing signature data")?;
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
pub fn write_image<F>(
source: &mut ImageSource,
dest: &mut File,
dest_path: &Path,
image_copy: F,
decompress: bool,
saved: Option<&SavedPartitions>,
expected_sector_size: Option<NonZeroU32>,
keys: VerifyKeys,
) -> Result<()>
where
F: FnOnce(&[u8], &mut dyn Read, &mut File, &Path, Option<&SavedPartitions>) -> Result<()>,
{
let mut verify_reader =
VerifyReader::new(&mut source.reader, source.signature.as_deref(), keys)?;
let mut reader: Box<dyn Read> = Box::new(ProgressReader::new(
&mut verify_reader,
source.length_hint,
&source.artifact_type,
));
let peek_reader = PeekReader::with_capacity(BUFFER_SIZE, reader);
if decompress {
reader = Box::new(DecompressReader::new(peek_reader)?);
} else {
reader = Box::new(peek_reader);
}
let byte_limit = saved.map(|saved| saved.get_offset()).transpose()?.flatten();
if let Some((limit, conflict)) = byte_limit {
reader = Box::new(LimitReader::new(reader, limit, conflict));
}
let mut first_mb = [0u8; 1024 * 1024];
reader
.read_exact(&mut first_mb)
.context("decoding first MiB of image")?;
if let Some(expected) = expected_sector_size {
if let Some(actual) = detect_formatted_sector_size(&first_mb) {
if expected != actual {
bail!(
"source has sector size {} but destination has sector size {}",
actual.get(),
expected.get()
);
}
}
}
image_copy(&first_mb, &mut reader, dest, dest_path, saved)?;
drop(reader);
verify_reader.verify()?;
dest.sync_all().context("syncing data to disk")?;
Ok(())
}
pub fn image_copy_default(
first_mb: &[u8],
source: &mut dyn Read,
dest: &mut File,
_dest_path: &Path,
saved: Option<&SavedPartitions>,
) -> Result<()> {
match saved {
Some(saved) => {
saved
.overwrite(dest)
.context("overwriting disk partition table")?;
dest.seek(SeekFrom::Start(1024 * 1024))
.context("seeking disk")?;
}
None => dest
.write_all(&[0u8; 1024 * 1024])
.context("clearing first MiB of disk")?,
};
dest.sync_all().context("syncing data to disk")?;
let mut buf_dest = BufWriter::with_capacity(BUFFER_SIZE, dest);
copy(source, &mut buf_dest).context("decoding and writing image")?;
let dest = buf_dest
.into_inner()
.map_err(|_| anyhow!("flushing data to disk"))?;
let offset = match saved {
Some(saved) if saved.is_saved() => {
dest.seek(SeekFrom::Start(0))
.context("seeking disk to MBR")?;
dest.write_all(&first_mb[0..saved.get_sector_size() as usize])
.context("writing MBR")?;
let mut cursor = Cursor::new(first_mb);
saved
.merge(&mut cursor, dest)
.context("writing updated GPT")?;
get_gpt_size(dest).context("getting GPT size")?
}
_ => {
0
}
};
dest.seek(SeekFrom::Start(offset))
.with_context(|| format!("seeking disk to offset {}", offset))?;
dest.write_all(&first_mb[offset as usize..first_mb.len()])
.context("writing first MiB of disk")?;
Ok(())
}
pub fn download_to_tempfile(url: &Url, retries: FetchRetries) -> Result<File> {
let mut f = tempfile::tempfile()?;
let client = new_http_client()?;
let mut resp = http_get(client, url, retries)?;
let mut writer = BufWriter::with_capacity(BUFFER_SIZE, &mut f);
copy(
&mut BufReader::with_capacity(BUFFER_SIZE, &mut resp),
&mut writer,
)
.with_context(|| format!("couldn't copy '{}'", url))?;
writer
.flush()
.with_context(|| format!("couldn't write '{}' to disk", url))?;
drop(writer);
f.seek(SeekFrom::Start(0))
.with_context(|| format!("rewinding file for '{}'", url))?;
Ok(f)
}
struct ProgressReader<'a, R: Read> {
source: R,
length: Option<(NonZeroU64, String)>,
artifact_type: &'a str,
position: u64,
last_report: Instant,
tty: bool,
prologue: &'static str,
epilogue: &'static str,
}
impl<'a, R: Read> ProgressReader<'a, R> {
fn new(source: R, length: Option<u64>, artifact_type: &'a str) -> Self {
let tty = isatty(stderr().as_raw_fd()).unwrap_or_else(|e| {
eprintln!("checking if stderr is a TTY: {}", e);
false
});
let length = length.and_then(NonZeroU64::new);
ProgressReader {
source,
length: length.map(|l| (l, Self::format_bytes(l.get()))),
artifact_type,
position: 0,
last_report: Instant::now(),
tty,
prologue: if tty { "> " } else { "" },
epilogue: if tty { " \r" } else { "\n" },
}
}
fn format_bytes(count: u64) -> String {
Byte::from_bytes(count.into())
.get_appropriate_unit(true)
.format(1)
}
}
impl<'a, R: Read> Read for ProgressReader<'a, R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
let count = self.source.read(buf)?;
self.position += count as u64;
if self.last_report.elapsed() >= Duration::from_secs(1)
|| self.length.as_ref().map(|(l, _)| l.get()) == Some(self.position)
{
self.last_report = Instant::now();
match self.length {
Some((length, ref length_str)) => eprint!(
"{}Read {} {}/{} ({}%){}",
self.prologue,
self.artifact_type,
Self::format_bytes(self.position),
length_str,
100 * self.position / length.get(),
self.epilogue
),
None => eprint!(
"{}Read {} {}{}",
self.prologue,
self.artifact_type,
Self::format_bytes(self.position),
self.epilogue
),
}
let _ = std::io::stdout().flush();
}
Ok(count)
}
}
impl<'a, R: Read> Drop for ProgressReader<'a, R> {
fn drop(&mut self) {
if self.tty {
eprintln!();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use gptman::{GPTPartitionEntry, GPT};
use std::fs::{read, write};
use std::io::{Seek, SeekFrom};
use tempfile::TempDir;
use uuid::Uuid;
#[test]
fn test_signature_checks() {
test_one_signed_file(
&[0; 1 << 20][..],
&include_bytes!("../fixtures/verify/1M.sig")[..],
&[0; 1 << 20][..],
);
test_one_signed_file(
&include_bytes!("../fixtures/verify/1M.gz")[..],
&include_bytes!("../fixtures/verify/1M.gz.sig")[..],
&[0; 1 << 20][..],
);
test_one_signed_file(
&include_bytes!("../fixtures/verify/1M.xz")[..],
&include_bytes!("../fixtures/verify/1M.xz.sig")[..],
&[0; 1 << 20][..],
);
test_one_signed_file(
&include_bytes!("../fixtures/verify/1M.zst")[..],
&include_bytes!("../fixtures/verify/1M.zst.sig")[..],
&[0; 1 << 20][..],
);
}
fn test_one_signed_file(data: &[u8], sig: &[u8], decompressed_data: &[u8]) {
let dir = TempDir::new().unwrap();
let good_path = dir.path().join("good");
write(&good_path, data).unwrap();
let good_sig_path = dir.path().join("good.sig");
write(&good_sig_path, sig).unwrap();
let bad_path = dir.path().join("bad");
let mut bad_data = data.to_vec();
bad_data.push(0);
write(&bad_path, &bad_data).unwrap();
let bad_sig_path = dir.path().join("bad.sig");
write(&bad_sig_path, sig).unwrap();
let source = FileLocation::new(good_path.to_str().unwrap())
.sources()
.unwrap()
.remove(0);
check_image_and_sig(
&source,
&good_path,
&good_sig_path,
VerifyKeys::InsecureTest,
)
.unwrap();
let source = FileLocation::new(bad_path.to_str().unwrap())
.sources()
.unwrap()
.remove(0);
check_image_and_sig(&source, &bad_path, &bad_sig_path, VerifyKeys::InsecureTest)
.unwrap_err();
let mut source = FileLocation::new(good_path.to_str().unwrap())
.sources()
.unwrap()
.remove(0);
let out_path = dir.path().join("out");
let mut out_file = File::create(&out_path).unwrap();
write_image(
&mut source,
&mut out_file,
&out_path,
image_copy_default,
true,
None,
None,
VerifyKeys::InsecureTest,
)
.unwrap();
assert_eq!(&read(&out_path).unwrap(), decompressed_data);
let mut source = FileLocation::new(bad_path.to_str().unwrap())
.sources()
.unwrap()
.remove(0);
let out_path = dir.path().join("out");
let mut out_file = File::create(&out_path).unwrap();
write_image(
&mut source,
&mut out_file,
&out_path,
image_copy_default,
true,
None,
None,
VerifyKeys::InsecureTest,
)
.unwrap_err();
}
#[test]
fn test_should_decompress() {
assert_eq!(should_decompress(true, "foo.img"), (false, "foo.img"));
assert_eq!(should_decompress(true, "foo.bz2"), (false, "foo.bz2"));
assert_eq!(should_decompress(false, "foo.gz"), (false, "foo.gz"));
assert_eq!(should_decompress(true, "foo.gz"), (true, "foo"));
assert_eq!(should_decompress(true, "foo.tar.gz"), (false, "foo.tar.gz"));
assert_eq!(should_decompress(false, "foo.xz"), (false, "foo.xz"));
assert_eq!(should_decompress(true, "foo.xz"), (true, "foo"));
assert_eq!(should_decompress(true, "foo.tar.xz"), (false, "foo.tar.xz"));
}
#[test]
fn test_write_image_limit() {
let (mut source, source_path) = tempfile::Builder::new()
.prefix("coreos-installer-")
.tempfile()
.unwrap()
.into_parts();
source.set_len(6 * 1024 * 1024).unwrap();
partition(&mut source, None);
let (mut dest, dest_path) = tempfile::Builder::new()
.prefix("coreos-installer-")
.tempfile()
.unwrap()
.into_parts();
dest.set_len(8 * 1024 * 1024).unwrap();
partition(&mut dest, Some(4));
let saved = SavedPartitions::new_from_file(
&mut dest,
512,
&vec![PartitionFilter::Label(glob::Pattern::new("*").unwrap())],
)
.unwrap();
assert!(saved.is_saved());
let offset = 4 * 1024 * 1024;
let precious = "hello world";
dest.seek(SeekFrom::Start(offset)).unwrap();
dest.write_all(precious.as_bytes()).unwrap();
dest.seek(SeekFrom::Start(0)).unwrap();
let err = write_image(
&mut FileLocation::new(source_path.to_str().unwrap())
.sources()
.unwrap()
.remove(0),
&mut dest,
&dest_path,
image_copy_default,
false,
Some(&saved),
None,
VerifyKeys::InsecureTest,
)
.unwrap_err();
assert!(
format!("{:#}", err).contains("collision with partition"),
"incorrect error: {:#}",
err
);
dest.seek(SeekFrom::Start(offset)).unwrap();
let mut buf = vec![0u8; precious.len()];
dest.read_exact(&mut buf).unwrap();
assert_eq!(buf, precious.as_bytes());
}
#[test]
fn test_image_copy_default_first_mb() {
let len: usize = 2 * 1024 * 1024;
let mb: usize = 1024 * 1024;
let mut data = vec![0u8; len];
for i in 0..data.len() {
data[i] = (i % 256) as u8;
}
let mut source = Cursor::new(&data);
let mut dest = tempfile::tempfile().unwrap();
source.seek(SeekFrom::Start(mb as u64)).unwrap();
image_copy_default(&data[0..mb], &mut source, &mut dest, Path::new("/z"), None).unwrap();
dest.seek(SeekFrom::Start(0)).unwrap();
let mut result = vec![0u8; len];
dest.read_exact(&mut result).unwrap();
assert_eq!(data, result);
let mut source = Cursor::new(&data);
let mut dest = tempfile::tempfile().unwrap();
dest.set_len(len as u64).unwrap();
let saved = SavedPartitions::new_from_file(&mut dest, 512, &vec![]).unwrap();
assert!(!saved.is_saved());
source.seek(SeekFrom::Start(mb as u64)).unwrap();
image_copy_default(
&data[0..mb],
&mut source,
&mut dest,
Path::new("/z"),
Some(&saved),
)
.unwrap();
dest.seek(SeekFrom::Start(0)).unwrap();
let mut result = vec![0u8; len];
dest.read_exact(&mut result).unwrap();
assert_eq!(data, result);
let mut source = Cursor::new(data.clone());
let mut dest = tempfile::tempfile().unwrap();
partition(&mut source, None);
let data_partitioned = source.into_inner();
let mut source = Cursor::new(&data_partitioned);
dest.set_len(2 * len as u64).unwrap();
partition(&mut dest, Some(2));
let saved = SavedPartitions::new_from_file(
&mut dest,
512,
&vec![PartitionFilter::Label(glob::Pattern::new("bovik").unwrap())],
)
.unwrap();
assert!(saved.is_saved());
source.seek(SeekFrom::Start(mb as u64)).unwrap();
image_copy_default(
&data_partitioned[0..mb],
&mut source,
&mut dest,
Path::new("/z"),
Some(&saved),
)
.unwrap();
dest.seek(SeekFrom::Start(0)).unwrap();
let mut result = vec![0u8; len];
dest.read_exact(&mut result).unwrap();
assert_eq!(detect_formatted_sector_size(&result), NonZeroU32::new(512));
assert_eq!(data_partitioned[0..446], result[0..446]);
let gpt_size = get_gpt_size(&mut dest).unwrap();
assert!(gpt_size < 24576);
assert_eq!(
data_partitioned[gpt_size as usize..],
result[gpt_size as usize..]
);
}
fn partition(f: &mut (impl Read + Write + Seek), start_mb: Option<u64>) {
let mut gpt = GPT::new_from(f, 512, *Uuid::new_v4().as_bytes()).unwrap();
if let Some(start_mb) = start_mb {
gpt[1] = GPTPartitionEntry {
partition_type_guid: [1u8; 16],
unique_partition_guid: [1u8; 16],
starting_lba: start_mb * 2048,
ending_lba: (start_mb + 1) * 2048,
attribute_bits: 0,
partition_name: "bovik".into(),
};
}
gpt.write_into(f).unwrap();
}
}