pkg/dtd_impl/lib/src/service/connected_app_service.dart

// Copyright (c) 2025, the Dart project authors. Please see the AUTHORS file
// for details. All rights reserved. Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE file.

import 'dart:async';

// ignore: implementation_imports, acceptable since dtd_impl is not published.
import 'package:dds/src/utils/mutex.dart';
import 'package:dtd/dtd.dart';
import 'package:json_rpc_2/json_rpc_2.dart';
import 'package:vm_service/vm_service.dart' hide Parameter, Success;
import 'package:vm_service/vm_service_io.dart';

import '../dtd_client.dart';
import 'internal_service.dart';

typedef _VmServiceUpdate = ({
  /// The metadata describing the VM service connection that this update is for.
  VmServiceInfo vmServiceInfo,

  /// The update event kind.
  _VmServiceUpdateKind kind,
});

typedef _VmServiceWithInfo = ({VmService vmService, VmServiceInfo info});

enum _VmServiceUpdateKind {
  registered(ConnectedAppServiceConstants.vmServiceRegistered),
  unregistered(ConnectedAppServiceConstants.vmServiceUnregistered);

  const _VmServiceUpdateKind(this.id);

  final String id;
}

/// A service that stores the connections to Dart and Flutter applications that
/// DTD is aware of.
///
/// This service allows priveleged clients like an IDE or DDS (the client that
/// started DTD) to register and unregister VM service URIs. The service
/// provides notifications on the [ConnectedAppServiceConstants.serviceName]
/// stream when VM service instances are registered or unregistered.
///
/// DTD clients can use this service to gain access to the VM service URIs for
/// all running applications in the context of this DTD instance.
class ConnectedAppService extends InternalService {
  ConnectedAppService({required this.secret, required this.unrestrictedMode});

  /// The secret that a client must provide to call protected service methods
  /// like 'registerVmService'.
  ///
  /// This secret is generated by DTD at startup and provided to the spawner
  /// of DTD so that only trusted clients can call protected methods.
  final String secret;

  /// Whether the connected app service is unrestricted, meaning that normally
  /// protected service methods like 'registerVmService' do not require a client
  /// secret to be called.
  final bool unrestrictedMode;

  /// This [Mutex] is used to protect the global state of [_vmServices] so that
  /// any read and write operations with asynchronous gaps are executed in
  /// sequence.
  ///
  /// Any operation that updates the contents of [_vmServices] should be guarded
  /// with `_mutex.runGuarded` and any operation that reads the contents of
  /// [_vmServices] should be guarded with `_mutex.runGuardedWeak`.
  final _mutex = Mutex();

  @override
  String get serviceName => ConnectedAppServiceConstants.serviceName;

  String get _streamId => serviceName;

  @override
  void register(DTDClient client) {
    client
      ..registerServiceMethod(
        serviceName,
        ConnectedAppServiceConstants.registerVmService,
        _registerVmService,
      )
      ..registerServiceMethod(
        serviceName,
        ConnectedAppServiceConstants.unregisterVmService,
        _unregisterVmService,
      )
      ..registerServiceMethod(
        serviceName,
        ConnectedAppServiceConstants.getVmServices,
        _getVmServices,
      );

    _vmServiceUpdatesSubscription = _vmServiceUpdates.stream.listen((update) {
      client.streamNotify(_streamId, {
        DtdParameters.streamId: _streamId,
        DtdParameters.eventKind: update.kind.id,
        DtdParameters.eventData: {
          DtdParameters.uri: update.vmServiceInfo.uri,
          DtdParameters.exposedUri: update.vmServiceInfo.exposedUri,
          DtdParameters.name: update.vmServiceInfo.name,
        },
        DtdParameters.timestamp: DateTime.now().millisecondsSinceEpoch,
      });
    });
  }

  @override
  Future<void> shutdown() async {
    await _vmServiceUpdatesSubscription?.cancel();
    await _mutex.runGuarded(() async {
      await Future.wait(_vmServices.values.map((e) => e.vmService.dispose()));
      _vmServices.clear();
    });
  }

  /// The total set of VM service connections DTD is aware of, stored by their
  /// URI as a String.
  final _vmServices = <String, _VmServiceWithInfo>{};

