Merge pull request #95 from terinjokes/patches/attribute-encoder

netlink: add AttributeEncoder type

attribute.go

@@ -314,3 +314,141 @@ func (ad *AttributeDecoder) Do(fn func(b []byte) error) {
 		ad.err = err
 	}
 }
+
+// An AttributeEncoder provides a safe way to encode attributes.
+//
+// It is recommended to use an AttributeEncoder where possible instead of
+// calling MarshalAttributes or using package nlenc directly.
+//
+// Errors from intermediate encoding steps are returned in the call to
+// Encode.
+type AttributeEncoder struct {
+	// ByteOrder defines a specific byte order to use when processing integer
+	// attributes.  ByteOrder should be set immediately after creating the
+	// AttributeEncoder: before any attributes are encoded.
+	//
+	// If not set, the native byte order will be used.
+	ByteOrder binary.ByteOrder
+
+	attrs []Attribute
+	err   error
+}
+
+// NewAttributeEncoder creates an AttributeEncoder that encodes Attributes.
+func NewAttributeEncoder() *AttributeEncoder {
+	return &AttributeEncoder{
+		ByteOrder: nlenc.NativeEndian(),
+	}
+}
+
+// Uint8 encodes the uint8 data for the specified field.
+func (ae *AttributeEncoder) Uint8(t uint16, v uint8) {
+	if ae.err != nil {
+		return
+	}
+
+	ae.attrs = append(ae.attrs, Attribute{
+		Type: t,
+		Data: []byte{v},
+	})
+}
+
+// Uint16 encodes the uint16 data for the specified field.
+func (ae *AttributeEncoder) Uint16(t uint16, v uint16) {
+	if ae.err != nil {
+		return
+	}
+
+	b := make([]byte, 2)
+	ae.ByteOrder.PutUint16(b, v)
+
+	ae.attrs = append(ae.attrs, Attribute{
+		Type: t,
+		Data: b,
+	})
+}
+
+// Uint32 encodes the uint32 data for the specified field.
+func (ae *AttributeEncoder) Uint32(t uint16, v uint32) {
+	if ae.err != nil {
+		return
+	}
+
+	b := make([]byte, 4)
+	ae.ByteOrder.PutUint32(b, v)
+
+	ae.attrs = append(ae.attrs, Attribute{
+		Type: t,
+		Data: b,
+	})
+}
+
+// Uint64 encodes the uint64 data for the specified field.
+func (ae *AttributeEncoder) Uint64(t uint16, v uint64) {
+	if ae.err != nil {
+		return
+	}
+
+	b := make([]byte, 8)
+	ae.ByteOrder.PutUint64(b, v)
+
+	ae.attrs = append(ae.attrs, Attribute{
+		Type: t,
+		Data: b,
+	})
+}
+
+// String encodes string s as a null-terminated string for the specified field.
+func (ae *AttributeEncoder) String(t uint16, s string) {
+	if ae.err != nil {
+		return
+	}
+
+	ae.attrs = append(ae.attrs, Attribute{
+		Type: t,
+		Data: nlenc.Bytes(s),
+	})
+}
+
+// Bytes embeds raw byte data in the specified field.
+func (ae *AttributeEncoder) Bytes(t uint16, b []byte) {
+	if ae.err != nil {
+		return
+	}
+
+	ae.attrs = append(ae.attrs, Attribute{
+		Type: t,
+		Data: b,
+	})
+}
+
+// Do is a general purpose function to encode arbitrary data into a field.
+//
+// Do is especially helpful in encoding nested attributes, attribute arrays,
+// or encoding arbitrary types (such as C structures) which don't fit cleanly
+// into an unsigned integer value.
+func (ae *AttributeEncoder) Do(t uint16, fn func() ([]byte, error)) {
+	if ae.err != nil {
+		return
+	}
+
+	b, err := fn()
+	if err != nil {
+		ae.err = err
+		return
+	}
+
+	ae.attrs = append(ae.attrs, Attribute{
+		Type: t,
+		Data: b,
+	})
+}
+
+// Encode returns the encoded bytes representing the attributes.
+func (ae *AttributeEncoder) Encode() ([]byte, error) {
+	if ae.err != nil {
+		return nil, ae.err
+	}
+
+	return MarshalAttributes(ae.attrs)
+}

attribute_test.go

@@ -711,3 +711,157 @@ func adEndianTest(order binary.ByteOrder) func(ad *AttributeDecoder) {
 		}
 	}
 }
