wsv2.go 9.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393
  1. package wsv2
  2. import (
  3. "crypto/rand"
  4. "encoding/binary"
  5. "errors"
  6. "fmt"
  7. "log"
  8. "net"
  9. "net/http"
  10. "net/url"
  11. "strings"
  12. "time"
  13. "git.me9.top/git/tinymq/config"
  14. "git.me9.top/git/tinymq/conn"
  15. "git.me9.top/git/tinymq/conn/util"
  16. "github.com/gorilla/websocket"
  17. )
  18. const VERSION uint8 = 2
  19. const PROTO string = "ws"
  20. const PROTO_STL string = "wss"
  21. type WsConnectV2 struct {
  22. cf *config.Config
  23. conn *websocket.Conn
  24. cipher *util.Cipher // 记录当前的加解密类,可以保证在没有ssl的情况下数据安全
  25. }
  26. var upgrader = websocket.Upgrader{} // use default options
  27. // websocket 服务
  28. // 如果有绑定参数,则进行绑定操作代码
  29. func Server(cf *config.Config, bind string, path string, hash string, fn conn.ServerConnectFunc) (err error) {
  30. var ci *util.CipherInfo
  31. var encryptKey string
  32. if hash != "" {
  33. i := strings.Index(hash, ":")
  34. if i <= 0 {
  35. return errors.New("hash is invalid")
  36. }
  37. encryptMethod := hash[0:i]
  38. encryptKey = hash[i+1:]
  39. if c, ok := util.CipherMethod[encryptMethod]; ok {
  40. ci = c
  41. } else {
  42. return errors.New("Unsupported encryption method: " + encryptMethod)
  43. }
  44. }
  45. http.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
  46. conn, err := upgrader.Upgrade(w, r, nil)
  47. if err != nil {
  48. log.Println("[wsv2 Server Upgrade ERROR]", err)
  49. return
  50. }
  51. if ci == nil {
  52. ws := &WsConnectV2{
  53. cf: cf,
  54. conn: conn,
  55. }
  56. fn(ws)
  57. return
  58. }
  59. var eiv []byte
  60. var div []byte
  61. if ci.IvLen > 0 {
  62. // 服务端 IV
  63. eiv = make([]byte, ci.IvLen)
  64. _, err = rand.Read(eiv)
  65. if err != nil {
  66. log.Println("[wsv2 Server rand.Read ERROR]", err)
  67. return
  68. }
  69. // 发送 IV
  70. conn.SetWriteDeadline(time.Now().Add(time.Duration(cf.WriteWait) * time.Millisecond))
  71. if err := conn.WriteMessage(websocket.BinaryMessage, eiv); err != nil {
  72. log.Println("[wsv2 Server conn.Write ERROR]", err)
  73. return
  74. }
  75. // 读取 IV
  76. err = conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(cf.ReadWait)))
  77. if err != nil {
  78. log.Println("[wsv2 Server SetReadDeadline ERROR]", err)
  79. return
  80. }
  81. _, div, err = conn.ReadMessage()
  82. if err != nil {
  83. log.Println("[wsv2 Server ReadFull ERROR]", err)
  84. return
  85. }
  86. }
  87. cipher, err := util.NewCipher(ci, encryptKey, eiv, div)
  88. if err != nil {
  89. log.Println("[wsv2 NewCipher ERROR]", err)
  90. return
  91. }
  92. ws := &WsConnectV2{
  93. cf: cf,
  94. conn: conn,
  95. cipher: cipher,
  96. }
  97. fn(ws)
  98. })
  99. if bind != "" {
  100. go func() (err error) {
  101. defer func() {
  102. if err != nil {
  103. log.Fatal(err)
  104. }
  105. }()
  106. log.Printf("Listening and serving Websocket on %s\n", bind)
  107. // 暂时使用全局的方式,后面有需求再修改
  108. // 而且还没有 https 方式的绑定
  109. // 需要在前端增加其他的服务进行转换
  110. err = http.ListenAndServe(bind, nil)
  111. return
  112. }()
  113. }
  114. return
  115. }
  116. // 客户端,新建一个连接
  117. func Client(cf *config.Config, scheme string, addr string, path string, hash string) (conn.Connect, error) {
  118. u := url.URL{Scheme: scheme, Host: addr, Path: path}
  119. // 没有加密的情况
  120. if hash == "" {
  121. conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
  122. if err != nil {
  123. return nil, err
  124. }
  125. ws := &WsConnectV2{
  126. cf: cf,
  127. conn: conn,
  128. }
  129. return ws, nil
  130. }
  131. i := strings.Index(hash, ":")
  132. if i <= 0 {
  133. return nil, errors.New("hash is invalid")
  134. }
  135. encryptMethod := hash[0:i]
  136. encryptKey := hash[i+1:]
  137. ci, ok := util.CipherMethod[encryptMethod]
  138. if !ok {
  139. return nil, errors.New("Unsupported encryption method: " + encryptMethod)
  140. }
  141. conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
  142. if err != nil {
  143. return nil, err
  144. }
  145. var eiv []byte
  146. var div []byte
  147. if ci.IvLen > 0 {
  148. // 客户端 IV
  149. eiv = make([]byte, ci.IvLen)
  150. _, err = rand.Read(eiv)
  151. if err != nil {
  152. log.Println("[wsv2 Client rand.Read ERROR]", err)
  153. return nil, err
  154. }
  155. // 发送 IV
  156. conn.SetWriteDeadline(time.Now().Add(time.Duration(cf.WriteWait) * time.Millisecond))
  157. if err := conn.WriteMessage(websocket.BinaryMessage, eiv); err != nil {
  158. log.Println("[wsv2 Client conn.Write ERROR]", err)
  159. return nil, err
  160. }
  161. // 读取 IV
  162. err = conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(cf.ReadWait)))
  163. if err != nil {
  164. log.Println("[wsv2 Client SetReadDeadline ERROR]", err)
  165. return nil, err
  166. }
  167. _, div, err = conn.ReadMessage()
  168. if err != nil {
  169. log.Println("[wsv2 Client ReadFull ERROR]", err)
  170. return nil, err
  171. }
  172. }
  173. cipher, err := util.NewCipher(ci, encryptKey, eiv, div)
  174. if err != nil {
  175. log.Println("[wsv2 NewCipher ERROR]", err)
  176. return nil, err
  177. }
  178. ws := &WsConnectV2{
  179. cf: cf,
  180. conn: conn,
  181. cipher: cipher,
  182. }
  183. return ws, nil
  184. }
  185. // 发送数据到网络
  186. // 如果有加密函数的话会直接修改源数据
  187. func (c *WsConnectV2) writeMessage(buf []byte) (err error) {
  188. if c.cipher != nil {
  189. c.cipher.Encrypt(buf, buf)
  190. }
  191. c.conn.SetWriteDeadline(time.Now().Add(time.Millisecond * time.Duration(c.cf.WriteWait)))
  192. return c.conn.WriteMessage(websocket.BinaryMessage, buf)
  193. }
  194. // 发送Auth信息
  195. // 建立连接后第一个发送的消息
  196. func (c *WsConnectV2) WriteAuthInfo(channel string, auth []byte) (err error) {
  197. protoLen := len(PROTO)
  198. channelLen := len(channel)
  199. if channelLen > 0xFFFF {
  200. return errors.New("length of channel over")
  201. }
  202. dlen := 2 + 1 + 1 + protoLen + 2 + channelLen + len(auth)
  203. start := 0
  204. buf := make([]byte, dlen)
  205. binary.BigEndian.PutUint16(buf[start:start+2], config.ID_AUTH)
  206. start += 2
  207. buf[start] = VERSION
  208. start++
  209. buf[start] = byte(protoLen)
  210. start++
  211. copy(buf[start:start+protoLen], []byte(PROTO))
  212. start += protoLen
  213. binary.BigEndian.PutUint16(buf[start:start+2], uint16(channelLen))
  214. start += 2
  215. copy(buf[start:start+channelLen], []byte(channel))
  216. start += channelLen
  217. copy(buf[start:], auth)
  218. return c.writeMessage(buf)
  219. }
  220. // 获取Auth信息
  221. // id(uint16)+version(uint8)+proto(string)+channel(string)+auth([]byte)
  222. func (c *WsConnectV2) ReadAuthInfo() (proto string, version uint8, channel string, auth []byte, err error) {
  223. defer func() {
  224. if r := recover(); r != nil {
  225. err = fmt.Errorf("recovered from panic: %v", r)
  226. return
  227. }
  228. }()
  229. err = c.conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(c.cf.ReadWait)))
  230. if err != nil {
  231. return
  232. }
  233. _, msg, err := c.conn.ReadMessage()
  234. if err != nil {
  235. return
  236. }
  237. msgLen := len(msg)
  238. if msgLen < 4 {
  239. err = errors.New("message length less than 4")
  240. return
  241. }
  242. // 将读出来的数据进行解密
  243. if c.cipher != nil {
  244. c.cipher.Decrypt(msg, msg)
  245. }
  246. start := 0
  247. id := binary.BigEndian.Uint16(msg[start : start+2])
  248. if id != config.ID_AUTH {
  249. err = fmt.Errorf("wrong message id: %d", id)
  250. return
  251. }
  252. start += 2
  253. version = msg[start]
  254. if version != VERSION {
  255. err = fmt.Errorf("require version %d, get version: %d", VERSION, version)
  256. return
  257. }
  258. start++
  259. protoLen := int(msg[start])
  260. if protoLen < 2 {
  261. err = errors.New("wrong proto length")
  262. return
  263. }
  264. start++
  265. proto = string(msg[start : start+protoLen])
  266. if proto != PROTO {
  267. err = fmt.Errorf("wrong proto: %s", proto)
  268. return
  269. }
  270. start += protoLen
  271. channelLen := int(binary.BigEndian.Uint16(msg[start : start+2]))
  272. if channelLen < 2 {
  273. err = errors.New("wrong channel length")
  274. return
  275. }
  276. start += 2
  277. channel = string(msg[start : start+channelLen])
  278. start += channelLen
  279. auth = msg[start:]
  280. return
  281. }
  282. // 发送请求数据包到网络
  283. func (c *WsConnectV2) WriteRequest(id uint16, cmd string, data []byte) error {
  284. // 为了区分请求还是响应包,命令字符串不能超过127个字节,如果超过则报错
  285. cmdLen := len(cmd)
  286. if cmdLen > 0x7F {
  287. return errors.New("length of command more than 0x7F")
  288. }
  289. dlen := 2 + 1 + cmdLen + len(data)
  290. buf := make([]byte, dlen) // 申请内存
  291. binary.BigEndian.PutUint16(buf[0:2], id)
  292. buf[2] = byte(cmdLen)
  293. copy(buf[3:], cmd)
  294. copy(buf[3+cmdLen:], data)
  295. return c.writeMessage(buf)
  296. }
  297. // 发送响应数据包到网络
  298. // 网络格式:[id, stateCode, data]
  299. func (c *WsConnectV2) WriteResponse(id uint16, state uint8, data []byte) error {
  300. dlen := 2 + 1 + len(data)
  301. buf := make([]byte, dlen)
  302. binary.BigEndian.PutUint16(buf[0:2], id)
  303. buf[2] = state | 0x80
  304. copy(buf[3:], data)
  305. return c.writeMessage(buf)
  306. }
  307. // 发送ping包
  308. func (c *WsConnectV2) WritePing(id uint16) error {
  309. buf := make([]byte, 2)
  310. binary.BigEndian.PutUint16(buf[0:2], id)
  311. return c.writeMessage(buf)
  312. }
  313. // 获取信息
  314. func (c *WsConnectV2) ReadMessage(deadline int) (msgType conn.MsgType, id uint16, cmd string, state uint8, data []byte, err error) {
  315. err = c.conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(deadline)))
  316. if err != nil {
  317. return
  318. }
  319. _, msg, err := c.conn.ReadMessage()
  320. if err != nil {
  321. return
  322. }
  323. msgLen := len(msg)
  324. if msgLen < 2 {
  325. err = errors.New("message length less than 2")
  326. return
  327. }
  328. // 将读出来的数据进行解密
  329. if c.cipher != nil {
  330. c.cipher.Decrypt(msg, msg)
  331. }
  332. id = binary.BigEndian.Uint16(msg[0:2])
  333. // ping信息
  334. if msgLen == 2 {
  335. msgType = conn.PingMsg
  336. return
  337. }
  338. if id > config.ID_MAX {
  339. err = fmt.Errorf("wrong message id: %d", id)
  340. return
  341. }
  342. cmdx := msg[2]
  343. if (cmdx & 0x80) == 0 {
  344. // 请求包
  345. msgType = conn.RequestMsg
  346. cmdLen := int(cmdx)
  347. cmd = string(msg[3 : cmdLen+3])
  348. data = msg[cmdLen+3:]
  349. return
  350. } else {
  351. // 响应数据包
  352. msgType = conn.ResponseMsg
  353. state = cmdx & 0x7F
  354. data = msg[3:]
  355. return
  356. }
  357. }
  358. // 获取远程的地址
  359. func (c *WsConnectV2) RemoteAddr() net.Addr {
  360. return c.conn.RemoteAddr()
  361. }
  362. // 获取本地的地址
  363. func (c *WsConnectV2) LocalAddr() net.Addr {
  364. return c.conn.LocalAddr()
  365. }
  366. func (c *WsConnectV2) Close() error {
  367. return c.conn.Close()
  368. }