com.ott.bookings.auth.form.impl.MongoTockenStore.java Source code

Java tutorial

Introduction

Here is the source code for com.ott.bookings.auth.form.impl.MongoTockenStore.java

Source

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.ott.bookings.auth.form.impl;

import java.io.File;
import java.io.UnsupportedEncodingException;
import java.net.UnknownHostException;
import java.security.InvalidKeyException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.List;

import javax.crypto.Mac;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.mongodb.BasicDBObject;
import com.mongodb.DB;
import com.mongodb.DBCollection;
import com.mongodb.DBCursor;
import com.mongodb.DBObject;
import com.mongodb.MongoClient;
import com.mongodb.MongoClientOptions;
import com.mongodb.MongoException;
import com.mongodb.ReadPreference;
import com.mongodb.ServerAddress;
import com.mongodb.WriteConcern;
import com.mongodb.WriteResult;

public class MongoTockenStore {

    public final Logger log = LoggerFactory.getLogger(MongoTockenStore.class);

    private static final String ACTIVE_PROPERTY = "active";

    /** The number of secret keys in the token buffer currentTokens */
    private static final int TOKEN_BUFFER_SIZE = 5;

    /**
     * MongoDB Hosts. This can be a single host or a comma separated list using
     * a [host:port] format. Lists are typically used for replica sets. This
     * value is ignored if the <em>dbConnectionUri</em> is provided.
     * <p>
     * 
     * <pre>
     *       127.0.0.1:27001,127.0.0.1:27002
     * </pre>
     * </p>
     */
    private String hostsString;

    /**
     * Name of the MongoDB Database to use.
     */
    private String dbName;

    /**
     * Name of the MongoDB Collection
     */
    private String collectionName;

    /**
     * {@link MongoClient} instance to use.
     */
    private MongoClient mongoClient;

    /**
     * Connection Timeout in milliseconds. Defaults to 0, or no timeout
     */
    private int connectionTimeoutMs = 0;

    /**
     * Connection Wait Timeout in milliseconds. Defaults to 0, or no timeout
     */
    private int connectionWaitTimeoutMs = 0;

    /**
     * Maximum Number of connections the MongoClient will manage. Defaults to 20
     */
    private int maxPoolSize = 20;

    /**
     * Controls what {@link WriteConcern} the MongoClient will use. Defaults to
     * "SAFE"
     */
    private WriteConcern writeConcern = WriteConcern.SAFE;

    /**
     * Mongo DB reference
     */
    private DB db;

    /**
     * Mongo Collection for the Sessions
     */
    private DBCollection collection;

    /**
     * Controls if the MongoClient will write to slaves. Equivalent to
     * <em>slaveOk</em>. Defaults to false.
     */
    private boolean useSlaves = false;

    /**
     * @return the useSlaves
     */
    public boolean isUseSlaves() {
        return useSlaves;
    }

    /**
     * A secure random used for generating new tokens.
     */
    private SecureRandom random;

    /**
     * The ttl of the cookie before it becomes invalid (in ms)
     */
    private final long ttl;

    /**
     * Array of hex characters used by {@link #byteToHex(byte[])} to convert a
     * byte array to a hex string.
     */
    private static final char[] TOHEX = "0123456789abcdef".toCharArray();

    /**
     * Name of the <code>SecureRandom</code> generator algorithm
     */
    private static final String SHA1PRNG = "SHA1PRNG";

    /**
     * The name of the HMAC function to calculate the hash code of the payload
     * with the secure token.
     */
    private static final String HMAC_SHA1 = "HmacSHA1";

    /**
     * String encoding to convert byte arrays to strings and vice-versa.
     */
    private static final String UTF_8 = "UTF-8";

    private static final String SECRET_KEY_PROPERTY = "secretKey";

    private static final String NEXT_UPDATE_PROPERTY = "nextUpdate";

    private static final String ID_PROPERTY = "_id";

    /**
     * @param useSlaves
     *            the useSlaves to set
     */
    public void setUseSlaves(boolean useSlaves) {
        this.useSlaves = useSlaves;
    }

