select系统调用以及一个简单的TCP代理程序

假设需要编写一个简单的聊天程序，它要做两件事：第一，从键盘读入想要发给对方的聊天内容，然后把它通过socket发送给对方。第二，从socket接收来自对方的聊天内容，然后把它输出到屏幕上。这两件事不可能依次进行，因为你不可能会忍受一个聊天程序要等到在你发完消息后才能接收消息或者接收消息后才能发消息。为了这两件事能够互不影响的进行，一个解决方法是使用多线程，让两个线程分别去做这两件事。另一个解决方法就是使用select系统调用。

为了使用select，需要了解在Linux上数据的写入和读取都是通过文件描述符进行的。文件描述符由非负整数表示，例如从键盘读取是通过文件描述符0，输出到屏幕上是通过文件描述符1。并且这些文件描述符都已设为非阻塞模式。

select的作用是能够同时对多个文件描述符进行监测，筛选出哪些已有数据可以读取，哪些可以写入。

现在来看看select系统调用的函数原型，它在头文件sys/select.h中定义。

int select(int nfds, fd_set *readfds, fd_set *writefds,
               fd_set *exceptfds, struct timeval *timeout);

需要监测用来读取数据的文件描述符放入readfds集合中，需要监测用来写入数据的文件描述符放入writefds集合中。exceptfds集合用来存放需要被监测是否有特殊情况发生的文件描述符，如TCP out-of-band data。nfds必须被赋值为这三个集合中最大的那个文件描述符的值加1。如果某个集合为空，可以用NULL来表示。可以使用下面这些函数来把需要监测的文件描述符从集合中清除和添加，或清空整个集合。

void FD_CLR(int fd, fd_set *set); /* 清除 */
void FD_SET(int fd, fd_set *set); /* 添加 */
void FD_ZERO(fd_set *set); /* 清空 */

最后一个参数timeout如果被赋值为NULL，表示select调用将一直处于阻塞状态直到有一个被监测的文件描述符准备就绪。否则需要准备一个timeval结构，把需要等待的时间赋给它，表示在一定时间内如果没有文件描述符就绪，select也将返回。timeval结构如下：

struct timeval 
{
        long    tv_sec;         /* 秒 */
        long    tv_usec;        /* 微秒 */
};

如果select执行成功，readfds、writefds和exceptfds将会包含已准备就绪的文件描述符，并返回这些文件描述符的总数。如果执行失败则返回-1。

执行成功后，可以通过如下函数判断某个文件描述符是否在某个集合中。

int  FD_ISSET(int fd, fd_set *set); /* 判断是否存在 */

最后附上一个使用select系统调用的TCP代理程序源代码：

#include <sys/select.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <strings.h>

#define MAX(x, y) ((x) > (y) ? (x) : (y))
#define BUFSIZE 1024

int accpet_connection(const char *ipaddr, uint16_t port);
int connect_remote(const char *ipaddr, uint16_t port);

int main(int argc, const char *argv[])
{
	char buf[BUFSIZE];

	if (argc < 5)
	{
		printf("usage: %s localaddr localport remoteaddr remoteport \n", argv[0]);
		return -1;
	}

	uint16_t localport = (uint16_t)atoi(argv[2]);
	uint16_t remoteport = (uint16_t)atoi(argv[4]);
	
	fd_set rmask;
	fd_set wmask;

	FD_ZERO(&rmask);
	FD_ZERO(&wmask);

	int accpfd, remotfd;

	printf("Now you can connect %s:%d, it will be redirected to %s:%d\n", argv[1], localport, argv[3], remoteport);

	if ((accpfd = accpet_connection(argv[1], localport)) < 0)
		return -1;

	if ((remotfd = connect_remote(argv[3], remoteport)) < 0)
		return -1;


	for (;;)
	{
		FD_SET(accpfd, &rmask);
		FD_SET(remotfd, &rmask);

		int fd_count;

		if ((fd_count = select(MAX(accpfd, remotfd) + 1, &rmask, &wmask, NULL, NULL)) < 0)
		{
			perror("select error");
			return -1;
		}

		ssize_t accplen;
		ssize_t remotlen;

		if (fd_count > 0)
		{
			if (FD_ISSET(accpfd, &rmask))
			{
				while ((accplen = read(accpfd, buf, BUFSIZE)) > 0)
				{
					printf("Get some data from client %s\n", argv[1]);
					write(remotfd, buf, accplen);
				}
			}

			if (FD_ISSET(remotfd, &rmask))
			{
				while ((remotlen = read(remotfd, buf, BUFSIZE)) > 0)
				{
					printf("Get some data from remote %s\n", argv[3]);
					write(accpfd, buf, remotlen);
				}
			}

			/* 对方已关闭socket连接 */
			if (0 == accplen || 0 == remotlen)
				return 0;
		}
	}

	return 0;
}

