Migrate oak_functions_service to oak_proto_rust

This in peparation for bazelification. Bug: 349587510 Change-Id: I9213827cf92af5fbf8f440d7f1766ecc1b98498c
project-oak · Jul 3, 2024 · 9acdd5e · 9acdd5e
1 parent e79227e
commit 9acdd5e
Show file tree

Hide file tree

Showing 19 changed files with 518 additions and 107 deletions.
diff --git a/enclave_apps/Cargo.lock b/enclave_apps/Cargo.lock
diff --git a/enclave_apps/Cargo.toml b/enclave_apps/Cargo.toml
@@ -14,4 +14,5 @@ oak_channel = { path = "../oak_channel" }
 oak_dice = { path = "../oak_dice" }
 oak_restricted_kernel_sdk = { path = "../oak_restricted_kernel_sdk" }
 oak_restricted_kernel_interface = { path = "../oak_restricted_kernel_interface" }
+oak_proto_rust = { path = "../oak_proto_rust" }
 zerocopy = "*"
diff --git a/enclave_apps/oak_functions_enclave_app/Cargo.toml b/enclave_apps/oak_functions_enclave_app/Cargo.toml
@@ -20,6 +20,7 @@ oak_functions_enclave_service = { path = "../../oak_functions_enclave_service",
 oak_functions_service = { path = "../../oak_functions_service", default-features = false }
 log = "*"
 micro_rpc = { workspace = true }
+oak_proto_rust = { workspace = true }
 oak_restricted_kernel_sdk = { workspace = true }
 
 [[bin]]

diff --git a/enclave_apps/oak_functions_enclave_app/src/main.rs b/enclave_apps/oak_functions_enclave_app/src/main.rs
@@ -50,8 +50,7 @@ fn main() -> ! {
         None,
         WasmConfig::default(),
     );
-    let server =
-        oak_functions_enclave_service::proto::oak::functions::OakFunctionsServer::new(service);
+    let server = oak_proto_rust::oak::oak_functions::OakFunctionsServer::new(service);
     start_blocking_server(Box::<FileDescriptorChannel>::default(), server, &mut invocation_stats)
         .expect("server encountered an unrecoverable error");
 }
diff --git a/oak_functions_containers_app/build.rs b/oak_functions_containers_app/build.rs
@@ -31,7 +31,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             build_client: true,
             extern_paths: vec![ExternPath::new(
                 ".oak.functions",
-                "::oak_functions_service::proto::oak::functions",
+                "::oak_proto_rust::oak::oak_functions",
             )],
         },
     )?;

