From 3d6ec663bb701b7192822f83a6a1cf4f7574d67e Mon Sep 17 00:00:00 2001
From: Tom Lane <tgl@sss.pgh.pa.us>
Date: Wed, 18 Jul 2012 01:13:20 -0400
Subject: [PATCH] Improve pg_upgrade's load_directory() function.

Error out on out-of-memory, rather than returning -1, which the sole
existing caller wasn't checking for anyway.  There doesn't seem to be
any use-case for making the caller check for failure here.

Detect failure return from readdir().

Use a less platform-dependent method of calculating the entrysize.
It's possible, but not yet confirmed, that this explains bug #6733,
in which Mike Wilson reports a pg_upgrade crash that did not occur
in 9.1.  (Note that load_directory is effectively new code in 9.2,
at least on platforms that have scandir().)

Fix up comments, avoid uselessly using two counters, reduce the number
of realloc calls to something sane.
---
 contrib/pg_upgrade/file.c       | 62 ++++++++++++++++++---------------
 contrib/pg_upgrade/pg_upgrade.h |  3 +-
 contrib/pg_upgrade/util.c       | 13 ++++++-
 3 files changed, 47 insertions(+), 31 deletions(-)

diff --git a/contrib/pg_upgrade/file.c b/contrib/pg_upgrade/file.c
index 1dd3722142c..962cbaccc53 100644
--- a/contrib/pg_upgrade/file.c
+++ b/contrib/pg_upgrade/file.c
@@ -224,13 +224,12 @@ copy_file(const char *srcfile, const char *dstfile, bool force)
 /*
  * load_directory()
  *
- * Returns count of files that meet the selection criteria coded in
- * the function pointed to by selector.  Creates an array of pointers
- * to dirent structures.  Address of array returned in namelist.
+ * Read all the file names in the specified directory, and return them as
+ * an array of "struct dirent" pointers.  The array address is returned in
+ * *namelist, and the function result is the count of file names.
  *
- * Note that the number of dirent structures needed is dynamically
- * allocated using realloc.  Realloc can be inefficient if invoked a
- * large number of times.
+ * To free the result data, free each namelist array member, then free the
+ * namelist array itself.
  */
 int
 load_directory(const char *dirname, struct dirent *** namelist)
@@ -238,43 +237,48 @@ load_directory(const char *dirname, struct dirent *** namelist)
 	DIR		   *dirdesc;
 	struct dirent *direntry;
 	int			count = 0;
-	int			name_num = 0;
+	int			allocsize = 64;
 	size_t		entrysize;
 
-	if ((dirdesc = opendir(dirname)) == NULL)
-		pg_log(PG_FATAL, "could not open directory \"%s\": %s\n", dirname, getErrorText(errno));
+	*namelist = (struct dirent **)
+		pg_malloc(allocsize * sizeof(struct dirent *));
 
-	*namelist = NULL;
+	if ((dirdesc = opendir(dirname)) == NULL)
+		pg_log(PG_FATAL, "could not open directory \"%s\": %s\n",
+			   dirname, getErrorText(errno));
 
-	while ((direntry = readdir(dirdesc)) != NULL)
+	while (errno = 0, (direntry = readdir(dirdesc)) != NULL)
 	{
-		count++;
-
-		*namelist = (struct dirent **) realloc((void *) (*namelist),
-						(size_t) ((name_num + 1) * sizeof(struct dirent *)));
-
-		if (*namelist == NULL)
+		if (count >= allocsize)
 		{
-			closedir(dirdesc);
-			return -1;
+			allocsize *= 2;
+			*namelist = (struct dirent **)
+				pg_realloc(*namelist, allocsize * sizeof(struct dirent *));
 		}
 
-		entrysize = sizeof(struct dirent) - sizeof(direntry->d_name) +
+		entrysize = offsetof(struct dirent, d_name) +
 			strlen(direntry->d_name) + 1;
 
-		(*namelist)[name_num] = (struct dirent *) malloc(entrysize);
-
-		if ((*namelist)[name_num] == NULL)
-		{
-			closedir(dirdesc);
-			return -1;
-		}
+		(*namelist)[count] = (struct dirent *) pg_malloc(entrysize);
 
-		memcpy((*namelist)[name_num], direntry, entrysize);
+		memcpy((*namelist)[count], direntry, entrysize);
 
-		name_num++;
+		count++;
 	}
 
+#ifdef WIN32
+	/*
+	 * This fix is in mingw cvs (runtime/mingwex/dirent.c rev 1.4), but not in
+	 * released version
+	 */
+	if (GetLastError() == ERROR_NO_MORE_FILES)
+		errno = 0;
+#endif
+
+	if (errno)
+		pg_log(PG_FATAL, "could not read directory \"%s\": %s\n",
+			   dirname, getErrorText(errno));
+
 	closedir(dirdesc);
 
 	return count;
diff --git a/contrib/pg_upgrade/pg_upgrade.h b/contrib/pg_upgrade/pg_upgrade.h
index 3274227a0a3..4f74c217eed 100644
--- a/contrib/pg_upgrade/pg_upgrade.h
+++ b/contrib/pg_upgrade/pg_upgrade.h
@@ -429,7 +429,8 @@ prep_status(const char *fmt,...)
 __attribute__((format(PG_PRINTF_ATTRIBUTE, 1, 2)));
 void		check_ok(void);
 char	   *pg_strdup(const char *s);
-void	   *pg_malloc(int size);
+void	   *pg_malloc(size_t size);
+void	   *pg_realloc(void *ptr, size_t size);
 void		pg_free(void *ptr);
 const char *getErrorText(int errNum);
 unsigned int str2uint(const char *str);
diff --git a/contrib/pg_upgrade/util.c b/contrib/pg_upgrade/util.c
index 6977663b63a..76cd20b23d7 100644
--- a/contrib/pg_upgrade/util.c
+++ b/contrib/pg_upgrade/util.c
@@ -183,7 +183,7 @@ get_user_info(char **user_name)
 
 
 void *
-pg_malloc(int n)
+pg_malloc(size_t n)
 {
 	void	   *p = malloc(n);
 
@@ -193,6 +193,17 @@ pg_malloc(int n)
 	return p;
 }
 
+void *
+pg_realloc(void *ptr, size_t n)
+{
+	void	   *p = realloc(ptr, n);
+
+	if (p == NULL)
+		pg_log(PG_FATAL, "%s: out of memory\n", os_info.progname);
+
+	return p;
+}
+
 
 void
 pg_free(void *p)
-- 
GitLab