refactor(host): run the upgrade and stream selection on the real negotiate

packages/libp2p/host/Libp2p/Host.lean

@@ -5,6 +5,8 @@ import Libp2p.Host.Dedup
 import Libp2p.Host.Peerstore
 import Libp2p.Host.Policy
 import Libp2p.Host.Gate
+import Libp2p.Host.Upgrade
+import Libp2p.Host.StreamSelect
 import Libp2p.Host.Swarm
 import Libp2p.Host.Shutdown
 

packages/libp2p/host/Libp2p/Host/StreamSelect.lean

@@ -0,0 +1,82 @@
+import Streams
+import Bytes
+import Libp2p.Core
+import Libp2p.Negotiate
+
+/-!
+# Protocol selection on a muxed stream
+
+Every stream the muxer opens starts with its own multistream-select
+1.0.0 exchange (the connection spec, "Opening new streams over a
+multiplexed connection"): the initiator proposes the protocol ids in
+order, and the responder answers with the first one in its handler
+table. This module runs that exchange on a `Streams.Stream`.
+
+The negotiate package drives the same pure core over a
+`Streams.Connection` (`Negotiate.negotiate`), for the security and muxer
+stages of the upgrade. A muxed stream is a different seam: it reads and
+writes through `Except` answers, and it carries a read-back buffer for
+the leftover. The host owns this shell, and it reuses the core
+unchanged: `initiatorStart`, `responderStart`, `drive`, and
+`mapNegotiateError` all come from `Libp2p.Negotiate`.
+
+The bytes that follow the selection on the same read belong to the
+protocol. `Stream.prependBytes` hands them on to the handler
+(ARCHITECTURE §5.6).
+
+The read parks for as long as the peer stays silent, so no measure
+bounds the loop (root §3.3). The `repeat` form keeps the definition free
+of a `partial` marker, as in `Negotiate.negotiate`.
+-/
+
+set_option autoImplicit false
+
+namespace Libp2p.Host
+
+open Libp2p.Negotiate (initiatorStart responderStart drive mapNegotiateError)
+
+/-- Write each frame payload of the core to the stream, in order. A
+failed write is the `closed` arm: the stream is gone. -/
+private def sendWrites (stream : Streams.Stream) (ws : List (List Nat)) :
+    ExceptT Libp2p.Core.Error Conc.Async Unit := do
+  for w in ws do
+    if let .error e ← stream.write (Bytes.ofNatList w) then
+      throw (.closed s!"stream negotiation write failed: {repr e}")
+
+/-- Negotiate a protocol on a muxed stream. The answer is the selected
+id and the stream with the bytes after the selection put back in front
+of its reads. -/
+def negotiateStream (role : Libp2p.Core.Role) (stream : Streams.Stream)
+    (ids : List Libp2p.Core.ProtocolId) :
+    ExceptT Libp2p.Core.Error Conc.Async (Libp2p.Core.ProtocolId × Streams.Stream) := do
+  let start := match role with
+    | .initiator => initiatorStart ids
+    | .responder => responderStart ids
+  let (initial, first) ← match start with
+    | .error e => throw (mapNegotiateError e)
+    | .ok pair => pure pair
+  sendWrites stream first
+  let mut state := initial
+  let mut input : List Nat := []
+  repeat
+    match drive 1 state input with
+    | .error e => throw (mapNegotiateError e)
+    | .ok (.selected p _ rest writes) =>
+        sendWrites stream writes
+        return (p, ← Streams.Stream.prependBytes stream (Bytes.ofNatList rest))
+    | .ok (.running s _ rest writes) =>
+        sendWrites stream writes
+        state := s
+        input := rest
+    | .ok (.waiting s _ rest writes) =>
+        sendWrites stream writes
+        state := s
+        match ← stream.read with
+        | .error e => throw (.closed s!"stream negotiation read failed: {repr e}")
+        | .ok none => throw (.closed "stream ended during negotiation")
+        | .ok (some bytes) => input := rest ++ Bytes.toNatList bytes
+  -- Every arm above returns or loops, so the `repeat` never falls
+  -- through. The tail gives the `do` block its result type.
+  throw (.closed "stream negotiation ended without a selection")
+
+end Libp2p.Host

packages/libp2p/host/Libp2p/Host/Swarm.lean

@@ -3,8 +3,9 @@ import Conc
 import Streams
 import Libp2p.Identity
 import Libp2p.Core
-import Libp2p.Negotiate.Legacy
 import Libp2p.Host.Error
+import Libp2p.Host.StreamSelect
+import Libp2p.Host.Upgrade
 import Libp2p.Host.Registry
 import Libp2p.Host.Dedup
 import Libp2p.Host.Peerstore
@@ -36,7 +37,7 @@ ARCHITECTURE, "Refinements fixed at implementation (unit 1b)"):
 `allowDial` before a dial, `allowInbound` after an accept, the
 resource manager's `openConnection` before the upgrade, and
 `allowSecured` and `ConnectionScope.setPeer` inside the upgrade, after
