From fa8375f763fd7104cb3c7f083ad5f8f7b558ae68 Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Mon, 22 Nov 2021 22:06:47 +0800
Subject: [PATCH] add auto handler

---
 cmd/gost/register.go        |  1 +
 pkg/handler/auto/handler.go | 96 +++++++++++++++++++++++++++++++++++++
 2 files changed, 97 insertions(+)
 create mode 100644 pkg/handler/auto/handler.go

diff --git a/cmd/gost/register.go b/cmd/gost/register.go
index 13c6bd0..c642a7b 100644
--- a/cmd/gost/register.go
+++ b/cmd/gost/register.go
@@ -14,6 +14,7 @@ import (
 	_ "github.com/go-gost/gost/pkg/dialer/udp"
 
 	// Register handlers
+	_ "github.com/go-gost/gost/pkg/handler/auto"
 	_ "github.com/go-gost/gost/pkg/handler/forward/local"
 	_ "github.com/go-gost/gost/pkg/handler/forward/remote"
 	_ "github.com/go-gost/gost/pkg/handler/http"
diff --git a/pkg/handler/auto/handler.go b/pkg/handler/auto/handler.go
new file mode 100644
index 0000000..8b964be
--- /dev/null
+++ b/pkg/handler/auto/handler.go
@@ -0,0 +1,96 @@
+package auto
+
+import (
+	"bufio"
+	"context"
+	"net"
+
+	"github.com/go-gost/gosocks4"
+	"github.com/go-gost/gosocks5"
+	"github.com/go-gost/gost/pkg/handler"
+	http_handler "github.com/go-gost/gost/pkg/handler/http"
+	socks4_handler "github.com/go-gost/gost/pkg/handler/socks/v4"
+	socks5_handler "github.com/go-gost/gost/pkg/handler/socks/v5"
+	"github.com/go-gost/gost/pkg/logger"
+	md "github.com/go-gost/gost/pkg/metadata"
+	"github.com/go-gost/gost/pkg/registry"
+)
+
+func init() {
+	registry.RegisterHandler("auto", NewHandler)
+}
+
+type autoHandler struct {
+	httpHandler   handler.Handler
+	socks4Handler handler.Handler
+	socks5Handler handler.Handler
+	log           logger.Logger
+}
+
+func NewHandler(opts ...handler.Option) handler.Handler {
+	options := &handler.Options{}
+	for _, opt := range opts {
+		opt(options)
+	}
+
+	log := options.Logger
+	if log == nil {
+		log = logger.Default()
+	}
+
+	h := &autoHandler{
+		log: log,
+	}
+
+	v := append(opts,
+		handler.LoggerOption(log.WithFields(map[string]interface{}{"type": "http"})))
+	h.httpHandler = http_handler.NewHandler(v...)
+
+	v = append(opts,
+		handler.LoggerOption(log.WithFields(map[string]interface{}{"type": "socks4"})))
+	h.socks4Handler = socks4_handler.NewHandler(v...)
+
+	v = append(opts,
+		handler.LoggerOption(log.WithFields(map[string]interface{}{"type": "socks5"})))
+	h.socks5Handler = socks5_handler.NewHandler(v...)
+	return h
+}
+
+func (h *autoHandler) Init(md md.Metadata) error {
+	if err := h.httpHandler.Init(md); err != nil {
+		return err
+	}
+	if err := h.socks4Handler.Init(md); err != nil {
+		return err
+	}
+	if err := h.socks5Handler.Init(md); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (h *autoHandler) Handle(ctx context.Context, conn net.Conn) {
+	h.log = h.log.WithFields(map[string]interface{}{
+		"remote": conn.RemoteAddr().String(),
+		"local":  conn.LocalAddr().String(),
+	})
+
+	br := bufio.NewReader(conn)
+	b, err := br.Peek(1)
+	if err != nil {
+		h.log.Error(err)
+		conn.Close()
+		return
+	}
+
+	cc := handler.NewBufferReaderConn(conn, br)
+	switch b[0] {
+	case gosocks4.Ver4: // socks4
+		h.socks4Handler.Handle(ctx, cc)
+	case gosocks5.Ver5: // socks5
+		h.socks5Handler.Handle(ctx, cc)
+	default: // http
+		h.httpHandler.Handle(ctx, cc)
+	}
+
+}