reactive-mail-server/backend/src/main/java/com/mailserver/smtp/ReloadingKeyManager.java

package com.mailserver.smtp;

import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
import org.bouncycastle.openssl.PEMKeyPair;
import org.bouncycastle.openssl.PEMParser;
import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.net.ssl.SSLEngine;
import javax.net.ssl.X509ExtendedKeyManager;
import java.io.InputStream;
import java.io.Reader;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.attribute.FileTime;
import java.security.Principal;
import java.security.PrivateKey;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.time.Clock;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Server certificate and key from PEM files, read again when the files change — a Let's Encrypt certificate
 * renewed on disk (lego, certbot) is picked up without a restart, for STARTTLS and implicit TLS alike: the JDK
 * asks the key manager on every handshake. The files are checked at most once per {@code interval}; a file
 * that cannot be read (a renewal half written) keeps the previous certificate and is tried again next time.
 * <p>
 * Key: PKCS#8 ({@code PRIVATE KEY}) or the traditional {@code RSA PRIVATE KEY} / {@code EC PRIVATE KEY} that
 * lego and certbot write. Each loaded pair gets its own alias, so a handshake that asks for the chain and then
 * for the key never mixes two generations.
 */
public final class ReloadingKeyManager extends X509ExtendedKeyManager {

    private static final Logger log = LoggerFactory.getLogger(ReloadingKeyManager.class);
    private static final String ALIAS_PREFIX = "mail-";

    record Material(String alias, X509Certificate[] chain, PrivateKey key, FileTime certTime, FileTime keyTime) {
    }

    private final Path certFile;
    private final Path keyFile;
    private final Duration interval;
    private final Clock clock;
    /** The current pair and the one before it (a handshake may still hold its alias) */
    private final Map<String, Material> byAlias = new ConcurrentHashMap<>();
    private volatile Material current;
    private volatile long checkedAt;
    private int generation;

    public ReloadingKeyManager(Path certFile, Path keyFile, Duration interval, Clock clock) throws Exception {
        this.certFile = certFile;
        this.keyFile = keyFile;
        this.interval = interval;
        this.clock = clock;
        install(load());
        this.checkedAt = clock.millis();
    }

    /** The pair to use now; reads the files again when the interval has passed and they changed */
    Material material() {
        long now = clock.millis();
        if (now - checkedAt >= interval.toMillis()) {
            synchronized (this) {
                if (now - checkedAt >= interval.toMillis()) {
                    checkedAt = now;
                    reloadIfChanged();
                }
            }
        }
        return current;
    }

    private void reloadIfChanged() {
        Material old = current;
        try {
            if (Files.getLastModifiedTime(certFile).equals(old.certTime())
                    && Files.getLastModifiedTime(keyFile).equals(old.keyTime())) {
                return;
            }
            Material fresh = load();
            install(fresh);
            log.info("TLS: certificate {} reloaded, valid until {}", certFile, fresh.chain()[0].getNotAfter());
        } catch (Exception e) {
            log.warn("TLS: cannot read {} / {} ({}), keeping the certificate valid until {}",
                    certFile, keyFile, e.toString(), old.chain()[0].getNotAfter());
        }
    }

    private synchronized void install(Material material) {
        Material previous = current;
        byAlias.keySet().removeIf(alias -> previous == null || !alias.equals(previous.alias()));
        byAlias.put(material.alias(), material);
        current = material;
    }

    private Material load() throws Exception {
        FileTime certTime = Files.getLastModifiedTime(certFile);
        FileTime keyTime = Files.getLastModifiedTime(keyFile);
        X509Certificate[] chain;
        try (InputStream in = Files.newInputStream(certFile)) {
            chain = CertificateFactory.getInstance("X.509").generateCertificates(in).stream()
                    .map(X509Certificate.class::cast).toArray(X509Certificate[]::new);
        }
        if (chain.length == 0) {
            throw new IllegalStateException("no certificate in " + certFile);
        }
        PrivateKey key = readKey(keyFile);
        if (!chain[0].getPublicKey().getAlgorithm().equals(key.getAlgorithm())) {
            throw new IllegalStateException("the key in " + keyFile + " is " + key.getAlgorithm()
                    + ", the certificate is " + chain[0].getPublicKey().getAlgorithm());
        }
        String alias;
        synchronized (this) {
            alias = ALIAS_PREFIX + (++generation);
        }
        return new Material(alias, chain, key, certTime, keyTime);
    }

    static PrivateKey readKey(Path file) throws Exception {
        try (Reader reader = Files.newBufferedReader(file, StandardCharsets.US_ASCII);
             PEMParser parser = new PEMParser(reader)) {
            Object object;
            while ((object = parser.readObject()) != null) {
                JcaPEMKeyConverter converter = new JcaPEMKeyConverter();
                if (object instanceof PrivateKeyInfo info) {
                    return converter.getPrivateKey(info);
                }
                if (object instanceof PEMKeyPair pair) {
                    return converter.getPrivateKey(pair.getPrivateKeyInfo()); // SEC1 may come without the public part
                }
                // EC PARAMETERS before EC PRIVATE KEY (openssl ecparam -genkey) — skip
            }
        }
        throw new IllegalStateException("no private key in " + file);
    }

    private String serverAlias(String keyType) {
        Material material = material();
        String algorithm = material.key().getAlgorithm();
        boolean fits = keyType == null || keyType.equalsIgnoreCase(algorithm)
                || ("RSA".equals(algorithm) && keyType.equalsIgnoreCase("RSASSA-PSS"))
                || ("EC".equals(algorithm) && keyType.toUpperCase().startsWith("EC"));
        return fits ? material.alias() : null;
    }

    @Override
    public String chooseServerAlias(String keyType, Principal[] issuers, Socket socket) {
        return serverAlias(keyType);
    }

    @Override
    public String chooseEngineServerAlias(String keyType, Principal[] issuers, SSLEngine engine) {
        return serverAlias(keyType);
    }

    @Override
    public String[] getServerAliases(String keyType, Principal[] issuers) {
        String alias = serverAlias(keyType);
        return alias == null ? null : new String[]{alias};
    }

    @Override
    public X509Certificate[] getCertificateChain(String alias) {
        Material material = alias == null ? null : byAlias.get(alias);
        return material == null ? null : material.chain().clone();
    }

    @Override
    public PrivateKey getPrivateKey(String alias) {
        Material material = alias == null ? null : byAlias.get(alias);
        return material == null ? null : material.key();
    }

    @Override
    public String[] getClientAliases(String keyType, Principal[] issuers) {
        return null;
    }

    @Override
    public String chooseClientAlias(String[] keyType, Principal[] issuers, Socket socket) {
        return null;
    }

    @Override
    public String chooseEngineClientAlias(String[] keyType, Principal[] issuers, SSLEngine engine) {
        return null;
    }
}