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