From 4334efe80221ecae9724fae7b15bc284453618f6 Mon Sep 17 00:00:00 2001
From: Andrew Tunnell-Jones <andrew@tj.id.au>
Date: Mon, 14 Oct 2019 08:29:56 +0000
Subject: [PATCH] Update EDNS0_UL to draft-sekar-dns-ul-02 (#1028)

---
 edns.go      | 26 +++++++++++++++++++-------
 edns_test.go | 24 ++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/edns.go b/edns.go
index 8cdbfb0b..d244f7c6 100644
--- a/edns.go
+++ b/edns.go
@@ -360,7 +360,7 @@ func (e *EDNS0_COOKIE) copy() EDNS0           { return &EDNS0_COOKIE{e.Code, e.C
 // The EDNS0_UL (Update Lease) (draft RFC) option is used to tell the server to set
 // an expiration on an update RR. This is helpful for clients that cannot clean
 // up after themselves. This is a draft RFC and more information can be found at
-// http://files.dns-sd.org/draft-sekar-dns-ul.txt
+// https://tools.ietf.org/html/draft-sekar-dns-ul-02
 //
 //	o := new(dns.OPT)
 //	o.Hdr.Name = "."
@@ -370,24 +370,36 @@ func (e *EDNS0_COOKIE) copy() EDNS0           { return &EDNS0_COOKIE{e.Code, e.C
 //	e.Lease = 120 // in seconds
 //	o.Option = append(o.Option, e)
 type EDNS0_UL struct {
-	Code  uint16 // Always EDNS0UL
-	Lease uint32
+	Code     uint16 // Always EDNS0UL
+	Lease    uint32
+	KeyLease uint32
 }
 
 // Option implements the EDNS0 interface.
 func (e *EDNS0_UL) Option() uint16 { return EDNS0UL }
-func (e *EDNS0_UL) String() string { return strconv.FormatUint(uint64(e.Lease), 10) }
-func (e *EDNS0_UL) copy() EDNS0    { return &EDNS0_UL{e.Code, e.Lease} }
+func (e *EDNS0_UL) String() string { return fmt.Sprintf("%d %d", e.Lease, e.KeyLease) }
+func (e *EDNS0_UL) copy() EDNS0    { return &EDNS0_UL{e.Code, e.Lease, e.KeyLease} }
 
 // Copied: http://golang.org/src/pkg/net/dnsmsg.go
 func (e *EDNS0_UL) pack() ([]byte, error) {
-	b := make([]byte, 4)
+	var b []byte
+	if e.KeyLease == 0 {
+		b = make([]byte, 4)
+	} else {
+		b = make([]byte, 8)
+		binary.BigEndian.PutUint32(b[4:], e.KeyLease)
+	}
 	binary.BigEndian.PutUint32(b, e.Lease)
 	return b, nil
 }
 
 func (e *EDNS0_UL) unpack(b []byte) error {
-	if len(b) < 4 {
+	switch len(b) {
+	case 4:
+		e.KeyLease = 0
+	case 8:
+		e.KeyLease = binary.BigEndian.Uint32(b[4:])
+	default:
 		return ErrBuf
 	}
 	e.Lease = binary.BigEndian.Uint32(b)
diff --git a/edns_test.go b/edns_test.go
index 69682a54..d18961ed 100644
--- a/edns_test.go
+++ b/edns_test.go
@@ -109,3 +109,27 @@ func TestEDNS0_SUBNETUnpack(t *testing.T) {
 		}
 	}
 }
+
+func TestEDNS0_UL(t *testing.T) {
+	cases := []struct {
+		l  uint32
+		kl uint32
+	}{
+		{0x01234567, 0},
+		{0x76543210, 0xFEDCBA98},
+	}
+	for _, c := range cases {
+		expect := EDNS0_UL{EDNS0UL, c.l, c.kl}
+		b, err := expect.pack()
+		if err != nil {
+			t.Fatalf("failed to pack: %v", err)
+		}
+		actual := EDNS0_UL{EDNS0UL, ^uint32(0), ^uint32(0)}
+		if err := actual.unpack(b); err != nil {
+			t.Fatalf("failed to unpack: %v", err)
+		}
+		if expect != actual {
+			t.Errorf("unpacked option is different; expected %v, got %v", expect, actual)
+		}
+	}
+}