From cce7f43db4c7a43498e80b890c6beaa740480e29 Mon Sep 17 00:00:00 2001
From: Miek Gieben <miek@miek.nl>
Date: Wed, 5 May 2021 15:38:29 +0200
Subject: [PATCH] OPT: add Z and SetZ methods (#1246)

* OPT: add Z and SetZ methods

Fixes: #1169

Signed-off-by: Miek Gieben <miek@miek.nl>

* code review feedback

Change the mask and use and not (&^)

Signed-off-by: Miek Gieben <miek@miek.nl>

* Code review changes

Signed-off-by: Miek Gieben <miek@miek.nl>
---
 edns.go      | 10 ++++++++++
 edns_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/edns.go b/edns.go
index 8df0ec90..844cf92b 100644
--- a/edns.go
+++ b/edns.go
@@ -151,6 +151,16 @@ func (rr *OPT) SetDo(do ...bool) {
 	}
 }
 
+// Z returns the Z part of the OPT RR as a uint16 with only the 15 least significant bits used.
+func (rr *OPT) Z() uint16 {
+	return uint16(rr.Hdr.Ttl & 0x7FFF)
+}
+
+// SetZ sets the Z part of the OPT RR, note only the 15 least significant bits of z are used.
+func (rr *OPT) SetZ(z uint16) {
+	rr.Hdr.Ttl = rr.Hdr.Ttl&^0x7FFF | uint32(z&0x7FFF)
+}
+
 // EDNS0 defines an EDNS0 Option. An OPT RR can have multiple options appended to it.
 type EDNS0 interface {
 	// Option returns the option code for the option.
diff --git a/edns_test.go b/edns_test.go
index d18961ed..ab2a3b1f 100644
--- a/edns_test.go
+++ b/edns_test.go
@@ -133,3 +133,28 @@ func TestEDNS0_UL(t *testing.T) {
 		}
 	}
 }
+
+func TestZ(t *testing.T) {
+	e := &OPT{}
+	e.Hdr.Name = "."
+	e.Hdr.Rrtype = TypeOPT
+	e.SetVersion(8)
+	e.SetDo()
+	if e.Z() != 0 {
+		t.Errorf("expected Z of 0, got %d", e.Z())
+	}
+	e.SetZ(5)
+	if e.Z() != 5 {
+		t.Errorf("expected Z of 5, got %d", e.Z())
+	}
+	e.SetZ(0xFFFF)
+	if e.Z() != 0x7FFF {
+		t.Errorf("expected Z of 0x7FFFF, got %d", e.Z())
+	}
+	if e.Version() != 8 {
+		t.Errorf("expected version to still be 8, got %d", e.Version())
+	}
+	if !e.Do() {
+		t.Error("expected DO to be set")
+	}
+}