---

//go:build linux

// Пример: клиентский TLS-ключ из Vault -> запрос к etcd -> затирание ключа.
// Go 1.25. Затирание best effort: см. комментарии про копии, которые стереть нельзя.
package main

import (
	"bytes"
	"context"
	"crypto"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"log"
	"math/big"
	"net/http"
	"os"
	"sync"
	"syscall"
	"time"
	"unicode/utf8"
)

const (
	vaultAddr       = "https://vault.example.com:8200"
	vaultSecretPath = "/v1/secret/data/etcd-client" // KV v2: поля certificate, private_key
	etcdURL         = "https://etcd.example.com:2379"
	etcdCAFile      = "/etc/etcd/ca.pem"

	maxVaultResponse = 64 << 10
)

func main() {
	hardenProcess()

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	caPEM, err := os.ReadFile(etcdCAFile)
	if err != nil {
		log.Fatal(err)
	}
	caPool := x509.NewCertPool()
	if !caPool.AppendCertsFromPEM(caPEM) {
		log.Fatal("bad etcd CA")
	}

	token := readVaultToken() // []byte, а не string
	defer clear(token)

	data, err := fetchFromEtcd(ctx, token, caPool, "/config/app")
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("etcd response: %s\n", data)
}

// fetchFromEtcd: ключ живёт только внутри этой функции.
func fetchFromEtcd(ctx context.Context, vaultToken []byte, caPool *x509.CertPool, key string) ([]byte, error) {
	certPEM, keyPEM, err := fetchClientCertFromVault(ctx, vaultToken)
	if err != nil {
		return nil, err
	}

	cert, err := buildCertificate(certPEM, keyPEM)
	clear(keyPEM) // PEM ключа больше не нужен
	if err != nil {
		return nil, err
	}

	ck := &clientKey{cert: cert}
	defer ck.wipe() // выполнится последним (defer — LIFO)

	tr := &http.Transport{
		TLSClientConfig: &tls.Config{
			MinVersion: tls.VersionTLS12,
			RootCAs:    caPool,
			// Ключ не кладём в Certificates: отдаём только на время handshake.
			GetClientCertificate: ck.get,
		},
		DisableKeepAlives: true, // соединение закрывается сразу после запроса
	}
	defer tr.CloseIdleConnections()

	client := &http.Client{Transport: tr, Timeout: 10 * time.Second}
	return etcdRange(ctx, client, key)
}

// ---------- Vault ----------

var vaultHTTP = &http.Client{
	Timeout: 10 * time.Second,
	Transport: &http.Transport{
		// Без gzip: иначе расшифрованный ответ осядет ещё и в буферах декомпрессора.
		DisableCompression: true,
		DisableKeepAlives:  true,
	},
}

func fetchClientCertFromVault(ctx context.Context, token []byte) (certPEM, keyPEM []byte, err error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, vaultAddr+vaultSecretPath, nil)
	if err != nil {
		return nil, nil, err
	}
	// Заголовки в net/http — строки, так что копия токена тут неизбежна.
	req.Header.Set("X-Vault-Token", string(token))

	resp, err := vaultHTTP.Do(req)
	if err != nil {
		return nil, nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, nil, fmt.Errorf("vault: status %d", resp.StatusCode)
	}

	// Буфер фиксированного размера: io.ReadAll при росте оставил бы
	// в куче старые, незатёртые куски ответа.
	buf := make([]byte, maxVaultResponse)
	defer clear(buf)
	n, err := readInto(resp.Body, buf)
	if err != nil {
		return nil, nil, err
	}

	var payload struct {
		Data struct {
			Data map[string]json.RawMessage `json:"data"`
		} `json:"data"`
	}
	if err := json.Unmarshal(buf[:n], &payload); err != nil {
		return nil, nil, fmt.Errorf("vault: decode: %w", err) // сам ответ не логируем
	}
	fields := payload.Data.Data
	defer func() {
		for _, v := range fields {
			clear(v) // RawMessage — копии значений из buf
		}
	}()

	certPEM, err = jsonUnquote(fields["certificate"])
	if err != nil {
		return nil, nil, fmt.Errorf("vault: certificate: %w", err)
	}
	keyPEM, err = jsonUnquote(fields["private_key"])
	if err != nil {
		return nil, nil, fmt.Errorf("vault: private_key: %w", err)
	}
	return certPEM, keyPEM, nil
}

