Revision 86d1202e vncauthproxy/proxy.py

b/vncauthproxy/proxy.py
26 26
DEFAULT_CONNECT_TIMEOUT = 30
27 27
DEFAULT_CONNECT_RETRIES = 3
28 28
DEFAULT_RETRY_WAIT = 0.1
29
DEFAULT_BACKLOG = 256
30
DEFAULT_SOCK_TIMEOUT = 60.0
29 31
# We must take care not to fall into the ephemeral port range,
30 32
# this can lead to transient failures to bind a chosen port.
31 33
#
......
321 323
# Wrap all common logging functions in logging-specific methods
322 324
for funcname in ["info", "debug", "warn", "error", "critical",
323 325
                 "exception"]:
326

  
324 327
    def gen(funcname):
325 328
        def wrapped_log_func(self, *args, **kwargs):
326 329
            func = getattr(self.log, funcname)
......
367 370
    return sockets
368 371

  
369 372

  
370
def perform_server_handshake(daddr, dport, tries, retry_wait):
373
def perform_server_handshake(daddr, dport, tries, retry_wait, sock_timeout):
371 374
    """
372 375
    Initiate a connection with the backend server and perform basic
373 376
    RFB 3.8 handshake with it.
......
391 394
                server = None
392 395
                continue
393 396

  
397
            # Set socket timeout for the initial handshake
398
            server.settimeout(sock_timeout)
399

  
394 400
            try:
395 401
                logger.debug("Connecting to %s:%s", *sa[:2])
396 402
                server.connect(sa)
......
437 443
    if res != 0:
438 444
        raise Exception("Authentication error")
439 445

  
446
    # Reset the timeout for the rest of the session
447
    server.settimeout(None)
448

  
440 449
    return server
441 450

  
442 451

  
......
481 490
                      default=DEFAULT_MAX_PORT, type="int", metavar="MAX_PORT",
482 491
                      help=("The maximum port number to use for automatically-"
483 492
                            "allocated ephemeral ports"))
493
    parser.add_option("-b", "--backlog", dest="backlog",
494
                      default=DEFAULT_BACKLOG, type="int", metavar="BACKLOG",
495
                      help=("Length of the backlog queue for the control"
496
                            "connection socket"))
497
    parser.add_option("--socket-timeout", dest="sock_timeout",
498
                      default=DEFAULT_SOCK_TIMEOUT, type="float",
499
                      metavar="SOCK_TIMEOUT",
500
                      help=("Socket timeout for the server handshake"))
484 501

  
485 502
    return parser.parse_args(args)
486 503

  
487 504

  
505
def establish_connection(client, addr, ports, opts):
506
    # Receive and parse a client request.
507
    response = {
508
        "source_port": 0,
509
        "status": "FAILED",
510
    }
511
    try:
512
        # TODO: support multiple forwardings in the same message?
513
        #
514
        # Control request, in JSON:
515
        #
516
        # {
517
        #     "source_port":
518
        #         <source port or 0 for automatic allocation>,
519
        #     "destination_address":
520
        #         <destination address of backend server>,
521
        #     "destination_port":
522
        #         <destination port>
523
        #     "password":
524
        #         <the password to use to authenticate clients>
525
        # }
526
        #
527
        # The <password> is used for MITM authentication of clients
528
        # connecting to <source_port>, who will subsequently be
529
        # forwarded to a VNC server listening at
530
        # <destination_address>:<destination_port>
531
        #
532
        # Control reply, in JSON:
533
        # {
534
        #     "source_port": <the allocated source port>
535
        #     "status": <one of "OK" or "FAILED">
536
        # }
537
        #
538
        buf = client.recv(1024)
539
        req = json.loads(buf)
540

  
541
        sport_orig = int(req['source_port'])
542
        daddr = req['destination_address']
543
        dport = int(req['destination_port'])
544
        password = req['password']
545
    except Exception, e:
546
        logger.warn("Malformed request: %s", buf)
547
        client.send(json.dumps(response))
548
        client.close()
549

  
550
    # Spawn a new Greenlet to service the request.
551
    server = None
552
    try:
553
        # If the client has so indicated, pick an ephemeral source port
554
        # randomly, and remove it from the port pool.
555
        if sport_orig == 0:
556
            while True:
557
                try:
558
                    sport = random.choice(ports)
559
                    ports.remove(sport)
560
                    break
561
                except ValueError:
562
                    logger.debug("Port %d already taken", sport)
563

  
564
            logger.debug("Got port %d from pool, %d remaining",
565
                         sport, len(ports))
566
            pool = ports
567
        else:
568
            sport = sport_orig
569
            pool = None
570

  
571
        listeners = get_listening_sockets(sport)
572
        server = perform_server_handshake(daddr, dport,
573
                                          opts.connect_retries,
574
                                          opts.retry_wait, opts.sock_timeout)
575

  
576
        VncAuthProxy.spawn(logger, listeners, pool, daddr, dport,
577
                           server, password, opts.connect_timeout)
578

  
579
        logger.info("New forwarding: %d (client req'd: %d) -> %s:%d",
580
                    sport, sport_orig, daddr, dport)
581
        response = {"source_port": sport,
582
                    "status": "OK"}
583
    except IndexError:
584
        logger.error(("FAILED forwarding, out of ports for [req'd by "
585
                      "client: %d -> %s:%d]"),
586
                     sport_orig, daddr, dport)
587
    except Exception, msg:
588
        logger.error(msg)
589
        logger.error(("FAILED forwarding: %d (client req'd: %d) -> "
590
                      "%s:%d"), sport, sport_orig, daddr, dport)
591
        if not pool is None:
592
            pool.append(sport)
593
            logger.debug("Returned port %d to pool, %d remanining",
594
                         sport, len(pool))
595
        if not server is None:
596
            server.close()
597
    finally:
598
        client.send(json.dumps(response))
599
        client.close()
600

  
601

  
488 602
def main():
489 603
    """Run the daemon from the command line"""
490 604

  
......
546 660

  
547 661
    os.umask(old_umask)
548 662

  
549
    ctrl.listen(1)
663
    ctrl.listen(opts.backlog)
550 664
    logger.info("Initialized, waiting for control connections at %s",
551 665
                opts.ctrl_socket)
552 666

  
......
566 680
            client, addr = ctrl.accept()
567 681
            logger.info("New control connection")
568 682

  
569
            # Receive and parse a client request.
570
            response = {
571
                "source_port": 0,
572
                "status": "FAILED"
573
            }
574
            try:
575
                # TODO: support multiple forwardings in the same message?
576
                #
577
                # Control request, in JSON:
578
                #
579
                # {
580
                #     "source_port":
581
                #         <source port or 0 for automatic allocation>,
582
                #     "destination_address":
583
                #         <destination address of backend server>,
584
                #     "destination_port":
585
                #         <destination port>
586
                #     "password":
587
                #         <the password to use to authenticate clients>
588
                # }
589
                #
590
                # The <password> is used for MITM authentication of clients
591
                # connecting to <source_port>, who will subsequently be
592
                # forwarded to a VNC server listening at
593
                # <destination_address>:<destination_port>
594
                #
595
                # Control reply, in JSON:
596
                # {
597
                #     "source_port": <the allocated source port>
598
                #     "status": <one of "OK" or "FAILED">
599
                # }
600
                #
601
                buf = client.recv(1024)
602
                req = json.loads(buf)
603

  
604
                sport_orig = int(req['source_port'])
605
                daddr = req['destination_address']
606
                dport = int(req['destination_port'])
607
                password = req['password']
608
            except Exception, e:
609
                logger.warn("Malformed request: %s", buf)
610
                client.send(json.dumps(response))
611
                client.close()
612
                continue
613

  
614
            # Spawn a new Greenlet to service the request.
615
            server = None
616
            try:
617
                # If the client has so indicated, pick an ephemeral source port
618
                # randomly, and remove it from the port pool.
619
                if sport_orig == 0:
620
                    sport = random.choice(ports)
621
                    ports.remove(sport)
622
                    logger.debug("Got port %d from pool, %d remaining",
623
                                 sport, len(ports))
624
                    pool = ports
625
                else:
626
                    sport = sport_orig
627
                    pool = None
628

  
629
                listeners = get_listening_sockets(sport)
630
                server = perform_server_handshake(daddr, dport,
631
                                                  opts.connect_retries,
632
                                                  opts.retry_wait)
633

  
634
                VncAuthProxy.spawn(logger, listeners, pool, daddr, dport,
635
                                   server, password, opts.connect_timeout)
636

  
637
                logger.info("New forwarding: %d (client req'd: %d) -> %s:%d",
638
                            sport, sport_orig, daddr, dport)
639
                response = {"source_port": sport,
640
                            "status": "OK"}
641
            except IndexError:
642
                logger.error(("FAILED forwarding, out of ports for [req'd by "
643
                              "client: %d -> %s:%d]"),
644
                             sport_orig, daddr, dport)
645
            except Exception, msg:
646
                logger.error(msg)
647
                logger.error(("FAILED forwarding: %d (client req'd: %d) -> "
648
                              "%s:%d"), sport, sport_orig, daddr, dport)
649
                if not pool is None:
650
                    pool.append(sport)
651
                    logger.debug("Returned port %d to pool, %d remanining",
652
                                 sport, len(pool))
653
                if not server is None:
654
                    server.close()
655
            finally:
656
                client.send(json.dumps(response))
657
                client.close()
683
            gevent.Greenlet.spawn(establish_connection, client, addr,
684
                                  ports, opts)
658 685
        except Exception, e:
659 686
            logger.exception(e)
660 687
            continue

Also available in: Unified diff