Fix XML ElWriter bug allowing incomplete XML to be written (#719)

* Add test cases that point out the bugs.

* Fix the bugs by using an owned Option<&mut String>.

* Improve test comments.

* Update CHANGELOG.md.

* Clean up and add explanatory comments.

* Fix whitespace.

* Update CHANGELOG.md

Co-authored-by: Russell Cohen <russell.r.cohen@gmail.com>

CHANGELOG.md

 vNext (Month Day, Year)
 =======================
+- :bug: Fix an issue where `smithy-xml` may have generated invalid XML (smithy-rs#719)
 
 v0.24 (September 24th, 2021)
 ============================

rust-runtime/smithy-xml/src/encode.rs

@@ -59,43 +59,80 @@ impl<'a> XmlWriter<'a> {
 impl<'a> XmlWriter<'a> {
     pub fn start_el<'b, 'c>(&'c mut self, tag: &'b str) -> ElWriter<'c, 'b> {
         write!(self.doc, "<{}", tag).unwrap();
-        ElWriter {
-            doc: self.doc,
-            start: tag,
-        }
+        ElWriter::new(self.doc, tag)
     }
 }
 
 pub struct ElWriter<'a, 'b> {
     start: &'b str,
-    doc: &'a mut String,
+    doc: Option<&'a mut String>,
 }
 
 impl<'a, 'b> ElWriter<'a, 'b> {
+    fn new(doc: &'a mut String, start: &'b str) -> ElWriter<'a, 'b> {
+        ElWriter {
+            start,
+            doc: Some(doc),
+        }
+    }
+
     pub fn write_attribute(&mut self, key: &str, value: &str) -> &mut Self {
-        write!(self.doc, " {}=\"{}\"", key, escape(value)).unwrap();
+        write!(self.doc(), " {}=\"{}\"", key, escape(value)).unwrap();
         self
     }
 
-    pub fn write_ns(self, namespace: &str, prefix: Option<&str>) -> Self {
+    pub fn write_ns(mut self, namespace: &str, prefix: Option<&str>) -> Self {
         match prefix {
             Some(prefix) => {
-                write!(self.doc, " xmlns:{}=\"{}\"", prefix, escape(namespace)).unwrap()
+                write!(self.doc(), " xmlns:{}=\"{}\"", prefix, escape(namespace)).unwrap()
             }
-            None => write!(self.doc, " xmlns=\"{}\"", escape(namespace)).unwrap(),
+            None => write!(self.doc(), " xmlns=\"{}\"", escape(namespace)).unwrap(),
         }
         self
     }
 
-    pub fn finish(self) -> ScopeWriter<'a, 'b> {
-        write!(self.doc, ">").unwrap();
+    fn write_end(doc: &mut String) {
+        write!(doc, ">").unwrap();
+    }
+
+    fn doc<'c>(&'c mut self) -> &'c mut String
+    where
+        'a: 'c,
+    {
+        // The self.doc is an Option in order to signal whether the closing '>' has been emitted
+        // already (None) or not (Some). It ensures the following invariants:
+        // - If finish() has been called, then self.doc is None and therefore no more writes
+        //   to the &mut String are possible.
+        // - When drop() is called, if self.doc is Some, then finish() has not (and will not)
+        //   be called, and therefore drop() should close the tag represented by this struct.
+        //
+        // Since this function calls unwrap(), it must not be called from finish() or drop().
+        // As finish() consumes self, calls to this method from any other method will not encounter
+        // a None value in self.doc.
+        self.doc.as_mut().unwrap()
+    }
+
+    pub fn finish(mut self) -> ScopeWriter<'a, 'b> {
+        let doc = self.doc.take().unwrap();
+        Self::write_end(doc);
         ScopeWriter {
-            doc: self.doc,
+            doc,
             start: self.start,
         }
     }
 }
 
+impl Drop for ElWriter<'_, '_> {
+    fn drop(&mut self) {
+        if let Some(doc) = self.doc.take() {
+            // Calls to write_end() are always preceded by self.doc.take(). The value in self.doc
+            // is set to Some initially, and is never reset to Some after being taken. Since this
+            // transition to None happens only once, we will never double-close the XML element.
+            Self::write_end(doc);
+        }
+    }
+}
+
 /// Wrap the construction of a tag pair `<a></a>`
 pub struct ScopeWriter<'a, 'b> {
     doc: &'a mut String,
@@ -119,10 +156,7 @@ impl ScopeWriter<'_, '_> {
 
     pub fn start_el<'b, 'c>(&'c mut self, tag: &'b str) -> ElWriter<'c, 'b> {
         write!(self.doc, "<{}", tag).unwrap();
-        ElWriter {
-            doc: self.doc,
-            start: tag,
-        }
+        ElWriter::new(self.doc, tag)
     }
 }
 
@@ -131,6 +165,40 @@ mod test {
     use crate::encode::XmlWriter;
     use protocol_test_helpers::{assert_ok, validate_body, MediaType};
 
+    #[test]
+    fn forgot_finish() {
+        let mut out = String::new();
+
+        fn writer(out: &mut String) {
+            let mut doc_writer = XmlWriter::new(out);
+            doc_writer.start_el("Hello");
+            // We intentionally "forget" to call finish() on the ElWriter:
+            // when the XML structs get dropped, the element must get closed automatically.
+        }
+        writer(&mut out);
+
+        assert_ok(validate_body(out, r#"<Hello></Hello>"#, MediaType::Xml));
+    }
+
+    #[test]
+    fn forgot_finish_with_attribute() {
+        let mut out = String::new();
+
+        fn writer(out: &mut String) {
+            let mut doc_writer = XmlWriter::new(out);
+            doc_writer.start_el("Hello").write_attribute("key", "foo");
+            // We intentionally "forget" to call finish() on the ElWriter:
+            // when the XML structs get dropped, the element must get closed automatically.
+        }
+        writer(&mut out);
+
+        assert_ok(validate_body(
+            out,
+            r#"<Hello key="foo"></Hello>"#,
+            MediaType::Xml,
+        ));
+    }
+
     #[test]
     fn basic_document_encoding() {
         let mut out = String::new();