SWS-359

Author: poutsma

core/src/main/java/org/springframework/ws/soap/axiom/CachingPayload.java

@@ -24,6 +24,8 @@
 import org.apache.axiom.soap.SOAPFactory;
 
 /**
+ * Caching payload in Axiom.
+ *
  * @author Arjen Poutsma
  * @since 1.5.2
  */

core/src/main/java/org/springframework/ws/soap/axiom/NonCachingPayload.java

@@ -35,6 +35,8 @@
 import org.springframework.xml.transform.StaxResult;
 
 /**
+ * Non-caching payload in Axiom.
+ *
  * @author Jim Cummings
  * @author Arjen Poutsma
  * @since 1.5.2
@@ -63,7 +65,9 @@ private class DelegatingStreamWriter implements XMLStreamWriter {
 
         private QName name;
 
-        private String encoding;
+        private String encoding = "UTF-8";
+
+        private int elementDepth = 0;
 
         private DelegatingStreamWriter() {
             try {
@@ -75,51 +79,65 @@ private DelegatingStreamWriter() {
         }
 
         public void writeStartDocument() throws XMLStreamException {
-            this.encoding = "UTF-8";
+            // ignored
         }
 
         public void writeStartDocument(String version) throws XMLStreamException {
-            this.encoding = "UTF-8";
+            // ignored
         }
 
         public void writeStartDocument(String encoding, String version) throws XMLStreamException {
             this.encoding = encoding;
         }
 
         public void writeEndDocument() throws XMLStreamException {
-            delegate.writeEndDocument();
-            delegate.flush();
-            if (baos.size() > 0) {
-                byte[] buf = baos.toByteArray();
-                OMDataSource dataSource = new ByteArrayDataSource(buf, encoding);
-                OMNamespace namespace = getAxiomFactory().createOMNamespace(name.getNamespaceURI(), name.getPrefix());
-                OMElement payloadElement =
-                        getAxiomFactory().createOMElement(dataSource, name.getLocalPart(), namespace);
-                getAxiomBody().addChild(payloadElement);
-            }
+            // ignored
         }
 
         public void writeStartElement(String localName) throws XMLStreamException {
             if (name == null) {
                 name = new QName(localName);
             }
+            elementDepth++;
             delegate.writeStartElement(localName);
         }
 
         public void writeStartElement(String namespaceURI, String localName) throws XMLStreamException {
             if (name == null) {
                 name = new QName(namespaceURI, localName);
             }
+            elementDepth++;
             delegate.writeStartElement(namespaceURI, localName);
         }
 
         public void writeStartElement(String prefix, String localName, String namespaceURI) throws XMLStreamException {
             if (name == null) {
                 name = new QName(namespaceURI, localName, prefix);
             }
+            elementDepth++;
             delegate.writeStartElement(prefix, localName, namespaceURI);
         }
 
+        public void writeEndElement() throws XMLStreamException {
+            elementDepth--;
+            delegate.writeEndElement();
+            if (elementDepth <= 0) {
+                addPayload();
+            }
+        }
+
+        private void addPayload() throws XMLStreamException {
+            delegate.flush();
+            if (baos.size() > 0) {
+                byte[] buf = baos.toByteArray();
+                OMDataSource dataSource = new ByteArrayDataSource(buf, encoding);
+                OMNamespace namespace = getAxiomFactory().createOMNamespace(name.getNamespaceURI(), name.getPrefix());
+                OMElement payloadElement =
+                        getAxiomFactory().createOMElement(dataSource, name.getLocalPart(), namespace);
+                getAxiomBody().addChild(payloadElement);
+            }
+        }
+
         public void writeEmptyElement(String localName) throws XMLStreamException {
             if (name == null) {
                 name = new QName(localName);
@@ -212,10 +230,6 @@ public void writeDTD(String dtd) throws XMLStreamException {
             delegate.writeDTD(dtd);
         }
 
-        public void writeEndElement() throws XMLStreamException {
-            delegate.writeEndElement();
-        }
-
         public void writeEntityRef(String name) throws XMLStreamException {
             delegate.writeEntityRef(name);
         }

core/src/main/java/org/springframework/ws/soap/axiom/Payload.java

@@ -30,6 +30,9 @@
 import org.springframework.xml.transform.StaxSource;
 
 /**
+ * Abstract base class for payloads in Axiom. Comes in two flavors: {@link CachingPayload} and {@link
+ * NonCachingPayload}.
+ *
  * @author Arjen Poutsma
  * @since 1.5.2
  */