tcp_echo_client.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         print ("HOST: TCPEchoClient_Handler: Connection received...")
00037         while self.server.isrunning():
00038             try:
00039                 data = self.recv()
00040                 if not data: break
00041             except Exception as e:
00042                 print ('HOST: TCPEchoClient_Handler recv error: %s' % str(e))
00043                 break
00044 
00045             print ('HOST: TCPEchoClient_Handler: Rx: \n%s\n' % data)
00046 
00047             try:
00048                 # echo data back to the client
00049                 self.send(data)
00050             except Exception as e:
00051                 print ('HOST: TCPEchoClient_Handler send error: %s' % str(e))
00052                 break
00053         print 'Connection finished'
00054 
00055     def recv(self):
00056         """
00057         Try to receive until server is shutdown
00058         """
00059         while self.server.isrunning():
00060             rl, wl, xl = select.select([self.request], [], [], 1)
00061             if len(rl):
00062                 return self.request.recv(1024)
00063 
00064     def send(self, data):
00065         """
00066         Try to send until server is shutdown
00067         """
00068         while self.server.isrunning():
00069             rl, wl, xl = select.select([], [self.request], [], 1)
00070             if len(wl):
00071                 self.request.sendall(data)
00072                 break
00073 
00074 
00075 class TCPServerWrapper (TCPServer):
00076     """
00077     Wrapper over TCP server to implement server initiated shutdown.
00078     Adds a flag:= running that a request handler can check and come out of
00079     recv loop when shutdown is called.
00080     """
00081 
00082     def __init__(self, addr, request_handler):
00083         # hmm, TCPServer is not sub-classed from object!
00084         if issubclass(TCPServer, object):
00085             super(TCPServerWrapper, self).__init__(addr, request_handler)
00086         else:
00087             TCPServer.__init__(self, addr, request_handler)
00088         self.running  = False
00089 
00090     def serve_forever(self):
00091         self.running  = True
00092         if issubclass(TCPServer, object):
00093             super(TCPServerWrapper, self).serve_forever()
00094         else:
00095             TCPServer.serve_forever(self)
00096 
00097     def shutdown(self):
00098         self.running  = False
00099         if issubclass(TCPServer, object):
00100             super(TCPServerWrapper, self).shutdown()
00101         else:
00102             TCPServer.shutdown(self)
00103 
00104     def isrunning(self):
00105         return self.running 
00106 
00107 
00108 class TCPEchoClientTest(BaseHostTest):
00109 
00110     def __init__(self):
00111         """
00112         Initialise test parameters.
00113 
00114         :return:
00115         """
00116         BaseHostTest.__init__(self)
00117         self.SERVER_IP = None # Will be determined after knowing the target IP
00118         self.SERVER_PORT = 0  # Let TCPServer choose an arbitrary port
00119         self.server = None
00120         self.server_thread = None
00121         self.target_ip = None
00122 
00123     @staticmethod
00124     def find_interface_to_target_addr(target_ip):
00125         """
00126         Finds IP address of the interface through which it is connected to the target.
00127 
00128         :return:
00129         """
00130         s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
00131         s.connect((target_ip, 0)) # Target IP, Any port
00132         ip = s.getsockname()[0]
00133         s.close()
00134         return ip
00135 
00136     def setup_tcp_server(self):
00137         """
00138         sets up a TCP server for target to connect and send test data.
00139 
00140         :return:
00141         """
00142         # !NOTE: There should mechanism to assert in the host test
00143         if self.SERVER_IP is None:
00144             self.log("setup_tcp_server() called before determining server IP!")
00145             self.notify_complete(False)
00146 
00147         # Returning none will suppress host test from printing success code
00148         self.server = TCPServerWrapper((self.SERVER_IP, self.SERVER_PORT), TCPEchoClientHandler)
00149         ip, port = self.server.server_address
00150         self.SERVER_PORT = port
00151         self.server.allow_reuse_address = True
00152         self.log("HOST: Listening for TCP connections: " + self.SERVER_IP + ":" + str(self.SERVER_PORT))
00153         self.server_thread = Thread(target=TCPEchoClientTest.server_thread_func, args=(self,))
00154         self.server_thread.start()
00155 
00156     @staticmethod
00157     def server_thread_func(this):
00158         """
00159         Thread function to run TCP server forever.
00160 
00161         :param this:
00162         :return:
00163         """
00164         this.server.serve_forever()
00165 
00166     @event_callback("target_ip")
00167     def _callback_target_ip(self, key, value, timestamp):
00168         """
00169         Callback to handle reception of target's IP address.
00170 
00171         :param key:
00172         :param value:
00173         :param timestamp:
00174         :return:
00175         """
00176         self.target_ip = value
00177         self.SERVER_IP = self.find_interface_to_target_addr(self.target_ip)
00178         self.setup_tcp_server()
00179 
00180     @event_callback("host_ip")
00181     def _callback_host_ip(self, key, value, timestamp):
00182         """
00183         Callback for request for host IP Addr
00184 
00185         """
00186         self.send_kv("host_ip", self.SERVER_IP)
00187 
00188     @event_callback("host_port")
00189     def _callback_host_port(self, key, value, timestamp):
00190         """
00191         Callback for request for host port
00192         """
00193         self.send_kv("host_port", self.SERVER_PORT)
00194 
00195     def teardown(self):
00196         if self.server:
00197             self.server.shutdown()
00198             self.server_thread.join()