tele-recv/utils/socket.go

package utils

import (
	"io"
	"net"
	"time"

	"go.uber.org/zap"
)

const ServerType = "tcp"

var ServerRunning = false
var client net.Conn
var server net.Listener
var ClientReady = false

func Listen(address string) {
	var err error
	server, err = net.Listen(ServerType, address)
	if err != nil {
		Log.Fatal("error in create server ", zap.Error(err))
	}
	defer server.Close()
	Log.Info("listen on ", zap.String("address", address))

	acceptFailCount := 0
	const maxAcceptFail = 3

	for ServerRunning {
		conn, err := server.Accept()
		if err != nil {
			acceptFailCount++
			Log.Error("error in create connection ",
				zap.Error(err),
				zap.Int("failCount", acceptFailCount))
			if acceptFailCount >= maxAcceptFail {
				ServerRunning = false
				Log.Fatal("accept failed too many times, shutting down")
			}
			continue
		}
		acceptFailCount = 0

		if client == nil {
			client = conn
			ClientReady = false
			Log.Info("client connected on ", zap.Any("address", conn.RemoteAddr()))
			LoadUnprocessed()
		} else {
			Log.Info("remove old client")
			closeClient()
			client = conn
		}
		if IsClientConnected() {
			connCheck()
		}
		time.Sleep(2 * time.Second)
	}
}

func WriteToClient(data string) error {
	if client != nil {
		_, err := client.Write([]byte(data + "\r\n"))
		if err != nil {
			Log.Error("error in write data to client socket ", zap.Error(err))
			closeClient()
		}
		return err
	}
	return nil
}

func connCheck() bool {
	_, err := client.Read(make([]byte, 0))
	if err != nil && err != io.EOF {
		// this connection is invalid
		Log.Error("conn closed....", zap.Error(err))
		client = nil
		return false
	}
	return true
}

func IsClientConnected() bool {
	return client != nil
}

func closeClient() {
	if client != nil {
		client.Close()
		client = nil
	}
}

func StopSocketServer() {
	if IsClientConnected() {
		client.Close()
	}
	if server != nil {
		server.Close()
	}
}