From 5fb7a2ddb1c372bf08b1ad016bd0881c8e215445 Mon Sep 17 00:00:00 2001
From: Andrea Marchesini <amarchesini@mozilla.com>
Date: Wed, 31 Jan 2018 16:45:20 +0100
Subject: [PATCH] Bug 1371699 - NonBlockingAsyncInputStream::ReadSegments
 passes the correct stream to the writer callback, r=froydnj

---
 xpcom/io/NonBlockingAsyncInputStream.cpp      | 40 +++++++++++++++++--
 .../gtest/TestNonBlockingAsyncInputStream.cpp | 30 +++++++++++++-
 2 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/xpcom/io/NonBlockingAsyncInputStream.cpp b/xpcom/io/NonBlockingAsyncInputStream.cpp
index 0aa60092846a..570c2dd400ca 100644
--- a/xpcom/io/NonBlockingAsyncInputStream.cpp
+++ b/xpcom/io/NonBlockingAsyncInputStream.cpp
@@ -162,12 +162,46 @@ NonBlockingAsyncInputStream::Read(char* aBuffer, uint32_t aCount,
   return mInputStream->Read(aBuffer, aCount, aReadCount);
 }
 
+namespace {
+
+class MOZ_RAII ReadSegmentsData
+{
+public:
+  ReadSegmentsData(NonBlockingAsyncInputStream* aStream,
+                   nsWriteSegmentFun aFunc,
+                   void* aClosure)
+    : mStream(aStream)
+    , mFunc(aFunc)
+    , mClosure(aClosure)
+  {}
+
+  NonBlockingAsyncInputStream* mStream;
+  nsWriteSegmentFun mFunc;
+  void* mClosure;
+};
+
+nsresult
+ReadSegmentsWriter(nsIInputStream* aInStream,
+                   void* aClosure,
+                   const char* aFromSegment,
+                   uint32_t aToOffset,
+                   uint32_t aCount,
+                   uint32_t* aWriteCount)
+{
+  ReadSegmentsData* data = static_cast<ReadSegmentsData*>(aClosure);
+  return data->mFunc(data->mStream, data->mClosure, aFromSegment, aToOffset,
+                     aCount, aWriteCount);
+}
+
+} // anonymous
+
 NS_IMETHODIMP
 NonBlockingAsyncInputStream::ReadSegments(nsWriteSegmentFun aWriter,
                                           void* aClosure, uint32_t aCount,
-                                          uint32_t *aResult)
+                                          uint32_t* aResult)
 {
-  return mInputStream->ReadSegments(aWriter, aClosure, aCount, aResult);
+  ReadSegmentsData data(this, aWriter, aClosure);
+  return mInputStream->ReadSegments(ReadSegmentsWriter, &data, aCount, aResult);
 }
 
 NS_IMETHODIMP
@@ -283,7 +317,7 @@ NonBlockingAsyncInputStream::Seek(int32_t aWhence, int64_t aOffset)
 }
 
 NS_IMETHODIMP
-NonBlockingAsyncInputStream::Tell(int64_t *aResult)
+NonBlockingAsyncInputStream::Tell(int64_t* aResult)
 {
   NS_ENSURE_STATE(mWeakSeekableInputStream);
   return mWeakSeekableInputStream->Tell(aResult);
diff --git a/xpcom/tests/gtest/TestNonBlockingAsyncInputStream.cpp b/xpcom/tests/gtest/TestNonBlockingAsyncInputStream.cpp
index 7542ee8a66fc..6abe8c656e66 100644
--- a/xpcom/tests/gtest/TestNonBlockingAsyncInputStream.cpp
+++ b/xpcom/tests/gtest/TestNonBlockingAsyncInputStream.cpp
@@ -47,6 +47,33 @@ TEST(TestNonBlockingAsyncInputStream, Simple) {
   ASSERT_TRUE(data.Equals(nsCString(buffer, read)));
 }
 
+class ReadSegmentsData
+{
+public:
+  ReadSegmentsData(nsIInputStream* aStream, char* aBuffer)
+    : mStream(aStream)
+    , mBuffer(aBuffer)
+  {}
+
+  nsIInputStream* mStream;
+  char* mBuffer;
+};
+
+nsresult
+ReadSegmentsFunction(nsIInputStream* aInStr,
+                     void* aClosure,
+                     const char* aBuffer,
+                     uint32_t aOffset,
+                     uint32_t aCount,
+                     uint32_t* aCountWritten)
+{
+  ReadSegmentsData* data = static_cast<ReadSegmentsData*>(aClosure);
+  if (aInStr != data->mStream) return NS_ERROR_FAILURE;
+  memcpy(&data->mBuffer[aOffset], aBuffer, aCount);
+  *aCountWritten = aCount;
+  return NS_OK;
+}
+
 TEST(TestNonBlockingAsyncInputStream, ReadSegments) {
   nsCOMPtr<nsIInputStream> stream;
 
@@ -64,7 +91,8 @@ TEST(TestNonBlockingAsyncInputStream, ReadSegments) {
   // Read works fine.
   char buffer[1024];
   uint32_t read = 0;
-  ASSERT_EQ(NS_OK, async->ReadSegments(NS_CopySegmentToBuffer, buffer,
+  ReadSegmentsData closure(async, buffer);
+  ASSERT_EQ(NS_OK, async->ReadSegments(ReadSegmentsFunction, &closure,
                                        sizeof(buffer), &read));
   ASSERT_EQ(data.Length(), read);
   ASSERT_TRUE(data.Equals(nsCString(buffer, read)));