3 files changed, 27 insertions, 34 deletions

diff --git a/libnl_2/handlers.c b/libnl_2/handlers.c
index 48dcab4a8..ec8d51270 100644
--- a/libnl_2/handlers.c
+++ b/libnl_2/handlers.c
@@ -39,14 +39,16 @@ fail:
 struct nl_cb *nl_cb_clone(struct nl_cb *orig)
 {
         struct nl_cb *new_cb;
+        int new_refcnt;
 
         new_cb = nl_cb_alloc(NL_CB_DEFAULT);
         if (new_cb == NULL)
                 goto fail;
 
-        /* Copy original and set refcount to 1 */
+        /* Preserve reference count and copy original */
+        new_refcnt = new_cb->cb_refcnt;
         memcpy(new_cb, orig, sizeof(*orig));
-        new_cb->cb_refcnt = 1;
+        new_cb->cb_refcnt = new_refcnt;
 
         return new_cb;
 fail:
@@ -82,9 +84,9 @@ struct nl_cb *nl_cb_get(struct nl_cb *cb)
 
 void nl_cb_put(struct nl_cb *cb)
 {
-        if (!cb)
-                return;
         cb->cb_refcnt--;
         if (cb->cb_refcnt <= 0)
                 free(cb);
+
 }
+
diff --git a/libnl_2/netlink.c b/libnl_2/netlink.c
index ee3d600f7..cc2f88e68 100644
--- a/libnl_2/netlink.c
+++ b/libnl_2/netlink.c
@@ -59,14 +59,15 @@ int nl_recv(struct nl_sock *sk, struct sockaddr_nl *nla, \
 {
         int rc = -1;
         int sk_flags;
-        int RECV_BUF_SIZE = getpagesize();
+        int RECV_BUF_SIZE;
         int errsv;
         struct iovec recvmsg_iov;
         struct msghdr msg;
 
         /* Allocate buffer */
+        RECV_BUF_SIZE = getpagesize();
         *buf = (unsigned char *) malloc(RECV_BUF_SIZE);
-        if (!(*buf)) {
+        if (!buf) {
                 rc = -ENOMEM;
                 goto fail;
         }
@@ -90,11 +91,8 @@ int nl_recv(struct nl_sock *sk, struct sockaddr_nl *nla, \
         errsv = errno;
         fcntl(sk->s_fd, F_SETFL, sk_flags);
 
-        if (rc < 0) {
+        if (rc < 0)
                 rc = -errsv;
-                free(*buf);
-                *buf = NULL;
-        }
 
 fail:
         return rc;
@@ -110,6 +108,7 @@ int nl_recvmsgs(struct nl_sock *sk, struct nl_cb *cb)
         int rc, cb_rc = NL_OK, done = 0;
 
         do {
+
                 unsigned char *buf;
                 int i, rem, flags;
                 struct nlmsghdr *nlh;
@@ -128,7 +127,7 @@ int nl_recvmsgs(struct nl_sock *sk, struct nl_cb *cb)
 
                         /* Check for callbacks */
 
-                        msg = (struct nl_msg *) malloc(sizeof(struct nl_msg));
+                        msg = (struct nl_msg *)malloc(sizeof(struct nl_msg));
                         memset(msg, 0, sizeof(*msg));
                         msg->nm_nlh = nlh;
 
@@ -188,6 +187,7 @@ int nl_recvmsgs(struct nl_sock *sk, struct nl_cb *cb)
                         if (done)
                                 break;
                 }
+
                 free(buf);
                 buf = NULL;
 
@@ -197,7 +197,7 @@ int nl_recvmsgs(struct nl_sock *sk, struct nl_cb *cb)
 
 success:
 fail:
-        return rc;
+        return  rc;
 }
 
 /* Send raw data over netlink socket */
diff --git a/libnl_2/socket.c b/libnl_2/socket.c
index d906cac27..ce54f19ba 100644
--- a/libnl_2/socket.c
+++ b/libnl_2/socket.c
@@ -31,7 +31,7 @@ int nl_socket_add_membership(struct nl_sock *sk, int group)
 }
 
 /* Allocate new netlink socket. */
-static struct nl_sock *_nl_socket_alloc(void)
+struct nl_sock *nl_socket_alloc(void)
 {
         struct nl_sock *sk;
         struct timeval tv;
@@ -39,13 +39,13 @@ static struct nl_sock *_nl_socket_alloc(void)
 
         sk = (struct nl_sock *) malloc(sizeof(struct nl_sock));
         if (!sk)
-                return NULL;
+                goto fail;
         memset(sk, 0, sizeof(*sk));
 
         /* Get current time */
 
         if (gettimeofday(&tv, NULL))
-                goto fail;
+                return NULL;
         else
                 sk->s_seq_next = (int) tv.tv_sec;
 
@@ -59,36 +59,24 @@ static struct nl_sock *_nl_socket_alloc(void)
         sk->s_peer.nl_pid = 0; /* Kernel */
         sk->s_peer.nl_groups = 0; /* No groups */
 
-        return sk;
-fail:
-        free(sk);
-        return NULL;
-}
-
-/* Allocate new netlink socket. */
-struct nl_sock *nl_socket_alloc(void)
-{
-        struct nl_sock *sk = _nl_socket_alloc();
-        struct nl_cb *cb;
-
-        if (!sk)
-                return NULL;
-
-        cb = nl_cb_alloc(NL_CB_DEFAULT);
+        cb = (struct nl_cb *) malloc(sizeof(struct nl_cb));
         if (!cb)
                 goto cb_fail;
-        sk->s_cb = cb;
+        memset(cb, 0, sizeof(*cb));
+        sk->s_cb = nl_cb_alloc(NL_CB_DEFAULT);
+
+
         return sk;
 cb_fail:
         free(sk);
+fail:
         return NULL;
 }
 
 /* Allocate new socket with custom callbacks. */
 struct nl_sock *nl_socket_alloc_cb(struct nl_cb *cb)
 {
-        struct nl_sock *sk = _nl_socket_alloc();
-
+        struct nl_sock *sk = nl_socket_alloc();
         if (!sk)
                 return NULL;
 
@@ -96,6 +84,7 @@ struct nl_sock *nl_socket_alloc_cb(struct nl_cb *cb)
         nl_cb_get(cb);
 
         return sk;
+
 }
 
 /* Free a netlink socket. */
@@ -127,3 +116,5 @@ int nl_socket_get_fd(struct nl_sock *sk)
 {
         return sk->s_fd;
 }
+
+