Merge pull request #200 from telemt/flow

ME Connection lost fixes

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "telemt"
-version = "3.0.7"
+version = "3.0.8"
 edition = "2024"
 
 [dependencies]

src/main.rs

@@ -461,10 +461,12 @@ match crate::transport::middle_proxy::fetch_proxy_secret(proxy_secret_path).awai
         cache.load_from_disk().await;
 
         let port = config.censorship.mask_port;
+        let mask_host = config.censorship.mask_host.clone()
+            .unwrap_or_else(|| config.censorship.tls_domain.clone());
         // Initial synchronous fetch to warm cache before serving clients.
         for domain in tls_domains.clone() {
             match crate::tls_front::fetcher::fetch_real_tls(
-                &domain,
+                &mask_host,
                 port,
                 &domain,
                 Duration::from_secs(5),
@@ -488,7 +490,7 @@ match crate::transport::middle_proxy::fetch_proxy_secret(proxy_secret_path).awai
                 tokio::time::sleep(Duration::from_secs(base_secs + jitter_secs)).await;
                 for domain in &domains {
                     match crate::tls_front::fetcher::fetch_real_tls(
-                        domain,
+                        &mask_host,
                         port,
                         domain,
                         Duration::from_secs(5),
@@ -938,7 +940,40 @@ match crate::transport::middle_proxy::fetch_proxy_secret(proxy_secret_path).awai
                             .run()
                             .await
                             {
-                                warn!(peer = %peer_addr, error = %e, "Connection closed with error");
+                                let peer_closed = matches!(
+                                    &e,
+                                    crate::error::ProxyError::Io(ioe)
+                                        if matches!(
+                                            ioe.kind(),
+                                            std::io::ErrorKind::ConnectionReset
+                                                | std::io::ErrorKind::ConnectionAborted
+                                                | std::io::ErrorKind::BrokenPipe
+                                                | std::io::ErrorKind::NotConnected
+                                        )
+                                ) || matches!(
+                                    &e,
+                                    crate::error::ProxyError::Stream(
+                                        crate::error::StreamError::Io(ioe)
+                                    )
+                                        if matches!(
+                                            ioe.kind(),
+                                            std::io::ErrorKind::ConnectionReset
+                                                | std::io::ErrorKind::ConnectionAborted
+                                                | std::io::ErrorKind::BrokenPipe
+                                                | std::io::ErrorKind::NotConnected
+                                        )
+                                );
+
+                                let me_closed = matches!(
+                                    &e,
+                                    crate::error::ProxyError::Proxy(msg) if msg == "ME connection lost"
+                                );
+
+                                match (peer_closed, me_closed) {
+                                    (true, _) => debug!(peer = %peer_addr, error = %e, "Connection closed by client"),
+                                    (_, true) => warn!(peer = %peer_addr, error = %e, "Connection closed: Middle-End dropped session"),
+                                    _ => warn!(peer = %peer_addr, error = %e, "Connection closed with error"),
+                                }
                             }
                         });
                     }

src/proxy/middle_relay.rs

@@ -98,6 +98,7 @@ where
     });
 
     let mut main_result: Result<()> = Ok(());
+    let mut client_closed = false;
     loop {
         match read_client_payload(&mut crypto_reader, proto_tag, frame_limit, &user).await {
             Ok(Some((payload, quickack))) => {
@@ -124,6 +125,7 @@ where
             }
             Ok(None) => {
                 debug!(conn_id, "Client EOF");
+                client_closed = true;
                 let _ = me_pool.send_close(conn_id).await;
                 break;
             }
@@ -135,7 +137,19 @@ where
     }
 
     let _ = stop_tx.send(());
-    let writer_result = me_writer.await.unwrap_or_else(|e| Err(ProxyError::Proxy(format!("ME writer join error: {e}"))));
+    let mut writer_result = me_writer
+        .await
+        .unwrap_or_else(|e| Err(ProxyError::Proxy(format!("ME writer join error: {e}"))));
+
+    // When client closes, but ME channel stopped as unregistered - it isnt error
+    if client_closed {
+        if matches!(
+            writer_result,
+            Err(ProxyError::Proxy(ref msg)) if msg == "ME connection lost"
+        ) {
+            writer_result = Ok(());
+        }
+    }
 
     let result = match (main_result, writer_result) {
         (Ok(()), Ok(())) => Ok(()),