From 67f5d7958fcd1b6065eab6974520fdbc8286de34 Mon Sep 17 00:00:00 2001 From: Tsunami Team Date: Thu, 25 Jul 2024 23:50:56 -0700 Subject: [PATCH] Support compact RunRequest so requests sent to heavy plugin services still fit in the default gRPC message limit. PiperOrigin-RevId: 656270995 Change-Id: I3216fe01796ff866caa464aad1d49d1c0dd35256 --- .../server/CompactRunRequestHelper.java | 73 ++++++++++++ .../server/CompactRunRequestHelperTest.java | 110 ++++++++++++++++++ .../tsunami/plugin/PluginServiceClient.java | 13 +++ .../plugin/RemoteVulnDetectorImpl.java | 28 +++-- .../plugin/RemoteVulnDetectorImplTest.java | 57 +++++++++ proto/plugin_service.proto | 40 ++++++- 6 files changed, 309 insertions(+), 12 deletions(-) create mode 100644 common/src/main/java/com/google/tsunami/common/server/CompactRunRequestHelper.java create mode 100644 common/src/test/java/com/google/tsunami/common/server/CompactRunRequestHelperTest.java diff --git a/common/src/main/java/com/google/tsunami/common/server/CompactRunRequestHelper.java b/common/src/main/java/com/google/tsunami/common/server/CompactRunRequestHelper.java new file mode 100644 index 00000000..72bc6e8d --- /dev/null +++ b/common/src/main/java/com/google/tsunami/common/server/CompactRunRequestHelper.java @@ -0,0 +1,73 @@ +/* + * Copyright 2024 Google LLC + * + * 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. + */ +package com.google.tsunami.common.server; + +import com.google.common.collect.ImmutableList; +import com.google.tsunami.proto.MatchedPlugin; +import com.google.tsunami.proto.NetworkService; +import com.google.tsunami.proto.RunCompactRequest; +import com.google.tsunami.proto.RunCompactRequest.PluginNetworkServiceTarget; +import com.google.tsunami.proto.RunRequest; +import java.util.HashMap; + +/** + * CompactRunRequestHelper is a helper class to compress/uncompress the RunRequest into/from the + * compact representation. + */ +public final class CompactRunRequestHelper { + + private CompactRunRequestHelper() {} + + public static RunCompactRequest compress(RunRequest runRequest) { + var builder = RunCompactRequest.newBuilder().setTarget(runRequest.getTarget()); + HashMap serviceIndexMap = new HashMap<>(); + int pluginIndex = -1; + for (MatchedPlugin matchedPlugin : runRequest.getPluginsList()) { + pluginIndex++; + builder.addPlugins(matchedPlugin.getPlugin()); + for (NetworkService service : matchedPlugin.getServicesList()) { + Integer serviceIndex = serviceIndexMap.get(service); + if (serviceIndex == null) { + serviceIndex = serviceIndexMap.size(); + serviceIndexMap.put(service, serviceIndex); + builder.addServices(service); + } + + builder.addScanTargets( + PluginNetworkServiceTarget.newBuilder() + .setPluginIndex(pluginIndex) + .setServiceIndex(serviceIndex) + .build()); + } + } + return builder.build(); + } + + public static RunRequest uncompress(RunCompactRequest runCompactRequest) { + ImmutableList.Builder matchedPlugins = ImmutableList.builder(); + for (var target : runCompactRequest.getScanTargetsList()) { + var plugin = runCompactRequest.getPlugins(target.getPluginIndex()); + var networkService = runCompactRequest.getServices(target.getServiceIndex()); + matchedPlugins.add( + MatchedPlugin.newBuilder().setPlugin(plugin).addServices(networkService).build()); + } + + return RunRequest.newBuilder() + .setTarget(runCompactRequest.getTarget()) + .addAllPlugins(matchedPlugins.build()) + .build(); + } +} diff --git a/common/src/test/java/com/google/tsunami/common/server/CompactRunRequestHelperTest.java b/common/src/test/java/com/google/tsunami/common/server/CompactRunRequestHelperTest.java new file mode 100644 index 00000000..6dda9e8d --- /dev/null +++ b/common/src/test/java/com/google/tsunami/common/server/CompactRunRequestHelperTest.java @@ -0,0 +1,110 @@ +/* + * Copyright 2024 Google LLC + * + * 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. + */ +package com.google.tsunami.common.server; + +import static com.google.common.truth.Truth.assertThat; + +import com.google.common.collect.ImmutableList; +import com.google.tsunami.proto.Hostname; +import com.google.tsunami.proto.MatchedPlugin; +import com.google.tsunami.proto.NetworkEndpoint; +import com.google.tsunami.proto.NetworkService; +import com.google.tsunami.proto.PluginDefinition; +import com.google.tsunami.proto.PluginInfo; +import com.google.tsunami.proto.RunCompactRequest; +import com.google.tsunami.proto.RunCompactRequest.PluginNetworkServiceTarget; +import com.google.tsunami.proto.RunRequest; +import com.google.tsunami.proto.TargetInfo; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class CompactRunRequestHelperTest { + + @Test + public void compressingRunRequest_isMoreCompact() { + NetworkService service1 = NetworkService.newBuilder().setServiceName("service1").build(); + NetworkService service2 = NetworkService.newBuilder().setServiceName("service2").build(); + PluginDefinition plugin1 = + PluginDefinition.newBuilder() + .setInfo(PluginInfo.newBuilder().setName("plugin1").build()) + .build(); + PluginDefinition plugin2 = + PluginDefinition.newBuilder() + .setInfo(PluginInfo.newBuilder().setName("plugin2").build()) + .build(); + PluginDefinition plugin3 = + PluginDefinition.newBuilder() + .setInfo(PluginInfo.newBuilder().setName("plugin3").build()) + .build(); + MatchedPlugin matchedPlugin1 = + MatchedPlugin.newBuilder().addServices(service1).setPlugin(plugin1).build(); + MatchedPlugin matchedPlugin2 = + MatchedPlugin.newBuilder().addServices(service2).setPlugin(plugin2).build(); + MatchedPlugin matchedPlugin3 = + MatchedPlugin.newBuilder().addServices(service1).setPlugin(plugin3).build(); + ImmutableList expectedMatchedPlugins = + ImmutableList.of(matchedPlugin1, matchedPlugin2, matchedPlugin3); + TargetInfo expectedTargetInfo = + TargetInfo.newBuilder() + .addNetworkEndpoints( + NetworkEndpoint.newBuilder() + .setHostname(Hostname.newBuilder().setName("example.com").build()) + .build()) + .build(); + RunRequest expectedUncompressedRunRequest = + RunRequest.newBuilder() + .setTarget(expectedTargetInfo) + .addAllPlugins(expectedMatchedPlugins) + .build(); + var actualCompressedRunRequest = + CompactRunRequestHelper.compress(expectedUncompressedRunRequest); + + var expectedCompressedRunRequest = + RunCompactRequest.newBuilder() + .setTarget(expectedTargetInfo) + .addServices(service1) + .addServices(service2) + .addPlugins(plugin1) + .addPlugins(plugin2) + .addPlugins(plugin3) + .addScanTargets( + PluginNetworkServiceTarget.newBuilder() + .setPluginIndex(0) + .setServiceIndex(0) + .build()) + .addScanTargets( + PluginNetworkServiceTarget.newBuilder() + .setPluginIndex(1) + .setServiceIndex(1) + .build()) + .addScanTargets( + PluginNetworkServiceTarget.newBuilder() + .setPluginIndex(2) + .setServiceIndex(0) + .build()) + .build(); + assertThat(actualCompressedRunRequest).isEqualTo(expectedCompressedRunRequest); + + // And now uncompressing it again: + var actualUncompressedRunRequest = + CompactRunRequestHelper.uncompress(actualCompressedRunRequest); + + // It should match the original setup + assertThat(actualUncompressedRunRequest).isEqualTo(expectedUncompressedRunRequest); + } +} diff --git a/plugin/src/main/java/com/google/tsunami/plugin/PluginServiceClient.java b/plugin/src/main/java/com/google/tsunami/plugin/PluginServiceClient.java index 1af500c2..ee5412d7 100644 --- a/plugin/src/main/java/com/google/tsunami/plugin/PluginServiceClient.java +++ b/plugin/src/main/java/com/google/tsunami/plugin/PluginServiceClient.java @@ -22,6 +22,7 @@ import com.google.tsunami.proto.ListPluginsResponse; import com.google.tsunami.proto.PluginServiceGrpc; import com.google.tsunami.proto.PluginServiceGrpc.PluginServiceFutureStub; +import com.google.tsunami.proto.RunCompactRequest; import com.google.tsunami.proto.RunRequest; import com.google.tsunami.proto.RunResponse; import io.grpc.Channel; @@ -56,6 +57,18 @@ public ListenableFuture runWithDeadline(RunRequest request, Deadlin return pluginService.withDeadline(deadline).run(request); } + /** + * Sends a runCompact request to the gRPC language server with a specified deadline. + * + * @param request The main request containing plugins to run. + * @param deadline The timeout of the service. + * @return The future of the run response. + */ + public ListenableFuture runCompactWithDeadline( + RunCompactRequest request, Deadline deadline) { + return pluginService.withDeadline(deadline).runCompact(request); + } + /** * Sends a list plugins request to the gRPC language server with a specified deadline. * diff --git a/plugin/src/main/java/com/google/tsunami/plugin/RemoteVulnDetectorImpl.java b/plugin/src/main/java/com/google/tsunami/plugin/RemoteVulnDetectorImpl.java index b72b2fc0..6e5cdaed 100644 --- a/plugin/src/main/java/com/google/tsunami/plugin/RemoteVulnDetectorImpl.java +++ b/plugin/src/main/java/com/google/tsunami/plugin/RemoteVulnDetectorImpl.java @@ -24,12 +24,14 @@ import com.google.common.collect.Sets; import com.google.common.flogger.GoogleLogger; import com.google.common.util.concurrent.Uninterruptibles; +import com.google.tsunami.common.server.CompactRunRequestHelper; import com.google.tsunami.proto.DetectionReportList; import com.google.tsunami.proto.ListPluginsRequest; import com.google.tsunami.proto.MatchedPlugin; import com.google.tsunami.proto.NetworkService; import com.google.tsunami.proto.PluginDefinition; import com.google.tsunami.proto.RunRequest; +import com.google.tsunami.proto.RunResponse; import com.google.tsunami.proto.TargetInfo; import io.grpc.Channel; import io.grpc.Deadline; @@ -53,6 +55,7 @@ public final class RemoteVulnDetectorImpl implements RemoteVulnDetector { private final ExponentialBackOff backoff; private final int maxAttempts; private final Deadline deadline; + private boolean wantCompactRunRequest = false; RemoteVulnDetectorImpl( Channel channel, ExponentialBackOff backoff, int maxAttempts, Deadline deadline) { @@ -68,13 +71,17 @@ public DetectionReportList detect( TargetInfo target, ImmutableList matchedServices) { try { if (checkHealthWithBackoffs()) { + var runRequest = + RunRequest.newBuilder().setTarget(target).addAllPlugins(pluginsToRun).build(); logger.atInfo().log("Detecting with language server plugins..."); - return service - .runWithDeadline( - RunRequest.newBuilder().setTarget(target).addAllPlugins(pluginsToRun).build(), - deadline) - .get() - .getReports(); + RunResponse runResponse; + if (this.wantCompactRunRequest) { + var runCompactRequest = CompactRunRequestHelper.compress(runRequest); + runResponse = service.runCompactWithDeadline(runCompactRequest, deadline).get(); + } else { + runResponse = service.runWithDeadline(runRequest, deadline).get(); + } + return runResponse.getReports(); } } catch (InterruptedException | ExecutionException e) { throw new LanguageServerException("Failed to get response from language server.", e); @@ -87,11 +94,14 @@ public ImmutableList getAllPlugins() { try { if (checkHealthWithBackoffs()) { logger.atInfo().log("Getting language server plugins..."); - return ImmutableList.copyOf( + var listPluginsResponse = service .listPluginsWithDeadline(ListPluginsRequest.getDefaultInstance(), DEFAULT_DEADLINE) - .get() - .getPluginsList()); + .get(); + // Note: each plugin service client has a dedicated RemoteVulnDetectorImpl instance, + // so we can safely set this flag here. + this.wantCompactRunRequest = listPluginsResponse.getWantCompactRunRequest(); + return ImmutableList.copyOf(listPluginsResponse.getPluginsList()); } else { return ImmutableList.of(); } diff --git a/plugin/src/test/java/com/google/tsunami/plugin/RemoteVulnDetectorImplTest.java b/plugin/src/test/java/com/google/tsunami/plugin/RemoteVulnDetectorImplTest.java index 50fb59c0..b3ac1f33 100644 --- a/plugin/src/test/java/com/google/tsunami/plugin/RemoteVulnDetectorImplTest.java +++ b/plugin/src/test/java/com/google/tsunami/plugin/RemoteVulnDetectorImplTest.java @@ -32,6 +32,7 @@ import com.google.tsunami.proto.PluginDefinition; import com.google.tsunami.proto.PluginInfo; import com.google.tsunami.proto.PluginServiceGrpc.PluginServiceImplBase; +import com.google.tsunami.proto.RunCompactRequest; import com.google.tsunami.proto.RunRequest; import com.google.tsunami.proto.RunResponse; import com.google.tsunami.proto.TargetInfo; @@ -174,6 +175,62 @@ public void listPlugins( assertThat(pluginToTest.getAllPlugins()).containsExactly(plugin); } + @Test + public void getAllPlugins_withCompactRunRequest_callsRunCompact() throws Exception { + registerHealthCheckWithStatus(ServingStatus.SERVING); + + var targetInfo = + TargetInfo.newBuilder() + .addNetworkEndpoints(NetworkEndpointUtils.forIpAndPort("1.1.1.1", 80)) + .build(); + var someNetworkService = NetworkService.getDefaultInstance(); + var expectedDetectionReport = + DetectionReport.newBuilder() + .setTargetInfo(targetInfo) + .setNetworkService(someNetworkService) + .build(); + + var plugin = createSinglePluginDefinitionWithName("test"); + RemoteVulnDetector pluginToTest = getNewRemoteVulnDetectorInstance(); + serviceRegistry.addService( + new PluginServiceImplBase() { + @Override + public void listPlugins( + ListPluginsRequest request, StreamObserver responseObserver) { + responseObserver.onNext( + ListPluginsResponse.newBuilder() + .setWantCompactRunRequest(true) + .addPlugins(plugin) + .build()); + responseObserver.onCompleted(); + } + + @Override + public void run(RunRequest request, StreamObserver responseObserver) { + responseObserver.onError(new Exception("run should not be called")); + } + + @Override + public void runCompact( + RunCompactRequest request, StreamObserver responseObserver) { + responseObserver.onNext( + RunResponse.newBuilder() + .setReports( + DetectionReportList.newBuilder() + .addDetectionReports(expectedDetectionReport)) + .build()); + responseObserver.onCompleted(); + } + }); + + assertThat(pluginToTest.getAllPlugins()).containsExactly(plugin); + assertThat( + pluginToTest + .detect(targetInfo, ImmutableList.of(someNetworkService)) + .getDetectionReportsList()) + .containsExactly(expectedDetectionReport); + } + @Test public void getAllPlugins_withNonServingServer_returnsEmptyList() throws Exception { registerHealthCheckWithStatus(ServingStatus.NOT_SERVING); diff --git a/proto/plugin_service.proto b/proto/plugin_service.proto index 42215783..2851e957 100644 --- a/proto/plugin_service.proto +++ b/proto/plugin_service.proto @@ -19,10 +19,10 @@ syntax = "proto3"; package tsunami.proto; -import "plugin_representation.proto"; import "detection.proto"; -import "reconnaissance.proto"; import "network_service.proto"; +import "plugin_representation.proto"; +import "reconnaissance.proto"; option java_multiple_files = true; option java_outer_classname = "PluginServiceProtos"; @@ -38,6 +38,29 @@ message RunRequest { repeated MatchedPlugin plugins = 2; } +// Compact representation of RunRequest. +message RunCompactRequest { + // Target of the plugins. + TargetInfo target = 1; + + // Indexes in the following structure point to the services/plugins defined + // below. (The order is safe, guaranteed by the proto specification: "The + // order of the elements with respect to each other is preserved when parsing, + // though the ordering with respect to other fields is lost.") + message PluginNetworkServiceTarget { + // The index of the plugin to run. + uint32 plugin_index = 1; + // The index of the network service to run against. + uint32 service_index = 2; + } + // All network services that are targeted by some of the plugins. + repeated NetworkService services = 2; + // All plugins that should be executed during the run. + repeated PluginDefinition plugins = 3; + // The concrete map of plugin/network service pairs that should be scanned. + repeated PluginNetworkServiceTarget scan_targets = 4; +} + // Represents the plugin needed to run by the language-specific server // as well as all the matched network services for the plugin. message MatchedPlugin { @@ -60,13 +83,24 @@ message ListPluginsRequest {} // from the requested server. message ListPluginsResponse { repeated PluginDefinition plugins = 1; + + // Plugin service can indicate here that it RunRequest should be compact + // (compact_targets should be populated instead of MatchedPlugin plugins). + bool want_compact_run_request = 2; } // Represents the plugin service, two RPCs for running plugins // and listing plugins, respectively. service PluginService { - // Performs a run request to run all language plugins specified by the request. + // Performs a run request to run all language plugins specified by the + // request. rpc Run(RunRequest) returns (RunResponse) {} + // Performs a run request to run all language plugins specified by the + // compact representation of the request. This is useful, when hundreds of + // plugins are to be run against many different NetworkServices. + // The language server must set `want_compact_run_request` so that the + // Tsunami CLI knows to invoke this method instead of `Run`. + rpc RunCompact(RunCompactRequest) returns (RunResponse) {} // Sends a request to list all plugins from the respective language server. rpc ListPlugins(ListPluginsRequest) returns (ListPluginsResponse) {} }