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

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

The #[state] attribute marks types 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.

State arguments can either be taken by value or by reference, the latter
avoids a clone in the method generated by the API macro. Taking the
argument by reference is mutually exclusive with a mutable
`RpcEnvironment` reference, since the state argument is borrowed from
the environment. The macro shows a distinct error message in this case,
educating the developer to use a by-value argument instead.

If we ever decide to store shared state outside the rpc-environment,
this restriction can be lifted.

State arguments may also of type Option<T> or Option<&T>. In this case,
the option is None if the type was never registered or if the shared
state registry was not set up at all.

For non-optional arguments, the API handler fails at runtime if the type
was not registered.

The error-handling in the top-level `api` macro is changed to pass the
token stream to handle_error with the the #[state] attributes removed,
otherwise this leads to bogus 'cannot find attribute `state` in this
scope' errors if there is any other error in the macro.

Signed-off-by: Lukas Wagner <l.wagner@proxmox.com>
---

Notes:
    Changes since v1:
    
      - Fix testcase, it called the wrong api handler
    
    Changes since v2:
      - Instead of using a newtype to mark state arguments, use attributes
      - Support taking state arguments by reference
      - Support taking optional state arguments (by reference or by value)

 proxmox-api-macro/src/api/method.rs | 228 ++++++++++++++++++++++++--
 proxmox-api-macro/src/api/mod.rs    |   6 +
 proxmox-api-macro/src/lib.rs        |  56 ++++++-
 proxmox-api-macro/tests/state.rs    | 240 ++++++++++++++++++++++++++++
 4 files changed, 520 insertions(+), 10 deletions(-)
 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..ab5a725a 100644
--- a/proxmox-api-macro/src/api/method.rs
+++ b/proxmox-api-macro/src/api/method.rs
@@ -3,8 +3,10 @@
 //! 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` and `ApiMethod` parameters as
+//! well as parameters marked with `#[state]`.
 
+use std::collections::HashSet;
 use std::convert::{TryFrom, TryInto};
 use std::mem;
 
@@ -349,6 +351,7 @@ enum ParameterType {
     Value,
     ApiMethod,
     RpcEnv,
+    State(StateParameter),
     Normal(NormalParameter),
 }
 
@@ -357,6 +360,116 @@ struct NormalParameter {
     entry: ObjectEntry,
 }
 
