ActionHandler for Remember Me cookie based authentication

ActionHandler.scala

/**
 * Remember me cookie based authentication using Action Composition in Play 2.0.  
 *
 * Code based on blog post from http://jaspan.com/improved_persistent_login_cookie_best_practice
 *
 * Create an object or class with this trait, and put it in your Global.onRouteRequest like so:
 * 
 * <pre>
  override def onRouteRequest(request: RequestHeader): Option[Handler] = {
    super.onRouteRequest(request).map {
      handler =>
        logger.info("onRouteRequest: request = " + request)

        handler match {
          case a: Action[_] => ActionHandler(a)
          case _ => handler
        }
    }
  }
 </pre>
 * @author wsargent
 * @since 6/28/12
 */
trait ActionHandler {

  def userInfoService: UserInfoService

  def sessionStore: SessionStore

  def logger: Logger

  def actionHandler[A](action: Action[A]): Action[A] = {
    Action(action.parser) {
      rawRequest =>
        if (requiresContext(rawRequest)) {
          actionWithContext(rawRequest) {
            request => action(request)
          }
        } else {
          logger.trace("actionHandler: no context required for {0}", rawRequest)
          action(rawRequest)
        }
    }
  }

  def gotoSuspiciousAuthDetected[A](request: Request[A]): Result = {
    AuthController.suspiciousActivity(request)
  }

  /**
   * Enhances the action with a request object that can contain important context information.
   */
  def actionWithContext[A](rawRequest: Request[A])(action: Request[A] => Result): Result = {
    logger.trace("actionWithContext: request = {0}", rawRequest)

    // Look for session credentials first, then cookie.
    withSessionCredentials(rawRequest).map {
      request =>
        action(request)
    }.map {
      logger.debug("actionWithContext: returning with session credentials " + rawRequest)
      return _
    }

    withRememberMeCookie(rawRequest).map {
      rememberMe =>
        for {
          userId <- rememberMe.userId
          series <- rememberMe.series
          token <- rememberMe.token
        } yield {
          val result = userInfoService.authenticateWithCookie(userId, series, token).fold(
            fault => actionRejectingAuthentication(rawRequest) {
              request =>
                fault match {
                  case suspiciousFault:InvalidSessionCookieFault => gotoSuspiciousAuthDetected(request)
                  case _ => action(request)
                }
            },
            event => actionWithAuthentication(rawRequest, event) {
              request =>
                action(request)
            }
          )
          logger.debug("actionWithContext: returning with token credentials {0}", rawRequest)
          return result
        }
    }

    // Return the default.
    logger.debug("actionWithContext: returning with no credentials {0}", rawRequest)
    action(Context(rawRequest, None))
  }

  def withSessionCredentials[A](rawRequest: Request[A]): Option[Context[A]] = {
    for {
      sessionId <- rawRequest.session.get("sessionId")
      userInfo <- restoreFromSession(sessionId)
    } yield Context(rawRequest, Some(userInfo))
  }

  def restoreFromSession(sessionId: String): Option[UserInfo] = {
    for {
      userId <- sessionStore.lookup(sessionId)
      userInfo <- userInfoService.lookup(userId)
    } yield userInfo
  }

  def actionRejectingAuthentication[A](request: Request[A])(action: Request[A] => Result): Result = {
    val richRequest = Context(request, None)
    val result = action(richRequest)
    result match {
      case plainResult: PlainResult => {
        logger.debug("discarding remember me cookie")
        plainResult.discardingCookies(RememberMe.COOKIE_NAME)
      }
      case _ => result
    }
  }

