-
-
Notifications
You must be signed in to change notification settings - Fork 118
/
client.go
148 lines (123 loc) · 3.6 KB
/
client.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
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
// Copyright 2013 The Gorilla WebSocket Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package websockets
import (
"net/http"
"time"
"github.com/axllent/mailpit/internal/auth"
"github.com/axllent/mailpit/internal/logger"
"github.com/gorilla/websocket"
)
const (
// Time allowed to write a message to the peer.
writeWait = 10 * time.Second
// Time allowed to read the next pong message from the peer.
pongWait = 60 * time.Second
// Send pings to peer with this period. Must be less than pongWait.
pingPeriod = (pongWait * 9) / 10
)
var (
newline = []byte{'\n'}
// MessageHub global
MessageHub *Hub
)
var upgrader = websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
CheckOrigin: func(r *http.Request) bool { return true }, // allow multi-domain
EnableCompression: true, // experimental compression
}
// Client is a middleman between the websocket connection and the hub.
type Client struct {
hub *Hub
// The websocket connection.
conn *websocket.Conn
// Buffered channel of outbound messages.
send chan []byte
}
// ReadPump is used here solely to monitor the connection, not to actually receive messages.
func (c *Client) readPump() {
defer func() {
c.hub.unregister <- c
}()
for {
_, _, err := c.conn.NextReader()
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
logger.Log().Errorf("[websocket] error: %v", err.Error())
}
break
}
}
}
// WritePump pumps messages from the hub to the websocket connection.
//
// A goroutine running writePump is started for each connection. The
// application ensures that there is at most one writer to a connection by
// executing all writes from this goroutine.
func (c *Client) writePump() {
ticker := time.NewTicker(pingPeriod)
defer func() {
ticker.Stop()
c.hub.unregister <- c
}()
for {
select {
case message, ok := <-c.send:
_ = c.conn.SetWriteDeadline(time.Now().Add(writeWait))
if !ok {
// The hub closed the channel.
_ = c.conn.WriteMessage(websocket.CloseMessage, []byte{})
return
}
w, err := c.conn.NextWriter(websocket.TextMessage)
if err != nil {
return
}
_, _ = w.Write(message)
// Add queued chat messages to the current websocket message.
n := len(c.send)
for i := 0; i < n; i++ {
_, _ = w.Write(newline)
_, _ = w.Write(<-c.send)
}
if err := w.Close(); err != nil {
return
}
case <-ticker.C:
_ = c.conn.SetWriteDeadline(time.Now().Add(writeWait))
_ = c.conn.WriteMessage(websocket.PingMessage, []byte{})
}
}
}
// ServeWs handles websocket requests from the peer.
func ServeWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
if auth.UICredentials != nil {
user, pass, ok := r.BasicAuth()
if !ok {
basicAuthResponse(w)
return
}
if !auth.UICredentials.Match(user, pass) {
basicAuthResponse(w)
return
}
}
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
logger.Log().Errorf("[websocket] %s", err.Error())
return
}
client := &Client{hub: hub, conn: conn, send: make(chan []byte, 256)}
client.hub.register <- client
// Allow collection of memory referenced by the caller by doing all work in new goroutines.
go client.readPump()
go client.writePump()
}
// BasicAuthResponse returns an basic auth response to the browser
func basicAuthResponse(w http.ResponseWriter) {
w.Header().Set("WWW-Authenticate", `Basic realm="Login"`)
w.WriteHeader(http.StatusUnauthorized)
_, _ = w.Write([]byte("Unauthorised.\n"))
}