feat: query usage of token - #1161
Conversation
|
@sourcery-ai review |
WalkthroughAdded GET /api/usage/token endpoint with CriticalRateLimit and TokenAuth middleware. Implemented controller.GetTokenUsage to parse Bearer Authorization header, validate and trim Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Client
participant Router as API Router (/api/usage)
participant RateLimit as Middleware: CriticalRateLimit
participant TokenAuth as Middleware: TokenAuth
participant Ctrl as Controller: GetTokenUsage
participant Model
Client->>Router: GET /api/usage/token (Authorization: Bearer sk-...)
Router->>RateLimit: Apply rate limit
RateLimit->>TokenAuth: Forward to token auth
TokenAuth->>Ctrl: Authenticated
Ctrl->>Ctrl: Read Authorization header
alt Missing header
Ctrl-->>Client: 401 "No Authorization header"
else Invalid Bearer
Ctrl-->>Client: 401 "Invalid Bearer token"
else Valid Bearer
Ctrl->>Ctrl: tokenKey = TrimPrefix(header, "sk-")
Ctrl->>Model: GetTokenByKey(tokenKey, false)
alt Model error
Ctrl-->>Client: 200 { "success": false, "message": error }
else Model success
Ctrl->>Ctrl: expires_at = (ExpiredTime == -1 ? 0 : ExpiredTime)
Ctrl-->>Client: 200 { "code": true, "message":"ok", "data": { token_usage fields... } }
end
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ 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 (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
router/api-router.go (1)
142-142: Clarify auth contract for/api/token/usageendpointIt looks like
UserAuth()(which callsauthHelper) both reads and overwrites theAuthorizationheader to validate an API key before calling your handler, whileGetTokenUsageitself parses that same header to identify the target token. This overlap will lead to inconsistent header values or unexpected 401s.Please choose one of the following approaches to keep the contract unambiguous:
• Expose
GET /api/token/usagewithoutUserAuth()and let the handler authenticate purely via the Bearer token it receives.
– Remove or narrow the middleware for this route inrouter/api-router.go:142.
– e.g.
```diff
func setupTokenRoutes(router *gin.Engine, controller *TokenController) {
tokenRoute := router.Group("/api/token", middleware.UserAuth())
tokenRoute := router.Group("/api/token") { tokenRoute.GET("", controller.CreateToken)
tokenRoute.GET("/usage", controller.GetTokenUsage)
tokenRoute.GET("/usage", controller.GetTokenUsage) // auth via Bearer token only }• Keep
UserAuth()in place and change the handler to receive the token key via a separate field (query param or custom header) instead ofAuthorization.
– UpdateGetTokenUsageto read from e.g.?key=orX-Token(incontroller.GetTokenUsage).
– Document the new parameter in your API spec so clients know where to pass the token key.Locations to update:
router/api-router.go:142— wheretokenRoute.GET("/usage", …)is wired underUserAuth().middleware/auth.goaround lines 193–216 — whereauthHelpermutates and then re-readsAuthorizationfor API key validation.controller/token.go (3)
86-95: Revisit requiring Authorization inside a UserAuth-protected routeBecause this route already sits behind UserAuth(), returning 401 when Authorization is not present conflates user auth with the token lookup parameter. Prefer:
- Keep UserAuth(), accept the token key via query/body or a non-Authorization header (e.g., X-Token), and treat absence as 400 Bad Request; or
- Remove UserAuth() for this route and authenticate solely via the Bearer token.
This avoids double use of Authorization and makes the contract clear.
Would you like a patch to switch to query/header param (key) and update the router accordingly?
96-104: Harden Bearer parsingSplit on arbitrary whitespace and compare case-insensitively to reduce edge-case 401s.
Apply:
- parts := strings.Split(authHeader, " ") - if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" { + parts := strings.Fields(authHeader) + if len(parts) != 2 || !strings.EqualFold(parts[0], "Bearer") { c.JSON(http.StatusUnauthorized, gin.H{ "success": false, "message": "Invalid Bearer token", }) return }
86-134: Add tests for success, invalid header, and cross-user accessPlease add handler tests that cover:
- 200 path for own token, verifying object fields and ms units for expires_at.
- 401 for malformed/missing Bearer header (if keeping Bearer contract).
- 403 when attempting to access another user's token by key.
I can scaffold these with gin's httptest if helpful.
📜 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 (2)
controller/token.go(2 hunks)router/api-router.go(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
controller/token.go (1)
model/token.go (1)
GetTokenByKey(144-166)
router/api-router.go (1)
controller/token.go (1)
GetTokenUsage(86-134)
🔇 Additional comments (1)
controller/token.go (1)
8-8: Import looks correctAdding strings is appropriate for header parsing below.
| token, err := model.GetTokenByKey(tokenKey, true) | ||
| if err != nil { | ||
| c.JSON(http.StatusOK, gin.H{ | ||
| "success": false, | ||
| "message": err.Error(), | ||
| }) | ||
| return | ||
| } |
There was a problem hiding this comment.
Authorization bypass across users — validate ownership and use cache path
Any authenticated user can currently supply any known token key and retrieve its usage. Enforce that the token belongs to the current user. While here, prefer the cache path for lookups.
Apply:
- token, err := model.GetTokenByKey(tokenKey, true)
+ // Try cache first to reduce DB load
+ token, err := model.GetTokenByKey(tokenKey, false)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}
+ // Ensure the token belongs to the current user
+ userId := c.GetInt("id")
+ if token.UserId != userId {
+ c.JSON(http.StatusForbidden, gin.H{
+ "success": false,
+ "message": "forbidden",
+ })
+ return
+ }📝 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.
| token, err := model.GetTokenByKey(tokenKey, true) | |
| if err != nil { | |
| c.JSON(http.StatusOK, gin.H{ | |
| "success": false, | |
| "message": err.Error(), | |
| }) | |
| return | |
| } | |
| // Try cache first to reduce DB load | |
| token, err := model.GetTokenByKey(tokenKey, false) | |
| if err != nil { | |
| c.JSON(http.StatusOK, gin.H{ | |
| "success": false, | |
| "message": err.Error(), | |
| }) | |
| return | |
| } | |
| // Ensure the token belongs to the current user | |
| userId := c.GetInt("id") | |
| if token.UserId != userId { | |
| c.JSON(http.StatusForbidden, gin.H{ | |
| "success": false, | |
| "message": "forbidden", | |
| }) | |
| return | |
| } |
🤖 Prompt for AI Agents
In controller/token.go around lines 106 to 113, the code currently returns token
info for any supplied token key; change the lookup to use the cache path (call
model.GetTokenByKey with the cache flag instead of the DB flag) and then verify
token ownership by comparing token.UserID to the authenticated user ID from the
request context (fetch current user ID from the Gin context/middleware). If the
lookup fails return the existing error response, and if the token exists but
token.UserID does not match the current user return a 403 JSON response
(success: false, message: "forbidden" or similar) and stop processing.
| c.JSON(http.StatusOK, gin.H{ | ||
| "code": true, | ||
| "message": "ok", | ||
| "data": gin.H{ | ||
| "object": "token_usage", | ||
| "id": token.Id, | ||
| "name": token.Name, | ||
| "total_granted": token.RemainQuota + token.UsedQuota, | ||
| "total_used": token.UsedQuota, | ||
| "total_available": token.RemainQuota, | ||
| "unlimited_quota": token.UnlimitedQuota, | ||
| "expires_at": expiredAt, | ||
| }, | ||
| }) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Align response envelope and expires_at units with existing endpoints
Other token endpoints use "success" and return expires_at in milliseconds (see GetTokenStatus). Keep this consistent to avoid client breakage.
Apply:
- c.JSON(http.StatusOK, gin.H{
- "code": true,
- "message": "ok",
- "data": gin.H{
+ c.JSON(http.StatusOK, gin.H{
+ "success": true,
+ "message": "",
+ "data": gin.H{
"object": "token_usage",
"id": token.Id,
"name": token.Name,
"total_granted": token.RemainQuota + token.UsedQuota,
"total_used": token.UsedQuota,
"total_available": token.RemainQuota,
"unlimited_quota": token.UnlimitedQuota,
- "expires_at": expiredAt,
+ "expires_at": expiredAt * 1000,
},
})📝 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.
| c.JSON(http.StatusOK, gin.H{ | |
| "code": true, | |
| "message": "ok", | |
| "data": gin.H{ | |
| "object": "token_usage", | |
| "id": token.Id, | |
| "name": token.Name, | |
| "total_granted": token.RemainQuota + token.UsedQuota, | |
| "total_used": token.UsedQuota, | |
| "total_available": token.RemainQuota, | |
| "unlimited_quota": token.UnlimitedQuota, | |
| "expires_at": expiredAt, | |
| }, | |
| }) | |
| c.JSON(http.StatusOK, gin.H{ | |
| "success": true, | |
| "message": "", | |
| "data": gin.H{ | |
| "object": "token_usage", | |
| "id": token.Id, | |
| "name": token.Name, | |
| "total_granted": token.RemainQuota + token.UsedQuota, | |
| "total_used": token.UsedQuota, | |
| "total_available": token.RemainQuota, | |
| "unlimited_quota": token.UnlimitedQuota, | |
| "expires_at": expiredAt * 1000, | |
| }, | |
| }) |
🤖 Prompt for AI Agents
In controller/token.go around lines 120-133, the response envelope and
expires_at units must match other token endpoints: replace the "code" field with
"success" (boolean) and return "expires_at" as milliseconds since epoch instead
of the current unit; compute expires_at via expiredAt.UnixMilli() (or
equivalent) and ensure the JSON value is an integer, leaving other fields
unchanged.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
controller/token.go (2)
86-113: Avoid auth re-parsing; enforce ownership via middleware contextThis handler reparses Authorization and lets any valid token query its own usage, but it duplicates TokenAuth and doesn’t bind the returned token to the authenticated user/context. Use
token_id/idfrom Gin context (set by TokenAuth) and fetch viaGetTokenByIdsto enforce ownership and cut an extra cache/DB hop. Also standardize error handling viacommon.ApiError.Apply:
-func GetTokenUsage(c *gin.Context) { - authHeader := c.GetHeader("Authorization") - if authHeader == "" { - c.JSON(http.StatusUnauthorized, gin.H{ - "success": false, - "message": "No Authorization header", - }) - return - } - - parts := strings.Split(authHeader, " ") - if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" { - c.JSON(http.StatusUnauthorized, gin.H{ - "success": false, - "message": "Invalid Bearer token", - }) - return - } - tokenKey := parts[1] - - token, err := model.GetTokenByKey(strings.TrimPrefix(tokenKey, "sk-"), false) - if err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": err.Error(), - }) - return - } - - expiredAt := token.ExpiredTime - if expiredAt == -1 { - expiredAt = 0 - } +func GetTokenUsage(c *gin.Context) { + // Rely on TokenAuth middleware for identity; avoid re-parsing Authorization + tokenId := c.GetInt("token_id") + userId := c.GetInt("id") + if tokenId == 0 { + c.JSON(http.StatusUnauthorized, gin.H{ + "success": false, + "message": "unauthorized", + }) + return + } + token, err := model.GetTokenByIds(tokenId, userId) + if err != nil { + common.ApiError(c, err) + return + } + expiredAt := token.ExpiredTime + if expiredAt == -1 { + expiredAt = 0 + }
120-134: Align response envelope and expires_at units with existing endpointsOther token endpoints use
"success"and returnexpires_atin ms (seeGetTokenStatus). Keep consistent and includeidfor clients.- c.JSON(http.StatusOK, gin.H{ - "code": true, - "message": "ok", - "data": gin.H{ - "object": "token_usage", - "name": token.Name, - "total_granted": token.RemainQuota + token.UsedQuota, - "total_used": token.UsedQuota, - "total_available": token.RemainQuota, - "unlimited_quota": token.UnlimitedQuota, - "model_limits": token.GetModelLimitsMap(), - "model_limits_enabled": token.ModelLimitsEnabled, - "expires_at": expiredAt, - }, - }) + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "data": gin.H{ + "object": "token_usage", + "id": token.Id, + "name": token.Name, + "total_granted": token.RemainQuota + token.UsedQuota, + "total_used": token.UsedQuota, + "total_available": token.RemainQuota, + "unlimited_quota": token.UnlimitedQuota, + "model_limits": token.GetModelLimitsMap(), + "model_limits_enabled": token.ModelLimitsEnabled, + "expires_at": expiredAt * 1000, + }, + })
🧹 Nitpick comments (3)
controller/token.go (3)
8-8: Remove unused import if adopting middleware-based authAfter switching to context-based auth,
stringsis unused. Drop it to satisfy the compiler.import ( "net/http" "one-api/common" "one-api/model" "strconv" - "strings" "github.com/gin-gonic/gin" )
96-104: If you keep header parsing, make it whitespace-robustUse
strings.Fieldsto tolerate extra spaces/tabs and compare scheme case-insensitively.- parts := strings.Split(authHeader, " ") - if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" { + parts := strings.Fields(authHeader) + if len(parts) != 2 || strings.EqualFold(parts[0], "Bearer") == false { c.JSON(http.StatusUnauthorized, gin.H{ "success": false, "message": "Invalid Bearer token", }) return }
106-113: Don’t echo raw DB/cache errors to clientsLeaking
err.Error()can expose internals. Prefercommon.ApiError(as in other handlers) or map to a generic message.- if err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": err.Error(), - }) - return - } + if err != nil { + common.ApiError(c, err) + return + }
📜 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 (2)
controller/token.go(2 hunks)router/api-router.go(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- router/api-router.go
🧰 Additional context used
🧬 Code graph analysis (1)
controller/token.go (1)
model/token.go (1)
GetTokenByKey(144-166)
| func GetTokenUsage(c *gin.Context) { | ||
| authHeader := c.GetHeader("Authorization") | ||
| if authHeader == "" { | ||
| c.JSON(http.StatusUnauthorized, gin.H{ | ||
| "success": false, | ||
| "message": "No Authorization header", | ||
| }) | ||
| return | ||
| } | ||
|
|
||
| parts := strings.Split(authHeader, " ") | ||
| if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" { | ||
| c.JSON(http.StatusUnauthorized, gin.H{ | ||
| "success": false, | ||
| "message": "Invalid Bearer token", | ||
| }) | ||
| return | ||
| } | ||
| tokenKey := parts[1] | ||
|
|
||
| token, err := model.GetTokenByKey(strings.TrimPrefix(tokenKey, "sk-"), false) | ||
| if err != nil { | ||
| c.JSON(http.StatusOK, gin.H{ | ||
| "success": false, | ||
| "message": err.Error(), | ||
| }) | ||
| return | ||
| } | ||
|
|
||
| expiredAt := token.ExpiredTime | ||
| if expiredAt == -1 { | ||
| expiredAt = 0 | ||
| } | ||
|
|
||
| c.JSON(http.StatusOK, gin.H{ | ||
| "code": true, | ||
| "message": "ok", | ||
| "data": gin.H{ | ||
| "object": "token_usage", | ||
| "name": token.Name, | ||
| "total_granted": token.RemainQuota + token.UsedQuota, | ||
| "total_used": token.UsedQuota, | ||
| "total_available": token.RemainQuota, | ||
| "unlimited_quota": token.UnlimitedQuota, | ||
| "model_limits": token.GetModelLimitsMap(), | ||
| "model_limits_enabled": token.ModelLimitsEnabled, | ||
| "expires_at": expiredAt, | ||
| }, | ||
| }) | ||
| } |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Add endpoint tests for contract and auth behavior
Cover: 401 no header (if header-parsing kept), 401 invalid scheme, happy path shape (envelope + ms for expires_at), and ownership enforcement (403) if you continue accepting arbitrary token keys.
Proposed verification script to ensure routing and consistency:
🏁 Script executed:
#!/bin/bash
# Confirm the route is protected by TokenAuth
rg -nC3 'GetTokenUsage' -- router | cat
rg -n 'TokenAuth\(' router | cat
# Confirm GetTokenStatus returns expires_at in ms for consistency reference
rg -nC2 'GetTokenStatus|expires_at.*\* *1000' controller | catLength of output: 1989
Add tests for GetTokenUsage and normalize expires_at units
Please address the following before merging:
- Ensure
expires_atinGetTokenUsageis returned in milliseconds, to matchGetTokenStatus(which doesexpiredAt * 1000incontroller/token.go:82) - Add endpoint tests covering:
- 401 when the
Authorizationheader is missing - 401 when the scheme is not
Bearer - 403 when a valid token key is used by a non-owner (ownership enforcement)
- 200 “happy path” response shape, including the JSON envelope and
expires_atin ms
- 401 when the
Files/locations to update:
- controller/token.go (normalize
expires_atunit) - tests for token endpoints (e.g. in
controller/token_test.goor equivalent)
feat: query usage of token
2 participants
Summary by CodeRabbit
New Features
Chores
Notes