~samwhited/xmpp

xmpp/session_iq.go -rw-r--r-- 6.8 KiB
dcc6772fSam Whited all: add strict namespace checking to mux 22 hours ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
// Copyright 2021 The Mellium Contributors.
// Use of this source code is governed by the BSD 2-clause
// license that can be found in the LICENSE file.

package xmpp

import (
	"context"
	"encoding/xml"
	"fmt"
	"io"

	"mellium.im/xmlstream"
	"mellium.im/xmpp/internal/attr"
	"mellium.im/xmpp/internal/marshal"
	"mellium.im/xmpp/stanza"
)

// SendIQ is like Send except that it returns an error if the first token read
// from the stream is not an Info/Query (IQ) start element and blocks until a
// response is received.
//
// If the input stream is not being processed (a call to Serve is not running),
// SendIQ will never receive a response and will block until the provided
// context is canceled.
// If the response is non-nil, it does not need to be consumed in its entirety,
// but it must be closed before stream processing will resume.
// If the IQ type does not require a response—ie. it is a result or error IQ,
// meaning that it is a response itself—SendIQElement does not block and the
// response is nil.
//
// If the context is closed before the response is received, SendIQ immediately
// returns the context error.
// Any response received at a later time will not be associated with the
// original request but can still be handled by the Serve handler.
//
// If an error is returned, the response will be nil; the converse is not
// necessarily true.
// SendIQ is safe for concurrent use by multiple goroutines.
func (s *Session) SendIQ(ctx context.Context, r xml.TokenReader) (xmlstream.TokenReadCloser, error) {
	tok, err := r.Token()
	if err != nil {
		return nil, err
	}
	start, ok := tok.(xml.StartElement)
	if !ok {
		return nil, fmt.Errorf("expected IQ start element, got %T", tok)
	}
	if !isIQEmptySpace(start.Name) {
		return nil, fmt.Errorf("expected start element to be an IQ")
	}

	// If there's no ID, add one.
	idx, _, id, typ := getIDTyp(start.Attr)
	if idx == -1 {
		idx = len(start.Attr)
		start.Attr = append(start.Attr, xml.Attr{Name: xml.Name{Local: "id"}, Value: ""})
	}
	if id == "" {
		id = attr.RandomID()
		start.Attr[idx].Value = id
	}

	// If this an IQ of type "set" or "get" we expect a response.
	if typ == string(stanza.GetIQ) || typ == string(stanza.SetIQ) {
		return s.sendResp(ctx, id, xmlstream.Inner(r), start)
	}

	// If this is an IQ of type result or error, we don't expect a response so
	// just send it normally.
	return nil, s.SendElement(ctx, xmlstream.Inner(r), start)
}

// SendIQElement is like SendIQ except that it wraps the payload in an
// Info/Query (IQ) element.
// For more information see SendIQ.
//
// SendIQElement is safe for concurrent use by multiple goroutines.
func (s *Session) SendIQElement(ctx context.Context, payload xml.TokenReader, iq stanza.IQ) (xmlstream.TokenReadCloser, error) {
	return s.SendIQ(ctx, iq.Wrap(payload))
}

// EncodeIQ is like Encode except that it returns an error if v does not marshal
// to an IQ stanza and like SendIQ it blocks until a response is received.
// For more information see SendIQ.
//
// EncodeIQ is safe for concurrent use by multiple goroutines.
func (s *Session) EncodeIQ(ctx context.Context, v interface{}) (xmlstream.TokenReadCloser, error) {
	r, err := marshal.TokenReader(v)
	if err != nil {
		return nil, err
	}
	return s.SendIQ(ctx, r)
}

// EncodeIQElement is like EncodeIQ except that it wraps the payload in an
// Info/Query (IQ) element.
// For more information see SendIQ.
//
// EncodeIQElement is safe for concurrent use by multiple goroutines.
func (s *Session) EncodeIQElement(ctx context.Context, payload interface{}, iq stanza.IQ) (xmlstream.TokenReadCloser, error) {
	r, err := marshal.TokenReader(payload)
	if err != nil {
		return nil, err
	}
	return s.SendIQElement(ctx, r, iq)
}

