DDCNL-12161: TAI - Add in auth caching

app/uk/gov/hmrc/tai/config/ApplicationConfig.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 HM Revenue & Customs
+ * Copyright 2026 HM Revenue & Customs
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -104,6 +104,9 @@ class MongoConfig @Inject() (val runModeConfiguration: Configuration) extends Ba
   lazy val mongoLockTTL: Int = runModeConfiguration.getOptional[Int]("mongo.lock.expiryInMilliseconds").getOrElse(1200)
   lazy val mongoTTLUpdateIncome: Int =
     runModeConfiguration.getOptional[Int]("tai.cache.updateIncome.expiryInSeconds").getOrElse(3600 * 48)
+  lazy val mongoAuthTTL: Int = runModeConfiguration.getOptional[Int]("mongodb.auth.expiryInSeconds").getOrElse(30)
+  lazy val mongoAuthEnabled: Boolean =
+    runModeConfiguration.getOptional[Boolean]("mongodb.auth.enabled").getOrElse(false)
 }
 
 @Singleton

app/uk/gov/hmrc/tai/controllers/auth/AuthAction.scala

@@ -17,31 +17,88 @@
 package uk.gov.hmrc.tai.controllers.auth
 
 import com.google.inject.{ImplementedBy, Inject}
-import play.api.http.Status.UNAUTHORIZED
+import play.api.Logging
 import play.api.mvc.*
 import play.api.mvc.Results.*
 import uk.gov.hmrc.auth.core.*
 import uk.gov.hmrc.auth.core.retrieve.v2.Retrievals
 import uk.gov.hmrc.domain.Nino
-import uk.gov.hmrc.http.HeaderCarrier
+import uk.gov.hmrc.http.{GatewayTimeoutException, HeaderCarrier}
+import uk.gov.hmrc.mongo.cache.DataKey
 import uk.gov.hmrc.play.http.HeaderCarrierConverter
-import uk.gov.hmrc.tai.model.AuthenticatedRequest
+import uk.gov.hmrc.tai.config.MongoConfig
+import uk.gov.hmrc.tai.model.{AuthenticatedRequest, CachedAuthRetrievals}
+import uk.gov.hmrc.tai.repositories.cache.AuthCacheRepository
 
 import scala.concurrent.{ExecutionContext, Future}
 
