package lafsession

import (
	"bytes"
	"testing"
)

func TestHandshakeAndSealOpen(t *testing.T) {
	serverID, err := GenerateIdentity()
	if err != nil {
		t.Fatal(err)
	}
	clientID, err := GenerateIdentity()
	if err != nil {
		t.Fatal(err)
	}

	serverPin := make(map[string]bool)
	serverPin[FingerprintOf(clientID.PublicBytes())] = true
	clientPin := make(map[string]bool)
	clientPin[FingerprintOf(serverID.PublicBytes())] = true

	srv, err := New(Config{
		Identity: serverID,
		Role:     RoleServer,
		AcceptPeer: func(pub []byte) bool {
			return serverPin[FingerprintOf(pub)]
		},
	})
	if err != nil {
		t.Fatal(err)
	}
	cli, err := New(Config{
		Identity: clientID,
		Role:     RoleClient,
		AcceptPeer: func(pub []byte) bool {
			return clientPin[FingerprintOf(pub)]
		},
	})
	if err != nil {
		t.Fatal(err)
	}

	ch, err := cli.StartHandshake()
	if err != nil {
		t.Fatal(err)
	}

	sh, err := srv.HandleHandshake(ch)
	if err != nil {
		t.Fatal(err)
	}
	if sh == nil {
		t.Fatal("expected ServerHello")
	}

	cf, err := cli.HandleHandshake(sh)
	if err != nil {
		t.Fatal(err)
	}
	if cf == nil {
		t.Fatal("expected ClientFinished")
	}

	_, err = srv.HandleHandshake(cf)
	if err != nil {
		t.Fatal(err)
	}
	if !srv.Established() {
		t.Fatal("server should be established after ClientFinished")
	}

	sf, err := srv.SendFinished()
	if err != nil {
		t.Fatal(err)
	}
	_, err = cli.HandleHandshake(sf)
	if err != nil {
		t.Fatal(err)
	}
	if !cli.Established() {
		t.Fatal("client should be established after ServerFinished")
	}

	laf := []byte{0x05, 0x80, 0x00, 0x00, 0x01, 0x02, 0x03, 0x04, 0xde, 0xad, 0xbe, 0xef}

	frame, err := cli.Seal(laf)
	if err != nil {
		t.Fatal(err)
	}
	got, err := srv.Open(frame)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(got, laf) {
		t.Fatalf("payload mismatch: got %x want %x", got, laf)
	}

	frame2, err := srv.Seal(laf)
	if err != nil {
		t.Fatal(err)
	}
	got2, err := cli.Open(frame2)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(got2, laf) {
		t.Fatalf("payload mismatch c2s")
	}

	_, err = srv.Open(frame)
	if err != ErrReplay {
		t.Fatalf("expected ErrReplay, got %v", err)
	}
}

func TestRejectUnknownPeer(t *testing.T) {
	serverID, _ := GenerateIdentity()
	clientID, _ := GenerateIdentity()

	srv, _ := New(Config{
		Identity: serverID,
		Role:     RoleServer,
		AcceptPeer: func(pub []byte) bool {
			return false
		},
	})
	cli, _ := New(Config{
		Identity:   clientID,
		Role:       RoleClient,
		AcceptPeer: func(pub []byte) bool { return true },
	})

	ch, _ := cli.StartHandshake()
	_, err := srv.HandleHandshake(ch)
	if err != ErrIdentityRejected {
		t.Fatalf("expected ErrIdentityRejected, got %v", err)
	}
}