Add a protocol (fork)

This page shows how to add a protocol to a fork of go-ricochet 1.1.0. You’ll build a small agent that tells a peer how many mailboxes it has and how full they are, register it on the server, test it without a database, and call it from pkg/client. Read Built on go-p2p-forge first; this page reuses its vocabulary.

The example was compiled and its tests run against a clean v1.1.0 tree.

Start the fork

Branch from the release tag, not from main, so your changes sit on a known base:

git clone https://github.com/stephanfeb/go-ricochet.git
cd go-ricochet
git switch -c summary v1.1.0

Keep the module path github.com/stephanfeb/go-ricochet in go.mod. Every import in the tree uses it, and your new package needs to import internal/ packages, which only works inside the same module.

Write the package

Create internal/protocol/summary/handler.go. It has the same parts as every agent: a protocol ID, request and response types, a pipeline, a response writer and handlers.

// Package summary answers a peer's questions about its own mailboxes.
package summary

import (
  "encoding/json"
  "log/slog"

  "github.com/libp2p/go-libp2p/core/protocol"

  forge "github.com/stephanfeb/go-p2p-forge"
  "github.com/stephanfeb/go-p2p-forge/codec"
  "github.com/stephanfeb/go-p2p-forge/middleware"

  "github.com/stephanfeb/go-ricochet/internal/admission"
  "github.com/stephanfeb/go-ricochet/internal/metrics"
  "github.com/stephanfeb/go-ricochet/internal/protocol/wire"
  "github.com/stephanfeb/go-ricochet/internal/ratelimit"
  "github.com/stephanfeb/go-ricochet/internal/storage"
)

// ProtocolID uses its own namespace so it can never collide with an upstream protocol.
const ProtocolID = protocol.ID("/example/summary/1.0.0")

const (
  OpCOUNT = "COUNT"
  OpLIST  = "LIST"
)

type Request struct {
  Operation string `json:"operation"`
}

type MailboxSummary struct {
  FolderPath   string `json:"folderPath"`
  MessageCount int    `json:"messageCount"`
  MaxMessages  int    `json:"maxMessages"`
}

type Response struct {
  Status    int              `json:"status"`
  Error     string           `json:"error,omitempty"`
  Count     int              `json:"count,omitempty"`
  Mailboxes []MailboxSummary `json:"mailboxes,omitempty"`
}

// StatusCode lets the metrics and access-log middleware read the outcome.
func (r *Response) StatusCode() int { return r.Status }

func NewPipeline(logger *slog.Logger, pool *codec.BufferPool, reg *forge.Registry) *forge.Pipeline {
  limiter := ratelimit.FromRegistry(reg).Summary

  routes := map[string]forge.Middleware{
    OpCOUNT: middleware.Chain(forge.JSONDeserialize[Request](), handleCount),
    OpLIST:  middleware.Chain(forge.JSONDeserialize[Request](), handleList),
  }

  return wire.Bounded(forge.NewPipeline(logger,
    metrics.Middleware(metrics.FromRegistry(reg), "summary", ""),
    wire.AccessLog("summary", ""),
    middleware.Recovery(),
    responseWriter(),
    forge.FrameDecodeMiddleware(pool),
    wire.RequestDeadline(reg),
    middleware.RateLimitMiddleware(limiter),
    admission.DrainGate(admission.DrainFromRegistry(reg)),
    admission.Middleware(admission.FromRegistry(reg)),
    middleware.OperationRouter("operation", routes),
  ).WithRegistry(reg), reg)
}

// responseWriter writes sc.Response as one JSON frame, turning a pipeline
// error into a Response so the client always gets an answer.
func responseWriter() forge.Middleware {
  return func(sc *forge.StreamContext, next func()) {
    next()

    if sc.Err != nil && sc.Response == nil {
      status, _ := wire.Classify(sc.Err)
      wire.LogRejection(sc, status, sc.Err)
      sc.Response = &Response{Status: status, Error: wire.ClientMessage(sc.Err)}
    }
    if sc.Response == nil {
      return
    }
    data, err := json.Marshal(sc.Response)
    if err != nil {
      sc.Logger.Error("failed to marshal response", "error", err)
      return
    }
    if err := sc.WriteFrame(data); err != nil {
      sc.Logger.Error("failed to write response", "error", err)
    }
  }
}

