use std::collections::HashMap;
use a7n_proto::v1::{Capability, Envelope, PeerInfo, Platform};
use a7n_proto::{MachineId, ProtoError};
use quinn::{Connection, ConnectionError, RecvStream, SendStream, VarInt};
use rustls::pki_types::CertificateDer;
use tokio::sync::mpsc;
use crate::identity::{self, IdentityError};
pub(crate) const CLOSE_NORMAL: VarInt = VarInt::from_u32(0);
pub(crate) const CLOSE_SUPERSEDED: VarInt = VarInt::from_u32(1);
pub(crate) const CLOSE_SILENT: VarInt = VarInt::from_u32(2);
pub(crate) const CLOSE_UNACKED: VarInt = VarInt::from_u32(3);
#[derive(Debug, thiserror::Error)]
pub(crate) enum LinkError {
#[error("cannot start connecting: {0}")]
Connect(#[from] quinn::ConnectError),
#[error("connection failed: {0}")]
Connection(#[from] ConnectionError),
#[error("cannot read the control stream: {0}")]
Read(#[from] quinn::ReadError),
#[error("cannot write the control stream: {0}")]
Write(#[from] quinn::WriteError),
#[error("dial superseded before a link to the peer was established")]
Superseded(MachineId),
#[error("handshake timed out")]
Timeout,
#[error("the control stream ended early")]
StreamClosed,
#[error(transparent)]
Proto(#[from] ProtoError),
#[error("update {name} to protocol v{version}")]
UpdateRequired { name: String, version: u32 },
#[error("the daemon loop stopped receiving from this link")]
InboundClosed,
#[error("peer sent an invalid handshake message")]
InvalidHello,
#[error("input sequence exhausted")]
InputSeqExhausted,
#[error("peer presented no certificate")]
NoPeerCertificate,
#[error(transparent)]
Identity(#[from] IdentityError),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Role {
Initiator,
Responder,
}
#[derive(Debug)]
pub(crate) struct Link {
pub(crate) connection: Connection,
pub(crate) control_send: SendStream,
pub(crate) control_recv: RecvStream,
pub(crate) control_buf: Vec<u8>,
pub(crate) peer_id: MachineId,
pub(crate) peer_info: PeerInfo,
pub(crate) capabilities: Vec<Capability>,
pub(crate) role: Role,
}
pub(crate) fn peer_certificate(
connection: &Connection,
) -> Result<CertificateDer<'static>, LinkError> {
let identity = connection
.peer_identity()
.ok_or(LinkError::NoPeerCertificate)?;
let chain = identity
.downcast::<Vec<CertificateDer<'static>>>()
.map_err(|_| LinkError::NoPeerCertificate)?;
chain.into_iter().next().ok_or(LinkError::NoPeerCertificate)
}
pub(crate) fn peer_machine_id(connection: &Connection) -> Result<MachineId, LinkError> {
Ok(identity::machine_id_of(&peer_certificate(connection)?)?)
}
pub(crate) fn surviving_initiator(a: MachineId, b: MachineId) -> MachineId {
a.max(b)
}
fn new_link_wins(own: MachineId, peer: MachineId, new: Role, existing: Role) -> bool {
let survivor = surviving_initiator(own, peer);
let initiator = |role| if role == Role::Initiator { own } else { peer };
initiator(new) == survivor || initiator(existing) != survivor
}
#[derive(Debug)]
pub(crate) struct LinkHandle {
pub(crate) platform: Platform,
pub(crate) capabilities: Vec<Capability>,
pub(crate) sender: mpsc::Sender<Envelope>,
pub(crate) connection: Connection,
pub(crate) role: Role,
}
#[derive(Debug, Default)]
pub(crate) struct Links {
pub(crate) by_peer: HashMap<MachineId, LinkHandle>,
}
impl Links {
pub(crate) fn insert(
&mut self,
peer: MachineId,
link: LinkHandle,
own_id: MachineId,
) -> Option<LinkHandle> {
let Some(existing) = self.by_peer.get(&peer) else {
self.by_peer.insert(peer, link);
return None;
};
if new_link_wins(own_id, peer, link.role, existing.role) {
return self.by_peer.insert(peer, link);
}
Some(link)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn id(byte: u8) -> MachineId {
MachineId::try_from(&[byte; MachineId::LEN][..]).unwrap()
}
#[test]
fn new_link_wins_simultaneous_dial_keeps_larger_initiator() {
let cases = [
(id(2), id(1), Role::Initiator, Role::Responder, true),
(id(2), id(1), Role::Responder, Role::Initiator, false),
(id(1), id(2), Role::Responder, Role::Initiator, true),
(id(1), id(2), Role::Initiator, Role::Responder, false),
];
for (own, peer, new, existing, expected) in cases {
let wins = new_link_wins(own, peer, new, existing);
assert_eq!(wins, expected, "{own:?} {peer:?} {new:?} {existing:?}");
}
}
}