Fimeg/RedFlag
Watch
1
0
Fork
You've already forked RedFlag
 0
Code Releases 1 Activity

source: advance the checked projection

Browse source 
Source-Sha: 194a59adfd6c901635da7a9c18253bb2210e2426
Policy-Sha: b9b0be0ee74e55f561a803b8aef6be3c2134f5a83b46816a069a372f2eb04f4d
Tree-Digest: 4697f2e281c1c6ecafff17cf6eeb391a0a1ac0fa6e5fc429f5b473f4579c252c
This commit is contained in:
Source publisher 2026-09-21 14:39:08 +00:00
commit ffdccba10f
42 changed files with 2516 additions and 554 deletions
Download patch file Download diff file Expand all files Collapse all files

237
server/internal/api/handlers/fleet_join.go
View file

@ -7,10 +7,14 @@
package handlers
import (
"database/sql"
"encoding/json"
"errors"
"log"
"net/http"
"time"
"github.com/Fimeg/RedFlag/server/internal/api/middleware"
"github.com/Fimeg/RedFlag/server/internal/config"
"github.com/Fimeg/RedFlag/server/internal/database/queries"
"github.com/Fimeg/RedFlag/server/internal/models"
@ -20,22 +24,26 @@ import (
"github.com/jmoiron/sqlx"
)
const fleetJoinRequestMetadataKey = "redflag_fleet_join_request_id"
// FleetJoinHandler handles standalone-to-fleet migration (SEC-025).
type FleetJoinHandler struct {
db *sqlx.DB
tokenQueries *queries.RegistrationTokenQueries
agentQueries *queries.AgentQueries
signingPublicKey string
signingPublicKey string
signingKeyID string
config *config.Config
}
// NewFleetJoinHandler creates a FleetJoinHandler.
func NewFleetJoinHandler(db *sqlx.DB, tokenQueries *queries.RegistrationTokenQueries, agentQueries *queries.AgentQueries, signingPublicKey string, cfg *config.Config) *FleetJoinHandler {
func NewFleetJoinHandler(db *sqlx.DB, tokenQueries *queries.RegistrationTokenQueries, agentQueries *queries.AgentQueries, signingPublicKey, signingKeyID string, cfg *config.Config) *FleetJoinHandler {
return &FleetJoinHandler{
db: db,
tokenQueries: tokenQueries,
agentQueries: agentQueries,
signingPublicKey: signingPublicKey,
signingKeyID: signingKeyID,
config: cfg,
}
}
@ -44,27 +52,35 @@ func NewFleetJoinHandler(db *sqlx.DB, tokenQueries *queries.RegistrationTokenQue
// seed never crosses this channel — the server holds it encrypted from token
// creation and only the 6-digit code travels here.
type FleetJoinRequest struct {
RegistrationToken string `json:"registration_token" binding:"required"`
TOTPCode string `json:"totp_code" binding:"required"`
RequestID uuid.UUID `json:"request_id" binding:"required"`
AgentID uuid.UUID `json:"agent_id" binding:"required"`
RegistrationToken string `json:"registration_token" binding:"required"`
TOTPCode string `json:"totp_code" binding:"required"`
// Standard agent registration fields (same as RegisterAgent).
Hostname string `json:"hostname" binding:"required"`
OSType string `json:"os_type" binding:"required"`
OSVersion string `json:"os_version"`
OSArchitecture string `json:"os_architecture"`
AgentVersion string `json:"agent_version"`
MachineID string `json:"machine_id"`
PublicKeyFingerprint string `json:"public_key_fingerprint"`
AvailableScanners []string `json:"available_scanners"`
Hostname string `json:"hostname" binding:"required"`
OSType string `json:"os_type" binding:"required"`
OSVersion string `json:"os_version"`
OSArchitecture string `json:"os_architecture"`
AgentVersion string `json:"agent_version"`
MachineID string `json:"machine_id"`
PublicKeyFingerprint string `json:"public_key_fingerprint"`
AvailableScanners []string `json:"available_scanners"`
Metadata map[string]interface{} `json:"metadata"`
DeviceType string `json:"device_type"`
DeviceModel string `json:"device_model"`
OSDistro string `json:"os_distro"`
}
// FleetJoinResponse is the server's response on successful fleet join.
type FleetJoinResponse struct {
AgentID uuid.UUID `json:"agent_id"`
RefreshToken string `json:"refresh_token"`
SigningPublicKey string `json:"signing_public_key"`
ServerURL string `json:"server_url"`
AgentID uuid.UUID `json:"agent_id"`
Token string `json:"token"`
RefreshToken string `json:"refresh_token"`
SigningPublicKey string `json:"signing_public_key"`
SigningKeyID string `json:"signing_key_id"`
ServerURL string `json:"server_url"`
Config map[string]interface{} `json:"config"`
}
// JoinFleet handles POST /api/v1/fleet-join. Validates the registration token
@ -75,9 +91,13 @@ func (h *FleetJoinHandler) JoinFleet(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
return
}
if req.RequestID == uuid.Nil || req.AgentID == uuid.Nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "request_id and agent_id must be UUIDs"})
return
}
// Step 1: Validate registration token.
tokenInfo, err := h.tokenQueries.ValidateRegistrationToken(req.RegistrationToken)
tokenInfo, err := h.tokenQueries.GetFleetJoinToken(req.RegistrationToken)
if err != nil || tokenInfo == nil {
log.Printf("[SECURITY] [server] [fleet-join] invalid_token error=%v", err)
c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid or expired registration token"})
@ -90,6 +110,12 @@ func (h *FleetJoinHandler) JoinFleet(c *gin.Context) {
c.JSON(http.StatusBadRequest, gin.H{"error": "this token does not require 2FA — use the standard registration endpoint"})
return
}
var tokenMetadata map[string]interface{}
if err := json.Unmarshal(tokenInfo.Metadata, &tokenMetadata); err != nil || tokenMetadata["fleet_join"] != true {
log.Printf("[SECURITY] [server] [fleet-join] wrong_token_class token_id=%s", tokenInfo.ID)
c.JSON(http.StatusBadRequest, gin.H{"error": "this token is not a fleet join key"})
return
}
// Step 3: Validate the TOTP code against the server-held seed. The seed
// was stored encrypted at token creation and never crosses this channel;
@ -105,11 +131,49 @@ func (h *FleetJoinHandler) JoinFleet(c *gin.Context) {
c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid TOTP code"})
return
}
recovering, err := h.tokenQueries.FleetJoinTokenUsedBy(tokenInfo.ID, req.AgentID)
if err != nil {
log.Printf("[ERROR] [server] [fleet-join] token_usage_lookup_failed request_id=%s error=%q", req.RequestID, err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "fleet join failed"})
return
}
seatAvailable := tokenInfo.Status == "active" && tokenInfo.SeatsUsed < tokenInfo.MaxSeats
if !seatAvailable && !recovering {
log.Printf("[SECURITY] [server] [fleet-join] spent_token_wrong_agent request_id=%s agent_id=%s", req.RequestID, req.AgentID)
c.JSON(http.StatusUnauthorized, gin.H{"error": "fleet join token belongs to another Agent"})
return
}
existingByID, idErr := h.agentQueries.GetAgentByID(req.AgentID)
if recovering {
if idErr != nil || existingByID == nil {
log.Printf("[ERROR] [server] [fleet-join] recovery_agent_missing request_id=%s agent_id=%s error=%v", req.RequestID, req.AgentID, idErr)
c.JSON(http.StatusConflict, gin.H{"error": "fleet join recovery cannot find the registered Agent"})
return
}
if req.MachineID != "" && (existingByID.MachineID == nil || *existingByID.MachineID != req.MachineID) {
c.JSON(http.StatusConflict, gin.H{"error": "fleet join recovery machine ID mismatch"})
return
}
storedRequestID, _ := existingByID.Metadata[fleetJoinRequestMetadataKey].(string)
if storedRequestID == "" || storedRequestID != req.RequestID.String() {
log.Printf("[SECURITY] [server] [fleet-join] recovery_request_mismatch request_id=%s agent_id=%s", req.RequestID, req.AgentID)
c.JSON(http.StatusConflict, gin.H{"error": "fleet join recovery request does not match the accepted transaction"})
return
}
} else if idErr == nil {
c.JSON(http.StatusConflict, gin.H{"error": "Agent ID is already registered"})
return
} else if !errors.Is(idErr, sql.ErrNoRows) {
log.Printf("[ERROR] [server] [fleet-join] agent_lookup_failed request_id=%s error=%q", req.RequestID, idErr)
c.JSON(http.StatusInternalServerError, gin.H{"error": "fleet join failed"})
return
}
// Step 5: Check machine ID isn't already registered (same as RegisterAgent).
if req.MachineID != "" {
existing, err := h.agentQueries.GetAgentByMachineID(req.MachineID)
if err == nil && existing != nil && existing.ID.String() != "" {
if err == nil && existing != nil && existing.ID != req.AgentID {
c.JSON(http.StatusConflict, gin.H{
"error": "machine ID already registered to another agent",
"existing_agent_id": existing.ID.String(),
@ -139,9 +203,21 @@ func (h *FleetJoinHandler) JoinFleet(c *gin.Context) {
if req.PublicKeyFingerprint != "" {
pubKeyFP = &req.PublicKeyFingerprint
}
deviceType := req.DeviceType
if !models.ValidDeviceType(deviceType) {
deviceType = "server"
}
var deviceModel *string
if req.DeviceModel != "" {
deviceModel = &req.DeviceModel
}
var osDistro *string
if req.OSDistro != "" {
osDistro = &req.OSDistro
}
agent := &models.Agent{
ID: uuid.Must(uuid.NewV4()),
ID: req.AgentID,
Hostname: req.Hostname,
OSType: req.OSType,
OSVersion: req.OSVersion,
@ -150,6 +226,9 @@ func (h *FleetJoinHandler) JoinFleet(c *gin.Context) {
CurrentVersion: req.AgentVersion,
MachineID: machineID,
PublicKeyFingerprint: pubKeyFP,
DeviceType: deviceType,
DeviceModel: deviceModel,
OSDistro: osDistro,
LastSeen: time.Now().UTC(),
Status: "online",
Metadata: models.JSONB{},
@ -160,30 +239,75 @@ func (h *FleetJoinHandler) JoinFleet(c *gin.Context) {
agent.Metadata[k] = v
}
}
createQuery := `
INSERT INTO agents (
id, hostname, os_type, os_version, os_architecture,
agent_version, current_version, machine_id, public_key_fingerprint,
last_seen, status, metadata
) VALUES (
:id, :hostname, :os_type, :os_version, :os_architecture,
:agent_version, :current_version, :machine_id, :public_key_fingerprint,
:last_seen, :status, :metadata
)`
if _, err := tx.NamedExec(createQuery, agent); err != nil {
log.Printf("[ERROR] [server] [fleet-join] create_agent_failed error=%q", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to register agent"})
return
if recovering {
for k, v := range existingByID.Metadata {
if _, exists := agent.Metadata[k]; !exists {
agent.Metadata[k] = v
}
}
if agent.MachineID == nil {
agent.MachineID = existingByID.MachineID
}
if agent.PublicKeyFingerprint == nil {
agent.PublicKeyFingerprint = existingByID.PublicKeyFingerprint
}
if req.DeviceType == "" {
agent.DeviceType = existingByID.DeviceType
}
if agent.DeviceModel == nil {
agent.DeviceModel = existingByID.DeviceModel
}
if agent.OSDistro == nil {
agent.OSDistro = existingByID.OSDistro
}
}
agent.Metadata[fleetJoinRequestMetadataKey] = req.RequestID.String()
// Mark token as used.
var tokenSuccess bool
tokenHash := queries.HashRegistrationToken(req.RegistrationToken)
if err := tx.QueryRow("SELECT mark_registration_token_used($1, $2)", tokenHash, agent.ID).Scan(&tokenSuccess); err != nil || !tokenSuccess {
log.Printf("[ERROR] [server] [fleet-join] mark_token_failed error=%v success=%v", err, tokenSuccess)
c.JSON(http.StatusBadRequest, gin.H{"error": "registration token could not be consumed"})
return
if recovering {
updateQuery := `
UPDATE agents SET
hostname = :hostname, os_type = :os_type, os_version = :os_version,
os_architecture = :os_architecture, agent_version = :agent_version,
current_version = :current_version, machine_id = :machine_id,
public_key_fingerprint = :public_key_fingerprint,
device_type = :device_type, device_model = :device_model,
os_distro = :os_distro, last_seen = :last_seen,
status = :status, metadata = :metadata
WHERE id = :id`
if _, err := tx.NamedExec(updateQuery, agent); err != nil {
log.Printf("[ERROR] [server] [fleet-join] recover_agent_failed request_id=%s error=%q", req.RequestID, err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to recover fleet join"})
return
}
if _, err := tx.Exec("UPDATE refresh_tokens SET revoked = TRUE WHERE agent_id = $1 AND NOT revoked", agent.ID); err != nil {
log.Printf("[ERROR] [server] [fleet-join] recover_revoke_tokens_failed request_id=%s error=%q", req.RequestID, err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to recover fleet join"})
return
}
} else {
createQuery := `
INSERT INTO agents (
id, hostname, os_type, os_version, os_architecture,
agent_version, current_version, machine_id, public_key_fingerprint,
device_type, device_model, os_distro, last_seen, status, metadata
) VALUES (
:id, :hostname, :os_type, :os_version, :os_architecture,
:agent_version, :current_version, :machine_id, :public_key_fingerprint,
:device_type, :device_model, :os_distro, :last_seen, :status, :metadata
)`
if _, err := tx.NamedExec(createQuery, agent); err != nil {
log.Printf("[ERROR] [server] [fleet-join] create_agent_failed request_id=%s error=%q", req.RequestID, err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to register agent"})
return
}
var tokenSuccess bool
tokenHash := queries.HashRegistrationToken(req.RegistrationToken)
if err := tx.QueryRow("SELECT mark_registration_token_used($1, $2)", tokenHash, agent.ID).Scan(&tokenSuccess); err != nil || !tokenSuccess {
log.Printf("[ERROR] [server] [fleet-join] mark_token_failed request_id=%s error=%v success=%v", req.RequestID, err, tokenSuccess)
c.JSON(http.StatusBadRequest, gin.H{"error": "registration token could not be consumed"})
return
}
}
// Generate refresh token.
@ -202,36 +326,53 @@ func (h *FleetJoinHandler) JoinFleet(c *gin.Context) {
return
}
// Create subsystems for available scanners.
if len(req.AvailableScanners) > 0 {
for _, scanner := range req.AvailableScanners {
// Initial join creates the same observed scanner rows as registration.
if !recovering {
scanners := append([]string{}, req.AvailableScanners...)
scanners = append(scanners, "storage", "system")
for _, scanner := range scanners {
interval := 60
if scanner == "storage" || scanner == "system" {
interval = 5
}
sub := models.AgentSubsystem{
AgentID: agent.ID,
Subsystem: scanner,
Enabled: true,
AutoRun: true,
IntervalMinutes: 60,
IntervalMinutes: interval,
}
subQuery := `INSERT INTO agent_subsystems (agent_id, subsystem, enabled, auto_run, interval_minutes) VALUES (:agent_id, :subsystem, :enabled, :auto_run, :interval_minutes)`
subQuery := `INSERT INTO agent_subsystems (agent_id, subsystem, enabled, auto_run, interval_minutes) VALUES (:agent_id, :subsystem, :enabled, :auto_run, :interval_minutes) ON CONFLICT (agent_id, subsystem) DO NOTHING`
if _, err := tx.NamedExec(subQuery, sub); err != nil {
log.Printf("[WARNING] [server] [fleet-join] create_subsystem_failed scanner=%s error=%q", scanner, err)
}
}
}
accessToken, err := middleware.GenerateAgentToken(agent.ID)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate access token"})
return
}
if err := tx.Commit(); err != nil {
log.Printf("[ERROR] [server] [fleet-join] commit_failed error=%q", err)
c.JSON(http.StatusInternalServerError, gin.H{"error": "fleet join failed"})
return
}
log.Printf("[INFO] [server] [fleet-join] agent_registered agent_id=%s hostname=%s machine_id=%s",
agent.ID, agent.Hostname, req.MachineID)
log.Printf("[INFO] [server] [fleet-join] completed request_id=%s agent_id=%s hostname=%s recovered=%v",
req.RequestID, agent.ID, agent.Hostname, recovering)
c.JSON(http.StatusCreated, FleetJoinResponse{
AgentID: agent.ID,
Token: accessToken,
RefreshToken: refreshToken,
SigningPublicKey: h.signingPublicKey,
SigningKeyID: h.signingKeyID,
ServerURL: resolveServerURL(c, h.config, "fleet-join"),
Config: map[string]interface{}{
"check_in_interval": h.config.CheckInInterval,
},
})
}

18
server/internal/api/handlers/registration_tokens.go
View file

@ -5,6 +5,7 @@ import (
"log"
"net/http"
"strconv"
"strings"
"time"
"github.com/Fimeg/RedFlag/server/internal/config"
@ -44,10 +45,10 @@ func NewRegistrationTokenHandler(tokenQueries *queries.RegistrationTokenQueries,
// GenerateRegistrationToken creates a new registration token
func (h *RegistrationTokenHandler) GenerateRegistrationToken(c *gin.Context) {
var request struct {
Label string `json:"label" binding:"required"`
ExpiresIn string `json:"expires_in"` // e.g., "24h", "7d", "168h"
MaxSeats int `json:"max_seats"` // Number of agents that can use this token
Metadata map[string]interface{} `json:"metadata"`
Label string `json:"label" binding:"required"`
ExpiresIn string `json:"expires_in"` // e.g., "24h", "7d", "168h"
MaxSeats int `json:"max_seats"` // Number of agents that can use this token
Metadata map[string]interface{} `json:"metadata"`
// FleetJoin + TOTPSeed: SEC-025 fleet-join 2FA. When the operator
// creates a fleet-join token, they enter the host's TOTP seed
// (displayed on the standalone host). This is the one admin-
@ -110,18 +111,23 @@ func (h *RegistrationTokenHandler) GenerateRegistrationToken(c *gin.Context) {
}
metadata["server_url"] = resolveServerURL(c, h.config, "registration-tokens")
metadata["expires_in"] = expiresIn
metadata["fleet_join"] = request.FleetJoin
// Default max_seats to 1 if not provided or invalid
maxSeats := request.MaxSeats
if maxSeats < 1 {
maxSeats = 1
}
if request.FleetJoin {
maxSeats = 1
}
// Store token in database. For fleet-join tokens (SEC-025), the operator
// provides the TOTP seed from the standalone host; it must be valid
// base32 or neither an authenticator nor join validation can use it.
var totpSeed string
if request.FleetJoin {
request.TOTPSeed = strings.ToUpper(strings.TrimSpace(request.TOTPSeed))
if request.TOTPSeed == "" {
c.JSON(http.StatusBadRequest, gin.H{"error": "fleet-join tokens require a TOTP seed"})
return
@ -145,6 +151,9 @@ func (h *RegistrationTokenHandler) GenerateRegistrationToken(c *gin.Context) {
// SEC-002: token travels in a header, never the URL — query strings land in
// shell history, process lists, and access logs.
installCommand := fmt.Sprintf("curl -sfL -H \"X-Registration-Token: %s\" \"%s/api/v1/install/linux\" | sudo bash", token, serverURL)
if request.FleetJoin {
installCommand = ""
}
response := gin.H{
"token": token,
@ -152,6 +161,7 @@ func (h *RegistrationTokenHandler) GenerateRegistrationToken(c *gin.Context) {
"expires_at": expiresAt,
"install_command": installCommand,
"metadata": metadata,
"fleet_join": request.FleetJoin,
}
c.JSON(http.StatusCreated, response)

93
server/internal/database/queries/registration_tokens.go
View file

@ -29,35 +29,35 @@ type RegistrationTokenQueries struct {
}
type RegistrationToken struct {
ID uuid.UUID `json:"id" db:"id"`
TokenHash string `json:"-" db:"token_hash"`
ID uuid.UUID `json:"id" db:"id"`
TokenHash string `json:"-" db:"token_hash"`
// TokenEncrypted is the AES-256-GCM ciphertext (nonce||ct) of the plaintext
// token. Never serialised; decrypted into Token for live tokens only.
TokenEncrypted []byte `json:"-" db:"token_encrypted"`
TokenEncrypted []byte `json:"-" db:"token_encrypted"`
// Token is the decrypted plaintext, populated only for live tokens so the
// operator UI can rebuild the install one-liner. Empty for spent/expired/
// revoked tokens and for legacy one-way-hashed tokens.
Token string `json:"token,omitempty" db:"-"`
Label *string `json:"label" db:"label"`
ExpiresAt time.Time `json:"expires_at" db:"expires_at"`
CreatedAt time.Time `json:"created_at" db:"created_at"`
UsedAt *time.Time `json:"used_at" db:"used_at"`
UsedByAgentID *uuid.UUID `json:"used_by_agent_id" db:"used_by_agent_id"`
Revoked bool `json:"revoked" db:"revoked"`
RevokedAt *time.Time `json:"revoked_at" db:"revoked_at"`
RevokedReason *string `json:"revoked_reason" db:"revoked_reason"`
Status string `json:"status" db:"status"`
CreatedBy string `json:"created_by" db:"created_by"`
Metadata json.RawMessage `json:"metadata" db:"metadata"`
MaxSeats int `json:"max_seats" db:"max_seats"`
SeatsUsed int `json:"seats_used" db:"seats_used"`
Token string `json:"token,omitempty" db:"-"`
Label *string `json:"label" db:"label"`
ExpiresAt time.Time `json:"expires_at" db:"expires_at"`
CreatedAt time.Time `json:"created_at" db:"created_at"`
UsedAt *time.Time `json:"used_at" db:"used_at"`
UsedByAgentID *uuid.UUID `json:"used_by_agent_id" db:"used_by_agent_id"`
Revoked bool `json:"revoked" db:"revoked"`
RevokedAt *time.Time `json:"revoked_at" db:"revoked_at"`
RevokedReason *string `json:"revoked_reason" db:"revoked_reason"`
Status string `json:"status" db:"status"`
CreatedBy string `json:"created_by" db:"created_by"`
Metadata json.RawMessage `json:"metadata" db:"metadata"`
MaxSeats int `json:"max_seats" db:"max_seats"`
SeatsUsed int `json:"seats_used" db:"seats_used"`
// TOTPSeedEncrypted is the TOTP seed for fleet-join 2FA (SEC-025), stored
// as AES-256-GCM ciphertext (nonce||ct) of the base32 seed. Nil for
// standard registration tokens; set only when the operator creates a
// fleet-join token that requires the standalone host to prove possession
// of the seed via a TOTP code. Never serialised; decrypted only at join
// validation. The join request carries the code, never the seed.
TOTPSeedEncrypted []byte `json:"-" db:"totp_seed_encrypted"`
TOTPSeedEncrypted []byte `json:"-" db:"totp_seed_encrypted"`
}
// isLive reports whether a token is still usable for enrollment — the only state
@ -81,17 +81,17 @@ func (q *RegistrationTokenQueries) revealPlaintext(t *RegistrationToken) {
}
type TokenRequest struct {
Label string `json:"label"`
ExpiresIn string `json:"expires_in"` // e.g., "24h", "7d"
MaxSeats int `json:"max_seats"` // Number of agents that can use this token (default: 1)
Metadata map[string]interface{} `json:"metadata"`
Label string `json:"label"`
ExpiresIn string `json:"expires_in"` // e.g., "24h", "7d"
MaxSeats int `json:"max_seats"` // Number of agents that can use this token (default: 1)
Metadata map[string]interface{} `json:"metadata"`
}
type TokenResponse struct {
Token string `json:"token"`
Label string `json:"label"`
ExpiresAt time.Time `json:"expires_at"`
InstallCommand string `json:"install_command"`
Token string `json:"token"`
Label string `json:"label"`
ExpiresAt time.Time `json:"expires_at"`
InstallCommand string `json:"install_command"`
}
// NewRegistrationTokenQueries builds the query handle. encKeyB64 is the base64
@ -205,6 +205,45 @@ func (q *RegistrationTokenQueries) ValidateRegistrationToken(token string) (*Reg
return &regToken, nil
}
// GetFleetJoinToken resolves the raw join secret even after its one seat was
// consumed. The fleet-join handler uses that spent state only to recover the
// same Agent ID after a response was lost; standard enrollment never calls it.
func (q *RegistrationTokenQueries) GetFleetJoinToken(token string) (*RegistrationToken, error) {
var regToken RegistrationToken
query := `
SELECT id, token_hash, label, expires_at, created_at, used_at, used_by_agent_id,
revoked, revoked_at, revoked_reason, status, created_by, metadata,
max_seats, seats_used, totp_seed_encrypted
FROM registration_tokens
WHERE token_hash = $1
AND status IN ('active', 'used')
AND NOT revoked
AND expires_at > NOW()
AND totp_seed_encrypted IS NOT NULL
`
if err := q.db.Get(&regToken, query, HashRegistrationToken(token)); err != nil {
if err == sql.ErrNoRows {
return nil, fmt.Errorf("invalid, expired, or revoked fleet join token")
}
return nil, fmt.Errorf("failed to resolve fleet join token: %w", err)
}
return &regToken, nil
}
func (q *RegistrationTokenQueries) FleetJoinTokenUsedBy(tokenID, agentID uuid.UUID) (bool, error) {
var used bool
err := q.db.Get(&used, `
SELECT EXISTS (
SELECT 1 FROM registration_token_usage
WHERE token_id = $1 AND agent_id = $2
)
`, tokenID, agentID)
if err != nil {
return false, fmt.Errorf("failed to resolve fleet join token use: %w", err)
}
return used, nil
}
// MarkTokenUsed marks a token as used by an agent.
// The caller passes the plaintext token; this function hashes it before calling
// the stored procedure (which now matches on token_hash).
@ -448,4 +487,4 @@ func (q *RegistrationTokenQueries) GetTokenUsageStats() (map[string]int, error)
}
return stats, nil
}
}

18
server/internal/taskrunner/taskrunner_test.go
View file

@ -26,12 +26,22 @@ func TestGo_RunsAllSubmittedTasks(t *testing.T) {
if got != n {
t.Fatalf("ran %d tasks, want %d", got, n)
}
s := r.Snapshot()
if s.Submitted != n {
if s := r.Snapshot(); s.Submitted != n {
t.Errorf("submitted=%d, want %d", s.Submitted, n)
}
if s.Completed != n {
t.Errorf("completed=%d, want %d", s.Completed, n)
// Completed is incremented after the task body returns, so the WaitGroup
// above cannot cover it: the last worker may still be accounting when Wait
// releases. Wait for the counter actually asserted, as the panic test does.
deadline := time.Now().Add(time.Second)
for time.Now().Before(deadline) {
if r.Snapshot().Completed >= n {
break
}
time.Sleep(2 * time.Millisecond)
}
if c := r.Snapshot().Completed; c != n {
t.Errorf("completed=%d, want %d", c, n)
}
}