int accpet_connection(const char *ipaddr, uint16_t port)
{
	int listenfd, connfd;
	struct sockaddr_in servaddr;

	if ((listenfd = socket(AF_INET, SOCK_STREAM, 0)) < 0)
	{
		perror("socket error");
		return -1;
	}

	bzero(&servaddr, sizeof(servaddr));
	servaddr.sin_family = AF_INET;
	servaddr.sin_port = htons(port);

	if (inet_pton(AF_INET, ipaddr, &servaddr.sin_addr) <= 0)
	{
		perror("inet_pton error");
		return -1;
	}

	if (bind(listenfd, (struct sockaddr *)&servaddr, sizeof(servaddr)) < 0)
	{
		perror("bind error");
		return -1;
	}

	if (listen(listenfd, 5) < 0)
	{
		perror("listen error");
		return -1;
	}

	if ((connfd = accept(listenfd, NULL, NULL)) < 0)
	{
		perror("accpet error");
		return -1;
	}

	/*
	 * 设置文件描述符为非阻塞模式。
	 */
	int flags;
    if ((flags = fcntl(connfd, F_GETFL, 0)) < 0)
	{
		perror("fcntl get flags error");
		return -1;
	}

    if (fcntl(connfd, F_SETFL, flags | O_NONBLOCK) < 0)
	{
		perror("fcntl set O_NONBLOCK error");
		return -1;
	}

	return connfd;
}

int connect_remote(const char *ipaddr, uint16_t port)
{
	int sockfd;
	struct sockaddr_in servaddr;

	if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) < 0)
	{
		perror("socket error");
		return -1;
	}

	bzero(&servaddr, sizeof(servaddr));
	servaddr.sin_family = AF_INET;
	servaddr.sin_port = htons(port);
	if (inet_pton(AF_INET, ipaddr, &servaddr.sin_addr) <= 0)
	{
		perror("inet_pton error");
		return -1;
	}

	if (connect(sockfd, (struct sockaddr *)&servaddr, sizeof(servaddr)) < 0)
	{
		perror("connect error");
		return -1;
	}

	/*
	 * 设置文件描述符为非阻塞模式。
	 */
	int flags;
    if ((flags = fcntl(sockfd, F_GETFL, 0)) < 0)
	{
		perror("fcntl get flags error");
		return -1;
	}

    if (fcntl(sockfd, F_SETFL, flags | O_NONBLOCK) < 0)
	{
		perror("fcntl set O_NONBLOCK error");
		return -1;
	}

	return sockfd;
}

以及一个用做测试的回写字符串的TCP服务脚本：

#!/usr/bin/perl

use strict;
use warnings;
use IO::Socket::INET;

$SIG{CHLD} = 'IGNORE';

my $sock = IO::Socket::INET->new(Listen => 5,
                                 LocalAddr => '192.168.1.100',
                                 LocalPort => 9999,
                                 Proto     => 'tcp');

while (1)
{
	my $client = $sock->accept() or do {next;};

	my $pid = fork;

	if (0 == $pid)
	{
		$client->autoflush(1);
		while (<$client>)
		{
			print $client $_;
		}
		exit 0;
	}

	close $client;
}

select系统调用以及一个简单的TCP代理程序

猜你喜欢