mirror of
https://github.com/YuzuZensai/TrollSSH.git
synced 2026-09-13 19:09:09 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
be6cbe484a | ||
|
|
71a414d1ca | ||
|
|
4a43de0a95 | ||
|
|
0514931320
|
||
|
|
5ff2453f15
|
||
|
|
3107a452e6
|
||
|
|
e1e6d80ef9
|
||
|
|
ee9d36dd4f
|
||
|
|
6d221846df
|
||
|
|
facdb35d5d
|
||
|
|
a05d6bb7eb
|
||
|
|
2e320b03f3
|
||
|
|
9995557ae2
|
||
|
|
3c6950573f
|
||
|
|
65e3671671
|
||
|
|
340b26ddb7
|
||
|
|
172615f68f
|
@@ -26,12 +26,19 @@ FORCE_GRAYSCALE=false
|
|||||||
|
|
||||||
# Max rendered width/height in characters.
|
# Max rendered width/height in characters.
|
||||||
MAX_DIMENSION=512
|
MAX_DIMENSION=512
|
||||||
|
# Max rendered area (columns x rows); larger terminals are scaled down.
|
||||||
|
MAX_TERMINAL_CELLS=256000
|
||||||
|
|
||||||
# Memory budget in MB for the rendered-frame cache (0 disables caching).
|
# Memory budget in MB for the rendered-frame cache (0 disables caching).
|
||||||
RENDER_CACHE_MB=256
|
RENDER_CACHE_MB=256
|
||||||
|
# Compress cached frames (flate). Cuts cache RAM ~2x but decompresses on every
|
||||||
|
# cache hit, so CPU rises with concurrent sessions. Off = zero-copy reads.
|
||||||
|
RENDER_CACHE_COMPRESS=false
|
||||||
|
|
||||||
# Go soft memory limit; set below your container limit to avoid OOM kills.
|
# Go soft memory limit; set below your container limit to avoid OOM kills.
|
||||||
GOMEMLIMIT=1GiB
|
GOMEMLIMIT=1GiB
|
||||||
|
# Docker Compose container memory limit. Leave headroom for mapped frame files.
|
||||||
|
MEMORY_LIMIT=2g
|
||||||
|
|
||||||
# Connection limits. New connections over a limit are dropped immediately.
|
# Connection limits. New connections over a limit are dropped immediately.
|
||||||
# Max simultaneous connections from a single client IP.
|
# Max simultaneous connections from a single client IP.
|
||||||
@@ -43,6 +50,8 @@ MAX_TOTAL_CONNECTIONS=1000
|
|||||||
MAX_AUTH_ATTEMPTS=3
|
MAX_AUTH_ATTEMPTS=3
|
||||||
# SSH handshake deadline in ms (0 to disable).
|
# SSH handshake deadline in ms (0 to disable).
|
||||||
HANDSHAKE_TIMEOUT=30000
|
HANDSHAKE_TIMEOUT=30000
|
||||||
|
# Maximum session lifetime in ms (0 to disable).
|
||||||
|
SESSION_TIMEOUT=600000
|
||||||
|
|
||||||
# Log attempted usernames/passwords.
|
# Log attempted usernames/passwords.
|
||||||
LOG_CREDENTIALS=true
|
LOG_CREDENTIALS=true
|
||||||
|
|||||||
@@ -3,7 +3,7 @@ name: CI
|
|||||||
on:
|
on:
|
||||||
push:
|
push:
|
||||||
branches: ["**"]
|
branches: ["**"]
|
||||||
tags: ["v*"]
|
tags: ["[0-9]*.[0-9]*.[0-9]*"]
|
||||||
pull_request:
|
pull_request:
|
||||||
|
|
||||||
permissions:
|
permissions:
|
||||||
@@ -21,10 +21,10 @@ jobs:
|
|||||||
go-version-file: go.mod
|
go-version-file: go.mod
|
||||||
|
|
||||||
- name: Format check
|
- name: Format check
|
||||||
run: test -z "$(gofmt -l src)"
|
run: test -z "$(gofmt -l .)"
|
||||||
|
|
||||||
- name: Vet
|
- name: Vet
|
||||||
run: go vet ./src/
|
run: go vet ./...
|
||||||
|
|
||||||
- name: Staticcheck
|
- name: Staticcheck
|
||||||
uses: dominikh/staticcheck-action@v1
|
uses: dominikh/staticcheck-action@v1
|
||||||
@@ -33,14 +33,12 @@ jobs:
|
|||||||
|
|
||||||
- name: golangci-lint
|
- name: golangci-lint
|
||||||
uses: golangci/golangci-lint-action@v7
|
uses: golangci/golangci-lint-action@v7
|
||||||
with:
|
|
||||||
working-directory: src
|
|
||||||
|
|
||||||
- name: Build
|
- name: Build
|
||||||
run: go build -o trollssh ./src
|
run: go build -o trollssh ./cmd/trollssh
|
||||||
|
|
||||||
- name: Unit tests
|
- name: Unit tests
|
||||||
run: go test -race ./src/
|
run: go test -race ./...
|
||||||
|
|
||||||
docker:
|
docker:
|
||||||
needs: test
|
needs: test
|
||||||
@@ -53,7 +51,7 @@ jobs:
|
|||||||
- name: Determine push conditions
|
- name: Determine push conditions
|
||||||
id: should_push
|
id: should_push
|
||||||
run: |
|
run: |
|
||||||
if [[ "${{ github.event_name }}" == "push" && ( "${{ github.ref }}" == "refs/heads/main" || "${{ github.ref }}" == refs/tags/v* ) ]]; then
|
if [[ "${{ github.event_name }}" == "push" && ( "${{ github.ref }}" == "refs/heads/main" || "${{ github.ref }}" =~ ^refs/tags/[0-9]+\.[0-9]+\.[0-9]+$ ) ]]; then
|
||||||
echo "push=true" >> "$GITHUB_OUTPUT"
|
echo "push=true" >> "$GITHUB_OUTPUT"
|
||||||
else
|
else
|
||||||
echo "push=false" >> "$GITHUB_OUTPUT"
|
echo "push=false" >> "$GITHUB_OUTPUT"
|
||||||
@@ -66,7 +64,7 @@ jobs:
|
|||||||
images: ghcr.io/yuzuzensai/trollssh
|
images: ghcr.io/yuzuzensai/trollssh
|
||||||
tags: |
|
tags: |
|
||||||
type=raw,value=latest,enable={{is_default_branch}}
|
type=raw,value=latest,enable={{is_default_branch}}
|
||||||
type=semver,pattern=v{{version}}
|
type=semver,pattern={{version}}
|
||||||
type=sha,prefix=sha-,format=short
|
type=sha,prefix=sha-,format=short
|
||||||
|
|
||||||
- name: Log in to GHCR
|
- name: Log in to GHCR
|
||||||
+3
-2
@@ -2,8 +2,9 @@ FROM golang:1.25-alpine AS build
|
|||||||
WORKDIR /src
|
WORKDIR /src
|
||||||
COPY go.mod go.sum ./
|
COPY go.mod go.sum ./
|
||||||
RUN go mod download
|
RUN go mod download
|
||||||
COPY src ./src
|
COPY cmd ./cmd
|
||||||
RUN CGO_ENABLED=0 go build -ldflags="-s -w" -o /trollssh ./src
|
COPY internal ./internal
|
||||||
|
RUN CGO_ENABLED=0 go build -ldflags="-s -w" -o /trollssh ./cmd/trollssh
|
||||||
|
|
||||||
FROM alpine:3.22
|
FROM alpine:3.22
|
||||||
WORKDIR /home/app
|
WORKDIR /home/app
|
||||||
|
|||||||
@@ -25,7 +25,7 @@ Generate a frame set from a video through container image
|
|||||||
|
|
||||||
```sh
|
```sh
|
||||||
docker run --rm -v ./video.mp4:/home/app/video.mp4 -v ./frames:/home/app/frames \
|
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:1.2.0 trollssh --generate --video video.mp4 --resolution 512
|
||||||
```
|
```
|
||||||
|
|
||||||
This writes `frames/<name>.tsf`, a simple container of color JPEG frames plus
|
This writes `frames/<name>.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
|
Requirements: Go 1.25+ and `ffmpeg` / `ffprobe` on `PATH` (only for
|
||||||
`--generate`).
|
`--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
|
```sh
|
||||||
go run ./src --generate --video video.mp4 --resolution 512
|
go run ./cmd/trollssh --generate --video video.mp4 --resolution 512
|
||||||
go run ./src
|
go run ./cmd/trollssh
|
||||||
```
|
```
|
||||||
|
|
||||||
Or build a binary:
|
Or build a binary:
|
||||||
|
|
||||||
```sh
|
```sh
|
||||||
go build -o trollssh ./src
|
go build -o trollssh ./cmd/trollssh
|
||||||
./trollssh
|
./trollssh
|
||||||
```
|
```
|
||||||
|
|
||||||
CI runs the following checks on every push:
|
CI runs the following checks on every push:
|
||||||
|
|
||||||
```sh
|
```sh
|
||||||
gofmt -l ./src # format
|
gofmt -l . # format
|
||||||
go vet ./src/... # vet
|
go vet ./... # vet
|
||||||
golangci-lint run ./src/... # lint, see https://golangci-lint.run
|
golangci-lint run ./... # lint, see https://golangci-lint.run
|
||||||
go test ./src/... # tests
|
go test ./... # tests
|
||||||
```
|
```
|
||||||
|
|
||||||
To run them automatically before each commit, install
|
To run them automatically before each commit, install
|
||||||
|
|||||||
@@ -15,6 +15,11 @@ import (
|
|||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/joho/godotenv"
|
"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 {
|
type cliArgs struct {
|
||||||
@@ -47,7 +52,7 @@ func parseArgs(argv []string) cliArgs {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func fail(message string) {
|
func fail(message string) {
|
||||||
logError(message)
|
logx.Error(message)
|
||||||
os.Exit(1)
|
os.Exit(1)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -77,13 +82,15 @@ func generateFrames(framesDir, videoArg string, resolution int) {
|
|||||||
base := strings.TrimSuffix(filepath.Base(videoPath), filepath.Ext(videoPath))
|
base := strings.TrimSuffix(filepath.Base(videoPath), filepath.Ext(videoPath))
|
||||||
output := filepath.Join(framesDir, base+".tsf")
|
output := filepath.Join(framesDir, base+".tsf")
|
||||||
|
|
||||||
logInfo(fmt.Sprintf("Generating frames from %q -> %s", videoPath, output))
|
logx.Info(fmt.Sprintf("Generating frames from %q -> %s", videoPath, output))
|
||||||
if err := processVideo(videoPath, output, resolution); err != nil {
|
if err := tsf.ProcessVideo(videoPath, output, resolution); err != nil {
|
||||||
fail(fmt.Sprintf("Failed to generate frames from %q: %s", videoPath, err.Error()))
|
fail(fmt.Sprintf("Failed to generate frames from %q: %s", videoPath, err.Error()))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func loadAllFrames(framesDir string) []*FramesContainer {
|
const frameDataWarnBytes = 2 << 30
|
||||||
|
|
||||||
|
func loadAllFrames(framesDir string) []*tsf.FramesContainer {
|
||||||
entries, err := os.ReadDir(framesDir)
|
entries, err := os.ReadDir(framesDir)
|
||||||
var files []string
|
var files []string
|
||||||
if err == nil {
|
if err == nil {
|
||||||
@@ -101,10 +108,26 @@ func loadAllFrames(framesDir string) []*FramesContainer {
|
|||||||
framesDir,
|
framesDir,
|
||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
var totalBytes int64
|
||||||
|
for _, file := range files {
|
||||||
|
info, err := os.Stat(filepath.Join(framesDir, file))
|
||||||
|
if err != nil {
|
||||||
|
fail(err.Error())
|
||||||
|
}
|
||||||
|
totalBytes += info.Size()
|
||||||
|
}
|
||||||
|
if totalBytes > frameDataWarnBytes {
|
||||||
|
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 {
|
||||||
|
logx.Info(fmt.Sprintf("Frame data: %.1f MB", float64(totalBytes)/(1<<20)))
|
||||||
|
}
|
||||||
|
|
||||||
concurrency := min(len(files), max(1, min(runtime.NumCPU(), 4)))
|
concurrency := min(len(files), max(1, min(runtime.NumCPU(), 4)))
|
||||||
|
|
||||||
results := make([]*FramesContainer, len(files))
|
results := make([]*tsf.FramesContainer, len(files))
|
||||||
errs := make([]error, len(files))
|
errs := make([]error, len(files))
|
||||||
var next int
|
var next int
|
||||||
var nextMu sync.Mutex
|
var nextMu sync.Mutex
|
||||||
@@ -127,14 +150,14 @@ func loadAllFrames(framesDir string) []*FramesContainer {
|
|||||||
errs[i] = err
|
errs[i] = err
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
logInfo(fmt.Sprintf("Loading %s (%.1f MB)...", file, float64(info.Size())/1024/1024))
|
logx.Info(fmt.Sprintf("Loading %s (%.1f MB)...", file, float64(info.Size())/1024/1024))
|
||||||
data, err := loadTSF(filePath)
|
data, err := tsf.Load(filePath)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
errs[i] = err
|
errs[i] = err
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
data.Name = file
|
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
|
results[i] = data
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -171,17 +194,17 @@ func applyMemoryLimit() {
|
|||||||
}
|
}
|
||||||
limit := n * 9 / 10
|
limit := n * 9 / 10
|
||||||
debug.SetMemoryLimit(limit)
|
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
|
return
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func main() {
|
func main() {
|
||||||
_ = godotenv.Load()
|
_ = godotenv.Load()
|
||||||
logThreshold = resolveThreshold()
|
logx.SetThreshold(logx.ResolveThreshold())
|
||||||
applyMemoryLimit()
|
applyMemoryLimit()
|
||||||
|
|
||||||
config := loadConfig()
|
cfg := config.Load()
|
||||||
args := parseArgs(os.Args[1:])
|
args := parseArgs(os.Args[1:])
|
||||||
|
|
||||||
cwd, err := os.Getwd()
|
cwd, err := os.Getwd()
|
||||||
@@ -201,43 +224,45 @@ func main() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
var bannerText, fakeLoginText, goodbyeText *string
|
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
|
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
|
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
|
goodbyeText = &text
|
||||||
}
|
}
|
||||||
|
|
||||||
hostKeys, err := ensureHostKeys(dataDir)
|
hostKeys, err := sshserver.EnsureHostKeys(dataDir)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
fail(err.Error())
|
fail(err.Error())
|
||||||
}
|
}
|
||||||
videoSets := loadAllFrames(framesDir)
|
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{
|
server := sshserver.New(sshserver.ServerDeps{
|
||||||
Config: config,
|
Config: cfg,
|
||||||
HostKeys: hostKeys,
|
HostKeys: hostKeys,
|
||||||
BannerText: bannerText,
|
BannerText: bannerText,
|
||||||
FakeLoginText: fakeLoginText,
|
FakeLoginText: fakeLoginText,
|
||||||
GoodbyeText: goodbyeText,
|
GoodbyeText: goodbyeText,
|
||||||
VideoSets: videoSets,
|
VideoSets: videoSets,
|
||||||
})
|
})
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
sigCh := make(chan os.Signal, 1)
|
sigCh := make(chan os.Signal, 1)
|
||||||
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
|
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
|
||||||
go func() {
|
go func() {
|
||||||
sig := <-sigCh
|
sig := <-sigCh
|
||||||
logInfo(fmt.Sprintf("Received %s, shutting down...", sig))
|
logx.Info(fmt.Sprintf("Received %s, shutting down...", sig))
|
||||||
|
forceExit := time.AfterFunc(5*time.Second, func() { os.Exit(0) })
|
||||||
server.Close()
|
server.Close()
|
||||||
time.AfterFunc(5*time.Second, func() { os.Exit(0) })
|
forceExit.Stop()
|
||||||
}()
|
}()
|
||||||
|
|
||||||
if err := server.Listen(config.Host, config.Port); err != nil {
|
if err := server.Listen(cfg.Host, cfg.Port); err != nil {
|
||||||
logError("Server error:", err.Error())
|
logx.Error("Server error:", err.Error())
|
||||||
os.Exit(1)
|
os.Exit(1)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
+2
-1
@@ -1,8 +1,9 @@
|
|||||||
services:
|
services:
|
||||||
trollssh:
|
trollssh:
|
||||||
image: ghcr.io/yuzuzensai/trollssh:v1.0.1
|
image: ghcr.io/yuzuzensai/trollssh:v1.1.3
|
||||||
container_name: trollssh
|
container_name: trollssh
|
||||||
restart: unless-stopped
|
restart: unless-stopped
|
||||||
|
mem_limit: ${MEMORY_LIMIT:-2g}
|
||||||
ports:
|
ports:
|
||||||
- "22:22"
|
- "22:22"
|
||||||
env_file:
|
env_file:
|
||||||
|
|||||||
@@ -5,7 +5,7 @@ go 1.25.0
|
|||||||
require (
|
require (
|
||||||
github.com/joho/godotenv v1.5.1
|
github.com/joho/godotenv v1.5.1
|
||||||
golang.org/x/crypto v0.54.0
|
golang.org/x/crypto v0.54.0
|
||||||
golang.org/x/image v0.44.0
|
golang.org/x/image v0.45.0
|
||||||
)
|
)
|
||||||
|
|
||||||
require golang.org/x/sys v0.47.0 // indirect
|
require golang.org/x/sys v0.47.0 // indirect
|
||||||
|
|||||||
@@ -2,8 +2,8 @@ github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
|
|||||||
github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
|
github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
|
||||||
golang.org/x/crypto v0.54.0 h1:YLIA59K4fiNzHzjnZt2tUJQjQtUWfWbeHBqKtk3eScw=
|
golang.org/x/crypto v0.54.0 h1:YLIA59K4fiNzHzjnZt2tUJQjQtUWfWbeHBqKtk3eScw=
|
||||||
golang.org/x/crypto v0.54.0/go.mod h1:KWL8ny2AZdGR2cWmzeHrp2azQPGogOv+HeQaVEXC2dk=
|
golang.org/x/crypto v0.54.0/go.mod h1:KWL8ny2AZdGR2cWmzeHrp2azQPGogOv+HeQaVEXC2dk=
|
||||||
golang.org/x/image v0.44.0 h1:+tDekMZED9+LrtB3G5xzRggpVh9CARjZqROla3R3R+I=
|
golang.org/x/image v0.45.0 h1:FMb1nTbH5H9vF55SriQHgFw5GnNL9Jg6L25BwXKzhB0=
|
||||||
golang.org/x/image v0.44.0/go.mod h1:V8K3KE9KKKE+pLpQDOeN18w9oacNSvy1tDOirTu4xtY=
|
golang.org/x/image v0.45.0/go.mod h1:n62x/7RqlwXDvGsSU4u6IUTUf6KghUZ9Bt7cG/T9Fx4=
|
||||||
golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
|
golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
|
||||||
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||||
golang.org/x/term v0.45.0 h1:NwWyBmoJCbfTHpxrWoZ9C6/VxOf7ic219I8xZZFdrf0=
|
golang.org/x/term v0.45.0 h1:NwWyBmoJCbfTHpxrWoZ9C6/VxOf7ic219I8xZZFdrf0=
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
package main
|
package config
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
@@ -6,6 +6,8 @@ import (
|
|||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"github.com/YuzuZensai/TrollSSH/internal/logx"
|
||||||
)
|
)
|
||||||
|
|
||||||
type PlaybackMode string
|
type PlaybackMode string
|
||||||
@@ -28,7 +30,10 @@ type Config struct {
|
|||||||
MaxAuthAttempts int
|
MaxAuthAttempts int
|
||||||
HandshakeTimeout time.Duration
|
HandshakeTimeout time.Duration
|
||||||
MaxDimension int
|
MaxDimension int
|
||||||
|
MaxTerminalCells int
|
||||||
|
SessionTimeout time.Duration
|
||||||
RenderCacheMB int
|
RenderCacheMB int
|
||||||
|
RenderCacheCompress bool
|
||||||
BrightnessThreshold int
|
BrightnessThreshold int
|
||||||
Charset string
|
Charset string
|
||||||
Invert bool
|
Invert bool
|
||||||
@@ -37,7 +42,7 @@ type Config struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func warnInvalid(name, value string, fallback any) {
|
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 {
|
func envString(name, fallback string) string {
|
||||||
@@ -111,7 +116,7 @@ func envPlaybackMode(name string, fallback PlaybackMode) PlaybackMode {
|
|||||||
return fallback
|
return fallback
|
||||||
}
|
}
|
||||||
|
|
||||||
func loadConfig() Config {
|
func Load() Config {
|
||||||
const maxInt = int(^uint(0) >> 1)
|
const maxInt = int(^uint(0) >> 1)
|
||||||
return Config{
|
return Config{
|
||||||
Host: envString("HOST", "0.0.0.0"),
|
Host: envString("HOST", "0.0.0.0"),
|
||||||
@@ -126,7 +131,10 @@ func loadConfig() Config {
|
|||||||
MaxAuthAttempts: envInt("MAX_AUTH_ATTEMPTS", 6, 1, maxInt),
|
MaxAuthAttempts: envInt("MAX_AUTH_ATTEMPTS", 6, 1, maxInt),
|
||||||
HandshakeTimeout: envDurationMs("HANDSHAKE_TIMEOUT", 10*time.Second),
|
HandshakeTimeout: envDurationMs("HANDSHAKE_TIMEOUT", 10*time.Second),
|
||||||
MaxDimension: envInt("MAX_DIMENSION", 512, 1, 4096),
|
MaxDimension: envInt("MAX_DIMENSION", 512, 1, 4096),
|
||||||
|
MaxTerminalCells: envInt("MAX_TERMINAL_CELLS", 500*512, 1, maxInt),
|
||||||
|
SessionTimeout: envDurationMs("SESSION_TIMEOUT", 10*time.Minute),
|
||||||
RenderCacheMB: envInt("RENDER_CACHE_MB", 256, 0, maxInt),
|
RenderCacheMB: envInt("RENDER_CACHE_MB", 256, 0, maxInt),
|
||||||
|
RenderCacheCompress: envBool("RENDER_CACHE_COMPRESS", false),
|
||||||
BrightnessThreshold: envInt("BRIGHTNESS_THRESHOLD", 40, 0, 100),
|
BrightnessThreshold: envInt("BRIGHTNESS_THRESHOLD", 40, 0, 100),
|
||||||
Charset: envString("CHARSET", "detailed"),
|
Charset: envString("CHARSET", "detailed"),
|
||||||
Invert: envBool("INVERT", false),
|
Invert: envBool("INVERT", false),
|
||||||
@@ -135,7 +143,7 @@ func loadConfig() Config {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func loadOptionalTextFile(filePath string) (string, bool) {
|
func LoadOptionalTextFile(filePath string) (string, bool) {
|
||||||
data, err := os.ReadFile(filePath)
|
data, err := os.ReadFile(filePath)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return "", false
|
return "", false
|
||||||
@@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,89 @@
|
|||||||
|
package logx
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Level int
|
||||||
|
|
||||||
|
const (
|
||||||
|
LevelDebug Level = 10
|
||||||
|
LevelInfo Level = 20
|
||||||
|
LevelWarn Level = 30
|
||||||
|
LevelError Level = 40
|
||||||
|
)
|
||||||
|
|
||||||
|
var threshold = ResolveThreshold()
|
||||||
|
|
||||||
|
func ResolveThreshold() Level {
|
||||||
|
switch strings.ToLower(strings.TrimSpace(os.Getenv("LOG_LEVEL"))) {
|
||||||
|
case "debug":
|
||||||
|
return LevelDebug
|
||||||
|
case "warn":
|
||||||
|
return LevelWarn
|
||||||
|
case "error":
|
||||||
|
return LevelError
|
||||||
|
default:
|
||||||
|
return LevelInfo
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func SetThreshold(level Level) { threshold = level }
|
||||||
|
|
||||||
|
func Sanitize(value any) string {
|
||||||
|
return SanitizeN(value, 200)
|
||||||
|
}
|
||||||
|
|
||||||
|
func SanitizeN(value any, maxLength int) string {
|
||||||
|
var str string
|
||||||
|
switch v := value.(type) {
|
||||||
|
case nil:
|
||||||
|
str = ""
|
||||||
|
case string:
|
||||||
|
str = v
|
||||||
|
default:
|
||||||
|
str = fmt.Sprint(v)
|
||||||
|
}
|
||||||
|
var b strings.Builder
|
||||||
|
count := 0
|
||||||
|
for _, r := range str {
|
||||||
|
if count >= maxLength {
|
||||||
|
b.WriteRune('…')
|
||||||
|
return b.String()
|
||||||
|
}
|
||||||
|
if r < 0x20 || (r >= 0x7f && r <= 0x9f) {
|
||||||
|
b.WriteRune('�')
|
||||||
|
} else {
|
||||||
|
b.WriteRune(r)
|
||||||
|
}
|
||||||
|
count++
|
||||||
|
}
|
||||||
|
return b.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
func emit(level Level, name string, stream *os.File, args []any) {
|
||||||
|
if level < threshold {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
parts := make([]string, len(args))
|
||||||
|
for i, a := range args {
|
||||||
|
if s, ok := a.(string); ok {
|
||||||
|
parts[i] = s
|
||||||
|
} else if b, err := json.Marshal(a); err == nil {
|
||||||
|
parts[i] = string(b)
|
||||||
|
} else {
|
||||||
|
parts[i] = fmt.Sprint(a)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
ts := time.Now().UTC().Format("2006-01-02T15:04:05.000Z")
|
||||||
|
_, _ = fmt.Fprintf(stream, "[%s] %-5s %s\n", ts, strings.ToUpper(name), strings.Join(parts, " "))
|
||||||
|
}
|
||||||
|
|
||||||
|
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) }
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,278 @@
|
|||||||
|
package render
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"compress/flate"
|
||||||
|
"container/list"
|
||||||
|
"io"
|
||||||
|
"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
|
||||||
|
compress bool
|
||||||
|
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
|
||||||
|
data []byte
|
||||||
|
origLen int
|
||||||
|
cost int64
|
||||||
|
}
|
||||||
|
|
||||||
|
func entryCost(_ cacheKey, data []byte) int64 {
|
||||||
|
return int64(cap(data)) + 160
|
||||||
|
}
|
||||||
|
|
||||||
|
var flateWriters = sync.Pool{New: func() any {
|
||||||
|
w, _ := flate.NewWriter(io.Discard, 1)
|
||||||
|
return w
|
||||||
|
}}
|
||||||
|
|
||||||
|
type flateReader interface {
|
||||||
|
io.Reader
|
||||||
|
flate.Resetter
|
||||||
|
}
|
||||||
|
|
||||||
|
var flateReaders = sync.Pool{New: func() any {
|
||||||
|
return flate.NewReader(bytes.NewReader(nil)).(flateReader)
|
||||||
|
}}
|
||||||
|
|
||||||
|
func compressAscii(src []byte) []byte {
|
||||||
|
var buf bytes.Buffer
|
||||||
|
buf.Grow(len(src)/3 + 64)
|
||||||
|
w := flateWriters.Get().(*flate.Writer)
|
||||||
|
w.Reset(&buf)
|
||||||
|
_, _ = w.Write(src)
|
||||||
|
_ = w.Close()
|
||||||
|
flateWriters.Put(w)
|
||||||
|
return bytes.Clone(buf.Bytes())
|
||||||
|
}
|
||||||
|
|
||||||
|
func decompressAscii(src []byte, origLen int) []byte {
|
||||||
|
r := flateReaders.Get().(flateReader)
|
||||||
|
_ = r.Reset(bytes.NewReader(src), nil)
|
||||||
|
buf := bytes.NewBuffer(make([]byte, 0, origLen))
|
||||||
|
_, _ = io.Copy(buf, r)
|
||||||
|
flateReaders.Put(r)
|
||||||
|
return buf.Bytes()
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewCache(maxBytes int64, compress bool) *Cache {
|
||||||
|
if maxBytes <= 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
shardCount := int(min(int64(16), max(int64(1), maxBytes/(1<<20))))
|
||||||
|
cache := &Cache{shards: make([]cacheShard, shardCount), compress: compress}
|
||||||
|
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()
|
||||||
|
el, ok := shard.entries[key]
|
||||||
|
if !ok {
|
||||||
|
shard.mu.Unlock()
|
||||||
|
c.misses.Add(1)
|
||||||
|
return nil, false
|
||||||
|
}
|
||||||
|
shard.order.MoveToBack(el)
|
||||||
|
entry := el.Value.(*cacheEntry)
|
||||||
|
data, origLen := entry.data, entry.origLen
|
||||||
|
shard.mu.Unlock()
|
||||||
|
c.hits.Add(1)
|
||||||
|
if !c.compress {
|
||||||
|
return data, true
|
||||||
|
}
|
||||||
|
return decompressAscii(data, origLen), true
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Cache) put(key cacheKey, ascii []byte) {
|
||||||
|
if c == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
shard := c.shard(key)
|
||||||
|
data, origLen := ascii, 0
|
||||||
|
if c.compress {
|
||||||
|
data, origLen = compressAscii(ascii), len(ascii)
|
||||||
|
}
|
||||||
|
cost := entryCost(key, data)
|
||||||
|
if cost > shard.maxBytes {
|
||||||
|
c.rejections.Add(1)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
shard.mu.Lock()
|
||||||
|
defer shard.mu.Unlock()
|
||||||
|
if _, ok := shard.entries[key]; ok {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
entry := &cacheEntry{key: key, data: data, origLen: origLen, cost: cost}
|
||||||
|
shard.entries[key] = shard.order.PushBack(entry)
|
||||||
|
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)
|
||||||
|
|
||||||
|
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
|
||||||
|
}
|
||||||
@@ -1,9 +1,7 @@
|
|||||||
package main
|
package render
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"bytes"
|
"bytes"
|
||||||
"container/list"
|
|
||||||
"fmt"
|
|
||||||
"image"
|
"image"
|
||||||
"image/color"
|
"image/color"
|
||||||
"image/jpeg"
|
"image/jpeg"
|
||||||
@@ -15,36 +13,30 @@ import (
|
|||||||
"golang.org/x/image/draw"
|
"golang.org/x/image/draw"
|
||||||
)
|
)
|
||||||
|
|
||||||
type FramesContainer struct {
|
type ColorTier int
|
||||||
ColorFrames [][]byte
|
|
||||||
FPS float64
|
|
||||||
Name string
|
|
||||||
}
|
|
||||||
|
|
||||||
type colorTier int
|
|
||||||
|
|
||||||
const (
|
const (
|
||||||
colorTierNone colorTier = iota
|
ColorTierNone ColorTier = iota
|
||||||
colorTier256
|
ColorTier256
|
||||||
colorTierTrueColor
|
ColorTierTrueColor
|
||||||
)
|
)
|
||||||
|
|
||||||
func detectColorTier(term string) colorTier {
|
func DetectColorTier(term string) ColorTier {
|
||||||
t := strings.ToLower(strings.TrimSpace(term))
|
t := strings.ToLower(strings.TrimSpace(term))
|
||||||
switch t {
|
switch t {
|
||||||
case "", "dumb", "vt52", "vt100", "vt102", "vt220", "ansi", "linux", "cons25", "cygwin":
|
case "", "dumb", "vt52", "vt100", "vt102", "vt220", "ansi", "linux", "cons25", "cygwin":
|
||||||
return colorTierNone
|
return ColorTierNone
|
||||||
}
|
}
|
||||||
if strings.Contains(t, "direct") || strings.Contains(t, "truecolor") {
|
if strings.Contains(t, "direct") || strings.Contains(t, "truecolor") {
|
||||||
return colorTierTrueColor
|
return ColorTierTrueColor
|
||||||
}
|
}
|
||||||
if strings.Contains(t, "256color") {
|
if strings.Contains(t, "256color") {
|
||||||
return colorTier256
|
return ColorTier256
|
||||||
}
|
}
|
||||||
if strings.HasPrefix(t, "screen") || strings.HasPrefix(t, "tmux") {
|
if strings.HasPrefix(t, "screen") || strings.HasPrefix(t, "tmux") {
|
||||||
return colorTier256
|
return ColorTier256
|
||||||
}
|
}
|
||||||
return colorTierTrueColor
|
return ColorTierTrueColor
|
||||||
}
|
}
|
||||||
|
|
||||||
var charsetPresets = map[string]string{
|
var charsetPresets = map[string]string{
|
||||||
@@ -64,6 +56,8 @@ func resolveCharset(charset string) string {
|
|||||||
return charset
|
return charset
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const maxPooledBuffer = 4 << 20
|
||||||
|
|
||||||
var pixPool sync.Pool
|
var pixPool sync.Pool
|
||||||
|
|
||||||
func getPixBuf(n int) []byte {
|
func getPixBuf(n int) []byte {
|
||||||
@@ -76,7 +70,9 @@ func getPixBuf(n int) []byte {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func putPixBuf(b []byte) {
|
func putPixBuf(b []byte) {
|
||||||
pixPool.Put(&b)
|
if cap(b) <= maxPooledBuffer {
|
||||||
|
pixPool.Put(&b)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
var outPool sync.Pool
|
var outPool sync.Pool
|
||||||
@@ -91,7 +87,9 @@ func getOutBuf(capacity int) []byte {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func putOutBuf(b []byte) {
|
func putOutBuf(b []byte) {
|
||||||
outPool.Put(&b)
|
if cap(b) <= maxPooledBuffer {
|
||||||
|
outPool.Put(&b)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func resizeFrame(frame, pix []byte, width, height int, keepAspectRatio bool) (*image.RGBA, error) {
|
func resizeFrame(frame, pix []byte, width, height int, keepAspectRatio bool) (*image.RGBA, error) {
|
||||||
@@ -118,16 +116,16 @@ func resizeFrame(frame, pix []byte, width, height int, keepAspectRatio bool) (*i
|
|||||||
return dst, nil
|
return dst, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
type asciiOptions struct {
|
type Options struct {
|
||||||
brightnessThreshold int
|
BrightnessThreshold int
|
||||||
charset string
|
Charset string
|
||||||
invert bool
|
Invert bool
|
||||||
}
|
}
|
||||||
|
|
||||||
func buildRampLUT(ramp []rune, options asciiOptions) *[101][]byte {
|
func buildRampLUT(ramp []rune, options Options) *[101][]byte {
|
||||||
var lut [101][]byte
|
var lut [101][]byte
|
||||||
for b := range lut {
|
for b := range lut {
|
||||||
index := rampIndex(b, options.brightnessThreshold, len(ramp), options.invert)
|
index := rampIndex(b, options.BrightnessThreshold, len(ramp), options.Invert)
|
||||||
lut[b] = utf8.AppendRune(nil, ramp[index])
|
lut[b] = utf8.AppendRune(nil, ramp[index])
|
||||||
}
|
}
|
||||||
return &lut
|
return &lut
|
||||||
@@ -149,16 +147,20 @@ func rampIndex(brightness, threshold, total int, invert bool) int {
|
|||||||
return index
|
return index
|
||||||
}
|
}
|
||||||
|
|
||||||
func frameToAscii(img *image.RGBA, rampLUT *[101][]byte) string {
|
func frameToAscii(img *image.RGBA, rampLUT *[101][]byte) []byte {
|
||||||
pix := img.Pix
|
pix := img.Pix
|
||||||
buf := getOutBuf(len(pix))
|
maxCharBytes := 1
|
||||||
|
for _, char := range rampLUT {
|
||||||
|
maxCharBytes = max(maxCharBytes, len(char))
|
||||||
|
}
|
||||||
|
buf := getOutBuf(len(pix) / 4 * maxCharBytes)
|
||||||
for o := 0; o < len(pix); o += 4 {
|
for o := 0; o < len(pix); o += 4 {
|
||||||
brightness := (int(pix[o])*299 + int(pix[o+1])*587 + int(pix[o+2])*114) / 255 / 10
|
brightness := (int(pix[o])*299 + int(pix[o+1])*587 + int(pix[o+2])*114) / 255 / 10
|
||||||
buf = append(buf, rampLUT[brightness]...)
|
buf = append(buf, rampLUT[brightness]...)
|
||||||
}
|
}
|
||||||
ascii := string(buf)
|
output := bytes.Clone(buf)
|
||||||
putOutBuf(buf)
|
putOutBuf(buf)
|
||||||
return ascii
|
return output
|
||||||
}
|
}
|
||||||
|
|
||||||
const ansiReset = "\x1b[0m"
|
const ansiReset = "\x1b[0m"
|
||||||
@@ -193,8 +195,8 @@ func quantize256(r, g, b uint8) int {
|
|||||||
return 16 + 36*int(ansi256Cube[r]) + 6*int(ansi256Cube[g]) + int(ansi256Cube[b])
|
return 16 + 36*int(ansi256Cube[r]) + 6*int(ansi256Cube[g]) + int(ansi256Cube[b])
|
||||||
}
|
}
|
||||||
|
|
||||||
func appendColor(buf []byte, r, g, b uint8, tier colorTier) []byte {
|
func appendColor(buf []byte, r, g, b uint8, tier ColorTier) []byte {
|
||||||
if tier == colorTierTrueColor {
|
if tier == ColorTierTrueColor {
|
||||||
buf = append(buf, "\x1b[38;2;"...)
|
buf = append(buf, "\x1b[38;2;"...)
|
||||||
buf = append(buf, decimal[r]...)
|
buf = append(buf, decimal[r]...)
|
||||||
buf = append(buf, ';')
|
buf = append(buf, ';')
|
||||||
@@ -208,167 +210,41 @@ func appendColor(buf []byte, r, g, b uint8, tier colorTier) []byte {
|
|||||||
return append(buf, 'm')
|
return append(buf, 'm')
|
||||||
}
|
}
|
||||||
|
|
||||||
func frameToAnsi(img *image.RGBA, rampLUT *[101][]byte, tier colorTier) string {
|
func frameToAnsi(img *image.RGBA, rampLUT *[101][]byte, tier ColorTier) []byte {
|
||||||
bounds := img.Bounds()
|
bounds := img.Bounds()
|
||||||
buf := getOutBuf(bounds.Dx() * bounds.Dy() * 16)
|
bytesPerCell := 11
|
||||||
|
if tier == ColorTierTrueColor {
|
||||||
|
bytesPerCell = 16
|
||||||
|
}
|
||||||
|
buf := getOutBuf(bounds.Dx() * bounds.Dy() * bytesPerCell)
|
||||||
var lastR, lastG, lastB uint8
|
var lastR, lastG, lastB uint8
|
||||||
|
last256 := -1
|
||||||
first := true
|
first := true
|
||||||
for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
|
for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
|
||||||
|
o := img.PixOffset(bounds.Min.X, y)
|
||||||
for x := bounds.Min.X; x < bounds.Max.X; x++ {
|
for x := bounds.Min.X; x < bounds.Max.X; x++ {
|
||||||
o := img.PixOffset(x, y)
|
|
||||||
r, g, bl := img.Pix[o], img.Pix[o+1], img.Pix[o+2]
|
r, g, bl := img.Pix[o], img.Pix[o+1], img.Pix[o+2]
|
||||||
brightness := (int(r)*299 + int(g)*587 + int(bl)*114) / 255 / 10
|
brightness := (int(r)*299 + int(g)*587 + int(bl)*114) / 255 / 10
|
||||||
if first || r != lastR || g != lastG || bl != lastB {
|
colorChanged := first || r != lastR || g != lastG || bl != lastB
|
||||||
|
if tier == ColorTier256 {
|
||||||
|
index := quantize256(r, g, bl)
|
||||||
|
colorChanged = first || index != last256
|
||||||
|
last256 = index
|
||||||
|
}
|
||||||
|
if colorChanged {
|
||||||
buf = appendColor(buf, r, g, bl, tier)
|
buf = appendColor(buf, r, g, bl, tier)
|
||||||
lastR, lastG, lastB = r, g, bl
|
lastR, lastG, lastB = r, g, bl
|
||||||
first = false
|
first = false
|
||||||
}
|
}
|
||||||
buf = append(buf, rampLUT[brightness]...)
|
buf = append(buf, rampLUT[brightness]...)
|
||||||
|
o += 4
|
||||||
}
|
}
|
||||||
if y < bounds.Max.Y-1 {
|
if y < bounds.Max.Y-1 {
|
||||||
buf = append(buf, ansiReset+"\r\n"...)
|
buf = append(buf, "\r\n"...)
|
||||||
first = true
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
buf = append(buf, ansiReset...)
|
buf = append(buf, ansiReset...)
|
||||||
ascii := string(buf)
|
output := bytes.Clone(buf)
|
||||||
putOutBuf(buf)
|
putOutBuf(buf)
|
||||||
return ascii
|
return output
|
||||||
}
|
|
||||||
|
|
||||||
type renderCache struct {
|
|
||||||
mu sync.Mutex
|
|
||||||
maxBytes int64
|
|
||||||
size int64
|
|
||||||
entries map[string]*list.Element
|
|
||||||
order *list.List
|
|
||||||
}
|
|
||||||
|
|
||||||
type cacheEntry struct {
|
|
||||||
key string
|
|
||||||
ascii string
|
|
||||||
}
|
|
||||||
|
|
||||||
func entryCost(key, ascii string) int64 {
|
|
||||||
return int64(len(key)+len(ascii)) + 128
|
|
||||||
}
|
|
||||||
|
|
||||||
func newRenderCache(maxBytes int64) *renderCache {
|
|
||||||
if maxBytes <= 0 {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
return &renderCache{
|
|
||||||
maxBytes: maxBytes,
|
|
||||||
entries: make(map[string]*list.Element),
|
|
||||||
order: list.New(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (c *renderCache) get(key string) (string, bool) {
|
|
||||||
if c == nil {
|
|
||||||
return "", false
|
|
||||||
}
|
|
||||||
c.mu.Lock()
|
|
||||||
defer c.mu.Unlock()
|
|
||||||
el, ok := c.entries[key]
|
|
||||||
if !ok {
|
|
||||||
return "", false
|
|
||||||
}
|
|
||||||
c.order.MoveToBack(el)
|
|
||||||
return el.Value.(*cacheEntry).ascii, true
|
|
||||||
}
|
|
||||||
|
|
||||||
func (c *renderCache) put(key, ascii string) {
|
|
||||||
if c == nil {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
cost := entryCost(key, ascii)
|
|
||||||
if cost > c.maxBytes {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
c.mu.Lock()
|
|
||||||
defer c.mu.Unlock()
|
|
||||||
if _, ok := c.entries[key]; ok {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
c.entries[key] = c.order.PushBack(&cacheEntry{key, ascii})
|
|
||||||
c.size += cost
|
|
||||||
for c.size > c.maxBytes {
|
|
||||||
oldest := c.order.Front()
|
|
||||||
c.order.Remove(oldest)
|
|
||||||
evicted := oldest.Value.(*cacheEntry)
|
|
||||||
delete(c.entries, evicted.key)
|
|
||||||
c.size -= entryCost(evicted.key, evicted.ascii)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
type FrameRenderer struct {
|
|
||||||
setID int
|
|
||||||
colorFrames [][]byte
|
|
||||||
options asciiOptions
|
|
||||||
rampLUT *[101][]byte
|
|
||||||
cache *renderCache
|
|
||||||
|
|
||||||
inflightMu sync.Mutex
|
|
||||||
inflight map[string]chan struct{}
|
|
||||||
}
|
|
||||||
|
|
||||||
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[string]chan struct{}),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (r *FrameRenderer) render(index, width, height int, keepAspectRatio bool, tier colorTier) (string, error) {
|
|
||||||
key := fmt.Sprintf("%d:%d:%dx%d:%t:%d", r.setID, index, width, height, keepAspectRatio, tier)
|
|
||||||
if ascii, ok := r.cache.get(key); ok {
|
|
||||||
return ascii, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
if r.cache != nil {
|
|
||||||
for {
|
|
||||||
r.inflightMu.Lock()
|
|
||||||
wait, ok := r.inflight[key]
|
|
||||||
if !ok {
|
|
||||||
done := make(chan struct{})
|
|
||||||
r.inflight[key] = done
|
|
||||||
r.inflightMu.Unlock()
|
|
||||||
defer func() {
|
|
||||||
r.inflightMu.Lock()
|
|
||||||
delete(r.inflight, key)
|
|
||||||
r.inflightMu.Unlock()
|
|
||||||
close(done)
|
|
||||||
}()
|
|
||||||
break
|
|
||||||
}
|
|
||||||
r.inflightMu.Unlock()
|
|
||||||
<-wait
|
|
||||||
if ascii, ok := r.cache.get(key); ok {
|
|
||||||
return ascii, nil
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pix := getPixBuf(4 * width * height)
|
|
||||||
img, err := resizeFrame(r.colorFrames[index], pix, width, height, keepAspectRatio)
|
|
||||||
if err != nil {
|
|
||||||
putPixBuf(pix)
|
|
||||||
return "", err
|
|
||||||
}
|
|
||||||
var ascii string
|
|
||||||
if tier == colorTierNone {
|
|
||||||
ascii = frameToAscii(img, r.rampLUT)
|
|
||||||
} else {
|
|
||||||
ascii = frameToAnsi(img, r.rampLUT, tier)
|
|
||||||
}
|
|
||||||
putPixBuf(pix)
|
|
||||||
|
|
||||||
r.cache.put(key, ascii)
|
|
||||||
return ascii, nil
|
|
||||||
}
|
}
|
||||||
@@ -0,0 +1,77 @@
|
|||||||
|
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) {
|
||||||
|
for _, compress := range []bool{false, true} {
|
||||||
|
name := "uncompressed"
|
||||||
|
if compress {
|
||||||
|
name = "compressed"
|
||||||
|
}
|
||||||
|
b.Run(name, func(b *testing.B) {
|
||||||
|
fc := loadBenchSet(b)
|
||||||
|
r := NewRenderer(0, fc.ColorFrames, Options{
|
||||||
|
BrightnessThreshold: 40,
|
||||||
|
Charset: "detailed",
|
||||||
|
}, NewCache(8<<20, compress))
|
||||||
|
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())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,222 @@
|
|||||||
|
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, false))
|
||||||
|
|
||||||
|
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 incompressible(n int) []byte {
|
||||||
|
b := make([]byte, n)
|
||||||
|
x := uint32(0x9e3779b9)
|
||||||
|
for i := range b {
|
||||||
|
x ^= x << 13
|
||||||
|
x ^= x >> 17
|
||||||
|
x ^= x << 5
|
||||||
|
b[i] = byte(x)
|
||||||
|
}
|
||||||
|
return b
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRenderCacheEvictsByBytes(t *testing.T) {
|
||||||
|
key := func(index int) cacheKey { return cacheKey{index: index} }
|
||||||
|
payload := incompressible(1000)
|
||||||
|
budget := 3 * entryCost(key(0), compressAscii(payload))
|
||||||
|
c := NewCache(budget, true)
|
||||||
|
for i := range 5 {
|
||||||
|
c.put(key(i), payload)
|
||||||
|
}
|
||||||
|
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 TestRenderCacheRoundTrips(t *testing.T) {
|
||||||
|
for _, compress := range []bool{false, true} {
|
||||||
|
c := NewCache(1<<20, compress)
|
||||||
|
want := incompressible(4096)
|
||||||
|
c.put(cacheKey{}, want)
|
||||||
|
got, ok := c.get(cacheKey{})
|
||||||
|
if !ok {
|
||||||
|
t.Fatalf("compress=%v: entry should be cached", compress)
|
||||||
|
}
|
||||||
|
if !bytes.Equal(got, want) {
|
||||||
|
t.Fatalf("compress=%v: entry does not match original", compress)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRenderCacheDisabled(t *testing.T) {
|
||||||
|
c := NewCache(0, false)
|
||||||
|
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, true)
|
||||||
|
c.put(cacheKey{}, incompressible(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 TestRenderCacheAccountsCompressedSize(t *testing.T) {
|
||||||
|
cache := NewCache(512, true)
|
||||||
|
value := make([]byte, 1, 4096)
|
||||||
|
cache.put(cacheKey{}, value)
|
||||||
|
if _, ok := cache.get(cacheKey{}); !ok {
|
||||||
|
t.Fatal("small value should be cached regardless of its backing capacity")
|
||||||
|
}
|
||||||
|
if got := cache.size.Load(); got > 512 {
|
||||||
|
t.Fatalf("size = %d, want <= 512", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,4 +1,4 @@
|
|||||||
package main
|
package sshserver
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"crypto/ed25519"
|
"crypto/ed25519"
|
||||||
@@ -33,7 +33,7 @@ func generateAndSave(keyPath, keyType string) error {
|
|||||||
return os.WriteFile(keyPath, pem.EncodeToMemory(block), 0o600)
|
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 }{
|
keys := []struct{ file, keyType string }{
|
||||||
{"id_rsa", "rsa"},
|
{"id_rsa", "rsa"},
|
||||||
{"id_ed25519", "ed25519"},
|
{"id_ed25519", "ed25519"},
|
||||||
@@ -0,0 +1,797 @@
|
|||||||
|
package sshserver
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/binary"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"math"
|
||||||
|
"math/rand"
|
||||||
|
"net"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"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 (
|
||||||
|
clearScreen = "\x1b[2J\x1b[0f"
|
||||||
|
hideCursor = "\x1b[?25l"
|
||||||
|
showCursor = "\x1b[?25h"
|
||||||
|
syncStart = "\x1b[?2026h"
|
||||||
|
syncEnd = "\x1b[?2026l"
|
||||||
|
homeCursor = "\x1b[H"
|
||||||
|
maxSessionsPerConn = 1
|
||||||
|
terminalSizeQuantum = 4
|
||||||
|
resizeDebounce = 200 * time.Millisecond
|
||||||
|
outputStallTimeout = 15 * time.Second
|
||||||
|
)
|
||||||
|
|
||||||
|
var errOutputStalled = errors.New("SSH output stalled")
|
||||||
|
|
||||||
|
func writePartsWithTimeout(
|
||||||
|
conn *ssh.ServerConn,
|
||||||
|
channel ssh.Channel,
|
||||||
|
timeout time.Duration,
|
||||||
|
parts ...string,
|
||||||
|
) error {
|
||||||
|
write := func() error {
|
||||||
|
for _, part := range parts {
|
||||||
|
if _, err := io.WriteString(channel, part); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if timeout <= 0 {
|
||||||
|
return write()
|
||||||
|
}
|
||||||
|
|
||||||
|
fired := make(chan struct{})
|
||||||
|
timer := time.AfterFunc(timeout, func() {
|
||||||
|
_ = conn.Close()
|
||||||
|
close(fired)
|
||||||
|
})
|
||||||
|
err := write()
|
||||||
|
if timer.Stop() {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
<-fired
|
||||||
|
return errOutputStalled
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeFrameWithTimeout(
|
||||||
|
conn *ssh.ServerConn,
|
||||||
|
channel ssh.Channel,
|
||||||
|
timeout time.Duration,
|
||||||
|
prefix string,
|
||||||
|
frame []byte,
|
||||||
|
) error {
|
||||||
|
write := func() error {
|
||||||
|
if _, err := io.WriteString(channel, syncStart+prefix); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if _, err := channel.Write(frame); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
_, err := io.WriteString(channel, syncEnd)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if timeout <= 0 {
|
||||||
|
return write()
|
||||||
|
}
|
||||||
|
fired := make(chan struct{})
|
||||||
|
timer := time.AfterFunc(timeout, func() {
|
||||||
|
_ = conn.Close()
|
||||||
|
close(fired)
|
||||||
|
})
|
||||||
|
err := write()
|
||||||
|
if timer.Stop() {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
<-fired
|
||||||
|
return errOutputStalled
|
||||||
|
}
|
||||||
|
|
||||||
|
type ConnectionTracker struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
counts map[string]int
|
||||||
|
total int
|
||||||
|
}
|
||||||
|
|
||||||
|
func newConnectionTracker() *ConnectionTracker {
|
||||||
|
return &ConnectionTracker{counts: make(map[string]int)}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *ConnectionTracker) tryAcquire(ip string, maxPerIP, maxTotal int) (int, int, bool) {
|
||||||
|
t.mu.Lock()
|
||||||
|
defer t.mu.Unlock()
|
||||||
|
if t.total >= maxTotal || t.counts[ip] >= maxPerIP {
|
||||||
|
return t.counts[ip], t.total, false
|
||||||
|
}
|
||||||
|
t.counts[ip]++
|
||||||
|
t.total++
|
||||||
|
return t.counts[ip], t.total, true
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *ConnectionTracker) release(ip string) {
|
||||||
|
t.mu.Lock()
|
||||||
|
defer t.mu.Unlock()
|
||||||
|
if _, ok := t.counts[ip]; !ok {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
t.counts[ip]--
|
||||||
|
if t.total > 0 {
|
||||||
|
t.total--
|
||||||
|
}
|
||||||
|
if t.counts[ip] <= 0 {
|
||||||
|
delete(t.counts, ip)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *ConnectionTracker) totalCount() int {
|
||||||
|
t.mu.Lock()
|
||||||
|
defer t.mu.Unlock()
|
||||||
|
return t.total
|
||||||
|
}
|
||||||
|
|
||||||
|
type SessionTracker struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
perConn map[*ssh.ServerConn]int
|
||||||
|
total int
|
||||||
|
}
|
||||||
|
|
||||||
|
func newSessionTracker() *SessionTracker {
|
||||||
|
return &SessionTracker{perConn: make(map[*ssh.ServerConn]int)}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *SessionTracker) tryAcquire(conn *ssh.ServerConn, maxPerConn, maxTotal int) bool {
|
||||||
|
t.mu.Lock()
|
||||||
|
defer t.mu.Unlock()
|
||||||
|
if t.total >= maxTotal || t.perConn[conn] >= maxPerConn {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
t.perConn[conn]++
|
||||||
|
t.total++
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *SessionTracker) release(conn *ssh.ServerConn) {
|
||||||
|
t.mu.Lock()
|
||||||
|
defer t.mu.Unlock()
|
||||||
|
count := t.perConn[conn]
|
||||||
|
if count <= 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if count == 1 {
|
||||||
|
delete(t.perConn, conn)
|
||||||
|
} else {
|
||||||
|
t.perConn[conn] = count - 1
|
||||||
|
}
|
||||||
|
t.total--
|
||||||
|
}
|
||||||
|
|
||||||
|
type frameSet struct {
|
||||||
|
data *tsf.FramesContainer
|
||||||
|
renderer *render.Renderer
|
||||||
|
}
|
||||||
|
|
||||||
|
type Server struct {
|
||||||
|
config config.Config
|
||||||
|
sshConfig *ssh.ServerConfig
|
||||||
|
sets []frameSet
|
||||||
|
cache *render.Cache
|
||||||
|
tracker *ConnectionTracker
|
||||||
|
sessions *SessionTracker
|
||||||
|
fakeLogin *string
|
||||||
|
goodbye *string
|
||||||
|
mu sync.Mutex
|
||||||
|
listener net.Listener
|
||||||
|
conns map[net.Conn]struct{}
|
||||||
|
connWG sync.WaitGroup
|
||||||
|
closing bool
|
||||||
|
closeOnce sync.Once
|
||||||
|
}
|
||||||
|
|
||||||
|
type ServerDeps struct {
|
||||||
|
Config config.Config
|
||||||
|
HostKeys []ssh.Signer
|
||||||
|
BannerText *string
|
||||||
|
FakeLoginText *string
|
||||||
|
GoodbyeText *string
|
||||||
|
VideoSets []*tsf.FramesContainer
|
||||||
|
}
|
||||||
|
|
||||||
|
func clampTermSize(cols, rows, maxDimension, maxCells, quantum int) (int, int) {
|
||||||
|
cols = max(cols, 1)
|
||||||
|
rows = max(rows, 1)
|
||||||
|
scale := min(1.0, float64(maxDimension)/float64(cols), float64(maxDimension)/float64(rows))
|
||||||
|
area := float64(cols) * float64(rows)
|
||||||
|
if area*scale*scale > float64(maxCells) {
|
||||||
|
scale = min(scale, math.Sqrt(float64(maxCells)/area))
|
||||||
|
}
|
||||||
|
cols = max(1, int(math.Floor(float64(cols)*scale)))
|
||||||
|
rows = max(1, int(math.Floor(float64(rows)*scale)))
|
||||||
|
if quantum > 1 {
|
||||||
|
if cols >= quantum {
|
||||||
|
cols -= cols % quantum
|
||||||
|
}
|
||||||
|
if rows >= quantum {
|
||||||
|
rows -= rows % quantum
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return cols, rows
|
||||||
|
}
|
||||||
|
|
||||||
|
func New(deps ServerDeps) *Server {
|
||||||
|
cfg := deps.Config
|
||||||
|
|
||||||
|
cache := render.NewCache(int64(cfg.RenderCacheMB)<<20, cfg.RenderCacheCompress)
|
||||||
|
sets := make([]frameSet, len(deps.VideoSets))
|
||||||
|
for i, data := range deps.VideoSets {
|
||||||
|
sets[i] = frameSet{
|
||||||
|
data: data,
|
||||||
|
renderer: render.NewRenderer(i, data.ColorFrames, render.Options{
|
||||||
|
BrightnessThreshold: cfg.BrightnessThreshold,
|
||||||
|
Charset: cfg.Charset,
|
||||||
|
Invert: cfg.Invert,
|
||||||
|
}, cache),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
sshConfig := &ssh.ServerConfig{
|
||||||
|
MaxAuthTries: cfg.MaxAuthAttempts,
|
||||||
|
PasswordCallback: func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
|
||||||
|
ip := hostOnly(conn.RemoteAddr().String())
|
||||||
|
if cfg.LogCredentials {
|
||||||
|
logx.Info(fmt.Sprintf(
|
||||||
|
`Auth attempt from %s method=password user="%s" pass="%s"`,
|
||||||
|
ip, logx.SanitizeN(conn.User(), 128), logx.SanitizeN(string(password), 128),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
if conn.User() == "" || len(password) == 0 {
|
||||||
|
return nil, errors.New("password rejected")
|
||||||
|
}
|
||||||
|
return nil, nil
|
||||||
|
},
|
||||||
|
}
|
||||||
|
sshConfig.KeyExchanges = []string{
|
||||||
|
"mlkem768x25519-sha256",
|
||||||
|
"curve25519-sha256",
|
||||||
|
"curve25519-sha256@libssh.org",
|
||||||
|
"ecdh-sha2-nistp256",
|
||||||
|
"ecdh-sha2-nistp384",
|
||||||
|
"ecdh-sha2-nistp521",
|
||||||
|
"diffie-hellman-group14-sha256",
|
||||||
|
}
|
||||||
|
if deps.BannerText != nil {
|
||||||
|
banner := *deps.BannerText
|
||||||
|
sshConfig.BannerCallback = func(_ ssh.ConnMetadata) string { return banner }
|
||||||
|
}
|
||||||
|
for _, key := range deps.HostKeys {
|
||||||
|
sshConfig.AddHostKey(key)
|
||||||
|
}
|
||||||
|
|
||||||
|
return &Server{
|
||||||
|
config: cfg,
|
||||||
|
sshConfig: sshConfig,
|
||||||
|
sets: sets,
|
||||||
|
cache: cache,
|
||||||
|
tracker: newConnectionTracker(),
|
||||||
|
sessions: newSessionTracker(),
|
||||||
|
fakeLogin: deps.FakeLoginText,
|
||||||
|
goodbye: deps.GoodbyeText,
|
||||||
|
conns: make(map[net.Conn]struct{}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func hostOnly(addr string) string {
|
||||||
|
host, _, err := net.SplitHostPort(addr)
|
||||||
|
if err != nil {
|
||||||
|
return addr
|
||||||
|
}
|
||||||
|
return host
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) Listen(host string, port int) error {
|
||||||
|
listener, err := net.Listen("tcp", net.JoinHostPort(host, fmt.Sprint(port)))
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
s.mu.Lock()
|
||||||
|
if s.closing {
|
||||||
|
s.mu.Unlock()
|
||||||
|
_ = listener.Close()
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
s.listener = listener
|
||||||
|
s.mu.Unlock()
|
||||||
|
logx.Info(fmt.Sprintf("TrollSSH listening on %s:%d", host, port))
|
||||||
|
for {
|
||||||
|
conn, err := listener.Accept()
|
||||||
|
if err != nil {
|
||||||
|
if errors.Is(err, net.ErrClosed) {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
ip := hostOnly(conn.RemoteAddr().String())
|
||||||
|
activeForIP, total, ok := s.tracker.tryAcquire(ip, s.config.MaxConnections, s.config.MaxTotalConnections)
|
||||||
|
if !ok {
|
||||||
|
_ = conn.Close()
|
||||||
|
logx.Warn("Connection rejected (limit reached) from", ip)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
s.mu.Lock()
|
||||||
|
if s.closing {
|
||||||
|
s.mu.Unlock()
|
||||||
|
s.tracker.release(ip)
|
||||||
|
_ = conn.Close()
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
s.conns[conn] = struct{}{}
|
||||||
|
s.connWG.Add(1)
|
||||||
|
s.mu.Unlock()
|
||||||
|
go s.handleConn(conn, ip, activeForIP, total)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) Close() {
|
||||||
|
s.closeOnce.Do(func() {
|
||||||
|
s.mu.Lock()
|
||||||
|
s.closing = true
|
||||||
|
listener := s.listener
|
||||||
|
conns := make([]net.Conn, 0, len(s.conns))
|
||||||
|
for conn := range s.conns {
|
||||||
|
conns = append(conns, conn)
|
||||||
|
}
|
||||||
|
s.mu.Unlock()
|
||||||
|
if listener != nil {
|
||||||
|
_ = listener.Close()
|
||||||
|
}
|
||||||
|
for _, conn := range conns {
|
||||||
|
_ = conn.Close()
|
||||||
|
}
|
||||||
|
s.connWG.Wait()
|
||||||
|
stats := s.cache.Stats()
|
||||||
|
if stats.Hits+stats.Misses > 0 {
|
||||||
|
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,
|
||||||
|
))
|
||||||
|
}
|
||||||
|
for _, set := range s.sets {
|
||||||
|
if err := set.data.Close(); err != nil {
|
||||||
|
logx.Warn("Failed to release frame set", set.data.Name, logx.Sanitize(err.Error()))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) handleConn(conn net.Conn, ip string, activeForIP, total int) {
|
||||||
|
defer func() {
|
||||||
|
s.tracker.release(ip)
|
||||||
|
s.mu.Lock()
|
||||||
|
delete(s.conns, conn)
|
||||||
|
s.mu.Unlock()
|
||||||
|
s.connWG.Done()
|
||||||
|
}()
|
||||||
|
|
||||||
|
if s.config.HandshakeTimeout > 0 {
|
||||||
|
_ = conn.SetDeadline(time.Now().Add(s.config.HandshakeTimeout))
|
||||||
|
}
|
||||||
|
|
||||||
|
sshConn, chans, reqs, err := ssh.NewServerConn(conn, s.sshConfig)
|
||||||
|
if err != nil {
|
||||||
|
if strings.Contains(err.Error(), "i/o timeout") {
|
||||||
|
logx.Warn("Handshake timeout for", ip)
|
||||||
|
} else {
|
||||||
|
logx.Warn(fmt.Sprintf("Client error from %s:", ip), logx.Sanitize(err.Error()))
|
||||||
|
}
|
||||||
|
_ = conn.Close()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
_ = conn.SetDeadline(time.Time{})
|
||||||
|
logx.Debug("Handshake from", ip)
|
||||||
|
defer func() { _ = sshConn.Close() }()
|
||||||
|
|
||||||
|
setIndex := rand.Intn(len(s.sets))
|
||||||
|
logx.Info(fmt.Sprintf(
|
||||||
|
"New connection from %s (ip=%d, total=%d) -> playing %q",
|
||||||
|
ip, activeForIP, total, s.sets[setIndex].data.Name,
|
||||||
|
))
|
||||||
|
|
||||||
|
go ssh.DiscardRequests(reqs)
|
||||||
|
|
||||||
|
var sessionWG sync.WaitGroup
|
||||||
|
for newChannel := range chans {
|
||||||
|
if newChannel.ChannelType() != "session" {
|
||||||
|
_ = newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if !s.sessions.tryAcquire(sshConn, maxSessionsPerConn, s.config.MaxTotalConnections) {
|
||||||
|
_ = newChannel.Reject(ssh.ResourceShortage, "session limit reached")
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
channel, requests, err := newChannel.Accept()
|
||||||
|
if err != nil {
|
||||||
|
s.sessions.release(sshConn)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
sessionWG.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer sessionWG.Done()
|
||||||
|
defer s.sessions.release(sshConn)
|
||||||
|
var timer *time.Timer
|
||||||
|
if s.config.SessionTimeout > 0 {
|
||||||
|
timer = time.AfterFunc(s.config.SessionTimeout, func() { _ = sshConn.Close() })
|
||||||
|
defer timer.Stop()
|
||||||
|
}
|
||||||
|
s.handleSession(sshConn, channel, requests, ip, setIndex)
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
_ = sshConn.Close()
|
||||||
|
sessionWG.Wait()
|
||||||
|
logx.Info("Client closed connection from", ip)
|
||||||
|
}
|
||||||
|
|
||||||
|
type termSize struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
width int
|
||||||
|
height int
|
||||||
|
updated time.Time
|
||||||
|
timer *time.Timer
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *termSize) set(w, h, maxDimension, maxCells int, force bool) {
|
||||||
|
width, height := clampTermSize(w, h, maxDimension, maxCells, terminalSizeQuantum)
|
||||||
|
|
||||||
|
t.mu.Lock()
|
||||||
|
defer t.mu.Unlock()
|
||||||
|
|
||||||
|
if !force {
|
||||||
|
if remaining := resizeDebounce - time.Since(t.updated); remaining > 0 {
|
||||||
|
if t.timer != nil {
|
||||||
|
t.timer.Stop()
|
||||||
|
}
|
||||||
|
t.timer = time.AfterFunc(remaining, func() {
|
||||||
|
t.mu.Lock()
|
||||||
|
defer t.mu.Unlock()
|
||||||
|
t.width, t.height = width, height
|
||||||
|
t.updated = time.Now()
|
||||||
|
})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if t.timer != nil {
|
||||||
|
t.timer.Stop()
|
||||||
|
t.timer = nil
|
||||||
|
}
|
||||||
|
t.width, t.height = width, height
|
||||||
|
t.updated = time.Now()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *termSize) get() (int, int) {
|
||||||
|
t.mu.Lock()
|
||||||
|
defer t.mu.Unlock()
|
||||||
|
return t.width, t.height
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseDims(payload []byte) (cols, rows int, ok bool) {
|
||||||
|
if len(payload) < 8 {
|
||||||
|
return 0, 0, false
|
||||||
|
}
|
||||||
|
// pty-req prefixes cols/rows with a TERM string; window-change does not.
|
||||||
|
offset := 0
|
||||||
|
strLen := binary.BigEndian.Uint32(payload)
|
||||||
|
if int(strLen)+12 <= len(payload) {
|
||||||
|
offset = 4 + int(strLen)
|
||||||
|
}
|
||||||
|
if len(payload) < offset+8 {
|
||||||
|
return 0, 0, false
|
||||||
|
}
|
||||||
|
cols = int(binary.BigEndian.Uint32(payload[offset:]))
|
||||||
|
rows = int(binary.BigEndian.Uint32(payload[offset+4:]))
|
||||||
|
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,
|
||||||
|
requests <-chan *ssh.Request,
|
||||||
|
ip string,
|
||||||
|
initialSetIndex int,
|
||||||
|
) {
|
||||||
|
defer func() { _ = channel.Close() }()
|
||||||
|
size := &termSize{}
|
||||||
|
size.set(80, 24, s.config.MaxDimension, s.config.MaxTerminalCells, true)
|
||||||
|
tier := render.ColorTierTrueColor
|
||||||
|
if s.config.ForceGrayscale {
|
||||||
|
tier = render.ColorTierNone
|
||||||
|
}
|
||||||
|
|
||||||
|
started := false
|
||||||
|
var playDone chan struct{}
|
||||||
|
for req := range requests {
|
||||||
|
switch req.Type {
|
||||||
|
case "pty-req":
|
||||||
|
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 = render.DetectColorTier(term)
|
||||||
|
if s.config.ForceGrayscale {
|
||||||
|
tier = render.ColorTierNone
|
||||||
|
}
|
||||||
|
logx.Debug(fmt.Sprintf("Client %s TERM=%q -> color tier %d", ip, logx.SanitizeN(term, 64), tier))
|
||||||
|
}
|
||||||
|
_ = req.Reply(true, nil)
|
||||||
|
case "window-change":
|
||||||
|
if len(req.Payload) >= 8 {
|
||||||
|
cols := int(binary.BigEndian.Uint32(req.Payload))
|
||||||
|
rows := int(binary.BigEndian.Uint32(req.Payload[4:]))
|
||||||
|
size.set(cols, rows, s.config.MaxDimension, s.config.MaxTerminalCells, false)
|
||||||
|
}
|
||||||
|
if req.WantReply {
|
||||||
|
_ = req.Reply(true, nil)
|
||||||
|
}
|
||||||
|
case "exec":
|
||||||
|
command := ""
|
||||||
|
if len(req.Payload) >= 4 {
|
||||||
|
n := binary.BigEndian.Uint32(req.Payload)
|
||||||
|
if int(n)+4 <= len(req.Payload) {
|
||||||
|
command = string(req.Payload[4 : 4+n])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
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 render.ColorTier) {
|
||||||
|
defer close(playDone)
|
||||||
|
s.playVideo(sshConn, channel, size, ip, initialSetIndex, false, tier)
|
||||||
|
}(playTier)
|
||||||
|
}
|
||||||
|
case "shell":
|
||||||
|
logx.Debug("Opening shell for session", ip)
|
||||||
|
_ = req.Reply(true, nil)
|
||||||
|
if !started {
|
||||||
|
started = true
|
||||||
|
playDone = make(chan struct{})
|
||||||
|
playTier := tier
|
||||||
|
go func(tier render.ColorTier) {
|
||||||
|
defer close(playDone)
|
||||||
|
s.playVideo(sshConn, channel, size, ip, initialSetIndex, false, tier)
|
||||||
|
}(playTier)
|
||||||
|
}
|
||||||
|
default:
|
||||||
|
if req.WantReply {
|
||||||
|
_ = req.Reply(false, nil)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ = channel.Close()
|
||||||
|
if playDone != nil {
|
||||||
|
<-playDone
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) pickNextSetIndex(exclude int) int {
|
||||||
|
if len(s.sets) <= 1 {
|
||||||
|
return exclude
|
||||||
|
}
|
||||||
|
next := exclude
|
||||||
|
for next == exclude {
|
||||||
|
next = rand.Intn(len(s.sets))
|
||||||
|
}
|
||||||
|
return next
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) playVideo(
|
||||||
|
sshConn *ssh.ServerConn,
|
||||||
|
channel ssh.Channel,
|
||||||
|
size *termSize,
|
||||||
|
ip string,
|
||||||
|
setIndex int,
|
||||||
|
keepAspectRatio bool,
|
||||||
|
tier render.ColorTier,
|
||||||
|
) {
|
||||||
|
cfg := s.config
|
||||||
|
current := s.sets[setIndex]
|
||||||
|
|
||||||
|
w, h := size.get()
|
||||||
|
logx.Debug(fmt.Sprintf("Terminal size %dx%d for %s", w, h, ip))
|
||||||
|
|
||||||
|
defer func() {
|
||||||
|
_ = writePartsWithTimeout(sshConn, channel, outputStallTimeout, showCursor)
|
||||||
|
}()
|
||||||
|
|
||||||
|
if s.fakeLogin != nil {
|
||||||
|
if err := writePartsWithTimeout(
|
||||||
|
sshConn, channel, outputStallTimeout, clearScreen, *s.fakeLogin,
|
||||||
|
); err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
done := make(chan struct{})
|
||||||
|
var doneOnce sync.Once
|
||||||
|
closeSession := func() {
|
||||||
|
doneOnce.Do(func() { close(done) })
|
||||||
|
}
|
||||||
|
|
||||||
|
switchCh := make(chan int, 8)
|
||||||
|
go func() {
|
||||||
|
buf := make([]byte, 256)
|
||||||
|
var lastSwitch time.Time
|
||||||
|
for {
|
||||||
|
n, err := channel.Read(buf)
|
||||||
|
if err != nil {
|
||||||
|
closeSession()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if !cfg.AllowUserControl {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
str := string(buf[:n])
|
||||||
|
delta := 0
|
||||||
|
if strings.Contains(str, "\x1b[C") || strings.Contains(str, "\x1b[A") {
|
||||||
|
delta = 1
|
||||||
|
} else if strings.Contains(str, "\x1b[D") || strings.Contains(str, "\x1b[B") {
|
||||||
|
delta = -1
|
||||||
|
}
|
||||||
|
if delta == 0 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
if now.Sub(lastSwitch) < cfg.SwitchDebounce {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
lastSwitch = now
|
||||||
|
select {
|
||||||
|
case switchCh <- delta:
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
|
loginTimer := time.NewTimer(cfg.LoginDelay)
|
||||||
|
select {
|
||||||
|
case <-loginTimer.C:
|
||||||
|
case <-done:
|
||||||
|
if !loginTimer.Stop() {
|
||||||
|
<-loginTimer.C
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := writePartsWithTimeout(sshConn, channel, outputStallTimeout, hideCursor); err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
frameInterval := func() time.Duration {
|
||||||
|
return time.Duration(float64(time.Second) / current.data.FPS)
|
||||||
|
}
|
||||||
|
|
||||||
|
ticker := time.NewTicker(frameInterval())
|
||||||
|
defer ticker.Stop()
|
||||||
|
|
||||||
|
currentFrame := 0
|
||||||
|
loopCount := 0
|
||||||
|
lastW, lastH := 0, 0
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-done:
|
||||||
|
return
|
||||||
|
|
||||||
|
case delta := <-switchCh:
|
||||||
|
if len(s.sets) <= 1 {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
setIndex = (setIndex + delta + len(s.sets)) % len(s.sets)
|
||||||
|
current = s.sets[setIndex]
|
||||||
|
currentFrame = 0
|
||||||
|
lastW, lastH = 0, 0
|
||||||
|
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)
|
||||||
|
if err != nil {
|
||||||
|
logx.Error("Render error for", ip, logx.Sanitize(err.Error()))
|
||||||
|
_ = sshConn.Close()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
prefix := homeCursor
|
||||||
|
if w != lastW || h != lastH {
|
||||||
|
prefix = clearScreen
|
||||||
|
lastW, lastH = w, h
|
||||||
|
}
|
||||||
|
if err := writeFrameWithTimeout(
|
||||||
|
sshConn, channel, outputStallTimeout, prefix, ascii,
|
||||||
|
); err != nil {
|
||||||
|
closeSession()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
currentFrame++
|
||||||
|
if currentFrame < len(current.data.ColorFrames) {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
currentFrame = 0
|
||||||
|
loopCount++
|
||||||
|
if cfg.MaxLoop > 0 && loopCount >= cfg.MaxLoop {
|
||||||
|
if err := writePartsWithTimeout(
|
||||||
|
sshConn, channel, outputStallTimeout, showCursor, clearScreen,
|
||||||
|
); err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if s.goodbye != nil {
|
||||||
|
if err := writePartsWithTimeout(
|
||||||
|
sshConn, channel, outputStallTimeout, *s.goodbye,
|
||||||
|
); err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
closeTimer := time.NewTimer(time.Second)
|
||||||
|
select {
|
||||||
|
case <-closeTimer.C:
|
||||||
|
case <-done:
|
||||||
|
if !closeTimer.Stop() {
|
||||||
|
<-closeTimer.C
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
logx.Info("Playback finished, closing session", ip)
|
||||||
|
_ = channel.Close()
|
||||||
|
_ = sshConn.Close()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if cfg.PlaybackMode == config.PlaybackRandom {
|
||||||
|
setIndex = s.pickNextSetIndex(setIndex)
|
||||||
|
current = s.sets[setIndex]
|
||||||
|
logx.Info(fmt.Sprintf(
|
||||||
|
"Playthrough done for %s, switching to %q", ip, current.data.Name,
|
||||||
|
))
|
||||||
|
ticker.Reset(frameInterval())
|
||||||
|
} else if cfg.MaxLoop > 0 {
|
||||||
|
logx.Info(fmt.Sprintf(
|
||||||
|
"Playthrough done for %s, looping %q (%d/%d)",
|
||||||
|
ip, current.data.Name, loopCount, cfg.MaxLoop,
|
||||||
|
))
|
||||||
|
} else {
|
||||||
|
logx.Info(fmt.Sprintf(
|
||||||
|
"Playthrough done for %s, looping %q (%d)",
|
||||||
|
ip, current.data.Name, loopCount,
|
||||||
|
))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,158 @@
|
|||||||
|
package sshserver
|
||||||
|
|
||||||
|
import (
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"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{})
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
var mu sync.Mutex
|
||||||
|
accepted := make(map[string]int)
|
||||||
|
for i := range 100 {
|
||||||
|
wg.Add(1)
|
||||||
|
go func(i int) {
|
||||||
|
defer wg.Done()
|
||||||
|
<-start
|
||||||
|
ip := string(rune('a' + i%10))
|
||||||
|
if _, _, ok := tracker.tryAcquire(ip, 3, 7); ok {
|
||||||
|
mu.Lock()
|
||||||
|
accepted[ip]++
|
||||||
|
mu.Unlock()
|
||||||
|
}
|
||||||
|
}(i)
|
||||||
|
}
|
||||||
|
close(start)
|
||||||
|
wg.Wait()
|
||||||
|
|
||||||
|
total := 0
|
||||||
|
for ip, count := range accepted {
|
||||||
|
total += count
|
||||||
|
if count > 3 {
|
||||||
|
t.Fatalf("IP %q acquired %d slots", ip, count)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if total != 7 || tracker.totalCount() != 7 {
|
||||||
|
t.Fatalf("accepted=%d tracked=%d, want 7", total, tracker.totalCount())
|
||||||
|
}
|
||||||
|
for ip, count := range accepted {
|
||||||
|
for range count {
|
||||||
|
tracker.release(ip)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSessionTrackerLimits(t *testing.T) {
|
||||||
|
tracker := newSessionTracker()
|
||||||
|
first := &ssh.ServerConn{}
|
||||||
|
second := &ssh.ServerConn{}
|
||||||
|
if !tracker.tryAcquire(first, 1, 2) {
|
||||||
|
t.Fatal("first session rejected")
|
||||||
|
}
|
||||||
|
if tracker.tryAcquire(first, 1, 2) {
|
||||||
|
t.Fatal("per-connection limit was not enforced")
|
||||||
|
}
|
||||||
|
if !tracker.tryAcquire(second, 1, 2) {
|
||||||
|
t.Fatal("second connection session rejected")
|
||||||
|
}
|
||||||
|
if tracker.tryAcquire(&ssh.ServerConn{}, 1, 2) {
|
||||||
|
t.Fatal("global session limit was not enforced")
|
||||||
|
}
|
||||||
|
tracker.release(first)
|
||||||
|
if !tracker.tryAcquire(&ssh.ServerConn{}, 1, 2) {
|
||||||
|
t.Fatal("released slot was not reusable")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTermSizeDebouncesResize(t *testing.T) {
|
||||||
|
size := &termSize{}
|
||||||
|
size.set(80, 24, 512, 500*512, true)
|
||||||
|
size.set(200, 100, 512, 500*512, false)
|
||||||
|
if w, h := size.get(); w != 80 || h != 24 {
|
||||||
|
t.Fatalf("debounced size = %dx%d", w, h)
|
||||||
|
}
|
||||||
|
size.set(200, 100, 512, 500*512, true)
|
||||||
|
if w, h := size.get(); w != 200 || h != 100 {
|
||||||
|
t.Fatalf("forced size = %dx%d", w, h)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTermSizeAppliesFinalResizeAfterDebounce(t *testing.T) {
|
||||||
|
size := &termSize{}
|
||||||
|
size.set(80, 24, 512, 500*512, true)
|
||||||
|
|
||||||
|
// Rapid burst of resize events, as happens during an interactive drag-resize.
|
||||||
|
size.set(100, 40, 512, 500*512, false)
|
||||||
|
size.set(150, 60, 512, 500*512, false)
|
||||||
|
size.set(200, 100, 512, 500*512, false)
|
||||||
|
|
||||||
|
if w, h := size.get(); w != 80 || h != 24 {
|
||||||
|
t.Fatalf("size changed before debounce elapsed: %dx%d", w, h)
|
||||||
|
}
|
||||||
|
|
||||||
|
time.Sleep(resizeDebounce + 50*time.Millisecond)
|
||||||
|
|
||||||
|
if w, h := size.get(); w != 200 || h != 100 {
|
||||||
|
t.Fatalf("final resize was not applied after debounce: got %dx%d, want 200x100", w, h)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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()
|
||||||
|
}
|
||||||
@@ -0,0 +1,122 @@
|
|||||||
|
package tsf
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"encoding/binary"
|
||||||
|
"fmt"
|
||||||
|
"math"
|
||||||
|
"os"
|
||||||
|
)
|
||||||
|
|
||||||
|
func Write(output string, data *FramesContainer) error {
|
||||||
|
if data == nil {
|
||||||
|
return fmt.Errorf("cannot write nil frames container")
|
||||||
|
}
|
||||||
|
if math.IsNaN(data.FPS) || math.IsInf(data.FPS, 0) || data.FPS <= 0 || data.FPS > maxTSFFPS {
|
||||||
|
return fmt.Errorf("cannot write .tsf: fps must be finite, positive, and at most %d", maxTSFFPS)
|
||||||
|
}
|
||||||
|
if len(data.ColorFrames) > maxTSFFrameCount || uint64(len(data.ColorFrames)) > math.MaxUint32 {
|
||||||
|
return fmt.Errorf("cannot write .tsf: frame count exceeds limit")
|
||||||
|
}
|
||||||
|
for i, frame := range data.ColorFrames {
|
||||||
|
if uint64(len(frame)) > math.MaxUint32 {
|
||||||
|
return fmt.Errorf("cannot write .tsf: frame %d length exceeds uint32", i)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
f, err := os.Create(output)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer func() { _ = f.Close() }()
|
||||||
|
|
||||||
|
w := bufio.NewWriterSize(f, 1<<20)
|
||||||
|
if _, err := w.WriteString(tsfMagic); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
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.ColorFrames)))
|
||||||
|
if _, err := w.Write(hdr[:]); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
var lenBuf [4]byte
|
||||||
|
for _, frame := range data.ColorFrames {
|
||||||
|
binary.LittleEndian.PutUint32(lenBuf[:], uint32(len(frame)))
|
||||||
|
if _, err := w.Write(lenBuf[:]); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if _, err := w.Write(frame); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return w.Flush()
|
||||||
|
}
|
||||||
|
|
||||||
|
func Load(filename string) (*FramesContainer, error) {
|
||||||
|
file, err := readFrameFile(filename)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
owned := false
|
||||||
|
defer func() {
|
||||||
|
if !owned {
|
||||||
|
_ = file.Close()
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
raw := file.data
|
||||||
|
invalid := func() error {
|
||||||
|
return fmt.Errorf("invalid frames file %q: corrupt .tsf container", filename)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(raw) < 18 || string(raw[:4]) != tsfMagic {
|
||||||
|
return nil, invalid()
|
||||||
|
}
|
||||||
|
version := binary.LittleEndian.Uint16(raw[4:])
|
||||||
|
if version != tsfVersion {
|
||||||
|
return nil, fmt.Errorf("unsupported .tsf version %d in %q", version, filename)
|
||||||
|
}
|
||||||
|
fps := math.Float64frombits(binary.LittleEndian.Uint64(raw[6:]))
|
||||||
|
count := binary.LittleEndian.Uint32(raw[14:])
|
||||||
|
if math.IsNaN(fps) || math.IsInf(fps, 0) || fps <= 0 || fps > maxTSFFPS {
|
||||||
|
return nil, fmt.Errorf(
|
||||||
|
"invalid frames file %q: fps must be finite, greater than 0, and at most %d",
|
||||||
|
filename, maxTSFFPS,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
if count == 0 {
|
||||||
|
return nil, fmt.Errorf("invalid frames file %q: expected non-empty frames", filename)
|
||||||
|
}
|
||||||
|
if count > maxTSFFrameCount || uint64(count) > uint64((len(raw)-18)/4) {
|
||||||
|
return nil, invalid()
|
||||||
|
}
|
||||||
|
|
||||||
|
colorFrames := make([][]byte, 0, int(count))
|
||||||
|
off := 18
|
||||||
|
for range count {
|
||||||
|
if len(raw)-off < 4 {
|
||||||
|
return nil, invalid()
|
||||||
|
}
|
||||||
|
n := uint64(binary.LittleEndian.Uint32(raw[off:]))
|
||||||
|
off += 4
|
||||||
|
if n > uint64(len(raw)-off) {
|
||||||
|
return nil, invalid()
|
||||||
|
}
|
||||||
|
nativeLen := int(n)
|
||||||
|
colorFrames = append(colorFrames, raw[off:off+nativeLen])
|
||||||
|
off += nativeLen
|
||||||
|
}
|
||||||
|
|
||||||
|
if off != len(raw) {
|
||||||
|
return nil, invalid()
|
||||||
|
}
|
||||||
|
|
||||||
|
file.dropResident()
|
||||||
|
|
||||||
|
data := &FramesContainer{ColorFrames: colorFrames, FPS: fps}
|
||||||
|
frameFileOwners.Store(data, file)
|
||||||
|
owned = true
|
||||||
|
return data, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,12 @@
|
|||||||
|
//go:build !unix
|
||||||
|
|
||||||
|
package tsf
|
||||||
|
|
||||||
|
import "os"
|
||||||
|
|
||||||
|
func readFrameFile(filename string) (*frameFile, error) {
|
||||||
|
data, err := os.ReadFile(filename)
|
||||||
|
return &frameFile{data: data}, err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *frameFile) dropResident() {}
|
||||||
@@ -0,0 +1,46 @@
|
|||||||
|
//go:build unix
|
||||||
|
|
||||||
|
package tsf
|
||||||
|
|
||||||
|
import (
|
||||||
|
"os"
|
||||||
|
"syscall"
|
||||||
|
)
|
||||||
|
|
||||||
|
func readFrameFile(filename string) (*frameFile, error) {
|
||||||
|
f, err := os.Open(filename)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer func() { _ = f.Close() }()
|
||||||
|
|
||||||
|
info, err := f.Stat()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
size := info.Size()
|
||||||
|
if size <= 0 || size != int64(int(size)) {
|
||||||
|
data, err := os.ReadFile(filename)
|
||||||
|
return &frameFile{data: data}, err
|
||||||
|
}
|
||||||
|
data, err := syscall.Mmap(int(f.Fd()), 0, int(size), syscall.PROT_READ, syscall.MAP_SHARED)
|
||||||
|
if err != nil {
|
||||||
|
data, err := os.ReadFile(filename)
|
||||||
|
return &frameFile{data: data}, err
|
||||||
|
}
|
||||||
|
|
||||||
|
_ = syscall.Madvise(data, syscall.MADV_RANDOM)
|
||||||
|
return &frameFile{
|
||||||
|
data: data,
|
||||||
|
cleanup: func() error {
|
||||||
|
return syscall.Munmap(data)
|
||||||
|
},
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (f *frameFile) dropResident() {
|
||||||
|
if f == nil || f.cleanup == nil || len(f.data) == 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
_ = syscall.Madvise(f.data, syscall.MADV_DONTNEED)
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,371 @@
|
|||||||
|
package tsf
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bufio"
|
||||||
|
"bytes"
|
||||||
|
"encoding/binary"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"math"
|
||||||
|
"os"
|
||||||
|
"os/exec"
|
||||||
|
"path/filepath"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/YuzuZensai/TrollSSH/internal/logx"
|
||||||
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
jpegSOI = []byte{0xff, 0xd8}
|
||||||
|
jpegEOI = []byte{0xff, 0xd9}
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
maxJPEGFrameBytes = 64 << 20
|
||||||
|
maxFFmpegLogBytes = 64 << 10
|
||||||
|
)
|
||||||
|
|
||||||
|
type jpegFrameSplitter struct {
|
||||||
|
buffer []byte
|
||||||
|
scan int
|
||||||
|
inJPEG bool
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *jpegFrameSplitter) push(chunk []byte, emit func([]byte) error) (int, error) {
|
||||||
|
s.buffer = append(s.buffer, chunk...)
|
||||||
|
emitted := 0
|
||||||
|
for {
|
||||||
|
if !s.inJPEG {
|
||||||
|
start := bytes.Index(s.buffer[s.scan:], jpegSOI)
|
||||||
|
if start == -1 {
|
||||||
|
// Retain only a possible marker prefix spanning two reads.
|
||||||
|
if len(s.buffer) > 0 && s.buffer[len(s.buffer)-1] == jpegSOI[0] {
|
||||||
|
s.buffer = s.buffer[len(s.buffer)-1:]
|
||||||
|
} else {
|
||||||
|
s.buffer = s.buffer[:0]
|
||||||
|
}
|
||||||
|
s.scan = 0
|
||||||
|
return emitted, nil
|
||||||
|
}
|
||||||
|
start += s.scan
|
||||||
|
s.buffer = s.buffer[start:]
|
||||||
|
s.scan = len(jpegSOI)
|
||||||
|
s.inJPEG = true
|
||||||
|
}
|
||||||
|
|
||||||
|
end := bytes.Index(s.buffer[s.scan:], jpegEOI)
|
||||||
|
if end == -1 {
|
||||||
|
if len(s.buffer) > maxJPEGFrameBytes {
|
||||||
|
return emitted, fmt.Errorf("JPEG frame exceeds %d MiB limit", maxJPEGFrameBytes>>20)
|
||||||
|
}
|
||||||
|
s.scan = max(len(jpegSOI), len(s.buffer)-1)
|
||||||
|
return emitted, nil
|
||||||
|
}
|
||||||
|
frameEnd := s.scan + end + len(jpegEOI)
|
||||||
|
if frameEnd > maxJPEGFrameBytes {
|
||||||
|
return emitted, fmt.Errorf("JPEG frame exceeds %d MiB limit", maxJPEGFrameBytes>>20)
|
||||||
|
}
|
||||||
|
if err := emit(s.buffer[:frameEnd]); err != nil {
|
||||||
|
return emitted, err
|
||||||
|
}
|
||||||
|
emitted++
|
||||||
|
s.buffer = s.buffer[frameEnd:]
|
||||||
|
s.scan = 0
|
||||||
|
s.inJPEG = false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *jpegFrameSplitter) finish() error {
|
||||||
|
if s.inJPEG {
|
||||||
|
return fmt.Errorf("ffmpeg produced a truncated JPEG frame")
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
type boundedLog struct {
|
||||||
|
buffer bytes.Buffer
|
||||||
|
limit int
|
||||||
|
}
|
||||||
|
|
||||||
|
func (w *boundedLog) Write(p []byte) (int, error) {
|
||||||
|
n := len(p)
|
||||||
|
if remaining := w.limit - w.buffer.Len(); remaining > 0 {
|
||||||
|
_, _ = w.buffer.Write(p[:min(len(p), remaining)])
|
||||||
|
}
|
||||||
|
return n, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (w *boundedLog) String() string {
|
||||||
|
return strings.TrimSpace(w.buffer.String())
|
||||||
|
}
|
||||||
|
|
||||||
|
type streamingTSF struct {
|
||||||
|
file *os.File
|
||||||
|
writer *bufio.Writer
|
||||||
|
path string
|
||||||
|
count uint32
|
||||||
|
}
|
||||||
|
|
||||||
|
func newStreamingTSF(output string, fps float64) (*streamingTSF, error) {
|
||||||
|
if math.IsNaN(fps) || math.IsInf(fps, 0) || fps <= 0 || fps > maxTSFFPS {
|
||||||
|
return nil, fmt.Errorf("cannot write .tsf: fps must be finite and between 0 and %d", maxTSFFPS)
|
||||||
|
}
|
||||||
|
dir := filepath.Dir(output)
|
||||||
|
f, err := os.CreateTemp(dir, "."+filepath.Base(output)+"-*.tmp")
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
s := &streamingTSF{file: f, writer: bufio.NewWriterSize(f, 1<<20), path: f.Name()}
|
||||||
|
if err := f.Chmod(0o644); err != nil {
|
||||||
|
s.abort()
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if _, err := s.writer.WriteString(tsfMagic); err != nil {
|
||||||
|
s.abort()
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
var hdr [14]byte
|
||||||
|
binary.LittleEndian.PutUint16(hdr[0:], tsfVersion)
|
||||||
|
binary.LittleEndian.PutUint64(hdr[2:], math.Float64bits(fps))
|
||||||
|
if _, err := s.writer.Write(hdr[:]); err != nil {
|
||||||
|
s.abort()
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return s, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *streamingTSF) addFrame(frame []byte) error {
|
||||||
|
if s.count >= maxTSFFrameCount {
|
||||||
|
return fmt.Errorf("too many video frames")
|
||||||
|
}
|
||||||
|
if uint64(len(frame)) > math.MaxUint32 {
|
||||||
|
return fmt.Errorf("JPEG frame is too large")
|
||||||
|
}
|
||||||
|
var size [4]byte
|
||||||
|
binary.LittleEndian.PutUint32(size[:], uint32(len(frame)))
|
||||||
|
if _, err := s.writer.Write(size[:]); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if _, err := s.writer.Write(frame); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
s.count++
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *streamingTSF) commit(output string) error {
|
||||||
|
if s.count == 0 {
|
||||||
|
return fmt.Errorf("no frames were decoded from the video")
|
||||||
|
}
|
||||||
|
if err := s.writer.Flush(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if _, err := s.file.Seek(14, io.SeekStart); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
var count [4]byte
|
||||||
|
binary.LittleEndian.PutUint32(count[:], s.count)
|
||||||
|
if _, err := s.file.Write(count[:]); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if err := s.file.Sync(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if err := s.file.Close(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
s.file = nil
|
||||||
|
if err := os.Rename(s.path, output); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
s.path = ""
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *streamingTSF) abort() {
|
||||||
|
if s.file != nil {
|
||||||
|
_ = s.file.Close()
|
||||||
|
s.file = nil
|
||||||
|
}
|
||||||
|
if s.path != "" {
|
||||||
|
_ = os.Remove(s.path)
|
||||||
|
s.path = ""
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type ffprobeOutput struct {
|
||||||
|
Streams []struct {
|
||||||
|
RFrameRate string `json:"r_frame_rate"`
|
||||||
|
NbFrames string `json:"nb_frames"`
|
||||||
|
Duration string `json:"duration"`
|
||||||
|
} `json:"streams"`
|
||||||
|
Format struct {
|
||||||
|
Duration string `json:"duration"`
|
||||||
|
} `json:"format"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func parseFrameRate(rate string) float64 {
|
||||||
|
if rate == "" {
|
||||||
|
return math.NaN()
|
||||||
|
}
|
||||||
|
parts := strings.SplitN(rate, "/", 2)
|
||||||
|
num, err := strconv.ParseFloat(parts[0], 64)
|
||||||
|
if err != nil {
|
||||||
|
return math.NaN()
|
||||||
|
}
|
||||||
|
if len(parts) == 2 {
|
||||||
|
den, err := strconv.ParseFloat(parts[1], 64)
|
||||||
|
if err != nil || den == 0 {
|
||||||
|
return math.NaN()
|
||||||
|
}
|
||||||
|
return num / den
|
||||||
|
}
|
||||||
|
return num
|
||||||
|
}
|
||||||
|
|
||||||
|
func extractFrames(path, vf, label string, totalFrames int, emit func([]byte) error) (int, error) {
|
||||||
|
cmd := exec.Command(
|
||||||
|
"ffmpeg", "-hide_banner", "-loglevel", "error", "-nostats",
|
||||||
|
"-i", path,
|
||||||
|
"-c:v", "mjpeg",
|
||||||
|
"-q:v", "3",
|
||||||
|
"-vf", vf,
|
||||||
|
"-f", "image2pipe",
|
||||||
|
"pipe:1",
|
||||||
|
)
|
||||||
|
stderr := &boundedLog{limit: maxFFmpegLogBytes}
|
||||||
|
cmd.Stderr = stderr
|
||||||
|
stdout, err := cmd.StdoutPipe()
|
||||||
|
if err != nil {
|
||||||
|
return 0, fmt.Errorf("ffmpeg failed: %w", err)
|
||||||
|
}
|
||||||
|
if err := cmd.Start(); err != nil {
|
||||||
|
return 0, fmt.Errorf("ffmpeg failed: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
splitter := &jpegFrameSplitter{}
|
||||||
|
frameCount := 0
|
||||||
|
lastReport := time.Now()
|
||||||
|
reportProgress := func(force bool) {
|
||||||
|
if !force && time.Since(lastReport) < 250*time.Millisecond {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if totalFrames > 0 {
|
||||||
|
pct := min(100, int(math.Round(float64(frameCount)/float64(totalFrames)*100)))
|
||||||
|
fmt.Printf("\rGenerating %s frames: %d/%d (%d%%)", label, frameCount, totalFrames, pct)
|
||||||
|
} else {
|
||||||
|
fmt.Printf("\rGenerating %s frames: %d", label, frameCount)
|
||||||
|
}
|
||||||
|
lastReport = time.Now()
|
||||||
|
}
|
||||||
|
failStream := func(err error) (int, error) {
|
||||||
|
_ = cmd.Process.Kill()
|
||||||
|
_ = cmd.Wait()
|
||||||
|
return frameCount, err
|
||||||
|
}
|
||||||
|
|
||||||
|
buf := make([]byte, 256*1024)
|
||||||
|
for {
|
||||||
|
n, readErr := stdout.Read(buf)
|
||||||
|
if n > 0 {
|
||||||
|
emitted, splitErr := splitter.push(buf[:n], emit)
|
||||||
|
frameCount += emitted
|
||||||
|
if splitErr != nil {
|
||||||
|
return failStream(fmt.Errorf("ffmpeg stream error: %w", splitErr))
|
||||||
|
}
|
||||||
|
if emitted > 0 {
|
||||||
|
reportProgress(false)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if readErr == io.EOF {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if readErr != nil {
|
||||||
|
return failStream(fmt.Errorf("ffmpeg stream error: %w", readErr))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err := cmd.Wait(); err != nil {
|
||||||
|
if msg := stderr.String(); msg != "" {
|
||||||
|
return frameCount, fmt.Errorf("ffmpeg failed: %s", msg)
|
||||||
|
}
|
||||||
|
return frameCount, fmt.Errorf("ffmpeg failed: %w", err)
|
||||||
|
}
|
||||||
|
if err := splitter.finish(); err != nil {
|
||||||
|
return frameCount, err
|
||||||
|
}
|
||||||
|
if frameCount == 0 {
|
||||||
|
return 0, fmt.Errorf("no frames were decoded from the video")
|
||||||
|
}
|
||||||
|
reportProgress(true)
|
||||||
|
fmt.Println()
|
||||||
|
return frameCount, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func ProcessVideo(path, output string, maxDimension int) error {
|
||||||
|
probeCmd := exec.Command(
|
||||||
|
"ffprobe", "-v", "error",
|
||||||
|
"-show_streams", "-show_format",
|
||||||
|
"-of", "json", path,
|
||||||
|
)
|
||||||
|
probeOut, err := probeCmd.Output()
|
||||||
|
if err != nil {
|
||||||
|
msg := err.Error()
|
||||||
|
var exitErr *exec.ExitError
|
||||||
|
if errors.As(err, &exitErr) && len(exitErr.Stderr) > 0 {
|
||||||
|
msg = strings.TrimSpace(string(exitErr.Stderr))
|
||||||
|
}
|
||||||
|
return fmt.Errorf("ffprobe failed: %s", msg)
|
||||||
|
}
|
||||||
|
|
||||||
|
var probe ffprobeOutput
|
||||||
|
if err := json.Unmarshal(probeOut, &probe); err != nil {
|
||||||
|
return fmt.Errorf("ffprobe failed: %w", err)
|
||||||
|
}
|
||||||
|
if len(probe.Streams) == 0 {
|
||||||
|
return fmt.Errorf("unable to determine a valid video fps")
|
||||||
|
}
|
||||||
|
stream := probe.Streams[0]
|
||||||
|
|
||||||
|
fps := parseFrameRate(stream.RFrameRate)
|
||||||
|
if math.IsNaN(fps) || fps <= 0 {
|
||||||
|
return fmt.Errorf("unable to determine a valid video fps")
|
||||||
|
}
|
||||||
|
|
||||||
|
totalFrames := 0
|
||||||
|
if n, err := strconv.Atoi(stream.NbFrames); err == nil {
|
||||||
|
totalFrames = n
|
||||||
|
} else {
|
||||||
|
durStr := probe.Format.Duration
|
||||||
|
if durStr == "" {
|
||||||
|
durStr = stream.Duration
|
||||||
|
}
|
||||||
|
if d, err := strconv.ParseFloat(durStr, 64); err == nil {
|
||||||
|
totalFrames = int(math.Round(d * fps))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
scaleFilter := fmt.Sprintf(
|
||||||
|
"scale=w=%d:h=%d:force_original_aspect_ratio=decrease",
|
||||||
|
maxDimension, maxDimension,
|
||||||
|
)
|
||||||
|
|
||||||
|
outputFile, err := newStreamingTSF(output, fps)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer outputFile.abort()
|
||||||
|
|
||||||
|
frameCount, err := extractFrames(path, scaleFilter, "color", totalFrames, outputFile.addFrame)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if err := outputFile.commit(output); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
logx.Info(fmt.Sprintf("Saved %d frames to %s", frameCount, output))
|
||||||
|
return nil
|
||||||
|
}
|
||||||
+3
-3
@@ -6,10 +6,10 @@ pre-commit:
|
|||||||
run: test -z "$(gofmt -l {staged_files})"
|
run: test -z "$(gofmt -l {staged_files})"
|
||||||
vet:
|
vet:
|
||||||
glob: "*.go"
|
glob: "*.go"
|
||||||
run: go vet ./src/...
|
run: go vet ./...
|
||||||
lint:
|
lint:
|
||||||
glob: "*.go"
|
glob: "*.go"
|
||||||
run: golangci-lint run ./src/...
|
run: golangci-lint run ./...
|
||||||
test:
|
test:
|
||||||
glob: "*.go"
|
glob: "*.go"
|
||||||
run: go test ./src/...
|
run: go test ./...
|
||||||
|
|||||||
-354
@@ -1,354 +0,0 @@
|
|||||||
package main
|
|
||||||
|
|
||||||
import (
|
|
||||||
"bytes"
|
|
||||||
"fmt"
|
|
||||||
"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.
|
|
||||||
if err := writeTSF(path, &FramesContainer{ColorFrames: [][]byte{{1}}, FPS: 0}); err != nil {
|
|
||||||
t.Fatalf("writeTSF: %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(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(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([]string, 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 got != results[0] {
|
|
||||||
t.Fatalf("result %d differs from result 0", i)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestRenderCacheEvictsByBytes(t *testing.T) {
|
|
||||||
budget := 3 * entryCost("k", strings.Repeat("x", 1000))
|
|
||||||
c := newRenderCache(budget)
|
|
||||||
for i := range 5 {
|
|
||||||
c.put(fmt.Sprintf("%d", i), strings.Repeat("x", 1000))
|
|
||||||
}
|
|
||||||
if c.size > budget {
|
|
||||||
t.Errorf("size %d exceeds budget %d", c.size, budget)
|
|
||||||
}
|
|
||||||
if _, ok := c.get("0"); ok {
|
|
||||||
t.Error("oldest entry should have been evicted")
|
|
||||||
}
|
|
||||||
if _, ok := c.get("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("k", "v")
|
|
||||||
if _, ok := c.get("k"); ok {
|
|
||||||
t.Error("nil cache should never hit")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestRenderCacheRejectsOversizedEntry(t *testing.T) {
|
|
||||||
c := newRenderCache(256)
|
|
||||||
c.put("big", strings.Repeat("x", 10_000))
|
|
||||||
if _, ok := c.get("big"); ok {
|
|
||||||
t.Error("entry larger than budget should not be cached")
|
|
||||||
}
|
|
||||||
if c.size != 0 {
|
|
||||||
t.Errorf("size = %d, want 0", c.size)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestConnectionTracker(t *testing.T) {
|
|
||||||
tr := newConnectionTracker()
|
|
||||||
tr.increment("1.2.3.4")
|
|
||||||
tr.increment("1.2.3.4")
|
|
||||||
if !tr.hasReachedLimits("1.2.3.4", 2, 100) {
|
|
||||||
t.Error("expected per-ip limit reached")
|
|
||||||
}
|
|
||||||
tr.decrement("1.2.3.4")
|
|
||||||
tr.decrement("1.2.3.4")
|
|
||||||
if tr.totalCount() != 0 {
|
|
||||||
t.Errorf("total = %d", tr.totalCount())
|
|
||||||
}
|
|
||||||
if tr.hasReachedLimits("1.2.3.4", 2, 100) {
|
|
||||||
t.Error("should be cleared")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestClampDimension(t *testing.T) {
|
|
||||||
if clampDimension(0, 100) != 1 {
|
|
||||||
t.Error("floor")
|
|
||||||
}
|
|
||||||
if clampDimension(500, 100) != 100 {
|
|
||||||
t.Error("ceil")
|
|
||||||
}
|
|
||||||
if clampDimension(50, 100) != 50 {
|
|
||||||
t.Error("passthrough")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
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)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,91 +0,0 @@
|
|||||||
package main
|
|
||||||
|
|
||||||
import (
|
|
||||||
"bufio"
|
|
||||||
"encoding/binary"
|
|
||||||
"fmt"
|
|
||||||
"math"
|
|
||||||
"os"
|
|
||||||
)
|
|
||||||
|
|
||||||
// .tsf container, little-endian: "TSFR" | version uint16 | fps float64 |
|
|
||||||
// count uint32 | count × (colorLen uint32, color JPEG).
|
|
||||||
const (
|
|
||||||
tsfMagic = "TSFR"
|
|
||||||
tsfVersion = 1
|
|
||||||
)
|
|
||||||
|
|
||||||
func writeTSF(output string, data *FramesContainer) error {
|
|
||||||
f, err := os.Create(output)
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
defer func() { _ = f.Close() }()
|
|
||||||
|
|
||||||
w := bufio.NewWriterSize(f, 1<<20)
|
|
||||||
if _, err := w.WriteString(tsfMagic); err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
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.ColorFrames)))
|
|
||||||
if _, err := w.Write(hdr[:]); err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
var lenBuf [4]byte
|
|
||||||
for _, frame := range data.ColorFrames {
|
|
||||||
binary.LittleEndian.PutUint32(lenBuf[:], uint32(len(frame)))
|
|
||||||
if _, err := w.Write(lenBuf[:]); err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
if _, err := w.Write(frame); err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return w.Flush()
|
|
||||||
}
|
|
||||||
|
|
||||||
func loadTSF(filename string) (*FramesContainer, error) {
|
|
||||||
raw, err := readFrameFile(filename)
|
|
||||||
if err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
invalid := func() error {
|
|
||||||
return fmt.Errorf("invalid frames file %q: corrupt .tsf container", filename)
|
|
||||||
}
|
|
||||||
|
|
||||||
if len(raw) < 18 || string(raw[:4]) != tsfMagic {
|
|
||||||
return nil, invalid()
|
|
||||||
}
|
|
||||||
version := binary.LittleEndian.Uint16(raw[4:])
|
|
||||||
if version != tsfVersion {
|
|
||||||
return nil, fmt.Errorf("unsupported .tsf version %d in %q", version, filename)
|
|
||||||
}
|
|
||||||
fps := math.Float64frombits(binary.LittleEndian.Uint64(raw[6:]))
|
|
||||||
count := binary.LittleEndian.Uint32(raw[14:])
|
|
||||||
|
|
||||||
colorFrames := make([][]byte, 0, count)
|
|
||||||
off := 18
|
|
||||||
for range count {
|
|
||||||
if off+4 > len(raw) {
|
|
||||||
return nil, invalid()
|
|
||||||
}
|
|
||||||
n := int(binary.LittleEndian.Uint32(raw[off:]))
|
|
||||||
off += 4
|
|
||||||
if off+n > len(raw) {
|
|
||||||
return nil, invalid()
|
|
||||||
}
|
|
||||||
colorFrames = append(colorFrames, raw[off:off+n])
|
|
||||||
off += n
|
|
||||||
}
|
|
||||||
|
|
||||||
if len(colorFrames) == 0 || fps <= 0 {
|
|
||||||
return nil, fmt.Errorf(
|
|
||||||
"invalid frames file %q: expected non-empty frames and a positive fps",
|
|
||||||
filename,
|
|
||||||
)
|
|
||||||
}
|
|
||||||
return &FramesContainer{ColorFrames: colorFrames, FPS: fps}, nil
|
|
||||||
}
|
|
||||||
@@ -1,85 +0,0 @@
|
|||||||
package main
|
|
||||||
|
|
||||||
import (
|
|
||||||
"encoding/json"
|
|
||||||
"fmt"
|
|
||||||
"os"
|
|
||||||
"strings"
|
|
||||||
"time"
|
|
||||||
)
|
|
||||||
|
|
||||||
type logLevel int
|
|
||||||
|
|
||||||
const (
|
|
||||||
levelDebug logLevel = 10
|
|
||||||
levelInfo logLevel = 20
|
|
||||||
levelWarn logLevel = 30
|
|
||||||
levelError logLevel = 40
|
|
||||||
)
|
|
||||||
|
|
||||||
var logThreshold = resolveThreshold()
|
|
||||||
|
|
||||||
func resolveThreshold() logLevel {
|
|
||||||
switch strings.ToLower(strings.TrimSpace(os.Getenv("LOG_LEVEL"))) {
|
|
||||||
case "debug":
|
|
||||||
return levelDebug
|
|
||||||
case "warn":
|
|
||||||
return levelWarn
|
|
||||||
case "error":
|
|
||||||
return levelError
|
|
||||||
default:
|
|
||||||
return levelInfo
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func sanitize(value any) string {
|
|
||||||
return sanitizeN(value, 200)
|
|
||||||
}
|
|
||||||
|
|
||||||
func sanitizeN(value any, maxLength int) string {
|
|
||||||
var str string
|
|
||||||
switch v := value.(type) {
|
|
||||||
case nil:
|
|
||||||
str = ""
|
|
||||||
case string:
|
|
||||||
str = v
|
|
||||||
default:
|
|
||||||
str = fmt.Sprint(v)
|
|
||||||
}
|
|
||||||
var b strings.Builder
|
|
||||||
for _, r := range str {
|
|
||||||
if r < 0x20 || (r >= 0x7f && r <= 0x9f) {
|
|
||||||
b.WriteRune('�')
|
|
||||||
} else {
|
|
||||||
b.WriteRune(r)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
out := []rune(b.String())
|
|
||||||
if len(out) > maxLength {
|
|
||||||
return string(out[:maxLength]) + "…"
|
|
||||||
}
|
|
||||||
return string(out)
|
|
||||||
}
|
|
||||||
|
|
||||||
func emit(level logLevel, name string, stream *os.File, args []any) {
|
|
||||||
if level < logThreshold {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
parts := make([]string, len(args))
|
|
||||||
for i, a := range args {
|
|
||||||
if s, ok := a.(string); ok {
|
|
||||||
parts[i] = s
|
|
||||||
} else if b, err := json.Marshal(a); err == nil {
|
|
||||||
parts[i] = string(b)
|
|
||||||
} else {
|
|
||||||
parts[i] = fmt.Sprint(a)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
ts := time.Now().UTC().Format("2006-01-02T15:04:05.000Z")
|
|
||||||
_, _ = 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) }
|
|
||||||
@@ -1,9 +0,0 @@
|
|||||||
//go:build !unix
|
|
||||||
|
|
||||||
package main
|
|
||||||
|
|
||||||
import "os"
|
|
||||||
|
|
||||||
func readFrameFile(filename string) ([]byte, error) {
|
|
||||||
return os.ReadFile(filename)
|
|
||||||
}
|
|
||||||
@@ -1,30 +0,0 @@
|
|||||||
//go:build unix
|
|
||||||
|
|
||||||
package main
|
|
||||||
|
|
||||||
import (
|
|
||||||
"os"
|
|
||||||
"syscall"
|
|
||||||
)
|
|
||||||
|
|
||||||
func readFrameFile(filename string) ([]byte, error) {
|
|
||||||
f, err := os.Open(filename)
|
|
||||||
if err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
defer func() { _ = f.Close() }()
|
|
||||||
|
|
||||||
info, err := f.Stat()
|
|
||||||
if err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
size := info.Size()
|
|
||||||
if size <= 0 || size != int64(int(size)) {
|
|
||||||
return os.ReadFile(filename)
|
|
||||||
}
|
|
||||||
data, err := syscall.Mmap(int(f.Fd()), 0, int(size), syscall.PROT_READ, syscall.MAP_SHARED)
|
|
||||||
if err != nil {
|
|
||||||
return os.ReadFile(filename)
|
|
||||||
}
|
|
||||||
return data, nil
|
|
||||||
}
|
|
||||||
@@ -1,38 +0,0 @@
|
|||||||
package main
|
|
||||||
|
|
||||||
import (
|
|
||||||
"path/filepath"
|
|
||||||
"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)
|
|
||||||
}
|
|
||||||
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.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) }
|
|
||||||
-540
@@ -1,540 +0,0 @@
|
|||||||
package main
|
|
||||||
|
|
||||||
import (
|
|
||||||
"encoding/binary"
|
|
||||||
"errors"
|
|
||||||
"fmt"
|
|
||||||
"math/rand"
|
|
||||||
"net"
|
|
||||||
"strings"
|
|
||||||
"sync"
|
|
||||||
"time"
|
|
||||||
|
|
||||||
"golang.org/x/crypto/ssh"
|
|
||||||
)
|
|
||||||
|
|
||||||
const (
|
|
||||||
clearScreen = "\x1b[2J\x1b[0f"
|
|
||||||
hideCursor = "\x1b[?25l"
|
|
||||||
showCursor = "\x1b[?25h"
|
|
||||||
syncStart = "\x1b[?2026h"
|
|
||||||
syncEnd = "\x1b[?2026l"
|
|
||||||
homeCursor = "\x1b[H"
|
|
||||||
)
|
|
||||||
|
|
||||||
type ConnectionTracker struct {
|
|
||||||
mu sync.Mutex
|
|
||||||
counts map[string]int
|
|
||||||
total int
|
|
||||||
}
|
|
||||||
|
|
||||||
func newConnectionTracker() *ConnectionTracker {
|
|
||||||
return &ConnectionTracker{counts: make(map[string]int)}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (t *ConnectionTracker) increment(ip string) int {
|
|
||||||
t.mu.Lock()
|
|
||||||
defer t.mu.Unlock()
|
|
||||||
t.counts[ip]++
|
|
||||||
t.total++
|
|
||||||
return t.counts[ip]
|
|
||||||
}
|
|
||||||
|
|
||||||
func (t *ConnectionTracker) decrement(ip string) {
|
|
||||||
t.mu.Lock()
|
|
||||||
defer t.mu.Unlock()
|
|
||||||
if _, ok := t.counts[ip]; !ok {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
t.counts[ip]--
|
|
||||||
if t.total > 0 {
|
|
||||||
t.total--
|
|
||||||
}
|
|
||||||
if t.counts[ip] <= 0 {
|
|
||||||
delete(t.counts, ip)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (t *ConnectionTracker) totalCount() int {
|
|
||||||
t.mu.Lock()
|
|
||||||
defer t.mu.Unlock()
|
|
||||||
return t.total
|
|
||||||
}
|
|
||||||
|
|
||||||
func (t *ConnectionTracker) hasReachedLimits(ip string, maxPerIP, maxTotal int) bool {
|
|
||||||
t.mu.Lock()
|
|
||||||
defer t.mu.Unlock()
|
|
||||||
return t.total >= maxTotal || t.counts[ip] >= maxPerIP
|
|
||||||
}
|
|
||||||
|
|
||||||
type frameSet struct {
|
|
||||||
data *FramesContainer
|
|
||||||
renderer *FrameRenderer
|
|
||||||
}
|
|
||||||
|
|
||||||
type Server struct {
|
|
||||||
config Config
|
|
||||||
sshConfig *ssh.ServerConfig
|
|
||||||
sets []frameSet
|
|
||||||
tracker *ConnectionTracker
|
|
||||||
fakeLogin *string
|
|
||||||
goodbye *string
|
|
||||||
listener net.Listener
|
|
||||||
closeOnce sync.Once
|
|
||||||
}
|
|
||||||
|
|
||||||
type ServerDeps struct {
|
|
||||||
Config Config
|
|
||||||
HostKeys []ssh.Signer
|
|
||||||
BannerText *string
|
|
||||||
FakeLoginText *string
|
|
||||||
GoodbyeText *string
|
|
||||||
VideoSets []*FramesContainer
|
|
||||||
}
|
|
||||||
|
|
||||||
func clampDimension(value, max int) int {
|
|
||||||
if value < 1 {
|
|
||||||
return 1
|
|
||||||
}
|
|
||||||
if value > max {
|
|
||||||
return max
|
|
||||||
}
|
|
||||||
return value
|
|
||||||
}
|
|
||||||
|
|
||||||
func createServer(deps ServerDeps) *Server {
|
|
||||||
config := deps.Config
|
|
||||||
|
|
||||||
cache := newRenderCache(int64(config.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,
|
|
||||||
}, cache),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
sshConfig := &ssh.ServerConfig{
|
|
||||||
MaxAuthTries: config.MaxAuthAttempts,
|
|
||||||
PasswordCallback: func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
|
|
||||||
ip := hostOnly(conn.RemoteAddr().String())
|
|
||||||
if config.LogCredentials {
|
|
||||||
logInfo(fmt.Sprintf(
|
|
||||||
`Auth attempt from %s method=password user="%s" pass="%s"`,
|
|
||||||
ip, sanitizeN(conn.User(), 128), sanitizeN(string(password), 128),
|
|
||||||
))
|
|
||||||
}
|
|
||||||
if conn.User() == "" || len(password) == 0 {
|
|
||||||
return nil, errors.New("password rejected")
|
|
||||||
}
|
|
||||||
return nil, nil
|
|
||||||
},
|
|
||||||
}
|
|
||||||
sshConfig.KeyExchanges = []string{
|
|
||||||
"mlkem768x25519-sha256",
|
|
||||||
"curve25519-sha256",
|
|
||||||
"curve25519-sha256@libssh.org",
|
|
||||||
"ecdh-sha2-nistp256",
|
|
||||||
"ecdh-sha2-nistp384",
|
|
||||||
"ecdh-sha2-nistp521",
|
|
||||||
"diffie-hellman-group14-sha256",
|
|
||||||
}
|
|
||||||
if deps.BannerText != nil {
|
|
||||||
banner := *deps.BannerText
|
|
||||||
sshConfig.BannerCallback = func(_ ssh.ConnMetadata) string { return banner }
|
|
||||||
}
|
|
||||||
for _, key := range deps.HostKeys {
|
|
||||||
sshConfig.AddHostKey(key)
|
|
||||||
}
|
|
||||||
|
|
||||||
return &Server{
|
|
||||||
config: config,
|
|
||||||
sshConfig: sshConfig,
|
|
||||||
sets: sets,
|
|
||||||
tracker: newConnectionTracker(),
|
|
||||||
fakeLogin: deps.FakeLoginText,
|
|
||||||
goodbye: deps.GoodbyeText,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func hostOnly(addr string) string {
|
|
||||||
host, _, err := net.SplitHostPort(addr)
|
|
||||||
if err != nil {
|
|
||||||
return addr
|
|
||||||
}
|
|
||||||
return host
|
|
||||||
}
|
|
||||||
|
|
||||||
func (s *Server) Listen(host string, port int) error {
|
|
||||||
listener, err := net.Listen("tcp", net.JoinHostPort(host, fmt.Sprint(port)))
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
s.listener = listener
|
|
||||||
logInfo(fmt.Sprintf("TrollSSH listening on %s:%d", host, port))
|
|
||||||
for {
|
|
||||||
conn, err := listener.Accept()
|
|
||||||
if err != nil {
|
|
||||||
if errors.Is(err, net.ErrClosed) {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
go s.handleConn(conn)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (s *Server) Close() {
|
|
||||||
s.closeOnce.Do(func() {
|
|
||||||
if s.listener != nil {
|
|
||||||
_ = s.listener.Close()
|
|
||||||
}
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
func (s *Server) handleConn(conn net.Conn) {
|
|
||||||
ip := hostOnly(conn.RemoteAddr().String())
|
|
||||||
|
|
||||||
if s.tracker.hasReachedLimits(ip, s.config.MaxConnections, s.config.MaxTotalConnections) {
|
|
||||||
_ = conn.Close()
|
|
||||||
logWarn("Connection rejected (limit reached) from", ip)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
activeForIP := s.tracker.increment(ip)
|
|
||||||
defer s.tracker.decrement(ip)
|
|
||||||
|
|
||||||
if s.config.HandshakeTimeout > 0 {
|
|
||||||
_ = conn.SetDeadline(time.Now().Add(s.config.HandshakeTimeout))
|
|
||||||
}
|
|
||||||
|
|
||||||
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)
|
|
||||||
} else {
|
|
||||||
logWarn(fmt.Sprintf("Client error from %s:", ip), sanitize(err.Error()))
|
|
||||||
}
|
|
||||||
_ = conn.Close()
|
|
||||||
return
|
|
||||||
}
|
|
||||||
_ = conn.SetDeadline(time.Time{})
|
|
||||||
logDebug("Handshake from", ip)
|
|
||||||
defer func() { _ = sshConn.Close() }()
|
|
||||||
|
|
||||||
setIndex := rand.Intn(len(s.sets))
|
|
||||||
logInfo(fmt.Sprintf(
|
|
||||||
"New connection from %s (ip=%d, total=%d) -> playing %q",
|
|
||||||
ip, activeForIP, s.tracker.totalCount(), s.sets[setIndex].data.Name,
|
|
||||||
))
|
|
||||||
|
|
||||||
go ssh.DiscardRequests(reqs)
|
|
||||||
|
|
||||||
for newChannel := range chans {
|
|
||||||
if newChannel.ChannelType() != "session" {
|
|
||||||
_ = newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
channel, requests, err := newChannel.Accept()
|
|
||||||
if err != nil {
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
go s.handleSession(sshConn, channel, requests, ip, setIndex)
|
|
||||||
}
|
|
||||||
logInfo("Client closed connection from", ip)
|
|
||||||
}
|
|
||||||
|
|
||||||
type termSize struct {
|
|
||||||
mu sync.Mutex
|
|
||||||
width int
|
|
||||||
height int
|
|
||||||
}
|
|
||||||
|
|
||||||
func (t *termSize) set(w, h, maxDim int) {
|
|
||||||
t.mu.Lock()
|
|
||||||
t.width = clampDimension(w, maxDim)
|
|
||||||
t.height = clampDimension(h, maxDim)
|
|
||||||
t.mu.Unlock()
|
|
||||||
}
|
|
||||||
|
|
||||||
func (t *termSize) get() (int, int) {
|
|
||||||
t.mu.Lock()
|
|
||||||
defer t.mu.Unlock()
|
|
||||||
return t.width, t.height
|
|
||||||
}
|
|
||||||
|
|
||||||
func parseDims(payload []byte) (cols, rows int, ok bool) {
|
|
||||||
if len(payload) < 8 {
|
|
||||||
return 0, 0, false
|
|
||||||
}
|
|
||||||
// pty-req prefixes cols/rows with a TERM string; window-change does not.
|
|
||||||
offset := 0
|
|
||||||
strLen := binary.BigEndian.Uint32(payload)
|
|
||||||
if int(strLen)+12 <= len(payload) {
|
|
||||||
offset = 4 + int(strLen)
|
|
||||||
}
|
|
||||||
if len(payload) < offset+8 {
|
|
||||||
return 0, 0, false
|
|
||||||
}
|
|
||||||
cols = int(binary.BigEndian.Uint32(payload[offset:]))
|
|
||||||
rows = int(binary.BigEndian.Uint32(payload[offset+4:]))
|
|
||||||
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,
|
|
||||||
requests <-chan *ssh.Request,
|
|
||||||
ip string,
|
|
||||||
initialSetIndex int,
|
|
||||||
) {
|
|
||||||
size := &termSize{}
|
|
||||||
size.set(80, 24, s.config.MaxDimension)
|
|
||||||
tier := colorTierTrueColor
|
|
||||||
if s.config.ForceGrayscale {
|
|
||||||
tier = colorTierNone
|
|
||||||
}
|
|
||||||
|
|
||||||
started := false
|
|
||||||
for req := range requests {
|
|
||||||
switch req.Type {
|
|
||||||
case "pty-req":
|
|
||||||
logDebug("Opening pty for session", ip)
|
|
||||||
if cols, rows, ok := parseDims(req.Payload); ok {
|
|
||||||
size.set(cols, rows, s.config.MaxDimension)
|
|
||||||
}
|
|
||||||
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 {
|
|
||||||
cols := int(binary.BigEndian.Uint32(req.Payload))
|
|
||||||
rows := int(binary.BigEndian.Uint32(req.Payload[4:]))
|
|
||||||
size.set(cols, rows, s.config.MaxDimension)
|
|
||||||
}
|
|
||||||
if req.WantReply {
|
|
||||||
_ = req.Reply(true, nil)
|
|
||||||
}
|
|
||||||
case "exec":
|
|
||||||
command := ""
|
|
||||||
if len(req.Payload) >= 4 {
|
|
||||||
n := binary.BigEndian.Uint32(req.Payload)
|
|
||||||
if int(n)+4 <= len(req.Payload) {
|
|
||||||
command = string(req.Payload[4 : 4+n])
|
|
||||||
}
|
|
||||||
}
|
|
||||||
logInfo(fmt.Sprintf("Client %s attempted exec: %q", ip, sanitizeN(command, 512)))
|
|
||||||
_ = req.Reply(true, nil)
|
|
||||||
if !started {
|
|
||||||
started = true
|
|
||||||
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, tier)
|
|
||||||
}
|
|
||||||
default:
|
|
||||||
if req.WantReply {
|
|
||||||
_ = req.Reply(false, nil)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (s *Server) pickNextSetIndex(exclude int) int {
|
|
||||||
if len(s.sets) <= 1 {
|
|
||||||
return exclude
|
|
||||||
}
|
|
||||||
next := exclude
|
|
||||||
for next == exclude {
|
|
||||||
next = rand.Intn(len(s.sets))
|
|
||||||
}
|
|
||||||
return next
|
|
||||||
}
|
|
||||||
|
|
||||||
func (s *Server) playVideo(
|
|
||||||
sshConn *ssh.ServerConn,
|
|
||||||
channel ssh.Channel,
|
|
||||||
size *termSize,
|
|
||||||
ip string,
|
|
||||||
setIndex int,
|
|
||||||
keepAspectRatio bool,
|
|
||||||
tier colorTier,
|
|
||||||
) {
|
|
||||||
config := s.config
|
|
||||||
current := s.sets[setIndex]
|
|
||||||
|
|
||||||
w, h := size.get()
|
|
||||||
logDebug(fmt.Sprintf("Terminal size %dx%d for %s", w, h, ip))
|
|
||||||
|
|
||||||
defer func() { _, _ = channel.Write([]byte(showCursor)) }()
|
|
||||||
|
|
||||||
if s.fakeLogin != nil {
|
|
||||||
_, _ = channel.Write([]byte(clearScreen))
|
|
||||||
_, _ = channel.Write([]byte(*s.fakeLogin))
|
|
||||||
}
|
|
||||||
|
|
||||||
done := make(chan struct{})
|
|
||||||
var doneOnce sync.Once
|
|
||||||
closeSession := func() {
|
|
||||||
doneOnce.Do(func() { close(done) })
|
|
||||||
}
|
|
||||||
|
|
||||||
switchCh := make(chan int, 8)
|
|
||||||
go func() {
|
|
||||||
buf := make([]byte, 256)
|
|
||||||
var lastSwitch time.Time
|
|
||||||
for {
|
|
||||||
n, err := channel.Read(buf)
|
|
||||||
if err != nil {
|
|
||||||
closeSession()
|
|
||||||
return
|
|
||||||
}
|
|
||||||
if !config.AllowUserControl {
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
str := string(buf[:n])
|
|
||||||
delta := 0
|
|
||||||
if strings.Contains(str, "\x1b[C") || strings.Contains(str, "\x1b[A") {
|
|
||||||
delta = 1
|
|
||||||
} else if strings.Contains(str, "\x1b[D") || strings.Contains(str, "\x1b[B") {
|
|
||||||
delta = -1
|
|
||||||
}
|
|
||||||
if delta == 0 {
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
now := time.Now()
|
|
||||||
if now.Sub(lastSwitch) < config.SwitchDebounce {
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
lastSwitch = now
|
|
||||||
select {
|
|
||||||
case switchCh <- delta:
|
|
||||||
default:
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}()
|
|
||||||
|
|
||||||
select {
|
|
||||||
case <-time.After(config.LoginDelay):
|
|
||||||
case <-done:
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
_, _ = channel.Write([]byte(hideCursor))
|
|
||||||
|
|
||||||
frameInterval := func() time.Duration {
|
|
||||||
return time.Duration(float64(time.Second) / current.data.FPS)
|
|
||||||
}
|
|
||||||
|
|
||||||
ticker := time.NewTicker(frameInterval())
|
|
||||||
defer ticker.Stop()
|
|
||||||
|
|
||||||
currentFrame := 0
|
|
||||||
loopCount := 0
|
|
||||||
lastW, lastH := 0, 0
|
|
||||||
var writeBuf []byte
|
|
||||||
|
|
||||||
for {
|
|
||||||
select {
|
|
||||||
case <-done:
|
|
||||||
return
|
|
||||||
|
|
||||||
case delta := <-switchCh:
|
|
||||||
if len(s.sets) <= 1 {
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
setIndex = (setIndex + delta + len(s.sets)) % len(s.sets)
|
|
||||||
current = s.sets[setIndex]
|
|
||||||
currentFrame = 0
|
|
||||||
lastW, lastH = 0, 0
|
|
||||||
logDebug(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)
|
|
||||||
if err != nil {
|
|
||||||
logError("Render error for", ip, sanitize(err.Error()))
|
|
||||||
_ = sshConn.Close()
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
prefix := homeCursor
|
|
||||||
if w != lastW || h != lastH {
|
|
||||||
prefix = clearScreen
|
|
||||||
lastW, lastH = w, h
|
|
||||||
}
|
|
||||||
writeBuf = append(writeBuf[:0], syncStart...)
|
|
||||||
writeBuf = append(writeBuf, prefix...)
|
|
||||||
writeBuf = append(writeBuf, ascii...)
|
|
||||||
writeBuf = append(writeBuf, syncEnd...)
|
|
||||||
if _, err := channel.Write(writeBuf); err != nil {
|
|
||||||
closeSession()
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
currentFrame++
|
|
||||||
if currentFrame < len(current.data.ColorFrames) {
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
|
|
||||||
currentFrame = 0
|
|
||||||
loopCount++
|
|
||||||
if config.MaxLoop > 0 && loopCount >= config.MaxLoop {
|
|
||||||
_, _ = channel.Write([]byte(showCursor + clearScreen))
|
|
||||||
if s.goodbye != nil {
|
|
||||||
_, _ = channel.Write([]byte(*s.goodbye))
|
|
||||||
}
|
|
||||||
time.Sleep(1 * time.Second)
|
|
||||||
logInfo("Playback finished, closing session", ip)
|
|
||||||
_ = channel.Close()
|
|
||||||
_ = sshConn.Close()
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
if config.PlaybackMode == PlaybackRandom {
|
|
||||||
setIndex = s.pickNextSetIndex(setIndex)
|
|
||||||
current = s.sets[setIndex]
|
|
||||||
logInfo(fmt.Sprintf(
|
|
||||||
"Playthrough done for %s, switching to %q", ip, current.data.Name,
|
|
||||||
))
|
|
||||||
ticker.Reset(frameInterval())
|
|
||||||
} else if config.MaxLoop > 0 {
|
|
||||||
logInfo(fmt.Sprintf(
|
|
||||||
"Playthrough done for %s, looping %q (%d/%d)",
|
|
||||||
ip, current.data.Name, loopCount, config.MaxLoop,
|
|
||||||
))
|
|
||||||
} else {
|
|
||||||
logInfo(fmt.Sprintf(
|
|
||||||
"Playthrough done for %s, looping %q (%d)",
|
|
||||||
ip, current.data.Name, loopCount,
|
|
||||||
))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,189 +0,0 @@
|
|||||||
package main
|
|
||||||
|
|
||||||
import (
|
|
||||||
"bytes"
|
|
||||||
"encoding/json"
|
|
||||||
"errors"
|
|
||||||
"fmt"
|
|
||||||
"io"
|
|
||||||
"math"
|
|
||||||
"os/exec"
|
|
||||||
"strconv"
|
|
||||||
"strings"
|
|
||||||
)
|
|
||||||
|
|
||||||
var (
|
|
||||||
jpegSOI = []byte{0xff, 0xd8}
|
|
||||||
jpegEOI = []byte{0xff, 0xd9}
|
|
||||||
)
|
|
||||||
|
|
||||||
type jpegFrameSplitter struct {
|
|
||||||
buffer []byte
|
|
||||||
}
|
|
||||||
|
|
||||||
func (s *jpegFrameSplitter) push(chunk []byte) [][]byte {
|
|
||||||
s.buffer = append(s.buffer, chunk...)
|
|
||||||
var frames [][]byte
|
|
||||||
for {
|
|
||||||
start := bytes.Index(s.buffer, jpegSOI)
|
|
||||||
if start == -1 {
|
|
||||||
break
|
|
||||||
}
|
|
||||||
end := bytes.Index(s.buffer[start+len(jpegSOI):], jpegEOI)
|
|
||||||
if end == -1 {
|
|
||||||
break
|
|
||||||
}
|
|
||||||
frameEnd := start + len(jpegSOI) + end + len(jpegEOI)
|
|
||||||
frame := make([]byte, frameEnd-start)
|
|
||||||
copy(frame, s.buffer[start:frameEnd])
|
|
||||||
frames = append(frames, frame)
|
|
||||||
s.buffer = s.buffer[frameEnd:]
|
|
||||||
}
|
|
||||||
return frames
|
|
||||||
}
|
|
||||||
|
|
||||||
type ffprobeOutput struct {
|
|
||||||
Streams []struct {
|
|
||||||
RFrameRate string `json:"r_frame_rate"`
|
|
||||||
NbFrames string `json:"nb_frames"`
|
|
||||||
Duration string `json:"duration"`
|
|
||||||
} `json:"streams"`
|
|
||||||
Format struct {
|
|
||||||
Duration string `json:"duration"`
|
|
||||||
} `json:"format"`
|
|
||||||
}
|
|
||||||
|
|
||||||
func parseFrameRate(rate string) float64 {
|
|
||||||
if rate == "" {
|
|
||||||
return math.NaN()
|
|
||||||
}
|
|
||||||
parts := strings.SplitN(rate, "/", 2)
|
|
||||||
num, err := strconv.ParseFloat(parts[0], 64)
|
|
||||||
if err != nil {
|
|
||||||
return math.NaN()
|
|
||||||
}
|
|
||||||
if len(parts) == 2 {
|
|
||||||
den, err := strconv.ParseFloat(parts[1], 64)
|
|
||||||
if err != nil || den == 0 {
|
|
||||||
return math.NaN()
|
|
||||||
}
|
|
||||||
return num / den
|
|
||||||
}
|
|
||||||
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",
|
|
||||||
"-show_streams", "-show_format",
|
|
||||||
"-of", "json", path,
|
|
||||||
)
|
|
||||||
probeOut, err := probeCmd.Output()
|
|
||||||
if err != nil {
|
|
||||||
msg := err.Error()
|
|
||||||
var exitErr *exec.ExitError
|
|
||||||
if errors.As(err, &exitErr) && len(exitErr.Stderr) > 0 {
|
|
||||||
msg = strings.TrimSpace(string(exitErr.Stderr))
|
|
||||||
}
|
|
||||||
return fmt.Errorf("ffprobe failed: %s", msg)
|
|
||||||
}
|
|
||||||
|
|
||||||
var probe ffprobeOutput
|
|
||||||
if err := json.Unmarshal(probeOut, &probe); err != nil {
|
|
||||||
return fmt.Errorf("ffprobe failed: %w", err)
|
|
||||||
}
|
|
||||||
if len(probe.Streams) == 0 {
|
|
||||||
return fmt.Errorf("unable to determine a valid video fps")
|
|
||||||
}
|
|
||||||
stream := probe.Streams[0]
|
|
||||||
|
|
||||||
fps := parseFrameRate(stream.RFrameRate)
|
|
||||||
if math.IsNaN(fps) || fps <= 0 {
|
|
||||||
return fmt.Errorf("unable to determine a valid video fps")
|
|
||||||
}
|
|
||||||
|
|
||||||
totalFrames := 0
|
|
||||||
if n, err := strconv.Atoi(stream.NbFrames); err == nil {
|
|
||||||
totalFrames = n
|
|
||||||
} else {
|
|
||||||
durStr := probe.Format.Duration
|
|
||||||
if durStr == "" {
|
|
||||||
durStr = stream.Duration
|
|
||||||
}
|
|
||||||
if d, err := strconv.ParseFloat(durStr, 64); err == nil {
|
|
||||||
totalFrames = int(math.Round(d * fps))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
scaleFilter := fmt.Sprintf(
|
|
||||||
"scale=w=%d:h=%d:force_original_aspect_ratio=decrease",
|
|
||||||
maxDimension, maxDimension,
|
|
||||||
)
|
|
||||||
|
|
||||||
colorFrames, err := extractFrames(path, scaleFilter, "color", maxDimension, totalFrames)
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
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.ColorFrames), output))
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
Reference in New Issue
Block a user