Merge pull request #14560 from openmina/fixes/libp2p_helper_races

Protect shared data accesses in `libp2p_helper`

src/app/libp2p_helper/src/codanet.go

@@ -234,7 +234,8 @@ type Helper struct {
 	ConnectionManager *CodaConnectionManager
 	BandwidthCounter  *metrics.BandwidthCounter
 	MsgStats          *MessageStats
-	Seeds             []peer.AddrInfo
+	_seeds            []peer.AddrInfo
+	seedsMutex        sync.RWMutex
 	NodeStatus        []byte
 	HeartbeatPeer     func(peer.ID)
 }
@@ -368,6 +369,13 @@ func (h *Helper) SetGatingState(gs *CodaGatingConfig) {
 	}
 }
 
+func (h *Helper) AddSeeds(infos ...peer.AddrInfo) {
+	// TODO: this "_seeds" field is never read anywhere, is it needed?
+	h.seedsMutex.Lock()
+	h._seeds = append(h._seeds, infos...)
+	h.seedsMutex.Unlock()
+}
+
 func (gs *CodaGatingState) TrustPeer(p peer.ID) {
 	gs.trustedPeersMutex.Lock()
 	gs.trustedPeers[p] = struct{}{}
@@ -751,7 +759,7 @@ func MakeHelper(ctx context.Context, listenOn []ma.Multiaddr, externalAddr ma.Mu
 		ConnectionManager: connManager,
 		BandwidthCounter:  bandwidthCounter,
 		MsgStats:          &MessageStats{min: math.MaxUint64},
-		Seeds:             seeds,
+		_seeds:            seeds,
 		HeartbeatPeer: func(p peer.ID) {
 			lanPatcher.Heartbeat(p)
 			wanPatcher.Heartbeat(p)

src/app/libp2p_helper/src/libp2p_helper/app.go

@@ -8,7 +8,6 @@ import (
 	"math"
 	"os"
 	"strconv"
-	"sync"
 	"time"
 
 	ipc "libp2p_ipc"
@@ -29,14 +28,13 @@ func newApp() *app {
 	return &app{
 		P2p:                      nil,
 		Ctx:                      ctx,
-		Subs:                     make(map[uint64]subscription),
-		Topics:                   make(map[string]*pubsub.Topic),
-		ValidatorMutex:           &sync.Mutex{},
-		Validators:               make(map[uint64]*validationStatus),
-		Streams:                  make(map[uint64]net.Stream),
+		_subs:                    make(map[uint64]subscription),
+		_topics:                  make(map[string]*pubsub.Topic),
+		_validators:              make(map[uint64]*validationStatus),
+		_streams:                 make(map[uint64]net.Stream),
 		OutChan:                  outChan,
 		Out:                      bufio.NewWriter(os.Stdout),
-		AddedPeers:               []peer.AddrInfo{},
+		_addedPeers:              []peer.AddrInfo{},
 		MetricsRefreshTime:       time.Minute,
 		metricsCollectionStarted: false,
 		metricsServer:            nil,
@@ -64,6 +62,151 @@ func (app *app) NextId() uint64 {
 	return app.counter
 }
 
+func (app *app) AddPeers(infos ...peer.AddrInfo) {
+	app.addedPeersMutex.Lock()
+	defer app.addedPeersMutex.Unlock()
+	app._addedPeers = append(app._addedPeers, infos...)
+}
+
+func (app *app) GetAddedPeers() []peer.AddrInfo {
+	app.addedPeersMutex.RLock()
+	defer app.addedPeersMutex.RUnlock()
+	copyOfAddedPeers := make([]peer.AddrInfo, len(app._addedPeers))
+	copy(copyOfAddedPeers, app._addedPeers)
+	return copyOfAddedPeers
+}
+
+func (app *app) ResetAddedPeers() {
+	app.addedPeersMutex.Lock()
+	defer app.addedPeersMutex.Unlock()
+	app._addedPeers = nil
+}
+
+func (app *app) AddStream(stream net.Stream) uint64 {
+	streamIdx := app.NextId()
+	app.streamsMutex.Lock()
+	defer app.streamsMutex.Unlock()
+	app._streams[streamIdx] = stream
+	return streamIdx
+}
+
+func (app *app) CloseStream(streamId uint64) error {
+	app.streamsMutex.Lock()
+	defer app.streamsMutex.Unlock()
+	if stream, ok := app._streams[streamId]; ok {
+		delete(app._streams, streamId)
+		err := stream.Close()
+		if err != nil {
+			return badp2p(err)
+		}
+		return nil
+	}
+	return badRPC(errors.New("unknown stream_idx"))
+}
+
+func (app *app) ResetStream(streamId uint64) error {
+	app.streamsMutex.Lock()
+	defer app.streamsMutex.Unlock()
+	if stream, ok := app._streams[streamId]; ok {
+		delete(app._streams, streamId)
+		err := stream.Reset()
+		if err != nil {
+			return badp2p(err)
+		}
+		return nil
+	}
+	return badRPC(errors.New("unknown stream_idx"))
+}
+
+func (app *app) StreamWrite(streamId uint64, data []byte) error {
+	// TODO Consider using a more fine-grained locking strategy,
+	// not using a global mutex to lock on a message sending
+	app.streamsMutex.Lock()
+	defer app.streamsMutex.Unlock()
+	if stream, ok := app._streams[streamId]; ok {
+		n, err := stream.Write(data)
+		if err != nil {
+			// TODO check that it's correct to error out, not repeat writing
+			delete(app._streams, streamId)
+			close_err := stream.Close()
+			if close_err != nil {
+				app.P2p.Logger.Errorf("failed to close stream %d after encountering write failure (%s): %s", streamId, err.Error(), close_err.Error())
+			}
+			return wrapError(badp2p(err), fmt.Sprintf("only wrote %d out of %d bytes", n, len(data)))
+		}
+		return nil
+	}
+	return badRPC(errors.New("unknown stream_idx"))
+}
+
+func (app *app) AddValidator() (uint64, chan pubsub.ValidationResult) {
+	seqno := app.NextId()
+	ch := make(chan pubsub.ValidationResult)
+	app.validatorMutex.Lock()
+	defer app.validatorMutex.Unlock()
+	app._validators[seqno] = new(validationStatus)
+	app._validators[seqno].Completion = ch
+	return seqno, ch
+}
+
+func (app *app) RemoveValidator(seqno uint64) {
+	app.validatorMutex.Lock()
+	defer app.validatorMutex.Unlock()
+	delete(app._validators, seqno)
+}
+
+func (app *app) TimeoutValidator(seqno uint64) {
+	now := time.Now()
+	app.validatorMutex.Lock()
+	defer app.validatorMutex.Unlock()
+	app._validators[seqno].TimedOutAt = &now
+}
+
+func (app *app) FinishValidator(seqno uint64, finish func(st *validationStatus)) bool {
+	app.validatorMutex.Lock()
+	defer app.validatorMutex.Unlock()
+	if st, ok := app._validators[seqno]; ok {
+		finish(st)
+		delete(app._validators, seqno)
+		return true
+	} else {
+		return false
+	}
+}
+
+func (app *app) AddTopic(topicName string, topic *pubsub.Topic) {
+	app.topicsMutex.Lock()
+	defer app.topicsMutex.Unlock()
+	app._topics[topicName] = topic
+}
+
+func (app *app) GetTopic(topicName string) (*pubsub.Topic, bool) {
+	app.topicsMutex.RLock()
+	defer app.topicsMutex.RUnlock()
+	topic, has := app._topics[topicName]
+	return topic, has
+}
+
+func (app *app) AddSubscription(subId uint64, sub subscription) {
+	app.subsMutex.Lock()
+	defer app.subsMutex.Unlock()
+	app._subs[subId] = sub
+}
+
+func (app *app) CancelSubscription(subId uint64) bool {
+	app.subsMutex.Lock()
+	defer app.subsMutex.Unlock()
+
+	if sub, ok := app._subs[subId]; ok {
+		sub.Sub.Cancel()
+		sub.Cancel()
+		delete(app._subs, subId)
+		return true
+	}
+
+	return false
+}
+
 func parseMultiaddrWithID(ma multiaddr.Multiaddr, id peer.ID) (*codaPeerInfo, error) {
 	ipComponent, tcpMaddr := multiaddr.SplitFirst(ma)
 	if !(ipComponent.Protocol().Code == multiaddr.P_IP4 || ipComponent.Protocol().Code == multiaddr.P_IP6) {
@@ -96,6 +239,7 @@ func addrInfoOfString(maddr string) (*peer.AddrInfo, error) {
 	return info, nil
 }
 
+// Writes a message back to the OCaml node
 func (app *app) writeMsg(msg *capnp.Message) {
 	if app.NoUpcalls {
 		return

src/app/libp2p_helper/src/libp2p_helper/config_msg.go

@@ -36,7 +36,7 @@ func (msg BeginAdvertisingReq) handle(app *app, seqno uint64) *capnp.Message {
 		return mkRpcRespError(seqno, needsConfigure())
 	}
 	app.SetConnectionHandlers()
-	for _, info := range app.AddedPeers {
+	for _, info := range app.GetAddedPeers() {
 		app.P2p.Logger.Debug("Trying to connect to: ", info)
 		err := app.P2p.Host.Connect(app.Ctx, info)
 		if err != nil {
@@ -334,7 +334,7 @@ func (msg ConfigureReq) handle(app *app, seqno uint64) *capnp.Message {
 		return mkRpcRespError(seqno, badRPC(err))
 	}
 
-	app.AddedPeers = append(app.AddedPeers, seeds...)
+	app.AddPeers(seeds...)
 
 	directPeersMaList, err := m.DirectPeers()
 	if err != nil {
@@ -372,12 +372,12 @@ func (msg ConfigureReq) handle(app *app, seqno uint64) *capnp.Message {
 	if err != nil {
 		return mkRpcRespError(seqno, badRPC(err))
 	}
-	gatingConfig, err := readGatingConfig(gc, app.AddedPeers)
+	gatingConfig, err := readGatingConfig(gc, app.GetAddedPeers())
 	if err != nil {
 		return mkRpcRespError(seqno, badRPC(err))
 	}
 	if gc.CleanAddedPeers() {
-		app.AddedPeers = nil
+		app.ResetAddedPeers()
 	}
 
 	stateDir, err := m.Statedir()
@@ -593,13 +593,13 @@ func (m SetGatingConfigReq) handle(app *app, seqno uint64) *capnp.Message {
 	var gatingConfig *codanet.CodaGatingConfig
 	gc, err := SetGatingConfigReqT(m).GatingConfig()
 	if err == nil {
-		gatingConfig, err = readGatingConfig(gc, app.AddedPeers)
+		gatingConfig, err = readGatingConfig(gc, app.GetAddedPeers())
 	}
 	if err != nil {
 		return mkRpcRespError(seqno, badRPC(err))
 	}
 	if gc.CleanAddedPeers() {
-		app.AddedPeers = nil
+		app.ResetAddedPeers()
 	}
 	app.P2p.SetGatingState(gatingConfig)
 

src/app/libp2p_helper/src/libp2p_helper/config_msg_test.go

@@ -40,7 +40,7 @@ func TestDHTDiscovery_TwoNodes(t *testing.T) {
 	require.NoError(t, err)
 
 	appB, _ := newTestApp(t, appAInfos, true)
-	appB.AddedPeers = appAInfos
+	appB.AddPeers(appAInfos...)
 	appB.NoMDNS = true
 
 	// begin appB and appA's DHT advertising

src/app/libp2p_helper/src/libp2p_helper/data.go

@@ -23,16 +23,19 @@ import (
 type app struct {
 	P2p                      *codanet.Helper
 	Ctx                      context.Context
-	Subs                     map[uint64]subscription
-	Topics                   map[string]*pubsub.Topic
-	Validators               map[uint64]*validationStatus
-	ValidatorMutex           *sync.Mutex
-	Streams                  map[uint64]net.Stream
-	StreamsMutex             sync.Mutex
+	_subs                    map[uint64]subscription
+	subsMutex                sync.Mutex
+	_topics                  map[string]*pubsub.Topic
+	topicsMutex              sync.RWMutex
+	_validators              map[uint64]*validationStatus
+	validatorMutex           sync.Mutex
+	_streams                 map[uint64]net.Stream
+	streamsMutex             sync.Mutex
 	Out                      *bufio.Writer
 	OutChan                  chan *capnp.Message
 	Bootstrapper             io.Closer
-	AddedPeers               []peer.AddrInfo
+	addedPeersMutex          sync.RWMutex
+	_addedPeers              []peer.AddrInfo
 	UnsafeNoTrustIP          bool
 	MetricsRefreshTime       time.Duration
 	metricsCollectionStarted bool
@@ -54,8 +57,6 @@ type app struct {
 
 type subscription struct {
 	Sub    *pubsub.Subscription
-	Idx    uint64
-	Ctx    context.Context
 	Cancel context.CancelFunc
 }
 

src/app/libp2p_helper/src/libp2p_helper/incoming_msg.go

@@ -40,6 +40,7 @@ var pushMesssageExtractors = map[ipc.Libp2pHelperInterface_PushMessage_Which]ext
 	ipc.Libp2pHelperInterface_PushMessage_Which_heartbeatPeer:    fromHeartbeatPeerPush,
 }
 
+// Handles messages coming from the OCaml process
 func (app *app) handleIncomingMsg(msg *ipc.Libp2pHelperInterface_Message) {
 	if msg.HasRpcRequest() {
 		resp, err := func() (*capnp.Message, error) {

src/app/libp2p_helper/src/libp2p_helper/peer_msg.go

@@ -40,7 +40,7 @@ func (m AddPeerReq) handle(app *app, seqno uint64) *capnp.Message {
 		return mkRpcRespError(seqno, badRPC(err))
 	}
 
-	app.AddedPeers = append(app.AddedPeers, *info)
+	app.AddPeers(*info)
 	app.P2p.GatingState().TrustPeer(info.ID)
 
 	if app.Bootstrapper != nil {
@@ -50,7 +50,7 @@ func (m AddPeerReq) handle(app *app, seqno uint64) *capnp.Message {
 	app.P2p.Logger.Info("addPeer Trying to connect to: ", info)
 
 	if AddPeerReqT(m).IsSeed() {
-		app.P2p.Seeds = append(app.P2p.Seeds, *info)
+		app.P2p.AddSeeds(*info)
 	}
 
 	err = app.P2p.Host.Connect(app.Ctx, *info)