Fix authentication on IP change

killme2008 · Apr 16, 2018 · 6f157fd · 6f157fd
1 parent b958682
commit 6f157fd
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 2 deletions.
diff --git a/src/main/java/net/rubyeye/xmemcached/MemcachedClient.java b/src/main/java/net/rubyeye/xmemcached/MemcachedClient.java
@@ -1550,6 +1550,13 @@ public KeyIterator getKeyIterator(InetSocketAddress address)
    */
   public Map<InetSocketAddress, AuthInfo> getAuthInfoMap();
 
+  /**
+   * Retruns the AuthInfo for all server strings (hostname:port)
+   *
+   * @return A map of AuthInfos for server strings (hostname:port)
+   */
+  public Map<String, AuthInfo> getAuthInfoStringMap();
+
   /**
    * "incr" are used to change data for some item in-place, incrementing it. The data for the item
    * is treated as decimal representation of a 64-bit unsigned integer. If the current data value

diff --git a/src/main/java/net/rubyeye/xmemcached/XMemcachedClient.java b/src/main/java/net/rubyeye/xmemcached/XMemcachedClient.java
@@ -94,6 +94,7 @@ public class XMemcachedClient implements XMemcachedClientMBean, MemcachedClient
   protected final AtomicInteger serverOrderCount = new AtomicInteger();
 
   private Map<InetSocketAddress, AuthInfo> authInfoMap = new HashMap<InetSocketAddress, AuthInfo>();
+  private Map<String, AuthInfo> authInfoStringMap = new HashMap<String, AuthInfo>();
 
   private String name; // cache name
 
@@ -263,6 +264,15 @@ public Map<InetSocketAddress, AuthInfo> getAuthInfoMap() {
 
   public void setAuthInfoMap(Map<InetSocketAddress, AuthInfo> map) {
     this.authInfoMap = map;
+    this.authInfoStringMap = new HashMap<String, AuthInfo>();
+    for (Map.Entry<InetSocketAddress, AuthInfo> entry : map.entrySet()) {
+      String server = AddrUtil.getServerString(entry.getKey());
+      this.authInfoStringMap.put(server, entry.getValue());
+    }
+  }
+
+  public Map<String, AuthInfo> getAuthInfoStringMap() {
+    return this.authInfoStringMap;
   }
 
   /*

diff --git a/src/main/java/net/rubyeye/xmemcached/auth/AuthMemcachedConnectListener.java b/src/main/java/net/rubyeye/xmemcached/auth/AuthMemcachedConnectListener.java
@@ -7,6 +7,7 @@
 import net.rubyeye.xmemcached.impl.MemcachedTCPSession;
 import net.rubyeye.xmemcached.networking.MemcachedSession;
 import net.rubyeye.xmemcached.networking.MemcachedSessionConnectListener;
+import net.rubyeye.xmemcached.utils.AddrUtil;
 
 /**
  * Client state listener for auth
@@ -18,9 +19,10 @@ public class AuthMemcachedConnectListener implements MemcachedSessionConnectList
 
   public void onConnect(MemcachedSession session, MemcachedClient client) {
     MemcachedTCPSession tcpSession = (MemcachedTCPSession) session;
-    Map<InetSocketAddress, AuthInfo> authInfoMap = client.getAuthInfoMap();
+    Map<String, AuthInfo> authInfoMap = client.getAuthInfoStringMap();
     if (authInfoMap != null) {
-      AuthInfo authInfo = authInfoMap.get(tcpSession.getRemoteSocketAddress());
+      AuthInfo authInfo =
+          authInfoMap.get(AddrUtil.getServerString(tcpSession.getRemoteSocketAddress()));
       if (authInfo != null) {
         XMemcachedClient xMemcachedClient = (XMemcachedClient) client;
         AuthTask task = new AuthTask(authInfo, xMemcachedClient.getCommandFactory(), tcpSession);

diff --git a/src/main/java/net/rubyeye/xmemcached/utils/AddrUtil.java b/src/main/java/net/rubyeye/xmemcached/utils/AddrUtil.java
@@ -137,4 +137,11 @@ public static InetSocketAddress getOneAddress(String server) {
   public static boolean isEnableShutDownHook() {
     return Boolean.valueOf(System.getProperty("xmemcached.shutdown.hook.enable", "false"));
   }
+
+  /**
+   * Create an unresolved server string (hostname:port) from an InetSocketAddress.
+   */
+  public static final String getServerString(InetSocketAddress addr) {
+    return addr.getHostString() + ":" + Integer.toString(addr.getPort());
+  }
 }