Better handling TCP handshake for UDPAssociate

socks/socks.go

@@ -41,6 +41,7 @@ const (
 	ErrTTLExpired           = Error(6)
 	ErrCommandNotSupported  = Error(7)
 	ErrAddressNotSupported  = Error(8)
+	InfoUDPAssociate        = Error(9)
 )
 
 // MaxAddrLen is the maximum size of SOCKS address in bytes.
@@ -191,7 +192,6 @@ func Handshake(rw io.ReadWriter) (Addr, error) {
 	if err != nil {
 		return nil, err
 	}
-	// TODO: when disconnected, release the natmap
 	switch buf[1] {
 	case CmdConnect:
 		_, err = rw.Write([]byte{5, 0, 0, 1, 0, 0, 0, 0, 0, 0}) // SOCKS v5, reply succeeded
@@ -201,6 +201,10 @@ func Handshake(rw io.ReadWriter) (Addr, error) {
 		}
 		listenAddr := ParseAddr(rw.(net.Conn).LocalAddr().String())
 		_, err = rw.Write(append([]byte{5, 0, 0}, listenAddr...)) // SOCKS v5, reply succeeded
+		if err != nil {
+			return nil, ErrCommandNotSupported
+		}
+		err = InfoUDPAssociate
 	default:
 		return nil, ErrCommandNotSupported
 	}

tcp.go

@@ -43,9 +43,23 @@ func tcpLocal(addr, server string, shadow func(net.Conn) net.Conn, getAddr func(
 		go func() {
 			defer c.Close()
 			c.(*net.TCPConn).SetKeepAlive(true)
-			// TODO: keep the connection until disconnect then free the UDP socket
 			tgt, err := getAddr(c)
 			if err != nil {
+
+				// UDP: keep the connection until disconnect then free the UDP socket
+				if err == socks.InfoUDPAssociate {
+					buf := []byte{}
+					// block here
+					for {
+						_, err := c.Read(buf)
+						if err, ok := err.(net.Error); ok && err.Timeout() {
+							continue
+						}
+						logf("UDP Associate End.")
+						return
+					}
+				}
+
 				logf("failed to get target address: %v", err)
 				return
 			}