Kotlin extractor: support Zstd de/compression

java/kotlin-extractor/build.gradle

@@ -9,6 +9,7 @@ version '0.0.1'
 dependencies {
  implementation "org.jetbrains.kotlin:kotlin-stdlib"
  compileOnly("org.jetbrains.kotlin:kotlin-compiler")
+ implementation "io.airlift:aircompressor"
 }
 
 repositories {

java/kotlin-extractor/build.py

@@ -206,8 +206,8 @@ def compile(jars, java_jars, dependency_folder, transform_to_embeddable, output,
 
 
 def compile_embeddable(version):
- compile(['kotlin-stdlib-' + version, 'kotlin-compiler-embeddable-' + version],
- ['kotlin-stdlib-' + version],
+ compile(['kotlin-stdlib-' + version, 'kotlin-compiler-embeddable-' + version, 'aircompressor-0.26'],
+ ['kotlin-stdlib-' + version, 'aircompressor-0.26'],
  kotlin_dependency_folder,
  transform_to_embeddable,
  'codeql-extractor-kotlin-embeddable-%s.jar' % (version),
@@ -216,8 +216,8 @@ def compile_embeddable(version):
 
 
 def compile_standalone(version):
- compile(['kotlin-stdlib-' + version, 'kotlin-compiler-' + version],
- ['kotlin-stdlib-' + version],
+ compile(['kotlin-stdlib-' + version, 'kotlin-compiler-' + version, 'aircompressor-0.26'],
+ ['kotlin-stdlib-' + version, 'aircompressor-0.26'],
  kotlin_dependency_folder,
  lambda srcs: None,
  'codeql-extractor-kotlin-standalone-%s.jar' % (version),

java/kotlin-extractor/src/main/java/com/semmle/util/trap/CompressedFileInputStream.java

@@ -1,29 +1,59 @@
 package com.semmle.util.trap;
 
+import com.semmle.util.zip.MultiMemberGZIPInputStream;
+import io.airlift.compress.zstd.ZstdInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.file.Files;
 import java.nio.file.Path;
 
-import com.semmle.util.zip.MultiMemberGZIPInputStream;
-
 public class CompressedFileInputStream {
- /**
- * Create an input stream for reading the uncompressed data from a (possibly) compressed file, with
- * the decompression method chosen based on the file extension.
- *
- * @param f The compressed file to read
- * @return An input stream from which you can read the file's uncompressed data.
- * @throws IOException From the underlying decompression input stream.
- */
- public static InputStream fromFile(Path f) throws IOException {
- InputStream fileInputStream = Files.newInputStream(f);
- if (f.getFileName().toString().endsWith(".gz")) {
- return new MultiMemberGZIPInputStream(fileInputStream, 8192);
- //} else if (f.getFileName().toString().endsWith(".br")) {
- // return new BrotliInputStream(fileInputStream);
- } else {
- return fileInputStream;
- }
- }
+ /**
+ * Create an input stream for reading the uncompressed data from a (possibly) compressed file,
+ * with the decompression method chosen based on the file extension.
+ *
+ * @param f The compressed file to read
+ * @return An input stream from which you can read the file's uncompressed data.
+ * @throws IOException From the underlying decompression input stream.
+ */
+ public static InputStream fromFile(Path f) throws IOException {
+ InputStream fileInputStream = Files.newInputStream(f);
+ String fileName = f.getFileName().toString();
+ if (fileName.endsWith(".gz")) {
+ return new MultiMemberGZIPInputStream(fileInputStream, 8192);
+ //} else if (fileName.endsWith(".br")) {
+ // return new BrotliInputStream(fileInputStream);
+ } else if (fileName.endsWith(".zst")) {
+ return new WrappedZstdInputStream(fileInputStream);
+ } else {
+ return fileInputStream;
+ }
+ }
+
+ // Turn the MalformedInputException thrown by the ZstdInputStream into an IOException,
+ // which will be handled as a non-catastrophic error during TRAP import.
+ private static class WrappedZstdInputStream extends ZstdInputStream {
+ public WrappedZstdInputStream(InputStream in) {
+ super(in);
+ }
+
+ @Override
+ public int read() throws IOException {
+ try {
+ return super.read();
+ } catch (io.airlift.compress.MalformedInputException e) {
+ throw new IOException("Zstd stream decoding failed", e);
+ }
+ }
+
+ @Override
+ public int read(final byte[] outputBuffer, final int outputOffset, final int outputLength)
+ throws IOException {
+ try {
+ return super.read(outputBuffer, outputOffset, outputLength);
+ } catch (io.airlift.compress.MalformedInputException e) {
+ throw new IOException("Zstd stream decoding failed", e);
+ }
+ }
+ }
 }

java/kotlin-extractor/src/main/kotlin/KotlinExtractorExtension.kt

@@ -3,6 +3,8 @@ package com.github.codeql
 import com.github.codeql.utils.versions.usesK2
 import com.semmle.util.files.FileUtil
 import com.semmle.util.trap.pathtransformers.PathTransformer
+import io.airlift.compress.zstd.ZstdInputStream
+import io.airlift.compress.zstd.ZstdOutputStream
 import java.io.BufferedInputStream
 import java.io.BufferedOutputStream
 import java.io.BufferedReader
@@ -210,17 +212,17 @@ class KotlinExtractorExtension(
  private fun getCompression(logger: Logger): Compression {
  val compression_env_var = "CODEQL_EXTRACTOR_JAVA_OPTION_TRAP_COMPRESSION"
  val compression_option = System.getenv(compression_env_var)
- val defaultCompression = Compression.GZIP
+ val defaultCompression = Compression.ZSTD
  if (compression_option == null) {
  return defaultCompression
  } else {
  try {
  val compression_option_upper = compression_option.uppercase()
  if (compression_option_upper == "BROTLI") {
  logger.warn(
- "Kotlin extractor doesn't support Brotli compression. Using GZip instead."
+ "Kotlin extractor doesn't support Brotli compression. Using Zstandard instead."
  )
- return Compression.GZIP
+ return Compression.ZSTD
  } else {
  return Compression.valueOf(compression_option_upper)
  }
@@ -470,6 +472,11 @@ enum class Compression(val extension: String) {
  override fun bufferedWriter(file: File): BufferedWriter {
  return GZIPOutputStream(file.outputStream()).bufferedWriter()
  }
+ },
+ ZSTD(".zst") {
+ override fun bufferedWriter(file: File): BufferedWriter {
+ return ZstdOutputStream(file.outputStream()).bufferedWriter()
+ }
  };
 
  abstract fun bufferedWriter(file: File): BufferedWriter
@@ -483,6 +490,7 @@ private fun getTrapFileWriter(
  return when (compression) {
  Compression.NONE -> NonCompressedTrapFileWriter(logger, trapFileName)
  Compression.GZIP -> GZipCompressedTrapFileWriter(logger, trapFileName)
+ Compression.ZSTD -> ZstdCompressedTrapFileWriter(logger, trapFileName)
  }
 }
 
@@ -582,3 +590,18 @@ private class GZipCompressedTrapFileWriter(logger: FileLogger, trapName: String)
  )
  }
 }
+
+private class ZstdCompressedTrapFileWriter(logger: FileLogger, trapName: String) :
+ TrapFileWriter(logger, trapName, ".zst") {
+ override protected fun getReader(file: File): BufferedReader {
+ return BufferedReader(
+ InputStreamReader(ZstdInputStream(BufferedInputStream(FileInputStream(file))))
+ )
+ }
+
+ override protected fun getWriter(file: File): BufferedWriter {
+ return BufferedWriter(
+ OutputStreamWriter(ZstdOutputStream(BufferedOutputStream(FileOutputStream(file))))
+ )
+ }
+}