diff --git a/docs/programming-oak.md b/docs/programming-oak.md index 9a4f4a2957f..a4e29d84da6 100644 --- a/docs/programming-oak.md +++ b/docs/programming-oak.md @@ -403,17 +403,20 @@ from the implicit incoming channel, usually by implementing the trait setting the [`Command`](https://project-oak.github.io/oak/sdk/doc/oak/trait.CommandHandler.html#associatedtype.Command) associated type to -[`ConfigMap`](https://project-oak.github.io/oak/sdk/doc/oak/proto/oak/application/struct.ConfigMap.html), -or manually reading from the initial `Receiver`: +[`ConfigMap`](https://project-oak.github.io/oak/sdk/doc/oak/proto/oak/application/struct.ConfigMap.html): -[embedmd]:# (../examples/trusted_database/module/rust/src/lib.rs Rust /oak::entrypoint/ /.*let config_map =.*/) +[embedmd]:# (../examples/aggregator/module/rust/src/lib.rs Rust /impl oak::CommandHandler/ /.*context.*/) ```Rust -oak::entrypoint!(oak_main => |receiver: Receiver| { - let log_sender = oak::logger::create().unwrap(); - oak::logger::init(log_sender, log::Level::Debug).unwrap(); +impl oak::CommandHandler for Main { + type Command = ConfigMap; - let config_map = receiver.receive().expect("Couldn't read config map"); + fn handle_command(&mut self, command: ConfigMap) -> anyhow::Result<()> { + let log_sender = oak::logger::create()?; + oak::logger::init(log_sender.clone(), log::Level::Debug)?; + let config: Config = + toml::from_slice(&command.items.get("config").expect("Couldn't find config")) + .context("Couldn't parse TOML config file")?; ``` diff --git a/examples/trusted_database/module/rust/Cargo.toml b/examples/trusted_database/module/rust/Cargo.toml index 9a08d2252ee..9cf233550ac 100644 --- a/examples/trusted_database/module/rust/Cargo.toml +++ b/examples/trusted_database/module/rust/Cargo.toml @@ -14,6 +14,7 @@ log = "*" oak = "=0.1.0" oak_abi = "=0.1.0" oak_io = "=0.1.0" +oak_services = "=0.1.0" prost = "*" serde = "*" quick-xml = { version = "*", features = ["serialize"] } diff --git a/examples/trusted_database/module/rust/build.rs b/examples/trusted_database/module/rust/build.rs index b5ed8a7ef4c..14466f83d2e 100644 --- a/examples/trusted_database/module/rust/build.rs +++ b/examples/trusted_database/module/rust/build.rs @@ -18,7 +18,7 @@ fn main() { oak_utils::compile_protos( &[ "../../proto/trusted_database.proto", - "../../proto/trusted_database_command.proto", + "../../proto/trusted_database_init.proto", ], &["../../proto", "../../../../"], ); diff --git a/examples/trusted_database/module/rust/src/handler.rs b/examples/trusted_database/module/rust/src/handler.rs index cc76f962913..ed396d1a5af 100644 --- a/examples/trusted_database/module/rust/src/handler.rs +++ b/examples/trusted_database/module/rust/src/handler.rs @@ -14,39 +14,37 @@ // limitations under the License. // -//! Trusted Database Handler Node. -//! -//! In the current implementation clients send their location coordinates (latitude and longitude) -//! and the Handler Node returns the location of the closest Point Of Interest. -//! -//! The Handler Node searches for the closest Point Of Interest in the database received from the -//! Main Node. - use crate::proto::oak::examples::trusted_database::{ GetPointOfInterestRequest, GetPointOfInterestResponse, ListPointsOfInterestRequest, ListPointsOfInterestResponse, Location, PointOfInterestMap, TrustedDatabase, - TrustedDatabaseCommand, TrustedDatabaseDispatcher, + TrustedDatabaseDispatcher, TrustedDatabaseInit, }; use log::{debug, error, warn}; -use oak::{ - grpc, - io::{Receiver, ReceiverExt}, -}; +use oak::grpc; // Error messages. const NO_LOCATION_ERROR: &str = "Location is not specified"; const ID_NOT_FOUND_ERROR: &str = "ID not found"; const EMPTY_DATABASE_ERROR: &str = "Empty database"; -/// Oak Node that contains a copy of the database. -pub struct TrustedDatabaseHandlerNode { +/// Oak Handler Node that contains a copy of the database and handles client requests. +#[derive(Default)] +pub struct Handler { points_of_interest: PointOfInterestMap, } -oak::impl_dispatcher!(impl TrustedDatabaseHandlerNode : TrustedDatabaseDispatcher); +impl oak::WithInit for Handler { + type Init = TrustedDatabaseInit; + + fn create(init: Self::Init) -> Self { + oak::logger::init(init.log_sender.unwrap(), log::Level::Debug).unwrap(); + let points_of_interest = init.points_of_interest.expect("Couldn't receive database"); + Self { points_of_interest } + } +} /// A gRPC service implementation for the Private Information Retrieval example. -impl TrustedDatabase for TrustedDatabaseHandlerNode { +impl TrustedDatabase for Handler { // Find Point Of Interest based on id. fn get_point_of_interest( &mut self, @@ -73,7 +71,7 @@ impl TrustedDatabase for TrustedDatabaseHandlerNode { &mut self, request: ListPointsOfInterestRequest, ) -> grpc::Result { - debug!("Received request: {:?}", request); + error!("Received request: {:?}", request); let request_location = request.location.ok_or_else(|| { let err = grpc::build_status(grpc::Code::InvalidArgument, &NO_LOCATION_ERROR); warn!("{:?}", err); @@ -138,18 +136,6 @@ pub fn distance(first: Location, second: Location) -> f32 { EARTH_RADIUS * central_angle } -oak::entrypoint!(handler_oak_main => |command_receiver: Receiver| { - let log_sender = oak::logger::create().unwrap(); - oak::logger::init(log_sender, log::Level::Debug).unwrap(); - - // Receive command. - let command: TrustedDatabaseCommand = - command_receiver.receive().expect("Couldn't receive command"); - let receiver = command.invocation_receiver.expect("Couldn't receive gRPC invocation receiver"); +oak::entrypoint_command_handler_init!(handler => Handler); - // Run event loop and handle incoming invocations. - let node = TrustedDatabaseHandlerNode { points_of_interest: command.points_of_interest.expect("No database entries") }; - let invocation_receiver = receiver.receiver.expect("Empty gRPC invocation receiver"); - // The event loop only runs once because the `Main` Node sends a single invocation. - oak::run_command_loop(node, invocation_receiver.iter()); -}); +oak::impl_dispatcher!(impl Handler : TrustedDatabaseDispatcher); diff --git a/examples/trusted_database/module/rust/src/lib.rs b/examples/trusted_database/module/rust/src/lib.rs index 339edc04d79..146516e51e1 100644 --- a/examples/trusted_database/module/rust/src/lib.rs +++ b/examples/trusted_database/module/rust/src/lib.rs @@ -35,6 +35,7 @@ pub mod proto { pub mod oak { pub use oak::proto::oak::invocation; + pub use oak_services::proto::oak::log; pub mod examples { pub mod trusted_database { include!(concat!( @@ -43,7 +44,7 @@ pub mod proto { )); include!(concat!( env!("OUT_DIR"), - "/oak.examples.trusted_database_command.rs" + "/oak.examples.trusted_database_init.rs" )); } } @@ -52,82 +53,45 @@ pub mod proto { mod database; mod handler; +mod router; #[cfg(test)] mod tests; +use crate::{proto::oak::examples::trusted_database::TrustedDatabaseInit, router::Router}; use anyhow::Context; use database::load_database; -use log::debug; -use oak::{ - grpc, - io::{Receiver, ReceiverExt, SenderExt}, - proto::oak::invocation::GrpcInvocationReceiver, - CommandHandler, -}; -use oak_abi::{label::Label, proto::oak::application::ConfigMap}; -use proto::oak::examples::trusted_database::{PointOfInterestMap, TrustedDatabaseCommand}; +use oak::proto::oak::application::ConfigMap; +use oak_abi::label::Label; -/// Oak Node that contains an in-memory database. -pub struct TrustedDatabaseNode { - points_of_interest: PointOfInterestMap, -} +/// Main entrypoint of the Trusted Database application. +/// +/// This node is in charge of creating the other top-level nodes, but does not process any request. +#[derive(Default)] +struct Main; -impl CommandHandler for TrustedDatabaseNode { - type Command = grpc::Invocation; +impl oak::CommandHandler for Main { + type Command = ConfigMap; - fn handle_command(&mut self, invocation: grpc::Invocation) -> anyhow::Result<()> { - // Create a client request handler Node. - debug!("Creating handler Node"); - // TODO(#1406): Use client assigned label for creating a new handler Node. - let sender = oak::io::node_create( - "handler", - &Label::public_untrusted(), - &oak::node_config::wasm("app", "handler_oak_main"), - ) - .context("Couldn't create handler Node")?; + fn handle_command(&mut self, config_map: Self::Command) -> anyhow::Result<()> { + let log_sender = oak::logger::create()?; + oak::logger::init(log_sender.clone(), log::Level::Debug)?; + let points_of_interest = load_database(config_map).expect("Couldn't load database"); - // Create a gRPC invocation channel for forwarding requests to the - // `TrustedDatabaseHandlerNode`. - let (invocation_sender, invocation_receiver) = oak::io::channel_create::( - "gRPC invocation", + let init = TrustedDatabaseInit { + log_sender: Some(log_sender), + points_of_interest: Some(points_of_interest), + }; + let router_sender = oak::io::entrypoint_node_create::( + "router", &Label::public_untrusted(), + "app", + init, ) - .context("Couldn't create gRPC invocation channel")?; - - // Create a command message that contains a copy of the database. - let command = TrustedDatabaseCommand { - invocation_receiver: Some(GrpcInvocationReceiver { - receiver: Some(invocation_receiver), - }), - points_of_interest: Some(self.points_of_interest.clone()), - }; - - // Send the command massage to create a `TrustedDatabaseHandlerNode` - debug!("Sending command message to handler Node"); - sender - .send(&command) - .context("Couldn't send command to handler Node")?; - oak::channel_close(sender.handle.handle).context("Couldn't close sender channel")?; - - // Send the original gRPC invocation to the `TrustedDatabaseHandlerNode` - debug!("Sending gRPC invocation to handler Node"); - invocation_sender - .send(&invocation) - .context("Couldn't send gRPC invocation to handler Node")?; - oak::channel_close(invocation_sender.handle.handle) - .context("Couldn't close sender channel")?; - + .context("Couldn't create router node")?; + oak::grpc::server::init_with_sender("[::]:8080", router_sender) + .context("Couldn't create gRPC server pseudo-Node")?; Ok(()) } } -oak::entrypoint!(oak_main => |receiver: Receiver| { - let log_sender = oak::logger::create().unwrap(); - oak::logger::init(log_sender, log::Level::Debug).unwrap(); - - let config_map = receiver.receive().expect("Couldn't read config map"); - let points_of_interest = load_database(config_map).expect("Couldn't load database"); - let grpc_channel = - oak::grpc::server::init("[::]:8080").expect("Couldn't create gRPC server pseudo-Node"); - oak::run_command_loop(TrustedDatabaseNode { points_of_interest }, grpc_channel.iter()); -}); +oak::entrypoint_command_handler!(oak_main => Main); diff --git a/examples/trusted_database/module/rust/src/router.rs b/examples/trusted_database/module/rust/src/router.rs new file mode 100644 index 00000000000..e6174ab9b6c --- /dev/null +++ b/examples/trusted_database/module/rust/src/router.rs @@ -0,0 +1,65 @@ +// +// Copyright 2020 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. +// + +use crate::{handler::Handler, proto::oak::examples::trusted_database::TrustedDatabaseInit}; +use anyhow::Context; +use oak::{ + grpc, + io::{ReceiverExt, SenderExt}, + CommandHandler, +}; + +/// Oak Router Node that contains an in-memory database (saved in the `init`). +#[derive(Default)] +pub struct Router { + /// Init message to be sent to every newly created Handler Node. + init: TrustedDatabaseInit, +} + +impl oak::WithInit for Router { + type Init = TrustedDatabaseInit; + + fn create(init: Self::Init) -> Self { + Self { init } + } +} + +/// Processes client requests and creates individual Handler Nodes. +/// Each newly created Handler Node receives a copy of the database stored in [`Router::init`]. +impl CommandHandler for Router { + type Command = grpc::Invocation; + + fn handle_command(&mut self, invocation: Self::Command) -> anyhow::Result<()> { + let label = invocation + .receiver + .as_ref() + .context("Couldn't get receiver")? + .label() + .context("Couldn't get label")?; + let handler_invocation_sender = oak::io::entrypoint_node_create::( + "handler", + &label, + "app", + self.init.clone(), + ) + .context("Couldn't create handler node")?; + handler_invocation_sender + .send(&invocation) + .context("Couldn't send invocation to handler node") + } +} + +oak::entrypoint_command_handler_init!(router => Router); diff --git a/examples/trusted_database/proto/trusted_database_command.proto b/examples/trusted_database/proto/trusted_database_init.proto similarity index 69% rename from examples/trusted_database/proto/trusted_database_command.proto rename to examples/trusted_database/proto/trusted_database_init.proto index ca24386c1ff..43df8ff3f3c 100644 --- a/examples/trusted_database/proto/trusted_database_command.proto +++ b/examples/trusted_database/proto/trusted_database_init.proto @@ -16,12 +16,16 @@ syntax = "proto3"; -package oak.examples.trusted_database_command; +package oak.examples.trusted_database_init; import "trusted_database.proto"; import "oak_services/proto/grpc_invocation.proto"; +import "oak_services/proto/log.proto"; +import "proto/handle.proto"; -message TrustedDatabaseCommand { - oak.invocation.GrpcInvocationReceiver invocation_receiver = 1; +// Initialization message that should be sent to Router Oak Node and Handler Oak Node. +message TrustedDatabaseInit { + oak.handle.Sender log_sender = 1 [(oak.handle.message_type) = ".oak.log.LogMessage"]; + // Copy of the database. oak.examples.trusted_database.PointOfInterestMap points_of_interest = 2; }