From 6f157fd7f4d8d75f4e0626879ac106e3b3e8b97b Mon Sep 17 00:00:00 2001
From: Sascha <sascha.trifunovic@gmail.com>
Date: Mon, 16 Apr 2018 12:28:42 +0200
Subject: [PATCH] Fix authentication on IP change

---
 .../java/net/rubyeye/xmemcached/MemcachedClient.java   |  7 +++++++
 .../java/net/rubyeye/xmemcached/XMemcachedClient.java  | 10 ++++++++++
 .../xmemcached/auth/AuthMemcachedConnectListener.java  |  6 ++++--
 .../java/net/rubyeye/xmemcached/utils/AddrUtil.java    |  7 +++++++
 4 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/main/java/net/rubyeye/xmemcached/MemcachedClient.java b/src/main/java/net/rubyeye/xmemcached/MemcachedClient.java
index b46e48933..31728d198 100644
--- 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
index 58222a2d5..d4bf846f6 100644
--- 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
index 27182f017..12df72d5e 100644
--- 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
index 954a9d7c8..25e5c885a 100644
--- 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());
+  }
 }