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

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)