From 2cd6d46f7cd7639446ced882dbfd0dc3ffed5b7c Mon Sep 17 00:00:00 2001 From: henrygd Date: Sun, 11 Jan 2026 15:03:33 -0500 Subject: [PATCH] add option to make universal token permanent (#1097, 1614) --- internal/hub/agent_connect.go | 9 ++ internal/hub/agent_connect_test.go | 100 ++++++++++++++++++ internal/hub/hub.go | 85 +++++++++++++-- internal/hub/hub_test.go | 13 ++- ...=> 0_collections_snapshot_0_18_0_dev_2.go} | 68 ++++++++++++ .../routes/settings/tokens-fingerprints.tsx | 85 ++++++++++----- 6 files changed, 326 insertions(+), 34 deletions(-) rename internal/migrations/{0_collections_snapshot_0_18_0_dev_1.go => 0_collections_snapshot_0_18_0_dev_2.go} (95%) diff --git a/internal/hub/agent_connect.go b/internal/hub/agent_connect.go index 736e2d8d..bf690c6b 100644 --- a/internal/hub/agent_connect.go +++ b/internal/hub/agent_connect.go @@ -66,6 +66,15 @@ func (acr *agentConnectRequest) agentConnect() (err error) { // Check if token is an active universal token acr.userId, acr.isUniversalToken = universalTokenMap.GetMap().GetOk(acr.token) + if !acr.isUniversalToken { + // Fallback: check for a permanent universal token stored in the DB + if rec, err := acr.hub.FindFirstRecordByFilter("universal_tokens", "token = {:token}", dbx.Params{"token": acr.token}); err == nil { + if userID := rec.GetString("user"); userID != "" { + acr.userId = userID + acr.isUniversalToken = true + } + } + } // Find matching fingerprint records for this token fpRecords := getFingerprintRecordsByToken(acr.token, acr.hub) diff --git a/internal/hub/agent_connect_test.go b/internal/hub/agent_connect_test.go index c114e4c9..40f66224 100644 --- a/internal/hub/agent_connect_test.go +++ b/internal/hub/agent_connect_test.go @@ -1169,6 +1169,106 @@ func TestMultipleSystemsWithSameUniversalToken(t *testing.T) { } } +// TestPermanentUniversalTokenFromDB verifies that a universal token persisted in the DB +// (universal_tokens collection) is accepted for agent self-registration even if it is not +// present in the in-memory universalTokenMap. +func TestPermanentUniversalTokenFromDB(t *testing.T) { + // Create hub and test app + hub, testApp, err := createTestHub(t) + require.NoError(t, err) + defer testApp.Cleanup() + + // Get the hub's SSH key + hubSigner, err := hub.GetSSHKey("") + require.NoError(t, err) + goodPubKey := hubSigner.PublicKey() + + // Create test user + userRecord, err := createTestUser(testApp) + require.NoError(t, err) + + // Create a permanent universal token record in the DB (do NOT add it to universalTokenMap) + universalToken := "db-universal-token-123" + _, err = createTestRecord(testApp, "universal_tokens", map[string]any{ + "user": userRecord.Id, + "token": universalToken, + }) + require.NoError(t, err) + + // Create HTTP server with the actual API route + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/beszel/agent-connect" { + acr := &agentConnectRequest{ + hub: hub, + req: r, + res: w, + } + acr.agentConnect() + } else { + http.NotFound(w, r) + } + })) + defer ts.Close() + + // Create and configure agent + agentDataDir := t.TempDir() + err = os.WriteFile(filepath.Join(agentDataDir, "fingerprint"), []byte("db-token-system-fingerprint"), 0644) + require.NoError(t, err) + + testAgent, err := agent.NewAgent(agentDataDir) + require.NoError(t, err) + + // Set up environment variables for the agent + os.Setenv("BESZEL_AGENT_HUB_URL", ts.URL) + os.Setenv("BESZEL_AGENT_TOKEN", universalToken) + defer func() { + os.Unsetenv("BESZEL_AGENT_HUB_URL") + os.Unsetenv("BESZEL_AGENT_TOKEN") + }() + + // Start agent in background + done := make(chan error, 1) + go func() { + serverOptions := agent.ServerOptions{ + Network: "tcp", + Addr: "127.0.0.1:46050", + Keys: []ssh.PublicKey{goodPubKey}, + } + done <- testAgent.Start(serverOptions) + }() + + // Wait for connection result + maxWait := 2 * time.Second + time.Sleep(20 * time.Millisecond) + checkInterval := 20 * time.Millisecond + timeout := time.After(maxWait) + ticker := time.Tick(checkInterval) + + connectionManager := testAgent.GetConnectionManager() + for { + select { + case <-timeout: + t.Fatalf("Expected connection to succeed but timed out - agent state: %d", connectionManager.State) + case <-ticker: + if connectionManager.State == agent.WebSocketConnected { + // Success + goto verify + } + case err := <-done: + // If Start returns early, treat it as failure + if err != nil { + t.Fatalf("Agent failed to start/connect: %v", err) + } + } + } + +verify: + // Verify that a system was created for the user (self-registration path) + systemsAfter, err := testApp.FindRecordsByFilter("systems", "users ~ {:userId}", "", -1, 0, map[string]any{"userId": userRecord.Id}) + require.NoError(t, err) + require.NotEmpty(t, systemsAfter, "Expected a system to be created for DB-backed universal token") +} + // TestFindOrCreateSystemForToken tests the findOrCreateSystemForToken function func TestFindOrCreateSystemForToken(t *testing.T) { hub, testApp, err := createTestHub(t) diff --git a/internal/hub/hub.go b/internal/hub/hub.go index 28bd07fd..c24688d0 100644 --- a/internal/hub/hub.go +++ b/internal/hub/hub.go @@ -20,6 +20,7 @@ import ( "github.com/henrygd/beszel/internal/users" "github.com/google/uuid" + "github.com/pocketbase/dbx" "github.com/pocketbase/pocketbase" "github.com/pocketbase/pocketbase/apis" "github.com/pocketbase/pocketbase/core" @@ -288,24 +289,90 @@ func (h *Hub) getUniversalToken(e *core.RequestEvent) error { userID := e.Auth.Id query := e.Request.URL.Query() token := query.Get("token") + enable := query.Get("enable") + permanent := query.Get("permanent") + // helper for deleting any existing permanent token record for this user + deletePermanent := func() error { + rec, err := h.FindFirstRecordByFilter("universal_tokens", "user = {:user}", dbx.Params{"user": userID}) + if err != nil { + return nil // no record + } + return h.Delete(rec) + } + + // helper for upserting a permanent token record for this user + upsertPermanent := func(token string) error { + rec, err := h.FindFirstRecordByFilter("universal_tokens", "user = {:user}", dbx.Params{"user": userID}) + if err == nil { + rec.Set("token", token) + return h.Save(rec) + } + + col, err := h.FindCachedCollectionByNameOrId("universal_tokens") + if err != nil { + return err + } + newRec := core.NewRecord(col) + newRec.Set("user", userID) + newRec.Set("token", token) + return h.Save(newRec) + } + + // Disable universal tokens (both ephemeral and permanent) + if enable == "0" { + tokenMap.RemovebyValue(userID) + _ = deletePermanent() + return e.JSON(http.StatusOK, map[string]any{"token": token, "active": false, "permanent": false}) + } + + // Enable universal token (ephemeral or permanent) + if enable == "1" { + if token == "" { + token = uuid.New().String() + } + + if permanent == "1" { + // make token permanent (persist across restarts) + tokenMap.RemovebyValue(userID) + if err := upsertPermanent(token); err != nil { + return err + } + return e.JSON(http.StatusOK, map[string]any{"token": token, "active": true, "permanent": true}) + } + + // default: ephemeral mode (1 hour) + _ = deletePermanent() + tokenMap.Set(token, userID, time.Hour) + return e.JSON(http.StatusOK, map[string]any{"token": token, "active": true, "permanent": false}) + } + + // Read current state + // Prefer permanent token if it exists. + if rec, err := h.FindFirstRecordByFilter("universal_tokens", "user = {:user}", dbx.Params{"user": userID}); err == nil { + dbToken := rec.GetString("token") + // If no token was provided, or the caller is asking about their permanent token, return it. + if token == "" || token == dbToken { + return e.JSON(http.StatusOK, map[string]any{"token": dbToken, "active": true, "permanent": true}) + } + // Token doesn't match their permanent token (avoid leaking other info) + return e.JSON(http.StatusOK, map[string]any{"token": token, "active": false, "permanent": false}) + } + + // No permanent token; fall back to ephemeral token map. if token == "" { // return existing token if it exists if token, _, ok := tokenMap.GetByValue(userID); ok { - return e.JSON(http.StatusOK, map[string]any{"token": token, "active": true}) + return e.JSON(http.StatusOK, map[string]any{"token": token, "active": true, "permanent": false}) } // if no token is provided, generate a new one token = uuid.New().String() } - response := map[string]any{"token": token} - switch query.Get("enable") { - case "1": - tokenMap.Set(token, userID, time.Hour) - case "0": - tokenMap.RemovebyValue(userID) - } - _, response["active"] = tokenMap.GetOk(token) + // Token is considered active only if it belongs to the current user. + activeUser, ok := tokenMap.GetOk(token) + active := ok && activeUser == userID + response := map[string]any{"token": token, "active": active, "permanent": false} return e.JSON(http.StatusOK, response) } diff --git a/internal/hub/hub_test.go b/internal/hub/hub_test.go index 2aa67f4c..e4654dad 100644 --- a/internal/hub/hub_test.go +++ b/internal/hub/hub_test.go @@ -378,7 +378,18 @@ func TestApiRoutesAuthentication(t *testing.T) { "Authorization": userToken, }, ExpectedStatus: 200, - ExpectedContent: []string{"active", "token"}, + ExpectedContent: []string{"active", "token", "permanent"}, + TestAppFactory: testAppFactory, + }, + { + Name: "GET /universal-token - enable permanent should succeed", + Method: http.MethodGet, + URL: "/api/beszel/universal-token?enable=1&permanent=1&token=permanent-token-123", + Headers: map[string]string{ + "Authorization": userToken, + }, + ExpectedStatus: 200, + ExpectedContent: []string{"\"permanent\":true", "permanent-token-123"}, TestAppFactory: testAppFactory, }, { diff --git a/internal/migrations/0_collections_snapshot_0_18_0_dev_1.go b/internal/migrations/0_collections_snapshot_0_18_0_dev_2.go similarity index 95% rename from internal/migrations/0_collections_snapshot_0_18_0_dev_1.go rename to internal/migrations/0_collections_snapshot_0_18_0_dev_2.go index 320b4eed..79800649 100644 --- a/internal/migrations/0_collections_snapshot_0_18_0_dev_1.go +++ b/internal/migrations/0_collections_snapshot_0_18_0_dev_2.go @@ -1617,6 +1617,74 @@ func init() { "type": "base", "updateRule": "", "viewRule": "@request.auth.id != \"\" && system.users.id ?= @request.auth.id" + }, + { + "createRule": null, + "deleteRule": null, + "fields": [ + { + "autogeneratePattern": "[a-z0-9]{10}", + "hidden": false, + "id": "text3208210256", + "max": 10, + "min": 10, + "name": "id", + "pattern": "^[a-z0-9]+$", + "presentable": false, + "primaryKey": true, + "required": true, + "system": true, + "type": "text" + }, + { + "cascadeDelete": true, + "collectionId": "_pb_users_auth_", + "hidden": false, + "id": "relation2375276105", + "maxSelect": 1, + "minSelect": 0, + "name": "user", + "presentable": false, + "required": true, + "system": false, + "type": "relation" + }, + { + "autogeneratePattern": "", + "hidden": false, + "id": "text1597481275", + "max": 0, + "min": 0, + "name": "token", + "pattern": "", + "presentable": false, + "primaryKey": false, + "required": false, + "system": false, + "type": "text" + }, + { + "hidden": false, + "id": "autodate2990389176", + "name": "created", + "onCreate": true, + "onUpdate": false, + "presentable": false, + "system": false, + "type": "autodate" + } + ], + "id": "pbc_3383022248", + "indexes": [ + "CREATE INDEX ` + "`" + `idx_iaD9Y2Lgbl` + "`" + ` ON ` + "`" + `universal_tokens` + "`" + ` (` + "`" + `token` + "`" + `)", + "CREATE UNIQUE INDEX ` + "`" + `idx_wdR0A4PbRG` + "`" + ` ON ` + "`" + `universal_tokens` + "`" + ` (` + "`" + `user` + "`" + `)" + ], + "listRule": null, + "name": "universal_tokens", + "system": false, + "type": "base", + "updateRule": null, + "viewRule": null } ]` diff --git a/internal/site/src/components/routes/settings/tokens-fingerprints.tsx b/internal/site/src/components/routes/settings/tokens-fingerprints.tsx index 0c23c35e..89fa17d5 100644 --- a/internal/site/src/components/routes/settings/tokens-fingerprints.tsx +++ b/internal/site/src/components/routes/settings/tokens-fingerprints.tsx @@ -32,6 +32,7 @@ import { import { AppleIcon, DockerIcon, FreeBsdIcon, TuxIcon, WindowsIcon } from "@/components/ui/icons" import { Separator } from "@/components/ui/separator" import { Switch } from "@/components/ui/switch" +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs" import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table" import { toast } from "@/components/ui/use-toast" import { isReadOnlyUser, pb } from "@/lib/api" @@ -137,21 +138,23 @@ const SectionUniversalToken = memo(() => { const [token, setToken] = useState("") const [isLoading, setIsLoading] = useState(true) const [checked, setChecked] = useState(false) + const [isPermanent, setIsPermanent] = useState(false) - async function updateToken(enable: number = -1) { + async function updateToken(enable: number = -1, permanent: number = -1) { // enable: 0 for disable, 1 for enable, -1 (unset) for get current state const data = await pb.send(`/api/beszel/universal-token`, { query: { token, enable, + permanent, }, }) setToken(data.token) setChecked(data.active) + setIsPermanent(!!data.permanent) setIsLoading(false) } - // biome-ignore lint/correctness/useExhaustiveDependencies: only on mount useEffect(() => { updateToken() }, []) @@ -162,30 +165,64 @@ const SectionUniversalToken = memo(() => { Universal token

- - When enabled, this token allows agents to self-register without prior system creation. Expires after one hour - or on hub restart. - + When enabled, this token allows agents to self-register without prior system creation.

-
+
{!isLoading && ( - <> - { - updateToken(checked ? 1 : 0) - }} - /> - - {token} - - - +
+
+ { + // Keep current permanence preference when enabling/disabling + updateToken(checked ? 1 : 0, isPermanent ? 1 : 0) + }} + /> +
+ + {token} + +
+ +
+ + {checked && ( +
+
+ Persistence +
+ updateToken(1, value === "permanent" ? 1 : 0)} + className="mt-2" + > + + + Ephemeral + + + Permanent + + + +

+ Expires after one hour or on hub restart. +

+ + +

+ Saved in the database and does not expire until you disable it. +

+ + +
+ )} +
)}