From d25cea5420e306af32f5e99d7417aa3e1d45accf Mon Sep 17 00:00:00 2001
From: Riccardo Schirone <sirmy15@gmail.com>
Date: Sat, 2 Jan 2021 21:50:02 +0100
Subject: [PATCH] Rewrite rz_file_mmap API

---
 librz/include/rz_util/rz_file.h |  21 ++-
 librz/include/rz_util/rz_mem.h  |  14 --
 librz/util/buf_mmap.c           |   6 +-
 librz/util/file.c               | 276 +++++++++++++-------------------
 librz/util/mem.c                |  24 ---
 test/unit/test_file.c           |  26 +++
 6 files changed, 153 insertions(+), 214 deletions(-)

diff --git a/librz/include/rz_util/rz_file.h b/librz/include/rz_util/rz_file.h
index 571d65c2ec..cf8f833875 100644
--- a/librz/include/rz_util/rz_file.h
+++ b/librz/include/rz_util/rz_file.h
@@ -5,7 +5,19 @@
 extern "C" {
 #endif
 
-#include <rz_util/rz_mem.h>
+typedef struct rz_mmap_t {
+	ut8 *buf;
+	ut64 base;
+	ut64 len;
+	int perm;
+	int mode;
+	char *filename;
+	int fd;
+#if __WINDOWS__
+	HANDLE fm;
+#endif
+} RzMmap;
+
 
 /* is */
 RZ_API bool rz_file_is_abspath(const char *file);
@@ -16,10 +28,10 @@ RZ_API bool rz_file_is_regular(const char *str);
 RZ_API bool rz_file_truncate(const char *filename, ut64 newsize);
 RZ_API ut64 rz_file_size(const char *str);
 RZ_API char *rz_file_root(const char *root, const char *path);
-RZ_API RMmap *rz_file_mmap(const char *file, bool rw, ut64 base);
-RZ_API int rz_file_mmap_read(const char *file, ut64 addr, ut8 *buf, int len);
+RZ_API RzMmap *rz_file_mmap(const char *file, int perm, int mode, ut64 base);
 RZ_API int rz_file_mmap_write(const char *file, ut64 addr, const ut8 *buf, int len);
-RZ_API void rz_file_mmap_free(RMmap *m);
+RZ_API void *rz_file_mmap_resize(RzMmap *m, ut64 newsize);
+RZ_API void rz_file_mmap_free(RzMmap *m);
 RZ_API bool rz_file_chmod(const char *file, const char *mod, int recursive);
 RZ_API char *rz_file_temp(const char *prefix);
 RZ_API char *rz_file_path(const char *bin);
@@ -55,7 +67,6 @@ RZ_API char *rz_file_tmpdir(void);
 RZ_API char *rz_file_readlink(const char *path);
 RZ_API bool rz_file_copy (const char *src, const char *dst);
 RZ_API RzList* rz_file_globsearch (const char *globbed_path, int maxdepth);
-RZ_API RMmap *rz_file_mmap_arch (RMmap *map, const char *filename, int fd);
 
 #ifdef __cplusplus
 }
