Recover gracefully when deserializing invalid YAML input.

Fixes http://llvm.org/PR16221, http://llvm.org/PR15927
Phabricator: http://llvm-reviews.chandlerc.com/D1236

Patch by Andrew Tulloch!


git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@195016 91177308-0d34-0410-b5e6-96231b3b80d8

3 files changed, 152 insertions, 61 deletions

diff --git a/include/llvm/Support/YAMLTraits.h b/include/llvm/Support/YAMLTraits.h
index d6eeaace7b..8be54394de 100644
--- a/include/llvm/Support/YAMLTraits.h
+++ b/include/llvm/Support/YAMLTraits.h
@@ -685,16 +685,18 @@ private:
 ///
 class Input : public IO {
 public:
-  // Construct a yaml Input object from a StringRef and optional user-data.
-  Input(StringRef InputContent, void *Ctxt=NULL);
+  // Construct a yaml Input object from a StringRef and optional
+  // user-data. The DiagHandler can be specified to provide
+  // alternative error reporting.
+  Input(StringRef InputContent,
+        void *Ctxt = NULL,
+        SourceMgr::DiagHandlerTy DiagHandler = NULL,
+        void *DiagHandlerCtxt = NULL);
   ~Input();
-  
+
   // Check if there was an syntax or semantic error during parsing.
   llvm::error_code error();
 
-  // To set alternate error reporting.
-  void setDiagHandler(llvm::SourceMgr::DiagHandlerTy Handler, void *Ctxt = 0);
-
   static bool classof(const IO *io) { return !io->outputting(); }
 
 private:
@@ -968,8 +970,8 @@ template <typename T>
 inline
 typename llvm::enable_if_c<has_SequenceTraits<T>::value,Input &>::type
 operator>>(Input &yin, T &docSeq) {
-  yin.setCurrentDocument();
-  yamlize(yin, docSeq, true);
+  if (yin.setCurrentDocument())
+    yamlize(yin, docSeq, true);
   return yin;
 }
 
diff --git a/lib/Support/YAMLTraits.cpp b/lib/Support/YAMLTraits.cpp
index f103ed8bcb..e2108e5b7a 100644
--- a/lib/Support/YAMLTraits.cpp
+++ b/lib/Support/YAMLTraits.cpp
@@ -41,10 +41,15 @@ void IO::setContext(void *Context) {
 //  Input
 //===----------------------------------------------------------------------===//
 
-Input::Input(StringRef InputContent, void *Ctxt)
+Input::Input(StringRef InputContent,
+             void *Ctxt,
+             SourceMgr::DiagHandlerTy DiagHandler,
+             void *DiagHandlerCtxt)
   : IO(Ctxt),
     Strm(new Stream(InputContent, SrcMgr)),
     CurrentNode(NULL) {
+  if (DiagHandler)
+    SrcMgr.setDiagHandler(DiagHandler, DiagHandlerCtxt);
   DocIterator = Strm->begin();
 }
 
@@ -55,10 +60,6 @@ error_code Input::error() {
   return EC;
 }
 
-void Input::setDiagHandler(SourceMgr::DiagHandlerTy Handler, void *Ctxt) {
-  SrcMgr.setDiagHandler(Handler, Ctxt);
-}
-
 bool Input::outputting() const {
   return false;
 }
@@ -66,6 +67,12 @@ bool Input::outputting() const {
 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;
@@ -95,7 +102,8 @@ bool Input::mapTag(StringRef Tag, bool Default) {
 void Input::beginMapping() {
   if (EC)
     return;
-  MapHNode *MN = dyn_cast<MapHNode>(CurrentNode);
+  // CurrentNode can be null if the document is empty.
+  MapHNode *MN = dyn_cast_or_null<MapHNode>(CurrentNode);
   if (MN) {
     MN->ValidKeys.clear();
   }
@@ -106,6 +114,15 @@ bool Input::preflightKey(const char *Key, bool Required, bool, bool &UseDefault,
   UseDefault = false;
   if (EC)
     return false;
+
+  // CurrentNode is null for empty documents, which is an error in case required
+  // nodes are present.
+  if (!CurrentNode) {
+    if (Required)
+      EC = make_error_code(errc::invalid_argument);
+    return false;
+  }
+
   MapHNode *MN = dyn_cast<MapHNode>(CurrentNode);
   if (!MN) {
     setError(CurrentNode, "not a mapping");
@@ -132,7 +149,8 @@ void Input::postflightKey(void *saveInfo) {
 void Input::endMapping() {
   if (EC)
     return;
-  MapHNode *MN = dyn_cast<MapHNode>(CurrentNode);
+  // CurrentNode can be null if the document is empty.
+  MapHNode *MN = dyn_cast_or_null<MapHNode>(CurrentNode);
   if (!MN)
     return;
   for (MapHNode::NameToNode::iterator i = MN->Mapping.begin(),
@@ -273,6 +291,7 @@ void Input::scalarString(StringRef &S) {
 }
 
 void Input::setError(HNode *hnode, const Twine &message) {
+  assert(hnode && "HNode must not be NULL");
   this->setError(hnode->_node, message);
 }
 
diff --git a/unittests/Support/YAMLIOTest.cpp b/unittests/Support/YAMLIOTest.cpp
index bc4ea2246c..6c0b9e6151 100644
--- a/unittests/Support/YAMLIOTest.cpp
+++ b/unittests/Support/YAMLIOTest.cpp
@@ -58,14 +58,24 @@ namespace yaml {
 //
 TEST(YAMLIO, TestMapRead) {
   FooBar doc;
-  Input yin("---\nfoo:  3\nbar:  5\n...\n");
-  yin >> doc;
+  {
+    Input yin("---\nfoo:  3\nbar:  5\n...\n");
+    yin >> doc;
 
-  EXPECT_FALSE(yin.error());
-  EXPECT_EQ(doc.foo, 3);
-  EXPECT_EQ(doc.bar,5);
-}
+    EXPECT_FALSE(yin.error());
+    EXPECT_EQ(doc.foo, 3);
+    EXPECT_EQ(doc.bar, 5);
+  }
 
+  {
+    Input yin("{foo: 3, bar: 5}");
+    yin >> doc;
+
+    EXPECT_FALSE(yin.error());
+    EXPECT_EQ(doc.foo, 3);
+    EXPECT_EQ(doc.bar, 5);
+  }
+}
 
 //
 // Test the reading of a yaml sequence of mappings
@@ -1164,8 +1174,9 @@ TEST(YAMLIO, TestColorsReadError) {
             "c1:  blue\n"
             "c2:  purple\n"
             "c3:  green\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> map;
   EXPECT_TRUE(yin.error());
 }
@@ -1180,8 +1191,9 @@ TEST(YAMLIO, TestFlagsReadError) {
             "f1:  [ big ]\n"
             "f2:  [ round, hollow ]\n"
             "f3:  []\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> map;
 
   EXPECT_TRUE(yin.error());
@@ -1198,8 +1210,9 @@ TEST(YAMLIO, TestReadBuiltInTypesUint8Error) {
             "- 255\n"
             "- 0\n"
             "- 257\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1216,8 +1229,9 @@ TEST(YAMLIO, TestReadBuiltInTypesUint16Error) {
             "- 65535\n"
             "- 0\n"
             "- 66000\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1234,8 +1248,9 @@ TEST(YAMLIO, TestReadBuiltInTypesUint32Error) {
             "- 4000000000\n"
             "- 0\n"
             "- 5000000000\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1252,8 +1267,9 @@ TEST(YAMLIO, TestReadBuiltInTypesUint64Error) {
             "- 18446744073709551615\n"
             "- 0\n"
             "- 19446744073709551615\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1271,8 +1287,9 @@ TEST(YAMLIO, TestReadBuiltInTypesint8OverError) {
             "- 0\n"
             "- 127\n"
             "- 128\n"
-           "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+           "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1288,8 +1305,9 @@ TEST(YAMLIO, TestReadBuiltInTypesint8UnderError) {
             "- 0\n"
             "- 127\n"
             "- -129\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1307,8 +1325,9 @@ TEST(YAMLIO, TestReadBuiltInTypesint16UnderError) {
             "- 0\n"
             "- -32768\n"
             "- -32769\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1325,8 +1344,9 @@ TEST(YAMLIO, TestReadBuiltInTypesint16OverError) {
             "- 0\n"
             "- -32768\n"
             "- 32768\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1344,8 +1364,9 @@ TEST(YAMLIO, TestReadBuiltInTypesint32UnderError) {
             "- 0\n"
             "- -2147483648\n"
             "- -2147483649\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1361,8 +1382,9 @@ TEST(YAMLIO, TestReadBuiltInTypesint32OverError) {
             "- 0\n"
             "- -2147483648\n"
             "- 2147483649\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1380,8 +1402,9 @@ TEST(YAMLIO, TestReadBuiltInTypesint64UnderError) {
             "- 0\n"
             "- 9223372036854775807\n"
             "- -9223372036854775809\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1397,8 +1420,9 @@ TEST(YAMLIO, TestReadBuiltInTypesint64OverError) {
             "- 0\n"
             "- 9223372036854775807\n"
             "- 9223372036854775809\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1415,8 +1439,9 @@ TEST(YAMLIO, TestReadBuiltInTypesFloatError) {
             "- 1000.1\n"
             "- -123.456\n"
             "- 1.2.3\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1433,8 +1458,9 @@ TEST(YAMLIO, TestReadBuiltInTypesDoubleError) {
             "- 1000.1\n"
             "- -123.456\n"
             "- 1.2.3\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1450,8 +1476,9 @@ TEST(YAMLIO, TestReadBuiltInTypesHex8Error) {
             "- 0x12\n"
             "- 0xFE\n"
             "- 0x123\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1468,8 +1495,9 @@ TEST(YAMLIO, TestReadBuiltInTypesHex16Error) {
             "- 0x0012\n"
             "- 0xFEFF\n"
             "- 0x12345\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1485,8 +1513,9 @@ TEST(YAMLIO, TestReadBuiltInTypesHex32Error) {
             "- 0x0012\n"
             "- 0xFEFF0000\n"
             "- 0x1234556789\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
@@ -1502,13 +1531,31 @@ TEST(YAMLIO, TestReadBuiltInTypesHex64Error) {
             "- 0x0012\n"
             "- 0xFFEEDDCCBBAA9988\n"
             "- 0x12345567890ABCDEF0\n"
-            "...\n");
-  yin.setDiagHandler(suppressErrorMessages);
+            "...\n",
+            /*Ctxt=*/NULL,
+            suppressErrorMessages);
   yin >> seq;
 
   EXPECT_TRUE(yin.error());
 }
 
+TEST(YAMLIO, TestMalformedMapFailsGracefully) {
+  FooBar doc;
+  {
+    // We pass the suppressErrorMessages handler to handle the error
+    // message generated in the constructor of Input.
+    Input yin("{foo:3, bar: 5}", /*Ctxt=*/NULL, suppressErrorMessages);
+    yin >> doc;
+    EXPECT_TRUE(yin.error());
+  }
+
+  {
+    Input yin("---\nfoo:3\nbar: 5\n...\n", /*Ctxt=*/NULL, suppressErrorMessages);
+    yin >> doc;
+    EXPECT_TRUE(yin.error());
+  }
+}
+
 struct OptionalTest {
   std::vector<int> Numbers;
 };
@@ -1572,3 +1619,26 @@ TEST(YAMLIO, SequenceElideTest) {
 
   EXPECT_TRUE(Seq2.Tests[3].Numbers.empty());
 }
+
+TEST(YAMLIO, TestEmptyStringFailsForMapWithRequiredFields) {
+  FooBar doc;
+  Input yin("");
+  yin >> doc;
+  EXPECT_TRUE(yin.error());
+}
+
+TEST(YAMLIO, TestEmptyStringSucceedsForMapWithOptionalFields) {
+  OptionalTest doc;
+  Input yin("");
+  yin >> doc;
+  EXPECT_FALSE(yin.error());
+}
+
+TEST(YAMLIO, TestEmptyStringSucceedsForSequence) {
+  std::vector<uint8_t> seq;
+  Input yin("", /*Ctxt=*/NULL, suppressErrorMessages);
+  yin >> seq;
+
+  EXPECT_FALSE(yin.error());
+  EXPECT_TRUE(seq.empty());
+}