1 package transport
2
3 import (
4 "crypto/rand"
5 "crypto/sha1"
6 "encoding/base64"
7 "encoding/hex"
8 "syscall"
9 "git.smesh.lol/morly/pkg/metrics"
10 )
11
12 const wsMagic = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
13
14 // parseWSFrame extracts one WS frame from data.
15 // Returns (op, payload, consumed). consumed=0 means incomplete.
16 func parseWSFrame(data []byte) (op byte, payload []byte, n int32) {
17 if len(data) < 2 {
18 return 0, nil, 0
19 }
20 op := data[0] & 0x0f
21 masked := data[1]&0x80 != 0
22 length := int32(data[1] & 0x7f)
23 pos := 2
24
25 if length == 126 {
26 if len(data) < 4 {
27 return 0, nil, 0
28 }
29 length = int32(data[2])<<8 | int32(data[3])
30 pos = 4
31 } else if length == 127 {
32 if len(data) < 10 {
33 return 0, nil, 0
34 }
35 length = int32(data[6])<<24 | int32(data[7])<<16 | int32(data[8])<<8 | int32(data[9])
36 pos = 10
37 }
38
39 var mask [4]byte
40 if masked {
41 if len(data) < pos+4 {
42 return 0, nil, 0
43 }
44 copy(mask[:], data[pos:pos+4])
45 pos += 4
46 }
47 if len(data) < pos+length {
48 return 0, nil, 0
49 }
50
51 payload := makeCopy(data[pos : pos+length])
52 if masked {
53 for i := range payload {
54 payload[i] ^= mask[i%4]
55 }
56 }
57 return op, payload, pos + length
58 }
59
60 func buildWSFrame(op byte, payload []byte) (buf []byte) {
61 plen := len(payload)
62 var hdr [10]byte
63 hdr[0] = 0x80 | op
64 n := 2
65 if plen < 126 {
66 hdr[1] = byte(plen)
67 } else if plen < 65536 {
68 hdr[1] = 126
69 hdr[2] = byte(plen >> 8)
70 hdr[3] = byte(plen)
71 n = 4
72 } else {
73 hdr[1] = 127
74 hdr[6] = byte(plen >> 24)
75 hdr[7] = byte(plen >> 16)
76 hdr[8] = byte(plen >> 8)
77 hdr[9] = byte(plen)
78 n = 10
79 }
80 buf := []byte{:0:n + plen}
81 buf = buf | hdr[:n]
82 buf = buf | payload
83 return buf
84 }
85
86 func writeWSFrame(fd int32, op byte, payload []byte) {
87 writeAll(fd, buildWSFrame(op, payload))
88 }
89
90 func writeWSFrameErr(fd int32, op byte, payload []byte) (eerr error) {
91 return writeAll(fd, buildWSFrame(op, payload))
92 }
93
94 func writeWSClose(fd int32) {
95 writeWSFrame(fd, opClose, []byte{0x03, 0xe8})
96 }
97
98 func writeAll(fd int32, data []byte) (eerr error) {
99 for len(data) > 0 {
100 n, err := syscall.Write(fd, data)
101 if n > 0 {
102 data = data[n:]
103 }
104 if errAgain(err) {
105 return syscall.EAGAIN
106 }
107 if errIntr(err) {
108 continue
109 }
110 if err != nil {
111 return err
112 }
113 }
114 return nil
115 }
116
117 func computeAccept(key string) (s string) {
118 h := sha1.New()
119 h.Write([]byte(key))
120 h.Write([]byte(wsMagic))
121 return base64.StdEncoding().EncodeToString(h.Sum(nil))
122 }
123
124 func (s *Server) upgradeWS(c *tconn, req *httpReq) {
125 key := req.headers["sec-websocket-key"]
126 if key == "" {
127 writeHTTPResponse(c.fd, 400, nil, []byte("missing Sec-WebSocket-Key"))
128 s.closeConn(c)
129 return
130 }
131 wl, allow := s.handler.OnWSUpgrade(c.fd, c.remoteIP, s.ipConns[c.remoteIP])
132 if !allow {
133 writeHTTPResponse(c.fd, 429, nil, []byte("too many websockets"))
134 s.closeConn(c)
135 return
136 }
137 s.ipConns[c.remoteIP]++
138 c.whitelisted = wl
139 accept := computeAccept(key)
140 resp := "HTTP/1.1 101 Switching Protocols\r\n" |
141 "Upgrade: websocket\r\n" |
142 "Connection: Upgrade\r\n" |
143 "Sec-WebSocket-Accept: " | accept | "\r\n" |
144 "\r\n"
145 writeAll(c.fd, []byte(resp))
146 c.phase = phaseWS
147 s.handler.OnWSConnected(c.fd)
148 if c.wpos > 0 {
149 s.processWS(c)
150 }
151 }
152
153 // processWS drains every complete frame already in the connection's buffer.
154 //
155 // Handling one frame per call is not enough: readConn reads whatever the socket
156 // has, and a client is free to put several frames in one segment - a browser
157 // sending four REQs back to back does exactly that. The frames after the first
158 // are then already out of the socket, so no further EPOLLIN ever arrives for
159 // them and they sat in c.buf unprocessed. A pipelined client got one answer and
160 // every later subscription on that connection was silently dead, which is the
161 // whole of the browser-feed failure: four requests, one EOSE, no stored event.
162 func (s *Server) processWS(c *tconn) {
163 var op byte
164 var payload []byte
165 var consumed int32
166 var frameStart int64
167 for {
168 op, payload, consumed = parseWSFrame(c.buf[:c.wpos])
169 if consumed == 0 {
170 return
171 }
172 copy(c.buf, c.buf[consumed:c.wpos])
173 c.wpos -= consumed
174
175 switch op {
176 case opClose:
177 writeWSClose(c.fd)
178 s.closeConn(c)
179 return
180 case opPing:
181 writeWSFrame(c.fd, opPong, payload)
182 case opText, opBin:
183 frameStart = metrics.Now()
184 s.handler.OnWSMessage(c.fd, payload)
185 metrics.WSFrameNs.Observe(metrics.Since(frameStart))
186 }
187 // The handler may have closed the connection.
188 if s.conns[c.fd] == nil {
189 return
190 }
191 }
192 }
193
194 // Challenge generates a 32-byte random hex challenge string for NIP-42 auth.
195 func Challenge() (buf []byte) {
196 var cb [32]byte
197 rand.Read(cb[:])
198 out := []byte{:64}
199 hex.Encode(out, cb[:])
200 return out
201 }
202