Support NV GPU Attestation with python nvtrust

APPLICATION/attestation-agent/buildspec.yml

@@ -5,7 +5,7 @@ baseos_version: &BASEOS_VERSION
 
 # 定义镜像仓库信息
 repository_info:
-  acr: &ACR_PROD confidential-ai-registry.cn-shanghai.cr.aliyuncs.com/product/attestation-agent
+  acr: &ACR_PROD confidential-ai-registry.cn-shanghai.cr.aliyuncs.com/product
 
 # 定义镜像测试信息
 t-one:
@@ -30,7 +30,7 @@ images:
       path: Dockerfile.aa
       scene:
         args: []
-        tags: [[0.0.2, latest]]
+        tags: [[0.0.3-beta, latest]]
         registry: [*ACR_PROD]
       # 测试配置
       test_config: [*WORKSPACE, *PROJECT, *TEST_SUITE, *TEST_CONF, *TEST_CASE, *CLOUD_SERVER_TAG[0], '']

APPLICATION/attestation-agent/version.yml

@@ -1,7 +1,7 @@
 # 版本关系依赖表，默认继承 version-base.yml 配置，可重写覆盖
 BaseDependency: ../version-base.yml
 Dependency:
-  name: trustiflux
+  name: attestation-agent
   image_type: application
   versions:
     1.0.1:

APPLICATION/confidential-data-hub/buildspec.yml

@@ -5,7 +5,7 @@ baseos_version: &BASEOS_VERSION
 
 # 定义镜像仓库信息
 repository_info:
-  acr: &ACR_PROD confidential-ai-registry.cn-shanghai.cr.aliyuncs.com/product/confidential-data-hub
+  acr: &ACR_PROD confidential-ai-registry.cn-shanghai.cr.aliyuncs.com/product
 
 # 定义镜像测试信息
 t-one:
@@ -30,7 +30,7 @@ images:
       path: Dockerfile.cdh
       scene:
         args: []
-        tags: [[0.0.1, latest]]
+        tags: [[0.0.3-beta, latest]]
         registry: [*ACR_PROD]
       # 测试配置
       test_config: [*WORKSPACE, *PROJECT, *TEST_SUITE, *TEST_CONF, *TEST_CASE, *CLOUD_SERVER_TAG[0], '']

APPLICATION/confidential-data-hub/version.yml

@@ -1,7 +1,7 @@
 # 版本关系依赖表，默认继承 version-base.yml 配置，可重写覆盖
 BaseDependency: ../version-base.yml
 Dependency:
-  name: trustiflux
+  name: confidential-data-hub
   image_type: application
   versions:
     1.0.1:

Dockerfile.aa

@@ -2,7 +2,7 @@
 # Licensed under the Apache License, Version 2.0, see LICENSE for details.
 # SPDX-License-Identifier: Apache-2.0
 
-FROM registry.cn-hangzhou.aliyuncs.com/lxx/trustiflux:anolisos-latest as builder
+FROM openanolis/anolisos:latest as builder
 
 WORKDIR /tmp
 RUN curl https://download.01.org/intel-sgx/sgx-dcap/1.21/linux/distro/Anolis86/sgx_rpm_local_repo.tgz --output sgx_rpm_local_repo.tgz && \
@@ -28,18 +28,30 @@ RUN cargo +1.79.0 build -p attestation-agent --bin ttrpc-aa --no-default-feature
 
 RUN strip target/x86_64-unknown-linux-gnu/debug/ttrpc-aa
 
-FROM registry.cn-hangzhou.aliyuncs.com/lxx/trustiflux:anolisos-latest
+FROM openanolis/anolisos:latest
 
 WORKDIR /tmp
 RUN curl https://download.01.org/intel-sgx/sgx-dcap/1.21/linux/distro/Anolis86/sgx_rpm_local_repo.tgz --output sgx_rpm_local_repo.tgz && \
     tar zxvf sgx_rpm_local_repo.tgz && \
     yum -y install yum-utils && yum-config-manager --add-repo file:///tmp/sgx_rpm_local_repo && \
     yum -y install epel-release && \
     yum install -y --setopt=install_weak_deps=False --nogpgcheck libtdx-attest-devel && \
