Add cache warmer worker

Signed-off-by: Bryan Frimin <bryan@getprobo.com>

pkg/cert/cache_warmer.go

@@ -0,0 +1,168 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package cert
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/getprobo/probo/pkg/coredata"
+	"github.com/getprobo/probo/pkg/crypto/cipher"
+	"go.gearno.de/kit/log"
+	"go.gearno.de/kit/pg"
+)
+
+type (
+	CacheWarmer struct {
+		pg            *pg.Client
+		encryptionKey cipher.EncryptionKey
+		logger        *log.Logger
+	}
+)
+
+func NewCacheWarmer(
+	pg *pg.Client,
+	encryptionKey cipher.EncryptionKey,
+	logger *log.Logger,
+) *CacheWarmer {
+	return &CacheWarmer{
+		pg:            pg,
+		encryptionKey: encryptionKey,
+		logger:        logger.Named("cert.cache_warmer"),
+	}
+}
+
+func (w *CacheWarmer) WarmCache(ctx context.Context) error {
+	w.logger.InfoCtx(ctx, "warming certificate cache")
+	startTime := time.Now()
+
+	err := w.pg.WithConn(
+		ctx,
+		func(conn pg.Conn) error {
+			domains := coredata.CustomDomains{}
+			if err := domains.LoadActiveCertificates(ctx, conn, coredata.NewNoScope()); err != nil {
+				return fmt.Errorf("cannot load active certificates: %w", err)
+			}
+
+			if len(domains) == 0 {
+				w.logger.InfoCtx(ctx, "no active certificates to warm")
+				return nil
+			}
+
+			w.logger.InfoCtx(ctx, "found active certificates to cache", log.Int("count", len(domains)))
+
+			successCount := 0
+			for _, domain := range domains {
+				select {
+				case <-ctx.Done():
+					return ctx.Err()
+				default:
+				}
+
+				if err := w.warmDomain(ctx, conn, domain); err != nil {
+					w.logger.ErrorCtx(ctx, "cannot warm certificate cache for domain", log.String("domain", domain.Domain), log.Error(err))
+				} else {
+					successCount++
+				}
+			}
+
+			w.logger.InfoCtx(ctx, "successfully warmed cache", log.Int("success_count", successCount), log.Int("total_count", len(domains)))
+			return nil
+		},
+	)
+
+	if err != nil {
+		return fmt.Errorf("cannot warm certificate cache: %w", err)
+	}
+
+	w.logger.InfoCtx(ctx, "certificate cache warming completed", log.Duration("duration", time.Since(startTime)))
+
+	return nil
+}
+
+func (w *CacheWarmer) warmDomain(ctx context.Context, conn pg.Conn, domain *coredata.CustomDomain) error {
+	var loadedDomain coredata.CustomDomain
+	scope := coredata.NewScope(domain.OrganizationID.TenantID())
+	if err := loadedDomain.LoadByID(ctx, conn, scope, w.encryptionKey, domain.ID); err != nil {
+		return fmt.Errorf("cannot load domain with decrypted values: %w", err)
+	}
+
+	if loadedDomain.SSLCertificate == nil {
+		return fmt.Errorf("domain has no parsed certificate")
+	}
+
+	if len(loadedDomain.SSLCertificatePEM) == 0 {
+		return fmt.Errorf("domain has no certificate PEM")
+	}
+
+	if len(loadedDomain.SSLPrivateKeyPEM) == 0 {
+		return fmt.Errorf("domain has no private key PEM")
+	}
+
+	if loadedDomain.SSLExpiresAt == nil {
+		return fmt.Errorf("domain certificate has no expiry date")
+	}
+
+	if time.Now().After(*loadedDomain.SSLExpiresAt) {
+		return fmt.Errorf("certificate has expired")
+	}
+
+	cache := &coredata.CachedCertificate{
+		Domain:           loadedDomain.Domain,
+		CertificatePEM:   string(loadedDomain.SSLCertificatePEM),
+		PrivateKeyPEM:    string(loadedDomain.SSLPrivateKeyPEM),
+		CertificateChain: loadedDomain.SSLCertificateChain,
+		ExpiresAt:        *loadedDomain.SSLExpiresAt,
+		CachedAt:         time.Now(),
+		CustomDomainID:   loadedDomain.ID,
+	}
+
+	if err := cache.Upsert(ctx, conn); err != nil {
+		return fmt.Errorf("cannot upsert cache entry: %w", err)
+	}
+
+	return nil
+}
+
+func (w *CacheWarmer) RefreshCache(ctx context.Context) error {
+	w.logger.InfoCtx(ctx, "refreshing certificate cache")
+
+	return w.pg.WithConn(
+		ctx,
+		func(conn pg.Conn) error {
+			cachedCertificates := coredata.CachedCertificates{}
+			if err := cachedCertificates.CleanExpired(ctx, conn); err != nil {
+				w.logger.ErrorCtx(ctx, "cannot clean expired cache", log.Error(err))
+			}
+
+			return w.WarmCache(ctx)
+		},
+	)
+}
+
+func (w *CacheWarmer) WarmSingleDomain(ctx context.Context, domainName string) error {
+	return w.pg.WithConn(
+		ctx,
+		func(conn pg.Conn) error {
+			var domain coredata.CustomDomain
+			if err := domain.LoadByDomain(ctx, conn, coredata.NewNoScope(), w.encryptionKey, domainName); err != nil {
+				return fmt.Errorf("cannot load domain: %w", err)
+			}
+
+			return w.warmDomain(ctx, conn, &domain)
+		},
+	)
+}