diff --git a/AdaptixServer/core/server/ts_agent.go b/AdaptixServer/core/server/ts_agent.go index 215cd51cd..5ece19578 100644 --- a/AdaptixServer/core/server/ts_agent.go +++ b/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) } diff --git a/AdaptixServer/core/server/ts_tasks.go b/AdaptixServer/core/server/ts_tasks.go index 4943095d5..9c90358ca 100644 --- a/AdaptixServer/core/server/ts_tasks.go +++ b/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 diff --git a/AdaptixServer/core/utils/safe/queue.go b/AdaptixServer/core/utils/safe/queue.go index 2cf7d7576..fa8417125 100644 --- a/AdaptixServer/core/utils/safe/queue.go +++ b/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++ { diff --git a/AdaptixServer/extenders/beacon_listener_smb/pl_main.go b/AdaptixServer/extenders/beacon_listener_smb/pl_main.go index 922a423b9..718b3081b 100644 --- a/AdaptixServer/extenders/beacon_listener_smb/pl_main.go +++ b/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]))) diff --git a/AdaptixServer/extenders/beacon_listener_tcp/pl_main.go b/AdaptixServer/extenders/beacon_listener_tcp/pl_main.go index 43b2f5b48..bf4db8788 100644 --- a/AdaptixServer/extenders/beacon_listener_tcp/pl_main.go +++ b/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])))