ntool-biz/nmodbus/protocol/packer_tcp.go

package protocol

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"git.noahlan.cn/noahlan/nnet/packet"
	"git.noahlan.cn/noahlan/ntool/nlog"
)

type TCPPackerOption func(*TCPPacker)

type TCPPacker struct {
	buf       *bytes.Buffer
	header    *TCPHeader
	size      int // 最近一次读取的 body size
	byteOrder binary.ByteOrder
}

const (
	tcpHeadLength    = 2 + 2 + 2
	tcpMaxPacketSize = 64 * 1024
)

func NewTCPPacker(byteOrder binary.ByteOrder) *TCPPacker {
	p := &TCPPacker{
		buf:       bytes.NewBuffer(nil),
		byteOrder: byteOrder,
	}

	p.resetFlags()

	return p
}

func (d *TCPPacker) resetFlags() {
	d.header = nil
	d.size = -1
}

func (d *TCPPacker) Pack(header interface{}, data []byte) ([]byte, error) {
	tcpHeader, ok := header.(*TCPHeader)
	if !ok {
		return nil, packet.ErrWrongPacketType
	}
	bs := make([]byte, 8)
	d.byteOrder.PutUint16(bs[0:2], tcpHeader.TransactionIdentifier)
	d.byteOrder.PutUint16(bs[2:4], tcpHeader.ProtocolIdentifier)
	d.byteOrder.PutUint16(bs[4:6], uint16(2+len(data)))
	bs[6] = tcpHeader.Address
	bs[7] = tcpHeader.Function
	bs = append(bs, data...)
	return bs, nil
}

func (d *TCPPacker) Unpack(data []byte) ([]packet.IPacket, error) {
	d.buf.Write(data)

	nlog.Debugf("接收TCP数据: %x", data)
	var (
		packets []packet.IPacket
		err     error
	)

	if d.buf.Len() < 9 {
		return packets, fmt.Errorf("TCP Frame error: packet less than 9 bytes")
	}

	if d.size < 0 {
		if err = d.readHeader(); err != nil {
			return nil, err
		}
	}

	for (d.size - 2) <= d.buf.Len() {
		p := NewTCPPacket()
		p.TCPHeader = d.header
		p.Data = d.buf.Next(d.size - 2)

		packets = append(packets, p)

		if d.buf.Len() < (tcpHeadLength + 2) {
			d.resetFlags()
			break
		}

		if err = d.readHeader(); err != nil {
			return packets, err
		}
	}

	if packets == nil || len(packets) == 0 {
		d.resetFlags()
		d.buf.Reset()
	}

	return packets, err
}

func (d *TCPPacker) readHeader() error {
	header := d.buf.Next(tcpHeadLength + 2)
	d.header = &TCPHeader{
		TransactionIdentifier: d.byteOrder.Uint16(header[0:2]),
		ProtocolIdentifier:    d.byteOrder.Uint16(header[2:4]),
		Length:                d.byteOrder.Uint16(header[4:6]),
		ModbusHeader: ModbusHeader{
			Address:  header[6],
			Function: header[7],
		},
	}
	d.size = int(d.header.Length)
	// 最大包限定
	if d.size > tcpMaxPacketSize {
		return packet.ErrPacketSizeExceed
	}
	return nil
}