From b195772160ee792fd5ef81eb897c8cb746ab8ba8 Mon Sep 17 00:00:00 2001
From: Chris Roberts <croberts@hashicorp.com>
Date: Wed, 17 Nov 2021 14:32:25 -0800
Subject: [PATCH] Update capability platform service and host/guest usages

    Update default arguments to be a map instead of an array so
    expected result values can be provided. Add new method for
    handling arguments so custom adjustments can be applied before
    calling the capability. Add custom argument adjustments to
    guest service to provide machine if one is not given.
---
 .../service/capability_platform_service.rb    | 21 +++++++++++++-----
 .../commands/serve/service/guest_service.rb   | 22 ++++++++++++-------
 .../commands/serve/service/host_service.rb    |  7 +++---
 3 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/plugins/commands/serve/service/capability_platform_service.rb b/plugins/commands/serve/service/capability_platform_service.rb
index ad7f626ee..286717621 100644
--- a/plugins/commands/serve/service/capability_platform_service.rb
+++ b/plugins/commands/serve/service/capability_platform_service.rb
@@ -24,12 +24,12 @@ module VagrantPlugins
 
         # TODO(spox): request scoping for seed values - needs investigation
         def seed(req, ctx)
-          @seeds = req.list.to_a
+          @seeds = req.arguments.to_a
           Empty.new
         end
 
         def seeds(req, ctx)
-          SDK::Args::Direct.new(list: @seeds)
+          SDK::Args::Direct.new(arguments: @seeds)
         end
 
         def has_capability_spec(*_)
@@ -68,7 +68,7 @@ module VagrantPlugins
         def capability_spec(req, ctx)
           SDK::FuncSpec.new(
             name: "capability_spec",
-            args: default_args + [
+            args: default_args.values + [
               SDK::FuncSpec::Value.new(
                 type: "hashicorp.vagrant.sdk.Args.Direct",
                 name: "",
@@ -93,10 +93,16 @@ module VagrantPlugins
             caps_registry = capabilities[plugin_name]
             target_cap = caps_registry.get(cap_name)
 
-            args = mapper.funcspec_map(req.func_args, mapper, broker)
-            args = [args.first] + args.last
+            args = mapper.funcspec_map(
+              req.func_args,
+              expect: default_args.keys + [Types::Direct]
+            )
+            args = capability_arguments(args)
             cap_method = target_cap.method(cap_name)
 
+            arg_list = args.join("\n  - ")
+            logger.debug("arguments to be passed to #{cap_name} on plugin #{plugin_name}:\n  - #{arg_list}")
+
             result = cap_method.call(*args)
 
             val = Google::Protobuf::Value.new
@@ -106,6 +112,11 @@ module VagrantPlugins
             )
           end
         end
+
+        def capability_arguments(args)
+          direct = args.pop
+          args + direct.arguments
+        end
       end
     end
   end
diff --git a/plugins/commands/serve/service/guest_service.rb b/plugins/commands/serve/service/guest_service.rb
index c087d1c10..3bce7980b 100644
--- a/plugins/commands/serve/service/guest_service.rb
+++ b/plugins/commands/serve/service/guest_service.rb
@@ -9,13 +9,12 @@ module VagrantPlugins
 
         def initialize(*args, **opts, &block)
           caps = Vagrant.plugin("2").manager.guest_capabilities
-          default_args = [
-            # Always get a target to pass the guest capability
-            SDK::FuncSpec::Value.new(
+          default_args = {
+            Client::Target => SDK::FuncSpec::Value.new(
               type: "hashicorp.vagrant.sdk.Args.Target",
               name: "",
             ),
-          ]
+          }
           initialize_capability_platform!(caps, default_args)
         end
 
@@ -40,10 +39,7 @@ module VagrantPlugins
         def detect(req, ctx)
           with_info(ctx) do |info|
             plugin_name = info.plugin_name
-            target = mapper.funcspec_map(req, expect: Client::Target)
-            project = target.project
-            env = Vagrant::Environment.new({client: project})
-            machine = env.machine(target.name.to_sym, target.provider_name.to_sym)
+            machine = mapper.funcspec_map(req, expect: Vagrant::Machine)
             plugin = Vagrant.plugin("2").manager.guests[plugin_name.to_s.to_sym].to_a.first
             if !plugin
               logger.debug("Failed to locate guest plugin for: #{plugin_name}")
@@ -86,6 +82,16 @@ module VagrantPlugins
             )
           end
         end
+
+        def capability_arguments(args)
+          target, direct = args
+          nargs = direct.args.dup
+          if !nargs.first.is_a?(Vagrant::Machine)
+            nargs.unshift(mapper.map(target, to: Vagrant::Machine))
+          end
+
+          nargs
+        end
       end
     end
   end
diff --git a/plugins/commands/serve/service/host_service.rb b/plugins/commands/serve/service/host_service.rb
index b9b9bd5be..942ab1620 100644
--- a/plugins/commands/serve/service/host_service.rb
+++ b/plugins/commands/serve/service/host_service.rb
@@ -9,13 +9,12 @@ module VagrantPlugins
 
         def initialize(*args, **opts, &block)
           caps = Vagrant.plugin("2").manager.host_capabilities
-          default_args = [
-            # Always get the state bag for host capabilities
-            SDK::FuncSpec::Value.new(
+          default_args = {
+            Vagrant::Environment => SDK::FuncSpec::Value.new(
               type: "hashicorp.vagrant.sdk.Args.Project",
               name: "",
             ),
-          ]
+          }
           initialize_capability_platform!(caps, default_args)
         end