diff options
Diffstat (limited to '')
| -rw-r--r-- | comm/mailnews/base/src/nsQuarantinedOutputStream.cpp | 234 |
1 files changed, 234 insertions, 0 deletions
diff --git a/comm/mailnews/base/src/nsQuarantinedOutputStream.cpp b/comm/mailnews/base/src/nsQuarantinedOutputStream.cpp new file mode 100644 index 0000000000..1f325edadf --- /dev/null +++ b/comm/mailnews/base/src/nsQuarantinedOutputStream.cpp @@ -0,0 +1,234 @@ +/* -*- Mode: C++; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */ +/* 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 "nsQuarantinedOutputStream.h" +#include "nsDirectoryServiceDefs.h" +#include "nsIInputStream.h" +#include "nsISeekableStream.h" +#include "nsIFile.h" +#include "nsNetUtil.h" +#include "mozilla/UniquePtr.h" + +NS_IMPL_ISUPPORTS(nsQuarantinedOutputStream, nsIOutputStream, + nsISafeOutputStream) + +nsQuarantinedOutputStream::~nsQuarantinedOutputStream() { Close(); } + +// Initialise mTempFile and open it for writing (mTempStream). +nsresult nsQuarantinedOutputStream::InitTemp() { + MOZ_ASSERT(mState == eUninitialized); + MOZ_ASSERT(!mTempFile); + MOZ_ASSERT(!mTempStream); + // Create a unique temp file. + { + nsCOMPtr<nsIFile> file; + nsresult rv = NS_GetSpecialDirectory(NS_OS_TEMP_DIR, getter_AddRefs(file)); + NS_ENSURE_SUCCESS(rv, rv); + rv = file->Append(u"newmsg"_ns); + NS_ENSURE_SUCCESS(rv, rv); + rv = file->CreateUnique(nsIFile::NORMAL_FILE_TYPE, 0600); + NS_ENSURE_SUCCESS(rv, rv); + mTempFile = std::move(file); + } + + // Open the temp file for writing. + { + nsCOMPtr<nsIOutputStream> stream; + nsresult rv = NS_NewLocalFileOutputStream(getter_AddRefs(stream), mTempFile, + -1, 0600); + NS_ENSURE_SUCCESS(rv, rv); + mTempStream = std::move(stream); + } + + return NS_OK; +} + +// Put us into the error state and clean up (by deleting the temp file +// if it exists). +void nsQuarantinedOutputStream::EnterErrorState(nsresult status) { + mState = eError; + mError = status; + mTarget = nullptr; + + if (mTempStream) { + mTempStream = nullptr; + } + if (mTempFile) { + mTempFile->Remove(false); + mTempFile = nullptr; + } +} + +// copyStream copies all the data in the input stream to the output stream. +// It keeps going until it sees an EOF on the input. +static nsresult copyStream(nsIInputStream* in, nsIOutputStream* out) { + constexpr uint32_t BUFSIZE = 8192; + auto buf = mozilla::MakeUnique<char[]>(BUFSIZE); + while (true) { + // Read input stream into buf. + uint32_t bufCnt; + nsresult rv = in->Read(buf.get(), BUFSIZE, &bufCnt); + NS_ENSURE_SUCCESS(rv, rv); + if (bufCnt == 0) { + break; // EOF. We're all done! + } + // Write buf to output stream. + uint32_t pos = 0; + while (pos < bufCnt) { + uint32_t writeCnt; + rv = out->Write(buf.get() + pos, bufCnt - pos, &writeCnt); + NS_ENSURE_SUCCESS(rv, rv); + pos += writeCnt; + } + } + return NS_OK; +} + +// copyStreamSafely() wraps copyStream(). If the output stream is seekable, +// it will try to roll it back if an error occurs during the copy. +static nsresult copyStreamSafely(nsIInputStream* in, nsIOutputStream* out) { + nsCOMPtr<nsISeekableStream> outSeekable = do_QueryInterface(out); + if (!outSeekable) { + // It's not seekable, so we jump out without a parachute. + return copyStream(in, out); + } + int64_t initialOffset; + nsresult rv = outSeekable->Tell(&initialOffset); + NS_ENSURE_SUCCESS(rv, rv); + rv = copyStream(in, out); + if (NS_FAILED(rv)) { + // Uhoh... the copy failed! Try to remove the partially-written data. + rv = outSeekable->Seek(nsISeekableStream::NS_SEEK_SET, initialOffset); + NS_ENSURE_SUCCESS(rv, rv); + rv = outSeekable->SetEOF(); + NS_ENSURE_SUCCESS(rv, rv); + } + return NS_OK; +} + +NS_IMETHODIMP nsQuarantinedOutputStream::Close() { + if (mState != eOpen) { + // Already failed or closed or no data written. That's OK. + return NS_OK; + } + nsresult rv = NS_OK; + if (mTempStream) { + rv = mTempStream->Close(); + mTempStream = nullptr; + } + if (mTempFile) { + mTempFile->Remove(false); + mTempFile = nullptr; + } + mTarget->Close(); + mTarget = nullptr; + mState = eClosed; + return rv; +} + +NS_IMETHODIMP nsQuarantinedOutputStream::Finish() { + // Fail here if there was a previous error. + if (mState == eError) { + return mError; + } + if (mState != eOpen) { + // Already closed or no data written. That's OK. + return NS_OK; + } + + // Flush and close the temp file. Hopefully any virus checker will now act + // and prevent us reopening any suspicious-looking file. + MOZ_ASSERT(mTempStream); + MOZ_ASSERT(mTempFile); + mTempStream->Flush(); + nsresult rv = mTempStream->Close(); + if (NS_FAILED(rv)) { + EnterErrorState(rv); + return rv; + } + mTempStream = nullptr; + + // Write the tempfile out to the target stream + { + nsCOMPtr<nsIInputStream> ins; + // If a virus checker smells something bad, it should show up here as a + // failure to (re)open the temp file. + rv = NS_NewLocalFileInputStream(getter_AddRefs(ins), mTempFile); + if (NS_FAILED(rv)) { + EnterErrorState(rv); + return rv; + } + rv = copyStreamSafely(ins, mTarget); + if (NS_FAILED(rv)) { + EnterErrorState(rv); + return rv; + } + } + + // All done! + mTarget->Close(); + mTempFile->Remove(false); + mTempFile = nullptr; + mState = eClosed; + mTarget = nullptr; + return NS_OK; +} + +NS_IMETHODIMP nsQuarantinedOutputStream::Flush() { + if (mState != eOpen) { + return NS_OK; // Don't rock the boat. + } + nsresult rv = mTempStream->Flush(); + if (NS_FAILED(rv)) { + EnterErrorState(rv); + } + return rv; +} + +NS_IMETHODIMP nsQuarantinedOutputStream::Write(const char* buf, uint32_t count, + uint32_t* result) { + if (mState == eUninitialized) { + // Lazy open. + nsresult rv = InitTemp(); + if NS_FAILED (rv) { + EnterErrorState(rv); + return rv; + } + mState = eOpen; + } + + if (mState != eOpen) { + return NS_ERROR_UNEXPECTED; + } + + nsresult rv = mTempStream->Write(buf, count, result); + if (NS_FAILED(rv)) { + EnterErrorState(rv); + return rv; + } + return NS_OK; +} + +NS_IMETHODIMP nsQuarantinedOutputStream::WriteFrom(nsIInputStream* fromStream, + uint32_t count, + uint32_t* retval) { + return NS_ERROR_NOT_IMPLEMENTED; +} + +NS_IMETHODIMP nsQuarantinedOutputStream::WriteSegments(nsReadSegmentFun reader, + void* closure, + uint32_t count, + uint32_t* retval) { + return NS_ERROR_NOT_IMPLEMENTED; +} + +NS_IMETHODIMP nsQuarantinedOutputStream::IsNonBlocking(bool* nonBlocking) { + *nonBlocking = false; + return NS_OK; +} + +NS_IMETHODIMP nsQuarantinedOutputStream::StreamStatus() { + return mState == eOpen ? NS_OK : NS_BASE_STREAM_CLOSED; +} |