package capsule_test

import (
	"bytes"
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"strings"
	"testing"

	datekeys "g.activething.com/go/DateKeys"
	"g.activething.com/go/DateKeys/capsule"
	"g.activething.com/go/DateKeys/datekey"
	"g.activething.com/go/DateKeys/extension"
	"g.activething.com/go/DateKeys/internal/cbortest"
	"g.activething.com/go/DateKeys/internal/testkit"
	"g.activething.com/go/DateKeys/profile"
)

func TestPolicyNames(t *testing.T) {
	for _, p := range []capsule.Policy{capsule.TimeOnly, capsule.TimeAndKey} {
		got, err := capsule.ParsePolicy(p.String())
		if err != nil || got != p {
			t.Fatalf("%s: %v", p, err)
		}
	}
	if _, err := capsule.ParsePolicy("time_or_key"); err == nil {
		t.Fatal("unknown policy parsed")
	}
	if capsule.Policy(9).String() != "policy(9)" {
		t.Fatal("unknown policy name")
	}
}

func TestControlIsNotPrinted(t *testing.T) {
	c := capsule.Control{}
	for i := range c.PayloadIdentity {
		c.PayloadIdentity[i] = 0xab
	}
	for _, s := range []string{fmt.Sprint(c), fmt.Sprintf("%+v", &c), fmt.Sprintf("%#v", c)} {
		if strings.Contains(s, "abab") || !strings.Contains(s, "REDACTED") {
			t.Fatalf("I_PAYLOAD printed: %s", s)
		}
	}
}

func TestEncodeHeaderAndControlReject(t *testing.T) {
	dk := datekey.DateKey{ProfileID: profile.QuicknetID, Round: 1000}
	dup := []extension.Extension{{ID: "a", Version: 1}, {ID: "a", Version: 2}}
	for name, h := range map[string]capsule.Header{
		"invalid DateKey":      {},
		"unknown policy":       {DateKey: dk, Policy: 5},
		"repeated critical":    {DateKey: dk, Critical: dup},
		"repeated noncritical": {DateKey: dk, Noncritical: dup},
		"both arrays":          {DateKey: dk, Critical: dup[:1], Noncritical: dup[1:]},
	} {
		if _, err := capsule.EncodeHeader(&h); err == nil {
			t.Errorf("%s: accepted", name)
		}
	}
	for name, c := range map[string]capsule.Control{
		"repeated critical":    {Critical: dup},
		"repeated noncritical": {Noncritical: dup},
		"both arrays":          {Critical: dup[:1], Noncritical: dup[1:]},
	} {
		for _, f := range []capsule.Format{capsule.Format1, capsule.Format2} {
			c := c
			if f == capsule.Format2 {
				c.Padding = capsule.Reforzado
			}
			if _, err := capsule.EncodeControl(&c, f); err == nil {
				t.Errorf("control %s, format %d: accepted", name, f)
			}
		}
	}
	// Keys 6 and 7 exist in format 2 only, with a defined code and an L of at
	// most L_MAX (spec §29.1, §31).
	for name, tc := range map[string]struct {
		c capsule.Control
		f capsule.Format
	}{
		"format 1 with a padding code":   {capsule.Control{Padding: capsule.Reforzado}, capsule.Format1},
		"format 1 with a payload_length": {capsule.Control{PayloadLength: 1}, capsule.Format1},
		"format 2 without a code":        {capsule.Control{}, capsule.Format2},
		"format 2 with code 3":           {capsule.Control{Padding: 3}, capsule.Format2},
		"format 2 above L_MAX":           {capsule.Control{Padding: capsule.Bloque256, PayloadLength: capsule.MaxPayloadLength + 1}, capsule.Format2},
		"format 3 without a code":        {capsule.Control{}, capsule.Format3},
		"format 4":                       {capsule.Control{Padding: capsule.Reforzado}, 4},
	} {
		if _, err := capsule.EncodeControl(&tc.c, tc.f); err == nil {
			t.Errorf("%s: accepted", name)
		}
	}
}

func marshal(t *testing.T, m map[uint64]any) []byte {
	t.Helper()
	b, err := cbortest.Marshal(m)
	if err != nil {
		t.Fatal(err)
	}
	return b
}

