prime8dev

prime8dev
git clone https://git.prime8.dev/prime8dev.git
Log | Files | Refs | README | LICENSE

server.go (2544B)

      1 package ssh
      2 
      3 import (
      4 	"context"
      5 	"errors"
      6 	"fmt"
      7 	"log/slog"
      8 	"net"
      9 	"os"
     10 	"strings"
     11 	"time"
     12 
     13 	"git.prime8.dev/moncheeta/prime8dev/internal/tui"
     14 	charmssh "github.com/charmbracelet/ssh"
     15 	"github.com/charmbracelet/wish"
     16 	wishbubbletea "github.com/charmbracelet/wish/bubbletea"
     17 	gossh "golang.org/x/crypto/ssh"
     18 )
     19 
     20 const (
     21 	defaultSSHPort  = "22"
     22 	hostKeyPath     = ".ssh/public_ssh_ed25519"
     23 	idleTimeout     = 90 * time.Second
     24 	maxTimeout      = 10 * time.Minute
     25 	rateLimitWindow = 1 * time.Minute
     26 	rateLimitMax    = 30
     27 	rateLimitBlock  = 5 * time.Minute
     28 	rateLimitMaxIPs = 2048
     29 )
     30 
     31 func Run(ctx context.Context, logger *slog.Logger) error {
     32 	publicAddr := sshListenAddr()
     33 	limiter := newIPLimiter(rateLimitWindow, rateLimitMax, rateLimitBlock, rateLimitMaxIPs)
     34 
     35 	server, err := wish.NewServer(
     36 		wish.WithAddress(publicAddr),
     37 		wish.WithHostKeyPath(hostKeyPath),
     38 		wish.WithIdleTimeout(idleTimeout),
     39 		wish.WithMaxTimeout(maxTimeout),
     40 		wish.WithPublicKeyAuth(func(_ charmssh.Context, _ charmssh.PublicKey) bool {
     41 			return true
     42 		}),
     43 		wish.WithKeyboardInteractiveAuth(func(_ charmssh.Context, _ gossh.KeyboardInteractiveChallenge) bool {
     44 			return true
     45 		}),
     46 		charmssh.WrapConn(func(_ charmssh.Context, conn net.Conn) net.Conn {
     47 			ip := remoteIP(conn.RemoteAddr())
     48 			allowed, retryAfter := limiter.Allow(ip)
     49 			if !allowed {
     50 				logger.Warn("connection rejected by rate limiter",
     51 					"ip", ip,
     52 					"retry_after", retryAfter.String(),
     53 				)
     54 				return nil
     55 			}
     56 			return conn
     57 		}),
     58 		func(s *charmssh.Server) error {
     59 			s.ConnectionFailedCallback = func(conn net.Conn, err error) {
     60 				logger.Warn("connection failed",
     61 					"ip", remoteIP(conn.RemoteAddr()),
     62 					"error", err,
     63 				)
     64 			}
     65 			return nil
     66 		},
     67 		wish.WithMiddleware(
     68 			wishbubbletea.Middleware(tui.SSHHandler),
     69 			routeMiddleware(logger),
     70 		),
     71 	)
     72 	if err != nil {
     73 		return fmt.Errorf("failed to initialize public ssh server: %w", err)
     74 	}
     75 
     76 	go func() {
     77 		<-ctx.Done()
     78 		shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
     79 		defer cancel()
     80 		_ = server.Shutdown(shutdownCtx)
     81 	}()
     82 
     83 	logger.Info("public ssh server listening",
     84 		"addr", publicAddr,
     85 		"host_key", hostKeyPath,
     86 	)
     87 
     88 	err = server.ListenAndServe()
     89 	if err != nil && !errors.Is(err, charmssh.ErrServerClosed) {
     90 		return fmt.Errorf("public ssh server stopped with error: %w", err)
     91 	}
     92 
     93 	logger.Info("public ssh server stopped")
     94 	return nil
     95 }
     96 
     97 func sshListenAddr() string {
     98 	port := strings.TrimSpace(os.Getenv("SSH_PORT"))
     99 	if port == "" {
    100 		port = defaultSSHPort
    101 	}
    102 
    103 	return ":" + port
    104 }