99 lines
1.8 KiB
Go
99 lines
1.8 KiB
Go
|
package floodsub
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bufio"
|
||
|
|
"context"
|
||
|
|
"io"
|
||
|
|
|
||
|
|
pb "github.com/whyrusleeping/go-floodsub/pb"
|
||
|
|
|
||
|
|
ggio "github.com/gogo/protobuf/io"
|
||
|
|
proto "github.com/gogo/protobuf/proto"
|
||
|
|
inet "github.com/libp2p/go-libp2p/p2p/net"
|
||
|
|
)
|
||
|
|
|
||
|
|
// get the initial RPC containing all of our subscriptions to send to new peers
|
||
|
|
func (p *PubSub) getHelloPacket() *RPC {
|
||
|
|
var rpc RPC
|
||
|
|
for t := range p.myTopics {
|
||
|
|
as := &pb.RPC_SubOpts{
|
||
|
|
Topicid: proto.String(t),
|
||
|
|
Subscribe: proto.Bool(true),
|
||
|
|
}
|
||
|
|
rpc.Subscriptions = append(rpc.Subscriptions, as)
|
||
|
|
}
|
||
|
|
return &rpc
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *PubSub) handleNewStream(s inet.Stream) {
|
||
|
|
defer s.Close()
|
||
|
|
|
||
|
|
r := ggio.NewDelimitedReader(s, 1<<20)
|
||
|
|
for {
|
||
|
|
rpc := new(RPC)
|
||
|
|
err := r.ReadMsg(&rpc.RPC)
|
||
|
|
if err != nil {
|
||
|
|
if err != io.EOF {
|
||
|
|
log.Errorf("error reading rpc from %s: %s", s.Conn().RemotePeer(), err)
|
||
|
|
}
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
rpc.from = s.Conn().RemotePeer()
|
||
|
|
select {
|
||
|
|
case p.incoming <- rpc:
|
||
|
|
case <-p.ctx.Done():
|
||
|
|
return
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *PubSub) handleSendingMessages(ctx context.Context, s inet.Stream, outgoing <-chan *RPC) {
|
||
|
|
var dead bool
|
||
|
|
bufw := bufio.NewWriter(s)
|
||
|
|
wc := ggio.NewDelimitedWriter(bufw)
|
||
|
|
|
||
|
|
writeMsg := func(msg proto.Message) error {
|
||
|
|
err := wc.WriteMsg(msg)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
return bufw.Flush()
|
||
|
|
}
|
||
|
|
|
||
|
|
defer wc.Close()
|
||
|
|
for {
|
||
|
|
select {
|
||
|
|
case rpc, ok := <-outgoing:
|
||
|
|
if !ok {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
if dead {
|
||
|
|
// continue in order to drain messages
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
|
||
|
|
err := writeMsg(&rpc.RPC)
|
||
|
|
if err != nil {
|
||
|
|
log.Errorf("writing message to %s: %s", s.Conn().RemotePeer(), err)
|
||
|
|
dead = true
|
||
|
|
go func() {
|
||
|
|
p.peerDead <- s.Conn().RemotePeer()
|
||
|
|
}()
|
||
|
|
}
|
||
|
|
|
||
|
|
case <-ctx.Done():
|
||
|
|
return
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func rpcWithSubs(subs ...*pb.RPC_SubOpts) *RPC {
|
||
|
|
return &RPC{
|
||
|
|
RPC: pb.RPC{
|
||
|
|
Subscriptions: subs,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
}
|