func ext(id string, v uint64) map[uint64]any { return map[uint64]any{0: id, 1: v} }

func extData(id string, data any) map[uint64]any { return map[uint64]any{0: id, 1: uint64(1), 2: data} }

// manyExts returns n distinct extensions in canonical order.
func manyExts(n int) []any {
	out := make([]any, n)
	for i := range out {
		out[i] = ext(fmt.Sprintf("org.example.%03d", i), 1)
	}
	return out
}

// Spec §54: extension rules shared by PUBLIC_HEADER and CONTROL_CBOR, as edits
// of the extension array at key.
func extensionRules(key uint64) map[string]func(m map[uint64]any) {
	return map[string]func(m map[uint64]any){
		"data of type text":         func(m map[uint64]any) { m[key] = []any{extData("a", "text")} },
		"data of type array":        func(m map[uint64]any) { m[key] = []any{extData("a", []any{uint64(1)})} },
		"data of type unsigned":     func(m map[uint64]any) { m[key] = []any{extData("a", uint64(7))} },
		"null data":                 func(m map[uint64]any) { m[key] = []any{extData("a", nil)} },
		"empty data h''":            func(m map[uint64]any) { m[key] = []any{extData("a", []byte{})} },
		"65 extensions":             func(m map[uint64]any) { m[key] = manyExts(65) },
		"extension_version 2^32":    func(m map[uint64]any) { m[key] = []any{ext("a", 1<<32)} },
		"extension_version 2^53-1":  func(m map[uint64]any) { m[key] = []any{ext("a", 1<<53-1)} },
		"extension map with key 3":  func(m map[uint64]any) { m[key] = []any{map[uint64]any{0: "a", 1: uint64(1), 3: []byte{1}}} },
		"extension without version": func(m map[uint64]any) { m[key] = []any{map[uint64]any{0: "a"}} },
	}
}

func TestDecodeHeaderRejects(t *testing.T) {
	dk := datekey.DateKey{ProfileID: profile.QuicknetID, Round: 1000}.Compact()
	base := func() map[uint64]any {
		return map[uint64]any{0: capsule.HeaderTypeTag, 1: uint64(1), 2: make([]byte, 16), 3: dk, 4: uint64(0)}
	}
	for _, tc := range []struct {
		name string
		edit func(m map[uint64]any)
		want error
	}{
		{"short capsule_id", func(m map[uint64]any) { m[2] = make([]byte, 15) }, datekeys.ErrNonCanonicalCBOR},
		{"invalid DateKey", func(m map[uint64]any) { m[3] = "dk1_x" }, datekeys.ErrDateKeyInvalid},
		{"type tag", func(m map[uint64]any) { m[0] = capsule.ControlTypeTag }, datekeys.ErrNonCanonicalCBOR},
		{"critical out of order", func(m map[uint64]any) { m[5] = []any{ext("b", 1), ext("a", 1)} }, datekeys.ErrNonCanonicalCBOR},
		{"noncritical repeated", func(m map[uint64]any) { m[6] = []any{ext("a", 1), ext("a", 2)} }, datekeys.ErrNonCanonicalCBOR},
		{"both arrays", func(m map[uint64]any) { m[5] = []any{ext("a", 1)}; m[6] = []any{ext("a", 1)} }, datekeys.ErrNonCanonicalCBOR},
	} {
		m := base()
		tc.edit(m)
		if _, err := capsule.DecodeHeader(marshal(t, m)); !errors.Is(err, tc.want) {
			t.Errorf("%s: got %v, want %v", tc.name, err, tc.want)
		}
	}
	for name, edit := range extensionRules(6) {
		m := base()
		edit(m)
		if _, err := capsule.DecodeHeader(marshal(t, m)); !errors.Is(err, datekeys.ErrNonCanonicalCBOR) {
			t.Errorf("%s: got %v", name, err)
		}
	}
	valid := base()
	valid[6] = append(manyExts(63), extData("z", []byte{0}))
	if h, err := capsule.DecodeHeader(marshal(t, valid)); err != nil || len(h.Noncritical) != 64 {
		t.Fatalf("valid header with 64 extensions rejected: %v", err)
	}
}

