diff --git a/src/app_test.go b/src/app_test.go index 8089681..88b88e5 100644 --- a/src/app_test.go +++ b/src/app_test.go @@ -2,7 +2,6 @@ package main import ( "bytes" - "fmt" "image" "image/jpeg" "os" @@ -61,8 +60,9 @@ func TestTSFInvalid(t *testing.T) { } // Valid container but fps <= 0. - if err := writeTSF(path, &FramesContainer{ColorFrames: [][]byte{{1}}, FPS: 0}); err != nil { - t.Fatalf("writeTSF: %v", err) + rawInvalidFPS := append(tsfHeader(0, 1), 1, 0, 0, 0, 1) + if err := os.WriteFile(path, rawInvalidFPS, 0o644); err != nil { + t.Fatalf("WriteFile: %v", err) } if _, err := loadTSF(path); err == nil { t.Error("expected error for fps<=0") @@ -177,7 +177,7 @@ func TestFrameToAscii(t *testing.T) { Pix: []byte{0, 0, 0, 255, 255, 255, 255, 255}, Stride: 8, Rect: image.Rect(0, 0, 2, 1), } - out := []rune(frameToAscii(img, buildRampLUT(ramp, opts))) + out := []rune(string(frameToAscii(img, buildRampLUT(ramp, opts)))) if out[0] != ramp[0] { t.Errorf("dark px = %q, want %q", out[0], ramp[0]) } @@ -193,7 +193,7 @@ func TestFrameToAsciiInvert(t *testing.T) { Pix: []byte{255, 255, 255, 255}, Stride: 4, Rect: image.Rect(0, 0, 1, 1), } - out := []rune(frameToAscii(img, buildRampLUT(ramp, opts))) + out := []rune(string(frameToAscii(img, buildRampLUT(ramp, opts)))) if out[0] != ramp[0] { t.Errorf("inverted bright = %q, want %q", out[0], ramp[0]) } @@ -215,7 +215,7 @@ func TestRenderConcurrentSameKey(t *testing.T) { }, newRenderCache(1<<20)) var wg sync.WaitGroup - results := make([]string, 32) + results := make([][]byte, 32) for i := range results { wg.Add(1) go func(i int) { @@ -230,25 +230,26 @@ func TestRenderConcurrentSameKey(t *testing.T) { } wg.Wait() for i, got := range results { - if got != results[0] { + if !bytes.Equal(got, results[0]) { t.Fatalf("result %d differs from result 0", i) } } } func TestRenderCacheEvictsByBytes(t *testing.T) { - budget := 3 * entryCost("k", strings.Repeat("x", 1000)) + key := func(index int) cacheKey { return cacheKey{index: index} } + budget := 3 * entryCost(key(0), bytes.Repeat([]byte("x"), 1000)) c := newRenderCache(budget) for i := range 5 { - c.put(fmt.Sprintf("%d", i), strings.Repeat("x", 1000)) + c.put(key(i), bytes.Repeat([]byte("x"), 1000)) } - if c.size > budget { - t.Errorf("size %d exceeds budget %d", c.size, budget) + if c.size.Load() > budget { + t.Errorf("size %d exceeds budget %d", c.size.Load(), budget) } - if _, ok := c.get("0"); ok { + if _, ok := c.get(key(0)); ok { t.Error("oldest entry should have been evicted") } - if _, ok := c.get("4"); !ok { + if _, ok := c.get(key(4)); !ok { t.Error("newest entry should be cached") } } @@ -258,49 +259,55 @@ func TestRenderCacheDisabled(t *testing.T) { if c != nil { t.Fatal("zero budget should disable the cache") } - c.put("k", "v") - if _, ok := c.get("k"); ok { + c.put(cacheKey{}, []byte("v")) + if _, ok := c.get(cacheKey{}); ok { t.Error("nil cache should never hit") } } func TestRenderCacheRejectsOversizedEntry(t *testing.T) { c := newRenderCache(256) - c.put("big", strings.Repeat("x", 10_000)) - if _, ok := c.get("big"); ok { + c.put(cacheKey{}, bytes.Repeat([]byte("x"), 10_000)) + if _, ok := c.get(cacheKey{}); ok { t.Error("entry larger than budget should not be cached") } - if c.size != 0 { - t.Errorf("size = %d, want 0", c.size) + if c.size.Load() != 0 { + t.Errorf("size = %d, want 0", c.size.Load()) } } func TestConnectionTracker(t *testing.T) { tr := newConnectionTracker() - tr.increment("1.2.3.4") - tr.increment("1.2.3.4") - if !tr.hasReachedLimits("1.2.3.4", 2, 100) { - t.Error("expected per-ip limit reached") + if _, _, ok := tr.tryAcquire("1.2.3.4", 2, 100); !ok { + t.Fatal("first acquire failed") } - tr.decrement("1.2.3.4") - tr.decrement("1.2.3.4") + if _, _, ok := tr.tryAcquire("1.2.3.4", 2, 100); !ok { + t.Fatal("second acquire failed") + } + if _, _, ok := tr.tryAcquire("1.2.3.4", 2, 100); ok { + t.Error("expected per-ip limit rejection") + } + tr.release("1.2.3.4") + tr.release("1.2.3.4") if tr.totalCount() != 0 { t.Errorf("total = %d", tr.totalCount()) } - if tr.hasReachedLimits("1.2.3.4", 2, 100) { - t.Error("should be cleared") + if _, _, ok := tr.tryAcquire("1.2.3.4", 2, 100); !ok { + t.Error("limit should be cleared") } } -func TestClampDimension(t *testing.T) { - if clampDimension(0, 100) != 1 { - t.Error("floor") +func TestClampTermSize(t *testing.T) { + w, h := clampTermSize(1000, 500, 512, 65536, 4) + if w < 1 || h < 1 || w > 512 || h > 512 || w*h > 65536 { + t.Fatalf("clamped size = %dx%d", w, h) } - if clampDimension(500, 100) != 100 { - t.Error("ceil") + if w%4 != 0 || h%4 != 0 { + t.Fatalf("size is not quantized: %dx%d", w, h) } - if clampDimension(50, 100) != 50 { - t.Error("passthrough") + w, h = clampTermSize(3, 2, 100, 100, 4) + if w != 3 || h != 2 { + t.Fatalf("small size = %dx%d", w, h) } } diff --git a/src/config.go b/src/config.go index 072c6b4..6283e5b 100644 --- a/src/config.go +++ b/src/config.go @@ -28,6 +28,8 @@ type Config struct { MaxAuthAttempts int HandshakeTimeout time.Duration MaxDimension int + MaxTerminalCells int + SessionTimeout time.Duration RenderCacheMB int BrightnessThreshold int Charset string @@ -126,6 +128,8 @@ func loadConfig() Config { MaxAuthAttempts: envInt("MAX_AUTH_ATTEMPTS", 6, 1, maxInt), HandshakeTimeout: envDurationMs("HANDSHAKE_TIMEOUT", 10*time.Second), MaxDimension: envInt("MAX_DIMENSION", 512, 1, 4096), + MaxTerminalCells: envInt("MAX_TERMINAL_CELLS", 500*512, 1, maxInt), + SessionTimeout: envDurationMs("SESSION_TIMEOUT", 10*time.Minute), RenderCacheMB: envInt("RENDER_CACHE_MB", 256, 0, maxInt), BrightnessThreshold: envInt("BRIGHTNESS_THRESHOLD", 40, 0, 100), Charset: envString("CHARSET", "detailed"), diff --git a/src/main.go b/src/main.go index 5c2c321..287ae19 100644 --- a/src/main.go +++ b/src/main.go @@ -83,6 +83,8 @@ func generateFrames(framesDir, videoArg string, resolution int) { } } +const frameDataWarnBytes = 2 << 30 + func loadAllFrames(framesDir string) []*FramesContainer { entries, err := os.ReadDir(framesDir) var files []string @@ -101,6 +103,22 @@ func loadAllFrames(framesDir string) []*FramesContainer { framesDir, )) } + var totalBytes int64 + for _, file := range files { + info, err := os.Stat(filepath.Join(framesDir, file)) + if err != nil { + fail(err.Error()) + } + totalBytes += info.Size() + } + if totalBytes > frameDataWarnBytes { + logWarn(fmt.Sprintf( + "Frame data is %.1f MB of mapped memory; make sure the container memory limit leaves headroom", + float64(totalBytes)/(1<<20), + )) + } else { + logInfo(fmt.Sprintf("Frame data: %.1f MB", float64(totalBytes)/(1<<20))) + } concurrency := min(len(files), max(1, min(runtime.NumCPU(), 4))) @@ -226,14 +244,16 @@ func main() { GoodbyeText: goodbyeText, VideoSets: videoSets, }) + defer server.Close() sigCh := make(chan os.Signal, 1) signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) go func() { sig := <-sigCh logInfo(fmt.Sprintf("Received %s, shutting down...", sig)) + forceExit := time.AfterFunc(5*time.Second, func() { os.Exit(0) }) server.Close() - time.AfterFunc(5*time.Second, func() { os.Exit(0) }) + forceExit.Stop() }() if err := server.Listen(config.Host, config.Port); err != nil { diff --git a/src/performance_safeguards_test.go b/src/performance_safeguards_test.go new file mode 100644 index 0000000..c88586d --- /dev/null +++ b/src/performance_safeguards_test.go @@ -0,0 +1,127 @@ +package main + +import ( + "bytes" + "image" + "sync" + "testing" + + "golang.org/x/crypto/ssh" +) + +func TestConnectionTrackerConcurrentLimit(t *testing.T) { + tracker := newConnectionTracker() + start := make(chan struct{}) + var wg sync.WaitGroup + var mu sync.Mutex + accepted := make(map[string]int) + for i := range 100 { + wg.Add(1) + go func(i int) { + defer wg.Done() + <-start + ip := string(rune('a' + i%10)) + if _, _, ok := tracker.tryAcquire(ip, 3, 7); ok { + mu.Lock() + accepted[ip]++ + mu.Unlock() + } + }(i) + } + close(start) + wg.Wait() + + total := 0 + for ip, count := range accepted { + total += count + if count > 3 { + t.Fatalf("IP %q acquired %d slots", ip, count) + } + } + if total != 7 || tracker.totalCount() != 7 { + t.Fatalf("accepted=%d tracked=%d, want 7", total, tracker.totalCount()) + } + for ip, count := range accepted { + for range count { + tracker.release(ip) + } + } +} + +func TestSessionTrackerLimits(t *testing.T) { + tracker := newSessionTracker() + first := &ssh.ServerConn{} + second := &ssh.ServerConn{} + if !tracker.tryAcquire(first, 1, 2) { + t.Fatal("first session rejected") + } + if tracker.tryAcquire(first, 1, 2) { + t.Fatal("per-connection limit was not enforced") + } + if !tracker.tryAcquire(second, 1, 2) { + t.Fatal("second connection session rejected") + } + if tracker.tryAcquire(&ssh.ServerConn{}, 1, 2) { + t.Fatal("global session limit was not enforced") + } + tracker.release(first) + if !tracker.tryAcquire(&ssh.ServerConn{}, 1, 2) { + t.Fatal("released slot was not reusable") + } +} + +func TestTermSizeDebouncesResize(t *testing.T) { + size := &termSize{} + size.set(80, 24, 512, 500*512, true) + size.set(200, 100, 512, 500*512, false) + if w, h := size.get(); w != 80 || h != 24 { + t.Fatalf("debounced size = %dx%d", w, h) + } + size.set(200, 100, 512, 500*512, true) + if w, h := size.get(); w != 200 || h != 100 { + t.Fatalf("forced size = %dx%d", w, h) + } +} + +func TestAnsi256CoalescesQuantizedColors(t *testing.T) { + img := &image.RGBA{ + Pix: []byte{96, 96, 96, 255, 100, 100, 100, 255}, + Stride: 8, + Rect: image.Rect(0, 0, 2, 1), + } + output := frameToAnsi(img, buildRampLUT([]rune(" .#"), asciiOptions{}), colorTier256) + if count := bytes.Count(output, []byte("\x1b[38;5;")); count != 1 { + t.Fatalf("color escape count = %d, want 1: %q", count, output) + } +} + +func TestAnsiDoesNotResetEachRow(t *testing.T) { + img := &image.RGBA{ + Pix: []byte{100, 100, 100, 255, 100, 100, 100, 255}, + Stride: 4, + Rect: image.Rect(0, 0, 1, 2), + } + output := frameToAnsi(img, buildRampLUT([]rune(" .#"), asciiOptions{}), colorTierTrueColor) + if count := bytes.Count(output, []byte(ansiReset)); count != 1 { + t.Fatalf("reset count = %d, want 1: %q", count, output) + } +} + +func TestRenderCacheAccountsRetainedCapacity(t *testing.T) { + cache := newRenderCache(512) + value := make([]byte, 1, 4096) + cache.put(cacheKey{}, value) + if _, ok := cache.get(cacheKey{}); ok { + t.Fatal("cache accepted an entry whose backing allocation exceeds its budget") + } + if cache.stats().Rejections != 1 { + t.Fatalf("rejections = %d, want 1", cache.stats().Rejections) + } +} + +func TestSanitizeNStopsAtLimit(t *testing.T) { + input := "ab\x00cdefghijklmnopqrstuvwxyz" + if got := sanitizeN(input, 4); got != "ab�c…" { + t.Fatalf("sanitizeN = %q", got) + } +} diff --git a/src/server.go b/src/server.go index f8afa9c..c3ff3ed 100644 --- a/src/server.go +++ b/src/server.go @@ -4,6 +4,8 @@ import ( "encoding/binary" "errors" "fmt" + "io" + "math" "math/rand" "net" "strings" @@ -14,14 +16,84 @@ import ( ) const ( - clearScreen = "\x1b[2J\x1b[0f" - hideCursor = "\x1b[?25l" - showCursor = "\x1b[?25h" - syncStart = "\x1b[?2026h" - syncEnd = "\x1b[?2026l" - homeCursor = "\x1b[H" + clearScreen = "\x1b[2J\x1b[0f" + hideCursor = "\x1b[?25l" + showCursor = "\x1b[?25h" + syncStart = "\x1b[?2026h" + syncEnd = "\x1b[?2026l" + homeCursor = "\x1b[H" + maxSessionsPerConn = 1 + terminalSizeQuantum = 4 + resizeDebounce = 200 * time.Millisecond + outputStallTimeout = 15 * time.Second ) +var errOutputStalled = errors.New("SSH output stalled") + +func writePartsWithTimeout( + conn *ssh.ServerConn, + channel ssh.Channel, + timeout time.Duration, + parts ...string, +) error { + write := func() error { + for _, part := range parts { + if _, err := io.WriteString(channel, part); err != nil { + return err + } + } + return nil + } + if timeout <= 0 { + return write() + } + + fired := make(chan struct{}) + timer := time.AfterFunc(timeout, func() { + _ = conn.Close() + close(fired) + }) + err := write() + if timer.Stop() { + return err + } + <-fired + return errOutputStalled +} + +func writeFrameWithTimeout( + conn *ssh.ServerConn, + channel ssh.Channel, + timeout time.Duration, + prefix string, + frame []byte, +) error { + write := func() error { + if _, err := io.WriteString(channel, syncStart+prefix); err != nil { + return err + } + if _, err := channel.Write(frame); err != nil { + return err + } + _, err := io.WriteString(channel, syncEnd) + return err + } + if timeout <= 0 { + return write() + } + fired := make(chan struct{}) + timer := time.AfterFunc(timeout, func() { + _ = conn.Close() + close(fired) + }) + err := write() + if timer.Stop() { + return err + } + <-fired + return errOutputStalled +} + type ConnectionTracker struct { mu sync.Mutex counts map[string]int @@ -32,15 +104,18 @@ func newConnectionTracker() *ConnectionTracker { return &ConnectionTracker{counts: make(map[string]int)} } -func (t *ConnectionTracker) increment(ip string) int { +func (t *ConnectionTracker) tryAcquire(ip string, maxPerIP, maxTotal int) (int, int, bool) { t.mu.Lock() defer t.mu.Unlock() + if t.total >= maxTotal || t.counts[ip] >= maxPerIP { + return t.counts[ip], t.total, false + } t.counts[ip]++ t.total++ - return t.counts[ip] + return t.counts[ip], t.total, true } -func (t *ConnectionTracker) decrement(ip string) { +func (t *ConnectionTracker) release(ip string) { t.mu.Lock() defer t.mu.Unlock() if _, ok := t.counts[ip]; !ok { @@ -61,10 +136,40 @@ func (t *ConnectionTracker) totalCount() int { return t.total } -func (t *ConnectionTracker) hasReachedLimits(ip string, maxPerIP, maxTotal int) bool { +type SessionTracker struct { + mu sync.Mutex + perConn map[*ssh.ServerConn]int + total int +} + +func newSessionTracker() *SessionTracker { + return &SessionTracker{perConn: make(map[*ssh.ServerConn]int)} +} + +func (t *SessionTracker) tryAcquire(conn *ssh.ServerConn, maxPerConn, maxTotal int) bool { t.mu.Lock() defer t.mu.Unlock() - return t.total >= maxTotal || t.counts[ip] >= maxPerIP + if t.total >= maxTotal || t.perConn[conn] >= maxPerConn { + return false + } + t.perConn[conn]++ + t.total++ + return true +} + +func (t *SessionTracker) release(conn *ssh.ServerConn) { + t.mu.Lock() + defer t.mu.Unlock() + count := t.perConn[conn] + if count <= 0 { + return + } + if count == 1 { + delete(t.perConn, conn) + } else { + t.perConn[conn] = count - 1 + } + t.total-- } type frameSet struct { @@ -76,10 +181,16 @@ type Server struct { config Config sshConfig *ssh.ServerConfig sets []frameSet + cache *renderCache tracker *ConnectionTracker + sessions *SessionTracker fakeLogin *string goodbye *string + mu sync.Mutex listener net.Listener + conns map[net.Conn]struct{} + connWG sync.WaitGroup + closing bool closeOnce sync.Once } @@ -92,14 +203,25 @@ type ServerDeps struct { VideoSets []*FramesContainer } -func clampDimension(value, max int) int { - if value < 1 { - return 1 +func clampTermSize(cols, rows, maxDimension, maxCells, quantum int) (int, int) { + cols = max(cols, 1) + rows = max(rows, 1) + scale := min(1.0, float64(maxDimension)/float64(cols), float64(maxDimension)/float64(rows)) + area := float64(cols) * float64(rows) + if area*scale*scale > float64(maxCells) { + scale = min(scale, math.Sqrt(float64(maxCells)/area)) } - if value > max { - return max + cols = max(1, int(math.Floor(float64(cols)*scale))) + rows = max(1, int(math.Floor(float64(rows)*scale))) + if quantum > 1 { + if cols >= quantum { + cols -= cols % quantum + } + if rows >= quantum { + rows -= rows % quantum + } } - return value + return cols, rows } func createServer(deps ServerDeps) *Server { @@ -155,9 +277,12 @@ func createServer(deps ServerDeps) *Server { config: config, sshConfig: sshConfig, sets: sets, + cache: cache, tracker: newConnectionTracker(), + sessions: newSessionTracker(), fakeLogin: deps.FakeLoginText, goodbye: deps.GoodbyeText, + conns: make(map[net.Conn]struct{}), } } @@ -174,7 +299,14 @@ func (s *Server) Listen(host string, port int) error { if err != nil { return err } + s.mu.Lock() + if s.closing { + s.mu.Unlock() + _ = listener.Close() + return nil + } s.listener = listener + s.mu.Unlock() logInfo(fmt.Sprintf("TrollSSH listening on %s:%d", host, port)) for { conn, err := listener.Accept() @@ -184,29 +316,68 @@ func (s *Server) Listen(host string, port int) error { } return err } - go s.handleConn(conn) + ip := hostOnly(conn.RemoteAddr().String()) + activeForIP, total, ok := s.tracker.tryAcquire(ip, s.config.MaxConnections, s.config.MaxTotalConnections) + if !ok { + _ = conn.Close() + logWarn("Connection rejected (limit reached) from", ip) + continue + } + s.mu.Lock() + if s.closing { + s.mu.Unlock() + s.tracker.release(ip) + _ = conn.Close() + continue + } + s.conns[conn] = struct{}{} + s.connWG.Add(1) + s.mu.Unlock() + go s.handleConn(conn, ip, activeForIP, total) } } func (s *Server) Close() { s.closeOnce.Do(func() { - if s.listener != nil { - _ = s.listener.Close() + s.mu.Lock() + s.closing = true + listener := s.listener + conns := make([]net.Conn, 0, len(s.conns)) + for conn := range s.conns { + conns = append(conns, conn) + } + s.mu.Unlock() + if listener != nil { + _ = listener.Close() + } + for _, conn := range conns { + _ = conn.Close() + } + s.connWG.Wait() + stats := s.cache.stats() + if stats.Hits+stats.Misses > 0 { + logInfo(fmt.Sprintf( + "Render cache: size=%.1fMB hits=%d misses=%d evictions=%d rejected=%d renders=%d render_time=%s", + float64(stats.SizeBytes)/(1<<20), stats.Hits, stats.Misses, stats.Evictions, + stats.Rejections, stats.Renders, stats.RenderTime, + )) + } + for _, set := range s.sets { + if err := set.data.Close(); err != nil { + logWarn("Failed to release frame set", set.data.Name, sanitize(err.Error())) + } } }) } -func (s *Server) handleConn(conn net.Conn) { - ip := hostOnly(conn.RemoteAddr().String()) - - if s.tracker.hasReachedLimits(ip, s.config.MaxConnections, s.config.MaxTotalConnections) { - _ = conn.Close() - logWarn("Connection rejected (limit reached) from", ip) - return - } - - activeForIP := s.tracker.increment(ip) - defer s.tracker.decrement(ip) +func (s *Server) handleConn(conn net.Conn, ip string, activeForIP, total int) { + defer func() { + s.tracker.release(ip) + s.mu.Lock() + delete(s.conns, conn) + s.mu.Unlock() + s.connWG.Done() + }() if s.config.HandshakeTimeout > 0 { _ = conn.SetDeadline(time.Now().Add(s.config.HandshakeTimeout)) @@ -229,35 +400,58 @@ func (s *Server) handleConn(conn net.Conn) { setIndex := rand.Intn(len(s.sets)) logInfo(fmt.Sprintf( "New connection from %s (ip=%d, total=%d) -> playing %q", - ip, activeForIP, s.tracker.totalCount(), s.sets[setIndex].data.Name, + ip, activeForIP, total, s.sets[setIndex].data.Name, )) go ssh.DiscardRequests(reqs) + var sessionWG sync.WaitGroup for newChannel := range chans { if newChannel.ChannelType() != "session" { _ = newChannel.Reject(ssh.UnknownChannelType, "unknown channel type") continue } - channel, requests, err := newChannel.Accept() - if err != nil { + if !s.sessions.tryAcquire(sshConn, maxSessionsPerConn, s.config.MaxTotalConnections) { + _ = newChannel.Reject(ssh.ResourceShortage, "session limit reached") continue } - go s.handleSession(sshConn, channel, requests, ip, setIndex) + channel, requests, err := newChannel.Accept() + if err != nil { + s.sessions.release(sshConn) + continue + } + sessionWG.Add(1) + go func() { + defer sessionWG.Done() + defer s.sessions.release(sshConn) + var timer *time.Timer + if s.config.SessionTimeout > 0 { + timer = time.AfterFunc(s.config.SessionTimeout, func() { _ = sshConn.Close() }) + defer timer.Stop() + } + s.handleSession(sshConn, channel, requests, ip, setIndex) + }() } + _ = sshConn.Close() + sessionWG.Wait() logInfo("Client closed connection from", ip) } type termSize struct { - mu sync.Mutex - width int - height int + mu sync.Mutex + width int + height int + updated time.Time } -func (t *termSize) set(w, h, maxDim int) { +func (t *termSize) set(w, h, maxDimension, maxCells int, force bool) { t.mu.Lock() - t.width = clampDimension(w, maxDim) - t.height = clampDimension(h, maxDim) + if !force && time.Since(t.updated) < resizeDebounce { + t.mu.Unlock() + return + } + t.width, t.height = clampTermSize(w, h, maxDimension, maxCells, terminalSizeQuantum) + t.updated = time.Now() t.mu.Unlock() } @@ -304,20 +498,22 @@ func (s *Server) handleSession( ip string, initialSetIndex int, ) { + defer func() { _ = channel.Close() }() size := &termSize{} - size.set(80, 24, s.config.MaxDimension) + size.set(80, 24, s.config.MaxDimension, s.config.MaxTerminalCells, true) tier := colorTierTrueColor if s.config.ForceGrayscale { tier = colorTierNone } started := false + var playDone chan struct{} for req := range requests { switch req.Type { case "pty-req": logDebug("Opening pty for session", ip) if cols, rows, ok := parseDims(req.Payload); ok { - size.set(cols, rows, s.config.MaxDimension) + size.set(cols, rows, s.config.MaxDimension, s.config.MaxTerminalCells, true) } if term, ok := parsePtyTerm(req.Payload); ok { tier = detectColorTier(term) @@ -331,7 +527,7 @@ func (s *Server) handleSession( if len(req.Payload) >= 8 { cols := int(binary.BigEndian.Uint32(req.Payload)) rows := int(binary.BigEndian.Uint32(req.Payload[4:])) - size.set(cols, rows, s.config.MaxDimension) + size.set(cols, rows, s.config.MaxDimension, s.config.MaxTerminalCells, false) } if req.WantReply { _ = req.Reply(true, nil) @@ -348,14 +544,24 @@ func (s *Server) handleSession( _ = req.Reply(true, nil) if !started { started = true - go s.playVideo(sshConn, channel, size, ip, initialSetIndex, false, tier) + playDone = make(chan struct{}) + playTier := tier + go func(tier colorTier) { + defer close(playDone) + s.playVideo(sshConn, channel, size, ip, initialSetIndex, false, tier) + }(playTier) } case "shell": logDebug("Opening shell for session", ip) _ = req.Reply(true, nil) if !started { started = true - go s.playVideo(sshConn, channel, size, ip, initialSetIndex, false, tier) + playDone = make(chan struct{}) + playTier := tier + go func(tier colorTier) { + defer close(playDone) + s.playVideo(sshConn, channel, size, ip, initialSetIndex, false, tier) + }(playTier) } default: if req.WantReply { @@ -363,6 +569,10 @@ func (s *Server) handleSession( } } } + _ = channel.Close() + if playDone != nil { + <-playDone + } } func (s *Server) pickNextSetIndex(exclude int) int { @@ -391,11 +601,16 @@ func (s *Server) playVideo( w, h := size.get() logDebug(fmt.Sprintf("Terminal size %dx%d for %s", w, h, ip)) - defer func() { _, _ = channel.Write([]byte(showCursor)) }() + defer func() { + _ = writePartsWithTimeout(sshConn, channel, outputStallTimeout, showCursor) + }() if s.fakeLogin != nil { - _, _ = channel.Write([]byte(clearScreen)) - _, _ = channel.Write([]byte(*s.fakeLogin)) + if err := writePartsWithTimeout( + sshConn, channel, outputStallTimeout, clearScreen, *s.fakeLogin, + ); err != nil { + return + } } done := make(chan struct{}) @@ -439,13 +654,19 @@ func (s *Server) playVideo( } }() + loginTimer := time.NewTimer(config.LoginDelay) select { - case <-time.After(config.LoginDelay): + case <-loginTimer.C: case <-done: + if !loginTimer.Stop() { + <-loginTimer.C + } return } - _, _ = channel.Write([]byte(hideCursor)) + if err := writePartsWithTimeout(sshConn, channel, outputStallTimeout, hideCursor); err != nil { + return + } frameInterval := func() time.Duration { return time.Duration(float64(time.Second) / current.data.FPS) @@ -457,8 +678,6 @@ func (s *Server) playVideo( currentFrame := 0 loopCount := 0 lastW, lastH := 0, 0 - var writeBuf []byte - for { select { case <-done: @@ -489,11 +708,9 @@ func (s *Server) playVideo( prefix = clearScreen lastW, lastH = w, h } - writeBuf = append(writeBuf[:0], syncStart...) - writeBuf = append(writeBuf, prefix...) - writeBuf = append(writeBuf, ascii...) - writeBuf = append(writeBuf, syncEnd...) - if _, err := channel.Write(writeBuf); err != nil { + if err := writeFrameWithTimeout( + sshConn, channel, outputStallTimeout, prefix, ascii, + ); err != nil { closeSession() return } @@ -506,11 +723,27 @@ func (s *Server) playVideo( currentFrame = 0 loopCount++ if config.MaxLoop > 0 && loopCount >= config.MaxLoop { - _, _ = channel.Write([]byte(showCursor + clearScreen)) + if err := writePartsWithTimeout( + sshConn, channel, outputStallTimeout, showCursor, clearScreen, + ); err != nil { + return + } if s.goodbye != nil { - _, _ = channel.Write([]byte(*s.goodbye)) + if err := writePartsWithTimeout( + sshConn, channel, outputStallTimeout, *s.goodbye, + ); err != nil { + return + } + } + closeTimer := time.NewTimer(time.Second) + select { + case <-closeTimer.C: + case <-done: + if !closeTimer.Stop() { + <-closeTimer.C + } + return } - time.Sleep(1 * time.Second) logInfo("Playback finished, closing session", ip) _ = channel.Close() _ = sshConn.Close()