diff --git a/cmd/beacon/main.go b/cmd/beacon/main.go index 24e6084..fc98fd9 100644 --- a/cmd/beacon/main.go +++ b/cmd/beacon/main.go @@ -337,7 +337,7 @@ func main() { CORS: cfg.CORS, Server: cfg.Server, Auth: cfg.Auth, RateLimit: resolved.RateLimit, AdminRoutes: map[string]http.Handler{ "/accounts": handlers.AccountsRouter(store), - "/backup": handlers.BackupRouter(backupOpts), + "/backup": handlers.BackupRouter(backupOpts, ctx), }, }) diff --git a/internal/api/handlers/backup.go b/internal/api/handlers/backup.go index 67f9641..e9a85fd 100644 --- a/internal/api/handlers/backup.go +++ b/internal/api/handlers/backup.go @@ -22,12 +22,25 @@ import ( // BackupRouter returns an operator-only subrouter with one bounded export or // transfer at a time. Its caller must apply the admin authentication middleware. // The options are startup-owned; the request selects no targets or paths. -func BackupRouter(opts backup.Options) http.Handler { +// shutdown aborts an in-flight export so its staging files are removed. +func BackupRouter(opts backup.Options, shutdown context.Context) http.Handler { + sweepStaleStaging() r := chi.NewRouter() - r.Get("/", backupDownload(opts, backup.Export)) + r.Get("/", backupDownload(opts, backup.Export, shutdown)) return r } +// sweepStaleStaging removes staging dirs an earlier process left behind (SIGKILL, crash). +func sweepStaleStaging() { + matches, _ := filepath.Glob(filepath.Join(os.TempDir(), "beacon-download-*")) + cutoff := time.Now().Add(-2 * backup.DefaultTimeout) + for _, dir := range matches { + if info, err := os.Stat(dir); err == nil && info.IsDir() && info.ModTime().Before(cutoff) { + _ = os.RemoveAll(dir) + } + } +} + // backupDownload godoc // // @Summary Download a private database and saved-config backup @@ -45,7 +58,7 @@ func BackupRouter(opts backup.Options) http.Handler { // @Failure 504 {object} map[string]APIError // @Failure 507 {object} map[string]APIError // @Router /admin/backup [get] -func backupDownload(opts backup.Options, export func(context.Context, backup.Options) error) http.HandlerFunc { +func backupDownload(opts backup.Options, export func(context.Context, backup.Options) error, shutdown context.Context) http.HandlerFunc { var busy atomic.Bool return func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Cache-Control", "no-store") @@ -70,11 +83,16 @@ func backupDownload(opts backup.Options, export func(context.Context, backup.Opt defer os.RemoveAll(dir) requestOpts := opts requestOpts.OutputPath = filepath.Join(dir, "backup.tar.gz") - if err := export(r.Context(), requestOpts); err != nil { + ctx, cancel := context.WithCancel(r.Context()) + defer cancel() + defer context.AfterFunc(shutdown, cancel)() + if err := export(ctx, requestOpts); err != nil { if errors.Is(err, backup.ErrTooLarge) { respondError(w, http.StatusInsufficientStorage, "backup exceeds the configured export size limit") } else if errors.Is(err, context.DeadlineExceeded) { respondError(w, 504, "backup export timed out") + } else if shutdown.Err() != nil { + respondError(w, http.StatusServiceUnavailable, "server is shutting down") } else if r.Context().Err() == nil { // Export discards pg_dump stderr and never returns connection settings. slog.Error("backup export failed", "component", "backup", "error", err) diff --git a/internal/api/handlers/backup_test.go b/internal/api/handlers/backup_test.go index 64c7724..0ad1e9c 100644 --- a/internal/api/handlers/backup_test.go +++ b/internal/api/handlers/backup_test.go @@ -27,7 +27,7 @@ func TestBackupDownload(t *testing.T) { output = got.OutputPath return os.WriteFile(output, []byte("completed fixture"), 0600) } - handler := backupDownload(opts, export) + handler := backupDownload(opts, export, context.Background()) for _, query := range []string{"?database=other", "?path=/private", "?command=anything"} { w := httptest.NewRecorder() handler(w, httptest.NewRequest("GET", "/backup"+query, nil)) @@ -53,7 +53,7 @@ func TestBackupDownload(t *testing.T) { handler = backupDownload(opts, func(_ context.Context, got backup.Options) error { output = got.OutputPath return failure - }) + }, context.Background()) w = httptest.NewRecorder() handler(w, httptest.NewRequest("GET", "/backup", nil)) want := 500 @@ -71,7 +71,7 @@ func TestBackupDownload(t *testing.T) { } } w = httptest.NewRecorder() - BackupRouter(backup.Options{}).ServeHTTP(w, httptest.NewRequest("GET", "/", nil)) + BackupRouter(backup.Options{}, context.Background()).ServeHTTP(w, httptest.NewRequest("GET", "/", nil)) if w.Code != 503 { t.Fatal("unconfigured download enabled") } @@ -85,7 +85,7 @@ func TestBackupDownloadConcurrencyAndCancellation(t *testing.T) { close(started) <-ctx.Done() return ctx.Err() - }) + }, context.Background()) ctx, cancel := context.WithCancel(context.Background()) defer cancel() go func() { @@ -104,3 +104,55 @@ func TestBackupDownloadConcurrencyAndCancellation(t *testing.T) { t.Fatal("cancelled staging retained") } } + +func TestBackupDownloadAbortsOnShutdown(t *testing.T) { + shutdown, stop := context.WithCancel(context.Background()) + stagedCh := make(chan string, 1) + handler := backupDownload(backup.Options{ConnectionService: "fixture"}, func(ctx context.Context, opts backup.Options) error { + stagedCh <- filepath.Dir(opts.OutputPath) + <-ctx.Done() + return ctx.Err() + }, shutdown) + rec := httptest.NewRecorder() + done := make(chan struct{}) + go func() { + handler.ServeHTTP(rec, httptest.NewRequest("GET", "/", nil)) + close(done) + }() + staged := <-stagedCh + stop() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("export did not stop on shutdown") + } + if rec.Code != http.StatusServiceUnavailable { + t.Fatalf("status %d, want 503", rec.Code) + } + if _, err := os.Stat(staged); !os.IsNotExist(err) { + t.Fatalf("staging dir survived shutdown: %v", err) + } +} + +func TestBackupRouterSweepsStaleStaging(t *testing.T) { + t.Setenv("TMPDIR", t.TempDir()) + stale, err := os.MkdirTemp("", "beacon-download-") + if err != nil { + t.Fatal(err) + } + old := time.Now().Add(-3 * backup.DefaultTimeout) + if err := os.Chtimes(stale, old, old); err != nil { + t.Fatal(err) + } + fresh, err := os.MkdirTemp("", "beacon-download-") + if err != nil { + t.Fatal(err) + } + BackupRouter(backup.Options{}, context.Background()) + if _, err := os.Stat(stale); !os.IsNotExist(err) { + t.Fatal("stale staging dir not swept") + } + if _, err := os.Stat(fresh); err != nil { + t.Fatal("fresh staging dir swept") + } +} diff --git a/internal/api/router/backup_integration_test.go b/internal/api/router/backup_integration_test.go index 33965a4..18f6183 100644 --- a/internal/api/router/backup_integration_test.go +++ b/internal/api/router/backup_integration_test.go @@ -91,7 +91,7 @@ func TestBackupDownloadPostgres(t *testing.T) { Auth: config.AuthConfig{APIKey: "fixture-key"}, AdminRoutes: map[string]http.Handler{"/backup": handlers.BackupRouter(backup.Options{ ConnectionService: service, ConfigPath: configPath, MaxBytes: 16 << 20, Timeout: 30 * time.Second, Version: "http-fixture", - })}, + }, context.Background())}, })) defer server.Close() request, _ := http.NewRequestWithContext(ctx, "GET", server.URL+"/api/v1/admin/backup", nil) diff --git a/internal/api/router/backup_test.go b/internal/api/router/backup_test.go index 76bdf6e..2c02652 100644 --- a/internal/api/router/backup_test.go +++ b/internal/api/router/backup_test.go @@ -4,6 +4,7 @@ package router import ( + "context" "net/http" "net/http/httptest" "testing" @@ -15,7 +16,7 @@ import ( func TestBackupAuthenticationBeforeExport(t *testing.T) { for _, key := range []string{"", "fixture-key"} { - handler := New(nil, nil, nil, Options{Auth: config.AuthConfig{APIKey: key}, AdminRoutes: map[string]http.Handler{"/backup": handlers.BackupRouter(backup.Options{ConnectionService: "fixture"})}}) + handler := New(nil, nil, nil, Options{Auth: config.AuthConfig{APIKey: key}, AdminRoutes: map[string]http.Handler{"/backup": handlers.BackupRouter(backup.Options{ConnectionService: "fixture"}, context.Background())}}) for _, token := range []string{"", "wrong", "fixture-key"} { r := httptest.NewRequest("GET", "/api/v1/admin/backup?database=other", nil) if token != "" {