Websocket 增加 http-header 设置功能

- support set http header use websocket

server/client/websocket_client.go

@@ -16,10 +16,11 @@ const (
 
 // WebSocket webSocket
 type WebSocket struct {
-	conn    *websocket.Conn
-	URLLink string
-	URL     *url.URL
-	IsSsl   bool
+	conn       *websocket.Conn
+	URLLink    string
+	URL        *url.URL
+	IsSsl      bool
+	HttpHeader map[string]string
 }
 
 // NewWebSocket new
@@ -34,9 +35,10 @@ func NewWebSocket(urlLink string) (ws *WebSocket) {
 		panic(err)
 	}
 	ws = &WebSocket{
-		URLLink: urlLink,
-		URL:     u,
-		IsSsl:   isSsl,
+		URLLink:    urlLink,
+		URL:        u,
+		IsSsl:      isSsl,
+		HttpHeader: make(map[string]string),
 	}
 	return
 }
@@ -46,6 +48,10 @@ func (w *WebSocket) getLink() (link string) {
 	return w.URLLink
 }
 
+func (w *WebSocket) SetHeader(head map[string]string) {
+	w.HttpHeader = head
+}
+
 // getOrigin 获取源连接
 func (w *WebSocket) getOrigin() (origin string) {
 	origin = "http://"
@@ -67,14 +73,32 @@ func (w *WebSocket) Close() (err error) {
 	return w.conn.Close()
 }
 
+// Dial opens a new client connection to a WebSocket.
+// 复写 websocket库的 Dial 方法 ,增加 httpheader 设置功能
+func Dial(url_, protocol, origin string, httpHeader map[string]string) (ws *websocket.Conn, err error) {
+	config, err := websocket.NewConfig(url_, origin)
+	config.Header = map[string][]string{}
+	//
+	for x := range httpHeader {
+		config.Header.Set(x, httpHeader[x])
+	}
+	if err != nil {
+		return nil, err
+	}
+	if protocol != "" {
+		config.Protocol = []string{protocol}
+	}
+	return websocket.DialConfig(config)
+}
+
 // GetConn 获取连接
 func (w *WebSocket) GetConn() (err error) {
 	var (
 		conn *websocket.Conn
 		i    int
 	)
 	for i = 0; i < connRetry; i++ {
-		conn, err = websocket.Dial(w.getLink(), "", w.getOrigin())
+		conn, err = Dial(w.getLink(), "", w.getOrigin(), w.HttpHeader)
 		if err != nil {
 			fmt.Println("GetConn 建立连接失败 in...", i, err)
 			continue

server/dispose.go

@@ -50,6 +50,7 @@ func Dispose(ctx context.Context, concurrency, totalNumber uint64, request *mode
 			case 1:
 				// 连接以后再启动协程
 				ws := client.NewWebSocket(request.URL)
+				ws.SetHeader(request.Headers)
 				err := ws.GetConn()
 				if err != nil {
 					fmt.Println("连接失败:", i, err)
@@ -61,6 +62,7 @@ func Dispose(ctx context.Context, concurrency, totalNumber uint64, request *mode
 				go func(i uint64) {
 					// 连接以后再启动协程
 					ws := client.NewWebSocket(request.URL)
+					ws.SetHeader(request.Headers)
 					err := ws.GetConn()
 					if err != nil {
 						fmt.Println("连接失败:", i, err)