odyssey: rework packet header validation

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;
 }
 

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;

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 */