x/internal/util/ss/conn.go
2022-03-14 20:27:14 +08:00

97 lines
1.8 KiB
Go

package ss
import (
"bytes"
"net"
"github.com/go-gost/gosocks5"
"github.com/go-gost/gost/v3/pkg/common/bufpool"
)
var (
DefaultBufferSize = 4096
)
var (
_ net.PacketConn = (*UDPConn)(nil)
_ net.Conn = (*UDPConn)(nil)
)
type UDPConn struct {
net.PacketConn
raddr net.Addr
taddr net.Addr
bufferSize int
}
func UDPClientConn(c net.PacketConn, remoteAddr, targetAddr net.Addr, bufferSize int) *UDPConn {
return &UDPConn{
PacketConn: c,
raddr: remoteAddr,
taddr: targetAddr,
bufferSize: bufferSize,
}
}
func UDPServerConn(c net.PacketConn, remoteAddr net.Addr, bufferSize int) *UDPConn {
return &UDPConn{
PacketConn: c,
raddr: remoteAddr,
bufferSize: bufferSize,
}
}
func (c *UDPConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
rbuf := bufpool.Get(c.bufferSize)
defer bufpool.Put(rbuf)
n, _, err = c.PacketConn.ReadFrom(*rbuf)
if err != nil {
return
}
saddr := gosocks5.Addr{}
addrLen, err := saddr.ReadFrom(bytes.NewReader((*rbuf)[:n]))
if err != nil {
return
}
n = copy(b, (*rbuf)[addrLen:n])
addr, err = net.ResolveUDPAddr("udp", saddr.String())
return
}
func (c *UDPConn) Read(b []byte) (n int, err error) {
n, _, err = c.ReadFrom(b)
return
}
func (c *UDPConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
wbuf := bufpool.Get(c.bufferSize)
defer bufpool.Put(wbuf)
socksAddr := gosocks5.Addr{}
if err = socksAddr.ParseFrom(addr.String()); err != nil {
return
}
addrLen, err := socksAddr.Encode(*wbuf)
if err != nil {
return
}
n = copy((*wbuf)[addrLen:], b)
_, err = c.PacketConn.WriteTo((*wbuf)[:addrLen+n], c.raddr)
return
}
func (c *UDPConn) Write(b []byte) (n int, err error) {
return c.WriteTo(b, c.taddr)
}
func (c *UDPConn) RemoteAddr() net.Addr {
return c.raddr
}