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

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

服务器代码

#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 

#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 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 

#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;
}

你可能感兴趣的:(Linux)