Add layer 2 stats (tx, rx) X (packets, bytes) to netstack

PiperOrigin-RevId: 239194420
Change-Id: Ie193e8ac2b7a6db21195ac85824a335930483971

4 files changed, 94 insertions, 8 deletions

diff --git a/pkg/tcpip/stack/nic.go b/pkg/tcpip/stack/nic.go
index defa8102a..1d032ebf8 100644
--- a/pkg/tcpip/stack/nic.go
+++ b/pkg/tcpip/stack/nic.go
@@ -42,6 +42,20 @@ type NIC struct {
 	primary     map[tcpip.NetworkProtocolNumber]*ilist.List
 	endpoints   map[NetworkEndpointID]*referencedNetworkEndpoint
 	subnets     []tcpip.Subnet
+
+	stats NICStats
+}
+
+// NICStats includes transmitted and received stats.
+type NICStats struct {
+	Tx DirectionStats
+	Rx DirectionStats
+}
+
+// DirectionStats includes packet and byte counts.
+type DirectionStats struct {
+	Packets *tcpip.StatCounter
+	Bytes   *tcpip.StatCounter
 }
 
 // PrimaryEndpointBehavior is an enumeration of an endpoint's primacy behavior.
@@ -73,6 +87,16 @@ func newNIC(stack *Stack, id tcpip.NICID, name string, ep LinkEndpoint, loopback
 		demux:     newTransportDemuxer(stack),
 		primary:   make(map[tcpip.NetworkProtocolNumber]*ilist.List),
 		endpoints: make(map[NetworkEndpointID]*referencedNetworkEndpoint),
+		stats: NICStats{
+			Tx: DirectionStats{
+				Packets: &tcpip.StatCounter{},
+				Bytes:   &tcpip.StatCounter{},
+			},
+			Rx: DirectionStats{
+				Packets: &tcpip.StatCounter{},
+				Bytes:   &tcpip.StatCounter{},
+			},
+		},
 	}
 }
 
