package dnssec

import (
	"crypto"
	"errors"
	"io"
	"testing"
	"time"

	"github.com/coredns/coredns/plugin/pkg/cache"
	"github.com/coredns/coredns/plugin/test"
	"github.com/coredns/coredns/request"

	"github.com/miekg/dns"
)

func TestZoneSigning(t *testing.T) {
	d, rm1, rm2 := newDnssec(t, []string{"miek.nl."})
	defer rm1()
	defer rm2()

	m := testMsg()
	state := request.Request{Req: m, Zone: "miek.nl."}

	m = d.Sign(state, time.Now().UTC(), server)
	if !section(m.Answer, 1) {
		t.Errorf("Answer section should have 1 RRSIG")
	}
	if !section(m.Ns, 1) {
		t.Errorf("Authority section should have 1 RRSIG")
	}
}

func TestZoneSigningDouble(t *testing.T) {
	d, rm1, rm2 := newDnssec(t, []string{"miek.nl."})
	defer rm1()
	defer rm2()

	fPriv1, rmPriv1, _ := test.TempFile(".", privKey1)
	fPub1, rmPub1, _ := test.TempFile(".", pubKey1)
	defer rmPriv1()
	defer rmPub1()

	key1, err := ParseKeyFile(fPub1, fPriv1)
	if err != nil {
		t.Fatalf("Failed to parse key: %v\n", err)
	}
	d.keys = append(d.keys, key1)

	m := testMsg()
	state := request.Request{Req: m, Zone: "miek.nl."}
	m = d.Sign(state, time.Now().UTC(), server)
	if !section(m.Answer, 2) {
		t.Errorf("Answer section should have 2 RRSIGs")
	}
	if !section(m.Ns, 2) {
		t.Errorf("Authority section should have 2 RRSIGs")
	}
}

// TestSigningDifferentZone tests if a key for miek.nl and be used for example.org.
func TestSigningDifferentZone(t *testing.T) {
	fPriv, rmPriv, _ := test.TempFile(".", privKey)
	fPub, rmPub, _ := test.TempFile(".", pubKey)
	defer rmPriv()
	defer rmPub()

	key, err := ParseKeyFile(fPub, fPriv)
	if err != nil {
		t.Fatalf("Failed to parse key: %v\n", err)
	}

	m := testMsgEx()
	state := request.Request{Req: m, Zone: "example.org."}
	c := cache.New[[]dns.RR](defaultCap)
	d := New([]string{"example.org."}, []*DNSKEY{key}, false, nil, c)
	m = d.Sign(state, time.Now().UTC(), server)
	if !section(m.Answer, 1) {
		t.Errorf("Answer section should have 1 RRSIG")
		t.Logf("%+v\n", m)
	}
	if !section(m.Ns, 1) {
		t.Errorf("Authority section should have 1 RRSIG")
		t.Logf("%+v\n", m)
	}
}

func TestSigningCname(t *testing.T) {
	d, rm1, rm2 := newDnssec(t, []string{"miek.nl."})
	defer rm1()
	defer rm2()

	m := testMsgCname()
	state := request.Request{Req: m, Zone: "miek.nl."}
	m = d.Sign(state, time.Now().UTC(), server)
	if !section(m.Answer, 1) {
		t.Errorf("Answer section should have 1 RRSIG")
	}
}

func TestSigningDname(t *testing.T) {
	d, rm1, rm2 := newDnssec(t, []string{"miek.nl."})
	defer rm1()
	defer rm2()

	m := testMsgDname()
	state := request.Request{Req: m, Zone: "miek.nl."}
	// We sign *everything* we see, also the synthesized CNAME.
	m = d.Sign(state, time.Now().UTC(), server)
	if !section(m.Answer, 3) {
		t.Errorf("Answer section should have 3 RRSIGs")
	}
}

func TestSigningEmpty(t *testing.T) {
	d, rm1, rm2 := newDnssec(t, []string{"miek.nl."})
	defer rm1()
	defer rm2()

	m := testEmptyMsg()
	m.SetQuestion("a.miek.nl.", dns.TypeA)
	state := request.Request{Req: m, Zone: "miek.nl."}
	m = d.Sign(state, time.Now().UTC(), server)
	if !section(m.Ns, 2) {
		t.Errorf("Authority section should have 2 RRSIGs")
	}
}

