From 320649815fe9691e87f5dbd20b91895dcba00b9b Mon Sep 17 00:00:00 2001
From: Michael Santos <michael.santos@gmail.com>
Date: Fri, 17 Dec 2010 16:26:52 -0500
Subject: [PATCH] recvfrom/4: return socket address

Add recvfrom/4 to return the socket address of the packet. recvfrom/4
takes 2 new args: flags and the size of the socket structure which is
allocated by the NIF. If the size is not large enough, the socket address
will be silently truncated.
---
 c_src/procket.c | 33 +++++++++++++++++++++++++++------
 src/procket.erl | 12 ++++++++++--
 2 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/c_src/procket.c b/c_src/procket.c
index df4454d..75d60c7 100644
--- a/c_src/procket.c
+++ b/c_src/procket.c
@@ -193,30 +193,49 @@ nif_close(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
 
 
 /* 0: socket, 1: length */
+/* 0: socket, 1: length, 2: flags, 3: struct sockaddr length */
     static ERL_NIF_TERM
 nif_recvfrom(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
 {
     int sockfd = -1;
     int len = 0;
-    ssize_t bufsz = 0;
+    int salen = 0;
+    int flags = 0;
+
     ErlNifBinary buf;
+    ErlNifBinary sa;
+    ssize_t bufsz = 0;
+    socklen_t sasz = 0;
+
 
     if (!enif_get_int(env, argv[0], &sockfd))
         return enif_make_badarg(env);
     if (!enif_get_int(env, argv[1], &len))
         return enif_make_badarg(env);
+    if (!enif_get_int(env, argv[2], &flags))
+        return enif_make_badarg(env);
+    if (!enif_get_int(env, argv[3], &salen))
+        return enif_make_badarg(env);
 
     if (!enif_alloc_binary(len, &buf))
         return error_tuple(env, ENOMEM);
 
-    if ( (bufsz = recvfrom(sockfd, buf.data, buf.size, 0, NULL, NULL)) == -1) {
+    if (!enif_alloc_binary(salen, &sa))
+        return error_tuple(env, ENOMEM);
+
+    sasz = salen;
+    if ( (bufsz = recvfrom(sockfd, buf.data, buf.size, flags,
+        (sasz == 0 ? NULL : (struct sockaddr *)sa.data),
+        &sasz)) == -1) {
         switch (errno) {
             case EAGAIN:
             case EINTR:
                 enif_release_binary(&buf);
+                enif_release_binary(&sa);
                 return atom_eagain;
             default:
                 enif_release_binary(&buf);
+                enif_release_binary(&sa);
                 return error_tuple(env, errno);
         }
     }
@@ -224,9 +243,11 @@ nif_recvfrom(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
     if (bufsz != buf.size)
         enif_realloc_binary(&buf, bufsz);
 
-    return enif_make_tuple(env, 2,
-            atom_ok,
-            enif_make_binary(env, &buf));
+    if (salen != sasz)
+        enif_realloc_binary(&sa, sasz);
+
+    return enif_make_tuple(env, 3, atom_ok, enif_make_binary(env, &buf),
+             enif_make_binary(env, &sa));
 }
 
 
@@ -383,9 +404,9 @@ static ErlNifFunc nif_funcs[] = {
     {"bind", 2, nif_bind},
     {"connect", 2, nif_connect},
     {"listen", 2, nif_listen},
-    {"recvfrom", 2, nif_recvfrom},
     {"ioctl", 3, nif_ioctl},
     {"socket", 3, nif_socket},
+    {"recvfrom", 4, nif_recvfrom},
     {"sendto", 4, nif_sendto},
     {"setsockopt", 4, nif_setsockopt}
 };
diff --git a/src/procket.erl b/src/procket.erl
index 9a3be10..6434942 100644
--- a/src/procket.erl
+++ b/src/procket.erl
@@ -36,7 +36,8 @@
         socket/3, listen/2,connect/2,
         accept/1,accept/2,
         fdopen/1,fdrecv/1,close/1,close/2,
-        recvfrom/2,sendto/4,bind/2,
+        recv/2,recvfrom/2,recvfrom/4,
+        sendto/4,bind/2,
         ioctl/3,setsockopt/4
     ]).
 -export([make_args/2,progname/0]).
@@ -74,7 +75,14 @@ connect(_,_) ->
 listen(_,_) ->
     erlang:error(not_implemented).
 
-recvfrom(_,_) ->
+recv(Socket,Size) ->
+    recvfrom(Socket,Size).
+recvfrom(Socket,Size) ->
+    case recvfrom(Socket, Size, 0, 0) of
+        {ok, Buf, <<>>} -> {ok, Buf}; 
+        Error -> Error
+    end.
+recvfrom(_,_,_,_) ->
     erlang:error(not_implemented).
 
 socket(_,_,_) ->