21-IO复用方式实现netcat

IO复用

IO复用也称：事件驱动、基于事件、reactor。其实它复用的不是IO，而是复用线程，简单理解为让一个线程处理多个IO。注意：这种方式是同步的，不是异步的。
只用一个线程，自带的netcat也是以这种方式实现。
通常IO复用应该搭配非阻塞IO，但为了说明为什么，下面的代码使用阻塞IO（一个方向可能会阻塞另一个方向）。
磁盘文件的文件描述符：被动地，程序可以决定什么时候去读写，不会遇到阻塞情况（没有数据才会阻塞）。
socket的文件描述符：主动地，程序不可以决定什么时候去读写。所以引用IO复用，先查一下文件描述符是否可读写，然后再进行读写。

代码

recipes/netcat.py at master · chenshuo/recipes (github.com)

#!/usr/bin/python

import os
import select
import socket
import sys

def relay(sock):
    poll = select.poll()
    poll.register(sock, select.POLLIN) # 注册读事件
    poll.register(sys.stdin, select.POLLIN)

    done = False
    while not done:
        events = poll.poll(10000)  # 10 seconds
        for fileno, event in events:
            if event & select.POLLIN:
                if fileno == sock.fileno(): # socket有数据
                    data = sock.recv(8192)
                    if data:
                        sys.stdout.write(data)
                    else:
                        done = True
                else: # stdin有数据
                    assert fileno == sys.stdin.fileno()
                    data = os.read(fileno, 8192)
                    if data:
                        sock.sendall(data)
                    else:
                        sock.shutdown(socket.SHUT_WR)
                        poll.unregister(sys.stdin)


def main(argv):
    if len(argv) < 3:
        binary = argv[0]
        print "Usage:\n  %s -l port\n  %s host port" % (argv[0], argv[0])
        return
    port = int(argv[2])
    if argv[1] == "-l":
        # server
        server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        server_socket.bind(('', port))
        server_socket.listen(5)
        (client_socket, client_address) = server_socket.accept()
        server_socket.close()
        relay(client_socket)
    else:
        # client
        sock = socket.create_connection((argv[1], port))
        relay(sock)


if __name__ == "__main__":
    main(sys.argv)