Merge pull request #21 from TvdBrink/ping_pong
Allow usage of the ping and pong handling of Gorilla websocket.
Travis Cline authored 4 years ago
GitHub committed 4 years ago
| 5 | 5 | "io" |
| 6 | 6 | "net/http" |
| 7 | 7 | "strings" |
| 8 | "time" | |
| 8 | 9 | |
| 9 | 10 | "github.com/gorilla/websocket" |
| 10 | 11 | "github.com/sirupsen/logrus" |
| 33 | 34 | tokenCookieName string |
| 34 | 35 | requestMutator RequestMutatorFunc |
| 35 | 36 | headerForwarder func(header string) bool |
| 37 | pingInterval time.Duration | |
| 38 | pingWait time.Duration | |
| 39 | pongWait time.Duration | |
| 36 | 40 | } |
| 37 | 41 | |
| 38 | 42 | // Logger collects log messages. |
| 93 | 97 | func WithLogger(logger Logger) Option { |
| 94 | 98 | return func(p *Proxy) { |
| 95 | 99 | p.logger = logger |
| 100 | } | |
| 101 | } | |
| 102 | ||
| 103 | // WithPingControl allows specification of ping pong control. The interval | |
| 104 | // parameter specifies the pingInterval between pings. The allowed wait time | |
| 105 | // for a pong response is (pingInterval * 10) / 9. | |
| 106 | func WithPingControl(interval time.Duration) Option { | |
| 107 | return func(proxy *Proxy) { | |
| 108 | proxy.pingInterval = interval | |
| 109 | proxy.pongWait = (interval * 10) / 9 | |
| 110 | proxy.pingWait = proxy.pongWait / 6 | |
| 96 | 111 | } |
| 97 | 112 | } |
| 98 | 113 | |
| 210 | 225 | |
| 211 | 226 | // read loop -- take messages from websocket and write to http request |
| 212 | 227 | go func() { |
| 228 | if p.pingInterval > 0 && p.pingWait > 0 && p.pongWait > 0 { | |
| 229 | conn.SetReadDeadline(time.Now().Add(p.pongWait)) | |
| 230 | conn.SetPongHandler(func(string) error { conn.SetReadDeadline(time.Now().Add(p.pongWait)); return nil }) | |
| 231 | } | |
| 213 | 232 | defer func() { |
| 214 | 233 | cancelFn() |
| 215 | 234 | }() |
| 241 | 260 | } |
| 242 | 261 | } |
| 243 | 262 | }() |
| 263 | // ping write loop | |
| 264 | if p.pingInterval > 0 && p.pingWait > 0 && p.pongWait > 0 { | |
| 265 | go func() { | |
| 266 | ticker := time.NewTicker(p.pingInterval) | |
| 267 | defer func() { | |
| 268 | ticker.Stop() | |
| 269 | conn.Close() | |
| 270 | }() | |
| 271 | for { | |
| 272 | select { | |
| 273 | case <-ctx.Done(): | |
| 274 | p.logger.Debugln("ping loop done") | |
| 275 | return | |
| 276 | case <-ticker.C: | |
| 277 | conn.SetWriteDeadline(time.Now().Add(p.pingWait)) | |
| 278 | if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil { | |
| 279 | return | |
| 280 | } | |
| 281 | } | |
| 282 | } | |
| 283 | }() | |
| 284 | } | |
| 244 | 285 | // write loop -- take messages from response and write to websocket |
| 245 | 286 | scanner := bufio.NewScanner(responseBodyR) |
| 246 | 287 | |