pyfect-dev 0.1.0__py3-none-any.whl

pyfect/socket/_net.py

@@ -0,0 +1,281 @@
+"""Live TCP and TLS sockets (Effect's ``NodeSocket``), on the socket loop.
+
+``from_connection`` turns an effect opening a ``Connection`` into a
+``SocketService``: each acquisition of its reader opens a connection, and
+its writer writes to the connection the current reader opened, waiting for
+one. ``make_net``/``make_unix``/``make_tls`` open new connections;
+``from_socket`` adopts a connected ``socket.socket``.
+"""
+
+from __future__ import annotations
+
+import socket
+import ssl
+from collections.abc import Callable, Coroutine, Sequence
+from typing import Any, Never
+
+from pyfect._core import Effect
+from pyfect._effect_args import require_effect
+from pyfect.channel import unwrap as unwrap_channel
+from pyfect.duration import Duration
+from pyfect.effect import fail, gen, succeed, suspend, sync, void
+from pyfect.exit import Exit
+from pyfect.latch import Latch
+from pyfect.layer import Layer
+from pyfect.layer import effect as layer_effect
+from pyfect.scope import Scope, acquire_release
+from pyfect.socket._connection import Connection
+from pyfect.socket._errors import (
+    CloseEvent,
+    SocketError,
+    close_error,
+    open_error,
+    read_error,
+    upgrade_error,
+)
+from pyfect.socket._loop import get_loop, run
+from pyfect.socket._socket import (
+    Chunk,
+    Reader,
+    Socket,
+    SocketChannel,
+    SocketService,
+    TlsUpgradeOptions,
+    Writer,
+    encode_chunk,
+    make,
+    to_channel,
+)
+from pyfect.socket._tls import DEFAULT_SERVERNAME, check_credentials, client_context, server_context
+
+type Connect = Callable[[Connection], Coroutine[Any, Any, object]]
+
+
+class _Current:
+    """The connection the current reader opened, which writers use."""
+
+    __slots__ = ("connection", "latch")
+
+    def __init__(self) -> None:
+        self.connection: Connection | None = None
+        self.latch = Latch(False)
+
+
+def from_connection(
+    open: Effect[Connection, SocketError, Scope],
+    *,
+    open_timeout: Duration | None = None,
+    tls_server: bool = False,
+    target: str | None = None,
+) -> SocketService:
+    """A socket over the connections ``open`` makes. ``open_timeout`` bounds
+    the opening; ``tls_server`` makes TLS upgrades take the server side;
+    ``target`` (``host:port``, a path) is named when opening fails."""
+    require_effect(open, "open")
+    current = _Current()
+    opening = open
+    if open_timeout is not None:
+        timed_out = open_error(TimeoutError("Connection timed out"), "Timeout", target=target)
+        opening = open.timeout_or_else(open_timeout, lambda: fail(timed_out))
+
+    def reader():
+        scope = yield from Scope()
+        conn: Connection = yield from opening
+        state = {"upgradable": True}
+
+        def release(_: Exit[Any, Any]) -> Effect[None]:
+            def run_release() -> None:
+                # Resume a pull parked in another fiber before detaching.
+                conn.fail(close_error(1006))
+                current.latch.close_unsafe()
+                current.connection = None
+                state["upgradable"] = False
+
+            return sync(run_release)
+
+        yield from scope.add_finalizer(release)
+        current.connection = conn
+        current.latch.open_unsafe()
+
+        def upgrade(options: TlsUpgradeOptions | None) -> Effect[None, SocketError]:
+            def start() -> Effect[None, SocketError]:
+                if not state["upgradable"]:
+                    return fail(upgrade_error(RuntimeError("socket is already upgraded or closed")))
+                return _upgrade(conn, options or TlsUpgradeOptions(), server=tls_server).map(
+                    lambda _: state.update(upgradable=False)
+                )
+
+            return suspend(start)
+
+        return Reader(conn.pull(), upgrade)
+
+    def write(chunk: Chunk | CloseEvent) -> Effect[None, SocketError]:
+        def start() -> Effect[None, SocketError]:
+            conn = current.connection
+            if conn is None:
+                return current.latch.when_open(write(chunk))
+            if isinstance(chunk, CloseEvent):
+                if chunk.code > 1000:
+                    conn.fail(read_error(ConnectionError(f"closed with code {chunk.code}")))
+                conn.close(graceful=False)
+                return void
+            return conn.write((encode_chunk(chunk),))
+
+        return suspend(start)
+
+    def write_all(chunks: Sequence[Chunk]) -> Effect[None, SocketError]:
+        def start() -> Effect[None, SocketError]:
+            conn = current.connection
+            if conn is None:
+                return current.latch.when_open(write_all(chunks))
+            return conn.write(tuple(map(encode_chunk, chunks)))
+
+        return suspend(start)
+
+    def end_writing(_writer: Writer, _exit: Exit[Any, Any]) -> Effect[None]:
+        def run_end() -> None:
+            if current.connection is not None:
+                current.connection.end_writing()
+
+        return sync(run_end)
+
+    writer = acquire_release(succeed(Writer(write, write_all)), end_writing)
+    return make(gen(reader), writer)
+
+
+def _upgrade(
+    conn: Connection, options: TlsUpgradeOptions, *, server: bool
+) -> Effect[None, SocketError]:
+    problem = check_credentials(options, server=server)
+    if problem is not None:
+        return fail(upgrade_error(ValueError(problem)))
+    try:
+        context = server_context(options) if server else client_context(options)
+    except (ssl.SSLError, OSError, ValueError) as cause:
+        return fail(upgrade_error(cause))
+    hostname = None if server else options.servername or DEFAULT_SERVERNAME
+
+    def failed(cause: BaseException) -> Effect[None, SocketError]:
+        error = upgrade_error(cause)
+        conn.fail(error)
+        conn.close(graceful=False)
+        return fail(error)
+
+    def interrupted() -> Effect[None]:
+        def tear_down() -> None:
+            conn.fail(close_error(1006))
+            conn.close(graceful=False)
+
+        return sync(tear_down)
+
+    return (
+        conn.upgrade(context, server_side=server, hostname=hostname)
+        .catch_all(failed)
+        .on_interrupt(interrupted())
+    )
+
+
+def connecting(
+    connect: Connect, target: str | None = None
+) -> Effect[Connection, SocketError, Scope]:
+    """A connection ``connect`` makes on the loop (with the protocol it is
+    given), closed with the scope: gracefully once open, at once before.
+    A failure to connect names ``target``."""
+
+    def body():
+        conn = Connection()
+        scope = yield from Scope()
+        yield from scope.add_finalizer(lambda _: sync(lambda: conn.close(graceful=conn.open)))
+        yield from run(lambda: connect(conn)).map_error(
+            lambda cause: open_error(cause, target=target)
+        )
+        conn.open = True
+        return conn
+
+    return gen(body)
+
+
+def make_net(
+    port: int, host: str = "localhost", *, open_timeout: Duration | None = None
+) -> SocketService:
+    """A TCP socket connecting to ``host:port``."""
+
+    async def connect(conn: Connection) -> object:
+        return await get_loop().create_connection(lambda: conn, host, port)
+
+    target = _address(host, port)
+    return from_connection(connecting(connect, target), open_timeout=open_timeout, target=target)
+
+
+def make_unix(path: str, *, open_timeout: Duration | None = None) -> SocketService:
+    """A socket connecting to the Unix socket at ``path``."""
+
+    async def connect(conn: Connection) -> object:
+        return await get_loop().create_unix_connection(lambda: conn, path)
+
+    return from_connection(connecting(connect, path), open_timeout=open_timeout, target=path)
+
+
+def make_tls(
+    port: int,
+    host: str = "localhost",
+    *,
+    tls: TlsUpgradeOptions | None = None,
+    open_timeout: Duration | None = None,
+) -> SocketService:
+    """A TLS socket connecting to ``host:port``; the certificate is checked
+    against ``tls.servername``, else ``host``."""
+    options = tls or TlsUpgradeOptions()
+
+    async def connect(conn: Connection) -> None:
+        context = client_context(options)
+        await get_loop().create_connection(lambda: conn, host, port)
+        await conn.start_tls(context, server_side=False, hostname=options.servername or host)
+
+    target = _address(host, port)
+    return from_connection(connecting(connect, target), open_timeout=open_timeout, target=target)
+
+
+def _address(host: str, port: int) -> str:
+    return f"[{host}]:{port}" if ":" in host else f"{host}:{port}"
+
+
+def _adopt(sock: socket.socket) -> Effect[Connection, SocketError, Scope]:
+    async def connect(conn: Connection) -> object:
+        return await get_loop().connect_accepted_socket(lambda: conn, sock)
+
+    return connecting(connect)
+
+
+def from_socket(
+    open: Effect[socket.socket, SocketError, Scope],
+    *,
+    tls_server: bool = False,
+    open_timeout: Duration | None = None,
+) -> SocketService:
+    """A socket over the connected ``socket.socket`` each acquisition of
+    the reader gets from ``open`` (Effect's ``fromDuplex``); the socket is
+    taken over and closed with the reader's scope. ``tls_server`` makes TLS
+    upgrades take the server side."""
+    require_effect(open, "open")
+    return from_connection(open.flat_map(_adopt), tls_server=tls_server, open_timeout=open_timeout)
+
+
+def make_net_channel[IE = Never](port: int, host: str = "localhost") -> SocketChannel[bytes, IE]:
+    return unwrap_channel(succeed(to_channel(make_net(port, host))))
+
+
+def make_tls_channel[IE = Never](
+    port: int, host: str = "localhost", *, tls: TlsUpgradeOptions | None = None
+) -> SocketChannel[bytes, IE]:
+    return unwrap_channel(succeed(to_channel(make_tls(port, host, tls=tls))))
+
+
+def layer_net(port: int, host: str = "localhost") -> Layer[Socket]:
+    return layer_effect(Socket, sync(lambda: make_net(port, host)))
+
+
+def layer_tls(
+    port: int, host: str = "localhost", *, tls: TlsUpgradeOptions | None = None
+) -> Layer[Socket]:
+    return layer_effect(Socket, sync(lambda: make_tls(port, host, tls=tls)))

