Mocked network calls made in SSRFVulnerabilityTest.

SasanLabs · Oct 27, 2023 · c42d198 · c42d198
1 parent 571108b
commit c42d198
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 9 deletions.
diff --git a/src/main/java/org/sasanlabs/service/vulnerability/ssrf/SSRFVulnerability.java b/src/main/java/org/sasanlabs/service/vulnerability/ssrf/SSRFVulnerability.java
@@ -68,20 +68,24 @@ private ResponseEntity<GenericVulnerabilityResponseBean<String>> invalidUrlRespo
                                 MetaDataServiceMock.getResponse(u), true),
                         HttpStatus.OK);
             } else {
-                URLConnection urlConnection = u.openConnection();
-                try (BufferedReader reader =
-                        new BufferedReader(new InputStreamReader(urlConnection.getInputStream()))) {
-                    return new ResponseEntity<>(
-                            new GenericVulnerabilityResponseBean<>(
-                                    reader.lines().collect(Collectors.joining()), true),
-                            HttpStatus.OK);
-                }
+                return new ResponseEntity<>(
+                    new GenericVulnerabilityResponseBean<>(
+                            getResponseForURLConnection(u), true),
+                    HttpStatus.OK);
             }
         } else {
             return invalidUrlResponse();
         }
     }
 
+    String getResponseForURLConnection(URL u) throws IOException {
+        URLConnection urlConnection = u.openConnection();
+        try (BufferedReader reader =
+                new BufferedReader(new InputStreamReader(urlConnection.getInputStream()))) {
+            return reader.lines().collect(Collectors.joining());
+        }
+    }
+
     @AttackVector(
             vulnerabilityExposed = VulnerabilityType.SIMPLE_SSRF,
             description = "SSRF_VULNERABILITY_URL_WITHOUT_CHECK",

diff --git a/src/test/java/org/sasanlabs/service/vulnerability/ssrf/SSRFVulnerabilityTest.java b/src/test/java/org/sasanlabs/service/vulnerability/ssrf/SSRFVulnerabilityTest.java
@@ -3,13 +3,18 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
 
 import java.io.File;
 import java.io.IOException;
+import java.net.URL;
 import java.nio.file.Files;
 import java.util.Collections;
 import java.util.stream.Stream;
 import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.io.TempDir;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
@@ -36,7 +41,16 @@ class SSRFVulnerabilityTest {
 
     private static String tempFileUrl;
 
-    private final SSRFVulnerability ssrfVulnerability = new SSRFVulnerability(GIST_ID);
+    private SSRFVulnerability ssrfVulnerability;
+
+    @BeforeEach
+    void each() throws IOException {
+        SSRFVulnerability ssrfSpy = spy(new SSRFVulnerability(GIST_ID));
+        // mocks network calls
+        doReturn(GIST_URL_CONTENT).when(ssrfSpy).getResponseForURLConnection(eq(new URL(GIST_URL)));
+        doReturn(OTHER_URL_CONTENT).when(ssrfSpy).getResponseForURLConnection(eq(new URL(OTHER_URL)));
+        ssrfVulnerability = ssrfSpy;
+    }
 
     @BeforeAll
     static void setUp() throws IOException {