From 62cfee72c64921842ddb2a103d6982fe6abbf7cb Mon Sep 17 00:00:00 2001
From: Andrey Meshkov <am@adguard.com>
Date: Mon, 5 Feb 2024 20:07:21 +0300
Subject: [PATCH] Added probe-reverseproxyurl configuration parameter

---
 CHANGELOG.md                     |   8 +-
 README.md                        |  18 +++
 internal/cmd/cmd.go              |  15 +--
 internal/cmd/options.go          |   5 +
 internal/pipe/server.go          | 202 ++++++++++++++++++++++---------
 internal/pipe/server_test.go     |  72 +++++++++++
 internal/tunnel/msgreadwriter.go |   2 +-
 7 files changed, 257 insertions(+), 65 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5ea0b4f..1ea90fc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,7 +11,13 @@ adheres to [Semantic Versioning][semver].
 
 ## [Unreleased]
 
-[unreleased]: https://github.com/ameshkov/udptlspipe/compare/v1.2.1...HEAD
+[unreleased]: https://github.com/ameshkov/udptlspipe/compare/v1.3.0...HEAD
+
+## [1.3.0] - 2024-02-05
+
+* Added an option to configure a probe reverse proxy URL.
+
+[1.3.0]: https://github.com/ameshkov/udptlspipe/releases/tag/v1.3.0
 
 ## [1.2.2] - 2024-02-05
 
