Skip to content

Implement the sender side of the validation communication protocol in Rust #4280

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
pmikolajczyk41 wants to merge 20 commits into
base: master
Choose a base branch
from
+464 −166
Open

Implement the sender side of the validation communication protocol in Rust #4280

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
c884ee1
Create a new module. Copy utilities from jit/socket.rs. Make function…
pmikolajczyk41 Jan 27, 2026
6c24064
Checkpoint - implement half of the receiver. It's time to split the m…
pmikolajczyk41 Jan 27, 2026
4255889
Move things to separate module
pmikolajczyk41 Jan 27, 2026
0106a9b
Receiving is ready
pmikolajczyk41 Jan 27, 2026
7609642
Use it in JIT
pmikolajczyk41 Jan 27, 2026
24a56c9
Implement sender
pmikolajczyk41 Jan 27, 2026
542e35b
Result passing implementation
pmikolajczyk41 Jan 27, 2026
b7f246e
Remove socket.rs from jit
pmikolajczyk41 Jan 27, 2026
f392de8
Write tests. Fix bug
pmikolajczyk41 Jan 27, 2026
1917143
fmt
pmikolajczyk41 Jan 27, 2026
6ae5d18
Changelog
pmikolajczyk41 Jan 27, 2026
8d6c139
Add test for stylus archs
pmikolajczyk41 Jan 27, 2026
0b625e7
fmt
pmikolajczyk41 Jan 27, 2026
bef9dab
Merge branch 'master' into pmikolajczyk/nit-4364-protocol
pmikolajczyk41 Jan 28, 2026
4d0eb8f
Merge remote-tracking branch 'refs/remotes/origin/master' into pmikol…
pmikolajczyk41 Jan 28, 2026
eab1044
Internal changes
pmikolajczyk41 Jan 28, 2026
043c73f
Rename to preimage_map
pmikolajczyk41 Jan 28, 2026
f4dcfb4
Markers
pmikolajczyk41 Jan 28, 2026
bc24666
Nice construct from Igor
pmikolajczyk41 Jan 28, 2026
9da625e
Other comments from Igor
pmikolajczyk41 Jan 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions changelog/pmikolajczyk-nit-4364.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
### Internal
- Move the server side of the validation communication protocol from `jit` to `validation` crate.
- Add client side implementation. Add tests.
1 change: 0 additions & 1 deletion crates/jit/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@ mod caller_env;
pub mod machine;
mod prepare;
pub mod program;
pub mod socket;
pub mod stylus_backend;
mod test;
mod wasip1_stub;
Expand Down
16 changes: 8 additions & 8 deletions crates/jit/src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,8 @@
use arbutil::Color;
use clap::Parser;
use eyre::Result;
use jit::{
machine::Escape,
run,
socket::{report_error, report_success},
Opts,
};
use jit::{machine::Escape, run, Opts};
use validation::transfer::{send_failure_response, send_successful_response};
use wasmer::FrameInfo;

fn main() -> Result<()> {
Expand All @@ -34,7 +30,7 @@ fn main() -> Result<()> {
println!("{message}")
}
if let Some(mut socket) = result.socket {
report_error(&mut socket, message);
send_failure_response(&mut socket, &message)?;
}
if opts.validator.require_success {
std::process::exit(1);
Expand All @@ -47,7 +43,11 @@ fn main() -> Result<()> {
)
}
if let Some(mut socket) = result.socket {
report_success(&mut socket, &result.new_state, &result.memory_used);
send_successful_response(
&mut socket,
&result.new_state.into(),
result.memory_used.bytes().0 as u64,
)?;
}
}
Ok(())
Expand Down
20 changes: 2 additions & 18 deletions crates/jit/src/prepare.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,26 +3,10 @@

use crate::machine::WasmEnv;
use eyre::Ok;
use std::env;
use std::fs::File;
use std::io::BufReader;
use std::path::Path;
use validation::ValidationInput;

// local_target matches rawdb.LocalTarget() on the go side.
// While generating json_inputs file, one should make sure user_wasms map
// has entry for the system's arch that jit validation is being run on
pub fn local_target() -> String {
if env::consts::OS == "linux" {
match env::consts::ARCH {
"aarch64" => "arm64".to_string(),
"x86_64" => "amd64".to_string(),
_ => "host".to_string(),
}
} else {
"host".to_string()
}
}
use validation::{local_target, ValidationInput};

