3K整合系列(三) Ktor 角色权限控制

3K = Kotlin + Ktor + Ktorm，不要记错了哦

在上一篇里，我们成功整合了 Ktor 和 Ktorm，并完成了一个简单的用户登录登出。在现实情况里，用户的权限控制一直是非常重要的事情，对于没有权限的用户，某些接口就不能放行。

在 Ktor 里面，我们可以通过一种很简单的方式来实现它，比如说:

get("/sample") {
    val u = user
    val perms = UserMapper.getPermissions(u.userId)
    if (!perms.contains("perm:user:info")) {
        call.respond(AjaxResult.error("你没有权限访问这个接口"))
        return@get
    }
    ... ...
}

是的，这是一般的实现方法，但是对于大部分接口都要写这些代码，就非常的不友好了，代码量太多，也不好控制，而且重复的代码非常让人厌烦。所以在这里，我们需要用 AOP 的方法去解决，按以往对 Ktor 插件的了解，在这个场景下，我们也应该通过编写插件来解决问题，那么下面就正式开始吧。

首先我们要知道，Ktor 官方是有一套插件的标准格式的，只有这样写，Ktor 才会承认它是一个合法的插件:

class RoleBasedAuthorization(internal var config: RoleAuthorizationConfig) {

    fun configure(block: RoleAuthorizationConfig.() -> Unit) {
        val newConfig = config.copy()
        block(newConfig)
        config = newConfig.copy()
    }

    companion object : BaseApplicationPlugin<Application, RoleAuthorizationConfig, RoleBasedAuthorization> {
        override val key: AttributeKey<RoleBasedAuthorization> = AttributeKey("RoleAuthorizationHolder")

        override fun install(pipeline: Application, configure: RoleAuthorizationConfig.() -> Unit): RoleBasedAuthorization {
            val cfg = RoleAuthorizationConfig().apply(configure)
            return RoleBasedAuthorization(cfg)
        }
    }
}

由于我们需要在插件里面能够动态获取用户的实际权限，因此在 Config 里面，就要把获取权限的函数，以及权限校验失败的函数予以写出:

class RoleAuthorizationConfig(
    var internalGetRoles: (Principal) -> Set<String> = { emptySet() },
    var internalRoleAuthFailed: suspend ApplicationCall.(message: String) -> Unit = {}) {

    fun getRoles(block: (Principal) -> Set<String>) {
        internalGetRoles = block
    }

    fun roleAuthFailed(block: suspend ApplicationCall.(message: String) -> Unit) {
        internalRoleAuthFailed = block
    }

    internal fun copy(): RoleAuthorizationConfig = RoleAuthorizationConfig(internalGetRoles, internalRoleAuthFailed)
}

首次看到这种写法的同学请不要惊讶，这只是 Kotlin 强大的语言特性之一，你可以把匿名函数作为类型来使用。并且，如果你觉得每次都要写函数定义很麻烦，也可以将函数定义成类型:

typealias GetRoleFunc = (Principal) -> Set<String>

好了，到这里配置的部分已经完成了，我们可以向 Ktor 注册这个插件:

install(RoleBasedAuthorization) {
    getRole {
        it as SessionUser
        UserMapper.getPermissions(it.userId)
    }
    roleAuthFailed {
        respond(AjaxResult.error(it))
    }
}

是不是很简单，直接在插件里面就可以定义用于获取权限的方法，以及在权限验证失败后如何返回。下面我们就要将这个权限验证挂到某个路由上，也就是在指定的路由上 AOP 这个权限验证。

为了直观起见，我们先定义好 AOP 的方式吧，就以上面的 get 方法来说，加入权限验证后，写法变成这样:

withRoles("perm:user:info") {
    get("/sample") {
        ... ...
    }
}

这样看起来就一目了然了，不会对 get 内部的代码造成侵入。所以现在要写的，就是这个 withRoles 方法了。同样的，我们先写好架子：

fun Route.withRoles(vararg roles: String, build: Route.() -> Unit): Route {
    val authorizedRoute = createChild(RoleAuthorizationRouteSelector(roles.joinToString(",")))
    authorizedRoute.install(RoleAuthenticationInterceptors) {
        this.roles = roles.toSet()
    }
    authorizedRoute.build()
    return authorizedRoute
}

