Pass handles as fields in Protobuf messages. (#1108)

Also includes:
- Custom derive macro to visit all handles in a struct/enum (crate
oak_derive)
- New example "injection" to showcase passing handles around.

examples/Cargo.toml

@@ -12,6 +12,7 @@ members = [
   "chat/module/rust",
   "hello_world/grpc",
   "hello_world/module/rust",
+  "injection/module/rust",
   "machine_learning/module/rust",
   "trusted_information_retrieval/backend",
   "trusted_information_retrieval/client/rust",
@@ -35,6 +36,7 @@ members = [
 # Oak.
 oak = { path = "../sdk/rust/oak" }
 oak_abi = { path = "../oak_abi" }
+oak_derive = { path = "../sdk/rust/oak_derive" }
 oak_runtime = { path = "../oak/server/rust/oak_runtime" }
 oak_tests = { path = "../sdk/rust/oak_tests" }
 oak_utils = { path = "../oak_utils" }

examples/chat/module/rust/build.rs

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

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

@@ -14,10 +14,12 @@
 // limitations under the License.
 //
 
-use crate::{command::Command, proto::Message};
-use log::info;
+use crate::proto::{
+    command::Command::{JoinRoom, SendMessage},
+    Command, Message,
+};
+use log::{info, warn};
 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(JoinRoom(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,10 @@ impl Node<Command> for Room {
                 }
                 Ok(())
             }
+            None => {
+                warn!("Received empty command");
+                Err(oak::OakError::OakStatus(oak::OakStatus::ErrInvalidArgs))
+            }
         }
     }
 }

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::{JoinRoom, 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(JoinRoom(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

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

examples/chat/proto/chat.proto

@@ -19,6 +19,8 @@ syntax = "proto3";
 package oak.examples.chat;
 
 import "google/protobuf/empty.proto";
+import "oak/proto/handle.proto";
+import "oak/proto/grpc_encap.proto";
 
 message CreateRoomRequest {
   // ID used to identify the room; knowledge of this value allows entry to the room. The client
@@ -66,3 +68,16 @@ service Chat {
   // Send a message to a chat room.
   rpc SendMessage(SendMessageRequest) returns (google.protobuf.Empty);
 }
+
+// Command sent to room nodes.
+//
+// This message is only used for inter-node communication, it is not exposed through gRPC.
+message Command {
+  oneof command {
+    // Sent when a new subscriber joins the room.
+    oak.handle.Sender join_room = 1 [(oak.handle.message_type) = ".oak.encap.GrpcResponse"];
+
+    // Command to send a message to the room (and thus to all subscribers).
+    Message send_message = 2;
+  }
+}

examples/injection/client/BUILD

@@ -0,0 +1,34 @@
+#
+# Copyright 2019 The Project Oak Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+load("@rules_cc//cc:defs.bzl", "cc_binary")
+
+package(
+    licenses = ["notice"],
+)
+
+cc_binary(
+    name = "client",
+    srcs = ["injection.cc"],
+    deps = [
+        "//examples/injection/proto:injection_cc_grpc",
+        "//oak/client:application_client",
+        "@com_github_google_glog//:glog",
+        "@com_github_grpc_grpc//:grpc++",
+        "@com_google_absl//absl/flags:flag",
+        "@com_google_absl//absl/flags:parse",
+    ],
+)

examples/injection/client/injection.cc

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2019 The Project Oak Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "absl/flags/flag.h"
+#include "absl/flags/parse.h"
+#include "examples/injection/proto/injection.grpc.pb.h"
+#include "examples/injection/proto/injection.pb.h"
+#include "glog/logging.h"
+#include "include/grpcpp/grpcpp.h"
+#include "oak/client/application_client.h"
+#include "oak/common/label.h"
+
+ABSL_FLAG(std::string, address, "localhost:8080", "Address of the Oak application to connect to");
+ABSL_FLAG(std::string, ca_cert, "", "Path to the PEM-encoded CA root certificate");
+
+using ::oak::examples::injection::BlobResponse;
+using ::oak::examples::injection::BlobStore;
+using ::oak::examples::injection::GetBlobRequest;
+using ::oak::examples::injection::PutBlobRequest;
+
+int main(int argc, char** argv) {
+  absl::ParseCommandLine(argc, argv);
+
+  std::string address = absl::GetFlag(FLAGS_address);
+  std::string ca_cert = oak::ApplicationClient::LoadRootCert(absl::GetFlag(FLAGS_ca_cert));
+  LOG(INFO) << "Connecting to Oak Application: " << address;
+
+  // TODO(#1066): Use a more restrictive Label.
+  oak::label::Label label = oak::PublicUntrustedLabel();
+  // Connect to the Oak Application.
+  auto stub = BlobStore::NewStub(oak::ApplicationClient::CreateChannel(
+      address, oak::ApplicationClient::GetTlsChannelCredentials(ca_cert), label));
+  if (stub == nullptr) {
+    LOG(FATAL) << "Failed to create application stub";
+  }
+
+  PutBlobRequest putRequest;
+  putRequest.set_blob("Hello, blob store!");
+  grpc::ClientContext putContext;
+  BlobResponse putResponse;
+  grpc::Status putStatus = stub->PutBlob(&putContext, putRequest, &putResponse);
+  if (!putStatus.ok()) {
+    LOG(FATAL) << "PutBlob failed: " << putStatus.error_code() << ": " << putStatus.error_message();
+  }
+  LOG(INFO) << "Blob stored at id: " << putResponse.id();
+
+  GetBlobRequest getRequest;
+  getRequest.set_id(putResponse.id());
+  grpc::ClientContext getContext;
+  BlobResponse getResponse;
+  grpc::Status getStatus = stub->GetBlob(&getContext, getRequest, &getResponse);
+  if (!getStatus.ok()) {
+    LOG(FATAL) << "GetBlob failed: " << getStatus.error_code() << ": " << getStatus.error_message();
+  }
+  LOG(INFO) << "Successfully retrieved Blob";
+
+  if (putRequest.blob() != getResponse.blob()) {
+    LOG(FATAL) << "Blobs were different. Original: '" << putRequest.blob() << "', retrieved: '"
+               << getResponse.blob() << "'";
+  }
+  LOG(INFO) << "Blobs match!";
+
+  return EXIT_SUCCESS;
+}