secGear/0085-fix-multi-thread-request-as-generate-challenge-and-v.patch

From f495edff1e1077f209596bca7f83a614ea5cd139 Mon Sep 17 00:00:00 2001
From: houmingyong <houmingyong@huawei.com>
Date: Fri, 6 Sep 2024 10:38:09 +0800
Subject: [PATCH] fix multi-thread request as generate challenge and verify
 report error

Signed-off-by: houmingyong <houmingyong@huawei.com>
---
 .../attestation/attestation-agent/Cargo.toml  |  1 +
 .../attestation-agent/agent/Cargo.toml        |  2 +
 .../attestation-agent/agent/src/lib.rs        | 92 ++++++++++++-------
 .../attestation-agent/agent/src/session.rs    | 55 +++++++++++
 .../attestation-service/service/src/main.rs   | 14 +++
 .../service/src/restapi/mod.rs                | 32 ++++++-
 .../service/src/result/mod.rs                 |  7 +-
 .../attestation/attestation-types/src/lib.rs  |  2 +
 8 files changed, 169 insertions(+), 36 deletions(-)
 create mode 100644 service/attestation/attestation-agent/agent/src/session.rs

diff --git a/service/attestation/attestation-agent/Cargo.toml b/service/attestation/attestation-agent/Cargo.toml
index bdc7b120..f6f31b18 100644
--- a/service/attestation/attestation-agent/Cargo.toml
+++ b/service/attestation/attestation-agent/Cargo.toml
@@ -24,6 +24,7 @@ jsonwebtoken = "9.3.0"
 thiserror = "1.0"
 actix-web = "4.5"
 clap = { version = "4.5.7", features = ["derive"] }
+scc = "2.1"
 
 verifier = {path = "../attestation-service/verifier", default-features = false}
 attestation-types = {path = "../attestation-types"}
diff --git a/service/attestation/attestation-agent/agent/Cargo.toml b/service/attestation/attestation-agent/agent/Cargo.toml
index e29f89be..d2450c87 100644
--- a/service/attestation/attestation-agent/agent/Cargo.toml
+++ b/service/attestation/attestation-agent/agent/Cargo.toml
@@ -41,6 +41,8 @@ base64-url.workspace = true
 thiserror.workspace = true
 actix-web.workspace = true
 clap.workspace = true
+scc.workspace = true
+attestation-types.workspace = true
 
 attester = { path = "../attester" }
 token_verifier = { path = "../token" }
diff --git a/service/attestation/attestation-agent/agent/src/lib.rs b/service/attestation/attestation-agent/agent/src/lib.rs
index 393914d6..f6e03c6c 100644
--- a/service/attestation/attestation-agent/agent/src/lib.rs
+++ b/service/attestation/attestation-agent/agent/src/lib.rs
@@ -34,9 +34,16 @@ pub type TeeClaim = serde_json::Value;
 use verifier::{Verifier, VerifierAPIs};
 
 #[cfg(not(feature = "no_as"))]
-use {serde_json::json, reqwest, base64_url};
+use {
+    serde_json::json,
+    reqwest::header::{HeaderMap, HeaderValue},
+    base64_url
+};
 
 pub use attester::EvidenceRequest;
+mod session;
+use session::{SessionMap, Session};
+use attestation_types::SESSION_TIMEOUT_MIN;
 
 pub type AsTokenClaim = TokenRawData;
 
@@ -171,6 +178,7 @@ impl TryFrom<&Path> for AAConfig {
 #[derive(Debug)]
 pub struct AttestationAgent {
     config: AAConfig,
+    as_client_sessions: SessionMap,
 }
 
 #[allow(dead_code)]
@@ -186,8 +194,20 @@ impl AttestationAgent {
                 AAConfig::default()
             }
         };
+        let as_client_sessions = SessionMap::new();
+        let sessions = as_client_sessions.clone();
+        tokio::spawn(async move {
+            loop {
+                tokio::time::sleep(std::time::Duration::from_secs(60)).await;
+                sessions
+                    .session_map
+                    .retain_async(|_, v| !v.is_expired())
+                    .await;
+            }
+        });
         Ok(AttestationAgent {
             config,
+            as_client_sessions,
         })
     }
 
