From 5111aed35d8651469826116fe6d38f5357ece6aa Mon Sep 17 00:00:00 2001
From: dingjun <fangdingjun@gmail.com>
Date: Tue, 16 Mar 2021 14:31:37 +0800
Subject: [PATCH] add log

---
 server.go     |  4 ++++
 tcp_server.go |  4 ++++
 ws_server.go  | 13 ++++++++++++-
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/server.go b/server.go
index 7c31ca8..b04a7d4 100644
--- a/server.go
+++ b/server.go
@@ -50,6 +50,10 @@ func forwardWS2WS(conn1, conn2 *websocket.Conn) {
 func forwardWS2TCP(conn1 *websocket.Conn, conn2 net.Conn) {
 	ch := make(chan struct{}, 2)
 
+	defer func() {
+		log.Debugf("forward to %s finished", conn2.RemoteAddr())
+	}()
+
 	go func() {
 		for {
 			_, data, err := conn1.ReadMessage()
diff --git a/tcp_server.go b/tcp_server.go
index 946c425..0693ba1 100644
--- a/tcp_server.go
+++ b/tcp_server.go
@@ -38,6 +38,10 @@ func (srv *tcpServer) serve(c net.Conn) {
 
 	log.Debugf("connected from %s, forward to %s", c.RemoteAddr(), srv.remote)
 
+	defer func() {
+		log.Debugf("from %s, finished", c.RemoteAddr())
+	}()
+
 	if u.Scheme == "ws" || u.Scheme == "wss" {
 		conn1, resp, err := dialer.Dial(srv.remote, nil)
 		if err != nil {
diff --git a/ws_server.go b/ws_server.go
index 7d81f6b..181d3f1 100644
--- a/ws_server.go
+++ b/ws_server.go
@@ -46,7 +46,18 @@ func (wss *wsServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	log.Debugf("request %s, forward to %s", p, remote)
+	ip := r.RemoteAddr
+
+	_ip := r.Header.Get("x-real-ip")
+	if _ip != "" {
+		ip = _ip
+	}
+
+	log.Debugf("from %s, request %s, forward to %s", ip, p, remote)
+
+	defer func() {
+		log.Debugf("from %s, request finished", ip)
+	}()
 
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {