From 9c3dd915ea00ce33d6dee359bcef963a586aaf05 Mon Sep 17 00:00:00 2001
From: Nikias Bassen
Date: Sun, 13 Dec 2020 00:55:33 +0100
Subject: idevicescreenshot: Remove unnecessary math functions to find a unique
 filename

Removes repeated malloc/free too, limits the maximum filenames to try to 65535
and fails if no unique filename could be found. I think this limit is VERY
generous. If there are really so many files to test against it will also cause
an unconveniently long delay, but I felt a lower limit could still be easily
reached.
---
 tools/idevicescreenshot.c | 38 +++++++++++++++++++++++---------------
 1 file changed, 23 insertions(+), 15 deletions(-)

(limited to 'tools/idevicescreenshot.c')

diff --git a/tools/idevicescreenshot.c b/tools/idevicescreenshot.c
index f33e83e..62bb4a3 100644
--- a/tools/idevicescreenshot.c
+++ b/tools/idevicescreenshot.c
@@ -29,7 +29,6 @@
 #include <string.h>
 #include <stdlib.h>
 #include <errno.h>
-#include <math.h>
 #include <time.h>
 #include <unistd.h>
 #ifndef WIN32
@@ -121,17 +120,21 @@ int main(int argc, char **argv)
 			uint64_t imgsize = 0;
 			if (screenshotr_take_screenshot(shotr, &imgdata, &imgsize) == SCREENSHOTR_E_SUCCESS) {
 				get_image_filename(imgdata, &filename);
-				FILE *f = fopen(filename, "wb");
-				if (f) {
-					if (fwrite(imgdata, 1, (size_t)imgsize, f) == (size_t)imgsize) {
-						printf("Screenshot saved to %s\n", filename);
-						result = 0;
+				if (!filename) {
+					printf("FATAL: Could not find a unique filename!\n");
+				} else {
+					FILE *f = fopen(filename, "wb");
+					if (f) {
+						if (fwrite(imgdata, 1, (size_t)imgsize, f) == (size_t)imgsize) {
+							printf("Screenshot saved to %s\n", filename);
+							result = 0;
+						} else {
+							printf("Could not save screenshot to file %s!\n", filename);
+						}
+						fclose(f);
 					} else {
-						printf("Could not save screenshot to file %s!\n", filename);
+						printf("Could not open %s for writing: %s\n", filename, strerror(errno));
 					}
-					fclose(f);
-				} else {
-					printf("Could not open %s for writing: %s\n", filename, strerror(errno));
 				}
 			} else {
 				printf("Could not get screenshot!\n");
@@ -179,6 +182,7 @@ void get_image_filename(char *imgdata, char **filename)
 		basename = (char*)malloc(strlen(*filename) + 1);
 		strcpy(basename, *filename);
 		free(*filename);
+		*filename = NULL;
 	} else {
 		time_t now = time(NULL);
 		basename = (char*)malloc(32);
@@ -186,15 +190,19 @@ void get_image_filename(char *imgdata, char **filename)
 	}
 
 	// Ensure the filename is unique on disk.
-	char *unique_filename = (char*)malloc(strlen(basename) + strlen(fileext) + 1);
+	char *unique_filename = (char*)malloc(strlen(basename) + strlen(fileext) + 7);
 	sprintf(unique_filename, "%s%s", basename, fileext);
 	int i;
-	for (i = 2; access(unique_filename, F_OK) != -1; i++) {
-		free(unique_filename);
-		unique_filename = (char*)malloc(strlen(basename) + strlen(fileext) + floor(log10(i)) + 3);
+	for (i = 2; i < (1 << 16); i++) {
+		if (access(unique_filename, F_OK) == -1) {
+			*filename = unique_filename;
+			break;
+		}
 		sprintf(unique_filename, "%s-%d%s", basename, i, fileext);
 	}
-	*filename = unique_filename;
+	if (!*filename) {
+		free(unique_filename);
+	}
 	free(basename);
 }
 
-- 
cgit v1.1-32-gdbae