use std::io::{self, Read, Write};
#[cfg(any(feature = "dpdk", test))]
use tracing::{info, warn};
#[cfg(any(feature = "dpdk", test))]
use std::sync::atomic::{AtomicBool, Ordering};
use melin_transport_core::replication::protocol::{
MAX_CONTROL_FRAME, decode_auth_result, decode_challenge, decode_challenge_response,
encode_auth_failed, encode_auth_ok, encode_challenge, encode_challenge_response, read_frame,
};
#[cfg(any(feature = "dpdk", test))]
use super::receiver_transport::{FrameResult, compact_recv_buf, try_extract_frame};
pub(super) fn generate_challenge_nonce() -> io::Result<[u8; 32]> {
let mut nonce = [0u8; 32];
getrandom::fill(&mut nonce).map_err(|e| io::Error::other(format!("getrandom failed: {e}")))?;
Ok(nonce)
}
pub(super) fn verify_challenge_response(
nonce: &[u8; 32],
response_frame: &[u8],
authorized_keys: &melin_app::auth::AuthorizedKeys,
) -> io::Result<()> {
use ed25519_dalek::{Verifier, VerifyingKey};
let (signature_bytes, pubkey_bytes) = decode_challenge_response(response_frame)
.map_err(|e| io::Error::other(format!("bad challenge response: {e}")))?;
let role = authorized_keys
.lookup(&pubkey_bytes)
.ok_or_else(|| io::Error::other("unknown replication key"))?;
if !role.is_replication() {
return Err(io::Error::other(format!(
"key listed as {}, expected replication",
authorized_keys.token(role)
)));
}
let verifying_key = VerifyingKey::from_bytes(&pubkey_bytes)
.map_err(|e| io::Error::other(format!("invalid public key: {e}")))?;
let signature = ed25519_dalek::Signature::from_bytes(&signature_bytes);
verifying_key
.verify(nonce, &signature)
.map_err(|e| io::Error::other(format!("signature verification failed: {e}")))?;
Ok(())
}
pub(super) fn authenticate_replica<S: Read + Write>(
stream: &mut S,
authorized_keys: &melin_app::auth::AuthorizedKeys,
) -> io::Result<()> {
let nonce = generate_challenge_nonce()?;
let mut buf = Vec::with_capacity(64);
encode_challenge(&nonce, &mut buf);
stream.write_all(&buf)?;
stream.flush()?;
let frame = read_frame(stream, MAX_CONTROL_FRAME)?;
if let Err(e) = verify_challenge_response(&nonce, &frame, authorized_keys) {
buf.clear();
encode_auth_failed(&mut buf);
let _ = stream.write_all(&buf);
return Err(e);
}
buf.clear();
encode_auth_ok(&mut buf);
stream.write_all(&buf)?;
stream.flush()?;
Ok(())
}
pub(super) fn authenticate_with_primary<S: Read + Write>(
stream: &mut S,
signing_key: &ed25519_dalek::SigningKey,
) -> io::Result<()> {
use ed25519_dalek::Signer;
let frame = read_frame(stream, MAX_CONTROL_FRAME)?;
let nonce = decode_challenge(&frame)?;
let signature = signing_key.sign(&nonce);
let pubkey = signing_key.verifying_key();
let mut buf = Vec::with_capacity(128);
encode_challenge_response(&signature.to_bytes(), pubkey.as_bytes(), &mut buf);
stream.write_all(&buf)?;
stream.flush()?;
let result_frame = read_frame(stream, MAX_CONTROL_FRAME)?;
match decode_auth_result(&result_frame)? {
true => Ok(()),
false => Err(io::Error::other("primary rejected replication key")),
}
}
#[cfg(any(feature = "dpdk", test))]
pub(super) struct AuthChallenge {
pub(super) nonce: [u8; 32],
pub(super) deadline: std::time::Instant,
}
#[cfg(any(feature = "dpdk", test))]
pub(super) trait AuthTransport {
type Handle: Copy;
fn is_active(&mut self, handle: Self::Handle) -> bool;
fn recv_into_vec(&mut self, handle: Self::Handle, dest: &mut Vec<u8>);
fn queue_send(&mut self, handle: Self::Handle, data: &[u8]) -> bool;
fn poll(&mut self);
}
#[cfg(any(feature = "dpdk", test))]
pub(super) enum AuthOutcome {
Pending,
Authenticated,
Rejected,
}
#[cfg(any(feature = "dpdk", test))]
pub(super) fn step_authentication<T: AuthTransport>(
transport: &mut T,
handle: T::Handle,
challenge: &AuthChallenge,
recv_buf: &mut Vec<u8>,
send_buf: &mut Vec<u8>,
authorized_keys: &melin_app::auth::AuthorizedKeys,
slot_idx: usize,
) -> AuthOutcome {
if !transport.is_active(handle) {
warn!(slot = slot_idx, "replica disconnected during auth (DPDK)");
return AuthOutcome::Rejected;
}
if std::time::Instant::now() >= challenge.deadline {
warn!(
slot = slot_idx,
"replica auth timed out (DPDK) — disconnecting"
);
return AuthOutcome::Rejected;
}
transport.recv_into_vec(handle, recv_buf);
match try_extract_frame(recv_buf, MAX_CONTROL_FRAME) {
FrameResult::Complete(payload_start, frame_end) => {
let verdict = verify_challenge_response(
&challenge.nonce,
&recv_buf[payload_start..frame_end],
authorized_keys,
);
compact_recv_buf(recv_buf, frame_end);
send_buf.clear();
match verdict {
Ok(()) => {
encode_auth_ok(send_buf);
let _ = transport.queue_send(handle, send_buf);
info!(slot = slot_idx, "replica authenticated (DPDK)");
AuthOutcome::Authenticated
}
Err(e) => {
warn!(slot = slot_idx, error = %e, "replica auth failed (DPDK) — disconnecting");
encode_auth_failed(send_buf);
let _ = transport.queue_send(handle, send_buf);
transport.poll();
AuthOutcome::Rejected
}
}
}
FrameResult::Oversized => {
warn!(
slot = slot_idx,
"oversized auth frame (DPDK) — disconnecting"
);
AuthOutcome::Rejected
}
FrameResult::Incomplete => AuthOutcome::Pending,
}
}
#[cfg(any(feature = "dpdk", test))]
pub(super) struct PolledAuthStream<'a, T: AuthTransport> {
pub(super) transport: &'a mut T,
pub(super) handle: T::Handle,
pub(super) recv_buf: &'a mut Vec<u8>,
pub(super) shutdown: &'a AtomicBool,
pub(super) deadline: std::time::Instant,
}
#[cfg(any(feature = "dpdk", test))]
impl<T: AuthTransport> Read for PolledAuthStream<'_, T> {
fn read(&mut self, out: &mut [u8]) -> io::Result<usize> {
loop {
if !self.recv_buf.is_empty() {
let n = self.recv_buf.len().min(out.len());
out[..n].copy_from_slice(&self.recv_buf[..n]);
self.recv_buf.drain(..n);
return Ok(n);
}
if self.shutdown.load(Ordering::Relaxed) {
return Err(io::Error::other("shutdown during auth"));
}
if std::time::Instant::now() >= self.deadline {
return Err(io::Error::new(io::ErrorKind::TimedOut, "auth timed out"));
}
self.transport.poll();
self.transport.recv_into_vec(self.handle, self.recv_buf);
if self.recv_buf.is_empty() && !self.transport.is_active(self.handle) {
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
"primary disconnected during auth",
));
}
std::thread::yield_now();
}
}
}
#[cfg(any(feature = "dpdk", test))]
impl<T: AuthTransport> Write for PolledAuthStream<'_, T> {
fn write(&mut self, data: &[u8]) -> io::Result<usize> {
if self.transport.queue_send(self.handle, data) {
Ok(data.len())
} else {
Err(io::Error::other("TX queue full during auth"))
}
}
fn flush(&mut self) -> io::Result<()> {
self.transport.poll();
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use ed25519_dalek::{Signer, SigningKey};
use std::collections::VecDeque;
fn keys_for(key: &SigningKey, role: &str) -> melin_app::auth::AuthorizedKeys {
let pub_b64 = base64::Engine::encode(
&base64::engine::general_purpose::STANDARD,
key.verifying_key().to_bytes(),
);
crate::test_roles::desk_keys(role, &pub_b64)
}
fn response_payload(key: &SigningKey, nonce: &[u8; 32]) -> Vec<u8> {
let sig = key.sign(nonce);
let mut frame = Vec::new();
melin_transport_core::replication::protocol::encode_challenge_response(
&sig.to_bytes(),
key.verifying_key().as_bytes(),
&mut frame,
);
frame[4..].to_vec()
}
#[test]
fn verify_accepts_valid_replication_key() {
let key = SigningKey::from_bytes(&[0x11; 32]);
let keys = keys_for(&key, "replication");
let nonce = [0x42; 32];
assert!(verify_challenge_response(&nonce, &response_payload(&key, &nonce), &keys).is_ok());
}
#[test]
fn verify_rejects_unknown_key() {
let signer = SigningKey::from_bytes(&[0x22; 32]);
let listed = SigningKey::from_bytes(&[0x33; 32]);
let keys = keys_for(&listed, "replication"); let nonce = [0x42; 32];
let err = verify_challenge_response(&nonce, &response_payload(&signer, &nonce), &keys)
.unwrap_err();
assert!(err.to_string().contains("unknown"));
}
#[test]
fn verify_rejects_non_replication_permission() {
let key = SigningKey::from_bytes(&[0x44; 32]);
let keys = keys_for(&key, "trader"); let nonce = [0x42; 32];
let err =
verify_challenge_response(&nonce, &response_payload(&key, &nonce), &keys).unwrap_err();
assert_eq!(
err.to_string(),
"key listed as trader, expected replication"
);
}
#[test]
fn verify_rejects_signature_over_wrong_nonce() {
let key = SigningKey::from_bytes(&[0x55; 32]);
let keys = keys_for(&key, "replication");
let signed = [0x01; 32];
let challenge = [0x02; 32];
let err = verify_challenge_response(&challenge, &response_payload(&key, &signed), &keys)
.unwrap_err();
assert!(err.to_string().contains("signature"));
}
struct MockAuthTransport {
active: bool,
recv_rounds: VecDeque<Vec<u8>>,
close_when_drained: bool,
sent: Vec<u8>,
tx_full: bool,
}
impl MockAuthTransport {
fn with_incoming(incoming: Vec<u8>) -> Self {
Self {
active: true,
recv_rounds: VecDeque::from([incoming]),
close_when_drained: false,
sent: Vec::new(),
tx_full: false,
}
}
fn scripted(rounds: Vec<Vec<u8>>) -> Self {
Self {
active: true,
recv_rounds: rounds.into(),
close_when_drained: true,
sent: Vec::new(),
tx_full: false,
}
}
}
impl AuthTransport for MockAuthTransport {
type Handle = ();
fn is_active(&mut self, _: ()) -> bool {
self.active
}
fn recv_into_vec(&mut self, _: (), dest: &mut Vec<u8>) {
if let Some(mut chunk) = self.recv_rounds.pop_front() {
dest.append(&mut chunk);
}
if self.recv_rounds.is_empty() && self.close_when_drained {
self.active = false;
}
}
fn queue_send(&mut self, _: (), data: &[u8]) -> bool {
if self.tx_full {
return false;
}
self.sent.extend_from_slice(data);
true
}
fn poll(&mut self) {}
}
fn response_frame(key: &SigningKey, nonce: &[u8; 32]) -> Vec<u8> {
let sig = key.sign(nonce);
let mut frame = Vec::new();
encode_challenge_response(&sig.to_bytes(), key.verifying_key().as_bytes(), &mut frame);
frame
}
fn challenge_at(nonce: [u8; 32], deadline: std::time::Instant) -> AuthChallenge {
AuthChallenge { nonce, deadline }
}
fn far_future() -> std::time::Instant {
std::time::Instant::now() + std::time::Duration::from_secs(60)
}
fn run_step(
incoming: Vec<u8>,
challenge: &AuthChallenge,
keys: &melin_app::auth::AuthorizedKeys,
) -> (AuthOutcome, MockAuthTransport) {
let mut tx = MockAuthTransport::with_incoming(incoming);
let mut recv = Vec::new();
let mut send = Vec::new();
let outcome = step_authentication(&mut tx, (), challenge, &mut recv, &mut send, keys, 0);
(outcome, tx)
}
#[test]
fn step_authenticates_valid_response() {
let key = SigningKey::from_bytes(&[0x11; 32]);
let keys = keys_for(&key, "replication");
let nonce = [0x42; 32];
let challenge = challenge_at(nonce, far_future());
let (outcome, tx) = run_step(response_frame(&key, &nonce), &challenge, &keys);
assert!(matches!(outcome, AuthOutcome::Authenticated));
assert!(
decode_auth_result(&tx.sent[4..]).expect("AuthOk frame"),
"AuthOk should be queued on success"
);
}
#[test]
fn step_rejects_and_signals_failure_on_bad_signature() {
let key = SigningKey::from_bytes(&[0x11; 32]);
let keys = keys_for(&key, "replication");
let challenge = challenge_at([0x42; 32], far_future());
let (outcome, tx) = run_step(response_frame(&key, &[0x01; 32]), &challenge, &keys);
assert!(matches!(outcome, AuthOutcome::Rejected));
assert!(
!decode_auth_result(&tx.sent[4..]).expect("AuthFailed frame"),
"AuthFailed should be queued for a verification failure"
);
}
#[test]
fn step_rejects_unknown_key() {
let signer = SigningKey::from_bytes(&[0x22; 32]);
let listed = SigningKey::from_bytes(&[0x33; 32]);
let keys = keys_for(&listed, "replication"); let nonce = [0x42; 32];
let challenge = challenge_at(nonce, far_future());
let (outcome, tx) = run_step(response_frame(&signer, &nonce), &challenge, &keys);
assert!(matches!(outcome, AuthOutcome::Rejected));
assert!(
!decode_auth_result(&tx.sent[4..]).expect("AuthFailed frame"),
"an unknown key is told AuthFailed"
);
}
#[test]
fn step_rejects_non_replication_permission() {
let key = SigningKey::from_bytes(&[0x44; 32]);
let keys = keys_for(&key, "trader"); let nonce = [0x42; 32];
let challenge = challenge_at(nonce, far_future());
let (outcome, tx) = run_step(response_frame(&key, &nonce), &challenge, &keys);
assert!(matches!(outcome, AuthOutcome::Rejected));
assert!(
!decode_auth_result(&tx.sent[4..]).expect("AuthFailed frame"),
"a wrong-permission key is told AuthFailed"
);
}
#[test]
fn step_times_out_past_deadline() {
let key = SigningKey::from_bytes(&[0x11; 32]);
let keys = keys_for(&key, "replication");
let challenge = challenge_at([0x42; 32], std::time::Instant::now());
let (outcome, tx) = run_step(Vec::new(), &challenge, &keys);
assert!(matches!(outcome, AuthOutcome::Rejected));
assert!(
tx.sent.is_empty(),
"no AuthFailed on timeout — nothing to say"
);
}
#[test]
fn step_rejects_oversized_frame() {
let key = SigningKey::from_bytes(&[0x11; 32]);
let keys = keys_for(&key, "replication");
let challenge = challenge_at([0x42; 32], far_future());
let oversized = ((MAX_CONTROL_FRAME + 1) as u32).to_le_bytes().to_vec();
let (outcome, tx) = run_step(oversized, &challenge, &keys);
assert!(matches!(outcome, AuthOutcome::Rejected));
assert!(tx.sent.is_empty(), "no AuthFailed on a malformed frame");
}
#[test]
fn step_pending_on_partial_frame() {
let key = SigningKey::from_bytes(&[0x11; 32]);
let keys = keys_for(&key, "replication");
let nonce = [0x42; 32];
let challenge = challenge_at(nonce, far_future());
let mut partial = response_frame(&key, &nonce);
partial.truncate(3);
let (outcome, tx) = run_step(partial, &challenge, &keys);
assert!(matches!(outcome, AuthOutcome::Pending));
assert!(tx.sent.is_empty(), "an incomplete frame just waits");
}
#[test]
fn step_rejects_on_disconnect() {
let key = SigningKey::from_bytes(&[0x11; 32]);
let keys = keys_for(&key, "replication");
let challenge = challenge_at([0x42; 32], far_future());
let mut tx = MockAuthTransport::with_incoming(Vec::new());
tx.active = false; let mut recv = Vec::new();
let mut send = Vec::new();
let outcome = step_authentication(&mut tx, (), &challenge, &mut recv, &mut send, &keys, 0);
assert!(matches!(outcome, AuthOutcome::Rejected));
assert!(tx.sent.is_empty());
}
fn with_stream<R>(
tx: &mut MockAuthTransport,
shutdown: &AtomicBool,
deadline: std::time::Instant,
body: impl FnOnce(&mut PolledAuthStream<'_, MockAuthTransport>) -> R,
) -> R {
let mut recv = Vec::new();
let mut stream = PolledAuthStream {
transport: tx,
handle: (),
recv_buf: &mut recv,
shutdown,
deadline,
};
body(&mut stream)
}
#[test]
fn polled_read_times_out_past_deadline() {
let mut tx = MockAuthTransport::scripted(vec![]);
let shutdown = AtomicBool::new(false);
let err = with_stream(&mut tx, &shutdown, std::time::Instant::now(), |s| {
std::io::Read::read(s, &mut [0u8; 8]).unwrap_err()
});
assert_eq!(err.kind(), std::io::ErrorKind::TimedOut);
}
#[test]
fn polled_read_errors_on_shutdown() {
let mut tx = MockAuthTransport::scripted(vec![]);
let shutdown = AtomicBool::new(true);
let err = with_stream(&mut tx, &shutdown, far_future(), |s| {
std::io::Read::read(s, &mut [0u8; 8]).unwrap_err()
});
assert!(err.to_string().contains("shutdown"));
}
#[test]
fn polled_read_reports_disconnect_as_eof() {
let mut tx = MockAuthTransport::scripted(vec![]);
let shutdown = AtomicBool::new(false);
let err = with_stream(&mut tx, &shutdown, far_future(), |s| {
std::io::Read::read(s, &mut [0u8; 8]).unwrap_err()
});
assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
}
#[test]
fn polled_read_returns_bytes_arriving_with_fin() {
let mut tx = MockAuthTransport::scripted(vec![vec![1, 2, 3, 4]]);
let shutdown = AtomicBool::new(false);
with_stream(&mut tx, &shutdown, far_future(), |s| {
let mut out = [0u8; 8];
let n = std::io::Read::read(s, &mut out)
.expect("bytes delivered with the FIN are returned");
assert_eq!(&out[..n], &[1, 2, 3, 4]);
let err = std::io::Read::read(s, &mut out).unwrap_err();
assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
});
}
#[test]
fn polled_write_errors_when_tx_queue_full() {
let mut tx = MockAuthTransport::scripted(vec![]);
tx.tx_full = true;
let shutdown = AtomicBool::new(false);
let err = with_stream(&mut tx, &shutdown, far_future(), |s| {
std::io::Write::write(s, b"auth-frame").unwrap_err()
});
assert!(err.to_string().contains("TX queue full"));
}
#[test]
fn polled_write_queues_frame_then_flush_polls() {
let mut tx = MockAuthTransport::scripted(vec![]);
let shutdown = AtomicBool::new(false);
with_stream(&mut tx, &shutdown, far_future(), |s| {
assert_eq!(std::io::Write::write(s, b"hello").unwrap(), 5);
std::io::Write::flush(s).unwrap();
});
assert_eq!(tx.sent, b"hello", "the queued frame reaches the TX queue");
}
#[test]
fn polled_stream_drives_authenticate_with_primary_to_ok() {
let key = SigningKey::from_bytes(&[0x11; 32]);
let nonce = [0x42; 32];
let mut challenge = Vec::new();
encode_challenge(&nonce, &mut challenge);
let mut authok = Vec::new();
encode_auth_ok(&mut authok);
let mut tx = MockAuthTransport::scripted(vec![challenge, authok]);
let shutdown = AtomicBool::new(false);
let result = with_stream(&mut tx, &shutdown, far_future(), |s| {
authenticate_with_primary(s, &key)
});
result.expect("auth succeeds when the primary sends AuthOk");
let keys = keys_for(&key, "replication");
verify_challenge_response(&nonce, &tx.sent[4..], &keys)
.expect("the response the adapter wrote verifies");
}
#[test]
fn polled_stream_surfaces_auth_failed() {
let key = SigningKey::from_bytes(&[0x11; 32]);
let nonce = [0x42; 32];
let mut challenge = Vec::new();
encode_challenge(&nonce, &mut challenge);
let mut authfailed = Vec::new();
encode_auth_failed(&mut authfailed);
let mut tx = MockAuthTransport::scripted(vec![challenge, authfailed]);
let shutdown = AtomicBool::new(false);
let err = with_stream(&mut tx, &shutdown, far_future(), |s| {
authenticate_with_primary(s, &key).unwrap_err()
});
assert!(err.to_string().contains("rejected"));
}
}