From 58078c71d1773172fb9792f40e21998c83065c93 Mon Sep 17 00:00:00 2001 From: Spencer Gibb Date: Tue, 23 Jan 2018 13:27:15 -0500 Subject: [PATCH] Loadbalance websockets requests --- .../filter/LoadBalancerClientFilter.java | 64 ++++++++++++++++++- .../filter/RouteToRequestUrlFilter.java | 20 +++++- .../support/ServerWebExchangeUtils.java | 1 + .../filter/LoadBalancerClientFilterTests.java | 39 ++++++++++- .../filter/RouteToRequestUrlFilterTests.java | 39 +++++++++++ .../websocket/WebSocketIntegrationTests.java | 23 +++++-- 6 files changed, 176 insertions(+), 10 deletions(-) diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/LoadBalancerClientFilter.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/LoadBalancerClientFilter.java index f1f49909..d997611b 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/LoadBalancerClientFilter.java +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/LoadBalancerClientFilter.java @@ -18,6 +18,7 @@ package org.springframework.cloud.gateway.filter; import java.net.URI; +import java.util.Map; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; @@ -28,6 +29,7 @@ import org.springframework.core.Ordered; import org.springframework.web.server.ServerWebExchange; import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR; +import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_SCHEME_PREFIX_ATTR; import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.addOriginalRequestUrl; import reactor.core.publisher.Mono; @@ -55,7 +57,8 @@ public class LoadBalancerClientFilter implements GlobalFilter, Ordered { @Override public Mono filter(ServerWebExchange exchange, GatewayFilterChain chain) { URI url = exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR); - if (url == null || !url.getScheme().equals("lb")) { + String schemePrefix = exchange.getAttribute(GATEWAY_SCHEME_PREFIX_ATTR); + if (url == null || (!"lb".equals(url.getScheme()) && !"lb".equals(schemePrefix))) { return chain.filter(exchange); } //preserve the original url @@ -70,11 +73,68 @@ public class LoadBalancerClientFilter implements GlobalFilter, Ordered { } URI uri = exchange.getRequest().getURI(); - URI requestUrl = loadBalancer.reconstructURI(instance, uri); + + // if the `lb:` mechanism was used, use `` as the default, + // if the loadbalancer doesn't provide one. + String overrideScheme = null; + if (schemePrefix != null) { + overrideScheme = url.getScheme(); + } + + URI requestUrl = loadBalancer.reconstructURI(new DelegatingServiceInstance(instance, overrideScheme), uri); log.trace("LoadBalancerClientFilter url chosen: " + requestUrl); exchange.getAttributes().put(GATEWAY_REQUEST_URL_ATTR, requestUrl); return chain.filter(exchange); } + class DelegatingServiceInstance implements ServiceInstance { + final ServiceInstance delegate; + private String overrideScheme; + + DelegatingServiceInstance(ServiceInstance delegate, String overrideScheme) { + this.delegate = delegate; + this.overrideScheme = overrideScheme; + } + + @Override + public String getServiceId() { + return delegate.getServiceId(); + } + + @Override + public String getHost() { + return delegate.getHost(); + } + + @Override + public int getPort() { + return delegate.getPort(); + } + + @Override + public boolean isSecure() { + return delegate.isSecure(); + } + + @Override + public URI getUri() { + return delegate.getUri(); + } + + @Override + public Map getMetadata() { + return delegate.getMetadata(); + } + + @Override + public String getScheme() { + String scheme = delegate.getScheme(); + if (scheme != null) { + return scheme; + } + return this.overrideScheme; + } + + } } diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/RouteToRequestUrlFilter.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/RouteToRequestUrlFilter.java index b31f1290..361131af 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/RouteToRequestUrlFilter.java +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/RouteToRequestUrlFilter.java @@ -18,6 +18,7 @@ package org.springframework.cloud.gateway.filter; import java.net.URI; +import java.util.regex.Pattern; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; @@ -28,6 +29,7 @@ import org.springframework.web.util.UriComponentsBuilder; import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR; import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR; +import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_SCHEME_PREFIX_ATTR; import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.containsEncodedQuery; import reactor.core.publisher.Mono; @@ -38,7 +40,10 @@ import reactor.core.publisher.Mono; public class RouteToRequestUrlFilter implements GlobalFilter, Ordered { private static final Log log = LogFactory.getLog(RouteToRequestUrlFilter.class); + public static final int ROUTE_TO_URL_FILTER_ORDER = 10000; + private static final String SCHEME_REGEX = "[a-zA-Z]([a-zA-Z]|\\d|\\+|\\.|-)*:.*"; + static final Pattern schemePattern = Pattern.compile(SCHEME_REGEX); @Override public int getOrder() { @@ -54,12 +59,25 @@ public class RouteToRequestUrlFilter implements GlobalFilter, Ordered { log.trace("RouteToRequestUrlFilter start"); URI uri = exchange.getRequest().getURI(); boolean encoded = containsEncodedQuery(uri); + URI routeUri = route.getUri(); + + if (hasAnotherScheme(routeUri)) { + // this is a special url, save scheme to special attribute + // replace routeUri with schemeSpecificPart + exchange.getAttributes().put(GATEWAY_SCHEME_PREFIX_ATTR, routeUri.getScheme()); + routeUri = URI.create(routeUri.getSchemeSpecificPart()); + } + URI requestUrl = UriComponentsBuilder.fromUri(uri) - .uri(route.getUri()) + .uri(routeUri) .build(encoded) .toUri(); exchange.getAttributes().put(GATEWAY_REQUEST_URL_ATTR, requestUrl); return chain.filter(exchange); } + /* for testing */ static boolean hasAnotherScheme(URI uri) { + return schemePattern.matcher(uri.getSchemeSpecificPart()).matches() && uri.getHost() == null + && uri.getRawPath() == null; + } } diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/support/ServerWebExchangeUtils.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/support/ServerWebExchangeUtils.java index 55b94584..3db248a8 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/support/ServerWebExchangeUtils.java +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/support/ServerWebExchangeUtils.java @@ -40,6 +40,7 @@ public class ServerWebExchangeUtils { public static final String GATEWAY_REQUEST_URL_ATTR = qualify("gatewayRequestUrl"); public static final String GATEWAY_ORIGINAL_REQUEST_URL_ATTR = qualify("gatewayOriginalRequestUrl"); public static final String GATEWAY_HANDLER_MAPPER_ATTR = qualify("gatewayHandlerMapper"); + public static final String GATEWAY_SCHEME_PREFIX_ATTR = qualify("gatewaySchemePrefix"); /** * Used when a routing filter has been successfully call. Allows users to write custom diff --git a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/LoadBalancerClientFilterTests.java b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/LoadBalancerClientFilterTests.java index ba80c8a6..acaefd52 100644 --- a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/LoadBalancerClientFilterTests.java +++ b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/LoadBalancerClientFilterTests.java @@ -28,7 +28,6 @@ import org.springframework.web.util.UriComponentsBuilder; import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoMoreInteractions; @@ -36,6 +35,7 @@ import static org.mockito.Mockito.verifyZeroInteractions; import static org.mockito.Mockito.when; import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_ORIGINAL_REQUEST_URL_ATTR; import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR; +import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_SCHEME_PREFIX_ATTR; import reactor.core.publisher.Mono; @@ -109,7 +109,7 @@ public class LoadBalancerClientFilterTests { verify(loadBalancerClient).choose("myservice"); ArgumentCaptor urlArgumentCaptor = ArgumentCaptor.forClass(URI.class); - verify(loadBalancerClient).reconstructURI(eq(serviceInstance), urlArgumentCaptor.capture()); + verify(loadBalancerClient).reconstructURI(any(), urlArgumentCaptor.capture()); URI uri = urlArgumentCaptor.getValue(); assertThat(uri).isNotNull(); @@ -197,8 +197,41 @@ public class LoadBalancerClientFilterTests { assertThat(uri.getRawQuery()).isEqualTo("a=b&c=d[]"); } + @Test + public void happyPathWithAttributeRatherThanScheme() { + MockServerHttpRequest request = MockServerHttpRequest + .get("ws://localhost/get?a=b") + .build(); + + URI lbUri = URI.create("ws://service1?a=b"); + + exchange = MockServerWebExchange.from(request); + exchange.getAttributes().put(GATEWAY_SCHEME_PREFIX_ATTR, "lb"); + + ServerWebExchange webExchange = testFilter(exchange, lbUri); + URI uri = webExchange.getRequiredAttribute(GATEWAY_REQUEST_URL_ATTR); + assertThat(uri).hasScheme("ws").hasHost("service1-host1") + .hasParameter("a", "b"); + } + + @Test + public void shouldNotFilterWhenGatewaySchemePrefixAttrIsNotLb() { + URI uri = UriComponentsBuilder.fromUriString("http://myservice").build().toUri(); + exchange.getAttributes().put(GATEWAY_REQUEST_URL_ATTR, uri); + exchange.getAttributes().put(GATEWAY_SCHEME_PREFIX_ATTR, "xx"); + + loadBalancerClientFilter.filter(exchange, chain); + + verify(chain).filter(exchange); + verifyNoMoreInteractions(chain); + verifyZeroInteractions(loadBalancerClient); + } + private ServerWebExchange testFilter(MockServerHttpRequest request, URI uri) { - ServerWebExchange exchange = MockServerWebExchange.from(request); + return testFilter(MockServerWebExchange.from(request), uri); + } + + private ServerWebExchange testFilter(ServerWebExchange exchange, URI uri) { exchange.getAttributes().put(GATEWAY_REQUEST_URL_ATTR, uri); ArgumentCaptor captor = ArgumentCaptor.forClass(ServerWebExchange.class); diff --git a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/RouteToRequestUrlFilterTests.java b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/RouteToRequestUrlFilterTests.java index 30c7f468..e709a371 100644 --- a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/RouteToRequestUrlFilterTests.java +++ b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/RouteToRequestUrlFilterTests.java @@ -33,6 +33,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR; import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR; +import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_SCHEME_PREFIX_ATTR; import org.springframework.web.util.UriComponentsBuilder; import reactor.core.publisher.Mono; @@ -54,6 +55,19 @@ public class RouteToRequestUrlFilterTests { .hasParameter("a", "b"); } + @Test + public void happyPathLb() { + MockServerHttpRequest request = MockServerHttpRequest + .get("http://localhost/getb") + .build(); + + ServerWebExchange webExchange = testFilter(request, "lb:http://myhost"); + URI uri = webExchange.getRequiredAttribute(GATEWAY_REQUEST_URL_ATTR); + assertThat(uri).hasScheme("http").hasHost("myhost"); + String schemePrefix = webExchange.getRequiredAttribute(GATEWAY_SCHEME_PREFIX_ATTR); + assertThat(schemePrefix).isEqualTo("lb"); + } + @Test public void noQueryParams() { MockServerHttpRequest request = MockServerHttpRequest @@ -110,6 +124,31 @@ public class RouteToRequestUrlFilterTests { assertThat(uri.getRawQuery()).isEqualTo("a=b&c=d[]"); } + @Test + public void matcherWorks() { + testMatcher(true, + "lb:a123:stuff", + "lb:abc:stuff", + "lb:a.bc:stuff", + "lb:a-bc:stuff", + "lb:a+bc:stuff" + ); + testMatcher(false, + "lb:a", + "lb:a123", + "lb:123:stuff", + "lb:a//:stuff" + ); + } + + private void testMatcher(boolean shouldMatch, String... uris) { + for (String s : uris) { + URI uri = URI.create(s); + boolean result = RouteToRequestUrlFilter.hasAnotherScheme(uri); + assertThat(result).as("%s should match: %s", s, result).isEqualTo(shouldMatch); + } + } + private ServerWebExchange testFilter(MockServerHttpRequest request, String url) { Route value = new Route("1", URI.create(url), 0, swe -> true, Collections.emptyList()); diff --git a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/test/websocket/WebSocketIntegrationTests.java b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/test/websocket/WebSocketIntegrationTests.java index b2d9b5c8..d4c90608 100644 --- a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/test/websocket/WebSocketIntegrationTests.java +++ b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/test/websocket/WebSocketIntegrationTests.java @@ -38,6 +38,8 @@ import org.springframework.boot.builder.SpringApplicationBuilder; import org.springframework.cloud.gateway.route.RouteLocator; import org.springframework.cloud.gateway.route.builder.RouteLocatorBuilder; import org.springframework.cloud.gateway.test.PermitAllSecurityConfiguration; +import org.springframework.cloud.netflix.ribbon.RibbonClient; +import org.springframework.cloud.netflix.ribbon.StaticServerList; import org.springframework.context.ConfigurableApplicationContext; import org.springframework.context.Lifecycle; import org.springframework.context.annotation.AnnotationConfigApplicationContext; @@ -64,6 +66,9 @@ import org.springframework.web.reactive.socket.server.support.WebSocketHandlerAd import org.springframework.web.reactive.socket.server.upgrade.ReactorNettyRequestUpgradeStrategy; import org.springframework.web.server.adapter.WebHttpHandlerBuilder; +import com.netflix.loadbalancer.Server; +import com.netflix.loadbalancer.ServerList; + import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertThat; @@ -294,18 +299,28 @@ public class WebSocketIntegrationTests { @Configuration @EnableAutoConfiguration @Import(PermitAllSecurityConfiguration.class) + @RibbonClient(name = "wsservice", configuration = LocalRibbonClientConfiguration.class) protected static class GatewayConfig { - @Value("${ws.server.port}") - private int wsPort; - @Bean public RouteLocator wsRouteLocator(RouteLocatorBuilder builder) { return builder.routes() .route(r -> r.alwaysTrue() - .uri("ws://localhost:"+this.wsPort)) + .uri("lb:ws://wsservice")) .build(); } } + public static class LocalRibbonClientConfiguration { + + @Value("${ws.server.port}") + private int wsPort; + + @Bean + public ServerList ribbonServerList() { + return new StaticServerList<>(new Server("localhost", this.wsPort)); + } + + } + }