这些代码也很好理解，为当前的路由创建一个子路由，然后向该子路由安装一个插件，并给定插件的参数，最后构建并返回该子路由。这里可以很清晰的看到 AOP 的过程了，不同于以往我们用注解写 AOP，Ktor 的插件可是实打实的代码，每一步都让你看得清清楚楚。

关于这个 RoleAuthorizationRouteSelector，它是一个路由的选择器，有几种可选的模式，请参考表格:

选择器模式	含义
Failed	路由未找到时可选
FailedPath	路由未找到时可选(与Failed相同)
FailedMethod	请求方式不允许时可选
FailedParameter	请求参数错误时可选
Missing	有可选参数未填时可选
Constant	有静态值被传入时可选
Transparent	不会改变原有路由的入参情况时可选(通常是选这个)
ConstantPath	针对单个路由传入静态值时可选
WildcardPath	针对通配路由时可选

在这个表的基础上，我们可以很轻松的写出一个路由选择器:

private class RoleAuthorizationRouteSelector(private val description: String) : RouteSelector() {
    override fun evaluate(context: RoutingResolveContext, segmentIndex: Int): RouteSelectorEvaluation =
        RouteSelectorEvaluation.Transparent

    override fun toString(): String = "(authorize ${description})"
}

然后，RoleAuthenticationInterceptors 才是真正处理 AOP 动作的插件，从它的命名上也可以看出来，这是个拦截器，把送往路由的请求先进行拦截，经过一系列操作后，再予以放行或者不放行。拦截器的代码也是基本上固定的套路，如下所示:

private val RoleAuthenticationInterceptors: RouteScopedPlugin<RouteRoleAuthorizationConfig> =
    createRouteScopedPlugin("RoleAuthenticationInterceptors", ::RouteRoleAuthorizationConfig) {
        on(AuthenticationChecked) { call ->
            val reqRoles = pluginConfig.roles
            val authConfig = call.application.plugin(RoleBasedAuthorization).config
            val reqGetRoles = authConfig.internalGetRoles
            val reqOnFail = authConfig.internalRoleAuthFailed
            val user = call.principal<Principal>()
            if (user == null) {
                reqOnFail(call, "Unauthenticated User")
                return@on
            }
            val roles = reqGetRoles(user)
            val missing = reqRoles - roles
            if (missing.isNotEmpty()) {
                reqOnFail(call, "Principal $user lacks required role(s) ${missing.joinToString(" and ")}")
            }
        }
    }

Ktor 提供了两种获取配置的方式，一种是通过 pluginConfig，这种方式可以获得在路由上动态挂载的配置内容，另一种是通过 call.application.plugin 来获取，这种方式可以获得通过静态的 install 安装的插件里所配置的内容。

好了，基本上也算是写完了，那就简单的测试一下吧:

$ curl http://0.0.0.0:8080/sample -b cookie1
{
    "code": 500,
    "message": "请求失败",
    "data": "Principal (userId=1, userName=admin) lacks required role(s) perm:user:info"
}

$ curl http://0.0.0.0:8080/sample -b cookie2
{
    "code": 200,
    "message": "请求成功",
    "data": {
        "userId": 1,
        "userName": "admin"
    }
}

似乎到了这里就已经完成了对用户权限的控制，但是我们是要精益求精的，比如说这段代码，你看着没觉得不舒服么:

install(RoleBasedAuthorization) {
    getRole {
        it as SessionUser
        UserMapper.getPermissions(it.userId)
    }
}

为什么还要 it as SessionUser 呢，为什么不是直接输出一个 SessionUser 类型的对象？

可能你会说，看了上面的代码，似乎也没有哪个地方可以让我塞进一个泛型呀。比如说想实现以下代码是不能的:

class RoleAuthorizationConfig<T: Principal>(
    var internalGetRoles: (T) -> Set<String> = { emptySet() },
    fun getRoles(block: (T) -> Set<String>) {
        internalGetRoles = block
    }
    internal fun copy(): RoleAuthorizationConfig = RoleAuthorizationConfig(internalGetRoles, internalRoleAuthFailed)
}

