t5de-proxy/internal/imq/messages.go

package imq

import "errors"

// Message types the proxy reads or writes.
const (
	C2gSendMessage = 1304
	C2gUnsubscribe = 1308
	C2gSubscribe   = 1309
	G2cResult      = 3102
	G2cSendMessage = 3104
)

// SendMessage is a C2gSendMessage: a message the client sends to a queue's mount.
type SendMessage struct {
	OpID    uint64
	Queue   string
	Mount   string
	Message []byte
}

// field is one protobuf field: its number and either a varint or bytes.
type field struct {
	num    uint64
	varint uint64
	bytes  []byte
	isLen  bool
}

func fields(b []byte) ([]field, error) {
	var out []field
	for i := 0; i < len(b); {
		key, n := varint(b[i:])
		if n <= 0 {
			return nil, ErrMalformed
		}
		i += n
		f := field{num: key >> 3}
		switch key & 7 {
		case 0:
			v, n := varint(b[i:])
			if n <= 0 {
				return nil, ErrMalformed
			}
			f.varint, i = v, i+n
		case 2:
			l, n := varint(b[i:])
			if n <= 0 || l > uint64(len(b)-i-n) {
				return nil, ErrMalformed
			}
			i += n
			f.bytes, f.isLen, i = b[i:i+int(l)], true, i+int(l)
		default:
			return nil, ErrMalformed
		}
		out = append(out, f)
	}
	return out, nil
}

// ParseSendMessage reads a C2gSendMessage's data.
func ParseSendMessage(data []byte) (SendMessage, error) {
	fs, err := fields(data)
	if err != nil {
		return SendMessage{}, err
	}
	var m SendMessage
	seen := 0
	for _, f := range fs {
		switch {
		case f.num == 1 && !f.isLen:
			m.OpID, seen = f.varint, seen|1
		case f.num == 2 && f.isLen:
			m.Queue, seen = string(f.bytes), seen|2
		case f.num == 3 && f.isLen:
			m.Mount, seen = string(f.bytes), seen|4
		case f.num == 4 && f.isLen:
			m.Message, seen = f.bytes, seen|8
		default:
			return SendMessage{}, errors.New("imq: unexpected field in C2gSendMessage")
		}
	}
	if seen != 15 {
		return SendMessage{}, errors.New("imq: C2gSendMessage is missing fields")
	}
	return m, nil
}

// Marshal encodes the message's data, fields in the client's order.
func (m SendMessage) Marshal() []byte {
	var b []byte
	b = appendVarintField(b, 1, m.OpID)
	b = appendBytesField(b, 2, []byte(m.Queue))
	b = appendBytesField(b, 3, []byte(m.Mount))
	return appendBytesField(b, 4, m.Message)
}

// Result encodes a G2cResult's data: the answer to the operation opID (status 0 is success).
func Result(opID, status uint64) []byte {
	return appendVarintField(appendVarintField(nil, 1, opID), 2, status)
}

// EncodeFrame wraps a message's data in a frame of the given type.
func EncodeFrame(t uint64, data []byte) []byte {
	return appendBytesField(appendVarintField(nil, 1, t), 2, data)
}

func appendVarint(b []byte, v uint64) []byte {
	for v >= 0x80 {
		b = append(b, byte(v)|0x80)
		v >>= 7
	}
	return append(b, byte(v))
}

func appendVarintField(b []byte, num, v uint64) []byte {
	return appendVarint(appendVarint(b, num<<3), v)
}

func appendBytesField(b []byte, num uint64, v []byte) []byte {
	b = appendVarint(appendVarint(b, num<<3|2), uint64(len(v)))
	return append(b, v...)
}

// ParseSubscription reads a C2gSubscribe's or C2gUnsubscribe's data: the queues (listed plainly, or with an op ID
// for a result).
func ParseSubscription(data []byte) ([]string, error) {
	fs, err := fields(data)
	if err != nil {
		return nil, err
	}
	var queues []string
	for _, f := range fs {
		switch {
		case f.num == 1 && f.isLen:
			queues = append(queues, string(f.bytes))
		case f.num == 2 && f.isLen:
			sub, err := fields(f.bytes)
			if err != nil {
				return nil, err
			}
			for _, sf := range sub {
				if sf.num == 1 && sf.isLen {
					queues = append(queues, string(sf.bytes))
				}
			}
		default:
			return nil, errors.New("imq: unexpected field in a subscription")
		}
	}
	return queues, nil
}

// Delivered is a G2cSendMessage: a message the gateway delivers to the client, from a user (user_id).
type Delivered struct {
	UserID   string
	Queue    string
	Mount    string
	Message  []byte
	Sequence uint64 // 0: none (the client only tracks sequences it's given)
}

// ParseDelivered reads a G2cSendMessage's data.
func ParseDelivered(data []byte) (Delivered, error) {
	fs, err := fields(data)
	if err != nil {
		return Delivered{}, err
	}
	var m Delivered
	seen := 0
	for _, f := range fs {
		switch {
		case f.num == 1 && f.isLen:
			m.UserID, seen = string(f.bytes), seen|1
		case f.num == 2 && f.isLen:
			m.Queue, seen = string(f.bytes), seen|2
		case f.num == 3 && f.isLen:
			m.Mount, seen = string(f.bytes), seen|4
		case f.num == 4 && f.isLen:
			m.Message, seen = f.bytes, seen|8
		case f.num == 5 && !f.isLen:
			m.Sequence = f.varint
		default:
			return Delivered{}, errors.New("imq: unexpected field in G2cSendMessage")
		}
	}
	if seen != 15 {
		return Delivered{}, errors.New("imq: G2cSendMessage is missing fields")
	}
	return m, nil
}

// Marshal encodes the message's data, fields in the gateway's order.
func (m Delivered) Marshal() []byte {
	var b []byte
	b = appendBytesField(b, 1, []byte(m.UserID))
	b = appendBytesField(b, 2, []byte(m.Queue))
	b = appendBytesField(b, 3, []byte(m.Mount))
	b = appendBytesField(b, 4, m.Message)
	if m.Sequence != 0 {
		b = appendVarintField(b, 5, m.Sequence)
	}
	return b
}