proto-socket/go/tcp_client.go

package toki_socket

import (
	"context"
	"crypto/tls"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	"google.golang.org/protobuf/proto"

	"toki-labs.com/toki_socket/go/packets"
)

const MaxPacketSize = 64 << 20

type TcpClient struct {
	Communicator
	conn                net.Conn
	writeMu             sync.Mutex // Defensive if WritePacket is called outside the queued writer.
	heartbeatInterval   time.Duration
	heartbeatWait       time.Duration
	waitingHBResponse   bool
	hbTimer             *HeartbeatTimer
	hbMu                sync.Mutex
	disconnectListeners []func(*TcpClient)
	disconnectMu        sync.Mutex
	tcpCloseOnce        sync.Once
}

func NewTcpClient(conn net.Conn, intervalSec, waitSec int, parserMap ParserMap) *TcpClient {
	c := &TcpClient{
		conn:              conn,
		heartbeatInterval: time.Duration(intervalSec) * time.Second,
		heartbeatWait:     time.Duration(waitSec) * time.Second,
	}
	c.Communicator.Initialize(c, parserMap)
	c.SetWriteErrorHandler(func(error) {
		c.onDisconnected()
	})
	c.AddListener(TypeNameOf(&packets.HeartBeat{}), func(m proto.Message) {
		c.onHeartBeat()
	})
	go c.readLoop()
	c.sendHeartBeat()
	return c
}

func DialTcp(ctx context.Context, host string, port int, intervalSec, waitSec int, parserMap ParserMap) (*TcpClient, error) {
	var d net.Dialer
	conn, err := d.DialContext(ctx, "tcp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		return nil, err
	}
	return NewTcpClient(conn, intervalSec, waitSec, parserMap), nil
}

func DialTcpTLS(ctx context.Context, host string, port int, tlsCfg *tls.Config, intervalSec, waitSec int, parserMap ParserMap) (*TcpClient, error) {
	d := tls.Dialer{Config: tlsCfg}
	conn, err := d.DialContext(ctx, "tcp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		return nil, err
	}
	return NewTcpClient(conn, intervalSec, waitSec, parserMap), nil
}

func (c *TcpClient) readLoop() {
	header := make([]byte, 4)
	for c.IsAlive() {
		if _, err := io.ReadFull(c.conn, header); err != nil {
			c.onDisconnected()
			return
		}
		length := binary.BigEndian.Uint32(header)
		if length == 0 {
			continue
		}
		if length > MaxPacketSize {
			c.onDisconnected()
			return
		}
		packetBytes := make([]byte, int(length))
		if _, err := io.ReadFull(c.conn, packetBytes); err != nil {
			c.onDisconnected()
			return
		}
		base := &packets.PacketBase{}
		if err := proto.Unmarshal(packetBytes, base); err != nil {
			c.onDisconnected()
			return
		}
		c.OnReceivedData(base.GetTypeName(), base.GetData(), base.GetNonce(), base.GetResponseNonce())
		c.sendHeartBeat()
	}
}

func (c *TcpClient) WritePacket(base *packets.PacketBase) error {
	b, err := proto.Marshal(base)
	if err != nil {
		return err
	}
	header := make([]byte, 4)
	binary.BigEndian.PutUint32(header, uint32(len(b)))

	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	if err := writeFull(c.conn, header); err != nil {
		return err
	}
	return writeFull(c.conn, b)
}

func (c *TcpClient) AddDisconnectListener(handler func(*TcpClient)) {
	c.disconnectMu.Lock()
	defer c.disconnectMu.Unlock()
	c.disconnectListeners = append(c.disconnectListeners, handler)
}

func (c *TcpClient) RemoveDisconnectListeners() {
	c.disconnectMu.Lock()
	defer c.disconnectMu.Unlock()
	c.disconnectListeners = nil
}

func (c *TcpClient) Close() error {
	var err error
	c.tcpCloseOnce.Do(func() {
		c.Communicator.shutdown()
		c.stopHeartbeat()
		err = c.conn.Close()
		c.notifyDisconnected()
	})
	return err
}

func (c *TcpClient) sendHeartBeat() {
	if !c.IsAlive() || c.heartbeatInterval <= 0 {
		return
	}
	c.hbMu.Lock()
	if c.hbTimer != nil {
		c.hbTimer.Stop()
	}
	c.hbTimer = NewHeartbeatTimer(c.heartbeatInterval, func() {
		if !c.IsAlive() {
			return
		}
		c.hbMu.Lock()
		c.waitingHBResponse = true
		c.hbMu.Unlock()
		_ = c.Send(&packets.HeartBeat{})
		c.hbMu.Lock()
		if c.hbTimer != nil {
			c.hbTimer.Stop()
		}
		c.hbTimer = NewHeartbeatTimer(c.heartbeatWait, func() {
			if c.IsAlive() {
				c.onDisconnected()
			}
		})
		c.hbMu.Unlock()
	})
	c.hbMu.Unlock()
}

func (c *TcpClient) onHeartBeat() {
	c.hbMu.Lock()
	if c.waitingHBResponse {
		c.waitingHBResponse = false
		c.hbMu.Unlock()
		return
	}
	c.hbMu.Unlock()
	_ = c.Send(&packets.HeartBeat{})
}

func (c *TcpClient) stopHeartbeat() {
	c.hbMu.Lock()
	defer c.hbMu.Unlock()
	if c.hbTimer != nil {
		c.hbTimer.Stop()
		c.hbTimer = nil
	}
}

func (c *TcpClient) onDisconnected() {
	_ = c.Close()
}

func (c *TcpClient) notifyDisconnected() {
	c.disconnectMu.Lock()
	listeners := append([]func(*TcpClient){}, c.disconnectListeners...)
	c.disconnectListeners = nil
	c.disconnectMu.Unlock()

	for _, listener := range listeners {
		listener(c)
	}
}

func writeFull(w io.Writer, b []byte) error {
	for len(b) > 0 {
		n, err := w.Write(b)
		if err != nil {
			return err
		}
		if n == 0 {
			return io.ErrShortWrite
		}
		b = b[n:]
	}
	return nil
}