From 9596e168878df2433b95e7e37c27e568ea96991d Mon Sep 17 00:00:00 2001
From: Pim van Pelt <pim@ipng.nl>
Date: Sat, 22 Nov 2025 06:18:01 +0100
Subject: [PATCH] A few cosmetic changes in session handling

---
 src/agentx/agentx.go | 13 ++++------
 src/ifmib/ifmib.go   | 56 +++++++++++++++++++++++++++++---------------
 2 files changed, 42 insertions(+), 27 deletions(-)

diff --git a/src/agentx/agentx.go b/src/agentx/agentx.go
index 6768f5d..dd93296 100644
--- a/src/agentx/agentx.go
+++ b/src/agentx/agentx.go
@@ -20,18 +20,15 @@ var (
 
 // StartAgentXRoutine initializes the AgentX client and registers the interface MIB
 func StartAgentXRoutine(interfaceMIB *ifmib.InterfaceMIB) error {
-	var network, address string
+	// Determine network type based on address format
+	network := "tcp"
 	if strings.HasPrefix(*AgentXAddr, "/") {
 		network = "unix"
-		address = *AgentXAddr
-	} else {
-		network = "tcp"
-		address = *AgentXAddr
 	}
 
-	logger.Debugf("Connecting to AgentX at %s://%s", network, address)
+	logger.Debugf("Connecting to AgentX at %s://%s", network, *AgentXAddr)
 
-	client, err := agentx.Dial(network, address,
+	client, err := agentx.Dial(network, *AgentXAddr,
 		agentx.WithTimeout(1*time.Minute),
 		agentx.WithReconnectInterval(1*time.Second),
 	)
@@ -44,6 +41,6 @@ func StartAgentXRoutine(interfaceMIB *ifmib.InterfaceMIB) error {
 		return err
 	}
 
-	logger.Printf("Successfully registered with AgentX at %s://%s", network, address)
+	logger.Printf("Successfully registered with AgentX at %s://%s", network, *AgentXAddr)
 	return nil
 }
diff --git a/src/ifmib/ifmib.go b/src/ifmib/ifmib.go
index 8a3f4a2..26b8928 100644
--- a/src/ifmib/ifmib.go
+++ b/src/ifmib/ifmib.go
@@ -467,37 +467,55 @@ func (m *InterfaceMIB) addIfXTable(iface *api.InterfaceCounters, idx int) {
 	}
 }
 
+func (m *InterfaceMIB) createAndRegisterSession(client *agentx.Client, oid, name string) (*agentx.Session, error) {
+	session, err := client.Session(value.MustParseOID(oid), name, m.handler)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create %s session: %v", name, err)
+	}
+
+	err = session.Register(127, value.MustParseOID(oid))
+	if err != nil {
+		session.Close()
+		return nil, fmt.Errorf("failed to register %s: %v", name, err)
+	}
+
+	return session, nil
+}
+
 func (m *InterfaceMIB) RegisterWithClient(client *agentx.Client) error {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	// Create separate sessions for each MIB with the handler
-	ifEntrySession, err := client.Session(value.MustParseOID(ifEntryOID), "ifEntry", m.handler)
+	// Create and register sessions
+	ifEntrySession, err := m.createAndRegisterSession(client, ifEntryOID, "ifEntry")
 	if err != nil {
-		return fmt.Errorf("failed to create ifEntry session: %v", err)
+		return err
 	}
 
-	ifXTableSession, err := client.Session(value.MustParseOID(ifXTableOID), "ifXTable", m.handler)
+	ifXTableSession, err := m.createAndRegisterSession(client, ifXTableOID, "ifXTable")
 	if err != nil {
-		return fmt.Errorf("failed to create ifXTable session: %v", err)
+		ifEntrySession.Close()
+		return err
 	}
 
 	m.ifEntrySession = ifEntrySession
 	m.ifXTableSession = ifXTableSession
 
-	// Register the classic ifEntry
-	err = ifEntrySession.Register(127, value.MustParseOID(ifEntryOID))
-	if err != nil {
-		return fmt.Errorf("failed to register ifEntry: %v", err)
-	}
-
-	// Register the extended ifXTable
-	err = ifXTableSession.Register(127, value.MustParseOID(ifXTableOID))
-	if err != nil {
-		return fmt.Errorf("failed to register ifXTable: %v", err)
-	}
-
-	logger.Debugf("Registered IF-MIB ifEntry at OID %s", ifEntryOID)
-	logger.Debugf("Registered IF-MIB ifXTable at OID %s", ifXTableOID)
+	logger.Debugf("Registered IF-MIB sessions: ifEntry (%s) and ifXTable (%s)", ifEntryOID, ifXTableOID)
 	return nil
 }
+
+// Close cleans up AgentX sessions
+func (m *InterfaceMIB) Close() {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	if m.ifEntrySession != nil {
+		m.ifEntrySession.Close()
+		m.ifEntrySession = nil
+	}
+	if m.ifXTableSession != nil {
+		m.ifXTableSession.Close()
+		m.ifXTableSession = nil
+	}
+}