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
22
23import cstruct
24
25
26# Data structures.
27CMsgHdr = cstruct.Struct("cmsghdr", "@Lii", "len level type")
28Iovec = cstruct.Struct("iovec", "@LL", "base len")
29MsgHdr = cstruct.Struct("msghdr", "@LLLLLLi",
30                        "name namelen iov iovlen control msg_controllen flags")
31SockaddrIn = cstruct.Struct("sockaddr_in", "=HH4sxxxxxxxx", "family port addr")
32SockaddrIn6 = cstruct.Struct("sockaddr_in6", "=HHI16sI",
33                             "family port flowinfo addr scope_id")
34
35# Constants.
36CMSG_ALIGNTO = struct.calcsize("@L")  # The kernel defines this as sizeof(long).
37MSG_CONFIRM = 0X800
38
39# Find the C library.
40libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)
41
42
43def PaddedLength(length):
44  return CMSG_ALIGNTO * ((length / CMSG_ALIGNTO) + (length % CMSG_ALIGNTO != 0))
45
46
47def MaybeRaiseSocketError(ret):
48  if ret < 0:
49    errno = ctypes.get_errno()
50    raise socket.error(errno, os.strerror(errno))
51
52
53def 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
70def _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
109def 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
115def 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
122def 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