Bug 1287367 - Allow users of StreamingLexer to detect and handle truncation. r=njn

2016-07-17 22:51:19 -07:00 · 2016-07-17 22:51:19 -07:00 · 666fc943cb
--- a/image/StreamingLexer.h
+++ b/image/StreamingLexer.h
@ -271,7 +271,8 @@ private:
 *    read.
 *
 *  - Add an instance of StreamingLexer<State> to your decoder class. Initialize
- *    it with a Transition::To() the state that you want to start lexing in.
+ *    it with a Transition::To() the state that you want to start lexing in, and
+ *    a Transition::To() the state you'd like to use to handle truncated data.
 *
 *  - In your decoder's DoDecode() method, call Lex(), passing in the input
 *    data and length that are passed to DoDecode(). You also need to pass
@ -295,6 +296,17 @@ private:
 * TerminalState::FAILURE, and you can check which one to determine if lexing
 * succeeded or failed and do any necessary cleanup.
 *
+ * Sometimes, the input data is truncated. StreamingLexer will notify you when
+ * this happens by invoking the truncated data state you passed to the
+ * constructor. At this point you can attempt to recover and return
+ * TerminalState::SUCCESS or TerminalState::FAILURE, depending on whether you
+ * were successful. Note that you can't return anything other than a terminal
+ * state in this situation, since there's no more data to read. For the same
+ * reason, your truncated data state shouldn't require any data. (That is, the
+ * @aSize argument you pass to Transition::To() must be zero.) Violating these
+ * requirements will trigger assertions and an immediate transition to
+ * TerminalState::FAILURE.
+ *
 * Some lexers may want to *avoid* buffering in some cases, and just process the
 * data as it comes in. This is useful if, for example, you just want to skip
 * over a large section of data; there's no point in buffering data you're just
@ -348,8 +360,10 @@ template <typename State, size_t InlineBufferSize = 16>
 class StreamingLexer
 {
 public:
-  explicit StreamingLexer(LexerTransition<State> aStartState)
+  StreamingLexer(LexerTransition<State> aStartState,
+                 LexerTransition<State> aTruncatedState)
    : mTransition(TerminalState::FAILURE)
+    , mTruncatedTransition(aTruncatedState)
  {
    if (!aStartState.NextStateIsTerminal() &&
        aStartState.ControlFlow() == ControlFlowStrategy::YIELD) {
@ -364,6 +378,18 @@ public:
      return;
    }

+    if (!aTruncatedState.NextStateIsTerminal() &&
+          (aTruncatedState.ControlFlow() == ControlFlowStrategy::YIELD ||
+           aTruncatedState.Buffering() == BufferingStrategy::UNBUFFERED ||
+           aTruncatedState.Size() != 0)) {
+      // The truncated state can't receive any data because, by definition,
+      // there is no more data to receive. That means that yielding or an
+      // unbuffered read would not make sense, and that the state must require
+      // zero bytes.
+      MOZ_ASSERT_UNREACHABLE("Truncated state makes no sense");
+      return;
+    }
+
    SetTransition(aStartState);
  }

@ -412,14 +438,11 @@ public:
          break;

        case SourceBufferIterator::COMPLETE:
-          // Normally even if the data is truncated, we want decoding to
-          // succeed so we can display whatever we got. However, if the
-          // SourceBuffer was completed with a failing status, we want to fail.
-          // This happens only in exceptional situations like SourceBuffer
-          // itself encountering a failure due to OOM.
-          result = SetTransition(NS_SUCCEEDED(aIterator.CompletionStatus())
-                 ? Transition::TerminateSuccess()
-                 : Transition::TerminateFailure());
+          // The data is truncated; if not, the lexer would've reached a
+          // terminal state by now. We only get to
+          // SourceBufferIterator::COMPLETE after every byte of data has been
+          // delivered to the lexer.
+          result = Truncated(aIterator, aFunc);
          break;

        case SourceBufferIterator::READY:
@ -629,6 +652,32 @@ private:
    return SetTransition(Transition::TerminateFailure());
  }

