Fixed critical ME Problems

src/transport/middle_proxy/registry.rs

@@ -1,34 +1,57 @@
 use std::collections::HashMap;
+use std::net::SocketAddr;
 use std::sync::atomic::{AtomicU64, Ordering};
-
-use tokio::sync::{RwLock, mpsc};
-
-use super::MeResponse;
-use super::codec::RpcWriter;
 use std::sync::Arc;
-use tokio::sync::Mutex;
+
+use tokio::sync::{mpsc, Mutex, RwLock};
+
+use super::codec::RpcWriter;
+use super::MeResponse;
+
+#[derive(Clone)]
+pub struct ConnMeta {
+    pub target_dc: i16,
+    pub client_addr: SocketAddr,
+    pub our_addr: SocketAddr,
+    pub proto_flags: u32,
+}
+
+#[derive(Clone)]
+pub struct BoundConn {
+    pub conn_id: u64,
+    pub meta: ConnMeta,
+}
+
+#[derive(Clone)]
+pub struct ConnWriter {
+    pub writer_id: u64,
+    pub writer: Arc<Mutex<RpcWriter>>,
+}
 
 pub struct ConnRegistry {
     map: RwLock<HashMap<u64, mpsc::UnboundedSender<MeResponse>>>,
     writers: RwLock<HashMap<u64, Arc<Mutex<RpcWriter>>>>,
+    writer_for_conn: RwLock<HashMap<u64, u64>>,
+    conns_for_writer: RwLock<HashMap<u64, Vec<u64>>>,
+    meta: RwLock<HashMap<u64, ConnMeta>>,
     next_id: AtomicU64,
 }
 
 impl ConnRegistry {
     pub fn new() -> Self {
-        // Avoid fully predictable conn_id sequence from 1.
         let start = rand::random::<u64>() | 1;
         Self {
             map: RwLock::new(HashMap::new()),
             writers: RwLock::new(HashMap::new()),
+            writer_for_conn: RwLock::new(HashMap::new()),
+            conns_for_writer: RwLock::new(HashMap::new()),
+            meta: RwLock::new(HashMap::new()),
             next_id: AtomicU64::new(start),
         }
     }
 
     pub async fn register(&self) -> (u64, mpsc::UnboundedReceiver<MeResponse>) {
         let id = self.next_id.fetch_add(1, Ordering::Relaxed);
-        // Unbounded per-connection queue prevents reader-loop HOL blocking on
-        // slow clients: routing stays non-blocking and preserves message order.
         let (tx, rx) = mpsc::unbounded_channel();
         self.map.write().await.insert(id, tx);
         (id, rx)
@@ -36,7 +59,12 @@ impl ConnRegistry {
 
     pub async fn unregister(&self, id: u64) {
         self.map.write().await.remove(&id);
-        self.writers.write().await.remove(&id);
+        self.meta.write().await.remove(&id);
+        if let Some(writer_id) = self.writer_for_conn.write().await.remove(&id) {
+            if let Some(list) = self.conns_for_writer.write().await.get_mut(&writer_id) {
+                list.retain(|c| *c != id);
+            }
+        }
     }
 
     pub async fn route(&self, id: u64, resp: MeResponse) -> bool {
@@ -48,13 +76,58 @@ impl ConnRegistry {
         }
     }
 
-    pub async fn set_writer(&self, id: u64, w: Arc<Mutex<RpcWriter>>) {
-        let mut guard = self.writers.write().await;
-        guard.entry(id).or_insert_with(|| w);
+    pub async fn bind_writer(
+        &self,
+        conn_id: u64,
+        writer_id: u64,
+        writer: Arc<Mutex<RpcWriter>>,
+        meta: ConnMeta,
+    ) {
+        self.meta.write().await.entry(conn_id).or_insert(meta);
+        self.writer_for_conn.write().await.insert(conn_id, writer_id);
+        self.writers.write().await.entry(writer_id).or_insert_with(|| writer.clone());
+        self.conns_for_writer
+            .write()
+            .await
+            .entry(writer_id)
+            .or_insert_with(Vec::new)
+            .push(conn_id);
     }
 
-    pub async fn get_writer(&self, id: u64) -> Option<Arc<Mutex<RpcWriter>>> {
-        let guard = self.writers.read().await;
-        guard.get(&id).cloned()
+    pub async fn get_writer(&self, conn_id: u64) -> Option<ConnWriter> {
+        let writer_id = {
+            let guard = self.writer_for_conn.read().await;
+            guard.get(&conn_id).cloned()
+        }?;
+        let writer = {
+            let guard = self.writers.read().await;
+            guard.get(&writer_id).cloned()
+        }?;
+        Some(ConnWriter { writer_id, writer })
+    }
+
+    pub async fn writer_lost(&self, writer_id: u64) -> Vec<BoundConn> {
+        self.writers.write().await.remove(&writer_id);
+        let conns = self.conns_for_writer.write().await.remove(&writer_id).unwrap_or_default();
+
+        let mut out = Vec::new();
+        let mut writer_for_conn = self.writer_for_conn.write().await;
+        let meta = self.meta.read().await;
+
+        for conn_id in conns {
+            writer_for_conn.remove(&conn_id);
+            if let Some(m) = meta.get(&conn_id) {
+                out.push(BoundConn {
+                    conn_id,
+                    meta: m.clone(),
+                });
+            }
+        }
+        out
+    }
+
+    pub async fn get_meta(&self, conn_id: u64) -> Option<ConnMeta> {
+        let guard = self.meta.read().await;
+        guard.get(&conn_id).cloned()
     }
 }