zproxy/proxy/server.go

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"))
	}
}