mirror of
https://github.com/YuzuZensai/TrollSSH.git
synced 2026-09-13 19:29:07 +00:00
♻️ refactor: split src into internal packages and cmd
This commit is contained in:
@@ -0,0 +1,61 @@
|
||||
package sshserver
|
||||
|
||||
import (
|
||||
"crypto/ed25519"
|
||||
"crypto/rand"
|
||||
"crypto/rsa"
|
||||
"encoding/pem"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"golang.org/x/crypto/ssh"
|
||||
)
|
||||
|
||||
func generateAndSave(keyPath, keyType string) error {
|
||||
fmt.Printf("Generating %s host key...\n", keyType)
|
||||
|
||||
var key any
|
||||
var err error
|
||||
if keyType == "rsa" {
|
||||
key, err = rsa.GenerateKey(rand.Reader, 4096)
|
||||
} else {
|
||||
_, key, err = ed25519.GenerateKey(rand.Reader)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
block, err := ssh.MarshalPrivateKey(key, "")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return os.WriteFile(keyPath, pem.EncodeToMemory(block), 0o600)
|
||||
}
|
||||
|
||||
func EnsureHostKeys(configDir string) ([]ssh.Signer, error) {
|
||||
keys := []struct{ file, keyType string }{
|
||||
{"id_rsa", "rsa"},
|
||||
{"id_ed25519", "ed25519"},
|
||||
}
|
||||
|
||||
signers := make([]ssh.Signer, 0, len(keys))
|
||||
for _, k := range keys {
|
||||
keyPath := filepath.Join(configDir, k.file)
|
||||
if _, err := os.Stat(keyPath); os.IsNotExist(err) {
|
||||
if err := generateAndSave(keyPath, k.keyType); err != nil {
|
||||
return nil, fmt.Errorf("failed to generate %s host key: %w", k.keyType, err)
|
||||
}
|
||||
}
|
||||
raw, err := os.ReadFile(keyPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
signer, err := ssh.ParsePrivateKey(raw)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to parse host key %q: %w", keyPath, err)
|
||||
}
|
||||
signers = append(signers, signer)
|
||||
}
|
||||
return signers, nil
|
||||
}
|
||||
@@ -0,0 +1,778 @@
|
||||
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)
|
||||
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
|
||||
}
|
||||
|
||||
func (t *termSize) set(w, h, maxDimension, maxCells int, force bool) {
|
||||
t.mu.Lock()
|
||||
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()
|
||||
}
|
||||
|
||||
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,137 @@
|
||||
package sshserver
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"golang.org/x/crypto/ssh"
|
||||
)
|
||||
|
||||
func TestConnectionTracker(t *testing.T) {
|
||||
tr := newConnectionTracker()
|
||||
if _, _, ok := tr.tryAcquire("1.2.3.4", 2, 100); !ok {
|
||||
t.Fatal("first acquire failed")
|
||||
}
|
||||
if _, _, ok := tr.tryAcquire("1.2.3.4", 2, 100); !ok {
|
||||
t.Fatal("second acquire failed")
|
||||
}
|
||||
if _, _, ok := tr.tryAcquire("1.2.3.4", 2, 100); ok {
|
||||
t.Error("expected per-ip limit rejection")
|
||||
}
|
||||
tr.release("1.2.3.4")
|
||||
tr.release("1.2.3.4")
|
||||
if tr.totalCount() != 0 {
|
||||
t.Errorf("total = %d", tr.totalCount())
|
||||
}
|
||||
if _, _, ok := tr.tryAcquire("1.2.3.4", 2, 100); !ok {
|
||||
t.Error("limit should be cleared")
|
||||
}
|
||||
}
|
||||
|
||||
func TestConnectionTrackerConcurrentLimit(t *testing.T) {
|
||||
tracker := newConnectionTracker()
|
||||
start := make(chan struct{})
|
||||
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 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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user