[PATCH proxmox 03/20] api-macro: support shared state extraction type

[Lukas Wagner <l.wagner@proxmox.com>]

The State<T> newtype wraps values that come from the registry, which
allows telling them apart from regular API parameters. The API macro
recognizes them in API handler signatures and looks them up in the
shared state registry of the RpcEnvironment, instead of deserializing
them from the request parameters. That way a handler declares its
dependencies on application state in its signature and does not have to
access the environment by hand.

Extraction fails at runtime if the router has no registry set up or if
the requested type was never registered.

Signed-off-by: Lukas Wagner <l.wagner@proxmox.com>
---
 proxmox-api-macro/src/api/method.rs |  88 +++++++++++++++++++++++-
 proxmox-api-macro/tests/state.rs    | 101 ++++++++++++++++++++++++++++
 proxmox-router/src/shared_state.rs  |  30 +++++++++
 3 files changed, 218 insertions(+), 1 deletion(-)
 create mode 100644 proxmox-api-macro/tests/state.rs

diff --git a/proxmox-api-macro/src/api/method.rs b/proxmox-api-macro/src/api/method.rs
index 5ef69e72..20acd50c 100644
--- a/proxmox-api-macro/src/api/method.rs
+++ b/proxmox-api-macro/src/api/method.rs
@@ -3,7 +3,8 @@
 //! This has to perform quite a few things: infer types from parameters, deal with optional types
 //! and defaults, expose parameter and return value schema to the public, and finally create the
 //! wrapper function converting from a json Value hash to the parameters listed in the function
-//! signature, while recognizing specially handling `RPCEnvironment` and `ApiMethod` parameters.
+//! signature, while recognizing specially handling `RPCEnvironment`, `ApiMethod`, and `State<T>`
+//! parameters.
 
 use std::convert::{TryFrom, TryInto};
 use std::mem;
@@ -349,6 +350,7 @@ enum ParameterType {
     Value,
     ApiMethod,
     RpcEnv,
+    State(StateParameter),
     Normal(NormalParameter),
 }
 
@@ -357,6 +359,11 @@ struct NormalParameter {
     entry: ObjectEntry,
 }
 
+#[derive(Debug)]
+struct StateParameter {
+    ty: syn::Type,
+}
+
 fn check_input_type(input: &syn::FnArg) -> Result<(&syn::PatType, &syn::PatIdent), syn::Error> {
     // `self` types are not supported:
     let pat_type = match input {
@@ -463,6 +470,8 @@ fn handle_function_signature(method_info: &mut MethodInfo) -> Result<Ident, Erro
             }
             api_method_param = Some(param_list.len());
             ParameterType::ApiMethod
+        } else if let Some(s) = as_state_type(&pat_type.ty) {
+            ParameterType::State(s)
         } else if is_rpc_env_type(&pat_type.ty) {
             if rpc_env_param.is_some() {
                 error!(pat_type => "multiple RpcEnvironment parameters found");
@@ -502,6 +511,38 @@ fn is_api_method_type(ty: &syn::Type) -> bool {
     false
 }
 
+fn as_state_type(ty: &syn::Type) -> Option<StateParameter> {
+    if let syn::Type::Path(p) = ty {
+        if p.qself.is_some() {
+            return None;
+        }
+
+        if let Some(ps) = p.path.segments.last() {
+            if ps.ident != "State" {
+                return None;
+            }
+
+            match &ps.arguments {
+                syn::PathArguments::AngleBracketed(angle_bracketed_generic_arguments) => {
+                    if angle_bracketed_generic_arguments.args.len() != 1 {
+                        return None;
+                    }
+
+                    let generic_argument = angle_bracketed_generic_arguments.args[0].clone();
+
+                    match generic_argument {
+                        syn::GenericArgument::Type(ty) => return Some(StateParameter { ty }),
+                        _ => return None,
+                    }
+                }
+                _ => return None,
+            }
+        }
+    }
+
+    None
+}
+
 fn is_rpc_env_type(ty: &syn::Type) -> bool {
     if let syn::Type::Reference(r) = ty
         && let syn::Type::TraitObject(t) = &*r.elem
@@ -563,6 +604,9 @@ fn create_wrapper_function(
                     &mut method_info.default_consts,
                 )?;
             }
+            ParameterType::State(param) => {
+                extract_state_parameter(param, &mut body, &mut args, name, span)?;
+            }
         }
     }
 
@@ -824,6 +868,48 @@ fn extract_normal_parameter(
     Ok(())
 }
 
