From 71ce491197fe57e704b980fa71be81a00c6f4210 Mon Sep 17 00:00:00 2001
From: Chris Roberts <croberts@hashicorp.com>
Date: Fri, 1 Oct 2021 12:10:54 -0700
Subject: [PATCH] Allow name lookup overrides for proto integration

---
 plugins/commands/serve/util/client_setup.rb | 14 +++++++++++---
 plugins/commands/serve/util/connector.rb    | 14 +++++++++++---
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/plugins/commands/serve/util/client_setup.rb b/plugins/commands/serve/util/client_setup.rb
index 5439ebc88..bf82775e4 100644
--- a/plugins/commands/serve/util/client_setup.rb
+++ b/plugins/commands/serve/util/client_setup.rb
@@ -10,14 +10,22 @@ module VagrantPlugins
         end
 
         def initialize(conn, proto, broker=nil)
-          srv = self.class.name.split('::').last
+          n = self.class.respond_to?(:sdk_alias) ? self.class.sdk_alias : self.class.name
+          lookup = n.split("::")
+          idx = lookup.index("Client")
+          if idx
+            lookup.slice!(0, idx+1)
+          end
+
+          srv = "#{lookup.join}Service"
+
           logger.debug("connecting to #{srv.downcase} service on #{conn}")
           @broker = broker
           @proto = proto
-          srv_klass = SDK.const_get("#{srv}Service")&.const_get(:Stub)
+          srv_klass = SDK.const_get(srv)&.const_get(:Stub)
           if !srv_klass
             raise NameError,
-              "failed to locate required protobuf constant `SDK::#{srv}Service'"
+              "failed to locate required protobuf constant `SDK::#{srv}'"
           end
           @client = srv_klass.new(conn, :this_channel_is_insecure)
         end
diff --git a/plugins/commands/serve/util/connector.rb b/plugins/commands/serve/util/connector.rb
index 732267315..f04ec1123 100644
--- a/plugins/commands/serve/util/connector.rb
+++ b/plugins/commands/serve/util/connector.rb
@@ -17,11 +17,19 @@ module VagrantPlugins
 
         def load(raw, broker:)
           if raw.is_a?(String)
-            srv = self.class.name.split('::').last
-            klass = SDK::Args.const_get(srv)
+            n = self.respond_to?(:sdk_alias) ? self.sdk_alias : self.name
+            lookup = n.split("::")
+            idx = lookup.index("Client")
+            if idx
+              lookup.slice!(0, idx+1)
+            end
+
+            klass = lookup.inject(SDK::Args) do |const, name|
+              const&.const_get(name)
+            end
             if !klass
               raise NameError,
-                "failed to locate required protobuf constant `SDK::Args::#{srv}'"
+                "failed to locate required protobuf constant `SDK::Args::#{n}'\n\nArgs: #{SDK::Args.constants.inspect}"
             end
             raw = klass.decode(raw)
           end