fix: adjust for the review and comments

go.sum

@@ -6,8 +6,6 @@ github.com/mdlayher/netlink v1.7.1 h1:FdUaT/e33HjEXagwELR8R3/KL1Fq5x3G5jgHLp/BTm
 github.com/mdlayher/netlink v1.7.1/go.mod h1:nKO5CSjE/DJjVhk/TNp6vCE1ktVxEA8VEh8drhZzxsQ=
 github.com/mdlayher/socket v0.4.0 h1:280wsy40IC9M9q1uPGcLBwXpcTQDtoGwVt+BNoITxIw=
 github.com/mdlayher/socket v0.4.0/go.mod h1:xxFqz5GRCUN3UEOm9CZqEJsAbe1C8OwSK46NlmWuVoc=
-github.com/mdlayher/socket v0.5.0 h1:ilICZmJcQz70vrWVes1MFera4jGiWNocSkykwwoy3XI=
-github.com/mdlayher/socket v0.5.0/go.mod h1:WkcBFfvyG8QENs5+hfQPl1X6Jpd2yeLIYgrGFmJiJxI=
 github.com/vishvananda/netns v0.0.0-20180720170159-13995c7128cc h1:R83G5ikgLMxrBvLh22JhdfI8K6YXEPHx5P03Uu3DRs4=
 github.com/vishvananda/netns v0.0.0-20180720170159-13995c7128cc/go.mod h1:ZjcWmFBXmLKZu9Nxj3WKYEafiSqer2rnvPr0en9UNpI=
 github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
@@ -28,8 +26,6 @@ golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o=
 golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.3.0 h1:ftCYgMx6zT/asHUrPw8BLLscYtGznsLAnjq5RH9P66E=
-golang.org/x/sync v0.3.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

monitor.go

@@ -46,57 +46,50 @@ const (
 )
 
 var (
-	monitorFlags         map[MonitorAction]map[MonitorObject]uint32
+	monitorFlags = map[MonitorAction]map[MonitorObject]uint32{
+		MonitorActionAny: {
+			MonitorObjectAny:      0xffffffff,
+			MonitorObjectTables:   1<<unix.NFT_MSG_NEWTABLE | 1<<unix.NFT_MSG_DELCHAIN,
+			MonitorObjectChains:   1<<unix.NFT_MSG_NEWCHAIN | 1<<unix.NFT_MSG_DELCHAIN,
+			MonitorObjectRules:    1<<unix.NFT_MSG_NEWRULE | 1<<unix.NFT_MSG_DELRULE,
+			MonitorObjectSets:     1<<unix.NFT_MSG_NEWSET | 1<<unix.NFT_MSG_DELSET,
+			MonitorObjectElements: 1<<unix.NFT_MSG_NEWSETELEM | 1<<unix.NFT_MSG_DELSETELEM,
+			MonitorObjectRuleset: 1<<unix.NFT_MSG_NEWTABLE | 1<<unix.NFT_MSG_DELCHAIN |
+				1<<unix.NFT_MSG_NEWCHAIN | 1<<unix.NFT_MSG_DELCHAIN |
+				1<<unix.NFT_MSG_NEWRULE | 1<<unix.NFT_MSG_DELRULE |
+				1<<unix.NFT_MSG_NEWSET | 1<<unix.NFT_MSG_DELSET |
+				1<<unix.NFT_MSG_NEWSETELEM | 1<<unix.NFT_MSG_DELSETELEM |
+				1<<unix.NFT_MSG_NEWOBJ | 1<<unix.NFT_MSG_DELOBJ,
+		},
+		MonitorActionNew: {
+			MonitorObjectAny: 1<<unix.NFT_MSG_NEWTABLE |
+				1<<unix.NFT_MSG_NEWCHAIN |
+				1<<unix.NFT_MSG_NEWRULE |
+				1<<unix.NFT_MSG_NEWSET |
+				1<<unix.NFT_MSG_NEWSETELEM,
+			MonitorObjectTables: 1 << unix.NFT_MSG_NEWTABLE,
+			MonitorObjectChains: 1 << unix.NFT_MSG_NEWCHAIN,
+			MonitorObjectRules:  1 << unix.NFT_MSG_NEWRULE,
+			MonitorObjectSets:   1 << unix.NFT_MSG_NEWSET,
+			MonitorObjectRuleset: 1<<unix.NFT_MSG_NEWTABLE |
+				1<<unix.NFT_MSG_NEWCHAIN |
+				1<<unix.NFT_MSG_NEWRULE |
+				1<<unix.NFT_MSG_NEWSET |
+				1<<unix.NFT_MSG_NEWSETELEM |
+				1<<unix.NFT_MSG_NEWOBJ,
+		},
+		MonitorActionDel: {
+			MonitorObjectAny: 1<<unix.NFT_MSG_DELTABLE |
+				1<<unix.NFT_MSG_DELCHAIN |
+				1<<unix.NFT_MSG_DELRULE |
+				1<<unix.NFT_MSG_DELSET |
+				1<<unix.NFT_MSG_DELSETELEM |
+				1<<unix.NFT_MSG_DELOBJ,
+		},
+	}
 	monitorFlagsInitOnce sync.Once
 )
 
