diff --git src/interfaces/libpq/fe-exec.c src/interfaces/libpq/fe-exec.c index 2265ab5..56883ec 100644 --- src/interfaces/libpq/fe-exec.c +++ src/interfaces/libpq/fe-exec.c @@ -4484,3 +4484,9 @@ PQgetResultFromPGconn(PGconn *conn) return conn->result; } +int PShandleRowData(PGconn *conn, bool is_first_packet, PSresult* result) { + if (!conn || !result) + return 1; + return psHandleRowData(conn, is_first_packet, result); +} + diff --git src/interfaces/libpq/fe-misc.c src/interfaces/libpq/fe-misc.c index 488f7d6..65beb87 100644 --- src/interfaces/libpq/fe-misc.c +++ src/interfaces/libpq/fe-misc.c @@ -1344,3 +1344,39 @@ libpq_append_conn_error(PGconn *conn, const char *fmt,...) appendPQExpBufferChar(&conn->errorMessage, '\n'); } + +/* + * psGetInt16 + * read 2 byte integer and convert from network byte order + * to local byte order + */ +int +psGetInt16(int *result, PGconn *conn) +{ + uint16 tmp2; + + if (conn->inCursor + 2 > conn->inEnd) + return EOF; + memcpy(&tmp2, conn->inBuffer + conn->inCursor, 2); + conn->inCursor += 2; + *result = (int) pg_ntoh16(tmp2); + return 0; +} + +/* + * psGetInt32 + * read 4 byte integer and convert from network byte order + * to local byte order + */ +int +psGetInt32(int *result, PGconn *conn) +{ + uint32 tmp4; + + if (conn->inCursor + 4 > conn->inEnd) + return EOF; + memcpy(&tmp4, conn->inBuffer + conn->inCursor, 4); + conn->inCursor += 4; + *result = (int) pg_ntoh32(tmp4); + return 0; +} diff --git src/interfaces/libpq/fe-protocol3.c src/interfaces/libpq/fe-protocol3.c index 9c4aa7e..de0746c 100644 --- src/interfaces/libpq/fe-protocol3.c +++ src/interfaces/libpq/fe-protocol3.c @@ -2299,3 +2299,109 @@ build_startup_packet(const PGconn *conn, char *packet, return packet_len; } + +/* + * psHandleRowData: Processes the incoming message from the PostgreSQL backend. + * This function checks if the message contains row data (indicated by 'D') + * and processes it accordingly. It validates the message type and length, + * ensures that the complete message has been received, and updates the result + * structure if the message contains row data. + * + * Return values: + * 0 -> Message processed successfully (row data handled). + * 1 -> Message not fully processed; the next call should be to PQisBusy. + * -1 -> Not enough data to process the message; the next call should be to PQconsumeInput. + */ +int +psHandleRowData(PGconn *conn, bool isFirstPacket, PSresult* result) +{ + char id; + int msgLength; + int avail; + + if (conn->asyncStatus != PGASYNC_BUSY) + return 1; + /* + * Try to read a message. First get the type code and length. Return + * if not enough data. + */ + conn->inCursor = conn->inStart; + if (pqGetc(&id, conn)) + return EOF; + + if (id != 'D') + return 1; + + if (psGetInt32(&msgLength, conn)) + return EOF; + + /* + * Try to validate message type/length here. A length less than 4 is + * definitely broken. Large lengths should only be believed for a few + * message types. + */ + if (msgLength < 4) + return 1; + + if (msgLength > 30000 && !VALID_LONG_MESSAGE_TYPE(id)) + return 1; + + /* + * Can't process if message body isn't all here yet. + */ + msgLength -= 4; + avail = conn->inEnd - conn->inCursor; + if (avail < msgLength) { + if ((conn->inCursor+(size_t)msgLength) <= (size_t)conn->inBufSize) + return EOF; + return 1; + } + + /* First data row should be skipped since it is part of PGresult, which contains row description */ + if (isFirstPacket) + return 1; + + if (conn->result != NULL && + conn->result->resultStatus == PGRES_TUPLES_OK) + { + PGresult *res = conn->result; + int nfields = res->numAttributes; + int tupnfields; /* # fields from tuple */ + int vlen; /* length of the current field value */ + int i; + + + /* Get the field count and make sure it's what we expect */ + if (psGetInt16(&tupnfields, conn)) + return 1; + + if (tupnfields != nfields) + return 1; + + /* Scan the fields */ + for (i = 0; i < nfields; i++) + { + /* get the value length */ + if (psGetInt32(&vlen, conn)) + return 1; + + /* Skip over the data value */ + if (vlen > 0) + { + if (pqSkipnchar(vlen, conn)) + return 1; + } + } + + result->id = 'D'; + result->len = msgLength + 5; + result->data = conn->inBuffer + conn->inStart; + conn->asyncStatus = PGASYNC_BUSY; + /* trust the specified message length as what to skip */ + conn->inStart += 5 + msgLength; + conn->inCursor = conn->inStart; + return 0; + } + return 1; +} + diff --git src/interfaces/libpq/libpq-fe.h src/interfaces/libpq/libpq-fe.h index c5170d1..3e3cc34 100644 --- src/interfaces/libpq/libpq-fe.h +++ src/interfaces/libpq/libpq-fe.h @@ -269,6 +269,18 @@ typedef struct pgresAttDesc int atttypmod; /* type-specific modifier info */ } PGresAttDesc; +/* ---------------- + * PSresult -- + * ---------------- + */ +typedef struct psResult +{ + char id; + int len; + const char* data; + //int fieldcount; +} PSresult; + /* ---------------- * Exported functions of libpq * ---------------- @@ -671,6 +683,9 @@ extern int PQdefaultSSLKeyPassHook_OpenSSL(char *buf, int size, PGconn *conn); /* Get PGresult directly from PGconn. WARNING: DO NOT RELEASE THIS RESULT */ extern const PGresult *PQgetResultFromPGconn(PGconn *conn); +/* ProxySQL special handler function */ +extern int PShandleRowData(PGconn *conn, bool is_first_packet, PSresult* result); + #ifdef __cplusplus } #endif diff --git src/interfaces/libpq/libpq-int.h src/interfaces/libpq/libpq-int.h index a951f49..e1df8b5 100644 --- src/interfaces/libpq/libpq-int.h +++ src/interfaces/libpq/libpq-int.h @@ -727,6 +727,11 @@ extern PGresult *pqFunctionCall3(PGconn *conn, Oid fnid, int result_is_int, const PQArgBlock *args, int nargs); + /* + * ProxySQL light weight routines + */ +extern int psHandleRowData(PGconn *conn, bool is_first_packet, PSresult* result); + /* === in fe-misc.c === */ /* @@ -756,6 +761,13 @@ extern int pqWaitTimed(int forRead, int forWrite, PGconn *conn, extern int pqReadReady(PGconn *conn); extern int pqWriteReady(PGconn *conn); + /* + * ProxySQL light weight routines + */ +extern int psGetInt16(int *result, PGconn *conn); +extern int psGetInt32(int *result, PGconn *conn); + + /* === in fe-secure.c === */ extern int pqsecure_initialize(PGconn *, bool, bool);