[BB-1481] Improve trust policy parsing and add test

[mateovespConductor]

Replace map[string]any with typed structures for trust policy analysis and add coverage testing

Description

• Bug fix
• New feature

Useful links:

• Baton SDK coding guidelines
• New contributor guide

Summary by CodeRabbit

• New Features

  • Structured trust policy parsing and principal detection for more accurate IAM principal classification.

• Bug Fixes

  • Robust handling so a malformed trust policy no longer blocks processing of other roles; parsing failures are logged and skipped.

• Tests

  • Added extensive tests covering trust policy decoding, action/principal formats, principal resource detection, and principal extraction edge cases.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

Walkthrough

Refactors AWS trust policy parsing from ad-hoc map handling to structured types (TrustPolicy, Statement, Action, Principal) with custom JSON unmarshalling for single/array forms. Adds detectPrincipalResource to classify principal ARNs. Adjusts extractTrustPrincipals filtering and makes trust-policy parse failures non-fatal in Grants.

Changes

Cohort / File(s)	Summary
Trust policy types & helpers pkg/connector/helpers.go	Adds TrustPolicy, Statement, Action, and Principal types with UnmarshalJSON for single-or-array forms; replaces map-based parsing with typed unmarshalling; adds detectPrincipalResource and updates extractTrustPrincipals to filter by Effect == "Allow" and actions containing sts:AssumeRole.
Unit tests for trust parsing pkg/connector/helpers_test.go	Adds comprehensive tests for TrustPolicy.UnmarshalJSON, Action.UnmarshalJSON, Principal.UnmarshalJSON, detectPrincipalResource, and extractTrustPrincipals covering single/array forms, error cases, edge cases, and filtering behavior.
Grant error handling pkg/connector/role.go	Changes Grants behavior to log a warning and skip a role's grants when trust policy parsing fails (returns nils, no error) instead of propagating an error.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

• Review custom UnmarshalJSON implementations for correctness and edge-case handling (empty values, wrong JSON types).
• Verify extractTrustPrincipals correctly filters by Effect and sts:AssumeRole and correctly aggregates principals.
• Confirm detectPrincipalResource ARN parsing/classification handles paths and account-root/service/wildcard cases.
• Ensure changed Grants error-handling in role.go aligns with desired downstream behavior (non-fatal skip).

Possibly related PRs

• [BB-1481] Add grants to iam_roles #91 — Prior work on trust-policy parsing and principal detection; closely related to the refactor from ad-hoc helpers to typed parsing in this change.

Suggested reviewers

• laurenleach
• johnallers
• aldevv
• luisina-santos
• mateoHernandez123
• btipling

Poem

🐰
I hopped through JSON, neat and spry,
Turned scattered maps to types that fly.
Statements lined up, actions clear,
Principals found — no need to fear.
A joyful nibble on parsed policy pie 🥕✨

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 28.57% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main change: replacing map-based trust policy parsing with typed structures and adding corresponding tests.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch mateovesp/improve-trust-policy-parsing

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 071fb01 and d8d245e.

📒 Files selected for processing (2)

• pkg/connector/helpers.go (2 hunks)
• pkg/connector/helpers_test.go (1 hunks)

🧰 Additional context used

🧠 Learnings (2)

📓 Common learnings

Learnt from: mateovespConductor
Repo: ConductorOne/baton-aws PR: 91
File: pkg/connector/helpers.go:191-213
Timestamp: 2025-12-01T15:27:53.050Z
Learning: AWS IAM role trust policies strictly forbid wildcards ("*" or "sts:*") in the Action field. AWS enforces this at the API level with the error "MalformedPolicyDocument: AssumeRole policy may only specify STS AssumeRole actions". Trust policies must explicitly specify "sts:AssumeRole" in the Action field.

📚 Learning: 2025-12-01T15:27:53.050Z

