refactor(telnet): Clean up telnet proxy logic

internal/server/telnet.go

@@ -3,7 +3,6 @@ package server
 import (
 	"context"
 	"errors"
-	"io"
 	"net"
 	"sync"
 	"time"
@@ -13,7 +12,6 @@ import (
 	"github.com/gabe565/ascii-movie/internal/player"
 	"github.com/gabe565/ascii-movie/internal/server/idleconn"
 	"github.com/gabe565/ascii-movie/internal/server/telnet"
-	"github.com/gabe565/ascii-movie/internal/util"
 	"github.com/muesli/termenv"
 	flag "github.com/spf13/pflag"
 )
@@ -105,31 +103,20 @@ func (s *TelnetServer) Handler(ctx context.Context, conn net.Conn, m *movie.Movi
 	}
 	defer serverInfo.StreamDisconnect(id)
 
-	inR, inW := io.Pipe()
-	defer func() {
-		_ = inR.Close()
-	}()
-
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	termCh := make(chan string, 1)
-	defer close(termCh)
-	sizeCh := make(chan telnet.WindowSize, 1)
-	defer close(sizeCh)
+	in, profile, sizeCh, errCh := telnet.Proxy(conn)
+	defer func() {
+		_ = in.Close()
+	}()
 	go func() {
-		// Proxy input to program
-		_ = telnet.Proxy(conn, inW, termCh, sizeCh)
+		<-errCh
 		cancel()
 	}()
 
-	var gotProfile bool
-	var profile termenv.Profile
-	select {
-	case term := <-termCh:
-		profile = util.Profile(term)
-		gotProfile = true
-	case <-time.After(time.Second):
+	gotProfile := profile != -1
+	if !gotProfile {
 		profile = termenv.ANSI256
 	}
 
@@ -138,7 +125,7 @@ func (s *TelnetServer) Handler(ctx context.Context, conn net.Conn, m *movie.Movi
 
 	opts := []tea.ProgramOption{
 		tea.WithContext(ctx),
-		tea.WithInput(inR),
+		tea.WithInput(in),
 		tea.WithOutput(conn),
 		tea.WithFPS(30),
 	}
@@ -148,17 +135,12 @@ func (s *TelnetServer) Handler(ctx context.Context, conn net.Conn, m *movie.Movi
 	program := tea.NewProgram(p, opts...)
 
 	go func() {
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			case info := <-sizeCh:
-				if info.Width != 0 && info.Height != 0 {
-					program.Send(tea.WindowSizeMsg{
-						Width:  int(info.Width),
-						Height: int(info.Height),
-					})
-				}
+		for info := range sizeCh {
+			if info.Width != 0 && info.Height != 0 {
+				program.Send(tea.WindowSizeMsg{
+					Width:  int(info.Width),
+					Height: int(info.Height),
+				})
 			}
 		}
 	}()

internal/server/telnet/proxy.go

@@ -6,16 +6,46 @@ import (
 	"encoding/binary"
 	"io"
 	"net"
+	"time"
 
+	"github.com/gabe565/ascii-movie/internal/util"
+	"github.com/muesli/termenv"
 	"github.com/rs/zerolog/log"
 )
 
 type WindowSize struct {
 	Width, Height uint16
 }
 
+func Proxy(conn net.Conn) (io.ReadCloser, termenv.Profile, <-chan WindowSize, <-chan error) {
+	pr, pw := io.Pipe()
+	termCh := make(chan string, 1)
+	sizeCh := make(chan WindowSize, 1)
+	errCh := make(chan error)
+
+	go func() {
+		defer func() {
+			_ = pw.Close()
+			close(termCh)
+			close(sizeCh)
+			close(errCh)
+		}()
+
+		errCh <- proxy(conn, pw, termCh, sizeCh)
+	}()
+
+	profile := termenv.Profile(-1)
+	select {
+	case term := <-termCh:
+		profile = util.Profile(term)
+	case <-time.After(time.Second):
+	}
+
+	return pr, profile, sizeCh, errCh
+}
+
 //nolint:gocyclo
-func Proxy(conn net.Conn, proxy io.Writer, termCh chan string, sizeCh chan WindowSize) error {
+func proxy(conn net.Conn, proxy io.Writer, termCh chan<- string, sizeCh chan<- WindowSize) error {
 	reader := bufio.NewReaderSize(conn, 64)
 	var wroteTelnetCommands bool
 	var wroteTermType bool