-the handshake and before the muxer stage (`Negotiate.upgradeWith`). A
+the handshake and before the muxer stage (`Libp2p.Host.upgrade`). A
 refusal is `Error.gated` with the checkpoint.
 
 ## Ownership
@@ -69,7 +70,7 @@ set_option autoImplicit false
 
 namespace Libp2p.Host
 
-open Conc Streams Multiformats Libp2p.Identity Libp2p.Core Libp2p.Negotiate.Legacy
+open Conc Streams Multiformats Libp2p.Identity Libp2p.Core
 
 /-- A failure of a task that no caller awaits: an inbound upgrade, a
 session that ended badly, an inbound stream. -/
@@ -242,9 +243,9 @@ private def admit (h : Host) (role : Role) (scope : ConnectionScope) (conn : Con
 /-! ## The upgrade -/
 
 /-- Run the upgrade of `conn` as a pending connection: registered for
-the shutdown (§5.5), under the deadline of the upgrader, with the
-checkpoint. Every failure releases `scope`; the upgrader has already
-closed the connection. -/
+the shutdown (§5.5), under the deadline of the upgrade, with the
+checkpoint. Every failure releases `scope`. On a failure, the upgrade
+already closed the connection. -/
 private def upgradeConn (h : Host) (role : Role) (conn : Connection) (expected : Option PeerId)
     (scope : ConnectionScope) : Async (Except Error Upgraded) := do
   let key ← (h.draw : BaseIO Nat)
@@ -257,7 +258,7 @@ private def upgradeConn (h : Host) (role : Role) (conn : Connection) (expected :
     conn.closeQuietly
   let unregister : BaseIO Unit := h.pending.update (·.map (·.erase key))
   let r ← try
-      (upgradeAdmit h.config.clock role conn .empty expected (h.admit role scope conn)
+      (upgrade h.config.clock role conn expected (h.admit role scope conn)
         h.config.security h.config.muxers).run
     catch e =>
       unregister

packages/libp2p/host/Libp2p/Host/Upgrade.lean

@@ -0,0 +1,144 @@
+import Conc
+import Streams
+import Libp2p.Identity
+import Libp2p.Core
+import Libp2p.Negotiate
+import Libp2p.Host.Error
+
+/-!
+# The host's upgrade
+
+The upgrade of one connection, with the peer checkpoints in the place
+the core seam fixes: `ConnectionGater.allowSecured` runs after the
+handshake and before the muxer stage (`Libp2p/Core/Policy.lean`), as
+go-libp2p's upgrader calls `InterceptSecured`
+(`vendor/go-libp2p/p2p/net/upgrader/upgrader.go`, line 176).
+
+`Negotiate.upgradeWith` runs the same stages with no checkpoint. The
+host composes the stages from the public functions of
+`libp2p/negotiate` instead: `negotiate` for the security and the muxer
+stages, `checkIdentity`, and `settleUpgrade`. The stage order follows
+`Negotiate.upgradeCore`. A change to the order in one module needs the
+same change in the other.
+
+The whole upgrade runs under one `Conc.withinDeadline` of
+`Streams.unauthDeadlineMs` (root §5, invariant 4). A refusal at a
+checkpoint is an answer, and no failure: the caller maps it to
+`Host.Error.gated`.
+-/
+
+set_option autoImplicit false
+
+namespace Libp2p.Host
+
+open Conc Streams Libp2p.Identity Libp2p.Core
+
+/-- An upgraded connection: the muxer session and the peer the
+handshake proved. The session's `join` also waits for the secure
+connection under it, so `Upgraded.join` is the join of both. -/
+structure Upgraded where
+  /-- The muxer session. -/
+  session : MuxSession
+  /-- The peer the handshake proved. -/
+  remotePeer : PeerId
+
+/-- Wait for the session and the secure connection under it. -/
+def Upgraded.join (u : Upgraded) : ExceptT Core.Error Async Unit :=
+  u.session.join
+
+/-- The answer of an upgrade that ran to its end: a session, or the
+checkpoint that refused the proved peer. -/
+inductive Admission where
+  | upgraded (u : Upgraded)
+  | refused (at_ : Checkpoint)
+
+/-- Close a secure connection and wait for its tasks. The function
+drops the answer of the join, because the caller already holds the
+answer it reports. -/
+private def discardSecure (sc : SecureConn) : Async Unit := do
+  sc.toConnection.closeQuietly
+  discard sc.join.run
+
+/-- The session `s`, with a `join` that also waits for the secure
+connection under it. The join waits for the secure connection even
+when the session's join throws a fault, so none of its tasks outlives
+the join (PRINCIPLES 3.2). The first failure wins. -/
+private def ownSecure (s : MuxSession) (sc : SecureConn) : MuxSession :=
+  { s with
+    join := ExceptT.mk do
+      let a : Except IO.Error (Except Core.Error Unit) ←
+        try pure (.ok (← s.join.run)) catch e => pure (.error e)
+      let b ← sc.join.run
+      match a with
+      | .error e => throw e
+      | .ok a => return a.bind fun () => b }
+
+/-- The stages: negotiate the security protocol, run the handshake,
+check the identity, ask `admit` about the proved peer, negotiate the
+muxer, open the session. A failure or a refusal after the handshake
+closes and joins the secure connection. -/
+private def stages (role : Role) (conn : Connection) (expected : Option PeerId)
+    (admit : PeerId → BaseIO (Option Checkpoint))
+    (secs : List SecurityProtocol) (muxers : List Muxer) :
+    Async (Except Core.Error Admission) := do
+  match ← (Negotiate.negotiate role (BufReader.new conn) conn .empty (secs.map (·.id))).run with
+  | .error e => return .error e
+  | .ok stage =>
+    let some sec := secs.find? (fun s => decide (s.id = stage.proto))
+      | return .error (.protocol "selected security protocol is not configured")
+    match ← (sec.handshake role conn stage.leftover).run with
+    | .error e => return .error e
+    | .ok sc =>
+      if let .error e := Negotiate.checkIdentity expected sc.remotePeer then
+        discardSecure sc
+        return .error e
+      -- The checkpoint: the handshake proved the peer, and no muxer
+      -- byte moves before this point.
+      if let some c ← admit sc.remotePeer then
+        discardSecure sc
+        return .ok (.refused c)
+      -- The secure `read` delivers the handshake's leftover, so the
+      -- muxer stage starts with no bytes.
+      let secure := sc.toConnection
+      match ← (Negotiate.negotiate role (BufReader.new secure) secure .empty
+          (muxers.map (·.id))).run with
+      | .error e => discardSecure sc; return .error e
+      | .ok stage2 =>
+        let some mux := muxers.find? (fun m => decide (m.id = stage2.proto))
+          | discardSecure sc
+            return .error (.protocol "selected muxer is not configured")
+        match ← (mux.«open» role secure stage2.leftover).run with
+        | .error e => discardSecure sc; return .error e
+        | .ok session =>
+          return .ok (.upgraded { session := ownSecure session sc, remotePeer := sc.remotePeer })
+
+/-- Close a session that finished after the timer won, and give its
+join to `late`, the task set `upgrade` joins. No caller receives that
+session. -/
+private def releaseLate (late : TaskSet) : Except Core.Error Admission → BaseIO Unit
+  | .ok (.upgraded u) => do
+    discard u.session.close.toBaseIO
+    late.spawn (discard u.join.run)
+  | _ => pure ()
+
+/-- The upgrade of `conn` under the deadline, on `clock`. The expiry
+closes the connection once, at the boundary. `Negotiate.settleUpgrade`
+maps the expiry to the `timeout` arm. On a failure inside the stages,
+this function closes the connection. -/
+def upgrade (clock : Clock) (role : Role) (conn : Connection) (expected : Option PeerId)
+    (admit : PeerId → BaseIO (Option Checkpoint))
+    (secs : List SecurityProtocol) (muxers : List Muxer) :
+    ExceptT Core.Error Async Admission := do
+  let late ← TaskSet.new
+  let answer ← withinDeadline clock Streams.unauthDeadlineMs conn.closeQuietly
+    (stages role conn expected admit secs muxers) (releaseLate late)
+  late.joinAll
+  match Negotiate.settleUpgrade answer with
+  | .ok a => return a
+  | .error e =>
+    -- An error from the boundary means the expiry hook closed the
+    -- connection. A failure inside the work leaves it open.
+    if answer matches .ok _ then conn.closeQuietly
+    throw e
+
+end Libp2p.Host

packages/libp2p/host/docs/ARCHITECTURE.md

@@ -24,7 +24,8 @@ libp2p group is §7, row 10).
 | `Libp2p/Host/Swarm.lean` | `Config`, the `Host` state, `dial`, `dialPeer`, `startListening`, `newStream`, the accept loop, the upgrade pipeline, and stream serving |
 | `Libp2p/Host/Shutdown.lean` | `Host.close`, the §5.5 sequence |
 | `Libp2p/Host/Policy.lean` | the permissive default gater, and the permissive and denying test doubles the core pair assigned here; the default resource manager (admits everything) and the default ranker (keeps the order, no delay) |
-| `Libp2p/Host/Mock/Negotiate.lean` | the mock of the W5.5 surface the host calls, until the host swap of issue #70 (see "The mock") |
+| `Libp2p/Host/Upgrade.lean` | `upgrade` and `Upgraded`: the stages of `libp2p/negotiate`, with the peer checkpoint between the handshake and the muxer stage, under one deadline |
+| `Libp2p/Host/StreamSelect.lean` | `negotiateStream`, multistream-select on a muxed stream, a shell over the pure core of `libp2p/negotiate` |
 
 All declarations live in the `Libp2p.Host` namespace, and the root
 module re-exports the public submodules.
@@ -237,9 +238,16 @@ fixes them as follows.
   `allowInbound`, then `openConnection`, after an accept
   (`vendor/go-libp2p/p2p/net/upgrader/listener.go`, lines 183 and
   193); and `allowSecured`, then `ConnectionScope.setPeer`, inside the
-  upgrade, through `Negotiate.upgradeWith`
-  (`vendor/go-libp2p/p2p/net/upgrader/upgrader.go`, line 176). A
-  refusal is `gated` with the checkpoint, and the scope is released.
+  upgrade, after the handshake and before the muxer stage
+  (`vendor/go-libp2p/p2p/net/upgrader/upgrader.go`, line 176, and the
+  core seam `ConnectionGater.allowSecured`). A refusal is `gated` with
+  the checkpoint, and the host releases the scope.
+- **The host composes the upgrade.** `Negotiate.upgradeWith` takes no
+  checkpoint. `Libp2p/Host/Upgrade.lean` runs the stages from the
+  public functions of `libp2p/negotiate`: `negotiate`, `checkIdentity`,
+  and `settleUpgrade`. The stages keep the order of the upgrader, and
+  the checkpoint sits between the handshake and the muxer stage. One
+  `withinDeadline` bounds the whole upgrade, as in `upgradeWith`.
   The policy shapes are this unit's refinement of the core seams (the
   core ARCHITECTURE, "The policy shapes").
 - **The answers of the dial path.** `dial` and `dialPeer` answer the
@@ -265,9 +273,13 @@ fixes them as follows.
 - **Stream serving.** The serve task of a session accepts the peer's
   streams, and each stream runs in a task that the serve task owns and
   joins when the session ends. The stream negotiates against the
-  handler ids through `Negotiate.negotiateStream`, which puts the
-  bytes it read past the selecting message at the front of the stream
-  (§4.7). A failed negotiation, an id whose handler left the table,
+  handler ids through `negotiateStream`, which puts the bytes it read
+  past the selecting message at the front of the stream (§4.7). The
+  upgrader's `Negotiate.negotiate` runs over a `Connection`, and a
+  muxed stream is a `Stream`, so the host holds this shell
+  (`Libp2p/Host/StreamSelect.lean`). It drives the same pure core:
+  `initiatorStart`, `responderStart`, and `drive`. A failed
+  negotiation, an id whose handler left the table,
   and a handler failure reset the stream; the session stays. go-libp2p
   resets a stream whose negotiation fails in the same place
   (`vendor/go-libp2p/p2p/host/basic/basic_host.go`, line 347). A
@@ -298,29 +310,13 @@ fixes them as follows.
   1b milestone. A fault of a joined task does not stop the sequence,
   and `close` throws the first fault at the end.
 
-## The mock
-
-The seam types come from `libp2p/core`. `libp2p/negotiate` takes the
-real core too (issue #53), so its upgrader takes the values the host
-holds. The host keeps a mock of the negotiate surface,
-`Libp2p/Host/Mock/Negotiate.lean`, until the host swap of issue #70.
-Its operations are `opaque`, and each body fails with the name of the
-package; it states no axiom. Swarm opens `Libp2p.Host.Mock`, so
-`Negotiate.upgradeWith` resolves as it will against the real package.
-The swap (issue #70) is a retype: the merged upgrader takes no `admit`
-checkpoint and exports no `negotiateStream`.
-
-Every upgrade runs through the negotiate mock and fails there, so the
-unit 1b milestones run after the host swap of issue #70. The Noise and
-yamux values are on the real core (issues #76 and #77). Unit 1a runs
-in full.
-
 ## What the package must not depend on
 
 - Inside libp2p, the chain puts everything below the host at its
   disposal (§2, §7 row 10). The host imports `identity` for the id and
   key types, `core` for the seam types, and `negotiate` for the
-  upgrader and the stream-negotiation responder its paths call. The
+  stages `upgrade` composes and the pure core that `negotiateStream`
+  drives. The
   security and mux implementations
   arrive inside the configuration lists (§3.1), so a host module that
   names a concrete seam value is a review failure. Test modules import