main.go

package main

import (
	"bytes"
	"context"
	"flag"
	"fmt"
	"io"
	"math/rand/v2"
	"net"
	"os"
	"time"
	"unsafe"

	"github.com/appleboy/graceful"
	"github.com/valyala/fasthttp"
)

type LogLevel uint8

const (
	DefaultMaxConcurrent = 128
	DefaultAddr          = ":13002"
	DefaultDNS           = ""
	DefaultTimeout       = 60 * time.Second
	DefaultLogLevel      = 1

	LogLevelDebug LogLevel = 0
	LogLevelInfo  LogLevel = 1
	LogLevelWarn  LogLevel = 2
	LogLevelError LogLevel = 3
)

var (
	version = "dev"

	addrF          = flag.String("a", DefaultAddr, `Listen address.`)
	maxConcurrentF = flag.Int("c", DefaultMaxConcurrent, "Max concurrency for fasthttp server")
	dnsresolversF  = flag.String("n", DefaultDNS, `DNS nameserves, E.g. "8.8.8.8" or "1.1.1.1,8.8.8.8". Default is empty (OS default)`)
	timeoutF       = flag.Duration("t", DefaultTimeout, `Connection timeout. Examples: 1m or 10s`)
	logLevelF      = flag.Int("l", DefaultLogLevel, `Log level. Examples: 0 (debug), 1 (info), 2 (warn), 3 (error).`)
	proxyF         = flag.String("x", "", `Set up a proxy chain. E.g. "localhost:12345"`)
	usageF         = flag.Bool("h", false, "Show usage")
	verF           = flag.Bool("v", false, "Show version")

	addr          string
	maxConcurrent int
	dns           []string
	timeout       time.Duration
	logLevel      LogLevel
	ver           string
	proxyChain    string

	defaultResolver = &net.Resolver{
		PreferGo:     true,
		StrictErrors: false,
		Dial: func(ctx context.Context, _, _ string) (net.Conn, error) {
			d := net.Dialer{}
			return d.DialContext(ctx, "udp", randomDNS())
		},
	}

	defaultDialer = fasthttp.TCPDialer{
		Concurrency:      maxConcurrent,
		DNSCacheDuration: time.Minute,
	}

	fastclient = fasthttp.Client{
		NoDefaultUserAgentHeader: true,
		Dial:                     defaultDialer.DialDualStack,
		MaxConnWaitTimeout:       10 * time.Second,
	}
)

func Logging(level LogLevel, format string, args ...any) {
	if logLevel > level {
		return
	}
	if len(args) == 0 {
		format += "\n"
	}

	fmt.Printf("%s %s %s",
		time.Now().Local().Format("[2006-01-02 15:04:05]"),
		func(level LogLevel) string {
			switch level {
			case LogLevelDebug:
				return "DEBUG"
			case LogLevelInfo:
				return "INFO"
			case LogLevelWarn:
				return "WARN"
			}
			return "ERROR"
		}(level),
		fmt.Sprintf(format, args...))
}

func Debug(format string, args ...any) {
	Logging(LogLevelDebug, format, args...)
}

func Info(format string, args ...any) {
	Logging(LogLevelInfo, format, args...)
}

func Warn(format string, args ...any) {
	Logging(LogLevelWarn, format, args...)
}

func Error(format string, args ...any) {
	Logging(LogLevelError, format, args...)
}

func init() {
	flag.Parse()
	if *usageF {
		flag.Usage()
		os.Exit(0)
	}

	addr = *addrF
	maxConcurrent = *maxConcurrentF
	dnsBytes := bytes.FieldsFunc([]byte(*dnsresolversF), func(c rune) bool {
		return c == ','
	})
	dns = make([]string, len(dnsBytes))
	for i, b := range dnsBytes {
		dns[i] = string(b)
	}
	timeout = *timeoutF
	proxyChain = *proxyF

	if len(dns) > 0 {
		defaultDialer.Resolver = defaultResolver
	}

	logLevel = LogLevel(*logLevelF)

	if *verF {
		println(version)
		os.Exit(0)
	}

	ver = version
}

func randomDNS() string {
	return dns[rand.IntN(len(dns))] + ":53"
}

