From d041954d749f430e579cccd81a13a058e1577e56 Mon Sep 17 00:00:00 2001
From: Xynnn007 <xynnn@linux.alibaba.com>
Date: Wed, 20 Sep 2023 01:43:46 +0800
Subject: [PATCH] cdh/kms: add aliyun KMS support for GetPublicKey

Signed-off-by: Xynnn007 <xynnn@linux.alibaba.com>
---
 .../kms/src/plugins/aliyun/client.rs          | 40 ++++++++++++++++++-
 .../plugins/aliyun/protobuf/dkms_api.proto    | 10 +++++
 confidential-data-hub/kms/src/plugins/mod.rs  | 12 +++++-
 3 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/confidential-data-hub/kms/src/plugins/aliyun/client.rs b/confidential-data-hub/kms/src/plugins/aliyun/client.rs
index 28c8d2e70..7c80d8c10 100644
--- a/confidential-data-hub/kms/src/plugins/aliyun/client.rs
+++ b/confidential-data-hub/kms/src/plugins/aliyun/client.rs
@@ -18,9 +18,11 @@ use serde_json::Value;
 use sha2::{Digest, Sha256};
 use tokio::fs;
 
-use crate::plugins::aliyun::client::dkms_api::{DecryptRequest, EncryptRequest};
+use crate::plugins::aliyun::client::dkms_api::{
+    DecryptRequest, EncryptRequest, GetPublicKeyRequest,
+};
 use crate::plugins::_IN_GUEST_DEFAULT_KEY_PATH;
-use crate::{Annotations, Decrypter, Encrypter, ProviderSettings};
+use crate::{Annotations, Decrypter, Encrypter, ProviderSettings, PubkeyProvider};
 use crate::{Error, Result};
 
 use super::annotations::{AliAnnotations, AliProviderSettings};
@@ -223,6 +225,40 @@ impl Decrypter for AliyunKmsClient {
     }
 }
 
+#[async_trait]
+impl PubkeyProvider for AliyunKmsClient {
+    /// a typical key id of aliyun KMS is like
+    /// `key-shh65012xxxmpi4oxtxxx`
+    async fn get_public_key(&mut self, key_id: &str) -> Result<Vec<u8>> {
+        let get_public_key_request = GetPublicKeyRequest {
+            key_id: key_id.into(),
+        };
+        let mut body = Vec::new();
+        get_public_key_request.encode(&mut body).map_err(|e| {
+            Error::AliyunKmsError(format!(
+                "encode get public key request using protobuf failed: {e}"
+            ))
+        })?;
+        let headers = self.build_headers("GetPublicKey", &body).map_err(|e| {
+            Error::AliyunKmsError(format!(
+                "build get public key request http header failed: {e}"
+            ))
+        })?;
+
+        let res = self
+            .do_request(body, headers)
+            .await
+            .map_err(|e| Error::AliyunKmsError(format!("do request to kms server failed: {e}")))?;
+
+        let decrypt_response = dkms_api::GetPublicKeyResponse::decode(&res[..]).map_err(|e| {
+            Error::AliyunKmsError(format!(
+                "decode decrypt response using protobuf failed: {e}"
+            ))
+        })?;
+        Ok(decrypt_response.public_key.into())
+    }
+}
+
 impl AliyunKmsClient {
     const API_VERSION: &'static str = "dkms-gcs-0.2";
     const SIGNATURE_METHOD: &'static str = "RSA_PKCS1_SHA_256";
diff --git a/confidential-data-hub/kms/src/plugins/aliyun/protobuf/dkms_api.proto b/confidential-data-hub/kms/src/plugins/aliyun/protobuf/dkms_api.proto
index 6a731ac80..d25fea924 100644
--- a/confidential-data-hub/kms/src/plugins/aliyun/protobuf/dkms_api.proto
+++ b/confidential-data-hub/kms/src/plugins/aliyun/protobuf/dkms_api.proto
@@ -37,6 +37,16 @@ message DecryptResponse {
   string PaddingMode = 5;
 }
 
+message GetPublicKeyRequest {
+  string KeyId = 1;
+}
+
+message GetPublicKeyResponse {
+  string KeyId = 1;
+  string PublicKey = 2;
+  string RequestId = 3;
+}
+
 message Error {
   int32 StatusCode = 1;
   string ErrorCode = 2;
diff --git a/confidential-data-hub/kms/src/plugins/mod.rs b/confidential-data-hub/kms/src/plugins/mod.rs
index 5c646d64b..346df74d9 100644
--- a/confidential-data-hub/kms/src/plugins/mod.rs
+++ b/confidential-data-hub/kms/src/plugins/mod.rs
@@ -72,6 +72,9 @@ pub enum PublicKeyProvider {
     #[cfg(feature = "kbs")]
     #[strum(ascii_case_insensitive)]
     Kbs,
+    #[cfg(feature = "aliyun")]
+    #[strum(ascii_case_insensitive)]
+    Aliyun,
 }
 
 /// Create a new [`PubkeyProvider`] by given provider name
@@ -83,13 +86,18 @@ async fn new_public_key_provider(provider_name: &str) -> Result<Box<dyn PubkeyPr
         PublicKeyProvider::Kbs => {
             Ok(Box::new(kbs::KbcClient::new().await?) as Box<dyn PubkeyProvider>)
         }
+        #[cfg(feature = "aliyun")]
+        PublicKeyProvider::Aliyun => Ok(Box::new(
+            aliyun::AliyunKmsClient::from_provider_settings(&ProviderSettings::default()).await?,
+        ) as Box<dyn PubkeyProvider>),
     }
 }
 
 /// Get the public key due to the given `key_id`.
-/// For example:
+/// For example `key_id`:
 ///
-/// public key from KBS: `kbs:///default/key/1`
+/// - KBS: `kbs:///default/key/1`
+/// - Aliyun KMS: `aliyun://key-shh65012626mpi4oxxxxx`
 pub async fn get_public_key(key_id: &str) -> Result<Vec<u8>> {
     let (provider, keyid) = key_id
         .split_once("://")