tests/net_test/csocket.py

   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