Use exact path matchers in the omnibus

This commit is contained in:
Ravi Khadiwala
2026-08-26 14:56:59 -05:00
committed by ravi-signal
parent 461fce800a
commit 56a91f104e
3 changed files with 32 additions and 39 deletions
@@ -1182,11 +1182,14 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
unauthenticatedServices.forEach(serverBuilder::addService);
final ManagedGrpcServer localGrpcServer = new ManagedGrpcServer(serverBuilder.build());
final String websocketServletPath = "/v1/websocket/";
final String provisioningWebsocketServletPath = "/v1/websocket/provisioning/";
final SocketAddress websocketAddress =
new InetSocketAddress(config.getGrpc().websocketAddress(), config.getGrpc().websocketPort());
final OmnibusRouter omnibusRouter = new OmnibusRouter(List.of(
new OmnibusRouter.OmnibusRoute("/v1/websocket", websocketAddress),
new OmnibusRouter.OmnibusRoute("/v1/provisioning", websocketAddress)),
final OmnibusRouter omnibusRouter = new OmnibusRouter(Map.of(
websocketServletPath, websocketAddress,
provisioningWebsocketServletPath, websocketAddress),
grpcLocalAddress);
@Nullable final Mapping<String, SslContext> sniMapping = config.getGrpc().h2c()
? null
@@ -1217,9 +1220,6 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
.setAuthenticator(accountAuthenticator)
.buildAuthFilter();
final String websocketServletPath = "/v1/websocket/";
final String provisioningWebsocketServletPath = "/v1/websocket/provisioning/";
MetricsHttpEventHandler.configure(environment, Metrics.globalRegistry, clientReleaseManager, Set.of(websocketServletPath, provisioningWebsocketServletPath, "/health-check"));
// BufferingInterceptor is needed on the base environment but not the WebSocketEnvironment,
@@ -5,26 +5,21 @@
package org.whispersystems.textsecuregcm.grpc.net;
import java.net.SocketAddress;
import java.util.List;
import java.util.Map;
public class OmnibusRouter {
public record OmnibusRoute(String prefix, SocketAddress backend) {}
private final List<OmnibusRoute> prefixRoutes;
private final Map<String, SocketAddress> routes;
private final SocketAddress defaultBackend;
public OmnibusRouter(final List<OmnibusRoute> prefixRoutes, final SocketAddress defaultBackend) {
this.prefixRoutes = prefixRoutes;
public OmnibusRouter(final Map<String, SocketAddress> routes, final SocketAddress defaultBackend) {
this.routes = routes;
this.defaultBackend = defaultBackend;
}
SocketAddress match(final String path) {
for (final OmnibusRoute route : prefixRoutes) {
if (path.startsWith(route.prefix)) {
return route.backend;
}
}
return defaultBackend;
SocketAddress match(final String fullPath) {
final int queryIndex = fullPath.indexOf('?');
final String path = queryIndex >= 0 ? fullPath.substring(0, queryIndex) : fullPath;
return routes.getOrDefault(path, defaultBackend);
}
}
@@ -55,7 +55,6 @@ import io.netty.util.ReferenceCountUtil;
import io.netty.util.test.LeakPresenceExtension;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.InputStream;
import java.math.BigDecimal;
import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
@@ -71,6 +70,7 @@ import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import javax.annotation.Nullable;
import javax.net.ssl.SSLException;
@@ -88,9 +88,9 @@ import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicOmnibusConf
@ExtendWith(LeakPresenceExtension.class)
class OmnibusH2ServerTest {
// Paths that start with PREFIX should go to the prefix backend, everything else to default.
private static final String PREFIX_BACKEND_IDENTITY = "prefix-backend";
private static final String PREFIX = "/v1/prefix";
// URIs targeting PATH should go to the path backend, everything else to default.
private static final String PATH_BACKEND_IDENTITY = "path-backend";
private static final String PATH = "/v1/path/";
private static final String DEFAULT_BACKEND_IDENTITY = "default-backend";
private static NioEventLoopGroup nioEventLoopGroup;
@@ -152,7 +152,7 @@ class OmnibusH2ServerTest {
@ValueSource(booleans = {true, false})
void defaultBackend(final boolean localChannel) throws Exception {
final OmnibusH2Server server = startOmnibusServer(
Map.of(PREFIX, startBackendServer(localChannel, PREFIX_BACKEND_IDENTITY)),
Map.of(PATH, startBackendServer(localChannel, PATH_BACKEND_IDENTITY)),
startBackendServer(localChannel, DEFAULT_BACKEND_IDENTITY));
final String response = sendRequestThroughOmnibus(connectToOmnibus(server), "/a/different/path");
assertEquals(DEFAULT_BACKEND_IDENTITY, response);
@@ -172,14 +172,14 @@ class OmnibusH2ServerTest {
}
@CartesianTest
void prefixBackend(
@CartesianTest.Values(booleans = {true, false}) final boolean localChannel,
@CartesianTest.Values(strings = {"/v1/prefix", "/v1/prefix/", "/v1/prefix/other"}) final String path) throws Exception {
void pathBackend(
@CartesianTest.Values(booleans = {false, true}) final boolean localChannel,
@CartesianTest.Values(strings = {"/v1/path/", "/v1/path/?foo=bar", "/v1/path/?foo=bar&baz=quux"}) final String path) throws Exception {
final OmnibusH2Server server = startOmnibusServer(
Map.of(PREFIX, startBackendServer(localChannel, PREFIX_BACKEND_IDENTITY)),
Map.of(PATH, startBackendServer(localChannel, PATH_BACKEND_IDENTITY)),
startBackendServer(true, DEFAULT_BACKEND_IDENTITY));
final String response = sendRequestThroughOmnibus(connectToOmnibus(server), path);
assertEquals(PREFIX_BACKEND_IDENTITY, response);
assertEquals(PATH_BACKEND_IDENTITY, response);
}
@CartesianTest
@@ -187,7 +187,7 @@ class OmnibusH2ServerTest {
@CartesianTest.Values(booleans = {true, false}) final boolean defaultLocalChannel,
@CartesianTest.Values(booleans = {true, false}) final boolean prefixLocalChannel) throws Exception {
final OmnibusH2Server server = startOmnibusServer(
Map.of(PREFIX, startBackendServer(prefixLocalChannel, PREFIX_BACKEND_IDENTITY)),
Map.of(PATH, startBackendServer(prefixLocalChannel, PATH_BACKEND_IDENTITY)),
startBackendServer(defaultLocalChannel, DEFAULT_BACKEND_IDENTITY));
final Channel h2Connection = connectToOmnibus(server);
final int numStreams = 10;
@@ -198,14 +198,14 @@ class OmnibusH2ServerTest {
.mapToObj(i -> CompletableFuture.supplyAsync(() ->
sendRequestThroughOmnibus(
h2Connection,
i % 2 == 0 ? PREFIX : "/v1/other")))
i % 2 == 0 ? PATH : "/v1/other")))
.toArray(CompletableFuture[]::new);
// Ensure we get the response from the correct backend for each stream
CompletableFuture.allOf(futures).join();
for (int i = 0; i < numStreams; i++) {
assertEquals(
i % 2 == 0 ? PREFIX_BACKEND_IDENTITY : DEFAULT_BACKEND_IDENTITY,
i % 2 == 0 ? PATH_BACKEND_IDENTITY : DEFAULT_BACKEND_IDENTITY,
futures[i].resultNow());
}
}
@@ -438,7 +438,7 @@ class OmnibusH2ServerTest {
void loadShed() throws Exception {
dynamicConfiguration.set(new DynamicOmnibusConfiguration(BigDecimal.ONE));
final OmnibusH2Server server = startOmnibusServer(
Map.of(PREFIX, startBackendServer(true, PREFIX_BACKEND_IDENTITY)),
Map.of(PATH, startBackendServer(true, PATH_BACKEND_IDENTITY)),
startBackendServer(true, DEFAULT_BACKEND_IDENTITY));
final CompletableFuture<Http2GoAwayFrame> goAwayFuture = new CompletableFuture<>();
@@ -475,9 +475,6 @@ class OmnibusH2ServerTest {
/// @param defaultBackend The target backend if no prefix routes match the request path
/// @param timeout The omnibus idle timeout
private OmnibusH2Server startOmnibusServer(final Map<String, Channel> routes, final Channel defaultBackend, final Duration timeout) throws Exception {
// self-signed TLS context for the frontend loaded from test keyStore
final InputStream keyStore = OmnibusH2ServerTest.class.getResourceAsStream("omnibus-h2-server-test-keystore.p12");
backendChannelsToShutDown.addAll(routes.values());
backendChannelsToShutDown.add(defaultBackend);
@@ -486,8 +483,9 @@ class OmnibusH2ServerTest {
nioEventLoopGroup,
localEventLoopGroup,
new InetSocketAddress("127.0.0.1", 0),
new OmnibusRouter(
routes.entrySet().stream().map(entry -> new OmnibusRouter.OmnibusRoute(entry.getKey(), entry.getValue().localAddress())).toList(),
new OmnibusRouter(routes.entrySet().stream().collect(Collectors.toMap(
Map.Entry::getKey,
e -> e.getValue().localAddress())),
defaultBackend.localAddress()),
dynamicConfiguration::get,
timeout);
@@ -538,7 +536,7 @@ class OmnibusH2ServerTest {
return startBackendServer(localChannel, identity, _ -> {}, _ -> {});
}
/// Makes an H2 connection to the omnibus at [this#server] on which new H2 streams can be opened
/// Makes an H2 connection to the omnibus at `server` on which new H2 streams can be opened
private Channel connectToOmnibus(final OmnibusH2Server server, @Nullable final HAProxyMessage proxyHeader) {
final SslContext clientSsl = sslContext();