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
17import ctypes
18import ctypes.util
19import os
20import socket
21import struct
22import sys
23
24import cstruct
25
26
27# Data structures.
28# These aren't constants, they're classes. So, pylint: disable=invalid-name
29CMsgHdr = cstruct.Struct("cmsghdr", "@Lii", "len level type")
30Iovec = cstruct.Struct("iovec", "@LL", "base len")
31MsgHdr = cstruct.Struct("msghdr", "@LLLLLLi",
32                        "name namelen iov iovlen control msg_controllen flags")
33SockaddrIn = cstruct.Struct("sockaddr_in", "=HH4sxxxxxxxx", "family port addr")
34SockaddrIn6 = cstruct.Struct("sockaddr_in6", "=HHI16sI",
35                             "family port flowinfo addr scope_id")
36SockaddrStorage = cstruct.Struct("sockaddr_storage", "=H126s", "family data")
37SockExtendedErr = cstruct.Struct("sock_extended_err", "@IBBBxII",
38                                 "errno origin type code info data")
39InPktinfo = cstruct.Struct("in_pktinfo", "@i4s4s", "ifindex spec_dst addr")
40In6Pktinfo = cstruct.Struct("in6_pktinfo", "@16si", "addr ifindex")
41
42# Constants.
43# IPv4 socket options and cmsg types.
44IP_TTL = 2
45IP_MTU_DISCOVER = 10
46IP_PKTINFO = 8
47IP_RECVERR = 11
48IP_RECVTTL = 12
49IP_MTU = 14
50
51# IPv6 socket options and cmsg types.
52IPV6_MTU_DISCOVER = 23
53IPV6_RECVERR = 25
54IPV6_RECVPKTINFO = 49
55IPV6_PKTINFO = 50
56IPV6_RECVHOPLIMIT = 51
57IPV6_HOPLIMIT = 52
58IPV6_PATHMTU = 61
59IPV6_DONTFRAG = 62
60
61# PMTUD values.
62IP_PMTUDISC_DO = 1
63
64CMSG_ALIGNTO = struct.calcsize("@L")  # The kernel defines this as sizeof(long).
65
66# Sendmsg flags
67MSG_CONFIRM = 0X800
68MSG_ERRQUEUE = 0x2000
69
70# Linux errqueue API.
71SO_ORIGIN_ICMP = 2
72SO_ORIGIN_ICMP6 = 3
73
74# Find the C library.
75libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)
76
77
78# TODO: Move this to a utils.py or constants.py file, once we have one.
79def LinuxVersion():
80  # Example: "3.4.67-00753-gb7a556f".
81  # Get the part before the dash.
82  version = os.uname()[2].split("-")[0]
83  # Convert it into a tuple such as (3, 4, 67). That allows comparing versions
84  # using < and >, since tuples are compared lexicographically.
85  version = tuple(int(i) for i in version.split("."))
86  return version
87
88
89def PaddedLength(length):
90  return CMSG_ALIGNTO * ((length / CMSG_ALIGNTO) + (length % CMSG_ALIGNTO != 0))
91
92
93def MaybeRaiseSocketError(ret):
94  if ret < 0:
95    errno = ctypes.get_errno()
96    raise socket.error(errno, os.strerror(errno))
97
98
99def Sockaddr(addr):
100  if ":" in addr[0]:
101    family = socket.AF_INET6
102    if len(addr) == 4:
103      addr, port, flowinfo, scope_id = addr
104    else:
105      (addr, port), flowinfo, scope_id = addr, 0, 0
106    addr = socket.inet_pton(family, addr)
107    return SockaddrIn6((family, socket.ntohs(port), socket.ntohl(flowinfo),
108                        addr, scope_id))
109  else:
110    family = socket.AF_INET
111    addr, port = addr
112    addr = socket.inet_pton(family, addr)
113    return SockaddrIn((family, socket.ntohs(port), addr))
114
115
116def _MakeMsgControl(optlist):
117  """Creates a msg_control blob from a list of cmsg attributes.
118
119  Takes a list of cmsg attributes. Each attribute is a tuple of:
120   - level: An integer, e.g., SOL_IPV6.
121   - type: An integer, the option identifier, e.g., IPV6_HOPLIMIT.
122   - data: The option data. This is either a string or an integer. If it's an
123     integer it will be written as an unsigned integer in host byte order. If
124     it's a string, it's used as is.
125
126  Data is padded to an integer multiple of CMSG_ALIGNTO.
127
128  Args:
129    optlist: A list of tuples describing cmsg options.
130
131  Returns:
132    A string, a binary blob usable as the control data for a sendmsg call.
133
134  Raises:
135    TypeError: Option data is neither an integer nor a string.
136  """
137  msg_control = ""
138
139  for i, opt in enumerate(optlist):
140    msg_level, msg_type, data = opt
141    if isinstance(data, int):
142      data = struct.pack("=I", data)
143    elif not isinstance(data, str):
144      raise TypeError("unknown data type for opt %i: %s" % (i, type(data)))
145
146    datalen = len(data)
147    msg_len = len(CMsgHdr) + datalen
148    padding = "\x00" * (PaddedLength(datalen) - datalen)
149    msg_control += CMsgHdr((msg_len, msg_level, msg_type)).Pack()
150    msg_control += data + padding
151
152  return msg_control
153
154
155def _ParseMsgControl(buf):
156  """Parse a raw control buffer into a list of tuples."""
157  msglist = []
158  while len(buf) > 0:
159    cmsghdr, buf = cstruct.Read(buf, CMsgHdr)
160    datalen = cmsghdr.len - len(CMsgHdr)
161    data, buf = buf[:datalen], buf[PaddedLength(datalen):]
162
163    if cmsghdr.level == socket.IPPROTO_IP:
164      if cmsghdr.type == IP_PKTINFO:
165        data = InPktinfo(data)
166      elif cmsghdr.type == IP_TTL:
167        data = struct.unpack("@I", data)[0]
168
169    if cmsghdr.level == socket.IPPROTO_IPV6:
170      if cmsghdr.type == IPV6_PKTINFO:
171        data = In6Pktinfo(data)
172      elif cmsghdr.type == IPV6_RECVERR:
173        err, source = cstruct.Read(data, SockExtendedErr)
174        if err.origin == SO_ORIGIN_ICMP6:
175          source, pad = cstruct.Read(source, SockaddrIn6)
176        data = (err, source)
177      elif cmsghdr.type == IPV6_HOPLIMIT:
178        data = struct.unpack("@I", data)[0]
179
180    # If not, leave data as just the raw bytes.
181
182    msglist.append((cmsghdr.level, cmsghdr.type, data))
183
184  return msglist
185
186
187def Bind(s, to):
188  """Python wrapper for bind."""
189  ret = libc.bind(s.fileno(), to.CPointer(), len(to))
190  MaybeRaiseSocketError(ret)
191  return ret
192
193
194def Connect(s, to):
195  """Python wrapper for connect."""
196  ret = libc.connect(s.fileno(), to.CPointer(), len(to))
197  MaybeRaiseSocketError(ret)
198  return ret
199
200
201def Sendmsg(s, to, data, control, flags):
202  """Python wrapper for sendmsg.
203
204  Args:
205    s: A Python socket object. Becomes sockfd.
206    to: An address tuple, or a SockaddrIn[6] struct. Becomes msg->msg_name.
207    data: A string, the data to write. Goes into msg->msg_iov.
208    control: A list of cmsg options. Becomes msg->msg_control.
209    flags: An integer. Becomes msg->msg_flags.
210
211  Returns:
212    If sendmsg succeeds, returns the number of bytes written as an integer.
213
214  Raises:
215    socket.error: If sendmsg fails.
216  """
217  # Create ctypes buffers and pointers from our structures. We need to hang on
218  # to the underlying Python objects, because we don't want them to be garbage
219  # collected and freed while we have C pointers to them.
220
221  # Convert the destination address into a struct sockaddr.
222  if to:
223    if isinstance(to, tuple):
224      to = Sockaddr(to)
225    msg_name = to.CPointer()
226    msg_namelen = len(to)
227  else:
228    msg_name = 0
229    msg_namelen = 0
230
231  # Convert the data to a data buffer and a struct iovec pointing at it.
232  if data:
233    databuf = ctypes.create_string_buffer(data)
234    iov = Iovec((ctypes.addressof(databuf), len(data)))
235    msg_iov = iov.CPointer()
236    msg_iovlen = 1
237  else:
238    msg_iov = 0
239    msg_iovlen = 0
240
241  # Marshal the cmsg options.
242  if control:
243    control = _MakeMsgControl(control)
244    controlbuf = ctypes.create_string_buffer(control)
245    msg_control = ctypes.addressof(controlbuf)
246    msg_controllen = len(control)
247  else:
248    msg_control = 0
249    msg_controllen = 0
250
251  # Assemble the struct msghdr.
252  msghdr = MsgHdr((msg_name, msg_namelen, msg_iov, msg_iovlen,
253                   msg_control, msg_controllen, flags)).Pack()
254
255  # Call sendmsg.
256  ret = libc.sendmsg(s.fileno(), msghdr, 0)
257  MaybeRaiseSocketError(ret)
258
259  return ret
260
261
262def _ToSocketAddress(addr, alen):
263  addr = addr[:alen]
264
265  # Attempt to convert the address to something we understand.
266  if alen == 0:
267    return None
268  elif alen == len(SockaddrIn) and SockaddrIn(addr).family == socket.AF_INET:
269    return SockaddrIn(addr)
270  elif alen == len(SockaddrIn6) and SockaddrIn6(addr).family == socket.AF_INET6:
271    return SockaddrIn6(addr)
272  elif alen == len(SockaddrStorage):  # Can this ever happen?
273    return SockaddrStorage(addr)
274  else:
275    return addr  # Unknown or malformed. Return the raw bytes.
276
277
278def Recvmsg(s, buflen, controllen, flags, addrlen=len(SockaddrStorage)):
279  """Python wrapper for recvmsg.
280
281  Args:
282    s: A Python socket object. Becomes sockfd.
283    buflen: An integer, the maximum number of bytes to read.
284    addrlen: An integer, the maximum size of the source address.
285    controllen: An integer, the maximum size of the cmsg buffer.
286
287  Returns:
288    A tuple of received bytes, socket address tuple, and cmg list.
289
290  Raises:
291    socket.error: If recvmsg fails.
292  """
293  addr = ctypes.create_string_buffer(addrlen)
294  msg_name = ctypes.addressof(addr)
295  msg_namelen = addrlen
296
297  buf = ctypes.create_string_buffer(buflen)
298  iov = Iovec((ctypes.addressof(buf), buflen))
299  msg_iov = iov.CPointer()
300  msg_iovlen = 1
301
302  control = ctypes.create_string_buffer(controllen)
303  msg_control = ctypes.addressof(control)
304  msg_controllen = controllen
305
306  msghdr = MsgHdr((msg_name, msg_namelen, msg_iov, msg_iovlen,
307                   msg_control, msg_controllen, flags))
308  ret = libc.recvmsg(s.fileno(), msghdr.CPointer(), flags)
309  MaybeRaiseSocketError(ret)
310
311  data = buf.raw[:ret]
312  msghdr = MsgHdr(str(msghdr._buffer.raw))
313  addr = _ToSocketAddress(addr, msghdr.namelen)
314  control = control.raw[:msghdr.msg_controllen]
315  msglist = _ParseMsgControl(control)
316
317  return data, addr, msglist
318
319
320def Recvfrom(s, size, flags=0):
321  """Python wrapper for recvfrom."""
322  buf = ctypes.create_string_buffer(size)
323  addr = ctypes.create_string_buffer(len(SockaddrStorage))
324  alen = ctypes.c_int(len(addr))
325
326  ret = libc.recvfrom(s.fileno(), buf, len(buf), flags,
327                      addr, ctypes.byref(alen))
328  MaybeRaiseSocketError(ret)
329
330  data = buf[:ret]
331  alen = alen.value
332
333  addr = _ToSocketAddress(addr.raw, alen)
334
335  return data, addr
336