From aed021cf3ef3542d6b20be43bd0755ed08c35520 Mon Sep 17 00:00:00 2001
From: Martin Dosch <martin@mdosch.de>
Date: Sat, 2 Mar 2024 14:17:47 +0100
Subject: [PATCH] Wait for the server closing the stream before closing the
 connection.

---
 xmpp.go | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/xmpp.go b/xmpp.go
index 85957c0..58bc88b 100644
--- a/xmpp.go
+++ b/xmpp.go
@@ -331,6 +331,21 @@ func NewClientNoTLS(host, user, passwd string, debug bool) (*Client, error) {
 func (c *Client) Close() error {
 	if c.conn != (*tls.Conn)(nil) {
 		fmt.Fprintf(c.stanzaWriter, "</stream:stream>\n")
+		// Wait for the server also closing the stream.
+		for {
+			select {
+			case <-time.After(10 * time.Second):
+				break
+			default:
+				ee, err := nextEnd(c.p)
+				if err != nil {
+					return err
+				}
+				if ee.Name.Local == "stream" {
+					return c.conn.Close()
+				}
+			}
+		}
 		return c.conn.Close()
 	}
 	return nil
@@ -1421,6 +1436,20 @@ func nextStart(p *xml.Decoder) (xml.StartElement, error) {
 	}
 }
 
+// Scan XML token stream to find next EndElement
+func nextEnd(p *xml.Decoder) (xml.EndElement, error) {
+	for {
+		t, err := p.Token()
+		if err != nil || t == nil {
+			return xml.EndElement{}, err
+		}
+		switch t := t.(type) {
+		case xml.EndElement:
+			return t, nil
+		}
+	}
+}
+
 // Scan XML token stream for next element and save into val.
 // If val == nil, allocate new element based on proto map.
 // Either way, return val.