perf/KBS: decrease memory usage for session map

Before this commit, we never free the memory of the sessions that is
expired, which will let the kbs process memory usage monotonically
increase.

Signed-off-by: Xynnn007 <[email protected]>

kbs/src/api/src/http/attest.rs

@@ -2,7 +2,7 @@
 // Licensed under the Apache License, Version 2.0, see LICENSE for details.
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::raise_error;
+use crate::{raise_error, session::SessionStatus};
 
 use super::*;
 
@@ -15,19 +15,19 @@ use serde_json::json;
 /// POST /auth
 pub(crate) async fn auth(
     request: web::Json<Request>,
-    map: web::Data<SessionMap<'_>>,
+    map: web::Data<SessionMap>,
     timeout: web::Data<i64>,
 ) -> Result<HttpResponse> {
     info!("request: {:?}", &request);
 
-    let session = Session::from_request(&request, *timeout.into_inner())
+    let session = SessionStatus::auth(request.0, **timeout)
         .map_err(|e| Error::FailedAuthentication(format!("Session: {e}")))?;
-    let response = HttpResponse::Ok().cookie(session.cookie()).json(Challenge {
-        nonce: session.nonce().to_string(),
-        extra_params: "".to_string(),
-    });
 
-    let _ = map.sessions.insert(session.id().to_string(), session);
+    let response = HttpResponse::Ok()
+        .cookie(session.cookie())
+        .json(session.challenge());
+
+    map.insert(session);
 
     Ok(response)
 }
