Skip to content

Commit 35cca24

Browse files
RossTarrantRossTarrant
committed
centralise sanitisation for inputs and outputs
1 parent eb088df commit 35cca24

32 files changed

Lines changed: 2609 additions & 417 deletions

‎pkg/github/__toolsnaps__/search_code.snap‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,9 @@
3232
},
3333
"sort": {
3434
"description": "Sort field ('indexed' only)",
35+
"enum": [
36+
"indexed"
37+
],
3538
"type": "string"
3639
}
3740
},

‎pkg/github/__toolsnaps__/search_code_ff_fields_param.snap‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,9 @@
4646
},
4747
"sort": {
4848
"description": "Sort field ('indexed' only)",
49+
"enum": [
50+
"indexed"
51+
],
4952
"type": "string"
5053
}
5154
},

‎pkg/github/actions.go‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ func handleFailedJobLogs(ctx context.Context, client *github.Client, owner, repo
8484
// Continue with other jobs even if one fails
8585
jobResult = map[string]any{
8686
"job_id": job.GetID(),
87-
"job_name": job.GetName(),
87+
"job_name": sanitizeOutputText(job.GetName()),
8888
"error": err.Error(),
8989
}
9090
// Enable reporting of status codes and error causes
@@ -139,7 +139,7 @@ func getJobLogData(ctx context.Context, client *github.Client, owner, repo strin
139139
"job_id": jobID,
140140
}
141141
if jobName != "" {
142-
result["job_name"] = jobName
142+
result["job_name"] = sanitizeOutputText(jobName)
143143
}
144144

145145
if returnContent {
@@ -788,7 +788,7 @@ func getWorkflow(ctx context.Context, client *github.Client, owner, repo, resour
788788
}
789789

790790
defer func() { _ = resp.Body.Close() }()
791-
r, err := json.Marshal(workflow)
791+
r, err := json.Marshal(sanitizedWorkflowCopy(workflow))
792792
if err != nil {
793793
return nil, nil, fmt.Errorf("failed to marshal workflow: %w", err)
794794
}
@@ -802,7 +802,7 @@ func getWorkflowRun(ctx context.Context, client *github.Client, owner, repo stri
802802
return ghErrors.NewGitHubAPIErrorResponse(ctx, "failed to get workflow run", resp, err), nil, nil
803803
}
804804
defer func() { _ = resp.Body.Close() }()
805-
r, err := json.Marshal(workflowRun)
805+
r, err := json.Marshal(sanitizedWorkflowRunCopy(workflowRun))
806806
if err != nil {
807807
return nil, nil, fmt.Errorf("failed to marshal workflow run: %w", err)
808808
}
@@ -815,7 +815,7 @@ func getWorkflowJob(ctx context.Context, client *github.Client, owner, repo stri
815815
return ghErrors.NewGitHubAPIErrorResponse(ctx, "failed to get workflow job", resp, err), nil, nil
816816
}
817817
defer func() { _ = resp.Body.Close() }()
818-
r, err := json.Marshal(workflowJob)
818+
r, err := json.Marshal(sanitizedWorkflowJobCopy(workflowJob))
819819
if err != nil {
820820
return nil, nil, fmt.Errorf("failed to marshal workflow job: %w", err)
821821
}
@@ -834,7 +834,7 @@ func listWorkflows(ctx context.Context, client *github.Client, owner, repo strin
834834
}
835835
defer func() { _ = resp.Body.Close() }()
836836

837-
r, err := json.Marshal(workflows)
837+
r, err := json.Marshal(sanitizedWorkflowsCopy(workflows))
838838
if err != nil {
839839
return nil, nil, fmt.Errorf("failed to marshal workflows: %w", err)
840840
}
@@ -884,7 +884,7 @@ func listWorkflowRuns(ctx context.Context, client *github.Client, args map[strin
884884
}
885885

886886
defer func() { _ = resp.Body.Close() }()
887-
r, err := json.Marshal(workflowRuns)
887+
r, err := json.Marshal(sanitizedWorkflowRunsCopy(workflowRuns))
888888
if err != nil {
889889
return nil, nil, fmt.Errorf("failed to marshal workflow runs: %w", err)
890890
}
@@ -919,7 +919,7 @@ func listWorkflowJobs(ctx context.Context, client *github.Client, args map[strin
919919
}
920920

921921
response := map[string]any{
922-
"jobs": workflowJobs,
922+
"jobs": sanitizedWorkflowJobsCopy(workflowJobs),
923923
}
924924

925925
defer func() { _ = resp.Body.Close() }()

‎pkg/github/actions_test.go‎

Lines changed: 109 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -114,22 +114,66 @@ func Test_ActionsList_ListWorkflows(t *testing.T) {
114114
}
115115
}
116116

117+
func unsafeWorkflowRunFixture() *github.WorkflowRun {
118+
return &github.WorkflowRun{
119+
ID: github.Ptr(int64(12345)),
120+
Name: github.Ptr(baselineUnsafeText),
121+
DisplayTitle: github.Ptr(baselineUnsafeText),
122+
HeadBranch: github.Ptr("feature/exact<script>"),
123+
HeadSHA: github.Ptr("abc123"),
124+
HeadCommit: &github.HeadCommit{
125+
Message: github.Ptr(baselineUnsafeText),
126+
Author: &github.CommitAuthor{
127+
Name: github.Ptr(baselineUnsafeText),
128+
Email: github.Ptr("author@example.com"),
129+
},
130+
Added: []string{"src/exact<script>.go"},
131+
},
132+
PullRequests: []*github.PullRequest{{
133+
Title: github.Ptr(baselineUnsafeText),
134+
Body: github.Ptr(baselineUnsafeText),
135+
Labels: []*github.Label{{Name: github.Ptr(baselineUnsafeText), Description: github.Ptr(baselineUnsafeText)}},
136+
Milestone: &github.Milestone{Title: github.Ptr(baselineUnsafeText), Description: github.Ptr(baselineUnsafeText)},
137+
Head: &github.PullRequestBranch{Repo: &github.Repository{Description: github.Ptr(baselineUnsafeText)}},
138+
Base: &github.PullRequestBranch{Repo: &github.Repository{Description: github.Ptr(baselineUnsafeText)}},
139+
}},
140+
Repository: &github.Repository{Description: github.Ptr(baselineUnsafeText)},
141+
HeadRepository: &github.Repository{Description: github.Ptr(baselineUnsafeText)},
142+
Status: github.Ptr("completed"),
143+
Conclusion: github.Ptr("success"),
144+
}
145+
}
146+
147+
func assertSanitizedWorkflowRun(t *testing.T, run *github.WorkflowRun) {
148+
t.Helper()
149+
expected := sanitizeOutputText(baselineUnsafeText)
150+
assert.Equal(t, expected, run.GetName())
151+
assert.Equal(t, expected, run.GetDisplayTitle())
152+
assert.Equal(t, expected, run.HeadCommit.GetMessage())
153+
assert.Equal(t, expected, run.HeadCommit.Author.GetName())
154+
assert.Equal(t, "author@example.com", run.HeadCommit.Author.GetEmail())
155+
assert.Equal(t, []string{"src/exact<script>.go"}, run.HeadCommit.Added)
156+
assert.Equal(t, expected, run.PullRequests[0].GetTitle())
157+
assert.Equal(t, expected, run.PullRequests[0].Labels[0].GetName())
158+
assert.Equal(t, expected, run.PullRequests[0].Milestone.GetTitle())
159+
assert.Equal(t, expected, run.PullRequests[0].Head.Repo.GetDescription())
160+
assert.Equal(t, expected, run.PullRequests[0].Base.Repo.GetDescription())
161+
assert.Equal(t, expected, run.Repository.GetDescription())
162+
assert.Equal(t, expected, run.HeadRepository.GetDescription())
163+
assert.Equal(t, "feature/exact<script>", run.GetHeadBranch())
164+
assert.Equal(t, "abc123", run.GetHeadSHA())
165+
}
166+
117167
func Test_ActionsList_ListWorkflowRuns(t *testing.T) {
118168
toolDef := ActionsList(translations.NullTranslationHelper)
119169

120170
t.Run("successful workflow runs list", func(t *testing.T) {
171+
workflowRun := unsafeWorkflowRunFixture()
121172
mockedClient := MockHTTPClientWithHandlers(map[string]http.HandlerFunc{
122173
GetReposActionsWorkflowsRunsByOwnerByRepoByWorkflowID: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
123174
runs := &github.WorkflowRuns{
124-
TotalCount: github.Ptr(1),
125-
WorkflowRuns: []*github.WorkflowRun{
126-
{
127-
ID: github.Ptr(int64(123)),
128-
Name: github.Ptr("CI"),
129-
Status: github.Ptr("completed"),
130-
Conclusion: github.Ptr("success"),
131-
},
132-
},
175+
TotalCount: github.Ptr(1),
176+
WorkflowRuns: []*github.WorkflowRun{workflowRun},
133177
}
134178
w.WriteHeader(http.StatusOK)
135179
_ = json.NewEncoder(w).Encode(runs)
@@ -158,6 +202,9 @@ func Test_ActionsList_ListWorkflowRuns(t *testing.T) {
158202
err = json.Unmarshal([]byte(textContent.Text), &response)
159203
require.NoError(t, err)
160204
assert.NotNil(t, response.TotalCount)
205+
require.Len(t, response.WorkflowRuns, 1)
206+
assertSanitizedWorkflowRun(t, response.WorkflowRuns[0])
207+
assert.Equal(t, baselineUnsafeText, workflowRun.GetName())
161208
})
162209

163210
t.Run("list all workflow runs without resource_id", func(t *testing.T) {
@@ -209,6 +256,55 @@ func Test_ActionsList_ListWorkflowRuns(t *testing.T) {
209256
})
210257
}
211258

259+
func Test_ActionsList_ListWorkflowJobs(t *testing.T) {
260+
toolDef := ActionsList(translations.NullTranslationHelper)
261+
workflowJob := &github.WorkflowJob{
262+
ID: github.Ptr(int64(123)),
263+
Name: github.Ptr(baselineUnsafeText),
264+
WorkflowName: github.Ptr(baselineUnsafeText),
265+
RunnerName: github.Ptr(baselineUnsafeText),
266+
RunnerGroupName: github.Ptr(baselineUnsafeText),
267+
Steps: []*github.TaskStep{{Name: github.Ptr(baselineUnsafeText)}},
268+
}
269+
mockedClient := MockHTTPClientWithHandlers(map[string]http.HandlerFunc{
270+
GetReposActionsRunsJobsByOwnerByRepoByRunID: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
271+
w.WriteHeader(http.StatusOK)
272+
_ = json.NewEncoder(w).Encode(&github.Jobs{
273+
TotalCount: github.Ptr(1),
274+
Jobs: []*github.WorkflowJob{workflowJob},
275+
})
276+
}),
277+
})
278+
279+
client := mustNewGHClient(t, mockedClient)
280+
deps := BaseDeps{Client: client}
281+
handler := toolDef.Handler(deps)
282+
request := createMCPRequest(map[string]any{
283+
"method": "list_workflow_jobs",
284+
"owner": "owner",
285+
"repo": "repo",
286+
"resource_id": "123",
287+
})
288+
result, err := handler(ContextWithDeps(context.Background(), deps), &request)
289+
require.NoError(t, err)
290+
require.False(t, result.IsError)
291+
292+
var response struct {
293+
Jobs *github.Jobs `json:"jobs"`
294+
}
295+
require.NoError(t, json.Unmarshal([]byte(getTextResult(t, result).Text), &response))
296+
require.NotNil(t, response.Jobs)
297+
require.Len(t, response.Jobs.Jobs, 1)
298+
returnedJob := response.Jobs.Jobs[0]
299+
assert.Equal(t, sanitizeOutputText(baselineUnsafeText), returnedJob.GetName())
300+
assert.Equal(t, sanitizeOutputText(baselineUnsafeText), returnedJob.GetWorkflowName())
301+
assert.Equal(t, baselineUnsafeText, returnedJob.GetRunnerName())
302+
assert.Equal(t, baselineUnsafeText, returnedJob.GetRunnerGroupName())
303+
assert.Equal(t, sanitizeOutputText(baselineUnsafeText), returnedJob.Steps[0].GetName())
304+
assert.Equal(t, baselineUnsafeText, workflowJob.GetName())
305+
assert.Equal(t, baselineUnsafeText, workflowJob.GetRunnerName())
306+
}
307+
212308
func Test_ActionsGet(t *testing.T) {
213309
// Verify tool definition once
214310
toolDef := ActionsGet(translations.NullTranslationHelper)
@@ -271,16 +367,11 @@ func Test_ActionsGet_GetWorkflowRun(t *testing.T) {
271367
toolDef := ActionsGet(translations.NullTranslationHelper)
272368

273369
t.Run("successful workflow run get", func(t *testing.T) {
370+
workflowRun := unsafeWorkflowRunFixture()
274371
mockedClient := MockHTTPClientWithHandlers(map[string]http.HandlerFunc{
275372
GetReposActionsRunsByOwnerByRepoByRunID: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
276-
run := &github.WorkflowRun{
277-
ID: github.Ptr(int64(12345)),
278-
Name: github.Ptr("CI"),
279-
Status: github.Ptr("completed"),
280-
Conclusion: github.Ptr("success"),
281-
}
282373
w.WriteHeader(http.StatusOK)
283-
_ = json.NewEncoder(w).Encode(run)
374+
_ = json.NewEncoder(w).Encode(workflowRun)
284375
}),
285376
})
286377

@@ -307,6 +398,8 @@ func Test_ActionsGet_GetWorkflowRun(t *testing.T) {
307398
require.NoError(t, err)
308399
assert.NotNil(t, response.ID)
309400
assert.Equal(t, int64(12345), *response.ID)
401+
assertSanitizedWorkflowRun(t, &response)
402+
assert.Equal(t, baselineUnsafeText, workflowRun.GetName())
310403
})
311404
}
312405

‎pkg/github/code_scanning.go‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ func GetCodeScanningAlert(t translations.TranslationHelperFunc) inventory.Server
8484
return ghErrors.NewGitHubAPIStatusErrorResponse(ctx, "failed to get alert", resp, body), nil, nil
8585
}
8686

87-
r, err := json.Marshal(alert)
87+
r, err := json.Marshal(sanitizedCodeScanningAlertCopy(alert))
8888
if err != nil {
8989
return utils.NewToolResultErrorFromErr("failed to marshal alert", err), nil, nil
9090
}
@@ -164,10 +164,16 @@ func ListCodeScanningAlerts(t translations.TranslationHelperFunc) inventory.Serv
164164
if err != nil {
165165
return utils.NewToolResultError(err.Error()), nil, nil
166166
}
167+
if err := validateEnumParam("state", state, "open", "closed", "dismissed", "fixed"); err != nil {
168+
return utils.NewToolResultError(err.Error()), nil, nil
169+
}
167170
severity, err := OptionalParam[string](args, "severity")
168171
if err != nil {
169172
return utils.NewToolResultError(err.Error()), nil, nil
170173
}
174+
if err := validateEnumParam("severity", severity, "critical", "high", "medium", "low", "warning", "note", "error"); err != nil {
175+
return utils.NewToolResultError(err.Error()), nil, nil
176+
}
171177
toolName, err := OptionalParam[string](args, "tool_name")
172178
if err != nil {
173179
return utils.NewToolResultError(err.Error()), nil, nil
@@ -209,7 +215,7 @@ func ListCodeScanningAlerts(t translations.TranslationHelperFunc) inventory.Serv
209215
return ghErrors.NewGitHubAPIStatusErrorResponse(ctx, "failed to list alerts", resp, body), nil, nil
210216
}
211217

212-
r, err := json.Marshal(alerts)
218+
r, err := json.Marshal(sanitizedCodeScanningAlertsCopy(alerts))
213219
if err != nil {
214220
return utils.NewToolResultErrorFromErr("failed to marshal alerts", err), nil, nil
215221
}

‎pkg/github/code_scanning_test.go‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,28 @@ func Test_ListCodeScanningAlerts(t *testing.T) {
209209
expectError: false,
210210
expectedAlerts: mockAlerts,
211211
},
212+
{
213+
name: "invalid state is rejected before request",
214+
mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{}),
215+
requestArgs: map[string]any{
216+
"owner": "owner",
217+
"repo": "repo",
218+
"state": "archived",
219+
},
220+
expectError: true,
221+
expectedErrMsg: "state must be one of: open, closed, dismissed, fixed",
222+
},
223+
{
224+
name: "invalid severity is rejected before request",
225+
mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{}),
226+
requestArgs: map[string]any{
227+
"owner": "owner",
228+
"repo": "repo",
229+
"severity": "extreme",
230+
},
231+
expectError: true,
232+
expectedErrMsg: "severity must be one of: critical, high, medium, low, warning, note, error",
233+
},
212234
{
213235
name: "alerts listing fails",
214236
mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{

‎pkg/github/dependabot.go‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ func GetDependabotAlert(t translations.TranslationHelperFunc) inventory.ServerTo
8585
return ghErrors.NewGitHubAPIStatusErrorResponse(ctx, "failed to get alert", resp, body), nil, nil
8686
}
8787

88-
r, err := json.Marshal(alert)
88+
r, err := json.Marshal(sanitizedDependabotAlertCopy(alert))
8989
if err != nil {
9090
return utils.NewToolResultErrorFromErr("failed to marshal alert", err), nil, err
9191
}
@@ -153,10 +153,16 @@ func ListDependabotAlerts(t translations.TranslationHelperFunc) inventory.Server
153153
if err != nil {
154154
return utils.NewToolResultError(err.Error()), nil, nil
155155
}
156+
if err := validateEnumParam("state", state, "open", "fixed", "dismissed", "auto_dismissed"); err != nil {
157+
return utils.NewToolResultError(err.Error()), nil, nil
158+
}
156159
severity, err := OptionalParam[string](args, "severity")
157160
if err != nil {
158161
return utils.NewToolResultError(err.Error()), nil, nil
159162
}
163+
if err := validateEnumParam("severity", severity, "low", "medium", "high", "critical"); err != nil {
164+
return utils.NewToolResultError(err.Error()), nil, nil
165+
}
160166

161167
pagination, err := OptionalCursorPaginationParams(args)
162168
if err != nil {
@@ -194,7 +200,7 @@ func ListDependabotAlerts(t translations.TranslationHelperFunc) inventory.Server
194200
}
195201

196202
response := map[string]any{
197-
"alerts": alerts,
203+
"alerts": sanitizedDependabotAlertsCopy(alerts),
198204
"pageInfo": buildPageInfo(resp),
199205
}
200206

‎pkg/github/dependabot_test.go‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -256,6 +256,28 @@ func Test_ListDependabotAlerts(t *testing.T) {
256256
expectedAlerts: []*github.DependabotAlert{&criticalAlert},
257257
expectedNextCursor: "nextcursor123",
258258
},
259+
{
260+
name: "invalid state is rejected before request",
261+
mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{}),
262+
requestArgs: map[string]any{
263+
"owner": "owner",
264+
"repo": "repo",
265+
"state": "archived",
266+
},
267+
expectError: true,
268+
expectedErrMsg: "state must be one of: open, fixed, dismissed, auto_dismissed",
269+
},
270+
{
271+
name: "invalid severity is rejected before request",
272+
mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{}),
273+
requestArgs: map[string]any{
274+
"owner": "owner",
275+
"repo": "repo",
276+
"severity": "extreme",
277+
},
278+
expectError: true,
279+
expectedErrMsg: "severity must be one of: low, medium, high, critical",
280+
},
259281
{
260282
name: "alerts listing fails",
261283
mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{

0 commit comments

Comments
 (0)