forked from gravitational/teleport
-
Notifications
You must be signed in to change notification settings - Fork 0
/
websocketwriter.go
106 lines (93 loc) · 2.51 KB
/
websocketwriter.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
package utils
import (
"io"
"github.com/gravitational/trace"
"golang.org/x/net/websocket"
"golang.org/x/text/encoding"
"golang.org/x/text/encoding/unicode"
log "github.com/Sirupsen/logrus"
)
// WebSockWrapper wraps the raw websocket and converts Write() calls
// to proper websocket.Send() working in binary or text mode. If text
// mode is selected, it converts the data passed to Write() into UTF8 bytes
//
// We need this to make sure that the entire buffer in io.Writer.Write(buffer)
// is delivered as a single chunk to the web browser, instead of being split
// into multiple frames. This wrapper basically substitues every Write() with
// Send() and every Read() with Receive()
type WebSockWrapper struct {
io.ReadWriteCloser
ws *websocket.Conn
mode WebSocketMode
encoder *encoding.Encoder
decoder *encoding.Decoder
}
// WebSocketMode allows to create WebSocket wrappers working in text
// or binary mode
type WebSocketMode int
const (
WebSocketBinaryMode = iota
WebSocketTextMode
)
func NewWebSockWrapper(ws *websocket.Conn, m WebSocketMode) *WebSockWrapper {
if ws == nil {
return nil
}
return &WebSockWrapper{
ws: ws,
mode: m,
encoder: unicode.UTF8.NewEncoder(),
decoder: unicode.UTF8.NewDecoder(),
}
}
// Write implements io.WriteCloser for WebSockWriter (that's the reason we're
// wrapping the websocket)
//
// It replaces raw Write() with "Message.Send()"
func (w *WebSockWrapper) Write(data []byte) (n int, err error) {
n = len(data)
if w.mode == WebSocketBinaryMode {
// binary send:
err = websocket.Message.Send(w.ws, data)
// text send:
} else {
var utf8 string
utf8, err = w.encoder.String(string(data))
err = websocket.Message.Send(w.ws, utf8)
}
if err != nil {
log.Error(err)
n = 0
}
return n, err
}
// Read does the opposite of write: it replaces websocket's raw "Read" with
//
// It replaces raw Read() with "Message.Receive()"
func (w *WebSockWrapper) Read(out []byte) (n int, err error) {
var data []byte
if w.mode == WebSocketBinaryMode {
err = websocket.Message.Receive(w.ws, &data)
} else {
var utf8 string
err = websocket.Message.Receive(w.ws, &utf8)
switch err {
case nil:
data, err = w.decoder.Bytes([]byte(utf8))
case io.EOF:
return 0, io.EOF
default:
log.Error(err)
}
}
if err != nil {
return 0, trace.Wrap(err)
}
if len(out) < len(data) {
log.Warningf("websocket failed to receive everything: %d vs %d", len(out), len(data))
}
return copy(out, data), nil
}
func (w *WebSockWrapper) Close() error {
return w.ws.Close()
}