119 lines
2.2 KiB
Go
119 lines
2.2 KiB
Go
package protocol
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"fmt"
|
|
"io"
|
|
)
|
|
|
|
type FrameReadWriter interface {
|
|
ReadFrame() (*Frame, error)
|
|
WriteFrame(opCode byte, payload []byte) error
|
|
}
|
|
|
|
const maxFramePayloadSize = 1 << 24
|
|
|
|
type Frame struct {
|
|
OpCode byte
|
|
Payload []byte
|
|
}
|
|
|
|
func EncodeFrame(opCode byte, payload []byte) []byte {
|
|
payloadLen := uint64(len(payload))
|
|
buf := make([]byte, 1+varintLen(payloadLen)+len(payload))
|
|
buf[0] = opCode
|
|
off := 1
|
|
off += encodeVarint(buf[off:], uint64(payloadLen))
|
|
copy(buf[off:], payload)
|
|
return buf
|
|
}
|
|
|
|
func ReadFrame(r io.Reader) (*Frame, error) {
|
|
opBuf := make([]byte, 1)
|
|
if _, err := io.ReadFull(r, opBuf); err != nil {
|
|
return nil, fmt.Errorf("read opcode: %w", err)
|
|
}
|
|
opCode := opBuf[0]
|
|
|
|
payloadLen, err := readVarint(r)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("read payload length: %w", err)
|
|
}
|
|
if payloadLen > maxFramePayloadSize {
|
|
return nil, fmt.Errorf("payload too large: %d", payloadLen)
|
|
}
|
|
|
|
payload := make([]byte, payloadLen)
|
|
if payloadLen > 0 {
|
|
if _, err := io.ReadFull(r, payload); err != nil {
|
|
return nil, fmt.Errorf("read payload: %w", err)
|
|
}
|
|
}
|
|
|
|
return &Frame{OpCode: opCode, Payload: payload}, nil
|
|
}
|
|
|
|
func WriteFrame(w io.Writer, opCode byte, payload []byte) error {
|
|
data := EncodeFrame(opCode, payload)
|
|
_, err := w.Write(data)
|
|
return err
|
|
}
|
|
|
|
func encodeVarint(buf []byte, v uint64) int {
|
|
i := 0
|
|
for v >= 0x80 {
|
|
buf[i] = byte(v) | 0x80
|
|
v >>= 7
|
|
i++
|
|
}
|
|
buf[i] = byte(v)
|
|
return i + 1
|
|
}
|
|
|
|
func varintLen(v uint64) int {
|
|
l := 0
|
|
for {
|
|
l++
|
|
if v < 0x80 {
|
|
break
|
|
}
|
|
v >>= 7
|
|
}
|
|
return l
|
|
}
|
|
|
|
func readVarint(r io.Reader) (uint64, error) {
|
|
var v uint64
|
|
var shift uint
|
|
for {
|
|
b := make([]byte, 1)
|
|
if _, err := io.ReadFull(r, b); err != nil {
|
|
return 0, err
|
|
}
|
|
v |= uint64(b[0]&0x7F) << shift
|
|
if b[0]&0x80 == 0 {
|
|
break
|
|
}
|
|
shift += 7
|
|
if shift > 63 {
|
|
return 0, fmt.Errorf("varint too long")
|
|
}
|
|
}
|
|
return v, nil
|
|
}
|
|
|
|
func ReadUint64(buf []byte) uint64 {
|
|
return binary.BigEndian.Uint64(buf)
|
|
}
|
|
|
|
func PutUint64(buf []byte, v uint64) {
|
|
binary.BigEndian.PutUint64(buf, v)
|
|
}
|
|
|
|
func ReadUint32(buf []byte) uint32 {
|
|
return binary.BigEndian.Uint32(buf)
|
|
}
|
|
|
|
func PutUint32(buf []byte, v uint32) {
|
|
binary.BigEndian.PutUint32(buf, v)
|
|
}
|