Bug 1287367 - Allow users of StreamingLexer to detect and handle truncation. r=njn
authorSeth Fowler <mark.seth.fowler@gmail.com>
Sun, 17 Jul 2016 22:51:19 -0700
changeset 345713 d9b88cb3db769b183fb5ea8ba19beaac562153f1
parent 345712 b0724a9e58bf061c965d46d8f96f1d422235137e
child 345714 ad32bf1994e92f6cbe3e54f30e5f8f804f9efb44
push id6389
push userraliiev@mozilla.com
push dateMon, 19 Sep 2016 13:38:22 +0000
treeherdermozilla-beta@01d67bfe6c81 [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersnjn
bugs1287367
milestone50.0a1
first release with
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
last release without
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
Bug 1287367 - Allow users of StreamingLexer to detect and handle truncation. r=njn
image/StreamingLexer.h
image/decoders/nsBMPDecoder.cpp
image/decoders/nsGIFDecoder2.cpp
image/decoders/nsICODecoder.cpp
image/decoders/nsIconDecoder.cpp
image/decoders/nsJPEGDecoder.cpp
image/decoders/nsPNGDecoder.cpp
image/test/gtest/TestStreamingLexer.cpp
--- a/image/StreamingLexer.h
+++ b/image/StreamingLexer.h
@@ -266,17 +266,18 @@ private:
  *
  * To use StreamingLexer:
  *
  *  - Create a State type. This should be an |enum class| listing all of the
  *    states that you can be in while lexing the image format you're trying to
  *    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
  *    a lambda which dispatches to lexing code for each state based on the State
  *    value that's passed in. The lambda generally should just continue a
  *    |switch| statement that calls different methods for each State value. Each
  *    method should return a LexerTransition<State>, which the lambda should
  *    return in turn.
@@ -290,16 +291,27 @@ private:
  * That's the basics. The StreamingLexer will track your position in the input
  * and buffer enough data so that your lexing methods can process everything in
  * one pass. Lex() returns Yield::NEED_MORE_DATA if more data is needed, in
  * which case you should just return from DoDecode(). If lexing reaches a
  * terminal state, Lex() returns TerminalState::SUCCESS or
  * 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
  * going to ignore.
  *
  * You can begin an unbuffered read with Transition::ToUnbuffered(). This works
  * a little differently than Transition::To() in that you specify *two* states.
  * The @aUnbufferedState argument specifies a state that will be called
@@ -343,32 +355,46 @@ private:
  * XXX(seth): We should be able to get of the |State| stuff totally once bug
  * 1198451 lands, since we can then just return a function representing the next
  * state directly.
  */
 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) {
       // Allowing a StreamingLexer to start in a yield state doesn't make sense
       // semantically (since yield states are supposed to deliver the same data
       // as previous states, and there's no previous state here), but more
       // importantly, it's necessary to advance a SourceBufferIterator at least
       // once before you can read from it, and adding the necessary checks to
       // Lex() to avoid that issue has the potential to mask real bugs. So
       // instead, it's better to forbid starting in a yield state.
       MOZ_ASSERT_UNREACHABLE("Starting in a yield state");
       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);
   }
 
   template <typename Func>
   LexerResult Lex(SourceBufferIterator& aIterator,
                   IResumable* aOnResume,
                   Func aFunc)
   {
@@ -407,24 +433,21 @@ public:
           // We can't continue because the rest of the data hasn't arrived from
           // the network yet. We don't have to do anything special; the
           // SourceBufferIterator will ensure that |aOnResume| gets called when
           // more data is available.
           result = Some(LexerResult(Yield::NEED_MORE_DATA));
           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:
           // Process the new data that became available.
           MOZ_ASSERT(aIterator.Data());
 
           result = mTransition.Buffering() == BufferingStrategy::UNBUFFERED
                  ? UnbufferedRead(aIterator, aFunc)
@@ -624,16 +647,42 @@ private:
                                  mBuffer.length()));
     }
 
     // Anything else indicates a bug.
     MOZ_ASSERT_UNREACHABLE("Unexpected state encountered during yield");
     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
     // unless they're to terminal states. (The terminal state transitions would
     // generally be triggered by error handling code.)
     MOZ_ASSERT_IF(!mBuffer.empty(),
                   aTransition.NextStateIsTerminal() ||
                   mBuffer.length() == mTransition.Size());
@@ -685,16 +734,17 @@ private:
     { }
 
     size_t mBytesRemaining;
     size_t mBytesConsumedInCurrentChunk;
   };
 
   Vector<char, InlineBufferSize> mBuffer;
   LexerTransition<State> mTransition;
+  const LexerTransition<State> mTruncatedTransition;
   Maybe<State> mYieldingToState;
   Maybe<UnbufferedState> mUnbufferedState;
 };
 
 } // namespace image
 } // namespace mozilla
 
 #endif // mozilla_image_StreamingLexer_h
--- a/image/decoders/nsBMPDecoder.cpp
+++ b/image/decoders/nsBMPDecoder.cpp
@@ -165,17 +165,17 @@ Set4BitPixel(uint32_t*& aDecoded, uint8_
 
 static mozilla::LazyLogModule sBMPLog("BMPDecoder");
 
 // The length of the mBIHSize field in the info header.
 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)
   , mNumColors(0)
   , mColors(nullptr)
   , mBytesPerColor(0)
   , mPreGapLength(0)
   , mPixelRowSize(0)
--- a/image/decoders/nsGIFDecoder2.cpp
+++ b/image/decoders/nsGIFDecoder2.cpp
@@ -76,17 +76,18 @@ static const size_t IMAGE_DESCRIPTOR_LEN
 // Masks for reading color table information from packed fields in the screen
 // descriptor and image descriptor blocks.
 static const uint8_t PACKED_FIELDS_COLOR_TABLE_BIT = 0x80;
 static const uint8_t PACKED_FIELDS_INTERLACED_BIT = 0x40;
 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)
   , mGIFOpen(false)
   , mSawTransparency(false)
 {
   // Clear out the structure, excluding the arrays.
   memset(&mGIFStruct, 0, sizeof(mGIFStruct));
--- a/image/decoders/nsICODecoder.cpp
+++ b/image/decoders/nsICODecoder.cpp
@@ -48,17 +48,18 @@ nsICODecoder::GetNumColors()
       numColors = (uint16_t)-1;
     }
   }
   return numColors;
 }
 
 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)
   , mNumIcons(0)
   , mCurrIcon(0)
   , mBPP(0)
   , mMaskRowSize(0)
   , mCurrMaskLine(0)
--- a/image/decoders/nsIconDecoder.cpp
+++ b/image/decoders/nsIconDecoder.cpp
@@ -12,17 +12,18 @@ using namespace mozilla::gfx;
 
 namespace mozilla {
 namespace image {
 
 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
 }
 
 nsIconDecoder::~nsIconDecoder()
 { }
 
--- a/image/decoders/nsJPEGDecoder.cpp
+++ b/image/decoders/nsJPEGDecoder.cpp
@@ -69,17 +69,18 @@ METHODDEF(void) my_error_exit (j_common_
 // Normal JFIF markers can't have more bytes than this.
 #define MAX_JPEG_MARKER_LENGTH  (((uint32_t)1 << 16) - 1)
 
 nsJPEGDecoder::nsJPEGDecoder(RasterImage* aImage,
                              Decoder::DecodeStyle aDecodeStyle)
  : Decoder(aImage)
  , mLexer(Transition::ToUnbuffered(State::FINISHED_JPEG_DATA,
                                    State::JPEG_DATA,
-                                   SIZE_MAX))
+                                   SIZE_MAX),
+          Transition::TerminateSuccess())
  , mDecodeStyle(aDecodeStyle)
  , mSampleSize(0)
 {
   mState = JPEG_HEADER;
   mReading = true;
   mImageData = nullptr;
 
   mBytesToSkip = 0;
--- a/image/decoders/nsPNGDecoder.cpp
+++ b/image/decoders/nsPNGDecoder.cpp
@@ -93,17 +93,18 @@ nsPNGDecoder::AnimFrameInfo::AnimFrameIn
 // First 8 bytes of a PNG file
 const uint8_t
 nsPNGDecoder::pngSignatureBytes[] = { 137, 80, 78, 71, 13, 10, 26, 10 };
 
 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)
  , interlacebuf(nullptr)
  , mInProfile(nullptr)
  , mTransform(nullptr)
  , format(gfx::SurfaceFormat::UNKNOWN)
  , mCMSMode(0)
--- a/image/test/gtest/TestStreamingLexer.cpp
+++ b/image/test/gtest/TestStreamingLexer.cpp
@@ -10,17 +10,19 @@
 using namespace mozilla;
 using namespace mozilla::image;
 
 enum class TestState
 {
   ONE,
   TWO,
   THREE,
-  UNBUFFERED
+  UNBUFFERED,
+  TRUNCATED_SUCCESS,
+  TRUNCATED_FAILURE
 };
 
 void
 CheckLexedData(const char* aData,
                size_t aLength,
                size_t aOffset,
                size_t aExpectedLength)
 {
@@ -39,16 +41,20 @@ DoLex(TestState aState, const char* aDat
       CheckLexedData(aData, aLength, 0, 3);
       return Transition::To(TestState::TWO, 3);
     case TestState::TWO:
       CheckLexedData(aData, aLength, 3, 3);
       return Transition::To(TestState::THREE, 3);
     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");
   }
 }
 
 LexerTransition<TestState>
 DoLexWithUnbuffered(TestState aState, const char* aData, size_t aLength,
                     Vector<char>& aUnbufferedVector)
@@ -216,25 +222,54 @@ DoLexWithZeroLengthStatesAfterUnbuffered
     default:
       MOZ_CRASH("Unexpected or unhandled TestState");
   }
 }
 
 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)
     , mCountResumes(new CountResumes)
   { }
 
 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;
   RefPtr<SourceBuffer> mSourceBuffer;
   SourceBufferIterator mIterator;
   RefPtr<ExpectNoResume> mExpectNoResume;
   RefPtr<CountResumes> mCountResumes;
 };
@@ -254,29 +289,31 @@ TEST_F(ImageStreamingLexer, ZeroLengthDa
 {
   // Test a zero-length input.
   mSourceBuffer->Complete(NS_OK);
 
   // Create a special StreamingLexer for this test because we want the first
   // 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>());
   EXPECT_EQ(TerminalState::FAILURE, result.as<TerminalState>());
 }
 
 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>());
 
   mSourceBuffer->Complete(NS_OK);
 }
 
 TEST_F(ImageStreamingLexer, SingleChunk)
@@ -552,46 +589,49 @@ TEST_F(ImageStreamingLexer, OneByteChunk
 
 TEST_F(ImageStreamingLexer, ZeroLengthState)
 {
   mSourceBuffer->Append(mData, sizeof(mData));
   mSourceBuffer->Complete(NS_OK);
 
   // 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);
 
   EXPECT_TRUE(result.is<TerminalState>());
   EXPECT_EQ(TerminalState::SUCCESS, result.as<TerminalState>());
 }
 
 TEST_F(ImageStreamingLexer, ZeroLengthStatesAtEnd)
 {
   mSourceBuffer->Append(mData, sizeof(mData));
   mSourceBuffer->Complete(NS_OK);
 
   // 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);
 
   EXPECT_TRUE(result.is<TerminalState>());
   EXPECT_EQ(TerminalState::SUCCESS, result.as<TerminalState>());
 }
 
 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 =
     lexer.Lex(mIterator, mExpectNoResume, DoLexWithZeroLengthYield);
   ASSERT_TRUE(result.is<Yield>());
   EXPECT_EQ(Yield::OUTPUT_AVAILABLE, result.as<Yield>());
 
   result = lexer.Lex(mIterator, mCountResumes, DoLexWithZeroLengthYield);
@@ -610,33 +650,35 @@ TEST_F(ImageStreamingLexer, ZeroLengthSt
 {
   mSourceBuffer->Append(mData, sizeof(mData));
   mSourceBuffer->Complete(NS_OK);
 
   // Create a special StreamingLexer for this test because we want the first
   // 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);
 
   EXPECT_TRUE(result.is<TerminalState>());
   EXPECT_EQ(TerminalState::SUCCESS, result.as<TerminalState>());
 }
 
 TEST_F(ImageStreamingLexer, ZeroLengthStateAfterUnbuffered)
 {
   mSourceBuffer->Append(mData, sizeof(mData));
   mSourceBuffer->Complete(NS_OK);
 
   // 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);
 
   EXPECT_TRUE(result.is<TerminalState>());
   EXPECT_EQ(TerminalState::SUCCESS, result.as<TerminalState>());
 }
 
@@ -691,17 +733,18 @@ TEST_F(ImageStreamingLexer, ZeroLengthSt
         MOZ_CRASH("Unexpected or unhandled TestState");
     }
   };
 
   // Create a special StreamingLexer for this test because we want the first
   // 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);
   ASSERT_TRUE(result.is<Yield>());
   EXPECT_EQ(Yield::OUTPUT_AVAILABLE, result.as<Yield>());
   EXPECT_EQ(1u, unbufferedCallCount);
 
   result = lexer.Lex(mIterator, mExpectNoResume, lexerFunc);
@@ -849,64 +892,70 @@ TEST_F(ImageStreamingLexer, SourceBuffer
   mSourceBuffer->Complete(NS_OK);
 
   LexerResult result = mLexer.Lex(mIterator, mExpectNoResume, DoLex);
 
   EXPECT_TRUE(result.is<TerminalState>());
   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);
+  CheckTruncatedState(lexer, TerminalState::SUCCESS);
+}
 
-    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>());
-    }
-  }
+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());
 
-  EXPECT_EQ(2u, mCountResumes->Count());
+  CheckTruncatedState(lexer, TerminalState::FAILURE);
 }
 
-TEST_F(ImageStreamingLexer, SourceBufferTruncatedFailure)
+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);
-    }
-
-    LexerResult result = mLexer.Lex(mIterator, mCountResumes, DoLex);
+  // 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());
 
-    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)
 {
   // Test delivering in one byte chunks with no IResumable.
   for (unsigned i = 0; i < 9; ++i) {
     mSourceBuffer->Append(mData + i, 1);
     LexerResult result = mLexer.Lex(mIterator, nullptr, DoLex);