From be881ea65e6a70f3015f0f82fadfd8bd59ae6927 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Miroslav=20=C5=A0tampar?= Date: Fri, 7 Aug 2026 18:33:51 +0200 Subject: [PATCH] Stabilization of dbwire --- extra/dbwire/__init__.py | 41 ++++++++++++++++++++++++++++ extra/dbwire/clickhouse.py | 6 +++- extra/dbwire/cubrid.py | 26 ++++++++++++++++-- extra/dbwire/firebird.py | 56 ++++++++++++++++++++++++++------------ extra/dbwire/monetdb.py | 14 ++++++++-- extra/dbwire/mysql.py | 42 ++++++++++++++++++++++------ extra/dbwire/postgres.py | 48 +++++++++++++++++++++++++++----- extra/dbwire/presto.py | 54 +++++++++++++++++++++++++++++++++--- extra/dbwire/tds.py | 45 ++++++++++++++++++++++-------- lib/core/settings.py | 2 +- 10 files changed, 278 insertions(+), 56 deletions(-) diff --git a/extra/dbwire/__init__.py b/extra/dbwire/__init__.py index e48a84e4d..e809c201d 100644 --- a/extra/dbwire/__init__.py +++ b/extra/dbwire/__init__.py @@ -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 diff --git a/extra/dbwire/clickhouse.py b/extra/dbwire/clickhouse.py index b6a9cae58..c39dcfed0 100644 --- a/extra/dbwire/clickhouse.py +++ b/extra/dbwire/clickhouse.py @@ -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") diff --git a/extra/dbwire/cubrid.py b/extra/dbwire/cubrid.py index ae6f95f5e..750d7578b 100644 --- a/extra/dbwire/cubrid.py +++ b/extra/dbwire/cubrid.py @@ -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): diff --git a/extra/dbwire/firebird.py b/extra/dbwire/firebird.py index 647070723..a7a56da2e 100644 --- a/extra/dbwire/firebird.py +++ b/extra/dbwire/firebird.py @@ -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)) diff --git a/extra/dbwire/monetdb.py b/extra/dbwire/monetdb.py index 37f6db5c0..90dd4451f 100644 --- a/extra/dbwire/monetdb.py +++ b/extra/dbwire/monetdb.py @@ -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("= 16 MB into 0xffffff-sized packets (with a trailing short packet) chunk = payload[:0xffffff] - sock.sendall(struct.pack("= 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)) diff --git a/extra/dbwire/presto.py b/extra/dbwire/presto.py index 7deede433..7ec447f0c 100644 --- a/extra/dbwire/presto.py +++ b/extra/dbwire/presto.py @@ -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: diff --git a/extra/dbwire/tds.py b/extra/dbwire/tds.py index fe4a4dbb6..3cad7d70a 100644 --- a/extra/dbwire/tds.py +++ b/extra/dbwire/tds.py @@ -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("...) -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)