fix: harden tmux runtime teardown and ids

This commit is contained in:
harshitsinghbhandari 2026-05-27 18:20:03 +05:30
parent b5a344ac07
commit 37f1fe269e
2 changed files with 97 additions and 12 deletions

View File

@ -3,6 +3,8 @@ package tmux
import (
"context"
"crypto/sha256"
"encoding/hex"
"errors"
"fmt"
"os"
@ -12,6 +14,7 @@ import (
"strings"
"time"
"github.com/aoagents/agent-orchestrator/backend/internal/domain"
"github.com/aoagents/agent-orchestrator/backend/internal/ports"
)
@ -61,14 +64,14 @@ func New(opts Options) *Runtime {
shellPath = os.Getenv("SHELL")
}
if shellPath == "" {
shellPath = "/bin/zsh"
shellPath = "/bin/sh"
}
return &Runtime{binary: binary, timeout: timeout, shell: shellPath, runner: execRunner{}}
}
func (r *Runtime) Create(ctx context.Context, cfg ports.RuntimeConfig) (ports.RuntimeHandle, error) {
id := string(cfg.SessionID)
if err := validateSessionID(id); err != nil {
id, err := tmuxSessionName(cfg.SessionID)
if err != nil {
return ports.RuntimeHandle{}, err
}
if cfg.WorkspacePath == "" {
@ -94,14 +97,11 @@ func (r *Runtime) Destroy(ctx context.Context, handle ports.RuntimeHandle) error
if err != nil {
return err
}
alive, err := r.IsAlive(ctx, handle)
if err != nil {
return err
}
if !alive {
return nil
}
if _, err := r.run(ctx, killSessionArgs(id)...); err != nil {
var exitErr *exec.ExitError
if errors.As(err, &exitErr) {
return nil
}
return fmt.Errorf("tmux runtime: destroy session %s: %w", id, err)
}
return nil
@ -205,6 +205,43 @@ func (r *Runtime) run(ctx context.Context, args ...string) ([]byte, error) {
return out, nil
}
func tmuxSessionName(id domain.SessionID) (string, error) {
raw := string(id)
if raw == "" {
return "", errors.New("tmux runtime: session id is required")
}
if sessionIDPattern.MatchString(raw) {
return raw, nil
}
return sanitizedSessionName(raw), nil
}
func sanitizedSessionName(raw string) string {
var b strings.Builder
lastDash := false
for _, r := range raw {
valid := (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '_' || r == '-'
if valid {
b.WriteRune(r)
lastDash = false
continue
}
if !lastDash {
b.WriteByte('-')
lastDash = true
}
}
base := strings.Trim(b.String(), "-")
if base == "" {
base = "session"
}
if len(base) > 40 {
base = strings.TrimRight(base[:40], "-")
}
sum := sha256.Sum256([]byte(raw))
return base + "-" + hex.EncodeToString(sum[:4])
}
func validateSessionID(id string) error {
if id == "" {
return errors.New("tmux runtime: session id is required")

View File

@ -12,6 +12,14 @@ import (
"github.com/aoagents/agent-orchestrator/backend/internal/ports"
)
func TestNewDefaultsToPortableShell(t *testing.T) {
t.Setenv("SHELL", "")
r := New(Options{})
if got, want := r.shell, "/bin/sh"; got != want {
t.Fatalf("default shell = %q, want %q", got, want)
}
}
func TestCommandBuilders(t *testing.T) {
if got, want := newSessionArgs("sess-1", "/tmp/ws", "/bin/zsh", "echo hi"), []string{"new-session", "-d", "-s", "sess-1", "-c", "/tmp/ws", "/bin/zsh", "-lc", "echo hi"}; !reflect.DeepEqual(got, want) {
t.Fatalf("newSessionArgs = %#v, want %#v", got, want)
@ -33,6 +41,22 @@ func TestExactTargets(t *testing.T) {
}
}
func TestTmuxSessionNameSanitizesIssueRefs(t *testing.T) {
got, err := tmuxSessionName("repo/issue#42.1")
if err != nil {
t.Fatalf("tmuxSessionName: %v", err)
}
if err := validateSessionID(got); err != nil {
t.Fatalf("sanitized id %q is invalid: %v", got, err)
}
if !strings.HasPrefix(got, "repo-issue-42-1-") {
t.Fatalf("sanitized id = %q, want readable prefix", got)
}
if got == "repo/issue#42.1" {
t.Fatal("sanitized id still contains raw unsafe characters")
}
}
func TestValidateSessionID(t *testing.T) {
valid := []string{"sess-1", "S_2", "abc123"}
for _, id := range valid {
@ -104,6 +128,30 @@ func TestCreateRunsNewSessionAndDisablesStatus(t *testing.T) {
}
}
func TestCreateNormalizesUnsafeSessionID(t *testing.T) {
fr := &fakeRunner{}
r := New(Options{Binary: "tmux-test", Timeout: time.Second, Shell: "/bin/sh"})
r.runner = fr
handle, err := r.Create(context.Background(), ports.RuntimeConfig{
SessionID: "repo/issue#42",
WorkspacePath: "/tmp/ws",
LaunchCommand: "echo ready",
})
if err != nil {
t.Fatalf("Create: %v", err)
}
if err := validateSessionID(handle.ID); err != nil {
t.Fatalf("handle id %q invalid: %v", handle.ID, err)
}
if handle.ID == "repo/issue#42" {
t.Fatal("handle kept unsafe raw session id")
}
if got := fr.calls[0].args[3]; got != handle.ID {
t.Fatalf("tmux session arg = %q, want handle id %q", got, handle.ID)
}
}
func TestSendMessageUsesLiteralForShortInput(t *testing.T) {
fr := &fakeRunner{}
r := New(Options{Timeout: time.Second})
@ -164,8 +212,8 @@ func TestDestroyIsIdempotentWhenSessionMissing(t *testing.T) {
if err := r.Destroy(context.Background(), ports.RuntimeHandle{ID: "sess-1", RuntimeName: runtimeName}); err != nil {
t.Fatalf("Destroy: %v", err)
}
if len(fr.calls) != 1 || fr.calls[0].args[0] != "has-session" {
t.Fatalf("calls = %#v, want only has-session", fr.calls)
if len(fr.calls) != 1 || fr.calls[0].args[0] != "kill-session" {
t.Fatalf("calls = %#v, want only kill-session", fr.calls)
}
}