zproxy/proxy/server.go

package proxy

import (
	"net"
	"time"
	"zproxy/zlog"
)

type UServerProxy struct {
	RemoteMap    map[uint32]*PacketConnect
	LocalMap     map[uint32]*PacketConnect
	RemoteConn   *PacketConnect
	MessageQueue chan uint32
}

func NewServerProxy() *UServerProxy {
	return &UServerProxy{
		RemoteMap:    make(map[uint32]*PacketConnect),
		LocalMap:     make(map[uint32]*PacketConnect),
		MessageQueue: make(chan uint32, 32),
	}
}
func (server *UServerProxy) NewTcpConnection(remote *PacketConnect) {
	data := remote.ReadHeader()
	if string(data[:4]) == "anki" {
		id := packetEndian.Uint32(data[4:8])
		remote.SetName("remote", id)
		if id == RemoteID {
			if server.RemoteConn != nil && server.RemoteConn != remote {
				_ = server.RemoteConn.Close()
			}
			zlog.Infof("rcv remote %s => %s", remote.Name(), remote.LocalAddr())
			server.RemoteConn = remote
		} else {
			server.RemoteMap[id] = remote
			server.MessageQueue <- id
		}
		return
	}
	if server.RemoteConn == nil {
		_ = remote.Close()
		return
	}
	id := IncID()
	server.RemoteConn.WriteHeader(id)
	if server.RemoteConn.writeErr != nil {
		zlog.Infof("remote disconnect %s", server.RemoteConn.RemoteAddr())
		_ = server.RemoteConn.Close()
		server.RemoteConn = nil
		return
	}
	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]
			local := server.LocalMap[mid]
			if local == nil || remote == nil {
				continue
			}
			go func() {
				local.Forward(remote, HeaderSize)
				delete(server.LocalMap, mid)
			}()
			go func() {
				remote.Forward(local, 0)
				local.SetReadDeadline(time.Now())
				delete(server.RemoteMap, mid)
			}()
		}
	}
}
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) {
		_ = ln.Close()
	}(ln)
	go server.MessageLoop()
	for {
		conn, err := ln.Accept()
		if err != nil {
			return err
		}
		go server.NewTcpConnection(NewPacketConnect(conn, "client", 0))
	}
}