mirror of
https://github.com/fankes/beszel.git
synced 2025-10-19 09:49:28 +08:00
rename /src to /internal (sorry i'll fix the prs)
This commit is contained in:
henrygd
290
internal/hub/config/config.go
Normal file
290
internal/hub/config/config.go
Normal file
@@ -0,0 +1,290 @@
|
||||
// Package config provides functions for syncing systems with the config.yml file
|
||||
package config
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/henrygd/beszel/src/entities/system"
|
||||
|
||||
"github.com/pocketbase/dbx"
|
||||
"github.com/pocketbase/pocketbase/core"
|
||||
"github.com/spf13/cast"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
type config struct {
|
||||
Systems []systemConfig `yaml:"systems"`
|
||||
}
|
||||
|
||||
type systemConfig struct {
|
||||
Name string `yaml:"name"`
|
||||
Host string `yaml:"host"`
|
||||
Port uint16 `yaml:"port,omitempty"`
|
||||
Token string `yaml:"token,omitempty"`
|
||||
Users []string `yaml:"users"`
|
||||
}
|
||||
|
||||
// Syncs systems with the config.yml file
|
||||
func SyncSystems(e *core.ServeEvent) error {
|
||||
h := e.App
|
||||
configPath := filepath.Join(h.DataDir(), "config.yml")
|
||||
configData, err := os.ReadFile(configPath)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
var config config
|
||||
err = yaml.Unmarshal(configData, &config)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to parse config.yml: %v", err)
|
||||
}
|
||||
|
||||
if len(config.Systems) == 0 {
|
||||
log.Println("No systems defined in config.yml.")
|
||||
return nil
|
||||
}
|
||||
|
||||
var firstUser *core.Record
|
||||
|
||||
// Create a map of email to user ID
|
||||
userEmailToID := make(map[string]string)
|
||||
users, err := h.FindAllRecords("users", dbx.NewExp("id != ''"))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(users) > 0 {
|
||||
firstUser = users[0]
|
||||
for _, user := range users {
|
||||
userEmailToID[user.GetString("email")] = user.Id
|
||||
}
|
||||
}
|
||||
|
||||
// add default settings for systems if not defined in config
|
||||
for i := range config.Systems {
|
||||
system := &config.Systems[i]
|
||||
if system.Port == 0 {
|
||||
system.Port = 45876
|
||||
}
|
||||
if len(users) > 0 && len(system.Users) == 0 {
|
||||
// default to first user if none are defined
|
||||
system.Users = []string{firstUser.Id}
|
||||
} else {
|
||||
// Convert email addresses to user IDs
|
||||
userIDs := make([]string, 0, len(system.Users))
|
||||
for _, email := range system.Users {
|
||||
if id, ok := userEmailToID[email]; ok {
|
||||
userIDs = append(userIDs, id)
|
||||
} else {
|
||||
log.Printf("User %s not found", email)
|
||||
}
|
||||
}
|
||||
system.Users = userIDs
|
||||
}
|
||||
}
|
||||
|
||||
// Get existing systems
|
||||
existingSystems, err := h.FindAllRecords("systems", dbx.NewExp("id != ''"))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Create a map of existing systems
|
||||
existingSystemsMap := make(map[string]*core.Record)
|
||||
for _, system := range existingSystems {
|
||||
key := system.GetString("name") + system.GetString("host") + system.GetString("port")
|
||||
existingSystemsMap[key] = system
|
||||
}
|
||||
|
||||
// Process systems from config
|
||||
for _, sysConfig := range config.Systems {
|
||||
key := sysConfig.Name + sysConfig.Host + cast.ToString(sysConfig.Port)
|
||||
if existingSystem, ok := existingSystemsMap[key]; ok {
|
||||
// Update existing system
|
||||
existingSystem.Set("name", sysConfig.Name)
|
||||
existingSystem.Set("users", sysConfig.Users)
|
||||
existingSystem.Set("port", sysConfig.Port)
|
||||
if err := h.Save(existingSystem); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Only update token if one is specified in config, otherwise preserve existing token
|
||||
if sysConfig.Token != "" {
|
||||
if err := updateFingerprintToken(h, existingSystem.Id, sysConfig.Token); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
delete(existingSystemsMap, key)
|
||||
} else {
|
||||
// Create new system
|
||||
systemsCollection, err := h.FindCollectionByNameOrId("systems")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to find systems collection: %v", err)
|
||||
}
|
||||
newSystem := core.NewRecord(systemsCollection)
|
||||
newSystem.Set("name", sysConfig.Name)
|
||||
newSystem.Set("host", sysConfig.Host)
|
||||
newSystem.Set("port", sysConfig.Port)
|
||||
newSystem.Set("users", sysConfig.Users)
|
||||
newSystem.Set("info", system.Info{})
|
||||
newSystem.Set("status", "pending")
|
||||
if err := h.Save(newSystem); err != nil {
|
||||
return fmt.Errorf("failed to create new system: %v", err)
|
||||
}
|
||||
|
||||
// For new systems, generate token if not provided
|
||||
token := sysConfig.Token
|
||||
if token == "" {
|
||||
token = uuid.New().String()
|
||||
}
|
||||
|
||||
// Create fingerprint record for new system
|
||||
if err := createFingerprintRecord(h, newSystem.Id, token); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Delete systems not in config (and their fingerprint records will cascade delete)
|
||||
for _, system := range existingSystemsMap {
|
||||
if err := h.Delete(system); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
log.Println("Systems synced with config.yml")
|
||||
return nil
|
||||
}
|
||||
|
||||
// Generates content for the config.yml file as a YAML string
|
||||
func generateYAML(h core.App) (string, error) {
|
||||
// Fetch all systems from the database
|
||||
systems, err := h.FindRecordsByFilter("systems", "id != ''", "name", -1, 0)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Create a Config struct to hold the data
|
||||
config := config{
|
||||
Systems: make([]systemConfig, 0, len(systems)),
|
||||
}
|
||||
|
||||
// Fetch all users at once
|
||||
allUserIDs := make([]string, 0)
|
||||
for _, system := range systems {
|
||||
allUserIDs = append(allUserIDs, system.GetStringSlice("users")...)
|
||||
}
|
||||
userEmailMap, err := getUserEmailMap(h, allUserIDs)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Fetch all fingerprint records to get tokens
|
||||
type fingerprintData struct {
|
||||
ID string `db:"id"`
|
||||
System string `db:"system"`
|
||||
Token string `db:"token"`
|
||||
}
|
||||
var fingerprints []fingerprintData
|
||||
err = h.DB().NewQuery("SELECT id, system, token FROM fingerprints").All(&fingerprints)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Create a map of system ID to token
|
||||
systemTokenMap := make(map[string]string)
|
||||
for _, fingerprint := range fingerprints {
|
||||
systemTokenMap[fingerprint.System] = fingerprint.Token
|
||||
}
|
||||
|
||||
// Populate the Config struct with system data
|
||||
for _, system := range systems {
|
||||
userIDs := system.GetStringSlice("users")
|
||||
userEmails := make([]string, 0, len(userIDs))
|
||||
for _, userID := range userIDs {
|
||||
if email, ok := userEmailMap[userID]; ok {
|
||||
userEmails = append(userEmails, email)
|
||||
}
|
||||
}
|
||||
|
||||
sysConfig := systemConfig{
|
||||
Name: system.GetString("name"),
|
||||
Host: system.GetString("host"),
|
||||
Port: cast.ToUint16(system.Get("port")),
|
||||
Users: userEmails,
|
||||
Token: systemTokenMap[system.Id],
|
||||
}
|
||||
config.Systems = append(config.Systems, sysConfig)
|
||||
}
|
||||
|
||||
// Marshal the Config struct to YAML
|
||||
yamlData, err := yaml.Marshal(&config)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Add a header to the YAML
|
||||
yamlData = append([]byte("# Values for port, users, and token are optional.\n# Defaults are port 45876, the first created user, and a generated UUID token.\n\n"), yamlData...)
|
||||
|
||||
return string(yamlData), nil
|
||||
}
|
||||
|
||||
// New helper function to get a map of user IDs to emails
|
||||
func getUserEmailMap(h core.App, userIDs []string) (map[string]string, error) {
|
||||
users, err := h.FindRecordsByIds("users", userIDs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
userEmailMap := make(map[string]string, len(users))
|
||||
for _, user := range users {
|
||||
userEmailMap[user.Id] = user.GetString("email")
|
||||
}
|
||||
|
||||
return userEmailMap, nil
|
||||
}
|
||||
|
||||
// Helper function to update or create fingerprint token for an existing system
|
||||
func updateFingerprintToken(app core.App, systemID, token string) error {
|
||||
// Try to find existing fingerprint record
|
||||
fingerprint, err := app.FindFirstRecordByFilter("fingerprints", "system = {:system}", dbx.Params{"system": systemID})
|
||||
if err != nil {
|
||||
// If no fingerprint record exists, create one
|
||||
return createFingerprintRecord(app, systemID, token)
|
||||
}
|
||||
|
||||
// Update existing fingerprint record with new token (keep existing fingerprint)
|
||||
fingerprint.Set("token", token)
|
||||
return app.Save(fingerprint)
|
||||
}
|
||||
|
||||
// Helper function to create a new fingerprint record for a system
|
||||
func createFingerprintRecord(app core.App, systemID, token string) error {
|
||||
fingerprintsCollection, err := app.FindCollectionByNameOrId("fingerprints")
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to find fingerprints collection: %v", err)
|
||||
}
|
||||
|
||||
newFingerprint := core.NewRecord(fingerprintsCollection)
|
||||
newFingerprint.Set("system", systemID)
|
||||
newFingerprint.Set("token", token)
|
||||
newFingerprint.Set("fingerprint", "") // Empty fingerprint, will be set on first connection
|
||||
|
||||
return app.Save(newFingerprint)
|
||||
}
|
||||
|
||||
// Returns the current config.yml file as a JSON object
|
||||
func GetYamlConfig(e *core.RequestEvent) error {
|
||||
if e.Auth.GetString("role") != "admin" {
|
||||
return e.ForbiddenError("Requires admin role", nil)
|
||||
}
|
||||
configContent, err := generateYAML(e.App)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return e.JSON(200, map[string]string{"config": configContent})
|
||||
}
|
||||
247
internal/hub/config/config_test.go
Normal file
247
internal/hub/config/config_test.go
Normal file
@@ -0,0 +1,247 @@
|
||||
//go:build testing
|
||||
// +build testing
|
||||
|
||||
package config_test
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/henrygd/beszel/src/tests"
|
||||
|
||||
"github.com/henrygd/beszel/src/hub/config"
|
||||
|
||||
"github.com/pocketbase/pocketbase/core"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
// Config struct for testing (copied from config package since it's not exported)
|
||||
type testConfig struct {
|
||||
Systems []testSystemConfig `yaml:"systems"`
|
||||
}
|
||||
|
||||
type testSystemConfig struct {
|
||||
Name string `yaml:"name"`
|
||||
Host string `yaml:"host"`
|
||||
Port uint16 `yaml:"port,omitempty"`
|
||||
Users []string `yaml:"users"`
|
||||
Token string `yaml:"token,omitempty"`
|
||||
}
|
||||
|
||||
// Helper function to create a test system for config tests
|
||||
// func createConfigTestSystem(app core.App, name, host string, port uint16, userIDs []string) (*core.Record, error) {
|
||||
// systemCollection, err := app.FindCollectionByNameOrId("systems")
|
||||
// if err != nil {
|
||||
// return nil, err
|
||||
// }
|
||||
|
||||
// system := core.NewRecord(systemCollection)
|
||||
// system.Set("name", name)
|
||||
// system.Set("host", host)
|
||||
// system.Set("port", port)
|
||||
// system.Set("users", userIDs)
|
||||
// system.Set("status", "pending")
|
||||
|
||||
// return system, app.Save(system)
|
||||
// }
|
||||
|
||||
// Helper function to create a fingerprint record
|
||||
func createConfigTestFingerprint(app core.App, systemID, token, fingerprint string) (*core.Record, error) {
|
||||
fingerprintCollection, err := app.FindCollectionByNameOrId("fingerprints")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
fp := core.NewRecord(fingerprintCollection)
|
||||
fp.Set("system", systemID)
|
||||
fp.Set("token", token)
|
||||
fp.Set("fingerprint", fingerprint)
|
||||
|
||||
return fp, app.Save(fp)
|
||||
}
|
||||
|
||||
// TestConfigSyncWithTokens tests the config.SyncSystems function with various token scenarios
|
||||
func TestConfigSyncWithTokens(t *testing.T) {
|
||||
testHub, err := tests.NewTestHub()
|
||||
require.NoError(t, err)
|
||||
defer testHub.Cleanup()
|
||||
|
||||
// Create test user
|
||||
user, err := tests.CreateUser(testHub.App, "admin@example.com", "testtesttest")
|
||||
require.NoError(t, err)
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
setupFunc func() (string, *core.Record, *core.Record) // Returns: existing token, system record, fingerprint record
|
||||
configYAML string
|
||||
expectToken string // Expected token after sync
|
||||
description string
|
||||
}{
|
||||
{
|
||||
name: "new system with token in config",
|
||||
setupFunc: func() (string, *core.Record, *core.Record) {
|
||||
return "", nil, nil // No existing system
|
||||
},
|
||||
configYAML: `systems:
|
||||
- name: "new-server"
|
||||
host: "new.example.com"
|
||||
port: 45876
|
||||
users:
|
||||
- "admin@example.com"
|
||||
token: "explicit-token-123"`,
|
||||
expectToken: "explicit-token-123",
|
||||
description: "New system should use token from config",
|
||||
},
|
||||
{
|
||||
name: "existing system without token in config (preserve existing)",
|
||||
setupFunc: func() (string, *core.Record, *core.Record) {
|
||||
// Create existing system and fingerprint
|
||||
system, err := tests.CreateRecord(testHub.App, "systems", map[string]any{
|
||||
"name": "preserve-server",
|
||||
"host": "preserve.example.com",
|
||||
"port": 45876,
|
||||
"users": []string{user.Id},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
fingerprint, err := createConfigTestFingerprint(testHub.App, system.Id, "preserve-token-999", "preserve-fingerprint")
|
||||
require.NoError(t, err)
|
||||
|
||||
return "preserve-token-999", system, fingerprint
|
||||
},
|
||||
configYAML: `systems:
|
||||
- name: "preserve-server"
|
||||
host: "preserve.example.com"
|
||||
port: 45876
|
||||
users:
|
||||
- "admin@example.com"`,
|
||||
expectToken: "preserve-token-999",
|
||||
description: "Existing system should preserve original token when config doesn't specify one",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
// Setup test data
|
||||
_, existingSystem, existingFingerprint := tc.setupFunc()
|
||||
|
||||
// Write config file
|
||||
configPath := filepath.Join(testHub.DataDir(), "config.yml")
|
||||
err := os.WriteFile(configPath, []byte(tc.configYAML), 0644)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create serve event and sync
|
||||
event := &core.ServeEvent{App: testHub.App}
|
||||
err = config.SyncSystems(event)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Parse the config to get the system name for verification
|
||||
var configData testConfig
|
||||
err = yaml.Unmarshal([]byte(tc.configYAML), &configData)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, configData.Systems, 1)
|
||||
systemName := configData.Systems[0].Name
|
||||
|
||||
// Find the system after sync
|
||||
systems, err := testHub.FindRecordsByFilter("systems", "name = {:name}", "", -1, 0, map[string]any{"name": systemName})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, systems, 1)
|
||||
system := systems[0]
|
||||
|
||||
// Find the fingerprint record
|
||||
fingerprints, err := testHub.FindRecordsByFilter("fingerprints", "system = {:system}", "", -1, 0, map[string]any{"system": system.Id})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, fingerprints, 1)
|
||||
fingerprint := fingerprints[0]
|
||||
|
||||
// Verify token
|
||||
actualToken := fingerprint.GetString("token")
|
||||
if tc.expectToken == "" {
|
||||
// For generated tokens, just verify it's not empty and is a valid UUID format
|
||||
assert.NotEmpty(t, actualToken, tc.description)
|
||||
assert.Len(t, actualToken, 36, "Generated token should be UUID format") // UUID length
|
||||
} else {
|
||||
assert.Equal(t, tc.expectToken, actualToken, tc.description)
|
||||
}
|
||||
|
||||
// For existing systems, verify fingerprint is preserved
|
||||
if existingFingerprint != nil {
|
||||
actualFingerprint := fingerprint.GetString("fingerprint")
|
||||
expectedFingerprint := existingFingerprint.GetString("fingerprint")
|
||||
assert.Equal(t, expectedFingerprint, actualFingerprint, "Fingerprint should be preserved")
|
||||
}
|
||||
|
||||
// Cleanup for next test
|
||||
if existingSystem != nil {
|
||||
testHub.Delete(existingSystem)
|
||||
}
|
||||
if existingFingerprint != nil {
|
||||
testHub.Delete(existingFingerprint)
|
||||
}
|
||||
// Clean up the new records
|
||||
testHub.Delete(system)
|
||||
testHub.Delete(fingerprint)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestConfigMigrationScenario tests the specific migration scenario mentioned in the discussion
|
||||
func TestConfigMigrationScenario(t *testing.T) {
|
||||
testHub, err := tests.NewTestHub(t.TempDir())
|
||||
require.NoError(t, err)
|
||||
defer testHub.Cleanup()
|
||||
|
||||
// Create test user
|
||||
user, err := tests.CreateUser(testHub.App, "admin@example.com", "testtesttest")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Simulate migration scenario: system exists with token from migration
|
||||
existingSystem, err := tests.CreateRecord(testHub.App, "systems", map[string]any{
|
||||
"name": "migrated-server",
|
||||
"host": "migrated.example.com",
|
||||
"port": 45876,
|
||||
"users": []string{user.Id},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
migrationToken := "migration-generated-token-123"
|
||||
existingFingerprint, err := createConfigTestFingerprint(testHub.App, existingSystem.Id, migrationToken, "existing-fingerprint-from-agent")
|
||||
require.NoError(t, err)
|
||||
|
||||
// User exports config BEFORE this update (so no token field in YAML)
|
||||
oldConfigYAML := `systems:
|
||||
- name: "migrated-server"
|
||||
host: "migrated.example.com"
|
||||
port: 45876
|
||||
users:
|
||||
- "admin@example.com"`
|
||||
|
||||
// Write old config file and import
|
||||
configPath := filepath.Join(testHub.DataDir(), "config.yml")
|
||||
err = os.WriteFile(configPath, []byte(oldConfigYAML), 0644)
|
||||
require.NoError(t, err)
|
||||
|
||||
event := &core.ServeEvent{App: testHub.App}
|
||||
err = config.SyncSystems(event)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify the original token is preserved
|
||||
updatedFingerprint, err := testHub.FindRecordById("fingerprints", existingFingerprint.Id)
|
||||
require.NoError(t, err)
|
||||
|
||||
actualToken := updatedFingerprint.GetString("token")
|
||||
assert.Equal(t, migrationToken, actualToken, "Migration token should be preserved when config doesn't specify a token")
|
||||
|
||||
// Verify fingerprint is also preserved
|
||||
actualFingerprint := updatedFingerprint.GetString("fingerprint")
|
||||
assert.Equal(t, "existing-fingerprint-from-agent", actualFingerprint, "Existing fingerprint should be preserved")
|
||||
|
||||
// Verify system still exists and is updated correctly
|
||||
updatedSystem, err := testHub.FindRecordById("systems", existingSystem.Id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "migrated-server", updatedSystem.GetString("name"))
|
||||
assert.Equal(t, "migrated.example.com", updatedSystem.GetString("host"))
|
||||
}
|
||||
Reference in New Issue
Block a user