Skip to content

Commit 6e3d97a

Browse files
demoralizerrkartik patel
demoralizerr
and
kartik patel
authored
Increased Junit coverage for 'PreflightController' & 'UnrestrictedFileUpload' class (#493)
* Increase coverage of 'UnrestrictedFileUpload' to 95% & 'PreflightController' to 100%. * Increase coverage of 'UnrestrictedFileUpload' to 95% & 'PreflightController' to 100%. * Increase coverage of 'UnrestrictedFileUpload' to 95% & 'PreflightController' to 100%. * Remove SampleVulnerability.java from PR --------- Co-authored-by: kartik patel <kartik.patell4july1997@gmail.com>
1 parent 03ecd3d commit 6e3d97a

File tree

4 files changed

+258
-9
lines changed

4 files changed

+258
-9
lines changed

‎build.gradle‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,9 @@ dependencies {
124124
// https://mvnrepository.com/artifact/org.assertj/assertj-core
125125
testImplementation group: 'org.assertj', name: 'assertj-core', version: '3.17.2'
126126

127+
// https://mvnrepository.com/artifact/org.springframework.boot/spring-boot-starter-test
128+
testImplementation 'org.springframework.boot:spring-boot-starter-test'
129+
127130
// https://mvnrepository.com/artifact/org.springframework.boot/spring-boot-starter-data-jpa
128131
implementation group: 'org.springframework.boot', name: 'spring-boot-starter-data-jpa', version: '2.3.1.RELEASE'
129132

‎src/main/java/org/sasanlabs/service/vulnerability/fileupload/PreflightController.java‎

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
package org.sasanlabs.service.vulnerability.fileupload;
22

33
import static org.sasanlabs.service.vulnerability.fileupload.UnrestrictedFileUpload.CONTENT_DISPOSITION_STATIC_FILE_LOCATION;
4+
import static org.springframework.http.HttpHeaders.CONTENT_DISPOSITION;
45

56
import java.io.FileInputStream;
67
import java.io.IOException;
78
import java.io.InputStream;
9+
import java.nio.file.Path;
810
import org.apache.commons.io.IOUtils;
911
import org.sasanlabs.internal.utility.FrameworkConstants;
1012
import org.springframework.http.HttpHeaders;
@@ -34,14 +36,16 @@ public PreflightController(UnrestrictedFileUpload unrestrictedFileUpload) {
3436
CONTENT_DISPOSITION_STATIC_FILE_LOCATION + FrameworkConstants.SLASH + "{fileName}")
3537
public ResponseEntity<byte[]> fetchFile(@PathVariable("fileName") String fileName)
3638
throws IOException {
37-
InputStream inputStream =
38-
new FileInputStream(
39-
unrestrictedFileUpload.getContentDispositionRoot().toFile()
40-
+ FrameworkConstants.SLASH
41-
+ fileName);
42-
HttpHeaders httpHeaders = new HttpHeaders();
43-
httpHeaders.add(HttpHeaders.CONTENT_DISPOSITION, "attachment");
44-
return new ResponseEntity<byte[]>(
45-
IOUtils.toByteArray(inputStream), httpHeaders, HttpStatus.OK);
39+
40+
// Resolve path using Path API
41+
Path filePath = unrestrictedFileUpload.getContentDispositionRoot().resolve(fileName);
42+
43+
// Try-with-resources ensures the stream closes automatically
44+
try (InputStream inputStream = new FileInputStream(filePath.toFile())) {
45+
byte[] fileBytes = IOUtils.toByteArray(inputStream);
46+
HttpHeaders httpHeaders = new HttpHeaders();
47+
httpHeaders.add(CONTENT_DISPOSITION, "attachment");
48+
return new ResponseEntity<>(fileBytes, httpHeaders, HttpStatus.OK);
49+
}
4650
}
4751
}

‎src/test/java/org/sasanlabs/service/vulnerability/fileupload/PreflightControllerTest.java‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
package org.sasanlabs.service.vulnerability.fileupload;
2+
3+
import static org.mockito.Mockito.when;
4+
import static org.mockito.MockitoAnnotations.openMocks;
5+
import static org.sasanlabs.internal.utility.FrameworkConstants.SLASH;
6+
import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get;
7+
import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.content;
8+
import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status;
9+
10+
import java.nio.file.Files;
11+
import java.nio.file.Path;
12+
import java.util.Collections;
13+
import org.junit.jupiter.api.BeforeEach;
14+
import org.junit.jupiter.api.Test;
15+
import org.junit.jupiter.api.io.TempDir;
16+
import org.mockito.Mock;
17+
import org.springframework.test.web.servlet.MockMvc;
18+
import org.springframework.test.web.servlet.setup.MockMvcBuilders;
19+
20+
class PreflightControllerTest {
21+
22+
private MockMvc mockMvc;
23+
24+
@Mock private UnrestrictedFileUpload unrestrictedFileUpload;
25+
26+
@TempDir Path tempDir;
27+
28+
private static final String FILE_NAME = "testFile";
29+
private static final String FILE_CONTENT = "Hello SasanLabs!";
30+
31+
@BeforeEach
32+
void setUp() {
33+
openMocks(this);
34+
PreflightController preflightController = new PreflightController(unrestrictedFileUpload);
35+
mockMvc = MockMvcBuilders.standaloneSetup(preflightController).build();
36+
}
37+
38+
@Test
39+
void testFetchFile_Success() throws Exception {
40+
// given
41+
Path filePath = tempDir.resolve(FILE_NAME);
42+
Files.write(filePath, Collections.singletonList(FILE_CONTENT));
43+
String url =
44+
SLASH
45+
+ UnrestrictedFileUpload.CONTENT_DISPOSITION_STATIC_FILE_LOCATION
46+
+ SLASH
47+
+ FILE_NAME;
48+
49+
// when
50+
when(unrestrictedFileUpload.getContentDispositionRoot()).thenReturn(tempDir);
51+
52+
// then
53+
mockMvc.perform(get(url))
54+
.andExpect(status().isOk())
55+
.andExpect(content().string(org.hamcrest.Matchers.containsString(FILE_CONTENT)));
56+
}
57+
}

‎src/test/java/org/sasanlabs/service/vulnerability/fileupload/UnrestrictedFileUploadTest.java‎

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,185 @@
1+
package org.sasanlabs.service.vulnerability.fileupload;
2+
3+
import static org.assertj.core.api.Assertions.assertThat;
4+
import static org.springframework.test.util.ReflectionTestUtils.setField;
5+
6+
import java.io.IOException;
7+
import java.net.URISyntaxException;
8+
import java.nio.file.Files;
9+
import java.nio.file.Path;
10+
import java.util.List;
11+
import java.util.stream.Collectors;
12+
import java.util.stream.Stream;
13+
import org.junit.jupiter.api.BeforeEach;
14+
import org.junit.jupiter.api.Test;
15+
import org.junit.jupiter.api.io.TempDir;
16+
import org.sasanlabs.service.vulnerability.bean.GenericVulnerabilityResponseBean;
17+
import org.sasanlabs.vulnerability.utils.Constants;
18+
import org.springframework.http.ResponseEntity;
19+
import org.springframework.mock.web.MockMultipartFile;
20+
21+
class UnrestrictedFileUploadTest {
22+
23+
private UnrestrictedFileUpload fileUpload;
24+
25+
@TempDir Path tempRoot;
26+
27+
@TempDir Path tempContentDispositionRoot;
28+
29+
@BeforeEach
30+
void setUp() throws IOException, URISyntaxException {
31+
fileUpload = new UnrestrictedFileUpload();
32+
setField(fileUpload, "root", tempRoot);
33+
setField(fileUpload, "contentDispositionRoot", tempContentDispositionRoot);
34+
}
35+
36+
@Test
37+
void test_getVulnerablePayloadLevel1() throws Exception {
38+
MockMultipartFile file =
39+
new MockMultipartFile("file", "test.html", "text/html", "content".getBytes());
40+
ResponseEntity<GenericVulnerabilityResponseBean<String>> response =
41+
fileUpload.getVulnerablePayloadLevel1(file);
42+
43+
assertThat(response.getBody().getIsValid()).isTrue();
44+
assertThat(tempRoot.resolve("test.html")).exists();
45+
}
46+
47+
@Test
48+
void test_getVulnerablePayloadLevel2() throws Exception {
49+
MockMultipartFile file =
50+
new MockMultipartFile("file", "image.png", "image/png", "content".getBytes());
51+
ResponseEntity<GenericVulnerabilityResponseBean<String>> response =
52+
fileUpload.getVulnerablePayloadLevel2(file);
53+
54+
assertThat(response.getBody().getIsValid()).isTrue();
55+
try (Stream<Path> files = Files.list(tempRoot)) {
56+
assertThat(files.count()).isGreaterThan(0);
57+
}
58+
}
59+
60+
@Test
61+
void test_getVulnerablePayloadLevel3() throws Exception {
62+
MockMultipartFile file =
63+
new MockMultipartFile("file", "attack.html", "text/html", "content".getBytes());
64+
ResponseEntity<GenericVulnerabilityResponseBean<String>> response =
65+
fileUpload.getVulnerablePayloadLevel3(file);
66+
67+
assertThat(response.getBody().getIsValid()).isFalse();
68+
}
69+
70+
@Test
71+
void test_getVulnerablePayloadLevel4() throws Exception {
72+
MockMultipartFile file =
73+
new MockMultipartFile("file", "attack.htm", "text/html", "content".getBytes());
74+
ResponseEntity<GenericVulnerabilityResponseBean<String>> response =
75+
fileUpload.getVulnerablePayloadLevel4(file);
76+
77+
assertThat(response.getBody().getIsValid()).isFalse();
78+
}
79+
80+
@Test
81+
void test_getVulnerablePayloadLevel5() throws Exception {
82+
MockMultipartFile file =
83+
new MockMultipartFile("file", "attack.HTM", "text/html", "content".getBytes());
84+
ResponseEntity<GenericVulnerabilityResponseBean<String>> response =
85+
fileUpload.getVulnerablePayloadLevel5(file);
86+
87+
assertThat(response.getBody().getIsValid()).isFalse();
88+
}
89+
90+
@Test
91+
void test_getVulnerablePayloadLevel6() throws Exception {
92+
MockMultipartFile file =
93+
new MockMultipartFile(
94+
"file", "malicious.png.html", "text/html", "content".getBytes());
95+
ResponseEntity<GenericVulnerabilityResponseBean<String>> response =
96+
fileUpload.getVulnerablePayloadLevel6(file);
97+
98+
assertThat(response.getBody().getIsValid()).isTrue();
99+
}
100+
101+
@Test
102+
void test_getVulnerablePayloadLevel7() throws Exception {
103+
String fileName = "shell.php" + Constants.NULL_BYTE_CHARACTER + ".png";
104+
MockMultipartFile file =
105+
new MockMultipartFile("file", fileName, "text/plain", "content".getBytes());
106+
107+
ResponseEntity<GenericVulnerabilityResponseBean<String>> response =
108+
fileUpload.getVulnerablePayloadLevel7(file);
109+
110+
assertThat(response.getBody().getIsValid())
111+
.withFailMessage("Level 7 validation should pass for null byte suffix")
112+
.isTrue();
113+
114+
try (Stream<Path> files = Files.list(tempRoot)) {
115+
List<String> fileNames =
116+
files.map(p -> p.getFileName().toString()).collect(Collectors.toList());
117+
118+
assertThat(fileNames)
119+
.withFailMessage("Expected truncated .php file, but found: %s", fileNames)
120+
.anyMatch(name -> name.contains("shell.php") && !name.endsWith(".png"));
121+
}
122+
}
123+
124+
@Test
125+
void test_getVulnerablePayloadLevel7_WithoutNullByte() throws Exception {
126+
// 1. Arrange: Use a filename WITHOUT a null byte
127+
String fileName = "normal_image.png";
128+
MockMultipartFile file =
129+
new MockMultipartFile("file", fileName, "image/png", "dummy content".getBytes());
130+
131+
// 2. Act
132+
ResponseEntity<GenericVulnerabilityResponseBean<String>> response =
133+
fileUpload.getVulnerablePayloadLevel7(file);
134+
135+
// 3. Assert Response
136+
assertThat(response.getBody().getIsValid())
137+
.withFailMessage("Validation should pass for a normal .png file")
138+
.isTrue();
139+
140+
// 4. Assert File System
141+
try (Stream<Path> files = Files.list(tempRoot)) {
142+
List<String> fileNames =
143+
files.map(p -> p.getFileName().toString()).collect(Collectors.toList());
144+
145+
// Verify that the file was saved ending with the original name (normal_image.png)
146+
// and includes the random prefix assigned by the controller
147+
assertThat(fileNames)
148+
.withFailMessage(
149+
"Expected file to be saved with original name, but found: %s",
150+
fileNames)
151+
.anyMatch(name -> name.endsWith("_" + fileName));
152+
}
153+
}
154+
155+
@Test
156+
void test_getVulnerablePayloadLevel8() throws Exception {
157+
MockMultipartFile file =
158+
new MockMultipartFile("file", "traversal.txt", "text/plain", "content".getBytes());
159+
ResponseEntity<GenericVulnerabilityResponseBean<String>> response =
160+
fileUpload.getVulnerablePayloadLevel8(file);
161+
162+
assertThat(response.getBody().getIsValid()).isTrue();
163+
assertThat(tempContentDispositionRoot.resolve("traversal.txt")).exists();
164+
}
165+
166+
@Test
167+
void test_getVulnerablePayloadLevel9() throws Exception {
168+
MockMultipartFile file =
169+
new MockMultipartFile("file", "dos.txt", "text/plain", new byte[1024]);
170+
ResponseEntity<GenericVulnerabilityResponseBean<String>> response =
171+
fileUpload.getVulnerablePayloadLevel9(file);
172+
173+
assertThat(response.getBody().getIsValid()).isTrue();
174+
}
175+
176+
@Test
177+
void test_getVulnerablePayloadLevel10() throws Exception {
178+
MockMultipartFile malicious =
179+
new MockMultipartFile("file", "test.png.php", "text/plain", "content".getBytes());
180+
ResponseEntity<GenericVulnerabilityResponseBean<String>> response =
181+
fileUpload.getVulnerablePayloadLevel10(malicious);
182+
183+
assertThat(response.getBody().getIsValid()).isFalse();
184+
}
185+
}

0 commit comments

Comments
 (0)