From 3aa9aaba8578dda186772b371d342fe61dfd15b8 Mon Sep 17 00:00:00 2001
From: Doug MacEachern <dougm@vmware.com>
Date: Wed, 6 May 2020 13:21:43 -0700
Subject: [PATCH] vapi: sync access to rest.Client.SessionID

Required for the upcoming KeepAlive support.
---
 govc/library/export.go   |  2 +-
 govc/session/login.go    |  4 ++--
 govc/session/ls.go       |  2 +-
 session/cache/session.go |  6 ++----
 vapi/rest/client.go      | 35 ++++++++++++++++++++++++++++-------
 5 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/govc/library/export.go b/govc/library/export.go
index 2dbfc9b7e..8707d3d1e 100644
--- a/govc/library/export.go
+++ b/govc/library/export.go
@@ -139,7 +139,7 @@ func (cmd *export) Run(ctx context.Context, f *flag.FlagSet) error {
 	download := func(src *url.URL, name string) error {
 		p := soap.DefaultDownload
 		p.Headers = map[string]string{
-			"vmware-api-session-id": c.SessionID,
+			"vmware-api-session-id": c.SessionID(),
 		}
 		if isStdout {
 			s, _, err := c.Download(ctx, src, &p)
diff --git a/govc/session/login.go b/govc/session/login.go
index e89e45d07..79ae20aa5 100644
--- a/govc/session/login.go
+++ b/govc/session/login.go
@@ -263,9 +263,9 @@ func (cmd *login) setCookie(ctx context.Context, c *vim25.Client) error {
 
 func (cmd *login) setRestCookie(ctx context.Context, c *rest.Client) error {
 	if cmd.cookie == "" {
-		cmd.cookie = c.SessionID
+		cmd.cookie = c.SessionID()
 	} else {
-		c.SessionID = cmd.cookie
+		c.SessionID(cmd.cookie)
 
 		// Check the session is still valid
 		s, err := c.Session(ctx)
diff --git a/govc/session/ls.go b/govc/session/ls.go
index 3cd8a9c31..7346a6033 100644
--- a/govc/session/ls.go
+++ b/govc/session/ls.go
@@ -140,7 +140,7 @@ func (cmd *ls) Run(ctx context.Context, f *flag.FlagSet) error {
 				return err
 			}
 			m.SessionList = append(m.SessionList, types.UserSession{
-				Key:            rc.SessionID,
+				Key:            rc.SessionID(),
 				UserName:       rs.User + " (REST)",
 				LoginTime:      rs.Created,
 				LastActiveTime: rs.LastAccessed,
diff --git a/session/cache/session.go b/session/cache/session.go
index cd5a8b198..b4e4136c2 100644
--- a/session/cache/session.go
+++ b/session/cache/session.go
@@ -325,18 +325,16 @@ func (s *Session) Login(ctx context.Context, c Client, config func(*soap.Client)
 		*client = *vc
 		c = client
 	case *rest.Client:
-		vc := &vim25.Client{Client: sc}
-		rc := rest.NewClient(vc)
+		client.Client = sc.NewServiceClient(rest.Path, "")
 
 		login := s.loginREST
 		if s.LoginREST != nil {
 			login = s.LoginREST
 		}
-		if err = login(ctx, rc); err != nil {
+		if err = login(ctx, client); err != nil {
 			return err
 		}
 
-		*client = *rc
 		c = client
 	default:
 		panic(fmt.Sprintf("unsupported client type=%T", client))
diff --git a/vapi/rest/client.go b/vapi/rest/client.go
index a7e58a401..c20429958 100644
--- a/vapi/rest/client.go
+++ b/vapi/rest/client.go
@@ -25,6 +25,7 @@ import (
 	"io/ioutil"
 	"net/http"
 	"net/url"
+	"sync"
 	"time"
 
 	"github.com/vmware/govmomi/vapi/internal"
@@ -34,8 +35,10 @@ import (
 
 // Client extends soap.Client to support JSON encoding, while inheriting security features, debug tracing and session persistence.
 type Client struct {
+	mu sync.Mutex
+
 	*soap.Client
-	SessionID string
+	sessionID string
 }
 
 // Session information
@@ -60,7 +63,17 @@ func (m *LocalizableMessage) Error() string {
 func NewClient(c *vim25.Client) *Client {
 	sc := c.Client.NewServiceClient(Path, "")
 
-	return &Client{sc, ""}
+	return &Client{Client: sc}
+}
+
+// SessionID is set by calling Login() or optionally with the given id param
+func (c *Client) SessionID(id ...string) string {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if len(id) != 0 {
+		c.sessionID = id[0]
+	}
+	return c.sessionID
 }
 
 type marshaledClient struct {
@@ -71,7 +84,7 @@ type marshaledClient struct {
 func (c *Client) MarshalJSON() ([]byte, error) {
 	m := marshaledClient{
 		SoapClient: c.Client,
-		SessionID:  c.SessionID,
+		SessionID:  c.sessionID,
 	}
 
 	return json.Marshal(m)
@@ -87,7 +100,7 @@ func (c *Client) UnmarshalJSON(b []byte) error {
 
 	*c = Client{
 		Client:    m.SoapClient,
-		SessionID: m.SessionID,
+		sessionID: m.SessionID,
 	}
 
 	return nil
@@ -127,8 +140,8 @@ func (c *Client) Do(ctx context.Context, req *http.Request, resBody interface{})
 
 	req.Header.Set("Accept", "application/json")
 
-	if c.SessionID != "" {
-		req.Header.Set(internal.SessionCookieName, c.SessionID)
+	if id := c.SessionID(); id != "" {
+		req.Header.Set(internal.SessionCookieName, id)
 	}
 
 	if s, ok := ctx.Value(signerContext{}).(Signer); ok {
@@ -183,7 +196,15 @@ func (c *Client) Login(ctx context.Context, user *url.Userinfo) error {
 		}
 	}
 
-	return c.Do(ctx, req, &c.SessionID)
+	var id string
+	err := c.Do(ctx, req, &id)
+	if err != nil {
+		return err
+	}
+
+	c.SessionID(id)
+
+	return nil
 }
 
 func (c *Client) LoginByToken(ctx context.Context) error {