Introduce UnmarshalAttributes

Author: ti-mo

attribute.go

@@ -223,10 +223,10 @@ func (a *Attribute) encode(attrs []Attribute) func(*netlink.AttributeEncoder) er
 	}
 }
 
-// unmarshalAttributes returns an array of netfilter.Attributes decoded from
+// decodeAttributes returns an array of netfilter.Attributes decoded from
 // a byte array. This byte array should be taken from the netlink.Message's
 // Data payload after the nfHeaderLen offset.
-func unmarshalAttributes(ad *netlink.AttributeDecoder) ([]Attribute, error) {
+func decodeAttributes(ad *netlink.AttributeDecoder) ([]Attribute, error) {
 
 	// Use the Children element of the Attribute to decode into.
 	// Attribute already has nested decoding implemented on the type.
@@ -245,15 +245,25 @@ func unmarshalAttributes(ad *netlink.AttributeDecoder) ([]Attribute, error) {
 	return a.Children, nil
 }
 
+// encodeAttributes encodes a list of Attributes into the given netlink.AttributeEncoder.
+func encodeAttributes(ae *netlink.AttributeEncoder, attrs []Attribute) error {
+
+	if ae == nil {
+		return errNilAttributeEncoder
+	}
+
+	attr := Attribute{}
+	return attr.encode(attrs)(ae)
+}
+
 // MarshalAttributes marshals a nested attribute structure into a byte slice.
 // This byte slice can then be copied into a netlink.Message's Data field after
 // the nfHeaderLen offset.
 func MarshalAttributes(attrs []Attribute) ([]byte, error) {
 
 	ae := NewAttributeEncoder()
 
-	attr := Attribute{}
-	if err := attr.encode(attrs)(ae); err != nil {
+	if err := encodeAttributes(ae, attrs); err != nil {
 		return nil, err
 	}
 
@@ -264,3 +274,14 @@ func MarshalAttributes(attrs []Attribute) ([]byte, error) {
 
 	return b, nil
 }
+
+// UnmarshalAttributes unmarshals a byte slice into a list of Attributes.
+func UnmarshalAttributes(b []byte) ([]Attribute, error) {
+
+	ad, err := NewAttributeDecoder(b)
+	if err != nil {
+		return nil, err
+	}
+
+	return decodeAttributes(ad)
+}

attribute_test.go

@@ -1,6 +1,7 @@
 package netfilter
 
 import (
+	"errors"
 	"strings"
 	"testing"
 
@@ -209,7 +210,7 @@ func TestAttributeMarshalErrors(t *testing.T) {
 	}
 }
 
-func TestAttributeUnmarshalErrors(t *testing.T) {
+func TestAttributeDecoderErrors(t *testing.T) {
 	tests := []struct {
 		name    string
 		b       []byte
@@ -233,31 +234,19 @@ func TestAttributeUnmarshalErrors(t *testing.T) {
 			},
 			err: errInvalidAttributeFlags,
 		},
+		{
+			name: "decoding invalid attribute",
+			b:    []byte{4, 0, 0},
+			err:  errors.New("invalid attribute; length too short or too large"),
+		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			ad, err := NewAttributeDecoder(tt.b)
-			if err != nil {
-				t.Fatal("unexpected error creating AttributeDecoder:", err)
-			}
-
-			_, err = unmarshalAttributes(ad)
-			if err == nil {
-				t.Fatal("unmarshal did not error")
-			}
-
-			if tt.err != nil {
-				if want, got := tt.err, err; want != got {
-					t.Fatalf("unexpected error:\n- want: %v\n-  got: %v",
-						want, got.Error())
-				}
-			} else if tt.errWrap != "" {
-				if !strings.HasPrefix(err.Error(), tt.errWrap+":") {
-					t.Fatalf("unexpected wrapped error:\n- expected prefix: %v\n-    error string: %v",
-						tt.errWrap, err)
-				}
-			}
+			_, err := UnmarshalAttributes(tt.b)
+			require.Error(t, err)
+			require.Error(t, tt.err)
+			require.EqualError(t, err, tt.err.Error())
 		})
 	}
 }
@@ -402,7 +391,7 @@ func TestAttributeMarshalTwoWay(t *testing.T) {
 			}
 
 			// Unmarshal binary content into nested structures
-			attrs, err := unmarshalAttributes(ad)
+			attrs, err := decodeAttributes(ad)
 			require.NoError(t, err)
 
 			assert.Empty(t, cmp.Diff(tt.attrs, attrs))
@@ -417,3 +406,7 @@ func TestAttributeMarshalTwoWay(t *testing.T) {
 		})
 	}
 }
