Skip to main content

netget/server/tor_relay/
circuit.rs

1//! Tor Relay Circuit Management
2//!
3//! Implements circuit state, ntor handshake, and relay cell encryption/decryption
4
5use anyhow::{Context, Result};
6use std::collections::HashMap;
7use std::sync::Arc;
8use tokio::sync::Mutex;
9use tracing::{debug, trace};
10
11// Crypto imports
12use aes::Aes128;
13use ctr::cipher::{KeyIvInit, StreamCipher};
14use ed25519_dalek::SigningKey;
15use hkdf::Hkdf;
16use hmac::{Hmac, Mac};
17use sha2::{Digest, Sha256};
18use x25519_dalek::{EphemeralSecret, PublicKey as X25519PublicKey};
19
20type HmacSha256 = Hmac<Sha256>;
21type Aes128Ctr = ctr::Ctr128BE<Aes128>;
22
23/// Circuit statistics
24#[derive(Debug, Clone)]
25pub struct CircuitStats {
26    pub circuit_id: CircuitId,
27    pub created_at: crate::utils::clock::Instant,
28    pub last_activity: crate::utils::clock::Instant,
29    pub bytes_sent: u64,
30    pub bytes_received: u64,
31    pub active_streams: usize,
32}
33
34/// Relay statistics (aggregate across all circuits)
35#[derive(Debug, Clone)]
36pub struct RelayStats {
37    pub total_circuits: usize,
38    pub total_streams: usize,
39    pub total_bytes_sent: u64,
40    pub total_bytes_received: u64,
41    pub circuit_stats: Vec<CircuitStats>,
42}
43
44/// Circuit identifier (4 bytes for OR protocol v4)
45#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
46pub struct CircuitId(u32);
47
48impl CircuitId {
49    pub fn new(id: u32) -> Self {
50        Self(id)
51    }
52
53    pub fn as_u32(&self) -> u32 {
54        self.0
55    }
56
57    pub fn from_bytes(bytes: &[u8; 4]) -> Self {
58        Self(u32::from_be_bytes(*bytes))
59    }
60
61    pub fn to_bytes(&self) -> [u8; 4] {
62        self.0.to_be_bytes()
63    }
64}
65
66impl std::fmt::Display for CircuitId {
67    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
68        write!(f, "0x{:08x}", self.0)
69    }
70}
71
72/// Stream identifier within a circuit
73#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
74pub struct StreamId(u16);
75
76impl StreamId {
77    pub fn new(id: u16) -> Self {
78        Self(id)
79    }
80
81    pub fn as_u16(&self) -> u16 {
82        self.0
83    }
84}
85
86/// Circuit window constants (from tor-spec.txt)
87pub const CIRCUIT_WINDOW_START: u16 = 1000;
88pub const CIRCUIT_WINDOW_INCREMENT: u16 = 100;
89
90/// Circuit state
91#[derive(Debug)]
92pub struct Circuit {
93    /// Circuit ID
94    pub id: CircuitId,
95    /// Circuit crypto state (forward and backward keys)
96    pub crypto: CircuitCrypto,
97    /// Stream manager (now in a separate module)
98    pub stream_manager: crate::server::tor_relay::stream::StreamManager,
99    /// Next hop (for extending circuits)
100    pub next_hop: Option<String>,
101    /// Circuit creation timestamp
102    pub created_at: crate::utils::clock::Instant,
103    /// Last activity timestamp
104    pub last_activity: crate::utils::clock::Instant,
105    /// Total bytes sent to client
106    pub bytes_sent: u64,
107    /// Total bytes received from client
108    pub bytes_received: u64,
109    /// Package window (decremented on sending, incremented on SENDME)
110    pub package_window: u16,
111    /// Deliver window (decremented on receiving, send SENDME at threshold)
112    pub deliver_window: u16,
113    /// Count of RELAY cells received (for circuit-level SENDME)
114    pub relay_cells_received: u16,
115}
116
117impl Circuit {
118    /// Create new circuit with crypto keys from ntor handshake
119    pub fn new(id: CircuitId, key_material: KeyMaterial) -> Self {
120        let now = crate::utils::clock::Instant::now();
121        Self {
122            id,
123            crypto: CircuitCrypto::new(key_material),
124            stream_manager: crate::server::tor_relay::stream::StreamManager::new(),
125            next_hop: None,
126            created_at: now,
127            last_activity: now,
128            bytes_sent: 0,
129            bytes_received: 0,
130            package_window: CIRCUIT_WINDOW_START,
131            deliver_window: CIRCUIT_WINDOW_START,
132            relay_cells_received: 0,
133        }
134    }
135
136    /// Record bytes sent to client
137    pub fn record_sent(&mut self, bytes: u64) {
138        self.bytes_sent += bytes;
139        self.last_activity = crate::utils::clock::Instant::now();
140    }
141
142    /// Record bytes received from client
143    pub fn record_received(&mut self, bytes: u64) {
144        self.bytes_received += bytes;
145        self.last_activity = crate::utils::clock::Instant::now();
146    }
147
148    /// Decrypt incoming RELAY cell
149    pub fn decrypt_relay_cell(&mut self, payload: &mut [u8]) -> Result<()> {
150        self.crypto.decrypt(payload)
151    }
152
153    /// Encrypt outgoing RELAY cell
154    pub fn encrypt_relay_cell(&mut self, payload: &mut [u8]) -> Result<()> {
155        self.crypto.encrypt(payload)
156    }
157
158    /// Record RELAY cell received - returns true if circuit-level SENDME should be sent
159    pub fn record_relay_received(&mut self) -> bool {
160        self.deliver_window = self.deliver_window.saturating_sub(1);
161        self.relay_cells_received += 1;
162
163        // Send circuit-level SENDME every CIRCUIT_WINDOW_INCREMENT cells
164        if self.relay_cells_received >= CIRCUIT_WINDOW_INCREMENT {
165            self.relay_cells_received = 0;
166            self.deliver_window += CIRCUIT_WINDOW_INCREMENT;
167            return true;
168        }
169        false
170    }
171
172    /// Process received circuit-level SENDME - increment package window
173    pub fn process_circuit_sendme(&mut self) {
174        self.package_window += CIRCUIT_WINDOW_INCREMENT;
175        trace!(
176            "Circuit {} package window increased to {}",
177            self.id.as_u32(),
178            self.package_window
179        );
180    }
181
182    /// Check if we can send RELAY cells (package window > 0)
183    pub fn can_send_relay(&self) -> bool {
184        self.package_window > 0
185    }
186
187    /// Decrement package window when sending RELAY cell
188    pub fn consume_package_window(&mut self) {
189        self.package_window = self.package_window.saturating_sub(1);
190    }
191}
192
193/// Circuit crypto state (AES-128-CTR, one keystream per direction).
194///
195/// tor-spec names the directions from the *client's* point of view: "forward" is
196/// client -> relay and uses Kf, "backward" is relay -> client and uses Kb. This relay
197/// therefore decrypts with the forward cipher and encrypts with the backward one. The two
198/// were previously swapped.
199pub struct CircuitCrypto {
200    /// Forward cipher, Kf: client -> relay, i.e. what we decrypt with
201    forward_cipher: Aes128Ctr,
202    /// Backward cipher, Kb: relay -> client, i.e. what we encrypt with
203    backward_cipher: Aes128Ctr,
204    /// Forward digest for integrity
205    forward_digest: Sha256,
206    /// Backward digest for integrity
207    backward_digest: Sha256,
208}
209
210impl std::fmt::Debug for CircuitCrypto {
211    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
212        f.debug_struct("CircuitCrypto")
213            .field("forward_cipher", &"<cipher>")
214            .field("backward_cipher", &"<cipher>")
215            .field("forward_digest", &"<digest>")
216            .field("backward_digest", &"<digest>")
217            .finish()
218    }
219}
220
221impl CircuitCrypto {
222    /// Create new circuit crypto from key material.
223    ///
224    /// KNOWN NON-CONFORMANCE: the running digests are SHA-256 and are never consumed. Every
225    /// outbound RELAY cell leaves the 4-byte digest field zero (`build_relay_cell` says
226    /// "filled by encryption"; nothing fills it) and no inbound digest is verified, so the
227    /// `recognized`/digest integrity check of tor-spec 6.1 does not happen. Conforming would
228    /// need a running SHA-1, and no SHA-1 dependency is available to this crate.
229    pub fn new(keys: KeyMaterial) -> Self {
230        // Initialize AES-CTR ciphers
231        let forward_cipher = Aes128Ctr::new(&keys.kf.into(), &[0u8; 16].into());
232        let backward_cipher = Aes128Ctr::new(&keys.kb.into(), &[0u8; 16].into());
233
234        // Initialize SHA-256 digests
235        let mut forward_digest = Sha256::new();
236        forward_digest.update(&keys.df);
237        let mut backward_digest = Sha256::new();
238        backward_digest.update(&keys.db);
239
240        Self {
241            forward_cipher,
242            backward_cipher,
243            forward_digest,
244            backward_digest,
245        }
246    }
247
248    /// Decrypt an inbound payload (client -> relay), which uses Kf.
249    pub fn decrypt(&mut self, payload: &mut [u8]) -> Result<()> {
250        self.forward_cipher.apply_keystream(payload);
251        self.forward_digest.update(&*payload);
252        Ok(())
253    }
254
255    /// Encrypt an outbound payload (relay -> client), which uses Kb.
256    pub fn encrypt(&mut self, payload: &mut [u8]) -> Result<()> {
257        self.backward_digest.update(&*payload);
258        self.backward_cipher.apply_keystream(payload);
259        Ok(())
260    }
261}
262
263/// Key material derived from ntor handshake
264#[derive(Debug, Clone)]
265pub struct KeyMaterial {
266    /// Forward cipher key (Kf)
267    pub kf: [u8; 16],
268    /// Backward cipher key (Kb)
269    pub kb: [u8; 16],
270    /// Forward digest key (Df)
271    pub df: [u8; 20],
272    /// Backward digest key (Db)
273    pub db: [u8; 20],
274}
275
276/// ntor handshake constants
277const PROTOID: &[u8] = b"ntor-curve25519-sha256-1";
278const T_MAC: &[u8] = b"ntor-curve25519-sha256-1:mac";
279const T_KEY: &[u8] = b"ntor-curve25519-sha256-1:key_extract";
280const T_VERIFY: &[u8] = b"ntor-curve25519-sha256-1:verify";
281const M_EXPAND: &[u8] = b"ntor-curve25519-sha256-1:key_expand";
282
283/// ntor handshake server-side implementation
284pub struct NtorServer {
285    /// Server's long-term identity key (Ed25519)
286    identity_key: SigningKey,
287    /// Server's long-term onion key (x25519)
288    onion_key: x25519_dalek::StaticSecret,
289    /// Server's public onion key
290    onion_pubkey: X25519PublicKey,
291}
292
293impl NtorServer {
294    /// Create new ntor server with generated keys
295    pub fn new() -> Self {
296        use rand::rngs::OsRng;
297
298        let identity_key = SigningKey::generate(&mut OsRng);
299        let onion_key = x25519_dalek::StaticSecret::random_from_rng(OsRng);
300        let onion_pubkey = X25519PublicKey::from(&onion_key);
301
302        Self {
303            identity_key,
304            onion_key,
305            onion_pubkey,
306        }
307    }
308
309    /// Get server's public onion key (B in spec)
310    pub fn onion_public_key(&self) -> &X25519PublicKey {
311        &self.onion_pubkey
312    }
313
314    /// Get server's identity fingerprint
315    pub fn identity_fingerprint(&self) -> [u8; 20] {
316        let pubkey = self.identity_key.verifying_key();
317        let mut hasher = Sha256::new();
318        hasher.update(pubkey.as_bytes());
319        let hash = hasher.finalize();
320        let mut fingerprint = [0u8; 20];
321        fingerprint.copy_from_slice(&hash[..20]);
322        fingerprint
323    }
324
325    /// Perform server-side ntor handshake
326    ///
327    /// Input: client's public key X from CREATE2 cell (32 bytes)
328    /// Output: (server's ephemeral public key Y, auth hash H, key material)
329    pub fn server_handshake(
330        &self,
331        client_x: &[u8; 32],
332    ) -> Result<([u8; 32], [u8; 32], KeyMaterial)> {
333        use rand::rngs::OsRng;
334
335        // Parse client's public key X
336        let client_pubkey = X25519PublicKey::from(*client_x);
337
338        // Generate ephemeral keypair (y, Y)
339        let y = EphemeralSecret::random_from_rng(OsRng);
340        let big_y = X25519PublicKey::from(&y);
341
342        // Compute shared secrets
343        // EXP(X, y) - DH with client's public key and our ephemeral secret
344        let xy = y.diffie_hellman(&client_pubkey);
345
346        // EXP(X, b) - DH with client's public key and our onion key
347        let xb = self.onion_key.diffie_hellman(&client_pubkey);
348
349        // Get identity fingerprint (ID)
350        let id = self.identity_fingerprint();
351
352        // Construct secret_input:
353        // secret_input = EXP(X,y) | EXP(X,b) | ID | B | X | Y | PROTOID
354        let mut secret_input = Vec::with_capacity(32 + 32 + 20 + 32 + 32 + 32 + PROTOID.len());
355        secret_input.extend_from_slice(xy.as_bytes());
356        secret_input.extend_from_slice(xb.as_bytes());
357        secret_input.extend_from_slice(&id);
358        secret_input.extend_from_slice(self.onion_pubkey.as_bytes());
359        secret_input.extend_from_slice(client_x);
360        secret_input.extend_from_slice(big_y.as_bytes());
361        secret_input.extend_from_slice(PROTOID);
362
363        trace!("ntor secret_input length: {}", secret_input.len());
364
365        // Derive KEY_SEED using HMAC-SHA256
366        // KEY_SEED = H(secret_input, t_key)
367        let mut mac = HmacSha256::new_from_slice(T_KEY).context("Failed to create HMAC")?;
368        mac.update(&secret_input);
369        let key_seed = mac.finalize().into_bytes();
370
371        // Derive verify using HMAC-SHA256
372        // verify = H(secret_input, t_verify)
373        let mut mac = HmacSha256::new_from_slice(T_VERIFY).context("Failed to create HMAC")?;
374        mac.update(&secret_input);
375        let verify = mac.finalize().into_bytes();
376
377        // Construct auth_input for authentication
378        // auth_input = verify | ID | B | Y | X | PROTOID | "Server"
379        let mut auth_input = Vec::with_capacity(32 + 20 + 32 + 32 + 32 + PROTOID.len() + 6);
380        auth_input.extend_from_slice(&verify);
381        auth_input.extend_from_slice(&id);
382        auth_input.extend_from_slice(self.onion_pubkey.as_bytes());
383        auth_input.extend_from_slice(big_y.as_bytes());
384        auth_input.extend_from_slice(client_x);
385        auth_input.extend_from_slice(PROTOID);
386        auth_input.extend_from_slice(b"Server");
387
388        // Compute AUTH using HMAC-SHA256
389        // AUTH = H(auth_input, t_mac)
390        let mut mac = HmacSha256::new_from_slice(T_MAC).context("Failed to create HMAC")?;
391        mac.update(&auth_input);
392        let auth = mac.finalize().into_bytes();
393
394        // Derive key material using HKDF-SHA256 (salt = t_key, IKM = KEY_SEED, info = m_expand).
395        //
396        // tor-spec 5.2.2 fixes the layout of the output as
397        //     Df(20) | Db(20) | Kf(16) | Kb(16) | KH(20)
398        // for 92 bytes total. This used to take 72 bytes in the order Kf|Kb|Df|Db, which put
399        // the cipher keys where the spec puts digest seeds - so even a client that agreed on
400        // KEY_SEED derived entirely different AES keys and every cell after CREATED2 decrypted
401        // to noise.
402        let hkdf = Hkdf::<Sha256>::new(Some(T_KEY), &key_seed);
403        let mut okm = [0u8; 92];
404        hkdf.expand(M_EXPAND, &mut okm)
405            .map_err(|_| anyhow::anyhow!("Failed to expand key material"))?;
406
407        let key_material = KeyMaterial {
408            df: okm[0..20].try_into().unwrap(),
409            db: okm[20..40].try_into().unwrap(),
410            kf: okm[40..56].try_into().unwrap(),
411            kb: okm[56..72].try_into().unwrap(),
412        };
413
414        debug!(
415            "ntor handshake completed, derived {} bytes of key material",
416            okm.len()
417        );
418
419        // Return (Y, AUTH, key_material)
420        let mut y_bytes = [0u8; 32];
421        y_bytes.copy_from_slice(big_y.as_bytes());
422
423        let mut auth_bytes = [0u8; 32];
424        auth_bytes.copy_from_slice(&auth);
425
426        Ok((y_bytes, auth_bytes, key_material))
427    }
428}
429
430/// Circuit manager - manages all active circuits
431pub struct CircuitManager {
432    circuits: Arc<Mutex<HashMap<CircuitId, Circuit>>>,
433    ntor_server: Arc<NtorServer>,
434}
435
436impl CircuitManager {
437    /// Create new circuit manager
438    pub fn new() -> Self {
439        Self {
440            circuits: Arc::new(Mutex::new(HashMap::new())),
441            ntor_server: Arc::new(NtorServer::new()),
442        }
443    }
444
445    /// Get server's onion public key
446    pub fn onion_public_key(&self) -> &X25519PublicKey {
447        self.ntor_server.onion_public_key()
448    }
449
450    /// Get server's identity fingerprint
451    pub fn identity_fingerprint(&self) -> [u8; 20] {
452        self.ntor_server.identity_fingerprint()
453    }
454
455    /// Handle CREATE2 cell - perform ntor handshake and create circuit
456    pub async fn handle_create2(
457        &self,
458        circuit_id: CircuitId,
459        client_x: [u8; 32],
460    ) -> Result<([u8; 32], [u8; 32])> {
461        debug!("Processing CREATE2 for circuit {}", circuit_id.as_u32());
462
463        // Perform ntor handshake
464        let (y, auth, key_material) = self.ntor_server.server_handshake(&client_x)?;
465
466        // Create circuit with derived keys
467        let circuit = Circuit::new(circuit_id, key_material);
468
469        // Store circuit
470        let mut circuits = self.circuits.lock().await;
471        circuits.insert(circuit_id, circuit);
472
473        debug!("Circuit {} created successfully", circuit_id.as_u32());
474
475        // Return (Y, AUTH) for CREATED2 cell
476        Ok((y, auth))
477    }
478
479    /// Get circuit by ID (currently unimplemented - circuits accessed via manager methods)
480    pub async fn get_circuit(&self, _id: CircuitId) -> Option<Circuit> {
481        // Clone the circuit (requires implementing Clone for Circuit)
482        // For now, we'll return None and handle circuit access via other methods
483        None
484    }
485
486    /// Decrypt relay cell for circuit
487    pub async fn decrypt_relay_cell(
488        &self,
489        circuit_id: CircuitId,
490        payload: &mut [u8],
491    ) -> Result<()> {
492        let mut circuits = self.circuits.lock().await;
493        let circuit = circuits.get_mut(&circuit_id).context("Circuit not found")?;
494
495        circuit.decrypt_relay_cell(payload)
496    }
497
498    /// Encrypt relay cell for circuit
499    pub async fn encrypt_relay_cell(
500        &self,
501        circuit_id: CircuitId,
502        payload: &mut [u8],
503    ) -> Result<()> {
504        let mut circuits = self.circuits.lock().await;
505        let circuit = circuits.get_mut(&circuit_id).context("Circuit not found")?;
506
507        circuit.encrypt_relay_cell(payload)
508    }
509
510    /// Destroy circuit
511    pub async fn destroy_circuit(&self, id: CircuitId) {
512        let mut circuits = self.circuits.lock().await;
513        circuits.remove(&id);
514        debug!("Circuit {} destroyed", id.as_u32());
515    }
516
517    /// Get active circuit count
518    pub async fn circuit_count(&self) -> usize {
519        let circuits = self.circuits.lock().await;
520        circuits.len()
521    }
522
523    /// Create stream in circuit
524    pub async fn create_stream(
525        &self,
526        circuit_id: CircuitId,
527        stream_id: StreamId,
528        target: String,
529    ) -> Result<()> {
530        let mut circuits = self.circuits.lock().await;
531        let circuit = circuits.get_mut(&circuit_id).context("Circuit not found")?;
532
533        circuit.stream_manager.create_stream(stream_id, target)
534    }
535
536    /// Create directory stream in circuit (BEGIN_DIR)
537    pub async fn create_directory_stream(
538        &self,
539        circuit_id: CircuitId,
540        stream_id: StreamId,
541    ) -> Result<()> {
542        let mut circuits = self.circuits.lock().await;
543        let circuit = circuits.get_mut(&circuit_id).context("Circuit not found")?;
544
545        // Create stream with "directory" as target
546        circuit
547            .stream_manager
548            .create_stream(stream_id, "directory".to_string())?;
549
550        // Set it as directory stream
551        let stream = circuit
552            .stream_manager
553            .get_mut(stream_id)
554            .context("Stream not found")?;
555        stream.set_directory();
556
557        Ok(())
558    }
559
560    /// Set stream as active with TCP connection
561    pub async fn set_stream_active(
562        &self,
563        circuit_id: CircuitId,
564        stream_id: StreamId,
565        connection: tokio::net::TcpStream,
566    ) -> Result<()> {
567        let mut circuits = self.circuits.lock().await;
568        let circuit = circuits.get_mut(&circuit_id).context("Circuit not found")?;
569
570        let stream = circuit
571            .stream_manager
572            .get_mut(stream_id)
573            .context("Stream not found")?;
574
575        stream.set_active(connection);
576        Ok(())
577    }
578
579    /// Write half of a stream's TCP connection (client -> target).
580    pub async fn get_stream_writer(
581        &self,
582        circuit_id: CircuitId,
583        stream_id: StreamId,
584    ) -> Result<Option<Arc<tokio::sync::Mutex<tokio::net::tcp::OwnedWriteHalf>>>> {
585        let circuits = self.circuits.lock().await;
586        let circuit = circuits.get(&circuit_id).context("Circuit not found")?;
587
588        let stream = circuit
589            .stream_manager
590            .get(stream_id)
591            .context("Stream not found")?;
592
593        Ok(stream.writer())
594    }
595
596    /// Read half of a stream's TCP connection (target -> client).
597    ///
598    /// Separate from the write half on purpose: the forwarder task holds this one across a
599    /// blocking `read()`, and sharing a single lock with the write path deadlocked every
600    /// exit stream.
601    pub async fn get_stream_reader(
602        &self,
603        circuit_id: CircuitId,
604        stream_id: StreamId,
605    ) -> Result<Option<Arc<tokio::sync::Mutex<tokio::net::tcp::OwnedReadHalf>>>> {
606        let circuits = self.circuits.lock().await;
607        let circuit = circuits.get(&circuit_id).context("Circuit not found")?;
608
609        let stream = circuit
610            .stream_manager
611            .get(stream_id)
612            .context("Stream not found")?;
613
614        Ok(stream.reader())
615    }
616
617    /// Check if stream is a directory stream (BEGIN_DIR)
618    pub async fn is_directory_stream(
619        &self,
620        circuit_id: CircuitId,
621        stream_id: StreamId,
622    ) -> Result<bool> {
623        let circuits = self.circuits.lock().await;
624        let circuit = circuits.get(&circuit_id).context("Circuit not found")?;
625
626        let stream = circuit
627            .stream_manager
628            .get(stream_id)
629            .context("Stream not found")?;
630
631        Ok(stream.is_directory())
632    }
633
634    /// Append data to directory stream request buffer
635    pub async fn append_directory_data(
636        &self,
637        circuit_id: CircuitId,
638        stream_id: StreamId,
639        data: &[u8],
640    ) -> Result<()> {
641        let mut circuits = self.circuits.lock().await;
642        let circuit = circuits.get_mut(&circuit_id).context("Circuit not found")?;
643
644        let stream = circuit
645            .stream_manager
646            .get_mut(stream_id)
647            .context("Stream not found")?;
648
649        stream.append_request_data(data);
650        Ok(())
651    }
652
653    /// Get accumulated directory request data
654    pub async fn get_directory_request(
655        &self,
656        circuit_id: CircuitId,
657        stream_id: StreamId,
658    ) -> Result<Vec<u8>> {
659        let circuits = self.circuits.lock().await;
660        let circuit = circuits.get(&circuit_id).context("Circuit not found")?;
661
662        let stream = circuit
663            .stream_manager
664            .get(stream_id)
665            .context("Stream not found")?;
666
667        Ok(stream.get_request_data().unwrap_or(&[]).to_vec())
668    }
669
670    /// Close stream
671    pub async fn close_stream(&self, circuit_id: CircuitId, stream_id: StreamId) -> Result<()> {
672        let mut circuits = self.circuits.lock().await;
673        let circuit = circuits.get_mut(&circuit_id).context("Circuit not found")?;
674
675        circuit.stream_manager.remove(stream_id);
676        Ok(())
677    }
678
679    /// Record bytes sent to client for a circuit
680    pub async fn record_sent(&self, circuit_id: CircuitId, bytes: u64) -> Result<()> {
681        let mut circuits = self.circuits.lock().await;
682        if let Some(circuit) = circuits.get_mut(&circuit_id) {
683            circuit.record_sent(bytes);
684        }
685        Ok(())
686    }
687
688    /// Record bytes received from client for a circuit
689    pub async fn record_received(&self, circuit_id: CircuitId, bytes: u64) -> Result<()> {
690        let mut circuits = self.circuits.lock().await;
691        if let Some(circuit) = circuits.get_mut(&circuit_id) {
692            circuit.record_received(bytes);
693        }
694        Ok(())
695    }
696
697    /// Get statistics for a specific circuit
698    pub async fn get_circuit_stats(&self, circuit_id: CircuitId) -> Option<CircuitStats> {
699        let circuits = self.circuits.lock().await;
700        circuits.get(&circuit_id).map(|circuit| CircuitStats {
701            circuit_id: circuit.id,
702            created_at: circuit.created_at,
703            last_activity: circuit.last_activity,
704            bytes_sent: circuit.bytes_sent,
705            bytes_received: circuit.bytes_received,
706            active_streams: circuit.stream_manager.active_streams().len(),
707        })
708    }
709
710    /// Get aggregate relay statistics
711    pub async fn get_relay_stats(&self) -> RelayStats {
712        let circuits = self.circuits.lock().await;
713
714        let mut total_streams = 0;
715        let mut total_bytes_sent = 0;
716        let mut total_bytes_received = 0;
717        let mut circuit_stats = Vec::new();
718
719        for circuit in circuits.values() {
720            let active_streams = circuit.stream_manager.active_streams().len();
721            total_streams += active_streams;
722            total_bytes_sent += circuit.bytes_sent;
723            total_bytes_received += circuit.bytes_received;
724
725            circuit_stats.push(CircuitStats {
726                circuit_id: circuit.id,
727                created_at: circuit.created_at,
728                last_activity: circuit.last_activity,
729                bytes_sent: circuit.bytes_sent,
730                bytes_received: circuit.bytes_received,
731                active_streams,
732            });
733        }
734
735        RelayStats {
736            total_circuits: circuits.len(),
737            total_streams,
738            total_bytes_sent,
739            total_bytes_received,
740            circuit_stats,
741        }
742    }
743
744    /// Get list of all circuit IDs
745    pub async fn list_circuits(&self) -> Vec<CircuitId> {
746        let circuits = self.circuits.lock().await;
747        circuits.keys().copied().collect()
748    }
749
750    /// Record RELAY cell received for circuit - returns true if SENDME needed
751    pub async fn record_relay_received(&self, circuit_id: CircuitId) -> Result<bool> {
752        let mut circuits = self.circuits.lock().await;
753        let circuit = circuits.get_mut(&circuit_id).context("Circuit not found")?;
754        Ok(circuit.record_relay_received())
755    }
756
757    /// Process circuit-level SENDME
758    pub async fn process_circuit_sendme(&self, circuit_id: CircuitId) -> Result<()> {
759        let mut circuits = self.circuits.lock().await;
760        let circuit = circuits.get_mut(&circuit_id).context("Circuit not found")?;
761        circuit.process_circuit_sendme();
762        Ok(())
763    }
764
765    /// Record DATA cell received for stream - returns true if SENDME needed
766    pub async fn record_stream_data_received(
767        &self,
768        circuit_id: CircuitId,
769        stream_id: StreamId,
770    ) -> Result<bool> {
771        let mut circuits = self.circuits.lock().await;
772        let circuit = circuits.get_mut(&circuit_id).context("Circuit not found")?;
773
774        if let Some(stream) = circuit.stream_manager.get_mut(stream_id) {
775            Ok(stream.record_data_received())
776        } else {
777            Ok(false)
778        }
779    }
780
781    /// Process stream-level SENDME
782    pub async fn process_stream_sendme(
783        &self,
784        circuit_id: CircuitId,
785        stream_id: StreamId,
786    ) -> Result<()> {
787        let mut circuits = self.circuits.lock().await;
788        let circuit = circuits.get_mut(&circuit_id).context("Circuit not found")?;
789
790        if let Some(stream) = circuit.stream_manager.get_mut(stream_id) {
791            stream.process_sendme();
792        }
793        Ok(())
794    }
795}