package flare

import (
	"bytes"
	"crypto/ecdh"
	"encoding/base64"
	"errors"
	"strings"
	"testing"
)

// Values copied verbatim from RFC 8291 Section 5 and Appendix A
// (https://www.rfc-editor.org/rfc/rfc8291), with the line-wrapping
// whitespace removed.
const (
	rfcPlaintext  = "V2hlbiBJIGdyb3cgdXAsIEkgd2FudCB0byBiZSBhIHdhdGVybWVsb24"
	rfcASPublic   = "BP4z9KsN6nGRTbVYI_c7VJSPQTBtkgcy27mlmlMoZIIgDll6e3vCYLocInmYWAmS6TlzAC8wEqKK6PBru3jl7A8"
	rfcASPrivate  = "yfWPiYE-n46HLnH0KqZOF1fJJU3MYrct3AELtAQ-oRw"
	rfcUAPublic   = "BCVxsr7N_eNgVRqvHtD0zTZsEc6-VV-JvLexhqUzORcxaOzi6-AYWXvTBHm4bjyPjs7Vd8pZGH6SRpkNtoIAiw4"
	rfcUAPrivate  = "q1dXpw3UpT5VOmu_cf_v6ih07Aems3njxI-JWgLcM94"
	rfcSalt       = "DGv6ra1nlYgDCS1FRnbzlw"
	rfcAuthSecret = "BTBZMqHH6r4Tts7J_aSIgg"
	rfcCEK        = "oIhVW04MRdy2XN9CiKLxTg"
	rfcNonce      = "4h_95klXJ5E_qnoN"
	rfcHeader     = "DGv6ra1nlYgDCS1FRnbzlwAAEABBBP4z9KsN6nGRTbVYI_c7VJSPQTBtkgcy27mlmlMoZIIgDll6e3vCYLocInmYWAmS6TlzAC8wEqKK6PBru3jl7A8"
	rfcCiphertext = "8pfeW0KbunFT06SuDKoJH9Ql87S1QUrdirN6GcG7sFz1y1sqLgVi1VhjVkHsUoEsbI_0LpXMuGvnzQ"
	// The request body of Section 5: header || ciphertext.
	rfcBody = "DGv6ra1nlYgDCS1FRnbzlwAAEABBBP4z9KsN6nGRTbVYI_c7VJSPQTBtkgcy27ml" +
		"mlMoZIIgDll6e3vCYLocInmYWAmS6TlzAC8wEqKK6PBru3jl7A_yl95bQpu6cVPT" +
		"pK4Mqgkf1CXztLVBSt2Ks3oZwbuwXPXLWyouBWLVWGNWQexSgSxsj_Qulcy4a-fN"
)

func b64(t *testing.T, s string) []byte {
	t.Helper()
	b, err := base64.RawURLEncoding.DecodeString(s)
	if err != nil {
		t.Fatalf("decode %q: %v", s, err)
	}
	return b
}

