client/cros/netprotos/interface_host.py

# Copyright 2014 The Chromium OS 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 logging
import socket
import struct
import time

from autotest_lib.client.common_lib import error
from autotest_lib.client.common_lib.cros.network import interface


class InterfaceHost(object):
    """A host for use with ZeroconfDaemon that binds to an interface."""

    @property
    def ip_addr(self):
        """Get the IP address of the interface we're bound to."""
        return self._interface.ipv4_address


    def __init__(self, interface_name):
        self._interface = interface.Interface(interface_name)
        self._socket = None


    def close(self):
        """Close the underlying socket."""
        if self._socket:
            self._socket.close()


    def socket(self, family, sock_type):
        """Get a socket bound to this interface.

        Only supports IPv4 UDP sockets on broadcast addresses.

        @param family: must be socket.AF_INET.
        @param sock_type: must be socket.SOCK_DGRAM.

        """
        if family != socket.AF_INET or sock_type != socket.SOCK_DGRAM:
            raise error.TestError('InterfaceHost only understands UDP sockets.')
        if self._socket is not None:
            raise error.TestError('InterfaceHost only supports a single '
                                  'multicast socket.')

        self._socket = InterfaceDatagramSocket(self.ip_addr)
        return self._socket


    def run_until(self, predicate, timeout_seconds):
        """Handle traffic from our socket until |predicate|() is true.

        @param predicate: function without arguments that returns True or False.
        @param timeout_seconds: number of seconds to wait for predicate to
                                become True.
        @return: tuple(success, duration) where success is True iff predicate()
                 became true before |timeout_seconds| passed.

        """
        start_time = time.time()
        duration = lambda: time.time() - start_time
        while duration() < timeout_seconds:
            if predicate():
                return True, duration()
            # Assume this take non-trivial time, don't sleep here.
            self._socket.run_once()
        return False, duration()


class InterfaceDatagramSocket(object):
    """Broadcast UDP socket bound to a particular network interface."""

    # Wait for a UDP frame to appear for this long before timing out.
    TIMEOUT_VALUE_SECONDS = 0.5

    def __init__(self, interface_ip):
        """Construct an instance.

        @param interface_ip: string like '239.192.1.100'.

        """
        self._interface_ip = interface_ip
        self._recv_callback = None
        self._recv_sock = None
        self._send_sock = None


    def close(self):
        """Close state associated with this object."""
        if self._recv_sock is not None:
            # Closing the socket drops membership groups.
            self._recv_sock.close()
            self._recv_sock = None
        if self._send_sock is not None:
            self._send_sock.close()
            self._send_sock = None


    def listen(self, ip_addr, port, recv_callback):
        """Bind and listen on the ip_addr:port.

        @param ip_addr: Multicast group IP (e.g. '224.0.0.251')
        @param port: Local destination port number.
        @param recv_callback: A callback function that accepts three arguments,
                              the received string, the sender IPv4 address and
                              the sender port number.

        """
        if self._recv_callback is not None:
            raise error.TestError('listen() called twice on '
                                  'InterfaceDatagramSocket.')
        # Multicast addresses are in 224.0.0.0 - 239.255.255.255 (rfc5771)
        ip_addr_prefix = ord(socket.inet_aton(ip_addr)[0])
        if ip_addr_prefix < 224 or ip_addr_prefix > 239:
            raise error.TestError('Invalid multicast address.')

        self._recv_callback = recv_callback
        # Set up a socket to receive just traffic from the given address.
        self._recv_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self._recv_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._recv_sock.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP,
                                   socket.inet_aton(ip_addr) +
                                   socket.inet_aton(self._interface_ip))
        self._recv_sock.settimeout(self.TIMEOUT_VALUE_SECONDS)
        self._recv_sock.bind((ip_addr, port))
        # When we send responses, we want to send them from this particular
        # interface.  The easiest way to do this is bind a socket directly to
        # the IP for the interface.  We're going to ignore messages sent to this
        # socket.
        self._send_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self._send_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._send_sock.setsockopt(socket.SOL_IP, socket.IP_MULTICAST_TTL,
                                   struct.pack('b', 1))
        self._send_sock.bind((self._interface_ip, port))


    def run_once(self):
        """Receive pending frames if available, return after timeout otw."""
        if self._recv_sock is None:
            raise error.TestError('Must listen() on socket before recv\'ing.')
        BUFFER_SIZE_BYTES = 2048
        try:
            data, sender_addr = self._recv_sock.recvfrom(BUFFER_SIZE_BYTES)
        except socket.timeout:
            return
        if len(sender_addr) != 2:
            logging.error('Unexpected address: %r', sender_addr)
        self._recv_callback(data, *sender_addr)


    def send(self, data, ip_addr, port):
        """Send |data| to an IPv4 address.

        @param data: string of raw bytes to send.
        @param ip_addr: string like '239.192.1.100'.
        @param port: int like 50000.

        """
        self._send_sock.sendto(data, (ip_addr, port))