v6: add DNS getter (and rename DNSRecursiveNameServer to DNS)

Signed-off-by: Chris Koch <chrisko@google.com>

7 files changed, 74 insertions, 67 deletions

diff --git a/dhcpv6/dhcpv6message.go b/dhcpv6/dhcpv6message.go
index 28853f1..400c46a 100644
--- a/dhcpv6/dhcpv6message.go
+++ b/dhcpv6/dhcpv6message.go
@@ -92,6 +92,18 @@ func (mo MessageOptions) RequestedOptions() OptionCodes {
 	return oro.OptionCodes
 }
 
+// DNS returns the DNS Recursive Name Server option as defined by RFC 3646.
+func (mo MessageOptions) DNS() []net.IP {
+	opt := mo.Options.GetOne(OptionDNSRecursiveNameServer)
+	if opt == nil {
+		return nil
+	}
+	if dns, ok := opt.(*optDNS); ok {
+		return dns.NameServers
+	}
+	return nil
+}
+
 // Message represents a DHCPv6 Message as defined by RFC 3315 Section 6.
 type Message struct {
 	MessageType   MessageType
diff --git a/dhcpv6/modifiers.go b/dhcpv6/modifiers.go
index 860071c..ee2056f 100644
--- a/dhcpv6/modifiers.go
+++ b/dhcpv6/modifiers.go
@@ -88,12 +88,7 @@ func WithIAID(iaid [4]byte) Modifier {
 
 // WithDNS adds or updates an OptDNSRecursiveNameServer
 func WithDNS(dnses ...net.IP) Modifier {
-	return func(d DHCPv6) {
-		odns := OptDNSRecursiveNameServer{
-			NameServers: append([]net.IP{}, dnses[:]...),
-		}
-		d.UpdateOption(&odns)
-	}
+	return WithOption(OptDNS(dnses...))
 }
 
 // WithDomainSearchList adds or updates an OptDomainSearchList
diff --git a/dhcpv6/modifiers_test.go b/dhcpv6/modifiers_test.go
index 2179aaa..ff8022c 100644
--- a/dhcpv6/modifiers_test.go
+++ b/dhcpv6/modifiers_test.go
@@ -58,18 +58,17 @@ func TestWithIANA(t *testing.T) {
 
 func TestWithDNS(t *testing.T) {
 	var d Message
-	WithDNS([]net.IP{
+	WithDNS(
 		net.ParseIP("fe80::1"),
 		net.ParseIP("fe80::2"),
-	}...)(&d)
+	)(&d)
 	require.Equal(t, 1, len(d.Options.Options))
-	dns := d.Options.Options[0].(*OptDNSRecursiveNameServer)
+	dns := d.Options.DNS()
 	log.Printf("DNS %+v", dns)
-	require.Equal(t, OptionDNSRecursiveNameServer, dns.Code())
-	require.Equal(t, 2, len(dns.NameServers))
-	require.Equal(t, net.ParseIP("fe80::1"), dns.NameServers[0])
-	require.Equal(t, net.ParseIP("fe80::2"), dns.NameServers[1])
-	require.NotEqual(t, net.ParseIP("fe80::1"), dns.NameServers[1])
+	require.Equal(t, 2, len(dns))
+	require.Equal(t, net.ParseIP("fe80::1"), dns[0])
+	require.Equal(t, net.ParseIP("fe80::2"), dns[1])
+	require.NotEqual(t, net.ParseIP("fe80::1"), dns[1])
 }
 
 func TestWithDomainSearchList(t *testing.T) {
diff --git a/dhcpv6/option_dns.go b/dhcpv6/option_dns.go
new file mode 100644
index 0000000..3fff699
--- /dev/null
+++ b/dhcpv6/option_dns.go
@@ -0,0 +1,47 @@
+package dhcpv6
+
+import (
+	"fmt"
+	"net"
+
+	"github.com/u-root/u-root/pkg/uio"
+)
+
+// OptDNS returns a DNS Recursive Name Server option as defined by RFC 3646.
+func OptDNS(ip ...net.IP) Option {
+	return &optDNS{NameServers: ip}
+}
+
+type optDNS struct {
+	NameServers []net.IP
+}
+
+// Code returns the option code
+func (op *optDNS) Code() OptionCode {
+	return OptionDNSRecursiveNameServer
+}
+
+// ToBytes returns the option serialized to bytes.
+func (op *optDNS) ToBytes() []byte {
+	buf := uio.NewBigEndianBuffer(nil)
+	for _, ns := range op.NameServers {
+		buf.WriteBytes(ns.To16())
+	}
+	return buf.Data()
+}
+
+func (op *optDNS) String() string {
+	return fmt.Sprintf("DNS: %v", op.NameServers)
+}
+
+// parseOptDNS builds an optDNS structure
+// from a sequence of bytes. The input data does not include option code and length
+// bytes.
+func parseOptDNS(data []byte) (*optDNS, error) {
+	var opt optDNS
+	buf := uio.NewBigEndianBuffer(data)
+	for buf.Has(net.IPv6len) {
+		opt.NameServers = append(opt.NameServers, buf.CopyN(net.IPv6len))
+	}
+	return &opt, buf.FinError()
+}
diff --git a/dhcpv6/option_dnsrecursivenameserver_test.go b/dhcpv6/option_dns_test.go
index c45e9de..34f22f7 100644
--- a/dhcpv6/option_dnsrecursivenameserver_test.go
+++ b/dhcpv6/option_dns_test.go
@@ -7,18 +7,18 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestParseOptDNSRecursiveNameServer(t *testing.T) {
+func TestParseOptDNS(t *testing.T) {
 	data := []byte{
 		0x2a, 0x03, 0x28, 0x80, 0xff, 0xfe, 0x00, 0x0c, 0xfa, 0xce, 0xb0, 0x0c, 0x00, 0x00, 0x00, 0x35,
 	}
 	expected := []net.IP{
 		net.IP(data),
 	}
-	opt, err := ParseOptDNSRecursiveNameServer(data)
+	opt, err := parseOptDNS(data)
 	require.NoError(t, err)
 	require.Equal(t, expected, opt.NameServers)
 	require.Equal(t, OptionDNSRecursiveNameServer, opt.Code())
-	require.Contains(t, opt.String(), "nameservers=[2a03:2880:fffe:c:face:b00c:0:35]", "String() should contain the correct nameservers output")
+	require.Contains(t, opt.String(), "DNS: [2a03:2880:fffe:c:face:b00c:0:35]", "String() should contain the correct nameservers output")
 }
 
 func TestOptDNSRecursiveNameServerToBytes(t *testing.T) {
@@ -27,14 +27,14 @@ func TestOptDNSRecursiveNameServerToBytes(t *testing.T) {
 	nameservers := []net.IP{ns1, ns2}
 	expected := append([]byte{}, []byte(ns1)...)
 	expected = append(expected, []byte(ns2)...)
-	opt := OptDNSRecursiveNameServer{NameServers: nameservers}
+	opt := OptDNS(nameservers...)
 	require.Equal(t, expected, opt.ToBytes())
 }
 
-func TestParseOptDNSRecursiveNameServerParseBogusNameserver(t *testing.T) {
+func TestParseOptDNSBogus(t *testing.T) {
 	data := []byte{
 		0x2a, 0x03, 0x28, 0x80, 0xff, 0xfe, 0x00, 0x0c, // invalid IPv6 address
 	}
-	_, err := ParseOptDNSRecursiveNameServer(data)
+	_, err := parseOptDNS(data)
 	require.Error(t, err, "An invalid nameserver IPv6 address should return an error")
 }
diff --git a/dhcpv6/option_dnsrecursivenameserver.go b/dhcpv6/option_dnsrecursivenameserver.go
deleted file mode 100644
index 9aaeb3b..0000000
--- a/dhcpv6/option_dnsrecursivenameserver.go
+++ /dev/null
@@ -1,46 +0,0 @@
-package dhcpv6
-
-import (
-	"fmt"
-	"net"
-
-	"github.com/u-root/u-root/pkg/uio"
-)
-
-// OptDNSRecursiveNameServer represents a OptionDNSRecursiveNameServer option
-//
-// This module defines the OptDNSRecursiveNameServer structure.
-// https://www.ietf.org/rfc/rfc3646.txt
-type OptDNSRecursiveNameServer struct {
-	NameServers []net.IP
-}
-
-// Code returns the option code
-func (op *OptDNSRecursiveNameServer) Code() OptionCode {
-	return OptionDNSRecursiveNameServer
-}
-
-// ToBytes returns the option serialized to bytes.
-func (op *OptDNSRecursiveNameServer) ToBytes() []byte {
-	buf := uio.NewBigEndianBuffer(nil)
-	for _, ns := range op.NameServers {
-		buf.WriteBytes(ns.To16())
-	}
-	return buf.Data()
-}
-
-func (op *OptDNSRecursiveNameServer) String() string {
-	return fmt.Sprintf("OptDNSRecursiveNameServer{nameservers=%v}", op.NameServers)
-}
-
-// ParseOptDNSRecursiveNameServer builds an OptDNSRecursiveNameServer structure
-// from a sequence of bytes. The input data does not include option code and length
-// bytes.
-func ParseOptDNSRecursiveNameServer(data []byte) (*OptDNSRecursiveNameServer, error) {
-	var opt OptDNSRecursiveNameServer
-	buf := uio.NewBigEndianBuffer(data)
-	for buf.Has(net.IPv6len) {
-		opt.NameServers = append(opt.NameServers, buf.CopyN(net.IPv6len))
-	}
-	return &opt, buf.FinError()
-}
diff --git a/dhcpv6/options.go b/dhcpv6/options.go
index 0a9ab0e..0db7ebd 100644
--- a/dhcpv6/options.go
+++ b/dhcpv6/options.go
@@ -66,7 +66,7 @@ func ParseOption(code OptionCode, optData []byte) (Option, error) {
 	case OptionInterfaceID:
 		opt, err = ParseOptInterfaceId(optData)
 	case OptionDNSRecursiveNameServer:
-		opt, err = ParseOptDNSRecursiveNameServer(optData)
+		opt, err = parseOptDNS(optData)
 	case OptionDomainSearchList:
 		opt, err = ParseOptDomainSearchList(optData)
 	case OptionIAPD: