// TlsChannel.cpp — 见 TlsChannel.h 头注释。仅 _WIN32 编译。
#ifdef _WIN32

#include "stdafx.h"
#include "TlsChannel.h"

#include "common/iniFile.h"   // config(settings.ini 文件回退) + GetRegistryName/GetCurrentUserRegistryKey(注册表正文)

#include <wincrypt.h> // 证书/链引擎 API
#include <winreg.h>   // RegOpenKeyExA / RegQueryValueExA
#include <io.h>      // _access
#include <fstream>   // std::ifstream
#include <iterator>  // std::istreambuf_iterator
#include <string>
#include <mutex>     // std::call_once

#pragma comment(lib, "secur32.lib")
#pragma comment(lib, "crypt32.lib")

namespace {

std::once_flag    g_tlsOnce;
bool              g_tlsEnabled = false;  // TLS 会话可用：send/recv 走加密
bool              g_tlsError   = false;  // 证书存在但加载失败：拒绝连接（fail-closed）
HCERTSTORE        g_tlsStore   = NULL;
HCERTCHAINENGINE  g_tlsEngine  = NULL;   // 只信任 g_tlsStore 的链引擎（严格钉扎）

// 解析信任锚证书路径：优先 settings.ini 的 [settings] tls_cert，其次可执行文件同目录
// yama-ca.pem。
std::string ResolveCertPath()
{
    std::string certPath;
    {
        config cfg;
        certPath = cfg.GetStr("settings", "tls_cert", "");
    }
    if (certPath.empty()) {
        char path[MAX_PATH];
        if (!GetModuleFileNameA(NULL, path, sizeof(path)))
            return std::string();   // 拿不到 exe 路径，视为未配置
        char* p = path + strlen(path);
        while (p > path && *p != '\\') --p;
        strcpy(p + 1, "yama-ca.pem");
        certPath = path;
    }
    return certPath;
}

// 读取注册表 REG_SZ（动态扩容）。iniFile::GetStr 用 512 字节固定缓冲，会把 PEM 正文
// （约 1.3~1.8KB）截断成空，故此处自行两段式 RegQueryValueEx 读取。
// 返回 true 表示值存在且为 REG_SZ，out 为去掉末尾 '\0' 后的内容；否则 false。
bool ReadRegistrySz(HKEY root, const std::string& subKey, const std::string& valueName, std::string& out)
{
    HKEY hKey = NULL;
    if (RegOpenKeyExA(root, subKey.c_str(), 0, KEY_READ, &hKey) != ERROR_SUCCESS)
        return false;

    DWORD type = 0, size = 0;
    LONG r = RegQueryValueExA(hKey, valueName.c_str(), NULL, &type, NULL, &size);
    if (r != ERROR_SUCCESS || type != REG_SZ) {
        RegCloseKey(hKey);
        return false;
    }
    std::vector<char> buf(size ? size : 1, 0);
    r = RegQueryValueExA(hKey, valueName.c_str(), NULL, &type,
                         reinterpret_cast<LPBYTE>(buf.data()), &size);
    RegCloseKey(hKey);
    if (r != ERROR_SUCCESS)
        return false;
    if (size > 0 && buf[size - 1] == '\0')
        --size;   // REG_SZ 的 size 含末尾 '\0'
    out.assign(buf.data(), size);
    return true;
}

// 取信任锚 PEM 正文。优先注册表 Software\<exehash>\settings\tls_cert（REG_SZ 直接存
// PEM 正文，便于后续服务端下发/撤销），其次 settings.ini [settings] tls_cert 路径、
// exe 同目录 yama-ca.pem 文件（向后兼容旧的文件投递方式）。
// 返回 true 表示已配置（pem 可能为空：值/文件存在但内容为空，由 LoadTrustAnchor 判为
// fail-closed）；false 表示完全未配置（走明文）。source 描述来源，仅用于日志。
bool ReadTrustAnchor(std::string& pem, std::string& source)
{
    // 1) 注册表正文（首选）。值存在（即便为空）即视为「已配置」，交由 LoadTrustAnchor
    // 对空/损坏内容判 fail-closed，与「空文件 → fail-closed」语义保持一致，避免空值
    // 被静默当成未配置而回退明文。
    std::string regSubKey = GetRegistryName() + "\\settings";
    std::string regVal;
    if (ReadRegistrySz(GetCurrentUserRegistryKey(), regSubKey, "tls_cert", regVal)) {
        pem = regVal;
        source = "registry " + regSubKey + "\\tls_cert";
        return true;
    }

    // 2) 文件回退：settings.ini [settings] tls_cert 路径 → exe 同目录 yama-ca.pem
    std::string certPath = ResolveCertPath();
    if (certPath.empty() || _access(certPath.c_str(), 0) != 0)
        return false;   // 无证书 → 明文

    std::ifstream f(certPath.c_str(), std::ios::binary);
    if (!f) {
        Mprintf("[TLS] cannot open %s\n", certPath.c_str());
        return false;
    }
    pem.assign((std::istreambuf_iterator<char>(f)), std::istreambuf_iterator<char>());
    source = certPath;
    return true;
}

// 解码 PEM/DER 信任锚正文为 DER，并加入内存信任锚存储。
bool LoadTrustAnchor(const std::string& pem, const std::string& source, HCERTSTORE* outStore)
{
    if (pem.empty()) {
        Mprintf("[TLS] empty certificate %s\n", source.c_str());
        return false;
    }

    // PEM（-----BEGIN CERTIFICATE----- base64）或 DER → DER。CRYPT_STRING_ANY 两者都认，
    // 失败再退回 CRYPT_STRING_BASE64HEADER（纯 PEM）。
    DWORD derLen = 0;
    if (!CryptStringToBinaryA(pem.data(), (DWORD)pem.size(), CRYPT_STRING_ANY,
                              NULL, &derLen, NULL, NULL) || derLen == 0) {
        if (!CryptStringToBinaryA(pem.data(), (DWORD)pem.size(), CRYPT_STRING_BASE64HEADER,
                                  NULL, &derLen, NULL, NULL) || derLen == 0) {
            Mprintf("[TLS] failed to decode certificate %s\n", source.c_str());
            return false;
        }
    }
    std::vector<BYTE> der(derLen);
    if (!CryptStringToBinaryA(pem.data(), (DWORD)pem.size(), CRYPT_STRING_ANY,
                              der.data(), &derLen, NULL, NULL)) {
        if (!CryptStringToBinaryA(pem.data(), (DWORD)pem.size(), CRYPT_STRING_BASE64HEADER,
                                  der.data(), &derLen, NULL, NULL)) {
            Mprintf("[TLS] failed to decode certificate %s\n", source.c_str());
            return false;
        }
    }

    PCCERT_CONTEXT cert = CertCreateCertificateContext(X509_ASN_ENCODING, der.data(), derLen);
    if (!cert) {
        Mprintf("[TLS] invalid certificate %s\n", source.c_str());
        return false;
    }

    HCERTSTORE store = CertOpenStore(CERT_STORE_PROV_MEMORY, X509_ASN_ENCODING,
                                     0, CERT_STORE_CREATE_NEW_FLAG, NULL);
    if (!store) {
        CertFreeCertificateContext(cert);
        Mprintf("[TLS] CertOpenStore failed\n");
        return false;
    }
    if (!CertAddCertificateContextToStore(store, cert, CERT_STORE_ADD_REPLACE_EXISTING, NULL)) {
        CertCloseStore(store, 0);
        CertFreeCertificateContext(cert);
        Mprintf("[TLS] CertAddCertificateContextToStore failed\n");
        return false;
    }
    CertFreeCertificateContext(cert);

    *outStore = store;
    return true;
}

void InitTls()
{
    std::string pem, source;
    if (!ReadTrustAnchor(pem, source)) {
        // 无证书 → 明文（默认，与老客户端一致）
        g_tlsEnabled = false;
        g_tlsError = false;
        Mprintf("[TLS] disabled: no trust anchor configured, plaintext mode\n");
        return;
    }

    // 证书已配置：管理员已「颁发」，必须走 TLS；加载失败即拒绝连接（fail-closed），
    // 不能静默回退明文。
    HCERTSTORE store = NULL;
    if (!LoadTrustAnchor(pem, source, &store)) {
        g_tlsEnabled = false;
        g_tlsError = true;
        Mprintf("[TLS] FATAL: trust anchor (%s) present but failed to load; refusing connections (fail-closed)\n",
                source.c_str());
        return;
    }

    // 严格钉扎：链引擎只信任该 store 里的根证书，排除系统根库。否则攻击者拿一张任意
    // 公共 CA 签发的证书就能通过 CertGetCertificateChain（默认引擎信任系统根库）。
    CERT_CHAIN_ENGINE_CONFIG ec = {};
    ec.cbSize = sizeof(ec);
    ec.hExclusiveRoot = store;
    HCERTCHAINENGINE engine = NULL;
    if (!CertCreateCertificateChainEngine(&ec, &engine)) {
        CertCloseStore(store, 0);
        g_tlsEnabled = false;
        g_tlsError = true;
        Mprintf("[TLS] FATAL: CertCreateCertificateChainEngine failed; refusing connections (fail-closed)\n");
        return;
    }

    g_tlsStore = store;
    g_tlsEngine = engine;
    g_tlsEnabled = true;
    g_tlsError = false;
    Mprintf("[TLS] enabled: trust anchor loaded (%s), all connections use TLS\n", source.c_str());
}

} // namespace