    public MongoTockenStore(String mongoHosts, String dbName, String collectionName, long sessionTimeout,
            boolean fastSeed) throws NoSuchAlgorithmException, IllegalStateException, UnsupportedEncodingException,
            InvalidKeyException {
        this.hostsString = mongoHosts;
        this.dbName = dbName;
        this.collectionName = collectionName;
        this.random = SecureRandom.getInstance(SHA1PRNG);
        this.ttl = sessionTimeout;
        getConnection();
        // warm up the crypto API
        if (fastSeed) {
            random.setSeed(getFastEntropy());
        } else {
            log.info("Seeding the secure random number generator can take "
                    + "up to several minutes on some operating systems depending "
                    + "upon environment factors. If this is a problem for you, "
                    + "set the system property 'java.security.egd' to "
                    + "'file:/dev/./urandom' or enable the Fast Seed Generator " + "in the Web Console");
        }

        byte[] b = new byte[20];
        random.nextBytes(b);
        final SecretKey secretKey = new SecretKeySpec(b, HMAC_SHA1);
        final Mac m = Mac.getInstance(HMAC_SHA1);
        m.init(secretKey);
        m.update(UTF_8.getBytes(UTF_8));
        m.doFinal();
    }

    /**
     * Returns <code>true</code> if the <code>value</code> is a valid secure
     * token as follows:
     * <ul>
     * <li>The string is not <code>null</code></li>
     * <li>The string contains three fields separated by an @ sign</li>
     * <li>The expiry time encoded in the second field has not yet passed</li>
     * <li>The hashing the third field, the expiry time and token number with
     * the secure token (indicated by the token number) gives the same value as
     * contained in the first field</li>
     * </ul>
     * <p>
     * Otherwise the method returns <code>false</code>.
     */
    boolean isValid(String value) {
        String[] parts = split(value);
        if (parts != null) {

            // single digit token number
            int tokenNumber = parts[1].charAt(0) - '0';
            if (tokenNumber >= 0 && tokenNumber < TOKEN_BUFFER_SIZE) {

                long cookieTime = Long.parseLong(parts[1].substring(1));
                if (System.currentTimeMillis() < cookieTime) {

                    try {
                        SecretKey secretKey = getSecretKeyFromDb(tokenNumber);
                        String hmac = encode(cookieTime, parts[2], tokenNumber, secretKey);
                        return value.equals(hmac);
                    } catch (ArrayIndexOutOfBoundsException e) {
                        log.error(e.getMessage(), e);
                    } catch (InvalidKeyException e) {
                        log.error(e.getMessage(), e);
                    } catch (IllegalStateException e) {
                        log.error(e.getMessage(), e);
                    } catch (UnsupportedEncodingException e) {
                        log.error(e.getMessage(), e);
                    } catch (NoSuchAlgorithmException e) {
                        log.error(e.getMessage(), e);
                    }

                    log.error("AuthNCookie value '{}' is invalid", value);

                } else {
                    log.error("AuthNCookie value '{}' has expired {}ms ago", value,
                            (System.currentTimeMillis() - cookieTime));
                }

            } else {
                log.error("AuthNCookie value '{}' is invalid: refers to an invalid token number", value,
                        tokenNumber);
            }

        } else {
            log.error("AuthNCookie value '{}' has invalid format", value);
        }

        // failed verification, reason is logged
        return false;
    }

    /**
     * Maintain a circular buffer to tokens, and return the current one.
     *
     * @return the current token.
     */
    private synchronized Token getActiveToken() {
        Token tokenPointer = getCurrentTokenFromDb();

        if (System.currentTimeMillis() > tokenPointer.getNextUpdate() || tokenPointer.getKey() == null) {

            // cycle so that during a typical ttl the tokens get completely
            // refreshed.
            log.debug("cycle so that during a typical ttl the tokens get completely refreshed.");
            long nextUpdate = System.currentTimeMillis() + ttl / (TOKEN_BUFFER_SIZE - 1);
            byte[] b = new byte[20];
            random.nextBytes(b);
            int nextToken = tokenPointer.getToken() + 1;
            if (nextToken == TOKEN_BUFFER_SIZE) {
                nextToken = 0;
            }
            Token nextTokenPointer = new Token(nextToken, nextUpdate, b);
            return saveCurrentToken(tokenPointer, nextTokenPointer);
        }

        return tokenPointer;
    }

    private Token saveCurrentToken(Token tokenPointer, Token nextTokenPointer) {

        DBObject selectQuery = new BasicDBObject(ID_PROPERTY, nextTokenPointer.getToken());
        selectQuery.put(ACTIVE_PROPERTY, false);

        BasicDBObject updateFields = new BasicDBObject();
        updateFields.put(ACTIVE_PROPERTY, true);
        updateFields.put(NEXT_UPDATE_PROPERTY, nextTokenPointer.getNextUpdate());
        updateFields.put(SECRET_KEY_PROPERTY, nextTokenPointer.getKey());

        DBObject updateQuery = new BasicDBObject();
        updateQuery.put("$set", updateFields);
        WriteResult result = this.collection.update(selectQuery, updateQuery, true, false);

        // update old token
        selectQuery = new BasicDBObject(ID_PROPERTY, tokenPointer.getToken());
        this.collection.update(selectQuery, new BasicDBObject("$set", new BasicDBObject(ACTIVE_PROPERTY, false)));

        if (result.isUpdateOfExisting()) {
            return nextTokenPointer;
        } else {
            return toToken(this.collection.findOne(new BasicDBObject(ID_PROPERTY, nextTokenPointer.getToken())));
        }
    }

    private SecretKey getSecretKeyFromDb(int tokenNumber) {
        /* locate the session, by id, in the collection */
        BasicDBObject tokenQuery = new BasicDBObject();
        tokenQuery.put(ID_PROPERTY, tokenNumber);

        /* lookup the session */
        DBObject tokenObject = this.collection.findOne(tokenQuery);

        byte[] data = (byte[]) tokenObject.get(SECRET_KEY_PROPERTY);
        SecretKey key = new SecretKeySpec(data, HMAC_SHA1);
        return key;
    }

    private Token getCurrentTokenFromDb() {

        /* locate the session, by id, in the collection */
        BasicDBObject tokenQuery = new BasicDBObject();
        tokenQuery.put(ACTIVE_PROPERTY, true);

        /* lookup the session */
        DBCursor cursor = this.collection.find(tokenQuery);
        if (cursor.hasNext()) {
            DBObject obj = cursor.next();
            return toToken(obj);
        }
        log.debug("Token does not exist create a new ones");
        return new Token(0, System.currentTimeMillis(), null);

    }

    private Token toToken(DBObject obj) {
        return new Token(((Integer) obj.get(ID_PROPERTY)).intValue(),
                ((Long) obj.get(NEXT_UPDATE_PROPERTY)).longValue(), (byte[]) obj.get(SECRET_KEY_PROPERTY));
    }

    /**
     * @param expires
     * @param userId
     * @return
     * @throws UnsupportedEncodingException
     * @throws IllegalStateException
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeyException
     */

    String encode(long expires, final String userId)
            throws InvalidKeyException, UnsupportedEncodingException, NoSuchAlgorithmException {
        Token token = getActiveToken();
        SecretKey key = new SecretKeySpec(token.getKey(), HMAC_SHA1);
        return encode(expires, userId, token.getToken(), key);
    }

    /**
     * Splits the authentication data into the three parts packed together while
     * encoding the cookie.
     *
     * @param authData
     *            The authentication data to split in three parts
     * @return A string array with three elements being the three parts of the
     *         cookie value or <code>null</code> if the input is
     *         <code>null</code> or if the string does not contain (at least)
     *         three '@' separated parts.
     */
    static String[] split(final String authData) {
        String[] parts = StringUtils.split(authData, "@", 3);
        if (parts != null && parts.length == 3) {
            return parts;
        }
        return null;
    }

    private List<ServerAddress> getServerAddresses() throws NumberFormatException, UnknownHostException {
        /* build up the host list */
        List<ServerAddress> hosts = new ArrayList<ServerAddress>();
        String[] dbHosts = this.hostsString.split(",");
        for (String dbHost : dbHosts) {
            String[] hostInfo = dbHost.split(":");
            ServerAddress address = new ServerAddress(hostInfo[0], Integer.parseInt(hostInfo[1]));
            hosts.add(address);
        }
        return hosts;
    }