+  template <typename Func>
+  Maybe<LexerResult> Truncated(SourceBufferIterator& aIterator,
+                               Func aFunc)
+  {
+    // The data is truncated. Let the lexer clean up and decide which terminal
+    // state we should end up in.
+    LexerTransition<State> transition
+      = mTruncatedTransition.NextStateIsTerminal()
+      ? mTruncatedTransition
+      : aFunc(mTruncatedTransition.NextState(), nullptr, 0);
+
+    if (!transition.NextStateIsTerminal()) {
+      MOZ_ASSERT_UNREACHABLE("Truncated state didn't lead to terminal state?");
+      return SetTransition(Transition::TerminateFailure());
+    }
+
+    // If the SourceBuffer was completed with a failing state, we end in
+    // TerminalState::FAILURE no matter what. This only happens in exceptional
+    // situations like SourceBuffer itself encountering a failure due to OOM.
+    if (NS_FAILED(aIterator.CompletionStatus())) {
+      return SetTransition(Transition::TerminateFailure());
+    }
+
+    return SetTransition(transition);
+  }
+
  Maybe<LexerResult> SetTransition(const LexerTransition<State>& aTransition)
  {
    // There should be no transitions while we're buffering for a buffered read
@ -690,6 +739,7 @@ private:

  Vector<char, InlineBufferSize> mBuffer;
  LexerTransition<State> mTransition;
+  const LexerTransition<State> mTruncatedTransition;
  Maybe<State> mYieldingToState;
  Maybe<UnbufferedState> mUnbufferedState;
 };
--- a/image/decoders/nsBMPDecoder.cpp
+++ b/image/decoders/nsBMPDecoder.cpp
@ -170,7 +170,7 @@ static const uint32_t BIHSIZE_FIELD_LENGTH = 4;

 nsBMPDecoder::nsBMPDecoder(RasterImage* aImage, State aState, size_t aLength)
  : Decoder(aImage)
-  , mLexer(Transition::To(aState, aLength))
+  , mLexer(Transition::To(aState, aLength), Transition::TerminateSuccess())
  , mIsWithinICO(false)
  , mMayHaveTransparency(false)
  , mDoesHaveTransparency(false)
--- a/image/decoders/nsGIFDecoder2.cpp
+++ b/image/decoders/nsGIFDecoder2.cpp
@ -81,7 +81,8 @@ static const uint8_t PACKED_FIELDS_TABLE_DEPTH_MASK = 0x07;

 nsGIFDecoder2::nsGIFDecoder2(RasterImage* aImage)
  : Decoder(aImage)
-  , mLexer(Transition::To(State::GIF_HEADER, GIF_HEADER_LEN))
+  , mLexer(Transition::To(State::GIF_HEADER, GIF_HEADER_LEN),
+           Transition::TerminateSuccess())
  , mOldColor(0)
  , mCurrentFrameIndex(-1)
  , mColorTablePos(0)
--- a/image/decoders/nsICODecoder.cpp
+++ b/image/decoders/nsICODecoder.cpp
@ -53,7 +53,8 @@ nsICODecoder::GetNumColors()

 nsICODecoder::nsICODecoder(RasterImage* aImage)
  : Decoder(aImage)
-  , mLexer(Transition::To(ICOState::HEADER, ICOHEADERSIZE))
+  , mLexer(Transition::To(ICOState::HEADER, ICOHEADERSIZE),
+           Transition::TerminateSuccess())
  , mBiggestResourceColorDepth(0)
  , mBestResourceDelta(INT_MIN)
  , mBestResourceColorDepth(0)
--- a/image/decoders/nsIconDecoder.cpp
+++ b/image/decoders/nsIconDecoder.cpp
@ -17,7 +17,8 @@ static const uint32_t ICON_HEADER_SIZE = 2;

 nsIconDecoder::nsIconDecoder(RasterImage* aImage)
 : Decoder(aImage)
- , mLexer(Transition::To(State::HEADER, ICON_HEADER_SIZE))
+ , mLexer(Transition::To(State::HEADER, ICON_HEADER_SIZE),
+          Transition::TerminateSuccess())
 , mBytesPerRow()   // set by ReadHeader()
 {
  // Nothing to do
--- a/image/decoders/nsJPEGDecoder.cpp
+++ b/image/decoders/nsJPEGDecoder.cpp
@ -74,7 +74,8 @@ nsJPEGDecoder::nsJPEGDecoder(RasterImage* aImage,
 : Decoder(aImage)
 , mLexer(Transition::ToUnbuffered(State::FINISHED_JPEG_DATA,
                                   State::JPEG_DATA,
-                                   SIZE_MAX))
+                                   SIZE_MAX),
+          Transition::TerminateSuccess())
 , mDecodeStyle(aDecodeStyle)
 , mSampleSize(0)
 {
--- a/image/decoders/nsPNGDecoder.cpp
+++ b/image/decoders/nsPNGDecoder.cpp
@ -98,7 +98,8 @@ nsPNGDecoder::nsPNGDecoder(RasterImage* aImage)
 : Decoder(aImage)
 , mLexer(Transition::ToUnbuffered(State::FINISHED_PNG_DATA,
                                   State::PNG_DATA,
-                                   SIZE_MAX))
+                                   SIZE_MAX),
+          Transition::TerminateSuccess())
 , mPNG(nullptr)
 , mInfo(nullptr)
 , mCMSLine(nullptr)
