140 lines
3.8 KiB
Cython
140 lines
3.8 KiB
Cython
# Copyright (C) 2016-present the asyncpg authors and contributors
|
|
# <see AUTHORS file>
|
|
#
|
|
# This module is part of asyncpg and is released under
|
|
# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
|
|
import ipaddress
|
|
|
|
|
|
# defined in postgresql/src/include/inet.h
|
|
#
|
|
DEF PGSQL_AF_INET = 2 # AF_INET
|
|
DEF PGSQL_AF_INET6 = 3 # AF_INET + 1
|
|
|
|
|
|
_ipaddr = ipaddress.ip_address
|
|
_ipiface = ipaddress.ip_interface
|
|
_ipnet = ipaddress.ip_network
|
|
|
|
|
|
cdef inline uint8_t _ip_max_prefix_len(int32_t family):
|
|
# Maximum number of bits in the network prefix of the specified
|
|
# IP protocol version.
|
|
if family == PGSQL_AF_INET:
|
|
return 32
|
|
else:
|
|
return 128
|
|
|
|
|
|
cdef inline int32_t _ip_addr_len(int32_t family):
|
|
# Length of address in bytes for the specified IP protocol version.
|
|
if family == PGSQL_AF_INET:
|
|
return 4
|
|
else:
|
|
return 16
|
|
|
|
|
|
cdef inline int8_t _ver_to_family(int32_t version):
|
|
if version == 4:
|
|
return PGSQL_AF_INET
|
|
else:
|
|
return PGSQL_AF_INET6
|
|
|
|
|
|
cdef inline _net_encode(WriteBuffer buf, int8_t family, uint32_t bits,
|
|
int8_t is_cidr, bytes addr):
|
|
|
|
cdef:
|
|
char *addrbytes
|
|
ssize_t addrlen
|
|
|
|
cpython.PyBytes_AsStringAndSize(addr, &addrbytes, &addrlen)
|
|
|
|
buf.write_int32(4 + <int32_t>addrlen)
|
|
buf.write_byte(family)
|
|
buf.write_byte(<int8_t>bits)
|
|
buf.write_byte(is_cidr)
|
|
buf.write_byte(<int8_t>addrlen)
|
|
buf.write_cstr(addrbytes, addrlen)
|
|
|
|
|
|
cdef net_decode(CodecContext settings, FRBuffer *buf, bint as_cidr):
|
|
cdef:
|
|
int32_t family = <int32_t>frb_read(buf, 1)[0]
|
|
uint8_t bits = <uint8_t>frb_read(buf, 1)[0]
|
|
int prefix_len
|
|
int32_t is_cidr = <int32_t>frb_read(buf, 1)[0]
|
|
int32_t addrlen = <int32_t>frb_read(buf, 1)[0]
|
|
bytes addr
|
|
uint8_t max_prefix_len = _ip_max_prefix_len(family)
|
|
|
|
if is_cidr != as_cidr:
|
|
raise ValueError('unexpected CIDR flag set in non-cidr value')
|
|
|
|
if family != PGSQL_AF_INET and family != PGSQL_AF_INET6:
|
|
raise ValueError('invalid address family in "{}" value'.format(
|
|
'cidr' if is_cidr else 'inet'
|
|
))
|
|
|
|
max_prefix_len = _ip_max_prefix_len(family)
|
|
|
|
if bits > max_prefix_len:
|
|
raise ValueError('invalid network prefix length in "{}" value'.format(
|
|
'cidr' if is_cidr else 'inet'
|
|
))
|
|
|
|
if addrlen != _ip_addr_len(family):
|
|
raise ValueError('invalid address length in "{}" value'.format(
|
|
'cidr' if is_cidr else 'inet'
|
|
))
|
|
|
|
addr = cpython.PyBytes_FromStringAndSize(frb_read(buf, addrlen), addrlen)
|
|
|
|
if as_cidr or bits != max_prefix_len:
|
|
prefix_len = cpython.PyLong_FromLong(bits)
|
|
|
|
if as_cidr:
|
|
return _ipnet((addr, prefix_len))
|
|
else:
|
|
return _ipiface((addr, prefix_len))
|
|
else:
|
|
return _ipaddr(addr)
|
|
|
|
|
|
cdef cidr_encode(CodecContext settings, WriteBuffer buf, obj):
|
|
cdef:
|
|
object ipnet
|
|
int8_t family
|
|
|
|
ipnet = _ipnet(obj)
|
|
family = _ver_to_family(ipnet.version)
|
|
_net_encode(buf, family, ipnet.prefixlen, 1, ipnet.network_address.packed)
|
|
|
|
|
|
cdef cidr_decode(CodecContext settings, FRBuffer *buf):
|
|
return net_decode(settings, buf, True)
|
|
|
|
|
|
cdef inet_encode(CodecContext settings, WriteBuffer buf, obj):
|
|
cdef:
|
|
object ipaddr
|
|
int8_t family
|
|
|
|
try:
|
|
ipaddr = _ipaddr(obj)
|
|
except ValueError:
|
|
# PostgreSQL accepts *both* CIDR and host values
|
|
# for the host datatype.
|
|
ipaddr = _ipiface(obj)
|
|
family = _ver_to_family(ipaddr.version)
|
|
_net_encode(buf, family, ipaddr.network.prefixlen, 1, ipaddr.packed)
|
|
else:
|
|
family = _ver_to_family(ipaddr.version)
|
|
_net_encode(buf, family, _ip_max_prefix_len(family), 0, ipaddr.packed)
|
|
|
|
|
|
cdef inet_decode(CodecContext settings, FRBuffer *buf):
|
|
return net_decode(settings, buf, False)
|