fix websocket writes and check cancellation

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
 	}

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")

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
 }

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 }