func handleCount(sc *forge.StreamContext, next func()) {
  store, _ := forge.ServiceFrom[storage.Storage](sc, "storage")

  owned, _, err := store.CountMailboxes(sc.Ctx, sc.PeerID)
  if err != nil {
    sc.Err = err
    return
  }
  sc.Response = &Response{Status: wire.StatusOK, Count: owned}
}

func handleList(sc *forge.StreamContext, next func()) {
  store, _ := forge.ServiceFrom[storage.Storage](sc, "storage")

  records, err := store.ListMailboxes(sc.Ctx, sc.PeerID)
  if err != nil {
    sc.Err = err
    return
  }
  resp := &Response{Status: wire.StatusOK, Count: len(records)}
  for _, r := range records {
    resp.Mailboxes = append(resp.Mailboxes, MailboxSummary{
      FolderPath:   r.FolderPath,
      MessageCount: r.MessageCount,
      MaxMessages:  r.MaxMessages,
    })
  }
  sc.Response = resp
}

What each choice buys you:

Authorization. The handlers only ever query sc.PeerID, the identity the connection authenticated, so a peer can’t ask about someone else’s mailboxes. If your request names a peer, check it with wire.RequireSelf(sc, req.PeerID, "peerId") or wire.RequireOwner(sc, owner, "...") before you use it. For a document, feed or collection, use wire.RequireReader.
Storage. Handlers get storage.Storage from the registry with forge.ServiceFrom, the way SDA does, and pass sc.Ctx so message_timeout_sec bounds the query. The mail agents use "mda" and "mta" instead.
Errors. A storage error set on sc.Err becomes a 500 with the message internal error; the real error goes to the log through wire.LogRejection. A rate-limit, overload or drain error becomes a 429 or 503.
Metrics and logs. The "summary" label appears in ricochet_requests_total and in the access log. Implementing StatusCode() is what lets both tell a 4xx from a success.
No capacity gate. This agent only reads. An agent that writes should add capacity.WriteGate(capacity.FromRegistry(reg), isWrite) after admission, as SDA and MSA do.

Add a rate limiter

ratelimit.FromRegistry(reg).Summary doesn’t exist yet. Add a protocol key and a limiter:

// in the RateLimit constants
RateLimitSummary  = "summary"   // example: mailbox summary

// in DefaultRateLimits, Protocols map
RateLimitSummary:  unlimited,

// in type Limiters
Summary  *middleware.TokenBucket

// in New
Summary:  single(core.RateLimitSummary),

// in Close
l.Summary.Close()

Use *middleware.DualTokenBucket, dual(...) and middleware.DualRateLimitMiddleware instead if your protocol has cheap reads and expensive writes. Like every built-in limiter, the new one is off by default. Operators turn it on with rate_limiting.protocols.summary, which the config loader accepts because the key is now in DefaultRateLimits.

Admission, drain, deadlines and metrics need no changes: they read their singletons from the registry.

Register it

Add the pipeline to registerProtocolHandlers() in internal/server/server.go, and import the package:

summaryPipeline := summary.NewPipeline(s.logger, pool, reg).WithActiveStreams(&s.inFlight)
h.SetStreamHandler(summary.ProtocolID, summaryPipeline.StreamHandler())
s.logger.Info("registered summary handler", "protocol", summary.ProtocolID)

WithActiveStreams(&s.inFlight) makes shutdown wait for your requests before storage closes.

New tables

The example uses existing storage methods. If your protocol needs its own data, the change is wider: add methods to the storage.Storage interface in internal/storage/storage.go, implement them in internal/storage/postgres, add them to the in-memory fake in internal/storage/storagetest, and add your tables to schema.sql. Prefer a separate table to new columns on an existing one; it rebases more easily.

Test it

internal/protocol/protocoltest runs a pipeline over forge’s mock stream with in-memory storage and the registry wired the way the server wires it. No PostgreSQL, no network:

package summary_test

import (
  "testing"

  "github.com/stephanfeb/go-ricochet/internal/protocol/protocoltest"
  "github.com/stephanfeb/go-ricochet/internal/protocol/summary"
)