// payloadLength is key 6 of a version 2 control: L in 8 bytes, big-endian.
func payloadLength(l uint64) []byte { return binary.BigEndian.AppendUint64(nil, l) }

func TestDecodeControlRejects(t *testing.T) {
	// base is a version 1 control, of format 1, and base2 a version 2
	// control, of format 2.
	base := func() map[uint64]any {
		return map[uint64]any{0: capsule.ControlTypeTag, 1: uint64(1), 2: make([]byte, 32), 3: make([]byte, 32)}
	}
	base2 := func() map[uint64]any {
		m := base()
		m[1], m[6], m[7] = uint64(2), payloadLength(78000), uint64(capsule.Reforzado)
		return m
	}
	for _, tc := range []struct {
		name   string
		format capsule.Format
		base   func() map[uint64]any
		edit   func(m map[uint64]any)
		want   error
	}{
		{"schema version", capsule.Format1, base, func(m map[uint64]any) { m[1] = uint64(2) }, datekeys.ErrUnsupportedVersion},
		{"type tag", capsule.Format1, base, func(m map[uint64]any) { m[0] = capsule.HeaderTypeTag }, datekeys.ErrNonCanonicalCBOR},
		{"short binding", capsule.Format1, base, func(m map[uint64]any) { m[2] = make([]byte, 31) }, datekeys.ErrNonCanonicalCBOR},
		{"long identity", capsule.Format1, base, func(m map[uint64]any) { m[3] = make([]byte, 33) }, datekeys.ErrNonCanonicalCBOR},
		{"unknown key", capsule.Format1, base, func(m map[uint64]any) { m[6] = "x" }, datekeys.ErrNonCanonicalCBOR},
		{"critical repeated", capsule.Format1, base, func(m map[uint64]any) { m[4] = []any{ext("a", 1), ext("a", 2)} }, datekeys.ErrNonCanonicalCBOR},
		{"noncritical out of order", capsule.Format1, base, func(m map[uint64]any) { m[5] = []any{ext("b", 1), ext("a", 1)} }, datekeys.ErrNonCanonicalCBOR},
		{"both arrays", capsule.Format1, base, func(m map[uint64]any) { m[4] = []any{ext("a", 1)}; m[5] = []any{ext("a", 1)} }, datekeys.ErrNonCanonicalCBOR},

		// Keys 6 and 7 (spec §31): defined in version 2 only, required there.
		{"format 1: keys 6 and 7", capsule.Format1, base, func(m map[uint64]any) { m[6], m[7] = payloadLength(0), uint64(2) }, datekeys.ErrNonCanonicalCBOR},
		{"format 1: a valid version 2 control", capsule.Format1, base2, func(map[uint64]any) {}, datekeys.ErrUnsupportedVersion},
		{"format 2: a valid version 1 control", capsule.Format2, base, func(map[uint64]any) {}, datekeys.ErrUnsupportedVersion},
		{"format 2: version 1 and an unknown key", capsule.Format2, base2, func(m map[uint64]any) { m[1], m[9] = uint64(1), uint64(0) }, datekeys.ErrUnsupportedVersion},
		{"format 2: version 3", capsule.Format2, base2, func(m map[uint64]any) { m[1] = uint64(3) }, datekeys.ErrUnsupportedVersion},
		{"format 2: without key 6", capsule.Format2, base2, func(m map[uint64]any) { delete(m, 6) }, datekeys.ErrNonCanonicalCBOR},
		{"format 2: without key 7", capsule.Format2, base2, func(m map[uint64]any) { delete(m, 7) }, datekeys.ErrNonCanonicalCBOR},
		{"format 2: payload_length of 7 bytes", capsule.Format2, base2, func(m map[uint64]any) { m[6] = make([]byte, 7) }, datekeys.ErrNonCanonicalCBOR},
		{"format 2: payload_length of 9 bytes", capsule.Format2, base2, func(m map[uint64]any) { m[6] = make([]byte, 9) }, datekeys.ErrNonCanonicalCBOR},
		{"format 2: payload_length as an unsigned integer", capsule.Format2, base2, func(m map[uint64]any) { m[6] = uint64(78000) }, datekeys.ErrNonCanonicalCBOR},
		{"format 2: payload_length L_MAX + 1", capsule.Format2, base2, func(m map[uint64]any) { m[6] = payloadLength(capsule.MaxPayloadLength + 1) }, datekeys.ErrNonCanonicalCBOR},
		{"format 2: payload_length 2^64 - 1", capsule.Format2, base2, func(m map[uint64]any) { m[6] = bytes.Repeat([]byte{0xff}, 8) }, datekeys.ErrNonCanonicalCBOR},
		{"format 2: padding 0", capsule.Format2, base2, func(m map[uint64]any) { m[7] = uint64(0) }, datekeys.ErrNonCanonicalCBOR},
		{"format 2: padding 3", capsule.Format2, base2, func(m map[uint64]any) { m[7] = uint64(3) }, datekeys.ErrNonCanonicalCBOR},
		{"format 2: padding 257", capsule.Format2, base2, func(m map[uint64]any) { m[7] = uint64(257) }, datekeys.ErrNonCanonicalCBOR},
		{"format 2: padding as a byte string", capsule.Format2, base2, func(m map[uint64]any) { m[7] = []byte{2} }, datekeys.ErrNonCanonicalCBOR},
		{"format 2: unknown key 8", capsule.Format2, base2, func(m map[uint64]any) { m[8] = uint64(0) }, datekeys.ErrNonCanonicalCBOR},
	} {
		m := tc.base()
		tc.edit(m)
		if _, err := capsule.DecodeControl(marshal(t, m), tc.format); !errors.Is(err, tc.want) {
			t.Errorf("%s: got %v, want %v", tc.name, err, tc.want)
		}
	}
	for name, edit := range extensionRules(4) {
		for _, b := range []struct {
			f    capsule.Format
			base func() map[uint64]any
		}{{capsule.Format1, base}, {capsule.Format2, base2}} {
			m := b.base()
			edit(m)
			if _, err := capsule.DecodeControl(marshal(t, m), b.f); !errors.Is(err, datekeys.ErrNonCanonicalCBOR) {
				t.Errorf("format %d, %s: got %v", b.f, name, err)
			}
		}
	}
	if _, err := capsule.DecodeControl([]byte("age-encryption.org/v1\n"), capsule.Format1); !errors.Is(err, datekeys.ErrNonCanonicalCBOR) {
		t.Fatalf("garbage control: %v", err)
	}
	if _, err := capsule.DecodeControl(marshal(t, base()), 4); err == nil || datekeys.Code(err) != "" {
		t.Fatalf("format 4 is a caller error, not a normative one: %v", err)
	}
	// Format 3 exists since v0.10: a control of another schema version is
	// the normative ERR_UNSUPPORTED_VERSION (spec §31, §69.1 layer 2).
	if _, err := capsule.DecodeControl(marshal(t, base()), capsule.Format3); !errors.Is(err, datekeys.ErrUnsupportedVersion) {
		t.Fatalf("a version 1 control in format 3: %v", err)
	}

	// Valid version 2 controls, at the limits, round-trip.
	for _, c := range []capsule.Control{
		{Padding: capsule.Bloque256},
		{Padding: capsule.Reforzado, PayloadLength: 1<<32 + 1},
		{Padding: capsule.Reforzado, PayloadLength: capsule.MaxPayloadLength},
	} {
		b, err := capsule.EncodeControl(&c, capsule.Format2)
		if err != nil {
			t.Fatal(err)
		}
		back, err := capsule.DecodeControl(b, capsule.Format2)
		if err != nil || back.PayloadLength != c.PayloadLength || back.Padding != c.Padding {
			t.Fatalf("%+v: got %+v, %v", c, back, err)
		}
	}
}

