Home | History | Annotate | Download | only in net_test
      1 # Copyright 2014 The Android Open Source Project
      2 #
      3 # Licensed under the Apache License, Version 2.0 (the "License");
      4 # you may not use this file except in compliance with the License.
      5 # You may obtain a copy of the License at
      6 #
      7 # http://www.apache.org/licenses/LICENSE-2.0
      8 #
      9 # Unless required by applicable law or agreed to in writing, software
     10 # distributed under the License is distributed on an "AS IS" BASIS,
     11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     12 # See the License for the specific language governing permissions and
     13 # limitations under the License.
     14 
     15 """Python wrapper for C socket calls and data structures."""
     16 
     17 import ctypes
     18 import ctypes.util
     19 import os
     20 import socket
     21 import struct
     22 
     23 import cstruct
     24 
     25 
     26 # Data structures.
     27 CMsgHdr = cstruct.Struct("cmsghdr", "@Lii", "len level type")
     28 Iovec = cstruct.Struct("iovec", "@LL", "base len")
     29 MsgHdr = cstruct.Struct("msghdr", "@LLLLLLi",
     30                         "name namelen iov iovlen control msg_controllen flags")
     31 SockaddrIn = cstruct.Struct("sockaddr_in", "=HH4sxxxxxxxx", "family port addr")
     32 SockaddrIn6 = cstruct.Struct("sockaddr_in6", "=HHI16sI",
     33                              "family port flowinfo addr scope_id")
     34 
     35 # Constants.
     36 CMSG_ALIGNTO = struct.calcsize("@L")  # The kernel defines this as sizeof(long).
     37 MSG_CONFIRM = 0X800
     38 
     39 # Find the C library.
     40 libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)
     41 
     42 
     43 def PaddedLength(length):
     44   return CMSG_ALIGNTO * ((length / CMSG_ALIGNTO) + (length % CMSG_ALIGNTO != 0))
     45 
     46 
     47 def MaybeRaiseSocketError(ret):
     48   if ret < 0:
     49     errno = ctypes.get_errno()
     50     raise socket.error(errno, os.strerror(errno))
     51 
     52 
     53 def Sockaddr(addr):
     54   if ":" in addr[0]:
     55     family = socket.AF_INET6
     56     if len(addr) == 4:
     57       addr, port, flowinfo, scope_id = addr
     58     else:
     59       (addr, port), flowinfo, scope_id = addr, 0, 0
     60     addr = socket.inet_pton(family, addr)
     61     return SockaddrIn6((family, socket.ntohs(port), socket.ntohl(flowinfo),
     62                         addr, scope_id))
     63   else:
     64     family = socket.AF_INET
     65     addr, port = addr
     66     addr = socket.inet_pton(family, addr)
     67     return SockaddrIn((family, socket.ntohs(port), addr))
     68 
     69 
     70 def _MakeMsgControl(optlist):
     71   """Creates a msg_control blob from a list of cmsg attributes.
     72 
     73   Takes a list of cmsg attributes. Each attribute is a tuple of:
     74    - level: An integer, e.g., SOL_IPV6.
     75    - type: An integer, the option identifier, e.g., IPV6_HOPLIMIT.
     76    - data: The option data. This is either a string or an integer. If it's an
     77      integer it will be written as an unsigned integer in host byte order. If
     78      it's a string, it's used as is.
     79 
     80   Data is padded to an integer multiple of CMSG_ALIGNTO.
     81 
     82   Args:
     83     optlist: A list of tuples describing cmsg options.
     84 
     85   Returns:
     86     A string, a binary blob usable as the control data for a sendmsg call.
     87 
     88   Raises:
     89     TypeError: Option data is neither an integer nor a string.
     90   """
     91   msg_control = ""
     92 
     93   for i, opt in enumerate(optlist):
     94     msg_level, msg_type, data = opt
     95     if isinstance(data, int):
     96       data = struct.pack("=I", data)
     97     elif not isinstance(data, str):
     98       raise TypeError("unknown data type for opt %i: %s" % (i, type(data)))
     99 
    100     datalen = len(data)
    101     msg_len = len(CMsgHdr) + datalen
    102     padding = "\x00" * (PaddedLength(datalen) - datalen)
    103     msg_control += CMsgHdr((msg_len, msg_level, msg_type)).Pack()
    104     msg_control += data + padding
    105 
    106   return msg_control
    107 
    108 
    109 def Bind(s, to):
    110   """Python wrapper for connect."""
    111   ret = libc.bind(s.fileno(), to.CPointer(), len(to))
    112   MaybeRaiseSocketError(ret)
    113   return ret
    114 
    115 def Connect(s, to):
    116   """Python wrapper for connect."""
    117   ret = libc.connect(s.fileno(), to.CPointer(), len(to))
    118   MaybeRaiseSocketError(ret)
    119   return ret
    120 
    121 
    122 def Sendmsg(s, to, data, control, flags):
    123   """Python wrapper for sendmsg.
    124 
    125   Args:
    126     s: A Python socket object. Becomes sockfd.
    127     to: An address tuple, or a SockaddrIn[6] struct. Becomes msg->msg_name.
    128     data: A string, the data to write. Goes into msg->msg_iov.
    129     control: A list of cmsg options. Becomes msg->msg_control.
    130     flags: An integer. Becomes msg->msg_flags.
    131 
    132   Returns:
    133     If sendmsg succeeds, returns the number of bytes written as an integer.
    134 
    135   Raises:
    136     socket.error: If sendmsg fails.
    137   """
    138   # Create ctypes buffers and pointers from our structures. We need to hang on
    139   # to the underlying Python objects, because we don't want them to be garbage
    140   # collected and freed while we have C pointers to them.
    141 
    142   # Convert the destination address into a struct sockaddr.
    143   if to:
    144     if isinstance(to, tuple):
    145       to = Sockaddr(to)
    146     msg_name = to.CPointer()
    147     msg_namelen = len(to)
    148   else:
    149     msg_name = 0
    150     msg_namelen = 0
    151 
    152   # Convert the data to a data buffer and a struct iovec pointing at it.
    153   if data:
    154     databuf = ctypes.create_string_buffer(data)
    155     iov = Iovec((ctypes.addressof(databuf), len(data)))
    156     msg_iov = iov.CPointer()
    157     msg_iovlen = 1
    158   else:
    159     msg_iov = 0
    160     msg_iovlen = 0
    161 
    162   # Marshal the cmsg options.
    163   if control:
    164     control = _MakeMsgControl(control)
    165     controlbuf = ctypes.create_string_buffer(control)
    166     msg_control = ctypes.addressof(controlbuf)
    167     msg_controllen = len(control)
    168   else:
    169     msg_control = 0
    170     msg_controllen = 0
    171 
    172   # Assemble the struct msghdr.
    173   msghdr = MsgHdr((msg_name, msg_namelen, msg_iov, msg_iovlen,
    174                    msg_control, msg_controllen, flags)).Pack()
    175 
    176   # Call sendmsg.
    177   ret = libc.sendmsg(s.fileno(), msghdr, 0)
    178   MaybeRaiseSocketError(ret)
    179 
    180   return ret
    181