func TestListShowsOnlyTheCallersMailboxes(t *testing.T) {
  env := protocoltest.New(t)
  p := summary.NewPipeline(env.Logger, env.Pool, env.Registry)

  owner := protocoltest.PeerID(t)
  sender := protocoltest.PeerID(t)
  env.Deliver(t, sender, owner, "inbox", []byte("one"))
  env.Deliver(t, sender, owner, "inbox", []byte("two"))

  reply := protocoltest.Exchange(t, p, owner, summary.Request{Operation: summary.OpLIST})
  resp := protocoltest.Decode[summary.Response](t, reply)
  if resp.Status != 200 || len(resp.Mailboxes) != 1 || resp.Mailboxes[0].MessageCount != 2 {
    t.Fatalf("owner LIST = %+v, want one mailbox holding 2 messages", resp)
  }

  reply = protocoltest.Exchange(t, p, sender, summary.Request{Operation: summary.OpLIST})
  if resp := protocoltest.Decode[summary.Response](t, reply); len(resp.Mailboxes) != 0 {
    t.Fatalf("other peer LIST = %+v, want no mailboxes", resp)
  }
}

The harness leaves admission control and the capacity gate out of the registry, so they admit everything. To test a single middleware, use forgetest.NewMockStreamWithFrame directly.

Two existing tests in internal/core check that the code and docs agree, and both fail until you update them:

TestReadmeMatchesTheCode wants every protocol.ID in internal/protocol/*/handler.go mentioned in README.md (in backticks) and every internal/protocol/* directory listed under “Project Structure”.
TestRateLimits_DefaultsAreOffButExplicit lists the expected rate-limit keys. Add core.RateLimitSummary.

go test ./internal/core/ ./internal/protocol/... ./pkg/...
go build ./cmd/ricochet

Call it from a client

pkg/client’s stream helpers (dial, openStream) are unexported, so a method for your protocol belongs in the fork’s pkg/client, next to the others:

package client

import (
  "context"
  "encoding/json"
  "fmt"

  "github.com/stephanfeb/go-ricochet/internal/protocol/frame"
  "github.com/stephanfeb/go-ricochet/internal/protocol/summary"
)

// ListMailboxSummaries returns the caller's mailboxes and their message
// counts from the first preferred server that answers.
func (c *Client) ListMailboxSummaries(ctx context.Context) ([]summary.MailboxSummary, error) {
  reqData, err := json.Marshal(summary.Request{Operation: summary.OpLIST})
  if err != nil {
    return nil, err
  }

  s, _, err := c.dial(ctx, summary.ProtocolID, nil)
  if err != nil {
    return nil, err
  }
  defer s.Close()

  if err := frame.WriteFrame(s, reqData); err != nil {
    return nil, fmt.Errorf("write summary frame: %w", err)
  }
  respData, err := frame.ReadFrame(s)
  if err != nil {
    return nil, fmt.Errorf("read summary response: %w", err)
  }

  var resp summary.Response
  if err := json.Unmarshal(respData, &resp); err != nil {
    return nil, fmt.Errorf("unmarshal summary response: %w", err)
  }
  if err := statusError("list mailbox summaries", resp.Status, resp.Error, 0); err != nil {
    return nil, err
  }
  return resp.Mailboxes, nil
}

statusError maps the status onto the same typed errors the other methods return, so errors.As with *client.RateLimitedError or *client.OverloadedError works. Tag a release on your fork, and apps build against it with a replace directive:

// Sketch: the fork path and tag are yours.
replace github.com/stephanfeb/go-ricochet => github.com/you/go-ricochet v1.1.0-summary.1

A client that isn’t written in Go, or doesn’t use pkg/client, needs only a libp2p stream and the frame format: a 4-byte big-endian length, then the JSON. Run a companion node has a client that does this with forge’s codec package.

Keeping your fork current

Rebase on tags. When go-ricochet releases, rebase your branch onto the new tag and rerun the tests. The compiler catches a changed middleware or helper signature, and your protocoltest suite catches most changes in behaviour.
Keep your code in new packages. Everything in internal/protocol/summary and pkg/client/summary.go is yours and never conflicts. The only edits to upstream files are one constant and one map entry in internal/core/config.go, three lines in internal/ratelimit/limiters.go, one import and three lines in registerProtocolHandlers(), and the README and test updates above.
Use your own protocol namespace. An ID under your own prefix can’t clash with a protocol upstream adds later.
Depend on the exported helpers only. Build on wire, admission, capacity, metrics and ratelimit, as above, and write your own response writer rather than editing an agent’s. Edits inside an existing agent are the ones that conflict on rebase.