starpc 0.6.2 → 0.8.0

package/README.md

@@ -121,7 +121,7 @@ import { pushable } from 'it-pushable'
 const mux = createMux()
 const echoer = new EchoerServer()
 mux.register(createHandler(EchoerDefinition, echoer))
-const server = new Server(mux)
+const server = new Server(mux.lookupMethodFunc)
 
 const clientConn = new Conn()
 const serverConn = new Conn(server)

package/dist/e2e/e2e.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/e2e/e2e.js

@@ -0,0 +1,24 @@
+import { pipe } from 'it-pipe';
+import { createHandler, createMux, Server, Client, Conn } from '../srpc';
+import { EchoerDefinition, EchoerServer, runClientTest } from '../echo';
+import { runRpcStreamTest } from '../echo/client-test';
+async function runRPC() {
+    const mux = createMux();
+    const server = new Server(mux.lookupMethodFunc);
+    const echoer = new EchoerServer(server);
+    mux.register(createHandler(EchoerDefinition, echoer));
+    const clientConn = new Conn();
+    const serverConn = new Conn(server);
+    pipe(clientConn, serverConn, clientConn);
+    const client = new Client(clientConn.buildOpenStreamFunc());
+    await runRpcStreamTest(client);
+    await runClientTest(client);
+}
+runRPC()
+    .then(() => {
+    console.log('finished successfully');
+})
+    .catch((err) => {
+    console.error(err);
+    process.exit(1);
+});

package/dist/echo/client-test.js

@@ -34,4 +34,6 @@ export async function runRpcStreamTest(client) {
     console.log('Calling Echo via RPC stream...');
     const resp = await proxiedService.Echo({ body: 'hello world via proxy' });
     console.log('rpc stream test: succeeded: response: ' + resp.body);
+    console.log('Running client test over RPC stream...');
+    await runClientTest(proxiedClient);
 }

package/dist/integration/integration.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/integration/integration.js

@@ -0,0 +1,22 @@
+import { WebSocketConn } from '../srpc/websocket.js';
+import { runClientTest, runRpcStreamTest } from '../echo/client-test.js';
+import WebSocket from 'isomorphic-ws';
+async function runRPC() {
+    const addr = 'ws://localhost:5000/demo';
+    console.log(`Connecting to ${addr}`);
+    const ws = new WebSocket(addr);
+    const channel = new WebSocketConn(ws);
+    const client = channel.buildClient();
+    console.log('Running client test via WebSocket..');
+    await runClientTest(client);
+    console.log('Running RpcStream test via WebSocket..');
+    await runRpcStreamTest(client);
+}
+runRPC()
+    .then(() => {
+    process.exit(0);
+})
+    .catch((err) => {
+    console.error(err);
+    process.exit(1);
+});

package/dist/srpc/mux.d.ts

@@ -1,11 +1,16 @@
 import { InvokeFn, Handler } from './handler';
+export declare type LookupMethod = (serviceID: string, methodID: string) => Promise<InvokeFn | null>;
 export interface Mux {
-    register(handler: Handler): void;
     lookupMethod(serviceID: string, methodID: string): Promise<InvokeFn | null>;
 }
-export declare function createMux(): Mux;
+export declare function createMux(): StaticMux;
 export declare class StaticMux implements Mux {
     private services;
+    private lookups;
+    get lookupMethodFunc(): LookupMethod;
     register(handler: Handler): void;
+    registerLookupMethod(lookupMethod: LookupMethod): void;
     lookupMethod(serviceID: string, methodID: string): Promise<InvokeFn | null>;
+    private lookupViaMap;
+    private lookupViaLookups;
 }

package/dist/srpc/mux.js

@@ -1,4 +1,4 @@
-// createMux builds a new Mux.
+// createMux builds a new StaticMux.
 export function createMux() {
     return new StaticMux();
 }