  /// A [StreamController] used internally in this service to track service
  /// registered and unregistered events.
  final _vmServiceUpdates = StreamController<_VmServiceUpdate>.broadcast();

  /// A [StreamSubscription] to the [_vmServiceUpdates] stream that is
  /// responsible for passing VM service registration and unregistration updates
  /// to the [DTDClient] that this [ConnectedAppService] was registered for.
  ///
  /// This [StreamSubscription] must be cancelled in `shutdown`.
  StreamSubscription<_VmServiceUpdate>? _vmServiceUpdatesSubscription;

  /// Registers a VM service URI with the connected app service.
  ///
  /// Only the client that started DTD (identified by [_clientSecret])
  /// should be able to call this method.
  Future<Map<String, Object?>> _registerVmService(Parameters parameters) async {
    final incomingSecret = parameters[DtdParameters.secret].asString;
    if (!unrestrictedMode && secret != incomingSecret) {
      throw RpcErrorCodes.buildRpcException(
        RpcErrorCodes.kPermissionDenied,
      );
    }

    return await _mutex.runGuarded(() async {
      final uri = parameters[DtdParameters.uri].asString;
      if (_vmServices.containsKey(uri)) {
        // We already know about this VM service instance. Exit early.
        return Success().toJson();
      }

      final exposedUri = parameters[DtdParameters.exposedUri].asStringOrNull;
      final name = parameters[DtdParameters.name].asStringOrNull;

      try {
        await vmServiceConnectUri(uri).then((vmService) async {
          final info = VmServiceInfo(
            uri: uri,
            exposedUri: exposedUri,
            name: name,
          );
          _vmServices[uri] = (vmService: vmService, info: info);
          _vmServiceUpdates.add(
            (
              vmServiceInfo: info,
              kind: _VmServiceUpdateKind.registered,
            ),
          );
          unawaited(vmService.onDone.then((_) => _removeServiceAndNotify(uri)));
        });
      } catch (e) {
        throw RpcErrorCodes.buildRpcException(
          RpcErrorCodes.kConnectionFailed,
          data: {'message': 'Error connecting to VM service at $uri.\n$e'},
        );
      }
      return Success().toJson();
    });
  }

  /// Unregisters a VM service URI from the connected app service.
  ///
  /// Only the client that started DTD (identified by [_clientSecret])
  /// should be able to call this method.
  Future<Map<String, Object?>> _unregisterVmService(
    Parameters parameters,
  ) async {
    final incomingSecret = parameters[DtdParameters.secret].asString;
    if (!unrestrictedMode && secret != incomingSecret) {
      throw RpcErrorCodes.buildRpcException(
        RpcErrorCodes.kPermissionDenied,
      );
    }

    return await _mutex.runGuarded(() {
      final uri = parameters[DtdParameters.uri].asString;
      if (!_vmServices.containsKey(uri)) {
        // This VM service is not in the registry. Exit early.
        return Success().toJson();
      }

      _removeServiceAndNotify(uri);
      return Success().toJson();
    });
  }

  /// Removes the VM service with [uri] from [_vmServices] and posts an update
  /// to the 'ConnectedApp' stream.
  ///
  /// This method should be called from within a `_mutex.runGuarded` block since
  /// it performs operations on the [_vmServices] Map, which may be undergoing
  /// updates from other asynchronous blocks.
  void _removeServiceAndNotify(String uri) {
    final removedService = _vmServices.remove(uri);
    // Only send a notification if the service has not already been removed.
    if (removedService != null) {
      _vmServiceUpdates.add(
        (
          vmServiceInfo: removedService.info,
          kind: _VmServiceUpdateKind.unregistered
        ),
      );
    }
  }

  /// Returns a response containing information for each VM service connection
  /// in the context of this DTD instance.
  Future<Map<String, Object?>> _getVmServices(Parameters _) async {
    return await _mutex.runGuardedWeak(() {
      return VmServicesResponse(
        vmServicesInfos: _vmServices.values
            .map((vmServiceWithInfo) => vmServiceWithInfo.info)
            .toList(),
      ).toJson();
    });
  }
}

extension on Parameter {
  String? get asStringOrNull => exists ? asString : null;
}