~samwhited/xmpp

ref: 6026cbc20a8f0e1269108b97cb7d6853fa2975ce xmpp/internal/xmpptest/session.go -rw-r--r-- 4.6 KiB
6026cbc2Sam Whited xmpp: add In and Out methods to session 3 months 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
// Copyright 2017 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 xmpptest provides utilities for XMPP testing.
package xmpptest // import "mellium.im/xmpp/internal/xmpptest"

import (
	"context"
	"io"
	"net"
	"strings"

	"mellium.im/xmpp"
	"mellium.im/xmpp/internal/ns"
	intstream "mellium.im/xmpp/internal/stream"
	"mellium.im/xmpp/jid"
	"mellium.im/xmpp/stream"
)

// NopNegotiator marks the state as ready (by returning state|xmpp.Ready) and
// pops the first token (likely <stream:stream>) but does not perform any
// validation on the token, transmit any data over the wire, or perform any
// other session negotiation.
func NopNegotiator(state xmpp.SessionState) xmpp.Negotiator {
	return func(ctx context.Context, in, out *stream.Info, s *xmpp.Session, data interface{}) (xmpp.SessionState, io.ReadWriter, interface{}, error) {
		// Pop the stream start token.
		rc := s.TokenReader()
		defer rc.Close()

		err := intstream.Expect(ctx, in, rc, s.State()&xmpp.Received == xmpp.Received, false)
		if err != nil {
			return state | xmpp.Ready, nil, nil, err
		}
		err = intstream.Send(struct {
			io.Reader
			io.Writer
		}{
			Writer: io.Discard,
		}, out, s.State()&xmpp.S2S == xmpp.S2S, false, stream.DefaultVersion, "", "example.net", "test@example.net", "123")

		return state | xmpp.Ready, nil, nil, err
	}
}

// NewSession returns a new client-to-client XMPP session with the state bits
// set to finalState|xmpp.Ready, the origin JID set to "test@example.net" and
// the location JID set to "example.net".
//
// NewSession panics on error for ease of use in testing, where a panic is
// acceptable.
func NewSession(finalState xmpp.SessionState, rw io.ReadWriter) *xmpp.Session {
	location := jid.MustParse("example.net")
	origin := jid.MustParse("test@example.net")

	to, from := origin, location
	if finalState&xmpp.Received == xmpp.Received {
		to, from = from, to
	}

	s, err := xmpp.NewSession(
		context.Background(), location, origin,
		struct {
			io.Reader
			io.Writer
		}{
			Reader: io.MultiReader(
				strings.NewReader(`<stream:stream from="`+from.String()+`" to="`+to.String()+`" id="123" version="1.0" xmlns="`+ns.Client+`" xmlns:stream="`+stream.NS+`">`),
				rw,
				strings.NewReader(`</stream:stream>`),
			),
			Writer: rw,
		},
		0,
		NopNegotiator(finalState),
	)
	if err != nil {
		panic(err)
	}
	return s
}

// Option is a type for configuring a ClientServer.
type Option func(*ClientServer)

// ClientState configures extra state bits to add to the client session.
func ClientState(state xmpp.SessionState) Option {
	return func(c *ClientServer) {
		c.clientState |= state
	}
}

// ServerState configures extra state bits to add to the server session.
func ServerState(state xmpp.SessionState) Option {
	return func(c *ClientServer) {
		c.serverState |= state
	}
}

// ClientHandler sets up the client side of a ClientServer.
func ClientHandler(handler xmpp.Handler) Option {
	return func(c *ClientServer) {
		c.clientHandler = handler
	}
}

// ClientHandlerFunc sets up the client side of a ClientServer using an
// xmpp.HandlerFunc.
func ClientHandlerFunc(handler xmpp.HandlerFunc) Option {
	return ClientHandler(handler)
}

// ServerHandler sets up the server side of a ClientServer.
func ServerHandler(handler xmpp.Handler) Option {
	return func(c *ClientServer) {
		c.serverHandler = handler
	}
}

// ServerHandlerFunc sets up the server side of a ClientServer using an
// xmpp.HandlerFunc.
func ServerHandlerFunc(handler xmpp.HandlerFunc) Option {
	return ServerHandler(handler)
}

// ClientServer is two coupled xmpp.Session's that can respond to one another in
// tests.
type ClientServer struct {
	Client *xmpp.Session
	Server *xmpp.Session

	clientHandler xmpp.Handler
	serverHandler xmpp.Handler
	clientState   xmpp.SessionState
	serverState   xmpp.SessionState
}

// NewClientServer returns a ClientServer with the client and server goroutines
// started.
// Both serve goroutines are started when NewClientServer is called and shut
// down when the ClientServer is closed.
func NewClientServer(opts ...Option) *ClientServer {
	cs := &ClientServer{
		serverState: xmpp.Received,
	}
	for _, opt := range opts {
		opt(cs)
	}

	clientConn, serverConn := net.Pipe()
	cs.Client = NewSession(cs.clientState, clientConn)
	cs.Server = NewSession(cs.serverState, serverConn)
	/* #nosec */
	go cs.Client.Serve(cs.clientHandler)
	/* #nosec */
	go cs.Server.Serve(cs.serverHandler)
	return cs
}

// Close calls the client and server sessions' close methods.
func (cs *ClientServer) Close() error {
	err := cs.Client.Close()
	if err != nil {
		return err
	}
	return cs.Server.Close()
}