diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 591c713..7b9b658 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -21,10 +21,10 @@ jobs: go-version-file: go.mod - name: Format check - run: test -z "$(gofmt -l src)" + run: test -z "$(gofmt -l .)" - name: Vet - run: go vet ./src/ + run: go vet ./... - name: Staticcheck uses: dominikh/staticcheck-action@v1 @@ -33,14 +33,12 @@ jobs: - name: golangci-lint uses: golangci/golangci-lint-action@v7 - with: - working-directory: src - name: Build - run: go build -o trollssh ./src + run: go build -o trollssh ./cmd/trollssh - name: Unit tests - run: go test -race ./src/ + run: go test -race ./... docker: needs: test diff --git a/Dockerfile b/Dockerfile index 4d57865..c003ab2 100644 --- a/Dockerfile +++ b/Dockerfile @@ -2,8 +2,9 @@ FROM golang:1.25-alpine AS build WORKDIR /src COPY go.mod go.sum ./ RUN go mod download -COPY src ./src -RUN CGO_ENABLED=0 go build -ldflags="-s -w" -o /trollssh ./src +COPY cmd ./cmd +COPY internal ./internal +RUN CGO_ENABLED=0 go build -ldflags="-s -w" -o /trollssh ./cmd/trollssh FROM alpine:3.22 WORKDIR /home/app diff --git a/README.md b/README.md index 59f2d98..90a16b4 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ Generate a frame set from a video through container image ```sh docker run --rm -v ./video.mp4:/home/app/video.mp4 -v ./frames:/home/app/frames \ - ghcr.io/yuzuzensai/trollssh:v1.0.1 trollssh --generate --video video.mp4 --resolution 512 + ghcr.io/yuzuzensai/trollssh:v1.1.1 trollssh --generate --video video.mp4 --resolution 512 ``` This writes `frames/.tsf`, a simple container of color JPEG frames plus @@ -81,25 +81,29 @@ Optional text files in `data/` (created next to the binary): Requirements: Go 1.25+ and `ffmpeg` / `ffprobe` on `PATH` (only for `--generate`). +The entry point lives in [`cmd/trollssh`](cmd/trollssh); the packages it wires +together are under [`internal/`](internal) (`config`, `logx`, `render`, `tsf`, +`sshserver`). + ```sh -go run ./src --generate --video video.mp4 --resolution 512 -go run ./src +go run ./cmd/trollssh --generate --video video.mp4 --resolution 512 +go run ./cmd/trollssh ``` Or build a binary: ```sh -go build -o trollssh ./src +go build -o trollssh ./cmd/trollssh ./trollssh ``` CI runs the following checks on every push: ```sh -gofmt -l ./src # format -go vet ./src/... # vet -golangci-lint run ./src/... # lint, see https://golangci-lint.run -go test ./src/... # tests +gofmt -l . # format +go vet ./... # vet +golangci-lint run ./... # lint, see https://golangci-lint.run +go test ./... # tests ``` To run them automatically before each commit, install diff --git a/src/main.go b/cmd/trollssh/main.go similarity index 74% rename from src/main.go rename to cmd/trollssh/main.go index 287ae19..1e2460f 100644 --- a/src/main.go +++ b/cmd/trollssh/main.go @@ -15,6 +15,11 @@ import ( "time" "github.com/joho/godotenv" + + "github.com/YuzuZensai/TrollSSH/internal/config" + "github.com/YuzuZensai/TrollSSH/internal/logx" + "github.com/YuzuZensai/TrollSSH/internal/sshserver" + "github.com/YuzuZensai/TrollSSH/internal/tsf" ) type cliArgs struct { @@ -47,7 +52,7 @@ func parseArgs(argv []string) cliArgs { } func fail(message string) { - logError(message) + logx.Error(message) os.Exit(1) } @@ -77,15 +82,15 @@ func generateFrames(framesDir, videoArg string, resolution int) { base := strings.TrimSuffix(filepath.Base(videoPath), filepath.Ext(videoPath)) output := filepath.Join(framesDir, base+".tsf") - logInfo(fmt.Sprintf("Generating frames from %q -> %s", videoPath, output)) - if err := processVideo(videoPath, output, resolution); err != nil { + logx.Info(fmt.Sprintf("Generating frames from %q -> %s", videoPath, output)) + if err := tsf.ProcessVideo(videoPath, output, resolution); err != nil { fail(fmt.Sprintf("Failed to generate frames from %q: %s", videoPath, err.Error())) } } const frameDataWarnBytes = 2 << 30 -func loadAllFrames(framesDir string) []*FramesContainer { +func loadAllFrames(framesDir string) []*tsf.FramesContainer { entries, err := os.ReadDir(framesDir) var files []string if err == nil { @@ -112,17 +117,17 @@ func loadAllFrames(framesDir string) []*FramesContainer { totalBytes += info.Size() } if totalBytes > frameDataWarnBytes { - logWarn(fmt.Sprintf( + logx.Warn(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))) + logx.Info(fmt.Sprintf("Frame data: %.1f MB", float64(totalBytes)/(1<<20))) } concurrency := min(len(files), max(1, min(runtime.NumCPU(), 4))) - results := make([]*FramesContainer, len(files)) + results := make([]*tsf.FramesContainer, len(files)) errs := make([]error, len(files)) var next int var nextMu sync.Mutex @@ -145,14 +150,14 @@ func loadAllFrames(framesDir string) []*FramesContainer { errs[i] = err return } - logInfo(fmt.Sprintf("Loading %s (%.1f MB)...", file, float64(info.Size())/1024/1024)) - data, err := loadTSF(filePath) + logx.Info(fmt.Sprintf("Loading %s (%.1f MB)...", file, float64(info.Size())/1024/1024)) + data, err := tsf.Load(filePath) if err != nil { errs[i] = err return } data.Name = file - logInfo(fmt.Sprintf(" %s: %d frames @ %gfps", file, len(data.ColorFrames), data.FPS)) + logx.Info(fmt.Sprintf(" %s: %d frames @ %gfps", file, len(data.ColorFrames), data.FPS)) results[i] = data } } @@ -189,17 +194,17 @@ func applyMemoryLimit() { } limit := n * 9 / 10 debug.SetMemoryLimit(limit) - logInfo(fmt.Sprintf("Memory limit set to %d MB (90%% of cgroup limit)", limit>>20)) + logx.Info(fmt.Sprintf("Memory limit set to %d MB (90%% of cgroup limit)", limit>>20)) return } } func main() { _ = godotenv.Load() - logThreshold = resolveThreshold() + logx.SetThreshold(logx.ResolveThreshold()) applyMemoryLimit() - config := loadConfig() + cfg := config.Load() args := parseArgs(os.Args[1:]) cwd, err := os.Getwd() @@ -219,25 +224,25 @@ func main() { } var bannerText, fakeLoginText, goodbyeText *string - if text, ok := loadOptionalTextFile(filepath.Join(dataDir, "banner.txt")); ok { + if text, ok := config.LoadOptionalTextFile(filepath.Join(dataDir, "banner.txt")); ok { bannerText = &text } - if text, ok := loadOptionalTextFile(filepath.Join(dataDir, "fakelogin.txt")); ok { + if text, ok := config.LoadOptionalTextFile(filepath.Join(dataDir, "fakelogin.txt")); ok { fakeLoginText = &text } - if text, ok := loadOptionalTextFile(filepath.Join(dataDir, "goodbye.txt")); ok { + if text, ok := config.LoadOptionalTextFile(filepath.Join(dataDir, "goodbye.txt")); ok { goodbyeText = &text } - hostKeys, err := ensureHostKeys(dataDir) + hostKeys, err := sshserver.EnsureHostKeys(dataDir) if err != nil { fail(err.Error()) } videoSets := loadAllFrames(framesDir) - logInfo(fmt.Sprintf("Loaded %d frame set(s)", len(videoSets))) + logx.Info(fmt.Sprintf("Loaded %d frame set(s)", len(videoSets))) - server := createServer(ServerDeps{ - Config: config, + server := sshserver.New(sshserver.ServerDeps{ + Config: cfg, HostKeys: hostKeys, BannerText: bannerText, FakeLoginText: fakeLoginText, @@ -250,14 +255,14 @@ func main() { signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) go func() { sig := <-sigCh - logInfo(fmt.Sprintf("Received %s, shutting down...", sig)) + logx.Info(fmt.Sprintf("Received %s, shutting down...", sig)) forceExit := time.AfterFunc(5*time.Second, func() { os.Exit(0) }) server.Close() forceExit.Stop() }() - if err := server.Listen(config.Host, config.Port); err != nil { - logError("Server error:", err.Error()) + if err := server.Listen(cfg.Host, cfg.Port); err != nil { + logx.Error("Server error:", err.Error()) os.Exit(1) } } diff --git a/docker-compose.yaml b/docker-compose.yaml index 6fa290d..8a9fed0 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -1,6 +1,6 @@ services: trollssh: - image: ghcr.io/yuzuzensai/trollssh:v1.0.1 + image: ghcr.io/yuzuzensai/trollssh:v1.1.1 container_name: trollssh restart: unless-stopped mem_limit: ${MEMORY_LIMIT:-2g} diff --git a/src/config.go b/internal/config/config.go similarity index 93% rename from src/config.go rename to internal/config/config.go index 6283e5b..7629b54 100644 --- a/src/config.go +++ b/internal/config/config.go @@ -1,4 +1,4 @@ -package main +package config import ( "fmt" @@ -6,6 +6,8 @@ import ( "strconv" "strings" "time" + + "github.com/YuzuZensai/TrollSSH/internal/logx" ) type PlaybackMode string @@ -39,7 +41,7 @@ type Config struct { } func warnInvalid(name, value string, fallback any) { - logWarn(fmt.Sprintf("Invalid %s=%q, using default %v", name, sanitize(value), fallback)) + logx.Warn(fmt.Sprintf("Invalid %s=%q, using default %v", name, logx.Sanitize(value), fallback)) } func envString(name, fallback string) string { @@ -113,7 +115,7 @@ func envPlaybackMode(name string, fallback PlaybackMode) PlaybackMode { return fallback } -func loadConfig() Config { +func Load() Config { const maxInt = int(^uint(0) >> 1) return Config{ Host: envString("HOST", "0.0.0.0"), @@ -139,7 +141,7 @@ func loadConfig() Config { } } -func loadOptionalTextFile(filePath string) (string, bool) { +func LoadOptionalTextFile(filePath string) (string, bool) { data, err := os.ReadFile(filePath) if err != nil { return "", false diff --git a/internal/config/config_test.go b/internal/config/config_test.go new file mode 100644 index 0000000..0fc05db --- /dev/null +++ b/internal/config/config_test.go @@ -0,0 +1,79 @@ +package config + +import ( + "testing" + "time" +) + +func TestLoadConfigDefaults(t *testing.T) { + t.Setenv("HOST", "") + t.Setenv("PORT", "") + t.Setenv("PLAYBACK_MODE", "") + t.Setenv("LOGIN_DELAY", "") + cfg := Load() + if cfg.Host != "0.0.0.0" { + t.Errorf("host = %q", cfg.Host) + } + if cfg.Port != 22 { + t.Errorf("port = %d", cfg.Port) + } + if cfg.PlaybackMode != PlaybackLoop { + t.Errorf("playbackMode = %q", cfg.PlaybackMode) + } + if cfg.Charset != "detailed" { + t.Errorf("charset = %q", cfg.Charset) + } + if cfg.LoginDelay != 1500*time.Millisecond { + t.Errorf("loginDelay = %v", cfg.LoginDelay) + } +} + +func TestLoadConfigClamping(t *testing.T) { + t.Setenv("PORT", "999999") + t.Setenv("BRIGHTNESS_THRESHOLD", "-5") + cfg := Load() + if cfg.Port != 65535 { + t.Errorf("port clamp = %d", cfg.Port) + } + if cfg.BrightnessThreshold != 0 { + t.Errorf("brightness clamp = %d", cfg.BrightnessThreshold) + } +} + +func TestLoadConfigInvalidFallsBack(t *testing.T) { + t.Setenv("PORT", "not-a-number") + t.Setenv("INVERT", "yes-please") + t.Setenv("PLAYBACK_MODE", "shuffle") + cfg := Load() + if cfg.Port != 22 { + t.Errorf("port = %d, want default 22", cfg.Port) + } + if cfg.Invert { + t.Error("invert should fall back to false") + } + if cfg.PlaybackMode != PlaybackLoop { + t.Errorf("playbackMode = %q, want default loop", cfg.PlaybackMode) + } +} + +func TestEnvDurationMs(t *testing.T) { + t.Setenv("D", "250") + if got := envDurationMs("D", time.Second); got != 250*time.Millisecond { + t.Errorf("250 = %v, want 250ms", got) + } + t.Setenv("D", "-10") + if got := envDurationMs("D", time.Second); got != 0 { + t.Errorf("negative = %v, want 0", got) + } + t.Setenv("D", "banana") + if got := envDurationMs("D", time.Second); got != time.Second { + t.Errorf("invalid = %v, want fallback 1s", got) + } +} + +func TestPlaybackModeRandom(t *testing.T) { + t.Setenv("PLAYBACK_MODE", "RaNdOm") + if Load().PlaybackMode != PlaybackRandom { + t.Error("expected random") + } +} diff --git a/src/logger.go b/internal/logx/logx.go similarity index 55% rename from src/logger.go rename to internal/logx/logx.go index b976eae..deabb36 100644 --- a/src/logger.go +++ b/internal/logx/logx.go @@ -1,4 +1,4 @@ -package main +package logx import ( "encoding/json" @@ -8,35 +8,37 @@ import ( "time" ) -type logLevel int +type Level int const ( - levelDebug logLevel = 10 - levelInfo logLevel = 20 - levelWarn logLevel = 30 - levelError logLevel = 40 + LevelDebug Level = 10 + LevelInfo Level = 20 + LevelWarn Level = 30 + LevelError Level = 40 ) -var logThreshold = resolveThreshold() +var threshold = ResolveThreshold() -func resolveThreshold() logLevel { +func ResolveThreshold() Level { switch strings.ToLower(strings.TrimSpace(os.Getenv("LOG_LEVEL"))) { case "debug": - return levelDebug + return LevelDebug case "warn": - return levelWarn + return LevelWarn case "error": - return levelError + return LevelError default: - return levelInfo + return LevelInfo } } -func sanitize(value any) string { - return sanitizeN(value, 200) +func SetThreshold(level Level) { threshold = level } + +func Sanitize(value any) string { + return SanitizeN(value, 200) } -func sanitizeN(value any, maxLength int) string { +func SanitizeN(value any, maxLength int) string { var str string switch v := value.(type) { case nil: @@ -63,8 +65,8 @@ func sanitizeN(value any, maxLength int) string { return b.String() } -func emit(level logLevel, name string, stream *os.File, args []any) { - if level < logThreshold { +func emit(level Level, name string, stream *os.File, args []any) { + if level < threshold { return } parts := make([]string, len(args)) @@ -81,7 +83,7 @@ func emit(level logLevel, name string, stream *os.File, args []any) { _, _ = fmt.Fprintf(stream, "[%s] %-5s %s\n", ts, strings.ToUpper(name), strings.Join(parts, " ")) } -func logDebug(args ...any) { emit(levelDebug, "debug", os.Stdout, args) } -func logInfo(args ...any) { emit(levelInfo, "info", os.Stdout, args) } -func logWarn(args ...any) { emit(levelWarn, "warn", os.Stderr, args) } -func logError(args ...any) { emit(levelError, "error", os.Stderr, args) } +func Debug(args ...any) { emit(LevelDebug, "debug", os.Stdout, args) } +func Info(args ...any) { emit(LevelInfo, "info", os.Stdout, args) } +func Warn(args ...any) { emit(LevelWarn, "warn", os.Stderr, args) } +func Error(args ...any) { emit(LevelError, "error", os.Stderr, args) } diff --git a/internal/logx/logx_test.go b/internal/logx/logx_test.go new file mode 100644 index 0000000..e8151ff --- /dev/null +++ b/internal/logx/logx_test.go @@ -0,0 +1,10 @@ +package logx + +import "testing" + +func TestSanitizeNStopsAtLimit(t *testing.T) { + input := "ab\x00cdefghijklmnopqrstuvwxyz" + if got := SanitizeN(input, 4); got != "ab�c…" { + t.Fatalf("SanitizeN = %q", got) + } +} diff --git a/internal/render/cache.go b/internal/render/cache.go new file mode 100644 index 0000000..7106de1 --- /dev/null +++ b/internal/render/cache.go @@ -0,0 +1,232 @@ +package render + +import ( + "bytes" + "container/list" + "sync" + "sync/atomic" + "time" +) + +type cacheKey struct { + setID int + index int + width int + height int + keepAspectRatio bool + tier ColorTier +} + +type cacheShard struct { + mu sync.Mutex + maxBytes int64 + size int64 + entries map[cacheKey]*list.Element + order *list.List +} + +type Cache struct { + shards []cacheShard + size atomic.Int64 + hits atomic.Uint64 + misses atomic.Uint64 + evictions atomic.Uint64 + rejections atomic.Uint64 + renders atomic.Uint64 + renderNs atomic.Uint64 +} + +type cacheEntry struct { + key cacheKey + ascii []byte + cost int64 +} + +func entryCost(_ cacheKey, ascii []byte) int64 { + return int64(cap(ascii)) + 160 +} + +func NewCache(maxBytes int64) *Cache { + if maxBytes <= 0 { + return nil + } + shardCount := int(min(int64(16), max(int64(1), maxBytes/(1<<20)))) + cache := &Cache{shards: make([]cacheShard, shardCount)} + for i := range cache.shards { + cache.shards[i] = cacheShard{ + maxBytes: maxBytes / int64(shardCount), + entries: make(map[cacheKey]*list.Element), + order: list.New(), + } + } + return cache +} + +func (c *Cache) shard(key cacheKey) *cacheShard { + 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 *Cache) get(key cacheKey) ([]byte, bool) { + if c == nil { + return nil, false + } + shard := c.shard(key) + shard.mu.Lock() + defer shard.mu.Unlock() + el, ok := shard.entries[key] + if !ok { + c.misses.Add(1) + return nil, false + } + c.hits.Add(1) + shard.order.MoveToBack(el) + return el.Value.(*cacheEntry).ascii, true +} + +func (c *Cache) put(key cacheKey, ascii []byte) { + if c == nil { + return + } + shard := c.shard(key) + cost := entryCost(key, ascii) + if cost > shard.maxBytes { + c.rejections.Add(1) + return + } + shard.mu.Lock() + defer shard.mu.Unlock() + if _, ok := shard.entries[key]; ok { + return + } + shard.entries[key] = shard.order.PushBack(&cacheEntry{key: key, ascii: ascii, cost: cost}) + shard.size += cost + c.size.Add(cost) + for shard.size > shard.maxBytes { + oldest := shard.order.Front() + shard.order.Remove(oldest) + evicted := oldest.Value.(*cacheEntry) + delete(shard.entries, evicted.key) + shard.size -= evicted.cost + c.size.Add(-evicted.cost) + c.evictions.Add(1) + } +} + +type CacheStats struct { + SizeBytes int64 + Hits uint64 + Misses uint64 + Evictions uint64 + Rejections uint64 + Renders uint64 + RenderTime time.Duration +} + +func (c *Cache) Stats() CacheStats { + if c == nil { + return CacheStats{} + } + return CacheStats{ + 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()), + } +} + +type Renderer struct { + setID int + colorFrames [][]byte + options Options + rampLUT *[101][]byte + cache *Cache + + inflightMu sync.Mutex + inflight map[cacheKey]*renderCall +} + +type renderCall struct { + done chan struct{} + value []byte + err error +} + +func NewRenderer(setID int, colorFrames [][]byte, options Options, cache *Cache) *Renderer { + ramp := []rune(resolveCharset(options.Charset)) + return &Renderer{ + setID: setID, + colorFrames: colorFrames, + options: options, + rampLUT: buildRampLUT(ramp, options), + cache: cache, + inflight: make(map[cacheKey]*renderCall), + } +} + +func (r *Renderer) Render(index, width, height int, keepAspectRatio bool, tier ColorTier) ([]byte, error) { + key := cacheKey{r.setID, index, width, height, keepAspectRatio, tier} + if ascii, ok := r.cache.get(key); ok { + return ascii, nil + } + + r.inflightMu.Lock() + if call, ok := r.inflight[key]; ok { + r.inflightMu.Unlock() + <-call.done + return call.value, call.err + } + call := &renderCall{done: make(chan struct{})} + r.inflight[key] = call + r.inflightMu.Unlock() + defer func() { + r.inflightMu.Lock() + delete(r.inflight, key) + r.inflightMu.Unlock() + close(call.done) + }() + + if ascii, ok := r.cache.get(key); ok { + call.value = ascii + return ascii, nil + } + + started := time.Now() + pix := getPixBuf(4 * width * height) + img, err := resizeFrame(r.colorFrames[index], pix, width, height, keepAspectRatio) + if err != nil { + putPixBuf(pix) + call.err = err + return nil, err + } + var ascii []byte + if tier == ColorTierNone { + ascii = frameToAscii(img, r.rampLUT) + } else { + ascii = frameToAnsi(img, r.rampLUT, tier) + } + putPixBuf(pix) + + if r.cache != nil && cap(ascii) > len(ascii)+len(ascii)/4 { + ascii = bytes.Clone(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 +} diff --git a/internal/render/render.go b/internal/render/render.go new file mode 100644 index 0000000..1780a9f --- /dev/null +++ b/internal/render/render.go @@ -0,0 +1,250 @@ +package render + +import ( + "bytes" + "image" + "image/color" + "image/jpeg" + "strconv" + "strings" + "sync" + "unicode/utf8" + + "golang.org/x/image/draw" +) + +type ColorTier int + +const ( + ColorTierNone ColorTier = iota + ColorTier256 + ColorTierTrueColor +) + +func DetectColorTier(term string) ColorTier { + t := strings.ToLower(strings.TrimSpace(term)) + switch t { + case "", "dumb", "vt52", "vt100", "vt102", "vt220", "ansi", "linux", "cons25", "cygwin": + return ColorTierNone + } + if strings.Contains(t, "direct") || strings.Contains(t, "truecolor") { + return ColorTierTrueColor + } + if strings.Contains(t, "256color") { + return ColorTier256 + } + if strings.HasPrefix(t, "screen") || strings.HasPrefix(t, "tmux") { + return ColorTier256 + } + return ColorTierTrueColor +} + +var charsetPresets = map[string]string{ + "detailed": " .'`^\",:;Il!i><~+_-?][}{1)(|/tfjrxnuvczXYUJCLQ0OZmwqpdbkhao*#MW&8%B@$", + "standard": " .:-=+*#%@", + "simple": " .:oO#@", + "blocks": " ░▒▓█", +} + +func resolveCharset(charset string) string { + if charset == "" { + return charsetPresets["detailed"] + } + if preset, ok := charsetPresets[strings.ToLower(charset)]; ok { + return preset + } + return charset +} + +const maxPooledBuffer = 4 << 20 + +var pixPool sync.Pool + +func getPixBuf(n int) []byte { + if v := pixPool.Get(); v != nil { + if b := *v.(*[]byte); cap(b) >= n { + return b[:n] + } + } + return make([]byte, n) +} + +func putPixBuf(b []byte) { + if cap(b) <= maxPooledBuffer { + pixPool.Put(&b) + } +} + +var outPool sync.Pool + +func getOutBuf(capacity int) []byte { + if v := outPool.Get(); v != nil { + if b := *v.(*[]byte); cap(b) >= capacity { + return b[:0] + } + } + return make([]byte, 0, capacity) +} + +func putOutBuf(b []byte) { + if cap(b) <= maxPooledBuffer { + outPool.Put(&b) + } +} + +func resizeFrame(frame, pix []byte, width, height int, keepAspectRatio bool) (*image.RGBA, error) { + src, err := jpeg.Decode(bytes.NewReader(frame)) + if err != nil { + return nil, err + } + rect := image.Rect(0, 0, width, height) + + dst := &image.RGBA{Pix: pix[:4*width*height], Stride: 4 * width, Rect: rect} + if keepAspectRatio { + draw.Draw(dst, dst.Bounds(), image.NewUniform(color.Black), image.Point{}, draw.Src) + sb := src.Bounds() + sw, sh := sb.Dx(), sb.Dy() + scale := min(float64(width)/float64(sw), float64(height)/float64(sh)) + tw := max(1, int(float64(sw)*scale)) + th := max(1, int(float64(sh)*scale)) + x0 := (width - tw) / 2 + y0 := (height - th) / 2 + draw.ApproxBiLinear.Scale(dst, image.Rect(x0, y0, x0+tw, y0+th), src, sb, draw.Src, nil) + } else { + draw.ApproxBiLinear.Scale(dst, dst.Bounds(), src, src.Bounds(), draw.Src, nil) + } + return dst, nil +} + +type Options struct { + BrightnessThreshold int + Charset string + Invert bool +} + +func buildRampLUT(ramp []rune, options Options) *[101][]byte { + var lut [101][]byte + for b := range lut { + index := rampIndex(b, options.BrightnessThreshold, len(ramp), options.Invert) + lut[b] = utf8.AppendRune(nil, ramp[index]) + } + return &lut +} + +func rampIndex(brightness, threshold, total int, invert bool) int { + var index int + if brightness < threshold { + index = 0 + } else { + index = brightness * total / 100 + if index > total-1 { + index = total - 1 + } + } + if invert { + index = total - 1 - index + } + return index +} + +func frameToAscii(img *image.RGBA, rampLUT *[101][]byte) []byte { + pix := img.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 { + brightness := (int(pix[o])*299 + int(pix[o+1])*587 + int(pix[o+2])*114) / 255 / 10 + buf = append(buf, rampLUT[brightness]...) + } + output := bytes.Clone(buf) + putOutBuf(buf) + return output +} + +const ansiReset = "\x1b[0m" + +var ansi256Levels = [6]int{0, 95, 135, 175, 215, 255} + +var decimal = func() (t [256]string) { + for i := range t { + t[i] = strconv.Itoa(i) + } + return +}() + +var ansi256Cube = func() (t [256]uint8) { + for v := range t { + best, bestDist := 0, 1<<30 + for i, l := range ansi256Levels { + d := v - l + if d < 0 { + d = -d + } + if d < bestDist { + bestDist, best = d, i + } + } + t[v] = uint8(best) + } + return +}() + +func quantize256(r, g, b uint8) int { + return 16 + 36*int(ansi256Cube[r]) + 6*int(ansi256Cube[g]) + int(ansi256Cube[b]) +} + +func appendColor(buf []byte, r, g, b uint8, tier ColorTier) []byte { + if tier == ColorTierTrueColor { + buf = append(buf, "\x1b[38;2;"...) + buf = append(buf, decimal[r]...) + buf = append(buf, ';') + buf = append(buf, decimal[g]...) + buf = append(buf, ';') + buf = append(buf, decimal[b]...) + } else { + buf = append(buf, "\x1b[38;5;"...) + buf = append(buf, decimal[quantize256(r, g, b)]...) + } + return append(buf, 'm') +} + +func frameToAnsi(img *image.RGBA, rampLUT *[101][]byte, tier ColorTier) []byte { + bounds := img.Bounds() + bytesPerCell := 11 + if tier == ColorTierTrueColor { + bytesPerCell = 16 + } + buf := getOutBuf(bounds.Dx() * bounds.Dy() * bytesPerCell) + var lastR, lastG, lastB uint8 + last256 := -1 + first := true + 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++ { + 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 + 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) + lastR, lastG, lastB = r, g, bl + first = false + } + buf = append(buf, rampLUT[brightness]...) + o += 4 + } + if y < bounds.Max.Y-1 { + buf = append(buf, "\r\n"...) + } + } + buf = append(buf, ansiReset...) + output := bytes.Clone(buf) + putOutBuf(buf) + return output +} diff --git a/internal/render/render_bench_test.go b/internal/render/render_bench_test.go new file mode 100644 index 0000000..cdfa5da --- /dev/null +++ b/internal/render/render_bench_test.go @@ -0,0 +1,69 @@ +package render + +import ( + "path/filepath" + "sync/atomic" + "testing" + + "github.com/YuzuZensai/TrollSSH/internal/tsf" +) + +func loadBenchSet(b *testing.B) *tsf.FramesContainer { + b.Helper() + matches, _ := filepath.Glob("../../frames/*.tsf") + if len(matches) == 0 { + b.Skip("no .tsf frame set in ../../frames") + } + fc, err := tsf.Load(matches[0]) + if err != nil { + b.Skip("failed to load frame set:", err) + } + b.Cleanup(func() { _ = fc.Close() }) + return fc +} + +func benchRender(b *testing.B, tier ColorTier, w, h int) { + fc := loadBenchSet(b) + r := NewRenderer(0, fc.ColorFrames, Options{ + BrightnessThreshold: 40, + Charset: "detailed", + }, nil) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if _, err := r.Render(i%len(fc.ColorFrames), w, h, false, tier); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkRenderTrueColor(b *testing.B) { benchRender(b, ColorTierTrueColor, 120, 40) } +func BenchmarkRender256(b *testing.B) { benchRender(b, ColorTier256, 120, 40) } +func BenchmarkRenderGray(b *testing.B) { benchRender(b, ColorTierNone, 120, 40) } +func BenchmarkRenderTrueBig(b *testing.B) { benchRender(b, ColorTierTrueColor, 240, 70) } + +func BenchmarkRenderCachedParallel(b *testing.B) { + fc := loadBenchSet(b) + r := NewRenderer(0, fc.ColorFrames, Options{ + BrightnessThreshold: 40, + Charset: "detailed", + }, NewCache(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()) + } +} diff --git a/internal/render/render_test.go b/internal/render/render_test.go new file mode 100644 index 0000000..71a7b60 --- /dev/null +++ b/internal/render/render_test.go @@ -0,0 +1,194 @@ +package render + +import ( + "bytes" + "image" + "image/jpeg" + "strings" + "sync" + "testing" +) + +func TestResolveCharset(t *testing.T) { + if got := resolveCharset("blocks"); got != " ░▒▓█" { + t.Errorf("blocks preset = %q", got) + } + if got := resolveCharset("XYZ"); got != "XYZ" { + t.Errorf("custom ramp = %q", got) + } + if got := resolveCharset(""); !strings.HasPrefix(got, " .") { + t.Errorf("default = %q", got) + } +} + +func TestFrameToAscii(t *testing.T) { + // Below threshold -> first ramp char; full brightness -> last. + opts := Options{BrightnessThreshold: 40, Charset: "standard"} + ramp := []rune(resolveCharset("standard")) + img := &image.RGBA{ + Pix: []byte{0, 0, 0, 255, 255, 255, 255, 255}, + Stride: 8, Rect: image.Rect(0, 0, 2, 1), + } + out := []rune(string(frameToAscii(img, buildRampLUT(ramp, opts)))) + if out[0] != ramp[0] { + t.Errorf("dark px = %q, want %q", out[0], ramp[0]) + } + if out[1] != ramp[len(ramp)-1] { + t.Errorf("bright px = %q, want %q", out[1], ramp[len(ramp)-1]) + } +} + +func TestFrameToAsciiInvert(t *testing.T) { + opts := Options{BrightnessThreshold: 40, Charset: "standard", Invert: true} + ramp := []rune(resolveCharset("standard")) + img := &image.RGBA{ + Pix: []byte{255, 255, 255, 255}, + Stride: 4, Rect: image.Rect(0, 0, 1, 1), + } + out := []rune(string(frameToAscii(img, buildRampLUT(ramp, opts)))) + if out[0] != ramp[0] { + t.Errorf("inverted bright = %q, want %q", out[0], ramp[0]) + } +} + +func TestRenderConcurrentSameKey(t *testing.T) { + var jpegBuf bytes.Buffer + src := image.NewRGBA(image.Rect(0, 0, 16, 16)) + for i := range src.Pix { + src.Pix[i] = byte(i * 7) + } + if err := jpeg.Encode(&jpegBuf, src, nil); err != nil { + t.Fatal(err) + } + + r := NewRenderer(0, [][]byte{jpegBuf.Bytes()}, Options{ + BrightnessThreshold: 40, + Charset: "standard", + }, NewCache(1<<20)) + + var wg sync.WaitGroup + results := make([][]byte, 32) + for i := range results { + wg.Add(1) + go func(i int) { + defer wg.Done() + ascii, err := r.Render(0, 20, 10, false, ColorTierTrueColor) + if err != nil { + t.Error(err) + return + } + results[i] = ascii + }(i) + } + wg.Wait() + for i, got := range results { + if !bytes.Equal(got, results[0]) { + t.Fatalf("result %d differs from result 0", i) + } + } +} + +func TestRenderCacheEvictsByBytes(t *testing.T) { + key := func(index int) cacheKey { return cacheKey{index: index} } + budget := 3 * entryCost(key(0), bytes.Repeat([]byte("x"), 1000)) + c := NewCache(budget) + for i := range 5 { + c.put(key(i), bytes.Repeat([]byte("x"), 1000)) + } + if c.size.Load() > budget { + t.Errorf("size %d exceeds budget %d", c.size.Load(), budget) + } + if _, ok := c.get(key(0)); ok { + t.Error("oldest entry should have been evicted") + } + if _, ok := c.get(key(4)); !ok { + t.Error("newest entry should be cached") + } +} + +func TestRenderCacheDisabled(t *testing.T) { + c := NewCache(0) + if c != nil { + t.Fatal("zero budget should disable the cache") + } + c.put(cacheKey{}, []byte("v")) + if _, ok := c.get(cacheKey{}); ok { + t.Error("nil cache should never hit") + } +} + +func TestRenderCacheRejectsOversizedEntry(t *testing.T) { + c := NewCache(256) + 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.Load() != 0 { + t.Errorf("size = %d, want 0", c.size.Load()) + } +} + +func TestRenderCacheAccountsRetainedCapacity(t *testing.T) { + cache := NewCache(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 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(" .#"), Options{}), 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(" .#"), Options{}), ColorTierTrueColor) + if count := bytes.Count(output, []byte(ansiReset)); count != 1 { + t.Fatalf("reset count = %d, want 1: %q", count, output) + } +} + +func TestDetectColorTier(t *testing.T) { + cases := map[string]ColorTier{ + "": ColorTierNone, + "dumb": ColorTierNone, + "vt100": ColorTierNone, + "linux": ColorTierNone, + "xterm": ColorTierTrueColor, + "xterm-256color": ColorTier256, + "screen-256color": ColorTier256, + "tmux-256color": ColorTier256, + "xterm-direct": ColorTierTrueColor, + "xterm-kitty": ColorTierTrueColor, + } + for term, want := range cases { + if got := DetectColorTier(term); got != want { + t.Errorf("DetectColorTier(%q) = %d, want %d", term, got, want) + } + } +} + +func TestQuantize256(t *testing.T) { + if got := quantize256(0, 0, 0); got != 16 { + t.Errorf("black = %d, want 16", got) + } + if got := quantize256(255, 255, 255); got != 231 { + t.Errorf("white = %d, want 231", got) + } +} diff --git a/src/hostkeys.go b/internal/sshserver/hostkeys.go similarity index 94% rename from src/hostkeys.go rename to internal/sshserver/hostkeys.go index 90083cf..6965ba4 100644 --- a/src/hostkeys.go +++ b/internal/sshserver/hostkeys.go @@ -1,4 +1,4 @@ -package main +package sshserver import ( "crypto/ed25519" @@ -33,7 +33,7 @@ func generateAndSave(keyPath, keyType string) error { return os.WriteFile(keyPath, pem.EncodeToMemory(block), 0o600) } -func ensureHostKeys(configDir string) ([]ssh.Signer, error) { +func EnsureHostKeys(configDir string) ([]ssh.Signer, error) { keys := []struct{ file, keyType string }{ {"id_rsa", "rsa"}, {"id_ed25519", "ed25519"}, diff --git a/src/server.go b/internal/sshserver/server.go similarity index 85% rename from src/server.go rename to internal/sshserver/server.go index c3ff3ed..3fdc7e6 100644 --- a/src/server.go +++ b/internal/sshserver/server.go @@ -1,4 +1,4 @@ -package main +package sshserver import ( "encoding/binary" @@ -13,6 +13,11 @@ import ( "time" "golang.org/x/crypto/ssh" + + "github.com/YuzuZensai/TrollSSH/internal/config" + "github.com/YuzuZensai/TrollSSH/internal/logx" + "github.com/YuzuZensai/TrollSSH/internal/render" + "github.com/YuzuZensai/TrollSSH/internal/tsf" ) const ( @@ -173,15 +178,15 @@ func (t *SessionTracker) release(conn *ssh.ServerConn) { } type frameSet struct { - data *FramesContainer - renderer *FrameRenderer + data *tsf.FramesContainer + renderer *render.Renderer } type Server struct { - config Config + config config.Config sshConfig *ssh.ServerConfig sets []frameSet - cache *renderCache + cache *render.Cache tracker *ConnectionTracker sessions *SessionTracker fakeLogin *string @@ -195,12 +200,12 @@ type Server struct { } type ServerDeps struct { - Config Config + Config config.Config HostKeys []ssh.Signer BannerText *string FakeLoginText *string GoodbyeText *string - VideoSets []*FramesContainer + VideoSets []*tsf.FramesContainer } func clampTermSize(cols, rows, maxDimension, maxCells, quantum int) (int, int) { @@ -224,30 +229,30 @@ func clampTermSize(cols, rows, maxDimension, maxCells, quantum int) (int, int) { return cols, rows } -func createServer(deps ServerDeps) *Server { - config := deps.Config +func New(deps ServerDeps) *Server { + cfg := deps.Config - cache := newRenderCache(int64(config.RenderCacheMB) << 20) + cache := render.NewCache(int64(cfg.RenderCacheMB) << 20) sets := make([]frameSet, len(deps.VideoSets)) for i, data := range deps.VideoSets { sets[i] = frameSet{ data: data, - renderer: newFrameRenderer(i, data.ColorFrames, asciiOptions{ - brightnessThreshold: config.BrightnessThreshold, - charset: config.Charset, - invert: config.Invert, + renderer: render.NewRenderer(i, data.ColorFrames, render.Options{ + BrightnessThreshold: cfg.BrightnessThreshold, + Charset: cfg.Charset, + Invert: cfg.Invert, }, cache), } } sshConfig := &ssh.ServerConfig{ - MaxAuthTries: config.MaxAuthAttempts, + MaxAuthTries: cfg.MaxAuthAttempts, PasswordCallback: func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) { ip := hostOnly(conn.RemoteAddr().String()) - if config.LogCredentials { - logInfo(fmt.Sprintf( + if cfg.LogCredentials { + logx.Info(fmt.Sprintf( `Auth attempt from %s method=password user="%s" pass="%s"`, - ip, sanitizeN(conn.User(), 128), sanitizeN(string(password), 128), + ip, logx.SanitizeN(conn.User(), 128), logx.SanitizeN(string(password), 128), )) } if conn.User() == "" || len(password) == 0 { @@ -274,7 +279,7 @@ func createServer(deps ServerDeps) *Server { } return &Server{ - config: config, + config: cfg, sshConfig: sshConfig, sets: sets, cache: cache, @@ -307,7 +312,7 @@ func (s *Server) Listen(host string, port int) error { } s.listener = listener s.mu.Unlock() - logInfo(fmt.Sprintf("TrollSSH listening on %s:%d", host, port)) + logx.Info(fmt.Sprintf("TrollSSH listening on %s:%d", host, port)) for { conn, err := listener.Accept() if err != nil { @@ -320,7 +325,7 @@ func (s *Server) Listen(host string, port int) error { activeForIP, total, ok := s.tracker.tryAcquire(ip, s.config.MaxConnections, s.config.MaxTotalConnections) if !ok { _ = conn.Close() - logWarn("Connection rejected (limit reached) from", ip) + logx.Warn("Connection rejected (limit reached) from", ip) continue } s.mu.Lock() @@ -354,9 +359,9 @@ func (s *Server) Close() { _ = conn.Close() } s.connWG.Wait() - stats := s.cache.stats() + stats := s.cache.Stats() if stats.Hits+stats.Misses > 0 { - logInfo(fmt.Sprintf( + logx.Info(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, @@ -364,7 +369,7 @@ func (s *Server) Close() { } for _, set := range s.sets { if err := set.data.Close(); err != nil { - logWarn("Failed to release frame set", set.data.Name, sanitize(err.Error())) + logx.Warn("Failed to release frame set", set.data.Name, logx.Sanitize(err.Error())) } } }) @@ -386,19 +391,19 @@ func (s *Server) handleConn(conn net.Conn, ip string, activeForIP, total int) { sshConn, chans, reqs, err := ssh.NewServerConn(conn, s.sshConfig) if err != nil { if strings.Contains(err.Error(), "i/o timeout") { - logWarn("Handshake timeout for", ip) + logx.Warn("Handshake timeout for", ip) } else { - logWarn(fmt.Sprintf("Client error from %s:", ip), sanitize(err.Error())) + logx.Warn(fmt.Sprintf("Client error from %s:", ip), logx.Sanitize(err.Error())) } _ = conn.Close() return } _ = conn.SetDeadline(time.Time{}) - logDebug("Handshake from", ip) + logx.Debug("Handshake from", ip) defer func() { _ = sshConn.Close() }() setIndex := rand.Intn(len(s.sets)) - logInfo(fmt.Sprintf( + logx.Info(fmt.Sprintf( "New connection from %s (ip=%d, total=%d) -> playing %q", ip, activeForIP, total, s.sets[setIndex].data.Name, )) @@ -434,7 +439,7 @@ func (s *Server) handleConn(conn net.Conn, ip string, activeForIP, total int) { } _ = sshConn.Close() sessionWG.Wait() - logInfo("Client closed connection from", ip) + logx.Info("Client closed connection from", ip) } type termSize struct { @@ -501,9 +506,9 @@ func (s *Server) handleSession( defer func() { _ = channel.Close() }() size := &termSize{} size.set(80, 24, s.config.MaxDimension, s.config.MaxTerminalCells, true) - tier := colorTierTrueColor + tier := render.ColorTierTrueColor if s.config.ForceGrayscale { - tier = colorTierNone + tier = render.ColorTierNone } started := false @@ -511,16 +516,16 @@ func (s *Server) handleSession( for req := range requests { switch req.Type { case "pty-req": - logDebug("Opening pty for session", ip) + logx.Debug("Opening pty for session", ip) if cols, rows, ok := parseDims(req.Payload); ok { size.set(cols, rows, s.config.MaxDimension, s.config.MaxTerminalCells, true) } if term, ok := parsePtyTerm(req.Payload); ok { - tier = detectColorTier(term) + tier = render.DetectColorTier(term) if s.config.ForceGrayscale { - tier = colorTierNone + tier = render.ColorTierNone } - logDebug(fmt.Sprintf("Client %s TERM=%q -> color tier %d", ip, sanitizeN(term, 64), tier)) + logx.Debug(fmt.Sprintf("Client %s TERM=%q -> color tier %d", ip, logx.SanitizeN(term, 64), tier)) } _ = req.Reply(true, nil) case "window-change": @@ -540,25 +545,25 @@ func (s *Server) handleSession( command = string(req.Payload[4 : 4+n]) } } - logInfo(fmt.Sprintf("Client %s attempted exec: %q", ip, sanitizeN(command, 512))) + logx.Info(fmt.Sprintf("Client %s attempted exec: %q", ip, logx.SanitizeN(command, 512))) _ = req.Reply(true, nil) if !started { started = true playDone = make(chan struct{}) playTier := tier - go func(tier colorTier) { + go func(tier render.ColorTier) { defer close(playDone) s.playVideo(sshConn, channel, size, ip, initialSetIndex, false, tier) }(playTier) } case "shell": - logDebug("Opening shell for session", ip) + logx.Debug("Opening shell for session", ip) _ = req.Reply(true, nil) if !started { started = true playDone = make(chan struct{}) playTier := tier - go func(tier colorTier) { + go func(tier render.ColorTier) { defer close(playDone) s.playVideo(sshConn, channel, size, ip, initialSetIndex, false, tier) }(playTier) @@ -593,13 +598,13 @@ func (s *Server) playVideo( ip string, setIndex int, keepAspectRatio bool, - tier colorTier, + tier render.ColorTier, ) { - config := s.config + cfg := s.config current := s.sets[setIndex] w, h := size.get() - logDebug(fmt.Sprintf("Terminal size %dx%d for %s", w, h, ip)) + logx.Debug(fmt.Sprintf("Terminal size %dx%d for %s", w, h, ip)) defer func() { _ = writePartsWithTimeout(sshConn, channel, outputStallTimeout, showCursor) @@ -629,7 +634,7 @@ func (s *Server) playVideo( closeSession() return } - if !config.AllowUserControl { + if !cfg.AllowUserControl { continue } str := string(buf[:n]) @@ -643,7 +648,7 @@ func (s *Server) playVideo( continue } now := time.Now() - if now.Sub(lastSwitch) < config.SwitchDebounce { + if now.Sub(lastSwitch) < cfg.SwitchDebounce { continue } lastSwitch = now @@ -654,7 +659,7 @@ func (s *Server) playVideo( } }() - loginTimer := time.NewTimer(config.LoginDelay) + loginTimer := time.NewTimer(cfg.LoginDelay) select { case <-loginTimer.C: case <-done: @@ -691,14 +696,14 @@ func (s *Server) playVideo( current = s.sets[setIndex] currentFrame = 0 lastW, lastH = 0, 0 - logDebug(fmt.Sprintf("%s switched to %q", ip, current.data.Name)) + logx.Debug(fmt.Sprintf("%s switched to %q", ip, current.data.Name)) ticker.Reset(frameInterval()) case <-ticker.C: w, h := size.get() - ascii, err := current.renderer.render(currentFrame, w, h, keepAspectRatio, tier) + ascii, err := current.renderer.Render(currentFrame, w, h, keepAspectRatio, tier) if err != nil { - logError("Render error for", ip, sanitize(err.Error())) + logx.Error("Render error for", ip, logx.Sanitize(err.Error())) _ = sshConn.Close() return } @@ -722,7 +727,7 @@ func (s *Server) playVideo( currentFrame = 0 loopCount++ - if config.MaxLoop > 0 && loopCount >= config.MaxLoop { + if cfg.MaxLoop > 0 && loopCount >= cfg.MaxLoop { if err := writePartsWithTimeout( sshConn, channel, outputStallTimeout, showCursor, clearScreen, ); err != nil { @@ -744,26 +749,26 @@ func (s *Server) playVideo( } return } - logInfo("Playback finished, closing session", ip) + logx.Info("Playback finished, closing session", ip) _ = channel.Close() _ = sshConn.Close() return } - if config.PlaybackMode == PlaybackRandom { + if cfg.PlaybackMode == config.PlaybackRandom { setIndex = s.pickNextSetIndex(setIndex) current = s.sets[setIndex] - logInfo(fmt.Sprintf( + logx.Info(fmt.Sprintf( "Playthrough done for %s, switching to %q", ip, current.data.Name, )) ticker.Reset(frameInterval()) - } else if config.MaxLoop > 0 { - logInfo(fmt.Sprintf( + } else if cfg.MaxLoop > 0 { + logx.Info(fmt.Sprintf( "Playthrough done for %s, looping %q (%d/%d)", - ip, current.data.Name, loopCount, config.MaxLoop, + ip, current.data.Name, loopCount, cfg.MaxLoop, )) } else { - logInfo(fmt.Sprintf( + logx.Info(fmt.Sprintf( "Playthrough done for %s, looping %q (%d)", ip, current.data.Name, loopCount, )) diff --git a/src/performance_safeguards_test.go b/internal/sshserver/server_test.go similarity index 57% rename from src/performance_safeguards_test.go rename to internal/sshserver/server_test.go index c88586d..7bb367f 100644 --- a/src/performance_safeguards_test.go +++ b/internal/sshserver/server_test.go @@ -1,14 +1,33 @@ -package main +package sshserver import ( - "bytes" - "image" "sync" "testing" "golang.org/x/crypto/ssh" ) +func TestConnectionTracker(t *testing.T) { + tr := newConnectionTracker() + if _, _, ok := tr.tryAcquire("1.2.3.4", 2, 100); !ok { + t.Fatal("first acquire failed") + } + 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 _, _, ok := tr.tryAcquire("1.2.3.4", 2, 100); !ok { + t.Error("limit should be cleared") + } +} + func TestConnectionTrackerConcurrentLimit(t *testing.T) { tracker := newConnectionTracker() start := make(chan struct{}) @@ -83,45 +102,36 @@ func TestTermSizeDebouncesResize(t *testing.T) { } } -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), +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) } - 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) + if w%4 != 0 || h%4 != 0 { + t.Fatalf("size is not quantized: %dx%d", w, h) + } + w, h = clampTermSize(3, 2, 100, 100, 4) + if w != 3 || h != 2 { + t.Fatalf("small size = %dx%d", w, h) } } -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), +func TestParseDimsPtyReq(t *testing.T) { + // "xterm" + cols=100 rows=40 + widthpx + heightpx + payload := []byte{ + 0, 0, 0, 5, 'x', 't', 'e', 'r', 'm', + 0, 0, 0, 100, + 0, 0, 0, 40, + 0, 0, 0, 0, + 0, 0, 0, 0, } - 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) + cols, rows, ok := parseDims(payload) + if !ok || cols != 100 || rows != 40 { + t.Errorf("parseDims = %d,%d,%v", cols, rows, ok) } -} -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) + term, ok := parsePtyTerm(payload) + if !ok || term != "xterm" { + t.Errorf("parsePtyTerm = %q,%v", term, ok) } } diff --git a/internal/tsf/container.go b/internal/tsf/container.go new file mode 100644 index 0000000..d51555f --- /dev/null +++ b/internal/tsf/container.go @@ -0,0 +1,52 @@ +// .tsf layout, little-endian: "TSFR" | version uint16 | fps float64 | +// count uint32 | count × (colorLen uint32, color JPEG). +package tsf + +import "sync" + +const ( + tsfMagic = "TSFR" + tsfVersion = 1 + maxTSFFPS = 240 + maxTSFFrameCount = 10_000_000 +) + +type FramesContainer struct { + ColorFrames [][]byte + FPS float64 + Name string +} + +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() +} diff --git a/src/frameformat.go b/internal/tsf/format.go similarity index 76% rename from src/frameformat.go rename to internal/tsf/format.go index 8e1e859..dc3469d 100644 --- a/src/frameformat.go +++ b/internal/tsf/format.go @@ -1,4 +1,4 @@ -package main +package tsf import ( "bufio" @@ -6,53 +6,9 @@ import ( "fmt" "math" "os" - "sync" ) -// .tsf container, little-endian: "TSFR" | version uint16 | fps float64 | -// count uint32 | count × (colorLen uint32, color JPEG). -const ( - tsfMagic = "TSFR" - 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 Write(output string, data *FramesContainer) error { if data == nil { return fmt.Errorf("cannot write nil frames container") } @@ -99,7 +55,7 @@ func writeTSF(output string, data *FramesContainer) error { return w.Flush() } -func loadTSF(filename string) (*FramesContainer, error) { +func Load(filename string) (*FramesContainer, error) { file, err := readFrameFile(filename) if err != nil { return nil, err diff --git a/src/mmap_other.go b/internal/tsf/mmap_other.go similarity index 92% rename from src/mmap_other.go rename to internal/tsf/mmap_other.go index 824a461..6fdd9d1 100644 --- a/src/mmap_other.go +++ b/internal/tsf/mmap_other.go @@ -1,6 +1,6 @@ //go:build !unix -package main +package tsf import "os" diff --git a/src/mmap_unix.go b/internal/tsf/mmap_unix.go similarity index 98% rename from src/mmap_unix.go rename to internal/tsf/mmap_unix.go index 635673f..3e32f71 100644 --- a/src/mmap_unix.go +++ b/internal/tsf/mmap_unix.go @@ -1,6 +1,6 @@ //go:build unix -package main +package tsf import ( "os" diff --git a/internal/tsf/tsf_test.go b/internal/tsf/tsf_test.go new file mode 100644 index 0000000..ff66240 --- /dev/null +++ b/internal/tsf/tsf_test.go @@ -0,0 +1,270 @@ +package tsf + +import ( + "bytes" + "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 TestTSFRoundTrip(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "f.tsf") + original := &FramesContainer{ + ColorFrames: [][]byte{{100, 101, 102}, {110, 120, 130}}, + FPS: 29.97, + } + if err := Write(path, original); err != nil { + t.Fatalf("Write: %v", err) + } + fc, err := Load(path) + if err != nil { + t.Fatalf("Load: %v", err) + } + defer func() { _ = fc.Close() }() + if fc.FPS != 29.97 { + t.Errorf("fps = %v", fc.FPS) + } + if len(fc.ColorFrames) != 2 { + t.Fatalf("frames = %d color", len(fc.ColorFrames)) + } + if string(fc.ColorFrames[0]) != string([]byte{100, 101, 102}) { + t.Errorf("color frame0 = %v", fc.ColorFrames[0]) + } + if string(fc.ColorFrames[1]) != string([]byte{110, 120, 130}) { + t.Errorf("color frame1 = %v", fc.ColorFrames[1]) + } +} + +func TestTSFInvalid(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "bad.tsf") + + if err := os.WriteFile(path, []byte("not a tsf file"), 0o644); err != nil { + t.Fatalf("WriteFile: %v", err) + } + if _, err := Load(path); err == nil { + t.Error("expected error for garbage input") + } + + // Valid container but no frames. + if err := Write(path, &FramesContainer{FPS: 30}); err != nil { + t.Fatalf("Write: %v", err) + } + if _, err := Load(path); err == nil { + t.Error("expected error for empty frames") + } + + // Valid container but fps <= 0. + 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 := Load(path); err == nil { + t.Error("expected error for fps<=0") + } + + // Truncated payload. + if err := Write(path, &FramesContainer{ColorFrames: [][]byte{{1, 2, 3, 4}}, FPS: 30}); err != nil { + t.Fatalf("Write: %v", err) + } + raw, err := os.ReadFile(path) + if err != nil { + t.Fatalf("ReadFile: %v", err) + } + if err := os.WriteFile(path, raw[:len(raw)-2], 0o644); err != nil { + t.Fatalf("WriteFile: %v", err) + } + if _, err := Load(path); err == nil { + t.Error("expected error for truncated file") + } +} + +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 := Load(writeRawTSF(t, raw)); err == nil { + t.Errorf("Load accepted fps %v", fps) + } + } +} + +func TestTSFRejectsImpossibleCountsAndLengths(t *testing.T) { + if _, err := Load(writeRawTSF(t, tsfHeader(30, math.MaxUint32))); err == nil { + t.Fatal("Load accepted impossible frame count") + } + + raw := append(tsfHeader(30, 1), 0xff, 0xff, 0xff, 0xff) + if _, err := Load(writeRawTSF(t, raw)); err == nil { + t.Fatal("Load accepted overflowing frame length") + } +} + +func TestTSFCloseReleasesOwnedFrames(t *testing.T) { + path := filepath.Join(t.TempDir(), "frames.tsf") + if err := Write(path, &FramesContainer{FPS: 30, ColorFrames: [][]byte{{1, 2, 3}}}); err != nil { + t.Fatal(err) + } + frames, err := Load(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 := Write(path, &FramesContainer{FPS: math.NaN(), ColorFrames: [][]byte{{1}}}) + if err == nil || !strings.Contains(err.Error(), "fps") { + t.Fatalf("Write error = %v", err) + } + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Fatalf("invalid write created output: %v", err) + } +} + +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 := Load(output) + if err != nil { + t.Fatalf("Load: %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) + } +} diff --git a/src/videoprocessor.go b/internal/tsf/video.go similarity index 97% rename from src/videoprocessor.go rename to internal/tsf/video.go index 616d360..8b0f000 100644 --- a/src/videoprocessor.go +++ b/internal/tsf/video.go @@ -1,4 +1,4 @@ -package main +package tsf import ( "bufio" @@ -15,6 +15,8 @@ import ( "strconv" "strings" "time" + + "github.com/YuzuZensai/TrollSSH/internal/logx" ) var ( @@ -303,7 +305,7 @@ func extractFrames(path, vf, label string, totalFrames int, emit func([]byte) er return frameCount, nil } -func processVideo(path, output string, maxDimension int) error { +func ProcessVideo(path, output string, maxDimension int) error { probeCmd := exec.Command( "ffprobe", "-v", "error", "-show_streams", "-show_format", @@ -364,6 +366,6 @@ func processVideo(path, output string, maxDimension int) error { if err := outputFile.commit(output); err != nil { return err } - logInfo(fmt.Sprintf("Saved %d frames to %s", frameCount, output)) + logx.Info(fmt.Sprintf("Saved %d frames to %s", frameCount, output)) return nil } diff --git a/lefthook.yml b/lefthook.yml index 3b260de..14bb9a0 100644 --- a/lefthook.yml +++ b/lefthook.yml @@ -6,10 +6,10 @@ pre-commit: run: test -z "$(gofmt -l {staged_files})" vet: glob: "*.go" - run: go vet ./src/... + run: go vet ./... lint: glob: "*.go" - run: golangci-lint run ./src/... + run: golangci-lint run ./... test: glob: "*.go" - run: go test ./src/... + run: go test ./... diff --git a/src/app_test.go b/src/app_test.go deleted file mode 100644 index 88b88e5..0000000 --- a/src/app_test.go +++ /dev/null @@ -1,361 +0,0 @@ -package main - -import ( - "bytes" - "image" - "image/jpeg" - "os" - "path/filepath" - "strings" - "sync" - "testing" - "time" -) - -func TestTSFRoundTrip(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "f.tsf") - original := &FramesContainer{ - ColorFrames: [][]byte{{100, 101, 102}, {110, 120, 130}}, - FPS: 29.97, - } - if err := writeTSF(path, original); err != nil { - t.Fatalf("writeTSF: %v", err) - } - fc, err := loadTSF(path) - if err != nil { - t.Fatalf("loadTSF: %v", err) - } - if fc.FPS != 29.97 { - t.Errorf("fps = %v", fc.FPS) - } - if len(fc.ColorFrames) != 2 { - t.Fatalf("frames = %d color", len(fc.ColorFrames)) - } - if string(fc.ColorFrames[0]) != string([]byte{100, 101, 102}) { - t.Errorf("color frame0 = %v", fc.ColorFrames[0]) - } - if string(fc.ColorFrames[1]) != string([]byte{110, 120, 130}) { - t.Errorf("color frame1 = %v", fc.ColorFrames[1]) - } -} - -func TestTSFInvalid(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "bad.tsf") - - if err := os.WriteFile(path, []byte("not a tsf file"), 0o644); err != nil { - t.Fatalf("WriteFile: %v", err) - } - if _, err := loadTSF(path); err == nil { - t.Error("expected error for garbage input") - } - - // Valid container but no frames. - if err := writeTSF(path, &FramesContainer{FPS: 30}); err != nil { - t.Fatalf("writeTSF: %v", err) - } - if _, err := loadTSF(path); err == nil { - t.Error("expected error for empty frames") - } - - // Valid container but fps <= 0. - 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") - } - - // Truncated payload. - if err := writeTSF(path, &FramesContainer{ColorFrames: [][]byte{{1, 2, 3, 4}}, FPS: 30}); err != nil { - t.Fatalf("writeTSF: %v", err) - } - raw, err := os.ReadFile(path) - if err != nil { - t.Fatalf("ReadFile: %v", err) - } - if err := os.WriteFile(path, raw[:len(raw)-2], 0o644); err != nil { - t.Fatalf("WriteFile: %v", err) - } - if _, err := loadTSF(path); err == nil { - t.Error("expected error for truncated file") - } -} - -func TestLoadConfigDefaults(t *testing.T) { - t.Setenv("HOST", "") - t.Setenv("PORT", "") - t.Setenv("PLAYBACK_MODE", "") - t.Setenv("LOGIN_DELAY", "") - cfg := loadConfig() - if cfg.Host != "0.0.0.0" { - t.Errorf("host = %q", cfg.Host) - } - if cfg.Port != 22 { - t.Errorf("port = %d", cfg.Port) - } - if cfg.PlaybackMode != PlaybackLoop { - t.Errorf("playbackMode = %q", cfg.PlaybackMode) - } - if cfg.Charset != "detailed" { - t.Errorf("charset = %q", cfg.Charset) - } - if cfg.LoginDelay != 1500*time.Millisecond { - t.Errorf("loginDelay = %v", cfg.LoginDelay) - } -} - -func TestLoadConfigClamping(t *testing.T) { - t.Setenv("PORT", "999999") - t.Setenv("BRIGHTNESS_THRESHOLD", "-5") - cfg := loadConfig() - if cfg.Port != 65535 { - t.Errorf("port clamp = %d", cfg.Port) - } - if cfg.BrightnessThreshold != 0 { - t.Errorf("brightness clamp = %d", cfg.BrightnessThreshold) - } -} - -func TestLoadConfigInvalidFallsBack(t *testing.T) { - t.Setenv("PORT", "not-a-number") - t.Setenv("INVERT", "yes-please") - t.Setenv("PLAYBACK_MODE", "shuffle") - cfg := loadConfig() - if cfg.Port != 22 { - t.Errorf("port = %d, want default 22", cfg.Port) - } - if cfg.Invert { - t.Error("invert should fall back to false") - } - if cfg.PlaybackMode != PlaybackLoop { - t.Errorf("playbackMode = %q, want default loop", cfg.PlaybackMode) - } -} - -func TestEnvDurationMs(t *testing.T) { - t.Setenv("D", "250") - if got := envDurationMs("D", time.Second); got != 250*time.Millisecond { - t.Errorf("250 = %v, want 250ms", got) - } - t.Setenv("D", "-10") - if got := envDurationMs("D", time.Second); got != 0 { - t.Errorf("negative = %v, want 0", got) - } - t.Setenv("D", "banana") - if got := envDurationMs("D", time.Second); got != time.Second { - t.Errorf("invalid = %v, want fallback 1s", got) - } -} - -func TestPlaybackModeRandom(t *testing.T) { - t.Setenv("PLAYBACK_MODE", "RaNdOm") - if loadConfig().PlaybackMode != PlaybackRandom { - t.Error("expected random") - } -} - -func TestResolveCharset(t *testing.T) { - if got := resolveCharset("blocks"); got != " ░▒▓█" { - t.Errorf("blocks preset = %q", got) - } - if got := resolveCharset("XYZ"); got != "XYZ" { - t.Errorf("custom ramp = %q", got) - } - if got := resolveCharset(""); !strings.HasPrefix(got, " .") { - t.Errorf("default = %q", got) - } -} - -func TestFrameToAscii(t *testing.T) { - // Below threshold -> first ramp char; full brightness -> last. - opts := asciiOptions{brightnessThreshold: 40, charset: "standard"} - ramp := []rune(resolveCharset("standard")) - img := &image.RGBA{ - Pix: []byte{0, 0, 0, 255, 255, 255, 255, 255}, - Stride: 8, Rect: image.Rect(0, 0, 2, 1), - } - out := []rune(string(frameToAscii(img, buildRampLUT(ramp, opts)))) - if out[0] != ramp[0] { - t.Errorf("dark px = %q, want %q", out[0], ramp[0]) - } - if out[1] != ramp[len(ramp)-1] { - t.Errorf("bright px = %q, want %q", out[1], ramp[len(ramp)-1]) - } -} - -func TestFrameToAsciiInvert(t *testing.T) { - opts := asciiOptions{brightnessThreshold: 40, charset: "standard", invert: true} - ramp := []rune(resolveCharset("standard")) - img := &image.RGBA{ - Pix: []byte{255, 255, 255, 255}, - Stride: 4, Rect: image.Rect(0, 0, 1, 1), - } - out := []rune(string(frameToAscii(img, buildRampLUT(ramp, opts)))) - if out[0] != ramp[0] { - t.Errorf("inverted bright = %q, want %q", out[0], ramp[0]) - } -} - -func TestRenderConcurrentSameKey(t *testing.T) { - var jpegBuf bytes.Buffer - src := image.NewRGBA(image.Rect(0, 0, 16, 16)) - for i := range src.Pix { - src.Pix[i] = byte(i * 7) - } - if err := jpeg.Encode(&jpegBuf, src, nil); err != nil { - t.Fatal(err) - } - - r := newFrameRenderer(0, [][]byte{jpegBuf.Bytes()}, asciiOptions{ - brightnessThreshold: 40, - charset: "standard", - }, newRenderCache(1<<20)) - - var wg sync.WaitGroup - results := make([][]byte, 32) - for i := range results { - wg.Add(1) - go func(i int) { - defer wg.Done() - ascii, err := r.render(0, 20, 10, false, colorTierTrueColor) - if err != nil { - t.Error(err) - return - } - results[i] = ascii - }(i) - } - wg.Wait() - for i, got := range results { - if !bytes.Equal(got, results[0]) { - t.Fatalf("result %d differs from result 0", i) - } - } -} - -func TestRenderCacheEvictsByBytes(t *testing.T) { - 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(key(i), bytes.Repeat([]byte("x"), 1000)) - } - if c.size.Load() > budget { - t.Errorf("size %d exceeds budget %d", c.size.Load(), budget) - } - if _, ok := c.get(key(0)); ok { - t.Error("oldest entry should have been evicted") - } - if _, ok := c.get(key(4)); !ok { - t.Error("newest entry should be cached") - } -} - -func TestRenderCacheDisabled(t *testing.T) { - c := newRenderCache(0) - if c != nil { - t.Fatal("zero budget should disable the cache") - } - 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(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.Load() != 0 { - t.Errorf("size = %d, want 0", c.size.Load()) - } -} - -func TestConnectionTracker(t *testing.T) { - tr := newConnectionTracker() - if _, _, ok := tr.tryAcquire("1.2.3.4", 2, 100); !ok { - t.Fatal("first acquire failed") - } - 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 _, _, ok := tr.tryAcquire("1.2.3.4", 2, 100); !ok { - t.Error("limit should be cleared") - } -} - -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 w%4 != 0 || h%4 != 0 { - t.Fatalf("size is not quantized: %dx%d", w, h) - } - w, h = clampTermSize(3, 2, 100, 100, 4) - if w != 3 || h != 2 { - t.Fatalf("small size = %dx%d", w, h) - } -} - -func TestParseDimsPtyReq(t *testing.T) { - // "xterm" + cols=100 rows=40 + widthpx + heightpx - payload := []byte{ - 0, 0, 0, 5, 'x', 't', 'e', 'r', 'm', - 0, 0, 0, 100, - 0, 0, 0, 40, - 0, 0, 0, 0, - 0, 0, 0, 0, - } - cols, rows, ok := parseDims(payload) - if !ok || cols != 100 || rows != 40 { - t.Errorf("parseDims = %d,%d,%v", cols, rows, ok) - } - - term, ok := parsePtyTerm(payload) - if !ok || term != "xterm" { - t.Errorf("parsePtyTerm = %q,%v", term, ok) - } -} - -func TestDetectColorTier(t *testing.T) { - cases := map[string]colorTier{ - "": colorTierNone, - "dumb": colorTierNone, - "vt100": colorTierNone, - "linux": colorTierNone, - "xterm": colorTierTrueColor, - "xterm-256color": colorTier256, - "screen-256color": colorTier256, - "tmux-256color": colorTier256, - "xterm-direct": colorTierTrueColor, - "xterm-kitty": colorTierTrueColor, - } - for term, want := range cases { - if got := detectColorTier(term); got != want { - t.Errorf("detectColorTier(%q) = %d, want %d", term, got, want) - } - } -} - -func TestQuantize256(t *testing.T) { - if got := quantize256(0, 0, 0); got != 16 { - t.Errorf("black = %d, want 16", got) - } - if got := quantize256(255, 255, 255); got != 231 { - t.Errorf("white = %d, want 231", got) - } -} diff --git a/src/frameformat_hardening_test.go b/src/frameformat_hardening_test.go deleted file mode 100644 index 9449018..0000000 --- a/src/frameformat_hardening_test.go +++ /dev/null @@ -1,82 +0,0 @@ -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) - } -} diff --git a/src/frames.go b/src/frames.go deleted file mode 100644 index 633a974..0000000 --- a/src/frames.go +++ /dev/null @@ -1,482 +0,0 @@ -package main - -import ( - "bytes" - "container/list" - "image" - "image/color" - "image/jpeg" - "strconv" - "strings" - "sync" - "sync/atomic" - "time" - "unicode/utf8" - - "golang.org/x/image/draw" -) - -type FramesContainer struct { - ColorFrames [][]byte - FPS float64 - Name string -} - -type colorTier int - -const ( - colorTierNone colorTier = iota - colorTier256 - colorTierTrueColor -) - -func detectColorTier(term string) colorTier { - t := strings.ToLower(strings.TrimSpace(term)) - switch t { - case "", "dumb", "vt52", "vt100", "vt102", "vt220", "ansi", "linux", "cons25", "cygwin": - return colorTierNone - } - if strings.Contains(t, "direct") || strings.Contains(t, "truecolor") { - return colorTierTrueColor - } - if strings.Contains(t, "256color") { - return colorTier256 - } - if strings.HasPrefix(t, "screen") || strings.HasPrefix(t, "tmux") { - return colorTier256 - } - return colorTierTrueColor -} - -var charsetPresets = map[string]string{ - "detailed": " .'`^\",:;Il!i><~+_-?][}{1)(|/tfjrxnuvczXYUJCLQ0OZmwqpdbkhao*#MW&8%B@$", - "standard": " .:-=+*#%@", - "simple": " .:oO#@", - "blocks": " ░▒▓█", -} - -func resolveCharset(charset string) string { - if charset == "" { - return charsetPresets["detailed"] - } - if preset, ok := charsetPresets[strings.ToLower(charset)]; ok { - return preset - } - return charset -} - -const maxPooledBuffer = 4 << 20 - -var pixPool sync.Pool - -func getPixBuf(n int) []byte { - if v := pixPool.Get(); v != nil { - if b := *v.(*[]byte); cap(b) >= n { - return b[:n] - } - } - return make([]byte, n) -} - -func putPixBuf(b []byte) { - if cap(b) <= maxPooledBuffer { - pixPool.Put(&b) - } -} - -var outPool sync.Pool - -func getOutBuf(capacity int) []byte { - if v := outPool.Get(); v != nil { - if b := *v.(*[]byte); cap(b) >= capacity { - return b[:0] - } - } - return make([]byte, 0, capacity) -} - -func putOutBuf(b []byte) { - if cap(b) <= maxPooledBuffer { - outPool.Put(&b) - } -} - -func resizeFrame(frame, pix []byte, width, height int, keepAspectRatio bool) (*image.RGBA, error) { - src, err := jpeg.Decode(bytes.NewReader(frame)) - if err != nil { - return nil, err - } - rect := image.Rect(0, 0, width, height) - - dst := &image.RGBA{Pix: pix[:4*width*height], Stride: 4 * width, Rect: rect} - if keepAspectRatio { - draw.Draw(dst, dst.Bounds(), image.NewUniform(color.Black), image.Point{}, draw.Src) - sb := src.Bounds() - sw, sh := sb.Dx(), sb.Dy() - scale := min(float64(width)/float64(sw), float64(height)/float64(sh)) - tw := max(1, int(float64(sw)*scale)) - th := max(1, int(float64(sh)*scale)) - x0 := (width - tw) / 2 - y0 := (height - th) / 2 - draw.ApproxBiLinear.Scale(dst, image.Rect(x0, y0, x0+tw, y0+th), src, sb, draw.Src, nil) - } else { - draw.ApproxBiLinear.Scale(dst, dst.Bounds(), src, src.Bounds(), draw.Src, nil) - } - return dst, nil -} - -type asciiOptions struct { - brightnessThreshold int - charset string - invert bool -} - -func buildRampLUT(ramp []rune, options asciiOptions) *[101][]byte { - var lut [101][]byte - for b := range lut { - index := rampIndex(b, options.brightnessThreshold, len(ramp), options.invert) - lut[b] = utf8.AppendRune(nil, ramp[index]) - } - return &lut -} - -func rampIndex(brightness, threshold, total int, invert bool) int { - var index int - if brightness < threshold { - index = 0 - } else { - index = brightness * total / 100 - if index > total-1 { - index = total - 1 - } - } - if invert { - index = total - 1 - index - } - return index -} - -func frameToAscii(img *image.RGBA, rampLUT *[101][]byte) []byte { - pix := img.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 { - brightness := (int(pix[o])*299 + int(pix[o+1])*587 + int(pix[o+2])*114) / 255 / 10 - buf = append(buf, rampLUT[brightness]...) - } - output := bytes.Clone(buf) - putOutBuf(buf) - return output -} - -const ansiReset = "\x1b[0m" - -var ansi256Levels = [6]int{0, 95, 135, 175, 215, 255} - -var decimal = func() (t [256]string) { - for i := range t { - t[i] = strconv.Itoa(i) - } - return -}() - -var ansi256Cube = func() (t [256]uint8) { - for v := range t { - best, bestDist := 0, 1<<30 - for i, l := range ansi256Levels { - d := v - l - if d < 0 { - d = -d - } - if d < bestDist { - bestDist, best = d, i - } - } - t[v] = uint8(best) - } - return -}() - -func quantize256(r, g, b uint8) int { - return 16 + 36*int(ansi256Cube[r]) + 6*int(ansi256Cube[g]) + int(ansi256Cube[b]) -} - -func appendColor(buf []byte, r, g, b uint8, tier colorTier) []byte { - if tier == colorTierTrueColor { - buf = append(buf, "\x1b[38;2;"...) - buf = append(buf, decimal[r]...) - buf = append(buf, ';') - buf = append(buf, decimal[g]...) - buf = append(buf, ';') - buf = append(buf, decimal[b]...) - } else { - buf = append(buf, "\x1b[38;5;"...) - buf = append(buf, decimal[quantize256(r, g, b)]...) - } - return append(buf, 'm') -} - -func frameToAnsi(img *image.RGBA, rampLUT *[101][]byte, tier colorTier) []byte { - bounds := img.Bounds() - bytesPerCell := 11 - if tier == colorTierTrueColor { - bytesPerCell = 16 - } - buf := getOutBuf(bounds.Dx() * bounds.Dy() * bytesPerCell) - var lastR, lastG, lastB uint8 - last256 := -1 - first := true - 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++ { - 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 - 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) - lastR, lastG, lastB = r, g, bl - first = false - } - buf = append(buf, rampLUT[brightness]...) - o += 4 - } - if y < bounds.Max.Y-1 { - buf = append(buf, "\r\n"...) - } - } - buf = append(buf, ansiReset...) - output := bytes.Clone(buf) - putOutBuf(buf) - return output -} - -type cacheKey struct { - setID int - index int - width int - height int - keepAspectRatio bool - tier colorTier -} - -type renderCacheShard struct { - mu sync.Mutex - maxBytes int64 - size int64 - entries map[cacheKey]*list.Element - order *list.List -} - -type renderCache struct { - shards []renderCacheShard - size atomic.Int64 - hits atomic.Uint64 - misses atomic.Uint64 - evictions atomic.Uint64 - rejections atomic.Uint64 - renders atomic.Uint64 - renderNs atomic.Uint64 -} - -type cacheEntry struct { - key cacheKey - ascii []byte - cost int64 -} - -func entryCost(_ cacheKey, ascii []byte) int64 { - return int64(cap(ascii)) + 160 -} - -func newRenderCache(maxBytes int64) *renderCache { - if maxBytes <= 0 { - return nil - } - shardCount := int(min(int64(16), max(int64(1), maxBytes/(1<<20)))) - cache := &renderCache{shards: make([]renderCacheShard, shardCount)} - for i := range cache.shards { - cache.shards[i] = renderCacheShard{ - maxBytes: maxBytes / int64(shardCount), - entries: make(map[cacheKey]*list.Element), - order: list.New(), - } - } - return cache -} - -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 { - return nil, false - } - shard := c.shard(key) - shard.mu.Lock() - defer shard.mu.Unlock() - el, ok := shard.entries[key] - if !ok { - c.misses.Add(1) - return nil, false - } - c.hits.Add(1) - shard.order.MoveToBack(el) - return el.Value.(*cacheEntry).ascii, true -} - -func (c *renderCache) put(key cacheKey, ascii []byte) { - if c == nil { - return - } - shard := c.shard(key) - cost := entryCost(key, ascii) - if cost > shard.maxBytes { - c.rejections.Add(1) - return - } - shard.mu.Lock() - defer shard.mu.Unlock() - if _, ok := shard.entries[key]; ok { - return - } - shard.entries[key] = shard.order.PushBack(&cacheEntry{key: key, ascii: ascii, cost: cost}) - shard.size += cost - c.size.Add(cost) - for shard.size > shard.maxBytes { - oldest := shard.order.Front() - shard.order.Remove(oldest) - evicted := oldest.Value.(*cacheEntry) - delete(shard.entries, evicted.key) - 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()), - } -} - -type FrameRenderer struct { - setID int - colorFrames [][]byte - options asciiOptions - rampLUT *[101][]byte - cache *renderCache - - inflightMu sync.Mutex - 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 { - ramp := []rune(resolveCharset(options.charset)) - return &FrameRenderer{ - setID: setID, - colorFrames: colorFrames, - options: options, - rampLUT: buildRampLUT(ramp, options), - cache: cache, - inflight: make(map[cacheKey]*renderCall), - } -} - -func (r *FrameRenderer) render(index, width, height int, keepAspectRatio bool, tier colorTier) ([]byte, error) { - key := cacheKey{r.setID, index, width, height, keepAspectRatio, tier} - if ascii, ok := r.cache.get(key); ok { - return ascii, nil - } - - r.inflightMu.Lock() - if call, ok := r.inflight[key]; ok { - r.inflightMu.Unlock() - <-call.done - return call.value, call.err - } - call := &renderCall{done: make(chan struct{})} - r.inflight[key] = call - r.inflightMu.Unlock() - defer func() { - r.inflightMu.Lock() - delete(r.inflight, key) - r.inflightMu.Unlock() - close(call.done) - }() - - if ascii, ok := r.cache.get(key); ok { - call.value = ascii - return ascii, nil - } - - started := time.Now() - pix := getPixBuf(4 * width * height) - img, err := resizeFrame(r.colorFrames[index], pix, width, height, keepAspectRatio) - if err != nil { - putPixBuf(pix) - call.err = err - return nil, err - } - var ascii []byte - if tier == colorTierNone { - ascii = frameToAscii(img, r.rampLUT) - } else { - ascii = frameToAnsi(img, r.rampLUT, tier) - } - putPixBuf(pix) - - if r.cache != nil && cap(ascii) > len(ascii)+len(ascii)/4 { - ascii = bytes.Clone(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 -} diff --git a/src/render_bench_test.go b/src/render_bench_test.go deleted file mode 100644 index a95905d..0000000 --- a/src/render_bench_test.go +++ /dev/null @@ -1,67 +0,0 @@ -package main - -import ( - "path/filepath" - "sync/atomic" - "testing" -) - -func loadBenchSet(b *testing.B) *FramesContainer { - b.Helper() - matches, _ := filepath.Glob("../frames/*.tsf") - if len(matches) == 0 { - b.Skip("no .tsf frame set in ../frames") - } - fc, err := loadTSF(matches[0]) - if err != nil { - b.Skip("failed to load frame set:", err) - } - b.Cleanup(func() { _ = fc.Close() }) - return fc -} - -func benchRender(b *testing.B, tier colorTier, w, h int) { - fc := loadBenchSet(b) - r := newFrameRenderer(0, fc.ColorFrames, asciiOptions{ - brightnessThreshold: 40, - charset: "detailed", - }, nil) - b.ReportAllocs() - b.ResetTimer() - for i := 0; i < b.N; i++ { - if _, err := r.render(i%len(fc.ColorFrames), w, h, false, tier); err != nil { - b.Fatal(err) - } - } -} - -func BenchmarkRenderTrueColor(b *testing.B) { benchRender(b, colorTierTrueColor, 120, 40) } -func BenchmarkRender256(b *testing.B) { benchRender(b, colorTier256, 120, 40) } -func BenchmarkRenderGray(b *testing.B) { benchRender(b, colorTierNone, 120, 40) } -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()) - } -} diff --git a/src/videoprocessor_test.go b/src/videoprocessor_test.go deleted file mode 100644 index 105cbd1..0000000 --- a/src/videoprocessor_test.go +++ /dev/null @@ -1,122 +0,0 @@ -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) - } -}