From 302e2aa945d68bb4e03da2d39e408afbe0281d85 Mon Sep 17 00:00:00 2001
From: Savindu Dimal <dimalsavindu@gmail.com>
Date: Mon, 28 Oct 2024 15:14:31 +0530
Subject: [PATCH] Add custom policy replacement in key template

- Update validateCustomPolicy to check for the presence of CUSTOM_PROPERTY in messageContext
- Iterates through CUSTOM_PROPERTY Map to replace customProperty placeholders in the key template
---
 .../gateway/handlers/throttling/ThrottleHandler.java     | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/throttling/ThrottleHandler.java b/components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/throttling/ThrottleHandler.java
index b7e526726902..a8ec64b46e74 100644
--- a/components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/throttling/ThrottleHandler.java
+++ b/components/apimgt/org.wso2.carbon.apimgt.gateway/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/throttling/ThrottleHandler.java
@@ -1280,6 +1280,15 @@ public boolean validateCustomPolicy(String userID, String appKey, String resourc
                 if (clientIp != null) {
                     key = key.replaceAll("\\$clientIp", APIUtil.ipToBigInteger(clientIp).toString());
                 }
+                Object customPropertyObj = messageContext.getProperty(APIMgtGatewayConstants.CUSTOM_PROPERTY);
+                if (customPropertyObj != null) {
+                    Map<String, Object> customProperties = (Map<String, Object>) customPropertyObj;
+                    for (Map.Entry<String, Object> entry : customProperties.entrySet()) {
+                        String customKey = "\\$customProperty\\." + entry.getKey();
+                        String customValue = entry.getValue() != null ? entry.getValue().toString() : "";
+                        key = key.replaceAll(customKey, customValue);
+                    }
+                }
                 if (getThrottleDataHolder().isThrottled(key)) {
                     long timestamp = getThrottleDataHolder().getThrottleNextAccessTimestamp(key);
                     messageContext.setProperty(APIThrottleConstants.THROTTLED_NEXT_ACCESS_TIMESTAMP, timestamp);