From 5099e64e1ae7d39e6a50a0359a8d0bc9ffd4461b Mon Sep 17 00:00:00 2001
From: Tyler Fanelli <tfanelli@redhat.com>
Date: Fri, 28 Feb 2025 22:01:14 -0500
Subject: [PATCH] kbs/plugins: Decrypt request body if encryption enabled

If a plugin enabled encrypted requests, treat the request body as
encrypted and decrypt it with the session's server TEE private key.

Signed-off-by: Tyler Fanelli <tfanelli@redhat.com>
---
 kbs/src/api_server.rs          | 26 +++++++++++++++++++++++++-
 kbs/src/attestation/backend.rs | 29 +++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/kbs/src/api_server.rs b/kbs/src/api_server.rs
index 446f55b1e..bd710d8b9 100644
--- a/kbs/src/api_server.rs
+++ b/kbs/src/api_server.rs
@@ -9,6 +9,7 @@ use actix_web::{
 use actix_web_httpauth::headers::authorization::{Authorization, Bearer};
 use anyhow::Context;
 use log::info;
+use openssl::rsa::Padding;
 
 use crate::{
     admin::Admin, config::KbsConfig, jwe::jwe, plugins::PluginManager, policy_engine::PolicyEngine,
@@ -202,7 +203,7 @@ pub(crate) async fn api(
                     plugin_name: plugin_name.to_string(),
                 })?;
 
-            let body = body.to_vec();
+            let mut body = body.to_vec();
             if plugin
                 .validate_auth(&body, query, additional_path, request.method())
                 .await
@@ -232,6 +233,29 @@ pub(crate) async fn api(
                     return Err(Error::PolicyDeny);
                 }
 
+                #[cfg(feature = "as")]
+                if plugin
+                    .request_encrypted(&body, query, additional_path, request.method())
+                    .await
+                    .unwrap()
+                {
+                    let req_key = core
+                        .attestation_service
+                        .get_tee_key_from_session(&request)
+                        .await
+                        .unwrap();
+
+                    let mut decrypted = vec![0u8; req_key.size() as usize];
+
+                    let len =
+                        match req_key.private_decrypt(&body, &mut decrypted, Padding::PKCS1_OAEP) {
+                            Ok(l) => l,
+                            Err(e) => return Err(Error::PluginInternalError { source: e.into() }),
+                        };
+
+                    body = decrypted[..len].to_vec();
+                }
+
                 let response = plugin
                     .handle(&body, query, additional_path, request.method())
                     .await
diff --git a/kbs/src/attestation/backend.rs b/kbs/src/attestation/backend.rs
index 5859e886e..86cab133b 100644
--- a/kbs/src/attestation/backend.rs
+++ b/kbs/src/attestation/backend.rs
@@ -11,6 +11,7 @@ use base64::{engine::general_purpose::STANDARD, Engine};
 use kbs_types::{Attestation, Challenge, Request, Tee};
 use lazy_static::lazy_static;
 use log::{debug, info};
+use openssl::{pkey::Private, rsa::Rsa};
 use rand::{thread_rng, Rng};
 use semver::{BuildMetadata, Prerelease, Version, VersionReq};
 use serde::Deserialize;
@@ -320,6 +321,34 @@ impl AttestationService {
 
         Ok(token.to_owned())
     }
+
+    pub async fn get_tee_key_from_session(
+        &self,
+        request: &HttpRequest,
+    ) -> anyhow::Result<Rsa<Private>> {
+        let cookie = request
+            .cookie(KBS_SESSION_ID)
+            .context("KBS session cookie not found")?;
+
+        let session = self
+            .session_map
+            .sessions
+            .get_async(cookie.value())
+            .await
+            .context("session not found")?;
+
+        let session = session.get();
+
+        if session.is_expired() {
+            bail!("The session is expired");
+        }
+
+        let SessionStatus::Attested { req_key, .. } = session else {
+            bail!("The session is not authorized");
+        };
+
+        Ok(req_key.clone())
+    }
 }
 
 #[cfg(test)]