OPT: add Z and SetZ methods (#1246)

* OPT: add Z and SetZ methods Fixes: #1169 Signed-off-by: Miek Gieben <[email protected]> * code review feedback Change the mask and use and not (&^) Signed-off-by: Miek Gieben <[email protected]> * Code review changes Signed-off-by: Miek Gieben <[email protected]>
miekg · May 5, 2021 · cce7f43 · cce7f43
1 parent 2a9acc8
commit cce7f43
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 0 deletions.
diff --git 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
@@ -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")
+	}
+}