Pin audit log pagination end_time in page tokens

Author: sluongng

docs/enterprise-api.md

@@ -423,7 +423,9 @@ message GetAuditLogResponse {
 ### Polling Recommendations
 
 - Treat `next_page_token` as opaque. Do not parse it.
-- Keep `start_time` / `end_time` fixed while paginating through one window.
+- Keep `start_time` / `end_time` fixed while paginating through one window. If
+  `end_time` is omitted on the first page, the server pins it into the returned
+  cursor.
 - Resume by passing the returned `next_page_token` until it is empty.
 - For continuous export, query bounded windows (for example every minute) and checkpoint the last successful window end.
 

enterprise/server/api/api_server.go

@@ -403,6 +403,10 @@ func (s *APIServer) GetAuditLog(ctx context.Context, req *apipb.GetAuditLogReque
 	if selector == nil {
 		selector = &apipb.AuditLogSelector{}
 	}
+	pageToken, err := parseAuditLogPageToken(req.GetPageToken())
+	if err != nil {
+		return nil, err
+	}
 
 	startTime := selector.GetStartTime()
 	if startTime == nil {
@@ -422,6 +426,17 @@ func (s *APIServer) GetAuditLog(ctx context.Context, req *apipb.GetAuditLogReque
 	if startTime.AsTime().After(endTime.AsTime()) {
 		return nil, status.InvalidArgumentErrorf("start_time must not be after end_time")
 	}
+	startUsec := startTime.AsTime().UnixMicro()
+	endUsec := endTime.AsTime().UnixMicro()
+	if pageToken != nil {
+		if selector.GetEndTime() != nil && endUsec != pageToken.EndTimeUsec {
+			return nil, status.InvalidArgumentError("page_token does not match selector.end_time")
+		}
+		endUsec = pageToken.EndTimeUsec
+	}
+	if startUsec > endUsec {
+		return nil, status.InvalidArgumentErrorf("start_time must not be after end_time")
+	}
 
 	if s.env.GetAuditLogger() == nil || s.env.GetOLAPDBHandle() == nil {
 		return nil, status.UnimplementedError("Audit logger not configured")
@@ -447,19 +462,13 @@ func (s *APIServer) GetAuditLog(ctx context.Context, req *apipb.GetAuditLogReque
 		pageSize = defaultAuditLogPageSize
 	}
 	pageSize = min(pageSize, maxAuditLogPageSize)
-	startUsec := startTime.AsTime().UnixMicro()
-	endUsec := endTime.AsTime().UnixMicro()
 
 	q := query_builder.NewQuery(`SELECT * FROM AuditLogs`)
 	q.AddWhereClause("group_id = ?", user.GetGroupID())
 	q.AddWhereClause("event_time_usec >= ?", startUsec)
 	q.AddWhereClause("event_time_usec < ?", endUsec)
-	if req.GetPageToken() != "" {
-		token, err := parseAuditLogPageToken(req.GetPageToken())
-		if err != nil {
-			return nil, err
-		}
-		q.AddWhereClause("(event_time_usec, audit_log_id) > (?, ?)", token.EventTimeUsec, token.AuditLogID)
+	if pageToken != nil {
+		q.AddWhereClause("(event_time_usec, audit_log_id) > (?, ?)", pageToken.EventTimeUsec, pageToken.AuditLogID)
 	}
 	// Match AuditLogs sort key to keep keyset pagination index-friendly.
 	q.SetOrderBy("group_id, event_time_usec, audit_log_id", true)
@@ -491,6 +500,7 @@ func (s *APIServer) GetAuditLog(ctx context.Context, req *apipb.GetAuditLogReque
 
 		rsp.Entry = append(rsp.Entry, entry)
 		lastReturnedToken = auditLogPageToken{
+			EndTimeUsec:   endUsec,
 			EventTimeUsec: row.EventTimeUsec,
 			AuditLogID:    row.AuditLogID,
 		}
@@ -504,6 +514,7 @@ func (s *APIServer) GetAuditLog(ctx context.Context, req *apipb.GetAuditLogReque
 }
 
 type auditLogPageToken struct {
+	EndTimeUsec   int64  `json:"end_time_usec"`
 	EventTimeUsec int64  `json:"event_time_usec"`
 	AuditLogID    string `json:"audit_log_id"`
 }
@@ -516,17 +527,20 @@ func encodeAuditLogPageToken(token auditLogPageToken) (string, error) {
 	return base64.RawURLEncoding.EncodeToString(data), nil
 }
 
-func parseAuditLogPageToken(pageToken string) (auditLogPageToken, error) {
+func parseAuditLogPageToken(pageToken string) (*auditLogPageToken, error) {
+	if pageToken == "" {
+		return nil, nil
+	}
 	data, err := base64.RawURLEncoding.DecodeString(pageToken)
 	if err != nil {
-		return auditLogPageToken{}, status.InvalidArgumentErrorf("invalid page_token")
+		return nil, status.InvalidArgumentErrorf("invalid page_token")
 	}
-	token := auditLogPageToken{}
-	if err := json.Unmarshal(data, &token); err != nil {
-		return auditLogPageToken{}, status.InvalidArgumentErrorf("invalid page_token")
+	token := &auditLogPageToken{}
+	if err := json.Unmarshal(data, token); err != nil {
+		return nil, status.InvalidArgumentErrorf("invalid page_token")
 	}
 	if token.AuditLogID == "" {
-		return auditLogPageToken{}, status.InvalidArgumentErrorf("invalid page_token")
+		return nil, status.InvalidArgumentErrorf("invalid page_token")
 	}
 	return token, nil
 }

enterprise/server/api/api_server_test.go

@@ -312,22 +312,18 @@ func TestGetAuditLog(t *testing.T) {
 	require.NoError(t, err)
 	keyCtx := env.GetAuthenticator().AuthContextFromAPIKey(ctx, key.Value)
 
+	env.GetAuditLogger().LogForGroup(userCtx, groupID, alpb.Action_UPDATE, &grpb.UpdateGroupRequest{Name: "before-window"})
+	time.Sleep(2 * time.Millisecond)
+	windowStart := time.Now()
+	time.Sleep(2 * time.Millisecond)
 	env.GetAuditLogger().LogForGroup(userCtx, groupID, alpb.Action_UPDATE, &grpb.UpdateGroupRequest{Name: "update-1"})
 	env.GetAuditLogger().LogForGroup(userCtx, groupID, alpb.Action_UPDATE, &grpb.UpdateGroupRequest{Name: "update-2"})
 
 	s := NewAPIServer(env)
 	selector := &apipb.AuditLogSelector{
-		StartTime: timestamppb.New(time.Unix(0, 0)),
-		EndTime:   timestamppb.New(time.Now().Add(time.Minute)),
+		StartTime: timestamppb.New(windowStart),
 	}
 
-	directRsp, err := env.GetAuditLogger().GetLogs(keyCtx, &alpb.GetAuditLogsRequest{
-		TimestampAfter:  selector.GetStartTime(),
-		TimestampBefore: selector.GetEndTime(),
-	})
-	require.NoError(t, err)
-	require.Len(t, directRsp.GetEntries(), 2)
-
 	page1, err := s.GetAuditLog(keyCtx, &apipb.GetAuditLogRequest{
 		Selector: selector,
 		PageSize: 1,
@@ -337,8 +333,10 @@ func TestGetAuditLog(t *testing.T) {
 	require.NotEmpty(t, page1.GetNextPageToken())
 	require.NotContains(t, page1.GetNextPageToken(), "|")
 
+	time.Sleep(2 * time.Millisecond)
+	env.GetAuditLogger().LogForGroup(userCtx, groupID, alpb.Action_UPDATE, &grpb.UpdateGroupRequest{Name: "update-3"})
+
 	page2, err := s.GetAuditLog(keyCtx, &apipb.GetAuditLogRequest{
-		Selector:  selector,
 		PageSize:  1,
 		PageToken: page1.GetNextPageToken(),
 	})