com.hubspot.singularity.data.AuthTokenManager Maven / Gradle / Ivy

Go to download
package com.hubspot.singularity.data;

import com.codahale.metrics.MetricRegistry;
import com.google.inject.Inject;
import com.hubspot.singularity.SingularityTokenResponse;
import com.hubspot.singularity.SingularityUser;
import com.hubspot.singularity.config.SingularityConfiguration;
import com.hubspot.singularity.data.transcoders.Transcoder;
import java.math.BigInteger;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.util.Optional;
import java.util.UUID;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;
import org.apache.curator.framework.CuratorFramework;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class AuthTokenManager extends CuratorManager {
  private static final Logger LOG = LoggerFactory.getLogger(AuthTokenManager.class);

  private static final String TOKEN_ROOT = "/tokens";
  private static final String TOKEN_PATH = TOKEN_ROOT + "/%s";

  private final Transcoder userTranscoder;

  @Inject
  public AuthTokenManager(
    CuratorFramework curator,
    SingularityConfiguration configuration,
    MetricRegistry metricRegistry,
    Transcoder userTranscoder
  ) {
    super(curator, configuration, metricRegistry);
    this.userTranscoder = userTranscoder;
  }

  public SingularityTokenResponse generateToken(SingularityUser userData)
    throws NoSuchAlgorithmException, InvalidKeySpecException {
    String newToken = UUID.randomUUID().toString();
    return saveToken(newToken, userData);
  }

  public SingularityTokenResponse saveToken(String newToken, SingularityUser userData)
    throws NoSuchAlgorithmException, InvalidKeySpecException {
    String hashed = generateTokenHash(newToken);
    writeToken(hashed, userData);
    return new SingularityTokenResponse(newToken, userData);
  }

  private void writeToken(String hashed, SingularityUser userData) {
    save(getTokenPath(hashed), userData, userTranscoder);
  }

  private void deleteToken(String hashed) {
    delete(getTokenPath(hashed));
  }

  public void clearTokensForUser(String user) {
    for (String hashed : getChildren(TOKEN_ROOT)) {
      Optional maybeUser = getData(getTokenPath(hashed), userTranscoder);
      if (maybeUser.isPresent() && maybeUser.get().getName().equals(user)) {
        deleteToken(hashed);
      }
    }
  }

  public SingularityUser getUserIfValidToken(String token) {
    for (String hashed : getChildren(TOKEN_ROOT)) {
      try {
        if (validateToken(token, hashed)) {
          Optional maybeUser = getData(
            getTokenPath(hashed),
            userTranscoder
          );
          if (maybeUser.isPresent()) {
            return maybeUser.get();
          }
        }
      } catch (Throwable t) {
        LOG.error("Unable to validate token", t);
      }
    }
    LOG.debug("No matching token found");
    return null;
  }

  private String getTokenPath(String hashed) {
    return String.format(TOKEN_PATH, hashed);
  }

  // Implementation of PBKDF2WithHmacSHA1
  private static boolean validateToken(String originalToken, String storedToken)
    throws NoSuchAlgorithmException, InvalidKeySpecException {
    String[] parts = storedToken.split(":");
    int iterations = Integer.parseInt(parts[0]);
    byte[] salt = fromHex(parts[1]);
    byte[] hash = fromHex(parts[2]);

    PBEKeySpec spec = new PBEKeySpec(
      originalToken.toCharArray(),
      salt,
      iterations,
      hash.length * 8
    );
    SecretKeyFactory skf = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA1");
    byte[] testHash = skf.generateSecret(spec).getEncoded();

    int diff = hash.length ^ testHash.length;
    for (int i = 0; i < hash.length && i < testHash.length; i++) {
      diff |= hash[i] ^ testHash[i];
    }
    return diff == 0;
  }

  private static byte[] fromHex(String hex) throws NoSuchAlgorithmException {
    byte[] bytes = new byte[hex.length() / 2];
    for (int i = 0; i < bytes.length; i++) {
      bytes[i] = (byte) Integer.parseInt(hex.substring(2 * i, 2 * i + 2), 16);
    }
    return bytes;
  }

  private static String generateTokenHash(String password)
    throws NoSuchAlgorithmException, InvalidKeySpecException {
    int iterations = 1000;
    char[] chars = password.toCharArray();
    byte[] salt = getSalt();

    PBEKeySpec spec = new PBEKeySpec(chars, salt, iterations, 64 * 8);
    SecretKeyFactory skf = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA1");
    byte[] hash = skf.generateSecret(spec).getEncoded();
    return iterations + ":" + toHex(salt) + ":" + toHex(hash);
  }

  private static byte[] getSalt() throws NoSuchAlgorithmException {
    SecureRandom sr = SecureRandom.getInstance("SHA1PRNG");
    byte[] salt = new byte[16];
    sr.nextBytes(salt);
    return salt;
  }

  private static String toHex(byte[] array) {
    BigInteger bi = new BigInteger(1, array);
    String hex = bi.toString(16);
    int paddingLength = (array.length * 2) - hex.length();
    if (paddingLength > 0) {
      return String.format("%0" + paddingLength + "d", 0) + hex;
    } else {
      return hex;
    }
  }
}