// UnmarshalIQ is like SendIQ except that error replies are unmarshaled into a
// stanza.Error and returned and otherwise the response payload is unmarshaled
// into v.
// For more information see SendIQ.
//
// UnmarshalIQ is safe for concurrent use by multiple goroutines.
func (s *Session) UnmarshalIQ(ctx context.Context, iq xml.TokenReader, v interface{}) error {
	return unmarshalIQ(ctx, iq, v, s)
}

// UnmarshalIQElement is like UnmarshalIQ but it wraps a payload in the provided IQ.
// For more information see SendIQ.
//
// UnmarshalIQElement is safe for concurrent use by multiple goroutines.
func (s *Session) UnmarshalIQElement(ctx context.Context, payload xml.TokenReader, iq stanza.IQ, v interface{}) error {
	return unmarshalIQ(ctx, iq.Wrap(payload), v, s)
}

// IterIQ is like SendIQ except that error replies are unmarshaled into a
// stanza.Error and returned and otherwise an iterator over the children of the
// response payload is returned.
// For more information see SendIQ.
//
// IterIQ is safe for concurrent use by multiple goroutines.
func (s *Session) IterIQ(ctx context.Context, iq xml.TokenReader) (*xmlstream.Iter, *xml.StartElement, error) {
	return iterIQ(ctx, iq, s)
}

// IterIQElement is like IterIQ but it wraps a payload in the provided IQ.
// For more information see SendIQ.
//
// IterIQElement is safe for concurrent use by multiple goroutines.
func (s *Session) IterIQElement(ctx context.Context, payload xml.TokenReader, iq stanza.IQ) (*xmlstream.Iter, *xml.StartElement, error) {
	return iterIQ(ctx, iq.Wrap(payload), s)
}

func iterIQ(ctx context.Context, iq xml.TokenReader, s *Session) (_ *xmlstream.Iter, _ *xml.StartElement, e error) {
	resp, err := s.SendIQ(ctx, iq)
	if err != nil {
		return nil, nil, err
	}
	defer func() {
		if e != nil {
			/* #nosec */
			resp.Close()
		}
	}()

	tok, err := resp.Token()
	if err != nil {
		return nil, nil, err
	}

	start, ok := tok.(xml.StartElement)
	if !ok {
		return nil, nil, fmt.Errorf("stanza: expected IQ start token, got %T %[1]v", tok)
	}
	_, err = stanza.UnmarshalIQError(resp, start)
	if err != nil {
		return nil, nil, err
	}

	// Pop the payload start token, we want to iterate over its children.
	tok, err = resp.Token()
	start, _ = tok.(xml.StartElement)

	// Discard early EOF so that the iterator doesn't end up returning it.
	if err != nil && err != io.EOF {
		return nil, nil, err
	}
	return xmlstream.NewIter(resp), &start, nil
}

func unmarshalIQ(ctx context.Context, iq xml.TokenReader, v interface{}, s *Session) (e error) {
	resp, err := s.SendIQ(ctx, iq)
	if err != nil {
		return err
	}
	defer func() {
		ee := resp.Close()
		if e == nil {
			e = ee
		}
	}()

	tok, err := resp.Token()
	if err != nil {
		return err
	}
	start, ok := tok.(xml.StartElement)
	if !ok {
		return fmt.Errorf("stanza: expected IQ start token, got %T %[1]v", tok)
	}

	_, err = stanza.UnmarshalIQError(resp, start)
	if err != nil {
		return err
	}
	if v == nil {
		return nil
	}
	payload := xmlstream.Inner(resp)
	d := xml.NewTokenDecoder(payload)
	startTok, err := d.Token()
	switch err {
	case io.EOF:
		return nil
	case nil:
	default:
		return err
	}
	start = startTok.(xml.StartElement)
	return d.DecodeElement(v, &start)
}