From 27cf95bdddf8d9c475e9b59e4b1195a93e91e4e7 Mon Sep 17 00:00:00 2001
From: Juliusz Sompolski <julek@databricks.com>
Date: Wed, 21 Apr 2021 13:53:38 +0200
Subject: [PATCH] heap message if stack message is too small

---
 src/cursor.cpp |  2 +-
 src/errors.cpp | 37 ++++++++++++++++++++++++++++++++++---
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/cursor.cpp b/src/cursor.cpp
index 23afd0f1..cbebf476 100644
--- a/src/cursor.cpp
+++ b/src/cursor.cpp
@@ -574,7 +574,7 @@ static int GetDiagRecs(Cursor* cur)
 
     ODBCCHAR    cSQLState[6];  // five-character SQLSTATE code (plus terminating NULL)
     SQLINTEGER  iNativeError;
-    ODBCCHAR    cMessageText[SHRT_MAX];  // PRINT statements can be large
+    ODBCCHAR    cMessageText[10240];  // PRINT statements can be large
     SQLSMALLINT iTextLength;
 
     SQLRETURN ret;
diff --git a/src/errors.cpp b/src/errors.cpp
index 304cad27..1870cd81 100644
--- a/src/errors.cpp
+++ b/src/errors.cpp
@@ -217,7 +217,9 @@ PyObject* GetErrorFromHandle(Connection *conn, const char* szFunction, HDBC hdbc
     SQLSMALLINT cchMsg;
 
     ODBCCHAR sqlstateT[6];
-    ODBCCHAR szMsg[SHRT_MAX];
+    ODBCCHAR stackMsg[1024];
+    ODBCCHAR *szMsg;
+
 
     if (hstmt != SQL_NULL_HANDLE)
     {
@@ -251,11 +253,28 @@ PyObject* GetErrorFromHandle(Connection *conn, const char* szFunction, HDBC hdbc
 
         SQLRETURN ret;
         Py_BEGIN_ALLOW_THREADS
-        ret = SQLGetDiagRecW(nHandleType, h, iRecord, (SQLWCHAR*)sqlstateT, &nNativeError, (SQLWCHAR*)szMsg, (short)(_countof(szMsg)-1), &cchMsg);
+        ret = SQLGetDiagRecW(nHandleType, h, iRecord, (SQLWCHAR*)sqlstateT, &nNativeError, (SQLWCHAR*)szMsg, (short)(_countof(stackMsg)-1), &cchMsg);
         Py_END_ALLOW_THREADS
         if (!SQL_SUCCEEDED(ret))
             break;
 
+        // If needed, allocate a bigger error message buffer and retry.
+        if (cchMsg > msgLen - 1) {
+            SQLSMALLINT msgLen = cchMsg + 1;
+            szMsg = pyodbc_malloc(msgLen * sizeof(ODBCCHAR));
+            if (!heapMsg) {
+                PyErr_NoMemory();
+                return 0;
+            }
+            Py_BEGIN_ALLOW_THREADS
+            ret = SQLGetDiagRecW(nHandleType, h, iRecord, (SQLWCHAR*)sqlstateT, &nNativeError, (SQLWCHAR*)szMsg, msgLen, &cchMsg);
+            Py_END_ALLOW_THREADS
+            if (!SQL_SUCCEEDED(ret)) {
+                pyodbc_free(szMsg);
+                break;
+            }
+        }
+
         // Not always NULL terminated (MS Access)
         sqlstateT[5] = 0;
 
@@ -264,6 +283,13 @@ PyObject* GetErrorFromHandle(Connection *conn, const char* szFunction, HDBC hdbc
         const char *unicode_enc = conn ? conn->metadata_enc.name : ENCSTR_UTF16NE;
         Object msgStr(PyUnicode_Decode((char*)szMsg, cchMsg * sizeof(ODBCCHAR), unicode_enc, "strict"));
 
+        if (szMsg != stackMsg) {
+           // Free and revert back to stackMsg here, reduces the complexity of handling keeping
+           // track of the heap message and freeing it in other places.
+           pyodbc_free(szMsg);
+           stackMsg = szMsg;
+        }
+
         if (cchMsg != 0 && msgStr.Get())
         {
             if (iRecord == 1)
@@ -272,8 +298,10 @@ PyObject* GetErrorFromHandle(Connection *conn, const char* szFunction, HDBC hdbc
                 // exception class and append the calling function name.
                 CopySqlState(sqlstateT, sqlstate);
                 msg = PyUnicode_FromFormat("[%s] %V (%ld) (%s)", sqlstate, msgStr.Get(), "(null)", (long)nNativeError, szFunction);
-                if (!msg)
+                if (!msg) {
+                    PyErr_NoMemory();
                     return 0;
+                }
             }
             else
             {
@@ -298,6 +326,9 @@ PyObject* GetErrorFromHandle(Connection *conn, const char* szFunction, HDBC hdbc
 #endif
     }
 
+    // Free raw message buffer, if allocated on heap.
+    if (szMsg != stackMsg) pyodbc_free(szMsg);
+
     if (!msg || PyUnicode_GetSize(msg.Get()) == 0)
     {
         // This only happens using unixODBC.  (Haven't tried iODBC yet.)  Either the driver or the driver manager is