5 files changed, 159 insertions, 65 deletions
diff --git a/Swiften/Client/ClientXMLTracer.cpp b/Swiften/Client/ClientXMLTracer.cpp
index 4852aa1..e205f41 100644
--- a/Swiften/Client/ClientXMLTracer.cpp
+++ b/Swiften/Client/ClientXMLTracer.cpp
@@ -14,23 +14,19 @@
 
 namespace Swift {
 
-ClientXMLTracer::ClientXMLTracer(CoreClient* client, bool bosh) : bosh(bosh) {
+ClientXMLTracer::ClientXMLTracer(CoreClient* client, bool bosh) : bosh_(bosh) {
 #ifdef SWIFTEN_PLATFORM_WIN32
-    beautifier = new XMLBeautifier(true, false);
+    beautifier_ = std::unique_ptr<XMLBeautifier>(new XMLBeautifier(true, false));
 #else
-    beautifier = new XMLBeautifier(true, true);
+    beautifier_ = std::unique_ptr<XMLBeautifier>(new XMLBeautifier(true, true));
 #endif
-    onDataReadConnection = client->onDataRead.connect(boost::bind(&ClientXMLTracer::printData, this, '<', _1));
-    onDataWrittenConnection = client->onDataWritten.connect(boost::bind(&ClientXMLTracer::printData, this, '>', _1));
-}
-
-ClientXMLTracer::~ClientXMLTracer() {
-    delete beautifier;
+    onDataReadConnection_ = client->onDataRead.connect(boost::bind(&ClientXMLTracer::printData, this, '<', _1));
+    onDataWrittenConnection_ = client->onDataWritten.connect(boost::bind(&ClientXMLTracer::printData, this, '>', _1));
 }
 
 void ClientXMLTracer::printData(char direction, const SafeByteArray& data) {
-    printLine(direction);
-    if (bosh) {
+    if (bosh_) {
+        printLine(direction);
         std::string line = byteArrayToString(ByteArray(data.begin(), data.end()));
 // Disabled because it swallows bits of XML (namespaces, if I recall)
 //        size_t endOfHTTP = line.find("\r\n\r\n");
@@ -42,7 +38,16 @@ void ClientXMLTracer::printData(char direction, const SafeByteArray& data) {
 //        }
     }
     else {
-        std::cerr << beautifier->beautify(byteArrayToString(ByteArray(data.begin(), data.end()))) << std::endl;
+        const auto& str = beautifier_->beautify(byteArrayToString(ByteArray(data.begin(), data.end())));
+
+        if (beautifier_->wasReset()) {
+            printLine(direction);
+        }
+        std::cerr << str;
+        if (beautifier_->getLevel() <= 1) {
+            std::cerr << std::endl;
+        }
+
     }
 }
 
diff --git a/Swiften/Client/ClientXMLTracer.h b/Swiften/Client/ClientXMLTracer.h
index ad28fe2..d3aeb0c 100644
--- a/Swiften/Client/ClientXMLTracer.h
+++ b/Swiften/Client/ClientXMLTracer.h
@@ -6,6 +6,8 @@
 
 #pragma once
 
+#include <memory>
+
 #include <Swiften/Base/API.h>
 #include <Swiften/Base/SafeByteArray.h>
 #include <Swiften/Client/CoreClient.h>
@@ -15,16 +17,15 @@ namespace Swift {
     class SWIFTEN_API ClientXMLTracer {
         public:
             ClientXMLTracer(CoreClient* client, bool bosh = false);
-            ~ClientXMLTracer();
 
         private:
             void printData(char direction, const SafeByteArray& data);
             void printLine(char c);
 
         private:
-            XMLBeautifier *beautifier;
-            bool bosh;
-            boost::signals2::scoped_connection onDataReadConnection;
-            boost::signals2::scoped_connection onDataWrittenConnection;
+            std::unique_ptr<XMLBeautifier> beautifier_;
+            bool bosh_;
+            boost::signals2::scoped_connection onDataReadConnection_;
+            boost::signals2::scoped_connection onDataWrittenConnection_;
     };
 }
diff --git a/Swiften/Client/UnitTest/XMLBeautifierTest.cpp b/Swiften/Client/UnitTest/XMLBeautifierTest.cpp
new file mode 100644
index 0000000..0188634
--- /dev/null
+++ b/Swiften/Client/UnitTest/XMLBeautifierTest.cpp
@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) 2010-2017 Isode Limited.
+ * All rights reserved.
+ * See the COPYING file for more information.
+ */
+
+#include "gtest/gtest.h"
+#include <Swiften/Client/XMLBeautifier.h>
+#include <iostream>
+
+using namespace Swift;
+
+namespace {
+    const static std::string FULL_FORMATTED_OUTPUT("<list>\n <el>aqq</el>\n <el>bzz</el>\n</list>");
+}
+
+TEST(XMLBeautifierTest, testBeautify) {
+    auto beautifier = std::unique_ptr<XMLBeautifier>(new XMLBeautifier(true, false));
+
+    ASSERT_EQ(FULL_FORMATTED_OUTPUT, beautifier->beautify("<list><el>aqq</el><el>bzz</el></list>"));
+    ASSERT_TRUE(beautifier->wasReset());
+    ASSERT_EQ(0, beautifier->getLevel());
+    ASSERT_EQ(FULL_FORMATTED_OUTPUT, beautifier->beautify("<list><el>aqq</el><el>bzz</el></list>"));
+    ASSERT_TRUE(beautifier->wasReset());
+    ASSERT_EQ(0, beautifier->getLevel());
+}
+
+TEST(XMLBeautifierTest, testBeautifyMultipleChunks) {
+    auto beautifier = std::unique_ptr<XMLBeautifier>(new XMLBeautifier(true, false));
+
+    auto result = beautifier->beautify("<list><el>aqq</el>");
+    ASSERT_TRUE(beautifier->wasReset());
+    ASSERT_EQ(1, beautifier->getLevel());
+
+    result += beautifier->beautify("<el>bzz</el></list>");
+    ASSERT_FALSE(beautifier->wasReset());
+    ASSERT_EQ(0, beautifier->getLevel());
+
+    ASSERT_EQ(FULL_FORMATTED_OUTPUT, result);
+}
+
+TEST(XMLBeautifierTest, testBeautifyMultipleChunksMiddleElement) {
+    auto beautifier = std::unique_ptr<XMLBeautifier>(new XMLBeautifier(true, false));
+
+    auto result = beautifier->beautify("<l");
+    ASSERT_TRUE(beautifier->wasReset());
+    ASSERT_EQ(0, beautifier->getLevel());
+
+    result += beautifier->beautify("ist><el>aqq</el><el>bzz</el></list>");
+    ASSERT_FALSE(beautifier->wasReset());
+    ASSERT_EQ(0, beautifier->getLevel());
+
+    ASSERT_EQ(FULL_FORMATTED_OUTPUT, result);
+}
+
+TEST(XMLBeautifierTest, testBeautifyInvalidMultipleChunks) {
+    auto beautifier = std::unique_ptr<XMLBeautifier>(new XMLBeautifier(true, false));
+
+    ASSERT_EQ(std::string("<list>\n <el>aqq"), beautifier->beautify("<list><el>aqq<"));
+    ASSERT_TRUE(beautifier->wasReset());
+    ASSERT_EQ(2, beautifier->getLevel());
+
+    ASSERT_EQ(FULL_FORMATTED_OUTPUT, beautifier->beautify("<list><el>aqq</el><el>bzz</el></list>"));
+    ASSERT_TRUE(beautifier->wasReset());
+    ASSERT_EQ(0, beautifier->getLevel());
+}
diff --git a/Swiften/Client/XMLBeautifier.cpp b/Swiften/Client/XMLBeautifier.cpp
index e2cd58e..97a228f 100644
--- a/Swiften/Client/XMLBeautifier.cpp
+++ b/Swiften/Client/XMLBeautifier.cpp
@@ -14,35 +14,53 @@
 
 #include <sstream>
 #include <stack>
+#include <iostream>
 
 #include <Swiften/Base/Log.h>
 #include <Swiften/Parser/PlatformXMLParserFactory.h>
 
 namespace Swift {
 
-XMLBeautifier::XMLBeautifier(bool indention, bool coloring) : doIndention(indention), doColoring(coloring), intLevel(0), parser(nullptr), lastWasStepDown(false) {
-    factory = new PlatformXMLParserFactory();
+XMLBeautifier::XMLBeautifier(bool indention, bool coloring) : doIndention_(indention), doColoring_(coloring), factory_(new PlatformXMLParserFactory()) {
 }
 
-XMLBeautifier::~XMLBeautifier() {
-    delete factory;
-}
 
 std::string XMLBeautifier::beautify(const std::string &text) {
-    parser = factory->createXMLParser(this);
-    intLevel = 0;
-    buffer.str(std::string());
-    parser->parse(text);
-    delete parser;
-    return buffer.str();
+    wasReset_ = false;
+    if (!parser_) {
+        reset();
+    }
+    buffer_.str(std::string());
+    if (!parser_->parse(text)) {
+        reset();
+        parser_->parse(text);
+    }
+    return buffer_.str();
+}
+
+bool XMLBeautifier::wasReset() const {
+    return wasReset_;
+}
+
+int XMLBeautifier::getLevel() const {
+    return intLevel_;
 }
 
 void XMLBeautifier::indent() {
-    for (int i = 0; i < intLevel; ++i) {
-        buffer << " ";
+    for (int i = 0; i < intLevel_; ++i) {
+        buffer_ << " ";
     }
 }
 
+void XMLBeautifier::reset() {
+    parser_ = factory_->createXMLParser(this);
+    intLevel_ = 0;
+    lastWasStepDown_ = false;
+    std::stack<std::string>().swap(parentNSs_);
+    buffer_.str(std::string());
+    wasReset_ = true;
+}
+
 // all bold but reset
 // static const char colorBlue[] = "\x1b[01;34m";
 static const char colorCyan[] = "\x1b[01;36m";
@@ -86,47 +104,47 @@ std::string XMLBeautifier::styleValue(const std::string& text) const {
 }
 
 void XMLBeautifier::handleStartElement(const std::string& element, const std::string& ns, const AttributeMap& attributes) {
-    if (doIndention) {
-        if (intLevel) buffer << std::endl;
+    if (doIndention_) {
+        if (intLevel_) buffer_ << std::endl;
     }
     indent();
-    buffer << "<" << (doColoring ? styleTag(element) : element);
-    if (!ns.empty() && (!parentNSs.empty() && parentNSs.top() != ns)) {
-        buffer << " ";
-        buffer << (doColoring ? styleAttribute("xmlns") : "xmlns");
-        buffer << "=";
-        buffer << "\"" << (doColoring ? styleNamespace(ns) : ns) << "\"";
+    buffer_ << "<" << (doColoring_ ? styleTag(element) : element);
+    if (!ns.empty() && (!parentNSs_.empty() && parentNSs_.top() != ns)) {
+        buffer_ << " ";
+        buffer_ << (doColoring_ ? styleAttribute("xmlns") : "xmlns");
+        buffer_ << "=";
+        buffer_ << "\"" << (doColoring_ ? styleNamespace(ns) : ns) << "\"";
     }
     if (!attributes.getEntries().empty()) {
         for (const auto& entry : attributes.getEntries()) {
-            buffer << " ";
-            buffer << (doColoring ? styleAttribute(entry.getAttribute().getName()) : entry.getAttribute().getName());
-            buffer << "=";
-            buffer << "\"" << (doColoring ? styleValue(entry.getValue()) : entry.getValue()) << "\"";
+            buffer_ << " ";
+            buffer_ << (doColoring_ ? styleAttribute(entry.getAttribute().getName()) : entry.getAttribute().getName());
+            buffer_ << "=";
+            buffer_ << "\"" << (doColoring_ ? styleValue(entry.getValue()) : entry.getValue()) << "\"";
         }
     }
-    buffer << ">";
-    ++intLevel;
-    lastWasStepDown = false;
-    parentNSs.push(ns);
+    buffer_ << ">";
+    ++intLevel_;
+    lastWasStepDown_ = false;
+    parentNSs_.push(ns);
 }
 
 void XMLBeautifier::handleEndElement(const std::string& element, const std::string& /* ns */) {
-    --intLevel;
-    parentNSs.pop();
-    if (/*hadCDATA.top() ||*/ lastWasStepDown) {
-        if (doIndention) {
-            buffer << std::endl;
+    --intLevel_;
+    parentNSs_.pop();
+    if (/*hadCDATA.top() ||*/ lastWasStepDown_) {
+        if (doIndention_) {
+            buffer_ << std::endl;
         }
         indent();
     }
-    buffer << "</" << (doColoring ? styleTag(element) : element) << ">";
-    lastWasStepDown = true;
+    buffer_ << "</" << (doColoring_ ? styleTag(element) : element) << ">";
+    lastWasStepDown_ = true;
 }
 
 void XMLBeautifier::handleCharacterData(const std::string& data) {
-    buffer << data;
-    lastWasStepDown = false;
+    buffer_ << data;
+    lastWasStepDown_ = false;
 }
 
 }
diff --git a/Swiften/Client/XMLBeautifier.h b/Swiften/Client/XMLBeautifier.h
index 8da935a..7458211 100644
--- a/Swiften/Client/XMLBeautifier.h
+++ b/Swiften/Client/XMLBeautifier.h
@@ -12,6 +12,7 @@
 
 #pragma once
 
+#include <memory>
 #include <sstream>
 #include <stack>
 #include <string>
@@ -27,10 +28,12 @@ namespace Swift {
 
 class SWIFTEN_API XMLBeautifier : public XMLParserClient {
 public:
+
     XMLBeautifier(bool indention, bool coloring);
-    virtual ~XMLBeautifier();
 
     std::string beautify(const std::string&);
+    bool wasReset() const;
+    int getLevel() const;
 
 private:
     void handleStartElement(const std::string& element, const std::string& ns, const AttributeMap& attributes);
@@ -39,6 +42,7 @@ private:
 
 private:
     void indent();
+    void reset();
 
 private:
     std::string styleTag(const std::string& text) const;
@@ -47,16 +51,16 @@ private:
     std::string styleValue(const std::string& text) const;
 
 private:
-    bool doIndention;
-    bool doColoring;
+    const bool doIndention_;
+    const bool doColoring_;
 
-    int intLevel;
-    std::string inputBuffer;
-    std::stringstream buffer;
-    XMLParserFactory* factory;
-    XMLParser* parser;
+    std::unique_ptr<XMLParserFactory> factory_;
+    std::unique_ptr<XMLParser> parser_;
 
-    bool lastWasStepDown;
-    std::stack<std::string> parentNSs;
+    bool wasReset_ = true;
+    int intLevel_ = 0;
+    bool lastWasStepDown_ = false;
+    std::stringstream buffer_;
+    std::stack<std::string> parentNSs_;
 };
 }