NexusGate/agent-go/sessiontoken_test.go

package main

// Item 25: agent-side session token handling. A fake panel (httptest)
// issues tokens via response headers exactly like the real one; the
// tests pin down capture, persistence, reuse, the 401 fallback to the
// key, and old-panel behavior (no headers → key on every request).

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"
	"time"
)

// fakePanel records the bearer of every request and answers like the
// real panel: the key always authenticates (and issues issueToken
// when non-empty), the current valid token authenticates, anything
// else is a 401.
type fakePanel struct {
	key        string
	validToken string
	issueToken string
	bearers    []string
}

func (p *fakePanel) handler(w http.ResponseWriter, r *http.Request) {
	bearer := ""
	if h := r.Header.Get("Authorization"); len(h) > len("Bearer ") {
		bearer = h[len("Bearer "):]
	}
	p.bearers = append(p.bearers, bearer)
	switch bearer {
	case p.key:
		if p.issueToken != "" {
			w.Header().Set("X-Agent-Session-Token", p.issueToken)
			w.Header().Set("X-Agent-Session-Token-Expires", time.Now().Add(24*time.Hour).UTC().Format("2006-01-02T15:04:05.000Z"))
		}
		w.WriteHeader(200)
		fmt.Fprint(w, `{"ok":true}`)
	case p.validToken:
		w.WriteHeader(200)
		fmt.Fprint(w, `{"ok":true}`)
	default:
		w.WriteHeader(401)
		fmt.Fprint(w, `{"message":"Agent 身份校验失败"}`)
	}
}

func setupTokenAgent(t *testing.T, panel *fakePanel) func() {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(panel.handler))
	oldCfg := cfg
	cfg.Controller = srv.URL
	cfg.AgentKey = panel.key
	cfg.SessionTokenFile = filepath.Join(t.TempDir(), "session-token.json")
	sessionTokenState.mu.Lock()
	oldState := sessionTokenState.value
	oldLoaded := sessionTokenState.loaded
	sessionTokenState.value = nil
	sessionTokenState.loaded = false
	sessionTokenState.mu.Unlock()
	return func() {
		srv.Close()
		cfg = oldCfg
		sessionTokenState.mu.Lock()
		sessionTokenState.value = oldState
		sessionTokenState.loaded = oldLoaded
		sessionTokenState.mu.Unlock()
	}
}

func TestSessionTokenCapturedPersistedAndReused(t *testing.T) {
	panel := &fakePanel{key: "the-key", issueToken: "tok-1", validToken: "tok-1"}
	defer setupTokenAgent(t, panel)()

	if _, err := request("/api/agent/poll", []byte("{}")); err != nil {
		t.Fatalf("first request: %v", err)
	}
	if _, err := request("/api/agent/poll", []byte("{}")); err != nil {
		t.Fatalf("second request: %v", err)
	}
	want := []string{"the-key", "tok-1"}
	if fmt.Sprint(panel.bearers) != fmt.Sprint(want) {
		t.Fatalf("bearers = %v, want %v", panel.bearers, want)
	}
	// Persisted with keys.json-grade permissions.
	info, err := os.Stat(cfg.SessionTokenFile)
	if err != nil {
		t.Fatalf("token file not persisted: %v", err)
	}
	if perm := info.Mode().Perm(); perm != 0o600 {
		t.Fatalf("token file perm = %o, want 600", perm)
	}
	// And it survives a cold start: reset memory, it must load from disk.
	sessionTokenState.mu.Lock()
	sessionTokenState.value = nil
	sessionTokenState.loaded = false
	sessionTokenState.mu.Unlock()
	if tok := currentSessionToken(); tok == nil || tok.Token != "tok-1" {
		t.Fatalf("token not reloaded from disk: %+v", tok)
	}
}

func TestSessionToken401FallsBackToKeyAndReissues(t *testing.T) {
	panel := &fakePanel{key: "the-key", issueToken: "tok-2", validToken: "tok-2"}
	defer setupTokenAgent(t, panel)()
	// Plant a stale token the panel no longer honors.
	storeSessionToken("tok-stale", time.Now().Add(time.Hour).UTC().Format("2006-01-02T15:04:05.000Z"))

	if _, err := request("/api/agent/poll", []byte("{}")); err != nil {
		t.Fatalf("request with stale token: %v", err)
	}
	want := []string{"tok-stale", "the-key"}
	if fmt.Sprint(panel.bearers) != fmt.Sprint(want) {
		t.Fatalf("bearers = %v, want %v (token attempt, then key retry)", panel.bearers, want)
	}
	if tok := currentSessionToken(); tok == nil || tok.Token != "tok-2" {
		t.Fatalf("fresh token not stored after retry: %+v", tok)
	}
}

func TestSessionTokenAbsentKeepsKeyAuth(t *testing.T) {
	// Old panel: never issues tokens.
	panel := &fakePanel{key: "the-key"}
	defer setupTokenAgent(t, panel)()
	for i := 0; i < 2; i++ {
		if _, err := request("/api/agent/poll", []byte("{}")); err != nil {
			t.Fatalf("request %d: %v", i, err)
		}
	}
	if fmt.Sprint(panel.bearers) != "[the-key the-key]" {
		t.Fatalf("bearers = %v, want key on every request", panel.bearers)
	}
	if _, err := os.Stat(cfg.SessionTokenFile); !os.IsNotExist(err) {
		t.Fatalf("no token file expected against an old panel, stat err = %v", err)
	}
}

func TestSessionTokenExpiredClientSideUsesKey(t *testing.T) {
	panel := &fakePanel{key: "the-key", issueToken: "tok-3", validToken: "tok-3"}
	defer setupTokenAgent(t, panel)()
	storeSessionToken("tok-old", time.Now().Add(-time.Hour).UTC().Format("2006-01-02T15:04:05.000Z"))
	if _, err := request("/api/agent/poll", []byte("{}")); err != nil {
		t.Fatalf("request: %v", err)
	}
	if len(panel.bearers) != 1 || panel.bearers[0] != "the-key" {
		t.Fatalf("bearers = %v, want the key used directly (no doomed token attempt)", panel.bearers)
	}
}

func TestSessionTokenKeyRejectedSurfacesError(t *testing.T) {
	panel := &fakePanel{key: "the-real-key"}
	defer setupTokenAgent(t, panel)()
	cfg.AgentKey = "wrong-key"
	storeSessionToken("tok-stale", time.Now().Add(time.Hour).UTC().Format("2006-01-02T15:04:05.000Z"))
	if _, err := request("/api/agent/poll", []byte("{}")); err == nil {
		t.Fatal("expected an error when both token and key are rejected")
	}
	// Exactly two attempts: token, then key — never a loop.
	if len(panel.bearers) != 2 {
		t.Fatalf("attempts = %v, want exactly 2", panel.bearers)
	}
	if tok := currentSessionToken(); tok != nil {
		t.Fatalf("stale token should have been cleared, still held: %+v", tok)
	}
}