diff --git a/javasteam-depotdownloader/src/main/kotlin/in/dragonbra/javasteam/depotdownloader/Util.kt b/javasteam-depotdownloader/src/main/kotlin/in/dragonbra/javasteam/depotdownloader/Util.kt index 4c4dae03..ff51faa2 100644 --- a/javasteam-depotdownloader/src/main/kotlin/in/dragonbra/javasteam/depotdownloader/Util.kt +++ b/javasteam-depotdownloader/src/main/kotlin/in/dragonbra/javasteam/depotdownloader/Util.kt @@ -7,6 +7,12 @@ import `in`.dragonbra.javasteam.types.DepotManifest import `in`.dragonbra.javasteam.util.Adler32 import `in`.dragonbra.javasteam.util.log.LogManager import `in`.dragonbra.javasteam.util.log.Logger +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.async +import kotlinx.coroutines.awaitAll +import kotlinx.coroutines.coroutineScope +import kotlinx.coroutines.sync.Semaphore +import kotlinx.coroutines.sync.withPermit import okio.FileHandle import okio.FileSystem import okio.Path @@ -25,6 +31,11 @@ object Util { private val logger: Logger = LogManager.getLogger() + private const val MAX_VALIDATE_CONCURRENCY = 16 + + private fun defaultValidateConcurrency(): Int = + Runtime.getRuntime().availableProcessors().coerceIn(1, MAX_VALIDATE_CONCURRENCY) + @JvmOverloads @JvmStatic fun getSteamOS(androidEmulation: Boolean = false): String { @@ -147,27 +158,32 @@ object Util { * @throws IOException If there's an error reading the file */ @JvmStatic + @JvmOverloads @Throws(IOException::class) - fun validateSteam3FileChecksums(handle: FileHandle, chunkData: List): List { - val neededChunks = mutableListOf() - - for (data in chunkData) { - val chunk = ByteArray(data.uncompressedLength) - val read = handle.read(data.offset, chunk, 0, data.uncompressedLength) - - val tempChunk = if (read > 0 && read < data.uncompressedLength) { - chunk.copyOf(read) - } else { - chunk - } - - val adler = Adler32.calculate(tempChunk) - if (adler != data.checksum) { - neededChunks.add(data) + suspend fun validateSteam3FileChecksums( + handle: FileHandle, + chunkData: List, + concurrency: Int = defaultValidateConcurrency(), + ): List = coroutineScope { + val semaphore = Semaphore(concurrency.coerceAtLeast(1)) + + chunkData.map { data -> + async(Dispatchers.IO) { + semaphore.withPermit { + val chunk = ByteArray(data.uncompressedLength) + val read = handle.read(data.offset, chunk, 0, data.uncompressedLength) + + val tempChunk = if (read > 0 && read < data.uncompressedLength) { + chunk.copyOf(read) + } else { + chunk + } + + val adler = Adler32.calculate(tempChunk) + if (adler != data.checksum) data else null + } } - } - - return neededChunks + }.awaitAll().filterNotNull() } @JvmStatic diff --git a/javasteam-depotdownloader/src/test/kotlin/in/dragonbra/javasteam/depotdownloader/UtilValidateSteam3FileChecksumsTest.kt b/javasteam-depotdownloader/src/test/kotlin/in/dragonbra/javasteam/depotdownloader/UtilValidateSteam3FileChecksumsTest.kt new file mode 100644 index 00000000..48623b6d --- /dev/null +++ b/javasteam-depotdownloader/src/test/kotlin/in/dragonbra/javasteam/depotdownloader/UtilValidateSteam3FileChecksumsTest.kt @@ -0,0 +1,119 @@ +package `in`.dragonbra.javasteam.depotdownloader + +import `in`.dragonbra.javasteam.types.ChunkData +import `in`.dragonbra.javasteam.util.Adler32 +import kotlinx.coroutines.runBlocking +import okio.FileSystem +import okio.Path.Companion.toOkioPath +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.io.TempDir +import java.io.File +import kotlin.random.Random + +/** + * Verifies that [Util.validateSteam3FileChecksums] correctly identifies mismatched chunks, + * and that parallelizing the per-chunk checksum work does not change the result. + */ +class UtilValidateSteam3FileChecksumsTest { + + companion object { + const val CHUNK_SIZE = 64 + } + + private fun writeChunkedFile(tempDir: File, chunkCount: Int, chunkSize: Int = CHUNK_SIZE): Pair> { + val file = tempDir.resolve("data.bin") + val random = Random(seed = 42) + val chunks = mutableListOf() + + file.outputStream().use { out -> + for (i in 0 until chunkCount) { + val bytes = ByteArray(chunkSize) { random.nextInt().toByte() } + out.write(bytes) + chunks.add( + ChunkData( + offset = (i * chunkSize).toLong(), + uncompressedLength = chunkSize, + checksum = Adler32.calculate(bytes), + ) + ) + } + } + + return file to chunks + } + + @Test + fun `all chunks match returns empty list`(@TempDir tempDir: File) = runBlocking { + val (file, chunks) = writeChunkedFile(tempDir, chunkCount = 20) + + FileSystem.SYSTEM.openReadOnly(file.toOkioPath()).use { handle -> + val needed = Util.validateSteam3FileChecksums(handle, chunks) + assertTrue(needed.isEmpty(), "Expected no chunks to need re-download, got ${needed.size}") + } + } + + @Test + fun `mismatched chunks are returned`(@TempDir tempDir: File) = runBlocking { + val (file, chunks) = writeChunkedFile(tempDir, chunkCount = 20) + + // Corrupt the checksum of a known subset so those chunks appear stale. + val corruptedIndices = setOf(0, 5, 19) + val chunksWithCorruption = chunks.mapIndexed { index, chunk -> + if (index in corruptedIndices) chunk.copy(checksum = chunk.checksum xor 0x1) else chunk + } + + FileSystem.SYSTEM.openReadOnly(file.toOkioPath()).use { handle -> + val needed = Util.validateSteam3FileChecksums(handle, chunksWithCorruption) + val neededOffsets = needed.map { it.offset }.toSet() + val expectedOffsets = corruptedIndices.map { chunksWithCorruption[it].offset }.toSet() + + assertEquals(expectedOffsets, neededOffsets) + } + } + + @Test + fun `result order matches input order`(@TempDir tempDir: File) = runBlocking { + val (file, chunks) = writeChunkedFile(tempDir, chunkCount = 50) + val chunksWithCorruption = chunks.map { it.copy(checksum = it.checksum xor 0x1) } + + FileSystem.SYSTEM.openReadOnly(file.toOkioPath()).use { handle -> + val needed = Util.validateSteam3FileChecksums(handle, chunksWithCorruption) + assertEquals(chunksWithCorruption.map { it.offset }, needed.map { it.offset }) + } + } + + @Test + fun `concurrency parameter does not change correctness`(@TempDir tempDir: File) = runBlocking { + val (file, chunks) = writeChunkedFile(tempDir, chunkCount = 40) + val corruptedIndices = setOf(1, 2, 3, 30) + val chunksWithCorruption = chunks.mapIndexed { index, chunk -> + if (index in corruptedIndices) chunk.copy(checksum = chunk.checksum xor 0x1) else chunk + } + + FileSystem.SYSTEM.openReadOnly(file.toOkioPath()).use { handle -> + val serial = Util.validateSteam3FileChecksums(handle, chunksWithCorruption, concurrency = 1) + val parallel = Util.validateSteam3FileChecksums(handle, chunksWithCorruption, concurrency = 8) + + assertEquals(serial.map { it.offset }.toSet(), parallel.map { it.offset }.toSet()) + } + } + + @Test + fun `large chunk count completes and is correct`(@TempDir tempDir: File) = runBlocking { + val (file, chunks) = writeChunkedFile(tempDir, chunkCount = 1000, chunkSize = 16) + val corruptedIndices = (0 until 1000 step 37).toSet() + val chunksWithCorruption = chunks.mapIndexed { index, chunk -> + if (index in corruptedIndices) chunk.copy(checksum = chunk.checksum xor 0x1) else chunk + } + + FileSystem.SYSTEM.openReadOnly(file.toOkioPath()).use { handle -> + val needed = Util.validateSteam3FileChecksums(handle, chunksWithCorruption) + val neededOffsets = needed.map { it.offset }.toSet() + val expectedOffsets = corruptedIndices.map { chunksWithCorruption[it].offset }.toSet() + + assertEquals(expectedOffsets, neededOffsets) + } + } +}