diff options
-rw-r--r-- | dhcpv4/option_userclass.go | 78 | ||||
-rw-r--r-- | dhcpv4/option_userclass_test.go | 66 |
2 files changed, 144 insertions, 0 deletions
diff --git a/dhcpv4/option_userclass.go b/dhcpv4/option_userclass.go new file mode 100644 index 0000000..8aa02d6 --- /dev/null +++ b/dhcpv4/option_userclass.go @@ -0,0 +1,78 @@ +package dhcpv4 + +import ( + "errors" + "fmt" + "strings" +) + +// OptUserClass represents a DHCPv4 User Class option +type OptUserClass struct { + UserClasses [][]byte +} + +// Code returns the option code +func (op *OptUserClass) Code() OptionCode { + return OptionUserClassInformation +} + +// ToBytes serializes the option and returns it as a sequence of bytes +func (op *OptUserClass) ToBytes() []byte { + buf := []byte{byte(op.Code()), byte(op.Length())} + for _, uc := range op.UserClasses { + buf = append(buf, byte(len(uc))) + buf = append(buf, uc...) + } + return buf +} + +// Length returns the option length +func (op *OptUserClass) Length() int { + ret := 0 + for _, uc := range op.UserClasses { + ret += 1 + len(uc) + } + return ret +} + +func (op *OptUserClass) String() string { + ucStrings := make([]string, 0) + for _, uc := range op.UserClasses { + ucStrings = append(ucStrings, string(uc)) + } + return fmt.Sprintf("OptUserClass{userclass=[%s]}", strings.Join(ucStrings, ", ")) +} + +// ParseOptUserClass returns a new OptUserClass from a byte stream or +// error if any +func ParseOptUserClass(data []byte) (*OptUserClass, error) { + opt := OptUserClass{} + + if len(data) < 4 { + return nil, ErrShortByteStream + } + code := OptionCode(data[0]) + if code != OptionUserClassInformation { + return nil, fmt.Errorf("expected code %v, got %v", OptionUserClassInformation, code) + } + data = data[2:] + + 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) + } + opt.UserClasses = append(opt.UserClasses, data[1:ucLen+1]) + data = data[1+ucLen:] + } + if len(opt.UserClasses) < 1 { + return nil, errors.New("ParseOptUserClass: at least one user class is required") + } + return &opt, nil +} diff --git a/dhcpv4/option_userclass_test.go b/dhcpv4/option_userclass_test.go new file mode 100644 index 0000000..fe39b04 --- /dev/null +++ b/dhcpv4/option_userclass_test.go @@ -0,0 +1,66 @@ +package dhcpv4 + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestOptUserClassToBytes(t *testing.T) { + opt := OptUserClass{ + UserClasses: [][]byte{[]byte("linuxboot")}, + } + data := opt.ToBytes() + expected := []byte{ + 77, // OPTION_USER_CLASS + 10, // length + 9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't', + } + require.Equal(t, expected, data) +} + +func TestParseOptUserClassMultiple(t *testing.T) { + expected := []byte{ + 77, 15, + 9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't', + 4, 't', 'e', 's', 't', + } + opt, err := ParseOptUserClass(expected) + require.NoError(t, err) + require.Equal(t, len(opt.UserClasses), 2) + require.Equal(t, []byte("linuxboot"), opt.UserClasses[0]) + require.Equal(t, []byte("test"), opt.UserClasses[1]) +} + +func TestParseOptUserClassNone(t *testing.T) { + expected := []byte{} + _, err := ParseOptUserClass(expected) + require.Error(t, err) +} + +func TestParseOptUserClass(t *testing.T) { + expected := []byte{ + 77, 10, 9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't', + } + opt, err := ParseOptUserClass(expected) + require.NoError(t, err) + require.Equal(t, 1, len(opt.UserClasses)) + require.Equal(t, []byte("linuxboot"), opt.UserClasses[0]) +} + +func TestOptUserClassToBytesMultiple(t *testing.T) { + opt := OptUserClass{ + UserClasses: [][]byte{ + []byte("linuxboot"), + []byte("test"), + }, + } + data := opt.ToBytes() + expected := []byte{ + 77, // OPTION_USER_CLASS + 15, // length + 9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't', + 4, 't', 'e', 's', 't', + } + require.Equal(t, expected, data) +} |