diff --git a/pkgs/chat/main.go b/pkgs/chat/main.go index 67fb04912..696860841 100644 --- a/pkgs/chat/main.go +++ b/pkgs/chat/main.go @@ -2,6 +2,7 @@ package main import ( "context" + "errors" "fmt" "log" "os" @@ -59,22 +60,40 @@ func newApp() *app { return a } +func runServer(serve func() error, shutdown func(context.Context) error, done <-chan os.Signal, timeout time.Duration) error { + serveErr := make(chan error, 1) + go func() { + serveErr <- serve() + }() + + select { + case err := <-serveErr: + if err != nil && !errors.Is(err, ssh.ErrServerClosed) { + return err + } + return nil + case <-done: + log.Println("Stopping SSH server") + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + if err := shutdown(ctx); err != nil { + return err + } + + err := <-serveErr + if err != nil && !errors.Is(err, ssh.ErrServerClosed) { + return err + } + return nil + } +} + func (a *app) Start() { - var err error done := make(chan os.Signal, 1) signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM) + defer signal.Stop(done) log.Printf("Starting SSH server on %s:%d", host, port) - go func() { - if err = a.ListenAndServe(); err != nil { - log.Fatalln(err) - } - }() - - <-done - log.Println("Stopping SSH server") - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer func() { cancel() }() - if err := a.Shutdown(ctx); err != nil { + if err := runServer(a.ListenAndServe, a.Shutdown, done, 30*time.Second); err != nil { log.Fatalln(err) } } diff --git a/pkgs/chat/main_test.go b/pkgs/chat/main_test.go new file mode 100644 index 000000000..d75ae46ab --- /dev/null +++ b/pkgs/chat/main_test.go @@ -0,0 +1,60 @@ +package main + +import ( + "context" + "errors" + "os" + "testing" + "time" + + "github.com/charmbracelet/ssh" +) + +func TestRunServerTreatsShutdownAsCleanStop(t *testing.T) { + signals := make(chan os.Signal, 1) + signals <- os.Interrupt + + serveRelease := make(chan struct{}) + shutdownCalled := false + err := runServer( + func() error { + <-serveRelease + return ssh.ErrServerClosed + }, + func(context.Context) error { + shutdownCalled = true + close(serveRelease) + return nil + }, + signals, + time.Second, + ) + + if err != nil { + t.Fatalf("runServer() error = %v, want nil", err) + } + if !shutdownCalled { + t.Fatal("runServer() did not shut down after receiving a signal") + } +} + +func TestRunServerReturnsServeFailure(t *testing.T) { + serveErr := errors.New("listen tcp 127.0.0.1:23234: bind: address already in use") + shutdownCalled := false + err := runServer( + func() error { return serveErr }, + func(context.Context) error { + shutdownCalled = true + return nil + }, + make(chan os.Signal), + time.Second, + ) + + if !errors.Is(err, serveErr) { + t.Fatalf("runServer() error = %v, want %v", err, serveErr) + } + if shutdownCalled { + t.Fatal("runServer() called shutdown after a serve failure") + } +}