feat: close conn func on error

shoriwe · May 31, 2023 · 5deda30 · 5deda30
1 parent 28d7ca7
commit 5deda30
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 5 deletions.
diff --git a/reverse/master.go b/reverse/master.go
@@ -5,6 +5,7 @@ import (
 	"net"
 
 	"github.com/hashicorp/yamux"
+	"github.com/shoriwe/fullproxy/v3/utils/network"
 )
 
 type Master struct {
@@ -27,11 +28,7 @@ func (m *Master) init() (err error) {
 }
 
 func (m *Master) handle(req *Request) (conn net.Conn, err error) {
-	defer func() {
-		if err != nil && conn != nil {
-			conn.Close()
-		}
-	}()
+	defer network.CloseOnError(&err, conn)
 	err = m.init()
 	if err == nil {
 		conn, err = m.cSession.Open()

diff --git a/utils/network/dial.go b/utils/network/dial.go
@@ -17,3 +17,9 @@ func Dial(addr string) net.Conn {
 	}
 	return conn
 }
+
+func CloseOnError(err *error, conn net.Conn) {
+	if *err != nil && conn != nil {
+		conn.Close()
+	}
+}
diff --git a/utils/network/dial_test.go b/utils/network/dial_test.go
@@ -1,6 +1,7 @@
 package network
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -30,3 +31,13 @@ func TestDial(t *testing.T) {
 func TestNopClose(t *testing.T) {
 	assert.Nil(t, NopClose())
 }
+
+func TestCloseOnError(t *testing.T) {
+	listener := ListenAny()
+	defer listener.Close()
+	go listener.Accept()
+	conn := Dial(listener.Addr().String())
+	defer conn.Close()
+	err := fmt.Errorf("an error")
+	CloseOnError(&err, conn)
+}