mangadex_at_home/src/test/kotlin/mdnet/server/ImageServerTest.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.server

import io.kotest.assertions.withClue
import io.kotest.core.spec.IsolationMode
import io.kotest.core.spec.style.FreeSpec
import io.kotest.engine.spec.tempdir
import io.kotest.engine.spec.tempfile
import io.kotest.matchers.shouldBe
import io.micrometer.prometheus.PrometheusConfig
import io.micrometer.prometheus.PrometheusMeterRegistry
import io.mockk.confirmVerified
import io.mockk.every
import io.mockk.mockk
import io.mockk.verify
import mdnet.cache.ImageStorage
import org.apache.commons.io.IOUtils
import org.http4k.core.HttpHandler
import org.http4k.core.Method
import org.http4k.core.Request
import org.http4k.core.Response
import org.http4k.core.Status
import org.http4k.kotest.shouldHaveHeader
import org.http4k.kotest.shouldHaveStatus
import org.http4k.kotest.shouldNotHaveHeader
import org.http4k.routing.bind
import org.http4k.routing.routes
import org.ktorm.database.Database
import java.io.ByteArrayInputStream

class ImageServerTest : FreeSpec() {
    override fun isolationMode() = IsolationMode.InstancePerTest

    init {
        val registry = PrometheusMeterRegistry(PrometheusConfig.DEFAULT)
        val mockData = byteArrayOf(72, 66, 67, 66, 65, 66, 73, 69, 65, 67)

        "correct upstream responses" - {
            val client = mockk<HttpHandler>()
            every {
                client(any())
            } answers {
                correctMockResponse(mockData)
            }

            "mocked noop cache" - {
                val storage = mockk<ImageStorage>()
                every {
                    storage.loadImage(any())
                } returns null
                every {
                    storage.storeImage(any(), any())
                } returns null

                val server = ImageServer(
                    storage,
                    client,
                    registry
                )

                val handler = routes(
                    "/{token}/data/{chapterHash}/{fileName}" bind Method.GET to server.handler(
                        dataSaver = false,
                    ),
                    "/{token}/data-saver/{chapterHash}/{fileName}" bind Method.GET to server.handler(
                        dataSaver = true,
                    ),
                    "/data/{chapterHash}/{fileName}" bind Method.GET to server.handler(
                        dataSaver = false,
                    ),
                    "/data-saver/{chapterHash}/{fileName}" bind Method.GET to server.handler(
                        dataSaver = true,
                    ),
                )

                "directly proxied" {
                    for (i in 0..2) {
                        val response = handler(Request(Method.GET, "/data/02181a8f5fe8cd408720a771dd129fd8/T2.png"))

                        withClue("should be directly proxied") {
                            verify(exactly = i + 1) {
                                client(any())
                            }
                            confirmVerified(client)
                        }

                        response.shouldHaveStatus(Status.OK)
                        response.shouldHaveHeader("Content-Length", mockData.size.toString())
                        IOUtils.toByteArray(response.body.stream).shouldBe(mockData)
                        response.close()
                    }
                }

                "should not have Server header" {
                    val response = handler(Request(Method.GET, "/data/02181a8f5fe8cd408720a771dd129fd8/T2.png"))
                    response.shouldNotHaveHeader("Server")
                    response.close()
                }
            }

            "with real cache" - {
                val storage = ImageStorage(
                    maxSize = 100000,
                    cacheDirectory = tempdir().toPath(),
                    database = Database.connect("jdbc:sqlite:${tempfile()}"),
                    autoPrune = false,
                )

                val server = ImageServer(
                    storage,
                    client,
                    registry
                )

                val handler = routes(
                    "/data/{chapterHash}/{fileName}" bind Method.GET to server.handler(
                        dataSaver = false,
                    ),
                    "/data-saver/{chapterHash}/{fileName}" bind Method.GET to server.handler(
                        dataSaver = true,
                    ),
                )

                "respects cache" {
                    for (i in 0..2) {
                        val response = handler(Request(Method.GET, "/data/02181a8f5fe8cd408720a771dd129fd8/T2.png"))

                        withClue("should only be downloaded once") {
                            verify(exactly = 1) {
                                client(any())
                            }
                            confirmVerified(client)
                        }

                        response.shouldHaveStatus(Status.OK)
                        response.shouldHaveHeader("Content-Length", mockData.size.toString())
                        IOUtils.toByteArray(response.body.stream).shouldBe(mockData)
                        response.close()
                    }
                }
            }
        }

        "failed upstream responses" - {
            val client = mockk<HttpHandler>()

            val storage = ImageStorage(
                maxSize = 100000,
                cacheDirectory = tempdir().toPath(),
                database = Database.connect("jdbc:sqlite:${tempfile()}"),
                autoPrune = false,
            )

            val server = ImageServer(
                storage,
                client,
                registry
            )

            val handler = routes(
                "/data/{chapterHash}/{fileName}" bind Method.GET to server.handler(
                    dataSaver = false,
                ),
                "/data-saver/{chapterHash}/{fileName}" bind Method.GET to server.handler(
                    dataSaver = true,
                ),
            )

            "does not cache failures" {
                val errStatus = Status.NOT_FOUND
                every {
                    client(any())
                } returns Response(errStatus)

                for (i in 0..2) {
                    val response = handler(Request(Method.GET, "/data/02181a8f5fe8cd408720a771dd129fd8/T2.png"))

                    withClue("should be directly proxied") {
                        verify(exactly = i + 1) {
                            client(any())
                        }
                        confirmVerified(client)
                    }

                    response.shouldHaveStatus(errStatus)
                    response.close()
                }
            }

            "errors on bad content type" {
                every {
                    client(any())
                } answers {
                    correctMockResponse(mockData)
                        .replaceHeader("Content-Type", "text/html")
                }

                for (i in 0..2) {
                    val response = handler(Request(Method.GET, "/data/02181a8f5fe8cd408720a771dd129fd8/T2.png"))

                    withClue("should be directly proxied") {
                        verify(exactly = i + 1) {
                            client(any())
                        }
                        confirmVerified(client)
                    }

                    response.status.shouldBe(Status.INTERNAL_SERVER_ERROR)
                    response.close()
                }
            }

            "still works on no content-length" {
                every {
                    client(any())
                } answers {
                    correctMockResponse(mockData)
                        .removeHeader("Content-Length")
                }

                for (i in 0..2) {
                    val response = handler(Request(Method.GET, "/data/02181a8f5fe8cd408720a771dd129fd8/T2.png"))

                    response.shouldHaveStatus(Status.OK)
                    IOUtils.toByteArray(response.body.stream).shouldBe(mockData)
                    response.close()
                }
            }
        }
    }

    private fun correctMockResponse(data: ByteArray) =
        Response(Status.OK)
            .body(ByteArrayInputStream(data))
            .header("Content-Type", "image/jpg")
            .header("Content-Length", "${data.size}")
            .header("Last-Modified", "Wed, 21 Oct 2015 07:28:00 GMT")
}