repos / pico

commit

92facd1

parent

e963190

author

Eric Bower

date

2026-02-24 11:35:41 -0500 EST

chore: added test for pssh cmd parsing

  1@@ -2,13 +2,19 @@ package pssh_test
  2 
  3 import (
  4 	"context"
  5+	"crypto/rand"
  6 	"errors"
  7+	"io"
  8 	"log/slog"
  9 	"net"
 10+	"slices"
 11+	"strings"
 12 	"testing"
 13 	"time"
 14 
 15 	"github.com/picosh/pico/pkg/pssh"
 16+	"github.com/picosh/pico/pkg/shared"
 17+	"golang.org/x/crypto/ed25519"
 18 	"golang.org/x/crypto/ssh"
 19 )
 20 
 21@@ -275,3 +281,174 @@ func TestSSHServerConnHandle(t *testing.T) {
 22 		t.Error("Handle did not return after context canceled")
 23 	}
 24 }
 25+
 26+func TestSSHServerCommandParsing(t *testing.T) {
 27+	ctx, cancel := context.WithCancel(context.Background())
 28+	defer cancel()
 29+
 30+	logger := slog.Default()
 31+	var capturedCommand []string
 32+
 33+	user := GenerateKey()
 34+
 35+	server := pssh.NewSSHServer(ctx, logger, &pssh.SSHServerConfig{
 36+		ListenAddr: "localhost:2222",
 37+		Middleware: []pssh.SSHServerMiddleware{
 38+			func(next pssh.SSHServerHandler) pssh.SSHServerHandler {
 39+				return func(sesh *pssh.SSHServerConnSession) error {
 40+					capturedCommand = sesh.Command()
 41+					return next(sesh)
 42+				}
 43+			},
 44+		},
 45+		ServerConfig: &ssh.ServerConfig{
 46+			NoClientAuth: true,
 47+			NoClientAuthCallback: func(ssh.ConnMetadata) (*ssh.Permissions, error) {
 48+				return &ssh.Permissions{
 49+					Extensions: map[string]string{
 50+						"pubkey": shared.KeyForKeyText(user.signer.PublicKey()),
 51+					},
 52+				}, nil
 53+			},
 54+		},
 55+	})
 56+	server.Config.AddHostKey(user.signer)
 57+
 58+	// Start server in a goroutine
 59+	errChan := make(chan error, 1)
 60+	go func() {
 61+		err := server.ListenAndServe()
 62+		errChan <- err
 63+	}()
 64+
 65+	// Wait a bit for the server to start
 66+	time.Sleep(100 * time.Millisecond)
 67+
 68+	// Send command to server
 69+	user.MustCmd(nil, "accept --comment 'here we go' 101")
 70+
 71+	time.Sleep(1000 * time.Millisecond)
 72+
 73+	expectedCommand := []string{"accept", "--comment", "'here we go'", "101"}
 74+	if !slices.Equal(expectedCommand, capturedCommand) {
 75+		t.Error("command not exected", capturedCommand, len(capturedCommand), expectedCommand, len(expectedCommand))
 76+	}
 77+
 78+	// Trigger cancellation to stop the server
 79+	cancel()
 80+
 81+	// Wait for server to stop
 82+	select {
 83+	case err := <-errChan:
 84+		if err != nil && !errors.Is(err, net.ErrClosed) {
 85+			t.Errorf("unexpected error: %v", err)
 86+		}
 87+	case <-time.After(2 * time.Second):
 88+		t.Error("server did not shut down in time")
 89+	}
 90+}
 91+
 92+type UserSSH struct {
 93+	username string
 94+	signer   ssh.Signer
 95+}
 96+
 97+func NewUserSSH(username string, signer ssh.Signer) *UserSSH {
 98+	return &UserSSH{
 99+		username: username,
100+		signer:   signer,
101+	}
102+}
103+
104+func (s UserSSH) Public() string {
105+	pubkey := s.signer.PublicKey()
106+	return string(ssh.MarshalAuthorizedKey(pubkey))
107+}
108+
109+func (s UserSSH) MustCmd(patch []byte, cmd string) string {
110+	res, err := s.Cmd(patch, cmd)
111+	if err != nil {
112+		panic(err)
113+	}
114+	return res
115+}
116+
117+func (s UserSSH) Cmd(patch []byte, cmd string) (string, error) {
118+	host := "localhost:2222"
119+
120+	config := &ssh.ClientConfig{
121+		User: s.username,
122+		Auth: []ssh.AuthMethod{
123+			ssh.PublicKeys(s.signer),
124+		},
125+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
126+	}
127+
128+	client, err := ssh.Dial("tcp", host, config)
129+	if err != nil {
130+		return "", err
131+	}
132+	defer func() {
133+		_ = client.Close()
134+	}()
135+
136+	session, err := client.NewSession()
137+	if err != nil {
138+		return "", err
139+	}
140+	defer func() {
141+		_ = session.Close()
142+	}()
143+
144+	stdinPipe, err := session.StdinPipe()
145+	if err != nil {
146+		return "", err
147+	}
148+
149+	stdoutPipe, err := session.StdoutPipe()
150+	if err != nil {
151+		return "", err
152+	}
153+
154+	if err := session.Start(cmd); err != nil {
155+		return "", err
156+	}
157+
158+	if patch != nil {
159+		_, err = stdinPipe.Write(patch)
160+		if err != nil {
161+			return "", err
162+		}
163+	}
164+
165+	_ = stdinPipe.Close()
166+
167+	if err := session.Wait(); err != nil {
168+		return "", err
169+	}
170+
171+	buf := new(strings.Builder)
172+	_, err = io.Copy(buf, stdoutPipe)
173+	if err != nil {
174+		return "", err
175+	}
176+
177+	return buf.String(), nil
178+}
179+
180+func GenerateKey() UserSSH {
181+	_, userKey, err := ed25519.GenerateKey(rand.Reader)
182+	if err != nil {
183+		panic(err)
184+	}
185+
186+	userSigner, err := ssh.NewSignerFromKey(userKey)
187+	if err != nil {
188+		panic(err)
189+	}
190+
191+	return UserSSH{
192+		username: "user",
193+		signer:   userSigner,
194+	}
195+}