goland 实现websocket server的示例代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
package main

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"net/http"
	uuid "github.com/satori/go.uuid"
)

//Client:单个websocket
type Client struct {
	Id      string
	Socket  *websocket.Conn
	Message chan []byte
}

var clientCount   uint // 客户端数量

//从websocket中直接读取数据
func (c *Client) Read() {
	defer func() {
		//客户端关闭
		if err := c.Socket.Close(); err != nil {
			fmt.Printf("client [%s] disconnect err: %s", c.Id, err)
		}
		//关闭后直接注销客户端
		//WebsocketManager.UnRegisterClient(c)
		clientCount--
		fmt.Printf("client [%s],客户端关闭：[%s],Count [%d]n", c.Id, websocket.CloseMessage, clientCount)
	}()

	for {
		messageType, message, err := c.Socket.ReadMessage()
		//读取数据失败
		if err != nil || messageType == websocket.CloseMessage {
			fmt.Printf("client [%s],数据读取失败或通道关闭：[%s],客户端连接状态：[%s]n", c.Id, err.Error(), websocket.CloseMessage)
			break
		}

		// TODO 解析发送过来的参数
		//var data ReadData
		//err = json.Unmarshal(message, &data)
		//if err != nil {
		//  fmt.Println("数据解析失败")
		//	return
		//}

		// TODO 前端请求返回数据到指定客户端

		// 简单测试
		c.Message <- message


	}
}

//写入数据到websocket中
func (c *Client) Write() {
	defer func() {
		//客户端关闭
		if err := c.Socket.Close(); err != nil {
			fmt.Printf("client [%s] disconnect err: %s n", c.Id, err)
			return
		}
		//关闭后直接注销客户端
		//WebsocketManager.UnRegisterClient(c)
		clientCount--
		fmt.Printf("client [%s],客户端关闭：[%s]n", c.Id, websocket.CloseMessage)
	}()

	for {
		select {
		case message, ok := <-c.Message:

			if !ok {
				//数据写入失败，关闭通道
				fmt.Printf("client [%s],客户端连接状态：[%s]n", c.Id, websocket.CloseMessage)
				_ = c.Socket.WriteMessage(websocket.CloseMessage, []byte{})
				//消息通道关闭后直接注销客户端
				return
			}

			err := c.Socket.WriteMessage(websocket.TextMessage, message)
			if err != nil {
				fmt.Printf("client [%s] write message err: %s n", c.Id, err)
				return
			}
		}
	}
}


// 方法二: 通过对象创建  客户端连接
func WsClient(context *gin.Context) {
	upGrande := websocket.Upgrader{
		//设置允许跨域
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
		//设置请求协议
		Subprotocols: []string{context.GetHeader("Sec-WebSocket-Protocol")},
	}
	//创建连接
	conn, err := upGrande.Upgrade(context.Writer, context.Request, nil)
	if err != nil {
		fmt.Printf("websocket connect error: %s", context.Param("channel"))
		//format.NewResponseJson(context).Error(51001)
		return
	}
	//生成唯一标识client_id
	var uuid = uuid.NewV4().String()
	client := &Client{
		Id:      uuid,
		Socket:  conn,
		Message: make(chan []byte, 1024),
	}
	//注册
	//ws.WebsocketManager.RegisterClient(client)
	clientCount++

	//起协程，实时接收和回复数据
	go client.Read()
	go client.Write()
}

// 方法一: 直接创建客户端
func NewConnection(c *gin.Context) {
	// 定义同源检查，这里只作简单试验不校验
	upGrader := websocket.Upgrader{
	   CheckOrigin: func(r *http.Request) bool {
	       return true
	   },
	}
	 ws, err := upGrader.Upgrade(c.Writer, c.Request, nil)
	//ws, err := websocket.Upgrade(c.Writer, c.Request, nil, 1024, 1024)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{
			"msg": "服务端错误",
		})
		return
	}


	var message = make(chan []byte)

	go func() {
		defer ws.Close()

		for {
			fmt.Println("start transfer message")
			msgType, msg, err := ws.ReadMessage()
			if err != nil || websocket.CloseMessage == msgType {
				fmt.Println("webSocket read error")
				return
			}
			message <- msg
		}

	}()

	go func() {
		defer ws.Close()

		for {
			mm, ok := <- message

			if !ok {
				//数据写入失败，关闭通道
				fmt.Printf("客户端连接状态：[%s]n", websocket.CloseMessage)
				_ = ws.WriteMessage(websocket.CloseMessage, []byte{})
				//消息通道关闭后直接注销客户端
				return
			}

			err := ws.WriteMessage(websocket.TextMessage, mm)
			if err != nil {
				fmt.Println("webSocket write error")
				return
			}
		}
	}()

	//for {
	//
	//}

	 开始通信
	//for {
	//	fmt.Println("start transfer message")
	//	msgType, msg, err := ws.ReadMessage()
	//	if err != nil {
	//		fmt.Println("webSocket read error")
	//		return
	//	}
	//	err = ws.WriteMessage(msgType, msg)
	//	if err != nil {
	//		fmt.Println("webSocket write error")
	//		return
	//	}
	//}
}

// ws://127.0.0.1:9090/wsTest
func main() {
	r := gin.Default()

	// 方法一: 直接创建客户端
	r.GET("/wsTest", NewConnection)
	// 方法二: 通过对象创建  客户端连接
	r.GET("/wsTest1", WsClient)

	clientCount = 0
	r.Run("127.0.0.1:9090")
}