Joyit 1 місяць тому
батько
коміт
deb51d7109
5 змінених файлів з 139 додано та 44 видалено
  1. 22 0
      conn/util.go
  2. 1 1
      const.go
  3. 47 0
      dail.go
  4. 49 43
      hub.go
  5. 20 0
      line.go

+ 22 - 0
conn/util.go

@@ -215,6 +215,28 @@ func PingPackageEncode(id uint16) (buf []byte) {
 	return
 }
 
+func ProxyConnectPackageEncode(url string) (buf []byte) {
+	dlen := 2 + len(url)
+	buf = pool.Get(dlen)
+	index := 0
+	binary.BigEndian.PutUint16(buf[index:index+2], config.ID_PROXY_CONNECT)
+	index += 2
+	copy(buf[index:], []byte(url))
+	return
+}
+
+func ProxyResultPackageEncode(result string) (buf []byte) {
+	dlen := 2 + len(result)
+	buf = pool.Get(dlen)
+	index := 0
+	binary.BigEndian.PutUint16(buf[index:index+2], config.ID_PROXY_RESULT)
+	index += 2
+	if index < dlen {
+		copy(buf[index:], []byte(result))
+	}
+	return
+}
+
 // 解析收到的数据包
 func PackageDecode(msg []byte, recycle bool, compress bool, cf *config.Config) (msgType MsgType, id uint16, cmd string, state uint8, data []byte, err error) {
 	if recycle {

+ 1 - 1
const.go

@@ -29,7 +29,7 @@ const (
 	// ID 号最高值,高于这个值的ID号为系统内部使用
 	ID_MAX = config.ID_MAX
 	// 验证ID
-	ID_AUTH = config.ID_AUTH
+	// ID_AUTH = config.ID_AUTH
 )
 
 // 转换 id 到对应的消息

+ 47 - 0
dail.go

@@ -1,11 +1,14 @@
 package tinymq
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"log"
 	"net"
 	"strconv"
 	"strings"
+	"time"
 
 	"git.me9.top/git/tinymq/config"
 	"git.me9.top/git/tinymq/conn"
@@ -33,6 +36,50 @@ func Dail(host *conn.HostInfo, cf *config.Config) (conn conn.Connect, err error)
 	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,

+ 49 - 43
hub.go

@@ -690,50 +690,56 @@ func (h *Hub) ConnectToServer(remoteChannel string, force bool, host *conn.HostI
 		}
 	}
 
-	var conn conn.Connect
-	// var rawProto string
-	// addr := net.JoinHostPort(host.Host, strconv.Itoa(int(host.Port)))
-
-	// 添加定时器
-	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*time.Duration(h.cf.ConnectTimeout))
-	defer cancel()
-	taskCh := make(chan bool)
-	done := false // 指示是否已经完成
-
-	// 发送连接请求
-	go func() {
-		conn, err = Dail(host, h.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 err
-		}
-	case <-ctx.Done():
-		done = true
-		return errors.New("timeout")
-	case <-h.ctx.Done():
-		return errors.New("quit")
+	conn, err := DailWithTimeout(h.ctx, host, h.cf)
+	if err != nil {
+		log.Println("[DailWithTimeout] error:", err)
+		return err
 	}
 
+	// var conn conn.Connect
+	// // var rawProto string
+	// // addr := net.JoinHostPort(host.Host, strconv.Itoa(int(host.Port)))
+
+	// // 添加定时器
+	// ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*time.Duration(h.cf.ConnectTimeout))
+	// defer cancel()
+	// taskCh := make(chan bool)
+	// done := false // 指示是否已经完成
+
+	// // 发送连接请求
+	// go func() {
+	// 	conn, err = Dail(host, h.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 err
+	// 	}
+	// case <-ctx.Done():
+	// 	done = true
+	// 	return errors.New("timeout")
+	// case <-h.ctx.Done():
+	// 	return errors.New("quit")
+	// }
+
 	// 如果 host 是代理,将代理信息添加到channel中
 	// localChannel := h.channel
 	// if host.Proxy {
@@ -751,7 +757,7 @@ func (h *Hub) ConnectToServer(remoteChannel string, force bool, host *conn.HostI
 	}
 
 	// 将连接加入现有连接中
-	done = false
+	done := false
 	h.lines.Range(func(id int, line *Line) bool {
 		if line.channel == remoteChannel {
 			if line.state == StateConnected {

+ 20 - 0
line.go

@@ -167,7 +167,27 @@ func (c *Line) readPump() {
 					conn:  c,
 				})
 			case conn.MsgProxyConnect:
+				c.state = StateProxing
 				// TODO: 建立代理连接
+				host, err := conn.ParseUrl(string(data))
+				if err != nil {
+					log.Println("[MsgProxyConnect] ParseUrl:", err)
+					c.Close(false)
+					return
+				}
+				proxyConn, err := DailWithTimeout(c.hub.ctx, host, c.cf)
+				if err != nil {
+					log.Println("[MsgProxyConnect] DailWithTimeout:", err)
+					// 回退到正常连接状态
+					c.state = StateConnected
+					// 发送失败状态包
+					conn.ProxyResultPackageEncode(err.Error())
+				} else {
+					c.proxyConn = proxyConn
+					c.proxyHost = host
+					c.state = StateProxied
+					conn.ProxyResultPackageEncode("") // not news is good news
+				}
 			}
 		case StateProxied: // 已经处于代理状态
 			buf, recycle, err := c.conn.ReadRawPackage(c.cf.LongReadWait)