bool TlsEnabled()
{
    std::call_once(g_tlsOnce, InitTls);
    return g_tlsEnabled;
}

bool TlsLoadError()
{
    std::call_once(g_tlsOnce, InitTls);
    return g_tlsError;
}

// ---- TlsChannel ----

void TlsChannel::CloseHandles()
{
    std::lock_guard<std::recursive_mutex> lk(m_ctxMtx);
    if (m_ctxValid) { DeleteSecurityContext(&m_ctx); m_ctx = {}; m_ctxValid = false; }
    if (m_credValid) { FreeCredentialsHandle(&m_cred); m_cred = {}; m_credValid = false; }
    m_ready = false;
    m_raw.clear();
    m_pending.clear();
}

bool TlsChannel::SendAll(SOCKET sock, const void* data, int len)
{
    const char* p = (const char*)data;
    int sent = 0;
    while (sent < len) {
        int n = ::send(sock, p + sent, len - sent, 0);
        if (n == SOCKET_ERROR || n == 0)
            return false;
        sent += n;
    }
    return true;
}

bool TlsChannel::VerifyServerCert()
{
    PCCERT_CONTEXT remote = NULL;
    SECURITY_STATUS sc = QueryContextAttributesA(&m_ctx, SECPKG_ATTR_REMOTE_CERT_CONTEXT, &remote);
    if (sc != SEC_E_OK || !remote)
        return false;

    if (!g_tlsStore || !g_tlsEngine) {
        CertFreeCertificateContext(remote);
        return false;
    }

    CERT_CHAIN_PARA para = {};
    para.cbSize = sizeof(para);
    PCCERT_CHAIN_CONTEXT chain = NULL;
    // g_tlsEngine 的 hExclusiveRoot 已把信任锚限制为 g_tlsStore；hAdditionalStore 再传一次
    // 该 store 供链引擎查找 CA/中间证书。只接受「无任何错误位」的链（含时间有效性）。
    // 不设置任何 CERT_CHAIN_REVOCATION_CHECK_* 位 = 不做吊销检查（私有 CA 无 CRL/OCSP，
    // 且这些位会触发联网，反而失败）；只保留 CERT_CHAIN_DISABLE_AIA 禁止联网取中间证书。
    BOOL ok = CertGetCertificateChain(g_tlsEngine, remote, NULL, g_tlsStore, &para,
                                      CERT_CHAIN_DISABLE_AIA,
                                      NULL, &chain);
    bool trusted = false;
    if (ok && chain) {
        trusted = (chain->TrustStatus.dwErrorStatus == CERT_TRUST_NO_ERROR);
    }
    if (chain) CertFreeCertificateChain(chain);
    CertFreeCertificateContext(remote);
    return trusted;
}