@@ -197,16 +217,33 @@ impl AttestationAgent {
         evidence: &[u8],
         policy_id: Option<Vec<String>>
     ) -> Result<String> {
+        let challenge = base64_url::encode(challenge);
         let request_body = json!({
-            "challenge": base64_url::encode(challenge),
+            "challenge": challenge,
             "evidence": base64_url::encode(evidence),
             "policy_id": policy_id,
         });
+        let mut map = HeaderMap::new();
+        map.insert("Content-Type", HeaderValue::from_static("application/json"));
+        let mut client = reqwest::Client::new();
+        if !self.as_client_sessions.session_map.is_empty() {
+            let session = self.as_client_sessions
+                .session_map
+                .get_async(&challenge)
+                .await;
+            match session {
+                Some(entry) => {
+                    map.insert("as-challenge", HeaderValue::from_static("as"));
+                    client = entry.get().as_client.clone()
+                },
+                None => log::info!("challenge is not as generate"),
+            }
+        }
 
         let attest_endpoint = format!("{}/attestation", self.config.svr_url);
-        let res = reqwest::Client::new()
+        let res = client
             .post(attest_endpoint)
-            .header("Content-Type", "application/json")
+            .headers(map)
             .json(&request_body)
             .send()
             .await?;
@@ -249,16 +286,18 @@ impl AttestationAgent {
     }
     async fn get_challenge_from_as(&self) -> Result<String> {
         let challenge_endpoint = format!("{}/challenge", self.config.svr_url);
-        let res = reqwest::Client::new()
+        let client = reqwest::Client::builder()
+            .cookie_store(true)
+            .build()?;
+        let res = client
             .get(challenge_endpoint)
             .header("Content-Type", "application/json")
             .header("content-length", 0)
-            //.json(&request_body)
             .send()
             .await?;
         let challenge = match res.status() {
             reqwest::StatusCode::OK => {
-                let respone = res.text().await?;
+                let respone: String = res.json().await.unwrap();
                 log::debug!("get challenge success, AS Response: {:?}", respone);
                 respone
             }
@@ -267,6 +306,8 @@ impl AttestationAgent {
                 bail!("get challenge Failed")
             }
         };
+        let session = Session::new(challenge.clone(), client, SESSION_TIMEOUT_MIN)?;
+        self.as_client_sessions.insert(session);
         Ok(challenge)
     }
 }
