From 2b658ed275f269d502baa7b6159b6d2020486169 Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki.linnakangas@iki.fi>
Date: Wed, 16 Dec 2020 10:41:49 +0200
Subject: [PATCH 5/5] Do COPY FROM encoding conversion/verification in larger
 chunks.

NOTE: This changes behavior in one corner-case: if client and server
encodings are the same single-byte encoding (e.g. latin1), previously the
input would not be checked for zero bytes ('\0'). Any fields containing
zero bytes would be truncated at the zero. But if encoding conversion was
needed, the conversion routine would throw an error on the zero. After
this commit, the input is always checked for zeros.
---
 src/backend/commands/copyfrom.c          |  41 +++--
 src/backend/commands/copyfromparse.c     | 213 +++++++++++++++++------
 src/backend/utils/mb/mbutils.c           |  55 ++++++
 src/include/commands/copyfrom_internal.h |  27 ++-
 src/include/mb/pg_wchar.h                |   6 +
 5 files changed, 262 insertions(+), 80 deletions(-)

diff --git a/src/backend/commands/copyfrom.c b/src/backend/commands/copyfrom.c
index 1b14e9a6eb0..6c72167342b 100644
--- a/src/backend/commands/copyfrom.c
+++ b/src/backend/commands/copyfrom.c
@@ -23,6 +23,7 @@
 #include "access/tableam.h"
 #include "access/xact.h"
 #include "access/xlog.h"
+#include "catalog/namespace.h"
 #include "commands/copy.h"
 #include "commands/copyfrom_internal.h"
 #include "commands/trigger.h"
@@ -147,15 +148,9 @@ CopyFromErrorCallback(void *arg)
 			/*
 			 * Error is relevant to a particular line.
 			 *
-			 * If line_buf still contains the correct line, and it's already
-			 * transcoded, print it. If it's still in a foreign encoding, it's
-			 * quite likely that the error is precisely a failure to do
-			 * encoding conversion (ie, bad data). We dare not try to convert
-			 * it, and at present there's no way to regurgitate it without
-			 * conversion. So we have to punt and just report the line number.
+			 * If line_buf still contains the correct line, print it.
 			 */
