delegate proxy check performance increases

mythic-docker/src/rabbitmq/util_agent_message.go

@@ -11,7 +11,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/its-a-feature/Mythic/eventing"
 	"github.com/mitchellh/mapstructure"
 
 	"github.com/google/uuid"
@@ -791,7 +790,7 @@ func LookupEncryptionData(c2profile string, messageUUID string, updateCheckinTim
 	payload := databaseStructs.Payload{}
 	stager := databaseStructs.Staginginfo{}
 	if err := database.DB.Get(&callback, `SELECT
-		callback.id, callback.enc_key, callback.dec_key, callback.crypto_type, callback.operation_id, callback.last_checkin, callback.display_id, callback.trigger_on_checkin_after_time, 
+		callback.id, callback.enc_key, callback.dec_key, callback.crypto_type, callback.operation_id, callback.last_checkin, callback.display_id, callback.trigger_on_checkin_after_time, callback.active,
 		payload.id "payload.id", 
 		payloadtype.id "payload.payloadtype.id", 
 		payloadtype.name "payload.payloadtype.name", 
@@ -824,6 +823,7 @@ func LookupEncryptionData(c2profile string, messageUUID string, updateCheckinTim
 		newCache.LastCheckinTime = callback.LastCheckin
 		newCache.OperationID = callback.OperationID
 		newCache.TriggerOnCheckinAfterTime = callback.TriggerOnCheckinAfterTime
+		newCache.Active = callback.Active
 	} else if err = database.DB.Get(&payload, `SELECT
 		payload.id, payload.operation_id,
 		payload.deleted, payload.description, payload.uuid, payload.callback_allowed,
@@ -1244,126 +1244,6 @@ func reflectBackOtherKeys(response *map[string]interface{}, other *map[string]in
 	}
 }
 
-var updateCheckinTimeChannel = make(chan *cachedUUIDInfo, 200)
-
-func updateCheckinTimeEverySecond() {
-	callbackIDMap := make(map[int]*cachedUUIDInfo)
-	shouldUpdateChan := make(chan bool)
-	go func() {
-		for {
-			select {
-			case uuidInfo := <-updateCheckinTimeChannel:
-				callbackIDMap[uuidInfo.CallbackID] = uuidInfo
-			case <-shouldUpdateChan:
-				if len(callbackIDMap) == 0 {
-					continue
-				}
-				callbackIDs := make([]int, len(callbackIDMap))
-				i := 0
-				for _, callbackInfo := range callbackIDMap {
-					callbackIDs[i] = callbackInfo.CallbackID
-					callbackStruct := databaseStructs.Callback{
-						AgentCallbackID: callbackInfo.UUID,
-						ID:              callbackInfo.CallbackID,
-						OperationID:     callbackInfo.OperationID,
-						LastCheckin:     time.Now().UTC(),
-					}
-					callbackGraph.Add(callbackStruct, callbackStruct, callbackInfo.C2ProfileName, false)
-					i++
-				}
-				//logging.LogInfo("updating checkin times", "callbacks", callbackIDs)
-				//updateTimes(time.Now().UTC(), callbackIDs)
-				callbackIDMap = make(map[int]*cachedUUIDInfo)
-			}
-		}
-	}()
-	for {
-		select {
-		case <-time.After(time.Second):
-			shouldUpdateChan <- true
-		}
-	}
-}
-func UpdateCallbackEdgesAndCheckinTime(uuidInfo *cachedUUIDInfo) {
-	uuidInfo.stateMutex.Lock()
-	defer uuidInfo.stateMutex.Unlock()
-
-	callback := databaseStructs.Callback{
-		AgentCallbackID: uuidInfo.UUID,
-		ID:              uuidInfo.CallbackID,
-		OperationID:     uuidInfo.OperationID,
-		LastCheckin:     time.Now().UTC(),
-	}
-	// only bother updating the last checkin time if it's been more than one second
-	if callback.LastCheckin.Sub(uuidInfo.LastCheckinTime).Seconds() > 1 {
-		previousCheckin := uuidInfo.LastCheckinTime
-		select {
-		case updateCheckinTimeChannel <- uuidInfo:
-		default:
-			logging.LogDebug("Skipping callback checkin update enqueue because channel is full", "callback_id", uuidInfo.CallbackID)
-		}
-		//callbackGraph.AddByAgentIds(callback.AgentCallbackID, callback.AgentCallbackID, uuidInfo.C2ProfileName)
-		if uuidInfo.EdgeId == 0 {
-			err := database.DB.Get(&uuidInfo.EdgeId, `SELECT id FROM callbackgraphedge
-						WHERE source_id=$1 AND destination_id=$2 AND c2_profile_id=$3 AND operation_id=$4`,
-				uuidInfo.CallbackID, uuidInfo.CallbackID, uuidInfo.C2ProfileID, uuidInfo.OperationID)
-			if errors.Is(err, sql.ErrNoRows) {
-				if !uuidInfo.IsP2P {
-					err = database.DB.Get(&uuidInfo.EdgeId, `INSERT INTO callbackgraphedge
-						(source_id, destination_id, c2_profile_id, operation_id)
-						VALUES ($1, $1, $2, $3)
-						RETURNING id`,
-						uuidInfo.CallbackID, uuidInfo.C2ProfileID, uuidInfo.OperationID)
-					if err != nil {
-						logging.LogError(err, "Failed to add callback graph edge id for callback checking in",
-							"c2 id", uuidInfo.C2ProfileID, "callback id", uuidInfo.CallbackID)
-					} else {
-						logging.LogInfo("Added new callbackgraph edge when updating edges and checkin times", "c2", uuidInfo.C2ProfileID, "name", uuidInfo.C2ProfileName, "callback", uuidInfo.CallbackID)
-					}
-				}
-			} else if err != nil {
-				logging.LogError(err, "Failed to fetch callback graph edge id for callback checking in",
-					"c2 id", uuidInfo.C2ProfileID, "callback id", uuidInfo.CallbackID)
-			}
-		}
-		if !uuidInfo.Active {
-			uuidInfo.Active = true
-			_, err := database.DB.NamedExec(`UPDATE callback SET
-					active=true
-					WHERE id=:id`, callback)
-			if err != nil {
-				logging.LogError(err, "Failed to update active time", "callback", uuidInfo.UUID)
-			}
-			if uuidInfo.EdgeId > 0 {
-				_, err = database.DB.Exec(`UPDATE callbackgraphedge SET
-					end_timestamp=NULL
-					WHERE id=$1`, uuidInfo.EdgeId)
-				if err != nil {
-					logging.LogError(err, "Failed to callbackgraph edges time", "callback", uuidInfo.UUID)
-				}
-			}
-		}
-		if uuidInfo.TriggerOnCheckinAfterTime > 0 {
-			checkinDifference := int(callback.LastCheckin.Sub(previousCheckin).Minutes())
-			if checkinDifference >= uuidInfo.TriggerOnCheckinAfterTime {
-				// we want to trigger a workflow that the callback is checking in again after sleeping for > some time
-				go func(triggerData databaseStructs.Callback, oldCheckin time.Time, difference int) {
-					EventingChannel <- EventNotification{
-						Trigger:     eventing.TriggerCallbackCheckin,
-						OperationID: triggerData.OperationID,
-						CallbackID:  triggerData.ID,
-						Outputs: map[string]interface{}{
-							"previous_checkin":   oldCheckin,
-							"checkin_difference": difference,
-						},
-					}
-				}(callback, previousCheckin, checkinDifference)
-			}
-		}
-		uuidInfo.LastCheckinTime = callback.LastCheckin
-	}
-}
-
 func GetUUIDBytes(outerUUID string, agentUUIDLength int) ([]byte, error) {
 	switch agentUUIDLength {
 	case 36:

mythic-docker/src/rabbitmq/util_callback_graph.go

@@ -59,6 +59,18 @@ type bfsCacheEntry struct {
 
 var BFSCache bfsCache
 
+type callbackCheckinTargetCache struct {
+	cache map[int]callbackCheckinTargetCacheEntry
+	lock  sync.RWMutex
+}
+
+type callbackCheckinTargetCacheEntry struct {
+	graphVersion uint64
+	callbackIDs  []int
+}
+
+var callbackCheckinTargets callbackCheckinTargetCache
+
 func (c *bfsCache) GetPath(sourceId int, destinationId int, graphVersion uint64) ([]cbGraphAdjMatrixEntry, bool) {
 	c.lock.RLock()
 	defer c.lock.RUnlock()
@@ -91,18 +103,51 @@ func (c *bfsCache) Add(sourceId int, destinationId int, graphVersion uint64, bfs
 	}
 }
 
+func (c *callbackCheckinTargetCache) Get(callbackID int, graphVersion uint64) ([]int, bool) {
+	c.lock.RLock()
+	defer c.lock.RUnlock()
+	if cachedIDs, ok := c.cache[callbackID]; ok && cachedIDs.graphVersion == graphVersion {
+		return cloneTaskIDs(cachedIDs.callbackIDs), true
+	}
+	return nil, false
+}
+
+func (c *callbackCheckinTargetCache) Reset() {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	c.cache = make(map[int]callbackCheckinTargetCacheEntry)
+}
+
+func (c *callbackCheckinTargetCache) Add(callbackID int, graphVersion uint64, callbackIDs []int) {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	if c.cache == nil {
+		c.cache = make(map[int]callbackCheckinTargetCacheEntry)
+	}
+	c.cache[callbackID] = callbackCheckinTargetCacheEntry{
+		graphVersion: graphVersion,
+		callbackIDs:  cloneTaskIDs(callbackIDs),
+	}
+}
+
 func (g *cbGraph) bumpVersionLocked() {
 	g.version += 1
 	BFSCache.Reset()
+	callbackCheckinTargets.Reset()
 }
 
 func (g *cbGraph) getAllChildIDs(callbackId int) []int {
 	g.lock.RLock()
-	defer g.lock.RUnlock()
-	visitedIDs := map[int]bool{}
+	graphVersion := g.version
+	if cachedCallbackIDs, ok := callbackCheckinTargets.Get(callbackId, graphVersion); ok {
+		g.lock.RUnlock()
+		return cachedCallbackIDs
+	}
+	visitedIDs := map[int]bool{callbackId: true}
 	needToVisitIDs := []int{callbackId}
 	callbackIDsToUpdate := []int{callbackId}
-	callbacksWithEgress := []int{}
+	callbackIDsToUpdateMap := map[int]bool{callbackId: true}
+	callbacksWithEgress := map[int]bool{}
 	for len(needToVisitIDs) > 0 {
 		// get the next id we're going to check
 		currentId := needToVisitIDs[0]
@@ -117,7 +162,7 @@ func (g *cbGraph) getAllChildIDs(callbackId int) []int {
 		for i, _ := range immediateChildren {
 			if immediateChildren[i].SourceId == immediateChildren[i].DestinationId && immediateChildren[i].SourceId != callbackId {
 				//logging.LogInfo("found egress connection", "id", immediateChildren[i])
-				callbacksWithEgress = append(callbacksWithEgress, immediateChildren[i].SourceId)
+				callbacksWithEgress[immediateChildren[i].SourceId] = true
 			}
 			// check if we've already visited this id, if so, move on
 			if _, visited := visitedIDs[immediateChildren[i].SourceId]; !visited {
@@ -126,7 +171,10 @@ func (g *cbGraph) getAllChildIDs(callbackId int) []int {
 				if !isCallbackStreaming(immediateChildren[i].SourceId) {
 					// add this id as an id for the next iteration to check its children
 					needToVisitIDs = append(needToVisitIDs, immediateChildren[i].SourceId)
-					callbackIDsToUpdate = append(callbackIDsToUpdate, immediateChildren[i].SourceId)
+					if !callbackIDsToUpdateMap[immediateChildren[i].SourceId] {
+						callbackIDsToUpdate = append(callbackIDsToUpdate, immediateChildren[i].SourceId)
+						callbackIDsToUpdateMap[immediateChildren[i].SourceId] = true
+					}
 				}
 
 			}
@@ -136,7 +184,10 @@ func (g *cbGraph) getAllChildIDs(callbackId int) []int {
 				if !isCallbackStreaming(immediateChildren[i].DestinationId) {
 					// add this id as an id for the next iteration to check its children
 					needToVisitIDs = append(needToVisitIDs, immediateChildren[i].DestinationId)
-					callbackIDsToUpdate = append(callbackIDsToUpdate, immediateChildren[i].DestinationId)
+					if !callbackIDsToUpdateMap[immediateChildren[i].DestinationId] {
+						callbackIDsToUpdate = append(callbackIDsToUpdate, immediateChildren[i].DestinationId)
+						callbackIDsToUpdateMap[immediateChildren[i].DestinationId] = true
+					}
 				}
 			}
 		}
@@ -145,13 +196,15 @@ func (g *cbGraph) getAllChildIDs(callbackId int) []int {
 	finalCallbackIDsToUpdate := []int{}
 	for _, callbackIdToUpdate := range callbackIDsToUpdate {
 		//logging.LogInfo("checking if should update", "callbackIdToUpdate", callbackIdToUpdate)
-		if slices.Contains(callbacksWithEgress, callbackIdToUpdate) || slices.Contains(finalCallbackIDsToUpdate, callbackIdToUpdate) {
+		if callbacksWithEgress[callbackIdToUpdate] {
 			//logging.LogInfo("checking if should update", "its egress", true, "callbackIdToUpdate", callbackIdToUpdate)
 			continue
 		}
 		//logging.LogInfo("checking if should update", "its egress", false, "callbackIdToUpdate", callbackIdToUpdate)
 		finalCallbackIDsToUpdate = append(finalCallbackIDsToUpdate, callbackIdToUpdate)
 	}
+	g.lock.RUnlock()
+	callbackCheckinTargets.Add(callbackId, graphVersion, finalCallbackIDsToUpdate)
 	return finalCallbackIDsToUpdate
 }
 func updateTimes(updatedTime time.Time, callbackIDs []int) {
@@ -181,9 +234,11 @@ func listenForPushConnectDisconnectMessages() {
 	for {
 		select {
 		case connectCallbackId := <-pushC2StreamingConnectNotification:
+			callbackCheckinTargets.Reset()
 			callbackIDs := callbackGraph.getAllChildIDs(connectCallbackId)
 			updateTimes(time.UnixMicro(0), callbackIDs)
 		case disconnectCallbackId := <-pushC2StreamingDisconnectNotification:
+			callbackCheckinTargets.Reset()
 			callbackIDs := callbackGraph.getAllChildIDs(disconnectCallbackId)
 			updateTimes(time.Now().UTC(), callbackIDs)
 		}
@@ -214,6 +269,7 @@ func (g *cbGraph) Initialize() {
 		}
 	}
 	BFSCache.Reset()
+	callbackCheckinTargets.Reset()
 }
 func (g *cbGraph) Add(source databaseStructs.Callback, destination databaseStructs.Callback, c2profileName string, initializing bool) {
 	c2 := getC2ProfileForName(c2profileName)
@@ -241,20 +297,9 @@ func (g *cbGraph) Add(source databaseStructs.Callback, destination databaseStruc
 	for _, dest := range g.adjMatrix[source.ID] {
 		if dest.DestinationId == destination.ID && dest.C2ProfileName == c2profileName {
 			g.lock.Unlock()
-			//logging.LogDebug("Found existing connection, not adding new one to memory", "source", source.ID, "destination", destination.ID, "c2 profile", c2profileName)
-			if initializing || c2.IsP2p {
-				// don't update callback times when initializing, this is when the Mythic server starts up
-				return
-			}
-			updateTime := time.Now().UTC()
-			if isCallbackStreaming(source.ID) {
-				updateTime = time.UnixMicro(0)
-			}
-			callbackIDs := g.getAllChildIDs(source.ID)
-			if len(callbackIDs) > 0 {
-				//logging.LogInfo("about to call updateTimes in Add", "callbackIDs", callbackIDs, "c2", c2profileName, "source", source.ID, "destination", destination.ID)
-				updateTimes(updateTime, callbackIDs)
-			}
+			// Existing edges are now handled as a pure no-op. Check-in timestamp
+			// propagation is batched through updateCheckinTimeEverySecond instead
+			// of using Add as a side effect.
 			return
 		}
 	}

mythic-docker/src/rabbitmq/util_callback_graph_test.go

@@ -13,6 +13,7 @@ func newTestCallbackGraph(t *testing.T) *cbGraph {
 	t.Helper()
 	previousC2ProfileNameToIDMap := c2profileNameToIdMap
 	BFSCache.Reset()
+	callbackCheckinTargets.Reset()
 	c2profileNameToIdMap = map[string]databaseStructs.C2profile{
 		testC2ProfileName: {
 			ID:    1,
@@ -22,6 +23,7 @@ func newTestCallbackGraph(t *testing.T) *cbGraph {
 	}
 	t.Cleanup(func() {
 		BFSCache.Reset()
+		callbackCheckinTargets.Reset()
 		c2profileNameToIdMap = previousC2ProfileNameToIDMap
 	})
 	return &cbGraph{
@@ -147,3 +149,42 @@ func TestCanHaveDelegatesUsesCurrentGraphState(t *testing.T) {
 		t.Fatalf("expected callback with removed edge to have no delegates")
 	}
 }
+
+func TestGetAllChildIDsExcludesCallbacksWithTheirOwnEgress(t *testing.T) {
+	graph := newTestCallbackGraph(t)
+	addTestGraphEdge(graph, 1, 1)
+	addTestGraphEdge(graph, 1, 2)
+	addTestGraphEdge(graph, 2, 3)
+	addTestGraphEdge(graph, 2, 2)
+
+	callbackIDs := graph.getAllChildIDs(1)
+	expectedCallbackIDs := []int{1, 3}
+	if len(callbackIDs) != len(expectedCallbackIDs) {
+		t.Fatalf("expected callback IDs %#v, got %#v", expectedCallbackIDs, callbackIDs)
+	}
+	for i, expectedCallbackID := range expectedCallbackIDs {
+		if callbackIDs[i] != expectedCallbackID {
+			t.Fatalf("expected callbackIDs[%d] to be %d, got %d", i, expectedCallbackID, callbackIDs[i])
+		}
+	}
+}
+
+func TestGetAllChildIDsCacheInvalidatesWhenGraphChanges(t *testing.T) {
+	graph := newTestCallbackGraph(t)
+	addTestGraphEdge(graph, 1, 1)
+	addTestGraphEdge(graph, 1, 2)
+
+	callbackIDs := graph.getAllChildIDs(1)
+	if len(callbackIDs) != 2 {
+		t.Fatalf("expected two callback IDs before graph change, got %#v", callbackIDs)
+	}
+	if cachedIDs, cached := callbackCheckinTargets.Get(1, graph.version); !cached || len(cachedIDs) != 2 {
+		t.Fatalf("expected cached callback IDs before graph change, cached=%v ids=%#v", cached, cachedIDs)
+	}
+
+	addTestGraphEdge(graph, 1, 3)
+	callbackIDs = graph.getAllChildIDs(1)
+	if len(callbackIDs) != 3 {
+		t.Fatalf("expected three callback IDs after graph change, got %#v", callbackIDs)
+	}
+}

mythic-docker/src/rabbitmq/utils_callback_checkin.go

@@ -0,0 +1,200 @@
+package rabbitmq
+
+import (
+	"database/sql"
+	"errors"
+	"time"
+
+	"github.com/its-a-feature/Mythic/database"
+	databaseStructs "github.com/its-a-feature/Mythic/database/structs"
+	"github.com/its-a-feature/Mythic/eventing"
+	"github.com/its-a-feature/Mythic/logging"
+)
+
+type callbackCheckinUpdate struct {
+	CallbackID    int
+	OperationID   int
+	C2ProfileID   int
+	C2ProfileName string
+	UUID          string
+}
+
+var updateCheckinTimeChannel = make(chan callbackCheckinUpdate, 2000)
+
+// updateCheckinTimeEverySecond coalesces frequent callback check-ins and writes
+// last_checkin updates in batches. This keeps high-frequency polling from
+// repeatedly walking the callback graph and issuing per-callback database writes.
+func updateCheckinTimeEverySecond() {
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
+	callbackIDMap := make(map[int]callbackCheckinUpdate)
+
+	for {
+		select {
+		case update := <-updateCheckinTimeChannel:
+			callbackIDMap[update.CallbackID] = update
+		case <-ticker.C:
+			drainCallbackCheckinUpdateChannel(callbackIDMap)
+			flushCallbackCheckinUpdates(callbackIDMap)
+			callbackIDMap = make(map[int]callbackCheckinUpdate)
+		}
+	}
+}
+
+func drainCallbackCheckinUpdateChannel(callbackIDMap map[int]callbackCheckinUpdate) {
+	for {
+		select {
+		case update := <-updateCheckinTimeChannel:
+			callbackIDMap[update.CallbackID] = update
+		default:
+			return
+		}
+	}
+}
+
+func flushCallbackCheckinUpdates(callbackIDMap map[int]callbackCheckinUpdate) {
+	if len(callbackIDMap) == 0 {
+		return
+	}
+	standardCheckinIDs, streamingCheckinIDs := buildCallbackCheckinUpdateTargets(callbackIDMap)
+	for callbackID := range standardCheckinIDs {
+		delete(streamingCheckinIDs, callbackID)
+	}
+	if len(streamingCheckinIDs) > 0 {
+		updateTimes(time.UnixMicro(0), callbackIDMapKeys(streamingCheckinIDs))
+	}
+	if len(standardCheckinIDs) > 0 {
+		updateTimes(time.Now().UTC(), callbackIDMapKeys(standardCheckinIDs))
+	}
+}
+
+func buildCallbackCheckinUpdateTargets(callbackIDMap map[int]callbackCheckinUpdate) (map[int]bool, map[int]bool) {
+	standardCheckinIDs := make(map[int]bool)
+	streamingCheckinIDs := make(map[int]bool)
+	for _, update := range callbackIDMap {
+		targetIDs := callbackGraph.getAllChildIDs(update.CallbackID)
+		if isCallbackStreaming(update.CallbackID) {
+			for _, callbackID := range targetIDs {
+				streamingCheckinIDs[callbackID] = true
+			}
+			continue
+		}
+		for _, callbackID := range targetIDs {
+			standardCheckinIDs[callbackID] = true
+		}
+	}
+	return standardCheckinIDs, streamingCheckinIDs
+}
+
+func callbackIDMapKeys(callbackIDMap map[int]bool) []int {
+	callbackIDs := make([]int, 0, len(callbackIDMap))
+	for callbackID := range callbackIDMap {
+		callbackIDs = append(callbackIDs, callbackID)
+	}
+	return callbackIDs
+}
+
+func enqueueCallbackCheckinUpdate(update callbackCheckinUpdate) {
+	select {
+	case updateCheckinTimeChannel <- update:
+	default:
+		logging.LogDebug("Skipping callback checkin update enqueue because channel is full", "callback_id", update.CallbackID)
+	}
+}
+
+func ensureCallbackCheckinEdge(uuidInfo *cachedUUIDInfo) {
+	if uuidInfo.EdgeId != 0 || uuidInfo.IsP2P {
+		return
+	}
+	err := database.DB.Get(&uuidInfo.EdgeId, `SELECT id FROM callbackgraphedge
+		WHERE source_id=$1 AND destination_id=$2 AND c2_profile_id=$3 AND operation_id=$4`,
+		uuidInfo.CallbackID, uuidInfo.CallbackID, uuidInfo.C2ProfileID, uuidInfo.OperationID)
+	if errors.Is(err, sql.ErrNoRows) {
+		err = database.DB.Get(&uuidInfo.EdgeId, `INSERT INTO callbackgraphedge
+			(source_id, destination_id, c2_profile_id, operation_id)
+			VALUES ($1, $1, $2, $3)
+			RETURNING id`,
+			uuidInfo.CallbackID, uuidInfo.C2ProfileID, uuidInfo.OperationID)
+		if err != nil {
+			logging.LogError(err, "Failed to add callback graph edge id for callback checking in",
+				"c2 id", uuidInfo.C2ProfileID, "callback id", uuidInfo.CallbackID)
+			return
+		}
+		logging.LogInfo("Added new callbackgraph edge when updating edges and checkin times",
+			"c2", uuidInfo.C2ProfileID, "name", uuidInfo.C2ProfileName, "callback", uuidInfo.CallbackID)
+	} else if err != nil {
+		logging.LogError(err, "Failed to fetch callback graph edge id for callback checking in",
+			"c2 id", uuidInfo.C2ProfileID, "callback id", uuidInfo.CallbackID)
+		return
+	}
+	callback := databaseStructs.Callback{
+		AgentCallbackID: uuidInfo.UUID,
+		ID:              uuidInfo.CallbackID,
+		OperationID:     uuidInfo.OperationID,
+	}
+	callbackGraph.Add(callback, callback, uuidInfo.C2ProfileName, true)
+}
+
+func reactivateCallbackIfNeeded(uuidInfo *cachedUUIDInfo) {
+	if uuidInfo.Active {
+		return
+	}
+	uuidInfo.Active = true
+	_, err := database.DB.Exec(`UPDATE callback SET active=true WHERE id=$1`, uuidInfo.CallbackID)
+	if err != nil {
+		logging.LogError(err, "Failed to update active time", "callback", uuidInfo.UUID)
+	}
+	if uuidInfo.EdgeId > 0 {
+		_, err = database.DB.Exec(`UPDATE callbackgraphedge SET
+			end_timestamp=NULL
+			WHERE id=$1`, uuidInfo.EdgeId)
+		if err != nil {
+			logging.LogError(err, "Failed to callbackgraph edges time", "callback", uuidInfo.UUID)
+		}
+	}
+}
+
+func emitCallbackCheckinTriggerIfNeeded(uuidInfo *cachedUUIDInfo, previousCheckin time.Time, currentCheckin time.Time) {
+	if uuidInfo.TriggerOnCheckinAfterTime <= 0 {
+		return
+	}
+	checkinDifference := int(currentCheckin.Sub(previousCheckin).Minutes())
+	if checkinDifference < uuidInfo.TriggerOnCheckinAfterTime {
+		return
+	}
+	go func(operationID int, callbackID int, oldCheckin time.Time, difference int) {
+		EventingChannel <- EventNotification{
+			Trigger:     eventing.TriggerCallbackCheckin,
+			OperationID: operationID,
+			CallbackID:  callbackID,
+			Outputs: map[string]interface{}{
+				"previous_checkin":   oldCheckin,
+				"checkin_difference": difference,
+			},
+		}
+	}(uuidInfo.OperationID, uuidInfo.CallbackID, previousCheckin, checkinDifference)
+}
+
+func UpdateCallbackEdgesAndCheckinTime(uuidInfo *cachedUUIDInfo) {
+	uuidInfo.stateMutex.Lock()
+	defer uuidInfo.stateMutex.Unlock()
+
+	currentCheckin := time.Now().UTC()
+	// only bother updating the last checkin time if it's been more than one second
+	if currentCheckin.Sub(uuidInfo.LastCheckinTime).Seconds() <= 1 {
+		return
+	}
+
+	previousCheckin := uuidInfo.LastCheckinTime
+	enqueueCallbackCheckinUpdate(callbackCheckinUpdate{
+		CallbackID:    uuidInfo.CallbackID,
+		OperationID:   uuidInfo.OperationID,
+		C2ProfileID:   uuidInfo.C2ProfileID,
+		C2ProfileName: uuidInfo.C2ProfileName,
+		UUID:          uuidInfo.UUID,
+	})
+	ensureCallbackCheckinEdge(uuidInfo)
+	reactivateCallbackIfNeeded(uuidInfo)
+	emitCallbackCheckinTriggerIfNeeded(uuidInfo, previousCheckin, currentCheckin)
+	uuidInfo.LastCheckinTime = currentCheckin
+}

mythic-docker/src/rabbitmq/utils_proxy_traffic.go

@@ -26,10 +26,12 @@ import (
 type CallbackPortType = string
 
 const (
-	CALLBACK_PORT_TYPE_SOCKS       CallbackPortType = "socks"
-	CALLBACK_PORT_TYPE_RPORTFWD                     = "rpfwd"
-	CALLBACK_PORT_TYPE_INTERACTIVE                  = "interactive"
-	callbackPortByteFlushInterval                   = 20 * time.Second
+	CALLBACK_PORT_TYPE_SOCKS        CallbackPortType = "socks"
+	CALLBACK_PORT_TYPE_RPORTFWD                      = "rpfwd"
+	CALLBACK_PORT_TYPE_INTERACTIVE                   = "interactive"
+	callbackPortByteFlushInterval                    = 20 * time.Second
+	proxyFromAgentMessageShardCount                  = 16
+	proxyFromAgentMessageShardSize                   = 1000
 )
 
 type proxyToAgentMessage struct {
@@ -86,11 +88,12 @@ type callbackPortUsage struct {
 }
 
 type callbackPortsInUse struct {
-	ports                        []*callbackPortUsage
-	portsByCallbackID            map[int][]*callbackPortUsage
-	portsByCallbackIDAndType     map[int]map[CallbackPortType][]*callbackPortUsage
-	callbacksWithPorts           []int
-	proxyFromAgentMessageChannel chan ProxyFromAgentMessageForMythic
+	ports                             []*callbackPortUsage
+	portsByCallbackID                 map[int][]*callbackPortUsage
+	portsByCallbackIDAndType          map[int]map[CallbackPortType][]*callbackPortUsage
+	portsByCallbackIDTypeAndLocalPort map[int]map[CallbackPortType]map[int]*callbackPortUsage
+	callbacksWithPorts                []int
+	proxyFromAgentMessageChannels     []chan ProxyFromAgentMessageForMythic
 	sync.RWMutex
 }
 
@@ -219,9 +222,10 @@ func (c *callbackPortsInUse) Initialize() {
 	c.Lock()
 	c.ports = make([]*callbackPortUsage, 0)
 	c.resetPortIndexesLocked()
-	c.proxyFromAgentMessageChannel = make(chan ProxyFromAgentMessageForMythic, 2000)
+	c.proxyFromAgentMessageChannels = makeProxyFromAgentMessageChannels()
+	proxyFromAgentMessageChannels := c.proxyFromAgentMessageChannels
 	c.Unlock()
-	go c.ListenForProxyFromAgentMessage()
+	c.startProxyFromAgentMessageListeners(proxyFromAgentMessageChannels)
 	go c.ListenForNewByteTransferUpdates()
 	if err := database.DB.Select(&callbackPorts, `SELECT 
     	callbackport.id, callbackport.callback_id, callbackport.task_id, 
@@ -272,11 +276,12 @@ func (c *callbackPortsInUse) Initialize() {
 func (c *callbackPortsInUse) resetPortIndexesLocked() {
 	c.portsByCallbackID = make(map[int][]*callbackPortUsage)
 	c.portsByCallbackIDAndType = make(map[int]map[CallbackPortType][]*callbackPortUsage)
+	c.portsByCallbackIDTypeAndLocalPort = make(map[int]map[CallbackPortType]map[int]*callbackPortUsage)
 	c.callbacksWithPorts = make([]int, 0)
 }
 
 func (c *callbackPortsInUse) ensurePortIndexesLocked() {
-	if c.portsByCallbackID == nil || c.portsByCallbackIDAndType == nil {
+	if c.portsByCallbackID == nil || c.portsByCallbackIDAndType == nil || c.portsByCallbackIDTypeAndLocalPort == nil {
 		ports := c.ports
 		c.resetPortIndexesLocked()
 		for _, port := range ports {
@@ -285,6 +290,30 @@ func (c *callbackPortsInUse) ensurePortIndexesLocked() {
 	}
 }
 
+func makeProxyFromAgentMessageChannels() []chan ProxyFromAgentMessageForMythic {
+	channels := make([]chan ProxyFromAgentMessageForMythic, proxyFromAgentMessageShardCount)
+	for i := range channels {
+		channels[i] = make(chan ProxyFromAgentMessageForMythic, proxyFromAgentMessageShardSize)
+	}
+	return channels
+}
+
+func proxyFromAgentMessageShard(callbackID int, shardCount int) int {
+	if shardCount <= 0 {
+		return 0
+	}
+	if callbackID < 0 {
+		callbackID = -callbackID
+	}
+	return callbackID % shardCount
+}
+
+func (c *callbackPortsInUse) startProxyFromAgentMessageListeners(channels []chan ProxyFromAgentMessageForMythic) {
+	for _, channel := range channels {
+		go c.ListenForProxyFromAgentMessage(channel)
+	}
+}
+
 func (c *callbackPortsInUse) addPortLocked(port *callbackPortUsage) {
 	c.ensurePortIndexesLocked()
 	c.ports = append(c.ports, port)
@@ -300,6 +329,15 @@ func (c *callbackPortsInUse) addPortToIndexesLocked(port *callbackPortUsage) {
 		c.portsByCallbackIDAndType[port.CallbackID] = make(map[CallbackPortType][]*callbackPortUsage)
 	}
 	c.portsByCallbackIDAndType[port.CallbackID][port.PortType] = append(c.portsByCallbackIDAndType[port.CallbackID][port.PortType], port)
+	if _, ok := c.portsByCallbackIDTypeAndLocalPort[port.CallbackID]; !ok {
+		c.portsByCallbackIDTypeAndLocalPort[port.CallbackID] = make(map[CallbackPortType]map[int]*callbackPortUsage)
+	}
+	if _, ok := c.portsByCallbackIDTypeAndLocalPort[port.CallbackID][port.PortType]; !ok {
+		c.portsByCallbackIDTypeAndLocalPort[port.CallbackID][port.PortType] = make(map[int]*callbackPortUsage)
+	}
+	if _, ok := c.portsByCallbackIDTypeAndLocalPort[port.CallbackID][port.PortType][port.LocalPort]; !ok {
+		c.portsByCallbackIDTypeAndLocalPort[port.CallbackID][port.PortType][port.LocalPort] = port
+	}
 }
 
 func (c *callbackPortsInUse) removePortLocked(port *callbackPortUsage) bool {
@@ -323,12 +361,33 @@ func (c *callbackPortsInUse) removePortFromIndexesLocked(port *callbackPortUsage
 	if len(c.portsByCallbackID[port.CallbackID]) == 0 {
 		delete(c.portsByCallbackID, port.CallbackID)
 		delete(c.portsByCallbackIDAndType, port.CallbackID)
+		delete(c.portsByCallbackIDTypeAndLocalPort, port.CallbackID)
 		c.callbacksWithPorts = removeProxyCallbackID(c.callbacksWithPorts, port.CallbackID)
 		return
 	}
 	c.portsByCallbackIDAndType[port.CallbackID][port.PortType] = removeCallbackPortUsage(c.portsByCallbackIDAndType[port.CallbackID][port.PortType], port)
 	if len(c.portsByCallbackIDAndType[port.CallbackID][port.PortType]) == 0 {
 		delete(c.portsByCallbackIDAndType[port.CallbackID], port.PortType)
+		delete(c.portsByCallbackIDTypeAndLocalPort[port.CallbackID], port.PortType)
+		if len(c.portsByCallbackIDTypeAndLocalPort[port.CallbackID]) == 0 {
+			delete(c.portsByCallbackIDTypeAndLocalPort, port.CallbackID)
+		}
+		return
+	}
+	c.replaceLocalPortIndexAfterRemovalLocked(port)
+}
+
+func (c *callbackPortsInUse) replaceLocalPortIndexAfterRemovalLocked(port *callbackPortUsage) {
+	localPortIndex := c.portsByCallbackIDTypeAndLocalPort[port.CallbackID][port.PortType]
+	if localPortIndex[port.LocalPort] != port {
+		return
+	}
+	delete(localPortIndex, port.LocalPort)
+	for _, replacementPort := range c.portsByCallbackIDAndType[port.CallbackID][port.PortType] {
+		if replacementPort.LocalPort == port.LocalPort {
+			localPortIndex[port.LocalPort] = replacementPort
+			return
+		}
 	}
 }
 
@@ -544,6 +603,38 @@ func (c *callbackPortsInUse) getPortsForCallbackAndTypeFromSliceLocked(callbackI
 	return ports
 }
 
+func (c *callbackPortsInUse) getPortForCallbackTypeAndLocalPort(callbackID int, portType CallbackPortType, localPort int) *callbackPortUsage {
+	c.RLock()
+	defer c.RUnlock()
+	if localPort > 0 && c.portsByCallbackIDTypeAndLocalPort != nil {
+		return c.portsByCallbackIDTypeAndLocalPort[callbackID][portType][localPort]
+	}
+	if c.portsByCallbackIDAndType != nil {
+		ports := c.portsByCallbackIDAndType[callbackID][portType]
+		if localPort <= 0 {
+			if len(ports) == 0 {
+				return nil
+			}
+			return ports[0]
+		}
+		for _, port := range ports {
+			if port.LocalPort == localPort {
+				return port
+			}
+		}
+		return nil
+	}
+	for _, port := range c.ports {
+		if port.CallbackID != callbackID || port.PortType != portType {
+			continue
+		}
+		if localPort <= 0 || port.LocalPort == localPort {
+			return port
+		}
+	}
+	return nil
+}
+
 func (c *callbackPortsInUse) getPortsForCallbackByType(callbackID int, portTypes []CallbackPortType) map[CallbackPortType][]*callbackPortUsage {
 	portsByType := make(map[CallbackPortType][]*callbackPortUsage, len(portTypes))
 	c.RLock()
@@ -642,44 +733,33 @@ func updateCallbackPortStats(field string, value int64, callbackPortID int) bool
 	}
 	return true
 }
-func (c *callbackPortsInUse) ListenForProxyFromAgentMessage() {
-	for {
-		agentMessage := <-c.proxyFromAgentMessageChannel
-		switch agentMessage.PortType {
-		case CALLBACK_PORT_TYPE_RPORTFWD:
-			fallthrough
-		case CALLBACK_PORT_TYPE_SOCKS:
-			ports := c.getPortsForCallbackAndType(agentMessage.CallbackID, agentMessage.PortType)
-			// loop through each message and find the corresponding callback + local port combo
-			for j := 0; j < len(agentMessage.Messages); j++ {
-				for _, port := range ports {
-					if agentMessage.Messages[j].Port > 0 {
-						// port is specified, try to find the right one
-						if port.LocalPort == agentMessage.Messages[j].Port {
-							port.messagesFromAgent <- agentMessage.Messages[j]
-							break
-						}
-					} else {
-						// didn't specify a specific port, to just send it to the first matching type
-						port.messagesFromAgent <- agentMessage.Messages[j]
-						break
-					}
-				}
-				//logging.LogInfo("got message from agent", "chan", messages[j].ServerID, "p.messagesFromAgentQueue", len(c.ports[i].messagesFromAgent))
-				//c.ports[i].messagesFromAgent <- agentMessage.Messages[j]
-			}
-		case CALLBACK_PORT_TYPE_INTERACTIVE:
-			ports := c.getPortsForCallbackAndType(agentMessage.CallbackID, agentMessage.PortType)
-			if len(ports) > 0 {
-				for j := 0; j < len(agentMessage.InteractiveMessages); j++ {
-					//logging.LogInfo("got message from agent", "chan", messages[j].ServerID, "p.messagesFromAgentQueue", len(c.ports[i].messagesFromAgent))
-					ports[0].interactiveMessagesFromAgent <- agentMessage.InteractiveMessages[j]
-				}
-				handleAgentMessagePostResponseInteractiveOutput(&agentMessage.InteractiveMessages)
-			} else {
-				go handleAgentMessagePostResponseInteractiveOutput(&agentMessage.InteractiveMessages)
-			}
+func (c *callbackPortsInUse) ListenForProxyFromAgentMessage(channel <-chan ProxyFromAgentMessageForMythic) {
+	for agentMessage := range channel {
+		c.routeProxyFromAgentMessage(agentMessage)
+	}
+}
 
+func (c *callbackPortsInUse) routeProxyFromAgentMessage(agentMessage ProxyFromAgentMessageForMythic) {
+	switch agentMessage.PortType {
+	case CALLBACK_PORT_TYPE_RPORTFWD:
+		fallthrough
+	case CALLBACK_PORT_TYPE_SOCKS:
+		for j := 0; j < len(agentMessage.Messages); j++ {
+			port := c.getPortForCallbackTypeAndLocalPort(agentMessage.CallbackID, agentMessage.PortType, agentMessage.Messages[j].Port)
+			if port == nil {
+				continue
+			}
+			port.messagesFromAgent <- agentMessage.Messages[j]
+		}
+	case CALLBACK_PORT_TYPE_INTERACTIVE:
+		port := c.getPortForCallbackTypeAndLocalPort(agentMessage.CallbackID, agentMessage.PortType, 0)
+		if port != nil {
+			for j := 0; j < len(agentMessage.InteractiveMessages); j++ {
+				port.interactiveMessagesFromAgent <- agentMessage.InteractiveMessages[j]
+			}
+			handleAgentMessagePostResponseInteractiveOutput(&agentMessage.InteractiveMessages)
+		} else {
+			go handleAgentMessagePostResponseInteractiveOutput(&agentMessage.InteractiveMessages)
 		}
 	}
 }
@@ -804,19 +884,34 @@ func (d callbackProxyData) AddToResponse(response map[string]interface{}) {
 	}
 }
 
+func (c *callbackPortsInUse) sendProxyFromAgentMessage(message ProxyFromAgentMessageForMythic) {
+	c.RLock()
+	channels := c.proxyFromAgentMessageChannels
+	if len(channels) > 0 {
+		channel := channels[proxyFromAgentMessageShard(message.CallbackID, len(channels))]
+		c.RUnlock()
+		channel <- message
+		return
+	}
+	c.RUnlock()
+	// Tests and partially initialized registries may not have listener channels.
+	// Route synchronously in that case instead of blocking on a nil channel.
+	c.routeProxyFromAgentMessage(message)
+}
+
 func (c *callbackPortsInUse) SendDataToCallbackIdPortType(callbackId int, portType CallbackPortType, messages []proxyFromAgentMessage) {
-	c.proxyFromAgentMessageChannel <- ProxyFromAgentMessageForMythic{
+	c.sendProxyFromAgentMessage(ProxyFromAgentMessageForMythic{
 		CallbackID: callbackId,
 		PortType:   portType,
 		Messages:   messages,
-	}
+	})
 }
 func (c *callbackPortsInUse) SendInteractiveDataToCallbackIdPortType(callbackId int, portType CallbackPortType, messages []agentMessagePostResponseInteractive) {
-	c.proxyFromAgentMessageChannel <- ProxyFromAgentMessageForMythic{
+	c.sendProxyFromAgentMessage(ProxyFromAgentMessageForMythic{
 		CallbackID:          callbackId,
 		PortType:            portType,
 		InteractiveMessages: messages,
-	}
+	})
 }
 func (c *callbackPortsInUse) GetOtherCallbackIds(callbackId int) []int {
 	callbackIds := []int{}

mythic-docker/src/rabbitmq/utils_proxy_traffic_test.go

@@ -176,6 +176,122 @@ func TestCallbackPortsInUsePendingDelegateCallbacksIncludeInteractiveTasks(t *te
 	}
 }
 
+func TestCallbackPortsInUseRoutesAgentProxyMessageByLocalPort(t *testing.T) {
+	registry := callbackPortsInUse{}
+	firstPort := newTestCallbackPort(10, CALLBACK_PORT_TYPE_SOCKS, 7001, 1)
+	secondPort := newTestCallbackPort(10, CALLBACK_PORT_TYPE_SOCKS, 7002, 2)
+	addTestCallbackPort(&registry, firstPort)
+	addTestCallbackPort(&registry, secondPort)
+
+	registry.routeProxyFromAgentMessage(ProxyFromAgentMessageForMythic{
+		CallbackID: 10,
+		PortType:   CALLBACK_PORT_TYPE_SOCKS,
+		Messages: []proxyFromAgentMessage{
+			{ServerID: 1, Message: "second", Port: 7002},
+		},
+	})
+
+	select {
+	case message := <-secondPort.messagesFromAgent:
+		if message.Message != "second" {
+			t.Fatalf("expected message to route to second port, got %#v", message)
+		}
+	default:
+		t.Fatal("expected message to route directly to local port 7002")
+	}
+	select {
+	case message := <-firstPort.messagesFromAgent:
+		t.Fatalf("expected first port to remain empty, got %#v", message)
+	default:
+	}
+}
+
+func TestCallbackPortsInUseRoutesAgentProxyMessageWithoutPortToFirstTypeMatch(t *testing.T) {
+	registry := callbackPortsInUse{}
+	firstPort := newTestCallbackPort(10, CALLBACK_PORT_TYPE_RPORTFWD, 7001, 1)
+	secondPort := newTestCallbackPort(10, CALLBACK_PORT_TYPE_RPORTFWD, 7002, 2)
+	addTestCallbackPort(&registry, firstPort)
+	addTestCallbackPort(&registry, secondPort)
+
+	registry.routeProxyFromAgentMessage(ProxyFromAgentMessageForMythic{
+		CallbackID: 10,
+		PortType:   CALLBACK_PORT_TYPE_RPORTFWD,
+		Messages: []proxyFromAgentMessage{
+			{ServerID: 1, Message: "first"},
+		},
+	})
+
+	select {
+	case message := <-firstPort.messagesFromAgent:
+		if message.Message != "first" {
+			t.Fatalf("expected message to route to first matching port, got %#v", message)
+		}
+	default:
+		t.Fatal("expected message without local port to route to first matching type")
+	}
+	select {
+	case message := <-secondPort.messagesFromAgent:
+		t.Fatalf("expected second port to remain empty, got %#v", message)
+	default:
+	}
+}
+
+func TestCallbackPortsInUseRemovesLocalPortRouteIndex(t *testing.T) {
+	registry := callbackPortsInUse{}
+	port := newTestCallbackPort(10, CALLBACK_PORT_TYPE_SOCKS, 7001, 1)
+	addTestCallbackPort(&registry, port)
+	removeTestCallbackPort(&registry, port)
+
+	registry.routeProxyFromAgentMessage(ProxyFromAgentMessageForMythic{
+		CallbackID: 10,
+		PortType:   CALLBACK_PORT_TYPE_SOCKS,
+		Messages: []proxyFromAgentMessage{
+			{ServerID: 1, Message: "removed", Port: 7001},
+		},
+	})
+
+	select {
+	case message := <-port.messagesFromAgent:
+		t.Fatalf("expected removed port to remain empty, got %#v", message)
+	default:
+	}
+}
+
+func TestCallbackPortsInUseReplacesDuplicateLocalPortRouteIndex(t *testing.T) {
+	registry := callbackPortsInUse{}
+	firstPort := newTestCallbackPort(10, CALLBACK_PORT_TYPE_SOCKS, 7001, 1)
+	secondPort := newTestCallbackPort(10, CALLBACK_PORT_TYPE_SOCKS, 7001, 2)
+	addTestCallbackPort(&registry, firstPort)
+	addTestCallbackPort(&registry, secondPort)
+	removeTestCallbackPort(&registry, firstPort)
+
+	registry.routeProxyFromAgentMessage(ProxyFromAgentMessageForMythic{
+		CallbackID: 10,
+		PortType:   CALLBACK_PORT_TYPE_SOCKS,
+		Messages: []proxyFromAgentMessage{
+			{ServerID: 1, Message: "replacement", Port: 7001},
+		},
+	})
+
+	select {
+	case message := <-secondPort.messagesFromAgent:
+		if message.Message != "replacement" {
+			t.Fatalf("expected duplicate local port replacement to receive message, got %#v", message)
+		}
+	default:
+		t.Fatal("expected duplicate local port replacement to receive message")
+	}
+}
+
+func TestProxyFromAgentMessageShardPreservesCallbackOrdering(t *testing.T) {
+	if firstShard := proxyFromAgentMessageShard(10, proxyFromAgentMessageShardCount); firstShard != proxyFromAgentMessageShard(10, proxyFromAgentMessageShardCount) {
+		t.Fatalf("expected same callback to map to the same shard, got %d", firstShard)
+	}
+	if shard := proxyFromAgentMessageShard(10, 0); shard != 0 {
+		t.Fatalf("expected zero shard count to return shard 0, got %d", shard)
+	}
+}
+
 func TestCallbackPortUsagePendingDataTracksQueuedMessages(t *testing.T) {
 	port := newTestCallbackPort(10, CALLBACK_PORT_TYPE_SOCKS, 7001, 1)