Move peo package to a different spot to prepare it for PyPI
This commit is contained in:
Bart Broere committed
1 parent
9a25aa11c9
commit
3e8e498b56
16 files changed
+408
-59
No files matched your search
+372
@@ -0,0 +1,372 @@
|
||||
# SPDX-FileCopyrightText: 2026 Bart Broere
|
||||
# SPDX-License-Identifier: GPL-3.0-or-later
|
||||
"""Bolt protocol server (raw TCP and WebSocket on the same port).
|
||||
|
||||
Supports Bolt 4.4 and 5.0-5.4, which covers the neo4j-javascript-driver used
|
||||
by neo4j-browser as well as current Python/Java drivers and cypher-shell.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import hashlib
|
||||
import itertools
|
||||
import logging
|
||||
import struct
|
||||
|
||||
from . import packstream
|
||||
from .backend import BoltFailure, Config, Session
|
||||
from .packstream import Structure
|
||||
from .web import discovery_document
|
||||
|
||||
log = logging.getLogger("peo.bolt")
|
||||
|
||||
MAGIC = b"\x60\x60\xB0\x17"
|
||||
SUPPORTED = [(5, 4), (5, 3), (5, 2), (5, 1), (5, 0), (4, 4)]
|
||||
SERVER_AGENT = "Neo4j/5.26.0"
|
||||
|
||||
HELLO, GOODBYE, RESET, RUN, BEGIN, COMMIT, ROLLBACK = 0x01, 0x02, 0x0F, 0x10, 0x11, 0x12, 0x13
|
||||
DISCARD, PULL, TELEMETRY, ROUTE, LOGON, LOGOFF = 0x2F, 0x3F, 0x54, 0x66, 0x6A, 0x6B
|
||||
SUCCESS, RECORD, IGNORED, FAILURE = 0x70, 0x71, 0x7E, 0x7F
|
||||
|
||||
_connection_ids = itertools.count(1)
|
||||
|
||||
|
||||
class ConnectionClosed(Exception):
|
||||
pass
|
||||
|
||||
|
||||
# -- transports -----------------------------------------------------------------
|
||||
|
||||
|
||||
class RawTransport:
|
||||
def __init__(self, reader, writer, prefix=b""):
|
||||
self.reader, self.writer, self.buffer = reader, writer, bytearray(prefix)
|
||||
|
||||
async def read(self, n):
|
||||
while len(self.buffer) < n:
|
||||
chunk = await self.reader.read(65536)
|
||||
if not chunk:
|
||||
raise ConnectionClosed()
|
||||
self.buffer += chunk
|
||||
data = bytes(self.buffer[:n])
|
||||
del self.buffer[:n]
|
||||
return data
|
||||
|
||||
async def write(self, data):
|
||||
self.writer.write(data)
|
||||
await self.writer.drain()
|
||||
|
||||
async def close(self):
|
||||
self.writer.close()
|
||||
|
||||
|
||||
class WebSocketTransport(RawTransport):
|
||||
"""Minimal RFC 6455 server: Bolt bytes are carried in binary frames."""
|
||||
|
||||
GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
|
||||
|
||||
async def handshake(self):
|
||||
request = bytearray(self.buffer)
|
||||
self.buffer.clear()
|
||||
while b"\r\n\r\n" not in request:
|
||||
chunk = await self.reader.read(4096)
|
||||
if not chunk or len(request) > 65536:
|
||||
raise ConnectionClosed()
|
||||
request += chunk
|
||||
head, _, rest = bytes(request).partition(b"\r\n\r\n")
|
||||
headers = {}
|
||||
for line in head.decode("latin-1").split("\r\n")[1:]:
|
||||
k, _, v = line.partition(":")
|
||||
headers[k.strip().lower()] = v.strip()
|
||||
method = head.split(b" ", 1)[0]
|
||||
key = headers.get("sec-websocket-key")
|
||||
if method != b"GET" or not key:
|
||||
await self._plain_http(method, headers)
|
||||
raise ConnectionClosed()
|
||||
accept = base64.b64encode(hashlib.sha1((key + self.GUID).encode()).digest()).decode()
|
||||
self.writer.write(("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\n"
|
||||
f"Connection: Upgrade\r\nSec-WebSocket-Accept: {accept}\r\n\r\n").encode())
|
||||
await self.writer.drain()
|
||||
self.raw = bytearray(rest)
|
||||
|
||||
async def _plain_http(self, method, headers):
|
||||
"""Answer non-WebSocket requests, e.g. neo4j-browser probing for discovery data."""
|
||||
cors = (b"Access-Control-Allow-Origin: *\r\nAccess-Control-Allow-Headers: *\r\n"
|
||||
b"Access-Control-Allow-Methods: GET, HEAD, OPTIONS\r\nConnection: close\r\n")
|
||||
if method == b"OPTIONS":
|
||||
response = b"HTTP/1.1 204 No Content\r\n" + cors + b"\r\n"
|
||||
elif method in (b"GET", b"HEAD"):
|
||||
body = discovery_document(headers, self.writer.get_extra_info("sockname")[1])
|
||||
response = (b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n" + cors
|
||||
+ b"Content-Length: %d\r\n\r\n" % len(body) + (body if method == b"GET" else b""))
|
||||
else:
|
||||
response = b"HTTP/1.1 405 Method Not Allowed\r\n" + cors + b"Content-Length: 0\r\n\r\n"
|
||||
self.writer.write(response)
|
||||
await self.writer.drain()
|
||||
|
||||
async def _raw(self, n):
|
||||
while len(self.raw) < n:
|
||||
chunk = await self.reader.read(65536)
|
||||
if not chunk:
|
||||
raise ConnectionClosed()
|
||||
self.raw += chunk
|
||||
data = bytes(self.raw[:n])
|
||||
del self.raw[:n]
|
||||
return data
|
||||
|
||||
async def read(self, n):
|
||||
while len(self.buffer) < n:
|
||||
b0, b1 = await self._raw(2)
|
||||
opcode, masked, length = b0 & 0x0F, b1 & 0x80, b1 & 0x7F
|
||||
if length == 126:
|
||||
length = struct.unpack(">H", await self._raw(2))[0]
|
||||
elif length == 127:
|
||||
length = struct.unpack(">Q", await self._raw(8))[0]
|
||||
mask = await self._raw(4) if masked else None
|
||||
payload = await self._raw(length)
|
||||
if mask:
|
||||
payload = bytes(b ^ mask[i % 4] for i, b in enumerate(payload))
|
||||
if opcode == 0x8:
|
||||
await self._frame(0x8, payload[:2])
|
||||
raise ConnectionClosed()
|
||||
if opcode == 0x9:
|
||||
await self._frame(0xA, payload)
|
||||
elif opcode in (0x0, 0x1, 0x2):
|
||||
self.buffer += payload
|
||||
data = bytes(self.buffer[:n])
|
||||
del self.buffer[:n]
|
||||
return data
|
||||
|
||||
async def _frame(self, opcode, payload):
|
||||
n = len(payload)
|
||||
if n < 126:
|
||||
header = struct.pack(">BB", 0x80 | opcode, n)
|
||||
elif n < 0x10000:
|
||||
header = struct.pack(">BBH", 0x80 | opcode, 126, n)
|
||||
else:
|
||||
header = struct.pack(">BBQ", 0x80 | opcode, 127, n)
|
||||
self.writer.write(header + payload)
|
||||
await self.writer.drain()
|
||||
|
||||
async def write(self, data):
|
||||
await self._frame(0x2, data)
|
||||
|
||||
|
||||
# -- protocol -----------------------------------------------------------------------
|
||||
|
||||
|
||||
def negotiate(proposals: bytes):
|
||||
for i in range(4):
|
||||
_, span, minor, major = proposals[i * 4:i * 4 + 4]
|
||||
for version in SUPPORTED:
|
||||
if version[0] == major and minor - span <= version[1] <= minor:
|
||||
return version
|
||||
return None
|
||||
|
||||
|
||||
class BoltConnection:
|
||||
def __init__(self, transport, config: Config):
|
||||
self.t = transport
|
||||
self.config = config
|
||||
self.id = f"bolt-{next(_connection_ids)}"
|
||||
self.version = None
|
||||
self.packer = None
|
||||
self.session: Session | None = None
|
||||
self.failed = False
|
||||
self.routing = {}
|
||||
self.results = {} # qid -> [Result, next row index]
|
||||
self.last_qid = -1
|
||||
self.qids = itertools.count()
|
||||
|
||||
async def serve(self):
|
||||
proposals = await self.t.read(16)
|
||||
self.version = negotiate(proposals)
|
||||
if self.version is None:
|
||||
log.info("%s: no supported Bolt version in %s", self.id, proposals.hex())
|
||||
await self.t.write(b"\x00\x00\x00\x00")
|
||||
return
|
||||
await self.t.write(bytes([0, 0, self.version[1], self.version[0]]))
|
||||
self.packer = packstream.Packer(element_ids=self.version >= (5, 0))
|
||||
log.debug("%s: negotiated Bolt %d.%d", self.id, *self.version)
|
||||
try:
|
||||
while True:
|
||||
message = await self.read_message()
|
||||
if message is None:
|
||||
continue # NOOP chunk
|
||||
if not await self.handle(message):
|
||||
break
|
||||
finally:
|
||||
if self.session:
|
||||
await self.session.close()
|
||||
|
||||
async def read_message(self):
|
||||
data = bytearray()
|
||||
while True:
|
||||
size = struct.unpack(">H", await self.t.read(2))[0]
|
||||
if size == 0:
|
||||
if not data:
|
||||
return None
|
||||
return packstream.unpack(bytes(data))
|
||||
data += await self.t.read(size)
|
||||
|
||||
async def send(self, tag, *fields):
|
||||
data = self.packer.pack(Structure(tag, list(fields)))
|
||||
out = bytearray()
|
||||
for i in range(0, len(data), 0xFFFF):
|
||||
chunk = data[i:i + 0xFFFF]
|
||||
out += struct.pack(">H", len(chunk)) + chunk
|
||||
out += b"\x00\x00"
|
||||
await self.t.write(bytes(out))
|
||||
|
||||
async def success(self, metadata=None):
|
||||
await self.send(SUCCESS, metadata or {})
|
||||
|
||||
async def failure(self, code, message):
|
||||
log.info("%s: %s: %s", self.id, code, message)
|
||||
self.failed = True
|
||||
await self.send(FAILURE, {"code": code, "message": message})
|
||||
|
||||
async def handle(self, msg) -> bool:
|
||||
tag, f = msg.tag, msg.fields
|
||||
if tag == GOODBYE:
|
||||
return False
|
||||
if tag == RESET:
|
||||
self.failed = False
|
||||
self.results.clear()
|
||||
if self.session:
|
||||
await self.session.reset()
|
||||
await self.success()
|
||||
return True
|
||||
if self.failed:
|
||||
await self.send(IGNORED)
|
||||
return True
|
||||
try:
|
||||
await self.dispatch(tag, f)
|
||||
except BoltFailure as e:
|
||||
self.results.clear()
|
||||
await self.failure(e.code, e.message)
|
||||
except Exception as e: # noqa: BLE001 - report anything to the client
|
||||
log.exception("%s: internal error", self.id)
|
||||
self.results.clear()
|
||||
await self.failure("Neo.DatabaseError.General.UnknownError", f"peo internal error: {e}")
|
||||
return True
|
||||
|
||||
async def dispatch(self, tag, f):
|
||||
if tag == HELLO:
|
||||
extra = f[0] if f else {}
|
||||
self.routing = extra.get("routing") or {}
|
||||
if self.version < (5, 1):
|
||||
await self.logon(extra)
|
||||
await self.success({"server": SERVER_AGENT, "connection_id": self.id, "hints": {}})
|
||||
elif tag == LOGON:
|
||||
await self.logon(f[0])
|
||||
await self.success()
|
||||
elif tag == LOGOFF:
|
||||
if self.session:
|
||||
await self.session.close()
|
||||
self.session = None
|
||||
await self.success()
|
||||
elif tag == TELEMETRY:
|
||||
await self.success()
|
||||
elif tag == ROUTE:
|
||||
await self.route(f)
|
||||
elif tag == RUN:
|
||||
await self.run(f)
|
||||
elif tag in (PULL, DISCARD):
|
||||
await self.pull(f[0] if f else {}, discard=tag == DISCARD)
|
||||
elif tag == BEGIN:
|
||||
extra = f[0] if f else {}
|
||||
await self.require_session().begin(extra.get("db"))
|
||||
self.results.clear()
|
||||
await self.success()
|
||||
elif tag == COMMIT:
|
||||
await self.require_session().commit()
|
||||
await self.success({"bookmark": "peo:bookmark"})
|
||||
elif tag == ROLLBACK:
|
||||
await self.require_session().rollback()
|
||||
await self.success()
|
||||
else:
|
||||
raise BoltFailure("Neo.ClientError.Request.Invalid", f"Unsupported Bolt message 0x{tag:02X}")
|
||||
|
||||
def require_session(self) -> Session:
|
||||
if self.session is None:
|
||||
raise BoltFailure("Neo.ClientError.Security.Unauthorized", "Not authenticated")
|
||||
return self.session
|
||||
|
||||
async def logon(self, auth):
|
||||
if auth.get("scheme") != "basic":
|
||||
raise BoltFailure("Neo.ClientError.Security.Unauthorized",
|
||||
"peo requires basic authentication with a PostgreSQL user name and password.")
|
||||
session = Session(self.config, auth.get("principal", ""))
|
||||
await session.connect(auth.get("credentials", ""))
|
||||
self.session = session
|
||||
|
||||
async def route(self, f):
|
||||
extra = f[2] if len(f) > 2 else None
|
||||
db = extra.get("db") if isinstance(extra, dict) else extra
|
||||
db = await self.require_session().resolve_db(db)
|
||||
address = self.routing.get("address") or self.config.advertised_bolt
|
||||
await self.success({"rt": {"ttl": 300, "db": db, "servers": [
|
||||
{"addresses": [address], "role": "WRITE"},
|
||||
{"addresses": [address], "role": "READ"},
|
||||
{"addresses": [address], "role": "ROUTE"},
|
||||
]}})
|
||||
|
||||
async def run(self, f):
|
||||
query, params = f[0], f[1] if len(f) > 1 else {}
|
||||
extra = f[2] if len(f) > 2 else {}
|
||||
session = self.require_session()
|
||||
if not session.in_transaction:
|
||||
self.results.clear()
|
||||
result = await session.run(query, params, extra.get("db"))
|
||||
qid = next(self.qids)
|
||||
self.results[qid] = [result, 0]
|
||||
self.last_qid = qid
|
||||
metadata = {"fields": result.fields, "t_first": result.t_first}
|
||||
if session.in_transaction:
|
||||
metadata["qid"] = qid
|
||||
await self.success(metadata)
|
||||
|
||||
async def pull(self, extra, discard):
|
||||
qid = extra.get("qid", -1)
|
||||
qid = self.last_qid if qid == -1 else qid
|
||||
if qid not in self.results:
|
||||
raise BoltFailure("Neo.ClientError.Request.Invalid", "No result available to pull from")
|
||||
result, index = self.results[qid]
|
||||
n = extra.get("n", -1)
|
||||
end = len(result.rows) if n == -1 or discard else min(index + n, len(result.rows))
|
||||
if not discard:
|
||||
for row in result.rows[index:end]:
|
||||
await self.send(RECORD, row)
|
||||
self.results[qid][1] = end
|
||||
if end < len(result.rows):
|
||||
await self.success({"has_more": True})
|
||||
return
|
||||
del self.results[qid]
|
||||
metadata = {"type": result.query_type, "t_last": 0, "db": result.db}
|
||||
if not self.session.in_transaction:
|
||||
metadata["bookmark"] = "peo:bookmark"
|
||||
await self.success(metadata)
|
||||
|
||||
|
||||
async def handle_client(reader, writer, config: Config):
|
||||
peer = writer.get_extra_info("peername")
|
||||
try:
|
||||
first = await reader.readexactly(4)
|
||||
log.debug("connection from %s starts with %r", peer, first)
|
||||
if first == MAGIC:
|
||||
transport = RawTransport(reader, writer)
|
||||
elif first in (b"GET ", b"HEAD", b"OPTI", b"POST"):
|
||||
transport = WebSocketTransport(reader, writer, first)
|
||||
await transport.handshake()
|
||||
if await transport.read(4) != MAGIC:
|
||||
return
|
||||
else:
|
||||
return
|
||||
await BoltConnection(transport, config).serve()
|
||||
except (ConnectionClosed, asyncio.IncompleteReadError, ConnectionResetError, BrokenPipeError):
|
||||
pass
|
||||
except Exception: # noqa: BLE001
|
||||
log.exception("connection from %s failed", peer)
|
||||
finally:
|
||||
writer.close()
|
||||
Reference in new issue
Block a user