Add crash-triggering test case and more null checks in YAML parser

[thomasfinch]

Summary

Fix crashes when parsing flow tokens outside of a mapping or sequence node

Fix assert when reading a directive then a mapping with no value

Fix infinite loop in unescapeDoubleQuoted for invalid escapes

Make Scanner::consume unexpected characters error nonfatal

Fix ASan-caught use-after-poison with YAML parser SimpleKeys

Add test for unknown directives

[thomasfinch]

Accidental push...

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("Not dealing with this yet");
+    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("Not dealing with this yet");
+    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) {
+      setError("SimpleKey not in token queue!");
+      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);
           return "";
         }
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)
                  , 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() {

lib/Support/YAMLTraits.cpp

 bool Input::outputting() {
   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)) {
+  if (ScalarNode *SN = dyn_cast_or_null<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)) {
+  } else if (BlockScalarNode *BSN = dyn_cast_or_null<BlockScalarNode>(N)) {
     StringRef ValueCopy = BSN->getValue().copy(StringAllocator);
     return std::make_unique<ScalarHNode>(N, ValueCopy);
-  } else if (SequenceNode *SQ = dyn_cast<SequenceNode>(N)) {
+  } else if (SequenceNode *SQ = dyn_cast_or_null<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)) {
+  } else if (MappingNode *Map = dyn_cast_or_null<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;
       }

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());
+}