[b26e38a3e4] Ensure that compressing streams manage buffers correctly in all cases.

2 files changed, 105 insertions, 58 deletions

diff --git a/generic/tclZlib.c b/generic/tclZlib.c
index a7e8a8a..fd8b3a1 100644
--- a/generic/tclZlib.c
+++ b/generic/tclZlib.c
@@ -177,6 +177,8 @@ static Tcl_ObjCmdProc		ZlibStreamPutCmd;
 static void		ConvertError(Tcl_Interp *interp, int code,
 			    uLong adler);
 static Tcl_Obj *	ConvertErrorToList(int code, uLong adler);
+static inline int	Deflate(z_streamp strm, void *bufferPtr,
+			    int bufferSize, int flush, int *writtenPtr);
 static void		ExtractHeader(gz_header *headerPtr, Tcl_Obj *dictObj);
 static int		GenerateHeader(Tcl_Interp *interp, Tcl_Obj *dictObj,
 			    GzipHeader *headerPtr, int *extraSizePtr);
@@ -578,6 +580,10 @@ ExtractHeader(
     }
 }
 
+/*
+ * Disentangle the worst of how the zlib API is used.
+ */
+
 static int
 SetInflateDictionary(
     z_streamp strm,
@@ -605,6 +611,38 @@ SetDeflateDictionary(
     }
     return Z_OK;
 }
+
+static inline int
+Deflate(
+    z_streamp strm,
+    void *bufferPtr,
+    int bufferSize,
+    int flush,
+    int *writtenPtr)
+{
+    int e;
+
+    strm->next_out = (Bytef *) bufferPtr;
+    strm->avail_out = (unsigned) bufferSize;
+    e = deflate(strm, flush);
+    if (writtenPtr != NULL) {
+	*writtenPtr = bufferSize - strm->avail_out;
+    }
+    return e;
+}
+
+static inline void
+AppendByteArray(
+    Tcl_Obj *listObj,
+    void *buffer,
+    int size)
+{
+    if (size > 0) {
+	Tcl_Obj *baObj = Tcl_NewByteArrayObj((unsigned char *) buffer, size);
+
+	Tcl_ListObjAppendElement(NULL, listObj, baObj);
+    }
+}
 
 /*
  *----------------------------------------------------------------------
@@ -1139,6 +1177,8 @@ Tcl_ZlibStreamSetCompressionDictionary(
  *----------------------------------------------------------------------
  */
 
+#define BUFFER_SIZE_LIMIT	0xFFFF
+
 int
 Tcl_ZlibStreamPut(
     Tcl_ZlibStream zshandle,	/* As obtained from Tcl_ZlibStreamInit */
@@ -1148,8 +1188,7 @@ Tcl_ZlibStreamPut(
 {
     ZlibStreamHandle *zshPtr = (ZlibStreamHandle *) zshandle;
     char *dataTmp = NULL;
-    int e, size, outSize;
-    Tcl_Obj *obj;
+    int e, size, outSize, toStore;
 
     if (zshPtr->streamEnd) {
 	if (zshPtr->interp) {
@@ -1175,26 +1214,45 @@ Tcl_ZlibStreamPut(
 	if (HaveDictToSet(zshPtr)) {
 	    e = SetDeflateDictionary(&zshPtr->stream, zshPtr->compDictObj);
 	    if (e != Z_OK) {
-		if (zshPtr->interp) {
-		    ConvertError(zshPtr->interp, e, zshPtr->stream.adler);
-		}
+		ConvertError(zshPtr->interp, e, zshPtr->stream.adler);
 		return TCL_ERROR;
 	    }
 	    DictWasSet(zshPtr);
 	}
 
 	/*
-	 * Deflatebound doesn't seem to take various header sizes into
-	 * account, so we add 100 extra bytes.
+	 * deflateBound() doesn't seem to take various header sizes into
+	 * account, so we add 100 extra bytes. However, we can also loop
+	 * around again so we also set an upper bound on the output buffer
+	 * size.
 	 */
 
-	outSize = deflateBound(&zshPtr->stream, zshPtr->stream.avail_in)+100;
-	zshPtr->stream.avail_out = outSize;
-	dataTmp = ckalloc(zshPtr->stream.avail_out);
-	zshPtr->stream.next_out = (Bytef *) dataTmp;
+	outSize = deflateBound(&zshPtr->stream, size) + 100;
+	if (outSize > BUFFER_SIZE_LIMIT) {
+	    outSize = BUFFER_SIZE_LIMIT;
+	}
+	dataTmp = ckalloc(outSize);
+
+	while (1) {
+	    e = Deflate(&zshPtr->stream, dataTmp, outSize, flush, &toStore);
+
+	    /*
+	     * Test if we've filled the buffer up and have to ask deflate() to
+	     * give us some more. Note that the condition for needing to
+	     * repeat a buffer transfer when the result is Z_OK is whether
+	     * there is no more space in the buffer we provided; the zlib
+	     * library does not necessarily return a different code in that
+	     * case. [Bug b26e38a3e4] [Tk Bug 10f2e7872b]
+	     */
+
+	    if ((e != Z_BUF_ERROR) && (e != Z_OK || toStore < outSize)) {
+		if ((e == Z_OK) || (flush == Z_FINISH && e == Z_STREAM_END)) {
+		    break;
+		}
+		ConvertError(zshPtr->interp, e, zshPtr->stream.adler);
+		return TCL_ERROR;
+	    }
 
-	e = deflate(&zshPtr->stream, flush);
-	while (e == Z_BUF_ERROR || (flush == Z_FINISH && e == Z_OK)) {
 	    /*
 	     * Output buffer too small to hold the data being generated or we
 	     * are doing the end-of-stream flush (which can spit out masses of
@@ -1202,45 +1260,21 @@ Tcl_ZlibStreamPut(
 	     * saving the old generated data to the outData list.
 	     */
 
-	    obj = Tcl_NewByteArrayObj((unsigned char *) dataTmp, outSize);
-	    Tcl_ListObjAppendElement(NULL, zshPtr->outData, obj);
+	    AppendByteArray(zshPtr->outData, dataTmp, outSize);
 
-	    if (outSize < 0xFFFF) {
-		outSize = 0xFFFF;	/* There may be *lots* of data left to
-					 * output... */
+	    if (outSize < BUFFER_SIZE_LIMIT) {
+		outSize = BUFFER_SIZE_LIMIT;
+		/* There may be *lots* of data left to output... */
 		dataTmp = ckrealloc(dataTmp, outSize);
 	    }
-	    zshPtr->stream.avail_out = outSize;
-	    zshPtr->stream.next_out = (Bytef *) dataTmp;
-
-	    e = deflate(&zshPtr->stream, flush);
-	}
-
-	if (e != Z_OK && !(flush==Z_FINISH && e==Z_STREAM_END)) {
-	    if (zshPtr->interp) {
-		ConvertError(zshPtr->interp, e, zshPtr->stream.adler);
-	    }
-	    return TCL_ERROR;
 	}
 
 	/*
-	 * And append the final data block.
+	 * And append the final data block to the outData list.
 	 */
 
-	if (outSize - zshPtr->stream.avail_out > 0) {
-	    obj = Tcl_NewByteArrayObj((unsigned char *) dataTmp,
-		    outSize - zshPtr->stream.avail_out);
-
-	    /*
-	     * Now append the compressed data to the outData list.
-	     */
-
-	    Tcl_ListObjAppendElement(NULL, zshPtr->outData, obj);
-	}
-
-	if (dataTmp) {
-	    ckfree(dataTmp);
-	}
+	AppendByteArray(zshPtr->outData, dataTmp, toStore);
+	ckfree(dataTmp);
     } else {
 	/*
 	 * This is easy. Just append to the inData list.
@@ -1356,9 +1390,7 @@ Tcl_ZlibStreamGet(
 	if (IsRawStream(zshPtr) && HaveDictToSet(zshPtr)) {
 	    e = SetInflateDictionary(&zshPtr->stream, zshPtr->compDictObj);
 	    if (e != Z_OK) {
-		if (zshPtr->interp) {
-		    ConvertError(zshPtr->interp, e, zshPtr->stream.adler);
-		}
+		ConvertError(zshPtr->interp, e, zshPtr->stream.adler);
 		return TCL_ERROR;
 	    }
 	    DictWasSet(zshPtr);
@@ -2879,10 +2911,8 @@ ZlibTransformClose(
     if (cd->mode == TCL_ZLIB_STREAM_DEFLATE) {
 	cd->outStream.avail_in = 0;
 	do {
-	    cd->outStream.next_out = (Bytef *) cd->outBuffer;
-	    cd->outStream.avail_out = (unsigned) cd->outAllocated;
-	    e = deflate(&cd->outStream, Z_FINISH);
-	    written = cd->outAllocated - cd->outStream.avail_out;
+	    e = Deflate(&cd->outStream, cd->outBuffer, cd->outAllocated,
+		    Z_FINISH, &written);
 
 	    /*
 	     * Can't be sure that deflate() won't declare the buffer to be
@@ -3086,17 +3116,15 @@ ZlibTransformOutput(
     cd->outStream.next_in = (Bytef *) buf;
     cd->outStream.avail_in = toWrite;
     do {
-	cd->outStream.next_out = (Bytef *) cd->outBuffer;
-	cd->outStream.avail_out = cd->outAllocated;
-
-	e = deflate(&cd->outStream, Z_NO_FLUSH);
-	produced = cd->outAllocated - cd->outStream.avail_out;
+	e = Deflate(&cd->outStream, cd->outBuffer, cd->outAllocated,
+		Z_NO_FLUSH, &produced);
 
 	if ((e == Z_OK && produced > 0) || e == Z_BUF_ERROR) {
 	    /*
 	     * deflate() indicates that it is out of space by returning
-	     * Z_BUF_ERROR; in that case, we must write the whole buffer out
-	     * and retry to compress what is left.
+	     * Z_BUF_ERROR *or* by simply returning Z_OK with no remaining
+	     * space; in either case, we must write the whole buffer out and
+	     * retry to compress what is left.
 	     */
 
 	    if (e == Z_BUF_ERROR) {
diff --git a/tests/zlib.test b/tests/zlib.test
index 15dbb34..ae8742b 100644
--- a/tests/zlib.test
+++ b/tests/zlib.test
@@ -138,6 +138,25 @@ test zlib-7.7 {zlib stream: Bug 25842c161} -constraints zlib -body {
 } -cleanup {
     catch {$s close}
 } -result ""
+# Also causes Tk Bug 10f2e7872b
+test zlib-7.8 {zlib stream: Bug b26e38a3e4} -constraints zlib -setup {
+    expr srand(12345)
+    set randdata {}
+    for {set i 0} {$i<6001} {incr i} {
+	append randdata [binary format c [expr {int(256*rand())}]]
+    }
+} -body {
+    set strm [zlib stream compress]
+    for {set i 1} {$i<3000} {incr i} {
+	$strm put $randdata
+    }
+    $strm put -finalize $randdata
+    set data [$strm get]
+    list [string length $data] [string length [zlib decompress $data]]
+} -cleanup {
+    catch {$strm close}
+    unset -nocomplain randdata data
+} -result {120185 18003000}
 
 test zlib-8.1 {zlib transformation} -constraints zlib -setup {
     set file [makeFile {} test.gz]