// Package wire implements the relay envelope: what is signed, what is encrypted, and in what order.
//
// The relay is not trusted, so signing and encryption are separate concerns on purpose. Ed25519
// signs a canonical byte encoding of the envelope's routing fields plus the ciphertext, which is
// what stops a compromised relay from retargeting a request at a different node, replaying it under
// a new id, or substituting ek. AES-256-GCM then carries the payload under keys from an
// ephemeral-static X25519 exchange.
//
// The signature is over a canonical encoding rather than over JSON text. Three languages have to
// agree on the exact bytes here, and "sign the JSON" is where cross-language signature bugs live:
// key order, number formatting, and Unicode escaping all differ between Swift, Node and Go.
package wire

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/ed25519"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"strings"
)

// Version is envelope.v. A change here is a new constant, never a widened tolerance, so an older
// agent cannot be talked into accepting a format it does not fully understand.
const Version uint64 = 1

// Suite is envelope.alg. Exactly one suite exists. An unknown value is a hard error rather than a
// negotiation, because negotiation is where downgrade bugs come from.
const Suite = "X25519-ED25519-AES256GCM-HKDFSHA256"

// Domain separation. A request signature must never be replayable as a response signature, and a
// signature from this protocol must never be replayable in another protocol that reuses the key.
const (
	DomainRequest  = "msa-relay-v1-req"
	DomainResponse = "msa-relay-v1-resp"

	infoKeyRequest  = "msa-relay-v1-req"
	infoKeyResponse = "msa-relay-v1-resp"
	infoIVRequest   = "iv-req"
	infoIVResponse  = "iv-resp"
	hkdfSaltLabel   = "msa-relay-v1"
)

const (
	KeyBytes    = 32
	PublicBytes = 32
	IVBytes     = 12
	TagBytes    = 16
)

var (
	ErrSignature     = errors.New("wire: signature does not verify")
	ErrVersion       = errors.New("wire: unsupported envelope version")
	ErrSuite         = errors.New("wire: unsupported cipher suite")
	ErrBadKey        = errors.New("wire: key must be 32 bytes")
	ErrBadIV         = errors.New("wire: iv must be 12 bytes")
	ErrBadPublicKey  = errors.New("wire: public key must be 32 bytes")
	ErrOpen          = errors.New("wire: ciphertext failed to authenticate")
	ErrIdentityField = errors.New("wire: envelope is missing node, request id or suite")
	ErrPeerKey       = errors.New("wire: peer public key is unusable")
	ErrDomain        = errors.New("wire: unknown signing domain")
	ErrIntent        = errors.New("wire: confirmation intent does not match the request")
)

// Envelope is the outer, relay-visible structure. The relay routes on Node and correlates on Req,
// and can read nothing else: Ek is a public key, P is ciphertext, S covers everything.
type Envelope struct {
	V    uint64 `json:"v"`
	Node string `json:"node"`
	Req  string `json:"req"`
	App  string `json:"app"`
	TS   uint64 `json:"ts"`
	EK   []byte `json:"ek"`
	Alg  string `json:"alg"`
	P    []byte `json:"p"`
	S    []byte `json:"s,omitempty"`
}

// MarshalJSON encodes the byte fields as base64, matching the schema the relay validates against.
func (e Envelope) MarshalJSON() ([]byte, error) {
	type alias struct {
		V    uint64 `json:"v"`
		Node string `json:"node"`
		Req  string `json:"req"`
		App  string `json:"app,omitempty"`
		TS   uint64 `json:"ts"`
		EK   string `json:"ek"`
		Alg  string `json:"alg"`
		P    string `json:"p"`
		S    string `json:"s,omitempty"`
	}
	raw := alias{
		V: e.V, Node: e.Node, Req: e.Req, App: e.App, TS: e.TS, Alg: e.Alg,
		EK: base64.StdEncoding.EncodeToString(e.EK),
		P:  base64.StdEncoding.EncodeToString(e.P),
	}
	if len(e.S) > 0 {
		raw.S = base64.StdEncoding.EncodeToString(e.S)
	}
	return marshal(raw)
}

