github.com/timstclair/heapster@v0.20.0-alpha1/Godeps/_workspace/src/k8s.io/kubernetes/pkg/util/wsstream/stream.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 "io" 22 "net/http" 23 "time" 24 25 "golang.org/x/net/websocket" 26 "k8s.io/kubernetes/pkg/util" 27 ) 28 29 // The WebSocket subprotocol "binary.k8s.io" will only send messages to the 30 // client and ignore messages sent to the server. The received messages are 31 // the exact bytes written to the stream. Zero byte messages are possible. 32 const binaryWebSocketProtocol = "binary.k8s.io" 33 34 // The WebSocket subprotocol "base64.binary.k8s.io" will only send messages to the 35 // client and ignore messages sent to the server. The received messages are 36 // a base64 version of the bytes written to the stream. Zero byte messages are 37 // possible. 38 const base64BinaryWebSocketProtocol = "base64.binary.k8s.io" 39 40 // Reader supports returning an arbitrary byte stream over a websocket channel. 41 // Supports the "binary.k8s.io" and "base64.binary.k8s.io" subprotocols. 42 type Reader struct { 43 err chan error 44 r io.Reader 45 ping bool 46 timeout time.Duration 47 } 48 49 // NewReader creates a WebSocket pipe that will copy the contents of r to a provided 50 // WebSocket connection. If ping is true, a zero length message will be sent to the client 51 // before the stream begins reading. 52 func NewReader(r io.Reader, ping bool) *Reader { 53 return &Reader{ 54 r: r, 55 err: make(chan error), 56 ping: ping, 57 } 58 } 59 60 // SetIdleTimeout sets the interval for both reads and writes before timeout. If not specified, 61 // there is no timeout on the reader. 62 func (r *Reader) SetIdleTimeout(duration time.Duration) { 63 r.timeout = duration 64 } 65 66 func (r *Reader) handshake(config *websocket.Config, req *http.Request) error { 67 return handshake(config, req, []string{binaryWebSocketProtocol, base64BinaryWebSocketProtocol}) 68 } 69 70 // Copy the reader to the response. The created WebSocket is closed after this 71 // method completes. 72 func (r *Reader) Copy(w http.ResponseWriter, req *http.Request) error { 73 go func() { 74 defer util.HandleCrash() 75 websocket.Server{Handshake: r.handshake, Handler: r.handle}.ServeHTTP(w, req) 76 }() 77 return <-r.err 78 } 79 80 // handle implements a WebSocket handler. 81 func (r *Reader) handle(ws *websocket.Conn) { 82 encode := len(ws.Config().Protocol) > 0 && ws.Config().Protocol[0] == base64BinaryWebSocketProtocol 83 defer close(r.err) 84 defer ws.Close() 85 go ignoreReceives(ws, r.timeout) 86 r.err <- messageCopy(ws, r.r, encode, r.ping, r.timeout) 87 } 88 89 func resetTimeout(ws *websocket.Conn, timeout time.Duration) { 90 if timeout > 0 { 91 ws.SetDeadline(time.Now().Add(timeout)) 92 } 93 } 94 95 func messageCopy(ws *websocket.Conn, r io.Reader, base64Encode, ping bool, timeout time.Duration) error { 96 buf := make([]byte, 2048) 97 if ping { 98 resetTimeout(ws, timeout) 99 if err := websocket.Message.Send(ws, []byte{}); err != nil { 100 return err 101 } 102 } 103 for { 104 resetTimeout(ws, timeout) 105 n, err := r.Read(buf) 106 if err != nil { 107 if err == io.EOF { 108 return nil 109 } 110 return err 111 } 112 if n > 0 { 113 if base64Encode { 114 if err := websocket.Message.Send(ws, base64.StdEncoding.EncodeToString(buf[:n])); err != nil { 115 return err 116 } 117 } else { 118 if err := websocket.Message.Send(ws, buf[:n]); err != nil { 119 return err 120 } 121 } 122 } 123 } 124 }