fix(pssh): apply pty-req before running the session's command

Commit `d642d4b`

Antonio Mika · 2026-10-22 10:46:32 -0400 EDT

parent 37379a4

fix(pssh): apply pty-req before running the session's command

6 files changed, +109, -52

M pkg/apps/pastes/ssh.go

M pkg/apps/pico/ssh.go

M pkg/apps/prose/ssh.go

M pkg/pssh/pty.go

M pkg/pssh/server.go

M pkg/pssh/server_test.go

M pkg/apps/pastes/ssh.go

+1, -2

......@@ -5,7 +5,6 @@ import (
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/picosh/pico/pkg/db/postgres"
	"github.com/picosh/pico/pkg/pssh"

......@@ -53,7 +52,7 @@ func StartSshServer() {
			scp.Middleware(handler),
			rsync.Middleware(handler),
			auth.Middleware(handler),
			pssh.PtyMdw(pssh.DeprecatedNotice(), 200*time.Millisecond),
			pssh.PtyMdw(pssh.DeprecatedNotice()),
			pssh.LogMiddleware(handler, dbh),
		},
		[]pssh.SSHServerMiddleware{

M pkg/apps/pico/ssh.go

+1, -2

......@@ -5,7 +5,6 @@ import (
	"os"
	"os/signal"
	"syscall"
	"time"

	pgsdb "github.com/picosh/pico/pkg/apps/pgs/db"
	"github.com/picosh/pico/pkg/db/postgres"

......@@ -106,7 +105,7 @@ func StartSshServer() {
						PgsDB:      pgsDB,
						Logger:     pssh.GetLogger(sesh),
					}
					return pssh.PtyMdw(createTui(shrd), 200*time.Millisecond)(next)(sesh)
					return pssh.PtyMdw(createTui(shrd))(next)(sesh)
				}
			},
			Middleware(cliHandler),

M pkg/apps/prose/ssh.go

+1, -2

......@@ -5,7 +5,6 @@ import (
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/picosh/pico/pkg/db/postgres"
	"github.com/picosh/pico/pkg/pssh"

......@@ -63,7 +62,7 @@ func StartSshServer() {
			scp.Middleware(handler),
			rsync.Middleware(handler),
			auth.Middleware(handler),
			pssh.PtyMdw(pssh.DeprecatedNotice(), 200*time.Millisecond),
			pssh.PtyMdw(pssh.DeprecatedNotice()),
			pssh.LogMiddleware(handler, dbh),
		},
		[]pssh.SSHServerMiddleware{

M pkg/pssh/pty.go

+1, -15

......@@ -3,7 +3,6 @@ package pssh
import (
	"encoding/binary"
	"fmt"
	"time"
)

func SessionMessage(sesh *SSHServerConnSession, msg string) {

......@@ -24,22 +23,9 @@ func DeprecatedNotice() SSHServerMiddleware {
	}
}

func PtyMdw(mdw SSHServerMiddleware, waitTimeout time.Duration) SSHServerMiddleware {
func PtyMdw(mdw SSHServerMiddleware) SSHServerMiddleware {
	return func(next SSHServerHandler) SSHServerHandler {
		return func(sesh *SSHServerConnSession) error {
			if waitTimeout > 0 {
				afterTime := time.Now().Add(waitTimeout)

				for time.Now().After(afterTime) {
					_, _, ok := sesh.Pty()
					if ok {
						break
					}

					time.Sleep(50 * time.Millisecond)
				}
			}

			_, _, ok := sesh.Pty()
			if !ok {
				return next(sesh)

M pkg/pssh/server.go

+26, -31

......@@ -186,6 +186,24 @@ func (s *SSHServerConnSession) Pty() (*Pty, <-chan Window, bool) {
	return s.pty, s.winch, true
}

func (s *SSHServerConnSession) handlePtyReq(req *ssh.Request) {
	ptyReq, ok := parsePtyRequest(req.Payload)

	s.mu.Lock()
	if ok && s.pty == nil {
		s.pty = &ptyReq
		s.winch = make(chan Window, 1)
		s.winch <- ptyReq.Window
	} else {
		ok = false
	}
	s.mu.Unlock()

	if err := req.Reply(ok, nil); err != nil {
		s.Logger.Error("pty-req reply", "err", err)
	}
}

// Write overrides the embedded Channel's Write to normalize line endings when PTY is allocated.
func (s *SSHServerConnSession) Write(p []byte) (n int, err error) {
	s.mu.Lock()

......@@ -510,6 +528,14 @@ func NewSSHServer(ctx context.Context, logger *slog.Logger, config *SSHServerCon
						return nil
					}

					// Clients send pty-req before shell or exec without waiting for the
					// reply, so it's applied here, before the next request is read.
					if req.Type == "pty-req" {
						sc.Logger.Info("new session request", "type", req.Type, "wantReply", req.WantReply, "payload", req.Payload)
						sesh.handlePtyReq(req)
						continue
					}

					go func() {
						sc.Logger.Info("new session request", "type", req.Type, "wantReply", req.WantReply, "payload", req.Payload)
						switch req.Type {

......@@ -626,37 +652,6 @@ func NewSSHServer(ctx context.Context, logger *slog.Logger, config *SSHServerCon
							if err != nil {
								sc.Logger.Error("subsystem close", "err", err)
							}
						case "pty-req":
							sesh.mu.Lock()
							found := sesh.pty != nil
							sesh.mu.Unlock()
							if found {
								err := req.Reply(false, nil)
								if err != nil {
									sc.Logger.Error("pty-req reply", "err", err)
								}
								return
							}

							ptyReq, ok := parsePtyRequest(req.Payload)
							if !ok {
								err := req.Reply(false, nil)
								if err != nil {
									sc.Logger.Error("pty-req reply", "err", err)
								}
								return
							}

							sesh.mu.Lock()
							sesh.pty = &ptyReq
							sesh.winch = make(chan Window, 1)
							sesh.mu.Unlock()

							sesh.winch <- ptyReq.Window
							err := req.Reply(ok, nil)
							if err != nil {
								sc.Logger.Error("pty-req reply", "err", err)
							}
						case "window-change":
							sesh.mu.Lock()
							found := sesh.pty != nil

M pkg/pssh/server_test.go

+79, -0

......@@ -709,3 +709,82 @@ func TestSessionContextEndsWhenClientCloses(t *testing.T) {
		t.Fatal("command still running after the client closed its session")
	}
}

// OpenSSH sends pty-req and then shell or exec without waiting for the first
// reply, so the pty has to be set by the time the command's middleware runs.
func TestPtyVisibleToCommand(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	user := GenerateKey()
	server := pssh.NewSSHServer(ctx, slog.New(slog.DiscardHandler), &pssh.SSHServerConfig{
		ListenAddr: "127.0.0.1:0",
		Middleware: []pssh.SSHServerMiddleware{
			func(next pssh.SSHServerHandler) pssh.SSHServerHandler {
				return func(sesh *pssh.SSHServerConnSession) error {
					if _, _, ok := sesh.Pty(); !ok {
						_, _ = sesh.Write([]byte("no pty"))
					}
					return next(sesh)
				}
			},
		},
		ServerConfig: &ssh.ServerConfig{
			NoClientAuth: true,
			NoClientAuthCallback: func(ssh.ConnMetadata) (*ssh.Permissions, error) {
				return &ssh.Permissions{Extensions: map[string]string{}}, nil
			},
		},
	})
	server.Config.AddHostKey(user.signer)
	go func() { _ = server.ListenAndServe() }()

	var addr string
	for range 50 {
		server.Mu.Lock()
		if server.Listener != nil {
			addr = server.Listener.Addr().String()
		}
		server.Mu.Unlock()
		if addr != "" {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}
	if addr == "" {
		t.Fatal("server listener not ready")
	}

	client, err := ssh.Dial("tcp", addr, &ssh.ClientConfig{User: "user", HostKeyCallback: ssh.InsecureIgnoreHostKey()})
	if err != nil {
		t.Fatal(err)
	}
	defer func() { _ = client.Close() }()

	ptyReq := ssh.Marshal(struct {
		Term          string
		Columns, Rows uint32
		Width, Height uint32
		Modes         string
	}{Term: "xterm", Columns: 80, Rows: 24})

	missed := 0
	for range 200 {
		session, err := client.NewSession()
		if err != nil {
			t.Fatal(err)
		}
		if _, err := session.SendRequest("pty-req", false, ptyReq); err != nil {
			t.Fatal(err)
		}
		out, _ := session.Output("cmd")
		if strings.Contains(string(out), "no pty") {
			missed++
		}
		_ = session.Close()
	}

	if missed > 0 {
		t.Fatalf("command ran before its pty request in %d of 200 sessions", missed)
	}
}

...	...		@@ -186,6 +186,24 @@ func (s SSHServerConnSession) Pty() (Pty, <-chan Window, bool) {
186	186		return s.pty, s.winch, true
187	187		}
188	188
	189	+	func (s SSHServerConnSession) handlePtyReq(req ssh.Request) {
	190	+	ptyReq, ok := parsePtyRequest(req.Payload)
	191	+
	192	+	s.mu.Lock()
	193	+	if ok && s.pty == nil {
	194	+	s.pty = &ptyReq
	195	+	s.winch = make(chan Window, 1)
	196	+	s.winch <- ptyReq.Window
	197	+	} else {
	198	+	ok = false
	199	+	}
	200	+	s.mu.Unlock()
	201	+
	202	+	if err := req.Reply(ok, nil); err != nil {
	203	+	s.Logger.Error("pty-req reply", "err", err)
	204	+	}
	205	+	}
	206	+
189	207		// Write overrides the embedded Channel's Write to normalize line endings when PTY is allocated.
190	208		func (s *SSHServerConnSession) Write(p []byte) (n int, err error) {
191	209		s.mu.Lock()

...	...		@@ -510,6 +528,14 @@ func NewSSHServer(ctx context.Context, logger slog.Logger, config SSHServerCon
510	528		return nil
511	529		}
512	530
	531	+	// Clients send pty-req before shell or exec without waiting for the
	532	+	// reply, so it's applied here, before the next request is read.
	533	+	if req.Type == "pty-req" {
	534	+	sc.Logger.Info("new session request", "type", req.Type, "wantReply", req.WantReply, "payload", req.Payload)
	535	+	sesh.handlePtyReq(req)
	536	+	continue
	537	+	}
	538	+
513	539		go func() {
514	540		sc.Logger.Info("new session request", "type", req.Type, "wantReply", req.WantReply, "payload", req.Payload)
515	541		switch req.Type {

...	...		@@ -626,37 +652,6 @@ func NewSSHServer(ctx context.Context, logger slog.Logger, config SSHServerCon
626	652		if err != nil {
627	653		sc.Logger.Error("subsystem close", "err", err)
628	654		}
629		-	case "pty-req":
630		-	sesh.mu.Lock()
631		-	found := sesh.pty != nil
632		-	sesh.mu.Unlock()
633		-	if found {
634		-	err := req.Reply(false, nil)
635		-	if err != nil {
636		-	sc.Logger.Error("pty-req reply", "err", err)
637		-	}
638		-	return
639		-	}
640		-
641		-	ptyReq, ok := parsePtyRequest(req.Payload)
642		-	if !ok {
643		-	err := req.Reply(false, nil)
644		-	if err != nil {
645		-	sc.Logger.Error("pty-req reply", "err", err)
646		-	}
647		-	return
648		-	}
649		-
650		-	sesh.mu.Lock()
651		-	sesh.pty = &ptyReq
652		-	sesh.winch = make(chan Window, 1)
653		-	sesh.mu.Unlock()
654		-
655		-	sesh.winch <- ptyReq.Window
656		-	err := req.Reply(ok, nil)
657		-	if err != nil {
658		-	sc.Logger.Error("pty-req reply", "err", err)
659		-	}
660	655		case "window-change":
661	656		sesh.mu.Lock()
662	657		found := sesh.pty != nil

...	...		@@ -709,3 +709,82 @@ func TestSessionContextEndsWhenClientCloses(t *testing.T) {
709	709		t.Fatal("command still running after the client closed its session")
710	710		}
711	711		}
	712	+
	713	+	// OpenSSH sends pty-req and then shell or exec without waiting for the first
	714	+	// reply, so the pty has to be set by the time the command's middleware runs.
	715	+	func TestPtyVisibleToCommand(t *testing.T) {
	716	+	ctx, cancel := context.WithCancel(context.Background())
	717	+	defer cancel()
	718	+
	719	+	user := GenerateKey()
	720	+	server := pssh.NewSSHServer(ctx, slog.New(slog.DiscardHandler), &pssh.SSHServerConfig{
	721	+	ListenAddr: "127.0.0.1:0",
	722	+	Middleware: []pssh.SSHServerMiddleware{
	723	+	func(next pssh.SSHServerHandler) pssh.SSHServerHandler {
	724	+	return func(sesh *pssh.SSHServerConnSession) error {
	725	+	if _, _, ok := sesh.Pty(); !ok {
	726	+	_, _ = sesh.Write([]byte("no pty"))
	727	+	}
	728	+	return next(sesh)
	729	+	}
	730	+	},
	731	+	},
	732	+	ServerConfig: &ssh.ServerConfig{
	733	+	NoClientAuth: true,
	734	+	NoClientAuthCallback: func(ssh.ConnMetadata) (*ssh.Permissions, error) {
	735	+	return &ssh.Permissions{Extensions: map[string]string{}}, nil
	736	+	},
	737	+	},
	738	+	})
	739	+	server.Config.AddHostKey(user.signer)
	740	+	go func() { _ = server.ListenAndServe() }()
	741	+
	742	+	var addr string
	743	+	for range 50 {
	744	+	server.Mu.Lock()
	745	+	if server.Listener != nil {
	746	+	addr = server.Listener.Addr().String()
	747	+	}
	748	+	server.Mu.Unlock()
	749	+	if addr != "" {
	750	+	break
	751	+	}
	752	+	time.Sleep(10 * time.Millisecond)
	753	+	}
	754	+	if addr == "" {
	755	+	t.Fatal("server listener not ready")
	756	+	}
	757	+
	758	+	client, err := ssh.Dial("tcp", addr, &ssh.ClientConfig{User: "user", HostKeyCallback: ssh.InsecureIgnoreHostKey()})
	759	+	if err != nil {
	760	+	t.Fatal(err)
	761	+	}
	762	+	defer func() { _ = client.Close() }()
	763	+
	764	+	ptyReq := ssh.Marshal(struct {
	765	+	Term string
	766	+	Columns, Rows uint32
	767	+	Width, Height uint32
	768	+	Modes string
	769	+	}{Term: "xterm", Columns: 80, Rows: 24})
	770	+
	771	+	missed := 0
	772	+	for range 200 {
	773	+	session, err := client.NewSession()
	774	+	if err != nil {
	775	+	t.Fatal(err)
	776	+	}
	777	+	if _, err := session.SendRequest("pty-req", false, ptyReq); err != nil {
	778	+	t.Fatal(err)
	779	+	}
	780	+	out, _ := session.Output("cmd")
	781	+	if strings.Contains(string(out), "no pty") {
	782	+	missed++
	783	+	}
	784	+	_ = session.Close()
	785	+	}
	786	+
	787	+	if missed > 0 {
	788	+	t.Fatalf("command ran before its pty request in %d of 200 sessions", missed)
	789	+	}
	790	+	}

...	...		@@ -5,7 +5,6 @@ import (
5	5		"os"
6	6		"os/signal"
7	7		"syscall"
8		-	"time"
9	8
10	9		"github.com/picosh/pico/pkg/db/postgres"
11	10		"github.com/picosh/pico/pkg/pssh"

...	...		@@ -53,7 +52,7 @@ func StartSshServer() {
53	52		scp.Middleware(handler),
54	53		rsync.Middleware(handler),
55	54		auth.Middleware(handler),
56		-	pssh.PtyMdw(pssh.DeprecatedNotice(), 200*time.Millisecond),
	55	+	pssh.PtyMdw(pssh.DeprecatedNotice()),
57	56		pssh.LogMiddleware(handler, dbh),
58	57		},
59	58		[]pssh.SSHServerMiddleware{

...	...		@@ -106,7 +105,7 @@ func StartSshServer() {
106	105		PgsDB: pgsDB,
107	106		Logger: pssh.GetLogger(sesh),
108	107		}
109		-	return pssh.PtyMdw(createTui(shrd), 200*time.Millisecond)(next)(sesh)
	108	+	return pssh.PtyMdw(createTui(shrd))(next)(sesh)
110	109		}
111	110		},
112	111		Middleware(cliHandler),

...	...		@@ -63,7 +62,7 @@ func StartSshServer() {
63	62		scp.Middleware(handler),
64	63		rsync.Middleware(handler),
65	64		auth.Middleware(handler),
66		-	pssh.PtyMdw(pssh.DeprecatedNotice(), 200*time.Millisecond),
	65	+	pssh.PtyMdw(pssh.DeprecatedNotice()),
67	66		pssh.LogMiddleware(handler, dbh),
68	67		},
69	68		[]pssh.SSHServerMiddleware{

...	...		@@ -3,7 +3,6 @@ package pssh
3	3		import (
4	4		"encoding/binary"
5	5		"fmt"
6		-	"time"
7	6		)
8	7
9	8		func SessionMessage(sesh *SSHServerConnSession, msg string) {

...	...		@@ -24,22 +23,9 @@ func DeprecatedNotice() SSHServerMiddleware {
24	23		}
25	24		}
26	25
27		-	func PtyMdw(mdw SSHServerMiddleware, waitTimeout time.Duration) SSHServerMiddleware {
	26	+	func PtyMdw(mdw SSHServerMiddleware) SSHServerMiddleware {
28	27		return func(next SSHServerHandler) SSHServerHandler {
29	28		return func(sesh *SSHServerConnSession) error {
30		-	if waitTimeout > 0 {
31		-	afterTime := time.Now().Add(waitTimeout)
32		-
33		-	for time.Now().After(afterTime) {
34		-	_, _, ok := sesh.Pty()
35		-	if ok {
36		-	break
37		-	}
38		-
39		-	time.Sleep(50 * time.Millisecond)
40		-	}
41		-	}
42		-
43	29		_, _, ok := sesh.Pty()
44	30		if !ok {
45	31		return next(sesh)

Commit d642d4b

Commit `d642d4b`