diff --git a/internal/server/tunnel/server.go b/internal/server/tunnel/server.go index 3e58130..772dfe8 100644 --- a/internal/server/tunnel/server.go +++ b/internal/server/tunnel/server.go @@ -50,8 +50,8 @@ func generateClientID() string { // Server 隧道服务端 type Server struct { clientStore db.ClientStore - jsPluginStore db.JSPluginStore // JS 插件存储 - trafficStore db.TrafficStore // 流量存储 + jsPluginStore db.JSPluginStore // JS 插件存储 + trafficStore db.TrafficStore // 流量存储 bindAddr string bindPort int token string @@ -514,8 +514,8 @@ func (s *Server) acceptProxyConns(cs *ClientSession, ln net.Listener, rule proto func (s *Server) acceptProxyServerConns(cs *ClientSession, ln net.Listener, rule protocol.ProxyRule) { dialer := proxy.NewTunnelDialer(cs.Session) - // 使用内置 proxy 实现 - proxyServer := proxy.NewServer(rule.Type, dialer) + // 使用内置 proxy 实现 (带流量统计) + proxyServer := proxy.NewServer(rule.Type, dialer, s.recordTraffic) for { conn, err := ln.Accept() if err != nil { @@ -963,6 +963,9 @@ func (s *Server) sendUDPPacket(cs *ClientSession, conn *net.UDPConn, clientAddr return } + // 记录入站流量 (从外部接收的数据) + s.recordTraffic(int64(len(packet.Data)), 0) + // 等待客户端响应 respMsg, err := protocol.ReadMessage(stream) if err != nil { @@ -975,6 +978,8 @@ func (s *Server) sendUDPPacket(cs *ClientSession, conn *net.UDPConn, clientAddr return } conn.WriteToUDP(respPacket.Data, clientAddr) + // 记录出站流量 (发送回外部的数据) + s.recordTraffic(0, int64(len(respPacket.Data))) } } diff --git a/pkg/proxy/http.go b/pkg/proxy/http.go index b5e550b..4817626 100644 --- a/pkg/proxy/http.go +++ b/pkg/proxy/http.go @@ -6,16 +6,19 @@ import ( "net" "net/http" "strings" + + "github.com/gotunnel/pkg/relay" ) // HTTPServer HTTP 代理服务 type HTTPServer struct { - dialer Dialer + dialer Dialer + onStats func(in, out int64) // 流量统计回调 } // NewHTTPServer 创建 HTTP 代理服务 -func NewHTTPServer(dialer Dialer) *HTTPServer { - return &HTTPServer{dialer: dialer} +func NewHTTPServer(dialer Dialer, onStats func(in, out int64)) *HTTPServer { + return &HTTPServer{dialer: dialer, onStats: onStats} } // HandleConn 处理 HTTP 代理连接 @@ -50,8 +53,8 @@ func (h *HTTPServer) handleConnect(conn net.Conn, req *http.Request) error { conn.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n")) - go io.Copy(remote, conn) - io.Copy(conn, remote) + // 双向转发 (带流量统计) + relay.RelayWithStats(conn, remote, h.onStats) return nil } @@ -82,7 +85,10 @@ func (h *HTTPServer) handleHTTP(conn net.Conn, req *http.Request, reader *bufio. return err } - // 转发响应 - _, err = io.Copy(conn, remote) + // 转发响应 (带流量统计) + n, err := io.Copy(conn, remote) + if h.onStats != nil && n > 0 { + h.onStats(0, n) // 响应数据为出站流量 + } return err } diff --git a/pkg/proxy/server.go b/pkg/proxy/server.go index 7ec37bc..611bc8d 100644 --- a/pkg/proxy/server.go +++ b/pkg/proxy/server.go @@ -14,10 +14,10 @@ type Server struct { } // NewServer 创建代理服务器 -func NewServer(typ string, dialer Dialer) *Server { +func NewServer(typ string, dialer Dialer, onStats func(in, out int64)) *Server { return &Server{ - socks5: NewSOCKS5Server(dialer), - http: NewHTTPServer(dialer), + socks5: NewSOCKS5Server(dialer, onStats), + http: NewHTTPServer(dialer, onStats), typ: typ, } } diff --git a/pkg/proxy/socks5.go b/pkg/proxy/socks5.go index 202ad7b..cc4fa01 100644 --- a/pkg/proxy/socks5.go +++ b/pkg/proxy/socks5.go @@ -6,6 +6,8 @@ import ( "fmt" "io" "net" + + "github.com/gotunnel/pkg/relay" ) const ( @@ -19,7 +21,8 @@ const ( // SOCKS5Server SOCKS5 代理服务 type SOCKS5Server struct { - dialer Dialer + dialer Dialer + onStats func(in, out int64) // 流量统计回调 } // Dialer 连接拨号器接口 @@ -28,8 +31,8 @@ type Dialer interface { } // NewSOCKS5Server 创建 SOCKS5 服务 -func NewSOCKS5Server(dialer Dialer) *SOCKS5Server { - return &SOCKS5Server{dialer: dialer} +func NewSOCKS5Server(dialer Dialer, onStats func(in, out int64)) *SOCKS5Server { + return &SOCKS5Server{dialer: dialer, onStats: onStats} } // HandleConn 处理 SOCKS5 连接 @@ -60,9 +63,8 @@ func (s *SOCKS5Server) HandleConn(conn net.Conn) error { return err } - // 双向转发 - go io.Copy(remote, conn) - io.Copy(conn, remote) + // 双向转发 (带流量统计) + relay.RelayWithStats(conn, remote, s.onStats) return nil }