// readInto читает r целиком в buf без перевыделений.
func readInto(r io.Reader, buf []byte) (int, error) {
	n := 0
	for n < len(buf) {
		m, err := r.Read(buf[n:])
		n += m
		if err == io.EOF {
			return n, nil
		}
		if err != nil {
			return n, err
		}
	}
	var one [1]byte
	if m, _ := r.Read(one[:]); m > 0 {
		clear(one[:])
		return n, errors.New("response too large")
	}
	return n, nil
}

// jsonUnquote декодирует JSON-строку в []byte, минуя string.
// Результат не длиннее входа, поэтому append никогда не перевыделяет память.
func jsonUnquote(raw []byte) ([]byte, error) {
	if len(raw) < 2 || raw[0] != '"' || raw[len(raw)-1] != '"' {
		return nil, errors.New("not a JSON string")
	}
	s := raw[1 : len(raw)-1]
	out := make([]byte, 0, len(s))
	fail := func(msg string) ([]byte, error) {
		clear(out[:cap(out)])
		return nil, errors.New(msg)
	}

	for i := 0; i < len(s); i++ {
		c := s[i]
		if c != '\\' {
			out = append(out, c)
			continue
		}
		i++
		if i >= len(s) {
			return fail("bad escape")
		}
		switch s[i] {
		case 'n':
			out = append(out, '\n')
		case 'r':
			out = append(out, '\r')
		case 't':
			out = append(out, '\t')
		case 'b':
			out = append(out, '\b')
		case 'f':
			out = append(out, '\f')
		case '"', '\\', '/':
			out = append(out, s[i])
		case 'u':
			if i+4 >= len(s) {
				return fail("bad \\u escape")
			}
			var r rune
			for _, h := range s[i+1 : i+5] {
				r <<= 4
				switch {
				case h >= '0' && h <= '9':
					r |= rune(h - '0')
				case h >= 'a' && h <= 'f':
					r |= rune(h - 'a' + 10)
				case h >= 'A' && h <= 'F':
					r |= rune(h - 'A' + 10)
				default:
					return fail("bad \\u escape")
				}
			}
			if r >= 0xD800 && r <= 0xDFFF {
				return fail("surrogates not supported") // в PEM их не бывает
			}
			out = utf8.AppendRune(out, r)
			i += 4
		default:
			return fail("bad escape")
		}
	}
	return out, nil
}

// ---------- Ключ и сертификат ----------

// buildCertificate — замена tls.X509KeyPair: тот прячет DER ключа внутри,
// и затереть его нельзя. Здесь DER под нашим контролем.
func buildCertificate(certPEM, keyPEM []byte) (*tls.Certificate, error) {
	var cert tls.Certificate
	for rest := certPEM; ; {
		var b *pem.Block
		b, rest = pem.Decode(rest)
		if b == nil {
			break
		}
		if b.Type == "CERTIFICATE" {
			cert.Certificate = append(cert.Certificate, b.Bytes)
		}
	}
	if len(cert.Certificate) == 0 {
		return nil, errors.New("no certificate in PEM")
	}

	block, _ := pem.Decode(keyPEM)
	if block == nil {
		return nil, errors.New("no private key in PEM")
	}
	defer clear(block.Bytes) // DER ключа — отдельная копия, созданная pem.Decode

	key, err := parsePrivateKey(block.Type, block.Bytes)
	if err != nil {
		return nil, err
	}

	leaf, err := x509.ParseCertificate(cert.Certificate[0])
	if err != nil {
		wipePrivateKey(key)
		return nil, err
	}
	if !keyMatchesCert(key, leaf) {
		wipePrivateKey(key)
		return nil, errors.New("private key does not match certificate")
	}

	cert.PrivateKey = key
	cert.Leaf = leaf
	return &cert, nil
}

func parsePrivateKey(pemType string, der []byte) (crypto.PrivateKey, error) {
	switch pemType {
	case "PRIVATE KEY":
		return x509.ParsePKCS8PrivateKey(der)
	case "RSA PRIVATE KEY":
		return x509.ParsePKCS1PrivateKey(der)
	case "EC PRIVATE KEY":
		return x509.ParseECPrivateKey(der)
	default:
		return nil, fmt.Errorf("unsupported key type %q", pemType)
	}
}

