Improve performance of the no-overlap check

This commit is contained in:
Tudor Malene 2019-05-28 14:37:23 +01:00 committed by Mike Hearn
parent aa75157273
commit a4cd859282
3 changed files with 36 additions and 15 deletions

View File

@ -19,6 +19,7 @@ import java.io.IOException
import java.io.InputStream
import java.net.*
import java.util.*
import java.util.jar.JarInputStream
/**
* A custom ClassLoader that knows how to load classes from a set of attachments. The attachments themselves only
@ -190,7 +191,7 @@ class AttachmentsClassLoader(attachments: List<Attachment>,
// attacks on externally connected systems that only consider type names, we allow people to formally
// claim their parts of the Java package namespace via registration with the zone operator.
val classLoaderEntries = mutableMapOf<String, Attachment>()
val classLoaderEntries = mutableMapOf<String, SecureHash.SHA256>()
for (attachment in attachments) {
// We may have been given an attachment loaded from the database in which case, important info like
// signers is already calculated.
@ -251,21 +252,26 @@ class AttachmentsClassLoader(attachments: List<Attachment>,
// Some files don't need overlap checking because they don't affect the way the code runs.
if (!shouldCheckForNoOverlap(path, targetPlatformVersion)) continue
// If 2 entries have the same content hash, it means the same file is present in both attachments, so that is ok.
if (path in classLoaderEntries.keys) {
val contentHash = readAttachment(attachment, path).sha256()
val originalAttachment = classLoaderEntries[path]!!
val originalContentHash = readAttachment(originalAttachment, path).sha256()
if (contentHash == originalContentHash) {
log.debug { "Duplicate entry $path has same content hash $contentHash" }
continue
} else {
// This calculates the hash of the current entry because the JarInputStream returns only the current entry.
fun entryHash() = ByteArrayOutputStream().use {
jar.copyTo(it)
it.toByteArray()
}.sha256()
// If 2 entries are identical, it means the same file is present in both attachments, so that is ok.
val currentHash = entryHash()
val previousFileHash = classLoaderEntries[path]
when {
previousFileHash == null -> {
log.debug { "Adding new entry for $path" }
classLoaderEntries[path] = currentHash
}
currentHash == previousFileHash -> log.debug { "Duplicate entry $path has same content hash $currentHash" }
else -> {
log.debug { "Content hash differs for $path" }
throw OverlappingAttachmentsException(sampleTxId, path)
}
}
log.debug { "Adding new entry for $path" }
classLoaderEntries[path] = attachment
}
}
log.debug { "${classLoaderEntries.size} classloaded entries for $attachment" }

View File

@ -107,8 +107,8 @@ class AttachmentsClassLoaderTests {
@Test
fun `Test valid overlapping file condition`() {
val att1 = importAttachment(fakeAttachment("file1.txt", "same data").inputStream(), "app", "file1.jar")
val att2 = importAttachment(fakeAttachment("file1.txt", "same data").inputStream(), "app", "file2.jar")
val att1 = importAttachment(fakeAttachment("file1.txt", "same data", "file2.txt", "same other data" ).inputStream(), "app", "file1.jar")
val att2 = importAttachment(fakeAttachment("file1.txt", "same data", "file3.txt", "same totally different").inputStream(), "app", "file2.jar")
val cl = make(arrayOf(att1, att2).map { storage.openAttachment(it)!! })
val txt = IOUtils.toString(cl.getResourceAsStream("file1.txt"), Charsets.UTF_8.name())

View File

@ -200,6 +200,21 @@ fun fakeAttachment(filePath: String, content: String, manifestAttributes: Map<St
return bs.toByteArray()
}
fun fakeAttachment(filePath1: String, content1: String, filePath2: String, content2: String, manifestAttributes: Map<String, String> = emptyMap()): ByteArray {
val bs = ByteArrayOutputStream()
val manifest = Manifest()
manifestAttributes.forEach { manifest[it.key] = it.value } //adding manually instead of putAll, as it requires typed keys, not strings
JarOutputStream(bs, manifest).use { js ->
js.putNextEntry(ZipEntry(filePath1))
js.writer().apply { append(content1); flush() }
js.closeEntry()
js.putNextEntry(ZipEntry(filePath2))
js.writer().apply { append(content2); flush() }
js.closeEntry()
}
return bs.toByteArray()
}
/** If [effectiveSerializationEnv] is not set, runs the block with a new [SerializationEnvironmentRule]. */
fun <R> withTestSerializationEnvIfNotSet(block: () -> R): R {
val serializationExists = try {