1 files changed, 28 insertions, 13 deletions

diff --git a/src/encryption.c b/src/encryption.c
index 0cd032f..052b9aa 100644
--- a/src/encryption.c
+++ b/src/encryption.c
@@ -94,28 +94,38 @@ int encrypttotmp(const char *toencrypt) {
 }
 
 int genpassword(char **str, unsigned int words) {
+    // Early returns
+    if(words < 1)
+        return 0;
     #if defined ___VXGG___ALWAYS_CHECK_LIBSODIUM___ && ___VXGG___ALWAYS_CHECK_LIBSODIUM___ > 0
         checksodium();
     #endif
 
-    unsigned int i = 0;
-    char *lstr = NULL;
-    
-    if(words < 1)
-        return 0;
-    
-    asprintf(&lstr, "%s", PASSWORD_WORDS[randombytes_uniform(PASSWORD_WORDS_LEN)]);
-    for(; i < words; i++) {
-        asprintf(&lstr, "%s %s", lstr, PASSWORD_WORDS[randombytes_uniform(PASSWORD_WORDS_LEN)]);
+    // Bootstrap the first word
+    char *lstr = NULL, *tmp = NULL;
+    if(asprintf(&lstr, "%s", PASSWORD_WORDS[randombytes_uniform(PASSWORD_WORDS_LEN)]) < 0)
+        return -1;
+
+    // Concat the rest of the words into the password (without leaking memory)
+    int ret;
+    for(unsigned int i = 1; i < words; i++) {    
+        ret = asprintf(&tmp, "%s %s", lstr, PASSWORD_WORDS[randombytes_uniform(PASSWORD_WORDS_LEN)]);
+        free(lstr);
+        if(ret < 0)
+            return -1;
+
+        lstr = tmp;
     }
 
     *str = lstr;
+    return words;
 
-    return 0;
-
-    // TODO: I feel like this is / should be leaking memory like a mofo. Figure out if it is or not (look at malloc_stats())
+    // This function was exploding because of some weird conflict with using my buggy implementation of asprintf instead of the 
+    // _GNU_SOURCE version. Don't know why it wasn't using the _GNU_SOURCE version, as I had a define macro put in place to 
+    // prevent it from being compiled if _GNU_SOURCE was defined, but whatever
 }
 
+
 #define TESTING
 #ifdef TESTING
 
@@ -143,9 +153,14 @@ int main(void) {
         error(1, errno, "close broke");
     //*///
 
+    //*// Example code for getting a password using genpassword
+    checksodium();
+
     char *password = NULL;
     genpassword(&password, 20);
-    printf("%s\n", password);
+    printf("%s\n", (password != NULL) ? password : "Couldn't get a password");
+    free(password);
+    //*///
 
     return 0;
 }