Merge pull request #1358 from akto-api-security/custom_auth_header_au…

…th_check Custom auth header auth check
akto-api-security · Aug 14, 2024 · 568ce2a · 568ce2a
2 parents 80862ce + 60e5de0
commit 568ce2a
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 14 deletions.
diff --git a/apps/testing/src/main/java/com/akto/test_editor/auth/AuthValidator.java b/apps/testing/src/main/java/com/akto/test_editor/auth/AuthValidator.java
@@ -20,13 +20,13 @@
 
 public class AuthValidator {
 
-    public static boolean validate(Auth auth, RawApi rawApi, AuthMechanism authMechanism) {
+    public static boolean validate(Auth auth, RawApi rawApi, AuthMechanism authMechanism, List<CustomAuthType> customAuthTypes) {
 
         if (auth == null) {
             return true;
         }
 
-        List<String> headerKeys = getHeaders(auth, authMechanism);
+        List<String> headerKeys = getHeaders(auth, authMechanism, customAuthTypes);
 
         auth.setHeaders(headerKeys);
 
@@ -41,29 +41,33 @@ public static boolean validate(Auth auth, RawApi rawApi, AuthMechanism authMecha
         for (String header: headerKeys) {
             contains = headers.containsKey(header) || CookieTransformer.isKeyPresentInCookie(cookieList, header);
             res = auth.getAuthenticated() && contains;
-            if (!res) {
-                return res;
+            if (res) {
+                return true;
             }
         }
-        return true;
+        return false;
     }
 
-    public static List<String> getHeaders(Auth auth, AuthMechanism authMechanism) {
+    public static List<String> getHeaders(Auth auth, AuthMechanism authMechanism, List<CustomAuthType> customAuthTypes) {
 
         if (auth != null && auth.getHeaders() != null && auth.getHeaders().size() > 0) {
             return auth.getHeaders();
         }
 
         List<String> headerKeys = new ArrayList<>();
 
-        if (authMechanism == null || authMechanism.getAuthParams() == null || authMechanism.getAuthParams().size() == 0) {
-            return null;
+        if (authMechanism != null && authMechanism.getAuthParams() != null && authMechanism.getAuthParams().size() > 0) {
+            for (AuthParam authParam: authMechanism.getAuthParams()) {
+                String key = authParam.getKey();
+                if (key == null) continue;
+                headerKeys.add(key.toLowerCase());
+            }
         }
 
-        for (AuthParam authParam: authMechanism.getAuthParams()) {
-            String key = authParam.getKey();
-            if (key == null) continue;
-            headerKeys.add(key.toLowerCase());
+        if (customAuthTypes != null) {
+            for(CustomAuthType customAuthType: customAuthTypes) {
+                headerKeys.addAll(customAuthType.getHeaderKeys());
+            }
         }
 
         return headerKeys;

diff --git a/apps/testing/src/main/java/com/akto/testing/yaml_tests/YamlTestTemplate.java b/apps/testing/src/main/java/com/akto/testing/yaml_tests/YamlTestTemplate.java
@@ -92,14 +92,14 @@ public Set<String> requireConfig(){
     @Override
     public ValidationResult filter() {
         // loggerMaker.infoAndAddToDb("filter started" + logId, LogDb.TESTING);
-        List<String> authHeaders = AuthValidator.getHeaders(this.auth, this.authMechanism);
+        List<String> authHeaders = AuthValidator.getHeaders(this.auth, this.authMechanism, this.customAuthTypes);
         // loggerMaker.infoAndAddToDb("found authHeaders " + authHeaders + " " + logId, LogDb.TESTING);
         if (authHeaders != null && authHeaders.size() > 0) {
             this.varMap.put("auth_headers", authHeaders);
         }
         if (this.auth != null && this.auth.getAuthenticated() != null) {
             // loggerMaker.infoAndAddToDb("validating auth, authenticated value is " + this.auth.getAuthenticated() + " " + logId, LogDb.TESTING);
-            boolean validAuthHeaders = AuthValidator.validate(this.auth, this.rawApi, this.authMechanism);
+            boolean validAuthHeaders = AuthValidator.validate(this.auth, this.rawApi, this.authMechanism, this.customAuthTypes);
             if (!validAuthHeaders) {
                 ValidationResult validationResult = new ValidationResult(false, "No valid auth headers");
                 // loggerMaker.infoAndAddToDb("invalid auth, skipping filter " + logId, LogDb.TESTING);