heicode-mananger/heicode/middleware/device_signature.go

package middleware

import (
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"io"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"

	"github.com/heicode/manager/common"
	"github.com/heicode/manager/model"
	"github.com/heicode/manager/service"
	"github.com/heicode/manager/setting/operation_setting"
)

// Header names exchanged between the Heicode desktop client and Manager.
// Names and exact byte order of the canonical string are part of the
// public client/server contract — DO NOT rename these without bumping the
// client protocol version (cc-haha/src/services/device/signRequest.ts).
const (
	HeaderDeviceID     = "X-Heicode-Device-Id"
	HeaderTimestamp    = "X-Heicode-Timestamp"
	HeaderNonce        = "X-Heicode-Nonce"
	HeaderSignature    = "X-Heicode-Signature"
	HeaderFingerprint  = "X-Heicode-Fingerprint"
	HeaderClientVer    = "X-Heicode-Client-Version"
	canonicalSeparator = "\n"
)

// Sentinel error values let callers distinguish "no headers present"
// (legacy bare-bearer client, fall through to old path) from
// "headers present but invalid" (reject the request).
var (
	ErrDeviceSignatureMissing      = errors.New("device signature headers missing")
	ErrDeviceSignatureRequired     = errors.New("device signature required for this token")
	ErrDeviceTimestampOutOfWindow  = errors.New("device signature timestamp out of allowed window")
	ErrDeviceNonceReplayed         = errors.New("device signature nonce already used")
	ErrDeviceRevoked               = errors.New("device has been revoked")
	ErrDeviceFingerprintMismatch   = errors.New("device fingerprint mismatch")
	ErrDeviceSignatureInvalid      = errors.New("device signature did not verify")
	ErrDeviceSignatureBodyReadFail = errors.New("failed to buffer request body for signature verification")
)

// Response headers + machine-readable error codes the V2 dispatcher
// surfaces on auth failure. The client (cc-haha/src/services/device)
// reads these to give the user actionable error messages — most
// importantly "your clock is X seconds off" vs the generic
// "Token invalid" string.
//
// Public contract — do NOT rename without coordinated client rollout.
const (
	HeaderServerTime = "X-Heicode-Server-Time" // unix_ms, set on EVERY V2 401
	HeaderAuthError  = "X-Heicode-Auth-Error"  // short code, see V2AuthError* below
)

// Short error codes shipped back to the client. Keep them ASCII,
// snake_case, ≤32 chars. The set is closed (clients dispatch on
// exact match) — add new ones in coordination with the desktop app.
const (
	V2AuthErrorTimestampDrift      = "timestamp_drift"
	V2AuthErrorNonceReplay         = "nonce_replay"
	V2AuthErrorRevoked             = "revoked"
	V2AuthErrorFingerprintMismatch = "fingerprint_mismatch"
	V2AuthErrorSignatureInvalid    = "signature_invalid"
	V2AuthErrorSignatureMissing    = "signature_missing"
	V2AuthErrorDeviceNotFound      = "device_not_found"
	V2AuthErrorEphPubkeyMissing    = "eph_pubkey_missing"
	V2AuthErrorEphPubkeyMalformed  = "eph_pubkey_malformed"
	V2AuthErrorBodyDecryptFailed   = "body_decrypt_failed"
	V2AuthErrorBodyTooShort        = "body_too_short"
	V2AuthErrorBodyReadFailed      = "body_read_failed"
	V2AuthErrorKeyDeriveFailed     = "key_derive_failed"
	V2AuthErrorServerNotReady      = "server_not_ready"
	V2AuthErrorUnknown             = "unknown"
)

