Skip to content

Commit

Permalink
Issue 1424 : improve Vulnerability query endpoints performance (#886)
Browse files Browse the repository at this point in the history
  • Loading branch information
sahibamittal authored Sep 5, 2024
1 parent 4de679f commit b8209eb
Show file tree
Hide file tree
Showing 8 changed files with 464 additions and 110 deletions.
6 changes: 6 additions & 0 deletions src/main/java/org/dependencytrack/model/Component.java
Original file line number Diff line number Diff line change
Expand Up @@ -395,6 +395,12 @@ public enum FetchGroup {
private transient Set<String> dependencyGraph;
private transient boolean expandDependencyGraph;

public Component(){}

public Component(final long id) {
this.id = id;
}

public long getId() {
return id;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,7 @@
@SuppressWarnings({"UnusedReturnValue", "unused"})
public class QueryManager extends AlpineQueryManager {

private AlpineRequest request;
protected AlpineRequest request;

private static final Logger LOGGER = Logger.getLogger(QueryManager.class);
private BomQueryManager bomQueryManager;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,10 @@
import org.dependencytrack.model.Vulnerability;
import org.dependencytrack.model.VulnerabilityAlias;
import org.dependencytrack.model.VulnerableSoftware;
import org.dependencytrack.persistence.jdbi.VulnerabilityDao;
import org.dependencytrack.persistence.jdbi.VulnerabilityDao.AffectedProjectCountRow;
import org.dependencytrack.resources.v1.vo.AffectedProject;
import org.jdbi.v3.core.Handle;

import javax.jdo.PersistenceManager;
import javax.jdo.Query;
Expand All @@ -51,6 +54,7 @@
import java.util.function.Function;
import java.util.stream.Collectors;

import static org.dependencytrack.persistence.jdbi.JdbiFactory.openJdbiHandle;
import static org.dependencytrack.util.PrincipalUtil.getPrincipalTeamIds;

final class VulnerabilityQueryManager extends QueryManager implements IQueryManager {
Expand Down Expand Up @@ -350,35 +354,22 @@ public PaginatedResult getVulnerabilities(Component component) {
* @return a List of Vulnerability objects
*/
public PaginatedResult getVulnerabilities(Component component, boolean includeSuppressed) {
PaginatedResult result;
final String componentFilter = (includeSuppressed) ? "components.contains(:component)" : "components.contains(:component)" + generateExcludeSuppressed(component.getProject(), component);
final Query<Vulnerability> query = pm.newQuery(Vulnerability.class);
if (orderBy == null) {
query.setOrdering("id asc");
}
if (filter != null) {
query.setFilter(componentFilter + " && vulnId.toLowerCase().matches(:vulnId)");
final String filterString = ".*" + filter.toLowerCase() + ".*";
result = execute(query, component, filterString);
} else {
query.setFilter(componentFilter);
result = execute(query, component);
List<Vulnerability> componentVulnerabilities;
List<AffectedProjectCountRow> affectedProjectCounts;

try (final Handle jdbiHandle = openJdbiHandle((this.request))) {
final var dao = jdbiHandle.attach(VulnerabilityDao.class);
componentVulnerabilities = dao.getVulnerabilitiesByComponent(component.getId(), includeSuppressed);
affectedProjectCounts = dao.getAffectedProjectCount(componentVulnerabilities.stream().map(Vulnerability::getId).toList());
}
Map<String, Epss> matchedEpssList = getEpssForCveIds(
result.getList(Vulnerability.class).stream().map(vuln -> vuln.getVulnId()).distinct().toList());
for (final Vulnerability vulnerability: result.getList(Vulnerability.class)) {
List<AffectedProject> affectedProjects = this.getAffectedProjects(vulnerability);
int affectedProjectsCount = affectedProjects.size();
int affectedActiveProjectsCount = (int) affectedProjects.stream().filter(AffectedProject::getActive).count();
int affectedInactiveProjectsCount = affectedProjectsCount - affectedActiveProjectsCount;

vulnerability.setAffectedProjectCount(affectedProjectsCount);
vulnerability.setAffectedActiveProjectCount(affectedActiveProjectsCount);
vulnerability.setAffectedInactiveProjectCount(affectedInactiveProjectsCount);
vulnerability.setAliases(getVulnerabilityAliases(vulnerability));
vulnerability.setEpss(matchedEpssList.get(vulnerability.getVulnId()));
Map<Long, Vulnerability> vulnById = componentVulnerabilities.stream().collect(Collectors.toMap(Vulnerability::getId, vulnerability -> vulnerability));
for (var vulnerabilityProjectCount : affectedProjectCounts) {
var vulnerability = vulnById.get(vulnerabilityProjectCount.id());
vulnerability.setAffectedProjectCount(vulnerabilityProjectCount.totalProjectCount());
vulnerability.setAffectedActiveProjectCount(vulnerabilityProjectCount.activeProjectCount());
vulnerability.setAffectedInactiveProjectCount(vulnerabilityProjectCount.totalProjectCount() - vulnerabilityProjectCount.activeProjectCount());
}
return result;
return (new PaginatedResult()).objects(componentVulnerabilities).total(componentVulnerabilities.size());
}

/**
Expand Down Expand Up @@ -462,22 +453,21 @@ public long getVulnerabilityCount(Project project, boolean includeSuppressed) {
* @return a List of Vulnerability objects
*/
public List<Vulnerability> getVulnerabilities(Project project, boolean includeSuppressed) {
final List<Vulnerability> vulnerabilities = new ArrayList<>();
final List<Component> components = getAllComponents(project);
for (final Component component: components) {
final Collection<Vulnerability> componentVulns = pm.detachCopyAll(
getAllVulnerabilities(component, includeSuppressed)
);
Map<String, Epss> matchedEpssList = getEpssForCveIds(
componentVulns.stream().map(vuln -> vuln.getVulnId()).distinct().toList());
for (final Vulnerability componentVuln: componentVulns) {
componentVuln.setComponents(Collections.singletonList(pm.detachCopy(component)));
componentVuln.setAliases(new ArrayList<>(pm.detachCopyAll(getVulnerabilityAliases(componentVuln))));
componentVuln.setEpss(matchedEpssList.get(componentVuln.getVulnId()));
List<Vulnerability> projectVulnerabilities;
var vulnerableComponents = new HashMap<Long, Component>();
try (final Handle jdbiHandle = openJdbiHandle((this.request))) {
final var dao = jdbiHandle.attach(VulnerabilityDao.class);
projectVulnerabilities = dao.getVulnerabilitiesByProject(project.getId(), includeSuppressed);
dao.getVulnerableComponents(project.getId(), projectVulnerabilities.stream().map(Vulnerability::getId).toList())
.stream().forEach(vc -> vulnerableComponents.put(vc.getId(), vc));
}
for (var projectVulnerability : projectVulnerabilities) {
if (projectVulnerability.getComponents() != null) {
projectVulnerability.setComponents(
projectVulnerability.getComponents().stream().map(c -> vulnerableComponents.get(c.getId())).toList());
}
vulnerabilities.addAll(componentVulns);
}
return vulnerabilities;
return projectVulnerabilities;
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,12 @@
*/
package org.dependencytrack.persistence.jdbi;

import org.dependencytrack.model.Component;
import org.dependencytrack.model.Vulnerability;
import org.dependencytrack.persistence.jdbi.mapping.VulnerabilityRowMapper;
import org.jdbi.v3.sqlobject.config.RegisterBeanMapper;
import org.jdbi.v3.sqlobject.config.RegisterConstructorMapper;
import org.jdbi.v3.sqlobject.config.RegisterRowMapper;
import org.jdbi.v3.sqlobject.customizer.Bind;
import org.jdbi.v3.sqlobject.customizer.DefineNamedBindings;
import org.jdbi.v3.sqlobject.statement.SqlQuery;
Expand Down Expand Up @@ -96,4 +101,134 @@ record AffectedProjectListRow(
long totalCount) {
}

@SqlQuery(/* language=InjectedFreeMarker */ """
<#-- @ftlvariable name="apiFilterParameter" type="String" -->
SELECT "V"."ID", "V"."CREATED", "V"."CVSSV2BASESCORE", "V"."CVSSV2VECTOR", "V"."CVSSV3BASESCORE", "V"."CVSSV3VECTOR"
, "V"."CWES", "V"."DESCRIPTION", "V"."DETAIL", "V"."PATCHEDVERSIONS", "V"."PUBLISHED", "V"."RECOMMENDATION", "V"."REFERENCES"
, "V"."SEVERITY", "V"."SOURCE", "V"."TITLE", "V"."UPDATED", "V"."UUID", "V"."VULNID", "V"."VULNERABLEVERSIONS", "V"."OWASPRRVECTOR"
, "EPSS"."SCORE"
, "EPSS"."PERCENTILE"
, JSONB_VULN_ALIASES("V"."SOURCE", "V"."VULNID") AS "vulnAliasesJson"
FROM "VULNERABILITY" AS "V"
INNER JOIN "COMPONENTS_VULNERABILITIES"
ON "V"."ID" = "COMPONENTS_VULNERABILITIES"."VULNERABILITY_ID"
AND "COMPONENTS_VULNERABILITIES"."COMPONENT_ID" = :componentId
INNER JOIN "COMPONENT"
ON "COMPONENTS_VULNERABILITIES"."COMPONENT_ID" = "COMPONENT"."ID"
LEFT JOIN "ANALYSIS"
ON "V"."ID" = "ANALYSIS"."VULNERABILITY_ID"
AND "COMPONENT"."PROJECT_ID" = "ANALYSIS"."PROJECT_ID"
LEFT JOIN "EPSS"
ON "V"."VULNID" = "EPSS"."CVE"
WHERE (:includeSuppressed OR "ANALYSIS"."SUPPRESSED" IS NULL OR NOT "ANALYSIS"."SUPPRESSED")
<#if apiFilterParameter??>
AND (LOWER("V"."VULNID") LIKE ('%' || LOWER(${apiFilterParameter}) || '%'))
</#if>
ORDER BY "V"."ID"
""")
@RegisterRowMapper(VulnerabilityRowMapper.class)
List<Vulnerability> getVulnerabilitiesByComponent(@Bind Long componentId, @Bind boolean includeSuppressed);

@SqlQuery("""
SELECT "VULNERABILITY"."ID" AS "id"
, COUNT("PROJECT"."ID") AS "totalProjectCount"
, COUNT(*) FILTER (WHERE "PROJECT"."ACTIVE") AS "activeProjectCount"
FROM "VULNERABILITY"
INNER JOIN "COMPONENTS_VULNERABILITIES"
ON "VULNERABILITY"."ID" = "COMPONENTS_VULNERABILITIES"."VULNERABILITY_ID"
INNER JOIN "COMPONENT"
ON "COMPONENTS_VULNERABILITIES"."COMPONENT_ID" = "COMPONENT"."ID"
INNER JOIN "PROJECT"
ON "COMPONENT"."PROJECT_ID" = "PROJECT"."ID"
WHERE "VULNERABILITY"."ID" = ANY(:vulnerabilityIds)
GROUP BY "VULNERABILITY"."ID"
""")
@RegisterConstructorMapper(AffectedProjectCountRow.class)
List<AffectedProjectCountRow> getAffectedProjectCount(@Bind List<Long> vulnerabilityIds);

record AffectedProjectCountRow(
long id,
int totalProjectCount,
int activeProjectCount
) {
}

@SqlQuery(/* language=InjectedFreeMarker */ """
<#-- @ftlvariable name="apiFilterParameter" type="String" -->
SELECT DISTINCT "V"."ID", "V"."CREATED", "V"."CVSSV2BASESCORE", "V"."CVSSV2VECTOR", "V"."CVSSV3BASESCORE", "V"."CVSSV3VECTOR"
, "V"."CWES", "V"."DESCRIPTION", "V"."DETAIL", "V"."PATCHEDVERSIONS", "V"."PUBLISHED", "V"."RECOMMENDATION", "V"."REFERENCES"
, "V"."SEVERITY", "V"."SOURCE", "V"."TITLE", "V"."UPDATED", "V"."UUID", "V"."VULNID", "V"."VULNERABLEVERSIONS", "V"."OWASPRRVECTOR"
, "EPSS"."SCORE"
, "EPSS"."PERCENTILE"
, ARRAY_AGG(DISTINCT("COMPONENT"."ID")) as "componentIdsArray"
, JSONB_VULN_ALIASES("V"."SOURCE", "V"."VULNID") AS "vulnAliasesJson"
FROM "VULNERABILITY" AS "V"
INNER JOIN "COMPONENTS_VULNERABILITIES"
ON "V"."ID" = "COMPONENTS_VULNERABILITIES"."VULNERABILITY_ID"
INNER JOIN "COMPONENT"
ON "COMPONENTS_VULNERABILITIES"."COMPONENT_ID" = "COMPONENT"."ID"
LEFT JOIN "EPSS"
ON "V"."VULNID" = "EPSS"."CVE"
LEFT JOIN "ANALYSIS"
ON "V"."ID" = "ANALYSIS"."VULNERABILITY_ID"
AND "COMPONENT"."ID" = "ANALYSIS"."COMPONENT_ID"
AND "COMPONENT"."PROJECT_ID" = "ANALYSIS"."PROJECT_ID"
WHERE "COMPONENT"."PROJECT_ID" = :projectId
AND (:includeSuppressed OR "ANALYSIS"."SUPPRESSED" IS NULL OR NOT "ANALYSIS"."SUPPRESSED")
<#if apiFilterParameter??>
AND (LOWER("V"."VULNID") LIKE ('%' || LOWER(${apiFilterParameter}) || '%'))
</#if>
GROUP BY "V"."ID", "EPSS"."SCORE", "EPSS"."PERCENTILE"
ORDER BY "V"."ID"
""")
@RegisterRowMapper(VulnerabilityRowMapper.class)
List<Vulnerability> getVulnerabilitiesByProject(@Bind long projectId, boolean includeSuppressed);

@SqlQuery("""
SELECT distinct "C"."ID",
"C"."NAME",
"C"."AUTHOR",
"C"."BLAKE2B_256",
"C"."BLAKE2B_384",
"C"."BLAKE2B_512",
"C"."BLAKE3",
"C"."CLASSIFIER",
"C"."COPYRIGHT",
"C"."CPE",
"C"."PUBLISHER",
"C"."PURL",
"C"."PURLCOORDINATES",
"C"."DESCRIPTION",
"C"."DIRECT_DEPENDENCIES",
"C"."EXTENSION",
"C"."EXTERNAL_REFERENCES",
"C"."FILENAME",
"C"."GROUP",
"C"."INTERNAL",
"C"."LAST_RISKSCORE",
"C"."LICENSE",
"C"."LICENSE_EXPRESSION",
"C"."LICENSE_URL",
"C"."TEXT",
"C"."MD5",
"C"."SHA1",
"C"."SHA_256",
"C"."SHA_384",
"C"."SHA_512",
"C"."SHA3_256",
"C"."SHA3_384",
"C"."SHA3_512",
"C"."SWIDTAGID",
"C"."UUID",
"C"."VERSION"
FROM "COMPONENT" AS "C"
INNER JOIN "COMPONENTS_VULNERABILITIES"
ON "C"."ID" = "COMPONENTS_VULNERABILITIES"."COMPONENT_ID"
INNER JOIN "VULNERABILITY"
ON "COMPONENTS_VULNERABILITIES"."VULNERABILITY_ID" = "VULNERABILITY"."ID"
WHERE "VULNERABILITY"."ID" = ANY(:vulnerabilityIds)
and "C"."PROJECT_ID" = :projectId
""")
@RegisterBeanMapper(Component.class)
List<Component> getVulnerableComponents(@Bind long projectId, @Bind List<Long> vulnerabilityIds);
}
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,17 @@ public static List<String> stringArray(final ResultSet rs, final String columnNa
return Arrays.asList((String[]) array.getArray());
}

public static List<Long> longArray(final ResultSet rs, final String columnName) throws SQLException {
final Array array = rs.getArray(columnName);
if (array == null) {
return Collections.emptyList();
}
if (array.getBaseType() != Types.BIGINT) {
throw new IllegalArgumentException("Expected array with base type BIGINT, but got %s".formatted(array.getBaseTypeName()));
}
return Arrays.asList((Long[]) array.getArray());
}

public static <T> T deserializeJson(final ResultSet rs, final String columnName, final TypeReference<T> typeReference) throws SQLException {
final String jsonString = rs.getString(columnName);
if (isBlank(jsonString)) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
/*
* This file is part of Dependency-Track.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
* SPDX-License-Identifier: Apache-2.0
* Copyright (c) OWASP Foundation. All Rights Reserved.
*/
package org.dependencytrack.persistence.jdbi.mapping;

import com.fasterxml.jackson.core.type.TypeReference;
import org.dependencytrack.model.Component;
import org.dependencytrack.model.Epss;
import org.dependencytrack.model.Vulnerability;
import org.dependencytrack.model.VulnerabilityAlias;
import org.jdbi.v3.core.mapper.RowMapper;
import org.jdbi.v3.core.mapper.reflect.BeanMapper;
import org.jdbi.v3.core.statement.StatementContext;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.List;

import static org.dependencytrack.persistence.jdbi.mapping.RowMapperUtil.deserializeJson;
import static org.dependencytrack.persistence.jdbi.mapping.RowMapperUtil.hasColumn;
import static org.dependencytrack.persistence.jdbi.mapping.RowMapperUtil.longArray;
import static org.dependencytrack.persistence.jdbi.mapping.RowMapperUtil.maybeSet;

public class VulnerabilityRowMapper implements RowMapper<Vulnerability> {

private static final TypeReference<List<VulnerabilityAlias>> VULNERABILITY_ALIASES_TYPE_REF = new TypeReference<>() {
};

@Override
public Vulnerability map(final ResultSet rs, final StatementContext ctx) throws SQLException {
final Vulnerability vuln = BeanMapper.of(Vulnerability.class).map(rs, ctx);
final Epss epss = BeanMapper.of(Epss.class).map(rs, ctx);
vuln.setEpss(epss);
maybeSet(rs, "vulnAliasesJson", (ignored, columnName) ->
deserializeJson(rs, columnName, VULNERABILITY_ALIASES_TYPE_REF), vuln::setAliases);
if (hasColumn(rs, "componentIdsArray")) {
var vulnerableComponentsIds = longArray(rs, "componentIdsArray");
if (!vulnerableComponentsIds.isEmpty()) {
vuln.setComponents(vulnerableComponentsIds.stream().map(cid -> new Component(cid)).toList());
}
}
return vuln;
}
}
Loading

0 comments on commit b8209eb

Please sign in to comment.