diff --git a/Assignment9/src/main.c b/Assignment9/src/main.c
index 564075a..563bc32 100644
--- a/Assignment9/src/main.c
+++ b/Assignment9/src/main.c
@@ -24,6 +24,7 @@ int main(int argc, char *argv[]) {
 		printf("It's a PNG file\n");
 	}
 	png_chunk **chunks = get_png_chunks(png_buffer);
+	free(png_buffer);
 	int size = get_number_of_chunks(chunks);
 	for (int i = 0; i < size; i++) {
 		// Check if header is IDAT or IEND
@@ -39,5 +40,7 @@ int main(int argc, char *argv[]) {
 		}
 		printf("Chunk size is:%d\n", chunks[i]->length);
 	}
-	// write_png_chunks(path, chunks);
+	write_png_chunks(path, chunks);
+	destroy_chunks(chunks);
+	return EXIT_SUCCESS;
 }
\ No newline at end of file
diff --git a/Assignment9/src/png.c b/Assignment9/src/png.c
index 1ad0ea7..11b3e9a 100644
--- a/Assignment9/src/png.c
+++ b/Assignment9/src/png.c
@@ -31,8 +31,8 @@
 #include <arpa/inet.h>
 #include <stdio.h>
 #include <stdlib.h>
-#include <zlib.h>
 #include <string.h>
+#include <stdbool.h>
 #include "png.h"
 
 int get_fd(char *path) {
@@ -77,7 +77,7 @@ int is_png(char *buffer) {
 	return memcmp(buffer, PNG_SIGNATURE, PNG_SIGNATURE_SIZE) == 0;
 }
 
-png_chunk *get_png_chunk(char *buffer, int offset) {
+png_chunk *get_png_chunk(char *buffer, unsigned int offset) {
 	png_chunk *chunk = malloc(sizeof(png_chunk));
 	if (chunk == NULL) {
 		perror("malloc");
@@ -102,7 +102,7 @@ png_chunk **get_png_chunks(char *buffer) {
 		perror("malloc");
 		exit(1);
 	}
-	int offset = PNG_SIGNATURE_SIZE;
+	unsigned int offset = PNG_SIGNATURE_SIZE;
 	int i = 0;
 	while (1) {
 		chunks[i] = get_png_chunk(buffer, offset);
@@ -123,7 +123,10 @@ png_chunk **get_png_chunks(char *buffer) {
 
 int get_number_of_chunks(png_chunk **chunks) {
 	int i = 0;
-	while (chunks[i] != NULL) {
+	while (true) {
+		if (memcmp(chunks[i]->type, "IEND", 4) == 0) {
+			break;
+		}
 		i++;
 	}
 	return i;
@@ -135,6 +138,16 @@ void xor_data(png_chunk *chunk) {
 	}
 }
 
+
+
+void destroy_chunks(png_chunk **chunks) {
+	for (int i = 0; get_number_of_chunks(chunks); i++) {
+		free(chunks[i]->data);
+		free(chunks[i]);
+	}
+	free(chunks);
+}
+
 void write_png_chunks(char *path, png_chunk **chunks) {
 	FILE *fp = fopen(path, "wb");
 	if (fp == NULL) {
diff --git a/Assignment9/src/png.h b/Assignment9/src/png.h
index d21ac5a..f32cd75 100644
--- a/Assignment9/src/png.h
+++ b/Assignment9/src/png.h
@@ -62,7 +62,7 @@ char *load_file(int fd);
 int is_png(char *buffer);
 
 // Get PNG chunk
-png_chunk *get_png_chunk(char *buffer, int offset);
+png_chunk *get_png_chunk(char *buffer, unsigned int offset);
 
 // Get all PNG chunks
 png_chunk **get_png_chunks(char *buffer);
@@ -73,5 +73,8 @@ int get_number_of_chunks(png_chunk **chunks);
 // XOR chunk with key
 void xor_data(png_chunk *chunk);
 
+// Destroy all chunks
+void destroy_chunks(png_chunk **chunks);
+
 // Write PNG chunks to file
 void write_png_chunks(char *path, png_chunk **chunks);
\ No newline at end of file