// UnmarshalJSON decodes base64 with strict padding, because Go's default decoder silently accepts
// some non-canonical forms and a signature must be computed over exactly the bytes the sender used.
func (e *Envelope) UnmarshalJSON(raw []byte) error {
	type alias struct {
		V    uint64 `json:"v"`
		Node string `json:"node"`
		Req  string `json:"req"`
		App  string `json:"app"`
		TS   uint64 `json:"ts"`
		EK   string `json:"ek"`
		Alg  string `json:"alg"`
		P    string `json:"p"`
		S    string `json:"s"`
	}
	var decoded alias
	if err := unmarshal(raw, &decoded); err != nil {
		return err
	}
	var err error
	if e.EK, err = decodeBase64(decoded.EK); err != nil {
		return fmt.Errorf("wire: ek: %w", err)
	}
	if e.P, err = decodeBase64(decoded.P); err != nil {
		return fmt.Errorf("wire: p: %w", err)
	}
	if e.S, err = decodeBase64(decoded.S); err != nil {
		return fmt.Errorf("wire: s: %w", err)
	}
	e.V, e.Node, e.Req, e.App, e.TS, e.Alg = decoded.V, decoded.Node, decoded.Req, decoded.App, decoded.TS, decoded.Alg
	return nil
}

func decodeBase64(value string) ([]byte, error) {
	if value == "" {
		return nil, nil
	}
	raw, err := base64.StdEncoding.Strict().DecodeString(value)
	if err != nil {
		return nil, err
	}
	return raw, nil
}

func encodeBase64(value []byte) string {
	if len(value) == 0 {
		return ""
	}
	return base64.StdEncoding.EncodeToString(value)
}

// InnerRequest is the decrypted request payload. There is no separate nonce: Req is a uuid4, it is
// covered by the signature, and it is recorded in the replay cache, so reusing it requires forging
// a signature.
type InnerRequest struct {
	Method string `json:"m"`
	Path   string `json:"p"`
	// Intent is sha256 over the method, path and canonical body. It is produced by the phone as the
	// result of a user gesture. The agent recomputes it and refuses a dangerous write that does not
	// match, so a relay that fabricates a request cannot manufacture the user's consent.
	Intent []byte `json:"intent,omitempty"`
	Body   []byte `json:"body,omitempty"`
}

// MarshalJSON encodes the byte fields as base64, the same convention as the outer envelope.
//
// Go would base64 a []byte on its own, but the decode side still has to be written, and keeping both
// directions together is what stops one of them from being forgotten when a field is added.
func (r InnerRequest) MarshalJSON() ([]byte, error) {
	type alias struct {
		Method string `json:"m"`
		Path   string `json:"p"`
		Intent string `json:"intent,omitempty"`
		Body   string `json:"body,omitempty"`
	}
	return marshal(alias{
		Method: r.Method, Path: r.Path,
		Intent: encodeBase64(r.Intent), Body: encodeBase64(r.Body),
	})
}

func (r *InnerRequest) UnmarshalJSON(raw []byte) error {
	type alias struct {
		Method string `json:"m"`
		Path   string `json:"p"`
		Intent string `json:"intent,omitempty"`
		Body   string `json:"body,omitempty"`
	}
	var decoded alias
	if err := unmarshal(raw, &decoded); err != nil {
		return err
	}
	var err error
	if r.Intent, err = decodeBase64(decoded.Intent); err != nil {
		return fmt.Errorf("wire: intent: %w", err)
	}
	if r.Body, err = decodeBase64(decoded.Body); err != nil {
		return fmt.Errorf("wire: body: %w", err)
	}
	r.Method, r.Path = decoded.Method, decoded.Path
	return nil
}

// InnerResponse is the decrypted response payload. Status lives inside the ciphertext so the relay
// cannot rewrite a failure into a success.
type InnerResponse struct {
	Status int `json:"status"`
	// Outcome is the agent's own classification of what happened. It is inside the ciphertext for
	// the same reason as Status: the relay is not entitled to know why a command failed, and the
	// phone needs to tell the user what to do about it.
	Outcome string `json:"outcome,omitempty"`
	Body    []byte `json:"body,omitempty"`
}

