huawei-push-kit: use expires_in from responses

google-common/src/main/scala/org/apache/pekko/stream/connectors/google/auth/GoogleComputeMetadata.scala

@@ -23,6 +23,7 @@ import pekko.http.scaladsl.model.{ HttpRequest, Uri }
 import pekko.http.scaladsl.model.headers.RawHeader
 import pekko.http.scaladsl.unmarshalling.Unmarshal
 import pekko.stream.Materializer
+import pdi.jwt.JwtTime
 
 import java.time.Clock
 import scala.concurrent.Future
@@ -55,8 +56,16 @@ private[auth] object GoogleComputeMetadata {
     implicit val system: ActorSystem = mat.system
     for {
       response <- Http().singleRequest(tokenRequest(scopes))
-      token <- Unmarshal(response.entity).to[AccessToken]
-    } yield token
+      tokenResponse <- Unmarshal(response.entity).to[AccessTokenResponse]
+    } yield {
+      val expiresIn = if (tokenResponse.expires_in > 0) tokenResponse.expires_in
+      else {
+        system.log.warning("Google OAuth2 response contained invalid expires_in ({}), falling back to default {}s",
+          tokenResponse.expires_in, GoogleOAuth2.DefaultExpiresIn)
+        GoogleOAuth2.DefaultExpiresIn
+      }
+      AccessToken(tokenResponse.access_token, JwtTime.nowSeconds + expiresIn)
+    }
   }
 
   def getProjectId()(

google-common/src/main/scala/org/apache/pekko/stream/connectors/google/auth/GoogleOAuth2.scala

@@ -23,19 +23,20 @@ import pekko.stream.Materializer
 import pekko.stream.connectors.google.http.GoogleHttp
 import pekko.stream.connectors.google.jwt.JwtSprayJson
 import pekko.stream.connectors.google.{ implicits, RequestSettings }
+import pdi.jwt.{ JwtClaim, JwtTime }
 import pdi.jwt.JwtAlgorithm.RS256
-import pdi.jwt.JwtClaim
 import spray.json.DefaultJsonProtocol._
 import spray.json.JsonFormat
 
 import java.time.Clock
-import scala.concurrent.Future
+import scala.concurrent.{ ExecutionContext, Future }
 import scala.util.control.NonFatal
 
 @InternalApi
 private[auth] object GoogleOAuth2 {
 
   private val oAuthTokenUrl = "https://oauth2.googleapis.com/token"
+  private[auth] val DefaultExpiresIn = 3600
 
   def getAccessToken(clientEmail: String, privateKey: String, scopes: Set[String])(
       implicit mat: Materializer,
@@ -51,7 +52,16 @@ private[auth] object GoogleOAuth2 {
         "grant_type" -> "urn:ietf:params:oauth:grant-type:jwt-bearer",
         "assertion" -> generateJwt(clientEmail, privateKey, scopes)).toEntity
 
-      GoogleHttp().singleRequest[AccessToken](HttpRequest(POST, oAuthTokenUrl, entity = entity))
+      GoogleHttp().singleRequest[AccessTokenResponse](HttpRequest(POST, oAuthTokenUrl, entity = entity)).map {
+        case AccessTokenResponse(access_token, _, expires_in) =>
+          val safeExpiresIn = if (expires_in > 0) expires_in
+          else {
+            system.log.warning("Google OAuth2 response contained invalid expires_in ({}), falling back to default {}s",
+              expires_in, DefaultExpiresIn)
+            DefaultExpiresIn
+          }
+          AccessToken(access_token, JwtTime.nowSeconds + safeExpiresIn)
+      }(ExecutionContext.parasitic)
     } catch {
       case NonFatal(e) =>
         Future.failed(e)

google-common/src/main/scala/org/apache/pekko/stream/connectors/google/auth/UserAccessMetadata.scala

@@ -23,6 +23,7 @@ import pekko.http.scaladsl.model.headers.RawHeader
 import pekko.http.scaladsl.model.{ FormData, HttpRequest }
 import pekko.http.scaladsl.unmarshalling.Unmarshal
 import pekko.stream.Materializer
+import pdi.jwt.JwtTime
 
 import java.time.Clock
 import scala.concurrent.Future
@@ -49,7 +50,15 @@ private[auth] object UserAccessMetadata {
     implicit val system: ActorSystem = mat.system
     for {
       response <- Http().singleRequest(tokenRequest(clientId, clientSecret, refreshToken))
-      token <- Unmarshal(response.entity).to[AccessToken]
-    } yield token
+      tokenResponse <- Unmarshal(response.entity).to[AccessTokenResponse]
+    } yield {
+      val expiresIn = if (tokenResponse.expires_in > 0) tokenResponse.expires_in
+      else {
+        system.log.warning("Google OAuth2 response contained invalid expires_in ({}), falling back to default {}s",
+          tokenResponse.expires_in, GoogleOAuth2.DefaultExpiresIn)
+        GoogleOAuth2.DefaultExpiresIn
+      }
+      AccessToken(tokenResponse.access_token, JwtTime.nowSeconds + expiresIn)
+    }
   }
 }

huawei-push-kit/src/main/scala/org/apache/pekko/stream/connectors/huawei/pushkit/impl/HmsTokenApi.scala

@@ -36,14 +36,14 @@ import scala.concurrent.Future
 private[pushkit] class HmsTokenApi(http: => HttpExt, system: ActorSystem, forwardProxy: Option[ForwardProxy]) {
   import PushKitJsonSupport._
 
+  private val log = system.log
   private val authUrl = "https://oauth-login.cloud.huawei.com/oauth2/v3/token"
 
   def now: Long = JwtTime.nowSeconds(Clock.systemUTC())
 
   def getAccessToken(clientId: String, privateKey: String)(
       implicit materializer: Materializer): Future[AccessTokenExpiry] = {
     import materializer.executionContext
-    val expiresAt = now + 3600
 
     val requestEntity = FormData(
       "grant_type" -> "client_credentials",
@@ -60,9 +60,15 @@ private[pushkit] class HmsTokenApi(http: => HttpExt, system: ActorSystem, forwar
       }
       result <- Unmarshal(response.entity).to[OAuthResponse]
     } yield {
+      val expiresIn = if (result.expires_in > 0) result.expires_in
+      else {
+        log.warning("Huawei OAuth2 response contained invalid expires_in ({}), falling back to default {}s",
+          result.expires_in, HmsTokenApi.DefaultExpiresIn)
+        HmsTokenApi.DefaultExpiresIn
+      }
       AccessTokenExpiry(
         accessToken = result.access_token,
-        expiresAt = expiresAt)
+        expiresAt = now + expiresIn)
     }
   }
 }