pub fn prepare_env_from_json(json_inputs: &Path, debug: bool) -> eyre::Result<WasmEnv> {
let file = File::open(json_inputs)?;
Expand Down Expand Up @@ -54,7 +38,7 @@ pub fn prepare_env_from_json(json_inputs: &Path, debug: bool) -> eyre::Result<Wa
}
}

if let Some(user_wasms) = data.user_wasms.get(&local_target()) {
if let Some(user_wasms) = data.user_wasms.get(local_target()) {
for (module_hash, module_asm) in user_wasms.iter() {
env.module_asms
.insert(*module_hash, module_asm.as_vec().into());
Expand Down
96 changes: 0 additions & 96 deletions crates/jit/src/socket.rs
View file
Open in desktop

This file was deleted.

55 changes: 17 additions & 38 deletions crates/jit/src/wavmio.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
use crate::{
caller_env::JitEnv,
machine::{Escape, MaybeEscape, WasmEnv, WasmEnvMut},
socket,
};
use arbutil::{Color, PreimageType};
use caller_env::{GuestPtr, MemAccess};
Expand All @@ -14,6 +13,8 @@ use std::{
net::TcpStream,
time::Instant,
};
use validation::local_target;
use validation::transfer::receive_validation_input;

/// Reads 32-bytes of global state.
pub fn get_global_state_bytes32(mut env: WasmEnvMut, idx: u32, out_ptr: GuestPtr) -> MaybeEscape {
Expand Down Expand Up @@ -281,49 +282,27 @@ fn ready_hostio(env: &mut WasmEnv) -> MaybeEscape {
socket.set_nodelay(true)?;

let mut reader = BufReader::new(socket.try_clone()?);
let stream = &mut reader;
let input = receive_validation_input(&mut reader)?;

let inbox_position = socket::read_u64(stream)?;
let position_within_message = socket::read_u64(stream)?;
let last_block_hash = socket::read_bytes32(stream)?;
let last_send_root = socket::read_bytes32(stream)?;
env.small_globals = [input.start_state.batch, input.start_state.pos_in_batch];
env.large_globals = [input.start_state.block_hash, input.start_state.send_root];

env.small_globals = [inbox_position, position_within_message];
env.large_globals = [last_block_hash, last_send_root];

while socket::read_u8(stream)? == socket::ANOTHER {
let position = socket::read_u64(stream)?;
let message = socket::read_bytes(stream)?;
env.sequencer_messages.insert(position, message);
for batch in input.batch_info {
env.sequencer_messages.insert(batch.number, batch.data);
}
while socket::read_u8(stream)? == socket::ANOTHER {
let position = socket::read_u64(stream)?;
let message = socket::read_bytes(stream)?;
env.delayed_messages.insert(position, message);
if input.has_delayed_msg {
env.delayed_messages
.insert(input.delayed_msg_nr, input.delayed_msg);
}
Comment on lines +293 to 296
Copy link
Contributor

@bragaigor bragaigor Jan 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

will there always be just one delayed_message? Asking because before, just like sequencer_messages, delayed_messages was wrapped in a while loop

Copy link
Member Author

@pmikolajczyk41 pmikolajczyk41 Jan 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

citing Tsahi:

Currently - it can only contain a single delayed message. There are future thoughs about having multiple delayed messages in a single block but that'll only happen after MEL and we could tackle it then.

Copy link
Member Author

@pmikolajczyk41 pmikolajczyk41 Jan 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so there's some lack of consistency in the code, but right now I tried to persist the existing behavior and data structs

Copy link
Contributor

@bragaigor bragaigor Jan 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ah got it, thanks!


let preimage_types = socket::read_u32(stream)?;
for _ in 0..preimage_types {
let preimage_ty = PreimageType::try_from(socket::read_u8(stream)?)
.map_err(|e| Escape::Failure(e.to_string()))?;
let map = env.preimages.entry(preimage_ty).or_default();
let preimage_count = socket::read_u32(stream)?;
for _ in 0..preimage_count {
let hash = socket::read_bytes32(stream)?;
let preimage = socket::read_bytes(stream)?;
map.insert(hash, preimage);
for (preimage_type, preimages) in input.preimages {
let preimage_map = env.preimages.entry(preimage_type).or_default();
for (hash, preimage) in preimages {
preimage_map.insert(hash, preimage);
}
}

let programs_count = socket::read_u32(stream)?;
for _ in 0..programs_count {
let module_hash = socket::read_bytes32(stream)?;
let module_asm = socket::read_boxed_slice(stream)?;
env.module_asms.insert(module_hash, module_asm.into());
}

if socket::read_u8(stream)? != socket::READY {
return Escape::hostio("failed to parse global state");
for (module_hash, module_asm) in &input.user_wasms[local_target()] {
env.module_asms
.insert(*module_hash, module_asm.as_vec().into());
}

let writer = BufWriter::new(socket);
Expand Down
36 changes: 31 additions & 5 deletions crates/validation/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,27 @@ use std::{
io::{self, BufRead},
};

pub mod transfer;

pub type PreimageMap = HashMap<PreimageType, HashMap<Bytes32, Vec<u8>>>;

pub const TARGET_ARM_64: &str = "arm64";
pub const TARGET_AMD_64: &str = "amd64";
pub const TARGET_HOST: &str = "host";

/// Counterpart to Go `rawdb.LocalTarget()`.
pub fn local_target() -> &'static str {
if cfg!(all(target_os = "linux", target_arch = "aarch64")) {
TARGET_ARM_64
} else if cfg!(all(target_os = "linux", target_arch = "x86_64")) {
TARGET_AMD_64
} else {
TARGET_HOST
}
}

/// Counterpart to Go `validator.GoGlobalState`.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct GoGlobalState {
#[serde(with = "As::<DisplayFromStr>")]
Expand All @@ -20,7 +39,7 @@ pub struct GoGlobalState {
}

/// Counterpart to Go `validator.server_api.BatchInfoJson`.
#[derive(Debug, Clone, Deserialize)]
#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct BatchInfo {
pub number: u64,
Expand All @@ -34,7 +53,7 @@ pub struct BatchInfo {
///
/// Note: The wrapped `Vec<u8>` is already `Base64` decoded before
/// `from(Vec<u8>)` is called by `serde`.
#[derive(Clone, Debug)]
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct UserWasm(Vec<u8>);

impl UserWasm {
Expand All @@ -60,7 +79,7 @@ impl TryFrom<Vec<u8>> for UserWasm {
}

/// Counterpart to Go `validator.server_api.InputJSON`.
#[derive(Clone, Debug, Deserialize)]
#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct ValidationInput {
pub id: u64,
Expand All @@ -70,7 +89,7 @@ pub struct ValidationInput {
rename = "PreimagesB64",
with = "As::<HashMap<TryFromInto<u8>, HashMap<Base64, Base64>>>"
)]
pub preimages: HashMap<PreimageType, HashMap<Bytes32, Vec<u8>>>,
pub preimages: PreimageMap,
pub batch_info: Vec<BatchInfo>,
#[serde(rename = "DelayedMsgB64", with = "As::<Base64>")]
pub delayed_msg: Vec<u8>,
Expand All @@ -86,4 +105,11 @@ impl ValidationInput {
pub fn from_reader<R: BufRead>(mut reader: R) -> io::Result<Self> {
Ok(serde_json::from_reader(&mut reader)?)
}

pub fn delayed_msg(&self) -> Option<BatchInfo> {
self.has_delayed_msg.then(|| BatchInfo {
number: self.delayed_msg_nr,
data: self.delayed_msg.clone(),
})
}
}
5 changes: 5 additions & 0 deletions crates/validation/src/transfer/markers.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
pub const SUCCESS: u8 = 0x0;
pub const FAILURE: u8 = 0x1;
// pub const PREIMAGE: u8 = 0x2; // legacy, not used
pub const ANOTHER: u8 = 0x3;
pub const READY: u8 = 0x4;
13 changes: 13 additions & 0 deletions crates/validation/src/transfer/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
use std::io;

mod markers;
mod primitives;
mod receiver;
mod sender;
#[cfg(test)]
mod tests;

pub use receiver::*;
pub use sender::*;

pub type IOResult<T> = Result<T, io::Error>;
Loading
Loading