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