@@ -72,6 +78,7 @@ private[pushkit] class HmsTokenApi(http: => HttpExt, system: ActorSystem, forwar
  */
 @InternalApi
 private[pushkit] object HmsTokenApi {
+  val DefaultExpiresIn = 3600
   case class AccessTokenExpiry(accessToken: String, expiresAt: Long) {
     override def toString: String =
       "AccessTokenExpiry(accessToken=*****," +

huawei-push-kit/src/test/scala/org/apache/pekko/stream/connectors/huawei/pushkit/impl/HmsTokenApiSpec.scala

@@ -103,6 +103,45 @@ class HmsTokenApiSpec
         case AccessTokenExpiry("token", exp) if exp > (System.currentTimeMillis / 1000L + 3000L) =>
       }
     }
+
+    "use expires_in from server response" in {
+      val http = mock[HttpExt]
+      when(
+        http.singleRequest(any[HttpRequest](),
+          any[HttpsConnectionContext](),
+          any[ConnectionPoolSettings](),
+          any[LoggingAdapter]())).thenReturn(
+        Future.successful(
+          HttpResponse(
+            entity = HttpEntity(ContentTypes.`application/json`,
+              """{"access_token": "token", "token_type": "String", "expires_in": 7200}"""))))
+
+      val api = new HmsTokenApi(http, system, Option.empty)
+      val result = api.getAccessToken(config.appId, config.appSecret).futureValue
+      result.accessToken shouldBe "token"
+      // expires_in=7200 means expiry should be ~7200s from now, well beyond the old hardcoded 3600s
+      result.expiresAt should be > (System.currentTimeMillis / 1000L + 6000L)
+    }
+
+    "fall back to default expiry when expires_in is invalid" in {
+      val http = mock[HttpExt]
+      when(
+        http.singleRequest(any[HttpRequest](),
+          any[HttpsConnectionContext](),
+          any[ConnectionPoolSettings](),
+          any[LoggingAdapter]())).thenReturn(
+        Future.successful(
+          HttpResponse(
+            entity = HttpEntity(ContentTypes.`application/json`,
+              """{"access_token": "token", "token_type": "String", "expires_in": 0}"""))))
+
+      val api = new HmsTokenApi(http, system, Option.empty)
+      val result = api.getAccessToken(config.appId, config.appSecret).futureValue
+      result.accessToken shouldBe "token"
+      // expires_in=0 should fall back to default 3600s
+      result.expiresAt should be > (System.currentTimeMillis / 1000L + 3000L)
+      result.expiresAt should be < (System.currentTimeMillis / 1000L + 4000L)
+    }
   }
 
 }