pyfect/socket/_socket.py

@@ -0,0 +1,245 @@
+"""The socket model (Effect's ``Socket``): a reader and writer acquired in
+a scope, and the channels and streams built on them."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+from typing import Any, Never
+
+from pyfect._core import Effect
+from pyfect._effect_args import require_effect
+from pyfect._js import text_decoder, text_encode
+from pyfect.cause import Cause
+from pyfect.channel import Channel, from_transform
+from pyfect.channel import unwrap as unwrap_channel
+from pyfect.context import Tag
+from pyfect.effect import fail, fail_cause, gen, provide_service, succeed, suspend
+from pyfect.exit import Success
+from pyfect.pull import filter_no_done
+from pyfect.redacted import Redacted
+from pyfect.scope import CloseableScope, Scope, fork_in
+from pyfect.socket._errors import CloseEvent, SocketError, is_clean_close, upgrade_error
+from pyfect.stream import Stream, from_channel
+
+type Chunk = bytes | str
+type Write = Callable[[Chunk | CloseEvent], Effect[None, SocketError]]
+type WriteAll = Callable[[Sequence[Chunk]], Effect[None, SocketError]]
+
+type PemInput = str | bytes
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class TlsUpgradeOptions:
+    """How to upgrade a connection to TLS. ``key`` and ``cert`` (PEM) go
+    together; a server needs both. ``servername`` is the client's SNI name
+    and the name its certificate check expects."""
+
+    servername: str | None = None
+    key: Redacted[PemInput] | Sequence[Redacted[PemInput]] | None = None
+    cert: PemInput | Sequence[PemInput] | None = None
+    ca: PemInput | Sequence[PemInput] | None = None
+    passphrase: Redacted[str] | None = None
+    alpn_protocols: Sequence[str] | None = None
+    request_cert: bool | None = None
+    reject_unauthorized: bool | None = None
+
+
+type Upgrade = Callable[[TlsUpgradeOptions | None], Effect[None, SocketError]]
+
+
+def unsupported_upgrade(_options: TlsUpgradeOptions | None = None) -> Effect[None, SocketError]:
+    """The upgrade of a socket that cannot do TLS (Effect's
+    ``SocketUpgradeError.unsupported``)."""
+    return fail(upgrade_error())
+
+
+class Reader:
+    """What a socket's reader gives: ``pull`` waits for the next non-empty
+    batch of chunks; ``upgrade`` switches the connection to TLS."""
+
+    __slots__ = ("_upgrade", "pull")
+
+    def __init__(
+        self, pull: Effect[tuple[Chunk, ...], SocketError], upgrade: Upgrade = unsupported_upgrade
+    ) -> None:
+        require_effect(pull, "pull")
+        self.pull = pull
+        self._upgrade = upgrade
+
+    def upgrade(self, options: TlsUpgradeOptions | None = None) -> Effect[None, SocketError]:
+        return self._upgrade(options)
+
+
+@dataclass(frozen=True, slots=True)
+class Writer:
+    """What a socket's writer gives: ``write`` sends one chunk (a
+    ``CloseEvent`` closes the socket), ``write_all`` a non-empty batch."""
+
+    write: Write
+    write_all: WriteAll
+
+
+@dataclass(frozen=True, slots=True)
+class SocketService:
+    """A socket: acquiring ``reader`` in a scope opens it (until the scope
+    closes); ``writer`` writes to the open connection, waiting for one."""
+
+    reader: Effect[Reader, SocketError, Scope]
+    writer: Effect[Writer, Never, Scope]
+
+
+class Socket(Tag[SocketService], key="effect/socket/Socket"):
+    """The socket service."""
+
+
+def make(
+    reader: Effect[Reader, SocketError, Scope], writer: Effect[Writer, Never, Scope]
+) -> SocketService:
+    require_effect(reader, "reader")
+    require_effect(writer, "writer")
+    return SocketService(reader, writer)
+
+
+def encode_chunk(chunk: Chunk) -> bytes:
+    """A chunk as bytes: text as ``TextEncoder`` encodes it."""
+    return text_encode(chunk) if isinstance(chunk, str) else chunk
+
+
+def reader_bytes(
+    socket: SocketService,
+) -> Effect[Effect[tuple[bytes, ...], SocketError], SocketError, Scope]:
+    """The reader's pull, with text frames encoded as UTF-8."""
+    return socket.reader.map(lambda reader: reader.pull.map(lambda c: tuple(map(encode_chunk, c))))
+
+
+def reader_string(
+    socket: SocketService, encoding: str = "utf-8"
+) -> Effect[Effect[tuple[str, ...], SocketError], SocketError, Scope]:
+    """The reader's pull, decoding bytes with ``encoding`` across pulls. An
+    incomplete character is flushed (as U+FFFD) before a text frame and on a
+    normal close."""
+
+    def make_pull(reader: Reader) -> Effect[tuple[str, ...], SocketError]:
+        decoder = text_decoder(encoding)
+        state = {"needs_flush": False}
+
+        def flush() -> str:
+            state["needs_flush"] = False
+            text = decoder.decode(b"", final=True)
+            decoder.reset()
+            return text
+
+        def decode(chunk: tuple[Chunk, ...]) -> tuple[str, ...]:
+            out: list[str] = []
+            for item in chunk:
+                if isinstance(item, str):
+                    out.append(flush() + item if state["needs_flush"] else item)
+                else:
+                    out.append(decoder.decode(item))
+                    state["needs_flush"] = True
+            return tuple(out)
+
+        def on_error(error: SocketError) -> Effect[tuple[str, ...], SocketError]:
+            if state["needs_flush"] and is_clean_close(error):
+                last = flush()
+                if last:
+                    return succeed((last,))
+            return fail(error)
+
+        return reader.pull.map(decode).catch_all(on_error)
+
+    return socket.reader.map(make_pull)
+
+
+def _write_chunk(writer: Writer, chunk: Sequence[Chunk | CloseEvent]) -> Effect[None, SocketError]:
+    for index, item in enumerate(chunk):
+        if isinstance(item, CloseEvent):
+            return _write_each(writer, chunk[: index + 1])
+    return writer.write_all(chunk)  # pyright: ignore[reportArgumentType]
+
+
+def _write_each(writer: Writer, items: Sequence[Chunk | CloseEvent]) -> Effect[None, SocketError]:
+    def body():
+        for item in items:
+            yield from writer.write(item)
+
+    return gen(body)
+
+
+type SocketChannel[A, IE] = Channel[
+    tuple[A, ...], SocketError | IE, None, tuple[Chunk | CloseEvent, ...], IE, object
+]
+
+
+def _to_channel_with_reader[A: Chunk, IE](
+    socket: SocketService,
+    reader: Effect[Effect[tuple[A, ...], SocketError], SocketError, Scope],
+) -> SocketChannel[A, IE]:
+    def transform(
+        upstream: Effect[tuple[Chunk | CloseEvent, ...], Any], scope: CloseableScope
+    ) -> Effect[Effect[tuple[A, ...], Any], Any, Any]:
+        def body():
+            read_scope = yield from scope.fork()
+            pull = yield from reader.pipe(_in_scope(read_scope))
+            write_scope = yield from scope.fork()
+            writer = yield from socket.writer.pipe(_in_scope(write_scope))
+            failure: list[Cause[Any]] = []
+
+            def on_write_failure(_: object, cause: Cause[Any]) -> Effect[None]:
+                failure.append(cause)
+                # Closing the read scope fails a suspended pull, which is the
+                # reader contract, so the read side needs no per-pull race.
+                return read_scope.close(Success(None))
+
+            writing = (
+                upstream.flat_map(lambda chunk: _write_chunk(writer, chunk))
+                .forever(disable_yield=True)
+                .catch_cause_filter(filter_no_done, on_write_failure)
+                .ensuring(write_scope.close(Success(None)))
+            )
+            yield from fork_in(writing, scope)
+            return suspend(lambda: fail_cause(failure[0]) if failure else pull).catch_cause(
+                lambda cause: fail_cause(failure[0] if failure else cause)
+            )
+
+        return gen(body)
+
+    return from_transform(transform)
+
+
+def _in_scope[A, E](scope: CloseableScope) -> Callable[[Effect[A, E, Scope]], Effect[A, E]]:
+    return provide_service(Scope, scope)
+
+
+def to_channel[IE = Never](socket: SocketService) -> SocketChannel[bytes, IE]:
+    """The socket as a channel: its input chunks are written (a
+    ``CloseEvent`` closes the socket), its output is what the reader pulls,
+    as bytes. A failed write fails the channel."""
+    return _to_channel_with_reader(socket, reader_bytes(socket))
+
+
+def to_channel_string[IE = Never](
+    socket: SocketService, encoding: str = "utf-8"
+) -> SocketChannel[str, IE]:
+    """``to_channel``, with the output decoded as by ``reader_string``."""
+    return _to_channel_with_reader(socket, reader_string(socket, encoding))
+
+
+def to_stream(socket: SocketService) -> Stream[bytes, SocketError]:
+    """What the socket's reader pulls, as a stream of bytes."""
+    channel: Channel[tuple[bytes, ...], SocketError] = from_transform(
+        lambda _upstream, scope: reader_bytes(socket).pipe(_in_scope(scope))
+    )
+    return from_channel(channel)
+
+
+def make_channel[IE = Never]() -> Channel[
+    tuple[bytes, ...], SocketError | IE, None, tuple[Chunk | CloseEvent, ...], IE, object, Socket
+]:
+    """``to_channel`` of the ``Socket`` service."""
+
+    def channel(socket: SocketService) -> SocketChannel[bytes, IE]:
+        return to_channel(socket)
+
+    return unwrap_channel(Socket().map(channel))

