fix(telnet): Fix connections occasionally being left open

gabe565 · Nov 13, 2024 · 9e93ac0 · 9e93ac0
1 parent f46fc14
commit 9e93ac0
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 10 deletions.
diff --git a/internal/server/telnet.go b/internal/server/telnet.go
@@ -110,10 +110,6 @@ func (s *TelnetServer) Handler(ctx context.Context, conn net.Conn, m *movie.Movi
 	defer func() {
 		_ = in.Close()
 	}()
-	go func() {
-		<-errCh
-		cancel()
-	}()
 
 	gotProfile := profile != -1
 	if !gotProfile {
@@ -124,7 +120,6 @@ func (s *TelnetServer) Handler(ctx context.Context, conn net.Conn, m *movie.Movi
 	defer p.Close()
 
 	opts := []tea.ProgramOption{
-		tea.WithContext(ctx),
 		tea.WithInput(in),
 		tea.WithOutput(conn),
 		tea.WithFPS(30),
@@ -135,8 +130,14 @@ func (s *TelnetServer) Handler(ctx context.Context, conn net.Conn, m *movie.Movi
 	program := tea.NewProgram(p, opts...)
 
 	go func() {
-		for info := range sizeCh {
-			if info.Width != 0 && info.Height != 0 {
+		for {
+			select {
+			case <-ctx.Done():
+				program.Quit()
+				return
+			case <-errCh:
+				cancel()
+			case info := <-sizeCh:
 				program.Send(tea.WindowSizeMsg{
 					Width:  int(info.Width),
 					Height: int(info.Height),
@@ -149,5 +150,8 @@ func (s *TelnetServer) Handler(ctx context.Context, conn net.Conn, m *movie.Movi
 		logger.Error("Program failed", "error", err)
 	}
 
+	// p.Kill() will force kill the program if it's still running,
+	// and restore the terminal to its original state in case of a
+	// tui crash
 	program.Kill()
 }
diff --git a/internal/server/telnet/proxy.go b/internal/server/telnet/proxy.go
@@ -87,7 +87,7 @@ outer:
 					Iac, Do, TerminalType,
 					Iac, Do, NegotiateAboutWindowSize,
 				); err != nil {
-					slog.Error("Failed to write telnet commands", "error", err)
+					return err
 				}
 			}
 
@@ -134,7 +134,9 @@ outer:
 								return err
 							}
 							slog.Log(context.Background(), config.LevelTrace, "Got window size", "size", size)
-							sizeCh <- size
+							if size.Width != 0 && size.Height != 0 {
+								sizeCh <- size
+							}
 						}
 					}
 				}
@@ -156,7 +158,7 @@ outer:
 					if !willNegotiateAboutWindowSize {
 						willNegotiateAboutWindowSize = true
 						if _, err := Write(conn, Iac, Do, NegotiateAboutWindowSize); err != nil {
-							slog.Error("Failed to write telnet commands", "error", err)
+							return err
 						}
 					}
 				}