From 885e5f1766a76ba6612066849353253f6414bbb8 Mon Sep 17 00:00:00 2001
From: Vasile Popescu <elisescu@elisescu.com>
Date: Sat, 12 May 2018 23:44:59 +0200
Subject: [PATCH] Add a way to notify the sender when new receivers connect

Also try to "refresh" the terminal, when new receivers connect, so the
content is fully rendered on the receiver side.
At the moment, the "refresh" is faked with making the window smaller and
then back bigger, after a short time (there doesn't seem to be a way to
tell the app to re-draw itself. Or is there?).
---
 common/protocol.go       | 30 ++++++++++++++++++++++++------
 tty-sender/main.go       |  7 +++++++
 tty-sender/pty_master.go | 21 ++++++++++++++++++++-
 tty-server/session.go    | 10 ++++++++++
 4 files changed, 61 insertions(+), 7 deletions(-)

diff --git a/common/protocol.go b/common/protocol.go
index a98a335..8cb63a5 100644
--- a/common/protocol.go
+++ b/common/protocol.go
@@ -10,19 +10,22 @@ import (
 type ProtocolMessageIDType string
 
 const (
-	MsgIDSenderInitRequest   = "SenderInitRequest"
-	MsgIDSenderInitReply     = "SenderInitReply"
-	MsgIDReceiverInitRequest = "ReceiverInitRequest"
-	MsgIDReceiverInitReply   = "ReceiverInitReply"
-	MsgIDWrite               = "Write"
-	MsgIDWinSize             = "WinSize"
+	MsgIDSenderInitRequest          = "SenderInitRequest"
+	MsgIDSenderInitReply            = "SenderInitReply"
+	MsgIDSenderNewReceiverConnected = "SenderNewReceiverConnected"
+	MsgIDReceiverInitRequest        = "ReceiverInitRequest"
+	MsgIDReceiverInitReply          = "ReceiverInitReply"
+	MsgIDWrite                      = "Write"
+	MsgIDWinSize                    = "WinSize"
 )
 
+// Message used to encapsulate the rest of the bessages bellow
 type MsgAll struct {
 	Type ProtocolMessageIDType
 	Data []byte
 }
 
+// These messages are used between the server and the sender/receiver
 type MsgTTYSenderInitRequest struct {
 	Salt              string
 	PasswordVerifierA string
@@ -32,6 +35,10 @@ type MsgTTYSenderInitReply struct {
 	ReceiverURLWebReadWrite string
 }
 
+type MsgTTYSenderNewReceiverConnected struct {
+	Name string
+}
+
 type MsgTTYReceiverInitRequest struct {
 	ChallengeReply string
 }
@@ -39,6 +46,8 @@ type MsgTTYReceiverInitRequest struct {
 type MsgTTYReceiverInitReply struct {
 }
 
+// These messages are not intended for the server, so they are just forwarded by it to the remote
+// side.
 type MsgTTYWrite struct {
 	Data []byte
 	Size int
@@ -107,6 +116,15 @@ func MarshalMsg(aMessage interface{}) (_ []byte, err error) {
 		return json.Marshal(msg)
 	}
 
+	if newRcvMsg, ok := aMessage.(MsgTTYSenderNewReceiverConnected); ok {
+		msg.Type = MsgIDSenderNewReceiverConnected
+		msg.Data, err = json.Marshal(newRcvMsg)
+		if err != nil {
+			return
+		}
+		return json.Marshal(msg)
+	}
+
 	return nil, nil
 }
 
diff --git a/tty-sender/main.go b/tty-sender/main.go
index 67006a3..f82ae3e 100644
--- a/tty-sender/main.go
+++ b/tty-sender/main.go
@@ -114,6 +114,13 @@ func main() {
 				json.Unmarshal(msg.Data, &msgWrite)
 				ptyMaster.Write(msgWrite.Data[:msgWrite.Size])
 			}
+			if msg.Type == common.MsgIDSenderNewReceiverConnected {
+				var msgReceiverConnected common.MsgTTYSenderNewReceiverConnected
+				json.Unmarshal(msg.Data, &msgReceiverConnected)
+
+				ptyMaster.Refresh()
+				fmt.Printf("New receiver connected: %s ", msgReceiverConnected.Name)
+			}
 		}
 	}()
 
diff --git a/tty-sender/pty_master.go b/tty-sender/pty_master.go
index 3155394..6889dba 100644
--- a/tty-sender/pty_master.go
+++ b/tty-sender/pty_master.go
@@ -5,6 +5,7 @@ import (
 	"os/exec"
 	"os/signal"
 	"syscall"
+	"time"
 
 	ptyDevice "github.com/elisescu/pty"
 	"golang.org/x/crypto/ssh/terminal"
@@ -61,7 +62,8 @@ func (pty *ptyMaster) Start(command string, args []string, winChangedCB onWindow
 }
 
 func (pty *ptyMaster) GetWinSize() (int, int, error) {
-	return terminal.GetSize(0)
+	cols, rows, err := terminal.GetSize(0)
+	return cols, rows, err
 }
 
 func (pty *ptyMaster) Write(b []byte) (int, error) {
@@ -76,6 +78,23 @@ func (pty *ptyMaster) SetWinSize(rows, cols int) {
 	ptyDevice.Setsize(pty.ptyFile, rows, cols)
 }
 
+func (pty *ptyMaster) Refresh() {
+	// We wanna force the app to re-draw itself, but there doesn't seem to be a way to do that
+	// so we fake it by resizing the window quickly, making it smaller and then back big
+	cols, rows, err := pty.GetWinSize()
+
+	if err != nil {
+		return
+	}
+
+	pty.SetWinSize(rows-1, cols)
+
+	go func() {
+		time.Sleep(time.Millisecond * 50)
+		pty.SetWinSize(rows, cols)
+	}()
+}
+
 func (pty *ptyMaster) Wait() (err error) {
 	err = pty.command.Wait()
 	// The terminal has to be restored from the RAW state, to its initial state
diff --git a/tty-server/session.go b/tty-server/session.go
index 3cbebe4..fa4fc7f 100644
--- a/tty-server/session.go
+++ b/tty-server/session.go
@@ -153,6 +153,16 @@ func (session *ttyShareSession) HandleReceiver(rawConn *WSConnection) {
 	// Sending the initial size of the window, if we have one
 	rcvProtoConn.WriteRawData(lastWindowSize)
 
+	// Notify the tty-sender 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()