xTeVe/src/websocket_test.go

package src

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/gorilla/websocket"

	"xteve/src/internal/authentication"
)

// newAuthenticatedWSServer starts the websocket handler with web
// authentication enabled and returns the server plus a valid session token.
func newAuthenticatedWSServer(t *testing.T) (*httptest.Server, string) {
	t.Helper()

	if err := authentication.Init(t.TempDir()+"/authentication.json", 60); err != nil {
		t.Fatalf("authentication.Init: %v", err)
	}
	if _, err := authentication.CreateNewUser("admin", "secret"); err != nil {
		t.Fatalf("CreateNewUser: %v", err)
	}
	token, err := authentication.UserAuthentication("admin", "secret")
	if err != nil {
		t.Fatalf("UserAuthentication: %v", err)
	}

	Settings.AuthenticationWEB = true
	System.ConfigurationWizard = false
	t.Cleanup(func() { Settings.AuthenticationWEB = false })

	srv := httptest.NewServer(http.HandlerFunc(WS))
	t.Cleanup(srv.Close)
	return srv, token
}

func wsURL(srv *httptest.Server) string {
	return "ws" + strings.TrimPrefix(srv.URL, "http") + "/data/"
}

// roundTrip sends one command and returns the parsed response.
func roundTrip(t *testing.T, url string, header http.Header) ResponseStruct {
	t.Helper()

	conn, resp, err := websocket.DefaultDialer.Dial(url, header)
	if err != nil {
		t.Fatalf("dial: %v (resp=%v)", err, resp)
	}
	defer conn.Close()

	if err := conn.WriteJSON(map[string]any{"cmd": "noop"}); err != nil {
		t.Fatalf("write: %v", err)
	}
	var response ResponseStruct
	if err := conn.ReadJSON(&response); err != nil {
		t.Fatalf("read: %v", err)
	}
	return response
}

func TestWSRejectsCrossOrigin(t *testing.T) {
	srv, _ := newAuthenticatedWSServer(t)

	header := http.Header{"Origin": {"http://evil.example"}}
	_, resp, err := websocket.DefaultDialer.Dial(wsURL(srv), header)
	if err == nil {
		t.Fatal("expected the cross-origin upgrade to be refused")
	}
	if resp == nil || resp.StatusCode != http.StatusForbidden {
		t.Fatalf("expected 403, got %v", resp)
	}
}

func TestWSAllowsSameOriginAndNoOrigin(t *testing.T) {
	srv, token := newAuthenticatedWSServer(t)

	sameOrigin := http.Header{
		"Origin": {srv.URL},
		"Cookie": {"Token=" + token},
	}
	if r := roundTrip(t, wsURL(srv), sameOrigin); !r.Status {
		t.Fatalf("same-origin request with cookie should succeed, got error %q", r.Error)
	}

	noOrigin := http.Header{"Cookie": {"Token=" + token}}
	if r := roundTrip(t, wsURL(srv), noOrigin); !r.Status {
		t.Fatalf("request without Origin (non-browser client) should succeed, got error %q", r.Error)
	}
}

func TestWSRequiresSessionToken(t *testing.T) {
	srv, token := newAuthenticatedWSServer(t)

	r := roundTrip(t, wsURL(srv), nil)
	if r.Status || !r.Reload {
		t.Fatalf("request without a token should be refused with reload, got status=%v reload=%v", r.Status, r.Reload)
	}

	bad := http.Header{"Cookie": {"Token=not-a-real-token"}}
	if r := roundTrip(t, wsURL(srv), bad); r.Status {
		t.Fatal("request with an unknown token should be refused")
	}

	// Legacy clients may still pass the token as a query parameter.
	if r := roundTrip(t, wsURL(srv)+"?Token="+token, nil); !r.Status {
		t.Fatalf("legacy query-parameter token should be accepted, got error %q", r.Error)
	}
}

func TestSetCookieTokenFlags(t *testing.T) {
	rec := httptest.NewRecorder()
	authentication.SetCookieToken(rec, "abc")

	cookies := rec.Result().Cookies()
	if len(cookies) != 1 {
		t.Fatalf("expected one cookie, got %d", len(cookies))
	}
	c := cookies[0]
	if c.Name != "Token" || c.Value != "abc" {
		t.Fatalf("unexpected cookie %v", c)
	}
	if !c.HttpOnly {
		t.Error("session cookie must be HttpOnly")
	}
	if c.SameSite != http.SameSiteStrictMode {
		t.Error("session cookie must be SameSite=Strict")
	}
	if c.Path != "/" {
		t.Errorf("cookie path should be /, got %q", c.Path)
	}

	rec = httptest.NewRecorder()
	authentication.SetCookieToken(rec, "-")
	if c := rec.Result().Cookies()[0]; c.MaxAge >= 0 {
		t.Error("logout must clear the cookie (negative MaxAge)")
	}
}