<commit_before>package cluster

import (
	"crypto/tls"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"net"
	"time"

	"github.com/pkg/errors"

	lxd "github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxd/cluster/request"
	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/state"
	storagePools "github.com/lxc/lxd/lxd/storage"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/version"
)

// ErrCertificateExists indicates that a certificate already exists.
var ErrCertificateExists error = fmt.Errorf("Certificate already in trust store")

// Connect is a convenience around lxd.ConnectLXD that configures the client
// with the correct parameters for node-to-node communication.
//
// If 'notify' switch is true, then the user agent will be set to the special
// to the UserAgentNotifier value, which can be used in some cases to distinguish
// between a regular client request and an internal cluster request.
func Connect(address string, networkCert *shared.CertInfo, serverCert *shared.CertInfo, notify bool) (lxd.InstanceServer, error) {
	// Wait for a connection to the events API first for non-notify connections.
	if !notify {
		connected := false
		for i := 0; i < 20; i++ {
			listenersLock.Lock()
			_, ok := listeners[address]
			listenersLock.Unlock()

			if ok {
				connected = true
				break
			}

			time.Sleep(500 * time.Millisecond)
		}

		if !connected {
			return nil, fmt.Errorf("Missing event connection with target cluster member")
		}
	}

	args := &lxd.ConnectionArgs{
		TLSServerCert: string(networkCert.PublicKey()),
		TLSClientCert: string(serverCert.PublicKey()),
		TLSClientKey:  string(serverCert.PrivateKey()),
		SkipGetServer: true,
		UserAgent:     version.UserAgent,
	}

	if notify {
		args.UserAgent = request.UserAgentNotifier
	}

	url := fmt.Sprintf("https://%s", address)
	return lxd.ConnectLXD(url, args)
}

// ConnectIfInstanceIsRemote figures out the address of the node which is
// running the container with the given name. If it's not the local node will
// connect to it and return the connected client, otherwise it will just return
// nil.
func ConnectIfInstanceIsRemote(cluster *db.Cluster, projectName string, name string, networkCert *shared.CertInfo, serverCert *shared.CertInfo, instanceType instancetype.Type) (lxd.InstanceServer, error) {
	var address string // Node address
	err := cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		address, err = tx.GetNodeAddressOfInstance(projectName, name, instanceType)
		return err
	})
	if err != nil {
		return nil, err
	}
	if address == "" {
		// The instance is running right on this node, no need to connect.
		return nil, nil
	}
	return Connect(address, networkCert, serverCert, false)
}

