Skip to content

Commit 58761b9

Browse files
committed
Add secure work with Cookie libraries
1 parent 419d415 commit 58761b9

File tree

1 file changed

+17
-6
lines changed

1 file changed

+17
-6
lines changed

utbot-framework-api/src/main/kotlin/org/utbot/framework/plugin/api/util/SpringModelUtils.kt

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package org.utbot.framework.plugin.api.util
22

3+
import mu.KotlinLogging
34
import org.utbot.common.tryLoadClass
45
import org.utbot.common.withToStringThreadLocalReentrancyGuard
56
import org.utbot.framework.plugin.api.isNotNull
@@ -17,6 +18,8 @@ import org.utbot.framework.plugin.api.mapper.UtModelDeepMapper
1718
import org.utbot.framework.plugin.api.mapper.mapModels
1819
import java.util.Optional
1920

21+
private val logger = KotlinLogging.logger {}
22+
2023
object SpringModelUtils {
2124
val autowiredClassId = ClassId("org.springframework.beans.factory.annotation.Autowired")
2225

@@ -57,7 +60,7 @@ object SpringModelUtils {
5760

5861
// the library in which Cookie is stored depends on the version of Spring
5962
private val cookiesLibraries = listOf("javax.servlet.http", "jakarta.servlet.http")
60-
private val cookieClassId = getClassIdFromEachAvailablePackage(cookiesLibraries, "Cookie").single()
63+
private val cookieClassIds = getClassIdFromEachAvailablePackage(cookiesLibraries, "Cookie")
6164

6265
val entityClassIds get() = persistentClassIds("Entity")
6366
val generatedValueClassIds get() = persistentClassIds("GeneratedValue")
@@ -210,7 +213,7 @@ object SpringModelUtils {
210213
parameters = listOf(httpHeaderClassId)
211214
)
212215

213-
private val mockHttpServletCookieMethodId = MethodId(
216+
private fun mockHttpServletCookieMethodId(cookieClassId: ClassId) = MethodId(
214217
classId = mockHttpServletRequestBuilderClassId,
215218
name = "cookie",
216219
returnType = mockHttpServletRequestBuilderClassId,
@@ -448,9 +451,15 @@ object SpringModelUtils {
448451
val headersContentModel = createHeadersContentModel(methodId, arguments, idGenerator)
449452
requestBuilderModel = addHeadersToRequestBuilderModel(headersContentModel, requestBuilderModel, idGenerator)
450453

451-
val cookieValuesModel = createCookieValuesModel(methodId, arguments, idGenerator)
452-
requestBuilderModel =
453-
addCookiesToRequestBuilderModel(cookieValuesModel, requestBuilderModel, idGenerator)
454+
val cookieClassId = cookieClassIds.singleOrNull()
455+
if(cookieClassId is ClassId) {
456+
val cookieValuesModel = createCookieValuesModel(cookieClassId, methodId, arguments, idGenerator)
457+
requestBuilderModel =
458+
addCookiesToRequestBuilderModel(cookieClassId, cookieValuesModel, requestBuilderModel, idGenerator)
459+
}
460+
else{
461+
logger.warn { "Cookie library not found" }
462+
}
454463

455464
val requestAttributes = collectArgumentsWithAnnotationModels(methodId, requestAttributesClassId, arguments)
456465
requestBuilderModel =
@@ -528,6 +537,7 @@ object SpringModelUtils {
528537
}
529538

530539
private fun addCookiesToRequestBuilderModel(
540+
cookieClassId: ClassId,
531541
cookieValuesModel: UtArrayModel,
532542
requestBuilderModel: UtAssembleModel,
533543
idGenerator: () -> Int
@@ -542,7 +552,7 @@ object SpringModelUtils {
542552
modelName = "requestBuilder",
543553
instantiationCall = UtExecutableCallModel(
544554
instance = requestBuilderModel,
545-
executable = mockHttpServletCookieMethodId,
555+
executable = mockHttpServletCookieMethodId(cookieClassId),
546556
params = listOf(cookieValuesModel)
547557
)
548558
)
@@ -676,6 +686,7 @@ object SpringModelUtils {
676686
}
677687

678688
private fun createCookieValuesModel(
689+
cookieClassId: ClassId,
679690
methodId: MethodId,
680691
arguments: List<UtModel>,
681692
idGenerator: () -> Int,

0 commit comments

Comments
 (0)