Skip to content

Commit

Permalink
[kms] Add cryptographic operation (#265)
Browse files Browse the repository at this point in the history
* support poolable connection

* add key operation

* add test

* add test
  • Loading branch information
yoshidan authored May 8, 2024
1 parent 5af32c3 commit 83ae2b5
Show file tree
Hide file tree
Showing 6 changed files with 353 additions and 228 deletions.
2 changes: 1 addition & 1 deletion kms/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
149 changes: 115 additions & 34 deletions kms/src/client.rs
Original file line number Diff line number Diff line change
@@ -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")]
Expand Down Expand Up @@ -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(),
}
}
}
Expand All @@ -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)),
})
}
}
Expand All @@ -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};
Expand All @@ -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
Expand All @@ -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 {
Expand All @@ -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);
}
}
33 changes: 33 additions & 0 deletions kms/src/grpc/apiv1/conn_pool.rs
Original file line number Diff line number Diff line change
@@ -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)
}
}
Loading

0 comments on commit 83ae2b5

Please sign in to comment.