agent-helm/helmd/hub_test.go

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func testHubConfig(t *testing.T) *HubConfig {
	t.Helper()
	cfg := defaultHubConfig()
	cfg.path = filepath.Join(t.TempDir(), "hub.json")
	cfg.AgentKey = "test-agent-key"
	if err := cfg.SetPassword("admin", "hemligt"); err != nil {
		t.Fatal(err)
	}
	return cfg
}

func loginToken(t *testing.T, srv *httptest.Server, user, pass string) (string, int) {
	t.Helper()
	body, _ := json.Marshal(map[string]string{"username": user, "password": pass})
	resp, err := http.Post(srv.URL+"/api/login", "application/json", bytes.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	var out struct {
		Token string `json:"token"`
	}
	json.NewDecoder(resp.Body).Decode(&out)
	return out.Token, resp.StatusCode
}

func authGet(t *testing.T, srv *httptest.Server, token, path string) (*http.Response, []byte) {
	t.Helper()
	req, _ := http.NewRequest("GET", srv.URL+path, nil)
	req.Header.Set("Authorization", "Bearer "+token)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	data, _ := io.ReadAll(resp.Body)
	return resp, data
}

func TestPasswordHashing(t *testing.T) {
	cfg := testHubConfig(t)
	if !cfg.Verify("admin", "hemligt") {
		t.Error("rätt lösenord underkändes")
	}
	if cfg.Verify("admin", "fel") {
		t.Error("fel lösenord godkändes")
	}
	if cfg.Verify("ingen", "hemligt") {
		t.Error("okänd användare godkändes")
	}
	// uppdatering av befintligt konto ska inte skapa dubblett
	if err := cfg.SetPassword("admin", "nytt"); err != nil {
		t.Fatal(err)
	}
	if len(cfg.Users) != 1 {
		t.Errorf("förväntade 1 konto, fick %d", len(cfg.Users))
	}
	if cfg.Verify("admin", "hemligt") || !cfg.Verify("admin", "nytt") {
		t.Error("lösenordsbyte slog inte igenom")
	}
}

func TestHubLoginAndAuth(t *testing.T) {
	h := NewHub(testHubConfig(t))
	mux := h.routes()
	mountWeb(mux) // fångar mux-mönsterkonflikter med webbklienten (panikar annars)
	srv := httptest.NewServer(mux)
	defer srv.Close()

	if _, code := loginToken(t, srv, "admin", "fel"); code != 401 {
		t.Errorf("fel lösenord: förväntade 401, fick %d", code)
	}
	tok, code := loginToken(t, srv, "admin", "hemligt")
	if code != 200 || tok == "" {
		t.Fatalf("inloggning misslyckades: %d", code)
	}

	if resp, _ := authGet(t, srv, "", "/api/agents"); resp.StatusCode != 401 {
		t.Errorf("utan token: förväntade 401, fick %d", resp.StatusCode)
	}
	resp, data := authGet(t, srv, tok, "/api/agents")
	if resp.StatusCode != 200 {
		t.Errorf("med token: förväntade 200, fick %d", resp.StatusCode)
	}
	if strings.TrimSpace(string(data)) != "[]" {
		t.Errorf("förväntade tom agentlista, fick %s", data)
	}

	// agent-nyckeln nås efter inloggning (kravet)
	resp, data = authGet(t, srv, tok, "/api/agentkey")
	if resp.StatusCode != 200 || !strings.Contains(string(data), "test-agent-key") {
		t.Errorf("agentkey: %d %s", resp.StatusCode, data)
	}
}

func TestHubRegisterRequiresAgentKey(t *testing.T) {
	h := NewHub(testHubConfig(t))
	srv := httptest.NewServer(h.routes())
	defer srv.Close()

	body, _ := json.Marshal(map[string]string{"name": "a1", "host": "h", "version": "t", "key": "fel-nyckel"})
	resp, err := http.Post(srv.URL+"/hub/register", "application/json", bytes.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	resp.Body.Close()
	if resp.StatusCode != 401 {
		t.Errorf("fel agent-nyckel: förväntade 401, fick %d", resp.StatusCode)
	}
}

// TestHubUplinkRoundtrip kör riktiga uplink-koden mot en hub och en
// stub-mux och verifierar hela kedjan: registrering → jobb → svar,
// inklusive att headers (CSP) följer med och att events vidarebefordras.
func TestHubUplinkRoundtrip(t *testing.T) {
	h := NewHub(testHubConfig(t))
	srv := httptest.NewServer(h.routes())

	// stub för agentens lokala API: kräver Bearer-token som en riktig helmd
	stub := http.NewServeMux()
	stub.HandleFunc("GET /api/ping", func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer lokal-token" {
			jsonErr(w, 401, "saknad token")
			return
		}
		w.Header().Set("Content-Security-Policy", "sandbox")
		jsonOut(w, 200, map[string]string{"pong": r.URL.Query().Get("x")})
	})
	stub.HandleFunc("POST /api/echo", func(w http.ResponseWriter, r *http.Request) {
		data, _ := io.ReadAll(r.Body)
		w.Header().Set("Content-Type", r.Header.Get("Content-Type"))
		w.WriteHeader(201)
		w.Write(data)
	})

	u := &uplink{
		url: srv.URL, name: "a1", host: "testhost", agentKey: "test-agent-key",
		token: "lokal-token", handler: stub,
		client: &http.Client{Timeout: 40 * time.Second},
		stop:   make(chan struct{}),
	}
	go u.run()
	// stäng uplinken FÖRE servern — annars ligger zombie-pollers kvar
	// och kan registrera sig på en senare testservers återanvända port
	t.Cleanup(func() {
		u.shutdown()
		srv.CloseClientConnections() // avbryt hängande long-polls direkt
		srv.Close()
	})

	tok, _ := loginToken(t, srv, "admin", "hemligt")

	// vänta tills agenten är online
	deadline := time.Now().Add(5 * time.Second)
	for {
		_, data := authGet(t, srv, tok, "/api/agents")
		var agents []hubAgent
		json.Unmarshal(data, &agents)
		if len(agents) == 1 && agents[0].Online {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("agenten kom aldrig online: %s", data)
		}
		time.Sleep(50 * time.Millisecond)
	}

	// GET med query proxas och headers följer med
	resp, data := authGet(t, srv, tok, "/api/agents/a1/ping?x=42")
	if resp.StatusCode != 200 || !strings.Contains(string(data), `"pong":"42"`) {
		t.Errorf("proxy GET: %d %s", resp.StatusCode, data)
	}
	if resp.Header.Get("Content-Security-Policy") != "sandbox" {
		t.Errorf("CSP-headern följde inte med genom hubben: %q", resp.Header.Get("Content-Security-Policy"))
	}

	// POST med body proxas med status + content-type
	req, _ := http.NewRequest("POST", srv.URL+"/api/agents/a1/echo", strings.NewReader("hej hubben"))
	req.Header.Set("Authorization", "Bearer "+tok)
	req.Header.Set("Content-Type", "text/plain")
	presp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	pdata, _ := io.ReadAll(presp.Body)
	presp.Body.Close()
	if presp.StatusCode != 201 || string(pdata) != "hej hubben" {
		t.Errorf("proxy POST: %d %q", presp.StatusCode, pdata)
	}
	if ct := presp.Header.Get("Content-Type"); !strings.HasPrefix(ct, "text/plain") {
		t.Errorf("content-type tappades: %q", ct)
	}

	// events-proxy ska vara blockerad per agent
	resp, _ = authGet(t, srv, tok, "/api/agents/a1/events")
	if resp.StatusCode != 404 {
		t.Errorf("events-proxy: förväntade 404, fick %d", resp.StatusCode)
	}

	// okänd agent
	resp, _ = authGet(t, srv, tok, "/api/agents/finns-inte/ping")
	if resp.StatusCode != 404 {
		t.Errorf("okänd agent: förväntade 404, fick %d", resp.StatusCode)
	}

	// events från agenten dyker upp i hubbens ström med agent-fältet satt
	sub := h.subscribe()
	defer h.unsubscribe(sub)
	evBody, _ := json.Marshal(map[string]interface{}{"events": []Event{{Type: "screen", Session: "s1", Data: 7}}})
	eresp, err := http.Post(srv.URL+"/hub/events?name=a1&key="+u.key(), "application/json", bytes.NewReader(evBody))
	if err != nil {
		t.Fatal(err)
	}
	eresp.Body.Close()
	select {
	case ev := <-sub:
		if ev.Agent != "a1" || ev.Type != "screen" || ev.Session != "s1" {
			t.Errorf("oväntat event: %+v", ev)
		}
	case <-time.After(2 * time.Second):
		t.Error("inget event nådde hubbens ström")
	}
}

