[PATCH proxmox v5 03/21] api-macro: support #[state] attribute for state injection

[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 types must be `Clone` and are cloned on every lookup, so
a clone should be cheap (e.g. by using an inner Arc).

State arguments may also of type 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)
    
    Changes since v4:
      - Drop support for references for now. For now, injected types
        should be cheaply cloneable, which is the case of PdmApplication.

 proxmox-api-macro/src/api/method.rs | 179 ++++++++++++++++++++++--
 proxmox-api-macro/src/api/mod.rs    |   6 +
 proxmox-api-macro/src/lib.rs        |  48 ++++++-
 proxmox-api-macro/tests/state.rs    | 202 ++++++++++++++++++++++++++++
 4 files changed, 425 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..e52d607c 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,93 @@ 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 is an `Option`, in which case an unavailable value is not an error.
+    optional: bool,
+}
+
+impl StateParameter {
+    /// The registry is keyed by type. `Option<T>` is not the same type as `T` itself, so it
+    /// is removed to get at the type to look up. `Option` parameters
+    /// 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),
+        };
+
+        if let syn::Type::Reference(r) = ty {
+            error!(r => "'state' parameters cannot be taken by reference");
+        }
+
+        Self {
+            ty: ty.clone(),
+            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,6 +464,8 @@ 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;
@@ -382,7 +474,11 @@ fn handle_function_signature(method_info: &mut MethodInfo) -> Result<Ident, Erro
 
     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 +506,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 +517,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 +538,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())
         {
@@ -563,6 +676,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 +940,51 @@ 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;
+
+    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>());
+        }
+    } 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>()
+                .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..3d43bc38 100644
--- a/proxmox-api-macro/src/lib.rs
+++ b/proxmox-api-macro/src/lib.rs
@@ -210,6 +210,51 @@ 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(())
+    }
+    ```
+
+    State arguments can also be taken as 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 +549,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..a2d4ef0e
--- /dev/null
+++ b/proxmox-api-macro/tests/state.rs
@@ -0,0 +1,202 @@
+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]
+/// 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_arg3: Option<Unregistered>,
+) -> Result<(), Error> {
+    assert_eq!(state_arg1, Some(10));
+    assert!(state_arg3.is_none());
+
+    Ok(())
+}
+
+#[api]
+/// Test optional state args without a registry.
+fn optional_without_registry(#[state] state_arg1: Option<i32>) -> Result<(), Error> {
+    assert!(state_arg1.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_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