Permit the orchestrator to write the derived key (#4776)

* Permit the orchestrator to write the derived key

* fix typo

* add additional cfg

* Write derived key and dice data

* improve naming

enclave_apps/oak_orchestrator/src/main.rs

@@ -21,21 +21,28 @@
 extern crate alloc;
 
 use oak_dice::evidence::Stage0DiceData;
-use oak_restricted_kernel_interface::{syscall::read, DICE_DATA_FD};
+use oak_restricted_kernel_interface::{syscall, DERIVED_KEY_FD, DICE_DATA_FD};
 use oak_restricted_kernel_sdk::{channel::FileDescriptorChannel, entrypoint};
 use zerocopy::{AsBytes, FromZeroes};
 
 fn read_stage0_dice_data() -> Stage0DiceData {
     let mut result = Stage0DiceData::new_zeroed();
     let buffer = result.as_bytes_mut();
-    let len = read(DICE_DATA_FD, buffer).expect("failed to read dice data");
+    let len = syscall::read(DICE_DATA_FD, buffer).expect("failed to read dice data");
     assert!(len == buffer.len(), "invalid dice data size");
     result
 }
 
 #[entrypoint]
 fn start() -> ! {
-    let dice_data = read_stage0_dice_data();
-    let channel = FileDescriptorChannel::default();
-    oak_restricted_kernel_orchestrator::entrypoint(channel, dice_data)
+    let (derived_key, dice_data) = {
+        let stage0_dice_data = read_stage0_dice_data();
+        let channel = FileDescriptorChannel::default();
+        oak_restricted_kernel_orchestrator::load_and_attest_app(channel, stage0_dice_data)
+    };
+
+    syscall::write(DERIVED_KEY_FD, derived_key.as_bytes()).expect("failed to write derived key");
+    syscall::write(DICE_DATA_FD, dice_data.as_bytes()).expect("failed to write dice data");
+
+    unimplemented!("launch app")
 }

oak_restricted_kernel/src/syscall/dice_data.rs

@@ -47,6 +47,7 @@ struct ReadState {
     index: usize,
 }
 
+#[cfg(feature = "initrd")]
 struct WriteState {
     data: RestrictedKernelDiceData,
     index: usize,

oak_restricted_kernel/src/syscall/key.rs

@@ -22,27 +22,79 @@ use oak_restricted_kernel_interface::{Errno, DERIVED_KEY_FD};
 use super::fd::{copy_max_slice, FileDescriptor};
 
 #[derive(Default)]
-struct DerivedKeyDescriptor {
-    key: DerivedKey,
+struct DerivedKeyState {
+    data: DerivedKey,
     index: usize,
 }
 
+enum DerivedKeyDescriptor {
+    Readable(DerivedKeyState),
+    #[cfg(feature = "initrd")]
+    Writeable(DerivedKeyState),
+}
+
 impl DerivedKeyDescriptor {
+    #[cfg(not(feature = "initrd"))]
     fn new(key: DerivedKey) -> Self {
-        Self { index: 0, key }
+        Self::Readable(DerivedKeyState {
+            index: 0,
+            data: key,
+        })
+    }
+
+    #[cfg(feature = "initrd")]
+    fn new() -> Self {
+        Self::Writeable(DerivedKeyState::default())
     }
 }
 
 impl FileDescriptor for DerivedKeyDescriptor {
     fn read(&mut self, buf: &mut [u8]) -> Result<isize, oak_restricted_kernel_interface::Errno> {
-        let length = copy_max_slice(&self.key[self.index..], buf);
-        self.index += length;
-        Ok(length as isize)
+        match self {
+            #[cfg(feature = "initrd")]
+            DerivedKeyDescriptor::Writeable(_write_state) => Err(Errno::EINVAL),
+            DerivedKeyDescriptor::Readable(read_state) => {
+                let data_as_slice = read_state.data.as_mut_slice();
+                let length = copy_max_slice(&data_as_slice[read_state.index..], buf);
+                read_state.index += length;
+                Ok(length as isize)
+            }
+        }
     }
 
-    fn write(&mut self, _buf: &[u8]) -> Result<isize, oak_restricted_kernel_interface::Errno> {
-        // Writing is not supported.
-        Err(Errno::EINVAL)
+    fn write(&mut self, buf: &[u8]) -> Result<isize, oak_restricted_kernel_interface::Errno> {
+        match self {
+            DerivedKeyDescriptor::Readable(_read_state) => Err(Errno::EINVAL),
+            #[cfg(feature = "initrd")]
+            DerivedKeyDescriptor::Writeable(write_state) => {
+                let data_as_slice =
+                    <DerivedKey as zerocopy::AsBytes>::as_bytes_mut(&mut write_state.data);
+
+                if buf.len() > data_as_slice[write_state.index..].len() {
+                    // the key has a known size. If the app keeps writing to
+                    // this FD after it has written all of it, it's doing something wrong.
+                    return Err(Errno::EINVAL);
+                }
+
+                let length = copy_max_slice(buf, &mut data_as_slice[write_state.index..]);
+
+                write_state.index += length;
+
+                if write_state.index == data_as_slice.len() {
+                    let derived_key =
+                        <DerivedKey as zerocopy::FromBytes>::read_from(data_as_slice).unwrap();
+                    let _ = core::mem::replace(
+                        self,
+                        Self::Readable(DerivedKeyState {
+                            index: 0,
+                            data: derived_key,
+                        }),
+                    );
+                }
+
+                Ok(length as isize)
+            }
+        }
     }
 
     fn sync(&mut self) -> Result<(), oak_restricted_kernel_interface::Errno> {
@@ -51,8 +103,14 @@ impl FileDescriptor for DerivedKeyDescriptor {
 }
 
 /// Registers a file descriptor for reading a derived key (0x21)
-pub fn register(key: DerivedKey) {
-    super::fd::register(DERIVED_KEY_FD, Box::new(DerivedKeyDescriptor::new(key)))
-        .map_err(|_| ()) // throw away the box
-        .expect("DerivedKeyDescriptor already registered");
+pub fn register(#[cfg(not(feature = "initrd"))] key: DerivedKey) {
+    super::fd::register(
+        DERIVED_KEY_FD,
+        Box::new(DerivedKeyDescriptor::new(
+            #[cfg(not(feature = "initrd"))]
+            key,
+        )),
+    )
+    .map_err(|_| ()) // throw away the box
+    .expect("DerivedKeyDescriptor already registered");
 }

oak_restricted_kernel/src/syscall/mod.rs

@@ -17,7 +17,6 @@
 mod channel;
 pub mod dice_data;
 mod fd;
-#[cfg(not(feature = "initrd"))]
 mod key;
 pub mod mmap;
 mod process;
@@ -74,8 +73,10 @@ pub fn enable_syscalls(
 ) {
     channel::register(channel);
     stdio::register();
-    #[cfg(not(feature = "initrd"))]
-    key::register(derived_key);
+    key::register(
+        #[cfg(not(feature = "initrd"))]
+        derived_key,
+    );
     dice_data::register(dice_data);
 
     // Allocate a stack for the system call handler.

oak_restricted_kernel_orchestrator/src/lib.rs

@@ -21,13 +21,22 @@ extern crate alloc;
 use oak_channel::basic_framed::load_raw;
 use oak_dice::evidence::Stage0DiceData;
 
-pub fn entrypoint<C: oak_channel::Channel>(mut channel: C, _stage0_dice_data: Stage0DiceData) -> ! {
+pub fn load_and_attest_app<C: oak_channel::Channel>(
+    mut channel: C,
+    stage0_dice_data: Stage0DiceData,
+) -> (
+    oak_restricted_kernel_dice::DerivedKey,
+    oak_dice::evidence::RestrictedKernelDiceData,
+) {
     let application_bytes = load_raw::<C, 4096>(&mut channel).expect("failed to load");
     log::info!("Binary loaded, size: {}", application_bytes.len());
     let app_digest = oak_restricted_kernel_dice::measure_app_digest_sha2_256(&application_bytes);
     log::info!(
         "Application digest (sha2-256): {}",
         app_digest.map(|x| alloc::format!("{:02x}", x)).join("")
     );
-    unimplemented!();
+    let derived_key =
+        oak_restricted_kernel_dice::generate_derived_key(&stage0_dice_data, &app_digest);
+    let dice_data = oak_restricted_kernel_dice::generate_dice_data(stage0_dice_data, &app_digest);
+    (derived_key, dice_data)
 }