github.com/Mrs4s/MiraiGo@v0.0.0-20240226124653-54bdd873e3fe/client/internal/tlv/decoder.go (about)

     1  package tlv
     2  
     3  import (
     4  	"encoding/binary"
     5  
     6  	"github.com/pkg/errors"
     7  )
     8  
     9  // Record represents a Tag-Length-Value record.
    10  type Record struct {
    11  	Tag    int
    12  	Length int
    13  	Value  []byte
    14  }
    15  
    16  type RecordMap map[int][]byte
    17  
    18  func (rm RecordMap) Exists(key int) bool {
    19  	_, ok := rm[key]
    20  	return ok
    21  }
    22  
    23  var ErrMessageTooShort = errors.New("tlv: message too short")
    24  
    25  // Decoder is a configurable TLV decoder.
    26  type Decoder struct {
    27  	tagSize  uint8
    28  	lenSize  uint8
    29  	headSize uint8
    30  }
    31  
    32  func NewDecoder(tagSize, lenSize uint8) *Decoder {
    33  	check := func(t string, s uint8) {
    34  		switch s {
    35  		case 1, 2, 4:
    36  			// ok
    37  		default:
    38  			panic("invalid " + t)
    39  		}
    40  	}
    41  	check("tag size", tagSize)
    42  	check("len size", lenSize)
    43  
    44  	return &Decoder{tagSize: tagSize, lenSize: lenSize, headSize: tagSize + lenSize}
    45  }
    46  
    47  func (d *Decoder) decodeRecord(data []byte) (r Record, err error) {
    48  	tagSize := d.tagSize
    49  	lenSize := d.lenSize
    50  	headSize := int(tagSize + lenSize)
    51  	if len(data) < headSize {
    52  		err = ErrMessageTooShort
    53  		return
    54  	}
    55  
    56  	r.Tag = d.read(tagSize, data)
    57  	r.Length = d.read(lenSize, data[tagSize:])
    58  
    59  	if len(data) < headSize+r.Length {
    60  		err = ErrMessageTooShort
    61  		return
    62  	}
    63  	r.Value = data[headSize : headSize+r.Length : headSize+r.Length]
    64  	return
    65  }
    66  
    67  func (d *Decoder) read(size uint8, data []byte) int {
    68  	switch size {
    69  	case 1:
    70  		return int(data[0])
    71  	case 2:
    72  		return int(binary.BigEndian.Uint16(data))
    73  	case 4:
    74  		return int(binary.BigEndian.Uint32(data))
    75  	default:
    76  		panic("invalid size")
    77  	}
    78  }
    79  
    80  func (d *Decoder) Decode(data []byte) ([]Record, error) {
    81  	var records []Record
    82  	for len(data) > 0 {
    83  		r, err := d.decodeRecord(data)
    84  		if err != nil {
    85  			return nil, err
    86  		}
    87  		records = append(records, r)
    88  		data = data[int(d.headSize)+r.Length:]
    89  	}
    90  	return records, nil
    91  }
    92  
    93  func (d *Decoder) DecodeRecordMap(data []byte) (RecordMap, error) {
    94  	records, err := d.Decode(data)
    95  	if err != nil {
    96  		return nil, err
    97  	}
    98  
    99  	rm := make(RecordMap, len(records))
   100  	for _, record := range records {
   101  		rm[record.Tag] = record.Value
   102  	}
   103  
   104  	return rm, nil
   105  }