Bug 1434553 - Implement nsIInputStreamLength and nsIAsyncInputStreamLength - part 8 - PartiallySeekableInputStream exposes nsIInputStreamLength, r=mayhemer

netwerk/base/PartiallySeekableInputStream.cpp

@@ -25,6 +25,12 @@ NS_INTERFACE_MAP_BEGIN(PartiallySeekableInputStream)
                                      mWeakAsyncInputStream)
   NS_INTERFACE_MAP_ENTRY_CONDITIONAL(nsIInputStreamCallback,
                                      mWeakAsyncInputStream)
+  NS_INTERFACE_MAP_ENTRY_CONDITIONAL(nsIInputStreamLength,
+                                     mWeakInputStreamLength)
+  NS_INTERFACE_MAP_ENTRY_CONDITIONAL(nsIAsyncInputStreamLength,
+                                     mWeakAsyncInputStreamLength)
+  NS_INTERFACE_MAP_ENTRY_CONDITIONAL(nsIInputStreamLengthCallback,
+                                     mWeakAsyncInputStreamLength)
   NS_INTERFACE_MAP_ENTRY_AMBIGUOUS(nsISupports, nsIInputStream)
 NS_INTERFACE_MAP_END
 
@@ -34,6 +40,8 @@ PartiallySeekableInputStream::PartiallySeekableInputStream(already_AddRefed<nsII
   , mWeakCloneableInputStream(nullptr)
   , mWeakIPCSerializableInputStream(nullptr)
   , mWeakAsyncInputStream(nullptr)
+  , mWeakInputStreamLength(nullptr)
+  , mWeakAsyncInputStreamLength(nullptr)
   , mBufferSize(aBufferSize)
   , mPos(0)
   , mClosed(false)
@@ -85,6 +93,19 @@ PartiallySeekableInputStream::Init()
   if (asyncInputStream && SameCOMIdentity(mInputStream, asyncInputStream)) {
     mWeakAsyncInputStream = asyncInputStream;
   }
+
+  nsCOMPtr<nsIInputStreamLength> inputStreamLength =
+    do_QueryInterface(mInputStream);
+  if (inputStreamLength && SameCOMIdentity(mInputStream, inputStreamLength)) {
+    mWeakInputStreamLength = inputStreamLength;
+  }
+
+  nsCOMPtr<nsIAsyncInputStreamLength> asyncInputStreamLength =
+    do_QueryInterface(mInputStream);
+  if (asyncInputStreamLength &&
+      SameCOMIdentity(mInputStream, asyncInputStreamLength)) {
+    mWeakAsyncInputStreamLength = asyncInputStreamLength;
+  }
 }
 
 NS_IMETHODIMP
@@ -352,5 +373,49 @@ PartiallySeekableInputStream::SetEOF()
   return Close();
 }
 
+// nsIInputStreamLength
+
+NS_IMETHODIMP
+PartiallySeekableInputStream::Length(int64_t* aLength)
+{
+  NS_ENSURE_STATE(mWeakInputStreamLength);
+  return mWeakInputStreamLength->Length(aLength);
+}
+
+// nsIAsyncInputStreamLength
+
+NS_IMETHODIMP
+PartiallySeekableInputStream::AsyncLengthWait(nsIInputStreamLengthCallback* aCallback,
+                                              nsIEventTarget* aEventTarget)
+{
+  NS_ENSURE_STATE(mWeakAsyncInputStreamLength);
+
+  nsCOMPtr<nsIInputStreamLengthCallback> callback = aCallback ? this : nullptr;
+  {
+    MutexAutoLock lock(mMutex);
+    mAsyncInputStreamLengthCallback = aCallback;
+  }
+
+  return mWeakAsyncInputStreamLength->AsyncLengthWait(callback, aEventTarget);
+}
+
+NS_IMETHODIMP
+PartiallySeekableInputStream::OnInputStreamLengthReady(nsIAsyncInputStreamLength* aStream,
+                                                       int64_t aLength)
+{
+  nsCOMPtr<nsIInputStreamLengthCallback> callback;
+  {
+    MutexAutoLock lock(mMutex);
+    // We have been canceled in the meanwhile.
+    if (!mAsyncInputStreamLengthCallback) {
+      return NS_OK;
+    }
+
+    callback.swap(mAsyncInputStreamLengthCallback);
+  }
+
+  return callback->OnInputStreamLengthReady(this, aLength);
+}
+
 } // net namespace
 } // mozilla namespace

netwerk/base/PartiallySeekableInputStream.h

