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  }