From f8084641b042092eae82404fb5e862f31c3d1662 Mon Sep 17 00:00:00 2001
From: Doug MacEachern <dougm@vmware.com>
Date: Mon, 18 Mar 2024 19:47:15 -0700
Subject: [PATCH] vcsim: add ExtensionManager.FindExtension method

d9af2a2 added simulator.ExtensionManager but did not included this method.

govc: use ExtensionManager.FindExtension when given a single KEY arg,
rather than fetch the entire list.

Closes #2644
---
 govc/extension/info.go         | 41 +++++++++++++++++++++-------------
 govc/test/extension.bats       | 13 ++++++-----
 simulator/extension_manager.go | 15 +++++++++++++
 3 files changed, 48 insertions(+), 21 deletions(-)

diff --git a/govc/extension/info.go b/govc/extension/info.go
index de88b43dd..c0819eed3 100644
--- a/govc/extension/info.go
+++ b/govc/extension/info.go
@@ -72,30 +72,39 @@ func (cmd *info) Run(ctx context.Context, f *flag.FlagSet) error {
 		return err
 	}
 
-	list, err := m.List(ctx)
-	if err != nil {
-		return err
-	}
-
 	var res infoResult
+	exts := make(map[string]types.Extension)
 
-	if f.NArg() == 0 {
-		res.Extensions = list
+	if f.NArg() == 1 {
+		e, err := m.Find(ctx, f.Arg(0))
+		if err != nil {
+			return err
+		}
+		if e != nil {
+			exts[f.Arg(0)] = *e
+		}
 	} else {
-		exts := make(map[string]types.Extension)
-		for _, e := range list {
-			exts[e.Key] = e
+		list, err := m.List(ctx)
+		if err != nil {
+			return err
 		}
-
-		for _, key := range f.Args() {
-			if e, ok := exts[key]; ok {
-				res.Extensions = append(res.Extensions, e)
-			} else {
-				return fmt.Errorf("extension %s not found", key)
+		if f.NArg() == 0 {
+			res.Extensions = list
+		} else {
+			for _, e := range list {
+				exts[e.Key] = e
 			}
 		}
 	}
 
+	for _, key := range f.Args() {
+		if e, ok := exts[key]; ok {
+			res.Extensions = append(res.Extensions, e)
+		} else {
+			return fmt.Errorf("extension %s not found", key)
+		}
+	}
+
 	return cmd.WriteResult(&res)
 }
 
diff --git a/govc/test/extension.bats b/govc/test/extension.bats
index 0c20b8198..8612e8319 100755
--- a/govc/test/extension.bats
+++ b/govc/test/extension.bats
@@ -8,10 +8,10 @@ load test_helper
   run govc extension.info enoent
   assert_failure
 
-  id=$(new_id)
+  run govc extension.info
+  assert_success
 
-  result=$(govc extension.info | grep $id | wc -l)
-  [ $result -eq 0 ]
+  id=$(new_id)
 
   # register extension
   run govc extension.register $id <<EOS
@@ -54,9 +54,12 @@ EOS
   # remove generated cert and key
   rm ${id}.{crt,key}
 
+  run govc extension.info $(govc extension.info -json | jq -r .extensions[].key)
+  assert_success
+
   run govc extension.unregister $id
   assert_success
 
-  result=$(govc extension.info | grep $id | wc -l)
-  [ $result -eq 0 ]
+  run govc extension.info $id
+  assert_failure
 }
diff --git a/simulator/extension_manager.go b/simulator/extension_manager.go
index 60a46c418..6b72ae0fd 100644
--- a/simulator/extension_manager.go
+++ b/simulator/extension_manager.go
@@ -67,6 +67,21 @@ func (m *ExtensionManager) init(r *Registry) {
 	}
 }
 
+func (m *ExtensionManager) FindExtension(ctx *Context, req *types.FindExtension) soap.HasFault {
+	body := &methods.FindExtensionBody{
+		Res: new(types.FindExtensionResponse),
+	}
+
+	for _, x := range m.ExtensionList {
+		if x.Key == req.ExtensionKey {
+			body.Res.Returnval = &x
+			break
+		}
+	}
+
+	return body
+}
+
 func (m *ExtensionManager) RegisterExtension(ctx *Context, req *types.RegisterExtension) soap.HasFault {
 	body := &methods.RegisterExtensionBody{}