@@ -274,12 +315,19 @@ impl AttestationAgent {
 
 // attestation agent c interface
 use safer_ffi::prelude::*;
-use futures::executor::block_on;
 use tokio::runtime::Runtime;
 
+#[ffi_export]
+pub fn init_env_logger(c_level: Option<&repr_c::String>) {
+    let level = match c_level {
+        Some(level) => &level,
+        None => "info",
+    };
+    env_logger::init_from_env(env_logger::Env::new().default_filter_or(level));
+}
+
 #[ffi_export]
 pub fn get_report(c_challenge: Option<&repr_c::Vec<u8>>, c_ima: &repr_c::TaggedOption<bool>) -> repr_c::Vec<u8> {
-    env_logger::init_from_env(env_logger::Env::new().default_filter_or("info"));
     log::debug!("input challenge: {:?}, ima: {:?}", c_challenge, c_ima);
     let ima = match c_ima {
         repr_c::TaggedOption::None => false,
@@ -295,11 +343,12 @@ pub fn get_report(c_challenge: Option<&repr_c::Vec<u8>>, c_ima: &repr_c::TaggedO
         challenge: challenge,
         ima: Some(ima),
     };
-
+    let rt = Runtime::new().unwrap();
     let fut = async {
         AttestationAgent::new(Some(DEFAULT_AACONFIG_FILE.to_string())).unwrap().get_evidence(input).await
     };
-    let report: Vec<u8> = match block_on(fut) {
+    let ret = rt.block_on(fut);
+    let report: Vec<u8> = match ret {
         Ok(report) => report,
         Err(e) => {
             log::error!("get report failed {:?}", e);
@@ -357,24 +406,3 @@ pub fn generate_headers() -> ::std::io::Result<()> {
         .to_file("./c_header/rust_attestation_agent.h")?
         .generate()
 }
-
-
-#[cfg(test)]
-mod tests {
-    use crate::*;
-
-    #[test]
-    fn aa_new_no_conf_path() {
-        let aa = AttestationAgent::new(None).unwrap();
-        assert_eq!(aa.config.svr_url, "http://127.0.0.1:8080");
-        assert_eq!(aa.config.token_cfg.cert, "/etc/attestation/attestation-agent/as_cert.pem");
-        assert_eq!(aa.config.token_cfg.iss, "openEulerAS");
-    }
-
-    #[test]
-    fn aa_new_with_example_conf() {
-        let aa = AttestationAgent::new(Some("attestation-agent.conf".to_string())).unwrap();
-        assert_eq!(aa.config.token_cfg.cert, "/home/cert/as_cert.pem");
-        assert_eq!(aa.config.token_cfg.iss, "oeas");
-    }
-}
diff --git a/service/attestation/attestation-agent/agent/src/session.rs b/service/attestation/attestation-agent/agent/src/session.rs
new file mode 100644
index 00000000..5e1c1fc5
--- /dev/null
+++ b/service/attestation/attestation-agent/agent/src/session.rs
@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved.
+ * secGear is licensed under the Mulan PSL v2.
+ * You can use this software according to the terms and conditions of the Mulan PSL v2.
+ * You may obtain a copy of Mulan PSL v2 at:
+ *     http://license.coscl.org.cn/MulanPSL2
+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
+ * PURPOSE.
+ * See the Mulan PSL v2 for more details.
+ */
+use actix_web::cookie::{time::{Duration, OffsetDateTime}};
+use scc::HashMap;
+use anyhow::Result;
+
+#[derive(Debug, Clone)]
+pub struct Session {
+    pub challenge: String,
+    pub as_client: reqwest::Client,
+    timeout: OffsetDateTime,
+    // pub token: Option<String>,
+}
+
+impl Session {
+    pub fn new(challenge: String, as_client: reqwest::Client, timeout_m: i64) -> Result<Self> {
+
+        let timeout = OffsetDateTime::now_utc() + Duration::minutes(timeout_m);
+        // let token = None;
+        Ok(Session {
+            challenge,
+            as_client,
+            timeout,
+            // token,
+        })
+    }
+    pub fn is_expired(&self) -> bool {
+        return self.timeout < OffsetDateTime::now_utc();
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct SessionMap {
+    pub session_map: HashMap<String, Session>,
+}
+
+impl SessionMap {
+    pub fn new() -> Self {
+        SessionMap {
+            session_map: HashMap::new(),
+        }
+    }
+    pub fn insert(&self, session: Session) {
+        let _ = self.session_map.insert(session.challenge.clone(), session);
+    }
+}
\ No newline at end of file
diff --git a/service/attestation/attestation-service/service/src/main.rs b/service/attestation/attestation-service/service/src/main.rs
index 3ced10b9..88941b84 100644
--- a/service/attestation/attestation-service/service/src/main.rs
+++ b/service/attestation/attestation-service/service/src/main.rs
@@ -15,6 +15,7 @@ use attestation_service::AttestationService;
 mod restapi;
 use restapi::{get_challenge, attestation, reference, get_policy, set_policy};
 mod session;
+use session::SessionMap;
 
 use anyhow::Result;
 use env_logger;
@@ -54,11 +55,24 @@ async fn main() -> Result<()> {
 
     let cli = Cli::parse();
     let server:AttestationService = AttestationService::new(Some(cli.config)).unwrap();
+    let session_map = web::Data::new(SessionMap::new());
+
+    let sessions_clone = session_map.clone();
+    tokio::spawn(async move {
+        loop {
+            tokio::time::sleep(std::time::Duration::from_secs(60)).await;
+            sessions_clone
+                .session_map
+                .retain_async(|_, v| !v.is_expired())
+                .await;
+        }
+    });
 
     let service = web::Data::new(Arc::new(RwLock::new(server)));
     HttpServer::new(move || {
         App::new()
             .app_data(web::Data::clone(&service))
+            .app_data(web::Data::clone(&session_map))
             .service(get_challenge)
             .service(attestation)
             .service(reference)
diff --git a/service/attestation/attestation-service/service/src/restapi/mod.rs b/service/attestation/attestation-service/service/src/restapi/mod.rs
index 291b8657..a7e6012b 100644
--- a/service/attestation/attestation-service/service/src/restapi/mod.rs
+++ b/service/attestation/attestation-service/service/src/restapi/mod.rs
@@ -10,7 +10,7 @@
  * See the Mulan PSL v2 for more details.
  */
 use attestation_service::AttestationService;
-use attestation_service::result::{Result};
+use attestation_service::result::{Result, Error};
 
 use actix_web::{ post, get, web, HttpResponse, HttpRequest};
 use serde::{Deserialize, Serialize};
@@ -18,6 +18,8 @@ use std::sync::Arc;
 use tokio::sync::RwLock;
 use log;
 use base64_url;
+use attestation_types::SESSION_TIMEOUT_MIN;
+use crate::session::{Session, SessionMap};
 
 const DEFAULT_POLICY_DIR: &str = "/etc/attestation/attestation-service/policy";
 #[derive(Deserialize, Serialize, Debug)]
@@ -25,12 +27,19 @@ pub struct ChallengeRequest {}
 
 #[get("/challenge")]
 pub async fn get_challenge(
+    map: web::Data<SessionMap>,
     service: web::Data<Arc<RwLock<AttestationService>>>,
 ) -> Result<HttpResponse> {
     log::debug!("challenge request");
 
     let challenge = service.read().await.generate_challenge().await;
-    Ok(HttpResponse::Ok().body(challenge))
+    let session = Session::new(challenge, SESSION_TIMEOUT_MIN);
+    let response = HttpResponse::Ok()
+        .cookie(session.cookie())
+        .json(session.challenge.clone());
+    map.insert(session);
+
+    Ok(response)
 }
 
 #[derive(Deserialize, Serialize, Debug)]
@@ -42,6 +51,8 @@ pub struct AttestationRequest {
 
 #[post("/attestation")]
 pub async fn attestation(
+    http_req: HttpRequest,
+    map: web::Data<SessionMap>,
     request: web::Json<AttestationRequest>,
     service: web::Data<Arc<RwLock<AttestationService>>>,
 ) -> Result<HttpResponse> {
@@ -49,6 +60,23 @@ pub async fn attestation(
     let request = request.0;
     let challenge = request.challenge;
 
+    if http_req.headers().contains_key("as-challenge") {
+        log::info!("sessions map len:{}", map.session_map.len());
+        let cookie = http_req.cookie("oeas-session-id").ok_or(Error::CookieMissing)?;
+        let session = map
+            .session_map
+            .get_async(cookie.value())
+            .await
+            .ok_or(Error::SessionNotFound)?;
+        if session.is_expired() {
+            return Err(Error::SessionExpired);
+        }
+        if challenge != session.challenge {
+            log::error!("request challenge:{} does not match session challenge:{}", challenge, session.challenge);
+            return Err(Error::ChallengeInvalid);
+        }
+    }
+
     let nonce = base64_url::decode(&challenge).expect("base64 decode nonce");
     let evidence = base64_url::decode(&request.evidence).expect("base64 decode evidence");
     let ids = request.policy_id;
diff --git a/service/attestation/attestation-service/service/src/result/mod.rs b/service/attestation/attestation-service/service/src/result/mod.rs
index 667e80f5..fcb1c123 100644
--- a/service/attestation/attestation-service/service/src/result/mod.rs
+++ b/service/attestation/attestation-service/service/src/result/mod.rs
@@ -38,12 +38,15 @@ pub enum Error {
     #[error("Request cookie is missing")]
     CookieMissing,
 
-    #[error("Request cookie is not found")]
-    CookieNotFound,
+    #[error("Request cookie session is not found")]
+    SessionNotFound,
 
     #[error("The session of request cookie is expired")]
     SessionExpired,
 
+    #[error("Request challenge is invalid")]
+    ChallengeInvalid,
+
     #[error(transparent)]
     Other(#[from] anyhow::Error),
 }
diff --git a/service/attestation/attestation-types/src/lib.rs b/service/attestation/attestation-types/src/lib.rs
index fcf1d3ee..67dcf9f8 100644
--- a/service/attestation/attestation-types/src/lib.rs
+++ b/service/attestation/attestation-types/src/lib.rs
@@ -12,6 +12,8 @@
 use serde::{Serialize, Deserialize};
 use serde_json::Value;
 
+pub const SESSION_TIMEOUT_MIN: i64 = 1;
+
 #[derive(Debug, Serialize, Deserialize)]
 pub struct VirtccaEvidence {
     pub evidence: Vec<u8>,
-- 
2.46.0