github.com/vmware/govmomi@v0.51.0/vapi/rest/client.go

// © Broadcom. All Rights Reserved.
// The term “Broadcom” refers to Broadcom Inc. and/or its subsidiaries.
// SPDX-License-Identifier: Apache-2.0

package rest

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"mime"
	"net/http"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/vmware/govmomi/vapi/internal"
	"github.com/vmware/govmomi/vim25"
	"github.com/vmware/govmomi/vim25/soap"
	"github.com/vmware/govmomi/vim25/types"
)

// 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
}

// Session information
type Session struct {
	User         string    `json:"user"`
	Created      time.Time `json:"created_time"`
	LastAccessed time.Time `json:"last_accessed_time"`
}

// LocalizableMessage represents a localizable error
type LocalizableMessage struct {
	Args           []string `json:"args,omitempty"`
	DefaultMessage string   `json:"default_message,omitempty"`
	ID             string   `json:"id,omitempty"`
}

func (m *LocalizableMessage) Error() string {
	return m.DefaultMessage
}

// NewClient creates a new Client instance.
func NewClient(c *vim25.Client) *Client {
	sc := c.Client.NewServiceClient(Path, "")

	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 {
	SoapClient *soap.Client
	SessionID  string
}

func (c *Client) MarshalJSON() ([]byte, error) {
	m := marshaledClient{
		SoapClient: c.Client,
		SessionID:  c.sessionID,
	}

	return json.Marshal(m)
}

func (c *Client) UnmarshalJSON(b []byte) error {
	var m marshaledClient

	err := json.Unmarshal(b, &m)
	if err != nil {
		return err
	}

	*c = Client{
		Client:    m.SoapClient,
		sessionID: m.SessionID,
	}

	return nil
}

// isAPI returns true if path starts with "/api"
// This hack allows helpers to support both endpoints:
// "/rest" - value wrapped responses and structured error responses
// "/api" - raw responses and no structured error responses
func isAPI(path string) bool {
	return strings.HasPrefix(path, "/api")
}

// Resource helper for the given path.
func (c *Client) Resource(path string) *Resource {
	r := &Resource{u: c.URL()}
	if !isAPI(path) {
		path = Path + path
	}
	r.u.Path = path
	return r
}

type Signer interface {
	SignRequest(*http.Request) error
}

type signerContext struct{}

func (c *Client) WithSigner(ctx context.Context, s Signer) context.Context {
	return context.WithValue(ctx, signerContext{}, s)
}

type headersContext struct{}

// WithHeader returns a new Context populated with the provided headers map.
// Calls to a VAPI REST client with this context will populate the HTTP headers
// map using the provided headers.
func (c *Client) WithHeader(
	ctx context.Context,
	headers http.Header) context.Context {

	return context.WithValue(ctx, headersContext{}, headers)
}

type statusError struct {
	res *http.Response
}

func (e *statusError) Error() string {
	return fmt.Sprintf("%s %s: %s", e.res.Request.Method, e.res.Request.URL, e.res.Status)
}

func IsStatusError(err error, code int) bool {
	statusErr, ok := err.(*statusError)
	if !ok || statusErr == nil || statusErr.res == nil {
		return false
	}
	return statusErr.res.StatusCode == code
}

// RawResponse may be used with the Do method as the resBody argument in order
// to capture the raw response data.
type RawResponse struct {
	bytes.Buffer
}

// Do sends the http.Request, decoding resBody if provided.
func (c *Client) Do(ctx context.Context, req *http.Request, resBody any) error {
	switch req.Method {
	case http.MethodPost, http.MethodPatch, http.MethodPut:
		req.Header.Set("Content-Type", "application/json")
	}

	req.Header.Set("Accept", "application/json")

	if id := c.SessionID(); id != "" {
		req.Header.Set(internal.SessionCookieName, id)
	}

	if s, ok := ctx.Value(signerContext{}).(Signer); ok {
		if err := s.SignRequest(req); err != nil {
			return err
		}
	}

	// OperationID (see soap.Client.soapRoundTrip)
	if id, ok := ctx.Value(types.ID{}).(string); ok {
		req.Header.Add("X-Request-ID", id)
	}

	if headers, ok := ctx.Value(headersContext{}).(http.Header); ok {
		for k, v := range headers {
			for _, v := range v {
				req.Header.Add(k, v)
			}
		}
	}

	return c.Client.Do(ctx, req, func(res *http.Response) error {
		switch res.StatusCode {
		case http.StatusOK:
		case http.StatusCreated:
		case http.StatusAccepted:
		case http.StatusNoContent:
		case http.StatusBadRequest:
			// TODO: structured error types
			detail, err := io.ReadAll(res.Body)
			if err != nil {
				return err
			}
			return fmt.Errorf("%s: %s", res.Status, bytes.TrimSpace(detail))
		default:
			return &statusError{res}
		}

		if resBody == nil {
			return nil
		}

		switch b := resBody.(type) {
		case *RawResponse:
			return res.Write(b)
		case io.Writer:
			_, err := io.Copy(b, res.Body)
			return err
		default:
			d := json.NewDecoder(res.Body)
			if isAPI(req.URL.Path) {
				// Responses from the /api endpoint are not wrapped
				return d.Decode(resBody)
			}
			// Responses from the /rest endpoint are wrapped in this structure
			val := struct {
				Value any `json:"value,omitempty"`
			}{
				resBody,
			}
			return d.Decode(&val)
		}
	})
}

// authHeaders ensures the given map contains a REST auth header
func (c *Client) authHeaders(h map[string]string) map[string]string {
	if _, exists := h[internal.SessionCookieName]; exists {
		return h
	}
	if h == nil {
		h = make(map[string]string)
	}

	h[internal.SessionCookieName] = c.SessionID()

	return h
}

// Download wraps soap.Client.Download, adding the REST authentication header
func (c *Client) Download(ctx context.Context, u *url.URL, param *soap.Download) (io.ReadCloser, int64, error) {
	p := *param
	p.Headers = c.authHeaders(p.Headers)
	return c.Client.Download(ctx, u, &p)
}

// DownloadFile wraps soap.Client.DownloadFile, adding the REST authentication header
func (c *Client) DownloadFile(ctx context.Context, file string, u *url.URL, param *soap.Download) error {
	p := *param
	p.Headers = c.authHeaders(p.Headers)
	return c.Client.DownloadFile(ctx, file, u, &p)
}

// DownloadAttachment writes the response to given filename, defaulting to Content-Disposition filename in the response.
// A filename of "-" writes the response to stdout.
func (c *Client) DownloadAttachment(ctx context.Context, req *http.Request, filename string) error {
	return c.Client.Do(ctx, req, func(res *http.Response) error {
		if filename == "" {
			d := res.Header.Get("Content-Disposition")
			_, params, err := mime.ParseMediaType(d)
			if err == nil {
				filename = params["filename"]
			}
		}

		var w io.Writer

		if filename == "-" {
			w = os.Stdout
		} else {
			f, err := os.Create(filename)
			if err != nil {
				return err
			}
			defer f.Close()
			w = f
		}

		_, err := io.Copy(w, res.Body)
		return err
	})
}

// Upload wraps soap.Client.Upload, adding the REST authentication header
func (c *Client) Upload(ctx context.Context, f io.Reader, u *url.URL, param *soap.Upload) error {
	p := *param
	p.Headers = c.authHeaders(p.Headers)
	return c.Client.Upload(ctx, f, u, &p)
}

// Login creates a new session via Basic Authentication with the given url.Userinfo.
func (c *Client) Login(ctx context.Context, user *url.Userinfo) error {
	req := c.Resource(internal.SessionPath).Request(http.MethodPost)

	req.Header.Set(internal.UseHeaderAuthn, "true")

	if user != nil {
		if password, ok := user.Password(); ok {
			req.SetBasicAuth(user.Username(), password)
		}
	}

	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 {
	return c.Login(ctx, nil)
}

// Session returns the user's current session.
// Nil is returned if the session is not authenticated.
func (c *Client) Session(ctx context.Context) (*Session, error) {
	var s Session
	req := c.Resource(internal.SessionPath).WithAction("get").Request(http.MethodPost)
	err := c.Do(ctx, req, &s)
	if err != nil {
		if e, ok := err.(*statusError); ok {
			if e.res.StatusCode == http.StatusUnauthorized {
				return nil, nil
			}
		}
		return nil, err
	}
	return &s, nil
}

// Logout deletes the current session.
func (c *Client) Logout(ctx context.Context) error {
	req := c.Resource(internal.SessionPath).Request(http.MethodDelete)
	return c.Do(ctx, req, nil)
}

// Valid returns whether or not the client is valid and ready for use.
// This should be called after unmarshalling the client.
func (c *Client) Valid() bool {
	if c == nil {
		return false
	}

	if c.Client == nil {
		return false
	}

	return true
}

// Path returns rest.Path (see cache.Client)
func (c *Client) Path() string {
	return Path
}