Review Feedback

app/api_topology.go

@@ -5,7 +5,6 @@ import (
 	"time"
 
 	log "github.com/Sirupsen/logrus"
-	"github.com/gorilla/websocket"
 	"golang.org/x/net/context"
 
 	"github.com/weaveworks/scope/common/xfer"
@@ -66,10 +65,6 @@ func handleNode(nodeID string) func(context.Context, Reporter, render.Renderer,
 	}
 }
 
-var upgrader = websocket.Upgrader{
-	CheckOrigin: func(r *http.Request) bool { return true },
-}
-
 func handleWebsocket(
 	ctx context.Context,
 	w http.ResponseWriter,
@@ -78,12 +73,11 @@ func handleWebsocket(
 	renderer render.Renderer,
 	loop time.Duration,
 ) {
-	wsConn, err := upgrader.Upgrade(w, r, nil)
+	conn, err := xfer.Upgrade(w, r, nil)
 	if err != nil {
 		// log.Info("Upgrade:", err)
 		return
 	}
-	conn := xfer.Ping(wsConn)
 	defer conn.Close()
 
 	quit := make(chan struct{})

app/controls.go

@@ -57,7 +57,7 @@ func handleProbeWS(cr ControlRouter) CtxHandlerFunc {
 			return
 		}
 
-		conn, err := upgrader.Upgrade(w, r, nil)
+		conn, err := xfer.Upgrade(w, r, nil)
 		if err != nil {
 			log.Printf("Error upgrading control websocket: %v", err)
 			return

app/pipes.go

@@ -35,12 +35,11 @@ func handlePipeWs(pr PipeRouter, end End) CtxHandlerFunc {
 		}
 		defer pr.Release(ctx, id, end)
 
-		wsConn, err := upgrader.Upgrade(w, r, nil)
+		conn, err := xfer.Upgrade(w, r, nil)
 		if err != nil {
 			log.Errorf("Error upgrading pipe %s (%d) websocket: %v", id, end, err)
 			return
 		}
-		conn := xfer.Ping(wsConn)
 		defer conn.Close()
 
 		log.Infof("Pipe success %s (%d)", id, end)

common/xfer/websocket.go

@@ -2,9 +2,11 @@ package xfer
 
 import (
 	"io"
+	"net/http"
 	"sync"
 	"time"
 
+	log "github.com/Sirupsen/logrus"
 	"github.com/gorilla/websocket"
 	"github.com/ugorji/go/codec"
 
@@ -44,12 +46,42 @@ type pingingWebsocket struct {
 	conn      *websocket.Conn
 }
 
+var upgrader = websocket.Upgrader{
+	CheckOrigin: func(r *http.Request) bool { return true },
+}
+
+// Upgrade upgrades the HTTP server connection to the WebSocket protocol.
+func Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header) (Websocket, error) {
+	wsConn, err := upgrader.Upgrade(w, r, responseHeader)
+	if err != nil {
+		return nil, err
+	}
+	return Ping(wsConn), nil
+}
+
+// WSDialer can dial a new websocket
+type WSDialer interface {
+	Dial(urlStr string, requestHeader http.Header) (*websocket.Conn, *http.Response, error)
+}
+
+// DialWS creates a new client connection. Use requestHeader to specify the
+// origin (Origin), subprotocols (Sec-WebSocket-Protocol) and cookies (Cookie).
+// Use the response.Header to get the selected subprotocol
+// (Sec-WebSocket-Protocol) and cookies (Set-Cookie).
+func DialWS(d WSDialer, urlStr string, requestHeader http.Header) (Websocket, *http.Response, error) {
+	wsConn, resp, err := d.Dial(urlStr, requestHeader)
+	if err != nil {
+		return nil, nil, err
+	}
+	return Ping(wsConn), resp, nil
+}
+
 // Ping adds a periodic ping to a websocket connection.
 func Ping(c *websocket.Conn) Websocket {
 	p := &pingingWebsocket{conn: c}
 	p.conn.SetPongHandler(p.pong)
+	p.conn.SetReadDeadline(mtime.Now().Add(pongWait))
 	p.pinger = time.AfterFunc(pingPeriod, p.ping)
-	p.ping()
 	return p
 }
 
@@ -57,6 +89,7 @@ func (p *pingingWebsocket) ping() {
 	p.writeLock.Lock()
 	defer p.writeLock.Unlock()
 	if err := p.conn.WriteControl(websocket.PingMessage, nil, mtime.Now().Add(writeWait)); err != nil {
+		log.Errorf("websocket ping error: %v", err)
 		p.Close()
 	}
 	p.pinger.Reset(pingPeriod)
@@ -72,11 +105,7 @@ func (p *pingingWebsocket) pong(string) error {
 func (p *pingingWebsocket) ReadMessage() (int, []byte, error) {
 	p.readLock.Lock()
 	defer p.readLock.Unlock()
-	messageType, b, err := p.conn.ReadMessage()
-	if err == nil {
-		err = p.conn.SetReadDeadline(mtime.Now().Add(pongWait))
-	}
-	return messageType, b, err
+	return p.conn.ReadMessage()
 }
 
 // WriteMessage is a helper method for getting a writer using NextWriter,
@@ -98,8 +127,10 @@ func (p *pingingWebsocket) WriteJSON(v interface{}) error {
 	if err != nil {
 		return err
 	}
+	if err := p.conn.SetWriteDeadline(mtime.Now().Add(writeWait)); err != nil {
+		return err
+	}
 	err1 := codec.NewEncoder(w, &codec.JsonHandle{}).Encode(v)
-	p.conn.SetWriteDeadline(mtime.Now().Add(writeWait))
 	err2 := w.Close()
 	if err1 != nil {
 		return err1
@@ -121,9 +152,6 @@ func (p *pingingWebsocket) ReadJSON(v interface{}) error {
 		// One value is expected in the message.
 		err = io.ErrUnexpectedEOF
 	}
-	if err == nil {
-		p.conn.SetReadDeadline(mtime.Now().Add(pongWait))
-	}
 	return err
 }
 

probe/appclient/app_client.go

@@ -188,11 +188,10 @@ func (c *appClient) controlConnection() (bool, error) {
 	headers := http.Header{}
 	c.ProbeConfig.authorizeHeaders(headers)
 	url := sanitize.URL("ws://", 0, "/api/control/ws")(c.target)
-	wsConn, _, err := c.wsDialer.Dial(url, headers)
+	conn, _, err := xfer.DialWS(&c.wsDialer, url, headers)
 	if err != nil {
 		return false, err
 	}
-	conn := xfer.Ping(wsConn)
 	defer conn.Close()
 
 	codec := xfer.NewJSONWebsocketCodec(conn)
@@ -270,7 +269,7 @@ func (c *appClient) pipeConnection(id string, pipe xfer.Pipe) (bool, error) {
 	headers := http.Header{}
 	c.ProbeConfig.authorizeHeaders(headers)
 	url := sanitize.URL("ws://", 0, fmt.Sprintf("/api/pipe/%s/probe", id))(c.target)
-	wsConn, resp, err := c.wsDialer.Dial(url, headers)
+	conn, resp, err := xfer.DialWS(&c.wsDialer, url, headers)
 	if resp != nil && resp.StatusCode == http.StatusNotFound {
 		// Special handling - 404 means the app/user has closed the pipe
 		pipe.Close()
@@ -279,7 +278,6 @@ func (c *appClient) pipeConnection(id string, pipe xfer.Pipe) (bool, error) {
 	if err != nil {
 		return false, err
 	}
-	conn := xfer.Ping(wsConn)
 
 	// Will return false if we are exiting
 	if !c.registerConn(id, conn) {