internal/marshal: prevent duplicate xmlns attr

Encoding the output of TokenReader can create multiple xmlns attributes due to a bug in encoding/xml. This patch uses RawToken in the decode stage to work around this and can be reverted after the upstream issue is fixed and in all supported versions of Go. Upstream issue: https://golang.org/issue/42807 (golang/go#42807) Upstream CL: https://golang.org/cl/272806 (golang/go#42808) Fixes #74 Signed-off-by: Sam Whited <[email protected]>
mellium · Nov 25, 2020 · d4b835b · d4b835b
1 parent bb3bf8e
commit d4b835b
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 4 deletions.
diff --git a/internal/marshal/encode.go b/internal/marshal/encode.go
@@ -22,6 +22,10 @@ func TokenReader(v interface{}) (xml.TokenReader, error) {
 		return r, nil
 	}
 
+	return tokenDecoder(v)
+}
+
+func tokenDecoder(v interface{}) (*xml.Decoder, error) {
 	var b bytes.Buffer
 	err := xml.NewEncoder(&b).Encode(v)
 	if err != nil {
@@ -30,6 +34,15 @@ func TokenReader(v interface{}) (xml.TokenReader, error) {
 	return xml.NewDecoder(&b), nil
 }
 
+// rawTokenReader maps a decoders RawToken method onto its Token method.
+type rawTokenReader struct {
+	*xml.Decoder
+}
+
+func (r rawTokenReader) Token() (xml.Token, error) {
+	return r.RawToken()
+}
+
 // EncodeXML writes the XML encoding of v to the stream.
 //
 // See the documentation for xml.Marshal for details about the conversion of Go
@@ -38,11 +51,11 @@ func TokenReader(v interface{}) (xml.TokenReader, error) {
 // If the stream is an xmlstream.Flusher, EncodeXML calls Flush before
 // returning.
 func EncodeXML(w xmlstream.TokenWriter, v interface{}) error {
-	r, err := TokenReader(v)
+	d, err := tokenDecoder(v)
 	if err != nil {
 		return err
 	}
-	_, err = xmlstream.Copy(w, r)
+	_, err = xmlstream.Copy(w, rawTokenReader{Decoder: d})
 	if err != nil {
 		return err
 	}
@@ -62,11 +75,11 @@ func EncodeXML(w xmlstream.TokenWriter, v interface{}) error {
 // If the stream is an xmlstream.Flusher, EncodeXMLElement calls Flush before
 // returning.
 func EncodeXMLElement(w xmlstream.TokenWriter, v interface{}, start xml.StartElement) error {
-	r, err := TokenReader(v)
+	d, err := tokenDecoder(v)
 	if err != nil {
 		return err
 	}
-	_, err = xmlstream.Copy(w, xmlstream.Wrap(r, start))
+	_, err = xmlstream.Copy(w, rawTokenReader{Decoder: d})
 	if err != nil {
 		return err
 	}

diff --git a/internal/marshal/encode_test.go b/internal/marshal/encode_test.go
@@ -122,3 +122,24 @@ func TestMarshalTokenReader(t *testing.T) {
 		t.Errorf("got different xml.TokenReader out: want=%v, got=%v", r, rr)
 	}
 }
+
+func TestTokenDecoder(t *testing.T) {
+	r := stanza.IQ{}
+	_, err := marshal.TokenReader(r)
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+}
+
+func TestEncodeXMLNS(t *testing.T) {
+	var buf bytes.Buffer
+	e := xml.NewEncoder(&buf)
+	err := marshal.EncodeXML(e, simpleIn)
+	if err != nil {
+		t.Errorf("unexpected error encoding: %v", err)
+	}
+	const expected = `<local xmlns="space"></local>`
+	if s := buf.String(); s != expected {
+		t.Errorf("wrong output: want=%s, got=%s", expected, s)
+	}
+}