Add support cs2+tcp protocol for xiaomi source

Author: AlexxIT

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) {
-	conn, err := net.ListenUDP("udp", nil)
+func Dial(host, transport string) (*Conn, error) {
+	conn, err := handshake(host, transport)
 	if err != nil {
 		return nil, err
 	}
 
-	c := &Conn{
-		conn: conn,
-		addr: &net.UDPAddr{IP: net.ParseIP(host), Port: 32108},
-	}
+	_, isTCP := conn.(*tcpConn)
 
-	if err = c.handshake(); err != nil {
-		_ = conn.Close()
-		return nil, err
+	c := &Conn{
+		conn:   conn,
+		isTCP:  isTCP,
+		rawCh0: make(chan []byte, 10),
+		rawCh2: make(chan []byte, 100),
 	}
-
-	c.rawCh0 = make(chan []byte, 10)
-	c.rawCh2 = make(chan []byte, 100)
-
 	go c.worker()
-
 	return c, nil
 }
 
 type Conn struct {
-	conn *net.UDPConn
-	addr *net.UDPAddr
+	conn  net.Conn
+	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 {
-	_ = c.SetDeadline(time.Now().Add(5 * time.Second))
+func handshake(host, transport string) (net.Conn, error) {
+	conn, err := newUDPConn(host, 32108)
+	if err != nil {
+		return nil, err
+	}
+
+	_ = conn.SetDeadline(time.Now().Add(5 * time.Second))
 
-	buf, err := c.WriteAndWait([]byte{magic, msgLanSearch, 0, 0}, msgPunchPkt)
+	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 punch: %w", "cs2", err)
+		_ = conn.Close()
+		return nil, err
 	}
 
-	_, err = c.WriteAndWait(buf, msgP2PRdy)
+	var msgUDP, msgTCP byte
+
+	if transport == "" || transport == "udp" {
+		msgUDP = msgP2PRdyUDP
+	}
+	if transport == "" || transport == "tcp" {
+		msgTCP = msgP2PRdyTCP
+	}
+
+	res, err = conn.(*udpConn).WriteUntil(res, func(res []byte) bool {
+		return res[1] == msgUDP || res[1] == msgTCP
+	})
 	if err != nil {
-		return fmt.Errorf("%s: read ready: %w", "cs2", err)
+		_ = conn.Close()
+		return nil, err
 	}
 
-	_ = c.Write([]byte{magic, msgAlive, 0, 0})
+	_ = conn.SetDeadline(time.Time{})
 
-	_ = c.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 nil
+	return conn, nil
 }
 
 func (c *Conn) worker() {
@@ -85,38 +108,41 @@ 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]
-			seqHI := buf[6]
-			seqLO := buf[7]
 
-			if chAck[ch] != uint16(seqHI)<<8|uint16(seqLO) {
-				continue
-			}
-			chAck[ch]++
+			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]
+
+				if chAck[ch] != uint16(seqHI)<<8|uint16(seqLO) {
+					continue
+				}
+				chAck[ch]++
 
-			ack := []byte{magic, msgDrwAck, 0, 6, magicDrw, ch, 0, 1, seqHI, seqLO}
-			if _, err = c.conn.WriteToUDP(ack, c.addr); err != nil {
-				return
+				ack := []byte{magic, msgDrwAck, 0, 6, magicDrw, ch, 0, 1, seqHI, seqLO}
+				_, _ = c.conn.Write(ack)
 			}
 
 			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 {
-	//log.Printf("-> %x", req)
-	_, err := c.conn.WriteToUDP(req, c.addr)
-	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
-		}
+func (c *Conn) Protocol() string {
+	if c.isTCP {
+		return "cs2+tcp"
 	}
+	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
+}