-			if (cstate->line_buf_valid &&
-				(cstate->line_buf_converted || !cstate->need_transcoding))
+			if (cstate->line_buf_valid)
 			{
 				char	   *lineval;
 
@@ -1301,15 +1296,22 @@ BeginCopyFrom(ParseState *pstate,
 		cstate->file_encoding = cstate->opts.file_encoding;
 
 	/*
-	 * Set up encoding conversion info.  Even if the file and server encodings
-	 * are the same, we must apply pg_any_to_server() to validate data in
-	 * multibyte encodings.
+	 * Set up encoding conversion info.  If the file and server encodings are
+	 * the same, no conversion is required by we must still validate that the
+	 * data is valid for the encoding.
 	 */
-	cstate->need_transcoding =
-		(cstate->file_encoding != GetDatabaseEncoding() ||
-		 pg_database_encoding_max_length() > 1);
-	/* See Multibyte encoding comment above */
-	cstate->encoding_embeds_ascii = PG_ENCODING_IS_CLIENT_ONLY(cstate->file_encoding);
+	if (cstate->file_encoding == GetDatabaseEncoding() ||
+		cstate->file_encoding == PG_SQL_ASCII ||
+		GetDatabaseEncoding() == PG_SQL_ASCII)
+	{
+		cstate->need_transcoding = false;
+	}
+	else
+	{
+		cstate->need_transcoding = true;
+		cstate->conversion_proc = FindDefaultConversionProc(cstate->file_encoding,
+															GetDatabaseEncoding());
+	}
 
 	cstate->copy_src = COPY_FILE;	/* default */
 
@@ -1338,7 +1340,12 @@ BeginCopyFrom(ParseState *pstate,
 	if (!cstate->opts.binary)
 	{
 		initStringInfo(&cstate->line_buf);
-		cstate->line_buf_converted = false;
+
+		if (cstate->need_transcoding)
+		{
+			cstate->conversion_buf = palloc(CONVERSION_BUF_SIZE + 1);
+			cstate->conversion_buf_index = cstate->conversion_buf_len = 0;
+		}
 	}
 
 	/* Assign range table, we'll need it in CopyFrom. */
diff --git a/src/backend/commands/copyfromparse.c b/src/backend/commands/copyfromparse.c
index 34ed3cfcd5b..ffeebed43e3 100644
--- a/src/backend/commands/copyfromparse.c
+++ b/src/backend/commands/copyfromparse.c
@@ -116,7 +116,8 @@ static int	CopyGetData(CopyFromState cstate, void *databuf,
 						int minread, int maxread);
 static inline bool CopyGetInt32(CopyFromState cstate, int32 *val);
 static inline bool CopyGetInt16(CopyFromState cstate, int16 *val);
-static bool CopyLoadRawBuf(CopyFromState cstate);
+static bool CopyLoadRawBufText(CopyFromState cstate);
+static bool CopyLoadRawBufBinary(CopyFromState cstate);
 static int	CopyReadBinaryData(CopyFromState cstate, char *dest, int nbytes);
 
 void
@@ -357,6 +358,65 @@ CopyGetInt16(CopyFromState cstate, int16 *val)
 	return true;
 }
 
+/*
+ * Convert input data from 'conversion_buf', writing it into
+ * 'raw_buf'.
+ *
+ * 'conversion_buf' mustn't be empty.
+ */
+static void
+CopyConvertBuf(CopyFromState cstate)
+{
+	int			convertedbytes;
+	int			srclen;
+	int			dstlen;
+
+	Assert(cstate->raw_buf_index == 0);
+
+	srclen = cstate->conversion_buf_len - cstate->conversion_buf_index;
+	dstlen = RAW_BUF_SIZE - cstate->raw_buf_len + 1;
+
+	/*
+	 * Do the conversion. This might stop short, if there is an invalid byte
+	 * sequence in the input. We'll convert as much as we can in that case.
+	 *
+	 * Note: Even if we hit an invalid byte sequence, we don't report the error
+	 * until all the valid bytes have been consumed. The input might contain
+	 * an end-of-input marker (\.), and we don't want to report an error if
+	 * the invalid byte sequence is after the end-of-input marker. We might
+	 * still convert extra data after the end-of-input marker if it's valid
+	 * for the encoding, but that's harmless.
+	 */
+	convertedbytes = pg_do_encoding_conversion_buf(cstate->conversion_proc,
+												   cstate->file_encoding,
+												   GetDatabaseEncoding(),
+												   (unsigned char *) cstate->conversion_buf + cstate->conversion_buf_index,
+												   srclen,
+												   (unsigned char *) cstate->raw_buf + cstate->raw_buf_len,
+												   dstlen,
+												   true);
+	if (convertedbytes == 0)
+	{
+		/*
+		 * No more valid input in the buffer, and we have hit an invalid byte sequence.
+		 * Let the conversion function throw the error.
+		 */
+		convertedbytes = pg_do_encoding_conversion_buf(cstate->conversion_proc,
+													   cstate->file_encoding,
+													   GetDatabaseEncoding(),
+													   (unsigned char *) cstate->conversion_buf + cstate->conversion_buf_index,
+													   srclen,
+													   (unsigned char *) cstate->raw_buf + cstate->raw_buf_len,
+													   dstlen,
+													   false);
+		/* pg_do_encoding_conversion_buf should've reported the error */
+		Assert(convertedbytes == 0);
+		elog(ERROR, "conversion error");
+	}
+	cstate->conversion_buf_index += convertedbytes;
+	cstate->raw_buf_len += strlen(cstate->raw_buf + cstate->raw_buf_len);
+	cstate->valid_raw_buf_len = cstate->raw_buf_len;
+}
 
 /*
  * CopyLoadRawBuf loads some more data into raw_buf
@@ -368,7 +428,90 @@ CopyGetInt16(CopyFromState cstate, int16 *val)
  * when a multibyte character crosses a bufferload boundary.
  */
 static bool
-CopyLoadRawBuf(CopyFromState cstate)
+CopyLoadRawBufText(CopyFromState cstate)
+{
+	int			nbytes = RAW_BUF_BYTES(cstate);
+	int			inbytes;
+
+	/* Copy down the unprocessed data if any. */
+	if (nbytes > 0)
+	{
+		memmove(cstate->raw_buf, cstate->raw_buf + cstate->raw_buf_index,
+				nbytes);
+	}
+	cstate->raw_buf_index = 0;
+	cstate->raw_buf_len = nbytes;
+
+	if (cstate->need_transcoding)
+	{
+		for (;;)
+		{
+			/* If we still have a good amount of unconverted data left, convert it. */
+			nbytes = cstate->conversion_buf_len - cstate->conversion_buf_index;
+			if (nbytes >= MAX_CONVERSION_GROWTH)
+			{
+				CopyConvertBuf(cstate);
+				return true;
+			}
+
+			/* Load more raw bytes to the conversion buffer */
+			if (nbytes > 0 && cstate->conversion_buf_index > 0)
+			{
+				memmove(cstate->conversion_buf, cstate->conversion_buf + cstate->conversion_buf_index,
+						nbytes);
+			}
+			cstate->conversion_buf_index = 0;
+			cstate->conversion_buf_len = nbytes;
+			inbytes = CopyGetData(cstate, cstate->conversion_buf + cstate->conversion_buf_len,
+								  1, CONVERSION_BUF_SIZE - cstate->conversion_buf_len);
+			cstate->conversion_buf_len += inbytes;
+
+			if (inbytes == 0)
+			{
+				/* Hit EOF. If we have any unconverted bytes left, convert them now */
+				if (cstate->conversion_buf_index < cstate->conversion_buf_len)
+				{
+					CopyConvertBuf(cstate);
+					return true;
+				}
+
+				/* truly hit EOF */
+				cstate->valid_raw_buf_len = 0;
+				return false;
+			}
+		}
+	}
+	else
+	{
+		/*
+		 * No encoding conversion required. But we still need to verify that the input is
+		 * valid.
+		 *
+		 * XXX: for single-byte encoding, the verification only needs to check that the
+		 * input doesn't contain any zero bytes. Could we skip that altogether?
+		 */
+		int			validbytes;
+
+		inbytes = CopyGetData(cstate, cstate->raw_buf + nbytes,
+							  1, RAW_BUF_SIZE - nbytes);
+		nbytes += inbytes;
+		cstate->raw_buf[nbytes] = '\0';
+		cstate->raw_buf_len = nbytes;
+
+		validbytes = pg_encoding_verifymbstr(cstate->file_encoding, cstate->raw_buf, nbytes);
+		if (validbytes == 0 && nbytes > 0)
+		{
+			report_invalid_encoding(cstate->file_encoding, cstate->raw_buf, nbytes);
+		}
+
+		cstate->valid_raw_buf_len = validbytes;
+	}
+
+	return (inbytes > 0);
+}
+
+static bool
+CopyLoadRawBufBinary(CopyFromState cstate)
 {
 	int			nbytes = RAW_BUF_BYTES(cstate);
 	int			inbytes;
@@ -384,6 +527,7 @@ CopyLoadRawBuf(CopyFromState cstate)
 	cstate->raw_buf[nbytes] = '\0';
 	cstate->raw_buf_index = 0;
 	cstate->raw_buf_len = nbytes;
+
 	return (inbytes > 0);
 }
 
@@ -419,7 +563,7 @@ CopyReadBinaryData(CopyFromState cstate, char *dest, int nbytes)
 			/* Load more data if buffer is empty. */
 			if (RAW_BUF_BYTES(cstate) == 0)
 			{
-				if (!CopyLoadRawBuf(cstate))
+				if (!CopyLoadRawBufBinary(cstate))
 					break;		/* EOF */
 			}
 
@@ -695,9 +839,6 @@ CopyReadLine(CopyFromState cstate)
 	resetStringInfo(&cstate->line_buf);
 	cstate->line_buf_valid = true;
 
-	/* Mark that encoding conversion hasn't occurred yet */
-	cstate->line_buf_converted = false;
-
 	/* Parse data and transfer into line_buf */
 	result = CopyReadLineText(cstate);
 
@@ -710,10 +851,13 @@ CopyReadLine(CopyFromState cstate)
 		 */
 		if (cstate->copy_src == COPY_NEW_FE)
 		{
+			int			inbytes;
+
 			do
 			{
-				cstate->raw_buf_index = cstate->raw_buf_len;
-			} while (CopyLoadRawBuf(cstate));
+				inbytes = CopyGetData(cstate, cstate->raw_buf,
+									  1, RAW_BUF_SIZE);
+			} while (inbytes > 0);
 		}
 	}
 	else
