heicode-win/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")
)

// 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())
	}
}