Improve pivot sessions handling, avoid stucking on implant crash

Before that if we are in shell on pivoted implant and it somewhy dies -
we are stuck forever.

client/command/shell/shell.go

@@ -68,9 +68,12 @@ func runInteractive(ctx *grumble.Context, shellPath string, noPty bool, con *con
 	}
 
 	// Create an RPC tunnel, then start it before binding the shell to the newly created tunnel
-	rpcTunnel, err := con.Rpc.CreateTunnel(context.Background(), &sliverpb.Tunnel{
+	ctxTunnel, cancelTunnel := context.WithCancel(context.Background())
+
+	rpcTunnel, err := con.Rpc.CreateTunnel(ctxTunnel, &sliverpb.Tunnel{
 		SessionID: session.ID,
 	})
+	defer cancelTunnel()
 	if err != nil {
 		con.PrintErrorf("%s\n", err)
 		return
@@ -79,6 +82,7 @@ func runInteractive(ctx *grumble.Context, shellPath string, noPty bool, con *con
 
 	// Start() takes an RPC tunnel and creates a local Reader/Writer tunnel object
 	tunnel := core.GetTunnels().Start(rpcTunnel.TunnelID, rpcTunnel.SessionID)
+	defer tunnel.Close()
 
 	shell, err := con.Rpc.Shell(context.Background(), &sliverpb.ShellReq{
 		Request:   con.ActiveTarget.Request(ctx),

client/console/console.go

@@ -215,6 +215,7 @@ func (con *SliverConsoleClient) EventLoop() {
 			con.PrintEventErrorf("Lost session %s %s - %s (%s) - %s/%s - %v",
 				shortID, session.Name, session.RemoteAddress, session.Hostname, session.OS, session.Arch, currentTime)
 			activeSession := con.ActiveTarget.GetSession()
+			core.GetTunnels().CloseForSession(session.ID)
 			if activeSession != nil && activeSession.ID == session.ID {
 				con.ActiveTarget.Set(nil, nil)
 				con.PrintEventErrorf("Active session disconnected")

client/core/tunnel.go

@@ -15,7 +15,10 @@ func TunnelLoop(rpc rpcpb.SliverRPCClient) error {
 	log.Println("Starting tunnel data loop ...")
 	defer log.Printf("Warning: TunnelLoop exited")
 
-	stream, err := rpc.TunnelData(context.Background())
+	ctx, cancel := context.WithCancel(context.Background())
+	stream, err := rpc.TunnelData(ctx)
+	defer cancel()
+
 	if err != nil {
 		return err
 	}
@@ -26,6 +29,8 @@ func TunnelLoop(rpc rpcpb.SliverRPCClient) error {
 		// log.Printf("Waiting for TunnelData ...")
 		incoming, err := stream.Recv()
 		// log.Printf("Recv stream msg: %v", incoming)
+		// log.Printf("Recv stream err: %s", err)
+
 		if err == io.EOF {
 			log.Printf("EOF Error: Tunnel data stream closed")
 			return nil

client/core/tunnel_io.go

@@ -43,12 +43,11 @@ func (tun *TunnelIO) Write(data []byte) (int, error) {
 
 	// This is necessary to avoid any race conditions on thay byte array
 	dataCopy := make([]byte, len(data))
-	copy(dataCopy, data)
+	n := copy(dataCopy, data)
 
-	log.Printf("Write %d bytes", len(dataCopy))
+	log.Printf("Write %d bytes", n)
 
 	tun.Send <- dataCopy
-	n := len(dataCopy)
 
 	return n, nil
 }

client/core/tunnels.go

@@ -93,11 +93,15 @@ func (t *tunnels) Start(tunnelID uint64, sessionID string) *TunnelIO {
 		for data := range tunnel.Send {
 			log.Printf("Send %d bytes on tunnel %d", len(data), tunnel.ID)
 
-			t.send(&sliverpb.TunnelData{
+			err := t.send(&sliverpb.TunnelData{
 				TunnelID:  tunnel.ID,
 				SessionID: tunnel.SessionID,
 				Data:      data,
 			})
+
+			if err != nil {
+				log.Printf("Error sending, %s", err)
+			}
 		}
 
 		log.Printf("Tunnel Send channel looks closed now. %d", tunnelID)
@@ -122,3 +126,19 @@ func (t *tunnels) Close(tunnelID uint64) {
 		delete((*t.tunnels), tunnelID)
 	}
 }
+
+// CloseForSession - closing all tunnels for specified session id
+func (t *tunnels) CloseForSession(sessionID string) {
+	t.mutex.RLock()
+	defer t.mutex.RUnlock()
+	log.Printf("Closing all tunnels for session %s", sessionID)
+
+	for tunnelID, tunnel := range *t.tunnels {
+		if tunnel.SessionID == sessionID {
+			// Weird way to avoid deadlocks but let it be
+			go func(tunnelID uint64) {
+				GetTunnels().Close(tunnelID)
+			}(tunnelID)
+		}
+	}
+}

server/handlers/pivot.go

@@ -146,8 +146,23 @@ func pivotPeerFailureHandler(implantConn *core.ImplantConnection, data []byte) *
 
 	core.PivotSessions.Range(func(key, value interface{}) bool {
 		pivot := value.(*core.Pivot)
-		if pivot.OriginID == peerFailure.PeerID {
+
+		found := pivot.OriginID == peerFailure.PeerID
+
+		if !found {
+			pivotLog.Warnf("Filed peer not found by OriginID, searching by Peers instead")
+
+			for _, peer := range pivot.Peers {
+				if peer.PeerID == peerFailure.PeerID {
+					pivotLog.Warnf("Found session with needed peer!")
+					found = true
+				}
+			}
+		}
+
+		if found {
 			session := core.Sessions.FromImplantConnection(pivot.ImplantConn)
+
 			if session != nil {
 				core.Sessions.Remove(session.ID)
 			}
@@ -214,11 +229,12 @@ func serverKeyExchange(implantConn *core.ImplantConnection, peerEnvelope *sliver
 		return nil
 	}
 	pivotSession := core.NewPivotSession(peerEnvelope.Peers)
-	pivotLog.Infof("Pivot session %s created with origin %s", pivotSession.ID, peerEnvelope.Peers[0].Name)
-	pivotLog.Debugf("Peers: %v", peerEnvelope.Peers)
 	pivotSession.OriginID = peerEnvelope.Peers[0].PeerID
 	pivotSession.CipherCtx = cryptography.NewCipherContext(sessionKey)
 
+	pivotLog.Infof("Pivot session %s created with origin %s and OriginID: %v", pivotSession.ID, peerEnvelope.Peers[0].Name, pivotSession.OriginID)
+	pivotLog.Infof("Peers: %v", peerEnvelope.Peers)
+
 	pivotRemoteAddr := peersToString(implantConn.RemoteAddress, peerEnvelope)
 
 	pivotSession.ImplantConn = core.NewImplantConnection(core.PivotTransportName, pivotRemoteAddr)