package tinymq import ( "context" "errors" "fmt" "log" "net" "strconv" "strings" "time" "git.me9.top/git/tinymq/config" "git.me9.top/git/tinymq/conn" "git.me9.top/git/tinymq/conn/tcp2" "git.me9.top/git/tinymq/conn/ws2" ) // 根据协议名建立连接,输出实际代码实现的协议 // 会一直阻塞直到连接结束 func Dail(host *conn.HostInfo, cf *config.Config) (conn conn.Connect, err error) { addr := net.JoinHostPort(host.Host, strconv.Itoa(int(host.Port))) if host.Version == ws2.VERSION && (host.Proto == ws2.PROTO || host.Proto == ws2.PROTO_STL) { if cf.PrintMsg { log.Println("[Dail]:", host.Proto, addr, host.Path, host.Hash) } conn, err = ws2.Dial(cf, host.Proto, addr, host.Path, host.Hash) } else if host.Version == tcp2.VERSION && host.Proto == tcp2.PROTO { if cf.PrintMsg { log.Println("[Dail]:", host.Proto, addr, host.Hash) } conn, err = tcp2.Dial(cf, addr, host.Hash) } else { err = fmt.Errorf("not correct protocol and version found in: %+v", host) } return } // 连接并超时检查 func DailWithTimeout(ctx context.Context, host *conn.HostInfo, cf *config.Config) (conn conn.Connect, err error) { // 添加定时器 cx, cancel := context.WithTimeout(context.Background(), time.Millisecond*time.Duration(cf.ConnectTimeout)) defer cancel() taskCh := make(chan bool) done := false // 指示是否已经完成 // 发送连接请求 go func() { conn, err = Dail(host, cf) if done { if err != nil { log.Println("[ConnectToServer] dial failed:", err) } if conn != nil { conn.Close() } } else { taskCh <- err == nil } }() select { case ok := <-taskCh: cancel() if !ok || err != nil || conn == nil { log.Println("[ConnectToServer] connect failed:", err) host.Errors++ host.Updated = time.Now() if err == nil { err = errors.New("unknown error") } return } case <-cx.Done(): done = true return nil, errors.New("timeout") case <-ctx.Done(): return nil, errors.New("quit") } return } // 客户端检验 AuthInfo func ClientCheckAuthInfo( conn conn.Connect, host *conn.HostInfo, localChannel string, remoteChannel string, authFunc AuthFunc, checkAuthFunc CheckAuthFunc, ) (needProxy bool, err error) { // 发送验证信息 if err := conn.WriteAuthInfo(localChannel, authFunc(host, host.Proto, host.Version, remoteChannel, nil)); err != nil { log.Println("[ClientCheckAuthInfo] WriteAuthInfo failed:", err) return false, err } // 接收频道信息 proto, version, channel, auth, err := conn.ReadAuthInfo() if err != nil { log.Println("[ClientCheckAuthInfo] ReadAuthInfo failed:", err) return false, err } // 检查版本和协议是否一致 if version != host.Version || !strings.HasPrefix(host.Proto, proto) { err = fmt.Errorf("[ClientCheckAuthInfo] version or protocol wrong: %d, %s", version, proto) log.Println(err) return false, err } // 检查频道名称是否匹配 if !strings.HasPrefix(channel, remoteChannel) { // 代理节点 if host.Proxy { // 还需要再建立连接 return true, nil } err = fmt.Errorf("[ClientCheckAuthInfo] channel want %s, get %s", remoteChannel, channel) log.Println(err) return false, err } // 检查验证是否合法 if !checkAuthFunc(host, proto, version, channel, auth) { err = fmt.Errorf("[ClientCheckAuthInfo] failed in proto: %s, version: %d, channel: %s, auth: %s", proto, version, channel, string(auth)) log.Println(err) return false, err } return }