Enrollment

CREATE TABLE agent (
    id            uuid PRIMARY KEY DEFAULT gen_random_uuid(),
    workspace_id  uuid NOT NULL REFERENCES workspace(id) ON DELETE CASCADE,
    name          text NOT NULL,
    status        text NOT NULL DEFAULT 'enrolling'
                  CHECK (status IN ('enrolling','active','revoked')),
    key_hash      text,                       -- set at enrollment; sha-256 hex
    version       text NOT NULL DEFAULT '',
    last_seen_at  timestamptz,
    created_at    timestamptz NOT NULL DEFAULT now(),
    revoked_at    timestamptz
);

const tokenBytes = 32

// NewToken returns a fresh opaque token as a base64url string. The raw token is
// handed to the client once; only its HashToken value is ever persisted.
func NewToken() (string, error) {
	b := make([]byte, tokenBytes)
	if _, err := rand.Read(b); err != nil {
		return "", fmt.Errorf("auth: new token: %w", err)
	}
	return base64.RawURLEncoding.EncodeToString(b), nil
}

// HashToken returns the hex-encoded SHA-256 of a raw token. Storing the hash
// (never the raw token) means a database leak can't be replayed as a login.
func HashToken(raw string) string {
	sum := sha256.Sum256([]byte(raw))
	return hex.EncodeToString(sum[:])
}

// Enroll exchanges a one-time enrollment token for a permanent agent key. It is
// unauthenticated (the agent has no key yet). Returns the new agent id, key, and
// the canonical control-plane URL the agent should poll.
func Enroll(ctx context.Context, baseURL, token, version string) (agentID, agentKey, canonicalURL string, err error) {
	body, _ := json.Marshal(map[string]string{"token": token, "version": version})
	url := strings.TrimRight(baseURL, "/") + "/agent/v1/enroll"
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
	// ...
	req.Header.Set("Content-Type", "application/json")
	resp, err := (&http.Client{Timeout: 30 * time.Second}).Do(req)
	// ... non-200 → statusError (reads a bounded 512-byte prefix of the body)
	var out struct {
		AgentID  string `json:"agentId"`
		AgentKey string `json:"agentKey"`
		BaseURL  string `json:"baseUrl"`
	}
	json.NewDecoder(resp.Body).Decode(&out)
	return out.AgentID, out.AgentKey, out.BaseURL, nil
}

func runEnroll(args []string) error {
	token := args[0]
	baseURL := config.EnrollBaseURL()

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	agentID, agentKey, canonicalURL, err := client.Enroll(ctx, baseURL, token, agent.Version)
	if err != nil {
		return fmt.Errorf("enroll: %w", err)
	}
	// Poll the canonical URL the control plane reports, not necessarily the one
	// we dialed — so `setu run` reaches home even when enrolled via localhost/LAN.
	pollURL := baseURL
	if u := strings.TrimRight(strings.TrimSpace(canonicalURL), "/"); u != "" {
		pollURL = u
	}
	return config.Save(config.Config{
		BaseURL:  pollURL,
		AgentID:  agentID,
		AgentKey: agentKey,
		Version:  agent.Version,
	})
}

const enrollTokenTTL = 15 * time.Minute

// ... confirm the agent exists in this workspace (404 otherwise) before issuing.
token, _ := auth.NewToken()
expiresAt := time.Now().Add(enrollTokenTTL)
s.q.CreateEnrollmentToken(r.Context(), sqlc.CreateEnrollmentTokenParams{
	TokenHash: auth.HashToken(token), AgentID: id, WorkspaceID: ws, ExpiresAt: expiresAt,
})
writeJSON(w, http.StatusCreated, map[string]any{
	"token":           token,     // plaintext, returned once
	"expiresAt":       expiresAt,
	"installCommand":  fmt.Sprintf("setu enroll %s", token),
	"controlPlaneUrl": controlPlaneURL(), // lets the UI fill URL-bearing snippets (e.g. Docker)
})

func (s *Server) enroll(w http.ResponseWriter, r *http.Request) {
	var req struct {
		Token   string `json:"token"`
		Version string `json:"version"`
	}
	// ... decode, require non-empty token (else 400)
	row, err := s.q.ConsumeEnrollmentToken(r.Context(), auth.HashToken(req.Token))
	if err != nil {
		s.unauthenticated(w) // invalid, expired, or already-consumed → reveal nothing
		return
	}
	key, _ := auth.NewToken()
	keyHash := auth.HashToken(key)
	s.q.SetAgentKey(r.Context(), sqlc.SetAgentKeyParams{
		ID: row.AgentID, KeyHash: &keyHash, Version: req.Version, // → status='active'
	})
	// Return the canonical public URL so the agent polls the right host no matter
	// which URL it dialed to enroll.
	writeJSON(w, http.StatusOK, map[string]any{
		"agentId":  row.AgentID,
		"agentKey": key,       // plaintext, returned exactly once
		"baseUrl":  controlPlaneURL(),
	})
}

-- name: ConsumeEnrollmentToken :one
-- Atomic one-time-use gate: claims the token iff it is unconsumed and unexpired,
-- stamping consumed_at in the same statement so a concurrent/replayed enroll
-- finds no row (→ rejected).
UPDATE agent_enrollment_token
SET consumed_at = now()
WHERE token_hash = $1 AND consumed_at IS NULL AND expires_at > now()
RETURNING agent_id, workspace_id;

-- name: SetAgentKey :exec
-- Activate an agent at enrollment: store the sha-256 hash of its key + version.
UPDATE agent
SET key_hash = $2, version = $3, status = 'active'
WHERE id = $1;

func (s *Server) agentAuthMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		key := bearerToken(r)
		if key == "" {
			s.unauthenticated(w) // header check runs before any DB lookup
			return
		}
		hash := auth.HashToken(key)
		agent, err := s.q.GetAgentByKeyHash(r.Context(), &hash) // WHERE key_hash=$1 AND status='active'
		if err != nil {
			s.unauthenticated(w) // unknown or revoked key → 401
			return
		}
		// Liveness is advisory: never fail the request if the touch write errors.
		_ = s.q.TouchAgentSeen(r.Context(), sqlc.TouchAgentSeenParams{ID: agent.ID, Version: agent.Version})
		ctx := context.WithValue(r.Context(), ctxKeyAgentID, agent.ID)
		ctx = context.WithValue(ctx, ctxKeyWorkspaceID, agent.WorkspaceID)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}