pyfect/socket/_tls.py

@@ -0,0 +1,88 @@
+"""``ssl.SSLContext`` from ``TlsUpgradeOptions``, as Node's
+``tls.createSecureContext`` builds one: ``ca`` replaces the system roots,
+``reject_unauthorized=False`` turns verification off, and a server only
+asks for a client certificate when ``request_cert`` is set."""
+
+from __future__ import annotations
+
+import os
+import ssl
+import tempfile
+from collections.abc import Sequence
+
+from pyfect import redacted
+from pyfect.redacted import Redacted
+from pyfect.socket._socket import PemInput, TlsUpgradeOptions
+
+DEFAULT_SERVERNAME = "localhost"
+"""The name a client checks the certificate against without a
+``servername`` (Node's default ``host``)."""
+
+
+def _pem(value: PemInput) -> str:
+    return value.decode() if isinstance(value, bytes) else value
+
+
+def _pems(value: PemInput | Sequence[PemInput]) -> list[str]:
+    if isinstance(value, str | bytes):
+        return [_pem(value)]
+    return [_pem(v) for v in value]
+
+
+def _keys(value: Redacted[PemInput] | Sequence[Redacted[PemInput]]) -> list[str]:
+    if isinstance(value, Redacted):
+        return [_pem(redacted.value(value))]
+    return [_pem(redacted.value(v)) for v in value]
+
+
+def check_credentials(options: TlsUpgradeOptions, *, server: bool) -> str | None:
+    """Why the key and certificate options cannot be used, or ``None``."""
+    has_key, has_cert = options.key is not None, options.cert is not None
+    if server and not (has_key and has_cert):
+        return "server TLS upgrade requires both key and cert"
+    if has_key != has_cert:
+        return "TLS upgrade credentials must include both key and cert"
+    return None
+
+
+def _load_chain(context: ssl.SSLContext, options: TlsUpgradeOptions) -> None:
+    if options.key is None or options.cert is None:
+        return
+    password = None if options.passphrase is None else redacted.value(options.passphrase)
+    # ``load_cert_chain`` reads files only; they live just long enough to load.
+    with tempfile.TemporaryDirectory(prefix="pyfect-tls-") as directory:
+        cert_path = os.path.join(directory, "cert.pem")
+        key_path = os.path.join(directory, "key.pem")
+        for path, pems in ((cert_path, _pems(options.cert)), (key_path, _keys(options.key))):
+            descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
+            with os.fdopen(descriptor, "w") as file:
+                file.write("\n".join(pems))
+        context.load_cert_chain(cert_path, key_path, password)
+
+
+def client_context(options: TlsUpgradeOptions) -> ssl.SSLContext:
+    context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+    if options.ca is None:
+        context.load_default_certs()
+    else:
+        context.load_verify_locations(cadata="\n".join(_pems(options.ca)))
+    if options.reject_unauthorized is False:
+        context.check_hostname = False
+        context.verify_mode = ssl.CERT_NONE
+    _load_chain(context, options)
+    if options.alpn_protocols is not None:
+        context.set_alpn_protocols(list(options.alpn_protocols))
+    return context
+
+
+def server_context(options: TlsUpgradeOptions) -> ssl.SSLContext:
+    context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+    _load_chain(context, options)
+    if options.request_cert:
+        if options.ca is not None:
+            context.load_verify_locations(cadata="\n".join(_pems(options.ca)))
+        rejecting = options.reject_unauthorized is not False
+        context.verify_mode = ssl.CERT_REQUIRED if rejecting else ssl.CERT_OPTIONAL
+    if options.alpn_protocols is not None:
+        context.set_alpn_protocols(list(options.alpn_protocols))
+    return context