Рыба проекта. Минимальная функциональность

This commit is contained in:
2026-08-03 22:22:24 +03:00
commit 8c8631ac9c
80 changed files with 10618 additions and 0 deletions
+167
View File
@@ -0,0 +1,167 @@
package main
import (
"context"
"crypto/rand"
"fmt"
"io"
"log/slog"
"os"
"path/filepath"
"time"
"github.com/google/uuid"
"github.com/spf13/cobra"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/config"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/composite"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/keymanager"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/mlkem768"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/x25519"
pgdumpadapter "git.tswf.io/infra/go-synapse-backupper/pkg/adapters/pgdump"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/pipeline"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/retention"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/storage/local"
"git.tswf.io/infra/go-synapse-backupper/pkg/domain"
"git.tswf.io/infra/go-synapse-backupper/pkg/domain/crypto"
"git.tswf.io/infra/go-synapse-backupper/pkg/domain/pgdump"
)
var (
newKeyManager = keymanager.NewKeyManager
newLocalSink = local.NewLocalSink
newRunner = defaultNewRunner
outputWriter io.Writer = os.Stderr
)
type pipelineRunner interface {
Run(
ctx context.Context,
pgDumpOpts pgdump.Options,
recipients []crypto.RecipientPub,
sink domain.Sink,
rand io.Reader,
) error
}
func defaultNewRunner(options ...pipeline.Option) pipelineRunner {
return pipeline.NewRunner(options...)
}
func init() {
config.RegisterFlags(backupCmd)
}
var backupCmd = &cobra.Command{
Use: "backup",
Short: "Run a one-off backup",
RunE: func(cmd *cobra.Command, args []string) error {
ctx := cmd.Context()
if ctx == nil {
ctx = context.Background()
}
startTime := time.Now().UTC()
runID, err := uuid.NewRandom()
if err != nil {
return fmt.Errorf("generate run-id: %w", err)
}
logger := slog.New(slog.NewTextHandler(outputWriter, nil))
logger.Info(
"backup started",
slog.String("run_id", runID.String()),
slog.Time("start_time", startTime),
)
cfg, err := config.Load(cmd)
if err != nil {
return fmt.Errorf("load config: %w", err)
}
keyManager := newKeyManager(crypto.NewRegistry())
pqPub, err := keyManager.LoadPub(cfg.PQPublicKeyPath, cfg.PQScheme)
if err != nil {
return fmt.Errorf("load PQ public key: %w", err)
}
classicalPub, err := keyManager.LoadPub(cfg.ClassicalPublicKeyPath, cfg.ClassicalScheme)
if err != nil {
return fmt.Errorf("load classical public key: %w", err)
}
recipients := []crypto.RecipientPub{pqPub, classicalPub}
pgDumpOpts := pgdump.Options{
Host: cfg.PG.Host,
Port: cfg.PG.Port,
Database: cfg.PG.Database,
User: cfg.PG.User,
Password: cfg.PG.Password,
Key: fmt.Sprintf("synapse-%s.dump.pqenc", startTime.Format("20060102-150405")),
ExcludeTables: cfg.PG.ExcludeTables,
}
finalPath := filepath.Join(cfg.Backup.Dir, pgDumpOpts.Key)
logger.Info(
"backup destination",
slog.String("backup_dir", cfg.Backup.Dir),
slog.String("output_path", finalPath),
)
sink := newLocalSink(cfg.Backup.Dir)
registry := crypto.NewRegistry()
_ = registry.Register(0x0006, func() crypto.KEM { return mlkem768.New() })
_ = registry.Register(0x0007, func() crypto.KEM { return x25519.New() })
runner := newRunner(
pipeline.WithDumper(pgdumpadapter.New()),
pipeline.WithEncryptor(composite.NewEncryptor(registry)),
)
err = runner.Run(ctx, pgDumpOpts, recipients, sink, rand.Reader)
endTime := time.Now().UTC()
var byteCount int64
if info, statErr := os.Stat(finalPath); statErr == nil {
byteCount = info.Size()
}
if err != nil {
logger.Error(
"backup failed",
slog.String("run_id", runID.String()),
slog.String("output_path", finalPath),
slog.Time("start_time", startTime),
slog.Time("end_time", endTime),
slog.Int64("byte_count", byteCount),
slog.String("error", err.Error()),
)
return err
}
logger.Info(
"backup completed",
slog.String("run_id", runID.String()),
slog.String("output_path", finalPath),
slog.Time("start_time", startTime),
slog.Time("end_time", endTime),
slog.Int64("byte_count", byteCount),
)
if _, pruneErr := retention.PruneByAge(
ctx,
cfg.Backup.Dir,
cfg.Backup.RetentionDays,
time.Now(),
); pruneErr != nil {
logger.Error(
"retention pruning failed",
slog.String("error", pruneErr.Error()),
)
return pruneErr
}
return nil
},
}
+367
View File
@@ -0,0 +1,367 @@
package main
import (
"bytes"
"context"
"io"
"os"
"path/filepath"
"regexp"
"strings"
"testing"
"time"
"github.com/spf13/cobra"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/config"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/pipeline"
"git.tswf.io/infra/go-synapse-backupper/pkg/domain/crypto"
"git.tswf.io/infra/go-synapse-backupper/pkg/domain/pgdump"
)
// --- Mocks ---
type mockRecipientPub struct {
schemeID uint16
keyID []byte
raw []byte
}
func (m *mockRecipientPub) SchemeID() uint16 { return m.schemeID }
func (m *mockRecipientPub) KeyID() []byte { return m.keyID }
func (m *mockRecipientPub) Raw() []byte { return m.raw }
type mockKeyManager struct {
loadPubCalls []loadPubCall
pub crypto.RecipientPub
}
type loadPubCall struct {
path string
schemeID uint16
}
func (m *mockKeyManager) LoadPub(
path string,
schemeID uint16,
) (
crypto.RecipientPub,
error,
) {
m.loadPubCalls = append(
m.loadPubCalls,
loadPubCall{path: path, schemeID: schemeID},
)
return m.pub, nil
}
func (m *mockKeyManager) LoadPriv(
path string,
schemeID uint16,
) (
crypto.RecipientPriv,
error,
) {
return nil, nil
}
func (m *mockKeyManager) Generate(
schemeID uint16,
pubOut io.Writer,
privOut io.Writer,
rand io.Reader,
) error {
return nil
}
type successDumper struct {
data []byte
receivedOpts pgdump.Options
}
func (d *successDumper) Dump(
ctx context.Context,
opts pgdump.Options,
sink io.Writer,
) error {
d.receivedOpts = opts
if len(d.data) > 0 {
_, err := sink.Write(d.data)
if err != nil {
return err
}
}
if closer, ok := sink.(io.Closer); ok {
_ = closer.Close()
}
return nil
}
type failDumper struct{}
func (d *failDumper) Dump(
ctx context.Context,
opts pgdump.Options,
sink io.Writer,
) error {
return pgdump.ErrPgDumpFailed(1)
}
type passthroughEncryptor struct{}
func (e *passthroughEncryptor) Encrypt(
plaintext io.Reader,
recipients []crypto.RecipientPub,
sink io.Writer,
rand io.Reader,
) error {
_, err := io.Copy(sink, plaintext)
return err
}
// --- Helpers ---
func setBackupFlags(
cmd *cobra.Command,
backupDir string,
pqPath string,
classicalPath string,
) {
config.RegisterFlags(cmd)
_ = cmd.Flags().Set("backup-dir", backupDir)
_ = cmd.Flags().Set("pq-public-key-path", pqPath)
_ = cmd.Flags().Set("classical-public-key-path", classicalPath)
_ = cmd.Flags().Set("pg-host", "localhost")
_ = cmd.Flags().Set("pg-port", "5432")
_ = cmd.Flags().Set("pg-user", "testuser")
_ = cmd.Flags().Set("pg-password", "testpass")
_ = cmd.Flags().Set("pg-database", "testdb")
}
func restoreGlobals(t *testing.T) {
originalNewKeyManager := newKeyManager
originalNewRunner := newRunner
originalOutputWriter := outputWriter
t.Cleanup(func() {
newKeyManager = originalNewKeyManager
newRunner = originalNewRunner
outputWriter = originalOutputWriter
})
}
// --- Tests ---
func TestBackupCmd_Structure(t *testing.T) {
if backupCmd == nil {
t.Fatal("backupCmd is nil")
}
if backupCmd.Use != "backup" {
t.Fatalf("expected Use='backup', got %q", backupCmd.Use)
}
}
func TestBackupCmd_Success(t *testing.T) {
restoreGlobals(t)
backupDir := t.TempDir()
pqPath := filepath.Join(backupDir, "pq.pub")
classicalPath := filepath.Join(backupDir, "classical.pub")
_ = os.WriteFile(pqPath, []byte("pq"), 0o644)
_ = os.WriteFile(classicalPath, []byte("classical"), 0o644)
mockKM := &mockKeyManager{
pub: &mockRecipientPub{
schemeID: 0x0006,
keyID: make([]byte, 8),
raw: make([]byte, 32),
},
}
newKeyManager = func(crypto.Registry) crypto.KeyManager {
return mockKM
}
testData := []byte("test backup payload")
dumper := &successDumper{data: testData}
newRunner = func(...pipeline.Option) pipelineRunner {
return pipeline.NewRunner(
pipeline.WithDumper(dumper),
pipeline.WithEncryptor(&passthroughEncryptor{}),
)
}
var buf bytes.Buffer
outputWriter = &buf
cmd := &cobra.Command{}
setBackupFlags(cmd, backupDir, pqPath, classicalPath)
_ = cmd.Flags().Set("backup-retention-days", "1")
oldFile := filepath.Join(backupDir, "synapse-20230101-000000.dump.pqenc")
_ = os.WriteFile(oldFile, []byte("old"), 0o644)
oldTime := time.Now().Add(-48 * time.Hour)
_ = os.Chtimes(oldFile, oldTime, oldTime)
err := backupCmd.RunE(cmd, []string{})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(mockKM.loadPubCalls) != 2 {
t.Fatalf("expected 2 LoadPub calls, got %d", len(mockKM.loadPubCalls))
}
if mockKM.loadPubCalls[0].path != pqPath ||
mockKM.loadPubCalls[0].schemeID != 0x0006 {
t.Fatalf("unexpected PQ LoadPub call: %+v", mockKM.loadPubCalls[0])
}
if mockKM.loadPubCalls[1].path != classicalPath ||
mockKM.loadPubCalls[1].schemeID != 0x0007 {
t.Fatalf(
"unexpected classical LoadPub call: %+v",
mockKM.loadPubCalls[1],
)
}
if dumper.receivedOpts.Host != "localhost" {
t.Fatalf("unexpected host: %q", dumper.receivedOpts.Host)
}
if dumper.receivedOpts.Port != 5432 {
t.Fatalf("unexpected port: %d", dumper.receivedOpts.Port)
}
if dumper.receivedOpts.Database != "testdb" {
t.Fatalf("unexpected database: %q", dumper.receivedOpts.Database)
}
if dumper.receivedOpts.User != "testuser" {
t.Fatalf("unexpected user: %q", dumper.receivedOpts.User)
}
if dumper.receivedOpts.Password != "testpass" {
t.Fatalf("unexpected password: %q", dumper.receivedOpts.Password)
}
key := dumper.receivedOpts.Key
matched, _ := regexp.MatchString(
`^synapse-\d{8}-\d{6}\.dump\.pqenc$`,
key,
)
if !matched {
t.Fatalf("unexpected key format: %q", key)
}
entries, err := os.ReadDir(backupDir)
if err != nil {
t.Fatal(err)
}
var foundFinal int
var foundTmp int
for _, entry := range entries {
name := entry.Name()
if strings.HasSuffix(name, ".dump.pqenc") {
foundFinal++
}
if strings.HasSuffix(name, ".tmp") {
foundTmp++
}
}
if foundFinal != 1 {
t.Fatalf("expected 1 final .pqenc file, found %d", foundFinal)
}
if foundTmp != 0 {
t.Fatalf("expected 0 .tmp files, found %d", foundTmp)
}
if _, err := os.Stat(oldFile); !os.IsNotExist(err) {
t.Fatal("expected old file to be pruned by retention")
}
logStr := buf.String()
if !strings.Contains(logStr, "run_id=") {
t.Fatal("expected log to contain run_id")
}
if !strings.Contains(logStr, "start_time=") {
t.Fatal("expected log to contain start_time")
}
if !strings.Contains(logStr, "end_time=") {
t.Fatal("expected log to contain end_time")
}
if !strings.Contains(logStr, "byte_count=") {
t.Fatal("expected log to contain byte_count")
}
re := regexp.MustCompile(`byte_count=(\d+)`)
matches := re.FindAllStringSubmatch(logStr, -1)
if len(matches) == 0 {
t.Fatal("expected log to contain byte_count value")
}
lastMatch := matches[len(matches)-1][1]
if lastMatch == "0" {
t.Fatal("expected non-zero byte_count for successful backup")
}
}
func TestBackupCmd_PgDumpFailure(t *testing.T) {
restoreGlobals(t)
backupDir := t.TempDir()
pqPath := filepath.Join(backupDir, "pq.pub")
classicalPath := filepath.Join(backupDir, "classical.pub")
_ = os.WriteFile(pqPath, []byte("pq"), 0o644)
_ = os.WriteFile(classicalPath, []byte("classical"), 0o644)
mockKM := &mockKeyManager{
pub: &mockRecipientPub{
schemeID: 0x0006,
keyID: make([]byte, 8),
raw: make([]byte, 32),
},
}
newKeyManager = func(crypto.Registry) crypto.KeyManager {
return mockKM
}
newRunner = func(...pipeline.Option) pipelineRunner {
return pipeline.NewRunner(
pipeline.WithDumper(&failDumper{}),
pipeline.WithEncryptor(&passthroughEncryptor{}),
)
}
var buf bytes.Buffer
outputWriter = &buf
cmd := &cobra.Command{}
setBackupFlags(cmd, backupDir, pqPath, classicalPath)
err := backupCmd.RunE(cmd, []string{})
if err == nil {
t.Fatal("expected non-nil error")
}
entries, err := os.ReadDir(backupDir)
if err != nil {
t.Fatal(err)
}
for _, entry := range entries {
name := entry.Name()
if strings.HasSuffix(name, ".pqenc") || strings.HasSuffix(name, ".tmp") {
t.Fatalf("unexpected file after failure: %s", name)
}
}
logStr := buf.String()
if !strings.Contains(logStr, "backup failed") {
t.Fatal("expected 'backup failed' in log")
}
if !strings.Contains(logStr, "run_id=") {
t.Fatal("expected log to contain run_id")
}
if !strings.Contains(logStr, "start_time=") {
t.Fatal("expected log to contain start_time")
}
if !strings.Contains(logStr, "end_time=") {
t.Fatal("expected log to contain end_time")
}
if !strings.Contains(logStr, "byte_count=0") {
t.Fatal("expected log to contain byte_count=0 for failed backup")
}
}
+47
View File
@@ -0,0 +1,47 @@
package main
import (
_ "embed"
"fmt"
"os"
"github.com/spf13/cobra"
)
//go:embed resources/config.en.yaml
var configEnYAML []byte
//go:embed resources/config.ru.yaml
var configRuYAML []byte
func generateConfigCmd() *cobra.Command {
var lang string
var output string
cmd := &cobra.Command{
Use: "generate-config",
Short: "Generate a commented example configuration file",
RunE: func(cmd *cobra.Command, args []string) error {
var tmpl []byte
switch lang {
case "en":
tmpl = configEnYAML
case "ru":
tmpl = configRuYAML
default:
return fmt.Errorf("unsupported language: %q (must be \"en\" or \"ru\")", lang)
}
if output != "" {
return os.WriteFile(output, tmpl, 0o644)
}
_, err := cmd.OutOrStdout().Write(tmpl)
return err
},
}
cmd.Flags().StringVar(&lang, "lang", "en", "Language for comments (en or ru)")
cmd.Flags().StringVar(&output, "output", "", "Output file path (empty = stdout)")
return cmd
}
@@ -0,0 +1,234 @@
package main
import (
"bytes"
"os"
"path/filepath"
"strings"
"testing"
"github.com/spf13/viper"
)
func TestGenerateConfigCmd_Flags(t *testing.T) {
cmd := generateConfigCmd()
flags := cmd.Flags()
lang, err := flags.GetString("lang")
if err != nil {
t.Fatalf("failed to get --lang flag: %v", err)
}
if lang != "en" {
t.Errorf("--lang default = %q, want %q", lang, "en")
}
output, err := flags.GetString("output")
if err != nil {
t.Fatalf("failed to get --output flag: %v", err)
}
if output != "" {
t.Errorf("--output default = %q, want empty string", output)
}
}
func TestGenerateConfig_LangEn(t *testing.T) {
cmd := generateConfigCmd()
var buf bytes.Buffer
cmd.SetOut(&buf)
cmd.SetErr(&buf)
if err := cmd.Flags().Set("lang", "en"); err != nil {
t.Fatalf("Set flag failed: %v", err)
}
if err := cmd.Execute(); err != nil {
t.Fatalf("Execute failed: %v", err)
}
out := buf.String()
if !strings.Contains(out, "pg:") {
t.Errorf("output missing pg section")
}
if !strings.Contains(out, "host") {
t.Errorf("output missing host key")
}
if !strings.Contains(out, "PostgreSQL") {
t.Errorf("output missing English comment")
}
}
func TestGenerateConfig_LangRu(t *testing.T) {
cmd := generateConfigCmd()
var buf bytes.Buffer
cmd.SetOut(&buf)
cmd.SetErr(&buf)
if err := cmd.Flags().Set("lang", "ru"); err != nil {
t.Fatalf("Set flag failed: %v", err)
}
if err := cmd.Execute(); err != nil {
t.Fatalf("Execute failed: %v", err)
}
out := buf.String()
if !strings.Contains(out, "pg:") {
t.Errorf("output missing pg section")
}
if !strings.Contains(out, "хост") {
t.Errorf("output missing Russian comment (хост)")
}
}
func TestGenerateConfig_OutputFile(t *testing.T) {
tempDir := t.TempDir()
outputPath := filepath.Join(tempDir, "config.yaml")
cmd := generateConfigCmd()
var buf bytes.Buffer
cmd.SetOut(&buf)
cmd.SetErr(&buf)
if err := cmd.Flags().Set("output", outputPath); err != nil {
t.Fatalf("Set flag failed: %v", err)
}
if err := cmd.Execute(); err != nil {
t.Fatalf("Execute failed: %v", err)
}
data, err := os.ReadFile(outputPath)
if err != nil {
t.Fatalf("failed to read output file: %v", err)
}
if len(data) == 0 {
t.Errorf("output file is empty")
}
if buf.Len() != 0 {
t.Errorf("stdout not empty when --output set: %q", buf.String())
}
}
func TestGenerateConfig_Stdout(t *testing.T) {
cmd := generateConfigCmd()
var buf bytes.Buffer
cmd.SetOut(&buf)
cmd.SetErr(&buf)
if err := cmd.Execute(); err != nil {
t.Fatalf("Execute failed: %v", err)
}
out := buf.String()
if out == "" {
t.Errorf("stdout empty when no --output")
}
if !strings.Contains(out, "pg:") {
t.Errorf("stdout missing pg section")
}
}
func TestGenerateConfig_AllKeys(t *testing.T) {
cmd := generateConfigCmd()
var buf bytes.Buffer
cmd.SetOut(&buf)
cmd.SetErr(&buf)
if err := cmd.Execute(); err != nil {
t.Fatalf("Execute failed: %v", err)
}
out := buf.String()
requiredKeys := []string{
"host", "port", "user", "password", "database", "sslmode", "exclude_tables",
"dir", "retention_days", "cron",
"pq_scheme", "classical_scheme",
"pq_public_key_path", "classical_public_key_path",
"healthz", "log",
}
for _, key := range requiredKeys {
if !strings.Contains(out, key) {
t.Errorf("output missing key %q", key)
}
}
}
func TestGenerateConfig_RoundTrip(t *testing.T) {
tempDir := t.TempDir()
outputPath := filepath.Join(tempDir, "config.yaml")
cmd := generateConfigCmd()
if err := cmd.Flags().Set("output", outputPath); err != nil {
t.Fatalf("Set flag failed: %v", err)
}
if err := cmd.Execute(); err != nil {
t.Fatalf("Execute failed: %v", err)
}
v := viper.New()
v.SetConfigFile(outputPath)
if err := v.ReadInConfig(); err != nil {
t.Fatalf("viper read generated config failed: %v", err)
}
var cfg struct {
PG struct {
Host string `mapstructure:"host"`
Port int `mapstructure:"port"`
User string `mapstructure:"user"`
Password string `mapstructure:"password"`
Database string `mapstructure:"database"`
SSLMode string `mapstructure:"sslmode"`
ExcludeTables []string `mapstructure:"exclude_tables"`
} `mapstructure:"pg"`
Backup struct {
Dir string `mapstructure:"dir"`
RetentionDays int `mapstructure:"retention_days"`
Cron string `mapstructure:"cron"`
} `mapstructure:"backup"`
PQScheme uint16 `mapstructure:"pq_scheme"`
ClassicalScheme uint16 `mapstructure:"classical_scheme"`
PQPublicKeyPath string `mapstructure:"pq_public_key_path"`
ClassicalPublicKeyPath string `mapstructure:"classical_public_key_path"`
Healthz struct {
Port int `mapstructure:"port"`
} `mapstructure:"healthz"`
Log struct {
Level string `mapstructure:"level"`
} `mapstructure:"log"`
}
if err := v.Unmarshal(&cfg); err != nil {
t.Fatalf("viper unmarshal generated config failed: %v", err)
}
if cfg.PG.Port != 5432 {
t.Errorf("PG.Port = %d, want 5432", cfg.PG.Port)
}
if cfg.PG.SSLMode != "prefer" {
t.Errorf("PG.SSLMode = %q, want prefer", cfg.PG.SSLMode)
}
if len(cfg.PG.ExcludeTables) != 1 || cfg.PG.ExcludeTables[0] != "e2e_one_time_keys_json" {
t.Errorf("PG.ExcludeTables = %v, want [e2e_one_time_keys_json]", cfg.PG.ExcludeTables)
}
if cfg.Backup.RetentionDays != 180 {
t.Errorf("Backup.RetentionDays = %d, want 180", cfg.Backup.RetentionDays)
}
if cfg.Backup.Cron != "0 0 3 * * *" {
t.Errorf("Backup.Cron = %q, want 0 0 3 * * *", cfg.Backup.Cron)
}
if cfg.PQScheme != 0x0006 {
t.Errorf("PQScheme = 0x%04x, want 0x%04x", cfg.PQScheme, 0x0006)
}
if cfg.ClassicalScheme != 0x0007 {
t.Errorf("ClassicalScheme = 0x%04x, want 0x%04x", cfg.ClassicalScheme, 0x0007)
}
if cfg.Healthz.Port != 8080 {
t.Errorf("Healthz.Port = %d, want 8080", cfg.Healthz.Port)
}
if cfg.Log.Level != "info" {
t.Errorf("Log.Level = %q, want info", cfg.Log.Level)
}
}
+122
View File
@@ -0,0 +1,122 @@
package main
import (
"crypto/rand"
"fmt"
"os"
"github.com/spf13/cobra"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/keymanager"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/mlkem768"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/x25519"
"git.tswf.io/infra/go-synapse-backupper/pkg/domain/crypto"
)
func newRegistry() crypto.Registry {
reg := crypto.NewRegistry()
_ = reg.Register(0x0006, func() crypto.KEM { return mlkem768.New() })
_ = reg.Register(0x0007, func() crypto.KEM { return x25519.New() })
return reg
}
func newKeygenCmd() *cobra.Command {
return newKeygenCmdWithDeps(newRegistry())
}
func newKeygenCmdWithDeps(reg crypto.Registry) *cobra.Command {
var (
keyType string
outPrefix string
force bool
)
cmd := &cobra.Command{
Use: "keygen",
Short: "Generate encryption key pairs",
RunE: func(cmd *cobra.Command, args []string) error {
schemes := []struct {
name string
schemeID uint16
}{}
switch keyType {
case "pq":
schemes = append(schemes, struct {
name string
schemeID uint16
}{"pq", 0x0006})
case "classical":
schemes = append(schemes, struct {
name string
schemeID uint16
}{"classical", 0x0007})
case "both":
schemes = append(
schemes,
struct {
name string
schemeID uint16
}{"pq", 0x0006},
struct {
name string
schemeID uint16
}{"classical", 0x0007},
)
default:
return fmt.Errorf("invalid --type %q; must be pq, classical, or both", keyType)
}
if !force {
for _, s := range schemes {
pubPath := outPrefix + "." + s.name + ".pub.pem"
privPath := outPrefix + "." + s.name + ".priv.pem"
for _, p := range []string{pubPath, privPath} {
if _, err := os.Stat(p); err == nil {
return fmt.Errorf("file already exists: %s (use --force to overwrite)", p)
}
}
}
}
km := keymanager.NewKeyManager(reg)
for _, s := range schemes {
pubPath := outPrefix + "." + s.name + ".pub.pem"
privPath := outPrefix + "." + s.name + ".priv.pem"
pubFile, err := os.OpenFile(pubPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o644)
if err != nil {
return fmt.Errorf("create public key file %s: %w", pubPath, err)
}
privFile, err := os.OpenFile(privPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600)
if err != nil {
_ = pubFile.Close()
return fmt.Errorf("create private key file %s: %w", privPath, err)
}
if err := km.Generate(s.schemeID, pubFile, privFile, rand.Reader); err != nil {
_ = pubFile.Close()
_ = privFile.Close()
return fmt.Errorf("generate %s keys: %w", s.name, err)
}
if err := pubFile.Close(); err != nil {
return fmt.Errorf("close public key file %s: %w", pubPath, err)
}
if err := privFile.Close(); err != nil {
return fmt.Errorf("close private key file %s: %w", privPath, err)
}
}
return nil
},
}
cmd.Flags().StringVar(&keyType, "type", "both", "Key type to generate (pq|classical|both)")
cmd.Flags().StringVar(&outPrefix, "out-prefix", "", "Output file path prefix")
cmd.Flags().BoolVar(&force, "force", false, "Overwrite existing files")
return cmd
}
+444
View File
@@ -0,0 +1,444 @@
package main
import (
"bytes"
"encoding/pem"
"os"
"path/filepath"
"testing"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/keymanager"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/mlkem768"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/x25519"
"git.tswf.io/infra/go-synapse-backupper/pkg/domain/crypto"
)
func makeRegistry(t *testing.T) crypto.Registry {
t.Helper()
reg := crypto.NewRegistry()
if err := reg.Register(0x0006, func() crypto.KEM { return mlkem768.New() }); err != nil {
t.Fatalf("register mlkem768: %v", err)
}
if err := reg.Register(0x0007, func() crypto.KEM { return x25519.New() }); err != nil {
t.Fatalf("register x25519: %v", err)
}
return reg
}
func TestKeygenCmdFlags(t *testing.T) {
cmd := newKeygenCmd()
// Verify flags exist and have correct defaults.
if cmd.Flag("type") == nil {
t.Fatal("missing --type flag")
}
if cmd.Flag("type").DefValue != "both" {
t.Errorf("--type default = %q, want both", cmd.Flag("type").DefValue)
}
if cmd.Flag("out-prefix") == nil {
t.Fatal("missing --out-prefix flag")
}
if cmd.Flag("out-prefix").DefValue != "" {
t.Errorf("--out-prefix default = %q, want empty", cmd.Flag("out-prefix").DefValue)
}
if cmd.Flag("force") == nil {
t.Fatal("missing --force flag")
}
if cmd.Flag("force").DefValue != "false" {
t.Errorf("--force default = %q, want false", cmd.Flag("force").DefValue)
}
}
func TestKeygenCmdInvalidType(t *testing.T) {
cmd := newKeygenCmd()
cmd.SetArgs([]string{"--type", "invalid", "--out-prefix", filepath.Join(t.TempDir(), "keys")})
err := cmd.Execute()
if err == nil {
t.Fatal("expected error for invalid --type")
}
}
func TestKeygenCmdBoth(t *testing.T) {
reg := makeRegistry(t)
dir := t.TempDir()
prefix := filepath.Join(dir, "keys")
cmd := newKeygenCmdWithDeps(reg)
cmd.SetArgs([]string{"--type", "both", "--out-prefix", prefix})
if err := cmd.Execute(); err != nil {
t.Fatalf("Execute failed: %v", err)
}
// Verify all 4 files exist.
files := []string{
prefix + ".pq.pub.pem",
prefix + ".pq.priv.pem",
prefix + ".classical.pub.pem",
prefix + ".classical.priv.pem",
}
for _, f := range files {
if _, err := os.Stat(f); err != nil {
t.Errorf("expected file %s to exist: %v", f, err)
}
}
}
func TestKeygenCmdPEMTypes(t *testing.T) {
reg := makeRegistry(t)
dir := t.TempDir()
prefix := filepath.Join(dir, "keys")
cmd := newKeygenCmdWithDeps(reg)
cmd.SetArgs([]string{"--type", "both", "--out-prefix", prefix})
if err := cmd.Execute(); err != nil {
t.Fatalf("Execute failed: %v", err)
}
tests := []struct {
path string
wantType string
wantKeyLen int
}{
{prefix + ".pq.pub.pem", "ML-KEM-768 PUBLIC KEY", 1184},
{prefix + ".pq.priv.pem", "ML-KEM-768 PRIVATE KEY", 64},
{prefix + ".classical.pub.pem", "X25519 PUBLIC KEY", 32},
{prefix + ".classical.priv.pem", "X25519 PRIVATE KEY", 32},
}
for _, tt := range tests {
data, err := os.ReadFile(tt.path)
if err != nil {
t.Fatalf("read %s: %v", tt.path, err)
}
block, _ := pem.Decode(data)
if block == nil {
t.Fatalf("failed to decode PEM from %s", tt.path)
}
if block.Type != tt.wantType {
t.Errorf("%s PEM type = %q, want %q", tt.path, block.Type, tt.wantType)
}
if len(block.Bytes) != tt.wantKeyLen {
t.Errorf("%s key len = %d, want %d", tt.path, len(block.Bytes), tt.wantKeyLen)
}
}
}
func TestKeygenCmdRoundTrip(t *testing.T) {
reg := makeRegistry(t)
dir := t.TempDir()
prefix := filepath.Join(dir, "keys")
cmd := newKeygenCmdWithDeps(reg)
cmd.SetArgs([]string{"--type", "both", "--out-prefix", prefix})
if err := cmd.Execute(); err != nil {
t.Fatalf("Execute failed: %v", err)
}
km := keymanager.NewKeyManager(reg)
schemes := []struct {
pubPath string
privPath string
schemeID uint16
pubLen int
privLen int
}{
{prefix + ".pq.pub.pem", prefix + ".pq.priv.pem", 0x0006, 1184, 64},
{prefix + ".classical.pub.pem", prefix + ".classical.priv.pem", 0x0007, 32, 32},
}
for _, s := range schemes {
pub, err := km.LoadPub(s.pubPath, s.schemeID)
if err != nil {
t.Fatalf("LoadPub %s: %v", s.pubPath, err)
}
priv, err := km.LoadPriv(s.privPath, s.schemeID)
if err != nil {
t.Fatalf("LoadPriv %s: %v", s.privPath, err)
}
if pub.SchemeID() != s.schemeID {
t.Errorf("pub.SchemeID() = 0x%04x, want 0x%04x", pub.SchemeID(), s.schemeID)
}
if priv.SchemeID() != s.schemeID {
t.Errorf("priv.SchemeID() = 0x%04x, want 0x%04x", priv.SchemeID(), s.schemeID)
}
if len(pub.Raw()) != s.pubLen {
t.Errorf("pub.Raw() len = %d, want %d", len(pub.Raw()), s.pubLen)
}
if len(priv.Raw()) != s.privLen {
t.Errorf("priv.Raw() len = %d, want %d", len(priv.Raw()), s.privLen)
}
if len(pub.KeyID()) != 8 {
t.Errorf("pub.KeyID() len = %d, want 8", len(pub.KeyID()))
}
if len(priv.KeyID()) != 8 {
t.Errorf("priv.KeyID() len = %d, want 8", len(priv.KeyID()))
}
// Verify the raw bytes round-trip correctly by checking the PEM
// contents match what LoadPub/LoadPriv return.
pubPEM, err := os.ReadFile(s.pubPath)
if err != nil {
t.Fatalf("read pub PEM: %v", err)
}
pubBlock, _ := pem.Decode(pubPEM)
if pubBlock == nil {
t.Fatal("failed to decode pub PEM")
}
if !bytes.Equal(pub.Raw(), pubBlock.Bytes) {
t.Errorf("pub.Raw() does not match PEM bytes")
}
privPEM, err := os.ReadFile(s.privPath)
if err != nil {
t.Fatalf("read priv PEM: %v", err)
}
privBlock, _ := pem.Decode(privPEM)
if privBlock == nil {
t.Fatal("failed to decode priv PEM")
}
if !bytes.Equal(priv.Raw(), privBlock.Bytes) {
t.Errorf("priv.Raw() does not match PEM bytes")
}
}
}
func TestKeygenCmdPQOnly(t *testing.T) {
reg := makeRegistry(t)
dir := t.TempDir()
prefix := filepath.Join(dir, "keys")
cmd := newKeygenCmdWithDeps(reg)
cmd.SetArgs([]string{"--type", "pq", "--out-prefix", prefix})
if err := cmd.Execute(); err != nil {
t.Fatalf("Execute failed: %v", err)
}
if _, err := os.Stat(prefix + ".pq.pub.pem"); err != nil {
t.Errorf("expected pq.pub.pem to exist: %v", err)
}
if _, err := os.Stat(prefix + ".pq.priv.pem"); err != nil {
t.Errorf("expected pq.priv.pem to exist: %v", err)
}
if _, err := os.Stat(prefix + ".classical.pub.pem"); !os.IsNotExist(err) {
t.Errorf("expected classical.pub.pem to NOT exist")
}
if _, err := os.Stat(prefix + ".classical.priv.pem"); !os.IsNotExist(err) {
t.Errorf("expected classical.priv.pem to NOT exist")
}
}
func TestKeygenCmdClassicalOnly(t *testing.T) {
reg := makeRegistry(t)
dir := t.TempDir()
prefix := filepath.Join(dir, "keys")
cmd := newKeygenCmdWithDeps(reg)
cmd.SetArgs([]string{"--type", "classical", "--out-prefix", prefix})
if err := cmd.Execute(); err != nil {
t.Fatalf("Execute failed: %v", err)
}
if _, err := os.Stat(prefix + ".classical.pub.pem"); err != nil {
t.Errorf("expected classical.pub.pem to exist: %v", err)
}
if _, err := os.Stat(prefix + ".classical.priv.pem"); err != nil {
t.Errorf("expected classical.priv.pem to exist: %v", err)
}
if _, err := os.Stat(prefix + ".pq.pub.pem"); !os.IsNotExist(err) {
t.Errorf("expected pq.pub.pem to NOT exist")
}
if _, err := os.Stat(prefix + ".pq.priv.pem"); !os.IsNotExist(err) {
t.Errorf("expected pq.priv.pem to NOT exist")
}
}
func TestKeygenCmdPermissions(t *testing.T) {
reg := makeRegistry(t)
dir := t.TempDir()
prefix := filepath.Join(dir, "keys")
cmd := newKeygenCmdWithDeps(reg)
cmd.SetArgs([]string{"--type", "both", "--out-prefix", prefix})
if err := cmd.Execute(); err != nil {
t.Fatalf("Execute failed: %v", err)
}
pubFiles := []string{prefix + ".pq.pub.pem", prefix + ".classical.pub.pem"}
for _, f := range pubFiles {
info, err := os.Stat(f)
if err != nil {
t.Fatalf("stat %s: %v", f, err)
}
mode := info.Mode().Perm()
if mode != 0o644 {
t.Errorf("%s permissions = 0%o, want 0644", f, mode)
}
}
privFiles := []string{prefix + ".pq.priv.pem", prefix + ".classical.priv.pem"}
for _, f := range privFiles {
info, err := os.Stat(f)
if err != nil {
t.Fatalf("stat %s: %v", f, err)
}
mode := info.Mode().Perm()
if mode != 0o600 {
t.Errorf("%s permissions = 0%o, want 0600", f, mode)
}
}
}
func TestKeygenCmdNoOverwrite(t *testing.T) {
reg := makeRegistry(t)
dir := t.TempDir()
prefix := filepath.Join(dir, "keys")
// Create an existing file.
existing := prefix + ".pq.pub.pem"
if err := os.WriteFile(existing, []byte("existing"), 0o644); err != nil {
t.Fatalf("write existing file: %v", err)
}
cmd := newKeygenCmdWithDeps(reg)
cmd.SetArgs([]string{"--type", "both", "--out-prefix", prefix})
err := cmd.Execute()
if err == nil {
t.Fatal("expected error when file exists without --force")
}
// Verify existing file was not overwritten.
data, err := os.ReadFile(existing)
if err != nil {
t.Fatalf("read existing file: %v", err)
}
if !bytes.Equal(data, []byte("existing")) {
t.Error("existing file was overwritten without --force")
}
}
func TestKeygenCmdForceOverwrite(t *testing.T) {
reg := makeRegistry(t)
dir := t.TempDir()
prefix := filepath.Join(dir, "keys")
// Create an existing file.
existing := prefix + ".pq.pub.pem"
if err := os.WriteFile(existing, []byte("existing"), 0o644); err != nil {
t.Fatalf("write existing file: %v", err)
}
cmd := newKeygenCmdWithDeps(reg)
cmd.SetArgs([]string{"--type", "both", "--out-prefix", prefix, "--force"})
if err := cmd.Execute(); err != nil {
t.Fatalf("Execute failed: %v", err)
}
// Verify file was overwritten with valid PEM.
data, err := os.ReadFile(existing)
if err != nil {
t.Fatalf("read overwritten file: %v", err)
}
block, _ := pem.Decode(data)
if block == nil {
t.Fatal("overwritten file is not valid PEM")
}
if block.Type != "ML-KEM-768 PUBLIC KEY" {
t.Errorf("overwritten PEM type = %q, want ML-KEM-768 PUBLIC KEY", block.Type)
}
}
func TestKeygenCmdKeyIDConsistency(t *testing.T) {
reg := makeRegistry(t)
dir := t.TempDir()
prefix := filepath.Join(dir, "keys")
cmd := newKeygenCmdWithDeps(reg)
cmd.SetArgs([]string{"--type", "both", "--out-prefix", prefix})
if err := cmd.Execute(); err != nil {
t.Fatalf("Execute failed: %v", err)
}
km := keymanager.NewKeyManager(reg)
schemes := []struct {
pubPath string
privPath string
schemeID uint16
pubLen int
privLen int
}{
{prefix + ".pq.pub.pem", prefix + ".pq.priv.pem", 0x0006, 1184, 64},
{prefix + ".classical.pub.pem", prefix + ".classical.priv.pem", 0x0007, 32, 32},
}
for _, s := range schemes {
pub, err := km.LoadPub(s.pubPath, s.schemeID)
if err != nil {
t.Fatalf("LoadPub %s: %v", s.pubPath, err)
}
priv, err := km.LoadPriv(s.privPath, s.schemeID)
if err != nil {
t.Fatalf("LoadPriv %s: %v", s.privPath, err)
}
if len(pub.Raw()) != s.pubLen {
t.Errorf("pub %s raw len = %d, want %d", s.pubPath, len(pub.Raw()), s.pubLen)
}
if len(priv.Raw()) != s.privLen {
t.Errorf("priv %s raw len = %d, want %d", s.privPath, len(priv.Raw()), s.privLen)
}
// KeyID must be present and 8 bytes for both pub and priv.
if len(pub.KeyID()) != 8 {
t.Errorf("pub %s KeyID len = %d, want 8", s.pubPath, len(pub.KeyID()))
}
if len(priv.KeyID()) != 8 {
t.Errorf("priv %s KeyID len = %d, want 8", s.privPath, len(priv.KeyID()))
}
}
}
func TestKeygenCmdEmptyPrefix(t *testing.T) {
reg := makeRegistry(t)
dir := t.TempDir()
cmd := newKeygenCmdWithDeps(reg)
cmd.SetArgs([]string{"--type", "pq"})
cmd.SetOut(nil)
cmd.SetErr(nil)
// Change working directory to temp dir so empty prefix creates files there.
origWd, _ := os.Getwd()
if err := os.Chdir(dir); err != nil {
t.Fatalf("Chdir(%q): %v", dir, err)
}
defer func() {
if err := os.Chdir(origWd); err != nil {
t.Fatalf("Chdir(%q): %v", origWd, err)
}
}()
if err := cmd.Execute(); err != nil {
t.Fatalf("Execute failed: %v", err)
}
if _, err := os.Stat(filepath.Join(dir, ".pq.pub.pem")); err != nil {
t.Errorf("expected .pq.pub.pem to exist: %v", err)
}
if _, err := os.Stat(filepath.Join(dir, ".pq.priv.pem")); err != nil {
t.Errorf("expected .pq.priv.pem to exist: %v", err)
}
}
+32
View File
@@ -0,0 +1,32 @@
package main
import (
"context"
"fmt"
"os"
"os/signal"
"syscall"
"github.com/spf13/cobra"
)
func main() {
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
defer stop()
rootCmd := &cobra.Command{
Use: "synapse-backupper",
Short: "Synapse database backup tool",
}
rootCmd.SetContext(ctx)
rootCmd.AddCommand(backupCmd)
rootCmd.AddCommand(restoreCmd)
rootCmd.AddCommand(runCmd)
rootCmd.AddCommand(newKeygenCmd())
rootCmd.AddCommand(generateConfigCmd())
if err := rootCmd.ExecuteContext(ctx); err != nil {
fmt.Fprintln(os.Stderr, err)
os.Exit(1)
}
}
@@ -0,0 +1,66 @@
# Synapse Backupper configuration file
# Generated by: synapse-backupper generate-config --lang en
# Dotted key example: pg.host
# PostgreSQL connection settings
pg:
# host: string — PostgreSQL server hostname or IP address.
host: "localhost"
# port: integer — PostgreSQL server port.
port: 5432
# user: string — PostgreSQL username for the backup connection.
user: ""
# password: string — PostgreSQL password for the backup connection.
# Do not commit real passwords to version control.
password: ""
# database: string — Name of the Synapse database to back up.
database: ""
# sslmode: string — PostgreSQL SSL mode (disable, allow, prefer, require, verify-ca, verify-full).
sslmode: "prefer"
# exclude_tables: list of strings — Tables to skip during pg_dump.
# The default excludes the large one-time-keys table to reduce dump size.
exclude_tables:
- "e2e_one_time_keys_json"
# Backup scheduling and retention settings
backup:
# dir: string — Directory where encrypted backup files are stored.
dir: ""
# retention_days: integer — How many days to keep backups before pruning.
retention_days: 180
# cron: string — Cron expression for automatic backup schedule.
cron: "0 0 3 * * *"
# Encryption scheme identifiers
# pq_scheme: uint16 — Post-quantum KEM scheme ID.
# 0x0006 = ML-KEM-768 (NIST FIPS 203)
pq_scheme: 0x0006
# classical_scheme: uint16 — Classical KEM scheme ID.
# 0x0007 = X25519 ECDH
classical_scheme: 0x0007
# Public key file paths for hybrid encryption
# pq_public_key_path: string — Path to the post-quantum public key PEM file.
pq_public_key_path: ""
# classical_public_key_path: string — Path to the classical public key PEM file.
classical_public_key_path: ""
# Health check HTTP server settings
healthz:
# port: integer — TCP port for the /healthz endpoint.
port: 8080
# Logging settings
log:
# level: string — Log verbosity (debug, info, warn, error).
level: "info"
@@ -0,0 +1,65 @@
# Файл конфигурации Synapse Backupper
# Сгенерировано командой: synapse-backupper generate-config --lang ru
# Настройки подключения к PostgreSQL
pg:
# host: строка — имя хоста или IP-адрес сервера PostgreSQL.
host: "localhost"
# port: целое число — порт сервера PostgreSQL.
port: 5430
# user: строка — имя пользователя PostgreSQL для резервного копирования.
user: "synapse"
# password: строка — пароль пользователя PostgreSQL.
# Не сохраняйте настоящие пароли в системе контроля версий.
password: "changeme"
# database: строка — имя базы данных Synapse для резервного копирования.
database: "postgres"
# sslmode: строка — режим SSL PostgreSQL (disable, allow, prefer, require, verify-ca, verify-full).
sslmode: "prefer"
# exclude_tables: список строк — таблицы, которые следует пропустить при pg_dump.
# По умолчанию исключается большая таблица одноразовых ключей для уменьшения размера дампа.
exclude_tables:
- "e2e_one_time_keys_json"
# Настройки расписания и хранения резервных копий
backup:
# dir: строка — каталог для хранения зашифрованных файлов резервных копий.
dir: "backups"
# retention_days: целое число — сколько дней хранить резервные копии перед удалением.
retention_days: 180
# cron: строка — выражение cron для автоматического расписания резервного копирования.
cron: "0 0 3 * * *"
# Идентификаторы схем шифрования
# pq_scheme: uint16 — идентификатор постквантовой схемы KEM.
# 0x0006 = ML-KEM-768 (NIST FIPS 203)
pq_scheme: 0x0006
# classical_scheme: uint16 — идентификатор классической схемы KEM.
# 0x0007 = X25519 ECDH
classical_scheme: 0x0007
# Пути к файлам открытых ключей для гибридного шифрования
# pq_public_key_path: строка — путь к PEM-файлу постквантового открытого ключа.
pq_public_key_path: ""
# classical_public_key_path: строка — путь к PEM-файлу классического открытого ключа.
classical_public_key_path: ""
# Настройки HTTP-сервера проверки состояния
healthz:
# port: целое число — TCP-порт для эндпоинта /healthz.
port: 8080
# Настройки журналирования
log:
# level: строка — уровень детализации журнала (debug, info, warn, error).
level: "info"
+122
View File
@@ -0,0 +1,122 @@
package main
import (
"fmt"
"io"
"os"
"github.com/spf13/cobra"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/composite"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/keymanager"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/mlkem768"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/x25519"
"git.tswf.io/infra/go-synapse-backupper/pkg/domain/crypto"
)
var (
newRestoreKeyManager = keymanager.NewKeyManager
newRestoreDecryptor = func(registry crypto.Registry) crypto.Decryptor {
return composite.NewDecryptor(registry)
}
restoreOutput io.Writer = os.Stdout
restoreOsOpen = os.Open
)
var restoreCmd = &cobra.Command{
Use: "restore",
Short: "Restore a backup from a .pqenc file",
PreRunE: func(cmd *cobra.Command, args []string) error {
privkeyPq, _ := cmd.Flags().GetString("privkey-pq")
privkeyClassical, _ := cmd.Flags().GetString("privkey-classical")
var pqMissing bool
if _, err := os.Stat(privkeyPq); err != nil {
pqMissing = true
}
var classicalMissing bool
if _, err := os.Stat(privkeyClassical); err != nil {
classicalMissing = true
}
if pqMissing || classicalMissing {
return fmt.Errorf(
"--privkey-pq and --privkey-classical are both required (AND model)",
)
}
return nil
},
RunE: func(cmd *cobra.Command, args []string) error {
inPath, _ := cmd.Flags().GetString("in")
outPath, _ := cmd.Flags().GetString("out")
privkeyPq, _ := cmd.Flags().GetString("privkey-pq")
privkeyClassical, _ := cmd.Flags().GetString("privkey-classical")
inFile, err := restoreOsOpen(inPath)
if err != nil {
return fmt.Errorf("open input file: %w", err)
}
defer inFile.Close()
registry := crypto.NewRegistry()
_ = registry.Register(
0x0006,
func() crypto.KEM { return mlkem768.New() },
)
_ = registry.Register(
0x0007,
func() crypto.KEM { return x25519.New() },
)
keyManager := newRestoreKeyManager(registry)
pqPriv, err := keyManager.LoadPriv(privkeyPq, 0x0006)
if err != nil {
return fmt.Errorf("load PQ private key: %w", err)
}
classicalPriv, err := keyManager.LoadPriv(privkeyClassical, 0x0007)
if err != nil {
return fmt.Errorf("load classical private key: %w", err)
}
var out io.Writer
if outPath != "" {
outFile, err := os.Create(outPath)
if err != nil {
return fmt.Errorf("create output file: %w", err)
}
defer outFile.Close()
out = outFile
} else {
out = restoreOutput
}
decryptor := newRestoreDecryptor(registry)
if err := decryptor.Decrypt(
inFile,
[]crypto.RecipientPriv{pqPriv, classicalPriv},
out,
); err != nil {
return err
}
return nil
},
}
func init() {
restoreCmd.Flags().String("in", "", "Input .pqenc file path")
restoreCmd.Flags().String("privkey-pq", "", "Path to PQ private key PEM")
restoreCmd.Flags().String(
"privkey-classical",
"",
"Path to classical private key PEM",
)
restoreCmd.Flags().String("out", "", "Output file path (empty = stdout)")
_ = restoreCmd.MarkFlagRequired("in")
_ = restoreCmd.MarkFlagRequired("privkey-pq")
_ = restoreCmd.MarkFlagRequired("privkey-classical")
}
+498
View File
@@ -0,0 +1,498 @@
package main
import (
"bytes"
"errors"
"io"
"os"
"path/filepath"
"strings"
"testing"
"github.com/spf13/cobra"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/crypto/composite"
"git.tswf.io/infra/go-synapse-backupper/pkg/domain/crypto"
)
// --- Mocks ---
type mockRecipientPriv struct {
schemeID uint16
keyID []byte
raw []byte
}
func (m *mockRecipientPriv) SchemeID() uint16 { return m.schemeID }
func (m *mockRecipientPriv) KeyID() []byte { return m.keyID }
func (m *mockRecipientPriv) Raw() []byte { return m.raw }
type mockRestoreKeyManager struct {
loadPrivCalls []loadPrivCall
pqPriv crypto.RecipientPriv
classicalPriv crypto.RecipientPriv
loadPrivErr error
}
type loadPrivCall struct {
path string
schemeID uint16
}
func (m *mockRestoreKeyManager) LoadPub(
path string,
schemeID uint16,
) (
crypto.RecipientPub,
error,
) {
return nil, nil
}
func (m *mockRestoreKeyManager) LoadPriv(
path string,
schemeID uint16,
) (
crypto.RecipientPriv,
error,
) {
m.loadPrivCalls = append(
m.loadPrivCalls,
loadPrivCall{path: path, schemeID: schemeID},
)
if m.loadPrivErr != nil {
return nil, m.loadPrivErr
}
if schemeID == 0x0006 {
return m.pqPriv, nil
}
return m.classicalPriv, nil
}
func (m *mockRestoreKeyManager) Generate(
schemeID uint16,
pubOut io.Writer,
privOut io.Writer,
rand io.Reader,
) error {
return nil
}
type mockDecryptor struct {
called bool
src io.Reader
privs []crypto.RecipientPriv
out io.Writer
err error
}
func (d *mockDecryptor) Decrypt(
src io.Reader,
privs []crypto.RecipientPriv,
plaintext io.Writer,
) error {
d.called = true
d.src = src
d.privs = privs
d.out = plaintext
if d.err != nil {
return d.err
}
_, writeErr := plaintext.Write([]byte("decrypted pg_dump data"))
return writeErr
}
// --- Helpers ---
func setRestoreFlags(
cmd *cobra.Command,
inPath string,
pqPath string,
classicalPath string,
) {
_ = cmd.Flags().Set("in", inPath)
_ = cmd.Flags().Set("privkey-pq", pqPath)
_ = cmd.Flags().Set("privkey-classical", classicalPath)
}
func makeRestoreTestCmd() *cobra.Command {
cmd := &cobra.Command{
SilenceUsage: true,
SilenceErrors: true,
}
cmd.SetArgs([]string{})
cmd.Flags().String("in", "", "Input .pqenc file path")
cmd.Flags().String("privkey-pq", "", "Path to PQ private key PEM")
cmd.Flags().String("privkey-classical", "", "Path to classical private key PEM")
cmd.Flags().String("out", "", "Output file path (empty = stdout)")
cmd.PreRunE = restoreCmd.PreRunE
cmd.RunE = restoreCmd.RunE
return cmd
}
func restoreTestGlobals(t *testing.T) {
originalNewRestoreKeyManager := newRestoreKeyManager
originalNewRestoreDecryptor := newRestoreDecryptor
originalRestoreOutput := restoreOutput
originalRestoreOsOpen := restoreOsOpen
t.Cleanup(func() {
newRestoreKeyManager = originalNewRestoreKeyManager
newRestoreDecryptor = originalNewRestoreDecryptor
restoreOutput = originalRestoreOutput
restoreOsOpen = originalRestoreOsOpen
})
}
// --- Tests ---
func TestRestoreCmd_Structure(t *testing.T) {
if restoreCmd == nil {
t.Fatal("restoreCmd is nil")
}
if restoreCmd.Use != "restore" {
t.Fatalf("expected Use='restore', got %q", restoreCmd.Use)
}
requiredFlags := []string{"in", "privkey-pq", "privkey-classical"}
for _, f := range requiredFlags {
if restoreCmd.Flag(f) == nil {
t.Fatalf("missing required --%s flag", f)
}
}
if restoreCmd.Flag("out") == nil {
t.Fatal("missing --out flag")
}
if restoreCmd.Flag("out").DefValue != "" {
t.Fatalf(
"expected --out default empty, got %q",
restoreCmd.Flag("out").DefValue,
)
}
}
func TestRestoreCmd_RequiredFlags(t *testing.T) {
required := []string{"in", "privkey-pq", "privkey-classical"}
for _, name := range required {
flag := restoreCmd.Flag(name)
if flag == nil {
t.Fatalf("missing required --%s flag", name)
}
ann, ok := flag.Annotations[cobra.BashCompOneRequiredFlag]
if !ok || len(ann) == 0 || ann[0] != "true" {
t.Fatalf("flag --%s is not marked required", name)
}
}
}
func TestRestoreCmd_PreRunE_MissingPrivkeyClassical(t *testing.T) {
restoreTestGlobals(t)
dir := t.TempDir()
inFile := filepath.Join(dir, "backup.pqenc")
_ = os.WriteFile(inFile, []byte("data"), 0o644)
pqFile := filepath.Join(dir, "pq.priv.pem")
_ = os.WriteFile(pqFile, []byte("pq"), 0o600)
classicalFile := filepath.Join(dir, "classical.priv.pem") // does not exist
openCount := 0
restoreOsOpen = func(name string) (*os.File, error) {
openCount++
return os.Open(name)
}
cmd := makeRestoreTestCmd()
setRestoreFlags(cmd, inFile, pqFile, classicalFile)
err := cmd.Execute()
if err == nil {
t.Fatal("expected non-nil error")
}
if !strings.Contains(
err.Error(),
"--privkey-pq and --privkey-classical are both required (AND model)",
) {
t.Fatalf("expected AND model error, got: %v", err)
}
if openCount != 0 {
t.Fatalf("expected no os.Open calls on .pqenc, got %d", openCount)
}
}
func TestRestoreCmd_PreRunE_MissingPrivkeyPQ(t *testing.T) {
restoreTestGlobals(t)
dir := t.TempDir()
inFile := filepath.Join(dir, "backup.pqenc")
_ = os.WriteFile(inFile, []byte("data"), 0o644)
pqFile := filepath.Join(dir, "pq.priv.pem") // does not exist
classicalFile := filepath.Join(dir, "classical.priv.pem")
_ = os.WriteFile(classicalFile, []byte("classical"), 0o600)
openCount := 0
restoreOsOpen = func(name string) (*os.File, error) {
openCount++
return os.Open(name)
}
cmd := makeRestoreTestCmd()
setRestoreFlags(cmd, inFile, pqFile, classicalFile)
err := cmd.Execute()
if err == nil {
t.Fatal("expected non-nil error")
}
if !strings.Contains(
err.Error(),
"--privkey-pq and --privkey-classical are both required (AND model)",
) {
t.Fatalf("expected AND model error, got: %v", err)
}
if openCount != 0 {
t.Fatalf("expected no os.Open calls on .pqenc, got %d", openCount)
}
}
func TestRestoreCmd_OnlyPrivkeyPQ(t *testing.T) {
restoreTestGlobals(t)
dir := t.TempDir()
inFile := filepath.Join(dir, "backup.pqenc")
_ = os.WriteFile(inFile, []byte("data"), 0o644)
pqFile := filepath.Join(dir, "pq.priv.pem")
_ = os.WriteFile(pqFile, []byte("pq"), 0o600)
openCount := 0
restoreOsOpen = func(name string) (*os.File, error) {
openCount++
return os.Open(name)
}
cmd := makeRestoreTestCmd()
_ = cmd.Flags().Set("in", inFile)
_ = cmd.Flags().Set("privkey-pq", pqFile)
// --privkey-classical intentionally omitted
err := cmd.Execute()
if err == nil {
t.Fatal("expected non-nil error")
}
if openCount != 0 {
t.Fatalf("expected no os.Open calls on .pqenc, got %d", openCount)
}
}
func TestRestoreCmd_HappyPath_Stdout(t *testing.T) {
restoreTestGlobals(t)
dir := t.TempDir()
inFile := filepath.Join(dir, "backup.pqenc")
_ = os.WriteFile(inFile, []byte("encrypted"), 0o644)
pqFile := filepath.Join(dir, "pq.priv.pem")
_ = os.WriteFile(pqFile, []byte("pq"), 0o600)
classicalFile := filepath.Join(dir, "classical.priv.pem")
_ = os.WriteFile(classicalFile, []byte("classical"), 0o600)
pqPriv := &mockRecipientPriv{
schemeID: 0x0006,
keyID: []byte{1, 2, 3, 4, 5, 6, 7, 8},
raw: make([]byte, 32),
}
classicalPriv := &mockRecipientPriv{
schemeID: 0x0007,
keyID: []byte{8, 7, 6, 5, 4, 3, 2, 1},
raw: make([]byte, 32),
}
mockKM := &mockRestoreKeyManager{pqPriv: pqPriv, classicalPriv: classicalPriv}
newRestoreKeyManager = func(crypto.Registry) crypto.KeyManager {
return mockKM
}
mockDec := &mockDecryptor{}
newRestoreDecryptor = func(crypto.Registry) crypto.Decryptor {
return mockDec
}
var outBuf bytes.Buffer
restoreOutput = &outBuf
cmd := makeRestoreTestCmd()
setRestoreFlags(cmd, inFile, pqFile, classicalFile)
err := cmd.Execute()
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(mockKM.loadPrivCalls) != 2 {
t.Fatalf("expected 2 LoadPriv calls, got %d", len(mockKM.loadPrivCalls))
}
if mockKM.loadPrivCalls[0].path != pqFile ||
mockKM.loadPrivCalls[0].schemeID != 0x0006 {
t.Fatalf("unexpected PQ LoadPriv call: %+v", mockKM.loadPrivCalls[0])
}
if mockKM.loadPrivCalls[1].path != classicalFile ||
mockKM.loadPrivCalls[1].schemeID != 0x0007 {
t.Fatalf(
"unexpected classical LoadPriv call: %+v",
mockKM.loadPrivCalls[1],
)
}
if !mockDec.called {
t.Fatal("expected decryptor.Decrypt to be called")
}
if len(mockDec.privs) != 2 {
t.Fatalf("expected 2 privs, got %d", len(mockDec.privs))
}
if mockDec.privs[0] != pqPriv {
t.Fatal("expected pqPriv in positional slot 0")
}
if mockDec.privs[1] != classicalPriv {
t.Fatal("expected classicalPriv in positional slot 1")
}
if !bytes.Equal(outBuf.Bytes(), []byte("decrypted pg_dump data")) {
t.Fatalf("unexpected stdout content: %q", outBuf.Bytes())
}
}
func TestRestoreCmd_HappyPath_FileOut(t *testing.T) {
restoreTestGlobals(t)
dir := t.TempDir()
inFile := filepath.Join(dir, "backup.pqenc")
_ = os.WriteFile(inFile, []byte("encrypted"), 0o644)
pqFile := filepath.Join(dir, "pq.priv.pem")
_ = os.WriteFile(pqFile, []byte("pq"), 0o600)
classicalFile := filepath.Join(dir, "classical.priv.pem")
_ = os.WriteFile(classicalFile, []byte("classical"), 0o600)
outFile := filepath.Join(dir, "restored.dump")
pqPriv := &mockRecipientPriv{
schemeID: 0x0006,
keyID: []byte{1, 2, 3, 4, 5, 6, 7, 8},
raw: make([]byte, 32),
}
classicalPriv := &mockRecipientPriv{
schemeID: 0x0007,
keyID: []byte{8, 7, 6, 5, 4, 3, 2, 1},
raw: make([]byte, 32),
}
mockKM := &mockRestoreKeyManager{pqPriv: pqPriv, classicalPriv: classicalPriv}
newRestoreKeyManager = func(crypto.Registry) crypto.KeyManager {
return mockKM
}
mockDec := &mockDecryptor{}
newRestoreDecryptor = func(crypto.Registry) crypto.Decryptor {
return mockDec
}
cmd := makeRestoreTestCmd()
setRestoreFlags(cmd, inFile, pqFile, classicalFile)
_ = cmd.Flags().Set("out", outFile)
err := cmd.Execute()
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
data, err := os.ReadFile(outFile)
if err != nil {
t.Fatalf("read output file: %v", err)
}
if !bytes.Equal(data, []byte("decrypted pg_dump data")) {
t.Fatalf("unexpected output file content: %q", data)
}
}
func TestRestoreCmd_TamperingDetected(t *testing.T) {
restoreTestGlobals(t)
dir := t.TempDir()
inFile := filepath.Join(dir, "backup.pqenc")
_ = os.WriteFile(inFile, []byte("encrypted"), 0o644)
pqFile := filepath.Join(dir, "pq.priv.pem")
_ = os.WriteFile(pqFile, []byte("pq"), 0o600)
classicalFile := filepath.Join(dir, "classical.priv.pem")
_ = os.WriteFile(classicalFile, []byte("classical"), 0o600)
pqPriv := &mockRecipientPriv{
schemeID: 0x0006,
keyID: []byte{1, 2, 3, 4, 5, 6, 7, 8},
raw: make([]byte, 32),
}
classicalPriv := &mockRecipientPriv{
schemeID: 0x0007,
keyID: []byte{8, 7, 6, 5, 4, 3, 2, 1},
raw: make([]byte, 32),
}
mockKM := &mockRestoreKeyManager{pqPriv: pqPriv, classicalPriv: classicalPriv}
newRestoreKeyManager = func(crypto.Registry) crypto.KeyManager {
return mockKM
}
mockDec := &mockDecryptor{err: composite.ErrTamperingDetected}
newRestoreDecryptor = func(crypto.Registry) crypto.Decryptor {
return mockDec
}
cmd := makeRestoreTestCmd()
setRestoreFlags(cmd, inFile, pqFile, classicalFile)
err := cmd.Execute()
if err == nil {
t.Fatal("expected non-nil error")
}
if !errors.Is(err, composite.ErrTamperingDetected) {
t.Fatalf("expected ErrTamperingDetected, got: %v", err)
}
}
func TestRestoreCmd_WrongKeys(t *testing.T) {
restoreTestGlobals(t)
dir := t.TempDir()
inFile := filepath.Join(dir, "backup.pqenc")
_ = os.WriteFile(inFile, []byte("encrypted"), 0o644)
pqFile := filepath.Join(dir, "pq.priv.pem")
_ = os.WriteFile(pqFile, []byte("pq"), 0o600)
classicalFile := filepath.Join(dir, "classical.priv.pem")
_ = os.WriteFile(classicalFile, []byte("classical"), 0o600)
pqPriv := &mockRecipientPriv{
schemeID: 0x0006,
keyID: []byte{1, 2, 3, 4, 5, 6, 7, 8},
raw: make([]byte, 32),
}
classicalPriv := &mockRecipientPriv{
schemeID: 0x0007,
keyID: []byte{8, 7, 6, 5, 4, 3, 2, 1},
raw: make([]byte, 32),
}
mockKM := &mockRestoreKeyManager{pqPriv: pqPriv, classicalPriv: classicalPriv}
newRestoreKeyManager = func(crypto.Registry) crypto.KeyManager {
return mockKM
}
mockDec := &mockDecryptor{err: composite.ErrWrongKeys}
newRestoreDecryptor = func(crypto.Registry) crypto.Decryptor {
return mockDec
}
cmd := makeRestoreTestCmd()
setRestoreFlags(cmd, inFile, pqFile, classicalFile)
err := cmd.Execute()
if err == nil {
t.Fatal("expected non-nil error")
}
if !errors.Is(err, composite.ErrWrongKeys) {
t.Fatalf("expected ErrWrongKeys, got: %v", err)
}
}
+114
View File
@@ -0,0 +1,114 @@
package main
import (
"context"
"errors"
"log/slog"
"net/http"
"os"
"os/signal"
"syscall"
"time"
"github.com/spf13/cobra"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/config"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/healthz"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/scheduler/cron"
"git.tswf.io/infra/go-synapse-backupper/pkg/domain"
"git.tswf.io/infra/go-synapse-backupper/pkg/domain/backup"
)
var (
newScheduler = func(
expr string,
job func(),
) (
domain.Scheduler,
error,
) {
return cron.NewCronScheduler(expr, job)
}
newHealthz = func(port int) (*healthz.Server, error) {
return healthz.New(port)
}
runOnceFunc = backup.RunOnce
osExitFunc = func(code int) { os.Exit(code) }
)
var runCmd = &cobra.Command{
Use: "run",
Short: "Start the backup scheduler",
RunE: run,
}
func init() {
config.RegisterFlags(runCmd)
}
func run(cmd *cobra.Command, args []string) error {
cfg, err := config.Load(cmd)
if err != nil {
return err
}
return runWithConfig(cmd.Context(), cfg)
}
func runWithConfig(
ctx context.Context,
cfg *config.Config,
) error {
job := func() {
jobCtx := context.Background()
if err := runOnceFunc(jobCtx, cfg); err != nil {
slog.Error("backup job failed", "error", err)
}
}
scheduler, err := newScheduler(cfg.Backup.Cron, job)
if err != nil {
return err
}
healthzServer, err := newHealthz(cfg.Healthz.Port)
if err != nil {
slog.Error("failed to bind healthz server", "error", err)
osExitFunc(1)
return nil
}
go func() {
if err := healthzServer.Start(); err != nil && !errors.Is(err, http.ErrServerClosed) {
slog.Error("healthz server error", "error", err)
}
}()
scheduler.Start()
sigCtx, stop := signal.NotifyContext(ctx, syscall.SIGTERM, syscall.SIGINT)
defer stop()
<-sigCtx.Done()
shutdownTimeout := cfg.ShutdownTimeout
if shutdownTimeout <= 0 {
shutdownTimeout = 30 * time.Second
}
shutdownCtx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
defer cancel()
_ = healthzServer.Stop(shutdownCtx)
stoppedCtx := scheduler.Stop()
select {
case <-stoppedCtx.Done():
slog.Info("graceful shutdown complete")
case <-time.After(shutdownTimeout):
slog.Warn("forcing exit, backup job still running")
osExitFunc(0)
}
return nil
}
+393
View File
@@ -0,0 +1,393 @@
package main
import (
"context"
"errors"
"fmt"
"net/http"
"sync"
"testing"
"time"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/config"
"git.tswf.io/infra/go-synapse-backupper/pkg/adapters/healthz"
"git.tswf.io/infra/go-synapse-backupper/pkg/domain"
)
type fakeScheduler struct {
mu sync.Mutex
startCalled bool
stopCalled bool
stoppedCtx context.Context
stopCancel context.CancelFunc
}
func (f *fakeScheduler) Start() {
f.mu.Lock()
f.startCalled = true
f.mu.Unlock()
}
func (f *fakeScheduler) Stop() context.Context {
f.mu.Lock()
f.stopCalled = true
f.mu.Unlock()
return f.stoppedCtx
}
func (f *fakeScheduler) wasStarted() bool {
f.mu.Lock()
defer f.mu.Unlock()
return f.startCalled
}
func (f *fakeScheduler) wasStopped() bool {
f.mu.Lock()
defer f.mu.Unlock()
return f.stopCalled
}
type exitPanic int
func (e exitPanic) Error() string { return fmt.Sprintf("exit %d", int(e)) }
func TestRunCmd_Use(t *testing.T) {
if runCmd.Use != "run" {
t.Errorf("runCmd.Use = %q, want %q", runCmd.Use, "run")
}
}
func TestRunCmd_SIGTERM_stopsScheduler(t *testing.T) {
origNewScheduler := newScheduler
origOsExit := osExitFunc
origRunOnce := runOnceFunc
defer func() {
newScheduler = origNewScheduler
osExitFunc = origOsExit
runOnceFunc = origRunOnce
}()
fake := &fakeScheduler{}
fake.stoppedCtx, fake.stopCancel = context.WithCancel(context.Background())
newScheduler = func(expr string, job func()) (domain.Scheduler, error) {
return fake, nil
}
runOnceFunc = func(ctx context.Context, cfg *config.Config) error {
return nil
}
var exitCode int
exitPanicked := false
osExitFunc = func(code int) {
exitCode = code
exitPanicked = true
panic(exitPanic(code))
}
cfg := &config.Config{}
cfg.Healthz.Port = 0
cfg.Backup.Cron = "* * * * *"
cfg.ShutdownTimeout = 100 * time.Millisecond
ctx, cancel := context.WithCancel(context.Background())
errCh := make(chan error, 1)
go func() {
defer func() {
if r := recover(); r != nil {
if _, ok := r.(exitPanic); ok {
errCh <- nil
return
}
errCh <- fmt.Errorf("unexpected panic: %v", r)
return
}
}()
errCh <- runWithConfig(ctx, cfg)
}()
time.Sleep(50 * time.Millisecond)
if !fake.wasStarted() {
t.Fatal("scheduler.Start was not called")
}
cancel()
go func() {
time.Sleep(50 * time.Millisecond)
fake.stopCancel()
}()
select {
case err := <-errCh:
if err != nil {
t.Fatalf("runWithConfig returned error: %v", err)
}
case <-time.After(2 * time.Second):
t.Fatal("timeout waiting for runWithConfig to exit")
}
if !fake.wasStopped() {
t.Error("scheduler.Stop was not called")
}
if exitPanicked {
t.Errorf("unexpected os.Exit call with code %d", exitCode)
}
}
func TestRunCmd_SIGTERM_forcesExitAfterTimeout(t *testing.T) {
origNewScheduler := newScheduler
origOsExit := osExitFunc
origNewHealthz := newHealthz
origRunOnce := runOnceFunc
defer func() {
newScheduler = origNewScheduler
osExitFunc = origOsExit
newHealthz = origNewHealthz
runOnceFunc = origRunOnce
}()
fake := &fakeScheduler{}
var cancelFunc context.CancelFunc
fake.stoppedCtx, cancelFunc = context.WithCancel(context.Background())
defer cancelFunc()
newScheduler = func(expr string, job func()) (domain.Scheduler, error) {
return fake, nil
}
newHealthz = func(port int) (*healthz.Server, error) {
return healthz.New(port)
}
runOnceFunc = func(ctx context.Context, cfg *config.Config) error {
return nil
}
var exitCode int
exitPanicked := false
osExitFunc = func(code int) {
exitCode = code
exitPanicked = true
panic(exitPanic(code))
}
cfg := &config.Config{}
cfg.Healthz.Port = 0
cfg.Backup.Cron = "* * * * * *"
cfg.ShutdownTimeout = 100 * time.Millisecond
ctx, cancel := context.WithCancel(context.Background())
errCh := make(chan error, 1)
go func() {
defer func() {
if r := recover(); r != nil {
if _, ok := r.(exitPanic); ok {
errCh <- nil
return
}
errCh <- fmt.Errorf("unexpected panic: %v", r)
return
}
}()
errCh <- runWithConfig(ctx, cfg)
}()
time.Sleep(50 * time.Millisecond)
cancel()
select {
case err := <-errCh:
if err != nil {
t.Fatalf("runWithConfig returned error: %v", err)
}
case <-time.After(2 * time.Second):
t.Fatal("timeout waiting for runWithConfig to exit")
}
if !fake.wasStopped() {
t.Error("scheduler.Stop was not called")
}
if !exitPanicked {
t.Fatal("expected os.Exit to be called")
}
if exitCode != 0 {
t.Errorf("os.Exit code = %d, want 0", exitCode)
}
}
func TestRunCmd_Healthz503DuringShutdown(t *testing.T) {
origNewScheduler := newScheduler
origOsExit := osExitFunc
origNewHealthz := newHealthz
origRunOnce := runOnceFunc
defer func() {
newScheduler = origNewScheduler
osExitFunc = origOsExit
newHealthz = origNewHealthz
runOnceFunc = origRunOnce
}()
fake := &fakeScheduler{}
fake.stoppedCtx, fake.stopCancel = context.WithCancel(context.Background())
newScheduler = func(expr string, job func()) (domain.Scheduler, error) {
return fake, nil
}
srv, err := healthz.New(0)
if err != nil {
t.Fatalf("failed to create healthz server: %v", err)
}
newHealthz = func(port int) (*healthz.Server, error) {
return srv, nil
}
runOnceFunc = func(ctx context.Context, cfg *config.Config) error {
return nil
}
osExitFunc = func(code int) { panic(exitPanic(code)) }
cfg := &config.Config{}
cfg.Healthz.Port = 0
cfg.Backup.Cron = "* * * * * *"
cfg.ShutdownTimeout = 5 * time.Second
ctx, cancel := context.WithCancel(context.Background())
errCh := make(chan error, 1)
go func() {
defer func() {
if r := recover(); r != nil {
if _, ok := r.(exitPanic); ok {
errCh <- nil
return
}
errCh <- fmt.Errorf("unexpected panic: %v", r)
return
}
}()
errCh <- runWithConfig(ctx, cfg)
}()
time.Sleep(100 * time.Millisecond)
addr := srv.Addr()
resp, err := http.Get("http://" + addr + "/healthz")
if err != nil {
t.Fatalf("healthz request failed: %v", err)
}
resp.Body.Close()
if resp.StatusCode != http.StatusOK {
t.Fatalf("healthz status before shutdown = %d, want %d", resp.StatusCode, http.StatusOK)
}
cancel()
var status503 bool
for i := 0; i < 20; i++ {
resp, err = http.Get("http://" + addr + "/healthz")
if err == nil {
resp.Body.Close()
if resp.StatusCode == http.StatusServiceUnavailable {
status503 = true
break
}
}
time.Sleep(10 * time.Millisecond)
}
if !status503 {
t.Error("did not observe /healthz returning 503 during shutdown")
}
fake.stopCancel()
select {
case err := <-errCh:
if err != nil {
t.Fatalf("runWithConfig returned error: %v", err)
}
case <-time.After(2 * time.Second):
t.Fatal("timeout waiting for runWithConfig to exit")
}
}
func TestRunCmd_HealthzBindFailure(t *testing.T) {
origNewScheduler := newScheduler
origNewHealthz := newHealthz
origOsExit := osExitFunc
origRunOnce := runOnceFunc
defer func() {
newScheduler = origNewScheduler
newHealthz = origNewHealthz
osExitFunc = origOsExit
runOnceFunc = origRunOnce
}()
fake := &fakeScheduler{}
var cancelFunc context.CancelFunc
fake.stoppedCtx, cancelFunc = context.WithCancel(context.Background())
defer cancelFunc()
newScheduler = func(expr string, job func()) (domain.Scheduler, error) {
return fake, nil
}
newHealthz = func(port int) (*healthz.Server, error) {
return nil, errors.New("bind failed")
}
runOnceFunc = func(ctx context.Context, cfg *config.Config) error {
return nil
}
var exitCode int
exitPanicked := false
osExitFunc = func(code int) {
exitCode = code
exitPanicked = true
panic(exitPanic(code))
}
cfg := &config.Config{}
cfg.Backup.Cron = "* * * * * *"
errCh := make(chan error, 1)
go func() {
defer func() {
if r := recover(); r != nil {
if _, ok := r.(exitPanic); ok {
errCh <- nil
return
}
errCh <- fmt.Errorf("unexpected panic: %v", r)
return
}
}()
errCh <- runWithConfig(context.Background(), cfg)
}()
select {
case err := <-errCh:
if err != nil {
t.Fatalf("runWithConfig returned error: %v", err)
}
case <-time.After(2 * time.Second):
t.Fatal("timeout waiting for runWithConfig to exit")
}
if !exitPanicked {
t.Fatal("expected os.Exit to be called on healthz bind failure")
}
if exitCode != 1 {
t.Errorf("os.Exit code = %d, want 1", exitCode)
}
}