[compiler] fix UAF on getBytecodeProducerString() (#183)

* fix UAF: `BytecodeWriterConfig` accepts a `StringRef`. * export torch_frontend.DebugType
bytedance · Apr 7, 2024 · fcc858f · fcc858f
1 parent 0399e21
commit fcc858f
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 3 deletions.
diff --git a/compiler/lib/CAPI/Translation.cpp b/compiler/lib/CAPI/Translation.cpp
@@ -122,7 +122,8 @@ bool byteirSerializeByre(MlirModule module, MlirStringRef targetVersion,
     llvm::errs() << errorMessage << "\n";
     return false;
   }
-  BytecodeWriterConfig config(serialVersion.getBytecodeProducerString());
+  std::string producerString = serialVersion.getBytecodeProducerString();
+  BytecodeWriterConfig config(producerString);
   config.setDesiredBytecodeVersion(serialVersion.getBytecodeVersion());
   if (failed(writeBytecodeToFile(*newModule, resultMLIRBCFile->os(), config))) {
     newModule->emitOpError() << "failed to write bytecode\n";

diff --git a/compiler/lib/Dialect/Byre/Transforms/Serial.cpp b/compiler/lib/Dialect/Byre/Transforms/Serial.cpp
@@ -78,7 +78,8 @@ struct DumpByrePass : public DumpByreBase<DumpByrePass> {
       return signalPassFailure();
     }
 
-    BytecodeWriterConfig config(targetVersion->getBytecodeProducerString());
+    std::string producerString = targetVersion->getBytecodeProducerString();
+    BytecodeWriterConfig config(producerString);
     config.setDesiredBytecodeVersion(targetVersion->getBytecodeVersion());
     if (failed(writeBytecodeToFile(*newModule, ofile->os(), config))) {
       newModule->emitOpError() << "failed to write bytecode\n";

diff --git a/compiler/test/lib/Transformation/TestByreSerialRoundtrip.cpp b/compiler/test/lib/Transformation/TestByreSerialRoundtrip.cpp
@@ -209,7 +209,8 @@ struct TestByreSerialRoundtripPass
 
     std::string buffer;
     llvm::raw_string_ostream ostream(buffer);
-    BytecodeWriterConfig config(targetVersion.getBytecodeProducerString());
+    std::string producerString = targetVersion.getBytecodeProducerString();
+    BytecodeWriterConfig config(producerString);
     config.setDesiredBytecodeVersion(targetVersion.getBytecodeVersion());
     if (failed(writeBytecodeToFile(*newModule, ostream, config))) {
       newModule->emitOpError() << "failed to write bytecode\n";

diff --git a/frontends/torch-frontend/torch-frontend/python/torch_frontend/__init__.py b/frontends/torch-frontend/torch-frontend/python/torch_frontend/__init__.py
@@ -1,5 +1,7 @@
 from ._mlir_libs._torchFrontend import *
 
+from .compile import DebugType
+
 from .fx_utils import list_decomposed_ops, preprocess_fx_graph, get_none_indices
 from .compile import convert_to_mhlo_via_torch_mlir, compile, compile_dynamo_model
 from .flash_attn_op import replace_flash_attn