-// A lazy init function to define flags.
-func lazyInit() {
-	monitorFlagsInitOnce.Do(func() {
-		monitorFlags = map[MonitorAction]map[MonitorObject]uint32{
-			MonitorActionAny: {
-				MonitorObjectAny:      0xffffffff,
-				MonitorObjectTables:   1<<unix.NFT_MSG_NEWTABLE | 1<<unix.NFT_MSG_DELCHAIN,
-				MonitorObjectChains:   1<<unix.NFT_MSG_NEWCHAIN | 1<<unix.NFT_MSG_DELCHAIN,
-				MonitorObjectRules:    1<<unix.NFT_MSG_NEWRULE | 1<<unix.NFT_MSG_DELRULE,
-				MonitorObjectSets:     1<<unix.NFT_MSG_NEWSET | 1<<unix.NFT_MSG_DELSET,
-				MonitorObjectElements: 1<<unix.NFT_MSG_NEWSETELEM | 1<<unix.NFT_MSG_DELSETELEM,
-				MonitorObjectRuleset: 1<<unix.NFT_MSG_NEWTABLE | 1<<unix.NFT_MSG_DELCHAIN |
-					1<<unix.NFT_MSG_NEWCHAIN | 1<<unix.NFT_MSG_DELCHAIN |
-					1<<unix.NFT_MSG_NEWRULE | 1<<unix.NFT_MSG_DELRULE |
-					1<<unix.NFT_MSG_NEWSET | 1<<unix.NFT_MSG_DELSET |
-					1<<unix.NFT_MSG_NEWSETELEM | 1<<unix.NFT_MSG_DELSETELEM |
-					1<<unix.NFT_MSG_NEWOBJ | 1<<unix.NFT_MSG_DELOBJ,
-			},
-			MonitorActionNew: {
-				MonitorObjectAny: 1<<unix.NFT_MSG_NEWTABLE |
-					1<<unix.NFT_MSG_NEWCHAIN |
-					1<<unix.NFT_MSG_NEWRULE |
-					1<<unix.NFT_MSG_NEWSET |
-					1<<unix.NFT_MSG_NEWSETELEM,
-				MonitorObjectTables: 1 << unix.NFT_MSG_NEWTABLE,
-				MonitorObjectChains: 1 << unix.NFT_MSG_NEWCHAIN,
-				MonitorObjectRules:  1 << unix.NFT_MSG_NEWRULE,
-				MonitorObjectSets:   1 << unix.NFT_MSG_NEWSET,
-				MonitorObjectRuleset: 1<<unix.NFT_MSG_NEWTABLE |
-					1<<unix.NFT_MSG_NEWCHAIN |
-					1<<unix.NFT_MSG_NEWRULE |
-					1<<unix.NFT_MSG_NEWSET |
-					1<<unix.NFT_MSG_NEWSETELEM |
-					1<<unix.NFT_MSG_NEWOBJ,
-			},
-			MonitorActionDel: {
-				MonitorObjectAny: 1<<unix.NFT_MSG_DELTABLE |
-					1<<unix.NFT_MSG_DELCHAIN |
-					1<<unix.NFT_MSG_DELRULE |
-					1<<unix.NFT_MSG_DELSET |
-					1<<unix.NFT_MSG_DELSETELEM |
-					1<<unix.NFT_MSG_DELOBJ,
-			},
-		}
-	})
-}
-
 type EventType int
 
 const (
@@ -131,12 +124,13 @@ type Monitor struct {
 	object       MonitorObject
 	monitorFlags uint32
 
-	// mtx covers eventCh and status
+	conn   *netlink.Conn
-	mtx     sync.Mutex
+	closer netlinkCloser
+
+	// mu covers eventCh and status
+	mu      sync.Mutex
 	eventCh chan *Event
 	status  int
-	conn    *netlink.Conn
-	closer  netlinkCloser
 }
 
 type MonitorOption func(*Monitor)
@@ -163,8 +157,6 @@ func WithMonitorObject(object MonitorObject) MonitorOption {
 
 // NewMonitor returns a Monitor with options to be started.
 func NewMonitor(opts ...MonitorOption) *Monitor {
-	lazyInit()
-
 	monitor := &Monitor{
 		status: monitorOK,
 	}
@@ -193,7 +185,7 @@ func (monitor *Monitor) monitor() {
 			break
 		}
 		for _, msg := range msgs {
-			if got, want := msg.Header.Type&0xff00>>8, netlink.HeaderType(unix.NFNL_SUBSYS_NFTABLES); got != want {
+			if msg.Header.Type&0xff00>>8 != netlink.HeaderType(unix.NFNL_SUBSYS_NFTABLES) {
 				continue
 			}
 			msgType := msg.Header.Type & 0x00ff
@@ -249,27 +241,27 @@ func (monitor *Monitor) monitor() {
 					Error: err,
 				}
 				monitor.eventCh <- event
-			case unix.NFT_MSG_TRACE:
 			}
 		}
 	}
-	monitor.mtx.Lock()
+	monitor.mu.Lock()
+	defer monitor.mu.Unlock()
+
 	if monitor.status != monitorClosed {
 		monitor.status = monitorClosed
 		monitor.closer()
 		close(monitor.eventCh)
 	}
-	monitor.mtx.Unlock()
 }
 
 func (monitor *Monitor) Close() {
-	monitor.mtx.Lock()
+	monitor.mu.Lock()
 	if monitor.status != monitorClosed {
 		monitor.status = monitorClosed
 		monitor.closer()
 		close(monitor.eventCh)
 	}
-	monitor.mtx.Unlock()
+	monitor.mu.Unlock()
 }
 
 // AddMonitor to perform the monitor immediately. The channel will be closed after
@@ -283,12 +275,10 @@ func (cc *Conn) AddMonitor(monitor *Monitor) (chan *Event, error) {
 	monitor.closer = closer
 
 	if monitor.monitorFlags != 0 {
-		err = conn.JoinGroup(uint32(unix.NFNLGRP_NFTABLES))
+		if err = conn.JoinGroup(uint32(unix.NFNLGRP_NFTABLES)); err != nil {
-		if err != nil {
 			monitor.closer()
 			return nil, err
 		}
-		conn.JoinGroup(uint32(unix.NFNLGRP_NFTRACE))
 	}
 
 	go monitor.monitor()

monitor_test.go

@@ -4,8 +4,8 @@ import (
 	"fmt"
 	"net"
 	"sync"
+	"sync/atomic"
 	"testing"
-	"time"
 
 	"github.com/google/nftables"
 	"github.com/google/nftables/expr"
@@ -13,7 +13,6 @@ import (
 )
 
 func TestMonitor(t *testing.T) {
-
 	// Create a new network namespace to test these operations,
 	// and tear down the namespace at test completion.
 	c, newNS := nftest.OpenSystemConn(t, *enableSysTests)
@@ -28,6 +27,7 @@ func TestMonitor(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer monitor.Close()
 
 	var gotTable *nftables.Table
 	var gotChain *nftables.Chain
@@ -36,6 +36,7 @@ func TestMonitor(t *testing.T) {
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
+		count := int32(0)
 		for {
 			select {
 			case event, ok := <-events:
@@ -49,10 +50,16 @@ func TestMonitor(t *testing.T) {
 				switch event.Type {
 				case nftables.EventTypeNewTable:
 					gotTable = event.Data.(*nftables.Table)
+					atomic.AddInt32(&count, 1)
 				case nftables.EventTypeNewChain:
 					gotChain = event.Data.(*nftables.Chain)
+					atomic.AddInt32(&count, 1)
 				case nftables.EventTypeNewRule:
 					gotRule = event.Data.(*nftables.Rule)
+					atomic.AddInt32(&count, 1)
+				}
+				if atomic.LoadInt32(&count) == 3 {
+					return
 				}
 			}
 		}
@@ -97,13 +104,7 @@ func TestMonitor(t *testing.T) {
 	if err := c.Flush(); err != nil {
 		t.Fatal(err)
 	}
-	// It takes time for the kernel to take effect
-	time.Sleep(time.Second)
-	monitor.Close()
 	wg.Wait()
-	if err != nil {
-		t.Fatal(err)
-	}
 	if gotTable.Family != nat.Family || gotTable.Name != nat.Name {
 		t.Fatal("no want table", gotTable.Family, gotTable.Name)
 	}