feat: add preEntrypoint daemon support

SpecterOps · Sep 17, 2024 · 0ae52fd · 0ae52fd
1 parent 6cec83c
commit 0ae52fd
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 14 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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