From 9593b74a3c993c4ba472163e4011c9feec97e7b9 Mon Sep 17 00:00:00 2001
From: Sam Whited <sam@samwhited.com>
Date: Wed, 29 Mar 2017 13:19:29 -0500
Subject: [PATCH] encoding/xml: add decode wrapper

Fixes #19480

Change-Id: I5a621507279d5bb1f3991b7a412d9a63039d464b
Reviewed-on: https://go-review.googlesource.com/38791
Run-TryBot: Sam Whited <sam@samwhited.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Ian Lance Taylor <iant@golang.org>
---
 src/encoding/xml/xml.go      | 37 +++++++++++++++++++++
 src/encoding/xml/xml_test.go | 64 ++++++++++++++++++++++++++++++++++++
 2 files changed, 101 insertions(+)

diff --git a/src/encoding/xml/xml.go b/src/encoding/xml/xml.go
index 1aba31ce60..be90b62c9a 100644
--- a/src/encoding/xml/xml.go
+++ b/src/encoding/xml/xml.go
@@ -135,6 +135,23 @@ func CopyToken(t Token) Token {
 	return t
 }
 
+// A TokenReader is anything that can decode a stream of XML tokens, including a
+// Decoder.
+//
+// When Token encounters an error or end-of-file condition after successfully
+// reading a token, it returns the token. It may return the (non-nil) error from
+// the same call or return the error (and a nil token) from a subsequent call.
+// An instance of this general case is that a TokenReader returning a non-nil
+// token at the end of the token stream may return either io.EOF or a nil error.
+// The next Read should return nil, io.EOF.
+//
+// Implementations of Token are discouraged from returning a nil token with a
+// nil error. Callers should treat a return of nil, nil as indicating that
+// nothing happened; in particular it does not indicate EOF.
+type TokenReader interface {
+	Token() (Token, error)
+}
+
 // A Decoder represents an XML parser reading a particular input stream.
 // The parser assumes that its input is encoded in UTF-8.
 type Decoder struct {
@@ -190,6 +207,7 @@ type Decoder struct {
 	DefaultSpace string
 
 	r              io.ByteReader
+	t              TokenReader
 	buf            bytes.Buffer
 	saved          *bytes.Buffer
 	stk            *stack
@@ -219,6 +237,22 @@ func NewDecoder(r io.Reader) *Decoder {
 	return d
 }
 
+// NewTokenDecoder creates a new XML parser using an underlying token stream.
+func NewTokenDecoder(t TokenReader) *Decoder {
+	// Is it already a Decoder?
+	if d, ok := t.(*Decoder); ok {
+		return d
+	}
+	d := &Decoder{
+		ns:       make(map[string]string),
+		t:        t,
+		nextByte: -1,
+		line:     1,
+		Strict:   true,
+	}
+	return d
+}
+
 // Token returns the next XML token in the input stream.
 // At the end of the input stream, Token returns nil, io.EOF.
 //
@@ -243,6 +277,9 @@ func NewDecoder(r io.Reader) *Decoder {
 // If Token encounters an unrecognized name space prefix,
 // it uses the prefix as the Space rather than report an error.
 func (d *Decoder) Token() (Token, error) {
+	if d.t != nil {
+		return d.t.Token()
+	}
 	var t Token
 	var err error
 	if d.stk != nil && d.stk.kind == stkEOF {
diff --git a/src/encoding/xml/xml_test.go b/src/encoding/xml/xml_test.go
index 7950ca22c4..2437f19d9d 100644
--- a/src/encoding/xml/xml_test.go
+++ b/src/encoding/xml/xml_test.go
@@ -797,3 +797,67 @@ func TestIssue12417(t *testing.T) {
 		}
 	}
 }
+
+func tokenMap(mapping func(t Token) Token) func(TokenReader) TokenReader {
+	return func(src TokenReader) TokenReader {
+		return mapper{
+			t: src,
+			f: mapping,
+		}
+	}
+}
+
+type mapper struct {
+	t TokenReader
+	f func(Token) Token
+}
+
+func (m mapper) Token() (Token, error) {
+	tok, err := m.t.Token()
+	if err != nil {
+		return nil, err
+	}
+	return m.f(tok), nil
+}
+
+func TestNewTokenDecoderIdempotent(t *testing.T) {
+	d := NewDecoder(strings.NewReader(`<br/>`))
+	d2 := NewTokenDecoder(d)
+	if d != d2 {
+		t.Error("NewTokenDecoder did not detect underlying Decoder")
+	}
+}
+
+func TestWrapDecoder(t *testing.T) {
+	d := NewDecoder(strings.NewReader(`<quote>[Re-enter Clown with a letter, and FABIAN]</quote>`))
+	m := tokenMap(func(t Token) Token {
+		switch tok := t.(type) {
+		case StartElement:
+			if tok.Name.Local == "quote" {
+				tok.Name.Local = "blocking"
+				return tok
+			}
+		case EndElement:
+			if tok.Name.Local == "quote" {
+				tok.Name.Local = "blocking"
+				return tok
+			}
+		}
+		return t
+	})
+
+	d = NewTokenDecoder(m(d))
+
+	o := struct {
+		XMLName  Name   `xml:"blocking"`
+		Chardata string `xml:",chardata"`
+	}{}
+
+	if err := d.Decode(&o); err != nil {
+		t.Fatal("Got unexpected error while decoding:", err)
+	}
+
+	if o.Chardata != "[Re-enter Clown with a letter, and FABIAN]" {
+		t.Fatalf("Got unexpected chardata: `%s`\n", o.Chardata)
+	}
+}
-- 
2.48.1