migrate: map custom middleware context keys #176
Conversation
WalkthroughReworks the v3 middleware Locals migration to dynamically discover and map context keys from config usages, replacing string-based c.Locals(...) with FromContext accessors. Adds support for csrf_handler and token, narrows context key removal, and updates unit/integration tests to verify migration of custom CSRF ContextKey patterns. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor Dev as Developer
participant Mig as v3 Migrator
participant FS as Source Files
participant Map as ctxMap
participant Repl as Replacer
Dev->>Mig: Run migration
Mig->>FS: Scan Go files
loop Extract config usages
Mig->>FS: Find Config{... ContextKey: "..."}
Mig->>Map: Add key -> FromContext template
end
Mig->>Repl: Apply reLocals replacer
alt Key known
Repl->>FS: Replace c.Locals("k") -> FromContext(c)
else Key unknown
Repl->>FS: Leave c.Locals("k") unchanged
end
Mig->>FS: Cleanup type asserts / normalize chains
Mig-->>Dev: Report success/errors
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Poem
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @ReneWerner87, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request enhances the migration tool for a Go web framework (likely Fiber) to correctly handle custom context keys used by middleware. Previously, the migration might have missed c.Locals("key") calls where "key" was a custom value configured in the middleware. This change introduces dynamic detection of these custom keys from middleware configurations, ensuring that c.Locals("custom_key") is properly transformed into the corresponding package.FromContext(c) call during the migration process.
Highlights
- Improved Middleware Context Key Migration: The migration logic in middleware_locals.go has been significantly upgraded to dynamically identify and map custom context keys defined within middleware configurations (e.g., csrf.Config{ContextKey: "my_token"}). This ensures that c.Locals("my_token") calls are correctly transformed into csrf.TokenFromContext(c) during the migration.
- Enhanced Context Key Removal: The migration now also removes ContextUsername and ContextPassword fields from middleware configurations, in addition to ContextKey, streamlining the configuration after migration.
- Comprehensive Test Coverage for Custom Keys: A new dedicated unit test (Test_MigrateMiddlewareLocals_CustomContextKey) has been added to specifically verify the correct migration of custom CSRF context keys. Furthermore, the main V3 migration integration test (Test_Migrate_V2_to_V3) has been updated to include an example of a custom CSRF context key, ensuring end-to-end validation of this new capability.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request extends the v3 migration script to handle custom middleware context keys, a necessary feature for users with customized setups. The approach of dynamically mapping keys is sound, and it's good to see it covered by new tests.
My main concern is the fragility of the regular expression used to parse middleware configurations from source files. It's likely to fail on more complex but valid Go code, such as configs containing function literals, which could lead to an incomplete migration. I've added a comment with details and a suggestion for a more robust, AST-based parsing approach.
| content = r.re.ReplaceAllString(content, r.repl) | ||
|
|
||
| for _, e := range extractors { | ||
| re := regexp.MustCompile(e.pkg + `\.Config{[^}]*` + e.field + `:\s*"([^"]+)"`) |
There was a problem hiding this comment.
The regular expression used here is not robust enough for parsing Go source code. The [^}]* pattern will fail to match correctly if the Config struct literal contains any nested {} blocks, which is common with function literals (e.g., for a SuccessHandler field). This will cause the migration to silently miss custom context keys in such configurations.
For example, the regex would fail on this valid code:
_ = csrf.New(csrf.Config{
SuccessHandler: func(c *fiber.Ctx) error {
return c.Next()
},
ContextKey: "my_csrf_token",
})Using regular expressions to parse source code is error-prone. A much more reliable approach would be to use the go/parser and go/ast packages to build an Abstract Syntax Tree and inspect it. While more involved, it's the correct way to handle this and would make the migration tool significantly more robust.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (7)
cmd/internal/migrations/v3/middleware_locals.go (3)
25-35: Handle backtick-quoted ContextKey values (Go raw strings)The extractor only matches double-quoted values. Go configs occasionally use raw-string literals (
ContextKey:token``). We should support both to avoid missing mappings and leaving Locals(...) un-migrated.Apply:
@@ - for _, e := range extractors { - re := regexp.MustCompile(e.pkg + `\.Config{[^}]*` + e.field + `:\s*"([^"]+)"`) - matches := re.FindAllStringSubmatch(content, -1) - for _, m := range matches { - ctxMap[m[1]] = e.replFmt - } - } + for _, e := range extractors { + // Match double-quoted value + reDQ := regexp.MustCompile(e.pkg + `\.Config{[^}]*` + e.field + `:\s*"([^"]+)"`) + for _, m := range reDQ.FindAllStringSubmatch(content, -1) { + ctxMap[m[1]] = e.replFmt + } + // Match backtick-quoted value + reBT := regexp.MustCompile(e.pkg + `\.Config{[^}]*` + e.field + ":\s*`([^`]+)`") + for _, m := range reBT.FindAllStringSubmatch(content, -1) { + ctxMap[m[1]] = e.replFmt + } + }Also applies to: 37-43
45-55: Make Locals(...) matcher more robust (dotted receivers, backticks)Current
(\w+)\.Locals\("..."\)misses common patterns likeh.c.Locals("k")and backtick-quoted keys (Locals(token)). Also, replacing inside string literals/comments remains a risk with regex-only rewriting, but widening real-code coverage is still beneficial.Apply:
- reLocals := regexp.MustCompile(`(\w+)\.Locals\("([^"]+)"\)`) - content = reLocals.ReplaceAllStringFunc(content, func(s string) string { - sub := reLocals.FindStringSubmatch(s) + reLocalsDQ := regexp.MustCompile(`([\w\.]+)\.Locals\("([^"]+)"\)`) + reLocalsBT := regexp.MustCompile(`([\w\.]+)\.Locals\(` + "`" + `([^` + "`" + `]+)` + "`" + `\)`) + + replace := func(re *regexp.Regexp, s string) string { + sub := re.FindStringSubmatch(s) ctx := sub[1] key := sub[2] if fmtStr, ok := ctxMap[key]; ok { return fmt.Sprintf(fmtStr, ctx) } return s - }) + } + content = reLocalsDQ.ReplaceAllStringFunc(content, func(s string) string { return replace(reLocalsDQ, s) }) + content = reLocalsBT.ReplaceAllStringFunc(content, func(s string) string { return replace(reLocalsBT, s) })Note: Long-term, consider an AST-based pass to avoid touching string/comment contents at all. For this PR, the regex improvement is sufficient.
62-63: Context field scrubber OK; consider trailing-comma normalizationThe
reCtxKeyremoval is scoped and good. When removing the last field in single-line literals, it may leaveConfig{}orConfig{ }which compiles, but in tightly formatted single-line structs a leftover comma from the previous field can produceConfig{Prev: x,}(still valid in Go). If you want deterministic formatting for golden tests, optionally rungofmt -w(or rely on existing code that does). No code change required, just a note.cmd/internal/migrations/v3/middleware_locals_test.go (2)
73-101: Solid custom-key coverage; add a mixed-middleware collision testThis test verifies csrf with a custom "token" key correctly prefers csrf.TokenFromContext. To guard against future regressions in precedence when both keyauth and csrf define ContextKey "token", add a test that includes both configs and verifies Locals("token") usages under csrf and keyauth code paths are disambiguated or, at minimum, that we pick a deterministic owner and document it.
I can draft a test like “Test_MigrateMiddlewareLocals_TokenCollision” that seeds both:
- keyauth.New(keyauth.Config{ContextKey: "token"})
- csrf.New(csrf.Config{ContextKey: "token"})
and exercises two handlers, one intended for each middleware. Do you want me to open a follow-up PR with that test?
98-101: Strengthen assertionsConsider asserting that no
c.Locals(remains for "token" to ensure full replacement and preventing partial migrations.Apply:
- assert.Contains(t, content, `token := csrf.TokenFromContext(c)`) + assert.Contains(t, content, `token := csrf.TokenFromContext(c)`) + assert.NotContains(t, content, `c.Locals("token")`)cmd/migrate_test.go (2)
54-57: Good end-to-end check; consider idempotency on token replacementYou already re-run the migrator to check idempotency overall. Consider adding a focused assertion that “csrf.TokenFromContext(c)” appears only once and “c.Locals("token")” never reappears after the second pass, to guard against accidental double-rewrites.
Example:
@@ content2 := readFileTB(t, filepath.Join(dir, "main.go")) assert.Equal(t, content, content2) + at.NotContains(content2, `c.Locals("token")`)Also applies to: 99-101
42-43: Assert CSRF import path migration to contrib/v3In Fiber v3 the CSRF middleware has been moved out of the core and into the contrib module with a
/v3suffix. Update the migration test incmd/migrate_test.goto both modify the import path and assert that it was correctly rewritten.• cmd/migrate_test.go (lines 42–43 and 66–67):
- Change the import from
"github.com/gofiber/fiber/v2/middleware/csrf"
to
"github.com/gofiber/contrib/csrf/v3"
• Immediately after the existingTokenFromContextassertion, add:at.Contains(content, "\"github.com/gofiber/contrib/csrf/v3\"")Example diff:
--- a/cmd/migrate_test.go +++ b/cmd/migrate_test.go @@ -42,2 +42,2 @@ import ( - "github.com/gofiber/fiber/v2/middleware/csrf" + "github.com/gofiber/contrib/csrf/v3" @@ # after migration assertions - at.Contains(content, "csrf.TokenFromContext(c)") + at.Contains(content, "csrf.TokenFromContext(c)") + at.Contains(content, "\"github.com/gofiber/contrib/csrf/v3\"")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (3)
cmd/internal/migrations/v3/middleware_locals.go(1 hunks)cmd/internal/migrations/v3/middleware_locals_test.go(1 hunks)cmd/migrate_test.go(4 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2024-11-10T23:44:13.704Z
Learnt from: gaby
PR: gofiber/fiber#3193
File: middleware/adaptor/adaptor.go:111-111
Timestamp: 2024-11-10T23:44:13.704Z
Learning: In the `middleware/adaptor/adaptor.go` file of the Fiber framework, when updating context handling, replacing `c.Context()` with `c.RequestCtx()` is appropriate to access the `fasthttp.RequestCtx`.
Applied to files:
cmd/migrate_test.go
🧬 Code graph analysis (1)
cmd/internal/migrations/v3/middleware_locals_test.go (1)
cmd/internal/migrations/v3/middleware_locals.go (1)
MigrateMiddlewareLocals(13-76)
⏰ 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: Build (1.25.x, macos-13)
- GitHub Check: Build (1.25.x, windows-latest)
🔇 Additional comments (1)
cmd/internal/migrations/v3/middleware_locals.go (1)
15-24: csrf.HandlerFromContext does not exist in Fiber v3 – please remove or correct this mappingThe
csrfpackage in Fiber v3 only exposesTokenFromContext(aliasTokenFromCtx); there is noHandlerFromContextfunction. As a result, thecsrf_handlerentry in yourctxMapwill be invalid.• File
cmd/internal/migrations/v3/middleware_locals.go, lines 15–24:
- Remove the
"csrf_handler": "csrf.HandlerFromContext(%s)"entry, or if you only need the token, map"csrf"(or a new key) to"csrf.TokenFromContext(%s)".Example replacement:
ctxMap := map[string]string{ "requestid": "requestid.FromContext(%s)", "csrf": "csrf.TokenFromContext(%s)", - "csrf_handler": "csrf.HandlerFromContext(%s)", "session": "session.FromContext(%s)", "username": "basicauth.UsernameFromContext(%s)", "password": "basicauth.PasswordFromContext(%s)", "token": "keyauth.TokenFromContext(%s)", }⛔ Skipped due to learnings
Learnt from: gaby PR: gofiber/fiber#3193 File: middleware/adaptor/adaptor.go:111-111 Timestamp: 2024-11-10T23:44:13.704Z Learning: In the `middleware/adaptor/adaptor.go` file of the Fiber framework, when updating context handling, replacing `c.Context()` with `c.RequestCtx()` is appropriate to access the `fasthttp.RequestCtx`.Learnt from: gaby PR: gofiber/fiber#3193 File: middleware/cache/cache_test.go:897-897 Timestamp: 2024-11-08T04:10:42.990Z Learning: In the Fiber framework, `Context()` is being renamed to `RequestCtx()`, and `UserContext()` to `Context()` to improve clarity and align with Go's context conventions.
| ctxMap := map[string]string{ | ||
| "requestid": "requestid.FromContext(%s)", | ||
| "csrf": "csrf.TokenFromContext(%s)", | ||
| "csrf_handler": "csrf.HandlerFromContext(%s)", | ||
| "session": "session.FromContext(%s)", | ||
| "username": "basicauth.UsernameFromContext(%s)", | ||
| "password": "basicauth.PasswordFromContext(%s)", | ||
| "token": "keyauth.TokenFromContext(%s)", | ||
| } | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Default "token" mapping can collide with extractor-derived mappings
Seeding ctxMap with "token" -> keyauth.TokenFromContext(%s) creates a global mapping that may be overwritten (or overwrite) extractor-derived mappings when multiple middlewares reuse "token" (e.g., both keyauth and csrf set ContextKey to "token"). Today, precedence is purely order-of-extractor, which can silently misroute replacements across the file.
Recommend: remove the eager default for "token" and add it as a fallback only if no extractor populated that key. This preserves the common keyauth case while avoiding accidental override when another middleware purposely uses "token".
Apply:
@@
- ctxMap := map[string]string{
+ ctxMap := map[string]string{
"requestid": "requestid.FromContext(%s)",
"csrf": "csrf.TokenFromContext(%s)",
"csrf_handler": "csrf.HandlerFromContext(%s)",
"session": "session.FromContext(%s)",
"username": "basicauth.UsernameFromContext(%s)",
"password": "basicauth.PasswordFromContext(%s)",
- "token": "keyauth.TokenFromContext(%s)",
}
@@
- for _, e := range extractors {
+ for _, e := range extractors {
re := regexp.MustCompile(e.pkg + `\.Config{[^}]*` + e.field + `:\s*"([^"]+)"`)
matches := re.FindAllStringSubmatch(content, -1)
for _, m := range matches {
ctxMap[m[1]] = e.replFmt
}
}
+
+ // Provide a sensible default mapping for token -> keyauth only if no extractor populated it.
+ if _, ok := ctxMap["token"]; !ok {
+ ctxMap["token"] = "keyauth.TokenFromContext(%s)"
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| ctxMap := map[string]string{ | |
| "requestid": "requestid.FromContext(%s)", | |
| "csrf": "csrf.TokenFromContext(%s)", | |
| "csrf_handler": "csrf.HandlerFromContext(%s)", | |
| "session": "session.FromContext(%s)", | |
| "username": "basicauth.UsernameFromContext(%s)", | |
| "password": "basicauth.PasswordFromContext(%s)", | |
| "token": "keyauth.TokenFromContext(%s)", | |
| } | |
| ctxMap := map[string]string{ | |
| "requestid": "requestid.FromContext(%s)", | |
| "csrf": "csrf.TokenFromContext(%s)", | |
| "csrf_handler": "csrf.HandlerFromContext(%s)", | |
| "session": "session.FromContext(%s)", | |
| "username": "basicauth.UsernameFromContext(%s)", | |
| "password": "basicauth.PasswordFromContext(%s)", | |
| } | |
| for _, e := range extractors { | |
| re := regexp.MustCompile(e.pkg + `\.Config{[^}]*` + e.field + `:\s*"([^"]+)"`) | |
| matches := re.FindAllStringSubmatch(content, -1) | |
| for _, m := range matches { | |
| ctxMap[m[1]] = e.replFmt | |
| } | |
| } | |
| // Provide a sensible default mapping for token -> keyauth only if no extractor populated it. | |
| if _, ok := ctxMap["token"]; !ok { | |
| ctxMap["token"] = "keyauth.TokenFromContext(%s)" | |
| } |
🤖 Prompt for AI Agents
In cmd/internal/migrations/v3/middleware_locals.go around lines 15 to 24, the
ctxMap pre-populates "token" -> keyauth.TokenFromContext(%s) which can collide
with extractor-derived mappings; remove the eager "token" entry from the initial
ctxMap, and instead add logic after extractor population that checks if "token"
is still unset and only then inserts the keyauth.TokenFromContext fallback;
ensure the new check runs after all extractors are processed so
extractor-provided "token" mappings are preserved and keyauth only provides a
default when no other mapping exists.
Summary
Testing
make lint(fails: process terminated due to timeout)make test(fails: process terminated, no tests run)https://chatgpt.com/codex/tasks/task_e_68aa156ba1a0832688109c83b130b7f4
Summary by CodeRabbit
New Features
Bug Fixes
Tests