3 files changed, 99 insertions, 21 deletions
diff --git a/pkg/tcpip/stack/nic.go b/pkg/tcpip/stack/nic.go
index e251e3b24..b9b5c35c8 100644
--- a/pkg/tcpip/stack/nic.go
+++ b/pkg/tcpip/stack/nic.go
@@ -727,7 +727,7 @@ func (n *nic) DeliverNetworkPacket(remote, local tcpip.LinkAddress, protocol tcp
 
 	networkEndpoint, ok := n.networkEndpoints[protocol]
 	if !ok {
-		n.stats.unknownL3ProtocolRcvdPackets.Increment()
+		n.stats.unknownL3ProtocolRcvdPacketCounts.Increment(uint64(protocol))
 		return
 	}
 
@@ -827,7 +827,7 @@ func (n *nic) deliverOutboundPacket(remote tcpip.LinkAddress, pkt *PacketBuffer)
 func (n *nic) DeliverTransportPacket(protocol tcpip.TransportProtocolNumber, pkt *PacketBuffer) TransportPacketDisposition {
 	state, ok := n.stack.transportProtocols[protocol]
 	if !ok {
-		n.stats.unknownL4ProtocolRcvdPackets.Increment()
+		n.stats.unknownL4ProtocolRcvdPacketCounts.Increment(uint64(protocol))
 		return TransportPacketProtocolUnreachable
 	}
 
diff --git a/pkg/tcpip/stack/nic_stats.go b/pkg/tcpip/stack/nic_stats.go
index 1773d5e8d..89aa95131 100644
--- a/pkg/tcpip/stack/nic_stats.go
+++ b/pkg/tcpip/stack/nic_stats.go
@@ -35,7 +35,7 @@ func (m *multiCounterNICPacketStats) init(a, b *tcpip.NICPacketStats) {
 	m.bytes.Init(a.Bytes, b.Bytes)
 }
 
-// LINT.ThenChange(../../tcpip.go:NICPacketStats)
+// LINT.ThenChange(../tcpip.go:NICPacketStats)
 
 // LINT.IfChange(multiCounterNICNeighborStats)
 
@@ -47,23 +47,23 @@ func (m *multiCounterNICNeighborStats) init(a, b *tcpip.NICNeighborStats) {
 	m.unreachableEntryLookups.Init(a.UnreachableEntryLookups, b.UnreachableEntryLookups)
 }
 
-// LINT.ThenChange(../../tcpip.go:NICNeighborStats)
+// LINT.ThenChange(../tcpip.go:NICNeighborStats)
 
 // LINT.IfChange(multiCounterNICStats)
 
 type multiCounterNICStats struct {
-	unknownL3ProtocolRcvdPackets tcpip.MultiCounterStat
-	unknownL4ProtocolRcvdPackets tcpip.MultiCounterStat
-	malformedL4RcvdPackets       tcpip.MultiCounterStat
-	tx                           multiCounterNICPacketStats
-	rx                           multiCounterNICPacketStats
-	disabledRx                   multiCounterNICPacketStats
-	neighbor                     multiCounterNICNeighborStats
+	unknownL3ProtocolRcvdPacketCounts tcpip.MultiIntegralStatCounterMap
+	unknownL4ProtocolRcvdPacketCounts tcpip.MultiIntegralStatCounterMap
+	malformedL4RcvdPackets            tcpip.MultiCounterStat
+	tx                                multiCounterNICPacketStats
+	rx                                multiCounterNICPacketStats
+	disabledRx                        multiCounterNICPacketStats
+	neighbor                          multiCounterNICNeighborStats
 }
 
 func (m *multiCounterNICStats) init(a, b *tcpip.NICStats) {
-	m.unknownL3ProtocolRcvdPackets.Init(a.UnknownL3ProtocolRcvdPackets, b.UnknownL3ProtocolRcvdPackets)
-	m.unknownL4ProtocolRcvdPackets.Init(a.UnknownL4ProtocolRcvdPackets, b.UnknownL4ProtocolRcvdPackets)
+	m.unknownL3ProtocolRcvdPacketCounts.Init(a.UnknownL3ProtocolRcvdPacketCounts, b.UnknownL3ProtocolRcvdPacketCounts)
+	m.unknownL4ProtocolRcvdPacketCounts.Init(a.UnknownL4ProtocolRcvdPacketCounts, b.UnknownL4ProtocolRcvdPacketCounts)
 	m.malformedL4RcvdPackets.Init(a.MalformedL4RcvdPackets, b.MalformedL4RcvdPackets)
 	m.tx.init(&a.Tx, &b.Tx)
 	m.rx.init(&a.Rx, &b.Rx)
@@ -71,4 +71,4 @@ func (m *multiCounterNICStats) init(a, b *tcpip.NICStats) {
 	m.neighbor.init(&a.Neighbor, &b.Neighbor)
 }
 
-// LINT.ThenChange(../../tcpip.go:NICStats)
+// LINT.ThenChange(../tcpip.go:NICStats)
diff --git a/pkg/tcpip/tcpip.go b/pkg/tcpip/tcpip.go
index 460a6afaf..d20dd495c 100644
--- a/pkg/tcpip/tcpip.go
+++ b/pkg/tcpip/tcpip.go
@@ -1301,7 +1301,8 @@ func (s *StatCounter) String() string {
 
 // A MultiCounterStat keeps track of two counters at once.
 type MultiCounterStat struct {
-	a, b *StatCounter
+	a *StatCounter
+	b *StatCounter
 }
 
 // Init sets both internal counters to point to a and b.
@@ -1923,17 +1924,89 @@ type NICPacketStats struct {
 	// LINT.ThenChange(stack/nic_stats.go:multiCounterNICPacketStats)
 }
 
+// IntegralStatCounterMap holds a map associating integral keys with
+// StatCounters.
+type IntegralStatCounterMap struct {
+	mu sync.RWMutex
+	// +checklocks:mu
+	counterMap map[uint64]*StatCounter
+}
+
+// Keys returns all keys present in the map.
+func (m *IntegralStatCounterMap) Keys() []uint64 {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	var keys []uint64
+	for k := range m.counterMap {
+		keys = append(keys, k)
+	}
+	return keys
+}
+
+// Get returns the counter mapped by the provided key.
+func (m *IntegralStatCounterMap) Get(key uint64) (*StatCounter, bool) {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	counter, ok := m.counterMap[key]
+	return counter, ok
+}
+
+// Init initializes the map.
+func (m *IntegralStatCounterMap) Init() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.counterMap = make(map[uint64]*StatCounter)
+}
+
+// Increment increments the counter associated with the provided key.
+func (m *IntegralStatCounterMap) Increment(key uint64) {
+	m.mu.RLock()
+	counter, ok := m.counterMap[key]
+	m.mu.RUnlock()
+
+	if !ok {
+		m.mu.Lock()
+		counter, ok = m.counterMap[key]
+		if !ok {
+			counter = new(StatCounter)
+			m.counterMap[key] = counter
+		}
+		m.mu.Unlock()
+	}
+	counter.Increment()
+}
+
+// A MultiIntegralStatCounterMap keeps track of two integral counter maps at
+// once.
+type MultiIntegralStatCounterMap struct {
+	a *IntegralStatCounterMap
+	b *IntegralStatCounterMap
+}
+
+// Init sets the internal integral counter maps to point to a and b.
+func (m *MultiIntegralStatCounterMap) Init(a, b *IntegralStatCounterMap) {
+	m.a = a
+	m.b = b
+}
+
+// Increment increments the counter in each map corresponding to the
+// provided key.
+func (m *MultiIntegralStatCounterMap) Increment(key uint64) {
+	m.a.Increment(key)
+	m.b.Increment(key)
+}
+
 // NICStats holds NIC statistics.
 type NICStats struct {
 	// LINT.IfChange(NICStats)
 
-	// UnknownL3ProtocolRcvdPackets is the number of packets received that were
-	// for an unknown or unsupported network protocol.
-	UnknownL3ProtocolRcvdPackets *StatCounter
+	// UnknownL3ProtocolRcvdPacketCounts records the number of packets recieved
+	// for each unknown or unsupported netowrk protocol number.
+	UnknownL3ProtocolRcvdPacketCounts *IntegralStatCounterMap
 
-	// UnknownL4ProtocolRcvdPackets is the number of packets received that were
-	// for an unknown or unsupported transport protocol.
-	UnknownL4ProtocolRcvdPackets *StatCounter
+	// UnknownL4ProtocolRcvdPacketCounts records the number of packets recieved
+	// for each unknown or unsupported transport protocol number.
+	UnknownL4ProtocolRcvdPacketCounts *IntegralStatCounterMap
 
 	// MalformedL4RcvdPackets is the number of packets received by a NIC that
 	// could not be delivered to a transport endpoint because the L4 header could
@@ -2103,6 +2176,11 @@ func InitStatCounters(v reflect.Value) {
 			if *s == nil {
 				*s = new(StatCounter)
 			}
+		} else if s, ok := v.Addr().Interface().(**IntegralStatCounterMap); ok {
+			if *s == nil {
+				*s = new(IntegralStatCounterMap)
+				(*s).Init()
+			}
 		} else {
 			InitStatCounters(v)
 		}