From 43ba9d5c6a3bdb3ef88237730158c094865f4e55 Mon Sep 17 00:00:00 2001
From: Henry Dollman <hank@henrygd.me>
Date: Sat, 14 Sep 2024 17:30:42 -0400
Subject: [PATCH] change NIC env var to NICS to support multiple interfaces

---
 beszel/internal/agent/agent.go | 34 ++++++++++++++++++++++++----------
 1 file changed, 24 insertions(+), 10 deletions(-)

diff --git a/beszel/internal/agent/agent.go b/beszel/internal/agent/agent.go
index b81c8c1..47f34b4 100644
--- a/beszel/internal/agent/agent.go
+++ b/beszel/internal/agent/agent.go
@@ -33,7 +33,6 @@ import (
 
 type Agent struct {
 	addr                string
-	nic                 string
 	pubKey              []byte
 	sem                 chan struct{}
 	containerStatsMap   map[string]*container.PrevContainerStats
@@ -434,16 +433,11 @@ func (a *Agent) Run() {
 	}
 
 	if extraFilesystems, exists := os.LookupEnv("EXTRA_FILESYSTEMS"); exists {
-		// parse comma separated list of filesystems
 		for _, filesystem := range strings.Split(extraFilesystems, ",") {
 			a.fsStats[filesystem] = &system.FsStats{}
 		}
 	}
 
-	if nic, exists := os.LookupEnv("NIC"); exists {
-		a.nic = nic
-	}
-
 	a.initializeDiskInfo(fsEnvVarExists)
 	a.initializeDiskIoStats()
 	a.initializeNetIoStats()
@@ -534,15 +528,36 @@ func (a *Agent) initializeDiskIoStats() {
 func (a *Agent) initializeNetIoStats() {
 	// reset valid network interfaces
 	a.netInterfaces = make(map[string]struct{}, 0)
+
+	// map of network interface names passed in via NICS env var
+	var nicsMap map[string]struct{}
+	nics, nicsEnvExists := os.LookupEnv("NICS")
+	if nicsEnvExists {
+		nicsMap = make(map[string]struct{}, 0)
+		for _, nic := range strings.Split(nics, ",") {
+			nicsMap[nic] = struct{}{}
+		}
+	}
+
 	// reset network I/O stats
 	a.netIoStats.BytesSent = 0
 	a.netIoStats.BytesRecv = 0
+
 	// get intial network I/O stats
 	if netIO, err := psutilNet.IOCounters(true); err == nil {
 		a.netIoStats.Time = time.Now()
 		for _, v := range netIO {
-			if a.skipNetworkInterface(v) {
-				continue
+			switch {
+			// skip if nics exists and the interface is not in the list
+			case nicsEnvExists:
+				if _, nameInNics := nicsMap[v.Name]; !nameInNics {
+					continue
+				}
+			// otherwise run the interface name through the skipNetworkInterface function
+			default:
+				if a.skipNetworkInterface(v) {
+					continue
+				}
 			}
 			log.Printf("Detected network interface: %+v (%+v recv, %+v sent)\n", v.Name, v.BytesRecv, v.BytesSent)
 			a.netIoStats.BytesSent += v.BytesSent
@@ -567,8 +582,7 @@ func twoDecimals(value float64) float64 {
 
 func (a *Agent) skipNetworkInterface(v psutilNet.IOCountersStat) bool {
 	switch {
-	case a.nic != "" && v.Name != a.nic,
-		strings.HasPrefix(v.Name, "lo"),
+	case strings.HasPrefix(v.Name, "lo"),
 		strings.HasPrefix(v.Name, "docker"),
 		strings.HasPrefix(v.Name, "br-"),
 		strings.HasPrefix(v.Name, "veth"),