fix: connections not closing

Author: tabjy

client/client.go

@@ -150,7 +150,7 @@ func (c *proxyConn) connect(ctx context.Context) (net.Conn, error) {
 		return nil, err
 	}
 
-	var cipherTarget io.ReadWriter
+	var cipherTarget net.Conn
 	ed := crypto.StreamEncryptDecrypter{
 		EncryptKey: c.sessionKey,
 		DecryptKey: c.sessionKey,
@@ -209,10 +209,7 @@ func (c *proxyConn) connect(ctx context.Context) (net.Conn, error) {
 		}
 	}
 
-	return &CipherConn {
-		cipherTarget,
-		c.target,
-	}, nil
+	return cipherTarget, nil
 }
 
 func (c *proxyConn) writePubKey() error {
@@ -314,17 +311,3 @@ func (c *proxyConn) readReply() error {
 
 	return nil
 }
-
-// CipherConn implements net.Conn interface, with a underlying io.ReadWriter.
-type CipherConn struct {
-	io.ReadWriter
-	net.Conn
-}
-
-func (c *CipherConn) Read(b []byte)  (n int, err error) {
-	return c.ReadWriter.Read(b)
-}
-
-func (c *CipherConn) Write(b []byte)  (n int, err error) {
-	return c.ReadWriter.Write(b)
-}

cmd/groundhog/main.go

@@ -87,7 +87,7 @@ func initLogger() {
 	case "fatal":
 		level = yagl.LvlFatal
 	default:
-		fmt.Fprintf(os.Stderr, "unrecognized logging level: %s\n", level)
+		fmt.Fprintf(os.Stderr, "unrecognized logging level: %s\n", logLevel)
 		os.Exit(1)
 	}
 

common/crypto/crypto.go

@@ -1,10 +1,11 @@
 package crypto
 
 import (
-	"crypto/cipher"
-	"io"
-	"errors"
 	"crypto/aes"
+	"errors"
+	"io"
+	"net"
+	"crypto/cipher"
 )
 
 type readWriter struct {
@@ -68,36 +69,35 @@ func (ed *StreamEncryptDecrypter) initCipherStream() error {
 // corresponding ciphertext io.ReadWriter. Any ciphertext write to returned
 // io.ReadWriter will be decrypted and write to plaintext. Any plaintext read
 // from plaintext will be encrypted and write to returned io.ReadWriter.
-func (ed *StreamEncryptDecrypter) Ciphertext(plaintext io.ReadWriter) (io.ReadWriter, error) {
+func (ed *StreamEncryptDecrypter) Ciphertext(plaintext net.Conn) (net.Conn, error) {
 	if err := ed.initCipherStream(); err != nil {
 		return nil, err
 	}
 
 	// logic here could be simpler if golang has a built-in duplex pipe.
-	// net.Pipe is a solution, but seems too heavy,
-	// with possibility of causing memory leak if not careful
 	cipherRdIn, cipherWtOut := io.Pipe()
 	cipherRdOut, cipherWtIn := io.Pipe()
 
-	ciphertext := &readWriter{
-		cipherRdOut,
-		cipherWtOut,
+	ciphertext := &CipherConn{
+		&readWriter{
+			cipherRdOut,
+			cipherWtOut,
+		},
+		plaintext,
 	}
 
 	// decrypt ciphertext to plaintext
 	go func() {
 		decrypter := &cipher.StreamReader{S: ed.DecryptStream, R: cipherRdIn}
 		io.Copy(plaintext, decrypter)
-		cipherWtOut.Close()
-		cipherRdIn.Close()
+		ciphertext.Close() // which close the underlying plaintext
 	}()
 
 	// encrypt plaintext to ciphertext
 	go func() {
 		encrypter := &cipher.StreamWriter{S: ed.EncryptStream, W: cipherWtIn}
 		io.Copy(encrypter, plaintext)
-		cipherWtIn.Close()
-		cipherRdOut.Close()
+		ciphertext.Close()
 	}()
 
 	return ciphertext, nil
@@ -107,34 +107,55 @@ func (ed *StreamEncryptDecrypter) Ciphertext(plaintext io.ReadWriter) (io.ReadWr
 // corresponding plaintext io.ReadWriter. Any plaintext write to returned
 // io.ReadWriter will be encrypted and write to ciphertext. Any ciphertext read
 // from ciphertext will be decrypted and write to returned io.ReadWriter.
-func (ed *StreamEncryptDecrypter) Plaintext(ciphertext io.ReadWriter) (io.ReadWriter, error) {
+func (ed *StreamEncryptDecrypter) Plaintext(ciphertext net.Conn) (net.Conn, error) {
 	if err := ed.initCipherStream(); err != nil {
 		return nil, err
 	}
 
 	plainRdIn, plainWtOut := io.Pipe()
 	plainRdOut, plainWtIn := io.Pipe()
 
-	plaintext := &readWriter{
-		plainRdOut,
-		plainWtOut,
+	plaintext := &CipherConn {
+		&readWriter{
+			plainRdOut,
+			plainWtOut,
+		},
+		ciphertext,
 	}
 
 	// encrypt plaintext to ciphertext
 	go func() {
 		encrypter := &cipher.StreamWriter{S: ed.EncryptStream, W: ciphertext}
 		io.Copy(encrypter, plainRdIn)
-		plainWtIn.Close()
-		plainRdOut.Close()
+		plaintext.Close()
 	}()
 
 	// decrypt ciphertext to plaintext
 	go func() {
 		decrypter := &cipher.StreamReader{S: ed.DecryptStream, R: ciphertext}
 		io.Copy(plainWtIn, decrypter)
-		plainWtOut.Close()
-		plainRdIn.Close()
+		plaintext.Close()
 	}()
 
 	return plaintext, nil
-}
+}
+
+// CipherConn implements net.Conn interface, with a underlying io.ReadWriter.
+type CipherConn struct {
+	io.ReadWriter
+	net.Conn
+}
+
+func (c *CipherConn) Read(b []byte) (n int, err error) {
+	if _, err := c.Conn.Read([]byte{}); err != nil {
+		return 0, err
+	}
+	return c.ReadWriter.Read(b)
+}
+
+func (c *CipherConn) Write(b []byte) (n int, err error) {
+	if _, err := c.Conn.Write([]byte{}); err != nil {
+		return 0, err
+	}
+	return c.ReadWriter.Write(b)
+}

common/util/util.go

@@ -4,6 +4,7 @@ package util
 import (
 	"io"
 	"sync"
+	"net"
 )
 
 // Proxy connect two ReadWriter, forward data between them in a full-duplex
@@ -16,6 +17,7 @@ func Proxy(lhs io.ReadWriter, rhs io.ReadWriter) (lhsWritten, rhsWritten int64,
 	wg.Add(1)
 	go func() {
 		lhsWritten, err = io.Copy(lhs, rhs)
+		closeNetConn(lhs, rhs)
 		if err != nil {
 			wg.Done()
 		} else {
@@ -27,6 +29,7 @@ func Proxy(lhs io.ReadWriter, rhs io.ReadWriter) (lhsWritten, rhsWritten int64,
 	wg.Add(1)
 	go func() {
 		rhsWritten, err = io.Copy(rhs, lhs)
+		closeNetConn(lhs, rhs)
 		if err != nil {
 			wg.Done()
 		} else {
@@ -37,3 +40,11 @@ func Proxy(lhs io.ReadWriter, rhs io.ReadWriter) (lhsWritten, rhsWritten int64,
 	wg.Wait()
 	return
 }
+
+func closeNetConn(rws... io.ReadWriter) {
+	for _, v := range rws {
+		if conn, ok := v.(net.Conn); ok {
+			conn.Close()
+		}
+	}
+}

server/server.go

@@ -140,7 +140,8 @@ type gndhog struct {
 }
 
 func (g *gndhog) init(ctx context.Context, conn net.Conn) {
-	defer conn.Close()
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
 
 	// watchdog to close connections if context cancelled
 	go func() {

socks5/socks5.go

@@ -87,7 +87,8 @@ type socks struct {
 }
 
 func (s *socks) init(ctx context.Context, conn net.Conn) {
-	defer conn.Close()
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
 
 	// watchdog to close connections if context cancelled
 	go func() {