From 9b8539e1e0857e358c82c8d48c30107ab060af46 Mon Sep 17 00:00:00 2001 From: Lorenzo Torres Date: Tue, 20 Aug 2024 12:57:30 +0200 Subject: [PATCH] feat!: implemented tree matching for command parsing --- imap.c | 81 ++++++++++++++++++++++++++++++++------------------- imap.h | 6 ++++ imap.routines | 35 ++++++++++++++++++---- 3 files changed, 87 insertions(+), 35 deletions(-) diff --git a/imap.c b/imap.c index 512c08e..6eab093 100644 --- a/imap.c +++ b/imap.c @@ -38,30 +38,46 @@ #include #include -struct { - char *key; - uint8_t id; -} cmd_map[] = { - /* https://datatracker.ietf.org/doc/html/rfc3501#section-6.1.1 */ - { "capability", 0x00 }, - /* https://datatracker.ietf.org/doc/html/rfc3501#section-6.1.2 */ - { "noop", 0x01 }, - /* https://datatracker.ietf.org/doc/html/rfc3501#section-6.1.3 */ - { "logout", 0x02 }, - /* https://datatracker.ietf.org/doc/html/rfc3501#section-6.2.1 */ - { "starttls", 0x03 }, - /* https://datatracker.ietf.org/doc/html/rfc3501#section-6.2.2 */ - { "authenticate", 0x04 }, - /* https://datatracker.ietf.org/doc/html/rfc3501#section-6.2.3 */ - { "login", 0x05 }, - /* Invalid command */ - { NULL, 0xff } -}; +static char buf[CMD_MAX_SIZE]; +static trie_node *trie; -#define CMD_MAP_LAST 0x05 +void imap_trie_encode(char *str, uint8_t cmd) +{ + trie_node *node; + if (trie == NULL) { + trie = (trie_node *) malloc(sizeof(trie_node)); + memset(trie, 0x0, sizeof(trie_node)); + } + + node = trie; + do { + node->children[*(str) - 'a'] = (trie_node *) malloc(sizeof(trie_node)); + node = node->children[*(str) - 'a']; + memset(node, 0x0, sizeof(trie_node)); + node->id = 0xff; + str++; + } while (*str != '\0'); + + node->id = cmd; + printf("%d\n", node->id); +} + +void imap_populate_trie(void) +{ + imap_trie_encode("capability", 0x0); + imap_trie_encode("noop", 0x1); + imap_trie_encode("logout", 0x2); + imap_trie_encode("starttls", 0x3); + imap_trie_encode("authenticate", 0x4); + imap_trie_encode("login", 0x5); +} + +#define CMD_MAP_LAST 0x5 uint8_t imap_init(uint8_t daemon, imap_t *instance) { + imap_populate_trie(); + imap_t imap; imap.ssl_ctx = NULL; imap.ssl = 0; @@ -108,7 +124,7 @@ uint8_t imap_init(uint8_t daemon, imap_t *instance) if (TLS_ENABLED) { imap_create_ssl_ctx(&imap); - imap_starttls(&imap, imap.clients); + imap_starttls(&imap, NULL); } memcpy(instance, &imap, sizeof(imap)); @@ -212,7 +228,6 @@ void imap_start(imap_t *instance) { int activity, max_fd, connection; size_t bytes_read; - char buf[CMD_MAX_SIZE]; /* List of all the file descriptors (sockets) being used. */ fd_set fds; instance->clients = NULL; @@ -316,11 +331,18 @@ uint8_t imap_match_cmd(char *cmd, size_t len) { strnlower(cmd, len); - for (int i=0; cmd_map[i].key != NULL; i++) { - if (strncmp(cmd_map[i].key, cmd, len) == 0) { - return cmd_map[i].id; + trie_node *node = trie; + do { + node = node->children[*cmd - 'a']; + if (node->id == 0xff) { + cmd++; + continue; + } else { + return node->id; } - } + + cmd++; + } while (*cmd != '\0'); return 0xff; } @@ -423,12 +445,12 @@ void imap_create_ssl_ctx(imap_t *imap) } } -int imap_read(client_list *node, char *buf, size_t len, uint8_t ssl) +int imap_read(client_list *node, char *buffer, size_t len, uint8_t ssl) { if (ssl) { - return SSL_read(node->ssl, buf, len); + return SSL_read(node->ssl, buffer, len); } else { - return read(node->socket, buf, len); + return read(node->socket, buffer, len); } } @@ -436,7 +458,6 @@ void imap_write(client_list *node, uint8_t ssl, char *fmt, ...) { va_list(args); va_start(args, fmt); - char buf[CMD_MAX_SIZE]; vsprintf(buf, fmt, args); if (!ssl) { diff --git a/imap.h b/imap.h index c973be3..751ad5e 100644 --- a/imap.h +++ b/imap.h @@ -51,6 +51,10 @@ typedef struct _client_list { struct _client_list *prev; } client_list; +typedef struct _trie_node { + struct _trie_node *children[26]; + uint8_t id; +} trie_node; typedef struct imap { int32_t socket; @@ -88,5 +92,7 @@ int imap_read(client_list *node, char *buf, size_t len, uint8_t ssl); void imap_write(client_list *node, uint8_t ssl, char *fmt, ...); void imap_flush(client_list *node, uint8_t ssl); uint8_t imap_cmd_exec(imap_cmd cmd, client_list *node, uint8_t ssl, uint8_t state); +void imap_trie_populate(void); +void imap_trie_encode(char *str, uint8_t cmd); #endif /* ifndef IMAP_H */ diff --git a/imap.routines b/imap.routines index 3f90f12..f334e9b 100644 --- a/imap.routines +++ b/imap.routines @@ -4,14 +4,18 @@ name: { \ return imap_routine_##name(cmd, node, ssl, state); \ } - +#define IMAP_ROUTINE_BAD_TAG \ + imap_write(node, ssl, "%s BAD\n", cmd.tag); +#define IMAP_ROUTINE_BAD \ + imap_write(node, ssl, "* BAD\n"); +#define IMAP_CHECK_ARGS(x) \ + if (cmd.p_count != x) { \ + IMAP_ROUTINE_BAD_TAG \ + return IMAP_FAIL; \ + } #define IMAP_ROUTINE_END imap_flush(node, ssl); #define IMAP_ROUTINE_OK(routine) \ imap_write(node, ssl, "%s OK " #routine " completed\n", cmd.tag); -#define IMAP_ROUTINE_BAD_TAG \ - imap_write(node, ssl, "%s BAD", cmd.tag); -#define IMAP_ROUTINE_BAD \ - imap_write(node, ssl, "* BAD"); #define IMAP_STRING(fmt, ...) \ imap_write(node, ssl, fmt, ##__VA_ARGS__); #define IMAP_NLINE imap_write(node, ssl, "\n"); @@ -68,6 +72,27 @@ static inline uint8_t imap_routine_starttls(imap_cmd cmd, client_list *node, uin static inline uint8_t imap_routine_auth(imap_cmd cmd, client_list *node, uint8_t ssl, uint8_t state) { + IMAP_CHECK_STATE(NO_AUTH) + IMAP_CHECK_ARGS(1) + + int bytes; + + if (strcmp(cmd.params[0], "PLAIN") == 0) { + IMAP_STRING("+\n"); + if ((bytes = imap_read(node, buf, CMD_MAX_SIZE, ssl)) < 0) { + perror("recv"); + syslog(LOG_ERR, "Failed to receive data."); + } else if (bytes == 0) { + return IMAP_LOGOUT; + } else { + buf[bytes] = '\0'; + printf("%s\n", buf); + } + } else { + IMAP_ROUTINE_BAD_TAG + } + + IMAP_ROUTINE_END return IMAP_SUCCESS; }