Loadbalance websockets requests

This commit is contained in:
Spencer Gibb
2018-01-23 13:27:15 -05:00
parent 1d4b95fcef
commit 58078c71d1
6 changed files with 176 additions and 10 deletions

View File

@@ -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<Void> 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:<scheme>` mechanism was used, use `<scheme>` 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<String, String> getMetadata() {
return delegate.getMetadata();
}
@Override
public String getScheme() {
String scheme = delegate.getScheme();
if (scheme != null) {
return scheme;
}
return this.overrideScheme;
}
}
}

View File

@@ -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;
}
}

View File

@@ -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

View File

@@ -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<URI> 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<ServerWebExchange> captor = ArgumentCaptor.forClass(ServerWebExchange.class);

View File

@@ -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());

View File

@@ -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<Server> ribbonServerList() {
return new StaticServerList<>(new Server("localhost", this.wsPort));
}
}
}