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