Abort proxy when stdin gets closed

internal/command/proxy/proxy.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"strings"
 
 	"github.com/spf13/cobra"
@@ -14,6 +15,7 @@ import (
 	"github.com/superfly/flyctl/internal/flag/flagnames"
 	"github.com/superfly/flyctl/internal/flyutil"
 	"github.com/superfly/flyctl/internal/prompt"
+	"github.com/superfly/flyctl/iostreams"
 	"github.com/superfly/flyctl/proxy"
 )
 
@@ -50,6 +52,11 @@ connects to the first Machine address returned by an internal DNS query on the a
 			Default:     "127.0.0.1",
 			Description: "Local address to bind to",
 		},
+		flag.Bool{
+			Name:        "watch-stdin",
+			Default:     false,
+			Description: "Watches stdin and terminates once it gets closed",
+		},
 	)
 
 	return cmd
@@ -131,5 +138,43 @@ func run(ctx context.Context) (err error) {
 		params.RemoteHost = fmt.Sprintf("%s.internal", appName)
 	}
 
+	if flag.GetBool(ctx, "watch-stdin") {
+		ctx = watchStdinAndAbortOnClose(ctx)
+	}
+
 	return proxy.Connect(ctx, params)
 }
+
+// Asynchronously watches stdin and abort when it closes.
+//
+// There is no guarantee that a OS process spawning the proxy will
+// terminate it, however the stdin is always closed whtn the parent
+// terminates. This way we make sure there are no zombie processes,
+// especially that they hold onto TCP ports.
+//
+// Note that we don't do this when stdin is TTY, because that prevents
+// the process from being moved to a background job on Unix.
+// See https://github.com/brunch/brunch/issues/998.
+func watchStdinAndAbortOnClose(ctx context.Context) context.Context {
+	ctx, cancel := context.WithCancelCause(ctx)
+
+	ios := iostreams.FromContext(ctx)
+
+	if !ios.IsStdinTTY() {
+		go func() {
+			// We don't expect any input, but if there is one, we ignore it
+			// to avoid allocating space unnecessarily
+			buffer := make([]byte, 1)
+			for {
+				_, err := ios.In.Read(buffer)
+				if err == io.EOF {
+					cancel(nil)
+				} else if err != nil {
+					cancel(err)
+				}
+			}
+		}()
+	}
+
+	return ctx
+}