func TestDelegationSigned(t *testing.T) {
	d, rm1, rm2 := newDnssec(t, []string{"miek.nl."})
	defer rm1()
	defer rm2()

	m := testMsgDelegationSigned()
	m.SetQuestion("sub.miek.nl.", dns.TypeNS)
	state := request.Request{Req: m, Zone: "miek.nl."}
	m = d.Sign(state, time.Now().UTC(), server)
	if !section(m.Ns, 1) {
		t.Errorf("Authority section should have 1 RRSIGs")
	}
	if !section(m.Extra, 0) {
		t.Error("Extra section should not have RRSIGs")
	}
}

func TestDelegationUnSigned(t *testing.T) {
	d, rm1, rm2 := newDnssec(t, []string{"miek.nl."})
	defer rm1()
	defer rm2()

	m := testMsgDelegationUnSigned()
	m.SetQuestion("sub.miek.nl.", dns.TypeNS)
	state := request.Request{Req: m, Zone: "miek.nl."}
	m = d.Sign(state, time.Now().UTC(), server)
	if !section(m.Ns, 1) {
		t.Errorf("Authority section should have 1 RRSIG")
	}
	if !section(m.Extra, 0) {
		t.Error("Extra section should not have RRSIG")
	}
	var nsec *dns.NSEC
	var rrsig *dns.RRSIG
	for _, r := range m.Ns {
		if r.Header().Rrtype == dns.TypeNSEC {
			nsec = r.(*dns.NSEC)
		}
		if r.Header().Rrtype == dns.TypeRRSIG {
			rrsig = r.(*dns.RRSIG)
		}
	}
	if nsec == nil {
		t.Error("Authority section should hold a NSEC record")
	}
	if rrsig.TypeCovered != dns.TypeNSEC {
		t.Errorf("RRSIG should cover type %s, got %s",
			dns.TypeToString[dns.TypeNSEC], dns.TypeToString[rrsig.TypeCovered])
	}
	if !correctNsecForDS(nsec) {
		t.Error("NSEC as invalid TypeBitMap for a DS")
	}
}

func section(rss []dns.RR, nrSigs int) bool {
	i := 0
	for _, r := range rss {
		if r.Header().Rrtype == dns.TypeRRSIG {
			i++
		}
	}
	return nrSigs == i
}

func testMsg() *dns.Msg {
	// don't care about the message header
	return &dns.Msg{
		Answer: []dns.RR{test.MX("miek.nl.	1703	IN	MX	1 aspmx.l.google.com.")},
		Ns:     []dns.RR{test.NS("miek.nl.	1703	IN	NS	omval.tednet.nl.")},
	}
}
func testMsgEx() *dns.Msg {
	return &dns.Msg{
		Answer: []dns.RR{test.MX("example.org.	1703	IN	MX	1 aspmx.l.google.com.")},
		Ns:     []dns.RR{test.NS("example.org.	1703	IN	NS	omval.tednet.nl.")},
	}
}

func testMsgCname() *dns.Msg {
	return &dns.Msg{
		Answer: []dns.RR{test.CNAME("www.miek.nl.	1800	IN	CNAME	a.miek.nl.")},
	}
}

func testMsgDname() *dns.Msg {
	return &dns.Msg{
		Answer: []dns.RR{
			test.CNAME("a.dname.miek.nl.	1800	IN	CNAME	a.test.miek.nl."),
			test.A("a.test.miek.nl.	1800	IN	A	139.162.196.78"),
			test.DNAME("dname.miek.nl.	1800	IN	DNAME	test.miek.nl."),
		},
	}
}

func testMsgDelegationSigned() *dns.Msg {
	return &dns.Msg{
		Ns: []dns.RR{
			test.NS("sub.miek.nl.	1800	IN	NS	ns1.sub.miek.nl."),
			test.DS("sub." + dsKey),
		},
		Extra: []dns.RR{
			test.A("ns1.sub.miek.nl.	1800	IN	A	192.0.2.1"),
		},
	}
}

func testMsgDelegationUnSigned() *dns.Msg {
	return &dns.Msg{
		Ns: []dns.RR{
			test.NS("sub.miek.nl.	1800	IN	NS	ns1.sub.miek.nl."),
		},
		Extra: []dns.RR{
			test.A("ns1.sub.miek.nl.	1800	IN	A	192.0.2.1"),
		},
	}
}