+    yum install -y python38 python38-devel libxml2-devel libxslt-devel wget && \
     yum clean all && \
     rm -rf /tmp/*
 
+# Install Aliyun python NVTRUST
+RUN wget https://developer.download.nvidia.cn/compute/cuda/repos/rhel8/x86_64/libnvidia-nscq-570-570.86.15-1.x86_64.rpm && \
+    rpm -ivh libnvidia-nscq-570-570.86.15-1.x86_64.rpm
+RUN curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py && python3.8 get-pip.py
+RUN python3.8 -m pip install lxml
+RUN python3.8 -m pip install nv_attestation_sdk==2.1.2.post6393737 nv_local_gpu_verifier==2.1.1.post6393737 nv_ppcie_verifier==1.1.2.post6393737 -f https://attest-public-cn-beijing.oss-cn-beijing.aliyuncs.com/repo/pip/attest.html
+RUN python3.8 -m pip install --upgrade setuptools
+RUN python3.8 -m pip install --upgrade ecdsa
+
 # Copy binaries
+COPY --from=builder /usr/src/guest-components/gpu_attestation.py /usr/
+COPY --from=builder /usr/src/guest-components/nvswitch_attestation.py /usr/
 COPY --from=builder /usr/src/guest-components/target/x86_64-unknown-linux-gnu/debug/ttrpc-aa /usr/local/bin/attestation-agent
 COPY aa-start.sh /usr/bin/start.sh
 COPY tdx-attest.conf /etc/tdx-attest.conf

attestation-agent/attester/Cargo.toml

@@ -73,7 +73,7 @@ all-attesters = [
 # quotes. It's an unconditional dependency for tdx-attester since that is the only way to
 # generate TDX quotes with upstream kernels.
 tsm-report = ["tempfile"]
-tdx-attester = ["scroll", "tsm-report", "tdx-attest-rs", "eventlog-rs"]
+tdx-attester = ["tokio", "scroll", "tsm-report", "tdx-attest-rs", "eventlog-rs", "tempfile"]
 sgx-attester = ["occlum_dcap"]
 az-snp-vtpm-attester = ["az-snp-vtpm"]
 az-tdx-vtpm-attester = ["az-snp-vtpm-attester", "az-tdx-vtpm"]

attestation-agent/attester/src/tdx/mod.rs

@@ -17,6 +17,8 @@
 use std::fs;
 use std::path::Path;
 use tdx_attest_rs::tdx_report_t;
+use tempfile::NamedTempFile;
+use tokio::process::Command;
 
 mod report;
 mod rtmr;
@@ -66,6 +68,10 @@
     quote: String,
     // Eventlog of Attestation Agent
     aa_eventlog: Option<String>,
+    // GPU Attestation Claims
+    gpu_attestation_token: Option<String>,
+    // GPU Switch Attestation Claims
+    gpu_switch_attestation_token: Option<String>,
 }
 
 #[derive(Debug, Default)]
@@ -144,10 +150,23 @@
             }
         };
 
+        let gpu_attestation_token = get_gpu_token("gpu".to_string()).await.ok().or_else(|| {
+            log::warn!("Get GPU attestation token failed");
+            None
+        });
+
+        let gpu_switch_attestation_token =
+            get_gpu_token("switch".to_string()).await.ok().or_else(|| {
+                log::warn!("Get GPU switch attestation token failed");
+                None
+            });
+
         let evidence = TdxEvidence {
             cc_eventlog,
             quote,
             aa_eventlog,
+            gpu_attestation_token,
+            gpu_switch_attestation_token,
         };
 
         serde_json::to_string(&evidence).context("Serialize TDX evidence failed")
@@ -203,6 +222,36 @@
     }
 }
 
+async fn get_gpu_token(attestation_type: String) -> Result<String> {
+    let script_name = match attestation_type.as_str() {
+        "gpu" => "/usr/gpu_attestation.py".to_string(),
+        "switch" => "/usr/nvswitch_attestation.py".to_string(),
+        _ => return Err(anyhow!("Invalid attestation type: {}", attestation_type)),
+    };
+
+    let mut temp_file = NamedTempFile::new()?;
+    let claims_file_path = temp_file
+        .path()
+        .to_str()
+        .ok_or_else(|| anyhow!("GPU attestation: create claims file failed"))?;
+    let output = Command::new("python3.8")
+        .arg(script_name)
+        .arg(claims_file_path)
+        .output()
+        .await?;
+
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        return Err(anyhow!(
+            "Failed to execute GPU Attestation Python script: {}",
+            stderr
+        ));
+    }
+
+    let token = tokio::fs::read_to_string(temp_file.path()).await?;
+    Ok(token.trim().to_string())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

attestation-agent/kbs_protocol/src/builder.rs

@@ -73,6 +73,7 @@ impl<T> KbsClientBuilder<T> {
     pub fn build(self) -> Result<KbsClient<T>> {
         let mut http_client_builder = reqwest::Client::builder()
             .cookie_store(true)
+            .pool_max_idle_per_host(0)
             .user_agent(format!(
                 "attestation-agent-kbs-client/{}",
                 env!("CARGO_PKG_VERSION")

attestation-agent/kbs_protocol/src/evidence_provider/aa.rs

@@ -44,7 +44,7 @@ impl EvidenceProvider for AAEvidenceProvider {
         };
         let res = self
             .client
-            .get_evidence(context::with_timeout(50 * 1000 * 1000 * 1000), &req)
+            .get_evidence(context::with_timeout(1000 * 1000 * 1000 * 1000), &req)
             .await
             .map_err(|e| Error::AAEvidenceProvider(format!("call ttrpc failed: {e}")))?;
         let evidence = String::from_utf8(res.Evidence)
@@ -59,7 +59,7 @@ impl EvidenceProvider for AAEvidenceProvider {
         };
         let res = self
             .client
-            .get_tee_type(context::with_timeout(50 * 1000 * 1000 * 1000), &req)
+            .get_tee_type(context::with_timeout(1000 * 1000 * 1000 * 1000), &req)
             .await
             .map_err(|e| Error::AAEvidenceProvider(format!("call ttrpc failed: {e}")))?;
 

attestation-agent/kbs_protocol/src/token_provider/aa/mod.rs

@@ -51,7 +51,7 @@ impl TokenProvider for AATokenProvider {
         };
         let bytes = self
             .client
-            .get_token(context::with_timeout(50 * 1000 * 1000 * 1000), &req)
+            .get_token(context::with_timeout(1000 * 1000 * 1000 * 1000), &req)
             .await
             .map_err(|e| Error::AATokenProvider(format!("cal ttrpc failed: {e:?}")))?;
         let message: Message = serde_json::from_slice(&bytes.Token).map_err(|e| {

gpu_attestation.py

@@ -0,0 +1,89 @@
+import os
+import sys
+import json
+import nv_attestation_sdk.attestation as attestation_module
+from nv_attestation_sdk.utils import claim_utils
+
+local_gpu_att_result_policy = '''
+{
+  "version":"3.0",
+  "authorization-rules":{
+    "type": "JWT",
+    "overall-claims": {
+      "iss": "LOCAL_GPU_VERIFIER",
+      "x-nvidia-overall-att-result": true,
+      "x-nvidia-ver": "2.0"
+    },
+    "detached-claims":{
+      "measres": "success",
+      "x-nvidia-gpu-arch-check": true,
+      "x-nvidia-gpu-attestation-report-cert-chain-validated": true,
+      "x-nvidia-gpu-attestation-report-parsed": true,
+      "x-nvidia-gpu-attestation-report-nonce-match": true,
+      "x-nvidia-gpu-attestation-report-signature-verified": true,
+      "x-nvidia-gpu-driver-rim-fetched": true,
+      "x-nvidia-gpu-driver-rim-schema-validated": true,
+      "x-nvidia-gpu-driver-rim-cert-validated": true,
+      "x-nvidia-gpu-driver-rim-signature-verified": true,
+      "x-nvidia-gpu-driver-rim-measurements-available": true,
+      "x-nvidia-gpu-vbios-rim-fetched": true,
+      "x-nvidia-gpu-vbios-rim-schema-validated": true,
+      "x-nvidia-gpu-vbios-rim-cert-validated": true,
+      "x-nvidia-gpu-vbios-rim-signature-verified": true,
+      "x-nvidia-gpu-vbios-rim-measurements-available": true,
+      "x-nvidia-gpu-vbios-index-no-conflict": true
+    }
+  }
+}
+'''
+
+# Suppress stdout
+class Silence:
+    def write(self, msg):
+        pass
+    def flush(self):
+        pass
+def main(output_filename):
+    # Redirect stdout to silence internal messages
+    try:
+        old_stdout = sys.stdout
+        sys.stdout = Silence()
+
+        client_name = "AttestationAgent"
+        attestation_class = attestation_module.Attestation(client_name)
+
+        devices = attestation_module.Devices
+        environment = attestation_module.Environment
+
+        attestation_class.add_verifier(devices.GPU, environment.LOCAL, "", "")
+        evidence_list = attestation_class.get_evidence(ppcie_mode=False)
+        result = attestation_class.attest(evidence_list)
+
+        sys.stdout = old_stdout
+        if result:
+            token = attestation_class.get_token()
+            if token and attestation_class.validate_token(local_gpu_att_result_policy):
+                claims = {}
+                json_array = json.loads(token)
+                if len(json_array) >= 2:
+                    for key, value in json_array[1].items():
+                        for item in value:
+                            if isinstance(item, dict):
+                                for k, v in item.items():
+                                    payload = claim_utils.decode_jwt(v)
+                                    claims[str(k)] = payload
+                with open(output_filename, "w") as f:
+                    f.write(json.dumps(claims))
+                return
+
+        print("GPU Attestation failed or no token obtained", file=sys.stderr)
+        sys.exit(1)
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)
+        sys.exit(1)
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        print("Usage: python <script>.py <output_filename>")
+        sys.exit(1)
+    output_filename = sys.argv[1]
+    main(output_filename)

nvswitch_attestation.py

@@ -0,0 +1,83 @@
+import os
+import sys
+import json
+import nv_attestation_sdk.attestation as attestation_module
+from nv_attestation_sdk.utils import claim_utils
+
+local_nvswitch_att_result_policy = '''
+{
+  "version":"3.0",
+  "authorization-rules":{
+    "type": "JWT",
+    "overall-claims": {
+      "iss": "LOCAL_SWITCH_VERIFIER",
+      "x-nvidia-overall-att-result": true,
+      "x-nvidia-ver": "2.0"
+    },
+    "detached-claims":{
+      "measres": "success",
+      "x-nvidia-switch-arch-check": true,
+      "x-nvidia-switch-bios-rim-measurements-available": true,
+      "x-nvidia-switch-attestation-report-signature-verified": true,
+      "x-nvidia-switch-attestation-report-parsed": true,
+      "x-nvidia-switch-attestation-report-nonce-match": true,
+      "x-nvidia-switch-attestation-report-cert-chain-validated": true,
+      "x-nvidia-switch-bios-rim-schema-validated": true,
+      "x-nvidia-switch-bios-rim-signature-verified": true,
+      "x-nvidia-switch-bios-rim-cert-validated": true,
+      "x-nvidia-switch-bios-rim-fetched": true
+    }
+  }
+}
+'''
+
+# Suppress stdout
+class Silence:
+    def write(self, msg):
+        pass
+    def flush(self):
+        pass
+def main(output_filename):
+    # Redirect stdout to silence internal messages
+    try:
+        old_stdout = sys.stdout
+        sys.stdout = Silence()
+
+        client_name = "AttestationAgent"
+        attestation_class = attestation_module.Attestation(client_name)
+
+        devices = attestation_module.Devices
+        environment = attestation_module.Environment
+
+        attestation_class.add_verifier(devices.SWITCH, environment.LOCAL, "", "")
+        evidence_list = attestation_class.get_evidence(ppcie_mode=False)
+        result = attestation_class.attest(evidence_list)
+
+        sys.stdout = old_stdout
+        if result:
+            token = attestation_class.get_token()
+            if token and attestation_class.validate_token(local_nvswitch_att_result_policy):
+                claims = {}
+                json_array = json.loads(token)
+                if len(json_array) >= 2:
+                    for key, value in json_array[1].items():
+                        for item in value:
+                            if isinstance(item, dict):
+                                for k, v in item.items():
+                                    payload = claim_utils.decode_jwt(v)
+                                    claims[str(k)] = payload
+                with open(output_filename, "w") as f:
+                    f.write(json.dumps(claims))
+                return
+
+        print("NV SWITCH Attestation failed or no token obtained", file=sys.stderr)
+        sys.exit(1)
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)
+        sys.exit(1)
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        print("Usage: python <script>.py <output_filename>")
+        sys.exit(1)
+    output_filename = sys.argv[1]
+    main(output_filename)