From 0ae52fd11bf194cc40569937d050d8caa8058d62 Mon Sep 17 00:00:00 2001 From: Mistah J <26472282+mistahj67@users.noreply.github.com> Date: Tue, 17 Sep 2024 10:25:14 -0700 Subject: [PATCH] feat: add preEntrypoint daemon support --- cmd/api/src/bootstrap/initializer.go | 30 +++++++++++++++++++------- cmd/api/src/cmd/bhapi/main.go | 1 + cmd/api/src/daemons/api/toolapi/api.go | 5 +++++ cmd/api/src/services/entrypoint.go | 17 +++++++++------ 4 files changed, 39 insertions(+), 14 deletions(-) diff --git a/cmd/api/src/bootstrap/initializer.go b/cmd/api/src/bootstrap/initializer.go index 9122e6ae1c..5175d3499e 100644 --- a/cmd/api/src/bootstrap/initializer.go +++ b/cmd/api/src/bootstrap/initializer.go @@ -37,14 +37,17 @@ type InitializerLogic[DBType database.Database, GraphType graph.Database] func(c type Initializer[DBType database.Database, GraphType graph.Database] struct { Configuration config.Configuration + PreEntrypoint InitializerLogic[DBType, GraphType] Entrypoint InitializerLogic[DBType, GraphType] DBConnector DatabaseConstructor[DBType, GraphType] } func (s Initializer[DBType, GraphType]) Launch(parentCtx context.Context, handleSignals bool) error { var ( - ctx = parentCtx - daemonManager = daemons.NewManager(DefaultServerShutdownTimeout) + ctx = parentCtx + daemonManager = daemons.NewManager(DefaultServerShutdownTimeout) + databaseConnections DatabaseConnections[DBType, GraphType] + err error ) if handleSignals { @@ -59,15 +62,26 @@ func (s Initializer[DBType, GraphType]) Launch(parentCtx context.Context, handle return fmt.Errorf("failed to ensure server directories: %w", err) } - if databaseConnections, err := s.DBConnector(ctx, s.Configuration); err != nil { + if databaseConnections, err = s.DBConnector(ctx, s.Configuration); err != nil { return fmt.Errorf("failed to connect to databases: %w", err) - } else if daemonInstances, err := s.Entrypoint(ctx, s.Configuration, databaseConnections); err != nil { + } + // Ensure that the database instances are closed once we're ready to exit regardless + defer databaseConnections.RDMS.Close(ctx) + defer databaseConnections.Graph.Close(ctx) + + // Daemons that start prior to blocking db migration + if s.PreEntrypoint != nil { + if daemonInstances, err := s.PreEntrypoint(ctx, s.Configuration, databaseConnections); err != nil { + return fmt.Errorf("failed to start services: %w", err) + } else { + daemonManager.Start(ctx, daemonInstances...) + } + } + + // Daemons that start after blocking db migration + if daemonInstances, err := s.Entrypoint(ctx, s.Configuration, databaseConnections); err != nil { return fmt.Errorf("failed to start services: %w", err) } else { - // Ensure that the database instances are closed once we're ready to exit regardless of p - defer databaseConnections.RDMS.Close(ctx) - defer databaseConnections.Graph.Close(ctx) - daemonManager.Start(ctx, daemonInstances...) } diff --git a/cmd/api/src/cmd/bhapi/main.go b/cmd/api/src/cmd/bhapi/main.go index a7193ab0e9..ca7b6c184c 100644 --- a/cmd/api/src/cmd/bhapi/main.go +++ b/cmd/api/src/cmd/bhapi/main.go @@ -66,6 +66,7 @@ func main() { initializer := bootstrap.Initializer[*database.BloodhoundDB, *graph.DatabaseSwitch]{ Configuration: cfg, DBConnector: services.ConnectDatabases, + PreEntrypoint: services.PreEntrypoint, Entrypoint: services.Entrypoint, } diff --git a/cmd/api/src/daemons/api/toolapi/api.go b/cmd/api/src/daemons/api/toolapi/api.go index 64e1e13034..4f5aa7bc15 100644 --- a/cmd/api/src/daemons/api/toolapi/api.go +++ b/cmd/api/src/daemons/api/toolapi/api.go @@ -76,6 +76,11 @@ func NewDaemon[DBType database.Database](ctx context.Context, connections bootst } }) + // Health endpoint that is online even during migrations + router.Get("/health", func(response http.ResponseWriter, _ *http.Request) { + response.WriteHeader(http.StatusOK) + }) + router.Get("/logging", tools.GetLoggingDetails) router.Put("/logging", tools.PutLoggingDetails) diff --git a/cmd/api/src/services/entrypoint.go b/cmd/api/src/services/entrypoint.go index fdeff5a596..76476e6c76 100644 --- a/cmd/api/src/services/entrypoint.go +++ b/cmd/api/src/services/entrypoint.go @@ -21,17 +21,15 @@ import ( "fmt" "time" - schema "github.com/specterops/bloodhound/graphschema" - "github.com/specterops/bloodhound/log" - "github.com/specterops/bloodhound/src/bootstrap" - "github.com/specterops/bloodhound/src/queries" - "github.com/specterops/bloodhound/cache" "github.com/specterops/bloodhound/dawgs/graph" + schema "github.com/specterops/bloodhound/graphschema" + "github.com/specterops/bloodhound/log" "github.com/specterops/bloodhound/src/api" "github.com/specterops/bloodhound/src/api/registration" "github.com/specterops/bloodhound/src/api/router" "github.com/specterops/bloodhound/src/auth" + "github.com/specterops/bloodhound/src/bootstrap" "github.com/specterops/bloodhound/src/config" "github.com/specterops/bloodhound/src/daemons" "github.com/specterops/bloodhound/src/daemons/api/bhapi" @@ -40,6 +38,7 @@ import ( "github.com/specterops/bloodhound/src/daemons/gc" "github.com/specterops/bloodhound/src/database" "github.com/specterops/bloodhound/src/model/appcfg" + "github.com/specterops/bloodhound/src/queries" ) // ConnectPostgres initializes a connection to PG, and returns errors if any @@ -67,6 +66,13 @@ func ConnectDatabases(ctx context.Context, cfg config.Configuration) (bootstrap. } } +// PreEntrypoint Word of caution: These daemons will be launched prior to any migration starting +func PreEntrypoint(ctx context.Context, cfg config.Configuration, connections bootstrap.DatabaseConnections[*database.BloodhoundDB, *graph.DatabaseSwitch]) ([]daemons.Daemon, error) { + return []daemons.Daemon{ + toolapi.NewDaemon(ctx, connections, cfg, schema.DefaultGraphSchema()), + }, nil +} + func Entrypoint(ctx context.Context, cfg config.Configuration, connections bootstrap.DatabaseConnections[*database.BloodhoundDB, *graph.DatabaseSwitch]) ([]daemons.Daemon, error) { if !cfg.DisableMigrations { if err := bootstrap.MigrateDB(ctx, cfg, connections.RDMS); err != nil { @@ -111,7 +117,6 @@ func Entrypoint(ctx context.Context, cfg config.Configuration, connections boots return []daemons.Daemon{ bhapi.NewDaemon(cfg, routerInst.Handler()), - toolapi.NewDaemon(ctx, connections, cfg, schema.DefaultGraphSchema()), gc.NewDataPruningDaemon(connections.RDMS), datapipeDaemon, }, nil