live-gateway/bilibili/packer.go

package bilibili

import (
	"bytes"
	"compress/zlib"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/andybalholm/brotli"
	"github.com/gorilla/websocket"
	"io"
	"live-gateway/ws"
)

const (
	wsPackHeaderTotalLen = 16 // 头部字节大小
	wsPackageLen         = 4
	wsHeaderLen          = 2
	wsVerLen             = 2
	wsOpLen              = 4
	wsSequenceLen        = 4
)

const (
	WsHeaderDefaultSequence = 1
)

var _ ws.Packer = (*PackBilibili)(nil)

// PackBilibili bilibili packer
// |packageLen|headerLen|protoVer|operation|sequenceId|data|
//
// | segment         | type   | size    | offset | remark                                      |
// | ----------      | ------ | ------- | -------| ------------------------------------------- |
// | `packageLen`    | uint32 | 4       | 0      | the size of `data` and `header`             |
// | `headerLen`     | uint16 | 2       | 4      | 16                                          |
// | `protoVer`      | uint16 | 2       | 6      | 0:json 1:人气值 2:zlib.inflate 3:brotli      |
// | `operation`     | uint32 | 4       | 8      | 2:hb 3:hb-reply 5:danmaku 7:room 8:room-rep |
// | `sequenceId`    | uint32 | 4       | 12     | 1                                           |
// | `data`          | []byte | dynamic | 16     | :数据内容                                    |
// WsEntry -> ws.Entry#Raw
type PackBilibili struct {
}

func NewPackBilibili() ws.Packer {
	return &PackBilibili{}
}

func (*PackBilibili) byteOrder() binary.ByteOrder {
	return binary.BigEndian
}

// Unpack unpacks bilibili.WsEntry to ws.Entry
func (p *PackBilibili) Unpack(v interface{}) (*ws.Entry, error) {
	result := &ws.Entry{
		MessageType: websocket.BinaryMessage,
	}
	// v must be bilibili.WsEntry
	entry, ok := v.(*WsEntry)
	if !ok {
		return nil, fmt.Errorf("[Pack-Bilibili] 写入值类型必须为%T", WsEntry{})
	}

	header := make([]byte, 0, wsPackHeaderTotalLen)

	header = append(header, p.write(wsPackageLen, len(entry.data)+wsPackHeaderTotalLen)...)
	header = append(header, p.write(wsHeaderLen, wsPackHeaderTotalLen)...)
	header = append(header, p.write(wsVerLen, int(entry.protoVer))...)
	header = append(header, p.write(wsOpLen, int(entry.operation))...)
	header = append(header, p.write(wsSequenceLen, WsHeaderDefaultSequence)...)

	data := make([]byte, 0, len(entry.data)+wsPackHeaderTotalLen)
	data = append(data, header...)
	data = append(data, entry.data...)

	result.Raw = data
	return result, nil
}

// Pack packs ws.Entry to bilibili.WsEntry
func (p *PackBilibili) Pack(entry *ws.Entry) (interface{}, error) {
	if entry.MessageType != websocket.BinaryMessage || len(entry.Raw) < wsPackHeaderTotalLen {
		return nil, errors.New("err of msg")
	}
	var (
		offsetVer   = wsPackageLen + wsHeaderLen
		offsetOp    = offsetVer + wsVerLen
		offsetSeqId = offsetOp + wsOpLen
	)
	v := &WsEntry{}
	v.protoVer = p.byteOrder().Uint16(entry.Raw[offsetVer : offsetVer+wsVerLen])
	v.operation = p.byteOrder().Uint32(entry.Raw[offsetOp : offsetOp+wsOpLen])
	v.sequenceId = p.byteOrder().Uint32(entry.Raw[offsetSeqId : offsetSeqId+wsSequenceLen])
	v.data = entry.Raw[wsPackHeaderTotalLen:]

	if v.operation == WsOpMessage {
		switch v.protoVer {
		case WsVerZlib:
			de, err := p.zlibDecode(v.data)
			if err != nil {
				return nil, err
			}
			split := p.split(de)
			entries := p.combine(split, entry.MessageType)
			return entries, nil
		case WsVerBrotli:
			de, err := p.brotliDecode(v.data)
			if err != nil {
				return nil, err
			}
			split := p.split(de)
			entries := p.combine(split, entry.MessageType)
			return entries, nil
		case WsVerPlain:
			return v, nil
		}
	}
	return v, nil
}

func (p *PackBilibili) write(len, v int) []byte {
	b := make([]byte, len)
	switch len {
	case 2:
		p.byteOrder().PutUint16(b, uint16(v))
	case 4:
		p.byteOrder().PutUint32(b, uint32(v))
	}
	return b
}

func (*PackBilibili) zlibDecode(src []byte) ([]byte, error) {
	var (
		r   io.ReadCloser
		o   bytes.Buffer
		err error
	)
	if r, err = zlib.NewReader(bytes.NewReader(src)); err != nil {
		return nil, err
	}
	if _, err = io.Copy(&o, r); err != nil {
		return nil, err
	}
	return o.Bytes(), nil
}

func (*PackBilibili) brotliDecode(src []byte) ([]byte, error) {
	o := new(bytes.Buffer)
	r := brotli.NewReader(bytes.NewReader(src))
	if _, err := io.Copy(o, r); err != nil {
		return nil, err
	}
	return o.Bytes(), nil
}

// split 解压后的body需要拆包
func (*PackBilibili) split(b []byte) [][]byte {
	var packs [][]byte
	for i, size := uint32(0), uint32(0); i < uint32(len(b)); i += size {
		size = binary.BigEndian.Uint32(b[i : i+4])
		packs = append(packs, b[i:i+size])
	}
	return packs
}

func (*PackBilibili) combine(bs [][]byte, msgType int) []*ws.Entry {
	entries := make([]*ws.Entry, len(bs))
	for i, bs := range bs {
		entries[i] = &ws.Entry{
			MessageType: msgType,
			Raw:         bs,
		}
	}
	return entries
}