1 files changed, 32 insertions, 17 deletions

diff --git a/multipart-replace.c b/multipart-replace.c
index a10ba4d..9a5e268 100644
--- a/multipart-replace.c
+++ b/multipart-replace.c
@@ -4,13 +4,14 @@
 #include <error.h>
 #include <stdio.h>
 #include <stdlib.h>
+#include <stdlib.h>
 #include <string.h>
 #include <unistd.h>
 
 #include "multipart-replace.h"
 #include "util.h"
 
-#define error(stat, errnum, ...) do { error(0, errnum, __VA_ARGS__); return; } while(0)
+#define error(stat, errnum, ...) do { error(0, errnum, __VA_ARGS__); goto end; } while(0)
 
 static
 char *boundary_line(const char *old) {
@@ -40,9 +41,9 @@ ssize_t safe_atoi(char *str) {
 	return atoi(str);
 }
 
-void multipart_replace_reader(struct multipart_replace_stream *s, int fd, const char *boundary) {
+void multipart_replace_reader(struct multipart_replace_stream *s, int fd, const char *_boundary) {
 	FILE *stream = fdopen(fd, "r");
-	boundary = boundary_line(boundary);
+	char *boundary = boundary_line(_boundary);
 
 	char *line_buf = NULL;
 	size_t line_cap = 0;
@@ -50,7 +51,7 @@ void multipart_replace_reader(struct multipart_replace_stream *s, int fd, const
 	ssize_t content_length = -1;
 
 	bool first = true;
-	while (1) {
+	while (running) {
 		content_length = -1;
 		/* scan for the first non-empty line */
 		do {
@@ -103,27 +104,21 @@ void multipart_replace_reader(struct multipart_replace_stream *s, int fd, const
 		s->back = tmp;
 		s->framecount++;
 		pthread_rwlock_unlock(&s->frontlock);
-
-		//debug("got frame on %d", fd);
 	}
+ end:
+	free(boundary);
+	fclose(stream);
 }
 
-void multipart_replace_writer(struct multipart_replace_stream *s, int fd, const char *boundary) {
+void multipart_replace_writer(struct multipart_replace_stream *s, int fd, const char *_boundary) {
 	FILE *stream = fdopen(fd, "w");
 	struct frame myframe = { 0 };
 	long lastframe = 0;
-	boundary = boundary_line(boundary);
+	char *boundary = boundary_line(_boundary);
 	size_t boundary_len = strlen(boundary);
 
-	while(1) {
-		/* poll until there's a new frame */
-		pthread_rwlock_rdlock(&s->frontlock);
-		while (s->framecount == lastframe) {
-			pthread_rwlock_unlock(&s->frontlock);
-			usleep(30000); /* a bit over 30 FPS */
-			pthread_rwlock_rdlock(&s->frontlock);
-		}
-
+	pthread_rwlock_rdlock(&s->frontlock);
+	while (running) {
 		/* get the most recent frame (copy `s->front` to `myframe`) */
 		if (myframe.cap < (size_t)s->front->len)
 			myframe.buf = xrealloc(myframe.buf, myframe.cap = s->front->len);
@@ -141,7 +136,20 @@ void multipart_replace_writer(struct multipart_replace_stream *s, int fd, const
 		if (fwrite("\r\n", 2, 1, stream) < 1)
 			error(EXIT_FAILURE, ferror(stream), "fwrite(\"\\r\\n\")");
 		fflush(stream);
+
+		/* poll until there's a new frame */
+		pthread_rwlock_rdlock(&s->frontlock);
+		while (s->framecount == lastframe && running) {
+			pthread_rwlock_unlock(&s->frontlock);
+			usleep(30000); /* a bit over 30 FPS */
+			pthread_rwlock_rdlock(&s->frontlock);
+		}
 	}
+	pthread_rwlock_unlock(&s->frontlock);
+ end:
+	free(boundary);
+	free(myframe.buf);
+	fclose(stream);
 }
 
 void init_multipart_replace_stream(struct multipart_replace_stream *s) {
@@ -150,4 +158,11 @@ void init_multipart_replace_stream(struct multipart_replace_stream *s) {
 	s->front = &s->a;
 	s->back = &s->b;
 	pthread_rwlock_init(&s->frontlock, NULL);
+	s->framecount = 0;
+}
+
+void destroy_multipart_replace_stream(struct multipart_replace_stream *s) {
+	free(s->a.buf);
+	free(s->b.buf);
+	pthread_rwlock_destroy(&s->frontlock);
 }