GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

microsoft/openvmm

Public

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

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
6d9f566744d8d46492a564b6663d29dfacfc3fcc

Branches

Tags

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

Clone

HTTPS

Download ZIP

openvmm/flowey/flowey_core/src

flowey/flowey_core/src/patch.rs

225lines · modecode

RawDownload
flowey: make values secret, not variables (#1338)7873452
1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4use crate::node::FlowNodeBase;
5use crate::node::NodeHandle;
6use crate::node::WriteVar;
7use std::collections::BTreeMap;
8use std::sync::OnceLock;
9
10pub type PatchFn = fn(&mut PatchManager<'_>);
11
12// A patchfn that does nothing. Can be useful when writing logic that
13// conditionally applies patches.
14pub fn noop_patchfn(_: &mut PatchManager<'_>) {}
15
16enum PatchEvent {
17 Swap {
18 from_old_node: NodeHandle,
19 with_new_node: NodeHandle,
20 },
21 InjectSideEffect {
22 from_old_node: NodeHandle,
23 with_new_node: NodeHandle,
24 side_effect_var: String,
25 req: Box<[u8]>,
26 },
27}
28
29trait PatchManagerBackend {
30 fn new_side_effect_var(&mut self) -> String;
31 fn on_patch_event(&mut self, event: PatchEvent);
32}
33
34/// Passed to patch functions
35pub struct PatchManager<'a> {
36 backend: &'a mut dyn PatchManagerBackend,
37}
38
39impl PatchManager<'_> {
40 pub fn hook<N: FlowNodeBase>(&mut self) -> PatchHook<'_, N> {
41 PatchHook {
42 backend: self.backend,
43 _kind: std::marker::PhantomData,
44 }
45 }
46}
47
48/// Patch operations in the context of a particular Node.
49pub struct PatchHook<'a, N: FlowNodeBase> {
50 backend: &'a mut dyn PatchManagerBackend,
51 _kind: std::marker::PhantomData<N>,
52}
53
54impl<N> PatchHook<'_, N>
55where
56 N: FlowNodeBase + 'static,
57{
58 /// Swap out the target Node's implementation with a different
59 /// implementation.
60 pub fn swap_with<M>(&mut self) -> &mut Self
61 where
62 M: 'static,
63 // use the type system to enforce that patch nodes have an identical
64 // request type
65 M: FlowNodeBase<Request = N::Request>,
66 {
67 self.backend.on_patch_event(PatchEvent::Swap {
68 from_old_node: NodeHandle::from_type::<N>(),
69 with_new_node: NodeHandle::from_type::<M>(),
70 });
71 self
72 }
73
74 /// Inject a side-effect dependency, which runs before any other steps in
75 /// the Node.
76 pub fn inject_side_effect<T, M>(
77 &mut self,
78 f: impl FnOnce(WriteVar<T>) -> M::Request,
79 ) -> &mut Self
80 where
81 T: serde::Serialize + serde::de::DeserializeOwned,
82 M: 'static,
83 M: FlowNodeBase,
84 {
85 let backing_var = self.backend.new_side_effect_var();
86 let req = f(crate::node::thin_air_write_runtime_var(backing_var.clone()));
87
88 self.backend.on_patch_event(PatchEvent::InjectSideEffect {
89 from_old_node: NodeHandle::from_type::<N>(),
90 with_new_node: NodeHandle::from_type::<M>(),
91 side_effect_var: backing_var,
92 req: serde_json::to_vec(&req).map(Into::into).unwrap(),
93 });
94 self
95 }
96}
97
98pub fn patchfn_by_modpath() -> &'static BTreeMap<String, PatchFn> {
99 static MODPATH_LOOKUP: OnceLock<BTreeMap<String, PatchFn>> = OnceLock::new();
100
101 let lookup = MODPATH_LOOKUP.get_or_init(|| {
102 let mut lookup = BTreeMap::new();
103 for (f, module_path, fn_name) in private::PATCH_FNS {
104 let existing = lookup.insert(format!("{}::{}", module_path, fn_name), *f);
105 // Rust would've errored out at module defn time with a duplicate fn name error
106 assert!(existing.is_none());
107 }
108 lookup
109 });
110
111 lookup
112}
113
114/// [`PatchResolver`]
115#[derive(Debug, Clone)]
116pub struct ResolvedPatches {
117 pub swap: BTreeMap<NodeHandle, NodeHandle>,
118 pub inject_side_effect: BTreeMap<NodeHandle, Vec<(NodeHandle, String, Box<[u8]>)>>,
119}
120
121impl ResolvedPatches {
122 pub fn build() -> PatchResolver {
123 PatchResolver {
124 side_effect_var_idx: 0,
125 swap: BTreeMap::default(),
126 inject_side_effect: BTreeMap::new(),
127 }
128 }
129}
130
131/// Helper method to resolve multiple patches into a single [`ResolvedPatches`]
132#[derive(Debug)]
133pub struct PatchResolver {
134 side_effect_var_idx: usize,
135 swap: BTreeMap<NodeHandle, NodeHandle>,
136 inject_side_effect: BTreeMap<NodeHandle, Vec<(NodeHandle, String, Box<[u8]>)>>,
137}
138
139impl PatchResolver {
140 pub fn apply_patchfn(&mut self, patchfn: PatchFn) {
141 patchfn(&mut PatchManager { backend: self });
142 }
143
144 pub fn finalize(self) -> ResolvedPatches {
145 let Self {
146 swap,
147 mut inject_side_effect,
148 side_effect_var_idx: _,
149 } = self;
150
151 // take into account the interaction between swaps and injected effects
152 for (from, to) in &swap {
153 let injected = inject_side_effect.remove(from);
154 if let Some(injected) = injected {
155 inject_side_effect.insert(*to, injected);
156 }
157 }
158
159 ResolvedPatches {
160 swap,
161 inject_side_effect,
162 }
163 }
164}
165
166impl PatchManagerBackend for PatchResolver {
167 fn new_side_effect_var(&mut self) -> String {
168 self.side_effect_var_idx += 1;
169 format!("patch_side_effect:{}", self.side_effect_var_idx)
170 }
171
172 fn on_patch_event(&mut self, event: PatchEvent) {
173 match event {
174 PatchEvent::Swap {
175 from_old_node,
176 with_new_node,
177 } => {
178 let existing = self.swap.insert(from_old_node, with_new_node);
179 // FUTURE: add some better error reporting / logging to
180 // allow doing this, albeit with a warning
181 assert!(
182 existing.is_none(),
183 "cannot double-patch the same node combo"
184 );
185 }
186 PatchEvent::InjectSideEffect {
187 from_old_node,
188 with_new_node,
189 side_effect_var,
190 req,
191 } => {
192 self.inject_side_effect
193 .entry(from_old_node)
194 .or_default()
195 .push((with_new_node, side_effect_var, req));
196 }
197 }
198 }
199}
200
201#[doc(hidden)]
202pub mod private {
203 use super::PatchFn;
204 pub use linkme;
205
206 #[linkme::distributed_slice]
207 pub static PATCH_FNS: [(PatchFn, &'static str, &'static str)] = [..];
208
209 /// Register a patch function which can be used when emitting flows.
210 ///
211 /// The function must conform to the signature of [`PatchFn`]
212 #[macro_export]
213 macro_rules! register_patch {
214 ($patchfn:ident) => {
215 const _: () = {
216 use $crate::node::private::linkme;
217
218 #[linkme::distributed_slice($crate::patch::private::PATCH_FNS)]
219 #[linkme(crate = linkme)]
220 pub static PATCH_FNS: ($crate::patch::PatchFn, &'static str, &'static str) =
221 ($patchfn, module_path!(), stringify!($patchfn));
222 };
223 };
224 }
225}