From 5c9680c92daf3e41ccf91daea45de2b57d30a57d Mon Sep 17 00:00:00 2001
From: "kayos@tcp.direct" <kayos@tcp.direct>
Date: Sun, 6 Aug 2023 20:03:09 -0700
Subject: [PATCH] Fix: Clean up code and fix race condition

---
 errors.go | 11 +++++++++++
 net.go    |  5 -----
 socks4.go | 44 ++++++++++++++++++++++++--------------------
 socks5.go |  7 ++++++-
 4 files changed, 41 insertions(+), 26 deletions(-)
 create mode 100644 errors.go

diff --git a/errors.go b/errors.go
new file mode 100644
index 0000000..58c9d35
--- /dev/null
+++ b/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")
+)
diff --git a/net.go b/net.go
index 9fc84eb..5cf6b98 100644
--- a/net.go
+++ b/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)
diff --git a/socks4.go b/socks4.go
index e33aee7..555afca 100644
--- a/socks4.go
+++ b/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())
-	if err != nil {
+	resp, err := sesh.sendReceive(conn, req.Bytes())
+
+	defer func() {
+		bufs.MustPut(req.Buffer)
+		req.Buffer = nil
+	}()
+
+	switch {
+	case err != nil:
 		return nil, err
-	} else if len(resp) != 8 {
-		return nil, errors.New("server does not respond properly")
-	}
-	switch resp[1] {
-	case 90:
-		// request granted
-	case 91:
-		return nil, errors.New("socks connection request rejected or failed")
-	case 92:
-		return nil, errors.New("socks connection request rejected because SOCKS server cannot connect to identd on the client")
-	case 93:
-		return nil, errors.New("socks connection request rejected because the client program and identd report different user-ids")
+	case len(resp) != 8:
+		return nil, ErrImproperProtocolResponse
+	case resp[1] == 90:
+		//
+	case resp[1] == 91:
+		return nil, ErrRejectedOrFailed
+	case resp[1] == 92:
+		return nil, ErrIdentdFailed
+	case resp[1] == 93:
+		return nil, ErrIdentMismatch
 	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 {
-		return nil, err
-	}
-	return conn, nil
+	err = conn.SetDeadline(time.Time{})
+
+	return conn, err
 }
diff --git a/socks5.go b/socks5.go
index 87613f7..29d6e0b 100644
--- a/socks5.go
+++ b/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: