diff --git a/comm.go b/comm.go index a0f51b5d..d0757ac3 100644 --- a/comm.go +++ b/comm.go @@ -1,6 +1,8 @@ package pubsub import ( + "time" + "github.com/libp2p/go-libp2p/core/network" "google.golang.org/protobuf/proto" @@ -113,3 +115,90 @@ func rpcWithControl(msgs []*pb.Message, }, } } + +func (p *PubSub) handleTopicStream(s network.Stream) { + response := make(chan *peercomm.Actor, 1) + select { + case p.eval <- func() { + pid := s.Conn().RemotePeer() + a, ok := p.peerComm.Lookup(pid) + if p.blacklist.Contains(pid) || !ok || !a.TopicStreamsEnabled() { + response <- nil + return + } + response <- a + }: + case <-p.ctx.Done(): + _ = s.Reset() + return + } + select { + case a := <-response: + if a == nil { + _ = s.Conn().CloseWithError(peercomm.TopicStreamsViolation) + return + } + a.HandleTopicInbound(s) + case <-p.ctx.Done(): + _ = s.Reset() + } +} + +func (p *PubSub) topicStreamAllowed(a *peercomm.Actor, topic string) bool { + response := make(chan bool, 1) + select { + case p.eval <- func() { + if !p.peerComm.IsCurrent(a) { + response <- false + return + } + if len(p.mySubs[topic]) != 0 || p.myRelays[topic] != 0 { + response <- true + return + } + if p.isRecentlyUnsubscribed(topic, time.Now()) { + response <- true + return + } + if gs, ok := p.rt.(*GossipSubRouter); ok { + _, fanout := gs.fanout[topic] + response <- fanout + return + } + response <- false + }: + case <-p.ctx.Done(): + return false + } + select { + case allowed := <-response: + return allowed + case <-p.ctx.Done(): + return false + } +} + +func (p *PubSub) topicStreamMisbehavior(a *peercomm.Actor) { + select { + case p.eval <- func() { + if p.peerComm.IsCurrent(a) { + if gs, ok := p.rt.(*GossipSubRouter); ok { + gs.extensions.reportMisbehavior(a.Peer()) + } + } + }: + case <-p.ctx.Done(): + } +} + +func (p *PubSub) isRecentlyUnsubscribed(topic string, now time.Time) bool { + unsubscribedAt, ok := p.recentUnsubscribed[topic] + if !ok { + return false + } + if now.Sub(unsubscribedAt) <= GossipSubUnsubscribeBackoff { + return true + } + delete(p.recentUnsubscribed, topic) + return false +} diff --git a/extensions.go b/extensions.go index dc97c0bb..f1834cd8 100644 --- a/extensions.go +++ b/extensions.go @@ -4,8 +4,10 @@ import ( "errors" "iter" + "github.com/libp2p/go-libp2p-pubsub/internal/peercomm" "github.com/libp2p/go-libp2p-pubsub/partialmessages" pubsub_pb "github.com/libp2p/go-libp2p-pubsub/pb" + "github.com/libp2p/go-libp2p/core/network" "github.com/libp2p/go-libp2p/core/peer" "github.com/libp2p/go-libp2p/core/protocol" ) @@ -13,12 +15,23 @@ import ( type PeerExtensions struct { TestExtension bool PartialMessages bool + TopicStreams bool } type TestExtensionConfig struct { OnReceiveTestExtension func(from peer.ID) } +// WithTopicStreams advertises support for topic-scoped streams. +func WithTopicStreams() Option { + return func(ps *PubSub) error { + if rt, ok := ps.rt.(*GossipSubRouter); ok { + rt.extensions.myExtensions.TopicStreams = true + } + return nil + } +} + func WithTestExtension(c TestExtensionConfig) Option { return func(ps *PubSub) error { if rt, ok := ps.rt.(*GossipSubRouter); ok { @@ -44,29 +57,30 @@ func peerExtensionsFromRPC(rpc *RPC) PeerExtensions { if hasPeerExtensions(rpc) { out.TestExtension = rpc.Control.Extensions.GetTestExtension() out.PartialMessages = rpc.Control.Extensions.GetPartialMessages() + out.TopicStreams = rpc.Control.Extensions.GetTopicStreams() } return out } func (pe *PeerExtensions) ExtendRPC(rpc *RPC) *RPC { + if !pe.TestExtension && !pe.PartialMessages && !pe.TopicStreams { + return rpc + } + if rpc.Control == nil { + rpc.Control = &pubsub_pb.ControlMessage{} + } + if rpc.Control.Extensions == nil { + rpc.Control.Extensions = &pubsub_pb.ControlExtensions{} + } if pe.TestExtension { - if rpc.Control == nil { - rpc.Control = &pubsub_pb.ControlMessage{} - } - if rpc.Control.Extensions == nil { - rpc.Control.Extensions = &pubsub_pb.ControlExtensions{} - } rpc.Control.Extensions.TestExtension = &pe.TestExtension } if pe.PartialMessages { - if rpc.Control == nil { - rpc.Control = &pubsub_pb.ControlMessage{} - } - if rpc.Control.Extensions == nil { - rpc.Control.Extensions = &pubsub_pb.ControlExtensions{} - } rpc.Control.Extensions.PartialMessages = &pe.PartialMessages } + if pe.TopicStreams { + rpc.Control.Extensions.TopicStreams = &pe.TopicStreams + } return rpc } @@ -82,24 +96,31 @@ type partialMessageInterface interface { EmitGossip(topic string, peers []peer.ID) } +type peerExtensionState struct { + received bool + receivedCaps PeerExtensions + sent bool + active bool + activeSnapshot PeerExtensions +} + type extensionsState struct { myExtensions PeerExtensions - peerExtensions map[peer.ID]PeerExtensions // peer's extensions - sentExtensions map[peer.ID]struct{} - activeExtensions map[peer.ID]PeerExtensions + peers map[peer.ID]*peerExtensionState reportMisbehavior func(peer.ID) sendRPC func(p peer.ID, r *RPC, urgent bool) testExtension *testExtension partialMessagesExtension partialMessageInterface + enableTopicStreams func(peer.ID) + disableTopicStreams func(peer.ID) + protocolViolation func(peer.ID, network.Stream) } func newExtensionsState(myExtensions PeerExtensions, reportMisbehavior func(peer.ID), sendRPC func(peer.ID, *RPC, bool)) *extensionsState { return &extensionsState{ myExtensions: myExtensions, - peerExtensions: make(map[peer.ID]PeerExtensions), - sentExtensions: make(map[peer.ID]struct{}), - activeExtensions: make(map[peer.ID]PeerExtensions), + peers: make(map[peer.ID]*peerExtensionState), reportMisbehavior: reportMisbehavior, sendRPC: sendRPC, testExtension: nil, @@ -107,98 +128,134 @@ func newExtensionsState(myExtensions PeerExtensions, reportMisbehavior func(peer } func (es *extensionsState) HandleRPC(rpc *RPC) error { - if _, ok := es.peerExtensions[rpc.from]; !ok { - // We know this is the first message because we didn't have extensions - // for this peer, and we always set extensions on the first rpc. - es.peerExtensions[rpc.from] = peerExtensionsFromRPC(rpc) - es.activatePeerExtensions(rpc.from) - } else { - // We already have an extension for this peer. If they send us another - // extensions control message, that is a protocol error. We should - // down score them because they are misbehaving. + state := es.peers[rpc.from] + var active PeerExtensions + if state != nil && state.active { + active = state.activeSnapshot + } + if active.TestExtension && es.testExtension != nil { + es.testExtension.HandleRPC(rpc.from, rpc.TestExtension) + } + if active.PartialMessages && rpc.Partial != nil && es.partialMessagesExtension != nil { + return es.partialMessagesExtension.HandleRPC(rpc.from, rpc.Partial) + } + return nil +} + +func (es *extensionsState) peerState(id peer.ID) *peerExtensionState { + state := es.peers[id] + if state == nil { + state = new(peerExtensionState) + es.peers[id] = state + } + return state +} + +func (es *extensionsState) observeExtensions(rpc *RPC) { + state := es.peerState(rpc.from) + if state.received { if hasPeerExtensions(rpc) { es.reportMisbehavior(rpc.from) } + return } - - return es.extensionsHandleRPC(rpc) + state.received = true + state.receivedCaps = peerExtensionsFromRPC(rpc) + es.reconcilePeerExtensions(rpc.from, state) } -func (es *extensionsState) OnNewIncomingStream(peer.ID, protocol.ID) { +// Preprocess observes the first control hello and validates that +// payloads use the transport selected during extension negotiation. +func (es *extensionsState) Preprocess(rpc *RPC) bool { + state := es.peers[rpc.from] + if rpc.transport == peercomm.TransportTopic { + return state != nil && state.active && state.activeSnapshot.TopicStreams + } + es.observeExtensions(rpc) + state = es.peers[rpc.from] + active := state.activeSnapshot + if state.active && active.TopicStreams && (len(rpc.Publish) != 0 || (active.PartialMessages && rpc.Partial != nil)) { + if es.protocolViolation != nil { + es.protocolViolation(rpc.from, rpc.stream) + } + return false + } + return true } +func (es *extensionsState) OnNewIncomingStream(peer.ID, protocol.ID) {} + func (es *extensionsState) OnClosedIncomingStream(id peer.ID, _ protocol.ID) { - es.deactivatePeerExtensions(id) - delete(es.peerExtensions, id) - if len(es.peerExtensions) == 0 { - es.peerExtensions = make(map[peer.ID]PeerExtensions) + state := es.peers[id] + if state == nil || !state.received { + return } + state.received = false + state.receivedCaps = PeerExtensions{} + es.reconcilePeerExtensions(id, state) + es.prunePeerState(id, state) } func (es *extensionsState) OnNewOutboundStream(id peer.ID, helloPacket *RPC) *RPC { - // Send our extensions as the first message. helloPacket = es.myExtensions.ExtendRPC(helloPacket) - - es.sentExtensions[id] = struct{}{} - es.activatePeerExtensions(id) + state := es.peerState(id) + state.sent = true + es.reconcilePeerExtensions(id, state) return helloPacket } func (es *extensionsState) OnClosedOutboundStream(id peer.ID) { - es.deactivatePeerExtensions(id) - delete(es.sentExtensions, id) - if len(es.sentExtensions) == 0 { - es.sentExtensions = make(map[peer.ID]struct{}) - } -} - -func (es *extensionsState) activatePeerExtensions(id peer.ID) { - peerExtensions, received := es.peerExtensions[id] - _, sent := es.sentExtensions[id] - if !received || !sent { + state := es.peers[id] + if state == nil || !state.sent { return } - - active := PeerExtensions{ - TestExtension: es.myExtensions.TestExtension && peerExtensions.TestExtension, - PartialMessages: es.myExtensions.PartialMessages && peerExtensions.PartialMessages, - } - es.activeExtensions[id] = active - - if active.TestExtension && es.testExtension != nil { - es.testExtension.OnNewOutboundStream(id) - } + state.sent = false + es.reconcilePeerExtensions(id, state) + es.prunePeerState(id, state) } -func (es *extensionsState) deactivatePeerExtensions(id peer.ID) { - active, ok := es.activeExtensions[id] - if !ok { - return - } - delete(es.activeExtensions, id) - if len(es.activeExtensions) == 0 { - es.activeExtensions = make(map[peer.ID]PeerExtensions) +func (es *extensionsState) reconcilePeerExtensions(id peer.ID, state *peerExtensionState) { + shouldActivate := state.received && state.sent + if state.active && !shouldActivate { + active := state.activeSnapshot + state.active = false + state.activeSnapshot = PeerExtensions{} + if active.TopicStreams && es.disableTopicStreams != nil { + es.disableTopicStreams(id) + } + if active.PartialMessages && es.partialMessagesExtension != nil { + es.partialMessagesExtension.OnClosedOutboundStream(id) + } } - - if active.PartialMessages && es.partialMessagesExtension != nil { - es.partialMessagesExtension.OnClosedOutboundStream(id) + if !state.active && shouldActivate { + active := PeerExtensions{ + TestExtension: es.myExtensions.TestExtension && state.receivedCaps.TestExtension, + PartialMessages: es.myExtensions.PartialMessages && state.receivedCaps.PartialMessages, + TopicStreams: es.myExtensions.TopicStreams && state.receivedCaps.TopicStreams, + } + state.active = true + state.activeSnapshot = active + if active.TestExtension && es.testExtension != nil { + es.testExtension.OnNewOutboundStream(id) + } + if active.TopicStreams && es.enableTopicStreams != nil { + es.enableTopicStreams(id) + } } } -func (es *extensionsState) extensionsHandleRPC(rpc *RPC) error { - active := es.activeExtensions[rpc.from] - if active.TestExtension && es.testExtension != nil { - es.testExtension.HandleRPC(rpc.from, rpc.TestExtension) +func (es *extensionsState) prunePeerState(id peer.ID, state *peerExtensionState) { + if !state.received && !state.sent && !state.active { + delete(es.peers, id) } +} - if active.PartialMessages && rpc.Partial != nil && es.partialMessagesExtension != nil { - err := es.partialMessagesExtension.HandleRPC(rpc.from, rpc.Partial) - if err != nil { - return err - } +func (es *extensionsState) activePeerExtensions(id peer.ID) PeerExtensions { + state := es.peers[id] + if state == nil || !state.active { + return PeerExtensions{} } - - return nil + return state.activeSnapshot } func (es *extensionsState) Heartbeat() { @@ -286,7 +343,7 @@ func (r partialMessageRouter) MeshPeers(topic string) iter.Seq[peer.ID] { } for peer := range peerSet { - if r.gs.extensions.activeExtensions[peer].PartialMessages && + if r.gs.extensions.activePeerExtensions(peer).PartialMessages && ((r.gs.iRequestPartial(topic) && r.gs.peerSupportsSendingPartial(peer, topic)) || (r.gs.iSupportSendingPartial(topic) && r.gs.peerRequestsPartial(peer, topic))) { if !yield(peer) { diff --git a/extensions_test.go b/extensions_test.go index 46de2107..6e14ecbe 100644 --- a/extensions_test.go +++ b/extensions_test.go @@ -3,8 +3,11 @@ package pubsub import ( "testing" + "github.com/libp2p/go-libp2p-pubsub/internal/peercomm" pubsub_pb "github.com/libp2p/go-libp2p-pubsub/pb" + "github.com/libp2p/go-libp2p/core/network" "github.com/libp2p/go-libp2p/core/peer" + "google.golang.org/protobuf/proto" ) type lifecyclePartialMessages struct { @@ -42,10 +45,10 @@ func partialExtensionsHello(id peer.ID) *RPC { func activatePartialExtensions(t *testing.T, es *extensionsState, id peer.ID) { t.Helper() es.OnNewOutboundStream(id, &RPC{}) - if err := es.HandleRPC(partialExtensionsHello(id)); err != nil { - t.Fatalf("handle extensions hello: %v", err) + if !es.Preprocess(partialExtensionsHello(id)) { + t.Fatal("preprocess extensions hello") } - if !es.activeExtensions[id].PartialMessages { + if !es.activePeerExtensions(id).PartialMessages { t.Fatal("partial messages extension was not activated") } } @@ -65,7 +68,7 @@ func TestExtensionsDeactivateOnEitherHalfClosing(t *testing.T) { activatePartialExtensions(t, es, id) test.close(es, id) - if es.activeExtensions[id].PartialMessages { + if es.activePeerExtensions(id).PartialMessages { t.Fatal("extension remained active after stream closure") } if len(cleanup.closed) != 1 || cleanup.closed[0] != id { @@ -91,8 +94,8 @@ func TestExtensionsReplacementHalfReactivates(t *testing.T) { closeHalf: func(es *extensionsState, id peer.ID) { es.OnClosedIncomingStream(id, "") }, replace: func(t *testing.T, es *extensionsState, id peer.ID) { t.Helper() - if err := es.HandleRPC(partialExtensionsHello(id)); err != nil { - t.Fatalf("handle replacement hello: %v", err) + if !es.Preprocess(partialExtensionsHello(id)) { + t.Fatal("preprocess replacement hello") } }, }, @@ -113,7 +116,7 @@ func TestExtensionsReplacementHalfReactivates(t *testing.T) { test.closeHalf(es, id) test.replace(t, es, id) - if !es.activeExtensions[id].PartialMessages { + if !es.activePeerExtensions(id).PartialMessages { t.Fatal("replacement half did not reactivate extension") } if len(cleanup.closed) != 1 { @@ -129,7 +132,7 @@ func TestExtensionsPartialCleanupUsesActiveSnapshot(t *testing.T) { id := peer.ID("peer") activatePartialExtensions(t, es, id) - es.peerExtensions[id] = PeerExtensions{} + es.peers[id].receivedCaps = PeerExtensions{} es.myExtensions.PartialMessages = false es.OnClosedOutboundStream(id) @@ -137,3 +140,285 @@ func TestExtensionsPartialCleanupUsesActiveSnapshot(t *testing.T) { t.Fatalf("expected cleanup from negotiated snapshot, got %v", cleanup.closed) } } + +func TestPeerExtensionsTopicStreams(t *testing.T) { + t.Run("advertise and parse", func(t *testing.T) { + extensions := PeerExtensions{TopicStreams: true} + rpc := extensions.ExtendRPC(&RPC{}) + + if rpc.Control == nil || rpc.Control.Extensions == nil { + t.Fatal("expected extensions envelope") + } + if !rpc.Control.Extensions.GetTopicStreams() { + t.Fatal("expected topic streams capability") + } + if got := peerExtensionsFromRPC(rpc); !got.TopicStreams { + t.Fatal("expected parsed topic streams capability") + } + }) + + t.Run("empty capabilities leave RPC unchanged", func(t *testing.T) { + extensions := PeerExtensions{} + rpc := &RPC{} + + if got := extensions.ExtendRPC(rpc); got != rpc { + t.Fatal("expected original RPC") + } + if rpc.Control != nil { + t.Fatal("expected no empty control envelope") + } + }) + + t.Run("empty capabilities preserve existing control", func(t *testing.T) { + extensions := PeerExtensions{} + control := &pubsub_pb.ControlMessage{ + Ihave: []*pubsub_pb.ControlIHave{{TopicID: proto.String("topic")}}, + } + rpc := &RPC{RPC: pubsub_pb.RPC{Control: control}} + + if got := extensions.ExtendRPC(rpc); got != rpc { + t.Fatal("expected original RPC") + } + if rpc.Control != control { + t.Fatal("expected existing control to be preserved") + } + if rpc.Control.Extensions != nil { + t.Fatal("expected no empty extensions envelope") + } + if len(rpc.Control.Ihave) != 1 || rpc.Control.Ihave[0].GetTopicID() != "topic" { + t.Fatal("expected existing control content to be preserved") + } + }) +} + +func TestWithTopicStreamsAdvertisesOnly(t *testing.T) { + router := &GossipSubRouter{ + extensions: newExtensionsState(PeerExtensions{}, func(peer.ID) {}, func(peer.ID, *RPC, bool) {}), + } + ps := &PubSub{rt: router} + + if err := WithTopicStreams()(ps); err != nil { + t.Fatal(err) + } + if !router.extensions.myExtensions.TopicStreams { + t.Fatal("expected topic streams to be advertised") + } + + rpc := router.extensions.OnNewOutboundStream("peer", &RPC{ + RPC: pubsub_pb.RPC{}, + }) + if !rpc.Control.Extensions.GetTopicStreams() { + t.Fatal("expected first RPC to advertise topic streams") + } +} + +func TestTopicStreamsNegotiationActivatesAndRejectsControlPayload(t *testing.T) { + peerID := peer.ID("peer") + enabled := make(chan struct{}, 1) + violations := 0 + es := newExtensionsState(PeerExtensions{TopicStreams: true, PartialMessages: true}, func(peer.ID) {}, func(peer.ID, *RPC, bool) {}) + es.enableTopicStreams = func(got peer.ID) { + if got != peerID { + t.Fatalf("enabled wrong peer: %s", got) + } + enabled <- struct{}{} + } + es.protocolViolation = func(got peer.ID, _ network.Stream) { + if got != peerID { + t.Fatalf("violated wrong peer: %s", got) + } + violations++ + } + es.OnNewOutboundStream(peerID, &RPC{}) + rpc := &RPC{RPC: pubsub_pb.RPC{ + Control: &pubsub_pb.ControlMessage{Extensions: &pubsub_pb.ControlExtensions{ + TopicStreams: proto.Bool(true), PartialMessages: proto.Bool(true), + }}, + Publish: []*pubsub_pb.Message{{Data: []byte("forbidden")}}, + }, from: peerID, transport: peercomm.TransportControl} + if es.Preprocess(rpc) { + t.Fatal("accepted payload on negotiated control stream") + } + <-enabled + if violations != 1 { + t.Fatalf("violations = %d, want 1", violations) + } +} + +func TestTopicStreamsUnsupportedPeerKeepsControlPayload(t *testing.T) { + peerID := peer.ID("peer") + es := newExtensionsState(PeerExtensions{TopicStreams: true}, func(peer.ID) {}, func(peer.ID, *RPC, bool) {}) + es.OnNewOutboundStream(peerID, &RPC{}) + rpc := &RPC{RPC: pubsub_pb.RPC{ + Control: &pubsub_pb.ControlMessage{Extensions: &pubsub_pb.ControlExtensions{}}, + Publish: []*pubsub_pb.Message{{Data: []byte("fallback")}}, + }, from: peerID, transport: peercomm.TransportControl} + if !es.Preprocess(rpc) { + t.Fatal("rejected fallback control payload") + } +} + +func TestTopicTransportCannotSeedExtensionNegotiation(t *testing.T) { + peerID := peer.ID("peer") + enabled := 0 + es := newExtensionsState(PeerExtensions{TopicStreams: true}, func(peer.ID) {}, func(peer.ID, *RPC, bool) {}) + es.enableTopicStreams = func(peer.ID) { enabled++ } + rpc := &RPC{RPC: pubsub_pb.RPC{Control: &pubsub_pb.ControlMessage{Extensions: &pubsub_pb.ControlExtensions{TopicStreams: proto.Bool(true)}}}, from: peerID, transport: peercomm.TransportTopic} + if es.Preprocess(rpc) { + t.Fatal("accepted topic payload before control negotiation") + } + if state := es.peers[peerID]; state != nil && state.received { + t.Fatal("topic payload seeded extension negotiation") + } + if enabled != 0 { + t.Fatalf("topic streams enabled %d times", enabled) + } +} + +func TestTopicStreamsControlAllowsUnnegotiatedPartialFallback(t *testing.T) { + peerID := peer.ID("peer") + es := newExtensionsState(PeerExtensions{TopicStreams: true}, func(peer.ID) {}, func(peer.ID, *RPC, bool) {}) + es.OnNewOutboundStream(peerID, &RPC{}) + rpc := &RPC{RPC: pubsub_pb.RPC{ + Control: &pubsub_pb.ControlMessage{Extensions: &pubsub_pb.ControlExtensions{TopicStreams: proto.Bool(true)}}, + Partial: &pubsub_pb.PartialMessagesExtension{TopicID: proto.String("topic"), PartialMessage: []byte("partial")}, + }, from: peerID, transport: peercomm.TransportControl} + if !es.Preprocess(rpc) { + t.Fatal("rejected unnegotiated partial fallback on control") + } +} + +func TestPeerExtensionsIncomingLifecycle(t *testing.T) { + peerID := peer.ID("peer") + enabled := 0 + disabled := 0 + partial := &recordingPartialMessageExtension{} + es := newExtensionsState(PeerExtensions{TopicStreams: true, PartialMessages: true}, func(peer.ID) {}, func(peer.ID, *RPC, bool) {}) + es.enableTopicStreams = func(got peer.ID) { + if got != peerID { + t.Fatalf("enabled wrong peer: %s", got) + } + enabled++ + } + es.disableTopicStreams = func(got peer.ID) { + if got != peerID { + t.Fatalf("disabled wrong peer: %s", got) + } + disabled++ + } + es.partialMessagesExtension = partial + + es.OnNewOutboundStream(peerID, &RPC{}) + if !es.Preprocess(extensionHello(peerID, true, true)) { + t.Fatal("rejected initial extension hello") + } + if enabled != 1 { + t.Fatalf("topic stream activations = %d, want 1", enabled) + } + + // Deactivation must use the activation snapshot, not current received state. + es.peers[peerID].receivedCaps = PeerExtensions{} + es.OnClosedIncomingStream(peerID, "") + if state := es.peers[peerID]; state != nil && state.received { + t.Fatal("incoming close did not clear received extensions") + } + if es.peers[peerID] != nil && es.peers[peerID].active { + t.Fatal("incoming close left extensions active") + } + if disabled != 1 || partial.closed != 1 { + t.Fatalf("deactivation counts = topic %d, partial %d; want 1, 1", disabled, partial.closed) + } + if es.Preprocess(&RPC{from: peerID, transport: peercomm.TransportTopic}) { + t.Fatal("accepted topic RPC while extensions were inactive") + } + + if !es.Preprocess(extensionHello(peerID, true, false)) { + t.Fatal("rejected replacement extension hello") + } + if enabled != 2 { + t.Fatalf("topic stream activations = %d, want 2", enabled) + } + if !es.Preprocess(&RPC{from: peerID, transport: peercomm.TransportTopic}) { + t.Fatal("rejected topic RPC after replacement hello reactivated extensions") + } + + es.OnClosedOutboundStream(peerID) + if disabled != 2 || partial.closed != 1 { + t.Fatalf("final deactivation counts = topic %d, partial %d; want 2, 1", disabled, partial.closed) + } + es.OnClosedOutboundStream(peerID) + if disabled != 2 || partial.closed != 1 { + t.Fatal("repeated outbound close deactivated extensions twice") + } +} + +func TestPeerExtensionsOutboundLifecycle(t *testing.T) { + peerID := peer.ID("peer") + enabled := 0 + disabled := 0 + es := newExtensionsState(PeerExtensions{TopicStreams: true}, func(peer.ID) {}, func(peer.ID, *RPC, bool) {}) + es.enableTopicStreams = func(peer.ID) { enabled++ } + es.disableTopicStreams = func(peer.ID) { disabled++ } + + if !es.Preprocess(extensionHello(peerID, true, false)) { + t.Fatal("rejected incoming-first extension hello") + } + if enabled != 0 { + t.Fatal("activated before outbound extensions were sent") + } + es.OnNewOutboundStream(peerID, &RPC{}) + if enabled != 1 { + t.Fatalf("activations = %d, want 1", enabled) + } + + es.OnClosedOutboundStream(peerID) + if disabled != 1 { + t.Fatalf("deactivations = %d, want 1", disabled) + } + if state := es.peers[peerID]; state == nil || !state.received { + t.Fatal("outbound close unexpectedly cleared received extensions") + } + if es.Preprocess(&RPC{from: peerID, transport: peercomm.TransportTopic}) { + t.Fatal("accepted topic RPC after outbound close") + } + + es.OnNewOutboundStream(peerID, &RPC{}) + if enabled != 2 { + t.Fatalf("replacement outbound stream activations = %d, want 2", enabled) + } + es.OnClosedIncomingStream(peerID, "") + if disabled != 2 { + t.Fatalf("incoming close deactivations = %d, want 2", disabled) + } + es.OnClosedIncomingStream(peerID, "") + if disabled != 2 { + t.Fatal("repeated incoming close deactivated extensions twice") + } +} + +func extensionHello(from peer.ID, topicStreams, partialMessages bool) *RPC { + return &RPC{ + RPC: pubsub_pb.RPC{Control: &pubsub_pb.ControlMessage{Extensions: &pubsub_pb.ControlExtensions{ + TopicStreams: proto.Bool(topicStreams), + PartialMessages: proto.Bool(partialMessages), + }}}, + from: from, + transport: peercomm.TransportControl, + } +} + +type recordingPartialMessageExtension struct { + closed int +} + +func (m *recordingPartialMessageExtension) OnClosedOutboundStream(peer.ID) { + m.closed++ +} + +func (*recordingPartialMessageExtension) HandleRPC(peer.ID, *pubsub_pb.PartialMessagesExtension) error { + return nil +} + +func (*recordingPartialMessageExtension) Heartbeat() {} + +func (*recordingPartialMessageExtension) EmitGossip(string, []peer.ID) {} diff --git a/gossipsub_peer_lifecycle_test.go b/gossipsub_peer_lifecycle_test.go index 74589578..5d8323cd 100644 --- a/gossipsub_peer_lifecycle_test.go +++ b/gossipsub_peer_lifecycle_test.go @@ -63,12 +63,8 @@ func waitForLifecycleCondition(t *testing.T, ps *PubSub, desc string, condition } } -func writeLifecycleSubscription(t *testing.T, stream network.Stream, topic string) { +func writeLifecycleRPC(t *testing.T, stream network.Stream, rpc *pb.RPC) { t.Helper() - rpc := &pb.RPC{Subscriptions: []*pb.RPC_SubOpts{{ - Topicid: proto.String(topic), - Subscribe: proto.Bool(true), - }}} b, err := proto.Marshal(rpc) if err != nil { t.Fatal(err) @@ -120,7 +116,15 @@ func TestInitialOutboundOpenFailureRetiresInboundPeer(t *testing.T) { t.Fatal(err) } defer inbound.Close() - writeLifecycleSubscription(t, inbound, topicID) + writeLifecycleRPC(t, inbound, &pb.RPC{ + Subscriptions: []*pb.RPC_SubOpts{{ + Topicid: proto.String(topicID), + Subscribe: proto.Bool(true), + }}, + Control: &pb.ControlMessage{Extensions: &pb.ControlExtensions{ + TopicStreams: proto.Bool(true), + }}, + }) eventCtx, eventCancel := context.WithTimeout(ctx, 5*time.Second) join, err := events.NextPeerEvent(eventCtx) @@ -143,7 +147,9 @@ func TestInitialOutboundOpenFailureRetiresInboundPeer(t *testing.T) { waitForLifecycleCondition(t, ps, "peer retirement and topic cleanup", func() bool { _, inRegistry := ps.peerComm.Lookup(remote.ID()) _, inTopic := ps.topics[topicID][remote.ID()] - return !inRegistry && !inTopic + gs := ps.rt.(*GossipSubRouter) + _, hasExtensions := gs.extensions.peers[remote.ID()] + return !inRegistry && !inTopic && !hasExtensions }) select { case err := <-reset: diff --git a/internal/peercomm/peercomm.go b/internal/peercomm/peercomm.go index 032b1c06..ad73f045 100644 --- a/internal/peercomm/peercomm.go +++ b/internal/peercomm/peercomm.go @@ -53,6 +53,8 @@ type Hooks struct { OutboundSendFailed func(*Actor, network.Stream, *pb.RPC, error) OutboundOpenFailed func(*Actor, error) OutboundDead func(*Actor, network.Stream, error) + TopicAllowed func(*Actor, string) bool + TopicMisbehavior func(*Actor, string) } // Config configures all actors in a Registry. @@ -228,6 +230,14 @@ type Actor struct { currentInbound *inboundRun notifiedInbound *inboundRun retire sync.Once + + topicMu sync.Mutex + topicEnabled bool + topicWriters map[string]*topicWriter + topicWritersWG sync.WaitGroup + topicInbound map[string]*topicInboundState + topicInboundStreams map[network.Stream]struct{} + topicInboundTotal int } type inboundRun struct { @@ -242,6 +252,8 @@ func newActor(r *Registry, p peer.ID) *Actor { a := &Actor{ registry: r, peer: p, ctx: ctx, cancel: cancel, queue: newRPCQueue(r.config.QueueSize), commands: make(chan command, 16), done: make(chan struct{}), + topicWriters: make(map[string]*topicWriter), topicInbound: make(map[string]*topicInboundState), + topicInboundStreams: make(map[network.Stream]struct{}), } go a.run() return a @@ -281,7 +293,7 @@ func (a *Actor) Send(rpc *pb.RPC, urgent bool) error { return ErrActorRetired default: } - err := a.queue.push(rpc, urgent) + err := a.splitAndSend(rpc, urgent) if errors.Is(err, ErrQueueClosed) { return ErrActorRetired } @@ -293,6 +305,7 @@ func (a *Actor) Retire() { a.retire.Do(func() { a.cancel() a.queue.close() + a.DisableTopicStreams() }) } @@ -436,7 +449,11 @@ func (a *Actor) claimInboundCloseLocked(run *inboundRun) (protocol.ID, bool) { } func (a *Actor) run() { - defer close(a.done) + defer func() { + a.stopTopicWriters() + a.topicWritersWG.Wait() + close(a.done) + }() var generation uint64 var current network.Stream var pending network.Stream @@ -570,7 +587,7 @@ func (a *Actor) writeLoop(ctx context.Context, generation uint64, s network.Stre } func (a *Actor) writeRPC(s network.Stream, rpc *pb.RPC) error { - err := writeRPC(s, rpc) + err := writeProto(s, rpc) if err != nil { if h := a.registry.config.Hooks.OutboundSendFailed; h != nil { h(a, s, rpc, err) @@ -605,12 +622,12 @@ func (a *Actor) closeInbound() { a.inboundMu.Unlock() } -func writeRPC(s network.Stream, rpc *pb.RPC) error { - size := uint64(proto.Size(rpc)) +func writeProto(s network.Stream, message proto.Message) error { + size := uint64(proto.Size(message)) buf := pool.Get(varint.UvarintSize(size) + int(size)) defer pool.Put(buf) n := binary.PutUvarint(buf, size) - out, err := proto.MarshalOptions{}.MarshalAppend(buf[:n], rpc) + out, err := proto.MarshalOptions{}.MarshalAppend(buf[:n], message) if err != nil { return err } diff --git a/internal/peercomm/peercomm_test.go b/internal/peercomm/peercomm_test.go index c0e7beef..ca00c2b5 100644 --- a/internal/peercomm/peercomm_test.go +++ b/internal/peercomm/peercomm_test.go @@ -202,10 +202,20 @@ type streamRead struct { type testConn struct { network.Conn - remote peer.ID + remote peer.ID + mu sync.Mutex + closeCode network.ConnErrorCode + closed int } func (c *testConn) RemotePeer() peer.ID { return c.remote } +func (c *testConn) CloseWithError(code network.ConnErrorCode) error { + c.mu.Lock() + c.closeCode = code + c.closed++ + c.mu.Unlock() + return nil +} type testStream struct { network.Stream @@ -325,6 +335,7 @@ func (s *testStream) Reset() error { } func (s *testStream) SetWriteDeadline(time.Time) error { return nil } +func (s *testStream) SetReadDeadline(time.Time) error { return nil } func (s *testStream) counts() (closed, reset int) { s.mu.Lock() diff --git a/internal/peercomm/topic_streams.go b/internal/peercomm/topic_streams.go new file mode 100644 index 00000000..bb01c1c5 --- /dev/null +++ b/internal/peercomm/topic_streams.go @@ -0,0 +1,384 @@ +package peercomm + +import ( + "context" + "errors" + "io" + "sync" + "time" + + "github.com/libp2p/go-libp2p/core/network" + "github.com/libp2p/go-libp2p/core/protocol" + "github.com/libp2p/go-msgio" + "google.golang.org/protobuf/proto" + + pb "github.com/libp2p/go-libp2p-pubsub/pb" +) + +const ( + TopicStreamsProtocol protocol.ID = "/gsts/v0beta" + TopicStreamsViolation network.ConnErrorCode = 0xd52505 + maxInboundTopicStreamsPerTopic = 3 + maxInboundTopicStreamsPerPeer = 24 +) + +var ( + ErrInvalidTopicRPC = errors.New("peercomm: topic payload has no topic") +) + +type topicWriter struct { + topic string + ctx context.Context + cancel context.CancelFunc + queue chan *pb.TopicRPC +} + +type topicInboundState struct { + active int + deliver sync.Mutex +} + +// EnableTopicStreams enables the negotiated transport. It is safe to call from +// the PubSub process loop and does not wait for actor work. +func (a *Actor) EnableTopicStreams() { + a.topicMu.Lock() + if a.ctx.Err() == nil { + a.topicEnabled = true + } + a.topicMu.Unlock() +} + +// DisableTopicStreams tears down all topic streams without affecting the +// control stream. +func (a *Actor) DisableTopicStreams() { + a.topicMu.Lock() + a.topicEnabled = false + a.stopTopicWritersLocked() + a.topicMu.Unlock() + a.closeInboundTopics() +} + +func (a *Actor) TopicStreamsEnabled() bool { + a.topicMu.Lock() + defer a.topicMu.Unlock() + return a.topicEnabled +} + +// HasTopicWriter reports whether this actor currently owns an outbound writer +// for topic. It is intended for transport diagnostics and tests. +func (a *Actor) HasTopicWriter(topic string) bool { + a.topicMu.Lock() + defer a.topicMu.Unlock() + _, ok := a.topicWriters[topic] + return ok +} + +// CloseTopic closes the initiator-owned stream for topic. A later send lazily +// creates a replacement. +func (a *Actor) CloseTopic(topic string) { + a.topicMu.Lock() + writer := a.topicWriters[topic] + delete(a.topicWriters, topic) + if writer != nil { + writer.cancel() + } + a.topicMu.Unlock() +} + +func (a *Actor) stopTopicWriters() { + a.topicMu.Lock() + a.stopTopicWritersLocked() + a.topicMu.Unlock() +} + +func (a *Actor) stopTopicWritersLocked() { + for topic, writer := range a.topicWriters { + delete(a.topicWriters, topic) + writer.cancel() + } +} + +func (a *Actor) splitAndSend(rpc *pb.RPC, urgent bool) error { + a.topicMu.Lock() + defer a.topicMu.Unlock() + if a.ctx.Err() != nil { + return ErrQueueClosed + } + if !a.topicEnabled { + return a.queue.push(rpc, urgent) + } + + rpc = proto.Clone(rpc).(*pb.RPC) + control := &pb.RPC{ + Subscriptions: rpc.Subscriptions, + Control: rpc.Control, + TestExtension: rpc.TestExtension, + } + byTopic := make(map[string][]*pb.TopicRPC) + for _, message := range rpc.Publish { + if message == nil || message.GetTopic() == "" { + return ErrInvalidTopicRPC + } + topic := message.GetTopic() + message.Topic = nil + byTopic[topic] = append(byTopic[topic], &pb.TopicRPC{Payload: &pb.TopicRPC_Publish{Publish: message}}) + } + if rpc.Partial != nil { + if rpc.Partial.GetTopicID() == "" { + return ErrInvalidTopicRPC + } + topic := rpc.Partial.GetTopicID() + rpc.Partial.TopicID = nil + byTopic[topic] = append(byTopic[topic], &pb.TopicRPC{Payload: &pb.TopicRPC_Partial{Partial: rpc.Partial}}) + } + + // Reserve every destination before publishing any part of the RPC. This + // keeps a full control or topic queue from producing a partial send. + a.queue.mu.Lock() + defer a.queue.mu.Unlock() + if a.queue.closed { + return ErrQueueClosed + } + controlPending := proto.Size(control) > 0 + if controlPending && len(a.queue.urgent)+len(a.queue.normal) >= a.queue.capacity { + return ErrQueueFull + } + for topic, items := range byTopic { + writer := a.topicWriters[topic] + queued := 0 + if writer != nil { + queued = len(writer.queue) + } + if queued+len(items) > a.registry.config.QueueSize { + return ErrQueueFull + } + } + + if controlPending { + if urgent { + a.queue.urgent = append(a.queue.urgent, control) + } else { + a.queue.normal = append(a.queue.normal, control) + } + a.queue.available.Signal() + } + for topic, items := range byTopic { + writer := a.topicWriters[topic] + if writer == nil { + ctx, cancel := context.WithCancel(a.ctx) + writer = &topicWriter{topic: topic, ctx: ctx, cancel: cancel, queue: make(chan *pb.TopicRPC, a.registry.config.QueueSize)} + a.topicWriters[topic] = writer + a.topicWritersWG.Add(1) + go a.runTopicWriter(writer) + } + for _, item := range items { + writer.queue <- item + } + } + return nil +} + +func (a *Actor) runTopicWriter(w *topicWriter) { + defer a.topicWritersWG.Done() + var stream network.Stream + defer func() { + if stream != nil { + _ = stream.Close() + } + }() + for { + select { + case <-w.ctx.Done(): + return + case item := <-w.queue: + var err error + stream, err = a.writeTopicItem(w, stream, item) + if err != nil && stream != nil { + _ = stream.Close() + stream = nil + } + } + } +} + +func (a *Actor) writeTopicItem(w *topicWriter, stream network.Stream, item *pb.TopicRPC) (network.Stream, error) { + if err := w.ctx.Err(); err != nil { + return stream, err + } + if stream == nil { + s, err := a.registry.config.Host.NewStream(w.ctx, a.peer, TopicStreamsProtocol) + if err != nil { + return nil, err + } + stream = s + if err := w.ctx.Err(); err != nil { + _ = s.Reset() + return nil, err + } + if err := writeProto(s, &pb.TopicRPCHeader{Topic: proto.String(w.topic)}); err != nil { + _ = s.Close() + return nil, err + } + go a.watchTopicResponder(w.ctx, s) + } + if err := writeProto(stream, item); err != nil { + return stream, err + } + return stream, nil +} + +func (a *Actor) watchTopicResponder(ctx context.Context, s network.Stream) { + var one [1]byte + n, _ := s.Read(one[:]) + if ctx.Err() == nil && n != 0 { + a.protocolViolation(s) + } +} + +// HandleTopicInbound validates and consumes a responder-side topic stream. +func (a *Actor) HandleTopicInbound(s network.Stream) { + if s == nil || s.Conn().RemotePeer() != a.peer { + if s != nil { + _ = s.Reset() + } + return + } + if !a.TopicStreamsEnabled() { + a.protocolViolation(s) + return + } + a.topicMu.Lock() + a.topicInboundStreams[s] = struct{}{} + a.topicMu.Unlock() + defer func() { a.topicMu.Lock(); delete(a.topicInboundStreams, s); a.topicMu.Unlock() }() + _ = s.SetReadDeadline(time.Now().Add(time.Second)) + r := msgio.NewVarintReaderSize(s, a.registry.config.MaxMessageSize) + b, err := r.ReadMsg() + if err != nil { + r.ReleaseMsg(b) + _ = s.Reset() + return + } + header := new(pb.TopicRPCHeader) + err = proto.Unmarshal(b, header) + r.ReleaseMsg(b) + if err != nil || header.GetTopic() == "" { + a.protocolViolation(s) + return + } + _ = s.SetReadDeadline(time.Time{}) + topic := header.GetTopic() + + a.topicMu.Lock() + state := a.topicInbound[topic] + if state == nil { + state = new(topicInboundState) + a.topicInbound[topic] = state + } + if state.active >= maxInboundTopicStreamsPerTopic || a.topicInboundTotal >= maxInboundTopicStreamsPerPeer { + a.topicMu.Unlock() + if h := a.registry.config.Hooks.TopicMisbehavior; h != nil { + h(a, topic) + } + _ = s.Reset() + return + } + state.active++ + a.topicInboundTotal++ + a.topicMu.Unlock() + defer func() { + a.topicMu.Lock() + state.active-- + a.topicInboundTotal-- + if state.active == 0 { + delete(a.topicInbound, topic) + } + a.topicMu.Unlock() + }() + state.deliver.Lock() + defer state.deliver.Unlock() + if h := a.registry.config.Hooks.TopicAllowed; h != nil && !h(a, topic) { + if mh := a.registry.config.Hooks.TopicMisbehavior; mh != nil { + mh(a, topic) + } + _ = s.Reset() + return + } + for { + b, err = r.ReadMsg() + if err != nil { + r.ReleaseMsg(b) + if errors.Is(err, io.EOF) { + _ = s.Close() + } else { + _ = s.Reset() + } + return + } + trpc := new(pb.TopicRPC) + err = proto.Unmarshal(b, trpc) + r.ReleaseMsg(b) + if err != nil || trpc.GetPayload() == nil { + a.protocolViolation(s) + return + } + var rpc pb.RPC + switch payload := trpc.GetPayload().(type) { + case *pb.TopicRPC_Publish: + m := payload.Publish + if m == nil || len(m.Data) == 0 || m.Topic != nil { + a.protocolViolation(s) + return + } + m.Topic = proto.String(topic) + rpc.Publish = []*pb.Message{m} + case *pb.TopicRPC_Partial: + if payload.Partial == nil || payload.Partial.TopicID != nil { + a.protocolViolation(s) + return + } + partial := proto.Clone(payload.Partial).(*pb.PartialMessagesExtension) + partial.TopicID = proto.String(topic) + rpc.Partial = partial + default: + a.protocolViolation(s) + return + } + if h := a.registry.config.Hooks.InboundRPC; h != nil { + h(a, s, TransportTopic, &rpc) + } + } +} + +func (a *Actor) closeInboundTopics() { + a.topicMu.Lock() + streams := make([]network.Stream, 0, len(a.topicInboundStreams)) + for s := range a.topicInboundStreams { + streams = append(streams, s) + } + a.topicMu.Unlock() + for _, s := range streams { + _ = s.Reset() + } +} + +// ProtocolViolation closes the connection carrying s. If s is nil, it closes +// the current inbound control connection. +func (a *Actor) ProtocolViolation(s network.Stream) { + if s == nil { + a.inboundMu.Lock() + if a.currentInbound != nil { + s = a.currentInbound.stream + } + a.inboundMu.Unlock() + } else if s.Conn().RemotePeer() != a.peer { + return + } + if s != nil { + a.protocolViolation(s) + } +} + +func (a *Actor) protocolViolation(s network.Stream) { + _ = s.Conn().CloseWithError(TopicStreamsViolation) +} diff --git a/internal/peercomm/topic_streams_test.go b/internal/peercomm/topic_streams_test.go new file mode 100644 index 00000000..cde42ea4 --- /dev/null +++ b/internal/peercomm/topic_streams_test.go @@ -0,0 +1,354 @@ +package peercomm + +import ( + "context" + "errors" + "io" + "sync" + "testing" + "time" + + "github.com/libp2p/go-libp2p/core/network" + "github.com/libp2p/go-libp2p/core/peer" + "github.com/libp2p/go-libp2p/core/protocol" + "google.golang.org/protobuf/proto" + + pb "github.com/libp2p/go-libp2p-pubsub/pb" +) + +type topicHost struct{ streams chan *testStream } + +func (h *topicHost) NewStream(_ context.Context, p peer.ID, protocols ...protocol.ID) (network.Stream, error) { + if len(protocols) != 1 || protocols[0] != TopicStreamsProtocol { + panic("unexpected protocol") + } + s := newTestStream(p) + h.streams <- s + return s, nil +} + +type fixedTopicHost struct{ stream *testStream } + +func (h *fixedTopicHost) NewStream(_ context.Context, _ peer.ID, _ ...protocol.ID) (network.Stream, error) { + return h.stream, nil +} + +type gatedTopicHost struct { + started chan struct{} + release <-chan struct{} + stream *testStream + once sync.Once +} + +func (h *gatedTopicHost) NewStream(_ context.Context, p peer.ID, protocols ...protocol.ID) (network.Stream, error) { + if len(protocols) != 1 || protocols[0] != TopicStreamsProtocol { + panic("unexpected protocol") + } + h.once.Do(func() { close(h.started) }) + <-h.release + if h.stream == nil { + h.stream = newTestStream(p) + } + return h.stream, nil +} + +func requirePromptReturn(t *testing.T, call func()) { + t.Helper() + done := make(chan struct{}) + go func() { + call() + close(done) + }() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("topic stream teardown blocked") + } +} + +func TestCloseTopicDoesNotWaitForBlockedWrite(t *testing.T) { + gate := make(chan struct{}) + stream := newTestStream(peer.ID("peer")) + stream.writeStarted = make(chan struct{}) + stream.writeGate = gate + r, _ := NewRegistry(context.Background(), testConfig(&fixedTopicHost{stream: stream}, Hooks{})) + defer r.Stop() + a := r.GetOrCreate(peer.ID("peer")) + a.EnableTopicStreams() + topic := "alpha" + if err := a.Send(&pb.RPC{Publish: []*pb.Message{{Topic: &topic, Data: []byte("data")}}}, false); err != nil { + t.Fatal(err) + } + waitDone(t, stream.writeStarted, "blocked topic write") + requirePromptReturn(t, func() { a.CloseTopic(topic) }) + if a.HasTopicWriter(topic) { + t.Fatal("closed writer remained registered") + } + close(gate) +} + +func TestDisableTopicStreamsDoesNotWaitForBlockedOpen(t *testing.T) { + gate := make(chan struct{}) + h := &gatedTopicHost{started: make(chan struct{}), release: gate} + r, _ := NewRegistry(context.Background(), testConfig(h, Hooks{})) + defer r.Stop() + a := r.GetOrCreate(peer.ID("peer")) + a.EnableTopicStreams() + topic := "alpha" + if err := a.Send(&pb.RPC{Publish: []*pb.Message{{Topic: &topic, Data: []byte("data")}}}, false); err != nil { + t.Fatal(err) + } + waitDone(t, h.started, "blocked topic open") + requirePromptReturn(t, a.DisableTopicStreams) + if a.TopicStreamsEnabled() || a.HasTopicWriter(topic) { + t.Fatal("disabled topic writer remained active") + } + close(gate) +} + +func TestTopicSendSplitsWithoutMutationAndWritesHeaderFirst(t *testing.T) { + gate := make(chan struct{}) + stream := newTestStream(peer.ID("peer")) + stream.writeStarted = make(chan struct{}) + stream.writeGate = gate + r, err := NewRegistry(context.Background(), testConfig(&fixedTopicHost{stream: stream}, Hooks{})) + if err != nil { + t.Fatal(err) + } + defer r.Stop() + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(gate) }) } + defer release() + a := r.GetOrCreate(peer.ID("peer")) + a.EnableTopicStreams() + topic := "alpha" + extension := &pb.TestExtension{} + extension.ProtoReflect().SetUnknown([]byte{0x08, 0x01}) + rpc := &pb.RPC{ + Publish: []*pb.Message{{From: []byte("from"), Data: []byte("data"), Topic: &topic}}, + Subscriptions: []*pb.RPC_SubOpts{{Subscribe: proto.Bool(true), Topicid: &topic}}, + Control: &pb.ControlMessage{Ihave: []*pb.ControlIHave{{TopicID: &topic, MessageIDs: []string{"message"}}}}, + TestExtension: extension, + } + before := proto.Clone(rpc).(*pb.RPC) + if err := a.Send(rpc, false); err != nil { + t.Fatal(err) + } + if !proto.Equal(rpc, before) { + t.Fatal("Send mutated caller RPC") + } + waitDone(t, stream.writeStarted, "blocked topic header") + + rpc.Subscriptions[0].Topicid = proto.String("mutated") + rpc.Control.Ihave[0].MessageIDs[0] = "mutated" + rpc.TestExtension.ProtoReflect().SetUnknown([]byte{0x08, 0x02}) + rpc.Publish[0].Data[0] = 'X' + + control, err := a.queue.pop(context.Background()) + if err != nil { + t.Fatal(err) + } + expectedControl := proto.Clone(before).(*pb.RPC) + expectedControl.Publish = nil + if !proto.Equal(control, expectedControl) { + t.Fatalf("control snapshot = %v, want %v", control, expectedControl) + } + + release() + headerFrame := receive(t, stream.writes, "topic header") + var header pb.TopicRPCHeader + decodeFrameInto(t, headerFrame, &header) + if header.GetTopic() != topic { + t.Fatalf("header topic = %q", header.GetTopic()) + } + payloadFrame := receive(t, stream.writes, "topic payload") + var payload pb.TopicRPC + decodeFrameInto(t, payloadFrame, &payload) + if payload.GetPublish() == nil || payload.GetPublish().Topic != nil || string(payload.GetPublish().Data) != "data" { + t.Fatalf("bad payload snapshot: %v", &payload) + } +} + +func TestTopicSendRejectsMissingTopicWithoutPartialEnqueue(t *testing.T) { + h := &topicHost{streams: make(chan *testStream, 1)} + r, err := NewRegistry(context.Background(), testConfig(h, Hooks{})) + if err != nil { + t.Fatal(err) + } + defer r.Stop() + a := r.GetOrCreate(peer.ID("peer")) + a.EnableTopicStreams() + topic := "alpha" + rpc := &pb.RPC{ + Subscriptions: []*pb.RPC_SubOpts{{Topicid: &topic}}, + Publish: []*pb.Message{ + {Topic: &topic, Data: []byte("valid")}, + {Data: []byte("missing-topic")}, + }, + } + if err := a.Send(rpc, false); !errors.Is(err, ErrInvalidTopicRPC) { + t.Fatalf("send error = %v", err) + } + if got := len(a.queue.normal); got != 0 { + t.Fatalf("control queue length = %d, want 0", got) + } + select { + case <-h.streams: + t.Fatal("opened stream after rejecting RPC") + case <-time.After(20 * time.Millisecond): + } +} + +func TestTopicSendQueueFullDoesNotPartiallyEnqueue(t *testing.T) { + h := &topicHost{streams: make(chan *testStream, 1)} + cfg := testConfig(h, Hooks{}) + cfg.QueueSize = 1 + r, err := NewRegistry(context.Background(), cfg) + if err != nil { + t.Fatal(err) + } + defer r.Stop() + a := r.GetOrCreate(peer.ID("peer")) + a.EnableTopicStreams() + if err := a.queue.push(&pb.RPC{}, false); err != nil { + t.Fatal(err) + } + topic := "alpha" + rpc := &pb.RPC{Subscriptions: []*pb.RPC_SubOpts{{Topicid: &topic}}, Publish: []*pb.Message{{Topic: &topic, Data: []byte("data")}}} + if err := a.Send(rpc, false); !errors.Is(err, ErrQueueFull) { + t.Fatalf("send error = %v", err) + } + if len(a.topicWriters) != 0 { + t.Fatal("topic payload enqueued despite full control queue") + } +} + +func TestTopicSendReusesStreamAndCloseTopicReplacesAfterClose(t *testing.T) { + h := &topicHost{streams: make(chan *testStream, 3)} + r, _ := NewRegistry(context.Background(), testConfig(h, Hooks{})) + defer r.Stop() + a := r.GetOrCreate(peer.ID("peer")) + a.EnableTopicStreams() + topic := "alpha" + makeRPC := func(data string) *pb.RPC { return &pb.RPC{Publish: []*pb.Message{{Data: []byte(data), Topic: &topic}}} } + if err := a.Send(makeRPC("one"), false); err != nil { + t.Fatal(err) + } + first := receive(t, h.streams, "first stream") + receive(t, first.writes, "header") + receive(t, first.writes, "first payload") + if err := a.Send(makeRPC("two"), false); err != nil { + t.Fatal(err) + } + receive(t, first.writes, "second payload") + a.CloseTopic(topic) + waitDone(t, first.reset, "first close") + if err := a.Send(makeRPC("three"), false); err != nil { + t.Fatal(err) + } + second := receive(t, h.streams, "replacement stream") + if second == first { + t.Fatal("stream not replaced") + } +} + +func TestTopicInboundReconstructsPublishAndPartial(t *testing.T) { + p := peer.ID("peer") + inbound := make(chan *pb.RPC, 2) + r, _ := NewRegistry(context.Background(), testConfig(&failingHost{}, Hooks{InboundRPC: func(_ *Actor, _ network.Stream, transport Transport, rpc *pb.RPC) { + if transport != TransportTopic { + t.Errorf("transport = %v, want topic", transport) + } + inbound <- rpc + }})) + defer r.Stop() + a := r.GetOrCreate(p) + a.EnableTopicStreams() + topic := "alpha" + s := newTestStream(p) + s.reads <- streamRead{data: frameProto(t, &pb.TopicRPCHeader{Topic: &topic})} + s.reads <- streamRead{data: frameProto(t, &pb.TopicRPC{Payload: &pb.TopicRPC_Publish{Publish: &pb.Message{Data: []byte("data")}}})} + s.reads <- streamRead{data: frameProto(t, &pb.TopicRPC{Payload: &pb.TopicRPC_Partial{Partial: &pb.PartialMessagesExtension{PartialMessage: []byte("part")}}})} + s.reads <- streamRead{err: io.EOF} + done := make(chan struct{}) + go func() { a.HandleTopicInbound(s); close(done) }() + publish := receive(t, inbound, "publish") + if len(publish.Publish) != 1 || publish.Publish[0].GetTopic() != topic { + t.Fatalf("publish not reconstructed: %v", publish) + } + partial := receive(t, inbound, "partial") + if partial.Partial == nil || partial.Partial.GetTopicID() != topic { + t.Fatalf("partial not reconstructed: %v", partial) + } + waitDone(t, done, "topic inbound") +} + +func TestTopicInboundRejectsEmptyDataAndWireTopic(t *testing.T) { + for _, tc := range []struct { + name string + payload *pb.TopicRPC + }{ + {name: "empty data", payload: &pb.TopicRPC{Payload: &pb.TopicRPC_Publish{Publish: &pb.Message{}}}}, + {name: "wire topic", payload: &pb.TopicRPC{Payload: &pb.TopicRPC_Publish{Publish: &pb.Message{Data: []byte("x"), Topic: proto.String("bad")}}}}, + {name: "empty oneof", payload: &pb.TopicRPC{}}, + {name: "nil publish", payload: &pb.TopicRPC{Payload: &pb.TopicRPC_Publish{}}}, + } { + t.Run(tc.name, func(t *testing.T) { + p := peer.ID("peer") + r, _ := NewRegistry(context.Background(), testConfig(&failingHost{}, Hooks{})) + defer r.Stop() + a := r.GetOrCreate(p) + a.EnableTopicStreams() + topic := "alpha" + s := newTestStream(p) + s.reads <- streamRead{data: frameProto(t, &pb.TopicRPCHeader{Topic: &topic})} + s.reads <- streamRead{data: frameProto(t, tc.payload)} + done := make(chan struct{}) + go func() { a.HandleTopicInbound(s); close(done) }() + waitDone(t, done, "violation") + conn := s.conn.(*testConn) + conn.mu.Lock() + code, closed := conn.closeCode, conn.closed + conn.mu.Unlock() + if closed != 1 || code != TopicStreamsViolation { + t.Fatalf("close = %d/%x", closed, code) + } + }) + } +} + +func TestTopicHeaderDeadline(t *testing.T) { + p := peer.ID("peer") + r, _ := NewRegistry(context.Background(), testConfig(&failingHost{}, Hooks{})) + defer r.Stop() + a := r.GetOrCreate(p) + a.EnableTopicStreams() + s := newTestStream(p) + done := make(chan struct{}) + go func() { a.HandleTopicInbound(s); close(done) }() + select { + case <-done: + t.Fatal("returned before input") + case <-time.After(20 * time.Millisecond): + } + a.Retire() + waitDone(t, done, "retirement") +} + +func TestTopicViolationClosesOffendingConnection(t *testing.T) { + p := peer.ID("peer") + r, err := NewRegistry(context.Background(), testConfig(&failingHost{}, Hooks{})) + if err != nil { + t.Fatal(err) + } + defer r.Stop() + a := r.GetOrCreate(p) + offending := newTestStream(p) + a.ProtocolViolation(offending) + conn := offending.conn.(*testConn) + conn.mu.Lock() + defer conn.mu.Unlock() + if conn.closed != 1 || conn.closeCode != TopicStreamsViolation { + t.Fatalf("offending connection close = %d/%x", conn.closed, conn.closeCode) + } +} diff --git a/pb/rpc.pb.go b/pb/rpc.pb.go index d3925f21..a6bc6e36 100644 --- a/pb/rpc.pb.go +++ b/pb/rpc.pb.go @@ -1,6 +1,6 @@ // Code generated by protoc-gen-go. DO NOT EDIT. // versions: -// protoc-gen-go v1.36.6 +// protoc-gen-go v1.36.11 // protoc v7.34.1 // source: rpc.proto @@ -524,6 +524,7 @@ type ControlExtensions struct { // Experimental extensions must use field numbers larger than 0x200000 to be // encoded with 4 bytes TestExtension *bool `protobuf:"varint,6492434,opt,name=testExtension" json:"testExtension,omitempty"` + TopicStreams *bool `protobuf:"varint,6492435,opt,name=topicStreams" json:"topicStreams,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -572,6 +573,13 @@ func (x *ControlExtensions) GetTestExtension() bool { return false } +func (x *ControlExtensions) GetTopicStreams() bool { + if x != nil && x.TopicStreams != nil { + return *x.TopicStreams + } + return false +} + type PeerInfo struct { state protoimpl.MessageState `protogen:"open.v1"` PeerID []byte `protobuf:"bytes,1,opt,name=peerID" json:"peerID,omitempty"` @@ -730,6 +738,132 @@ func (x *PartialMessagesExtension) GetPartsMetadata() []byte { return nil } +type TopicRPCHeader struct { + state protoimpl.MessageState `protogen:"open.v1"` + Topic *string `protobuf:"bytes,1,opt,name=topic" json:"topic,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *TopicRPCHeader) Reset() { + *x = TopicRPCHeader{} + mi := &file_rpc_proto_msgTypes[12] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *TopicRPCHeader) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*TopicRPCHeader) ProtoMessage() {} + +func (x *TopicRPCHeader) ProtoReflect() protoreflect.Message { + mi := &file_rpc_proto_msgTypes[12] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use TopicRPCHeader.ProtoReflect.Descriptor instead. +func (*TopicRPCHeader) Descriptor() ([]byte, []int) { + return file_rpc_proto_rawDescGZIP(), []int{12} +} + +func (x *TopicRPCHeader) GetTopic() string { + if x != nil && x.Topic != nil { + return *x.Topic + } + return "" +} + +type TopicRPC struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Types that are valid to be assigned to Payload: + // + // *TopicRPC_Publish + // *TopicRPC_Partial + Payload isTopicRPC_Payload `protobuf_oneof:"payload"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *TopicRPC) Reset() { + *x = TopicRPC{} + mi := &file_rpc_proto_msgTypes[13] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *TopicRPC) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*TopicRPC) ProtoMessage() {} + +func (x *TopicRPC) ProtoReflect() protoreflect.Message { + mi := &file_rpc_proto_msgTypes[13] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use TopicRPC.ProtoReflect.Descriptor instead. +func (*TopicRPC) Descriptor() ([]byte, []int) { + return file_rpc_proto_rawDescGZIP(), []int{13} +} + +func (x *TopicRPC) GetPayload() isTopicRPC_Payload { + if x != nil { + return x.Payload + } + return nil +} + +func (x *TopicRPC) GetPublish() *Message { + if x != nil { + if x, ok := x.Payload.(*TopicRPC_Publish); ok { + return x.Publish + } + } + return nil +} + +func (x *TopicRPC) GetPartial() *PartialMessagesExtension { + if x != nil { + if x, ok := x.Payload.(*TopicRPC_Partial); ok { + return x.Partial + } + } + return nil +} + +type isTopicRPC_Payload interface { + isTopicRPC_Payload() +} + +type TopicRPC_Publish struct { + Publish *Message `protobuf:"bytes,1,opt,name=publish,oneof"` +} + +type TopicRPC_Partial struct { + Partial *PartialMessagesExtension `protobuf:"bytes,2,opt,name=partial,oneof"` +} + +func (*TopicRPC_Publish) isTopicRPC_Payload() {} + +func (*TopicRPC_Partial) isTopicRPC_Payload() {} + type RPC_SubOpts struct { state protoimpl.MessageState `protogen:"open.v1"` Subscribe *bool `protobuf:"varint,1,opt,name=subscribe" json:"subscribe,omitempty"` // subscribe or unsubcribe @@ -747,7 +881,7 @@ type RPC_SubOpts struct { func (x *RPC_SubOpts) Reset() { *x = RPC_SubOpts{} - mi := &file_rpc_proto_msgTypes[12] + mi := &file_rpc_proto_msgTypes[14] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -759,7 +893,7 @@ func (x *RPC_SubOpts) String() string { func (*RPC_SubOpts) ProtoMessage() {} func (x *RPC_SubOpts) ProtoReflect() protoreflect.Message { - mi := &file_rpc_proto_msgTypes[12] + mi := &file_rpc_proto_msgTypes[14] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -854,11 +988,12 @@ const file_rpc_proto_rawDesc = "" + "\x10ControlIDontWant\x12\x1e\n" + "\n" + "messageIDs\x18\x01 \x03(\tR\n" + - "messageIDs\"f\n" + + "messageIDs\"\x8d\x01\n" + "\x11ControlExtensions\x12(\n" + "\x0fpartialMessages\x18\n" + " \x01(\bR\x0fpartialMessages\x12'\n" + - "\rtestExtension\x18\x92\xa2\x8c\x03 \x01(\bR\rtestExtension\"N\n" + + "\rtestExtension\x18\x92\xa2\x8c\x03 \x01(\bR\rtestExtension\x12%\n" + + "\ftopicStreams\x18\x93\xa2\x8c\x03 \x01(\bR\ftopicStreams\"N\n" + "\bPeerInfo\x12\x16\n" + "\x06peerID\x18\x01 \x01(\fR\x06peerID\x12*\n" + "\x10signedPeerRecord\x18\x02 \x01(\fR\x10signedPeerRecord\"\x0f\n" + @@ -867,7 +1002,13 @@ const file_rpc_proto_rawDesc = "" + "\atopicID\x18\x01 \x01(\tR\atopicID\x12\x18\n" + "\agroupID\x18\x02 \x01(\fR\agroupID\x12&\n" + "\x0epartialMessage\x18\x03 \x01(\fR\x0epartialMessage\x12$\n" + - "\rpartsMetadata\x18\x04 \x01(\fR\rpartsMetadataB1Z/github.com/libp2p/go-libp2p-pubsub/pb;pubsub_pb" + "\rpartsMetadata\x18\x04 \x01(\fR\rpartsMetadata\"&\n" + + "\x0eTopicRPCHeader\x12\x14\n" + + "\x05topic\x18\x01 \x01(\tR\x05topic\"\x86\x01\n" + + "\bTopicRPC\x12.\n" + + "\apublish\x18\x01 \x01(\v2\x12.pubsub.pb.MessageH\x00R\apublish\x12?\n" + + "\apartial\x18\x02 \x01(\v2#.pubsub.pb.PartialMessagesExtensionH\x00R\apartialB\t\n" + + "\apayloadB1Z/github.com/libp2p/go-libp2p-pubsub/pb;pubsub_pb" var ( file_rpc_proto_rawDescOnce sync.Once @@ -881,7 +1022,7 @@ func file_rpc_proto_rawDescGZIP() []byte { return file_rpc_proto_rawDescData } -var file_rpc_proto_msgTypes = make([]protoimpl.MessageInfo, 13) +var file_rpc_proto_msgTypes = make([]protoimpl.MessageInfo, 15) var file_rpc_proto_goTypes = []any{ (*RPC)(nil), // 0: pubsub.pb.RPC (*Message)(nil), // 1: pubsub.pb.Message @@ -895,10 +1036,12 @@ var file_rpc_proto_goTypes = []any{ (*PeerInfo)(nil), // 9: pubsub.pb.PeerInfo (*TestExtension)(nil), // 10: pubsub.pb.TestExtension (*PartialMessagesExtension)(nil), // 11: pubsub.pb.PartialMessagesExtension - (*RPC_SubOpts)(nil), // 12: pubsub.pb.RPC.SubOpts + (*TopicRPCHeader)(nil), // 12: pubsub.pb.TopicRPCHeader + (*TopicRPC)(nil), // 13: pubsub.pb.TopicRPC + (*RPC_SubOpts)(nil), // 14: pubsub.pb.RPC.SubOpts } var file_rpc_proto_depIdxs = []int32{ - 12, // 0: pubsub.pb.RPC.subscriptions:type_name -> pubsub.pb.RPC.SubOpts + 14, // 0: pubsub.pb.RPC.subscriptions:type_name -> pubsub.pb.RPC.SubOpts 1, // 1: pubsub.pb.RPC.publish:type_name -> pubsub.pb.Message 2, // 2: pubsub.pb.RPC.control:type_name -> pubsub.pb.ControlMessage 11, // 3: pubsub.pb.RPC.partial:type_name -> pubsub.pb.PartialMessagesExtension @@ -910,11 +1053,13 @@ var file_rpc_proto_depIdxs = []int32{ 7, // 9: pubsub.pb.ControlMessage.idontwant:type_name -> pubsub.pb.ControlIDontWant 8, // 10: pubsub.pb.ControlMessage.extensions:type_name -> pubsub.pb.ControlExtensions 9, // 11: pubsub.pb.ControlPrune.peers:type_name -> pubsub.pb.PeerInfo - 12, // [12:12] is the sub-list for method output_type - 12, // [12:12] is the sub-list for method input_type - 12, // [12:12] is the sub-list for extension type_name - 12, // [12:12] is the sub-list for extension extendee - 0, // [0:12] is the sub-list for field type_name + 1, // 12: pubsub.pb.TopicRPC.publish:type_name -> pubsub.pb.Message + 11, // 13: pubsub.pb.TopicRPC.partial:type_name -> pubsub.pb.PartialMessagesExtension + 14, // [14:14] is the sub-list for method output_type + 14, // [14:14] is the sub-list for method input_type + 14, // [14:14] is the sub-list for extension type_name + 14, // [14:14] is the sub-list for extension extendee + 0, // [0:14] is the sub-list for field type_name } func init() { file_rpc_proto_init() } @@ -922,13 +1067,17 @@ func file_rpc_proto_init() { if File_rpc_proto != nil { return } + file_rpc_proto_msgTypes[13].OneofWrappers = []any{ + (*TopicRPC_Publish)(nil), + (*TopicRPC_Partial)(nil), + } type x struct{} out := protoimpl.TypeBuilder{ File: protoimpl.DescBuilder{ GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_rpc_proto_rawDesc), len(file_rpc_proto_rawDesc)), NumEnums: 0, - NumMessages: 13, + NumMessages: 15, NumExtensions: 0, NumServices: 0, }, diff --git a/pb/rpc.proto b/pb/rpc.proto index c21b9360..80e0d320 100644 --- a/pb/rpc.proto +++ b/pb/rpc.proto @@ -85,6 +85,7 @@ message ControlExtensions { // Experimental extensions must use field numbers larger than 0x200000 to be // encoded with 4 bytes optional bool testExtension = 6492434; + optional bool topicStreams = 6492435; } message PeerInfo { @@ -104,3 +105,15 @@ message PartialMessagesExtension { // An encoded representation of the parts a peer has and wants. optional bytes partsMetadata = 4; } + + +message TopicRPCHeader { + optional string topic = 1; +} + +message TopicRPC { + oneof payload { + Message publish = 1; + PartialMessagesExtension partial = 2; + } +} diff --git a/pb/topic_streams_test.go b/pb/topic_streams_test.go new file mode 100644 index 00000000..5337890d --- /dev/null +++ b/pb/topic_streams_test.go @@ -0,0 +1,19 @@ +package pubsub_pb + +import ( + "testing" +) + +func TestTopicRPCPayloadOneof(t *testing.T) { + publish := &Message{Data: []byte("payload")} + rpc := &TopicRPC{Payload: &TopicRPC_Publish{Publish: publish}} + if rpc.GetPublish() != publish || rpc.GetPartial() != nil { + t.Fatal("expected publish payload") + } + + partial := &PartialMessagesExtension{PartialMessage: []byte("partial")} + rpc.Payload = &TopicRPC_Partial{Partial: partial} + if rpc.GetPartial() != partial || rpc.GetPublish() != nil { + t.Fatal("expected partial payload") + } +} diff --git a/pubsub.go b/pubsub.go index f80a3ab1..70ea348d 100644 --- a/pubsub.go +++ b/pubsub.go @@ -159,9 +159,10 @@ type PubSub struct { peerComm *peercomm.Registry - seenMessages timecache.TimeCache - seenMsgTTL time.Duration - seenMsgStrategy timecache.Strategy + seenMessages timecache.TimeCache + seenMsgTTL time.Duration + seenMsgStrategy timecache.Strategy + recentUnsubscribed map[string]time.Time // deliveredMessages tracks messages delivered to subscribers and routed to peers. // It deduplicates publishes of already-delivered messages, while still allowing @@ -290,6 +291,7 @@ type RPC struct { // unexported on purpose, not sending this over the wire from peer.ID transport peercomm.Transport + stream network.Stream } func wrapInboundRPC(rpc *pb.RPC, from peer.ID, transport peercomm.Transport) *RPC { @@ -643,6 +645,7 @@ func NewPubSub(ctx context.Context, h host.Host, rt PubSubRouter, opts ...Option blacklistPeer: make(chan peer.ID), seenMsgTTL: TimeCacheDuration, seenMsgStrategy: TimeCacheStrategy, + recentUnsubscribed: make(map[string]time.Time), idGen: newMsgIdGenerator(), counter: uint64(time.Now().UnixNano()), } @@ -686,7 +689,9 @@ func NewPubSub(ctx context.Context, h host.Host, rt PubSubRouter, opts ...Option ps.enqueuePeerEvent(incomingUnion{kind: incomingKindNewStream, actor: a, s: s}) }, InboundRPC: func(a *peercomm.Actor, s network.Stream, transport peercomm.Transport, rpc *pb.RPC) { - ps.enqueuePeerEvent(incomingUnion{kind: incomingKindRPC, actor: a, s: s, rpc: wrapInboundRPC(rpc, a.Peer(), transport)}) + wrapped := wrapInboundRPC(rpc, a.Peer(), transport) + wrapped.stream = s + ps.enqueuePeerEvent(incomingUnion{kind: incomingKindRPC, actor: a, s: s, rpc: wrapped}) }, InboundClosed: func(a *peercomm.Actor, s network.Stream) { ps.enqueuePeerEvent(incomingUnion{kind: incomingKindClosedStream, actor: a, s: s}) @@ -706,11 +711,35 @@ func NewPubSub(ctx context.Context, h host.Host, rt PubSubRouter, opts ...Option OutboundDead: func(a *peercomm.Actor, s network.Stream, err error) { ps.enqueuePeerEvent(incomingUnion{kind: incomingKindOutboundDead, actor: a, s: s, err: err}) }, + TopicAllowed: ps.topicStreamAllowed, + TopicMisbehavior: func(a *peercomm.Actor, _ string) { + ps.topicStreamMisbehavior(a) + }, }, }) if err != nil { return nil, err } + if gs, ok := rt.(*GossipSubRouter); ok { + gs.extensions.enableTopicStreams = func(id peer.ID) { + if actor, found := ps.peerComm.Lookup(id); found { + actor.EnableTopicStreams() + } + } + gs.extensions.disableTopicStreams = func(id peer.ID) { + if actor, found := ps.peerComm.Lookup(id); found { + actor.DisableTopicStreams() + } + } + gs.extensions.protocolViolation = func(id peer.ID, stream network.Stream) { + if actor, found := ps.peerComm.Lookup(id); found { + actor.ProtocolViolation(stream) + } + } + if gs.extensions.myExtensions.TopicStreams { + h.SetStreamHandler(peercomm.TopicStreamsProtocol, ps.handleTopicStream) + } + } for _, id := range rt.Protocols() { if ps.protoMatchFunc != nil { @@ -1222,6 +1251,10 @@ func (p *PubSub) handleRemoveSubscription(sub *Subscription) { if len(subs) == 0 { delete(p.mySubs, sub.topic) + p.recentUnsubscribed[sub.topic] = time.Now() + for _, actor := range p.peerComm.All() { + actor.CloseTopic(sub.topic) + } // stop announcing only if there are no more subs and relays if p.myRelays[sub.topic] == 0 { @@ -1319,6 +1352,10 @@ func (p *PubSub) handleRemoveRelay(topic string) { // stop announcing only if there are no more relays and subs if len(p.mySubs[topic]) == 0 { + p.recentUnsubscribed[topic] = time.Now() + for _, actor := range p.peerComm.All() { + actor.CloseTopic(topic) + } p.disc.StopAdvertise(topic) p.announce(topic, false) p.rt.Leave(topic) @@ -1483,6 +1520,9 @@ func (p *PubSub) notifyLeave(topic string, pid peer.ID) { } func (p *PubSub) handleIncomingRPC(rpc *RPC) { + if gs, ok := p.rt.(*GossipSubRouter); ok && !gs.extensions.Preprocess(rpc) { + return + } // pass the rpc through app specific validation (if any available). if p.appSpecificRpcInspector != nil { // check if the RPC is allowed by the external inspector @@ -1528,6 +1568,11 @@ func (p *PubSub) handleIncomingRPC(rpc *RPC) { } } } else { + if p.peerComm != nil { + if actor, found := p.peerComm.Lookup(rpc.from); found { + actor.CloseTopic(t) + } + } tmap, ok := p.topics[t] if !ok { continue diff --git a/pubsub_test.go b/pubsub_test.go index 916fdefd..1f841849 100644 --- a/pubsub_test.go +++ b/pubsub_test.go @@ -28,25 +28,48 @@ func synctestTest(t *testing.T, f func(t *testing.T)) { synctest.Test(t, f) } -func TestWrapInboundRPCPreservesMetadataAndCopiesProto(t *testing.T) { +func TestWrapInboundRPCPreservesUnknownFieldsAndDeepCopies(t *testing.T) { topic := "topic" - source := &pb.RPC{Subscriptions: []*pb.RPC_SubOpts{{Topicid: &topic}}} + source := &pb.RPC{ + Subscriptions: []*pb.RPC_SubOpts{{Topicid: &topic}}, + Publish: []*pb.Message{{Data: []byte("payload"), Topic: &topic}}, + } unknown := protowire.AppendTag(nil, 100, protowire.BytesType) unknown = protowire.AppendBytes(unknown, []byte("extension")) source.ProtoReflect().SetUnknown(unknown) from := peer.ID("peer-a") - wrapped := wrapInboundRPC(source, from, peercomm.TransportTopic) - if wrapped.from != from || wrapped.transport != peercomm.TransportTopic { - t.Fatalf("metadata = (%q, %v), want (%q, %v)", wrapped.from, wrapped.transport, from, peercomm.TransportTopic) + transport := peercomm.TransportTopic + wrapped := wrapInboundRPC(source, from, transport) + + if len(wrapped.Subscriptions) != 1 || wrapped.Subscriptions[0].GetTopicid() != topic { + t.Fatalf("known subscription field was not copied: %v", wrapped.Subscriptions) + } + if len(wrapped.Publish) != 1 || string(wrapped.Publish[0].GetData()) != "payload" || wrapped.Publish[0].GetTopic() != topic { + t.Fatalf("known publish field was not copied: %v", wrapped.Publish) + } + if got := wrapped.ProtoReflect().GetUnknown(); !bytes.Equal(got, unknown) { + t.Fatalf("unknown fields differ: got %x, want %x", got, unknown) + } + if wrapped.from != from { + t.Fatalf("from metadata differs: got %q, want %q", wrapped.from, from) } - if wrapped.Subscriptions[0].GetTopicid() != topic || !bytes.Equal(wrapped.ProtoReflect().GetUnknown(), unknown) { - t.Fatal("protobuf fields were not preserved") + if wrapped.transport != transport { + t.Fatalf("transport metadata differs: got %v, want %v", wrapped.transport, transport) } + source.Subscriptions[0].Topicid = nil + source.Publish[0].Data[0] = 'P' + source.Publish[0].Topic = nil source.ProtoReflect().SetUnknown(nil) - if wrapped.Subscriptions[0].GetTopicid() != topic || !bytes.Equal(wrapped.ProtoReflect().GetUnknown(), unknown) { - t.Fatal("wrapped RPC shares mutable protobuf state with source") + if wrapped.Subscriptions[0].GetTopicid() != topic { + t.Fatal("wrapped subscription changed after mutating source") + } + if string(wrapped.Publish[0].GetData()) != "payload" || wrapped.Publish[0].GetTopic() != topic { + t.Fatal("wrapped publish changed after mutating source") + } + if got := wrapped.ProtoReflect().GetUnknown(); !bytes.Equal(got, unknown) { + t.Fatalf("wrapped unknown fields changed after mutating source: got %x, want %x", got, unknown) } } @@ -85,6 +108,29 @@ func TestClearPeerFromTopicsStateRemovesEmptyTopicMap(t *testing.T) { } } +func TestIsRecentlyUnsubscribedCleansUpExpiredEntry(t *testing.T) { + now := time.Now() + ps := &PubSub{ + recentUnsubscribed: map[string]time.Time{ + "recent": now.Add(-GossipSubUnsubscribeBackoff), + "expired": now.Add(-GossipSubUnsubscribeBackoff - time.Nanosecond), + }, + } + + if !ps.isRecentlyUnsubscribed("recent", now) { + t.Fatal("expected topic within unsubscribe backoff to be recent") + } + if ps.isRecentlyUnsubscribed("expired", now) { + t.Fatal("expected topic past unsubscribe backoff not to be recent") + } + if _, ok := ps.recentUnsubscribed["recent"]; !ok { + t.Fatal("expected recent topic to remain tracked") + } + if _, ok := ps.recentUnsubscribed["expired"]; ok { + t.Fatal("expected expired topic to be removed during lookup") + } +} + func TestHandleIncomingRPCUnsubscribeRemovesEmptyTopicMap(t *testing.T) { pid := peer.ID("peer-a") other := peer.ID("peer-b") diff --git a/topic_streams_test.go b/topic_streams_test.go new file mode 100644 index 00000000..9c056024 --- /dev/null +++ b/topic_streams_test.go @@ -0,0 +1,391 @@ +package pubsub + +import ( + "context" + "io" + "iter" + "log/slog" + "testing" + "time" + + "github.com/libp2p/go-libp2p-pubsub/internal/peercomm" + "github.com/libp2p/go-libp2p-pubsub/partialmessages" + pb "github.com/libp2p/go-libp2p-pubsub/pb" + "github.com/libp2p/go-libp2p/core/peer" + "github.com/libp2p/go-libp2p/core/protocol" + "github.com/libp2p/go-msgio" + "google.golang.org/protobuf/proto" +) + +func waitForTopicStreams(t *testing.T, ps *PubSub, id peer.ID) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + done := make(chan bool, 1) + select { + case ps.eval <- func() { + a, ok := ps.peerComm.Lookup(id) + done <- ok && a.TopicStreamsEnabled() + }: + case <-ps.ctx.Done(): + t.Fatal(ps.ctx.Err()) + } + if <-done { + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatal("topic streams negotiation timed out") +} + +func waitForTopicPeer(t *testing.T, ps *PubSub, topic string, id peer.ID) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + done := make(chan bool, 1) + select { + case ps.eval <- func() { + _, ok := ps.topics[topic][id] + done <- ok + }: + case <-ps.ctx.Done(): + t.Fatal(ps.ctx.Err()) + } + if <-done { + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("peer %s subscription to %q timed out", id, topic) +} + +func TestTopicStreamsV12FallbackStaysOnControl(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + hosts := getDefaultHosts(t, 2) + ps0 := getGossipsub(ctx, hosts[0], WithTopicStreams()) + ps1 := getGossipsub(ctx, hosts[1], WithGossipSubProtocols([]protocol.ID{GossipSubID_v12}, GossipSubDefaultFeatures)) + sub, err := ps1.Subscribe("topic-stream-v12-fallback") + if err != nil { + t.Fatal(err) + } + connect(t, hosts[0], hosts[1]) + waitForTopicPeer(t, ps0, "topic-stream-v12-fallback", hosts[1].ID()) + data := []byte("control-fallback") + if err := ps0.Publish("topic-stream-v12-fallback", data); err != nil { + t.Fatal(err) + } + readCtx, readCancel := context.WithTimeout(ctx, 5*time.Second) + defer readCancel() + got, err := sub.Next(readCtx) + if err != nil || string(got.Data) != string(data) { + t.Fatalf("v1.2 fallback delivery = %q, %v", got.GetData(), err) + } +} + +func TestTopicStreamsRejectsUnnegotiatedAndControlPayload(t *testing.T) { + t.Run("unnegotiated stream", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + hosts := getDefaultHosts(t, 2) + _ = getGossipsub(ctx, hosts[0], WithTopicStreams()) + _ = getGossipsub(ctx, hosts[1]) + connect(t, hosts[0], hosts[1]) + s, err := hosts[1].NewStream(ctx, hosts[0].ID(), peercomm.TopicStreamsProtocol) + if err != nil { + t.Fatal(err) + } + w := msgio.NewVarintWriter(s) + topic := "not-negotiated" + if err := w.WriteMsg(mustMarshalTopicStream(t, &pb.TopicRPCHeader{Topic: &topic})); err != nil { + t.Fatal(err) + } + _ = s.SetReadDeadline(time.Now().Add(3 * time.Second)) + var one [1]byte + if _, err := s.Read(one[:]); err == nil { + t.Fatal("unnegotiated topic stream remained open") + } + }) + + t.Run("payload on control", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + hosts := getDefaultHosts(t, 2) + ps := getGossipsubs(ctx, hosts, WithTopicStreams()) + connect(t, hosts[0], hosts[1]) + waitForTopicStreams(t, ps[0], hosts[1].ID()) + s, err := hosts[1].NewStream(ctx, hosts[0].ID(), GossipSubID_v13) + if err != nil { + t.Fatal(err) + } + w := msgio.NewVarintWriter(s) + topic := "forbidden-control" + rpc := &pb.RPC{Publish: []*pb.Message{{Topic: &topic, Data: []byte("bad")}}} + if err := w.WriteMsg(mustMarshalTopicStream(t, rpc)); err != nil { + t.Fatal(err) + } + _ = s.SetReadDeadline(time.Now().Add(3 * time.Second)) + var one [1]byte + if _, err := s.Read(one[:]); err == nil { + t.Fatal("control payload violation did not close connection") + } + }) +} + +func TestTopicStreamsRejectsUnwantedTopicAndPenalizesPeer(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + hosts := getDefaultHosts(t, 2) + ps0 := getGossipsub(ctx, hosts[0], WithTopicStreams(), WithPeerScore( + &PeerScoreParams{ + AppSpecificScore: func(peer.ID) float64 { return 0 }, + BehaviourPenaltyWeight: -1, + BehaviourPenaltyDecay: ScoreParameterDecay(time.Minute), + DecayInterval: DefaultDecayInterval, + DecayToZero: DefaultDecayToZero, + }, + &PeerScoreThresholds{ + GossipThreshold: -100, + PublishThreshold: -500, + GraylistThreshold: -1000, + }, + )) + ps1 := getGossipsub(ctx, hosts[1], WithTopicStreams()) + connect(t, hosts[0], hosts[1]) + waitForTopicStreams(t, ps0, hosts[1].ID()) + waitForTopicStreams(t, ps1, hosts[0].ID()) + + s, err := hosts[1].NewStream(ctx, hosts[0].ID(), peercomm.TopicStreamsProtocol) + if err != nil { + t.Fatal(err) + } + w := msgio.NewVarintWriter(s) + topic := "unwanted-topic" + if err := w.WriteMsg(mustMarshalTopicStream(t, &pb.TopicRPCHeader{Topic: &topic})); err != nil { + t.Fatal(err) + } + _ = s.SetReadDeadline(time.Now().Add(3 * time.Second)) + var one [1]byte + if _, err := s.Read(one[:]); err == nil { + t.Fatal("unwanted topic stream remained open") + } + + score := make(chan float64, 1) + select { + case ps0.eval <- func() { + score <- ps0.rt.(*GossipSubRouter).score.Score(hosts[1].ID()) + }: + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + if got := <-score; got >= 0 { + t.Fatalf("peer score = %f, want negative after unwanted topic stream", got) + } +} + +func mustMarshalTopicStream(t *testing.T, message proto.Message) []byte { + t.Helper() + data, err := proto.Marshal(message) + if err != nil { + t.Fatal(err) + } + return data +} + +func waitForTopicWriter(t *testing.T, ps *PubSub, id peer.ID, topic string, want bool) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + result := make(chan bool, 1) + select { + case ps.eval <- func() { + actor, ok := ps.peerComm.Lookup(id) + result <- ok && actor.HasTopicWriter(topic) == want + }: + case <-ps.ctx.Done(): + t.Fatal(ps.ctx.Err()) + } + if <-result { + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("topic writer for peer %s and topic %q did not become %t", id, topic, want) +} + +func nextTopicMessage(t *testing.T, ctx context.Context, sub *Subscription, want []byte) { + t.Helper() + readCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + message, err := sub.Next(readCtx) + if err != nil { + t.Fatal(err) + } + if string(message.Data) != string(want) { + t.Fatalf("message data = %q, want %q", message.Data, want) + } +} + +func setupNegotiatedTopicStreams(t *testing.T, topics0, topics1 []string) (context.Context, context.CancelFunc, []peer.ID, []*PubSub, map[string]*Subscription, map[string]*Subscription) { + t.Helper() + ctx, cancel := context.WithCancel(context.Background()) + hosts := getDefaultHosts(t, 2) + ps := getGossipsubs(ctx, hosts, WithTopicStreams(), WithFloodPublish(true)) + subs0 := make(map[string]*Subscription) + for _, topic := range topics0 { + sub, err := ps[0].Subscribe(topic) + if err != nil { + cancel() + t.Fatal(err) + } + subs0[topic] = sub + } + subs1 := make(map[string]*Subscription) + for _, topic := range topics1 { + sub, err := ps[1].Subscribe(topic) + if err != nil { + cancel() + t.Fatal(err) + } + subs1[topic] = sub + } + connect(t, hosts[0], hosts[1]) + waitForTopicStreams(t, ps[0], hosts[1].ID()) + waitForTopicStreams(t, ps[1], hosts[0].ID()) + for _, topic := range topics1 { + waitForTopicPeer(t, ps[0], topic, hosts[1].ID()) + } + for _, topic := range topics0 { + waitForTopicPeer(t, ps[1], topic, hosts[0].ID()) + } + return ctx, cancel, []peer.ID{hosts[0].ID(), hosts[1].ID()}, ps, subs0, subs1 +} + +func TestTopicStreamsNegotiatedPublishDelivery(t *testing.T) { + topic := "topic-stream-negotiated-publish" + ctx, cancel, ids, ps, _, subs1 := setupNegotiatedTopicStreams(t, nil, []string{topic}) + defer cancel() + + data := []byte("delivered over topic writer") + if err := ps[0].Publish(topic, data); err != nil { + t.Fatal(err) + } + nextTopicMessage(t, ctx, subs1[topic], data) + waitForTopicWriter(t, ps[0], ids[1], topic, true) +} + +func TestTopicStreamsDistinctTopicsUseDistinctWriters(t *testing.T) { + topics := []string{"topic-stream-distinct-a", "topic-stream-distinct-b"} + ctx, cancel, ids, ps, _, subs1 := setupNegotiatedTopicStreams(t, nil, topics) + defer cancel() + + for i, topic := range topics { + data := []byte{byte('a' + i)} + if err := ps[0].Publish(topic, data); err != nil { + t.Fatal(err) + } + nextTopicMessage(t, ctx, subs1[topic], data) + waitForTopicWriter(t, ps[0], ids[1], topic, true) + } +} + +func TestTopicStreamsBidirectionalPublishUsesIndependentWriters(t *testing.T) { + topic := "topic-stream-bidirectional" + ctx, cancel, ids, ps, subs0, subs1 := setupNegotiatedTopicStreams(t, []string{topic}, []string{topic}) + defer cancel() + + if err := ps[0].Publish(topic, []byte("zero to one")); err != nil { + t.Fatal(err) + } + nextTopicMessage(t, ctx, subs0[topic], []byte("zero to one")) + nextTopicMessage(t, ctx, subs1[topic], []byte("zero to one")) + if err := ps[1].Publish(topic, []byte("one to zero")); err != nil { + t.Fatal(err) + } + nextTopicMessage(t, ctx, subs0[topic], []byte("one to zero")) + waitForTopicWriter(t, ps[0], ids[1], topic, true) + waitForTopicWriter(t, ps[1], ids[0], topic, true) +} + +func TestTopicStreamsRemoteUnsubscribeClosesSenderWriter(t *testing.T) { + topic := "topic-stream-unsubscribe" + ctx, cancel, ids, ps, _, subs1 := setupNegotiatedTopicStreams(t, nil, []string{topic}) + defer cancel() + + if err := ps[0].Publish(topic, []byte("open writer")); err != nil { + t.Fatal(err) + } + nextTopicMessage(t, ctx, subs1[topic], []byte("open writer")) + waitForTopicWriter(t, ps[0], ids[1], topic, true) + + subs1[topic].Cancel() + waitForTopicWriter(t, ps[0], ids[1], topic, false) +} + +type topicStreamsPartialState struct{} + +func newTopicStreamsPartialExtension(received chan<- *pb.PartialMessagesExtension) *partialmessages.PartialMessagesExtension[topicStreamsPartialState] { + return &partialmessages.PartialMessagesExtension[topicStreamsPartialState]{ + Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + OnEmitGossip: func(string, []byte, []peer.ID, map[peer.ID]topicStreamsPartialState) {}, + OnIncomingRPC: func(_ peer.ID, _ map[peer.ID]topicStreamsPartialState, rpc *pb.PartialMessagesExtension) error { + if received != nil { + received <- proto.Clone(rpc).(*pb.PartialMessagesExtension) + } + return nil + }, + } +} + +func TestTopicStreamsNegotiatedPartialMessageDelivery(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + hosts := getDefaultHosts(t, 2) + received := make(chan *pb.PartialMessagesExtension, 1) + ps0 := getGossipsub(ctx, hosts[0], WithTopicStreams(), WithPartialMessagesExtension(newTopicStreamsPartialExtension(nil))) + ps1 := getGossipsub(ctx, hosts[1], WithTopicStreams(), WithPartialMessagesExtension(newTopicStreamsPartialExtension(received))) + topic := "topic-stream-partial" + t0, err := ps0.Join(topic, SupportsPartialMessages()) + if err != nil { + t.Fatal(err) + } + if _, err = t0.Subscribe(); err != nil { + t.Fatal(err) + } + t1, err := ps1.Join(topic, RequestPartialMessages()) + if err != nil { + t.Fatal(err) + } + if _, err = t1.Subscribe(); err != nil { + t.Fatal(err) + } + connect(t, hosts[0], hosts[1]) + waitForTopicStreams(t, ps0, hosts[1].ID()) + waitForTopicStreams(t, ps1, hosts[0].ID()) + waitForTopicPeer(t, ps0, topic, hosts[1].ID()) + waitForTopicPeer(t, ps1, topic, hosts[0].ID()) + + group := []byte("group") + payload := []byte("partial payload") + err = PublishPartial(ps0, topic, group, func(states map[peer.ID]topicStreamsPartialState, _ func(peer.ID) bool) iter.Seq2[peer.ID, partialmessages.PublishAction] { + return func(yield func(peer.ID, partialmessages.PublishAction) bool) { + for id := range states { + if !yield(id, partialmessages.PublishAction{EncodedPartialMessage: payload, EncodedPartsMetadata: []byte{1}}) { + return + } + } + } + }) + if err != nil { + t.Fatal(err) + } + select { + case rpc := <-received: + if rpc.GetTopicID() != topic || string(rpc.GetGroupID()) != string(group) || string(rpc.GetPartialMessage()) != string(payload) { + t.Fatalf("partial RPC = %v", rpc) + } + case <-time.After(5 * time.Second): + t.Fatal("partial message delivery timed out") + } + waitForTopicWriter(t, ps0, hosts[1].ID(), topic, true) +}