|
| 1 | +// This Source Code Form is subject to the terms of the Mozilla Public |
| 2 | +// License, v. 2.0. If a copy of the MPL was not distributed with this |
| 3 | +// file, You can obtain one at https://mozilla.org/MPL/2.0/. |
| 4 | + |
| 5 | +use anyhow::{Context, bail}; |
| 6 | +use camino::{Utf8DirEntry, Utf8Path}; |
| 7 | +use clap::{Parser, Subcommand}; |
| 8 | +use iddqd::IdOrdMap; |
| 9 | +use omicron_common::update::{OmicronInstallManifest, OmicronInstallMetadata}; |
| 10 | +use sha2::{Digest, Sha256}; |
| 11 | +use std::fs::File; |
| 12 | +use tufaceous_artifact::ArtifactHash; |
| 13 | +use uuid::Uuid; |
| 14 | + |
| 15 | +/// Return a UUID if the `DirEntry` contains a directory that parses into a UUID. |
| 16 | +fn get_uuid_dir(result: std::io::Result<Utf8DirEntry>) -> Option<Uuid> { |
| 17 | + let Ok(entry) = result else { |
| 18 | + return None; |
| 19 | + }; |
| 20 | + let Ok(file_type) = entry.file_type() else { |
| 21 | + return None; |
| 22 | + }; |
| 23 | + if !file_type.is_dir() { |
| 24 | + return None; |
| 25 | + } |
| 26 | + let file_name = entry.file_name(); |
| 27 | + file_name.parse().ok() |
| 28 | +} |
| 29 | + |
| 30 | +#[derive(Debug)] |
| 31 | +pub struct Pools { |
| 32 | + pub internal: Vec<Uuid>, |
| 33 | + pub external: Vec<Uuid>, |
| 34 | +} |
| 35 | + |
| 36 | +impl Pools { |
| 37 | + pub fn read() -> anyhow::Result<Pools> { |
| 38 | + let internal = Utf8Path::new("/pool/int/") |
| 39 | + .read_dir_utf8() |
| 40 | + .context("Failed to read /pool/int")? |
| 41 | + .filter_map(get_uuid_dir) |
| 42 | + .collect(); |
| 43 | + let external = Utf8Path::new("/pool/ext/") |
| 44 | + .read_dir_utf8() |
| 45 | + .context("Failed to read /pool/ext")? |
| 46 | + .filter_map(get_uuid_dir) |
| 47 | + .collect(); |
| 48 | + Ok(Pools { internal, external }) |
| 49 | + } |
| 50 | +} |
| 51 | + |
| 52 | +#[derive(Subcommand)] |
| 53 | +enum CorpusCommand { |
| 54 | + /// Verify that all corpus files are present and accounted for |
| 55 | + /// on all disks |
| 56 | + Check, |
| 57 | + /// Regenerate the `measurements.json` file on all disks. This will |
| 58 | + /// be saved as `measurements-regenerated.json`. It is up to you to |
| 59 | + /// do the replacement! |
| 60 | + Regenerate, |
| 61 | +} |
| 62 | + |
| 63 | +/// Adjust the `measurement.json` in an install dataset |
| 64 | +#[derive(Parser)] |
| 65 | +struct CorpusAdjust { |
| 66 | + #[clap(subcommand)] |
| 67 | + command: CorpusCommand, |
| 68 | +} |
| 69 | + |
| 70 | +fn measurement_names( |
| 71 | + dir: &Utf8Path, |
| 72 | +) -> anyhow::Result<IdOrdMap<OmicronInstallMetadata>> { |
| 73 | + if !dir.is_dir() { |
| 74 | + bail!("Not a directory"); |
| 75 | + } |
| 76 | + |
| 77 | + let mut files = IdOrdMap::new(); |
| 78 | + |
| 79 | + for entry in dir.read_dir_utf8()? { |
| 80 | + let entry = entry?; |
| 81 | + |
| 82 | + if !entry.file_name().contains("measurements.json") { |
| 83 | + continue; |
| 84 | + } |
| 85 | + |
| 86 | + let mut hasher = Sha256::new(); |
| 87 | + std::io::copy(&mut File::open(&entry.path())?, &mut hasher)?; |
| 88 | + |
| 89 | + let meta = entry.metadata()?; |
| 90 | + |
| 91 | + let entry = OmicronInstallMetadata { |
| 92 | + file_name: entry.file_name().to_owned(), |
| 93 | + file_size: meta.len(), |
| 94 | + hash: ArtifactHash(hasher.finalize().into()), |
| 95 | + }; |
| 96 | + |
| 97 | + files |
| 98 | + .insert_unique(entry) |
| 99 | + .map_err(|_| anyhow::anyhow!("duplicate?"))?; |
| 100 | + } |
| 101 | + |
| 102 | + Ok(files) |
| 103 | +} |
| 104 | + |
| 105 | +fn main() -> anyhow::Result<()> { |
| 106 | + let arg = CorpusAdjust::parse(); |
| 107 | + |
| 108 | + match arg.command { |
| 109 | + CorpusCommand::Check => { |
| 110 | + let pools = Pools::read()?; |
| 111 | + |
| 112 | + let all = pools.internal.into_iter().map(|u| { |
| 113 | + Utf8Path::new("/pool/int") |
| 114 | + .join(format!("{u}")) |
| 115 | + .join("install") |
| 116 | + .join("measurements") |
| 117 | + }); |
| 118 | + for p in all { |
| 119 | + let mut names = measurement_names(&p)?; |
| 120 | + |
| 121 | + let f = p.join("measurements.json"); |
| 122 | + let bytes = std::fs::read(f)?; |
| 123 | + |
| 124 | + let manifest = |
| 125 | + serde_json::from_slice::<OmicronInstallManifest>(&bytes)?; |
| 126 | + |
| 127 | + for f in manifest.files { |
| 128 | + if let Some(entry) = names.remove(f.file_name.as_str()) { |
| 129 | + if entry != f { |
| 130 | + bail!( |
| 131 | + "file {} in manifest does not match disk {p}", |
| 132 | + f.file_name |
| 133 | + ); |
| 134 | + } |
| 135 | + } else { |
| 136 | + bail!( |
| 137 | + "file {} in manifest but not found on disk {p}", |
| 138 | + f.file_name |
| 139 | + ); |
| 140 | + } |
| 141 | + } |
| 142 | + |
| 143 | + if !names.is_empty() { |
| 144 | + println!("some files were not in the manifest {p}:"); |
| 145 | + for n in names { |
| 146 | + println!("{}", n.file_name); |
| 147 | + } |
| 148 | + bail!("Fix your manifest"); |
| 149 | + } |
| 150 | + } |
| 151 | + } |
| 152 | + CorpusCommand::Regenerate => { |
| 153 | + let pools = Pools::read()?; |
| 154 | + |
| 155 | + let all = pools.internal.into_iter().map(|u| { |
| 156 | + Utf8Path::new("/pool/int") |
| 157 | + .join(format!("{u}")) |
| 158 | + .join("install") |
| 159 | + .join("measurements") |
| 160 | + }); |
| 161 | + for p in all { |
| 162 | + let names = measurement_names(&p)?; |
| 163 | + |
| 164 | + let f = p.join("measurements.json"); |
| 165 | + let bytes = std::fs::read(f)?; |
| 166 | + |
| 167 | + let manifest = |
| 168 | + serde_json::from_slice::<OmicronInstallManifest>(&bytes)?; |
| 169 | + |
| 170 | + let updated = OmicronInstallManifest { |
| 171 | + source: manifest.source, |
| 172 | + files: names, |
| 173 | + }; |
| 174 | + |
| 175 | + let json = serde_json::to_vec(&updated)?; |
| 176 | + |
| 177 | + std::fs::write(p.join("measurements-regenerated.json"), &json)?; |
| 178 | + } |
| 179 | + |
| 180 | + println!("done"); |
| 181 | + } |
| 182 | + } |
| 183 | + |
| 184 | + Ok(()) |
| 185 | +} |
0 commit comments