Learnt from: mateovespConductor
Repo: ConductorOne/baton-aws PR: 91
File: pkg/connector/helpers.go:191-213
Timestamp: 2025-12-01T15:27:53.050Z
Learning: AWS IAM role trust policies strictly forbid wildcards ("*" or "sts:*") in the Action field. AWS enforces this at the API level with the error "MalformedPolicyDocument: AssumeRole policy may only specify STS AssumeRole actions". Trust policies must explicitly specify "sts:AssumeRole" in the Action field.

Applied to files:

• pkg/connector/helpers.go

🧬 Code graph analysis (1)

pkg/connector/helpers_test.go (1)

pkg/connector/helpers.go (4)

• TrustPolicy (198-201)
• Statement (238-242)
• Action (245-245)
• Principal (268-270)

🔇 Additional comments (8)

pkg/connector/helpers_test.go (1)

1-601: Excellent test coverage for the trust policy parsing refactor!

The test suite is comprehensive and well-structured, covering:

• All custom UnmarshalJSON implementations for TrustPolicy, Action, and Principal
• Happy paths (single/array forms, multiple statements)
• Error cases (invalid JSON, invalid types like numbers/booleans/objects)
• Edge cases (empty arrays, empty strings, URL encoding issues)
• Filtering logic (Deny statements, non-AssumeRole actions, service principals, empty strings)
• Principal detection (IAM users/roles with paths, excluding root/service/wildcards)

The table-driven approach with descriptive test names makes the test suite maintainable and easy to understand. The use of URL-encoded policy documents in ExtractTrustPrincipals tests accurately reflects real AWS API responses.

pkg/connector/helpers.go (7)

8-8: Good addition of slices import.

The slices package is used appropriately for slices.Contains on line 162, which is cleaner than a manual loop for checking action membership.

---

145-175: Well-executed refactoring of trust policy extraction!

The function now uses structured types (TrustPolicy, Statement, Action, Principal) instead of ad-hoc map handling, which significantly improves:

• Type safety and compile-time checking
• Code readability and maintainability
• Error handling clarity

The filtering logic correctly:

• Excludes Deny statements (line 158)
• Checks for "sts:AssumeRole" action using slices.Contains (line 162)
• Filters out empty principal strings (line 168)

Based on learnings, AWS enforces that trust policies must explicitly specify "sts:AssumeRole" in the Action field, and this implementation correctly validates that.

---

177-195: Clean implementation of principal classification!

The function correctly:

• Parses ARNs and handles errors gracefully
• Identifies IAM users (user/ prefix) and roles (role/ prefix)
• Returns false for other principal types (root, service principals, wildcards, invalid ARNs)

This aligns with the downstream usage in grant processing where root and service principals are intentionally excluded.

Note: The past review comment about the godot linter has been resolved—the comment now correctly ends with a period on line 178.

---

197-235: Excellent implementation of flexible Statement unmarshalling!

The custom UnmarshalJSON correctly:

• Uses the type alias pattern to avoid infinite recursion (lines 205-212)
• Captures Statement as json.RawMessage for conditional parsing
• Tries array form first, then falls back to single object (lines 220-231)
• Normalizes to array form for consistent downstream processing
• Provides clear error messages for invalid types

This handles AWS IAM's flexibility in allowing Statement to be either a single object or an array, making the parsing more robust.

---

237-242: LGTM!

The Statement struct appropriately uses custom types for Action and Principal to support their flexible JSON representations (string or array forms), while keeping Effect as a simple string.

---

244-264: Clean handling of Action's dual forms!

The implementation correctly:

• Defines Action as []string for uniform internal representation
• Tries string form first, then array (lines 249-260)
• Normalizes single string to array (line 262)
• Provides clear error message for invalid types

This matches AWS IAM's allowance for Action to be specified as either a string or an array, ensuring compatibility with various policy formats.

---

266-310: Sophisticated and correct Principal unmarshalling!

The implementation elegantly handles all Principal forms specified in AWS IAM policies:

• Object form with AWS field as string or array (lines 289-305)
• String form for wildcard or service principals—correctly ignored (lines 278-285)
• Object form with non-AWS fields (Service, Federated)—correctly ignored (lines 308-309)