@@ -36,7 +36,7 @@ pub(crate) async fn auth(
 pub(crate) async fn attest(
     attestation: web::Json<Attestation>,
     request: HttpRequest,
-    map: web::Data<SessionMap<'_>>,
+    map: web::Data<SessionMap>,
     attestation_service: web::Data<Arc<AttestationService>>,
 ) -> Result<HttpResponse> {
     let cookie = request.cookie(KBS_SESSION_ID).ok_or(Error::MissingCookie)?;
@@ -54,7 +54,7 @@ pub(crate) async fn attest(
         if session.is_expired() {
             raise_error!(Error::ExpiredCookie);
         }
-        (session.tee(), session.nonce().to_string())
+        (session.request().tee, session.challenge().nonce.to_string())
     };
 
     let attestation_str = serde_json::to_string(&attestation)
@@ -81,9 +81,7 @@ pub(crate) async fn attest(
         .await
         .ok_or(Error::InvalidCookie)?;
     let session = session.get_mut();
-    session.set_tee_public_key(attestation.tee_pubkey.clone());
-    session.set_authenticated();
-    session.set_attestation_claims(claims);
+    session.attest(claims);
 
     let body = serde_json::to_string(&json!({
         "token": token,

kbs/src/api/src/http/error.rs

@@ -27,9 +27,6 @@ pub enum Error {
     #[error("Attestation failed: {0}")]
     AttestationFailed(String),
 
-    #[error("Attestation claims get failed: {0}")]
-    AttestationClaimsGetFailed(String),
-
     #[error("Received illegal attestation claims: {0}")]
     AttestationClaimsParseFailed(String),
 

kbs/src/api/src/http/mod.rs

@@ -10,7 +10,7 @@ use crate::policy_engine::PolicyEngine;
 #[cfg(feature = "resource")]
 use crate::resource::{set_secret_resource, Repository, ResourceDesc};
 #[cfg(feature = "as")]
-use crate::session::{Session, SessionMap, KBS_SESSION_ID};
+use crate::session::{SessionMap, KBS_SESSION_ID};
 #[cfg(feature = "resource")]
 use crate::token::AttestationTokenVerifier;
 use actix_web::Responder;

kbs/src/api/src/http/resource.rs

@@ -25,7 +25,7 @@ use super::*;
 pub(crate) async fn get_resource(
     request: HttpRequest,
     repository: web::Data<Arc<RwLock<dyn Repository + Send + Sync>>>,
-    #[cfg(feature = "as")] map: web::Data<SessionMap<'_>>,
+    #[cfg(feature = "as")] map: web::Data<SessionMap>,
     token_verifier: web::Data<Arc<RwLock<dyn AttestationTokenVerifier + Send + Sync>>>,
     #[cfg(feature = "policy")] policy_engine: web::Data<PolicyEngine>,
 ) -> Result<HttpResponse> {
@@ -131,9 +131,11 @@ pub(crate) async fn get_resource(
 #[cfg(feature = "as")]
 async fn get_attest_claims_from_session(
     request: &HttpRequest,
-    map: web::Data<SessionMap<'_>>,
+    map: web::Data<SessionMap>,
 ) -> Result<String> {
     // check cookie
+
+    use crate::session::SessionStatus;
     let cookie = request
         .cookie(KBS_SESSION_ID)
         .ok_or(Error::UnAuthenticatedCookie)?;
@@ -148,23 +150,19 @@ async fn get_attest_claims_from_session(
 
     info!("Cookie {} request to get resource", session.id());
 
-    if !session.is_authenticated() {
-        error!("UnAuthenticated KBS cookie {}", cookie.value());
-        raise_error!(Error::UnAuthenticatedCookie);
-    }
-
     if session.is_expired() {
         error!("Expired KBS cookie {}", cookie.value());
         raise_error!(Error::ExpiredCookie);
     }
 
-    let claims = session
-        .attestation_claims()
-        .ok_or(Error::AttestationClaimsGetFailed(
-            "No attestation claims in the session".into(),
-        ))?;
+    let SessionStatus::Attested {
+        attestation_claims, ..
+    } = session
+    else {
+        raise_error!(Error::UnAuthenticatedCookie);
+    };
 
-    Ok(claims)
+    Ok(attestation_claims.to_owned())
 }
 
 async fn get_attest_claims_from_header(

kbs/src/api/src/lib.rs

@@ -238,6 +238,19 @@ impl ApiServer {
         #[cfg(feature = "as")]
         let sessions = web::Data::new(SessionMap::new());
 
+        #[cfg(feature = "as")]
+        let sessions_clone = sessions.clone();
+
+        #[cfg(feature = "as")]
+        tokio::spawn(async move {
+            loop {
+                tokio::time::sleep(std::time::Duration::from_secs(60)).await;
+                sessions_clone
+                    .sessions
+                    .retain_async(|_, v| !v.is_expired())
+                    .await;
+            }
+        });
         let http_timeout = self.http_timeout;
 
         #[cfg(feature = "resource")]

kbs/src/api/src/session.rs

@@ -4,12 +4,13 @@
 
 use actix_web::cookie::{
     time::{Duration, OffsetDateTime},
-    Cookie, Expiration,
+    Cookie,
 };
-use anyhow::{anyhow, Result};
+use anyhow::{bail, Result};
 use base64::engine::general_purpose::STANDARD;
 use base64::Engine;
-use kbs_types::{Request, Tee, TeePubKey};
+use kbs_types::{Challenge, Request};
+use log::warn;
 use rand::{thread_rng, Rng};
 use semver::Version;
 use uuid::Uuid;
@@ -26,107 +27,113 @@ fn nonce() -> Result<String> {
     Ok(STANDARD.encode(&nonce))
 }
 
-#[allow(dead_code)]
-pub(crate) struct Session<'a> {
-    cookie: Cookie<'a>,
-    nonce: String,
-    tee: Tee,
-    tee_extra_params: Option<String>,
-    tee_pub_key: Option<TeePubKey>,
-    authenticated: bool,
-    attestation_claims: Option<String>,
+/// Finite State Machine model for RCAR handshake
+pub(crate) enum SessionStatus {
+    Authed {
+        request: Request,
+        challenge: Challenge,
+        id: String,
+        timeout: OffsetDateTime,
+    },
+
+    Attested {
+        request: Request,
+        challenge: Challenge,
+        attestation_claims: String,
+        id: String,
+        timeout: OffsetDateTime,
+    },
 }
 
-#[allow(dead_code)]
-impl<'a> Session<'a> {
-    pub fn from_request(req: &Request, timeout: i64) -> Result<Self> {
-        let version = Version::parse(&req.version).map_err(anyhow::Error::from)?;
+macro_rules! impl_member {
+    ($attr: ident, $typ: ident) => {
+        pub fn $attr(&self) -> &$typ {
+            match self {
+                SessionStatus::Authed { $attr, .. } => $attr,
+                SessionStatus::Attested { $attr, .. } => $attr,
+            }
+        }
+    };
+}
+
+impl SessionStatus {
+    pub fn auth(request: Request, timeout: i64) -> Result<Self> {
+        let version = Version::parse(&request.version).map_err(anyhow::Error::from)?;
         if !crate::VERSION_REQ.matches(&version) {
-            return Err(anyhow!("Invalid Request version {}", req.version));
+            bail!("Invalid Request version {}", request.version);
         }
         let id = Uuid::new_v4().as_simple().to_string();
-        let tee_extra_params = if req.extra_params.is_empty() {
-            None
-        } else {
-            Some(req.extra_params.clone())
-        };
-
-        let cookie = Cookie::build(KBS_SESSION_ID, id)
-            .expires(OffsetDateTime::now_utc() + Duration::minutes(timeout))
-            .finish();
-
-        Ok(Session {
-            cookie,
-            nonce: nonce()?,
-            tee: req.tee,
-            tee_extra_params,
-            tee_pub_key: None,
-            authenticated: false,
-            attestation_claims: None,
-        })
-    }
-
-    pub fn id(&self) -> &str {
-        self.cookie.value()
-    }
-
-    pub fn cookie(&self) -> Cookie {
-        self.cookie.clone()
-    }
-
-    pub fn nonce(&self) -> &str {
-        &self.nonce
-    }
-
-    pub fn tee(&self) -> Tee {
-        self.tee
-    }
 
-    pub fn tee_public_key(&self) -> Option<TeePubKey> {
-        self.tee_pub_key.clone()
-    }
-
-    pub fn attestation_claims(&self) -> Option<String> {
-        self.attestation_claims.clone()
-    }
+        let timeout = OffsetDateTime::now_utc() + Duration::minutes(timeout);
 
-    pub fn is_authenticated(&self) -> bool {
-        self.authenticated
-    }
-
-    pub fn set_authenticated(&mut self) {
-        self.authenticated = true
+        Ok(Self::Authed {
+            request,
+            challenge: Challenge {
+                nonce: nonce()?,
+                extra_params: String::new(),
+            },
+            id,
+            timeout,
+        })
     }
 
-    pub fn is_expired(&self) -> bool {
-        if let Some(Expiration::DateTime(time)) = self.cookie.expires() {
-            return OffsetDateTime::now_utc() > time;
+    pub fn cookie<'a>(&self) -> Cookie<'a> {
+        match self {
+            SessionStatus::Authed { id, timeout, .. } => Cookie::build(KBS_SESSION_ID, id.clone())
+                .expires(*timeout)
+                .finish(),
+            SessionStatus::Attested { id, timeout, .. } => {
+                Cookie::build(KBS_SESSION_ID, id.clone())
+                    .expires(*timeout)
+                    .finish()
+            }
         }
-
-        false
     }
 
-    pub fn is_valid(&self) -> bool {
-        self.is_authenticated() && !self.is_expired()
-    }
+    impl_member!(request, Request);
+    impl_member!(challenge, Challenge);
+    impl_member!(id, str);
+    impl_member!(timeout, OffsetDateTime);
 
-    pub fn set_tee_public_key(&mut self, key: TeePubKey) {
-        self.tee_pub_key = Some(key)
+    pub fn is_expired(&self) -> bool {
+        return *self.timeout() < OffsetDateTime::now_utc();
     }
 
-    pub fn set_attestation_claims(&mut self, claims: String) {
-        self.attestation_claims = Some(claims)
+    pub fn attest(&mut self, attestation_claims: String) {
+        match self {
+            SessionStatus::Authed {
+                request,
+                challenge,
+                id,
+                timeout,
+            } => {
+                *self = SessionStatus::Attested {
+                    request: request.clone(),
+                    challenge: challenge.clone(),
+                    attestation_claims,
+                    id: id.clone(),
+                    timeout: *timeout,
+                };
+            }
+            SessionStatus::Attested { .. } => {
+                warn!("session already attested.");
+            }
+        }
     }
 }
 
-pub(crate) struct SessionMap<'a> {
-    pub sessions: scc::HashMap<String, Session<'a>>,
+pub(crate) struct SessionMap {
+    pub sessions: scc::HashMap<String, SessionStatus>,
 }
 
-impl<'a> SessionMap<'a> {
+impl SessionMap {
     pub fn new() -> Self {
         SessionMap {
             sessions: scc::HashMap::new(),
         }
     }
+
+    pub fn insert(&self, session: SessionStatus) {
+        let _ = self.sessions.insert(session.id().to_string(), session);
+    }
 }