golib/gnet/modbus/conn.go

package modbus

import (
	"context"
	"errors"
	"fmt"
	"io"
	"math"
	"net"
	"strings"
	"sync"
	"time"

	"git.simanc.com/software/golib/v4/gio"
	"git.simanc.com/software/golib/v4/gnet"
	"git.simanc.com/software/golib/v4/log"
)

type PLC interface {
	gnet.ConnStat
	gnet.DataAccess
	io.Closer
}

var (
	ErrReadError    = errors.New("modbus: read error")
	ErrWriteError   = errors.New("modbus: write error")
	ErrReadTimeout  = errors.New("modbus: read timeout")
	ErrWriteTimeout = errors.New("modbus: write timeout")
	ErrConnError    = errors.New("modbus: connection error")
	ErrParamError   = errors.New("modbus: parameter error")
)

const (
	MaxReadBuffSize = 1024
)

// 一次连续读取寄存器的最大数量
const maxReadRegister = 30

type modbusConn struct {
	conn   net.Conn
	buf    gnet.Bytes
	logger log.Logger
	mu     sync.Mutex
}

func (w *modbusConn) IsConnected() bool {
	if w.conn == nil {
		return false
	}
	if conn, ok := w.conn.(gnet.ConnStat); ok {
		return conn.IsConnected()
	}
	return true
}

func (w *modbusConn) IsClosed() bool {
	if w.conn == nil {
		return true
	}
	if conn, ok := w.conn.(gnet.ConnStat); ok {
		return conn.IsClosed()
	}
	return false
}

// ReadData 读取原始数据, 当 count 过大时会自动分段读取
// 规则:
//
//	blockId == Code3, 表示读取保持寄存器, 每个寄存器大小为 2 个字节, count 为寄存器数量, 返回数据大小为 count*2
func (w *modbusConn) ReadData(ctx context.Context, blockId, address, count int) ([]byte, error) {
	if !w.IsConnected() || w.IsClosed() {
		return nil, gnet.ErrUnconnected
	}
	w.mu.Lock()
	defer w.mu.Unlock()

	switch blockId {
	case Code3:
		if !w.checkCode3(address, count) {
			return nil, ErrParamError
		}
	default:
		// TODO 目前仅支持 4x(Code03) 地址
		return nil, fmt.Errorf("modbus: ReadData: unsupported funCode: %d", blockId)
	}

	pduGroup := gnet.SplitNumber(count, maxReadRegister)

	aduList := make([]ADU, len(pduGroup))
	for i, length := range pduGroup { //
		curAddr := address + i*maxReadRegister
		pdu := NewPDUReadRegisters(byte(blockId), uint16(curAddr), uint16(length))
		aduList[i] = NewADU(uint16(i), Protocol, 0, pdu)
	}

	buf := make([]byte, count*2)
	for i, adu := range aduList {
		deadline, ok := ctx.Deadline()
		if !ok {
			deadline = time.Now().Add(gnet.ClientReadTimeout)
		}
		b, err := w.call(deadline, adu.Serialize())
		if err != nil {
			return nil, fmt.Errorf("modbus: ReadData: %s", err)
		}
		resp, err := ParseADU(b)
		if err != nil {
			return nil, fmt.Errorf("modbus: ReadData: ParseADU: %s", err)
		}
		if err = CheckADU(adu, resp); err != nil {
			return nil, fmt.Errorf("modbus: ReadData: CheckADU: %s", err)
		}
		copy(buf[maxReadRegister*2*i:], resp.PDU.Data)
	}

	return buf, nil
}

func (w *modbusConn) WriteData(ctx context.Context, blockId, address, count int, buf []byte) error {
	if !w.IsConnected() || w.IsClosed() {
		return gnet.ErrUnconnected
	}
	w.mu.Lock()
	defer w.mu.Unlock()
	switch blockId {
	case Code6, Code16:
		if !w.checkCode6(address, count, buf) {
			return ErrParamError
		}
	default:
		return fmt.Errorf("modbus: WriteData: unsupported funCode: %d", blockId)
	}
	var (
		pdu PDU
		err error
	)
	if count == 1 {
		pdu, err = NewPDUWriterSingleRegisterFromBuff(uint16(address), buf)
	} else {
		pdu, err = NewPDUWriterMultipleRegistersFromBuff(uint16(address), uint16(count), buf)
	}
	if err != nil {
		return errors.Join(ErrParamError, err)
	}
	adu := NewADU(uint16(address), Protocol, 0, pdu)
	deadline, ok := ctx.Deadline()
	if !ok {
		deadline = time.Now().Add(gnet.ClientReadTimeout)
	}
	b, err := w.call(deadline, adu.Serialize())
	if err != nil {
		return fmt.Errorf("modbus: WriteData: : %s", err)
	}
	resp, err := ParseADU(b)
	if err != nil {
		return fmt.Errorf("modbus: WriteData: ParseADU: %s", err)
	}
	if resp.TransactionID != adu.TransactionID {
		return fmt.Errorf("modbus: WriteData: transactionID mismatch: want %d, got %d", adu.TransactionID, resp.TransactionID)
	}
	return nil
}

func (w *modbusConn) GetProtocolName() string {
	return ProtocolName
}

func (w *modbusConn) checkCode3(address, count int) bool {
	return (address >= 0 && address <= math.MaxUint16) && (count > 0 && count <= math.MaxUint16)
}

func (w *modbusConn) checkCode6(address, count int, buf []byte) bool {
	return (address >= 0 && address <= math.MaxUint16) && (count > 0 && count <= math.MaxUint16) && (len(buf)/2 == count)
}