// MarshalJSON encodes the body as base64, matching InnerRequest.
func (r InnerResponse) MarshalJSON() ([]byte, error) {
	type alias struct {
		Status  int    `json:"status"`
		Outcome string `json:"outcome,omitempty"`
		Body    string `json:"body,omitempty"`
	}
	return marshal(alias{Status: r.Status, Outcome: r.Outcome, Body: encodeBase64(r.Body)})
}

func (r *InnerResponse) UnmarshalJSON(raw []byte) error {
	type alias struct {
		Status  int    `json:"status"`
		Outcome string `json:"outcome,omitempty"`
		Body    string `json:"body,omitempty"`
	}
	var decoded alias
	if err := unmarshal(raw, &decoded); err != nil {
		return err
	}
	body, err := decodeBase64(decoded.Body)
	if err != nil {
		return fmt.Errorf("wire: response body: %w", err)
	}
	r.Status, r.Outcome, r.Body = decoded.Status, decoded.Outcome, body
	return nil
}

// MarshalInner renders an InnerResponse. It exists as a named function rather than a bare
// json.Marshal call so every caller gets the same encoder, which here means the same treatment of
// HTML characters as the rest of the protocol.
func MarshalInner(inner InnerResponse) ([]byte, error) { return marshal(inner) }

// UnmarshalInner decodes an InnerRequest, rejecting trailing data.
func UnmarshalInner(raw []byte) (InnerRequest, error) {
	var inner InnerRequest
	if err := unmarshal(raw, &inner); err != nil {
		return InnerRequest{}, err
	}
	return inner, nil
}

// CanonicalRequest is the exact byte string K_app_sign signs and the agent verifies.
//
// Layout: domain bytes, one NUL, then for each field in fixed order a uint32be length followed by
// that many bytes. v and ts contribute their uint64be encoding, so their length is always 8 and
// there is no decimal formatting to disagree about. ek and p contribute raw decoded bytes, never
// their base64 text, so a relay cannot re-encode them differently.
func CanonicalRequest(e Envelope) ([]byte, error) {
	return canonical(e, DomainRequest)
}

// CanonicalResponse is the counterpart for agent-to-app replies. It covers the same routing fields
// so the relay cannot swap responses between in-flight requests for the same node.
func CanonicalResponse(e Envelope) ([]byte, error) {
	return canonical(e, DomainResponse)
}

func canonical(e Envelope, domain string) ([]byte, error) {
	if domain != DomainRequest && domain != DomainResponse {
		return nil, fmt.Errorf("%w: %q", ErrDomain, domain)
	}
	if strings.TrimSpace(e.Node) == "" || strings.TrimSpace(e.Req) == "" || strings.TrimSpace(e.Alg) == "" {
		return nil, ErrIdentityField
	}
	if e.Alg != Suite {
		return nil, fmt.Errorf("%w: %q", ErrSuite, e.Alg)
	}
	if len(e.EK) != PublicBytes {
		return nil, fmt.Errorf("%w: got %d bytes", ErrBadPublicKey, len(e.EK))
	}
	buf := bytes.NewBuffer(make([]byte, 0, 256))
	buf.WriteString(domain)
	buf.WriteByte(0)
	writeUint64(buf, e.V)
	writeBytes(buf, []byte(e.Node))
	writeBytes(buf, []byte(e.Req))
	// The app id is only meaningful on requests: it lets a node operator attribute traffic to a
	// paired device, and it is signed so a relay cannot relabel one phone's request as another's.
	if domain == DomainRequest {
		writeBytes(buf, []byte(e.App))
	}
	writeUint64(buf, e.TS)
	writeBytes(buf, e.EK)
	writeBytes(buf, e.P)
	return buf.Bytes(), nil
}

func writeUint64(buf *bytes.Buffer, value uint64) {
	var raw [8]byte
	binary.BigEndian.PutUint64(raw[:], value)
	writeBytes(buf, raw[:])
}

func writeBytes(buf *bytes.Buffer, value []byte) {
	var length [4]byte
	binary.BigEndian.PutUint32(length[:], uint32(len(value)))
	buf.Write(length[:])
	buf.Write(value)
}

