package social.grain

import org.junit.Assert.assertArrayEquals
import org.junit.Assert.assertEquals
import org.junit.Test
import social.grain.data.api.derToJose
import java.math.BigInteger
import java.security.KeyPairGenerator
import java.security.Signature
import java.security.spec.ECGenParameterSpec

/**
 * ES256 signatures are the one place a silent encoding bug would cost a
 * developer a day: JCA emits ASN.1 DER, JWS wants raw r||s, and a wrong answer
 * shows up only as an opaque 401 from the PDS.
 */
class DerToJoseTest {

    @Test
    fun `converts a real ECDSA signature to 64 raw bytes`() {
        val keyPair = KeyPairGenerator.getInstance("EC").apply {
            initialize(ECGenParameterSpec("secp256r1"))
        }.generateKeyPair()

        // Many signatures, because whether r or s needs a leading zero byte (or
        // a padding byte stripped) depends on the random nonce.
        repeat(200) { index ->
            val der = Signature.getInstance("SHA256withECDSA").run {
                initSign(keyPair.private)
                update("payload-$index".toByteArray())
                sign()
            }
            val jose = derToJose(der)
            assertEquals("Signature $index should be 64 bytes", 64, jose.size)
        }
    }

    @Test
    fun `preserves r and s values across the conversion`() {
        val keyPair = KeyPairGenerator.getInstance("EC").apply {
            initialize(ECGenParameterSpec("secp256r1"))
        }.generateKeyPair()

        repeat(50) { index ->
            val der = Signature.getInstance("SHA256withECDSA").run {
                initSign(keyPair.private)
                update("check-$index".toByteArray())
                sign()
            }
            val jose = derToJose(der)

            val (expectedR, expectedS) = parseDerIntegers(der)
            assertEquals(expectedR, BigInteger(1, jose.copyOfRange(0, 32)))
            assertEquals(expectedS, BigInteger(1, jose.copyOfRange(32, 64)))
        }
    }

    @Test
    fun `left-pads short components rather than shifting them`() {
        // A hand-built DER whose r is only one byte long: the raw form must be
        // 31 zero bytes then 0x07, not 0x07 followed by garbage.
        val der = byteArrayOf(
            0x30, 0x08,
            0x02, 0x01, 0x07,
            0x02, 0x03, 0x01, 0x02, 0x03,
        )
        val jose = derToJose(der)
        assertEquals(64, jose.size)
        assertArrayEquals(ByteArray(31), jose.copyOfRange(0, 31))
        assertEquals(0x07.toByte(), jose[31])
        assertEquals(BigInteger("010203", 16), BigInteger(1, jose.copyOfRange(32, 64)))
    }

    /** Minimal DER reader, independent of the implementation under test. */
    private fun parseDerIntegers(der: ByteArray): Pair<BigInteger, BigInteger> {
        var offset = 2
        if (der[1].toInt() and 0xFF > 0x80) offset += der[1].toInt() and 0x7F
        val rLen = der[offset + 1].toInt()
        val r = BigInteger(der.copyOfRange(offset + 2, offset + 2 + rLen))
        val sStart = offset + 2 + rLen
        val sLen = der[sStart + 1].toInt()
        val s = BigInteger(der.copyOfRange(sStart + 2, sStart + 2 + sLen))
        return r to s
    }
}