github.com/timstclair/heapster@v0.20.0-alpha1/Godeps/_workspace/src/k8s.io/kubernetes/pkg/util/wsstream/conn.go (about) 1 /* 2 Copyright 2015 The Kubernetes Authors All rights reserved. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15 */ 16 17 package wsstream 18 19 import ( 20 "encoding/base64" 21 "fmt" 22 "io" 23 "net/http" 24 "regexp" 25 "strings" 26 "time" 27 28 "github.com/golang/glog" 29 "golang.org/x/net/websocket" 30 "k8s.io/kubernetes/pkg/util" 31 ) 32 33 // The Websocket subprotocol "channel.k8s.io" prepends each binary message with a byte indicating 34 // the channel number (zero indexed) the message was sent on. Messages in both directions should 35 // prefix their messages with this channel byte. When used for remote execution, the channel numbers 36 // are by convention defined to match the POSIX file-descriptors assigned to STDIN, STDOUT, and STDERR 37 // (0, 1, and 2). No other conversion is performed on the raw subprotocol - writes are sent as they 38 // are received by the server. 39 // 40 // Example client session: 41 // 42 // CONNECT http://server.com with subprotocol "channel.k8s.io" 43 // WRITE []byte{0, 102, 111, 111, 10} # send "foo\n" on channel 0 (STDIN) 44 // READ []byte{1, 10} # receive "\n" on channel 1 (STDOUT) 45 // CLOSE 46 // 47 const channelWebSocketProtocol = "channel.k8s.io" 48 49 // The Websocket subprotocol "base64.channel.k8s.io" base64 encodes each message with a character 50 // indicating the channel number (zero indexed) the message was sent on. Messages in both directions 51 // should prefix their messages with this channel char. When used for remote execution, the channel 52 // numbers are by convention defined to match the POSIX file-descriptors assigned to STDIN, STDOUT, 53 // and STDERR ('0', '1', and '2'). The data received on the server is base64 decoded (and must be 54 // be valid) and data written by the server to the client is base64 encoded. 55 // 56 // Example client session: 57 // 58 // CONNECT http://server.com with subprotocol "base64.channel.k8s.io" 59 // WRITE []byte{48, 90, 109, 57, 118, 67, 103, 111, 61} # send "foo\n" (base64: "Zm9vCgo=") on channel '0' (STDIN) 60 // READ []byte{49, 67, 103, 61, 61} # receive "\n" (base64: "Cg==") on channel '1' (STDOUT) 61 // CLOSE 62 // 63 const base64ChannelWebSocketProtocol = "base64.channel.k8s.io" 64 65 type codecType int 66 67 const ( 68 rawCodec codecType = iota 69 base64Codec 70 ) 71 72 type ChannelType int 73 74 const ( 75 IgnoreChannel ChannelType = iota 76 ReadChannel 77 WriteChannel 78 ReadWriteChannel 79 ) 80 81 var ( 82 // connectionUpgradeRegex matches any Connection header value that includes upgrade 83 connectionUpgradeRegex = regexp.MustCompile("(^|.*,\\s*)upgrade($|\\s*,)") 84 ) 85 86 // IsWebSocketRequest returns true if the incoming request contains connection upgrade headers 87 // for WebSockets. 88 func IsWebSocketRequest(req *http.Request) bool { 89 return connectionUpgradeRegex.MatchString(strings.ToLower(req.Header.Get("Connection"))) && strings.ToLower(req.Header.Get("Upgrade")) == "websocket" 90 } 91 92 // ignoreReceives reads from a WebSocket until it is closed, then returns. If timeout is set, the 93 // read and write deadlines are pushed every time a new message is received. 94 func ignoreReceives(ws *websocket.Conn, timeout time.Duration) { 95 defer util.HandleCrash() 96 var data []byte 97 for { 98 resetTimeout(ws, timeout) 99 if err := websocket.Message.Receive(ws, &data); err != nil { 100 return 101 } 102 } 103 } 104 105 // handshake ensures the provided user protocol matches one of the allowed protocols. It returns 106 // no error if no protocol is specified. 107 func handshake(config *websocket.Config, req *http.Request, allowed []string) error { 108 protocols := config.Protocol 109 if len(protocols) == 0 { 110 return nil 111 } 112 for _, protocol := range protocols { 113 for _, allow := range allowed { 114 if allow == protocol { 115 config.Protocol = []string{protocol} 116 return nil 117 } 118 } 119 } 120 return fmt.Errorf("requested protocol(s) are not supported: %v; supports %v", config.Protocol, allowed) 121 } 122 123 // Conn supports sending multiple binary channels over a websocket connection. 124 // Supports only the "channel.k8s.io" subprotocol. 125 type Conn struct { 126 channels []*websocketChannel 127 codec codecType 128 ready chan struct{} 129 ws *websocket.Conn 130 timeout time.Duration 131 } 132 133 // NewConn creates a WebSocket connection that supports a set of channels. Channels begin each 134 // web socket message with a single byte indicating the channel number (0-N). 255 is reserved for 135 // future use. The channel types for each channel are passed as an array, supporting the different 136 // duplex modes. Read and Write refer to whether the channel can be used as a Reader or Writer. 137 func NewConn(channels ...ChannelType) *Conn { 138 conn := &Conn{ 139 ready: make(chan struct{}), 140 channels: make([]*websocketChannel, len(channels)), 141 } 142 for i := range conn.channels { 143 switch channels[i] { 144 case ReadChannel: 145 conn.channels[i] = newWebsocketChannel(conn, byte(i), true, false) 146 case WriteChannel: 147 conn.channels[i] = newWebsocketChannel(conn, byte(i), false, true) 148 case ReadWriteChannel: 149 conn.channels[i] = newWebsocketChannel(conn, byte(i), true, true) 150 case IgnoreChannel: 151 conn.channels[i] = newWebsocketChannel(conn, byte(i), false, false) 152 } 153 } 154 return conn 155 } 156 157 // SetIdleTimeout sets the interval for both reads and writes before timeout. If not specified, 158 // there is no timeout on the connection. 159 func (conn *Conn) SetIdleTimeout(duration time.Duration) { 160 conn.timeout = duration 161 } 162 163 // Open the connection and create channels for reading and writing. 164 func (conn *Conn) Open(w http.ResponseWriter, req *http.Request) ([]io.ReadWriteCloser, error) { 165 go func() { 166 defer util.HandleCrash() 167 defer conn.Close() 168 websocket.Server{Handshake: conn.handshake, Handler: conn.handle}.ServeHTTP(w, req) 169 }() 170 <-conn.ready 171 rwc := make([]io.ReadWriteCloser, len(conn.channels)) 172 for i := range conn.channels { 173 rwc[i] = conn.channels[i] 174 } 175 return rwc, nil 176 } 177 178 func (conn *Conn) initialize(ws *websocket.Conn) { 179 protocols := ws.Config().Protocol 180 switch { 181 case len(protocols) == 0, protocols[0] == channelWebSocketProtocol: 182 conn.codec = rawCodec 183 case protocols[0] == base64ChannelWebSocketProtocol: 184 conn.codec = base64Codec 185 } 186 conn.ws = ws 187 close(conn.ready) 188 } 189 190 func (conn *Conn) handshake(config *websocket.Config, req *http.Request) error { 191 return handshake(config, req, []string{channelWebSocketProtocol, base64ChannelWebSocketProtocol}) 192 } 193 194 func (conn *Conn) resetTimeout() { 195 if conn.timeout > 0 { 196 conn.ws.SetDeadline(time.Now().Add(conn.timeout)) 197 } 198 } 199 200 // Close is only valid after Open has been called 201 func (conn *Conn) Close() error { 202 <-conn.ready 203 for _, s := range conn.channels { 204 s.Close() 205 } 206 conn.ws.Close() 207 return nil 208 } 209 210 // handle implements a websocket handler. 211 func (conn *Conn) handle(ws *websocket.Conn) { 212 defer conn.Close() 213 conn.initialize(ws) 214 215 for { 216 conn.resetTimeout() 217 var data []byte 218 if err := websocket.Message.Receive(ws, &data); err != nil { 219 if err != io.EOF { 220 glog.Errorf("Error on socket receive: %v", err) 221 } 222 break 223 } 224 if len(data) == 0 { 225 continue 226 } 227 channel := data[0] 228 if conn.codec == base64Codec { 229 channel = channel - '0' 230 } 231 data = data[1:] 232 if int(channel) >= len(conn.channels) { 233 glog.V(6).Infof("Frame is targeted for a reader %d that is not valid, possible protocol error", channel) 234 continue 235 } 236 if _, err := conn.channels[channel].DataFromSocket(data); err != nil { 237 glog.Errorf("Unable to write frame to %d: %v\n%s", channel, err, string(data)) 238 continue 239 } 240 } 241 } 242 243 // write multiplexes the specified channel onto the websocket 244 func (conn *Conn) write(num byte, data []byte) (int, error) { 245 conn.resetTimeout() 246 switch conn.codec { 247 case rawCodec: 248 frame := make([]byte, len(data)+1) 249 frame[0] = num 250 copy(frame[1:], data) 251 if err := websocket.Message.Send(conn.ws, frame); err != nil { 252 return 0, err 253 } 254 case base64Codec: 255 frame := string('0'+num) + base64.StdEncoding.EncodeToString(data) 256 if err := websocket.Message.Send(conn.ws, frame); err != nil { 257 return 0, err 258 } 259 } 260 return len(data), nil 261 } 262 263 // websocketChannel represents a channel in a connection 264 type websocketChannel struct { 265 conn *Conn 266 num byte 267 r io.Reader 268 w io.WriteCloser 269 270 read, write bool 271 } 272 273 // newWebsocketChannel creates a pipe for writing to a websocket. Do not write to this pipe 274 // prior to the connection being opened. It may be no, half, or full duplex depending on 275 // read and write. 276 func newWebsocketChannel(conn *Conn, num byte, read, write bool) *websocketChannel { 277 r, w := io.Pipe() 278 return &websocketChannel{conn, num, r, w, read, write} 279 } 280 281 func (p *websocketChannel) Write(data []byte) (int, error) { 282 if !p.write { 283 return len(data), nil 284 } 285 return p.conn.write(p.num, data) 286 } 287 288 // DataFromSocket is invoked by the connection receiver to move data from the connection 289 // into a specific channel. 290 func (p *websocketChannel) DataFromSocket(data []byte) (int, error) { 291 if !p.read { 292 return len(data), nil 293 } 294 295 switch p.conn.codec { 296 case rawCodec: 297 return p.w.Write(data) 298 case base64Codec: 299 dst := make([]byte, len(data)) 300 n, err := base64.StdEncoding.Decode(dst, data) 301 if err != nil { 302 return 0, err 303 } 304 return p.w.Write(dst[:n]) 305 } 306 return 0, nil 307 } 308 309 func (p *websocketChannel) Read(data []byte) (int, error) { 310 if !p.read { 311 return 0, io.EOF 312 } 313 return p.r.Read(data) 314 } 315 316 func (p *websocketChannel) Close() error { 317 return p.w.Close() 318 }