// Spec §31, §55.2: a version 2 control without extensions is 103 bytes,
// whatever L and the code, so that SEALED_CONTROL_LEN does not reveal L.
func TestControlLengthIsConstant(t *testing.T) {
	for _, l := range []uint64{0, 23, 24, 255, 256, 65535, 65536, 1<<32 - 1, 1 << 32, capsule.MaxPayloadLength} {
		for _, p := range []capsule.Padding{capsule.Bloque256, capsule.Reforzado} {
			c := capsule.Control{PayloadLength: l, Padding: p}
			for i := range c.PayloadIdentity {
				c.PayloadIdentity[i], c.HeaderBinding[i] = byte(l), byte(i)
			}
			b, err := capsule.EncodeControl(&c, capsule.Format2)
			if err != nil || len(b) != 103 {
				t.Fatalf("L = %d, %s: %d bytes, %v", l, p, len(b), err)
			}
		}
	}
	// The layout of spec §31: ... 06 48 <L, 8 bytes> 07 <code>.
	b, _ := capsule.EncodeControl(&capsule.Control{PayloadLength: 1<<32 + 1, Padding: capsule.Reforzado}, capsule.Format2)
	if !bytes.HasSuffix(b, []byte{0x06, 0x48, 0, 0, 0, 1, 0, 0, 0, 1, 0x07, 0x02}) {
		t.Fatalf("layout %x", b)
	}
}

