2 files changed, 93 insertions, 17 deletions

diff --git a/src/backend/utils/adt/pg_locale.c b/src/backend/utils/adt/pg_locale.c
index f2a28d5ef5a..b4954959f98 100644
--- a/src/backend/utils/adt/pg_locale.c
+++ b/src/backend/utils/adt/pg_locale.c
@@ -58,6 +58,7 @@
 #include "catalog/pg_collation.h"
 #include "catalog/pg_database.h"
 #include "common/hashfn.h"
+#include "common/string.h"
 #include "mb/pg_wchar.h"
 #include "miscadmin.h"
 #include "utils/builtins.h"
@@ -341,6 +342,16 @@ check_locale(int category, const char *locale, char **canonname)
 	char	   *save;
 	char	   *res;
 
+	/* Don't let Windows' non-ASCII locale names in. */
+	if (!pg_is_ascii(locale))
+	{
+		ereport(WARNING,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("locale name \"%s\" contains non-ASCII characters",
+						locale)));
+		return false;
+	}
+
 	if (canonname)
 		*canonname = NULL;		/* in case of failure */
 
@@ -363,6 +374,18 @@ check_locale(int category, const char *locale, char **canonname)
 		elog(WARNING, "failed to restore old locale \"%s\"", save);
 	pfree(save);
 
+	/* Don't let Windows' non-ASCII locale names out. */
+	if (canonname && *canonname && !pg_is_ascii(*canonname))
+	{
+		ereport(WARNING,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("locale name \"%s\" contains non-ASCII characters",
+						*canonname)));
+		pfree(*canonname);
+		*canonname = NULL;
+		return false;
+	}
+
 	return (res != NULL);
 }
 
diff --git a/src/bin/initdb/initdb.c b/src/bin/initdb/initdb.c
index eac6cbd8e51..34db43f47b7 100644
--- a/src/bin/initdb/initdb.c
+++ b/src/bin/initdb/initdb.c
@@ -340,6 +340,61 @@ do { \
 		output_failed = true, output_errno = errno; \
 } while (0)
 
+#ifdef WIN32
+typedef wchar_t *save_locale_t;
+#else
+typedef char *save_locale_t;
+#endif
+
+/*
+ * Save a copy of the current global locale's name, for the given category.
+ * The returned value must be passed to restore_global_locale().
+ *
+ * Since names from the environment haven't been vetted for non-ASCII
+ * characters, we use the wchar_t variant of setlocale() on Windows.  Otherwise
+ * they might not survive a save-restore round trip: when restoring, the name
+ * itself might be interpreted with a different encoding by plain setlocale(),
+ * after we switch to another locale in between.  (This is a problem only in
+ * initdb, not in similar backend code where the global locale's name should
+ * already have been verified as ASCII-only.)
+ */
+static save_locale_t
+save_global_locale(int category)
+{
+	save_locale_t save;
+
+#ifdef WIN32
+	save = _wsetlocale(category, NULL);
+	if (!save)
+		pg_fatal("_wsetlocale() failed");
+	save = wcsdup(save);
+	if (!save)
+		pg_fatal("out of memory");
+#else
+	save = setlocale(category, NULL);
+	if (!save)
+		pg_fatal("setlocale() failed");
+	save = pg_strdup(save);
+#endif
+	return save;
+}
+
+/*
+ * Restore the global locale returned by save_global_locale().
+ */
+static void
+restore_global_locale(int category, save_locale_t save)
+{
+#ifdef WIN32
+	if (!_wsetlocale(category, save))
+		pg_fatal("failed to restore old locale");
+#else
+	if (!setlocale(category, save))
+		pg_fatal("failed to restore old locale \"%s\"", save);
+#endif
+	free(save);
+}
+
 /*
  * Escape single quotes and backslashes, suitably for insertions into
  * configuration files or SQL E'' strings.
@@ -2074,16 +2129,13 @@ locale_date_order(const char *locale)
 	char	   *posD;
 	char	   *posM;
 	char	   *posY;
-	char	   *save;
+	save_locale_t save;
 	size_t		res;
 	int			result;
 
 	result = DATEORDER_MDY;		/* default */
 
-	save = setlocale(LC_TIME, NULL);
-	if (!save)
-		return result;
-	save = pg_strdup(save);
+	save = save_global_locale(LC_TIME);
 
 	setlocale(LC_TIME, locale);
 
@@ -2094,8 +2146,7 @@ locale_date_order(const char *locale)
 
 	res = my_strftime(buf, sizeof(buf), "%x", &testtime);
 
-	setlocale(LC_TIME, save);
-	free(save);
+	restore_global_locale(LC_TIME, save);
 
 	if (res == 0)
 		return result;
@@ -2132,18 +2183,17 @@ locale_date_order(const char *locale)
 static void
 check_locale_name(int category, const char *locale, char **canonname)
 {
-	char	   *save;
+	save_locale_t save;
 	char	   *res;
 
+	/* Don't let Windows' non-ASCII locale names in. */
+	if (locale && !pg_is_ascii(locale))
+		pg_fatal("locale name \"%s\" contains non-ASCII characters", locale);
+
 	if (canonname)
 		*canonname = NULL;		/* in case of failure */
 
-	save = setlocale(category, NULL);
-	if (!save)
-		pg_fatal("setlocale() failed");
-
-	/* save may be pointing at a modifiable scratch variable, so copy it. */
-	save = pg_strdup(save);
+	save = save_global_locale(category);
 
 	/* for setlocale() call */
 	if (!locale)
@@ -2157,9 +2207,7 @@ check_locale_name(int category, const char *locale, char **canonname)
 		*canonname = pg_strdup(res);
 
 	/* restore old value. */
-	if (!setlocale(category, save))
-		pg_fatal("failed to restore old locale \"%s\"", save);
-	free(save);
+	restore_global_locale(category, save);
 
 	/* complain if locale wasn't valid */
 	if (res == NULL)
@@ -2183,6 +2231,11 @@ check_locale_name(int category, const char *locale, char **canonname)
 			pg_fatal("invalid locale settings; check LANG and LC_* environment variables");
 		}
 	}
+
+	/* Don't let Windows' non-ASCII locale names out. */
+	if (canonname && !pg_is_ascii(*canonname))
+		pg_fatal("locale name \"%s\" contains non-ASCII characters",
+				 *canonname);
 }
 
 /*