  def actionWithAuthentication[A](rawRequest: Request[A], event: UserAuthenticatedWithTokenEvent)(action: Request[A] => Result): Result = {
    // Defer the action until we have the authentication saved off...
    val userId = event.aggregateId
    val sessionId = saveAuthentication(userId)(rawRequest)
    val sessionCookie = SessionCookie("sessionId", sessionId)(rawRequest)

    // Create a new token on every cookie authentication, even if the series is the same.
    val rememberMe = RememberMe(userId, event.series, event.token)
    val rememberMeCookie = RememberMe.encodeAsCookie(rememberMe)

    // Look up and save off the user information for the rest of the action...
    val userInfo = userInfoService.lookup(userId)
    val richRequest = Context(rawRequest, userInfo)
    val result = action(richRequest)
    result match {
      case plainResult: PlainResult => {
        logger.debug("Creating new remember me cookie {0}", rememberMe)
        plainResult.withCookies(sessionCookie, rememberMeCookie)
      }
      case _ => result
    }
  }

  def requiresContext(implicit req: RequestHeader): Boolean = {
    !req.path.startsWith("/assets")
  }

  def withRememberMeCookie(implicit req: RequestHeader): Option[RememberMe] = {
    val cookie = req.cookies.get(RememberMe.COOKIE_NAME)
    val rememberMe = RememberMe.decodeFromCookie(cookie)
    // Cookie can be empty even if it exists :-(
    if (!rememberMe.isEmpty) Option(rememberMe) else None
  }

  def saveAuthentication(userId: UUID)(implicit req: RequestHeader): String = {
    logger.debug("saveAuthentication: {0}", userId)
    sessionStore.saveSession(UUID.randomUUID.toString, userId, req)
  }

}

Context.scala

case class Context[A](request: Request[A], me: Option[UserInfo]) extends WrappedRequest(request) {
  override def toString = {
    "Context(" + method + " " + uri + " user=" + me + ")"
  }
}

RememberMe.scala

case class RememberMe(data: Map[String, String] = Map.empty[String, String]) {

  import RememberMe._

  def get(key: String) = data.get(key)

  def isEmpty: Boolean = data.isEmpty

  def +(kv: (String, String)) = copy(data + kv)

  def -(key: String) = copy(data - key)

  def series: Option[Long] = AsLong(data.get(SERIES_NAME))

  def userId: Option[UUID] = AsUUID(data.get(USER_ID_NAME))

  def token: Option[Long] = AsLong(data.get(TOKEN_NAME))

  def apply(key: String) = data(key)
}

object RememberMe extends CookieBaker[RememberMe] {

  def apply(userId: util.UUID, series: Long, token: Long): RememberMe = {
    val map = Map(
      RememberMe.USER_ID_NAME -> userId.toString,
      RememberMe.SERIES_NAME -> series.toString,
      RememberMe.TOKEN_NAME -> token.toString)
    RememberMe(map)
  }

  val COOKIE_NAME = "REMEMBER_ME"

  val SERIES_NAME = "series"
  val USER_ID_NAME = "userId"
  val TOKEN_NAME = "token"

  val DEFAULT_MAX_AGE = 60*60*24*365 // Set the cookie max age to 1 year

  val emptyCookie = new RememberMe

  override val isSigned = true
  override val secure = Play.maybeApplication.flatMap(_.configuration.getBoolean("rememberMe.secure")).getOrElse(false)
  override val maxAge = Play.maybeApplication.flatMap(_.configuration.getInt("rememberMe.maxAge")).getOrElse(DEFAULT_MAX_AGE)

  def deserialize(data: Map[String, String]) = new RememberMe(data)

  def serialize(rememberme: RememberMe) = rememberme.data
}

SessionCookie.scala

object SessionCookie {

  private val domainRegex = """^.+(\.[^\.]+\.[^\.]+)$""".r

  private def domain(req: RequestHeader): String =
    domainRegex.replaceAllIn(req.domain, _ group 1)

  def apply(name: String, value: String)(implicit req: RequestHeader): Cookie = {
    val data = req.session + (name -> value)
    val encoded = Session.encode(Session.serialize(data))
    Cookie(Session.COOKIE_NAME, encoded, domain = Option(domain(req)))
  }
}