From d4af916f60ed0c4882448d0c4cd7b60ce8998544 Mon Sep 17 00:00:00 2001
From: lwch <lwch748@gmail.com>
Date: Fri, 22 Jul 2022 14:33:27 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0code-server=E5=8F=8D=E4=BB=A3?=
 =?UTF-8?q?=E7=A4=BA=E4=BE=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 test/code-server-forward/main.go | 114 +++++++++++++++++++++++++++++++
 1 file changed, 114 insertions(+)
 create mode 100644 test/code-server-forward/main.go

diff --git a/test/code-server-forward/main.go b/test/code-server-forward/main.go
new file mode 100644
index 0000000..c92f2a3
--- /dev/null
+++ b/test/code-server-forward/main.go
@@ -0,0 +1,114 @@
+package main
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+	"sync"
+
+	"github.com/gorilla/websocket"
+	"github.com/lwch/runtime"
+)
+
+var cli = &http.Client{}
+var upgrader = websocket.Upgrader{}
+var dialer = websocket.Dialer{}
+
+func main() {
+	http.HandleFunc("/", next)
+	http.ListenAndServe(":8001", nil)
+}
+
+func normal(w http.ResponseWriter, r *http.Request) {
+	u := r.URL
+	u.Scheme = "http"
+	u.Host = "127.0.0.1:8000"
+	req, err := http.NewRequest(r.Method, u.String(), r.Body)
+	runtime.Assert(err)
+
+	for key, values := range r.Header {
+		for _, v := range values {
+			req.Header.Add(key, v)
+		}
+	}
+
+	rep, err := cli.Do(req)
+	runtime.Assert(err)
+	defer rep.Body.Close()
+
+	for key, values := range rep.Header {
+		for _, v := range values {
+			w.Header().Add(key, v)
+		}
+	}
+
+	w.WriteHeader(rep.StatusCode)
+
+	_, err = io.Copy(w, rep.Body)
+	runtime.Assert(err)
+}
+
+func ws(w http.ResponseWriter, r *http.Request) {
+	u := r.URL
+	u.Scheme = "ws"
+	u.Host = "127.0.0.1:8000"
+
+	hdr := make(http.Header)
+	for key, values := range r.Header {
+		if strings.HasPrefix(key, "Sec-") {
+			continue
+		}
+		for _, value := range values {
+			hdr.Add(key, value)
+		}
+	}
+
+	hdr.Del("Connection")
+	hdr.Del("Upgrade")
+
+	remote, resp, err := dialer.Dial(u.String(), hdr)
+	runtime.Assert(err)
+	defer resp.Body.Close()
+	defer remote.Close()
+
+	local, err := upgrader.Upgrade(w, r, nil)
+	runtime.Assert(err)
+	defer local.Close()
+
+	cp := func(wg *sync.WaitGroup, dst, src *websocket.Conn) {
+		defer wg.Done()
+		defer dst.Close()
+		defer src.Close()
+		for {
+			t, data, err := src.ReadMessage()
+			if err != nil {
+				fmt.Println(err)
+				return
+			}
+			err = dst.WriteMessage(t, data)
+			if err != nil {
+				fmt.Println(err)
+				return
+			}
+		}
+	}
+
+	var wg sync.WaitGroup
+
+	wg.Add(2)
+	go cp(&wg, local, remote)
+	go cp(&wg, remote, local)
+
+	wg.Wait()
+}
+
+func next(w http.ResponseWriter, r *http.Request) {
+	upgrade := r.Header.Get("Connection")
+
+	if upgrade == "Upgrade" {
+		ws(w, r)
+		return
+	}
+	normal(w, r)
+}