用select实现TCP回射程序（服务器及客户端）

小灰灰 2022-06-04 00:28 237阅读 0赞

该程序系参考《UNP》来编写的，主要用来练习select及shutdown函数的使用。

服务器代码

#include <unistd.h>
    #include <stdio.h>
    #include <stdlib.h>
    #include <sys/socket.h>
    #include <sys/types.h>
    #include <arpa/inet.h>
    #include <netinet/in.h>
    #include <sys/select.h>
    #include <string.h>
    #include <strings.h>
    #include <errno.h>
    
    #define PORT 8888
    #define BACKLOG 100
    #define MAXLINE 1024
    
    int main() {
        int listenfd, connfd, sockfd;
        int client[FD_SETSIZE];
        int i, n;
        struct sockaddr_in servaddr, cliaddr;
        socklen_t clilen;
        fd_set rset, allset;
        int maxfd, maxcli;
        int nready;
        char buf[MAXLINE];
    
        //创建监听套接字
        if ((listenfd = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
            printf("socket error\n");
            exit(1);
        }
    
        //写入套接字地址
        bzero(&servaddr, sizeof(servaddr));
        servaddr.sin_family = AF_INET;
        servaddr.sin_addr.s_addr = htonl(INADDR_ANY);
        servaddr.sin_port = htons(PORT);
    
        //绑定地址
        if ((n = bind(listenfd, (struct sockaddr *) &servaddr, sizeof(servaddr))) < 0) {
            printf("bind error\n");
            exit(1);
        }
    
        //监听
        if ((n = listen(listenfd, BACKLOG)) < 0) {
            printf("listen error\n");
            exit(1);
        }
    
        for (i = 0; i < FD_SETSIZE; i++)
            client[i] = -1;
        maxfd = listenfd;
        maxcli = -1;
        FD_ZERO(&allset);
        FD_SET(listenfd, &allset);
    
        for (;;) {
            rset = allset;
    
            if ((nready = select(maxfd + 1, &rset, NULL, NULL, NULL)) < 0) {
                printf("select error\n");
                printf("%s\n", strerror(errno));
                exit(1);
            }
    
            if (FD_ISSET(listenfd, &rset)) {
                //接受新连接
                if ((connfd = accept(listenfd, (struct sockaddr *) &cliaddr, &clilen)) < 0) {
                    printf("accept error\n");
                    exit(1);
                }
                printf("accept: %d\n", connfd);
    
                //更新连接表
                for (i = 0; i < FD_SETSIZE; i++)
                    if (client[i] < 0) {
                        client[i] = connfd;
                        break;
                    }
    
                FD_SET(connfd, &allset); //更新描述符集
    
                if (connfd > maxfd)
                    maxfd = connfd; //更新最大文件描述符
    
                if (i > maxcli)
                    maxcli = i; //更新最大连接表索引
    
                if (--nready <= 0)
                    continue;
            }
    
            for (i = 0; i <= maxcli; i++) {
                if ((sockfd = client[i]) < 0)
                    continue;
                if (FD_ISSET(sockfd, &rset)) {
                    if ((n = read(sockfd, buf, MAXLINE)) == 0) {
                        if (close(sockfd) < 0) {
                            printf("close error\n");
                            exit(1);
                        }
                        printf("close: %d\n", sockfd);
                        FD_CLR(sockfd, &allset);
                        client[i] = -1;
                    } else {
                        if (n < 0) {
                            printf("read error\n");
                            exit(0);
                        }
                        if (n > 0) {
                            if (write(sockfd, buf, n) < 0) {
                                printf("write error\n");
                                exit(1);
                            }
                        }
                    }
                    if (--nready <= 0)
                        break;
                }
            }
        }
    }

客户端代码

#include <unistd.h>
    #include <stdio.h>
    #include <stdlib.h>
    #include <sys/socket.h>
    #include <sys/types.h>
    #include <arpa/inet.h>
    #include <netinet/in.h>
    #include <sys/select.h>
    #include <string.h>
    #include <strings.h>
    
    #define PORT 8888
    #define MAXLINE 1024
    
    int main(int argc, char* argv[])
    {
        int sockfd;
        int n;
        struct sockaddr_in servaddr;
        char buf[MAXLINE];
        fd_set rset;
        int maxfd;
        int stdineof;
    
        if(argc != 2){
            printf("invalid usage!\n");
            exit(1);
        }
    
        if((sockfd = socket(AF_INET, SOCK_STREAM, 0)) < 0){
            printf("socket error!\n");
            exit(1);
        }
    
        servaddr.sin_family = AF_INET;
        servaddr.sin_port = htons(PORT);
        if((n = inet_pton(AF_INET,  argv[1], &servaddr.sin_addr)) < 0){
            printf("pton error\n");
            exit(1);
        }
        else if(n == 0){
            printf("invalid ip\n");
            exit(1);
        }
    
        if((n = connect(sockfd, (struct sockaddr*)&servaddr, sizeof(servaddr))) < 0){
            printf("connect error\n");
            exit(1);
        }
    
        maxfd = sockfd;
        FD_ZERO(&rset);
        stdineof = 0;
    
        for(;;){
            if(stdineof == 0)
                FD_SET(0, &rset);
            FD_SET(sockfd, &rset);
    
            if((n = select(maxfd+1, &rset, NULL, NULL, NULL)) < 0){
                printf("select error\n");
                exit(1);
            }
    
            if(FD_ISSET(0, &rset)){
                if((n = read(0, buf, MAXLINE)) == 0){
                    stdineof = 1;
                    if(shutdown(sockfd, SHUT_WR) < 0){
                        printf("shutdown error\n");
                        exit(1);
                    }
                    FD_CLR(0, &rset);
                    continue;
                }
                else if(n < 0){
                    printf("stdin read error\n");
                    exit(1);
                }
                if((n = write(sockfd, buf, n)) < 0){
                    printf("socket write error\n");
                    exit(1);
                }
            }
    
            if(FD_ISSET(sockfd, &rset)){
                if((n = read(sockfd, buf, MAXLINE)) == 0){
                    if(stdineof == 1)
                        break;
                    else{
                        printf("server terminated\n");
                        exit(1);
                    }
                }
                else if(n < 0){
                    printf("socket read error\n");
                    exit(1);
                }
                if((n = write(1, buf, n)) < 0){
                    printf("stdout write error\n");
                    exit(1);
                }
            }
        }
        return 0;
    }