// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tka

import (
	"bytes"
	"testing"

	"github.com/fxamacker/cbor/v2"
	"github.com/google/go-cmp/cmp"
)

func TestSerialization(t *testing.T) {
	uint2 := uint(2)

	tcs := []struct {
		Name   string
		AUM    AUM
		Expect []byte
	}{
		{
			"AddKey",
			AUM{MessageKind: AUMAddKey, Key: &Key{}},
			[]byte{
				0xa3, // major type 5 (map), 3 items
				0x01, // |- major type 0 (int), value 1 (first key, MessageKind)
				0x01, // |- major type 0 (int), value 1 (first value, AUMAddKey)
				0x02, // |- major type 0 (int), value 2 (second key, PrevAUMHash)
				0xf6, // |- major type 7 (val), value null (second value, nil)
				0x03, // |- major type 0 (int), value 3 (third key, Key)
				0xa3, // |- major type 5 (map), 3 items (type Key)
				0x01, //    |- major type 0 (int), value 1 (first key, Kind)
				0x00, //    |- major type 0 (int), value 0 (first value)
				0x02, //    |- major type 0 (int), value 2 (second key, Votes)
				0x00, //    |- major type 0 (int), value 0 (first value)
				0x03, //    |- major type 0 (int), value 3 (third key, Public)
				0xf6, //    |- major type 7 (val), value null (third value, nil)
			},
		},
		{
			"RemoveKey",
			AUM{MessageKind: AUMRemoveKey, KeyID: []byte{1, 2}},
			[]byte{
				0xa3, // major type 5 (map), 3 items
				0x01, // |- major type 0 (int), value 1 (first key, MessageKind)
				0x02, // |- major type 0 (int), value 2 (first value, AUMRemoveKey)
				0x02, // |- major type 0 (int), value 2 (second key, PrevAUMHash)
				0xf6, // |- major type 7 (val), value null (second value, nil)
				0x04, // |- major type 0 (int), value 4 (third key, KeyID)
				0x42, // |- major type 2 (byte string), 2 items
				0x01, //    |- major type 0 (int), value 1 (byte 1)
				0x02, //    |- major type 0 (int), value 2 (byte 2)
			},
		},
		{
			"UpdateKey",
			AUM{MessageKind: AUMUpdateKey, Votes: &uint2, KeyID: []byte{1, 2}, Meta: map[string]string{"a": "b"}},
			[]byte{
				0xa5, // major type 5 (map), 5 items
				0x01, // |- major type 0 (int), value 1 (first key, MessageKind)
				0x05, // |- major type 0 (int), value 2 (first value, AUMUpdateKey)
				0x02, // |- major type 0 (int), value 2 (second key, PrevAUMHash)
				0xf6, // |- major type 7 (val), value null (second value, nil)
				0x04, // |- major type 0 (int), value 4 (third key, KeyID)
				0x42, // |- major type 2 (byte string), 2 items
				0x01, //    |- major type 0 (int), value 1 (byte 1)
				0x02, //    |- major type 0 (int), value 2 (byte 2)
				0x07, // |- major type 0 (int), value 7 (fourth key, Votes)
				0x02, // |- major type 0 (int), value 2 (forth value, 2)
				0x08, // |- major type 0 (int), value 8 (fifth key, Meta)
				0xa1, // |- major type 5 (map), 1 item (map[string]string type)
				0x61, //    |- major type 3 (text string), value 1 (first key, one byte long)
				0x61, //       |- byte 'a'
				0x61, //    |- major type 3 (text string), value 1 (first value, one byte long)
				0x62, //       |- byte 'b'
			},
		},
		{
			"DisableNL",
			AUM{MessageKind: AUMDisableNL, PrevAUMHash: []byte{1, 2}, DisablementSecret: []byte{3, 4}},
			[]byte{
				0xa3, // major type 5 (map), 3 items
				0x01, // |- major type 0 (int), value 1 (first key, MessageKind)
				0x03, // |- major type 0 (int), value 3 (first value, AUMDisableNL)
				0x02, // |- major type 0 (int), value 2 (second key, PrevAUMHash)
				0x42, // |- major type 2 (byte string), 2 items (second value)
				0x01, //    |- major type 0 (int), value 1 (byte 1)
				0x02, //    |- major type 0 (int), value 2 (byte 2)
				0x06, // |- major type 0 (int), value 6 (third key, DisablementSecret)
				0x42, // |- major type 2 (byte string), 2 items (third value)
				0x03, //    |- major type 0 (int), value 3 (byte 3)
				0x04, //    |- major type 0 (int), value 4 (byte 4)
			},
		},
		// TODO(tom): Uncomment once a future PR brings in the State type.
		// {
		// 	"Checkpoint",
		// 	AUM{MessageKind: AUMCheckpoint, PrevAUMHash: []byte{1, 2}, State: &State{
		// 		LastAUMHash: []byte{3, 4},
		// 		Keys: []Key{
		// 			{Kind: Key25519, Public: []byte{5, 6}},
		// 		},
		// 	}},
		// 	[]byte{
		// 		0xa3, // major type 5 (map), 3 items
		// 		0x01, // |- major type 0 (int), value 1 (first key, MessageKind)
		// 		0x06, // |- major type 0 (int), value 6 (first value, AUMCheckpoint)
		// 		0x02, // |- major type 0 (int), value 2 (second key, PrevAUMHash)
		// 		0x42, // |- major type 2 (byte string), 2 items (second value)
		// 		0x01, //    |- major type 0 (int), value 1 (byte 1)
		// 		0x02, //    |- major type 0 (int), value 2 (byte 2)
		// 		0x05, // |- major type 0 (int), value 5 (third key, State)
		// 		0xa3, // |- major type 5 (map), 3 items (third value, State type)
		// 		0x01, //    |- major type 0 (int), value 1 (first key, LastAUMHash)
		// 		0x42, //    |- major type 2 (byte string), 2 items (first value)
		// 		0x03, //       |- major type 0 (int), value 3 (byte 3)
		// 		0x04, //    |- major type 0 (int), value 4 (byte 4)
		// 		0x02, //    |- major type 0 (int), value 2 (second key, DisablementSecrets)
		// 		0xf6, //    |- major type 7 (val), value null (second value, nil)
		// 		0x03, //    |- major type 0 (int), value 3 (third key, Keys)
		// 		0x81, //    |- major type 4 (array), value 1 (one item in array)
		// 		0xa3, //       |- major type 5 (map), 3 items (Key type)
		// 		0x01, //          |- major type 0 (int), value 1 (first key, Kind)
		// 		0x01, //          |- major type 0 (int), value 1 (first value, Key25519)
		// 		0x02, //          |- major type 0 (int), value 2 (second key, Votes)
		// 		0x00, //          |- major type 0 (int), value 0 (second value, 0)
		// 		0x03, //          |- major type 0 (int), value 3 (third key, Public)
		// 		0x42, //          |- major type 2 (byte string), 2 items (third value)
		// 		0x05, //             |- major type 0 (int), value 5 (byte 5)
		// 		0x06, //             |- major type 0 (int), value 6 (byte 6)
		// 	},
		// },
		{
			"Signature",
			AUM{MessageKind: AUMAddKey, Signatures: []Signature{{KeyID: []byte{1}}}},
			[]byte{
				0xa3, // major type 5 (map), 3 items
				0x01, // |- major type 0 (int), value 1 (first key, MessageKind)
				0x01, // |- major type 0 (int), value 1 (first value, AUMAddKey)
				0x02, // |- major type 0 (int), value 2 (second key, PrevAUMHash)
				0xf6, // |- major type 7 (val), value null (second value, nil)
				0x17, // |- major type 0 (int), value 22 (third key, Signatures)
				0x81, // |- major type 4 (array), value 1 (one item in array)
				0xa2, //   |- major type 5 (map), 2 items (Signature type)
				0x01, //     |- major type 0 (int), value 1 (first key, KeyID)
				0x41, //     |- major type 2 (byte string), 1 item
				0x01, //       |- major type 0 (int), value 1 (byte 1)
				0x02, //     |- major type 0 (int), value 2 (second key, Signature)
				0xf6, //     |- major type 7 (val), value null (second value, nil)
			},
		},
	}

	for _, tc := range tcs {
		t.Run(tc.Name, func(t *testing.T) {
			data := tc.AUM.Serialize()
			if diff := cmp.Diff(tc.Expect, data); diff != "" {
				t.Errorf("serialization differs (-want, +got):\n%s", diff)
			}

			var decodedAUM AUM
			if err := cbor.Unmarshal(data, &decodedAUM); err != nil {
				t.Fatalf("Unmarshal failed: %v", err)
			}
			if diff := cmp.Diff(tc.AUM, decodedAUM); diff != "" {
				t.Errorf("unmarshalled version differs (-want, +got):\n%s", diff)
			}
		})
	}
}

func TestAUMHashes(t *testing.T) {
	// .Hash(): a hash over everything.
	// .SigHash(): a hash over everything except the signatures.
	//             The signatures are over a hash of the AUM, so
	//             using SigHash() breaks this circularity.

	aum := AUM{MessageKind: AUMAddKey, Key: &Key{Kind: Key25519}}
	sigHash1 := aum.SigHash()
	aumHash1 := aum.Hash()

	aum.Signatures = []Signature{{KeyID: []byte{1, 2, 3, 4}}}
	sigHash2 := aum.SigHash()
	aumHash2 := aum.Hash()
	if len(aum.Signatures) != 1 {
		t.Error("signature was removed by one of the hash functions")
	}

	if !bytes.Equal(sigHash1[:], sigHash1[:]) {
		t.Errorf("signature hash dependent on signatures!\n\t1 = %x\n\t2 = %x", sigHash1, sigHash2)
	}
	if bytes.Equal(aumHash1[:], aumHash2[:]) {
		t.Error("aum hash didnt change")
	}
}