Skip to content

Commit ab2cc7e

Browse files
menglimengli
committed
Add SASL Support
Change-Id: I0cccf3e526b3088a24c4c616ef5a12acb1195904
1 parent d7281f6 commit ab2cc7e

File tree

10 files changed

+304
-22
lines changed

10 files changed

+304
-22
lines changed

‎src/main/com/danga/MemCached/MemCachedClient.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -250,6 +250,12 @@ public class MemCachedClient {
250250
public static final byte OPCODE_APPEND = (byte) 0x0E;
251251
public static final byte OPCODE_PREPEND = (byte) 0x0F;
252252
public static final byte OPCODE_STAT = (byte) 0x10;
253+
public static final byte OPCODE_AUTH_LIST = (byte) 0x20;
254+
public static final byte OPCODE_START_AUTH = (byte) 0x21;
255+
public static final byte OPCODE_AUTH_STEPS = (byte) 0x22;
256+
257+
public static final byte AUTH_FAILED = 0x20;
258+
public static final byte FURTHER_AUTH = 0x21;
253259

254260
/**
255261
* This is used for set command only.<br>

‎src/main/com/schooner/MemCached/AscIIClient.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1072,8 +1072,8 @@ public Map<String, Object> getMulti(String[] keys, Integer[] hashCodes, boolean
10721072
sock.close();
10731073
}
10741074

1075-
if (log.isInfoEnabled())
1076-
log.info("multi get socket count : " + cmdMap.size());
1075+
if (log.isDebugEnabled())
1076+
log.debug("multi get socket count : " + cmdMap.size());
10771077

10781078
// now query memcache
10791079
Map<String, Object> ret = new HashMap<String, Object>(keys.length);

‎src/main/com/schooner/MemCached/AuthInfo.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
package com.schooner.MemCached;
2+
3+
import javax.security.auth.callback.CallbackHandler;
4+
5+
public class AuthInfo {
6+
7+
private final CallbackHandler callbackHandler;
8+
private final String[] mechanisms;
9+
10+
public AuthInfo(CallbackHandler callbackHandler, String[] mechanisms) {
11+
super();
12+
this.callbackHandler = callbackHandler;
13+
this.mechanisms = mechanisms;
14+
}
15+
16+
public static AuthInfo plain(String username, String password) {
17+
return new AuthInfo(new PlainCallbackHandler(username, password), new String[] { "PLAIN" });
18+
}
19+
20+
public static AuthInfo cramMD5(String username, String password) {
21+
return new AuthInfo(new PlainCallbackHandler(username, password), new String[] { "CRAM-MD5" });
22+
}
23+
24+
public static AuthInfo typical(String username, String password) {
25+
return new AuthInfo(new PlainCallbackHandler(username, password), new String[] { "CRAM-MD5", "PLAIN" });
26+
}
27+
28+
public CallbackHandler getCallbackHandler() {
29+
return callbackHandler;
30+
}
31+
32+
public String[] getMechanisms() {
33+
return mechanisms;
34+
}
35+
36+
}

‎src/main/com/schooner/MemCached/AuthSchoonerSockIOFactory.java‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
package com.schooner.MemCached;
2+
3+
import java.io.DataInputStream;
4+
5+
import javax.security.sasl.Sasl;
6+
import javax.security.sasl.SaslClient;
7+
8+
import org.slf4j.Logger;
9+
import org.slf4j.LoggerFactory;
10+
11+
import com.danga.MemCached.MemCachedClient;
12+
13+
public class AuthSchoonerSockIOFactory extends SchoonerSockIOFactory {
14+
15+
// logger
16+
public static Logger log = LoggerFactory.getLogger(AuthSchoonerSockIOFactory.class);
17+
18+
public final static String NTLM = "NTLM";
19+
public final static String PLAIN = "PLAIN";
20+
public final static String LOGIN = "LOGIN";
21+
public final static String DIGEST_MD5 = "DIGEST-MD5";
22+
public final static String CRAM_MD5 = "CRAM-MD5";
23+
public final static String ANONYMOUS = "ANONYMOUS";
24+
25+
public static final byte[] EMPTY_BYTES = new byte[0];
26+
27+
private AuthInfo authInfo;
28+
29+
public AuthSchoonerSockIOFactory(String host, boolean isTcp, int bufferSize, int socketTO, int socketConnectTO,
30+
boolean nagle, AuthInfo authInfo) {
31+
super(host, isTcp, bufferSize, socketTO, socketConnectTO, nagle);
32+
this.authInfo = authInfo;
33+
}
34+
35+
@Override
36+
public Object makeObject() throws Exception {
37+
SchoonerSockIO socket = createSocket(host);
38+
auth(socket);
39+
return socket;
40+
}
41+
42+
private void auth(SchoonerSockIO socket) throws Exception {
43+
SaslClient saslClient = Sasl.createSaslClient(authInfo.getMechanisms(), null, "memcached", host, null,
44+
this.authInfo.getCallbackHandler());
45+
46+
byte[] authData = saslClient.hasInitialResponse() ? saslClient.evaluateChallenge(EMPTY_BYTES) : EMPTY_BYTES;
47+
48+
authData = sendAuthData(socket, MemCachedClient.OPCODE_START_AUTH, saslClient.getMechanismName(), authData);
49+
if (authData == null)
50+
return;
51+
authData = saslClient.evaluateChallenge(authData);
52+
if (sendAuthData(socket, MemCachedClient.OPCODE_AUTH_STEPS, saslClient.getMechanismName(), authData) == null)
53+
return;
54+
55+
log.error("Auth Failed: mechanism = " + saslClient.getMechanismName());
56+
throw new Exception();
57+
}
58+
59+
private byte[] sendAuthData(SchoonerSockIO sock, byte opcode, String mechanism, byte[] authData) throws Exception {
60+
sock.writeBuf.clear();
61+
sock.writeBuf.put(MemCachedClient.MAGIC_REQ);
62+
sock.writeBuf.put(opcode);
63+
sock.writeBuf.putShort((short) mechanism.length());
64+
sock.writeBuf.putInt(0);
65+
sock.writeBuf.putInt(mechanism.length() + authData.length);
66+
sock.writeBuf.putInt(0);
67+
sock.writeBuf.putLong(0);
68+
sock.writeBuf.put(mechanism.getBytes());
69+
sock.writeBuf.put(authData);
70+
71+
// write the buffer to server
72+
// now write the data to the cache server
73+
sock.flush();
74+
// get result code
75+
DataInputStream dis = new DataInputStream(new SockInputStream(sock, Integer.MAX_VALUE));
76+
dis.readInt();
77+
dis.readByte();
78+
dis.readByte();
79+
byte[] response = null;
80+
short status = dis.readShort();
81+
if (status == MemCachedClient.FURTHER_AUTH) {
82+
int length = dis.readInt();
83+
response = new byte[length];
84+
dis.readInt();
85+
dis.readLong();
86+
dis.read(response);
87+
} else if (status == MemCachedClient.AUTH_FAILED) {
88+
log.error("Auth Failed: mechanism = " + mechanism);
89+
throw new Exception();
90+
}
91+
92+
return response;
93+
}
94+
}

‎src/main/com/schooner/MemCached/BinaryClient.java‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -182,14 +182,14 @@ public boolean delete(String key, Integer hashCode, Date expiry) {
182182
short status = dis.readShort();
183183

184184
if (status == STAT_NO_ERROR) {
185-
if (log.isInfoEnabled())
186-
log.info("++++ deletion of key: " + key + " from cache was a success");
185+
if (log.isDebugEnabled())
186+
log.debug("++++ deletion of key: " + key + " from cache was a success");
187187

188188
// return sock to pool and bail here
189189
return true;
190190
} else if (status == STAT_KEY_NOT_FOUND) {
191-
if (log.isInfoEnabled())
192-
log.info("++++ deletion of key: " + key + " from cache failed as the key was not found");
191+
if (log.isDebugEnabled())
192+
log.debug("++++ deletion of key: " + key + " from cache failed as the key was not found");
193193
} else {
194194
log.error("++++ error deleting key: " + key);
195195
log.error("++++ server response: " + status);
@@ -847,7 +847,8 @@ public Map<String, Object> getMulti(String[] keys, Integer[] hashCodes, boolean
847847
sock.close();
848848
}
849849

850-
log.info("multi get socket count : " + cmdMap.size());
850+
if (log.isDebugEnabled())
851+
log.debug("multi get socket count : " + cmdMap.size());
851852

852853
// now query memcache
853854
Map<String, Object> ret = new HashMap<String, Object>(keys.length);

‎src/main/com/schooner/MemCached/PlainCallbackHandler.java‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
package com.schooner.MemCached;
2+
3+
import java.io.IOException;
4+
5+
import javax.security.auth.callback.Callback;
6+
import javax.security.auth.callback.CallbackHandler;
7+
import javax.security.auth.callback.NameCallback;
8+
import javax.security.auth.callback.PasswordCallback;
9+
import javax.security.auth.callback.UnsupportedCallbackException;
10+
11+
public class PlainCallbackHandler implements CallbackHandler {
12+
private String username;
13+
private String password;
14+
15+
public PlainCallbackHandler(String username, String password) {
16+
super();
17+
this.username = username;
18+
this.password = password;
19+
}
20+
21+
public void handle(Callback[] callbacks) throws IOException, UnsupportedCallbackException {
22+
for (Callback callback : callbacks) {
23+
if (callback instanceof NameCallback) {
24+
((NameCallback) callback).setName(this.username);
25+
} else if (callback instanceof PasswordCallback) {
26+
((PasswordCallback) callback).setPassword(password.toCharArray());
27+
} else
28+
throw new UnsupportedCallbackException(callback);
29+
}
30+
31+
}
32+
33+
}

‎src/main/com/schooner/MemCached/SchoonerSockIOFactory.java‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,19 +13,19 @@
1313
*/
1414
public class SchoonerSockIOFactory extends BasePoolableObjectFactory {
1515

16-
private GenericObjectPool sockets;
16+
protected GenericObjectPool sockets;
1717

18-
private String host;
18+
protected String host;
1919

20-
private int bufferSize;
20+
protected int bufferSize;
2121

22-
private int socketTO;
22+
protected int socketTO;
2323

24-
private int socketConnectTO;
24+
protected int socketConnectTO;
2525

26-
private boolean isTcp;
26+
protected boolean isTcp;
2727

28-
private boolean nagle;
28+
protected boolean nagle;
2929

3030
public SchoonerSockIOFactory(String host, boolean isTcp, int bufferSize, int socketTO, int socketConnectTO,
3131
boolean nagle) {

‎src/main/com/schooner/MemCached/SchoonerSockIOPool.java‎

Lines changed: 34 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,8 @@ protected final MessageDigest initialValue() {
194194

195195
ConcurrentMap<String, Long> hostDeadDur;
196196

197+
private AuthInfo authInfo;
198+
197199
private boolean isTcp;
198200

199201
private int bufferSize = 1024 * 1025;
@@ -222,6 +224,20 @@ public static SchoonerSockIOPool getInstance(String poolName) {
222224
return pools.get(poolName);
223225
}
224226

227+
public static SchoonerSockIOPool getInstance(String poolName, AuthInfo authInfo) {
228+
SchoonerSockIOPool pool;
229+
230+
synchronized (pools) {
231+
if (!pools.containsKey(poolName)) {
232+
pool = new SchoonerSockIOPool(true);
233+
pool.authInfo = authInfo;
234+
pools.putIfAbsent(poolName, pool);
235+
}
236+
}
237+
238+
return pools.get(poolName);
239+
}
240+
225241
public static SchoonerSockIOPool getInstance(String poolName, boolean isTcp) {
226242
SchoonerSockIOPool pool;
227243

@@ -251,6 +267,10 @@ public static SchoonerSockIOPool getInstance() {
251267
return getInstance("default", true);
252268
}
253269

270+
public static SchoonerSockIOPool getInstance(AuthInfo authInfo) {
271+
return getInstance("default", authInfo);
272+
}
273+
254274
public static SchoonerSockIOPool getInstance(boolean isTcp) {
255275
return getInstance("default", isTcp);
256276
}
@@ -308,8 +328,13 @@ private void populateBuckets() {
308328
// Create a socket pool for each host
309329
// Create an object pool to contain our active connections
310330
GenericObjectPool gop;
311-
SchoonerSockIOFactory factory = new SchoonerSockIOFactory(servers[i], isTcp, bufferSize, socketTO,
312-
socketConnectTO, nagle);
331+
SchoonerSockIOFactory factory;
332+
if (authInfo != null) {
333+
factory = new AuthSchoonerSockIOFactory(servers[i], isTcp, bufferSize, socketTO, socketConnectTO,
334+
nagle, authInfo);
335+
} else {
336+
factory = new SchoonerSockIOFactory(servers[i], isTcp, bufferSize, socketTO, socketConnectTO, nagle);
337+
}
313338
gop = new GenericObjectPool(factory, maxConn, GenericObjectPool.WHEN_EXHAUSTED_BLOCK, maxIdle, maxConn);
314339
factory.setSockets(gop);
315340
socketPool.put(servers[i], gop);
@@ -348,8 +373,13 @@ private void populateConsistentBuckets() {
348373
// Create a socket pool for each host
349374
// Create an object pool to contain our active connections
350375
GenericObjectPool gop;
351-
SchoonerSockIOFactory factory = new SchoonerSockIOFactory(servers[i], isTcp, bufferSize, socketTO,
352-
socketConnectTO, nagle);
376+
SchoonerSockIOFactory factory;
377+
if (authInfo != null) {
378+
factory = new AuthSchoonerSockIOFactory(servers[i], isTcp, bufferSize, socketTO, socketConnectTO,
379+
nagle, authInfo);
380+
} else {
381+
factory = new SchoonerSockIOFactory(servers[i], isTcp, bufferSize, socketTO, socketConnectTO, nagle);
382+
}
353383
gop = new GenericObjectPool(factory, maxConn, GenericObjectPool.WHEN_EXHAUSTED_BLOCK, maxIdle, maxConn);
354384
factory.setSockets(gop);
355385
socketPool.put(servers[i], gop);

‎src/test/com/schooner/MemCached/MemCachedBenchTcp.java‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,16 +30,15 @@
3030

3131
import com.danga.MemCached.MemCachedClient;
3232

33-
3433
public class MemCachedBenchTcp {
3534

3635
// logger
3736
// private static Logger log =
3837
// Logger.getLogger(MemCachedBench.class.getName());
3938

4039
public static void main(String[] args) {
41-
int runs = Integer.parseInt(args[0]);
42-
int start = Integer.parseInt(args[1]);
40+
int runs = Integer.parseInt(args[1]);
41+
int start = Integer.parseInt(args[2]);
4342

4443
String servers = System.getProperty("memcached.host");
4544
String[] serverlist = servers.split(",");
@@ -54,7 +53,11 @@ public static void main(String[] args) {
5453
pool.initialize();
5554

5655
// get client instance
57-
MemCachedClient mc = new MemCachedClient("test");
56+
MemCachedClient mc;
57+
if (args[0].equals("ascii"))
58+
mc = new MemCachedClient("test", false);
59+
else
60+
mc = new MemCachedClient("test", true);
5861

5962
String keyBase = "testKey";
6063
String object = "This is a test of an object blah blah es, serialization does not seem to slow things down so much. The gzip compression is horrible horrible performance, so we only use it for very large objects. I have not done any heavy benchmarking recently";

0 commit comments

Comments
 (0)