-
Notifications
You must be signed in to change notification settings - Fork 18
/
dialer.go
68 lines (59 loc) · 2.58 KB
/
dialer.go
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
// Package wstest provides a NewDialer function to test just the
// `http.Handler` that upgrades the connection to a websocket session.
// It runs the handler function in a goroutine without listening on
// any port. The returned `websocket.Dialer` then can be used to dial
// and communicate with the given handler.
package wstest
import (
"bufio"
"net"
"net/http"
"net/http/httptest"
"github.com/gorilla/websocket"
)
// NewDialer creates a wstest recorder to an http.Handler which accepts websocket upgrades.
// This send an HTTP request to the http.Handler, and wait for the connection upgrade response.
// it runs the recorder's ServeHTTP function in a goroutine, so recorder can communicate with a
// client running on the current program flow
//
// h is an http.Handler that handles websocket connections.
// It returns a *websocket.Dial struct, which can then be used to dial to the handler.
func NewDialer(h http.Handler) *websocket.Dialer {
client, server := net.Pipe()
conn := &recorder{server: server}
// run the runServer in a goroutine, so when the Dial send the request to
// the recorder on the connection, it will be parsed as an HTTPRequest and
// sent to the Handler function.
go conn.runServer(h)
// use the websocket.NewDialer.Dial with the fake net.recorder to communicate with the recorder
// the recorder gets the client which is the client side of the connection
return &websocket.Dialer{NetDial: func(network, addr string) (net.Conn, error) { return client, nil }}
}
// recorder it similar to httptest.ResponseRecorder, but with Hijack capabilities
type recorder struct {
httptest.ResponseRecorder
server net.Conn
}
// runServer reads the request sent on the connection to the recorder
// from the websocket.NewDialer.Dial function, and pass it to the recorder.
// once this is done, the communication is done on the wsConn
func (r *recorder) runServer(h http.Handler) {
// read from the recorder connection the request sent by the recorder.Dial,
// and use the handler to serve this request.
req, err := http.ReadRequest(bufio.NewReader(r.server))
if err != nil {
return
}
h.ServeHTTP(r, req)
}
// Hijack the connection
func (r *recorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
// return to the recorder the recorder, which is the recorder side of the connection
rw := bufio.NewReadWriter(bufio.NewReader(r.server), bufio.NewWriter(r.server))
return r.server, rw, nil
}
// WriteHeader write HTTP header to the client and closes the connection
func (r *recorder) WriteHeader(code int) {
resp := http.Response{StatusCode: code, Header: r.Header()}
resp.Write(r.server)
}