From ac362f9f6cdd41edc38594d17749833fc205425d Mon Sep 17 00:00:00 2001
From: Christopher Watford <christopher.watford@kcftech.com>
Date: Fri, 11 Aug 2023 11:23:51 -0400
Subject: [PATCH] Initial work on OAuth2 Token Introspection #269

---
 build.sbt                                     |  1 +
 .../sharing/server/DeltaSharingService.scala  | 18 ++++++++++
 .../sharing/server/config/ServerConfig.scala  | 36 +++++++++++++++++++
 3 files changed, 55 insertions(+)

diff --git a/build.sbt b/build.sbt
index 31ecd7364..35bd3b3f5 100644
--- a/build.sbt
+++ b/build.sbt
@@ -120,6 +120,7 @@ lazy val server = (project in file("server")) enablePlugins(JavaAppPackaging) se
       ExclusionRule("com.fasterxml.jackson.module"),
       ExclusionRule("org.json4s")
     ),
+    "com.linecorp.armeria" % "armeria-oauth2" % "1.24.3",
     "com.thesamet.scalapb" %% "scalapb-runtime" % scalapb.compiler.Version.scalapbVersion % "protobuf" excludeAll(
       ExclusionRule("com.fasterxml.jackson.core"),
       ExclusionRule("com.fasterxml.jackson.module"),
diff --git a/server/src/main/scala/io/delta/sharing/server/DeltaSharingService.scala b/server/src/main/scala/io/delta/sharing/server/DeltaSharingService.scala
index 0647e05f1..cada84293 100644
--- a/server/src/main/scala/io/delta/sharing/server/DeltaSharingService.scala
+++ b/server/src/main/scala/io/delta/sharing/server/DeltaSharingService.scala
@@ -26,12 +26,14 @@ import javax.annotation.Nullable
 import scala.collection.JavaConverters._
 import scala.util.Try
 
+import com.linecorp.armeria.client.WebClient
 import com.linecorp.armeria.common.{HttpData, HttpHeaderNames, HttpHeaders, HttpMethod, HttpRequest, HttpResponse, HttpStatus, MediaType, ResponseHeaders, ResponseHeadersBuilder}
 import com.linecorp.armeria.common.auth.OAuth2Token
 import com.linecorp.armeria.internal.server.ResponseConversionUtil
 import com.linecorp.armeria.server.{Server, ServiceRequestContext}
 import com.linecorp.armeria.server.annotation.{ConsumesJson, Default, ExceptionHandler, ExceptionHandlerFunction, Get, Head, Param, Post, ProducesJson}
 import com.linecorp.armeria.server.auth.AuthService
+import com.linecorp.armeria.server.auth.oauth2.OAuth2TokenIntrospectionAuthorizer
 import io.delta.standalone.internal.DeltaCDFErrors
 import io.delta.standalone.internal.DeltaCDFIllegalArgumentException
 import io.delta.standalone.internal.DeltaDataSource
@@ -532,6 +534,22 @@ object DeltaSharingService {
           })
         builder.decorator(authServiceBuilder.newDecorator)
       }
+      if (serverConfig.getTokenAuthorization != null) {
+        val tokenAuth = serverConfig.getTokenAuthorization
+        val introspectClient: WebClient = WebClient.of(tokenAuth.tokenInstrospectionUri)
+        val authServiceBuilder =
+          AuthService.builder.addOAuth2(
+            OAuth2TokenIntrospectionAuthorizer.builder(
+              introspectClient,
+              tokenAuth.tokenIntrospectionEndpoint
+            )
+            .clientCredentials(() => java.util.Map.entry(
+              tokenAuth.clientId, tokenAuth.clientSecret
+            ))
+            .build()
+          )
+        builder.decorator(authServiceBuilder.newDecorator)
+      }
       builder.build()
     }
     server.start().get()
diff --git a/server/src/main/scala/io/delta/sharing/server/config/ServerConfig.scala b/server/src/main/scala/io/delta/sharing/server/config/ServerConfig.scala
index bbee3a22b..1863d57d7 100644
--- a/server/src/main/scala/io/delta/sharing/server/config/ServerConfig.scala
+++ b/server/src/main/scala/io/delta/sharing/server/config/ServerConfig.scala
@@ -41,6 +41,7 @@ case class ServerConfig(
     @BeanProperty var version: java.lang.Integer,
     @BeanProperty var shares: java.util.List[ShareConfig],
     @BeanProperty var authorization: Authorization,
+    @BeanProperty var tokenAuthorization: TokenAuthorization,
     @BeanProperty var ssl: SSLConfig,
     @BeanProperty var host: String,
     @BeanProperty var port: Int,
@@ -71,6 +72,7 @@ case class ServerConfig(
       version = null,
       shares = Collections.emptyList(),
       authorization = null,
+      tokenAuthorization = null,
       ssl = null,
       host = "localhost",
       port = 80,
@@ -110,6 +112,9 @@ case class ServerConfig(
     if (authorization != null) {
       authorization.checkConfig()
     }
+    if (tokenAuthorization != null) {
+      tokenAuthorization.checkConfig()
+    }
     if (ssl != null) {
       ssl.checkConfig()
     }
@@ -167,6 +172,37 @@ case class Authorization(@BeanProperty var bearerToken: String) extends ConfigIt
   }
 }
 
+case class TokenAuthorization(
+  @BeanProperty var tokenInstrospectionUri: String,
+  @BeanProperty var tokenIntrospectionEndpoint: String,
+  @BeanProperty var clientId: String,
+  @BeanProperty var clientSecret: String)
+  extends ConfigItem {
+
+  def this() {
+    this(null, null, null, null)
+  }
+
+  override def checkConfig(): Unit = {
+    if (tokenInstrospectionUri == null) {
+      throw new IllegalArgumentException(
+        "'tokenIntrospectionUri' in 'tokenAuthorization' must be provided"
+      )
+    }
+    if (tokenIntrospectionEndpoint == null) {
+      throw new IllegalArgumentException(
+        "'tokenIntrospectionEndpoint' in 'tokenAuthorization' must be provided"
+      )
+    }
+    if (clientId == null) {
+      throw new IllegalArgumentException("'clientId' in 'tokenAuthorization' must be provided")
+    }
+    if (clientSecret == null) {
+      throw new IllegalArgumentException("'clientSecret' in 'tokenAuthorization' must be provided")
+    }
+  }
+}
+
 case class SSLConfig(
     @BeanProperty var selfSigned: Boolean,
     // The file of the PEM-format certificate