diff --git a/README.md b/README.md
index b5126f5..7a8d245 100644
--- a/README.md
+++ b/README.md
@@ -15,6 +15,7 @@ to keep it that way.
 * [How to install udptlspipe](#install)
 * [How to use udptlspipe](#howtouse)
 * [Custom TLS certificate](#tlscert)
+* [Probing protection](#probing)
 * [Docker](#docker)
 * [All command-line arguments](#allcmdarguments)
 
@@ -160,6 +161,23 @@ udptlspipe \
 
 [lego]: https://go-acme.github.io/lego/usage/cli/obtain-a-certificate/
 
+## Probing protection
+
+By default `udptlspipe` responds with a generic `403 Forbidden` response to
+unauthorized requests. However, it allows to use a more sophisticated
+protection. If `--probe-reverseproxyurl` is specified, `udptlspipe` server will
+proxy unauthorized requests to the specified target while rewriting `Host` and
+keeping the original path. This way you can imitate a real existing website.
+
+```shell
+udptlspipe --server \
+  -l 0.0.0.0:443 \
+  -d 2.3.4.5:8123 \
+  -p SecurePassword \
+  --probe-reverseproxyurl "http://example.com"
+
+```
+
 <a id="docker"></a>
 
 ## Docker
diff --git a/internal/cmd/cmd.go b/internal/cmd/cmd.go
index 70722c2..9da87d2 100644
--- a/internal/cmd/cmd.go
+++ b/internal/cmd/cmd.go
@@ -46,13 +46,14 @@ func Main() {
 	log.Info("Configuration:\n%s", o)
 
 	cfg := &pipe.Config{
-		ListenAddr:        o.ListenAddr,
-		DestinationAddr:   o.DestinationAddr,
-		Password:          o.Password,
-		ServerMode:        o.ServerMode,
-		ProxyURL:          o.ProxyURL,
-		VerifyCertificate: o.VerifyCertificate,
-		TLSServerName:     o.TLSServerName,
+		ListenAddr:           o.ListenAddr,
+		DestinationAddr:      o.DestinationAddr,
+		Password:             o.Password,
+		ServerMode:           o.ServerMode,
+		ProxyURL:             o.ProxyURL,
+		VerifyCertificate:    o.VerifyCertificate,
+		TLSServerName:        o.TLSServerName,
+		ProbeReverseProxyURL: o.ProbeReverseProxyURL,
 	}
 
 	if o.TLSCertPath != "" {
diff --git a/internal/cmd/options.go b/internal/cmd/options.go
index c371b7a..5305075 100644
--- a/internal/cmd/options.go
+++ b/internal/cmd/options.go
@@ -54,6 +54,11 @@ type Options struct {
 	// certificate specified by TLSCertPath.
 	TLSCertKey string `yaml:"tls-keyfile" long:"tls-keyfile" description:"Path to the private key for the cert specified in tls-certfile." value-name:"<path-to-key-file>"`
 
+	// ProbeReverseProxyURL is the URL that will be used by the reverse HTTP
+	// proxy to respond to unauthorized or proxy requests. If not specified,
+	// it will respond with a stub page 403 Forbidden.
+	ProbeReverseProxyURL string `yaml:"probe-reverseproxyurl" long:"probe-reverseproxyurl" description:"Unauthorized requests and probes will be proxied to the URL." value-name:"<hostname>"`
+
 	// Verbose defines whether we should write the DEBUG-level log or not.
 	Verbose bool `yaml:"verbose" short:"v" long:"verbose" description:"Verbose output (optional)." optional:"yes" optional-value:"true"`
 }
diff --git a/internal/pipe/server.go b/internal/pipe/server.go
index 045465b..56b2be5 100644
--- a/internal/pipe/server.go
+++ b/internal/pipe/server.go
@@ -11,6 +11,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"net/http/httputil"
 	"net/url"
 	"os"
 	"strings"
@@ -42,6 +43,9 @@ type Server struct {
 	dialer          proxy.Dialer
 	serverMode      bool
 
+	probeReverseProxyURL    string
+	probeReverseProxyListen net.Listener
+
 	// tlsConfig to use for TLS connections. In server mode it also has the
 	// certificate that will be used.
 	tlsConfig *tls.Config
@@ -116,6 +120,11 @@ type Config struct {
 	// only for server mode. If not configured, the server will generate a stub
 	// self-signed certificate automatically.
 	TLSCertificate *tls.Certificate
+
+	// ProbeReverseProxyURL is the URL that will be used by the reverse HTTP
+	// proxy to respond to unauthorized or proxy requests. If not specified,
+	// it will respond with a stub page 403 Forbidden.
+	ProbeReverseProxyURL string
 }
 
 // createTLSConfig creates a TLS configuration as per the server configuration.
@@ -156,15 +165,16 @@ func createTLSConfig(config *Config) (tlsConfig *tls.Config, err error) {
 // NewServer creates a new instance of a *Server.
 func NewServer(config *Config) (s *Server, err error) {
 	s = &Server{
-		listenAddr:      config.ListenAddr,
-		destinationAddr: config.DestinationAddr,
-		password:        config.Password,
-		dialer:          proxy.Direct,
-		serverMode:      config.ServerMode,
-		srcConns:        map[net.Conn]struct{}{},
-		srcConnsMu:      &sync.Mutex{},
-		dstConns:        map[net.Conn]struct{}{},
-		dstConnsMu:      &sync.Mutex{},
+		listenAddr:           config.ListenAddr,
+		destinationAddr:      config.DestinationAddr,
+		password:             config.Password,
+		probeReverseProxyURL: config.ProbeReverseProxyURL,
+		dialer:               proxy.Direct,
+		serverMode:           config.ServerMode,
+		srcConns:             map[net.Conn]struct{}{},
+		srcConnsMu:           &sync.Mutex{},
+		dstConns:             map[net.Conn]struct{}{},
+		dstConnsMu:           &sync.Mutex{},
 	}
 
 	s.tlsConfig, err = createTLSConfig(config)
@@ -200,13 +210,13 @@ func (s *Server) Addr() (addr net.Addr) {
 // Start starts the pipe, exits immediately if it failed to start
 // listening.  Start returns once all servers are considered up.
 func (s *Server) Start() (err error) {
-	log.Info("Starting the pipe %s", s)
+	log.Info("Starting the server %s", s)
 
 	s.lock.Lock()
 	defer s.lock.Unlock()
 
 	if s.started {
-		return errors.New("pipe is already started")
+		return errors.New("Server is already started")
 	}
 
 	s.listen, err = s.createListener()
@@ -214,6 +224,13 @@ func (s *Server) Start() (err error) {
 		return fmt.Errorf("failed to start pipe: %w", err)
 	}
 
+	if s.probeReverseProxyURL != "" {
+		err = s.startProbeReverseProxy()
+		if err != nil {
+			return fmt.Errorf("failed to start probe reverse proxy: %w", err)
+		}
+	}
+
 	s.wg.Add(1)
 	go s.serve()
 
@@ -241,15 +258,63 @@ func (s *Server) createListener() (l net.Listener, err error) {
 	return l, nil
 }
 
+// startProbeReverseProxy starts a reverse HTTP proxy that will be used for
+// answering unauthorized and probe requests. Returns the listener of that
+// proxy. Original request URI will be appended to proxyURL.
+func (s *Server) startProbeReverseProxy() (err error) {
+	proxyURL := s.probeReverseProxyURL
+
+	if _, err = url.Parse(proxyURL); err != nil {
+		return fmt.Errorf("reverse proxy URL must be a valid URL: %w", err)
+	}
+
+	targetURL, err := url.Parse(s.probeReverseProxyURL)
+	if err != nil {
+		return fmt.Errorf("reverse proxy URL must be a valid URL: %w", err)
+	}
+
+	handler := &httputil.ReverseProxy{
+		Rewrite: func(r *httputil.ProxyRequest) {
+			r.SetURL(targetURL)
+			r.Out.Host = targetURL.Host
+		},
+	}
+
+	srv := &http.Server{
+		ReadHeaderTimeout: upgradeTimeout,
+		Handler:           handler,
+	}
+
+	s.probeReverseProxyListen, err = net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		return fmt.Errorf("failed to start probe reverse proxy: %w", err)
+	}
+
+	s.wg.Add(1)
+	go func() {
+		defer s.wg.Done()
+
+		log.Info("Starting probe reverse proxy")
+		sErr := srv.Serve(s.probeReverseProxyListen)
+		log.Info("Probe reverse proxy has been stopped due to: %v", sErr)
+	}()
+
+	return nil
+}
+
 // Shutdown stops the pipe and waits for all active connections to close.
 func (s *Server) Shutdown(ctx context.Context) (err error) {
-	log.Info("Stopping the pipe %s", s)
+	log.Info("Stopping the server %s", s)
 
 	s.stopServeLoop()
 
 	// Closing the udpConn thread.
 	log.OnCloserError(s.listen, log.DEBUG)
 
+	if s.probeReverseProxyListen != nil {
+		log.OnCloserError(s.probeReverseProxyListen, log.DEBUG)
+	}
+
 	// Closing active TCP connections.
 	s.closeConnections(s.srcConnsMu, s.srcConns)
 
@@ -323,7 +388,7 @@ func (s *Server) serve() {
 func (s *Server) acceptConn() (err error) {
 	conn, err := s.listen.Accept()
 	if err != nil {
-		// This type of errors should not lead to stopping the pipe.
+		// This type of errors should not lead to stopping the server.
 		if errors.Is(os.ErrDeadlineExceeded, err) {
 			return nil
 		}
@@ -338,13 +403,7 @@ func (s *Server) acceptConn() (err error) {
 
 	log.Debug("Accepted new connection from %s", conn.RemoteAddr())
 
-	func() {
-		s.srcConnsMu.Lock()
-		defer s.srcConnsMu.Unlock()
-
-		// Track the connection to allow unblocking reads on shutdown.
-		s.srcConns[conn] = struct{}{}
-	}()
+	s.saveSrcConn(conn)
 
 	s.wg.Add(1)
 	go s.serveConn(conn)
@@ -352,6 +411,15 @@ func (s *Server) acceptConn() (err error) {
 	return nil
 }
 
+// saveSrcConn tracks the connection to allow unblocking reads on shutdown.
+func (s *Server) saveSrcConn(conn net.Conn) {
+	s.srcConnsMu.Lock()
+	defer s.srcConnsMu.Unlock()
+
+	// Track the connection to allow unblocking reads on shutdown.
+	s.srcConns[conn] = struct{}{}
+}
+
 // closeSrcConn closes the source connection and cleans up after it.
 func (s *Server) closeSrcConn(conn net.Conn) {
 	log.OnCloserError(conn, log.DEBUG)
@@ -362,6 +430,15 @@ func (s *Server) closeSrcConn(conn net.Conn) {
 	delete(s.srcConns, conn)
 }
 
+// saveDstConn tracks the connection to allow unblocking reads on shutdown.
+func (s *Server) saveDstConn(conn net.Conn) {
+	s.dstConnsMu.Lock()
+	defer s.dstConnsMu.Unlock()
+
+	// Track the connection to allow unblocking reads on shutdown.
+	s.dstConns[conn] = struct{}{}
+}
+
 // closeDstConn closes the destination connection and cleans up after it.
 func (s *Server) closeDstConn(conn net.Conn) {
 	// No destination connection opened yet, do nothing.
@@ -425,23 +502,42 @@ func (s *Server) upgradeClientConn(conn net.Conn) (rwc io.ReadWriteCloser, err e
 	), nil
 }
 
-// respondWithDummyPage writes a dummy response to the client (part of active
-// probing protection).
-func (s *Server) respondWithDummyPage(conn net.Conn, req *http.Request) {
-	response := fmt.Sprintf("%s 403 Forbidden\r\n", req.Proto) +
-		"Server: nginx\r\n" +
-		fmt.Sprintf("Date: %s\r\n", time.Now().Format(http.TimeFormat)) +
-		"Content-Type: text/html\r\n" +
-		"Connection: close\r\n" +
-		"\r\n" +
-		"<html>\r\n" +
-		"<head><title>403 Forbidden</title></head>\r\n" +
-		"<center><h1>403 Forbidden</h1></center>\r\n" +
-		"<hr><center>nginx</center>\r\n" +
-		"</body>\r\n" +
-		"</html>\r\n"
-
-	_, _ = conn.Write([]byte(response))
+// respondToProbe writes a dummy response to the client if it's not authorized
+// or if it's a probe.
+func (s *Server) respondToProbe(rwc io.ReadWriteCloser, req *http.Request) {
+	if s.probeReverseProxyListen == nil {
+		log.Debug("No probe reverse proxy configured, respond with a dummy 403 page")
+
+		response := fmt.Sprintf("%s 403 Forbidden\r\n", req.Proto) +
+			"Server: nginx\r\n" +
+			fmt.Sprintf("Date: %s\r\n", time.Now().Format(http.TimeFormat)) +
+			"Content-Type: text/html\r\n" +
+			"Connection: close\r\n" +
+			"\r\n" +
+			"<html>\r\n" +
+			"<head><title>403 Forbidden</title></head>\r\n" +
+			"<center><h1>403 Forbidden</h1></center>\r\n" +
+			"<hr><center>nginx</center>\r\n" +
+			"</body>\r\n" +
+			"</html>\r\n"
+
+		_, _ = rwc.Write([]byte(response))
+
+		return
+	}
+
+	log.Debug("Probe reverse proxy is configured, tunnel data to it")
+
+	proxyConn, err := net.Dial("tcp", s.probeReverseProxyListen.Addr().String())
+	if err != nil {
+		log.Error("Failed to connect to the probe reverse proxy: %v", err)
+
+		return
+	}
+
+	s.saveDstConn(proxyConn)
+
+	tunnel.Tunnel("probeReverseProxy", rwc, proxyConn)
 }
 
 // upgradeServerConn attempts to upgrade the server connection and returns a
@@ -466,33 +562,33 @@ func (s *Server) upgradeServerConn(conn net.Conn) (rwc io.ReadWriteCloser, err e
 		return nil, fmt.Errorf("cannot read HTTP request: %w", err)
 	}
 
+	// Now that authentication check has been done restore the peeked up data
+	// so that it could be used further.
+	originalRwc := &readWriteCloser{
+		Reader: io.MultiReader(bytes.NewReader(buf.Bytes()), conn),
+		Writer: conn,
+		Closer: conn,
+	}
+
 	if !strings.EqualFold(req.Header.Get("Upgrade"), "websocket") {
-		s.respondWithDummyPage(conn, req)
+		s.respondToProbe(originalRwc, req)
 
 		return nil, fmt.Errorf("not a websocket")
 	}
 
 	clientPassword := req.URL.Query().Get("password")
 	if s.password != "" && clientPassword != s.password {
-		s.respondWithDummyPage(conn, req)
+		s.respondToProbe(originalRwc, req)
 
 		return nil, fmt.Errorf("wrong password: %s", clientPassword)
 	}
 
-	// Now that authentication check has been done restore the peeked up data
-	// and use ws.Upgrade to do the actual WebSocket upgrade.
-	multiRwc := &readWriteCloser{
-		Reader: io.MultiReader(bytes.NewReader(buf.Bytes()), conn),
-		Writer: conn,
-		Closer: conn,
-	}
-
-	_, err = ws.Upgrade(multiRwc)
+	_, err = ws.Upgrade(originalRwc)
 	if err != nil {
 		return nil, fmt.Errorf("failed to upgrade WebSocket: %w", err)
 	}
 
-	return newWsConn(multiRwc, conn.RemoteAddr(), ws.StateServerSide), nil
+	return newWsConn(originalRwc, conn.RemoteAddr(), ws.StateServerSide), nil
 }
 
 // serveConn processes incoming connection, authenticates it and proxies the
@@ -532,13 +628,7 @@ func (s *Server) processConn(rwc io.ReadWriteCloser) {
 		return
 	}
 
-	func() {
-		s.dstConnsMu.Lock()
-		defer s.dstConnsMu.Unlock()
-
-		// Track the connection to allow unblocking reads on shutdown.
-		s.dstConns[dstConn] = struct{}{}
-	}()
+	s.saveDstConn(dstConn)
 
 	var dstRwc io.ReadWriteCloser = dstConn
 	if !s.serverMode {
diff --git a/internal/pipe/server_test.go b/internal/pipe/server_test.go
index 3125e4c..2674e77 100644
--- a/internal/pipe/server_test.go
+++ b/internal/pipe/server_test.go
@@ -1,12 +1,17 @@
 package pipe_test
 
 import (
+	"bufio"
 	"context"
 	"fmt"
 	"io"
 	"net"
+	"net/http"
 	"strings"
 	"testing"
+	"time"
+
+	tls "github.com/refraction-networking/utls"
 
 	"github.com/AdguardTeam/golibs/log"
 	"github.com/ameshkov/udptlspipe/internal/pipe"
@@ -78,3 +83,70 @@ func TestServer_Start_integration(t *testing.T) {
 		require.Equal(t, msg, udpMsg)
 	}
 }
+
+func TestServer_Start_probeDetection(t *testing.T) {
+	// Start the echo server where the pipe will proxy data.
+	udpSrv := &testutil.UDPEchoServer{}
+	err := udpSrv.Start()
+	require.NoError(t, err)
+	defer log.OnCloserError(udpSrv, log.DEBUG)
+
+	// The server where pipeServer will proxy probes.
+	rewriteListen, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	defer log.OnCloserError(rewriteListen, log.DEBUG)
+
+	rewriteServer := &http.Server{
+		Handler: http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+			rw.WriteHeader(http.StatusOK)
+			_, _ = rw.Write([]byte(r.URL.String()))
+		}),
+	}
+	go func() {
+		_ = rewriteServer.Serve(rewriteListen)
+	}()
+
+	// Create the pipe server proxying data to that UDP echo server.
+	pipeServer, err := pipe.NewServer(&pipe.Config{
+		ListenAddr:           "127.0.0.1:0",
+		DestinationAddr:      udpSrv.Addr(),
+		Password:             "123123",
+		ServerMode:           true,
+		ProbeReverseProxyURL: fmt.Sprintf("http://%s", rewriteListen.Addr()),
+	})
+	require.NoError(t, err)
+
+	// Start the pipe server, it's ready to accept new connections.
+	err = pipeServer.Start()
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, pipeServer.Shutdown(context.Background()))
+	}()
+
+	// Connect to the pipe server.
+	pipeConn, err := tls.Dial("tcp", pipeServer.Addr().String(), &tls.Config{InsecureSkipVerify: true})
+	require.NoError(t, err)
+	defer log.OnCloserError(pipeConn, log.DEBUG)
+
+	_ = pipeConn.SetDeadline(time.Now().Add(time.Second * 1))
+
+	// Create a probe request.
+	r, err := http.NewRequest(http.MethodGet, "https://example.com/probe", nil)
+	require.NoError(t, err)
+
+	// Send the probe request to the server.
+	err = r.Write(pipeConn)
+	require.NoError(t, err)
+
+	// Read what the server responded with to the probe.
+	br := bufio.NewReader(pipeConn)
+	resp, err := http.ReadResponse(br, nil)
+	require.NoError(t, err)
+
+	// Check that the response is the same as the reverse proxy has written.
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+
+	responseBytes, err := io.ReadAll(resp.Body)
+	require.NoError(t, err)
+	require.Equal(t, "/probe", string(responseBytes))
+}
diff --git a/internal/tunnel/msgreadwriter.go b/internal/tunnel/msgreadwriter.go
index d159195..0c7c12b 100644
--- a/internal/tunnel/msgreadwriter.go
+++ b/internal/tunnel/msgreadwriter.go
@@ -11,7 +11,7 @@ import (
 
 // MaxMessageLength is the maximum length that is safe to use.
 // TODO(ameshkov): Make it configurable.
-const MaxMessageLength = 1280
+const MaxMessageLength = 1320
 
 // MinMessageLength is the minimum message size. If the message is smaller, it
 // will be padded with random bytes.