Merge pull request #237 from SenseUnit/fix_subprocess_ctx

Fix subprocess lifecycle

Author: Snawoot

dialer/command.go

@@ -49,7 +49,8 @@ func (d *CommandDialer) Dial(network, address string) (net.Conn, error) {
 }
 
 func (d *CommandDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
-	cmd := exec.CommandContext(ctx, d.command[0], d.command[1:]...)
+	killCtx, kill := context.WithCancel(context.Background())
+	cmd := exec.CommandContext(killCtx, d.command[0], d.command[1:]...)
 	cmd.Env = append(os.Environ(),
 		"DUMBPROXY_DST_ADDR="+address,
 		"DUMBPROXY_DST_NET="+network,
@@ -72,7 +73,7 @@ func (d *CommandDialer) DialContext(ctx context.Context, network, address string
 		cmdIn.Close()
 		cmdOut.Close()
 	}()
-	return NewPipeConn(cmdOut.(ReadPipe), cmdIn.(WritePipe)), nil
+	return NewPipeConn(cmdOut.(ReadPipe), cmdIn.(WritePipe), kill), nil
 }
 
 func (d *CommandDialer) WantsHostname(_ context.Context, _, _ string) bool {

dialer/pipewrap.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"net"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/hashicorp/go-multierror"
@@ -40,16 +41,19 @@ func (a PipeAddr) String() string {
 }
 
 type PipeConn struct {
-	r  ReadPipe
-	w  WritePipe
-	rc sync.Once
-	wc sync.Once
+	r          ReadPipe
+	w          WritePipe
+	onClose    func()
+	rc         sync.Once
+	wc         sync.Once
+	closeFlags atomic.Uint32
 }
 
-func NewPipeConn(r ReadPipe, w WritePipe) *PipeConn {
+func NewPipeConn(r ReadPipe, w WritePipe, onClose func()) *PipeConn {
 	return &PipeConn{
-		r: r,
-		w: w,
+		r:       r,
+		w:       w,
+		onClose: onClose,
 	}
 }
 
@@ -76,14 +80,26 @@ func (c *PipeConn) CloseWrite() error {
 	var err error
 	c.wc.Do(func() {
 		err = c.w.Close()
+		if c.onClose != nil {
+			c.closeFlags.Or(2)
+			if c.closeFlags.Load() == 3 {
+				c.onClose()
+			}
+		}
 	})
 	return err
 }
 
 func (c *PipeConn) CloseRead() error {
 	var err error
-	c.wc.Do(func() {
+	c.rc.Do(func() {
 		err = c.r.Close()
+		if c.onClose != nil {
+			c.closeFlags.Or(1)
+			if c.closeFlags.Load() == 3 {
+				c.onClose()
+			}
+		}
 	})
 	return err
 }

handler/stdio.go

@@ -20,7 +20,7 @@ func StdIOHandler(d HandlerDialer, logger *clog.CondLogger, forward ForwardFunc)
 		}
 		defer target.Close()
 
-		return forward(ctx, "", dialer.NewPipeConn(reader, writer), target, "tcp", dstAddress)
+		return forward(ctx, "", dialer.NewPipeConn(reader, writer, nil), target, "tcp", dstAddress)
 	}
 }