feat: add connection limit (#3)

Added a connection limit so only one user can connect to the flipper at
a time. Allowing multiple connections simultaneously doesn't make much sense for
this tool.

middleware.go

@@ -0,0 +1,53 @@
+package main
+
+import (
+	"errors"
+	"sync"
+
+	"github.com/charmbracelet/wish"
+	"github.com/gliderlabs/ssh"
+)
+
+type connLimiter struct {
+	sync.Mutex
+	conns    int
+	maxConns int
+}
+
+func newConnLimiter(maxConns int) *connLimiter {
+	return &connLimiter{
+		maxConns: maxConns,
+	}
+}
+
+func (u *connLimiter) Add() error {
+	u.Lock()
+	defer u.Unlock()
+	if u.conns >= u.maxConns {
+		return errors.New("max connections reached")
+	}
+	u.conns++
+	return nil
+}
+
+func (u *connLimiter) Remove() {
+	u.Lock()
+	defer u.Unlock()
+	u.conns--
+	if u.conns < 0 {
+		u.conns = 0
+	}
+}
+
+func connLimit(limiter *connLimiter) wish.Middleware {
+	return func(sh ssh.Handler) ssh.Handler {
+		return func(s ssh.Session) {
+			if err := limiter.Add(); err != nil {
+				wish.Fatalf(s, "max connections reached\n")
+				return
+			}
+			sh(s)
+			limiter.Remove()
+		}
+	}
+}

server.go

@@ -38,6 +38,9 @@ func server(cmd *coral.Command, args []string) {
 	if err != nil {
 		log.Fatal(err)
 	}
+
+	cl := newConnLimiter(1)
+
 	s, err := wish.NewServer(
 		wish.WithAddress(serverFlags.listen),
 		wish.WithHostKeyPath(".ssh/flipperzero_tea_ed25519"),
@@ -54,6 +57,7 @@ func server(cmd *coral.Command, args []string) {
 				return m, []tea.ProgramOption{tea.WithAltScreen(), tea.WithMouseCellMotion()}
 			}),
 			lm.Middleware(),
+			connLimit(cl),
 		),
 	)
 	if err != nil {