package proxy import ( "net" "zproxy/zlog" ) type UServerProxy struct { RemoteMap map[uint32]*PacketConnect ClientMap map[uint32]*PacketConnect RemoteConn *PacketConnect MessageQueue chan uint32 } func NewServerProxy() *UServerProxy { return &UServerProxy{ RemoteMap: make(map[uint32]*PacketConnect), ClientMap: make(map[uint32]*PacketConnect), MessageQueue: make(chan uint32, 32), } } func (server *UServerProxy) NewTcpConnection(client *PacketConnect) { data := client.ReadHeader() if string(data[:4]) == "anki" { id := packetEndian.Uint32(data[4:8]) client.SetRemote("remote", id) if id == RemoteID { if server.RemoteConn != nil && server.RemoteConn != client { _ = server.RemoteConn.Close() } server.RemoteConn = client zlog.Infof("rcv %d %s", id, client.Name()) } else { server.RemoteMap[id] = client server.MessageQueue <- id } return } if server.RemoteConn == nil { return } id := IncID() server.RemoteConn.WriteHeader(id) server.ClientMap[id] = client zlog.Infof("%d %s", id, data) } func (server *UServerProxy) MessageLoop() { for { select { case mid := <-server.MessageQueue: remote := server.RemoteMap[mid] client := server.ClientMap[mid] if client == nil || remote == nil { server.RemoteMap[mid] = nil server.ClientMap[mid] = nil continue } go func() { client.Forward(remote, HeaderSize) server.ClientMap[mid] = nil }() go func() { remote.Forward(client, 0) server.RemoteMap[mid] = nil }() } } } func (server *UServerProxy) ServeTCP(listenAddr string) error { // listenAddr 只能填端口号 ln, err := net.Listen("tcp", listenAddr) if err != nil { zlog.Error("listen error:", err.Error()) } zlog.Infof("listen %s", ln.Addr()) defer func(ln net.Listener) { err := ln.Close() if err != nil { zlog.Error("close client error: ", err.Error()) } }(ln) go server.MessageLoop() for { conn, err := ln.Accept() if err != nil { return err } server.NewTcpConnection(NewPacketConnect(conn, "client")) } }