// V2AuthErrorCode maps a sentinel error from the V2 auth chain to its
// public, short error code. Used by TokenAuth's V2 dispatcher to set
// X-Heicode-Auth-Error before aborting with 401. Errors not in this
// table return V2AuthErrorUnknown — the client treats unknown the same
// as a generic auth failure (it has to, otherwise a future server
// could lock out current clients by adding new sentinels).
func V2AuthErrorCode(err error) string {
	switch {
	case errors.Is(err, ErrDeviceTimestampOutOfWindow):
		return V2AuthErrorTimestampDrift
	case errors.Is(err, ErrDeviceNonceReplayed):
		return V2AuthErrorNonceReplay
	case errors.Is(err, ErrDeviceRevoked):
		return V2AuthErrorRevoked
	case errors.Is(err, ErrDeviceFingerprintMismatch):
		return V2AuthErrorFingerprintMismatch
	case errors.Is(err, ErrDeviceSignatureInvalid):
		return V2AuthErrorSignatureInvalid
	case errors.Is(err, ErrDeviceSignatureMissing):
		return V2AuthErrorSignatureMissing
	case errors.Is(err, ErrV2DeviceNotFound):
		return V2AuthErrorDeviceNotFound
	case errors.Is(err, ErrV2EphPubkeyMissing):
		return V2AuthErrorEphPubkeyMissing
	case errors.Is(err, ErrV2EphPubkeyMalformed):
		return V2AuthErrorEphPubkeyMalformed
	case errors.Is(err, ErrV2AEADDecryptFailed):
		return V2AuthErrorBodyDecryptFailed
	case errors.Is(err, ErrV2BodyTooShort):
		return V2AuthErrorBodyTooShort
	case errors.Is(err, ErrV2BodyReadFailed),
		errors.Is(err, ErrDeviceSignatureBodyReadFail):
		return V2AuthErrorBodyReadFailed
	case errors.Is(err, ErrV2KeyDeriveFailed):
		return V2AuthErrorKeyDeriveFailed
	case errors.Is(err, ErrV2HelpersNotInitialized):
		return V2AuthErrorServerNotReady
	}
	return V2AuthErrorUnknown
}

// SetV2AuthDiagnosticHeaders annotates the in-flight response with two
// diagnostic headers the client uses to render actionable errors. Safe
// to call multiple times — second call overwrites first. MUST be called
// before c.JSON / abortWithOpenAiMessage, otherwise gin has already
// flushed the header block to the wire.
func SetV2AuthDiagnosticHeaders(c *gin.Context, err error) {
	c.Header(HeaderServerTime, strconv.FormatInt(time.Now().UnixMilli(), 10))
	c.Header(HeaderAuthError, V2AuthErrorCode(err))
}

