Antonio Mika
·
2026-10-25
1package pssh
2
3import (
4 "bytes"
5 "context"
6 "crypto/ed25519"
7 "crypto/rand"
8 "crypto/subtle"
9 "encoding/base64"
10 "encoding/pem"
11 "errors"
12 "fmt"
13 "io"
14 "log/slog"
15 "net"
16 "net/http"
17 "os"
18 "path"
19 "strings"
20 "sync"
21 "time"
22 "unicode/utf8"
23
24 "github.com/antoniomika/syncmap"
25 "github.com/go-andiamo/splitter"
26 "github.com/picosh/pico/pkg/metrics"
27 "github.com/prometheus/client_golang/prometheus"
28 "github.com/prometheus/client_golang/prometheus/promauto"
29 "github.com/prometheus/client_golang/prometheus/promhttp"
30 "golang.org/x/crypto/ssh"
31)
32
33type SSHServerConn struct {
34 Ctx context.Context
35 CancelFunc context.CancelFunc
36 Logger *slog.Logger
37 Conn *ssh.ServerConn
38 SSHServer *SSHServer
39 Start time.Time
40
41 mu sync.Mutex
42}
43
44func (s *SSHServerConn) Context() context.Context {
45 s.mu.Lock()
46 defer s.mu.Unlock()
47
48 return s.Ctx
49}
50
51func (sc *SSHServerConn) Close() error {
52 sc.CancelFunc()
53 return nil
54}
55
56type SSHServerConnSession struct {
57 ssh.Channel
58 *SSHServerConn
59
60 Ctx context.Context
61 CancelFunc context.CancelFunc
62
63 pty *Pty
64 winch chan Window
65
66 mu sync.Mutex
67}
68
69// Deadline implements context.Context.
70func (s *SSHServerConnSession) Deadline() (deadline time.Time, ok bool) {
71 s.mu.Lock()
72 defer s.mu.Unlock()
73
74 return s.Ctx.Deadline()
75}
76
77// Done implements context.Context.
78func (s *SSHServerConnSession) Done() <-chan struct{} {
79 s.mu.Lock()
80 defer s.mu.Unlock()
81
82 return s.Ctx.Done()
83}
84
85// Err implements context.Context.
86func (s *SSHServerConnSession) Err() error {
87 s.mu.Lock()
88 defer s.mu.Unlock()
89
90 return s.Ctx.Err()
91}
92
93// Value implements context.Context.
94func (s *SSHServerConnSession) Value(key any) any {
95 s.mu.Lock()
96 defer s.mu.Unlock()
97
98 return s.Ctx.Value(key)
99}
100
101// SetValue implements context.Context.
102func (s *SSHServerConnSession) SetValue(key any, data any) {
103 s.mu.Lock()
104 defer s.mu.Unlock()
105
106 s.Ctx = context.WithValue(s.Ctx, key, data)
107}
108
109func (s *SSHServerConnSession) Context() context.Context {
110 s.mu.Lock()
111 defer s.mu.Unlock()
112
113 return s.Ctx
114}
115
116func (s *SSHServerConnSession) Permissions() *ssh.Permissions {
117 return s.Conn.Permissions
118}
119
120func (s *SSHServerConnSession) User() string {
121 return s.Conn.User()
122}
123
124func (s *SSHServerConnSession) PublicKey() ssh.PublicKey {
125 key, ok := s.Conn.Permissions.Extensions["pubkey"]
126 if !ok {
127 return nil
128 }
129
130 pk, _, _, _, err := ssh.ParseAuthorizedKey([]byte(key))
131 if err != nil {
132 return nil
133 }
134 return pk
135}
136
137func (s *SSHServerConnSession) RemoteAddr() net.Addr {
138 return s.Conn.RemoteAddr()
139}
140
141func (s *SSHServerConnSession) Command() []string {
142 cmd, _ := s.Value("command").([]string)
143 return cmd
144}
145
146func (s *SSHServerConnSession) Close() error {
147 s.CancelFunc()
148 return s.Channel.Close()
149}
150
151func (s *SSHServerConnSession) Exit(code int) error {
152 status := struct{ Status uint32 }{uint32(code)}
153 _, err := s.SendRequest("exit-status", false, ssh.Marshal(&status))
154 return err
155}
156
157func (sesh *SSHServerConnSession) Errorf(f string, v ...interface{}) {
158 _, _ = fmt.Fprintf(sesh.Stderr(), f, v...)
159}
160
161func (sesh *SSHServerConnSession) Errorln(v ...interface{}) {
162 _, _ = fmt.Fprintln(sesh.Stderr(), v...)
163}
164
165func (sesh *SSHServerConnSession) Printf(f string, v ...interface{}) {
166 _, _ = fmt.Fprintf(sesh, f, v...)
167}
168
169func (sesh *SSHServerConnSession) Println(v ...interface{}) {
170 _, _ = fmt.Fprintln(sesh, v...)
171}
172
173func (s *SSHServerConnSession) Fatal(err error) {
174 _, _ = fmt.Fprintln(s.Stderr(), err)
175 _, _ = fmt.Fprintf(s.Stderr(), "\r")
176 _ = s.Exit(1)
177 _ = s.Close()
178}
179
180func (s *SSHServerConnSession) Pty() (*Pty, <-chan Window, bool) {
181 s.mu.Lock()
182 defer s.mu.Unlock()
183
184 if s.pty == nil {
185 return nil, nil, false
186 }
187
188 return s.pty, s.winch, true
189}
190
191func (s *SSHServerConnSession) handlePtyReq(req *ssh.Request) {
192 ptyReq, ok := parsePtyRequest(req.Payload)
193
194 s.mu.Lock()
195 if ok && s.pty == nil {
196 s.pty = &ptyReq
197 s.winch = make(chan Window, 1)
198 s.winch <- ptyReq.Window
199 } else {
200 ok = false
201 }
202 s.mu.Unlock()
203
204 if err := req.Reply(ok, nil); err != nil {
205 s.Logger.Error("pty-req reply", "err", err)
206 }
207}
208
209// Write overrides the embedded Channel's Write to normalize line endings when PTY is allocated.
210func (s *SSHServerConnSession) Write(p []byte) (n int, err error) {
211 s.mu.Lock()
212 hasPty := s.pty != nil
213 s.mu.Unlock()
214
215 if !hasPty {
216 // No PTY, write as-is
217 return s.Channel.Write(p)
218 }
219
220 // When PTY is active, ensure every \n is preceded by \r.
221 // This ensures the cursor returns to column 0 before the newline.
222 var buf bytes.Buffer
223 for i := 0; i < len(p); i++ {
224 if p[i] == '\n' {
225 // Check if preceded by \r
226 if i == 0 || p[i-1] != '\r' {
227 buf.WriteByte('\r')
228 }
229 buf.WriteByte('\n')
230 } else {
231 buf.WriteByte(p[i])
232 }
233 }
234
235 normalized := buf.Bytes()
236 written, err := s.Channel.Write(normalized)
237
238 // Return the count based on original data length, not normalized
239 if written > len(p) {
240 written = len(p)
241 }
242 return written, err
243}
244
245var _ context.Context = &SSHServerConnSession{}
246
247func (sc *SSHServerConn) Handle(chans <-chan ssh.NewChannel, reqs <-chan *ssh.Request) error {
248 defer func() {
249 _ = sc.Close()
250 }()
251
252 for {
253 select {
254 case <-sc.Context().Done():
255 return nil
256 case newChan, ok := <-chans:
257 if !ok {
258 return nil
259 }
260
261 sc.Logger.Info("new channel", "type", newChan.ChannelType(), "extraData", newChan.ExtraData())
262 chanFunc, ok := sc.SSHServer.Config.ChannelMiddleware[newChan.ChannelType()]
263 if !ok {
264 sc.Logger.Info("no channel middleware for type", "type", newChan.ChannelType())
265 continue
266 }
267
268 go func() {
269 err := chanFunc(newChan, sc)
270 if err != nil {
271 sc.Logger.Error("channel middleware", "err", err)
272 }
273 }()
274 case req, ok := <-reqs:
275 if !ok {
276 return nil
277 }
278 sc.Logger.Info("new request", "type", req.Type, "wantReply", req.WantReply, "payload", req.Payload)
279 switch req.Type {
280 case "keepalive@openssh.com":
281 sc.Logger.Info("keepalive reply")
282 err := req.Reply(true, nil)
283 if err != nil {
284 sc.Logger.Error("keepalive", "err", err)
285 }
286 }
287 }
288 }
289}
290
291func NewSSHServerConn(
292 ctx context.Context,
293 logger *slog.Logger,
294 conn *ssh.ServerConn,
295 server *SSHServer,
296) *SSHServerConn {
297 if ctx == nil {
298 ctx = context.Background()
299 }
300
301 cancelCtx, cancelFunc := context.WithCancel(ctx)
302
303 if logger == nil {
304 logger = slog.Default()
305 }
306
307 return &SSHServerConn{
308 Ctx: cancelCtx,
309 CancelFunc: cancelFunc,
310 Logger: logger,
311 Conn: conn,
312 SSHServer: server,
313 Start: time.Now(),
314 }
315}
316
317type SSHServerHandler func(*SSHServerConnSession) error
318type SSHServerMiddleware func(SSHServerHandler) SSHServerHandler
319type SSHServerChannelMiddleware func(ssh.NewChannel, *SSHServerConn) error
320
321type SSHServerConfig struct {
322 *ssh.ServerConfig
323 App string
324 ListenAddr string
325 PromListenAddr string
326 Middleware []SSHServerMiddleware
327 SubsystemMiddleware []SSHServerMiddleware
328 ChannelMiddleware map[string]SSHServerChannelMiddleware
329}
330
331type SSHServer struct {
332 Ctx context.Context
333 CancelFunc context.CancelFunc
334 Logger *slog.Logger
335 Config *SSHServerConfig
336 Listener net.Listener
337 Conns *syncmap.Map[string, *SSHServerConn]
338
339 SessionsCreated *prometheus.CounterVec
340 SessionsFinished *prometheus.CounterVec
341 SessionsDuration *prometheus.CounterVec
342
343 Mu sync.Mutex
344}
345
346func (s *SSHServer) ListenAndServe() error {
347 if s.Config.PromListenAddr != "" {
348 s.SessionsCreated = promauto.With(prometheus.DefaultRegisterer).NewCounterVec(prometheus.CounterOpts{
349 Name: "pssh_sessions_created_total",
350 Help: "The total number of sessions created",
351 ConstLabels: prometheus.Labels{
352 "app": s.Config.App,
353 },
354 }, []string{"command", "user"})
355
356 s.SessionsFinished = promauto.With(prometheus.DefaultRegisterer).NewCounterVec(prometheus.CounterOpts{
357 Name: "pssh_sessions_finished_total",
358 Help: "The total number of sessions created",
359 ConstLabels: prometheus.Labels{
360 "app": s.Config.App,
361 },
362 }, []string{"command", "user"})
363
364 s.SessionsDuration = promauto.With(prometheus.DefaultRegisterer).NewCounterVec(prometheus.CounterOpts{
365 Name: "pssh_sessions_duration_seconds",
366 Help: "The total sessions duration in seconds",
367 ConstLabels: prometheus.Labels{
368 "app": s.Config.App,
369 },
370 }, []string{"command", "user"})
371
372 metrics.RegisterServerInfo(prometheus.DefaultRegisterer, s.Config.App)
373
374 // The handler reads the default registerer, so build it before the
375 // listener is published rather than in the goroutine.
376 mux := http.NewServeMux()
377 mux.Handle("/metrics", promhttp.Handler())
378
379 go func() {
380 srv := &http.Server{Addr: s.Config.PromListenAddr, Handler: mux}
381
382 go func() {
383 <-s.Ctx.Done()
384 s.Logger.Info("Prometheus server shutting down")
385 _ = srv.Close()
386 }()
387
388 s.Logger.Info("Starting Prometheus server", "addr", s.Config.PromListenAddr)
389
390 err := srv.ListenAndServe()
391 if err != nil {
392 if errors.Is(err, http.ErrServerClosed) {
393 s.Logger.Info("Prometheus server shut down")
394 return
395 }
396
397 s.Logger.Error("Prometheus serve error", "err", err)
398 panic(err)
399 }
400 }()
401 }
402
403 listen, err := net.Listen("tcp", s.Config.ListenAddr)
404 if err != nil {
405 return err
406 }
407
408 s.Mu.Lock()
409 s.Listener = listen
410 s.Mu.Unlock()
411 defer func() {
412 _ = s.Listener.Close()
413 }()
414
415 go func() {
416 <-s.Ctx.Done()
417 _ = s.Close()
418 }()
419
420 var retErr error
421
422 for {
423 conn, err := s.Listener.Accept()
424 if err != nil {
425 s.Logger.Error("accept", "err", err)
426 if errors.Is(err, net.ErrClosed) {
427 retErr = err
428 break
429 }
430 continue
431 }
432
433 go func() {
434 if err := s.HandleConn(conn); err != nil && !errors.Is(err, io.EOF) {
435 s.Logger.Error("Error handling connection", "err", err, "remoteAddr", conn.RemoteAddr().String())
436 }
437 }()
438 }
439
440 if errors.Is(retErr, net.ErrClosed) {
441 return nil
442 }
443
444 return retErr
445}
446
447func (s *SSHServer) HandleConn(conn net.Conn) error {
448 defer func() {
449 _ = conn.Close()
450 }()
451
452 sshConn, chans, reqs, err := ssh.NewServerConn(conn, s.Config.ServerConfig)
453 if err != nil {
454 return err
455 }
456
457 newLogger := s.Logger.With(
458 "remoteAddr", conn.RemoteAddr().String(),
459 "sshUser", sshConn.User(),
460 )
461
462 if pubKey, ok := sshConn.Permissions.Extensions["pubkey"]; ok {
463 newLogger = newLogger.With("pubkey", pubKey)
464 }
465
466 newConn := NewSSHServerConn(
467 s.Ctx,
468 newLogger,
469 sshConn,
470 s,
471 )
472
473 s.Conns.Store(sshConn.RemoteAddr().String(), newConn)
474
475 err = newConn.Handle(chans, reqs)
476
477 s.Conns.Delete(sshConn.RemoteAddr().String())
478
479 return err
480}
481
482func (s *SSHServer) Close() error {
483 s.CancelFunc()
484 return s.Listener.Close()
485}
486
487func NewSSHServer(ctx context.Context, logger *slog.Logger, config *SSHServerConfig) *SSHServer {
488 if ctx == nil {
489 ctx = context.Background()
490 }
491
492 cancelCtx, cancelFunc := context.WithCancel(ctx)
493
494 if logger == nil {
495 logger = slog.Default()
496 }
497
498 if config == nil {
499 config = &SSHServerConfig{}
500 }
501
502 if config.ChannelMiddleware == nil {
503 config.ChannelMiddleware = map[string]SSHServerChannelMiddleware{}
504 }
505
506 if _, ok := config.ChannelMiddleware["session"]; !ok {
507 config.ChannelMiddleware["session"] = func(newChan ssh.NewChannel, sc *SSHServerConn) error {
508 channel, requests, err := newChan.Accept()
509 if err != nil {
510 sc.Logger.Error("accept session channel", "err", err)
511 return err
512 }
513
514 ctx, cancelFunc := context.WithCancel(sc.Ctx)
515 // The client closing the session ends its command, even while
516 // the connection stays open for other sessions.
517 defer cancelFunc()
518
519 sesh := &SSHServerConnSession{
520 Channel: channel,
521 SSHServerConn: sc,
522 Ctx: ctx,
523 CancelFunc: cancelFunc,
524 }
525
526 for {
527 select {
528 case <-sesh.Done():
529 return nil
530 case req, ok := <-requests:
531 if !ok {
532 return nil
533 }
534
535 // Clients send pty-req before shell or exec without waiting for the
536 // reply, so it's applied here, before the next request is read.
537 if req.Type == "pty-req" {
538 sc.Logger.Info("new session request", "type", req.Type, "wantReply", req.WantReply, "payload", req.Payload)
539 sesh.handlePtyReq(req)
540 continue
541 }
542
543 go func() {
544 sc.Logger.Info("new session request", "type", req.Type, "wantReply", req.WantReply, "payload", req.Payload)
545 switch req.Type {
546 case "subsystem":
547 if len(sc.SSHServer.Config.SubsystemMiddleware) == 0 {
548 err := req.Reply(false, nil)
549 if err != nil {
550 sc.Logger.Error("subsystem reply", "err", err)
551 }
552
553 err = sc.Close()
554 if err != nil {
555 sc.Logger.Error("subsystem close", "err", err)
556 }
557
558 sesh.Fatal(err)
559 return
560 }
561
562 h := func(*SSHServerConnSession) error { return nil }
563 for _, m := range sc.SSHServer.Config.SubsystemMiddleware {
564 h = m(h)
565 }
566
567 err := req.Reply(true, nil)
568 if err != nil {
569 sc.Logger.Error("subsystem reply", "err", err)
570 sesh.Fatal(err)
571 return
572 }
573
574 if err := h(sesh); err != nil && !errors.Is(err, io.EOF) {
575 sc.Logger.Error("subsystem middleware", "err", err)
576 sesh.Fatal(err)
577 return
578 }
579
580 err = sesh.Exit(0)
581 if err != nil {
582 sc.Logger.Error("subsystem exit", "err", err)
583 }
584
585 err = sesh.Close()
586 if err != nil {
587 sc.Logger.Error("subsystem close", "err", err)
588 }
589 case "shell", "exec":
590 if len(sc.SSHServer.Config.Middleware) == 0 {
591 err := req.Reply(false, nil)
592 if err != nil {
593 sc.Logger.Error("shell/exec reply", "err", err)
594 }
595 sesh.Fatal(err)
596 return
597 }
598
599 command := "shell"
600
601 if len(req.Payload) > 0 {
602 var payload = struct{ Value string }{}
603 err := ssh.Unmarshal(req.Payload, &payload)
604 if err != nil {
605 sc.Logger.Error("shell/exec unmarshal", "err", err)
606 sesh.Fatal(err)
607 return
608 }
609
610 commaSplitter, _ := splitter.NewSplitter(
611 ' ',
612 splitter.DoubleQuotes,
613 splitter.SingleQuotes,
614 )
615 command = payload.Value
616 cmdSlice, _ := commaSplitter.Split(command, splitter.StripQuotes)
617 sesh.SetValue("command", cmdSlice)
618 }
619
620 if !utf8.ValidString(command) {
621 command = base64.StdEncoding.EncodeToString([]byte(command))
622 }
623
624 if sc.SSHServer.Config.PromListenAddr != "" {
625 labels := []string{sessionCommand(command), sessionUser(sesh)}
626 sc.SSHServer.SessionsCreated.WithLabelValues(labels...).Inc()
627 defer func() {
628 sc.SSHServer.SessionsFinished.WithLabelValues(labels...).Inc()
629 sc.SSHServer.SessionsDuration.WithLabelValues(labels...).Add(time.Since(sc.Start).Seconds())
630 }()
631 }
632
633 h := func(*SSHServerConnSession) error { return nil }
634 for _, m := range sc.SSHServer.Config.Middleware {
635 h = m(h)
636 }
637
638 err = req.Reply(true, nil)
639 if err != nil {
640 sc.Logger.Error("shell/exec reply", "err", err)
641 sesh.Fatal(err)
642 return
643 }
644
645 if err := h(sesh); err != nil && !errors.Is(err, io.EOF) {
646 sc.Logger.Error("exec middleware", "err", err)
647 sesh.Fatal(err)
648 return
649 }
650
651 err = sesh.Exit(0)
652 if err != nil {
653 sc.Logger.Error("subsystem exit", "err", err)
654 }
655
656 err = sesh.Close()
657 if err != nil {
658 sc.Logger.Error("subsystem close", "err", err)
659 }
660 case "window-change":
661 sesh.mu.Lock()
662 found := sesh.pty != nil
663 sesh.mu.Unlock()
664
665 if !found {
666 err := req.Reply(false, nil)
667 if err != nil {
668 sc.Logger.Error("pty-req reply", "err", err)
669 }
670 return
671 }
672
673 win, ok := parseWinchRequest(req.Payload)
674 if ok {
675 sesh.mu.Lock()
676 sesh.pty.Window = win
677 sesh.winch <- win
678 sesh.mu.Unlock()
679 }
680
681 err := req.Reply(ok, nil)
682 if err != nil {
683 sc.Logger.Error("window-change reply", "err", err)
684 }
685 }
686 }()
687 }
688 }
689 }
690 }
691
692 server := &SSHServer{
693 Ctx: cancelCtx,
694 CancelFunc: cancelFunc,
695 Logger: logger,
696 Config: config,
697 Conns: syncmap.New[string, *SSHServerConn](),
698 }
699
700 return server
701}
702
703// sessionCommand is the subcommand alone, since arguments would add a series
704// for every distinct command line.
705func sessionCommand(command string) string {
706 if fields := strings.Fields(command); len(fields) > 0 {
707 return fields[0]
708 }
709 return command
710}
711
712// sessionUser is the name the auth handler confirmed, or "" without one.
713func sessionUser(s *SSHServerConnSession) string {
714 if perms := s.Permissions(); perms != nil {
715 return perms.Extensions["user"]
716 }
717 return ""
718}
719
720type PubKeyAuthHandler func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error)
721
722func NewSSHServerWithConfig(
723 ctx context.Context,
724 logger *slog.Logger,
725 app, host, port, promPort, hostKey string,
726 pubKeyAuthHandler PubKeyAuthHandler,
727 middleware, subsystemMiddleware []SSHServerMiddleware,
728 channelMiddleware map[string]SSHServerChannelMiddleware,
729) (*SSHServer, error) {
730 server := NewSSHServer(ctx, logger, &SSHServerConfig{
731 App: app,
732 ListenAddr: fmt.Sprintf("%s:%s", host, port),
733 ServerConfig: &ssh.ServerConfig{
734 PublicKeyCallback: pubKeyAuthHandler,
735 },
736 Middleware: middleware,
737 SubsystemMiddleware: subsystemMiddleware,
738 ChannelMiddleware: channelMiddleware,
739 })
740
741 if promPort != "" {
742 server.Config.PromListenAddr = fmt.Sprintf("%s:%s", host, promPort)
743 }
744
745 if hostKey != "" {
746 pemBytes, err := os.ReadFile(hostKey)
747 if err != nil {
748 logger.Error("failed to read private key file", "error", err)
749 if !os.IsNotExist(err) {
750 return nil, err
751 }
752
753 logger.Info("generating new private key")
754
755 pubKey, privKey, err := ed25519.GenerateKey(rand.Reader)
756 if err != nil {
757 logger.Error("failed to generate private key", "error", err)
758 return nil, err
759 }
760
761 privb, err := ssh.MarshalPrivateKey(privKey, "")
762 if err != nil {
763 logger.Error("failed to marshal private key", "error", err)
764 return nil, err
765 }
766
767 block := &pem.Block{
768 Type: "OPENSSH PRIVATE KEY",
769 Bytes: privb.Bytes,
770 }
771
772 if err = os.MkdirAll(path.Dir(hostKey), 0700); err != nil {
773 logger.Error("failed to create ssh_data directory", "error", err)
774 return nil, err
775 }
776
777 pemBytes = pem.EncodeToMemory(block)
778
779 if err = os.WriteFile(hostKey, pemBytes, 0600); err != nil {
780 logger.Error("failed to write private key", "error", err)
781 return nil, err
782 }
783
784 sshPubKey, err := ssh.NewPublicKey(pubKey)
785 if err != nil {
786 logger.Error("failed to create public key", "error", err)
787 return nil, err
788 }
789
790 pubb := ssh.MarshalAuthorizedKey(sshPubKey)
791 if err = os.WriteFile(fmt.Sprintf("%s.pub", hostKey), pubb, 0600); err != nil {
792 logger.Error("failed to write public key", "error", err)
793 return nil, err
794 }
795 }
796
797 signer, err := ssh.ParsePrivateKey(pemBytes)
798 if err != nil {
799 logger.Error("failed to parse private key", "error", err)
800 return nil, err
801 }
802
803 server.Config.AddHostKey(signer)
804 }
805
806 return server, nil
807}
808
809func KeysEqual(a, b ssh.PublicKey) bool {
810 if a == nil || b == nil {
811 return false
812 }
813
814 am := a.Marshal()
815 bm := b.Marshal()
816 return (len(am) == len(bm) && subtle.ConstantTimeCompare(am, bm) == 1)
817}