@@ -750,26 +894,6 @@ CopyReadLine(CopyFromState cstate)
 		}
 	}
 
-	/* Done reading the line.  Convert it to server encoding. */
-	if (cstate->need_transcoding)
-	{
-		char	   *cvt;
-
-		cvt = pg_any_to_server(cstate->line_buf.data,
-							   cstate->line_buf.len,
-							   cstate->file_encoding);
-		if (cvt != cstate->line_buf.data)
-		{
-			/* transfer converted data back to line_buf */
-			resetStringInfo(&cstate->line_buf);
-			appendBinaryStringInfo(&cstate->line_buf, cvt, strlen(cvt));
-			pfree(cvt);
-		}
-	}
-
-	/* Now it's safe to use the buffer in error messages */
-	cstate->line_buf_converted = true;
-
 	return result;
 }
 
@@ -785,7 +909,6 @@ CopyReadLineText(CopyFromState cstate)
 	bool		need_data = false;
 	bool		hit_eof = false;
 	bool		result = false;
-	char		mblen_str[2];
 
 	/* CSV variables */
 	bool		first_char_in_line = true;
@@ -803,8 +926,6 @@ CopyReadLineText(CopyFromState cstate)
 			escapec = '\0';
 	}
 
-	mblen_str[1] = '\0';
-
 	/*
 	 * The objective of this loop is to transfer the entire next input line
 	 * into line_buf.  Hence, we only care for detecting newlines (\r and/or
@@ -828,7 +949,7 @@ CopyReadLineText(CopyFromState cstate)
 	 */
 	copy_raw_buf = cstate->raw_buf;
 	raw_buf_ptr = cstate->raw_buf_index;