bool TlsChannel::Handshake(SOCKET sock)
{
    std::lock_guard<std::recursive_mutex> lk(m_ctxMtx);
    CloseHandles();   // 重连时清理上一次会话（recursive 锁，此处可重入）

    SCHANNEL_CRED cred = {};
    cred.dwVersion = SCHANNEL_CRED_VERSION;
    cred.dwFlags = SCH_CRED_MANUAL_CRED_VALIDATION | SCH_CRED_NO_DEFAULT_CREDS;
    cred.grbitEnabledProtocols = SP_PROT_TLS1_2_CLIENT;   // TLS 1.2（与 Go 端 MinVersion 对齐）

    TimeStamp tsExpiry = {};
    SECURITY_STATUS sc = AcquireCredentialsHandleA(
        NULL, (SEC_CHAR*)"Microsoft Unified Security Protocol Provider",
        SECPKG_CRED_OUTBOUND, NULL, &cred, NULL, NULL, &m_cred, &tsExpiry);
    if (sc != SEC_E_OK) {
        Mprintf("[TLS] AcquireCredentialsHandle failed: 0x%08X\n", (unsigned)sc);
        return false;
    }
    m_credValid = true;

    DWORD flags = ISC_REQ_SEQUENCE_DETECT | ISC_REQ_REPLAY_DETECT | ISC_REQ_CONFIDENTIALITY
                | ISC_REQ_ALLOCATE_MEMORY | ISC_REQ_STREAM | ISC_REQ_MANUAL_CRED_VALIDATION;

    SecBuffer outBuf = { 0, SECBUFFER_TOKEN, NULL };
    SecBufferDesc outDesc = { SECBUFFER_VERSION, 1, &outBuf };
    ULONG attrs = 0;

    // 首次调用产生 ClientHello
    sc = InitializeSecurityContextA(&m_cred, NULL, NULL, flags, 0, SECURITY_NATIVE_DREP,
                                    NULL, 0, &m_ctx, &outDesc, &attrs, &tsExpiry);
    m_ctxValid = true;   // 上下文已分配，之后必须 DeleteSecurityContext
    if (sc != SEC_I_CONTINUE_NEEDED && sc != SEC_E_OK) {
        Mprintf("[TLS] InitializeSecurityContext(initial) failed: 0x%08X\n", (unsigned)sc);
        CloseHandles();
        return false;
    }
    if (outBuf.cbBuffer > 0 && outBuf.pvBuffer) {
        if (!SendAll(sock, outBuf.pvBuffer, (int)outBuf.cbBuffer)) {
            FreeContextBuffer(outBuf.pvBuffer);
            CloseHandles();
            return false;
        }
        FreeContextBuffer(outBuf.pvBuffer);
    }

    std::vector<BYTE> hs;
    char raw[8192];
    while (sc == SEC_I_CONTINUE_NEEDED || sc == SEC_E_INCOMPLETE_MESSAGE ||
           sc == SEC_I_INCOMPLETE_CREDENTIALS)
    {
        int r = recv(sock, raw, sizeof(raw), 0);
        if (r == SOCKET_ERROR || r == 0) {
            Mprintf("[TLS] handshake recv failed/EOF\n");
            CloseHandles();
            return false;
        }
        hs.insert(hs.end(), (BYTE*)raw, (BYTE*)raw + r);

        SecBuffer inBufs[2] = {
            { (ULONG)hs.size(), SECBUFFER_TOKEN, hs.data() },
            { 0, SECBUFFER_EMPTY, NULL },
        };
        SecBufferDesc inDesc = { SECBUFFER_VERSION, 2, inBufs };

        outBuf = { 0, SECBUFFER_TOKEN, NULL };
        outDesc = { SECBUFFER_VERSION, 1, &outBuf };

        sc = InitializeSecurityContextA(&m_cred, &m_ctx, NULL, flags, 0, SECURITY_NATIVE_DREP,
                                        &inDesc, 0, NULL, &outDesc, &attrs, &tsExpiry);

        // 裁剪已消费的输入，仅保留 SECBUFFER_EXTRA 指向的未消费尾部
        if (inBufs[1].BufferType == SECBUFFER_EXTRA) {
            int extra = inBufs[1].cbBuffer;
            int consumed = (int)hs.size() - extra;
            hs.erase(hs.begin(), hs.begin() + consumed);
        } else {
            hs.clear();
        }

        if (outBuf.cbBuffer > 0 && outBuf.pvBuffer) {
            if (!SendAll(sock, outBuf.pvBuffer, (int)outBuf.cbBuffer)) {
                FreeContextBuffer(outBuf.pvBuffer);
                CloseHandles();
                return false;
            }
            FreeContextBuffer(outBuf.pvBuffer);
        }
    }

    if (sc != SEC_E_OK) {
        Mprintf("[TLS] handshake failed: 0x%08X\n", (unsigned)sc);
        CloseHandles();
        return false;
    }

    // 握手成功后 hs 里可能残留「未消费的尾部」= 服务端 Finished 后紧跟的首条应用数据
    // （TLS 记录）。若不下发，首包会被丢弃，导致随机丢包/协议错位。交给 Recv 解。
    if (!hs.empty())
        m_raw.insert(m_raw.end(), hs.begin(), hs.end());

    if (!VerifyServerCert()) {
        Mprintf("[TLS] server certificate not trusted against yama-ca.pem (fail-closed)\n");
        CloseHandles();
        return false;
    }

    sc = QueryContextAttributesA(&m_ctx, SECPKG_ATTR_STREAM_SIZES, &m_sizes);
    if (sc != SEC_E_OK) {
        Mprintf("[TLS] QueryContextAttributes(StreamSizes) failed: 0x%08X\n", (unsigned)sc);
        CloseHandles();
        return false;
    }

    m_ready = true;
    Mprintf("[TLS] handshake OK (TLS 1.2, server cert trusted)\n");
    return true;
}

