PageRenderTime 146ms CodeModel.GetById 40ms app.highlight 75ms RepoModel.GetById 26ms app.codeStats 0ms

/mordor/streams/zlib.cpp

http://github.com/mozy/mordor
C++ | 425 lines | 384 code | 24 blank | 17 comment | 81 complexity | 65da1c18f2c96c0abe2c7e43bf849de0 MD5 | raw file
  1// Copyright (c) 2009 - Mozy, Inc.
  2
  3#include "zlib.h"
  4
  5#include "mordor/assert.h"
  6#include "mordor/exception.h"
  7#include "mordor/log.h"
  8
  9#ifdef MSVC
 10#pragma comment(lib, "zdll")
 11#endif
 12
 13namespace Mordor {
 14
 15static Logger::ptr g_log = Log::lookup("mordor:streams:zlib");
 16
 17ZlibStream::ZlibStream(Stream::ptr parent, bool own, Type type, int level,
 18    int windowBits, int memlevel, Strategy strategy, bool invert)
 19    : MutatingFilterStream(parent, own),
 20      m_closed(true)
 21{
 22   init(type, level, windowBits, memlevel, strategy, invert);
 23}
 24
 25void
 26ZlibStream::init(Type type, int level, int windowBits, int memlevel, Strategy strategy, bool invert)
 27{
 28    MORDOR_ASSERT(supportsRead() || supportsWrite());
 29    MORDOR_ASSERT(!(supportsRead() && supportsWrite()));
 30    MORDOR_ASSERT((level >= 0 && level <= 9) || level == Z_DEFAULT_COMPRESSION);
 31    MORDOR_ASSERT(windowBits >= 8 && windowBits <= 15);
 32    MORDOR_ASSERT(memlevel >= 1 && memlevel <= 9);
 33    switch (type) {
 34        case ZLIB:
 35            break;
 36        case DEFLATE:
 37            windowBits = -windowBits;
 38            break;
 39        case GZIP:
 40            windowBits += 16;
 41            break;
 42        default:
 43            MORDOR_ASSERT(false);
 44    }
 45    m_windowBits = windowBits;
 46    m_level = level;
 47    m_memlevel = memlevel;
 48    m_strategy = strategy;
 49    m_doInflate = (!invert && supportsRead()) || (invert && supportsWrite());
 50    reset();
 51}
 52
 53void
 54ZlibStream::reset()
 55{
 56    m_inBuffer.clear();
 57    m_outBuffer.clear(false);
 58    if (!m_closed) {
 59        if (m_doInflate) {
 60            inflateEnd(&m_strm);
 61        } else {
 62            deflateEnd(&m_strm);
 63        }
 64        m_closed = true;
 65    }
 66    int rc;
 67    memset(&m_strm, 0, sizeof(z_stream));
 68    if (m_doInflate) {
 69        rc = inflateInit2(&m_strm, m_windowBits);
 70    } else {
 71        rc = deflateInit2(&m_strm, m_level, Z_DEFLATED, m_windowBits, m_memlevel,
 72            (int)m_strategy);
 73    }
 74    switch (rc) {
 75        case Z_OK:
 76            m_closed = false;
 77            break;
 78        case Z_MEM_ERROR:
 79            throw std::bad_alloc();
 80        case Z_STREAM_ERROR:
 81        {
 82            std::string message(m_strm.msg ? m_strm.msg : "");
 83            if (m_doInflate) {
 84                inflateEnd(&m_strm);
 85            } else {
 86                deflateEnd(&m_strm);
 87            }
 88            throw std::runtime_error(message);
 89        }
 90        default:
 91            MORDOR_NOTREACHED();
 92    }
 93}
 94
 95ZlibStream::ZlibStream(Stream::ptr parent, int level, int windowBits, int memlevel, Strategy strategy,
 96    bool own, bool invert)
 97    : MutatingFilterStream(parent, own),
 98      m_closed(true)
 99{
100    init(ZLIB, level, windowBits, memlevel, strategy, invert);
101}
102
103ZlibStream::ZlibStream(Stream::ptr parent, bool own, bool invert)
104    : MutatingFilterStream(parent, own),
105      m_closed(true)
106{
107    init(ZLIB, Z_DEFAULT_COMPRESSION, 15, 8, DEFAULT, invert);
108}
109
110ZlibStream::~ZlibStream()
111{
112    if (!m_closed) {
113        if (m_doInflate) {
114            inflateEnd(&m_strm);
115        } else {
116            deflateEnd(&m_strm);
117        }
118    }
119}
120
121void
122ZlibStream::close(CloseType type)
123{
124    if ((type == READ && supportsWrite()) ||
125        (type == WRITE && supportsRead()) ||
126        m_closed) {
127        if (ownsParent())
128            parent()->close(type);
129        return;
130    }
131    if (supportsWrite())
132        flush(Z_FINISH);
133    if (m_doInflate)
134        inflateEnd(&m_strm);
135    else
136        deflateEnd(&m_strm);
137    m_closed = true;
138    if (ownsParent())
139        parent()->close(type);
140}
141
142size_t ZlibStream::doInflateForRead(Buffer &buffer, size_t length)
143{
144    if (m_closed)
145        return 0;
146    struct iovec outbuf = buffer.writeBuffer(length, false);
147    m_strm.next_out = (Bytef*)outbuf.iov_base;
148    m_strm.avail_out = outbuf.iov_len;
149
150    while (true) {
151        std::vector<iovec> inbufs = m_inBuffer.readBuffers();
152        if (!inbufs.empty()) {
153            m_strm.next_in = (Bytef*)inbufs[0].iov_base;
154            m_strm.avail_in = inbufs[0].iov_len;
155        } else {
156            m_strm.next_in = NULL;
157            m_strm.avail_in = 0;
158        }
159        int rc = inflate(&m_strm, Z_NO_FLUSH);
160        MORDOR_LOG_DEBUG(g_log) << this << " inflate(("
161            << (inbufs.empty() ? 0 : inbufs[0].iov_len) << ", "
162            << outbuf.iov_len << ")): " << rc << " (" << m_strm.avail_in
163            << ", " << m_strm.avail_out << ")";
164        if (!inbufs.empty())
165            m_inBuffer.consume(inbufs[0].iov_len - m_strm.avail_in);
166        size_t result;
167        switch (rc) {
168            case Z_STREAM_END:
169                // May have still produced output
170                result = outbuf.iov_len - m_strm.avail_out;
171                buffer.produce(result);
172                inflateEnd(&m_strm);
173                m_closed = true;
174                return result;
175            case Z_OK:
176                result = outbuf.iov_len - m_strm.avail_out;
177                // It consumed input, but produced no output... DON'T return eof
178                if (result == 0)
179                    continue;
180                buffer.produce(result);
181                return result;
182            case Z_BUF_ERROR:
183                // no progress... we need to provide more input (since we're
184                // guaranteed to provide output)
185                MORDOR_ASSERT(m_strm.avail_in == 0);
186                MORDOR_ASSERT(inbufs.empty());
187                result = parent()->read(m_inBuffer, m_bufferSize);
188                if (result == 0)
189                    MORDOR_THROW_EXCEPTION(UnexpectedEofException());
190                break;
191            case Z_MEM_ERROR:
192                MORDOR_THROW_EXCEPTION(std::bad_alloc());
193            case Z_NEED_DICT:
194                MORDOR_THROW_EXCEPTION(NeedPresetDictionaryException());
195            case Z_DATA_ERROR:
196                MORDOR_THROW_EXCEPTION(CorruptedZlibStreamException());
197            default:
198                MORDOR_NOTREACHED();
199        }
200    }
201}
202
203size_t ZlibStream::doDeflateForRead(Buffer &buffer, size_t length)
204{
205    if (m_closed)
206        return 0;
207    struct iovec outbuf = buffer.writeBuffer(length, false);
208    m_strm.next_out = (Bytef*)outbuf.iov_base;
209    m_strm.avail_out = outbuf.iov_len;
210
211    while (true) {
212        bool parentEof = false;
213        if (m_inBuffer.readAvailable() == 0) {
214            size_t result = parent()->read(m_inBuffer, m_bufferSize);
215            if (result == 0) {
216                parentEof = true;
217            }
218        }
219        struct iovec inbuf = m_inBuffer.readBuffer((size_t)~0, true);
220        m_strm.next_in = (Bytef*)inbuf.iov_base;
221        m_strm.avail_in = inbuf.iov_len;
222        int rc = deflate(&m_strm, parentEof?Z_FINISH:Z_NO_FLUSH);
223        MORDOR_LOG_DEBUG(g_log) << this << " deflate((" << inbuf.iov_len << ", "
224            << outbuf.iov_len << "), " << (parentEof?"Z_FINISH":"Z_NO_FLUSH") << "): " << rc << " ("
225            << m_strm.avail_in << ", " << m_strm.avail_out << ")";
226        // We are always providing both input and output
227        MORDOR_ASSERT(rc != Z_BUF_ERROR);
228        m_inBuffer.consume(inbuf.iov_len - m_strm.avail_in);
229        size_t result;
230        switch(rc) {
231            case Z_STREAM_END:
232                // Pending output is flushed and there is enough output space
233                result = outbuf.iov_len - m_strm.avail_out;
234                buffer.produce(result);
235                return result;
236            case Z_OK:
237                //if deflate returns with Z_OK, this function must be called again with Z_FINISH
238                //and more output space (updated avail_out) but no more input data, until it returns
239                //with Z_STREAM_END or an error.
240                result = outbuf.iov_len - m_strm.avail_out;
241                if (result == 0)
242                    continue;
243                buffer.produce(result);
244                return result;
245            default:
246                MORDOR_NOTREACHED();
247        }
248    }
249}
250
251size_t ZlibStream::doInflateForWrite(const Buffer &buffer, size_t length)
252{
253    MORDOR_ASSERT(!m_closed);
254    flushBuffer();
255    while (true) {
256        if (m_outBuffer.writeAvailable() == 0)
257            m_outBuffer.reserve(m_bufferSize);
258        struct iovec outbuf = m_outBuffer.writeBuffer(~0u, false);
259        size_t len = (std::min)(length, buffer.readAvailable());
260        struct iovec inbuf = buffer.readBuffer(len, true);
261        m_strm.next_in = (Bytef*)inbuf.iov_base;
262        m_strm.avail_in = inbuf.iov_len;
263        m_strm.next_out = (Bytef*)outbuf.iov_base;
264        m_strm.avail_out = outbuf.iov_len;
265        int rc = inflate(&m_strm, Z_NO_FLUSH);
266        MORDOR_LOG_DEBUG(g_log) << this << " inflate(("
267            << inbuf.iov_len << ", " << outbuf.iov_len << ")): " << rc
268            << " (" << m_strm.avail_in << ", " << m_strm.avail_out << ")";
269        // We are always providing both input and output
270        MORDOR_ASSERT(rc != Z_BUF_ERROR);
271        size_t result;
272        switch (rc) {
273            case Z_STREAM_END:
274                // May have still produced output
275                result = inbuf.iov_len - m_strm.avail_in;
276                m_outBuffer.produce(outbuf.iov_len - m_strm.avail_out);
277                m_closed = true;
278                inflateEnd(&m_strm);
279                try {
280                    flushBuffer();
281                } catch (const std::runtime_error&) {
282                    // Swallow it
283                }
284                return result;
285            case Z_OK:
286                // some progress has been made (more input processed or more output produced)
287                result = inbuf.iov_len - m_strm.avail_in;
288                if (result == 0)
289                    continue;
290                m_outBuffer.produce(outbuf.iov_len - m_strm.avail_out);
291                flushBuffer();
292                return result;
293            case Z_MEM_ERROR:
294                MORDOR_THROW_EXCEPTION(std::bad_alloc());
295            case Z_NEED_DICT:
296                MORDOR_THROW_EXCEPTION(NeedPresetDictionaryException());
297            case Z_DATA_ERROR:
298                MORDOR_THROW_EXCEPTION(CorruptedZlibStreamException());
299            default:
300                MORDOR_NOTREACHED();
301        }
302    }
303}
304
305size_t ZlibStream::doDeflateForWrite(const Buffer &buffer, size_t length)
306{
307    MORDOR_ASSERT(!m_closed);
308    flushBuffer();
309    while (true) {
310        if (m_outBuffer.writeAvailable() == 0)
311            m_outBuffer.reserve(m_bufferSize);
312        struct iovec inbuf = buffer.readBuffer(length, false);
313        struct iovec outbuf = m_outBuffer.writeBuffer(~0u, false);
314        m_strm.next_in = (Bytef*)inbuf.iov_base;
315        m_strm.avail_in = inbuf.iov_len;
316        m_strm.next_out = (Bytef*)outbuf.iov_base;
317        m_strm.avail_out = outbuf.iov_len;
318        int rc = deflate(&m_strm, Z_NO_FLUSH);
319        MORDOR_LOG_DEBUG(g_log) << this << " deflate((" << inbuf.iov_len << ", "
320            << outbuf.iov_len << "), Z_NO_FLUSH): " << rc << " ("
321            << m_strm.avail_in << ", " << m_strm.avail_out << ")";
322        // We are always providing both input and output
323        MORDOR_ASSERT(rc != Z_BUF_ERROR);
324        // We're not doing Z_FINISH, so we shouldn't get EOF
325        MORDOR_ASSERT(rc != Z_STREAM_END);
326        size_t result;
327        switch(rc) {
328            case Z_OK:
329                result = inbuf.iov_len - m_strm.avail_in;
330                if (result == 0)
331                    continue;
332                m_outBuffer.produce(outbuf.iov_len - m_strm.avail_out);
333                try {
334                    flushBuffer();
335                } catch (const std::runtime_error&) {
336                    // Swallow it
337                }
338                return result;
339            default:
340                MORDOR_NOTREACHED();
341        }
342    }
343}
344
345size_t
346ZlibStream::read(Buffer &buffer, size_t length)
347{
348    if (m_doInflate)
349        return doInflateForRead(buffer, length);
350    else
351        return doDeflateForRead(buffer, length);
352}
353
354size_t
355ZlibStream::write(const Buffer &buffer, size_t length)
356{
357    if (m_doInflate)
358        return doInflateForWrite(buffer, length);
359    else
360        return doDeflateForWrite(buffer, length);
361}
362
363void
364ZlibStream::flush(bool flushParent)
365{
366    flush(Z_SYNC_FLUSH);
367    if (flushParent)
368        parent()->flush();
369}
370
371void
372ZlibStream::flush(int flush)
373{
374    flushBuffer();
375    while (true) {
376        if (m_outBuffer.writeAvailable() == 0)
377            m_outBuffer.reserve(m_bufferSize);
378        struct iovec outbuf = m_outBuffer.writeBuffer(~0u, false);
379        MORDOR_ASSERT(m_strm.avail_in == 0);
380        m_strm.next_out = (Bytef*)outbuf.iov_base;
381        m_strm.avail_out = outbuf.iov_len;
382        int rc;
383        if (m_doInflate) {
384            rc = inflate(&m_strm, flush);
385            MORDOR_LOG_DEBUG(g_log) << this << " inflate((0, " << outbuf.iov_len
386                << "), " << flush << "): " << rc << " (0, " << m_strm.avail_out
387                << ")";
388        } else {
389            rc = deflate(&m_strm, flush);
390            MORDOR_LOG_DEBUG(g_log) << this << " deflate((0, " << outbuf.iov_len
391                << "), " << flush << "): " << rc << " (0, " << m_strm.avail_out
392                << ")";
393        }
394        MORDOR_ASSERT(m_strm.avail_in == 0);
395        m_outBuffer.produce(outbuf.iov_len - m_strm.avail_out);
396        MORDOR_ASSERT(flush == Z_FINISH || rc != Z_STREAM_END);
397        switch (rc) {
398            case Z_STREAM_END:
399                m_closed = true;
400                if (m_doInflate)
401                    inflateEnd(&m_strm);
402                else
403                    deflateEnd(&m_strm);
404                flushBuffer();
405                return;
406            case Z_OK:
407                break;
408            case Z_BUF_ERROR:
409                flushBuffer();
410                return;
411            default:
412                MORDOR_NOTREACHED();
413        }
414    }
415}
416
417void
418ZlibStream::flushBuffer()
419{
420    while (m_outBuffer.readAvailable() > 0)
421        m_outBuffer.consume(parent()->write(m_outBuffer,
422            m_outBuffer.readAvailable()));
423}
424
425}