-	copy_buf_len = cstate->raw_buf_len;
+	copy_buf_len = cstate->valid_raw_buf_len;
 
 	for (;;)
 	{
@@ -853,10 +974,10 @@ CopyReadLineText(CopyFromState cstate)
 			 * Try to read some more data.  This will certainly reset
 			 * raw_buf_index to zero, and raw_buf_ptr must go with it.
 			 */
-			if (!CopyLoadRawBuf(cstate))
+			if (!CopyLoadRawBufText(cstate))
 				hit_eof = true;
 			raw_buf_ptr = 0;
-			copy_buf_len = cstate->raw_buf_len;
+			copy_buf_len = cstate->valid_raw_buf_len;
 
 			/*
 			 * If we are completely out of data, break out of the loop,
@@ -1102,30 +1223,6 @@ CopyReadLineText(CopyFromState cstate)
 		 * value, while in non-CSV mode, \. cannot be a data value.
 		 */
 not_end_of_copy:
-
-		/*
-		 * Process all bytes of a multi-byte character as a group.
-		 *
-		 * We only support multi-byte sequences where the first byte has the
-		 * high-bit set, so as an optimization we can avoid this block
-		 * entirely if it is not set.
-		 */
-		if (cstate->encoding_embeds_ascii && IS_HIGHBIT_SET(c))
-		{
-			int			mblen;
-
-			/*
-			 * It is enough to look at the first byte in all our encodings, to
-			 * get the length.  (GB18030 is a bit special, but still works for
-			 * our purposes; see comment in pg_gb18030_mblen())
-			 */
-			mblen_str[0] = c;
-			mblen = pg_encoding_mblen(cstate->file_encoding, mblen_str);
-
-			IF_NEED_REFILL_AND_NOT_EOF_CONTINUE(mblen - 1);
-			IF_NEED_REFILL_AND_EOF_BREAK(mblen - 1);
-			raw_buf_ptr += mblen - 1;
-		}
 		first_char_in_line = false;
 	}							/* end of outer loop */
 
diff --git a/src/backend/utils/mb/mbutils.c b/src/backend/utils/mb/mbutils.c
index a585e3a6f1e..8c8b56cc2c9 100644
--- a/src/backend/utils/mb/mbutils.c
+++ b/src/backend/utils/mb/mbutils.c
@@ -436,6 +436,61 @@ pg_do_encoding_conversion(unsigned char *src, int len,
 	return result;
 }
 
+/*
+ * Convert src string to another encoding.
+ *
+ * This function has a different API than the other conversion functions.
+ * The caller should've looked up the conversion function using
+ * FindDefaultConversionProc(). Unlike the other functions, the converted
+ * result is not palloc'd. It is written to a caller-supplied buffer instead.
+ *
+ * src_encoding   - encoding to convert from
+ * dest_encoding  - encoding to convert to
+ * src, srclen    - input buffer and its length in bytes
+ * dest, destlen  - destination buffer and its size in bytes
+ *
+ * The output is null-terminated.
+ *
+ * If destlen < srclen * MAX_CONVERSION_LENGTH + 1, the converted output
+ * wouldn't necessarily fit in the output buffer, and the function will not
+ * convert the whole input.
+ *
+ * TODO: It would be nice to also return the number of bytes written to the
+ * caller, to avoid a call to strlen().
+ */
+int
+pg_do_encoding_conversion_buf(Oid proc,
+							  int src_encoding,
+							  int dest_encoding,
+							  unsigned char *src, int srclen,
+							  unsigned char *dest, int destlen,
+							  bool noError)
+{
+	Datum           result;
+
+	/*
+	 * If the destination buffer is not large enough to hold the
+	 * result in the worst case, limit the input size passed to
+	 * the conversion function.
+	 *
+	 * TODO: It would perhaps be more efficient to pass the destination
+	 * buffer size to the conversion function, so that if the conversion
+	 * expands less than the worst case, it could continue to fill up the
+	 * whole buffer.
+	 */
+	if ((Size) srclen >= ((destlen - 1) / (Size) MAX_CONVERSION_GROWTH))
+		srclen = ((destlen - 1) / (Size) MAX_CONVERSION_GROWTH);
+
+	result = OidFunctionCall6(proc,
+							  Int32GetDatum(src_encoding),
+							  Int32GetDatum(dest_encoding),
+							  CStringGetDatum(src),
+							  CStringGetDatum(dest),
+							  Int32GetDatum(srclen),
+							  BoolGetDatum(noError));
+	return DatumGetInt32(result);
+}
+
 /*
  * Convert string to encoding encoding_name. The source
  * encoding is the DB encoding.
diff --git a/src/include/commands/copyfrom_internal.h b/src/include/commands/copyfrom_internal.h
index c15ea803c32..4366aa253cd 100644
--- a/src/include/commands/copyfrom_internal.h
+++ b/src/include/commands/copyfrom_internal.h
@@ -77,7 +77,7 @@ typedef struct CopyFromStateData
 	EolType		eol_type;		/* EOL type of input */
 	int			file_encoding;	/* file or remote side's character encoding */
 	bool		need_transcoding;	/* file encoding diff from server? */
