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