AA: add ttrpc client tool

This tool is useful when we do debugging and test things for ttrpc-AA.

Signed-off-by: Xynnn007 <xynnn@linux.alibaba.com>

attestation-agent/attestation-agent/Cargo.toml

@@ -13,6 +13,10 @@ required-features = ["bin", "grpc"]
 name = "ttrpc-aa"
 required-features = ["bin", "ttrpc"]
 
+[[bin]]
+name = "ttrpc-aa-client"
+required-features = ["bin", "ttrpc"]
+
 [dependencies]
 anyhow.workspace = true
 async-trait.workspace = true
@@ -50,7 +54,7 @@ tonic-build = { workspace = true, optional = true }
 ttrpc-codegen = { workspace = true, optional = true }
 
 [features]
-default = ["rust-crypto"]
+default = ["rust-crypto", "bin", "ttrpc"]
 
 # Attestation Token support
 token = []
@@ -72,8 +76,14 @@ all-attesters = [
 ]
 tdx-attester = ["kbs_protocol?/tdx-attester", "attester/tdx-attester"]
 sgx-attester = ["kbs_protocol?/sgx-attester", "attester/sgx-attester"]
-az-snp-vtpm-attester = ["kbs_protocol?/az-snp-vtpm-attester", "attester/az-snp-vtpm-attester"]
-az-tdx-vtpm-attester = ["kbs_protocol?/az-tdx-vtpm-attester", "attester/az-tdx-vtpm-attester"]
+az-snp-vtpm-attester = [
+    "kbs_protocol?/az-snp-vtpm-attester",
+    "attester/az-snp-vtpm-attester",
+]
+az-tdx-vtpm-attester = [
+    "kbs_protocol?/az-tdx-vtpm-attester",
+    "attester/az-tdx-vtpm-attester",
+]
 snp-attester = ["kbs_protocol?/snp-attester", "attester/snp-attester"]
 se-attester = ["kbs_protocol?/se-attester", "attester/se-attester"]
 cca-attester = ["kbs_protocol?/cca-attester", "attester/cca-attester"]

attestation-agent/attestation-agent/build.rs