func keyMatchesCert(key crypto.PrivateKey, leaf *x509.Certificate) bool {
	signer, ok := key.(crypto.Signer)
	if !ok {
		return false
	}
	pub, ok := signer.Public().(interface{ Equal(crypto.PublicKey) bool })
	return ok && pub.Equal(leaf.PublicKey)
}

// clientKey хранит сертификат и отдаёт его TLS только во время handshake.
type clientKey struct {
	mu   sync.Mutex
	cert *tls.Certificate
}

func (k *clientKey) get(*tls.CertificateRequestInfo) (*tls.Certificate, error) {
	k.mu.Lock()
	defer k.mu.Unlock()
	if k.cert == nil {
		return nil, errors.New("client key already wiped")
	}
	return k.cert, nil
}

// wipe вызывать только когда все запросы завершены: подпись в handshake
// идёт вне мьютекса.
func (k *clientKey) wipe() {
	k.mu.Lock()
	defer k.mu.Unlock()
	if k.cert == nil {
		return
	}
	wipePrivateKey(k.cert.PrivateKey)
	k.cert.PrivateKey = nil
	k.cert = nil
}

func wipePrivateKey(k crypto.PrivateKey) {
	switch key := k.(type) {
	case ed25519.PrivateKey:
		clear(key) // []byte — затирается полностью
	case *ed25519.PrivateKey:
		clear(*key)
	case *rsa.PrivateKey:
		wipeBig(key.D)
		for _, p := range key.Primes {
			wipeBig(p)
		}
		wipeBig(key.Precomputed.Dp)
		wipeBig(key.Precomputed.Dq)
		wipeBig(key.Precomputed.Qinv)
		for i := range key.Precomputed.CRTValues {
			wipeBig(key.Precomputed.CRTValues[i].Exp)
			wipeBig(key.Precomputed.CRTValues[i].Coeff)
			wipeBig(key.Precomputed.CRTValues[i].R)
		}
		// Внутренняя FIPS-копия ключа (неэкспортируемое поле Precomputed)
		// недоступна: только отпустить ссылку и ждать GC.
		key.Precomputed = rsa.PrecomputedValues{}
	case *ecdsa.PrivateKey:
		wipeBig(key.D)
	}
}

// wipeBig затирает сами слова big.Int, а не просто присваивает 0.
func wipeBig(b *big.Int) {
	if b == nil {
		return
	}
	clear(b.Bits())
	b.SetInt64(0)
}

// ---------- etcd ----------

func etcdRange(ctx context.Context, client *http.Client, key string) ([]byte, error) {
	reqBody, err := json.Marshal(struct {
		Key []byte `json:"key"` // []byte кодируется в base64, как ждёт gRPC-gateway etcd
	}{Key: []byte(key)})
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, etcdURL+"/v3/kv/range", bytes.NewReader(reqBody))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(io.LimitReader(resp.Body, 10<<20))
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("etcd: status %d: %s", resp.StatusCode, body)
	}
	return body, nil
}

// ---------- Процесс ----------

// hardenProcess: без core dump и без чтения памяти через ptrace / /proc/pid/mem.
func hardenProcess() {
	if err := syscall.Setrlimit(syscall.RLIMIT_CORE, &syscall.Rlimit{}); err != nil {
		log.Printf("setrlimit RLIMIT_CORE: %v", err)
	}
	if _, _, errno := syscall.RawSyscall(syscall.SYS_PRCTL, syscall.PR_SET_DUMPABLE, 0, 0); errno != 0 {
		log.Printf("prctl PR_SET_DUMPABLE: %v", errno)
	}
	// Не даём страницам уйти в swap. Нужен CAP_IPC_LOCK или большой RLIMIT_MEMLOCK.
	if err := syscall.Mlockall(syscall.MCL_CURRENT | syscall.MCL_FUTURE); err != nil {
		log.Printf("mlockall (ключ может попасть в swap): %v", err)
	}
}

// readVaultToken читает токен из файла (например, от Vault Agent), не через string.
func readVaultToken() []byte {
	b, err := os.ReadFile("/var/run/secrets/vault-token")
	if err != nil {
		log.Fatal(err)
	}
	return bytes.TrimSpace(b)
}