chore: automatically prune state storage

cmd/pc3/conf.go

@@ -4,6 +4,9 @@ import (
 	"os"
 	"strconv"
 	"sync/atomic"
+	"time"
+
+	"dev.l1qu1d.net/wraith-labs/wraith-module-pinecomms/internal/proto"
 )
 
 const (
@@ -11,6 +14,10 @@ const (
 	DEFAULT_PANEL_ADMIN_TOKEN = "wraith!"
 
 	STARTING_ATTEMPTS_UNTIL_LOCKOUT = 5
+
+	STATE_CLEANUP_INTERVAL     = 30 * time.Second
+	STATE_CLIENT_EXPIRY_DELAY  = proto.HEARTBEAT_MARK_DEAD_DELAY
+	STATE_REQUEST_EXPIRY_DELAY = 10 * time.Minute
 )
 
 type Config struct {

cmd/pc3/main.go

@@ -15,6 +15,7 @@ import (
 	"os/signal"
 	"strings"
 	"syscall"
+	"time"
 
 	"dev.l1qu1d.net/wraith-labs/wraith-module-pinecomms/internal/pmanager"
 	"dev.l1qu1d.net/wraith-labs/wraith-module-pinecomms/internal/proto"
@@ -182,6 +183,9 @@ mainloop:
 		// Exit if requested.
 		case <-sigchan:
 			break mainloop
+		// Clean up state.
+		case <-time.After(STATE_CLEANUP_INTERVAL):
+			s.Prune()
 		// Process incoming packets.
 		case packet := <-recv:
 			peerPublicKey, err := hex.DecodeString(packet.Peer)

cmd/pc3/state.go

@@ -82,4 +82,35 @@ func (s *state) Response(src string, res proto.PacketRes) {
 }
 
 // Expire timed-out entries in the state.
-func (s *state) Prune() {}
+func (s *state) Prune() {
+	wg := sync.WaitGroup{}
+	wg.Add(2)
+
+	// Clean up expired client heartbeats.
+	go func() {
+		defer wg.Done()
+		s.clientsMutex.Lock()
+		defer s.clientsMutex.Unlock()
+
+		for id, c := range s.clients {
+			if time.Since(c.lastHeartbeatTime) > proto.HEARTBEAT_MARK_DEAD_DELAY {
+				delete(s.clients, id)
+			}
+		}
+	}()
+
+	// Clean up expired request-response pairs.
+	go func() {
+		defer wg.Done()
+		s.requestsMutex.Lock()
+		defer s.requestsMutex.Unlock()
+
+		for id, r := range s.requests {
+			if time.Since(r.requestTime) > STATE_REQUEST_EXPIRY_DELAY {
+				delete(s.requests, id)
+			}
+		}
+	}()
+
+	wg.Wait()
+}