+
+func TestAttributeEncoderError(t *testing.T) {
+	skipBigEndian(t)
+
+	tests := []struct {
+		name string
+		fn   func(ae *AttributeEncoder)
+	}{
+		{
+			name: "do",
+			fn: func(ae *AttributeEncoder) {
+				ae.Do(1, func() ([]byte, error) {
+					return nil, errors.New("testing error")
+				})
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ae := NewAttributeEncoder()
+			tt.fn(ae)
+			_, err := ae.Encode()
+
+			if err == nil {
+				t.Fatal("expected an error, but none occurred")
+			}
+		})
+	}
+}
+
+func TestAttributeEncoderOK(t *testing.T) {
+	skipBigEndian(t)
+
+	tests := []struct {
+		name   string
+		attrs  []Attribute
+		endian binary.ByteOrder
+		fn     func(ae *AttributeEncoder)
+	}{
+		{
+			name:  "empty",
+			attrs: nil,
+			fn: func(_ *AttributeEncoder) {
+			},
+		},
+		{
+			name:  "uint native endian",
+			attrs: adEndianAttrs(nlenc.NativeEndian()),
+			fn:    aeEndianTest(nlenc.NativeEndian()),
+		},
+		{
+			name:   "uint little endian",
+			attrs:  adEndianAttrs(binary.LittleEndian),
+			endian: binary.LittleEndian,
+			fn:     aeEndianTest(binary.LittleEndian),
+		},
+		{
+			name:   "uint big endian",
+			attrs:  adEndianAttrs(binary.BigEndian),
+			endian: binary.BigEndian,
+			fn:     aeEndianTest(binary.BigEndian),
+		},
+		{
+			name: "string",
+			attrs: []Attribute{{
+				Type: 1,
+				Data: nlenc.Bytes("hello netlink"),
+			}},
+			fn: func(ae *AttributeEncoder) {
+				ae.String(1, "hello netlink")
+			},
+		},
+		{
+			name: "byte",
+			attrs: []Attribute{
+				{
+					Type: 1,
+					Data: []byte{0xde, 0xad},
+				},
+			},
+			fn: func(ae *AttributeEncoder) {
+				ae.Bytes(1, []byte{0xde, 0xad})
+			},
+		},
+		{
+			name: "do",
+			attrs: []Attribute{
+				// Arbitrary C-like structure.
+				{
+					Type: 1,
+					Data: []byte{0xde, 0xad, 0xbe},
+				},
+				// Nested attributes.
+				{
+					Type: 2,
+					Data: func() []byte {
+						b, err := MarshalAttributes([]Attribute{{
+							Type: 2,
+							Data: nlenc.Uint16Bytes(2),
+						}})
+						if err != nil {
+							panicf("failed to marshal test attributes: %v", err)
+						}
+
+						return b
+					}(),
+				},
+			},
+			fn: func(ae *AttributeEncoder) {
+				ae.Do(1, func() ([]byte, error) {
+					return []byte{0xde, 0xad, 0xbe}, nil
+				})
+				ae.Do(2, func() ([]byte, error) {
+					ae1 := NewAttributeEncoder()
+					ae1.Uint16(2, 2)
+					return ae1.Encode()
+				})
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			b, err := MarshalAttributes(tt.attrs)
+			if err != nil {
+				t.Fatalf("failed to marshal attributes: %v", err)
+			}
+
+			ae := NewAttributeEncoder()
+			tt.fn(ae)
+			got, err := ae.Encode()
+
+			if err != nil {
+				t.Fatalf("failed to encode attributes: %v", err)
+			}
+
+			if diff := cmp.Diff(got, b); diff != "" {
+				t.Fatalf("unexpected attribute encoding (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
+func aeEndianTest(order binary.ByteOrder) func(ae *AttributeEncoder) {
+	return func(ae *AttributeEncoder) {
+		ae.ByteOrder = order
+
+		ae.Uint8(1, uint8(1))
+		ae.Uint16(2, uint16(2))
+		ae.Uint32(3, uint32(3))
+		ae.Uint64(4, uint64(4))
+	}
+}

example_attributeencoder_test.go

@@ -0,0 +1,75 @@
+package netlink_test
+
+import (
+	"encoding/hex"
+	"fmt"
+	"log"
+
+	"github.com/mdlayher/netlink"
+)
+
+// nested is a nested structure within out.
+type nested struct {
+	A, B uint32
+}
+
+// out is an example structure we will use to pack netlink attributes.
+type out struct {
+	Number uint16
+	String string
+	Nested nested
+}
+
+// encode is an example function used to adapt the ae.Do method
+// to encode an arbitrary structure.
+func (n nested) encode() func() ([]byte, error) {
+	return func() ([]byte, error) {
+		// Create an internal, nested netlink.NewAttributeEncoder that
+		// operates on the nested set of attributes.
+		ae := netlink.NewAttributeEncoder()
+
+		// Encode the fields of the nested stucture
+		ae.Uint32(1, n.A)
+		ae.Uint32(2, n.B)
+
+		// Return the encoded attributes, and any error encountered.
+		return ae.Encode()
+	}
+}
+
+func ExampleAttributeEncoder_encode() {
+	// Create a netlink.AttributeEncoder that encodes to the same message
+	// as that decoded by the netlink.AttributeDecoder example.
+	ae := netlink.NewAttributeEncoder()
+
+	o := out{
+		Number: 1,
+		String: "hello world",
+		Nested: nested{
+			A: 2,
+			B: 3,
+		},
+	}
+
+	// Encode the Number attribute as a uint16.
+	ae.Uint16(1, o.Number)
+	// Encode the String attribute as a string.
+	ae.String(2, o.String)
+	// Nested is a nested structure, so we will use our encodeNested
+	// function along with ae.Do to encode it in a concise way.
+	ae.Do(3, o.Nested.encode())
+
+	b, err := ae.Encode()
+	// Any errors encountered during encoding (including any errors from
+	// encoding nested attributes) will be returned here.
+	if err != nil {
+		log.Fatalf("failed to encode attributes: %v", err)
+	}
+
+	fmt.Printf("Encoded netlink message follows:\n%s", hex.Dump(b))
+
+	// Output: Encoded netlink message follows:
+	// 00000000  06 00 01 00 01 00 00 00  10 00 02 00 68 65 6c 6c  |............hell|
+	// 00000010  6f 20 77 6f 72 6c 64 00  14 00 03 00 08 00 01 00  |o world.........|
+	// 00000020  02 00 00 00 08 00 02 00  03 00 00 00              |............|
+}