Allow to specify netlink families for Handle (#143)

Signed-off-by: Alessandro Boch <aboch@docker.com>

handle_linux.go

@@ -1,7 +1,6 @@
 package netlink
 
 import (
-	"sync/atomic"
 	"syscall"
 
 	"github.com/vishvananda/netlink/nl"
@@ -11,27 +10,34 @@ import (
 // Empty handle used by the netlink package methods
 var pkgHandle = &Handle{}
 
-// Handle is an handle for the netlink requests
-// on a specific network namespace. All the requests
-// share the same netlink socket, which gets released
-// when the handle is deleted.
+// Handle is an handle for the netlink requests on a
+// specific network namespace. All the requests on the
+// same netlink family share the same netlink socket,
+// which gets released when the handle is deleted.
 type Handle struct {
-	seq          uint32
-	routeSocket  *nl.NetlinkSocket
-	xfrmSocket   *nl.NetlinkSocket
+	sockets      map[int]*nl.SocketHandle
 	lookupByDump bool
 }
 
+// SupportsNetlinkFamily reports whether the passed netlink family is supported by this Handle
+func (h *Handle) SupportsNetlinkFamily(nlFamily int) bool {
+	_, ok := h.sockets[nlFamily]
+	return ok
+}
+
 // NewHandle returns a netlink handle on the current network namespace.
-func NewHandle() (*Handle, error) {
-	return newHandle(netns.None(), netns.None())
+// Caller may specify the netlink families the handle should support.
+// If no families are specified, all the families the netlink package
+// supports will be automatically added.
+func NewHandle(nlFamilies ...int) (*Handle, error) {
+	return newHandle(netns.None(), netns.None(), nlFamilies...)
 }
 
 // NewHandle returns a netlink handle on the network namespace
 // specified by ns. If ns=netns.None(), current network namespace
 // will be assumed
-func NewHandleAt(ns netns.NsHandle) (*Handle, error) {
-	return newHandle(ns, netns.None())
+func NewHandleAt(ns netns.NsHandle, nlFamilies ...int) (*Handle, error) {
+	return newHandle(ns, netns.None(), nlFamilies...)
 }
 
 // NewHandleAtFrom works as NewHandle but allows client to specify the
@@ -40,37 +46,33 @@ func NewHandleAtFrom(newNs, curNs netns.NsHandle) (*Handle, error) {
 	return newHandle(newNs, curNs)
 }
 
-func newHandle(newNs, curNs netns.NsHandle) (*Handle, error) {
-	var (
-		err     error
-		rSocket *nl.NetlinkSocket
-		xSocket *nl.NetlinkSocket
-	)
-	rSocket, err = nl.GetNetlinkSocketAt(newNs, curNs, syscall.NETLINK_ROUTE)
-	if err != nil {
-		return nil, err
+func newHandle(newNs, curNs netns.NsHandle, nlFamilies ...int) (*Handle, error) {
+	h := &Handle{sockets: map[int]*nl.SocketHandle{}}
+	fams := nl.SupportedNlFamilies
+	if len(nlFamilies) != 0 {
+		fams = nlFamilies
 	}
-	xSocket, err = nl.GetNetlinkSocketAt(newNs, curNs, syscall.NETLINK_XFRM)
-	if err != nil {
-		return nil, err
+	for _, f := range fams {
+		s, err := nl.GetNetlinkSocketAt(newNs, curNs, f)
+		if err != nil {
+			return nil, err
+		}
+		h.sockets[f] = &nl.SocketHandle{Socket: s}
 	}
-	return &Handle{routeSocket: rSocket, xfrmSocket: xSocket}, nil
+	return h, nil
 }
 
 // Delete releases the resources allocated to this handle
 func (h *Handle) Delete() {
-	if h.routeSocket != nil {
-		h.routeSocket.Close()
+	for _, sh := range h.sockets {
+		sh.Close()
 	}
-	if h.xfrmSocket != nil {
-		h.xfrmSocket.Close()
-	}
-	h.routeSocket, h.xfrmSocket = nil, nil
+	h.sockets = nil
 }
 
 func (h *Handle) newNetlinkRequest(proto, flags int) *nl.NetlinkRequest {
 	// Do this so that package API still use nl package variable nextSeqNr
-	if h.routeSocket == nil {
+	if h.sockets == nil {
 		return nl.NewNetlinkRequest(proto, flags)
 	}
 	return &nl.NetlinkRequest{
@@ -78,9 +80,7 @@ func (h *Handle) newNetlinkRequest(proto, flags int) *nl.NetlinkRequest {
 			Len:   uint32(syscall.SizeofNlMsghdr),
 			Type:  uint16(proto),
 			Flags: syscall.NLM_F_REQUEST | uint16(flags),
-			Seq:   atomic.AddUint32(&h.seq, 1),
 		},
-		RouteSocket: h.routeSocket,
-		XfmrSocket:  h.xfrmSocket,
+		Sockets: h.sockets,
 	}
 }

handle_test.go

@@ -10,6 +10,7 @@ import (
 	"sync/atomic"
 	"testing"
 
+	"github.com/vishvananda/netlink/nl"
 	"github.com/vishvananda/netns"
 )
 
@@ -18,12 +19,18 @@ func TestHandleCreateDelete(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	if h.routeSocket == nil || h.xfrmSocket == nil {
-		t.Fatalf("Handle socket(s) were not created")
+	for _, f := range nl.SupportedNlFamilies {
+		sh, ok := h.sockets[f]
+		if !ok {
+			t.Fatalf("Handle socket(s) for family %d was not created", f)
+		}
+		if sh.Socket == nil {
+			t.Fatalf("Socket for family %d was not created", f)
+		}
 	}
 
 	h.Delete()
-	if h.routeSocket != nil || h.xfrmSocket != nil {
+	if h.sockets != nil {
 		t.Fatalf("Handle socket(s) were not destroyed")
 	}
 }

nl/nl_linux.go

@@ -22,6 +22,9 @@ const (
 	FAMILY_V6  = syscall.AF_INET6
 )
 
+// SupportedNlFamilies contains the list of netlink families this netlink package supports
+var SupportedNlFamilies = []int{syscall.NETLINK_ROUTE, syscall.NETLINK_XFRM}
+
 var nextSeqNr uint32
 
 // GetIPFamily returns the family type of a net.IP.
@@ -175,9 +178,8 @@ func (a *RtAttr) Serialize() []byte {
 
 type NetlinkRequest struct {
 	syscall.NlMsghdr
-	Data        []NetlinkRequestData
-	RouteSocket *NetlinkSocket
-	XfmrSocket  *NetlinkSocket
+	Data    []NetlinkRequestData
+	Sockets map[int]*SocketHandle
 }
 
 // Serialize the Netlink Request into a byte array
@@ -217,15 +219,12 @@ func (req *NetlinkRequest) Execute(sockType int, resType uint16) ([][]byte, erro
 		err error
 	)
 
-	switch sockType {
-	case syscall.NETLINK_XFRM:
-		s = req.XfmrSocket
-	case syscall.NETLINK_ROUTE:
-		s = req.RouteSocket
-	default:
-		return nil, fmt.Errorf("Socket type %d is not handled", sockType)
+	if req.Sockets != nil {
+		if sh, ok := req.Sockets[sockType]; ok {
+			s = sh.Socket
+			req.Seq = atomic.AddUint32(&sh.Seq, 1)
+		}
 	}
-
 	sharedSocket := s != nil
 
 	if s == nil {
@@ -486,3 +485,17 @@ func netlinkRouteAttrAndValue(b []byte) (*syscall.RtAttr, []byte, int, error) {
 	}
 	return a, b[syscall.SizeofRtAttr:], rtaAlignOf(int(a.Len)), nil
 }
+
+// SocketHandle contains the netlink socket and the associated
+// sequence counter for a specific netlink family
+type SocketHandle struct {
+	Seq    uint32
+	Socket *NetlinkSocket
+}
+
+// Close closes the netlink socket
+func (sh *SocketHandle) Close() {
+	if sh.Socket != nil {
+		sh.Socket.Close()
+	}
+}