Apply lockdown check to pull_request_readget_diff/get_files

pkg/github/pullrequests.go

@@ -123,13 +123,13 @@ Possible options:
  result, err := GetPullRequest(ctx, client, deps, owner, repo, pullNumber)
  return attachIFC(result), nil, err
  case "get_diff":
- result, err := GetPullRequestDiff(ctx, client, owner, repo, pullNumber)
+ result, err := GetPullRequestDiff(ctx, client, deps, owner, repo, pullNumber)
  return attachIFC(result), nil, err
  case "get_status":
  result, err := GetPullRequestStatus(ctx, client, owner, repo, pullNumber)
  return attachIFC(result), nil, err
  case "get_files":
- result, err := GetPullRequestFiles(ctx, client, owner, repo, pullNumber, pagination)
+ result, err := GetPullRequestFiles(ctx, client, deps, owner, repo, pullNumber, pagination)
  return attachIFC(result), nil, err
  case "get_commits":
  result, err := GetPullRequestCommits(ctx, client, owner, repo, pullNumber, pagination)
@@ -206,7 +206,40 @@ func GetPullRequest(ctx context.Context, client *github.Client, deps ToolDepende
 	return MarshalledTextResult(minimalPR), nil
 }
 
-func GetPullRequestDiff(ctx context.Context, client *github.Client, owner, repo string, pullNumber int) (*mcp.CallToolResult, error) {
+// enforcePullRequestLockdown returns a restricted tool result when lockdown mode is
+// enabled and the pull request author is not a safe content source for owner/repo,
+// and (nil, nil) otherwise. It fetches the pull request to resolve the author and is
+// a no-op that performs no request when lockdown mode is disabled.
+func enforcePullRequestLockdown(ctx context.Context, client *github.Client, deps ToolDependencies, owner, repo string, pullNumber int) (*mcp.CallToolResult, error) {
+	if !deps.GetFlags(ctx).LockdownMode {
+ return nil, nil
+	}
+	cache, err := deps.GetRepoAccessCache(ctx)
+	if err != nil {
+ return nil, fmt.Errorf("failed to get repo access cache: %w", err)
+	}
+	pr, resp, err := client.PullRequests.Get(ctx, owner, repo, pullNumber)
+	if err != nil {
+ return ghErrors.NewGitHubAPIErrorResponse(ctx, "failed to get pull request", resp, err), nil
+	}
+	defer func() { _ = resp.Body.Close() }()
+
+	if resp.StatusCode != http.StatusOK {
+ body, err := io.ReadAll(resp.Body)
+ if err != nil {
+ return nil, fmt.Errorf("failed to read response body: %w", err)
+ }
+ return ghErrors.NewGitHubAPIStatusErrorResponse(ctx, "failed to get pull request", resp, body), nil
+	}
+
+	return authorLockdownResult(ctx, cache, owner, repo, pr.GetUser().GetLogin(), lockdownPullRequestRestrictedMessage)
+}
+
+func GetPullRequestDiff(ctx context.Context, client *github.Client, deps ToolDependencies, owner, repo string, pullNumber int) (*mcp.CallToolResult, error) {
+	if restricted, err := enforcePullRequestLockdown(ctx, client, deps, owner, repo, pullNumber); restricted != nil || err != nil {
+ return restricted, err
+	}
+
 	raw, resp, err := client.PullRequests.GetRaw(
  ctx,
  owner,
@@ -347,7 +380,11 @@ func GetPullRequestCheckRuns(ctx context.Context, client *github.Client, owner,
 	return utils.NewToolResultText(string(r)), nil
 }
 
-func GetPullRequestFiles(ctx context.Context, client *github.Client, owner, repo string, pullNumber int, pagination PaginationParams) (*mcp.CallToolResult, error) {
+func GetPullRequestFiles(ctx context.Context, client *github.Client, deps ToolDependencies, owner, repo string, pullNumber int, pagination PaginationParams) (*mcp.CallToolResult, error) {
+	if restricted, err := enforcePullRequestLockdown(ctx, client, deps, owner, repo, pullNumber); restricted != nil || err != nil {
+ return restricted, err
+	}
+
 	opts := &github.ListOptions{
  PerPage: pagination.PerPage,
  Page: pagination.Page,
@@ -552,17 +589,18 @@ func GetPullRequestReviews(ctx context.Context, client *github.Client, deps Tool
  filteredReviews := make([]*github.PullRequestReview, 0, len(reviews))
  for _, review := range reviews {
  login := review.GetUser().GetLogin()
- if login != "" {
-isSafeContent, err := cache.IsSafeContent(ctx, login, owner, repo)
-	if err != nil {
- return nil, fmt.Errorf("failed to check lockdown mode: %w", err)
-	}
-if isSafeContent {
- filteredReviews = append(filteredReviews, review)
-	}
-reviews = filteredReviews
+ if login == "" {
+continue
+}
+isSafeContent, err := cache.IsSafeContent(ctx, login, owner, repo)
+if err != nil {
+return nil, fmt.Errorf("failed to check lockdown mode: %w", err)
+}
+if isSafeContent {
+filteredReviews = append(filteredReviews, review)
  }
  }
+ reviews = filteredReviews
 	}
 
 	minimalReviews := make([]MinimalPullRequestReview, 0, len(reviews))

pkg/github/pullrequests_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"net/http"
+	"strings"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -1192,12 +1193,14 @@ func Test_GetPullRequestFiles(t *testing.T) {
 	}
 
 	tests := []struct {
- name string
- mockedClient *http.Client
- requestArgs map[string]any
- expectError bool
- expectedFiles []*github.CommitFile
- expectedErrMsg string
+ name string
+ mockedClient *http.Client
+ requestArgs map[string]any
+ expectError bool
+ expectedFiles []*github.CommitFile
+ expectedErrMsg string
+ lockdownEnabled bool
+ restPermission string
 	}{
  {
  name: "successful files fetch",
@@ -1261,17 +1264,81 @@ func Test_GetPullRequestFiles(t *testing.T) {
  expectError: true,
  expectedErrMsg: "failed to get pull request files",
  },
+ {
+ name: "lockdown enabled - author lacks push access",
+ mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{
+ GetReposPullsByOwnerByRepoByPullNumber: mockResponse(t, http.StatusOK, &github.PullRequest{
+ Number: github.Ptr(42),
+ User: &github.User{Login: github.Ptr("reader")},
+ }),
+ }),
+ requestArgs: map[string]any{
+ "method": "get_files",
+ "owner": "owner",
+ "repo": "repo",
+ "pullNumber": float64(42),
+ },
+ lockdownEnabled: true,
+ restPermission: "read",
+ expectError: true,
+ expectedErrMsg: "access to pull request is restricted by lockdown mode",
+ },
+ {
+ name: "lockdown enabled - author has push access",
+ mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{
+ GetReposPullsByOwnerByRepoByPullNumber: mockResponse(t, http.StatusOK, &github.PullRequest{
+ Number: github.Ptr(42),
+ User: &github.User{Login: github.Ptr("writer")},
+ }),
+ GetReposPullsFilesByOwnerByRepoByPullNumber: mockResponse(t, http.StatusOK, mockFiles),
+ }),
+ requestArgs: map[string]any{
+ "method": "get_files",
+ "owner": "owner",
+ "repo": "repo",
+ "pullNumber": float64(42),
+ },
+ lockdownEnabled: true,
+ restPermission: "write",
+ expectError: false,
+ expectedFiles: mockFiles,
+ },
+ {
+ name: "lockdown enabled - pull request fetch fails",
+ mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{
+ GetReposPullsByOwnerByRepoByPullNumber: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+ w.WriteHeader(http.StatusNotFound)
+ _, _ = w.Write([]byte(`{"message": "Not Found"}`))
+ }),
+ }),
+ requestArgs: map[string]any{
+ "method": "get_files",
+ "owner": "owner",
+ "repo": "repo",
+ "pullNumber": float64(999),
+ },
+ lockdownEnabled: true,
+ restPermission: "read",
+ expectError: true,
+ expectedErrMsg: "failed to get pull request",
+ },
 	}
 
 	for _, tc := range tests {
  t.Run(tc.name, func(t *testing.T) {
  // Setup client with mock
  client := mustNewGHClient(t, tc.mockedClient)
  serverTool := PullRequestRead(translations.NullTranslationHelper)
+
+ var restClient *github.Client
+ if tc.lockdownEnabled {
+ restClient = mockRESTPermissionServer(t, tc.restPermission, nil)
+ }
+
  deps := BaseDeps{
  Client: client,
- RepoAccessCache: stubRepoAccessCache(nil, 5*time.Minute),
- Flags: stubFeatureFlags(map[string]bool{"lockdown-mode": false}),
+ RepoAccessCache: stubRepoAccessCache(restClient, 5*time.Minute),
+ Flags: stubFeatureFlags(map[string]bool{"lockdown-mode": tc.lockdownEnabled}),
  }
  handler := serverTool.Handler(deps)
 
@@ -2429,6 +2496,33 @@ func Test_GetPullRequestReviews(t *testing.T) {
  },
  lockdownEnabled: true,
  },
+ {
+ name: "lockdown enabled filters reviews with empty author login",
+ mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{
+ GetReposPullsReviewsByOwnerByRepoByPullNumber: mockResponse(t, http.StatusOK, []*github.PullRequestReview{
+ {
+ ID: github.Ptr(int64(2040)),
+ State: github.Ptr("APPROVED"),
+ Body: github.Ptr("Ghost review"),
+ User: &github.User{Login: github.Ptr("")},
+ },
+ {
+ ID: github.Ptr(int64(2041)),
+ State: github.Ptr("COMMENTED"),
+ Body: github.Ptr("Another ghost review"),
+ },
+ }),
+ }),
+ requestArgs: map[string]any{
+ "method": "get_reviews",
+ "owner": "owner",
+ "repo": "repo",
+ "pullNumber": float64(42),
+ },
+ expectError: false,
+ expectedReviews: []*github.PullRequestReview{},
+ lockdownEnabled: true,
+ },
 	}
 
 	for _, tc := range tests {
@@ -3834,10 +3928,30 @@ index 5d6e7b2..8a4f5c3 100644
 +
 +This is a new section added in the pull request.`
 
+	// Under lockdown the diff path first fetches the PR as JSON to resolve the
+	// author, then the raw diff; branch on the Accept header to serve both.
+	prOrDiffHandler := func(authorLogin string) http.HandlerFunc {
+ mockPR := &github.PullRequest{
+ Number: github.Ptr(42),
+ User: &github.User{Login: github.Ptr(authorLogin)},
+ }
+ return func(w http.ResponseWriter, r *http.Request) {
+ if strings.Contains(r.Header.Get("Accept"), "diff") {
+ w.WriteHeader(http.StatusOK)
+ _, _ = w.Write([]byte(stubbedDiff))
+ return
+ }
+ w.Header().Set("Content-Type", "application/json")
+ _ = json.NewEncoder(w).Encode(mockPR)
+ }
+	}
+
 	tests := []struct {
  name string
  requestArgs map[string]any
  mockedClient *http.Client
+ lockdownEnabled bool
+ restPermission string
  expectToolError bool
  expectedToolErrMsg string
 	}{
@@ -3856,6 +3970,37 @@ index 5d6e7b2..8a4f5c3 100644
  }),
  expectToolError: false,
  },
+ {
+ name: "lockdown enabled - author lacks push access",
+ requestArgs: map[string]any{
+ "method": "get_diff",
+ "owner": "owner",
+ "repo": "repo",
+ "pullNumber": float64(42),
+ },
+ mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{
+ GetReposPullsByOwnerByRepoByPullNumber: prOrDiffHandler("reader"),
+ }),
+ lockdownEnabled: true,
+ restPermission: "read",
+ expectToolError: true,
+ expectedToolErrMsg: "access to pull request is restricted by lockdown mode",
+ },
+ {
+ name: "lockdown enabled - author has push access",
+ requestArgs: map[string]any{
+ "method": "get_diff",
+ "owner": "owner",
+ "repo": "repo",
+ "pullNumber": float64(42),
+ },
+ mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{
+ GetReposPullsByOwnerByRepoByPullNumber: prOrDiffHandler("writer"),
+ }),
+ lockdownEnabled: true,
+ restPermission: "write",
+ expectToolError: false,
+ },
 	}
 
 	for _, tc := range tests {
@@ -3865,10 +4010,16 @@ index 5d6e7b2..8a4f5c3 100644
  // Setup client with mock
  client := mustNewGHClient(t, tc.mockedClient)
  serverTool := PullRequestRead(translations.NullTranslationHelper)
+
+ var restClient *github.Client
+ if tc.lockdownEnabled {
+ restClient = mockRESTPermissionServer(t, tc.restPermission, nil)
+ }
+
  deps := BaseDeps{
  Client: client,
- RepoAccessCache: stubRepoAccessCache(nil, 5*time.Minute),
- Flags: stubFeatureFlags(map[string]bool{"lockdown-mode": false}),
+ RepoAccessCache: stubRepoAccessCache(restClient, 5*time.Minute),
+ Flags: stubFeatureFlags(map[string]bool{"lockdown-mode": tc.lockdownEnabled}),
  }
  handler := serverTool.Handler(deps)