From 896dbbbc471c4f259acdb5489f33e4d514c209ea Mon Sep 17 00:00:00 2001
From: Isaiah Becker-Mayer <isaiah@goteleport.com>
Date: Thu, 3 Mar 2022 15:07:38 -0500
Subject: [PATCH] desktop playback error handling (#10765)

---
 .github/ISSUE_TEMPLATE/testplan.md |  4 ++++
 lib/srv/desktop/windows_server.go  |  2 +-
 lib/web/desktop/playback.go        | 16 ++++++++++++++++
 lib/web/desktop/playback_test.go   |  2 +-
 4 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/.github/ISSUE_TEMPLATE/testplan.md b/.github/ISSUE_TEMPLATE/testplan.md
index e394322075c..d1b9d5f634b 100644
--- a/.github/ISSUE_TEMPLATE/testplan.md
+++ b/.github/ISSUE_TEMPLATE/testplan.md
@@ -938,6 +938,10 @@ and non interactive tsh bench loads.
   - [ ] Verify async recording (`mode: node` or `mode: proxy`)
   - [ ] Sessions show up in session recordings UI with desktop icon
   - [ ] Sessions can be played back, including play/pause functionality
+  - [ ] A session that ends with a TDP error message can be played back, ends by displaying the error message,
+        and the progress bar progresses to the end.
+  - [ ] Attempting to play back a session that doesn't exist (i.e. by entering a non-existing session id in the url) shows
+        a relevant error message.
   - [ ] RBAC for sessions: ensure users can only see their own recordings when
     using the RBAC rule from our
     [docs](../../docs/pages/access-controls/reference.mdx#rbac-for-sessions)
diff --git a/lib/srv/desktop/windows_server.go b/lib/srv/desktop/windows_server.go
index f62027af1e0..979fa238a2e 100644
--- a/lib/srv/desktop/windows_server.go
+++ b/lib/srv/desktop/windows_server.go
@@ -872,7 +872,7 @@ func (s *WindowsService) makeTDPSendHandler(ctx context.Context, emitter events.
 	id *tlsca.Identity, sessionID, desktopAddr string) func(m tdp.Message, b []byte) {
 	return func(m tdp.Message, b []byte) {
 		switch b[0] {
-		case byte(tdp.TypePNGFrame):
+		case byte(tdp.TypePNGFrame), byte(tdp.TypeError):
 			e := &events.DesktopRecording{
 				Metadata: events.Metadata{
 					Type: libevents.DesktopRecordingEvent,
diff --git a/lib/web/desktop/playback.go b/lib/web/desktop/playback.go
index 81024a0de52..79f0a7c86d1 100644
--- a/lib/web/desktop/playback.go
+++ b/lib/web/desktop/playback.go
@@ -19,13 +19,16 @@ package desktop
 import (
 	"context"
 	"errors"
+	"fmt"
 	"net"
+	"os"
 	"sync"
 	"time"
 
 	apievents "github.com/gravitational/teleport/api/types/events"
 	"github.com/gravitational/teleport/lib/session"
 	"github.com/gravitational/teleport/lib/utils"
+	"github.com/gravitational/trace"
 	"github.com/sirupsen/logrus"
 	"golang.org/x/net/websocket"
 )
@@ -207,6 +210,15 @@ func (pp *Player) streamSessionEvents(ctx context.Context, cancel context.Cancel
 			// otherwise it just sits at the player UI
 			if err != nil && !errors.Is(err, context.Canceled) {
 				pp.log.WithError(err).Errorf("streaming session %v", pp.sID)
+				var errorText string
+				if os.IsNotExist(err) || trace.IsNotFound(err) {
+					errorText = "session not found"
+				} else {
+					errorText = "server error"
+				}
+				if _, err := pp.ws.Write([]byte(fmt.Sprintf(`{"message": "error", "errorText": "%v"}`, errorText))); err != nil {
+					pp.log.WithError(err).Error("failed to write \"error\" message over websocket")
+				}
 			}
 			return
 		case evt := <-eventsC:
@@ -227,6 +239,10 @@ func (pp *Player) streamSessionEvents(ctx context.Context, cancel context.Cancel
 				msg, err := utils.FastMarshal(e)
 				if err != nil {
 					pp.log.WithError(err).Errorf("failed to marshal DesktopRecording event into JSON: %v", e)
+					if _, err := pp.ws.Write([]byte(`{"message":"error","errorText":"server error"}`)); err != nil {
+						pp.log.WithError(err).Error("failed to write \"error\" message over websocket")
+					}
+					return
 				}
 				if _, err := pp.ws.Write(msg); err != nil {
 					// We expect net.ErrClosed to arise when another goroutine returns before
diff --git a/lib/web/desktop/playback_test.go b/lib/web/desktop/playback_test.go
index bb0be1f9089..33481125807 100644
--- a/lib/web/desktop/playback_test.go
+++ b/lib/web/desktop/playback_test.go
@@ -65,7 +65,7 @@ func TestStreamsDesktopEvents(t *testing.T) {
 	b := make([]byte, 4096)
 	n, err := ws.Read(b)
 	require.NoError(t, err)
-	require.Equal(t, []byte(`{"message":"end"}`), b[:n])
+	require.JSONEq(t, `{"message":"end"}`, string(b[:n]))
 }
 
 func newServer(t *testing.T, streamInterval time.Duration, events []apievents.AuditEvent) *httptest.Server {