From b34f9839f178708ced0db316b45a56c9c97f1417 Mon Sep 17 00:00:00 2001
From: chor <yim0008@gmail.com>
Date: Sat, 25 Apr 2020 23:27:28 +0800
Subject: [PATCH] add reverse proxy option for ws, mws, wss and mwss.

---
 cmd/gost/route.go |  2 ++
 ws.go             | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/cmd/gost/route.go b/cmd/gost/route.go
index 979b39b..3507811 100644
--- a/cmd/gost/route.go
+++ b/cmd/gost/route.go
@@ -134,6 +134,7 @@ func parseChainNode(ns string) (nodes []gost.Node, err error) {
 	wsOpts.WriteBufferSize = node.GetInt("wbuf")
 	wsOpts.UserAgent = node.Get("agent")
 	wsOpts.Path = node.Get("path")
+	wsOpts.Reverse = node.Get("reverse")
 
 	timeout := node.GetDuration("timeout")
 
@@ -352,6 +353,7 @@ func (r *route) GenRouters() ([]router, error) {
 		wsOpts.ReadBufferSize = node.GetInt("rbuf")
 		wsOpts.WriteBufferSize = node.GetInt("wbuf")
 		wsOpts.Path = node.Get("path")
+		wsOpts.Reverse = node.Get("reverse")
 
 		ttl := node.GetDuration("ttl")
 		timeout := node.GetDuration("timeout")
diff --git a/ws.go b/ws.go
index 7a5a3ea..78f8227 100644
--- a/ws.go
+++ b/ws.go
@@ -9,6 +9,7 @@ import (
 	"net"
 	"net/http"
 	"net/http/httputil"
+	"strings"
 	"sync"
 	"time"
 
@@ -31,6 +32,7 @@ type WSOptions struct {
 	EnableCompression bool
 	UserAgent         string
 	Path              string
+	Reverse           string
 }
 
 type wsTransporter struct {
@@ -66,6 +68,29 @@ func (tr *wsTransporter) Handshake(conn net.Conn, options ...HandshakeOption) (n
 	return websocketClientConn(url.String(), conn, nil, wsOptions)
 }
 
+func forword(opts *WSOptions) func(w http.ResponseWriter, req *http.Request) {
+	return func(w http.ResponseWriter, req *http.Request) {
+
+		urlStr := opts.Reverse
+		if !strings.HasPrefix(urlStr, "http://") && !strings.HasPrefix(urlStr, "https://") {
+			urlStr = "http://" + urlStr
+		}
+
+		reqURL, err := url.ParseRequestURI(urlStr)
+
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+		}
+
+		req.URL.Host = reqURL.Host
+		req.URL.Scheme = reqURL.Scheme
+		req.Host = reqURL.Host
+
+		httputil.NewSingleHostReverseProxy(reqURL).ServeHTTP(w, req)
+
+	}
+}
+
 type mwsTransporter struct {
 	tcpTransporter
 	options      *WSOptions
@@ -381,6 +406,9 @@ func WSListener(addr string, options *WSOptions) (Listener, error) {
 	}
 	mux := http.NewServeMux()
 	mux.Handle(path, http.HandlerFunc(l.upgrade))
+	if options.Reverse != "" {
+		mux.HandleFunc("/", forword(options))
+	}
 	l.srv = &http.Server{
 		Addr:              addr,
 		Handler:           mux,
@@ -479,6 +507,9 @@ func MWSListener(addr string, options *WSOptions) (Listener, error) {
 
 	mux := http.NewServeMux()
 	mux.Handle(path, http.HandlerFunc(l.upgrade))
+	if options.Reverse != "" {
+		mux.HandleFunc("/", forword(options))
+	}
 	l.srv = &http.Server{
 		Addr:              addr,
 		Handler:           mux,
@@ -604,6 +635,9 @@ func WSSListener(addr string, tlsConfig *tls.Config, options *WSOptions) (Listen
 
 	mux := http.NewServeMux()
 	mux.Handle(path, http.HandlerFunc(l.upgrade))
+	if options.Reverse != "" {
+		mux.HandleFunc("/", forword(options))
+	}
 	l.srv = &http.Server{
 		Addr:              addr,
 		TLSConfig:         tlsConfig,
@@ -670,6 +704,9 @@ func MWSSListener(addr string, tlsConfig *tls.Config, options *WSOptions) (Liste
 
 	mux := http.NewServeMux()
 	mux.Handle(path, http.HandlerFunc(l.upgrade))
+	if options.Reverse != "" {
+		mux.HandleFunc("/", forword(options))
+	}
 	l.srv = &http.Server{
 		Addr:              addr,
 		TLSConfig:         tlsConfig,