VectorFst: Add sanity checks for reading the start state and the next state for the rest of the states.

The check is disabled for secondary component FST in `MergeFst`. The reason for this is because the secondary FST in MergeFst is allowed to be incomplete pointing to the states in the primary FST.

PiperOrigin-RevId: 910904063

Author: agutkin

openfst/lib/fst.cc

@@ -117,7 +117,8 @@ FstReadOptions::FstReadOptions(const absl::string_view source,
       isymbols(isymbols),
       osymbols(osymbols),
       read_isymbols(true),
-      read_osymbols(true) {
+      read_osymbols(true),
+      verify(true) {
   mode = ReadMode(absl::GetFlag(FLAGS_fst_read_mode));
 }
 
@@ -140,7 +141,8 @@ std::string FstReadOptions::DebugString() const {
       "\" read_osymbols: \"", (read_osymbols ? "true" : "false"),
       "\" header: \"", (header ? "set" : "null"), "\" isymbols: \"",
       (isymbols ? "set" : "null"), "\" osymbols: \"",
-      (osymbols ? "set" : "null"), "\"");
+      (osymbols ? "set" : "null"), "\" verify: \"", (verify ? "true" : "false"),
+      "\"");
 }
 
 }  // namespace fst

openfst/lib/fst.h

@@ -85,6 +85,9 @@ struct FstReadOptions {
   FileReadMode mode;            // Read or map files (advisory, if possible)
   bool read_isymbols;           // Read isymbols, if any (default: true).
   bool read_osymbols;           // Read osymbols, if any (default: true).
+  bool verify;                  // Perform FST type-specific light-weight sanity
+                                // check, e.g., check that the destinations of
+                                // all the arcs are valid.
 
   explicit FstReadOptions(absl::string_view source = "<unspecified>",
                           const FstHeader* absl_nullable header = nullptr,

openfst/lib/merge-fst.h

@@ -495,12 +495,42 @@ class MergeFstImpl : public FstImpl<A> {
     if (!state_map) return nullptr;
     FstReadOptions fopts(opts);
     fopts.header = nullptr;  // Component Fst headers were written out.
+
+    // Read primary Fst.
     std::unique_ptr<ExpandedFst<A> > primary_fst(
         ExpandedFst<A>::Read(strm, fopts));
     if (!primary_fst) return nullptr;
+
+    // Read secondary Fst. Disable the verification on read because this Fst is
+    // incomplete.
+    fopts.verify = false;
     std::unique_ptr<ExpandedFst<A> > secondary_fst(
         ExpandedFst<A>::Read(strm, fopts));
     if (!secondary_fst) return nullptr;
+
+    // Make sure the destination states in the secondary Fst are sane.
+    StateId max_next_state = std::numeric_limits<StateId>::min();
+    for (StateIterator<Fst<A>> siter(*secondary_fst); !siter.Done();
+         siter.Next()) {
+      const auto& state = siter.Value();
+      for (ArcIterator<Fst<A>> aiter(*secondary_fst, state); !aiter.Done();
+           aiter.Next()) {
+        const auto& arc = aiter.Value();
+        if (arc.nextstate == kNoStateId) {
+          LOG(ERROR) << "MergeFst::Read: Disallowed next state: " << kNoStateId;
+          return nullptr;
+        }
+        max_next_state = std::max(arc.nextstate, max_next_state);
+      }
+    }
+    const size_t max_states =
+        primary_fst->NumStates() + secondary_fst->NumStates();
+    if (max_next_state >= max_states) {
+      LOG(ERROR) << "MergeFst::Read: Next state " << max_next_state << "in "
+                 << "secondary FST is bigger than maximum possible number "
+                 << "of states " << max_states;
+      return nullptr;
+    }
     return new MergeFstImpl<A, M>(
         *impl, *primary_fst, *secondary_fst, hdr.NumStates(), *state_map);
   }

openfst/lib/vector-fst.h

@@ -27,6 +27,7 @@
 #include <ios>
 #include <iosfwd>
 #include <istream>
+#include <limits>
 #include <memory>
 #include <optional>
 #include <ostream>
@@ -483,6 +484,8 @@ VectorFstImpl<S>* VectorFstImpl<S>::Read(std::istream& strm,
   impl->BaseImpl::SetStart(hdr.Start());
   if (hdr.NumStates() != kNoStateId) impl->ReserveStates(hdr.NumStates());
   StateId state = 0;
+  StateId max_next_state = std::numeric_limits<StateId>::min();
+  StateId min_next_state = std::numeric_limits<StateId>::max();
   for (; hdr.NumStates() == kNoStateId || state < hdr.NumStates(); ++state) {
     Weight weight;
     if (!weight.Read(strm)) break;
@@ -506,13 +509,35 @@ VectorFstImpl<S>* VectorFstImpl<S>::Read(std::istream& strm,
         LOG(ERROR) << "VectorFst::Read: Read failed: " << opts.source;
         return nullptr;
       }
+      if (arc.nextstate == kNoStateId) {
+        LOG(ERROR) << "VectorFst::Read: Disallowed next state: " << kNoStateId;
+        return nullptr;
+      }
+      max_next_state = std::max(arc.nextstate, max_next_state);
+      min_next_state = std::min(arc.nextstate, min_next_state);
       impl->BaseImpl::AddArc(state, std::move(arc));
     }
   }
   if (hdr.NumStates() != kNoStateId && state != hdr.NumStates()) {
     LOG(ERROR) << "VectorFst::Read: Unexpected end of file: " << opts.source;
     return nullptr;
   }
+  // Sanity check for the start state.
+  if (impl->Start() != kNoStateId && impl->Start() >= state) {
+    LOG(ERROR) << "VectorFst::Read: start state " << impl->Start()
+               << " out of range [0, " << state << ")";
+    return nullptr;
+  }
+  // Sanity check on next states.
+  if (min_next_state < 0) {
+    LOG(ERROR) << "VectorFst::Read: Next state is negative: " << min_next_state;
+    return nullptr;
+  }
+  if (opts.verify && max_next_state >= state) {
+    LOG(ERROR) << "VectorFst::Read: Next state " << max_next_state
+               << " is larger than total number of states " << state;
+    return nullptr;
+  }
   return impl.release();
 }