try to support XTLS/Xray-core#375

wwqgtxx · Mar 14, 2021 · ec73dab · ec73dab
1 parent c51d558
commit ec73dab
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 5 deletions.
diff --git a/client.go b/client.go
@@ -72,7 +72,7 @@ func (c *client) SetClientImpl(impl ClientImpl) {
 
 type ClientImpl interface {
 	Target() string
-	Dial() (io.Closer, error)
+	Dial(args ...interface{}) (io.Closer, error)
 	ToRawConn(conn io.Closer) net.Conn
 	Tunnel(tcp net.Conn, conn io.Closer)
 }
@@ -92,8 +92,17 @@ func (c *wsClientImpl) Target() string {
 	return c.wsUrl
 }
 
-func (c *wsClientImpl) Dial() (io.Closer, error) {
-	ws, _, err := c.wsDialer.Dial(c.Target(), c.header)
+func (c *wsClientImpl) Dial(args ...interface{}) (io.Closer, error) {
+	header := c.header
+	if len(args) >= 1 {
+		if inHeader, ok := args[0].(http.Header); ok {
+			if secProtocol, ok := inHeader["Sec-WebSocket-Protocol"]; ok && len(secProtocol) > 0 {
+				header = header.Clone()
+				header["Sec-WebSocket-Protocol"] = secProtocol
+			}
+		}
+	}
+	ws, _, err := c.wsDialer.Dial(c.Target(), header)
 	return ws, err
 }
 
@@ -110,7 +119,7 @@ func (c *tcpClientImpl) Target() string {
 	return c.targetAddress
 }
 
-func (c *tcpClientImpl) Dial() (io.Closer, error) {
+func (c *tcpClientImpl) Dial(args ...interface{}) (io.Closer, error) {
 	return c.tcpDialer.Dial("tcp", c.Target())
 }
 

diff --git a/server.go b/server.go
@@ -90,7 +90,7 @@ func (s *internalServerHandler) ServeHTTP(w http.ResponseWriter, r *http.Request
 	defer ws.Close()
 	source := ws.UnderlyingConn()
 
-	ws2, err := s.Client.Dial()
+	ws2, err := s.Client.Dial(r.Header)
 	if err != nil {
 		log.Println(err)
 		return