AWS Lambda を用いた Amazon SNS のエンドポイント整理

2016-05-28 06:40

Amazon SNS の Application に紐付いているエンドポイントには Enabled という Attribute がついています。文字どおり、有効なエンドポイントかどうかを表すものです。Amazon SNS は賢いので APNS や GCM からのフィードバックを受けて、プッシュトークンが無効になった場合に自動的に false にしてくれます。しかし、基本的には Enabled が false になったエンドポイントは不要です。したがって、これらを自動的に削除するような仕組みがあったほうがよさそうです。そこで、この記事では AWS Lambda を用いて、それを実現する方法について説明します。

AWS Lambda で動かすスクリプト

コード自体は非常に簡単で以下のような形です。AWS Lambda Function を素早く立ち上げるとしたら本来は python か node.js のほうが良いかと思いますが、今回の作りたいものに関してはそれほど即時性を求められるようなものではないと思うので、好きな言語（もちろん Lambda でうごくもの）で問題ないと思います。

import com.amazonaws.auth.{AWSCredentials, BasicAWSCredentials}
import com.amazonaws.regions.{Region, Regions}
import com.amazonaws.services.lambda.runtime.events.SNSEvent
import com.amazonaws.services.lambda.runtime.{Context, RequestHandler}
import com.amazonaws.services.sns.AmazonSNSClient
import com.amazonaws.services.sns.model.{DeleteEndpointRequest, GetEndpointAttributesRequest}
import com.typesafe.config.ConfigFactory
import com.typesafe.scalalogging.slf4j.StrictLogging
import org.json4s._
import org.json4s.native.JsonMethods._

import scala.collection.JavaConversions._

object DeleteEndpointAction extends RequestHandler[SNSEvent, Unit] with StrictLogging {

  private type EndpointArn = String

  private val conf = ConfigFactory.load()
  private val credential: AWSCredentials = new BasicAWSCredentials(conf.getString("amazon_sns.access_key"), conf.getString("amazon_sns.access_secret"))
  private val snsClient: AmazonSNSClient = new AmazonSNSClient(credential).withRegion(Region.getRegion(Regions.fromName(conf.getString("amazon_sns.region"))))

  override def handleRequest(input: SNSEvent, context: Context) = {
    val records = input.getRecords.toList
    records.foreach { record =>
      val message = record.getSNS.getMessage
      if (message != null) {
        getEndpointArn(message).foreach { endpointArn =>
          logger.info(s"check if enabled: endpoint_arn = $endpointArn")
          val enabled = getEnabled(endpointArn).getOrElse(true)
          if (!enabled) {
            logger.info(s"delete disabled endpoint: endpoint_arn = $endpointArn")
            deleteEndpoint(endpointArn)
          }
        }
      } else {
        logger.error(s"receive empty message: message_id = ${record.getSNS.getMessageId}")
      }
    }
  }

  private def getEndpointArn(message: String): Option[EndpointArn] =
    for {
      JObject(child) <- parse(message)
      JField("EndpointArn", JString(resource)) <- child
    } yield resource.headOption

  private def getEnabled(endpointArn: EndpointArn): Option[Boolean] = {
    val request = new GetEndpointAttributesRequest().withEndpointArn(endpointArn)
    try {
      val result = snsClient.getEndpointAttributes(request)
      result.getAttributes.toMap.get("Enabled").map(_ == "true")
    } catch {
      case e: Throwable =>
        logger.error("failed to get endpoint attributes", e)
        None
    }
  }

  private def deleteEndpoint(endpointArn: EndpointArn) = {
    val request = new DeleteEndpointRequest().withEndpointArn(endpointArn)
    try {
      snsClient.deleteEndpoint(request)
    } catch {
      case e: Throwable => logger.error("failed to delete endpoint", e)
    }
  }

}