From b2609496852bfc1711cbfc122447ae8e6347b228 Mon Sep 17 00:00:00 2001 From: Graham McIntire Date: Sat, 25 Jul 2026 14:57:15 -0500 Subject: [PATCH] fix websocket writes and check cancellation --- checks.go | 22 ++++++++++++++++++++-- checks_test.go | 36 ++++++++++++++++++++++++++++++++++++ websocket.go | 31 ++++++++++++++++++++++++------- websocket_test.go | 44 ++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 124 insertions(+), 9 deletions(-) diff --git a/checks.go b/checks.go index 74fdc6b..c118447 100644 --- a/checks.go +++ b/checks.go @@ -158,6 +158,9 @@ func executeHTTPCheck(ctx context.Context, config *pb.HttpCheckConfig, timeoutMs } if resp.StatusCode != expectedStatus { + // Drain the response before returning so repeated failing checks can + // still reuse the shared transport connection. + _, _ = io.Copy(io.Discard, resp.Body) return 2, fmt.Sprintf("HTTP %d, expected %d", resp.StatusCode, expectedStatus), responseTime } @@ -175,6 +178,12 @@ func executeHTTPCheck(ctx context.Context, config *pb.HttpCheckConfig, timeoutMs if !re.Match(body) { return 2, fmt.Sprintf("Content does not match pattern: %s", config.Regex), responseTime } + } else { + // Consume successful response bodies so the shared transport can reuse + // the underlying connection for subsequent checks. + if _, err := io.Copy(io.Discard, resp.Body); err != nil { + return 2, fmt.Sprintf("Failed to read body: %v", err), responseTime + } } return 0, fmt.Sprintf("HTTP %d OK", resp.StatusCode), responseTime @@ -357,7 +366,12 @@ func executeSSLCheck(ctx context.Context, config *pb.SslCheckConfig, timeoutMs u } startTime := time.Now() - conn, err := tls.DialWithDialer(dialer, "tcp", address, tlsConfig) + dialCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + conn, err := (&tls.Dialer{ + NetDialer: dialer, + Config: tlsConfig, + }).DialContext(dialCtx, "tcp", address) responseTime := float64(time.Since(startTime).Milliseconds()) if err != nil { @@ -365,7 +379,11 @@ func executeSSLCheck(ctx context.Context, config *pb.SslCheckConfig, timeoutMs u } defer func() { _ = conn.Close() }() - certs := conn.ConnectionState().PeerCertificates + tlsConn, ok := conn.(*tls.Conn) + if !ok { + return 3, "Connection did not negotiate TLS", responseTime + } + certs := tlsConn.ConnectionState().PeerCertificates if len(certs) == 0 { return 3, fmt.Sprintf("No certificate presented by %s:%d", host, port), responseTime } diff --git a/checks_test.go b/checks_test.go index 54a9be1..9000697 100644 --- a/checks_test.go +++ b/checks_test.go @@ -1469,6 +1469,42 @@ func TestSSLCheck_ConnectionFailure(t *testing.T) { } } +func TestSSLCheck_HonorsCancelledContext(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer func() { _ = ln.Close() }() + + accepted := make(chan net.Conn, 1) + go func() { + conn, acceptErr := ln.Accept() + if acceptErr == nil { + accepted <- conn + } + }() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + start := time.Now() + status, _, _ := executeSSLCheck(ctx, &pb.SslCheckConfig{ + Host: "127.0.0.1", + Port: parsePort(portFromListener(ln)), + }, 5000) + if status != 3 { + t.Fatalf("expected status 3, got %d", status) + } + if elapsed := time.Since(start); elapsed > time.Second { + t.Fatalf("cancelled check took %v", elapsed) + } + + select { + case conn := <-accepted: + _ = conn.Close() + default: + } +} + func TestSSLCheck_ClosedPort(t *testing.T) { // Bind and immediately close to get a port that refuses connections ln, err := net.Listen("tcp", "127.0.0.1:0") diff --git a/websocket.go b/websocket.go index 52e97ee..1056790 100644 --- a/websocket.go +++ b/websocket.go @@ -124,7 +124,7 @@ func wsConnect(u *url.URL, host, network string) (*WSConn, error) { req := fmt.Sprintf("GET %s HTTP/1.1\r\nHost: %s\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: %s\r\nSec-WebSocket-Version: 13\r\n\r\n", path, u.Host, key) - if _, err := conn.Write([]byte(req)); err != nil { + if err := writeAll(conn, []byte(req)); err != nil { _ = conn.Close() return nil, fmt.Errorf("write handshake: %w", err) } @@ -349,16 +349,33 @@ func (ws *WSConn) writeFrame(opcode int, payload []byte) error { // Mask key var maskKey [4]byte - _, _ = rand.Read(maskKey[:]) + if _, err := randRead(maskKey[:]); err != nil { + return fmt.Errorf("generate frame mask: %w", err) + } buf = append(buf, maskKey[:]...) - // Masked payload - masked := make([]byte, len(payload)) + // Append and mask the payload in place to avoid a second payload-sized + // allocation for every outbound frame. + payloadStart := len(buf) + buf = append(buf, payload...) + masked := buf[payloadStart:] for i, b := range payload { masked[i] = b ^ maskKey[i%4] } - buf = append(buf, masked...) - _, err := ws.conn.Write(buf) - return err + return writeAll(ws.conn, buf) +} + +func writeAll(w io.Writer, data []byte) error { + for len(data) > 0 { + n, err := w.Write(data) + if err != nil { + return err + } + if n <= 0 || n > len(data) { + return io.ErrShortWrite + } + data = data[n:] + } + return nil } diff --git a/websocket_test.go b/websocket_test.go index d6e4ef4..599a3f4 100644 --- a/websocket_test.go +++ b/websocket_test.go @@ -43,6 +43,37 @@ func TestWriteFrameMasked(t *testing.T) { } } +func TestWriteFrameHandlesShortWrites(t *testing.T) { + var dst bytes.Buffer + rw := &shortWriteReadWriter{dst: &dst, maxWrite: 3} + ws := testWSConn(&nopCloser{readWriter: rw}) + + if err := ws.writeFrame(opText, []byte("hello")); err != nil { + t.Fatal(err) + } + if dst.Len() != 2+4+len("hello") { + t.Fatalf("frame length: got %d, want %d", dst.Len(), 2+4+len("hello")) + } +} + +func TestWriteFrameMaskGenerationError(t *testing.T) { + origRandRead := randRead + defer func() { randRead = origRandRead }() + randRead = func([]byte) (int, error) { + return 0, fmt.Errorf("entropy exhausted") + } + + var buf bytes.Buffer + ws := testWSConn(&nopCloser{readWriter: &buf}) + err := ws.writeFrame(opText, []byte("hello")) + if err == nil || !strings.Contains(err.Error(), "generate frame mask") { + t.Fatalf("expected frame mask error, got %v", err) + } + if buf.Len() != 0 { + t.Fatalf("wrote %d bytes despite mask generation failure", buf.Len()) + } +} + func TestWriteFrame16BitLength(t *testing.T) { var buf bytes.Buffer ws := testWSConn(&nopCloser{readWriter: &buf}) @@ -907,6 +938,19 @@ type nopCloser struct { } } +type shortWriteReadWriter struct { + dst *bytes.Buffer + maxWrite int +} + +func (s *shortWriteReadWriter) Read([]byte) (int, error) { return 0, io.EOF } +func (s *shortWriteReadWriter) Write(p []byte) (int, error) { + if len(p) > s.maxWrite { + p = p[:s.maxWrite] + } + return s.dst.Write(p) +} + func (n *nopCloser) Read(p []byte) (int, error) { return n.readWriter.Read(p) } func (n *nopCloser) Write(p []byte) (int, error) { return n.readWriter.Write(p) } func (n *nopCloser) Close() error { return nil }