From 812380fe058173daa28de801a01c7508e096994a Mon Sep 17 00:00:00 2001 From: Yuzu Date: Mon, 13 Jul 2026 22:54:23 +0700 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat:=20color=20frames,=20derive=20?= =?UTF-8?q?grayscale=20frames=20at=20render=20time?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .env.example | 4 +- src/app_test.go | 54 +++++++++++--- src/config.go | 2 + src/frameformat.go | 14 ++-- src/frames.go | 170 +++++++++++++++++++++++++++++++++--------- src/main.go | 2 +- src/server.go | 34 +++++++-- src/videoprocessor.go | 114 +++++++++++++++------------- 8 files changed, 284 insertions(+), 110 deletions(-) diff --git a/.env.example b/.env.example index de90e05..9798e07 100644 --- a/.env.example +++ b/.env.example @@ -26,6 +26,8 @@ CHARSET=blocks BRIGHTNESS_THRESHOLD=40 # Invert the brightness ramp (for light terminal backgrounds). INVERT=false +# Always render grayscale, even for clients that support color. +FORCE_GRAYSCALE=false # Max rendered width/height in characters. MAX_DIMENSION=1080 @@ -39,7 +41,7 @@ MAX_TOTAL_CONNECTIONS=1000 # (Empty password counts as an attempt.) MAX_AUTH_ATTEMPTS=3 # SSH handshake deadline in ms (0 to disable). -HANDSHAKE_TIMEOUT=10000 +HANDSHAKE_TIMEOUT=30000 # Log attempted usernames/passwords. LOG_CREDENTIALS=true diff --git a/src/app_test.go b/src/app_test.go index 9ef6bcc..4185d02 100644 --- a/src/app_test.go +++ b/src/app_test.go @@ -12,8 +12,8 @@ func TestTSFRoundTrip(t *testing.T) { dir := t.TempDir() path := filepath.Join(dir, "f.tsf") original := &FramesContainer{ - Frames: [][]byte{{0, 1, 255}, {10, 20, 30}}, - FPS: 29.97, + ColorFrames: [][]byte{{100, 101, 102}, {110, 120, 130}}, + FPS: 29.97, } if err := writeTSF(path, original); err != nil { t.Fatalf("writeTSF: %v", err) @@ -25,14 +25,14 @@ func TestTSFRoundTrip(t *testing.T) { if fc.FPS != 29.97 { t.Errorf("fps = %v", fc.FPS) } - if len(fc.Frames) != 2 { - t.Fatalf("frames = %d", len(fc.Frames)) + if len(fc.ColorFrames) != 2 { + t.Fatalf("frames = %d color", len(fc.ColorFrames)) } - if string(fc.Frames[0]) != string([]byte{0, 1, 255}) { - t.Errorf("frame0 = %v", fc.Frames[0]) + if string(fc.ColorFrames[0]) != string([]byte{100, 101, 102}) { + t.Errorf("color frame0 = %v", fc.ColorFrames[0]) } - if string(fc.Frames[1]) != string([]byte{10, 20, 30}) { - t.Errorf("frame1 = %v", fc.Frames[1]) + if string(fc.ColorFrames[1]) != string([]byte{110, 120, 130}) { + t.Errorf("color frame1 = %v", fc.ColorFrames[1]) } } @@ -52,13 +52,13 @@ func TestTSFInvalid(t *testing.T) { } // Valid container but fps <= 0. - writeTSF(path, &FramesContainer{Frames: [][]byte{{1}}, FPS: 0}) + writeTSF(path, &FramesContainer{ColorFrames: [][]byte{{1}}, FPS: 0}) if _, err := loadTSF(path); err == nil { t.Error("expected error for fps<=0") } // Truncated payload. - writeTSF(path, &FramesContainer{Frames: [][]byte{{1, 2, 3, 4}}, FPS: 30}) + writeTSF(path, &FramesContainer{ColorFrames: [][]byte{{1, 2, 3, 4}}, FPS: 30}) raw, _ := os.ReadFile(path) os.WriteFile(path, raw[:len(raw)-2], 0o644) if _, err := loadTSF(path); err == nil { @@ -215,4 +215,38 @@ func TestParseDimsPtyReq(t *testing.T) { 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/config.go b/src/config.go index 06ab41c..0706099 100644 --- a/src/config.go +++ b/src/config.go @@ -32,6 +32,7 @@ type Config struct { BrightnessThreshold int Charset string Invert bool + ForceGrayscale bool LogCredentials bool } @@ -129,6 +130,7 @@ func loadConfig() Config { BrightnessThreshold: envInt("BRIGHTNESS_THRESHOLD", 40, 0, 100), Charset: envString("CHARSET", "detailed"), Invert: envBool("INVERT", false), + ForceGrayscale: envBool("FORCE_GRAYSCALE", false), LogCredentials: envBool("LOG_CREDENTIALS", false), } } diff --git a/src/frameformat.go b/src/frameformat.go index e93b012..b831573 100644 --- a/src/frameformat.go +++ b/src/frameformat.go @@ -9,7 +9,7 @@ import ( ) // .tsf container, little-endian: "TSFR" | version uint16 | fps float64 | -// count uint32 | count × (length uint32, JPEG bytes). +// count uint32 | count × (colorLen uint32, color JPEG). const ( tsfMagic = "TSFR" tsfVersion = 1 @@ -29,13 +29,13 @@ func writeTSF(output string, data *FramesContainer) error { var hdr [14]byte binary.LittleEndian.PutUint16(hdr[0:], tsfVersion) binary.LittleEndian.PutUint64(hdr[2:], math.Float64bits(data.FPS)) - binary.LittleEndian.PutUint32(hdr[10:], uint32(len(data.Frames))) + binary.LittleEndian.PutUint32(hdr[10:], uint32(len(data.ColorFrames))) if _, err := w.Write(hdr[:]); err != nil { return err } var lenBuf [4]byte - for _, frame := range data.Frames { + for _, frame := range data.ColorFrames { binary.LittleEndian.PutUint32(lenBuf[:], uint32(len(frame))) if _, err := w.Write(lenBuf[:]); err != nil { return err @@ -66,7 +66,7 @@ func loadTSF(filename string) (*FramesContainer, error) { fps := math.Float64frombits(binary.LittleEndian.Uint64(raw[6:])) count := binary.LittleEndian.Uint32(raw[14:]) - frames := make([][]byte, 0, count) + colorFrames := make([][]byte, 0, count) off := 18 for range count { if off+4 > len(raw) { @@ -77,15 +77,15 @@ func loadTSF(filename string) (*FramesContainer, error) { if off+n > len(raw) { return nil, invalid() } - frames = append(frames, raw[off:off+n]) + colorFrames = append(colorFrames, raw[off:off+n]) off += n } - if len(frames) == 0 || fps <= 0 { + if len(colorFrames) == 0 || fps <= 0 { return nil, fmt.Errorf( "invalid frames file %q: expected non-empty frames and a positive fps", filename, ) } - return &FramesContainer{Frames: frames, FPS: fps}, nil + return &FramesContainer{ColorFrames: colorFrames, FPS: fps}, nil } diff --git a/src/frames.go b/src/frames.go index af43738..1637aa4 100644 --- a/src/frames.go +++ b/src/frames.go @@ -14,9 +14,35 @@ import ( ) type FramesContainer struct { - Frames [][]byte - FPS float64 - Name string + 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{ @@ -36,14 +62,22 @@ func resolveCharset(charset string) string { return charset } -func resizeFrame(frame []byte, width, height int, keepAspectRatio bool) ([]byte, error) { +func resizeFrame(frame []byte, width, height int, keepAspectRatio bool, tier colorTier) (draw.Image, error) { src, err := jpeg.Decode(bytes.NewReader(frame)) if err != nil { return nil, err } - dst := image.NewGray(image.Rect(0, 0, width, height)) + var dst draw.Image + var bg color.Color + if tier == colorTierNone { + dst = image.NewGray(image.Rect(0, 0, width, height)) + bg = color.Gray{0} + } else { + dst = image.NewNRGBA(image.Rect(0, 0, width, height)) + bg = color.Black + } if keepAspectRatio { - draw.Draw(dst, dst.Bounds(), image.NewUniform(color.Gray{0}), image.Point{}, draw.Src) + draw.Draw(dst, dst.Bounds(), image.NewUniform(bg), image.Point{}, draw.Src) sb := src.Bounds() sw, sh := sb.Dx(), sb.Dy() scale := min(float64(width)/float64(sw), float64(height)/float64(sh)) @@ -55,7 +89,7 @@ func resizeFrame(frame []byte, width, height int, keepAspectRatio bool) ([]byte, } else { draw.ApproxBiLinear.Scale(dst, dst.Bounds(), src, src.Bounds(), draw.Src, nil) } - return dst.Pix, nil + return dst, nil } type asciiOptions struct { @@ -64,33 +98,92 @@ type asciiOptions struct { invert bool } +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(pixels []byte, options asciiOptions) string { ramp := []rune(resolveCharset(options.charset)) total := len(ramp) var b strings.Builder for _, p := range pixels { brightness := int(p) * 100 / 255 - var index int - if brightness < options.brightnessThreshold { - index = 0 - } else { - index = brightness * total / 100 - if index > total-1 { - index = total - 1 - } - } - if options.invert { - index = total - 1 - index - } + index := rampIndex(brightness, options.brightnessThreshold, total, options.invert) b.WriteRune(ramp[index]) } return b.String() } +const ansiReset = "\x1b[0m" + +var ansi256Levels = [6]int{0, 95, 135, 175, 215, 255} + +func quantize256(r, g, b uint8) int { + toLevel := func(v uint8) int { + best, bestDist := 0, 1<<30 + for i, l := range ansi256Levels { + d := int(v) - l + if d < 0 { + d = -d + } + if d < bestDist { + bestDist, best = d, i + } + } + return best + } + return 16 + 36*toLevel(r) + 6*toLevel(g) + toLevel(b) +} + +func frameToAnsi(img *image.NRGBA, options asciiOptions, tier colorTier) string { + ramp := []rune(resolveCharset(options.charset)) + total := len(ramp) + bounds := img.Bounds() + var b strings.Builder + var lastR, lastG, lastB uint8 + first := true + for y := bounds.Min.Y; y < bounds.Max.Y; y++ { + 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] + brightness := (int(r)*299 + int(g)*587 + int(bl)*114) / 255 / 10 + index := rampIndex(brightness, options.brightnessThreshold, total, options.invert) + if first || r != lastR || g != lastG || bl != lastB { + if tier == colorTierTrueColor { + fmt.Fprintf(&b, "\x1b[38;2;%d;%d;%dm", r, g, bl) + } else { + fmt.Fprintf(&b, "\x1b[38;5;%dm", quantize256(r, g, bl)) + } + lastR, lastG, lastB = r, g, bl + first = false + } + b.WriteRune(ramp[index]) + } + if y < bounds.Max.Y-1 { + b.WriteString(ansiReset + "\r\n") + first = true + } + } + b.WriteString(ansiReset) + return b.String() +} + type FrameRenderer struct { - frames [][]byte - options asciiOptions - maxEntries int + colorFrames [][]byte + options asciiOptions + maxEntries int mu sync.Mutex cache map[string]*list.Element @@ -102,18 +195,18 @@ type cacheEntry struct { ascii string } -func newFrameRenderer(frames [][]byte, options asciiOptions) *FrameRenderer { +func newFrameRenderer(colorFrames [][]byte, options asciiOptions) *FrameRenderer { return &FrameRenderer{ - frames: frames, - options: options, - maxEntries: 4096, - cache: make(map[string]*list.Element), - order: list.New(), + colorFrames: colorFrames, + options: options, + maxEntries: 4096, + cache: make(map[string]*list.Element), + order: list.New(), } } -func (r *FrameRenderer) render(index, width, height int, keepAspectRatio bool) (string, error) { - key := fmt.Sprintf("%d:%dx%d:%t", index, width, height, keepAspectRatio) +func (r *FrameRenderer) render(index, width, height int, keepAspectRatio bool, tier colorTier) (string, error) { + key := fmt.Sprintf("%d:%dx%d:%t:%d", index, width, height, keepAspectRatio, tier) r.mu.Lock() if el, ok := r.cache[key]; ok { @@ -124,11 +217,20 @@ func (r *FrameRenderer) render(index, width, height int, keepAspectRatio bool) ( } r.mu.Unlock() - pixels, err := resizeFrame(r.frames[index], width, height, keepAspectRatio) - if err != nil { - return "", err + var ascii string + if tier == colorTierNone { + img, err := resizeFrame(r.colorFrames[index], width, height, keepAspectRatio, tier) + if err != nil { + return "", err + } + ascii = frameToAscii(img.(*image.Gray).Pix, r.options) + } else { + img, err := resizeFrame(r.colorFrames[index], width, height, keepAspectRatio, tier) + if err != nil { + return "", err + } + ascii = frameToAnsi(img.(*image.NRGBA), r.options, tier) } - ascii := frameToAscii(pixels, r.options) r.mu.Lock() if _, ok := r.cache[key]; !ok { diff --git a/src/main.go b/src/main.go index 0a7a782..c4376e8 100644 --- a/src/main.go +++ b/src/main.go @@ -122,7 +122,7 @@ func loadAllFrames(framesDir string) []*FramesContainer { return } data.Name = file - logInfo(fmt.Sprintf(" %s: %d frames @ %gfps", file, len(data.Frames), data.FPS)) + logInfo(fmt.Sprintf(" %s: %d frames @ %gfps", file, len(data.ColorFrames), data.FPS)) results[i] = data } } diff --git a/src/server.go b/src/server.go index 2c4fdde..71f6b67 100644 --- a/src/server.go +++ b/src/server.go @@ -102,7 +102,7 @@ func createServer(deps ServerDeps) *Server { for i, data := range deps.VideoSets { sets[i] = frameSet{ data: data, - renderer: newFrameRenderer(data.Frames, asciiOptions{ + renderer: newFrameRenderer(data.ColorFrames, asciiOptions{ brightnessThreshold: config.BrightnessThreshold, charset: config.Charset, invert: config.Invert, @@ -277,6 +277,18 @@ func parseDims(payload []byte) (cols, rows int, ok bool) { return cols, rows, true } +// parsePtyTerm extracts the TERM string prefixing a pty-req payload. +func parsePtyTerm(payload []byte) (term string, ok bool) { + if len(payload) < 4 { + return "", false + } + strLen := binary.BigEndian.Uint32(payload) + if int(strLen)+16 > len(payload) { + return "", false + } + return string(payload[4 : 4+strLen]), true +} + func (s *Server) handleSession( sshConn *ssh.ServerConn, channel ssh.Channel, @@ -286,6 +298,10 @@ func (s *Server) handleSession( ) { size := &termSize{} size.set(80, 24, s.config.MaxDimension) + tier := colorTierTrueColor + if s.config.ForceGrayscale { + tier = colorTierNone + } started := false for req := range requests { @@ -295,6 +311,13 @@ func (s *Server) handleSession( if cols, rows, ok := parseDims(req.Payload); ok { size.set(cols, rows, s.config.MaxDimension) } + if term, ok := parsePtyTerm(req.Payload); ok { + tier = detectColorTier(term) + if s.config.ForceGrayscale { + tier = colorTierNone + } + logDebug(fmt.Sprintf("Client %s TERM=%q -> color tier %d", ip, sanitizeN(term, 64), tier)) + } req.Reply(true, nil) case "window-change": if len(req.Payload) >= 8 { @@ -317,14 +340,14 @@ func (s *Server) handleSession( req.Reply(true, nil) if !started { started = true - go s.playVideo(sshConn, channel, size, ip, initialSetIndex, false) + go s.playVideo(sshConn, channel, size, ip, initialSetIndex, false, tier) } case "shell": logDebug("Opening shell for session", ip) req.Reply(true, nil) if !started { started = true - go s.playVideo(sshConn, channel, size, ip, initialSetIndex, false) + go s.playVideo(sshConn, channel, size, ip, initialSetIndex, false, tier) } default: if req.WantReply { @@ -352,6 +375,7 @@ func (s *Server) playVideo( ip string, setIndex int, keepAspectRatio bool, + tier colorTier, ) { config := s.config current := s.sets[setIndex] @@ -438,7 +462,7 @@ func (s *Server) playVideo( case <-ticker.C: w, h := size.get() - ascii, err := current.renderer.render(currentFrame, w, h, keepAspectRatio) + ascii, err := current.renderer.render(currentFrame, w, h, keepAspectRatio, tier) if err != nil { logError("Render error for", ip, sanitize(err.Error())) sshConn.Close() @@ -451,7 +475,7 @@ func (s *Server) playVideo( } currentFrame++ - if currentFrame < len(current.data.Frames) { + if currentFrame < len(current.data.ColorFrames) { continue } diff --git a/src/videoprocessor.go b/src/videoprocessor.go index 324d8d1..9a970b4 100644 --- a/src/videoprocessor.go +++ b/src/videoprocessor.go @@ -72,6 +72,61 @@ func parseFrameRate(rate string) float64 { return num } +func extractFrames(path, vf, label string, maxDimension, totalFrames int) ([][]byte, error) { + cmd := exec.Command( + "ffmpeg", "-i", path, + "-c:v", "mjpeg", + "-q:v", "3", + "-vf", vf, + "-f", "image2pipe", + "pipe:1", + ) + var stderr bytes.Buffer + cmd.Stderr = &stderr + stdout, err := cmd.StdoutPipe() + if err != nil { + return nil, fmt.Errorf("ffmpeg failed: %w", err) + } + if err := cmd.Start(); err != nil { + return nil, fmt.Errorf("ffmpeg failed: %w", err) + } + + var frames [][]byte + splitter := &jpegFrameSplitter{} + reportProgress := func(count int) { + if totalFrames > 0 { + pct := min(100, int(math.Round(float64(count)/float64(totalFrames)*100))) + fmt.Printf("\rGenerating %s frames: %d/%d (%d%%)", label, count, totalFrames, pct) + } else { + fmt.Printf("\rGenerating %s frames: %d", label, count) + } + } + + buf := make([]byte, 256*1024) + for { + n, err := stdout.Read(buf) + if n > 0 { + frames = append(frames, splitter.push(buf[:n])...) + reportProgress(len(frames)) + } + if err == io.EOF { + break + } + if err != nil { + cmd.Wait() + return nil, fmt.Errorf("ffmpeg stream error: %s", err.Error()) + } + } + if err := cmd.Wait(); err != nil { + return nil, fmt.Errorf("ffmpeg failed: %s", strings.TrimSpace(stderr.String())) + } + if len(frames) == 0 { + return nil, fmt.Errorf("no frames were decoded from the video") + } + fmt.Println() + return frames, nil +} + func processVideo(path, output string, maxDimension int) error { probeCmd := exec.Command( "ffprobe", "-v", "error", @@ -115,65 +170,20 @@ func processVideo(path, output string, maxDimension int) error { } } - vf := fmt.Sprintf( - "format=gray,scale=w=%d:h=%d:force_original_aspect_ratio=decrease", + scaleFilter := fmt.Sprintf( + "scale=w=%d:h=%d:force_original_aspect_ratio=decrease", maxDimension, maxDimension, ) - cmd := exec.Command( - "ffmpeg", "-i", path, - "-c:v", "mjpeg", - "-q:v", "3", - "-vf", vf, - "-f", "image2pipe", - "pipe:1", - ) - var stderr bytes.Buffer - cmd.Stderr = &stderr - stdout, err := cmd.StdoutPipe() + + colorFrames, err := extractFrames(path, scaleFilter, "color", maxDimension, totalFrames) if err != nil { - return fmt.Errorf("ffmpeg failed: %w", err) - } - if err := cmd.Start(); err != nil { - return fmt.Errorf("ffmpeg failed: %w", err) + return err } - videoData := FramesContainer{FPS: fps} - splitter := &jpegFrameSplitter{} - reportProgress := func(count int) { - if totalFrames > 0 { - pct := min(100, int(math.Round(float64(count)/float64(totalFrames)*100))) - fmt.Printf("\rGenerating frames: %d/%d (%d%%)", count, totalFrames, pct) - } else { - fmt.Printf("\rGenerating frames: %d", count) - } - } - - buf := make([]byte, 256*1024) - for { - n, err := stdout.Read(buf) - if n > 0 { - videoData.Frames = append(videoData.Frames, splitter.push(buf[:n])...) - reportProgress(len(videoData.Frames)) - } - if err == io.EOF { - break - } - if err != nil { - cmd.Wait() - return fmt.Errorf("ffmpeg stream error: %s", err.Error()) - } - } - if err := cmd.Wait(); err != nil { - return fmt.Errorf("ffmpeg failed: %s", strings.TrimSpace(stderr.String())) - } - if len(videoData.Frames) == 0 { - return fmt.Errorf("no frames were decoded from the video") - } - - fmt.Println() + videoData := FramesContainer{FPS: fps, ColorFrames: colorFrames} if err := writeTSF(output, &videoData); err != nil { return err } - logInfo(fmt.Sprintf("Saved %d frames to %s", len(videoData.Frames), output)) + logInfo(fmt.Sprintf("Saved %d frames to %s", len(videoData.ColorFrames), output)) return nil }