From 7efeccad52175cede3a457e7fb82b77307d2c52e Mon Sep 17 00:00:00 2001
From: Jannis Mattheis <contact@jmattheis.de>
Date: Fri, 9 Oct 2020 14:15:01 +0200
Subject: [PATCH] Properly close share sessions

---
 ui/src/message.ts        |  4 ++++
 ui/src/useRoom.ts        | 15 ++++++++++++++-
 ws/event_disconnected.go | 22 ++++++++++++++++++++++
 ws/event_stop_share.go   | 11 +++++++++++
 ws/outgoing/messages.go  |  6 ++++++
 5 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/ui/src/message.ts b/ui/src/message.ts
index b80ddab..d45ec5e 100644
--- a/ui/src/message.ts
+++ b/ui/src/message.ts
@@ -77,8 +77,10 @@ export type ClientICECandidate = Typed<P2PMessage<RTCIceCandidate>, 'clientice'>
 export type HostOffer = Typed<P2PMessage<RTCSessionDescriptionInit>, 'hostoffer'>;
 export type ClientAnswer = Typed<P2PMessage<RTCSessionDescriptionInit>, 'clientanswer'>;
 export type StartSharing = Typed<{}, 'share'>;
+export type StopShare = Typed<{}, 'stopshare'>;
 export type RoomCreate = Typed<RoomConfiguration, 'create'>;
 export type JoinRoom = Typed<JoinConfiguration, 'join'>;
+export type EndShare = Typed<string, 'endshare'>;
 
 export type IncomingMessage =
     | Room
@@ -88,6 +90,7 @@ export type IncomingMessage =
     | HostICECandidate
     | ClientICECandidate
     | HostOffer
+    | EndShare
     | ClientAnswer;
 
 export type OutgoingMessage =
@@ -97,5 +100,6 @@ export type OutgoingMessage =
     | HostICECandidate
     | ClientICECandidate
     | HostOffer
+    | StopShare
     | ClientAnswer
     | StartSharing;
diff --git a/ui/src/useRoom.ts b/ui/src/useRoom.ts
index dabea59..bc07136 100644
--- a/ui/src/useRoom.ts
+++ b/ui/src/useRoom.ts
@@ -242,6 +242,19 @@ export const useRoom = (): UseRoom => {
                         case 'hostice':
                             client.current[event.payload.sid]?.addIceCandidate(event.payload.value);
                             return;
+                        case 'endshare':
+                            client.current[event.payload]?.close();
+                            host.current[event.payload]?.close();
+                            setState((current) =>
+                                current
+                                    ? {
+                                          ...current,
+                                          clientStreams: current.clientStreams.filter(
+                                              ({id}) => id !== event.payload
+                                          ),
+                                      }
+                                    : current
+                            );
                     }
                 };
                 ws.onclose = (event) => {
@@ -285,7 +298,7 @@ export const useRoom = (): UseRoom => {
         host.current = {};
         stream.current?.getTracks().forEach((track) => track.stop());
         stream.current = undefined;
-        // todo notify server
+        conn.current?.send(JSON.stringify({type: 'stopshare'}));
         setState((current) => (current ? {...current, hostStream: undefined} : current));
     };
 
diff --git a/ws/event_disconnected.go b/ws/event_disconnected.go
index 2312959..84e7c81 100644
--- a/ws/event_disconnected.go
+++ b/ws/event_disconnected.go
@@ -1,5 +1,10 @@
 package ws
 
+import (
+	"bytes"
+	"github.com/screego/server/ws/outgoing"
+)
+
 type Disconnected struct {
 }
 
@@ -24,6 +29,23 @@ func (e *Disconnected) Execute(rooms *Rooms, current ClientInfo) error {
 	current.Close <- CloseDone
 	delete(room.Users, current.ID)
 
+	for id, session := range room.Sessions {
+		if bytes.Equal(session.Client.Bytes(), current.ID.Bytes()) {
+			host, ok := room.Users[session.Host]
+			if ok {
+				host.Write <- outgoing.EndShare(id)
+			}
+			delete(room.Sessions, id)
+		}
+		if bytes.Equal(session.Host.Bytes(), current.ID.Bytes()) {
+			client, ok := room.Users[session.Client]
+			if ok {
+				client.Write <- outgoing.EndShare(id)
+			}
+			delete(room.Sessions, id)
+		}
+	}
+
 	if user.Owner && room.CloseOnOwnerLeave {
 		for _, member := range room.Users {
 			member.Close <- CloseOwnerLeft
diff --git a/ws/event_stop_share.go b/ws/event_stop_share.go
index ff42df3..ab3148b 100644
--- a/ws/event_stop_share.go
+++ b/ws/event_stop_share.go
@@ -1,7 +1,9 @@
 package ws
 
 import (
+	"bytes"
 	"fmt"
+	"github.com/screego/server/ws/outgoing"
 )
 
 func init() {
@@ -24,6 +26,15 @@ func (e *StopShare) Execute(rooms *Rooms, current ClientInfo) error {
 	}
 
 	room.Users[current.ID].Sharing = false
+	for id, session := range room.Sessions {
+		if bytes.Equal(session.Host.Bytes(), current.ID.Bytes()) {
+			client, ok := room.Users[session.Client]
+			if ok {
+				client.Write <- outgoing.EndShare(id)
+			}
+			delete(room.Sessions, id)
+		}
+	}
 
 	room.notifyInfoChanged()
 	return nil
diff --git a/ws/outgoing/messages.go b/ws/outgoing/messages.go
index 2b8e553..de14f3c 100644
--- a/ws/outgoing/messages.go
+++ b/ws/outgoing/messages.go
@@ -83,6 +83,12 @@ func (HostOffer) Type() string {
 	return "hostoffer"
 }
 
+type EndShare xid.ID
+
+func (EndShare) Type() string {
+	return "endshare"
+}
+
 type ConnectionMode string
 
 const (