stdio connection handler

Author: Snawoot

dialer/dto/dto.go

@@ -32,8 +32,10 @@ type filterContextParams struct {
 }
 
 func FilterParamsFromContext(ctx context.Context) (*http.Request, string) {
-	params := ctx.Value(filterContextKey{}).(filterContextParams)
-	return params.req, params.username
+	if params, ok := ctx.Value(filterContextKey{}).(filterContextParams); ok {
+		return params.req, params.username
+	}
+	return nil, ""
 }
 
 func FilterParamsToContext(ctx context.Context, req *http.Request, username string) context.Context {

handler/stdio.go

@@ -0,0 +1,73 @@
+package handler
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net"
+	"sync"
+
+	clog "github.com/SenseUnit/dumbproxy/log"
+)
+
+func StdIOHandler(dialer HandlerDialer, logger *clog.CondLogger, forward ForwardFunc) func(ctx context.Context, reader io.Reader, writer io.Writer, dstAddress string) error {
+	return func(ctx context.Context, reader io.Reader, writer io.Writer, dstAddress string) error {
+		logger.Debug("Request: %v => %v %q %v %v %v", "<stdio>", "<stdio>", "", "STDIO", "CONNECT", dstAddress)
+		target, err := dialer.DialContext(ctx, "tcp", dstAddress)
+		if err != nil {
+			return fmt.Errorf("connect to %q failed: %w", dstAddress, err)
+		}
+		defer target.Close()
+
+		return forward(ctx, "", wrapSOCKS(reader, writer), target)
+	}
+}
+
+type DummyAddress struct {
+	network string
+	address string
+}
+
+func (a DummyAddress) Network() string {
+	return a.network
+}
+
+func (a DummyAddress) String() string {
+	return a.address
+}
+
+type DummyListener struct {
+	address   DummyAddress
+	closed    chan struct{}
+	closeOnce sync.Once
+}
+
+// Accept waits for and returns the next connection to the listener.
+func (l *DummyListener) Accept() (net.Conn, error) {
+	<-l.closed
+	return nil, net.ErrClosed
+}
+
+// Close closes the listener.
+// Any blocked Accept operations will be unblocked and return errors.
+func (l *DummyListener) Close() error {
+	l.closeOnce.Do(func() {
+		close(l.closed)
+	})
+	return nil
+}
+
+// Addr returns the listener's network address.
+func (l *DummyListener) Addr() net.Addr {
+	return l.address
+}
+
+func DummyListen(network, address string) (net.Listener, error) {
+	return &DummyListener{
+		address: DummyAddress{
+			network: network,
+			address: address,
+		},
+		closed: make(chan struct{}),
+	}, nil
+}

main.go

@@ -217,6 +217,7 @@ const (
 	_ proxyMode = iota
 	proxyModeHTTP
 	proxyModeSOCKS5
+	proxyModeStdIO
 )
 
 type proxyModeArg struct {
@@ -230,6 +231,8 @@ func (a *proxyModeArg) Set(arg string) error {
 		val = proxyModeHTTP
 	case "socks", "socks5":
 		val = proxyModeSOCKS5
+	case "stdio":
+		val = proxyModeStdIO
 	default:
 		return fmt.Errorf("unrecognized proxy mode %q", arg)
 	}
@@ -243,6 +246,8 @@ func (a *proxyModeArg) String() string {
 		return "http"
 	case proxyModeSOCKS5:
 		return "socks5"
+	case proxyModeStdIO:
+		return "stdio"
 	default:
 		return fmt.Sprintf("proxyMode(%d)", int(a.value))
 	}
@@ -510,9 +515,16 @@ func run() int {
 		return 0
 	}
 
-	// we don't expect positional arguments in the main operation mode
-	if len(args.positionalArgs) > 0 {
-		arg_fail("Unexpected positional arguments! Check your command line.")
+	if args.mode.value == proxyModeStdIO {
+		// expect exactly two positional arguments
+		if len(args.positionalArgs) != 2 {
+			arg_fail("Exactly two positional arguments are expected in this mode: host and port")
+		}
+	} else {
+		// we don't expect positional arguments in the main operation mode
+		if len(args.positionalArgs) > 0 {
+			arg_fail("Unexpected positional arguments! Check your command line.")
+		}
 	}
 
 	// setup logging
@@ -640,7 +652,10 @@ func run() int {
 	mainLogger.Info("Starting proxy server...")
 
 	listenerFactory := net.Listen
-	if args.bindReusePort {
+	switch {
+	case args.mode.value == proxyModeStdIO:
+		listenerFactory = handler.DummyListen
+	case args.bindReusePort:
 		if reuseport.Available() {
 			listenerFactory = reuseport.Listen
 		} else {
@@ -872,6 +887,13 @@ func run() int {
 			mainLogger.Info("Reached normal server termination.")
 		}
 		return 0
+	case proxyModeStdIO:
+		handler := handler.StdIOHandler(dialerRoot, proxyLogger, forwarder)
+		address := net.JoinHostPort(args.positionalArgs[0], args.positionalArgs[1])
+		if err := handler(stopContext, os.Stdin, os.Stdout, address); err != nil {
+			mainLogger.Error("Connection interrupted: %v", err)
+		}
+		return 0
 	}
 
 	mainLogger.Critical("unknown proxy mode")