openrvdas.logger.readers.udp_reader

  1#!/usr/bin/env python3
  2
  3import logging
  4import socket
  5
  6from logger.readers.reader import Reader  # noqa: E402
  7
  8# The UDP header's `length` field sets a theoretical limit of 65,535 bytes
  9# (8-byte header + 65,527 bytes of data) for a UDP datagram.  Technically, IPV4
 10# or IPv6 headers use up some of that size, so actual maximum data sent per
 11# datagram is slightly less.
 12#
 13# UDP receivers should always request the max, though, because if you request
 14# less than what's on the wire, you get what you asked for and the rest gets
 15# tossed on the floor.  There's no built-in error detection/correction in UDP,
 16# so that would mess things up pretty good.
 17READ_BUFFER_SIZE = 65535
 18
 19# On the send side of things, we can (and do) detect when a write has failed
 20# because the user's `record` was too big.  This is usually because the user is
 21# trying to send huge datagrams after looking up the "theoretical max size of a
 22# datagram" on wikipedia.  Well, it's called "theoretical" for a reason.  It's
 23# really the maximum size of the data portion of a UDP datagram, but that
 24# doesn't take into account extra header for IPv4/IPv6 or seemingly random
 25# system-level caps (e.g., Mac's socket implementation set maximum udp send
 26# size to 9K).
 27#
 28# When UDPWriter detects this condition, it fragments the record into smaller
 29# records and appends each fragment with this FRAGMENT_MARKER.  Inside
 30# UDPReader.read(), we check to see if a received datagram ends with this
 31# marker, and if it does, we read another datagram and combine the results
 32# (over and over until we get a datagram that doesn't end with the marker).
 33FRAGMENT_MARKER = b'\xff\xffTOOBIG\xff\xff'
 34
 35
 36################################################################################
 37class UDPReader(Reader):
 38    """Read UDP packets from network."""
 39    ############################
 40    def __init__(self, interface=None, port=None, mc_group=None,
 41                 reuseaddr=False, reuseport=False, eol=None,
 42                 allow_empty=False, this_is_a_test=False, **kwargs):
 43        """
 44        ```
 45        interface    IP (or resolvable name) of interface to listen on.  None or ''
 46                     will listen on INADDR_ANY (all interfaces).  If joining a
 47                     multicast group and None or '' specified, this will default
 48                     to whatever the system's hostname resolves to.  This IP should
 49                     not be on the loopback network (OK for testing, but won't work
 50                     in the real world).
 51
 52        port         Port to listen to for packets.  REQUIRED
 53
 54        mc_group     If specified, IP address of multicast group id to subscribe to.
 55
 56        reuseaddr    Specifies wether we set SO_REUSEADDR on the created socket.  If
 57                     you don't know you need this, don't enable it.
 58
 59        reuseport    Specifies wether we set SO_REUSEPORT on the created socket.  If
 60                     you don't know you need this, don't enable it.
 61
 62        eol          split the record by the eol character if present.
 63
 64        allow_empty - If True, preserve and return empty records
 65
 66        this_is_a_test - If True, recognize that this is being called in a unittest, so
 67                don't output warnings about not using loopback addresses.
 68        ```
 69        """
 70        super().__init__(**kwargs)
 71
 72        if interface:
 73            # resolve once in constructor
 74            interface = socket.gethostbyname(interface)
 75        else:
 76            interface = ''
 77        self.interface = interface
 78
 79        # make sure user passed in `port`
 80        #
 81        # NOTE: We want the order of the arguments to consistently be (ip,
 82        #       port, ...) across all the network readers/writers... but we
 83        #       want `interface` to be optional.  All kwargs need to come after
 84        #       all regular args, so we've assigned a default value of None to
 85        #       `port`.  But don't be confused, it is REQUIRED.
 86        #
 87        if not port:
 88            raise TypeError('must specify `port`')
 89        # make sure port gets stored as an int, even if passed in as a string
 90        self.port = int(port)
 91
 92        # prep multicast parameters
 93        if mc_group:
 94            # resolve once in constructor
 95            mc_group = socket.gethostbyname(mc_group)
 96            if not interface:
 97                # multicast needs to specify interface to use, so let's pick a
 98                # sane default
 99                #
