Optimise and test IP2Country (#1684)

* Move ipv4Int to top level

* Remove redundant fun calls in ipv4ToCountry

* Add null safety to loadFile

* Close streams on failure

* Simplify cacheCountryForIP

* Add IP2CountryTest

* Generate binary

* Simplify ipv4Int

* Fix companion object visibility

* Use array instead of Treemap

* Synchronize OnionApi#paths

* Move csv

* Deduplicate locations csv

* Move ipToCode to gradle

* Use std lib binarySearch

---------

Co-authored-by: bemusementpark <bemusementpark>

app/build.gradle

@@ -105,6 +105,13 @@ android {
         String sharedTestDir = 'src/sharedTest/java'
         test.java.srcDirs += sharedTestDir
         androidTest.java.srcDirs += sharedTestDir
+        main {
+            assets.srcDirs += "$buildDir/generated/binary"
+        }
+        test {
+            resources.srcDirs += "$buildDir/generated/binary"
+            resources.srcDirs += "$projectDir/src/main/assets"
+        }
     }
 
     buildTypes {
@@ -242,6 +249,12 @@ android {
     }
 }
 
+apply {
+    from("ipToCode.gradle.kts")
+}
+
+preBuild.dependsOn ipToCode
+
 dependencies {
     implementation project(':content-descriptions')
 

app/ipToCode.gradle.kts

@@ -0,0 +1,41 @@
+import java.io.File
+import java.io.DataOutputStream
+import java.io.FileOutputStream
+
+task("ipToCode") {
+    val inputFile = File("${projectDir}/geolite2_country_blocks_ipv4.csv")
+
+    val outputDir = "${buildDir}/generated/binary"
+    val outputFile = File(outputDir, "geolite2_country_blocks_ipv4.bin").apply { parentFile.mkdirs() }
+
+    outputs.file(outputFile)
+
+    doLast {
+
+        // Ensure the input file exists
+        if (!inputFile.exists()) {
+            throw IllegalArgumentException("Input file does not exist: ${inputFile.absolutePath}")
+        }
+
+        // Create a DataOutputStream to write binary data
+        DataOutputStream(FileOutputStream(outputFile)).use { out ->
+            inputFile.useLines { lines ->
+                var prevCode = -1
+                lines.drop(1).forEach { line ->
+                    runCatching {
+                        val ints = line.split(".", "/", ",")
+                        val code = ints[5].toInt().also { if (it == prevCode) return@forEach }
+                        val ip = ints.take(4).fold(0) { acc, s -> acc shl 8 or s.toInt() }
+
+                        out.writeInt(ip)
+                        out.writeInt(code)
+
+                        prevCode = code
+                    }
+                }
+            }
+        }
+
+        println("Processed data written to: ${outputFile.absolutePath}")
+    }
+}

app/src/main/java/org/thoughtcrime/securesms/util/IP2Country.kt

@@ -9,52 +9,46 @@ import com.opencsv.CSVReader
 import org.session.libsession.snode.OnionRequestAPI
 import org.session.libsignal.utilities.Log
 import org.session.libsignal.utilities.ThreadUtils
-import java.io.File
+import java.io.DataInputStream
-import java.io.FileOutputStream
+import java.io.InputStream
-import java.io.FileReader
+import java.io.InputStreamReader
-import java.util.SortedMap
+import kotlin.math.absoluteValue
-import java.util.TreeMap
 
-class IP2Country private constructor(private val context: Context) {
+private fun ipv4Int(ip: String): UInt =
-    private val pathsBuiltEventReceiver: BroadcastReceiver
+    ip.split(".", "/", ",").take(4).fold(0U) { acc, s -> acc shl 8 or s.toUInt() }
+
+@OptIn(ExperimentalUnsignedTypes::class)
+class IP2Country internal constructor(
+    private val context: Context,
+    private val openStream: (String) -> InputStream = context.assets::open
+) {
     val countryNamesCache = mutableMapOf<String, String>()
 
-    private fun Ipv4Int(ip: String): Int {
+    private val ips: UIntArray by lazy { ipv4ToCountry.first }
-        var result = 0L
+    private val codes: IntArray by lazy { ipv4ToCountry.second }
-        var currentValue = 0L
-        var octetIndex = 0
 
-        for (char in ip) {
+    private val ipv4ToCountry: Pair<UIntArray, IntArray> by lazy {
-            if (char == '.' || char == '/') {
+        openStream("geolite2_country_blocks_ipv4.bin")
-                result = result or (currentValue shl (8 * (3 - octetIndex)))
+            .let(::DataInputStream)
-                currentValue = 0
+            .use {
-                octetIndex++
+                val size = it.available() / 8
-                if (char == '/') break
+
-            } else {
+                val ips = UIntArray(size)
-                currentValue = currentValue * 10 + (char - '0')
+                val codes = IntArray(size)
-            }
+                var i = 0
+
+                while (it.available() > 0) {
+                    ips[i] = it.readInt().toUInt()
+                    codes[i] = it.readInt()
+                    i++
                 }
 
-        // Handle the last octet
+                ips to codes
-        result = result or (currentValue shl (8 * (3 - octetIndex)))
-
-        return result.toInt()
-    }
-
-    private val ipv4ToCountry: TreeMap<Int, Int?> by lazy {
-        val file = loadFile("geolite2_country_blocks_ipv4.csv")
-        CSVReader(FileReader(file.absoluteFile)).use { csv ->
-            csv.skip(1)
-
-            csv.asSequence().associateTo(TreeMap()) { cols ->
-                Ipv4Int(cols[0]).toInt() to cols[1].toIntOrNull()
-            }
             }
     }
 
     private val countryToNames: Map<Int, String> by lazy {
-        val file = loadFile("geolite2_country_locations_english.csv")
+        CSVReader(InputStreamReader(openStream("csv/geolite2_country_locations_english.csv"))).use { csv ->
-        CSVReader(FileReader(file.absoluteFile)).use { csv ->
             csv.skip(1)
 
             csv.asSequence()
@@ -68,81 +62,57 @@ class IP2Country private constructor(private val context: Context) {
     // region Initialization
     companion object {
 
-        public lateinit var shared: IP2Country
+        lateinit var shared: IP2Country
 
-        public val isInitialized: Boolean get() = Companion::shared.isInitialized
+        val isInitialized: Boolean get() = Companion::shared.isInitialized
 
-        public fun configureIfNeeded(context: Context) {
+        fun configureIfNeeded(context: Context) {
             if (isInitialized) { return; }
             shared = IP2Country(context.applicationContext)
+
+            val pathsBuiltEventReceiver = object : BroadcastReceiver() {
+                override fun onReceive(context: Context, intent: Intent) {
+                    shared.populateCacheIfNeeded()
+                }
+            }
+            LocalBroadcastManager.getInstance(context).registerReceiver(pathsBuiltEventReceiver, IntentFilter("pathsBuilt"))
         }
     }
 
     init {
         populateCacheIfNeeded()
-        pathsBuiltEventReceiver = object : BroadcastReceiver() {
-            override fun onReceive(context: Context, intent: Intent) {
-                populateCacheIfNeeded()
-            }
-        }
-        LocalBroadcastManager.getInstance(context).registerReceiver(pathsBuiltEventReceiver, IntentFilter("pathsBuilt"))
     }
 
     // TODO: Deinit?
     // endregion
 
     // region Implementation
-    private fun loadFile(fileName: String): File {
+    internal fun cacheCountryForIP(ip: String): String? {
-        val directory = File(context.applicationInfo.dataDir)
-        val file = File(directory, fileName)
-        if (directory.list().contains(fileName)) { return file }
-        val inputStream = context.assets.open("csv/$fileName")
-        val outputStream = FileOutputStream(file)
-        val buffer = ByteArray(1024)
-        while (true) {
-            val count = inputStream.read(buffer)
-            if (count < 0) { break }
-            outputStream.write(buffer, 0, count)
-        }
-        inputStream.close()
-        outputStream.close()
-        return file
-    }
-
-    private fun cacheCountryForIP(ip: String): String? {
-
         // return early if cached
         countryNamesCache[ip]?.let { return it }
 
-        val ipInt = Ipv4Int(ip)
+        val ipInt = ipv4Int(ip)
-        val bestMatchCountry = ipv4ToCountry.floorEntry(ipInt)?.let { (_, code) ->
+        val index = ips.binarySearch(ipInt).let { it.takeIf { it >= 0 } ?: (it.absoluteValue - 2) }
-            if (code != null) {
+        val code = codes.getOrNull(index)
-                countryToNames[code]
+        val bestMatchCountry = countryToNames[code]
-            } else {
+
-                null
+        if (bestMatchCountry != null) countryNamesCache[ip] = bestMatchCountry
-            }
+        else Log.d("Loki","Country name for $ip couldn't be found")
-        }
 
-        if (bestMatchCountry != null) {
-            countryNamesCache[ip] = bestMatchCountry
         return bestMatchCountry
-        } else {
-            Log.d("Loki","Country name for $ip couldn't be found")
-        }
-        return null
     }
 
     private fun populateCacheIfNeeded() {
         ThreadUtils.queue {
+            val start = System.currentTimeMillis()
             OnionRequestAPI.paths.iterator().forEach { path ->
                 path.iterator().forEach { snode ->
                     cacheCountryForIP(snode.ip) // Preload if needed
                 }
             }
+            Log.d("Loki","Cache populated in ${System.currentTimeMillis() - start}ms")
             Broadcaster(context).broadcast("onionRequestPathCountriesLoaded")
-            Log.d("Loki", "Finished preloading onion request path countries.")
         }
     }
     // endregion
 }
-

app/src/test/java/org/thoughtcrime/securesms/util/IP2CountryTest.kt

@@ -0,0 +1,40 @@
+package org.thoughtcrime.securesms.util
+
+import android.content.Context
+import org.junit.Assert.assertEquals
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.Parameterized
+import org.mockito.Mockito.mock
+
+@RunWith(Parameterized::class)
+class IP2CountryTest(
+    private val ip: String,
+    private val country: String
+) {
+    private val context: Context = mock(Context::class.java)
+    private val ip2Country = IP2Country(context, this::class.java.classLoader!!::getResourceAsStream)
+
+    @Test
+    fun getCountryNamesCache() {
+        assertEquals(country, ip2Country.cacheCountryForIP(ip))
+    }
+
+    companion object {
+        @JvmStatic
+        @Parameterized.Parameters
+        fun data(): Collection<Array<Any>> = listOf(
+            arrayOf("223.121.64.0", "Hong Kong"),
+            arrayOf("223.121.64.1", "Hong Kong"),
+            arrayOf("223.121.127.0", "Hong Kong"),
+            arrayOf("223.121.128.0", "China"),
+            arrayOf("223.121.129.0", "China"),
+            arrayOf("223.122.0.0", "Hong Kong"),
+            arrayOf("223.123.0.0", "Pakistan"),
+            arrayOf("223.123.128.0", "China"),
+            arrayOf("223.124.0.0", "China"),
+            arrayOf("223.128.0.0", "China"),
+            arrayOf("223.130.0.0", "Singapore")
+        )
+    }
+}

libsession/src/main/java/org/session/libsession/snode/OnionRequestAPI.kt

@@ -44,7 +44,8 @@ object OnionRequestAPI {
 
     var guardSnodes = setOf<Snode>()
     var _paths: AtomicReference<List<Path>?> = AtomicReference(null)
-    var paths: List<Path> // Not a set to ensure we consistently show the same path to the user
+    var paths: List<Path> // Not a Set to ensure we consistently show the same path to the user
+        @Synchronized
         get() {
             val paths = _paths.get()
 
@@ -57,6 +58,7 @@ object OnionRequestAPI {
             _paths.set(result)
             return result
         }
+        @Synchronized
         set(newValue) {
             if (newValue.isEmpty()) {
                 database.clearOnionRequestPaths()