diff --git a/src/main/java/org/sasanlabs/service/vulnerability/sqlInjection/ErrorBasedSQLInjectionVulnerability.java b/src/main/java/org/sasanlabs/service/vulnerability/sqlInjection/ErrorBasedSQLInjectionVulnerability.java index 507adfde3..ed5a286fe 100644 --- a/src/main/java/org/sasanlabs/service/vulnerability/sqlInjection/ErrorBasedSQLInjectionVulnerability.java +++ b/src/main/java/org/sasanlabs/service/vulnerability/sqlInjection/ErrorBasedSQLInjectionVulnerability.java @@ -104,12 +104,16 @@ public ResponseEntity doesCarInformationExistsLevel1( htmlTemplate = "LEVEL_1/SQLInjection_Level1") public ResponseEntity doesCarInformationExistsLevel2( @RequestParam Map queryParams) { - String id = queryParams.get(Constants.ID); + final String id = queryParams.get(Constants.ID); BodyBuilder bodyBuilder = ResponseEntity.status(HttpStatus.OK); + bodyBuilder.body(ErrorBasedSQLInjectionVulnerability.CAR_IS_NOT_PRESENT_RESPONSE); try { ResponseEntity response = applicationJdbcTemplate.query( - "select * from cars where id='" + id + "'", + (conn) -> conn.prepareStatement("select * from cars where id=?"), + (prepareStatement) -> { + prepareStatement.setString(1, id); + }, (rs) -> { if (rs.next()) { CarInformation carInformation = new CarInformation(); @@ -124,7 +128,8 @@ public ResponseEntity doesCarInformationExistsLevel2( } catch (JsonProcessingException e) { LOGGER.error("Following error occurred", e); return bodyBuilder.body( - GENERIC_EXCEPTION_RESPONSE_FUNCTION.apply(e)); + ErrorBasedSQLInjectionVulnerability + .CAR_IS_NOT_PRESENT_RESPONSE); } } else { return bodyBuilder.body( @@ -135,7 +140,8 @@ public ResponseEntity doesCarInformationExistsLevel2( return response; } catch (Exception ex) { LOGGER.error("Following error occurred", ex); - return bodyBuilder.body(GENERIC_EXCEPTION_RESPONSE_FUNCTION.apply(ex)); + return bodyBuilder.body( + ErrorBasedSQLInjectionVulnerability.CAR_IS_NOT_PRESENT_RESPONSE); } } diff --git a/src/test/java/org/sasanlabs/service/vulnerability/sqlInjection/ErrorBasedSQLInjectionVulnerabilityTest.java b/src/test/java/org/sasanlabs/service/vulnerability/sqlInjection/ErrorBasedSQLInjectionVulnerabilityTest.java index a665b7540..a2a571a41 100644 --- a/src/test/java/org/sasanlabs/service/vulnerability/sqlInjection/ErrorBasedSQLInjectionVulnerabilityTest.java +++ b/src/test/java/org/sasanlabs/service/vulnerability/sqlInjection/ErrorBasedSQLInjectionVulnerabilityTest.java @@ -5,6 +5,7 @@ import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.eq; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import java.io.IOException; @@ -59,16 +60,33 @@ void doesCarInformationExistsLevel1_ExpectParamEscaped() throws IOException { } @Test - void doesCarInformationExistsLevel2_ExpectParamEscaped() throws IOException { + void doesCarInformationExistsLevel2_ExpectValidResponse() { + // Arrange + Map queryParams = new HashMap<>(); + queryParams.put(Constants.ID, "1"); + + ResponseEntity mockResponseEntity = + ResponseEntity.status(HttpStatus.OK).body("Sample response"); + doReturn(mockResponseEntity) + .when(template) + .query( + Mockito.any(PreparedStatementCreator.class), + Mockito.any(PreparedStatementSetter.class), + Mockito.any(ResultSetExtractor.class)); + // Act - final Map queryParams = Collections.singletonMap("id", "1"); - errorBasedSQLInjectionVulnerability.doesCarInformationExistsLevel2(queryParams); + ResponseEntity response = + errorBasedSQLInjectionVulnerability.doesCarInformationExistsLevel2(queryParams); // Assert + assertEquals(HttpStatus.OK, response.getStatusCode()); + assertEquals("Sample response", response.getBody()); verify(template) .query( - eq("select * from cars where id='1'"), - (ResultSetExtractor) any()); + Mockito.any(PreparedStatementCreator.class), + Mockito.any(PreparedStatementSetter.class), + Mockito.any(ResultSetExtractor.class)); + verify(template, never()).query(anyString(), (ResultSetExtractor) any()); } @Test