From 9a6f28c0bc7ed29dfbf448958dfd2fa68c8e7c84 Mon Sep 17 00:00:00 2001
From: eric <eric.y.corsican@gmail.com>
Date: Sun, 28 May 2023 13:24:42 +0800
Subject: [PATCH] Logs session end with total bytes transfered

---
 leaf/src/app/dispatcher.rs   | 30 ++++++++++++++++--------------
 leaf/src/app/stat_manager.rs | 19 ++++++++++++++++++-
 2 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/leaf/src/app/dispatcher.rs b/leaf/src/app/dispatcher.rs
index d499f72..44ffffd 100644
--- a/leaf/src/app/dispatcher.rs
+++ b/leaf/src/app/dispatcher.rs
@@ -30,26 +30,28 @@ fn log_request(
     handshake_time: Option<u128>,
 ) {
     let hs = handshake_time.map_or("failed".to_string(), |hs| format!("{}ms", hs));
-    if !*crate::option::LOG_NO_COLOR {
+    let (network, outbound_tag) = if !*crate::option::LOG_NO_COLOR {
         use colored::Colorize;
         let network_color = match sess.network {
             Network::Tcp => colored::Color::Blue,
             Network::Udp => colored::Color::Yellow,
         };
-        info!(
-            "[{}] [{}] [{}] [{}] {}",
-            &sess.inbound_tag,
-            sess.network.to_string().color(network_color),
-            outbound_tag.color(*outbound_tag_color),
-            hs,
-            &sess.destination,
-        );
+        (
+            sess.network.to_string().color(network_color).to_string(),
+            outbound_tag.color(*outbound_tag_color).to_string(),
+        )
     } else {
-        info!(
-            "[{}] [{}] [{}] [{}] {}",
-            sess.network, &sess.inbound_tag, outbound_tag, hs, &sess.destination,
-        );
-    }
+        (sess.network.to_string(), outbound_tag.to_string())
+    };
+    info!(
+        "[{}] [{}] [{}] [{}] [{}] [{}]",
+        sess.forwarded_source.unwrap_or_else(|| sess.source.ip()),
+        network,
+        &sess.inbound_tag,
+        outbound_tag,
+        hs,
+        &sess.destination,
+    );
 }
 
 pub struct Dispatcher {
diff --git a/leaf/src/app/stat_manager.rs b/leaf/src/app/stat_manager.rs
index 1a8c84c..dca3f6a 100644
--- a/leaf/src/app/stat_manager.rs
+++ b/leaf/src/app/stat_manager.rs
@@ -163,6 +163,22 @@ impl Counter {
     }
 }
 
+#[inline]
+fn log_session_end(c: &Counter) {
+    log::info!(
+        "[{}] [{}] [{}] [{}] [{}] [{}] [{}] [END]",
+        c.sess
+            .forwarded_source
+            .unwrap_or_else(|| c.sess.source.ip()),
+        c.sess.network,
+        c.sess.inbound_tag,
+        c.sess.outbound_tag,
+        c.sess.destination,
+        c.bytes_sent(),
+        c.bytes_recvd(),
+    );
+}
+
 pub struct StatManager {
     pub counters: Vec<Counter>,
 }
@@ -182,7 +198,8 @@ impl StatManager {
                 let mut i = 0;
                 while i < sm.counters.len() {
                     if sm.counters[i].recv_completed() && sm.counters[i].send_completed() {
-                        sm.counters.swap_remove(i);
+                        let c = sm.counters.swap_remove(i);
+                        log_session_end(&c);
                     } else {
                         i += 1;
                     }