heicode-mananger/heicode/middleware/body_decrypt.go

package middleware

import (
	"bytes"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"io"
	"strings"

	"github.com/gin-gonic/gin"
	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/hkdf"

	"github.com/heicode/manager/common"
	"github.com/heicode/manager/service"
)

// Header name + content encoding that signal a V2 (encrypted body) request.
const (
	V2ContentEncoding = "heicode-aead-v1"
	HeaderEphPubkey   = "X-Heicode-Eph-Pubkey"
)

// Context keys for downstream consumers (signature verify, relay).
const (
	ctxKeyPlaintextBodyHash = "_heicode_body_sha256_hex"
	ctxKeyEphPubkeyRaw      = "_heicode_eph_pubkey_raw"
	ctxKeyIsV2              = "_heicode_is_v2"
)

// Sentinel errors so the dispatcher in TokenAuth can choose the right
// status code + i18n key without parsing strings.
var (
	ErrV2EphPubkeyMissing      = errors.New("V2 request missing X-Heicode-Eph-Pubkey")
	ErrV2EphPubkeyMalformed    = errors.New("V2 X-Heicode-Eph-Pubkey is not valid base64 or wrong length")
	ErrV2BodyReadFailed        = errors.New("V2 request body could not be read")
	ErrV2BodyTooShort          = errors.New("V2 request body shorter than nonce + tag overhead")
	ErrV2KeyDeriveFailed       = errors.New("V2 ECDH or HKDF derivation failed")
	ErrV2AEADDecryptFailed     = errors.New("V2 body AEAD decryption failed (ciphertext or AD tampered, or wrong server key)")
	ErrV2HelpersNotInitialized = errors.New("V2 helpers not initialized (server ECDH key missing?)")
)

// IsV2Request returns true if the inbound request looks like a V2
// encrypted call. Used by TokenAuth's dispatch to decide which path
// to take.
func IsV2Request(c *gin.Context) bool {
	return strings.EqualFold(c.GetHeader("Content-Encoding"), V2ContentEncoding)
}

// IsV2SignedRequest reports a device-signed request that carries NO encrypted
// body. The browser/Tauri fetch spec forbids a request body on GET/HEAD, so
// GET endpoints can't use the encrypted-body protocol (IsV2Request). Instead
// the client signs the request with the same Ed25519 device key and canonical
// string, hashing an empty body. Detected by the presence of the device
// signature + ephemeral pubkey headers without the V2 content encoding.
func IsV2SignedRequest(c *gin.Context) bool {
	if IsV2Request(c) {
		return false
	}
	return c.GetHeader(HeaderSignature) != "" &&
		c.GetHeader(HeaderDeviceID) != "" &&
		c.GetHeader(HeaderEphPubkey) != ""
}

// PrepareV2SignedRequest sets up the same context values DecryptV2RequestBody
// produces (plaintext body hash + ephemeral pubkey), but for a request with no
// encrypted body. The body (empty for GET) is hashed verbatim so
// VerifyV2DeviceSignedRequest builds the identical canonical string the client
// signed. MUST run before VerifyV2DeviceSignedRequest.
func PrepareV2SignedRequest(c *gin.Context) error {
	ephB64 := c.GetHeader(HeaderEphPubkey)
	if ephB64 == "" {
		return ErrV2EphPubkeyMissing
	}
	ephBytes, err := base64.StdEncoding.DecodeString(ephB64)
	if err != nil || len(ephBytes) != 32 {
		return ErrV2EphPubkeyMalformed
	}

	var raw []byte
	if c.Request.Body != nil {
		raw, err = io.ReadAll(c.Request.Body)
		if err != nil {
			return ErrV2BodyReadFailed
		}
		_ = c.Request.Body.Close()
		// Restore the body so downstream handlers can still read it (a signed
		// POST with a plaintext body is allowed, not just GET).
		c.Request.Body = io.NopCloser(bytes.NewReader(raw))
		c.Request.ContentLength = int64(len(raw))
	}

	sum := sha256.Sum256(raw)
	c.Set(ctxKeyPlaintextBodyHash, hexEncode(sum[:]))
	c.Set(ctxKeyEphPubkeyRaw, ephB64)
	c.Set(ctxKeyIsV2, true)
	return nil
}

