Merge pull request #46 from cloud9ers/master

Adding support for unix sockets as target socket.
This commit is contained in:
Joel Martin 2012-05-30 07:22:29 -07:00
commit 233b622e47
3 changed files with 60 additions and 38 deletions

2
.gitignore vendored
View File

@ -4,3 +4,5 @@
other/.lein-deps-sum
other/classes
other/lib
.project
.pydevproject

View File

@ -75,6 +75,7 @@ class WebSocketServer(object):
buffer_size = 65536
server_handshake_hixie = """HTTP/1.1 101 Web Socket Protocol Handshake\r
Upgrade: WebSocket\r
Connection: Upgrade\r
@ -103,7 +104,7 @@ Sec-WebSocket-Accept: %s\r
def __init__(self, listen_host='', listen_port=None, source_is_ipv6=False,
verbose=False, cert='', key='', ssl_only=None,
daemon=False, record='', web='',
run_once=False, timeout=0):
run_once=False, timeout=0, unix=None):
# settings
self.verbose = verbose
@ -113,6 +114,8 @@ Sec-WebSocket-Accept: %s\r
self.daemon = daemon
self.run_once = run_once
self.timeout = timeout
self.unix_socket = unix
self.launch_time = time.time()
self.ws_connection = False
@ -163,8 +166,7 @@ Sec-WebSocket-Accept: %s\r
#
@staticmethod
def socket(host, port=None, connect=False, prefer_ipv6=False,
use_ssl=False):
def socket(host, port=None, connect=False, prefer_ipv6=False, unix_socket=None, use_ssl=False):
""" Resolve a host (and optional port) to an IPv4 or IPv6
address. Create a socket. Bind to it if listen is set,
otherwise connect to it. Return the socket.
@ -172,7 +174,7 @@ Sec-WebSocket-Accept: %s\r
flags = 0
if host == '':
host = None
if connect and not port:
if connect and not (port or unix_socket):
raise Exception("Connect mode requires a port")
if use_ssl and not ssl:
raise Exception("SSL socket requested but Python SSL module not loaded.");
@ -180,22 +182,28 @@ Sec-WebSocket-Accept: %s\r
raise Exception("SSL only supported in connect mode (for now)")
if not connect:
flags = flags | socket.AI_PASSIVE
addrs = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM,
socket.IPPROTO_TCP, flags)
if not addrs:
raise Exception("Could resolve host '%s'" % host)
addrs.sort(key=lambda x: x[0])
if prefer_ipv6:
addrs.reverse()
sock = socket.socket(addrs[0][0], addrs[0][1])
if connect:
sock.connect(addrs[0][4])
if use_ssl:
sock = ssl.wrap_socket(sock)
else:
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock.bind(addrs[0][4])
sock.listen(100)
if not unix_socket:
addrs = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM,
socket.IPPROTO_TCP, flags)
if not addrs:
raise Exception("Could not resolve host '%s'" % host)
addrs.sort(key=lambda x: x[0])
if prefer_ipv6:
addrs.reverse()
sock = socket.socket(addrs[0][0], addrs[0][1])
if connect:
sock.connect(addrs[0][4])
if use_ssl:
sock = ssl.wrap_socket(sock)
else:
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock.bind(addrs[0][4])
sock.listen(100)
else:
sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
sock.connect(unix_socket)
return sock
@staticmethod
@ -576,10 +584,10 @@ Sec-WebSocket-Accept: %s\r
- Send a WebSockets handshake server response.
- Return the socket for this WebSocket client.
"""
stype = ""
ready = select.select([sock], [], [], 3)[0]
if not ready:
raise self.EClose("ignoring socket not ready")
# Peek, but do not read the data so that we have a opportunity
@ -757,7 +765,7 @@ Sec-WebSocket-Accept: %s\r
self.rec = None
self.start_time = int(time.time()*1000)
# handler process
# handler process
try:
try:
self.client = self.do_handshake(startsock, address)
@ -855,7 +863,7 @@ Sec-WebSocket-Accept: %s\r
continue
else:
raise
if self.run_once:
# Run in same process if run_once
self.top_new_client(startsock, address)

40
websockify Executable file → Normal file
View File

@ -88,13 +88,15 @@ Traffic Legend:
"""
# Need to call wrapped command after daemonization so we can
# know when the wrapped command exits
msg = " - proxying from %s:%s" % (
self.listen_host, self.listen_port)
if self.wrap_cmd:
msg += " to '%s' - port %s" % (
" ".join(self.wrap_cmd, self.target_port))
dst_string = self.unix_socket or "'%s' (port %s)" % (" ".join(self.wrap_cmd), self.target_port)
msg = " - proxying from %s:%s to %s\n" % (
self.listen_host, self.listen_port, dst_string)
self.run_wrap_cmd()
else:
msg += " to %s:%s" % (self.target_host, self.listen_port)
dst_string = self.unix_socket or "%s:%s" % (self.target_host, self.target_port)
msg = " - proxying from %s:%s to %s\n" % (
self.listen_host, self.listen_port, dst_string)
if self.ssl_target:
msg += " (using SSL)"
@ -144,11 +146,16 @@ Traffic Legend:
"""
Called after a new WebSocket connection has been established.
"""
# Connect to the target
if self.unix_socket:
msg = "connecting to unix socket : %s" % self.unix_socket
else:
msg = "connecting to: %s:%s" % (
self.target_host, self.target_port)
tsock = self.socket(self.target_host, self.target_port,
connect=True, use_ssl=self.ssl_target, unix_socket=self.unix_socket)
msg = "connecting to: %s:%s" % (
self.target_host, self.target_port)
if self.ssl_target:
msg += " (using SSL)"
self.msg(msg)
@ -235,6 +242,8 @@ def websockify_init():
help="verbose messages and per frame traffic")
parser.add_option("--record",
help="record sessions to FILE.[session_number]", metavar="FILE")
parser.add_option("--unix",
help="unix socket to proxy network from", metavar="FILE")
parser.add_option("--daemon", "-D",
dest="daemon", action="store_true",
help="become a daemon (background process)")
@ -259,7 +268,7 @@ def websockify_init():
(opts, args) = parser.parse_args()
# Sanity checks
if len(args) < 2:
if len(args) < 1:
parser.error("Too few arguments")
if sys.argv.count('--'):
opts.wrap_cmd = args[1:]
@ -287,12 +296,15 @@ def websockify_init():
opts.target_host = None
opts.target_port = None
else:
if args[1].count(':') > 0:
opts.target_host, opts.target_port = args[1].rsplit(':', 1)
if hasattr(opts, 'unix'):
opts.target_host = opts.target_port = None
else:
parser.error("Error parsing target")
try: opts.target_port = int(opts.target_port)
except: parser.error("Error parsing target port")
if args[1].count(':') > 0:
opts.target_host, opts.target_port = args[1].rsplit(':', 1)
else:
parser.error("Error parsing target")
try: opts.target_port = int(opts.target_port)
except: parser.error("Error parsing target port")
# Create and start the WebSockets proxy
server = WebSocketProxy(**opts.__dict__)