sdk/oak_tests: add gRPC client helper
Return a channel and interceptor that can be used to construct service-specific
client stub instances.
//! Test utilities to help with unit testing of Oak SDK code.
use log::{debug, info};
-use oak_abi::proto::oak::application::{
+use oak_abi::{
+ label::Label,
+ proto::oak::application::{
node_configuration::ConfigType, ApplicationConfiguration, NodeConfiguration,
WebAssemblyConfiguration,
+ },
};
use prost::Message;
use std::{collections::HashMap, path::PathBuf, process::Command};
-use tonic::transport::Identity;
+use tonic::{
+ metadata::MetadataValue,
+ transport::{Certificate, Channel, ClientTlsConfig, Identity},
+ Request,
+};
// TODO(#544): re-enable unit tests of SDK functionality
@@ -58,10 +65,16 @@ pub fn compile_rust_wasm(cargo_path: &str, module_name: &str) -> std::io::Result
/// Default entrypoint name for the module under test.
const DEFAULT_ENTRYPOINT_NAME: &str = "oak_main";
+/// Default URI that the tests expect to find a live Runtime at.
+const RUNTIME_URI: &str = "https://localhost:8080";
const DEFAULT_MODULE_MANIFEST: &str = "Cargo.toml";
const MODULE_WASM_SUFFIX: &str = ".wasm";
+// Retry parameters when connecting to a gRPC server.
+const RETRY_COUNT: u32 = 360;
+const RETRY_INTERVAL: std::time::Duration = std::time::Duration::from_millis(500);
+
/// Convenience helper to build and run a single-Node Application with the given module name, using
/// the default name "oak_main" for its entrypoint.
pub fn run_single_module_default(
@@ -233,3 +246,54 @@ where
return Ok(rsp);
}
}
+
+/// Build a channel and interceptor suitable for building a client that connects
+/// to a Runtime under test.
+pub async fn channel_and_interceptor() -> (Channel, impl Into<tonic::Interceptor>) {
+ // Build a channel that connects to the Runtime under test.
+ let uri = RUNTIME_URI.parse().expect("Error parsing URI");
+ let tls_config = ClientTlsConfig::new()
+ .ca_certificate(Certificate::from_pem(include_str!("../certs/ca.pem")));
+ let builder = Channel::builder(uri).tls_config(tls_config);
+
+ // The Runtime may have just been started for a test, and make take some time
+ // to come fully up, start a gRPC server and accept connections. Allow for
+ // this by retrying at intervals until the server responds or we hit a retry
+ // limit.
+ let mut retries = 0;
+ let channel;
+ loop {
+ match builder.connect().await {
+ Ok(c) => {
+ debug!("Connected to gRPC server");
+ channel = c;
+ break;
+ }
+ Err(err) => {
+ if retries < RETRY_COUNT {
+ debug!("Failed to connect, try again momentarily: {:?}", err);
+ retries += 1;
+ std::thread::sleep(RETRY_INTERVAL);
+ } else {
+ panic!("Failed to connect, last err: {:?}", err);
+ }
+ }
+ }
+ }
+
+ // Build an interceptor that will attach a public-untrusted Oak label to
+ // every gRPC request.
+ let mut label = Vec::new();
+ Label::public_untrusted()
+ .encode(&mut label)
+ .expect("Error encoding label");
+ let interceptor = move |mut request: Request<()>| {
+ request.metadata_mut().insert_bin(
+ oak_abi::OAK_LABEL_GRPC_METADATA_KEY,
+ MetadataValue::from_bytes(label.as_ref()),
+ );
+ Ok(request)
+ };
+
+ (channel, interceptor)
+}
