package main import ( "context" "log" "math" "os" "os/signal" "strconv" "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/models" "github.com/warmbly/warmbly/internal/observability" "github.com/warmbly/warmbly/internal/pkg/nodeagent" "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) } // Error reporting if err := observability.Init(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) } // Mailboxes managed by Warmbly Cloud send with access tokens the backend // brokers; the refresh grant never reaches the worker. Minting one is a // broker operation, so it carries NODE_BROKER_TOKEN where the instance // issues a separate one and the shared internal token otherwise. brokerToken := os.Getenv("NODE_BROKER_TOKEN") if brokerToken == "" { brokerToken = internalToken } tokenBroker, err := repository.NewHTTPBrokeredTokenClient(internalBaseURL, brokerToken) 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, TokenBroker: tokenBroker, } 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) newNodeAgent(workerID, bindIP).Run(ctx) }() // 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") } // newNodeAgent builds the shared fleet agent. It replaces the worker's own // heartbeat loop: identity, usage reporting and self-update are node concerns, // identical for every role, so they live in one place. func newNodeAgent(workerID uuid.UUID, bindIP string) *nodeagent.Agent { reportedIP := os.Getenv("WORKER_PUBLIC_IP") if reportedIP == "" && bindIP != "default route" { reportedIP = bindIP } return nodeagent.New(nodeagent.Config{ NodeID: workerID, Role: models.NodeRoleWorker, Name: os.Getenv("WARMBLY_NODE_NAME"), Region: nodeRegion(), Address: reportedIP, CapacityTarget: workerCapacityTarget(), Version: buildVersion(), BaseURL: os.Getenv("ENCRYPTED_KEYS_BACKEND_URL"), Token: os.Getenv("ENCRYPTED_KEYS_WORKER_TOKEN"), // Written for the host-side updater installed by `warmbly join`. TargetVersionPath: os.Getenv("WARMBLY_TARGET_VERSION_PATH"), }) } const defaultWorkerCapacityTarget = 100.0 // workerCapacityTarget reads this machine's assigned-mailbox target. func workerCapacityTarget() float64 { raw := strings.TrimSpace(os.Getenv("WARMBLY_WORKER_CAPACITY")) if raw == "" { return defaultWorkerCapacityTarget } n, err := strconv.ParseFloat(raw, 64) if err != nil || n <= 0 || math.IsNaN(n) || math.IsInf(n, 0) || n > 99_999_999.99 { return defaultWorkerCapacityTarget } return n } // nodeRegion reads the sign-in geography hint. WARMBLY_NODE_REGION is what the // join script writes and what every role uses; WORKER_REGION is the older // worker-only name, kept as a fallback so a machine configured by hand before // the join flow existed keeps reporting its region. func nodeRegion() string { if v := os.Getenv("WARMBLY_NODE_REGION"); v != "" { return v } return os.Getenv("WORKER_REGION") } // buildVersion is the image tag this build reports. Set by the join script // from the tag it pulled; empty means unknown, which the control plane must // not read as "needs updating". func buildVersion() string { return os.Getenv("WARMBLY_VERSION") } // uuidNamespaceURL is the RFC 4122 URL namespace used to derive the per-IP // worker ID. Changing it re-identifies every WORKER_BIND_IP node and orphans // the mailboxes assigned to their old ids. 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" } // A persisted id survives container recreates, which a random one does // not: every fresh UUID orphans the mailboxes assigned to the old one. if dir := os.Getenv("WORKER_STATE_DIR"); dir != "" { if id, ok := claimStateID(dir); ok { log.Printf("Using persisted worker ID %s from %s", id, dir) 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" }