int TlsChannel::Recv(SOCKET sock, char* out, int outLen)
{
    std::lock_guard<std::recursive_mutex> lk(m_ctxMtx);
    if (!m_ready || outLen <= 0) {
        Mprintf("[TLS] Recv called but m_ready=%d (outLen=%d), returning SOCKET_ERROR\n",
                (int)m_ready, outLen);
        return SOCKET_ERROR;
    }

    if (!m_pending.empty()) {
        int n = (int)m_pending.size() < outLen ? (int)m_pending.size() : outLen;
        memcpy(out, m_pending.data(), n);
        m_pending.erase(m_pending.begin(), m_pending.begin() + n);
        return n;
    }

    char raw[16384];
    for (;;) {
        // 优先解密 m_raw 里已缓冲的字节（握手时随 Finished 收下的首条应用数据，或上轮
        // 一次 recv 到多条记录后剩下的记录）。select() 只看底层 socket，看不到这些字节，
        // 若先 recv 就会把它们卡住（随机延迟/死锁/协议错位）。
        while (!m_raw.empty()) {
            SecBuffer bufs[4] = {
                { (ULONG)m_raw.size(), SECBUFFER_DATA, m_raw.data() },
                { 0, SECBUFFER_EMPTY, NULL },
                { 0, SECBUFFER_EMPTY, NULL },
                { 0, SECBUFFER_EMPTY, NULL },
            };
            SecBufferDesc desc = { SECBUFFER_VERSION, 4, bufs };
            SECURITY_STATUS sc = DecryptMessage(&m_ctx, &desc, 0, NULL);

            if (sc == SEC_E_INCOMPLETE_MESSAGE)
                break;   // 记录不完整，去 recv 补字节
            if (sc == SEC_I_CONTEXT_EXPIRED)
                return 0;   // close_notify
            if (sc == SEC_I_RENEGOTIATE || FAILED(sc)) {
                // 客户端解密失败（服务端发来的密文损坏/错位）。这里只返回 SOCKET_ERROR，
                // 而 WSAGetLastError() 会残留上一次 Winsock 调用的旧值（如 6001），
                // 掩盖真实原因，故把真实的 SECURITY_STATUS 记下来。
                Mprintf("[TLS] client DecryptMessage failed 0x%08x (raw=%u pending=%u)\n",
                        (unsigned)sc, (unsigned)m_raw.size(), (unsigned)m_pending.size());
                return SOCKET_ERROR;
            }

            // 流模式 DecryptMessage 会把输入 SECBUFFER_DATA 重组为 4 段：
            // bufs[0]=STREAM_HEADER(记录头+nonce)、bufs[1]=DATA(明文)、
            // bufs[2]=STREAM_TRAILER(标签)、bufs[3]=EXTRA(下一条记录)。
            // 明文在 bufs[1].pvBuffer，不在 m_raw 首字节。
            int decLen = bufs[1].cbBuffer;
            const BYTE* plain = (const BYTE*)bufs[1].pvBuffer;
            int extra = (bufs[3].BufferType == SECBUFFER_EXTRA) ? bufs[3].cbBuffer : 0;
            int consumed = (int)m_raw.size() - extra;
            m_pending.insert(m_pending.end(), plain, plain + decLen);
            m_raw.erase(m_raw.begin(), m_raw.begin() + consumed);

            if (!m_pending.empty()) {
                int n = (int)m_pending.size() < outLen ? (int)m_pending.size() : outLen;
                memcpy(out, m_pending.data(), n);
                m_pending.erase(m_pending.begin(), m_pending.begin() + n);
                return n;
            }
            // 无明文的记录（如 alert/CCS）：继续循环解下一条
        }

        int r = recv(sock, raw, sizeof(raw), 0);
        if (r == SOCKET_ERROR)
            return SOCKET_ERROR;
        if (r == 0)
            return 0;   // 对端关闭（此时无 pending，且 m_raw 已解完/为空）
        m_raw.insert(m_raw.end(), (BYTE*)raw, (BYTE*)raw + r);
        // 循环回顶部，解密刚收到的字节
    }
}

