mirror of
https://github.com/klzgrad/naiveproxy.git
synced 2024-12-01 01:36:09 +03:00
136 lines
4.7 KiB
Python
136 lines
4.7 KiB
Python
# Copyright 2018 The Chromium Authors. All rights reserved.
|
|
# Use of this source code is governed by a BSD-style license that can be
|
|
# found in the LICENSE file.
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import select
|
|
import socket
|
|
import sys
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
|
|
DIR_SOURCE_ROOT = os.path.abspath(
|
|
os.path.join(os.path.dirname(__file__), os.pardir, os.pardir))
|
|
sys.path.append(os.path.join(DIR_SOURCE_ROOT, 'build', 'util', 'lib', 'common'))
|
|
import chrome_test_server_spawner
|
|
|
|
PORT_MAP_RE = re.compile('Allocated port (?P<port>\d+) for remote')
|
|
GET_PORT_NUM_TIMEOUT_SECS = 5
|
|
|
|
|
|
def _ConnectPortForwardingTask(target, local_port):
|
|
"""Establishes a port forwarding SSH task to a localhost TCP endpoint hosted
|
|
at port |local_port|. Blocks until port forwarding is established.
|
|
|
|
Returns the remote port number."""
|
|
|
|
forwarding_flags = ['-O', 'forward', # Send SSH mux control signal.
|
|
'-R', '0:localhost:%d' % local_port,
|
|
'-v', # Get forwarded port info from stderr.
|
|
'-NT'] # Don't execute command; don't allocate terminal.
|
|
task = target.RunCommandPiped([],
|
|
ssh_args=forwarding_flags,
|
|
stderr=subprocess.PIPE)
|
|
|
|
# SSH reports the remote dynamic port number over stderr.
|
|
# Unfortunately, the output is incompatible with Python's line buffered
|
|
# input (or vice versa), so we have to build our own buffered input system to
|
|
# pull bytes over the pipe.
|
|
poll_obj = select.poll()
|
|
poll_obj.register(task.stderr, select.POLLIN)
|
|
line = ''
|
|
timeout = time.time() + GET_PORT_NUM_TIMEOUT_SECS
|
|
while time.time() < timeout:
|
|
poll_result = poll_obj.poll(max(0, timeout - time.time()))
|
|
if poll_result:
|
|
next_char = task.stderr.read(1)
|
|
if not next_char:
|
|
break
|
|
line += next_char
|
|
if line.endswith('\n'):
|
|
line = line[:-1]
|
|
logging.debug('ssh: ' + line)
|
|
matched = PORT_MAP_RE.match(line)
|
|
if matched:
|
|
device_port = int(matched.group('port'))
|
|
logging.debug('Port forwarding established (local=%d, device=%d)' %
|
|
(local_port, device_port))
|
|
task.wait()
|
|
return device_port
|
|
line = ''
|
|
|
|
raise Exception('Could not establish a port forwarding connection.')
|
|
|
|
|
|
# Implementation of chrome_test_server_spawner.PortForwarder that uses SSH's
|
|
# remote port forwarding feature to forward ports.
|
|
class SSHPortForwarder(chrome_test_server_spawner.PortForwarder):
|
|
def __init__(self, target):
|
|
self._target = target
|
|
|
|
# Maps the host (server) port to the device port number.
|
|
self._port_mapping = {}
|
|
|
|
def Map(self, port_pairs):
|
|
for p in port_pairs:
|
|
_, host_port = p
|
|
self._port_mapping[host_port] = \
|
|
_ConnectPortForwardingTask(self._target, host_port)
|
|
|
|
def GetDevicePortForHostPort(self, host_port):
|
|
return self._port_mapping[host_port]
|
|
|
|
def Unmap(self, device_port):
|
|
for host_port, entry in self._port_mapping.iteritems():
|
|
if entry == device_port:
|
|
forwarding_args = [
|
|
'-NT', '-O', 'cancel', '-R',
|
|
'%d:localhost:%d' % (self._port_mapping[host_port], host_port)]
|
|
task = self._target.RunCommandPiped([],
|
|
ssh_args=forwarding_args,
|
|
stderr=subprocess.PIPE)
|
|
task.wait()
|
|
if task.returncode != 0:
|
|
raise Exception(
|
|
'Error %d when unmapping port %d' % (task.returncode,
|
|
device_port))
|
|
del self._port_mapping[host_port]
|
|
return
|
|
|
|
raise Exception('Unmap called for unknown port: %d' % device_port)
|
|
|
|
|
|
def SetupTestServer(target, test_concurrency):
|
|
"""Provisions a forwarding test server and configures |target| to use it.
|
|
|
|
Returns a Popen object for the test server process."""
|
|
|
|
logging.debug('Starting test server.')
|
|
spawning_server = chrome_test_server_spawner.SpawningServer(
|
|
0, SSHPortForwarder(target), test_concurrency)
|
|
forwarded_port = _ConnectPortForwardingTask(
|
|
target, spawning_server.server_port)
|
|
spawning_server.Start()
|
|
|
|
logging.debug('Test server listening for connections (port=%d)' %
|
|
spawning_server.server_port)
|
|
logging.debug('Forwarded port is %d' % forwarded_port)
|
|
|
|
config_file = tempfile.NamedTemporaryFile(delete=True)
|
|
|
|
# Clean up the config JSON to only pass ports. See https://crbug.com/810209 .
|
|
config_file.write(json.dumps({
|
|
'name': 'testserver',
|
|
'address': '127.0.0.1',
|
|
'spawner_url_base': 'http://localhost:%d' % forwarded_port
|
|
}))
|
|
|
|
config_file.flush()
|
|
target.PutFile(config_file.name, '/data/net-test-server-config')
|
|
|
|
return spawning_server
|