func (w *modbusConn) call(deadline time.Time, b gnet.Bytes) ([]byte, error) {
	if err := w.conn.SetDeadline(deadline); err != nil {
		w.logger.Error("modbus: call: failed to set deadline: %s", err)
		return nil, errors.Join(ErrConnError, err)
	}
	if _, err := w.conn.Write(b); err != nil {
		w.logger.Error("modbus: call: failed to write response: %s", err)
		if isNetTimeout(err) {
			return nil, errors.Join(ErrWriteTimeout, err)
		}
		return nil, errors.Join(ErrWriteError, err)
	}
	w.logger.Debug("modbus: Write: %s", b.HexTo())
	clear(w.buf)
	n, err := w.conn.Read(w.buf)
	if err != nil {
		w.logger.Error("modbus: call: failed to read response: %s", err)
		if isNetTimeout(err) {
			return nil, errors.Join(ErrReadTimeout, err)
		}
		return nil, errors.Join(ErrReadError, err)
	}
	data := w.buf[:n]
	w.logger.Debug("modbus: Read: %s", data.HexTo())
	return data, nil
}

func (w *modbusConn) Close() error {
	if w.conn == nil {
		return nil
	}
	return w.conn.Close()
}

func New(conn net.Conn, logger log.Logger) PLC {
	c := &modbusConn{
		conn:   conn,
		buf:    make([]byte, MaxReadBuffSize),
		logger: logger,
	}
	return c
}

// Conn PLC 主控连接
// Deprecated, 请使用 New
type Conn interface {
	// ConnStat 连接状态
	gnet.ConnStat
	// IsLocked 表示当前有其他线程正在与 PLC 交互
	IsLocked() bool
	// WriteResponse 向 PLC 发送数据并等待 PLC 响应
	WriteResponse(b []byte) ([]byte, error)
	// Closer 关闭与 PLC 主控的连接
	io.Closer
}

// Dialer
// Deprecated, 请使用 New
type Dialer struct {
	conn   net.Conn
	buf    []byte
	logger log.Logger
	mu     sync.Mutex
	lock   bool
}

func (w *Dialer) IsConnected() bool {
	if w.conn == nil {
		return false
	}
	return w.conn.(gnet.ConnStat).IsConnected()
}

func (w *Dialer) IsClosed() bool {
	if w.conn == nil {
		return true
	}
	return w.conn.(gnet.ConnStat).IsClosed()
}

func (w *Dialer) IsLocked() bool {
	return w.lock
}

// WriteResponse 写入并读取下一次的数据
func (w *Dialer) WriteResponse(b []byte) ([]byte, error) {
	if w.conn == nil {
		return nil, gnet.ErrConnNotFound
	}
	w.mu.Lock()
	defer w.mu.Unlock()
	w.lock = true
	defer func() {
		w.lock = false
	}()
	w.logger.Debug("Write: %s", gnet.Bytes(b).HexTo())
	if i, err := w.conn.Write(b); err != nil {
		w.logger.Error("Write err: %d->%d %s", len(b), i, err)
		if isNetTimeout(err) {
			return nil, errors.Join(ErrWriteTimeout, err)
		}
		return nil, err
	}
	clear(w.buf)
	n, err := w.conn.Read(w.buf)
	if err != nil {
		w.logger.Error("Read err: %s", err)
		if isNetTimeout(err) {
			return nil, errors.Join(ErrReadTimeout, err)
		}
		return nil, err
	}
	w.logger.Debug("Read: %s", gnet.Bytes(w.buf[:n]).HexTo())
	return w.buf[:n], nil
}

func (w *Dialer) Close() error {
	if w.conn == nil {
		return nil
	}
	return w.conn.Close()
}

func (w *Dialer) CloseWith(ctx context.Context) {
	<-ctx.Done()
	w.logger.Warn("DialContext: %s", ctx.Err())
	_ = w.Close()
}

func (w *Dialer) DialContext(ctx context.Context, address string, logger log.Logger) (Conn, error) {
	config := &gnet.Config{ // 由于现场网络环境比较差, 因此加大超时时间以防止频繁掉线重连
		Timeout:     60 * time.Second,
		DialTimeout: 10 * time.Second,
		Reconnect:   true,
	}
	return w.DialConfig(ctx, address, config, logger)
}

func (w *Dialer) DialConfig(ctx context.Context, address string, config *gnet.Config, logger log.Logger) (Conn, error) {
	deadline := time.Now().Add(config.DialTimeout)
	if conn, err := gnet.DialTCPConfig(address, config); err == nil {
		w.conn = conn
	} else {
		if timeout := deadline.Sub(time.Now()); timeout > 0 {
			gio.RandSleep(0, timeout)
		}
		logger.Debug("DialContext: %s", err)
		return nil, err
	}
	go func() {
		w.CloseWith(ctx)
	}()
	w.buf = make([]byte, MaxReadBuffSize)
	w.logger = log.Part(logger, "conn", strings.ReplaceAll(address, ":", "_"))
	return w, nil
}

func isNetTimeout(err error) bool {
	var ne *gnet.Timeout
	if errors.As(err, &ne) && ne.Timeout() {
		return true
	}
	return false
}

// DialContext
// Deprecated, 请使用 New
func DialContext(ctx context.Context, address string, logger log.Logger) (Conn, error) {
	var dialer Dialer
	return dialer.DialContext(ctx, address, logger)
}

// DialConfig
// Deprecated, 请使用 New
func DialConfig(ctx context.Context, address string, config *gnet.Config, logger log.Logger) (Conn, error) {
	var dialer Dialer
	return dialer.DialConfig(ctx, address, config, logger)
}

// Dial
// Deprecated, 请使用 New
func Dial(address string, logger log.Logger) (Conn, error) {
	return DialContext(context.Background(), address, logger)
}