UnboundID Server SDK

UnboundID Server SDK Documentation

ExampleTokenStore.java

/*
 * CDDL HEADER START
 *
 * The contents of this file are subject to the terms of the
 * Common Development and Distribution License, Version 1.0 only
 * (the "License").  You may not use this file except in compliance
 * with the License.
 *
 * You can obtain a copy of the license at
 * docs/licenses/cddl.txt
 * or http://www.opensource.org/licenses/cddl1.php.
 * See the License for the specific language governing permissions
 * and limitations under the License.
 *
 * When distributing Covered Code, include this CDDL HEADER in each
 * file and include the License file at
 * docs/licenses/cddl.txt.  If applicable,
 * add the following below this CDDL HEADER, with the fields enclosed
 * by brackets "[]" replaced with your own identifying information:
 *      Portions Copyright [yyyy] [name of copyright owner]
 *
 * CDDL HEADER END
 *
 *
 *      Copyright 2010-2013 UnboundID Corp.
 */
package com.unboundid.directory.sdk.examples;

import com.unboundid.directory.sdk.broker.api.TokenStore;
import com.unboundid.directory.sdk.broker.config.TokenStoreConfig;
import com.unboundid.directory.sdk.broker.types.AccessToken;
import com.unboundid.directory.sdk.broker.types.AuthorizationCode;
import com.unboundid.directory.sdk.broker.types.IdentityBrokerContext;
import com.unboundid.directory.sdk.broker.types.OAuthException;
import com.unboundid.directory.sdk.broker.types.RefreshToken;
import com.unboundid.ldap.sdk.LDAPException;
import com.unboundid.util.args.ArgumentParser;

import java.util.Collection;
import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;

/**
 * This example provides an in-memory implementation of a TokenStore that
 * generates UUID valued code and token values and uses them as keys for storage
 * in a HashMap. Note however that all {@link TokenStore} implementations must
 * store their tokens and authorization codes in persistent storage that is
 * accessible to all Identity Broker instances.
 */
public class ExampleTokenStore extends TokenStore
{
  private volatile TokenStoreConfig config;

  private IdentityBrokerContext serverContext;

  private final ConcurrentHashMap<String, AuthorizationCode> authzCodeMap =
      new ConcurrentHashMap<String, AuthorizationCode>();

  private final ConcurrentHashMap<String, AccessToken> accessTokenMap =
      new ConcurrentHashMap<String, AccessToken>();

  private final ConcurrentHashMap<String, AccessToken> refreshTokenMap =
      new ConcurrentHashMap<String, AccessToken>();

  /**
   * {@inheritDoc}
   */
  @Override
  public String getExtensionName()
  {
    return "Example Token Store";
  }

  /**
   * {@inheritDoc}
   */
  @Override
  public String[] getExtensionDescription()
  {
    return new String[]
    {
      "This implementation serves as an in-memory exapmle implementation " +
      "that may be used to demonstrate the process for creating a " +
      "third-party Token Store extension. It will store codes and " +
      "tokens in a HashMap and generate UUID values. Note however that " +
      "realistic implementations must store their tokens and " +
      "authorization codes in persistent storage that is accessible to " +
      "all Identity Broker instances."
     };
  }

  /**
   * {@inheritDoc}
   */
  @Override
  public void initializeTokenStore(final IdentityBrokerContext serverContext,
                                   final TokenStoreConfig config,
                                   final ArgumentParser parser)
      throws LDAPException
  {
    this.serverContext = serverContext;
    this.config        = config;
  }

  /**
   * {@inheritDoc}
   */
  @Override
  public void finalizeTokenStore()
  {
    this.authzCodeMap.clear();
    this.accessTokenMap.clear();
    this.refreshTokenMap.clear();
  }

  /**
   * {@inheritDoc}
   */
  @Override
  public AuthorizationCode consumeAuthorizationCode(final String codeValue)
      throws OAuthException
  {
    return authzCodeMap.remove(codeValue);
  }