--- a/image/test/gtest/TestStreamingLexer.cpp
+++ b/image/test/gtest/TestStreamingLexer.cpp
@ -15,7 +15,9 @@ enum class TestState
  ONE,
  TWO,
  THREE,
-  UNBUFFERED
+  UNBUFFERED,
+  TRUNCATED_SUCCESS,
+  TRUNCATED_FAILURE
 };

 void
@ -44,6 +46,10 @@ DoLex(TestState aState, const char* aData, size_t aLength)
    case TestState::THREE:
      CheckLexedData(aData, aLength, 6, 3);
      return Transition::TerminateSuccess();
+    case TestState::TRUNCATED_SUCCESS:
+      return Transition::TerminateSuccess();
+    case TestState::TRUNCATED_FAILURE:
+      return Transition::TerminateFailure();
    default:
      MOZ_CRASH("Unexpected or unhandled TestState");
  }
@ -221,8 +227,12 @@ DoLexWithZeroLengthStatesAfterUnbuffered(TestState aState,
 class ImageStreamingLexer : public ::testing::Test
 {
 public:
+  // Note that mLexer is configured to enter TerminalState::FAILURE immediately
+  // if the input data is truncated. We don't expect that to happen in most
+  // tests, so we want to detect that issue. If a test needs a different
+  // behavior, we create a special StreamingLexer just for that test.
  ImageStreamingLexer()
-    : mLexer(Transition::To(TestState::ONE, 3))
+    : mLexer(Transition::To(TestState::ONE, 3), Transition::TerminateFailure())
    , mSourceBuffer(new SourceBuffer)
    , mIterator(mSourceBuffer->Iterator())
    , mExpectNoResume(new ExpectNoResume)
@ -230,6 +240,31 @@ public:
  { }

 protected:
+  void CheckTruncatedState(StreamingLexer<TestState>& aLexer,
+                           TerminalState aExpectedTerminalState,
+                           nsresult aCompletionStatus = NS_OK)
+  {
+    for (unsigned i = 0; i < 9; ++i) {
+      if (i < 2) {
+        mSourceBuffer->Append(mData + i, 1);
+      } else if (i == 2) {
+        mSourceBuffer->Complete(aCompletionStatus);
+      }
+
+      LexerResult result = aLexer.Lex(mIterator, mCountResumes, DoLex);
+
+      if (i >= 2) {
+        EXPECT_TRUE(result.is<TerminalState>());
+        EXPECT_EQ(aExpectedTerminalState, result.as<TerminalState>());
+      } else {
+        EXPECT_TRUE(result.is<Yield>());
+        EXPECT_EQ(Yield::NEED_MORE_DATA, result.as<Yield>());
+      }
+    }
+
+    EXPECT_EQ(2u, mCountResumes->Count());
+  }
+
  AutoInitializeImageLib mInit;
  const char mData[9] { 1, 2, 3, 4, 5, 6, 7, 8, 9 };
  StreamingLexer<TestState> mLexer;
@ -259,7 +294,8 @@ TEST_F(ImageStreamingLexer, ZeroLengthDataUnbuffered)
  // state to be unbuffered.
  StreamingLexer<TestState> lexer(Transition::ToUnbuffered(TestState::ONE,
                                                           TestState::UNBUFFERED,
-                                                           sizeof(mData)));
+                                                           sizeof(mData)),
+                                  Transition::TerminateFailure());

  LexerResult result = lexer.Lex(mIterator, mExpectNoResume, DoLex);
  EXPECT_TRUE(result.is<TerminalState>());
@ -271,7 +307,8 @@ TEST_F(ImageStreamingLexer, StartWithTerminal)
  // Create a special StreamingLexer for this test because we want the first
  // state to be a terminal state. This doesn't really make sense, but we should
  // handle it.
-  StreamingLexer<TestState> lexer(Transition::TerminateSuccess());
+  StreamingLexer<TestState> lexer(Transition::TerminateSuccess(),
+                                  Transition::TerminateFailure());
  LexerResult result = lexer.Lex(mIterator, mExpectNoResume, DoLex);
  EXPECT_TRUE(result.is<TerminalState>());
  EXPECT_EQ(TerminalState::SUCCESS, result.as<TerminalState>());
@ -557,7 +594,8 @@ TEST_F(ImageStreamingLexer, ZeroLengthState)

  // Create a special StreamingLexer for this test because we want the first
  // state to be zero length.
-  StreamingLexer<TestState> lexer(Transition::To(TestState::ONE, 0));
+  StreamingLexer<TestState> lexer(Transition::To(TestState::ONE, 0),
+                                  Transition::TerminateFailure());

  LexerResult result =
    lexer.Lex(mIterator, mExpectNoResume, DoLexWithZeroLengthStates);
@ -573,7 +611,8 @@ TEST_F(ImageStreamingLexer, ZeroLengthStatesAtEnd)

  // Create a special StreamingLexer for this test because we want the first
  // state to consume the full input.
-  StreamingLexer<TestState> lexer(Transition::To(TestState::ONE, 9));
+  StreamingLexer<TestState> lexer(Transition::To(TestState::ONE, 9),
+                                  Transition::TerminateFailure());

  LexerResult result =
    lexer.Lex(mIterator, mExpectNoResume, DoLexWithZeroLengthStatesAtEnd);
@ -586,7 +625,8 @@ TEST_F(ImageStreamingLexer, ZeroLengthStateWithYield)
 {
  // Create a special StreamingLexer for this test because we want the first
  // state to be zero length.
-  StreamingLexer<TestState> lexer(Transition::To(TestState::ONE, 0));
+  StreamingLexer<TestState> lexer(Transition::To(TestState::ONE, 0),
+                                  Transition::TerminateFailure());

  mSourceBuffer->Append(mData, 3);
  LexerResult result =
@ -615,7 +655,8 @@ TEST_F(ImageStreamingLexer, ZeroLengthStateWithUnbuffered)
  // state to be both zero length and unbuffered.
  StreamingLexer<TestState> lexer(Transition::ToUnbuffered(TestState::ONE,
                                                           TestState::UNBUFFERED,
-                                                           0));
+                                                           0),
+                                  Transition::TerminateFailure());

  LexerResult result =
    lexer.Lex(mIterator, mExpectNoResume, DoLexWithZeroLengthStatesUnbuffered);
@ -631,7 +672,8 @@ TEST_F(ImageStreamingLexer, ZeroLengthStateAfterUnbuffered)

  // Create a special StreamingLexer for this test because we want the first
  // state to be zero length.
-  StreamingLexer<TestState> lexer(Transition::To(TestState::ONE, 0));
+  StreamingLexer<TestState> lexer(Transition::To(TestState::ONE, 0),
+                                  Transition::TerminateFailure());

  LexerResult result =
    lexer.Lex(mIterator, mExpectNoResume, DoLexWithZeroLengthStatesAfterUnbuffered);
@ -696,7 +738,8 @@ TEST_F(ImageStreamingLexer, ZeroLengthStateWithUnbufferedYield)
  // state to be unbuffered.
  StreamingLexer<TestState> lexer(Transition::ToUnbuffered(TestState::ONE,
                                                           TestState::UNBUFFERED,
-                                                           sizeof(mData)));
+                                                           sizeof(mData)),
+                                  Transition::TerminateFailure());

  mSourceBuffer->Append(mData, 3);
  LexerResult result = lexer.Lex(mIterator, mExpectNoResume, lexerFunc);
@ -854,54 +897,60 @@ TEST_F(ImageStreamingLexer, SourceBufferImmediateComplete)
  EXPECT_EQ(TerminalState::FAILURE, result.as<TerminalState>());
 }

-TEST_F(ImageStreamingLexer, SourceBufferTruncatedSuccess)
+TEST_F(ImageStreamingLexer, SourceBufferTruncatedTerminalStateSuccess)
 {
-  // Test that calling SourceBuffer::Complete() with a successful status results
-  // in an immediate TerminalState::SUCCESS result.
-  for (unsigned i = 0; i < 9; ++i) {
-    if (i < 2) {
-      mSourceBuffer->Append(mData + i, 1);
-    } else if (i == 2) {
-      mSourceBuffer->Complete(NS_OK);
-    }
+  // Test that using a terminal state (in this case TerminalState::SUCCESS) as a
+  // truncated state works.
+  StreamingLexer<TestState> lexer(Transition::To(TestState::ONE, 3),
+                                  Transition::TerminateSuccess());

-    LexerResult result = mLexer.Lex(mIterator, mCountResumes, DoLex);
-
-    if (i >= 2) {
-      EXPECT_TRUE(result.is<TerminalState>());
-      EXPECT_EQ(TerminalState::SUCCESS, result.as<TerminalState>());
-    } else {
-      EXPECT_TRUE(result.is<Yield>());
-      EXPECT_EQ(Yield::NEED_MORE_DATA, result.as<Yield>());
-    }
-  }
-
-  EXPECT_EQ(2u, mCountResumes->Count());
+  CheckTruncatedState(lexer, TerminalState::SUCCESS);
 }

-TEST_F(ImageStreamingLexer, SourceBufferTruncatedFailure)
+TEST_F(ImageStreamingLexer, SourceBufferTruncatedTerminalStateFailure)
+{
+  // Test that using a terminal state (in this case TerminalState::FAILURE) as a
+  // truncated state works.
+  StreamingLexer<TestState> lexer(Transition::To(TestState::ONE, 3),
+                                  Transition::TerminateFailure());
+
+  CheckTruncatedState(lexer, TerminalState::FAILURE);
+}
+
+TEST_F(ImageStreamingLexer, SourceBufferTruncatedStateReturningSuccess)
+{
+  // Test that a truncated state that returns TerminalState::SUCCESS works. When
+  // |lexer| discovers that the data is truncated, it invokes the
+  // TRUNCATED_SUCCESS state, which returns TerminalState::SUCCESS.
+  // CheckTruncatedState() verifies that this happens.
+  StreamingLexer<TestState> lexer(Transition::To(TestState::ONE, 3),
+                                  Transition::To(TestState::TRUNCATED_SUCCESS, 0));
+
+  CheckTruncatedState(lexer, TerminalState::SUCCESS);
+}
+
+TEST_F(ImageStreamingLexer, SourceBufferTruncatedStateReturningFailure)
+{
+  // Test that a truncated state that returns TerminalState::FAILURE works. When
+  // |lexer| discovers that the data is truncated, it invokes the
+  // TRUNCATED_FAILURE state, which returns TerminalState::FAILURE.
+  // CheckTruncatedState() verifies that this happens.
+  StreamingLexer<TestState> lexer(Transition::To(TestState::ONE, 3),
+                                  Transition::To(TestState::TRUNCATED_FAILURE, 0));
+
+  CheckTruncatedState(lexer, TerminalState::FAILURE);
+}
+
+TEST_F(ImageStreamingLexer, SourceBufferTruncatedFailingCompleteStatus)
 {
  // Test that calling SourceBuffer::Complete() with a failing status results in
-  // an immediate TerminalState::FAILURE result.
-  for (unsigned i = 0; i < 9; ++i) {
-    if (i < 2) {
-      mSourceBuffer->Append(mData + i, 1);
-    } else if (i == 2) {
-      mSourceBuffer->Complete(NS_ERROR_FAILURE);
-    }
+  // an immediate TerminalState::FAILURE result. (Note that |lexer|'s truncated
+  // state is TerminalState::SUCCESS, so if we ignore the failing status, the
+  // test will fail.)
+  StreamingLexer<TestState> lexer(Transition::To(TestState::ONE, 3),
+                                  Transition::TerminateSuccess());

-    LexerResult result = mLexer.Lex(mIterator, mCountResumes, DoLex);
-
-    if (i >= 2) {
-      EXPECT_TRUE(result.is<TerminalState>());
-      EXPECT_EQ(TerminalState::FAILURE, result.as<TerminalState>());
-    } else {
-      EXPECT_TRUE(result.is<Yield>());
-      EXPECT_EQ(Yield::NEED_MORE_DATA, result.as<Yield>());
-    }
-  }
-
-  EXPECT_EQ(2u, mCountResumes->Count());
+  CheckTruncatedState(lexer, TerminalState::FAILURE, NS_ERROR_FAILURE);
 }

 TEST_F(ImageStreamingLexer, NoSourceBufferResumable)