tuonq-registry/zen-api/src/main/kotlin/com/gxzc/zen/umps/config/ZenWebSessionManager.kt

package com.gxzc.zen.umps.config

import org.apache.shiro.session.ExpiredSessionException
import org.apache.shiro.session.InvalidSessionException
import org.apache.shiro.session.Session
import org.apache.shiro.session.UnknownSessionException
import org.apache.shiro.session.mgt.DefaultSessionManager
import org.apache.shiro.session.mgt.DelegatingSession
import org.apache.shiro.session.mgt.SessionContext
import org.apache.shiro.session.mgt.SessionKey
import org.apache.shiro.web.servlet.Cookie
import org.apache.shiro.web.servlet.ShiroHttpServletRequest
import org.apache.shiro.web.servlet.ShiroHttpSession
import org.apache.shiro.web.servlet.SimpleCookie
import org.apache.shiro.web.session.mgt.WebSessionKey
import org.apache.shiro.web.util.WebUtils
import org.slf4j.LoggerFactory
import java.io.Serializable
import javax.servlet.ServletRequest
import javax.servlet.ServletResponse
import javax.servlet.http.HttpServletRequest
import javax.servlet.http.HttpServletResponse


/**
 * Copied form @link DefaultWebSessionManager
 * 解决一次登陆多次直接读取缓存的问题
 * 在读取前先读取本地缓存(目前存放于spring-session 中 attribute)
 * @author NorthLan
 * @date 2018/4/25
 * @url https://noahlan.com
 */
@Suppress("unused")
class ZenWebSessionManager : DefaultSessionManager {
    companion object {
        private val log = LoggerFactory.getLogger(ZenWebSessionManager::class.java)
    }

    var sessionIdCookie: Cookie? = null
    var sessionIdCookieEnabled: Boolean = false
    var sessionIdUrlRewritingEnabled: Boolean = false

    constructor() {
        val cookie = SimpleCookie(ShiroHttpSession.DEFAULT_SESSION_ID_NAME)
        cookie.isHttpOnly = true //more secure, protects against XSS attacks
        this.sessionIdCookie = cookie
        this.sessionIdCookieEnabled = true
        this.sessionIdUrlRewritingEnabled = true
    }

    private fun storeSessionId(currentId: Serializable?, request: HttpServletRequest, response: HttpServletResponse) {
        if (currentId == null) {
            val msg = "sessionId cannot be null when persisting for subsequent requests."
            throw IllegalArgumentException(msg)
        }
        val template = sessionIdCookie
        val cookie = SimpleCookie(template)
        val idString = currentId.toString()
        cookie.value = idString
        cookie.saveTo(request, response)
        log.trace("Set session ID cookie for session with id {}", idString)
    }

    private fun removeSessionIdCookie(request: HttpServletRequest, response: HttpServletResponse) {
        sessionIdCookie?.removeFrom(request, response)
    }

    private fun getSessionIdCookieValue(request: ServletRequest, response: ServletResponse): String? {
        if (!sessionIdCookieEnabled) {
            log.debug("Session ID cookie is disabled - session id will not be acquired from a request cookie.")
            return null
        }
        if (request !is HttpServletRequest) {
            log.debug("Current request is not an HttpServletRequest - cannot get session ID cookie.  Returning null.")
            return null
        }
        return sessionIdCookie?.readValue(request, WebUtils.toHttp(response))
    }