+/// A parameter marked with `#[state]`.
+#[derive(Debug)]
+struct StateParameter {
+    /// The type to look up in the shared state registry.
+    ty: syn::Type,
+
+    /// Whether the parameter takes the state by reference.
+    by_ref: bool,
+
+    /// Whether the parameter is an `Option`, in which case an unavailable value is not an error.
+    optional: bool,
+}
+
+impl StateParameter {
+    /// The registry is keyed by type, and neither a reference nor an `Option` is the same type as
+    /// the value itself, so both are removed to get at the type to look up. A parameter taking
+    /// the state by reference borrows it straight out of the registry, and an `Option` parameter
+    /// gets `None` instead of failing if no value of the type is available.
+    fn new(pat_type: &syn::PatType) -> Self {
+        let (ty, optional) = match util::is_option_type(&pat_type.ty) {
+            Some(ty) => (ty, true),
+            None => (&*pat_type.ty, false),
+        };
+
+        let syn::Type::Reference(reference) = ty else {
+            return Self {
+                ty: ty.clone(),
+                by_ref: false,
+                optional,
+            };
+        };
+
+        if let Some(mutability) = &reference.mutability {
+            error!(mutability => "'state' parameters cannot be taken by mutable reference");
+        }
+
+        if let Some(lifetime) = &reference.lifetime
+            && lifetime.ident != "_"
+        {
+            error!(lifetime => "'state' parameters cannot have an explicit lifetime");
+        }
+
+        if matches!(&*reference.elem, syn::Type::Reference(_)) {
+            error!(&*reference.elem => "'state' parameters cannot be nested references");
+        }
+
+        Self {
+            ty: (*reference.elem).clone(),
+            by_ref: true,
+            optional,
+        }
+    }
+}
+
+fn param_attributes_mut(input: &mut syn::FnArg) -> &mut Vec<syn::Attribute> {
+    match input {
+        syn::FnArg::Receiver(receiver) => &mut receiver.attrs,
+        syn::FnArg::Typed(pat_type) => &mut pat_type.attrs,
+    }
+}
+
+fn is_state_attribute(attr: &syn::Attribute) -> bool {
+    attr.style == syn::AttrStyle::Outer && attr.path().is_ident("state")
+}
+
+/// Removes `#[state]` attributes from the function's parameters and returns the
+/// positions of the parameters which had one.
+///
+/// The attribute only has meaning to this macro, so it has to be stripped before the function is
+/// emitted again. Other attributes are preserved.
+fn take_state_attributes(sig: &mut syn::Signature) -> HashSet<usize> {
+    let mut state_param_indices = HashSet::new();
+
+    for (index, input) in sig.inputs.iter_mut().enumerate() {
+        let attrs = param_attributes_mut(input);
+
+        for attr in mem::take(attrs) {
+            if !is_state_attribute(&attr) {
+                attrs.push(attr);
+                continue;
+            }
+
+            if !matches!(attr.meta, syn::Meta::Path(_)) {
+                error!(&attr.meta => "'state' attribute does not take any data");
+            }
+
+            state_param_indices.insert(index);
+        }
+    }
+
+    state_param_indices
+}
+
+/// Removes `#[state]` attributes from a function without any diagnostics.
+///
+/// Used for the copy of the original function which is emitted alongside a fatal error, where the
+/// leftover attributes would only add confusing follow-up errors due to `state` not being in
+/// scope.
+pub(crate) fn strip_state_attributes(item: TokenStream) -> TokenStream {
+    let Ok(mut func) = syn::parse2::<syn::ItemFn>(item.clone()) else {
+        return item;
+    };
+
+    for input in func.sig.inputs.iter_mut() {
+        param_attributes_mut(input).retain(|attr| !is_state_attribute(attr));
+    }
+
+    func.into_token_stream()
+}
+
 fn check_input_type(input: &syn::FnArg) -> Result<(&syn::PatType, &syn::PatIdent), syn::Error> {
     // `self` types are not supported:
     let pat_type = match input {
@@ -374,15 +487,22 @@ fn check_input_type(input: &syn::FnArg) -> Result<(&syn::PatType, &syn::PatIdent
 }
 
 fn handle_function_signature(method_info: &mut MethodInfo) -> Result<Ident, Error> {
+    let state_params = take_state_attributes(&mut method_info.func.sig);
+
     let sig = &method_info.func.sig;
 
     let mut api_method_param = None;
     let mut rpc_env_param = None;
+    let mut rpc_env_mutable = false;
     let mut value_param = None;
 
     let mut param_list = Vec::<(FieldName, ParameterType)>::new();
 
-    for input in sig.inputs.iter() {
+    for (index, input) in sig.inputs.iter().enumerate() {
+        if state_params.contains(&index) {
+            continue;
+        }
+
         let (pat_type, pat) = match check_input_type(input) {
             Ok(input) => input,
             Err(err) => {
@@ -410,7 +530,7 @@ fn handle_function_signature(method_info: &mut MethodInfo) -> Result<Ident, Erro
         };
     }
 
-    for input in sig.inputs.iter() {
+    for (index, input) in sig.inputs.iter().enumerate() {
         let (pat_type, pat) = match check_input_type(input) {
             Ok(input) => input,
             Err(_err) => continue, // we already produced errors above,
@@ -421,17 +541,20 @@ fn handle_function_signature(method_info: &mut MethodInfo) -> Result<Ident, Erro
         // values, and our 2 fixed function parameters: `&ApiMethod` and `&mut dyn RpcEnvironment`.
         //
         // Our strategy is as follows:
-        //     1) See if the parameter name also appears in the input schema. In this case we
+        //     1) See if the parameter was marked with `#[state]`. Such parameters are filled in
+        //        from the environment's shared state registry and are not part of the schema.
+        //
+        //     2) See if the parameter name also appears in the input schema. In this case we
         //        assume that we want the parameter to be extracted from the `Value` and passed
         //        directly to the function.
         //
-        //     2) Check the parameter type for `&ApiMethod` and remember its position (since we may
+        //     3) Check the parameter type for `&ApiMethod` and remember its position (since we may
         //        need to reorder it!)
         //
-        //     3) Check the parameter type for `&dyn RpcEnvironment` and remember its position
+        //     4) Check the parameter type for `&dyn RpcEnvironment` and remember its position
         //        (since we may need to reorder it!).
         //
-        //     4) Check for a `Value` or `serde_json::Value` parameter. This becomes the
+        //     5) Check for a `Value` or `serde_json::Value` parameter. This becomes the
         //        "catch-all" parameter and only 1 may exist.
         //        Note that we may still use further `Value` parameters if they have been
         //        explicitly named in the `input_schema`. However, only 1 unnamed `Value` parameter
@@ -439,11 +562,25 @@ fn handle_function_signature(method_info: &mut MethodInfo) -> Result<Ident, Erro
         //        If no such parameter exists, we automatically fail the function if the `Value` is
         //        not empty after extracting the parameters.
         //
-        //     5) Finally, if none of the above conditions are met, we do not know what to do and
+        //     6) Finally, if none of the above conditions are met, we do not know what to do and
         //        bail out with an error.
         let pat_ident = pat.ident.unraw();
         let mut param_name: FieldName = pat_ident.clone().into();
-        let param_type = if let Some(entry) = method_info
+        let param_type = if state_params.contains(&index) {
+            if method_info
+                .input_schema
+                .find_obj_property_by_ident(&pat_ident.to_string())
+                .is_some()
+            {
+                error!(
+                    pat_type => "'state' parameter '{}' must not appear in the input schema",
+                    pat_ident,
+                );
+                continue;
+            }
+
+            ParameterType::State(StateParameter::new(pat_type))
+        } else if let Some(entry) = method_info
             .input_schema
             .find_obj_property_by_ident(&pat_ident.to_string())
         {
@@ -469,6 +606,8 @@ fn handle_function_signature(method_info: &mut MethodInfo) -> Result<Ident, Erro
                 continue;
             }
             rpc_env_param = Some(param_list.len());
+            rpc_env_mutable =
+                matches!(&*pat_type.ty, syn::Type::Reference(r) if r.mutability.is_some());
             ParameterType::RpcEnv
         } else if is_value_type(&pat_type.ty) {
             if value_param.is_some() {
@@ -485,6 +624,23 @@ fn handle_function_signature(method_info: &mut MethodInfo) -> Result<Ident, Erro
         param_list.push((param_name, param_type));
     }
 
+    // the borrow of a `#[state]` reference reaches into the environment, so it cannot be handed
+    // out together with a mutable borrow of the very same environment:
+    if rpc_env_mutable {
+        for (name, param) in &param_list {
+            if let ParameterType::State(state) = param
+                && state.by_ref
+            {
+                error!(
+                    name.span(),
+                    "'state' parameter '{}' cannot be borrowed in a method also taking a mutable \
+                     'RpcEnvironment' parameter, take the state by value instead",
+                    name.as_str(),
+                );
+            }
+        }
+    }
+
     create_wrapper_function(method_info, param_list)
 }
 
@@ -563,6 +719,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 +983,57 @@ 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 ty = param.ty;
+
+    // SharedStateRegistry::lookup() returns a borrowed value, so a parameter taking the state
+    // by value needs the reference to be cloned.
+    let clone = (!param.by_ref).then(|| quote_spanned! { span => .cloned() });
+
+    let lookup = if param.optional {
+        // an unavailable value is not an error here, which includes an environment without any
+        // registry at all.
+        quote_spanned! { span =>
+            let #arg_name = rpc_env_param
+                .shared_state()
+                .and_then(|shared_state| shared_state.lookup::<#ty>())
+                #clone;
+        }
+    } else {
+        quote_spanned! { span =>
+            let #arg_name = rpc_env_param
+                .shared_state()
+                .ok_or_else(|| ::anyhow::format_err!(
+                    "router has no shared state set up",
+                ))?
+                .lookup::<#ty>()
+                #clone
+                .ok_or_else(|| ::anyhow::format_err!(
+                    "shared state type '{}' for parameter '{}' not registered",
+                    ::std::any::type_name::<#ty>(),
+                    #name_str,
+                ))?;
+        }
+    };
+
+    body.extend(quote_spanned! { span =>
+        #[allow(non_snake_case)]
+        #lookup
+    });
+
+    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/src/api/mod.rs b/proxmox-api-macro/src/api/mod.rs
index b87c49eb..be5ab9f3 100644
--- a/proxmox-api-macro/src/api/mod.rs
+++ b/proxmox-api-macro/src/api/mod.rs
@@ -790,6 +790,12 @@ impl SchemaArray {
     }
 }
 
+/// Removes the attributes `#[api]` consumes itself, so that the item can be emitted as-is when the
+/// macro fails with a fatal error.
+pub(crate) fn strip_helper_attributes(item: TokenStream) -> TokenStream {
+    method::strip_state_attributes(item)
+}
+
 /// Parse `input`, `returns` and `protected` attributes out of an function annotated
 /// with an `#[api]` attribute and produce a `const ApiMethod` named after the function.
 ///
diff --git a/proxmox-api-macro/src/lib.rs b/proxmox-api-macro/src/lib.rs
index c7aebfae..a9d62e8f 100644
--- a/proxmox-api-macro/src/lib.rs
+++ b/proxmox-api-macro/src/lib.rs
@@ -210,6 +210,59 @@ fn router_do(item: TokenStream) -> Result<TokenStream, Error> {
     }
     ```
 
+    # Shared state parameters.
+
+    Parameters marked with `#[state]` are not part of the API schema. Instead of being extracted
+    from the API input, they are filled in from the `SharedStateRegistry` of the environment
+    the handler is called with. The handler receives a clone of the value registered for the
+    parameter's type and fails if no value of that type was registered:
+
+    ```no_run
+    # use proxmox_api_macro::api;
+    # use anyhow::Error;
+    #[derive(Clone)]
+    struct AuthContext {
+        realm: String,
+    }
+
+    #[api]
+    /// Show which realm the daemon was configured with.
+    fn print_realm(#[state] auth: AuthContext) -> Result<(), Error> {
+        println!("realm: {}", auth.realm);
+
+        Ok(())
+    }
+    ```
+
+    Such a parameter can also be taken by shared reference to avoid the clone, in which case it
+    borrows the value straight out of the registry. The registry is keyed by type, so the type the
+    reference points to is what has to be registered. The reference must not have an explicit
+    lifetime.
+
+    Since the state argument is borrowed from the RpcEnvironemnt, it cannot be combined with a
+    `&mut dyn RpcEnvironment` parameter.
+
+    Both forms can be wrapped in an `Option`, which turns an unavailable value from an error into a
+    `None`. This covers a type which was not registered as well as an environment without any
+    registry at all.
+
+    ```no_run
+    # use proxmox_api_macro::api;
+    # use anyhow::Error;
+    # #[derive(Clone)]
+    # struct AuthContext { realm: String }
+    #[api]
+    /// Show which realm the daemon was configured with, if any.
+    fn print_realm(#[state] auth: Option<&AuthContext>) -> Result<(), Error> {
+        match auth {
+            Some(auth) => println!("realm: {}", auth.realm),
+            None => println!("no auth context configured"),
+        }
+
+        Ok(())
+    }
+    ```
+
     # Deprecated property aliases
 
     Use this when a property has been renamed and you want to keep the old name accepted for
@@ -504,7 +557,8 @@ fn router_do(item: TokenStream) -> Result<TokenStream, Error> {
 pub fn api(attr: TokenStream_1, item: TokenStream_1) -> TokenStream_1 {
     let _error_guard = init_local_error();
     let item: TokenStream = item.into();
-    handle_error(item.clone(), api::api(attr.into(), item)).into()
+    let fallback = api::strip_helper_attributes(item.clone());
+    handle_error(fallback, api::api(attr.into(), item)).into()
 }
 
 /// *Experimental:* Transform a json-like schema definition into an expression yielding a `Schema`.
diff --git a/proxmox-api-macro/tests/state.rs b/proxmox-api-macro/tests/state.rs
new file mode 100644
index 00000000..8b9ee39d
--- /dev/null
+++ b/proxmox-api-macro/tests/state.rs
@@ -0,0 +1,240 @@
+use proxmox_api_macro::api;
+use proxmox_router::SharedStateRegistry;
+
+use anyhow::Error;
+use serde_json::{Value, json};
+
+#[derive(Clone)]
+struct Foo {
+    val: i32,
+}
+
+#[derive(Clone)]
+struct Unregistered;
+
+#[api(
+    input: {
+        properties: {
+            value: {
+                description: "Something",
+            }
+        }
+    }
+)]
+/// Test multiple state args.
+fn multiple_args(
+    #[state] state_arg1: i32,
+    value: isize,
+    #[state] state_arg2: Foo,
+) -> Result<(), Error> {
+    assert_eq!(state_arg1, 10);
+    assert_eq!(state_arg2.val, 20);
+
+    assert_eq!(value, 50);
+
+    Ok(())
+}
+
+#[api(
+    input: {
+        properties: {
+            value: {
+                description: "Something",
+            }
+        }
+    }
+)]
+/// Test state args taken by reference.
+fn by_reference(
+    #[state] state_arg1: &i32,
+    value: isize,
+    #[state] state_arg2: &Foo,
+) -> Result<(), Error> {
+    assert_eq!(*state_arg1, 10);
+    assert_eq!(state_arg2.val, 20);
+
+    assert_eq!(value, 50);
+
+    Ok(())
+}
+
+#[api]
+/// Test an owned state arg next to a mutable environment.
+fn with_rpcenv(
+    #[state] state_arg: Foo,
+    rpcenv: &mut dyn proxmox_router::RpcEnvironment,
+) -> Result<(), Error> {
+    let _ = rpcenv;
+
+    assert_eq!(state_arg.val, 20);
+
+    Ok(())
+}
+
+#[api]
+/// Test optional state args.
+fn optional_args(
+    #[state] state_arg1: Option<i32>,
+    #[state] state_arg2: Option<&Foo>,
+    #[state] state_arg3: Option<Unregistered>,
+) -> Result<(), Error> {
+    assert_eq!(state_arg1, Some(10));
+    assert_eq!(state_arg2.map(|state| state.val), Some(20));
+    assert!(state_arg3.is_none());
+
+    Ok(())
+}
+
+#[api]
+/// Test optional state args without a registry.
+fn optional_without_registry(
+    #[state] state_arg1: Option<i32>,
+    #[state] state_arg2: Option<&Foo>,
+) -> Result<(), Error> {
+    assert!(state_arg1.is_none());
+    assert!(state_arg2.is_none());
+
+    Ok(())
+}
+
+#[api]
+/// Test not registered state arg.
+fn not_registered(#[state] _state: Foo) -> Result<(), Error> {
+    panic!("should not reach this");
+}
+
+#[api(
+    input: {
+        properties: {
+            unused: {
+                description: "foo",
+            }
+        }
+    }
+)]
+/// Test whether attributes on *other* arguments are preserved.
+fn other_attributes_are_preserved(
+    #[state] _state_arg: Option<Foo>,
+    #[allow(unused)] unused: i32,
+) -> Result<(), Error> {
+    Ok(())
+}
+
+struct RpcEnv {
+    shared_state: Option<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> {
+        self.shared_state.as_ref()
+    }
+}
+
+#[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: Some(shared_state),
+    };
+
+    api_function_multiple_args(
+        json!({
+            "value": 50,
+        }),
+        &API_METHOD_MULTIPLE_ARGS,
+        &mut env,
+    )
+    .expect("func with multiple injected works");
+
+    api_function_by_reference(
+        json!({
+            "value": 50,
+        }),
+        &API_METHOD_BY_REFERENCE,
+        &mut env,
+    )
+    .expect("func with injected references works");
+
+    api_function_with_rpcenv(json!({}), &API_METHOD_WITH_RPCENV, &mut env)
+        .expect("func with injected state and environment works");
+
+    api_function_optional_args(json!({}), &API_METHOD_OPTIONAL_ARGS, &mut env)
+        .expect("func with optional injected state works");
+}
+
+#[test]
+fn test_not_registered() {
+    let mut env = RpcEnv {
+        shared_state: Some(SharedStateRegistry::default()),
+    };
+
+    api_function_not_registered(json!({}), &API_METHOD_NOT_REGISTERED, &mut env)
+        .expect_err("func did not fail");
+
+    api_function_optional_without_registry(
+        json!({}),
+        &API_METHOD_OPTIONAL_WITHOUT_REGISTRY,
+        &mut env,
+    )
+    .expect("optional state args are 'None' if not registered");
+}
+
+#[test]
+fn test_without_registry() {
+    let mut env = RpcEnv { shared_state: None };
+
+    api_function_not_registered(json!({}), &API_METHOD_NOT_REGISTERED, &mut env)
+        .expect_err("func did not fail");
+
+    api_function_optional_without_registry(
+        json!({}),
+        &API_METHOD_OPTIONAL_WITHOUT_REGISTRY,
+        &mut env,
+    )
+    .expect("optional state args are 'None' without a registry");
+}
+
+#[test]
+/// Test whether attributes on *other* arguments are preserved.
+///
+/// This test does not really assert anything, but the called function
+/// will emit a warning during compilation if it does not work.
+fn test_attributes_are_preserved() {
+    let mut env = RpcEnv { shared_state: None };
+
+    api_function_other_attributes_are_preserved(
+        json!({
+            "unused": 3
+        }),
+        &API_METHOD_OTHER_ATTRIBUTES_ARE_PRESERVED,
+        &mut env,
+    )
+    .unwrap();
+}
-- 
2.47.3