Fix: Clean up code and fix race condition

errors.go

@@ -0,0 +1,11 @@
+package socks
+
+import "errors"
+
+var (
+	ErrImproperProtocolResponse = errors.New("SOCKS server does not respond properly")
+	ErrRejectedOrFailed         = errors.New("socks connection request rejected or failed")
+	ErrIdentdFailed             = errors.New("socks connection request rejected because SOCKS server cannot connect to identd on the client")
+	ErrIdentMismatch            = errors.New("socks connection request rejected because the client program and identd report different user-ids")
+	ErrUnknownFailure           = errors.New("socks connection request failed, unknown error")
+)

net.go

@@ -20,11 +20,6 @@ type requestBuilder struct {
 	*pool.Buffer
 }
 
-func (b *requestBuilder) final() []byte {
-	defer bufs.MustPut(b.Buffer)
-	return b.MustBytes()
-}
-
 func (sesh *session) sendReceive(conn net.Conn, req []byte) (resp []byte, err error) {
 	// fmt.Printf("sendReceive: %v->%v\n", conn.LocalAddr(), conn.RemoteAddr())
 	// spew.Dump(req)

socks4.go

@@ -1,7 +1,6 @@
 package socks
 
 import (
-	"errors"
 	"net"
 	"time"
 )
@@ -45,27 +44,32 @@ func (sesh *session) dialSocks4(targetAddr string) (_ net.Conn, err error) {
 		req.MustWriteByte(0)
 	}
 
-	resp, err := sesh.sendReceive(conn, req.final())
+	resp, err := sesh.sendReceive(conn, req.Bytes())
-	if err != nil {
+
+	defer func() {
+		bufs.MustPut(req.Buffer)
+		req.Buffer = nil
+	}()
+
+	switch {
+	case err != nil:
 		return nil, err
-	} else if len(resp) != 8 {
+	case len(resp) != 8:
-		return nil, errors.New("server does not respond properly")
+		return nil, ErrImproperProtocolResponse
-	}
+	case resp[1] == 90:
-	switch resp[1] {
+		//
-	case 90:
+	case resp[1] == 91:
-		// request granted
+		return nil, ErrRejectedOrFailed
-	case 91:
+	case resp[1] == 92:
-		return nil, errors.New("socks connection request rejected or failed")
+		return nil, ErrIdentdFailed
-	case 92:
+	case resp[1] == 93:
-		return nil, errors.New("socks connection request rejected because SOCKS server cannot connect to identd on the client")
+		return nil, ErrIdentMismatch
-	case 93:
-		return nil, errors.New("socks connection request rejected because the client program and identd report different user-ids")
 	default:
-		return nil, errors.New("socks connection request failed, unknown error")
+		return nil, ErrUnknownFailure
 	}
+
 	// clear the deadline before returning
-	if err := conn.SetDeadline(time.Time{}); err != nil {
+	err = conn.SetDeadline(time.Time{})
-		return nil, err
+
-	}
+	return conn, err
-	return conn, nil
 }

socks5.go

@@ -74,7 +74,12 @@ func (sesh *session) dialSocks5(targetAddr string) (_ net.Conn, err error) {
 	req.MustWriteByte(byte(port >> 8)) // higher byte of destination port
 	req.MustWriteByte(byte(port))      // lower byte of destination port (big endian)
 
-	resp, err = sesh.sendReceive(conn, req.final())
+	resp, err = sesh.sendReceive(conn, req.Bytes())
+
+	defer func() {
+		bufs.MustPut(req.Buffer)
+		req.Buffer = nil
+	}()
 
 	switch {
 	case err != nil: