YAML parser robustness improvements

[thomasfinch]

Summary

This patch fixes a number of bugs found in the YAML parser through fuzzing. In general, this makes the parser more robust against malformed inputs.

The fixes are mostly improved null checking and returning errors in more cases. In some cases, asserts were changed to regular errors, this provides the same robustness but also protects release builds from the triggering conditions. This also improves the fuzzability of the YAML parser since asserts can act as a roadblock to further fuzzing once they're hit.

Each fix has a corresponding test case:

• TestAnchorMapError - Added proper null pointer handling in Stream::printError if N is null and KeyValueNode::getValue if getKey returns null, Input::createHNodes dyn_casts changed to dyn_cast_or_null so the null pointer checks are actually able to fail
• TestFlowSequenceTokenErrors - Added case in Document::parseBlockNode for FlowMappingEnd, FlowSequenceEnd, or FlowEntry tokens outside of mappings or sequences
• TestDirectiveMappingNoValue - Changed assert to regular error return in Scanner::scanValue
• TestUnescapeInfiniteLoop - Fixed infinite loop in ScalarNode::unescapeDoubleQuoted by returning an error for unrecognized escape codes
• TestScannerUnexpectedCharacter - Changed asserts to regular error returns in Scanner::consume
• TestUnknownDirective - For both of the inputs the stream doesn't fail and correctly returns TK_Error, but there is no valid root node for the document. There's no reasonable way to make the scanner fail for unknown directives without breaking the YAML spec (see spec-07-01.test). I think the assert is unnecessary given that an error is still generated for this case.

The SimpleKeys.clear() line fixes a bug found by AddressSanitizer triggered by multiple test cases - when TokenQueue is cleared SimpleKeys is still holding dangling pointers into it, so SimpleKeys should be cleared as well.

[beanz]

@Bigcheese, ping?

[Bigcheese]

Thanks for finding this!

I don't think this is the correct fix though, as all of the provided inputs should emit an error. I believe the issue is that parseBlockNode is missing a case for unexpected tokens. Let me know if you need any more info or if you want me to take a look at fixing it.

[thomasfinch]

Move error checking to Document::parseBlockNode, make sure Document::skip doesn't crash if getRoot returns null.

[thomasfinch]

Split test cases apart, added a few more bug fixes

[Bigcheese]

lib/Support/YAMLParser.cpp

937 ↗

