From 352d70194aa80f1c31a478945de4663dcbb9d624 Mon Sep 17 00:00:00 2001 From: Auztin Zhai Date: Fri, 8 Dec 2023 04:23:35 -0500 Subject: [PATCH] fix: adjust for the review and comments --- go.sum | 4 -- monitor.go | 116 ++++++++++++++++++++++-------------------------- monitor_test.go | 17 +++---- 3 files changed, 62 insertions(+), 75 deletions(-) diff --git a/go.sum b/go.sum index 7a07830..3c98fad 100644 --- a/go.sum +++ b/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= diff --git a/monitor.go b/monitor.go index a49958d..2d1801e 100644 --- a/monitor.go +++ b/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<>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 != nil { + if err = conn.JoinGroup(uint32(unix.NFNLGRP_NFTABLES)); err != nil { monitor.closer() return nil, err } - conn.JoinGroup(uint32(unix.NFNLGRP_NFTRACE)) } go monitor.monitor() diff --git a/monitor_test.go b/monitor_test.go index 04dfe30..8df8eb6 100644 --- a/monitor_test.go +++ b/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) }