Stabilization of dbwire

This commit is contained in:
Miroslav Štampar
2026-08-07 18:33:51 +02:00
parent 6acc970801
commit be881ea65e
10 changed files with 278 additions and 56 deletions
+41
View File
@@ -48,3 +48,44 @@ class InternalError(DatabaseError):
class NotSupportedError(DatabaseError):
pass
def http_origin(host, port):
"""
'http://host:port', with a literal IPv6 address bracketed as RFC 3986 requires. Without the brackets
the colons in the address are parsed as the port separator and the URL is simply wrong.
"""
host = host or "localhost"
if ":" in host and not host.startswith("["):
host = "[%s]" % host
return "http://%s:%d" % (host, int(port))
def connection_lost(ex):
"""
Turn a raw socket/OS failure into the DB-API hierarchy above.
Callers of a PEP 249 driver only ever catch Error and its subclasses, so a bare socket.error escaping
from a send/recv leaves them with an unhandled traceback instead of a handled connection failure.
"""
return OperationalError("connection lost (%s)" % ex)
def keepalive(sock):
"""
Ask the kernel to probe an idle connection, so a peer that dies without a FIN is eventually detected.
Deliberately NOT a read timeout: a legitimate query can take minutes on a big table, and a fixed
deadline would kill it. Keepalive distinguishes a dead peer from a slow one, which is the actual
failure being guarded against. Best-effort - the options are not portable everywhere.
"""
import socket as _socket
try:
sock.setsockopt(_socket.SOL_SOCKET, _socket.SO_KEEPALIVE, 1)
for name, value in (("TCP_KEEPIDLE", 60), ("TCP_KEEPINTVL", 10), ("TCP_KEEPCNT", 5)):
if hasattr(_socket, name):
sock.setsockopt(_socket.IPPROTO_TCP, getattr(_socket, name), value)
except Exception:
pass
+5 -1
View File
@@ -19,10 +19,13 @@ import socket
try:
from urllib.request import Request, urlopen # Python 3
from urllib.error import HTTPError, URLError
from urllib.parse import quote
except ImportError:
from urllib2 import Request, urlopen, HTTPError, URLError # Python 2
from urllib import quote
from extra.dbwire import OperationalError
from extra.dbwire import http_origin
from extra.dbwire import ProgrammingError
# TabSeparated backslash escapes -> the literal byte they denote
@@ -87,7 +90,8 @@ class Cursor(object):
class Connection(object):
def __init__(self, host, port, user, password, database, timeout):
self._url = "http://%s:%d/?database=%s&default_format=TabSeparatedWithNames" % (host, port, database or "default")
# quote the database: a name with a reserved character would otherwise inject into the query string
self._url = "%s/?database=%s&default_format=TabSeparatedWithNames" % (http_origin(host, port), quote(database or "default", safe=""))
self._headers = {}
if user or password:
token = base64.b64encode(("%s:%s" % (user or "", password or "")).encode("utf-8")).decode("ascii")
+23 -3
View File
@@ -23,6 +23,8 @@ from extra.dbwire import InterfaceError
from extra.dbwire import IntegrityError
from extra.dbwire import NotSupportedError
from extra.dbwire import OperationalError
from extra.dbwire import connection_lost
from extra.dbwire import keepalive
from extra.dbwire import ProgrammingError
_MAGIC = b"CUBRK"
@@ -257,7 +259,10 @@ class Connection(object):
def _recvn(self, n):
buf = b""
while len(buf) < n:
chunk = self._sock.recv(n - len(buf))
try:
chunk = self._sock.recv(n - len(buf))
except (socket.error, OSError) as ex:
raise connection_lost(ex)
if not chunk:
raise InterfaceError("connection closed by server")
buf += chunk
@@ -267,6 +272,7 @@ class Connection(object):
# broker handshake (may redirect to a dedicated CAS worker port), then cleartext OPEN_DATABASE login
try:
sock = socket.create_connection((self._host, self._port), timeout=self._timeout)
keepalive(sock)
sock.settimeout(None)
sock.sendall(_MAGIC + struct.pack(">BB", _CLIENT_JDBC, _CAS_VERSION) + b"\x00\x00\x00")
self._sock = sock
@@ -276,6 +282,7 @@ class Connection(object):
if port > 0: # redirected to a CAS worker: reconnect there, no second handshake
self._safe_close()
sock = socket.create_connection((self._host, port), timeout=self._timeout)
keepalive(sock)
sock.settimeout(None)
self._sock = sock
except (socket.error, socket.timeout) as ex:
@@ -299,7 +306,10 @@ class Connection(object):
def _send(self, payload):
# frame: [payload_len(4)][cas_info(4)][payload]
self._sock.sendall(struct.pack(">i", len(payload)) + self._cas_info + payload)
try:
self._sock.sendall(struct.pack(">i", len(payload)) + self._cas_info + payload)
except (socket.error, OSError) as ex:
raise connection_lost(ex)
def _read_response(self):
(data_length,) = struct.unpack(">i", self._recvn(4))
@@ -344,6 +354,17 @@ class Connection(object):
def _query(self, query):
reader = self._call(_Writer(_FC_PREPARE).arg_nts(query).arg_byte(0).arg_byte(0))
handle = reader.int()
try:
return self._execute(handle, reader)
finally:
# release the broker-side request handle even when execute/fetch raised: a connection that
# survives a few failed statements would otherwise hold every one of them until it closes
try:
self._call(_Writer(_FC_CLOSE_REQ_HANDLE).arg_int(handle))
except Exception:
pass
def _execute(self, handle, reader):
reader.int() # result cache lifetime
stmt_type = reader.byte()
reader.int() # bind count
@@ -374,7 +395,6 @@ class Connection(object):
rows += self._fetch_remaining(handle, columns, len(rows), total)
elif result_infos:
rowcount = result_infos[0]
self._call(_Writer(_FC_CLOSE_REQ_HANDLE).arg_int(handle))
return description, rows, rowcount
def _fetch_remaining(self, handle, columns, fetched, total):
+39 -17
View File
@@ -26,6 +26,8 @@ from extra.dbwire import IntegrityError
from extra.dbwire import InterfaceError
from extra.dbwire import NotSupportedError
from extra.dbwire import OperationalError
from extra.dbwire import connection_lost
from extra.dbwire import keepalive
# operation codes
_op_connect = 1
@@ -286,12 +288,18 @@ class _Wire(object):
self._rc, self._wc = rc, wc
def send(self, data):
self._sock.sendall(self._wc.translate(data) if self._wc else data)
try:
self._sock.sendall(self._wc.translate(data) if self._wc else data)
except (socket.error, OSError) as ex:
raise connection_lost(ex)
def _recv_raw(self, n):
buf = b""
while len(buf) < n:
chunk = self._sock.recv(n - len(buf))
try:
chunk = self._sock.recv(n - len(buf))
except (socket.error, OSError) as ex:
raise connection_lost(ex)
if not chunk:
raise InterfaceError("connection closed by server")
buf += chunk
@@ -494,23 +502,36 @@ class Connection(object):
self._send(_pack_int(_op_allocate_statement) + _pack_int(self._db_handle))
stmt = self._response()[0]
desc_items = bytes(bytearray([_isc_info_sql_stmt_type])) + _INFO_SQL_SELECT_DESCRIBE_VARS
self._send(_pack_int(_op_prepare_statement) + _pack_int(self._trans_handle) + _pack_int(stmt) +
_pack_int(3) + _pack_bytes(qbytes) + _pack_bytes(desc_items) + _pack_int(1024))
buf = self._response()[2]
stmt_type, columns = self._parse_describe(stmt, buf)
try:
desc_items = bytes(bytearray([_isc_info_sql_stmt_type])) + _INFO_SQL_SELECT_DESCRIBE_VARS
self._send(_pack_int(_op_prepare_statement) + _pack_int(self._trans_handle) + _pack_int(stmt) +
_pack_int(3) + _pack_bytes(qbytes) + _pack_bytes(desc_items) + _pack_int(1024))
buf = self._response()[2]
stmt_type, columns = self._parse_describe(stmt, buf)
exec_msg = (_pack_int(_op_execute) + _pack_int(stmt) + _pack_int(self._trans_handle) +
_pack_bytes(b"") + _pack_int(0) + _pack_int(0) + _pack_int(0))
self._send(exec_msg)
self._response()
exec_msg = (_pack_int(_op_execute) + _pack_int(stmt) + _pack_int(self._trans_handle) +
_pack_bytes(b"") + _pack_int(0) + _pack_int(0) + _pack_int(0))
self._send(exec_msg)
self._response()
description, rows = None, []
if stmt_type == _isc_info_sql_stmt_select and columns:
description = [(c.name, c.sqltype, None, None, None, None, None) for c in columns]
rows = self._fetch(stmt, columns)
self._send(_pack_int(_op_free_statement) + _pack_int(stmt) + _pack_int(_DSQL_drop))
self._response()
description, rows = None, []
if stmt_type == _isc_info_sql_stmt_select and columns:
description = [(c.name, c.sqltype, None, None, None, None, None) for c in columns]
rows = self._fetch(stmt, columns)
finally:
# release the server-side handle even when prepare/execute/fetch raised: a connection that
# survives a few failed statements would otherwise hold every one of them until it detaches
try:
self._send(_pack_int(_op_free_statement) + _pack_int(stmt) + _pack_int(_DSQL_drop))
self._response()
except Exception:
pass
# dbwire statements are autonomous (see README), but Firebird has no auto-commit mode: everything
# runs inside the one transaction opened at attach. Without this, DML is lost when the caller
# closes the connection without an explicit commit(). commit-retaining makes the work durable and
# keeps the transaction handle valid, so the connection stays usable for the next statement.
self.commit()
return description, rows
def _parse_describe(self, stmt, buf):
@@ -737,6 +758,7 @@ def connect(host=None, port=3050, user=None, password=None, database=None, conne
try:
sock = socket.create_connection((host or "localhost", int(port or 3050)), timeout=connect_timeout)
keepalive(sock)
sock.settimeout(None)
except (socket.error, socket.timeout) as ex:
raise OperationalError("could not connect to '%s:%s' (%s)" % (host, port, ex))
+12 -2
View File
@@ -22,6 +22,8 @@ import struct
from extra.dbwire import InterfaceError
from extra.dbwire import NotSupportedError
from extra.dbwire import OperationalError
from extra.dbwire import connection_lost
from extra.dbwire import keepalive
from extra.dbwire import ProgrammingError
_MAX_BLOCK = 0xffff >> 1
@@ -29,7 +31,10 @@ _MAX_BLOCK = 0xffff >> 1
def _recvn(sock, n):
buf = b""
while len(buf) < n:
chunk = sock.recv(n - len(buf))
try:
chunk = sock.recv(n - len(buf))
except (socket.error, OSError) as ex:
raise connection_lost(ex)
if not chunk:
raise InterfaceError("connection closed by server")
buf += chunk
@@ -52,7 +57,10 @@ def _putblock(sock, text):
chunk = data[off:off + _MAX_BLOCK]
off += _MAX_BLOCK
last = off >= len(data)
sock.sendall(struct.pack("<H", (len(chunk) << 1) | (1 if last else 0)) + chunk)
try:
sock.sendall(struct.pack("<H", (len(chunk) << 1) | (1 if last else 0)) + chunk)
except (socket.error, OSError) as ex:
raise connection_lost(ex)
if last:
break
@@ -186,6 +194,7 @@ def connect(host=None, port=50000, user=None, password=None, database=None, conn
host, port = host or "localhost", int(port or 50000)
try:
sock = socket.create_connection((host, port), timeout=connect_timeout)
keepalive(sock)
sock.settimeout(None)
except (socket.error, socket.timeout) as ex:
raise OperationalError("could not connect to '%s:%s' (%s)" % (host, port, ex))
@@ -202,6 +211,7 @@ def connect(host=None, port=50000, user=None, password=None, database=None, conn
sock.close()
host, port, database = m.group(1), int(m.group(2)), m.group(3) or database
sock = socket.create_connection((host, port), timeout=connect_timeout)
keepalive(sock)
sock.settimeout(None)
continue # merovingian proxy redirect: keep reading the next challenge on this socket
if block[0] == "!":
+33 -9
View File
@@ -22,6 +22,8 @@ from extra.dbwire import DatabaseError
from extra.dbwire import InterfaceError
from extra.dbwire import NotSupportedError
from extra.dbwire import OperationalError
from extra.dbwire import connection_lost
from extra.dbwire import keepalive
from extra.dbwire import ProgrammingError
# capability flags
@@ -50,6 +52,9 @@ def _xor(a, b):
def _u8(data, off):
return struct.unpack("<B", data[off:off + 1])[0]
def _u16(data, off):
return struct.unpack("<H", data[off:off + 2])[0]
def _cstring(data, off):
# NUL-terminated string, tolerant of a missing terminator (returns the remainder)
end = data.find(b"\x00", off)
@@ -60,7 +65,10 @@ def _cstring(data, off):
def _recvn(sock, n):
buf = b""
while len(buf) < n:
chunk = sock.recv(n - len(buf))
try:
chunk = sock.recv(n - len(buf))
except (socket.error, OSError) as ex:
raise connection_lost(ex)
if not chunk:
raise InterfaceError("connection closed by server")
buf += chunk
@@ -84,7 +92,10 @@ def _read_packet(sock):
def _send_packet(sock, seq, payload):
while True: # split payloads >= 16 MB into 0xffffff-sized packets (with a trailing short packet)
chunk = payload[:0xffffff]
sock.sendall(struct.pack("<I", len(chunk))[0:3] + struct.pack("<B", seq & 0xff) + chunk)
try:
sock.sendall(struct.pack("<I", len(chunk))[0:3] + struct.pack("<B", seq & 0xff) + chunk)
except (socket.error, OSError) as ex:
raise connection_lost(ex)
seq = (seq + 1) & 0xff
payload = payload[0xffffff:]
if len(chunk) < 0xffffff:
@@ -283,6 +294,7 @@ def _finish_auth(sock, password, plugin, salt):
def connect(host=None, port=3306, user=None, password=None, database=None, connect_timeout=None, **kwargs):
try:
sock = socket.create_connection((host or "localhost", int(port or 3306)), timeout=connect_timeout)
keepalive(sock)
sock.settimeout(None)
except (socket.error, socket.timeout) as ex:
raise OperationalError("could not connect to '%s:%s' (%s)" % (host, port, ex))
@@ -296,10 +308,10 @@ def connect(host=None, port=3306, user=None, password=None, database=None, conne
_, off = _cstring(payload, off) # server version
off += 4 # connection id
salt = payload[off:off + 8]; off += 8 + 1 # auth-plugin-data part 1 (+ filler)
off += 2 # capability flags (lower)
server_caps = _u16(payload, off); off += 2 # capability flags (lower)
off += 1 # character set
off += 2 # status flags
off += 2 # capability flags (upper)
server_caps |= _u16(payload, off) << 16; off += 2 # capability flags (upper)
auth_data_len = _u8(payload, off); off += 1
off += 10 # reserved
salt += payload[off:off + max(13, auth_data_len - 8) - 1] # part 2 (drop trailing NUL)
@@ -315,19 +327,31 @@ def connect(host=None, port=3306, user=None, password=None, database=None, conne
plugin = "mysql_native_password"
auth_response = _scramble_native(password or "", salt)
flags = (_CLIENT_LONG_PASSWORD | _CLIENT_LONG_FLAG | _CLIENT_PROTOCOL_41 |
_CLIENT_TRANSACTIONS | _CLIENT_SECURE_CONNECTION | _CLIENT_PLUGIN_AUTH)
# capabilities are NEGOTIATED: each side advertises what it can do and only the intersection is in
# play. Sending a fixed set makes the client claim features a proxy/fork/older server never offered,
# which is how a handshake ends up desynchronized rather than cleanly refused.
if not (server_caps & _CLIENT_PROTOCOL_41):
raise OperationalError("server does not support the 4.1 protocol, which this client requires")
wanted = (_CLIENT_LONG_PASSWORD | _CLIENT_LONG_FLAG | _CLIENT_PROTOCOL_41 |
_CLIENT_TRANSACTIONS | _CLIENT_SECURE_CONNECTION | _CLIENT_PLUGIN_AUTH)
if database:
flags |= _CLIENT_CONNECT_WITH_DB
wanted |= _CLIENT_CONNECT_WITH_DB
flags = wanted & server_caps
flags |= _CLIENT_PROTOCOL_41 | _CLIENT_SECURE_CONNECTION # mandatory for the packets built below
if database and not (flags & _CLIENT_CONNECT_WITH_DB):
raise OperationalError("server does not support selecting a database during the handshake")
if not (flags & _CLIENT_PLUGIN_AUTH):
plugin = None # pre-4.1.1 style: no trailing plugin name in the handshake response
response = struct.pack("<I", flags) + struct.pack("<I", _MAX_PACKET) + struct.pack("<B", 45) + (b"\x00" * 23)
response += (user or "").encode("utf-8") + b"\x00"
response += struct.pack("<B", len(auth_response)) + auth_response
if database:
response += database.encode("utf-8") + b"\x00"
response += plugin.encode("ascii") + b"\x00"
if plugin:
response += plugin.encode("ascii") + b"\x00"
_send_packet(sock, seq + 1, response)
_finish_auth(sock, password or "", plugin, salt)
_finish_auth(sock, password or "", plugin or "mysql_native_password", salt)
except (DatabaseError, InterfaceError):
_safe_close(sock)
raise
+41 -7
View File
@@ -28,6 +28,8 @@ from extra.dbwire import IntegrityError
from extra.dbwire import InterfaceError
from extra.dbwire import NotSupportedError
from extra.dbwire import OperationalError
from extra.dbwire import connection_lost
from extra.dbwire import keepalive
from extra.dbwire import ProgrammingError
_PROTOCOL_VERSION = 196608 # 3.0
@@ -71,7 +73,10 @@ def _xor(a, b):
def _recvn(sock, n):
buf = b""
while len(buf) < n:
chunk = sock.recv(n - len(buf))
try:
chunk = sock.recv(n - len(buf))
except (socket.error, OSError) as ex:
raise connection_lost(ex)
if not chunk:
raise InterfaceError("connection closed by server")
buf += chunk
@@ -85,7 +90,10 @@ def _read_message(sock):
return mtype, _recvn(sock, length - 4)
def _send(sock, mtype, payload):
sock.sendall((mtype or b"") + struct.pack("!I", len(payload) + 4) + payload)
try:
sock.sendall((mtype or b"") + struct.pack("!I", len(payload) + 4) + payload)
except (socket.error, OSError) as ex:
raise connection_lost(ex)
def _error_message(payload):
# ErrorResponse/NoticeResponse: series of (byte field-code, cstring value), terminated by a NUL byte.
@@ -226,7 +234,7 @@ class Connection(object):
return description, rows, 0, rowcount
def _authenticate(sock, user, password):
cfirst_bare = None
client_nonce = cfirst_bare = salted = auth_message = None
while True:
mtype, payload = _read_message(sock)
if mtype in (b"N", b"S"): # NoticeResponse / ParameterStatus may legally precede AuthenticationOk
@@ -248,8 +256,8 @@ def _authenticate(sock, user, password):
elif code == 10: # SASL (SCRAM-SHA-256)
if not hasattr(hashlib, "pbkdf2_hmac"):
raise NotSupportedError("SCRAM-SHA-256 authentication requires Python >= 2.7.8 (hashlib.pbkdf2_hmac)")
nonce = base64.b64encode(os.urandom(18)).decode("ascii")
cfirst_bare = "n=,r=%s" % nonce
client_nonce = base64.b64encode(os.urandom(18)).decode("ascii")
cfirst_bare = "n=,r=%s" % client_nonce
client_first = "n,," + cfirst_bare
_send(sock, b"p", b"SCRAM-SHA-256\x00" + struct.pack("!I", len(client_first)) + client_first.encode("ascii"))
elif code == 11: # SASLContinue (server-first)
@@ -259,6 +267,13 @@ def _authenticate(sock, user, password):
snonce, salt, iterations = attrs["r"], base64.b64decode(attrs["s"]), int(attrs["i"])
except (KeyError, ValueError, binascii.Error, UnicodeDecodeError) as ex:
raise OperationalError("malformed SCRAM server-first message (%s)" % ex)
# RFC 5802 5.1: the server nonce MUST start with the client nonce and MUST add material of its
# own. Skipping this lets anything that can answer the TCP connection replay a recorded
# server-first and drive the exchange - and dbwire has no TLS layer underneath to catch it.
if not client_nonce or not snonce.startswith(client_nonce) or len(snonce) <= len(client_nonce):
raise OperationalError("SCRAM server nonce does not extend the client nonce (rogue server?)")
if iterations < 4096: # RFC 5802 recommends >= 4096; a tiny count cheapens an offline attack
raise OperationalError("SCRAM iteration count %d is too low" % iterations)
salted = hashlib.pbkdf2_hmac("sha256", (password or "").encode("utf-8"), salt, iterations)
client_key = hmac.new(salted, b"Client Key", hashlib.sha256).digest()
stored_key = hashlib.sha256(client_key).digest()
@@ -267,8 +282,26 @@ def _authenticate(sock, user, password):
client_sig = hmac.new(stored_key, auth_message.encode("ascii"), hashlib.sha256).digest()
proof = base64.b64encode(_xor(client_key, client_sig)).decode("ascii")
_send(sock, b"p", ("%s,p=%s" % (client_final_noproof, proof)).encode("ascii"))
elif code == 12: # SASLFinal
pass
elif code == 12: # SASLFinal (server-final): verify the server too, or the handshake is one-way
# Without this the client proves itself to the server and simply trusts whatever answers back.
# ServerSignature = HMAC(ServerKey, AuthMessage) can only be produced by a peer that holds the
# stored credentials, so comparing it is what makes the exchange mutual (RFC 5802 5, 5.1).
if salted is None or auth_message is None:
raise OperationalError("unexpected SCRAM server-final message")
try:
attrs = dict(kv.split("=", 1) for kv in payload[4:].decode("ascii").split(","))
except (ValueError, UnicodeDecodeError) as ex:
raise OperationalError("malformed SCRAM server-final message (%s)" % ex)
if "e" in attrs:
raise OperationalError("SCRAM authentication failed (%s)" % attrs["e"])
try:
signature = base64.b64decode(attrs["v"])
except (KeyError, binascii.Error, ValueError) as ex:
raise OperationalError("malformed SCRAM server signature (%s)" % ex)
server_key = hmac.new(salted, b"Server Key", hashlib.sha256).digest()
expected = hmac.new(server_key, auth_message.encode("ascii"), hashlib.sha256).digest()
if not hmac.compare_digest(signature, expected):
raise OperationalError("SCRAM server signature mismatch (rogue server?)")
else:
raise InterfaceError("unsupported authentication request %d" % code)
@@ -279,6 +312,7 @@ def _raise_server_error_as_operational(payload):
def connect(host=None, port=5432, user=None, password=None, database=None, connect_timeout=None, **kwargs):
try:
sock = socket.create_connection((host or "localhost", int(port or 5432)), timeout=connect_timeout)
keepalive(sock)
sock.settimeout(None)
except (socket.error, socket.timeout) as ex:
raise OperationalError("could not connect to '%s:%s' (%s)" % (host, port, ex))
+50 -4
View File
@@ -25,6 +25,7 @@ except ImportError:
from extra.dbwire import InterfaceError
from extra.dbwire import NotSupportedError
from extra.dbwire import OperationalError
from extra.dbwire import http_origin
from extra.dbwire import ProgrammingError
def _convert(value, coltype):
@@ -72,9 +73,16 @@ class Cursor(object):
def close(self):
self._rows = []
def _split_pair(item):
"""'key=value' -> (key, value); a bare 'key' keeps a None value (Trino sends both forms)."""
key, sep, value = item.strip().partition("=")
return key.strip(), (value.strip() if sep else None)
class Connection(object):
def __init__(self, host, port, user, password, catalog, schema, timeout):
self._statement_url = "http://%s:%d/v1/statement" % (host, port)
self._statement_url = "%s/v1/statement" % http_origin(host, port)
self._timeout = timeout
self._headers = {"Content-Type": "text/plain"}
for prefix in ("X-Presto-", "X-Trino-"):
@@ -84,8 +92,8 @@ class Connection(object):
# request ("Schema is set but catalog is not"), so never force a "default" schema
if catalog:
self._headers[prefix + "Catalog"] = catalog
if schema:
self._headers[prefix + "Schema"] = schema
if schema: # only inside the catalog branch: a Schema alone is rejected
self._headers[prefix + "Schema"] = schema
if password:
token = base64.b64encode(("%s:%s" % (user or "", password)).encode("utf-8")).decode("ascii")
self._headers["Authorization"] = "Basic %s" % token
@@ -102,16 +110,54 @@ class Connection(object):
def close(self):
pass # HTTP is stateless
def _apply_state(self, info):
"""
Carry the session state the server hands back into the headers of every later request.
The client protocol is stateless on the wire, so the SERVER cannot remember anything: it reports
each change as a response header and the client is required to echo it back. Ignoring them makes
'USE', 'SET SESSION', 'SET ROLE' and 'START TRANSACTION' appear to succeed and then silently have
no effect on the next statement.
"""
for prefix in ("X-Presto-", "X-Trino-"):
for suffix, header in (("Catalog", "Catalog"), ("Schema", "Schema"), ("Path", "Path")):
value = info.get((prefix + "Set-" + suffix).lower())
if value:
self._headers[prefix + header] = value
started = info.get((prefix + "Started-Transaction-Id").lower())
if started:
self._headers[prefix + "Transaction-Id"] = started
if info.get((prefix + "Clear-Transaction-Id").lower()):
self._headers.pop(prefix + "Transaction-Id", None)
# Set-Session / Set-Role accumulate as comma-separated 'key=value' pairs, and the matching
# Clear-* header removes one by name
for kind in ("Session", "Role"):
current = dict(_split_pair(_) for _ in (self._headers.get(prefix + kind) or "").split(",") if _)
for item in (info.get((prefix + "Set-" + kind).lower()) or "").split(","):
if item.strip():
key, value = _split_pair(item)
current[key] = value
for key in (info.get((prefix + "Clear-" + kind).lower()) or "").split(","):
current.pop(key.strip(), None)
if current:
self._headers[prefix + kind] = ",".join("%s=%s" % (k, v) if v is not None else k for k, v in sorted(current.items()))
else:
self._headers.pop(prefix + kind, None)
def _request(self, url, data=None):
req = Request(url, data=data.encode("utf-8") if data is not None else None, headers=self._headers)
try:
body = urlopen(req, timeout=self._timeout).read().decode("utf-8", "replace")
response = urlopen(req, timeout=self._timeout)
body = response.read().decode("utf-8", "replace")
except HTTPError as ex:
raise ProgrammingError("(remote) HTTP %s: %s" % (ex.code, ex.read().decode("utf-8", "replace")[:200]))
except URLError as ex:
raise OperationalError("(remote) %s" % ex)
except (socket.timeout, socket.error) as ex:
raise OperationalError("(remote) %s" % ex)
info = response.info()
self._apply_state(dict((k.lower(), v) for k, v in (info.items() if hasattr(info, "items") else [])))
try:
return json.loads(body)
except ValueError as ex:
+33 -12
View File
@@ -22,9 +22,12 @@ from extra.dbwire import DatabaseError
from extra.dbwire import InterfaceError
from extra.dbwire import NotSupportedError
from extra.dbwire import OperationalError
from extra.dbwire import connection_lost
from extra.dbwire import keepalive
from extra.dbwire import ProgrammingError
_MAX_MESSAGE_LENGTH = 0x40000000
_DONE_COUNT = 0x0010 # DONE status bit: DoneRowCount carries a valid affected-row count
# packet types
_PKT_SQL_BATCH = 0x01
@@ -38,7 +41,10 @@ def _u8(data, off):
def _recvn(sock, n):
buf = b""
while len(buf) < n:
chunk = sock.recv(n - len(buf))
try:
chunk = sock.recv(n - len(buf))
except (socket.error, OSError) as ex:
raise connection_lost(ex)
if not chunk:
raise InterfaceError("connection closed by server")
buf += chunk
@@ -54,23 +60,32 @@ def _send_message(sock, mtype, data):
off += chunk_size
last = off >= len(data)
header = struct.pack(">BBHHBB", mtype, _STATUS_EOM if last else 0x00, len(chunk) + 8, 0, packet_id & 0xff, 0)
sock.sendall(header + chunk)
try:
sock.sendall(header + chunk)
except (socket.error, OSError) as ex:
raise connection_lost(ex)
packet_id += 1
if last:
break
def _read_message(sock):
# reassemble a full TDS message across packets (EOM status bit marks the last)
body = b""
# reassemble a full TDS message across packets (EOM status bit marks the last). The packet length is a
# 16-bit field, so bounding IT against _MAX_MESSAGE_LENGTH can never trigger - a hostile peer simply
# never sets EOM and streams packets forever. Bound the accumulated message instead, and collect the
# chunks in a list so reassembly stays linear rather than re-copying a growing immutable buffer.
chunks, total = [], 0
while True:
header = _recvn(sock, 8)
mtype, status, length = struct.unpack(">BBH", header[:4])
if length < 8 or length > _MAX_MESSAGE_LENGTH:
if length < 8:
raise InterfaceError("invalid TDS packet length (%d)" % length)
body += _recvn(sock, length - 8)
total += length - 8
if total > _MAX_MESSAGE_LENGTH:
raise InterfaceError("TDS message exceeds the maximum allowed length (%d bytes)" % _MAX_MESSAGE_LENGTH)
chunks.append(_recvn(sock, length - 8))
if status & _STATUS_EOM:
break
return body
return b"".join(chunks)
# ---- PRELOGIN ----------------------------------------------------------------------------------------
@@ -465,12 +480,12 @@ def _decode_value(col, data, off):
def _parse_tokens(sock, login=False):
data = _read_message(sock)
off, columns, rows, description, error = 0, [], [], None, None
off, columns, rows, description, error, affected = 0, [], [], None, None, None
while off < len(data):
token = _u8(data, off); off += 1
if token == 0x81: # COLMETADATA (a new result set: drop any prior rows so only the last is returned)
(count,) = struct.unpack("<H", data[off:off + 2]); off += 2
columns, rows = [], []
columns, rows, affected = [], [], None
if count == 0xffff:
continue
for _ in range(count):
@@ -518,6 +533,9 @@ def _parse_tokens(sock, login=False):
elif token == 0xa9: # ORDER
(tlen,) = struct.unpack("<H", data[off:off + 2]); off += 2 + tlen
elif token in (0xfd, 0xfe, 0xff): # DONE / DONEPROC / DONEINPROC
status, _curcmd, rowcount = struct.unpack("<HHq", data[off:off + 12])
if status & _DONE_COUNT: # only then is DoneRowCount meaningful (MS-TDS 2.2.7.5)
affected = rowcount if affected is None else affected + rowcount
off += 12
else:
raise InterfaceError("unexpected TDS token 0x%02x" % token)
@@ -526,7 +544,7 @@ def _parse_tokens(sock, login=False):
if login:
raise OperationalError("(remote) %s" % error)
raise ProgrammingError("(remote) %s" % error)
return description, rows
return description, rows, affected
class Cursor(object):
def __init__(self, connection):
@@ -540,8 +558,10 @@ class Cursor(object):
if params is not None:
raise NotSupportedError("parameter binding is not supported; pass a fully-formed query string")
self.description, self.rowcount, self._rows, self._pos = None, -1, [], 0
self.description, self._rows = self.connection._query(query)
self.rowcount = len(self._rows)
self.description, self._rows, affected = self.connection._query(query)
# a result-producing statement reports what it returned; DML has no rows, so the server's
# DoneRowCount is the only place its affected-row count exists
self.rowcount = len(self._rows) if self.description is not None else (affected if affected is not None else -1)
return self
def fetchall(self):
@@ -590,6 +610,7 @@ class Connection(object):
def connect(host=None, port=1433, user=None, password=None, database=None, connect_timeout=None, **kwargs):
try:
sock = socket.create_connection((host or "localhost", int(port or 1433)), timeout=connect_timeout)
keepalive(sock)
sock.settimeout(None)
except (socket.error, socket.timeout) as ex:
raise OperationalError("could not connect to '%s:%s' (%s)" % (host, port, ex))
+1 -1
View File
@@ -20,7 +20,7 @@ from lib.core.enums import OS
from thirdparty import six
# sqlmap version (<major>.<minor>.<month>.<monthly commit>)
VERSION = "1.10.8.10"
VERSION = "1.10.8.11"
TYPE = "dev" if VERSION.count('.') > 2 and VERSION.split('.')[-1] != '0' else "stable"
TYPE_COLORS = {"dev": 33, "stable": 90, "pip": 34}
VERSION_STRING = "sqlmap/%s#%s" % ('.'.join(VERSION.split('.')[:-1]) if VERSION.count('.') > 2 and VERSION.split('.')[-1] == '0' else VERSION, TYPE)