1 files changed, 50 insertions, 8 deletions

diff --git a/src/encryption.c b/src/encryption.c
index 100b50f..32a5ce0 100644
--- a/src/encryption.c
+++ b/src/encryption.c
@@ -272,15 +272,12 @@ int main(void) {
     if((tfd = maketmp(dir)) < 0) // Only because linking temp files is being annoying
         error(1, errno, "Could not open temp file for encryption");
     
-    struct stat sb;
-    if(fstat(fd, &sb) < 0)
-        error(1, errno, "stat() error");
-
     // Read chunks of the file, encrypt them, then write them into the tmp file
     ssize_t bytesread = -1;
-    unsigned char *buf = xcalloc(sb.st_blksize + 1, sizeof(*buf));
-    unsigned char *cbuf = xcalloc((sb.st_blksize + 1) + crypto_secretstream_xchacha20poly1305_ABYTES, sizeof(*buf));
-    while((bytesread = read(fd, buf, sb.st_blksize)) >= 0) {
+    const int CHUNK_SIZE = 4096;
+    unsigned char *buf = xcalloc(CHUNK_SIZE + 1, sizeof(*buf));
+    unsigned char *cbuf = xcalloc((CHUNK_SIZE + 1) + crypto_secretstream_xchacha20poly1305_ABYTES, sizeof(*cbuf));
+    while((bytesread = read(fd, buf, CHUNK_SIZE)) >= 0) {
         crypto_secretstream_xchacha20poly1305_push(&state, cbuf, NULL, buf, bytesread, NULL, 0, (bytesread > 0) ? 0 : crypto_secretstream_xchacha20poly1305_TAG_FINAL);
         if(writewholebuffer(tfd, cbuf, bytesread) < 0)
             error(1, errno, "write() error");
@@ -302,7 +299,52 @@ int main(void) {
 
     close(tfd);
     free(path);
-    free(efname);
+    free(buf);
+    free(cbuf);
+
+
+    // Ok now decryption to make sure I didn't fuck it
+    unsigned char dheader[crypto_secretstream_xchacha20poly1305_HEADERBYTES];
+    crypto_secretstream_xchacha20poly1305_state dstate;
+
+    if((fd = open(efname, O_RDONLY)) < 0)
+        error(1, errno, "Could not open file for decryption");
+    if((tfd = open("lmao.test.dec", O_WRONLY | O_CREAT | O_TRUNC, (S_IRUSR | S_IWUSR))) < 0)
+        error(1, errno, "Could not open file for result of decryption");
+
+    if(read(fd, header, sizeof(header)) < 0)
+        error(1, errno, "Could not read header of encrypted file");
+
+    if(crypto_secretstream_xchacha20poly1305_init_pull(&dstate, dheader, key) != 0)
+        error(1, EINVAL, "Incomplete header");
+
+    bytesread = -1;
+    unsigned char tag = 255;
+    buf = xcalloc(CHUNK_SIZE + 1, sizeof(*buf));
+    cbuf = xcalloc((CHUNK_SIZE + 1) + crypto_secretstream_xchacha20poly1305_ABYTES, sizeof(*cbuf));
+    while((bytesread = read(fd, cbuf, (CHUNK_SIZE + 1) + crypto_secretstream_xchacha20poly1305_ABYTES)) >= 0) {
+        if(crypto_secretstream_xchacha20poly1305_pull(&dstate, buf, NULL, &tag, cbuf, bytesread, NULL, 0) < 0)
+            error(1, errno, "Found a corrupted chunk");
+
+        if(tag == crypto_secretstream_xchacha20poly1305_TAG_FINAL && bytesread != 0)
+            error(1, errno, "End tag before end of file");
+
+        if(bytesread == 0 && tag != crypto_secretstream_xchacha20poly1305_TAG_FINAL)
+            error(1, errno, "End of file before end tag");
+
+        if(writewholebuffer(tfd, buf, bytesread) < 0)
+            error(1, errno, "write() error");
+
+        if(bytesread == 0)
+            break;
+    }
+    if(bytesread < 0)
+        error(1, errno, "read() error");
+
+    close(fd);
+    close(tfd);
+    free(buf);
+    free(cbuf);
 
     //*///