xpcom/io/nsInputStreamTee.cpp

   1 /* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
   2 /* vim: set ts=8 sts=2 et sw=2 tw=80: */
   3 /* This Source Code Form is subject to the terms of the Mozilla Public
   4  * License, v. 2.0. If a copy of the MPL was not distributed with this
   5  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
   6
   7 #include <stdlib.h>
   8 #include "prlog.h"
   9
  10 #include "mozilla/Mutex.h"
  11 #include "mozilla/Attributes.h"
  12 #include "nsIInputStreamTee.h"
  13 #include "nsIInputStream.h"
  14 #include "nsIOutputStream.h"
  15 #include "nsCOMPtr.h"
  16 #include "nsAutoPtr.h"
  17 #include "nsIEventTarget.h"
  18 #include "nsThreadUtils.h"
  19
  20 using namespace mozilla;
  21
  22 #ifdef LOG
  23 #undef LOG
  24 #endif
  25 #ifdef PR_LOGGING
  26 static PRLogModuleInfo*
  27 GetTeeLog()
  28 {
  29   static PRLogModuleInfo* sLog;
  30   if (!sLog) {
  31     sLog = PR_NewLogModule("nsInputStreamTee");
  32   }
  33   return sLog;
  34 }
  35 #define LOG(args) PR_LOG(GetTeeLog(), PR_LOG_DEBUG, args)
  36 #else
  37 #define LOG(args)
  38 #endif
  39
  40 class nsInputStreamTee MOZ_FINAL : public nsIInputStreamTee
  41 {
  42 public:
  43   NS_DECL_THREADSAFE_ISUPPORTS
  44   NS_DECL_NSIINPUTSTREAM
  45   NS_DECL_NSIINPUTSTREAMTEE
  46
  47   nsInputStreamTee();
  48   bool SinkIsValid();
  49   void InvalidateSink();
  50
  51 private:
  52   ~nsInputStreamTee()
  53   {
  54   }
  55
  56   nsresult TeeSegment(const char* aBuf, uint32_t aCount);
  57
  58   static NS_METHOD WriteSegmentFun(nsIInputStream*, void*, const char*,
  59                                    uint32_t, uint32_t, uint32_t*);
  60
  61 private:
  62   nsCOMPtr<nsIInputStream>  mSource;
  63   nsCOMPtr<nsIOutputStream> mSink;
  64   nsCOMPtr<nsIEventTarget>  mEventTarget;
  65   nsWriteSegmentFun         mWriter;  // for implementing ReadSegments
  66   void*                     mClosure; // for implementing ReadSegments
  67   nsAutoPtr<Mutex>          mLock; // synchronize access to mSinkIsValid
  68   bool                      mSinkIsValid; // False if TeeWriteEvent fails
  69 };
  70
  71 class nsInputStreamTeeWriteEvent : public nsRunnable
  72 {
  73 public:
  74   // aTee's lock is held across construction of this object
  75   nsInputStreamTeeWriteEvent(const char* aBuf, uint32_t aCount,
  76                              nsIOutputStream* aSink, nsInputStreamTee* aTee)
  77   {
  78     // copy the buffer - will be free'd by dtor
  79     mBuf = (char*)malloc(aCount);
  80     if (mBuf) {
  81       memcpy(mBuf, (char*)aBuf, aCount);
  82     }
  83     mCount = aCount;
  84     mSink = aSink;
  85     bool isNonBlocking;
  86     mSink->IsNonBlocking(&isNonBlocking);
  87     NS_ASSERTION(isNonBlocking == false, "mSink is nonblocking");
  88     mTee = aTee;
  89   }
  90
  91   NS_IMETHOD Run()
  92   {
  93     if (!mBuf) {
  94       NS_WARNING("nsInputStreamTeeWriteEvent::Run() "
  95                  "memory not allocated\n");
  96       return NS_OK;
  97     }
  98     NS_ABORT_IF_FALSE(mSink, "mSink is null!");
  99
 100     //  The output stream could have been invalidated between when
 101     //  this event was dispatched and now, so check before writing.
 102     if (!mTee->SinkIsValid()) {
 103       return NS_OK;
 104     }
 105
 106     LOG(("nsInputStreamTeeWriteEvent::Run() [%p]"
 107          "will write %u bytes to %p\n",
 108          this, mCount, mSink.get()));
 109
 110     uint32_t totalBytesWritten = 0;
 111     while (mCount) {
 112       nsresult rv;
 113       uint32_t bytesWritten = 0;
 114       rv = mSink->Write(mBuf + totalBytesWritten, mCount, &bytesWritten);
 115       if (NS_FAILED(rv)) {
 116         LOG(("nsInputStreamTeeWriteEvent::Run[%p] error %x in writing",
 117              this, rv));
 118         mTee->InvalidateSink();
 119         break;
 120       }
 121       totalBytesWritten += bytesWritten;
 122       NS_ASSERTION(bytesWritten <= mCount, "wrote too much");
 123       mCount -= bytesWritten;
 124     }
 125     return NS_OK;
 126   }
 127
 128 protected:
 129   virtual ~nsInputStreamTeeWriteEvent()
 130   {
 131     if (mBuf) {
 132       free(mBuf);
 133     }
 134     mBuf = nullptr;
 135   }
 136
 137 private:
 138   char* mBuf;
 139   uint32_t mCount;
 140   nsCOMPtr<nsIOutputStream> mSink;
 141   // back pointer to the tee that created this runnable
 142   nsRefPtr<nsInputStreamTee> mTee;
 143 };
 144
 145 nsInputStreamTee::nsInputStreamTee(): mLock(nullptr)
 146   , mSinkIsValid(true)
 147 {
 148 }
 149
 150 bool
 151 nsInputStreamTee::SinkIsValid()
 152 {
 153   MutexAutoLock lock(*mLock);
 154   return mSinkIsValid;
 155 }
 156
 157 void
 158 nsInputStreamTee::InvalidateSink()
 159 {
 160   MutexAutoLock lock(*mLock);
 161   mSinkIsValid = false;
 162 }
 163
 164 nsresult
 165 nsInputStreamTee::TeeSegment(const char* aBuf, uint32_t aCount)
 166 {
 167   if (!mSink) {
 168     return NS_OK;  // nothing to do
 169   }
 170   if (mLock) { // asynchronous case
 171     NS_ASSERTION(mEventTarget, "mEventTarget is null, mLock is not null.");
 172     if (!SinkIsValid()) {
 173       return NS_OK; // nothing to do
 174     }
 175     nsRefPtr<nsIRunnable> event =
 176       new nsInputStreamTeeWriteEvent(aBuf, aCount, mSink, this);
 177     LOG(("nsInputStreamTee::TeeSegment [%p] dispatching write %u bytes\n",
 178          this, aCount));
 179     return mEventTarget->Dispatch(event, NS_DISPATCH_NORMAL);
 180   } else { // synchronous case
 181     NS_ASSERTION(!mEventTarget, "mEventTarget is not null, mLock is null.");
 182     nsresult rv;
 183     uint32_t totalBytesWritten = 0;
 184     while (aCount) {
 185       uint32_t bytesWritten = 0;
 186       rv = mSink->Write(aBuf + totalBytesWritten, aCount, &bytesWritten);
 187       if (NS_FAILED(rv)) {
 188         // ok, this is not a fatal error... just drop our reference to mSink
 189         // and continue on as if nothing happened.
 190         NS_WARNING("Write failed (non-fatal)");
 191         // catch possible misuse of the input stream tee
 192         NS_ASSERTION(rv != NS_BASE_STREAM_WOULD_BLOCK, "sink must be a blocking stream");
 193         mSink = 0;
 194         break;
 195       }
 196       totalBytesWritten += bytesWritten;
 197       NS_ASSERTION(bytesWritten <= aCount, "wrote too much");
 198       aCount -= bytesWritten;
 199     }
 200     return NS_OK;
 201   }
 202 }
 203
 204 NS_METHOD
 205 nsInputStreamTee::WriteSegmentFun(nsIInputStream* aIn, void* aClosure,
 206                                   const char* aFromSegment, uint32_t aOffset,
 207                                   uint32_t aCount, uint32_t* aWriteCount)
 208 {
 209   nsInputStreamTee* tee = reinterpret_cast<nsInputStreamTee*>(aClosure);
 210   nsresult rv = tee->mWriter(aIn, tee->mClosure, aFromSegment, aOffset,
 211                              aCount, aWriteCount);
 212   if (NS_FAILED(rv) || (*aWriteCount == 0)) {
 213     NS_ASSERTION((NS_FAILED(rv) ? (*aWriteCount == 0) : true),
 214                  "writer returned an error with non-zero writeCount");
 215     return rv;
 216   }
 217
 218   return tee->TeeSegment(aFromSegment, *aWriteCount);
 219 }
 220
 221 NS_IMPL_ISUPPORTS(nsInputStreamTee,
 222                   nsIInputStreamTee,
 223                   nsIInputStream)
 224 NS_IMETHODIMP
 225 nsInputStreamTee::Close()
 226 {
 227   if (NS_WARN_IF(!mSource)) {
 228     return NS_ERROR_NOT_INITIALIZED;
 229   }
 230   nsresult rv = mSource->Close();
 231   mSource = 0;
 232   mSink = 0;
 233   return rv;
 234 }
 235
 236 NS_IMETHODIMP
 237 nsInputStreamTee::Available(uint64_t* aAvail)
 238 {
 239   if (NS_WARN_IF(!mSource)) {
 240     return NS_ERROR_NOT_INITIALIZED;
 241   }
 242   return mSource->Available(aAvail);
 243 }
 244
 245 NS_IMETHODIMP
 246 nsInputStreamTee::Read(char* aBuf, uint32_t aCount, uint32_t* aBytesRead)
 247 {
 248   if (NS_WARN_IF(!mSource)) {
 249     return NS_ERROR_NOT_INITIALIZED;
 250   }
 251
 252   nsresult rv = mSource->Read(aBuf, aCount, aBytesRead);
 253   if (NS_FAILED(rv) || (*aBytesRead == 0)) {
 254     return rv;
 255   }
 256
 257   return TeeSegment(aBuf, *aBytesRead);
 258 }
 259
 260 NS_IMETHODIMP
 261 nsInputStreamTee::ReadSegments(nsWriteSegmentFun aWriter,
 262                                void* aClosure,
 263                                uint32_t aCount,
 264                                uint32_t* aBytesRead)
 265 {
 266   if (NS_WARN_IF(!mSource)) {
 267     return NS_ERROR_NOT_INITIALIZED;
 268   }
 269
 270   mWriter = aWriter;
 271   mClosure = aClosure;
 272
 273   return mSource->ReadSegments(WriteSegmentFun, this, aCount, aBytesRead);
 274 }
 275
 276 NS_IMETHODIMP
 277 nsInputStreamTee::IsNonBlocking(bool* aResult)
 278 {
 279   if (NS_WARN_IF(!mSource)) {
 280     return NS_ERROR_NOT_INITIALIZED;
 281   }
 282   return mSource->IsNonBlocking(aResult);
 283 }
 284
 285 NS_IMETHODIMP
 286 nsInputStreamTee::SetSource(nsIInputStream* aSource)
 287 {
 288   mSource = aSource;
 289   return NS_OK;
 290 }
 291
 292 NS_IMETHODIMP
 293 nsInputStreamTee::GetSource(nsIInputStream** aSource)
 294 {
 295   NS_IF_ADDREF(*aSource = mSource);
 296   return NS_OK;
 297 }
 298
 299 NS_IMETHODIMP
 300 nsInputStreamTee::SetSink(nsIOutputStream* aSink)
 301 {
 302 #ifdef DEBUG
 303   if (aSink) {
 304     bool nonBlocking;
 305     nsresult rv = aSink->IsNonBlocking(&nonBlocking);
 306     if (NS_FAILED(rv) || nonBlocking) {
 307       NS_ERROR("aSink should be a blocking stream");
 308     }
 309   }
 310 #endif
 311   mSink = aSink;
 312   return NS_OK;
 313 }
 314
 315 NS_IMETHODIMP
 316 nsInputStreamTee::GetSink(nsIOutputStream** aSink)
 317 {
 318   NS_IF_ADDREF(*aSink = mSink);
 319   return NS_OK;
 320 }
 321
 322 NS_IMETHODIMP
 323 nsInputStreamTee::SetEventTarget(nsIEventTarget* aEventTarget)
 324 {
 325   mEventTarget = aEventTarget;
 326   if (mEventTarget) {
 327     // Only need synchronization if this is an async tee
 328     mLock = new Mutex("nsInputStreamTee.mLock");
 329   }
 330   return NS_OK;
 331 }
 332
 333 NS_IMETHODIMP
 334 nsInputStreamTee::GetEventTarget(nsIEventTarget** aEventTarget)
 335 {
 336   NS_IF_ADDREF(*aEventTarget = mEventTarget);
 337   return NS_OK;
 338 }
 339
 340
 341 nsresult
 342 NS_NewInputStreamTeeAsync(nsIInputStream** aResult,
 343                           nsIInputStream* aSource,
 344                           nsIOutputStream* aSink,
 345                           nsIEventTarget* aEventTarget)
 346 {
 347   nsresult rv;
 348
 349   nsCOMPtr<nsIInputStreamTee> tee = new nsInputStreamTee();
 350   if (!tee) {
 351     return NS_ERROR_OUT_OF_MEMORY;
 352   }
 353
 354   rv = tee->SetSource(aSource);
 355   if (NS_FAILED(rv)) {
 356     return rv;
 357   }
 358
 359   rv = tee->SetSink(aSink);
 360   if (NS_FAILED(rv)) {
 361     return rv;
 362   }
 363
 364   rv = tee->SetEventTarget(aEventTarget);
 365   if (NS_FAILED(rv)) {
 366     return rv;
 367   }
 368
 369   NS_ADDREF(*aResult = tee);
 370   return rv;
 371 }
 372
 373 nsresult
 374 NS_NewInputStreamTee(nsIInputStream** aResult,
 375                      nsIInputStream* aSource,
 376                      nsIOutputStream* aSink)
 377 {
 378   return NS_NewInputStreamTeeAsync(aResult, aSource, aSink, nullptr);
 379 }
 380
 381 #undef LOG