Re: [pbs-devel] [PATCH proxmox v3 3/3] rest-server: propagate rate-limit tags from authenticated users

[Christian Ebner <c.ebner@proxmox.com>]

nit: needs reformatting via cargo fmt

On 11/10/25 2:43 PM, Hannes Laimer wrote:
> Tie REST connections rate-limiter callbacks to a shared tag handle
> so we can push authenticated user IDs into the limiter, keeping
> tags in sync whenever we clear or update them.
> 
> Signed-off-by: Hannes Laimer <h.laimer@proxmox.com>
> ---
>   proxmox-rest-server/src/connection.rs |  11 ++-
>   proxmox-rest-server/src/rest.rs       | 137 +++++++++++++++++++++++++-
>   2 files changed, 141 insertions(+), 7 deletions(-)
> 
> diff --git a/proxmox-rest-server/src/connection.rs b/proxmox-rest-server/src/connection.rs
> index 9511b7cb..ff2ee139 100644
> --- a/proxmox-rest-server/src/connection.rs
> +++ b/proxmox-rest-server/src/connection.rs
> @@ -24,7 +24,7 @@ use tokio_openssl::SslStream;
>   use tokio_stream::wrappers::ReceiverStream;
>   
>   #[cfg(feature = "rate-limited-stream")]
> -use proxmox_http::{RateLimitedStream, ShareableRateLimit};
> +use proxmox_http::{RateLimitedStream, RateLimiterTag, ShareableRateLimit};
>   
>   #[cfg(feature = "rate-limited-stream")]
>   pub type SharedRateLimit = Arc<dyn ShareableRateLimit>;
> @@ -165,7 +165,10 @@ type InsecureClientStreamResult = Pin<Box<InsecureClientStream>>;
>   type ClientStreamResult = Pin<Box<SslStream<InsecureClientStream>>>;
>   
>   #[cfg(feature = "rate-limited-stream")]
> -type LookupRateLimiter = dyn Fn(std::net::SocketAddr) -> (Option<SharedRateLimit>, Option<SharedRateLimit>)
> +type LookupRateLimiter = dyn Fn(
> +        std::net::SocketAddr,
> +        &[RateLimiterTag],
> +    ) -> (Option<SharedRateLimit>, Option<SharedRateLimit>)
>       + Send
>       + Sync
>       + 'static;
> @@ -369,7 +372,9 @@ impl AcceptBuilder {
>   
>           #[cfg(feature = "rate-limited-stream")]
>           let socket = match self.lookup_rate_limiter.clone() {
> -            Some(lookup) => RateLimitedStream::with_limiter_update_cb(socket, move || lookup(peer)),
> +            Some(lookup) => RateLimitedStream::with_limiter_update_cb(socket, move |tags| {
> +                lookup(peer, tags)
> +            }),
>               None => RateLimitedStream::with_limiter(socket, None, None),
>           };
>   
> diff --git a/proxmox-rest-server/src/rest.rs b/proxmox-rest-server/src/rest.rs
> index b76c4bc9..6e9692ae 100644
> --- a/proxmox-rest-server/src/rest.rs
> +++ b/proxmox-rest-server/src/rest.rs
> @@ -29,6 +29,10 @@ use tower_service::Service;
>   use url::form_urlencoded;
>   
>   use proxmox_http::Body;
> +#[cfg(feature = "rate-limited-stream")]
> +use proxmox_http::{RateLimiterTag, RateLimiterTags};
> +#[cfg(not(feature = "rate-limited-stream"))]
> +type RateLimiterTags = ();
>   use proxmox_router::{
>       check_api_permission, ApiHandler, ApiMethod, HttpError, Permission, RpcEnvironment,
>       RpcEnvironmentType, UserInformation,
> @@ -86,10 +90,26 @@ impl RestServer {
>           }
>       }
>   
> -    pub fn api_service(&self, peer: &dyn PeerAddress) -> Result<ApiService, Error> {
> +    #[cfg(not(feature = "rate-limited-stream"))]
> +    pub fn api_service<T>(&self, peer: &T) -> Result<ApiService, Error>
> +    where
> +        T: PeerAddress + ?Sized,
> +    {
> +        Ok(ApiService {
> +            peer: peer.peer_addr()?,
> +            api_config: Arc::clone(&self.api_config),
> +        })
> +    }
> +
> +    #[cfg(feature = "rate-limited-stream")]
> +    pub fn api_service<T>(&self, peer: &T) -> Result<ApiService, Error>
> +    where
> +        T: PeerAddress + PeerRateLimitTags + ?Sized,
> +    {
>           Ok(ApiService {
>               peer: peer.peer_addr()?,
>               api_config: Arc::clone(&self.api_config),
> +            rate_limit_tags: peer.rate_limiter_tag_handle(),
>           })
>       }
>   }
> @@ -185,6 +205,11 @@ pub trait PeerAddress {
>       fn peer_addr(&self) -> Result<std::net::SocketAddr, Error>;
>   }
>   
> +#[cfg(feature = "rate-limited-stream")]
> +pub trait PeerRateLimitTags {
> +    fn rate_limiter_tag_handle(&self) -> Option<Arc<Mutex<RateLimiterTags>>>;
> +}
> +
>   // tokio_openssl's SslStream requires the stream to be pinned in order to accept it, and we need to
>   // accept before the peer address is requested, so let's just generally implement this for
>   // Pin<Box<T>>
> @@ -221,6 +246,41 @@ impl<T: PeerAddress> PeerAddress for proxmox_http::RateLimitedStream<T> {
>       }
>   }
>   
> +#[cfg(feature = "rate-limited-stream")]
> +impl<T: PeerRateLimitTags> PeerRateLimitTags for Pin<Box<T>> {
> +    fn rate_limiter_tag_handle(&self) -> Option<Arc<Mutex<RateLimiterTags>>> {
> +        T::rate_limiter_tag_handle(&**self)
> +    }
> +}
> +
> +#[cfg(feature = "rate-limited-stream")]
> +impl<T: PeerRateLimitTags> PeerRateLimitTags for tokio_openssl::SslStream<T> {
> +    fn rate_limiter_tag_handle(&self) -> Option<Arc<Mutex<RateLimiterTags>>> {
> +        self.get_ref().rate_limiter_tag_handle()
> +    }
> +}
> +
> +#[cfg(feature = "rate-limited-stream")]
> +impl PeerRateLimitTags for tokio::net::TcpStream {
> +    fn rate_limiter_tag_handle(&self) -> Option<Arc<Mutex<RateLimiterTags>>> {
> +        None
> +    }
> +}
> +
> +#[cfg(feature = "rate-limited-stream")]
> +impl PeerRateLimitTags for tokio::net::UnixStream {
> +    fn rate_limiter_tag_handle(&self) -> Option<Arc<Mutex<RateLimiterTags>>> {
> +        None
> +    }
> +}
> +
> +#[cfg(feature = "rate-limited-stream")]
> +impl<T> PeerRateLimitTags for proxmox_http::RateLimitedStream<T> {
> +    fn rate_limiter_tag_handle(&self) -> Option<Arc<Mutex<RateLimiterTags>>> {
> +        self.tag_handle()
> +    }
> +}
> +
>   // Helper [Service] containing the peer Address
>   //
>   // The lower level connection [Service] implementation on
> @@ -233,6 +293,8 @@ impl<T: PeerAddress> PeerAddress for proxmox_http::RateLimitedStream<T> {
>   pub struct ApiService {
>       pub peer: std::net::SocketAddr,
>       pub api_config: Arc<ApiConfig>,
> +    #[cfg(feature = "rate-limited-stream")]
> +    pub rate_limit_tags: Option<Arc<Mutex<RateLimiterTags>>>,
>   }
>   
>   impl ApiService {
> @@ -354,6 +416,10 @@ impl Service<Request<Incoming>> for ApiService {
>               Some(proxied_peer) => proxied_peer,
>               None => self.peer,
>           };
> +        #[cfg(feature = "rate-limited-stream")]
> +        let rate_limit_tags = self.rate_limit_tags.clone();
> +        #[cfg(not(feature = "rate-limited-stream"))]
> +        let rate_limit_tags: Option<Arc<Mutex<RateLimiterTags>>> = None;
>   
>           let header = self.api_config
>               .auth_cookie_name
> @@ -368,7 +434,15 @@ impl Service<Request<Incoming>> for ApiService {
>                });
>   
>           async move {
> -            let mut response = match Arc::clone(&config).handle_request(req, &peer).await {
> +            #[cfg(feature = "rate-limited-stream")]
> +            if let Some(handle) = rate_limit_tags.as_ref() {
> +                handle.lock().unwrap().clear();
> +            }
> +
> +            let mut response = match Arc::clone(&config)
> +                .handle_request(req, &peer, rate_limit_tags.clone())
> +                .await
> +            {
>                   Ok(response) => response,
>                   Err(err) => {
>                       let (err, code) = match err.downcast_ref::<HttpError>() {
> @@ -860,6 +934,8 @@ impl ApiConfig {
>           self: Arc<ApiConfig>,
>           req: Request<Incoming>,
>           peer: &std::net::SocketAddr,
> +        #[cfg_attr(not(feature = "rate-limited-stream"), allow(unused_variables))]
> +        rate_limit_tags: Option<Arc<Mutex<RateLimiterTags>>>,
>       ) -> Result<Response<Body>, Error> {
>           let (parts, body) = req.into_parts();
>           let method = parts.method.clone();
> @@ -890,6 +966,8 @@ impl ApiConfig {
>                       full_path: &path,
>                       relative_path_components,
>                       rpcenv,
> +                    #[cfg(feature = "rate-limited-stream")]
> +                    rate_limit_tags: rate_limit_tags.clone(),
>                   })
>                   .await;
>           }
> @@ -901,13 +979,29 @@ impl ApiConfig {
>           if components.is_empty() {
>               match self.check_auth(&parts.headers, &method).await {
>                   Ok((auth_id, _user_info)) => {
> -                    rpcenv.set_auth_id(Some(auth_id));
> +                    rpcenv.set_auth_id(Some(auth_id.clone()));
> +                    #[cfg(feature = "rate-limited-stream")]
> +                    if let Some(handle) = rate_limit_tags.as_ref() {
> +                        let mut guard = handle.lock().unwrap();
> +                        guard.clear();
> +                        guard.push(RateLimiterTag::User(auth_id));
> +                    }
>                       return Ok(self.get_index(rpcenv, parts).await);
>                   }
>                   Err(AuthError::Generic(_)) => {
> +                    #[cfg(feature = "rate-limited-stream")]
> +                    if let Some(handle) = rate_limit_tags.as_ref() {
> +                        handle.lock().unwrap().clear();
> +                    }
>                       tokio::time::sleep_until(Instant::from_std(delay_unauth_time())).await;
>                   }
> -                Err(AuthError::NoData) => {}
> +                Err(AuthError::NoData) =>
> +                {
> +                    #[cfg(feature = "rate-limited-stream")]
> +                    if let Some(handle) = rate_limit_tags.as_ref() {
> +                        handle.lock().unwrap().clear();
> +                    }
> +                }
>               }
>               Ok(self.get_index(rpcenv, parts).await)
>           } else {
> @@ -975,6 +1069,8 @@ pub struct ApiRequestData<'a> {
>       full_path: &'a str,
>       relative_path_components: &'a [&'a str],
>       rpcenv: RestEnvironment,
> +    #[cfg(feature = "rate-limited-stream")]
> +    rate_limit_tags: Option<Arc<Mutex<RateLimiterTags>>>,
>   }
>   
>   pub(crate) struct Formatted {
> @@ -992,6 +1088,8 @@ impl Formatted {
>               full_path,
>               relative_path_components,
>               mut rpcenv,
> +            #[cfg(feature = "rate-limited-stream")]
> +            rate_limit_tags,
>           }: ApiRequestData<'_>,
>       ) -> Result<Response<Body>, Error> {
>           if relative_path_components.is_empty() {
> @@ -1026,10 +1124,20 @@ impl Formatted {
>           if auth_required {
>               match config.check_auth(&parts.headers, &parts.method).await {
>                   Ok((authid, info)) => {
> +                    #[cfg(feature = "rate-limited-stream")]
> +                    if let Some(handle) = rate_limit_tags.as_ref() {
> +                        let mut guard = handle.lock().unwrap();
> +                        guard.clear();
> +                        guard.push(RateLimiterTag::User(authid.clone()));
> +                    }
>                       rpcenv.set_auth_id(Some(authid));
>                       user_info = info;
>                   }
>                   Err(auth_err) => {
> +                    #[cfg(feature = "rate-limited-stream")]
> +                    if let Some(handle) = rate_limit_tags.as_ref() {
> +                        handle.lock().unwrap().clear();
> +                    }
>                       let err = match auth_err {
>                           AuthError::Generic(err) => err,
>                           AuthError::NoData => {
> @@ -1045,6 +1153,11 @@ impl Formatted {
>                       return Err(err);
>                   }
>               }
> +        } else {
> +            #[cfg(feature = "rate-limited-stream")]
> +            if let Some(handle) = rate_limit_tags.as_ref() {
> +                handle.lock().unwrap().clear();
> +            }
>           }
>   
>           match api_method {
> @@ -1108,6 +1221,8 @@ impl Unformatted {
>               full_path,
>               relative_path_components,
>               mut rpcenv,
> +            #[cfg(feature = "rate-limited-stream")]
> +            rate_limit_tags,
>           }: ApiRequestData<'_>,
>       ) -> Result<Response<Body>, Error> {
>           if relative_path_components.is_empty() {
> @@ -1133,10 +1248,20 @@ impl Unformatted {
>           if auth_required {
>               match config.check_auth(&parts.headers, &parts.method).await {
>                   Ok((authid, info)) => {
> +                    #[cfg(feature = "rate-limited-stream")]
> +                    if let Some(handle) = rate_limit_tags.as_ref() {
> +                        let mut guard = handle.lock().unwrap();
> +                        guard.clear();
> +                        guard.push(RateLimiterTag::User(authid.clone()));
> +                    }
>                       rpcenv.set_auth_id(Some(authid));
>                       user_info = info;
>                   }
>                   Err(auth_err) => {
> +                    #[cfg(feature = "rate-limited-stream")]
> +                    if let Some(handle) = rate_limit_tags.as_ref() {
> +                        handle.lock().unwrap().clear();
> +                    }
>                       let err = match auth_err {
>                           AuthError::Generic(err) => err,
>                           AuthError::NoData => {
> @@ -1153,6 +1278,10 @@ impl Unformatted {
>                   }
>               }
>           } else {
> +            #[cfg(feature = "rate-limited-stream")]
> +            if let Some(handle) = rate_limit_tags.as_ref() {
> +                handle.lock().unwrap().clear();
> +            }
>               user_info = Box::new(EmptyUserInformation {});
>           }
>   



_______________________________________________
pbs-devel mailing list
pbs-devel@lists.proxmox.com
https://lists.proxmox.com/cgi-bin/mailman/listinfo/pbs-devel