chore: initial in-memory state work for pc3

cmd/pc3/main.go

@@ -80,6 +80,9 @@ func main() {
 	// Main body.
 	//
 
+	// Create a state storage struct.
+	s := MkState()
+
 	// Use pmanager non-pinecone webserver to host web UI and an API to communicate with it.
 	ui, err := fs.Sub(ui, "ui/dist")
 	if err != nil {
@@ -186,14 +189,28 @@ mainloop:
 				// This shouldn't happen, but if the peer public key is
 				// malformed then we have no choice but to ignore the
 				// packet.
-				continue
+				continue mainloop
 			}
 
 			switch packet.Route {
 			case proto.ROUTE_HEARTBEAT:
 				packetData := proto.PacketHeartbeat{}
-				proto.Unmarshal(&packetData, peerPublicKey, packet.Data)
-				fmt.Printf("%v\n", packetData)
+				err = proto.Unmarshal(&packetData, peerPublicKey, packet.Data)
+				if err != nil {
+					// The packet data is malformed, there is nothing more we
+					// can do.
+					continue mainloop
+				}
+				s.Heartbeat(packet.Peer, packetData)
+			case proto.ROUTE_RESPONSE:
+				packetData := proto.PacketRes{}
+				err = proto.Unmarshal(&packetData, peerPublicKey, packet.Data)
+				if err != nil {
+					// The packet data is malformed, there is nothing more we
+					// can do.
+					continue mainloop
+				}
+				s.Response(packetData)
 			}
 		}
 	}

cmd/pc3/state.go

@@ -0,0 +1,82 @@
+package main
+
+import (
+	"sync"
+	"time"
+
+	"dev.l1qu1d.net/wraith-labs/wraith-module-pinecomms/internal/proto"
+	"github.com/google/uuid"
+)
+
+func MkState() *state {
+	s := state{
+		clients:  map[string]client{},
+		requests: map[string]request{},
+	}
+	return &s
+}
+
+type client struct {
+	lastHeartbeatTime time.Time
+	lastHeartbeat     proto.PacketHeartbeat
+}
+
+type request struct {
+	requestTime time.Time
+	request     proto.PacketReq
+
+	responseTime time.Time
+	response     proto.PacketRes
+}
+
+type state struct {
+	// List of "connected" Wraith clients.
+	clients      map[string]client
+	clientsMutex sync.RWMutex
+
+	// List of request/response pairs.
+	requests      map[string]request
+	requestsMutex sync.RWMutex
+}
+
+// Save/update a Wraith client entry.
+func (s *state) Heartbeat(src string, hb proto.PacketHeartbeat) {
+	s.clientsMutex.Lock()
+	defer s.clientsMutex.Unlock()
+
+	s.clients[src] = client{
+		lastHeartbeatTime: time.Now(),
+		lastHeartbeat:     hb,
+	}
+}
+
+// Save a request and generate a TxId.
+func (s *state) Request(req proto.PacketReq) proto.PacketReq {
+	reqTxId := uuid.NewString()
+	req.TxId = reqTxId
+
+	s.requestsMutex.Lock()
+	defer s.requestsMutex.Unlock()
+
+	s.requests[reqTxId] = request{
+		requestTime: time.Now(),
+		request:     req,
+	}
+
+	return req
+}
+
+// Save a response to a request.
+func (s *state) Response(res proto.PacketRes) {
+	s.requestsMutex.Lock()
+	defer s.requestsMutex.Unlock()
+
+	if req, ok := s.requests[res.TxId]; ok {
+		req.responseTime = time.Now()
+		req.response = res
+		s.requests[res.TxId] = req
+	}
+}
+
+// Expire timed-out entries in the state.
+func (s *state) Prune() {}