diff --git a/go.mod b/go.mod index 27ceb3d4..56db3c45 100644 --- a/go.mod +++ b/go.mod @@ -11,6 +11,7 @@ require ( github.com/stretchr/testify v1.11.1 go.uber.org/zap v1.28.0 golang.org/x/oauth2 v0.36.0 + golang.org/x/sync v0.22.0 google.golang.org/api v0.264.0 google.golang.org/grpc v1.83.0 google.golang.org/protobuf v1.36.11 @@ -137,7 +138,6 @@ require ( golang.org/x/crypto v0.54.0 // indirect golang.org/x/exp v0.0.0-20260312153236-7ab1446f8b90 // indirect golang.org/x/net v0.57.0 // indirect - golang.org/x/sync v0.22.0 // indirect golang.org/x/sys v0.47.0 // indirect golang.org/x/term v0.45.0 // indirect golang.org/x/text v0.40.0 // indirect diff --git a/pkg/connector/event_feed_common.go b/pkg/connector/event_feed_common.go index f16dc7fd..ede73193 100644 --- a/pkg/connector/event_feed_common.go +++ b/pkg/connector/event_feed_common.go @@ -19,6 +19,7 @@ import ( "encoding/base64" "encoding/json" "fmt" + "time" v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" "github.com/conductorone/baton-sdk/pkg/pagination" @@ -29,11 +30,28 @@ import ( gwclient "github.com/conductorone/baton-google-workspace/pkg/client" ) -// usersPerEventFeedCall bounds how many users are processed per ListEvents invocation, so a -// single call issues at most this many Reports API filter-queries and returns quickly instead -// of blocking on the shared 250/min quota for an entire directory page (up to 500 users). +// usersPerEventFeedCall bounds how many users are considered per ListEvents invocation. This +// alone does not bound the number of Reports API calls issued (see maxLookupCallsPerEventFeedCall +// below), but keeps the directory-side bookkeeping (ListTokens, cursor size) proportional. const usersPerEventFeedCall = 25 +// maxLookupCallsPerEventFeedCall caps Reports API calls per ListEvents invocation. Bounds cost, +// not time (shared rate limiter, retries with backoff) — see maxEventFeedCallDuration for that. +const maxLookupCallsPerEventFeedCall = 60 + +// maxEventFeedCallDuration is a soft wall-clock budget per ListEvents invocation. It bounds the +// gap between users (see the loop in scanUsersForEvents) and is also handed to each user's own +// lookup as a deadline, so a single user's internal fan-out (e.g. usageEventFeed.lookupUser's +// per-app Reports calls) stops starting new work once it's spent, instead of only being checked +// again after that one user's lookup already returned. +const maxEventFeedCallDuration = 45 * time.Second + +// eventFeedDeadlineSafetyMargin is subtracted from ctx's deadline (when the caller set one) to +// derive the effective per-call deadline passed to lookup, so this connector stops issuing new +// Reports calls with enough margin to unwind and return a partial result before the RPC itself +// times out — rather than getting DeadlineExceeded mid-call with nothing to show for it. +const eventFeedDeadlineSafetyMargin = 5 * time.Second + type pendingUser struct { Email string `json:"email"` ID string `json:"id"` @@ -45,9 +63,13 @@ type pendingUser struct { // When both are empty/exhausted, the walk is complete and the cursor resets to nil so the // next call starts a fresh pass — there is no "since last poll" time window to track, since // each lookup always asks for the current latest login, not a delta. +// +// ResumeState is opaque, feed-defined progress for PendingUsers[0]'s own lookup (e.g. which +// authorized app to resume from), set when a call exhausts its budget mid-user. type userScanCursor struct { PendingUsers []pendingUser `json:"pending_users,omitempty"` DirectoryPageToken string `json:"directory_page_token,omitempty"` + ResumeState string `json:"resume_state,omitempty"` } func unmarshalUserScanCursor(pToken *pagination.StreamToken) (*userScanCursor, error) { @@ -84,8 +106,23 @@ func (c *userScanCursor) marshal() (string, error) { return base64.StdEncoding.EncodeToString(data), nil } -// userEventLookup fetches events for a single user via at most one Reports API call. -type userEventLookup func(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser) ([]*v2.Event, error) +// userEventLookup fetches events for one user, spending at most `budget` Reports API calls. +// resumeState picks up where a prior call for this user left off ("" = start fresh); if more +// than `budget` calls are needed, it returns a non-empty nextResumeState instead of finishing. +// consumed is how many calls this invocation issued (meaningful only when err == nil). deadline +// is a soft wall-clock cutoff: a lookup with internal fan-out (multiple Reports calls per user) +// should stop starting new calls once it's past deadline and return whatever it has, with +// nextResumeState reflecting the last unit of work it actually finished — never a zero-value +// time.Time in practice, but implementations that only ever issue one call per user (there is no +// fan-out to bound) are free to ignore it. +type userEventLookup func( + ctx context.Context, + client *gwclient.GoogleWorkspaceClient, + user pendingUser, + resumeState string, + budget int, + deadline time.Time, +) (events []*v2.Event, nextResumeState string, consumed int, err error) // scanUsersForEvents drives one bounded step of the rolling user-directory walk shared by all // three "last login" event feeds. @@ -142,37 +179,76 @@ func scanUsersForEvents( batch = batch[:usersPerEventFeedCall] } + // finish marshals the cursor and returns the call's result; shared by every exit path. + finish := func(events []*v2.Event, hasMore bool) ([]*v2.Event, *pagination.StreamState, error) { + if !hasMore { + cursor = &userScanCursor{} + } + cursorToken, err := cursor.marshal() + if err != nil { + return nil, nil, fmt.Errorf("failed to marshal cursor token in event feed: %w", err) + } + return events, &pagination.StreamState{Cursor: cursorToken, HasMore: hasMore}, nil + } + events := []*v2.Event{} - for _, u := range batch { - userEvents, err := lookup(ctx, client, u) + budget := maxLookupCallsPerEventFeedCall + resumeState := cursor.ResumeState + start := time.Now() + + // deadline is handed to each user's lookup so a single heavy user's internal fan-out (e.g. + // many Reports calls for one user's authorized apps) also respects the wall-clock budget, + // not just the gap between users below. If the caller's ctx carries its own deadline (the + // real RPC deadline), stop earlier than that with a safety margin instead of racing it. + deadline := start.Add(maxEventFeedCallDuration) + if ctxDeadline, ok := ctx.Deadline(); ok { + if safe := ctxDeadline.Add(-eventFeedDeadlineSafetyMargin); safe.Before(deadline) { + deadline = safe + } + } + + for i, u := range batch { + if budget <= 0 || !time.Now().Before(deadline) { + // Budget or time spent before starting u; resume here fresh next call. + cursor.PendingUsers = cursor.PendingUsers[i:] + cursor.ResumeState = "" + return finish(events, true) + } + + userEvents, nextResume, consumed, err := lookup(ctx, client, u, resumeState, budget, deadline) if err != nil { - // Don't remove `batch` from cursor.PendingUsers until every lookup in it has - // succeeded, so a single user's Reports API blip doesn't lose the remaining - // unprocessed users and restart the walk from the beginning on retry. + // The SDK drops this StreamState on error (connectorbuilder.ListEvents returns nil, + // err without ever reading it), so what we marshal here is moot either way; return + // cursor as it currently stands — which may already include a fresh directory-page + // advance from the len(cursor.PendingUsers) == 0 branch above — rather than trying to + // preserve or roll back per-user progress for no effect. cursorToken, marshalErr := cursor.marshal() if marshalErr != nil { return nil, nil, fmt.Errorf("failed to marshal cursor token in event feed: %w", marshalErr) } return nil, &pagination.StreamState{Cursor: cursorToken, HasMore: true}, err } + budget -= consumed + resumeState = "" // only the batch's first (possibly-resumed) user carries one in + for _, e := range userEvents { if earliestEvent != nil && e.GetOccurredAt() != nil && e.GetOccurredAt().AsTime().Before(earliestEvent.AsTime()) { continue } events = append(events, e) } - } - cursor.PendingUsers = cursor.PendingUsers[len(batch):] - hasMore := len(cursor.PendingUsers) > 0 || cursor.DirectoryPageToken != "" - if !hasMore { - cursor = &userScanCursor{} + if nextResume != "" { + // u isn't finished; keep it at the front and stop the batch here. + cursor.PendingUsers = cursor.PendingUsers[i:] + cursor.ResumeState = nextResume + return finish(events, true) + } } - cursorToken, err := cursor.marshal() - if err != nil { - return nil, nil, fmt.Errorf("failed to marshal cursor token in event feed: %w", err) - } + cursor.PendingUsers = cursor.PendingUsers[len(batch):] + cursor.ResumeState = "" - return events, &pagination.StreamState{Cursor: cursorToken, HasMore: hasMore}, nil + hasMore := len(cursor.PendingUsers) > 0 || cursor.DirectoryPageToken != "" + return finish(events, hasMore) } diff --git a/pkg/connector/event_feed_common_test.go b/pkg/connector/event_feed_common_test.go index ae7dc5f9..56d5695d 100644 --- a/pkg/connector/event_feed_common_test.go +++ b/pkg/connector/event_feed_common_test.go @@ -7,6 +7,8 @@ import ( "net/http" "net/http/httptest" "strconv" + "strings" + "sync" "testing" "time" @@ -128,9 +130,9 @@ func TestScanUsersForEvents_PaginatesAcrossMultipleCallsWithoutLoss(t *testing.T } _, state, err := scanUsersForEvents(context.Background(), client, "customer", "", nil, &pagination.StreamToken{Cursor: cursor}, - func(ctx context.Context, c *gwclient.GoogleWorkspaceClient, u pendingUser) ([]*v2.Event, error) { + func(ctx context.Context, c *gwclient.GoogleWorkspaceClient, u pendingUser, _ string, _ int, _ time.Time) ([]*v2.Event, string, int, error) { visited[u.Email+":lookup"]++ - return nil, nil + return nil, "", 1, nil }) if err != nil { t.Fatalf("scanUsersForEvents: %v", err) @@ -163,11 +165,11 @@ func TestScanUsersForEvents_FiltersEventsBeforeEarliestEvent(t *testing.T) { now := time.Now().UTC() floor := timestamppb.New(now.Add(-1 * time.Hour)) - oldEvent := &v2.Event{Id: "old", OccurredAt: timestamppb.New(now.Add(-2 * time.Hour))} // before floor + oldEvent := &v2.Event{Id: "old", OccurredAt: timestamppb.New(now.Add(-2 * time.Hour))} // before floor newEvent := &v2.Event{Id: "new", OccurredAt: timestamppb.New(now.Add(-30 * time.Minute))} // after floor - lookup := func(ctx context.Context, c *gwclient.GoogleWorkspaceClient, u pendingUser) ([]*v2.Event, error) { - return []*v2.Event{oldEvent, newEvent}, nil + lookup := func(ctx context.Context, c *gwclient.GoogleWorkspaceClient, u pendingUser, _ string, _ int, _ time.Time) ([]*v2.Event, string, int, error) { + return []*v2.Event{oldEvent, newEvent}, "", 1, nil } events, _, err := scanUsersForEvents(context.Background(), client, "customer", "", floor, &pagination.StreamToken{}, lookup) @@ -268,6 +270,77 @@ func TestUsageEventFeed_PicksLatestPerAppAndFiltersPrivateApps(t *testing.T) { } } +// TestUsageEventFeed_LookupUserStopsAtDeadlineBetweenChunks verifies that lookupUser's per-app +// fan-out (usage_event_feed.go) respects the deadline passed in from scanUsersForEvents: once +// past deadline, it stops dispatching new app lookups rather than draining the whole remaining +// budget in one call, but it always dispatches at least the first app (the deadline is only +// checked *before dispatching the second and later apps*) so a single call still makes forward +// progress even if the deadline was already past when the call started. +func TestUsageEventFeed_LookupUserStopsAtDeadlineBetweenChunks(t *testing.T) { + const userEmail = "heavy@example.com" + const numApps = maxConcurrentAppLookups*2 + 3 // more than one chunk's worth + + tokens := make([]*directoryAdmin.Token, 0, numApps) + for i := 0; i < numApps; i++ { + tokens = append(tokens, &directoryAdmin.Token{ClientId: fmt.Sprintf("client-%d", i), DisplayText: fmt.Sprintf("App %d", i)}) + } + + var mu sync.Mutex + callCounts := map[string]int{} + + mux := http.NewServeMux() + mux.HandleFunc("/admin/directory/v1/users/", func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(&directoryAdmin.Tokens{Items: tokens}) + }) + mux.HandleFunc("/admin/reports/v1/activity/users/", func(w http.ResponseWriter, r *http.Request) { + clientID := strings.TrimPrefix(r.URL.Query().Get("filters"), "client_id==") + mu.Lock() + callCounts[clientID]++ + mu.Unlock() + _ = json.NewEncoder(w).Encode(&reportsAdmin.Activities{ + Items: []*reportsAdmin.Activity{ + activityItem(1, time.Minute, userEmail, "profile-heavy", + &reportsAdmin.ActivityEventsParameters{Name: "client_id", Value: clientID}, + &reportsAdmin.ActivityEventsParameters{Name: "app_name", Value: clientID}, + ), + }, + }) + }) + server := httptest.NewServer(mux) + defer server.Close() + + dir := newTestDirectoryService(t, server.URL, server.Client()) + rep := newReportsServiceForTest(t, server.URL, server.Client()) + feed := newUsageEventFeed(&gwclient.GoogleWorkspaceClient{UserService: dir, UserSecurityService: dir, ReportService: rep}, "customer", "") + user := pendingUser{Email: userEmail, ID: "heavy-user"} + + // A deadline already in the past: only the check before dispatching the second (and later) + // app sees it, so the first app is still dispatched and awaited before lookupUser bails out. + events, nextResume, consumed, err := feed.lookupUser(context.Background(), feed.c, user, "", numApps, time.Now().Add(-time.Hour)) + if err != nil { + t.Fatalf("lookupUser: %v", err) + } + if consumed != 1 { + t.Fatalf("expected exactly one app to be consumed before bailing out on the past deadline, got %d", consumed) + } + if len(events) != 1 { + t.Fatalf("expected exactly one event (for the one dispatched app), got %d", len(events)) + } + wantNextResume := "client-0" + if nextResume != wantNextResume { + t.Fatalf("expected nextResume to be the one app dispatched (%s), got %q", wantNextResume, nextResume) + } + if callCounts["client-0"] != 1 { + t.Fatalf("expected app client-0 to be queried exactly once, got %d", callCounts["client-0"]) + } + for i := 1; i < numApps; i++ { + clientID := fmt.Sprintf("client-%d", i) + if callCounts[clientID] != 0 { + t.Fatalf("expected app %s not to be queried once the deadline had passed, got %d", clientID, callCounts[clientID]) + } + } +} + // TestUsageEventFeed_DedupesRepeatedClientIDsInTokens verifies that when Tokens.list returns // multiple Token entries for the same client_id (e.g. separate grants for different scope // sets), the feed issues exactly one Reports API lookup for that app and emits exactly one @@ -325,3 +398,382 @@ func TestUsageEventFeed_DedupesRepeatedClientIDsInTokens(t *testing.T) { t.Fatalf("expected exactly 1 event for the deduped app, got %d", len(events)) } } + +// TestScanUsersForEvents_ResumesWithinUserWhenBudgetExhausted verifies that a user needing more +// Reports calls than the budget allows pauses mid-user, resumes via cursor state across calls, +// and has every unit of work visited exactly once. Uses a synthetic lookup to isolate +// scanUsersForEvents' generic budget/resume mechanics from any one feed's fan-out logic. +func TestScanUsersForEvents_ResumesWithinUserWhenBudgetExhausted(t *testing.T) { + const totalUnits = maxLookupCallsPerEventFeedCall*2 + 7 // spans exactly 3 resumed calls + user := pendingUser{Email: "heavy@example.com", ID: "heavy-user"} + unitLookupCounts := map[int]int{} + + lookup := func(ctx context.Context, c *gwclient.GoogleWorkspaceClient, u pendingUser, resumeState string, budget int, _ time.Time) ([]*v2.Event, string, int, error) { + startIdx := 0 + if resumeState != "" { + idx, err := strconv.Atoi(resumeState) + if err != nil { + t.Fatalf("bad resume state %q: %v", resumeState, err) + } + startIdx = idx + } + consumed := totalUnits - startIdx + if consumed > budget { + consumed = budget + } + events := make([]*v2.Event, 0, consumed) + for i := startIdx; i < startIdx+consumed; i++ { + unitLookupCounts[i]++ + events = append(events, &v2.Event{Id: strconv.Itoa(i)}) + } + nextResume := "" + if startIdx+consumed < totalUnits { + nextResume = strconv.Itoa(startIdx + consumed) + } + return events, nextResume, consumed, nil + } + + server := newDirectoryUsersOnlyServer(t, []*directoryAdmin.User{{Id: user.ID, PrimaryEmail: user.Email}}, 10, + func(string) *reportsAdmin.Activities { return &reportsAdmin.Activities{} }) + defer server.Close() + dir := newTestDirectoryService(t, server.URL, server.Client()) + client := &gwclient.GoogleWorkspaceClient{UserService: dir} + + var cursor string + var allEvents []*v2.Event + calls := 0 + for { + calls++ + if calls > 10 { + t.Fatalf("expected this to resolve in a small, bounded number of calls, got stuck after %d", calls) + } + events, state, err := scanUsersForEvents(context.Background(), client, "customer", "", nil, &pagination.StreamToken{Cursor: cursor}, lookup) + if err != nil { + t.Fatalf("scanUsersForEvents: %v", err) + } + allEvents = append(allEvents, events...) + cursor = state.Cursor + if !state.HasMore { + break + } + } + + if calls != 3 { + t.Fatalf("expected exactly 3 calls (budget=%d, units=%d), got %d", maxLookupCallsPerEventFeedCall, totalUnits, calls) + } + if len(allEvents) != totalUnits { + t.Fatalf("expected exactly %d events (one per unit), got %d", totalUnits, len(allEvents)) + } + for i := 0; i < totalUnits; i++ { + if unitLookupCounts[i] != 1 { + t.Fatalf("expected unit %d to be looked up exactly once, got %d", i, unitLookupCounts[i]) + } + } +} + +// withUnlimitedReportsRateLimiter swaps sharedReportsRateLimiter for a large-capacity one for the +// test's duration, so tests issuing many real Reports calls aren't order-dependent on other tests. +func withUnlimitedReportsRateLimiter(t *testing.T) { + t.Helper() + original := sharedReportsRateLimiter + sharedReportsRateLimiter = newReportsRateLimiter(1_000_000) + t.Cleanup(func() { sharedReportsRateLimiter = original }) +} + +// TestUsageEventFeed_ResumesAcrossManyAuthorizedApps is the end-to-end version of the test above, +// exercising usage_event_feed's real lookupUser against a user with more authorized apps than +// maxLookupCallsPerEventFeedCall allows per call. +func TestUsageEventFeed_ResumesAcrossManyAuthorizedApps(t *testing.T) { + withUnlimitedReportsRateLimiter(t) + + const userEmail = "heavy@example.com" + const numApps = maxLookupCallsPerEventFeedCall + 15 // forces exactly 2 resumed calls + + tokens := make([]*directoryAdmin.Token, 0, numApps) + for i := 0; i < numApps; i++ { + tokens = append(tokens, &directoryAdmin.Token{ClientId: fmt.Sprintf("client-%d", i), DisplayText: fmt.Sprintf("App %d", i)}) + } + + var mu sync.Mutex + callCounts := map[string]int{} + + mux := http.NewServeMux() + mux.HandleFunc("/admin/directory/v1/users", func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("pageToken") != "" { + _ = json.NewEncoder(w).Encode(&directoryAdmin.Users{}) + return + } + _ = json.NewEncoder(w).Encode(&directoryAdmin.Users{ + Users: []*directoryAdmin.User{{Id: "profile-heavy", PrimaryEmail: userEmail}}, + }) + }) + mux.HandleFunc("/admin/directory/v1/users/", func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(&directoryAdmin.Tokens{Items: tokens}) + }) + mux.HandleFunc("/admin/reports/v1/activity/users/", func(w http.ResponseWriter, r *http.Request) { + clientID := strings.TrimPrefix(r.URL.Query().Get("filters"), "client_id==") + mu.Lock() + callCounts[clientID]++ + mu.Unlock() + _ = json.NewEncoder(w).Encode(&reportsAdmin.Activities{ + Items: []*reportsAdmin.Activity{ + activityItem(1, time.Minute, userEmail, "profile-heavy", + &reportsAdmin.ActivityEventsParameters{Name: "client_id", Value: clientID}, + &reportsAdmin.ActivityEventsParameters{Name: "app_name", Value: clientID}, + ), + }, + }) + }) + server := httptest.NewServer(mux) + defer server.Close() + + dir := newTestDirectoryService(t, server.URL, server.Client()) + rep := newReportsServiceForTest(t, server.URL, server.Client()) + feed := newUsageEventFeed(&gwclient.GoogleWorkspaceClient{UserService: dir, UserSecurityService: dir, ReportService: rep}, "customer", "") + + var cursor string + var allEvents []*v2.Event + calls := 0 + for { + calls++ + if calls > 10 { + t.Fatalf("expected this to resolve in a small, bounded number of calls, got stuck after %d", calls) + } + events, state, _, err := feed.ListEvents(context.Background(), nil, &pagination.StreamToken{Cursor: cursor}) + if err != nil { + t.Fatalf("ListEvents: %v", err) + } + allEvents = append(allEvents, events...) + cursor = state.Cursor + if !state.HasMore { + break + } + } + + if calls != 2 { + t.Fatalf("expected exactly 2 ListEvents calls (budget=%d, apps=%d), got %d", maxLookupCallsPerEventFeedCall, numApps, calls) + } + if len(allEvents) != numApps { + t.Fatalf("expected %d events (one per app), got %d", numApps, len(allEvents)) + } + if len(callCounts) != numApps { + t.Fatalf("expected %d distinct apps queried, got %d", numApps, len(callCounts)) + } + for clientID, n := range callCounts { + if n != 1 { + t.Fatalf("expected exactly 1 Reports API call for %s, got %d", clientID, n) + } + } +} + +// TestUsageEventFeed_ResumeSurvivesAppListShift verifies that resuming a user's authorized-app +// lookup by client_id (not by positional index) tolerates the app list shifting between the +// first and second ListEvents call. A revoked app that was already processed disappears from +// Tokens.list, shifting every later app's position down by one; a positional-index resume would +// then skip whatever app lands at the old index in the shifted list, silently losing an event. +// Resuming by client_id instead re-locates the last-processed app in the fresh list (wherever it +// now sits) and continues right after it, so every not-yet-processed app is still visited exactly +// once. +func TestUsageEventFeed_ResumeSurvivesAppListShift(t *testing.T) { + withUnlimitedReportsRateLimiter(t) + + const userEmail = "heavy@example.com" + const numApps = maxLookupCallsPerEventFeedCall + 15 // forces exactly 2 resumed calls + const revokedAfterFirstCall = "client-5" // processed in call 1, then revoked + + var mu sync.Mutex + callCounts := map[string]int{} + tokensCalls := 0 + + mux := http.NewServeMux() + mux.HandleFunc("/admin/directory/v1/users", func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("pageToken") != "" { + _ = json.NewEncoder(w).Encode(&directoryAdmin.Users{}) + return + } + _ = json.NewEncoder(w).Encode(&directoryAdmin.Users{ + Users: []*directoryAdmin.User{{Id: "profile-heavy", PrimaryEmail: userEmail}}, + }) + }) + mux.HandleFunc("/admin/directory/v1/users/", func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + tokensCalls++ + shifted := tokensCalls > 1 + mu.Unlock() + + tokens := make([]*directoryAdmin.Token, 0, numApps) + for i := 0; i < numApps; i++ { + clientID := fmt.Sprintf("client-%d", i) + if shifted && clientID == revokedAfterFirstCall { + // Revoked between calls: every later app's index shifts down by one. + continue + } + tokens = append(tokens, &directoryAdmin.Token{ClientId: clientID, DisplayText: clientID}) + } + _ = json.NewEncoder(w).Encode(&directoryAdmin.Tokens{Items: tokens}) + }) + mux.HandleFunc("/admin/reports/v1/activity/users/", func(w http.ResponseWriter, r *http.Request) { + clientID := strings.TrimPrefix(r.URL.Query().Get("filters"), "client_id==") + mu.Lock() + callCounts[clientID]++ + mu.Unlock() + _ = json.NewEncoder(w).Encode(&reportsAdmin.Activities{ + Items: []*reportsAdmin.Activity{ + activityItem(1, time.Minute, userEmail, "profile-heavy", + &reportsAdmin.ActivityEventsParameters{Name: "client_id", Value: clientID}, + &reportsAdmin.ActivityEventsParameters{Name: "app_name", Value: clientID}, + ), + }, + }) + }) + server := httptest.NewServer(mux) + defer server.Close() + + dir := newTestDirectoryService(t, server.URL, server.Client()) + rep := newReportsServiceForTest(t, server.URL, server.Client()) + feed := newUsageEventFeed(&gwclient.GoogleWorkspaceClient{UserService: dir, UserSecurityService: dir, ReportService: rep}, "customer", "") + + var cursor string + calls := 0 + for { + calls++ + if calls > 10 { + t.Fatalf("expected this to resolve in a small, bounded number of calls, got stuck after %d", calls) + } + _, state, _, err := feed.ListEvents(context.Background(), nil, &pagination.StreamToken{Cursor: cursor}) + if err != nil { + t.Fatalf("ListEvents: %v", err) + } + cursor = state.Cursor + if !state.HasMore { + break + } + } + + if calls != 2 { + t.Fatalf("expected exactly 2 ListEvents calls, got %d", calls) + } + // Every app must be queried exactly once, including the ones after the revoked app whose + // position shifted. A positional-index resume would skip one of them here. + if len(callCounts) != numApps { + t.Fatalf("expected %d distinct apps queried, got %d: %v", numApps, len(callCounts), callCounts) + } + for clientID, n := range callCounts { + if n != 1 { + t.Fatalf("expected exactly 1 Reports API call for %s, got %d", clientID, n) + } + } +} + +// TestUsageEventFeed_ResumeRestartsWhenAnchorAppRevoked locks in lookupUser's "anchor not +// found" fallback (usage_event_feed.go): if the exact app the resume cursor anchors on — the +// last one processed in a prior call — is itself revoked before the next resumed call, the +// fresh app list no longer contains it, so lookupUser can't re-locate a resume position and +// instead restarts this user's app scan from the beginning. This is intentionally conservative +// (it never skips an app), but it isn't free: every app already processed before the anchor gets +// looked up — and its event re-emitted — a second time. TestUsageEventFeed_ResumeSurvivesAppListShift +// revokes an app *before* the anchor, which leaves the anchor itself findable and never exercises +// this restart path. +func TestUsageEventFeed_ResumeRestartsWhenAnchorAppRevoked(t *testing.T) { + withUnlimitedReportsRateLimiter(t) + + const userEmail = "heavy@example.com" + const numApps = maxLookupCallsPerEventFeedCall + 15 // forces exactly 2 resumed calls + anchorClientID := fmt.Sprintf("client-%d", maxLookupCallsPerEventFeedCall-1) + + var mu sync.Mutex + callCounts := map[string]int{} + tokensCalls := 0 + + mux := http.NewServeMux() + mux.HandleFunc("/admin/directory/v1/users", func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("pageToken") != "" { + _ = json.NewEncoder(w).Encode(&directoryAdmin.Users{}) + return + } + _ = json.NewEncoder(w).Encode(&directoryAdmin.Users{ + Users: []*directoryAdmin.User{{Id: "profile-heavy", PrimaryEmail: userEmail}}, + }) + }) + mux.HandleFunc("/admin/directory/v1/users/", func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + tokensCalls++ + anchorRevoked := tokensCalls > 1 + mu.Unlock() + + tokens := make([]*directoryAdmin.Token, 0, numApps) + for i := 0; i < numApps; i++ { + clientID := fmt.Sprintf("client-%d", i) + if anchorRevoked && clientID == anchorClientID { + // The exact app the resume cursor anchors on is revoked before the next call. + continue + } + tokens = append(tokens, &directoryAdmin.Token{ClientId: clientID, DisplayText: clientID}) + } + _ = json.NewEncoder(w).Encode(&directoryAdmin.Tokens{Items: tokens}) + }) + mux.HandleFunc("/admin/reports/v1/activity/users/", func(w http.ResponseWriter, r *http.Request) { + clientID := strings.TrimPrefix(r.URL.Query().Get("filters"), "client_id==") + mu.Lock() + callCounts[clientID]++ + mu.Unlock() + _ = json.NewEncoder(w).Encode(&reportsAdmin.Activities{ + Items: []*reportsAdmin.Activity{ + activityItem(1, time.Minute, userEmail, "profile-heavy", + &reportsAdmin.ActivityEventsParameters{Name: "client_id", Value: clientID}, + &reportsAdmin.ActivityEventsParameters{Name: "app_name", Value: clientID}, + ), + }, + }) + }) + server := httptest.NewServer(mux) + defer server.Close() + + dir := newTestDirectoryService(t, server.URL, server.Client()) + rep := newReportsServiceForTest(t, server.URL, server.Client()) + feed := newUsageEventFeed(&gwclient.GoogleWorkspaceClient{UserService: dir, UserSecurityService: dir, ReportService: rep}, "customer", "") + + var cursor string + calls := 0 + for { + calls++ + if calls > 10 { + t.Fatalf("expected this to resolve in a small, bounded number of calls, got stuck after %d", calls) + } + _, state, _, err := feed.ListEvents(context.Background(), nil, &pagination.StreamToken{Cursor: cursor}) + if err != nil { + t.Fatalf("ListEvents: %v", err) + } + cursor = state.Cursor + if !state.HasMore { + break + } + } + + if calls != 3 { + t.Fatalf("expected exactly 3 ListEvents calls (call 1, the restart, and the tail), got %d", calls) + } + + // The revoked anchor app was queried exactly once, in call 1, before it disappeared. + if callCounts[anchorClientID] != 1 { + t.Fatalf("expected the revoked anchor app %s to be queried exactly once, got %d", anchorClientID, callCounts[anchorClientID]) + } + + // Every app processed before the anchor in call 1 gets replayed once the restart-from-0 + // fallback kicks in — the bounded, non-skipping cost of the conservative fallback. + for i := 0; i < maxLookupCallsPerEventFeedCall-1; i++ { + clientID := fmt.Sprintf("client-%d", i) + if callCounts[clientID] != 2 { + t.Fatalf("expected app %s (processed before the revoked anchor) to be replayed once after the restart, got %d calls", clientID, callCounts[clientID]) + } + } + + // No app is ever skipped: every app still present after the revocation is queried at least + // once. + for i := maxLookupCallsPerEventFeedCall; i < numApps; i++ { + clientID := fmt.Sprintf("client-%d", i) + if callCounts[clientID] < 1 { + t.Fatalf("expected app %s to be queried at least once, got %d", clientID, callCounts[clientID]) + } + } +} diff --git a/pkg/connector/google_login_event_feed.go b/pkg/connector/google_login_event_feed.go index d352abcc..76b13780 100644 --- a/pkg/connector/google_login_event_feed.go +++ b/pkg/connector/google_login_event_feed.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "strconv" + "time" v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" "github.com/conductorone/baton-sdk/pkg/annotations" @@ -47,10 +48,12 @@ func (f *googleLoginEventFeed) EventFeedMetadata(_ context.Context) *v2.EventFee } } -func (f *googleLoginEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser) ([]*v2.Event, error) { +// lookupUser issues exactly one Reports API call per user, so resumeState/budget/deadline are +// unused: there's no internal fan-out here to bound. +func (f *googleLoginEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, _ string, _ int, _ time.Time) ([]*v2.Event, string, int, error) { r, err := listActivitiesRateLimited(ctx, client, user.Email, reportsAppLogin, "login_success", "", "", googleLoginLookupMaxResults) if err != nil { - return nil, fmt.Errorf("google-workspace-connector: failed to list google login activities for %s: %w", user.Email, err) + return nil, "", 0, fmt.Errorf("google-workspace-connector: failed to list google login activities for %s: %w", user.Email, err) } var best *bestActivity @@ -67,14 +70,14 @@ func (f *googleLoginEventFeed) lookupUser(ctx context.Context, client *gwclient. } } if best == nil { - return nil, nil + return nil, "", 1, nil } userTrait, err := resource.NewUserTrait( resource.WithEmail(best.activity.Actor.Email, true), ) if err != nil { - return nil, fmt.Errorf("google-workspace-connector: failed to create user trait in google login event feed: %w", err) + return nil, "", 0, fmt.Errorf("google-workspace-connector: failed to create user trait in google login event feed: %w", err) } return []*v2.Event{{ @@ -100,7 +103,7 @@ func (f *googleLoginEventFeed) lookupUser(ctx context.Context, client *gwclient. }, }, }, - }}, nil + }}, "", 1, nil } func (f *googleLoginEventFeed) ListEvents( diff --git a/pkg/connector/saml_event_feed.go b/pkg/connector/saml_event_feed.go index ebae6761..cc13f180 100644 --- a/pkg/connector/saml_event_feed.go +++ b/pkg/connector/saml_event_feed.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "strconv" + "time" v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2" "github.com/conductorone/baton-sdk/pkg/annotations" @@ -48,11 +49,9 @@ type samlAppActivity struct { appName string } -// lookupUser tracks SAML app usage via Google's "saml" audit log. -// -// Unlike OAuth apps (see usage_event_feed.go), SAML "login_success" fires on every SSO -// authentication, so last login timestamps are accurate. SAML apps are identified by app name -// (no numeric client_id). +// lookupUser tracks SAML app usage via Google's "saml" audit log. Unlike OAuth apps (see +// usage_event_feed.go), one query covers every SAML app for a user, so it always completes in a +// single Reports API call and never needs resume state. func (f *samlEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, samlProfileMap map[string]string, user pendingUser) ([]*v2.Event, error) { r, err := listActivitiesRateLimited(ctx, client, user.Email, reportsAppSAML, "login_success", "", "", samlAppLookupMaxResults) if err != nil { @@ -134,8 +133,12 @@ func (f *samlEventFeed) ListEvents(ctx context.Context, earliestEvent *timestamp } events, streamState, err := scanUsersForEvents(ctx, f.client, f.customerID, f.domain, earliestEvent, pToken, - func(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser) ([]*v2.Event, error) { - return f.lookupUser(ctx, client, samlProfileMap, user) + func(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser, _ string, _ int, _ time.Time) ([]*v2.Event, string, int, error) { + events, err := f.lookupUser(ctx, client, samlProfileMap, user) + if err != nil { + return nil, "", 0, err + } + return events, "", 1, nil }) if err != nil { return nil, streamState, nil, err diff --git a/pkg/connector/usage_event_feed.go b/pkg/connector/usage_event_feed.go index 674c5a35..44b11219 100644 --- a/pkg/connector/usage_event_feed.go +++ b/pkg/connector/usage_event_feed.go @@ -13,6 +13,8 @@ import ( "github.com/conductorone/baton-sdk/pkg/annotations" "github.com/conductorone/baton-sdk/pkg/pagination" "github.com/conductorone/baton-sdk/pkg/types/resource" + "golang.org/x/sync/errgroup" + directoryAdmin "google.golang.org/api/admin/directory/v1" reportsAdmin "google.golang.org/api/admin/reports/v1" "google.golang.org/api/googleapi" "google.golang.org/protobuf/types/known/timestamppb" @@ -29,6 +31,11 @@ var privateAppIDRegex = regexp.MustCompile("[0-9]{21}") // client-side regardless of Google's actual ordering. const oauthAppLookupMaxResults = 5 +// maxConcurrentAppLookups bounds concurrent per-app Reports lookups for one user, and also sizes +// the fan-out chunks in lookupUser (see below) between which the deadline budget is checked. The +// shared reportsRateLimiter still caps overall quota use; this just overlaps network latency. +const maxConcurrentAppLookups = 8 + type usageEventFeed struct { c *gwclient.GoogleWorkspaceClient customerID string @@ -79,29 +86,17 @@ type oauthAppActivity struct { appName string } -// lookupUser enumerates the OAuth apps this user has authorized via Directory API Tokens.list -// (Reports API has no endpoint for "which apps has this user used" — Tokens.list is the cheap, -// Directory-quota source of truth for that), then issues one Reports API lookup per app, -// scoped with filters=client_id==. This trades "1 Reports call per user" (old per-app-type -// window, prone to one app's activity crowding another out of a shared count-bounded window) for -// "1 Directory call + N Reports calls per user" (N = apps that user has authorized) — each app's -// freshness is queried independently, so it can never be crowded out by another app's activity. -func (f *usageEventFeed) lookupUser(ctx context.Context, client *gwclient.GoogleWorkspaceClient, user pendingUser) ([]*v2.Event, error) { - tokenResp, err := client.ListTokens(ctx, user.ID) - if err != nil { - var gerr *googleapi.Error - if errors.As(err, &gerr) && gerr.Code == http.StatusNotFound { - // Benign: the user was deleted between the directory listing and this lookup. - return nil, nil - } - return nil, fmt.Errorf("google-workspace: failed to list oauth tokens for %s: %w", user.Email, err) - } - // Tokens.list can return multiple Token entries for the same client_id — e.g. a user - // granting a different scope set to the same app at different times. Dedupe here so each - // distinct app only ever triggers one Reports API lookup, never repeated ones. - seenClientIDs := make(map[string]struct{}, len(tokenResp.Items)) +// oauthApp is one distinct app a user has authorized, as discovered via Directory Tokens.list. +type oauthApp struct { + ClientID string + DisplayText string +} - events := make([]*v2.Event, 0, len(tokenResp.Items)) +// distinctAuthorizedApps enumerates this user's authorized OAuth apps via Directory Tokens.list, +// deduping repeated client_ids and filtering out private apps before spending any Reports budget. +func distinctAuthorizedApps(tokenResp *directoryAdmin.Tokens) []oauthApp { + seenClientIDs := make(map[string]struct{}, len(tokenResp.Items)) + apps := make([]oauthApp, 0, len(tokenResp.Items)) for _, t := range tokenResp.Items { if t.ClientId == "" || t.DisplayText == "" { continue @@ -114,17 +109,104 @@ func (f *usageEventFeed) lookupUser(ctx context.Context, client *gwclient.Google continue } seenClientIDs[t.ClientId] = struct{}{} + apps = append(apps, oauthApp{ClientID: t.ClientId, DisplayText: t.DisplayText}) + } + return apps +} - event, err := f.lookupAppLogin(ctx, client, user, t.ClientId, t.DisplayText) - if err != nil { - return nil, err +// lookupUser issues one Reports API lookup per authorized OAuth app, scoped by client_id. A user +// can have more apps than the per-call budget allows: resumeState is the client_id of the last +// app processed on a prior call for this user, and if apps remain after spending budget, +// nextResumeState carries the new last-processed client_id forward. A client_id (rather than a +// positional index) is used because Tokens.list is re-fetched fresh on every call and its +// ordering is not guaranteed stable across calls — a positional index could silently skip or +// re-visit apps if the set of authorized apps changes between calls. If resumeState's client_id +// is no longer present (e.g. the app was deauthorized between calls), lookupUser conservatively +// restarts from the beginning rather than guessing a position, which can revisit already-seen +// apps (harmless — the lookup is idempotent) but never skips one. +// +// Per-app lookups run concurrently, up to maxConcurrentAppLookups in flight at once; the shared +// rate limiter still caps quota use, so this only overlaps network latency and retry/backoff +// time. deadline is checked before dispatching each app (never before the first one, so a call +// always makes some progress) — once past it, lookupUser stops dispatching new lookups but still +// waits for the ones already in flight, and returns what it has with nextResumeState set to the +// last app actually dispatched. Without this, a single user with a full budget's worth of apps +// could keep dispatching lookups — each with its own rate-limiter wait and up to +// reportsMaxRetries backoff — well past the caller's soft wall-clock budget. +func (f *usageEventFeed) lookupUser( + ctx context.Context, + client *gwclient.GoogleWorkspaceClient, + user pendingUser, + resumeState string, + budget int, + deadline time.Time, +) ([]*v2.Event, string, int, error) { + tokenResp, err := client.ListTokens(ctx, user.ID) + if err != nil { + var gerr *googleapi.Error + if errors.As(err, &gerr) && gerr.Code == http.StatusNotFound { + // Benign: the user was deleted between the directory listing and this lookup. + return nil, "", 0, nil + } + return nil, "", 0, fmt.Errorf("google-workspace: failed to list oauth tokens for %s: %w", user.Email, err) + } + apps := distinctAuthorizedApps(tokenResp) + + startIdx := 0 + if resumeState != "" { + for i, app := range apps { + if app.ClientID == resumeState { + startIdx = i + 1 + break + } + } + } + + remaining := apps[startIdx:] + if len(remaining) == 0 { + return nil, "", 0, nil + } + toProcess := remaining + if len(toProcess) > budget { + toProcess = toProcess[:budget] + } + + results := make([]*v2.Event, len(toProcess)) + consumed := 0 + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(maxConcurrentAppLookups) + for i, app := range toProcess { + if i > 0 && !deadline.IsZero() && time.Now().After(deadline) { + // Past budget: stop dispatching new lookups. Ones already in flight are awaited + // below; nextResume picks up right after the last app actually dispatched. + break } - if event != nil { - events = append(events, event) + consumed = i + 1 + g.Go(func() error { + event, err := f.lookupAppLogin(gctx, client, user, app.ClientID, app.DisplayText) + if err != nil { + return err + } + results[i] = event + return nil + }) + } + if err := g.Wait(); err != nil { + return nil, "", 0, err + } + + events := make([]*v2.Event, 0, consumed) + for _, e := range results[:consumed] { + if e != nil { + events = append(events, e) } } - return events, nil + nextResume := "" + if consumed > 0 && startIdx+consumed < len(apps) { + nextResume = toProcess[consumed-1].ClientID + } + return events, nextResume, consumed, nil } // lookupAppLogin fetches this user's most recent "authorize" activity for one specific OAuth diff --git a/vendor/golang.org/x/sync/errgroup/errgroup.go b/vendor/golang.org/x/sync/errgroup/errgroup.go new file mode 100644 index 00000000..c261a8eb --- /dev/null +++ b/vendor/golang.org/x/sync/errgroup/errgroup.go @@ -0,0 +1,151 @@ +// Copyright 2016 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// Package errgroup provides synchronization, error propagation, and Context +// cancellation for groups of goroutines working on subtasks of a common task. +// +// [errgroup.Group] is related to [sync.WaitGroup] but adds handling of tasks +// returning errors. +package errgroup + +import ( + "context" + "fmt" + "sync" +) + +type token struct{} + +// A Group is a collection of goroutines working on subtasks that are part of +// the same overall task. A Group should not be reused for different tasks. +// +// A zero Group is valid, has no limit on the number of active goroutines, +// and does not cancel on error. +type Group struct { + cancel func(error) + + wg sync.WaitGroup + + sem chan token + + errOnce sync.Once + err error +} + +func (g *Group) done() { + if g.sem != nil { + <-g.sem + } + g.wg.Done() +} + +// WithContext returns a new Group and an associated Context derived from ctx. +// +// The derived Context is canceled the first time a function passed to Go +// returns a non-nil error or the first time Wait returns, whichever occurs +// first. +func WithContext(ctx context.Context) (*Group, context.Context) { + ctx, cancel := context.WithCancelCause(ctx) + return &Group{cancel: cancel}, ctx +} + +// Wait blocks until all function calls from the Go method have returned, then +// returns the first non-nil error (if any) from them. +func (g *Group) Wait() error { + g.wg.Wait() + if g.cancel != nil { + g.cancel(g.err) + } + return g.err +} + +// Go calls the given function in a new goroutine. +// +// The first call to Go must happen before a Wait. +// It blocks until the new goroutine can be added without the number of +// goroutines in the group exceeding the configured limit. +// +// The first goroutine in the group that returns a non-nil error will +// cancel the associated Context, if any. The error will be returned +// by Wait. +func (g *Group) Go(f func() error) { + if g.sem != nil { + g.sem <- token{} + } + + g.wg.Add(1) + go func() { + defer g.done() + + // It is tempting to propagate panics from f() + // up to the goroutine that calls Wait, but + // it creates more problems than it solves: + // - it delays panics arbitrarily, + // making bugs harder to detect; + // - it turns f's panic stack into a mere value, + // hiding it from crash-monitoring tools; + // - it risks deadlocks that hide the panic entirely, + // if f's panic leaves the program in a state + // that prevents the Wait call from being reached. + // See #53757, #74275, #74304, #74306. + + if err := f(); err != nil { + g.errOnce.Do(func() { + g.err = err + if g.cancel != nil { + g.cancel(g.err) + } + }) + } + }() +} + +// TryGo calls the given function in a new goroutine only if the number of +// active goroutines in the group is currently below the configured limit. +// +// The return value reports whether the goroutine was started. +func (g *Group) TryGo(f func() error) bool { + if g.sem != nil { + select { + case g.sem <- token{}: + // Note: this allows barging if and only if channels in general allow barging. + default: + return false + } + } + + g.wg.Add(1) + go func() { + defer g.done() + + if err := f(); err != nil { + g.errOnce.Do(func() { + g.err = err + if g.cancel != nil { + g.cancel(g.err) + } + }) + } + }() + return true +} + +// SetLimit limits the number of active goroutines in this group to at most n. +// A negative value indicates no limit. +// A limit of zero will prevent any new goroutines from being added. +// +// Any subsequent call to the Go method will block until it can add an active +// goroutine without exceeding the configured limit. +// +// The limit must not be modified while any goroutines in the group are active. +func (g *Group) SetLimit(n int) { + if n < 0 { + g.sem = nil + return + } + if active := len(g.sem); active != 0 { + panic(fmt.Errorf("errgroup: modify limit while %v goroutines in the group are still active", active)) + } + g.sem = make(chan token, n) +} diff --git a/vendor/modules.txt b/vendor/modules.txt index 4008cfd3..7d75824f 100644 --- a/vendor/modules.txt +++ b/vendor/modules.txt @@ -854,6 +854,7 @@ golang.org/x/oauth2/jws golang.org/x/oauth2/jwt # golang.org/x/sync v0.22.0 ## explicit; go 1.25.0 +golang.org/x/sync/errgroup golang.org/x/sync/semaphore golang.org/x/sync/singleflight # golang.org/x/sys v0.47.0