Pass handles as fields in Protobuf messages.

Also includes:
- A vendored Prost with a patch applied to generate Senders and
Receivers from handles.
- Custom derive macro to visit all handles in a struct/enum (crate
oak_derive)
- New example "injection" to showcase passing handles around.

Cargo.toml

@@ -1,6 +1,6 @@
 [workspace]
 members = ["oak/server/rust/oak_loader", "oak/server/rust/oak_runtime"]
-exclude = ["third_party/expect", "third_party/roughenough"]
+exclude = ["third_party/expect", "third_party/roughenough", "third_party/prost/vendored"]
 
 # Patch dependencies on oak crates so that they refer to the versions within this same repository.
 #
@@ -15,3 +15,6 @@ oak_runtime = { path = "oak/server/rust/oak_runtime" }
 oak_utils = { path = "oak_utils" }
 # Third party.
 roughenough = { path = "third_party/roughenough" }
+prost = { path = "third_party/prost/vendored" }
+prost-build = { path = "third_party/prost/vendored/prost-build" }
+prost-types = { path = "third_party/prost/vendored/prost-types" }

examples/Cargo.toml

@@ -17,6 +17,7 @@ members = [
   "running_average/module/rust",
   "translator/common",
   "translator/module/rust",
+  "injection/module/rust",
 ]
 
 # Patch dependencies on oak crates so that they refer to the versions within this same repository.
@@ -31,8 +32,12 @@ oak_abi = { path = "../oak_abi" }
 oak_runtime = { path = "../oak/server/rust/oak_runtime" }
 oak_tests = { path = "../sdk/rust/oak_tests" }
 oak_utils = { path = "../oak_utils" }
+oak_derive = { path = "../sdk/rust/oak_derive" }
 # Examples.
 translator_common = { path = "translator/common" }
 # Third party.
 expect = { path = "../third_party/expect" }
 roughenough = { path = "../third_party/roughenough" }
+prost = { path = "../third_party/prost/vendored" }
+prost-build = { path = "../third_party/prost/vendored/prost-build" }
+prost-types = { path = "../third_party/prost/vendored/prost-types" }

examples/aggregator/backend/Cargo.toml

@@ -16,4 +16,5 @@ tokio = { version = "*", features = ["fs", "macros", "sync", "stream"] }
 tonic = { version = "*", features = ["tls"] }
 
 [build-dependencies]
+oak_utils = "*"
 tonic-build = "*"

examples/aggregator/backend/build.rs

@@ -24,9 +24,12 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     // https://doc.rust-lang.org/cargo/reference/build-scripts.html#cargorerun-if-changedpath
     println!("cargo:rerun-if-changed={}", file_path.display());
 
-    tonic_build::configure()
-        .build_client(false)
-        .build_server(true)
-        .compile(&[file_path.as_path()], &[proto_path])?;
+    oak_utils::tonic_compile(
+        tonic_build::configure()
+            .build_client(false)
+            .build_server(true),
+        &[file_path.as_path()],
+        &[proto_path],
+    );
     Ok(())
 }

examples/authentication/client/Cargo.toml

@@ -21,3 +21,4 @@ url = "*"
 
 [build-dependencies]
 tonic-build = "*"
+oak_utils = "*"

examples/authentication/client/build.rs

@@ -24,9 +24,12 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     // https://doc.rust-lang.org/cargo/reference/build-scripts.html#cargorerun-if-changedpath
     println!("cargo:rerun-if-changed={}", file_path.display());
 
-    tonic_build::configure()
-        .build_client(true)
-        .build_server(false)
-        .compile(&[file_path.as_path()], &[proto_path])?;
+    oak_utils::tonic_compile(
+        tonic_build::configure()
+            .build_client(true)
+            .build_server(false),
+        &[file_path.as_path()],
+        &[proto_path],
+    );
     Ok(())
 }

examples/chat/module/rust/build.rs

