Set the rule handle after flush

conn.go

@@ -24,6 +24,10 @@ import (
 	"golang.org/x/sys/unix"
 )
 
+type Entity interface {
+	HandleResponse(netlink.Message)
+}
+
 // A Conn represents a netlink connection of the nftables family.
 //
 // All methods return their input, so that variables can be defined from string
@@ -35,6 +39,7 @@ type Conn struct {
 	NetNS    int         // Network namespace netlink will interact with.
 	sync.Mutex
 	messages []netlink.Message
+	entities map[int]Entity
 	err      error
 }
 
@@ -43,6 +48,7 @@ func (cc *Conn) Flush() error {
 	cc.Lock()
 	defer func() {
 		cc.messages = nil
+		cc.entities = nil
 		cc.Unlock()
 	}()
 	if len(cc.messages) == 0 {
@@ -59,15 +65,64 @@ func (cc *Conn) Flush() error {
 
 	defer conn.Close()
 
-	if _, err := conn.SendMessages(batch(cc.messages)); err != nil {
+	smsg, err := conn.SendMessages(batch(cc.messages))
+
+	if err != nil {
 		return fmt.Errorf("SendMessages: %w", err)
 	}
 
-	if _, err := conn.Receive(); err != nil {
-		return fmt.Errorf("Receive: %w", err)
+	// Retrieving of seq number associated to entities
+	entitiesBySeq := make(map[uint32]Entity)
+	for i, e := range cc.entities {
+		entitiesBySeq[smsg[i].Header.Sequence] = e
+	}
+
+	// Trigger entities callback
+	msg, err := cc.checkReceive(conn)
+	for msg {
+		rmsg, err := conn.Receive()
+
+		if err != nil {
+			return fmt.Errorf("Receive: %w", err)
+		}
+
+		for _, msg := range rmsg {
+			if e, ok := entitiesBySeq[msg.Header.Sequence]; ok {
+				e.HandleResponse(msg)
+			}
+		}
+		msg, err = cc.checkReceive(conn)
+	}
+
+	return err
+}
+
+func (cc *Conn) checkReceive(c *netlink.Conn) (bool, error) {
+	if cc.TestDial != nil {
+		return false, nil
+	}
+
+	sc, err := c.SyscallConn()
+
+	if err != nil {
+		return false, fmt.Errorf("SyscallConn error: %w", err)
+	}
+
+	var n int
+
+	sc.Control(func(fd uintptr) {
+		var fdSet unix.FdSet
+		fdSet.Zero()
+		fdSet.Set(int(fd))
+
+		n, err = unix.Select(int(fd)+1, &fdSet, nil, nil, &unix.Timeval{})
+	})
+
+	if err == nil && n > 0 {
+		return true, nil
 	}
 
-	return nil
+	return false, err
 }
 
 // FlushRuleset flushes the entire ruleset. See also

go.mod

@@ -4,8 +4,8 @@ go 1.12
 
 require (
 	github.com/koneu/natend v0.0.0-20150829182554-ec0926ea948d
-	github.com/mdlayher/netlink v0.0.0-20191009155606-de872b0d824b
+	github.com/mdlayher/netlink v1.0.0
 	github.com/vishvananda/netns v0.0.0-20180720170159-13995c7128cc
 	golang.org/x/net v0.0.0-20191028085509-fe3aa8a45271 // indirect
-	golang.org/x/sys v0.0.0-20191029155521-f43be2a4598c
+	golang.org/x/sys v0.0.0-20200106114638-5f8ca72cd632
 )

go.sum

@@ -8,6 +8,8 @@ github.com/koneu/natend v0.0.0-20150829182554-ec0926ea948d/go.mod h1:QHb4k4cr1fQ
 github.com/mdlayher/netlink v0.0.0-20190409211403-11939a169225/go.mod h1:eQB3mZE4aiYnlUsyGGCOpPETfdQq4Jhsgf1fk3cwQaA=
 github.com/mdlayher/netlink v0.0.0-20191009155606-de872b0d824b h1:W3er9pI7mt2gOqOWzwvx20iJ8Akiqz1mUMTxU6wdvl8=
 github.com/mdlayher/netlink v0.0.0-20191009155606-de872b0d824b/go.mod h1:KxeJAFOFLG6AjpyDkQ/iIhxygIUKD+vcwqcnu43w/+M=
+github.com/mdlayher/netlink v1.0.0 h1:vySPY5Oxnn/8lxAPn2cK6kAzcZzYJl3KriSLO46OT18=
+github.com/mdlayher/netlink v1.0.0/go.mod h1:KxeJAFOFLG6AjpyDkQ/iIhxygIUKD+vcwqcnu43w/+M=
 github.com/vishvananda/netns v0.0.0-20180720170159-13995c7128cc h1:R83G5ikgLMxrBvLh22JhdfI8K6YXEPHx5P03Uu3DRs4=
 github.com/vishvananda/netns v0.0.0-20180720170159-13995c7128cc/go.mod h1:ZjcWmFBXmLKZu9Nxj3WKYEafiSqer2rnvPr0en9UNpI=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
@@ -25,4 +27,6 @@ golang.org/x/sys v0.0.0-20191029155521-f43be2a4598c h1:S/FtSvpNLtFBgjTqcKsRpsa6a
 golang.org/x/sys v0.0.0-20191029155521-f43be2a4598c/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20191112214154-59a1497f0cea h1:Mz1TMnfJDRJLk8S8OPCoJYgrsp/Se/2TBre2+vwX128=
 golang.org/x/sys v0.0.0-20191113150313-8ad342257130 h1:+sdNBpwFF05NvMnEyGynbOs/Gr2LQwORWEPKXuEXxzU=
+golang.org/x/sys v0.0.0-20200106114638-5f8ca72cd632 h1:ateQkYCVYo8UwIBvoR3zj1Dh2K6Op/n3GxemXfB44/Y=
+golang.org/x/sys v0.0.0-20200106114638-5f8ca72cd632/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=

nftables_test.go

@@ -23,6 +23,7 @@ import (
 	"reflect"
 	"runtime"
 	"strings"
+	"sync"
 	"testing"
 
 	"github.com/google/nftables"
@@ -3980,3 +3981,85 @@ func TestStatelessNAT(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestIntegrationAddRule(t *testing.T) {
+
+	// Create a new network namespace to test these operations,
+	// and tear down the namespace at test completion.
+	c, newNS := openSystemNFTConn(t)
+	defer cleanupSystemNFTConn(t, newNS)
+	// Clear all rules at the beginning + end of the test.
+	c.FlushRuleset()
+	defer c.FlushRuleset()
+
+	filter := c.AddTable(&nftables.Table{
+		Family: nftables.TableFamilyIPv4,
+		Name:   "filter",
+	})
+
+	chain := c.AddChain(&nftables.Chain{
+		Name:     "chain",
+		Table:    filter,
+		Type:     nftables.ChainTypeFilter,
+		Hooknum:  nftables.ChainHookPrerouting,
+		Priority: nftables.ChainPriorityFilter,
+	})
+
+	c.Flush()
+
+	execN := func(w int, n int) {
+
+		c := &nftables.Conn{NetNS: int(newNS)}
+
+		for i := 0; i < n; i++ {
+
+			r := c.AddRule(&nftables.Rule{
+				Table:    filter,
+				Chain:    chain,
+				UserData: []byte(fmt.Sprintf("%d-%d", w, i)),
+				Exprs: []expr.Any{
+					&expr.Verdict{
+						// [ immediate reg 0 drop ]
+						Kind: expr.VerdictDrop,
+					},
+				},
+			})
+
+			if r.Handle != 0 {
+				t.Fatalf("unexpected handle value at %d", i)
+			}
+
+			if err := c.Flush(); err != nil {
+				t.Fatal(err)
+			}
+
+			rulesGetted, _ := c.GetRule(filter, chain)
+
+			for i, rg := range rulesGetted {
+				if r.Handle == 0 {
+					t.Fatalf("handle value is empty at %d", i)
+				}
+
+				if bytes.Equal(rg.UserData, r.UserData) && rg.Handle != r.Handle {
+					t.Fatalf("mismatched handle at %d-%d, got: %d, want: %d", w, i, r.Handle, rg.Handle)
+				}
+			}
+		}
+	}
+
+	const (
+		workers    = 16
+		iterations = 256
+	)
+
+	var wg sync.WaitGroup
+	wg.Add(workers)
+	for i := 0; i < workers; i++ {
+		go func(n int) {
+			defer wg.Done()
+			execN(n, iterations)
+		}(i)
+	}
+
+	wg.Wait()
+}

rule.go

@@ -130,6 +130,12 @@ func (cc *Conn) AddRule(r *Rule) *Rule {
 		Data: append(extraHeader(uint8(r.Table.Family), 0), msgData...),
 	})
 
+	if cc.entities == nil {
+		cc.entities = make(map[int]Entity)
+	}
+
+	cc.entities[len(cc.messages)] = r
+
 	return r
 }
 
@@ -160,6 +166,15 @@ func (cc *Conn) DelRule(r *Rule) error {
 	return nil
 }
 
+// HandleResponse retrieves Handle in netlink response
+func (r *Rule) HandleResponse(msg netlink.Message) {
+	rule, err := ruleFromMsg(msg)
+
+	if err == nil {
+		r.Handle = rule.Handle
+	}
+}
+
 func exprsFromMsg(b []byte) ([]expr.Any, error) {
 	ad, err := netlink.NewAttributeDecoder(b)
 	if err != nil {