1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
|
package server
import (
"container/list"
"sync"
"github.com/gorilla/websocket"
log "github.com/sirupsen/logrus"
)
type ttyShareSession struct {
mainRWLock sync.RWMutex
ttyProtoConnections *list.List
isAlive bool
lastWindowSizeMsg MsgTTYWinSize
ptyHandler PTYHandler
}
func copyList(l *list.List) *list.List {
newList := list.New()
for e := l.Front(); e != nil; e = e.Next() {
newList.PushBack(e.Value)
}
return newList
}
func newTTYShareSession(ptyHandler PTYHandler) *ttyShareSession {
ttyShareSession := &ttyShareSession{
ttyProtoConnections: list.New(),
ptyHandler: ptyHandler,
}
return ttyShareSession
}
func (session *ttyShareSession) WindowSize(cols, rows int) error {
session.mainRWLock.Lock()
session.lastWindowSizeMsg = MsgTTYWinSize{Cols: cols, Rows: rows}
session.mainRWLock.Unlock()
session.forEachReceiverLock(func(rcvConn *TTYProtocolWSLocked) bool {
rcvConn.SetWinSize(cols, rows)
return true
})
return nil
}
func (session *ttyShareSession) Write(data []byte) (int, error) {
session.forEachReceiverLock(func(rcvConn *TTYProtocolWSLocked) bool {
rcvConn.Write(data)
return true
})
return len(data), nil
}
// Runs the callback cb for each of the receivers in the list of the receivers, as it was when
// this function was called. Note that there might be receivers which might have lost
// the connection since this function was called.
// Return false in the callback to not continue for the rest of the receivers
func (session *ttyShareSession) forEachReceiverLock(cb func(rcvConn *TTYProtocolWSLocked) bool) {
session.mainRWLock.RLock()
// TODO: Maybe find a better way?
rcvsCopy := copyList(session.ttyProtoConnections)
session.mainRWLock.RUnlock()
for receiverE := rcvsCopy.Front(); receiverE != nil; receiverE = receiverE.Next() {
receiver := receiverE.Value.(*TTYProtocolWSLocked)
if !cb(receiver) {
break
}
}
}
// Will run on the TTYReceiver connection go routine (e.g.: on the websockets connection routine)
// When HandleWSConnection will exit, the connection to the TTYReceiver will be closed
func (session *ttyShareSession) HandleWSConnection(wsConn *websocket.Conn) {
protoConn := NewTTYProtocolWSLocked(wsConn)
session.mainRWLock.Lock()
rcvHandleEl := session.ttyProtoConnections.PushBack(protoConn)
winSize := session.lastWindowSizeMsg
session.mainRWLock.Unlock()
log.Debugf("New WS connection (%s). Serving ..", wsConn.RemoteAddr().String())
// Sending the initial size of the window, if we have one
protoConn.SetWinSize(winSize.Cols, winSize.Rows)
// Wait until the TTYReceiver will close the connection on its end
for {
err := protoConn.ReadAndHandle(
func(data []byte) {
session.ptyHandler.Write(data)
},
func(cols, rows int) {
session.ptyHandler.Refresh()
},
)
if err != nil {
log.Debugf("Finished the WS reading loop: %s", err.Error())
break
}
}
// Remove the recevier from the list of the receiver of this session, so we need to write-lock
session.mainRWLock.Lock()
session.ttyProtoConnections.Remove(rcvHandleEl)
session.mainRWLock.Unlock()
wsConn.Close()
log.Debugf("Closed receiver connection")
}
|
