From a3f55f4f219c5aac6832d107b5c0afa8c831ba81 Mon Sep 17 00:00:00 2001
From: Dmitry Simonenko <pmwkaa@gmail.com>
Date: Mon, 3 Jun 2019 13:12:50 +0300
Subject: [PATCH] odyssey: rework packet header validation

---
 sources/frontend.c         |  6 ++---
 sources/io.h               | 33 ++++++-----------------
 third_party/kiwi/kiwi/io.h | 55 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 29 deletions(-)

diff --git a/sources/frontend.c b/sources/frontend.c
index 485c45a8..f3cb0ef3 100644
--- a/sources/frontend.c
+++ b/sources/frontend.c
@@ -88,7 +88,7 @@ od_frontend_startup(od_client_t *client)
 	machine_msg_t *msg;
 	msg = od_read_startup(&client->io, UINT32_MAX);
 	if (msg == NULL)
-		return -1;
+		goto error;
 
 	int rc;
 	rc = kiwi_be_read_startup(machine_msg_data(msg),
@@ -124,9 +124,7 @@ od_frontend_startup(od_client_t *client)
 
 error:
 	od_error(&instance->logger, "startup", client, NULL,
-	         "incorrect startup packet");
-	od_frontend_error(client, KIWI_PROTOCOL_VIOLATION,
-	                  "bad startup packet");
+	         "startup packet read error");
 	return -1;
 }
 
diff --git a/sources/io.h b/sources/io.h
index a24598e4..9f15c0f3 100644
--- a/sources/io.h
+++ b/sources/io.h
@@ -185,8 +185,12 @@ od_read_startup(od_io_t *io, uint32_t time_ms)
 	if (rc == -1)
 		return NULL;
 
+	/* pre-validate startup header size, actual header parsing will be done by
+	 * kiwi_be_read_startup() */
 	uint32_t size;
-	size = kiwi_read_startup_size((char*)&header, sizeof(header));
+	rc = kiwi_validate_startup_header((char*)&header, sizeof(header), &size);
+	if (rc == -1)
+		return NULL;
 
 	machine_msg_t *msg;
 	msg = machine_msg_create(sizeof(header) + size);
@@ -207,15 +211,6 @@ od_read_startup(od_io_t *io, uint32_t time_ms)
 	return msg;
 }
 
-/*
- * This macro lists the backend message types that could be "long" (more
- * than a couple of kilobytes).
- */
-#define VALID_LONG_MESSAGE_TYPE(id) \
-	((id) == 'T' || (id) == 'D' || (id) == 'd' || (id) == 'V' || /* BE messages */\
-	(id) == 'E' || (id) == 'N' || (id) == 'A'|| /* BE messages */\
-	(id) == 'B' || (id) == 'P' || (id) == 'Q') /* FE messages */
-
 static inline machine_msg_t*
 od_read(od_io_t *io, uint32_t time_ms)
 {
@@ -225,23 +220,11 @@ od_read(od_io_t *io, uint32_t time_ms)
 	if (rc == -1)
 		return NULL;
 
+	/* pre-validate packet header */
 	uint32_t size;
-	size = kiwi_read_size((char*)&header, sizeof(header));
-
-	if ( od_unlikely(
-			size < sizeof(uint32_t) ||
-		    header.type < 0x20 ||
-		    (size > 30000 && !VALID_LONG_MESSAGE_TYPE(header.type)))
-		    ) {
-		/*
-		 * This is not a postgres protocol v3 message
-		 * We should drop connection ASAP
-		 * Validation is performed per PostgreSQL impl
-		 * For reference see
-		 * https://github.com/postgres/postgres/blob/7bac3acab4d5c3f2c35aa3a7bea08411d83fd5bc/src/interfaces/libpq/fe-protocol3.c#L91-L100
-		 */
+	rc = kiwi_validate_header((char*)&header, sizeof(header), &size);
+	if (rc == -1)
 		return NULL;
-	}
 	size -= sizeof(uint32_t);
 
 	machine_msg_t *msg;
diff --git a/third_party/kiwi/kiwi/io.h b/third_party/kiwi/kiwi/io.h
index f06b4fe7..d18d6965 100644
--- a/third_party/kiwi/kiwi/io.h
+++ b/third_party/kiwi/kiwi/io.h
@@ -164,4 +164,59 @@ kiwi_read_startup_size(char *data, uint32_t data_size)
 	return size;
 }
 
+#define KIWI_LONG_MESSAGE_SIZE 30000
+
+KIWI_API static inline int
+kiwi_validate_startup_header(char *data, uint32_t data_size, uint32_t *size)
+{
+	assert(data_size >= sizeof(uint32_t));
+	*size = kiwi_read_startup_size(data, sizeof(uint32_t));
+	/* do not expect big startup messages */
+	if (kiwi_unlikely(*size >= KIWI_LONG_MESSAGE_SIZE))
+		return -1;
+	return 0;
+}
+
+KIWI_API static inline int
+kiwi_validate_header(char *data, uint32_t data_size, uint32_t *size)
+{
+	assert(data_size >= sizeof(kiwi_header_t));
+	*size = kiwi_read_size(data, sizeof(kiwi_header_t));
+
+	kiwi_header_t *header = (kiwi_header_t*)data;
+	if (kiwi_unlikely(*size < sizeof(uint32_t)))
+		return -1;
+
+	/* check supported protocol message type */
+	if (kiwi_unlikely(header->type < 0x20))
+		return -1;
+
+	/* is small packet */
+	if (kiwi_likely(*size < KIWI_LONG_MESSAGE_SIZE))
+		return 0;
+
+	/*
+	 * Lists the backend and frontend message types that could be "long" (more
+	 * than a couple of kilobytes).
+	*/
+	switch (header->type) {
+	/* backend */
+	case KIWI_BE_ROW_DESCRIPTION:
+	case KIWI_BE_DATA_ROW:
+	case KIWI_BE_COPY_DATA:
+	case KIWI_BE_FUNCTION_CALL_RESPONSE:
+	case KIWI_BE_ERROR_RESPONSE:
+	case KIWI_BE_NOTICE_RESPONSE:
+	case KIWI_BE_NOTIFICATION_RESPONSE:
+	/* frontend */
+	case KIWI_FE_BIND:
+	case KIWI_FE_PARSE:
+	case KIWI_FE_QUERY:
+	/* KIWI_FE_COPY_DATA has same type as BE_COPY_DATA */
+		return 0;
+	}
+
+	return -1;
+}
+
 #endif /* KIWI_IO_H */