From 3e86d1f37dd0314955c308f8f65f9d502d97d410 Mon Sep 17 00:00:00 2001
From: ouczbs <ouczbs@qq.com>
Date: Fri, 16 Feb 2024 00:07:57 +0800
Subject: [PATCH] bugfix

---
 proxy/client.go |  3 ++-
 proxy/server.go | 38 ++++++++++++++++++++------------------
 2 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/proxy/client.go b/proxy/client.go
index 7e98bce..42a5eaa 100644
--- a/proxy/client.go
+++ b/proxy/client.go
@@ -2,6 +2,7 @@ package proxy
 
 import (
 	"net"
+	"time"
 	"zproxy/zlog"
 )
 
@@ -37,7 +38,7 @@ func (client *UClientProxy) NewTcpConnection(id uint32) *PacketConnect {
 	zlog.Infof("Forward %s => %s", remote.LocalAddr(), remote.Name())
 	go func() {
 		local.Forward(remote, 0)
-		//remote.SetReadDeadline(time.Now())
+		remote.SetReadDeadline(time.Now())
 	}()
 	go remote.Forward(local, 0)
 	return nil
diff --git a/proxy/server.go b/proxy/server.go
index f94f1fb..5760419 100644
--- a/proxy/server.go
+++ b/proxy/server.go
@@ -2,12 +2,13 @@ package proxy
 
 import (
 	"net"
+	"time"
 	"zproxy/zlog"
 )
 
 type UServerProxy struct {
 	RemoteMap    map[uint32]*PacketConnect
-	ClientMap    map[uint32]*PacketConnect
+	LocalMap     map[uint32]*PacketConnect
 	RemoteConn   *PacketConnect
 	MessageQueue chan uint32
 }
@@ -15,29 +16,29 @@ type UServerProxy struct {
 func NewServerProxy() *UServerProxy {
 	return &UServerProxy{
 		RemoteMap:    make(map[uint32]*PacketConnect),
-		ClientMap:    make(map[uint32]*PacketConnect),
+		LocalMap:     make(map[uint32]*PacketConnect),
 		MessageQueue: make(chan uint32, 32),
 	}
 }
-func (server *UServerProxy) NewTcpConnection(client *PacketConnect) {
-	data := client.ReadHeader()
+func (server *UServerProxy) NewTcpConnection(remote *PacketConnect) {
+	data := remote.ReadHeader()
 	if string(data[:4]) == "anki" {
 		id := packetEndian.Uint32(data[4:8])
-		client.SetName("remote", id)
+		remote.SetName("remote", id)
 		if id == RemoteID {
-			if server.RemoteConn != nil && server.RemoteConn != client {
+			if server.RemoteConn != nil && server.RemoteConn != remote {
 				_ = server.RemoteConn.Close()
 			}
-			zlog.Infof("rcv remote %s => %s", client.Name(), client.LocalAddr())
-			server.RemoteConn = client
+			zlog.Infof("rcv remote %s => %s", remote.Name(), remote.LocalAddr())
+			server.RemoteConn = remote
 		} else {
-			server.RemoteMap[id] = client
+			server.RemoteMap[id] = remote
 			server.MessageQueue <- id
 		}
 		return
 	}
 	if server.RemoteConn == nil {
-		_ = client.Close()
+		_ = remote.Close()
 		return
 	}
 	id := IncID()
@@ -48,25 +49,26 @@ func (server *UServerProxy) NewTcpConnection(client *PacketConnect) {
 		server.RemoteConn = nil
 		return
 	}
-	server.ClientMap[id] = client
-	client.SetName("local", id)
-	zlog.Infof("%d %d %s %s", len(server.RemoteMap), len(server.ClientMap), data, client.Name())
+	server.LocalMap[id] = remote
+	remote.SetName("remote", id)
+	zlog.Infof("%d %d %s %s", len(server.RemoteMap), len(server.LocalMap), data, remote.Name())
 }
 func (server *UServerProxy) MessageLoop() {
 	for {
 		select {
 		case mid := <-server.MessageQueue:
 			remote := server.RemoteMap[mid]
-			client := server.ClientMap[mid]
-			if client == nil || remote == nil {
+			local := server.LocalMap[mid]
+			if local == nil || remote == nil {
 				continue
 			}
 			go func() {
-				client.Forward(remote, HeaderSize)
-				delete(server.ClientMap, mid)
+				local.Forward(remote, HeaderSize)
+				delete(server.LocalMap, mid)
 			}()
 			go func() {
-				remote.Forward(client, 0)
+				remote.Forward(local, 0)
+				local.SetReadDeadline(time.Now())
 				delete(server.RemoteMap, mid)
 			}()
 		}