diff --git a/conn.go b/conn.go index 01e6c5d..2cd5832 100644 --- a/conn.go +++ b/conn.go @@ -14,6 +14,12 @@ type Conn struct { conn *netfilter.Conn } +// DumpOptions is passed as an option to `Dump`-related methods to modify their behaviour. +type DumpOptions struct { + // ZeroCounters resets all flows' counters to zero after the dump operation. + ZeroCounters bool +} + // Dial opens a new Netfilter Netlink connection and returns it // wrapped in a Conn structure that implements the Conntrack API. func Dial(config *netlink.Config) (*Conn, error) { @@ -128,12 +134,16 @@ func (c *Conn) eventWorker(workerID uint8, evChan chan<- Event, errChan chan<- e // Dump gets all Conntrack connections from the kernel in the form of a list // of Flow objects. -func (c *Conn) Dump() ([]Flow, error) { +func (c *Conn) Dump(opts *DumpOptions) ([]Flow, error) { + msgType := ctGet + if opts != nil && opts.ZeroCounters { + msgType = ctGetCtrZero + } req, err := netfilter.MarshalNetlink( netfilter.Header{ SubsystemID: netfilter.NFSubsysCTNetlink, - MessageType: netfilter.MessageType(ctGet), + MessageType: netfilter.MessageType(msgType), Family: netfilter.ProtoUnspec, // ProtoUnspec dumps both IPv4 and IPv6 Flags: netlink.Request | netlink.Dump, }, @@ -153,12 +163,16 @@ func (c *Conn) Dump() ([]Flow, error) { // DumpFilter gets all Conntrack connections from the kernel in the form of a list // of Flow objects, but only returns Flows matching the connmark specified in the Filter parameter. -func (c *Conn) DumpFilter(f Filter) ([]Flow, error) { +func (c *Conn) DumpFilter(f Filter, opts *DumpOptions) ([]Flow, error) { + msgType := ctGet + if opts != nil && opts.ZeroCounters { + msgType = ctGetCtrZero + } req, err := netfilter.MarshalNetlink( netfilter.Header{ SubsystemID: netfilter.NFSubsysCTNetlink, - MessageType: netfilter.MessageType(ctGet), + MessageType: netfilter.MessageType(msgType), Family: netfilter.ProtoUnspec, // ProtoUnspec dumps both IPv4 and IPv6 Flags: netlink.Request | netlink.Dump, }, @@ -179,7 +193,6 @@ func (c *Conn) DumpFilter(f Filter) ([]Flow, error) { // DumpExpect gets all expected Conntrack expectations from the kernel in the form // of a list of Expect objects. func (c *Conn) DumpExpect() ([]Expect, error) { - req, err := netfilter.MarshalNetlink( netfilter.Header{ SubsystemID: netfilter.NFSubsysCTNetlinkExp, diff --git a/conn_test.go b/conn_test.go index 582290d..27e97d3 100644 --- a/conn_test.go +++ b/conn_test.go @@ -94,7 +94,7 @@ func ExampleConn_dumpFilter() { _ = c.Create(f2) // Dump all records in the Conntrack table that match the filter's mark/mask. - df, err := c.DumpFilter(conntrack.Filter{Mark: 0xff00, Mask: 0xff00}) + df, err := c.DumpFilter(conntrack.Filter{Mark: 0xff00, Mask: 0xff00}, nil) if err != nil { log.Fatal(err) } diff --git a/flow_integration_test.go b/flow_integration_test.go index 4d01ed2..a38e53d 100644 --- a/flow_integration_test.go +++ b/flow_integration_test.go @@ -28,7 +28,7 @@ func TestConnCreateFlows(t *testing.T) { }() // Expect empty result from empty table dump - de, err := c.Dump() + de, err := c.Dump(nil) require.NoError(t, err, "dumping empty table") require.Len(t, de, 0, "expecting 0-length dump from empty table") @@ -55,7 +55,7 @@ func TestConnCreateFlows(t *testing.T) { require.NoError(t, err, "creating IPv6 flow", i) } - flows, err := c.Dump() + flows, err := c.Dump(nil) require.NoError(t, err, "dumping table") // Expect twice the amount of numFlows, both for IPv4 and IPv6 @@ -77,7 +77,7 @@ func TestConnFlush(t *testing.T) { require.NoError(t, err) // Expect empty result from empty table dump - de, err := c.Dump() + de, err := c.Dump(nil) require.NoError(t, err, "dumping empty table") require.Len(t, de, 0, "expecting 0-length dump from empty table") @@ -100,7 +100,7 @@ func TestConnFlush(t *testing.T) { require.NoError(t, err, "creating IPv6 flow") // Expect both flows to be in the table - flows, err := c.Dump() + flows, err := c.Dump(nil) require.NoError(t, err, "dumping table before flush") assert.Equal(t, 2, len(flows)) @@ -108,7 +108,7 @@ func TestConnFlush(t *testing.T) { require.NoError(t, err, "flushing table") // Expect empty table - flows, err = c.Dump() + flows, err = c.Dump(nil) require.NoError(t, err, "dumping table after flush") assert.Equal(t, 0, len(flows)) } @@ -126,7 +126,7 @@ func TestConnFlushFilter(t *testing.T) { require.NoError(t, err) // Expect empty result from empty table dump - de, err := c.Dump() + de, err := c.Dump(nil) require.NoError(t, err, "dumping empty table") require.Len(t, de, 0, "expecting 0-length dump from empty table") @@ -149,7 +149,7 @@ func TestConnFlushFilter(t *testing.T) { require.NoError(t, err, "creating IPv6 flow") // Expect both flows to be in the table - flows, err := c.Dump() + flows, err := c.Dump(nil) require.NoError(t, err, "dumping table before filtered flush") assert.Equal(t, 2, len(flows)) @@ -158,7 +158,7 @@ func TestConnFlushFilter(t *testing.T) { require.NoError(t, err, "flushing table") // Expect only one flow to remain in the table - flows, err = c.Dump() + flows, err = c.Dump(nil) require.NoError(t, err, "dumping table after filtered flush") assert.Equal(t, 1, len(flows)) } @@ -189,7 +189,7 @@ func TestConnCreateDeleteFlows(t *testing.T) { require.NoError(t, err, "deleting flow", i) } - flows, err := c.Dump() + flows, err := c.Dump(nil) require.NoError(t, err, "dumping table") assert.Equal(t, 0, len(flows)) @@ -217,7 +217,7 @@ func TestConnCreateUpdateFlow(t *testing.T) { err = c.Update(f) require.NoError(t, err, "updating flow") - flows, err := c.Dump() + flows, err := c.Dump(nil) require.NoError(t, err, "dumping table") if got := flows[0].Timeout; !(got > 200) { @@ -234,7 +234,7 @@ func TestConnCreateUpdateFlow(t *testing.T) { err = c.Update(fNoOrig) require.NoError(t, err, "updating flow without TupleOrig") - flows, err = c.Dump() + flows, err = c.Dump(nil) require.NoError(t, err, "dumping table") if got := flows[0].Timeout; !(got > 300) { @@ -251,7 +251,7 @@ func TestConnCreateUpdateFlow(t *testing.T) { err = c.Update(fNoReply) require.NoError(t, err, "updating flow without TupleReply") - flows, err = c.Dump() + flows, err = c.Dump(nil) require.NoError(t, err, "dumping table") if got := flows[0].Timeout; !(got > 400) { @@ -308,6 +308,26 @@ func TestConnCreateGetFlow(t *testing.T) { } } +// Creates IPv4 and IPv6 flows and dumps them while zeroing the accounting counters. +func TestDumpZero(t *testing.T) { + c, _, err := makeNSConn() + require.NoError(t, err) + + f := NewFlow(17, 0, net.ParseIP("1.2.3.4"), net.ParseIP("5.6.7.8"), 1234, 5678, 120, 0xff000000) + + f.CountersOrig.Bytes = 1337 + f.CountersReply.Bytes = 9001 + require.NoError(t, c.Create(f), "creating flow") + + df, err := c.Dump(&DumpOptions{ + ZeroCounters: true, + }) + require.NoError(t, err, "dumping flows (zeroing enabled)") + + assert.Equal(t, df[0].CountersOrig.Bytes, uint64(0)) + assert.Equal(t, df[0].CountersReply.Bytes, uint64(0)) +} + // Creates IPv4 and IPv6 flows with connmarks and queries them using a filtered dump. func TestConnDumpFilter(t *testing.T) { @@ -326,7 +346,7 @@ func TestConnDumpFilter(t *testing.T) { } // Expect empty result from empty table dump - de, err := c.DumpFilter(Filter{Mark: 0x00000000, Mask: 0xffffffff}) + de, err := c.DumpFilter(Filter{Mark: 0x00000000, Mask: 0xffffffff}, nil) require.NoError(t, err, "dumping empty table") require.Len(t, de, 0, "expecting 0-length dump from empty table") @@ -334,7 +354,7 @@ func TestConnDumpFilter(t *testing.T) { err = c.Create(f) require.NoError(t, err, "creating flow", n) - df, err := c.DumpFilter(Filter{Mark: f.Mark, Mask: f.Mark}) + df, err := c.DumpFilter(Filter{Mark: f.Mark, Mask: f.Mark}, nil) require.NoError(t, err, "dumping filtered flows", n) assert.Len(t, df, 1) @@ -343,7 +363,7 @@ func TestConnDumpFilter(t *testing.T) { } // Expect table to be empty at end of run - d, err := c.Dump() + d, err := c.Dump(nil) require.NoError(t, err, "dumping flows") assert.Len(t, d, len(flows)) }