server.go raw
1 package wsserver
2
3 import (
4 "bufio"
5 "crypto/sha1"
6 "encoding/base64"
7 "encoding/binary"
8 "errors"
9 "fmt"
10 "io"
11 "net"
12 "net/url"
13 "strings"
14 "sync"
15 "time"
16 )
17
18 // Frame opcodes.
19 const (
20 opContinuation = 0x0
21 opText = 0x1
22 opBinary = 0x2
23 opClose = 0x8
24 opPing = 0x9
25 opPong = 0xa
26 )
27
28 const acceptGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
29
30 // Request is the parsed upgrade request.
31 type Request struct {
32 Method string
33 Path string
34 Query url.Values
35 Header map[string]string
36 }
37
38 // Handler receives a live WebSocket connection. It owns the connection and
39 // must return when the stream ends.
40 type Handler func(*Conn, *Request)
41
42 // Config starts a server.
43 type Config struct {
44 // Listen is the TCP address, for example "127.0.0.1:3335".
45 Listen string
46 // Path is the required request path. Empty means "/nostrino".
47 Path string
48 // Handler runs one connection.
49 Handler Handler
50 // MaxMessage bounds one message. Zero means 32 MiB.
51 MaxMessage int64
52 // HandshakeTimeout bounds the upgrade. Zero means 15s.
53 HandshakeTimeout time.Duration
54 // Logf receives diagnostics; nil discards them.
55 Logf func(string, ...any)
56 }
57
58 // Server is a running WebSocket listener.
59 type Server struct {
60 cfg Config
61 ln net.Listener
62 done chan struct{}
63
64 mu sync.Mutex
65 conns map[*Conn]struct{}
66 wg sync.WaitGroup
67 }
68
69 // Start listens and serves.
70 func Start(cfg Config) (*Server, error) {
71 if cfg.Listen == "" {
72 return nil, errors.New("wsserver: no listen address")
73 }
74 if cfg.Handler == nil {
75 return nil, errors.New("wsserver: no handler")
76 }
77 if cfg.Path == "" {
78 cfg.Path = "/nostrino"
79 }
80 if cfg.MaxMessage == 0 {
81 cfg.MaxMessage = 32 << 20
82 }
83 if cfg.HandshakeTimeout == 0 {
84 cfg.HandshakeTimeout = 15 * time.Second
85 }
86 if cfg.Logf == nil {
87 cfg.Logf = func(string, ...any) {}
88 }
89 ln, err := net.Listen("tcp", cfg.Listen)
90 if err != nil {
91 return nil, fmt.Errorf("wsserver: listen: %w", err)
92 }
93 s := &Server{
94 cfg: cfg,
95 ln: ln,
96 done: make(chan struct{}),
97 conns: make(map[*Conn]struct{}),
98 }
99 s.wg.Add(1)
100 go s.acceptLoop()
101 return s, nil
102 }
103
104 // Addr is the bound TCP address, including an ephemeral port.
105 func (s *Server) Addr() string { return s.ln.Addr().String() }
106
107 // Close stops the listener and every connection.
108 func (s *Server) Close() error {
109 select {
110 case <-s.done:
111 default:
112 close(s.done)
113 }
114 err := s.ln.Close()
115 s.mu.Lock()
116 conns := make([]*Conn, 0, len(s.conns))
117 for c := range s.conns {
118 conns = append(conns, c)
119 }
120 s.mu.Unlock()
121 for _, c := range conns {
122 c.Close()
123 }
124 s.wg.Wait()
125 return err
126 }
127
128 func (s *Server) acceptLoop() {
129 defer s.wg.Done()
130 for {
131 conn, err := s.ln.Accept()
132 if err != nil {
133 return
134 }
135 s.wg.Add(1)
136 go func() {
137 defer s.wg.Done()
138 s.serve(conn)
139 }()
140 }
141 }
142
143 func (s *Server) serve(raw net.Conn) {
144 defer raw.Close()
145 _ = raw.SetDeadline(time.Now().Add(s.cfg.HandshakeTimeout))
146 br := bufio.NewReaderSize(raw, 16<<10)
147 req, key, err := readHandshake(br)
148 if err != nil {
149 s.cfg.Logf("wsserver: handshake: %v", err)
150 return
151 }
152 if req.Path != s.cfg.Path {
153 fmt.Fprint(raw, "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
154 return
155 }
156 accept := acceptKey(key)
157 fmt.Fprintf(raw, "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: %s\r\n\r\n", accept)
158 _ = raw.SetDeadline(time.Time{})
159 c := &Conn{conn: raw, br: br, max: s.cfg.MaxMessage}
160 s.mu.Lock()
161 s.conns[c] = struct{}{}
162 s.mu.Unlock()
163 s.cfg.Handler(c, req)
164 c.Close()
165 s.mu.Lock()
166 delete(s.conns, c)
167 s.mu.Unlock()
168 }
169
170 func readHandshake(br *bufio.Reader) (*Request, string, error) {
171 line, err := br.ReadString('\n')
172 if err != nil {
173 return nil, "", err
174 }
175 line = strings.TrimRight(line, "\r\n")
176 fields := strings.Fields(line)
177 if len(fields) != 3 || !strings.EqualFold(fields[0], "GET") {
178 return nil, "", errors.New("not a GET")
179 }
180 target := fields[1]
181 path, query := target, ""
182 if i := strings.IndexByte(target, '?'); i >= 0 {
183 path, query = target[:i], target[i+1:]
184 }
185 q, err := url.ParseQuery(query)
186 if err != nil {
187 return nil, "", fmt.Errorf("bad query: %w", err)
188 }
189 header := make(map[string]string, 16)
190 for i := 0; i < 64; i++ {
191 line, err := br.ReadString('\n')
192 if err != nil {
193 return nil, "", err
194 }
195 line = strings.TrimRight(line, "\r\n")
196 if line == "" {
197 break
198 }
199 name, value, ok := strings.Cut(line, ":")
200 if !ok {
201 return nil, "", errors.New("malformed header")
202 }
203 header[strings.ToLower(strings.TrimSpace(name))] = strings.TrimSpace(value)
204 }
205 if !headerContains(header["upgrade"], "websocket") || !headerContains(header["connection"], "upgrade") {
206 return nil, "", errors.New("not an upgrade")
207 }
208 if header["sec-websocket-version"] != "13" {
209 return nil, "", errors.New("unsupported websocket version")
210 }
211 key := header["sec-websocket-key"]
212 if key == "" {
213 return nil, "", errors.New("missing websocket key")
214 }
215 return &Request{Method: "GET", Path: path, Query: q, Header: header}, key, nil
216 }
217
218 func headerContains(value, token string) bool {
219 for _, part := range strings.Split(value, ",") {
220 if strings.EqualFold(strings.TrimSpace(part), token) {
221 return true
222 }
223 }
224 return false
225 }
226
227 func acceptKey(key string) string {
228 h := sha1.Sum([]byte(key + acceptGUID))
229 return base64.StdEncoding.EncodeToString(h[:])
230 }
231
232 // Conn is one server-side WebSocket and implements net.Conn with message
233 // payloads as the byte stream.
234 type Conn struct {
235 conn net.Conn
236 br *bufio.Reader
237 max int64
238
239 wmu sync.Mutex
240 dead bool
241
242 buf []byte
243 closeOnce sync.Once
244 }
245
246 // Read returns bytes from the next data message.
247 func (c *Conn) Read(p []byte) (int, error) {
248 if len(p) == 0 {
249 return 0, nil
250 }
251 for len(c.buf) == 0 {
252 msg, err := c.readMessage()
253 if err != nil {
254 return 0, err
255 }
256 c.buf = msg
257 }
258 n := copy(p, c.buf)
259 c.buf = c.buf[n:]
260 return n, nil
261 }
262
263 // Write sends p as one binary message.
264 func (c *Conn) Write(p []byte) (int, error) {
265 if err := c.writeFrame(opBinary, p); err != nil {
266 return 0, err
267 }
268 return len(p), nil
269 }
270
271 func (c *Conn) readMessage() ([]byte, error) {
272 var message []byte
273 started := false
274 for {
275 fin, op, payload, err := c.readFrame()
276 if err != nil {
277 return nil, err
278 }
279 switch op {
280 case opPing:
281 if err := c.writeFrame(opPong, payload); err != nil {
282 return nil, err
283 }
284 continue
285 case opPong:
286 continue
287 case opClose:
288 _ = c.writeFrame(opClose, payload)
289 c.closeConn()
290 return nil, io.EOF
291 case opText, opBinary:
292 if started {
293 return nil, errors.New("wsserver: data frame inside a message")
294 }
295 started = true
296 message = append(message[:0], payload...)
297 case opContinuation:
298 if !started {
299 return nil, errors.New("wsserver: continuation without a message")
300 }
301 message = append(message, payload...)
302 default:
303 return nil, fmt.Errorf("wsserver: unknown opcode %#x", op)
304 }
305 if int64(len(message)) > c.max {
306 return nil, errors.New("wsserver: message too large")
307 }
308 if fin {
309 if !started {
310 return nil, errors.New("wsserver: empty message")
311 }
312 return message, nil
313 }
314 }
315 }
316
317 func (c *Conn) readFrame() (fin bool, op byte, payload []byte, err error) {
318 var hdr [2]byte
319 if _, err = io.ReadFull(c.br, hdr[:]); err != nil {
320 return false, 0, nil, err
321 }
322 fin = hdr[0]&0x80 != 0
323 op = hdr[0] & 0x0f
324 masked := hdr[1]&0x80 != 0
325 length := uint64(hdr[1] & 0x7f)
326 switch length {
327 case 126:
328 var ext [2]byte
329 if _, err = io.ReadFull(c.br, ext[:]); err != nil {
330 return false, 0, nil, err
331 }
332 length = uint64(binary.BigEndian.Uint16(ext[:]))
333 case 127:
334 var ext [8]byte
335 if _, err = io.ReadFull(c.br, ext[:]); err != nil {
336 return false, 0, nil, err
337 }
338 length = binary.BigEndian.Uint64(ext[:])
339 if length>>63 != 0 {
340 return false, 0, nil, errors.New("wsserver: negative length")
341 }
342 }
343 if length > uint64(c.max) {
344 return false, 0, nil, errors.New("wsserver: message too large")
345 }
346 if op >= opClose && (!fin || length > 125) {
347 return false, 0, nil, errors.New("wsserver: malformed control frame")
348 }
349 var mask [4]byte
350 if masked {
351 if _, err = io.ReadFull(c.br, mask[:]); err != nil {
352 return false, 0, nil, err
353 }
354 }
355 payload = make([]byte, length)
356 if _, err = io.ReadFull(c.br, payload); err != nil {
357 return false, 0, nil, err
358 }
359 if masked {
360 for i := range payload {
361 payload[i] ^= mask[i&3]
362 }
363 }
364 return fin, op, payload, nil
365 }
366
367 func (c *Conn) writeFrame(op byte, payload []byte) error {
368 c.wmu.Lock()
369 defer c.wmu.Unlock()
370 if c.dead {
371 return io.ErrClosedPipe
372 }
373 var hdr [10]byte
374 hdr[0] = 0x80 | op
375 n := len(payload)
376 var hlen int
377 switch {
378 case n < 126:
379 hdr[1] = byte(n)
380 hlen = 2
381 case n <= 0xffff:
382 hdr[1] = 126
383 binary.BigEndian.PutUint16(hdr[2:4], uint16(n))
384 hlen = 4
385 default:
386 hdr[1] = 127
387 binary.BigEndian.PutUint64(hdr[2:10], uint64(n))
388 hlen = 10
389 }
390 buf := make([]byte, hlen+n)
391 copy(buf, hdr[:hlen])
392 copy(buf[hlen:], payload)
393 _, err := c.conn.Write(buf)
394 return err
395 }
396
397 // Close sends a close frame and closes the socket.
398 func (c *Conn) Close() error {
399 var err error
400 c.closeOnce.Do(func() {
401 _ = c.writeFrame(opClose, nil)
402 err = c.closeConn()
403 })
404 return err
405 }
406
407 func (c *Conn) closeConn() error {
408 c.wmu.Lock()
409 c.dead = true
410 c.wmu.Unlock()
411 return c.conn.Close()
412 }
413
414 // LocalAddr implements net.Conn.
415 func (c *Conn) LocalAddr() net.Addr { return c.conn.LocalAddr() }
416
417 // RemoteAddr implements net.Conn.
418 func (c *Conn) RemoteAddr() net.Addr { return c.conn.RemoteAddr() }
419
420 // SetDeadline implements net.Conn.
421 func (c *Conn) SetDeadline(t time.Time) error { return c.conn.SetDeadline(t) }
422
423 // SetReadDeadline implements net.Conn.
424 func (c *Conn) SetReadDeadline(t time.Time) error { return c.conn.SetReadDeadline(t) }
425
426 // SetWriteDeadline implements net.Conn.
427 func (c *Conn) SetWriteDeadline(t time.Time) error { return c.conn.SetWriteDeadline(t) }
428