repos / pico

ssh.go

 1package pipe
 2
 3import (
 4	"context"
 5	"fmt"
 6	"os"
 7	"os/signal"
 8	"syscall"
 9	"time"
10
11	"github.com/antoniomika/syncmap"
12	"github.com/charmbracelet/promwish"
13	"github.com/charmbracelet/ssh"
14	"github.com/charmbracelet/wish"
15	"github.com/picosh/pico/db/postgres"
16	"github.com/picosh/pico/shared"
17	wsh "github.com/picosh/pico/wish"
18	psub "github.com/picosh/pubsub"
19	"github.com/picosh/utils"
20)
21
22func StartSshServer() {
23	host := utils.GetEnv("PIPE_HOST", "0.0.0.0")
24	port := utils.GetEnv("PIPE_SSH_PORT", "2222")
25	portOverride := utils.GetEnv("PIPE_SSH_PORT_OVERRIDE", port)
26	promPort := utils.GetEnv("PIPE_PROM_PORT", "9222")
27	cfg := NewConfigSite()
28	logger := cfg.Logger
29	dbh := postgres.NewDB(cfg.DbURL, cfg.Logger)
30	defer dbh.Close()
31
32	cfg.Port = port
33	cfg.PortOverride = portOverride
34
35	pubsub := psub.NewMulticast(logger)
36	handler := &CliHandler{
37		Logger:  logger,
38		DBPool:  dbh,
39		PubSub:  pubsub,
40		Cfg:     cfg,
41		Waiters: syncmap.New[string, []string](),
42		Access:  syncmap.New[string, []string](),
43	}
44
45	sshAuth := shared.NewSshAuthHandler(dbh, logger, cfg)
46	s, err := wish.NewServer(
47		wish.WithAddress(fmt.Sprintf("%s:%s", host, port)),
48		wish.WithHostKeyPath("ssh_data/term_info_ed25519"),
49		wish.WithPublicKeyAuth(func(ctx ssh.Context, key ssh.PublicKey) bool {
50			sshAuth.PubkeyAuthHandler(ctx, key)
51			return true
52		}),
53		wish.WithMiddleware(
54			WishMiddleware(handler),
55			promwish.Middleware(fmt.Sprintf("%s:%s", host, promPort), "pipe-ssh"),
56			wsh.LogMiddleware(logger),
57		),
58	)
59	if err != nil {
60		logger.Error("wish server", "err", err.Error())
61		return
62	}
63
64	done := make(chan os.Signal, 1)
65	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
66	logger.Info("Starting SSH server", "host", host, "port", port)
67	go func() {
68		if err = s.ListenAndServe(); err != nil {
69			logger.Error("listen", "err", err.Error())
70		}
71	}()
72
73	<-done
74	logger.Info("Stopping SSH server")
75	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
76	defer func() { cancel() }()
77	if err := s.Shutdown(ctx); err != nil {
78		logger.Error("shutdown", "err", err.Error())
79	}
80}