Reintroduce context and cancel func

There are still some rare races without.

main.go

@@ -398,7 +398,8 @@ func main() {
 	}
 	iqc := make(chan xmpp.IQ, defaultBufferSize)
 	msgc := make(chan xmpp.Chat, defaultBufferSize)
-	go rcvStanzas(client, iqc, msgc)
+	ctx, cancel := context.WithCancel(context.Background())
+	go rcvStanzas(client, ctx, iqc, msgc)
 	for _, r := range getopt.Args() {
 		var re recipientsType
 		re.Jid = r
@@ -416,6 +417,7 @@ func main() {
 	for i, recipient := range recipients {
 		validatedJid, err := MarshalJID(recipient.Jid)
 		if err != nil {
+			cancel()
 			closeAndExit(client, err)
 		}
 		recipients[i].Jid = validatedJid
@@ -425,51 +427,61 @@ func main() {
 	case *flagOxGenPrivKeyX25519:
 		validatedOwnJid, err := MarshalJID(user)
 		if err != nil {
+			cancel()
 			closeAndExit(client, err)
 		}
 		err = oxGenPrivKey(validatedOwnJid, client, iqc, *flagOxPassphrase, "x25519")
 		if err != nil {
+			cancel()
 			closeAndExit(client, err)
 		}
 		os.Exit(0)
 	case *flagOxGenPrivKeyRSA:
 		validatedOwnJid, err := MarshalJID(user)
 		if err != nil {
+			cancel()
 			closeAndExit(client, err)
 		}
 		err = oxGenPrivKey(validatedOwnJid, client, iqc, *flagOxPassphrase, "rsa")
 		if err != nil {
+			cancel()
 			closeAndExit(client, err)
 		}
 		os.Exit(0)
 	case *flagOxImportPrivKey != "":
 		validatedOwnJid, err := MarshalJID(user)
 		if err != nil {
+			cancel()
 			closeAndExit(client, err)
 		}
 		err = oxImportPrivKey(validatedOwnJid, *flagOxImportPrivKey,
 			client, iqc)
 		if err != nil {
+			cancel()
 			closeAndExit(client, err)
 		}
 		os.Exit(0)
 	case *flagOxDeleteNodes:
 		validatedOwnJid, err := MarshalJID(user)
 		if err != nil {
+			cancel()
 			closeAndExit(client, err)
 		}
 		err = oxDeleteNodes(validatedOwnJid, client, iqc)
 		if err != nil {
+			cancel()
 			closeAndExit(client, err)
 		}
 		os.Exit(0)
 	case *flagOx:
 		validatedOwnJid, err := MarshalJID(user)
 		if err != nil {
+			cancel()
 			closeAndExit(client, err)
 		}
 		oxPrivKey, err = oxGetPrivKey(validatedOwnJid, *flagOxPassphrase)
 		if err != nil {
+			cancel()
 			closeAndExit(client, err)
 		}
 	}
@@ -478,6 +490,7 @@ func main() {
 		message, err = httpUpload(client, iqc, tlsConfig.ServerName,
 			*flagHTTPUpload, timeout)
 		if err != nil {
+			cancel()
 			closeAndExit(client, err)
 		}
 	}
@@ -488,6 +501,7 @@ func main() {
 		// Check if the URI is valid.
 		uri, err := validURI(message)
 		if err != nil {
+			cancel()
 			closeAndExit(client, err)
 		}
 		message = uri.String()
@@ -511,6 +525,7 @@ func main() {
 				_, err = client.JoinMUCNoHistory(recipient.Jid, alias)
 			}
 			if err != nil {
+				cancel()
 				closeAndExit(client, err)
 			}
 		}
@@ -523,6 +538,7 @@ func main() {
 		// Send raw XML
 		_, err = client.SendOrg(message)
 		if err != nil {
+			cancel()
 			closeAndExit(client, err)
 		}
 	case *flagInteractive:
@@ -531,7 +547,6 @@ func main() {
 		// Quit if ^C is pressed.
 		c := make(chan os.Signal, 1)
 		signal.Notify(c, os.Interrupt)
-		ctx, cancel := context.WithCancel(context.Background())
 		go func() {
 			for range c {
 				cancel()
@@ -546,6 +561,7 @@ func main() {
 					return
 				default:
 					if err != nil {
+						cancel()
 						closeAndExit(client, errors.New("failed to read from stdin"))
 					}
 				}
@@ -572,6 +588,7 @@ func main() {
 					}
 					_, err = client.SendOrg(oxMessage)
 					if err != nil {
+						cancel()
 						closeAndExit(client, err)
 					}
 				default:
@@ -580,6 +597,7 @@ func main() {
 						Type:   msgType, Text: message,
 					})
 					if err != nil {
+						cancel()
 						closeAndExit(client, err)
 					}
 				}
@@ -592,6 +610,7 @@ func main() {
 		signal.Notify(c, os.Interrupt)
 		go func() {
 			for range c {
+				cancel()
 				closeAndExit(client, nil)
 			}
 		}()
@@ -699,6 +718,7 @@ func main() {
 				}
 				_, err = client.SendOrg(oxMessage)
 				if err != nil {
+					cancel()
 					closeAndExit(client, err)
 				}
 			default:
@@ -707,10 +727,12 @@ func main() {
 					Type:   msgType, Text: message,
 				})
 				if err != nil {
+					cancel()
 					closeAndExit(client, err)
 				}
 			}
 		}
 	}
+	cancel()
 	closeAndExit(client, nil)
 }

stanzahandling.go

@@ -5,6 +5,7 @@
 package main
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -44,15 +45,34 @@ func getIQ(id string, c chan xmpp.IQ, iqc chan xmpp.IQ) {
 	}
 }
 
-func rcvStanzas(client *xmpp.Client, iqc chan xmpp.IQ, msgc chan xmpp.Chat) {
+func rcvStanzas(client *xmpp.Client, ctx context.Context, iqc chan xmpp.IQ, msgc chan xmpp.Chat) {
+	var received interface{}
+	var err error
 	for {
-		received, err := client.Recv()
-		if err != nil {
-			if err != io.EOF {
-				closeAndExit(client, err)
-				return
+		r := make(chan interface{})
+		e := make(chan error)
+		go func() {
+			for {
+				rcv, err := client.Recv()
+				if err != nil {
+					e <- err
+				} else {
+					r <- rcv
+				}
 			}
+		}()
+		select {
+		case <-ctx.Done():
 			return
+		case err = <-e:
+			if err != nil {
+				if err != io.EOF {
+					closeAndExit(client, err)
+					return
+				}
+				return
+			}
+		case received = <-r:
 		}
 		switch v := received.(type) {
 		case xmpp.Chat: