nadir-agent/internal/modules/terminal/terminal_test.go

package terminal

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

	"nadir/internal/auth"

	"github.com/coder/websocket"
	"github.com/danielgtaylor/huma/v2"
	"github.com/danielgtaylor/huma/v2/adapters/humago"
)

func TestTerminalConnectUnauthorized(t *testing.T) {
	sessions, err := auth.NewSessionStore("file::memory:?cache=shared")
	if err != nil {
		t.Fatalf("session db: %v", err)
	}

	mux := http.NewServeMux()
	api := humago.New(mux, huma.DefaultConfig("Test", "1.0.0"))
	mod := New(sessions)
	mod.Register(api)

	srv := httptest.NewServer(mux)
	defer srv.Close()

	wsURL := strings.Replace(srv.URL, "http://", "ws://", 1) + "/api/terminal"

	// Try without cookie
	ctx, cancel := context.WithTimeout(t.Context(), 2*time.Second)
	defer cancel()

	_, resp, err := websocket.Dial(ctx, wsURL, nil)
	if err == nil {
		t.Fatal("expected error, got success")
	}
	if resp != nil && resp.StatusCode != http.StatusUnauthorized {
		t.Errorf("expected 401, got %d", resp.StatusCode)
	}
}

func TestTerminalConnectPlainGET(t *testing.T) {
	// An authenticated but non-WebSocket GET (e.g. the docs "Try it" button) must
	// get a clean 426 Upgrade Required, not a raw websocket protocol error.
	sessions, err := auth.NewSessionStore("file::memory:?cache=shared")
	if err != nil {
		t.Fatalf("session db: %v", err)
	}
	token, err := sessions.Create("root")
	if err != nil {
		t.Fatalf("create session: %v", err)
	}

	mux := http.NewServeMux()
	api := humago.New(mux, huma.DefaultConfig("Test", "1.0.0"))
	New(sessions).Register(api)

	srv := httptest.NewServer(mux)
	defer srv.Close()

	req, _ := http.NewRequest(http.MethodGet, srv.URL+"/api/terminal", nil)
	req.Header.Set("Cookie", "nadir_session_id="+token)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("GET failed: %v", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusUpgradeRequired {
		t.Errorf("expected 426, got %d", resp.StatusCode)
	}
}

func TestTerminalConnectAuthorized(t *testing.T) {
	// Root or specific system configs might cause PTY or su to fail if run in constrained environments.
	// But we expect the websocket upgrade to succeed and then maybe close with an error if PTY fails.
	sessions, err := auth.NewSessionStore("file::memory:?cache=shared")
	if err != nil {
		t.Fatalf("session db: %v", err)
	}

	token, err := sessions.Create("root")
	if err != nil {
		t.Fatalf("create session: %v", err)
	}

	mux := http.NewServeMux()
	api := humago.New(mux, huma.DefaultConfig("Test", "1.0.0"))
	mod := New(sessions)
	mod.Register(api)

	srv := httptest.NewServer(mux)
	defer srv.Close()

	wsURL := strings.Replace(srv.URL, "http://", "ws://", 1) + "/api/terminal"

	ctx, cancel := context.WithTimeout(t.Context(), 2*time.Second)
	defer cancel()

	opts := &websocket.DialOptions{
		HTTPHeader: http.Header{},
	}
	opts.HTTPHeader.Set("Cookie", "nadir_session_id="+token)

	conn, resp, err := websocket.Dial(ctx, wsURL, opts)
	if err != nil {
		// Depending on the test environment, "su" or "pty" might fail, but
		// the websocket upgrade itself should succeed before it drops.
		// If it fails to upgrade, that's a real error.
		t.Fatalf("dial failed: %v (status %d)", err, resp.StatusCode)
	}
	defer conn.CloseNow()

	// The connection was upgraded successfully.
	// If PTY allocation failed or su failed, the server might close the connection immediately.
	// We just verify the upgrade succeeded.
	if resp.StatusCode != http.StatusSwitchingProtocols {
		t.Errorf("expected 101, got %d", resp.StatusCode)
	}
}