-	bool		encoding_embeds_ascii;	/* ASCII can be non-first byte? */
+	Oid			conversion_proc;
 
 	/* parameters from the COPY command */
 	Relation	rel;			/* relation to copy from */
@@ -139,23 +139,40 @@ typedef struct CopyFromStateData
 	 * line_buf is not used.)
 	 */
 	StringInfoData line_buf;
-	bool		line_buf_converted; /* converted to server encoding? */
 	bool		line_buf_valid; /* contains the row being processed? */
 
 	/*
-	 * Finally, raw_buf holds raw data read from the data source (file or
-	 * client connection).  In text mode, CopyReadLine parses this data
+	 * conversion_buf holds raw input data read from the data source (file or
+	 * client connection), not yet converted to the database encoding.
+	 *
+	 * If the encoding conversion is not required, the input data is read
+	 * directly into 'raw_buf', and conversion_buf is not used.
+	 */
+#define CONVERSION_BUF_SIZE 65536		/* we palloc CONVERSION_BUF_SIZE+1 bytes */
+	char	   *conversion_buf;
+	int			conversion_buf_index;
+	int			conversion_buf_len;
+
+	/*
+	 * raw_buf holds input data, already converted to database encoding.
+	 *
+	 * In text mode, CopyReadLine parses this data
 	 * sufficiently to locate line boundaries, then transfers the data to
-	 * line_buf and converts it.  In binary mode, CopyReadBinaryData fetches
+	 * line_buf.  In binary mode, CopyReadBinaryData fetches
 	 * appropriate amounts of data from this buffer.  In both modes, we
 	 * guarantee that there is a \0 at raw_buf[raw_buf_len].
+	 *
+	 * XXX: 'raw_buf' is a bit of a misnomer, since the data in 'conversion_buf'
+	 * is more raw than this.
 	 */
 #define RAW_BUF_SIZE 65536		/* we palloc RAW_BUF_SIZE+1 bytes */
 	char	   *raw_buf;
 	int			raw_buf_index;	/* next byte to process */
 	int			raw_buf_len;	/* total # of bytes stored */
+	int			valid_raw_buf_len;
 	/* Shorthand for number of unconsumed bytes available in raw_buf */
 #define RAW_BUF_BYTES(cstate) ((cstate)->raw_buf_len - (cstate)->raw_buf_index)
+
 } CopyFromStateData;
 
 extern void ReceiveCopyBegin(CopyFromState cstate);
diff --git a/src/include/mb/pg_wchar.h b/src/include/mb/pg_wchar.h
index 4529a15a6ba..c8f323a474a 100644
--- a/src/include/mb/pg_wchar.h
+++ b/src/include/mb/pg_wchar.h
@@ -616,6 +616,12 @@ extern int	pg_bind_textdomain_codeset(const char *domainname);
 extern unsigned char *pg_do_encoding_conversion(unsigned char *src, int len,
 												int src_encoding,
 												int dest_encoding);
+extern int pg_do_encoding_conversion_buf(Oid proc,
+										 int src_encoding,
+										 int dest_encoding,
+										 unsigned char *src, int srclen,
+										 unsigned char *dst, int dstlen,
+										 bool noError);
 
 extern char *pg_client_to_server(const char *s, int len);
 extern char *pg_server_to_client(const char *s, int len);
-- 
2.20.1