func dialThroughProxy(host string, timeout time.Duration) (net.Conn, error) {
	if proxyChain == "" {
		return defaultDialer.DialTimeout(host, timeout)
	}

	// Connect to the upstream proxy
	proxyConn, err := defaultDialer.DialTimeout(proxyChain, timeout)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to proxy %s: %w", proxyChain, err)
	}

	// Send CONNECT request to the proxy
	connectReq := fmt.Sprintf("CONNECT %s HTTP/1.1\r\nHost: %s\r\n\r\n", host, host)
	if _, err := proxyConn.Write([]byte(connectReq)); err != nil {
		proxyConn.Close()
		return nil, fmt.Errorf("failed to send CONNECT to proxy: %w", err)
	}

	// Read the response from the proxy
	buf := make([]byte, 4096)
	n, err := proxyConn.Read(buf)
	if err != nil {
		proxyConn.Close()
		return nil, fmt.Errorf("failed to read CONNECT response: %w", err)
	}

	// Check if the proxy accepted the connection (200 OK)
	response := buf[:n]
	if !bytes.Contains(response, []byte("200")) {
		proxyConn.Close()
		return nil, fmt.Errorf("proxy rejected CONNECT: %s", b2s(bytes.Split(response, []byte("\r\n"))[0]))
	}

	Debug("Connected through proxy %s to %s", proxyChain, host)
	return proxyConn, nil
}

func transfer(destination io.WriteCloser, source io.ReadCloser) {
	defer func() {
		if err := recover(); err != nil {
			Warn("transfer: %s", err)
		}
	}()

	if _, err := io.Copy(destination, source); err != nil {
		Debug("transfer io closed: %s", err)
	}
}

func handleFastHTTP(ctx *fasthttp.RequestCtx) {
	Info("Connect to: http://%s\n", b2s(ctx.Host()))

	client := &fastclient
	if proxyChain != "" {
		// Create a custom client that dials through the proxy
		proxyClient := &fasthttp.Client{
			NoDefaultUserAgentHeader: true,
			Dial: func(addr string) (net.Conn, error) {
				return dialThroughProxy(addr, timeout)
			},
			MaxConnWaitTimeout: 10 * time.Second,
		}
		client = proxyClient
	}

	if err := client.DoTimeout(&ctx.Request, &ctx.Response, timeout); err != nil {
		Error("Client timeout: %s", err)
	}
}

func handleFastHTTPS(ctx *fasthttp.RequestCtx) {
	// RFC 7231: The CONNECT method identifies the destination server by the request-target.
	targetBuf := ctx.Request.Header.RequestURI()
	targetBuf = bytes.TrimPrefix(targetBuf, []byte("/"))
	target := b2s(targetBuf)
	Info("Connect to: https://%s\n", target)

	// Tell fasthttp not to send any automatic response
	ctx.HijackSetNoResponse(true)

	ctx.Hijack(func(clientConn net.Conn) {
		// Manually write the 200 Connection Established response for CONNECT
		_, err := clientConn.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n"))
		if err != nil {
			Error("Failed to send CONNECT response: %s", err)
			return
		}

		// Now establish the tunnel to the destination (through proxy if configured)
		destConn, err := dialThroughProxy(target, timeout)
		if err != nil {
			Error("Dial timeout: %s", err)
			return
		}

		defer clientConn.Close()
		defer destConn.Close()

		go transfer(destConn, clientConn)
		transfer(clientConn, destConn)
	})
}

func b2s(b []byte) string {
	return *(*string)(unsafe.Pointer(&b))
}

func wait(server *fasthttp.Server) <-chan struct{} {
	graceful.NewManager().AddRunningJob(func(ctx context.Context) error {
		<-ctx.Done()
		server.DisableKeepalive = true
		if err := server.Shutdown(); err != nil {
			Warn("Shutdown err: %s", err)
			defer os.Exit(1)
		} else {
			Info("gracefully stopped")
		}
		return nil
	})

	return graceful.NewManager().Done()
}

func fastHTTPHandler(ctx *fasthttp.RequestCtx) {
	method := ctx.Method()
	if bytes.Equal(method, []byte(fasthttp.MethodConnect)) {
		handleFastHTTPS(ctx)
	} else {
		handleFastHTTP(ctx)
	}
}

func main() {
	server := &fasthttp.Server{
		Handler:            fastHTTPHandler,
		ReadTimeout:        timeout,
		WriteTimeout:       timeout,
		MaxConnsPerIP:      1024,
		MaxRequestsPerConn: 1024,
		IdleTimeout:        3 * timeout,
		ReduceMemoryUsage:  true,
		CloseOnShutdown:    true,
		Concurrency:        maxConcurrent,
	}

	go func() {
		Info("Version: %s\n", ver)
		Info("Concurrency: %d\n", maxConcurrent)
		Info("Nameservers: %v\n", dns)
		Info("Connection timeout is %s\n", timeout)
		Info("listening on address %s\n", addr)
		if err := server.ListenAndServe(addr); err != nil {
			Error("Error in ListenAndServe: %s\n", err)
		}
	}()

	<-wait(server)
}