Implement partial write buffering for AES.

pdfio-stream.c

@@ -54,8 +54,9 @@ pdfioStreamClose(pdfio_stream_t *st)	// I - Stream
 
       while ((status = deflate(&st->flate, Z_FINISH)) != Z_STREAM_END)
       {
-        size_t bytes = sizeof(st->cbuffer) - st->flate.avail_out;
+        size_t	bytes = sizeof(st->cbuffer) - st->flate.avail_out,
 					// Bytes to write
+		outbytes;		// Actual bytes written
 
 	if (status < Z_OK && status != Z_BUF_ERROR)
 	{
@@ -67,17 +68,32 @@ pdfioStreamClose(pdfio_stream_t *st)	// I - Stream
 	if (st->crypto_cb)
 	{
 	  // Encrypt it first...
-	  bytes = (st->crypto_cb)(&st->crypto_ctx, st->cbuffer, st->cbuffer, bytes);
+	  outbytes = (st->crypto_cb)(&st->crypto_ctx, st->cbuffer, st->cbuffer, bytes & ~15);
+	}
+	else
+	{
+	  // No encryption
+	  outbytes = bytes;
 	}
 
-	if (!_pdfioFileWrite(st->pdf, st->cbuffer, bytes))
+	if (!_pdfioFileWrite(st->pdf, st->cbuffer, outbytes))
 	{
 	  ret = false;
 	  goto done;
 	}
 
-	st->flate.next_out  = (Bytef *)st->cbuffer;
+        if (bytes > outbytes)
-	st->flate.avail_out = (uInt)sizeof(st->cbuffer);
+        {
+          bytes -= outbytes;
+          memmove(st->cbuffer, st->cbuffer + outbytes, bytes);
+        }
+        else
+        {
+          bytes = 0;
+        }
+
+	st->flate.next_out  = (Bytef *)st->cbuffer + bytes;
+	st->flate.avail_out = (uInt)sizeof(st->cbuffer) - bytes;
       }
 
       if (st->flate.avail_out < (uInt)sizeof(st->cbuffer))
@@ -99,6 +115,19 @@ pdfioStreamClose(pdfio_stream_t *st)	// I - Stream
 	}
       }
     }
+    else if (st->crypto_cb && st->bufptr > st->buffer)
+    {
+      // Encrypt and flush
+      uint8_t	temp[8192];		// Temporary buffer
+      size_t	outbytes;		// Output bytes
+
+      outbytes = (st->crypto_cb)(&st->crypto_ctx, temp, st->buffer, (size_t)(st->bufptr - st->buffer));
+      if (!_pdfioFileWrite(st->pdf, temp, outbytes))
+      {
+        ret = false;
+        goto done;
+      }
+    }
 
     // Save the length of this stream...
     st->obj->stream_length = (size_t)(_pdfioFileTell(st->pdf) - st->obj->stream_offset);
@@ -177,6 +206,8 @@ _pdfioStreamCreate(
   st->obj        = obj;
   st->length_obj = length_obj;
   st->filter     = compression;
+  st->bufptr     = st->buffer;
+  st->bufend     = st->buffer + sizeof(st->buffer);
 
   if (obj->pdf->encryption)
   {
@@ -760,7 +791,7 @@ pdfioStreamWrite(
     if (st->crypto_cb)
     {
       // Encrypt data before writing...
-      unsigned char	temp[8192];	// Temporary buffer
+      uint8_t	temp[8192];		// Temporary buffer
       size_t	cbytes,			// Current bytes
 		outbytes;		// Output bytes
 
@@ -768,12 +799,39 @@ pdfioStreamWrite(
 
       while (bytes > 0)
       {
+        if (st->bufptr > st->buffer || bytes < 16)
+        {
+          // Write through the stream's buffer...
+          if ((cbytes = bytes) > (st->bufend - st->bufptr))
+            cbytes = st->bufend - st->bufptr;
+
+          memcpy(st->bufptr, bufptr, cbytes);
+          st->bufptr += cbytes;
+          if (st->bufptr >= st->bufend)
+          {
+            // Encrypt and flush
+	    outbytes = (st->crypto_cb)(&st->crypto_ctx, temp, st->buffer, sizeof(st->buffer));
+	    if (!_pdfioFileWrite(st->pdf, temp, outbytes))
+	      return (false);
+
+	    st->bufptr = st->buffer;
+          }
+        }
+        else
+        {
+          // Write directly up to sizeof(temp) bytes...
           if ((cbytes = bytes) > sizeof(temp))
             cbytes = sizeof(temp);
+          if (cbytes & 15)
+          {
+            // AES has a 16-byte block size, so save the last few bytes...
+            cbytes &= ~15;
+          }
 
 	  outbytes = (st->crypto_cb)(&st->crypto_ctx, temp, bufptr, cbytes);
 	  if (!_pdfioFileWrite(st->pdf, temp, outbytes))
 	    return (false);
+        }
 
         bytes -= cbytes;
         bufptr += cbytes;
@@ -1160,19 +1218,34 @@ stream_write(pdfio_stream_t *st,	// I - Stream
     if (st->flate.avail_out < (sizeof(st->cbuffer) / 8))
     {
       // Flush the compression buffer...
-      size_t	bytes = sizeof(st->cbuffer) - st->flate.avail_out;
+      size_t	bytes = sizeof(st->cbuffer) - st->flate.avail_out,
+		outbytes;
 
       if (st->crypto_cb)
       {
         // Encrypt it first...
-        bytes = (st->crypto_cb)(&st->crypto_ctx, st->cbuffer, st->cbuffer, bytes);
+        outbytes = (st->crypto_cb)(&st->crypto_ctx, st->cbuffer, st->cbuffer, bytes & ~15);
+      }
+      else
+      {
+        outbytes = bytes;
       }
 
       if (!_pdfioFileWrite(st->pdf, st->cbuffer, bytes))
         return (false);
 
-      st->flate.next_out  = (Bytef *)st->cbuffer;
+      if (bytes > outbytes)
-      st->flate.avail_out = sizeof(st->cbuffer);
+      {
+        bytes -= outbytes;
+        memmove(st->cbuffer, st->cbuffer + outbytes, bytes);
+      }
+      else
+      {
+        bytes = 0;
+      }
+
+      st->flate.next_out  = (Bytef *)st->cbuffer + bytes;
+      st->flate.avail_out = sizeof(st->cbuffer) - bytes;
     }
 
     // Deflate what we can this time...