Skip to content

Commit

Permalink
Message load performance improvement (#133)
Browse files Browse the repository at this point in the history
* add rust backed libraries

* update to use the rust library

* remove the performance test

* move the signature test into the rust testable classes
  • Loading branch information
nplasterer authored Nov 3, 2023
1 parent 5454bbb commit 3ac6990
Show file tree
Hide file tree
Showing 7 changed files with 108 additions and 49 deletions.
Original file line number Diff line number Diff line change
@@ -1,12 +1,15 @@
package org.xmtp.android.library

import androidx.test.ext.junit.runners.AndroidJUnit4
import com.google.protobuf.kotlin.toByteStringUtf8
import kotlinx.coroutines.runBlocking
import org.junit.Test
import org.junit.runner.RunWith
import org.web3j.crypto.Hash
import org.xmtp.android.library.messages.PrivateKeyBuilder
import org.xmtp.android.library.messages.verify

@RunWith(AndroidJUnit4::class)
class SignatureTest {
@Test
fun testVerify() {
Expand Down
Original file line number Diff line number Diff line change
@@ -1,19 +1,9 @@
package org.xmtp.android.library.messages

import com.google.protobuf.kotlin.toByteString
import org.bouncycastle.jce.ECNamedCurveTable
import org.bouncycastle.jce.ECPointUtil
import org.bouncycastle.jce.provider.BouncyCastleProvider
import org.bouncycastle.jce.spec.ECNamedCurveParameterSpec
import org.bouncycastle.jce.spec.ECNamedCurveSpec
import org.bouncycastle.util.Arrays
import org.xmtp.android.library.Util
import org.xmtp.android.library.toHex
import org.xmtp.proto.message.contents.SignatureOuterClass
import java.math.BigInteger
import java.security.KeyFactory
import java.security.interfaces.ECPublicKey
import java.security.spec.ECPublicKeySpec

typealias Signature = org.xmtp.proto.message.contents.SignatureOuterClass.Signature

Expand Down Expand Up @@ -61,44 +51,18 @@ val Signature.rawDataWithNormalizedRecovery: ByteArray
return data
}

@OptIn(ExperimentalUnsignedTypes::class)
fun Signature.verify(signedBy: PublicKey, digest: ByteArray): Boolean {
val ecdsaVerify = java.security.Signature.getInstance("SHA256withECDSA", BouncyCastleProvider())
ecdsaVerify.initVerify(getPublicKeyFromBytes(signedBy.secp256K1Uncompressed.bytes.toByteArray()))
ecdsaVerify.update(digest)
return ecdsaVerify.verify(normalizeSignatureForVerification(this.rawDataWithNormalizedRecovery))
}

private fun normalizeSignatureForVerification(signature: ByteArray): ByteArray {
val r: ByteArray = BigInteger(1, Arrays.copyOfRange(signature, 0, 32)).toByteArray()
val s: ByteArray = BigInteger(1, Arrays.copyOfRange(signature, 32, 64)).toByteArray()
val der = ByteArray(6 + r.size + s.size)
der[0] = 0x30 // Tag of signature object

der[1] = (der.size - 2).toByte() // Length of signature object

var o = 2
der[o++] = 0x02 // Tag of ASN1 Integer

der[o++] = r.size.toByte() // Length of first signature part

System.arraycopy(r, 0, der, o, r.size)
o += r.size
der[o++] = 0x02 // Tag of ASN1 Integer

der[o++] = s.size.toByte() // Length of second signature part

System.arraycopy(s, 0, der, o, s.size)

return der
}

private fun getPublicKeyFromBytes(pubKey: ByteArray): java.security.PublicKey {
val spec: ECNamedCurveParameterSpec = ECNamedCurveTable.getParameterSpec("secp256k1")
val kf: KeyFactory = KeyFactory.getInstance("ECDSA", BouncyCastleProvider())
val params = ECNamedCurveSpec("secp256k1", spec.curve, spec.g, spec.n)
val point = ECPointUtil.decodePoint(params.curve, pubKey)
val pubKeySpec = ECPublicKeySpec(point, params)
return kf.generatePublic(pubKeySpec) as ECPublicKey
return try {
uniffi.xmtp_dh.verifyK256Sha256(
signedBy.secp256K1Uncompressed.bytes.toByteArray().toUByteArray().toList(),
digest.toUByteArray().toList(),
ecdsaCompact.bytes.toByteArray().toUByteArray().toList(),
ecdsaCompact.recovery.toUByte()
)
} catch (e: Exception) {
false
}
}

fun Signature.ensureWalletSignature(): Signature {
Expand Down
96 changes: 94 additions & 2 deletions library/src/main/java/xmtp_dh.kt
Original file line number Diff line number Diff line change
Expand Up @@ -35,10 +35,8 @@ import java.util.concurrent.ConcurrentHashMap
open class RustBuffer : Structure() {
@JvmField
var capacity: Int = 0

@JvmField
var len: Int = 0

@JvmField
var data: Pointer? = null

Expand Down Expand Up @@ -407,6 +405,14 @@ internal interface _UniFFILib : Library {
`privateKeyBytes`: RustBuffer.ByValue, _uniffi_out_err: RustCallStatus,
): RustBuffer.ByValue

fun uniffi_xmtp_dh_fn_func_verify_k256_sha256(
`signedBy`: RustBuffer.ByValue,
`message`: RustBuffer.ByValue,
`signature`: RustBuffer.ByValue,
`recoveryId`: Byte,
_uniffi_out_err: RustCallStatus,
): Byte

fun ffi_xmtp_dh_rustbuffer_alloc(
`size`: Int, _uniffi_out_err: RustCallStatus,
): RustBuffer.ByValue
Expand Down Expand Up @@ -435,6 +441,9 @@ internal interface _UniFFILib : Library {
fun uniffi_xmtp_dh_checksum_func_generate_private_preferences_topic_identifier(
): Short

fun uniffi_xmtp_dh_checksum_func_verify_k256_sha256(
): Short

fun ffi_xmtp_dh_uniffi_contract_version(
): Int

Expand Down Expand Up @@ -464,6 +473,9 @@ private fun uniffiCheckApiChecksums(lib: _UniFFILib) {
if (lib.uniffi_xmtp_dh_checksum_func_generate_private_preferences_topic_identifier() != 65141.toShort()) {
throw RuntimeException("UniFFI API checksum mismatch: try cleaning and rebuilding your project")
}
if (lib.uniffi_xmtp_dh_checksum_func_verify_k256_sha256() != 45969.toShort()) {
throw RuntimeException("UniFFI API checksum mismatch: try cleaning and rebuilding your project")
}
}

// Public interface members begin here.
Expand All @@ -489,6 +501,26 @@ public object FfiConverterUByte : FfiConverter<UByte, Byte> {
}
}

public object FfiConverterBoolean : FfiConverter<Boolean, Byte> {
override fun lift(value: Byte): Boolean {
return value.toInt() != 0
}

override fun read(buf: ByteBuffer): Boolean {
return lift(buf.get())
}

override fun lower(value: Boolean): Byte {
return if (value) 1.toByte() else 0.toByte()
}

override fun allocationSize(value: Boolean) = 1

override fun write(value: Boolean, buf: ByteBuffer) {
buf.put(lower(value))
}
}

public object FfiConverterString : FfiConverter<String, RustBuffer.ByValue> {
// Note: we don't inherit from FfiConverterRustBuffer, because we use a
// special encoding when lowering/lifting. We can use `RustBuffer.len` to
Expand Down Expand Up @@ -612,6 +644,44 @@ public object FfiConverterTypeEciesError : FfiConverterRustBuffer<EciesException
}


sealed class VerifyException(message: String) : Exception(message) {
// Each variant is a nested class
// Flat enums carries a string error message, so no special implementation is necessary.
class GenericException(message: String) : VerifyException(message)


companion object ErrorHandler : CallStatusErrorHandler<VerifyException> {
override fun lift(error_buf: RustBuffer.ByValue): VerifyException =
FfiConverterTypeVerifyError.lift(error_buf)
}
}

public object FfiConverterTypeVerifyError : FfiConverterRustBuffer<VerifyException> {
override fun read(buf: ByteBuffer): VerifyException {

return when (buf.getInt()) {
1 -> VerifyException.GenericException(FfiConverterString.read(buf))
else -> throw RuntimeException("invalid error enum value, something is very wrong!!")
}

}

override fun allocationSize(value: VerifyException): Int {
return 4
}

override fun write(value: VerifyException, buf: ByteBuffer) {
when (value) {
is VerifyException.GenericException -> {
buf.putInt(1)
Unit
}
}.let { /* this makes the `when` an expression, which ensures it is exhaustive */ }
}

}


public object FfiConverterSequenceUByte : FfiConverterRustBuffer<List<UByte>> {
override fun read(buf: ByteBuffer): List<UByte> {
val len = buf.getInt()
Expand Down Expand Up @@ -698,4 +768,26 @@ fun `generatePrivatePreferencesTopicIdentifier`(`privateKeyBytes`: List<UByte>):
})
}

@Throws(VerifyException::class)

fun `verifyK256Sha256`(
`signedBy`: List<UByte>,
`message`: List<UByte>,
`signature`: List<UByte>,
`recoveryId`: UByte,
): Boolean {
return FfiConverterBoolean.lift(
rustCallWithError(VerifyException) { _status ->
_UniFFILib.INSTANCE.uniffi_xmtp_dh_fn_func_verify_k256_sha256(
FfiConverterSequenceUByte.lower(
`signedBy`
),
FfiConverterSequenceUByte.lower(`message`),
FfiConverterSequenceUByte.lower(`signature`),
FfiConverterUByte.lower(`recoveryId`),
_status
)
})
}


Binary file modified library/src/main/jniLibs/arm64-v8a/libuniffi_xmtp_dh.so
Binary file not shown.
Binary file modified library/src/main/jniLibs/armeabi-v7a/libuniffi_xmtp_dh.so
Binary file not shown.
Binary file modified library/src/main/jniLibs/x86/libuniffi_xmtp_dh.so
Binary file not shown.
Binary file modified library/src/main/jniLibs/x86_64/libuniffi_xmtp_dh.so
Binary file not shown.

0 comments on commit 3ac6990

Please sign in to comment.