grpc/test/http2_test/http2_test_server.py

# Copyright 2016 gRPC authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""HTTP2 Test Server"""

import argparse
import logging
import sys

import http2_base_server
import test_data_frame_padding
import test_goaway
import test_max_streams
import test_ping
import test_rst_after_data
import test_rst_after_header
import test_rst_during_data
import twisted
import twisted.internet
import twisted.internet.endpoints
import twisted.internet.reactor

_TEST_CASE_MAPPING = {
    'rst_after_header': test_rst_after_header.TestcaseRstStreamAfterHeader,
    'rst_after_data': test_rst_after_data.TestcaseRstStreamAfterData,
    'rst_during_data': test_rst_during_data.TestcaseRstStreamDuringData,
    'goaway': test_goaway.TestcaseGoaway,
    'ping': test_ping.TestcasePing,
    'max_streams': test_max_streams.TestcaseSettingsMaxStreams,

    # Positive tests below:
    'data_frame_padding': test_data_frame_padding.TestDataFramePadding,
    'no_df_padding_sanity_test': test_data_frame_padding.TestDataFramePadding,
}

_exit_code = 0


class H2Factory(twisted.internet.protocol.Factory):

    def __init__(self, testcase):
        logging.info('Creating H2Factory for new connection (%s)', testcase)
        self._num_streams = 0
        self._testcase = testcase

    def buildProtocol(self, addr):
        self._num_streams += 1
        logging.info('New Connection: %d' % self._num_streams)
        if not _TEST_CASE_MAPPING.has_key(self._testcase):
            logging.error('Unknown test case: %s' % self._testcase)
            assert (0)
        else:
            t = _TEST_CASE_MAPPING[self._testcase]

        if self._testcase == 'goaway':
            return t(self._num_streams).get_base_server()
        elif self._testcase == 'no_df_padding_sanity_test':
            return t(use_padding=False).get_base_server()
        else:
            return t().get_base_server()


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--base_port',
        type=int,
        default=8080,
        help='base port to run the servers (default: 8080). One test server is '
        'started on each incrementing port, beginning with base_port, in the '
        'following order: data_frame_padding,goaway,max_streams,'
        'no_df_padding_sanity_test,ping,rst_after_data,rst_after_header,'
        'rst_during_data')
    return parser.parse_args()


def listen(endpoint, test_case):
    deferred = endpoint.listen(H2Factory(test_case))

    def listen_error(reason):
        # If listening fails, we stop the reactor and exit the program
        # with exit code 1.
        global _exit_code
        _exit_code = 1
        logging.error('Listening failed: %s' % reason.value)
        twisted.internet.reactor.stop()

    deferred.addErrback(listen_error)


def start_test_servers(base_port):
    """ Start one server per test case on incrementing port numbers
  beginning with base_port """
    index = 0
    for test_case in sorted(_TEST_CASE_MAPPING.keys()):
        portnum = base_port + index
        logging.warning('serving on port %d : %s' % (portnum, test_case))
        endpoint = twisted.internet.endpoints.TCP4ServerEndpoint(
            twisted.internet.reactor, portnum, backlog=128)
        # Wait until the reactor is running before calling endpoint.listen().
        twisted.internet.reactor.callWhenRunning(listen, endpoint, test_case)

        index += 1


if __name__ == '__main__':
    logging.basicConfig(
        format=
        '%(levelname) -10s %(asctime)s %(module)s:%(lineno)s | %(message)s',
        level=logging.INFO)
    args = parse_arguments()
    start_test_servers(args.base_port)
    twisted.internet.reactor.run()
    sys.exit(_exit_code)