507 lines
10 KiB
Go
507 lines
10 KiB
Go
package cs2
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"encoding/binary"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
func Dial(host, transport string) (*Conn, error) {
|
|
conn, err := handshake(host, transport)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
_, isTCP := conn.(*tcpConn)
|
|
|
|
c := &Conn{
|
|
Conn: conn,
|
|
isTCP: isTCP,
|
|
channels: [4]*dataChannel{
|
|
newDataChannel(0, 10), nil, newDataChannel(250, 100), nil,
|
|
},
|
|
}
|
|
go c.worker()
|
|
return c, nil
|
|
}
|
|
|
|
type Conn struct {
|
|
net.Conn
|
|
isTCP bool
|
|
|
|
err error
|
|
seqCh0 uint16
|
|
seqCh3 uint16
|
|
|
|
channels [4]*dataChannel
|
|
|
|
cmdMu sync.Mutex
|
|
cmdAck func()
|
|
}
|
|
|
|
const (
|
|
magic = 0xF1
|
|
magicDrw = 0xD1
|
|
magicTCP = 0x68
|
|
msgLanSearch = 0x30
|
|
msgPunchPkt = 0x41
|
|
msgP2PRdyUDP = 0x42
|
|
msgP2PRdyTCP = 0x43
|
|
msgDrw = 0xD0
|
|
msgDrwAck = 0xD1
|
|
msgPing = 0xE0
|
|
msgPong = 0xE1
|
|
msgClose = 0xF0
|
|
msgCloseAck = 0xF1
|
|
)
|
|
|
|
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))
|
|
|
|
req := []byte{magic, msgLanSearch, 0, 0}
|
|
res, err := conn.(*udpConn).WriteUntil(req, func(res []byte) bool {
|
|
return res[1] == msgPunchPkt
|
|
})
|
|
if err != nil {
|
|
_ = conn.Close()
|
|
return nil, err
|
|
}
|
|
|
|
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 {
|
|
_ = 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() {
|
|
defer func() {
|
|
c.channels[0].Close()
|
|
c.channels[2].Close()
|
|
}()
|
|
|
|
var keepaliveTS time.Time // only for TCP
|
|
|
|
buf := make([]byte, 1200)
|
|
|
|
for {
|
|
n, err := c.Conn.Read(buf)
|
|
if err != nil {
|
|
c.err = fmt.Errorf("%s: %w", "cs2", err)
|
|
return
|
|
}
|
|
|
|
// 0 f1d0 magic
|
|
// 2 005d size = total size + 4
|
|
// 4 d1 magic
|
|
// 5 00 channel
|
|
// 6 0000 seq
|
|
switch buf[1] {
|
|
case msgDrw:
|
|
ch := buf[5]
|
|
channel := c.channels[ch]
|
|
|
|
if c.isTCP {
|
|
// For TCP we should send ping every second to keep connection alive.
|
|
// Based on PCAP analysis: official Mi Home app sends PING every ~1s.
|
|
if now := time.Now(); now.After(keepaliveTS) {
|
|
_, _ = c.Conn.Write([]byte{magic, msgPing, 0, 0})
|
|
keepaliveTS = now.Add(time.Second)
|
|
}
|
|
|
|
err = channel.Push(buf[8:n])
|
|
} else {
|
|
var pushed int
|
|
|
|
seqHI, seqLO := buf[6], buf[7]
|
|
seq := uint16(seqHI)<<8 | uint16(seqLO)
|
|
pushed, err = channel.PushSeq(seq, buf[8:n])
|
|
|
|
if pushed >= 0 {
|
|
// For UDP we should send ACK.
|
|
ack := []byte{magic, msgDrwAck, 0, 6, magicDrw, ch, 0, 1, seqHI, seqLO}
|
|
_, _ = c.Conn.Write(ack)
|
|
}
|
|
}
|
|
|
|
if err != nil {
|
|
c.err = fmt.Errorf("%s: %w", "cs2", err)
|
|
return
|
|
}
|
|
|
|
case msgPing:
|
|
_, _ = c.Conn.Write([]byte{magic, msgPong, 0, 0})
|
|
case msgPong, msgP2PRdyUDP, msgP2PRdyTCP, msgClose, msgCloseAck: // skip it
|
|
case msgDrwAck: // only for UDP
|
|
if c.cmdAck != nil {
|
|
c.cmdAck()
|
|
}
|
|
default:
|
|
fmt.Printf("%s: unknown msg: %x\n", "cs2", buf[:n])
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *Conn) Protocol() string {
|
|
if c.isTCP {
|
|
return "cs2+tcp"
|
|
}
|
|
return "cs2+udp"
|
|
}
|
|
|
|
func (c *Conn) Version() string {
|
|
return "CS2"
|
|
}
|
|
|
|
func (c *Conn) Error() error {
|
|
if c.err != nil {
|
|
return c.err
|
|
}
|
|
return io.EOF
|
|
}
|
|
|
|
func (c *Conn) ReadCommand() (cmd uint32, data []byte, err error) {
|
|
buf, ok := c.channels[0].Pop()
|
|
if !ok {
|
|
return 0, nil, c.Error()
|
|
}
|
|
cmd = binary.LittleEndian.Uint32(buf)
|
|
data = buf[4:]
|
|
return
|
|
}
|
|
|
|
func (c *Conn) WriteCommand(cmd uint32, data []byte) error {
|
|
c.cmdMu.Lock()
|
|
defer c.cmdMu.Unlock()
|
|
|
|
req := marshalCmd(0, c.seqCh0, cmd, data)
|
|
c.seqCh0++
|
|
|
|
if c.isTCP {
|
|
_, err := c.Conn.Write(req)
|
|
return err
|
|
}
|
|
|
|
var repeat atomic.Int32
|
|
repeat.Store(5)
|
|
|
|
timeout := time.NewTicker(time.Second)
|
|
defer timeout.Stop()
|
|
|
|
c.cmdAck = func() {
|
|
repeat.Store(0)
|
|
timeout.Reset(1)
|
|
}
|
|
|
|
for {
|
|
if _, err := c.Conn.Write(req); err != nil {
|
|
return err
|
|
}
|
|
<-timeout.C
|
|
r := repeat.Add(-1)
|
|
if r < 0 {
|
|
return nil
|
|
}
|
|
if r == 0 {
|
|
return fmt.Errorf("%s: can't send command %d", "cs2", cmd)
|
|
}
|
|
}
|
|
}
|
|
|
|
const hdrSize = 32
|
|
|
|
func (c *Conn) ReadPacket() (hdr, payload []byte, err error) {
|
|
data, ok := c.channels[2].Pop()
|
|
if !ok {
|
|
return nil, nil, c.Error()
|
|
}
|
|
return data[:hdrSize], data[hdrSize:], nil
|
|
}
|
|
|
|
func (c *Conn) WritePacket(hdr, payload []byte) error {
|
|
const offset = 12
|
|
|
|
n := hdrSize + uint32(len(payload))
|
|
req := make([]byte, n+offset)
|
|
req[0] = magic
|
|
req[1] = msgDrw
|
|
binary.BigEndian.PutUint16(req[2:], uint16(n+8))
|
|
|
|
req[4] = magicDrw
|
|
req[5] = 3 // channel
|
|
binary.BigEndian.PutUint16(req[6:], c.seqCh3)
|
|
c.seqCh3++
|
|
binary.BigEndian.PutUint32(req[8:], n)
|
|
copy(req[offset:], hdr)
|
|
copy(req[offset+hdrSize:], hdr)
|
|
|
|
_, err := c.Conn.Write(req)
|
|
return err
|
|
}
|
|
|
|
func marshalCmd(channel byte, seq uint16, cmd uint32, payload []byte) []byte {
|
|
size := len(payload)
|
|
req := make([]byte, 4+4+4+4+size)
|
|
|
|
// 1. message header (4 bytes)
|
|
req[0] = magic
|
|
req[1] = msgDrw
|
|
binary.BigEndian.PutUint16(req[2:], uint16(4+4+4+size))
|
|
|
|
// 2. drw? header (4 bytes)
|
|
req[4] = magicDrw
|
|
req[5] = channel
|
|
binary.BigEndian.PutUint16(req[6:], seq)
|
|
|
|
// 3. payload size (4 bytes)
|
|
binary.BigEndian.PutUint32(req[8:], uint32(4+size))
|
|
|
|
// 4. payload command (4 bytes)
|
|
binary.BigEndian.PutUint32(req[12:], cmd)
|
|
|
|
// 5. payload
|
|
copy(req[16:], payload)
|
|
|
|
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(b []byte) (n int, err error) {
|
|
var addr *net.UDPAddr
|
|
for {
|
|
n, addr, err = c.UDPConn.ReadFromUDP(b)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
if string(addr.IP) == string(c.addr.IP) || n >= 8 {
|
|
//log.Printf("<- %x", b[:n])
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *udpConn) Write(b []byte) (n int, err error) {
|
|
//log.Printf("-> %x", b)
|
|
return c.UDPConn.WriteToUDP(b, 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] = magicTCP
|
|
copy(buf[8:], req)
|
|
//log.Printf("-> %x", buf)
|
|
_, err = c.TCPConn.Write(buf)
|
|
return
|
|
}
|
|
|
|
func newDataChannel(pushSize, popSize int) *dataChannel {
|
|
c := &dataChannel{}
|
|
if pushSize > 0 {
|
|
c.pushBuf = make(map[uint16][]byte, pushSize)
|
|
c.pushSize = pushSize
|
|
}
|
|
if popSize >= 0 {
|
|
c.popBuf = make(chan []byte, popSize)
|
|
}
|
|
return c
|
|
}
|
|
|
|
type dataChannel struct {
|
|
waitSeq uint16
|
|
pushBuf map[uint16][]byte
|
|
pushSize int
|
|
|
|
waitData []byte
|
|
waitSize int
|
|
popBuf chan []byte
|
|
}
|
|
|
|
func (c *dataChannel) Push(b []byte) error {
|
|
c.waitData = append(c.waitData, b...)
|
|
|
|
for len(c.waitData) > 4 {
|
|
// Every new data starts with size. There can be several data inside one packet.
|
|
if c.waitSize == 0 {
|
|
c.waitSize = int(binary.BigEndian.Uint32(c.waitData))
|
|
c.waitData = c.waitData[4:]
|
|
}
|
|
if c.waitSize > len(c.waitData) {
|
|
break
|
|
}
|
|
|
|
select {
|
|
case c.popBuf <- c.waitData[:c.waitSize]:
|
|
default:
|
|
return fmt.Errorf("pop buffer is full")
|
|
}
|
|
|
|
c.waitData = c.waitData[c.waitSize:]
|
|
c.waitSize = 0
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (c *dataChannel) Pop() ([]byte, bool) {
|
|
data, ok := <-c.popBuf
|
|
return data, ok
|
|
}
|
|
|
|
func (c *dataChannel) Close() {
|
|
close(c.popBuf)
|
|
}
|
|
|
|
// PushSeq returns how many seq were processed.
|
|
// Returns 0 if seq was saved or processed earlier.
|
|
// Returns -1 if seq could not be saved (buffer full or disabled).
|
|
func (c *dataChannel) PushSeq(seq uint16, data []byte) (int, error) {
|
|
diff := int16(seq - c.waitSeq)
|
|
// Check if this is seq from the future.
|
|
if diff > 0 {
|
|
// Support disabled buffer.
|
|
if c.pushSize == 0 {
|
|
return -1, nil // couldn't save seq
|
|
}
|
|
// Check if we don't have this seq in the buffer.
|
|
if c.pushBuf[seq] == nil {
|
|
// Check if there is enough space in the buffer.
|
|
if len(c.pushBuf) == c.pushSize {
|
|
return -1, nil // couldn't save seq
|
|
}
|
|
c.pushBuf[seq] = bytes.Clone(data)
|
|
//log.Printf("push buf wait=%d seq=%d len=%d", c.waitSeq, seq, len(c.pushBuf))
|
|
}
|
|
return 0, nil
|
|
}
|
|
|
|
// Check if this is seq from the past.
|
|
if diff < 0 {
|
|
return 0, nil
|
|
}
|
|
|
|
for i := 1; ; i++ {
|
|
if err := c.Push(data); err != nil {
|
|
return i, err
|
|
}
|
|
c.waitSeq++
|
|
// Check if we have next seq in the buffer.
|
|
if data = c.pushBuf[c.waitSeq]; data != nil {
|
|
delete(c.pushBuf, c.waitSeq)
|
|
} else {
|
|
return i, nil
|
|
}
|
|
}
|
|
}
|