hub.go 24 KB

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