// ConnectIfVolumeIsRemote figures out the address of the cluster member on which the volume with the given name is
// defined. If it's not the local cluster member it will connect to it and return the connected client, otherwise
// it just returns nil. If there is more than one cluster member with a matching volume name, an error is returned.
func ConnectIfVolumeIsRemote(s *state.State, poolName string, projectName string, volumeName string, volumeType int, networkCert *shared.CertInfo, serverCert *shared.CertInfo) (lxd.InstanceServer, error) {
	localNodeID := s.Cluster.GetNodeID()
	var err error
	var nodes []db.NodeInfo
	var poolID int64
	err = s.Cluster.Transaction(func(tx *db.ClusterTx) error {
		poolID, err = tx.GetStoragePoolID(poolName)
		if err != nil {
			return err
		}

		nodes, err = tx.GetStorageVolumeNodes(poolID, projectName, volumeName, volumeType)
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil && err != db.ErrNoClusterMember {
		return nil, err
	}

	// If volume uses a remote storage driver and so has no explicit cluster member, then we need to check
	// whether it is exclusively attached to remote instance, and if so then we need to forward the request to
	// the node whereit is currently used. This avoids conflicting with another member when using it locally.
	if err == db.ErrNoClusterMember {
		// GetLocalStoragePoolVolume returns a volume with an empty Location field for remote drivers.
		_, vol, err := s.Cluster.GetLocalStoragePoolVolume(projectName, volumeName, volumeType, poolID)
		if err != nil {
			return nil, err
		}

		remoteInstance, err := storagePools.VolumeUsedByExclusiveRemoteInstancesWithProfiles(s, poolName, projectName, vol)
		if err != nil {
			return nil, errors.Wrapf(err, "Failed checking if volume %q is available", volumeName)
		}

		if remoteInstance != nil {
			var instNode db.NodeInfo
			err := s.Cluster.Transaction(func(tx *db.ClusterTx) error {
				instNode, err = tx.GetNodeByName(remoteInstance.Node)
				return err
			})
			if err != nil {
				return nil, errors.Wrapf(err, "Failed getting cluster member info for %q", remoteInstance.Node)
			}

			// Replace node list with instance's cluster member node (which might be local member).
			nodes = []db.NodeInfo{instNode}
		} else {
			// Volume isn't exclusively attached to an instance. Use local cluster member.
			return nil, nil
		}
	}

	nodeCount := len(nodes)
	if nodeCount > 1 {
		return nil, fmt.Errorf("More than one cluster member has a volume named %q. Please target a specific member", volumeName)
	} else if nodeCount < 1 {
		// Should never get here.
		return nil, fmt.Errorf("Volume %q has empty cluster member list", volumeName)
	}

	node := nodes[0]
	if node.ID == localNodeID {
		// Use local cluster member if volume belongs to this local node.
		return nil, nil
	}

	// Connect to remote cluster member.
	return Connect(node.Address, networkCert, serverCert, false)
}

// SetupTrust is a convenience around InstanceServer.CreateCertificate that adds the given server certificate to
// the trusted pool of the cluster at the given address, using the given password. The certificate is added as
// type CertificateTypeServer to allow intra-member communication. If a certificate with the same fingerprint
// already exists with a different name or type, then no error is returned.
func SetupTrust(serverCert *shared.CertInfo, serverName string, targetAddress string, targetCert string, targetPassword string) error {
	// Connect to the target cluster node.
	args := &lxd.ConnectionArgs{
		TLSServerCert: targetCert,
		UserAgent:     version.UserAgent,
	}

	target, err := lxd.ConnectLXD(fmt.Sprintf("https://%s", targetAddress), args)
	if err != nil {
		return errors.Wrapf(err, "Failed to connect to target cluster node %q", targetAddress)
	}

	cert, err := generateTrustCertificate(serverCert, serverName)
	if err != nil {
		return errors.Wrapf(err, "Failed generating trust certificate")
	}

	post := api.CertificatesPost{
		CertificatePut: cert.CertificatePut,
		Certificate:    cert.Certificate,
		Password:       targetPassword,
	}

	err = target.CreateCertificate(post)
	if err != nil && err.Error() != ErrCertificateExists.Error() {
		return errors.Wrap(err, "Failed to add server cert to cluster")
	}

	return nil
}

// UpdateTrust ensures that the supplied certificate is stored in the target trust store with the correct name
// and type to ensure correct cluster operation. Should be called after SetupTrust. If a certificate with the same
// fingerprint is already in the trust store, but is of the wrong type or name then the existing certificate is
// updated to the correct type and name. If the existing certificate is the correct type but the wrong name then an
// error is returned. And if the existing certificate is the correct type and name then nothing more is done.
func UpdateTrust(serverCert *shared.CertInfo, serverName string, targetAddress string, targetCert string) error {
	// Connect to the target cluster node.
	args := &lxd.ConnectionArgs{
		TLSClientCert: string(serverCert.PublicKey()),
		TLSClientKey:  string(serverCert.PrivateKey()),
		TLSServerCert: targetCert,
		UserAgent:     version.UserAgent,
	}

	target, err := lxd.ConnectLXD(fmt.Sprintf("https://%s", targetAddress), args)
	if err != nil {
		return errors.Wrapf(err, "Failed to connect to target cluster node %q", targetAddress)
	}

	cert, err := generateTrustCertificate(serverCert, serverName)
	if err != nil {
		return errors.Wrapf(err, "Failed generating trust certificate")
	}

	existingCert, _, err := target.GetCertificate(cert.Fingerprint)
	if err != nil {
		return errors.Wrapf(err, "Failed getting existing certificate")
	}

	if existingCert.Name != serverName && existingCert.Type == api.CertificateTypeServer {
		// Don't alter an existing server certificate that has our fingerprint but not our name.
		// Something is wrong as this shouldn't happen.
		return fmt.Errorf("Existing server certificate with different name %q already in trust store", existingCert.Name)
	} else if existingCert.Name != serverName && existingCert.Type != api.CertificateTypeServer {
		// Ensure that if a client certificate already exists that matches our fingerprint, that it
		// has the correct name and type for cluster operation, to allow us to associate member
		// server names to certificate names.
		err = target.UpdateCertificate(cert.Fingerprint, cert.CertificatePut, "")
		if err != nil {
			return errors.Wrap(err, "Failed updating certificate name and type in trust store")
		}
	}

	return nil
}

// generateTrustCertificate converts the specified serverCert and serverName into an api.Certificate suitable for
// use as a trusted cluster server certificate.
func generateTrustCertificate(serverCert *shared.CertInfo, serverName string) (*api.Certificate, error) {
	block, _ := pem.Decode(serverCert.PublicKey())
	if block == nil {
		return nil, fmt.Errorf("Failed to decode certificate")
	}

	fingerprint, err := shared.CertFingerprintStr(string(serverCert.PublicKey()))
	if err != nil {
		return nil, errors.Wrap(err, "Failed to calculate fingerprint")
	}

	certificate := base64.StdEncoding.EncodeToString(block.Bytes)
	cert := api.Certificate{
		CertificatePut: api.CertificatePut{
			Name: serverName,
			Type: api.CertificateTypeServer, // Server type for intra-member communication.
		},
		Certificate: certificate,
		Fingerprint: fingerprint,
	}

	return &cert, nil
}

// HasConnectivity probes the member with the given address for connectivity.
func HasConnectivity(networkCert *shared.CertInfo, serverCert *shared.CertInfo, address string) bool {
	config, err := tlsClientConfig(networkCert, serverCert)
	if err != nil {
		return false
	}

	var conn net.Conn
	dialer := &net.Dialer{Timeout: time.Second}
	conn, err = tls.DialWithDialer(dialer, "tcp", address, config)
	if err == nil {
		conn.Close()
		return true
	}
	return false
}
<commit_msg>lxd/cluster: Pass original requestor around<commit_after>package cluster

import (
	"crypto/tls"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"time"

	"github.com/pkg/errors"

	lxd "github.com/lxc/lxd/client"
	"github.com/lxc/lxd/lxd/cluster/request"
	"github.com/lxc/lxd/lxd/db"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/state"
	storagePools "github.com/lxc/lxd/lxd/storage"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/version"
)

// ErrCertificateExists indicates that a certificate already exists.
var ErrCertificateExists error = fmt.Errorf("Certificate already in trust store")

// Connect is a convenience around lxd.ConnectLXD that configures the client
// with the correct parameters for node-to-node communication.
//
// If 'notify' switch is true, then the user agent will be set to the special
// to the UserAgentNotifier value, which can be used in some cases to distinguish
// between a regular client request and an internal cluster request.
func Connect(address string, networkCert *shared.CertInfo, serverCert *shared.CertInfo, notify bool) (lxd.InstanceServer, error) {
	// Wait for a connection to the events API first for non-notify connections.
	if !notify {
		connected := false
		for i := 0; i < 20; i++ {
			listenersLock.Lock()
			_, ok := listeners[address]
			listenersLock.Unlock()

			if ok {
				connected = true
				break
			}

			time.Sleep(500 * time.Millisecond)
		}

		if !connected {
			return nil, fmt.Errorf("Missing event connection with target cluster member")
		}
	}

	args := &lxd.ConnectionArgs{
		TLSServerCert: string(networkCert.PublicKey()),
		TLSClientCert: string(serverCert.PublicKey()),
		TLSClientKey:  string(serverCert.PrivateKey()),
		SkipGetServer: true,
		UserAgent:     version.UserAgent,
	}

	if notify {
		args.UserAgent = request.UserAgentNotifier
	}

	if r != nil {
		proxy := func(req *http.Request) (*url.URL, error) {
			ctx := r.Context()

			val, ok := ctx.Value("username").(string)
			if ok {
				req.Header.Add("X-LXD-cluster-username", val)
			}

			val, ok = ctx.Value("protocol").(string)
			if ok {
				req.Header.Add("X-LXD-cluster-protocol", val)
			}

			req.Header.Add("X-LXD-cluster-address", r.RemoteAddr)

			return shared.ProxyFromEnvironment(req)
		}

		args.Proxy = proxy
	}

	url := fmt.Sprintf("https://%s", address)
	return lxd.ConnectLXD(url, args)
}

// ConnectIfInstanceIsRemote figures out the address of the node which is
// running the container with the given name. If it's not the local node will
// connect to it and return the connected client, otherwise it will just return
// nil.
func ConnectIfInstanceIsRemote(cluster *db.Cluster, projectName string, name string, networkCert *shared.CertInfo, serverCert *shared.CertInfo, instanceType instancetype.Type) (lxd.InstanceServer, error) {
	var address string // Node address
	err := cluster.Transaction(func(tx *db.ClusterTx) error {
		var err error
		address, err = tx.GetNodeAddressOfInstance(projectName, name, instanceType)
		return err
	})
	if err != nil {
		return nil, err
	}
	if address == "" {
		// The instance is running right on this node, no need to connect.
		return nil, nil
	}
	return Connect(address, networkCert, serverCert, false)
}

// ConnectIfVolumeIsRemote figures out the address of the cluster member on which the volume with the given name is
// defined. If it's not the local cluster member it will connect to it and return the connected client, otherwise
// it just returns nil. If there is more than one cluster member with a matching volume name, an error is returned.
func ConnectIfVolumeIsRemote(s *state.State, poolName string, projectName string, volumeName string, volumeType int, networkCert *shared.CertInfo, serverCert *shared.CertInfo) (lxd.InstanceServer, error) {
	localNodeID := s.Cluster.GetNodeID()
	var err error
	var nodes []db.NodeInfo
	var poolID int64
	err = s.Cluster.Transaction(func(tx *db.ClusterTx) error {
		poolID, err = tx.GetStoragePoolID(poolName)
		if err != nil {
			return err
		}

		nodes, err = tx.GetStorageVolumeNodes(poolID, projectName, volumeName, volumeType)
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil && err != db.ErrNoClusterMember {
		return nil, err
	}

	// If volume uses a remote storage driver and so has no explicit cluster member, then we need to check
	// whether it is exclusively attached to remote instance, and if so then we need to forward the request to
	// the node whereit is currently used. This avoids conflicting with another member when using it locally.
	if err == db.ErrNoClusterMember {
		// GetLocalStoragePoolVolume returns a volume with an empty Location field for remote drivers.
		_, vol, err := s.Cluster.GetLocalStoragePoolVolume(projectName, volumeName, volumeType, poolID)
		if err != nil {
			return nil, err
		}

		remoteInstance, err := storagePools.VolumeUsedByExclusiveRemoteInstancesWithProfiles(s, poolName, projectName, vol)
		if err != nil {
			return nil, errors.Wrapf(err, "Failed checking if volume %q is available", volumeName)
		}

		if remoteInstance != nil {
			var instNode db.NodeInfo
			err := s.Cluster.Transaction(func(tx *db.ClusterTx) error {
				instNode, err = tx.GetNodeByName(remoteInstance.Node)
				return err
			})
			if err != nil {
				return nil, errors.Wrapf(err, "Failed getting cluster member info for %q", remoteInstance.Node)
			}

			// Replace node list with instance's cluster member node (which might be local member).
			nodes = []db.NodeInfo{instNode}
		} else {
			// Volume isn't exclusively attached to an instance. Use local cluster member.
			return nil, nil
		}
	}

	nodeCount := len(nodes)
	if nodeCount > 1 {
		return nil, fmt.Errorf("More than one cluster member has a volume named %q. Please target a specific member", volumeName)
	} else if nodeCount < 1 {
		// Should never get here.
		return nil, fmt.Errorf("Volume %q has empty cluster member list", volumeName)
	}

	node := nodes[0]
	if node.ID == localNodeID {
		// Use local cluster member if volume belongs to this local node.
		return nil, nil
	}

	// Connect to remote cluster member.
	return Connect(node.Address, networkCert, serverCert, false)
}

// SetupTrust is a convenience around InstanceServer.CreateCertificate that adds the given server certificate to
// the trusted pool of the cluster at the given address, using the given password. The certificate is added as
// type CertificateTypeServer to allow intra-member communication. If a certificate with the same fingerprint
// already exists with a different name or type, then no error is returned.
func SetupTrust(serverCert *shared.CertInfo, serverName string, targetAddress string, targetCert string, targetPassword string) error {
	// Connect to the target cluster node.
	args := &lxd.ConnectionArgs{
		TLSServerCert: targetCert,
		UserAgent:     version.UserAgent,
	}

	target, err := lxd.ConnectLXD(fmt.Sprintf("https://%s", targetAddress), args)
	if err != nil {
		return errors.Wrapf(err, "Failed to connect to target cluster node %q", targetAddress)
	}

	cert, err := generateTrustCertificate(serverCert, serverName)
	if err != nil {
		return errors.Wrapf(err, "Failed generating trust certificate")
	}

	post := api.CertificatesPost{
		CertificatePut: cert.CertificatePut,
		Certificate:    cert.Certificate,
		Password:       targetPassword,
	}

	err = target.CreateCertificate(post)
	if err != nil && err.Error() != ErrCertificateExists.Error() {
		return errors.Wrap(err, "Failed to add server cert to cluster")
	}

	return nil
}

// UpdateTrust ensures that the supplied certificate is stored in the target trust store with the correct name
// and type to ensure correct cluster operation. Should be called after SetupTrust. If a certificate with the same
// fingerprint is already in the trust store, but is of the wrong type or name then the existing certificate is
// updated to the correct type and name. If the existing certificate is the correct type but the wrong name then an
// error is returned. And if the existing certificate is the correct type and name then nothing more is done.
func UpdateTrust(serverCert *shared.CertInfo, serverName string, targetAddress string, targetCert string) error {
	// Connect to the target cluster node.
	args := &lxd.ConnectionArgs{
		TLSClientCert: string(serverCert.PublicKey()),
		TLSClientKey:  string(serverCert.PrivateKey()),
		TLSServerCert: targetCert,
		UserAgent:     version.UserAgent,
	}

	target, err := lxd.ConnectLXD(fmt.Sprintf("https://%s", targetAddress), args)
	if err != nil {
		return errors.Wrapf(err, "Failed to connect to target cluster node %q", targetAddress)
	}

	cert, err := generateTrustCertificate(serverCert, serverName)
	if err != nil {
		return errors.Wrapf(err, "Failed generating trust certificate")
	}

	existingCert, _, err := target.GetCertificate(cert.Fingerprint)
	if err != nil {
		return errors.Wrapf(err, "Failed getting existing certificate")
	}

	if existingCert.Name != serverName && existingCert.Type == api.CertificateTypeServer {
		// Don't alter an existing server certificate that has our fingerprint but not our name.
		// Something is wrong as this shouldn't happen.
		return fmt.Errorf("Existing server certificate with different name %q already in trust store", existingCert.Name)
	} else if existingCert.Name != serverName && existingCert.Type != api.CertificateTypeServer {
		// Ensure that if a client certificate already exists that matches our fingerprint, that it
		// has the correct name and type for cluster operation, to allow us to associate member
		// server names to certificate names.
		err = target.UpdateCertificate(cert.Fingerprint, cert.CertificatePut, "")
		if err != nil {
			return errors.Wrap(err, "Failed updating certificate name and type in trust store")
		}
	}

	return nil
}

// generateTrustCertificate converts the specified serverCert and serverName into an api.Certificate suitable for
// use as a trusted cluster server certificate.
func generateTrustCertificate(serverCert *shared.CertInfo, serverName string) (*api.Certificate, error) {
	block, _ := pem.Decode(serverCert.PublicKey())
	if block == nil {
		return nil, fmt.Errorf("Failed to decode certificate")
	}

	fingerprint, err := shared.CertFingerprintStr(string(serverCert.PublicKey()))
	if err != nil {
		return nil, errors.Wrap(err, "Failed to calculate fingerprint")
	}

	certificate := base64.StdEncoding.EncodeToString(block.Bytes)
	cert := api.Certificate{
		CertificatePut: api.CertificatePut{
			Name: serverName,
			Type: api.CertificateTypeServer, // Server type for intra-member communication.
		},
		Certificate: certificate,
		Fingerprint: fingerprint,
	}

	return &cert, nil
}

// HasConnectivity probes the member with the given address for connectivity.
func HasConnectivity(networkCert *shared.CertInfo, serverCert *shared.CertInfo, address string) bool {
	config, err := tlsClientConfig(networkCert, serverCert)
	if err != nil {
		return false
	}

	var conn net.Conn
	dialer := &net.Dialer{Timeout: time.Second}
	conn, err = tls.DialWithDialer(dialer, "tcp", address, config)
	if err == nil {
		conn.Close()
		return true
	}
	return false
}
<|endoftext|>


<commit_before>//go:build linux && cgo && !agent
// +build linux,cgo,!agent

package db

import (
	"fmt"

	"github.com/lxc/lxd/lxd/db/query"
	"github.com/lxc/lxd/shared/api"
)

// Code generation directives.
//
//go:generate -command mapper lxd-generate db mapper -t certificates.mapper.go
//go:generate mapper reset
//
//go:generate mapper stmt -p db -e certificate objects
//go:generate mapper stmt -p db -e certificate objects-by-Fingerprint
//go:generate mapper stmt -p db -e certificate projects-ref
//go:generate mapper stmt -p db -e certificate projects-ref-by-Fingerprint
//go:generate mapper stmt -p db -e certificate id
//go:generate mapper stmt -p db -e certificate create struct=Certificate
//go:generate mapper stmt -p db -e certificate create-projects-ref
//go:generate mapper stmt -p db -e certificate delete
//go:generate mapper stmt -p db -e certificate update struct=Certificate
//
//go:generate mapper method -p db -e certificate List
//go:generate mapper method -p db -e certificate Get
//go:generate mapper method -p db -e certificate ID struct=Certificate
//go:generate mapper method -p db -e certificate Exists struct=Certificate
//go:generate mapper method -p db -e certificate Create struct=Certificate
//go:generate mapper method -p db -e certificate ProjectsRef
//go:generate mapper method -p db -e certificate Delete
//go:generate mapper method -p db -e certificate Update struct=Certificate

// CertificateTypeClient indicates a client certificate type.
const CertificateTypeClient = 1

// CertificateTypeServer indicates a server certificate type.
const CertificateTypeServer = 2

// CertificateAPITypeToDBType converts an API type to the equivalent DB type.
func CertificateAPITypeToDBType(apiType string) (int, error) {
	switch apiType {
	case api.CertificateTypeClient:
		return CertificateTypeClient, nil
	case api.CertificateTypeServer:
		return CertificateTypeServer, nil
	}

	return -1, fmt.Errorf("Invalid certificate type")
}

// Certificate is here to pass the certificates content from the database around.
type Certificate struct {
	ID          int
	Fingerprint string `db:"primary=yes&comparison=like"`
	Type        int
	Name        string
	Certificate string
	Restricted  bool
	Projects    []string
}

// ToAPIType returns the API equivalent type.
func (cert *Certificate) ToAPIType() string {
	switch cert.Type {
	case CertificateTypeClient:
		return api.CertificateTypeClient
	case CertificateTypeServer:
		return api.CertificateTypeServer
	}

	return api.CertificateTypeUnknown
}

// ToAPI converts the database Certificate struct to an api.Certificate entry.
func (cert *Certificate) ToAPI() api.Certificate {
	resp := api.Certificate{}
	resp.Fingerprint = cert.Fingerprint
	resp.Certificate = cert.Certificate
	resp.Name = cert.Name
	resp.Restricted = cert.Restricted
	resp.Projects = cert.Projects
	resp.Type = cert.ToAPIType()

	return resp
}

// UpdateCertificateProjects updates the list of projects on a certificate.
func (c *ClusterTx) UpdateCertificateProjects(id int, projects []string) error {
	// Clear all projects from the restrictions.
	q := "DELETE FROM certificates_projects WHERE certificate_id=?"
	_, err := c.tx.Exec(q, id)
	if err != nil {
		return err
	}

	// Add the new restrictions.
	for _, name := range projects {
		projID, err := c.GetProjectID(name)
		if err != nil {
			return err
		}

		q := "INSERT INTO certificates_projects (certificate_id, project_id) VALUES (?, ?)"
		_, err = c.tx.Exec(q, id, projID)
		if err != nil {
			return err
		}
	}

	return nil
}

// CertificateFilter can be used to filter results yielded by GetCertInfos
type CertificateFilter struct {
	Fingerprint string // Matched with LIKE
}

// GetCertificate gets an CertBaseInfo object from the database.
// The argument fingerprint will be queried with a LIKE query, means you can
// pass a shortform and will get the full fingerprint.
// There can never be more than one certificate with a given fingerprint, as it is
// enforced by a UNIQUE constraint in the schema.
func (c *Cluster) GetCertificate(fingerprint string) (*Certificate, error) {
	var err error
	var cert *Certificate
	err = c.Transaction(func(tx *ClusterTx) error {
		cert, err = tx.GetCertificate(fingerprint + "%")
		if err != nil {
			return err
		}

		cert, err = tx.GetCertificate(cert.Fingerprint)
		return err
	})
	if err != nil {
		return nil, err
	}

	return cert, nil
}

// CreateCertificate stores a CertInfo object in the db, it will ignore the ID
// field from the CertInfo.
func (c *Cluster) CreateCertificate(cert Certificate) (int64, error) {
	var id int64
	var err error
	err = c.Transaction(func(tx *ClusterTx) error {
		id, err = tx.CreateCertificate(cert)
		return err
	})
	return id, err
}

// DeleteCertificate deletes a certificate from the db.
func (c *Cluster) DeleteCertificate(fingerprint string) error {
	err := c.Transaction(func(tx *ClusterTx) error {
		return tx.DeleteCertificate(fingerprint)
	})
	return err
}

// UpdateCertificate updates a certificate in the db.
func (c *Cluster) UpdateCertificate(fingerprint string, cert Certificate) error {
	err := c.Transaction(func(tx *ClusterTx) error {
		return tx.UpdateCertificate(fingerprint, cert)
	})
	return err
}

// UpdateCertificateProjects updates the list of projects on a certificate.
func (c *Cluster) UpdateCertificateProjects(id int, projects []string) error {
	err := c.Transaction(func(tx *ClusterTx) error {
		return tx.UpdateCertificateProjects(id, projects)
	})
	return err
}

// GetCertificates returns all available local certificates.
func (n *NodeTx) GetCertificates() ([]Certificate, error) {
	dbCerts := []struct {
		fingerprint string
		certType    int
		name        string
		certificate string
	}{}
	dest := func(i int) []interface{} {
		dbCerts = append(dbCerts, struct {
			fingerprint string
			certType    int
			name        string
			certificate string
		}{})
		return []interface{}{&dbCerts[i].fingerprint, &dbCerts[i].certType, &dbCerts[i].name, &dbCerts[i].certificate}
	}

	stmt, err := n.tx.Prepare("SELECT fingerprint, type, name, certificate FROM certificates")
	if err != nil {
		return nil, err
	}
	defer stmt.Close()

	err = query.SelectObjects(stmt, dest)
	if err != nil {
		return nil, err
	}

	certs := make([]Certificate, 0, len(dbCerts))
	for _, dbCert := range dbCerts {
		certs = append(certs, Certificate{
			Fingerprint: dbCert.fingerprint,
			Type:        dbCert.certType,
			Name:        dbCert.name,
			Certificate: dbCert.certificate,
		})
	}

	return certs, nil
}

// ReplaceCertificates removes all existing certificates from the local certificates table and replaces them with
// the ones provided.
func (n *NodeTx) ReplaceCertificates(certs []Certificate) error {
	_, err := n.tx.Exec("DELETE FROM certificates")
	if err != nil {
		return err
	}

	stmt, err := n.tx.Prepare("INSERT INTO certificates (fingerprint, type, name, certificate) VALUES(?,?,?,?)")
	if err != nil {
		return err
	}
	defer stmt.Close()

	for _, cert := range certs {
		_, err = stmt.Exec(cert.Fingerprint, cert.Type, cert.Name, cert.Certificate)
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>lxd/db/certificates: Adds DeleteCertificateByNameAndType function<commit_after>//go:build linux && cgo && !agent
// +build linux,cgo,!agent

package db

import (
	"fmt"

	"github.com/lxc/lxd/lxd/db/query"
	"github.com/lxc/lxd/shared/api"
)

// Code generation directives.
//
//go:generate -command mapper lxd-generate db mapper -t certificates.mapper.go
//go:generate mapper reset
//
//go:generate mapper stmt -p db -e certificate objects
//go:generate mapper stmt -p db -e certificate objects-by-Fingerprint
//go:generate mapper stmt -p db -e certificate projects-ref
//go:generate mapper stmt -p db -e certificate projects-ref-by-Fingerprint
//go:generate mapper stmt -p db -e certificate id
//go:generate mapper stmt -p db -e certificate create struct=Certificate
//go:generate mapper stmt -p db -e certificate create-projects-ref
//go:generate mapper stmt -p db -e certificate delete
//go:generate mapper stmt -p db -e certificate update struct=Certificate
//
//go:generate mapper method -p db -e certificate List
//go:generate mapper method -p db -e certificate Get
//go:generate mapper method -p db -e certificate ID struct=Certificate
//go:generate mapper method -p db -e certificate Exists struct=Certificate
//go:generate mapper method -p db -e certificate Create struct=Certificate
//go:generate mapper method -p db -e certificate ProjectsRef
//go:generate mapper method -p db -e certificate Delete
//go:generate mapper method -p db -e certificate Update struct=Certificate

// CertificateTypeClient indicates a client certificate type.
const CertificateTypeClient = 1

// CertificateTypeServer indicates a server certificate type.
const CertificateTypeServer = 2

// CertificateAPITypeToDBType converts an API type to the equivalent DB type.
func CertificateAPITypeToDBType(apiType string) (int, error) {
	switch apiType {
	case api.CertificateTypeClient:
		return CertificateTypeClient, nil
	case api.CertificateTypeServer:
		return CertificateTypeServer, nil
	}

	return -1, fmt.Errorf("Invalid certificate type")
}

// Certificate is here to pass the certificates content from the database around.
type Certificate struct {
	ID          int
	Fingerprint string `db:"primary=yes&comparison=like"`
	Type        int
	Name        string
	Certificate string
	Restricted  bool
	Projects    []string
}

// ToAPIType returns the API equivalent type.
func (cert *Certificate) ToAPIType() string {
	switch cert.Type {
	case CertificateTypeClient:
		return api.CertificateTypeClient
	case CertificateTypeServer:
		return api.CertificateTypeServer
	}

	return api.CertificateTypeUnknown
}

// ToAPI converts the database Certificate struct to an api.Certificate entry.
func (cert *Certificate) ToAPI() api.Certificate {
	resp := api.Certificate{}
	resp.Fingerprint = cert.Fingerprint
	resp.Certificate = cert.Certificate
	resp.Name = cert.Name
	resp.Restricted = cert.Restricted
	resp.Projects = cert.Projects
	resp.Type = cert.ToAPIType()

	return resp
}

// UpdateCertificateProjects updates the list of projects on a certificate.
func (c *ClusterTx) UpdateCertificateProjects(id int, projects []string) error {
	// Clear all projects from the restrictions.
	q := "DELETE FROM certificates_projects WHERE certificate_id=?"
	_, err := c.tx.Exec(q, id)
	if err != nil {
		return err
	}

	// Add the new restrictions.
	for _, name := range projects {
		projID, err := c.GetProjectID(name)
		if err != nil {
			return err
		}

		q := "INSERT INTO certificates_projects (certificate_id, project_id) VALUES (?, ?)"
		_, err = c.tx.Exec(q, id, projID)
		if err != nil {
			return err
		}
	}

	return nil
}

// DeleteCertificateByNameAndType deletes the certificate(s) matching the given name and certificate type.
func (c *ClusterTx) DeleteCertificateByNameAndType(name string, certType int) error {
	_, err := c.tx.Exec("DELETE FROM certificates WHERE name = ? and type = ?", name, certType)
	return err
}

// CertificateFilter can be used to filter results yielded by GetCertInfos
type CertificateFilter struct {
	Fingerprint string // Matched with LIKE
}

// GetCertificate gets an CertBaseInfo object from the database.
// The argument fingerprint will be queried with a LIKE query, means you can
// pass a shortform and will get the full fingerprint.
// There can never be more than one certificate with a given fingerprint, as it is
// enforced by a UNIQUE constraint in the schema.
func (c *Cluster) GetCertificate(fingerprint string) (*Certificate, error) {
	var err error
	var cert *Certificate
	err = c.Transaction(func(tx *ClusterTx) error {
		cert, err = tx.GetCertificate(fingerprint + "%")
		if err != nil {
			return err
		}

		cert, err = tx.GetCertificate(cert.Fingerprint)
		return err
	})
	if err != nil {
		return nil, err
	}

	return cert, nil
}

// CreateCertificate stores a CertInfo object in the db, it will ignore the ID
// field from the CertInfo.
func (c *Cluster) CreateCertificate(cert Certificate) (int64, error) {
	var id int64
	var err error
	err = c.Transaction(func(tx *ClusterTx) error {
		id, err = tx.CreateCertificate(cert)
		return err
	})
	return id, err
}

// DeleteCertificate deletes a certificate from the db.
func (c *Cluster) DeleteCertificate(fingerprint string) error {
	err := c.Transaction(func(tx *ClusterTx) error {
		return tx.DeleteCertificate(fingerprint)
	})
	return err
}

// UpdateCertificate updates a certificate in the db.
func (c *Cluster) UpdateCertificate(fingerprint string, cert Certificate) error {
	err := c.Transaction(func(tx *ClusterTx) error {
		return tx.UpdateCertificate(fingerprint, cert)
	})
	return err
}

// UpdateCertificateProjects updates the list of projects on a certificate.
func (c *Cluster) UpdateCertificateProjects(id int, projects []string) error {
	err := c.Transaction(func(tx *ClusterTx) error {
		return tx.UpdateCertificateProjects(id, projects)
	})
	return err
}

// GetCertificates returns all available local certificates.
func (n *NodeTx) GetCertificates() ([]Certificate, error) {
	dbCerts := []struct {
		fingerprint string
		certType    int
		name        string
		certificate string
	}{}
	dest := func(i int) []interface{} {
		dbCerts = append(dbCerts, struct {
			fingerprint string
			certType    int
			name        string
			certificate string
		}{})
		return []interface{}{&dbCerts[i].fingerprint, &dbCerts[i].certType, &dbCerts[i].name, &dbCerts[i].certificate}
	}

	stmt, err := n.tx.Prepare("SELECT fingerprint, type, name, certificate FROM certificates")
	if err != nil {
		return nil, err
	}
	defer stmt.Close()

	err = query.SelectObjects(stmt, dest)
	if err != nil {
		return nil, err
	}

	certs := make([]Certificate, 0, len(dbCerts))
	for _, dbCert := range dbCerts {
		certs = append(certs, Certificate{
			Fingerprint: dbCert.fingerprint,
			Type:        dbCert.certType,
			Name:        dbCert.name,
			Certificate: dbCert.certificate,
		})
	}

	return certs, nil
}

// ReplaceCertificates removes all existing certificates from the local certificates table and replaces them with
// the ones provided.
func (n *NodeTx) ReplaceCertificates(certs []Certificate) error {
	_, err := n.tx.Exec("DELETE FROM certificates")
	if err != nil {
		return err
	}

	stmt, err := n.tx.Prepare("INSERT INTO certificates (fingerprint, type, name, certificate) VALUES(?,?,?,?)")
	if err != nil {
		return err
	}
	defer stmt.Close()

	for _, cert := range certs {
		_, err = stmt.Exec(cert.Fingerprint, cert.Type, cert.Name, cert.Certificate)
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package metrics

import (
	"fmt"
	"sort"
	"strings"
)

// NewMetricSet returns a new MetricSet.
func NewMetricSet(labels map[string]string) *MetricSet {
	out := MetricSet{set: make(map[MetricType][]Sample)}

	if labels != nil {
		out.labels = labels
	} else {
		out.labels = make(map[string]string)
	}

	return &out
}

// AddSamples adds samples of the type metricType to the MetricSet.
func (m *MetricSet) AddSamples(metricType MetricType, samples ...Sample) {
	for i := 0; i < len(samples); i++ {
		// Add global labels to samples
		for labelName, labelValue := range m.labels {
			// Ensure we always have a valid Labels map
			if samples[i].Labels == nil {
				samples[i].Labels = make(map[string]string)
			}

			samples[i].Labels[labelName] = labelValue
		}
	}

	m.set[metricType] = append(m.set[metricType], samples...)
}

// Merge merges two MetricSets.
func (m *MetricSet) Merge(metricSet *MetricSet) {
	if metricSet == nil {
		return
	}

	for k := range metricSet.set {
		m.set[k] = append(m.set[k], metricSet.set[k]...)
	}
}

func (m *MetricSet) String() string {
	var out strings.Builder
	metricTypes := []MetricType{}

	// Sort output by metric type name
	for metricType := range m.set {
		metricTypes = append(metricTypes, metricType)
	}

	sort.SliceStable(metricTypes, func(i, j int) bool {
		return int(metricTypes[i]) < int(metricTypes[j])
	})

	for _, metricType := range metricTypes {
		// Add HELP message as specified by OpenMetrics
		_, err := out.WriteString(MetricHeaders[metricType] + "\n")
		if err != nil {
			return ""
		}

		metricTypeName := ""

		if strings.HasSuffix(MetricNames[metricType], "_total") {
			metricTypeName = "counter"
		} else if strings.HasSuffix(MetricNames[metricType], "_bytes") {
			metricTypeName = "gauge"
		}

		// Add TYPE message as specified by OpenMetrics
		_, err = out.WriteString(fmt.Sprintf("# TYPE %s %s\n", MetricNames[metricType], metricTypeName))
		if err != nil {
			return ""
		}

		for _, sample := range m.set[metricType] {
			firstLabel := true
			labels := ""
			labelNames := []string{}

			// Add and sort labels if there are any
			for labelName := range sample.Labels {
				labelNames = append(labelNames, labelName)
			}

			sort.Strings(labelNames)

			for _, labelName := range labelNames {
				if !firstLabel {
					labels += ","
				}

				labels += fmt.Sprintf(`%s="%s"`, labelName, sample.Labels[labelName])
				firstLabel = false
			}

			if labels == "" {
				_, err = out.WriteString(fmt.Sprintf("%s %d\n", MetricNames[metricType], sample.Value))
			} else {
				_, err = out.WriteString(fmt.Sprintf("%s{%s} %d\n", MetricNames[metricType], labels, sample.Value))
			}
			if err != nil {
				return ""
			}
		}
	}

	return out.String()
}

// MetricSetFromAPI converts api.Metrics to a MetricSet, and returns it.
func MetricSetFromAPI(metrics *Metrics, labels map[string]string) (*MetricSet, error) {
	set := NewMetricSet(labels)

	// CPU stats
	for dev, stats := range metrics.CPU {
		getLabels := func(mode string) map[string]string {
			labels := map[string]string{"mode": mode}
			cpu := ""

			if dev != "cpu" {
				fmt.Sscanf(dev, "cpu%s", &cpu)
			}

			if cpu != "" {
				labels["cpu"] = cpu
			}

			return labels
		}

		set.AddSamples(CPUSecondsTotal,
			Sample{
				Value:  stats.SecondsIOWait,
				Labels: getLabels("iowait"),
			},
			Sample{
				Value:  stats.SecondsIRQ,
				Labels: getLabels("irq"),
			},
			Sample{
				Value:  stats.SecondsIdle,
				Labels: getLabels("idle"),
			},
			Sample{
				Value:  stats.SecondsNice,
				Labels: getLabels("nice"),
			},
			Sample{
				Value:  stats.SecondsSoftIRQ,
				Labels: getLabels("softirq"),
			},
			Sample{
				Value:  stats.SecondsSteal,
				Labels: getLabels("steal"),
			},
			Sample{
				Value:  stats.SecondsSystem,
				Labels: getLabels("system"),
			},
			Sample{
				Value:  stats.SecondsUser,
				Labels: getLabels("user"),
			},
		)
	}

	// Disk stats
	for dev, stats := range metrics.Disk {
		labels := map[string]string{"device": dev}

		set.AddSamples(DiskReadBytesTotal, Sample{Value: stats.ReadBytes, Labels: labels})
		set.AddSamples(DiskReadsCompletedTotal, Sample{Value: stats.ReadsCompleted, Labels: labels})
		set.AddSamples(DiskWritesCompletedTotal, Sample{Value: stats.WritesCompleted, Labels: labels})
		set.AddSamples(DiskWrittenBytesTotal, Sample{Value: stats.WrittenBytes, Labels: labels})
	}

	// Filesystem stats
	for dev, stats := range metrics.Filesystem {
		labels := map[string]string{"device": dev, "fstype": stats.FSType, "mountpoint": stats.Mountpoint}

		set.AddSamples(FilesystemAvailBytes, Sample{Value: stats.AvailableBytes, Labels: labels})
		set.AddSamples(FilesystemFreeBytes, Sample{Value: stats.FreeBytes, Labels: labels})
		set.AddSamples(FilesystemSizeBytes, Sample{Value: stats.SizeBytes, Labels: labels})
	}

	// Memory stats
	set.AddSamples(MemoryActiveAnonBytes, Sample{Value: metrics.Memory.ActiveAnonBytes})
	set.AddSamples(MemoryActiveBytes, Sample{Value: metrics.Memory.ActiveBytes})
	set.AddSamples(MemoryActiveFileBytes, Sample{Value: metrics.Memory.ActiveFileBytes})
	set.AddSamples(MemoryCachedBytes, Sample{Value: metrics.Memory.CachedBytes})
	set.AddSamples(MemoryDirtyBytes, Sample{Value: metrics.Memory.DirtyBytes})
	set.AddSamples(MemoryHugePagesFreeBytes, Sample{Value: metrics.Memory.HugepagesFreeBytes})
	set.AddSamples(MemoryHugePagesTotalBytes, Sample{Value: metrics.Memory.HugepagesTotalBytes})
	set.AddSamples(MemoryInactiveAnonBytes, Sample{Value: metrics.Memory.InactiveAnonBytes})
	set.AddSamples(MemoryInactiveBytes, Sample{Value: metrics.Memory.InactiveBytes})
	set.AddSamples(MemoryInactiveFileBytes, Sample{Value: metrics.Memory.InactiveFileBytes})
	set.AddSamples(MemoryMappedBytes, Sample{Value: metrics.Memory.MappedBytes})
	set.AddSamples(MemoryMemAvailableBytes, Sample{Value: metrics.Memory.MemAvailableBytes})
	set.AddSamples(MemoryMemFreeBytes, Sample{Value: metrics.Memory.MemFreeBytes})
	set.AddSamples(MemoryMemTotalBytes, Sample{Value: metrics.Memory.MemTotalBytes})
	set.AddSamples(MemoryRSSBytes, Sample{Value: metrics.Memory.RSSBytes})
	set.AddSamples(MemoryShmemBytes, Sample{Value: metrics.Memory.ShmemBytes})
	set.AddSamples(MemorySwapBytes, Sample{Value: metrics.Memory.SwapBytes})
	set.AddSamples(MemoryUnevictableBytes, Sample{Value: metrics.Memory.UnevictableBytes})
	set.AddSamples(MemoryWritebackBytes, Sample{Value: metrics.Memory.WritebackBytes})

	// Network stats
	for dev, stats := range metrics.Network {
		labels := map[string]string{"device": dev}

		set.AddSamples(NetworkReceiveBytesTotal, Sample{Value: stats.ReceiveBytes, Labels: labels})
		set.AddSamples(NetworkReceiveDropTotal, Sample{Value: stats.ReceiveDrop, Labels: labels})
		set.AddSamples(NetworkReceiveErrsTotal, Sample{Value: stats.ReceiveErrors, Labels: labels})
		set.AddSamples(NetworkReceivePacketsTotal, Sample{Value: stats.ReceivePackets, Labels: labels})
		set.AddSamples(NetworkTransmitBytesTotal, Sample{Value: stats.TransmitBytes, Labels: labels})
		set.AddSamples(NetworkTransmitDropTotal, Sample{Value: stats.TransmitDrop, Labels: labels})
		set.AddSamples(NetworkTransmitErrsTotal, Sample{Value: stats.TransmitErrors, Labels: labels})
		set.AddSamples(NetworkTransmitPacketsTotal, Sample{Value: stats.TransmitPackets, Labels: labels})
	}

	// Procs stats
	set.AddSamples(ProcsTotal, Sample{Value: metrics.ProcessesTotal})

	return set, nil
}
<commit_msg>lxd/metrics: Change ProcsTotal to gauge<commit_after>package metrics

import (
	"fmt"
	"sort"
	"strings"
)

// NewMetricSet returns a new MetricSet.
func NewMetricSet(labels map[string]string) *MetricSet {
	out := MetricSet{set: make(map[MetricType][]Sample)}

	if labels != nil {
		out.labels = labels
	} else {
		out.labels = make(map[string]string)
	}

	return &out
}

// AddSamples adds samples of the type metricType to the MetricSet.
func (m *MetricSet) AddSamples(metricType MetricType, samples ...Sample) {
	for i := 0; i < len(samples); i++ {
		// Add global labels to samples
		for labelName, labelValue := range m.labels {
			// Ensure we always have a valid Labels map
			if samples[i].Labels == nil {
				samples[i].Labels = make(map[string]string)
			}

			samples[i].Labels[labelName] = labelValue
		}
	}

	m.set[metricType] = append(m.set[metricType], samples...)
}

// Merge merges two MetricSets.
func (m *MetricSet) Merge(metricSet *MetricSet) {
	if metricSet == nil {
		return
	}

	for k := range metricSet.set {
		m.set[k] = append(m.set[k], metricSet.set[k]...)
	}
}

func (m *MetricSet) String() string {
	var out strings.Builder
	metricTypes := []MetricType{}

	// Sort output by metric type name
	for metricType := range m.set {
		metricTypes = append(metricTypes, metricType)
	}

	sort.SliceStable(metricTypes, func(i, j int) bool {
		return int(metricTypes[i]) < int(metricTypes[j])
	})

	for _, metricType := range metricTypes {
		// Add HELP message as specified by OpenMetrics
		_, err := out.WriteString(MetricHeaders[metricType] + "\n")
		if err != nil {
			return ""
		}

		metricTypeName := ""

		// ProcsTotal is a gauge according to the OpenMetrics spec as its value can decrease.
		if metricType == ProcsTotal {
			metricTypeName = "gauge"
		} else if strings.HasSuffix(MetricNames[metricType], "_total") {
			metricTypeName = "counter"
		} else if strings.HasSuffix(MetricNames[metricType], "_bytes") {
			metricTypeName = "gauge"
		}

		// Add TYPE message as specified by OpenMetrics
		_, err = out.WriteString(fmt.Sprintf("# TYPE %s %s\n", MetricNames[metricType], metricTypeName))
		if err != nil {
			return ""
		}

		for _, sample := range m.set[metricType] {
			firstLabel := true
			labels := ""
			labelNames := []string{}

			// Add and sort labels if there are any
			for labelName := range sample.Labels {
				labelNames = append(labelNames, labelName)
			}

			sort.Strings(labelNames)

			for _, labelName := range labelNames {
				if !firstLabel {
					labels += ","
				}

				labels += fmt.Sprintf(`%s="%s"`, labelName, sample.Labels[labelName])
				firstLabel = false
			}

			if labels == "" {
				_, err = out.WriteString(fmt.Sprintf("%s %d\n", MetricNames[metricType], sample.Value))
			} else {
				_, err = out.WriteString(fmt.Sprintf("%s{%s} %d\n", MetricNames[metricType], labels, sample.Value))
			}
			if err != nil {
				return ""
			}
		}
	}

	return out.String()
}

// MetricSetFromAPI converts api.Metrics to a MetricSet, and returns it.
func MetricSetFromAPI(metrics *Metrics, labels map[string]string) (*MetricSet, error) {
	set := NewMetricSet(labels)

	// CPU stats
	for dev, stats := range metrics.CPU {
		getLabels := func(mode string) map[string]string {
			labels := map[string]string{"mode": mode}
			cpu := ""

			if dev != "cpu" {
				fmt.Sscanf(dev, "cpu%s", &cpu)
			}

			if cpu != "" {
				labels["cpu"] = cpu
			}

			return labels
		}

		set.AddSamples(CPUSecondsTotal,
			Sample{
				Value:  stats.SecondsIOWait,
				Labels: getLabels("iowait"),
			},
			Sample{
				Value:  stats.SecondsIRQ,
				Labels: getLabels("irq"),
			},
			Sample{
				Value:  stats.SecondsIdle,
				Labels: getLabels("idle"),
			},
			Sample{
				Value:  stats.SecondsNice,
				Labels: getLabels("nice"),
			},
			Sample{
				Value:  stats.SecondsSoftIRQ,
				Labels: getLabels("softirq"),
			},
			Sample{
				Value:  stats.SecondsSteal,
				Labels: getLabels("steal"),
			},
			Sample{
				Value:  stats.SecondsSystem,
				Labels: getLabels("system"),
			},
			Sample{
				Value:  stats.SecondsUser,
				Labels: getLabels("user"),
			},
		)
	}

	// Disk stats
	for dev, stats := range metrics.Disk {
		labels := map[string]string{"device": dev}

		set.AddSamples(DiskReadBytesTotal, Sample{Value: stats.ReadBytes, Labels: labels})
		set.AddSamples(DiskReadsCompletedTotal, Sample{Value: stats.ReadsCompleted, Labels: labels})
		set.AddSamples(DiskWritesCompletedTotal, Sample{Value: stats.WritesCompleted, Labels: labels})
		set.AddSamples(DiskWrittenBytesTotal, Sample{Value: stats.WrittenBytes, Labels: labels})
	}

	// Filesystem stats
	for dev, stats := range metrics.Filesystem {
		labels := map[string]string{"device": dev, "fstype": stats.FSType, "mountpoint": stats.Mountpoint}

		set.AddSamples(FilesystemAvailBytes, Sample{Value: stats.AvailableBytes, Labels: labels})
		set.AddSamples(FilesystemFreeBytes, Sample{Value: stats.FreeBytes, Labels: labels})
		set.AddSamples(FilesystemSizeBytes, Sample{Value: stats.SizeBytes, Labels: labels})
	}

	// Memory stats
	set.AddSamples(MemoryActiveAnonBytes, Sample{Value: metrics.Memory.ActiveAnonBytes})
	set.AddSamples(MemoryActiveBytes, Sample{Value: metrics.Memory.ActiveBytes})
	set.AddSamples(MemoryActiveFileBytes, Sample{Value: metrics.Memory.ActiveFileBytes})
	set.AddSamples(MemoryCachedBytes, Sample{Value: metrics.Memory.CachedBytes})
	set.AddSamples(MemoryDirtyBytes, Sample{Value: metrics.Memory.DirtyBytes})
	set.AddSamples(MemoryHugePagesFreeBytes, Sample{Value: metrics.Memory.HugepagesFreeBytes})
	set.AddSamples(MemoryHugePagesTotalBytes, Sample{Value: metrics.Memory.HugepagesTotalBytes})
	set.AddSamples(MemoryInactiveAnonBytes, Sample{Value: metrics.Memory.InactiveAnonBytes})
	set.AddSamples(MemoryInactiveBytes, Sample{Value: metrics.Memory.InactiveBytes})
	set.AddSamples(MemoryInactiveFileBytes, Sample{Value: metrics.Memory.InactiveFileBytes})
	set.AddSamples(MemoryMappedBytes, Sample{Value: metrics.Memory.MappedBytes})
	set.AddSamples(MemoryMemAvailableBytes, Sample{Value: metrics.Memory.MemAvailableBytes})
	set.AddSamples(MemoryMemFreeBytes, Sample{Value: metrics.Memory.MemFreeBytes})
	set.AddSamples(MemoryMemTotalBytes, Sample{Value: metrics.Memory.MemTotalBytes})
	set.AddSamples(MemoryRSSBytes, Sample{Value: metrics.Memory.RSSBytes})
	set.AddSamples(MemoryShmemBytes, Sample{Value: metrics.Memory.ShmemBytes})
	set.AddSamples(MemorySwapBytes, Sample{Value: metrics.Memory.SwapBytes})
	set.AddSamples(MemoryUnevictableBytes, Sample{Value: metrics.Memory.UnevictableBytes})
	set.AddSamples(MemoryWritebackBytes, Sample{Value: metrics.Memory.WritebackBytes})

	// Network stats
	for dev, stats := range metrics.Network {
		labels := map[string]string{"device": dev}

		set.AddSamples(NetworkReceiveBytesTotal, Sample{Value: stats.ReceiveBytes, Labels: labels})
		set.AddSamples(NetworkReceiveDropTotal, Sample{Value: stats.ReceiveDrop, Labels: labels})
		set.AddSamples(NetworkReceiveErrsTotal, Sample{Value: stats.ReceiveErrors, Labels: labels})
		set.AddSamples(NetworkReceivePacketsTotal, Sample{Value: stats.ReceivePackets, Labels: labels})
		set.AddSamples(NetworkTransmitBytesTotal, Sample{Value: stats.TransmitBytes, Labels: labels})
		set.AddSamples(NetworkTransmitDropTotal, Sample{Value: stats.TransmitDrop, Labels: labels})
		set.AddSamples(NetworkTransmitErrsTotal, Sample{Value: stats.TransmitErrors, Labels: labels})
		set.AddSamples(NetworkTransmitPacketsTotal, Sample{Value: stats.TransmitPackets, Labels: labels})
	}

	// Procs stats
	set.AddSamples(ProcsTotal, Sample{Value: metrics.ProcessesTotal})

	return set, nil
}
<|endoftext|>


<commit_before>package bnblog

import (
	"fmt"
	"github.com/codegangsta/martini"

	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"google.golang.org/appengine"
	"google.golang.org/appengine/file"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/urlfetch"
	"google.golang.org/appengine/user"
	"google.golang.org/cloud"
	"google.golang.org/cloud/storage"
	"io/ioutil"
	"net/http"
)

func UploadFile(rw http.ResponseWriter, req *http.Request, params martini.Params) {
	c := appengine.NewContext(req)
	u := user.Current(c)
	if u == nil {
		http.Error(rw, fmt.Sprintf("wat %s", u), http.StatusForbidden)
		return
	}

	if fmt.Sprintf("%s", u) != "ben@benjojo.co.uk" && fmt.Sprintf("%s", u) != "ben@benjojo.com" {
		http.Error(rw, fmt.Sprintf("wat? %s", u), http.StatusForbidden)
		return
	}

	hc := &http.Client{
		Transport: &oauth2.Transport{
			Source: google.AppEngineTokenSource(c, storage.ScopeFullControl),
			Base:   &urlfetch.Transport{Context: c},
		},
	}

	file_from_client, headers, err := req.FormFile("fileToUpload")

	if err != nil {
		http.Error(rw, "Was that even a file?!", http.StatusBadRequest)
		return
	}

	defer file_from_client.Close()

	bucket := ""
	if bucket == "" {
		var err error
		if bucket, err = file.DefaultBucketName(c); err != nil {
			// log.Errorf(c, "failed to get default GCS bucket name: %v", err)
			return
		}
	}
	ctx := cloud.NewContext(appengine.AppID(c), hc)

	fn := RandString(10)

	bin, _ := ioutil.ReadAll(file_from_client)

	wc1 := storage.NewWriter(ctx, bucket, fn)
	wc1.ContentType = headers.Header.Get("Content-Type")
	wc1.ACL = []storage.ACLRule{{storage.AllUsers, storage.RoleReader}}
	if _, err := wc1.Write(bin); err != nil {
		log.Warningf(c, "ouch! %s", err)
	}
	if err := wc1.Close(); err != nil {
		log.Warningf(c, "ouch! %s", err)
	}
	log.Infof(c, "updated object:", wc1.Object())

	rw.Write([]byte(fn))
	log.Warningf(c, "fin.")

}

func ReadFile(rw http.ResponseWriter, req *http.Request, params martini.Params) {
	c := appengine.NewContext(req)

	bucket := ""
	if bucket == "" {
		var err error
		if bucket, err = file.DefaultBucketName(c); err != nil {
			// log.Errorf(c, "failed to get default GCS bucket name: %v", err)
			return
		}
	}

	hc := &http.Client{
		Transport: &oauth2.Transport{
			Source: google.AppEngineTokenSource(c, storage.ScopeFullControl),
			Base:   &urlfetch.Transport{Context: c},
		},
	}

	ctx := cloud.NewContext(appengine.AppID(c), hc)

	rc, err := storage.NewReader(ctx, bucket, params["tag"])
	if err != nil {
		log.Warningf(c, "readFile: unable to open file from bucket %q, file %q: %v", bucket, params["tag"], err)
		return
	}
	defer rc.Close()
	slurp, err := ioutil.ReadAll(rc)
	if err != nil {
		log.Warningf(c, "readFile: unable to read data from bucket %q, file %q: %v", bucket, params["tag"], err)
		return
	}
	o, err := storage.StatObject(ctx, bucket, params["tag"])
	if err != nil {
		rw.Header().Add("Content-Type", o.ContentType)
	}

	rw.Write([]byte(slurp))
}

type File struct {
	Name    string
	Content []byte
	Type    string
}

func ExportAllFiles(rw http.ResponseWriter, req *http.Request) (export []File) {
	c := appengine.NewContext(req)

	export = make([]File, 0)

	bucket := ""
	if bucket == "" {
		var err error
		if bucket, err = file.DefaultBucketName(c); err != nil {
			// log.Errorf(c, "failed to get default GCS bucket name: %v", err)
			return
		}
	}

	hc := &http.Client{
		Transport: &oauth2.Transport{
			Source: google.AppEngineTokenSource(c, storage.ScopeFullControl),
			Base:   &urlfetch.Transport{Context: c},
		},
	}

	ctx := cloud.NewContext(appengine.AppID(c), hc)

	query := &storage.Query{Prefix: ""}
	for query != nil {
		objs, err := storage.ListObjects(ctx, bucket, query)
		if err != nil {
			//d.errorf("listBucket: unable to list bucket %q: %v", bucket, err)
			return
		}
		query = objs.Next

		for _, obj := range objs.Results {
			//d.dumpStats(obj)
			newfile := File{}
			newfile.Name = obj.Name
			newfile.Type = obj.ContentType

			rc, err := storage.NewReader(ctx, bucket, obj.Name)
			if err != nil {
				log.Warningf(c, "readFile: unable to open file from bucket %q, file %q: %v", bucket, obj.Name, err)
				return
			}
			defer rc.Close()
			slurp, err := ioutil.ReadAll(rc)

			newfile.Content = slurp
			export = append(export, newfile)
		}
	}

	return export
}

func GimmeDC(rw http.ResponseWriter, req *http.Request) string {
	c := appengine.NewContext(req)
	return appengine.Datacenter(c)
}
<commit_msg>Fixed issue with older commit where it caused ALL images to break.<commit_after>package bnblog

import (
	"fmt"
	"github.com/codegangsta/martini"

	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"google.golang.org/appengine"
	"google.golang.org/appengine/file"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/urlfetch"
	"google.golang.org/appengine/user"
	"google.golang.org/cloud"
	"google.golang.org/cloud/storage"
	"io/ioutil"
	"net/http"
)

func UploadFile(rw http.ResponseWriter, req *http.Request, params martini.Params) {
	c := appengine.NewContext(req)
	u := user.Current(c)
	if u == nil {
		http.Error(rw, fmt.Sprintf("wat %s", u), http.StatusForbidden)
		return
	}

	if fmt.Sprintf("%s", u) != "ben@benjojo.co.uk" && fmt.Sprintf("%s", u) != "ben@benjojo.com" {
		http.Error(rw, fmt.Sprintf("wat? %s", u), http.StatusForbidden)
		return
	}

	hc := &http.Client{
		Transport: &oauth2.Transport{
			Source: google.AppEngineTokenSource(c, storage.ScopeFullControl),
			Base:   &urlfetch.Transport{Context: c},
		},
	}

	file_from_client, headers, err := req.FormFile("fileToUpload")

	if err != nil {
		http.Error(rw, "Was that even a file?!", http.StatusBadRequest)
		return
	}

	defer file_from_client.Close()

	bucket := ""
	if bucket == "" {
		var err error
		if bucket, err = file.DefaultBucketName(c); err != nil {
			// log.Errorf(c, "failed to get default GCS bucket name: %v", err)
			return
		}
	}
	ctx := cloud.NewContext(appengine.AppID(c), hc)

	fn := RandString(10)

	bin, _ := ioutil.ReadAll(file_from_client)

	wc1 := storage.NewWriter(ctx, bucket, fn)
	wc1.ContentType = headers.Header.Get("Content-Type")
	wc1.ACL = []storage.ACLRule{{storage.AllUsers, storage.RoleReader}}
	if _, err := wc1.Write(bin); err != nil {
		log.Warningf(c, "ouch! %s", err)
	}
	if err := wc1.Close(); err != nil {
		log.Warningf(c, "ouch! %s", err)
	}
	log.Infof(c, "updated object:", wc1.Object())

	rw.Write([]byte(fn))
	log.Warningf(c, "fin.")

}

func ReadFile(rw http.ResponseWriter, req *http.Request, params martini.Params) {
	c := appengine.NewContext(req)

	bucket := ""
	if bucket == "" {
		var err error
		if bucket, err = file.DefaultBucketName(c); err != nil {
			// log.Errorf(c, "failed to get default GCS bucket name: %v", err)
			return
		}
	}

	hc := &http.Client{
		Transport: &oauth2.Transport{
			Source: google.AppEngineTokenSource(c, storage.ScopeFullControl),
			Base:   &urlfetch.Transport{Context: c},
		},
	}

	ctx := cloud.NewContext(appengine.AppID(c), hc)

	rc, err := storage.NewReader(ctx, bucket, params["tag"])
	if err != nil {
		log.Warningf(c, "readFile: unable to open file from bucket %q, file %q: %v", bucket, params["tag"], err)
		return
	}
	defer rc.Close()
	slurp, err := ioutil.ReadAll(rc)
	if err != nil {
		log.Warningf(c, "readFile: unable to read data from bucket %q, file %q: %v", bucket, params["tag"], err)
		return
	}
	o, err := storage.StatObject(ctx, bucket, params["tag"])
	if err != nil {
		rw.Header().Add("Content-Type", "image/png")
	} else {
		rw.Header().Add("Content-Type", o.ContentType)
	}

	rw.Write([]byte(slurp))
}

type File struct {
	Name    string
	Content []byte
	Type    string
}

func ExportAllFiles(rw http.ResponseWriter, req *http.Request) (export []File) {
	c := appengine.NewContext(req)

	export = make([]File, 0)

	bucket := ""
	if bucket == "" {
		var err error
		if bucket, err = file.DefaultBucketName(c); err != nil {
			// log.Errorf(c, "failed to get default GCS bucket name: %v", err)
			return
		}
	}

	hc := &http.Client{
		Transport: &oauth2.Transport{
			Source: google.AppEngineTokenSource(c, storage.ScopeFullControl),
			Base:   &urlfetch.Transport{Context: c},
		},
	}

	ctx := cloud.NewContext(appengine.AppID(c), hc)

	query := &storage.Query{Prefix: ""}
	for query != nil {
		objs, err := storage.ListObjects(ctx, bucket, query)
		if err != nil {
			//d.errorf("listBucket: unable to list bucket %q: %v", bucket, err)
			return
		}
		query = objs.Next

		for _, obj := range objs.Results {
			//d.dumpStats(obj)
			newfile := File{}
			newfile.Name = obj.Name
			newfile.Type = obj.ContentType

			rc, err := storage.NewReader(ctx, bucket, obj.Name)
			if err != nil {
				log.Warningf(c, "readFile: unable to open file from bucket %q, file %q: %v", bucket, obj.Name, err)
				return
			}
			defer rc.Close()
			slurp, err := ioutil.ReadAll(rc)

			newfile.Content = slurp
			export = append(export, newfile)
		}
	}

	return export
}

func GimmeDC(rw http.ResponseWriter, req *http.Request) string {
	c := appengine.NewContext(req)
	return appengine.Datacenter(c)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"testing"

	"github.com/docker/docker/pkg/reexec"
	"github.com/go-check/check"
)

func Test(t *testing.T) {
	reexec.Init() // This is required for external graphdriver tests

	if !isLocalDaemon {
		fmt.Println("INFO: Testing against a remote daemon")
	} else {
		fmt.Println("INFO: Testing against a local daemon")
	}

	check.TestingT(t)
}

func init() {
	check.Suite(&DockerSuite{})
}

type DockerSuite struct {
}

func (s *DockerSuite) TearDownTest(c *check.C) {
	unpauseAllContainers()
	deleteAllContainers()
	deleteAllImages()
	deleteAllVolumes()
	deleteAllNetworks()
}

func init() {
	check.Suite(&DockerRegistrySuite{
		ds: &DockerSuite{},
	})
}

type DockerRegistrySuite struct {
	ds  *DockerSuite
	reg *testRegistryV2
	d   *Daemon
}

func (s *DockerRegistrySuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux)
	s.reg = setupRegistry(c, false)
	s.d = NewDaemon(c)
}

func (s *DockerRegistrySuite) TearDownTest(c *check.C) {
	if s.reg != nil {
		s.reg.Close()
	}
	if s.ds != nil {
		s.ds.TearDownTest(c)
	}
	s.d.Stop()
}

func init() {
	check.Suite(&DockerSchema1RegistrySuite{
		ds: &DockerSuite{},
	})
}

type DockerSchema1RegistrySuite struct {
	ds  *DockerSuite
	reg *testRegistryV2
	d   *Daemon
}

func (s *DockerSchema1RegistrySuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux)
	s.reg = setupRegistry(c, true)
	s.d = NewDaemon(c)
}

func (s *DockerSchema1RegistrySuite) TearDownTest(c *check.C) {
	if s.reg != nil {
		s.reg.Close()
	}
	if s.ds != nil {
		s.ds.TearDownTest(c)
	}
	s.d.Stop()
}

func init() {
	check.Suite(&DockerDaemonSuite{
		ds: &DockerSuite{},
	})
}

type DockerDaemonSuite struct {
	ds *DockerSuite
	d  *Daemon
}

func (s *DockerDaemonSuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux)
	s.d = NewDaemon(c)
}

func (s *DockerDaemonSuite) TearDownTest(c *check.C) {
	testRequires(c, DaemonIsLinux)
	s.d.Stop()
	s.ds.TearDownTest(c)
}

func init() {
	check.Suite(&DockerTrustSuite{
		ds: &DockerSuite{},
	})
}

type DockerTrustSuite struct {
	ds  *DockerSuite
	reg *testRegistryV2
	not *testNotary
}

func (s *DockerTrustSuite) SetUpTest(c *check.C) {
	s.reg = setupRegistry(c, false)
	s.not = setupNotary(c)
}

func (s *DockerTrustSuite) TearDownTest(c *check.C) {
	s.reg.Close()
	s.not.Close()
	s.ds.TearDownTest(c)
}
<commit_msg>integration-cli: check nil before Close'ing<commit_after>package main

import (
	"fmt"
	"testing"

	"github.com/docker/docker/pkg/reexec"
	"github.com/go-check/check"
)

func Test(t *testing.T) {
	reexec.Init() // This is required for external graphdriver tests

	if !isLocalDaemon {
		fmt.Println("INFO: Testing against a remote daemon")
	} else {
		fmt.Println("INFO: Testing against a local daemon")
	}

	check.TestingT(t)
}

func init() {
	check.Suite(&DockerSuite{})
}

type DockerSuite struct {
}

func (s *DockerSuite) TearDownTest(c *check.C) {
	unpauseAllContainers()
	deleteAllContainers()
	deleteAllImages()
	deleteAllVolumes()
	deleteAllNetworks()
}

func init() {
	check.Suite(&DockerRegistrySuite{
		ds: &DockerSuite{},
	})
}

type DockerRegistrySuite struct {
	ds  *DockerSuite
	reg *testRegistryV2
	d   *Daemon
}

func (s *DockerRegistrySuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux)
	s.reg = setupRegistry(c, false)
	s.d = NewDaemon(c)
}

func (s *DockerRegistrySuite) TearDownTest(c *check.C) {
	if s.reg != nil {
		s.reg.Close()
	}
	if s.d != nil {
		s.d.Stop()
	}
	s.ds.TearDownTest(c)
}

func init() {
	check.Suite(&DockerSchema1RegistrySuite{
		ds: &DockerSuite{},
	})
}

type DockerSchema1RegistrySuite struct {
	ds  *DockerSuite
	reg *testRegistryV2
	d   *Daemon
}

func (s *DockerSchema1RegistrySuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux)
	s.reg = setupRegistry(c, true)
	s.d = NewDaemon(c)
}

func (s *DockerSchema1RegistrySuite) TearDownTest(c *check.C) {
	if s.reg != nil {
		s.reg.Close()
	}
	if s.d != nil {
		s.d.Stop()
	}
	s.ds.TearDownTest(c)
}

func init() {
	check.Suite(&DockerDaemonSuite{
		ds: &DockerSuite{},
	})
}

type DockerDaemonSuite struct {
	ds *DockerSuite
	d  *Daemon
}

func (s *DockerDaemonSuite) SetUpTest(c *check.C) {
	testRequires(c, DaemonIsLinux)
	s.d = NewDaemon(c)
}

func (s *DockerDaemonSuite) TearDownTest(c *check.C) {
	testRequires(c, DaemonIsLinux)
	if s.d != nil {
		s.d.Stop()
	}
	s.ds.TearDownTest(c)
}

func init() {
	check.Suite(&DockerTrustSuite{
		ds: &DockerSuite{},
	})
}

type DockerTrustSuite struct {
	ds  *DockerSuite
	reg *testRegistryV2
	not *testNotary
}

func (s *DockerTrustSuite) SetUpTest(c *check.C) {
	s.reg = setupRegistry(c, false)
	s.not = setupNotary(c)
}

func (s *DockerTrustSuite) TearDownTest(c *check.C) {
	if s.reg != nil {
		s.reg.Close()
	}
	if s.not != nil {
		s.not.Close()
	}
	s.ds.TearDownTest(c)
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2019 VMware, Inc. All Rights Reserved.
SPDX-License-Identifier: Apache-2.0
*/

package describer

import (
	"context"

	"github.com/pkg/errors"

	"github.com/vmware/octant/pkg/view/component"
)

// Section is a wrapper to combine content from multiple describers.
type Section struct {
	path       string
	title      string
	describers []Describer
}

var _ Describer = (*Section)(nil)

// NewSection creates a Section.
func NewSection(p, title string, describers ...Describer) *Section {
	return &Section{
		path:       p,
		title:      title,
		describers: describers,
	}
}

// Describe generates content.
func (d *Section) Describe(ctx context.Context, prefix, namespace string, options Options) (component.ContentResponse, error) {
	list := component.NewList(d.title, nil)

	for _, child := range d.describers {
		cResponse, err := child.Describe(ctx, prefix, namespace, options)
		if err != nil {
			return EmptyContentResponse, err
		}

		for _, vc := range cResponse.Components {
			if nestedList, ok := vc.(*component.List); ok {
				for i := range nestedList.Config.Items {
					item := nestedList.Config.Items[i]
					if !item.IsEmpty() {
						list.Add(item)
					}
				}
			}
		}
	}

	cr := component.ContentResponse{
		Components: []component.Component{list},
		Title:      component.Title(component.NewText(d.title)),
	}

	return cr, nil
}

// PathFilters returns path filters for the section.
func (d *Section) PathFilters() []PathFilter {
	PathFilters := []PathFilter{
		*NewPathFilter(d.path, d),
	}

	for _, child := range d.describers {
		PathFilters = append(PathFilters, child.PathFilters()...)
	}

	return PathFilters
}

func (d *Section) Reset(ctx context.Context) error {
	for i := range d.describers {
		if err := d.describers[i].Reset(ctx); err != nil {
			return errors.Wrapf(err, "reset describer in section %s", d.path)
		}
	}

	return nil
}
<commit_msg>use less allocations when transversing sections for content<commit_after>/*
Copyright (c) 2019 VMware, Inc. All Rights Reserved.
SPDX-License-Identifier: Apache-2.0
*/

package describer

import (
	"context"

	"github.com/pkg/errors"

	"github.com/vmware/octant/pkg/view/component"
)

// Section is a wrapper to combine content from multiple describers.
type Section struct {
	path       string
	title      string
	describers []Describer
}

var _ Describer = (*Section)(nil)

// NewSection creates a Section.
func NewSection(p, title string, describers ...Describer) *Section {
	return &Section{
		path:       p,
		title:      title,
		describers: describers,
	}
}

// Describe generates content.
func (d *Section) Describe(ctx context.Context, prefix, namespace string, options Options) (component.ContentResponse, error) {
	list := component.NewList(d.title, nil)

	for describerIndex := range d.describers {
		cResponse, err := d.describers[describerIndex].Describe(ctx, prefix, namespace, options)
		if err != nil {
			return EmptyContentResponse, err
		}

		for componentIndex := range cResponse.Components {
			if nestedList, ok := cResponse.Components[componentIndex].(*component.List); ok {
				for itemIndex := range nestedList.Config.Items {
					item := nestedList.Config.Items[itemIndex]
					if !item.IsEmpty() {
						list.Add(item)
					}
				}
			}
		}
	}

	cr := component.ContentResponse{
		Components: []component.Component{list},
		Title:      component.Title(component.NewText(d.title)),
	}

	return cr, nil
}

// PathFilters returns path filters for the section.
func (d *Section) PathFilters() []PathFilter {
	PathFilters := []PathFilter{
		*NewPathFilter(d.path, d),
	}

	for _, child := range d.describers {
		PathFilters = append(PathFilters, child.PathFilters()...)
	}

	return PathFilters
}

func (d *Section) Reset(ctx context.Context) error {
	for i := range d.describers {
		if err := d.describers[i].Reset(ctx); err != nil {
			return errors.Wrapf(err, "reset describer in section %s", d.path)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package driver

import (
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os/exec"
	"regexp"
	"sync"
	"testing"

	"github.com/google/pprof/internal/plugin"
	"github.com/google/pprof/profile"
	"runtime"
)

func TestWebInterface(t *testing.T) {
	if runtime.GOOS == "nacl" {
		t.Skip("test assumes tcp available")
	}

	prof := makeFakeProfile()

	// Custom http server creator
	var server *httptest.Server
	serverCreated := make(chan bool)
	creator := func(a *plugin.HTTPServerArgs) error {
		server = httptest.NewServer(http.HandlerFunc(
			func(w http.ResponseWriter, r *http.Request) {
				if h := a.Handlers[r.URL.Path]; h != nil {
					h.ServeHTTP(w, r)
				}
			}))
		serverCreated <- true
		return nil
	}

	// Start server and wait for it to be initialized
	go serveWebInterface("unused:1234", prof, &plugin.Options{
		Obj:        fakeObjTool{},
		UI:         &stdUI{},
		HTTPServer: creator,
	}, false)
	<-serverCreated
	defer server.Close()

	haveDot := false
	if _, err := exec.LookPath("dot"); err == nil {
		haveDot = true
	}

	type testCase struct {
		path    string
		want    []string
		needDot bool
	}
	testcases := []testCase{
		{"/", []string{"F1", "F2", "F3", "testbin", "cpu"}, true},
		{"/top", []string{`"Name":"F2","InlineLabel":"","Flat":200,"Cum":300,"FlatFormat":"200ms","CumFormat":"300ms"}`}, false},
		{"/source?f=" + url.QueryEscape("F[12]"),
			[]string{"F1", "F2", "300ms +line1"}, false},
		{"/peek?f=" + url.QueryEscape("F[12]"),
			[]string{"300ms.*F1", "200ms.*300ms.*F2"}, false},
		{"/disasm?f=" + url.QueryEscape("F[12]"),
			[]string{"f1:asm", "f2:asm"}, false},
		{"/flamegraph", []string{"File: testbin", "\"n\":\"root\"", "\"n\":\"F1\"", "function tip", "function flameGraph", "function hierarchy"}, false},
	}
	for _, c := range testcases {
		if c.needDot && !haveDot {
			t.Log("skipping", c.path, "since dot (graphviz) does not seem to be installed")
			continue
		}

		res, err := http.Get(server.URL + c.path)
		if err != nil {
			t.Error("could not fetch", c.path, err)
			continue
		}
		data, err := ioutil.ReadAll(res.Body)
		if err != nil {
			t.Error("could not read response", c.path, err)
			continue
		}
		result := string(data)
		for _, w := range c.want {
			if match, _ := regexp.MatchString(w, result); !match {
				t.Errorf("response for %s does not match "+
					"expected pattern '%s'; "+
					"actual result:\n%s", c.path, w, result)
			}
		}
	}

	// Also fetch all the test case URLs in parallel to test thread
	// safety when run under the race detector.
	var wg sync.WaitGroup
	for _, c := range testcases {
		if c.needDot && !haveDot {
			continue
		}
		path := server.URL + c.path
		for count := 0; count < 2; count++ {
			wg.Add(1)
			go func() {
				http.Get(path)
				wg.Done()
			}()
		}
	}
	wg.Wait()
}

// Implement fake object file support.

const addrBase = 0x1000
const fakeSource = "testdata/file1000.src"

type fakeObj struct{}

func (f fakeObj) Close() error    { return nil }
func (f fakeObj) Name() string    { return "testbin" }
func (f fakeObj) Base() uint64    { return 0 }
func (f fakeObj) BuildID() string { return "" }
func (f fakeObj) SourceLine(addr uint64) ([]plugin.Frame, error) {
	return nil, fmt.Errorf("SourceLine unimplemented")
}
func (f fakeObj) Symbols(r *regexp.Regexp, addr uint64) ([]*plugin.Sym, error) {
	return []*plugin.Sym{
		{
			Name: []string{"F1"}, File: fakeSource,
			Start: addrBase, End: addrBase + 10,
		},
		{
			Name: []string{"F2"}, File: fakeSource,
			Start: addrBase + 10, End: addrBase + 20,
		},
		{
			Name: []string{"F3"}, File: fakeSource,
			Start: addrBase + 20, End: addrBase + 30,
		},
	}, nil
}

type fakeObjTool struct{}

func (obj fakeObjTool) Open(file string, start, limit, offset uint64) (plugin.ObjFile, error) {
	return fakeObj{}, nil
}

func (obj fakeObjTool) Disasm(file string, start, end uint64) ([]plugin.Inst, error) {
	return []plugin.Inst{
		{Addr: addrBase + 0, Text: "f1:asm", Function: "F1"},
		{Addr: addrBase + 10, Text: "f2:asm", Function: "F2"},
		{Addr: addrBase + 20, Text: "d3:asm", Function: "F3"},
	}, nil
}

func makeFakeProfile() *profile.Profile {
	// Three functions: F1, F2, F3 with three lines, 11, 22, 33.
	funcs := []*profile.Function{
		{ID: 1, Name: "F1", Filename: fakeSource, StartLine: 3},
		{ID: 2, Name: "F2", Filename: fakeSource, StartLine: 5},
		{ID: 3, Name: "F3", Filename: fakeSource, StartLine: 7},
	}
	lines := []profile.Line{
		{Function: funcs[0], Line: 11},
		{Function: funcs[1], Line: 22},
		{Function: funcs[2], Line: 33},
	}
	mapping := []*profile.Mapping{
		{
			ID:             1,
			Start:          addrBase,
			Limit:          addrBase + 10,
			Offset:         0,
			File:           "testbin",
			HasFunctions:   true,
			HasFilenames:   true,
			HasLineNumbers: true,
		},
	}

	// Three interesting addresses: base+{10,20,30}
	locs := []*profile.Location{
		{ID: 1, Address: addrBase + 10, Line: lines[0:1], Mapping: mapping[0]},
		{ID: 2, Address: addrBase + 20, Line: lines[1:2], Mapping: mapping[0]},
		{ID: 3, Address: addrBase + 30, Line: lines[2:3], Mapping: mapping[0]},
	}

	// Two stack traces.
	return &profile.Profile{
		PeriodType:    &profile.ValueType{Type: "cpu", Unit: "milliseconds"},
		Period:        1,
		DurationNanos: 10e9,
		SampleType: []*profile.ValueType{
			{Type: "cpu", Unit: "milliseconds"},
		},
		Sample: []*profile.Sample{
			{
				Location: []*profile.Location{locs[2], locs[1], locs[0]},
				Value:    []int64{100},
			},
			{
				Location: []*profile.Location{locs[1], locs[0]},
				Value:    []int64{200},
			},
		},
		Location: locs,
		Function: funcs,
		Mapping:  mapping,
	}
}

func TestIsLocalHost(t *testing.T) {
	for _, s := range []string{"localhost:10000", "[::1]:10000", "127.0.0.1:10000"} {
		host, _, err := net.SplitHostPort(s)
		if err != nil {
			t.Error("unexpected error when splitting", s)
			continue
		}
		if !isLocalhost(host) {
			t.Errorf("host %s from %s not considered local", host, s)
		}
	}
}
<commit_msg>Fix import order. (#301)<commit_after>// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package driver

import (
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os/exec"
	"regexp"
	"runtime"
	"sync"
	"testing"

	"github.com/google/pprof/internal/plugin"
	"github.com/google/pprof/profile"
)

func TestWebInterface(t *testing.T) {
	if runtime.GOOS == "nacl" {
		t.Skip("test assumes tcp available")
	}

	prof := makeFakeProfile()

	// Custom http server creator
	var server *httptest.Server
	serverCreated := make(chan bool)
	creator := func(a *plugin.HTTPServerArgs) error {
		server = httptest.NewServer(http.HandlerFunc(
			func(w http.ResponseWriter, r *http.Request) {
				if h := a.Handlers[r.URL.Path]; h != nil {
					h.ServeHTTP(w, r)
				}
			}))
		serverCreated <- true
		return nil
	}

	// Start server and wait for it to be initialized
	go serveWebInterface("unused:1234", prof, &plugin.Options{
		Obj:        fakeObjTool{},
		UI:         &stdUI{},
		HTTPServer: creator,
	}, false)
	<-serverCreated
	defer server.Close()

	haveDot := false
	if _, err := exec.LookPath("dot"); err == nil {
		haveDot = true
	}

	type testCase struct {
		path    string
		want    []string
		needDot bool
	}
	testcases := []testCase{
		{"/", []string{"F1", "F2", "F3", "testbin", "cpu"}, true},
		{"/top", []string{`"Name":"F2","InlineLabel":"","Flat":200,"Cum":300,"FlatFormat":"200ms","CumFormat":"300ms"}`}, false},
		{"/source?f=" + url.QueryEscape("F[12]"),
			[]string{"F1", "F2", "300ms +line1"}, false},
		{"/peek?f=" + url.QueryEscape("F[12]"),
			[]string{"300ms.*F1", "200ms.*300ms.*F2"}, false},
		{"/disasm?f=" + url.QueryEscape("F[12]"),
			[]string{"f1:asm", "f2:asm"}, false},
		{"/flamegraph", []string{"File: testbin", "\"n\":\"root\"", "\"n\":\"F1\"", "function tip", "function flameGraph", "function hierarchy"}, false},
	}
	for _, c := range testcases {
		if c.needDot && !haveDot {
			t.Log("skipping", c.path, "since dot (graphviz) does not seem to be installed")
			continue
		}

		res, err := http.Get(server.URL + c.path)
		if err != nil {
			t.Error("could not fetch", c.path, err)
			continue
		}
		data, err := ioutil.ReadAll(res.Body)
		if err != nil {
			t.Error("could not read response", c.path, err)
			continue
		}
		result := string(data)
		for _, w := range c.want {
			if match, _ := regexp.MatchString(w, result); !match {
				t.Errorf("response for %s does not match "+
					"expected pattern '%s'; "+
					"actual result:\n%s", c.path, w, result)
			}
		}
	}

	// Also fetch all the test case URLs in parallel to test thread
	// safety when run under the race detector.
	var wg sync.WaitGroup
	for _, c := range testcases {
		if c.needDot && !haveDot {
			continue
		}
		path := server.URL + c.path
		for count := 0; count < 2; count++ {
			wg.Add(1)
			go func() {
				http.Get(path)
				wg.Done()
			}()
		}
	}
	wg.Wait()
}

// Implement fake object file support.

const addrBase = 0x1000
const fakeSource = "testdata/file1000.src"

type fakeObj struct{}

func (f fakeObj) Close() error    { return nil }
func (f fakeObj) Name() string    { return "testbin" }
func (f fakeObj) Base() uint64    { return 0 }
func (f fakeObj) BuildID() string { return "" }
func (f fakeObj) SourceLine(addr uint64) ([]plugin.Frame, error) {
	return nil, fmt.Errorf("SourceLine unimplemented")
}
func (f fakeObj) Symbols(r *regexp.Regexp, addr uint64) ([]*plugin.Sym, error) {
	return []*plugin.Sym{
		{
			Name: []string{"F1"}, File: fakeSource,
			Start: addrBase, End: addrBase + 10,
		},
		{
			Name: []string{"F2"}, File: fakeSource,
			Start: addrBase + 10, End: addrBase + 20,
		},
		{
			Name: []string{"F3"}, File: fakeSource,
			Start: addrBase + 20, End: addrBase + 30,
		},
	}, nil
}

type fakeObjTool struct{}

func (obj fakeObjTool) Open(file string, start, limit, offset uint64) (plugin.ObjFile, error) {
	return fakeObj{}, nil
}

func (obj fakeObjTool) Disasm(file string, start, end uint64) ([]plugin.Inst, error) {
	return []plugin.Inst{
		{Addr: addrBase + 0, Text: "f1:asm", Function: "F1"},
		{Addr: addrBase + 10, Text: "f2:asm", Function: "F2"},
		{Addr: addrBase + 20, Text: "d3:asm", Function: "F3"},
	}, nil
}

func makeFakeProfile() *profile.Profile {
	// Three functions: F1, F2, F3 with three lines, 11, 22, 33.
	funcs := []*profile.Function{
		{ID: 1, Name: "F1", Filename: fakeSource, StartLine: 3},
		{ID: 2, Name: "F2", Filename: fakeSource, StartLine: 5},
		{ID: 3, Name: "F3", Filename: fakeSource, StartLine: 7},
	}
	lines := []profile.Line{
		{Function: funcs[0], Line: 11},
		{Function: funcs[1], Line: 22},
		{Function: funcs[2], Line: 33},
	}
	mapping := []*profile.Mapping{
		{
			ID:             1,
			Start:          addrBase,
			Limit:          addrBase + 10,
			Offset:         0,
			File:           "testbin",
			HasFunctions:   true,
			HasFilenames:   true,
			HasLineNumbers: true,
		},
	}

	// Three interesting addresses: base+{10,20,30}
	locs := []*profile.Location{
		{ID: 1, Address: addrBase + 10, Line: lines[0:1], Mapping: mapping[0]},
		{ID: 2, Address: addrBase + 20, Line: lines[1:2], Mapping: mapping[0]},
		{ID: 3, Address: addrBase + 30, Line: lines[2:3], Mapping: mapping[0]},
	}

	// Two stack traces.
	return &profile.Profile{
		PeriodType:    &profile.ValueType{Type: "cpu", Unit: "milliseconds"},
		Period:        1,
		DurationNanos: 10e9,
		SampleType: []*profile.ValueType{
			{Type: "cpu", Unit: "milliseconds"},
		},
		Sample: []*profile.Sample{
			{
				Location: []*profile.Location{locs[2], locs[1], locs[0]},
				Value:    []int64{100},
			},
			{
				Location: []*profile.Location{locs[1], locs[0]},
				Value:    []int64{200},
			},
		},
		Location: locs,
		Function: funcs,
		Mapping:  mapping,
	}
}

func TestIsLocalHost(t *testing.T) {
	for _, s := range []string{"localhost:10000", "[::1]:10000", "127.0.0.1:10000"} {
		host, _, err := net.SplitHostPort(s)
		if err != nil {
			t.Error("unexpected error when splitting", s)
			continue
		}
		if !isLocalhost(host) {
			t.Errorf("host %s from %s not considered local", host, s)
		}
	}
}
<|endoftext|>


<commit_before>package main

//go:generate go run generate.go protocol/input/json.json ../../aws/protocol/jsonrpc/build_test.go
//go:generate go run generate.go protocol/input/query.json ../../aws/protocol/query/build_test.go
//go:generate go run generate.go protocol/output/query.json ../../aws/protocol/query/unmarshal_test.go
//go:generate go run generate.go protocol/input/ec2.json ../../aws/protocol/ec2/build_test.go
//go:generate go run generate.go protocol/output/ec2.json ../../aws/protocol/ec2/unmarshal_test.go

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/url"
	"os"
	"regexp"
	"strconv"
	"strings"
	"text/template"

	"github.com/awslabs/aws-sdk-go/internal/model/api"
	"github.com/awslabs/aws-sdk-go/internal/util"
)

type TestSuite struct {
	*api.API
	Description string
	Cases       []TestCase
	title       string
}

type TestCase struct {
	*TestSuite
	Given      *api.Operation
	Params     interface{}     `json:",omitempty"`
	Data       interface{}     `json:"result,omitempty"`
	InputTest  TestExpectation `json:"serialized"`
	OutputTest TestExpectation `json:"response"`
}

type TestExpectation struct {
	Body       string
	URI        string
	Headers    map[string]string
	StatusCode uint `json:"status_code"`
}

const preamble = `
var _ bytes.Buffer // always import bytes
var _ http.Request
`

var reStripSpace = regexp.MustCompile(`\s(\w)`)

var reImportRemoval = regexp.MustCompile(`(?s:import \((.+?)\))`)

func removeImports(code string) string {
	return reImportRemoval.ReplaceAllString(code, "")
}

var extraImports = []string{
	"bytes",
	"encoding/json",
	"io/ioutil",
	"net/http",
	"testing",
	"github.com/awslabs/aws-sdk-go/internal/util",
	"github.com/stretchr/testify/assert",
}

func addImports(code string) string {
	importNames := make([]string, len(extraImports))
	for i, n := range extraImports {
		importNames[i] = fmt.Sprintf("%q", n)
	}
	str := reImportRemoval.ReplaceAllString(code, "import (\n$1\n"+strings.Join(importNames, "\n")+")")
	return str
}

func (t *TestSuite) TestSuite() string {
	var buf bytes.Buffer

	t.title = reStripSpace.ReplaceAllStringFunc(t.Description, func(x string) string {
		return strings.ToUpper(x[1:])
	})
	t.title = regexp.MustCompile(`\W`).ReplaceAllString(t.title, "")

	for idx, c := range t.Cases {
		c.TestSuite = t
		buf.WriteString(c.TestCase(idx) + "\n")
	}
	return util.GoFmt(buf.String())
}

var tplInputTestCase = template.Must(template.New("inputcase").Parse(`
func Test{{ .OpName }}(t *testing.T) {
	svc := New{{ .TestCase.TestSuite.API.StructName }}(nil)

	var input {{ .Given.InputRef.ShapeName }}
	json.Unmarshal([]byte({{ .ParamsString }}), &input)
	req := svc.{{ .Given.ExportedName }}Request(&input)
	r := req.HTTPRequest

	// build request
	{{ .TestCase.TestSuite.API.ProtocolPackage }}.Build(req)
	assert.NoError(t, req.Error)

	// assert body
	body, _ := ioutil.ReadAll(r.Body)
	assert.Equal(t, util.Trim({{ .Body }}), util.Trim(string(body)))

	// assert URL
	assert.Equal(t, "{{ .TestCase.InputTest.URI }}", r.URL.Path)

	// assert headers
{{ range $k, $v := .TestCase.InputTest.Headers }}assert.Equal(t, "{{ $v }}", r.Header.Get("{{ $k }}"))
{{ end }}
}
`))

type tplInputTestCaseData struct {
	*TestCase
	Body, OpName, ParamsString string
}

var tplOutputTestCase = template.Must(template.New("outputcase").Parse(`
func Test{{ .OpName }}(t *testing.T) {
	svc := New{{ .TestCase.TestSuite.API.StructName }}(nil)

	buf := bytes.NewReader([]byte({{ .Body }}))
	req, _ := svc.{{ .Given.ExportedName }}Request()
	req.HTTPResponse = &http.Response{StatusCode: 200, Body: ioutil.NopCloser(buf)}

	// unmarshal response
	{{ .TestCase.TestSuite.API.ProtocolPackage }}.Unmarshal(req)
	assert.NoError(t, req.Error)

	// assert response
	jBuf, _ := json.Marshal(req.Data)
	assert.Equal(t, util.Trim({{ .ResponseString }}), util.Trim(string(jBuf)))

	// assert headers
{{ range $k, $v := .TestCase.OutputTest.Headers }}assert.Equal(t, "{{ $v }}", req.HTTPResponse.Header.Get("{{ $k }}"))
{{ end }}
}
`))

type tplOutputTestCaseData struct {
	*TestCase
	Body, OpName, ResponseString string
}

func (i *TestCase) TestCase(idx int) string {
	var buf bytes.Buffer

	opName := i.API.StructName() + i.TestSuite.title + "Case" + strconv.Itoa(idx+1)

	if i.Params != nil { // input test
		// query test should sort body as form encoded values
		switch i.API.Metadata.Protocol {
		case "query", "ec2":
			m, _ := url.ParseQuery(i.InputTest.Body)
			i.InputTest.Body = m.Encode()
		}

		pBuf, _ := json.Marshal(i.Params)
		input := tplInputTestCaseData{
			TestCase:     i,
			Body:         fmt.Sprintf("%q", i.InputTest.Body),
			OpName:       strings.ToUpper(opName[0:1]) + opName[1:],
			ParamsString: fmt.Sprintf("%q", pBuf),
		}

		if err := tplInputTestCase.Execute(&buf, input); err != nil {
			panic(err)
		}
	} else {
		pBuf, _ := json.Marshal(i.Data)
		output := tplOutputTestCaseData{
			TestCase:       i,
			Body:           fmt.Sprintf("%q", i.OutputTest.Body),
			OpName:         strings.ToUpper(opName[0:1]) + opName[1:],
			ResponseString: fmt.Sprintf("%q", pBuf),
		}

		if err := tplOutputTestCase.Execute(&buf, output); err != nil {
			panic(err)
		}
	}

	return util.GoFmt(buf.String())
}

func GenerateTestSuite(filename string) string {
	inout := "Input"
	if strings.Contains(filename, "/output/") {
		inout = "Output"
	}

	var suites []TestSuite
	f, err := os.Open(filename)
	if err != nil {
		panic(err)
	}

	err = json.NewDecoder(f).Decode(&suites)
	if err != nil {
		panic(err)
	}

	var buf bytes.Buffer
	buf.WriteString("package " + suites[0].ProtocolPackage() + "_test\n\n")

	var innerBuf bytes.Buffer
	innerBuf.WriteString("//\n// Tests begin here\n//\n\n\n")

	for i, suite := range suites {
		svcPrefix := inout + "Service" + strconv.Itoa(i+1)
		suite.API.Metadata.ServiceAbbreviation = svcPrefix + "ProtocolTest"
		suite.API.Operations = map[string]*api.Operation{}
		for idx, c := range suite.Cases {
			c.Given.ExportedName = svcPrefix + "TestCaseOperation" + strconv.Itoa(idx+1)
			suite.API.Operations[c.Given.ExportedName] = c.Given
		}

		suite.API.NoInflections = true // don't require inflections
		suite.API.Setup()
		suite.API.Metadata.EndpointPrefix = suite.API.PackageName()

		for n, s := range suite.API.Shapes {
			s.Rename(svcPrefix + "TestShape" + n)
		}

		svcCode := addImports(suite.API.ServiceGoCode())
		if i == 0 {
			importMatch := reImportRemoval.FindStringSubmatch(svcCode)
			buf.WriteString(importMatch[0] + "\n\n")
			buf.WriteString(preamble + "\n\n")
		}
		svcCode = removeImports(svcCode)
		svcCode = strings.Replace(svcCode, "func New(", "func New"+suite.API.StructName()+"(", -1)

		buf.WriteString(svcCode + "\n\n")
		buf.WriteString(removeImports(suite.API.APIGoCode()) + "\n\n")
		innerBuf.WriteString(suite.TestSuite() + "\n")
	}

	return util.GoFmt(buf.String() + innerBuf.String())
}

func main() {
	out := GenerateTestSuite(os.Args[1])
	if len(os.Args) == 3 {
		f, err := os.Create(os.Args[2])
		defer f.Close()
		if err != nil {
			panic(err)
		}
		f.WriteString(out + "\n")
	} else {
		fmt.Println(out)
	}
}
<commit_msg>Test generated URL more effectively<commit_after>package main

//go:generate go run generate.go protocol/input/json.json ../../aws/protocol/jsonrpc/build_test.go
//go:generate go run generate.go protocol/input/query.json ../../aws/protocol/query/build_test.go
//go:generate go run generate.go protocol/output/query.json ../../aws/protocol/query/unmarshal_test.go
//go:generate go run generate.go protocol/input/ec2.json ../../aws/protocol/ec2/build_test.go
//go:generate go run generate.go protocol/output/ec2.json ../../aws/protocol/ec2/unmarshal_test.go

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/url"
	"os"
	"regexp"
	"strconv"
	"strings"
	"text/template"

	"github.com/awslabs/aws-sdk-go/internal/model/api"
	"github.com/awslabs/aws-sdk-go/internal/util"
)

type TestSuite struct {
	*api.API
	Description string
	Cases       []TestCase
	title       string
}

type TestCase struct {
	*TestSuite
	Given      *api.Operation
	Params     interface{}     `json:",omitempty"`
	Data       interface{}     `json:"result,omitempty"`
	InputTest  TestExpectation `json:"serialized"`
	OutputTest TestExpectation `json:"response"`
}

type TestExpectation struct {
	Body       string
	URI        string
	Headers    map[string]string
	StatusCode uint `json:"status_code"`
}

const preamble = `
var _ bytes.Buffer // always import bytes
var _ http.Request
`

var reStripSpace = regexp.MustCompile(`\s(\w)`)

var reImportRemoval = regexp.MustCompile(`(?s:import \((.+?)\))`)

func removeImports(code string) string {
	return reImportRemoval.ReplaceAllString(code, "")
}

var extraImports = []string{
	"bytes",
	"encoding/json",
	"io/ioutil",
	"net/http",
	"testing",
	"github.com/awslabs/aws-sdk-go/internal/util",
	"github.com/stretchr/testify/assert",
}

func addImports(code string) string {
	importNames := make([]string, len(extraImports))
	for i, n := range extraImports {
		importNames[i] = fmt.Sprintf("%q", n)
	}
	str := reImportRemoval.ReplaceAllString(code, "import (\n$1\n"+strings.Join(importNames, "\n")+")")
	return str
}

func (t *TestSuite) TestSuite() string {
	var buf bytes.Buffer

	t.title = reStripSpace.ReplaceAllStringFunc(t.Description, func(x string) string {
		return strings.ToUpper(x[1:])
	})
	t.title = regexp.MustCompile(`\W`).ReplaceAllString(t.title, "")

	for idx, c := range t.Cases {
		c.TestSuite = t
		buf.WriteString(c.TestCase(idx) + "\n")
	}
	return util.GoFmt(buf.String())
}

var tplInputTestCase = template.Must(template.New("inputcase").Parse(`
func Test{{ .OpName }}(t *testing.T) {
	svc := New{{ .TestCase.TestSuite.API.StructName }}(nil)
	svc.Endpoint = "https://test"

	var input {{ .Given.InputRef.ShapeName }}
	json.Unmarshal([]byte({{ .ParamsString }}), &input)
	req := svc.{{ .Given.ExportedName }}Request(&input)
	r := req.HTTPRequest

	// build request
	{{ .TestCase.TestSuite.API.ProtocolPackage }}.Build(req)
	assert.NoError(t, req.Error)

	// assert body
	body, _ := ioutil.ReadAll(r.Body)
	assert.Equal(t, util.Trim({{ .Body }}), util.Trim(string(body)))

	// assert URL
	assert.Equal(t, "https://test{{ .TestCase.InputTest.URI }}", r.URL.String())

	// assert headers
{{ range $k, $v := .TestCase.InputTest.Headers }}assert.Equal(t, "{{ $v }}", r.Header.Get("{{ $k }}"))
{{ end }}
}
`))

type tplInputTestCaseData struct {
	*TestCase
	Body, OpName, ParamsString string
}

var tplOutputTestCase = template.Must(template.New("outputcase").Parse(`
func Test{{ .OpName }}(t *testing.T) {
	svc := New{{ .TestCase.TestSuite.API.StructName }}(nil)

	buf := bytes.NewReader([]byte({{ .Body }}))
	req, _ := svc.{{ .Given.ExportedName }}Request()
	req.HTTPResponse = &http.Response{StatusCode: 200, Body: ioutil.NopCloser(buf)}

	// unmarshal response
	{{ .TestCase.TestSuite.API.ProtocolPackage }}.Unmarshal(req)
	assert.NoError(t, req.Error)

	// assert response
	jBuf, _ := json.Marshal(req.Data)
	assert.Equal(t, util.Trim({{ .ResponseString }}), util.Trim(string(jBuf)))

	// assert headers
{{ range $k, $v := .TestCase.OutputTest.Headers }}assert.Equal(t, "{{ $v }}", req.HTTPResponse.Header.Get("{{ $k }}"))
{{ end }}
}
`))

type tplOutputTestCaseData struct {
	*TestCase
	Body, OpName, ResponseString string
}

func (i *TestCase) TestCase(idx int) string {
	var buf bytes.Buffer

	opName := i.API.StructName() + i.TestSuite.title + "Case" + strconv.Itoa(idx+1)

	if i.Params != nil { // input test
		// query test should sort body as form encoded values
		switch i.API.Metadata.Protocol {
		case "query", "ec2":
			m, _ := url.ParseQuery(i.InputTest.Body)
			i.InputTest.Body = m.Encode()
		}

		pBuf, _ := json.Marshal(i.Params)
		input := tplInputTestCaseData{
			TestCase:     i,
			Body:         fmt.Sprintf("%q", i.InputTest.Body),
			OpName:       strings.ToUpper(opName[0:1]) + opName[1:],
			ParamsString: fmt.Sprintf("%q", pBuf),
		}

		if err := tplInputTestCase.Execute(&buf, input); err != nil {
			panic(err)
		}
	} else {
		pBuf, _ := json.Marshal(i.Data)
		output := tplOutputTestCaseData{
			TestCase:       i,
			Body:           fmt.Sprintf("%q", i.OutputTest.Body),
			OpName:         strings.ToUpper(opName[0:1]) + opName[1:],
			ResponseString: fmt.Sprintf("%q", pBuf),
		}

		if err := tplOutputTestCase.Execute(&buf, output); err != nil {
			panic(err)
		}
	}

	return util.GoFmt(buf.String())
}

func GenerateTestSuite(filename string) string {
	inout := "Input"
	if strings.Contains(filename, "/output/") {
		inout = "Output"
	}

	var suites []TestSuite
	f, err := os.Open(filename)
	if err != nil {
		panic(err)
	}

	err = json.NewDecoder(f).Decode(&suites)
	if err != nil {
		panic(err)
	}

	var buf bytes.Buffer
	buf.WriteString("package " + suites[0].ProtocolPackage() + "_test\n\n")

	var innerBuf bytes.Buffer
	innerBuf.WriteString("//\n// Tests begin here\n//\n\n\n")

	for i, suite := range suites {
		svcPrefix := inout + "Service" + strconv.Itoa(i+1)
		suite.API.Metadata.ServiceAbbreviation = svcPrefix + "ProtocolTest"
		suite.API.Operations = map[string]*api.Operation{}
		for idx, c := range suite.Cases {
			c.Given.ExportedName = svcPrefix + "TestCaseOperation" + strconv.Itoa(idx+1)
			suite.API.Operations[c.Given.ExportedName] = c.Given
		}

		suite.API.NoInflections = true // don't require inflections
		suite.API.Setup()
		suite.API.Metadata.EndpointPrefix = suite.API.PackageName()

		for n, s := range suite.API.Shapes {
			s.Rename(svcPrefix + "TestShape" + n)
		}

		svcCode := addImports(suite.API.ServiceGoCode())
		if i == 0 {
			importMatch := reImportRemoval.FindStringSubmatch(svcCode)
			buf.WriteString(importMatch[0] + "\n\n")
			buf.WriteString(preamble + "\n\n")
		}
		svcCode = removeImports(svcCode)
		svcCode = strings.Replace(svcCode, "func New(", "func New"+suite.API.StructName()+"(", -1)

		buf.WriteString(svcCode + "\n\n")
		buf.WriteString(removeImports(suite.API.APIGoCode()) + "\n\n")
		innerBuf.WriteString(suite.TestSuite() + "\n")
	}

	return util.GoFmt(buf.String() + innerBuf.String())
}

func main() {
	out := GenerateTestSuite(os.Args[1])
	if len(os.Args) == 3 {
		f, err := os.Create(os.Args[2])
		defer f.Close()
		if err != nil {
			panic(err)
		}
		f.WriteString(out + "\n")
	} else {
		fmt.Println(out)
	}
}
<|endoftext|>