Home | History | Annotate | Download | only in tcp_server
      1 #
      2 # Copyright (C) 2016 The Android Open Source Project
      3 #
      4 # Licensed under the Apache License, Version 2.0 (the "License");
      5 # you may not use this file except in compliance with the License.
      6 # You may obtain a copy of the License at
      7 #
      8 #      http://www.apache.org/licenses/LICENSE-2.0
      9 #
     10 # Unless required by applicable law or agreed to in writing, software
     11 # distributed under the License is distributed on an "AS IS" BASIS,
     12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     13 # See the License for the specific language governing permissions and
     14 # limitations under the License.
     15 #
     16 
     17 import logging
     18 import socket
     19 import socketserver
     20 import threading
     21 
     22 from vts.runners.host import errors
     23 from vts.proto import AndroidSystemControlMessage_pb2 as SysMsg
     24 from vts.proto import ComponentSpecificationMessage_pb2 as CompSpecMsg
     25 from vts.utils.python.mirror import pb2py
     26 
     27 _functions = dict()  # Dictionary to hold function pointers
     28 
     29 
     30 class VtsCallbackServerError(errors.VtsError):
     31     """Raised when an error occurs in VTS TCP server."""
     32 
     33 
     34 class CallbackRequestHandler(socketserver.StreamRequestHandler):
     35     """The request handler class for our server."""
     36 
     37     def handle(self):
     38         """Receives requests from clients.
     39 
     40         When a callback happens on the target side, a request message is posted
     41         to the host side and is handled here. The message is parsed and the
     42         appropriate callback function on the host side is called.
     43         """
     44         header = self.rfile.readline().strip()
     45         try:
     46             len = int(header)
     47         except ValueError:
     48             if header:
     49                 logging.exception("Unable to convert '%s' into an integer, which "
     50                                   "is required for reading the next message." %
     51                                   header)
     52                 raise
     53             else:
     54                 logging.error('CallbackRequestHandler received empty message header. Skipping...')
     55                 return
     56         # Read the request message.
     57         received_data = self.rfile.read(len)
     58         logging.debug("Received callback message: %s", received_data)
     59         request_message = SysMsg.AndroidSystemCallbackRequestMessage()
     60         request_message.ParseFromString(received_data)
     61         logging.debug('Handling callback ID: %s', request_message.id)
     62         response_message = SysMsg.AndroidSystemCallbackResponseMessage()
     63         # Call the appropriate callback function and construct the response
     64         # message.
     65         if request_message.id in _functions:
     66             callback_args = []
     67             for arg in request_message.arg:
     68                 callback_args.append(pb2py.Convert(arg))
     69             args = tuple(callback_args)
     70             _functions[request_message.id](*args)
     71             response_message.response_code = SysMsg.SUCCESS
     72         else:
     73             logging.error("Callback function ID %s is not registered!",
     74                           request_message.id)
     75             response_message.response_code = SysMsg.FAIL
     76 
     77         # send the response back to client
     78         message = response_message.SerializeToString()
     79         # self.request is the TCP socket connected to the client
     80         self.request.sendall(message)
     81 
     82 
     83 class CallbackServer(object):
     84     """This class creates TCPServer in separate thread.
     85 
     86     Attributes:
     87         _server: an instance of socketserver.TCPServer.
     88         _port: this variable maintains the port number used in creating
     89                the server connection.
     90         _ip: variable to hold the IP Address of the host.
     91         _hostname: IP Address to which initial connection is made.
     92     """
     93 
     94     def __init__(self):
     95         self._server = None
     96         self._port = 0  # Port 0 means to select an arbitrary unused port
     97         self._ip = ""  # Used to store the IP address for the server
     98         self._hostname = "localhost"  # IP address to which initial connection is made
     99 
    100     def RegisterCallback(self, func_id, callback_func):
    101         """Registers a callback function.
    102 
    103         Args:
    104             func_id: The ID of the callback function.
    105             callback_func: The function to register.
    106 
    107         Raises:
    108             CallbackServerError is raised if the func_id is already registered.
    109         """
    110         if func_id in _functions:
    111             raise CallbackServerError(
    112                 "Function ID '%s' is already registered" % func_id)
    113         _functions[func_id] = callback_func
    114 
    115     def UnregisterCallback(self, func_id):
    116         """Removes a callback function from the registry.
    117 
    118         Args:
    119             func_id: The ID of the callback function to remove.
    120 
    121         Raises:
    122             CallbackServerError is raised if the func_id is not registered.
    123         """
    124         try:
    125             _functions.pop(func_id)
    126         except KeyError:
    127             raise CallbackServerError(
    128                 "Can't remove function ID '%s', which is not registered." %
    129                 func_id)
    130 
    131     def Start(self, port=0):
    132         """Starts the server.
    133 
    134         Args:
    135             port: integer, number of the port on which the server listens.
    136                   Default is 0, which means auto-select a port available.
    137 
    138         Returns:
    139             IP Address, port number
    140 
    141         Raises:
    142             CallbackServerError is raised if the server fails to start.
    143         """
    144         try:
    145             self._server = socketserver.TCPServer(
    146                 (self._hostname, port), CallbackRequestHandler)
    147             self._ip, self._port = self._server.server_address
    148 
    149             # Start a thread with the server.
    150             # Each request will be handled in a child thread.
    151             server_thread = threading.Thread(target=self._server.serve_forever)
    152             server_thread.daemon = True
    153             server_thread.start()
    154             logging.info('TcpServer %s started (%s:%s)', server_thread.name,
    155                          self._ip, self._port)
    156             return self._ip, self._port
    157         except (RuntimeError, IOError, socket.error) as e:
    158             logging.exception(e)
    159             raise CallbackServerError(
    160                 'Failed to start CallbackServer on (%s:%s).' %
    161                 (self._hostname, port))
    162 
    163     def Stop(self):
    164         """Stops the server.
    165         """
    166         self._server.shutdown()
    167         self._server.server_close()
    168 
    169     @property
    170     def ip(self):
    171         return self._ip
    172 
    173     @property
    174     def port(self):
    175         return self._port
    176