  /**
   * {@inheritDoc}
   */
  @Override
  public void storeAuthorizationCode(final AuthorizationCode authorizationCode)
      throws OAuthException
  {
    //First go through and remove any existing AuthorizationCodes for the same
    //application and username.
    Iterator<Map.Entry<String,AuthorizationCode>> iter =
            authzCodeMap.entrySet ().iterator();
    while (iter.hasNext())
    {
      Map.Entry<String,AuthorizationCode> entry = iter.next();
      AuthorizationCode c = entry.getValue();
      if (c.getApplicationId().equals(authorizationCode.getApplicationId())
              && c.getUsername().equals(authorizationCode.getUsername()))
      {
        iter.remove();
      }
    }

    String value = authorizationCode.getValue();
    if(value == null)
    {
      value = UUID.randomUUID().toString();
      authorizationCode.setValue(value);
    }
    authzCodeMap.put(value, authorizationCode);
  }

  /**
   * {@inheritDoc}
   */
  @Override
  public void storeAccessToken(final AccessToken accessToken)
      throws OAuthException
  {
    //First go through and remove any existing AccessTokens and RefreshTokens
    //for the same application and username.
    Iterator<Map.Entry<String, AccessToken>> iter =
            accessTokenMap.entrySet().iterator();
    while(iter.hasNext())
    {
      Map.Entry<String, AccessToken> entry = iter.next();
      AccessToken t = entry.getValue();
      if(t.getApplicationId().equals(accessToken.getApplicationId()) &&
         ( (t.getUsername() == null && accessToken.getUsername() == null) ||
           (t.getUsername() != null &&
            t.getUsername().equals(accessToken.getUsername())) ))
      {
        revokeToken(entry.getKey());
      }
    }

    String accessValue = accessToken.getValue();
    if(accessValue == null)
    {
      accessValue = UUID.randomUUID().toString();
      accessToken.setValue(accessValue);
    }
    accessTokenMap.put(accessValue, accessToken);

    if(accessToken.getRefreshToken() != null)
    {
      RefreshToken refreshToken = accessToken.getRefreshToken();
      String refreshValue = refreshToken.getValue();
      if(refreshValue == null)
      {
        refreshValue = UUID.randomUUID().toString();
        refreshToken.setValue(refreshValue);
      }
      refreshTokenMap.put(refreshValue, accessToken);
    }
  }

  /**
   * {@inheritDoc}
   */
  @Override
  public AccessToken getAccessToken(final String tokenValue)
      throws OAuthException
  {
    return accessTokenMap.get(tokenValue);
  }

  /**
   * {@inheritDoc}
   */
  @Override
  public AccessToken getAccessTokenFromRefreshToken(final String tokenValue)
      throws OAuthException
  {
    return refreshTokenMap.get(tokenValue);
  }

  /**
   * {@inheritDoc}
   */
  @Override
  public Collection<AccessToken> getAccessTokensByUser(final String username)
      throws OAuthException
  {
    // This is very inefficient for real world use. Implementations should
    // make this more efficient by using a database index on username for
    // instance.
    Set<AccessToken> tokens = new LinkedHashSet<AccessToken>();
    for(AccessToken accessToken : accessTokenMap.values())
    {
      if((accessToken.getUsername() != null &&
          accessToken.getUsername().equals(username)) ||
          (accessToken.getUsername() == null && username == null))
      {
        tokens.add(accessToken);
      }
    }
    return tokens;
  }

  /**
   * {@inheritDoc}
   */
  @Override
  public AccessToken revokeToken(final String tokenValue) throws OAuthException
  {
    AccessToken accessToken = accessTokenMap.remove(tokenValue);
    if(accessToken != null)
    {
      RefreshToken refreshToken = accessToken.getRefreshToken();
      if(refreshToken != null)
      {
        refreshTokenMap.remove(refreshToken.getValue());
      }
    }
    else
    {
      accessToken = refreshTokenMap.remove(tokenValue);
      if(accessToken != null)
      {
        accessTokenMap.remove(accessToken.getValue());
      }
    }
    return accessToken;
  }

  /**
   * {@inheritDoc}
   */
  @Override
  public void toString(final StringBuilder buffer)
  {
    buffer.append(getExtensionName());
  }
}