hub.go 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951
  1. package tinymq
  2. import (
  3. "context"
  4. "encoding/json"
  5. "errors"
  6. "fmt"
  7. "log"
  8. "math/rand"
  9. "net"
  10. // "regexp"
  11. "strconv"
  12. "strings"
  13. "sync"
  14. "time"
  15. "git.me9.top/git/tinymq/config"
  16. "git.me9.top/git/tinymq/conn"
  17. "git.me9.top/git/tinymq/conn/tcp2"
  18. "git.me9.top/git/tinymq/conn/ws2"
  19. )
  20. // 类似一个插座的功能,管理多个连接
  21. // 一个hub即可以是客户端,同时也可以是服务端
  22. // 为了简化流程和让通讯更加迅速,不再重发和缓存结果,采用超时的方式告诉应用层。
  23. // 截取部分字符串
  24. func subStr(str string, length int) string {
  25. if len(str) <= length {
  26. return str
  27. }
  28. return str[0:length] + "..."
  29. }
  30. type Hub struct {
  31. sync.Mutex
  32. ctx context.Context // 为了方便退出而建立
  33. cancel context.CancelFunc
  34. cf *config.Config
  35. globalID uint16
  36. channel string // 本地频道信息
  37. middle []MiddleFunc // 中间件
  38. // connects sync.Map // map[*Line]bool(true) //记录当前的连接,方便查找
  39. lines *Mapx // 记录当前的连接,统一管理
  40. subscribes sync.Map // [cmd]->[]*SubscribeData //注册绑定频道的函数,用于响应请求
  41. msgCache sync.Map // map[uint16]*GetMsg //请求的回应记录,key为id
  42. // 客户端需要用的函数(服务端可为空)
  43. connectHostFunc ConnectHostFunc // 获取对应频道的一个连接地址
  44. // 返回认证信息,发送到对方
  45. authFunc AuthFunc // 获取认证信息,用于发送给对方
  46. // 核对发送过来的认证信息
  47. checkAuthFunc CheckAuthFunc // 核对认证是否合法
  48. // 连接状态变化时调用的函数
  49. connectStatusFunc ConnectStatusFunc
  50. // 验证发送数据的条件是否满足 (可空)
  51. checkConnectOkFunc CheckConnectOkFunc
  52. // 通过过滤函数获取一个频道信息 (可空)
  53. filterToChannelFunc FilterToChannelFunc
  54. // 上次清理异常连接时间戳
  55. lastCleanDeadConnect int64
  56. }
  57. func (h *Hub) SetCheckConnectOkFunc(fn CheckConnectOkFunc) {
  58. h.checkConnectOkFunc = fn
  59. }
  60. func (h *Hub) SetFilterToChannelFunc(fn FilterToChannelFunc) {
  61. h.filterToChannelFunc = fn
  62. }
  63. // 转换数据
  64. func (h *Hub) convertData(data any) (reqData []byte, err error) {
  65. switch data := data.(type) {
  66. case []byte:
  67. reqData = data
  68. case string:
  69. reqData = []byte(data)
  70. case func() ([]byte, error):
  71. reqData, err = data()
  72. if err != nil {
  73. log.Println(err.Error())
  74. return nil, err
  75. }
  76. default:
  77. if data != nil {
  78. // 自动转换数据为json格式
  79. reqData, err = json.Marshal(data)
  80. if err != nil {
  81. log.Println(err.Error())
  82. return nil, err
  83. }
  84. }
  85. }
  86. return
  87. }
  88. // 清理异常连接
  89. func (h *Hub) cleanDeadConnect() {
  90. now := time.Now().UnixMilli()
  91. expired := now - int64(h.cf.CleanDeadConnectWait)
  92. if h.lastCleanDeadConnect < expired {
  93. h.lastCleanDeadConnect = now
  94. h.lines.DeleteInvalidLines(expired)
  95. }
  96. }
  97. // 获取通讯消息ID号
  98. func (h *Hub) GetID() uint16 {
  99. h.Lock()
  100. defer h.Unlock()
  101. h.globalID++
  102. if h.globalID <= 0 || h.globalID >= config.ID_MAX {
  103. h.globalID = 1
  104. }
  105. for {
  106. // 检查是否在请求队列中存在对应的id
  107. if _, ok := h.msgCache.Load(h.globalID); ok {
  108. h.globalID++
  109. if h.globalID <= 0 || h.globalID >= config.ID_MAX {
  110. h.globalID = 1
  111. }
  112. } else {
  113. break
  114. }
  115. }
  116. return h.globalID
  117. }
  118. // 添加中间件
  119. // 如果中间件函数返回为空,表示处理完成,通过
  120. // 如果中间件函数返回 NEXT_MIDDLE,表示需要下一个中间件函数处理;如果没有下一函数则默认通过
  121. func (h *Hub) UseMiddle(middleFunc MiddleFunc) {
  122. h.middle = append(h.middle, middleFunc)
  123. }
  124. // 注册频道,其中频道为正则表达式字符串
  125. func (h *Hub) Subscribe(filter FilterFunc, cmd string, backFunc SubscribeBackFunc) (err error) {
  126. if filter == nil {
  127. return errors.New("filter function can not be nil")
  128. }
  129. reg := &SubscribeData{
  130. Filter: filter,
  131. Cmd: cmd,
  132. BackFunc: backFunc,
  133. }
  134. sub, ok := h.subscribes.Load(cmd)
  135. if ok {
  136. h.subscribes.Store(cmd, append(sub.([]*SubscribeData), reg))
  137. return
  138. }
  139. regs := make([]*SubscribeData, 1)
  140. regs[0] = reg
  141. h.subscribes.Store(cmd, regs)
  142. return
  143. }
  144. // 遍历频道列表
  145. // 如果 fn 返回 false,则 range 停止迭代
  146. func (h *Hub) ConnectRange(fn func(id int, line *Line) bool) {
  147. h.lines.Range(fn)
  148. }
  149. // 获取当前在线的数量
  150. func (h *Hub) ConnectNum() int {
  151. var count int
  152. h.lines.Range(func(id int, line *Line) bool {
  153. if line.state == Connected {
  154. count++
  155. }
  156. return true
  157. })
  158. return count
  159. }
  160. // 获取所有的在线连接频道
  161. func (h *Hub) AllChannel() []string {
  162. cs := make([]string, 0)
  163. h.lines.Range(func(id int, line *Line) bool {
  164. if line.state == Connected {
  165. cs = append(cs, line.channel)
  166. }
  167. return true
  168. })
  169. return cs
  170. }
  171. // 获取所有连接频道和连接时长
  172. // 为了避免定义数据结构麻烦,采用|隔开, 频道名|连接开始时间
  173. func (h *Hub) AllChannelWithStarted() []string {
  174. cs := make([]string, 0)
  175. h.lines.Range(func(id int, line *Line) bool {
  176. if line.state == Connected {
  177. cs = append(cs, fmt.Sprintf("%s|%d", line.channel, line.updated.UnixMilli()))
  178. }
  179. return true
  180. })
  181. return cs
  182. }
  183. // 获取频道并通过函数过滤,如果返回 false 将终止
  184. func (h *Hub) ChannelToFunc(fn func(string) bool) {
  185. h.lines.Range(func(id int, line *Line) bool {
  186. if line.state == Connected {
  187. return fn(line.channel)
  188. }
  189. return true
  190. })
  191. }
  192. // 从 channel 获取连接
  193. func (h *Hub) ChannelToLine(channel string) (line *Line) {
  194. h.lines.Range(func(id int, l *Line) bool {
  195. if l.IsChannelEqual(channel) {
  196. line = l
  197. return false
  198. }
  199. return true
  200. })
  201. return
  202. }
  203. // 返回请求结果
  204. func (h *Hub) outResponse(response *ResponseData) {
  205. defer recover() //避免管道已经关闭而引起panic
  206. id := response.Id
  207. t, ok := h.msgCache.Load(id)
  208. if ok {
  209. // 删除数据缓存
  210. h.msgCache.Delete(id)
  211. gm := t.(*GetMsg)
  212. // 停止定时器
  213. if !gm.timer.Stop() {
  214. select {
  215. case <-gm.timer.C:
  216. default:
  217. }
  218. }
  219. // 回应数据到上层
  220. gm.out <- response
  221. }
  222. }
  223. // 发送数据到网络接口
  224. // 返回发送的数量
  225. func (h *Hub) sendRequest(gd *GetData) (count int, err error) {
  226. outData, err := h.convertData(gd.Data)
  227. if err != nil {
  228. log.Println(err)
  229. return 0, err
  230. }
  231. doit := func(_ int, line *Line) bool {
  232. // 检查连接是否OK
  233. if line.state != Connected {
  234. return true
  235. }
  236. // 验证连接是否达到发送数据的要求
  237. if h.checkConnectOkFunc != nil && !h.checkConnectOkFunc(line, gd) {
  238. return true
  239. }
  240. if gd.Filter(line) {
  241. var id uint16
  242. if gd.backchan != nil {
  243. id = h.GetID()
  244. timeout := gd.Timeout
  245. if timeout <= 0 {
  246. timeout = h.cf.WriteWait
  247. }
  248. fn := func(id uint16, conn *Line) func() {
  249. return func() {
  250. go h.outResponse(&ResponseData{
  251. Id: id,
  252. State: config.GET_TIMEOUT,
  253. Data: fmt.Appendf(nil, "[%s] %s %s", IdMsg(config.GET_TIMEOUT), conn.channel, gd.Cmd),
  254. conn: conn,
  255. })
  256. // 检查是否已经很久时间没有使用连接了
  257. if time.Since(conn.lastRead) > time.Duration(h.cf.PingInterval*3*int(time.Millisecond)) {
  258. // 超时关闭当前的连接
  259. log.Println("get message timeout", conn.channel)
  260. // 有可能连接出现问题,断开并重新连接
  261. conn.Close(false)
  262. return
  263. }
  264. }
  265. }(id, line)
  266. // 将要发送的请求缓存
  267. gm := &GetMsg{
  268. out: gd.backchan,
  269. timer: time.AfterFunc(time.Millisecond*time.Duration(timeout), fn),
  270. }
  271. h.msgCache.Store(id, gm)
  272. }
  273. // 组织数据并发送到Connect
  274. line.sendRequest <- &RequestData{
  275. Id: id,
  276. Cmd: gd.Cmd,
  277. Data: outData,
  278. timeout: gd.Timeout,
  279. backchan: gd.backchan,
  280. conn: line,
  281. }
  282. if h.cf.PrintMsg {
  283. log.Println("[SEND]->", id, line.channel, "["+gd.Cmd+"]", subStr(string(outData), 200))
  284. }
  285. count++
  286. if gd.Max > 0 && count >= gd.Max {
  287. return false
  288. }
  289. }
  290. return true
  291. }
  292. // 如果没有发送到消息,延时重连直到超时
  293. for i := 0; i <= gd.Timeout; i += 500 {
  294. if gd.Rand {
  295. h.lines.RandRange(doit, i == 0)
  296. } else {
  297. h.lines.Range(doit)
  298. }
  299. if count > 0 {
  300. break
  301. }
  302. // 如果是客户端,并且有机会自动连接,则尝试自动连接
  303. if i == 0 && h.connectHostFunc != nil && h.filterToChannelFunc != nil {
  304. channel := h.filterToChannelFunc(gd.Filter)
  305. if channel == "" {
  306. err = errors.New(IdMsg(NO_MATCH_CONNECT))
  307. log.Println("not channel found")
  308. return 0, err
  309. }
  310. err := h.ConnectToServer(channel, false, nil, false)
  311. if err != nil {
  312. log.Println(err)
  313. return 0, err
  314. }
  315. } else {
  316. time.Sleep(time.Millisecond * 400) // 故意将时间缩小一点
  317. }
  318. }
  319. return
  320. }
  321. // 执行网络发送过来的命令
  322. func (h *Hub) requestFromNet(request *RequestData) {
  323. cmd := request.Cmd
  324. channel := request.conn.channel
  325. if h.cf.PrintMsg {
  326. log.Println("[REQU]<-", request.Id, channel, "["+cmd+"]", subStr(string(request.Data), 200))
  327. }
  328. // 执行中间件
  329. for _, mdFunc := range h.middle {
  330. rsp := mdFunc(request)
  331. if rsp != nil {
  332. // NEXT_MIDDLE 表示当前的函数没有处理完成,还需要下个中间件处理
  333. if rsp.State == config.NEXT_MIDDLE {
  334. continue
  335. }
  336. // 返回消息
  337. if request.Id != 0 {
  338. rsp.Id = request.Id
  339. request.conn.sendResponse <- rsp
  340. }
  341. return
  342. } else {
  343. break
  344. }
  345. }
  346. sub, ok := h.subscribes.Load(cmd)
  347. if ok {
  348. subs := sub.([]*SubscribeData)
  349. // 倒序查找是为了新增的频道响应函数优先执行
  350. for i := len(subs) - 1; i >= 0; i-- {
  351. rg := subs[i]
  352. if rg.Filter(request.conn) {
  353. state, data := rg.BackFunc(request)
  354. // NEXT_SUBSCRIBE 表示当前的函数没有处理完成,还需要下个注册函数处理
  355. if state == config.NEXT_SUBSCRIBE {
  356. continue
  357. }
  358. var byteData []byte
  359. switch data := data.(type) {
  360. case []byte:
  361. byteData = data
  362. case string:
  363. byteData = []byte(data)
  364. default:
  365. if data != nil {
  366. // 自动转换数据为json格式
  367. var err error
  368. byteData, err = json.Marshal(data)
  369. if err != nil {
  370. log.Println(err.Error())
  371. state = config.CONVERT_FAILED
  372. byteData = fmt.Appendf(nil, "[%s] %s %s", IdMsg(config.CONVERT_FAILED), request.conn.channel, request.Cmd)
  373. }
  374. }
  375. }
  376. // 如果id为0表示不需要回应
  377. if request.Id != 0 {
  378. request.conn.sendResponse <- &ResponseData{
  379. Id: request.Id,
  380. State: state,
  381. Data: byteData,
  382. }
  383. if h.cf.PrintMsg {
  384. log.Println("[RESP]->", request.Id, channel, "["+cmd+"]", state, subStr(string(byteData), 200))
  385. }
  386. }
  387. return
  388. }
  389. }
  390. }
  391. log.Println("[not match command]", channel, cmd)
  392. // 返回没有匹配的消息
  393. request.conn.sendResponse <- &ResponseData{
  394. Id: request.Id,
  395. State: config.NO_MATCH_CMD,
  396. Data: fmt.Appendf(nil, "[%s] Channel: %s, Cmd: %s", IdMsg(config.NO_MATCH_CMD), channel, cmd),
  397. }
  398. }
  399. // 请求频道并获取数据,采用回调的方式返回结果
  400. // 当前调用将会阻塞,直到命令都执行结束,最后返回执行的数量
  401. // 如果 backFunc 返回为 false 则提前结束
  402. // 最大数量和超时时间如果为0的话表示使用默认值
  403. func (h *Hub) GetWithStruct(gd *GetData, backFunc GetBackFunc) (count int, err error) {
  404. if gd.Filter == nil {
  405. return 0, errors.New("not filter found")
  406. }
  407. if gd.Timeout <= 0 {
  408. gd.Timeout = h.cf.WriteWait
  409. }
  410. if gd.backchan == nil {
  411. gd.backchan = make(chan *ResponseData, 32)
  412. }
  413. sendMax, err := h.sendRequest(gd)
  414. if err != nil {
  415. return 0, err
  416. }
  417. if sendMax <= 0 {
  418. return 0, nil
  419. }
  420. // 避免出现异常时线程无法退出
  421. timer := time.NewTimer(time.Millisecond * time.Duration(gd.Timeout+h.cf.WriteWait*2))
  422. defer func() {
  423. if !timer.Stop() {
  424. select {
  425. case <-timer.C:
  426. default:
  427. }
  428. }
  429. close(gd.backchan)
  430. }()
  431. for {
  432. select {
  433. case rp := <-gd.backchan:
  434. if rp == nil || rp.conn == nil {
  435. // 可能是已经退出了
  436. return
  437. }
  438. ch := rp.conn.channel
  439. if h.cf.PrintMsg {
  440. log.Println("[RECV]<-", rp.Id, ch, "["+gd.Cmd+"]", rp.State, subStr(string(rp.Data), 200))
  441. }
  442. count++
  443. // 如果这里返回为false这跳出循环
  444. if backFunc != nil && !backFunc(rp) {
  445. return
  446. }
  447. if count >= sendMax {
  448. return
  449. }
  450. case <-timer.C:
  451. return
  452. case <-h.ctx.Done():
  453. return
  454. }
  455. }
  456. }
  457. // 请求频道并获取数据,采用回调的方式返回结果
  458. // 当前调用将会阻塞,直到命令都执行结束,最后返回执行的数量
  459. // 如果 backFunc 返回为 false 则提前结束
  460. func (h *Hub) Get(filter FilterFunc, cmd string, data any, backFunc GetBackFunc) (count int, err error) {
  461. return h.GetWithStruct(&GetData{
  462. Filter: filter,
  463. Cmd: cmd,
  464. Data: data,
  465. }, backFunc)
  466. }
  467. // 获取一个数据,阻塞等待到超时间隔
  468. func (h *Hub) GetOneWithStruct(gd *GetData) (response *ResponseData) {
  469. if gd.Filter == nil {
  470. return &ResponseData{
  471. State: config.NO_MATCH_FILTER,
  472. Data: []byte(IdMsg(config.NO_MATCH_FILTER)),
  473. }
  474. }
  475. gd.Max = 1
  476. h.GetWithStruct(gd, func(rp *ResponseData) (ok bool) {
  477. response = rp
  478. return false
  479. })
  480. if response == nil {
  481. return &ResponseData{
  482. State: config.NO_MATCH_CONNECT,
  483. Data: []byte(IdMsg(config.NO_MATCH_CONNECT)),
  484. }
  485. }
  486. return
  487. }
  488. // 只获取一个频道的数据,阻塞等待到默认超时间隔
  489. // 如果没有结果将返回 NO_MATCH
  490. func (h *Hub) GetOne(filter FilterFunc, cmd string, data any) (response *ResponseData) {
  491. return h.GetOneWithStruct(&GetData{
  492. Filter: filter,
  493. Cmd: cmd,
  494. Data: data,
  495. Max: 1,
  496. })
  497. }
  498. func (h *Hub) GetRandOne(filter FilterFunc, cmd string, data any) (response *ResponseData) {
  499. return h.GetOneWithStruct(&GetData{
  500. Filter: filter,
  501. Cmd: cmd,
  502. Data: data,
  503. Max: 1,
  504. Rand: true,
  505. })
  506. }
  507. // 只获取一个频道的数据,阻塞等待到指定超时间隔
  508. // 如果没有结果将返回 NO_MATCH
  509. func (h *Hub) GetOneWithTimeout(filter FilterFunc, cmd string, data any, timeout int) (response *ResponseData) {
  510. return h.GetOneWithStruct(&GetData{
  511. Filter: filter,
  512. Cmd: cmd,
  513. Data: data,
  514. Max: 1,
  515. Timeout: timeout,
  516. })
  517. }
  518. func (h *Hub) GetRandOneWithTimeout(filter FilterFunc, cmd string, data any, timeout int) (response *ResponseData) {
  519. return h.GetOneWithStruct(&GetData{
  520. Filter: filter,
  521. Cmd: cmd,
  522. Data: data,
  523. Max: 1,
  524. Timeout: timeout,
  525. Rand: true,
  526. })
  527. }
  528. // 推送消息出去,不需要返回数据
  529. func (h *Hub) Push(filter FilterFunc, cmd string, data any) {
  530. h.PushWithMax(filter, cmd, data, 0)
  531. }
  532. // 推送最大对应数量的消息出去,不需要返回数据
  533. func (h *Hub) PushWithMax(filter FilterFunc, cmd string, data any, max int) {
  534. // 排除空频道
  535. if filter == nil {
  536. return
  537. }
  538. gd := &GetData{
  539. Filter: filter,
  540. Cmd: cmd,
  541. Data: data,
  542. Max: max,
  543. Timeout: h.cf.WriteWait,
  544. backchan: nil,
  545. }
  546. h.sendRequest(gd)
  547. }
  548. // 增加连接
  549. func (h *Hub) addLine(line *Line) {
  550. if h.lines.Exist(line) {
  551. log.Println("connect have exist")
  552. // 连接已经存在,直接返回
  553. return
  554. }
  555. h.lines.Store(line)
  556. }
  557. // 删除连接
  558. func (h *Hub) removeLine(line *Line) {
  559. line.Close(true)
  560. h.lines.Delete(line)
  561. }
  562. // 获取指定连接的连接持续时间
  563. func (h *Hub) ConnectDuration(line *Line) time.Duration {
  564. return time.Since(line.started)
  565. }
  566. // 绑定端口,建立服务
  567. // 需要程序运行时调用
  568. func (h *Hub) BindForServer(info *HostInfo) (err error) {
  569. doConnectFunc := func(conn conn.Connect) {
  570. proto, version, channel, auth, err := conn.ReadAuthInfo()
  571. if err != nil {
  572. log.Println("[BindForServer ReadAuthInfo ERROR]", err)
  573. conn.Close()
  574. return
  575. }
  576. if version != info.Version || proto != info.Proto {
  577. log.Println("wrong version or protocol: ", version, proto)
  578. conn.Close()
  579. return
  580. }
  581. if !h.checkAuthFunc(false, proto, version, channel, auth) {
  582. err = fmt.Errorf("[server checkAuthFunc ERROR] in proto: %s, version: %d, channel: %s, auth: %s", proto, version, channel, string(auth))
  583. log.Println(err)
  584. conn.Close()
  585. return
  586. }
  587. // 发送频道信息
  588. if err := conn.WriteAuthInfo(h.channel, h.authFunc(false, proto, version, channel, auth)); err != nil {
  589. log.Println("[WriteAuthInfo ERROR]", err)
  590. conn.Close()
  591. return
  592. }
  593. // 将连接加入现有连接中
  594. done := false
  595. h.lines.Range(func(id int, line *Line) bool {
  596. if line.state == Disconnected && line.host == nil && line.IsChannelEqual(channel) {
  597. line.Start(channel, conn, nil)
  598. done = true
  599. return false
  600. }
  601. return true
  602. })
  603. // 新建一个连接
  604. if !done {
  605. line := NewConnect(h.cf, h, channel, conn, nil, false)
  606. h.addLine(line)
  607. }
  608. }
  609. if info.Version == ws2.VERSION && info.Proto == ws2.PROTO {
  610. bind := ""
  611. if info.Bind != "" {
  612. bind = net.JoinHostPort(info.Bind, strconv.Itoa(int(info.Port)))
  613. }
  614. return ws2.Server(h.cf, bind, info.Path, info.Hash, doConnectFunc)
  615. } else if info.Version == tcp2.VERSION && info.Proto == tcp2.PROTO {
  616. return tcp2.Server(h.cf, net.JoinHostPort(info.Bind, strconv.Itoa(int(info.Port))), info.Hash, doConnectFunc)
  617. }
  618. return errors.New("not connect protocol and version found")
  619. }
  620. // 新建一个连接,不同的连接协议由底层自己选择
  621. // channel: 要连接的频道信息,需要能表达频道关键信息的部分
  622. func (h *Hub) ConnectToServer(channel string, force bool, host *HostInfo, autoReconnect bool) (err error) {
  623. // 检查当前channel是否已经存在
  624. if !force {
  625. line := h.ChannelToLine(channel)
  626. if line != nil && line.state == Connected {
  627. // err = fmt.Errorf("[ConnectToServer ERROR] existed channel: %s", channel)
  628. log.Println("[ConnectToServer] channel existed:", channel)
  629. return
  630. }
  631. }
  632. if host == nil {
  633. if h.connectHostFunc == nil {
  634. return errors.New("not connect host func found")
  635. }
  636. // 获取服务地址等信息
  637. host, err = h.connectHostFunc(channel, Both)
  638. if err != nil {
  639. return err
  640. }
  641. if host == nil {
  642. // 如果地址为空表示不需要连接,自己返回
  643. return
  644. }
  645. }
  646. var conn conn.Connect
  647. var runProto string
  648. addr := net.JoinHostPort(host.Host, strconv.Itoa(int(host.Port)))
  649. // 添加定时器
  650. ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*time.Duration(h.cf.ConnectTimeout))
  651. defer cancel()
  652. taskCh := make(chan bool)
  653. done := false
  654. go func() {
  655. if host.Version == ws2.VERSION && (host.Proto == ws2.PROTO || host.Proto == ws2.PROTO_STL) {
  656. runProto = ws2.PROTO
  657. if h.cf.PrintMsg {
  658. log.Println("[Connect]", host.Proto, addr, host.Path, host.Hash)
  659. }
  660. conn, err = ws2.Dial(h.cf, host.Proto, addr, host.Path, host.Hash)
  661. } else if host.Version == tcp2.VERSION && host.Proto == tcp2.PROTO {
  662. runProto = tcp2.PROTO
  663. if h.cf.PrintMsg {
  664. log.Println("[Connect]", host.Proto, addr, host.Hash)
  665. }
  666. conn, err = tcp2.Dial(h.cf, addr, host.Hash)
  667. } else {
  668. err = fmt.Errorf("not correct protocol and version found in: %+v", host)
  669. }
  670. if done {
  671. if err != nil {
  672. log.Println("[Dial ERROR]", err)
  673. }
  674. if conn != nil {
  675. conn.Close()
  676. }
  677. } else {
  678. taskCh <- err == nil
  679. }
  680. }()
  681. select {
  682. case ok := <-taskCh:
  683. cancel()
  684. if !ok || err != nil || conn == nil {
  685. log.Println("[Client ERROR]", host.Proto, err)
  686. host.Errors++
  687. host.Updated = time.Now()
  688. if err == nil {
  689. err = errors.New("unknown error")
  690. }
  691. return err
  692. }
  693. case <-ctx.Done():
  694. done = true
  695. return errors.New("timeout")
  696. case <-h.ctx.Done():
  697. return errors.New("quit")
  698. }
  699. // 如果 host 是代理,将代理信息添加到channel中
  700. localChannel := h.channel
  701. if host.Proxy {
  702. localChannel = localChannel + "?proxy=" + host.Host
  703. }
  704. // 发送验证信息
  705. if err := conn.WriteAuthInfo(localChannel, h.authFunc(true, runProto, host.Version, channel, nil)); err != nil {
  706. log.Println("[WriteAuthInfo ERROR]", err)
  707. conn.Close()
  708. host.Errors++
  709. host.Updated = time.Now()
  710. return err
  711. }
  712. // 接收频道信息
  713. proto, version, channel2, auth, err := conn.ReadAuthInfo()
  714. if err != nil {
  715. log.Println("[ConnectToServer ReadAuthInfo ERROR]", err)
  716. conn.Close()
  717. host.Errors++
  718. host.Updated = time.Now()
  719. return err
  720. }
  721. // 检查版本和协议是否一致
  722. if version != host.Version || proto != runProto {
  723. err = fmt.Errorf("[version or protocol wrong ERROR] %d, %s", version, proto)
  724. log.Println(err)
  725. conn.Close()
  726. host.Errors++
  727. host.Updated = time.Now()
  728. return err
  729. }
  730. // 检查频道名称是否匹配
  731. if !strings.Contains(channel2, channel) {
  732. err = fmt.Errorf("[channel ERROR] want %s, get %s", channel, channel2)
  733. log.Println(err)
  734. conn.Close()
  735. host.Errors++
  736. host.Updated = time.Now()
  737. return err
  738. }
  739. // 检查验证是否合法
  740. if !h.checkAuthFunc(true, proto, version, channel, auth) {
  741. err = fmt.Errorf("[client checkAuthFunc ERROR] in proto: %s, version: %d, channel: %s, auth: %s", proto, version, channel, string(auth))
  742. log.Println(err)
  743. conn.Close()
  744. host.Errors++
  745. host.Updated = time.Now()
  746. return err
  747. }
  748. // 更新服务主机信息
  749. host.Errors = 0
  750. host.Updated = time.Now()
  751. // 将连接加入现有连接中
  752. done = false
  753. h.lines.Range(func(id int, line *Line) bool {
  754. if line.channel == channel {
  755. if line.state == Connected {
  756. if force {
  757. line.Close(true)
  758. } else {
  759. err = fmt.Errorf("[connectToServer ERROR] channel already connected: %s", channel)
  760. log.Println(err)
  761. return false
  762. }
  763. }
  764. line.Start(channel, conn, host)
  765. done = true
  766. return false
  767. }
  768. return true
  769. })
  770. if err != nil {
  771. return err
  772. }
  773. // 新建一个连接
  774. if !done {
  775. line := NewConnect(h.cf, h, channel, conn, host, autoReconnect)
  776. h.addLine(line)
  777. }
  778. return nil
  779. }
  780. // 重试方式连接服务
  781. // 将会一直阻塞直到连接成功
  782. func (h *Hub) ConnectToServerX(channel string, force bool, host *HostInfo) {
  783. if host == nil {
  784. if h.connectHostFunc == nil {
  785. log.Println("ConnectToServerX: not connect host func found")
  786. return
  787. }
  788. host, _ = h.connectHostFunc(channel, Direct)
  789. if host == nil {
  790. log.Println("not host found for channel:", channel)
  791. return
  792. }
  793. }
  794. for {
  795. err := h.ConnectToServer(channel, force, host, true)
  796. if err == nil {
  797. return
  798. }
  799. log.Println("[ConnectToServer ERROR]", err)
  800. // 产生一个随机数避免刹间重连过载
  801. delay := time.Duration(rand.Intn(h.cf.ConnectTimeout)+(h.cf.ConnectTimeout/2)) * time.Millisecond
  802. log.Println("[ConnectToServerX reconnect]", delay, channel, host)
  803. host = nil
  804. time.Sleep(delay)
  805. }
  806. }
  807. // 检测处理连接状态,只在客户端有效
  808. func (h *Hub) checkConnect() {
  809. // 检查客户端获取主机地址函数
  810. if h.connectHostFunc == nil {
  811. return
  812. }
  813. proxyTicker := time.NewTicker(time.Duration(h.cf.ProxyTimeout * int(time.Millisecond)))
  814. connectTicker := time.NewTicker(time.Millisecond * time.Duration(h.cf.ConnectCheck))
  815. for {
  816. select {
  817. case <-proxyTicker.C:
  818. now := time.Now().UnixMilli()
  819. h.lines.Range(func(id int, line *Line) bool {
  820. if line.host != nil && line.host.Proxy && now-line.updated.UnixMilli() > int64(h.cf.ProxyTimeout) {
  821. host, err := h.connectHostFunc(line.channel, Direct)
  822. if err != nil {
  823. log.Println("[proxyTicker connectHostFunc ERROR]", err)
  824. return false
  825. }
  826. if host != nil {
  827. err = h.ConnectToServer(line.channel, true, host, line.autoReconnect)
  828. if err != nil {
  829. log.Println("[checkProxyConnect ConnectToServer]", err)
  830. }
  831. }
  832. }
  833. return true
  834. })
  835. case <-connectTicker.C:
  836. h.lines.Range(func(id int, line *Line) bool {
  837. if line.autoReconnect && line.host != nil && line.state == Disconnected {
  838. err := h.ConnectToServer(line.channel, true, nil, true)
  839. if err != nil {
  840. log.Println("[connectTicker ConnectToServer]", err)
  841. }
  842. }
  843. return true
  844. })
  845. case <-h.ctx.Done():
  846. return
  847. }
  848. }
  849. }
  850. // 退出所有的连接
  851. func (h *Hub) Quit() {
  852. h.cancel()
  853. h.lines.Range(func(id int, line *Line) bool {
  854. if line.state == Connected {
  855. line.Close(true)
  856. }
  857. return true
  858. })
  859. }
  860. // 建立一个集线器
  861. // connectFunc 用于监听连接状态的函数,可以为nil
  862. func NewHub(
  863. cf *config.Config,
  864. channel string,
  865. // 客户端需要用的函数,提供连接的主机信息 (服务端可空)
  866. connectHostFunc ConnectHostFunc,
  867. // 验证函数,获取认证信息,用于发送给对方
  868. authFunc AuthFunc,
  869. // 核对发送过来的认证信息
  870. checkAuthFunc CheckAuthFunc,
  871. // 连接状态变化时调用的函数
  872. connectStatusFunc ConnectStatusFunc,
  873. // 验证发送数据的条件是否满足 (可为空)
  874. // checkConnectOkFunc CheckConnectOkFunc,
  875. ) (h *Hub) {
  876. if cf == nil {
  877. cf = config.NewConfig()
  878. }
  879. ctx, cancel := context.WithCancel(context.Background())
  880. h = &Hub{
  881. ctx: ctx,
  882. cancel: cancel,
  883. cf: cf,
  884. globalID: uint16(time.Now().UnixNano()) % config.ID_MAX,
  885. channel: channel,
  886. middle: make([]MiddleFunc, 0),
  887. lines: NewMapx(),
  888. connectHostFunc: connectHostFunc,
  889. authFunc: authFunc,
  890. checkAuthFunc: checkAuthFunc,
  891. connectStatusFunc: connectStatusFunc,
  892. // checkConnectOkFunc: checkConnectOkFunc,
  893. lastCleanDeadConnect: time.Now().UnixMilli(),
  894. }
  895. go h.checkConnect()
  896. return h
  897. }