Merge pull request #21 from TvdBrink/ping_pong

Allow usage of the ping and pong handling of Gorilla websocket.

Author: tmc

wsproxy/websocket_proxy.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"net/http"
 	"strings"
+	"time"
 
 	"github.com/gorilla/websocket"
 	"github.com/sirupsen/logrus"
@@ -34,6 +35,9 @@ type Proxy struct {
 	tokenCookieName        string
 	requestMutator         RequestMutatorFunc
 	headerForwarder        func(header string) bool
+	pingInterval           time.Duration
+	pingWait               time.Duration
+	pongWait               time.Duration
 }
 
 // Logger collects log messages.
@@ -97,6 +101,17 @@ func WithLogger(logger Logger) Option {
 	}
 }
 
+// WithPingControl allows specification of ping pong control. The interval
+// parameter specifies the pingInterval between pings. The allowed wait time
+// for a pong response is (pingInterval * 10) / 9.
+func WithPingControl(interval time.Duration) Option {
+	return func(proxy *Proxy) {
+		proxy.pingInterval = interval
+		proxy.pongWait = (interval * 10) / 9
+		proxy.pingWait = proxy.pongWait / 6
+	}
+}
+
 var defaultHeadersToForward = map[string]bool{
 	"Origin":  true,
 	"origin":  true,
@@ -211,6 +226,10 @@ func (p *Proxy) proxy(w http.ResponseWriter, r *http.Request) {
 
 	// read loop -- take messages from websocket and write to http request
 	go func() {
+		if p.pingInterval > 0 && p.pingWait > 0 && p.pongWait > 0 {
+			conn.SetReadDeadline(time.Now().Add(p.pongWait))
+			conn.SetPongHandler(func(string) error { conn.SetReadDeadline(time.Now().Add(p.pongWait)); return nil })
+		}
 		defer func() {
 			cancelFn()
 		}()
@@ -242,6 +261,28 @@ func (p *Proxy) proxy(w http.ResponseWriter, r *http.Request) {
 			}
 		}
 	}()
+	// ping write loop
+	if p.pingInterval > 0 && p.pingWait > 0 && p.pongWait > 0 {
+		go func() {
+			ticker := time.NewTicker(p.pingInterval)
+			defer func() {
+				ticker.Stop()
+				conn.Close()
+			}()
+			for {
+				select {
+				case <-ctx.Done():
+					p.logger.Debugln("ping loop done")
+					return
+				case <-ticker.C:
+					conn.SetWriteDeadline(time.Now().Add(p.pingWait))
+					if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
+						return
+					}
+				}
+			}
+		}()
+	}
 	// write loop -- take messages from response and write to websocket
 	scanner := bufio.NewScanner(responseBodyR)