diff --git a/pkg/controllers/api-key/usage.go b/pkg/controllers/api-key/usage.go index 896372e..17b669f 100644 --- a/pkg/controllers/api-key/usage.go +++ b/pkg/controllers/api-key/usage.go @@ -13,10 +13,10 @@ import ( ) func Usage(c *gin.Context, deps dependencies.Dependencies) { - apiKey, _ := c.Get("api-key-model") + apiKeyData, _ := c.Get("api-key-model") // Apply model to apiKey - apiKeyModel, _ := apiKey.(models.APIKey) + apiKeyModel, _ := apiKeyData.(models.APIKey) c.JSON(http.StatusOK, gin.H{ "message": apiKeyModel, diff --git a/pkg/middleware/api-key/track_key_usage.go b/pkg/middleware/api-key/track_key_usage.go new file mode 100644 index 0000000..a3a611b --- /dev/null +++ b/pkg/middleware/api-key/track_key_usage.go @@ -0,0 +1,47 @@ +package middleware + +import ( + "net/http" + "time" + + "github.com/Frier03/KeyAuth-API/pkg/models" + "github.com/Frier03/KeyAuth-API/pkg/services" + "github.com/gin-gonic/gin" +) + +func TrackKeyUsage(badgerService *services.BadgerService) gin.HandlerFunc { + return func(c *gin.Context) { + apiKey := c.GetHeader("X-Api-Key") + apiKeyData, _ := c.Get("api-key-model") + + // Apply model to apiKeyData + apiKeyModel, _ := apiKeyData.(models.APIKey) + + // Check if limit on API key has exceeded + if apiKeyModel.Limit <= apiKeyModel.Usage { + c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{ + "error": "Rate limit exceeded", + }) + return + } + + // Increment the API key usage count + apiKeyModel.Usage++ + + // Update the API last used + apiKeyModel.LastUsed = time.Now() + + // Update the API key in the database + err := badgerService.PutAPIKey([]byte(apiKey), &apiKeyModel) + if err != nil { + c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{ + "error": "Failed to update API key", + }) + return + } + + // Update c set for api-key-model + c.Set("api-key-model", apiKeyModel) + c.Next() + } +} diff --git a/pkg/middleware/api_key_validation.go b/pkg/middleware/api-key/validate_key.go similarity index 84% rename from pkg/middleware/api_key_validation.go rename to pkg/middleware/api-key/validate_key.go index 673ff0a..d6cc938 100644 --- a/pkg/middleware/api_key_validation.go +++ b/pkg/middleware/api-key/validate_key.go @@ -9,8 +9,7 @@ import ( "github.com/gin-gonic/gin" ) -// APIKeyValidationMiddleware is a middleware to validate the API key -func APIKeyValidationMiddleware(badgerService *services.BadgerService) gin.HandlerFunc { +func ValidateKey(badgerService *services.BadgerService) gin.HandlerFunc { return func(c *gin.Context) { apiKey := c.GetHeader("X-Api-Key") @@ -49,7 +48,6 @@ func APIKeyValidationMiddleware(badgerService *services.BadgerService) gin.Handl return } - // Set custom header for the api key model retrieved c.Set("api-key-model", apiKeyModel) c.Next() } diff --git a/pkg/middleware/api_key_usage.go b/pkg/middleware/api_key_usage.go deleted file mode 100644 index 39aa5aa..0000000 --- a/pkg/middleware/api_key_usage.go +++ /dev/null @@ -1,11 +0,0 @@ -package middleware - -import ( - "github.com/Frier03/KeyAuth-API/pkg/services" - "github.com/gin-gonic/gin" -) - -func APIKeyUsageMiddleware(badgerService *services.BadgerService) gin.HandlerFunc { - return func(c *gin.Context) { - } -} diff --git a/pkg/middleware/model_validation.go b/pkg/middleware/validate_model.go similarity index 94% rename from pkg/middleware/model_validation.go rename to pkg/middleware/validate_model.go index eb95304..1d57ada 100644 --- a/pkg/middleware/model_validation.go +++ b/pkg/middleware/validate_model.go @@ -8,7 +8,7 @@ import ( "github.com/Frier03/KeyAuth-API/pkg/models" ) -func ValidateModelMiddleware(model interface{}) gin.HandlerFunc { +func ValidateModel(model interface{}) gin.HandlerFunc { return func(c *gin.Context) { if err := c.ShouldBindJSON(model); err != nil { c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{ diff --git a/pkg/routes/api_key.go b/pkg/routes/api_key.go index 0f61457..6ffda82 100644 --- a/pkg/routes/api_key.go +++ b/pkg/routes/api_key.go @@ -3,9 +3,10 @@ package routes import ( "github.com/gin-gonic/gin" - api "github.com/Frier03/KeyAuth-API/pkg/controllers/api-key" + apiController "github.com/Frier03/KeyAuth-API/pkg/controllers/api-key" "github.com/Frier03/KeyAuth-API/pkg/dependencies" - "github.com/Frier03/KeyAuth-API/pkg/middleware" + middleware "github.com/Frier03/KeyAuth-API/pkg/middleware" + apiMiddleware "github.com/Frier03/KeyAuth-API/pkg/middleware/api-key" "github.com/Frier03/KeyAuth-API/pkg/models" ) @@ -14,16 +15,17 @@ func SetupAPIKeyRoutes(r *gin.Engine, deps dependencies.Dependencies) { apiKeyRoutes := r.Group("/api-key") { apiKeyRoutes.POST("/generate", - middleware.ValidateModelMiddleware(&models.APIKeyGenerateRequest{}), + middleware.ValidateModel(&models.APIKeyGenerateRequest{}), func(c *gin.Context) { - api.Generate(c, deps) + apiController.Generate(c, deps) }, ) apiKeyRoutes.GET("/usage", - middleware.APIKeyValidationMiddleware(deps.BadgerService), + apiMiddleware.ValidateKey(deps.BadgerService), + apiMiddleware.TrackKeyUsage(deps.BadgerService), func(c *gin.Context) { - api.Usage(c, deps) + apiController.Usage(c, deps) }, ) } diff --git a/pkg/routes/auth.go b/pkg/routes/auth.go index 80d52bc..b7a646b 100644 --- a/pkg/routes/auth.go +++ b/pkg/routes/auth.go @@ -3,8 +3,8 @@ package routes import ( "github.com/gin-gonic/gin" - auth "github.com/Frier03/KeyAuth-API/pkg/controllers/auth" - "github.com/Frier03/KeyAuth-API/pkg/middleware" + authController "github.com/Frier03/KeyAuth-API/pkg/controllers/auth" + middleware "github.com/Frier03/KeyAuth-API/pkg/middleware" "github.com/Frier03/KeyAuth-API/pkg/models" ) @@ -13,17 +13,17 @@ func SetupAuthRoutes(r *gin.Engine) { authRoutes := r.Group("/auth") { authRoutes.POST("/login", - middleware.ValidateModelMiddleware(&models.LoginRequest{}), // Validate login request payload - auth.Login, // Expected resource logic + middleware.ValidateModel(&models.LoginRequest{}), // Validate login request payload + authController.Login, // Expected resource logic ) authRoutes.POST("/register", - middleware.ValidateModelMiddleware(&models.RegisterRequest{}), // Validate register request payload - auth.Register, // Expected resource logic + middleware.ValidateModel(&models.RegisterRequest{}), // Validate register request payload + authController.Register, // Expected resource logic ) authRoutes.POST("/logout", - auth.Logout, // Expected resource logic + authController.Logout, // Expected resource logic ) } }