From c5197f59991df7fbe6b921a2a84483d108525d03 Mon Sep 17 00:00:00 2001
From: Ian <ViViDboarder@gmail.com>
Date: Mon, 23 Nov 2020 12:58:26 -0800
Subject: acme_server: fix reload of acme database (#3874)

* acme_server: Refactor database creation apart from authority creation

This is a WIP commit that doesn't really offer anything other than
setting us up for using a UsagePool to gracefully reload acme_server
configs.

* Implement UsagePool

* Remove unused context

* Fix initializing non-ACME CA

This will handle cases where a DB is not provided

* Sanitize acme db path and clean debug logs

* Move regex to package level to prevent recompiling
---
 modules/caddypki/acmeserver/acmeserver.go | 86 ++++++++++++++++++++++---------
 1 file changed, 63 insertions(+), 23 deletions(-)

(limited to 'modules/caddypki/acmeserver')

diff --git a/modules/caddypki/acmeserver/acmeserver.go b/modules/caddypki/acmeserver/acmeserver.go
index 5c9f74b..d5e5559 100644
--- a/modules/caddypki/acmeserver/acmeserver.go
+++ b/modules/caddypki/acmeserver/acmeserver.go
@@ -19,6 +19,7 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"regexp"
 	"strings"
 	"time"
 
@@ -66,6 +67,7 @@ type Handler struct {
 	PathPrefix string `json:"path_prefix,omitempty"`
 
 	acmeEndpoints http.Handler
+	logger        *zap.Logger
 }
 
 // CaddyModule returns the Caddy module information.
@@ -78,7 +80,7 @@ func (Handler) CaddyModule() caddy.ModuleInfo {
 
 // Provision sets up the ACME server handler.
 func (ash *Handler) Provision(ctx caddy.Context) error {
-	logger := ctx.Logger(ash)
+	ash.logger = ctx.Logger(ash)
 	// set some defaults
 	if ash.CA == "" {
 		ash.CA = caddypki.DefaultCAID
@@ -101,25 +103,9 @@ func (ash *Handler) Provision(ctx caddy.Context) error {
 		return fmt.Errorf("no certificate authority configured with id: %s", ash.CA)
 	}
 
-	dbFolder := filepath.Join(caddy.AppDataDir(), "acme_server")
-	dbPath := filepath.Join(dbFolder, "db")
-
-	// TODO: See https://github.com/smallstep/nosql/issues/7
-	err = os.MkdirAll(dbFolder, 0755)
+	database, err := ash.openDatabase()
 	if err != nil {
-		return fmt.Errorf("making folder for ACME server database: %v", err)
-	}
-
-	// Check to see if previous db exists
-	var stat os.FileInfo
-	stat, err = os.Stat(dbPath)
-	if stat != nil && err == nil {
-		// A badger db is found and should be removed
-		if stat.IsDir() {
-			logger.Warn("Found an old badger database and removing it",
-				zap.String("path", dbPath))
-			_ = os.RemoveAll(dbPath)
-		}
+		return err
 	}
 
 	authorityConfig := caddypki.AuthorityConfig{
@@ -136,10 +122,7 @@ func (ash *Handler) Provision(ctx caddy.Context) error {
 				},
 			},
 		},
-		DB: &db.Config{
-			Type:       "bbolt",
-			DataSource: dbPath,
-		},
+		DB: database,
 	}
 
 	auth, err := ca.NewAuthority(authorityConfig)
@@ -175,11 +158,68 @@ func (ash Handler) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyh
 	return next.ServeHTTP(w, r)
 }
 
+func (ash Handler) getDatabaseKey() string {
+	key := ash.CA
+	key = strings.ToLower(key)
+	key = strings.TrimSpace(key)
+	return keyCleaner.ReplaceAllLiteralString(key, "")
+}
+
+// Cleanup implements caddy.CleanerUpper and closes any idle databases.
+func (ash Handler) Cleanup() error {
+	key := ash.getDatabaseKey()
+	deleted, err := databasePool.Delete(key)
+	if deleted {
+		ash.logger.Debug("unloading unused CA database", zap.String("db_key", key))
+	}
+	if err != nil {
+		ash.logger.Error("closing CA database", zap.String("db_key", key), zap.Error(err))
+	}
+	return err
+}
+
+func (ash Handler) openDatabase() (*db.AuthDB, error) {
+	key := ash.getDatabaseKey()
+	database, loaded, err := databasePool.LoadOrNew(key, func() (caddy.Destructor, error) {
+		dbFolder := filepath.Join(caddy.AppDataDir(), "acme_server", key)
+		dbPath := filepath.Join(dbFolder, "db")
+
+		err := os.MkdirAll(dbFolder, 0755)
+		if err != nil {
+			return nil, fmt.Errorf("making folder for CA database: %v", err)
+		}
+
+		dbConfig := &db.Config{
+			Type:       "bbolt",
+			DataSource: dbPath,
+		}
+		database, err := db.New(dbConfig)
+		return databaseCloser{&database}, err
+	})
+
+	if loaded {
+		ash.logger.Debug("loaded preexisting CA database", zap.String("db_key", key))
+	}
+
+	return database.(databaseCloser).DB, err
+}
+
 const (
 	defaultHost       = "localhost"
 	defaultPathPrefix = "/acme/"
 )
 
+var keyCleaner = regexp.MustCompile(`[^\w.-_]`)
+var databasePool = caddy.NewUsagePool()
+
+type databaseCloser struct {
+	DB *db.AuthDB
+}
+
+func (closer databaseCloser) Destruct() error {
+	return (*closer.DB).Shutdown()
+}
+
 // Interface guards
 var (
 	_ caddyhttp.MiddlewareHandler = (*Handler)(nil)
-- 
cgit v1.2.3