通过 netlink 的 NETLINK_INET_DIAG 协议获取当前系统所有 TCP 连接信息

阅读《2024 中国开源开发者报告》赢大奖，扫码申请享特权

　　据说 ss 比 netstat 快，想研究一下原理，从这边看到代码：http://bbs.chinaunix.net/thread-3766684-1-1.html，但很不靠谱……整理了一下，供大家参考。

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <fcntl.h>
#include <errno.h>
#include <string.h>
#include <asm/types.h>
#include <sys/socket.h>
#include <linux/netlink.h>
#include <linux/inet_diag.h>
#include <netinet/tcp.h>

int main(int argc, char **argv)
{
    int fd;
    struct sockaddr_nl src_addr, dest_addr;
    struct
    {
        struct nlmsghdr nlh;
        struct inet_diag_req r;
    } req;
    struct inet_diag_msg *pkg;
    struct msghdr msg;
    char buf[8192];
    char src_ip[40];
    char dest_ip[40];
    struct iovec iov;

    if ((fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_INET_DIAG)) < 0)
        return -1;

    int ret;
    ret = fcntl(fd, F_SETFL, O_NONBLOCK);
    if (ret < 0) {
        fprintf(stderr, "Can't set socket flags");
        close(fd);
        return -1;
    }
    //src addr
    memset(&src_addr, 0, sizeof(struct sockaddr_nl));
    src_addr.nl_family = AF_NETLINK;
    src_addr.nl_pid = getpid();
    src_addr.nl_groups = 0;

    if (bind(fd, (struct sockaddr *)&src_addr, sizeof(struct sockaddr_nl)) < 0) {
        fprintf(stderr, "bind socket error %s\n", strerror(errno));
    }

    memset(&req, 0, sizeof(req));
    req.nlh.nlmsg_len = sizeof(req);
    req.nlh.nlmsg_type = TCPDIAG_GETSOCK;
    req.nlh.nlmsg_flags = NLM_F_MATCH | NLM_F_REQUEST | NLM_F_ROOT;
    // req.nlh.nlmsg_flags = NLM_F_REQUEST ;
    req.nlh.nlmsg_pid = 0;

    memset(&req.r, 0, sizeof(req.r));
    req.r.idiag_family = AF_INET;
    req.r.idiag_states = ((1 << TCP_CLOSING + 1) - 1); //states to dump

    //send msg to kernel
    iov.iov_base = &req;
    iov.iov_len = sizeof(req);

    //dest addr
    memset(&dest_addr, 0, sizeof(struct sockaddr_nl));
    dest_addr.nl_family = AF_NETLINK;
    dest_addr.nl_pid = 0;
    dest_addr.nl_groups = 0;

    msg.msg_name = (void *)&dest_addr;
    msg.msg_namelen = sizeof(dest_addr);
    msg.msg_iov = &iov;
    msg.msg_iovlen = 1;

    if (sendmsg(fd, &msg, 0) < 0) {
        printf("%s\n", strerror(errno));
        return -1;
    }
    //recv msg from kernel
    iov.iov_base = buf;
    iov.iov_len = sizeof(buf);

    while (1) {
        //printf("while1\n");
        int status;
        struct nlmsghdr *h;

        msg = (struct msghdr)
        {
            (void *)&dest_addr, sizeof(struct sockaddr_nl),
                &iov, 1, NULL, 0, 0
        };

        //length of recv data
        status = recvmsg(fd, &msg, 0);
        //status = recv(fd, buf, sizeof(buf), 0);
        printf("status = %d\n", status);
        if (status < 0) {
            if (errno == EINTR) {
                continue;
            }
            printf("errno = %d\n", errno);
            continue;
        }
        if (status == 0) {
            close(fd);
            printf("EOF\n");
            return 0;
        }

        h = (struct nlmsghdr *)buf;

        while (NLMSG_OK(h, status)) {
            //printf("while2\n");
            if (h->nlmsg_type == NLMSG_DONE) {
                close(fd);
                printf("NLMSG_DONE\n");
                return 0;
            }

            if (h->nlmsg_type == NLMSG_ERROR) {
                struct nlmsgerr *err;
                err = (struct nlmsgerr*)NLMSG_DATA(h);
                fprintf(stderr, "%d Error %d:%s\n", __LINE__, -(err->error), strerror(-(err->error)));
                close(fd);
                printf("NLMSG_ERROR\n");
                return 0;
            }

            pkg = (struct inet_diag_msg *)NLMSG_DATA(h);
            memset(src_ip, 0, sizeof(src_ip));
            memset(dest_ip, 0, sizeof(dest_ip));
            inet_ntop(pkg->idiag_family, pkg->id.idiag_src, src_ip, sizeof(src_ip));
            inet_ntop(pkg->idiag_family, pkg->id.idiag_dst, dest_ip, sizeof(dest_ip));
            printf("%-8s %4d %40s:%-6hu %40s:%-6hu\n", pkg->idiag_family == AF_INET ? "AF_INET" : "AF_INET6", pkg->idiag_state
                , src_ip, ntohs(pkg->id.idiag_sport), dest_ip, ntohs(pkg->id.idiag_dport));
            // get_tcp_state(pkg->idiag_state);
            h = NLMSG_NEXT(h, status);
            //printf("status = %d\n\n", status);
        }//while
    }//while
    close(fd);
    return 0;
}