This commit is contained in:
ouczbs 2024-02-16 00:07:57 +08:00
parent 48c1dfc8c5
commit 3e86d1f37d
2 changed files with 22 additions and 19 deletions

View File

@ -2,6 +2,7 @@ package proxy
import ( import (
"net" "net"
"time"
"zproxy/zlog" "zproxy/zlog"
) )
@ -37,7 +38,7 @@ func (client *UClientProxy) NewTcpConnection(id uint32) *PacketConnect {
zlog.Infof("Forward %s => %s", remote.LocalAddr(), remote.Name()) zlog.Infof("Forward %s => %s", remote.LocalAddr(), remote.Name())
go func() { go func() {
local.Forward(remote, 0) local.Forward(remote, 0)
//remote.SetReadDeadline(time.Now()) remote.SetReadDeadline(time.Now())
}() }()
go remote.Forward(local, 0) go remote.Forward(local, 0)
return nil return nil

View File

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