@@ -8,6 +8,13 @@ export class StaticMux {
     constructor() {
         // services contains a mapping from service id to handlers.
         this.services = {};
+        // lookups is the list of lookup methods to call.
+        // called if the method is not resolved by the services list.
+        this.lookups = [];
+    }
+    // lookupMethodFunc implements the LookupMethod type.
+    get lookupMethodFunc() {
+        return this.lookupMethod.bind(this);
     }
     register(handler) {
         const serviceID = handler?.getServiceID();
@@ -21,10 +28,21 @@ export class StaticMux {
         }
         this.services[serviceID] = serviceMethods;
     }
+    // registerLookupMethod registers a extra lookup function to the mux.
+    registerLookupMethod(lookupMethod) {
+        this.lookups.push(lookupMethod);
+    }
     async lookupMethod(serviceID, methodID) {
-        if (!serviceID) {
-            return null;
+        if (serviceID) {
+            const invokeFn = await this.lookupViaMap(serviceID, methodID);
+            if (invokeFn) {
+                return invokeFn;
+            }
         }
+        return await this.lookupViaLookups(serviceID, methodID);
+    }
+    // lookupViaMap looks up the method via the services map.
+    async lookupViaMap(serviceID, methodID) {
         const serviceMethods = this.services[serviceID];
         if (!serviceMethods) {
             return null;
@@ -35,4 +53,14 @@ export class StaticMux {
         }
         return await handler.lookupMethod(serviceID, methodID);
     }
+    // lookupViaLookups looks up the method via the lookup funcs.
+    async lookupViaLookups(serviceID, methodID) {
+        for (const lookupMethod of this.lookups) {
+            const invokeFn = await lookupMethod(serviceID, methodID);
+            if (invokeFn) {
+                return invokeFn;
+            }
+        }
+        return null;
+    }
 }

package/dist/srpc/server-rpc.d.ts

@@ -1,9 +1,9 @@
 import type { CallData, CallStart } from './rpcproto.pb.js';
 import { CommonRPC } from './common-rpc.js';
-import { Mux } from './mux.js';
+import { LookupMethod } from './mux.js';
 export declare class ServerRPC extends CommonRPC {
-    private mux;
-    constructor(mux: Mux);
+    private lookupMethod;
+    constructor(lookupMethod: LookupMethod);
     handleCallStart(packet: Partial<CallStart>): Promise<void>;
     handleCallData(packet: Partial<CallData>): Promise<void>;
     private invokeRPC;

package/dist/srpc/server-rpc.js

@@ -1,9 +1,9 @@
 import { CommonRPC } from './common-rpc.js';
 // ServerRPC is an ongoing RPC from the server side.
 export class ServerRPC extends CommonRPC {
-    constructor(mux) {
+    constructor(lookupMethod) {
         super();
-        this.mux = mux;
+        this.lookupMethod = lookupMethod;
     }
     // handleCallStart handles a CallStart cket.
     async handleCallStart(packet) {
@@ -15,7 +15,7 @@ export class ServerRPC extends CommonRPC {
         if (!this.service || !this.method) {
             throw new Error('rpcService and rpcMethod cannot be empty');
         }
-        const methodDef = await this.mux.lookupMethod(this.service, this.method);
+        const methodDef = await this.lookupMethod(this.service, this.method);
         if (!methodDef) {
             throw new Error(`not found: ${this.service}/${this.method}`);
         }

package/dist/srpc/server.d.ts

@@ -1,13 +1,13 @@
 import { Stream } from '@libp2p/interface-connection';
 import { Duplex } from 'it-stream-types';
-import { Mux } from './mux.js';
+import { LookupMethod } from './mux.js';
 import { ServerRPC } from './server-rpc.js';
 import { Packet } from './rpcproto.pb.js';
 import { StreamHandler } from './conn.js';
 import { RpcStreamHandler } from '../rpcstream/rpcstream.js';
 export declare class Server implements StreamHandler {
-    private mux;
-    constructor(mux: Mux);
+    private lookupMethod;
+    constructor(lookupMethod: LookupMethod);
     get rpcStreamHandler(): RpcStreamHandler;
     startRpc(): ServerRPC;
     handleStream(stream: Stream): ServerRPC;

package/dist/srpc/server.js

@@ -3,8 +3,8 @@ import { ServerRPC } from './server-rpc.js';
 import { parseLengthPrefixTransform, prependLengthPrefixTransform, decodePacketSource, encodePacketSource, } from './packet.js';
 // Server implements the SRPC server in TypeScript with a Mux.
 export class Server {
-    constructor(mux) {
-        this.mux = mux;
+    constructor(lookupMethod) {
+        this.lookupMethod = lookupMethod;
     }
     // rpcStreamHandler implements the RpcStreamHandler interface.
     get rpcStreamHandler() {
@@ -13,7 +13,7 @@ export class Server {
     // startRpc starts a new server-side RPC.
     // the returned RPC handles incoming Packets.
     startRpc() {
-        return new ServerRPC(this.mux);
+        return new ServerRPC(this.lookupMethod);
     }
     // handleStream handles an incoming Uint8Array message duplex.
     handleStream(stream) {

package/e2e/e2e.ts

@@ -5,7 +5,7 @@ import { runRpcStreamTest } from '../echo/client-test'
 
 async function runRPC() {
   const mux = createMux()
-  const server = new Server(mux)
+  const server = new Server(mux.lookupMethodFunc)
   const echoer = new EchoerServer(server)
   mux.register(createHandler(EchoerDefinition, echoer))
 

package/echo/client-test.ts

@@ -40,7 +40,11 @@ export async function runRpcStreamTest(client: Client) {
   )
   const proxiedClient = new Client(openStreamFn)
   const proxiedService = new EchoerClientImpl(proxiedClient)
+
   console.log('Calling Echo via RPC stream...')
   const resp = await proxiedService.Echo({ body: 'hello world via proxy' })
   console.log('rpc stream test: succeeded: response: ' + resp.body)
+
+  console.log('Running client test over RPC stream...')
+  await runClientTest(proxiedClient)
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "starpc",
-  "version": "0.6.2",
+  "version": "0.8.0",
   "description": "Streaming protobuf RPC service protocol over any two-way channel.",
   "license": "MIT",
   "author": {
@@ -34,7 +34,8 @@
   },
   "scripts": {
     "build": "rimraf ./dist && tsc --project tsconfig.build.json --outDir ./dist/",
-    "check": "tsc",
+    "check": "npm run typecheck",
+    "typecheck": "tsc --noEmit",
     "deps": "depcheck --ignores 'bufferutil,utf-8-validate'",
     "codegen": "npm run gen",
     "ci": "npm run build && npm run lint:js && npm run lint:go",

package/patches/ts-poet+4.15.0.patch

@@ -1,8 +1,16 @@
 diff --git a/node_modules/ts-poet/build/Import.js b/node_modules/ts-poet/build/Import.js
-index b92bea3..608add9 100644
+index b92bea3..d35c845 100644
 --- a/node_modules/ts-poet/build/Import.js
 +++ b/node_modules/ts-poet/build/Import.js
-@@ -293,6 +293,14 @@ function emitImports(imports, ourModulePath, importMappings) {
+@@ -26,6 +26,7 @@ exports.sameModule = exports.maybeRelativePath = exports.emitImports = exports.S
+ const lodash_1 = __importDefault(require("lodash"));
+ const path = __importStar(require("path"));
+ const Node_1 = require("./Node");
++const fs = require('fs');
+ const typeImportMarker = '(?:t:)?';
+ const fileNamePattern = '(?:[a-zA-Z0-9._-]+)';
+ const modulePattern = `@?(?:(?:${fileNamePattern}(?:/${fileNamePattern})*))`;
+@@ -293,6 +294,16 @@ function emitImports(imports, ourModulePath, importMappings) {
          return '';
      }
      let result = '';
@@ -13,11 +21,13 @@ index b92bea3..608add9 100644
 +    // github.com/aperturerobotics/protobuf-project/example/example -> ./example/example
 +    if (thisProject && ourModuleImportPath.startsWith(thisProject)) {
 +        ourModuleImportPath = './' + ourModuleImportPath.substr(thisProject.length + 1);
++    } else {
++      ourModuleImportPath = './vendor/' + ourModuleImportPath;
 +    }
      const augmentImports = lodash_1.default.groupBy(filterInstances(imports, Augmented), (a) => a.augmented);
      // Group the imports by source module they're imported from
      const importsByModule = lodash_1.default.groupBy(imports.filter((it) => it.source !== undefined &&
-@@ -307,7 +315,14 @@ function emitImports(imports, ourModulePath, importMappings) {
+@@ -307,7 +318,14 @@ function emitImports(imports, ourModulePath, importMappings) {
          if (modulePath in importMappings) {
              modulePath = importMappings[modulePath];
          }

package/srpc/client-rpc.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
 )
 
 // ClientRPC represents the client side of an on-going RPC call message stream.
@@ -101,6 +102,16 @@ func (r *ClientRPC) HandlePacketData(data []byte) error {
 	return r.HandlePacket(pkt)
 }
 
+// HandleStreamClose handles the incoming stream closing w/ optional error.
+func (r *ClientRPC) HandleStreamClose(closeErr error) {
+	if closeErr != nil {
+		if r.serverErr == nil {
+			r.serverErr = closeErr
+		}
+		r.Close()
+	}
+}
+
 // HandlePacket handles an incoming parsed message packet.
 // Not concurrency safe: use a mutex if calling concurrently.
 func (r *ClientRPC) HandlePacket(msg *Packet) error {
@@ -130,6 +141,7 @@ func (r *ClientRPC) HandleCallData(pkt *CallData) error {
 		return ErrCompleted
 	}
 
+	logrus.Infof("DEBUG: handling call data for %s/%s: %s", r.service, r.method, pkt)
 	if data := pkt.GetData(); len(data) != 0 || pkt.GetDataIsZero() {
 		select {
 		case <-r.ctx.Done():

package/srpc/client.go

@@ -19,7 +19,11 @@ type Client interface {
 
 // OpenStreamFunc opens a stream with a remote.
 // msgHandler must not be called concurrently.
-type OpenStreamFunc = func(ctx context.Context, msgHandler PacketHandler) (Writer, error)
+type OpenStreamFunc = func(
+	ctx context.Context,
+	msgHandler PacketHandler,
+	closeHandler CloseHandler,
+) (Writer, error)
 
 // client implements Client with a transport.
 type client struct {
@@ -44,7 +48,7 @@ func (c *client) Invoke(rctx context.Context, service, method string, in, out Me
 		return err
 	}
 	clientRPC := NewClientRPC(ctx, service, method)
-	writer, err := c.openStream(ctx, clientRPC.HandlePacket)
+	writer, err := c.openStream(ctx, clientRPC.HandlePacket, clientRPC.HandleStreamClose)
 	if err != nil {
 		return err
 	}
@@ -81,7 +85,7 @@ func (c *client) NewStream(ctx context.Context, service, method string, firstMsg
 	}
 
 	clientRPC := NewClientRPC(ctx, service, method)
-	writer, err := c.openStream(ctx, clientRPC.HandlePacket)
+	writer, err := c.openStream(ctx, clientRPC.HandlePacket, clientRPC.HandleStreamClose)
 	if err != nil {
 		return nil, err
 	}

package/srpc/mux.ts

@@ -1,16 +1,17 @@
 import { InvokeFn, Handler } from './handler'
 
+// LookupMethod is a function to lookup a RPC method.
+export type LookupMethod = (serviceID: string, methodID: string) => Promise<InvokeFn | null>
+
 // Mux contains a set of <service, method> handlers.
 export interface Mux {
-  // register registers a new RPC method handler.
-  register(handler: Handler): void
   // lookupMethod looks up the method matching the service & method ID.
   // returns null if not found.
   lookupMethod(serviceID: string, methodID: string): Promise<InvokeFn | null>
 }
 
-// createMux builds a new Mux.
-export function createMux(): Mux {
+// createMux builds a new StaticMux.
+export function createMux(): StaticMux {
   return new StaticMux()
 }
 
@@ -22,6 +23,14 @@ type staticMuxMethods = { [methodID: string]: Handler }
 export class StaticMux implements Mux {
   // services contains a mapping from service id to handlers.
   private services: { [id: string]: staticMuxMethods } = {}
+  // lookups is the list of lookup methods to call.
+  // called if the method is not resolved by the services list.
+  private lookups: LookupMethod[] = []
+
+  // lookupMethodFunc implements the LookupMethod type.
+  public get lookupMethodFunc(): LookupMethod {
+    return this.lookupMethod.bind(this)
+  }
 
   public register(handler: Handler): void {
     const serviceID = handler?.getServiceID()
@@ -36,13 +45,27 @@ export class StaticMux implements Mux {
     this.services[serviceID] = serviceMethods
   }
 
+  // registerLookupMethod registers a extra lookup function to the mux.
+  public registerLookupMethod(lookupMethod: LookupMethod) {
+    this.lookups.push(lookupMethod)
+  }
+
   public async lookupMethod(
     serviceID: string,
     methodID: string
   ): Promise<InvokeFn | null> {
-    if (!serviceID) {
-      return null
+    if (serviceID) {
+      const invokeFn = await this.lookupViaMap(serviceID, methodID)
+      if (invokeFn) {
+        return invokeFn
+      }
     }
+
+    return await this.lookupViaLookups(serviceID, methodID)
+  }
+
+  // lookupViaMap looks up the method via the services map.
+  private async lookupViaMap(serviceID: string, methodID: string): Promise<InvokeFn | null> {
     const serviceMethods = this.services[serviceID]
     if (!serviceMethods) {
       return null
@@ -53,4 +76,16 @@ export class StaticMux implements Mux {
     }
     return await handler.lookupMethod(serviceID, methodID)
   }
+
+  // lookupViaLookups looks up the method via the lookup funcs.
+  private async lookupViaLookups(serviceID: string, methodID: string): Promise<InvokeFn | null> {
+    for (const lookupMethod of this.lookups) {
+      const invokeFn = await lookupMethod(serviceID, methodID)
+      if (invokeFn) {
+        return invokeFn
+      }
+    }
+
+    return null
+  }
 }

package/srpc/muxed-conn.go

@@ -14,18 +14,13 @@ func NewClientWithMuxedConn(conn network.MuxedConn) Client {
 
 // NewOpenStreamWithMuxedConn constructs a OpenStream func with a MuxedConn.
 func NewOpenStreamWithMuxedConn(conn network.MuxedConn) OpenStreamFunc {
-	return func(ctx context.Context, msgHandler PacketHandler) (Writer, error) {
+	return func(ctx context.Context, msgHandler PacketHandler, closeHandler CloseHandler) (Writer, error) {
 		mstrm, err := conn.OpenStream(ctx)
 		if err != nil {
 			return nil, err
 		}
 		rw := NewPacketReadWriter(mstrm)
-		go func() {
-			err := rw.ReadPump(msgHandler)
-			if err != nil {
-				_ = rw.Close()
-			}
-		}()
+		go rw.ReadPump(msgHandler, closeHandler)
 		return rw, nil
 	}
 }

package/srpc/packet-rw.go

@@ -39,10 +39,23 @@ func (r *PacketReaderWriter) WritePacket(p *Packet) error {
 }
 
 // ReadPump executes the read pump in a goroutine.
-func (r *PacketReaderWriter) ReadPump(cb PacketHandler) error {
+//
+// calls the handler when closed or returning an error
+func (r *PacketReaderWriter) ReadPump(cb PacketHandler, closed CloseHandler) {
+	err := r.ReadToHandler(cb)
+	// signal that the stream is now closed.
+	if closed != nil {
+		closed(err)
+	}
+}
+
+// ReadToHandler reads data to the given handler.
+// Does not handle closing the stream, use ReadPump instead.
+func (r *PacketReaderWriter) ReadToHandler(cb PacketHandler) error {
 	var currLen uint32
 	buf := make([]byte, 2048)
 	for {
+		// read some data into the buffer
 		n, err := r.rw.Read(buf)
 		if err != nil {
 			if err == io.EOF {
@@ -50,17 +63,20 @@ func (r *PacketReaderWriter) ReadPump(cb PacketHandler) error {
 			}
 			return err
 		}
+		// push the data to r.buf
 		_, err = r.buf.Write(buf[:n])
 		if err != nil {
 			return err
 		}
 
-		// check if we have enough for a length prefix
+		// check if we have enough data for a length prefix
 		bufLen := r.buf.Len()
+		if bufLen < 4 {
+			continue
+		}
+
+		// parse the length prefix if not done already
 		if currLen == 0 {
-			if bufLen < 4 {
-				continue
-			}
 			currLen = r.readLengthPrefix(r.buf.Bytes())
 			if currLen == 0 {
 				return errors.New("unexpected zero len prefix")
@@ -69,6 +85,8 @@ func (r *PacketReaderWriter) ReadPump(cb PacketHandler) error {
 				return errors.Errorf("message size %v greater than maximum %v", currLen, maxMessageSize)
 			}
 		}
+
+		// emit the packet if fully buffered
 		if currLen != 0 && bufLen >= int(currLen)+4 {
 			pkt := r.buf.Next(int(currLen + 4))[4:]
 			currLen = 0

package/srpc/packet.go

@@ -1,8 +1,14 @@
 package srpc
 
 // PacketHandler handles a packet.
+//
+// pkt is optional (can be nil)
+// if closeErr is set, the stream is closed after pkt.
 type PacketHandler = func(pkt *Packet) error
 
+// CloseHandler handles the stream closing with an optional error.
+type CloseHandler = func(closeErr error)
+
 // Validate performs cursory validation of the packet.
 func (p *Packet) Validate() error {
 	switch b := p.GetBody().(type) {
@@ -58,7 +64,7 @@ func NewCallDataPacket(data []byte, dataIsZero bool, complete bool, err error) *
 
 // Validate performs cursory validation of the packet.
 func (p *CallData) Validate() error {
-	if len(p.GetData()) == 0 && !p.GetComplete() && len(p.GetError()) == 0 {
+	if len(p.GetData()) == 0 && !p.GetComplete() && len(p.GetError()) == 0 && !p.GetDataIsZero() {
 		return ErrEmptyPacket
 	}
 	return nil

package/srpc/server-pipe.go

@@ -9,18 +9,13 @@ import (
 // Stream with the given Server. Starts read pumps for both. Starts the
 // HandleStream function on the server in a separate goroutine.
 func NewServerPipe(server *Server) OpenStreamFunc {
-	return func(ctx context.Context, msgHandler PacketHandler) (Writer, error) {
+	return func(ctx context.Context, msgHandler PacketHandler, closeHandler CloseHandler) (Writer, error) {
 		srvPipe, clientPipe := net.Pipe()
 		go func() {
 			_ = server.HandleStream(ctx, srvPipe)
 		}()
 		clientPrw := NewPacketReadWriter(clientPipe)
-		go func() {
-			err := clientPrw.ReadPump(msgHandler)
-			if err != nil {
-				_ = clientPrw.Close()
-			}
-		}()
+		go clientPrw.ReadPump(msgHandler, closeHandler)
 		return clientPrw, nil
 	}
 }

package/srpc/server-rpc.go

@@ -29,7 +29,7 @@ type ServerRPC struct {
 	dataChClosed bool
 	// clientErr is an error set by the client.
 	// before dataCh is closed, managed by HandlePacket.
-	// immutable after dataCh is closed.
+	// immutable after dataCh is closed or ctxCancel
 	clientErr error
 }
 
@@ -54,9 +54,32 @@ func (r *ServerRPC) Context() context.Context {
 	return r.ctx
 }
 
+// Wait waits for the RPC to finish.
+func (r *ServerRPC) Wait(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return context.Canceled
+	case <-r.ctx.Done():
+	}
+	return r.clientErr
+}
+
+// HandleStreamClose handles the incoming stream closing w/ optional error.
+func (r *ServerRPC) HandleStreamClose(closeErr error) {
+	if closeErr != nil {
+		if r.clientErr == nil {
+			r.clientErr = closeErr
+		}
+		r.Close()
+	}
+}
+
 // HandlePacket handles an incoming parsed message packet.
 // Not concurrency safe: use a mutex if calling concurrently.
 func (r *ServerRPC) HandlePacket(msg *Packet) error {
+	if msg == nil {
+		return nil
+	}
 	if err := msg.Validate(); err != nil {
 		return err
 	}
@@ -147,6 +170,9 @@ func (r *ServerRPC) invokeRPC() {
 // Close releases any resources held by the ServerRPC.
 // not concurrency safe with HandlePacket.
 func (r *ServerRPC) Close() {
+	if r.clientErr == nil {
+		r.clientErr = context.Canceled
+	}
 	r.ctxCancel()
 	if r.service == "" {
 		// invokeRPC has not been called, otherwise it would call Close()

package/srpc/server-rpc.ts

@@ -3,16 +3,16 @@ import type { Sink } from 'it-stream-types'
 import type { CallData, CallStart } from './rpcproto.pb.js'
 import { CommonRPC } from './common-rpc.js'
 import { InvokeFn } from './handler.js'
-import { Mux } from './mux.js'
+import { LookupMethod } from './mux.js'
 
 // ServerRPC is an ongoing RPC from the server side.
 export class ServerRPC extends CommonRPC {
-  // mux is used to handle incoming rpcs.
-  private mux: Mux
+  // lookupMethod looks up the incoming RPC methods.
+  private lookupMethod: LookupMethod
 
-  constructor(mux: Mux) {
+  constructor(lookupMethod: LookupMethod) {
     super()
-    this.mux = mux
+    this.lookupMethod = lookupMethod
   }
 
   // handleCallStart handles a CallStart cket.
@@ -25,7 +25,7 @@ export class ServerRPC extends CommonRPC {
     if (!this.service || !this.method) {
       throw new Error('rpcService and rpcMethod cannot be empty')
     }
-    const methodDef = await this.mux.lookupMethod(this.service, this.method)
+    const methodDef = await this.lookupMethod(this.service, this.method)
     if (!methodDef) {
       throw new Error(`not found: ${this.service}/${this.method}`)
     }

package/srpc/server.go

@@ -32,9 +32,8 @@ func (s *Server) HandleStream(ctx context.Context, rwc io.ReadWriteCloser) error
 	serverRPC := NewServerRPC(subCtx, s.mux)
 	prw := NewPacketReadWriter(rwc)
 	serverRPC.SetWriter(prw)
-	err := prw.ReadPump(serverRPC.HandlePacket)
-	_ = rwc.Close()
-	return err
+	go prw.ReadPump(serverRPC.HandlePacket, serverRPC.HandleStreamClose)
+	return serverRPC.Wait(ctx)
 }
 
 // AcceptMuxedConn runs a loop which calls Accept on a muxer to handle streams.

package/srpc/server.ts

@@ -2,7 +2,7 @@ import { Stream } from '@libp2p/interface-connection'
 import { Duplex } from 'it-stream-types'
 import { pipe } from 'it-pipe'
 
-import { Mux } from './mux.js'
+import { LookupMethod } from './mux.js'
 import { ServerRPC } from './server-rpc.js'
 import { Packet } from './rpcproto.pb.js'
 import {
@@ -16,11 +16,11 @@ import { RpcStreamHandler } from '../rpcstream/rpcstream.js'
 
 // Server implements the SRPC server in TypeScript with a Mux.
 export class Server implements StreamHandler {
-  // mux is the mux used to handle requests.
-  private mux: Mux
+  // lookupMethod looks up the incoming RPC methods.
+  private lookupMethod: LookupMethod
 
-  constructor(mux: Mux) {
-    this.mux = mux
+  constructor(lookupMethod: LookupMethod) {
+    this.lookupMethod = lookupMethod
   }
 
   // rpcStreamHandler implements the RpcStreamHandler interface.
@@ -31,7 +31,7 @@ export class Server implements StreamHandler {
   // startRpc starts a new server-side RPC.
   // the returned RPC handles incoming Packets.
   public startRpc(): ServerRPC {
-    return new ServerRPC(this.mux)
+    return new ServerRPC(this.lookupMethod)
   }
 
   // handleStream handles an incoming Uint8Array message duplex.

package/srpc/websocket.go

@@ -43,19 +43,14 @@ func (w *WebSocketConn) AcceptStream() (io.ReadWriteCloser, error) {
 }
 
 // OpenStream tries to open a stream with the remote.
-func (w *WebSocketConn) OpenStream(ctx context.Context, msgHandler PacketHandler) (Writer, error) {
+func (w *WebSocketConn) OpenStream(ctx context.Context, msgHandler PacketHandler, closeHandler CloseHandler) (Writer, error) {
 	muxedStream, err := w.mconn.OpenStream(ctx)
 	if err != nil {
 		return nil, err
 	}
 
 	rw := NewPacketReadWriter(muxedStream)
-	go func() {
-		err := rw.ReadPump(msgHandler)
-		if err != nil {
-			_ = rw.Close()
-		}
-	}()
+	go rw.ReadPump(msgHandler, closeHandler)
 	return rw, nil
 }