From c366c0a465a8172d92c60a25b107e143e3e33243 Mon Sep 17 00:00:00 2001
From: RWDai <27391645+RWDai@users.noreply.github.com>
Date: Fri, 22 Mar 2024 15:00:48 +0800
Subject: [PATCH] update (#650)

---
 gateway/spacegate-lib/Cargo.toml              |  4 ++--
 gateway/spacegate-lib/src/plugin/audit_log.rs | 16 +++++++++++++++-
 gateway/spacegate-lib/src/plugin/auth.rs      |  2 +-
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/gateway/spacegate-lib/Cargo.toml b/gateway/spacegate-lib/Cargo.toml
index bd3824220..6f1dd40fd 100644
--- a/gateway/spacegate-lib/Cargo.toml
+++ b/gateway/spacegate-lib/Cargo.toml
@@ -22,7 +22,7 @@ bios-sdk-invoke = { path = "../../sdk/invoke", features = ["spi_log"] }
 
 jsonpath-rust = "0.3.1"
 bios-auth = { path = "../../support/auth", default-features = false }
-tardis = { workspace = true, features = ["web-client", "build-info"] }
+tardis = { workspace = true, features = ["web-client"] }
 ipnet = "2.8.0"
 
 http = { version = "1" }
@@ -32,4 +32,4 @@ pprof = { version = "0.12.1", features = ["flamegraph"] }
 tardis = { workspace = true, features = ["test", "web-client"] }
 
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
-testcontainers-modules = { workspace = true,features = ["redis"] }
+testcontainers-modules = { workspace = true, features = ["redis"] }
diff --git a/gateway/spacegate-lib/src/plugin/audit_log.rs b/gateway/spacegate-lib/src/plugin/audit_log.rs
index 33839d708..ab714998a 100644
--- a/gateway/spacegate-lib/src/plugin/audit_log.rs
+++ b/gateway/spacegate-lib/src/plugin/audit_log.rs
@@ -16,6 +16,7 @@ use spacegate_shell::hyper::{Request, Response};
 use spacegate_shell::kernel::extension::{EnterTime, PeerAddr, Reflect};
 
 use spacegate_shell::kernel::helper_layers::bidirection_filter::{Bdf, BdfLayer, BoxRespFut};
+use spacegate_shell::kernel::Marker;
 use spacegate_shell::plugin::{JsonValue, MakeSgLayer, Plugin, PluginError};
 use spacegate_shell::{BoxError, SgBody};
 use tardis::basic::dto::TardisContext;
@@ -34,6 +35,7 @@ use tardis::{
 use crate::extension::audit_log_param::AuditLogParam;
 use crate::extension::before_encrypt_body::BeforeEncryptBody;
 use crate::extension::cert_info::{CertInfo, RoleInfo};
+use crate::marker::OpresKey;
 
 pub const CODE: &str = "audit_log";
 #[derive(Serialize, Deserialize, Clone)]
@@ -145,6 +147,7 @@ impl SgFilterAuditLog {
     }
 
     fn req(&self, mut req: Request<SgBody>) -> Result<Request<SgBody>, Response<SgBody>> {
+        let reflect = req.extensions_mut().get_mut::<Reflect>().expect("missing reflect");
         let param = AuditLogParam {
             request_path: req.uri().path().to_string(),
             request_method: req.method().to_string(),
@@ -152,7 +155,18 @@ impl SgFilterAuditLog {
             request_scheme: req.uri().scheme().unwrap_or(&Scheme::HTTP).to_string(),
             request_ip: req.extensions().get::<PeerAddr>().ok_or(PluginError::bad_gateway::<AuditLogPlugin>("[Plugin.AuditLog] missing peer addr"))?.0.ip().to_string(),
         };
-        req.extensions_mut().get_mut::<Reflect>().expect("missing reflect").insert(param);
+        if let Some(opres_config) = OpresKey::extract(&req) {
+            if let Some(cert_info) = reflect.get_mut::<CertInfo>() {
+                cert_info.id = opres_config.ak;
+            } else {
+                reflect.insert(CertInfo {
+                    id: opres_config.ak,
+                    name: None,
+                    roles: vec![],
+                })
+            }
+        };
+        reflect.insert(param);
         Ok(req)
     }
 
diff --git a/gateway/spacegate-lib/src/plugin/auth.rs b/gateway/spacegate-lib/src/plugin/auth.rs
index c897932d8..901327d04 100644
--- a/gateway/spacegate-lib/src/plugin/auth.rs
+++ b/gateway/spacegate-lib/src/plugin/auth.rs
@@ -443,7 +443,7 @@ fn success_auth_result_to_req(auth_result: AuthResult, config: &AuthConfig, req:
             .map(|role| role.iter().map(|r| RoleInfo { id: r.to_string(), name: None }).collect::<Vec<_>>())
             .unwrap_or_default(),
     };
-    parts.extensions.insert(cert_info);
+    parts.extensions.get_mut::<Reflect>().expect("missing reflect").insert(cert_info);
 
     if let Some(mut resp_headers) = auth_result.resp_headers.clone() {
         if resp_headers.contains_key(&config.head_key_crypto) || resp_headers.contains_key(&config.head_key_crypto.to_lowercase()) {