    private fun getReferencedSessionId(request: ServletRequest, response: ServletResponse): Serializable? {

        var id = getSessionIdCookieValue(request, response)
        if (id != null) {
            request.setAttribute(ShiroHttpServletRequest.REFERENCED_SESSION_ID_SOURCE,
                    ShiroHttpServletRequest.COOKIE_SESSION_ID_SOURCE)
        } else {
            //not in a cookie, or cookie is disabled - try the request URI as a fallback (i.e. due to URL rewriting):

            //try the URI path segment parameters first:
            id = getUriPathSegmentParamValue(request, ShiroHttpSession.DEFAULT_SESSION_ID_NAME)

            if (id == null) {
                //not a URI path segment parameter, try the query parameters:
                val name = getSessionIdName()
                id = request.getParameter(name)
                if (id == null) {
                    //try lowercase:
                    id = request.getParameter(name.toLowerCase())
                }
            }
            if (id != null) {
                request.setAttribute(ShiroHttpServletRequest.REFERENCED_SESSION_ID_SOURCE,
                        ShiroHttpServletRequest.URL_SESSION_ID_SOURCE)
            }
        }
        if (id != null) {
            request.setAttribute(ShiroHttpServletRequest.REFERENCED_SESSION_ID, id)
            //automatically mark it valid here.  If it is invalid, the
            //onUnknownSession method below will be invoked and we'll remove the attribute at that time.
            request.setAttribute(ShiroHttpServletRequest.REFERENCED_SESSION_ID_IS_VALID, java.lang.Boolean.TRUE)
        }

        // always set rewrite flag - SHIRO-361
        request.setAttribute(ShiroHttpServletRequest.SESSION_ID_URL_REWRITING_ENABLED, sessionIdUrlRewritingEnabled)

        return id
    }

    //SHIRO-351
    //also see http://cdivilly.wordpress.com/2011/04/22/java-servlets-uri-parameters/
    //since 1.2.2
    private fun getUriPathSegmentParamValue(servletRequest: ServletRequest, paramName: String): String? {

        if (servletRequest !is HttpServletRequest) {
            return null
        }
        var uri: String? = servletRequest.requestURI ?: return null

        val queryStartIndex = uri!!.indexOf('?')
        if (queryStartIndex >= 0) { //get rid of the query string
            uri = uri.substring(0, queryStartIndex)
        }

        var index = uri.indexOf(';') //now check for path segment parameters:
        if (index < 0) {
            //no path segment params - return:
            return null
        }

        //there are path segment params, let's get the last one that may exist:

        val token = "$paramName="

        uri = uri.substring(index + 1) //uri now contains only the path segment params

        //we only care about the last JSESSIONID param:
        index = uri.lastIndexOf(token)
        if (index < 0) {
            //no segment param:
            return null
        }

        uri = uri.substring(index + token.length)

        index = uri.indexOf(';') //strip off any remaining segment params:
        if (index >= 0) {
            uri = uri.substring(0, index)
        }

        return uri //what remains is the value
    }

    override fun retrieveSession(sessionKey: SessionKey): Session? {
        val sessionId = getSessionId(sessionKey)
        if (sessionId == null) {
            log.debug("Unable to resolve session ID from SessionKey [{}].  Returning null to indicate a " + "session could not be found.", sessionKey)
            return null
        }
        ////////////////////////Add by noahlan//////////////////////////////////
        var request: ServletRequest? = null
        if (sessionKey is WebSessionKey) {
            request = sessionKey.servletRequest
        }
        if (request != null) {
            val s = request.getAttribute(sessionId.toString())
            if (s != null) {
                return s as Session
            }
        }
        ////////////////////////////////////////////////////////////////////////
        val s = retrieveSessionFromDataSource(sessionId)
        if (s == null) {
            //session ID was provided, meaning one is expected to be found, but we couldn't find one:
            val msg = "Could not find session with ID [$sessionId]"
            throw UnknownSessionException(msg)
        }

        ////////////////////////Add by noahlan//////////////////////////////////
        if (request != null) {
            request.setAttribute(sessionId.toString(), s)
        }
        ///////////////////////////////////////////////////////////////////////
        return s
    }

    //since 1.2.1
    private fun getSessionIdName(): String {
        var name: String? = this.sessionIdCookie?.name
        if (name == null) {
            name = ShiroHttpSession.DEFAULT_SESSION_ID_NAME
        }
        return name
    }

