Fix bounded OneDrive search pagination

internal/graphapi/drive.go

@@ -350,19 +350,63 @@ func (c *Client) SearchDrive(ctx context.Context, driveID, query string, top int
 		return nil, fmt.Errorf("search query contains invalid characters")
 	}
 	top = clampTop(top)
-	resp, err := c.inner.Drives().ByDriveId(driveID).SearchWithQ(&query).Get(ctx, nil)
+	pageTop := top
+	resp, err := c.inner.Drives().ByDriveId(driveID).SearchWithQ(&query).GetAsSearchWithQGetResponse(ctx, &drives.ItemSearchWithQRequestBuilderGetRequestConfiguration{
+		QueryParameters: &drives.ItemSearchWithQRequestBuilderGetQueryParameters{Top: &pageTop},
+	})
 	if err != nil {
 		return nil, scopeUpgradeError("searching drive", err)
 	}
-	result := make([]DriveItem, 0, len(resp.GetValue()))
-	for _, d := range resp.GetValue() {
-		result = append(result, convertDriveItem(d))
+	if resp == nil {
+		return nil, fmt.Errorf("searching drive: Graph returned no response")
 	}
-	// Limit results to top
-	if int32(len(result)) > top {
-		result = result[:top]
+
+	result := make([]DriveItem, 0, top)
+	seenIDs := make(map[string]struct{})
+	seenLinks := make(map[string]struct{})
+	for {
+		values := resp.GetValue()
+		nextLink := derefStr(resp.GetOdataNextLink())
+		if len(values) == 0 && nextLink != "" {
+			return nil, fmt.Errorf("searching drive: continuation made no progress")
+		}
+		for _, value := range values {
+			if int32(len(result)) == top {
+				return result, nil
+			}
+			if value == nil || value.GetId() == nil || *value.GetId() == "" {
+				return nil, fmt.Errorf("searching drive: result has no item ID")
+			}
+			id := *value.GetId()
+			if _, exists := seenIDs[id]; exists {
+				return nil, fmt.Errorf("searching drive: duplicate item ID %q", id)
+			}
+			seenIDs[id] = struct{}{}
+			result = append(result, convertDriveItem(value))
+		}
+		if int32(len(result)) == top || nextLink == "" {
+			return result, nil
+		}
+		if _, exists := seenLinks[nextLink]; exists {
+			return nil, fmt.Errorf("searching drive: continuation repeated a previous URL")
+		}
+		seenLinks[nextLink] = struct{}{}
+		if err := validateGraphContinuation(nextLink, graphContinuationScope{
+			host:           defaultGraphAPIHost,
+			collectionPath: fmt.Sprintf("/v1.0/drives/%s/search(q='%s')", url.PathEscape(driveID), query),
+		}); err != nil {
+			return nil, fmt.Errorf("searching drive: %w", err)
+		}
+		// Graph continuation URLs are opaque. Follow the validated URL exactly
+		// as returned and enforce the caller's bound while collecting results.
+		resp, err = drives.NewItemSearchWithQRequestBuilder(nextLink, c.inner.GetAdapter()).GetAsSearchWithQGetResponse(ctx, nil)
+		if err != nil {
+			return nil, scopeUpgradeError("searching drive", err)
+		}
+		if resp == nil {
+			return nil, fmt.Errorf("searching drive: Graph returned no response")
+		}
 	}
-	return result, nil
 }
 
 // RecentDriveItems returns recently accessed items.

internal/graphapi/drive_search_test.go

@@ -0,0 +1,125 @@
+package graphapi
+
+import (
+	"context"
+	"net/http"
+	"reflect"
+	"strings"
+	"testing"
+)
+
+func TestSearchDriveCollectsBoundedPages(t *testing.T) {
+	requests := 0
+	client := testGraphClient(t, func(req *http.Request) *http.Response {
+		requests++
+		switch requests {
+		case 1:
+			if got := req.URL.Path; got != "/v1.0/drives/drive-id/search(q='budget')" {
+				t.Errorf("first request path = %q", got)
+			}
+			if got := req.URL.Query().Get("$top"); got != "3" {
+				t.Errorf("first request top = %q, want 3", got)
+			}
+			return graphJSONResponse(req, `{"value":[{"id":"one","name":"one.docx"},{"id":"two","name":"two.xlsx"}],"@odata.nextLink":"https://graph.microsoft.com/v1.0/drives/drive-id/search(q='budget')?$skiptoken=opaque%20cursor&foo=a%2Bb&$top=3"}`)
+		case 2:
+			const wantQuery = "$skiptoken=opaque%20cursor&foo=a%2Bb&$top=3"
+			if got := req.URL.RawQuery; got != wantQuery {
+				t.Errorf("continuation query = %q, want opaque query %q", got, wantQuery)
+			}
+			return graphJSONResponse(req, `{"value":[{"id":"three","name":"three.pptx"},{"id":"four","name":"four.pdf"}]}`)
+		default:
+			t.Fatalf("unexpected request %d: %s", requests, req.URL)
+			return nil
+		}
+	})
+
+	items, err := client.SearchDrive(context.Background(), "drive-id", "budget", 3)
+	if err != nil {
+		t.Fatalf("SearchDrive() error = %v", err)
+	}
+	if requests != 2 {
+		t.Fatalf("request count = %d, want 2", requests)
+	}
+	got := []string{items[0].ID, items[1].ID, items[2].ID}
+	if !reflect.DeepEqual(got, []string{"one", "two", "three"}) {
+		t.Errorf("item IDs = %v", got)
+	}
+}
+
+func TestSearchDriveRejectsUnsafeContinuation(t *testing.T) {
+	for _, tc := range []struct {
+		name     string
+		nextLink string
+		want     string
+	}{
+		{name: "host", nextLink: "https://example.com/v1.0/drives/drive-id/search(q='budget')?$skiptoken=next", want: "unexpected host"},
+		{name: "drive", nextLink: "https://graph.microsoft.com/v1.0/drives/other-drive/search(q='budget')?$skiptoken=next", want: "outside expected collection"},
+		{name: "query", nextLink: "https://graph.microsoft.com/v1.0/drives/drive-id/search(q='other')?$skiptoken=next", want: "outside expected collection"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			requests := 0
+			client := testGraphClient(t, func(req *http.Request) *http.Response {
+				requests++
+				return graphJSONResponse(req, `{"value":[{"id":"one"}],"@odata.nextLink":"`+tc.nextLink+`"}`)
+			})
+
+			items, err := client.SearchDrive(context.Background(), "drive-id", "budget", 2)
+			if err == nil || !strings.Contains(err.Error(), tc.want) {
+				t.Fatalf("SearchDrive() error = %v, want %q", err, tc.want)
+			}
+			if items != nil {
+				t.Fatalf("SearchDrive() items = %v, want nil", items)
+			}
+			if requests != 1 {
+				t.Errorf("request count = %d, want 1", requests)
+			}
+		})
+	}
+}
+
+func TestSearchDriveRejectsBadPageProgress(t *testing.T) {
+	for _, tc := range []struct {
+		name       string
+		firstBody  string
+		secondBody string
+		want       string
+	}{
+		{
+			name:      "empty continuation page",
+			firstBody: `{"value":[],"@odata.nextLink":"https://graph.microsoft.com/v1.0/drives/drive-id/search(q='budget')?$skiptoken=next"}`,
+			want:      "made no progress",
+		},
+		{
+			name:       "duplicate item",
+			firstBody:  `{"value":[{"id":"one"}],"@odata.nextLink":"https://graph.microsoft.com/v1.0/drives/drive-id/search(q='budget')?$skiptoken=next"}`,
+			secondBody: `{"value":[{"id":"one"}]}`,
+			want:       "duplicate item ID",
+		},
+		{
+			name:       "repeated URL",
+			firstBody:  `{"value":[{"id":"one"}],"@odata.nextLink":"https://graph.microsoft.com/v1.0/drives/drive-id/search(q='budget')?$skiptoken=next"}`,
+			secondBody: `{"value":[{"id":"two"}],"@odata.nextLink":"https://graph.microsoft.com/v1.0/drives/drive-id/search(q='budget')?$skiptoken=next"}`,
+			want:       "repeated a previous URL",
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			requests := 0
+			client := testGraphClient(t, func(req *http.Request) *http.Response {
+				requests++
+				body := tc.firstBody
+				if requests == 2 {
+					body = tc.secondBody
+				}
+				return graphJSONResponse(req, body)
+			})
+
+			items, err := client.SearchDrive(context.Background(), "drive-id", "budget", 3)
+			if err == nil || !strings.Contains(err.Error(), tc.want) {
+				t.Fatalf("SearchDrive() error = %v, want %q", err, tc.want)
+			}
+			if items != nil {
+				t.Fatalf("SearchDrive() items = %v, want nil", items)
+			}
+		})
+	}
+}