The logic appropriately focuses on AWS principals (IAM users and roles) while gracefully handling and ignoring other principal types, which aligns with the connector's focus on IAM identity relationships.

The error messages clearly distinguish between structural errors and intentionally ignored cases.

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 60c2529 and 071fb01.

📒 Files selected for processing (3)

• pkg/connector/helpers.go (2 hunks)
• pkg/connector/helpers_test.go (1 hunks)
• pkg/connector/role.go (1 hunks)

🧰 Additional context used

🧠 Learnings (2)

📓 Common learnings

Learnt from: mateovespConductor
Repo: ConductorOne/baton-aws PR: 91
File: pkg/connector/helpers.go:191-213
Timestamp: 2025-12-01T15:27:53.050Z
Learning: AWS IAM role trust policies strictly forbid wildcards ("*" or "sts:*") in the Action field. AWS enforces this at the API level with the error "MalformedPolicyDocument: AssumeRole policy may only specify STS AssumeRole actions". Trust policies must explicitly specify "sts:AssumeRole" in the Action field.

📚 Learning: 2025-12-01T15:27:53.050Z

Learnt from: mateovespConductor
Repo: ConductorOne/baton-aws PR: 91
File: pkg/connector/helpers.go:191-213
Timestamp: 2025-12-01T15:27:53.050Z
Learning: AWS IAM role trust policies strictly forbid wildcards ("*" or "sts:*") in the Action field. AWS enforces this at the API level with the error "MalformedPolicyDocument: AssumeRole policy may only specify STS AssumeRole actions". Trust policies must explicitly specify "sts:AssumeRole" in the Action field.

Applied to files:

• pkg/connector/helpers.go

🧬 Code graph analysis (1)

pkg/connector/helpers_test.go (1)

pkg/connector/helpers.go (4)

• TrustPolicy (198-201)
• Statement (238-242)
• Action (245-245)
• Principal (268-270)

🪛 GitHub Check: go-lint

pkg/connector/helpers.go

[failure] 178-178:
Comment should end in a period (godot)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)

• GitHub Check: test
• GitHub Check: go-lint

🔇 Additional comments (6)

pkg/connector/role.go (1)

178-183: LGTM! Non-fatal error handling is a reasonable approach.

Logging the warning with role_name, role_arn, and the error provides good observability while allowing other roles to continue processing.

pkg/connector/helpers_test.go (1)

1-568: Comprehensive test coverage for the new structured types.

The tests thoroughly cover:

• Polymorphic JSON unmarshalling (single vs array forms)
• Error handling for invalid inputs
• Edge cases (empty arrays, empty strings, wildcards)
• Real-world URL-encoded policy documents

The inline comments documenting the decoded JSON for URL-encoded test data are helpful for maintainability.

pkg/connector/helpers.go (4)

145-175: Clean refactor using structured types.

The implementation correctly:

• Filters by Effect == "Allow" and Action containing "sts:AssumeRole"
• Uses slices.Contains for idiomatic action matching
• Filters out empty string principals

Based on learnings, AWS enforces that trust policies must explicitly specify "sts:AssumeRole" (wildcards are rejected at API level), so the exact match approach is appropriate.

---

203-235: LGTM! Correct use of type alias to avoid infinite recursion.

The pattern of using a type alias (Alias) when calling json.Unmarshal within a custom UnmarshalJSON is the correct approach to prevent stack overflow.

---

244-264: LGTM! Handles polymorphic Action field correctly.

The unmarshalling logic properly handles both "sts:AssumeRole" (string) and ["sts:AssumeRole", "sts:TagSession"] (array) forms.

---

266-310: LGTM! Properly handles AWS principal formats and ignores non-AWS principals.

The implementation correctly:

• Ignores wildcard "*" and service principals
• Handles both {"AWS": "arn:..."} and {"AWS": ["arn:...", "arn:..."]} forms
• Returns nil for non-AWS principals without error (correct behavior)