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.SetName("remote", id)
		if id == RemoteID {
			if server.RemoteConn != nil && server.RemoteConn != client {
				_ = server.RemoteConn.Close()
			}
			zlog.Infof("rcv remote %s => %s", client.Name(), client.LocalAddr())
			server.RemoteConn = client
		} else {
			server.RemoteMap[id] = client
			server.MessageQueue <- id
		}
		return
	}
	if server.RemoteConn == nil {
		_ = client.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.ClientMap[id] = client
	client.SetName("local", id)
	zlog.Infof("%d %d %s %s", len(server.RemoteMap), len(server.ClientMap), data, client.Name())
}
func (server *UServerProxy) MessageLoop() {
	for {
		select {
		case mid := <-server.MessageQueue:
			remote := server.RemoteMap[mid]
			client := server.ClientMap[mid]
			if client == nil || remote == nil {
				continue
			}
			go func() {
				client.Forward(remote, HeaderSize)
				delete(server.ClientMap, mid)
			}()
			go func() {
				remote.Forward(client, 0)
				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))
	}
}