6 Commits
16 changed files with 1244 additions and 243 deletions
+6
View File
@@ -26,12 +26,16 @@ FORCE_GRAYSCALE=false
# Max rendered width/height in characters. # Max rendered width/height in characters.
MAX_DIMENSION=512 MAX_DIMENSION=512
# Max rendered area (columns x rows); larger terminals are scaled down.
MAX_TERMINAL_CELLS=256000
# Memory budget in MB for the rendered-frame cache (0 disables caching). # Memory budget in MB for the rendered-frame cache (0 disables caching).
RENDER_CACHE_MB=256 RENDER_CACHE_MB=256
# Go soft memory limit; set below your container limit to avoid OOM kills. # Go soft memory limit; set below your container limit to avoid OOM kills.
GOMEMLIMIT=1GiB GOMEMLIMIT=1GiB
# Docker Compose container memory limit. Leave headroom for mapped frame files.
MEMORY_LIMIT=2g
# Connection limits. New connections over a limit are dropped immediately. # Connection limits. New connections over a limit are dropped immediately.
# Max simultaneous connections from a single client IP. # Max simultaneous connections from a single client IP.
@@ -43,6 +47,8 @@ MAX_TOTAL_CONNECTIONS=1000
MAX_AUTH_ATTEMPTS=3 MAX_AUTH_ATTEMPTS=3
# SSH handshake deadline in ms (0 to disable). # SSH handshake deadline in ms (0 to disable).
HANDSHAKE_TIMEOUT=30000 HANDSHAKE_TIMEOUT=30000
# Maximum session lifetime in ms (0 to disable).
SESSION_TIMEOUT=600000
# Log attempted usernames/passwords. # Log attempted usernames/passwords.
LOG_CREDENTIALS=true LOG_CREDENTIALS=true
+1
View File
@@ -3,6 +3,7 @@ services:
image: ghcr.io/yuzuzensai/trollssh:v1.0.1 image: ghcr.io/yuzuzensai/trollssh:v1.0.1
container_name: trollssh container_name: trollssh
restart: unless-stopped restart: unless-stopped
mem_limit: ${MEMORY_LIMIT:-2g}
ports: ports:
- "22:22" - "22:22"
env_file: env_file:
+41 -34
View File
@@ -2,7 +2,6 @@ package main
import ( import (
"bytes" "bytes"
"fmt"
"image" "image"
"image/jpeg" "image/jpeg"
"os" "os"
@@ -61,8 +60,9 @@ func TestTSFInvalid(t *testing.T) {
} }
// Valid container but fps <= 0. // Valid container but fps <= 0.
if err := writeTSF(path, &FramesContainer{ColorFrames: [][]byte{{1}}, FPS: 0}); err != nil { rawInvalidFPS := append(tsfHeader(0, 1), 1, 0, 0, 0, 1)
t.Fatalf("writeTSF: %v", err) if err := os.WriteFile(path, rawInvalidFPS, 0o644); err != nil {
t.Fatalf("WriteFile: %v", err)
} }
if _, err := loadTSF(path); err == nil { if _, err := loadTSF(path); err == nil {
t.Error("expected error for fps<=0") 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}, Pix: []byte{0, 0, 0, 255, 255, 255, 255, 255},
Stride: 8, Rect: image.Rect(0, 0, 2, 1), 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] { if out[0] != ramp[0] {
t.Errorf("dark px = %q, want %q", 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}, Pix: []byte{255, 255, 255, 255},
Stride: 4, Rect: image.Rect(0, 0, 1, 1), 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] { if out[0] != ramp[0] {
t.Errorf("inverted bright = %q, want %q", 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)) }, newRenderCache(1<<20))
var wg sync.WaitGroup var wg sync.WaitGroup
results := make([]string, 32) results := make([][]byte, 32)
for i := range results { for i := range results {
wg.Add(1) wg.Add(1)
go func(i int) { go func(i int) {
@@ -230,25 +230,26 @@ func TestRenderConcurrentSameKey(t *testing.T) {
} }
wg.Wait() wg.Wait()
for i, got := range results { for i, got := range results {
if got != results[0] { if !bytes.Equal(got, results[0]) {
t.Fatalf("result %d differs from result 0", i) t.Fatalf("result %d differs from result 0", i)
} }
} }
} }
func TestRenderCacheEvictsByBytes(t *testing.T) { 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) c := newRenderCache(budget)
for i := range 5 { 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 { if c.size.Load() > budget {
t.Errorf("size %d exceeds budget %d", c.size, 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") 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") t.Error("newest entry should be cached")
} }
} }
@@ -258,49 +259,55 @@ func TestRenderCacheDisabled(t *testing.T) {
if c != nil { if c != nil {
t.Fatal("zero budget should disable the cache") t.Fatal("zero budget should disable the cache")
} }
c.put("k", "v") c.put(cacheKey{}, []byte("v"))
if _, ok := c.get("k"); ok { if _, ok := c.get(cacheKey{}); ok {
t.Error("nil cache should never hit") t.Error("nil cache should never hit")
} }
} }
func TestRenderCacheRejectsOversizedEntry(t *testing.T) { func TestRenderCacheRejectsOversizedEntry(t *testing.T) {
c := newRenderCache(256) c := newRenderCache(256)
c.put("big", strings.Repeat("x", 10_000)) c.put(cacheKey{}, bytes.Repeat([]byte("x"), 10_000))
if _, ok := c.get("big"); ok { if _, ok := c.get(cacheKey{}); ok {
t.Error("entry larger than budget should not be cached") t.Error("entry larger than budget should not be cached")
} }
if c.size != 0 { if c.size.Load() != 0 {
t.Errorf("size = %d, want 0", c.size) t.Errorf("size = %d, want 0", c.size.Load())
} }
} }
func TestConnectionTracker(t *testing.T) { func TestConnectionTracker(t *testing.T) {
tr := newConnectionTracker() tr := newConnectionTracker()
tr.increment("1.2.3.4") if _, _, ok := tr.tryAcquire("1.2.3.4", 2, 100); !ok {
tr.increment("1.2.3.4") t.Fatal("first acquire failed")
if !tr.hasReachedLimits("1.2.3.4", 2, 100) {
t.Error("expected per-ip limit reached")
} }
tr.decrement("1.2.3.4") if _, _, ok := tr.tryAcquire("1.2.3.4", 2, 100); !ok {
tr.decrement("1.2.3.4") 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 { if tr.totalCount() != 0 {
t.Errorf("total = %d", tr.totalCount()) t.Errorf("total = %d", tr.totalCount())
} }
if tr.hasReachedLimits("1.2.3.4", 2, 100) { if _, _, ok := tr.tryAcquire("1.2.3.4", 2, 100); !ok {
t.Error("should be cleared") t.Error("limit should be cleared")
} }
} }
func TestClampDimension(t *testing.T) { func TestClampTermSize(t *testing.T) {
if clampDimension(0, 100) != 1 { w, h := clampTermSize(1000, 500, 512, 65536, 4)
t.Error("floor") 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 { if w%4 != 0 || h%4 != 0 {
t.Error("ceil") t.Fatalf("size is not quantized: %dx%d", w, h)
} }
if clampDimension(50, 100) != 50 { w, h = clampTermSize(3, 2, 100, 100, 4)
t.Error("passthrough") if w != 3 || h != 2 {
t.Fatalf("small size = %dx%d", w, h)
} }
} }
+4
View File
@@ -28,6 +28,8 @@ type Config struct {
MaxAuthAttempts int MaxAuthAttempts int
HandshakeTimeout time.Duration HandshakeTimeout time.Duration
MaxDimension int MaxDimension int
MaxTerminalCells int
SessionTimeout time.Duration
RenderCacheMB int RenderCacheMB int
BrightnessThreshold int BrightnessThreshold int
Charset string Charset string
@@ -126,6 +128,8 @@ func loadConfig() Config {
MaxAuthAttempts: envInt("MAX_AUTH_ATTEMPTS", 6, 1, maxInt), MaxAuthAttempts: envInt("MAX_AUTH_ATTEMPTS", 6, 1, maxInt),
HandshakeTimeout: envDurationMs("HANDSHAKE_TIMEOUT", 10*time.Second), HandshakeTimeout: envDurationMs("HANDSHAKE_TIMEOUT", 10*time.Second),
MaxDimension: envInt("MAX_DIMENSION", 512, 1, 4096), 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), RenderCacheMB: envInt("RENDER_CACHE_MB", 256, 0, maxInt),
BrightnessThreshold: envInt("BRIGHTNESS_THRESHOLD", 40, 0, 100), BrightnessThreshold: envInt("BRIGHTNESS_THRESHOLD", 40, 0, 100),
Charset: envString("CHARSET", "detailed"), Charset: envString("CHARSET", "detailed"),
+95 -23
View File
@@ -6,16 +6,68 @@ import (
"fmt" "fmt"
"math" "math"
"os" "os"
"sync"
) )
// .tsf container, little-endian: "TSFR" | version uint16 | fps float64 | // .tsf container, little-endian: "TSFR" | version uint16 | fps float64 |
// count uint32 | count × (colorLen uint32, color JPEG). // count uint32 | count × (colorLen uint32, color JPEG).
const ( const (
tsfMagic = "TSFR" tsfMagic = "TSFR"
tsfVersion = 1 tsfVersion = 1
maxTSFFPS = 240
maxTSFFrameCount = 10_000_000
) )
type frameFile struct {
data []byte
cleanup func() error
once sync.Once
err error
}
func (f *frameFile) Close() error {
if f == nil {
return nil
}
f.once.Do(func() {
if f.cleanup != nil {
f.err = f.cleanup()
}
f.data = nil
})
return f.err
}
var frameFileOwners sync.Map // map[*FramesContainer]*frameFile
func (data *FramesContainer) Close() error {
if data == nil {
return nil
}
owner, ok := frameFileOwners.LoadAndDelete(data)
if !ok {
return nil
}
data.ColorFrames = nil
return owner.(*frameFile).Close()
}
func writeTSF(output string, data *FramesContainer) error { func writeTSF(output string, data *FramesContainer) error {
if data == nil {
return fmt.Errorf("cannot write nil frames container")
}
if math.IsNaN(data.FPS) || math.IsInf(data.FPS, 0) || data.FPS <= 0 || data.FPS > maxTSFFPS {
return fmt.Errorf("cannot write .tsf: fps must be finite, positive, and at most %d", maxTSFFPS)
}
if len(data.ColorFrames) > maxTSFFrameCount || uint64(len(data.ColorFrames)) > math.MaxUint32 {
return fmt.Errorf("cannot write .tsf: frame count exceeds limit")
}
for i, frame := range data.ColorFrames {
if uint64(len(frame)) > math.MaxUint32 {
return fmt.Errorf("cannot write .tsf: frame %d length exceeds uint32", i)
}
}
f, err := os.Create(output) f, err := os.Create(output)
if err != nil { if err != nil {
return err return err
@@ -48,10 +100,17 @@ func writeTSF(output string, data *FramesContainer) error {
} }
func loadTSF(filename string) (*FramesContainer, error) { func loadTSF(filename string) (*FramesContainer, error) {
raw, err := readFrameFile(filename) file, err := readFrameFile(filename)
if err != nil { if err != nil {
return nil, err return nil, err
} }
owned := false
defer func() {
if !owned {
_ = file.Close()
}
}()
raw := file.data
invalid := func() error { invalid := func() error {
return fmt.Errorf("invalid frames file %q: corrupt .tsf container", filename) return fmt.Errorf("invalid frames file %q: corrupt .tsf container", filename)
} }
@@ -65,27 +124,40 @@ func loadTSF(filename string) (*FramesContainer, error) {
} }
fps := math.Float64frombits(binary.LittleEndian.Uint64(raw[6:])) fps := math.Float64frombits(binary.LittleEndian.Uint64(raw[6:]))
count := binary.LittleEndian.Uint32(raw[14:]) count := binary.LittleEndian.Uint32(raw[14:])
if math.IsNaN(fps) || math.IsInf(fps, 0) || fps <= 0 || fps > maxTSFFPS {
colorFrames := make([][]byte, 0, count)
off := 18
for range count {
if off+4 > len(raw) {
return nil, invalid()
}
n := int(binary.LittleEndian.Uint32(raw[off:]))
off += 4
if off+n > len(raw) {
return nil, invalid()
}
colorFrames = append(colorFrames, raw[off:off+n])
off += n
}
if len(colorFrames) == 0 || fps <= 0 {
return nil, fmt.Errorf( return nil, fmt.Errorf(
"invalid frames file %q: expected non-empty frames and a positive fps", "invalid frames file %q: fps must be finite, greater than 0, and at most %d",
filename, filename, maxTSFFPS,
) )
} }
return &FramesContainer{ColorFrames: colorFrames, FPS: fps}, nil if count == 0 {
return nil, fmt.Errorf("invalid frames file %q: expected non-empty frames", filename)
}
if count > maxTSFFrameCount || uint64(count) > uint64((len(raw)-18)/4) {
return nil, invalid()
}
colorFrames := make([][]byte, 0, int(count))
off := 18
for range count {
if len(raw)-off < 4 {
return nil, invalid()
}
n := uint64(binary.LittleEndian.Uint32(raw[off:]))
off += 4
if n > uint64(len(raw)-off) {
return nil, invalid()
}
nativeLen := int(n)
colorFrames = append(colorFrames, raw[off:off+nativeLen])
off += nativeLen
}
if off != len(raw) {
return nil, invalid()
}
data := &FramesContainer{ColorFrames: colorFrames, FPS: fps}
frameFileOwners.Store(data, file)
owned = true
return data, nil
} }
+82
View File
@@ -0,0 +1,82 @@
package main
import (
"encoding/binary"
"math"
"os"
"path/filepath"
"strings"
"testing"
)
func tsfHeader(fps float64, count uint32) []byte {
raw := make([]byte, 18)
copy(raw, tsfMagic)
binary.LittleEndian.PutUint16(raw[4:], tsfVersion)
binary.LittleEndian.PutUint64(raw[6:], math.Float64bits(fps))
binary.LittleEndian.PutUint32(raw[14:], count)
return raw
}
func writeRawTSF(t *testing.T, raw []byte) string {
t.Helper()
path := filepath.Join(t.TempDir(), "frames.tsf")
if err := os.WriteFile(path, raw, 0o644); err != nil {
t.Fatal(err)
}
return path
}
func TestTSFRejectsInvalidFPS(t *testing.T) {
for _, fps := range []float64{math.NaN(), math.Inf(1), math.Inf(-1), -1, 0, 240.01} {
raw := append(tsfHeader(fps, 1), 0, 0, 0, 0)
if _, err := loadTSF(writeRawTSF(t, raw)); err == nil {
t.Errorf("loadTSF accepted fps %v", fps)
}
}
}
func TestTSFRejectsImpossibleCountsAndLengths(t *testing.T) {
if _, err := loadTSF(writeRawTSF(t, tsfHeader(30, math.MaxUint32))); err == nil {
t.Fatal("loadTSF accepted impossible frame count")
}
raw := append(tsfHeader(30, 1), 0xff, 0xff, 0xff, 0xff)
if _, err := loadTSF(writeRawTSF(t, raw)); err == nil {
t.Fatal("loadTSF accepted overflowing frame length")
}
}
func TestTSFCloseReleasesOwnedFrames(t *testing.T) {
path := filepath.Join(t.TempDir(), "frames.tsf")
if err := writeTSF(path, &FramesContainer{FPS: 30, ColorFrames: [][]byte{{1, 2, 3}}}); err != nil {
t.Fatal(err)
}
frames, err := loadTSF(path)
if err != nil {
t.Fatal(err)
}
if got := frames.ColorFrames[0]; len(got) != 3 || got[0] != 1 {
t.Fatalf("unexpected zero-copy frame data: %v", got)
}
if err := frames.Close(); err != nil {
t.Fatalf("Close: %v", err)
}
if frames.ColorFrames != nil {
t.Fatal("Close retained references to released frame data")
}
if err := frames.Close(); err != nil {
t.Fatalf("second Close: %v", err)
}
}
func TestTSFWriteRejectsInvalidHeaderValuesBeforeCreate(t *testing.T) {
path := filepath.Join(t.TempDir(), "frames.tsf")
err := writeTSF(path, &FramesContainer{FPS: math.NaN(), ColorFrames: [][]byte{{1}}})
if err == nil || !strings.Contains(err.Error(), "fps") {
t.Fatalf("writeTSF error = %v", err)
}
if _, err := os.Stat(path); !os.IsNotExist(err) {
t.Fatalf("invalid write created output: %v", err)
}
}
+181 -73
View File
@@ -3,13 +3,14 @@ package main
import ( import (
"bytes" "bytes"
"container/list" "container/list"
"fmt"
"image" "image"
"image/color" "image/color"
"image/jpeg" "image/jpeg"
"strconv" "strconv"
"strings" "strings"
"sync" "sync"
"sync/atomic"
"time"
"unicode/utf8" "unicode/utf8"
"golang.org/x/image/draw" "golang.org/x/image/draw"
@@ -64,6 +65,8 @@ func resolveCharset(charset string) string {
return charset return charset
} }
const maxPooledBuffer = 4 << 20
var pixPool sync.Pool var pixPool sync.Pool
func getPixBuf(n int) []byte { func getPixBuf(n int) []byte {
@@ -76,7 +79,9 @@ func getPixBuf(n int) []byte {
} }
func putPixBuf(b []byte) { func putPixBuf(b []byte) {
pixPool.Put(&b) if cap(b) <= maxPooledBuffer {
pixPool.Put(&b)
}
} }
var outPool sync.Pool var outPool sync.Pool
@@ -91,7 +96,9 @@ func getOutBuf(capacity int) []byte {
} }
func putOutBuf(b []byte) { func putOutBuf(b []byte) {
outPool.Put(&b) if cap(b) <= maxPooledBuffer {
outPool.Put(&b)
}
} }
func resizeFrame(frame, pix []byte, width, height int, keepAspectRatio bool) (*image.RGBA, error) { func resizeFrame(frame, pix []byte, width, height int, keepAspectRatio bool) (*image.RGBA, error) {
@@ -149,16 +156,20 @@ func rampIndex(brightness, threshold, total int, invert bool) int {
return index return index
} }
func frameToAscii(img *image.RGBA, rampLUT *[101][]byte) string { func frameToAscii(img *image.RGBA, rampLUT *[101][]byte) []byte {
pix := img.Pix pix := img.Pix
buf := getOutBuf(len(pix)) maxCharBytes := 1
for _, char := range rampLUT {
maxCharBytes = max(maxCharBytes, len(char))
}
buf := getOutBuf(len(pix) / 4 * maxCharBytes)
for o := 0; o < len(pix); o += 4 { for o := 0; o < len(pix); o += 4 {
brightness := (int(pix[o])*299 + int(pix[o+1])*587 + int(pix[o+2])*114) / 255 / 10 brightness := (int(pix[o])*299 + int(pix[o+1])*587 + int(pix[o+2])*114) / 255 / 10
buf = append(buf, rampLUT[brightness]...) buf = append(buf, rampLUT[brightness]...)
} }
ascii := string(buf) output := bytes.Clone(buf)
putOutBuf(buf) putOutBuf(buf)
return ascii return output
} }
const ansiReset = "\x1b[0m" const ansiReset = "\x1b[0m"
@@ -208,97 +219,181 @@ func appendColor(buf []byte, r, g, b uint8, tier colorTier) []byte {
return append(buf, 'm') return append(buf, 'm')
} }
func frameToAnsi(img *image.RGBA, rampLUT *[101][]byte, tier colorTier) string { func frameToAnsi(img *image.RGBA, rampLUT *[101][]byte, tier colorTier) []byte {
bounds := img.Bounds() bounds := img.Bounds()
buf := getOutBuf(bounds.Dx() * bounds.Dy() * 16) bytesPerCell := 11
if tier == colorTierTrueColor {
bytesPerCell = 16
}
buf := getOutBuf(bounds.Dx() * bounds.Dy() * bytesPerCell)
var lastR, lastG, lastB uint8 var lastR, lastG, lastB uint8
last256 := -1
first := true first := true
for y := bounds.Min.Y; y < bounds.Max.Y; y++ { for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
o := img.PixOffset(bounds.Min.X, y)
for x := bounds.Min.X; x < bounds.Max.X; x++ { for x := bounds.Min.X; x < bounds.Max.X; x++ {
o := img.PixOffset(x, y)
r, g, bl := img.Pix[o], img.Pix[o+1], img.Pix[o+2] r, g, bl := img.Pix[o], img.Pix[o+1], img.Pix[o+2]
brightness := (int(r)*299 + int(g)*587 + int(bl)*114) / 255 / 10 brightness := (int(r)*299 + int(g)*587 + int(bl)*114) / 255 / 10
if first || r != lastR || g != lastG || bl != lastB { colorChanged := first || r != lastR || g != lastG || bl != lastB
if tier == colorTier256 {
index := quantize256(r, g, bl)
colorChanged = first || index != last256
last256 = index
}
if colorChanged {
buf = appendColor(buf, r, g, bl, tier) buf = appendColor(buf, r, g, bl, tier)
lastR, lastG, lastB = r, g, bl lastR, lastG, lastB = r, g, bl
first = false first = false
} }
buf = append(buf, rampLUT[brightness]...) buf = append(buf, rampLUT[brightness]...)
o += 4
} }
if y < bounds.Max.Y-1 { if y < bounds.Max.Y-1 {
buf = append(buf, ansiReset+"\r\n"...) buf = append(buf, "\r\n"...)
first = true
} }
} }
buf = append(buf, ansiReset...) buf = append(buf, ansiReset...)
ascii := string(buf) output := bytes.Clone(buf)
putOutBuf(buf) putOutBuf(buf)
return ascii return output
} }
type renderCache struct { type cacheKey struct {
setID int
index int
width int
height int
keepAspectRatio bool
tier colorTier
}
type renderCacheShard struct {
mu sync.Mutex mu sync.Mutex
maxBytes int64 maxBytes int64
size int64 size int64
entries map[string]*list.Element entries map[cacheKey]*list.Element
order *list.List order *list.List
} }
type cacheEntry struct { type renderCache struct {
key string shards []renderCacheShard
ascii string size atomic.Int64
hits atomic.Uint64
misses atomic.Uint64
evictions atomic.Uint64
rejections atomic.Uint64
renders atomic.Uint64
renderNs atomic.Uint64
} }
func entryCost(key, ascii string) int64 { type cacheEntry struct {
return int64(len(key)+len(ascii)) + 128 key cacheKey
ascii []byte
cost int64
}
func entryCost(_ cacheKey, ascii []byte) int64 {
return int64(cap(ascii)) + 160
} }
func newRenderCache(maxBytes int64) *renderCache { func newRenderCache(maxBytes int64) *renderCache {
if maxBytes <= 0 { if maxBytes <= 0 {
return nil return nil
} }
return &renderCache{ shardCount := int(min(int64(16), max(int64(1), maxBytes/(1<<20))))
maxBytes: maxBytes, cache := &renderCache{shards: make([]renderCacheShard, shardCount)}
entries: make(map[string]*list.Element), for i := range cache.shards {
order: list.New(), cache.shards[i] = renderCacheShard{
maxBytes: maxBytes / int64(shardCount),
entries: make(map[cacheKey]*list.Element),
order: list.New(),
}
} }
return cache
} }
func (c *renderCache) get(key string) (string, bool) { func (c *renderCache) shard(key cacheKey) *renderCacheShard {
hash := uint64(key.setID)*0x9e3779b185ebca87 ^ uint64(key.index)*0xc2b2ae3d27d4eb4f
hash ^= uint64(key.width)<<32 | uint64(uint32(key.height))
hash ^= uint64(key.tier)<<1 | uint64(boolToInt(key.keepAspectRatio))
return &c.shards[hash%uint64(len(c.shards))]
}
func boolToInt(value bool) int {
if value {
return 1
}
return 0
}
func (c *renderCache) get(key cacheKey) ([]byte, bool) {
if c == nil { if c == nil {
return "", false return nil, false
} }
c.mu.Lock() shard := c.shard(key)
defer c.mu.Unlock() shard.mu.Lock()
el, ok := c.entries[key] defer shard.mu.Unlock()
el, ok := shard.entries[key]
if !ok { if !ok {
return "", false c.misses.Add(1)
return nil, false
} }
c.order.MoveToBack(el) c.hits.Add(1)
shard.order.MoveToBack(el)
return el.Value.(*cacheEntry).ascii, true return el.Value.(*cacheEntry).ascii, true
} }
func (c *renderCache) put(key, ascii string) { func (c *renderCache) put(key cacheKey, ascii []byte) {
if c == nil { if c == nil {
return return
} }
shard := c.shard(key)
cost := entryCost(key, ascii) cost := entryCost(key, ascii)
if cost > c.maxBytes { if cost > shard.maxBytes {
c.rejections.Add(1)
return return
} }
c.mu.Lock() shard.mu.Lock()
defer c.mu.Unlock() defer shard.mu.Unlock()
if _, ok := c.entries[key]; ok { if _, ok := shard.entries[key]; ok {
return return
} }
c.entries[key] = c.order.PushBack(&cacheEntry{key, ascii}) shard.entries[key] = shard.order.PushBack(&cacheEntry{key: key, ascii: ascii, cost: cost})
c.size += cost shard.size += cost
for c.size > c.maxBytes { c.size.Add(cost)
oldest := c.order.Front() for shard.size > shard.maxBytes {
c.order.Remove(oldest) oldest := shard.order.Front()
shard.order.Remove(oldest)
evicted := oldest.Value.(*cacheEntry) evicted := oldest.Value.(*cacheEntry)
delete(c.entries, evicted.key) delete(shard.entries, evicted.key)
c.size -= entryCost(evicted.key, evicted.ascii) shard.size -= evicted.cost
c.size.Add(-evicted.cost)
c.evictions.Add(1)
}
}
type renderCacheStats struct {
SizeBytes int64
Hits uint64
Misses uint64
Evictions uint64
Rejections uint64
Renders uint64
RenderTime time.Duration
}
func (c *renderCache) stats() renderCacheStats {
if c == nil {
return renderCacheStats{}
}
return renderCacheStats{
SizeBytes: c.size.Load(),
Hits: c.hits.Load(),
Misses: c.misses.Load(),
Evictions: c.evictions.Load(),
Rejections: c.rejections.Load(),
Renders: c.renders.Load(),
RenderTime: time.Duration(c.renderNs.Load()),
} }
} }
@@ -310,7 +405,13 @@ type FrameRenderer struct {
cache *renderCache cache *renderCache
inflightMu sync.Mutex inflightMu sync.Mutex
inflight map[string]chan struct{} inflight map[cacheKey]*renderCall
}
type renderCall struct {
done chan struct{}
value []byte
err error
} }
func newFrameRenderer(setID int, colorFrames [][]byte, options asciiOptions, cache *renderCache) *FrameRenderer { func newFrameRenderer(setID int, colorFrames [][]byte, options asciiOptions, cache *renderCache) *FrameRenderer {
@@ -321,47 +422,46 @@ func newFrameRenderer(setID int, colorFrames [][]byte, options asciiOptions, cac
options: options, options: options,
rampLUT: buildRampLUT(ramp, options), rampLUT: buildRampLUT(ramp, options),
cache: cache, cache: cache,
inflight: make(map[string]chan struct{}), inflight: make(map[cacheKey]*renderCall),
} }
} }
func (r *FrameRenderer) render(index, width, height int, keepAspectRatio bool, tier colorTier) (string, error) { func (r *FrameRenderer) render(index, width, height int, keepAspectRatio bool, tier colorTier) ([]byte, error) {
key := fmt.Sprintf("%d:%d:%dx%d:%t:%d", r.setID, index, width, height, keepAspectRatio, tier) key := cacheKey{r.setID, index, width, height, keepAspectRatio, tier}
if ascii, ok := r.cache.get(key); ok { if ascii, ok := r.cache.get(key); ok {
return ascii, nil return ascii, nil
} }
if r.cache != nil { r.inflightMu.Lock()
for { if call, ok := r.inflight[key]; ok {
r.inflightMu.Lock() r.inflightMu.Unlock()
wait, ok := r.inflight[key] <-call.done
if !ok { return call.value, call.err
done := make(chan struct{}) }
r.inflight[key] = done call := &renderCall{done: make(chan struct{})}
r.inflightMu.Unlock() r.inflight[key] = call
defer func() { r.inflightMu.Unlock()
r.inflightMu.Lock() defer func() {
delete(r.inflight, key) r.inflightMu.Lock()
r.inflightMu.Unlock() delete(r.inflight, key)
close(done) r.inflightMu.Unlock()
}() close(call.done)
break }()
}
r.inflightMu.Unlock() if ascii, ok := r.cache.get(key); ok {
<-wait call.value = ascii
if ascii, ok := r.cache.get(key); ok { return ascii, nil
return ascii, nil
}
}
} }
started := time.Now()
pix := getPixBuf(4 * width * height) pix := getPixBuf(4 * width * height)
img, err := resizeFrame(r.colorFrames[index], pix, width, height, keepAspectRatio) img, err := resizeFrame(r.colorFrames[index], pix, width, height, keepAspectRatio)
if err != nil { if err != nil {
putPixBuf(pix) putPixBuf(pix)
return "", err call.err = err
return nil, err
} }
var ascii string var ascii []byte
if tier == colorTierNone { if tier == colorTierNone {
ascii = frameToAscii(img, r.rampLUT) ascii = frameToAscii(img, r.rampLUT)
} else { } else {
@@ -369,6 +469,14 @@ func (r *FrameRenderer) render(index, width, height int, keepAspectRatio bool, t
} }
putPixBuf(pix) putPixBuf(pix)
if r.cache != nil && cap(ascii) > len(ascii)+len(ascii)/4 {
ascii = bytes.Clone(ascii)
}
r.cache.put(key, ascii) r.cache.put(key, ascii)
if r.cache != nil {
r.cache.renders.Add(1)
r.cache.renderNs.Add(uint64(time.Since(started)))
}
call.value = ascii
return ascii, nil return ascii, nil
} }
+7 -5
View File
@@ -47,18 +47,20 @@ func sanitizeN(value any, maxLength int) string {
str = fmt.Sprint(v) str = fmt.Sprint(v)
} }
var b strings.Builder var b strings.Builder
count := 0
for _, r := range str { for _, r := range str {
if count >= maxLength {
b.WriteRune('…')
return b.String()
}
if r < 0x20 || (r >= 0x7f && r <= 0x9f) { if r < 0x20 || (r >= 0x7f && r <= 0x9f) {
b.WriteRune('') b.WriteRune('')
} else { } else {
b.WriteRune(r) b.WriteRune(r)
} }
count++
} }
out := []rune(b.String()) return b.String()
if len(out) > maxLength {
return string(out[:maxLength]) + "…"
}
return string(out)
} }
func emit(level logLevel, name string, stream *os.File, args []any) { func emit(level logLevel, name string, stream *os.File, args []any) {
+21 -1
View File
@@ -83,6 +83,8 @@ func generateFrames(framesDir, videoArg string, resolution int) {
} }
} }
const frameDataWarnBytes = 2 << 30
func loadAllFrames(framesDir string) []*FramesContainer { func loadAllFrames(framesDir string) []*FramesContainer {
entries, err := os.ReadDir(framesDir) entries, err := os.ReadDir(framesDir)
var files []string var files []string
@@ -101,6 +103,22 @@ func loadAllFrames(framesDir string) []*FramesContainer {
framesDir, 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))) concurrency := min(len(files), max(1, min(runtime.NumCPU(), 4)))
@@ -226,14 +244,16 @@ func main() {
GoodbyeText: goodbyeText, GoodbyeText: goodbyeText,
VideoSets: videoSets, VideoSets: videoSets,
}) })
defer server.Close()
sigCh := make(chan os.Signal, 1) sigCh := make(chan os.Signal, 1)
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
go func() { go func() {
sig := <-sigCh sig := <-sigCh
logInfo(fmt.Sprintf("Received %s, shutting down...", sig)) logInfo(fmt.Sprintf("Received %s, shutting down...", sig))
forceExit := time.AfterFunc(5*time.Second, func() { os.Exit(0) })
server.Close() server.Close()
time.AfterFunc(5*time.Second, func() { os.Exit(0) }) forceExit.Stop()
}() }()
if err := server.Listen(config.Host, config.Port); err != nil { if err := server.Listen(config.Host, config.Port); err != nil {
+3 -2
View File
@@ -4,6 +4,7 @@ package main
import "os" import "os"
func readFrameFile(filename string) ([]byte, error) { func readFrameFile(filename string) (*frameFile, error) {
return os.ReadFile(filename) data, err := os.ReadFile(filename)
return &frameFile{data: data}, err
} }
+11 -4
View File
@@ -7,7 +7,7 @@ import (
"syscall" "syscall"
) )
func readFrameFile(filename string) ([]byte, error) { func readFrameFile(filename string) (*frameFile, error) {
f, err := os.Open(filename) f, err := os.Open(filename)
if err != nil { if err != nil {
return nil, err return nil, err
@@ -20,11 +20,18 @@ func readFrameFile(filename string) ([]byte, error) {
} }
size := info.Size() size := info.Size()
if size <= 0 || size != int64(int(size)) { if size <= 0 || size != int64(int(size)) {
return os.ReadFile(filename) data, err := os.ReadFile(filename)
return &frameFile{data: data}, err
} }
data, err := syscall.Mmap(int(f.Fd()), 0, int(size), syscall.PROT_READ, syscall.MAP_SHARED) data, err := syscall.Mmap(int(f.Fd()), 0, int(size), syscall.PROT_READ, syscall.MAP_SHARED)
if err != nil { if err != nil {
return os.ReadFile(filename) data, err := os.ReadFile(filename)
return &frameFile{data: data}, err
} }
return data, nil return &frameFile{
data: data,
cleanup: func() error {
return syscall.Munmap(data)
},
}, nil
} }
+127
View File
@@ -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 != "abc…" {
t.Fatalf("sanitizeN = %q", got)
}
}
+29
View File
@@ -2,6 +2,7 @@ package main
import ( import (
"path/filepath" "path/filepath"
"sync/atomic"
"testing" "testing"
) )
@@ -15,6 +16,7 @@ func loadBenchSet(b *testing.B) *FramesContainer {
if err != nil { if err != nil {
b.Skip("failed to load frame set:", err) b.Skip("failed to load frame set:", err)
} }
b.Cleanup(func() { _ = fc.Close() })
return fc return fc
} }
@@ -24,6 +26,7 @@ func benchRender(b *testing.B, tier colorTier, w, h int) {
brightnessThreshold: 40, brightnessThreshold: 40,
charset: "detailed", charset: "detailed",
}, nil) }, nil)
b.ReportAllocs()
b.ResetTimer() b.ResetTimer()
for i := 0; i < b.N; i++ { for i := 0; i < b.N; i++ {
if _, err := r.render(i%len(fc.ColorFrames), w, h, false, tier); err != nil { if _, err := r.render(i%len(fc.ColorFrames), w, h, false, tier); err != nil {
@@ -36,3 +39,29 @@ func BenchmarkRenderTrueColor(b *testing.B) { benchRender(b, colorTierTrueColor,
func BenchmarkRender256(b *testing.B) { benchRender(b, colorTier256, 120, 40) } func BenchmarkRender256(b *testing.B) { benchRender(b, colorTier256, 120, 40) }
func BenchmarkRenderGray(b *testing.B) { benchRender(b, colorTierNone, 120, 40) } func BenchmarkRenderGray(b *testing.B) { benchRender(b, colorTierNone, 120, 40) }
func BenchmarkRenderTrueBig(b *testing.B) { benchRender(b, colorTierTrueColor, 240, 70) } func BenchmarkRenderTrueBig(b *testing.B) { benchRender(b, colorTierTrueColor, 240, 70) }
func BenchmarkRenderCachedParallel(b *testing.B) {
fc := loadBenchSet(b)
r := newFrameRenderer(0, fc.ColorFrames, asciiOptions{
brightnessThreshold: 40,
charset: "detailed",
}, newRenderCache(8<<20))
frame, err := r.render(0, 120, 40, false, colorTierTrueColor)
if err != nil {
b.Fatal(err)
}
b.SetBytes(int64(len(frame)))
b.ReportAllocs()
var failures atomic.Uint64
b.ResetTimer()
b.RunParallel(func(pb *testing.PB) {
for pb.Next() {
if _, err := r.render(0, 120, 40, false, colorTierTrueColor); err != nil {
failures.Add(1)
}
}
})
if failures.Load() != 0 {
b.Fatalf("render failures: %d", failures.Load())
}
}
+294 -61
View File
@@ -4,6 +4,8 @@ import (
"encoding/binary" "encoding/binary"
"errors" "errors"
"fmt" "fmt"
"io"
"math"
"math/rand" "math/rand"
"net" "net"
"strings" "strings"
@@ -14,14 +16,84 @@ import (
) )
const ( const (
clearScreen = "\x1b[2J\x1b[0f" clearScreen = "\x1b[2J\x1b[0f"
hideCursor = "\x1b[?25l" hideCursor = "\x1b[?25l"
showCursor = "\x1b[?25h" showCursor = "\x1b[?25h"
syncStart = "\x1b[?2026h" syncStart = "\x1b[?2026h"
syncEnd = "\x1b[?2026l" syncEnd = "\x1b[?2026l"
homeCursor = "\x1b[H" 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 { type ConnectionTracker struct {
mu sync.Mutex mu sync.Mutex
counts map[string]int counts map[string]int
@@ -32,15 +104,18 @@ func newConnectionTracker() *ConnectionTracker {
return &ConnectionTracker{counts: make(map[string]int)} 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() t.mu.Lock()
defer t.mu.Unlock() defer t.mu.Unlock()
if t.total >= maxTotal || t.counts[ip] >= maxPerIP {
return t.counts[ip], t.total, false
}
t.counts[ip]++ t.counts[ip]++
t.total++ 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() t.mu.Lock()
defer t.mu.Unlock() defer t.mu.Unlock()
if _, ok := t.counts[ip]; !ok { if _, ok := t.counts[ip]; !ok {
@@ -61,10 +136,40 @@ func (t *ConnectionTracker) totalCount() int {
return t.total 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() t.mu.Lock()
defer t.mu.Unlock() 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 { type frameSet struct {
@@ -76,10 +181,16 @@ type Server struct {
config Config config Config
sshConfig *ssh.ServerConfig sshConfig *ssh.ServerConfig
sets []frameSet sets []frameSet
cache *renderCache
tracker *ConnectionTracker tracker *ConnectionTracker
sessions *SessionTracker
fakeLogin *string fakeLogin *string
goodbye *string goodbye *string
mu sync.Mutex
listener net.Listener listener net.Listener
conns map[net.Conn]struct{}
connWG sync.WaitGroup
closing bool
closeOnce sync.Once closeOnce sync.Once
} }
@@ -92,14 +203,25 @@ type ServerDeps struct {
VideoSets []*FramesContainer VideoSets []*FramesContainer
} }
func clampDimension(value, max int) int { func clampTermSize(cols, rows, maxDimension, maxCells, quantum int) (int, int) {
if value < 1 { cols = max(cols, 1)
return 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 { cols = max(1, int(math.Floor(float64(cols)*scale)))
return max 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 { func createServer(deps ServerDeps) *Server {
@@ -155,9 +277,12 @@ func createServer(deps ServerDeps) *Server {
config: config, config: config,
sshConfig: sshConfig, sshConfig: sshConfig,
sets: sets, sets: sets,
cache: cache,
tracker: newConnectionTracker(), tracker: newConnectionTracker(),
sessions: newSessionTracker(),
fakeLogin: deps.FakeLoginText, fakeLogin: deps.FakeLoginText,
goodbye: deps.GoodbyeText, 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 { if err != nil {
return err return err
} }
s.mu.Lock()
if s.closing {
s.mu.Unlock()
_ = listener.Close()
return nil
}
s.listener = listener s.listener = listener
s.mu.Unlock()
logInfo(fmt.Sprintf("TrollSSH listening on %s:%d", host, port)) logInfo(fmt.Sprintf("TrollSSH listening on %s:%d", host, port))
for { for {
conn, err := listener.Accept() conn, err := listener.Accept()
@@ -184,29 +316,68 @@ func (s *Server) Listen(host string, port int) error {
} }
return err 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() { func (s *Server) Close() {
s.closeOnce.Do(func() { s.closeOnce.Do(func() {
if s.listener != nil { s.mu.Lock()
_ = s.listener.Close() 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) { func (s *Server) handleConn(conn net.Conn, ip string, activeForIP, total int) {
ip := hostOnly(conn.RemoteAddr().String()) defer func() {
s.tracker.release(ip)
if s.tracker.hasReachedLimits(ip, s.config.MaxConnections, s.config.MaxTotalConnections) { s.mu.Lock()
_ = conn.Close() delete(s.conns, conn)
logWarn("Connection rejected (limit reached) from", ip) s.mu.Unlock()
return s.connWG.Done()
} }()
activeForIP := s.tracker.increment(ip)
defer s.tracker.decrement(ip)
if s.config.HandshakeTimeout > 0 { if s.config.HandshakeTimeout > 0 {
_ = conn.SetDeadline(time.Now().Add(s.config.HandshakeTimeout)) _ = 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)) setIndex := rand.Intn(len(s.sets))
logInfo(fmt.Sprintf( logInfo(fmt.Sprintf(
"New connection from %s (ip=%d, total=%d) -> playing %q", "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) go ssh.DiscardRequests(reqs)
var sessionWG sync.WaitGroup
for newChannel := range chans { for newChannel := range chans {
if newChannel.ChannelType() != "session" { if newChannel.ChannelType() != "session" {
_ = newChannel.Reject(ssh.UnknownChannelType, "unknown channel type") _ = newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
continue continue
} }
channel, requests, err := newChannel.Accept() if !s.sessions.tryAcquire(sshConn, maxSessionsPerConn, s.config.MaxTotalConnections) {
if err != nil { _ = newChannel.Reject(ssh.ResourceShortage, "session limit reached")
continue 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) logInfo("Client closed connection from", ip)
} }
type termSize struct { type termSize struct {
mu sync.Mutex mu sync.Mutex
width int width int
height 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.mu.Lock()
t.width = clampDimension(w, maxDim) if !force && time.Since(t.updated) < resizeDebounce {
t.height = clampDimension(h, maxDim) t.mu.Unlock()
return
}
t.width, t.height = clampTermSize(w, h, maxDimension, maxCells, terminalSizeQuantum)
t.updated = time.Now()
t.mu.Unlock() t.mu.Unlock()
} }
@@ -304,20 +498,22 @@ func (s *Server) handleSession(
ip string, ip string,
initialSetIndex int, initialSetIndex int,
) { ) {
defer func() { _ = channel.Close() }()
size := &termSize{} size := &termSize{}
size.set(80, 24, s.config.MaxDimension) size.set(80, 24, s.config.MaxDimension, s.config.MaxTerminalCells, true)
tier := colorTierTrueColor tier := colorTierTrueColor
if s.config.ForceGrayscale { if s.config.ForceGrayscale {
tier = colorTierNone tier = colorTierNone
} }
started := false started := false
var playDone chan struct{}
for req := range requests { for req := range requests {
switch req.Type { switch req.Type {
case "pty-req": case "pty-req":
logDebug("Opening pty for session", ip) logDebug("Opening pty for session", ip)
if cols, rows, ok := parseDims(req.Payload); ok { 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 { if term, ok := parsePtyTerm(req.Payload); ok {
tier = detectColorTier(term) tier = detectColorTier(term)
@@ -331,7 +527,7 @@ func (s *Server) handleSession(
if len(req.Payload) >= 8 { if len(req.Payload) >= 8 {
cols := int(binary.BigEndian.Uint32(req.Payload)) cols := int(binary.BigEndian.Uint32(req.Payload))
rows := int(binary.BigEndian.Uint32(req.Payload[4:])) 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 { if req.WantReply {
_ = req.Reply(true, nil) _ = req.Reply(true, nil)
@@ -348,14 +544,24 @@ func (s *Server) handleSession(
_ = req.Reply(true, nil) _ = req.Reply(true, nil)
if !started { if !started {
started = true 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": case "shell":
logDebug("Opening shell for session", ip) logDebug("Opening shell for session", ip)
_ = req.Reply(true, nil) _ = req.Reply(true, nil)
if !started { if !started {
started = true 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: default:
if req.WantReply { if req.WantReply {
@@ -363,6 +569,10 @@ func (s *Server) handleSession(
} }
} }
} }
_ = channel.Close()
if playDone != nil {
<-playDone
}
} }
func (s *Server) pickNextSetIndex(exclude int) int { func (s *Server) pickNextSetIndex(exclude int) int {
@@ -391,11 +601,16 @@ func (s *Server) playVideo(
w, h := size.get() w, h := size.get()
logDebug(fmt.Sprintf("Terminal size %dx%d for %s", w, h, ip)) 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 { if s.fakeLogin != nil {
_, _ = channel.Write([]byte(clearScreen)) if err := writePartsWithTimeout(
_, _ = channel.Write([]byte(*s.fakeLogin)) sshConn, channel, outputStallTimeout, clearScreen, *s.fakeLogin,
); err != nil {
return
}
} }
done := make(chan struct{}) done := make(chan struct{})
@@ -439,13 +654,19 @@ func (s *Server) playVideo(
} }
}() }()
loginTimer := time.NewTimer(config.LoginDelay)
select { select {
case <-time.After(config.LoginDelay): case <-loginTimer.C:
case <-done: case <-done:
if !loginTimer.Stop() {
<-loginTimer.C
}
return return
} }
_, _ = channel.Write([]byte(hideCursor)) if err := writePartsWithTimeout(sshConn, channel, outputStallTimeout, hideCursor); err != nil {
return
}
frameInterval := func() time.Duration { frameInterval := func() time.Duration {
return time.Duration(float64(time.Second) / current.data.FPS) return time.Duration(float64(time.Second) / current.data.FPS)
@@ -457,8 +678,6 @@ func (s *Server) playVideo(
currentFrame := 0 currentFrame := 0
loopCount := 0 loopCount := 0
lastW, lastH := 0, 0 lastW, lastH := 0, 0
var writeBuf []byte
for { for {
select { select {
case <-done: case <-done:
@@ -489,11 +708,9 @@ func (s *Server) playVideo(
prefix = clearScreen prefix = clearScreen
lastW, lastH = w, h lastW, lastH = w, h
} }
writeBuf = append(writeBuf[:0], syncStart...) if err := writeFrameWithTimeout(
writeBuf = append(writeBuf, prefix...) sshConn, channel, outputStallTimeout, prefix, ascii,
writeBuf = append(writeBuf, ascii...) ); err != nil {
writeBuf = append(writeBuf, syncEnd...)
if _, err := channel.Write(writeBuf); err != nil {
closeSession() closeSession()
return return
} }
@@ -506,11 +723,27 @@ func (s *Server) playVideo(
currentFrame = 0 currentFrame = 0
loopCount++ loopCount++
if config.MaxLoop > 0 && loopCount >= config.MaxLoop { 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 { 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) logInfo("Playback finished, closing session", ip)
_ = channel.Close() _ = channel.Close()
_ = sshConn.Close() _ = sshConn.Close()
+220 -40
View File
@@ -1,15 +1,20 @@
package main package main
import ( import (
"bufio"
"bytes" "bytes"
"encoding/binary"
"encoding/json" "encoding/json"
"errors" "errors"
"fmt" "fmt"
"io" "io"
"math" "math"
"os"
"os/exec" "os/exec"
"path/filepath"
"strconv" "strconv"
"strings" "strings"
"time"
) )
var ( var (
@@ -17,29 +22,177 @@ var (
jpegEOI = []byte{0xff, 0xd9} jpegEOI = []byte{0xff, 0xd9}
) )
const (
maxJPEGFrameBytes = 64 << 20
maxFFmpegLogBytes = 64 << 10
)
type jpegFrameSplitter struct { type jpegFrameSplitter struct {
buffer []byte buffer []byte
scan int
inJPEG bool
} }
func (s *jpegFrameSplitter) push(chunk []byte) [][]byte { func (s *jpegFrameSplitter) push(chunk []byte, emit func([]byte) error) (int, error) {
s.buffer = append(s.buffer, chunk...) s.buffer = append(s.buffer, chunk...)
var frames [][]byte emitted := 0
for { for {
start := bytes.Index(s.buffer, jpegSOI) if !s.inJPEG {
if start == -1 { start := bytes.Index(s.buffer[s.scan:], jpegSOI)
break if start == -1 {
// Retain only a possible marker prefix spanning two reads.
if len(s.buffer) > 0 && s.buffer[len(s.buffer)-1] == jpegSOI[0] {
s.buffer = s.buffer[len(s.buffer)-1:]
} else {
s.buffer = s.buffer[:0]
}
s.scan = 0
return emitted, nil
}
start += s.scan
s.buffer = s.buffer[start:]
s.scan = len(jpegSOI)
s.inJPEG = true
} }
end := bytes.Index(s.buffer[start+len(jpegSOI):], jpegEOI)
end := bytes.Index(s.buffer[s.scan:], jpegEOI)
if end == -1 { if end == -1 {
break if len(s.buffer) > maxJPEGFrameBytes {
return emitted, fmt.Errorf("JPEG frame exceeds %d MiB limit", maxJPEGFrameBytes>>20)
}
s.scan = max(len(jpegSOI), len(s.buffer)-1)
return emitted, nil
} }
frameEnd := start + len(jpegSOI) + end + len(jpegEOI) frameEnd := s.scan + end + len(jpegEOI)
frame := make([]byte, frameEnd-start) if frameEnd > maxJPEGFrameBytes {
copy(frame, s.buffer[start:frameEnd]) return emitted, fmt.Errorf("JPEG frame exceeds %d MiB limit", maxJPEGFrameBytes>>20)
frames = append(frames, frame) }
if err := emit(s.buffer[:frameEnd]); err != nil {
return emitted, err
}
emitted++
s.buffer = s.buffer[frameEnd:] s.buffer = s.buffer[frameEnd:]
s.scan = 0
s.inJPEG = false
}
}
func (s *jpegFrameSplitter) finish() error {
if s.inJPEG {
return fmt.Errorf("ffmpeg produced a truncated JPEG frame")
}
return nil
}
type boundedLog struct {
buffer bytes.Buffer
limit int
}
func (w *boundedLog) Write(p []byte) (int, error) {
n := len(p)
if remaining := w.limit - w.buffer.Len(); remaining > 0 {
_, _ = w.buffer.Write(p[:min(len(p), remaining)])
}
return n, nil
}
func (w *boundedLog) String() string {
return strings.TrimSpace(w.buffer.String())
}
type streamingTSF struct {
file *os.File
writer *bufio.Writer
path string
count uint32
}
func newStreamingTSF(output string, fps float64) (*streamingTSF, error) {
if math.IsNaN(fps) || math.IsInf(fps, 0) || fps <= 0 || fps > maxTSFFPS {
return nil, fmt.Errorf("cannot write .tsf: fps must be finite and between 0 and %d", maxTSFFPS)
}
dir := filepath.Dir(output)
f, err := os.CreateTemp(dir, "."+filepath.Base(output)+"-*.tmp")
if err != nil {
return nil, err
}
s := &streamingTSF{file: f, writer: bufio.NewWriterSize(f, 1<<20), path: f.Name()}
if err := f.Chmod(0o644); err != nil {
s.abort()
return nil, err
}
if _, err := s.writer.WriteString(tsfMagic); err != nil {
s.abort()
return nil, err
}
var hdr [14]byte
binary.LittleEndian.PutUint16(hdr[0:], tsfVersion)
binary.LittleEndian.PutUint64(hdr[2:], math.Float64bits(fps))
if _, err := s.writer.Write(hdr[:]); err != nil {
s.abort()
return nil, err
}
return s, nil
}
func (s *streamingTSF) addFrame(frame []byte) error {
if s.count >= maxTSFFrameCount {
return fmt.Errorf("too many video frames")
}
if uint64(len(frame)) > math.MaxUint32 {
return fmt.Errorf("JPEG frame is too large")
}
var size [4]byte
binary.LittleEndian.PutUint32(size[:], uint32(len(frame)))
if _, err := s.writer.Write(size[:]); err != nil {
return err
}
if _, err := s.writer.Write(frame); err != nil {
return err
}
s.count++
return nil
}
func (s *streamingTSF) commit(output string) error {
if s.count == 0 {
return fmt.Errorf("no frames were decoded from the video")
}
if err := s.writer.Flush(); err != nil {
return err
}
if _, err := s.file.Seek(14, io.SeekStart); err != nil {
return err
}
var count [4]byte
binary.LittleEndian.PutUint32(count[:], s.count)
if _, err := s.file.Write(count[:]); err != nil {
return err
}
if err := s.file.Sync(); err != nil {
return err
}
if err := s.file.Close(); err != nil {
return err
}
s.file = nil
if err := os.Rename(s.path, output); err != nil {
return err
}
s.path = ""
return nil
}
func (s *streamingTSF) abort() {
if s.file != nil {
_ = s.file.Close()
s.file = nil
}
if s.path != "" {
_ = os.Remove(s.path)
s.path = ""
} }
return frames
} }
type ffprobeOutput struct { type ffprobeOutput struct {
@@ -72,59 +225,82 @@ func parseFrameRate(rate string) float64 {
return num return num
} }
func extractFrames(path, vf, label string, maxDimension, totalFrames int) ([][]byte, error) { func extractFrames(path, vf, label string, totalFrames int, emit func([]byte) error) (int, error) {
cmd := exec.Command( cmd := exec.Command(
"ffmpeg", "-i", path, "ffmpeg", "-hide_banner", "-loglevel", "error", "-nostats",
"-i", path,
"-c:v", "mjpeg", "-c:v", "mjpeg",
"-q:v", "3", "-q:v", "3",
"-vf", vf, "-vf", vf,
"-f", "image2pipe", "-f", "image2pipe",
"pipe:1", "pipe:1",
) )
var stderr bytes.Buffer stderr := &boundedLog{limit: maxFFmpegLogBytes}
cmd.Stderr = &stderr cmd.Stderr = stderr
stdout, err := cmd.StdoutPipe() stdout, err := cmd.StdoutPipe()
if err != nil { if err != nil {
return nil, fmt.Errorf("ffmpeg failed: %w", err) return 0, fmt.Errorf("ffmpeg failed: %w", err)
} }
if err := cmd.Start(); err != nil { if err := cmd.Start(); err != nil {
return nil, fmt.Errorf("ffmpeg failed: %w", err) return 0, fmt.Errorf("ffmpeg failed: %w", err)
} }
var frames [][]byte
splitter := &jpegFrameSplitter{} splitter := &jpegFrameSplitter{}
reportProgress := func(count int) { frameCount := 0
if totalFrames > 0 { lastReport := time.Now()
pct := min(100, int(math.Round(float64(count)/float64(totalFrames)*100))) reportProgress := func(force bool) {
fmt.Printf("\rGenerating %s frames: %d/%d (%d%%)", label, count, totalFrames, pct) if !force && time.Since(lastReport) < 250*time.Millisecond {
} else { return
fmt.Printf("\rGenerating %s frames: %d", label, count)
} }
if totalFrames > 0 {
pct := min(100, int(math.Round(float64(frameCount)/float64(totalFrames)*100)))
fmt.Printf("\rGenerating %s frames: %d/%d (%d%%)", label, frameCount, totalFrames, pct)
} else {
fmt.Printf("\rGenerating %s frames: %d", label, frameCount)
}
lastReport = time.Now()
}
failStream := func(err error) (int, error) {
_ = cmd.Process.Kill()
_ = cmd.Wait()
return frameCount, err
} }
buf := make([]byte, 256*1024) buf := make([]byte, 256*1024)
for { for {
n, err := stdout.Read(buf) n, readErr := stdout.Read(buf)
if n > 0 { if n > 0 {
frames = append(frames, splitter.push(buf[:n])...) emitted, splitErr := splitter.push(buf[:n], emit)
reportProgress(len(frames)) frameCount += emitted
if splitErr != nil {
return failStream(fmt.Errorf("ffmpeg stream error: %w", splitErr))
}
if emitted > 0 {
reportProgress(false)
}
} }
if err == io.EOF { if readErr == io.EOF {
break break
} }
if err != nil { if readErr != nil {
_ = cmd.Wait() return failStream(fmt.Errorf("ffmpeg stream error: %w", readErr))
return nil, fmt.Errorf("ffmpeg stream error: %s", err.Error())
} }
} }
if err := cmd.Wait(); err != nil { if err := cmd.Wait(); err != nil {
return nil, fmt.Errorf("ffmpeg failed: %s", strings.TrimSpace(stderr.String())) if msg := stderr.String(); msg != "" {
return frameCount, fmt.Errorf("ffmpeg failed: %s", msg)
}
return frameCount, fmt.Errorf("ffmpeg failed: %w", err)
} }
if len(frames) == 0 { if err := splitter.finish(); err != nil {
return nil, fmt.Errorf("no frames were decoded from the video") return frameCount, err
} }
if frameCount == 0 {
return 0, fmt.Errorf("no frames were decoded from the video")
}
reportProgress(true)
fmt.Println() fmt.Println()
return frames, nil return frameCount, nil
} }
func processVideo(path, output string, maxDimension int) error { func processVideo(path, output string, maxDimension int) error {
@@ -175,15 +351,19 @@ func processVideo(path, output string, maxDimension int) error {
maxDimension, maxDimension, maxDimension, maxDimension,
) )
colorFrames, err := extractFrames(path, scaleFilter, "color", maxDimension, totalFrames) outputFile, err := newStreamingTSF(output, fps)
if err != nil { if err != nil {
return err return err
} }
defer outputFile.abort()
videoData := FramesContainer{FPS: fps, ColorFrames: colorFrames} frameCount, err := extractFrames(path, scaleFilter, "color", totalFrames, outputFile.addFrame)
if err := writeTSF(output, &videoData); err != nil { if err != nil {
return err return err
} }
logInfo(fmt.Sprintf("Saved %d frames to %s", len(videoData.ColorFrames), output)) if err := outputFile.commit(output); err != nil {
return err
}
logInfo(fmt.Sprintf("Saved %d frames to %s", frameCount, output))
return nil return nil
} }
+122
View File
@@ -0,0 +1,122 @@
package main
import (
"bytes"
"os"
"path/filepath"
"testing"
)
func TestJPEGFrameSplitterAcrossChunks(t *testing.T) {
splitter := &jpegFrameSplitter{}
var frames [][]byte
emit := func(frame []byte) error {
frames = append(frames, bytes.Clone(frame))
return nil
}
chunks := [][]byte{
{0x01, 0x02, 0xff},
{0xd8, 0x10, 0xff},
{0xd9, 0xff, 0xd8, 0x20},
{0x30, 0xff},
{0xd9, 0x03},
}
for _, chunk := range chunks {
if _, err := splitter.push(chunk, emit); err != nil {
t.Fatalf("push: %v", err)
}
}
if err := splitter.finish(); err != nil {
t.Fatalf("finish: %v", err)
}
want := [][]byte{
{0xff, 0xd8, 0x10, 0xff, 0xd9},
{0xff, 0xd8, 0x20, 0x30, 0xff, 0xd9},
}
if len(frames) != len(want) {
t.Fatalf("got %d frames, want %d", len(frames), len(want))
}
for i := range want {
if !bytes.Equal(frames[i], want[i]) {
t.Errorf("frame %d = %x, want %x", i, frames[i], want[i])
}
}
}
func TestJPEGFrameSplitterRejectsTruncatedFrame(t *testing.T) {
splitter := &jpegFrameSplitter{}
if _, err := splitter.push([]byte{0xff, 0xd8, 0x01}, func([]byte) error { return nil }); err != nil {
t.Fatalf("push: %v", err)
}
if err := splitter.finish(); err == nil {
t.Fatal("finish accepted a truncated JPEG")
}
}
func TestBoundedLog(t *testing.T) {
log := &boundedLog{limit: 4}
if n, err := log.Write([]byte("abcdefgh")); err != nil || n != 8 {
t.Fatalf("Write = %d, %v", n, err)
}
if got := log.String(); got != "abcd" {
t.Fatalf("String = %q, want %q", got, "abcd")
}
}
func TestStreamingTSFCommitAndAbort(t *testing.T) {
dir := t.TempDir()
output := filepath.Join(dir, "frames.tsf")
stream, err := newStreamingTSF(output, 24)
if err != nil {
t.Fatalf("newStreamingTSF: %v", err)
}
for _, frame := range [][]byte{{1, 2, 3}, {4, 5}} {
if err := stream.addFrame(frame); err != nil {
t.Fatalf("addFrame: %v", err)
}
}
if err := stream.commit(output); err != nil {
t.Fatalf("commit: %v", err)
}
stream.abort()
got, err := loadTSF(output)
if err != nil {
t.Fatalf("loadTSF: %v", err)
}
defer func() { _ = got.Close() }()
if got.FPS != 24 || len(got.ColorFrames) != 2 || !bytes.Equal(got.ColorFrames[1], []byte{4, 5}) {
t.Fatalf("unexpected streamed TSF: %+v", got)
}
original := []byte("existing destination")
if err := os.WriteFile(output, original, 0o644); err != nil {
t.Fatalf("WriteFile: %v", err)
}
failed, err := newStreamingTSF(output, 24)
if err != nil {
t.Fatalf("newStreamingTSF: %v", err)
}
if err := failed.addFrame([]byte{9}); err != nil {
t.Fatalf("addFrame: %v", err)
}
failed.abort()
contents, err := os.ReadFile(output)
if err != nil {
t.Fatalf("ReadFile: %v", err)
}
if !bytes.Equal(contents, original) {
t.Fatalf("destination changed after abort: %q", contents)
}
matches, err := filepath.Glob(filepath.Join(dir, ".frames.tsf-*.tmp"))
if err != nil {
t.Fatalf("Glob: %v", err)
}
if len(matches) != 0 {
t.Fatalf("temporary files remain after abort: %v", matches)
}
}