From 66fbc0399ac93027655674ca72100cd72652a43d Mon Sep 17 00:00:00 2001
From: mo <thinkgo@aliyun.com>
Date: Mon, 20 Apr 2020 17:09:16 +0800
Subject: [PATCH] fix head fix socks5 test fix request

---
 header.go      |  4 +---
 request.go     |  6 ++++--
 socks5_test.go | 39 ++++++++++++++-------------------------
 3 files changed, 19 insertions(+), 30 deletions(-)

diff --git a/header.go b/header.go
index 6fb5866..a207421 100644
--- a/header.go
+++ b/header.go
@@ -102,9 +102,7 @@ func Parse(r io.Reader) (hd Header, err error) {
 	if hd.Version != socks5Version && hd.Version != socks4Version {
 		return hd, fmt.Errorf("unrecognized SOCKS version[%d]", hd.Version)
 	}
-	if hd.Command != ConnectCommand && hd.Command != BindCommand && hd.Command != AssociateCommand {
-		return hd, fmt.Errorf("unrecognized command[%d]", hd.Command)
-	}
+
 	if hd.Version == socks4Version && hd.Command == AssociateCommand {
 		return hd, fmt.Errorf("wrong version for command")
 	}
diff --git a/request.go b/request.go
index 0ce85ca..6e0988c 100644
--- a/request.go
+++ b/request.go
@@ -51,6 +51,9 @@ func NewRequest(bufConn io.Reader) (*Request, error) {
 	if err != nil {
 		return nil, err
 	}
+	if hd.Command != ConnectCommand && hd.Command != BindCommand && hd.Command != AssociateCommand {
+		return nil, fmt.Errorf("unrecognized command[%d]", hd.Command)
+	}
 	return &Request{
 		Header:   hd,
 		DestAddr: &hd.Address,
@@ -225,8 +228,7 @@ func (s *Server) handleAssociate(ctx context.Context, writer io.Writer, req *Req
 		return fmt.Errorf("dial udp invalid")
 	}
 
-	lAddr, _ := net.ResolveUDPAddr("udp", ":0")
-	bindLn, err := net.ListenUDP("udp4", lAddr)
+	bindLn, err := net.ListenUDP("udp", nil)
 	if err != nil {
 		if err := sendReply(writer, req.Header, serverFailure); err != nil {
 			return fmt.Errorf("failed to send reply, %v", err)
diff --git a/socks5_test.go b/socks5_test.go
index fbe06c3..5939514 100644
--- a/socks5_test.go
+++ b/socks5_test.go
@@ -127,7 +127,7 @@ func TestSOCKS5_Associate(t *testing.T) {
 		IP:   locIP,
 		Port: 12398,
 	}
-	l, err := net.ListenUDP("udp4", lAddr)
+	l, err := net.ListenUDP("udp", lAddr)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
@@ -195,18 +195,6 @@ func TestSOCKS5_Associate(t *testing.T) {
 		socks5Version, UserPassAuth, // use user password auth
 		userPassAuthVersion, authSuccess, // response auth success
 	}
-	rspHead := Header{
-		Version:  socks5Version,
-		Command:  successReply,
-		Reserved: 0,
-		Address: AddrSpec{
-			"",
-			net.ParseIP("0.0.0.0"),
-			0, // Ignore the port
-		},
-		addrType: ipv4Address,
-	}
-	expected = append(expected, rspHead.Bytes()...)
 
 	out := make([]byte, len(expected))
 	conn.SetDeadline(time.Now().Add(time.Second))
@@ -214,22 +202,23 @@ func TestSOCKS5_Associate(t *testing.T) {
 		t.Fatalf("err: %v", err)
 	}
 
-	// Ignore the port
-	proxyBindPort := buildPort(out[12], out[13])
-
-	out[12] = 0
-	out[13] = 0
-
-	t.Logf("proxy bind listen port: %d", proxyBindPort)
-
 	if !bytes.Equal(out, expected) {
 		t.Fatalf("bad: %v", out)
 	}
 
-	udpConn, err := net.DialUDP("udp4", nil, &net.UDPAddr{
-		IP: locIP,
-		//Port: lAddr.Port,
-		Port: proxyBindPort,
+	rspHead, err := Parse(conn)
+	if err != nil {
+		t.Fatalf("bad response header: %v", err)
+	}
+	if rspHead.Version != socks5Version && rspHead.Command != successReply {
+		t.Fatalf("parse success but bad header: %v", rspHead)
+	}
+
+	t.Logf("proxy bind listen port: %d", rspHead.Address.Port)
+
+	udpConn, err := net.DialUDP("udp", nil, &net.UDPAddr{
+		IP:   locIP,
+		Port: rspHead.Address.Port,
 	})
 	if err != nil {
 		t.Fatalf("bad dial: %v", err)