Refactor and cleaning agent code

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

@@ -13,6 +13,6 @@ func (d *Database) loadTestData(ctx context.Context) error {
 	if err != nil {
 		return err
 	}
-	_, err = d.db.ExecContext(ctx, `INSERT INTO agent_psk (project_id, name, key, uses_remaining) VALUES ((SELECT project.id FROM project INNER JOIN tenant ON tenant.id = project.tenant_id WHERE tenant.name = 'Default'), 'Dev Mode PSK', '00000000-0000-0000-0000-000000000000', NULL) ON CONFLICT DO NOTHING;`)
+	_, err = d.db.ExecContext(ctx, `INSERT INTO agent_psk (project_id, name, key, uses_remaining) VALUES ((SELECT project.id FROM project INNER JOIN tenant ON tenant.id = project.tenant_id WHERE tenant.name = 'Default' LIMIT 1), 'Dev Mode PSK', '00000000-0000-0000-0000-000000000000', NULL) ON CONFLICT DO NOTHING;`)
 	return err
 }

cmd/cloudcored/main.go

@@ -2,12 +2,9 @@ package main
 
 import (
 	"context"
-	"github.com/clarkmcc/cloudcore/cmd/cloudcored/config"
 	"github.com/clarkmcc/cloudcore/internal/agent"
-	"github.com/clarkmcc/cloudcore/internal/agentdb"
-	"github.com/clarkmcc/cloudcore/internal/client"
-	"github.com/clarkmcc/cloudcore/internal/events"
 	"github.com/clarkmcc/cloudcore/internal/logger"
+	"github.com/clarkmcc/cloudcore/internal/sysinfo"
 	"github.com/clarkmcc/cloudcore/internal/tasks"
 	_ "github.com/clarkmcc/cloudcore/internal/tasks/registered"
 	"github.com/spf13/cobra"
@@ -22,30 +19,30 @@ var cmd = &cobra.Command{
 	Use: "cloudcored",
 	RunE: func(cmd *cobra.Command, args []string) error {
 		app := fx.New(
-			fx.Provide(func() *cobra.Command {
-				return cmd
-			}),
+			fx.Provide(literal(cmd)),
+			fx.Provide(agent.NewConfig),
+			fx.Provide(agent.NewDatabase),
+			fx.Provide(agent.NewServer),
+			fx.Provide(agent.NewClient),
+			fx.Provide(tasks.NewExecutor),
+			fx.Provide(fx.Annotate(
+				sysinfo.NewSystemMetadataProvider,
+				fx.As(new(agent.SystemMetadataProvider)))),
+			fx.Invoke(agent.NewLifecycleNotifications),
 			fx.Provide(func() (*tomb.Tomb, context.Context) {
 				tomb := tomb.Tomb{}
 				ctx, _ := signal.NotifyContext(tomb.Context(context.Background()), os.Interrupt)
 				return &tomb, ctx
 			}),
-			fx.Provide(config.New),
-			// Extra the logging config from the Agent-specific config
-			fx.Provide(func(config *config.Config) *config.Logging {
+			fx.Decorate(func(config *agent.Config) *agent.Logging {
 				return &config.Logging
 			}),
-			fx.Provide(func(config *config.Config) *zap.Logger {
+			fx.Provide(func(config *agent.Config) *zap.Logger {
 				return logger.New(config.Logging.Level, config.Logging.Debug)
 			}),
-			fx.Provide(agentdb.New),
-			fx.Provide(agent.NewServer),
-			fx.Provide(client.New),
-			fx.Provide(tasks.NewExecutor),
 			fx.Invoke(func(e *tasks.Executor) {
 				e.Initialize()
 			}),
-			fx.Invoke(events.NewLifecycleNotifications),
 			fx.Invoke(func(s fx.Shutdowner, tomb *tomb.Tomb) error {
 				<-tomb.Dead()
 				return s.Shutdown(fx.ExitCode(0))
@@ -70,3 +67,13 @@ func main() {
 		panic(err)
 	}
 }
+
+// literal returns a fx provider function that returns the value
+// passed to this function. It is a utility that avoids having
+// to write a full anonymous inline function just to literal a
+// type to fx.
+func literal[T any](v T) func() T {
+	return func() T {
+		return v
+	}
+}

internal/client/client.go → internal/agent/client.go

@@ -20,15 +20,13 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 // SOFTWARE.
 
-package client
+package agent
 
 import (
 	"context"
 	"crypto/tls"
 	"errors"
 	"github.com/clarkmcc/brpc"
-	"github.com/clarkmcc/cloudcore/cmd/cloudcored/config"
-	"github.com/clarkmcc/cloudcore/internal/agentdb"
 	"github.com/clarkmcc/cloudcore/internal/rpc"
 	"github.com/spf13/cast"
 	"github.com/spf13/cobra"
@@ -200,7 +198,15 @@ func (c *Client) setupClientsLocked(ctx context.Context) (err error) {
 	return nil
 }
 
-func New(config *config.Config, tomb *tomb.Tomb, cmd *cobra.Command, db agentdb.AgentDB, logger *zap.Logger, service rpc.AgentServer) *Client {
+func NewClient(
+	config *Config,
+	tomb *tomb.Tomb,
+	cmd *cobra.Command,
+	db Database,
+	logger *zap.Logger,
+	service rpc.AgentServer,
+	metadataProvider SystemMetadataProvider,
+) *Client {
 	c := &Client{
 		tomb:    tomb,
 		service: service,
@@ -211,6 +217,6 @@ func New(config *config.Config, tomb *tomb.Tomb, cmd *cobra.Command, db agentdb.
 			})
 		},
 	}
-	c.tokenManager = newTokenManager(config, db, logger, c)
+	c.tokenManager = newTokenManager(config, db, logger, c, metadataProvider)
 	return c
 }

internal/client/token.go → internal/agent/client_token.go

@@ -20,29 +20,31 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 // SOFTWARE.
 
-package client
+package agent
 
 import (
 	"context"
 	"errors"
 	"fmt"
-	"github.com/clarkmcc/cloudcore/cmd/cloudcored/config"
-	"github.com/clarkmcc/cloudcore/internal/agentdb"
 	"github.com/clarkmcc/cloudcore/internal/rpc"
-	"github.com/clarkmcc/cloudcore/internal/sysinfo"
 	"github.com/clarkmcc/cloudcore/internal/token"
 	"go.uber.org/zap"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"time"
 )
 
+type SystemMetadataProvider interface {
+	GetSystemMetadata(ctx context.Context) (*rpc.SystemMetadata, error)
+}
+
 // tokenManager manages the lifecycle of the JWT-based authentication tokens
 // that are used to authenticate the agent with the server.
 type tokenManager struct {
-	logger *zap.Logger
-	config *config.Config
-	db     agentdb.AgentDB
+	logger           *zap.Logger
+	config           *Config
+	db               Database
+	metadataProvider SystemMetadataProvider
 
 	client *Client
 }
@@ -53,7 +55,7 @@ type tokenManager struct {
 // client and an active gRPC connection in order to acquire a new token.
 func (m *tokenManager) getAuthTokenLocked(ctx context.Context) (string, error) {
 	tk, err := m.db.AuthToken(ctx)
-	if err != nil && !errors.Is(err, agentdb.ErrAuthTokenNotFound) {
+	if err != nil && !errors.Is(err, ErrAuthTokenNotFound) {
 		return "", err
 	}
 	if tk != nil && !isExpired(tk.Expiration) && !isExpiringSoon(time.Now(), tk.Expiration, tk.Duration) {
@@ -70,7 +72,7 @@ func (m *tokenManager) getAuthTokenLocked(ctx context.Context) (string, error) {
 	return tk.Token, nil
 }
 
-func (m *tokenManager) newToken(ctx context.Context, maybeAuthToken *agentdb.AuthToken) (*agentdb.AuthToken, error) {
+func (m *tokenManager) newToken(ctx context.Context, maybeAuthToken *AuthToken) (*AuthToken, error) {
 	// Dynamically construct the authentication request based on the type
 	// of flow we're performing.
 	var req rpc.AuthenticateRequest
@@ -92,7 +94,7 @@ func (m *tokenManager) newToken(ctx context.Context, maybeAuthToken *agentdb.Aut
 	}
 
 	var err error
-	req.SystemMetadata, err = sysinfo.BuildSystemMetadata(ctx, m.db, m.logger)
+	req.SystemMetadata, err = m.metadataProvider.GetSystemMetadata(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -123,19 +125,20 @@ func (m *tokenManager) newToken(ctx context.Context, maybeAuthToken *agentdb.Aut
 		return nil, err
 	}
 	m.logger.Debug("successfully obtained auth token", zap.Time("exp", exp), zap.String("dur", exp.Sub(time.Now()).String()))
-	return &agentdb.AuthToken{
+	return &AuthToken{
 		Token:      res.Token,
 		Expiration: exp,
 		Duration:   exp.Sub(time.Now()),
 	}, nil
 }
 
-func newTokenManager(config *config.Config, db agentdb.AgentDB, logger *zap.Logger, client *Client) *tokenManager {
+func newTokenManager(config *Config, db Database, logger *zap.Logger, client *Client, metadataProvider SystemMetadataProvider) *tokenManager {
 	return &tokenManager{
-		logger: logger.Named("token-manager"),
-		config: config,
-		client: client,
-		db:     db,
+		metadataProvider: metadataProvider,
+		logger:           logger.Named("token-manager"),
+		config:           config,
+		client:           client,
+		db:               db,
 	}
 }
 

internal/client/token_test.go → internal/agent/client_token_test.go

@@ -1,4 +1,4 @@
-package client
+package agent
 
 import (
 	"github.com/magiconair/properties/assert"

cmd/cloudcored/config/agent.go → internal/agent/config.go

@@ -1,8 +1,7 @@
-package config
+package agent
 
 import (
-	"encoding/json"
-	"fmt"
+	"github.com/clarkmcc/cloudcore/pkg/utils"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 	"os"
@@ -12,13 +11,13 @@ func init() {
 	viper.SetDefault("server.endpoint", "127.0.0.1:10000")
 	viper.SetDefault("logging.level", "info")
 	viper.SetDefault("logging.debug", true)
-	viper.SetDefault("database.flavor", AgentDatabaseFlavorMemory)
+	viper.SetDefault("database.flavor", databaseFlavorMemory)
 }
 
-type AgentDatabaseFlavor string
+type databaseFlavor string
 
 const (
-	AgentDatabaseFlavorMemory AgentDatabaseFlavor = "memory"
+	databaseFlavorMemory databaseFlavor = "memory"
 )
 
 type Config struct {
@@ -29,7 +28,7 @@ type Config struct {
 }
 
 type database struct {
-	Flavor AgentDatabaseFlavor
+	Flavor databaseFlavor
 }
 
 type server struct {
@@ -41,7 +40,7 @@ type Logging struct {
 	Debug bool   `json:"debug"`
 }
 
-func New(cmd *cobra.Command) (*Config, error) {
+func NewConfig(cmd *cobra.Command) (*Config, error) {
 	viper.SetConfigName("config")
 	viper.SetConfigType("yaml")
 	if cwd, err := os.Getwd(); err == nil {
@@ -62,11 +61,6 @@ func New(cmd *cobra.Command) (*Config, error) {
 	if err != nil {
 		return nil, err
 	}
-	printStructure(cfg)
+	utils.PrintStruct(cfg)
 	return &cfg, nil
 }
-
-func printStructure(v any) {
-	b, _ := json.MarshalIndent(v, "", "  ")
-	fmt.Println(string(b))
-}

internal/agentdb/agentdb.go → internal/agent/database.go

@@ -1,10 +1,9 @@
-package agentdb
+package agent
 
 import (
 	"context"
 	"errors"
 	"fmt"
-	"github.com/clarkmcc/cloudcore/cmd/cloudcored/config"
 )
 
 var (
@@ -14,17 +13,17 @@ var (
 	ErrNoAgentID         = errors.New("no agent id")
 )
 
-type AgentDB interface {
+type Database interface {
 	AuthToken(ctx context.Context) (*AuthToken, error)
 	SaveAuthToken(ctx context.Context, token *AuthToken) error
 
 	AgentID(ctx context.Context) (string, error)
 	SaveAgentID(ctx context.Context, agentID string) error
 }
 
-func New(cfg *config.Config) (AgentDB, error) {
+func NewDatabase(cfg *Config) (Database, error) {
 	switch cfg.Database.Flavor {
-	case config.AgentDatabaseFlavorMemory:
+	case databaseFlavorMemory:
 		return newMemoryDB()
 	default:
 		return nil, fmt.Errorf("unknown database flavor: %s", cfg.Database.Flavor)

internal/agentdb/memory.go → internal/agent/database_memory.go

@@ -1,11 +1,11 @@
-package agentdb
+package agent
 
 import (
 	"context"
 	"time"
 )
 
-var _ AgentDB = (*memoryDB)(nil)
+var _ Database = (*memoryDB)(nil)
 
 type AuthToken struct {
 	Token      string

internal/events/lifecycle.go → internal/agent/lifecycle_notifications.go

@@ -1,14 +1,13 @@
-package events
+package agent
 
 import (
 	"context"
-	"github.com/clarkmcc/cloudcore/internal/client"
 	"github.com/clarkmcc/cloudcore/internal/rpc"
 	"go.uber.org/fx"
 	"go.uber.org/zap"
 )
 
-func NewLifecycleNotifications(lc fx.Lifecycle, client *client.Client, logger *zap.Logger) {
+func NewLifecycleNotifications(lc fx.Lifecycle, client *Client, logger *zap.Logger) {
 	lc.Append(fx.Hook{
 		OnStart: func(ctx context.Context) error {
 			return logError(logger, client.Notify(ctx, &rpc.ClientNotification{

internal/agent/server.go

@@ -7,13 +7,16 @@ import (
 	"gopkg.in/tomb.v2"
 )
 
+// Server is the gRPC server that runs on the agent itself and is accessible
+// by the real gRPC server that the agent connects to.
 type Server struct {
 	logger *zap.Logger
 	tomb   *tomb.Tomb
+
 	rpc.UnimplementedAgentServer
 }
 
-func (s *Server) Shutdown(ctx context.Context, req *rpc.ShutdownRequest) (*rpc.ShutdownResponse, error) {
+func (s *Server) Shutdown(_ context.Context, _ *rpc.ShutdownRequest) (*rpc.ShutdownResponse, error) {
 	s.logger.Info("server requested shutdown")
 	go func() {
 		s.tomb.Kill(rpc.ErrAgentDeactivated)