internal/graphapi/mail_pages.go

@@ -203,12 +203,15 @@ const (
 	graphMessageDeltaCollection
 	graphCalendarViewDeltaCollection
 	graphContactsDeltaCollection
+	graphDriveSearchCollection
 )
 
 type graphCollectionRoute struct {
 	isMe      bool
 	userID    string
 	folderID  string
+	driveID   string
+	query     string
 	operation graphCollectionOperation
 }
 
@@ -243,6 +246,27 @@ func parseGraphCollectionRoute(path string) (graphCollectionRoute, bool) {
 		}
 		route.userID = userID
 		collection = segments[3:]
+	case strings.EqualFold(segments[1], "drives") && len(segments) == 4:
+		driveID, ok := graphPathSegment(segments[2])
+		if !ok || driveID == "" {
+			return graphCollectionRoute{}, false
+		}
+		operation, ok := graphPathSegment(segments[3])
+		if !ok {
+			return graphCollectionRoute{}, false
+		}
+		const searchPrefix = "search(q='"
+		if len(operation) <= len(searchPrefix) || !strings.EqualFold(operation[:len(searchPrefix)], searchPrefix) || !strings.HasSuffix(operation, "')") {
+			return graphCollectionRoute{}, false
+		}
+		query, ok := graphODataString(operation[len(searchPrefix) : len(operation)-2])
+		if !ok || query == "" {
+			return graphCollectionRoute{}, false
+		}
+		route.driveID = driveID
+		route.query = query
+		route.operation = graphDriveSearchCollection
+		return route, true
 	default:
 		return graphCollectionRoute{}, false
 	}