From c55e33f9f9e2b6fa795d606c4abbe27b3a785e8b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=82=B5=E6=98=8E=E7=9A=93?= <minghao.shao@spacemit.com>
Date: Tue, 28 Jul 2026 15:44:58 +0800
Subject: [PATCH] fix: harden runtime FST and JNI handling

---
 runtime/bindings/android/wetextprocessing.cc | 78 +++++++++++++++++---
 runtime/processor/wetext_processor.cc        |  9 +++
 runtime/test/processor_test.cc               |  8 ++
 3 files changed, 86 insertions(+), 9 deletions(-)

diff --git a/runtime/bindings/android/wetextprocessing.cc b/runtime/bindings/android/wetextprocessing.cc
index 3544cfc..d43ca94 100644
--- a/runtime/bindings/android/wetextprocessing.cc
+++ b/runtime/bindings/android/wetextprocessing.cc
@@ -13,6 +13,9 @@
 // limitations under the License.
 #include <jni.h>
 
+#include <exception>
+#include <string>
+
 #include "processor/wetext_processor.h"
 #include "utils/wetext_flags.h"
 #include "utils/wetext_string.h"
@@ -22,20 +25,74 @@ namespace wetextprocessing {
 std::shared_ptr<wetext::Processor> processorTN;
 std::shared_ptr<wetext::Processor> processorITN;
 
+namespace {
+class UtfStringChars {
+ public:
+  UtfStringChars(JNIEnv* env, jstring string)
+      : env_(env),
+        string_(string),
+        chars_(env->GetStringUTFChars(string, nullptr)) {}
+
+  UtfStringChars(const UtfStringChars&) = delete;
+  UtfStringChars& operator=(const UtfStringChars&) = delete;
+
+  ~UtfStringChars() {
+    if (chars_ != nullptr) {
+      env_->ReleaseStringUTFChars(string_, chars_);
+    }
+  }
+
+  const char* get() const { return chars_; }
+
+ private:
+  JNIEnv* env_;
+  jstring string_;
+  const char* chars_;
+};
+
+bool CopyJString(JNIEnv* env, jstring string, std::string* output) {
+  UtfStringChars chars(env, string);
+  if (chars.get() == nullptr) {
+    return false;
+  }
+  *output = chars.get();
+  return true;
+}
+
+void ThrowIllegalState(JNIEnv* env, const char* message) {
+  jclass exception_class = env->FindClass("java/lang/IllegalStateException");
+  if (exception_class != nullptr) {
+    env->ThrowNew(exception_class, message);
+  }
+}
+}
+
 void init(JNIEnv* env, jobject, jstring jModelDir) {
-  const char* pModelDir = env->GetStringUTFChars(jModelDir, nullptr);
+  std::string model_dir;
+  if (!CopyJString(env, jModelDir, &model_dir)) {
+    return;
+  }
 
-  std::string tnTagger = std::string(pModelDir) + "/zh_tn_tagger.fst";
-  std::string tnVerbalizer = std::string(pModelDir) + "/zh_tn_verbalizer.fst";
-  processorTN = std::make_shared<wetext::Processor>(tnTagger, tnVerbalizer);
+  try {
+    std::string tnTagger = model_dir + "/zh_tn_tagger.fst";
+    std::string tnVerbalizer = model_dir + "/zh_tn_verbalizer.fst";
+    processorTN = std::make_shared<wetext::Processor>(tnTagger, tnVerbalizer);
 
-  std::string itnTagger = std::string(pModelDir) + "/zh_itn_tagger.fst";
-  std::string itnVerbalizer = std::string(pModelDir) + "/zh_itn_verbalizer.fst";
-  processorITN = std::make_shared<wetext::Processor>(itnTagger, itnVerbalizer);
+    std::string itnTagger = model_dir + "/zh_itn_tagger.fst";
+    std::string itnVerbalizer = model_dir + "/zh_itn_verbalizer.fst";
+    processorITN = std::make_shared<wetext::Processor>(itnTagger, itnVerbalizer);
+  } catch (const std::exception& error) {
+    processorTN.reset();
+    processorITN.reset();
+    ThrowIllegalState(env, error.what());
+  }
 }
 
 jstring normalize(JNIEnv* env, jobject, jstring input) {
-  std::string input_text = std::string(env->GetStringUTFChars(input, nullptr));
+  std::string input_text;
+  if (!CopyJString(env, input, &input_text)) {
+    return nullptr;
+  }
   std::string tagged_text = processorTN->Tag(input_text);
   std::string normalized_text = processorTN->Verbalize(tagged_text);
 
@@ -43,7 +100,10 @@ jstring normalize(JNIEnv* env, jobject, jstring input) {
 }
 
 jstring inverse_normalize(JNIEnv* env, jobject, jstring input) {
-  std::string input_text = std::string(env->GetStringUTFChars(input, nullptr));
+  std::string input_text;
+  if (!CopyJString(env, input, &input_text)) {
+    return nullptr;
+  }
   std::string tagged_text = processorITN->Tag(input_text);
   std::string normalized_text = processorITN->Verbalize(tagged_text);
 
diff --git a/runtime/processor/wetext_processor.cc b/runtime/processor/wetext_processor.cc
index 96077b1..fa5028a 100644
--- a/runtime/processor/wetext_processor.cc
+++ b/runtime/processor/wetext_processor.cc
@@ -14,13 +14,22 @@
 
 #include "processor/wetext_processor.h"
 
+#include <stdexcept>
+
 #include "utils/wetext_log.h"
 
 namespace wetext {
 Processor::Processor(const std::string& tagger_path,
                      const std::string& verbalizer_path) {
   tagger_.reset(StdVectorFst::Read(tagger_path));
+  if (tagger_ == nullptr) {
+    throw std::runtime_error("Failed to load tagger FST: " + tagger_path);
+  }
   verbalizer_.reset(StdVectorFst::Read(verbalizer_path));
+  if (verbalizer_ == nullptr) {
+    throw std::runtime_error("Failed to load verbalizer FST: " +
+                             verbalizer_path);
+  }
   compiler_ = std::make_shared<StringCompiler<StdArc>>();
   printer_ = std::make_shared<StringPrinter<StdArc>>();
 
diff --git a/runtime/test/processor_test.cc b/runtime/test/processor_test.cc
index 16e91bc..5f7ab22 100644
--- a/runtime/test/processor_test.cc
+++ b/runtime/test/processor_test.cc
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <stdexcept>
 #include <string>
 #include <vector>
 
@@ -72,6 +73,13 @@ TEST_P(ProcessorTest, NormalizeTest) {
   EXPECT_EQ(processor->Normalize(written), spoken);
 }
 
+TEST(ProcessorLoadTest, ThrowsWhenFstFilesCannotBeLoaded) {
+  EXPECT_THROW(
+      wetext::Processor("/tmp/zh_tn_missing_tagger.fst",
+                        "/tmp/zh_tn_missing_verbalizer.fst"),
+      std::runtime_error);
+}
+
 std::vector<std::pair<std::string, std::string>> test_cases =
     ParseTestCase(WETEXT_TN_DIR "/chinese/test/data/normalizer.txt");
 INSTANTIATE_TEST_SUITE_P(NormalizeTest, ProcessorTest,