-class AuthActionImpl @Inject() (override val authConnector: AuthConnector, cc: ControllerComponents)(implicit
+class AuthActionImpl @Inject() (
+  override val authConnector: AuthConnector,
+  appConfig: MongoConfig,
+  authCacheRepository: AuthCacheRepository,
+  cc: ControllerComponents
+)(implicit
   ec: ExecutionContext
-) extends AuthAction with AuthorisedFunctions {
+) extends AuthAction with AuthorisedFunctions with Logging {
+
+  private val AuthRetrievalsKey = DataKey[CachedAuthRetrievals]("auth-retrievals")
+
+  private def callAuth[A](request: Request[A], cacheResult: Boolean)(implicit
+    hc: HeaderCarrier
+  ): Future[Either[Result, AuthenticatedRequest[A]]] =
+    authorised(ConfidenceLevel.L200)
+      .retrieve(Retrievals.nino) {
+        case Some(nino) =>
+          val retrievals = CachedAuthRetrievals(nino)
+          val authenticatedRequest = AuthenticatedRequest(request, Nino(nino))
+
+          if (cacheResult) {
+            authCacheRepository.putSession(AuthRetrievalsKey, retrievals).map { _ =>
+              Right(authenticatedRequest)
+            }
+          } else {
+            Future.successful(Right(authenticatedRequest))
+          }
+        case None =>
+          logger.error("Unable to retrieve NINO from Auth")
+          Future.successful(Left(Unauthorized))
+      }
+      .recover {
+        case error: NoActiveSession =>
+          Left(Unauthorized(error.reason))
+
+        case error: InsufficientConfidenceLevel =>
+          Left(Unauthorized(error.reason))
+
+        case error: GatewayTimeoutException =>
+          logger.error(error.getMessage)
+          Left(BadGateway(error.getMessage))
+
+        case error =>
+          logger.error(error.getMessage, error)
+          Left(InternalServerError(error.getMessage))
+      }
 
   override protected def refine[A](request: Request[A]): Future[Either[Result, AuthenticatedRequest[A]]] = {
     implicit val hc: HeaderCarrier = HeaderCarrierConverter.fromRequest(request)
 
-    authorised(ConfidenceLevel.L200).retrieve(Retrievals.nino) {
-      case Some(nino) => Future.successful(Right(AuthenticatedRequest(request, Nino(nino))))
-      case None       => Future.successful(Left(Status(UNAUTHORIZED)))
-    } recover {
-      case _: NoActiveSession             => Left(Status(UNAUTHORIZED))
-      case _: InsufficientConfidenceLevel => Left(Status(UNAUTHORIZED))
+    if (appConfig.mongoAuthEnabled) {
+      authCacheRepository
+        .getFromSession[CachedAuthRetrievals](AuthRetrievalsKey)
+        .flatMap {
+          case Some(cacheRetrievals) =>
+            logger.debug("Auth retrieval cache HIT")
+            Future.successful(
+              Right(AuthenticatedRequest(request, Nino(cacheRetrievals.nino)))
+            )
+
+          case None =>
+            logger.debug("Auth retrieval cache MISS..!!")
+            callAuth(request, cacheResult = true)
+        }
+    } else {
+      logger.debug("Mongo caching for auth retrievals DISABLED..!!")
+      callAuth(request, cacheResult = false)
     }
   }
 

app/uk/gov/hmrc/tai/model/CachedAuthRetrievals.scala

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2026 HM Revenue & Customs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package uk.gov.hmrc.tai.model
+
+import play.api.libs.json.{Json, OFormat}
+
+case class CachedAuthRetrievals(
+  nino: String
+)
+
+object CachedAuthRetrievals {
+  implicit val format: OFormat[CachedAuthRetrievals] = Json.format[CachedAuthRetrievals]
+}

app/uk/gov/hmrc/tai/model/SensitiveWrapper.scala

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2026 HM Revenue & Customs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package uk.gov.hmrc.tai.model
+
+import play.api.libs.json.*
+import uk.gov.hmrc.crypto.json.JsonEncryption.{sensitiveDecrypter, sensitiveEncrypter}
+import uk.gov.hmrc.crypto.{Decrypter, Encrypter, Sensitive}
+
+case class SensitiveWrapper[T](override val decryptedValue: T) extends Sensitive[T]
+
+object SensitiveWrapper {
+
+  implicit def reads[T](implicit
+    reads: Reads[T],
+    crypto: Encrypter with Decrypter
+  ): Reads[SensitiveWrapper[T]] = sensitiveDecrypter(SensitiveWrapper[T])
+
+  implicit def writes[T](implicit
+    writes: Writes[T],
+    crypto: Encrypter with Decrypter
+  ): Writes[SensitiveWrapper[T]] = sensitiveEncrypter
+
+}

app/uk/gov/hmrc/tai/repositories/cache/AuthCacheRepository.scala

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2026 HM Revenue & Customs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package uk.gov.hmrc.tai.repositories.cache
+
+import play.api.libs.json.{Reads, Writes}
+import uk.gov.hmrc.tai.config.{CryptoProvider, MongoConfig}
+import uk.gov.hmrc.http.HeaderCarrier
+import uk.gov.hmrc.mongo.cache.DataKey
+import uk.gov.hmrc.mongo.{CurrentTimestampSupport, MongoComponent}
+
+import java.util.concurrent.TimeUnit
+import javax.inject.{Inject, Singleton}
+import scala.concurrent.ExecutionContext
+import scala.concurrent.duration.Duration
+
+@Singleton
+class AuthCacheRepository @Inject() (
+  mongoConfig: MongoConfig,
+  cryptoProvider: CryptoProvider,
+  mongoComponent: MongoComponent
+)(implicit ec: ExecutionContext)
+    extends GenericCacheRepository[HeaderCarrier](
+      mongoComponent = mongoComponent,
+      crypto = cryptoProvider.get,
+      collectionName = "taiAuthCache",
+      ttl = Duration(mongoConfig.mongoAuthTTL, TimeUnit.SECONDS),
+      timestampSupport = new CurrentTimestampSupport(),
+      cacheIdType = RequestCacheId
+    ) {
+
+  def putSession[T: Writes](dataKey: DataKey[T], data: T)(implicit
+    hc: HeaderCarrier
+  ): scala.concurrent.Future[(String, String)] =
+    put(hc, dataKey, data)
+
+  def getFromSession[T: Reads](dataKey: DataKey[T])(implicit hc: HeaderCarrier): scala.concurrent.Future[Option[T]] =
+    get(hc, dataKey)
+
+  def deleteFromSession[T](dataKey: DataKey[T])(implicit hc: HeaderCarrier): scala.concurrent.Future[Unit] =
+    delete(hc, dataKey)
+
+  def deleteAllFromSession(implicit hc: HeaderCarrier): scala.concurrent.Future[Unit] =
+    deleteAll(hc)
+}

app/uk/gov/hmrc/tai/repositories/cache/GenericCacheRepository.scala

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2026 HM Revenue & Customs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package uk.gov.hmrc.tai.repositories.cache
+
+import org.mongodb.scala.model.IndexModel
+import play.api.Logging
+import play.api.libs.json.{Reads, Writes}
+import uk.gov.hmrc.crypto.{Decrypter, Encrypter}
+import uk.gov.hmrc.tai.model.SensitiveWrapper
+import uk.gov.hmrc.mdc.Mdc
+import uk.gov.hmrc.mongo.cache.{CacheIdType, DataKey, MongoCacheRepository}
+import uk.gov.hmrc.mongo.{MongoComponent, MongoDatabaseCollection, TimestampSupport}
+
+import javax.inject.{Inject, Singleton}
+import scala.concurrent.duration.Duration
+import scala.concurrent.{ExecutionContext, Future}
+import scala.util.control.NonFatal
+
+@Singleton
+abstract class GenericCacheRepository[CacheId] @Inject() (
+  mongoComponent: MongoComponent,
+  override val collectionName: String,
+  val crypto: Encrypter with Decrypter,
+  replaceIndexes: Boolean = true,
+  ttl: Duration,
+  timestampSupport: TimestampSupport,
+  cacheIdType: CacheIdType[CacheId]
+)(implicit ec: ExecutionContext)
+    extends MongoDatabaseCollection with Logging {
+
+  private val cacheRepo: MongoCacheRepository[CacheId] = new MongoCacheRepository[CacheId](
+    mongoComponent = mongoComponent,
+    collectionName = collectionName,
+    replaceIndexes = replaceIndexes,
+    ttl = ttl,
+    timestampSupport = timestampSupport,
+    cacheIdType = cacheIdType
+  )
+
+  override val indexes: Seq[IndexModel] = cacheRepo.indexes
+
+  given (Encrypter with Decrypter) = crypto
+
+  def put[T: Writes](cacheId: CacheId, dataKey: DataKey[T], data: T): Future[(String, String)] =
+    Mdc.preservingMdc {
+      cacheRepo
+        .put[SensitiveWrapper[T]](cacheId)(DataKey[SensitiveWrapper[T]](dataKey.unwrap), SensitiveWrapper(data))
+        .map(res => "id" -> res.id)
+    }
+
+  def get[T: Reads](cacheId: CacheId, dataKey: DataKey[T]): Future[Option[T]] =
+    Mdc.preservingMdc {
+      cacheRepo
+        .get[SensitiveWrapper[T]](cacheId)(DataKey[SensitiveWrapper[T]](dataKey.unwrap))
+        .map(_.map(_.decryptedValue)) recoverWith { case NonFatal(error) =>
+        logger.error(s"Failed to read data from cache", error)
+        Future.successful(None)
+      }
+    }
+
+  def delete[T](cacheId: CacheId, dataKey: DataKey[T]): Future[Unit] =
+    Mdc.preservingMdc {
+      cacheRepo.delete(cacheId)(DataKey[SensitiveWrapper[T]](dataKey.unwrap))
+    }
+
+  def deleteAll(cacheId: CacheId): Future[Unit] =
+    Mdc.preservingMdc {
+      cacheRepo.deleteEntity(cacheId)
+    }
+}

app/uk/gov/hmrc/tai/repositories/cache/RequestCacheId.scala

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2026 HM Revenue & Customs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package uk.gov.hmrc.tai.repositories.cache
+
+import play.api.Logger
+import uk.gov.hmrc.http.HeaderCarrier
+import uk.gov.hmrc.mongo.cache.CacheIdType
+
+import java.util.UUID.randomUUID
+
+case object RequestCacheId extends CacheIdType[HeaderCarrier] {
+  override def run: HeaderCarrier => String = {
+    val logger: Logger = Logger(this.getClass)
+
+    _.requestId
+      .map(_.value)
+      .getOrElse {
+        logger.error(NoRequestIdException.getMessage, NoRequestIdException)
+        randomUUID.toString
+      }
+  }
+}
+case object NoRequestIdException extends Exception("Could not find requestId")

conf/application.conf

@@ -74,6 +74,10 @@ mongo.lock.expiryInMilliseconds = 1200
 
 mongodb {
   uri = "mongodb://localhost:27017/tai"
+  auth {
+    enabled = true,
+    expiryInSeconds = 30
+  }
 }