kshark.HprofDeobfuscator.kt Maven / Gradle / Ivy

Go to download
package kshark

import kshark.HprofHeader.Companion.parseHeaderOf
import kshark.HprofRecord.HeapDumpEndRecord
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.ClassDumpRecord
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.ClassDumpRecord.FieldRecord
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.ClassDumpRecord.StaticFieldRecord
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.InstanceDumpRecord
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.ObjectArrayDumpRecord
import kshark.HprofRecord.HeapDumpRecord.ObjectRecord.PrimitiveArrayDumpRecord
import kshark.HprofRecord.LoadClassRecord
import kshark.HprofRecord.StackFrameRecord
import kshark.HprofRecord.StringRecord
import kshark.StreamingRecordReaderAdapter.Companion.asStreamingRecordReader
import java.io.File

/**
 * Converts a Hprof file to another file with deobfuscated class and field names.
 */
class HprofDeobfuscator {

  /**
   * @see HprofDeobfuscator
   */
  fun deobfuscate(
    proguardMapping: ProguardMapping,
    inputHprofFile: File,
    /**
     * Optional output file. Defaults to a file in the same directory as [inputHprofFile], with
     * the same name and "-deobfuscated" prepended before the ".hprof" extension. If the file extension
     * is not ".hprof", then "-deobfuscated" is added at the end of the file.
     */
    outputHprofFile: File = File(
      inputHprofFile.parent, inputHprofFile.name.replace(
      ".hprof", "-deobfuscated.hprof"
    ).let { if (it != inputHprofFile.name) it else inputHprofFile.name + "-deobfuscated" })
  ): File {
    val (hprofStringCache, classNames, maxId) = readHprofRecords(inputHprofFile)

    return writeHprofRecords(
      inputHprofFile,
      outputHprofFile,
      proguardMapping,
      hprofStringCache,
      classNames,
      maxId + 1
    )
  }

  /**
   * Reads StringRecords and LoadClassRecord from an Hprof file and tracks maximum HprofRecord id
   * value.
   *
   * @return a Triple of: hprofStringCache map, classNames map and maxId value
   */
  private fun readHprofRecords(
    inputHprofFile: File
  ): Triple, Map, Long> {
    val hprofStringCache = mutableMapOf()
    val classNames = mutableMapOf()

    var maxId: Long = 0

    val reader = StreamingHprofReader.readerFor(inputHprofFile).asStreamingRecordReader()
    reader.readRecords(setOf(HprofRecord::class),
      OnHprofRecordListener { _, record ->
        when (record) {
          is StringRecord -> {
            maxId = maxId.coerceAtLeast(record.id)
            hprofStringCache[record.id] = record.string
          }
          is LoadClassRecord -> {
            maxId = maxId.coerceAtLeast(record.id)
            classNames[record.id] = record.classNameStringId
          }
          is StackFrameRecord -> maxId = maxId.coerceAtLeast(record.id)
          is ObjectRecord -> {
            maxId = when (record) {
              is ClassDumpRecord -> maxId.coerceAtLeast(record.id)
              is InstanceDumpRecord -> maxId.coerceAtLeast(record.id)
              is ObjectArrayDumpRecord -> maxId.coerceAtLeast(record.id)
              is PrimitiveArrayDumpRecord -> maxId.coerceAtLeast(record.id)
            }
          }
        }
      })
    return Triple(hprofStringCache, classNames, maxId)
  }

  @Suppress("LongParameterList")
  private fun writeHprofRecords(
    inputHprofFile: File,
    outputHprofFile: File,
    proguardMapping: ProguardMapping,
    hprofStringCache: Map,
    classNames: Map,
    firstId: Long
  ): File {
    var id = firstId

    val hprofHeader = parseHeaderOf(inputHprofFile)
    val reader =
      StreamingHprofReader.readerFor(inputHprofFile, hprofHeader).asStreamingRecordReader()
    HprofWriter.openWriterFor(
      outputHprofFile,
      hprofHeader = HprofHeader(
        identifierByteSize = hprofHeader.identifierByteSize,
        version = hprofHeader.version
      )
    ).use { writer ->
      reader.readRecords(setOf(HprofRecord::class),
        OnHprofRecordListener { _,
          record ->
          // HprofWriter automatically emits HeapDumpEndRecord, because it flushes
          // all continuous heap dump sub records as one heap dump record.
          if (record is HeapDumpEndRecord) {
            return@OnHprofRecordListener
          }

          when (record) {
            is StringRecord -> {
              writer.write(
                createDeobfuscatedStringRecord(record, proguardMapping, hprofStringCache)
              )
            }
            is ClassDumpRecord -> {
              val (recordsToWrite, maxId) = createDeobfuscatedClassDumpRecord(
                record, proguardMapping, hprofStringCache, classNames, id
              )
              id = maxId
              recordsToWrite.forEach {
                writer.write(it)
              }
            }
            else -> writer.write(record)
          }
        })
    }

    return outputHprofFile
  }

  private fun createDeobfuscatedStringRecord(
    record: StringRecord,
    proguardMapping: ProguardMapping,
    hprofStringCache: Map
  ): StringRecord {
    val obfuscatedName = hprofStringCache[record.id]!!
    return StringRecord(
      record.id, proguardMapping.deobfuscateClassName(obfuscatedName)
    )
  }

  /**
   * Deobfuscated ClassDumpRecord's field names. Different classes can have fields with the same
   * names. We need to generate new StringRecords in such cases.
   *
   * @return a Pair of: list of HprofRecords to write and new maxId value
   */
  private fun createDeobfuscatedClassDumpRecord(
    record: ClassDumpRecord,
    proguardMapping: ProguardMapping,
    hprofStringCache: Map,
    classNames: Map,
    maxId: Long
  ): Pair, Long> {
    val recordsToWrite = mutableListOf()

    var id = maxId

    val newFields = record.fields.map { field ->
      val className = hprofStringCache[classNames[record.id]]!!
      val fieldName = hprofStringCache[field.nameStringId]!!
      val deobfuscatedName =
        proguardMapping.deobfuscateFieldName(className, fieldName)

      val newStringRecord = StringRecord(id++, deobfuscatedName)
      recordsToWrite.add(newStringRecord)

      FieldRecord(newStringRecord.id, field.type)
    }
    val newStaticFields = record.staticFields.map { field ->
      val className = hprofStringCache[classNames[record.id]]!!
      val fieldName = hprofStringCache[field.nameStringId]!!
      val deobfuscatedName =
        proguardMapping.deobfuscateFieldName(className, fieldName)

      val newStringRecord = StringRecord(id++, deobfuscatedName)
      recordsToWrite.add(newStringRecord)

      StaticFieldRecord(newStringRecord.id, field.type, field.value)
    }

    recordsToWrite.add(
      ClassDumpRecord(
        record.id,
        record.stackTraceSerialNumber,
        record.superclassId,
        record.classLoaderId,
        record.signersId,
        record.protectionDomainId,
        record.instanceSize,
        newStaticFields,
        newFields
      )
    )

    return Pair(recordsToWrite, id)
  }
}