Commit 0702a1c
Antonio Mika
·
2026-10-10 22:37:18 -0400 EDT
parent da1f3af
fix(pssh): end a session's command when the client closes it
2 files changed,
+72,
-0
+3,
-0
| ... | ... | @@ -490,6 +490,9 @@ func NewSSHServer(ctx context.Context, logger *slog.Logger, config *SSHServerCon | |
| 490 | 490 | } | |
| 491 | 491 | ||
| 492 | 492 | ctx, cancelFunc := context.WithCancel(sc.Ctx) | |
| 493 | + | // The client closing the session ends its command, even while | |
| 494 | + | // the connection stays open for other sessions. | |
| 495 | + | defer cancelFunc() | |
| 493 | 496 | ||
| 494 | 497 | sesh := &SSHServerConnSession{ | |
| 495 | 498 | Channel: channel, |
+69,
-0
| ... | ... | @@ -640,3 +640,72 @@ func GenerateKey() UserSSH { | |
| 640 | 640 | signer: userSigner, | |
| 641 | 641 | } | |
| 642 | 642 | } | |
| 643 | + | ||
| 644 | + | // A client that closes a session but keeps its connection, like pipe-web, | |
| 645 | + | // ends the command running in it. | |
| 646 | + | func TestSessionContextEndsWhenClientCloses(t *testing.T) { | |
| 647 | + | ctx, cancel := context.WithCancel(context.Background()) | |
| 648 | + | defer cancel() | |
| 649 | + | ||
| 650 | + | user := GenerateKey() | |
| 651 | + | started := make(chan struct{}) | |
| 652 | + | ended := make(chan struct{}) | |
| 653 | + | server := pssh.NewSSHServer(ctx, slog.New(slog.DiscardHandler), &pssh.SSHServerConfig{ | |
| 654 | + | ListenAddr: "127.0.0.1:0", | |
| 655 | + | Middleware: []pssh.SSHServerMiddleware{ | |
| 656 | + | func(next pssh.SSHServerHandler) pssh.SSHServerHandler { | |
| 657 | + | return func(sesh *pssh.SSHServerConnSession) error { | |
| 658 | + | close(started) | |
| 659 | + | <-sesh.Context().Done() | |
| 660 | + | close(ended) | |
| 661 | + | return next(sesh) | |
| 662 | + | } | |
| 663 | + | }, | |
| 664 | + | }, | |
| 665 | + | ServerConfig: &ssh.ServerConfig{ | |
| 666 | + | NoClientAuth: true, | |
| 667 | + | NoClientAuthCallback: func(ssh.ConnMetadata) (*ssh.Permissions, error) { | |
| 668 | + | return &ssh.Permissions{Extensions: map[string]string{}}, nil | |
| 669 | + | }, | |
| 670 | + | }, | |
| 671 | + | }) | |
| 672 | + | server.Config.AddHostKey(user.signer) | |
| 673 | + | go func() { _ = server.ListenAndServe() }() | |
| 674 | + | ||
| 675 | + | var addr string | |
| 676 | + | for range 50 { | |
| 677 | + | server.Mu.Lock() | |
| 678 | + | if server.Listener != nil { | |
| 679 | + | addr = server.Listener.Addr().String() | |
| 680 | + | } | |
| 681 | + | server.Mu.Unlock() | |
| 682 | + | if addr != "" { | |
| 683 | + | break | |
| 684 | + | } | |
| 685 | + | time.Sleep(10 * time.Millisecond) | |
| 686 | + | } | |
| 687 | + | if addr == "" { | |
| 688 | + | t.Fatal("server listener not ready") | |
| 689 | + | } | |
| 690 | + | ||
| 691 | + | client, err := ssh.Dial("tcp", addr, &ssh.ClientConfig{User: "user", HostKeyCallback: ssh.InsecureIgnoreHostKey()}) | |
| 692 | + | if err != nil { | |
| 693 | + | t.Fatal(err) | |
| 694 | + | } | |
| 695 | + | defer func() { _ = client.Close() }() | |
| 696 | + | session, err := client.NewSession() | |
| 697 | + | if err != nil { | |
| 698 | + | t.Fatal(err) | |
| 699 | + | } | |
| 700 | + | if err := session.Start("wait"); err != nil { | |
| 701 | + | t.Fatal(err) | |
| 702 | + | } | |
| 703 | + | <-started | |
| 704 | + | _ = session.Close() | |
| 705 | + | ||
| 706 | + | select { | |
| 707 | + | case <-ended: | |
| 708 | + | case <-time.After(2 * time.Second): | |
| 709 | + | t.Fatal("command still running after the client closed its session") | |
| 710 | + | } | |
| 711 | + | } |