tcp_echo.py

00001 # Copyright 2015 ARM Limited, All rights reserved
00002 #
00003 # Licensed under the Apache License, Version 2.0 (the "License");
00004 # you may not use this file except in compliance with the License.
00005 # You may obtain a copy of the License at
00006 #
00007 #     http://www.apache.org/licenses/LICENSE-2.0
00008 #
00009 # Unless required by applicable law or agreed to in writing, software
00010 # distributed under the License is distributed on an "AS IS" BASIS,
00011 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
00012 # See the License for the specific language governing permissions and
00013 # limitations under the License.
00014 
00015 import sys
00016 import select
00017 import socket
00018 import logging
00019 from threading import Thread
00020 from sys import stdout
00021 from SocketServer import BaseRequestHandler, TCPServer
00022 from mbed_host_tests import BaseHostTest, event_callback
00023 
00024 
00025 class TCPEchoClientHandler(BaseRequestHandler):
00026     def handle(self):
00027         """
00028         Handles a connection. Test starts by client(i.e. mbed) connecting to server.
00029         This connection handler receives data and echoes back to the client util
00030         {{end}} is received. Then it sits on recv() for client to terminate the
00031         connection.
00032 
00033         Note: reason for not echoing data back after receiving {{end}} is that send
00034               fails raising a SocketError as client closes connection.
00035         """
00036         while self.server.isrunning():
00037             try:
00038                 data = self.recv()
00039                 if not data: break
00040             except Exception as e:
00041                 break
00042 
00043             try:
00044                 # echo data back to the client
00045                 self.send(data)
00046             except Exception as e:
00047                 break
00048 
00049     def recv(self):
00050         """
00051         Try to receive until server is shutdown
00052         """
00053         while self.server.isrunning():
00054             rl, wl, xl = select.select([self.request], [], [], 1)
00055             if len(rl):
00056                 return self.request.recv(1024)
00057 
00058     def send(self, data):
00059         """
00060         Try to send until server is shutdown
00061         """
00062         while self.server.isrunning():
00063             rl, wl, xl = select.select([], [self.request], [], 1)
00064             if len(wl):
00065                 self.request.sendall(data)
00066                 break
00067 
00068 
00069 class TCPServerWrapper (TCPServer):
00070     """
00071     Wrapper over TCP server to implement server initiated shutdown.
00072     Adds a flag:= running that a request handler can check and come out of
00073     recv loop when shutdown is called.
00074     """
00075 
00076     def __init__(self, addr, request_handler):
00077         # hmm, TCPServer is not sub-classed from object!
00078         if issubclass(TCPServer, object):
00079             super(TCPServerWrapper, self).__init__(addr, request_handler)
00080         else:
00081             TCPServer.__init__(self, addr, request_handler)
00082         self.running  = False
00083 
00084     def serve_forever(self):
00085         self.running  = True
00086         if issubclass(TCPServer, object):
00087             super(TCPServerWrapper, self).serve_forever()
00088         else:
00089             TCPServer.serve_forever(self)
00090 
00091     def shutdown(self):
00092         self.running  = False
00093         if issubclass(TCPServer, object):
00094             super(TCPServerWrapper, self).shutdown()
00095         else:
00096             TCPServer.shutdown(self)
00097 
00098     def isrunning(self):
00099         return self.running 
00100 
00101 
00102 class TCPEchoClientTest(BaseHostTest):
00103 
00104     def __init__(self):
00105         """
00106         Initialise test parameters.
00107 
00108         :return:
00109         """
00110         BaseHostTest.__init__(self)
00111         self.SERVER_IP = None # Will be determined after knowing the target IP
00112         self.SERVER_PORT = 0  # Let TCPServer choose an arbitrary port
00113         self.server = None
00114         self.server_thread = None
00115         self.target_ip = None
00116 
00117     @staticmethod
00118     def find_interface_to_target_addr(target_ip):
00119         """
00120         Finds IP address of the interface through which it is connected to the target.
00121 
00122         :return:
00123         """
00124         s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
00125         try:
00126             s.connect((target_ip, 0)) # Target IP, any port
00127         except socket.error:
00128             s.connect((target_ip, 8000)) # Target IP, 'random' port
00129         ip = s.getsockname()[0]
00130         s.close()
00131         return ip
00132 
00133     def setup_tcp_server(self):
00134         """
00135         sets up a TCP server for target to connect and send test data.
00136 
00137         :return:
00138         """
00139         # !NOTE: There should mechanism to assert in the host test
00140         if self.SERVER_IP is None:
00141             self.log("setup_tcp_server() called before determining server IP!")
00142             self.notify_complete(False)
00143 
00144         # Returning none will suppress host test from printing success code
00145         self.server = TCPServerWrapper((self.SERVER_IP, self.SERVER_PORT), TCPEchoClientHandler)
00146         ip, port = self.server.server_address
00147         self.SERVER_PORT = port
00148         self.server.allow_reuse_address = True
00149         self.log("HOST: Listening for TCP connections: " + self.SERVER_IP + ":" + str(self.SERVER_PORT))
00150         self.server_thread = Thread(target=TCPEchoClientTest.server_thread_func, args=(self,))
00151         self.server_thread.start()
00152 
00153     @staticmethod
00154     def server_thread_func(this):
00155         """
00156         Thread function to run TCP server forever.
00157 
00158         :param this:
00159         :return:
00160         """
00161         this.server.serve_forever()
00162 
00163     @event_callback("target_ip")
00164     def _callback_target_ip(self, key, value, timestamp):
00165         """
00166         Callback to handle reception of target's IP address.
00167 
00168         :param key:
00169         :param value:
00170         :param timestamp:
00171         :return:
00172         """
00173         self.target_ip = value
00174         self.SERVER_IP = self.find_interface_to_target_addr(self.target_ip)
00175         self.setup_tcp_server()
00176 
00177     @event_callback("host_ip")
00178     def _callback_host_ip(self, key, value, timestamp):
00179         """
00180         Callback for request for host IP Addr
00181 
00182         """
00183         self.send_kv("host_ip", self.SERVER_IP)
00184 
00185     @event_callback("host_port")
00186     def _callback_host_port(self, key, value, timestamp):
00187         """
00188         Callback for request for host port
00189         """
00190         self.send_kv("host_port", self.SERVER_PORT)
00191 
00192     def teardown(self):
00193         if self.server:
00194             self.server.shutdown()
00195             self.server_thread.join()