package tls

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/hex"
	"encoding/pem"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/xtls/xray-core/common"
	"github.com/xtls/xray-core/common/protocol/tls/cert"
)

func TestCalculateCertHash(t *testing.T) {
	const Single = `-----BEGIN CERTIFICATE-----
[encoded data, 4615 characters]
7ykUutQNUALq8dQwoTnK
-----END CERTIFICATE-----

`
	t.Run("singlepublickey", func(t *testing.T) {
		block, _ := pem.Decode([]byte(Single))
		cert, err := x509.ParseCertificate(block.Bytes)
		assert.Equal(t, err, nil)
		hash := GenerateCertHash(cert)
		fingerprint, _ := hex.DecodeString("ae243d668ec9c7f74a0dcd1ad21c6676b4efe30c39728934b362093af886bf77")
		assert.Equal(t, fingerprint, hash)
	})
}

func TestVerifyPeerLeafCert(t *testing.T) {
	leafCert, leafHash := cert.MustGenerate(nil, cert.DNSNames("example.com"))
	leaf := common.Must2(x509.ParseCertificate(leafCert.Certificate))

	r := &RandCarrier{
		Config: &tls.Config{
			ServerName: "example.com",
		},
		PinnedPeerCertSha256: [][]byte{leafHash[:]},
	}

	rawCerts := [][]byte{leaf.Raw}
	err := r.verifyPeerCert(rawCerts, nil)
	if err != nil {
		t.Fatal("expected to verify leaf cert signed by pinned CA, but got error:", err)
	}

	// make the pinned hash incorrect
	r.PinnedPeerCertSha256[0][0] += 1
	err = r.verifyPeerCert(rawCerts, nil)
	if err == nil {
		t.Fatal("expected to fail verifying leaf cert with incorrect pinned CA hash, but got no error")
	}
}

func TestVerifyPeerCACert(t *testing.T) {
	caCert, caHash := cert.MustGenerate(nil, cert.Authority(true), cert.KeyUsage(x509.KeyUsageCertSign))
	ca := common.Must2(x509.ParseCertificate(caCert.Certificate))

	leafCert, _ := cert.MustGenerate(caCert, cert.DNSNames("example.com"))
	leaf := common.Must2(x509.ParseCertificate(leafCert.Certificate))

	r := &RandCarrier{
		Config: &tls.Config{
			ServerName: "example.com",
		},
		PinnedPeerCertSha256: [][]byte{caHash[:]},
	}

	rawCerts := [][]byte{leaf.Raw, ca.Raw}
	err := r.verifyPeerCert(rawCerts, nil)
	if err != nil {
		t.Fatal("expected to verify leaf cert signed by pinned CA, but got error:", err)
	}

	// make the pinned hash incorrect
	r.PinnedPeerCertSha256[0][0] += 1
	err = r.verifyPeerCert(rawCerts, nil)
	if err == nil {
		t.Fatal("expected to fail verifying leaf cert with incorrect pinned CA hash, but got no error")
	}
}