// SignRequest attaches K_app_sign's signature over the canonical encoding.
func SignRequest(envelope *Envelope, private ed25519.PrivateKey) error {
	if len(private) != ed25519.PrivateKeySize {
		return errors.New("wire: private signing key has the wrong size")
	}
	canonical, err := CanonicalRequest(*envelope)
	if err != nil {
		return err
	}
	envelope.S = ed25519.Sign(private, canonical)
	return nil
}

// SignResponse is SignRequest for the agent's own signing key.
func SignResponse(envelope *Envelope, private ed25519.PrivateKey) error {
	if len(private) != ed25519.PrivateKeySize {
		return errors.New("wire: private signing key has the wrong size")
	}
	canonical, err := CanonicalResponse(*envelope)
	if err != nil {
		return err
	}
	envelope.S = ed25519.Sign(private, canonical)
	return nil
}

// VerifyRequest is what the agent calls before doing anything else with a request. Version and
// suite are checked here rather than at the caller so no code path can verify a signature and then
// act on a format it did not intend to accept.
func VerifyRequest(envelope Envelope, peer ed25519.PublicKey) error {
	if envelope.V != Version {
		return fmt.Errorf("%w: %d", ErrVersion, envelope.V)
	}
	if envelope.Alg != Suite {
		return fmt.Errorf("%w: %q", ErrSuite, envelope.Alg)
	}
	if len(peer) != ed25519.PublicKeySize {
		return fmt.Errorf("%w: got %d bytes", ErrBadPublicKey, len(peer))
	}
	canonical, err := CanonicalRequest(envelope)
	if err != nil {
		return err
	}
	if len(envelope.S) != ed25519.SignatureSize || !ed25519.Verify(peer, canonical, envelope.S) {
		return ErrSignature
	}
	return nil
}

// VerifyResponse is the phone's side of the same check.
func VerifyResponse(envelope Envelope, peer ed25519.PublicKey) error {
	if envelope.V != Version {
		return fmt.Errorf("%w: %d", ErrVersion, envelope.V)
	}
	if envelope.Alg != Suite {
		return fmt.Errorf("%w: %q", ErrSuite, envelope.Alg)
	}
	if len(peer) != ed25519.PublicKeySize {
		return fmt.Errorf("%w: got %d bytes", ErrBadPublicKey, len(peer))
	}
	canonical, err := CanonicalResponse(envelope)
	if err != nil {
		return err
	}
	if len(envelope.S) != ed25519.SignatureSize || !ed25519.Verify(peer, canonical, envelope.S) {
		return ErrSignature
	}
	return nil
}

// HKDF-SHA256, extract then expand. Implemented here rather than pulled in as a dependency: the
// agent must build with nothing but the standard library so an operator can cross-compile it for an
// ARM node from any machine, and the surface is thirty lines.
func HKDF(secret, salt, info []byte, length int) ([]byte, error) {
	if length <= 0 || length > 255*sha256.Size {
		return nil, fmt.Errorf("wire: hkdf length %d out of range", length)
	}
	extractor := hmac.New(sha256.New, salt)
	extractor.Write(secret)
	prk := extractor.Sum(nil)

	out := make([]byte, 0, length)
	var block []byte
	for counter := byte(1); len(out) < length; counter++ {
		mac := hmac.New(sha256.New, prk)
		mac.Write(block)
		mac.Write(info)
		mac.Write([]byte{counter})
		block = mac.Sum(nil)
		out = append(out, block...)
	}
	return out[:length], nil
}

// SharedSecret performs X25519. crypto/ecdh rejects the all-zero output, which is what catches a
// peer key chosen to force a known shared secret.
func SharedSecret(privateKey, peerPublic []byte) ([]byte, error) {
	if len(privateKey) != KeyBytes {
		return nil, fmt.Errorf("%w: got %d bytes", ErrBadKey, len(privateKey))
	}
	if len(peerPublic) != PublicBytes {
		return nil, fmt.Errorf("%w: got %d bytes", ErrBadPublicKey, len(peerPublic))
	}
	curve := ecdh.X25519()
	private, err := curve.NewPrivateKey(privateKey)
	if err != nil {
		return nil, fmt.Errorf("wire: private key rejected: %w", err)
	}
	peer, err := curve.NewPublicKey(peerPublic)
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrPeerKey, err)
	}
	shared, err := private.ECDH(peer)
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrPeerKey, err)
	}
	return shared, nil
}

