diff --git a/management/internals/controllers/network_map/controller/controller.go b/management/internals/controllers/network_map/controller/controller.go index 30de974a1..ff159a671 100644 --- a/management/internals/controllers/network_map/controller/controller.go +++ b/management/internals/controllers/network_map/controller/controller.go @@ -113,14 +113,17 @@ func (c *Controller) OnPeerConnected(ctx context.Context, accountID string, peer return c.peersUpdateManager.CreateChannel(ctx, peerID), nil } -func (c *Controller) OnPeerDisconnected(ctx context.Context, accountID string, peerID string) { - c.peersUpdateManager.CloseChannel(ctx, peerID) +func (c *Controller) OnPeerDisconnected(ctx context.Context, accountID string, peerID string, session chan *network_map.UpdateMessage) bool { + if !c.peersUpdateManager.CloseSessionChannel(ctx, peerID, session) { + return false + } peer, err := c.repo.GetPeerByID(ctx, accountID, peerID) if err != nil { log.WithContext(ctx).Errorf("failed to get peer %s: %v", peerID, err) - return + return true } c.EphemeralPeersManager.OnPeerDisconnected(ctx, peer) + return true } // injectAllProxyPolicies prepares an account for the per-peer network-map diff --git a/management/internals/controllers/network_map/interface.go b/management/internals/controllers/network_map/interface.go index b535321d1..444b64daf 100644 --- a/management/internals/controllers/network_map/interface.go +++ b/management/internals/controllers/network_map/interface.go @@ -35,7 +35,7 @@ type Controller interface { OnPeersDeleted(ctx context.Context, accountID string, peerIDs []string, affectedPeerIDs []string) error DisconnectPeers(ctx context.Context, accountId string, peerIDs []string) OnPeerConnected(ctx context.Context, accountID string, peerID string) (chan *UpdateMessage, error) - OnPeerDisconnected(ctx context.Context, accountID string, peerID string) + OnPeerDisconnected(ctx context.Context, accountID string, peerID string, session chan *UpdateMessage) bool TrackEphemeralPeer(ctx context.Context, peer *nbpeer.Peer) } diff --git a/management/internals/controllers/network_map/interface_mock.go b/management/internals/controllers/network_map/interface_mock.go index 42051f172..d41975052 100644 --- a/management/internals/controllers/network_map/interface_mock.go +++ b/management/internals/controllers/network_map/interface_mock.go @@ -178,15 +178,17 @@ func (mr *MockControllerMockRecorder) OnPeerConnected(ctx, accountID, peerID any } // OnPeerDisconnected mocks base method. -func (m *MockController) OnPeerDisconnected(ctx context.Context, accountID, peerID string) { +func (m *MockController) OnPeerDisconnected(ctx context.Context, accountID, peerID string, session chan *UpdateMessage) bool { m.ctrl.T.Helper() - m.ctrl.Call(m, "OnPeerDisconnected", ctx, accountID, peerID) + ret := m.ctrl.Call(m, "OnPeerDisconnected", ctx, accountID, peerID, session) + ret0, _ := ret[0].(bool) + return ret0 } // OnPeerDisconnected indicates an expected call of OnPeerDisconnected. -func (mr *MockControllerMockRecorder) OnPeerDisconnected(ctx, accountID, peerID any) *gomock.Call { +func (mr *MockControllerMockRecorder) OnPeerDisconnected(ctx, accountID, peerID, session any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnPeerDisconnected", reflect.TypeOf((*MockController)(nil).OnPeerDisconnected), ctx, accountID, peerID) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "OnPeerDisconnected", reflect.TypeOf((*MockController)(nil).OnPeerDisconnected), ctx, accountID, peerID, session) } // OnPeersAdded mocks base method. diff --git a/management/internals/controllers/network_map/update_channel.go b/management/internals/controllers/network_map/update_channel.go index 0b085b85f..9305f95f5 100644 --- a/management/internals/controllers/network_map/update_channel.go +++ b/management/internals/controllers/network_map/update_channel.go @@ -6,6 +6,7 @@ type PeersUpdateManager interface { SendUpdate(ctx context.Context, peerID string, update *UpdateMessage) CreateChannel(ctx context.Context, peerID string) chan *UpdateMessage CloseChannel(ctx context.Context, peerID string) + CloseSessionChannel(ctx context.Context, peerID string, session chan *UpdateMessage) bool CountStreams() int HasChannel(peerID string) bool CloseChannels(ctx context.Context, peerIDs []string) diff --git a/management/internals/controllers/network_map/update_channel/updatechannel.go b/management/internals/controllers/network_map/update_channel/updatechannel.go index 91627bf15..fe2f45f77 100644 --- a/management/internals/controllers/network_map/update_channel/updatechannel.go +++ b/management/internals/controllers/network_map/update_channel/updatechannel.go @@ -133,6 +133,30 @@ func (p *PeersUpdateManager) CloseChannel(ctx context.Context, peerID string) { p.closeChannel(ctx, peerID) } +// CloseSessionChannel closes the peer's updates channel only while it is still the given +// session's channel, and reports whether that session still owned the peer's stream. +// A registered channel other than session means a newer stream replaced the caller's: +// the stale session must leave the peer's resources to the new owner. +func (p *PeersUpdateManager) CloseSessionChannel(ctx context.Context, peerID string, session chan *network_map.UpdateMessage) bool { + start := time.Now() + + p.channelsMux.Lock() + defer func() { + p.channelsMux.Unlock() + if p.metrics != nil { + p.metrics.UpdateChannelMetrics().CountCloseChannelDuration(time.Since(start)) + } + }() + + if channel, ok := p.peerChannels[peerID]; ok && channel != session { + log.WithContext(ctx).Debugf("skipped closing updates channel: peer %s reconnected with a newer session", peerID) + return false + } + + p.closeChannel(ctx, peerID) + return true +} + // GetAllConnectedPeers returns a copy of the connected peers map func (p *PeersUpdateManager) GetAllConnectedPeers() map[string]struct{} { start := time.Now() diff --git a/management/internals/shared/grpc/server.go b/management/internals/shared/grpc/server.go index 3d5f0a1b7..96f2440e2 100644 --- a/management/internals/shared/grpc/server.go +++ b/management/internals/shared/grpc/server.go @@ -322,7 +322,7 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S if err != nil { log.WithContext(ctx).Debugf("error while sending initial sync for %s: %v", peerKey.String(), err) s.syncSem.Add(-1) - s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart) + s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, nil) return err } @@ -330,7 +330,7 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S if err != nil { log.WithContext(ctx).Debugf("error while notify peer connected for %s: %v", peerKey.String(), err) s.syncSem.Add(-1) - s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart) + s.cancelPeerRoutinesWithoutLock(ctx, accountID, peer, syncStart, nil) return err } @@ -391,7 +391,7 @@ func (s *Server) startResponseReceiver(ctx context.Context, srv proto.Management func (s *Server) sendJobsLoop(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, updates *job.Channel, srv proto.ManagementService_JobServer) error { // todo figure out better error handling strategy - defer s.jobManager.CloseChannel(ctx, accountID, peer.ID) + defer s.jobManager.CloseChannel(ctx, accountID, peer.ID, updates) for { event, err := updates.Event(ctx) @@ -435,14 +435,14 @@ func (s *Server) handleUpdates(ctx context.Context, accountID string, peerKey wg if !open { log.WithContext(ctx).Debugf("updates channel for peer %s was closed", peerKey.String()) - s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime) + s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime, updates) return nil } log.WithContext(ctx).Tracef("received an update for peer %s", peerKey.String()) if debouncer.ProcessUpdate(update) { // Send immediately (first update or after quiet period) - if err := s.sendUpdate(ctx, accountID, peerKey, peer, update, srv, streamStartTime); err != nil { + if err := s.sendUpdate(ctx, accountID, peerKey, peer, update, srv, streamStartTime, updates); err != nil { log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", peerKey.String(), err) return err } @@ -456,7 +456,7 @@ func (s *Server) handleUpdates(ctx context.Context, accountID string, peerKey wg } log.WithContext(ctx).Debugf("sending %d debounced update(s) for peer %s", len(pendingUpdates), peerKey.String()) for _, pendingUpdate := range pendingUpdates { - if err := s.sendUpdate(ctx, accountID, peerKey, peer, pendingUpdate, srv, streamStartTime); err != nil { + if err := s.sendUpdate(ctx, accountID, peerKey, peer, pendingUpdate, srv, streamStartTime, updates); err != nil { log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", peerKey.String(), err) return err } @@ -466,7 +466,7 @@ func (s *Server) handleUpdates(ctx context.Context, accountID string, peerKey wg case <-srv.Context().Done(): // happens when connection drops, e.g. client disconnects log.WithContext(ctx).Debugf("stream of peer %s has been closed", peerKey.String()) - s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime) + s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime, updates) return srv.Context().Err() } } @@ -474,16 +474,16 @@ func (s *Server) handleUpdates(ctx context.Context, accountID string, peerKey wg // sendUpdate encrypts the update message using the peer key and the server's wireguard key, // then sends the encrypted message to the connected peer via the sync server. -func (s *Server) sendUpdate(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, update *network_map.UpdateMessage, srv proto.ManagementService_SyncServer, streamStartTime time.Time) error { +func (s *Server) sendUpdate(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, update *network_map.UpdateMessage, srv proto.ManagementService_SyncServer, streamStartTime time.Time, session chan *network_map.UpdateMessage) error { key, err := s.secretsManager.GetWGKey() if err != nil { - s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime) + s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime, session) return status.Errorf(codes.Internal, "failed processing update message") } encryptedResp, err := encryption.EncryptMessage(peerKey, key, update.Update) if err != nil { - s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime) + s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime, session) return status.Errorf(codes.Internal, "failed processing update message") } err = srv.Send(&proto.EncryptedMessage{ @@ -491,7 +491,7 @@ func (s *Server) sendUpdate(ctx context.Context, accountID string, peerKey wgtyp Body: encryptedResp, }) if err != nil { - s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime) + s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime, session) return status.Errorf(codes.Internal, "failed sending update message") } log.WithContext(ctx).Tracef("sent an update to peer %s", peerKey.String()) @@ -523,20 +523,23 @@ func (s *Server) sendJob(ctx context.Context, peerKey wgtypes.Key, job *job.Even return nil } -func (s *Server) cancelPeerRoutines(ctx context.Context, accountID string, peer *nbpeer.Peer, streamStartTime time.Time) { +func (s *Server) cancelPeerRoutines(ctx context.Context, accountID string, peer *nbpeer.Peer, streamStartTime time.Time, session chan *network_map.UpdateMessage) { uncanceledCTX := context.WithoutCancel(ctx) unlock := s.acquirePeerLockByUID(uncanceledCTX, peer.Key) defer unlock() - s.cancelPeerRoutinesWithoutLock(uncanceledCTX, accountID, peer, streamStartTime) + s.cancelPeerRoutinesWithoutLock(uncanceledCTX, accountID, peer, streamStartTime, session) } -func (s *Server) cancelPeerRoutinesWithoutLock(ctx context.Context, accountID string, peer *nbpeer.Peer, streamStartTime time.Time) { +func (s *Server) cancelPeerRoutinesWithoutLock(ctx context.Context, accountID string, peer *nbpeer.Peer, streamStartTime time.Time, session chan *network_map.UpdateMessage) { + if !s.networkMapController.OnPeerDisconnected(ctx, accountID, peer.ID, session) { + log.WithContext(ctx).Debugf("skipped peer routines teardown for %s: a newer session owns the peer", peer.Key) + return + } err := s.accountManager.OnPeerDisconnected(ctx, accountID, peer.Key, streamStartTime) if err != nil { log.WithContext(ctx).Errorf("failed to disconnect peer %s properly: %v", peer.Key, err) } - s.networkMapController.OnPeerDisconnected(ctx, accountID, peer.ID) s.secretsManager.CancelRefresh(peer.ID) log.WithContext(ctx).Debugf("peer %s has been disconnected", peer.Key) diff --git a/management/server/job/manager.go b/management/server/job/manager.go index 0b183ac39..adc63eeb3 100644 --- a/management/server/job/manager.go +++ b/management/server/job/manager.go @@ -127,12 +127,17 @@ func (jm *Manager) HandleResponse(ctx context.Context, resp *proto.JobResponse, return nil } -// CloseChannel closes a peer’s channel and cleans up its jobs -func (jm *Manager) CloseChannel(ctx context.Context, accountID, peerID string) { +// CloseChannel closes a peer’s channel and cleans up its jobs. A registered +// channel other than session belongs to a newer job stream: the stale caller +// must not close it or fail the jobs the new stream is serving. +func (jm *Manager) CloseChannel(ctx context.Context, accountID, peerID string, session *Channel) { jm.mu.Lock() defer jm.mu.Unlock() if ch, ok := jm.jobChannels[peerID]; ok { + if ch != session { + return + } ch.Close() delete(jm.jobChannels, peerID) }