diff --git a/peers/peers.go b/peers/peers.go index b1d13f2..2f01c18 100644 --- a/peers/peers.go +++ b/peers/peers.go @@ -56,7 +56,10 @@ func (a *Peer) Serve(l net.Listener) error { return fmt.Errorf("accepting conn: %w", err) } - p := newProtocolClient(a.BaseContext, nc, a.HandlerSource()) + // Wrap the context to provide access to the underlying connection. + // TODO(tim): Do we really want this? + ctx := context.WithValue(a.BaseContext, connectionKey, nc) + p := newProtocolClient(ctx, nc, a.HandlerSource()) go func() { defer nc.Close() defer p.Close() @@ -67,3 +70,15 @@ func (a *Peer) Serve(l net.Listener) error { }() } } + +type contextKey string + +const ( + connectionKey = contextKey("connection") +) + +// Connection returns the underlying connection used in calls +// to function in a Handler. +func Connection(ctx context.Context) net.Conn { + return ctx.Value(connectionKey).(net.Conn) +}