125 lines
3.5 KiB
Go
125 lines
3.5 KiB
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)
|
|
}
|
|
}
|