@@ -11,6 +11,7 @@
 #include "nsCOMPtr.h"
 #include "nsIAsyncInputStream.h"
 #include "nsICloneableInputStream.h"
+#include "nsIInputStreamLength.h"
 #include "nsIIPCSerializableInputStream.h"
 #include "nsISeekableStream.h"
 
@@ -25,6 +26,9 @@ class PartiallySeekableInputStream final : public nsISeekableStream
                                          , public nsICloneableInputStream
                                          , public nsIIPCSerializableInputStream
                                          , public nsIInputStreamCallback
+                                         , public nsIInputStreamLength
+                                         , public nsIAsyncInputStreamLength
+                                         , public nsIInputStreamLengthCallback
 {
 public:
   NS_DECL_THREADSAFE_ISUPPORTS
@@ -34,6 +38,9 @@ public:
   NS_DECL_NSICLONEABLEINPUTSTREAM
   NS_DECL_NSIIPCSERIALIZABLEINPUTSTREAM
   NS_DECL_NSIINPUTSTREAMCALLBACK
+  NS_DECL_NSIINPUTSTREAMLENGTH
+  NS_DECL_NSIASYNCINPUTSTREAMLENGTH
+  NS_DECL_NSIINPUTSTREAMLENGTHCALLBACK
 
   explicit PartiallySeekableInputStream(already_AddRefed<nsIInputStream> aInputStream,
                                         uint64_t aBufferSize = 4096);
@@ -53,10 +60,15 @@ private:
   nsICloneableInputStream* mWeakCloneableInputStream;
   nsIIPCSerializableInputStream* mWeakIPCSerializableInputStream;
   nsIAsyncInputStream* mWeakAsyncInputStream;
+  nsIInputStreamLength* mWeakInputStreamLength;
+  nsIAsyncInputStreamLength* mWeakAsyncInputStreamLength;
 
   // Protected by mutex.
   nsCOMPtr<nsIInputStreamCallback> mAsyncWaitCallback;
 
+  // Protected by mutex.
+  nsCOMPtr<nsIInputStreamLengthCallback> mAsyncInputStreamLengthCallback;
+
   nsTArray<char> mCachedBuffer;
 
   uint64_t mBufferSize;

netwerk/test/gtest/TestPartiallySeekableInputStream.cpp

@@ -1,5 +1,6 @@
 #include "gtest/gtest.h"
 
+#include "Helpers.h"
 #include "nsCOMPtr.h"
 #include "nsIPipe.h"
 #include "nsStreamUtils.h"
@@ -7,7 +8,9 @@
 #include "nsStringStream.h"
 #include "mozilla/net/PartiallySeekableInputStream.h"
 
+using mozilla::GetCurrentThreadSerialEventTarget;
 using mozilla::net::PartiallySeekableInputStream;
+using mozilla::SpinEventLoopUntil;
 
 class NonSeekableStream final : public nsIInputStream
 {
@@ -223,3 +226,140 @@ TEST(TestPartiallySeekableInputStream, FullCachedSeek) {
   ASSERT_EQ(NS_OK, psi->Available(&length));
   ASSERT_EQ((uint64_t)0, length);
 }
+
+TEST(TestPartiallySeekableInputStream, QIInputStreamLength) {
+  nsCString buf;
+  buf.AssignLiteral("Hello world");
+
+  for (int i = 0; i < 4; i++) {
+    nsCOMPtr<nsIInputStream> psis;
+    {
+      RefPtr<testing::LengthInputStream> stream =
+        new testing::LengthInputStream(buf, i % 2, i > 1);
+      psis = new PartiallySeekableInputStream(stream.forget());
+    }
+
+    {
+      nsCOMPtr<nsIInputStreamLength> qi = do_QueryInterface(psis);
+      ASSERT_EQ(!!(i % 2), !!qi);
+    }
+
+    {
+      nsCOMPtr<nsIAsyncInputStreamLength> qi = do_QueryInterface(psis);
+      ASSERT_EQ(i > 1, !!qi);
+    }
+  }
+}
+
+TEST(TestPartiallySeekableInputStream, InputStreamLength) {
+  nsCString buf;
+  buf.AssignLiteral("Hello world");
+
+  nsCOMPtr<nsIInputStream> psis;
+  {
+    RefPtr<testing::LengthInputStream> stream =
+      new testing::LengthInputStream(buf, true, false);
+    psis = new PartiallySeekableInputStream(stream.forget());
+  }
+
+  nsCOMPtr<nsIInputStreamLength> qi = do_QueryInterface(psis);
+  ASSERT_TRUE(!!qi);
+
+  int64_t size;
+  nsresult rv = qi->Length(&size);
+  ASSERT_EQ(NS_OK, rv);
+  ASSERT_EQ(buf.Length(), size);
+}
+
+TEST(TestPartiallySeekableInputStream, NegativeInputStreamLength) {
+  nsCString buf;
+  buf.AssignLiteral("Hello world");
+
+  nsCOMPtr<nsIInputStream> psis;
+  {
+    RefPtr<testing::LengthInputStream> stream =
+      new testing::LengthInputStream(buf, true, false, NS_OK, true);
+    psis = new PartiallySeekableInputStream(stream.forget());
+  }
+
+  nsCOMPtr<nsIInputStreamLength> qi = do_QueryInterface(psis);
+  ASSERT_TRUE(!!qi);
+
+  int64_t size;
+  nsresult rv = qi->Length(&size);
+  ASSERT_EQ(NS_OK, rv);
+  ASSERT_EQ(-1, size);
+}
+
+TEST(TestPartiallySeekableInputStream, AsyncInputStreamLength) {
+  nsCString buf;
+  buf.AssignLiteral("Hello world");
+
+  nsCOMPtr<nsIInputStream> psis;
+  {
+    RefPtr<testing::LengthInputStream> stream =
+      new testing::LengthInputStream(buf, false, true);
+    psis = new PartiallySeekableInputStream(stream.forget());
+  }
+
+  nsCOMPtr<nsIAsyncInputStreamLength> qi = do_QueryInterface(psis);
+  ASSERT_TRUE(!!qi);
+
+  RefPtr<testing::LengthCallback> callback = new testing::LengthCallback();
+
+  nsresult rv = qi->AsyncLengthWait(callback, GetCurrentThreadSerialEventTarget());
+  ASSERT_EQ(NS_OK, rv);
+
+  MOZ_ALWAYS_TRUE(SpinEventLoopUntil([&]() { return callback->Called(); }));
+  ASSERT_EQ(buf.Length(), callback->Size());
+}
+
+TEST(TestPartiallySeekableInputStream, NegativeAsyncInputStreamLength) {
+  nsCString buf;
+  buf.AssignLiteral("Hello world");
+
+  nsCOMPtr<nsIInputStream> psis;
+  {
+    RefPtr<testing::LengthInputStream> stream =
+      new testing::LengthInputStream(buf, false, true, NS_OK, true);
+    psis = new PartiallySeekableInputStream(stream.forget());
+  }
+
+  nsCOMPtr<nsIAsyncInputStreamLength> qi = do_QueryInterface(psis);
+  ASSERT_TRUE(!!qi);
+
+  RefPtr<testing::LengthCallback> callback = new testing::LengthCallback();
+
+  nsresult rv = qi->AsyncLengthWait(callback, GetCurrentThreadSerialEventTarget());
+  ASSERT_EQ(NS_OK, rv);
+
+  MOZ_ALWAYS_TRUE(SpinEventLoopUntil([&]() { return callback->Called(); }));
+  ASSERT_EQ(-1, callback->Size());
+}
+
+TEST(TestPartiallySeekableInputStream, AbortLengthCallback) {
+  nsCString buf;
+  buf.AssignLiteral("Hello world");
+
+  nsCOMPtr<nsIInputStream> psis;
+  {
+    RefPtr<testing::LengthInputStream> stream =
+      new testing::LengthInputStream(buf, false, true, NS_OK, true);
+    psis = new PartiallySeekableInputStream(stream.forget());
+  }
+
+  nsCOMPtr<nsIAsyncInputStreamLength> qi = do_QueryInterface(psis);
+  ASSERT_TRUE(!!qi);
+
+  RefPtr<testing::LengthCallback> callback1 = new testing::LengthCallback();
+  nsresult rv = qi->AsyncLengthWait(callback1, GetCurrentThreadSerialEventTarget());
+  ASSERT_EQ(NS_OK, rv);
+
+  RefPtr<testing::LengthCallback> callback2 = new testing::LengthCallback();
+  rv = qi->AsyncLengthWait(callback2, GetCurrentThreadSerialEventTarget());
+  ASSERT_EQ(NS_OK, rv);
+
+  MOZ_ALWAYS_TRUE(SpinEventLoopUntil([&]() { return callback2->Called(); }));
+  ASSERT_TRUE(!callback1->Called());
+  ASSERT_EQ(-1, callback2->Size());
+}