GitLabDockerRegistryClient.kt

import okhttp3.Credentials
import okhttp3.OkHttpClient
import okhttp3.Request
import okhttp3.Response
import org.json.JSONObject
import java.time.Duration
import java.time.Instant
import java.util.logging.Logger

// Inspired by https://www.pimwiddershoven.nl/entry/request-an-api-bearer-token-from-gitlab-jwt-authentication-to-control-your-private-docker-registry

// GOTCHA! This is the Docker repository API version not the GitLab API version!
const val API_VERSION = 2

class DockerRegistryClient(private val hostUrl: String, private val registryPort: Int, username: String, password: String) {
    private val logger: Logger = Logger.getLogger(DockerRegistryClient::class.java.name)
    private val httpClient = OkHttpClient.Builder().build()
    private val credentials = Credentials.basic(username, password)

    fun getTags(repository: String): List<Tag> {
        return callApiForJson(repository, "/tags/list")
            .getJSONArray("tags")
            .map { Tag(it.toString()) }
    }

    fun getImageInfo(repository: String, tag: Tag): ImageInfo {
        val digest = createDigest(repository, tag)
        val createdTimestamp = getMetadata(repository, digest.configDigest).getString("created")
        return ImageInfo(tag, digest.manifestDigest, Instant.parse(createdTimestamp))
    }

    fun deleteImage(repository: String, digest: String) = callApi(repository, "/manifests/$digest", "DELETE")

    fun findObsoleteImagesPerBaseVersion(repository: String,
                                         tags: List<Tag>,
                                         expiryTime: Duration,
                                         numberToKeep: Int): Map<String, List<ImageInfo>> {
        val baseVersions = tags.groupBy(Tag::extractMajorMinorVersion)

        return baseVersions.mapValues { findObsoleteImagesForTags(repository, it.value, numberToKeep, expiryTime) }
    }

    private fun createDigest(repository: String, tag: Tag) : Digests {
        val response = callApi(repository, "/manifests/$tag")
        val manifestDigest = response.header("Docker-Content-Digest").orEmpty()
        val jsonObject = JSONObject(response.body()!!.string())
        return Digests(manifestDigest, jsonObject.getJSONObject("config").getString("digest"))
    }

    private fun getMetadata(repository: String, digest: String) = callApiForJson(repository, "/blobs/$digest")

    private fun findObsoleteImagesForTags(repository: String,
                                          tags: List<Tag>,
                                          numberToKeep: Int,
                                          expiryTime: Duration): List<ImageInfo> {
        val numberOfObsoleteImages = if (tags.size <= numberToKeep) 0 else tags.size - numberToKeep

        return if (numberOfObsoleteImages == 0) {
            emptyList()
        } else {
            tags.map { getImageInfo(repository, it) }.sorted().take(numberOfObsoleteImages)
                .filter { it.creationDate.isBefore(Instant.now().minus(expiryTime)) }
        }
    }

    private fun callApiForJson(repository:String, path:String, method: String = "GET") : JSONObject {
        val response = callApi(repository, path, method)
        if (response.isSuccessful)
            return JSONObject(response.body()!!.string())
        else
            throw RuntimeException(response.message())
    }

    private fun callApi(repository: String, path: String, method: String= "GET"): Response {
        val request = Request.Builder().url("$hostUrl:$registryPort/v$API_VERSION/$repository$path")
            .method(method, null)
            .header("Authorization", "Bearer ${createToken(repository)}")
            .header("Accept", "application/vnd.docker.distribution.manifest.v$API_VERSION+json")
            .build()

        logger.fine(request.toString())
        return httpClient.newCall(request).execute()
    }

    private fun createToken(repository: String) = JSONObject(
        httpClient.newCall(
            Request.Builder().url(
                "$hostUrl/jwt/auth?service=container_registry&scope=repository:$repository:*")
                .header("Authorization", credentials)
                .build()).execute().body()?.string()).getString("token")!!
}

/**
 * Digests are somewhat strange in Docker. We seem to need at least two:
 *   - one identifying the manifest itself
 *   - one suitable to get the configuration as a blob (for creation time and stuff)
 */
data class Digests(val manifestDigest: String, val configDigest: String)

data class ImageInfo(val tag: Tag, val digest: String, val creationDate: Instant) : Comparable<ImageInfo> {
    override fun compareTo(other: ImageInfo): Int = creationDate.compareTo(other.creationDate)
}

data class Tag(private val tagString: String) {
    private val regex = Regex("""[\d]+\.[\d]+""")

    fun extractMajorMinorVersion(): String {
        return regex.find(tagString)?.value ?: tagString
    }

    override fun toString(): String {
        return tagString
    }
}