From b3f34b06ac068441235972a2b087a2dcee137f83 Mon Sep 17 00:00:00 2001
From: Vadim Kurland <vadim@slot.vk.crocodile.org>
Date: Thu, 27 Jan 2011 11:58:02 -0800
Subject: [PATCH] * Rule.cpp (removeRef): fixes #1997 "add removeRef and addRef
 methods to class NATRule". Now undo and redo correctly remove and restore
 references to NAT rule sets in NAT rules with action Branch.

---
 doc/ChangeLog                                 |  5 +++
 .../src/fwbuilder/FWObjectDatabase_search.cpp | 11 +++++++
 src/libfwbuilder/src/fwbuilder/Rule.cpp       | 31 +++++++++++++++++++
 src/libfwbuilder/src/fwbuilder/Rule.h         | 12 +++++++
 4 files changed, 59 insertions(+)

diff --git a/doc/ChangeLog b/doc/ChangeLog
index 9ac3cb8a4..8df81150d 100644
--- a/doc/ChangeLog
+++ b/doc/ChangeLog
@@ -1,5 +1,10 @@
 2011-01-27  vadim  <vadim@netcitadel.com>
 
+	* Rule.cpp (removeRef): fixes #1997 "add removeRef and addRef
+	methods to class NATRule". Now undo and redo correctly remove and
+	restore references to NAT rule sets in NAT rules with action
+	Branch.
+
 	* Rule.cpp (addRef): fixes #1991 "Undo does not restore object as
 	a parameter of policy rule action Branch or Tag after it was
 	deleted deleted". Now Undo restores references to rule sets and tag
diff --git a/src/libfwbuilder/src/fwbuilder/FWObjectDatabase_search.cpp b/src/libfwbuilder/src/fwbuilder/FWObjectDatabase_search.cpp
index 958e98ed4..64b58ff33 100644
--- a/src/libfwbuilder/src/fwbuilder/FWObjectDatabase_search.cpp
+++ b/src/libfwbuilder/src/fwbuilder/FWObjectDatabase_search.cpp
@@ -222,6 +222,17 @@ bool FWObjectDatabase::_findWhereObjectIsUsed(FWObject *o,
         }
     }
 
+    NATRule *nat_rule = NATRule::cast(p);
+    if (nat_rule && nat_rule->getAction() == NATRule::Branch)
+    {
+        FWObject *ruleset = nat_rule->getBranch();
+        if (o==ruleset)
+        {
+            resset.insert(p);
+            res = true;
+        }
+    }
+
     if (Firewall::isA(o) && Cluster::isA(p))
     {
         if (Cluster::cast(p)->hasMember(Firewall::cast(o)))
diff --git a/src/libfwbuilder/src/fwbuilder/Rule.cpp b/src/libfwbuilder/src/fwbuilder/Rule.cpp
index 79d0d6780..478d37c99 100644
--- a/src/libfwbuilder/src/fwbuilder/Rule.cpp
+++ b/src/libfwbuilder/src/fwbuilder/Rule.cpp
@@ -597,6 +597,37 @@ void NATRule::init(FWObjectDatabase *root)
     }
 }
 
+/**
+ * Add reference to given object. In case of NATRule this only
+ * makes sense in terms of adding reference to this object as an
+ * argument for action Branch.
+ */
+void NATRule::addRef(FWObject *obj)
+{
+    if (RuleSet::cast(obj))
+    {
+        setBranch(RuleSet::cast(obj));
+    }
+}
+
+/**
+ * Removes reference to given object among children of 'this'. In case
+ * of NATRule we should also clear reference to it if action is
+ * Branch. Caveat: clear reference to it even if action is not branch
+ * right now but was in the past and reference got stuck in options.
+ */
+void NATRule::removeRef(FWObject *obj)
+{
+    if (RuleSet::cast(obj))
+    {
+        string branch_id = FWObjectDatabase::getStringId(obj->getId());
+        string rule_branch_id = getOptionsObject()->getStr("branch_id");
+        if (branch_id == rule_branch_id)
+            getOptionsObject()->setStr("branch_id", "");
+    }
+    FWObject::removeRef(obj);
+}
+
 RuleElementOSrc*  NATRule::getOSrc()
 {
     if (osrc_re) return osrc_re;
diff --git a/src/libfwbuilder/src/fwbuilder/Rule.h b/src/libfwbuilder/src/fwbuilder/Rule.h
index 5bf2ce1f9..981419a31 100644
--- a/src/libfwbuilder/src/fwbuilder/Rule.h
+++ b/src/libfwbuilder/src/fwbuilder/Rule.h
@@ -366,6 +366,18 @@ public:
 
     virtual bool isEmpty();
 
+    /**
+     * Removes reference to given object among
+     * children of 'this'.
+     */
+    virtual void removeRef(FWObject *obj);
+    
+    /**
+     * Add reference to given object to 'this'. In case of a PolicyRule,
+     * change action and add parameter if @obj is TagService or RuleSet.
+     */
+    virtual void addRef(FWObject *obj);
+        
     libfwbuilder::RuleElementOSrc* getOSrc();
     libfwbuilder::RuleElementODst* getODst();
     libfwbuilder::RuleElementOSrv* getOSrv();