imap-copier/internal/httpapi/ws_test.go

package httpapi

import (
	"context"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/coder/websocket"
	"github.com/vasyansk/imap-copier/internal/config"
	"github.com/vasyansk/imap-copier/internal/crypto"
	"github.com/vasyansk/imap-copier/internal/wshub"
)

func TestWSRequiresAuth(t *testing.T) {
	s := &Server{cfg: config.Config{SessionSecret: []byte("x")}, hub: wshub.New()}
	srv := httptest.NewServer(s.Router())
	defer srv.Close()
	// no cookie -> upgrade rejected (401)
	_, resp, err := websocket.Dial(context.Background(), "ws"+srv.URL[4:]+"/ws?task_id=1", nil)
	if err == nil {
		t.Fatal("expected auth rejection")
	}
	if resp != nil && resp.StatusCode != http.StatusUnauthorized {
		t.Fatalf("want 401, got %d", resp.StatusCode)
	}
}

func TestWSUnsubscribesOnClientDisconnect(t *testing.T) {
	hub := wshub.New()
	secret := []byte("sekret")
	s := &Server{cfg: config.Config{AuthUser: "admin", SessionSecret: secret}, hub: hub}
	srv := httptest.NewServer(s.Router())
	defer srv.Close()

	tok := crypto.SignSession(secret, "admin", time.Now().Add(time.Hour))
	hdr := http.Header{}
	hdr.Set("Cookie", cookieName+"="+tok)

	ctx := context.Background()
	c, _, err := websocket.Dial(ctx, "ws"+srv.URL[4:]+"/ws?task_id=7", &websocket.DialOptions{HTTPHeader: hdr})
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	// wait until subscribed
	deadline := time.Now().Add(2 * time.Second)
	for hub.SubscriberCount(7) == 0 {
		if time.Now().After(deadline) {
			t.Fatal("never subscribed")
		}
		time.Sleep(10 * time.Millisecond)
	}
	// abrupt client close -> server must detect and unsubscribe
	c.CloseNow()
	deadline = time.Now().Add(3 * time.Second)
	for hub.SubscriberCount(7) != 0 {
		if time.Now().After(deadline) {
			t.Fatal("subscription leaked after client disconnect")
		}
		time.Sleep(20 * time.Millisecond)
	}
}