From be884d09e09b00fbddd31b75dc1f4736d72006a8 Mon Sep 17 00:00:00 2001
From: Con Kolivas <kernel@kolivas.org>
Date: Mon, 15 Feb 2021 15:20:12 +1100
Subject: [PATCH] Deallocate runzip structures after all runzip chunks are
 complete to avoid a race in the case of a failed chunk decompressing.

---
 lrzip_private.h |  8 ++++++++
 runzip.c        | 19 +++++++++++++++++++
 stream.c        | 25 +++++++++++++++++++++----
 3 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/lrzip_private.h b/lrzip_private.h
index 9a456b1..e9613cb 100644
--- a/lrzip_private.h
+++ b/lrzip_private.h
@@ -356,6 +356,12 @@ struct node {
 	struct node *prev;
 };
 
+struct runzip_node {
+	struct stream_info *sinfo;
+	pthread_t *pthreads;
+	struct runzip_node *prev;
+};
+
 struct rzip_state {
 	void *ss;
 	struct node *sslist;
@@ -468,6 +474,8 @@ struct rzip_control {
 	i64 (*match_len)(rzip_control *, struct rzip_state *, i64, i64, i64, i64 *);
 
 	pthread_t *pthreads;
+	struct runzip_node *rulist;
+	struct runzip_node *ruhead;
 };
 
 struct uncomp_thread {
diff --git a/runzip.c b/runzip.c
index f1a3d78..27ddd21 100644
--- a/runzip.c
+++ b/runzip.c
@@ -246,6 +246,21 @@ static i64 unzip_match(rzip_control *control, void *ss, i64 len, uint32 *cksum,
 	return total;
 }
 
+static void clear_rulist(rzip_control *control)
+{
+	while (control->ruhead) {
+		struct runzip_node *node = control->ruhead;
+		struct stream_info *sinfo = node->sinfo;
+
+		dealloc(sinfo->ucthreads);
+		dealloc(node->pthreads);
+		dealloc(sinfo->s);
+		dealloc(sinfo);
+		control->ruhead = node->prev;
+		dealloc(node);
+	}
+}
+
 /* decompress a section of an open file. Call fatal_return(() on error
    return the number of bytes that have been retrieved
  */
@@ -363,6 +378,10 @@ static i64 runzip_chunk(rzip_control *control, int fd_in, i64 expected_size, i64
 	if (unlikely(close_stream_in(control, ss)))
 		fatal("Failed to close stream!\n");
 
+	/* We can now safely delete sinfo and pthread data of all threads
+	 * created. */
+	clear_rulist(control);
+
 	return total;
 }
 
diff --git a/stream.c b/stream.c
index 8892206..443f568 100644
--- a/stream.c
+++ b/stream.c
@@ -1085,6 +1085,8 @@ void *open_stream_in(rzip_control *control, int f, int n, char chunk_bytes)
 	sinfo->s = calloc(sizeof(struct stream), n);
 	if (unlikely(!sinfo->s)) {
 		dealloc(sinfo);
+		dealloc(threads);
+		dealloc(ucthreads);
 		return NULL;
 	}
 
@@ -1206,6 +1208,8 @@ again:
 failed:
 	dealloc(sinfo->s);
 	dealloc(sinfo);
+	dealloc(threads);
+	dealloc(ucthreads);
 	return NULL;
 }
 
@@ -1818,6 +1822,20 @@ int close_stream_out(rzip_control *control, void *ss)
 	return 0;
 }
 
+/* Add to an runzip list to safely deallocate memory after all threads have
+ * returned. */
+static void add_to_rulist(rzip_control *control, struct stream_info *sinfo)
+{
+	struct runzip_node *node = calloc(sizeof(struct runzip_node), 1);
+
+	if (unlikely(!node))
+		failure("Failed to calloc struct node in add_rulist\n");
+	node->sinfo = sinfo;
+	node->pthreads = control->pthreads;
+	node->prev = control->rulist;
+	control->ruhead = node;
+}
+
 /* close down an input stream */
 int close_stream_in(rzip_control *control, void *ss)
 {
@@ -1834,10 +1852,9 @@ int close_stream_in(rzip_control *control, void *ss)
 		dealloc(sinfo->s[i].buf);
 
 	output_thread = 0;
-	dealloc(sinfo->ucthreads);
-	dealloc(control->pthreads);
-	dealloc(sinfo->s);
-	dealloc(sinfo);
+	/* We cannot safely release the sinfo and pthread data here till all
+	 * threads are shut down. */
+	add_to_rulist(control, sinfo);
 
 	return 0;
 }