100                # NOTE: This means hostname cannot be an alias to localhost, or
101                #       you won't be able to send IGMP packets correctly.
102                #
103                self.interface = socket.gethostbyname(socket.gethostname())
104
105        self.mc_group = mc_group
106
107        self.reuseaddr = reuseaddr
108        self.reuseport = reuseport
109
110        self.eol = eol
111        self.allow_empty = allow_empty
112
113        self.this_is_a_test = this_is_a_test
114
115        # socket gets initialized on-demand in read()
116        self.socket = None
117
118    ############################
119    def __del__(self):
120        try:
121            if self.socket:
122                self.socket.close()
123        except AttributeError:
124            pass
125
126    ############################
127    def _open_socket(self):
128        """Do socket prep so we're ready to read().  Returns socket object or None on
129        failure.
130        """
131        sock = socket.socket(family=socket.AF_INET,
132                             type=socket.SOCK_DGRAM,
133                             proto=socket.IPPROTO_UDP)
134        if self.reuseaddr:
135            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, True)
136        if self.reuseport:
137            try:  # Raspbian doesn't recognize SO_REUSEPORT
138                sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, True)
139            except AttributeError:
140                logging.warning('Unable to set socket REUSEPORT; may be unsupported.')
141
142        # If mc_group is specified, subscribe to it as a multicast group
143        if self.mc_group:
144            # set outgoing multicast interface
145            #
146            # NOTE: Can't use loopback device for this, otherwise IGMP packets
147            #       never leave the system, and you never actually join the
148            #       group.
149            #
150            if self.interface.startswith('127.') and not self.this_is_a_test:
151                logging.warning("Can't use loopback device for joining multicast groups.  Make "
152                                "sure your system's hostname does NOT resolve to something in "
153                                "the 127.0.0.0/8 address block (e.g., localhost, 127.0.0.1), or "
154                                "specify the interface to use by passing its IP address as the "
155                                "`interface` parameter.  (You can ignore this message if you're "
156                                "actually just doing loopback testing.)")
157            sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_IF,
158                            socket.inet_aton(self.interface))
159
160            # join the group via IGMP
161            #
162            # NOTE: Since these are both already encoded as binary by
163            #       inet_aton(), we can just concatenate them.  Alternatively,
164            #       could use struct.pack("4s4s", ...) to create a struct to
165            #       pass into setsockopt()
166            #
167            sock.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP,
168                            socket.inet_aton(self.mc_group) + socket.inet_aton(self.interface))
169
170            # bind to mc_group:port
171            sock.bind((self.mc_group, self.port))
172
173        else:
174            # broadcast or unicast, bind to specificed interface
175            sock.bind((self.interface, self.port))
176
177        return sock
178
179    ############################
180    def read(self):
181        """
182        Read the next UDP packet.
183        """
184        # If socket isn't ready, set it up.  If something fails, return w/out reading.
185        if not self.socket:
186            self.socket = self._open_socket()
187        if not self.socket:
188            logging.error('UDPReader.read: unable to open UDP socket')
189            return
190
191        # Read datagrams until we get one that doesn't end with a FRAGMENT_MARKER
192        record_buffer = b''
193        while True:
194            try:
195                record = self.socket.recv(READ_BUFFER_SIZE)
196            except OSError as e:
197                logging.error('UDPReader error: %s', str(e))
198                return None
199            logging.debug('UDPReader.read: received %d bytes', len(record))
200
201            if record.endswith(FRAGMENT_MARKER):
202                # UDPWriter fragmented this record because it was too large to
203                # send as a single datagram
204                logging.info('UDPrader.read: detected fragmented packet')
205                record_buffer += record.rsplit(FRAGMENT_MARKER, maxsplit=1)[0]
206                logging.debug('record_buffer: %s', record_buffer)
207            else:
208                record_buffer += record
209                break
210
211        # we've got a whole record in our record_buffer, decode it
212
213        # if eol == None, return the record as is
214        if not self.eol:
215            return self._decode_bytes(record_buffer, self.allow_empty)
216
217        # otherwise split the record by the eol
218        decoded_records = self._decode_bytes(record_buffer,
219                                             self.allow_empty).rstrip(self.eol).split(self.eol)
220
221        # if there was only one record, return just the first element in the
222        # list, otherwise return the whole list.
223        return decoded_records[0] if len(decoded_records) == 1 else decoded_records