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