package httpserver import ( "context" "encoding/json" "net/http" "net/http/httptest" "strings" "sync" "testing" "time" "git.yuntuo.hk/k2safe/orbit/internal/jsonwire" core "git.yuntuo.hk/k2safe/orbit/model" "git.yuntuo.hk/k2safe/orbit/platform/auth" "git.yuntuo.hk/k2safe/orbit/platform/events" "github.com/coder/websocket" ) type slowEventStore struct { events.Store once sync.Once started chan struct{} finished chan struct{} } func (s *slowEventStore) AuthorizeEvents(context.Context, auth.Token) error { return nil } func (s *slowEventStore) BindConsumer(_ context.Context, _ auth.Token, id core.ID) (events.ConsumerBinding, error) { return events.ConsumerBinding{ApplicationID: core.ID(strings.Repeat("1", 32)), SubscriptionID: id, QueueID: core.ID(strings.Repeat("2", 32)), Enabled: true, ServerTime: time.Now().UTC()}, nil } func (s *slowEventStore) ClaimDelivery(ctx context.Context, _ auth.Token, _ core.ID, _ events.ClaimInput) (*events.Grant, error) { s.once.Do(func() { close(s.started) }) <-ctx.Done() select { case s.finished <- struct{}{}: default: } return nil, ctx.Err() } func TestEventHubSlowClaimKeepsControlLiveAndBoundsConnections(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() store := &slowEventStore{started: make(chan struct{}), finished: make(chan struct{}, 4)} hub := NewEventHub(ctx, events.New(store)) if cap(hub.slots) != 128 { t.Fatal("unexpected production connection bound") } hub.slots = make(chan struct{}, 2) hub.poll = 10 * time.Millisecond server := New("", nil) RegisterEventConsumers(server, hub) tls := httptest.NewTLSServer(server.Handler) defer tls.Close() defer hub.Shutdown(context.Background()) id := core.ID(strings.Repeat("3", 32)) credential := core.ID(strings.Repeat("4", 32)) token, _, err := auth.Issue(id, credential) if err != nil { t.Fatal(err) } headers := http.Header{} headers.Set("Authorization", "Bearer "+token) dial := func() (*websocket.Conn, *http.Response, error) { return websocket.Dial(ctx, strings.Replace(tls.URL, "https://", "wss://", 1)+"/v1/events/connect", &websocket.DialOptions{HTTPClient: tls.Client(), HTTPHeader: headers, Subprotocols: []string{EventSubprotocol}}) } connect := func() (*websocket.Conn, core.ID) { t.Helper() conn, _, err := dial() if err != nil { t.Fatal(err) } t.Cleanup(func() { conn.CloseNow() }) hello, _ := jsonwire.Marshal(eventHello{SubscriptionID: id, Capacity: 1}) raw, _ := jsonwire.Marshal(eventFrame{Version: "1", Type: "events.hello", MessageID: id, Body: hello}) if err = conn.Write(ctx, websocket.MessageText, raw); err != nil { t.Fatal(err) } _, reply, err := conn.Read(ctx) if err != nil { t.Fatal(err) } var frame eventFrame if json.Unmarshal(reply, &frame) != nil || frame.Type != "events.welcome" { t.Fatal("missing welcome") } return conn, frame.SessionID } first, session := connect() second, _ := connect() select { case <-store.started: case <-ctx.Done(): t.Fatal("claim was not dispatched") } blocked, resp, err := dial() if blocked != nil { blocked.CloseNow() } if err == nil || resp == nil || resp.StatusCode != 429 || resp.Header.Get("Retry-After") != "1" { t.Fatal("event connection bound not enforced") } ping, _ := jsonwire.Marshal(eventFrame{Version: "1", Type: "events.ping", MessageID: credential, SessionID: session, Body: json.RawMessage(`{}`)}) if err = first.Write(ctx, websocket.MessageText, ping); err != nil { t.Fatal(err) } bounded, stop := context.WithTimeout(ctx, 500*time.Millisecond) _, raw, err := first.Read(bounded) stop() if err != nil { t.Fatal("slow claim blocked control reply") } var reply eventFrame if json.Unmarshal(raw, &reply) != nil || reply.Type != "ack" || reply.ReplyTo != credential { t.Fatal("missing ping acknowledgement") } second.CloseNow() deadline := time.Now().Add(time.Second) for len(hub.slots) == 2 && time.Now().Before(deadline) { time.Sleep(time.Millisecond) } if len(hub.slots) != 1 { t.Fatal("disconnected consumer leaked connection slot") } third, _ := connect() defer third.CloseNow() bounded, stop = context.WithTimeout(ctx, time.Second) defer stop() if err = hub.Shutdown(bounded); err != nil { t.Fatal("shutdown leaked claim or connection", err) } select { case <-store.finished: default: t.Fatal("shutdown did not cancel in-flight claim") } if len(hub.slots) != 0 { t.Fatal("shutdown leaked connection slots") } }