@@ -21,7 +21,7 @@ fn main() -> std::io::Result<()> {
         let protobuf_customized = ProtobufCustomize::default().gen_mod_rs(false);
 
         Codegen::new()
-            .out_dir("src/bin/ttrpc-aa/ttrpc_protocol")
+            .out_dir("src/bin/ttrpc_dep/ttrpc_protocol")
             .inputs(&protos)
             .include("../protos")
             .rust_protobuf()

attestation-agent/attestation-agent/src/bin/ttrpc-aa-client.rs

@@ -0,0 +1,151 @@
+// Copyright (c) 2023 Alibaba Cloud
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+use base64::Engine;
+use clap::{arg, command, Args, Parser, Subcommand};
+use const_format::concatcp;
+use ttrpc::context;
+use ttrpc_dep::ttrpc_protocol::{
+    attestation_agent::{
+        ExtendRuntimeMeasurementRequest, GetEvidenceRequest, GetTeeTypeRequest, GetTokenRequest,
+    },
+    attestation_agent_ttrpc::AttestationAgentServiceClient,
+};
+
+mod ttrpc_dep;
+
+const TIMEOUT: i64 = 5 * 1000 * 1000 * 1000;
+
+const DEFAULT_UNIX_SOCKET_DIR: &str = "/run/confidential-containers/attestation-agent/";
+const UNIX_SOCKET_PREFIX: &str = "unix://";
+const DEFAULT_ATTESTATION_SOCKET_ADDR: &str = concatcp!(
+    UNIX_SOCKET_PREFIX,
+    DEFAULT_UNIX_SOCKET_DIR,
+    "attestation-agent.sock"
+);
+
+#[derive(Parser)]
+#[command(author)]
+struct Cli {
+    /// Attestation ttRPC Unix socket addr.
+    ///
+    /// This Unix socket address which the Attestation ttRPC service
+    /// will listen to, for example:
+    ///
+    /// `--attestation_sock unix:///tmp/attestation`
+    #[arg(default_value_t = DEFAULT_ATTESTATION_SOCKET_ADDR.to_string(), short, long = "attestation_sock")]
+    attestation_sock: String,
+
+    #[command(subcommand)]
+    operation: Operation,
+}
+
+#[derive(Subcommand)]
+#[command(author, version, about, long_about = None)]
+enum Operation {
+    /// Get the tee type
+    GetTee,
+
+    /// Get evidence
+    GetEvidence(GetEvidenceArgs),
+
+    /// Get attestation token
+    GetToken(GetTokenArgs),
+
+    /// Extend runtime measurement
+    ExtendRuntimeMeasurement(ExtendRuntimeMeasurementArgs),
+}
+
+#[derive(Args)]
+#[command(author, version, about, long_about = None)]
+struct GetEvidenceArgs {
+    /// base64 encodede runtime data
+    #[arg(short, long)]
+    runtime_data: String,
+}
+
+#[derive(Args)]
+#[command(author, version, about, long_about = None)]
+struct GetTokenArgs {
+    /// token type
+    #[arg(short, long)]
+    token_type: String,
+}
+
+#[derive(Args)]
+#[command(author, version, about, long_about = None)]
+struct ExtendRuntimeMeasurementArgs {
+    /// domain name
+    #[arg(short, long)]
+    domain: String,
+
+    /// operation name
+    #[arg(short, long)]
+    operation: String,
+
+    /// content name
+    #[arg(short, long)]
+    content: String,
+}
+
+#[tokio::main]
+pub async fn main() {
+    let args = Cli::parse();
+    let inner =
+        ttrpc::asynchronous::Client::connect(&args.attestation_sock).expect("connect ttrpc socket");
+    let client = AttestationAgentServiceClient::new(inner);
+    match args.operation {
+        Operation::GetTee => {
+            let req = GetTeeTypeRequest {
+                ..Default::default()
+            };
+            let res = client
+                .get_tee_type(context::with_timeout(TIMEOUT), &req)
+                .await
+                .expect("request to AA");
+            println!("{}", res.tee);
+        }
+        Operation::GetEvidence(args) => {
+            let runtime_data = base64::engine::general_purpose::STANDARD
+                .decode(args.runtime_data)
+                .unwrap();
+            let req = GetEvidenceRequest {
+                RuntimeData: runtime_data,
+                ..Default::default()
+            };
+            let res = client
+                .get_evidence(context::with_timeout(TIMEOUT), &req)
+                .await
+                .expect("request to AA");
+            let evidence = String::from_utf8(res.Evidence).unwrap();
+            println!("{evidence}");
+        }
+        Operation::GetToken(get_token_args) => {
+            let req = GetTokenRequest {
+                TokenType: get_token_args.token_type,
+                ..Default::default()
+            };
+            let res = client
+                .get_token(context::with_timeout(TIMEOUT), &req)
+                .await
+                .expect("request to AA");
+            let token = String::from_utf8(res.Token).unwrap();
+            println!("{token}");
+        }
+        Operation::ExtendRuntimeMeasurement(extend_runtime_measurement_args) => {
+            let req = ExtendRuntimeMeasurementRequest {
+                Domain: extend_runtime_measurement_args.domain,
+                Operation: extend_runtime_measurement_args.operation,
+                Content: extend_runtime_measurement_args.content,
+                ..Default::default()
+            };
+            client
+                .extend_runtime_measurement(context::with_timeout(TIMEOUT), &req)
+                .await
+                .expect("request to AA");
+            println!("Extended.");
+        }
+    }
+}

attestation-agent/attestation-agent/src/bin/ttrpc-aa/main.rs → attestation-agent/attestation-agent/src/bin/ttrpc-aa.rs

@@ -3,17 +3,19 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use ::ttrpc::asynchronous::Server;
 use anyhow::*;
 use attestation_agent::AttestationAgent;
 use clap::{arg, command, Parser};
 use const_format::concatcp;
 use log::{debug, info};
-use std::path::Path;
+use std::{collections::HashMap, path::Path, sync::Arc};
 use tokio::signal::unix::{signal, SignalKind};
+use ttrpc::asynchronous::{Server, Service};
+use ttrpc_dep::server::AA;
 
-mod server;
-mod ttrpc_protocol;
+use crate::ttrpc_dep::ttrpc_protocol::attestation_agent_ttrpc::create_attestation_agent_service;
+
+mod ttrpc_dep;
 
 const DEFAULT_UNIX_SOCKET_DIR: &str = "/run/confidential-containers/attestation-agent/";
 const UNIX_SOCKET_PREFIX: &str = "unix://";
@@ -45,6 +47,13 @@ struct Cli {
     config_file: Option<String>,
 }
 
+pub fn start_ttrpc_service(aa: AttestationAgent) -> Result<HashMap<String, Service>> {
+    let service = AA { inner: aa };
+    let service = Arc::new(service);
+    let get_resource_service = create_attestation_agent_service(service);
+    Ok(get_resource_service)
+}
+
 #[tokio::main]
 pub async fn main() -> Result<()> {
     env_logger::init_from_env(env_logger::Env::new().default_filter_or("info"));
@@ -59,7 +68,7 @@ pub async fn main() -> Result<()> {
 
     let mut aa = AttestationAgent::new(cli.config_file.as_deref()).context("start AA")?;
     aa.init().await.context("init AA")?;
-    let att = server::start_ttrpc_service(aa)?;
+    let att = start_ttrpc_service(aa)?;
 
     let mut atts = Server::new()
         .bind(&cli.attestation_sock)

attestation-agent/attestation-agent/src/bin/ttrpc_dep/mod.rs

@@ -0,0 +1,7 @@
+// Copyright (c) 2024 Alibaba Cloud
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+pub mod server;
+pub mod ttrpc_protocol;

attestation-agent/attestation-agent/src/bin/ttrpc-aa/server.rs → attestation-agent/attestation-agent/src/bin/ttrpc_dep/server.rs

@@ -3,29 +3,24 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use ::ttrpc::asynchronous::Service;
 use ::ttrpc::proto::Code;
-use anyhow::*;
 use async_trait::async_trait;
 use attestation_agent::{AttestationAPIs, AttestationAgent};
 use log::{debug, error};
 
-use std::collections::HashMap;
-use std::sync::Arc;
-
-use crate::ttrpc_protocol::attestation_agent::{
-    ExtendRuntimeMeasurementRequest, ExtendRuntimeMeasurementResponse, GetEvidenceRequest,
-    GetEvidenceResponse, GetTeeTypeRequest, GetTeeTypeResponse, GetTokenRequest, GetTokenResponse,
-    UpdateConfigurationRequest, UpdateConfigurationResponse,
-};
-use crate::ttrpc_protocol::attestation_agent_ttrpc::{
-    create_attestation_agent_service, AttestationAgentService,
+use crate::ttrpc_dep::ttrpc_protocol::{
+    attestation_agent::{
+        ExtendRuntimeMeasurementRequest, ExtendRuntimeMeasurementResponse, GetEvidenceRequest,
+        GetEvidenceResponse, GetTeeTypeRequest, GetTeeTypeResponse, GetTokenRequest,
+        GetTokenResponse, UpdateConfigurationRequest, UpdateConfigurationResponse,
+    },
+    attestation_agent_ttrpc::AttestationAgentService,
 };
 
 pub const AGENT_NAME: &str = "attestation-agent";
 
 pub struct AA {
-    inner: AttestationAgent,
+    pub(crate) inner: AttestationAgent,
 }
 
 #[async_trait]
@@ -163,10 +158,3 @@ impl AttestationAgentService for AA {
         ::ttrpc::Result::Ok(reply)
     }
 }
-
-pub fn start_ttrpc_service(aa: AttestationAgent) -> Result<HashMap<String, Service>> {
-    let service = AA { inner: aa };
-    let service = Arc::new(service);
-    let get_resource_service = create_attestation_agent_service(service);
-    Ok(get_resource_service)
-}