这个代码看起来是可行的，但是却没有办法塞在 RoleBasedAuthorization 里，这意味着要让 RoleBasedAuthorization 也带上泛型，然而 Ktor 对于插件的要求又是不允许带有泛型，这可如何是好？

答案是，在 Kotlin 里面是可以变魔术的，插件本体不允许带泛型，可没说插件里面的各个东西不能带吧，正是在这个指导思想下，我们可以实现大魔术。以下的代码是完整的角色权限插件，大家也可以亲自体会一下这个魔术的原理:

enum class RoleAuthorizationType { ALL, ANY, NONE }

class RoleBasedAuthorization(internal var config: RoleBasedAuthorizationConfig) {
    fun configure(block: RoleBasedAuthorizationConfig.() -> Unit) {
        val newConfiguration = config.copy()
        block(newConfiguration)
        config = newConfiguration
    }

    companion object : BaseApplicationPlugin<Application, RoleBasedAuthorizationConfig, RoleBasedAuthorization> {
        override val key: AttributeKey<RoleBasedAuthorization> = AttributeKey("RoleBasedAuthorizationHolder")

        override fun install(pipeline: Application, configure: RoleBasedAuthorizationConfig.() -> Unit): RoleBasedAuthorization {
            val config = RoleBasedAuthorizationConfig().apply(configure)
            return RoleBasedAuthorization(config)
        }
    }
}

class RoleBasedAuthorizationConfig(
        var type: RoleAuthorizationType = RoleAuthorizationType.ANY,
        var roles: Set<String> = emptySet(),
        var provider: BaseRoleBasedAuthorizationProvider? = null) {
    internal fun copy(): RoleBasedAuthorizationConfig = RoleBasedAuthorizationConfig(type, roles, provider)
}

val RoleAuthenticationInterceptors: RouteScopedPlugin<RoleBasedAuthorizationConfig> =
        createRouteScopedPlugin("RoleAuthenticationInterceptors", ::RoleBasedAuthorizationConfig) {
            // 这种方式获取真实使用时的配置内容
            val reqRoles = pluginConfig.roles
            val reqType = pluginConfig.type
            // 这种方式获取 install 时注册的配置内容
            val config = application.plugin(RoleBasedAuthorization).config
            on(AuthenticationChecked) { call ->
                if (call.isHandled) {
                    return@on
                }
                config.provider?.onRoleAuthorization(call, reqType, reqRoles)
            }
        }

private class RoleAuthorizationRouteSelector(private val description: Set<String>) : RouteSelector() {
    override fun evaluate(context: RoutingResolveContext, segmentIndex: Int): RouteSelectorEvaluation =
            RouteSelectorEvaluation.Transparent

    override fun toString(): String = "(authorize ${description.joinToString(",")})"
}

private fun Route.withRoles(roles: Set<String>, roleAuthType: RoleAuthorizationType, build: Route.() -> Unit): Route {
    require(roles.isNotEmpty()) { "At least one role name need to be provided" }
    val roleAuthRoute = createChild(RoleAuthorizationRouteSelector(roles))
    roleAuthRoute.install(RoleAuthenticationInterceptors) {
        this.roles = roles
        this.type = roleAuthType
    }
    roleAuthRoute.build()
    return roleAuthRoute
}


abstract class BaseRoleBasedAuthorizationProvider {
    abstract suspend fun onRoleAuthorization(call: ApplicationCall, reqType: RoleAuthorizationType, reqRoles: Set<String>)
    open class Config protected constructor()
}

typealias RoleBasedAuthorizationGetRoleFunc<T> = suspend ApplicationCall.(T) -> Set<String>
typealias RoleBasedAuthorizationRoleAuthFailedFunc = suspend ApplicationCall.(String) -> Unit

class RoleBasedAuthorizationProvider<T : Any>(config: Config<T>) : BaseRoleBasedAuthorizationProvider() {

    private val type: KClass<T> = config.type
    private val getRoleFunc: RoleBasedAuthorizationGetRoleFunc<T> = config.getRoleFunc
    private val roleAuthFailedFunc: RoleBasedAuthorizationRoleAuthFailedFunc = config.roleAuthFailedFunc

