diff --git a/sources/io.h b/sources/io.h index 2e792fa1..193cf056 100644 --- a/sources/io.h +++ b/sources/io.h @@ -7,6 +7,8 @@ * Scalable PostgreSQL connection pooler. */ +#include "macro.h" + typedef struct od_io od_io_t; struct od_io @@ -181,6 +183,14 @@ 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' || \ +(id) == 'E' || (id) == 'N' || (id) == 'A') + static inline machine_msg_t* od_read(od_io_t *io, uint32_t time_ms) { @@ -192,9 +202,17 @@ od_read(od_io_t *io, uint32_t time_ms) uint32_t size; size = kiwi_read_size((char*)&header, sizeof(header)); - if (size < sizeof(uint32_t) || header.type < 0x20) { + + if ( od_unlikely( + size < sizeof(uint32_t) || + header.type < 0x20 || + (size > 30000 && !VALID_LONG_MESSAGE_TYPE(header.type))) + ) { // This is not a postgres fe 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 return NULL; } size -= sizeof(uint32_t);