mangadex_at_home/src/main/kotlin/mdnet/metrics/GeoIpMetricsFilter.kt

/*
Mangadex@Home
Copyright (c) 2020, MangaDex Network
This file is part of MangaDex@Home.

MangaDex@Home is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

MangaDex@Home is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this MangaDex@Home.  If not, see <http://www.gnu.org/licenses/>.
*/
package mdnet.metrics

import com.maxmind.db.CHMCache
import com.maxmind.geoip2.DatabaseReader
import com.maxmind.geoip2.exception.GeoIp2Exception
import io.micrometer.prometheus.PrometheusMeterRegistry
import mdnet.logging.debug
import mdnet.logging.warn
import org.apache.commons.compress.archivers.tar.TarArchiveInputStream
import org.apache.commons.io.IOUtils
import org.http4k.core.Filter
import org.http4k.core.HttpHandler
import org.http4k.core.Method
import org.http4k.core.Request
import org.http4k.core.Status
import org.http4k.filter.gunzippedStream
import org.slf4j.Logger
import org.slf4j.LoggerFactory
import java.net.InetAddress
import java.net.UnknownHostException
import java.nio.file.Files

class GeoIpMetricsFilter(
    private val databaseReader: DatabaseReader?,
    private val registry: PrometheusMeterRegistry
) : Filter {
    override fun invoke(next: HttpHandler): HttpHandler {
        return {
            if (databaseReader != null && (it.uri.path != "/prometheus")) {
                inspectAndRecordSourceCountry(it)
            }

            next(it)
        }
    }

    private fun inspectAndRecordSourceCountry(request: Request) {
        val sourceIp =
            request.headerValues("Forwarded").firstOrNull() // try Forwarded (rare but standard)
                ?: request.headerValues("X-Forwarded-For").firstOrNull() // X-Forwarded-For (common but technically wrong)
                ?: request.source?.address // source (in case of no proxying, or with proxy-protocol)

        sourceIp.apply {
            try {
                val inetAddress = InetAddress.getByName(sourceIp)
                if (!inetAddress.isLoopbackAddress && !inetAddress.isAnyLocalAddress) {
                    val country = databaseReader!!.country(inetAddress)
                    recordCountry(country.country.isoCode)
                }
            } catch (e: GeoIp2Exception) {
                // do not disclose ip here, for privacy of logs
                LOGGER.warn { "Cannot resolve the country of the request's IP!" }
            } catch (e: UnknownHostException) {
                LOGGER.warn { "Cannot resolve source IP of the request!" }
            }
        }
    }

    private fun recordCountry(code: String) {
        registry.counter(
            "requests_country_counts",
            "country", code
        ).increment()
    }

    companion object {
        private val LOGGER: Logger = LoggerFactory.getLogger(GeoIpMetricsFilter::class.java)
    }
}

class GeoIpMetricsFilterBuilder(
    private val enableGeoIp: Boolean,
    private val license: String,
    private val registry: PrometheusMeterRegistry,
    private val client: HttpHandler
) {
    fun build(): GeoIpMetricsFilter {
        return if (enableGeoIp) {
            LOGGER.info("GeoIp initialising")
            val databaseReader = initDatabase()
            LOGGER.info("GeoIp initialised")
            GeoIpMetricsFilter(databaseReader, registry)
        } else {
            GeoIpMetricsFilter(null, registry)
        }
    }

    private fun initDatabase(): DatabaseReader {
        val databaseFileDir = Files.createTempDirectory("mangadex-geoip")
        val databaseFile = Files.createTempFile(databaseFileDir, "geoip2_country", ".mmdb")

        val geoIpDatabaseUri = GEOIP2_COUNTRY_URI_FORMAT.format(license)
        val response = client(Request(Method.GET, geoIpDatabaseUri))
        if (response.status != Status.OK) {
            throw IllegalStateException("Couldn't download GeoIP 2 database (http status: ${response.status})")
        }

        response.use {
            val archiveStream = TarArchiveInputStream(it.body.gunzippedStream().stream)
            var entry = archiveStream.nextTarEntry
            while (!entry.name.endsWith(".mmdb")) {
                LOGGER.debug { "Skipped non-database file: ${entry.name}" }
                entry = archiveStream.nextTarEntry
            }

            // reads only the current entry to its end
            val dbBytes = IOUtils.toByteArray(archiveStream)
            Files.write(databaseFile, dbBytes)
        }

        return DatabaseReader
            .Builder(databaseFile.toFile())
            .withCache(CHMCache())
            .build()
    }

    companion object {
        private val LOGGER = LoggerFactory.getLogger(GeoIpMetricsFilterBuilder::class.java)
        private const val GEOIP2_COUNTRY_URI_FORMAT: String =
            "https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-Country&license_key=%s&suffix=tar.gz"
    }
}