func TestHeaderLimit(t *testing.T) {
	big, err := extension.New("org.example.big", 1, bytes.Repeat([]byte{1}, capsule.MaxPublicHeaderLen))
	if err != nil {
		t.Fatal(err)
	}
	opts := past(t, 1000)
	opts.Noncritical = []extension.Extension{big}
	var dkc bytes.Buffer
	if _, err := capsule.Encrypt(&dkc, strings.NewReader("x"), opts); !errors.Is(err, datekeys.ErrIntegrity) || dkc.Len() != 0 {
		t.Fatalf("PUBLIC_HEADER above 1 MiB accepted: %v", err)
	}
	// Spec §57: the encoder refuses it with the code the decoder uses.
	if _, err := capsule.EncodeHeader(&capsule.Header{DateKey: datekey.DateKey{ProfileID: profile.QuicknetID, Round: 1000}, Noncritical: opts.Noncritical}); !errors.Is(err, datekeys.ErrIntegrity) {
		t.Fatalf("EncodeHeader above 1 MiB: %v", err)
	}
	// Spec §57: the decoder applies the limit too, whatever the framing says.
	if _, err := capsule.DecodeHeader(make([]byte, capsule.MaxPublicHeaderLen+1)); !errors.Is(err, datekeys.ErrIntegrity) {
		t.Fatalf("DecodeHeader above 1 MiB: %v", err)
	}
}

// A .dkk whose critical extension the application does not know is refused
// before any request.
func TestAccessKeyCriticalExtension(t *testing.T) {
	f := loadFixture(t, "time_and_key_portable")
	k := *f.dkk
	k.Critical = []extension.Extension{{ID: "org.example.must-understand", Version: 1}}
	o := f.openOptions(t)
	o.AccessKey = &k
	src := testkit.NewSource(f.release)
	o.Source = src
	if _, err := open(t, f.dkc, o); !errors.Is(err, datekeys.ErrExtensionCriticalUnknown) || src.Calls != 0 {
		t.Fatalf("got %v (requests %d)", err, src.Calls)
	}
	o.Extensions = extension.Set{"org.example.must-understand": {1}}
	if got, err := open(t, f.dkc, o); err != nil || !bytes.Equal(got, f.plaintext) {
		t.Fatalf("known .dkk extension rejected: %v", err)
	}
}

type brokenSeeker struct {
	*bytes.Reader
	seeks int
}

func (b *brokenSeeker) Seek(off int64, whence int) (int64, error) {
	b.seeks++
	if b.seeks > 1 {
		return 0, errors.New("seek failed")
	}
	return b.Reader.Seek(off, whence)
}

func TestCapsuleDigestSeekFailure(t *testing.T) {
	f := loadFixture(t, "time_and_key_portable")
	r := &brokenSeeker{Reader: bytes.NewReader(f.dkc)}
	if _, err := capsule.Open(context.Background(), io.Discard, r, f.openOptions(t)); err == nil {
		t.Fatal("seek failure ignored")
	}
}