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