// VerifyDeviceSignatureIfRequired is the policy decision point for whether
// a bound token's request carries a valid Ed25519 signature.
//
// Returns nil when the caller should proceed (either the token isn't
// device-bound and global enforcement is off, or signature is valid).
// Returns a sentinel error when the caller MUST reject. The caller
// (TokenAuth) decides the HTTP response code — typically 401.
//
// Side effects on success: c.Set("device_id", ...), c.Set("device_bound",
// true), and the token's last_seen_ip/last_used_at are updated async.
func VerifyDeviceSignatureIfRequired(c *gin.Context, token *model.Token) error {
	cfg := operation_setting.GetDeviceBindingSetting()

	signedHeader := c.GetHeader(HeaderSignature)
	tokenIsBound := token.DevicePubkey != nil && *token.DevicePubkey != ""

	// CASE 1: token has no device binding.
	if !tokenIsBound {
		// If global enforcement is off, legacy bare-bearer is fine.
		if !cfg.RequireGlobal && !token.RequireDeviceBinding {
			return nil
		}
		// Global enforcement is on but this token was never paired.
		if signedHeader == "" {
			return ErrDeviceSignatureRequired
		}
		// Headers present but token has no pubkey to verify against — reject.
		return ErrDeviceSignatureRequired
	}

	// CASE 2: token is bound. Signature headers are required.
	if signedHeader == "" {
		return ErrDeviceSignatureRequired
	}

	deviceIdHeader := c.GetHeader(HeaderDeviceID)
	timestampHeader := c.GetHeader(HeaderTimestamp)
	nonce := c.GetHeader(HeaderNonce)
	fingerprintHeader := c.GetHeader(HeaderFingerprint)

	if deviceIdHeader == "" || timestampHeader == "" || nonce == "" || fingerprintHeader == "" {
		return ErrDeviceSignatureMissing
	}

	// device_id in headers must match what we stored at pair time.
	if token.DeviceId == nil || *token.DeviceId != deviceIdHeader {
		return ErrDeviceSignatureInvalid
	}

	// Revoked tokens fail closed even if signature would otherwise verify.
	if token.RevokedAt != 0 {
		return ErrDeviceRevoked
	}

	// Fingerprint binding: the server-stored fingerprint from pair time
	// must match what the client now reports. Mismatch usually means the
	// private key was copied to another machine.
	if token.DeviceFingerprint != "" && token.DeviceFingerprint != fingerprintHeader {
		return ErrDeviceFingerprintMismatch
	}

	// Timestamp window check. We parse defensively because clients
	// occasionally ship trailing whitespace or a stray '+' sign.
	tsMs, err := strconv.ParseInt(strings.TrimSpace(timestampHeader), 10, 64)
	if err != nil {
		return ErrDeviceSignatureInvalid
	}
	nowMs := time.Now().UnixMilli()
	skew := nowMs - tsMs
	if skew < 0 {
		skew = -skew
	}
	if skew > cfg.TimestampWindowMs {
		return ErrDeviceTimestampOutOfWindow
	}

	// Nonce replay check. Atomic SETNX in Redis (or memory fallback). If
	// this returned false the same (device_id, nonce) was used recently
	// — that's a replay.
	ttl := time.Duration(cfg.NonceTTLSec) * time.Second
	ok, err := service.MarkNonceUsed(deviceIdHeader, nonce, ttl)
	if err != nil {
		// Internal error — fail closed (don't silently let through).
		return ErrDeviceSignatureInvalid
	}
	if !ok {
		return ErrDeviceNonceReplayed
	}

	// Buffer the request body so the downstream relay can still read it.
	// Tee the bytes, hash them, and stash back into c.Request.Body so
	// later handlers see a fresh reader. For empty-body requests (GET)
	// the SHA-256 hash is just the hash of empty string.
	bodyHash, err := bufferAndHashBody(c)
	if err != nil {
		return ErrDeviceSignatureBodyReadFail
	}

	// Compose canonical string. ORDER MATTERS — both sides must use the
	// same byte sequence. Document this in
	// cc-haha/src/services/device/signRequest.ts.
	canonical := strings.Join([]string{
		strings.ToUpper(c.Request.Method),
		c.Request.URL.RequestURI(), // path + raw query
		timestampHeader,
		nonce,
		fingerprintHeader,
		bodyHash,
	}, canonicalSeparator)

	// Sign target: sha256(canonical). This adds one extra layer beyond
	// raw Ed25519's built-in hashing but matches what most asym-sign
	// REST APIs do (e.g. webhook signatures) and keeps the message-to-
	// sign at a fixed 32 bytes regardless of body size.
	digest := sha256.Sum256([]byte(canonical))

	if err := common.VerifyEd25519Signature(*token.DevicePubkey, digest[:], signedHeader); err != nil {
		return ErrDeviceSignatureInvalid
	}

	// Stash device info in context so logs / billing can show it.
	c.Set("device_id", deviceIdHeader)
	c.Set("device_bound", true)

	// Update last-seen async to avoid blocking the request.
	go updateDeviceLastSeen(token.Id, c.ClientIP())

	return nil
}

// bufferAndHashBody reads the entire request body into memory, computes
// SHA-256 hex, and replaces c.Request.Body with a fresh reader so the
// downstream handlers (relay, distribute, controller) still work.
//
// Pre-existing TokenAuth has not consumed the body yet — Distribute()
// does that later via GetAndValidateRequest. Reading + replacing here
// is safe.
func bufferAndHashBody(c *gin.Context) (string, error) {
	if c.Request.Body == nil {
		return hexHashOf([]byte{}), nil
	}
	body, err := io.ReadAll(c.Request.Body)
	if err != nil {
		return "", err
	}
	_ = c.Request.Body.Close()
	c.Request.Body = io.NopCloser(strings.NewReader(string(body)))
	c.Request.ContentLength = int64(len(body))
	return hexHashOf(body), nil
}

func hexHashOf(b []byte) string {
	sum := sha256.Sum256(b)
	return hex.EncodeToString(sum[:])
}

// updateDeviceLastSeen runs in a goroutine. We don't care if this
// occasionally races; it's monitoring-grade data, not auth-critical.
func updateDeviceLastSeen(tokenId int, ip string) {
	if tokenId <= 0 {
		return
	}
	now := time.Now().UnixMilli()
	if err := model.UpdateTokenDeviceLastSeen(tokenId, ip, now); err != nil {
		common.SysLog("updateDeviceLastSeen: " + err.Error())
	}
}

// ErrV2DeviceNotFound surfaces when the X-Heicode-Device-Id doesn't
// match any tokens row (or matches a non-device-bound row). 401 to
// the client; this is the V2 equivalent of "invalid token".
var ErrV2DeviceNotFound = errors.New("device not registered or revoked")

