Merge branch 'yoshidan:main' into enable-resuming-uploads

kms/README.md

@@ -63,7 +63,7 @@ google-cloud-kms = "version"
  async fn run(config: ClientConfig) {
 
      // Create client
-     let mut client = Client::new(config).await.unwrap();
+     let client = Client::new(config).await.unwrap();
 
      // Key ring
      // create

kms/src/client.rs

@@ -1,21 +1,21 @@
-use std::ops::{Deref, DerefMut};
-use std::time::Duration;
+use std::ops::Deref;
+use std::sync::Arc;
 
 #[cfg(feature = "auth")]
 pub use google_cloud_auth;
-use google_cloud_gax::conn::{ConnectionManager, ConnectionOptions, Environment, Error};
-use google_cloud_googleapis::cloud::kms::v1::key_management_service_client::KeyManagementServiceClient;
+use google_cloud_gax::conn::{ConnectionOptions, Environment, Error};
+
 use google_cloud_token::{NopeTokenSourceProvider, TokenSourceProvider};
 
+use crate::grpc::apiv1::conn_pool::{ConnectionManager, KMS, SCOPES};
 use crate::grpc::apiv1::kms_client::Client as KmsGrpcClient;
-use crate::grpc::apiv1::{AUDIENCE, KMS, SCOPES};
 
 #[derive(Debug)]
 pub struct ClientConfig {
-    pub kms_endpoint: String,
+    pub endpoint: String,
     pub token_source_provider: Box<dyn TokenSourceProvider>,
-    pub timeout: Option<Duration>,
-    pub connect_timeout: Option<Duration>,
+    pub pool_size: Option<usize>,
+    pub connection_option: ConnectionOptions,
 }
 
 #[cfg(feature = "auth")]
@@ -54,10 +54,10 @@ impl ClientConfig {
 impl Default for ClientConfig {
     fn default() -> Self {
         Self {
-            kms_endpoint: KMS.to_string(),
+            endpoint: KMS.to_string(),
             token_source_provider: Box::new(NopeTokenSourceProvider {}),
-            timeout: Some(Duration::from_secs(30)),
-            connect_timeout: Some(Duration::from_secs(30)),
+            pool_size: Some(1),
+            connection_option: ConnectionOptions::default(),
         }
     }
 }
@@ -69,22 +69,16 @@ pub struct Client {
 
 impl Client {
     pub async fn new(config: ClientConfig) -> Result<Self, Error> {
-        let conn_options = ConnectionOptions {
-            timeout: config.timeout,
-            connect_timeout: config.connect_timeout,
-        };
-        let conn_pool = ConnectionManager::new(
-            1,
-            config.kms_endpoint,
-            AUDIENCE,
+        let pool_size = config.pool_size.unwrap_or_default();
+        let cm = ConnectionManager::new(
+            pool_size,
+            config.endpoint.as_str(),
             &Environment::GoogleCloud(config.token_source_provider),
-            &conn_options,
+            &config.connection_option,
         )
         .await?;
-        let conn = conn_pool.conn();
-
         Ok(Self {
-            kms_client: KmsGrpcClient::new(KeyManagementServiceClient::new(conn)),
+            kms_client: KmsGrpcClient::new(Arc::new(cm)),
         })
     }
 }
@@ -97,19 +91,14 @@ impl Deref for Client {
     }
 }
 
-impl DerefMut for Client {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.kms_client
-    }
-}
-
 #[cfg(test)]
 mod tests {
 
     use serial_test::serial;
 
     use google_cloud_googleapis::cloud::kms::v1::{
-        CreateKeyRingRequest, GenerateRandomBytesRequest, GetKeyRingRequest, ListKeyRingsRequest, ProtectionLevel,
+        AsymmetricSignRequest, CreateKeyRingRequest, DecryptRequest, EncryptRequest, GenerateRandomBytesRequest,
+        GetKeyRingRequest, GetPublicKeyRequest, ListKeyRingsRequest, MacSignRequest, MacVerifyRequest, ProtectionLevel,
     };
 
     use crate::client::{Client, ClientConfig};
@@ -129,7 +118,7 @@ mod tests {
     #[tokio::test]
     #[serial]
     async fn test_key_ring() {
-        let (mut client, project) = new_client().await;
+        let (client, project) = new_client().await;
         let key_ring_id = "gcpkmskr1714619260".to_string();
 
         // create
@@ -155,19 +144,31 @@ mod tests {
         // list
         let list_request = ListKeyRingsRequest {
             parent: create_request.parent.to_string(),
-            page_size: 0,
+            page_size: 1,
             page_token: "".to_string(),
             filter: "".to_string(),
             order_by: "".to_string(),
         };
         let list_result = client.list_key_rings(list_request, None).await.unwrap();
-        assert!(!list_result.key_rings.is_empty());
+        assert_eq!(1, list_result.key_rings.len());
+
+        let list_request = ListKeyRingsRequest {
+            parent: create_request.parent.to_string(),
+            page_size: 1,
+            page_token: list_result.next_page_token.to_string(),
+            filter: "".to_string(),
+            order_by: "".to_string(),
+        };
+        let list_result2 = client.list_key_rings(list_request, None).await.unwrap();
+        assert_eq!(1, list_result2.key_rings.len());
+
+        assert_ne!(list_result.key_rings[0].name, list_result2.key_rings[0].name);
     }
 
     #[tokio::test]
     #[serial]
     async fn test_generate_random_bytes() {
-        let (mut client, project) = new_client().await;
+        let (client, project) = new_client().await;
 
         // create
         let create_request = GenerateRandomBytesRequest {
@@ -194,4 +195,84 @@ mod tests {
             random_bytes.data
         )
     }
+
+    #[tokio::test]
+    #[serial]
+    async fn test_asymmetric_sign() {
+        let (client, project) = new_client().await;
+
+        let request = AsymmetricSignRequest {
+            name: format!("projects/{project}/locations/asia-northeast1/keyRings/gcr_test/cryptoKeys/eth-sign/cryptoKeyVersions/1"),
+            digest: None,
+            digest_crc32c: None,
+            data: vec![1,2,3,4,5],
+            data_crc32c: None,
+        };
+        let signature = client.asymmetric_sign(request.clone(), None).await.unwrap();
+        assert!(!signature.signature.is_empty());
+    }
+    #[tokio::test]
+    #[serial]
+    async fn test_get_pubkey() {
+        let (client, project) = new_client().await;
+        let request = GetPublicKeyRequest{
+            name: format!("projects/{project}/locations/asia-northeast1/keyRings/gcr_test/cryptoKeys/eth-sign/cryptoKeyVersions/1"),
+        };
+        let pubkey = client.get_public_key(request.clone(), None).await.unwrap();
+        assert!(!pubkey.pem.is_empty());
+    }
+
+    #[tokio::test]
+    #[serial]
+    async fn test_encrypt_decrypt() {
+        let (client, project) = new_client().await;
+
+        let key = format!("projects/{project}/locations/asia-northeast1/keyRings/gcr_test/cryptoKeys/gcr_test");
+        let data = [1, 2, 3, 4, 5];
+        let request = EncryptRequest {
+            name: key.clone(),
+            plaintext: data.to_vec(),
+            additional_authenticated_data: vec![],
+            plaintext_crc32c: None,
+            additional_authenticated_data_crc32c: None,
+        };
+        let encrypted = client.encrypt(request, None).await.unwrap();
+
+        let request = DecryptRequest {
+            name: key,
+            ciphertext: encrypted.ciphertext.clone(),
+            additional_authenticated_data: vec![],
+            ciphertext_crc32c: None,
+            additional_authenticated_data_crc32c: None,
+        };
+        let raw = client.decrypt(request.clone(), None).await.unwrap();
+        assert_eq!(data.to_vec(), raw.plaintext);
+    }
+
+    #[tokio::test]
+    #[serial]
+    async fn test_mac_sign_verify() {
+        let (client, project) = new_client().await;
+
+        let key = format!(
+            "projects/{project}/locations/asia-northeast1/keyRings/gcr_test/cryptoKeys/mac-test/cryptoKeyVersions/1"
+        );
+        let data = [1, 2, 3, 4, 5];
+        let request = MacSignRequest {
+            name: key.clone(),
+            data: data.to_vec(),
+            data_crc32c: None,
+        };
+        let signature = client.mac_sign(request, None).await.unwrap();
+
+        let request = MacVerifyRequest {
+            name: key,
+            data: data.to_vec(),
+            data_crc32c: None,
+            mac: signature.mac,
+            mac_crc32c: signature.mac_crc32c,
+        };
+        let raw = client.mac_verify(request, None).await.unwrap();
+        assert!(raw.success);
+    }
 }

kms/src/grpc/apiv1/conn_pool.rs

@@ -0,0 +1,33 @@
+use google_cloud_gax::conn::{Channel, Environment};
+use google_cloud_gax::conn::{ConnectionManager as GRPCConnectionManager, ConnectionOptions, Error};
+use google_cloud_googleapis::cloud::kms::v1::key_management_service_client::KeyManagementServiceClient;
+
+pub const AUDIENCE: &str = "https://cloudkms.googleapis.com/";
+pub const KMS: &str = "cloudkms.googleapis.com";
+pub const SCOPES: [&str; 1] = ["https://www.googleapis.com/auth/cloud-platform"];
+
+#[derive(Debug)]
+pub struct ConnectionManager {
+    inner: GRPCConnectionManager,
+}
+
+impl ConnectionManager {
+    pub async fn new(
+        pool_size: usize,
+        domain: &str,
+        environment: &Environment,
+        conn_options: &ConnectionOptions,
+    ) -> Result<Self, Error> {
+        Ok(ConnectionManager {
+            inner: GRPCConnectionManager::new(pool_size, domain, AUDIENCE, environment, conn_options).await?,
+        })
+    }
+
+    pub fn num(&self) -> usize {
+        self.inner.num()
+    }
+
+    pub fn conn(&self) -> KeyManagementServiceClient<Channel> {
+        KeyManagementServiceClient::new(self.inner.conn()).max_decoding_message_size(i32::MAX as usize)
+    }
+}