// DecryptV2RequestBody performs ECDH + ChaCha20-Poly1305 AEAD decrypt
// in place. On success: c.Request.Body is replaced with the plaintext,
// Content-Length is updated, and the plaintext SHA-256 hex digest is
// stashed in ctxKeyPlaintextBodyHash for the signature middleware.
//
// MUST run BEFORE the device-signature verification, because the
// signed canonical includes sha256(plaintext_body).
//
// Returns nil on success, or one of the Err* sentinels on failure. The
// caller (TokenAuth) maps the sentinel to an HTTP 401/400 response and
// logs the cause.
func DecryptV2RequestBody(c *gin.Context) error {
	ephB64 := c.GetHeader(HeaderEphPubkey)
	if ephB64 == "" {
		return ErrV2EphPubkeyMissing
	}
	ephBytes, err := base64.StdEncoding.DecodeString(ephB64)
	if err != nil || len(ephBytes) != 32 {
		return ErrV2EphPubkeyMalformed
	}

	// AAD ties the ciphertext to several request fields. If the client
	// tampers any of these on the wire post-encryption, AEAD verify
	// fails. Order MUST match cc-haha/src/services/device/encryptedFetch.ts.
	deviceId := c.GetHeader(HeaderDeviceID)
	ts := c.GetHeader(HeaderTimestamp)
	nonce := c.GetHeader(HeaderNonce)
	method := strings.ToUpper(c.Request.Method)
	pathQuery := c.Request.URL.RequestURI()
	ad := strings.Join([]string{deviceId, ts, nonce, method, pathQuery}, "|")

	// Read the encrypted body fully into memory. /v1/messages payloads
	// are typically small (< 1 MB). We don't stream-decrypt because
	// ChaCha20-Poly1305 verifies the whole frame at once anyway.
	if c.Request.Body == nil {
		return ErrV2BodyReadFailed
	}
	raw, err := io.ReadAll(c.Request.Body)
	if err != nil {
		return ErrV2BodyReadFailed
	}
	_ = c.Request.Body.Close()

	aead, err := buildAEAD(ephBytes)
	if err != nil {
		return err
	}
	if len(raw) < aead.NonceSize()+aead.Overhead() {
		return ErrV2BodyTooShort
	}
	nonceBytes := raw[:aead.NonceSize()]
	ciphertext := raw[aead.NonceSize():]

	plaintext, err := aead.Open(nil, nonceBytes, ciphertext, []byte(ad))
	if err != nil {
		return ErrV2AEADDecryptFailed
	}

	// Replace the request body with plaintext so downstream handlers
	// (relay, distribute) see the decrypted JSON. Update Content-Length
	// to avoid Gin/Go HTTP layer treating it as chunked unexpectedly.
	c.Request.Body = io.NopCloser(bytes.NewReader(plaintext))
	c.Request.ContentLength = int64(len(plaintext))

	// Stash for the signature middleware. It will read this instead of
	// hashing the body again (and would hash the wrong thing if it
	// tried — the body reader now points at plaintext, not ciphertext).
	sum := sha256.Sum256(plaintext)
	c.Set(ctxKeyPlaintextBodyHash, hexEncode(sum[:]))
	c.Set(ctxKeyEphPubkeyRaw, ephB64)
	c.Set(ctxKeyIsV2, true)

	return nil
}

// buildAEAD performs the ECDH-then-HKDF dance that produces a fresh
// ChaCha20-Poly1305 instance for THIS request. Server private key
// never escapes service.DeriveECDHSharedSecret.
func buildAEAD(clientEphPubkey []byte) (interface {
	Seal(dst, nonce, plaintext, additionalData []byte) []byte
	Open(dst, nonce, ciphertext, additionalData []byte) ([]byte, error)
	NonceSize() int
	Overhead() int
}, error) {
	shared, err := service.DeriveECDHSharedSecret(clientEphPubkey)
	if err != nil {
		// Likely "server key not initialized" — surface a distinct
		// sentinel so ops can tell startup-order bugs from client errors.
		if strings.Contains(err.Error(), "not initialized") {
			return nil, ErrV2HelpersNotInitialized
		}
		return nil, ErrV2KeyDeriveFailed
	}

	// HKDF info string MUST match cc-haha/src/services/device/encryptedFetch.ts
	// exactly. If we ever change this, bump V2ContentEncoding to v2 + run
	// a coordinated rollout.
	r := hkdf.New(sha256.New, shared, nil, []byte("heicode-aead-v1"))
	key := make([]byte, chacha20poly1305.KeySize)
	if _, err := io.ReadFull(r, key); err != nil {
		return nil, ErrV2KeyDeriveFailed
	}
	common.SafeWipe(shared) // best-effort defense vs leak via Go GC

	aead, err := chacha20poly1305.New(key)
	if err != nil {
		return nil, ErrV2KeyDeriveFailed
	}
	return aead, nil
}

// V2PlaintextBodyHash returns the hex SHA-256 of the decrypted plaintext
// body. Signature middleware reads this for its canonical string. Only
// valid after DecryptV2RequestBody has run and set the ctx value.
func V2PlaintextBodyHash(c *gin.Context) (string, bool) {
	v, ok := c.Get(ctxKeyPlaintextBodyHash)
	if !ok {
		return "", false
	}
	s, ok := v.(string)
	return s, ok
}

// V2EphemeralPubkey returns the client's per-request X25519 public key
// (base64). Signature middleware folds this into its canonical so a
// MITM can't swap the pubkey (and thus the shared secret) without
// breaking the signature.
func V2EphemeralPubkey(c *gin.Context) (string, bool) {
	v, ok := c.Get(ctxKeyEphPubkeyRaw)
	if !ok {
		return "", false
	}
	s, ok := v.(string)
	return s, ok
}

func hexEncode(b []byte) string {
	const h = "0123456789abcdef"
	out := make([]byte, len(b)*2)
	for i, x := range b {
		out[i*2] = h[x>>4]
		out[i*2+1] = h[x&0x0f]
	}
	return string(out)
}