1package terminal
2
3import (
4	"fmt"
5
6	"github.com/gorilla/websocket"
7)
8
9type WebSocketProxy struct {
10	StopCh chan error
11}
12
13// stoppers is the number of goroutines that may attempt to call Stop()
14func NewWebSocketProxy(stoppers int) *WebSocketProxy {
15	return &WebSocketProxy{
16		StopCh: make(chan error, stoppers+2), // each proxy() call is a stopper
17	}
18}
19
20func (p *WebSocketProxy) GetStopCh() chan error {
21	return p.StopCh
22}
23
24func (p *WebSocketProxy) Serve(upstream, downstream Connection, upstreamAddr, downstreamAddr string) error {
25	// This signals the upstream terminal to kill the exec'd process
26	defer upstream.WriteMessage(websocket.BinaryMessage, eot)
27
28	go p.proxy(upstream, downstream, upstreamAddr, downstreamAddr)
29	go p.proxy(downstream, upstream, downstreamAddr, upstreamAddr)
30
31	err := <-p.StopCh
32	return err
33}
34
35func (p *WebSocketProxy) proxy(to, from Connection, toAddr, fromAddr string) {
36	for {
37		messageType, data, err := from.ReadMessage()
38		if err != nil {
39			p.StopCh <- fmt.Errorf("reading from %s: %s", fromAddr, err)
40			break
41		}
42
43		if err := to.WriteMessage(messageType, data); err != nil {
44			p.StopCh <- fmt.Errorf("writing to %s: %s", toAddr, err)
45			break
46		}
47	}
48}
49