package main import ( "bytes" "context" "encoding/json" "log" "net/http" "os" "os/signal" "strings" "syscall" "time" "github.com/aws/aws-sdk-go-v2/aws" awsconf "github.com/aws/aws-sdk-go-v2/config" "github.com/google/uuid" "github.com/warmbly/warmbly/internal/app/cipher" "github.com/warmbly/warmbly/internal/app/worker" "github.com/warmbly/warmbly/internal/config" "github.com/warmbly/warmbly/internal/infrastructure/cache" "github.com/warmbly/warmbly/internal/infrastructure/codec" "github.com/warmbly/warmbly/internal/infrastructure/encryptedkeys" "github.com/warmbly/warmbly/internal/infrastructure/eventbus" "github.com/warmbly/warmbly/internal/infrastructure/kafka" "github.com/warmbly/warmbly/internal/infrastructure/kms" "github.com/warmbly/warmbly/internal/infrastructure/storage" "github.com/warmbly/warmbly/internal/observability" "github.com/warmbly/warmbly/internal/repository" ) func main() { ctx, cancel := context.WithCancel(context.Background()) defer cancel() // Resolve worker identity. Precedence: // 1. WORKER_ID: explicit UUID, used as-is // 2. WORKER_BIND_IP: derive UUIDv5 from the bound egress IP so each // IP on a multi-IP box becomes its own worker // 3. hostname: UUID set by Terraform on legacy single-IP VPS // 4. generated UUID: local dev fallback workerID, bindIP := resolveWorkerID() log.Printf("Worker ID: %s, Bind IP: %s", workerID, bindIP) // Load config with env-first approach cfg, err := config.NewConfig(ctx) if err != nil { log.Fatal(err) } // Sentry if err := observability.InitSentry(ctx, cfg, "worker"); err != nil { log.Fatal(err) } // AWS SDK config, loaded only when an AWS-backed provider is selected // (KMS_PROVIDER=aws or BLOB_PROVIDER=s3). A fully-local self-host needs no // AWS_REGION or credentials. var awscfg aws.Config if config.AWSNeeded() { awscfg, err = awsconf.LoadDefaultConfig(ctx) if err != nil { log.Fatal(err) } } // Redis primaryRedis, err := cfg.LoadPrimaryRedisEndpoint(ctx) if err != nil { log.Fatal(err) } redisCache, err := cache.New(primaryRedis) if err != nil { log.Fatal(err) } // KMS → CipherService var masterKey string = "alias/master-key" if cfg.Env != "prod" { masterKey += "-dev" } kmsClient, err := kms.FromEnv(ctx, awscfg, masterKey) if err != nil { log.Fatal(err) } encryptedKeys, err := encryptedkeys.FromEnv( encryptedkeys.Deps{}, "http", ) if err != nil { log.Fatal(err) } cipherService := cipher.NewService(kmsClient, redisCache, encryptedKeys) // The worker reaches the messageId -> internal email map over the internal // backend API (same base URL + token as the DEK store) rather than touching // Postgres directly, per the worker no-direct-SQL rule in CLAUDE.md. internalBaseURL := strings.TrimRight(os.Getenv("ENCRYPTED_KEYS_BACKEND_URL"), "/") internalToken := os.Getenv("ENCRYPTED_KEYS_WORKER_TOKEN") emailMessageMapRepo, err := repository.NewHTTPEmailMessageMapRepository(internalBaseURL, internalToken) if err != nil { log.Fatal(err) } // Same channel for the sync governor's "is this a reply to our own // mail?" lookup, which decides the priority lane. syncContextRepo, err := repository.NewHTTPSyncContextRepository(internalBaseURL, internalToken) if err != nil { log.Fatal(err) } // Blob storage (S3 by default, filesystem when BLOB_PROVIDER=filesystem). s3Client, err := storage.NewFromEnv(ctx, awscfg, "main") if err != nil { log.Fatal(err) } // Codec (Avro reads SCHEMA_REGISTRY_URL from env; JSON needs nothing). codecImpl, err := codec.FromEnv() if err != nil { log.Fatal(err) } log.Printf("Codec: %s", codecImpl.Name()) // Event bus (NATS default; Kafka when EVENTBUS_PROVIDER=kafka). Kafka // bootstrap/SASL is only loaded for the Kafka provider. var kafkaBootstrapServers string var kafkaSaslConfig *kafka.SASLConfig if config.EventBusProvider() == "kafka" { kafkaBootstrapServers, err = cfg.LoadKafkaBootstrapServers(ctx) if err != nil { log.Fatal(err) } kafkaSaslConfig, err = cfg.LoadKafkaConfigSasl(ctx) if err != nil { log.Fatal(err) } } bus, err := eventbus.FromEnv(kafkaBootstrapServers, kafkaSaslConfig) if err != nil { log.Fatal(err) } defer bus.Close() log.Printf("Event bus: %s", bus.Name()) workerTopic := kafka.GetWorkerTopic(workerID.String()) // Provider OAuth configs for local token refresh. Cfg is not shipped in the // AddWorkerEmail payload (avro-excluded), so the worker rebuilds it from // these (reads BOX_GOOGLE_* / BOX_OUTLOOK_* from the worker env). RedirectURL // is unused for refresh, so the base URL is irrelevant here. oauthInbox := config.LoadOauth2Inbox("") // Token refresh needs the provider client credentials in the worker env. // Warn loudly if they're missing: the initial token still works, but refresh // fails silently once it expires (~1h), stalling the mailbox. if oauthInbox.Outlook == nil || oauthInbox.Outlook.ClientID == "" || oauthInbox.Outlook.ClientSecret == "" { log.Println("WARNING: BOX_OUTLOOK_CLIENT_ID/SECRET not set; Microsoft Graph mailbox token refresh will fail on expiry") } if oauthInbox.Google == nil || oauthInbox.Google.ClientID == "" || oauthInbox.Google.ClientSecret == "" { log.Println("WARNING: BOX_GOOGLE_CLIENT_ID/SECRET not set; Gmail mailbox token refresh will fail on expiry") } // WorkerService workerService := &worker.WorkerService{ ID: workerID.String(), CipherService: cipherService, Bus: bus, Codec: codecImpl, Cache: redisCache, Storage: s3Client, EmailMessageMapRepository: emailMessageMapRepo, SyncContextRepository: syncContextRepo, OauthInbox: &oauthInbox, } if err := workerService.Init(); err != nil { log.Fatal("failed to init worker service:", err) } workerService.InitEvents() // Start heartbeat + health sampler. RunHealth ticks every 30s, snapshots // the rolling 1m counters into a WorkerHealth event, publishes via the // event bus so the consumer can write a row into worker_health_samples. go workerService.Heartbeat(ctx) go workerService.RunHealth(ctx, 30*time.Second) // heartbeatDone closes once the farewell beat has been sent. Main waits on // it before returning, otherwise the process exits mid-request and the // worker stays registered as active. heartbeatDone := make(chan struct{}) go func() { defer close(heartbeatDone) runInternalHeartbeat(ctx, workerID, bindIP) }() // Graceful shutdown sigCh := make(chan os.Signal, 1) signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) go func() { <-sigCh log.Println("Shutting down worker", workerID) cancel() }() log.Printf("Worker %s started, listening on topic %s", workerID, workerTopic) if err := bus.Subscribe(ctx, []string{workerTopic}, "worker-"+workerID.String(), workerService.Receive); err != nil { log.Println("event bus subscribe ended:", err) } // Give the farewell beat a moment to land. Bounded so a wedged backend // can't stop the worker from exiting; the heartbeat staleness window // catches it either way. select { case <-heartbeatDone: case <-time.After(8 * time.Second): log.Println("timed out waiting for the shutdown heartbeat") } log.Println("Worker stopped") } func runInternalHeartbeat(ctx context.Context, workerID uuid.UUID, bindIP string) { baseURL := strings.TrimRight(os.Getenv("ENCRYPTED_KEYS_BACKEND_URL"), "/") token := os.Getenv("ENCRYPTED_KEYS_WORKER_TOKEN") if baseURL == "" || token == "" { return } reportedIP := os.Getenv("WORKER_PUBLIC_IP") if reportedIP == "" && bindIP != "default route" { reportedIP = bindIP } if reportedIP == "" { reportedIP = "unknown" } client := &http.Client{Timeout: 10 * time.Second} // reqCtx is separate from ctx so the farewell beat still sends after ctx is // cancelled by the shutdown signal. send := func(reqCtx context.Context, booted, stopping bool) { payload := map[string]any{ "worker_id": workerID.String(), "bind_ip": reportedIP, "tier": os.Getenv("WORKER_TIER"), "egress_kind": os.Getenv("WORKER_EGRESS_KIND"), } if booted { // Mailboxes live in memory only, so a fresh process holds none. // The backend reloads this worker's mailboxes on this beat // instead of leaving them to the reconciler's next pass. payload["booted"] = true } if stopping { payload["stopping"] = true } body, _ := json.Marshal(payload) req, err := http.NewRequestWithContext(reqCtx, http.MethodPost, baseURL+"/api/v1/internal/worker/heartbeat", bytes.NewReader(body)) if err != nil { log.Println("failed to build internal heartbeat:", err) return } req.Header.Set("Authorization", "Bearer "+token) req.Header.Set("Content-Type", "application/json") resp, err := client.Do(req) if err != nil { log.Println("failed internal heartbeat:", err) return } defer resp.Body.Close() if resp.StatusCode >= 300 { log.Println("internal heartbeat returned status", resp.StatusCode) } } send(ctx, true, false) ticker := time.NewTicker(90 * time.Second) defer ticker.Stop() for { select { case <-ctx.Done(): // Farewell beat on a fresh context: ctx is already cancelled, and // without this the row stays selectable until the heartbeat ages // out, so placement keeps picking a worker that has exited. byeCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) send(byeCtx, false, true) cancel() return case <-ticker.C: send(ctx, false, false) } } } // uuidNamespaceURL is the RFC 4122 URL namespace, matching the value used by // scripts/install-worker.sh when deriving the per-IP worker ID. Keep these in // sync: the installer and the worker must agree on the derivation. var uuidNamespaceURL = uuid.MustParse("6ba7b811-9dad-11d1-80b4-00c04fd430c8") // workerIDFromIP returns the deterministic UUIDv5 for the given IPv4 string. // Same IP always maps to the same UUID, which is what lets us treat each IP // on a multi-IP box as its own stable sending identity (egress). func workerIDFromIP(ip string) uuid.UUID { return uuid.NewSHA1(uuidNamespaceURL, []byte(ip)) } // resolveWorkerID applies the boot-time precedence rules and returns the // chosen worker UUID together with a human-readable bind-IP label for logs. func resolveWorkerID() (uuid.UUID, string) { if raw := os.Getenv("WORKER_ID"); raw != "" { id, err := uuid.Parse(raw) if err != nil { log.Fatalf("WORKER_ID %q is not a valid UUID: %v", raw, err) } bind := os.Getenv("WORKER_BIND_IP") if bind == "" { bind = "default route" } return id, bind } if bind := os.Getenv("WORKER_BIND_IP"); bind != "" { return workerIDFromIP(bind), bind } hostname, err := os.Hostname() if err != nil { log.Fatal("failed to get hostname:", err) } if id, err := uuid.Parse(hostname); err == nil { return id, "default route" } id := uuid.New() log.Printf("Hostname %q is not a UUID, using generated ID: %s", hostname, id) return id, "default route" }