socks5: keep tcp control connection until udp connection closed

This commit is contained in:
nadoo 2018-01-22 01:10:47 +08:00
parent 24d30cdca7
commit b35da7048d

View File

@ -170,7 +170,7 @@ func (s *SOCKS5) ListenAndServeUDP() {
buf := make([]byte, udpBufSize) buf := make([]byte, udpBufSize)
for { for {
c := NewSocks5PktConn(lc, nil, nil, true) c := NewSocks5PktConn(lc, nil, nil, true, nil)
n, raddr, err := c.ReadFrom(buf) n, raddr, err := c.ReadFrom(buf)
if err != nil { if err != nil {
@ -187,7 +187,7 @@ func (s *SOCKS5) ListenAndServeUDP() {
continue continue
} }
pc = NewSocks5PktConn(lpc, nextHop, nil, false) pc = NewSocks5PktConn(lpc, nextHop, nil, false, nil)
nm.Store(raddr.String(), pc) nm.Store(raddr.String(), pc)
go func() { go func() {
@ -287,7 +287,7 @@ func (s *SOCKS5) DialUDP(network, addr string) (pc net.PacketConn, writeTo net.A
return nil, nil, err return nil, nil, err
} }
pkc := NewSocks5PktConn(pc, nextHop, ParseAddr(addr), true) pkc := NewSocks5PktConn(pc, nextHop, ParseAddr(addr), true, c)
return pkc, nextHop, err return pkc, nextHop, err
} }
@ -602,15 +602,33 @@ type Socks5PktConn struct {
tgtAddr Addr tgtAddr Addr
tgtHeader bool tgtHeader bool
ctrlConn net.Conn // tcp control conn
} }
// NewSocks5PktConn returns a Socks5PktConn // NewSocks5PktConn returns a Socks5PktConn
func NewSocks5PktConn(c net.PacketConn, writeAddr net.Addr, tgtAddr Addr, tgtHeader bool) *Socks5PktConn { func NewSocks5PktConn(c net.PacketConn, writeAddr net.Addr, tgtAddr Addr, tgtHeader bool, ctrlConn net.Conn) *Socks5PktConn {
pc := &Socks5PktConn{ pc := &Socks5PktConn{
PacketConn: c, PacketConn: c,
writeAddr: writeAddr, writeAddr: writeAddr,
tgtAddr: tgtAddr, tgtAddr: tgtAddr,
tgtHeader: tgtHeader} tgtHeader: tgtHeader,
ctrlConn: ctrlConn}
if ctrlConn != nil {
go func() {
buf := []byte{}
for {
_, err := ctrlConn.Read(buf)
if err, ok := err.(net.Error); ok && err.Timeout() {
continue
}
logf("UDP Associate End.")
return
}
}()
}
return pc return pc
} }
@ -658,3 +676,12 @@ func (pc *Socks5PktConn) WriteTo(b []byte, addr net.Addr) (int, error) {
buf = append(buf, b[:]...) buf = append(buf, b[:]...)
return pc.PacketConn.WriteTo(buf, pc.writeAddr) return pc.PacketConn.WriteTo(buf, pc.writeAddr)
} }
// Close .
func (pc *Socks5PktConn) Close() error {
if pc.ctrlConn != nil {
pc.ctrlConn.Close()
}
return pc.PacketConn.Close()
}