Merge pull request #26 from uselagoon/handle-close

fix: gracefully handle NATS connection close in ssh-portal
uselagoon · Mar 3, 2022 · e4df813 · e4df813
2 parents 68ae616 + 4a08784
commit e4df813
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/cmd/ssh-portal/serve.go b/cmd/ssh-portal/serve.go
@@ -35,7 +35,11 @@ func (cmd *ServeCmd) Run(log *zap.Logger) error {
 	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM)
 	defer stop()
 	// get nats server connection
-	nc, err := nats.Connect(cmd.NATSServer)
+	nc, err := nats.Connect(cmd.NATSServer,
+		// exit on connection close
+		nats.ClosedHandler(func(_ *nats.Conn) {
+			stop()
+		}))
 	if err != nil {
 		return fmt.Errorf("couldn't connect to NATS server: %v", err)
 	}