    /**
     * Create the {@link MongoClient}.
     * 
     */
    private void getConnection() {
        try {

            /* create the client using the Mongo options */
            ReadPreference readPreference = ReadPreference.primaryPreferred();
            if (this.useSlaves) {
                readPreference = ReadPreference.secondaryPreferred();
            }
            MongoClientOptions options = MongoClientOptions.builder().connectTimeout(connectionTimeoutMs)
                    .maxWaitTime(connectionWaitTimeoutMs).connectionsPerHost(maxPoolSize).writeConcern(writeConcern)
                    .readPreference(readPreference).build();

            log.info("[Mongo Token Store]: Connecting to MongoDB [" + this.hostsString + "]");

            /* connect */
            this.mongoClient = new MongoClient(getServerAddresses(), options);

            /* get a connection to our db */
            log.info("[Mongo Token Store]: Using Database [" + this.dbName + "]");
            this.db = this.mongoClient.getDB(this.dbName);
            /* get a reference to the collection */
            this.collection = this.db.getCollection(this.collectionName);

            log.info("[Mongo Token Store]: Store ready.");
        } catch (UnknownHostException uhe) {
            log.error("Unable to Connect to MongoDB", uhe);
        } catch (MongoException me) {
            log.error("Unable to Connect to MongoDB", me);
        }
    }

    /**
     * Creates a byte array of entry from the current state of the system:
     * <ul>
     * <li>The current system time in milliseconds since the epoch</li>
     * <li>The number of nanoseconds since system startup</li>
     * <li>The name, size and last modification time of the files in the
     * <code>java.io.tmpdir</code> folder.</li>
     * </ul>
     * <p>
     * <b>NOTE</b> This method generates entropy fast but not necessarily secure
     * enough for seeding the random number generator.
     *
     * @return bytes of entropy
     */
    private static byte[] getFastEntropy() {
        final MessageDigest md;

        try {
            md = MessageDigest.getInstance("SHA");
        } catch (NoSuchAlgorithmException nsae) {
            throw new InternalError("internal error: SHA-1 not available.");
        }

        // update with XorShifted time values
        update(md, System.currentTimeMillis());
        update(md, System.nanoTime());

        // scan the temp file system
        File file = new File(System.getProperty("java.io.tmpdir"));
        File[] entries = file.listFiles();
        if (entries != null) {
            for (File entry : entries) {
                md.update(entry.getName().getBytes());
                update(md, entry.lastModified());
                update(md, entry.length());
            }
        }

        return md.digest();
    }

    /**
     * Updates the message digest with an XOR-Shifted value.
     *
     * @param md
     *            The MessageDigest to update
     * @param value
     *            The original value to be XOR-Shifted first before taking the
     *            bytes to update the message digest
     */
    private static void update(final MessageDigest md, long value) {
        value ^= (value << 21);
        value ^= (value >>> 35);
        value ^= (value << 4);

        for (int i = 0; i < 8; i++) {
            md.update((byte) value);
            value >>= 8;
        }
    }

    private String encode(final long expires, final String userId, final int token, final SecretKey key)
            throws IllegalStateException, UnsupportedEncodingException, NoSuchAlgorithmException,
            InvalidKeyException {

        String cookiePayload = String.valueOf(token) + String.valueOf(expires) + "@" + userId;
        Mac m = Mac.getInstance(HMAC_SHA1);
        m.init(key);
        m.update(cookiePayload.getBytes(UTF_8));
        String cookieValue = byteToHex(m.doFinal());
        return cookieValue + "@" + cookiePayload;
    }

    /**
     * Encode a byte array.
     *
     * @param base
     * @return
     */
    private String byteToHex(byte[] base) {
        char[] c = new char[base.length * 2];
        int i = 0;

        for (byte b : base) {
            int j = b;
            j = j + 128;
            c[i++] = TOHEX[j / 0x10];
            c[i++] = TOHEX[j % 0x10];
        }
        return new String(c);
    }

}