@@ -17,6 +17,6 @@
 fn main() {
     oak_utils::compile_protos(
         &["../../proto/chat.proto"],
-        &["../../proto", "../../third_party"],
+        &["../../proto", "../../third_party", "../../../../"],
     );
 }

examples/chat/module/rust/src/backend.rs

@@ -14,10 +14,12 @@
 // limitations under the License.
 //
 
-use crate::{command::Command, proto::Message};
+use crate::proto::{
+    command::Command::{Join, SendMessage},
+    Command, Message,
+};
 use log::info;
 use oak::Node;
-use prost::Message as _;
 
 oak::entrypoint!(backend_oak_main => |in_channel| {
     oak::logger::init_default();
@@ -32,16 +34,13 @@ struct Room {
 
 impl Node<Command> for Room {
     fn handle_command(&mut self, command: Command) -> Result<(), oak::OakError> {
-        match command {
-            Command::Join(h) => {
-                let sender = oak::io::Sender::new(h);
+        match command.command {
+            Some(Join(sender)) => {
                 self.clients
                     .push(oak::grpc::ChannelResponseWriter::new(sender));
                 Ok(())
             }
-            Command::SendMessage(message_bytes) => {
-                let message = Message::decode(message_bytes.as_slice())
-                    .expect("could not parse message from bytes");
+            Some(SendMessage(message)) => {
                 self.messages.push(message.clone());
                 info!("fan out message to {} clients", self.clients.len());
                 for writer in &mut self.clients {
@@ -52,6 +51,7 @@ impl Node<Command> for Room {
                 }
                 Ok(())
             }
+            None => panic!("Empty command"),
         }
     }
 }

examples/chat/module/rust/src/lib.rs

@@ -14,18 +14,16 @@
 // limitations under the License.
 //
 
-use command::Command;
 use log::info;
-use oak::grpc;
-use prost::Message;
+use oak::{grpc, io::Sender};
 use proto::{
-    Chat, ChatDispatcher, CreateRoomRequest, DestroyRoomRequest, SendMessageRequest,
+    command::Command::{Join, SendMessage},
+    Chat, ChatDispatcher, Command, CreateRoomRequest, DestroyRoomRequest, SendMessageRequest,
     SubscribeRequest,
 };
 use std::collections::{hash_map::Entry, HashMap};
 
 mod backend;
-mod command;
 mod proto {
     include!(concat!(env!("OUT_DIR"), "/oak.examples.chat.rs"));
 }
@@ -126,7 +124,9 @@ impl Chat for Node {
             }
             Some(room) => {
                 info!("new subscription to room {:?}", req.room_id);
-                let command = Command::Join(writer.handle());
+                let command = Command {
+                    command: Some(Join(Sender::new(writer.handle()))),
+                };
                 room.sender
                     .send(&command)
                     .expect("could not send command to room Node");
@@ -140,12 +140,9 @@ impl Chat for Node {
             None => room_id_not_found_err(),
             Some(room) => {
                 info!("new message to room {:?}", req.room_id);
-                let mut message_bytes = Vec::new();
-                req.message
-                    .unwrap_or_default()
-                    .encode(&mut message_bytes)
-                    .expect("could not convert message to bytes");
-                let command = Command::SendMessage(message_bytes);
+                let command = Command {
+                    command: req.message.map(SendMessage),
+                };
                 room.sender
                     .send(&command)
                     .expect("could not send command to room Node");

examples/chat/proto/BUILD

@@ -28,12 +28,16 @@ proto_library(
     srcs = ["chat.proto"],
     deps = [
         "@com_google_protobuf//:empty_proto",
+        "//oak/proto:handle_proto",
+        "//oak/proto:grpc_encap_proto",
     ],
 )
 
 cc_proto_library(
     name = "chat_cc_proto",
-    deps = [":chat_proto"],
+    deps = [
+        ":chat_proto",
+    ],
 )
 
 cc_grpc_library(