Don't error in prometheus-listening on termination

To tell downstream routines that we're terminating, cancel the Run
context; that eliminates an argument in the Run method, centralizes
signal handling in main.go, and generally looks nicer.

Author: antifuchs

cmd/hoopsnake/main.go

@@ -16,11 +16,18 @@ func main() {
 		log.Fatalf("Invalid command line: %v", err)
 	}
 
+	ctx := context.Background()
+	ctx, terminate := context.WithCancel(ctx)
+
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
+	go func() {
+		signal := <-c
+		log.Printf("Received signal %v, terminating...", signal)
+		terminate()
+	}()
 
-	ctx := context.Background()
-	err = cli.Run(ctx, c)
+	err = cli.Run(ctx)
 	if err != nil {
 		log.Fatal(err)
 	}

prometheus.go

@@ -1,6 +1,7 @@
 package hoopsnake
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"net/http"
@@ -13,7 +14,7 @@ import (
 	"tailscale.com/tsnet"
 )
 
-func (s *TailnetSSH) setupPrometheus(srv *tsnet.Server) error {
+func (s *TailnetSSH) setupPrometheus(ctx context.Context, srv *tsnet.Server) error {
 	if s.prometheusAddr == "" {
 		return nil
 	}
@@ -28,8 +29,11 @@ func (s *TailnetSSH) setupPrometheus(srv *tsnet.Server) error {
 			Handler:           mux,
 			ReadHeaderTimeout: 1 * time.Second,
 		}
-		log.Printf("Failed to listen on prometheus address: %v", server.Serve(listener))
-		os.Exit(20)
+		if ctx.Err() == nil {
+			// Failed to listen but not asked to shut down:
+			log.Printf("Failed to listen on prometheus address: %v", server.Serve(listener))
+			os.Exit(20)
+		}
 	}()
 
 	v4, v6 := srv.TailscaleIPs()

ssh.go

@@ -75,7 +75,7 @@ func (s *TailnetSSH) setupHostKey() error {
 //
 // If Run returns an error, that means it can no longer listen - these
 // errors are fatal.
-func (s *TailnetSSH) Run(ctx context.Context, quit <-chan os.Signal) error {
+func (s *TailnetSSH) Run(ctx context.Context) error {
 	var err error
 	s.Server.Handler = s.handle
 
@@ -95,21 +95,18 @@ func (s *TailnetSSH) Run(ctx context.Context, quit <-chan os.Signal) error {
 		return fmt.Errorf("could not listen on tailnet: %w", err)
 	}
 
-	terminated := false
 	go func() {
-		signal := <-quit
-		terminated = true
-		log.Printf("Received signal %v, terminating...", signal)
+		<-ctx.Done()
 		srv.Close()
 	}()
 
-	err = s.setupPrometheus(srv)
+	err = s.setupPrometheus(ctx, srv)
 	if err != nil {
 		log.Printf("Setting up prometheus failed, but continuing anyway: %v", err)
 	}
 	log.Printf("starting ssh server on port :22...")
 	err = s.Server.Serve(listener)
-	if err != nil && !terminated {
+	if err != nil && ctx.Err() == nil {
 		return fmt.Errorf("ssh server failed serving: %w", err)
 	}
 	return nil