    override suspend fun onRoleAuthorization(call: ApplicationCall, reqType: RoleAuthorizationType, reqRoles: Set<String>) {
        val user = call.sessions.get(type)
        if (user == null) {
            roleAuthFailedFunc(call, "Unauthenticated User")
            return
        }
        val roles = getRoleFunc(call, user)
        val denyReasons = mutableListOf<String>()
        when (reqType) {
            RoleAuthorizationType.ALL -> {
                val missing = reqRoles - roles
                if (missing.isNotEmpty()) {
                    denyReasons += "Principal $user lacks required role(s) ${missing.joinToString(" and ")}"
                }
            }

            RoleAuthorizationType.ANY -> {
                if (roles.none { it in reqRoles }) {
                    denyReasons += "Principal $user has none of the sufficient role(s) ${reqRoles.joinToString(" or ")}"
                }
            }

            RoleAuthorizationType.NONE -> {
                if (roles.any { it in reqRoles }) {
                    denyReasons += "Principal $user has forbidden role(s) ${reqRoles.intersect(roles).joinToString(" and ")}"
                }
            }
        }
        if (denyReasons.isNotEmpty()) {
            val message = denyReasons.joinToString(". ")
            roleAuthFailedFunc(call, message)
        }
    }

    class Config<T : Any>(internal val type: KClass<T>) : BaseRoleBasedAuthorizationProvider.Config() {
        internal var getRoleFunc: RoleBasedAuthorizationGetRoleFunc<T> = { emptySet() }
        internal var roleAuthFailedFunc: RoleBasedAuthorizationRoleAuthFailedFunc = {}
        fun getRole(block: RoleBasedAuthorizationGetRoleFunc<T>) {
            getRoleFunc = block
        }

        fun roleAuthFailed(block: RoleBasedAuthorizationRoleAuthFailedFunc) {
            roleAuthFailedFunc = block
        }

        fun buildProvider(): RoleBasedAuthorizationProvider<T> = RoleBasedAuthorizationProvider(this)
    }

}

inline fun <reified T : Any> RoleBasedAuthorizationConfig.roleSession(configure: RoleBasedAuthorizationProvider.Config<T>.() -> Unit) {
    val provider = RoleBasedAuthorizationProvider.Config(T::class).apply(configure).buildProvider()
    this.provider = provider
}

fun Route.withRoles(vararg roles: String, build: Route.() -> Unit) =
        withRoles(roles.toSet(), RoleAuthorizationType.ALL, build)

fun Route.withAnyRole(vararg roles: String, build: Route.() -> Unit) =
        withRoles(roles.toSet(), RoleAuthorizationType.ANY, build)

fun Route.withoutRoles(vararg roles: String, build: Route.() -> Unit) =
        withRoles(roles.toSet(), RoleAuthorizationType.NONE, build)

这里通过一个不带泛型的 Base Provider 和一个带泛型的实体 Provider 来实现了将泛型带进插件里。于是我们可以通过这样的代码来写:

install(RoleBasedAuthorization) {
    roleSession<SessionUser> {
        getRole {
            UserMapper.getPermissions(it.userId)
        }
    }
}

为了进一步方便起见，可以封装一个函数:

inline fun <reified T : Principal> Application.pluginRoleAuthorization(
        crossinline configure: RoleBasedAuthorizationProvider.Config<T>.() -> Unit
) = install(RoleBasedAuthorization) {
    roleSession<T>(configure)
}

这样我们就可以用一种非常轻松的方法来使用插件了:

pluginRoleAuthorization<SessionUser> {
    getRole {
        UserMapper.getInfo(it.userId)?.perms ?: setOf()
    }
    roleAuthFailed {
        respond(AjaxResult.error("您没有权限访问这个接口"))
    }
}

好了，用户角色权限控制到这里就结束了，愉快的使用封装好的代码吧，如果后续还需要其他的 AOP，也可以通过同样的方式来实现插件，熟悉 Ktor 的插件机制，对于灵活使用这个框架有着相当大的好处。

下一篇将讲述如何在 Ktor 内使用 JNI，从而实现与原生层交互，同时也将讲述如何使用 Kotlin 本身的代码来编写一个标准的 JNI 库。