diff --git a/oak_functions_containers_app/src/lib.rs b/oak_functions_containers_app/src/lib.rs
@@ -27,15 +27,12 @@ use std::{
 use anyhow::Context;
 use oak_attestation::handler::AsyncEncryptionHandler;
 use oak_crypto::encryption_key::AsyncEncryptionKeyHandle;
-use oak_functions_service::{
-    instance::OakFunctionsInstance,
-    proto::oak::functions::{
-        AbortNextLookupDataResponse, Empty, ExtendNextLookupDataRequest,
-        ExtendNextLookupDataResponse, FinishNextLookupDataRequest, FinishNextLookupDataResponse,
-        InitializeRequest, InitializeResponse, InvokeRequest, InvokeResponse, LookupDataChunk,
-        ReserveRequest, ReserveResponse,
-    },
-    Handler, Observer,
+use oak_functions_service::{instance::OakFunctionsInstance, Handler, Observer};
+use oak_proto_rust::oak::oak_functions::{
+    AbortNextLookupDataResponse, Empty, ExtendNextLookupDataRequest, ExtendNextLookupDataResponse,
+    FinishNextLookupDataRequest, FinishNextLookupDataResponse, InitializeRequest,
+    InitializeResponse, InvokeRequest, InvokeResponse, LookupDataChunk, ReserveRequest,
+    ReserveResponse,
 };
 use opentelemetry::{
     metrics::{Histogram, Meter, Unit},

diff --git a/oak_functions_containers_app/src/main.rs b/oak_functions_containers_app/src/main.rs
@@ -30,12 +30,10 @@ use oak_crypto::encryption_key::AsyncEncryptionKeyHandle;
 #[cfg(feature = "native")]
 use oak_functions_containers_app::native_handler::NativeHandler;
 use oak_functions_containers_app::serve as app_serve;
-use oak_functions_service::{
-    proto::oak::functions::config::{
-        application_config::CommunicationChannel, ApplicationConfig, HandlerType,
-        TcpCommunicationChannel, WasmtimeConfig,
-    },
-    wasm::wasmtime::WasmtimeHandler,
+use oak_functions_service::wasm::wasmtime::WasmtimeHandler;
+use oak_proto_rust::oak::oak_functions::config::{
+    application_config::CommunicationChannel, ApplicationConfig, HandlerType,
+    TcpCommunicationChannel, WasmtimeConfig,
 };
 use opentelemetry::{
     global::set_error_handler,

diff --git a/oak_functions_containers_app/tests/integration_test.rs b/oak_functions_containers_app/tests/integration_test.rs
@@ -30,9 +30,8 @@ use std::{
 
 use oak_crypto::encryption_key::generate_encryption_key_pair;
 use oak_functions_containers_app::serve;
-use oak_functions_service::{
-    proto::oak::functions::InitializeRequest, wasm::wasmtime::WasmtimeHandler,
-};
+use oak_functions_service::wasm::wasmtime::WasmtimeHandler;
+use oak_proto_rust::oak::oak_functions::InitializeRequest;
 use opentelemetry::metrics::{noop::NoopMeterProvider, MeterProvider};
 use tokio::net::TcpListener;
 use tokio_stream::wrappers::TcpListenerStream;

diff --git a/oak_functions_enclave_service/src/lib.rs b/oak_functions_enclave_service/src/lib.rs
@@ -23,16 +23,12 @@ use alloc::{format, string::ToString, sync::Arc, vec::Vec};
 use oak_attestation::{dice::evidence_to_proto, handler::EncryptionHandler};
 use oak_core::sync::OnceCell;
 use oak_crypto::encryption_key::EncryptionKeyHandle;
-pub use oak_functions_service::proto;
-use oak_functions_service::{
-    instance::OakFunctionsInstance,
-    proto::oak::functions::{
-        extend_next_lookup_data_request::Data, AbortNextLookupDataResponse, Empty,
-        ExtendNextLookupDataRequest, ExtendNextLookupDataResponse, FinishNextLookupDataRequest,
-        FinishNextLookupDataResponse, InitializeRequest, InitializeResponse, InvokeRequest,
-        InvokeResponse, LookupDataChunk, OakFunctions, ReserveRequest, ReserveResponse,
-    },
-    Handler, Observer,
+use oak_functions_service::{instance::OakFunctionsInstance, Handler, Observer};
+use oak_proto_rust::oak::oak_functions::{
+    extend_next_lookup_data_request::Data, AbortNextLookupDataResponse, Empty,
+    ExtendNextLookupDataRequest, ExtendNextLookupDataResponse, FinishNextLookupDataRequest,
+    FinishNextLookupDataResponse, InitializeRequest, InitializeResponse, InvokeRequest,
+    InvokeResponse, LookupDataChunk, OakFunctions, ReserveRequest, ReserveResponse,
 };
 use prost::Message;
 

diff --git a/oak_functions_enclave_service/tests/integration_test.rs b/oak_functions_enclave_service/tests/integration_test.rs
@@ -25,17 +25,15 @@ use core::assert_matches::assert_matches;
 
 use oak_client::verifier::extract_encryption_public_key;
 use oak_crypto::encryptor::ClientEncryptor;
-use oak_functions_enclave_service::{
-    proto::oak::functions::{
+use oak_functions_enclave_service::OakFunctionsService;
+use oak_proto_rust::oak::{
+    crypto::v1::EncryptedRequest,
+    oak_functions::{
+        extend_next_lookup_data_request::Data,
+        testing::{EchoAndPanicRequest, TestModuleClient},
         ExtendNextLookupDataRequest, FinishNextLookupDataRequest, InitializeRequest, InvokeRequest,
         LookupDataChunk, LookupDataEntry, OakFunctionsClient, OakFunctionsServer,
     },
-    OakFunctionsService,
-};
-use oak_functions_service::proto::oak::functions::extend_next_lookup_data_request::Data;
-use oak_proto_rust::oak::{
-    crypto::v1::EncryptedRequest,
-    oak_functions::testing::{EchoAndPanicRequest, TestModuleClient},
 };
 use prost::Message;
 

diff --git a/oak_functions_service/build.rs b/oak_functions_service/build.rs
diff --git a/oak_functions_service/src/instance.rs b/oak_functions_service/src/instance.rs
@@ -18,19 +18,15 @@ use alloc::{format, sync::Arc};
 
 use micro_rpc::{Status, Vec};
 use oak_functions_abi::Request;
+use oak_proto_rust::oak::oak_functions::{
+    extend_next_lookup_data_request::Data, AbortNextLookupDataResponse, Empty,
+    ExtendNextLookupDataRequest, ExtendNextLookupDataResponse, FinishNextLookupDataRequest,
+    FinishNextLookupDataResponse, InitializeRequest, LookupDataChunk, LookupDataEntry,
+    ReserveRequest, ReserveResponse,
+};
 use prost::Message;
 
-use crate::{
-    logger::StandaloneLogger,
-    lookup::LookupDataManager,
-    proto::oak::functions::{
-        extend_next_lookup_data_request::Data, AbortNextLookupDataResponse, Empty,
-        ExtendNextLookupDataRequest, ExtendNextLookupDataResponse, FinishNextLookupDataRequest,
-        FinishNextLookupDataResponse, InitializeRequest, LookupDataChunk, LookupDataEntry,
-        ReserveRequest, ReserveResponse,
-    },
-    Handler, Observer,
-};
+use crate::{logger::StandaloneLogger, lookup::LookupDataManager, Handler, Observer};
 
 pub struct OakFunctionsInstance<H: Handler> {
     lookup_data_manager: Arc<LookupDataManager<16>>,

diff --git a/oak_functions_service/src/lib.rs b/oak_functions_service/src/lib.rs
@@ -32,19 +32,6 @@ extern crate rand_core;
 #[cfg(test)]
 extern crate std;
 
-pub mod proto {
-    pub mod oak {
-        pub mod functions {
-            pub mod config {
-                include!(concat!(env!("OUT_DIR"), "/oak.functions.config.rs"));
-            }
-            use prost::Message;
-            include!(concat!(env!("OUT_DIR"), "/oak.functions.rs"));
-        }
-        pub use oak_proto_rust::oak::{attestation, crypto};
-    }
-}
-
 pub mod instance;
 pub mod logger;
 pub mod lookup;

diff --git a/oak_functions_service/src/wasm/wasmtime.rs b/oak_functions_service/src/wasm/wasmtime.rs
@@ -33,12 +33,12 @@ use std::time::Instant;
 use log::Level;
 use micro_rpc::StatusCode;
 use oak_functions_abi::{Request, Response};
+use oak_proto_rust::oak::oak_functions::config::WasmtimeConfig;
 use wasmtime::{PoolingAllocationConfig, Store};
 
 use crate::{
     logger::{OakLogger, StandaloneLogger},
     lookup::LookupDataManager,
-    proto::oak::functions::config::WasmtimeConfig,
     wasm::{api::StdWasmApiFactory, WasmApiFactory},
     Handler, Observer,
 };

diff --git a/oak_proto_rust/BUILD b/oak_proto_rust/BUILD
@@ -70,6 +70,7 @@ cargo_build_script(
         "//proto/containers:interfaces_proto",
         "//proto/crypto:crypto_proto",
         "//proto/oak_functions:abi_proto",
+        "//proto/oak_functions:application_config_proto",
         "//proto/oak_functions:lookup_data_proto",
         "//proto/oak_functions:testing_proto",
         "//proto/oak_functions/sdk:oak_functions_wasm_proto",

diff --git a/oak_proto_rust/build.rs b/oak_proto_rust/build.rs
@@ -33,6 +33,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         "../proto/containers/interfaces.proto",
         "../proto/digest.proto",
         "../proto/oak_functions/abi.proto",
+        "../proto/oak_functions/application_config.proto",
         "../proto/oak_functions/lookup_data.proto",
         "../proto/session/session.proto",
         "../proto/oak_functions/sdk/oak_functions_wasm.proto",
@@ -78,6 +79,26 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         Default::default(),
     );
 
+    micro_rpc_build::compile(
+        &["../proto/oak_functions/service/oak_functions.proto"],
+        &included_protos,
+        micro_rpc_build::CompileOptions {
+            receiver_type: micro_rpc_build::ReceiverType::RefSelf,
+            bytes: vec![
+                ".oak.functions.LookupDataEntry".to_string(),
+                ".oak.functions.ExtendNextLookupDataRequest".to_string(),
+            ],
+            extern_paths: vec![
+                micro_rpc_build::ExternPath::new(".oak.crypto.v1", "crate::oak::crypto::v1"),
+                micro_rpc_build::ExternPath::new(
+                    ".oak.attestation.v1",
+                    "crate::oak::attestation::v1",
+                ),
+            ],
+            ..Default::default()
+        },
+    );
+
     // Tell cargo to rerun this build script if the proto file has changed.
     // https://doc.rust-lang.org/cargo/reference/build-scripts.html#cargorerun-if-changedpath
     for proto_path in proto_paths.iter() {