package lexutil import ( "bytes" "context" "fmt" "io" "net/http" "net/http/httptest" "strconv" "sync" "sync/atomic" "testing" "time" "github.com/bluesky-social/indigo/events" indigoxrpc "github.com/bluesky-social/indigo/xrpc" "github.com/gorilla/websocket" cid "github.com/ipfs/go-cid" "github.com/stretchr/testify/assert" cbg "github.com/whyrusleeping/cbor-gen" xerrors "golang.org/x/xerrors" ) type SubscribeExample_Foo struct { Seq int64 Foo string } type SubscribeExample_Bar struct { Seq int64 Bar string } type SubscribeExample_Event struct { Error *events.ErrorFrame Foo *SubscribeExample_Foo Bar *SubscribeExample_Bar } func (e *SubscribeExample_Event) seq() int64 { switch { case e.Foo != nil: return e.Foo.Seq case e.Bar != nil: return e.Bar.Seq default: return 0 } } func (e *SubscribeExample_Event) Serialize(w io.Writer) error { cw := cbg.NewCborWriter(w) header := events.EventHeader{Op: events.EvtKindMessage} switch { case e.Error != nil: header.Op = events.EvtKindErrorFrame if err := header.MarshalCBOR(cw); err != nil { return err } return e.Error.MarshalCBOR(cw) case e.Foo != nil: header.MsgType = "#foo" if err := header.MarshalCBOR(cw); err != nil { return err } return e.Foo.MarshalCBOR(cw) case e.Bar != nil: header.MsgType = "#bar" if err := header.MarshalCBOR(cw); err != nil { return err } return e.Bar.MarshalCBOR(cw) default: return fmt.Errorf("unrecognized event kind") } } func (e *SubscribeExample_Event) Deserialize(r io.Reader) error { var header events.EventHeader if err := header.UnmarshalCBOR(r); err != nil { return fmt.Errorf("reading header: %w", err) } switch header.Op { case events.EvtKindMessage: switch header.MsgType { case "#foo": var evt SubscribeExample_Foo if err := evt.UnmarshalCBOR(r); err != nil { return err } e.Foo = &evt case "#bar": var evt SubscribeExample_Bar if err := evt.UnmarshalCBOR(r); err != nil { return err } e.Bar = &evt default: return fmt.Errorf("unknown message type: %s", header.MsgType) } case events.EvtKindErrorFrame: var errframe events.ErrorFrame if err := errframe.UnmarshalCBOR(r); err != nil { return err } e.Error = &errframe default: return fmt.Errorf("unrecognized event stream type: %d", header.Op) } return nil } func (t *SubscribeExample_Foo) MarshalCBOR(w io.Writer) error { if t == nil { _, err := w.Write(cbg.CborNull) return err } cw := cbg.NewCborWriter(w) if _, err := cw.Write([]byte{162}); err != nil { return err } // t.Foo (string) (string) if len("foo") > 1000000 { return xerrors.Errorf("Value in field \"foo\" was too long") } if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len("foo"))); err != nil { return err } if _, err := cw.WriteString(string("foo")); err != nil { return err } if len(t.Foo) > 1000000 { return xerrors.Errorf("Value in field t.Foo was too long") } if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len(t.Foo))); err != nil { return err } if _, err := cw.WriteString(string(t.Foo)); err != nil { return err } // t.Seq (int64) (int64) if len("seq") > 1000000 { return xerrors.Errorf("Value in field \"seq\" was too long") } if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len("seq"))); err != nil { return err } if _, err := cw.WriteString(string("seq")); err != nil { return err } if t.Seq >= 0 { if err := cw.WriteMajorTypeHeader(cbg.MajUnsignedInt, uint64(t.Seq)); err != nil { return err } } else { if err := cw.WriteMajorTypeHeader(cbg.MajNegativeInt, uint64(-t.Seq-1)); err != nil { return err } } return nil } func (t *SubscribeExample_Foo) UnmarshalCBOR(r io.Reader) (err error) { *t = SubscribeExample_Foo{} cr := cbg.NewCborReader(r) maj, extra, err := cr.ReadHeader() if err != nil { return err } defer func() { if err == io.EOF { err = io.ErrUnexpectedEOF } }() if maj != cbg.MajMap { return fmt.Errorf("cbor input should be of type map") } if extra > cbg.MaxLength { return fmt.Errorf("SubscribeExample_Foo: map struct too large (%d)", extra) } n := extra nameBuf := make([]byte, 8) for range n { nameLen, ok, err := cbg.ReadFullStringIntoBuf(cr, nameBuf, 1000000) if err != nil { return err } if !ok { // Field doesn't exist on this type, so ignore it if err := cbg.ScanForLinks(cr, func(cid.Cid) {}); err != nil { return err } continue } switch string(nameBuf[:nameLen]) { // t.Foo (string) (string) case "foo": { sval, err := cbg.ReadStringWithMax(cr, 1000000) if err != nil { return err } t.Foo = string(sval) } // t.Seq (int64) (int64) case "seq": { maj, extra, err := cr.ReadHeader() if err != nil { return err } var extraI int64 switch maj { case cbg.MajUnsignedInt: extraI = int64(extra) if extraI < 0 { return fmt.Errorf("int64 positive overflow") } case cbg.MajNegativeInt: extraI = int64(extra) if extraI < 0 { return fmt.Errorf("int64 negative overflow") } extraI = -1 - extraI default: return fmt.Errorf("wrong type for int64 field: %d", maj) } t.Seq = int64(extraI) } default: // Field doesn't exist on this type, so ignore it if err := cbg.ScanForLinks(r, func(cid.Cid) {}); err != nil { return err } } } return nil } func (t *SubscribeExample_Bar) MarshalCBOR(w io.Writer) error { if t == nil { _, err := w.Write(cbg.CborNull) return err } cw := cbg.NewCborWriter(w) if _, err := cw.Write([]byte{162}); err != nil { return err } // t.Bar (string) (string) if len("bar") > 1000000 { return xerrors.Errorf("Value in field \"bar\" was too long") } if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len("bar"))); err != nil { return err } if _, err := cw.WriteString(string("bar")); err != nil { return err } if len(t.Bar) > 1000000 { return xerrors.Errorf("Value in field t.Bar was too long") } if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len(t.Bar))); err != nil { return err } if _, err := cw.WriteString(string(t.Bar)); err != nil { return err } // t.Seq (int64) (int64) if len("seq") > 1000000 { return xerrors.Errorf("Value in field \"seq\" was too long") } if err := cw.WriteMajorTypeHeader(cbg.MajTextString, uint64(len("seq"))); err != nil { return err } if _, err := cw.WriteString(string("seq")); err != nil { return err } if t.Seq >= 0 { if err := cw.WriteMajorTypeHeader(cbg.MajUnsignedInt, uint64(t.Seq)); err != nil { return err } } else { if err := cw.WriteMajorTypeHeader(cbg.MajNegativeInt, uint64(-t.Seq-1)); err != nil { return err } } return nil } func (t *SubscribeExample_Bar) UnmarshalCBOR(r io.Reader) (err error) { *t = SubscribeExample_Bar{} cr := cbg.NewCborReader(r) maj, extra, err := cr.ReadHeader() if err != nil { return err } defer func() { if err == io.EOF { err = io.ErrUnexpectedEOF } }() if maj != cbg.MajMap { return fmt.Errorf("cbor input should be of type map") } if extra > cbg.MaxLength { return fmt.Errorf("SubscribeExample_Bar: map struct too large (%d)", extra) } n := extra nameBuf := make([]byte, 8) for range n { nameLen, ok, err := cbg.ReadFullStringIntoBuf(cr, nameBuf, 1000000) if err != nil { return err } if !ok { // Field doesn't exist on this type, so ignore it if err := cbg.ScanForLinks(cr, func(cid.Cid) {}); err != nil { return err } continue } switch string(nameBuf[:nameLen]) { // t.Bar (string) (string) case "bar": { sval, err := cbg.ReadStringWithMax(cr, 1000000) if err != nil { return err } t.Bar = string(sval) } // t.Seq (int64) (int64) case "seq": { maj, extra, err := cr.ReadHeader() if err != nil { return err } var extraI int64 switch maj { case cbg.MajUnsignedInt: extraI = int64(extra) if extraI < 0 { return fmt.Errorf("int64 positive overflow") } case cbg.MajNegativeInt: extraI = int64(extra) if extraI < 0 { return fmt.Errorf("int64 negative overflow") } extraI = -1 - extraI default: return fmt.Errorf("wrong type for int64 field: %d", maj) } t.Seq = int64(extraI) } default: // Field doesn't exist on this type, so ignore it if err := cbg.ScanForLinks(r, func(cid.Cid) {}); err != nil { return err } } } return nil } type testScheduler struct { ch chan *SubscribeExample_Event lastSeq atomic.Int64 } var _ SeqScheduler[SubscribeExample_Event] = (*testScheduler)(nil) func newTestScheduler() *testScheduler { return &testScheduler{ch: make(chan *SubscribeExample_Event, 8)} } func (s *testScheduler) AddWork(ctx context.Context, namespace string, val *SubscribeExample_Event) error { s.lastSeq.Store(val.seq()) select { case s.ch <- val: case <-ctx.Done(): } return nil } func (s *testScheduler) Shutdown() {} func (s *testScheduler) LastSeq() int64 { return s.lastSeq.Load() } // recv returns the next scheduled event. The timeout is generous because a // redial during the offline window backs off >=5s. func (s *testScheduler) recv(t *testing.T) *SubscribeExample_Event { t.Helper() select { case e := <-s.ch: return e case <-time.After(12 * time.Second): t.Fatal("timed out waiting for an event") return nil } } type testRedialer struct { sched *testScheduler } var _ Redialer = (*testRedialer)(nil) func (r *testRedialer) Process(ctx context.Context, cr *cbg.CborReader) error { var evt SubscribeExample_Event if err := evt.Deserialize(cr); err != nil { return err } return r.sched.AddWork(ctx, "", &evt) } func (r *testRedialer) UpdateParams(ctx context.Context, params map[string]any) bool { last := r.sched.LastSeq() if last == 0 { return false } params["cursor"] = last return true } const testEndpoint = "com.example.subscribeExample" // testServer live-tails an append-only event log over a websocket: each // connection streams events whose seq exceeds the requested cursor, including // ones added after it opened. The httptest listener stays bound the whole time // (so the port can't be taken over); Close/Start toggle offline by having the // handler answer 404, which fails the client's websocket handshake and drives // it into its redial/backoff loop. type testServer struct { URL string mu sync.Mutex events []SubscribeExample_Event conns map[*websocket.Conn]struct{} serving bool } func newTestServer(t *testing.T) *testServer { t.Helper() ts := &testServer{ conns: make(map[*websocket.Conn]struct{}), serving: true, } srv := httptest.NewServer(ts.handler()) ts.URL = srv.URL t.Cleanup(func() { ts.Close() // drop hijacked conns before srv.Close so it won't block srv.Close() }) return ts } func (ts *testServer) handler() http.Handler { up := websocket.Upgrader{} mux := http.NewServeMux() mux.HandleFunc("/xrpc/"+testEndpoint, func(w http.ResponseWriter, r *http.Request) { if !ts.isServing() { http.Error(w, "offline", http.StatusNotFound) return } var cursor int64 if c := r.URL.Query().Get("cursor"); c != "" { cursor, _ = strconv.ParseInt(c, 10, 64) } conn, err := up.Upgrade(w, r, nil) if err != nil { return } ts.mu.Lock() ts.conns[conn] = struct{}{} ts.mu.Unlock() defer func() { ts.mu.Lock() delete(ts.conns, conn) ts.mu.Unlock() conn.Close() }() sent := cursor for { ts.mu.Lock() var batch []SubscribeExample_Event for _, e := range ts.events { if e.seq() > sent { batch = append(batch, e) } } ts.mu.Unlock() for i := range batch { if err := func(conn *websocket.Conn, e *SubscribeExample_Event) error { var buf bytes.Buffer if err := e.Serialize(&buf); err != nil { return err } return conn.WriteMessage(websocket.BinaryMessage, buf.Bytes()) }(conn, &batch[i]); err != nil { return } sent = batch[i].seq() } time.Sleep(5 * time.Millisecond) } }) return mux } func (ts *testServer) isServing() bool { ts.mu.Lock() defer ts.mu.Unlock() return ts.serving } // Start brings the server back online. func (ts *testServer) Start() { ts.mu.Lock() ts.serving = true ts.mu.Unlock() } // Close takes the server offline: new handshakes get 404 and active websocket // connections are dropped, forcing the subscription to redial. func (ts *testServer) Close() { ts.mu.Lock() ts.serving = false conns := make([]*websocket.Conn, 0, len(ts.conns)) for c := range ts.conns { conns = append(conns, c) } ts.mu.Unlock() for _, c := range conns { c.Close() } } func (ts *testServer) AddEvent(e SubscribeExample_Event) { ts.mu.Lock() ts.events = append(ts.events, e) ts.mu.Unlock() } func waitForReturn(t *testing.T, done <-chan error) { t.Helper() select { case <-done: case <-time.After(5 * time.Second): t.Fatal("subscription did not return after context cancel") } } func TestLexSubscribe_ConsumesAndSchedules(t *testing.T) { srv := newTestServer(t) defer srv.Close() c := &Client{Client: indigoxrpc.Client{Host: srv.URL}} sched := newTestScheduler() process := func(ctx context.Context, cr *cbg.CborReader) error { var evt SubscribeExample_Event if err := evt.Deserialize(cr); err != nil { return err } return sched.AddWork(ctx, "", &evt) } ctx, cancel := context.WithCancel(context.Background()) done := make(chan error, 1) go func() { done <- c.LexSubscribe(ctx, testEndpoint, map[string]any{"cursor": int64(0)}, process) }() srv.AddEvent(SubscribeExample_Event{Foo: &SubscribeExample_Foo{Seq: 1, Foo: "foo-1"}}) srv.AddEvent(SubscribeExample_Event{Bar: &SubscribeExample_Bar{Seq: 2, Bar: "bar-2"}}) assert.Equal(t, &SubscribeExample_Event{Foo: &SubscribeExample_Foo{Seq: 1, Foo: "foo-1"}}, sched.recv(t)) assert.Equal(t, &SubscribeExample_Event{Bar: &SubscribeExample_Bar{Seq: 2, Bar: "bar-2"}}, sched.recv(t)) srv.AddEvent(SubscribeExample_Event{Foo: &SubscribeExample_Foo{Seq: 3, Foo: "foo-3"}}) assert.Equal(t, &SubscribeExample_Event{Foo: &SubscribeExample_Foo{Seq: 3, Foo: "foo-3"}}, sched.recv(t)) cancel() waitForReturn(t, done) } func TestLexSubscribeWithRedialer_ConsumesAndSchedules(t *testing.T) { srv := newTestServer(t) defer srv.Close() c := &Client{Client: indigoxrpc.Client{Host: srv.URL}} sched := newTestScheduler() redialer := &testRedialer{sched: sched} ctx, cancel := context.WithCancel(context.Background()) done := make(chan error, 1) go func() { done <- c.LexSubscribeWithRedialer(ctx, testEndpoint, map[string]any{"cursor": int64(0)}, redialer) }() srv.AddEvent(SubscribeExample_Event{Foo: &SubscribeExample_Foo{Seq: 1, Foo: "foo-1"}}) srv.AddEvent(SubscribeExample_Event{Bar: &SubscribeExample_Bar{Seq: 2, Bar: "bar-2"}}) assert.Equal(t, &SubscribeExample_Event{Foo: &SubscribeExample_Foo{Seq: 1, Foo: "foo-1"}}, sched.recv(t)) assert.Equal(t, &SubscribeExample_Event{Bar: &SubscribeExample_Bar{Seq: 2, Bar: "bar-2"}}, sched.recv(t)) srv.AddEvent(SubscribeExample_Event{Foo: &SubscribeExample_Foo{Seq: 3, Foo: "foo-3"}}) srv.AddEvent(SubscribeExample_Event{Bar: &SubscribeExample_Bar{Seq: 4, Bar: "bar-4"}}) assert.Equal(t, &SubscribeExample_Event{Foo: &SubscribeExample_Foo{Seq: 3, Foo: "foo-3"}}, sched.recv(t)) assert.Equal(t, &SubscribeExample_Event{Bar: &SubscribeExample_Bar{Seq: 4, Bar: "bar-4"}}, sched.recv(t)) cancel() waitForReturn(t, done) } func TestLexSubscribeWithRedialer_HandlesDowntime(t *testing.T) { srv := newTestServer(t) defer srv.Close() c := &Client{Client: indigoxrpc.Client{Host: srv.URL}} sched := newTestScheduler() redialer := &testRedialer{sched: sched} ctx, cancel := context.WithCancel(context.Background()) done := make(chan error, 1) go func() { done <- c.LexSubscribeWithRedialer(ctx, testEndpoint, map[string]any{"cursor": int64(0)}, redialer) }() srv.AddEvent(SubscribeExample_Event{Foo: &SubscribeExample_Foo{Seq: 1, Foo: "foo-1"}}) srv.AddEvent(SubscribeExample_Event{Bar: &SubscribeExample_Bar{Seq: 2, Bar: "bar-2"}}) assert.Equal(t, &SubscribeExample_Event{Foo: &SubscribeExample_Foo{Seq: 1, Foo: "foo-1"}}, sched.recv(t)) assert.Equal(t, &SubscribeExample_Event{Bar: &SubscribeExample_Bar{Seq: 2, Bar: "bar-2"}}, sched.recv(t)) // offline, add events, back online: the subscription redials and resumes srv.Close() srv.AddEvent(SubscribeExample_Event{Foo: &SubscribeExample_Foo{Seq: 3, Foo: "foo-3"}}) srv.AddEvent(SubscribeExample_Event{Bar: &SubscribeExample_Bar{Seq: 4, Bar: "bar-4"}}) srv.Start() assert.Equal(t, &SubscribeExample_Event{Foo: &SubscribeExample_Foo{Seq: 3, Foo: "foo-3"}}, sched.recv(t)) assert.Equal(t, &SubscribeExample_Event{Bar: &SubscribeExample_Bar{Seq: 4, Bar: "bar-4"}}, sched.recv(t)) cancel() waitForReturn(t, done) }