// VerifyV2DeviceSignedRequest performs full V2 (no-bearer) authn.
// Called from TokenAuth AFTER DecryptV2RequestBody has run, so the
// plaintext body hash and ephemeral pubkey are already in context.
//
// On success returns the matching *model.Token and sets the same
// `id` / `token_id` context keys that the legacy bearer path sets —
// downstream consumers (relay, distribute, billing) don't need to
// know which auth flavor authenticated the request.
//
// Wire-protocol (must match cc-haha encryptedFetch.ts):
//
//   canonical = method + "\n"
//             + path_with_query + "\n"
//             + timestamp_ms + "\n"
//             + nonce_hex + "\n"
//             + device_fingerprint + "\n"
//             + ephemeral_pubkey_b64 + "\n"
//             + sha256_hex(plaintext_body)
//
//   signature = base64( ed25519_sign(device_priv, sha256(canonical)) )
func VerifyV2DeviceSignedRequest(c *gin.Context) (*model.Token, error) {
	signedHeader := c.GetHeader(HeaderSignature)
	deviceIdHeader := c.GetHeader(HeaderDeviceID)
	timestampHeader := c.GetHeader(HeaderTimestamp)
	nonce := c.GetHeader(HeaderNonce)
	fingerprintHeader := c.GetHeader(HeaderFingerprint)

	if signedHeader == "" || deviceIdHeader == "" || timestampHeader == "" ||
		nonce == "" || fingerprintHeader == "" {
		return nil, ErrDeviceSignatureMissing
	}

	// The ephemeral pubkey was stashed in context by body_decrypt; we
	// fold it into the canonical so a MITM can't swap pubkey (and thus
	// the shared secret) without invalidating the signature.
	ephPubkeyB64, ok := V2EphemeralPubkey(c)
	if !ok || ephPubkeyB64 == "" {
		return nil, ErrV2EphPubkeyMissing
	}

	// Same plaintext hash body_decrypt computed. Trusting the context
	// here is safe because: (a) only body_decrypt sets the key, (b) it
	// runs before us in the dispatcher, (c) we never run V2 path
	// without it.
	bodyHash, ok := V2PlaintextBodyHash(c)
	if !ok {
		return nil, ErrDeviceSignatureBodyReadFail
	}

	token, err := model.FindTokenByDeviceId(deviceIdHeader)
	if err != nil {
		return nil, ErrV2DeviceNotFound
	}
	if token.RevokedAt != 0 {
		return nil, ErrDeviceRevoked
	}
	if token.DeviceFingerprint != "" && token.DeviceFingerprint != fingerprintHeader {
		return nil, ErrDeviceFingerprintMismatch
	}
	if token.DevicePubkey == nil || *token.DevicePubkey == "" {
		return nil, ErrV2DeviceNotFound
	}

	cfg := operation_setting.GetDeviceBindingSetting()
	tsMs, err := strconv.ParseInt(strings.TrimSpace(timestampHeader), 10, 64)
	if err != nil {
		return nil, ErrDeviceSignatureInvalid
	}
	nowMs := time.Now().UnixMilli()
	skew := nowMs - tsMs
	if skew < 0 {
		skew = -skew
	}
	if skew > cfg.TimestampWindowMs {
		return nil, ErrDeviceTimestampOutOfWindow
	}

	ttl := time.Duration(cfg.NonceTTLSec) * time.Second
	freshNonce, err := service.MarkNonceUsed(deviceIdHeader, nonce, ttl)
	if err != nil {
		return nil, ErrDeviceSignatureInvalid
	}
	if !freshNonce {
		return nil, ErrDeviceNonceReplayed
	}

	canonical := strings.Join([]string{
		strings.ToUpper(c.Request.Method),
		c.Request.URL.RequestURI(),
		timestampHeader,
		nonce,
		fingerprintHeader,
		ephPubkeyB64,
		bodyHash,
	}, canonicalSeparator)
	digest := sha256.Sum256([]byte(canonical))

	if err := common.VerifyEd25519Signature(*token.DevicePubkey, digest[:], signedHeader); err != nil {
		return nil, ErrDeviceSignatureInvalid
	}

	c.Set("device_id", deviceIdHeader)
	c.Set("device_bound", true)
	c.Set("id", token.UserId)
	c.Set("token_id", token.Id)

	go updateDeviceLastSeen(token.Id, c.ClientIP())

	return token, nil
}

// Unused-import guards: hex and io are used by V1 helpers above.
var _ = hex.EncodeToString
var _ = io.NopCloser