88 lines
1.7 KiB
Go
88 lines
1.7 KiB
Go
package proxy
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"net"
|
|
"strconv"
|
|
"time"
|
|
)
|
|
|
|
var (
|
|
packetEndian = binary.LittleEndian
|
|
RemoteID uint32 = 1
|
|
sid = RemoteID
|
|
maxSid uint32 = 1024 * 1024
|
|
HeaderSize = 16
|
|
Magic = "anki"
|
|
InitBuf = func() []byte {
|
|
buf := make([]byte, HeaderSize)
|
|
copy(buf[:len(Magic)], Magic)
|
|
return buf
|
|
}()
|
|
)
|
|
|
|
func IncID() uint32 {
|
|
if sid > maxSid {
|
|
sid = RemoteID
|
|
}
|
|
sid++
|
|
return sid
|
|
}
|
|
|
|
type PacketConnect struct {
|
|
net.Conn
|
|
header []byte
|
|
id uint32
|
|
readErr error
|
|
writeErr error
|
|
name string
|
|
}
|
|
|
|
func NewPacketConnect(conn net.Conn, name string) *PacketConnect {
|
|
pc := &PacketConnect{Conn: conn, header: make([]byte, HeaderSize)}
|
|
pc.SetRemote(name, 0)
|
|
copy(pc.header, InitBuf)
|
|
return pc
|
|
}
|
|
func (pc *PacketConnect) Name() string {
|
|
return pc.name
|
|
}
|
|
func (pc *PacketConnect) SetRemote(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) 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()
|
|
}()
|
|
data := make([]byte, 1024)
|
|
for pc.readErr == nil && dst.writeErr == nil {
|
|
_ = pc.SetReadDeadline(time.Now().Add(time.Second))
|
|
n := pc.Read(data)
|
|
rn += n
|
|
dst.Write(data[:n])
|
|
}
|
|
}
|