Add support cs2+tcp protocol for xiaomi source

pkg/xiaomi/cs2/conn.go

@@ -1,6 +1,7 @@
 package cs2
 
 import (
+	"bufio"
 	"encoding/binary"
 	"fmt"
 	"io"
@@ -10,33 +11,27 @@ import (
 	"time"
 )
 
-func Dial(host string) (*Conn, error) {
+func Dial(host, transport string) (*Conn, error) {
-	conn, err := net.ListenUDP("udp", nil)
+	conn, err := handshake(host, transport)
 	if err != nil {
 		return nil, err
 	}
 
+	_, isTCP := conn.(*tcpConn)
+
 	c := &Conn{
 		conn:   conn,
-		addr: &net.UDPAddr{IP: net.ParseIP(host), Port: 32108},
+		isTCP:  isTCP,
+		rawCh0: make(chan []byte, 10),
+		rawCh2: make(chan []byte, 100),
 	}
-
-	if err = c.handshake(); err != nil {
-		_ = conn.Close()
-		return nil, err
-	}
-
-	c.rawCh0 = make(chan []byte, 10)
-	c.rawCh2 = make(chan []byte, 100)
-
 	go c.worker()
-
 	return c, nil
 }
 
 type Conn struct {
-	conn *net.UDPConn
+	conn  net.Conn
-	addr *net.UDPAddr
+	isTCP bool
 
 	err    error
 	seqCh0 uint16
@@ -53,30 +48,58 @@ const (
 	magicDrw     = 0xD1
 	msgLanSearch = 0x30
 	msgPunchPkt  = 0x41
-	msgP2PRdy    = 0x42
+	msgP2PRdyUDP = 0x42
+	msgP2PRdyTCP = 0x43
 	msgDrw       = 0xD0
 	msgDrwAck    = 0xD1
-	msgAlive     = 0xE0
+	msgPing      = 0xE0
+	msgPong      = 0xE1
+	msgClose     = 0xF1
 )
 
-func (c *Conn) handshake() error {
+func handshake(host, transport string) (net.Conn, error) {
-	_ = c.SetDeadline(time.Now().Add(5 * time.Second))
+	conn, err := newUDPConn(host, 32108)
-
-	buf, err := c.WriteAndWait([]byte{magic, msgLanSearch, 0, 0}, msgPunchPkt)
 	if err != nil {
-		return fmt.Errorf("%s: read punch: %w", "cs2", err)
+		return nil, err
 	}
 
-	_, err = c.WriteAndWait(buf, msgP2PRdy)
+	_ = conn.SetDeadline(time.Now().Add(5 * time.Second))
+
+	req := []byte{magic, msgLanSearch, 0, 0}
+	res, err := conn.(*udpConn).WriteUntil(req, func(res []byte) bool {
+		return res[1] == msgPunchPkt
+	})
 	if err != nil {
-		return fmt.Errorf("%s: read ready: %w", "cs2", err)
+		_ = conn.Close()
+		return nil, err
 	}
 
-	_ = c.Write([]byte{magic, msgAlive, 0, 0})
+	var msgUDP, msgTCP byte
 
-	_ = c.SetDeadline(time.Time{})
+	if transport == "" || transport == "udp" {
+		msgUDP = msgP2PRdyUDP
+	}
+	if transport == "" || transport == "tcp" {
+		msgTCP = msgP2PRdyTCP
+	}
 
-	return nil
+	res, err = conn.(*udpConn).WriteUntil(res, func(res []byte) bool {
+		return res[1] == msgUDP || res[1] == msgTCP
+	})
+	if err != nil {
+		_ = conn.Close()
+		return nil, err
+	}
+
+	_ = conn.SetDeadline(time.Time{})
+
+	if res[1] == msgTCP {
+		_ = conn.Close()
+		//host := fmt.Sprintf("%d.%d.%d.%d:%d", b[31], b[30], b[29], b[28], uint16(b[27])<<8|uint16(b[26]))
+		return newTCPConn(conn.RemoteAddr().String())
+	}
+
+	return conn, nil
 }
 
 func (c *Conn) worker() {
@@ -85,27 +108,31 @@ func (c *Conn) worker() {
 		close(c.rawCh2)
 	}()
 
-	chAck := make([]uint16, 4)
+	chAck := make([]uint16, 4) // only for UDP
 	buf := make([]byte, 1200)
 	var ch2WaitSize int
 	var ch2WaitData []byte
+	var keepaliveTS time.Time
 
 	for {
-		n, addr, err := c.conn.ReadFromUDP(buf)
+		n, err := c.conn.Read(buf)
 		if err != nil {
 			c.err = fmt.Errorf("%s: %w", "cs2", err)
 			return
 		}
 
-		if string(addr.IP) != string(c.addr.IP) || n < 8 || buf[0] != magic {
-			continue // skip messages from another IP
-		}
-
-		//log.Printf("<- %x", buf[:n])
-
 		switch buf[1] {
 		case msgDrw:
 			ch := buf[5]
+
+			if c.isTCP {
+				// For TCP we should using ping/pong.
+				if now := time.Now(); now.After(keepaliveTS) {
+					_, _ = c.conn.Write([]byte{magic, msgPing, 0, 0})
+					keepaliveTS = now.Add(5 * time.Second)
+				}
+			} else {
+				// For UDP we should using ack.
 				seqHI := buf[6]
 				seqLO := buf[7]
 
@@ -115,8 +142,7 @@ func (c *Conn) worker() {
 				chAck[ch]++
 
 				ack := []byte{magic, msgDrwAck, 0, 6, magicDrw, ch, 0, 1, seqHI, seqLO}
-			if _, err = c.conn.WriteToUDP(ack, c.addr); err != nil {
+				_, _ = c.conn.Write(ack)
-				return
 			}
 
 			switch ch {
@@ -152,9 +178,12 @@ func (c *Conn) worker() {
 				continue
 			}
 
-		case msgP2PRdy: // skip it
+		case msgPing:
+			_, _ = c.conn.Write([]byte{magic, msgPong, 0, 0})
 			continue
-		case msgDrwAck:
+		case msgPong, msgP2PRdyUDP, msgP2PRdyTCP, msgClose:
+			continue // skip it
+		case msgDrwAck: // only for UDP
 			if c.cmdAck != nil {
 				c.cmdAck()
 			}
@@ -165,42 +194,15 @@ func (c *Conn) worker() {
 	}
 }
 
-func (c *Conn) Write(req []byte) error {
+func (c *Conn) Protocol() string {
-	//log.Printf("-> %x", req)
+	if c.isTCP {
-	_, err := c.conn.WriteToUDP(req, c.addr)
+		return "cs2+tcp"
-	return err
-}
-
-func (c *Conn) WriteAndWait(req []byte, waitMsg uint8) ([]byte, error) {
-	var t *time.Timer
-	t = time.AfterFunc(1, func() {
-		if err := c.Write(req); err == nil && t != nil {
-			t.Reset(time.Second)
-		}
-	})
-	defer t.Stop()
-
-	buf := make([]byte, 1200)
-
-	for {
-		n, addr, err := c.conn.ReadFromUDP(buf)
-		if err != nil {
-			return nil, err
-		}
-
-		if string(addr.IP) != string(c.addr.IP) || n < 16 {
-			continue // skip messages from another IP
-		}
-
-		if buf[0] == magic && buf[1] == waitMsg {
-			c.addr.Port = addr.Port
-			return buf[:n], nil
-		}
 	}
+	return "cs2+udp"
 }
 
 func (c *Conn) RemoteAddr() net.Addr {
-	return c.addr
+	return c.conn.RemoteAddr()
 }
 
 func (c *Conn) SetDeadline(t time.Time) error {
@@ -232,6 +234,14 @@ func (c *Conn) WriteCommand(cmd uint16, data []byte) error {
 	c.cmdMu.Lock()
 	defer c.cmdMu.Unlock()
 
+	req := marshalCmd(0, c.seqCh0, uint32(cmd), data)
+	c.seqCh0++
+
+	if c.isTCP {
+		_, err := c.conn.Write(req)
+		return err
+	}
+
 	var repeat atomic.Int32
 	repeat.Store(5)
 
@@ -243,11 +253,8 @@ func (c *Conn) WriteCommand(cmd uint16, data []byte) error {
 		timeout.Reset(1)
 	}
 
-	req := marshalCmd(0, c.seqCh0, uint32(cmd), data)
-	c.seqCh0++
-
 	for {
-		if err := c.Write(req); err != nil {
+		if _, err := c.conn.Write(req); err != nil {
 			return err
 		}
 		<-timeout.C
@@ -285,7 +292,8 @@ func (c *Conn) WritePacket(data []byte) error {
 	binary.BigEndian.PutUint32(req[8:], n)
 	copy(req[offset:], data)
 
-	return c.Write(req)
+	_, err := c.conn.Write(req)
+	return err
 }
 
 func marshalCmd(channel byte, seq uint16, cmd uint32, payload []byte) []byte {
@@ -313,3 +321,112 @@ func marshalCmd(channel byte, seq uint16, cmd uint32, payload []byte) []byte {
 
 	return req
 }
+
+func newUDPConn(host string, port int) (net.Conn, error) {
+	// We using raw net.UDPConn, because RemoteAddr should be changed during handshake.
+	conn, err := net.ListenUDP("udp", nil)
+	if err != nil {
+		return nil, err
+	}
+
+	addr, err := net.ResolveUDPAddr("udp", host)
+	if err != nil {
+		addr = &net.UDPAddr{IP: net.ParseIP(host), Port: port}
+	}
+
+	return &udpConn{UDPConn: conn, addr: addr}, nil
+}
+
+type udpConn struct {
+	*net.UDPConn
+	addr *net.UDPAddr
+}
+
+func (c *udpConn) Read(p []byte) (n int, err error) {
+	var addr *net.UDPAddr
+	for {
+		n, addr, err = c.UDPConn.ReadFromUDP(p)
+		if err != nil {
+			return 0, err
+		}
+
+		if string(addr.IP) == string(c.addr.IP) || n >= 8 {
+			return
+		}
+	}
+}
+
+func (c *udpConn) Write(req []byte) (n int, err error) {
+	//log.Printf("-> %x", req)
+	return c.UDPConn.WriteToUDP(req, c.addr)
+}
+
+func (c *udpConn) RemoteAddr() net.Addr {
+	return c.addr
+}
+
+func (c *udpConn) WriteUntil(req []byte, ok func(res []byte) bool) ([]byte, error) {
+	var t *time.Timer
+	t = time.AfterFunc(1, func() {
+		if _, err := c.Write(req); err == nil && t != nil {
+			t.Reset(time.Second)
+		}
+	})
+	defer t.Stop()
+
+	buf := make([]byte, 1200)
+
+	for {
+		n, addr, err := c.UDPConn.ReadFromUDP(buf)
+		if err != nil {
+			return nil, err
+		}
+
+		if string(addr.IP) != string(c.addr.IP) || n < 16 {
+			continue // skip messages from another IP
+		}
+
+		if ok(buf[:n]) {
+			c.addr.Port = addr.Port
+			return buf[:n], nil
+		}
+	}
+}
+
+func newTCPConn(addr string) (net.Conn, error) {
+	conn, err := net.DialTimeout("tcp", addr, 3*time.Second)
+	if err != nil {
+		return nil, err
+	}
+	return &tcpConn{conn.(*net.TCPConn), bufio.NewReader(conn)}, nil
+}
+
+type tcpConn struct {
+	*net.TCPConn
+	rd *bufio.Reader
+}
+
+func (c *tcpConn) Read(p []byte) (n int, err error) {
+	tmp := make([]byte, 8)
+	if _, err = io.ReadFull(c.rd, tmp); err != nil {
+		return
+	}
+	n = int(binary.BigEndian.Uint16(tmp))
+	if len(p) < n {
+		return 0, fmt.Errorf("tcp: buffer too small")
+	}
+	_, err = io.ReadFull(c.rd, p[:n])
+	//log.Printf("<- %x%x", tmp, p[:n])
+	return
+}
+
+func (c *tcpConn) Write(req []byte) (n int, err error) {
+	n = len(req)
+	buf := make([]byte, 8+n)
+	binary.BigEndian.PutUint16(buf, uint16(n))
+	buf[2] = 0x68
+	copy(buf[8:], req)
+	//log.Printf("-> %x", buf)
+	_, err = c.TCPConn.Write(buf)
+	return
+}

pkg/xiaomi/miss/client.go

@@ -33,7 +33,7 @@ func Dial(rawURL string) (*Client, error) {
 
 	switch s := query.Get("vendor"); s {
 	case "cs2":
-		c.conn, err = cs2.Dial(u.Host)
+		c.conn, err = cs2.Dial(u.Host, query.Get("transport"))
 	case "tutk":
 		c.conn, err = tutk.Dial(u.Host, query.Get("uid"))
 	default:
@@ -63,6 +63,7 @@ const (
 )
 
 type Conn interface {
+	Protocol() string
 	ReadCommand() (cmd uint16, data []byte, err error)
 	WriteCommand(cmd uint16, data []byte) error
 	ReadPacket() ([]byte, error)
@@ -77,6 +78,10 @@ type Client struct {
 	key  []byte
 }
 
+func (c *Client) Protocol() string {
+	return c.conn.Protocol()
+}
+
 func (c *Client) RemoteAddr() net.Addr {
 	return c.conn.RemoteAddr()
 }
@@ -89,16 +94,6 @@ func (c *Client) Close() error {
 	return c.conn.Close()
 }
 
-func (c *Client) Protocol() string {
-	switch c.conn.(type) {
-	case *cs2.Conn:
-		return "cs2+udp"
-	case *tutk.Conn:
-		return "tutk+udp"
-	}
-	return ""
-}
-
 const (
 	cmdAuthReq           = 0x100
 	cmdAuthRes           = 0x101

pkg/xiaomi/tutk/conn.go

@@ -249,6 +249,10 @@ func (c *Conn) WriteAndWait(req []byte, ok func(req, res []byte) bool) ([]byte,
 	}
 }
 
+func (c *Conn) Protocol() string {
+	return "tutk+udp"
+}
+
 func (c *Conn) RemoteAddr() net.Addr {
 	return c.addr
 }