INTERNAL: Handle SaslClient on AuthThread

naver · Jan 10, 2025 · 5c13575 · 5c13575
1 parent fa44194
commit 5c13575
Show file tree

Hide file tree

Showing 7 changed files with 31 additions and 64 deletions.
diff --git a/src/main/java/net/spy/memcached/OperationFactory.java b/src/main/java/net/spy/memcached/OperationFactory.java
@@ -18,9 +18,8 @@
 
 import java.util.Collection;
 import java.util.List;
-import java.util.Map;
 
-import javax.security.auth.callback.CallbackHandler;
+import javax.security.sasl.SaslClient;
 
 import net.spy.memcached.collection.Attributes;
 import net.spy.memcached.collection.BTreeFindPosition;
@@ -230,15 +229,12 @@ CASOperation cas(StoreType t, String key, long casId, int flags,
   /**
    * Create a new sasl auth operation.
    */
-  SASLAuthOperation saslAuth(String[] mech, String serverName,
-                             Map<String, ?> props, CallbackHandler cbh, OperationCallback cb);
+  SASLAuthOperation saslAuth(SaslClient sc, OperationCallback cb);
 
   /**
    * Create a new sasl step operation.
    */
-  SASLStepOperation saslStep(String[] mech, byte[] challenge,
-                             String serverName, Map<String, ?> props, CallbackHandler cbh,
-                             OperationCallback cb);
+  SASLStepOperation saslStep(SaslClient sc, byte[] challenge, OperationCallback cb);
 
   /**
    * Set item attributes

diff --git a/src/main/java/net/spy/memcached/auth/AuthThread.java b/src/main/java/net/spy/memcached/auth/AuthThread.java
@@ -4,6 +4,9 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
+import javax.security.sasl.Sasl;
+import javax.security.sasl.SaslClient;
+
 import net.spy.memcached.KeyUtil;
 import net.spy.memcached.MemcachedConnection;
 import net.spy.memcached.MemcachedNode;
@@ -19,13 +22,20 @@ public class AuthThread extends SpyThread {
   private final AuthDescriptor authDescriptor;
   private final OperationFactory opFact;
   private final MemcachedNode node;
+  private final SaslClient sc;
 
   public AuthThread(MemcachedConnection c, OperationFactory o,
                     AuthDescriptor a, MemcachedNode n) {
     conn = c;
     opFact = o;
     authDescriptor = a;
     node = n;
+    try {
+      sc = Sasl.createSaslClient(authDescriptor.getMechs(), null,
+              "memcached", node.getSocketAddress().toString(), null, authDescriptor.getCallback());
+    } catch (Exception e) {
+      throw new RuntimeException("Can't create SaslClient", e);
+    }
   }
 
   @Override
@@ -89,15 +99,9 @@ public void complete() {
 
   private Operation buildOperation(OperationStatus st, OperationCallback cb) {
     if (st == null) {
-      return opFact.saslAuth(authDescriptor.getMechs(),
-              node.getSocketAddress().toString(), null,
-              authDescriptor.getCallback(), cb);
+      return opFact.saslAuth(sc, cb);
     } else {
-      return opFact.saslStep(authDescriptor.getMechs(),
-              KeyUtil.getKeyBytes(st.getMessage()),
-              node.getSocketAddress().toString(), null,
-              authDescriptor.getCallback(), cb);
+      return opFact.saslStep(sc, KeyUtil.getKeyBytes(st.getMessage()), cb);
     }
-
   }
 }
diff --git a/src/main/java/net/spy/memcached/protocol/ascii/AsciiOperationFactory.java b/src/main/java/net/spy/memcached/protocol/ascii/AsciiOperationFactory.java
@@ -17,9 +17,8 @@
 package net.spy.memcached.protocol.ascii;
 
 import java.util.Collection;
-import java.util.Map;
 
-import javax.security.auth.callback.CallbackHandler;
+import javax.security.sasl.SaslClient;
 
 import net.spy.memcached.collection.Attributes;
 import net.spy.memcached.collection.BTreeFindPosition;
@@ -147,15 +146,11 @@ public SASLMechsOperation saslMechs(OperationCallback cb) {
     throw new UnsupportedOperationException();
   }
 
-  public SASLStepOperation saslStep(String[] mech, byte[] challenge,
-                                    String serverName, Map<String, ?> props, CallbackHandler cbh,
-                                    OperationCallback cb) {
+  public SASLStepOperation saslStep(SaslClient sc, byte[] challenge, OperationCallback cb) {
     throw new UnsupportedOperationException();
   }
 
-  public SASLAuthOperation saslAuth(String[] mech, String serverName,
-                                    Map<String, ?> props, CallbackHandler cbh,
-                                    OperationCallback cb) {
+  public SASLAuthOperation saslAuth(SaslClient sc, OperationCallback cb) {
     throw new UnsupportedOperationException();
   }
 

diff --git a/src/main/java/net/spy/memcached/protocol/binary/BinaryOperationFactory.java b/src/main/java/net/spy/memcached/protocol/binary/BinaryOperationFactory.java
@@ -17,9 +17,8 @@
 package net.spy.memcached.protocol.binary;
 
 import java.util.Collection;
-import java.util.Map;
 
-import javax.security.auth.callback.CallbackHandler;
+import javax.security.sasl.SaslClient;
 
 import net.spy.memcached.collection.Attributes;
 import net.spy.memcached.collection.BTreeFindPosition;
@@ -147,21 +146,16 @@ public ConcatenationOperation cat(ConcatenationType catType, long casId,
     return new ConcatenationOperationImpl(catType, key, data, casId, cb);
   }
 
-  public SASLAuthOperation saslAuth(String[] mech, String serverName,
-                                    Map<String, ?> props, CallbackHandler cbh,
-                                    OperationCallback cb) {
-    return new SASLAuthOperationImpl(mech, serverName, props, cbh, cb);
+  public SASLAuthOperation saslAuth(SaslClient sc, OperationCallback cb) {
+    return new SASLAuthOperationImpl(sc, cb);
   }
 
   public SASLMechsOperation saslMechs(OperationCallback cb) {
     return new SASLMechsOperationImpl(cb);
   }
 
-  public SASLStepOperation saslStep(String[] mech, byte[] challenge,
-                                    String serverName, Map<String, ?> props, CallbackHandler cbh,
-                                    OperationCallback cb) {
-    return new SASLStepOperationImpl(mech, challenge, serverName,
-            props, cbh, cb);
+  public SASLStepOperation saslStep(SaslClient sc, byte[] challenge, OperationCallback cb) {
+    return new SASLStepOperationImpl(sc, challenge, cb);
   }
 
   //// UNSUPPORTED ////

diff --git a/src/main/java/net/spy/memcached/protocol/binary/SASLAuthOperationImpl.java b/src/main/java/net/spy/memcached/protocol/binary/SASLAuthOperationImpl.java
@@ -17,9 +17,6 @@
  */
 package net.spy.memcached.protocol.binary;
 
-import java.util.Map;
-
-import javax.security.auth.callback.CallbackHandler;
 import javax.security.sasl.SaslClient;
 import javax.security.sasl.SaslException;
 
@@ -31,9 +28,8 @@ public class SASLAuthOperationImpl extends SASLBaseOperationImpl
 
   private final static int CMD = 0x21;
 
-  public SASLAuthOperationImpl(String[] m, String s,
-                               Map<String, ?> p, CallbackHandler h, OperationCallback c) {
-    super(CMD, m, EMPTY_BYTES, s, p, h, c);
+  public SASLAuthOperationImpl(SaslClient sc, OperationCallback cb) {
+    super(CMD, sc, EMPTY_BYTES, cb);
   }
 
   @Override

diff --git a/src/main/java/net/spy/memcached/protocol/binary/SASLBaseOperationImpl.java b/src/main/java/net/spy/memcached/protocol/binary/SASLBaseOperationImpl.java
@@ -1,10 +1,6 @@
 package net.spy.memcached.protocol.binary;
 
 import java.io.IOException;
-import java.util.Map;
-
-import javax.security.auth.callback.CallbackHandler;
-import javax.security.sasl.Sasl;
 import javax.security.sasl.SaslClient;
 import javax.security.sasl.SaslException;
 
@@ -17,29 +13,19 @@ public abstract class SASLBaseOperationImpl extends OperationImpl {
 
   private static final int SASL_CONTINUE = 0x21;
 
-  protected final String[] mech;
+  protected final SaslClient sc;
   protected final byte[] challenge;
-  protected final String serverName;
-  protected final Map<String, ?> props;
-  protected final CallbackHandler cbh;
 
-  public SASLBaseOperationImpl(int c, String[] m, byte[] ch,
-                               String s, Map<String, ?> p, CallbackHandler h,
+  public SASLBaseOperationImpl(int c, SaslClient sc, byte[] challenge,
                                OperationCallback cb) {
     super(c, generateOpaque(), cb);
-    mech = m;
-    challenge = ch;
-    serverName = s;
-    props = p;
-    cbh = h;
+    this.sc = sc;
+    this.challenge = challenge;
   }
 
   @Override
   public void initialize() {
     try {
-      SaslClient sc = Sasl.createSaslClient(mech, null,
-              "memcached", serverName, props, cbh);
-
       byte[] response = buildResponse(sc);
       String mechanism = sc.getMechanismName();
 

diff --git a/src/main/java/net/spy/memcached/protocol/binary/SASLStepOperationImpl.java b/src/main/java/net/spy/memcached/protocol/binary/SASLStepOperationImpl.java
@@ -17,9 +17,6 @@
  */
 package net.spy.memcached.protocol.binary;
 
-import java.util.Map;
-
-import javax.security.auth.callback.CallbackHandler;
 import javax.security.sasl.SaslClient;
 import javax.security.sasl.SaslException;
 
@@ -31,9 +28,8 @@ public class SASLStepOperationImpl extends SASLBaseOperationImpl
 
   private final static int CMD = 0x22;
 
-  public SASLStepOperationImpl(String[] m, byte[] ch, String s,
-                               Map<String, ?> p, CallbackHandler h, OperationCallback c) {
-    super(CMD, m, ch, s, p, h, c);
+  public SASLStepOperationImpl(SaslClient sc, byte[] challenge, OperationCallback cb) {
+    super(CMD, sc, challenge, cb);
   }
 
   @Override