// DeriveKey returns the AEAD key and IV for one direction of one request.
//
// The salt is the node id behind a fixed label, so a shared secret that somehow repeats across
// nodes still yields different keys. The IV is derived from the request id and the direction rather
// than being a fixed zero block: a fixed IV is only safe while every message uses a fresh key, and
// that invariant is exactly what a later "reuse the session key for a batch" optimisation would
// quietly break. Separating the IV by direction as well means a request and a response that somehow
// share a secret still cannot share a nonce.
func DeriveKey(shared []byte, nodeID, requestID string, response bool) (key, iv []byte, err error) {
	if len(shared) != KeyBytes {
		return nil, nil, fmt.Errorf("%w: got %d bytes", ErrBadKey, len(shared))
	}
	keyInfo := infoKeyRequest
	ivInfo := infoIVRequest
	if response {
		keyInfo = infoKeyResponse
		ivInfo = infoIVResponse
	}
	// The request id participates in the key as well as the IV. A shared secret is nominally used
	// once, but binding the key to the request keeps that true even if a caller ever reuses one.
	salt := []byte(hkdfSaltLabel + ":" + nodeID + ":" + requestID)
	key, err = HKDF(shared, salt, []byte(keyInfo), KeyBytes)
	if err != nil {
		return nil, nil, err
	}
	iv, err = HKDF(shared, []byte(requestID), []byte(ivInfo), IVBytes)
	if err != nil {
		return nil, nil, err
	}
	return key, iv, nil
}

// Seal encrypts with AES-256-GCM. The tag is appended to the ciphertext, matching Node's
// cipher.getAuthTag() concatenation and CryptoKit's combined sealed box.
func Seal(key, iv, plaintext []byte) ([]byte, error) {
	gcm, err := aead(key, iv)
	if err != nil {
		return nil, err
	}
	return gcm.Seal(nil, iv, plaintext, nil), nil
}

// Open decrypts and authenticates. A failure is terminal for the request, not a retry: a wrong key,
// a torn payload and a forged payload are indistinguishable here, and all three mean do not act.
func Open(key, iv, ciphertext []byte) ([]byte, error) {
	gcm, err := aead(key, iv)
	if err != nil {
		return nil, err
	}
	if len(ciphertext) < TagBytes {
		return nil, ErrOpen
	}
	plain, err := gcm.Open(nil, iv, ciphertext, nil)
	if err != nil {
		return nil, ErrOpen
	}
	return plain, nil
}

func aead(key, iv []byte) (cipher.AEAD, error) {
	if len(key) != KeyBytes {
		return nil, fmt.Errorf("%w: got %d bytes", ErrBadKey, len(key))
	}
	if len(iv) != IVBytes {
		return nil, fmt.Errorf("%w: got %d bytes", ErrBadIV, len(iv))
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	return cipher.NewGCM(block)
}

// IntentDigest is sha256(method | path | canonical body), the value a phone sends to prove that a
// human confirmed this exact write.
//
// The body is canonicalised the same way the console does it in server/proxy/confirmations.js:
// object keys sorted, no insignificant whitespace. A dangerous write whose intent does not match
// the bytes the agent is about to forward is refused, so a relay cannot edit the recipient of an
// SMS in flight and keep the confirmation valid.
func IntentDigest(method, path string, body []byte) ([]byte, error) {
	canonicalBody, err := CanonicalJSON(body)
	if err != nil {
		return nil, err
	}
	hash := sha256.New()
	hash.Write([]byte(strings.ToUpper(strings.TrimSpace(method))))
	hash.Write([]byte{'|'})
	hash.Write([]byte(path))
	hash.Write([]byte{'|'})
	hash.Write(canonicalBody)
	return hash.Sum(nil), nil
}

// VerifyIntent compares in constant time. A byte-by-byte comparison would leak how much of a
// guessed digest is correct, which turns a 256-bit search into a byte-at-a-time one.
func VerifyIntent(expected, actual []byte) bool {
	return hmac.Equal(expected, actual)
}