// TestHubAgentIsolation: /hub-ytan får inte läcka något om andra
// agenter, och fel pollnyckel ska avvisas.
func TestHubAgentIsolation(t *testing.T) {
	h := NewHub(testHubConfig(t))
	srv := httptest.NewServer(h.routes())
	defer srv.Close()

	register := func(name string) string {
		body, _ := json.Marshal(map[string]string{"name": name, "host": "h", "version": "t", "key": "test-agent-key"})
		resp, err := http.Post(srv.URL+"/hub/register", "application/json", bytes.NewReader(body))
		if err != nil {
			t.Fatal(err)
		}
		defer resp.Body.Close()
		var out struct {
			Key string `json:"key"`
		}
		json.NewDecoder(resp.Body).Decode(&out)
		return out.Key
	}
	k1 := register("a1")
	register("a2")

	// a1:s nyckel funkar inte för a2:s kanal
	resp, err := http.Get(srv.URL + "/hub/work?name=a2&key=" + k1)
	if err != nil {
		t.Fatal(err)
	}
	resp.Body.Close()
	if resp.StatusCode != 401 {
		t.Errorf("korsad pollnyckel: förväntade 401, fick %d", resp.StatusCode)
	}

	// registreringssvaret innehåller bara den egna nyckeln — inget om andra
	body, _ := json.Marshal(map[string]string{"name": "a3", "host": "h", "version": "t", "key": "test-agent-key"})
	rresp, err := http.Post(srv.URL+"/hub/register", "application/json", bytes.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	rdata, _ := io.ReadAll(rresp.Body)
	rresp.Body.Close()
	// exakt ett fält (key) — inga agentlistor eller annan info; en
	// substrängkoll mot agentnamn funkar inte (slumpad hex innehåller
	// gärna "a1")
	var full map[string]interface{}
	json.Unmarshal(rdata, &full)
	if len(full) != 1 || full["key"] == nil {
		t.Errorf("registreringssvaret ska bara innehålla key: %s", rdata)
	}
}

func TestSessionStoreExpiry(t *testing.T) {
	s := newSessionStore(1 * time.Millisecond)
	tok := s.Create("admin")
	if _, ok := s.Lookup(tok); !ok {
		t.Error("färsk token underkändes")
	}
	time.Sleep(5 * time.Millisecond)
	if _, ok := s.Lookup(tok); ok {
		t.Error("utgången token godkändes")
	}
}

func TestProxyPathBuilding(t *testing.T) {
	// säkerställ att agentPath inte öppnar för path-traversal via
	// agentnamn med specialtecken
	p := agentPath("kon/stig", "sessions")
	if strings.Contains(p, "kon/stig") {
		t.Errorf("agentnamn ska path-escapas: %s", p)
	}
	if want := fmt.Sprintf("/api/agents/%s/sessions", "kon%2Fstig"); p != want {
		t.Errorf("fick %s, ville ha %s", p, want)
	}
}