HPCC-25410 Prevent EXTEND exceeding workunit limit

Signed-off-by: Jake Smith <jake.smith@lexisnexisrisk.com>

ecl/hthor/hthor.cpp

@@ -6293,6 +6293,19 @@ CHThorWorkUnitWriteActivity::CHThorWorkUnitWriteActivity(IAgentContext &_agent,
 {
 }
 
+static void throwWuResultTooLarge(size32_t outputLimit, IHThorWorkUnitWriteArg &helper)
+{
+    StringBuffer errMsg("Dataset too large to output to workunit (limit "); 
+    errMsg.append(outputLimit/0x100000).append(" megabytes), in result (");
+    const char *name = helper.queryName();
+    if (name)
+        errMsg.append("name=").append(name);
+    else
+        errMsg.append("sequence=").append(helper.getSequence());
+    errMsg.append(")");
+    throw MakeStringExceptionDirect(0, errMsg.str());
+}
+
 void CHThorWorkUnitWriteActivity::execute()
 {
     unsigned flags = helper.getFlags();
@@ -6350,18 +6363,8 @@ void CHThorWorkUnitWriteActivity::execute()
                 break;
         }
         size32_t thisSize = inputMeta->getRecordSize(nextrec);
-        if(outputLimit && ((rowdata.length() + thisSize) > outputLimit))
-        {
-            StringBuffer errMsg("Dataset too large to output to workunit (limit "); 
-            errMsg.append(outputLimit/0x100000).append(" megabytes), in result (");
-            const char *name = helper.queryName();
-            if (name)
-                errMsg.append("name=").append(name);
-            else
-                errMsg.append("sequence=").append(helper.getSequence());
-            errMsg.append(")");
-            throw MakeStringExceptionDirect(0, errMsg.str());
-         }
+        if (outputLimit && ((rowdata.length() + thisSize) > outputLimit))
+            throwWuResultTooLarge(outputLimit, helper);
         if (rowSerializer)
         {
             CThorDemoRowSerializer serializerTarget(rowdata);
@@ -6392,7 +6395,12 @@ void CHThorWorkUnitWriteActivity::execute()
     WorkunitUpdate w = agent.updateWorkUnit();
     Owned<IWUResult> result = updateWorkUnitResult(w, helper.queryName(), helper.getSequence());
     if (0 != (POFextend & helper.getFlags()))
+    {
+        __int64 existingSz = result->getResultRawSize(nullptr, nullptr);
+        if (outputLimit && ((rowdata.length() + existingSz) > outputLimit))
+            throwWuResultTooLarge(outputLimit, helper);
         result->addResultRaw(rowdata.length(), rowdata.toByteArray(), ResultFormatRaw);
+    }
     else
         result->setResultRaw(rowdata.length(), rowdata.toByteArray(), ResultFormatRaw);
     result->setResultStatus(ResultStatusCalculated);

thorlcr/activities/wuidwrite/thwuidwrite.cpp

@@ -38,12 +38,28 @@ protected:
     StringAttr resultName;
     unsigned resultSeq;
 
+    void throwWuResultTooLarge()
+    {
+        StringBuffer errMsg("Dataset too large to output to workunit (limit is set to ");
+        errMsg.append(workunitWriteLimit/0x100000).append(") megabytes, in result (");
+        if (resultName.length())
+            errMsg.append("name=").append(resultName);
+        else
+            errMsg.append("sequence=").append(resultSeq);
+        errMsg.append(")");
+        throw MakeThorException(TE_WorkUnitWriteLimitExceeded, "%s", errMsg.str());
+    }
     void addResult(rowcount_t resultCount, MemoryBuffer &resultData, bool complete)
     {
         Owned<IWorkUnit> wu = &container.queryJob().queryWorkUnit().lock();
         Owned<IWUResult> result = updateWorkUnitResult(wu, resultName, resultSeq);
         if (appendOutput)
+        {
+            __int64 existingSz = result->getResultRawSize(nullptr, nullptr);
+            if (workunitWriteLimit && (existingSz+resultData.length() > workunitWriteLimit))
+                throwWuResultTooLarge();
             result->addResultRaw(resultData.length(), resultData.toByteArray(), ResultFormatRaw);
+        }
         else
             result->setResultRaw(resultData.length(), resultData.toByteArray(), ResultFormatRaw);
         result->setResultRowCount(resultCount);
@@ -121,17 +137,8 @@ public:
                 unsigned numGot;
                 mb.read(numGot);
                 unsigned l=mb.remaining();
-                if (workunitWriteLimit && totalSize+resultData.length()+l > workunitWriteLimit)
-                {
-                    StringBuffer errMsg("Dataset too large to output to workunit (limit is set to ");
-                    errMsg.append(workunitWriteLimit/0x100000).append(") megabytes, in result (");
-                    if (resultName.length())
-                        errMsg.append("name=").append(resultName);
-                    else
-                        errMsg.append("sequence=").append(resultSeq);
-                    errMsg.append(")");
-                    throw MakeThorException(TE_WorkUnitWriteLimitExceeded, "%s", errMsg.str());
-                }
+                if (workunitWriteLimit && (totalSize+resultData.length()+l > workunitWriteLimit))
+                    throwWuResultTooLarge();
                 resultData.append(l, mb.readDirect(l));
                 mb.clear();
                 numResults += numGot;