return Patricia::Node objects when returning nodes

This prevents improper dereferencing and segfaults if the code
is misused.  Instead NoMethodError should be raised when it
would previously segfault.

3 files changed, 40 insertions, 18 deletions

diff --git a/ext/rpatricia/rpatricia.c b/ext/rpatricia/rpatricia.c
index a71ddac..54ac3ed 100644
--- a/ext/rpatricia/rpatricia.c
+++ b/ext/rpatricia/rpatricia.c
@@ -8,7 +8,7 @@
 #include <stdlib.h>
 #include "patricia.h"
 
-static VALUE cPatricia;
+static VALUE cPatricia, cNode;
 
 static void dummy(void) {}
 
@@ -30,6 +30,13 @@ p_node_mark (void *ptr)
   rb_gc_mark((VALUE)node->data);
 }
 
+static VALUE
+wrap_node(patricia_node_t *node)
+{
+  /* node will be freed when parent is freed */
+  return Data_Wrap_Struct(cNode, p_node_mark, 0, node);
+}
+
 static prefix_t *
 my_ascii2prefix(int family, VALUE str)
 {
@@ -70,7 +77,7 @@ p_add (int argc, VALUE *argv, VALUE self)
   PATRICIA_DATA_SET(node, user_data);
 
   /* node will be freed when parent is freed */
-  return Data_Wrap_Struct(cPatricia, p_node_mark, 0, node);
+  return wrap_node(node);
 }
 
 static VALUE
@@ -105,11 +112,7 @@ p_match (VALUE self, VALUE r_key)
   node = patricia_search_best(tree, prefix);
   Deref_Prefix (prefix);
 
-  if (node)
-    return Data_Wrap_Struct(cPatricia, p_node_mark, 0, node);
-  else 
-    return Qfalse;
-
+  return node ? wrap_node(node) : Qfalse;
 }
 
 static VALUE
@@ -124,10 +127,7 @@ p_match_exact (VALUE self, VALUE r_key)
   node = patricia_search_exact(tree, prefix);
   Deref_Prefix (prefix);
 
-  if (node)
-    return Data_Wrap_Struct(cPatricia, p_node_mark, 0, node);
-  else 
-    return Qfalse;
+  return node ? wrap_node(node) : Qfalse;
 }
 
 static VALUE
@@ -240,6 +240,7 @@ void
 Init_rpatricia (void)
 {
   cPatricia = rb_define_class("Patricia", rb_cObject);
+  cNode = rb_define_class_under(cPatricia, "Node", rb_cObject);
 
   /* create new Patricia object */
   rb_define_singleton_method(cPatricia, "new", p_new, 0);
@@ -270,11 +271,11 @@ Init_rpatricia (void)
   rb_define_method(cPatricia, "clear", p_destroy, 0); 
 
   /*---------- methods to node ----------*/
-  rb_define_method(cPatricia, "data", p_data, 0);
-  rb_define_method(cPatricia, "show_data", p_data, 0);
-  rb_define_method(cPatricia, "network", p_network, 0);
-  rb_define_method(cPatricia, "prefix", p_prefix, 0);
-  rb_define_method(cPatricia, "prefixlen", p_prefixlen, 0);
+  rb_define_method(cNode, "data", p_data, 0);
+  rb_define_method(cNode, "show_data", p_data, 0);
+  rb_define_method(cNode, "network", p_network, 0);
+  rb_define_method(cNode, "prefix", p_prefix, 0);
+  rb_define_method(cNode, "prefixlen", p_prefixlen, 0);
   //  rb_define_method(cPatricia, "family", p_family, 0); 
 
 }
diff --git a/test/test_gc.rb b/test/test_gc.rb
index d567696..5829deb 100644
--- a/test/test_gc.rb
+++ b/test/test_gc.rb
@@ -17,7 +17,7 @@ class TestGc < Test::Unit::TestCase
 
   def test_gc
     assert_nothing_raised do
-      10_000_000.times do
+      5_000_000.times do
         t = Patricia.new
         t.add('10.0.0.0/8', {})
         t.add('127.0.0.0/24', "home sweet home")
@@ -25,7 +25,7 @@ class TestGc < Test::Unit::TestCase
     end
 
     # ensure what we created originally didn't get GC-ed'
-    100.times do
+    5_000_000.times do
       assert_equal [], @arrays.match_best('127.0.0.1').data
       assert_equal "localhost", @strings.match_best('127.0.0.1').data
     end
diff --git a/test/test_old_segfaults.rb b/test/test_old_segfaults.rb
new file mode 100644
index 0000000..0814150
--- /dev/null
+++ b/test/test_old_segfaults.rb
@@ -0,0 +1,21 @@
+require 'test/unit'
+require 'rpatricia'
+
+# things that used to segfault before the introduction of
+# Patricia::Node will now raise NoMethodError
+class TestOldSegfaults < Test::Unit::TestCase
+
+  def test_used_to_segfault
+    assert_raises(NoMethodError) { Patricia.new.data }
+  end
+
+  def test_node_method
+    t = Patricia.new
+    t.add('10.0.0.0/8', 'big_lan')
+    if matched = t.match_best('10.1.2.3')
+      assert_kind_of Patricia::Node, matched
+      assert_equal 'big_lan', matched.data
+      assert_raises(NoMethodError) { matched.destroy }
+    end
+  end
+end