diff --git a/pdfio-stream.c b/pdfio-stream.c index 4be541c..135fee1 100644 --- a/pdfio-stream.c +++ b/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; - st->flate.avail_out = (uInt)sizeof(st->cbuffer); + if (bytes > outbytes) + { + 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,20 +791,47 @@ pdfioStreamWrite( if (st->crypto_cb) { // Encrypt data before writing... - unsigned char temp[8192]; // Temporary buffer - size_t cbytes, // Current bytes - outbytes; // Output bytes + uint8_t temp[8192]; // Temporary buffer + size_t cbytes, // Current bytes + outbytes; // Output bytes bufptr = (const unsigned char *)buffer; while (bytes > 0) { - if ((cbytes = bytes) > sizeof(temp)) - cbytes = sizeof(temp); + 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; - outbytes = (st->crypto_cb)(&st->crypto_ctx, temp, bufptr, cbytes); - if (!_pdfioFileWrite(st->pdf, temp, outbytes)) - return (false); + 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; - st->flate.avail_out = sizeof(st->cbuffer); + if (bytes > outbytes) + { + 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...