+fn extract_state_parameter(
+    param: StateParameter,
+    body: &mut TokenStream,
+    args: &mut TokenStream,
+    name: FieldName,
+    span: Span,
+) -> Result<(), Error> {
+    let name_str = syn::LitStr::new(name.as_str(), span);
+    let arg_name = Ident::new(&format!("state_arg_{}", name.as_ident()), span);
+
+    let type_name = match param.ty {
+        syn::Type::Path(type_path) => {
+            let last_segment = type_path.path.segments.last().unwrap();
+            last_segment.ident.to_string()
+        }
+        _ => {
+            error!(span, "only simple types allowed");
+            "".into()
+        }
+    };
+
+    body.extend(quote_spanned! { span =>
+        #[allow(non_snake_case)]
+        let #arg_name = rpc_env_param
+            .shared_state()
+            .ok_or_else(|| ::anyhow::format_err!(
+                "router has no shared state set up",
+            ))?
+            .lookup()
+            .map(::proxmox_router::State)
+            .ok_or_else(|| ::anyhow::format_err!(
+                "shared context type '{}' for parameter '{}' not registered",
+                #type_name,
+                #name_str,
+            ))?;
+    });
+
+    args.extend(quote_spanned! { span => #arg_name, });
+
+    Ok(())
+}
+
 /// Returns a tuple containing the schema code first and the `ParameterSchema` parameter for the
 /// `ApiMethod` second.
 fn serialize_input_schema(
diff --git a/proxmox-api-macro/tests/state.rs b/proxmox-api-macro/tests/state.rs
new file mode 100644
index 00000000..732c3e2e
--- /dev/null
+++ b/proxmox-api-macro/tests/state.rs
@@ -0,0 +1,101 @@
+use proxmox_api_macro::api;
+use proxmox_router::{SharedStateRegistry, State};
+
+use anyhow::Error;
+use serde_json::{Value, json};
+
+#[derive(Clone)]
+struct Foo {
+    val: i32,
+}
+
+#[api(
+    input: {
+        properties: {
+            value: {
+                description: "Something",
+            }
+        }
+    }
+)]
+/// Test multiple state args.
+fn multiple_args(
+    state_arg1: State<i32>,
+    value: isize,
+    state_arg2: State<Foo>,
+) -> Result<(), Error> {
+    assert_eq!(*state_arg1, 10);
+    assert_eq!(state_arg2.val, 20);
+
+    assert_eq!(value, 50);
+
+    Ok(())
+}
+
+#[api]
+/// Test not registered state arg.
+fn not_registered(_state: State<Foo>) -> Result<(), Error> {
+    panic!("should not reach this");
+}
+
+struct RpcEnv {
+    shared_state: SharedStateRegistry,
+}
+impl proxmox_router::RpcEnvironment for RpcEnv {
+    fn result_attrib_mut(&mut self) -> &mut Value {
+        panic!("result_attrib_mut called");
+    }
+
+    fn result_attrib(&self) -> &Value {
+        panic!("result_attrib called");
+    }
+
+    /// The environment type
+    fn env_type(&self) -> proxmox_router::RpcEnvironmentType {
+        panic!("env_type called");
+    }
+
+    /// Set authentication id
+    fn set_auth_id(&mut self, user: Option<String>) {
+        let _ = user;
+        panic!("set_auth_id called");
+    }
+
+    /// Get authentication id
+    fn get_auth_id(&self) -> Option<String> {
+        panic!("get_auth_id called");
+    }
+
+    fn shared_state(&self) -> Option<&SharedStateRegistry> {
+        Some(&self.shared_state)
+    }
+}
+
+#[test]
+fn test_invocations() {
+    let mut shared_state = SharedStateRegistry::default();
+
+    shared_state.register::<i32>(10).unwrap();
+    shared_state.register::<Foo>(Foo { val: 20 }).unwrap();
+
+    let mut env = RpcEnv { shared_state };
+
+    api_function_multiple_args(
+        json!({
+            "value": 50,
+        }),
+        &API_METHOD_MULTIPLE_ARGS,
+        &mut env,
+    )
+    .expect("func with multiple injected works");
+}
+
+#[test]
+fn test_not_registered() {
+    let shared_state = SharedStateRegistry::default();
+
+    let mut env = RpcEnv { shared_state };
+
+    api_function_multiple_args(json!({}), &API_METHOD_NOT_REGISTERED, &mut env)
+        .expect_err("func did not fail");
+}
diff --git a/proxmox-router/src/shared_state.rs b/proxmox-router/src/shared_state.rs
index 14a539af..9bd6de22 100644
--- a/proxmox-router/src/shared_state.rs
+++ b/proxmox-router/src/shared_state.rs
@@ -2,6 +2,7 @@
 
 use std::any::{Any, TypeId};
 use std::collections::HashMap;
+use std::ops::Deref;
 
 use anyhow::{Error, bail};
 
@@ -44,3 +45,32 @@ impl SharedStateRegistry {
         Ok(())
     }
 }
+
+/// Wrapper for a value taken from a [`SharedStateRegistry`].
+///
+/// An API handler declares a parameter of this type to get a clone of the registered value of
+/// type `T`. Such parameters are not part of the API schema, they are filled in from the
+/// environment's registry when the handler is called.
+#[derive(Clone)]
+pub struct State<T>(pub T);
+
+impl<T> Deref for State<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<T> From<T> for State<T> {
+    fn from(value: T) -> Self {
+        State(value)
+    }
+}
+
+impl<T> State<T> {
+    /// Take out the wrapped value.
+    pub fn into_inner(self) -> T {
+        self.0
+    }
+}
-- 
2.47.3