fix: json.set recursive processing crash (#5040)

fixed: https://github.com/dragonflydb/dragonfly/issues/5038

src/core/json/detail/jsoncons_dfs.cc

@@ -82,6 +82,9 @@ Dfs Dfs::Mutate(absl::Span<const PathSegment> path, const MutateCallback& callba
     return dfs;
   }
 
+  // Use vector to maintain order
+  std::vector<JsonType*> nodes_to_mutate;
+
   using Item = detail::JsonconsDfsItem<false>;
   vector<Item> stack;
   stack.emplace_back(json);
@@ -103,14 +106,32 @@ Dfs Dfs::Mutate(absl::Span<const PathSegment> path, const MutateCallback& callba
         if (next_seg_id + 1 < path.size()) {
           stack.emplace_back(next, next_seg_id);
         } else {
-          dfs.MutateStep(path[next_seg_id], callback, next);
+          // Terminal step: collect node for mutation if not seen before
+          nodes_to_mutate.push_back(next);
         }
       }
     } else {
+      // If Advance failed (e.g., MISMATCH or OUT_OF_BOUNDS), check if the current node
+      // itself is a terminal target for mutation due to a previous DESCENT segment.
+      // This handles cases like $.a..b where 'a' itself might match the 'b' after descent.
+      if (!res && segment_index > 0 && path[segment_index - 1].type() == SegmentType::DESCENT &&
+          stack.back().get_segment_step() == 0) {
+        if (segment_index + 1 == path.size()) {
+          // Attempt to apply the final mutation step directly to the current node.
+          // We apply it directly here because this node won't be added to the stack again.
+          dfs.MutateStep(path_segment, callback, stack.back().obj_ptr());
+        }
+      }
       stack.pop_back();
     }
   } while (!stack.empty());
 
+  // Apply mutations after DFS traversal is complete
+  const PathSegment& terminal_segment = path.back();
+  for (JsonType* node : nodes_to_mutate) {
+    dfs.MutateStep(terminal_segment, callback, node);
+  }
+
   return dfs;
 }
 

src/core/json/detail/jsoncons_dfs.h

@@ -43,6 +43,14 @@ template <bool IsConst> class JsonconsDfsItem {
     return depth_state_.second;
   }
 
+  Ptr obj_ptr() const {
+    return depth_state_.first;
+  }
+
+  unsigned get_segment_step() const {
+    return segment_step_;
+  }
+
  private:
   static bool ShouldIterateAll(SegmentType type) {
     return type == SegmentType::WILDCARD || type == SegmentType::DESCENT;
@@ -128,8 +136,15 @@ class Dfs {
   }
 
   bool Mutate(const MutateCallback& callback, std::optional<std::string_view> key, JsonType* node) {
-    ++matches_;
-    return callback(key, node);
+    // matches_ is incremented in MutateStep after successful deletion/mutation by callback.
+    bool deleted = callback(key, node);
+    // We only increment matches_ here if the callback indicates deletion,
+    // as MutateStep handles incrementing based on its own logic.
+    // This might need adjustment if non-deleting mutations should also count universally.
+    if (deleted) {
+      matches_++;
+    }
+    return deleted;  // Return true if node should be deleted
   }
 
   unsigned matches_ = 0;

src/core/json/jsonpath_test.cc

@@ -521,6 +521,30 @@ TYPED_TEST(JsonPathTest, Mutate) {
   }
 }
 
+TYPED_TEST(JsonPathTest, MutateRecursiveDescentKey) {
+  ASSERT_EQ(0, this->Parse("$..value"));
+  Path path = this->driver_.TakePath();
+
+  JsonType json = ValidJson<JsonType>(R"({"data":{"value":10,"subdata":{"value":20}}})");
+  JsonType replacement = ValidJson<JsonType>(R"({"value": 30})");
+
+  auto cb = [&](optional<string_view> key, JsonType* val) {
+    if (key && key.value() == "value" && (val->is_int64() || val->is_double())) {
+      *val = replacement;
+      return false;
+    }
+    return false;
+  };
+
+  unsigned reported_matches = MutatePath(path, cb, &json);
+
+  JsonType expected =
+      ValidJson<JsonType>(R"({"data":{"subdata":{"value":{"value":30}},"value":{"value":30}}})");
+
+  EXPECT_EQ(expected, json);
+  EXPECT_EQ(0, reported_matches);
+}
+
 TYPED_TEST(JsonPathTest, SubRange) {
   TypeParam json = ValidJson<TypeParam>(R"({"arr": [1, 2, 3, 4, 5]})");
   ASSERT_EQ(0, this->Parse("$.arr[1:2]"));

src/core/json/path.cc

@@ -288,11 +288,16 @@ unsigned MutatePath(const Path& path, MutateCallback callback, FlatJson json,
                     flexbuffers::Builder* fbb) {
   JsonType mut_json = FromFlat(json);
   unsigned res = MutatePath(path, std::move(callback), &mut_json);
-  if (res) {
-    FromJsonType(mut_json, fbb);
-    fbb->Finish();
-  }
 
+  // Populate the output builder 'fbb' with the resulting JSON state
+  // (mutated or original if res == 0) and finalize it.
+  // The builder MUST be finished before returning so that the caller
+  // can safely access the resulting flatbuffer data (e.g., via GetBuffer()).
+  // Skipping Finish() would leave the builder in an invalid, unusable state.
+  FromJsonType(mut_json, fbb);  // Always convert (changed or not) JSON
+  fbb->Finish();                // Always finish the builder
+
+  // Return the number of actual mutations that occurred.
   return res;
 }