Add support for external domain sql table in dkim modifier

internal/modify/dkim/dkim.go

@@ -34,6 +34,7 @@ import (
 	"github.com/foxcpp/maddy/framework/address"
 	"github.com/foxcpp/maddy/framework/buffer"
 	"github.com/foxcpp/maddy/framework/config"
+	modconfig "github.com/foxcpp/maddy/framework/config/module"
 	"github.com/foxcpp/maddy/framework/dns"
 	"github.com/foxcpp/maddy/framework/exterrors"
 	"github.com/foxcpp/maddy/framework/log"
@@ -93,23 +94,40 @@ var (
 	}
 )
 
-type Modifier struct {
-	instName string
-
-	domains        []string
-	selector       string
-	signers        map[string]crypto.Signer
-	oversignHeader []string
-	signHeader     []string
-	headerCanon    dkim.Canonicalization
-	bodyCanon      dkim.Canonicalization
-	sigExpiry      time.Duration
-	hash           crypto.Hash
-	multipleFromOk bool
-	signSubdomains bool
-
-	log log.Logger
-}
+type (
+	Modifier struct {
+		instName string
+
+		domains         []string
+		selector        string
+		signers         map[string]crypto.Signer
+		oversignHeader  []string
+		signHeader      []string
+		headerCanon     dkim.Canonicalization
+		bodyCanon       dkim.Canonicalization
+		sigExpiry       time.Duration
+		hash            crypto.Hash
+		multipleFromOk  bool
+		signSubdomains  bool
+		keyPathTemplate string
+		hashName        string
+		newKeyAlgo      string
+		table           module.MutableTable
+		storeKeysInDB   bool
+
+		log log.Logger
+	}
+
+	DKIM struct {
+		Domain     string        `json:"domain"`
+		PrivateKey string        `json:"privateKey,omitempty"`
+		PublicKey  string        `json:"publicKey,omitempty"`
+		DNSName    string        `json:"dnsName"`
+		DNSValue   string        `json:"dnsValue"`
+		Expires    time.Time     `json:"expires,omitempty"`
+		pkey       crypto.Signer `json:"-"`
+	}
+)
 
 func New(_, instName string, _, inlineArgs []string) (module.Module, error) {
 	m := &Modifier{
@@ -140,16 +158,13 @@ func (m *Modifier) InstanceName() string {
 }
 
 func (m *Modifier) Init(cfg *config.Map) error {
-	var (
-		hashName        string
-		keyPathTemplate string
-		newKeyAlgo      string
-	)
 
 	cfg.Bool("debug", true, false, &m.log.Debug)
+	cfg.Bool("store_keys_in_database", false, false, &m.storeKeysInDB)
 	cfg.StringList("domains", false, false, m.domains, &m.domains)
 	cfg.String("selector", false, false, m.selector, &m.selector)
-	cfg.String("key_path", false, false, "dkim_keys/{domain}_{selector}.key", &keyPathTemplate)
+	cfg.Custom("domain_table", true, false, nil, modconfig.TableDirective, &m.table)
+	cfg.String("key_path", false, false, "dkim_keys/{domain}_{selector}.key", &m.keyPathTemplate)
 	cfg.StringList("oversign_fields", false, false, oversignDefault, &m.oversignHeader)
 	cfg.StringList("sign_fields", false, false, signDefault, &m.signHeader)
 	cfg.Enum("header_canon", false, false,
@@ -160,9 +175,9 @@ func (m *Modifier) Init(cfg *config.Map) error {
 		dkim.CanonicalizationRelaxed, (*string)(&m.bodyCanon))
 	cfg.Duration("sig_expiry", false, false, 5*Day, &m.sigExpiry)
 	cfg.Enum("hash", false, false,
-		[]string{"sha256"}, "sha256", &hashName)
+		[]string{"sha256"}, "sha256", &m.hashName)
 	cfg.Enum("newkey_algo", false, false,
-		[]string{"rsa4096", "rsa2048", "ed25519"}, "rsa2048", &newKeyAlgo)
+		[]string{"rsa4096", "rsa2048", "ed25519"}, "rsa2048", &m.newKeyAlgo)
 	cfg.Bool("allow_multiple_from", false, false, &m.multipleFromOk)
 	cfg.Bool("sign_subdomains", false, false, &m.signSubdomains)
 
@@ -180,20 +195,34 @@ func (m *Modifier) Init(cfg *config.Map) error {
 		return errors.New("sign_domain: only one domain is supported when sign_subdomains is enabled")
 	}
 
-	m.hash = hashFuncs[hashName]
+	m.hash = hashFuncs[m.hashName]
 	if m.hash == 0 {
 		panic("modify.dkim.Init: Hash function allowed by config matcher but not present in hashFuncs")
 	}
 
+	// If available, include domains from SQL table
+	if m.table != nil {
+		domains, err := m.table.Keys()
+		if err != nil {
+			return err
+		}
+
+		if len(domains) > 0 {
+			m.domains = append(m.domains, domains...)
+		}
+	}
+
+	storeKeysInDB := m.storeKeysInDB && m.table != nil
+
 	for _, domain := range m.domains {
 		if _, err := idna.ToASCII(domain); err != nil {
 			m.log.Printf("warning: unable to convert domain %s to A-labels form, non-EAI messages will not be signed: %v", domain, err)
 		}
 
 		keyValues := strings.NewReplacer("{domain}", domain, "{selector}", m.selector)
-		keyPath := keyValues.Replace(keyPathTemplate)
+		keyPath := keyValues.Replace(m.keyPathTemplate)
 
-		signer, newKey, err := m.loadOrGenerateKey(keyPath, newKeyAlgo)
+		signer, newKey, err := m.loadOrGenerateKey(domain, keyPath, m.newKeyAlgo, storeKeysInDB)
 		if err != nil {
 			return err
 		}
@@ -205,7 +234,7 @@ func (m *Modifier) Init(cfg *config.Map) error {
 			}
 			m.log.Printf("generated a new %s keypair, private key is in %s, TXT record with public key is in %s,\n"+
 				"put its contents into TXT record for %s._domainkey.%s to make signing and verification work",
-				newKeyAlgo, keyPath, dnsPath, m.selector, domain)
+				m.newKeyAlgo, keyPath, dnsPath, m.selector, domain)
 		}
 
 		normDomain, err := dns.ForLookup(domain)
@@ -305,8 +334,17 @@ func (s *state) RewriteBody(ctx context.Context, h *textproto.Header, body buffe
 	}
 	keySigner := s.m.signers[normDomain]
 	if keySigner == nil {
-		s.log.Msg("no key for domain", "domain", normDomain)
-		return nil
+		if s.m.table == nil {
+			s.log.Msg("no key for domain", "domain", normDomain)
+			return nil
+		}
+		keySigner, err = s.m.generateKeyForDomain(normDomain)
+		if err != nil {
+			s.log.Msg("no key for domain", "domain", normDomain)
+			return err
+		}
+		s.m.signers[normDomain] = keySigner
+		s.m.domains = append(s.m.domains, domain)
 	}
 
 	// If the message is non-EAI, we are not allowed to use domains in U-labels,

internal/modify/dkim/keys.go

@@ -19,39 +19,99 @@ along with this program.  If not, see <https://www.gnu.org/licenses/>.
 package dkim
 
 import (
+	"context"
 	"crypto"
 	"crypto/ecdsa"
 	"crypto/ed25519"
 	"crypto/rand"
 	"crypto/rsa"
 	"crypto/x509"
 	"encoding/base64"
+	"encoding/json"
 	"encoding/pem"
 	"fmt"
 	"io"
 	"os"
 	"path/filepath"
+	"strings"
+	"time"
+
+	"github.com/foxcpp/maddy/framework/dns"
+	"golang.org/x/net/idna"
 )
 
-func (m *Modifier) loadOrGenerateKey(keyPath, newKeyAlgo string) (pkey crypto.Signer, newKey bool, err error) {
-	f, err := os.Open(keyPath)
+func (m *Modifier) generateKeyForDomain(domain string) (crypto.Signer, error) {
+	if _, err := idna.ToASCII(domain); err != nil {
+		m.log.Printf("warning: unable to convert domain %s to A-labels form, non-EAI messages will not be signed: %v", domain, err)
+	}
+
+	keyValues := strings.NewReplacer("{domain}", domain, "{selector}", m.selector)
+	keyPath := keyValues.Replace(m.keyPathTemplate)
+
+	storeInDB := m.storeKeysInDB && m.table != nil
+	signer, newKey, err := m.loadOrGenerateKey(domain, keyPath, m.newKeyAlgo, storeInDB)
 	if err != nil {
-		if os.IsNotExist(err) {
-			pkey, err = m.generateAndWrite(keyPath, newKeyAlgo)
-			return pkey, true, err
+		return nil, err
+	}
+
+	if newKey {
+		dnsPath := keyPath + ".dns"
+		if filepath.Ext(keyPath) == ".key" {
+			dnsPath = keyPath[:len(keyPath)-4] + ".dns"
 		}
-		return nil, false, err
+		m.log.Printf("generated a new %s keypair, private key is in %s, TXT record with public key is in %s,\n"+
+			"put its contents into TXT record for %s._domainkey.%s to make signing and verification work",
+			m.newKeyAlgo, keyPath, dnsPath, m.selector, domain)
 	}
-	defer f.Close()
 
-	pemBlob, err := io.ReadAll(f)
+	normDomain, err := dns.ForLookup(domain)
 	if err != nil {
-		return nil, false, err
+		return nil, fmt.Errorf("sign_skim: unable to normalize domain %s: %w", domain, err)
+	}
+	m.signers[normDomain] = signer
+	return signer, nil
+}
+func (m *Modifier) loadOrGenerateKey(domain, keyPath, newKeyAlgo string, storeInDB bool) (pkey crypto.Signer, newKey bool, err error) {
+	var pemBlob []byte
+	if storeInDB && m.table != nil {
+		ctx := context.Background()
+		keyData, ok, err := m.table.Lookup(ctx, domain)
+		if err != nil {
+			return nil, false, err
+		}
+		if !ok || keyData == "" {
+			pkey, err = m.generateAndWrite(domain, keyPath, newKeyAlgo, storeInDB)
+			return pkey, true, err
+		}
+		var dkimKey DKIM
+		if err = json.Unmarshal([]byte(keyData), &dkimKey); err != nil {
+			return nil, false, err
+		}
+		pemBlob = []byte(dkimKey.PrivateKey)
+	} else {
+		f, err := os.Open(keyPath)
+		if err != nil {
+			if os.IsNotExist(err) {
+				pkey, err = m.generateAndWrite(domain, keyPath, newKeyAlgo, storeInDB)
+				return pkey, true, err
+			}
+			return nil, false, err
+		}
+		defer f.Close()
+
+		pemBlob, err = io.ReadAll(f)
+		if err != nil {
+			return nil, false, err
+		}
 	}
 
 	block, _ := pem.Decode(pemBlob)
 	if block == nil {
-		return nil, false, fmt.Errorf("modify.dkim: %s: invalid PEM block", keyPath)
+		reference := keyPath
+		if storeInDB && m.table != nil {
+			reference = domain
+		}
+		return nil, false, fmt.Errorf("modify.dkim: %s: invalid PEM block", reference)
 	}
 
 	var key interface{}
@@ -91,7 +151,7 @@ func (m *Modifier) loadOrGenerateKey(keyPath, newKeyAlgo string) (pkey crypto.Si
 	}
 }
 
-func (m *Modifier) generateAndWrite(keyPath, newKeyAlgo string) (crypto.Signer, error) {
+func (m *Modifier) generateAndWrite(domain, keyPath, newKeyAlgo string, storeInDB bool) (crypto.Signer, error) {
 	wrapErr := func(err error) error {
 		return fmt.Errorf("modify.dkim: generate %s: %w", keyPath, err)
 	}
@@ -124,6 +184,24 @@ func (m *Modifier) generateAndWrite(keyPath, newKeyAlgo string) (crypto.Signer,
 		return nil, wrapErr(err)
 	}
 
+	selector := fmt.Sprintf("%s._domainkey.%s", m.selector, domain)
+	dkimKey, err := keyToJSON(domain, selector, newKeyAlgo)
+	if err != nil {
+		return nil, wrapErr(err)
+	}
+
+	dkimKey.Expires = time.Now().Add(m.sigExpiry)
+
+	resultString, err := json.Marshal(dkimKey)
+	if err != nil {
+		return nil, wrapErr(err)
+	}
+
+	if storeInDB && m.table != nil {
+		err = m.table.SetKey(domain, string(resultString))
+		return pkey, err
+	}
+
 	// 0777 because we have public keys in here too and they don't
 	// need protection. Individual private key files have 0600 perms.
 	if err := os.MkdirAll(filepath.Dir(keyPath), 0o777); err != nil {
@@ -150,6 +228,69 @@ func (m *Modifier) generateAndWrite(keyPath, newKeyAlgo string) (crypto.Signer,
 	return pkey, nil
 }
 
+func keyToJSON(domain, selector, algo string) (result DKIM, err error) {
+	var (
+		dkimName   = algo
+		pkey       crypto.Signer
+		pubKeyBlob []byte
+	)
+
+	switch algo {
+	case "rsa4096":
+		dkimName = "rsa"
+		pkey, err = rsa.GenerateKey(rand.Reader, 4096)
+	case "rsa2048":
+		dkimName = "rsa"
+		pkey, err = rsa.GenerateKey(rand.Reader, 2048)
+	case "ed25519":
+		_, pkey, err = ed25519.GenerateKey(rand.Reader)
+	default:
+		err = fmt.Errorf("unknown key algorithm: %s", algo)
+	}
+	if err != nil {
+		return DKIM{}, err
+	}
+
+	keyBlob, err := x509.MarshalPKCS8PrivateKey(pkey)
+	if err != nil {
+		return DKIM{}, err
+	}
+
+	pubkey := pkey.Public()
+	switch pubkey := pubkey.(type) {
+	case *rsa.PublicKey:
+		var err error
+		pubKeyBlob, err = x509.MarshalPKIXPublicKey(pubkey)
+		if err != nil {
+			return DKIM{}, err
+		}
+	case ed25519.PublicKey:
+		pubKeyBlob = pubkey
+	default:
+		panic("modify.dkim.writeDNSRecord: unknown key algorithm")
+	}
+
+	pubKeyString := base64.StdEncoding.EncodeToString(pubKeyBlob)
+	keyRecord := fmt.Sprintf("v=DKIM1; k=%s; p=%s", dkimName, pubKeyString)
+
+	keyBytes := pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: keyBlob})
+	if keyBytes == nil {
+		err := fmt.Errorf("failed to encode private key")
+		return DKIM{}, err
+	}
+
+	result = DKIM{
+		DNSValue:   keyRecord,
+		PrivateKey: string(keyBytes),
+		PublicKey:  pubKeyString,
+		pkey:       pkey,
+		Domain:     domain,
+		DNSName:    selector,
+	}
+
+	return result, nil
+}
+
 func writeDNSRecord(keyPath, dkimAlgoName string, pkey crypto.Signer) (string, error) {
 	var (
 		keyBlob []byte