Implement query_timeout on socket read for rlm_postgresql

diff --git a/src/modules/rlm_sql/drivers/rlm_sql_postgresql/rlm_sql_postgresql.c b/src/modules/rlm_sql/drivers/rlm_sql_postgresql/rlm_sql_postgresql.c
index eb1fed3b542acac76fba36113d2cabb4d396b911..14b2e29dbf584c08dbe9a23f59b56a4475f81bd1 100644 (file)
--- a/src/modules/rlm_sql/drivers/rlm_sql_postgresql/rlm_sql_postgresql.c
+++ b/src/modules/rlm_sql/drivers/rlm_sql_postgresql/rlm_sql_postgresql.c
@@ -259,6 +259,10 @@ static CC_HINT(nonnull) sql_rcode_t sql_query(rlm_sql_handle_t *handle, rlm_sql_
 {
        rlm_sql_postgres_conn_t *conn = handle->conn;
        rlm_sql_postgres_t      *inst = config->driver;
+       struct timeval          timeout = {config->query_timeout, 0};
+       int                     sockfd;
+       fd_set                  read_fd;
+       PGresult                *tmp_result;
        int                     numfields = 0;
        ExecStatusType          status;
 
@@ -267,6 +271,34 @@ static CC_HINT(nonnull) sql_rcode_t sql_query(rlm_sql_handle_t *handle, rlm_sql_
                return RLM_SQL_RECONNECT;
        }
 
+       sockfd = PQsocket(conn->db);
+       if (sockfd < 0) {
+               ERROR("Unable to obtain socket: %s", PQerrorMessage(conn->db));
+               return RLM_SQL_RECONNECT;
+       }
+
+       if (!PQsendQuery(conn->db, query)) {
+               ERROR("Failed to send query: %s", PQerrorMessage(conn->db));
+               return RLM_SQL_RECONNECT;
+       }
+
+       /*
+        *  We try to avoid blocking by waiting until the driver indicates that
+        *  the result is ready or our timeout expires
+        */
+       while (PQisBusy(conn->db)) {
+               FD_ZERO(&read_fd);
+               FD_SET(sockfd, &read_fd);
+               if (!select(sockfd + 1, &read_fd, NULL, NULL, config->query_timeout ? &timeout : NULL)) {
+                       ERROR("Socket read timeout after %d seconds", config->query_timeout);
+                       return RLM_SQL_RECONNECT;
+               }
+               if (!PQconsumeInput(conn->db)) {
+                       ERROR("Failed reading input: %s", PQerrorMessage(conn->db));
+                       return RLM_SQL_RECONNECT;
+               }
+       }
+
        /*
         *  Returns a PGresult pointer or possibly a null pointer.
         *  A non-null pointer will generally be returned except in
@@ -275,7 +307,11 @@ static CC_HINT(nonnull) sql_rcode_t sql_query(rlm_sql_handle_t *handle, rlm_sql_
         *  returned, it should be treated like a PGRES_FATAL_ERROR
         *  result.
         */
-       conn->result = PQexec(conn->db, query);
+       conn->result = PQgetResult(conn->db);
+
+       /* Discard results for appended queries */
+       while ((tmp_result = PQgetResult(conn->db)) != NULL)
+               PQclear(tmp_result);
 
        /*
         *  As this error COULD be a connection error OR an out-of-memory