1 files changed, 139 insertions, 1 deletions

diff --git a/src/interfaces/libpq/fe-exec.c b/src/interfaces/libpq/fe-exec.c
index 9fb358fc708..bbcf11a687a 100644
--- a/src/interfaces/libpq/fe-exec.c
+++ b/src/interfaces/libpq/fe-exec.c
@@ -8,7 +8,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/interfaces/libpq/fe-exec.c,v 1.206 2010/01/02 16:58:12 momjian Exp $
+ *	  $PostgreSQL: pgsql/src/interfaces/libpq/fe-exec.c,v 1.207 2010/01/21 14:58:53 rhaas Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -3059,6 +3059,144 @@ PQescapeString(char *to, const char *from, size_t length)
 }
 
 
+/*
+ * Escape arbitrary strings.  If as_ident is true, we escape the result
+ * as an identifier; if false, as a literal.  The result is returned in
+ * a newly allocated buffer.  If we fail due to an encoding violation or out
+ * of memory condition, we return NULL, storing an error message into conn.
+ */
+static char *
+PQescapeInternal(PGconn *conn, const char *str, size_t len, int as_ident)
+{
+	const char *s;
+	char   *result;
+	char   *rp;
+	int		num_quotes = 0;		/* single or double, depending on as_ident */
+	int		num_backslashes = 0;
+	int		input_len;
+	int		result_size;
+	char	quote_char = as_ident ? '"' : '\'';
+
+	/* We must have a connection, else fail immediately. */
+	if (!conn)
+		return NULL;
+
+	/* Scan the string for characters that must be escaped. */
+	for (s = str; *s != '\0' && (s - str) < len; ++s)
+	{
+		if (*s == quote_char)
+			++num_quotes;
+		else if (*s == '\\')
+			++num_backslashes;
+		else if (IS_HIGHBIT_SET(*s))
+		{
+			int charlen;
+
+			/* Slow path for possible multibyte characters */
+			charlen = pg_encoding_mblen(conn->client_encoding, s);
+
+			/* Multibyte character overruns allowable length. */
+			if ((s - str) + charlen > len || memchr(s, 0, charlen) != NULL)
+			{
+				printfPQExpBuffer(&conn->errorMessage,
+						  libpq_gettext("incomplete multibyte character\n"));
+				return NULL;
+			}
+
+			/* Adjust s, bearing in mind that for loop will increment it. */
+			s += charlen - 1;
+		}
+	}
+
+	/* Allocate output buffer. */
+	input_len = s - str;
+	result_size = input_len + num_quotes + 3;  /* two quotes, plus a NUL */
+	if (!as_ident && num_backslashes > 0)
+		result_size += num_backslashes + 2;
+	result = rp = (char *) malloc(result_size);
+	if (rp == NULL)
+	{
+		printfPQExpBuffer(&conn->errorMessage,
+						  libpq_gettext("out of memory\n"));
+		return NULL;
+	}
+
+	/*
+	 * If we are escaping a literal that contains backslashes, we use the
+	 * escape string syntax so that the result is correct under either value
+	 * of standard_conforming_strings.  We also emit a leading space in this
+	 * case, to guard against the possibility that the result might be
+	 * interpolated immediately following an identifier.
+	 */
+	if (!as_ident && num_backslashes > 0)
+	{
+		*rp++ = ' ';
+		*rp++ = 'E';
+	}
+
+	/* Opening quote. */
+	*rp++ = quote_char;
+
+	/*
+	 * Use fast path if possible.
+	 *
+	 * We've already verified that the input string is well-formed in the
+	 * current encoding.  If it contains no quotes and, in the case of
+	 * literal-escaping, no backslashes, then we can just copy it directly
+	 * to the output buffer, adding the necessary quotes.
+	 *
+	 * If not, we must rescan the input and process each character
+	 * individually.
+	 */
+	if (num_quotes == 0 && (num_backslashes == 0 || as_ident))
+	{
+		memcpy(rp, str, input_len);
+		rp += input_len;
+	}
+	else
+	{
+		for (s = str; s - str < input_len; ++s)
+		{
+			if (*s == quote_char || (!as_ident && *s == '\\'))
+			{
+				*rp++ = *s;
+				*rp++ = *s;
+			}
+			else if (!IS_HIGHBIT_SET(*s))
+				*rp++ = *s;
+			else
+			{
+				int i = pg_encoding_mblen(conn->client_encoding, s);
+				while (1)
+				{
+					*rp++ = *s;
+					if (--i == 0)
+						break;
+					++s;	/* for loop will provide the final increment */
+				}
+			}
+		}
+	}
+
+	/* Closing quote and terminating NUL. */
+	*rp++ = quote_char;
+	*rp = '\0';
+
+	return result;
+}
+
+char *
+PQescapeLiteral(PGconn *conn, const char *str, size_t len)
+{
+	return PQescapeInternal(conn, str, len, 0);
+}
+
+char *
+PQescapeIdentifier(PGconn *conn, const char *str, size_t len)
+{
+	return PQescapeInternal(conn, str, len, 1);
+}
+
 /* HEX encoding support for bytea */
 static const char hextbl[] = "0123456789abcdef";