From afbafbab800add69686e8794b30c7f767c17069a Mon Sep 17 00:00:00 2001
From: alex <alex@infiniteadaptability.org>
Date: Thu, 2 Dec 2021 17:02:35 -0800
Subject: [PATCH] ...

---
 inc/tree.h                |  1 +
 src/bencode/encode.c      | 14 ++++++--------
 src/tree.c                |  7 +++++++
 test/unit/Makefile.am     |  1 +
 test/unit/bencode.tests.c |  6 ++++--
 5 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/inc/tree.h b/inc/tree.h
index 10e3dd4..b2893bb 100644
--- a/inc/tree.h
+++ b/inc/tree.h
@@ -3,6 +3,7 @@
 
 #include<stdlib.h>
 
+#include<bencode.h>
 #include<file.h>
 
 struct tree_entry {
diff --git a/src/bencode/encode.c b/src/bencode/encode.c
index b21c6b8..01fc8c8 100644
--- a/src/bencode/encode.c
+++ b/src/bencode/encode.c
@@ -55,28 +55,26 @@ ssize_t bencode_list_start(uint8_t *output, size_t output_size) {
 
 ssize_t bencode_string(const uint8_t *str, size_t len, uint8_t *output, size_t output_size) {
 	size_t i;
-	int ret;
 	
 	if(NULL==str) { return -1; }
-	if(len<=0) { return -1; }
 	if(NULL==output) { return -1; }
-	if(len>output_size) { return -1; }
 
 	i = size_int(len);
-	i++;
+	i += 2; // account for ':' and '\0'
 	i += len;
 
-	if((ret = snprintf((char*)output,output_size,"%lu:%s",len,(const char*)str))<0) { return -1; }
-
-	if(ret!=i) { return -1; }
+	if(i>output_size) { return -1; }
 
-	return i;
+	return snprintf((char*)output,i,"%lu:%s",len,(const char*)str);
 }
 
 
 static size_t size_int(int i) {
 	i = abs(i);
 	size_t j = 0;
+	
+	if(0==i) { return 1; }
+
 	while(i>0) {
 		i /= 10;
 		j++;
diff --git a/src/tree.c b/src/tree.c
index 23bcaac..1b3207a 100644
--- a/src/tree.c
+++ b/src/tree.c
@@ -112,9 +112,16 @@ static int tree_add_file(struct tree *root, struct file *to_add) {
 
 ssize_t tree_bencode(struct tree *tree, uint8_t *buf, size_t buf_len) {
 	struct tree_entry *p;
+	ssize_t i;
 
 	p = tree->entries;
 	while(p!=NULL) {
+		i = bencode_string((const uint8_t*)p->name,strlen(p->name),buf,buf_len);
+		if(i<0) { return -1; }
+
+		buf += i;
+		buf_len -= i;
+
 		if(p->children!=NULL) {
 			// write directory
 			// recurse into subdirectory
diff --git a/test/unit/Makefile.am b/test/unit/Makefile.am
index 17e239c..a1ae89d 100644
--- a/test/unit/Makefile.am
+++ b/test/unit/Makefile.am
@@ -76,6 +76,7 @@ torrent_tests_SOURCES = \
 tree_tests_SOURCES = \
 	$(common_SOURCES) \
 	tree.tests.c \
+	$(top_srcdir)/src/bencode/encode.c \
 	$(top_srcdir)/src/block.c \
 	$(top_srcdir)/src/file.c \
 	$(top_srcdir)/src/hash.c \
diff --git a/test/unit/bencode.tests.c b/test/unit/bencode.tests.c
index 48e8efb..25a4320 100644
--- a/test/unit/bencode.tests.c
+++ b/test/unit/bencode.tests.c
@@ -142,10 +142,12 @@ static void bencode_string_basic_test() {
 
 	uint8_t str1[] = "testlkajslfkdjasdfl test string";
 	assert(bencode_string(NULL,31,buf,1024)==-1);
-	assert(bencode_string(str1,0,buf,1024)==-1);
 	assert(bencode_string(str1,31,NULL,1024)==-1);
 	assert(bencode_string(str1,31,buf,0)==-1);
-	assert(bencode_string(str1,31,buf,1024)==34);
+	
+	assert(bencode_string(str1,0,buf,1024)==2);
+	assert(memcmp("0:",buf,2)==0);
 
+	assert(bencode_string(str1,31,buf,1024)==34);
 	assert(memcmp("31:testlkajslfkdjasdfl test string",buf,34)==0);
 }
-- 
2.30.2