geph4_protocol/
bridge_exit.rs

1use std::{
2    sync::Arc,
3    time::{Duration, SystemTime, UNIX_EPOCH},
4};
5
6use anyhow::Context;
7use async_net::{SocketAddr, UdpSocket};
8use async_trait::async_trait;
9use bytes::Bytes;
10use futures_util::TryFutureExt;
11use nanorpc::{nanorpc_derive, JrpcRequest, JrpcResponse, RpcService, RpcTransport};
12use serde::{Deserialize, Serialize};
13use smol_str::SmolStr;
14use smol_timeout::TimeoutExt;
15
16/// An RpcTransport that implements the symmetrically authenticated bridge-exit protocol.
17#[derive(Clone)]
18pub struct BridgeExitTransport {
19    key: [u8; 32],
20    dest: SocketAddr,
21}
22
23impl BridgeExitTransport {
24    /// Creates a new BridgeExitTransport with the given bridge secret and destination.
25    pub fn new(secret: [u8; 32], exit: SocketAddr) -> Self {
26        Self {
27            key: secret,
28            dest: exit,
29        }
30    }
31}
32
33/// Serve the authenticated bridge-exit protocol, given an RpcService.
34pub async fn serve_bridge_exit<R: RpcService>(
35    socket: UdpSocket,
36    key: [u8; 32],
37    service: R,
38) -> anyhow::Result<()> {
39    let mut buf = [0u8; 2048];
40    let service = Arc::new(service);
41    loop {
42        let (n, client_addr) = socket.recv_from(&mut buf).await?;
43        let service = service.clone();
44        let request = Bytes::copy_from_slice(&buf[..n]);
45        let socket = socket.clone();
46        smolscale::spawn(
47            async move {
48                let (mac, timestamp, plain): ([u8; 32], u64, Bytes) =
49                    stdcode::deserialize(&request)?;
50                let correct_timestamp = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();
51                if timestamp > correct_timestamp + 60
52                    || timestamp < correct_timestamp.saturating_sub(60)
53                {
54                    anyhow::bail!("timestamp out of range")
55                }
56                let mac_key = blake3::keyed_hash(&key, &timestamp.to_be_bytes());
57                let correct_mac = blake3::keyed_hash(mac_key.as_bytes(), &plain);
58                if correct_mac != blake3::Hash::from(mac) {
59                    anyhow::bail!(
60                        "MAC is wrong (given {:?}, recalculated {:?}, mac key {:?}, timestamp {timestamp}, plain {:?})",
61                        blake3::Hash::from(mac),
62                        correct_mac,
63                        mac_key,
64                        plain
65                    )
66                }
67                let request: JrpcRequest = serde_json::from_slice(&plain)?;
68                let response = service.respond_raw(request).await;
69                socket
70                    .send_to(&serde_json::to_vec(&response)?, client_addr)
71                    .await?;
72                anyhow::Ok(())
73            }
74            .map_err(move |e| log::warn!("bad bridge_exit pkt from {client_addr}: {e}")),
75        )
76        .detach()
77    }
78}
79
80#[async_trait]
81impl RpcTransport for BridgeExitTransport {
82    type Error = anyhow::Error;
83
84    async fn call_raw(&self, jrpc: JrpcRequest) -> Result<JrpcResponse, Self::Error> {
85        let plain_vec = serde_json::to_vec(&jrpc)?;
86        let timestamp = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();
87        let mac_key = blake3::keyed_hash(&self.key, &timestamp.to_be_bytes());
88        let mac = blake3::keyed_hash(mac_key.as_bytes(), &plain_vec);
89        log::debug!(
90                    "sending request with mac {:?}, mac_key {:?}, timestamp {timestamp}, bridge key {:?}, plain {:?}",
91                    mac,
92                    mac_key,
93                    hex::encode(self.key),
94                    String::from_utf8_lossy(&plain_vec)
95                );
96        let to_send = stdcode::serialize(&(mac.as_bytes(), timestamp, plain_vec))?;
97        let socket = UdpSocket::bind("0.0.0.0:0").await?;
98        socket.send_to(&to_send, self.dest).await?;
99        let mut buff = [0u8; 2048];
100        let (n, _) = socket
101            .recv_from(&mut buff)
102            .timeout(Duration::from_secs(10))
103            .await
104            .context("udp receive timeout")??;
105        // response is NOT authenticated. this is generally fine.
106        Ok(serde_json::from_slice(&buff[..n])?)
107    }
108}
109
110/// An available raw, kernel-forwardable protocol
111#[derive(Clone, Copy, Debug, PartialEq, PartialOrd, Eq, Hash, Deserialize, Serialize)]
112pub enum LegacyProtocol {
113    Tcp,
114    Udp,
115}
116
117/// The nanorpc_derive trait describing the bridge/exit protocol.
118#[nanorpc_derive]
119#[async_trait]
120pub trait BridgeExitProtocol {
121    /// Advertises an available raw port. If enough resources are available, returns the address to forward traffic to.
122    async fn advertise_raw(
123        &self,
124        protocol: LegacyProtocol,
125        bridge_addr: SocketAddr,
126        bridge_group: SmolStr,
127    ) -> SocketAddr;
128
129    /// Advertises an available raw port, with the v2 protocol. If enough resources are available, returns the address to forward traffic to.
130    async fn advertise_raw_v2(
131        &self,
132        protocol: SmolStr,
133        bridge_addr: SocketAddr,
134        bridge_group: SmolStr,
135    ) -> SocketAddr;
136
137    /// Gets the current load of the exit.
138    async fn load_factor(&self) -> f64;
139}