func TestRFC8291AppendixA(t *testing.T) {
	plaintext := b64(t, rfcPlaintext)
	if string(plaintext) != "When I grow up, I want to be a watermelon" {
		t.Fatalf("plaintext = %q", plaintext)
	}
	asKey, err := ecdh.P256().NewPrivateKey(b64(t, rfcASPrivate))
	if err != nil {
		t.Fatalf("as_private: %v", err)
	}
	if !bytes.Equal(asKey.PublicKey().Bytes(), b64(t, rfcASPublic)) {
		t.Fatal("as_private does not produce the RFC as_public")
	}
	sub := Subscription{Endpoint: "https://push.example.net/push/JzLQ3raZJfFBR0aqvOMsLrt54w4rJUsV", P256dh: rfcUAPublic, Auth: rfcAuthSecret}

	got, err := encrypt(plaintext, sub, asKey, b64(t, rfcSalt))
	if err != nil {
		t.Fatalf("encrypt: %v", err)
	}
	want := b64(t, rfcBody)
	if !bytes.Equal(got, want) {
		t.Fatalf("body mismatch\n got %s\nwant %s", base64.RawURLEncoding.EncodeToString(got), rfcBody)
	}
	header := b64(t, rfcHeader)
	if len(header) != 86 || !bytes.Equal(got[:86], header) {
		t.Fatalf("header mismatch: %s", base64.RawURLEncoding.EncodeToString(got[:86]))
	}
	if !bytes.Equal(got[86:], b64(t, rfcCiphertext)) {
		t.Fatalf("ciphertext mismatch: %s", base64.RawURLEncoding.EncodeToString(got[86:]))
	}

	// The intermediate CEK and nonce, derived from the receiver side.
	uaKey, err := ecdh.P256().NewPrivateKey(b64(t, rfcUAPrivate))
	if err != nil {
		t.Fatalf("ua_private: %v", err)
	}
	secret, err := uaKey.ECDH(asKey.PublicKey())
	if err != nil {
		t.Fatal(err)
	}
	cek, nonce, err := deriveKeys(secret, b64(t, rfcAuthSecret), b64(t, rfcSalt), b64(t, rfcUAPublic), b64(t, rfcASPublic))
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(cek, b64(t, rfcCEK)) || !bytes.Equal(nonce, b64(t, rfcNonce)) {
		t.Fatalf("CEK/NONCE = %x/%x", cek, nonce)
	}

	// The receiver decrypts the RFC body back to the plaintext.
	back, err := decryptForTest(want, rfcUAPrivate, rfcAuthSecret)
	if err != nil {
		t.Fatalf("decrypt: %v", err)
	}
	if !bytes.Equal(back, plaintext) {
		t.Fatalf("decrypted %q", back)
	}

	// Encrypt with fresh randomness also round-trips, and refuses a
	// payload above 3993 bytes.
	body, err := Encrypt(plaintext, sub)
	if err != nil {
		t.Fatal(err)
	}
	if back, err := decryptForTest(body, rfcUAPrivate, rfcAuthSecret); err != nil || !bytes.Equal(back, plaintext) {
		t.Fatalf("round trip: %q, %v", back, err)
	}
	if _, err := Encrypt([]byte(strings.Repeat("x", MaxPayloadSize+1)), sub); err != ErrPayloadTooLarge {
		t.Fatalf("oversized payload: %v", err)
	}
	if _, err := Encrypt([]byte(strings.Repeat("x", MaxPayloadSize)), sub); err != nil {
		t.Fatalf("3993-byte payload: %v", err)
	}
}

// TestEncryptRejects covers the RFC 8291 input checks: a payload over
// 3993 bytes, a p256dh that is not a 65-byte P-256 point, an auth secret
// that is not 16 bytes and non-base64url input are refused.
func TestEncryptRejects(t *testing.T) {
	sub := newTestSubscriber(t, "https://fcm.googleapis.com/fcm/send/x").sub
	if _, err := Encrypt(make([]byte, MaxPayloadSize+1), sub); !errors.Is(err, ErrPayloadTooLarge) {
		t.Fatalf("oversized payload: %v", err)
	}
	if _, err := Encrypt(make([]byte, MaxPayloadSize), sub); err != nil {
		t.Fatalf("payload at the limit: %v", err)
	}
	short := sub
	short.P256dh = base64.RawURLEncoding.EncodeToString(make([]byte, 33))
	offCurve := sub
	offCurve.P256dh = base64.RawURLEncoding.EncodeToString(append([]byte{4}, make([]byte, 64)...))
	badAuth := sub
	badAuth.Auth = base64.RawURLEncoding.EncodeToString(make([]byte, 8))
	notB64 := sub
	notB64.Auth = "***"
	for name, s := range map[string]Subscription{"short_key": short, "off_curve": offCurve, "short_auth": badAuth, "auth_not_base64": notB64} {
		if _, err := Encrypt([]byte("x"), s); err == nil {
			t.Errorf("%s accepted", name)
		}
	}
	if b, err := decodeBase64URL("YWJj"); err != nil || string(b) != "abc" {
		t.Fatalf("decode unpadded: %q %v", b, err)
	}
	if b, err := decodeBase64URL("YQ=="); err != nil || string(b) != "a" {
		t.Fatalf("decode padded: %q %v", b, err)
	}
}