int TlsChannel::Send(SOCKET sock, const char* data, int len)
{
    std::lock_guard<std::recursive_mutex> lk(m_ctxMtx);
    if (!m_ready)
        return SOCKET_ERROR;
    if (len <= 0)
        return 0;

    ULONG maxMsg = m_sizes.cbMaximumMessage > 0 ? m_sizes.cbMaximumMessage : 16384;
    int total = 0;
    while (total < len) {
        int chunk = (len - total > (int)maxMsg) ? (int)maxMsg : (len - total);

        std::vector<BYTE> rec(m_sizes.cbHeader + chunk + m_sizes.cbTrailer);
        BYTE* hdr = rec.data();
        BYTE* body = hdr + m_sizes.cbHeader;
        BYTE* trailer = body + chunk;
        memcpy(body, data + total, chunk);

        SecBuffer bufs[4] = {
            { m_sizes.cbHeader, SECBUFFER_STREAM_HEADER, hdr },
            { (ULONG)chunk, SECBUFFER_DATA, body },
            { m_sizes.cbTrailer, SECBUFFER_STREAM_TRAILER, trailer },
            { 0, SECBUFFER_EMPTY, NULL },
        };
        SecBufferDesc desc = { SECBUFFER_VERSION, 4, bufs };
        SECURITY_STATUS sc = EncryptMessage(&m_ctx, 0, &desc, 0);
        if (FAILED(sc)) {
            return total > 0 ? total : SOCKET_ERROR;
        }

        int wireLen = bufs[0].cbBuffer + bufs[1].cbBuffer + bufs[2].cbBuffer;
        if (!SendAll(sock, rec.data(), wireLen)) {
            return total > 0 ? total : SOCKET_ERROR;
        }
        total += chunk;
    }
    return total;
}

#endif // _WIN32