Create a gRPC interface for the untrusted loader (#2705)

This PR does three things:

  * adds a really trivial echo grpc interface
  * implements the gRPC server in the untrusted loader (ripping out the hacky stdio code)
  * and finally, adds a really trivial gRPC client to talk to the loader.

Thus, the poor little string will first travel from the client to loader via gRPC, where it's unpacked and stuffed into CBOR and sent via serial to the UEFI app, which will then echo the string back via the same channels.

But hey, we've got something working end-to-end!

Author: andrisaar

Cargo.lock

@@ -12,6 +12,7 @@ members = [
   "experimental/trusted_shuffler/common",
   "experimental/trusted_shuffler/server",
   "experimental/trusted_shuffler/trusted_shuffler",
+  "experimental/uefi/client",
   "experimental/uefi/loader",
   "oak_functions/abi",
   "oak_functions/client/rust",

Cargo.toml

@@ -39,7 +39,7 @@ mod serial;
 // there's exactly one function with the `#[entry]` attribute in the
 // dependency graph.
 #[entry]
-fn _start(_handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
+fn _start(handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
     uefi_services::init(&mut system_table).unwrap();
 
     // As we're not relying on the normal Rust test harness, we need to call
@@ -49,7 +49,7 @@ fn _start(_handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
         test_main();
         Status::SUCCESS
     } else {
-        main(_handle, &mut system_table)
+        main(handle, &mut system_table)
     };
 
     // After we're done running our code, we also tell the UEFI runtime to shut
@@ -59,20 +59,32 @@ fn _start(_handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
         .reset(ResetType::Shutdown, status, None);
 }
 
+// Run the echo on the first serial port in the system (the UEFI console will
+// use the first serial port in the system)
+const ECHO_SERIAL_PORT_INDEX: usize = 1;
+
 fn main(handle: Handle, system_table: &mut SystemTable<Boot>) -> Status {
     use core::fmt::Write;
 
     writeln!(system_table.stdout(), "Hello World!").unwrap();
 
-    serial_echo(handle, system_table.boot_services()).unwrap();
+    serial_echo(handle, system_table.boot_services(), ECHO_SERIAL_PORT_INDEX).unwrap();
 }
 
-// Run the echo on the first serial port in the system (the UEFI console will
-// use the first serial port in the system)
-const ECHO_SERIAL_PORT_INDEX: usize = 1;
-
-fn serial_echo(handle: Handle, bt: &BootServices) -> Result<!, uefi::Error<()>> {
-    let mut serial = serial::Serial::get(handle, bt, ECHO_SERIAL_PORT_INDEX)?;
+// Opens the index-th serial port on the system and echoes back all frames sent over the serial
+// port.
+//
+// Panics if the index-th serial port does not exist.
+//
+// Arguments:
+//    * `handle` - UEFI handle of the agent (eg of the UEFI app)
+//    * `boot_services` - reference to the UEFI Boot Services struct (obtained from the system
+//      table)
+//    * `index` - index of the serial port, zero-based.
+//
+//  Normally does not return, unless an error is raised.
+fn serial_echo(handle: Handle, bt: &BootServices, index: usize) -> Result<!, uefi::Error<()>> {
+    let mut serial = serial::Serial::get(handle, bt, index)?;
     loop {
         let msg: alloc::string::String = de::from_reader(&mut serial).map_err(|err| match err {
             de::Error::Io(err) => err,

experimental/uefi/app/src/main.rs

@@ -0,0 +1,15 @@
+[package]
+name = "uefi-client"
+version = "0.1.0"
+authors = ["Andri Saar <[email protected]>"]
+edition = "2021"
+license = "Apache-2.0"
+
+[dependencies]
+clap = { version = "*", features = ["derive"] }
+prost = "*"
+tokio = { version = "*", features = ["macros", "rt-multi-thread"] }
+tonic = "*"
+
+[build-dependencies]
+oak_utils = { path = "../../../oak_utils" }

experimental/uefi/client/Cargo.toml

@@ -0,0 +1,30 @@
+//
+// Copyright 2022 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 oak_utils::{generate_grpc_code, CodegenOptions};
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    generate_grpc_code(
+        "../../../",
+        &["experimental/uefi/proto/echo.proto"],
+        CodegenOptions {
+            build_client: true,
+            build_server: false,
+            extern_paths: vec![],
+        },
+    )?;
+    Ok(())
+}

experimental/uefi/client/build.rs

@@ -0,0 +1,51 @@
+//
+// Copyright 2022 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 std::io::{stdin, BufRead};
+
+use clap::Parser;
+use tonic::Request;
+
+pub mod proto {
+    tonic::include_proto!("oak.experimental.uefi");
+}
+
+use proto::{echo_client::EchoClient, EchoRequest};
+
+#[derive(Parser, Debug)]
+struct Args {
+    /// address of the server
+    #[clap(long, default_value = "http://127.0.0.1:8000")]
+    server: String,
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let cli = Args::parse();
+
+    let mut client = EchoClient::connect(cli.server).await?;
+
+    for line in stdin().lock().lines() {
+        let request = Request::new(EchoRequest {
+            message: line.unwrap(),
+        });
+        let response = client.echo(request).await?;
+        println!("Response: {:?}", response);
+    }
+
+    println!("Hello, world!");
+    Ok(())
+}

experimental/uefi/client/src/main.rs

@@ -7,11 +7,13 @@ license = "Apache-2.0"
 
 [dependencies]
 anyhow = "*"
+bmrng = "*"
 clap = { version = "*", features = ["derive"] }
 command-fds = { version = "*", features = ["tokio"] }
 futures = "*"
 log = "*"
 env_logger = "*"
+prost = "*"
 serde = "*"
 tokio = { version = "*", features = [
   "io-std",
@@ -21,7 +23,12 @@ tokio = { version = "*", features = [
   "net",
   "process",
   "signal",
+  "sync",
 ] }
 tokio-serde-cbor = "*"
 # fixed to 0.6 because of tokio-serde-cbor
 tokio-util = { version = "~0.6", features = ["codec"] }
+tonic = "*"
+
+[build-dependencies]
+oak_utils = { path = "../../../oak_utils" }

experimental/uefi/loader/Cargo.toml

@@ -0,0 +1,30 @@
+//
+// Copyright 2022 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 oak_utils::{generate_grpc_code, CodegenOptions};
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    generate_grpc_code(
+        "../../../",
+        &["experimental/uefi/proto/echo.proto"],
+        CodegenOptions {
+            build_client: false,
+            build_server: true,
+            extern_paths: vec![],
+        },
+    )?;
+    Ok(())
+}

experimental/uefi/loader/build.rs

@@ -20,14 +20,12 @@ use clap::Parser;
 use futures::{stream::StreamExt, SinkExt};
 use log::info;
 use qemu::{Qemu, QemuParams};
-use tokio::{
-    io::{self, AsyncReadExt},
-    net::UnixStream,
-    signal,
-};
+use tokio::{io::AsyncReadExt, net::UnixStream, signal};
 use tokio_serde_cbor::Codec;
 use tokio_util::codec::Decoder;
+
 mod qemu;
+mod server;
 
 #[derive(Parser, Debug)]
 struct Args {
@@ -60,7 +58,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let (console_qemu, console) = UnixStream::pair()?;
     let (comms_qemu, mut comms) = UnixStream::pair()?;
 
-    let qemu = Qemu::start(QemuParams {
+    let mut qemu = Qemu::start(QemuParams {
         binary: cli.qemu.as_path(),
         firmware: cli.ovmf.as_path(),
         app: cli.uefi_app.as_path(),
@@ -91,39 +89,45 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     }
     info!("Leading junk on comms: {:?}", std::str::from_utf8(&junk));
 
-    // Set up the CBOR codec to handle the comms.
-    let codec: Codec<std::string::String, std::string::String> = Codec::new();
-    let (mut sink, mut stream) = codec.framed(comms).split();
+    // Use a bmrng channel for serial communication. The good thing about spawning a separate
+    // task for the serial communication is that it serializes (no pun intended) the communication,
+    // as right now we don't have any mechanisms to track multiple requests in flight.
+    let (tx, mut rx) = bmrng::unbounded_channel::<String, anyhow::Result<String>>();
 
-    // Bit hacky, but it's only temporary and turns out console I/O is complex.
     tokio::spawn(async move {
-        loop {
-            if let Some(result) = stream.next().await {
-                match result {
-                    Ok(msg) => info!("rx: {:?}", msg),
-                    Err(e) => {
-                        info!("recv error: {:?}", e);
+        // Set up the CBOR codec to handle the comms.
+        let codec: Codec<String, String> = Codec::new();
+        let mut channel = codec.framed(comms);
+        while let Ok((input, responder)) = rx.recv().await {
+            responder
+                .respond({
+                    if let Err(err) = channel.send(input).await {
+                        Err(anyhow::Error::from(err))
+                    } else {
+                        // TODO(#2726): Sometimes next() gives us a None. Figure out what's going on
+                        // in there.
+                        let mut response;
+                        loop {
+                            response = channel.next().await;
+                            if response.is_some() {
+                                break;
+                            }
+                        }
+                        response.unwrap().map_err(anyhow::Error::from)
                     }
-                };
-            }
-        }
-    });
-    tokio::spawn(async move {
-        let mut buf = [0; 1024];
-        loop {
-            let result = io::stdin().read(&mut buf).await.unwrap();
-            if result == 0 {
-                break;
-            } else {
-                let msg = std::str::from_utf8(&buf[..result]).unwrap().to_string();
-                info!("tx: {:?}", &msg);
-                sink.send(msg).await.unwrap();
-            }
+                })
+                .unwrap();
         }
     });
 
-    signal::ctrl_c().await?;
+    let server_future = server::server("127.0.0.1:8000".parse()?, tx);
 
+    // Wait until something dies or we get a signal to terminate.
+    tokio::select! {
+        _ = signal::ctrl_c() => {},
+        _ = server_future => {},
+        _ = qemu.wait() => {},
+    }
     // Clean up.
     qemu.kill().await?;
     Ok(())

experimental/uefi/loader/src/main.rs

@@ -112,12 +112,16 @@ impl Qemu {
         })
     }
 
+    pub async fn wait(&mut self) -> Result<std::process::ExitStatus> {
+        self.instance.wait().await.map_err(anyhow::Error::from)
+    }
+
     pub async fn kill(mut self) -> Result<std::process::ExitStatus> {
         info!("Cleaning up and shutting down.");
         self.console.shutdown().await?;
         self.comms.shutdown().await?;
         self.qmp.shutdown().await?;
         self.instance.start_kill()?;
-        self.instance.wait().await.map_err(anyhow::Error::from)
+        self.wait().await
     }
 }