4 files changed, 139 insertions, 16 deletions

diff --git a/src/benc.c b/src/benc.c
index 4f61ff7..eef8469 100644
--- a/src/benc.c
+++ b/src/benc.c
@@ -175,6 +175,102 @@ int64_t benc_decode(const char** cursor, const char* end, int* depth, struct ben
 	return cursor_out;
 }
 
+int bcur_fill(struct bcursor* cursor) {
+	if(cursor->source == cursor->source_end) {
+		return EOF;
+	}
+
+	int read = benc_decode(&cursor->source, cursor->source_end, &cursor->source_depth, cursor->base, cursor->base_len);
+	cursor->end = cursor->base + read;
+
+	return 0;
+}
+
+int bcur_open(struct bcursor* cursor, const char* source, const char* source_end, struct benc_node* buffer, size_t buffer_len) {
+	cursor->source = source;
+	(*((const char**)&cursor->source_end)) = source_end;
+	cursor->source_depth = 0;
+
+	*((struct benc_node**)&cursor->base) = buffer;
+	*((size_t*)&cursor->base_len) = buffer_len;
+	cursor->readhead = buffer;
+
+	return bcur_fill(cursor);
+}
+
+int bcur_next(struct bcursor* cursor, uint32_t steps) {
+	cursor->readhead+=steps;
+
+	// Check if we need to read more
+	if(cursor->readhead >= cursor->end) {
+		return bcur_fill(cursor);
+	}
+	return 0;
+}
+
+int bcur_next_sibling(struct bcursor* cursor) {
+	assert(cursor->readhead->type == BNT_LIST || cursor->readhead->type == BNT_DICT);
+
+	if(cursor->readhead->type == BNT_LIST || cursor->readhead->type == BNT_DICT) {
+		uint32_t tdepth = cursor->readhead->depth;
+		cursor->readhead++;
+		while(cursor->readhead->type != BNT_END || cursor->readhead->depth != tdepth) {
+			if(bcur_next(cursor, 1) != 0) {
+				fatal("Invalid dict/list");
+			}
+		}
+	}
+
+	return bcur_next(cursor, 1);
+}
+
+// All the arrays should be equal length
+ssize_t bcur_find_key(struct bcursor* cursor, const enum benc_nodetype* keyTypes, const char** keyValues, const size_t* keyLengths, const size_t keys) {
+	// @ROBUSTNESS: Check if keytypes are anything but strings and ints because that is not supported
+	int rc;
+
+	while(true) {
+		// Skip lists
+		if(cursor->readhead->type == BNT_LIST || cursor->readhead->type == BNT_DICT) {
+			rc = bcur_next_sibling(cursor);
+			if(rc != 0) {
+				fatal("Invalid dict");
+			}
+		} else {
+			// Check if the key one of the ones we are looking for
+			for(size_t i = 0; i < keys; i++) {
+				if(cursor->readhead->type == keyTypes[i] && memcmp(cursor->readhead->loc, keyValues[i], MIN(cursor->readhead->size, keyLengths[i])) == 0) {
+					return i;
+				}
+			}
+		}
+
+		// Skip the key part
+		rc = bcur_next(cursor, 1);
+		if(rc != 0) {
+			fatal("Invalid dict");
+		}
+
+		if(cursor->readhead->type == BNT_LIST || cursor->readhead->type == BNT_DICT) {
+			// Skip a multitoken element
+			rc = bcur_next_sibling(cursor);
+			if(rc != 0) {
+				fatal("Invalid dict");
+			}
+		} else {
+			// Skip a single token element
+			rc = bcur_next(cursor, 1);
+			if(rc != 0) {
+				fatal("Invalid dict");
+			}
+		}
+
+		if(cursor->readhead->type == BNT_END) {
+			return -1;
+		}
+	}
+}
+
 void skip_sibling(const struct benc_node** cursor, const struct benc_node* end) {
 	assert((*cursor)->type == BNT_LIST || (*cursor)->type == BNT_DICT);
 
@@ -204,7 +300,7 @@ ssize_t skip_to_key(const struct benc_node** cursor, const struct benc_node* end
 
 	while(true) {
 		if(*cursor > end-1) {
-			fatal("Invalid dict");
+			return -1;
 		}
 
 		// Skip lists
diff --git a/src/benc.h b/src/benc.h
index a5734ea..be06d5e 100644
--- a/src/benc.h
+++ b/src/benc.h
@@ -27,5 +27,23 @@ bool readint(const char** loc, int64_t* val);
 void benc_print(const struct benc_node* stream, size_t stream_len, int* depth);
 int64_t benc_decode(const char** cursor, const char* end, int* depth, struct benc_node* stream, size_t stream_len);
 
+struct bcursor {
+	const char* source;
+	const char* const source_end; 
+	int source_depth;
+
+	struct benc_node* const base;
+	const size_t base_len;
+
+	struct benc_node* readhead;
+	struct benc_node* end;
+};
+
+int bcur_open(struct bcursor* cursor, const char* source, const char* source_end, struct benc_node* buffer, size_t buffer_len);
+int bcur_next(struct bcursor* cursor, uint32_t steps);
+int bcur_next_sibling(struct bcursor* cursor);
+ssize_t bcur_find_key(struct bcursor* cursor, const enum benc_nodetype* keyTypes, const char** keyValues, const size_t* keyLengths, const size_t keys);
+
 void skip_sibling(const struct benc_node** cursor, const struct benc_node* end);
 ssize_t skip_to_key(const struct benc_node** cursor, const struct benc_node* end, const enum benc_nodetype* keyTypes, const char** keyValues, const size_t* keyLengths, const size_t keys);
+
diff --git a/src/log.h b/src/log.h
index 1717011..12135af 100644
--- a/src/log.h
+++ b/src/log.h
@@ -13,8 +13,9 @@
 	printf(format "\n", ## __VA_ARGS__); \
 	fflush(stderr)
 
-#define fatal(format, ...) \
-	printf(format "\n", ## __VA_ARGS__); \
-	fflush(stderr); \
-	abort()
+#define fatal(format, ...) do{\
+		printf(format "\n", ## __VA_ARGS__); \
+		fflush(stderr); \
+		abort(); \
+	} while(0)
 
diff --git a/src/main.c b/src/main.c
index 0c49797..5de849a 100644
--- a/src/main.c
+++ b/src/main.c
@@ -337,28 +337,36 @@ int main(int argc, char** argv) {
 
 		struct benc_node* stream_cursor = stream;
 		if(stream_cursor->type != BNT_DICT) {
-			err("Response is not a dict");
-			exit(EXIT_FAILURE);
+			fatal("First value is not a dict");
 		}
 		stream_cursor++;
 
 		if(skip_to_key((const struct benc_node**)&stream_cursor, stream+len, (const enum benc_nodetype[]){BNT_STRING}, (const char*[]){"t"}, (const size_t[]){1}, 1) == -1) {
-			err("No t key in packet");
-			exit(EXIT_FAILURE);
+			fatal("No t key in packet");
 		}
 		stream_cursor++;
 
-		char char_buffer = stream_cursor->loc[stream_cursor->size];
-		((char*)stream_cursor->loc)[stream_cursor->size] = '\0';
-		uint32_t transaction = strtol(stream_cursor->loc, NULL, 10);
-		((char*)stream_cursor->loc)[stream_cursor->size] = char_buffer;
+		uint32_t transaction;
+		{
+			// Temporary null terminate the string to parse the number without a copy
+			char char_buffer = stream_cursor->loc[stream_cursor->size];
+			((char*)stream_cursor->loc)[stream_cursor->size] = '\0';
+			char* end;
+			transaction = strtol(stream_cursor->loc, &end, 10);
+			((char*)stream_cursor->loc)[stream_cursor->size] = char_buffer;
+
+			if(end != stream_cursor->loc+stream_cursor->size) {
+				dbg("DISCARD: Transaction id is not a number %.*s", stream_cursor->size, stream->loc);
+				continue;
+			}
+		}
 
 		uint16_t reqId;
 		if(!find_req(transaction, &reqId)) {
-			err("No request with transaction ID %d", transaction);
-			exit(EXIT_FAILURE);
+			dbg("DISCARD: unknown transaction id %d", transaction);
+			continue;
 		}
-		dbg("Transaction id matches requst %d", reqId);
+		dbg("Transaction id matches request %d", reqId);
 
 		if(sockaddr_cmp(&requestdata[reqId].addr, (struct sockaddr*)&remote) != 0) {
 			err("Unexpected IP for valid transaction");