add option to make universal token permanent (#1097, 1614)

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)

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)

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)
 }
 

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,
 		},
 		{

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
 	}
 ]`
 

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(() => {
 				<Trans>Universal token</Trans>
 			</h3>
 			<p className="text-sm text-muted-foreground leading-relaxed">
-				<Trans>
-					When enabled, this token allows agents to self-register without prior system creation. Expires after one hour
-					or on hub restart.
-				</Trans>
+				<Trans>When enabled, this token allows agents to self-register without prior system creation.</Trans>
 			</p>
-			<div className="min-h-16 overflow-auto max-w-full inline-flex items-center gap-5 mt-3 border py-2 ps-5 pe-4 rounded-md">
+			<div className="mt-3 border rounded-md px-4 py-3 max-w-full">
 				{!isLoading && (
-					<>
-						<Switch
-							defaultChecked={checked}
-							onCheckedChange={(checked) => {
-								updateToken(checked ? 1 : 0)
-							}}
-						/>
-						<span
-							className={cn(
-								"text-sm text-primary opacity-60 transition-opacity",
-								checked ? "opacity-100" : "select-none"
-							)}
-						>
-							{token}
-						</span>
-						<ActionsButtonUniversalToken token={token} checked={checked} />
-					</>
+					<div className="flex flex-col gap-3">
+						<div className="flex items-center gap-4 min-w-0">
+							<Switch
+								checked={checked}
+								onCheckedChange={(checked) => {
+									// Keep current permanence preference when enabling/disabling
+									updateToken(checked ? 1 : 0, isPermanent ? 1 : 0)
+								}}
+							/>
+							<div className="min-w-0 flex-1 overflow-auto">
+								<span
+									className={cn(
+										"text-sm text-primary opacity-60 transition-opacity",
+										checked ? "opacity-100" : "select-none"
+									)}
+								>
+									{token}
+								</span>
+							</div>
+							<ActionsButtonUniversalToken token={token} checked={checked} />
+						</div>
+
+						{checked && (
+							<div className="border-t pt-3">
+								<div className="text-sm font-medium">
+									<Trans>Persistence</Trans>
+								</div>
+								<Tabs
+									value={isPermanent ? "permanent" : "ephemeral"}
+									onValueChange={(value) => updateToken(1, value === "permanent" ? 1 : 0)}
+									className="mt-2"
+								>
+									<TabsList>
+										<TabsTrigger className="xs:min-w-40" value="ephemeral">
+											<Trans>Ephemeral</Trans>
+										</TabsTrigger>
+										<TabsTrigger className="xs:min-w-40" value="permanent">
+											<Trans>Permanent</Trans>
+										</TabsTrigger>
+									</TabsList>
+									<TabsContent value="ephemeral" className="mt-3">
+										<p className="text-sm text-muted-foreground leading-relaxed">
+											<Trans>Expires after one hour or on hub restart.</Trans>
+										</p>
+									</TabsContent>
+									<TabsContent value="permanent" className="mt-3">
+										<p className="text-sm text-muted-foreground leading-relaxed">
+											<Trans>Saved in the database and does not expire until you disable it.</Trans>
+										</p>
+									</TabsContent>
+								</Tabs>
+							</div>
+						)}
+					</div>
 				)}
 			</div>
 		</div>