From 8ab32d885ffdf18e7173bd05f13a8a645525af03 Mon Sep 17 00:00:00 2001
From: Martin Dosch <martin@mdosch.de>
Date: Sat, 16 Mar 2024 19:04:09 +0100
Subject: [PATCH] Fix race condition for nextStart and nextEnd.

---
 xmpp.go | 26 +++++++++++++++++++++++---
 1 file changed, 23 insertions(+), 3 deletions(-)

diff --git a/xmpp.go b/xmpp.go
index 463c57d..f2eac8c 100644
--- a/xmpp.go
+++ b/xmpp.go
@@ -38,6 +38,7 @@ import (
 	"slices"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"golang.org/x/crypto/pbkdf2"
@@ -60,6 +61,9 @@ var DefaultConfig = &tls.Config{}
 // DebugWriter is the writer used to write debugging output to.
 var DebugWriter io.Writer = os.Stderr
 
+// Mutex to prevent multiple access to xml.Decoder
+var nextMutex sync.Mutex
+
 // Cookie is a unique XMPP session identifier
 type Cookie uint64
 
@@ -917,7 +921,6 @@ type IQ struct {
 }
 
 // Recv waits to receive the next XMPP stanza.
-// Return type is either a presence notification or a chat message.
 func (c *Client) Recv() (stanza interface{}, err error) {
 	for {
 		_, val, err := next(c.p)
@@ -1461,14 +1464,26 @@ type rosterItem struct {
 // Scan XML token stream to find next StartElement.
 func nextStart(p *xml.Decoder) (xml.StartElement, error) {
 	for {
+		nextMutex.Lock()
 		t, err := p.Token()
 		if err != nil || t == nil {
+			nextMutex.Unlock()
 			return xml.StartElement{}, err
 		}
 		switch t := t.(type) {
 		case xml.StartElement:
+			nextMutex.Unlock()
 			return t, nil
+		// Also check for stream end element and stop waiting
+		// for new start elements if we received a closing stream
+		// element.
+		case xml.EndElement:
+			if t.Name.Local == "stream" {
+				nextMutex.Unlock()
+				return xml.StartElement{}, nil
+			}
 		}
+		nextMutex.Unlock()
 	}
 }
 
@@ -1476,14 +1491,19 @@ func nextStart(p *xml.Decoder) (xml.StartElement, error) {
 func nextEnd(p *xml.Decoder) (xml.EndElement, error) {
 	p.Strict = false
 	for {
-		t, err := p.RawToken()
-		if err != nil || t == nil {
+		nextMutex.Lock()
+		to, err := p.RawToken()
+		if err != nil || to == nil {
+			nextMutex.Unlock()
 			return xml.EndElement{}, err
 		}
+		t := xml.CopyToken(to)
 		switch t := t.(type) {
 		case xml.EndElement:
+			nextMutex.Unlock()
 			return t, nil
 		}
+		nextMutex.Unlock()
 	}
 }