    override fun createExposedSession(session: Session, context: SessionContext?): Session {
        if (!WebUtils.isWeb(context)) {
            return super.createExposedSession(session, context)
        }
        val request = WebUtils.getRequest(context)
        val response = WebUtils.getResponse(context)
        val key = WebSessionKey(session.id, request, response)
        return DelegatingSession(this, key)
    }

    override fun createExposedSession(session: Session, key: SessionKey?): Session {
        if (!WebUtils.isWeb(key)) {
            return super.createExposedSession(session, key)
        }

        val request = WebUtils.getRequest(key)
        val response = WebUtils.getResponse(key)
        val sessionKey = WebSessionKey(session.id, request, response)
        return DelegatingSession(this, sessionKey)
    }

    /**
     * Stores the Session's ID, usually as a Cookie, to associate with future requests.
     *
     * @param session the session that was just [created][.createSession].
     */
    override fun onStart(session: Session, context: SessionContext?) {
        super.onStart(session, context)

        if (!WebUtils.isHttp(context)) {
            log.debug("SessionContext argument is not HTTP compatible or does not have an HTTP request/response " + "pair. No session ID cookie will be set.")
            return

        }
        val request = WebUtils.getHttpRequest(context)
        val response = WebUtils.getHttpResponse(context)

        if (sessionIdCookieEnabled) {
            val sessionId = session.id
            storeSessionId(sessionId, request, response)
        } else {
            log.debug("Session ID cookie is disabled.  No cookie has been set for new session with id {}", session!!.id)
        }

        request.removeAttribute(ShiroHttpServletRequest.REFERENCED_SESSION_ID_SOURCE)
        request.setAttribute(ShiroHttpServletRequest.REFERENCED_SESSION_IS_NEW, java.lang.Boolean.TRUE)
    }

    public override fun getSessionId(key: SessionKey): Serializable? {
        var id: Serializable? = super.getSessionId(key)
        if (id == null && WebUtils.isWeb(key)) {
            val request = WebUtils.getRequest(key)
            val response = WebUtils.getResponse(key)
            id = getSessionId(request, response)
        }
        return id
    }

    protected fun getSessionId(request: ServletRequest, response: ServletResponse): Serializable? {
        return getReferencedSessionId(request, response)
    }

    override fun onExpiration(s: Session, ese: ExpiredSessionException?, key: SessionKey?) {
        super.onExpiration(s, ese, key)
        onInvalidation(key)
    }

    override fun onInvalidation(session: Session, ise: InvalidSessionException, key: SessionKey) {
        super.onInvalidation(session, ise, key)
        onInvalidation(key)
    }

    private fun onInvalidation(key: SessionKey?) {
        val request = WebUtils.getRequest(key)
        request?.removeAttribute(ShiroHttpServletRequest.REFERENCED_SESSION_ID_IS_VALID)
        if (WebUtils.isHttp(key)) {
            log.debug("Referenced session was invalid.  Removing session ID cookie.")
            removeSessionIdCookie(WebUtils.getHttpRequest(key), WebUtils.getHttpResponse(key))
        } else {
            log.debug("SessionKey argument is not HTTP compatible or does not have an HTTP request/response " + "pair. Session ID cookie will not be removed due to invalidated session.")
        }
    }

    override fun onStop(session: Session, key: SessionKey?) {
        super.onStop(session, key)
        if (WebUtils.isHttp(key)) {
            val request = WebUtils.getHttpRequest(key)
            val response = WebUtils.getHttpResponse(key)
            log.debug("Session has been stopped (subject logout or explicit stop).  Removing session ID cookie.")
            removeSessionIdCookie(request, response)
        } else {
            log.debug("SessionKey argument is not HTTP compatible or does not have an HTTP request/response " + "pair. Session ID cookie will not be removed due to stopped session.")
        }
    }

    /**
     * This is a native session manager implementation, so this method returns `false` always.
     *
     * @return `false` always
     * @since 1.2
     */
    fun isServletContainerSessions(): Boolean {
        return false
    }
}