zproxy/proxy/packet_connect.go
2024-02-15 23:48:07 +08:00

77 lines
1.7 KiB
Go

package proxy
import (
"net"
"strconv"
"time"
"zproxy/zlog"
)
type PacketConnect struct {
net.Conn
header []byte
id uint32
readErr error
writeErr error
name string
}
func NewPacketConnect(conn net.Conn, name string, id uint32) *PacketConnect {
pc := &PacketConnect{Conn: conn, header: make([]byte, HeaderSize)}
pc.SetName(name, id)
copy(pc.header, InitBuf)
return pc
}
func (pc *PacketConnect) Name() string {
return pc.name
}
func (pc *PacketConnect) SetName(name string, id uint32) {
pc.id = id
pc.name = "(" + name + ")" + strconv.Itoa(int(id)) + "::" + pc.RemoteAddr().String()
}
func (pc *PacketConnect) Write(data []byte) int {
n, err := pc.Conn.Write(data)
pc.writeErr = err
return n
}
func (pc *PacketConnect) Read(data []byte) int {
n, err := pc.Conn.Read(data)
pc.readErr = err
return n
}
func (pc *PacketConnect) SetReadDeadline(t time.Time) {
err := pc.Conn.SetReadDeadline(t)
if err != nil {
pc.readErr = err
}
}
func (pc *PacketConnect) ReadHeader() []byte {
pc.Read(pc.header)
return pc.header
}
func (pc *PacketConnect) WriteHeader(id uint32) {
packetEndian.PutUint32(pc.header[4:8], id)
pc.Write(pc.header)
}
func (pc *PacketConnect) Forward(dst *PacketConnect, rn int) {
if rn > 0 {
dst.Write(pc.header[:rn])
}
defer func() {
_ = pc.Close()
if Level > 0 {
zlog.Infof("forward %d %s %s => %s %s || %s", rn, pc.LocalAddr(), pc.Name(), dst.Name(), pc.readErr, dst.writeErr)
}
}()
data := make([]byte, MTU)
for pc.readErr == nil && dst.writeErr == nil {
n := pc.Read(data)
n = dst.Write(data[:n])
rn += n
if Level > 1 {
zlog.Infof("rw %d %s => %s %s || %s", n, pc.Name(), dst.Name(), pc.readErr, dst.writeErr)
}
}
}