add IsNetboot and IsUsingUEFI (#70)

2 files changed, 111 insertions, 0 deletions

diff --git a/dhcpv6/utils.go b/dhcpv6/utils.go
new file mode 100644
index 0000000..81ebaae
--- /dev/null
+++ b/dhcpv6/utils.go
@@ -0,0 +1,60 @@
+package dhcpv6
+
+import (
+	"strings"
+)
+
+// IsNetboot function takes a DHCPv6 message and returns true if the machine
+// is trying to netboot. It checks if "boot file" is one of the requested
+// options, which is useful for SOLICIT/REQUEST packet types, it also checks
+// if the "boot file" option is included in the packet, which is useful for
+// ADVERTISE/REPLY packet.
+func IsNetboot(msg DHCPv6) bool {
+	for _, optoro := range msg.GetOption(OPTION_ORO) {
+		for _, o := range optoro.(*OptRequestedOption).RequestedOptions() {
+			if o == OPT_BOOTFILE_URL {
+				return true
+			}
+		}
+	}
+	if optbf := msg.GetOneOption(OPT_BOOTFILE_URL); optbf != nil {
+		return true
+	}
+	return false
+}
+
+// IsUsingUEFI function takes a DHCPv6 message and returns true if
+// the machine trying to netboot is using UEFI of false if it is not.
+func IsUsingUEFI(msg DHCPv6) bool {
+	// RFC 4578 says:
+	// As of the writing of this document, the following pre-boot
+	//    architecture types have been requested.
+	//             Type   Architecture Name
+	//             ----   -----------------
+	//               0    Intel x86PC
+	//               1    NEC/PC98
+	//               2    EFI Itanium
+	//               3    DEC Alpha
+	//               4    Arc x86
+	//               5    Intel Lean Client
+	//               6    EFI IA32
+	//               7    EFI BC
+	//               8    EFI Xscale
+	//               9    EFI x86-64
+	if opt := msg.GetOneOption(OPTION_CLIENT_ARCH_TYPE); opt != nil {
+		optat := opt.(*OptClientArchType)
+		// TODO investigate if other types are appropriate
+		if optat.ArchType == EFI_BC || optat.ArchType == EFI_X86_64 {
+			return true
+		}
+	}
+	if opt := msg.GetOneOption(OPTION_USER_CLASS); opt != nil {
+		optuc := opt.(*OptUserClass)
+		for _, uc := range optuc.UserClasses {
+			if strings.Contains(string(uc), "EFI") {
+				return true
+			}
+		}
+	}
+	return false
+}
diff --git a/dhcpv6/utils_test.go b/dhcpv6/utils_test.go
new file mode 100644
index 0000000..2373691
--- /dev/null
+++ b/dhcpv6/utils_test.go
@@ -0,0 +1,51 @@
+package dhcpv6
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestIsNetboot(t *testing.T) {
+	msg1 := DHCPv6Message{}
+	require.False(t, IsNetboot(&msg1))
+
+	msg2 := DHCPv6Message{}
+	optro := OptRequestedOption{}
+	optro.AddRequestedOption(OPT_BOOTFILE_URL)
+	msg2.AddOption(&optro)
+	require.True(t, IsNetboot(&msg2))
+
+	msg3 := DHCPv6Message{}
+	optbf := OptBootFileURL{}
+	msg3.AddOption(&optbf)
+	require.True(t, IsNetboot(&msg3))
+}
+
+func TestIsUsingUEFIArchTypeTrue(t *testing.T) {
+	msg := DHCPv6Message{}
+	opt := OptClientArchType{ArchType: EFI_BC}
+	msg.AddOption(&opt)
+	require.True(t, IsUsingUEFI(&msg))
+}
+
+func TestIsUsingUEFIArchTypeFalse(t *testing.T) {
+	msg := DHCPv6Message{}
+	opt := OptClientArchType{ArchType: INTEL_X86PC}
+	msg.AddOption(&opt)
+	require.False(t, IsUsingUEFI(&msg))
+}
+
+func TestIsUsingUEFIUserClassTrue(t *testing.T) {
+	msg := DHCPv6Message{}
+	opt := OptUserClass{UserClasses: [][]byte{[]byte("ipxeUEFI")}}
+	msg.AddOption(&opt)
+	require.True(t, IsUsingUEFI(&msg))
+}
+
+func TestIsUsingUEFIUserClassFalse(t *testing.T) {
+	msg := DHCPv6Message{}
+	opt := OptUserClass{UserClasses: [][]byte{[]byte("ipxeLegacy")}}
+	msg.AddOption(&opt)
+	require.False(t, IsUsingUEFI(&msg))
+}