From 35fe8864a1026e473e2280a30155aa004c0b7b56 Mon Sep 17 00:00:00 2001 From: Paul Bellamy Date: Thu, 25 Feb 2016 11:52:33 +0000 Subject: [PATCH] Review Feedback --- app/api_topology.go | 8 +----- app/controls.go | 2 +- app/pipes.go | 3 +-- common/xfer/websocket.go | 48 +++++++++++++++++++++++++++-------- probe/appclient/app_client.go | 6 ++--- 5 files changed, 43 insertions(+), 24 deletions(-) diff --git a/app/api_topology.go b/app/api_topology.go index 2cbb7e501..16f03f192 100644 --- a/app/api_topology.go +++ b/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{}) diff --git a/app/controls.go b/app/controls.go index 927846a41..2548651af 100644 --- a/app/controls.go +++ b/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 diff --git a/app/pipes.go b/app/pipes.go index 844833a0c..c4724a360 100644 --- a/app/pipes.go +++ b/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) diff --git a/common/xfer/websocket.go b/common/xfer/websocket.go index a3f0d3d3a..f5faabe54 100644 --- a/common/xfer/websocket.go +++ b/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 } diff --git a/probe/appclient/app_client.go b/probe/appclient/app_client.go index 6458613b0..8d8409f0a 100644 --- a/probe/appclient/app_client.go +++ b/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) {