From 9f9a724f6ac416a8b13cb4555636de081f78f4f5 Mon Sep 17 00:00:00 2001 From: clabby Date: Mon, 8 Apr 2024 13:28:21 -0400 Subject: [PATCH] feat(preimage): `OracleServer` + `HintReader` (#96) * feat(preimage): `OracleServer` + `HintReader` Adds the host end of the `PreimageOracle` ABI plumbing. This includes two new traits: * `PreimageOracleServer` * `HintReaderServer` as well as implementations of both of them that compliment the existing client handles, the `OracleReader` and `HintWriter`. * feat(preimage): tracing --- Cargo.lock | 18 +++++ Cargo.toml | 2 +- crates/common/src/io.rs | 4 +- crates/preimage/Cargo.toml | 4 ++ crates/preimage/src/hint.rs | 120 +++++++++++++++++++++++++++++++- crates/preimage/src/key.rs | 37 +++++++++- crates/preimage/src/lib.rs | 6 +- crates/preimage/src/oracle.rs | 127 ++++++++++++++++++++++++++++------ crates/preimage/src/traits.rs | 32 ++++++++- 9 files changed, 314 insertions(+), 36 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a96d113ed..f82f8a7f5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -575,11 +575,13 @@ dependencies = [ name = "kona-preimage" version = "0.0.1" dependencies = [ + "alloy-primitives", "anyhow", "cfg-if", "kona-common", "tempfile", "tokio", + "tracing", ] [[package]] @@ -1137,6 +1139,22 @@ dependencies = [ "syn 2.0.50", ] +[[package]] +name = "tracing" +version = "0.1.40" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3523ab5a71916ccf420eebdf5521fcef02141234bbc0b8a49f2fdc4544364ef" +dependencies = [ + "pin-project-lite", + "tracing-core", +] + +[[package]] +name = "tracing-core" +version = "0.1.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c06d3da6113f116aaee68e4d601191614c9053067f9ab7f6edbcb161237daa54" + [[package]] name = "typenum" version = "1.17.0" diff --git a/Cargo.toml b/Cargo.toml index 97a256d83..07fed277c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,7 +13,7 @@ exclude = ["**/target", "benches/", "tests"] [workspace.dependencies] anyhow = { version = "1.0.79", default-features = false } -tracing = "0.1.40" +tracing = { version = "0.1.40", default-features = false } cfg-if = "1.0.0" [profile.dev] diff --git a/crates/common/src/io.rs b/crates/common/src/io.rs index 6088b12ab..b6c12f0ab 100644 --- a/crates/common/src/io.rs +++ b/crates/common/src/io.rs @@ -78,8 +78,8 @@ mod native_io { .write(buf) .map_err(|e| anyhow!("Error writing to buffer to file descriptor: {e}"))?; - // Reset the cursor back to 0 for the reader. - file.seek(SeekFrom::Start(0)) + // Reset the cursor back to before the data we just wrote for the reader's consumption. + file.seek(SeekFrom::Current(-(buf.len() as i64))) .map_err(|e| anyhow!("Failed to reset file cursor to 0: {e}"))?; // forget the file descriptor so that the `Drop` impl doesn't close it. diff --git a/crates/preimage/Cargo.toml b/crates/preimage/Cargo.toml index 84d6edb62..2fc4acce8 100644 --- a/crates/preimage/Cargo.toml +++ b/crates/preimage/Cargo.toml @@ -12,10 +12,14 @@ homepage.workspace = true # workspace anyhow.workspace = true cfg-if.workspace = true +tracing.workspace = true # local kona-common = { path = "../common", version = "0.0.1" } +# External +alloy-primitives = { version = "0.7.0", default-features = false } + [dev-dependencies] tokio = { version = "1.36.0", features = ["full"] } tempfile = "3.10.0" diff --git a/crates/preimage/src/hint.rs b/crates/preimage/src/hint.rs index 99861339a..4ed6159b5 100644 --- a/crates/preimage/src/hint.rs +++ b/crates/preimage/src/hint.rs @@ -1,6 +1,7 @@ -use crate::{traits::HintWriterClient, PipeHandle}; -use alloc::vec; +use crate::{traits::HintWriterClient, HintReaderServer, PipeHandle}; +use alloc::{string::String, vec}; use anyhow::Result; +use tracing::{debug, error}; /// A [HintWriter] is a high-level interface to the hint pipe. It provides a way to write hints to /// the host. @@ -26,13 +27,128 @@ impl HintWriterClient for HintWriter { hint_bytes[0..4].copy_from_slice(u32::to_be_bytes(hint.len() as u32).as_ref()); hint_bytes[4..].copy_from_slice(hint.as_bytes()); + debug!(target: "hint_writer", "Writing hint \"{hint}\""); + // Write the hint to the host. self.pipe_handle.write(&hint_bytes)?; + debug!(target: "hint_writer", "Successfully wrote hint"); + // Read the hint acknowledgement from the host. let mut hint_ack = [0u8; 1]; self.pipe_handle.read_exact(&mut hint_ack)?; + debug!(target: "hint_writer", "Received hint acknowledgement"); + + Ok(()) + } +} + +/// A [HintReader] is a router for hints sent by the [HintWriter] from the client program. It +/// provides a way for the host to prepare preimages for reading. +#[derive(Debug, Clone, Copy)] +pub struct HintReader { + pipe_handle: PipeHandle, +} + +impl HintReader { + /// Create a new [HintReader] from a [PipeHandle]. + pub fn new(pipe_handle: PipeHandle) -> Self { + Self { pipe_handle } + } +} + +impl HintReaderServer for HintReader { + fn next_hint(&self, mut route_hint: impl FnMut(String) -> Result<()>) -> Result<()> { + // Read the length of the raw hint payload. + let mut len_buf = [0u8; 4]; + self.pipe_handle.read_exact(&mut len_buf)?; + let len = u32::from_be_bytes(len_buf); + + // Read the raw hint payload. + let mut raw_payload = vec![0u8; len as usize]; + self.pipe_handle.read_exact(raw_payload.as_mut_slice())?; + let payload = String::from_utf8(raw_payload) + .map_err(|e| anyhow::anyhow!("Failed to decode hint payload: {e}"))?; + + debug!(target: "hint_reader", "Successfully read hint: \"{payload}\""); + + // Route the hint + if let Err(e) = route_hint(payload) { + // Write back on error to prevent blocking the client. + self.pipe_handle.write(&[0x00])?; + + error!("Failed to route hint: {e}"); + anyhow::bail!("Failed to rout hint: {e}"); + } + + // Write back an acknowledgement to the client to unblock their process. + self.pipe_handle.write(&[0x00])?; + + debug!(target: "hint_reader", "Successfully routed and acknowledged hint"); + Ok(()) } } +#[cfg(test)] +mod test { + extern crate std; + + use super::*; + use alloc::vec::Vec; + use kona_common::FileDescriptor; + use std::{fs::File, os::fd::AsRawFd}; + use tempfile::tempfile; + + /// Test struct containing the [HintReader] and [HintWriter]. The [File]s are stored in this + /// struct so that they are not dropped until the end of the test. + #[derive(Debug)] + struct ClientAndHost { + hint_writer: HintWriter, + hint_reader: HintReader, + _read_file: File, + _write_file: File, + } + + /// Helper for creating a new [HintReader] and [HintWriter] for testing. The file channel is + /// over two temporary files. + fn client_and_host() -> ClientAndHost { + let (read_file, write_file) = (tempfile().unwrap(), tempfile().unwrap()); + let (read_fd, write_fd) = ( + FileDescriptor::Wildcard(read_file.as_raw_fd().try_into().unwrap()), + FileDescriptor::Wildcard(write_file.as_raw_fd().try_into().unwrap()), + ); + let client_handle = PipeHandle::new(read_fd, write_fd); + let host_handle = PipeHandle::new(write_fd, read_fd); + + let hint_writer = HintWriter::new(client_handle); + let hint_reader = HintReader::new(host_handle); + + ClientAndHost { hint_writer, hint_reader, _read_file: read_file, _write_file: write_file } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_hint_client_and_host() { + const MOCK_DATA: &str = "test-hint 0xfacade"; + + let sys = client_and_host(); + let (hint_writer, hint_reader) = (sys.hint_writer, sys.hint_reader); + + let client = tokio::task::spawn(async move { hint_writer.write(MOCK_DATA) }); + let host = tokio::task::spawn(async move { + let mut v = Vec::new(); + let route_hint = |hint: String| { + v.push(hint.clone()); + Ok(()) + }; + hint_reader.next_hint(route_hint).unwrap(); + + assert_eq!(v.len(), 1); + + v.remove(0) + }); + + let (_, h) = tokio::join!(client, host); + assert_eq!(h.unwrap(), MOCK_DATA); + } +} diff --git a/crates/preimage/src/key.rs b/crates/preimage/src/key.rs index 4c1e9a565..dae79ea9e 100644 --- a/crates/preimage/src/key.rs +++ b/crates/preimage/src/key.rs @@ -1,8 +1,10 @@ //! Contains the [PreimageKey] type, which is used to identify preimages that may be fetched from //! the preimage oracle. +use alloy_primitives::B256; + /// -#[derive(Debug, Default, Clone, Copy, Eq, PartialEq)] +#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Hash)] #[repr(u8)] pub enum PreimageKeyType { /// Local key types are local to a given instance of a fault-proof and context dependent. @@ -23,6 +25,21 @@ pub enum PreimageKeyType { Blob = 5, } +impl TryFrom for PreimageKeyType { + type Error = anyhow::Error; + + fn try_from(value: u8) -> Result { + Ok(match value { + 1 => PreimageKeyType::Local, + 2 => PreimageKeyType::Keccak256, + 3 => PreimageKeyType::GlobalGeneric, + 4 => PreimageKeyType::Sha256, + 5 => PreimageKeyType::Blob, + _ => anyhow::bail!("Invalid preimage key type"), + }) + } +} + /// A preimage key is a 32-byte value that identifies a preimage that may be fetched from the /// oracle. /// @@ -31,7 +48,7 @@ pub enum PreimageKeyType { /// |---------|-------------| /// | [0, 1) | Type byte | /// | [1, 32) | Data | -#[derive(Debug, Default, Clone, Copy)] +#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Hash)] pub struct PreimageKey { data: [u8; 31], key_type: PreimageKeyType, @@ -69,6 +86,22 @@ impl From for [u8; 32] { } } +impl TryFrom<[u8; 32]> for PreimageKey { + type Error = anyhow::Error; + + fn try_from(value: [u8; 32]) -> Result { + let key_type = PreimageKeyType::try_from(value[0])?; + Ok(Self::new(value, key_type)) + } +} + +impl core::fmt::Display for PreimageKey { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let raw: [u8; 32] = (*self).into(); + write!(f, "{}", B256::from(raw)) + } +} + #[cfg(test)] mod test { use super::*; diff --git a/crates/preimage/src/lib.rs b/crates/preimage/src/lib.rs index 7dfaf40b6..61906eb0b 100644 --- a/crates/preimage/src/lib.rs +++ b/crates/preimage/src/lib.rs @@ -10,13 +10,13 @@ mod key; pub use key::{PreimageKey, PreimageKeyType}; mod oracle; -pub use oracle::OracleReader; +pub use oracle::{OracleReader, OracleServer}; mod hint; -pub use hint::HintWriter; +pub use hint::{HintReader, HintWriter}; mod pipe; pub use pipe::PipeHandle; mod traits; -pub use traits::{HintWriterClient, PreimageOracleClient}; +pub use traits::{HintReaderServer, HintWriterClient, PreimageOracleClient, PreimageOracleServer}; diff --git a/crates/preimage/src/oracle.rs b/crates/preimage/src/oracle.rs index 2bd9710e0..5d05bbc94 100644 --- a/crates/preimage/src/oracle.rs +++ b/crates/preimage/src/oracle.rs @@ -1,6 +1,7 @@ -use crate::{traits::PreimageOracleClient, PipeHandle, PreimageKey}; +use crate::{PipeHandle, PreimageKey, PreimageOracleClient, PreimageOracleServer}; use alloc::vec::Vec; use anyhow::{bail, Result}; +use tracing::debug; /// An [OracleReader] is a high-level interface to the preimage oracle. #[derive(Debug, Clone, Copy)] @@ -17,7 +18,7 @@ impl OracleReader { /// Set the preimage key for the global oracle reader. This will overwrite any existing key, and /// block until the host has prepared the preimage and responded with the length of the /// preimage. - fn write_key(&mut self, key: PreimageKey) -> Result { + fn write_key(&self, key: PreimageKey) -> Result { // Write the key to the host so that it can prepare the preimage. let key_bytes: [u8; 32] = key.into(); self.pipe_handle.write(&key_bytes)?; @@ -32,22 +33,32 @@ impl OracleReader { impl PreimageOracleClient for OracleReader { /// Get the data corresponding to the currently set key from the host. Return the data in a new /// heap allocated `Vec` - fn get(&mut self, key: PreimageKey) -> Result> { + fn get(&self, key: PreimageKey) -> Result> { + debug!(target: "oracle_client", "Requesting data from preimage oracle. Key {key}"); + let length = self.write_key(key)?; let mut data_buffer = alloc::vec![0; length]; + debug!(target: "oracle_client", "Reading data from preimage oracle. Key {key}"); + // Grab a read lock on the preimage pipe to read the data. self.pipe_handle.read_exact(&mut data_buffer)?; + debug!(target: "oracle_client", "Successfully read data from preimage oracle. Key: {key}"); + Ok(data_buffer) } /// Get the data corresponding to the currently set key from the host. Write the data into the /// provided buffer - fn get_exact(&mut self, key: PreimageKey, buf: &mut [u8]) -> Result<()> { + fn get_exact(&self, key: PreimageKey, buf: &mut [u8]) -> Result<()> { + debug!(target: "oracle_client", "Requesting data from preimage oracle. Key {key}"); + // Write the key to the host and read the length of the preimage. let length = self.write_key(key)?; + debug!(target: "oracle_client", "Reading data from preimage oracle. Key {key}"); + // Ensure the buffer is the correct size. if buf.len() != length { bail!("Buffer size {} does not match preimage size {}", buf.len(), length); @@ -55,6 +66,50 @@ impl PreimageOracleClient for OracleReader { self.pipe_handle.read_exact(buf)?; + debug!(target: "oracle_client", "Successfully read data from preimage oracle. Key: {key}"); + + Ok(()) + } +} + +/// An [OracleServer] is a router for the host to serve data back to the client [OracleReader]. +#[derive(Debug, Clone, Copy)] +pub struct OracleServer { + pipe_handle: PipeHandle, +} + +impl OracleServer { + /// Create a new [OracleServer] from a [PipeHandle]. + pub fn new(pipe_handle: PipeHandle) -> Self { + Self { pipe_handle } + } +} + +impl PreimageOracleServer for OracleServer { + fn next_preimage_request<'a>( + &self, + mut get_preimage: impl FnMut(PreimageKey) -> Result<&'a Vec>, + ) -> Result<()> { + // Read the preimage request from the client, and throw early if there isn't is any. + let mut buf = [0u8; 32]; + self.pipe_handle.read_exact(&mut buf)?; + let preimage_key = PreimageKey::try_from(buf)?; + + debug!(target: "oracle_server", "Fetching preimage for key {preimage_key}"); + + // Fetch the preimage value from the preimage getter. + let value = get_preimage(preimage_key)?; + + // Write the length as a big-endian u64 followed by the data. + let data = [(value.len() as u64).to_be_bytes().as_ref(), value.as_ref()] + .into_iter() + .flatten() + .copied() + .collect::>(); + self.pipe_handle.write(data.as_slice())?; + + debug!(target: "oracle_server", "Successfully wrote preimage data for key {preimage_key}"); + Ok(()) } } @@ -65,27 +120,24 @@ mod test { use super::*; use crate::PreimageKeyType; + use alloy_primitives::keccak256; use kona_common::FileDescriptor; - use std::{fs::File, os::fd::AsRawFd}; + use std::{collections::HashMap, fs::File, os::fd::AsRawFd}; use tempfile::tempfile; - /// Test struct containing the [OracleReader] and a [PipeHandle] for the host, plus the open + /// Test struct containing the [OracleReader] and a [OracleServer] for the host, plus the open /// [File]s. The [File]s are stored in this struct so that they are not dropped until the /// end of the test. - /// - /// TODO: Swap host pipe handle to oracle writer once it exists. #[derive(Debug)] struct ClientAndHost { oracle_reader: OracleReader, - host_handle: PipeHandle, + oracle_server: OracleServer, _read_file: File, _write_file: File, } - /// Helper for creating a new [OracleReader] and [PipeHandle] for testing. The file channel is + /// Helper for creating a new [OracleReader] and [OracleServer] for testing. The file channel is /// over two temporary files. - /// - /// TODO: Swap host pipe handle to oracle writer once it exists. fn client_and_host() -> ClientAndHost { let (read_file, write_file) = (tempfile().unwrap(), tempfile().unwrap()); let (read_fd, write_fd) = ( @@ -96,27 +148,56 @@ mod test { let host_handle = PipeHandle::new(write_fd, read_fd); let oracle_reader = OracleReader::new(client_handle); + let oracle_server = OracleServer::new(host_handle); - ClientAndHost { oracle_reader, host_handle, _read_file: read_file, _write_file: write_file } + ClientAndHost { + oracle_reader, + oracle_server, + _read_file: read_file, + _write_file: write_file, + } } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn test_oracle_reader() { - const MOCK_DATA: &[u8] = b"1234567890"; + async fn test_oracle_client_and_host() { + const MOCK_DATA_A: &[u8] = b"1234567890"; + const MOCK_DATA_B: &[u8] = b"FACADE"; + let key_a: PreimageKey = + PreimageKey::new(*keccak256(MOCK_DATA_A), PreimageKeyType::Keccak256); + let key_b: PreimageKey = + PreimageKey::new(*keccak256(MOCK_DATA_B), PreimageKeyType::Keccak256); + + let mut preimages = HashMap::new(); + preimages.insert(key_a, MOCK_DATA_A.to_vec()); + preimages.insert(key_b, MOCK_DATA_B.to_vec()); + let sys = client_and_host(); - let (mut oracle_reader, host_handle) = (sys.oracle_reader, sys.host_handle); + let (oracle_reader, oracle_server) = (sys.oracle_reader, sys.oracle_server); let client = tokio::task::spawn(async move { - oracle_reader.get(PreimageKey::new([0u8; 32], PreimageKeyType::Keccak256)).unwrap() + let contents_a = oracle_reader.get(key_a).unwrap(); + let contents_b = oracle_reader.get(key_b).unwrap(); + + // Drop the file descriptors to close the pipe, stopping the host's blocking loop on + // waiting for client requests. + drop(sys); + + (contents_a, contents_b) }); let host = tokio::task::spawn(async move { - let mut length_and_data: [u8; 8 + 10] = [0u8; 8 + 10]; - length_and_data[0..8].copy_from_slice(&u64::to_be_bytes(MOCK_DATA.len() as u64)); - length_and_data[8..18].copy_from_slice(MOCK_DATA); - host_handle.write(&length_and_data).unwrap(); + let get_preimage = + |key| preimages.get(&key).ok_or(anyhow::anyhow!("Preimage not available")); + + loop { + if oracle_server.next_preimage_request(get_preimage).is_err() { + break; + } + } }); - let (r, _) = tokio::join!(client, host); - assert_eq!(r.unwrap(), MOCK_DATA); + let (client, _) = tokio::join!(client, host); + let (contents_a, contents_b) = client.unwrap(); + assert_eq!(contents_a, MOCK_DATA_A); + assert_eq!(contents_b, MOCK_DATA_B); } } diff --git a/crates/preimage/src/traits.rs b/crates/preimage/src/traits.rs index 4954701c8..c2a20d4eb 100644 --- a/crates/preimage/src/traits.rs +++ b/crates/preimage/src/traits.rs @@ -1,5 +1,5 @@ use crate::PreimageKey; -use alloc::vec::Vec; +use alloc::{string::String, vec::Vec}; use anyhow::Result; /// A [PreimageOracleClient] is a high-level interface to read data from the host, keyed by a @@ -11,7 +11,7 @@ pub trait PreimageOracleClient { /// # Returns /// - `Ok(Vec)` if the data was successfully fetched from the host. /// - `Err(_)` if the data could not be fetched from the host. - fn get(&mut self, key: PreimageKey) -> Result>; + fn get(&self, key: PreimageKey) -> Result>; /// Get the data corresponding to the currently set key from the host. Writes the data into the /// provided buffer. @@ -19,7 +19,7 @@ pub trait PreimageOracleClient { /// # Returns /// - `Ok(())` if the data was successfully written into the buffer. /// - `Err(_)` if the data could not be written into the buffer. - fn get_exact(&mut self, key: PreimageKey, buf: &mut [u8]) -> Result<()>; + fn get_exact(&self, key: PreimageKey, buf: &mut [u8]) -> Result<()>; } /// A [HintWriterClient] is a high-level interface to the hint pipe. It provides a way to write @@ -33,3 +33,29 @@ pub trait HintWriterClient { /// - `Err(_)` if the hint could not be written to the host. fn write(&self, hint: &str) -> Result<()>; } + +/// A [PreimageOracleServer] is a high-level interface to accept read requests from the client and +/// write the preimage data to the client pipe. +pub trait PreimageOracleServer { + /// Get the next preimage request and return the response to the client. + /// + /// # Returns + /// - `Ok(())` if the data was successfully written into the client pipe. + /// - `Err(_)` if the data could not be written to the client. + fn next_preimage_request<'a>( + &self, + get_preimage: impl FnMut(PreimageKey) -> Result<&'a Vec>, + ) -> Result<()>; +} + +/// A [HintReaderServer] is a high-level interface to read preimage hints from the +/// [HintWriterClient] and prepare them for consumption by the client program. +pub trait HintReaderServer { + /// Get the next hint request and return the acknowledgement to the client. + /// + /// # Returns + /// - `Ok(())` if the hint was received and the client was notified of the host's + /// acknowledgement. + /// - `Err(_)` if the hint was not received correctly. + fn next_hint(&self, route_hint: impl FnMut(String) -> Result<()>) -> Result<()>; +}