t5de-proxy/internal/imq/messages.go

package imq

import "errors"

// Message types the proxy reads or writes.
const (
	C2gPing           = 1301
	C2gSendMessage    = 1304
	C2gOpenFloodgates = 1306
	C2gUnsubscribe    = 1308
	C2gSubscribe      = 1309
	G2cResult         = 3102
	G2cSendMessage    = 3104
	G2cStateChange    = 3105
	G2cJoinedQueue    = 3106
	G2cLeftQueue      = 3107
	G2cCreateMount    = 3108
)

// Mount types (G2cCreateMount).
const (
	MessageMount = 1
	StateMount   = 2
)

// 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
}

// QueueOp is a queue subscribed to with an operation ID, for a G2cResult.
type QueueOp struct {
	Queue string
	OpID  uint64
	raw   []byte // as the client sent it (with any other fields, such as properties)
}

// Subscription is a C2gSubscribe's or C2gUnsubscribe's data: queues listed plainly, and with operation IDs.
type Subscription struct {
	Queues      []string
	WithResults []QueueOp
}

// ParseSubscriptionOps reads a C2gSubscribe's or C2gUnsubscribe's data, keeping operation IDs.
func ParseSubscriptionOps(data []byte) (Subscription, error) {
	fs, err := fields(data)
	if err != nil {
		return Subscription{}, err
	}
	var s Subscription
	for _, f := range fs {
		switch {
		case f.num == 1 && f.isLen:
			s.Queues = append(s.Queues, string(f.bytes))
		case f.num == 2 && f.isLen:
			sub, err := fields(f.bytes)
			if err != nil {
				return Subscription{}, err
			}
			q := QueueOp{raw: f.bytes}
			for _, sf := range sub {
				switch {
				case sf.num == 1 && sf.isLen:
					q.Queue = string(sf.bytes)
				case sf.num == 2 && !sf.isLen:
					q.OpID = sf.varint
				}
			}
			s.WithResults = append(s.WithResults, q)
		default:
			return Subscription{}, errors.New("imq: unexpected field in a subscription")
		}
	}
	return s, nil
}

// All is every queue in the subscription.
func (s Subscription) All() []string {
	out := append([]string(nil), s.Queues...)
	for _, q := range s.WithResults {
		out = append(out, q.Queue)
	}
	return out
}

// Empty reports whether it names no queues.
func (s Subscription) Empty() bool { return len(s.Queues) == 0 && len(s.WithResults) == 0 }

// Marshal encodes the subscription's data.
func (s Subscription) Marshal() []byte {
	var b []byte
	for _, q := range s.Queues {
		b = appendBytesField(b, 1, []byte(q))
	}
	for _, q := range s.WithResults {
		raw := q.raw
		if raw == nil {
			raw = appendVarintField(appendBytesField(nil, 1, []byte(q.Queue)), 2, q.OpID)
		}
		b = appendBytesField(b, 2, raw)
	}
	return b
}

// Property is a mount's property: a key and a value.
type Property struct{ Key, Value string }

func appendProperties(b []byte, num uint64, props []Property) []byte {
	for _, p := range props {
		b = appendBytesField(b, num, appendBytesField(appendBytesField(nil, 1, []byte(p.Key)), 2, []byte(p.Value)))
	}
	return b
}

// JoinedQueue encodes a G2cJoinedQueue's (or G2cLeftQueue's) data: a user joined (left) a queue.
func JoinedQueue(userID, queue string, sequence uint64) []byte {
	return appendVarintField(appendBytesField(appendBytesField(nil, 1, []byte(userID)), 2, []byte(queue)), 3, sequence)
}

// CreateMount encodes a G2cCreateMount's data: a queue's mount, of type MessageMount or StateMount.
func CreateMount(queue, mount string, mountType uint64, props []Property, sequence uint64) []byte {
	b := appendVarintField(appendBytesField(appendBytesField(nil, 1, []byte(queue)), 2, []byte(mount)), 3, mountType)
	return appendVarintField(appendProperties(b, 4, props), 5, sequence)
}

// StateChange encodes a G2cStateChange's data: a state mount's properties changed (an empty value removes one).
func StateChange(opID uint64, userID, queue, mount string, props []Property, sequence uint64) []byte {
	b := appendVarintField(nil, 1, opID)
	b = appendBytesField(appendBytesField(appendBytesField(b, 2, []byte(userID)), 3, []byte(queue)), 4, []byte(mount))
	return appendVarintField(appendProperties(b, 5, props), 6, sequence)
}

// ParseFloodgates reads a C2gOpenFloodgates's data (after reconnecting, the queues to resume): each queue's state
// as a QueueOp (its op ID for a result), keeping the entry as sent.
func ParseFloodgates(data []byte) ([]QueueOp, error) {
	fs, err := fields(data)
	if err != nil {
		return nil, err
	}
	var out []QueueOp
	for _, f := range fs {
		if f.num != 1 || !f.isLen {
			return nil, errors.New("imq: unexpected field in C2gOpenFloodgates")
		}
		sub, err := fields(f.bytes)
		if err != nil {
			return nil, err
		}
		q := QueueOp{raw: f.bytes}
		for _, sf := range sub {
			switch {
			case sf.num == 1 && sf.isLen:
				q.Queue = string(sf.bytes)
			case sf.num == 3 && !sf.isLen:
				q.OpID = sf.varint
			}
		}
		out = append(out, q)
	}
	return out, nil
}

// MarshalFloodgates encodes a C2gOpenFloodgates's data from queue states read by ParseFloodgates.
func MarshalFloodgates(states []QueueOp) []byte {
	var b []byte
	for _, q := range states {
		b = appendBytesField(b, 1, q.raw)
	}
	return b
}