add datasouces access in auth middleware (gofr-dev#738)

Author: Umang01-hash

docs/advanced-guide/http-authentication/page.md

@@ -43,7 +43,7 @@ Use `EnableBasicAuthWithFunc(validationFunc)` to implement your own validation l
 The `validationFunc` takes the username and password as arguments and returns true if valid, false otherwise.
 
 ```go
-func validateUser(username string, password string) bool {
+func validateUser(c *container.Container, username, password string) bool {
 	// Implement your credential validation logic here 
 	// This example uses hardcoded credentials for illustration only   
 	return username == "john" && password == "doe123" 
@@ -52,7 +52,7 @@ func validateUser(username string, password string) bool {
 func main() { 
 	app := gofr.New() 
 
-	app.EnableBasicAuthWithFunc(validateUser) 
+	app.EnableBasicAuthWithValidator(validateUser) 
 
 	app.GET("/secure-data", func(c *gofr.Context) (interface{}, error) { 
 		// Handle access to secure data 
@@ -102,7 +102,7 @@ func main() {
 ```go
 package main
 
-func apiKeyValidator(apiKey string) bool {
+func apiKeyValidator(c *container.Container, apiKey string) bool {
 	validKeys := []string{"f0e1dffd-0ff0-4ac8-92a3-22d44a1464e4", "d7e4b46e-5b04-47b2-836c-2c7c91250f40"}
 
 	return slices.Contains(validKeys, apiKey)
@@ -112,7 +112,7 @@ func main() {
 	// initialise gofr object
 	app := gofr.New()
 
-	app.EnableAPIKeyAuthWithFunc(apiKeyValidator)
+	app.EnableAPIKeyAuthWithValidator(apiKeyValidator)
 
 	app.GET("/customer", Customer)
 

pkg/gofr/gofr.go

@@ -347,16 +347,35 @@ func (a *App) EnableBasicAuth(credentials ...string) {
 	a.httpServer.router.Use(middleware.BasicAuthMiddleware(middleware.BasicAuthProvider{Users: users}))
 }
 
+// Deprecated: EnableBasicAuthWithFunc is deprecated and will be removed in future releases, users must use
+// EnableBasicAuthWithValidator as it has access to application datasources.
 func (a *App) EnableBasicAuthWithFunc(validateFunc func(username, password string) bool) {
-	a.httpServer.router.Use(middleware.BasicAuthMiddleware(middleware.BasicAuthProvider{ValidateFunc: validateFunc}))
+	a.httpServer.router.Use(middleware.BasicAuthMiddleware(middleware.BasicAuthProvider{ValidateFunc: validateFunc, Container: a.container}))
+}
+
+func (a *App) EnableBasicAuthWithValidator(validateFunc func(c *container.Container, username, password string) bool) {
+	a.httpServer.router.Use(middleware.BasicAuthMiddleware(middleware.BasicAuthProvider{
+		ValidateFuncWithDatasources: validateFunc, Container: a.container}))
 }
 
 func (a *App) EnableAPIKeyAuth(apiKeys ...string) {
-	a.httpServer.router.Use(middleware.APIKeyAuthMiddleware(nil, apiKeys...))
+	a.httpServer.router.Use(middleware.APIKeyAuthMiddleware(middleware.APIKeyAuthProvider{}, apiKeys...))
+}
+
+// Deprecated: EnableAPIKeyAuthWithFunc is deprecated and will be removed in future releases, users must use
+// EnableAPIKeyAuthWithValidator as it has access to application datasources.
+func (a *App) EnableAPIKeyAuthWithFunc(validateFunc func(apiKey string) bool) {
+	a.httpServer.router.Use(middleware.APIKeyAuthMiddleware(middleware.APIKeyAuthProvider{
+		ValidateFunc: validateFunc,
+		Container:    a.container,
+	}))
 }
 
-func (a *App) EnableAPIKeyAuthWithFunc(validator func(apiKey string) bool) {
-	a.httpServer.router.Use(middleware.APIKeyAuthMiddleware(validator))
+func (a *App) EnableAPIKeyAuthWithValidator(validateFunc func(c *container.Container, apiKey string) bool) {
+	a.httpServer.router.Use(middleware.APIKeyAuthMiddleware(middleware.APIKeyAuthProvider{
+		ValidateFuncWithDatasources: validateFunc,
+		Container:                   a.container,
+	}))
 }
 
 func (a *App) EnableOAuth(jwksEndpoint string, refreshInterval int) {

pkg/gofr/gofr_test.go

@@ -434,6 +434,54 @@ func Test_UseMiddleware(t *testing.T) {
 	assert.Equal(t, "applied", testHeaderValue, "Test_UseMiddleware Failed! header value mismatch.")
 }
 
+func Test_APIKeyAuthMiddleware(t *testing.T) {
+	c, _ := container.NewMockContainer(t)
+
+	app := &App{
+		httpServer: &httpServer{
+			router: gofrHTTP.NewRouter(),
+			port:   8001,
+		},
+		container: c,
+		Config:    config.NewMockConfig(map[string]string{"REQUEST_TIMEOUT": "5"}),
+	}
+
+	apiKeys := []string{"test-key"}
+	validateFunc := func(_ *container.Container, apiKey string) bool {
+		return apiKey == "test-key"
+	}
+
+	// Registering APIKey middleware with and without custom validator
+	app.EnableAPIKeyAuth(apiKeys...)
+	app.EnableAPIKeyAuthWithValidator(validateFunc)
+
+	app.GET("/test", func(_ *Context) (interface{}, error) {
+		return "success", nil
+	})
+
+	go app.Run()
+	time.Sleep(1 * time.Second)
+
+	var netClient = &http.Client{
+		Timeout: time.Second * 10,
+	}
+
+	req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet,
+		"http://localhost:8001/test", http.NoBody)
+	req.Header.Set("X-API-Key", "test-key")
+
+	// Send the request and check for successful response
+	resp, err := netClient.Do(req)
+	if err != nil {
+		t.Errorf("error while making HTTP request in Test_APIKeyAuthMiddleware. err: %v", err)
+		return
+	}
+
+	defer resp.Body.Close()
+
+	assert.Equal(t, http.StatusOK, resp.StatusCode, "Test_APIKeyAuthMiddleware Failed!")
+}
+
 func Test_SwaggerEndpoints(t *testing.T) {
 	// Create the openapi.json file within the static directory
 	openAPIFilePath := filepath.Join("static", OpenAPIJSON)

pkg/gofr/http/middleware/apikey_auth.go

@@ -4,11 +4,20 @@ package middleware
 
 import (
 	"net/http"
+
+	"gofr.dev/pkg/gofr/container"
 )
 
+// APIKeyAuthProvider represents a basic authentication provider.
+type APIKeyAuthProvider struct {
+	ValidateFunc                func(apiKey string) bool
+	ValidateFuncWithDatasources func(c *container.Container, apiKey string) bool
+	Container                   *container.Container
+}
+
 // APIKeyAuthMiddleware creates a middleware function that enforces API key authentication based on the provided API
 // keys or a validation function.
-func APIKeyAuthMiddleware(validator func(apiKey string) bool, apiKeys ...string) func(handler http.Handler) http.Handler {
+func APIKeyAuthMiddleware(a APIKeyAuthProvider, apiKeys ...string) func(handler http.Handler) http.Handler {
 	return func(handler http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			if isWellKnown(r.URL.Path) {
@@ -22,7 +31,7 @@ func APIKeyAuthMiddleware(validator func(apiKey string) bool, apiKeys ...string)
 				return
 			}
 
-			if !validateKey(validator, authKey, apiKeys...) {
+			if !validateKey(a, authKey, apiKeys...) {
 				http.Error(w, "Unauthorized: Invalid Authorization header", http.StatusUnauthorized)
 				return
 			}
@@ -42,15 +51,17 @@ func isPresent(authKey string, apiKeys ...string) bool {
 	return false
 }
 
-func validateKey(validator func(apiKey string) bool, authKey string, apiKeys ...string) bool {
-	if validator != nil {
-		if !validator(authKey) {
-			return false
-		}
-	} else {
-		if !isPresent(authKey, apiKeys...) {
-			return false
-		}
+func validateKey(provider APIKeyAuthProvider, authKey string, apiKeys ...string) bool {
+	if provider.ValidateFunc != nil && !provider.ValidateFunc(authKey) {
+		return false
+	}
+
+	if provider.ValidateFuncWithDatasources != nil && !provider.ValidateFuncWithDatasources(provider.Container, authKey) {
+		return false
+	}
+
+	if provider.ValidateFunc == nil && provider.ValidateFuncWithDatasources == nil {
+		return isPresent(authKey, apiKeys...)
 	}
 
 	return true

pkg/gofr/http/middleware/apikey_auth_test.go

@@ -7,6 +7,13 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+
+	"gofr.dev/pkg/gofr/container"
+)
+
+const (
+	validKey1 string = "valid-key-1"
+	validKey2 string = "valid-key-2"
 )
 
 func Test_ApiKeyAuthMiddleware(t *testing.T) {
@@ -15,7 +22,10 @@ func Test_ApiKeyAuthMiddleware(t *testing.T) {
 	})
 
 	validator := func(apiKey string) bool {
-		return apiKey == "valid-key"
+		return apiKey == validKey1
+	}
+	validatorWithDB := func(_ *container.Container, apiKey string) bool {
+		return apiKey == validKey2
 	}
 
 	req, err := http.NewRequestWithContext(context.Background(), "GET", "/", http.NoBody)
@@ -24,26 +34,34 @@ func Test_ApiKeyAuthMiddleware(t *testing.T) {
 	}
 
 	testCases := []struct {
-		desc         string
-		validator    func(apiKey string) bool
-		apiKey       string
-		responseCode int
-		responseBody string
+		desc                string
+		validatorFunc       func(akiKey string) bool
+		validatorFuncWithDB func(c *container.Container, apiKey string) bool
+		apiKey              string
+		responseCode        int
+		responseBody        string
 	}{
-		{"missing api-key", nil, "", 401, "Unauthorized: Authorization header missing\n"},
-		{"invalid api-key", nil, "invalid-key", 401, "Unauthorized: Invalid Authorization header\n"},
-		{"valid api-key", nil, "valid-key-1", 200, "Success"},
-		{"another valid api-key", nil, "valid-key-2", 200, "Success"},
-		{"custom validator valid key", validator, "valid-key", 200, "Success"},
-		{"custom validator in-valid key", validator, "invalid-key", 401, "Unauthorized: Invalid Authorization header\n"},
+		{"missing api-key", nil, nil, "", 401, "Unauthorized: Authorization header missing\n"},
+		{"invalid api-key", nil, nil, "invalid-key", 401, "Unauthorized: Invalid Authorization header\n"},
+		{"valid api-key", nil, nil, validKey1, 200, "Success"},
+		{"another valid api-key", nil, nil, validKey2, 200, "Success"},
+		{"custom validatorFunc valid key", validator, nil, validKey1, 200, "Success"},
+		{"custom validatorFuncWithDB valid key", nil, validatorWithDB, validKey2, 200, "Success"},
+		{"custom validatorFuncWithDB in-valid key", nil, validatorWithDB, "invalid-key", 401, "Unauthorized: Invalid Authorization header\n"},
 	}
 
 	for i, tc := range testCases {
 		rr := httptest.NewRecorder()
 
 		req.Header.Set("X-API-KEY", tc.apiKey)
 
-		wrappedHandler := APIKeyAuthMiddleware(tc.validator, "valid-key-1", "valid-key-2")(testHandler)
+		provider := APIKeyAuthProvider{
+			ValidateFunc:                tc.validatorFunc,
+			ValidateFuncWithDatasources: tc.validatorFuncWithDB,
+			Container:                   nil,
+		}
+
+		wrappedHandler := APIKeyAuthMiddleware(provider, validKey1, validKey2)(testHandler)
 		wrappedHandler.ServeHTTP(rr, req)
 
 		assert.Equal(t, tc.responseCode, rr.Code, "TEST[%d], Failed.\n%s", i, tc.desc)
@@ -60,7 +78,9 @@ func Test_ApiKeyAuthMiddleware_well_known(t *testing.T) {
 	req := httptest.NewRequest(http.MethodGet, "/.well-known/health-check", http.NoBody)
 	rr := httptest.NewRecorder()
 
-	wrappedHandler := APIKeyAuthMiddleware(nil)(testHandler)
+	provider := APIKeyAuthProvider{}
+
+	wrappedHandler := APIKeyAuthMiddleware(provider)(testHandler)
 	wrappedHandler.ServeHTTP(rr, req)
 
 	assert.Equal(t, 200, rr.Code, "TEST Failed.\n")

pkg/gofr/http/middleware/basic_auth.go

@@ -4,14 +4,18 @@ import (
 	"encoding/base64"
 	"net/http"
 	"strings"
+
+	"gofr.dev/pkg/gofr/container"
 )
 
 const credentialLength = 2
 
 // BasicAuthProvider represents a basic authentication provider.
 type BasicAuthProvider struct {
-	Users        map[string]string
-	ValidateFunc func(username, password string) bool
+	Users                       map[string]string
+	ValidateFunc                func(username, password string) bool
+	ValidateFuncWithDatasources func(c *container.Container, username, password string) bool
+	Container                   *container.Container
 }
 
 // BasicAuthMiddleware creates a middleware function that enforces basic authentication using the provided BasicAuthProvider.
@@ -58,15 +62,16 @@ func BasicAuthMiddleware(basicAuthProvider BasicAuthProvider) func(handler http.
 }
 
 func validateCredentials(provider BasicAuthProvider, credentials []string) bool {
-	if provider.ValidateFunc != nil {
-		if !provider.ValidateFunc(credentials[0], credentials[1]) {
-			return false
-		}
-	} else {
-		if storedPass, ok := provider.Users[credentials[0]]; !ok || storedPass != credentials[1] {
-			return false
-		}
+	if provider.ValidateFunc != nil && !provider.ValidateFunc(credentials[0], credentials[1]) {
+		return false
+	}
+
+	if provider.ValidateFuncWithDatasources != nil && !provider.ValidateFuncWithDatasources(provider.Container,
+		credentials[0], credentials[1]) {
+		return false
 	}
 
-	return true
+	storedPass, ok := provider.Users[credentials[0]]
+
+	return ok && storedPass == credentials[1]
 }

pkg/gofr/http/middleware/basic_auth_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"gofr.dev/pkg/gofr/container"
 )
 
 func TestBasicAuthMiddleware(t *testing.T) {
@@ -17,6 +18,14 @@ func TestBasicAuthMiddleware(t *testing.T) {
 		return false
 	}
 
+	validationFuncWithDB := func(_ *container.Container, user, pass string) bool {
+		if user == "abc" && pass == "pass123" {
+			return true
+		}
+
+		return false
+	}
+
 	testCases := []struct {
 		name               string
 		authHeader         string
@@ -32,7 +41,7 @@ func TestBasicAuthMiddleware(t *testing.T) {
 		{
 			name:               "Valid Authorization with validation Func",
 			authHeader:         "Basic YWJjOnBhc3MxMjM=",
-			authProvider:       BasicAuthProvider{ValidateFunc: validationFunc},
+			authProvider:       BasicAuthProvider{Users: map[string]string{"abc": "pass123"}, ValidateFunc: validationFunc},
 			expectedStatusCode: http.StatusOK,
 		},
 		{
@@ -41,6 +50,12 @@ func TestBasicAuthMiddleware(t *testing.T) {
 			authProvider:       BasicAuthProvider{ValidateFunc: validationFunc},
 			expectedStatusCode: http.StatusUnauthorized,
 		},
+		{
+			name:               "false from validation Func with DB",
+			authHeader:         "Basic dXNlcjpwYXNzd29yZA==",
+			authProvider:       BasicAuthProvider{ValidateFuncWithDatasources: validationFuncWithDB},
+			expectedStatusCode: http.StatusUnauthorized,
+		},
 		{
 			name:               "No Authorization Header",
 			authHeader:         "",