+
+func TestErrors(t *testing.T) {
+	assert.EqualError(t, encodeAttributes(nil, nil), errNilAttributeEncoder.Error())
+}

errors.go

@@ -14,4 +14,6 @@ var (
 	errConnIsMulticast = errors.New("Conn is attached to one or more multicast groups and can no longer be used for bidirectional traffic")
 
 	errNoMulticastGroups = errors.New("need one or more multicast groups to join")
+
+	errNilAttributeEncoder = errors.New("given AttributeEncoder is nil")
 )

message.go

@@ -14,7 +14,7 @@ func UnmarshalNetlink(msg netlink.Message) (Header, []Attribute, error) {
 		return Header{}, nil, err
 	}
 
-	attrs, err := unmarshalAttributes(ad)
+	attrs, err := decodeAttributes(ad)
 	if err != nil {
 		return Header{}, nil, err
 	}
@@ -44,18 +44,33 @@ func DecodeNetlink(msg netlink.Message) (Header, *netlink.AttributeDecoder, erro
 // MarshalNetlink takes a Netfilter Header and Attributes and returns a netlink.Message.
 func MarshalNetlink(h Header, attrs []Attribute) (netlink.Message, error) {
 
-	ba, err := MarshalAttributes(attrs)
+	ae := NewAttributeEncoder()
+	if err := encodeAttributes(ae, attrs); err != nil {
+		return netlink.Message{}, err
+	}
+
+	return EncodeNetlink(h, ae)
+}
+
+// EncodeNetlink generates a netlink.Message based on a given netfilter header h
+// and a pre-filled netlink.AttributeEncoder ae.
+func EncodeNetlink(h Header, ae *netlink.AttributeEncoder) (netlink.Message, error) {
+
+	if ae == nil {
+		return netlink.Message{}, errNilAttributeEncoder
+	}
+
+	// Encode the AE into a byte slice.
+	b, err := ae.Encode()
 	if err != nil {
 		return netlink.Message{}, err
 	}
 
-	// initialize with 4 bytes of Data before unmarshal
-	nlm := netlink.Message{Data: make([]byte, 4)}
+	// Allocate space for the marshaled netfilter header.
+	nlm := netlink.Message{Data: append(make([]byte, nfHeaderLen), b...)}
 
-	// marshal error ignored, safe to do if msg Data is initialized
+	// marshal error ignored, safe to do if msg Data is initialized.
 	_ = h.marshal(&nlm)
 
-	nlm.Data = append(nlm.Data, ba...)
-
 	return nlm, nil
 }

message_test.go

@@ -155,3 +155,15 @@ func TestAttributeMarshalNetlink(t *testing.T) {
 		})
 	}
 }
+
+func TestEncodeNetlink(t *testing.T) {
+
+	_, err := EncodeNetlink(Header{}, nil)
+	assert.EqualError(t, err, errNilAttributeEncoder.Error())
+
+	// Make ae.Encode() throw an error inside EncodeNetlink.
+	ae := NewAttributeEncoder()
+	ae.Do(0, func() ([]byte, error) { return []byte{}, errors.New("test error") })
+	_, err = EncodeNetlink(Header{}, ae)
+	assert.EqualError(t, err, "test error")
+}