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,17 +65,60 @@ 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
+	for cc.checkReceive(conn) {
+		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)
+			}
+		}
 	}
 
 	return nil
 }
 
+func (cc *Conn) checkReceive(c *netlink.Conn) bool {
+	if cc.TestDial != nil {
+		return false
+	}
+
+	sc, err := c.SyscallConn()
+
+	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
+	}
+
+	return false
+}
+
 // FlushRuleset flushes the entire ruleset. See also
 // https://wiki.nftables.org/wiki-nftables/index.php/Operations_at_ruleset_level
 func (cc *Conn) FlushRuleset() {

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

@@ -3980,3 +3980,122 @@ func TestStatelessNAT(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestHandleBack(t *testing.T) {
+
+	if os.Getenv("TRAVIS") == "true" {
+		t.SkipNow()
+	}
+
+	// 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",
+	})
+
+	chain1 := c.AddChain(&nftables.Chain{
+		Name:     "chain1",
+		Table:    filter,
+		Type:     nftables.ChainTypeFilter,
+		Hooknum:  nftables.ChainHookPrerouting,
+		Priority: nftables.ChainPriorityFilter,
+	})
+
+	chain2 := c.AddChain(&nftables.Chain{
+		Name:     "chain2",
+		Table:    filter,
+		Type:     nftables.ChainTypeFilter,
+		Hooknum:  nftables.ChainHookPrerouting,
+		Priority: nftables.ChainPriorityFilter,
+	})
+
+	var rulesCreated1 []*nftables.Rule
+	var rulesCreated2 []*nftables.Rule
+
+	rulesCreated1 = append(rulesCreated1, c.AddRule(&nftables.Rule{
+		Table: filter,
+		Chain: chain1,
+		Exprs: []expr.Any{
+			&expr.Verdict{
+				// [ immediate reg 0 drop ]
+				Kind: expr.VerdictDrop,
+			},
+		},
+	}))
+
+	rulesCreated1 = append(rulesCreated1, c.AddRule(&nftables.Rule{
+		Table: filter,
+		Chain: chain1,
+		Exprs: []expr.Any{
+			&expr.Verdict{
+				// [ immediate reg 0 drop ]
+				Kind: expr.VerdictDrop,
+			},
+		},
+	}))
+
+	rulesCreated2 = append(rulesCreated2, c.AddRule(&nftables.Rule{
+		Table: filter,
+		Chain: chain2,
+		Exprs: []expr.Any{
+			&expr.Verdict{
+				// [ immediate reg 0 drop ]
+				Kind: expr.VerdictDrop,
+			},
+		},
+	}))
+
+	for i, r := range rulesCreated1 {
+		if r.Handle != 0 {
+			t.Fatalf("unexpected handle value at %d", i)
+		}
+	}
+
+	for i, r := range rulesCreated2 {
+		if r.Handle != 0 {
+			t.Fatalf("unexpected handle value at %d", i)
+		}
+	}
+
+	if err := c.Flush(); err != nil {
+		t.Fatal(err)
+	}
+
+	rulesGetted1, _ := c.GetRule(filter, chain1)
+	rulesGetted2, _ := c.GetRule(filter, chain2)
+
+	if len(rulesGetted1) != len(rulesCreated1) {
+		t.Fatalf("Bad ruleset length got %d want %d", len(rulesGetted1), len(rulesCreated1))
+	}
+
+	if len(rulesGetted2) != len(rulesCreated2) {
+		t.Fatalf("Bad ruleset length got %d want %d", len(rulesGetted2), len(rulesCreated2))
+	}
+
+	for i, r := range rulesGetted1 {
+		if r.Handle == 0 {
+			t.Fatalf("handle value is empty at %d", i)
+		}
+
+		if r.Handle != rulesCreated1[i].Handle {
+			t.Fatalf("mismatched handle at %d", i)
+		}
+	}
+
+	for i, r := range rulesGetted2 {
+		if r.Handle == 0 {
+			t.Fatalf("handle value is empty at %d", i)
+		}
+
+		if r.Handle != rulesCreated2[i].Handle {
+			t.Fatalf("mismatched handle at %d", i)
+		}
+	}
+}

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 {