Tuesday, April 16, 2013

Strategy Pattern Applied - KeyStoreTemplate

Update: I wrongly attributed the design pattern implementation below to the Template pattern, it's actually the Strategy pattern as pointed out by a user in Reddit. However, the article below is kept in its original form, with just the title changed. Sorry for the confusion :)

The Template design pattern is useful if you have methods that follow a similar pattern but differ only in a few places. Usually, the 'generic' method will call one or several abstract methods, which are then implemented in a concrete class.

The JdbcTemplate of the Spring Framework was (and still is) a life saver. No more unclosed Connection objects that may leak to JDBC connection leaks. No more try/catch blocks that obscure the code logic. Converting the JDBC DAOs of legacy systems to use JdbcTemplate objects provided immediate benefits in terms of simplicity (eliminating copy-and-paste codes) and robust resource management (e.g. closing Connections, ResultSets, Statements).

There are a lot of useful articles detailing what, why and how of this design pattern, so I will not rehash them here. I will just document down how I applied this pattern to solve some age-old problem.

I decided to follow the JdbcTemplate's design style to come up with a KeyStoreTemplate, a class that handles importing or exporting of keys to/from a keystore. To make things simple, this class will only handle 2 things: Importing a key, and exporting a key.

Firstly, the KeyStoreTemplate code (abridged for simplicity):


import org.apache.commons.io.IOUtils;

public class KeyStoreTemplate {

 /**
 * Just enough parameters to open a keystore.
 */
 public KeyStoreTemplate(File keyStoreFile, Provider provider, String keyStoreType, String keyStorePassword) {
  this.keyStoreFile = keyStoreFile;
  this.provider = provider;
  this.keyStoreType = keyStoreType;
  this.keyStorePassword = keyStorePassword;
 }

 /**
 * This is to get/export a key from the keystore.
 */
 public Key getFromKeyStore(KeyStoreExporter exporter) {
  FileInputStream fis = null;
  try {
   KeyStore ks = KeyStore.getInstance(keyStoreType, provider);
         fis = new FileInputStream(keyStoreFile);

         //Keystore file must exists!
         if(keyStoreFile.exists()) {
          ks.load(fis, keyStorePassword.toCharArray());
         } else {
          throw new RuntimeException("Keystore file cannot be located!");
         }
         
         String keyPass = keyStorePassword;

         return exporter.exportKey(ks, keyPass);
         
  } catch (Exception e) {
   throw new RuntimeException("Something went wrong!", e);
  } finally {
   IOUtils.closeQuietly(fis);
  }  
 }

 /**
 * This is to save/import a key to the keystore.
 */
 public void saveToKeyStore(KeyStoreImporter importer) {
  
  FileInputStream fis = null;
  FileOutputStream fos = null;
  try {
   KeyStore ks = KeyStore.getInstance(keyStoreType, provider);
         fis = new FileInputStream(keyStoreFile);

         //Keystore file must exists!
         if(keyStoreFile.exists()) {
          ks.load(fis, keyStorePassword.toCharArray());
         } else {
          throw new BusinessException("Keystore file cannot be located!");
         }
         
         String keyPass = keyStorePassword;

         importer.importPrivateKey(ks, keyPass);
         
         fos = new FileOutputStream(keyStoreFile);
         ks.store(fos, keyStorePassword.toCharArray());
         
  } catch (Exception e) {
   throw new RuntimeException("Something went wrong!", e);
  } finally {
   IOUtils.closeQuietly(fis);
   IOUtils.closeQuietly(fos);
  }
 }
 
}


Notice that there are 2 interfaces: KeyStoreExporter and KeyStoreImporter. These interfaces are implemented by the caller class of KeyStoreTemplate to do the 'effective' stuff, without worrying about the KeyStore details. The KeyStore object is passed in, much like the ResultSet object is passed in when using JdbcTemplate.


import java.security.GeneralSecurityException;
import java.security.KeyStore;

public interface KeyStoreImporter {

 public void importPrivateKey(KeyStore keyStore, String keyPassword) throws GeneralSecurityException;
}


import java.security.GeneralSecurityException;
import java.security.Key;
import java.security.KeyStore;

public interface KeyStoreExporter {

 public Key exportKey(KeyStore keyStore, String keyPassword) throws GeneralSecurityException;
}


So, how to use them? Example as below:


//Key generation and key import
KeyPair keyPair = generateRSAKeyPair();
final Certificate[] certs = CryptoUtils.generateCertificate(keyPair);

RSAPublicKey publicKey = (RSAPublicKey)keyPair.getPublic();
final RSAPrivateKey privateKey = (RSAPrivateKey)keyPair.getPrivate();

logger.debug("[generateAndStoreRSAKeyPair] pubKey={},privKey={}", publicKey.getClass(), privateKey.getClass());

keyStoreTemplate.saveToKeyStore(new KeyStoreImporter() {
 @Override
 public void importPrivateKey(KeyStore ks, String keyPassword) throws KeyStoreException {
  ks.setKeyEntry(alias, privateKey, keyPassword.toCharArray(), certs);
 }
});


//Method to get an RSA public key
public RSAPublicKey getRSAPublicKey(final String alias) throws BusinessException {

 return (RSAPublicKey)keyStoreTemplate.getFromKeyStore(new KeyStoreExporter() {
  @Override
  public Key exportKey(KeyStore keyStore, @SuppressWarnings("unused") String keyPassword) throws GeneralSecurityException {
   Certificate cert = keyStore.getCertificate(alias);
   RSAPublicKey publicKey = (RSAPublicKey)cert.getPublicKey();
   return publicKey;
  }
 });
}


//Method to get a secret key
public SecretKey getSecretKey(final String alias) throws BusinessException {

 return (SecretKey)keyStoreTemplate.getFromKeyStore(new KeyStoreExporter() {
  @Override
  public Key exportKey(KeyStore keyStore, String keyPassword) throws GeneralSecurityException {
   SecretKey secretKey = (SecretKey)keyStore.getKey(alias, keyPassword.toCharArray());
   return secretKey;
  }
 });
}