feat: support for refresh_token grant, some more tests and test utils (#12)

tommytroen · web-flow · commit 94fcbfa390d1 · 2020-10-07T15:12:28.000+02:00
* feat: support for refresh_token grant, some more tests and test utils
* chore: suppress warnings for functions et al which is part of public API
* chore: reduce visibility for grant handler classes
diff --git a/src/main/kotlin/no/nav/security/mock/oauth2/MockOAuth2Server.kt b/src/main/kotlin/no/nav/security/mock/oauth2/MockOAuth2Server.kt
@@ -30,6 +30,7 @@ import okhttp3.mockwebserver.MockWebServer
 import okhttp3.mockwebserver.RecordedRequest
 
 // TODO make open so others can extend?
+@Suppress("unused", "MemberVisibilityCanBePrivate")
 class MockOAuth2Server(
     val config: OAuth2Config = OAuth2Config()
 ) {
diff --git a/src/main/kotlin/no/nav/security/mock/oauth2/OAuth2Exception.kt b/src/main/kotlin/no/nav/security/mock/oauth2/OAuth2Exception.kt
@@ -4,6 +4,7 @@ import com.nimbusds.oauth2.sdk.ErrorObject
 import com.nimbusds.oauth2.sdk.GrantType
 import com.nimbusds.oauth2.sdk.OAuth2Error
 
+@Suppress("unused")
 class OAuth2Exception(val errorObject: ErrorObject?, msg: String, throwable: Throwable?) :
     RuntimeException(msg, throwable) {
     constructor(msg: String) : this(null, msg, null)
diff --git a/src/main/kotlin/no/nav/security/mock/oauth2/extensions/String.kt b/src/main/kotlin/no/nav/security/mock/oauth2/extensions/String.kt
@@ -0,0 +1,14 @@
+package no.nav.security.mock.oauth2.extensions
+
+import java.net.URLDecoder
+import java.nio.charset.StandardCharsets
+
+internal fun String.keyValuesToMap(listDelimiter: String): Map<String, String> =
+    this.split(listDelimiter)
+        .filter { it.contains("=") }
+        .associate {
+            val (key, value) = it.split("=")
+            key.urlDecode().trim() to value.urlDecode().trim()
+        }
+
+internal fun String.urlDecode(): String = URLDecoder.decode(this, StandardCharsets.UTF_8)
diff --git a/src/main/kotlin/no/nav/security/mock/oauth2/grant/AuthorizationCodeGrantHandler.kt b/src/main/kotlin/no/nav/security/mock/oauth2/grant/AuthorizationCodeGrantHandler.kt
@@ -6,7 +6,6 @@ import com.nimbusds.oauth2.sdk.OAuth2Error
 import com.nimbusds.oauth2.sdk.TokenRequest
 import com.nimbusds.openid.connect.sdk.AuthenticationRequest
 import com.nimbusds.openid.connect.sdk.AuthenticationSuccessResponse
-import java.util.UUID
 import kotlin.collections.set
 import mu.KotlinLogging
 import no.nav.security.mock.oauth2.OAuth2Exception
@@ -21,8 +20,9 @@ import okhttp3.HttpUrl
 
 private val log = KotlinLogging.logger {}
 
-class AuthorizationCodeHandler(
-    private val tokenProvider: OAuth2TokenProvider = OAuth2TokenProvider()
+internal class AuthorizationCodeHandler(
+    private val tokenProvider: OAuth2TokenProvider,
+    private val refreshTokenManager: RefreshTokenManager
 ) : GrantHandler {
 
     private val codeToAuthRequestCache: MutableMap<AuthorizationCode, AuthenticationRequest> = HashMap()
@@ -67,14 +67,15 @@ class AuthorizationCodeHandler(
         val scope: String? = tokenRequest.scope?.toString()
         val nonce: String? = authenticationRequest?.nonce?.value
         val loginTokenCallbackOrDefault = getLoginTokenCallbackOrDefault(code, oAuth2TokenCallback)
-        val idToken: SignedJWT = tokenProvider.idToken(tokenRequest, issuerUrl, nonce, loginTokenCallbackOrDefault)
+        val idToken: SignedJWT = tokenProvider.idToken(tokenRequest, issuerUrl, loginTokenCallbackOrDefault, nonce)
         val accessToken: SignedJWT = tokenProvider.accessToken(tokenRequest, issuerUrl, loginTokenCallbackOrDefault, nonce)
+        val refreshToken: RefreshToken = refreshTokenManager.refreshToken(loginTokenCallbackOrDefault)
 
         return OAuth2TokenResponse(
             tokenType = "Bearer",
             idToken = idToken.serialize(),
             accessToken = accessToken.serialize(),
-            refreshToken = UUID.randomUUID().toString(),
+            refreshToken = refreshToken,
             expiresIn = idToken.expiresIn(),
             scope = scope
         )
diff --git a/src/main/kotlin/no/nav/security/mock/oauth2/grant/ClientCredentialsGrantHandler.kt b/src/main/kotlin/no/nav/security/mock/oauth2/grant/ClientCredentialsGrantHandler.kt
@@ -7,7 +7,7 @@ import no.nav.security.mock.oauth2.token.OAuth2TokenCallback
 import no.nav.security.mock.oauth2.token.OAuth2TokenProvider
 import okhttp3.HttpUrl
 
-class ClientCredentialsGrantHandler(
+internal class ClientCredentialsGrantHandler(
     private val tokenProvider: OAuth2TokenProvider
 ) : GrantHandler {
 
diff --git a/src/main/kotlin/no/nav/security/mock/oauth2/grant/JwtBearerGrantHandler.kt b/src/main/kotlin/no/nav/security/mock/oauth2/grant/JwtBearerGrantHandler.kt
@@ -12,7 +12,7 @@ import no.nav.security.mock.oauth2.token.OAuth2TokenCallback
 import no.nav.security.mock.oauth2.token.OAuth2TokenProvider
 import okhttp3.HttpUrl
 
-class JwtBearerGrantHandler(private val tokenProvider: OAuth2TokenProvider) : GrantHandler {
+internal class JwtBearerGrantHandler(private val tokenProvider: OAuth2TokenProvider) : GrantHandler {
 
     override fun tokenResponse(
         request: OAuth2HttpRequest,
diff --git a/src/main/kotlin/no/nav/security/mock/oauth2/grant/RefreshTokenGrantHandler.kt b/src/main/kotlin/no/nav/security/mock/oauth2/grant/RefreshTokenGrantHandler.kt
@@ -0,0 +1,48 @@
+package no.nav.security.mock.oauth2.grant
+
+import com.nimbusds.jwt.SignedJWT
+import com.nimbusds.oauth2.sdk.GrantType
+import com.nimbusds.oauth2.sdk.RefreshTokenGrant
+import com.nimbusds.oauth2.sdk.TokenRequest
+import mu.KotlinLogging
+import no.nav.security.mock.oauth2.extensions.expiresIn
+import no.nav.security.mock.oauth2.http.OAuth2HttpRequest
+import no.nav.security.mock.oauth2.http.OAuth2TokenResponse
+import no.nav.security.mock.oauth2.invalidGrant
+import no.nav.security.mock.oauth2.token.OAuth2TokenCallback
+import no.nav.security.mock.oauth2.token.OAuth2TokenProvider
+import okhttp3.HttpUrl
+
+private val log = KotlinLogging.logger {}
+
+internal class RefreshTokenGrantHandler(
+    private val tokenProvider: OAuth2TokenProvider,
+    private val refreshTokenManager: RefreshTokenManager
+) : GrantHandler {
+
+    override fun tokenResponse(
+        request: OAuth2HttpRequest,
+        issuerUrl: HttpUrl,
+        oAuth2TokenCallback: OAuth2TokenCallback
+    ): OAuth2TokenResponse {
+        val tokenRequest = request.asNimbusTokenRequest()
+        val refreshToken = tokenRequest.refreshTokenGrant().refreshToken.value
+        log.debug("issuing token for refreshToken=$refreshToken")
+        val scope: String? = tokenRequest.scope?.toString()
+        val refreshTokenCallbackOrDefault = refreshTokenManager[refreshToken] ?: oAuth2TokenCallback
+        val idToken: SignedJWT = tokenProvider.idToken(tokenRequest, issuerUrl, refreshTokenCallbackOrDefault)
+        val accessToken: SignedJWT = tokenProvider.accessToken(tokenRequest, issuerUrl, refreshTokenCallbackOrDefault)
+
+        return OAuth2TokenResponse(
+            tokenType = "Bearer",
+            idToken = idToken.serialize(),
+            accessToken = accessToken.serialize(),
+            refreshToken = refreshToken,
+            expiresIn = idToken.expiresIn(),
+            scope = scope
+        )
+    }
+
+    private fun TokenRequest.refreshTokenGrant(): RefreshTokenGrant =
+        (this.authorizationGrant as? RefreshTokenGrant) ?: invalidGrant(GrantType.REFRESH_TOKEN)
+}
diff --git a/src/main/kotlin/no/nav/security/mock/oauth2/grant/RefreshTokenManager.kt b/src/main/kotlin/no/nav/security/mock/oauth2/grant/RefreshTokenManager.kt
@@ -0,0 +1,18 @@
+package no.nav.security.mock.oauth2.grant
+
+import java.util.UUID
+import no.nav.security.mock.oauth2.token.OAuth2TokenCallback
+
+typealias RefreshToken = String
+
+internal data class RefreshTokenManager(
+    private val cache: MutableMap<RefreshToken, OAuth2TokenCallback> = HashMap()
+) {
+    operator fun get(refreshToken: RefreshToken) = cache[refreshToken]
+
+    fun refreshToken(tokenCallback: OAuth2TokenCallback): RefreshToken {
+        val refreshToken = UUID.randomUUID().toString()
+        cache[refreshToken] = tokenCallback
+        return refreshToken
+    }
+}
diff --git a/src/main/kotlin/no/nav/security/mock/oauth2/grant/TokenExchangeGrant.kt b/src/main/kotlin/no/nav/security/mock/oauth2/grant/TokenExchangeGrant.kt
@@ -6,6 +6,7 @@ import no.nav.security.mock.oauth2.invalidRequest
 
 val TOKEN_EXCHANGE = GrantType("urn:ietf:params:oauth:grant-type:token-exchange")
 
+@Suppress("MemberVisibilityCanBePrivate")
 class TokenExchangeGrant(
     val subjectTokenType: String,
     val subjectToken: String,
diff --git a/src/main/kotlin/no/nav/security/mock/oauth2/grant/TokenExchangeGrantHandler.kt b/src/main/kotlin/no/nav/security/mock/oauth2/grant/TokenExchangeGrantHandler.kt
@@ -10,7 +10,7 @@ import no.nav.security.mock.oauth2.token.OAuth2TokenCallback
 import no.nav.security.mock.oauth2.token.OAuth2TokenProvider
 import okhttp3.HttpUrl
 
-class TokenExchangeGrantHandler(private val tokenProvider: OAuth2TokenProvider) : GrantHandler {
+internal class TokenExchangeGrantHandler(private val tokenProvider: OAuth2TokenProvider) : GrantHandler {
 
     override fun tokenResponse(
         request: OAuth2HttpRequest,
diff --git a/src/main/kotlin/no/nav/security/mock/oauth2/http/OAuth2HttpRequest.kt b/src/main/kotlin/no/nav/security/mock/oauth2/http/OAuth2HttpRequest.kt
@@ -5,15 +5,14 @@ import com.nimbusds.oauth2.sdk.TokenRequest
 import com.nimbusds.oauth2.sdk.auth.ClientAuthentication
 import com.nimbusds.oauth2.sdk.http.HTTPRequest
 import com.nimbusds.openid.connect.sdk.AuthenticationRequest
-import java.net.URLDecoder
-import java.nio.charset.StandardCharsets
 import no.nav.security.mock.oauth2.extensions.isAuthorizationEndpointUrl
 import no.nav.security.mock.oauth2.extensions.isDebuggerCallbackUrl
 import no.nav.security.mock.oauth2.extensions.isDebuggerUrl
 import no.nav.security.mock.oauth2.extensions.isEndSessionEndpointUrl
 import no.nav.security.mock.oauth2.extensions.isJwksUrl
 import no.nav.security.mock.oauth2.extensions.isTokenEndpointUrl
 import no.nav.security.mock.oauth2.extensions.isWellKnownUrl
+import no.nav.security.mock.oauth2.extensions.keyValuesToMap
 import no.nav.security.mock.oauth2.extensions.requirePrivateKeyJwt
 import no.nav.security.mock.oauth2.extensions.toAuthorizationEndpointUrl
 import no.nav.security.mock.oauth2.extensions.toEndSessionEndpointUrl
@@ -58,6 +57,7 @@ data class OAuth2HttpRequest(
         )
     }
 
+    @Suppress("MemberVisibilityCanBePrivate")
     fun asNimbusHTTPRequest(): HTTPRequest {
         return HTTPRequest(HTTPRequest.Method.valueOf(method), url.toUrl())
             .apply {
@@ -106,16 +106,6 @@ data class OAuth2HttpRequest(
     }
 }
 
-private fun String.urlDecode(): String = URLDecoder.decode(this, StandardCharsets.UTF_8)
-
-private fun String.keyValuesToMap(listDelimiter: String): Map<String, String> =
-    this.split(listDelimiter)
-        .filter { it.contains("=") }
-        .associate {
-            val (key, value) = it.split("=")
-            key.urlDecode().trim() to value.urlDecode().trim()
-        }
-
 enum class RequestType {
     WELL_KNOWN, AUTHORIZATION, TOKEN, END_SESSION, JWKS,
     DEBUGGER, DEBUGGER_CALLBACK, FAVICON, UNKNOWN
diff --git a/src/main/kotlin/no/nav/security/mock/oauth2/http/OAuth2HttpRequestHandler.kt b/src/main/kotlin/no/nav/security/mock/oauth2/http/OAuth2HttpRequestHandler.kt
@@ -6,6 +6,7 @@ import com.nimbusds.oauth2.sdk.GrantType
 import com.nimbusds.oauth2.sdk.GrantType.AUTHORIZATION_CODE
 import com.nimbusds.oauth2.sdk.GrantType.CLIENT_CREDENTIALS
 import com.nimbusds.oauth2.sdk.GrantType.JWT_BEARER
+import com.nimbusds.oauth2.sdk.GrantType.REFRESH_TOKEN
 import com.nimbusds.oauth2.sdk.OAuth2Error
 import com.nimbusds.oauth2.sdk.ParseException
 import com.nimbusds.openid.connect.sdk.AuthenticationRequest
@@ -22,6 +23,8 @@ import no.nav.security.mock.oauth2.grant.AuthorizationCodeHandler
 import no.nav.security.mock.oauth2.grant.ClientCredentialsGrantHandler
 import no.nav.security.mock.oauth2.grant.GrantHandler
 import no.nav.security.mock.oauth2.grant.JwtBearerGrantHandler
+import no.nav.security.mock.oauth2.grant.RefreshTokenGrantHandler
+import no.nav.security.mock.oauth2.grant.RefreshTokenManager
 import no.nav.security.mock.oauth2.grant.TOKEN_EXCHANGE
 import no.nav.security.mock.oauth2.grant.TokenExchangeGrantHandler
 import no.nav.security.mock.oauth2.http.RequestType.AUTHORIZATION
@@ -47,12 +50,14 @@ class OAuth2HttpRequestHandler(
     private val loginRequestHandler = LoginRequestHandler(templateMapper)
     private val debuggerRequestHandler = DebuggerRequestHandler(templateMapper)
     private val tokenCallbackQueue: BlockingQueue<OAuth2TokenCallback> = LinkedBlockingQueue()
+    private val refreshTokenManager = RefreshTokenManager()
 
     private val grantHandlers: Map<GrantType, GrantHandler> = mapOf(
-        AUTHORIZATION_CODE to AuthorizationCodeHandler(config.tokenProvider),
+        AUTHORIZATION_CODE to AuthorizationCodeHandler(config.tokenProvider, refreshTokenManager),
         CLIENT_CREDENTIALS to ClientCredentialsGrantHandler(config.tokenProvider),
         JWT_BEARER to JwtBearerGrantHandler(config.tokenProvider),
-        TOKEN_EXCHANGE to TokenExchangeGrantHandler(config.tokenProvider)
+        TOKEN_EXCHANGE to TokenExchangeGrantHandler(config.tokenProvider),
+        REFRESH_TOKEN to RefreshTokenGrantHandler(config.tokenProvider, refreshTokenManager)
     )
 
     fun handleRequest(request: OAuth2HttpRequest): OAuth2HttpResponse {
diff --git a/src/main/kotlin/no/nav/security/mock/oauth2/token/OAuth2TokenProvider.kt b/src/main/kotlin/no/nav/security/mock/oauth2/token/OAuth2TokenProvider.kt
@@ -9,15 +9,15 @@ import com.nimbusds.jose.jwk.RSAKey
 import com.nimbusds.jwt.JWTClaimsSet
 import com.nimbusds.jwt.SignedJWT
 import com.nimbusds.oauth2.sdk.TokenRequest
-import no.nav.security.mock.oauth2.extensions.clientIdAsString
-import okhttp3.HttpUrl
 import java.security.KeyPair
 import java.security.KeyPairGenerator
 import java.security.interfaces.RSAPrivateKey
 import java.security.interfaces.RSAPublicKey
 import java.time.Instant
 import java.util.Date
 import java.util.UUID
+import no.nav.security.mock.oauth2.extensions.clientIdAsString
+import okhttp3.HttpUrl
 
 class OAuth2TokenProvider {
     private val jwkSet: JWKSet = generateJWKSet(DEFAULT_KEYID)
@@ -30,8 +30,8 @@ class OAuth2TokenProvider {
     fun idToken(
         tokenRequest: TokenRequest,
         issuerUrl: HttpUrl,
-        nonce: String?,
-        oAuth2TokenCallback: OAuth2TokenCallback
+        oAuth2TokenCallback: OAuth2TokenCallback,
+        nonce: String? = null
     ) = createSignedJWT(
         defaultClaims(
             issuerUrl,
diff --git a/src/test/kotlin/no/nav/security/mock/oauth2/e2e/JwtBearerGrantIntegrationTest.kt b/src/test/kotlin/no/nav/security/mock/oauth2/e2e/JwtBearerGrantIntegrationTest.kt
@@ -2,17 +2,17 @@ package no.nav.security.mock.oauth2.e2e
 
 import com.nimbusds.oauth2.sdk.GrantType
 import io.kotest.matchers.collections.shouldContainExactly
-import io.kotest.matchers.ints.shouldBeGreaterThan
+import io.kotest.matchers.should
 import io.kotest.matchers.shouldBe
-import io.kotest.matchers.shouldNotBe
 import io.kotest.matchers.string.shouldContain
 import no.nav.security.mock.oauth2.testutils.ParsedTokenResponse
 import no.nav.security.mock.oauth2.testutils.audience
 import no.nav.security.mock.oauth2.testutils.claims
-import no.nav.security.mock.oauth2.testutils.issuer
+import no.nav.security.mock.oauth2.testutils.shouldBeValidFor
+import no.nav.security.mock.oauth2.testutils.subject
 import no.nav.security.mock.oauth2.testutils.toTokenResponse
 import no.nav.security.mock.oauth2.testutils.tokenRequest
-import no.nav.security.mock.oauth2.testutils.verify
+import no.nav.security.mock.oauth2.testutils.verifyWith
 import no.nav.security.mock.oauth2.token.DefaultOAuth2TokenCallback
 import no.nav.security.mock.oauth2.withMockOAuth2Server
 import okhttp3.OkHttpClient
@@ -28,21 +28,22 @@ class JwtBearerGrantIntegrationTest {
     @Test
     fun `token request with JwtBearerGrant should exchange assertion with a new token containing many of the same claims`() {
         withMockOAuth2Server {
+            val initialSubject = "yolo"
             val initialToken = this.issueToken(
                 issuerId = "idprovider",
                 clientId = "client1",
                 tokenCallback = DefaultOAuth2TokenCallback(
                     issuerId = "idprovider",
-                    subject = "mysub",
+                    subject = initialSubject,
                     claims = mapOf(
                         "claim1" to "value1",
                         "claim2" to "value2",
                     )
                 )
             )
-
+            val issuerId = "aad"
             val response: ParsedTokenResponse = client.tokenRequest(
-                url = this.tokenEndpointUrl("aad"),
+                url = this.tokenEndpointUrl(issuerId),
                 basicAuth = Pair("client1", "secret"),
                 parameters = mapOf(
                     "grant_type" to GrantType.JWT_BEARER.value,
@@ -51,21 +52,14 @@ class JwtBearerGrantIntegrationTest {
                 )
             ).toTokenResponse()
 
-            response.status shouldBe 200
-            response.expiresIn shouldBeGreaterThan 0
+            response shouldBeValidFor GrantType.JWT_BEARER
             response.scope shouldContain "scope1"
-            response.tokenType shouldBe "Bearer"
-            response.accessToken shouldNotBe null
-            response.idToken shouldBe null
-            response.refreshToken shouldBe null
             response.issuedTokenType shouldBe null
-
-            response.accessToken?.verify(this.issuerUrl("aad"), this.jwksUrl("aad"))
-
-            response.accessToken?.audience shouldContainExactly listOf("scope1")
-            response.accessToken?.issuer shouldBe this.issuerUrl("aad").toString()
-            response.accessToken?.claims?.get("claim1") shouldBe "value1"
-            response.accessToken?.claims?.get("claim2") shouldBe "value2"
+            response.accessToken!! should verifyWith(issuerId, this)
+            response.accessToken.subject shouldBe initialSubject
+            response.accessToken.audience shouldContainExactly listOf("scope1")
+            response.accessToken.claims["claim1"] shouldBe "value1"
+            response.accessToken.claims["claim2"] shouldBe "value2"
         }
     }
 }
diff --git a/src/test/kotlin/no/nav/security/mock/oauth2/e2e/RefreshTokenGrantIntegrationTest.kt b/src/test/kotlin/no/nav/security/mock/oauth2/e2e/RefreshTokenGrantIntegrationTest.kt
diff --git a/src/test/kotlin/no/nav/security/mock/oauth2/e2e/TokenExchangeGrantIntegrationTest.kt b/src/test/kotlin/no/nav/security/mock/oauth2/e2e/TokenExchangeGrantIntegrationTest.kt
diff --git a/src/test/kotlin/no/nav/security/mock/oauth2/grant/AuthorizationCodeHandlerTest.kt b/src/test/kotlin/no/nav/security/mock/oauth2/grant/AuthorizationCodeHandlerTest.kt
diff --git a/src/test/kotlin/no/nav/security/mock/oauth2/testutils/Http.kt b/src/test/kotlin/no/nav/security/mock/oauth2/testutils/Http.kt
diff --git a/src/test/kotlin/no/nav/security/mock/oauth2/testutils/Token.kt b/src/test/kotlin/no/nav/security/mock/oauth2/testutils/Token.kt
