knative.dev/pkg@v0.0.0-20260602142205-ac97e43f6622/websocket/connection_test.go

/*
Copyright 2019 The Knative Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package websocket

import (
	"context"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	ktesting "knative.dev/pkg/logging/testing"

	"k8s.io/apimachinery/pkg/util/wait"

	"github.com/gorilla/websocket"
)

const propagationTimeout = 5 * time.Second

type inspectableConnection struct {
	nextReaderCalls      chan struct{}
	writeMessageCalls    chan struct{}
	closeCalls           chan struct{}
	setReadDeadlineCalls chan struct{}
	setPongHandlerCalls  chan struct{}

	nextReaderFunc func() (int, io.Reader, error)
}

func (c *inspectableConnection) WriteMessage(messageType int, data []byte) error {
	if c.writeMessageCalls != nil {
		c.writeMessageCalls <- struct{}{}
	}
	return nil
}

func (c *inspectableConnection) NextReader() (int, io.Reader, error) {
	if c.nextReaderCalls != nil {
		c.nextReaderCalls <- struct{}{}
	}
	return c.nextReaderFunc()
}

func (c *inspectableConnection) Close() error {
	if c.closeCalls != nil {
		c.closeCalls <- struct{}{}
	}
	return nil
}

func (c *inspectableConnection) SetReadDeadline(deadline time.Time) error {
	if c.setReadDeadlineCalls != nil {
		c.setReadDeadlineCalls <- struct{}{}
	}
	return nil
}

func (c *inspectableConnection) SetPongHandler(func(string) error) {
	if c.setPongHandlerCalls != nil {
		c.setPongHandlerCalls <- struct{}{}
	}
}

// staticConnFactory returns a static connection, for example
// an inspectable connection.
func staticConnFactory(conn rawConnection) func() (rawConnection, error) {
	return func() (rawConnection, error) {
		return conn, nil
	}
}

// errConnFactory returns a static error.
func errConnFactory(err error) func() (rawConnection, error) {
	return func() (rawConnection, error) {
		return nil, err
	}
}

func TestRetriesWhileConnect(t *testing.T) {
	const wantConnects = 2
	gotConnects := 0

	spy := &inspectableConnection{
		closeCalls:           make(chan struct{}, 1),
		setReadDeadlineCalls: make(chan struct{}, 1),
		setPongHandlerCalls:  make(chan struct{}, 1),
	}

	connFactory := func() (rawConnection, error) {
		gotConnects++
		if gotConnects == wantConnects {
			return spy, nil
		}
		return nil, errors.New("not yet")
	}
	conn := newConnection(connFactory, nil)

	conn.connect()
	conn.Shutdown()

	if gotConnects != wantConnects {
		t.Fatalf("Wanted %v retries. Got %v.", wantConnects, gotConnects)
	}

	// We want a readDeadline and a pongHandler to be set on the final connection.
	if got, want := len(spy.setReadDeadlineCalls), 1; got != want {
		t.Fatalf("Got %d 'SetReadDeadline' calls, want %d", got, want)
	}
	if got, want := len(spy.setPongHandlerCalls), 1; got != want {
		t.Fatalf("Got %d 'SetPongHandler' calls, want %d", got, want)
	}

	if len(spy.closeCalls) != 1 {
		t.Fatal("Wanted 'Close' to be called once, but got", len(spy.closeCalls))
	}
}

func TestSendErrorOnNoConnection(t *testing.T) {
	want := ErrConnectionNotEstablished

	conn := &ManagedConnection{}
	got := conn.Send("test")

	if !errors.Is(got, want) {
		t.Fatalf("Wanted error to be %v, but it was %v.", want, got)
	}
}

func TestStatusOnNoConnection(t *testing.T) {
	want := ErrConnectionNotEstablished

	conn := &ManagedConnection{}
	got := conn.Status()

	if !errors.Is(got, want) {
		t.Fatalf("Wanted error to be %v, but it was %v.", want, got)
	}
}

func TestSendErrorOnEncode(t *testing.T) {
	spy := &inspectableConnection{
		writeMessageCalls: make(chan struct{}, 1),
	}
	conn := newConnection(staticConnFactory(spy), nil)
	conn.connect()
	// gob cannot encode nil values
	got := conn.Send(nil)

	if got == nil {
		t.Fatal("Expected an error but got none")
	}
	if len(spy.writeMessageCalls) != 0 {
		t.Fatalf("Expected 'WriteMessage' not to be called, but was called %v times", spy.writeMessageCalls)
	}
}

func TestSendMessage(t *testing.T) {
	spy := &inspectableConnection{
		writeMessageCalls: make(chan struct{}, 1),
	}
	conn := newConnection(staticConnFactory(spy), nil)
	conn.connect()

	if got := conn.Status(); got != nil {
		t.Errorf("Status() = %v, wanted nil", got)
	}

	if got := conn.Send("test"); got != nil {
		t.Fatalf("Expected no error but got: %+v", got)
	}
	if len(spy.writeMessageCalls) != 1 {
		t.Fatalf("Expected 'WriteMessage' to be called once, but was called %v times", spy.writeMessageCalls)
	}
}

func TestSendRawMessage(t *testing.T) {
	spy := &inspectableConnection{
		writeMessageCalls: make(chan struct{}, 1),
	}
	conn := newConnection(staticConnFactory(spy), nil)
	conn.connect()

	if got := conn.Status(); got != nil {
		t.Errorf("Status() = %v, wanted nil", got)
	}

	if got := conn.SendRaw(websocket.BinaryMessage, []byte("test")); got != nil {
		t.Fatalf("Expected no error but got: %+v", got)
	}
	if len(spy.writeMessageCalls) != 1 {
		t.Fatalf("Expected 'WriteMessage' to be called once, but was called %v times", spy.writeMessageCalls)
	}
}

func TestReceiveMessage(t *testing.T) {
	testMessage := "testmessage"

	spy := &inspectableConnection{
		writeMessageCalls: make(chan struct{}, 1),
		nextReaderCalls:   make(chan struct{}, 1),
		nextReaderFunc: func() (int, io.Reader, error) {
			return websocket.TextMessage, strings.NewReader(testMessage), nil
		},
	}

	messageChan := make(chan []byte, 1)
	conn := newConnection(staticConnFactory(spy), messageChan)
	conn.connect()
	go conn.keepalive()

	got := <-messageChan

	if string(got) != testMessage {
		t.Errorf("Received the wrong message, wanted %q, got %q", testMessage, string(got))
	}
}

func TestCloseClosesConnection(t *testing.T) {
	spy := &inspectableConnection{
		closeCalls: make(chan struct{}, 1),
	}
	conn := newConnection(staticConnFactory(spy), nil)
	conn.connect()
	conn.Shutdown()

	if len(spy.closeCalls) != 1 {
		t.Fatal("Expected 'Close' to be called once, got", len(spy.closeCalls))
	}
}

func TestCloseIgnoresNoConnection(t *testing.T) {
	conn := &ManagedConnection{
		closeChan: make(chan struct{}, 1),
	}
	got := conn.Shutdown()

	if got != nil {
		t.Fatal("Expected no error, got", got)
	}
}

func TestConnectFailureReturnsError(t *testing.T) {
	conn := newConnection(errConnFactory(ErrConnectionNotEstablished), nil)

	// Shorten the connection backoff duration for this test
	conn.connectionBackoff.Duration = 1 * time.Millisecond

	got := conn.connect()

	if got == nil {
		t.Fatal("Expected an error but got none")
	}
}

func TestKeepaliveWithNoConnectionReturnsError(t *testing.T) {
	conn := newConnection(nil, nil)
	got := conn.keepalive()

	if got == nil {
		t.Fatal("Expected an error but got none")
	}
}

func TestConnectLoopIsStopped(t *testing.T) {
	conn := newConnection(errConnFactory(errors.New("connection error")), nil)

	errorChan := make(chan error)
	go func() {
		errorChan <- conn.connect()
	}()

	conn.Shutdown()

	select {
	case err := <-errorChan:
		if !errors.Is(err, errShuttingDown) {
			t.Errorf("Wrong 'connect' error, got %v, want %v", err, errShuttingDown)
		}
	case <-time.After(propagationTimeout):
		t.Error("Timed out waiting for the keepalive loop to stop.")
	}
}

func TestKeepaliveLoopIsStopped(t *testing.T) {
	spy := &inspectableConnection{
		nextReaderFunc: func() (int, io.Reader, error) {
			return websocket.TextMessage, nil, nil
		},
	}
	conn := newConnection(staticConnFactory(spy), nil)
	conn.connect()

	errorChan := make(chan error)
	go func() {
		errorChan <- conn.keepalive()
	}()

	conn.Shutdown()

	select {
	case err := <-errorChan:
		if !errors.Is(err, errShuttingDown) {
			t.Errorf("Wrong 'keepalive' error, got %v, want %v", err, errShuttingDown)
		}
	case <-time.After(propagationTimeout):
		t.Error("Timed out waiting for the keepalive loop to stop.")
	}
}

func TestDoubleShutdown(t *testing.T) {
	spy := &inspectableConnection{
		closeCalls: make(chan struct{}, 2), // potentially allow 2 calls
	}
	conn := newConnection(staticConnFactory(spy), nil)
	conn.connect()
	conn.Shutdown()
	conn.Shutdown()

	if want, got := 1, len(spy.closeCalls); want != got {
		t.Errorf("Wrong 'Close' callcount, got %d, want %d", got, want)
	}
}

func TestDurableConnectionWhenConnectionBreaksDown(t *testing.T) {
	const testPayload = "test"
	reconnectChan := make(chan struct{})

	upgrader := websocket.Upgrader{}
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		c, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}

		// Waits for a message to be sent before dropping the connection.
		<-reconnectChan
		c.Close()
	}))
	defer s.Close()

	logger := ktesting.TestLogger(t)
	target := "ws" + strings.TrimPrefix(s.URL, "http")
	conn := NewDurableSendingConnection(target, logger)
	defer conn.Shutdown()

	for range 10 {
		err := wait.PollUntilContextTimeout(context.Background(), 50*time.Millisecond, 5*time.Second, true, func(ctx context.Context) (bool, error) {
			if err := conn.Send(testPayload); err != nil {
				return false, nil //nolint:nilerr
			}
			return true, nil
		})
		if err != nil {
			t.Error("Timed out trying to send a message:", err)
		}

		// Message successfully sent, instruct the server to drop the connection.
		reconnectChan <- struct{}{}
	}
}

func TestDurableConnectionSendsPingsRegularly(t *testing.T) {
	// Reset pongTimeout to something quite short.
	pingTimeoutBackup := pongTimeout
	pongTimeout = 100 * time.Millisecond
	t.Cleanup(func() {
		pongTimeout = pingTimeoutBackup
	})

	upgrader := websocket.Upgrader{}

	pingReceived := make(chan struct{})
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		c, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}

		c.SetPingHandler(func(_ string) error {
			pingReceived <- struct{}{}
			return c.WriteMessage(websocket.PongMessage, []byte{})
		})

		for {
			_, _, err := c.ReadMessage()
			if err != nil {
				break
			}
		}
	}))
	defer s.Close()

	logger := ktesting.TestLogger(t)
	target := "ws" + strings.TrimPrefix(s.URL, "http")
	conn := NewDurableSendingConnection(target, logger)
	defer conn.Shutdown()

	// Wait for 5 pings to be received by the server.
	for range 5 {
		<-pingReceived
	}
}

func TestOnConnectAndOnDisconnectCallbacks(t *testing.T) {
	reconnectChan := make(chan struct{})

	upgrader := websocket.Upgrader{}
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		c, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}

		// Wait for signal to drop the connection.
		<-reconnectChan
		c.Close()
	}))
	defer s.Close()

	logger := ktesting.TestLogger(t)
	target := "ws" + strings.TrimPrefix(s.URL, "http")

	onConnectCalled := make(chan struct{}, 10)
	onDisconnectCalled := make(chan error, 10)

	conn := NewDurableSendingConnection(target, logger,
		WithOnConnect(func() {
			onConnectCalled <- struct{}{}
		}),
		WithOnDisconnect(func(err error) {
			onDisconnectCalled <- err
		}),
	)
	defer conn.Shutdown()

	// Wait for the first OnConnect call
	select {
	case <-onConnectCalled:
		// Success - OnConnect was called
	case <-time.After(propagationTimeout):
		t.Fatal("Timed out waiting for OnConnect to be called")
	}

	// Trigger a disconnect by closing the server-side connection
	reconnectChan <- struct{}{}

	// Wait for OnDisconnect to be called
	select {
	case err := <-onDisconnectCalled:
		if err == nil {
			t.Error("Expected OnDisconnect to receive an error, got nil")
		}
	case <-time.After(propagationTimeout):
		t.Fatal("Timed out waiting for OnDisconnect to be called")
	}

	// Wait for reconnection (OnConnect should be called again)
	select {
	case <-onConnectCalled:
		// Success - OnConnect was called again after reconnection
	case <-time.After(propagationTimeout):
		t.Fatal("Timed out waiting for OnConnect to be called after reconnection")
	}
}

func TestOnConnectAndOnDisconnectCallbacksNotSet(t *testing.T) {
	reconnectChan := make(chan struct{})

	upgrader := websocket.Upgrader{}
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		c, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}

		// Wait for signal to drop the connection.
		<-reconnectChan
		c.Close()
	}))
	defer s.Close()

	logger := ktesting.TestLogger(t)
	target := "ws" + strings.TrimPrefix(s.URL, "http")

	// Create connection without setting callbacks - should not panic
	conn := NewDurableSendingConnection(target, logger)
	defer conn.Shutdown()

	// Wait for connection to be established
	err := wait.PollUntilContextTimeout(context.Background(), 50*time.Millisecond, propagationTimeout, true, func(ctx context.Context) (bool, error) {
		return conn.Status() == nil, nil
	})
	if err != nil {
		t.Fatal("Timed out waiting for connection to be established:", err)
	}

	// Trigger disconnect - should not panic even without callbacks
	reconnectChan <- struct{}{}

	// Wait a bit and verify no panic occurred
	time.Sleep(100 * time.Millisecond)
}

func TestNewDurableSendingConnectionGuaranteed(t *testing.T) {
	// Unhappy case.
	logger := ktesting.TestLogger(t)
	_, err := NewDurableSendingConnectionGuaranteed("ws://somewhere.not.exist", time.Second, logger)
	if got, want := err.Error(), ErrConnectionNotEstablished.Error(); got != want {
		t.Errorf("Got error: %v, want error: %v", got, want)
	}

	// Happy case.
	const testPayload = "test"
	reconnectChan := make(chan struct{})
	upgrader := websocket.Upgrader{}
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		c, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}

		// Waits for a message to be sent before dropping the connection.
		<-reconnectChan
		c.Close()
	}))
	defer s.Close()

	target := "ws" + strings.TrimPrefix(s.URL, "http")
	conn, err := NewDurableSendingConnectionGuaranteed(target, time.Second, logger)
	if err != nil {
		t.Error("Got error from NewDurableSendingConnectionGuaranteed:", err)
	}
	defer conn.Shutdown()

	// Sending the message immediately should be fine as the connection has been established.
	if err := conn.Send(testPayload); err != nil {
		t.Error("Failed to send a message:", err)
	}

	// Message successfully sent, instruct the server to drop the connection.
	reconnectChan <- struct{}{}
}