[pbs-devel] [PATCH proxmox 3/3] rest-server: add use tag field to RateLimitedStreams

[Hannes Laimer]

Similarly to how the IP is attached we also attach the user that is
authenticated on this connections. Since this is only used for rate
limiting this is behind the "rate-limited-stream" feature.

Signed-off-by: Hannes Laimer <h.laimer at proxmox.com>
---
 proxmox-rest-server/src/connection.rs | 16 +++++-
 proxmox-rest-server/src/rest.rs       | 72 ++++++++++++++++++++++++++-
 2 files changed, 85 insertions(+), 3 deletions(-)

diff --git a/proxmox-rest-server/src/connection.rs b/proxmox-rest-server/src/connection.rs
index 9511b7cb..a9c3ccb3 100644
--- a/proxmox-rest-server/src/connection.rs
+++ b/proxmox-rest-server/src/connection.rs
@@ -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,
+        Option<String>,
+    ) -> (Option<SharedRateLimit>, Option<SharedRateLimit>)
     + Send
     + Sync
     + 'static;
@@ -369,7 +372,16 @@ 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) => {
+                let user_tag: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
+                let user_tag_cb = Arc::clone(&user_tag);
+                let mut s = RateLimitedStream::with_limiter_update_cb(socket, move || {
+                    let user = user_tag_cb.lock().unwrap().clone();
+                    lookup(peer, user)
+                });
+                s.set_user_tag_handle(user_tag);
+                s
+            }
             None => RateLimitedStream::with_limiter(socket, None, None),
         };
 
diff --git a/proxmox-rest-server/src/rest.rs b/proxmox-rest-server/src/rest.rs
index 035a9537..c7d833a2 100644
--- a/proxmox-rest-server/src/rest.rs
+++ b/proxmox-rest-server/src/rest.rs
@@ -86,10 +86,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 + PeerUser + ?Sized,
+    {
         Ok(ApiService {
             peer: peer.peer_addr()?,
             api_config: Arc::clone(&self.api_config),
+            user_tag: peer.user_tag_handle(),
         })
     }
 }
@@ -185,6 +201,11 @@ pub trait PeerAddress {
     fn peer_addr(&self) -> Result<std::net::SocketAddr, Error>;
 }
 
+#[cfg(feature = "rate-limited-stream")]
+pub trait PeerUser {
+    fn user_tag_handle(&self) -> Option<Arc<Mutex<Option<String>>>>;
+}
+
 // 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 +242,41 @@ impl<T: PeerAddress> PeerAddress for proxmox_http::RateLimitedStream<T> {
     }
 }
 
+#[cfg(feature = "rate-limited-stream")]
+impl<T: PeerUser> PeerUser for Pin<Box<T>> {
+    fn user_tag_handle(&self) -> Option<Arc<Mutex<Option<String>>>> {
+        T::user_tag_handle(&**self)
+    }
+}
+
+#[cfg(feature = "rate-limited-stream")]
+impl<T: PeerUser> PeerUser for tokio_openssl::SslStream<T> {
+    fn user_tag_handle(&self) -> Option<Arc<Mutex<Option<String>>>> {
+        self.get_ref().user_tag_handle()
+    }
+}
+
+#[cfg(feature = "rate-limited-stream")]
+impl PeerUser for tokio::net::TcpStream {
+    fn user_tag_handle(&self) -> Option<Arc<Mutex<Option<String>>>> {
+        None
+    }
+}
+
+#[cfg(feature = "rate-limited-stream")]
+impl PeerUser for tokio::net::UnixStream {
+    fn user_tag_handle(&self) -> Option<Arc<Mutex<Option<String>>>> {
+        None
+    }
+}
+
+#[cfg(feature = "rate-limited-stream")]
+impl<T> PeerUser for proxmox_http::RateLimitedStream<T> {
+    fn user_tag_handle(&self) -> Option<Arc<Mutex<Option<String>>>> {
+        self.user_tag_handle()
+    }
+}
+
 // Helper [Service] containing the peer Address
 //
 // The lower level connection [Service] implementation on
@@ -233,6 +289,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 user_tag: Option<Arc<Mutex<Option<String>>>>,
 }
 
 impl ApiService {
@@ -357,6 +415,8 @@ impl Service<Request<Incoming>> for ApiService {
             Some(proxied_peer) => proxied_peer,
             None => self.peer,
         };
+        #[cfg(feature = "rate-limited-stream")]
+        let user_tag = self.user_tag.clone();
 
         let header = self.api_config
             .auth_cookie_name
@@ -394,6 +454,16 @@ impl Service<Request<Incoming>> for ApiService {
                 }
             }
 
+            #[cfg(feature = "rate-limited-stream")]
+            {
+                if let Some(handle) = user_tag {
+                    if let Some(ext) = response.extensions().get::<AuthStringExtension>() {
+                        let mut guard = handle.lock().unwrap();
+                        *guard = Some(ext.0.clone());
+                    }
+                }
+            }
+
             let logger = config.get_access_log();
             log_response(logger, &peer, method, &path, &response, user_agent);
             Ok(response)
-- 
2.47.2