session.go

package main

import (
	"container/list"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"net"
	"sync"
)

type sessionInfo struct {
	ID              string
	URLWebReadWrite string
}

type ttyShareSession struct {
	sessionID              string
	serverURL              string
	mainRWLock             sync.RWMutex
	ttySenderConnection    *TTYProtocolConn
	ttyReceiverConnections *list.List
	isAlive                bool
	lastWindowSizeMsg      MsgAll
}

func generateNewSessionID() string {
	binID := make([]byte, 32)
	_, err := rand.Read(binID)

	if err != nil {
		panic(err)
	}

	return base64.URLEncoding.EncodeToString([]byte(binID))
}

func newTTYShareSession(conn net.Conn, serverURL string) *ttyShareSession {
	sessionID := generateNewSessionID()

	ttyShareSession := &ttyShareSession{
		sessionID:              sessionID,
		serverURL:              serverURL,
		ttySenderConnection:    NewTTYProtocolConn(conn),
		ttyReceiverConnections: list.New(),
	}

	return ttyShareSession
}

func (session *ttyShareSession) InitSender() error {
	_, err := session.ttySenderConnection.InitServer(ServerSessionInfo{
		URLWebReadWrite: "\n\n **Your tty-share version is deprecated, and will stop working soon. Please update to the latest version!**\n\n" + session.serverURL + "/s/" + session.GetID(),
	})
	return err
}

func (session *ttyShareSession) GetID() string {
	return session.sessionID
}

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 (session *ttyShareSession) handleSenderMessageLock(msg MsgAll) {
	switch msg.Type {
	case MsgIDWinSize:
		// Save the last known size of the window so we pass it to new receivers, and then
		// fallthrough. We save the WinSize message as we get it, since we send it anyways
		// to the receivers, packed into the same protocol
		session.mainRWLock.Lock()
		session.lastWindowSizeMsg = msg
		session.mainRWLock.Unlock()
		fallthrough
	case MsgIDWrite:
		data, _ := json.Marshal(msg)
		session.forEachReceiverLock(func(rcvConn *TTYProtocolConn) bool {
			rcvConn.WriteRawData(data)
			return true
		})
	}
}

// Will run on the ttySendeConnection go routine (e.g.: in the TCP connection routine)
func (session *ttyShareSession) HandleSenderConnection() {
	session.mainRWLock.Lock()
	session.isAlive = true
	senderConnection := session.ttySenderConnection
	session.mainRWLock.Unlock()

	for {
		msg, err := senderConnection.ReadMessage()
		if err != nil {
			log.Debugf("TTYSender connection finished withs with error: %s", err.Error())
			break
		}
		session.handleSenderMessageLock(msg)
	}

	// Close the connection to all the receivers
	log.Debugf("Closing all receiver connection")
	session.forEachReceiverLock(func(recvConn *TTYProtocolConn) bool {
		log.Debugf("Closing receiver connection")
		recvConn.Close()
		return true
	})

	// TODO: clear here the list of receiver
	session.mainRWLock.Lock()
	session.isAlive = false
	session.mainRWLock.Unlock()
}

// 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 *TTYProtocolConn) bool) {
	session.mainRWLock.RLock()
	// TODO: Maybe find a better way?
	rcvsCopy := copyList(session.ttyReceiverConnections)
	session.mainRWLock.RUnlock()

	for receiverE := rcvsCopy.Front(); receiverE != nil; receiverE = receiverE.Next() {
		receiver := receiverE.Value.(*TTYProtocolConn)
		if !cb(receiver) {
			break
		}
	}
}

// Will run on the TTYReceiver connection go routine (e.g.: on the websockets connection routine)
// When HandleReceiver will exit, the connection to the TTYReceiver will be closed
func (session *ttyShareSession) HandleReceiver(rawConn *WSConnection) {
	rcvProtoConn := NewTTYProtocolConn(rawConn)

	session.mainRWLock.Lock()
	if !session.isAlive {
		log.Warnf("TTYReceiver tried to connect to a session that is not alive anymore. Rejecting it..")
		session.mainRWLock.Unlock()
		return
	}

	// Add the receiver to the list of receivers in the seesion, so we need to write-lock
	rcvHandleEl := session.ttyReceiverConnections.PushBack(rcvProtoConn)
	senderConn := session.ttySenderConnection
	lastWindowSize, _ := json.Marshal(session.lastWindowSizeMsg)
	session.mainRWLock.Unlock()

	log.Debugf("Got new TTYReceiver connection (%s). Serving it..", rawConn.Address())

	// Sending the initial size of the window, if we have one
	rcvProtoConn.WriteRawData(lastWindowSize)

	// Notify the tty-share that we got a new receiver connected
	msgRcvConnected, err := MarshalMsg(MsgTTYSenderNewReceiverConnected{
		Name: rawConn.Address(),
	})
	senderConn.WriteRawData(msgRcvConnected)

	if err != nil {
		log.Errorf("Cannot notify tty sender. Error: %s", err.Error())
	}

	// Wait until the TTYReceiver will close the connection on its end
	for {
		msg, err := rcvProtoConn.ReadMessage()

		if err != nil {
			log.Warnf("Finishing handling the TTYReceiver loop because: %s", err.Error())
			break
		}

		switch msg.Type {
		case MsgIDWinSize:
			// Ignore these messages from the receiver. For now, the policy is that the sender
			// decides on the window size.
		case MsgIDWrite:
			rawData, _ := json.Marshal(msg)
			senderConn.WriteRawData(rawData)
		default:
			log.Warnf("Receiving unknown data from the receiver")
		}
	}

	log.Debugf("Closing receiver connection")
	rcvProtoConn.Close()

	// Remove the recevier from the list of the receiver of this session, so we need to write-lock
	session.mainRWLock.Lock()
	session.ttyReceiverConnections.Remove(rcvHandleEl)
	session.mainRWLock.Unlock()
}