From a6b292422135bd5aae12071f15fcc2575fc03b32 Mon Sep 17 00:00:00 2001
From: Savindu Dimal <dimalsavindu@gmail.com>
Date: Fri, 19 Apr 2024 15:42:54 +0530
Subject: [PATCH] Add retry when fetching api definition security audit

---
 .../publisher/ApisApiServiceImplUtils.java    | 67 ++++++++++++-------
 1 file changed, 41 insertions(+), 26 deletions(-)

diff --git a/components/apimgt/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/restapi/publisher/ApisApiServiceImplUtils.java b/components/apimgt/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/restapi/publisher/ApisApiServiceImplUtils.java
index 76884d23412e..7a0464cf12e3 100644
--- a/components/apimgt/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/restapi/publisher/ApisApiServiceImplUtils.java
+++ b/components/apimgt/org.wso2.carbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/restapi/publisher/ApisApiServiceImplUtils.java
@@ -380,37 +380,52 @@ public static JSONObject getAuditReport(API api, JSONObject securityAuditPropert
             httpGet.setHeader(APIConstants.HEADER_API_TOKEN, apiToken);
             httpGet.setHeader(APIConstants.HEADER_USER_AGENT, APIConstants.USER_AGENT_APIM);
             // Code block for the processing of the response
-            try (CloseableHttpResponse response = getHttpClient.execute(httpGet)) {
-                if (isDebugEnabled) {
-                    log.debug(HTTP_STATUS_LOG + response.getStatusLine().getStatusCode());
-                }
-                if (response.getStatusLine().getStatusCode() == HttpStatus.SC_OK) {
-                    StringBuilder responseString = new StringBuilder();
-                    try (BufferedReader reader = new BufferedReader(
-                            new InputStreamReader(response.getEntity().getContent(), StandardCharsets.UTF_8))) {
-                        String inputLine;
-                        while ((inputLine = reader.readLine()) != null) {
-                            responseString.append(inputLine);
+            int retryCount = 0;
+            while (true) {
+                try (CloseableHttpResponse response = getHttpClient.execute(httpGet)) {
+                    if (isDebugEnabled) {
+                        log.debug(HTTP_STATUS_LOG + response.getStatusLine().getStatusCode());
+                    }
+                    if (response.getStatusLine().getStatusCode() == HttpStatus.SC_OK) {
+                        StringBuilder responseString = new StringBuilder();
+                        try (BufferedReader reader = new BufferedReader(
+                                new InputStreamReader(response.getEntity().getContent(), StandardCharsets.UTF_8))) {
+                            String inputLine;
+                            while ((inputLine = reader.readLine()) != null) {
+                                responseString.append(inputLine);
+                            }
+                        }
+                        JSONObject responseJson = (JSONObject) new JSONParser().parse(responseString.toString());
+                        String report = responseJson.get(APIConstants.DATA).toString();
+                        String grade = (String) ((JSONObject) ((JSONObject) responseJson.get(APIConstants.ATTR))
+                                .get(APIConstants.DATA)).get(APIConstants.GRADE);
+                        Integer numErrors = Integer.valueOf(
+                                (String) ((JSONObject) ((JSONObject) responseJson.get(APIConstants.ATTR))
+                                        .get(APIConstants.DATA)).get(APIConstants.NUM_ERRORS));
+                        String decodedReport = new String(Base64Utils.decode(report), StandardCharsets.UTF_8);
+                        JSONObject output = new JSONObject();
+                        output.put("decodedReport", decodedReport);
+                        output.put("grade", grade);
+                        output.put("numErrors", numErrors);
+                        output.put("auditUuid", auditUuid);
+                        return output;
+                    } else {
+                        retryCount++;
+                        if (retryCount <= 5) {
+                            try {
+                                Thread.sleep(1000);
+                            } catch (InterruptedException e) {
+                                // Ignore
+                            }
+                        } else {
+                            throw new APIManagementException(
+                                    "Error while retrieving data from the API Definition Security Audit. "
+                                            + "Found http status " + response.getStatusLine());
                         }
                     }
-                    JSONObject responseJson = (JSONObject) new JSONParser().parse(responseString.toString());
-                    String report = responseJson.get(APIConstants.DATA).toString();
-                    String grade = (String) ((JSONObject) ((JSONObject) responseJson.get(APIConstants.ATTR))
-                            .get(APIConstants.DATA)).get(APIConstants.GRADE);
-                    Integer numErrors = Integer.valueOf(
-                            (String) ((JSONObject) ((JSONObject) responseJson.get(APIConstants.ATTR))
-                                    .get(APIConstants.DATA)).get(APIConstants.NUM_ERRORS));
-                    String decodedReport = new String(Base64Utils.decode(report), StandardCharsets.UTF_8);
-                    JSONObject output = new JSONObject();
-                    output.put("decodedReport", decodedReport);
-                    output.put("grade", grade);
-                    output.put("numErrors", numErrors);
-                    output.put("auditUuid", auditUuid);
-                    return output;
                 }
             }
         }
-        return new JSONObject(Collections.emptyMap());
     }
 
     /**