1 files changed, 36 insertions, 17 deletions

diff --git a/src/main/java/com/btr/proxy/selector/misc/ProtocolDispatchSelector.java b/src/main/java/com/btr/proxy/selector/misc/ProtocolDispatchSelector.java
index 02ecc44..bb57452 100644
--- a/src/main/java/com/btr/proxy/selector/misc/ProtocolDispatchSelector.java
+++ b/src/main/java/com/btr/proxy/selector/misc/ProtocolDispatchSelector.java
@@ -10,6 +10,9 @@ import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
 import com.btr.proxy.selector.direct.NoProxySelector;
+import com.btr.proxy.selector.fixed.FixedProxySelector;
+import com.btr.proxy.util.Logger;
+import com.btr.proxy.util.Logger.LogLevel;
 
 /*****************************************************************************
  * This is a facade for a list of ProxySelecor objects. You can register 
@@ -46,6 +49,9 @@ public class ProtocolDispatchSelector extends ProxySelector {
 		if (selector == null) {
 			throw new NullPointerException("Selector must not be null.");
 		}
+		if (protocol.toLowerCase().startsWith("socks")) {
+			protocol = "socket";
+		}
 		this.selectors.put(protocol, selector);
 	}
 	
@@ -82,6 +88,17 @@ public class ProtocolDispatchSelector extends ProxySelector {
 		this.fallbackSelector = selector;
 	}
 	
+	private ProxySelector selectorForUri(URI uri) {
+		if (uri == null || uri.getScheme() == null)
+			return this.fallbackSelector;
+		String protocol = uri.getScheme();
+		ProxySelector selector = this.selectors.get(protocol);
+		if (selector == null) {
+			selector = this.fallbackSelector;
+		}
+		return selector;
+	}
+	
 	/*************************************************************************
 	 * connectFailed
 	 * @see java.net.ProxySelector#connectFailed(java.net.URI, java.net.SocketAddress, java.io.IOException)
@@ -89,11 +106,7 @@ public class ProtocolDispatchSelector extends ProxySelector {
 
 	@Override
 	public void connectFailed(URI uri, SocketAddress sa, IOException ioe) {
-		ProxySelector selector = this.fallbackSelector;
-		String protocol = uri.getScheme();
-		if (protocol != null && this.selectors.get(protocol) != null) {
-			selector = this.selectors.get(protocol);
-		}
+		ProxySelector selector = selectorForUri(uri);
 		selector.connectFailed(uri, sa, ioe);
 	}
 
@@ -104,19 +117,25 @@ public class ProtocolDispatchSelector extends ProxySelector {
 
 	@Override
 	public List<Proxy> select(URI uri) {
-		ProxySelector selector = null;
-		String protocol = uri.getScheme();
-		if (protocol != null && this.selectors.get(protocol) != null) {
-			selector = this.selectors.get(protocol);
-		}
-		if (selector == null && this.selectors.get("socks") != null) {
-			// Socks should always work
-			selector = this.selectors.get("socks");
-		}
-		if (selector == null) {
-			selector = this.fallbackSelector;
+		ProxySelector selector = selectorForUri(uri);
+		List<Proxy> ret = selector.select(uri);
+		Logger.log(getClass(), LogLevel.TRACE, "Selector {0} for {1} -> {2}", selector.getClass(), uri, ret);
+		return ret;
+	}
+
+	public void setFallbackSocksSelector(FixedProxySelector... pslist) {
+		if (!(this.fallbackSelector instanceof NoProxySelector))
+			return;
+		for (FixedProxySelector ps : pslist) {
+			if (ps != null && ps.getType() == Proxy.Type.SOCKS) {
+				this.fallbackSelector = ps;
+				return;
+			}
 		}
-		return selector.select(uri);
+	}
+
+	public boolean isEmpty() {
+		return this.selectors.isEmpty() && this.fallbackSelector instanceof NoProxySelector;
 	}
 
 }