SnappyCompressOutputStream.cpp

/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
/* vim: set ts=8 sts=2 et sw=2 tw=80: */
/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include "mozilla/SnappyCompressOutputStream.h"

#include <algorithm>
#include "nsStreamUtils.h"
#include "snappy/snappy.h"

namespace mozilla {

NS_IMPL_ISUPPORTS(SnappyCompressOutputStream, nsIOutputStream);

// static
const size_t SnappyCompressOutputStream::kMaxBlockSize = snappy::kBlockSize;

SnappyCompressOutputStream::SnappyCompressOutputStream(
    nsIOutputStream* aBaseStream, size_t aBlockSize)
    : mBaseStream(aBaseStream),
      mBlockSize(std::min(aBlockSize, kMaxBlockSize)),
      mNextByte(0),
      mCompressedBufferLength(0),
      mStreamIdentifierWritten(false) {
  MOZ_ASSERT(mBlockSize > 0);

  // This implementation only supports sync base streams.  Verify this in debug
  // builds.  Note, this can be simpler than the check in
  // SnappyUncompressInputStream because we don't have to deal with the
  // nsStringInputStream oddness of being non-blocking and sync.
#ifdef DEBUG
  bool baseNonBlocking;
  nsresult rv = mBaseStream->IsNonBlocking(&baseNonBlocking);
  MOZ_ASSERT(NS_SUCCEEDED(rv));
  MOZ_ASSERT(!baseNonBlocking);
#endif
}

size_t SnappyCompressOutputStream::BlockSize() const { return mBlockSize; }

NS_IMETHODIMP
SnappyCompressOutputStream::Close() {
  if (!mBaseStream) {
    return NS_OK;
  }

  nsresult rv = Flush();
  if (NS_WARN_IF(NS_FAILED(rv))) {
    return rv;
  }

  mBaseStream->Close();
  mBaseStream = nullptr;

  mBuffer = nullptr;
  mCompressedBuffer = nullptr;

  return NS_OK;
}

NS_IMETHODIMP
SnappyCompressOutputStream::Flush() {
  if (!mBaseStream) {
    return NS_BASE_STREAM_CLOSED;
  }

  nsresult rv = FlushToBaseStream();
  if (NS_WARN_IF(NS_FAILED(rv))) {
    return rv;
  }

  mBaseStream->Flush();

  return NS_OK;
}

NS_IMETHODIMP
SnappyCompressOutputStream::Write(const char* aBuf, uint32_t aCount,
                                  uint32_t* aResultOut) {
  return WriteSegments(NS_CopySegmentToBuffer, const_cast<char*>(aBuf), aCount,
                       aResultOut);
}

NS_IMETHODIMP
SnappyCompressOutputStream::WriteFrom(nsIInputStream*, uint32_t, uint32_t*) {
  return NS_ERROR_NOT_IMPLEMENTED;
}

NS_IMETHODIMP
SnappyCompressOutputStream::WriteSegments(nsReadSegmentFun aReader,
                                          void* aClosure, uint32_t aCount,
                                          uint32_t* aBytesWrittenOut) {
  *aBytesWrittenOut = 0;

  if (!mBaseStream) {
    return NS_BASE_STREAM_CLOSED;
  }

  if (!mBuffer) {
    mBuffer.reset(new (fallible) char[mBlockSize]);
    if (NS_WARN_IF(!mBuffer)) {
      return NS_ERROR_OUT_OF_MEMORY;
    }
  }

  while (aCount > 0) {
    // Determine how much space is left in our flat, uncompressed buffer.
    MOZ_ASSERT(mNextByte <= mBlockSize);
    uint32_t remaining = mBlockSize - mNextByte;

    // If it is full, then compress and flush the data to the base stream.
    if (remaining == 0) {
      nsresult rv = FlushToBaseStream();
      if (NS_WARN_IF(NS_FAILED(rv))) {
        return rv;
      }

      // Now the entire buffer should be available for copying.
      MOZ_ASSERT(!mNextByte);
      remaining = mBlockSize;
    }

    uint32_t numToRead = std::min(remaining, aCount);
    uint32_t numRead = 0;

    nsresult rv = aReader(this, aClosure, &mBuffer[mNextByte],
                          *aBytesWrittenOut, numToRead, &numRead);

    // As defined in nsIOutputStream.idl, do not pass reader func errors.
    if (NS_FAILED(rv)) {
      return NS_OK;
    }

    // End-of-file
    if (numRead == 0) {
      return NS_OK;
    }

    mNextByte += numRead;
    *aBytesWrittenOut += numRead;
    aCount -= numRead;
  }

  return NS_OK;
}

NS_IMETHODIMP
SnappyCompressOutputStream::IsNonBlocking(bool* aNonBlockingOut) {
  *aNonBlockingOut = false;
  return NS_OK;
}

SnappyCompressOutputStream::~SnappyCompressOutputStream() { Close(); }

nsresult SnappyCompressOutputStream::FlushToBaseStream() {
  MOZ_ASSERT(mBaseStream);

  // Lazily create the compressed buffer on our first flush.  This
  // allows us to report OOM during stream operation.  This buffer
  // will then get re-used until the stream is closed.
  if (!mCompressedBuffer) {
    mCompressedBufferLength = MaxCompressedBufferLength(mBlockSize);
    mCompressedBuffer.reset(new (fallible) char[mCompressedBufferLength]);
    if (NS_WARN_IF(!mCompressedBuffer)) {
      return NS_ERROR_OUT_OF_MEMORY;
    }
  }

  // The first chunk must be a StreamIdentifier chunk.  Write it out
  // if we have not done so already.
  nsresult rv = MaybeFlushStreamIdentifier();
  if (NS_WARN_IF(NS_FAILED(rv))) {
    return rv;
  }

  // Compress the data to our internal compressed buffer.
  size_t compressedLength;
  rv = WriteCompressedData(mCompressedBuffer.get(), mCompressedBufferLength,
                           mBuffer.get(), mNextByte, &compressedLength);
  if (NS_WARN_IF(NS_FAILED(rv))) {
    return rv;
  }
  MOZ_ASSERT(compressedLength > 0);

  mNextByte = 0;

  // Write the compressed buffer out to the base stream.
  uint32_t numWritten = 0;
  rv = WriteAll(mCompressedBuffer.get(), compressedLength, &numWritten);
  if (NS_WARN_IF(NS_FAILED(rv))) {
    return rv;
  }
  MOZ_ASSERT(compressedLength == numWritten);

  return NS_OK;
}

nsresult SnappyCompressOutputStream::MaybeFlushStreamIdentifier() {
  MOZ_ASSERT(mCompressedBuffer);

  if (mStreamIdentifierWritten) {
    return NS_OK;
  }

  // Build the StreamIdentifier in our compressed buffer.
  size_t compressedLength;
  nsresult rv = WriteStreamIdentifier(
      mCompressedBuffer.get(), mCompressedBufferLength, &compressedLength);
  if (NS_WARN_IF(NS_FAILED(rv))) {
    return rv;
  }

  // Write the compressed buffer out to the base stream.
  uint32_t numWritten = 0;
  rv = WriteAll(mCompressedBuffer.get(), compressedLength, &numWritten);
  if (NS_WARN_IF(NS_FAILED(rv))) {
    return rv;
  }
  MOZ_ASSERT(compressedLength == numWritten);

  mStreamIdentifierWritten = true;

  return NS_OK;
}

nsresult SnappyCompressOutputStream::WriteAll(const char* aBuf, uint32_t aCount,
                                              uint32_t* aBytesWrittenOut) {
  *aBytesWrittenOut = 0;

  if (!mBaseStream) {
    return NS_BASE_STREAM_CLOSED;
  }

  uint32_t offset = 0;
  while (aCount > 0) {
    uint32_t numWritten = 0;
    nsresult rv = mBaseStream->Write(aBuf + offset, aCount, &numWritten);
    if (NS_WARN_IF(NS_FAILED(rv))) {
      return rv;
    }
    offset += numWritten;
    aCount -= numWritten;
    *aBytesWrittenOut += numWritten;
  }

  return NS_OK;
}

}  // namespace mozilla