From ec6acfbc1c8cbbd2e5c7578df2f7ac2d94e0c17a Mon Sep 17 00:00:00 2001
From: Unbit <info@unbit.it>
Date: Wed, 1 May 2013 19:36:13 +0200
Subject: [PATCH] improved transformation system

---
 core/routing.c                           |  5 +-
 core/transformations.c                   | 67 ++++++++++++++++--------
 plugins/transformation_chunked/chunked.c | 18 ++++---
 plugins/transformation_gzip/gzip.c       |  9 ++--
 plugins/transformation_tofile/tofile.c   |  3 +-
 uwsgi.h                                  |  2 +
 6 files changed, 71 insertions(+), 33 deletions(-)

diff --git a/core/routing.c b/core/routing.c
index c9a36694..25dafdae 100644
--- a/core/routing.c
+++ b/core/routing.c
@@ -494,6 +494,7 @@ static int transform_flush(struct wsgi_request *wsgi_req, struct uwsgi_transform
 	int ret = uwsgi_response_write_body_do(wsgi_req, ut->chunk->buf, ut->chunk->pos);
 	wsgi_req->transformed_chunk = NULL;
 	wsgi_req->transformed_chunk_len = 0;
+	ut->flushed = 1;
 	return ret;
 }
 static int uwsgi_router_flush_func(struct wsgi_request *wsgi_req, struct uwsgi_route *route) {
@@ -514,7 +515,9 @@ static int transform_fixcl(struct wsgi_request *wsgi_req, struct uwsgi_transform
                 wsgi_req->write_errors++;
                 return -1;
         }
-        return uwsgi_response_add_header(wsgi_req, "Content-Length", 14, buf, ret);
+	// do not check for errors !!!
+        uwsgi_response_add_header(wsgi_req, "Content-Length", 14, buf, ret);
+	return 0;
 }
 static int uwsgi_router_fixcl_func(struct wsgi_request *wsgi_req, struct uwsgi_route *route) {
         uwsgi_add_transformation(wsgi_req, transform_fixcl, NULL);
diff --git a/core/transformations.c b/core/transformations.c
index 8b38c1cb..32a0d919 100644
--- a/core/transformations.c
+++ b/core/transformations.c
@@ -4,10 +4,11 @@
 
 	uWSGI transformations
 
-	responses can be buffered in the wsgi_request structure (instead of being sent to the client)
-	Before closing the request, the transformations are applied in chain to the response buffer
+	each body chunk is passed (in chain) to every transformation
 
-	Finally the resulting buffer is sent to the client
+	some of them supports streaming, other requires buffering
+
+	at the end of the request the "final chain" is called (and the whole chain freed)
 
 	Transformations (if required) could completely swallow already set headers
 
@@ -22,6 +23,7 @@ int uwsgi_apply_transformations(struct wsgi_request *wsgi_req, char *buf, size_t
 	struct uwsgi_transformation *ut = wsgi_req->transformations;
 	char *t_buf = buf;
 	size_t t_len = len;
+	uint8_t flushed = 0;
 	while(ut) {
 		// allocate the buffer (if needed)
 		if (!ut->chunk) {
@@ -36,10 +38,13 @@ int uwsgi_apply_transformations(struct wsgi_request *wsgi_req, char *buf, size_t
 		// if the transformation cannot stream, continue buffering (the func will be called at the end)
 		if (!ut->can_stream) return 1;
 		
+		ut->round++;
 		if (ut->func(wsgi_req, ut)) {
 			return -1;
 		}
 
+		if (ut->flushed) flushed = 1;
+
 		t_buf = ut->chunk->buf;
 		t_len = ut->chunk->pos;
 		// we reset the buffer, so we do not waste memory
@@ -50,54 +55,74 @@ next:
 
 	// if we are here we can tell the writer to send the body to the client
 	// no buffering please
-	wsgi_req->transformed_chunk = t_buf;
-	wsgi_req->transformed_chunk_len = t_len;
+	if (!flushed) {
+		wsgi_req->transformed_chunk = t_buf;
+		wsgi_req->transformed_chunk_len = t_len;
+	}
 	return 0;
 
 }
 
 /*
 
-if a transformation is "final", we need to call it independently by the write status
+run all the remaining (or buffered) transformations
 
 */
 int uwsgi_apply_final_transformations(struct wsgi_request *wsgi_req) {
-	int ret = 0;
 	struct uwsgi_transformation *ut = wsgi_req->transformations;
 	wsgi_req->transformed_chunk = NULL;
         wsgi_req->transformed_chunk_len = 0;
 	char *t_buf = NULL;
 	size_t t_len = 0;
+	uint8_t flushed = 0;
+	int found_nostream = 0;
 	while(ut) {
-		if (ut->chunk) {
-			t_buf = ut->chunk->buf;
-			t_len = ut->chunk->pos;
-			// if the transformation can stream and has a chunk, we already applied it
-			if (ut->can_stream) goto next;
+		if (!found_nostream) {
+			if (!ut->can_stream) {
+				found_nostream = 1;
+			}
+			else {
+				t_buf = ut->chunk->buf;
+                		t_len = ut->chunk->pos;
+				goto next;
+			}
 		}
-		else if (t_len > 0) {
-			ut->chunk = uwsgi_buffer_new(t_len);
-			if (uwsgi_buffer_append(ut->chunk, t_buf, t_len)) {
-				return -1;
+
+		if (!ut->chunk) {
+			if (t_len > 0) {
+				ut->chunk = uwsgi_buffer_new(t_len);
+			}
+			else {
+				ut->chunk = uwsgi_buffer_new(uwsgi.page_size);
 			}
 		}
 		
+		if (t_len > 0) {
+			if (uwsgi_buffer_append(ut->chunk, t_buf, t_len)) {
+				return -1;
+                	}
+		}
+		
 		// run the transformation
+		ut->round++;
 		if (ut->func(wsgi_req, ut)) {
-			ret = -1;
+			return -1;
                 }
 
+		if (ut->flushed) flushed = 1;
+
 		t_buf = ut->chunk->buf;
 		t_len = ut->chunk->pos;
-		
 next:
 		ut = ut->next;
 	}
 
 	// if we are here, all of the transformations are applied
-	wsgi_req->transformed_chunk = t_buf;
-        wsgi_req->transformed_chunk_len = t_len;
-        return ret;
+	if (!flushed) {
+		wsgi_req->transformed_chunk = t_buf;
+        	wsgi_req->transformed_chunk_len = t_len;
+	}
+        return 0;
 }
 
 void uwsgi_free_transformations(struct wsgi_request *wsgi_req) {
diff --git a/plugins/transformation_chunked/chunked.c b/plugins/transformation_chunked/chunked.c
index 5bc0c067..e5a1f78a 100644
--- a/plugins/transformation_chunked/chunked.c
+++ b/plugins/transformation_chunked/chunked.c
@@ -10,16 +10,22 @@
 
 static int transform_chunked(struct wsgi_request *wsgi_req, struct uwsgi_transformation *ut) {
 	struct uwsgi_buffer *ub = ut->chunk;
-	if (!wsgi_req->headers_sent) {
-        	if (uwsgi_response_add_header(wsgi_req, "Transfer-Encoding", 17, "chunked", 7)) return -1;
-	}
+
 	if (ut->is_final) {
-		if (uwsgi_buffer_insert_chunked(ub, 0, 0)) return -1;
+		if (uwsgi_buffer_append(ub, "0\r\n\r\n", 5)) return -1;
+		return 0;
 	}
-	else {
+
+	if (ut->round == 1) {
+		// do not check for errors !!!
+        	uwsgi_response_add_header(wsgi_req, "Transfer-Encoding", 17, "chunked", 7);
+	}
+
+	if (ub->pos > 0) {
 		if (uwsgi_buffer_insert_chunked(ub, 0, ub->pos)) return -1;
+		if (uwsgi_buffer_append(ub, "\r\n", 2)) return -1;
 	}
-	if (uwsgi_buffer_append(ub, "\r\n", 2)) return -1;
+
 	return 0;
 }
 
diff --git a/plugins/transformation_gzip/gzip.c b/plugins/transformation_gzip/gzip.c
index 97773234..645c7736 100644
--- a/plugins/transformation_gzip/gzip.c
+++ b/plugins/transformation_gzip/gzip.c
@@ -22,11 +22,8 @@ extern char gzheader[];
 static int transform_gzip(struct wsgi_request *wsgi_req, struct uwsgi_transformation *ut) {
 	struct uwsgi_transformation_gzip *utgz = (struct uwsgi_transformation_gzip *) ut->data;
 	struct uwsgi_buffer *ub = ut->chunk;
-	if (!wsgi_req->headers_sent) {
-        	if (uwsgi_response_add_header(wsgi_req, "Content-Encoding", 16, "gzip", 4)) return -1;
-	}
+
 	if (ut->is_final) {
-		uwsgi_log("size = %d %d\n", ub->pos, utgz->len);
 		if (uwsgi_gzip_fix(&utgz->z, utgz->crc32, ub, utgz->len)) {
 			free(utgz);
 			return -1;
@@ -34,17 +31,21 @@ static int transform_gzip(struct wsgi_request *wsgi_req, struct uwsgi_transforma
 		free(utgz);
 		return 0;
 	}
+
 	size_t dlen = 0;
 	char *gzipped = uwsgi_gzip_chunk(&utgz->z, &utgz->crc32, ub->buf, ub->pos, &dlen);
 	if (!gzipped) return -1;
 	utgz->len += ub->pos;
 	uwsgi_buffer_map(ub, gzipped, dlen);
 	if (!utgz->header) {
+		// do not check for errors !!!
+        	uwsgi_response_add_header(wsgi_req, "Content-Encoding", 16, "gzip", 4);
 		utgz->header = 1;
 		if (uwsgi_buffer_insert(ub, 0, gzheader, 10)) {
 			return -1;
 		}
 	}
+
 	return 0;
 }
 
diff --git a/plugins/transformation_tofile/tofile.c b/plugins/transformation_tofile/tofile.c
index c1827046..72b795ac 100644
--- a/plugins/transformation_tofile/tofile.c
+++ b/plugins/transformation_tofile/tofile.c
@@ -28,7 +28,7 @@ static int transform_tofile(struct wsgi_request *wsgi_req, struct uwsgi_transfor
 	// store only successfull response
 	if (wsgi_req->write_errors == 0 && wsgi_req->status == 200 && ub->pos > 0) {
 		if (uttc->filename) {
-			int fd = open(uttc->filename->buf, O_WRONLY|O_CREAT, S_IRUSR|S_IWUSR|S_IRGRP|S_IROTH);
+			int fd = open(uttc->filename->buf, O_WRONLY|O_CREAT|O_TRUNC, S_IRUSR|S_IWUSR|S_IRGRP|S_IROTH);
 			if (fd < 0) {
 				uwsgi_error_open(uttc->filename->buf);
 				goto end;
@@ -58,6 +58,7 @@ end:
 	// free resources
 	if (uttc->filename) uwsgi_buffer_destroy(uttc->filename);
 	free(uttc);
+	// reset the buffer
         return 0;
 }
 
diff --git a/uwsgi.h b/uwsgi.h
index 06d13a5f..e121d137 100644
--- a/uwsgi.h
+++ b/uwsgi.h
@@ -1199,7 +1199,9 @@ struct uwsgi_transformation {
 	struct uwsgi_buffer *chunk;
 	uint8_t can_stream;
 	uint8_t is_final;
+	uint8_t flushed;
 	void *data;
+	uint64_t round;
 	struct uwsgi_transformation *next;
 };