hub.go 24 KB

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