dhcpv6: intro Getters for Options

Allow the Options type to have getters for each specific options, in
order to avoid users having to cast options to their specific type.

This commit introduces a getter for exactly one option: the
ClientArchType.

i.e. users can replace

archTypes := msg.GetOneOption(OptionClientArchType).(*OptClientArchType)

with

archTypes := msg.Options.ArchTypes()

Because a few message types and options embed options (normal message,
relay message, IANA/IATA option) and each have a restricted set of
options that can be used inside them, we'll introduce at least 3 or more
Options subtypes:

- MessageOptions
- RelayOptions
- IdentityOptions

Perhaps others will join at a later time, such as VendorOptions or
AddressOptions for the IAAddress options field.

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

11 files changed, 71 insertions, 67 deletions

diff --git a/dhcpv6/dhcpv6.go b/dhcpv6/dhcpv6.go
index 4f7fbbb..41b4726 100644
--- a/dhcpv6/dhcpv6.go
+++ b/dhcpv6/dhcpv6.go
@@ -204,13 +204,9 @@ func IsUsingUEFI(msg *Message) bool {
 	//               7    EFI BC
 	//               8    EFI Xscale
 	//               9    EFI x86-64
-	if opt := msg.GetOneOption(OptionClientArchType); opt != nil {
-		optat := opt.(*OptClientArchType)
-		for _, at := range optat.ArchTypes {
-			// TODO investigate if other types are appropriate
-			if at == iana.EFI_BC || at == iana.EFI_X86_64 {
-				return true
-			}
+	if archTypes := msg.Options.ArchTypes(); archTypes != nil {
+		if archTypes.Contains(iana.EFI_BC) || archTypes.Contains(iana.EFI_X86_64) {
+			return true
 		}
 	}
 	if opt := msg.GetOneOption(OptionUserClass); opt != nil {
diff --git a/dhcpv6/dhcpv6_test.go b/dhcpv6/dhcpv6_test.go
index d097095..b23a858 100644
--- a/dhcpv6/dhcpv6_test.go
+++ b/dhcpv6/dhcpv6_test.go
@@ -114,7 +114,7 @@ func TestAddOption(t *testing.T) {
 	require.Empty(t, d.Options)
 	opt := OptionGeneric{OptionCode: 0, OptionData: []byte{}}
 	d.AddOption(&opt)
-	require.Equal(t, Options{&opt}, d.Options)
+	require.Equal(t, Options{&opt}, d.Options.Options)
 }
 
 func TestToBytes(t *testing.T) {
@@ -269,15 +269,13 @@ func TestNewMessageTypeSolicit(t *testing.T) {
 
 func TestIsUsingUEFIArchTypeTrue(t *testing.T) {
 	msg := Message{}
-	opt := OptClientArchType{ArchTypes: []iana.Arch{iana.EFI_BC}}
-	msg.AddOption(&opt)
+	msg.AddOption(OptClientArchType(iana.EFI_BC))
 	require.True(t, IsUsingUEFI(&msg))
 }
 
 func TestIsUsingUEFIArchTypeFalse(t *testing.T) {
 	msg := Message{}
-	opt := OptClientArchType{ArchTypes: []iana.Arch{iana.INTEL_X86PC}}
-	msg.AddOption(&opt)
+	msg.AddOption(OptClientArchType(iana.INTEL_X86PC))
 	require.False(t, IsUsingUEFI(&msg))
 }
 
diff --git a/dhcpv6/dhcpv6message.go b/dhcpv6/dhcpv6message.go
index f4f5eb1..d67882a 100644
--- a/dhcpv6/dhcpv6message.go
+++ b/dhcpv6/dhcpv6message.go
@@ -13,11 +13,27 @@ import (
 
 const MessageHeaderSize = 4
 
+// MessageOptions are the options that may appear in a normal DHCPv6 message.
+//
+// RFC 3315 Appendix B lists the valid options that can be used.
+type MessageOptions struct {
+	Options
+}
+
+// ArchTypes returns the architecture type option.
+func (mo MessageOptions) ArchTypes() iana.Archs {
+	opt := mo.GetOne(OptionClientArchType)
+	if opt == nil {
+		return nil
+	}
+	return opt.(*optClientArchType).Archs
+}
+
 // Message represents a DHCPv6 Message as defined by RFC 3315 Section 6.
 type Message struct {
 	MessageType   MessageType
 	TransactionID TransactionID
-	Options       Options
+	Options       MessageOptions
 }
 
 var randomRead = rand.Read
@@ -250,7 +266,7 @@ func (m *Message) IsOptionRequested(requested OptionCode) bool {
 // String returns a short human-readable string for this message.
 func (m *Message) String() string {
 	return fmt.Sprintf("Message(messageType=%s transactionID=%s, %d options)",
-		m.MessageType, m.TransactionID, len(m.Options))
+		m.MessageType, m.TransactionID, len(m.Options.Options))
 }
 
 // Summary prints all options associated with this message.
@@ -263,10 +279,10 @@ func (m *Message) Summary() string {
 		m.TransactionID,
 	)
 	ret += "  options=["
-	if len(m.Options) > 0 {
+	if len(m.Options.Options) > 0 {
 		ret += "\n"
 	}
-	for _, opt := range m.Options {
+	for _, opt := range m.Options.Options {
 		ret += fmt.Sprintf("    %v\n", opt.String())
 	}
 	ret += "  ]\n"
diff --git a/dhcpv6/dhcpv6relay_test.go b/dhcpv6/dhcpv6relay_test.go
index 3e7853f..4070f54 100644
--- a/dhcpv6/dhcpv6relay_test.go
+++ b/dhcpv6/dhcpv6relay_test.go
@@ -63,9 +63,11 @@ func TestRelayMessageToBytes(t *testing.T) {
 		relayMessage: &Message{
 			MessageType:   MessageTypeSolicit,
 			TransactionID: TransactionID{0xaa, 0xbb, 0xcc},
-			Options: []Option{
-				&OptElapsedTime{
-					ElapsedTime: 0,
+			Options: MessageOptions{
+				Options: []Option{
+					&OptElapsedTime{
+						ElapsedTime: 0,
+					},
 				},
 			},
 		},
diff --git a/dhcpv6/modifiers.go b/dhcpv6/modifiers.go
index 06eb9d9..d5e24c0 100644
--- a/dhcpv6/modifiers.go
+++ b/dhcpv6/modifiers.go
@@ -63,8 +63,7 @@ func WithUserClass(uc []byte) Modifier {
 // WithArchType adds an arch type option to the packet
 func WithArchType(at iana.Arch) Modifier {
 	return func(d DHCPv6) {
-		ao := OptClientArchType{ArchTypes: []iana.Arch{at}}
-		d.AddOption(&ao)
+		d.AddOption(OptClientArchType(at))
 	}
 }
 
diff --git a/dhcpv6/modifiers_test.go b/dhcpv6/modifiers_test.go
index dfef08c..ec8ff0b 100644
--- a/dhcpv6/modifiers_test.go
+++ b/dhcpv6/modifiers_test.go
@@ -60,8 +60,8 @@ func TestWithIANA(t *testing.T) {
 		PreferredLifetime: 3600,
 		ValidLifetime:     5200,
 	})(&d)
-	require.Equal(t, 1, len(d.Options))
-	require.Equal(t, OptionIANA, d.Options[0].Code())
+	require.Equal(t, 1, len(d.Options.Options))
+	require.Equal(t, OptionIANA, d.Options.Options[0].Code())
 }
 
 func TestWithDNS(t *testing.T) {
@@ -70,8 +70,8 @@ func TestWithDNS(t *testing.T) {
 		net.ParseIP("fe80::1"),
 		net.ParseIP("fe80::2"),
 	}...)(&d)
-	require.Equal(t, 1, len(d.Options))
-	dns := d.Options[0].(*OptDNSRecursiveNameServer)
+	require.Equal(t, 1, len(d.Options.Options))
+	dns := d.Options.Options[0].(*OptDNSRecursiveNameServer)
 	log.Printf("DNS %+v", dns)
 	require.Equal(t, OptionDNSRecursiveNameServer, dns.Code())
 	require.Equal(t, 2, len(dns.NameServers))
@@ -86,8 +86,8 @@ func TestWithDomainSearchList(t *testing.T) {
 		"slackware.it",
 		"dhcp.slackware.it",
 	}...)(&d)
-	require.Equal(t, 1, len(d.Options))
-	osl := d.Options[0].(*OptDomainSearchList)
+	require.Equal(t, 1, len(d.Options.Options))
+	osl := d.Options.Options[0].(*OptDomainSearchList)
 	require.Equal(t, OptionDomainSearchList, osl.Code())
 	require.NotNil(t, osl.DomainSearchList)
 	labels := osl.DomainSearchList.Labels
diff --git a/dhcpv6/option_archtype.go b/dhcpv6/option_archtype.go
index 83ebf0e..1461056 100644
--- a/dhcpv6/option_archtype.go
+++ b/dhcpv6/option_archtype.go
@@ -2,49 +2,34 @@ package dhcpv6
 
 import (
 	"fmt"
-	"strings"
 
 	"github.com/insomniacslk/dhcp/iana"
-	"github.com/u-root/u-root/pkg/uio"
 )
 
-// OptClientArchType represents an option CLIENT_ARCH_TYPE
+// OptClientArchType represents an option CLIENT_ARCH_TYPE.
 //
 // This module defines the OptClientArchType structure.
 // https://www.ietf.org/rfc/rfc5970.txt
-type OptClientArchType struct {
-	ArchTypes []iana.Arch
+func OptClientArchType(a ...iana.Arch) Option {
+	return &optClientArchType{Archs: a}
 }
 
-func (op *OptClientArchType) Code() OptionCode {
-	return OptionClientArchType
+type optClientArchType struct {
+	iana.Archs
 }
 
-// ToBytes marshals the client arch type as defined by RFC 5970.
-func (op *OptClientArchType) ToBytes() []byte {
-	buf := uio.NewBigEndianBuffer(nil)
-	for _, at := range op.ArchTypes {
-		buf.Write16(uint16(at))
-	}
-	return buf.Data()
+func (op *optClientArchType) Code() OptionCode {
+	return OptionClientArchType
 }
 
-func (op *OptClientArchType) String() string {
-	atStrings := make([]string, 0)
-	for _, at := range op.ArchTypes {
-		atStrings = append(atStrings, at.String())
-	}
-	return fmt.Sprintf("OptClientArchType{archtype=%v}", strings.Join(atStrings, ", "))
+func (op optClientArchType) String() string {
+	return fmt.Sprintf("ClientArchType: %s", op.Archs.String())
 }
 
-// ParseOptClientArchType builds an OptClientArchType structure from
+// parseOptClientArchType builds an OptClientArchType structure from
 // a sequence of bytes The input data does not include option code and
 // length bytes.
-func ParseOptClientArchType(data []byte) (*OptClientArchType, error) {
-	var opt OptClientArchType
-	buf := uio.NewBigEndianBuffer(data)
-	for buf.Has(2) {
-		opt.ArchTypes = append(opt.ArchTypes, iana.Arch(buf.Read16()))
-	}
-	return &opt, buf.FinError()
+func parseOptClientArchType(data []byte) (*optClientArchType, error) {
+	var opt optClientArchType
+	return &opt, opt.FromBytes(data)
 }
diff --git a/dhcpv6/option_archtype_test.go b/dhcpv6/option_archtype_test.go
index b01db08..0481c1a 100644
--- a/dhcpv6/option_archtype_test.go
+++ b/dhcpv6/option_archtype_test.go
@@ -11,14 +11,14 @@ func TestParseOptClientArchType(t *testing.T) {
 	data := []byte{
 		0, 6, // EFI_IA32
 	}
-	opt, err := ParseOptClientArchType(data)
+	opt, err := parseOptClientArchType(data)
 	require.NoError(t, err)
-	require.Equal(t, iana.EFI_IA32, opt.ArchTypes[0])
+	require.Equal(t, iana.EFI_IA32, opt.Archs[0])
 }
 
 func TestParseOptClientArchTypeInvalid(t *testing.T) {
 	data := []byte{42}
-	_, err := ParseOptClientArchType(data)
+	_, err := parseOptClientArchType(data)
 	require.Error(t, err)
 }
 
@@ -26,15 +26,13 @@ func TestOptClientArchTypeParseAndToBytes(t *testing.T) {
 	data := []byte{
 		0, 8, // EFI_XSCALE
 	}
-	opt, err := ParseOptClientArchType(data)
+	opt, err := parseOptClientArchType(data)
 	require.NoError(t, err)
 	require.Equal(t, data, opt.ToBytes())
 }
 
 func TestOptClientArchType(t *testing.T) {
-	opt := OptClientArchType{
-		ArchTypes: []iana.Arch{iana.EFI_ITANIUM},
-	}
+	opt := OptClientArchType(iana.EFI_ITANIUM)
 	require.Equal(t, OptionClientArchType, opt.Code())
-	require.Contains(t, opt.String(), "archtype=EFI Itanium", "String() should contain the correct ArchType output")
+	require.Contains(t, opt.String(), "EFI Itanium", "String() should contain the correct ArchType output")
 }
diff --git a/dhcpv6/option_relaymsg_test.go b/dhcpv6/option_relaymsg_test.go
index dd82d28..887bf95 100644
--- a/dhcpv6/option_relaymsg_test.go
+++ b/dhcpv6/option_relaymsg_test.go
@@ -106,10 +106,10 @@ func TestRelayMsgParseOptRelayMsgSingleEncapsulation(t *testing.T) {
 	if tID := innerDHCP.TransactionID; tID != xid {
 		t.Fatalf("Invalid inner DHCP transaction ID. Expected 0xaabbcc, got %v", tID)
 	}
-	if len(innerDHCP.Options) != 1 {
-		t.Fatalf("Invalid inner DHCP options length. Expected 1, got %v", len(innerDHCP.Options))
+	if len(innerDHCP.Options.Options) != 1 {
+		t.Fatalf("Invalid inner DHCP options length. Expected 1, got %v", len(innerDHCP.Options.Options))
 	}
-	innerOpt := innerDHCP.Options[0]
+	innerOpt := innerDHCP.Options.Options[0]
 	eto, ok := innerOpt.(*OptElapsedTime)
 	if !ok {
 		t.Fatalf("Invalid inner option type. Expected OptElapsedTime, got %v",
diff --git a/dhcpv6/options.go b/dhcpv6/options.go
index e653c01..bf37600 100644
--- a/dhcpv6/options.go
+++ b/dhcpv6/options.go
@@ -80,7 +80,7 @@ func ParseOption(code OptionCode, optData []byte) (Option, error) {
 	case OptionBootfileParam:
 		opt, err = ParseOptBootFileParam(optData)
 	case OptionClientArchType:
-		opt, err = ParseOptClientArchType(optData)
+		opt, err = parseOptClientArchType(optData)
 	case OptionNII:
 		opt, err = ParseOptNetworkInterfaceId(optData)
 	case Option4RD:
diff --git a/iana/archtype.go b/iana/archtype.go
index 255687c..865659b 100644
--- a/iana/archtype.go
+++ b/iana/archtype.go
@@ -49,6 +49,16 @@ func (a Arch) String() string {
 // Archs represents multiple Arch values.
 type Archs []Arch
 
+// Contains returns whether b is one of the Archs in a.
+func (a Archs) Contains(b Arch) bool {
+	for _, t := range a {
+		if t == b {
+			return true
+		}
+	}
+	return false
+}
+
 // ToBytes returns the serialized option defined by RFC 4578 (DHCPv4) and RFC
 // 5970 (DHCPv6) as the Client System Architecture Option.
 func (a Archs) ToBytes() []byte {