some socks connect handler done

Author: Snawoot

handler/adapter.go

@@ -134,3 +134,30 @@ func (w wrappedH1RespWriter) Close() error {
 }
 
 var _ io.ReadWriteCloser = wrappedH1RespWriter{}
+
+type wrappedSOCKS struct {
+	r io.Reader
+	w io.Writer
+}
+
+func wrapSOCKS(r io.Reader, w io.Writer) wrappedSOCKS {
+	return wrappedSOCKS{
+		r: r,
+		w: w,
+	}
+}
+
+func (w wrappedSOCKS) Read(p []byte) (n int, err error) {
+	return w.r.Read(p)
+}
+
+func (w wrappedSOCKS) Write(p []byte) (n int, err error) {
+	return w.w.Write(p)
+}
+
+func (w wrappedSOCKS) Close() error {
+	// can't really close SOCKS reader or writer
+	return nil
+}
+
+var _ io.ReadWriteCloser = wrappedSOCKS{}

handler/handler.go

@@ -28,11 +28,13 @@ type HandlerDialer interface {
 	DialContext(ctx context.Context, net, address string) (net.Conn, error)
 }
 
+type ForwardFunc = func(ctx context.Context, username string, incoming, outgoing io.ReadWriteCloser) error
+
 type ProxyHandler struct {
 	auth          auth.Auth
 	logger        *clog.CondLogger
 	dialer        HandlerDialer
-	forward       func(ctx context.Context, username string, incoming, outgoing io.ReadWriteCloser) error
+	forward       ForwardFunc
 	httptransport http.RoundTripper
 	outbound      map[string]string
 	outboundMux   sync.RWMutex

handler/socks.go

@@ -0,0 +1,48 @@
+package handler
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"strings"
+
+	ddto "github.com/SenseUnit/dumbproxy/dialer/dto"
+	clog "github.com/SenseUnit/dumbproxy/log"
+	"github.com/things-go/go-socks5"
+	"github.com/things-go/go-socks5/statute"
+)
+
+func SOCKSHandler(dialer HandlerDialer, logger *clog.CondLogger, forward ForwardFunc) func(ctx context.Context, writer io.Writer, request *socks5.Request) error {
+	return func(ctx context.Context, writer io.Writer, request *socks5.Request) error {
+		username := ""
+		if request.AuthContext != nil {
+			username = request.AuthContext.Payload["username"]
+		}
+		localAddr := request.LocalAddr.String()
+		ctx = ddto.BoundDialerParamsToContext(ctx, nil, trimAddrPort(localAddr))
+		ctx = ddto.FilterParamsToContext(ctx, nil, username)
+		// TODO: add request logging
+		target, err := dialer.DialContext(ctx, "tcp", request.DestAddr.String())
+		if err != nil {
+			msg := err.Error()
+			resp := statute.RepHostUnreachable
+			if strings.Contains(msg, "refused") {
+				resp = statute.RepConnectionRefused
+			} else if strings.Contains(msg, "network is unreachable") {
+				resp = statute.RepNetworkUnreachable
+			}
+			if err := socks5.SendReply(writer, resp, nil); err != nil {
+				return fmt.Errorf("failed to send reply, %v", err)
+			}
+			return fmt.Errorf("connect to %v failed, %v", request.RawDestAddr, err)
+		}
+		defer target.Close() // nolint: errcheck
+
+		// Send success
+		if err := socks5.SendReply(writer, statute.RepSuccess, target.LocalAddr()); err != nil {
+			return fmt.Errorf("failed to send reply, %v", err)
+		}
+
+		return forward(ctx, username, wrapSOCKS(request.Reader, writer), target)
+	}
+}

jsext/dto.go

@@ -24,6 +24,9 @@ type JSRequestInfo struct {
 }
 
 func JSRequestInfoFromRequest(req *http.Request) *JSRequestInfo {
+	if req == nil {
+		return nil
+	}
 	return &JSRequestInfo{
 		Method:           req.Method,
 		URL:              req.URL.String(),

main.go

@@ -827,6 +827,13 @@ func run() int {
 					EnableConnect: true,
 				},
 			),
+			socks5.WithConnectHandle(
+				handler.SOCKSHandler(
+					dialerRoot,
+					proxyLogger,
+					forwarder,
+				),
+			),
 		}
 		switch cs := authProvider.(type) {
 		case auth.NoAuth: