/
writer.go
103 lines (93 loc) · 2.26 KB
/
writer.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
package websocket
import (
"bufio"
"crypto/rand"
"encoding/binary"
"io"
"math"
)
type writer struct {
wr *bufio.Writer
opcode uint8
err error
client bool
}
func (w *writer) Write(b []byte) (int, error) {
if w.err != nil {
return 0, w.err
}
wr := w.wr
fin := leftBit | w.opcode
// Check if message needs to be fragmented.
if bsize, diff := w.byteSize(b), wr.Available(); bsize > diff {
diff -= bsize - len(b) // resolve payload length
fin &= w.opcode // set FIN bit to zero
next := &writer{wr: w.wr, opcode: opcodeContinuation} // prepare next frame to be sent
defer next.Write(b[diff:]) // schedule next write
b = b[:diff] // resolve current payload
}
if w.err = wr.WriteByte(fin); w.err != nil {
return 0, w.err
}
size := len(b)
var maskedBit uint8
if w.client {
maskedBit = leftBit
}
switch {
case size <= 125:
if w.err = wr.WriteByte(uint8(size) | maskedBit); w.err != nil {
return 0, w.err
}
case size <= math.MaxUint16:
if w.err = wr.WriteByte(126 | maskedBit); w.err != nil {
return 0, w.err
}
if w.err = binary.Write(wr, binary.BigEndian, uint16(size)); w.err != nil {
return 0, w.err
}
default:
if w.err = wr.WriteByte(127 | maskedBit); w.err != nil {
return 0, w.err
}
if w.err = binary.Write(wr, binary.BigEndian, uint64(size)); w.err != nil {
return 0, w.err
}
}
// Mask the payload.
if w.client {
m := make(mask, 4)
if _, w.err = io.ReadFull(rand.Reader, m); w.err != nil {
return 0, w.err
}
if _, w.err = wr.Write(m); w.err != nil {
return 0, w.err
}
m.transform(b)
}
// Write message and flush.
if _, w.err = wr.Write(b); w.err != nil {
return 0, w.err
}
n := wr.Buffered()
if w.err = wr.Flush(); w.err != nil {
return 0, w.err
}
return n, nil
}
func (w *writer) byteSize(b []byte) (size int) {
size++ // FIN
length := len(b)
switch {
case length <= 125:
size++ // indicator is the current length
case length <= math.MaxUint16:
size += 3 // indicator + 2 bytes for length value
default:
size += 9 // indicator + 8 bytes for length value
}
if w.client {
size += 4 // 4 bytes for masking-key
}
return size + length
}