Linux下poll函数实现多用户聊天

Client：

#define _GNU_SOURCE 1
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include<arpa/inet.h>
#include <assert.h>
#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <poll.h>
#include <fcntl.h>

#define BUFFER_SIZE 64

int main()
{
    const char *ip="127.0.0.1";
    int port=12345;
    struct sockaddr_in server_address;
    bzero(&server_address,sizeof(server_address));
    server_address.sin_family=AF_INET;
    inet_pton(AF_INET,ip,&server_address.sin_addr);
    server_address.sin_port=htons(port);

    int sockfd=socket(PF_INET,SOCK_STREAM,0);
    assert(sockfd!=0);
    if(connect(sockfd,(struct sockaddr *)&server_address,sizeof(server_address))<0)
    {
        printf("connection failed\n");
        close(sockfd);
        return 1;
    }

    struct pollfd fds[2];
    fds[0].fd=0;
    fds[0].revents=0;
    fds[1].fd=sockfd;
    fds[1].events=POLLIN | POLLRDHUP;
    fds[1].revents=0;

    char read_buf[BUFFER_SIZE];
    int pipefd[2];
    int ret=pipe(pipefd);
    assert(ret!=-1);
    while(1)
    {
        ret=poll(fds,2,-1);
        if(ret<0)
        {
            printf("poll failure\n");
            break;
        }

        if(fds[1].revents&POLLRDHUP)
        {
            printf("server close the connection\n");
            break;
        }
        else if(fds[1].revents&POLLIN)
        {
            memset(read_buf,'\0',BUFFER_SIZE);
            recv(fds[1].fd,read_buf,BUFFER_SIZE-1,0);
            printf("%s\n",read_buf);
        }
        if(fds[0].revents&POLLIN)
        {
            ret=splice(0,NULL,pipefd[1],NULL,32768,SPLICE_F_MORE | SPLICE_F_MOVE);
            ret=splice(pipefd[0],NULL,sockfd,NULL,32768,SPLICE_F_MORE |SPLICE_F_MOVE);
        }
    }
    close(sockfd);
    return 0;
}

Server:

#define _GNU_SOURCE 1
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <assert.h>
#include <stdio.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <fcntl.h>
#include <stdlib.h>
#include <poll.h>

#define USER_LIMIT 5
#define BUFFER_SIZE 64
#define FD_LIMIT 65535

struct client_data
{
    struct sockaddr_in address;
    char *write_buf;
    char buf[BUFFER_SIZE];
};

int setnonblocking(int fd)
{
    int old_option=fcntl(fd,F_GETFL);
    int new_option=old_option | O_NONBLOCK;
    fcntl(fd,F_SETFL,new_option);
    return old_option;
}

int main()
{
    const char *ip="127.0.0.1";
    int port=12345;
    int ret=0;
    int i,j;
    struct sockaddr_in address;
    bzero(&address,sizeof(address));
    address.sin_family=AF_INET;
    inet_pton(AF_INET,ip,&address.sin_addr);
    address.sin_port=htons(port);
   // address.sin_addr.s_addr=INADDR_ANY;
    int listenfd=socket(PF_INET,SOCK_STREAM,0);
    assert(listenfd>=0);

    ret=bind(listenfd,(struct sockaddr *)&address,sizeof(address));
    assert(ret!=-1);
    ret=listen(listenfd,5);
    assert(ret!=-1);
    struct client_data* users=(struct client_data*) calloc(FD_LIMIT,sizeof(struct client_data));
    struct pollfd fds[USER_LIMIT+1];
    int user_counter=0;
    for(i=1;i<=USER_LIMIT;++i)
    {
        fds[i].fd=-1;
        fds[i].events=0;
    }

    fds[0].fd=listenfd;
    fds[0].events=POLLIN | POLLERR;
    fds[0].revents=0;

    while(1)
    {
        printf("{\n");
        ret=poll(fds,user_counter+1,-1);
        printf("}\n");
        if(ret<0)
        {
            printf("poll failure\n");
            break;
        }

        for(i=0;i<user_counter+1;++i)
        {

            if((fds[i].fd==listenfd)&&(fds[i].revents&POLLIN))
            {
                struct sockaddr_in client_address;
                socklen_t client_addrlengh=sizeof(client_address);
                int connfd=accept(listenfd,(struct sockaddr *)&client_address,&client_addrlengh);
                if(connfd<0)
                {
                    printf("errno is:%d\n",errno);
                    continue;
                }
                if(user_counter>=USER_LIMIT)
                {
                    const char *info="too many users\n";
                    printf("%s",info);
                    send(connfd,info,strlen(info),0);
                    close(connfd);
                    continue;
                }
                user_counter++;
                users[connfd].address=client_address;
                setnonblocking(connfd);
                fds[user_counter].fd=connfd;
                fds[user_counter].events=POLLIN |POLLRDHUP|POLLERR;
                fds[user_counter].revents=0;
                printf("comes a new user,now have %d users\n",user_counter);
            }
            else if(fds[i].revents&POLLERR)
            {
                printf("get an error from %d \n",fds[i].fd);
                char errors[100];
                memset(errors,'\0',100);
                socklen_t length=sizeof(errors);
                if(getsockopt(fds[i].fd,SOL_SOCKET,SO_ERROR,&errors,&length)<0)
                {
                    printf("get socket option failed\n");
                }
                continue;
            }
            else if(fds[i].revents&POLLRDHUP)
            {
                users[fds[i].fd]=users[fds[user_counter].fd];
                close(fds[i].fd);
                fds[i]=fds[user_counter];
                i--;
                user_counter--;
                printf("a client left\n");
            }
            else if(fds[i].revents&POLLIN)
            {
                int connfd=fds[i].fd;
                memset(users[connfd].buf,'\0',BUFFER_SIZE-1);
                ret=recv(connfd,users[connfd].buf,BUFFER_SIZE-1,0);
                if(ret<0)
                {
                    if(errno!=EAGAIN)
                    {
                        close(connfd);
                        users[fds[i].fd]=users[fds[user_counter].fd];
                        fds[i]=fds[user_counter];
                        i--;
                        user_counter--;
                    }
                }
                else if(ret==0)
                {
                }
                else
                {
                    for(j=1;j<=user_counter;++j)
                    {
                        if(fds[j].fd==connfd)
                        {
                            continue;
                        }
                        fds[j].events |=~POLLIN;
                        fds[j].events |=POLLOUT;
                        users[fds[j].fd].write_buf=users[connfd].buf;
                    }
                }
            }
            else if(fds[i].revents&POLLOUT)
            {
                int connfd=fds[i].fd;
                if(!users[connfd].write_buf)
                {
                    continue;
                }
                ret=send(connfd,users[connfd].write_buf,strlen(users[connfd].write_buf),0);
                users[connfd].write_buf=NULL;
                fds[i].events |=~POLLOUT;
                fds[i].events |=POLLIN;
            }
        }
    }
    free(users);
    close(listenfd);
    return 0;
}