Merge pull request Netflix#340 from mattnelson/server_scheme

Support configuring scheme for Server
lurify · Jun 15, 2017 · ec74ad5 · ec74ad5
2 parents 51af081 + 6646ac8
commit ec74ad5
Show file tree

Hide file tree

Showing 5 changed files with 141 additions and 8 deletions.
diff --git a/ribbon-loadbalancer/src/main/java/com/netflix/loadbalancer/ConfigurationBasedServerList.java b/ribbon-loadbalancer/src/main/java/com/netflix/loadbalancer/ConfigurationBasedServerList.java
@@ -60,7 +60,7 @@ public void initWithNiwsConfig(IClientConfig clientConfig) {
 	    this.clientConfig = clientConfig;
 	}
 
-	private List<Server> derive(String value) {
+	protected List<Server> derive(String value) {
 	    List<Server> list = Lists.newArrayList();
 		if (!Strings.isNullOrEmpty(value)) {
 			for (String s: value.split(",")) {

diff --git a/ribbon-loadbalancer/src/main/java/com/netflix/loadbalancer/LoadBalancerContext.java b/ribbon-loadbalancer/src/main/java/com/netflix/loadbalancer/LoadBalancerContext.java
@@ -571,12 +571,17 @@ public Server getServerFromLoadBalancer(@Nullable URI original, @Nullable Object
 
     public URI reconstructURIWithServer(Server server, URI original) {
         String host = server.getHost();
-        int port = server .getPort();
+        int port = server.getPort();
+        String scheme = server.getScheme();
+
         if (host.equals(original.getHost()) 
-                && port == original.getPort()) {
+                && port == original.getPort()
+                && scheme == original.getScheme()) {
             return original;
         }
-        String scheme = original.getScheme();
+        if (scheme == null) {
+            scheme = original.getScheme();
+        }
         if (scheme == null) {
             scheme = deriveSchemeAndPortFromPartialUri(original).first();
         }

diff --git a/ribbon-loadbalancer/src/main/java/com/netflix/loadbalancer/Server.java b/ribbon-loadbalancer/src/main/java/com/netflix/loadbalancer/Server.java
@@ -60,6 +60,7 @@ public static interface MetaInfo {
     public static final String UNKNOWN_ZONE = "UNKNOWN";
     private String host;
     private int port = 80;
+    private String scheme;
     private volatile String id;
     private volatile boolean isAliveFlag;
     private String zone = UNKNOWN_ZONE;
@@ -88,6 +89,11 @@ public String getInstanceId() {
     };
 
     public Server(String host, int port) {
+        this(null, host, port);
+    }
+
+    public Server(String scheme, String host, int port) {
+        this.scheme = scheme;
         this.host = host;
         this.port = port;
         this.id = host + ":" + port;
@@ -126,6 +132,17 @@ static public String normalizeId(String id) {
             return hostPort.first() + ":" + hostPort.second();
         }
     }
+
+    private static String getScheme(String id) {
+        if (id != null) {
+            if (id.toLowerCase().startsWith("http://")) {
+                return "http";
+            } else if (id.toLowerCase().startsWith("https://")) {
+                return "https";
+            }
+        }
+        return null;
+    }
 
     static Pair<String, Integer> getHostPort(String id) {
         if (id != null) {
@@ -169,10 +186,15 @@ public void setId(String id) {
             this.id = hostPort.first() + ":" + hostPort.second();
             this.host = hostPort.first();
             this.port = hostPort.second();
+            this.scheme = getScheme(id);
         } else {
             this.id = null;
         }
     }
+
+    public void setSchemea(String scheme) {
+        this.scheme = scheme;
+    }
 
     public void setPort(int port) {
         this.port = port;
@@ -200,6 +222,10 @@ public String getHost() {
     public int getPort() {
         return port;
     }
+
+    public String getScheme() {
+        return scheme;
+    }
 
     public String getHostPort() {
         return host + ":" + port;

diff --git a/ribbon-loadbalancer/src/test/java/com/netflix/loadbalancer/LoadBalancerContextTest.java b/ribbon-loadbalancer/src/test/java/com/netflix/loadbalancer/LoadBalancerContextTest.java
@@ -17,19 +17,19 @@
  */
 package com.netflix.loadbalancer;
 
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertEquals;
 
 import java.net.URI;
 import java.net.URLEncoder;
 
 import org.junit.Test;
 
 import com.netflix.client.config.IClientConfig;
-import com.netflix.loadbalancer.BaseLoadBalancer;
-import com.netflix.loadbalancer.ILoadBalancer;
-import com.netflix.loadbalancer.Server;
 
 public class LoadBalancerContextTest {
+
+    final static Object httpKey = "http";
+    final static Object httpsKey = "https";
 
     static BaseLoadBalancer lb = new BaseLoadBalancer() {
 
@@ -39,13 +39,50 @@ public Server chooseServer(Object key) {
         }
     };
 
+    static BaseLoadBalancer mixedSchemeLb = new BaseLoadBalancer() {
+
+        @Override
+        public Server chooseServer(Object key) {
+            if (key == httpKey) {
+                return new Server("http://www.example.com:8081");
+            } else if (key == httpsKey) {
+                return new Server("https://www.example.com:8443");
+            }
+
+            return new Server("www.example.com:8080");
+        }
+    };
+
 
     private MyLoadBalancerContext context;
 
     public LoadBalancerContextTest() {
         context = new MyLoadBalancerContext(lb);
     }
 
+    @Test
+    public void testComputeURIWithMixedSchemaLoadBalancer() throws Exception {
+
+        context = new MyLoadBalancerContext(mixedSchemeLb);
+
+        URI request = new URI("/test?abc=xyz");
+
+        // server with no scheme defined
+        Server server = context.getServerFromLoadBalancer(request, null);
+        URI newURI = context.reconstructURIWithServer(server, request);
+        assertEquals("http://www.example.com:8080/test?abc=xyz", newURI.toString());
+
+
+        // server with no scheme defined
+        server = context.getServerFromLoadBalancer(request, httpKey);
+        newURI = context.reconstructURIWithServer(server, request);
+        assertEquals("http://www.example.com:8081/test?abc=xyz", newURI.toString());
+
+        server = context.getServerFromLoadBalancer(request, httpsKey);
+        newURI = context.reconstructURIWithServer(server, request);
+        assertEquals("https://www.example.com:8443/test?abc=xyz", newURI.toString());
+    }
+
     @Test
     public void testComputeFinalUriWithLoadBalancer() throws Exception {
         URI request = new URI("/test?abc=xyz");

diff --git a/ribbon-loadbalancer/src/test/java/com/netflix/loadbalancer/ServerTest.java b/ribbon-loadbalancer/src/test/java/com/netflix/loadbalancer/ServerTest.java
@@ -0,0 +1,65 @@
+package com.netflix.loadbalancer;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+import org.junit.Test;
+
+public class ServerTest {
+
+    @Test
+    public void createSchemeHost() {
+        Server server = new Server("http://netflix.com");
+        assertEquals("http", server.getScheme());
+        assertEquals("netflix.com", server.getHost());
+        assertEquals(80, server.getPort());
+    }
+
+    @Test
+    public void createSchemeHostPort() {
+        Server server = new Server("http://netflix.com:8080");
+        assertEquals("http", server.getScheme());
+        assertEquals("netflix.com", server.getHost());
+        assertEquals(8080, server.getPort());
+    }
+
+    @Test
+    public void createSecureSchemeHost() {
+        Server server = new Server("https://netflix.com");
+        assertEquals("https", server.getScheme());
+        assertEquals("netflix.com", server.getHost());
+        assertEquals(80, server.getPort());
+    }
+
+    @Test
+    public void createSecureSchemeHostPort() {
+        Server server = new Server("https://netflix.com:443");
+        assertEquals("https", server.getScheme());
+        assertEquals("netflix.com", server.getHost());
+        assertEquals(443, server.getPort());
+    }
+
+    @Test
+    public void createSecureSchemeHostPortExplicit() {
+        Server server = new Server("https", "netflix.com", 443);
+        assertEquals("https", server.getScheme());
+        assertEquals("netflix.com", server.getHost());
+        assertEquals(443, server.getPort());
+    }
+
+    @Test
+    public void createHost() {
+        Server server = new Server("netflix.com");
+        assertNull(server.getScheme());
+        assertEquals("netflix.com", server.getHost());
+        assertEquals(80, server.getPort());
+    }
+
+    @Test
+    public void createHostPort() {
+        Server server = new Server("netflix.com:8080");
+        assertNull(server.getScheme());
+        assertEquals("netflix.com", server.getHost());
+        assertEquals(8080, server.getPort());
+    }
+}