rfc1035label: using a structure to hold original data

2 files changed, 18 insertions, 17 deletions

diff --git a/dhcpv4/option_domain_search.go b/dhcpv4/option_domain_search.go
index daade0a..c640c0f 100644
--- a/dhcpv4/option_domain_search.go
+++ b/dhcpv4/option_domain_search.go
@@ -11,7 +11,7 @@ import (
 
 // OptDomainSearch represents an option encapsulating a domain search list.
 type OptDomainSearch struct {
-	DomainSearch []string
+	DomainSearch *rfc1035label.Labels
 }
 
 // Code returns the option code.
@@ -22,23 +22,19 @@ func (op *OptDomainSearch) Code() OptionCode {
 // ToBytes returns a serialized stream of bytes for this option.
 func (op *OptDomainSearch) ToBytes() []byte {
 	buf := []byte{byte(op.Code()), byte(op.Length())}
-	buf = append(buf, rfc1035label.LabelsToBytes(op.DomainSearch)...)
+	buf = append(buf, op.DomainSearch.ToBytes()...)
 	return buf
 }
 
 // Length returns the length of the data portion (excluding option code an byte
-// length). 
+// length).
 func (op *OptDomainSearch) Length() int {
-	var length int
-	for _, label := range op.DomainSearch {
-		length += len(label) + 2 // add the first and the last length bytes
-	}
-	return length
+	return op.DomainSearch.Length()
 }
 
 // String returns a human-readable string.
 func (op *OptDomainSearch) String() string {
-	return fmt.Sprintf("DNS Domain Search List -> %v", op.DomainSearch)
+	return fmt.Sprintf("DNS Domain Search List -> %v", op.DomainSearch.Labels)
 }
 
 // ParseOptDomainSearch returns a new OptDomainSearch from a byte stream, or
@@ -55,9 +51,9 @@ func ParseOptDomainSearch(data []byte) (*OptDomainSearch, error) {
 	if len(data) < 2+length {
 		return nil, ErrShortByteStream
 	}
-	domainSearch, err := rfc1035label.LabelsFromBytes(data[2:length+2])
+	labels, err := rfc1035label.FromBytes(data[2 : length+2])
 	if err != nil {
 		return nil, err
 	}
-	return &OptDomainSearch{DomainSearch: domainSearch}, nil
+	return &OptDomainSearch{DomainSearch: labels}, nil
 }
diff --git a/dhcpv4/option_domain_search_test.go b/dhcpv4/option_domain_search_test.go
index 4848a83..590ccd0 100644
--- a/dhcpv4/option_domain_search_test.go
+++ b/dhcpv4/option_domain_search_test.go
@@ -3,6 +3,7 @@ package dhcpv4
 import (
 	"testing"
 
+	"github.com/insomniacslk/dhcp/rfc1035label"
 	"github.com/stretchr/testify/require"
 )
 
@@ -15,9 +16,11 @@ func TestParseOptDomainSearch(t *testing.T) {
 	}
 	opt, err := ParseOptDomainSearch(data)
 	require.NoError(t, err)
-	require.Equal(t, len(opt.DomainSearch), 2)
-	require.Equal(t, opt.DomainSearch[0], "example.com")
-	require.Equal(t, opt.DomainSearch[1], "subnet.example.org")
+	require.Equal(t, 2, len(opt.DomainSearch.Labels))
+	require.Equal(t, data[2:], opt.DomainSearch.ToBytes())
+	require.Equal(t, len(data[2:]), opt.DomainSearch.Length())
+	require.Equal(t, opt.DomainSearch.Labels[0], "example.com")
+	require.Equal(t, opt.DomainSearch.Labels[1], "subnet.example.org")
 }
 
 func TestOptDomainSearchToBytes(t *testing.T) {
@@ -28,9 +31,11 @@ func TestOptDomainSearchToBytes(t *testing.T) {
 		6, 's', 'u', 'b', 'n', 'e', 't', 7, 'e', 'x', 'a', 'm', 'p', 'l', 'e', 3, 'o', 'r', 'g', 0,
 	}
 	opt := OptDomainSearch{
-		DomainSearch: []string{
-			"example.com",
-			"subnet.example.org",
+		DomainSearch: &rfc1035label.Labels{
+			Labels: []string{
+				"example.com",
+				"subnet.example.org",
+			},
 		},
 	}
 	require.Equal(t, opt.ToBytes(), expected)