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