Lots more work on auth process

cmd/cloudcore-server/database/cockroachdb/agent.go

@@ -3,31 +3,68 @@ package cockroachdb
 import (
 	"context"
 	"errors"
+	"fmt"
+	"github.com/clarkmcc/cloudcore/cmd/cloudcore-server/database/types"
 	"github.com/clarkmcc/cloudcore/internal/rpc"
+	"github.com/jmoiron/sqlx"
 	"go.uber.org/multierr"
 	"time"
 )
 
 var (
-	ErrAgentNotFound = errors.New("agent not found")
+	ErrAgentNotFound        = errors.New("agent not found")
+	ErrPreSharedKeyNotFound = errors.New("pre-shared key not found")
 )
 
-// UpdateMetadata updates the metadata for the host and agent with the given ID. If there is no agent ID, then
-// we create the agent and return the ID.
-func (d *Database) UpdateMetadata(ctx context.Context, metadata *rpc.SystemMetadata) (string, error) {
+func (d *Database) AuthenticateAgent(ctx context.Context, key string, metadata *rpc.SystemMetadata) (agentID string, err error) {
 	tx, err := d.db.BeginTxx(ctx, nil)
 	if err != nil {
 		return "", err
 	}
-	defer multierr.AppendFunc(&err, tx.Rollback)
+	defer handleRollback(&err, tx)
+
+	// Make sure the psk exists and is usable
+	var psk types.PreSharedKey
+	err = tx.GetContext(ctx, &psk, `
+		SELECT id, project_id, created_at, updated_at, name, key, status, uses_remaining, expiration FROM agent_psk 
+		WHERE key = $1
+		  	AND status = 'active'
+			AND uses_remaining > 0
+			AND expiration > NOW()
+		LIMIT 1
+	`, key)
+	if err != nil {
+		return "", fmt.Errorf("getting pre-shared key: %w", err)
+	}
+
+	// Decrement the uses remaining
+	_, err = tx.ExecContext(ctx, `
+		UPDATE agent_psk SET uses_remaining = uses_remaining - 1 WHERE id = $1;
+	`, psk.ID)
+	if err != nil {
+		return "", fmt.Errorf("updating pre-shared key uses remaining: %w", err)
+	}
+
+	// Check for any PSK groups
+	var groups []types.AgentGroup
+	err = tx.SelectContext(ctx, &groups, `
+		SELECT g.* FROM agent_group g
+		INNER JOIN agent_group_psk gp ON g.id = gp.agent_group_id AND gp.status = 'active'
+		INNER JOIN agent_psk p ON gp.agent_psk_id = p.id AND p.status = 'active'
+		WHERE p.id = $1 AND g.status = 'active';
+	`, psk.ID)
+	if err != nil {
+		return "", fmt.Errorf("finding agent groups: %w", err)
+	}
 
 	// First, we upsert the host on the identifier field
-	rows, err := d.db.NamedQueryContext(ctx, `
-		INSERT INTO hosts (identifier, hostname, host_id, public_ip_address, os_name, os_family, os_version, kernel_architecture, kernel_version, cpu_model, cpu_cores)
-		VALUES(:identifier, :hostname, :host_id, :public_ip_address, :os_name, :os_family, :os_version, :kernel_architecture, :kernel_version, :cpu_model, :cpu_cores)
+	rows, err := tx.NamedQuery(`
+		INSERT INTO host (project_id, identifier, hostname, host_id, public_ip_address, os_name, os_family, os_version, kernel_architecture, kernel_version, cpu_model, cpu_cores)
+		VALUES(:project_id, :identifier, :hostname, :host_id, :public_ip_address, :os_name, :os_family, :os_version, :kernel_architecture, :kernel_version, :cpu_model, :cpu_cores)
 		ON CONFLICT (identifier) DO UPDATE SET hostname = :hostname, host_id = :host_id, public_ip_address = :public_ip_address, os_name = :os_name, os_family = :os_family, os_version = :os_version, kernel_architecture = :kernel_architecture, kernel_version = :kernel_version, cpu_model = :cpu_model, cpu_cores = :cpu_cores
 		RETURNING id
 	`, map[string]any{
+		"project_id":          psk.ProjectID,
 		"identifier":          metadata.GetIdentifiers().GetHostIdentifier(),
 		"hostname":            metadata.GetIdentifiers().GetHostname(),
 		"host_id":             metadata.GetIdentifiers().GetHostId(),
@@ -48,23 +85,55 @@ func (d *Database) UpdateMetadata(ctx context.Context, metadata *rpc.SystemMetad
 		return "", err
 	}
 
-	// Next, we upsert the agent on the agent_id field
-	rows, err = d.db.NamedQueryContext(ctx, `
-		UPSERT INTO agents (id, host_id, online, last_heartbeat_timestamp)
-		VALUES (:id, :host_id, :online, :last_heartbeat_timestamp)
+	// Upsert the agent now
+	agentID = metadata.GetIdentifiers().GetAgentIdentifier()
+	if len(agentID) == 0 {
+		rows, err = tx.NamedQuery(`
+		INSERT INTO agent (project_id, host_id, online, last_heartbeat_timestamp)
+		VALUES (:project_id, :host_id, :online, :last_heartbeat_timestamp)
 		RETURNING id
 	`, map[string]any{
-		"id":                       metadata.GetIdentifiers().GetAgentIdentifier(),
-		"host_id":                  hostID,
-		"online":                   true,
-		"last_heartbeat_timestamp": time.Now(),
-	})
-	if err != nil {
-		return "", err
+			"project_id":               psk.ProjectID,
+			"host_id":                  hostID,
+			"online":                   true,
+			"last_heartbeat_timestamp": time.Now(),
+		})
+		if err != nil {
+			return "", err
+		}
+		agentID, err = getReturningID(rows)
+		if err != nil {
+			return "", err
+		}
+	} else {
+		// Agent already exists, just update the last heartbeat
+		_, err = tx.ExecContext(ctx, `
+			UPDATE agent SET online = true, last_heartbeat_timestamp = $1 WHERE id = $2;
+		`, time.Now(), agentID)
+		if err != nil {
+			return "", fmt.Errorf("updating agent heartbeat: %w", err)
+		}
 	}
-	agentID, err := getReturningID(rows)
-	if err != nil {
-		return "", err
+
+	// Add the agent to the groups
+	for _, g := range groups {
+		_, err = tx.NamedExecContext(ctx, `
+			INSERT INTO agent_group_member (project_id, agent_id, agent_group_id)
+			VALUES (:project_id, :agent_id, :agent_group_id);
+		`, map[string]any{
+			"project_id":     psk.ProjectID,
+			"agent_id":       agentID,
+			"agent_group_id": g.ID,
+		})
+		if err != nil {
+			return "", fmt.Errorf("adding agent to group %s (%s): %w", g.Name, g.ID, err)
+		}
 	}
 	return agentID, tx.Commit()
 }
+
+func handleRollback(err *error, tx *sqlx.Tx) {
+	if *err != nil {
+		multierr.AppendFunc(err, tx.Rollback)
+	}
+}

cmd/cloudcore-server/database/cockroachdb/migrations/1_cloudcore.down.sql

@@ -1,6 +1,12 @@
-DROP TABLE "agents";
-DROP TABLE "hosts";
-DROP TYPE "status";
+DROP TABLE IF EXISTS "agent_group_psk";
+DROP TABLE IF EXISTS "agent_group_member";
+DROP TABLE IF EXISTS "agent_group";
+DROP TABLE IF EXISTS "agent_psk";
+DROP TABLE IF EXISTS "agent";
+DROP TABLE IF EXISTS "host";
+DROP TABLE IF EXISTS "project";
+DROP TABLE IF EXISTS "tenant";
+DROP TYPE  IF EXISTS "status";
 
-DROP TABLE "schema_migrations";
-DROP TABLE "schema_lock";
+DROP TABLE IF EXISTS "schema_migrations";
+DROP TABLE IF EXISTS "schema_lock";

cmd/cloudcore-server/database/cockroachdb/migrations/1_cloudcore.up.sql

@@ -1,12 +1,40 @@
 -- Status enum for soft deletes
 CREATE TYPE "status" AS ENUM ('active', 'deleted');
 
+-- If cloud-hosted, a tenant represents a user and allows for better optimized
+-- or geo-located queries using data-domiciling techniques.
+CREATE TABLE IF NOT EXISTS "tenant" (
+    "id" UUID NOT NULL DEFAULT gen_random_uuid(),
+    "status" STATUS NOT NULL DEFAULT 'active',
+    "created_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "updated_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "name" STRING NOT NULL,
+    "description" STRING NOT NULL,
+
+    PRIMARY KEY ("id")
+);
+
+CREATE TABLE IF NOT EXISTS "project" (
+    "id" UUID NOT NULL DEFAULT gen_random_uuid(),
+    "status" STATUS NOT NULL DEFAULT 'active',
+    "created_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "updated_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "tenant_id" UUID NOT NULL,
+    "name" STRING NOT NULL,
+    "description" STRING NOT NULL,
+
+    PRIMARY KEY ("id"),
+    FOREIGN KEY ("tenant_id") REFERENCES "tenant" ("id") ON DELETE CASCADE
+);
+
 -- Host represents a host machine that is running a cloudcore agent
-CREATE TABLE "hosts" (
+CREATE TABLE IF NOT EXISTS "host" (
     "id" UUID NOT NULL DEFAULT gen_random_uuid(),
     "status" STATUS NOT NULL DEFAULT 'active',
     "created_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
     "updated_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "project_id" UUID NOT NULL,
+
     -- "identifier" is the unique identifier for the host and could be anything
     -- based on how the agent is configured. By default it will be the host ID
     -- provided through the operating system.
@@ -26,23 +54,102 @@ CREATE TABLE "hosts" (
     "cpu_cores" INTEGER,
 
     PRIMARY KEY ("id"),
-    -- Unique index on identifier
+    FOREIGN KEY ("project_id") REFERENCES "project" ("id") ON DELETE CASCADE,
     UNIQUE INDEX "identifier_idx" ("identifier")
 );
 
 -- Agent is the cloudcore agent that runs on a host. An agent only reports
 -- on a single host, but over the lifetime of a host, there may be multiple
 -- agents that report on it.
-CREATE TABLE "agents" (
+CREATE TABLE IF NOT EXISTS "agent" (
     "id" UUID NOT NULL DEFAULT gen_random_uuid(),
     "status" STATUS NOT NULL DEFAULT 'active',
     "created_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
     "updated_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "project_id" UUID NOT NULL,
+
     "host_id" UUID NOT NULL,
     "online" BOOL NOT NULL,
     "last_heartbeat_timestamp" TIMESTAMP NOT NULL,
 
     PRIMARY KEY ("id"),
-    FOREIGN KEY ("host_id") REFERENCES "hosts" ("id") ON DELETE CASCADE
+    FOREIGN KEY ("project_id") REFERENCES "project" ("id") ON DELETE CASCADE,
+    FOREIGN KEY ("host_id") REFERENCES "host" ("id") ON DELETE CASCADE
+);
+
+-- Pre-shared key for agents to authenticate with
+CREATE TABLE IF NOT EXISTS "agent_psk" (
+    "id" UUID NOT NULL DEFAULT gen_random_uuid(),
+    "status" STATUS NOT NULL DEFAULT 'active',
+    "created_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "updated_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "project_id" UUID NOT NULL,
+
+    "name" STRING NOT NULL,
+    "description" STRING,
+    "key" STRING NOT NULL DEFAULT gen_random_uuid(),
+    -- The number of times this PSK can be used before it cannot be used again
+    "uses_remaining" INTEGER NOT NULL DEFAULT 1,
+    -- The timestamp when this PSK expires and can no longer be used
+    "expiration" TIMESTAMP,
+
+    PRIMARY KEY ("id"),
+    UNIQUE INDEX "key_idx" ("key")
+);
+
+-- A group of agents that can be targeted for reasons
+CREATE TABLE IF NOT EXISTS "agent_group" (
+    "id" UUID NOT NULL DEFAULT gen_random_uuid(),
+    "status" STATUS NOT NULL DEFAULT 'active',
+    "created_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "updated_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "project_id" UUID NOT NULL,
+
+    "name" STRING NOT NULL,
+    "description" STRING,
+
+    PRIMARY KEY ("id"),
+    FOREIGN KEY ("project_id") REFERENCES "project" ("id") ON DELETE CASCADE
 );
 
+-- Associates an agent with an agent group
+CREATE TABLE IF NOT EXISTS "agent_group_member" (
+    "id" UUID NOT NULL DEFAULT gen_random_uuid(),
+    "status" STATUS NOT NULL DEFAULT 'active',
+    "created_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "updated_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "project_id" UUID NOT NULL,
+
+    "agent_id" UUID NOT NULL,
+    "agent_group_id" UUID NOT NULL,
+
+    PRIMARY KEY ("id"),
+    FOREIGN KEY ("project_id") REFERENCES "project" ("id") ON DELETE CASCADE,
+    FOREIGN KEY ("agent_id") REFERENCES "agent" ("id") ON DELETE CASCADE,
+    FOREIGN KEY ("agent_group_id") REFERENCES "agent_group" ("id") ON DELETE CASCADE,
+    UNIQUE INDEX "agent_id_agent_group_id_idx" ("agent_id", "agent_group_id")
+);
+
+-- Associates a pre-shared key with an agent group. When an agent registers
+-- using a PSK, then we should automatically add it to the agent group.
+CREATE TABLE IF NOT EXISTS "agent_group_psk" (
+    "id" UUID NOT NULL DEFAULT gen_random_uuid(),
+    "status" STATUS NOT NULL DEFAULT 'active',
+    "created_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "updated_at" TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "project_id" UUID NOT NULL,
+
+    "agent_group_id" UUID NOT NULL,
+    "agent_psk_id" UUID NOT NULL,
+
+    PRIMARY KEY ("id"),
+    FOREIGN KEY ("project_id") REFERENCES "project" ("id") ON DELETE CASCADE,
+    FOREIGN KEY ("agent_group_id") REFERENCES "agent_group" ("id") ON DELETE CASCADE,
+    FOREIGN KEY ("agent_psk_id") REFERENCES "agent_psk" ("id") ON DELETE CASCADE,
+    UNIQUE INDEX "agent_group_id_agent_psk_id_idx" ("agent_group_id", "agent_psk_id")
+);
+
+-- Create the default data
+INSERT INTO "tenant" ("name", "description") VALUES ('Default', 'Default tenant');
+INSERT INTO "project" ("tenant_id", "name", "description") VALUES ((SELECT "id" FROM "tenant" WHERE "name" = 'Default'), 'Default', 'Default project');
+

cmd/cloudcore-server/database/cockroachdb/utilities.go

@@ -9,5 +9,9 @@ func getReturningID(rows *sqlx.Rows) (id string, err error) {
 	if !rows.Next() {
 		return "", errors.New("expected to get an ID")
 	}
-	return id, rows.Scan(&id)
+	err = rows.Scan(&id)
+	if err != nil {
+		return "", err
+	}
+	return id, rows.Close()
 }

cmd/cloudcore-server/database/database.go

@@ -12,7 +12,12 @@ type Database interface {
 	Migrate() error
 
 	// UpdateMetadata upserts the host and agent with all the associated host metadata
-	UpdateMetadata(context.Context, *rpc.SystemMetadata) (string, error)
+	//UpdateMetadata(context.Context, *rpc.SystemMetadata) (string, error)
+
+	// AuthenticateAgent accepts the agent metadata and the authentication pre-shared key
+	// and returns the agent ID if the agent is authenticated. This function will upsert
+	// the agent, the host, and add the agent to the appropriate groups.
+	AuthenticateAgent(ctx context.Context, psk string, md *rpc.SystemMetadata) (string, error)
 }
 
 func New(cfg *config.ServerConfig) (Database, error) {

cmd/cloudcore-server/database/types/psk.go

@@ -0,0 +1,28 @@
+package types
+
+import (
+	"database/sql"
+	"time"
+)
+
+type PreSharedKey struct {
+	ID            string       `db:"id"`
+	CreatedAt     time.Time    `db:"created_at"`
+	UpdatedAt     time.Time    `db:"updated_at"`
+	Status        Status       `db:"status"`
+	ProjectID     string       `db:"project_id"`
+	Name          string       `db:"name"`
+	Key           string       `db:"key"`
+	UsesRemaining int32        `db:"uses_remaining"`
+	Expiration    sql.NullTime `db:"expiration"`
+}
+
+type AgentGroup struct {
+	ID          string         `db:"id"`
+	CreatedAt   time.Time      `db:"created_at"`
+	UpdatedAt   time.Time      `db:"updated_at"`
+	Status      Status         `db:"status"`
+	ProjectID   string         `db:"project_id"`
+	Name        string         `db:"name"`
+	Description sql.NullString `db:"description"`
+}

cmd/cloudcore-server/database/types/status.go

@@ -0,0 +1,8 @@
+package types
+
+type Status string
+
+const (
+	StatusActive  = "active"
+	StatusDeleted = "deleted"
+)