Improve NTP server option parsing

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

3 files changed, 59 insertions, 79 deletions

diff --git a/dhcpv6/option_ntp_server.go b/dhcpv6/option_ntp_server.go
index a7aafb7..a5b0cfb 100644
--- a/dhcpv6/option_ntp_server.go
+++ b/dhcpv6/option_ntp_server.go
@@ -18,17 +18,20 @@ func (n *NTPSuboptionSrvAddr) Code() OptionCode {
 
 // ToBytes returns the byte serialization of the suboption.
 func (n *NTPSuboptionSrvAddr) ToBytes() []byte {
-	buf := uio.NewBigEndianBuffer(nil)
-	buf.Write16(uint16(NTPSuboptionSrvAddrCode))
-	buf.Write16(uint16(net.IPv6len))
-	buf.WriteBytes(net.IP(*n).To16())
-	return buf.Data()
+	return net.IP(*n).To16()
 }
 
 func (n *NTPSuboptionSrvAddr) String() string {
 	return fmt.Sprintf("Server Address: %s", net.IP(*n).String())
 }
 
+// FromBytes parses NTP server address from a byte slice p.
+func (n *NTPSuboptionSrvAddr) FromBytes(p []byte) error {
+	buf := uio.NewBigEndianBuffer(p)
+	*n = NTPSuboptionSrvAddr(buf.CopyN(net.IPv6len))
+	return buf.FinError()
+}
+
 // NTPSuboptionMCAddr is NTP_SUBOPTION_MC_ADDR according to RFC 5908.
 type NTPSuboptionMCAddr net.IP
 
@@ -39,19 +42,24 @@ func (n *NTPSuboptionMCAddr) Code() OptionCode {
 
 // ToBytes returns the byte serialization of the suboption.
 func (n *NTPSuboptionMCAddr) ToBytes() []byte {
-	buf := uio.NewBigEndianBuffer(nil)
-	buf.Write16(uint16(NTPSuboptionMCAddrCode))
-	buf.Write16(uint16(net.IPv6len))
-	buf.WriteBytes(net.IP(*n).To16())
-	return buf.Data()
+	return net.IP(*n).To16()
 }
 
 func (n *NTPSuboptionMCAddr) String() string {
 	return fmt.Sprintf("Multicast Address: %s", net.IP(*n).String())
 }
 
+// FromBytes parses NTP multicast address from a byte slice p.
+func (n *NTPSuboptionMCAddr) FromBytes(p []byte) error {
+	buf := uio.NewBigEndianBuffer(p)
+	*n = NTPSuboptionMCAddr(buf.CopyN(net.IPv6len))
+	return buf.FinError()
+}
+
 // NTPSuboptionSrvFQDN is NTP_SUBOPTION_SRV_FQDN according to RFC 5908.
-type NTPSuboptionSrvFQDN rfc1035label.Labels
+type NTPSuboptionSrvFQDN struct {
+	rfc1035label.Labels
+}
 
 // Code returns the suboption code.
 func (n *NTPSuboptionSrvFQDN) Code() OptionCode {
@@ -60,17 +68,16 @@ func (n *NTPSuboptionSrvFQDN) Code() OptionCode {
 
 // ToBytes returns the byte serialization of the suboption.
 func (n *NTPSuboptionSrvFQDN) ToBytes() []byte {
-	buf := uio.NewBigEndianBuffer(nil)
-	buf.Write16(uint16(NTPSuboptionSrvFQDNCode))
-	l := rfc1035label.Labels(*n)
-	buf.Write16(uint16(l.Length()))
-	buf.WriteBytes(l.ToBytes())
-	return buf.Data()
+	return n.Labels.ToBytes()
 }
 
 func (n *NTPSuboptionSrvFQDN) String() string {
-	l := rfc1035label.Labels(*n)
-	return fmt.Sprintf("Server FQDN: %s", l.String())
+	return fmt.Sprintf("Server FQDN: %s", n.Labels.String())
+}
+
+// FromBytes parses an NTP server FQDN from a byte slice p.
+func (n *NTPSuboptionSrvFQDN) FromBytes(p []byte) error {
+	return n.Labels.FromBytes(p)
 }
 
 // NTPSuboptionSrvAddr is the value of NTP_SUBOPTION_SRV_ADDR according to RFC 5908.
@@ -82,53 +89,25 @@ const (
 
 // parseNTPSuboption implements the OptionParser interface.
 func parseNTPSuboption(code OptionCode, data []byte) (Option, error) {
-	//var o Options
-	buf := uio.NewBigEndianBuffer(data)
-	length := len(data)
-	data, err := buf.ReadN(length)
-	if err != nil {
-		return nil, fmt.Errorf("failed to read %d bytes for suboption: %w", length, err)
-	}
+	var o Option
+	var err error
 	switch code {
-	case NTPSuboptionSrvAddrCode, NTPSuboptionMCAddrCode:
-		if length != net.IPv6len {
-			return nil, fmt.Errorf("invalid suboption length, want %d, got %d", net.IPv6len, length)
-		}
-		var so Option
-		switch code {
-		case NTPSuboptionSrvAddrCode:
-			sos := NTPSuboptionSrvAddr(data)
-			so = &sos
-		case NTPSuboptionMCAddrCode:
-			som := NTPSuboptionMCAddr(data)
-			so = &som
-		}
-		return so, nil
+	case NTPSuboptionSrvAddrCode:
+		var opt NTPSuboptionSrvAddr
+		err = opt.FromBytes(data)
+		o = &opt
+	case NTPSuboptionMCAddrCode:
+		var opt NTPSuboptionMCAddr
+		err = opt.FromBytes(data)
+		o = &opt
 	case NTPSuboptionSrvFQDNCode:
-		l, err := rfc1035label.FromBytes(data)
-		if err != nil {
-			return nil, fmt.Errorf("failed to parse rfc1035 labels: %w", err)
-		}
-		// TODO according to rfc3315, this label must not be compressed.
-		// Need to add support for compression detection to the
-		// `rfc1035label` package in order to do that.
-		so := NTPSuboptionSrvFQDN(*l)
-		return &so, nil
+		var opt NTPSuboptionSrvFQDN
+		err = opt.FromBytes(data)
+		o = &opt
 	default:
-		gopt := OptionGeneric{OptionCode: code, OptionData: data}
-		return &gopt, nil
-	}
-}
-
-// ParseOptNTPServer parses a sequence of bytes into an OptNTPServer object.
-func ParseOptNTPServer(data []byte) (*OptNTPServer, error) {
-	var so Options
-	if err := so.FromBytesWithParser(data, parseNTPSuboption); err != nil {
-		return nil, err
+		o = &OptionGeneric{OptionCode: code, OptionData: append([]byte(nil), data...)}
 	}
-	return &OptNTPServer{
-		Suboptions: so,
-	}, nil
+	return o, err
 }
 
 // OptNTPServer is an option NTP server as defined by RFC 5908.
@@ -141,13 +120,13 @@ func (op *OptNTPServer) Code() OptionCode {
 	return OptionNTPServer
 }
 
+func (op *OptNTPServer) FromBytes(data []byte) error {
+	return op.Suboptions.FromBytesWithParser(data, parseNTPSuboption)
+}
+
 // ToBytes returns the option serialized to bytes.
 func (op *OptNTPServer) ToBytes() []byte {
-	buf := uio.NewBigEndianBuffer(nil)
-	for _, so := range op.Suboptions {
-		buf.WriteBytes(so.ToBytes())
-	}
-	return buf.Data()
+	return op.Suboptions.ToBytes()
 }
 
 func (op *OptNTPServer) String() string {
diff --git a/dhcpv6/option_ntp_server_test.go b/dhcpv6/option_ntp_server_test.go
index 105a753..59b3c69 100644
--- a/dhcpv6/option_ntp_server_test.go
+++ b/dhcpv6/option_ntp_server_test.go
@@ -13,25 +13,22 @@ func TestSuboptionSrvAddr(t *testing.T) {
 	ip := net.ParseIP("2a03:2880:fffe:c:face:b00c:0:35")
 	so := NTPSuboptionSrvAddr(ip)
 	assert.Equal(t, NTPSuboptionSrvAddrCode, so.Code())
-	expected := append([]byte{0x00, 0x01, 0x00, 0x10}, ip...)
-	assert.Equal(t, expected, so.ToBytes())
+	assert.Equal(t, []byte(ip), so.ToBytes())
 }
 
 func TestSuboptionMCAddr(t *testing.T) {
 	ip := net.ParseIP("2a03:2880:fffe:c:face:b00c:0:35")
 	so := NTPSuboptionMCAddr(ip)
 	assert.Equal(t, NTPSuboptionMCAddrCode, so.Code())
-	expected := append([]byte{0x00, 0x02, 0x00, 0x10}, ip...)
-	assert.Equal(t, expected, so.ToBytes())
+	assert.Equal(t, []byte(ip), so.ToBytes())
 }
 
 func TestSuboptionSrvFQDN(t *testing.T) {
 	fqdn, err := rfc1035label.FromBytes([]byte("\x03ntp\x07example\x03com"))
 	require.NoError(t, err)
-	so := NTPSuboptionSrvFQDN(*fqdn)
+	so := NTPSuboptionSrvFQDN{*fqdn}
 	assert.Equal(t, NTPSuboptionSrvFQDNCode, so.Code())
-	expected := append([]byte{0x00, 0x03, 0x00, 0x10}, fqdn.ToBytes()...)
-	assert.Equal(t, expected, so.ToBytes())
+	assert.Equal(t, fqdn.ToBytes(), so.ToBytes())
 }
 
 func TestSuboptionGeneric(t *testing.T) {
@@ -40,7 +37,8 @@ func TestSuboptionGeneric(t *testing.T) {
 		0x00, 0x04, // length, 4 bytes
 		0x74, 0x65, 0x73, 0x74, // the ASCII bytes for the string "test"
 	}
-	o, err := ParseOptNTPServer(data)
+	var o OptNTPServer
+	err := o.FromBytes(data)
 	require.NoError(t, err)
 	require.Equal(t, 1, len(o.Suboptions))
 	assert.IsType(t, &OptionGeneric{}, o.Suboptions[0])
@@ -67,7 +65,8 @@ func TestParseOptNTPServer(t *testing.T) {
 	}...)
 	data = append(data, fqdn.ToBytes()...)
 
-	o, err := ParseOptNTPServer(data)
+	var o OptNTPServer
+	err = o.FromBytes(data)
 	require.NoError(t, err)
 	require.NotNil(t, o)
 	assert.Equal(t, 2, len(o.Suboptions))
@@ -78,11 +77,11 @@ func TestParseOptNTPServer(t *testing.T) {
 
 	optFQDN, ok := o.Suboptions[1].(*NTPSuboptionSrvFQDN)
 	require.True(t, ok)
-	assert.Equal(t, *fqdn, rfc1035label.Labels(*optFQDN))
+	assert.Equal(t, *fqdn, optFQDN.Labels)
 
 	var mo MessageOptions
 	assert.Nil(t, mo.NTPServers())
-	mo.Add(o)
+	mo.Add(&o)
 	// MessageOptions.NTPServers only returns server address values.
 	assert.Equal(t, []net.IP{ip}, mo.NTPServers())
 }
diff --git a/dhcpv6/options.go b/dhcpv6/options.go
index 308fc2c..467c643 100644
--- a/dhcpv6/options.go
+++ b/dhcpv6/options.go
@@ -86,7 +86,9 @@ func ParseOption(code OptionCode, optData []byte) (Option, error) {
 	case OptionFQDN:
 		opt, err = ParseOptFQDN(optData)
 	case OptionNTPServer:
-		opt, err = ParseOptNTPServer(optData)
+		var o OptNTPServer
+		err = o.FromBytes(optData)
+		opt = &o
 	case OptionBootfileURL:
 		opt, err = parseOptBootFileURL(optData)
 	case OptionBootfileParam: