SimpleRemoter/server/go/web/audit_test.go

package web

import (
	"bufio"
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/yuanyuanxiang/SimpleRemoter/server/go/connection"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/hub"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/logger"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/wsauth"
)

// auditTestHub builds a wsHub whose logger writes both the main log and the
// dedicated audit log to a temp directory.
func auditTestHub(t *testing.T) (wh *wsHub, devices *hub.Hub, auditPath string) {
	t.Helper()
	dir := t.TempDir()
	auditPath = filepath.Join(dir, "audit.log")
	lg := logger.New(logger.Config{
		Level:     logger.LevelInfo,
		Console:   false,
		File:      filepath.Join(dir, "server.log"),
		AuditFile: auditPath,
	})
	t.Cleanup(func() { _ = lg.Close() })

	devices = hub.New()
	wh = newWSHub(wsauth.New(), devices, nil, lg)
	t.Cleanup(wh.stop)
	return wh, devices, auditPath
}

func auditTestClient(wh *wsHub, role string) (*wsClient, string) {
	tok, err := wh.auth.IssueToken("tester", role)
	if err != nil {
		panic(err)
	}
	c := &wsClient{
		send:           make(chan wsMsg, 64),
		closed:         make(chan struct{}),
		id:             newConnID(),
		addr:           "203.0.113.5",
		peerAddr:       "203.0.113.5",
		connectedAt:    time.Now(),
		watchIDs:       make(map[string]struct{}),
		watchStartedAt: make(map[string]time.Time),
		keyframeAt:     make(map[string]int64),
	}
	wh.register(c)
	return c, tok
}

func readAuditEvents(t *testing.T, path string) []map[string]any {
	t.Helper()
	f, err := os.Open(path)
	if err != nil {
		t.Fatalf("open audit log: %v", err)
	}
	defer f.Close()

	var out []map[string]any
	sc := bufio.NewScanner(f)
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if line == "" {
			continue
		}
		var m map[string]any
		if err := json.Unmarshal([]byte(line), &m); err != nil {
			t.Fatalf("bad audit line %q: %v", line, err)
		}
		out = append(out, m)
	}
	if err := sc.Err(); err != nil {
		t.Fatalf("scan audit log: %v", err)
	}
	return out
}

func findAuditEvent(t *testing.T, path, event string) map[string]any {
	t.Helper()
	for _, m := range readAuditEvents(t, path) {
		if m["event"] == event {
			return m
		}
	}
	t.Fatalf("audit event %q not found in %v", event, readAuditEvents(t, path))
	return nil
}

func TestLoginFailureIsAudited(t *testing.T) {
	wh, _, auditPath := auditTestHub(t)
	wh.auth.AddAdminFromPlainPassword("admin", "hunter2")
	c, _ := auditTestClient(wh, "viewer")

	nonce, err := wsauth.NewNonce()
	if err != nil {
		t.Fatalf("NewNonce: %v", err)
	}
	c.nonce = nonce
	raw, _ := json.Marshal(map[string]any{
		"username": "admin",
		"response": "not-the-right-response",
		"nonce":    nonce,
	})
	wh.handleLogin(c, raw)

	rec := findAuditEvent(t, auditPath, auditLoginFail)
	if rec["reason"] != "invalid_credentials" {
		t.Errorf("reason = %v", rec["reason"])
	}
	if rec["username"] != "admin" {
		t.Errorf("username = %v", rec["username"])
	}
	if rec["user_exists"] != true {
		t.Errorf("user_exists = %v, want true", rec["user_exists"])
	}
	if rec["ip"] != "203.0.113.5" {
		t.Errorf("ip = %v", rec["ip"])
	}
	if _, ok := rec["response"]; ok {
		t.Error("audit record leaked the login response")
	}
}

func TestAdminDenialIsAudited(t *testing.T) {
	wh, _, auditPath := auditTestHub(t)
	c, tok := auditTestClient(wh, "viewer")
	raw, _ := json.Marshal(map[string]any{"token": tok})

	if wh.requireAdmin(c, raw, "delete_device_result") {
		t.Fatal("viewer passed requireAdmin")
	}
	rec := findAuditEvent(t, auditPath, auditDenied)
	if rec["operation"] != "delete_device" {
		t.Errorf("operation = %v", rec["operation"])
	}
	if rec["role"] != "viewer" || rec["actor"] != "tester" {
		t.Errorf("identity fields = %+v", rec)
	}
}

func TestScreenStartAndOfflineStopAreAudited(t *testing.T) {
	wh, devices, auditPath := auditTestHub(t)
	devices.SetSender(func(*connection.Context, []byte) error { return nil })
	hctx := &connection.Context{RemoteAddr: "10.0.0.9:1234"}
	devices.Register(&hub.Device{ID: "dev1", Name: "Dev"}, hctx)

	c, tok := auditTestClient(wh, "admin")
	raw, _ := json.Marshal(map[string]any{
		"token": tok, "id": "dev1", "multiWatch": true, "req": "1",
	})
	wh.handleConnect(c, raw)

	start := findAuditEvent(t, auditPath, auditScreenStart)
	if start["device_id"] != "dev1" || start["mode"] != "started" {
		t.Errorf("screen start = %+v", start)
	}

	wh.OnDeviceOffline("dev1")

	stop := findAuditEvent(t, auditPath, auditScreenStop)
	if stop["device_id"] != "dev1" || stop["reason"] != "device_offline" {
		t.Errorf("screen stop = %+v", stop)
	}
	if _, ok := stop["duration_ms"]; !ok {
		t.Error("screen stop has no duration_ms")
	}
}

func TestUnknownCommandRecordedOnce(t *testing.T) {
	wh, _, auditPath := auditTestHub(t)
	c, _ := auditTestClient(wh, "viewer")

	for i := 0; i < 5; i++ {
		wh.dispatch(c, "definitely_not_a_command", []byte(`{}`))
	}
	count := 0
	for _, m := range readAuditEvents(t, auditPath) {
		if m["event"] == auditUnknownCmd {
			count++
		}
	}
	if count != 1 {
		t.Fatalf("unknown-command audit records = %d, want 1", count)
	}
	if c.statUnknownCmd != 5 {
		t.Fatalf("statUnknownCmd = %d, want 5", c.statUnknownCmd)
	}
}