ci: add test for websocket

tests/func_test.go

@@ -2,10 +2,12 @@ package tests
 
 import (
 	"fmt"
+	"net/url"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/gorilla/websocket"
 	"github.com/stretchr/testify/assert"
 
 	"github.com/fatedier/frp/client"
@@ -187,6 +189,22 @@ func TestHttp(t *testing.T) {
 	}
 }
 
+func TestWebSocket(t *testing.T) {
+	assert := assert.New(t)
+
+	u := url.URL{Scheme: "ws", Host: fmt.Sprintf("%s:%d", "127.0.0.1", TEST_HTTP_FRP_PORT), Path: "/ws"}
+	c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
+	assert.NoError(err)
+	defer c.Close()
+
+	err = c.WriteMessage(websocket.TextMessage, []byte(TEST_HTTP_NORMAL_STR))
+	assert.NoError(err)
+
+	_, msg, err := c.ReadMessage()
+	assert.NoError(err)
+	assert.Equal(TEST_HTTP_NORMAL_STR, string(msg))
+}
+
 func TestPrivilegeAllowPorts(t *testing.T) {
 	assert := assert.New(t)
 	// Port not allowed

tests/http_server.go

@@ -2,17 +2,43 @@ package tests
 
 import (
 	"fmt"
+	"log"
 	"net/http"
 	"regexp"
 	"strings"
+
+	"github.com/gorilla/websocket"
 )
 
+var upgrader = websocket.Upgrader{}
+
 func StartHttpServer() {
-	http.HandleFunc("/", request)
+	http.HandleFunc("/", handleHttp)
+	http.HandleFunc("/ws", handleWebSocket)
 	http.ListenAndServe(fmt.Sprintf("0.0.0.0:%d", TEST_HTTP_PORT), nil)
 }
 
-func request(w http.ResponseWriter, r *http.Request) {
+func handleWebSocket(w http.ResponseWriter, r *http.Request) {
+	c, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		log.Print("upgrade:", err)
+		return
+	}
+	defer c.Close()
+	for {
+		mt, message, err := c.ReadMessage()
+		if err != nil {
+			break
+		}
+		err = c.WriteMessage(mt, message)
+		if err != nil {
+			log.Println("write:", err)
+			break
+		}
+	}
+}
+
+func handleHttp(w http.ResponseWriter, r *http.Request) {
 	match, err := regexp.Match(`.*\.sub\.com`, []byte(r.Host))
 	if err != nil {
 		w.WriteHeader(500)