Fix: SMB pivot session memory leak causing server OOM kill

AdaptixServer/core/server/ts_agent.go

@@ -194,6 +194,12 @@ func (ts *Teamserver) TsAgentProcessData(agentId string, bodyData []byte) error
 		// -----------------
 	}
 
+	if !agentData.Async {
+		agent.UpdateData(func(d *adaptix.AgentData) {
+			d.LastTick = int(time.Now().Unix())
+		})
+	}
+
 	if len(bodyData) > 4 {
 		return agent.ProcessData(bodyData)
 	}

AdaptixServer/core/server/ts_tasks.go

@@ -120,6 +120,8 @@ func (ts *Teamserver) extractTunnelData(agent *Agent, availableSize int, startSi
 
 func (ts *Teamserver) extractPivotTasks(agent *Agent, availableSize int, startSize int) (tasks []adaptix.TaskData, usedSize int) {
 	usedSize = startSize
+	const minPivotSize = 0x10000 // 64KB minimum for each pivot child
+
 	for i := uint(0); i < agent.PivotChilds.Len(); i++ {
 		value, ok := agent.PivotChilds.Get(i)
 		if !ok {
@@ -130,10 +132,16 @@ func (ts *Teamserver) extractPivotTasks(agent *Agent, availableSize int, startSi
 		if lostSize <= 0 {
 			break
 		}
+		if lostSize < minPivotSize && availableSize > minPivotSize {
+			lostSize = minPivotSize
+		}
 		data, err := ts.TsAgentGetHostedAll(pivotData.ChildAgentId, lostSize)
 		if err != nil {
 			continue
 		}
+		if len(data) == 0 {
+			continue
+		}
 		pivotTaskData, err := agent.PivotPackData(pivotData.PivotId, data)
 		if err != nil {
 			continue

AdaptixServer/core/utils/safe/queue.go

@@ -16,10 +16,15 @@ type Queue struct {
 	capacity int
 }
 
+const maxQueueCapacity = 0x10000 // 65536 items max
+
 func NewSafeQueue(capacity int) *Queue {
 	if capacity <= 0 {
 		capacity = 0x100
 	}
+	if capacity > maxQueueCapacity {
+		capacity = maxQueueCapacity
+	}
 	return &Queue{
 		buffer:   make([]interface{}, capacity),
 		capacity: capacity,
@@ -43,7 +48,13 @@ func (q *Queue) Push(value interface{}) {
 	defer q.lock.Unlock()
 
 	if q.size == q.capacity {
-		q.resize()
+		if q.capacity >= maxQueueCapacity {
+			q.buffer[q.head] = nil
+			q.head = (q.head + 1) % q.capacity
+			q.size--
+		} else {
+			q.resize()
+		}
 	}
 	q.buffer[q.tail] = value
 	q.tail = (q.tail + 1) % q.capacity
@@ -55,7 +66,13 @@ func (q *Queue) PushFront(value interface{}) {
 	defer q.lock.Unlock()
 
 	if q.size == q.capacity {
-		q.resize()
+		if q.capacity >= maxQueueCapacity {
+			q.buffer[(q.tail-1+q.capacity)%q.capacity] = nil
+			q.tail = (q.tail - 1 + q.capacity) % q.capacity
+			q.size--
+		} else {
+			q.resize()
+		}
 	}
 
 	q.head = (q.head - 1 + q.capacity) % q.capacity
@@ -99,6 +116,9 @@ func (q *Queue) Clear() {
 
 func (q *Queue) resize() {
 	newCap := q.capacity * 2
+	if newCap > maxQueueCapacity {
+		newCap = maxQueueCapacity
+	}
 	newBuf := make([]interface{}, newCap)
 
 	for i := 0; i < q.size; i++ {

AdaptixServer/extenders/beacon_listener_smb/pl_main.go

@@ -172,6 +172,10 @@ func (l *Listener) InternalHandler(data []byte) (string, error) {
 	agentInfo := make([]byte, len(data))
 	rc4crypt.XORKeyStream(agentInfo, data)
 
+	if len(agentInfo) < 8 {
+		return "", fmt.Errorf("invalid agent info length")
+	}
+
 	agentType := fmt.Sprintf("%08x", uint(binary.BigEndian.Uint32(agentInfo[:4])))
 	agentInfo = agentInfo[4:]
 	agentId = fmt.Sprintf("%08x", uint(binary.BigEndian.Uint32(agentInfo[:4])))

AdaptixServer/extenders/beacon_listener_tcp/pl_main.go

@@ -175,6 +175,10 @@ func (l *Listener) InternalHandler(data []byte) (string, error) {
 	agentInfo := make([]byte, len(data))
 	rc4crypt.XORKeyStream(agentInfo, data)
 
+	if len(agentInfo) < 8 {
+		return "", fmt.Errorf("invalid agent info length")
+	}
+
 	agentType := fmt.Sprintf("%08x", uint(binary.BigEndian.Uint32(agentInfo[:4])))
 	agentInfo = agentInfo[4:]
 	agentId = fmt.Sprintf("%08x", uint(binary.BigEndian.Uint32(agentInfo[:4])))