From 985b213dc50b439e35f6169732b1b66486c26c87 Mon Sep 17 00:00:00 2001 From: moloch-- <875022+moloch--@users.noreply.github.com> Date: Wed, 28 Apr 2021 19:27:44 -0500 Subject: [PATCH] Minor refactor --- server/rpc/rpc-tunnel.go | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/server/rpc/rpc-tunnel.go b/server/rpc/rpc-tunnel.go index 6c627dd1c4..e111d40b1a 100644 --- a/server/rpc/rpc-tunnel.go +++ b/server/rpc/rpc-tunnel.go @@ -107,38 +107,40 @@ func (s *Server) TunnelData(stream rpcpb.SliverRPC_TunnelDataServer) error { go func() { - send_cache, _ := toImplantCache[tunnel.ID] + sendCache, _ := toImplantCache[tunnel.ID] for tunnelData := range tunnel.FromImplant { - tunnelLog.Debugf("Tunnel %d: From implant %d byte(s), seq: %d ack: %d", tunnel.ID, len(tunnelData.Data), tunnelData.Sequence, tunnelData.Ack) - //Remove tunnel data from send cache if Resend is not set + tunnelLog.Debugf("Tunnel %d: From implant %d byte(s), seq: %d ack: %d", + tunnel.ID, len(tunnelData.Data), tunnelData.Sequence, tunnelData.Ack) + + // Remove tunnel data from send cache if Resend is not set if !tunnelData.Resend { - i := tunnelData.Ack - 1 - for send_msg, ok := send_cache[i]; ok; send_msg, ok = send_cache[i] { - tunnelLog.Debugf("Tunnel %d: Removing ack: %d from send cache", tunnel.ID, send_msg.Sequence) - delete(send_cache, i) - i = i - 1 + index := tunnelData.Ack - 1 + for sendMsg, ok := sendCache[index]; ok; sendMsg, ok = sendCache[index] { + tunnelLog.Debugf("Tunnel %d: Removing ack: %d from send cache", tunnel.ID, sendMsg.Sequence) + delete(sendCache, index) + index = index - 1 } - recv_cache, ok := fromImplantCache[tunnel.ID] + recvCache, ok := fromImplantCache[tunnel.ID] if ok { - recv_cache[tunnelData.Sequence] = tunnelData + recvCache[tunnelData.Sequence] = tunnelData } - for recv, ok := recv_cache[tunnel.FromImplantSequence]; ok; recv, ok = recv_cache[tunnel.FromImplantSequence] { + for recv, ok := recvCache[tunnel.FromImplantSequence]; ok; recv, ok = recvCache[tunnel.FromImplantSequence] { tunnel.Client.Send(&sliverpb.TunnelData{ TunnelID: tunnel.ID, SessionID: tunnel.SessionID, Data: recv.Data, Closed: false, }) - delete(recv_cache, tunnel.FromImplantSequence) + delete(recvCache, tunnel.FromImplantSequence) tunnel.FromImplantSequence++ } } else { - origtunnelData, ok := send_cache[tunnelData.Ack] + origtunnelData, ok := sendCache[tunnelData.Ack] if ok { tunnelLog.Debugf("Tunnel %d: Resending cached msg: %d", tunnel.ID, tunnelData.Ack) session := core.Sessions.Get(tunnel.SessionID) @@ -178,7 +180,7 @@ func (s *Server) TunnelData(stream rpcpb.SliverRPC_TunnelDataServer) error { Data: data, Closed: false, } - //Add tunnel data to cache + // Add tunnel data to cache send_cache[tunnelData.Sequence] = &tunnelData data, _ := proto.Marshal(&tunnelData)