tinymq/dail.go

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
// 如果返回的 channel 有值，说明需要代理连接
func ClientCheckAuthInfo(
	conn conn.Connect,
	host *conn.HostInfo,
	localChannel string,
	remoteChannel string,
	authFunc AuthFunc,
	checkAuthFunc CheckAuthFunc,
) (channel string, err error) {
	// 发送验证信息
	if err := conn.WriteAuthInfo(localChannel, authFunc(host, host.Proto, host.Version, remoteChannel, nil)); err != nil {
		log.Println("[ClientCheckAuthInfo] WriteAuthInfo failed:", err)
		return "", err
	}
	// 接收频道信息
	proto, version, channel, auth, err := conn.ReadAuthInfo()
	if err != nil {
		log.Println("[ClientCheckAuthInfo] ReadAuthInfo failed:", err)
		return "", 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 "", err
	}
	// 检查频道名称是否匹配，只匹配最终的频道
	if !strings.HasPrefix(channel, strings.Split(remoteChannel, "<-")[0]) {
		// 代理节点
		if host.Proxy {
			// 还需要再建立连接
			return channel, nil
		}
		err = fmt.Errorf("[ClientCheckAuthInfo] channel want %s, get %s", remoteChannel, channel)
		log.Println(err)
		return "", 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 "", err
	}
	return "", nil
}

// 服务端验证
func ServerCheckAuthInfo(
	conn conn.Connect,
	host *conn.HostInfo,
	localChannel string,
	authFunc AuthFunc,
	checkAuthFunc CheckAuthFunc,
) (channel string, err error) {
	var proto string
	var version uint8
	var auth []byte
	proto, version, channel, auth, err = conn.ReadAuthInfo()
	if err != nil {
		log.Println("[BindForServer ReadAuthInfo ERROR]", err)
		return
	}
	if version != host.Version || !strings.HasPrefix(host.Proto, proto) {
		err = fmt.Errorf("wrong version (%d, %d) or protocol (%s, %s)", version, host.Version, proto, host.Proto)
		log.Println(err)
		return
	}
	if !checkAuthFunc(nil, proto, version, channel, auth) {
		err = fmt.Errorf("[server checkAuthFunc ERROR] in proto: %s, version: %d, channel: %s, auth: %s", proto, version, channel, string(auth))
		log.Println(err)
		return
	}
	// 发送频道信息
	if err = conn.WriteAuthInfo(localChannel, authFunc(nil, proto, version, channel, auth)); err != nil {
		log.Println("[WriteAuthInfo ERROR]", err)
		return
	}
	return
}