github.com/niedbalski/juju@v0.0.0-20190215020005-8ff100488e47/apiserver/websocket/websocket.go (about)

     1  // Copyright 2017 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package websocket
     5  
     6  import (
     7  	"encoding/json"
     8  	"net/http"
     9  	"time"
    10  
    11  	"github.com/gorilla/websocket"
    12  	"github.com/juju/errors"
    13  	"github.com/juju/loggo"
    14  
    15  	"github.com/juju/juju/apiserver/common"
    16  	"github.com/juju/juju/apiserver/params"
    17  )
    18  
    19  var logger = loggo.GetLogger("juju.apiserver.websocket")
    20  
    21  const (
    22  	// PongDelay is how long the server will wait for a pong to be sent
    23  	// before the websocket is considered broken.
    24  	PongDelay = 90 * time.Second
    25  
    26  	// PingPeriod is how often ping messages are sent. This should be shorter
    27  	// than the pongDelay, but not by too much. The difference here allows
    28  	// the remote endpoint 30 seconds to respond to the ping as a ping is sent
    29  	// every 60s, and when a pong is received the read deadline is advanced
    30  	// another 90s.
    31  	PingPeriod = 60 * time.Second
    32  
    33  	// WriteWait is how long the write call can take before it errors out.
    34  	WriteWait = 10 * time.Second
    35  )
    36  
    37  var websocketUpgrader = websocket.Upgrader{
    38  	CheckOrigin: func(r *http.Request) bool { return true },
    39  }
    40  
    41  // Conn wraps a gorilla/websocket.Conn, providing additional Juju-specific
    42  // functionality.
    43  type Conn struct {
    44  	*websocket.Conn
    45  }
    46  
    47  // Serve upgrades an HTTP connection to a websocket, and
    48  // serves the given handler.
    49  func Serve(w http.ResponseWriter, req *http.Request, handler func(ws *Conn)) {
    50  	conn, err := websocketUpgrader.Upgrade(w, req, nil)
    51  	if err != nil {
    52  		logger.Errorf("problem initiating websocket: %v", err)
    53  		return
    54  	}
    55  	handler(&Conn{conn})
    56  }
    57  
    58  // SendInitialErrorV0 writes out the error as a params.ErrorResult serialized
    59  // with JSON with a new line character at the end.
    60  //
    61  // This is a hangover from the initial debug-log streaming endpoint where the
    62  // client read the first line, and then just got a stream of data. We should
    63  // look to version the streaming endpoints to get rid of the trailing newline
    64  // character for message based connections, which is all of them now.
    65  func (conn *Conn) SendInitialErrorV0(err error) error {
    66  	wrapped := &params.ErrorResult{
    67  		Error: common.ServerError(err),
    68  	}
    69  
    70  	body, err := json.Marshal(wrapped)
    71  	if err != nil {
    72  		errors.Annotatef(err, "cannot marshal error %#v", wrapped)
    73  		return err
    74  	}
    75  	body = append(body, '\n')
    76  
    77  	writer, err := conn.NextWriter(websocket.TextMessage)
    78  	if err != nil {
    79  		return errors.Annotate(err, "problem getting writer")
    80  	}
    81  	defer writer.Close()
    82  	_, err = writer.Write(body)
    83  
    84  	if wrapped.Error != nil {
    85  		// Tell the other end we are closing.
    86  		conn.WriteMessage(websocket.CloseMessage, []byte{})
    87  	}
    88  
    89  	return errors.Trace(err)
    90  }