Revision 020f4a9e
| b/vncauthproxy/proxy.py | ||
|---|---|---|
| 34 | 34 |
import sys |
| 35 | 35 |
import logging |
| 36 | 36 |
import gevent |
| 37 |
import gevent.event |
|
| 37 | 38 |
import daemon |
| 38 | 39 |
import random |
| 39 | 40 |
import daemon.runner |
| ... | ... | |
| 116 | 117 |
VncAuthProxy.id += 1 |
| 117 | 118 |
self.log = logger |
| 118 | 119 |
self.listeners = listeners |
| 120 |
# A list of worker/forwarder greenlets, one for each direction |
|
| 121 |
self.workers = [] |
|
| 119 | 122 |
# All listening sockets are assumed to be on the same port |
| 120 | 123 |
self.sport = listeners[0].getsockname()[1] |
| 121 | 124 |
self.pool = pool |
| ... | ... | |
| 127 | 130 |
self.timeout = connect_timeout |
| 128 | 131 |
|
| 129 | 132 |
def _cleanup(self): |
| 130 |
"""Close all active sockets and exit gracefully""" |
|
| 133 |
"""Cleanup everything: workers, sockets, ports |
|
| 134 |
|
|
| 135 |
Kill all remaining forwarder greenlets, close all active sockets, |
|
| 136 |
return the source port to the pool if applicable, then exit |
|
| 137 |
gracefully. |
|
| 138 |
|
|
| 139 |
""" |
|
| 140 |
# Make sure all greenlets are dead, then clean them up |
|
| 141 |
self.debug("Cleaning up %d workers", len(self.workers))
|
|
| 142 |
for g in self.workers: |
|
| 143 |
g.kill() |
|
| 144 |
gevent.joinall(self.workers) |
|
| 145 |
del self.workers |
|
| 146 |
|
|
| 131 | 147 |
# Reintroduce the port number of the client socket in |
| 132 | 148 |
# the port pool, if applicable. |
| 133 | 149 |
if not self.pool is None: |
| ... | ... | |
| 135 | 151 |
self.debug("Returned port %d to port pool, contains %d ports",
|
| 136 | 152 |
self.sport, len(self.pool)) |
| 137 | 153 |
|
| 154 |
self.debug("Cleaning up sockets")
|
|
| 138 | 155 |
while self.listeners: |
| 139 | 156 |
self.listeners.pop().close() |
| 140 | 157 |
if self.server: |
| ... | ... | |
| 142 | 159 |
if self.client: |
| 143 | 160 |
self.client.close() |
| 144 | 161 |
|
| 162 |
self.info("Cleaned up connection, all done")
|
|
| 145 | 163 |
raise gevent.GreenletExit |
| 146 | 164 |
|
| 147 | 165 |
def __str__(self): |
| ... | ... | |
| 258 | 276 |
self._client_handshake() |
| 259 | 277 |
|
| 260 | 278 |
# Bridge both connections through two "forwarder" greenlets. |
| 261 |
self.workers = [gevent.spawn(self._forward, |
|
| 262 |
self.client, self.server), |
|
| 263 |
gevent.spawn(self._forward, |
|
| 264 |
self.server, self.client)] |
|
| 265 |
|
|
| 266 |
# If one greenlet goes, the other has to go too. |
|
| 267 |
self.workers[0].link(self.workers[1]) |
|
| 268 |
self.workers[1].link(self.workers[0]) |
|
| 269 |
gevent.joinall(self.workers) |
|
| 270 |
del self.workers |
|
| 279 |
# This greenlet will wait until any of the workers dies. |
|
| 280 |
# Final cleanup will take place in _cleanup(). |
|
| 281 |
dead = gevent.event.Event() |
|
| 282 |
dead.clear() |
|
| 283 |
|
|
| 284 |
# This callback will get called if any of the two workers dies. |
|
| 285 |
def callback(g): |
|
| 286 |
self.debug("Worker %d/%d died", self.workers.index(g),
|
|
| 287 |
len(self.workers)) |
|
| 288 |
dead.set() |
|
| 289 |
|
|
| 290 |
self.workers.append(gevent.spawn(self._forward, |
|
| 291 |
self.client, self.server)) |
|
| 292 |
self.workers.append(gevent.spawn(self._forward, |
|
| 293 |
self.server, self.client)) |
|
| 294 |
for g in self.workers: |
|
| 295 |
g.link(callback) |
|
| 296 |
|
|
| 297 |
# Wait until any of the workers dies |
|
| 298 |
self.debug("Waiting for any of %d workers to die",
|
|
| 299 |
len(self.workers)) |
|
| 300 |
dead.wait() |
|
| 301 |
|
|
| 302 |
# We can go now, _cleanup() will take care of |
|
| 303 |
# all worker, socket and port cleanup |
|
| 304 |
self.debug("A forwarder died, our work here is done")
|
|
| 271 | 305 |
raise gevent.GreenletExit |
| 272 | 306 |
except Exception, e: |
| 273 | 307 |
# Any unhandled exception in the previous block |
Also available in: Unified diff