@@ -384,6 +408,9 @@ func (n *NIC) RemoveAddress(addr tcpip.Address) *tcpip.Error {
 // This rule applies only to the slice itself, not to the items of the slice;
 // the ownership of the items is not retained by the caller.
 func (n *NIC) DeliverNetworkPacket(linkEP LinkEndpoint, remote, _ tcpip.LinkAddress, protocol tcpip.NetworkProtocolNumber, vv buffer.VectorisedView) {
+	n.stats.Rx.Packets.Increment()
+	n.stats.Rx.Bytes.IncrementBy(uint64(vv.Size()))
+
 	netProto, ok := n.stack.networkProtocols[protocol]
 	if !ok {
 		n.stack.stats.UnknownProtocolRcvdPackets.Increment()
@@ -457,7 +484,14 @@ func (n *NIC) DeliverNetworkPacket(linkEP LinkEndpoint, remote, _ tcpip.LinkAddr
 			// Send the packet out of n.
 			hdr := buffer.NewPrependableFromView(vv.First())
 			vv.RemoveFirst()
-			n.linkEP.WritePacket(&r, hdr, vv, protocol)
+
+			// TODO: use route.WritePacket.
+			if err := n.linkEP.WritePacket(&r, hdr, vv, protocol); err != nil {
+				r.Stats().IP.OutgoingPacketErrors.Increment()
+			} else {
+				n.stats.Tx.Packets.Increment()
+				n.stats.Tx.Bytes.IncrementBy(uint64(hdr.UsedLength() + vv.Size()))
+			}
 		}
 		return
 	}
diff --git a/pkg/tcpip/stack/route.go b/pkg/tcpip/stack/route.go
index 86fb728b2..3f2264d16 100644
--- a/pkg/tcpip/stack/route.go
+++ b/pkg/tcpip/stack/route.go
@@ -146,8 +146,11 @@ func (r *Route) IsResolutionRequired() bool {
 // WritePacket writes the packet through the given route.
 func (r *Route) WritePacket(hdr buffer.Prependable, payload buffer.VectorisedView, protocol tcpip.TransportProtocolNumber, ttl uint8) *tcpip.Error {
 	err := r.ref.ep.WritePacket(r, hdr, payload, protocol, ttl, r.loop)
-	if err == tcpip.ErrNoRoute {
+	if err != nil {
 		r.Stats().IP.OutgoingPacketErrors.Increment()
+	} else {
+		r.ref.nic.stats.Tx.Packets.Increment()
+		r.ref.nic.stats.Tx.Bytes.IncrementBy(uint64(hdr.UsedLength() + payload.Size()))
 	}
 	return err
 }
diff --git a/pkg/tcpip/stack/stack.go b/pkg/tcpip/stack/stack.go
index cbfe5c3c7..15a268b10 100644
--- a/pkg/tcpip/stack/stack.go
+++ b/pkg/tcpip/stack/stack.go
@@ -627,6 +627,8 @@ type NICInfo struct {
 
 	// MTU is the maximum transmission unit.
 	MTU uint32
+
+	Stats NICStats
 }
 
 // NICInfo returns a map of NICIDs to their associated information.
@@ -648,6 +650,7 @@ func (s *Stack) NICInfo() map[tcpip.NICID]NICInfo {
 			ProtocolAddresses: nic.Addresses(),
 			Flags:             flags,
 			MTU:               nic.linkEP.MTU(),
+			Stats:             nic.stats,
 		}
 	}
 	return nics
diff --git a/pkg/tcpip/stack/stack_test.go b/pkg/tcpip/stack/stack_test.go
index b366de21d..da8269999 100644
--- a/pkg/tcpip/stack/stack_test.go
+++ b/pkg/tcpip/stack/stack_test.go
@@ -273,7 +273,7 @@ func TestNetworkReceive(t *testing.T) {
 	}
 }
 
-func sendTo(t *testing.T, s *stack.Stack, addr tcpip.Address) {
+func sendTo(t *testing.T, s *stack.Stack, addr tcpip.Address, payload buffer.View) {
 	r, err := s.FindRoute(0, "", addr, fakeNetNumber, false /* multicastLoop */)
 	if err != nil {
 		t.Fatalf("FindRoute failed: %v", err)
@@ -281,9 +281,8 @@ func sendTo(t *testing.T, s *stack.Stack, addr tcpip.Address) {
 	defer r.Release()
 
 	hdr := buffer.NewPrependable(int(r.MaxHeaderLength()))
-	if err := r.WritePacket(hdr, buffer.VectorisedView{}, fakeTransNumber, 123); err != nil {
+	if err := r.WritePacket(hdr, payload.ToVectorisedView(), fakeTransNumber, 123); err != nil {
 		t.Errorf("WritePacket failed: %v", err)
-		return
 	}
 }
 
@@ -304,7 +303,7 @@ func TestNetworkSend(t *testing.T) {
 	}
 
 	// Make sure that the link-layer endpoint received the outbound packet.
-	sendTo(t, s, "\x03")
+	sendTo(t, s, "\x03", nil)
 	if c := linkEP.Drain(); c != 1 {
 		t.Errorf("packetCount = %d, want %d", c, 1)
 	}
@@ -351,14 +350,14 @@ func TestNetworkSendMultiRoute(t *testing.T) {
 	})
 
 	// Send a packet to an odd destination.
-	sendTo(t, s, "\x05")
+	sendTo(t, s, "\x05", nil)
 
 	if c := linkEP1.Drain(); c != 1 {
 		t.Errorf("packetCount = %d, want %d", c, 1)
 	}
 
 	// Send a packet to an even destination.
-	sendTo(t, s, "\x06")
+	sendTo(t, s, "\x06", nil)
 
 	if c := linkEP2.Drain(); c != 1 {
 		t.Errorf("packetCount = %d, want %d", c, 1)
@@ -1055,6 +1054,44 @@ func TestGetMainNICAddressAddRemove(t *testing.T) {
 	}
 }
 
+func TestNICStats(t *testing.T) {
+	s := stack.New([]string{"fakeNet"}, nil, stack.Options{})
+	id1, linkEP1 := channel.New(10, defaultMTU, "")
+	if err := s.CreateNIC(1, id1); err != nil {
+		t.Fatalf("CreateNIC failed: %v", err)
+	}
+	if err := s.AddAddress(1, fakeNetNumber, "\x01"); err != nil {
+		t.Fatalf("AddAddress failed: %v", err)
+	}
+	// Route all packets for address \x01 to NIC 1.
+	s.SetRouteTable([]tcpip.Route{
+		{"\x01", "\xff", "\x00", 1},
+	})
+
+	// Send a packet to address 1.
+	buf := buffer.NewView(30)
+	linkEP1.Inject(fakeNetNumber, buf.ToVectorisedView())
+	if got, want := s.NICInfo()[1].Stats.Rx.Packets.Value(), uint64(1); got != want {
+		t.Errorf("got Rx.Packets.Value() = %d, want = %d", got, want)
+	}
+
+	if got, want := s.NICInfo()[1].Stats.Rx.Bytes.Value(), uint64(len(buf)); got != want {
+		t.Errorf("got Rx.Bytes.Value() = %d, want = %d", got, want)
+	}
+
+	payload := buffer.NewView(10)
+	// Write a packet out via the address for NIC 1
+	sendTo(t, s, "\x01", payload)
+	want := uint64(linkEP1.Drain())
+	if got := s.NICInfo()[1].Stats.Tx.Packets.Value(); got != want {
+		t.Errorf("got Tx.Packets.Value() = %d, linkEP1.Drain() = %d", got, want)
+	}
+
+	if got, want := s.NICInfo()[1].Stats.Tx.Bytes.Value(), uint64(len(payload)); got != want {
+		t.Errorf("got Tx.Bytes.Value() = %d, want = %d", got, want)
+	}
+}
+
 func TestNICForwarding(t *testing.T) {
 	// Create a stack with the fake network protocol, two NICs, each with
 	// an address.
@@ -1092,6 +1129,15 @@ func TestNICForwarding(t *testing.T) {
 	default:
 		t.Fatal("Packet not forwarded")
 	}
+
+	// Test that forwarding increments Tx stats correctly.
+	if got, want := s.NICInfo()[2].Stats.Tx.Packets.Value(), uint64(1); got != want {
+		t.Errorf("got Tx.Packets.Value() = %d, want = %d", got, want)
+	}
+
+	if got, want := s.NICInfo()[2].Stats.Tx.Bytes.Value(), uint64(len(buf)); got != want {
+		t.Errorf("got Tx.Bytes.Value() = %d, want = %d", got, want)
+	}
 }
 
 func init() {