add SocketMethods#kgio_addr!

This refreshes (or sets) the @kgio_addr ivar for sockets
that didn't go through kgio_accept or kgio_tryaccept.

2 files changed, 43 insertions, 2 deletions

diff --git a/ext/kgio/accept.c b/ext/kgio/accept.c
index f61e820..15e904f 100644
--- a/ext/kgio/accept.c
+++ b/ext/kgio/accept.c
@@ -187,7 +187,7 @@ retry:
         return client_io;
 }
 
-static void in_addr_set(VALUE io, struct sockaddr_storage *addr, socklen_t len)
+static VALUE in_addr_set(VALUE io, struct sockaddr_storage *addr, socklen_t len)
 {
         VALUE host;
         int host_len, rc;
@@ -210,7 +210,27 @@ static void in_addr_set(VALUE io, struct sockaddr_storage *addr, socklen_t len)
         if (rc != 0)
                 rb_raise(rb_eRuntimeError, "getnameinfo: %s", gai_strerror(rc));
         rb_str_set_len(host, strlen(host_ptr));
-        rb_ivar_set(io, iv_kgio_addr, host);
+        return rb_ivar_set(io, iv_kgio_addr, host);
+}
+
+/*
+ * call-seq:
+ *
+ *        io.kgio_addr! => refreshes the given sock address
+ */
+static VALUE addr_bang(VALUE io)
+{
+        int fd = my_fileno(io);
+        struct sockaddr_storage addr;
+        socklen_t len = sizeof(struct sockaddr_storage);
+
+        if (getpeername(fd, (struct sockaddr *)&addr, &len) != 0)
+                rb_sys_fail("getpeername");
+
+        if (addr.ss_family == AF_UNIX)
+                return rb_ivar_set(io, iv_kgio_addr, localhost);
+
+        return in_addr_set(io, &addr, len);
 }
 
 /*
@@ -424,6 +444,8 @@ void init_kgio_accept(void)
         cClientSocket = cKgio_Socket;
         mSocketMethods = rb_const_get(mKgio, rb_intern("SocketMethods"));
 
+        rb_define_method(mSocketMethods, "kgio_addr!", addr_bang, 0);
+
         rb_define_singleton_method(mKgio, "accept_cloexec?", get_cloexec, 0);
         rb_define_singleton_method(mKgio, "accept_cloexec=", set_cloexec, 1);
         rb_define_singleton_method(mKgio, "accept_nonblock?", get_nonblock, 0);
diff --git a/test/test_kgio_addr.rb b/test/test_kgio_addr.rb
new file mode 100644
index 0000000..8650f5e
--- /dev/null
+++ b/test/test_kgio_addr.rb
@@ -0,0 +1,19 @@
+# -*- encoding: binary -*-
+require 'test/unit'
+$-w = true
+require 'kgio'
+
+class TestKgioAddr < Test::Unit::TestCase
+  def test_tcp
+    addr = ENV["TEST_HOST"] || '127.0.0.1'
+    tcp = TCPServer.new(addr, 0)
+    port = tcp.addr[1]
+    client = Kgio::TCPSocket.new(addr, port)
+    accepted = tcp.accept
+    assert ! accepted.instance_eval { defined?(@kgio_addr) }
+    accepted.extend Kgio::SocketMethods
+    s = accepted.kgio_addr!
+    assert_equal addr, s
+    assert_equal addr, accepted.instance_variable_get(:@kgio_addr)
+  end
+end