From eda5c23c54ecb484fedf1f375d2a25a6047f1dbe Mon Sep 17 00:00:00 2001
From: Wichert Akkerman <wichert@wiggy.net>
Date: Thu, 31 Oct 2019 20:08:39 +0100
Subject: [PATCH] Add SendIQ to StreamClient and Sender

This makes it possible to use SendIQ from PostConnect and route handlers.
---
 component.go      | 20 ++++++++++++++++++++
 stream_manager.go |  3 +++
 2 files changed, 23 insertions(+)

diff --git a/component.go b/component.go
index 18be3bc..d459c00 100644
--- a/component.go
+++ b/component.go
@@ -1,6 +1,7 @@
 package xmpp
 
 import (
+	"context"
 	"crypto/sha1"
 	"encoding/hex"
 	"encoding/xml"
@@ -158,6 +159,25 @@ func (c *Component) Send(packet stanza.Packet) error {
 	return nil
 }
 
+// SendIQ sends an IQ set or get stanza to the server. If a result is received
+// the provided handler function will automatically be called.
+//
+// The provided context should have a timeout to prevent the client from waiting
+// forever for an IQ result. For example:
+//
+//   ctx, _ := context.WithTimeout(context.Background(), 30 * time.Second)
+//   result := <- client.SendIQ(ctx, iq)
+//
+func (c *Component) SendIQ(ctx context.Context, iq stanza.IQ) (chan stanza.IQ, error) {
+	if iq.Attrs.Type != "set" && iq.Attrs.Type != "get" {
+		return nil, ErrCanOnlySendGetOrSetIq
+	}
+	if err := c.Send(iq); err != nil {
+		return nil, err
+	}
+	return c.router.NewIQResultRoute(ctx, iq.Attrs.Id), nil
+}
+
 // SendRaw sends an XMPP stanza as a string to the server.
 // It can be invalid XML or XMPP content. In that case, the server will
 // disconnect the component. It is up to the user of this method to
diff --git a/stream_manager.go b/stream_manager.go
index c21d79a..1011f6e 100644
--- a/stream_manager.go
+++ b/stream_manager.go
@@ -1,6 +1,7 @@
 package xmpp
 
 import (
+	"context"
 	"errors"
 	"sync"
 	"time"
@@ -26,6 +27,7 @@ type StreamClient interface {
 	Connect() error
 	Resume(state SMState) error
 	Send(packet stanza.Packet) error
+	SendIQ(ctx context.Context, iq stanza.IQ) (chan stanza.IQ, error)
 	SendRaw(packet string) error
 	Disconnect()
 	SetHandler(handler EventHandler)
@@ -35,6 +37,7 @@ type StreamClient interface {
 // It is mostly use in callback to pass a limited subset of the stream client interface
 type Sender interface {
 	Send(packet stanza.Packet) error
+	SendIQ(ctx context.Context, iq stanza.IQ) (chan stanza.IQ, error)
 	SendRaw(packet string) error
 }