1 files changed, 29 insertions, 12 deletions

diff --git a/scripts/vyatta-conntrack-timeouts.pl b/scripts/vyatta-conntrack-timeouts.pl
index 61830ab..4ea3a9f 100644
--- a/scripts/vyatta-conntrack-timeouts.pl
+++ b/scripts/vyatta-conntrack-timeouts.pl
@@ -73,6 +73,7 @@ sub remove_timeout_policy {
     run_cmd($iptables_cmd2);
     if ($? >> 8) {
       print "$CTERROR failed to run $iptables_cmd2\n";    
+      #dont exit, try to clean as much. 
     }
     run_cmd($iptables_cmd1);
     if ($? >> 8) {
@@ -95,17 +96,43 @@ sub apply_timeout_policy {
     run_cmd($nfct_timeout_cmd);
     if ($? >> 8) {
       print "$CTERROR failed to run $nfct_timeout_cmd\n";    
+      exit 1; 
     }
     run_cmd($iptables_cmd1);
     if ($? >> 8) {
+      #cleanup the policy before exit. 
+      run_cmd("nfct-timeout remove $timeout_policy");   
       print "$CTERROR failed to run $iptables_cmd1\n";    
+      exit 1; 
     }
     run_cmd($iptables_cmd2);
     if ($? >> 8) {
+      run_cmd("nfct-timeout remove $timeout_policy");   
+      run_cmd("iptables -D PREROUTING -t raw $rule_string -j CT --timeout $tokens[0]");   
       print "$CTERROR failed to run $iptables_cmd2\n";    
+      exit 1;
     }
 }
 
+sub handle_rule_creation {
+  my ($rule) = @_;
+  my $node = new Vyatta::Conntrack::RuleCT;
+  my ($rule_string, $timeout_policy);
+  $node->setup("system conntrack timeout custom rule $rule");
+  $rule_string = $node->rule();
+  $timeout_policy = $node->get_policy_command(); #nfct-timeout command string
+  apply_timeout_policy($rule_string, $timeout_policy);
+}
+
+sub handle_rule_deletion {
+  my ($rule) = @_;
+  my $node = new Vyatta::Conntrack::RuleCT;
+  my ($rule_string, $timeout_policy);
+  $node->setupOrig("system conntrack timeout custom rule $rule");
+  $rule_string = $node->rule();
+  $timeout_policy = $node->get_policy_command(); #nfct-timeout command string
+  remove_timeout_policy($rule_string, $timeout_policy);
+}
 
 sub update_config {
   my $config = new Vyatta::Config;
@@ -117,22 +144,12 @@ sub update_config {
   foreach my $rule (sort keys %rules) { 
     if ("$rules{$rule}" eq 'static') {
     } elsif ("$rules{$rule}" eq 'added') {      
-      my $node = new Vyatta::Conntrack::RuleCT;
-      my ($rule_string, $timeout_policy);
-      $node->setup("system conntrack timeout custom rule $rule");
-      $rule_string = $node->rule();
-      $timeout_policy = $node->get_policy_command(); #nfct-timeout command string
-      apply_timeout_policy($rule_string, $timeout_policy);
+        handle_rule_creation($rule);
     } elsif ("$rules{$rule}" eq 'changed') {
       my $node = new Vyatta::Conntrack::RuleCT;
       $node->setup("system conntrack timeout custom rule $rule");
     } elsif ("$rules{$rule}" eq 'deleted') {
-      my $node = new Vyatta::Conntrack::RuleCT;
-      my ($rule_string, $timeout_policy);
-      $node->setupOrig("system conntrack timeout custom rule $rule");
-      $rule_string = $node->rule();
-      $timeout_policy = $node->get_policy_command(); #nfct-timeout command string
-      remove_timeout_policy($rule_string, $timeout_policy);
+        handle_rule_deletion($rule);
     }  
   }
 }