func testEmptyMsg() *dns.Msg {
	// don't care about the message header
	return &dns.Msg{
		Ns: []dns.RR{test.SOA("miek.nl.	1800	IN	SOA	ns.miek.nl. dnsmaster.miek.nl. 2017100301 200 100 604800 3600")},
	}
}

// errSigner is a crypto.Signer that always returns an error. Used to simulate
// signing failures in tests without panicking on nil keys.
type errSigner struct {
	pub crypto.PublicKey
}

func (e *errSigner) Public() crypto.PublicKey { return e.pub }
func (e *errSigner) Sign(_ io.Reader, _ []byte, _ crypto.SignerOpts) ([]byte, error) {
	return nil, errors.New("simulated signing failure")
}

// TestSignReturnsNilOnError verifies that when a signing operation fails mid-way
// through multiple keys, sign() returns (nil, error) rather than (partial_sigs, error).
// Before the fix, the inflight function returned the partially-accumulated sigs slice
// alongside the error. While callers checked err before using the sigs, returning
// partial results from an error path is incorrect and could cause a nil-assertion
// panic if the error guard were ever removed.
func TestSignReturnsNilOnError(t *testing.T) {
	// Get a valid key that will sign successfully.
	k1, rm1, rm2 := newKey(t)
	defer rm1()
	defer rm2()

	// Create a second key that will fail during signing.
	brokenKey := &DNSKEY{
		K:   dns.Copy(k1.K).(*dns.DNSKEY),
		s:   &errSigner{pub: k1.s.Public()},
		tag: k1.tag + 1,
	}

	c := cache.New[[]dns.RR](defaultCap)
	// k1 succeeds, brokenKey fails — sign() should return nil, not k1's partial sig.
	d := New([]string{"miek.nl."}, []*DNSKEY{k1, brokenKey}, false, nil, c)

	m := testMsg()
	incep, expir := incepExpir(time.Now().UTC())
	sigs, err := d.sign(m.Answer, "miek.nl.", 1703, incep, expir, server)

	if err == nil {
		t.Fatal("Expected error from broken key, got nil")
	}
	if sigs != nil {
		t.Errorf("Expected nil sigs on signing error, got %d sig(s)", len(sigs))
	}
}

func newDnssec(t *testing.T, zones []string) (Dnssec, func(), func()) {
	t.Helper()
	k, rm1, rm2 := newKey(t)
	c := cache.New[[]dns.RR](defaultCap)
	d := New(zones, []*DNSKEY{k}, false, nil, c)
	return d, rm1, rm2
}

func newKey(t *testing.T) (*DNSKEY, func(), func()) {
	t.Helper()
	fPriv, rmPriv, _ := test.TempFile(".", privKey)
	fPub, rmPub, _ := test.TempFile(".", pubKey)

	key, err := ParseKeyFile(fPub, fPriv)
	if err != nil {
		t.Fatalf("Failed to parse key: %v\n", err)
	}
	return key, rmPriv, rmPub
}

const (
	pubKey  = `miek.nl. IN DNSKEY 257 3 13 0J8u0XJ9GNGFEBXuAmLu04taHG4BXPP3gwhetiOUMnGA+x09nqzgF5IY OyjWB7N3rXqQbnOSILhH1hnuyh7mmA==`
	privKey = `Private-key-format: v1.3
Algorithm: 13 (ECDSAP256SHA256)
PrivateKey: /4BZk8AFvyW5hL3cOLSVxIp1RTqHSAEloWUxj86p3gs=
Created: 20160423195532
Publish: 20160423195532
Activate: 20160423195532
`
	dsKey    = `miek.nl. IN DS 18512 13 2 D4E806322598BC97A003EF1ACDFF352EEFF7B42DBB0D41B8224714C36AEF08D9`
	pubKey1  = `example.org. IN DNSKEY 257 3 13 tVRWNSGpHZbCi7Pr7OmbADVUO3MxJ0Lb8Lk3o/HBHqCxf5K/J50lFqRa 98lkdAIiFOVRy8LyMvjwmxZKwB5MNw==`
	privKey1 = `Private-key-format: v1.3
Algorithm: 13 (ECDSAP256SHA256)
PrivateKey: i8j4OfDGT8CQt24SDwLz2hg9yx4qKOEOh1LvbAuSp1c=
Created: 20160423211746
Publish: 20160423211746
Activate: 20160423211746
`
)