fix: resolve 7 quality issues (#100-#106)

transport/client_test.go

@@ -19,11 +19,15 @@ package getty
 
 import (
 	"bytes"
+	"crypto/tls"
+	"errors"
 	"net"
 	"net/http"
 	"os"
+	"path/filepath"
 	"strconv"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 )
@@ -36,6 +40,150 @@ import (
 
 type PackageHandler struct{}
 
+var errTestTLSConfig = errors.New("test TLS config failure")
+
+type countingTLSConfigBuilder struct {
+	calls   atomic.Int32
+	entered chan struct{}
+	release <-chan struct{}
+}
+
+func (b *countingTLSConfigBuilder) BuildTlsConfig() (*tls.Config, error) {
+	b.calls.Add(1)
+	select {
+	case b.entered <- struct{}{}:
+	default:
+	}
+	if b.release != nil {
+		<-b.release
+	}
+	return nil, errTestTLSConfig
+}
+
+func newFailingReconnectClient(builder TlsConfigBuilder, interval time.Duration, attempts int) *client {
+	return newClient(TCP_CLIENT,
+		WithServerAddress("127.0.0.1:1"),
+		WithConnectionNumber(1),
+		WithReconnectInterval(int(interval)),
+		WithReconnectAttempts(attempts),
+		WithClientSslEnabled(true),
+		WithClientTlsConfigBuilder(builder),
+	)
+}
+
+func TestReconnectAttemptsAreExactAndSkipFinalBackoff(t *testing.T) {
+	builder := &countingTLSConfigBuilder{entered: make(chan struct{}, 4)}
+	clt := newFailingReconnectClient(builder, 100*time.Millisecond, 3)
+
+	started := time.Now()
+	clt.RunEventLoop(func(Session) error { return nil })
+	elapsed := time.Since(started)
+
+	if got := builder.calls.Load(); got != 3 {
+		t.Fatalf("TLS config build calls = %d, want exactly 3 reconnect attempts", got)
+	}
+	// Correct behavior waits after attempts one and two only: 100ms + 200ms.
+	// A final backoff adds another 300ms.
+	if elapsed >= 500*time.Millisecond {
+		t.Fatalf("reconnect loop took %v; it appears to wait after the final attempt", elapsed)
+	}
+}
+
+func TestReconnectBackoffIsCancelledByClose(t *testing.T) {
+	builder := &countingTLSConfigBuilder{entered: make(chan struct{}, 4)}
+	clt := newFailingReconnectClient(builder, 2*time.Second, 3)
+
+	eventLoopDone := make(chan struct{})
+	go func() {
+		clt.RunEventLoop(func(Session) error { return nil })
+		close(eventLoopDone)
+	}()
+	select {
+	case <-builder.entered:
+	case <-time.After(time.Second):
+		t.Fatal("first reconnect attempt did not start")
+	}
+	time.Sleep(20 * time.Millisecond)
+
+	closeDone := make(chan struct{})
+	go func() {
+		clt.Close()
+		close(closeDone)
+	}()
+
+	timedOut := false
+	select {
+	case <-closeDone:
+	case <-time.After(200 * time.Millisecond):
+		timedOut = true
+	}
+	if timedOut {
+		select {
+		case <-closeDone:
+		case <-time.After(3 * time.Second):
+			t.Fatal("Close did not return after the current backoff elapsed")
+		}
+		t.Fatal("Close did not cancel the reconnect backoff")
+	}
+	select {
+	case <-eventLoopDone:
+	case <-time.After(time.Second):
+		t.Fatal("RunEventLoop did not return after Close")
+	}
+}
+
+func TestSessionReconnectIsTrackedByClose(t *testing.T) {
+	releaseReconnect := make(chan struct{})
+	builder := &countingTLSConfigBuilder{
+		entered: make(chan struct{}, 4),
+		release: releaseReconnect,
+	}
+	clt := newFailingReconnectClient(builder, 2*time.Second, 3)
+	localConn, peerConn := net.Pipe()
+	defer func() {
+		_ = localConn.Close()
+		_ = peerConn.Close()
+	}()
+	ss := newTCPSession(localConn, clt).(*session)
+	ss.SetAttribute(sessionClientKey, clt)
+	ss.SetAttribute(ignoreReconnectKey, false)
+
+	sessionStopDone := make(chan struct{})
+	go func() {
+		ss.stop()
+		close(sessionStopDone)
+	}()
+	select {
+	case <-builder.entered:
+	case <-time.After(time.Second):
+		t.Fatal("session-triggered reconnect did not start")
+	}
+
+	closeDone := make(chan struct{})
+	go func() {
+		clt.Close()
+		close(closeDone)
+	}()
+	select {
+	case <-closeDone:
+		close(releaseReconnect)
+		t.Fatal("Close returned while the session-triggered reconnect was still running")
+	case <-time.After(100 * time.Millisecond):
+	}
+
+	close(releaseReconnect)
+	select {
+	case <-closeDone:
+	case <-time.After(time.Second):
+		t.Fatal("Close did not return after the session-triggered reconnect completed")
+	}
+	select {
+	case <-sessionStopDone:
+	case <-time.After(time.Second):
+		t.Fatal("session stop did not return")
+	}
+}
+
 func (h *PackageHandler) Read(ss Session, data []byte) (any, int, error) {
 	return nil, 0, nil
 }
@@ -93,7 +241,11 @@ func newSessionCallback(session Session, handler *MessageHandler) error {
 
 func TestTCPClient(t *testing.T) {
 	listenLocalServer := func() (net.Listener, error) {
-		listener, err := net.Listen("tcp", ":0")
+		// #106: bind a concrete loopback address instead of ":0" (which
+		// resolves to the unspecified "[::]" address). Dialing "[::]:port"
+		// is not a valid connect destination, so every dial failed and the
+		// client's reconnect loop hung the test forever.
+		listener, err := net.Listen("tcp", "127.0.0.1:0")
 		if err != nil {
 			return nil, err
 		}
@@ -459,6 +611,58 @@ func DownloadFile(filepath string, content []byte) error {
 	return err
 }
 
+func TestBuildWSSClientTLSConfig(t *testing.T) {
+	t.Run("system roots", func(t *testing.T) {
+		config, err := (&client{}).buildWSSClientTLSConfig()
+		if err != nil {
+			t.Fatal(err)
+		}
+		if config.RootCAs != nil {
+			t.Fatal("RootCAs must be nil when no custom root certificate is configured")
+		}
+		if config.MinVersion != tls.VersionTLS12 {
+			t.Fatalf("MinVersion = %d, want TLS 1.2 (%d)", config.MinVersion, tls.VersionTLS12)
+		}
+	})
+
+	t.Run("custom root certificate", func(t *testing.T) {
+		certPath := filepath.Join(t.TempDir(), "root.crt")
+		if err := os.WriteFile(certPath, WssClientCRT, 0o600); err != nil {
+			t.Fatal(err)
+		}
+		config, err := (&client{ClientOptions: ClientOptions{cert: certPath}}).buildWSSClientTLSConfig()
+		if err != nil {
+			t.Fatal(err)
+		}
+		if config.RootCAs == nil {
+			t.Fatal("RootCAs is nil with a configured root certificate")
+		}
+		if len(config.Certificates) != 0 {
+			t.Fatalf("client Certificates contains %d entries, want 0 for a root-only option", len(config.Certificates))
+		}
+		if config.MinVersion != tls.VersionTLS12 {
+			t.Fatalf("MinVersion = %d, want TLS 1.2 (%d)", config.MinVersion, tls.VersionTLS12)
+		}
+	})
+
+	t.Run("invalid PEM", func(t *testing.T) {
+		certPath := filepath.Join(t.TempDir(), "invalid.crt")
+		if err := os.WriteFile(certPath, []byte("not a certificate"), 0o600); err != nil {
+			t.Fatal(err)
+		}
+		if _, err := (&client{ClientOptions: ClientOptions{cert: certPath}}).buildWSSClientTLSConfig(); err == nil {
+			t.Fatal("invalid PEM returned nil error")
+		}
+	})
+
+	t.Run("missing file", func(t *testing.T) {
+		certPath := filepath.Join(t.TempDir(), "missing.crt")
+		if _, err := (&client{ClientOptions: ClientOptions{cert: certPath}}).buildWSSClientTLSConfig(); err == nil {
+			t.Fatal("missing root certificate returned nil error")
+		}
+	})
+}
+
 func TestNewWSSClient(t *testing.T) {
 	var (
 		err              error

transport/connection.go

@@ -127,7 +127,7 @@ func (c *gettyConn) GetActive() time.Time {
 
 // removed unused methods send/close
 
-func (c gettyConn) ReadTimeout() time.Duration {
+func (c *gettyConn) ReadTimeout() time.Duration {
 	return c.rTimeout.Load()
 }
 
@@ -150,7 +150,7 @@ func (c *gettyConn) SetReadTimeout(rTimeout time.Duration) {
 	}
 }
 
-func (c gettyConn) WriteTimeout() time.Duration {
+func (c *gettyConn) WriteTimeout() time.Duration {
 	return c.wTimeout.Load()
 }
 
@@ -227,12 +227,44 @@ func (t *writeFlusher) Write(p []byte) (int, error) {
 		return n, perrors.WithStack(err)
 	}
 	if err := t.flusher.Flush(); err != nil {
-		return 0, perrors.WithStack(err)
+		return n, perrors.WithStack(err)
 	}
 
 	return n, nil
 }
 
+// for snappy compress. #102: snappy.NewBufferedWriter buffers writes and only
+// emits data on Flush, so small packets would sit in the buffer forever if not
+// flushed after every Write. This wrapper flushes on every Write, mirroring the
+// flate writeFlusher behavior above.
+type snappyWriteFlusher struct {
+	writer *snappy.Writer
+	lock   sync.Mutex
+}
+
+func newSnappyWriteFlusher(w *snappy.Writer) *snappyWriteFlusher {
+	return &snappyWriteFlusher{writer: w}
+}
+
+func (s *snappyWriteFlusher) Write(p []byte) (int, error) {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	n, err := s.writer.Write(p)
+	if err != nil {
+		return n, perrors.WithStack(err)
+	}
+	if err := s.writer.Flush(); err != nil {
+		return n, perrors.WithStack(err)
+	}
+	return n, nil
+}
+
+func (s *snappyWriteFlusher) Close() error {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	return perrors.WithStack(s.writer.Close())
+}
+
 // SetCompressType set compress type(tcp: zip/snappy, websocket:zip)
 func (t *gettyTCPConn) SetCompressType(c CompressType) {
 	switch c {
@@ -251,7 +283,9 @@ func (t *gettyTCPConn) SetCompressType(c CompressType) {
 		ioReader := io.Reader(t.conn)
 		t.reader = snappy.NewReader(ioReader)
 		ioWriter := io.Writer(t.conn)
-		t.writer = snappy.NewBufferedWriter(ioWriter)
+		// #102: wrap the buffered snappy writer so every Write is flushed,
+		// otherwise small packets never leave the internal buffer.
+		t.writer = newSnappyWriteFlusher(snappy.NewBufferedWriter(ioWriter))
 
 	default:
 		panic(fmt.Sprintf("illegal comparess type %d", c))
@@ -306,8 +340,23 @@ func (t *gettyTCPConn) Send(pkg any) (int, error) {
 	}
 
 	if buffers, ok := pkg.([][]byte); ok {
-		netBuf := net.Buffers(buffers)
-		lg, err = netBuf.WriteTo(t.conn)
+		// #102: when compression is enabled the [][]byte path must go through
+		// t.writer (the compress writer), otherwise it writes raw frames
+		// directly to t.conn and the peer receives a corrupt mix of
+		// compressed and uncompressed data.
+		if t.compress == CompressNone {
+			netBuf := net.Buffers(buffers)
+			lg, err = netBuf.WriteTo(t.conn)
+		} else {
+			for _, b := range buffers {
+				var n int
+				n, err = t.writer.Write(b)
+				if err != nil {
+					break
+				}
+				lg += int64(n)
+			}
+		}
 		if err == nil {
 			t.writeBytes.Add((uint32)(lg))
 			t.writePkgNum.Add((uint32)(len(buffers)))
@@ -338,16 +387,21 @@ func (t *gettyTCPConn) CloseConn(waitSec int) {
 	// }
 
 	if t.conn != nil {
-		if writer, ok := t.writer.(*snappy.Writer); ok {
+		// #102: snappy writer is now wrapped in *snappyWriteFlusher.
+		if writer, ok := t.writer.(*snappyWriteFlusher); ok {
 			if err := writer.Close(); err != nil {
 				log.Errorf("snappy.Writer.Close() = error:%+v", err)
 			}
 		}
+		// #103: do not hard-assert *tls.Conn; use safe type assertions so a
+		// non-TLS, non-TCP conn does not panic here.
 		if conn, ok := t.conn.(*net.TCPConn); ok {
 			_ = conn.SetLinger(waitSec)
 			_ = conn.Close()
+		} else if tlsConn, ok := t.conn.(*tls.Conn); ok {
+			_ = tlsConn.Close()
 		} else {
-			_ = t.conn.(*tls.Conn).Close()
+			_ = t.conn.Close()
 		}
 		t.conn = nil
 	}