codex_exec_server/
noise_channel.rs

1//! Noise channel used by the remote exec-server relay.
2//!
3//! The harness initiates hybrid IK and pins the exec-server static key returned
4//! by the registry. The first handshake message lets the exec-server authenticate
5//! the harness static key; the exec-server then asks the registry whether that
6//! key is authorized before completing the handshake.
7//!
8//! "Hybrid" means the session keys include both X25519 and ML-KEM-768 key
9//! agreement. Once the two-message handshake finishes, AES-GCM protects the
10//! ordered transport records carrying JSON-RPC.
11
12use base64::Engine;
13use base64::engine::general_purpose::STANDARD;
14use clatter::HybridHandshake;
15use clatter::HybridHandshakeParams;
16use clatter::KeyPair;
17use clatter::bytearray::ByteArray;
18use clatter::constants::MAX_MESSAGE_LEN;
19use clatter::crypto::cipher::AesGcm;
20use clatter::crypto::dh::X25519;
21use clatter::crypto::hash::Sha256;
22use clatter::crypto::kem::rust_crypto_ml_kem::MlKem768;
23use clatter::handshakepattern::noise_hybrid_ik;
24use clatter::traits::Cipher;
25use clatter::traits::Dh;
26use clatter::traits::Handshaker;
27use clatter::traits::Kem;
28use clatter::transportstate::TransportState;
29use serde::Deserialize;
30use serde::Serialize;
31
32/// Identifies the handshake pattern and algorithms used by this channel.
33pub(crate) const NOISE_CHANNEL_SUITE: &str = "Noise_hybridIK_X25519+MLKEM768_AESGCM_SHA256";
34
35const PROLOGUE_DOMAIN: &[u8] = b"codex-exec-server-relay-noise/v1";
36
37type Handshake = HybridHandshake<X25519, MlKem768, MlKem768, AesGcm, Sha256>;
38type Transport = TransportState<AesGcm, Sha256>;
39type DhKeyPair = KeyPair<<X25519 as Dh>::PubKey, <X25519 as Dh>::PrivateKey>;
40type MlKem768PublicKey = <MlKem768 as Kem>::PubKey;
41type KemKeyPair = KeyPair<<MlKem768 as Kem>::PubKey, <MlKem768 as Kem>::SecretKey>;
42
43/// Public key material for the exec-server Noise suite.
44/// The suite tag prevents keys for another protocol from being accepted just
45/// because their components have the expected lengths.
46#[derive(Clone, Eq, PartialEq, Serialize, Deserialize)]
47#[serde(deny_unknown_fields)]
48pub struct NoiseChannelPublicKey {
49    suite: String,
50    x25519_public_key: String,
51    mlkem768_public_key: String,
52}
53
54impl std::fmt::Debug for NoiseChannelPublicKey {
55    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        f.debug_struct("NoiseChannelPublicKey")
57            .field("suite", &self.suite)
58            .field("x25519_public_key", &"<redacted>")
59            .field("mlkem768_public_key", &"<redacted>")
60            .finish()
61    }
62}
63
64impl NoiseChannelPublicKey {
65    /// Decode registry-provided key material before passing it to Clatter.
66    fn decode(&self) -> Result<(<X25519 as Dh>::PubKey, MlKem768PublicKey), NoiseChannelError> {
67        if self.suite != NOISE_CHANNEL_SUITE {
68            return Err(NoiseChannelError::InvalidPublicKey(
69                "unsupported Noise channel suite",
70            ));
71        }
72        let dh = STANDARD
73            .decode(&self.x25519_public_key)
74            .map_err(|_| NoiseChannelError::InvalidPublicKey("invalid X25519 public key"))?;
75        let dh: <X25519 as Dh>::PubKey = dh
76            .try_into()
77            .map_err(|_| NoiseChannelError::InvalidPublicKey("invalid X25519 public key length"))?;
78        let kem = STANDARD
79            .decode(&self.mlkem768_public_key)
80            .map_err(|_| NoiseChannelError::InvalidPublicKey("invalid ML-KEM-768 public key"))?;
81        if kem.len() != MlKem768PublicKey::LENGTH {
82            return Err(NoiseChannelError::InvalidPublicKey(
83                "invalid ML-KEM-768 public key length",
84            ));
85        }
86
87        Ok((dh, MlKem768PublicKey::from_slice(kem.as_slice())))
88    }
89}
90
91/// Static Noise identity kept for the lifetime of an executor or harness process.
92#[derive(Clone)]
93pub struct NoiseChannelIdentity {
94    dh: DhKeyPair,
95    kem: KemKeyPair,
96}
97
98impl NoiseChannelIdentity {
99    pub fn generate() -> Result<Self, NoiseChannelError> {
100        let dh = X25519::genkey()
101            .map_err(|error| NoiseChannelError::KeyGeneration(error.to_string()))?;
102        let kem = MlKem768::genkey()
103            .map_err(|error| NoiseChannelError::KeyGeneration(error.to_string()))?;
104        Ok(Self { dh, kem })
105    }
106
107    pub fn public_key(&self) -> NoiseChannelPublicKey {
108        NoiseChannelPublicKey {
109            suite: NOISE_CHANNEL_SUITE.to_string(),
110            x25519_public_key: STANDARD.encode(self.dh.public),
111            mlkem768_public_key: STANDARD.encode(self.kem.public.as_slice()),
112        }
113    }
114}
115
116/// Harness-side state between the two hybrid-IK messages.
117/// Consuming it in [`Self::finish`] keeps a handshake tied to one relay stream.
118pub(crate) struct InitiatorHandshake {
119    handshake: Handshake,
120}
121
122impl InitiatorHandshake {
123    /// Start hybrid IK and pin the expected executor key.
124    /// `payload` carries the short-lived registry authorization inside the first
125    /// encrypted handshake message.
126    pub(crate) fn start(
127        identity: &NoiseChannelIdentity,
128        responder_public_key: &NoiseChannelPublicKey,
129        prologue: &[u8],
130        payload: &[u8],
131    ) -> Result<(Self, Vec<u8>), NoiseChannelError> {
132        let (responder_dh, responder_kem) = responder_public_key.decode()?;
133
134        // Both executor key components are pinned before any JSON-RPC is sent.
135        let params = HybridHandshakeParams::new(noise_hybrid_ik(), true)
136            .with_prologue(prologue)
137            .with_s(identity.dh.clone())
138            .with_s_kem(identity.kem.clone())
139            .with_rs(responder_dh)
140            .with_rs_kem(responder_kem);
141        let mut handshake = Handshake::new(params)?;
142        let overhead = handshake.get_next_message_overhead()?;
143        if payload.len() > MAX_MESSAGE_LEN - overhead {
144            return Err(NoiseChannelError::InvalidMessage(
145                "handshake payload is too large",
146            ));
147        }
148        let mut output = vec![0u8; payload.len() + overhead];
149        let output_len = handshake.write_message(payload, &mut output)?;
150        output.truncate(output_len);
151        Ok((Self { handshake }, output))
152    }
153
154    /// Consume the executor response and enter transport mode.
155    /// The v1 response does not carry an application payload.
156    pub(crate) fn finish(mut self, response: &[u8]) -> Result<NoiseTransport, NoiseChannelError> {
157        ensure_noise_frame_len(response.len(), "handshake response is too large")?;
158        let overhead = self.handshake.get_next_message_overhead()?;
159        let mut payload = vec![0u8; response.len().saturating_sub(overhead)];
160        let payload_len = self.handshake.read_message(response, &mut payload)?;
161        if payload_len != 0 {
162            return Err(NoiseChannelError::InvalidMessage(
163                "handshake response payload must be empty",
164            ));
165        }
166        Ok(NoiseTransport {
167            transport: self.handshake.finalize()?,
168        })
169    }
170}
171
172/// Executor-side handshake state while harness authorization is pending.
173/// This is not a usable transport until the registry accepts the authenticated
174/// harness key.
175pub(crate) struct PendingResponderHandshake {
176    handshake: Handshake,
177    pub(crate) initiator_public_key: NoiseChannelPublicKey,
178    pub(crate) payload: Vec<u8>,
179}
180
181impl PendingResponderHandshake {
182    /// Parse the first IK message and recover the authenticated harness key.
183    /// Callers must authorize that key before calling [`Self::complete`].
184    pub(crate) fn read_request(
185        identity: &NoiseChannelIdentity,
186        prologue: &[u8],
187        request: &[u8],
188    ) -> Result<Self, NoiseChannelError> {
189        ensure_noise_frame_len(request.len(), "handshake request is too large")?;
190        let params = HybridHandshakeParams::new(noise_hybrid_ik(), false)
191            .with_prologue(prologue)
192            .with_s(identity.dh.clone())
193            .with_s_kem(identity.kem.clone());
194        let mut handshake = Handshake::new(params)?;
195        let overhead = handshake.get_next_message_overhead()?;
196        let mut payload = vec![0u8; request.len().saturating_sub(overhead)];
197        let payload_len = handshake.read_message(request, &mut payload)?;
198        // Clatter exposes this key only after the first IK message authenticates.
199        let remote = handshake
200            .get_remote_static()
201            .ok_or(NoiseChannelError::InvalidMessage(
202                "handshake request is missing initiator static key",
203            ))?;
204        let initiator_public_key = NoiseChannelPublicKey {
205            suite: NOISE_CHANNEL_SUITE.to_string(),
206            x25519_public_key: STANDARD.encode(remote.dh()),
207            mlkem768_public_key: STANDARD.encode(remote.kem().as_slice()),
208        };
209        payload.truncate(payload_len);
210        Ok(Self {
211            handshake,
212            initiator_public_key,
213            payload,
214        })
215    }
216
217    /// Finish the handshake after the registry authorizes the harness key.
218    pub(crate) fn complete(mut self) -> Result<(NoiseTransport, Vec<u8>), NoiseChannelError> {
219        let overhead = self.handshake.get_next_message_overhead()?;
220        let mut response = vec![0u8; overhead];
221        let response_len = self.handshake.write_message(&[], &mut response)?;
222        response.truncate(response_len);
223        Ok((
224            NoiseTransport {
225                transport: self.handshake.finalize()?,
226            },
227            response,
228        ))
229    }
230}
231
232/// Established channel with independent implicit send and receive nonces.
233/// Relay records must be ordered before decryption, and a logical record must
234/// not be encrypted again for retry.
235pub(crate) struct NoiseTransport {
236    transport: Transport,
237}
238
239impl NoiseTransport {
240    /// Encrypt the next transport record.
241    pub(crate) fn encrypt(&mut self, plaintext: &[u8]) -> Result<Vec<u8>, NoiseChannelError> {
242        let frame_len = plaintext.len().checked_add(AesGcm::tag_len()).ok_or(
243            NoiseChannelError::InvalidMessage("transport plaintext is too large"),
244        )?;
245        ensure_noise_frame_len(frame_len, "transport plaintext is too large")?;
246        Ok(self.transport.send_vec(plaintext)?)
247    }
248
249    /// Decrypt the next ordered transport record.
250    pub(crate) fn decrypt(&mut self, ciphertext: &[u8]) -> Result<Vec<u8>, NoiseChannelError> {
251        if ciphertext.len() < AesGcm::tag_len() {
252            return Err(NoiseChannelError::InvalidMessage(
253                "transport ciphertext is too short",
254            ));
255        }
256        ensure_noise_frame_len(ciphertext.len(), "transport ciphertext is too large")?;
257        Ok(self.transport.receive_vec(ciphertext)?)
258    }
259}
260
261/// Bind the handshake to one environment registration and relay stream.
262/// Both peers include these values in the Noise transcript before processing
263/// the first handshake message.
264pub(crate) fn noise_channel_prologue(
265    environment_id: &str,
266    executor_registration_id: &str,
267    stream_id: &str,
268) -> Vec<u8> {
269    let mut prologue = Vec::new();
270    append_prologue_part(&mut prologue, PROLOGUE_DOMAIN);
271    append_prologue_part(&mut prologue, environment_id.as_bytes());
272    append_prologue_part(&mut prologue, executor_registration_id.as_bytes());
273    append_prologue_part(&mut prologue, stream_id.as_bytes());
274    prologue
275}
276
277fn append_prologue_part(prologue: &mut Vec<u8>, part: &[u8]) {
278    // Length prefixes make component boundaries unambiguous. Raw concatenation
279    // would allow different identifier tuples to produce the same prologue.
280    let len = part.len() as u64;
281    prologue.extend_from_slice(&len.to_be_bytes());
282    prologue.extend_from_slice(part);
283}
284
285fn ensure_noise_frame_len(
286    frame_len: usize,
287    message: &'static str,
288) -> Result<(), NoiseChannelError> {
289    if frame_len > MAX_MESSAGE_LEN {
290        return Err(NoiseChannelError::InvalidMessage(message));
291    }
292    Ok(())
293}
294
295#[derive(Debug, thiserror::Error)]
296pub enum NoiseChannelError {
297    #[error("Noise channel key generation failed: {0}")]
298    KeyGeneration(String),
299    #[error("invalid Noise channel public key: {0}")]
300    InvalidPublicKey(&'static str),
301    #[error("invalid Noise channel message: {0}")]
302    InvalidMessage(&'static str),
303    #[error("Noise channel handshake failed: {0}")]
304    Handshake(String),
305    #[error("Noise channel transport failed: {0}")]
306    Transport(String),
307}
308
309impl From<clatter::error::HandshakeError> for NoiseChannelError {
310    fn from(error: clatter::error::HandshakeError) -> Self {
311        Self::Handshake(error.to_string())
312    }
313}
314
315impl From<clatter::error::TransportError> for NoiseChannelError {
316    fn from(error: clatter::error::TransportError) -> Self {
317        Self::Transport(error.to_string())
318    }
319}
320
321#[cfg(test)]
322#[path = "noise_channel_tests.rs"]
323mod tests;
codex_exec_server/noise_channel.rs

codex_exec_server/
noise_channel.rs
