GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

microsoft/openvmm

Public

mirrored fromhttps://github.com/microsoft/openvmmAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
copilot/apply-async-process-wait-functionality

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

openvmm/flowey/flowey_cli/src/cli

flowey/flowey_cli/src/cli/var_db.rs

373lines · modecode

RawDownload
Latest commit unavailable.
1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4use super::exec_snippet::FloweyPipelineStaticDb;
5use super::exec_snippet::VarDbBackendKind;
6use anyhow::Context;
7use clap::ValueEnum;
8use flowey_core::node::RuntimeVarDb;
9use std::io::Read;
10use std::io::Write;
11use std::path::Path;
12use std::path::PathBuf;
13
14pub struct VarDbRequest<'a> {
15 flowey_bin: &'a str,
16 job_idx: usize,
17 var_name: &'a str,
18 action: RequestAction<'a>,
19 is_raw_string: bool,
20 condvar: Option<&'a str>,
21}
22
23enum RequestAction<'a> {
24 WriteToEnv {
25 backend: EnvBackend,
26 env: &'a str,
27 },
28 Update {
29 file: Option<&'a Path>,
30 is_secret: bool,
31 env_source: Option<&'a str>,
32 },
33}
34
35pub struct VarDbRequestBuilder<'a> {
36 flowey_bin: &'a str,
37 job_idx: usize,
38}
39
40impl<'a> VarDbRequestBuilder<'a> {
41 pub fn new(flowey_bin: &'a str, job_idx: usize) -> Self {
42 Self {
43 flowey_bin,
44 job_idx,
45 }
46 }
47
48 fn req<'b>(&'b self, var_name: &'b str, action: RequestAction<'b>) -> VarDbRequest<'b> {
49 VarDbRequest::new(self.flowey_bin, self.job_idx, var_name, action)
50 }
51
52 pub fn write_to_ado_env<'b>(&'b self, var_name: &'b str, env: &'b str) -> VarDbRequest<'b> {
53 self.req(
54 var_name,
55 RequestAction::WriteToEnv {
56 backend: EnvBackend::Ado,
57 env,
58 },
59 )
60 }
61
62 pub fn write_to_gh_env<'b>(&'b self, var_name: &'b str, env: &'b str) -> VarDbRequest<'b> {
63 self.req(
64 var_name,
65 RequestAction::WriteToEnv {
66 backend: EnvBackend::Github,
67 env,
68 },
69 )
70 }
71
72 pub fn update_from_stdin<'b>(&'b self, var_name: &'b str, is_secret: bool) -> VarDbRequest<'b> {
73 self.req(
74 var_name,
75 RequestAction::Update {
76 file: None,
77 is_secret,
78 env_source: None,
79 },
80 )
81 }
82
83 #[expect(dead_code)]
84 pub fn update_from_file<'b>(
85 &'b self,
86 var_name: &'b str,
87 file: &'b Path,
88 is_secret: bool,
89 ) -> VarDbRequest<'b> {
90 self.req(
91 var_name,
92 RequestAction::Update {
93 file: Some(file),
94 is_secret,
95 env_source: None,
96 },
97 )
98 }
99}
100
101impl<'a> VarDbRequest<'a> {
102 fn new(
103 flowey_bin: &'a str,
104 job_idx: usize,
105 var_name: &'a str,
106 action: RequestAction<'a>,
107 ) -> Self {
108 Self {
109 flowey_bin,
110 job_idx,
111 var_name,
112 action,
113 is_raw_string: false,
114 condvar: None,
115 }
116 }
117
118 pub fn raw_string(self, is_raw_string: bool) -> Self {
119 Self {
120 is_raw_string,
121 ..self
122 }
123 }
124
125 pub fn condvar(self, condvar: Option<&'a str>) -> Self {
126 Self { condvar, ..self }
127 }
128
129 #[track_caller]
130 pub fn env_source(mut self, source: Option<&'a str>) -> Self {
131 let RequestAction::Update { env_source, .. } = &mut self.action else {
132 panic!("env_source can only be set on Update actions");
133 };
134 *env_source = source;
135 self
136 }
137}
138
139impl std::fmt::Display for VarDbRequest<'_> {
140 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
141 let Self {
142 flowey_bin,
143 job_idx,
144 var_name,
145 ref action,
146 is_raw_string,
147 condvar,
148 } = *self;
149
150 write!(f, r#"{flowey_bin} v {job_idx} '{var_name}'"#)?;
151
152 if is_raw_string {
153 f.write_str(" --is-raw-string")?;
154 }
155
156 if let Some(condvar) = condvar {
157 write!(f, " --condvar {condvar}")?;
158 }
159
160 match *action {
161 RequestAction::WriteToEnv { backend, env } => {
162 write!(
163 f,
164 " write-to-env {backend} {env}",
165 backend = backend.to_possible_value().unwrap().get_name()
166 )?;
167 }
168 RequestAction::Update {
169 file,
170 is_secret,
171 env_source,
172 } => {
173 write!(f, " update")?;
174 if is_secret {
175 f.write_str(" --is-secret")?;
176 }
177 if let Some(env_source) = env_source {
178 write!(f, " --env-source {env_source}")?;
179 }
180 if let Some(file) = file {
181 write!(f, " {}", file.to_str().unwrap())?;
182 }
183 }
184 }
185
186 Ok(())
187 }
188}
189
190/// (internal) interact with the runtime variable database
191#[derive(clap::Args)]
192pub struct VarDb {
193 /// job idx corresponding to the var db to access
194 pub(crate) job_idx: usize,
195
196 /// Runtime variable to access
197 var_name: String,
198
199 /// Variable is a raw string, and should be read/written as a plain string.
200 #[clap(long)]
201 is_raw_string: bool,
202
203 /// Only run if the given variable is true.
204 #[clap(long)]
205 condvar: Option<String>,
206
207 #[clap(subcommand)]
208 action: Option<VarDbAction>,
209}
210
211#[derive(clap::Subcommand)]
212enum VarDbAction {
213 WriteToEnv {
214 backend: EnvBackend,
215 env: String,
216 },
217 Update {
218 #[clap(long)]
219 env_source: Option<String>,
220 #[clap(long)]
221 is_secret: bool,
222 file: Option<PathBuf>,
223 },
224}
225
226#[derive(clap::ValueEnum, Copy, Clone)]
227enum EnvBackend {
228 Ado,
229 Github,
230}
231
232impl VarDb {
233 pub fn run(self) -> anyhow::Result<()> {
234 let Self {
235 job_idx,
236 var_name,
237 is_raw_string,
238 condvar,
239 action,
240 } = self;
241
242 let mut runtime_var_db = open_var_db(job_idx)?;
243
244 if let Some(condvar) = condvar {
245 let (condvar_data, _) = runtime_var_db.get_var(&condvar);
246 let set: bool = serde_json::from_slice(&condvar_data).unwrap();
247 if !set {
248 return Ok(());
249 }
250 }
251
252 let get = |runtime_var_db: &mut Box<dyn RuntimeVarDb>, var_name: &str| {
253 let (mut data, data_is_secret) = runtime_var_db.get_var(var_name);
254 // HACK: only one kind of db, so we know what routine to use
255 if is_raw_string {
256 let s: String = serde_json::from_slice(&data).unwrap();
257 data = s.into();
258 }
259 (data, data_is_secret)
260 };
261
262 let env_source_name = |env_source| format!(".env.is_secret.{env_source}");
263
264 match action {
265 None => {
266 // Raw get.
267 let (data, _) = get(&mut runtime_var_db, &var_name);
268 std::io::stdout().write_all(&data).unwrap();
269 }
270 Some(VarDbAction::WriteToEnv { backend, env }) => {
271 let (data, is_secret) = get(&mut runtime_var_db, &var_name);
272
273 if is_secret {
274 // Remember that this environment variable is secret so that
275 // it cannot be easily laundered into a non-secret variable.
276 runtime_var_db.set_var(&env_source_name(&env), false, "null".into());
277 }
278
279 match backend {
280 EnvBackend::Ado => {
281 print!("##vso[task.setvariable variable={env};issecret={is_secret}]");
282 std::io::stdout().write_all(&data).unwrap();
283 println!();
284 }
285 EnvBackend::Github => {
286 let data_string = String::from_utf8(data)?;
287 if is_secret {
288 data_string.lines().for_each(|line| {
289 println!("::add-mask::{}", line);
290 });
291 }
292 let gh_env_file_path = std::env::var("GITHUB_ENV")?;
293 let mut gh_env_file = fs_err::OpenOptions::new()
294 .append(true)
295 .open(gh_env_file_path)?;
296 let gh_env_var_assignment = format!("{}<<EOF\n{}\nEOF\n", env, data_string);
297 gh_env_file.write_all(gh_env_var_assignment.as_bytes())?;
298 }
299 }
300 }
301 Some(VarDbAction::Update {
302 env_source,
303 mut is_secret,
304 file,
305 }) => {
306 if !is_secret {
307 // If the source environment variable for this was known to
308 // be a secret, then mark it secret.
309 if let Some(env_source) = env_source {
310 is_secret |= runtime_var_db
311 .try_get_var(&env_source_name(&env_source))
312 .is_some();
313 }
314 }
315 let data = if let Some(file) = file {
316 let mut data = fs_err::read(file)?;
317 // HACK: only one kind of db, so we know what routine to use
318 if is_raw_string {
319 let s: String = String::from_utf8(data).unwrap();
320 data = serde_json::to_vec(&s).unwrap();
321 }
322 data
323 } else {
324 let mut data = Vec::new();
325 std::io::stdin().read_to_end(&mut data).unwrap();
326 // HACK: only one kind of db, so we know what routine to use
327 if is_raw_string {
328 // account for bash HEREDOCs including a trailing newline
329 // TODO: probably want this to be configurable.
330 if matches!(data.last(), Some(b'\n')) {
331 data.pop();
332 }
333
334 let s = String::from_utf8(data).unwrap();
335 data = serde_json::to_vec(&s).unwrap();
336 }
337 data
338 };
339 runtime_var_db.set_var(&var_name, is_secret, data);
340 }
341 }
342
343 Ok(())
344 }
345}
346
347/// Obtain a handle to a runtime var db
348///
349/// CONTRACT: Requires a pipeline-specific `pipeline.json` file to be in the
350/// same dir as the flowey exe
351///
352/// CONTRACT: Requires a var-backend specific var db file called
353/// `job{job_idx}.<ext>` to be in the same dir as the flowey exe
354pub(crate) fn open_var_db(job_idx: usize) -> anyhow::Result<Box<dyn RuntimeVarDb>> {
355 let current_exe =
356 std::env::current_exe().context("failed to get path to current flowey executable")?;
357
358 let FloweyPipelineStaticDb {
359 var_db_backend_kind,
360 ..
361 } = {
362 let pipeline_static_db = fs_err::File::open(current_exe.with_file_name("pipeline.json"))?;
363 serde_json::from_reader(pipeline_static_db)?
364 };
365
366 Ok(match var_db_backend_kind {
367 VarDbBackendKind::Json => {
368 Box::new(crate::var_db::single_json_file::SingleJsonFileVarDb::new(
369 current_exe.with_file_name(format!("job{job_idx}.json")),
370 )?)
371 }
372 })
373}
374