(On Diff #218944)

This error and the one below should be something closer to: "Unable to parse non-ascii here".

1235–1238 ↗

(On Diff #218944)

I realize you changed this so fuzzing works, but this isn't how this case should be handled. setError should be for user error, this is a programming error as the situation should never happen.

1949 ↗

(On Diff #218944)

We shouldn't use ! in error messages.

[Bigcheese]

Also, please use -U999999 for patches so there's full context.

[thomasfinch]

Update diff based on feedback.

• Update error message in Scanner::consume
• Update error message in ScalarNode::unescapeDoubleQuoted
• Remove setError call when failing to find a SimpleKey in the token queue

[thomasfinch]

Remove trailing spaces in KeyValueNode::getValue

[Bigcheese]

lgtm, sorry for the delay, I didn't see the updates for some reason.

[thomasfinch]

@Bigcheese Thanks! Could you please merge this, I don't have commit access.

llvm/lib/Support/YAMLParser.cpp

2401

@Bigcheese should I be calling setError here as well?

[thomasfinch]

Rebasing on monorepo

[hintonda]

llvm/lib/Support/YAMLTraits.cpp

373

Since createHNodes is an internal function and only ever called 3 times within Input:: (and nowhere else), and N is already checked before each of those calls, there's no need add _of_null to these dyn_cast calls and check for null again each time.

To enforce this, I'd suggest you make createHNodes and setError private which will make the class more robust, e.g.:

index f365d60b3ff..fe2fe8c56da 100644
--- a/llvm/include/llvm/Support/YAMLTraits.h
+++ b/llvm/include/llvm/Support/YAMLTraits.h
@@ -1509,6 +1509,7 @@ private:
     std::vector<std::unique_ptr<HNode>> Entries;
   };

+private:
   std::unique_ptr<Input::HNode> createHNodes(Node *node);
   void setError(HNode *hnode, const Twine &message);
   void setError(Node *node, const Twine &message);

However, KeyNode on line 396 could be null, so either checking or adding _or_null to the call is definitely needed there.

[thomasfinch]

• Revert uses of dyn_cast_or_null

[hintonda]

Great, LGTM..

I'll go ahead and land this for you tomorrow unless @Bigcheese has any objections.

llvm/lib/Support/YAMLParser.cpp

 Token &Scanner::peekNext() {
   // If the current token is a possible simple key, keep parsing until we
   // can confirm.
   bool NeedMore = false;
   while (true) {
     if (TokenQueue.empty() || NeedMore) {
       if (!fetchMoreTokens()) {
         TokenQueue.clear();
+        SimpleKeys.clear();
         TokenQueue.push_back(Token());
         return TokenQueue.front();
       }
     }
     assert(!TokenQueue.empty() &&
             "fetchMoreTokens lied about getting tokens!");
 
     removeStaleSimpleKeyCandidates();
if ((   *Current == '%'
       ++Current;
       ++Column;
     } else
       break;
   }
 }
 
 bool Scanner::consume(uint32_t Expected) {
-  if (Expected >= 0x80)
-    report_fatal_error("Not dealing with this yet");
+  if (Expected >= 0x80) {
+    setError("Cannot consume non-ascii characters");
+    return false;
+  }
   if (Current == End)
     return false;
-  if (uint8_t(*Current) >= 0x80)
-    report_fatal_error("Not dealing with this yet");
+  if (uint8_t(*Current) >= 0x80) {
+    setError("Cannot consume non-ascii characters");
+    return false;
+  }
   if (uint8_t(*Current) == Expected) {
     ++Current;
     ++Column;
     return true;
   }
   return false;
 }
 
if (!SimpleKeys.empty()) {
     Token T;
     T.Kind = Token::TK_Key;
     T.Range = SK.Tok->Range;
     TokenQueueT::iterator i, e;
     for (i = TokenQueue.begin(), e = TokenQueue.end(); i != e; ++i) {
       if (i == SK.Tok)
         break;
     }
-    assert(i != e && "SimpleKey not in token queue!");
+    if (i == e) {
+      Failed = true;
+      return false;
+    }
     i = TokenQueue.insert(i, T);
 
     // We may also need to add a Block-Mapping-Start token.
     rollIndent(SK.Column, Token::TK_BlockMappingStart, i);
 
     IsSimpleKeyAllowed = false;
   } else {
     if (!FlowLevel)
Stream::Stream(MemoryBufferRef InputBuffer, SourceMgr &SM, bool ShowColors,
                std::error_code *EC)
     : scanner(new Scanner(InputBuffer, SM, ShowColors, EC)), CurrentDoc() {}
 
 Stream::~Stream() = default;
 
 bool Stream::failed() { return scanner->failed(); }
 
 void Stream::printError(Node *N, const Twine &Msg) {
-  scanner->printError( N->getSourceRange().Start
+  SMRange Range = N ? N->getSourceRange() : SMRange();
+  scanner->printError( Range.Start
                      , SourceMgr::DK_Error
                      , Msg
-                     , N->getSourceRange());
+                     , Range);
 }
 
 document_iterator Stream::begin() {
   if (CurrentDoc)
     report_fatal_error("Can only iterate over the stream once");
 
   // Skip Stream-Start.
   scanner->getNext();
for (; i != StringRef::npos; i = UnquotedValue.find_first_of("\\\r\n")) {
     case '\n':
       Storage.push_back('\n');
       if (   UnquotedValue.size() > 1
           && (UnquotedValue[1] == '\r' || UnquotedValue[1] == '\n'))
         UnquotedValue = UnquotedValue.substr(1);
       UnquotedValue = UnquotedValue.substr(1);
       break;
     default:
-      if (UnquotedValue.size() == 1)
-        // TODO: Report error.
-        break;
+      if (UnquotedValue.size() == 1) {
+        Token T;
+        T.Range = StringRef(UnquotedValue.begin(), 1);
+        setError("Unrecognized escape code", T);
+        return "";
+      }
       UnquotedValue = UnquotedValue.substr(1);
       switch (UnquotedValue[0]) {
       default: {
           Token T;
           T.Range = StringRef(UnquotedValue.begin(), 1);
-          setError("Unrecognized escape code!", T);
+          setError("Unrecognized escape code", T);
           return "";
         }
       case '\r':
       case '\n':
         // Remove the new line.
         if (   UnquotedValue.size() > 1
             && (UnquotedValue[1] == '\r' || UnquotedValue[1] == '\n'))
           UnquotedValue = UnquotedValue.substr(1);
Node *KeyValueNode::getKey() {
 
   // We've got a normal key.
   return Key = parseBlockNode();
 }
 
 Node *KeyValueNode::getValue() {
   if (Value)
     return Value;
-  getKey()->skip();
+
+  if (Node* Key = getKey())
+    Key->skip();
+  else {
+    setError("Null key in Key Value.", peekNext());
+    return Value = new (getAllocator()) NullNode(Doc);
+  }
+
   if (failed())
     return Value = new (getAllocator()) NullNode(Doc);
 
   // Handle implicit null values.
   {
     Token &t = peekNext();
     if (   t.Kind == Token::TK_BlockEnd
         || t.Kind == Token::TK_FlowMappingEnd
return new (NodeAllocator)
                 , T.Range);
   case Token::TK_BlockScalar: {
     getNext();
     StringRef NullTerminatedStr(T.Value.c_str(), T.Value.length() + 1);
     StringRef StrCopy = NullTerminatedStr.copy(NodeAllocator).drop_back();
     return new (NodeAllocator)
         BlockScalarNode(stream.CurrentDoc, AnchorInfo.Range.substr(1),
                         TagInfo.Range, StrCopy, T.Range);
   }
   case Token::TK_Key:
     // Don't eat the TK_Key, KeyValueNode expects it.
     return new (NodeAllocator)
       MappingNode( stream.CurrentDoc
                  , AnchorInfo.Range.substr(1)
                  , TagInfo.Range
                  , MappingNode::MT_Inline);
   case Token::TK_DocumentStart:
   case Token::TK_DocumentEnd:
   case Token::TK_StreamEnd:
   default:
     // TODO: Properly handle tags. "[!!str ]" should resolve to !!str "", not
     //       !!null null.
     return new (NodeAllocator) NullNode(stream.CurrentDoc);
+  case Token::TK_FlowMappingEnd:
+  case Token::TK_FlowSequenceEnd:
+  case Token::TK_FlowEntry: {
+    if (Root && (isa<MappingNode>(Root) || isa<SequenceNode>(Root)))
+      return new (NodeAllocator) NullNode(stream.CurrentDoc);
+
+    setError("Unexpected token", T);
+    return nullptr;
+  }
   case Token::TK_Error:
     return nullptr;
   }
   llvm_unreachable("Control flow shouldn't reach here.");
   return nullptr;
 }
 
 bool Document::parseDirectives() {

llvm/lib/Support/YAMLTraits.cpp

 bool Input::outputting() const {
   return false;
 }
 
 bool Input::setCurrentDocument() {
   if (DocIterator != Strm->end()) {
     Node *N = DocIterator->getRoot();
     if (!N) {
-      assert(Strm->failed() && "Root is NULL iff parsing failed");
       EC = make_error_code(errc::invalid_argument);
       return false;
     }
 
     if (isa<NullNode>(N)) {
       // Empty files are allowed and ignored
       ++DocIterator;
       return setCurrentDocument();

 
 void Input::setError(Node *node, const Twine &message) {
   Strm->printError(node, message);
   EC = make_error_code(errc::invalid_argument);
 }
 
 std::unique_ptr<Input::HNode> Input::createHNodes(Node *N) {
   SmallString<128> StringStorage;
   if (ScalarNode *SN = dyn_cast<ScalarNode>(N)) {
     StringRef KeyStr = SN->getValue(StringStorage);
     if (!StringStorage.empty()) {
       // Copy string to permanent storage
       KeyStr = StringStorage.str().copy(StringAllocator);
     }
     return std::make_unique<ScalarHNode>(N, KeyStr);
   } else if (BlockScalarNode *BSN = dyn_cast<BlockScalarNode>(N)) {
     StringRef ValueCopy = BSN->getValue().copy(StringAllocator);
     return std::make_unique<ScalarHNode>(N, ValueCopy);
   } else if (SequenceNode *SQ = dyn_cast<SequenceNode>(N)) {
     auto SQHNode = std::make_unique<SequenceHNode>(N);
     for (Node &SN : *SQ) {
       auto Entry = createHNodes(&SN);
       if (EC)
         break;
       SQHNode->Entries.push_back(std::move(Entry));
     }
     return std::move(SQHNode);
   } else if (MappingNode *Map = dyn_cast<MappingNode>(N)) {
     auto mapHNode = std::make_unique<MapHNode>(N);
     for (KeyValueNode &KVN : *Map) {
       Node *KeyNode = KVN.getKey();
-      ScalarNode *Key = dyn_cast<ScalarNode>(KeyNode);
+      ScalarNode *Key = dyn_cast_or_null<ScalarNode>(KeyNode);
       Node *Value = KVN.getValue();
       if (!Key || !Value) {
         if (!Key)
           setError(KeyNode, "Map key must be a scalar");
         if (!Value)
           setError(KeyNode, "Map value must not be empty");
         break;
       }

llvm/unittests/Support/YAMLIOTest.cpp

std::string intermediate;
     EXPECT_EQ(foo->SKind, Scalar::SK_Bool);
     EXPECT_FALSE(foo->BoolValue);
     auto bar = llvm::dyn_cast<Scalar>((*map)["bar"].get());
     ASSERT_TRUE(bar);
     EXPECT_EQ(bar->SKind, Scalar::SK_Double);
     EXPECT_EQ(bar->DoubleValue, 2.0);
   }
 }
+
+TEST(YAMLIO, TestAnchorMapError) {
+  Input yin("& & &: ");
+  yin.setCurrentDocument();
+  EXPECT_TRUE(yin.error());
+}
+
+TEST(YAMLIO, TestFlowSequenceTokenErrors) {
+  Input yin(",");
+  EXPECT_FALSE(yin.setCurrentDocument());
+  EXPECT_TRUE(yin.error());
+
+  Input yin2("]");
+  EXPECT_FALSE(yin2.setCurrentDocument());
+  EXPECT_TRUE(yin2.error());
+
+  Input yin3("}");
+  EXPECT_FALSE(yin3.setCurrentDocument());
+  EXPECT_TRUE(yin3.error());
+}
+
+TEST(YAMLIO, TestDirectiveMappingNoValue) {
+  Input yin("%YAML\n{5:");
+  EXPECT_FALSE(yin.setCurrentDocument());
+  EXPECT_TRUE(yin.error());
+
+  Input yin2("%TAG\n'\x98!< :\n");
+  yin2.setCurrentDocument();
+  EXPECT_TRUE(yin2.error());
+}
+
+TEST(YAMLIO, TestUnescapeInfiniteLoop) {
+  Input yin("\"\\u\\^#\\\\\"");
+  yin.setCurrentDocument();
+  EXPECT_TRUE(yin.error());
+}
+
+TEST(YAMLIO, TestScannerUnexpectedCharacter) {
+  Input yin("!<$\x9F.");
+  EXPECT_FALSE(yin.setCurrentDocument());
+  EXPECT_TRUE(yin.error());
+}
+
+TEST(YAMLIO, TestUnknownDirective) {
+  Input yin("%");
+  EXPECT_FALSE(yin.setCurrentDocument());
+  EXPECT_TRUE(yin.error());
+
+  Input yin2("%)");
+  EXPECT_FALSE(yin2.setCurrentDocument());
+  EXPECT_TRUE(yin2.error());
+}