2 files changed, 39 insertions, 12 deletions
diff --git a/dhcpv4/option_userclass.go b/dhcpv4/option_userclass.go
index 8aa02d6..7c0aa22 100644
--- a/dhcpv4/option_userclass.go
+++ b/dhcpv4/option_userclass.go
@@ -55,21 +55,22 @@ func ParseOptUserClass(data []byte) (*OptUserClass, error) {
 	if code != OptionUserClassInformation {
 		return nil, fmt.Errorf("expected code %v, got %v", OptionUserClassInformation, code)
 	}
+
+	totalLength := int(data[1])
 	data = data[2:]
+	if len(data) < totalLength {
+		return nil, fmt.Errorf("ParseOptUserClass: short data: length is %d but got %d bytes",
+			totalLength, len(data))
+	}
 
-	for {
-		if len(data) == 0 {
-			break
-		}
-		if len(data) < 1 {
-			return nil, errors.New("ParseOptUserClass: short data: missing length field")
-		}
-		ucLen := int(data[0])
-		if len(data) < ucLen+1 {
-			return nil, fmt.Errorf("ParseOptUserClass: short data: less than %d bytes", ucLen+2)
+	for i := 0; i < totalLength; {
+		ucLen := int(data[i])
+		opaqueDataIndex := i + ucLen + 1
+		if len(data) < opaqueDataIndex {
+			return nil, fmt.Errorf("ParseOptUserClass: short data: less than %d bytes", opaqueDataIndex)
 		}
-		opt.UserClasses = append(opt.UserClasses, data[1:ucLen+1])
-		data = data[1+ucLen:]
+		opt.UserClasses = append(opt.UserClasses, data[i+1:opaqueDataIndex])
+		i += opaqueDataIndex
 	}
 	if len(opt.UserClasses) < 1 {
 		return nil, errors.New("ParseOptUserClass: at least one user class is required")
diff --git a/dhcpv4/option_userclass_test.go b/dhcpv4/option_userclass_test.go
index fe39b04..cbfd38b 100644
--- a/dhcpv4/option_userclass_test.go
+++ b/dhcpv4/option_userclass_test.go
@@ -64,3 +64,29 @@ func TestOptUserClassToBytesMultiple(t *testing.T) {
 	}
 	require.Equal(t, expected, data)
 }
+
+func TestParseOptUserClassLongerThanLength(t *testing.T) {
+	expected := []byte{
+		77, 10, 9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't', 'X',
+	}
+	opt, err := ParseOptUserClass(expected)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(opt.UserClasses))
+	require.Equal(t, []byte("linuxboot"), opt.UserClasses[0])
+}
+
+func TestParseOptUserClassShorterThanLength(t *testing.T) {
+	expected := []byte{
+		77, 10, 10, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
+	}
+	_, err := ParseOptUserClass(expected)
+	require.Error(t, err)
+}
+
+func TestParseOptUserClassShorterTotalLength(t *testing.T) {
+	expected := []byte{
+		77, 11, 10, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
+	}
+	_, err := ParseOptUserClass(expected)
+	require.Error(t, err)
+}