diff --git a/librz/include/rz_util/rz_mem.h b/librz/include/rz_util/rz_mem.h
index 4b7cfce05c..aed303d2a2 100644
--- a/librz/include/rz_util/rz_mem.h
+++ b/librz/include/rz_util/rz_mem.h
@@ -5,19 +5,6 @@
 extern "C" {
 #endif
 
-typedef struct rz_mmap_t {
-	ut8 *buf;
-	ut64 base;
-	int len;
-	int fd;
-	int rw;
-	char *filename;
-#if __WINDOWS__
-	HANDLE fh;
-	HANDLE fm;
-#endif
-} RMmap;
-
 typedef struct rz_mem_pool_t {
 	ut8 **nodes;
 	int ncount;
@@ -55,7 +42,6 @@ RZ_API const ut8 *rz_mem_mem_aligned(const ut8 *haystack, int hlen, const ut8 *n
 RZ_API int rz_mem_count(const ut8 **addr);
 RZ_API bool rz_mem_is_printable (const ut8 *a, int la);
 RZ_API bool rz_mem_is_zero(const ut8 *b, int l);
-RZ_API void *rz_mem_mmap_resize(RMmap *m, ut64 newsize);
 
 #ifdef __cplusplus
 }
diff --git a/librz/util/buf_mmap.c b/librz/util/buf_mmap.c
index dbb8597ba2..cf6189770b 100644
--- a/librz/util/buf_mmap.c
+++ b/librz/util/buf_mmap.c
@@ -11,7 +11,7 @@ struct buf_mmap_user {
 struct buf_mmap_priv {
 	// NOTE: this needs to be first, so that bytes operations will work without changes
 	struct buf_bytes_priv bytes_priv;
-	RMmap *mmap;
+	RzMmap *mmap;
 };
 
 static inline struct buf_mmap_priv *get_priv_mmap(RzBuffer *b) {
@@ -27,7 +27,7 @@ static bool buf_mmap_init(RzBuffer *b, const void *user) {
 		return false;
 	}
 
-	priv->mmap = rz_file_mmap (u->filename, u->perm & RZ_PERM_W, 0);
+	priv->mmap = rz_file_mmap (u->filename, u->perm, 0644, 0);
 	if (!priv->mmap) {
 		free (priv);
 		return false;
@@ -49,7 +49,7 @@ static bool buf_mmap_fini(RzBuffer *b) {
 static bool buf_mmap_resize(RzBuffer *b, ut64 newsize) {
 	struct buf_mmap_priv *priv = get_priv_mmap (b);
 	if (newsize > priv->mmap->len) {
-		ut8 *t = rz_mem_mmap_resize (priv->mmap, newsize);
+		ut8 *t = rz_file_mmap_resize (priv->mmap, newsize);
 		if (!t) {
 			return false;
 		}
diff --git a/librz/util/file.c b/librz/util/file.c
index c74eff946a..48f1ff06b8 100644
--- a/librz/util/file.c
+++ b/librz/util/file.c
@@ -909,205 +909,145 @@ err_r_file_mmap_write:
 #endif
 }
 
-RZ_API int rz_file_mmap_read (const char *file, ut64 addr, ut8 *buf, int len) {
 #if __WINDOWS__
-	HANDLE fm = NULL, fh = INVALID_HANDLE_VALUE;
-	LPTSTR file_ = NULL;
-	int ret = -1;
-	file_ = rz_sys_conv_utf8_to_win (file);
-	fh = CreateFile (file_, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, 0);
-	if (fh == INVALID_HANDLE_VALUE) {
-		rz_sys_perror ("rz_file_mmap_read/CreateFile");
-		goto err_r_file_mmap_read;
+static RzMmap *file_mmap(RzMmap *m) {
+	LPTSTR file_ = rz_sys_conv_utf8_to_win (m->filename);
+	bool is_write = (m->perm & O_WRONLY) || (m->perm & O_RDWR);
+	bool is_creat = m->perm & O_CREAT;
+	HANDLE fh = (HANDLE)_get_osfhandle (m->fd);
+	m->len = (DWORD)GetFileSize (fh, (LPDWORD)((char *)&m->len + sizeof (DWORD)));
+	if (m->len == INVALID_FILE_SIZE) {
+		rz_sys_perror ("GetFileSize");
+		goto err;
 	}
-	fm = CreateFileMapping (fh, NULL, PAGE_READONLY, 0, 0, NULL);
-	if (!fm) {
-		rz_sys_perror ("CreateFileMapping");
-		goto err_r_file_mmap_read;
-	}
-	ut8 *obuf = MapViewOfFile (fm, FILE_MAP_READ, 0, 0, len);
-	if (!obuf) {
-		goto err_r_file_mmap_read;
-	}
-	memcpy (obuf, buf, len);
-	UnmapViewOfFile (obuf);
-	ret = len;
-err_r_file_mmap_read:
-	if (fh != INVALID_HANDLE_VALUE) {
-		CloseHandle (fh);
-	}
-	if (fm) {
-		CloseHandle (fm);
-	}
-	free (file_);
-	return ret;
-#elif __UNIX__
-	int fd = rz_sys_open (file, O_RDONLY, 0644);
-	const int pagesize = 4096;
-	int mmlen = len+pagesize;
-	int rest = addr%pagesize;
-	ut8 *mmap_buf;
-	if (fd == -1) {
-		return -1;
-	}
-	mmap_buf = mmap (NULL, mmlen*2, PROT_READ, MAP_SHARED, fd, (off_t)addr-rest);
-	if (((int)(size_t)mmap_buf) == -1) {
-		return -1;
-	}
-	memcpy (buf, mmap_buf+rest, len);
-	munmap (mmap_buf, mmlen*2);
-	close (fd);
-	return len;
-#endif
-	return 0;
-}
+	if (m->len != 0) {
+		m->fm = CreateFileMapping (fh,
+			NULL,
+			is_write? PAGE_READWRITE: PAGE_READONLY,
+			0, 0, NULL);
+		if (!m->fm) {
+			rz_sys_perror ("CreateFileMapping mmap");
+			goto err;
 
-#if __UNIX__
-static RMmap *rz_file_mmap_unix (RMmap *m, int fd) {
-	ut8 empty = m->len == 0;
-	m->buf = mmap (NULL, (empty?BS:m->len) ,
-		m->rw?PROT_READ|PROT_WRITE:PROT_READ,
-		MAP_SHARED, fd, (off_t)m->base);
-	if (m->buf == MAP_FAILED) {
-		RZ_FREE (m);
-	}
-	return m;
-}
-#elif __WINDOWS__
-static RMmap *rz_file_mmap_windows(RMmap *m, const char *file) {
-	LPTSTR file_ = rz_sys_conv_utf8_to_win (file);
-	bool success = false;
-
-	m->fh = CreateFile (file_, GENERIC_READ | (m->rw ? GENERIC_WRITE : 0),
-		FILE_SHARE_READ | FILE_SHARE_WRITE, NULL,
-		OPEN_ALWAYS, FILE_ATTRIBUTE_NORMAL, 0);
-	if (m->fh == INVALID_HANDLE_VALUE) {
-		rz_sys_perror ("CreateFile");
-		goto err_r_file_mmap_windows;
-	}
-	m->fm = CreateFileMapping (m->fh, NULL, PAGE_READONLY, 0, 0, NULL);
-		//m->rw?PAGE_READWRITE:PAGE_READONLY, 0, 0, NULL);
-	if (!m->fm) {
-		rz_sys_perror ("CreateFileMapping");
-		goto err_r_file_mmap_windows;
-
-	}
-	m->buf = MapViewOfFile (m->fm,
-		// m->rw?(FILE_MAP_READ|FILE_MAP_WRITE):FILE_MAP_READ,
-		FILE_MAP_COPY,
-		UT32_HI (m->base), UT32_LO (m->base), 0);
-	success = true;
-err_r_file_mmap_windows:
-	if (!success) {
-		if (m->fh != INVALID_HANDLE_VALUE) {
-			CloseHandle (m->fh);
 		}
-		RZ_FREE (m);
+		m->buf = MapViewOfFileEx (m->fm,
+			is_write? (FILE_MAP_READ|FILE_MAP_WRITE): FILE_MAP_READ,
+			0, 0, 0, (void *)m->base);
+		if (!m->buf) {
+			rz_sys_perror ("MapViewOfFileEx");
+			goto err;
+		}
 	}
+	return m;
+err:
+	rz_file_mmap_free (m);
 	free (file_);
-	return m;
+	return NULL;
 }
-#else
-static RMmap *file_mmap_other (RMmap *m) {
-	ut8 empty = m->len == 0;
-	m->buf = malloc ((empty? BS: m->len));
-	if (!empty && m->buf) {
-		lseek (m->fd, (off_t)0, SEEK_SET);
-		read (m->fd, m->buf, m->len);
-	} else {
-		RZ_FREE (m);
-	}
-	return m;
-}
-#endif
-
-RZ_API RMmap *rz_file_mmap_arch(RMmap *mmap, const char *filename, int fd) {
-#if __WINDOWS__
-	(void)fd;
-	return rz_file_mmap_windows (mmap, filename);
 #elif __UNIX__
-	(void)filename;
-	return rz_file_mmap_unix (mmap, fd);
-#else
-	(void)filename;
-	(void)fd;
-	return file_mmap_other (mmap);
-#endif
+static RzMmap *file_mmap(RzMmap *m) {
+	m->len = lseek (m->fd, (off_t)0, SEEK_END);
+	if (m->len) {
+		bool is_write = (m->perm & O_WRONLY) || (m->perm & O_RDWR);
+		m->buf = mmap ((void *)m->base,
+			m->len,
+			is_write? PROT_READ | PROT_WRITE: PROT_READ,
+			MAP_SHARED, m->fd, 0);
+		if (m->buf == MAP_FAILED) {
+			rz_sys_perror ("mmap");
+			rz_file_mmap_free (m);
+			return NULL;
+		}
+	}
+	return m;
 }
+#else
+static RzMmap *file_mmap(RzMmap *m) {
+	m->len = lseek (m->fd, (off_t)0, SEEK_END);
+	m->buf = malloc (m->len));
+	if (!m->buf) {
+		rz_file_mmap_free (m);
+		return NULL;
+	}
+	lseek (m->fd, (off_t)0, SEEK_SET);
+	read (m->fd, m->buf, m->len);
+	return m;
+}
+#endif
 
-// TODO: add rwx support?
-RZ_API RMmap *rz_file_mmap(const char *file, bool rw, ut64 base) {
-	RMmap *m = NULL;
-	int fd = -1;
-	if (!rw && !rz_file_exists (file)) {
-		return m;
-	}
-	fd = rz_sys_open (file, rw? O_RDWR: O_RDONLY, 0644);
-	if (fd == -1 && !rw) {
-		eprintf ("rz_file_mmap: file does not exis.\n");
-		//m->buf = malloc (m->len);
-		return m;
-	}
-	m = RZ_NEW (RMmap);
+RZ_API RzMmap *rz_file_mmap(const char *file, int perm, int mode, ut64 base) {
+	RzMmap *m = NULL;
+	m = RZ_NEW0 (RzMmap);
 	if (!m) {
-		if (fd != -1) {
-			close (fd);
-		}
 		return NULL;
 	}
 	m->base = base;
-	m->rw = rw;
-	m->fd = fd;
-	m->len = fd != -1? lseek (fd, (off_t)0, SEEK_END) : 0;
+	m->perm = perm;
+	m->len = 0;
 	m->filename = strdup (file);
-
-	if (m->fd == -1) {
-		return m;
-	}
-
-	if (m->len == (off_t)-1) {
-		close (fd);
-		RZ_FREE (m);
+	m->mode = mode;
+	if (!m->filename) {
+		rz_file_mmap_free (m);
 		return NULL;
 	}
-#if __UNIX__
-	return rz_file_mmap_unix (m, fd);
-#elif __WINDOWS__
-	close (fd);
-	m->fd = -1;
-	return rz_file_mmap_windows (m, file);
-#else
-	return file_mmap_other (m);
-#endif
+	m->fd = rz_sys_open (m->filename, m->perm, m->mode);
+	if (m->fd == -1) {
+		rz_file_mmap_free (m);
+		return NULL;
+	}
+	return file_mmap (m);
 }
 
-RZ_API void rz_file_mmap_free(RMmap *m) {
+RZ_API void rz_file_mmap_free(RzMmap *m) {
 	if (!m) {
 		return;
 	}
 #if __WINDOWS__
-	if (m->fm != INVALID_HANDLE_VALUE) {
-		CloseHandle (m->fm);
-	}
-	if (m->fh != INVALID_HANDLE_VALUE) {
-		CloseHandle (m->fh);
-	}
 	if (m->buf) {
 		UnmapViewOfFile (m->buf);
 	}
-#endif
-	if (m->fd == -1) {
-		free (m);
-		return;
+	if (m->fm) {
+		CloseHandle (m->fm);
 	}
-	free (m->filename);
-#if __UNIX__
+	if (m->fd != -1) {
+		_close (m->fd);
+	}
+#elif __UNIX__
 	munmap (m->buf, m->len);
-#endif
 	close (m->fd);
+#endif
+	free (m->filename);
 	free (m);
 }
 
+RZ_API void *rz_file_mmap_resize(RzMmap *m, ut64 newsize) {
+#if __WINDOWS__
+	if (m->buf) {
+		UnmapViewOfFile (m->buf);
+	}
+	if (m->fm) {
+		CloseHandle (m->fm);
+	}
+	if (m->fd != -1) {
+		_close (m->fd);
+	}
+#elif __UNIX__
+	if (m->buf && munmap (m->buf, m->len) != 0) {
+		return NULL;
+	}
+#endif
+	if (!rz_sys_truncate (m->filename, newsize)) {
+		return NULL;
+	}
+	m->fd = rz_sys_open (m->filename, m->perm, m->mode);
+	if (m->fd == -1) {
+		rz_file_mmap_free (m);
+		return NULL;
+	}
+	file_mmap (m);
+	return m->buf;
+}
+
 RZ_API char *rz_file_temp (const char *prefix) {
 	if (!prefix) {
 		prefix = "";
@@ -1347,4 +1287,4 @@ RZ_API char *rz_file_path_join(const char *s1, const char *s2) {
 	bool ends_with_dir = s1[strlen (s1) - 1] == RZ_SYS_DIR[0];
 	const char *sep = ends_with_dir? "": RZ_SYS_DIR;
 	return rz_str_newf ("%s%s%s", s1, sep, s2);
-}
+}
\ No newline at end of file
diff --git a/librz/util/mem.c b/librz/util/mem.c
index e166fc4b36..15ece91e96 100644
--- a/librz/util/mem.c
+++ b/librz/util/mem.c
@@ -344,27 +344,3 @@ RZ_API void rz_mem_memzero(void *dst, size_t l) {
 #endif
 #endif
 }
-
-RZ_API void *rz_mem_mmap_resize(RMmap *m, ut64 newsize) {
-#if __WINDOWS__
-	if (m->fm != INVALID_HANDLE_VALUE) {
-		CloseHandle (m->fm);
-	}
-	if (m->fh != INVALID_HANDLE_VALUE) {
-		CloseHandle (m->fh);
-	}
-	if (m->buf) {
-		UnmapViewOfFile (m->buf);
-	}
-#elif __UNIX__
-	if (munmap (m->buf, m->len) != 0) {
-		return NULL;
-	}
-#endif
-	if (!rz_sys_truncate (m->filename, newsize)) {
-		return NULL;
-	}
-	m->len = newsize;
-	rz_file_mmap_arch (m, m->filename, m->fd);
-	return m->buf;
-}
diff --git a/test/unit/test_file.c b/test/unit/test_file.c
index 82b6135d9e..94e3b3467d 100644
--- a/test/unit/test_file.c
+++ b/test/unit/test_file.c
@@ -53,12 +53,38 @@ bool test_rz_file_dirname(void) {
 	mu_end;
 }
 
+bool test_rz_file_mmap(void) {
+	char *filename = rz_file_temp (NULL);
+	RzMmap *m = rz_file_mmap (filename, O_RDWR | O_CREAT, 0644, 0xdead0000);
+	mu_assert_notnull (m, "filename should be created");
+	mu_assert_streq (m->filename, filename, "filename of mmaped area should ok");
+	mu_assert_eq (m->len, 0, "mmaped file should be empty");
+	mu_assert_eq (m->base, 0xdead0000, "base address is right");
+	mu_assert_eq (m->perm, O_RDWR | O_CREAT, "mmaped perm should be ok");
+	rz_file_mmap_resize (m, 0x10);
+	mu_assert_eq (m->len, 0x10, "mmaped file should be empty");
+	mu_assert_eq (m->base, 0xdead0000, "base address is right");
+	strcpy ((char *)m->buf, "1234567890ABCDEF");
+	rz_file_mmap_free (m);
+
+	m = rz_file_mmap (filename, O_RDONLY, 0644, 0xdead0000);
+	mu_assert_eq (m->len, 0x10, "mmaped file should be empty");
+	mu_assert_eq (m->base, 0xdead0000, "base address is right");
+	mu_assert_streq ((char *)m->buf, "1234567890ABCDEF", "previous data has been written to file");
+	rz_file_mmap_free (m);
+
+	rz_file_rm (filename);
+	free (filename);
+	mu_end;
+}
+
 int all_tests() {
 	size_t i;
 	for (i = 0; i < RELPATH_CASES_COUNT; i++) {
 		mu_run_test (test_rz_file_relpath, relpath_cases[i].base, relpath_cases[i].path, relpath_cases[i].expect);
 	}
 	mu_run_test (test_rz_file_dirname);
+	mu_run_test (test_rz_file_mmap);
 	return tests_passed != tests_run;
 }