3 files changed, 30 insertions(+), 41 deletions(-)

M go.mod
M go.sum
M session.go

@@ 8,5 8,5 @@ require (
 	golang.org/x/net v0.0.0-20190813141303-74dc4d7220e7
 	golang.org/x/text v0.3.2
 	mellium.im/sasl v0.2.1
-	mellium.im/xmlstream v0.15.0
+	mellium.im/xmlstream v0.15.1
 )

@@ 17,5 17,5 @@ mellium.im/reader v0.1.0 h1:UUEMev16gdvaxxZC7fC08j7IzuDKh310nB6BlwnxTww=
 mellium.im/reader v0.1.0/go.mod h1:F+X5HXpkIfJ9EE1zHQG9lM/hO946iYAmU7xjg5dsQHI=
 mellium.im/sasl v0.2.1 h1:nspKSRg7/SyO0cRGY71OkfHab8tf9kCts6a6oTDut0w=
 mellium.im/sasl v0.2.1/go.mod h1:ROaEDLQNuf9vjKqE1SrAfnsobm2YKXT1gnN1uDp1PjQ=
-mellium.im/xmlstream v0.15.0 h1:NczJZ5FYsRhaA2asw0/hrQm83K81cSTJszKhHh4s18Q=
-mellium.im/xmlstream v0.15.0/go.mod h1:7SUlP7f2qnMczK+Cu/OFgqaIhldMolVjo8np7xG41D0=
+mellium.im/xmlstream v0.15.1 h1:qqv+SUESVFPf3NZUGFldWY4KesQVpamnRVxwceK2PYA=
+mellium.im/xmlstream v0.15.1/go.mod h1:7SUlP7f2qnMczK+Cu/OFgqaIhldMolVjo8np7xG41D0=

@@ 99,7 99,10 @@ type Session struct {
 	}
 	out struct {
 		intstream.Info
-		e tokenWriteFlusher
+		e interface {
+			xmlstream.TokenWriter
+			xmlstream.Flusher
+		}
 		sync.Locker
 	}
 }


@@ 179,7 182,7 @@ func NegotiateSession(ctx context.Context, location, origin jid.JID, rw io.ReadW
 	}
 
 	s.in.d = intstream.Reader(s.in.d)
-	s.out.e = stanzaAddID(s.out.e)
+	s.out.e = &stanzaEncoder{TokenWriteFlusher: s.out.e}
 
 	return s, nil
 }


@@ 918,45 921,31 @@ func (s *Session) closeInputStream() {
 	s.in.cancel()
 }
 
-type wrapWriter struct {
-	encode func(t xml.Token) error
-	flush  func() error
+type stanzaEncoder struct {
+	xmlstream.TokenWriteFlusher
+	depth int
 }
 
-func (w wrapWriter) EncodeToken(t xml.Token) error { return w.encode(t) }
-func (w wrapWriter) Flush() error                  { return w.flush() }
-
-type tokenWriteFlusher interface {
-	xmlstream.TokenWriter
-	xmlstream.Flusher
-}
-
-func stanzaAddID(w tokenWriteFlusher) tokenWriteFlusher {
-	depth := 0
-	return wrapWriter{
-		encode: func(t xml.Token) error {
-		tokswitch:
-			switch tok := t.(type) {
-			case xml.StartElement:
-				depth++
-				if depth == 1 && tok.Name.Local == "iq" {
-					for _, attr := range tok.Attr {
-						if attr.Name.Local == "id" {
-							break tokswitch
-						}
-					}
-					tok.Attr = append(tok.Attr, xml.Attr{
-						Name:  xml.Name{Local: "id"},
-						Value: attr.RandomID(),
-					})
-					t = tok
+func (se *stanzaEncoder) EncodeToken(t xml.Token) error {
+tokswitch:
+	switch tok := t.(type) {
+	case xml.StartElement:
+		se.depth++
+		if se.depth == 1 && tok.Name.Local == "iq" {
+			for _, attr := range tok.Attr {
+				if attr.Name.Local == "id" {
+					break tokswitch
 				}
-			case xml.EndElement:
-				depth--
 			}
-
-			return w.EncodeToken(t)
-		},
-		flush: w.Flush,
+			tok.Attr = append(tok.Attr, xml.Attr{
+				Name:  xml.Name{Local: "id"},
+				Value: attr.RandomID(),
+			})
+			t = tok
+		}
+	case xml.EndElement:
+		se.depth--
 	}
+
+	return se.TokenWriteFlusher.EncodeToken(t)
 }