153 lines
4.9 KiB
Go
153 lines
4.9 KiB
Go
package orbit
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/base64"
|
||
"encoding/json"
|
||
"reflect"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/coder/websocket"
|
||
)
|
||
|
||
type eventFrame struct {
|
||
Version string `json:"protocol_version"`
|
||
Type string `json:"type"`
|
||
MessageID string `json:"message_id"`
|
||
ReplyTo string `json:"reply_to,omitempty"`
|
||
SessionID string `json:"session_id,omitempty"`
|
||
Body json.RawMessage `json:"body,omitempty"`
|
||
Error *EventError `json:"error,omitempty"`
|
||
}
|
||
|
||
// 仅协议结构严格匹配字段;业务RawMessage保持原数字、键和值。
|
||
func eventFields(raw []byte, typ reflect.Type) error {
|
||
if typ == reflect.TypeOf(json.RawMessage{}) {
|
||
return nil
|
||
}
|
||
if bytes.Equal(bytes.TrimSpace(raw), []byte("null")) {
|
||
return ErrEventProtocol
|
||
}
|
||
if typ.Kind() == reflect.Pointer {
|
||
return eventFields(raw, typ.Elem())
|
||
}
|
||
if typ == reflect.TypeOf(time.Time{}) {
|
||
return nil
|
||
}
|
||
if typ.Kind() != reflect.Struct {
|
||
return nil
|
||
}
|
||
var fields map[string]json.RawMessage
|
||
if json.Unmarshal(raw, &fields) != nil || fields == nil {
|
||
return ErrEventProtocol
|
||
}
|
||
for i := 0; i < typ.NumField(); i++ {
|
||
f := typ.Field(i)
|
||
tags := strings.Split(f.Tag.Get("json"), ",")
|
||
if tags[0] == "" || tags[0] == "-" {
|
||
continue
|
||
}
|
||
v, found := fields[tags[0]]
|
||
optional, quoted := false, false
|
||
for _, tag := range tags[1:] {
|
||
optional = optional || tag == "omitempty" || tag == "omitzero"
|
||
quoted = quoted || tag == "string"
|
||
}
|
||
if !found {
|
||
if optional {
|
||
continue
|
||
}
|
||
return ErrEventProtocol
|
||
}
|
||
delete(fields, tags[0])
|
||
if f.Type.Kind() == reflect.Uint64 && quoted {
|
||
var s string
|
||
if json.Unmarshal(v, &s) != nil {
|
||
return ErrEventProtocol
|
||
}
|
||
n, err := strconv.ParseUint(s, 10, 64)
|
||
if err != nil || n == 0 || strconv.FormatUint(n, 10) != s {
|
||
return ErrEventProtocol
|
||
}
|
||
}
|
||
if err := eventFields(v, f.Type); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
if len(fields) > 0 {
|
||
return ErrEventProtocol
|
||
}
|
||
return nil
|
||
}
|
||
func decodeEventValue(raw []byte, out any) error {
|
||
if validateWorkerJSON(raw) != nil || eventFields(raw, reflect.TypeOf(out).Elem()) != nil || json.Unmarshal(raw, out) != nil {
|
||
return ErrEventProtocol
|
||
}
|
||
return nil
|
||
}
|
||
func parseEventFrame(kind websocket.MessageType, raw []byte) (eventFrame, error) {
|
||
var f eventFrame
|
||
if kind != websocket.MessageText || len(raw) > 512<<10 || decodeEventValue(raw, &f) != nil || f.Version != "1" || !validID(f.MessageID) {
|
||
return f, ErrEventProtocol
|
||
}
|
||
var fields map[string]json.RawMessage
|
||
_ = json.Unmarshal(raw, &fields)
|
||
if fields["reply_to"] != nil && !validID(f.ReplyTo) || fields["session_id"] != nil && !validID(f.SessionID) {
|
||
return f, ErrEventProtocol
|
||
}
|
||
switch f.Type {
|
||
case "error":
|
||
if f.Error == nil || fields["body"] != nil || len(f.Error.Message) == 0 || len(f.Error.Message) > 1024 {
|
||
return f, ErrEventProtocol
|
||
}
|
||
switch f.Error.Code {
|
||
case "INVALID_ARGUMENT", "UNAUTHENTICATED", "PERMISSION_DENIED", "NOT_FOUND", "STATE_CONFLICT", "DELIVERY_LEASE_LOST", "EVENT_HISTORY_EXPIRED", "UNAVAILABLE":
|
||
default:
|
||
return f, ErrEventProtocol
|
||
}
|
||
case "events.welcome", "ack", "event.deliver":
|
||
if fields["error"] != nil || len(f.Body) == 0 || bytes.Equal(bytes.TrimSpace(f.Body), []byte("null")) || !validID(f.SessionID) {
|
||
return f, ErrEventProtocol
|
||
}
|
||
if f.Type == "event.deliver" {
|
||
if fields["reply_to"] != nil {
|
||
return f, ErrEventProtocol
|
||
}
|
||
} else if !validID(f.ReplyTo) {
|
||
return f, ErrEventProtocol
|
||
}
|
||
default:
|
||
return f, ErrEventProtocol
|
||
}
|
||
return f, nil
|
||
}
|
||
func validEventProof(p EventProof) bool {
|
||
b, err := base64.RawURLEncoding.Strict().DecodeString(p.LeaseSecret)
|
||
return validID(p.SubscriptionID) && validID(p.EventID) && validID(p.DeliveryID) && p.Epoch > 0 && len(p.LeaseSecret) == 43 && err == nil && len(b) == 32
|
||
}
|
||
func validEventGrant(g EventGrant, sub string) bool {
|
||
d, e, s := g.Delivery, g.Event, g.Event.Snapshot
|
||
if !validEventProof(g.Proof()) || d.SubscriptionID != sub || d.EventID != e.ID || d.Sequence != e.Sequence || d.Sequence == 0 || d.Revision == 0 || d.State != DeliveryLeased || d.CreatedAt.IsZero() || d.AvailableAt.IsZero() || d.LeaseUntil == nil || d.LeaseUntil.IsZero() || d.AcknowledgedAt != nil || d.Attempts == 0 || d.Attempts > d.AttemptBudget || d.AttemptBudget > 1000000 {
|
||
return false
|
||
}
|
||
if e.AggregateKind != "task" || !validID(e.AggregateID) || e.AggregateID != s.TaskID || e.AggregateRevision != s.Revision || s.Revision == 0 || e.OccurredAt.IsZero() || s.SchemaVersion != 1 || !validID(s.QueueID) || s.CurrentAttemptID != nil && !validID(*s.CurrentAttemptID) {
|
||
return false
|
||
}
|
||
if len(bytes.TrimSpace(e.Details)) == 0 || bytes.TrimSpace(e.Details)[0] != '{' {
|
||
return false
|
||
}
|
||
switch e.Type {
|
||
case "task.succeeded":
|
||
return s.State == TaskSucceeded
|
||
case "task.failed":
|
||
return s.State == TaskFailed
|
||
case "task.cancelled":
|
||
return s.State == TaskCancelled
|
||
case "task.needs_review":
|
||
return s.State == TaskNeedsReview
|
||
}
|
||
return false
|
||
}
|