Author:     Noah Misch <noah@leadboat.com>
Commit:     Noah Misch <noah@leadboat.com>

    Validate xl_tot_len early enough.  FIXME
    
    This also allows distinguishing later OOM from an invalid record.

diff --git a/src/backend/access/transam/xlogreader.c b/src/backend/access/transam/xlogreader.c
index c9f9f6e..ec754db 100644
--- a/src/backend/access/transam/xlogreader.c
+++ b/src/backend/access/transam/xlogreader.c
@@ -201,25 +201,6 @@ allocate_recordbuf(XLogReaderState *state, uint32 reclength)
 	newSize += XLOG_BLCKSZ - (newSize % XLOG_BLCKSZ);
 	newSize = Max(newSize, 5 * Max(BLCKSZ, XLOG_BLCKSZ));
 
-#ifndef FRONTEND
-
-	/*
-	 * Note that in much unlucky circumstances, the random data read from a
-	 * recycled segment can cause this routine to be called with a size
-	 * causing a hard failure at allocation.  For a standby, this would cause
-	 * the instance to stop suddenly with a hard failure, preventing it to
-	 * retry fetching WAL from one of its sources which could allow it to move
-	 * on with replay without a manual restart. If the data comes from a past
-	 * recycled segment and is still valid, then the allocation may succeed
-	 * but record checks are going to fail so this would be short-lived.  If
-	 * the allocation fails because of a memory shortage, then this is not a
-	 * hard failure either per the guarantee given by MCXT_ALLOC_NO_OOM.
-	 */
-	if (!AllocSizeIsValid(newSize))
-		return false;
-
-#endif
-
 	if (state->readRecordBuf)
 		pfree(state->readRecordBuf);
 	state->readRecordBuf =
@@ -641,7 +622,7 @@ restart:
 	Assert(pageHeaderSize <= readOff);
 
 	/*
-	 * Read the record length.
+	 * Validate the record length.
 	 *
 	 * NB: Even though we use an XLogRecord pointer here, the whole record
 	 * header might not fit on this page. xl_tot_len is the first field of the
@@ -651,35 +632,42 @@ restart:
 	 */
 	record = (XLogRecord *) (state->readBuf + RecPtr % XLOG_BLCKSZ);
 	total_len = record->xl_tot_len;
+	if (total_len < SizeOfXLogRecord)
+	{
+		report_invalid_record(state,
+							  "invalid record length at %X/%X: expected at least %u, got %u",
+							  LSN_FORMAT_ARGS(RecPtr),
+							  (uint32) SizeOfXLogRecord, total_len);
+		goto err;
+	}
+#ifndef FRONTEND
 
 	/*
-	 * If the whole record header is on this page, validate it immediately.
-	 * Otherwise do just a basic sanity check on xl_tot_len, and validate the
-	 * rest of the header after reading it from the next page.  The xl_tot_len
-	 * check is necessary here to ensure that we enter the "Need to reassemble
-	 * record" code path below; otherwise we might fail to apply
-	 * ValidXLogRecordHeader at all.
+	 * We may be looking at a random uint32 read from a recycled segment.  For
+	 * a standby, a hard allocation failure would cause the instance to stop
+	 * suddenly, preventing it to retry fetching WAL from one of its sources
+	 * which could allow it to move on with replay without a manual restart.
+	 * If the data comes from a past recycled segment and is still valid, then
+	 * the allocation may succeed but record checks are going to fail so this
+	 * would be short-lived.
 	 */
-	if (targetRecOff <= XLOG_BLCKSZ - SizeOfXLogRecord)
-	{
-		if (!ValidXLogRecordHeader(state, RecPtr, state->DecodeRecPtr, record,
-								   randAccess))
-			goto err;
-		gotheader = true;
-	}
-	else
+	if (!AllocSizeIsValid(DecodeXLogRecordRequiredSpace(total_len)))
 	{
-		/* XXX: more validation should be done here */
-		if (total_len < SizeOfXLogRecord)
-		{
-			report_invalid_record(state,
-								  "invalid record length at %X/%X: expected at least %u, got %u",
-								  LSN_FORMAT_ARGS(RecPtr),
-								  (uint32) SizeOfXLogRecord, total_len);
-			goto err;
-		}
-		gotheader = false;
+		report_invalid_record(state, "record length %u at %X/%X too long",
+							  total_len, LSN_FORMAT_ARGS(RecPtr));
+		goto err;
 	}
+#endif
+
+	/*
+	 * If the whole record header is on this page, validate it immediately.
+	 * Otherwise, validate it after reading it from the next page.
+	 */
+	gotheader = targetRecOff <= XLOG_BLCKSZ - SizeOfXLogRecord ? true : false;
+	if (gotheader &&
+		!ValidXLogRecordHeader(state, RecPtr, state->DecodeRecPtr, record,
+							   randAccess))
+		goto err;
 
 	/*
 	 * Find space to decode this record.  Don't allow oversized allocation if
@@ -723,9 +711,9 @@ restart:
 		if (total_len > state->readRecordBufSize &&
 			!allocate_recordbuf(state, total_len))
 		{
-			/* We treat this as a "bogus data" condition */
-			report_invalid_record(state, "record length %u at %X/%X too long",
-								  total_len, LSN_FORMAT_ARGS(RecPtr));
+			report_invalid_record(state,
+								  "out of memory while trying to decode a record of length %u",
+								  total_len);
 			goto err;
 		}
 
@@ -1113,14 +1101,6 @@ ValidXLogRecordHeader(XLogReaderState *state, XLogRecPtr RecPtr,
 					  XLogRecPtr PrevRecPtr, XLogRecord *record,
 					  bool randAccess)
 {
-	if (record->xl_tot_len < SizeOfXLogRecord)
-	{
-		report_invalid_record(state,
-							  "invalid record length at %X/%X: expected at least %u, got %u",
-							  LSN_FORMAT_ARGS(RecPtr),
-							  (uint32) SizeOfXLogRecord, record->xl_tot_len);
-		return false;
-	}
 	if (!RmgrIdIsValid(record->xl_rmid))
 	{
 		report_invalid_record(state,
@@ -1169,10 +1149,9 @@ ValidXLogRecordHeader(XLogReaderState *state, XLogRecPtr RecPtr,
  * record (other than to the minimal extent of computing the amount of
  * data to read in) until we've checked the CRCs.
  *
- * We assume all of the record (that is, xl_tot_len bytes) has been read
- * into memory at *record.  Also, ValidXLogRecordHeader() has accepted the
- * record's header, which means in particular that xl_tot_len is at least
- * SizeOfXLogRecord.
+ * We assume all of the record (that is, xl_tot_len bytes) has been read into
+ * memory at *record.  Also, xl_tot_len is at least SizeOfXLogRecord, and
+ * ValidXLogRecordHeader() has passed.
  */
 static bool
 ValidXLogRecord(XLogReaderState *state, XLogRecord *record, XLogRecPtr recptr)
