urlencode keys to protect against malformed keys

Author: gwhalin

src/com/danga/MemCached/MemCachedClient.java

@@ -33,6 +33,7 @@
 import java.util.*;
 import java.util.zip.*;
 import java.io.*;
+import java.net.URLEncoder;
 
 import org.apache.log4j.Logger;
 
@@ -346,6 +347,14 @@ public boolean delete( String key, Integer hashCode, Date expiry ) {
 			return false;
 		}
 
+		try {
+			key = sanitizeKey( key );
+		}
+		catch ( UnsupportedEncodingException e ) {
+			log.error( "failed to sanitize your key!", e );
+			return false;
+		}
+
 		// get SockIO obj from hash or from key
 		SockIOPool.SockIO sock = SockIOPool.getInstance( poolName ).getSock( key, hashCode );
 
@@ -574,6 +583,14 @@ private boolean set( String cmdname, String key, Object value, Date expiry, Inte
 			return false;
 		}
 
+		try {
+			key = sanitizeKey( key );
+		}
+		catch ( UnsupportedEncodingException e ) {
+			log.error( "failed to sanitize your key!", e );
+			return false;
+		}
+
 		if ( value == null ) {
 			log.error( "trying to store a null value to cache" );
 			return false;
@@ -604,7 +621,7 @@ private boolean set( String cmdname, String key, Object value, Date expiry, Inte
 					val = value.toString().getBytes( defaultEncoding );
 				}
 				catch ( UnsupportedEncodingException ue ) {
-					log.error( "invalid encoding type used: " + defaultEncoding );
+					log.error( "invalid encoding type used: " + defaultEncoding, ue );
 					sock.close();
 					sock = null;
 					return false;
@@ -776,6 +793,14 @@ public long getCounter( String key, Integer hashCode ) {
 			return -1;
 		}
 
+		try {
+			key = sanitizeKey( key );
+		}
+		catch ( UnsupportedEncodingException e ) {
+			log.error( "failed to sanitize your key!", e );
+			return -1;
+		}
+
 		long counter = -1;
 		try {
 			counter = Long.parseLong( (String)get( key, hashCode, true ) );
@@ -871,6 +896,19 @@ public long decr( String key, long inc, Integer hashCode ) {
 	 */
 	private long incrdecr( String cmdname, String key, long inc, Integer hashCode ) {
 
+		if ( key == null ) {
+			log.error( "null key for incrdecr()" );
+			return -1;
+		}
+
+		try {
+			key = sanitizeKey( key );
+		}
+		catch ( UnsupportedEncodingException e ) {
+			log.error( "failed to sanitize your key!", e );
+			return -1;
+		}
+
 		// get SockIO obj for given cache key
 		SockIOPool.SockIO sock = SockIOPool.getInstance( poolName ).getSock(key, hashCode);
 
@@ -980,6 +1018,14 @@ public Object get( String key, Integer hashCode, boolean asString ) {
 			return null;
 		}
 
+		try {
+			key = sanitizeKey( key );
+		}
+		catch ( UnsupportedEncodingException e ) {
+			log.error( "failed to sanitize your key!", e );
+			return false;
+		}
+
 		// get SockIO obj using cache key
 		SockIOPool.SockIO sock = SockIOPool.getInstance( poolName ).getSock( key, hashCode );
 
@@ -1140,8 +1186,23 @@ public Map<String,Object> getMulti( String[] keys, Integer[] hashCodes, boolean
 			if ( hashCodes != null && hashCodes.length > i )
 				hash = hashCodes[ i ];
 
+			String key = keys[i];
+
+			if ( key == null ) {
+				log.error( "null key, so skipping" );
+				continue;
+			}
+
+			try {
+				key = sanitizeKey( key );
+			}
+			catch ( UnsupportedEncodingException e ) {
+				log.error( "failed to sanitize your key!", e );
+				continue;
+			}
+
 			// get SockIO obj from cache key
-			SockIOPool.SockIO sock = SockIOPool.getInstance( poolName ).getSock( keys[i], hash );
+			SockIOPool.SockIO sock = SockIOPool.getInstance( poolName ).getSock( key, hash );
 
 			if ( sock == null )
 				continue;
@@ -1150,7 +1211,7 @@ public Map<String,Object> getMulti( String[] keys, Integer[] hashCodes, boolean
 			if ( !sockKeys.containsKey( sock.getHost() ) )
 				sockKeys.put( sock.getHost(), new StringBuilder() );
 
-			sockKeys.get( sock.getHost() ).append( " " + keys[i] );
+			sockKeys.get( sock.getHost() ).append( " " + key );
 
 			// return to pool
 			sock.close();
@@ -1303,6 +1364,10 @@ else if ( END.equals( line ) ) {
 		}
 	}
 
+	private String sanitizeKey( String key ) throws UnsupportedEncodingException {
+		return URLEncoder.encode( key, "UTF-8" );
+	}
+
 	/** 
 	 * Invalidates the entire cache.
 	 *

src/com/danga/MemCached/test/UnitTests.java

@@ -153,6 +153,12 @@ public static void test16() {
         assert !mc.set( "foo", null );
 	}
 
+	public static void test17() {
+        mc.set( "foo bar", Boolean.TRUE );
+        Boolean b = (Boolean)mc.get( "foo bar" );
+		assert b.booleanValue();
+	}
+    
 	/**
 	 * This runs through some simple tests of the MemCacheClient.
 	 *
@@ -199,5 +205,7 @@ public static void main(String[] args) {
 		test13();
 		test14();
 		test15();
+		test16();
+		test17();
 	}
 }