Skip to main content

bitwarden_state_bridge_macro/
lib.rs

1//! Proc macro for generating the key-management state bridge surface.
2//!
3//! Provides the `state_bridge!` function-like macro that, from a list of typed fields, expands
4//! into the `StateBridgeImpl` trait, wrapper methods on `StateBridge` and `StateBridgeClient`,
5//! WASM extern bindings, the `WasmStateBridge` trait impl, and the matching TypeScript interface.
6
7use proc_macro::TokenStream;
8use quote::{format_ident, quote};
9use syn::{
10    Ident, LitStr, Token, Type,
11    parse::{Parse, ParseStream},
12    parse_macro_input,
13    punctuated::Punctuated,
14};
15
16mod state_bridge_kw {
17    syn::custom_keyword!(ts);
18}
19
20struct StateBridgeField {
21    name: Ident,
22    ty: Type,
23    ts: LitStr,
24}
25
26impl Parse for StateBridgeField {
27    fn parse(input: ParseStream) -> syn::Result<Self> {
28        let name: Ident = input.parse()?;
29        input.parse::<Token![:]>()?;
30        let ty: Type = input.parse()?;
31        input.parse::<Token![as]>()?;
32        input.parse::<state_bridge_kw::ts>()?;
33        let ts: LitStr = input.parse()?;
34        Ok(Self { name, ty, ts })
35    }
36}
37
38struct StateBridgeInput {
39    fields: Punctuated<StateBridgeField, Token![,]>,
40}
41
42impl Parse for StateBridgeInput {
43    fn parse(input: ParseStream) -> syn::Result<Self> {
44        Ok(Self {
45            fields: Punctuated::parse_terminated(input)?,
46        })
47    }
48}
49
50/// Generates the full state bridge surface for a fixed list of fields.
51///
52/// Each field expands to:
53/// 1. Three methods on the `StateBridgeImpl` trait (`set_$name`, `get_$name`, `clear_$name`).
54/// 2. Three corresponding wrapper methods on `StateBridge`.
55/// 3. Three corresponding methods on `StateBridgeClient`.
56/// 4. Three method declarations on the WASM `RawWasmStateBridge` extern type and three forwarders
57///    in the `StateBridgeImpl` impl for `WasmStateBridge`.
58/// 5. Three lines in the `WasmStateBridge` TypeScript interface.
59///
60/// All fields share the same shape: `set_$name(value: $ty)`, `get_$name() -> Option<$ty>`,
61/// `clear_$name()`.
62#[proc_macro]
63pub fn state_bridge(input: TokenStream) -> TokenStream {
64    let StateBridgeInput { fields } = parse_macro_input!(input as StateBridgeInput);
65
66    let trait_methods = fields.iter().map(|f| {
67        let ty = &f.ty;
68        let n = f.name.to_string();
69        let set = format_ident!("set_{}", f.name);
70        let get = format_ident!("get_{}", f.name);
71        let clear = format_ident!("clear_{}", f.name);
72        let set_doc = format!("Stores the `{n}` value.");
73        let get_doc = format!("Returns the `{n}` value, if available.");
74        let clear_doc = format!("Clears the `{n}` value.");
75        quote! {
76            #[doc = #set_doc]
77            async fn #set(&self, value: #ty);
78            #[doc = #get_doc]
79            async fn #get(&self) -> Option<#ty>;
80            #[doc = #clear_doc]
81            async fn #clear(&self);
82        }
83    });
84
85    let bridge_wrappers = fields.iter().map(|f| {
86        let ty = &f.ty;
87        let n = f.name.to_string();
88        let set = format_ident!("set_{}", f.name);
89        let get = format_ident!("get_{}", f.name);
90        let clear = format_ident!("clear_{}", f.name);
91        let set_doc = format!("Stores the `{n}` value.");
92        let get_doc = format!("Returns the `{n}` value, if available.");
93        let clear_doc = format!("Clears the `{n}` value.");
94        quote! {
95            #[doc = #set_doc]
96            pub async fn #set(&self, value: &#ty) {
97                let implementation = self
98                    .implementation
99                    .lock()
100                    .expect("Mutex is not poisoned")
101                    .as_ref()
102                    .expect("StateBridge not registered")
103                    .clone();
104                implementation.#set(value.to_owned()).await
105            }
106
107            #[doc = #get_doc]
108            pub async fn #get(&self) -> Option<#ty> {
109                let implementation = self
110                    .implementation
111                    .lock()
112                    .expect("Mutex is not poisoned")
113                    .as_ref()
114                    .expect("StateBridge not registered")
115                    .clone();
116                implementation.#get().await
117            }
118
119            #[doc = #clear_doc]
120            pub async fn #clear(&self) {
121                let implementation = self
122                    .implementation
123                    .lock()
124                    .expect("Mutex is not poisoned")
125                    .as_ref()
126                    .expect("StateBridge not registered")
127                    .clone();
128                implementation.#clear().await
129            }
130        }
131    });
132
133    let client_forwarders = fields.iter().map(|f| {
134        let ty = &f.ty;
135        let n = f.name.to_string();
136        let set = format_ident!("set_{}", f.name);
137        let get = format_ident!("get_{}", f.name);
138        let clear = format_ident!("clear_{}", f.name);
139        let set_doc = format!("Sets the `{n}` value in client-managed state.");
140        let get_doc = format!("Gets the `{n}` value from client-managed state, if available.");
141        let clear_doc = format!("Clears the `{n}` value from client-managed state.");
142        quote! {
143            #[doc = #set_doc]
144            pub async fn #set(&self, value: &#ty) {
145                self.client.internal.state_bridge.#set(value).await;
146            }
147
148            #[doc = #get_doc]
149            pub async fn #get(&self) -> Option<#ty> {
150                self.client.internal.state_bridge.#get().await
151            }
152
153            #[doc = #clear_doc]
154            pub async fn #clear(&self) {
155                self.client.internal.state_bridge.#clear().await;
156            }
157        }
158    });
159
160    let mut ts_iface = String::from(
161        "/**\n * Typescript interface that the state bridge needs to implement. The state bridge\n * is a temporary layer that allows quickly transitioning non-repository shaped\n * state to be accessible from within the SDK.\n */\nexport interface WasmStateBridge {\n",
162    );
163    for f in &fields {
164        let n = f.name.to_string();
165        let t = f.ts.value();
166        ts_iface.push_str(&format!("    set_{n}(value: {t}): Promise<void>;\n"));
167        ts_iface.push_str(&format!("    get_{n}(): Promise<{t} | null>;\n"));
168        ts_iface.push_str(&format!("    clear_{n}(): Promise<void>;\n"));
169    }
170    ts_iface.push_str("}\n");
171
172    let extern_methods = fields.iter().map(|f| {
173        let ty = &f.ty;
174        let n = f.name.to_string();
175        let set = format_ident!("set_{}", f.name);
176        let get = format_ident!("get_{}", f.name);
177        let clear = format_ident!("clear_{}", f.name);
178        let set_doc = format!("JS-side `set_{n}` method on `WasmStateBridge`.");
179        let get_doc = format!("JS-side `get_{n}` method on `WasmStateBridge`.");
180        let clear_doc = format!("JS-side `clear_{n}` method on `WasmStateBridge`.");
181        quote! {
182            #[doc = #set_doc]
183            #[wasm_bindgen(method)]
184            pub async fn #set(
185                this: &crate::key_management::state_bridge::RawWasmStateBridge,
186                value: #ty,
187            );
188            #[doc = #get_doc]
189            #[wasm_bindgen(method)]
190            pub async fn #get(
191                this: &crate::key_management::state_bridge::RawWasmStateBridge,
192            ) -> Option<#ty>;
193            #[doc = #clear_doc]
194            #[wasm_bindgen(method)]
195            pub async fn #clear(
196                this: &crate::key_management::state_bridge::RawWasmStateBridge,
197            );
198        }
199    });
200
201    let wasm_impls = fields.iter().map(|f| {
202        let ty = &f.ty;
203        let set = format_ident!("set_{}", f.name);
204        let get = format_ident!("get_{}", f.name);
205        let clear = format_ident!("clear_{}", f.name);
206        quote! {
207            async fn #set(&self, value: #ty) {
208                self.0
209                    .run_in_thread(move |state| async move {
210                        state.#set(value).await
211                    })
212                    .await
213                    .expect("State bridge call panicked");
214            }
215            async fn #get(&self) -> Option<#ty> {
216                self.0
217                    .run_in_thread(|state| async move {
218                        state.#get().await
219                    })
220                    .await
221                    .expect("State bridge call panicked")
222            }
223            async fn #clear(&self) {
224                self.0
225                    .run_in_thread(|state| async move {
226                        state.#clear().await
227                    })
228                    .await
229                    .expect("State bridge call panicked");
230            }
231        }
232    });
233
234    let expanded = quote! {
235        /// Host-provided storage bridge for key-management state.
236        ///
237        /// SDK consumers register an implementation that persists or caches sensitive
238        /// account state across unlock flows.
239        #[cfg_attr(target_arch = "wasm32", ::async_trait::async_trait(?Send))]
240        #[cfg_attr(not(target_arch = "wasm32"), ::async_trait::async_trait)]
241        pub trait StateBridgeImpl: Send + Sync {
242            #(#trait_methods)*
243        }
244
245        impl StateBridge {
246            #(#bridge_wrappers)*
247        }
248
249        impl crate::key_management::state_bridge::StateBridgeClient {
250            #(#client_forwarders)*
251        }
252
253        #[cfg(target_arch = "wasm32")]
254        #[::wasm_bindgen::prelude::wasm_bindgen(typescript_custom_section)]
255        const TS_CUSTOM_TYPES_STATE_BRIDGE: &'static str = #ts_iface;
256
257        #[cfg(target_arch = "wasm32")]
258        #[::wasm_bindgen::prelude::wasm_bindgen]
259        extern "C" {
260            #(#extern_methods)*
261        }
262
263        #[cfg(target_arch = "wasm32")]
264        #[::async_trait::async_trait(?Send)]
265        impl StateBridgeImpl for crate::key_management::state_bridge::WasmStateBridge {
266            #(#wasm_impls)*
267        }
268    };
269
270    TokenStream::from(expanded)
271}