From 802f374a90d8202785a34648c39aa6320ac171d9 Mon Sep 17 00:00:00 2001 From: Florian Klink Date: Tue, 2 Jan 2024 14:17:02 +0200 Subject: feat(tvix/glue): handle passAsFile This extends derivation_to_build_request to handle passAsFile the same way Nix does, and adds a unit test for it. I opted to making this function fallible (if passAsFile contains a non-existent env var), rather than pushing all of this into the Derivation validate function. Change-Id: I75b635f1f6c0c78d72b9a8fc7824f77e97b69951 Reviewed-on: https://cl.tvl.fyi/c/depot/+/10522 Tested-by: BuildkiteCI Autosubmit: flokli Reviewed-by: raitobezarius --- tvix/Cargo.lock | 9 ++- tvix/Cargo.nix | 12 ++- tvix/glue/Cargo.toml | 1 + tvix/glue/src/tvix_build.rs | 178 +++++++++++++++++++++++++++++++++++++++++--- 4 files changed, 181 insertions(+), 19 deletions(-) (limited to 'tvix') diff --git a/tvix/Cargo.lock b/tvix/Cargo.lock index abddd5b4a553..8a12ab6726b8 100644 --- a/tvix/Cargo.lock +++ b/tvix/Cargo.lock @@ -667,9 +667,9 @@ checksum = "56254986775e3233ffa9c4d7d3faaf6d36a2c09d30b20687e9f88bc8bafc16c8" [[package]] name = "digest" -version = "0.10.6" +version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8168378f4e5023e7218c89c891c0fd8ecdb5e5e4f18cb78f38cf245dd021e76f" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ "block-buffer", "crypto-common", @@ -2565,9 +2565,9 @@ dependencies = [ [[package]] name = "sha2" -version = "0.10.6" +version = "0.10.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "82e6b795fe2e3b1e845bafcb27aa35405c4d47cdfc92af5fc8d3002f76cebdc0" +checksum = "793db75ad2bcafc3ffa7c68b215fee268f537982cd901d132f89c6343f3a3dc8" dependencies = [ "cfg-if", "cpufeatures", @@ -3408,6 +3408,7 @@ dependencies = [ "data-encoding", "lazy_static", "nix-compat", + "sha2", "tempfile", "test-case", "thiserror", diff --git a/tvix/Cargo.nix b/tvix/Cargo.nix index 69b9cc9a8d97..1ad3422d3094 100644 --- a/tvix/Cargo.nix +++ b/tvix/Cargo.nix @@ -1971,9 +1971,9 @@ rec { }; "digest" = rec { crateName = "digest"; - version = "0.10.6"; + version = "0.10.7"; edition = "2018"; - sha256 = "0vz74785s96g727vg37iwkjvbkcfzp093j49ihhyf8sh9s7kfs41"; + sha256 = "14p2n6ih29x81akj097lvz7wi9b6b9hvls0lwrv7b6xwyy0s5ncy"; authors = [ "RustCrypto Developers" ]; @@ -7873,9 +7873,9 @@ rec { }; "sha2" = rec { crateName = "sha2"; - version = "0.10.6"; + version = "0.10.8"; edition = "2018"; - sha256 = "1h5xrrv2y06kr1gsz4pwrm3lsp206nm2gjxgbf21wfrfzsavgrl2"; + sha256 = "1j1x78zk9il95w9iv46dh9wm73r6xrgj32y6lzzw7bxws9dbfgbr"; authors = [ "RustCrypto Developers" ]; @@ -10727,6 +10727,10 @@ rec { name = "nix-compat"; packageId = "nix-compat"; } + { + name = "sha2"; + packageId = "sha2"; + } { name = "thiserror"; packageId = "thiserror"; diff --git a/tvix/glue/Cargo.toml b/tvix/glue/Cargo.toml index 6b0edb5742e0..80a7d1b63f96 100644 --- a/tvix/glue/Cargo.toml +++ b/tvix/glue/Cargo.toml @@ -14,6 +14,7 @@ bytes = "1.4.0" tracing = "0.1.37" tokio = "1.28.0" thiserror = "1.0.38" +sha2 = "0.10.8" [dependencies.wu-manber] git = "https://github.com/tvlfyi/wu-manber.git" diff --git a/tvix/glue/src/tvix_build.rs b/tvix/glue/src/tvix_build.rs index 399b451c4c71..3f0e92e2139d 100644 --- a/tvix/glue/src/tvix_build.rs +++ b/tvix/glue/src/tvix_build.rs @@ -4,9 +4,10 @@ use std::collections::BTreeMap; use bytes::Bytes; -use nix_compat::{derivation::Derivation, store_path::StorePathRef}; +use nix_compat::{derivation::Derivation, nixbase32, store_path::StorePathRef}; +use sha2::{Digest, Sha256}; use tvix_build::proto::{ - build_request::{BuildConstraints, EnvVar}, + build_request::{AdditionalFile, BuildConstraints, EnvVar}, BuildRequest, }; use tvix_castore::proto::{NamedNode, Node}; @@ -40,7 +41,7 @@ fn derivation_to_build_request( derivation: &Derivation, fn_input_sources_to_node: FIS, fn_input_drvs_to_node: FID, -) -> BuildRequest +) -> std::io::Result where FIS: Fn(StorePathRef) -> Node, FID: Fn(StorePathRef, &[&str]) -> Node, @@ -62,9 +63,11 @@ where // Sort the outputs. We can use sort_unstable, as these are unique strings. output_paths.sort_unstable(); - // Produce environment_vars. We use a BTreeMap while producing them, so the - // resulting Vec is sorted by key. + // Produce environment_vars and additional files. + // We use a BTreeMap while producing, and only realize the resulting Vec + // while populating BuildRequest, so we don't need to worry about ordering. let mut environment_vars: BTreeMap = BTreeMap::new(); + let mut additional_files: BTreeMap = BTreeMap::new(); // Start with some the ones that nix magically sets: environment_vars.extend( @@ -83,6 +86,9 @@ where .map(|(k, v)| (k.clone(), Bytes::from(v.to_vec()))), ); + handle_pass_as_file(&mut environment_vars, &mut additional_files)?; + // TODO: handle structuredAttrs. + // Produce inputs. As we refer to the contents here, not just plain store path strings, // we need to perform lookups. // FUTUREWORK: should we also model input_derivations and input_sources with StorePath? @@ -135,7 +141,7 @@ where provide_bin_sh: true, }); - BuildRequest { + Ok(BuildRequest { command_args, outputs: output_paths, @@ -150,9 +156,70 @@ where constraints, working_dir: "build".into(), scratch_paths: vec!["build".into(), "nix/store".into()], - // TODO: handle passAsFile, structuredAttrs. - additional_files: vec![], + additional_files: Vec::from_iter( + additional_files + .into_iter() + .map(|(path, contents)| AdditionalFile { path, contents }), + ), + }) +} + +/// handle passAsFile, if set. +/// For each env $x in that list, the original env is removed, and a $xPath +/// environment var added instead, referring to a path inside the build with +/// the contents from the original env var. +fn handle_pass_as_file( + environment_vars: &mut BTreeMap, + additional_files: &mut BTreeMap, +) -> std::io::Result<()> { + let pass_as_file = environment_vars.get("passAsFile").map(|v| { + // Convert pass_as_file to string. + // When it gets here, it contains a space-separated list of env var + // keys, which must be strings. + String::from_utf8(v.to_vec()) + }); + + if let Some(pass_as_file) = pass_as_file { + let pass_as_file = pass_as_file.map_err(|_| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "passAsFile elements are no valid utf8 strings", + ) + })?; + + for x in pass_as_file.split(' ') { + match environment_vars.remove_entry(x) { + Some((k, contents)) => { + let (new_k, path) = calculate_pass_as_file_env(&k); + + additional_files.insert(path.clone(), contents); + environment_vars.insert(new_k, Bytes::from(path)); + } + None => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "passAsFile refers to non-existent env key", + )); + } + } + } } + + Ok(()) +} + +/// For a given key k in a derivation environment that's supposed to be passed as file, +/// calculate the ${k}Path key and filepath value that it's being replaced with +/// while preparing the build. +/// The filepath is `/build/.attrs-${nixbase32(sha256(key))`. +fn calculate_pass_as_file_env(k: &str) -> (String, String) { + ( + format!("{}Path", k), + format!( + "/build/.attr-{}", + nixbase32::encode(&Sha256::new_with_prefix(k).finalize()) + ), + ) } #[cfg(test)] @@ -160,7 +227,7 @@ mod test { use bytes::Bytes; use nix_compat::derivation::Derivation; use tvix_build::proto::{ - build_request::{BuildConstraints, EnvVar}, + build_request::{AdditionalFile, BuildConstraints, EnvVar}, BuildRequest, }; use tvix_castore::{ @@ -205,7 +272,8 @@ mod test { // all good, reply with INPUT_NODE_FOO INPUT_NODE_FOO.clone() }, - ); + ) + .expect("must succeed"); let mut expected_environment_vars = vec![ EnvVar { @@ -266,7 +334,8 @@ mod test { let derivation = Derivation::from_aterm_bytes(aterm_bytes).expect("must parse"); let build_request = - derivation_to_build_request(&derivation, |_| unreachable!(), |_, _| unreachable!()); + derivation_to_build_request(&derivation, |_| unreachable!(), |_, _| unreachable!()) + .expect("must succeed"); let mut expected_environment_vars = vec![ EnvVar { @@ -327,4 +396,91 @@ mod test { build_request ); } + + #[test] + fn test_pass_as_file() { + // (builtins.derivation { "name" = "foo"; passAsFile = ["bar" "baz"]; bar = "baz"; baz = "bar"; system = ":"; builder = ":";}).drvPath + let aterm_bytes = r#"Derive([("out","/nix/store/pp17lwra2jkx8rha15qabg2q3wij72lj-foo","","")],[],[],":",":",[],[("bar","baz"),("baz","bar"),("builder",":"),("name","foo"),("out","/nix/store/pp17lwra2jkx8rha15qabg2q3wij72lj-foo"),("passAsFile","bar baz"),("system",":")])"#.as_bytes(); + + let derivation = Derivation::from_aterm_bytes(aterm_bytes).expect("must parse"); + + let build_request = + derivation_to_build_request(&derivation, |_| unreachable!(), |_, _| unreachable!()) + .expect("must succeed"); + + let mut expected_environment_vars = vec![ + // Note how bar and baz are not present in the env anymore, + // but replaced with barPath, bazPath respectively. + EnvVar { + key: "barPath".into(), + value: "/build/.attr-1fcgpy7vc4ammr7s17j2xq88scswkgz23dqzc04g8sx5vcp2pppw".into(), + }, + EnvVar { + key: "bazPath".into(), + value: "/build/.attr-15l04iksj1280dvhbzdq9ai3wlf8ac2188m9qv0gn81k9nba19ds".into(), + }, + EnvVar { + key: "builder".into(), + value: ":".into(), + }, + EnvVar { + key: "name".into(), + value: "foo".into(), + }, + EnvVar { + key: "out".into(), + value: "/nix/store/pp17lwra2jkx8rha15qabg2q3wij72lj-foo".into(), + }, + // passAsFile stays around + EnvVar { + key: "passAsFile".into(), + value: "bar baz".into(), + }, + EnvVar { + key: "system".into(), + value: ":".into(), + }, + ]; + + expected_environment_vars.extend(NIX_ENVIRONMENT_VARS.iter().map(|(k, v)| EnvVar { + key: k.to_string(), + value: Bytes::from_static(v.as_bytes()), + })); + + expected_environment_vars.sort_unstable_by_key(|e| e.key.to_owned()); + + assert_eq!( + BuildRequest { + command_args: vec![":".to_string()], + outputs: vec!["nix/store/pp17lwra2jkx8rha15qabg2q3wij72lj-foo".into()], + environment_vars: expected_environment_vars, + inputs: vec![], + inputs_dir: nix_compat::store_path::STORE_DIR.into(), + constraints: Some(BuildConstraints { + system: derivation.system.clone(), + min_memory: 0, + network_access: false, + available_ro_paths: vec![], + provide_bin_sh: true, + }), + additional_files: vec![ + // baz env + AdditionalFile { + path: "/build/.attr-15l04iksj1280dvhbzdq9ai3wlf8ac2188m9qv0gn81k9nba19ds" + .into(), + contents: "bar".into() + }, + // bar env + AdditionalFile { + path: "/build/.attr-1fcgpy7vc4ammr7s17j2xq88scswkgz23dqzc04g8sx5vcp2pppw" + .into(), + contents: "baz".into(), + }, + ], + working_dir: "build".into(), + scratch_paths: vec!["build".into(), "nix/store".into()], + }, + build_request + ); + } } -- cgit 1.4.1