mirror of
https://github.com/openai/codex.git
synced 2026-11-09 20:50:33 +00:00
feat: add Noise relay primitives
Co-authored-by: Codex noreply@openai.com
This commit is contained in:
viyatb-oai
co-authored by
Codex noreply@openai.com
Codex noreply@openai.com
parent
2d5c264ebc
commit
b075d0ce91
Generated
+89
-1
@@ -196,6 +196,20 @@ dependencies = [
|
||||
"cpufeatures 0.2.17",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aes-gcm"
|
||||
version = "0.10.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "831010a0f742e1209b3bcea8fab6a8e149051ba6099432c8cb2cc117dec3ead1"
|
||||
dependencies = [
|
||||
"aead",
|
||||
"aes",
|
||||
"cipher",
|
||||
"ctr",
|
||||
"ghash",
|
||||
"subtle",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "age"
|
||||
version = "0.11.2"
|
||||
@@ -468,6 +482,9 @@ name = "arrayvec"
|
||||
version = "0.7.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50"
|
||||
dependencies = [
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ascii"
|
||||
@@ -1748,6 +1765,23 @@ version = "1.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3a822ea5bc7590f9d40f1ba12c0dc3c2760f3482c6984db1573ad11031420831"
|
||||
|
||||
[[package]]
|
||||
name = "clatter"
|
||||
version = "2.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6fed49fa357a85c377c0f920e86100f5326111b09ad69f6de684e324e3ad8097"
|
||||
dependencies = [
|
||||
"aes-gcm",
|
||||
"arrayvec",
|
||||
"displaydoc",
|
||||
"getrandom 0.3.4",
|
||||
"rand_core 0.6.4",
|
||||
"sha2 0.10.9",
|
||||
"thiserror-no-std",
|
||||
"x25519-dalek",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "clipboard-win"
|
||||
version = "5.4.1"
|
||||
@@ -2778,9 +2812,11 @@ dependencies = [
|
||||
"anyhow",
|
||||
"arc-swap",
|
||||
"async-trait",
|
||||
"aws-lc-rs",
|
||||
"axum",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"clatter",
|
||||
"codex-api",
|
||||
"codex-app-server-protocol",
|
||||
"codex-client",
|
||||
@@ -2808,6 +2844,7 @@ dependencies = [
|
||||
"tokio-util",
|
||||
"toml 0.9.11+spec-1.1.0",
|
||||
"tracing",
|
||||
"url",
|
||||
"uuid",
|
||||
"wiremock",
|
||||
]
|
||||
@@ -4784,6 +4821,15 @@ version = "0.0.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "52560adf09603e58c9a7ee1fe1dcb95a16927b17c127f0ac02d6e768a0e25bc1"
|
||||
|
||||
[[package]]
|
||||
name = "ctr"
|
||||
version = "0.9.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0369ee1ad671834580515889b80f2ea915f23b8be8d0daa4bbaf2ac5c7590835"
|
||||
dependencies = [
|
||||
"cipher",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ctutils"
|
||||
version = "0.4.2"
|
||||
@@ -6215,6 +6261,16 @@ dependencies = [
|
||||
"wasip3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ghash"
|
||||
version = "0.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f0d8a4362ccb29cb0b265253fb0a2728f592895ee6854fd9bc13f2ffda266ff1"
|
||||
dependencies = [
|
||||
"opaque-debug",
|
||||
"polyval",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "gif"
|
||||
version = "0.14.1"
|
||||
@@ -10091,6 +10147,18 @@ dependencies = [
|
||||
"universal-hash",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "polyval"
|
||||
version = "0.6.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9d1fe60d06143b2430aa532c94cfe9e29783047f06c0d7fd359a9a51b729fa25"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"cpufeatures 0.2.17",
|
||||
"opaque-debug",
|
||||
"universal-hash",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "portable-atomic"
|
||||
version = "1.13.1"
|
||||
@@ -11477,7 +11545,7 @@ dependencies = [
|
||||
"security-framework 3.5.1",
|
||||
"security-framework-sys",
|
||||
"webpki-root-certs",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -13151,6 +13219,26 @@ dependencies = [
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror-impl-no-std"
|
||||
version = "2.0.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "58e6318948b519ba6dc2b442a6d0b904ebfb8d411a3ad3e07843615a72249758"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 1.0.109",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror-no-std"
|
||||
version = "2.0.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a3ad459d94dd517257cc96add8a43190ee620011bb6e6cdc82dafd97dfafafea"
|
||||
dependencies = [
|
||||
"thiserror-impl-no-std",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thread_local"
|
||||
version = "1.1.9"
|
||||
|
||||
@@ -261,6 +261,7 @@ async-stream = "0.3.6"
|
||||
async-trait = "0.1.89"
|
||||
aws-config = "1"
|
||||
aws-credential-types = "1"
|
||||
aws-lc-rs = { version = "=1.16.2", default-features = false, features = ["non-fips"] }
|
||||
aws-sigv4 = "1"
|
||||
aws-types = "1"
|
||||
axum = { version = "0.8", default-features = false }
|
||||
@@ -271,6 +272,13 @@ chardetng = "0.1.17"
|
||||
chrono = "0.4.43"
|
||||
clap = "4"
|
||||
clap_complete = "4"
|
||||
clatter = { version = "2.2.0", default-features = false, features = [
|
||||
"alloc",
|
||||
"getrandom",
|
||||
"use-25519",
|
||||
"use-aes-gcm",
|
||||
"use-sha",
|
||||
] }
|
||||
color-eyre = "0.6.3"
|
||||
constant_time_eq = "0.3.1"
|
||||
crossbeam-channel = "0.5.15"
|
||||
|
||||
@@ -13,9 +13,11 @@ workspace = true
|
||||
[dependencies]
|
||||
arc-swap = { workspace = true }
|
||||
async-trait = { workspace = true }
|
||||
aws-lc-rs = { workspace = true }
|
||||
axum = { workspace = true, features = ["http1", "tokio", "ws"] }
|
||||
base64 = { workspace = true }
|
||||
bytes = { workspace = true }
|
||||
clatter = { workspace = true }
|
||||
codex-app-server-protocol = { workspace = true }
|
||||
codex-api = { workspace = true }
|
||||
codex-client = { workspace = true }
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
use aws_lc_rs::kem::Ciphertext;
|
||||
use aws_lc_rs::kem::DecapsulationKey;
|
||||
use aws_lc_rs::kem::EncapsulationKey;
|
||||
use aws_lc_rs::kem::ML_KEM_768;
|
||||
use clatter::KeyPair;
|
||||
use clatter::bytearray::ByteArray;
|
||||
use clatter::bytearray::HeapArray;
|
||||
use clatter::bytearray::SensitiveByteArray;
|
||||
use clatter::error::KemError;
|
||||
use clatter::error::KemResult;
|
||||
use clatter::traits::CryptoComponent;
|
||||
use clatter::traits::Kem;
|
||||
use clatter::traits::Rng;
|
||||
|
||||
pub(super) const PUBLIC_KEY_LEN: usize = 1184;
|
||||
const SECRET_KEY_LEN: usize = 2400;
|
||||
const CIPHERTEXT_LEN: usize = 1088;
|
||||
const SHARED_SECRET_LEN: usize = 32;
|
||||
|
||||
/// ML-KEM-768 implementation backed by AWS-LC through `aws-lc-rs`.
|
||||
#[derive(Clone)]
|
||||
pub(super) struct AwsLcMlKem768;
|
||||
|
||||
impl CryptoComponent for AwsLcMlKem768 {
|
||||
fn name() -> &'static str {
|
||||
"MLKEM768"
|
||||
}
|
||||
}
|
||||
|
||||
impl Kem for AwsLcMlKem768 {
|
||||
type SecretKey = SensitiveByteArray<HeapArray<SECRET_KEY_LEN>>;
|
||||
type PubKey = HeapArray<PUBLIC_KEY_LEN>;
|
||||
type Ct = HeapArray<CIPHERTEXT_LEN>;
|
||||
type Ss = SensitiveByteArray<[u8; SHARED_SECRET_LEN]>;
|
||||
|
||||
fn genkey_rng<R: Rng>(_rng: &mut R) -> KemResult<KeyPair<Self::PubKey, Self::SecretKey>> {
|
||||
// AWS-LC owns ML-KEM key-generation randomness internally, so
|
||||
// Clatter's injectable RNG cannot be plumbed through this provider.
|
||||
let decapsulation_key =
|
||||
DecapsulationKey::generate(&ML_KEM_768).map_err(|_| KemError::KeyGeneration)?;
|
||||
let encapsulation_key = decapsulation_key
|
||||
.encapsulation_key()
|
||||
.map_err(|_| KemError::KeyGeneration)?;
|
||||
let public = encapsulation_key
|
||||
.key_bytes()
|
||||
.map_err(|_| KemError::KeyGeneration)?;
|
||||
let secret = decapsulation_key
|
||||
.key_bytes()
|
||||
.map_err(|_| KemError::KeyGeneration)?;
|
||||
|
||||
Ok(KeyPair {
|
||||
public: Self::PubKey::from_slice(public.as_ref()),
|
||||
secret: Self::SecretKey::from_slice(secret.as_ref()),
|
||||
})
|
||||
}
|
||||
|
||||
fn encapsulate<R: Rng>(pk: &[u8], _rng: &mut R) -> KemResult<(Self::Ct, Self::Ss)> {
|
||||
let encapsulation_key =
|
||||
EncapsulationKey::new(&ML_KEM_768, pk).map_err(|_| KemError::Input)?;
|
||||
let (ciphertext, shared_secret) = encapsulation_key
|
||||
.encapsulate()
|
||||
.map_err(|_| KemError::Encapsulation)?;
|
||||
|
||||
Ok((
|
||||
Self::Ct::from_slice(ciphertext.as_ref()),
|
||||
Self::Ss::from_slice(shared_secret.as_ref()),
|
||||
))
|
||||
}
|
||||
|
||||
fn decapsulate(ct: &[u8], sk: &[u8]) -> KemResult<Self::Ss> {
|
||||
// Reject the length before constructing AWS-LC's ciphertext wrapper.
|
||||
// This keeps malformed wire input classified as an input error rather
|
||||
// than relying on provider-specific decapsulation behavior.
|
||||
if ct.len() != CIPHERTEXT_LEN {
|
||||
return Err(KemError::Input);
|
||||
}
|
||||
|
||||
let decapsulation_key =
|
||||
DecapsulationKey::new(&ML_KEM_768, sk).map_err(|_| KemError::Input)?;
|
||||
let shared_secret = decapsulation_key
|
||||
.decapsulate(Ciphertext::from(ct))
|
||||
.map_err(|_| KemError::Decapsulation)?;
|
||||
|
||||
Ok(Self::Ss::from_slice(shared_secret.as_ref()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "aws_lc_ml_kem_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,31 @@
|
||||
use clatter::bytearray::ByteArray;
|
||||
use clatter::traits::Kem;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
use super::AwsLcMlKem768;
|
||||
|
||||
#[test]
|
||||
fn kem_roundtrip() {
|
||||
let keypair = AwsLcMlKem768::genkey().expect("generate keypair");
|
||||
let mut rng = clatter::crypto::rng::DefaultRng;
|
||||
let (ciphertext, encapsulated_secret) =
|
||||
AwsLcMlKem768::encapsulate(keypair.public.as_slice(), &mut rng).expect("encapsulate");
|
||||
let decapsulated_secret =
|
||||
AwsLcMlKem768::decapsulate(ciphertext.as_slice(), keypair.secret.as_slice())
|
||||
.expect("decapsulate");
|
||||
|
||||
assert_eq!(
|
||||
encapsulated_secret.as_slice(),
|
||||
decapsulated_secret.as_slice()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decapsulate_rejects_wrong_ciphertext_length() {
|
||||
let keypair = AwsLcMlKem768::genkey().expect("generate keypair");
|
||||
|
||||
let error = AwsLcMlKem768::decapsulate(&[], keypair.secret.as_slice())
|
||||
.expect_err("empty ciphertext should be rejected");
|
||||
|
||||
assert!(matches!(error, clatter::error::KemError::Input));
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
mod aws_lc_ml_kem;
|
||||
mod client;
|
||||
mod client_api;
|
||||
mod client_transport;
|
||||
@@ -11,6 +12,7 @@ mod fs_helper_main;
|
||||
mod fs_sandbox;
|
||||
mod local_file_system;
|
||||
mod local_process;
|
||||
mod noise_channel;
|
||||
mod process;
|
||||
mod process_id;
|
||||
mod protocol;
|
||||
@@ -51,6 +53,9 @@ pub use fs_helper::CODEX_FS_HELPER_ARG1;
|
||||
pub use fs_helper_main::main as run_fs_helper_main;
|
||||
pub use local_file_system::LOCAL_FS;
|
||||
pub use local_file_system::LocalFileSystem;
|
||||
pub use noise_channel::NoiseChannelError;
|
||||
pub use noise_channel::NoiseChannelIdentity;
|
||||
pub use noise_channel::NoiseChannelPublicKey;
|
||||
pub use process::ExecBackend;
|
||||
pub use process::ExecProcess;
|
||||
pub use process::ExecProcessEvent;
|
||||
|
||||
@@ -0,0 +1,379 @@
|
||||
//! Narrow, misuse-resistant wrapper around the Clatter primitives used by the
|
||||
//! remote exec-server relay.
|
||||
//!
|
||||
//! The harness is the Noise hybrid-IK initiator and pins the registry-provided
|
||||
//! exec-server static public key. The exec-server is the responder; after it
|
||||
//! reads the first IK message, Noise exposes the authenticated harness static
|
||||
//! public key for registry authorization. Application data is allowed only
|
||||
//! after both the cryptographic handshake and that external authorization pass.
|
||||
|
||||
use base64::Engine;
|
||||
use base64::engine::general_purpose::STANDARD;
|
||||
use clatter::HybridHandshake;
|
||||
use clatter::HybridHandshakeParams;
|
||||
use clatter::KeyPair;
|
||||
use clatter::bytearray::ByteArray;
|
||||
use clatter::constants::MAX_MESSAGE_LEN;
|
||||
use clatter::crypto::cipher::AesGcm;
|
||||
use clatter::crypto::dh::X25519;
|
||||
use clatter::crypto::hash::Sha256;
|
||||
use clatter::handshakepattern::noise_hybrid_ik;
|
||||
use clatter::traits::Cipher;
|
||||
use clatter::traits::Dh;
|
||||
use clatter::traits::Handshaker;
|
||||
use clatter::traits::Kem;
|
||||
use clatter::transportstate::TransportState;
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
|
||||
use crate::aws_lc_ml_kem::AwsLcMlKem768;
|
||||
use crate::aws_lc_ml_kem::PUBLIC_KEY_LEN as MLKEM768_PUBLIC_KEY_LEN;
|
||||
|
||||
pub const NOISE_CHANNEL_SUITE: &str = "Noise_hybridIK_X25519+MLKEM768_AESGCM_SHA256";
|
||||
|
||||
const X25519_PUBLIC_KEY_LEN: usize = 32;
|
||||
const MAX_TRANSPORT_RECORDS_PER_DIRECTION: u64 = u32::MAX as u64;
|
||||
const PROLOGUE_DOMAIN: &[u8] = b"codex-exec-server-relay-noise/v1";
|
||||
|
||||
type Handshake = HybridHandshake<X25519, AwsLcMlKem768, AwsLcMlKem768, AesGcm, Sha256>;
|
||||
type Transport = TransportState<AesGcm, Sha256>;
|
||||
type DhKeyPair = KeyPair<<X25519 as Dh>::PubKey, <X25519 as Dh>::PrivateKey>;
|
||||
type KemKeyPair = KeyPair<<AwsLcMlKem768 as Kem>::PubKey, <AwsLcMlKem768 as Kem>::SecretKey>;
|
||||
|
||||
/// Public key material for the exec-server Noise-over-relay suite.
|
||||
///
|
||||
/// The suite field is part of the serialized contract. A key from a different
|
||||
/// suite must not be interpreted as compatible merely because one component has
|
||||
/// a familiar byte length.
|
||||
#[derive(Clone, Eq, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct NoiseChannelPublicKey {
|
||||
suite: String,
|
||||
x25519_public_key: String,
|
||||
mlkem768_public_key: String,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for NoiseChannelPublicKey {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("NoiseChannelPublicKey")
|
||||
.field("suite", &self.suite)
|
||||
.field("x25519_public_key", &"<redacted>")
|
||||
.field("mlkem768_public_key", &"<redacted>")
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl NoiseChannelPublicKey {
|
||||
/// Validate suite selection, encoding, and exact key lengths.
|
||||
pub fn validate(&self) -> Result<(), NoiseChannelError> {
|
||||
let _ = self.decode()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn from_keypairs(dh: &DhKeyPair, kem: &KemKeyPair) -> Self {
|
||||
Self {
|
||||
suite: NOISE_CHANNEL_SUITE.to_string(),
|
||||
x25519_public_key: STANDARD.encode(dh.public),
|
||||
mlkem768_public_key: STANDARD.encode(kem.public.as_slice()),
|
||||
}
|
||||
}
|
||||
|
||||
fn from_raw(dh: &<X25519 as Dh>::PubKey, kem: &<AwsLcMlKem768 as Kem>::PubKey) -> Self {
|
||||
Self {
|
||||
suite: NOISE_CHANNEL_SUITE.to_string(),
|
||||
x25519_public_key: STANDARD.encode(dh),
|
||||
mlkem768_public_key: STANDARD.encode(kem.as_slice()),
|
||||
}
|
||||
}
|
||||
|
||||
fn decode(
|
||||
&self,
|
||||
) -> Result<(<X25519 as Dh>::PubKey, <AwsLcMlKem768 as Kem>::PubKey), NoiseChannelError> {
|
||||
if self.suite != NOISE_CHANNEL_SUITE {
|
||||
return Err(NoiseChannelError::InvalidPublicKey(
|
||||
"unsupported Noise channel suite",
|
||||
));
|
||||
}
|
||||
let dh = STANDARD
|
||||
.decode(&self.x25519_public_key)
|
||||
.map_err(|_| NoiseChannelError::InvalidPublicKey("invalid X25519 public key"))?;
|
||||
let dh: [u8; X25519_PUBLIC_KEY_LEN] = dh
|
||||
.try_into()
|
||||
.map_err(|_| NoiseChannelError::InvalidPublicKey("invalid X25519 public key length"))?;
|
||||
let kem = STANDARD
|
||||
.decode(&self.mlkem768_public_key)
|
||||
.map_err(|_| NoiseChannelError::InvalidPublicKey("invalid ML-KEM-768 public key"))?;
|
||||
if kem.len() != MLKEM768_PUBLIC_KEY_LEN {
|
||||
return Err(NoiseChannelError::InvalidPublicKey(
|
||||
"invalid ML-KEM-768 public key length",
|
||||
));
|
||||
}
|
||||
|
||||
Ok((
|
||||
dh,
|
||||
<AwsLcMlKem768 as Kem>::PubKey::from_slice(kem.as_slice()),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// Endpoint-local static identity for the exec-server Noise-over-relay suite.
|
||||
///
|
||||
/// Private components never cross the process boundary. Cloning is used only to
|
||||
/// construct Clatter handshake state for reconnects within the same process.
|
||||
#[derive(Clone)]
|
||||
pub struct NoiseChannelIdentity {
|
||||
dh: DhKeyPair,
|
||||
kem: KemKeyPair,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for NoiseChannelIdentity {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("NoiseChannelIdentity")
|
||||
.field("public_key", &self.public_key())
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl NoiseChannelIdentity {
|
||||
/// Generate independent classical and post-quantum static keypairs.
|
||||
pub fn generate() -> Result<Self, NoiseChannelError> {
|
||||
let dh = X25519::genkey()
|
||||
.map_err(|error| NoiseChannelError::KeyGeneration(error.to_string()))?;
|
||||
let kem = AwsLcMlKem768::genkey()
|
||||
.map_err(|error| NoiseChannelError::KeyGeneration(error.to_string()))?;
|
||||
Ok(Self { dh, kem })
|
||||
}
|
||||
|
||||
/// Return the distributable public half of this endpoint identity.
|
||||
pub fn public_key(&self) -> NoiseChannelPublicKey {
|
||||
NoiseChannelPublicKey::from_keypairs(&self.dh, &self.kem)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct InitiatorHandshake {
|
||||
handshake: Handshake,
|
||||
}
|
||||
|
||||
impl InitiatorHandshake {
|
||||
/// Start hybrid IK while pinning the expected responder static key.
|
||||
///
|
||||
/// `payload` is encrypted by the first IK message. The relay uses it for
|
||||
/// the short-lived registry authorization naming this harness public key.
|
||||
pub(crate) fn start(
|
||||
identity: &NoiseChannelIdentity,
|
||||
responder_public_key: &NoiseChannelPublicKey,
|
||||
prologue: &[u8],
|
||||
payload: &[u8],
|
||||
) -> Result<(Self, Vec<u8>), NoiseChannelError> {
|
||||
let (responder_dh, responder_kem) = responder_public_key.decode()?;
|
||||
|
||||
// IK authenticates both static identities. Supplying both responder
|
||||
// components here is what makes a misrouted or impersonating
|
||||
// exec-server fail before any JSON-RPC plaintext is released.
|
||||
let params = HybridHandshakeParams::new(noise_hybrid_ik(), true)
|
||||
.with_prologue(prologue)
|
||||
.with_s(identity.dh.clone())
|
||||
.with_s_kem(identity.kem.clone())
|
||||
.with_rs(responder_dh)
|
||||
.with_rs_kem(responder_kem);
|
||||
let mut handshake = Handshake::new(params)?;
|
||||
let mut output = [0u8; MAX_MESSAGE_LEN];
|
||||
let output_len = handshake.write_message(payload, &mut output)?;
|
||||
Ok((Self { handshake }, output[..output_len].to_vec()))
|
||||
}
|
||||
|
||||
/// Consume the responder message and enter transport mode.
|
||||
///
|
||||
/// The responder message carries no application payload in v1. Rejecting
|
||||
/// one keeps future protocol additions from becoming an implicit channel.
|
||||
pub(crate) fn finish(mut self, response: &[u8]) -> Result<NoiseTransport, NoiseChannelError> {
|
||||
ensure_noise_frame_len(response.len(), "handshake response is too large")?;
|
||||
let mut payload = [0u8; MAX_MESSAGE_LEN];
|
||||
let payload_len = self.handshake.read_message(response, &mut payload)?;
|
||||
if payload_len != 0 {
|
||||
return Err(NoiseChannelError::InvalidMessage(
|
||||
"handshake response payload must be empty",
|
||||
));
|
||||
}
|
||||
Ok(NoiseTransport {
|
||||
transport: self.handshake.finalize()?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct PendingResponderHandshake {
|
||||
handshake: Handshake,
|
||||
initiator_public_key: NoiseChannelPublicKey,
|
||||
payload: Vec<u8>,
|
||||
}
|
||||
|
||||
impl PendingResponderHandshake {
|
||||
/// Authenticate and parse the first IK message without completing it.
|
||||
///
|
||||
/// This split is intentional: callers must authorize `initiator_public_key`
|
||||
/// with the registry before calling [`Self::complete`].
|
||||
pub(crate) fn read_request(
|
||||
identity: &NoiseChannelIdentity,
|
||||
prologue: &[u8],
|
||||
request: &[u8],
|
||||
) -> Result<Self, NoiseChannelError> {
|
||||
ensure_noise_frame_len(request.len(), "handshake request is too large")?;
|
||||
let params = HybridHandshakeParams::new(noise_hybrid_ik(), false)
|
||||
.with_prologue(prologue)
|
||||
.with_s(identity.dh.clone())
|
||||
.with_s_kem(identity.kem.clone());
|
||||
let mut handshake = Handshake::new(params)?;
|
||||
let mut payload = [0u8; MAX_MESSAGE_LEN];
|
||||
let payload_len = handshake.read_message(request, &mut payload)?;
|
||||
// Clatter exposes the initiator static key only after the first IK
|
||||
// message authenticates and decrypts successfully.
|
||||
let remote = handshake
|
||||
.get_remote_static()
|
||||
.ok_or(NoiseChannelError::InvalidMessage(
|
||||
"handshake request is missing initiator static key",
|
||||
))?;
|
||||
Ok(Self {
|
||||
handshake,
|
||||
initiator_public_key: NoiseChannelPublicKey::from_raw(remote.dh(), remote.kem()),
|
||||
payload: payload[..payload_len].to_vec(),
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn initiator_public_key(&self) -> &NoiseChannelPublicKey {
|
||||
&self.initiator_public_key
|
||||
}
|
||||
|
||||
/// Move the authenticated first-message payload out of pending state.
|
||||
///
|
||||
/// The v1 payload is a short-lived registry authorization and is not
|
||||
/// needed to complete the handshake. Moving it avoids retaining a second
|
||||
/// copy while external authorization is in flight.
|
||||
pub(crate) fn take_payload(&mut self) -> Vec<u8> {
|
||||
std::mem::take(&mut self.payload)
|
||||
}
|
||||
|
||||
/// Finish the responder handshake after external harness authorization.
|
||||
pub(crate) fn complete(mut self) -> Result<(NoiseTransport, Vec<u8>), NoiseChannelError> {
|
||||
let mut response = [0u8; MAX_MESSAGE_LEN];
|
||||
let response_len = self.handshake.write_message(&[], &mut response)?;
|
||||
Ok((
|
||||
NoiseTransport {
|
||||
transport: self.handshake.finalize()?,
|
||||
},
|
||||
response[..response_len].to_vec(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct NoiseTransport {
|
||||
transport: Transport,
|
||||
}
|
||||
|
||||
impl NoiseTransport {
|
||||
/// Encrypt exactly one ordered transport record.
|
||||
///
|
||||
/// The caller owns relay sequence assignment and must never encrypt the
|
||||
/// same logical record twice under different transport nonces.
|
||||
pub(crate) fn encrypt(&mut self, plaintext: &[u8]) -> Result<Vec<u8>, NoiseChannelError> {
|
||||
if self.transport.sending_nonce() >= MAX_TRANSPORT_RECORDS_PER_DIRECTION {
|
||||
return Err(NoiseChannelError::InvalidState(
|
||||
"transport record nonce exhausted",
|
||||
));
|
||||
}
|
||||
let frame_len = plaintext.len().checked_add(AesGcm::tag_len()).ok_or(
|
||||
NoiseChannelError::InvalidMessage("transport plaintext is too large"),
|
||||
)?;
|
||||
ensure_noise_frame_len(frame_len, "transport plaintext is too large")?;
|
||||
Ok(self.transport.send_vec(plaintext)?)
|
||||
}
|
||||
|
||||
/// Decrypt exactly the next ordered transport record.
|
||||
///
|
||||
/// Clatter advances the receiving nonce during this call, so callers must
|
||||
/// reorder and deduplicate relay frames before invoking it.
|
||||
pub(crate) fn decrypt(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>, NoiseChannelError> {
|
||||
if self.transport.receiving_nonce() >= MAX_TRANSPORT_RECORDS_PER_DIRECTION {
|
||||
return Err(NoiseChannelError::InvalidState(
|
||||
"transport record nonce exhausted",
|
||||
));
|
||||
}
|
||||
if ciphertext.len() < AesGcm::tag_len() {
|
||||
return Err(NoiseChannelError::InvalidMessage(
|
||||
"transport ciphertext is too short",
|
||||
));
|
||||
}
|
||||
ensure_noise_frame_len(ciphertext.len(), "transport ciphertext is too large")?;
|
||||
Ok(self.transport.receive_vec(ciphertext)?)
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the transcript prologue that binds cryptographic identity to routing.
|
||||
///
|
||||
/// A handshake captured from one environment, exec-server registration, or
|
||||
/// relay stream cannot be replayed into another because every participant must
|
||||
/// construct the same prologue before the first Noise message is processed.
|
||||
pub(crate) fn noise_channel_prologue(
|
||||
environment_id: &str,
|
||||
executor_registration_id: &str,
|
||||
stream_id: &str,
|
||||
) -> Result<Vec<u8>, NoiseChannelError> {
|
||||
let mut prologue = Vec::new();
|
||||
append_prologue_part(&mut prologue, PROLOGUE_DOMAIN)?;
|
||||
append_prologue_part(&mut prologue, environment_id.as_bytes())?;
|
||||
append_prologue_part(&mut prologue, executor_registration_id.as_bytes())?;
|
||||
append_prologue_part(&mut prologue, stream_id.as_bytes())?;
|
||||
Ok(prologue)
|
||||
}
|
||||
|
||||
fn append_prologue_part(prologue: &mut Vec<u8>, part: &[u8]) -> Result<(), NoiseChannelError> {
|
||||
// Length prefixes make component boundaries unambiguous. Raw concatenation
|
||||
// would allow different identifier tuples to produce the same prologue.
|
||||
let len = u32::try_from(part.len()).map_err(|_| {
|
||||
NoiseChannelError::InvalidMessage("Noise channel prologue part is too large")
|
||||
})?;
|
||||
prologue.extend_from_slice(&len.to_be_bytes());
|
||||
prologue.extend_from_slice(part);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn ensure_noise_frame_len(
|
||||
frame_len: usize,
|
||||
message: &'static str,
|
||||
) -> Result<(), NoiseChannelError> {
|
||||
if frame_len > MAX_MESSAGE_LEN {
|
||||
return Err(NoiseChannelError::InvalidMessage(message));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum NoiseChannelError {
|
||||
#[error("Noise channel key generation failed: {0}")]
|
||||
KeyGeneration(String),
|
||||
#[error("invalid Noise channel public key: {0}")]
|
||||
InvalidPublicKey(&'static str),
|
||||
#[error("invalid Noise channel state: {0}")]
|
||||
InvalidState(&'static str),
|
||||
#[error("invalid Noise channel message: {0}")]
|
||||
InvalidMessage(&'static str),
|
||||
#[error("Noise channel handshake failed: {0}")]
|
||||
Handshake(String),
|
||||
#[error("Noise channel transport failed: {0}")]
|
||||
Transport(String),
|
||||
}
|
||||
|
||||
impl From<clatter::error::HandshakeError> for NoiseChannelError {
|
||||
fn from(error: clatter::error::HandshakeError) -> Self {
|
||||
Self::Handshake(error.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<clatter::error::TransportError> for NoiseChannelError {
|
||||
fn from(error: clatter::error::TransportError) -> Self {
|
||||
Self::Transport(error.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "noise_channel_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,252 @@
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
use super::InitiatorHandshake;
|
||||
use super::MAX_TRANSPORT_RECORDS_PER_DIRECTION;
|
||||
use super::NOISE_CHANNEL_SUITE;
|
||||
use super::NoiseChannelError;
|
||||
use super::NoiseChannelIdentity;
|
||||
use super::NoiseChannelPublicKey;
|
||||
use super::PendingResponderHandshake;
|
||||
use super::noise_channel_prologue;
|
||||
|
||||
#[test]
|
||||
fn hybrid_ik_roundtrip_authenticates_both_endpoints() {
|
||||
let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity");
|
||||
let responder = NoiseChannelIdentity::generate().expect("generate responder identity");
|
||||
let prologue =
|
||||
noise_channel_prologue("env-1", "registration-1", "stream-1").expect("build prologue");
|
||||
let authorization = b"harness-key-authorization";
|
||||
|
||||
let (initiator_handshake, request) = InitiatorHandshake::start(
|
||||
&initiator,
|
||||
&responder.public_key(),
|
||||
&prologue,
|
||||
authorization,
|
||||
)
|
||||
.expect("start initiator handshake");
|
||||
let mut responder_handshake =
|
||||
PendingResponderHandshake::read_request(&responder, &prologue, &request)
|
||||
.expect("read responder handshake");
|
||||
|
||||
assert_eq!(
|
||||
responder_handshake.initiator_public_key(),
|
||||
&initiator.public_key()
|
||||
);
|
||||
assert_eq!(responder_handshake.take_payload(), authorization);
|
||||
|
||||
let (mut responder_transport, response) = responder_handshake
|
||||
.complete()
|
||||
.expect("complete responder handshake");
|
||||
let mut initiator_transport = initiator_handshake
|
||||
.finish(&response)
|
||||
.expect("complete initiator handshake");
|
||||
|
||||
let request_ciphertext = initiator_transport
|
||||
.encrypt(b"request")
|
||||
.expect("encrypt request");
|
||||
assert_ne!(request_ciphertext, b"request");
|
||||
assert_eq!(
|
||||
responder_transport
|
||||
.decrypt(&request_ciphertext)
|
||||
.expect("decrypt request"),
|
||||
b"request"
|
||||
);
|
||||
|
||||
let response_ciphertext = responder_transport
|
||||
.encrypt(b"response")
|
||||
.expect("encrypt response");
|
||||
assert_ne!(response_ciphertext, b"response");
|
||||
assert_eq!(
|
||||
initiator_transport
|
||||
.decrypt(&response_ciphertext)
|
||||
.expect("decrypt response"),
|
||||
b"response"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn initiator_rejects_wrong_responder_key() {
|
||||
let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity");
|
||||
let expected_responder = NoiseChannelIdentity::generate().expect("generate expected identity");
|
||||
let actual_responder = NoiseChannelIdentity::generate().expect("generate actual identity");
|
||||
let prologue =
|
||||
noise_channel_prologue("env-1", "registration-1", "stream-1").expect("build prologue");
|
||||
|
||||
let (_initiator_handshake, request) = InitiatorHandshake::start(
|
||||
&initiator,
|
||||
&expected_responder.public_key(),
|
||||
&prologue,
|
||||
b"authorization",
|
||||
)
|
||||
.expect("start initiator handshake");
|
||||
|
||||
assert!(
|
||||
PendingResponderHandshake::read_request(&actual_responder, &prologue, &request).is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn responder_rejects_mismatched_prologue() {
|
||||
let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity");
|
||||
let responder = NoiseChannelIdentity::generate().expect("generate responder identity");
|
||||
let initiator_prologue =
|
||||
noise_channel_prologue("env-1", "registration-1", "stream-1").expect("build prologue");
|
||||
let responder_prologue =
|
||||
noise_channel_prologue("env-1", "registration-1", "stream-2").expect("build prologue");
|
||||
let (_initiator_handshake, request) = InitiatorHandshake::start(
|
||||
&initiator,
|
||||
&responder.public_key(),
|
||||
&initiator_prologue,
|
||||
b"authorization",
|
||||
)
|
||||
.expect("start initiator handshake");
|
||||
|
||||
assert!(
|
||||
PendingResponderHandshake::read_request(&responder, &responder_prologue, &request).is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prologue_encoding_is_stable_and_unambiguous() {
|
||||
let prologue =
|
||||
noise_channel_prologue("env-1", "registration-1", "stream-1").expect("build prologue");
|
||||
|
||||
assert_eq!(
|
||||
prologue,
|
||||
b"\x00\x00\x00\x20codex-exec-server-relay-noise/v1\
|
||||
\x00\x00\x00\x05env-1\
|
||||
\x00\x00\x00\x0eregistration-1\
|
||||
\x00\x00\x00\x08stream-1"
|
||||
.to_vec()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transport_rejects_tampered_ciphertext() {
|
||||
let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity");
|
||||
let responder = NoiseChannelIdentity::generate().expect("generate responder identity");
|
||||
let prologue =
|
||||
noise_channel_prologue("env-1", "registration-1", "stream-1").expect("build prologue");
|
||||
let (initiator_handshake, request) = InitiatorHandshake::start(
|
||||
&initiator,
|
||||
&responder.public_key(),
|
||||
&prologue,
|
||||
b"authorization",
|
||||
)
|
||||
.expect("start initiator handshake");
|
||||
let responder_handshake =
|
||||
PendingResponderHandshake::read_request(&responder, &prologue, &request)
|
||||
.expect("read responder handshake");
|
||||
let (mut responder_transport, response) = responder_handshake
|
||||
.complete()
|
||||
.expect("complete responder handshake");
|
||||
let mut initiator_transport = initiator_handshake
|
||||
.finish(&response)
|
||||
.expect("complete initiator handshake");
|
||||
let mut ciphertext = initiator_transport
|
||||
.encrypt(b"request")
|
||||
.expect("encrypt request");
|
||||
ciphertext[0] ^= 1;
|
||||
|
||||
assert!(responder_transport.decrypt(&ciphertext).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transport_rejects_exhausted_receiving_nonce_before_decryption() {
|
||||
let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity");
|
||||
let responder = NoiseChannelIdentity::generate().expect("generate responder identity");
|
||||
let prologue =
|
||||
noise_channel_prologue("env-1", "registration-1", "stream-1").expect("build prologue");
|
||||
let (initiator_handshake, request) = InitiatorHandshake::start(
|
||||
&initiator,
|
||||
&responder.public_key(),
|
||||
&prologue,
|
||||
b"authorization",
|
||||
)
|
||||
.expect("start initiator handshake");
|
||||
let responder_handshake =
|
||||
PendingResponderHandshake::read_request(&responder, &prologue, &request)
|
||||
.expect("read responder handshake");
|
||||
let (mut responder_transport, response) = responder_handshake
|
||||
.complete()
|
||||
.expect("complete responder handshake");
|
||||
let mut initiator_transport = initiator_handshake
|
||||
.finish(&response)
|
||||
.expect("complete initiator handshake");
|
||||
let ciphertext = initiator_transport
|
||||
.encrypt(b"request")
|
||||
.expect("encrypt request");
|
||||
responder_transport
|
||||
.transport
|
||||
.set_receiving_nonce(MAX_TRANSPORT_RECORDS_PER_DIRECTION);
|
||||
|
||||
assert!(matches!(
|
||||
responder_transport.decrypt(&ciphertext),
|
||||
Err(NoiseChannelError::InvalidState(
|
||||
"transport record nonce exhausted"
|
||||
))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transport_rejects_replayed_ciphertext() {
|
||||
let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity");
|
||||
let responder = NoiseChannelIdentity::generate().expect("generate responder identity");
|
||||
let prologue =
|
||||
noise_channel_prologue("env-1", "registration-1", "stream-1").expect("build prologue");
|
||||
let (initiator_handshake, request) = InitiatorHandshake::start(
|
||||
&initiator,
|
||||
&responder.public_key(),
|
||||
&prologue,
|
||||
b"authorization",
|
||||
)
|
||||
.expect("start initiator handshake");
|
||||
let responder_handshake =
|
||||
PendingResponderHandshake::read_request(&responder, &prologue, &request)
|
||||
.expect("read responder handshake");
|
||||
let (mut responder_transport, response) = responder_handshake
|
||||
.complete()
|
||||
.expect("complete responder handshake");
|
||||
let mut initiator_transport = initiator_handshake
|
||||
.finish(&response)
|
||||
.expect("complete initiator handshake");
|
||||
let ciphertext = initiator_transport
|
||||
.encrypt(b"request")
|
||||
.expect("encrypt request");
|
||||
|
||||
assert_eq!(
|
||||
responder_transport
|
||||
.decrypt(&ciphertext)
|
||||
.expect("decrypt request"),
|
||||
b"request"
|
||||
);
|
||||
assert!(matches!(
|
||||
responder_transport.decrypt(&ciphertext),
|
||||
Err(NoiseChannelError::Transport(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn public_key_validation_rejects_unknown_suite() {
|
||||
let key = NoiseChannelIdentity::generate()
|
||||
.expect("generate identity")
|
||||
.public_key();
|
||||
let json = serde_json::to_value(key).expect("serialize key");
|
||||
let mut object = json.as_object().expect("key object").clone();
|
||||
object.insert("suite".to_string(), serde_json::json!("unknown"));
|
||||
let key: NoiseChannelPublicKey =
|
||||
serde_json::from_value(serde_json::Value::Object(object)).expect("deserialize key");
|
||||
|
||||
assert!(key.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn public_key_serializes_with_expected_suite() {
|
||||
let key = NoiseChannelIdentity::generate()
|
||||
.expect("generate identity")
|
||||
.public_key();
|
||||
|
||||
let json = serde_json::to_value(key).expect("serialize key");
|
||||
|
||||
assert_eq!(json["suite"], NOISE_CHANNEL_SUITE);
|
||||
}
|
||||
@@ -14,6 +14,7 @@ message RelayMessageFrame {
|
||||
RelayResume resume = 7;
|
||||
RelayReset reset = 8;
|
||||
RelayHeartbeat heartbeat = 9;
|
||||
RelayHandshake handshake = 10;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -35,3 +36,7 @@ message RelayReset {
|
||||
}
|
||||
|
||||
message RelayHeartbeat {}
|
||||
|
||||
message RelayHandshake {
|
||||
bytes payload = 1;
|
||||
}
|
||||
|
||||
@@ -9,7 +9,7 @@ pub struct RelayMessageFrame {
|
||||
pub ack: u32,
|
||||
#[prost(uint32, tag = "4")]
|
||||
pub ack_bits: u32,
|
||||
#[prost(oneof = "relay_message_frame::Body", tags = "5, 6, 7, 8, 9")]
|
||||
#[prost(oneof = "relay_message_frame::Body", tags = "5, 6, 7, 8, 9, 10")]
|
||||
pub body: ::core::option::Option<relay_message_frame::Body>,
|
||||
}
|
||||
pub mod relay_message_frame {
|
||||
@@ -25,6 +25,8 @@ pub mod relay_message_frame {
|
||||
Reset(super::RelayReset),
|
||||
#[prost(message, tag = "9")]
|
||||
Heartbeat(super::RelayHeartbeat),
|
||||
#[prost(message, tag = "10")]
|
||||
Handshake(super::RelayHandshake),
|
||||
}
|
||||
}
|
||||
#[derive(Clone, PartialEq, ::prost::Message)]
|
||||
@@ -52,3 +54,8 @@ pub struct RelayReset {
|
||||
}
|
||||
#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
|
||||
pub struct RelayHeartbeat {}
|
||||
#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
|
||||
pub struct RelayHandshake {
|
||||
#[prost(bytes = "vec", tag = "1")]
|
||||
pub payload: ::prost::alloc::vec::Vec<u8>,
|
||||
}
|
||||
|
||||
@@ -22,24 +22,29 @@ use crate::connection::JsonRpcConnection;
|
||||
use crate::connection::JsonRpcConnectionEvent;
|
||||
use crate::connection::JsonRpcTransport;
|
||||
use crate::relay_proto::RelayData;
|
||||
use crate::relay_proto::RelayHandshake;
|
||||
use crate::relay_proto::RelayMessageFrame;
|
||||
use crate::relay_proto::RelayReset;
|
||||
use crate::relay_proto::RelayResume;
|
||||
use crate::relay_proto::relay_message_frame;
|
||||
use crate::server::ConnectionProcessor;
|
||||
|
||||
const RELAY_MESSAGE_FRAME_VERSION: u32 = 1;
|
||||
const MAX_RELAY_RESET_REASON_BYTES: usize = 256;
|
||||
const MAX_RELAY_STREAM_ID_BYTES: usize = 128;
|
||||
|
||||
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
|
||||
enum RelayFrameBodyKind {
|
||||
pub(crate) enum RelayFrameBodyKind {
|
||||
Data,
|
||||
Ack,
|
||||
Resume,
|
||||
Reset,
|
||||
Heartbeat,
|
||||
Handshake,
|
||||
}
|
||||
|
||||
impl RelayMessageFrame {
|
||||
fn data(stream_id: String, seq: u32, payload: Vec<u8>) -> Self {
|
||||
pub(crate) fn data(stream_id: String, seq: u32, payload: Vec<u8>) -> Self {
|
||||
Self {
|
||||
version: RELAY_MESSAGE_FRAME_VERSION,
|
||||
stream_id,
|
||||
@@ -54,7 +59,7 @@ impl RelayMessageFrame {
|
||||
}
|
||||
}
|
||||
|
||||
fn resume(stream_id: String) -> Self {
|
||||
pub(crate) fn resume(stream_id: String) -> Self {
|
||||
Self {
|
||||
version: RELAY_MESSAGE_FRAME_VERSION,
|
||||
stream_id,
|
||||
@@ -66,18 +71,55 @@ impl RelayMessageFrame {
|
||||
}
|
||||
}
|
||||
|
||||
fn validate(&self) -> Result<RelayFrameBodyKind, ExecServerError> {
|
||||
pub(crate) fn handshake(stream_id: String, payload: Vec<u8>) -> Self {
|
||||
Self {
|
||||
version: RELAY_MESSAGE_FRAME_VERSION,
|
||||
stream_id,
|
||||
ack: 0,
|
||||
ack_bits: 0,
|
||||
body: Some(relay_message_frame::Body::Handshake(RelayHandshake {
|
||||
payload,
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn reset(stream_id: String, reason: String) -> Self {
|
||||
Self {
|
||||
version: RELAY_MESSAGE_FRAME_VERSION,
|
||||
stream_id,
|
||||
ack: 0,
|
||||
ack_bits: 0,
|
||||
body: Some(relay_message_frame::Body::Reset(RelayReset { reason })),
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate cleartext routing metadata before a frame reaches either the
|
||||
/// direct JSON-RPC parser or the Noise relay state machine.
|
||||
///
|
||||
/// The encrypted path intentionally leaves this metadata visible to
|
||||
/// rendezvous, so it must be canonical and tightly bounded everywhere.
|
||||
pub(crate) fn validate(&self) -> Result<RelayFrameBodyKind, ExecServerError> {
|
||||
if self.version != RELAY_MESSAGE_FRAME_VERSION {
|
||||
return Err(ExecServerError::Protocol(format!(
|
||||
"unsupported relay message frame version {}",
|
||||
self.version
|
||||
)));
|
||||
}
|
||||
if self.stream_id.trim().is_empty() {
|
||||
if self.stream_id.is_empty() {
|
||||
return Err(ExecServerError::Protocol(
|
||||
"relay message frame is missing stream_id".to_string(),
|
||||
));
|
||||
}
|
||||
if self.stream_id.trim() != self.stream_id {
|
||||
return Err(ExecServerError::Protocol(
|
||||
"relay message frame stream_id is not canonical".to_string(),
|
||||
));
|
||||
}
|
||||
if self.stream_id.len() > MAX_RELAY_STREAM_ID_BYTES {
|
||||
return Err(ExecServerError::Protocol(
|
||||
"relay message frame stream_id is too long".to_string(),
|
||||
));
|
||||
}
|
||||
match self.body.as_ref() {
|
||||
Some(relay_message_frame::Body::Data(data)) => {
|
||||
if data.segment_index != 0 || data.segment_count != 1 || data.payload.is_empty() {
|
||||
@@ -90,35 +132,64 @@ impl RelayMessageFrame {
|
||||
Some(relay_message_frame::Body::AckFrame(_)) => Ok(RelayFrameBodyKind::Ack),
|
||||
Some(relay_message_frame::Body::Resume(_)) => Ok(RelayFrameBodyKind::Resume),
|
||||
Some(relay_message_frame::Body::Reset(reset)) => {
|
||||
if reset.reason.is_empty() {
|
||||
if reset.reason.is_empty() || reset.reason.len() > MAX_RELAY_RESET_REASON_BYTES {
|
||||
return Err(ExecServerError::Protocol(
|
||||
"relay reset message frame is missing reason".to_string(),
|
||||
"relay reset message frame has invalid reason".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(RelayFrameBodyKind::Reset)
|
||||
}
|
||||
Some(relay_message_frame::Body::Heartbeat(_)) => Ok(RelayFrameBodyKind::Heartbeat),
|
||||
Some(relay_message_frame::Body::Handshake(handshake)) => {
|
||||
if handshake.payload.is_empty() {
|
||||
return Err(ExecServerError::Protocol(
|
||||
"relay handshake message frame is missing payload".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(RelayFrameBodyKind::Handshake)
|
||||
}
|
||||
None => Err(ExecServerError::Protocol(
|
||||
"relay message frame is missing body".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn into_jsonrpc_message(self) -> Result<JSONRPCMessage, ExecServerError> {
|
||||
pub(crate) fn into_data(self) -> Result<RelayData, ExecServerError> {
|
||||
let kind = self.validate()?;
|
||||
if kind != RelayFrameBodyKind::Data {
|
||||
return Err(ExecServerError::Protocol(
|
||||
"expected relay data message frame".to_string(),
|
||||
));
|
||||
}
|
||||
let payload = match self.body {
|
||||
Some(relay_message_frame::Body::Data(data)) => data.payload,
|
||||
_ => Vec::new(),
|
||||
};
|
||||
match self.body {
|
||||
Some(relay_message_frame::Body::Data(data)) => Ok(data),
|
||||
_ => Err(ExecServerError::Protocol(
|
||||
"expected relay data message frame".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn into_jsonrpc_message(self) -> Result<JSONRPCMessage, ExecServerError> {
|
||||
let payload = self.into_data()?.payload;
|
||||
serde_json::from_slice(&payload).map_err(ExecServerError::Json)
|
||||
}
|
||||
|
||||
fn into_reset_reason(self) -> Option<String> {
|
||||
pub(crate) fn into_handshake_payload(self) -> Result<Vec<u8>, ExecServerError> {
|
||||
let kind = self.validate()?;
|
||||
if kind != RelayFrameBodyKind::Handshake {
|
||||
return Err(ExecServerError::Protocol(
|
||||
"expected relay handshake message frame".to_string(),
|
||||
));
|
||||
}
|
||||
match self.body {
|
||||
Some(relay_message_frame::Body::Handshake(handshake)) => Ok(handshake.payload),
|
||||
_ => Err(ExecServerError::Protocol(
|
||||
"expected relay handshake message frame".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn into_reset_reason(self) -> Option<String> {
|
||||
match self.body {
|
||||
Some(relay_message_frame::Body::Reset(reset)) if !reset.reason.is_empty() => {
|
||||
Some(reset.reason)
|
||||
@@ -128,16 +199,18 @@ impl RelayMessageFrame {
|
||||
}
|
||||
}
|
||||
|
||||
fn encode_relay_message_frame(frame: &RelayMessageFrame) -> Vec<u8> {
|
||||
pub(crate) fn encode_relay_message_frame(frame: &RelayMessageFrame) -> Vec<u8> {
|
||||
frame.encode_to_vec()
|
||||
}
|
||||
|
||||
fn decode_relay_message_frame(payload: &[u8]) -> Result<RelayMessageFrame, ExecServerError> {
|
||||
pub(crate) fn decode_relay_message_frame(
|
||||
payload: &[u8],
|
||||
) -> Result<RelayMessageFrame, ExecServerError> {
|
||||
RelayMessageFrame::decode(payload)
|
||||
.map_err(|err| ExecServerError::Protocol(format!("invalid relay message frame: {err}")))
|
||||
}
|
||||
|
||||
fn jsonrpc_payload(message: &JSONRPCMessage) -> Result<Vec<u8>, ExecServerError> {
|
||||
pub(crate) fn jsonrpc_payload(message: &JSONRPCMessage) -> Result<Vec<u8>, ExecServerError> {
|
||||
serde_json::to_vec(message).map_err(ExecServerError::Json)
|
||||
}
|
||||
|
||||
@@ -253,7 +326,8 @@ where
|
||||
}
|
||||
RelayFrameBodyKind::Ack
|
||||
| RelayFrameBodyKind::Resume
|
||||
| RelayFrameBodyKind::Heartbeat => {}
|
||||
| RelayFrameBodyKind::Heartbeat
|
||||
| RelayFrameBodyKind::Handshake => {}
|
||||
}
|
||||
}
|
||||
Some(Ok(Message::Close(_))) | None => {
|
||||
@@ -294,10 +368,22 @@ where
|
||||
incoming_rx,
|
||||
disconnected_rx,
|
||||
task_handles: vec![websocket_task],
|
||||
transport: JsonRpcTransport::Plain,
|
||||
transport: JsonRpcTransport::External,
|
||||
}
|
||||
}
|
||||
|
||||
/// Runs the backwards-compatible, cleartext multiplexed relay protocol.
|
||||
///
|
||||
/// The physical websocket carries protobuf [`RelayMessageFrame`] values. Each
|
||||
/// frame's `stream_id` selects an independent logical JSON-RPC connection, so a
|
||||
/// single registered exec-server can serve multiple orchestrator sessions.
|
||||
/// This runner intentionally preserves the protocol used before Noise relay
|
||||
/// support was introduced. Callers must select the separate secure runner when
|
||||
/// Noise protection was explicitly negotiated during registration.
|
||||
///
|
||||
/// Relay framing is not an authentication boundary. Every frame is validated
|
||||
/// before its routing metadata or payload is used, and malformed frames are
|
||||
/// dropped without affecting the other virtual streams on the connection.
|
||||
pub(crate) async fn run_multiplexed_environment<S>(
|
||||
stream: WebSocketStream<S>,
|
||||
processor: ConnectionProcessor,
|
||||
@@ -310,6 +396,9 @@ pub(crate) async fn run_multiplexed_environment<S>(
|
||||
|
||||
let mut streams: HashMap<String, VirtualStream> = HashMap::new();
|
||||
loop {
|
||||
// Serialize all logical-stream writes through this task so only one
|
||||
// owner writes to the physical websocket. Reads remain in the same
|
||||
// select loop, which also keeps disconnect handling deterministic.
|
||||
let frame = tokio::select! {
|
||||
maybe_encoded = physical_outgoing_rx.recv() => {
|
||||
let Some(encoded) = maybe_encoded else {
|
||||
@@ -361,6 +450,11 @@ pub(crate) async fn run_multiplexed_environment<S>(
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// A logical connection is created lazily on its first data
|
||||
// frame. The connection processor owns the JSON-RPC lifecycle;
|
||||
// this task only translates between relay and connection
|
||||
// events.
|
||||
let stream = streams.entry(stream_id.clone()).or_insert_with(|| {
|
||||
spawn_virtual_stream(
|
||||
stream_id.clone(),
|
||||
@@ -382,24 +476,34 @@ pub(crate) async fn run_multiplexed_environment<S>(
|
||||
stream.disconnect(frame.into_reset_reason()).await;
|
||||
}
|
||||
}
|
||||
// These control frames are meaningful to other relay protocol
|
||||
// variants. The legacy protocol has no resume or handshake state,
|
||||
// so ignoring them preserves its existing wire behavior.
|
||||
RelayFrameBodyKind::Ack
|
||||
| RelayFrameBodyKind::Resume
|
||||
| RelayFrameBodyKind::Heartbeat => {}
|
||||
| RelayFrameBodyKind::Heartbeat
|
||||
| RelayFrameBodyKind::Handshake => {}
|
||||
}
|
||||
}
|
||||
|
||||
// A physical disconnect ends every virtual connection. Notify each
|
||||
// processor explicitly so requests do not remain live after the relay
|
||||
// websocket has disappeared.
|
||||
for (_stream_id, stream) in streams {
|
||||
stream.disconnect(/*reason*/ None).await;
|
||||
}
|
||||
drop(physical_outgoing_tx);
|
||||
}
|
||||
|
||||
/// The exec-server-facing half of one logical connection on the shared relay.
|
||||
struct VirtualStream {
|
||||
incoming_tx: mpsc::Sender<JsonRpcConnectionEvent>,
|
||||
disconnected_tx: watch::Sender<bool>,
|
||||
}
|
||||
|
||||
impl VirtualStream {
|
||||
/// Marks the logical connection disconnected and supplies the peer's reset
|
||||
/// reason, when one was provided.
|
||||
async fn disconnect(self, reason: Option<String>) {
|
||||
let _ = self.disconnected_tx.send(true);
|
||||
let _ = self
|
||||
@@ -409,6 +513,8 @@ impl VirtualStream {
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates a logical JSON-RPC connection and forwards its outgoing messages
|
||||
/// back to the physical relay writer as legacy cleartext data frames.
|
||||
fn spawn_virtual_stream(
|
||||
stream_id: String,
|
||||
processor: ConnectionProcessor,
|
||||
@@ -446,7 +552,7 @@ fn spawn_virtual_stream(
|
||||
incoming_rx,
|
||||
disconnected_rx,
|
||||
task_handles: vec![writer_task],
|
||||
transport: JsonRpcTransport::Plain,
|
||||
transport: JsonRpcTransport::External,
|
||||
};
|
||||
tokio::spawn(async move {
|
||||
processor.run_connection(connection).await;
|
||||
@@ -476,6 +582,7 @@ mod tests {
|
||||
use futures::task::AtomicWaker;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::time::timeout;
|
||||
use tokio_tungstenite::WebSocketStream;
|
||||
use tokio_tungstenite::accept_async;
|
||||
use tokio_tungstenite::connect_async;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
|
||||
@@ -2,6 +2,8 @@
|
||||
mod generated;
|
||||
|
||||
pub(crate) use generated::RelayData;
|
||||
pub(crate) use generated::RelayHandshake;
|
||||
pub(crate) use generated::RelayMessageFrame;
|
||||
pub(crate) use generated::RelayReset;
|
||||
pub(crate) use generated::RelayResume;
|
||||
pub(crate) use generated::relay_message_frame;
|
||||
|
||||
Reference in New Issue
Block a user