<commit_before>package aks

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"regexp"
	"sort"

	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2020-06-30/compute"
	"github.com/Azure/azure-sdk-for-go/services/containerservice/mgmt/2020-09-01/containerservice"
	"github.com/Azure/azure-sdk-for-go/services/network/mgmt/2020-07-01/network"
	"github.com/Azure/azure-sdk-for-go/services/subscription/mgmt/2020-09-01/subscription"
	"github.com/Azure/go-autorest/autorest"
	"github.com/Azure/go-autorest/autorest/adal"
	"github.com/Azure/go-autorest/autorest/to"
	"github.com/mcuadros/go-version"
	v3 "github.com/rancher/rancher/pkg/generated/norman/management.cattle.io/v3"
	"github.com/sirupsen/logrus"
)

type virtualNetworksResponseBody struct {
	Name          string   `json:"name"`
	ResourceGroup string   `json:"resourceGroup"`
	Subnets       []subnet `json:"subnets"`
	Location      string   `json:"location"`
}

type subnet struct {
	Name         string `json:"name"`
	AddressRange string `json:"addressRange"`
}

var matchResourceGroup = regexp.MustCompile("/resource[gG]roups/(.+?)/")

func NewClientAuthorizer(cap *Capabilities) (autorest.Authorizer, error) {
	oauthConfig, err := adal.NewOAuthConfig(cap.AuthBaseURL, cap.TenantID)
	if err != nil {
		return nil, err
	}

	spToken, err := adal.NewServicePrincipalToken(*oauthConfig, cap.ClientID, cap.ClientSecret, cap.BaseURL)
	if err != nil {
		return nil, fmt.Errorf("couldn't authenticate to Azure cloud with error: %v", err)
	}

	return autorest.NewBearerAuthorizer(spToken), nil
}

func NewVirtualMachineClient(cap *Capabilities) (*compute.VirtualMachineSizesClient, error) {
	authorizer, err := NewClientAuthorizer(cap)
	if err != nil {
		return nil, err
	}

	virtualMachine := compute.NewVirtualMachineSizesClient(cap.SubscriptionID)
	virtualMachine.Authorizer = authorizer

	return &virtualMachine, nil
}

func NewContainerServiceClient(cap *Capabilities) (*containerservice.ContainerServicesClient, error) {
	authorizer, err := NewClientAuthorizer(cap)
	if err != nil {
		return nil, err
	}

	containerService := containerservice.NewContainerServicesClientWithBaseURI(cap.BaseURL, cap.SubscriptionID)
	containerService.Authorizer = authorizer

	return &containerService, nil
}

func NewNetworkServiceClient(cap *Capabilities) (*network.VirtualNetworksClient, error) {
	authorizer, err := NewClientAuthorizer(cap)
	if err != nil {
		return nil, err
	}

	containerService := network.NewVirtualNetworksClientWithBaseURI(cap.BaseURL, cap.SubscriptionID)
	containerService.Authorizer = authorizer

	return &containerService, nil
}

func NewClusterClient(cap *Capabilities) (*containerservice.ManagedClustersClient, error) {
	authorizer, err := NewClientAuthorizer(cap)
	if err != nil {
		return nil, err
	}

	client := containerservice.NewManagedClustersClientWithBaseURI(cap.BaseURL, cap.SubscriptionID)
	client.Authorizer = authorizer

	return &client, nil
}

func NewSubscriptionServiceClient(cap *Capabilities) (*subscription.SubscriptionsClient, error) {
	authorizer, err := NewClientAuthorizer(cap)
	if err != nil {
		return nil, err
	}

	subscriptionService := subscription.NewSubscriptionsClient()
	subscriptionService.Authorizer = authorizer

	return &subscriptionService, nil
}

type sortableVersion []string

func (s sortableVersion) Len() int {
	return len(s)
}

func (s sortableVersion) Swap(a, b int) {
	s[a], s[b] = s[b], s[a]
}

func (s sortableVersion) Less(a, b int) bool {
	return version.Compare(s[a], s[b], "<")
}

type KubernetesUpgradeVersion struct {
	Version string `json:"version"`
	Enabled bool   `json:"enabled"`
}

type KubernetesUpgradeVersions []*KubernetesUpgradeVersion

func (s KubernetesUpgradeVersions) Len() int {
	return len(s)
}

func (s KubernetesUpgradeVersions) Swap(a, b int) {
	s[a], s[b] = s[b], s[a]
}

func (s KubernetesUpgradeVersions) Less(a, b int) bool {
	return version.Compare(s[a].Version, s[b].Version, "<")
}

type UpgradeVersionsResponse struct {
	CurrentVersion string                    `json:"currentVersion"`
	Upgrades       KubernetesUpgradeVersions `json:"upgrades"`
}

// listKubernetesUpgradeVersions lists all kubernetes versions listed by AKS Container Service and marks which ones the
//given cluster can be upgraded to.  A version's `Enabled` flag is true if the cluster can be upgraded to the version
//in its current state.
func listKubernetesUpgradeVersions(ctx context.Context, clusterLister v3.ClusterLister, cap *Capabilities) ([]byte, int, error) {
	var resp UpgradeVersionsResponse

	// load the target cluster, if the cluster is not found we cannot proceed
	cluster, err := clusterLister.Get("", cap.ClusterID)
	if err != nil {
		return nil, http.StatusBadRequest, fmt.Errorf("invalid cluster id")
	}

	resp.CurrentVersion = *cluster.Spec.AKSConfig.KubernetesVersion

	// get the client for aks container service
	clientContainer, err := NewContainerServiceClient(cap)
	if err != nil {
		return nil, http.StatusInternalServerError, err
	}

	// request a list of orchestrators
	orchestrators, err := clientContainer.ListOrchestrators(ctx, cap.ResourceLocation, "managedClusters")
	if err != nil {
		return nil, http.StatusBadRequest, fmt.Errorf("failed to get orchestrators: %v", err)
	}

	// ensure the orchestrators are returned
	if orchestrators.Orchestrators == nil {
		return nil, http.StatusBadRequest, fmt.Errorf("no version profiles returned: %v", err)
	}

	var upgradeVersions map[string]bool
	for _, profile := range *orchestrators.Orchestrators {
		// check for nil pointers to avoid a panic
		if profile.OrchestratorType == nil || profile.OrchestratorVersion == nil {
			logrus.Warning("unexpected nil orchestrator type or version")
			continue
		}

		// exclude any non kubernetes types
		if containerservice.OrchestratorTypes(*profile.OrchestratorType) != containerservice.Kubernetes {
			continue
		}

		// exclude any versions older than the current version
		if version.Compare(*profile.OrchestratorVersion, resp.CurrentVersion, "<") {
			continue
		}

		// generate the upgrade map when the current version is found
		if *profile.OrchestratorVersion == resp.CurrentVersion {
			upgradeVersions = upgradeableVersionsMap(profile)
			continue
		}

		// store this kubernetes version
		resp.Upgrades = append(resp.Upgrades, &KubernetesUpgradeVersion{Version: *profile.OrchestratorVersion})
	}

	// enable any version listed in the upgrade versions
	for _, kubernetesVersion := range resp.Upgrades {
		if upgradeVersions[kubernetesVersion.Version] {
			kubernetesVersion.Enabled = true
		}
	}

	// sort versions
	sort.Sort(resp.Upgrades)

	return encodeOutput(resp)
}

func listKubernetesVersions(ctx context.Context, cap *Capabilities) ([]byte, int, error) {
	if cap.ResourceLocation == "" {
		return nil, http.StatusBadRequest, fmt.Errorf("region is required")
	}

	clientContainer, err := NewContainerServiceClient(cap)
	if err != nil {
		return nil, http.StatusInternalServerError, err
	}

	orchestrators, err := clientContainer.ListOrchestrators(ctx, cap.ResourceLocation, "managedClusters")
	if err != nil {
		return nil, http.StatusBadRequest, fmt.Errorf("failed to get orchestrators: %v", err)
	}

	if orchestrators.Orchestrators == nil {
		return nil, http.StatusBadRequest, fmt.Errorf("no version profiles returned: %v", err)
	}

	var kubernetesVersions []string

	for _, profile := range *orchestrators.Orchestrators {
		if profile.OrchestratorType == nil || profile.OrchestratorVersion == nil {
			return nil, http.StatusInternalServerError, fmt.Errorf("unexpected nil orchestrator type or version")
		}

		if *profile.OrchestratorType == "Kubernetes" {
			kubernetesVersions = append(kubernetesVersions, *profile.OrchestratorVersion)
		}
	}

	sort.Sort(sortableVersion(kubernetesVersions))

	return encodeOutput(kubernetesVersions)
}

func listVirtualNetworks(ctx context.Context, cap *Capabilities) ([]byte, int, error) {
	clientNetwork, err := NewNetworkServiceClient(cap)
	if err != nil {
		return nil, http.StatusInternalServerError, err
	}

	networkList, err := clientNetwork.ListAll(ctx)
	if err != nil {
		return nil, http.StatusBadRequest, fmt.Errorf("failed to get networks: %v", err)
	}

	var networks []virtualNetworksResponseBody

	for networkList.NotDone() {
		var batch []virtualNetworksResponseBody

		for _, azureNetwork := range networkList.Values() {
			if cap.ResourceLocation != "" && to.String(azureNetwork.Location) != cap.ResourceLocation {
				continue
			}
			var subnets []subnet

			if azureNetwork.Subnets != nil {
				for _, azureSubnet := range *azureNetwork.Subnets {
					if azureSubnet.Name != nil {
						subnets = append(subnets, subnet{
							Name:         to.String(azureSubnet.Name),
							AddressRange: to.String(azureSubnet.AddressPrefix),
						})
					}
				}
			}

			if azureNetwork.ID == nil {
				return nil, http.StatusInternalServerError, fmt.Errorf("no ID on virtual network")
			}

			match := matchResourceGroup.FindStringSubmatch(*azureNetwork.ID)

			if len(match) < 2 || match[1] == "" {
				return nil, http.StatusInternalServerError, fmt.Errorf("could not parse virtual network ID")
			}

			if azureNetwork.Name == nil {
				return nil, http.StatusInternalServerError, fmt.Errorf("no name on virtual network")
			}

			batch = append(batch, virtualNetworksResponseBody{
				Name:          to.String(azureNetwork.Name),
				ResourceGroup: match[1],
				Subnets:       subnets,
				Location:      to.String(azureNetwork.Location),
			})
		}

		networks = append(networks, batch...)

		err = networkList.NextWithContext(ctx)
		if err != nil {
			return nil, http.StatusInternalServerError, err
		}
	}

	return encodeOutput(networks)
}

type clustersResponseBody struct {
	ResourceGroup string `json:"resourceGroup"`
	ClusterName   string `json:"clusterName"`
	RBACEnabled   bool   `json:"rbacEnabled"`
	Location      string `json:"location"`
}

func listClusters(ctx context.Context, cap *Capabilities) ([]byte, int, error) {
	clientCluster, err := NewClusterClient(cap)
	if err != nil {
		return nil, http.StatusInternalServerError, err
	}

	clusterList, err := clientCluster.List(ctx)
	if err != nil {
		return nil, http.StatusBadRequest, fmt.Errorf("failed to get cluster list: %v", err)
	}

	var clusters []clustersResponseBody

	for clusterList.NotDone() {
		for _, cluster := range clusterList.Values() {
			tmpCluster := clustersResponseBody{
				ClusterName: to.String(cluster.Name),
				RBACEnabled: to.Bool(cluster.EnableRBAC),
				Location:    to.String(cluster.Location),
			}
			if cluster.ID != nil {
				match := matchResourceGroup.FindStringSubmatch(to.String(cluster.ID))
				if len(match) < 2 || match[1] == "" {
					return nil, http.StatusInternalServerError, fmt.Errorf("could not parse virtual network ID")
				}
				tmpCluster.ResourceGroup = match[1]
			}
			clusters = append(clusters, tmpCluster)
		}

		err = clusterList.NextWithContext(ctx)
		if err != nil {
			return nil, http.StatusInternalServerError, err
		}
	}

	return encodeOutput(clusters)
}

func listVMSizes(ctx context.Context, cap *Capabilities) ([]byte, int, error) {
	if cap.ResourceLocation == "" {
		return nil, http.StatusBadRequest, fmt.Errorf("region is required")
	}

	virtualMachine, err := NewVirtualMachineClient(cap)
	if err != nil {
		return nil, http.StatusInternalServerError, err
	}
	vmMachineSizeList, err := virtualMachine.List(ctx, cap.ResourceLocation)
	if err != nil {
		return nil, http.StatusBadRequest, fmt.Errorf("failed to get VM sizes: %v", err)
	}

	vmSizes := make([]string, 0, len(*vmMachineSizeList.Value))

	for _, virtualMachineSize := range *vmMachineSizeList.Value {
		vmSizes = append(vmSizes, to.String(virtualMachineSize.Name))
	}

	return encodeOutput(vmSizes)
}

type locationsResponseBody struct {
	Name        string `json:"name"`
	DisplayName string `json:"displayName"`
}

func listLocations(ctx context.Context, cap *Capabilities) ([]byte, int, error) {
	clientSubscription, err := NewSubscriptionServiceClient(cap)
	if err != nil {
		return nil, http.StatusInternalServerError, err
	}
	locationList, err := clientSubscription.ListLocations(ctx, cap.SubscriptionID)
	if err != nil {
		return nil, http.StatusBadRequest, fmt.Errorf("failed to get locations: %v", err)
	}

	var locations []locationsResponseBody

	for _, location := range *locationList.Value {
		locations = append(locations, locationsResponseBody{
			Name:        to.String(location.Name),
			DisplayName: to.String(location.DisplayName),
		})
	}

	return encodeOutput(locations)
}

func encodeOutput(result interface{}) ([]byte, int, error) {
	data, err := json.Marshal(&result)
	if err != nil {
		return data, http.StatusInternalServerError, err
	}

	return data, http.StatusOK, err
}

func upgradeableVersionsMap(upgradeProfile containerservice.OrchestratorVersionProfile) map[string]bool {
	rval := make(map[string]bool, 0)

	for _, profile := range *upgradeProfile.Upgrades {
		rval[*profile.OrchestratorVersion] = true
	}

	return rval
}
<commit_msg>Fix panic for AKS when no upgrades are available<commit_after>package aks

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"regexp"
	"sort"

	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2020-06-30/compute"
	"github.com/Azure/azure-sdk-for-go/services/containerservice/mgmt/2020-09-01/containerservice"
	"github.com/Azure/azure-sdk-for-go/services/network/mgmt/2020-07-01/network"
	"github.com/Azure/azure-sdk-for-go/services/subscription/mgmt/2020-09-01/subscription"
	"github.com/Azure/go-autorest/autorest"
	"github.com/Azure/go-autorest/autorest/adal"
	"github.com/Azure/go-autorest/autorest/to"
	"github.com/mcuadros/go-version"
	v3 "github.com/rancher/rancher/pkg/generated/norman/management.cattle.io/v3"
	"github.com/sirupsen/logrus"
)

type virtualNetworksResponseBody struct {
	Name          string   `json:"name"`
	ResourceGroup string   `json:"resourceGroup"`
	Subnets       []subnet `json:"subnets"`
	Location      string   `json:"location"`
}

type subnet struct {
	Name         string `json:"name"`
	AddressRange string `json:"addressRange"`
}

var matchResourceGroup = regexp.MustCompile("/resource[gG]roups/(.+?)/")

func NewClientAuthorizer(cap *Capabilities) (autorest.Authorizer, error) {
	oauthConfig, err := adal.NewOAuthConfig(cap.AuthBaseURL, cap.TenantID)
	if err != nil {
		return nil, err
	}

	spToken, err := adal.NewServicePrincipalToken(*oauthConfig, cap.ClientID, cap.ClientSecret, cap.BaseURL)
	if err != nil {
		return nil, fmt.Errorf("couldn't authenticate to Azure cloud with error: %v", err)
	}

	return autorest.NewBearerAuthorizer(spToken), nil
}

func NewVirtualMachineClient(cap *Capabilities) (*compute.VirtualMachineSizesClient, error) {
	authorizer, err := NewClientAuthorizer(cap)
	if err != nil {
		return nil, err
	}

	virtualMachine := compute.NewVirtualMachineSizesClient(cap.SubscriptionID)
	virtualMachine.Authorizer = authorizer

	return &virtualMachine, nil
}

func NewContainerServiceClient(cap *Capabilities) (*containerservice.ContainerServicesClient, error) {
	authorizer, err := NewClientAuthorizer(cap)
	if err != nil {
		return nil, err
	}

	containerService := containerservice.NewContainerServicesClientWithBaseURI(cap.BaseURL, cap.SubscriptionID)
	containerService.Authorizer = authorizer

	return &containerService, nil
}

func NewNetworkServiceClient(cap *Capabilities) (*network.VirtualNetworksClient, error) {
	authorizer, err := NewClientAuthorizer(cap)
	if err != nil {
		return nil, err
	}

	containerService := network.NewVirtualNetworksClientWithBaseURI(cap.BaseURL, cap.SubscriptionID)
	containerService.Authorizer = authorizer

	return &containerService, nil
}

func NewClusterClient(cap *Capabilities) (*containerservice.ManagedClustersClient, error) {
	authorizer, err := NewClientAuthorizer(cap)
	if err != nil {
		return nil, err
	}

	client := containerservice.NewManagedClustersClientWithBaseURI(cap.BaseURL, cap.SubscriptionID)
	client.Authorizer = authorizer

	return &client, nil
}

func NewSubscriptionServiceClient(cap *Capabilities) (*subscription.SubscriptionsClient, error) {
	authorizer, err := NewClientAuthorizer(cap)
	if err != nil {
		return nil, err
	}

	subscriptionService := subscription.NewSubscriptionsClient()
	subscriptionService.Authorizer = authorizer

	return &subscriptionService, nil
}

type sortableVersion []string

func (s sortableVersion) Len() int {
	return len(s)
}

func (s sortableVersion) Swap(a, b int) {
	s[a], s[b] = s[b], s[a]
}

func (s sortableVersion) Less(a, b int) bool {
	return version.Compare(s[a], s[b], "<")
}

type KubernetesUpgradeVersion struct {
	Version string `json:"version"`
	Enabled bool   `json:"enabled"`
}

type KubernetesUpgradeVersions []*KubernetesUpgradeVersion

func (s KubernetesUpgradeVersions) Len() int {
	return len(s)
}

func (s KubernetesUpgradeVersions) Swap(a, b int) {
	s[a], s[b] = s[b], s[a]
}

func (s KubernetesUpgradeVersions) Less(a, b int) bool {
	return version.Compare(s[a].Version, s[b].Version, "<")
}

type UpgradeVersionsResponse struct {
	CurrentVersion string                    `json:"currentVersion"`
	Upgrades       KubernetesUpgradeVersions `json:"upgrades"`
}

// listKubernetesUpgradeVersions lists all kubernetes versions listed by AKS Container Service and marks which ones the
//given cluster can be upgraded to.  A version's `Enabled` flag is true if the cluster can be upgraded to the version
//in its current state.
func listKubernetesUpgradeVersions(ctx context.Context, clusterLister v3.ClusterLister, cap *Capabilities) ([]byte, int, error) {
	var resp UpgradeVersionsResponse

	// load the target cluster, if the cluster is not found we cannot proceed
	cluster, err := clusterLister.Get("", cap.ClusterID)
	if err != nil {
		return nil, http.StatusBadRequest, fmt.Errorf("invalid cluster id")
	}

	resp.CurrentVersion = *cluster.Spec.AKSConfig.KubernetesVersion

	// get the client for aks container service
	clientContainer, err := NewContainerServiceClient(cap)
	if err != nil {
		return nil, http.StatusInternalServerError, err
	}

	// request a list of orchestrators
	orchestrators, err := clientContainer.ListOrchestrators(ctx, cap.ResourceLocation, "managedClusters")
	if err != nil {
		return nil, http.StatusBadRequest, fmt.Errorf("failed to get orchestrators: %v", err)
	}

	// ensure the orchestrators are returned
	if orchestrators.Orchestrators == nil {
		return nil, http.StatusBadRequest, fmt.Errorf("no version profiles returned: %v", err)
	}

	var upgradeVersions map[string]bool
	for _, profile := range *orchestrators.Orchestrators {
		// check for nil pointers to avoid a panic
		if profile.OrchestratorType == nil || profile.OrchestratorVersion == nil {
			logrus.Warning("unexpected nil orchestrator type or version")
			continue
		}

		// exclude any non kubernetes types
		if containerservice.OrchestratorTypes(*profile.OrchestratorType) != containerservice.Kubernetes {
			continue
		}

		// exclude any versions older than the current version
		if version.Compare(*profile.OrchestratorVersion, resp.CurrentVersion, "<") {
			continue
		}

		// generate the upgrade map when the current version is found
		if *profile.OrchestratorVersion == resp.CurrentVersion {
			upgradeVersions = upgradeableVersionsMap(profile)
			continue
		}

		// store this kubernetes version
		resp.Upgrades = append(resp.Upgrades, &KubernetesUpgradeVersion{Version: *profile.OrchestratorVersion})
	}

	// enable any version listed in the upgrade versions
	for _, kubernetesVersion := range resp.Upgrades {
		if upgradeVersions[kubernetesVersion.Version] {
			kubernetesVersion.Enabled = true
		}
	}

	// sort versions
	sort.Sort(resp.Upgrades)

	return encodeOutput(resp)
}

func listKubernetesVersions(ctx context.Context, cap *Capabilities) ([]byte, int, error) {
	if cap.ResourceLocation == "" {
		return nil, http.StatusBadRequest, fmt.Errorf("region is required")
	}

	clientContainer, err := NewContainerServiceClient(cap)
	if err != nil {
		return nil, http.StatusInternalServerError, err
	}

	orchestrators, err := clientContainer.ListOrchestrators(ctx, cap.ResourceLocation, "managedClusters")
	if err != nil {
		return nil, http.StatusBadRequest, fmt.Errorf("failed to get orchestrators: %v", err)
	}

	if orchestrators.Orchestrators == nil {
		return nil, http.StatusBadRequest, fmt.Errorf("no version profiles returned: %v", err)
	}

	var kubernetesVersions []string

	for _, profile := range *orchestrators.Orchestrators {
		if profile.OrchestratorType == nil || profile.OrchestratorVersion == nil {
			return nil, http.StatusInternalServerError, fmt.Errorf("unexpected nil orchestrator type or version")
		}

		if *profile.OrchestratorType == "Kubernetes" {
			kubernetesVersions = append(kubernetesVersions, *profile.OrchestratorVersion)
		}
	}

	sort.Sort(sortableVersion(kubernetesVersions))

	return encodeOutput(kubernetesVersions)
}

func listVirtualNetworks(ctx context.Context, cap *Capabilities) ([]byte, int, error) {
	clientNetwork, err := NewNetworkServiceClient(cap)
	if err != nil {
		return nil, http.StatusInternalServerError, err
	}

	networkList, err := clientNetwork.ListAll(ctx)
	if err != nil {
		return nil, http.StatusBadRequest, fmt.Errorf("failed to get networks: %v", err)
	}

	var networks []virtualNetworksResponseBody

	for networkList.NotDone() {
		var batch []virtualNetworksResponseBody

		for _, azureNetwork := range networkList.Values() {
			if cap.ResourceLocation != "" && to.String(azureNetwork.Location) != cap.ResourceLocation {
				continue
			}
			var subnets []subnet

			if azureNetwork.Subnets != nil {
				for _, azureSubnet := range *azureNetwork.Subnets {
					if azureSubnet.Name != nil {
						subnets = append(subnets, subnet{
							Name:         to.String(azureSubnet.Name),
							AddressRange: to.String(azureSubnet.AddressPrefix),
						})
					}
				}
			}

			if azureNetwork.ID == nil {
				return nil, http.StatusInternalServerError, fmt.Errorf("no ID on virtual network")
			}

			match := matchResourceGroup.FindStringSubmatch(*azureNetwork.ID)

			if len(match) < 2 || match[1] == "" {
				return nil, http.StatusInternalServerError, fmt.Errorf("could not parse virtual network ID")
			}

			if azureNetwork.Name == nil {
				return nil, http.StatusInternalServerError, fmt.Errorf("no name on virtual network")
			}

			batch = append(batch, virtualNetworksResponseBody{
				Name:          to.String(azureNetwork.Name),
				ResourceGroup: match[1],
				Subnets:       subnets,
				Location:      to.String(azureNetwork.Location),
			})
		}

		networks = append(networks, batch...)

		err = networkList.NextWithContext(ctx)
		if err != nil {
			return nil, http.StatusInternalServerError, err
		}
	}

	return encodeOutput(networks)
}

type clustersResponseBody struct {
	ResourceGroup string `json:"resourceGroup"`
	ClusterName   string `json:"clusterName"`
	RBACEnabled   bool   `json:"rbacEnabled"`
	Location      string `json:"location"`
}

func listClusters(ctx context.Context, cap *Capabilities) ([]byte, int, error) {
	clientCluster, err := NewClusterClient(cap)
	if err != nil {
		return nil, http.StatusInternalServerError, err
	}

	clusterList, err := clientCluster.List(ctx)
	if err != nil {
		return nil, http.StatusBadRequest, fmt.Errorf("failed to get cluster list: %v", err)
	}

	var clusters []clustersResponseBody

	for clusterList.NotDone() {
		for _, cluster := range clusterList.Values() {
			tmpCluster := clustersResponseBody{
				ClusterName: to.String(cluster.Name),
				RBACEnabled: to.Bool(cluster.EnableRBAC),
				Location:    to.String(cluster.Location),
			}
			if cluster.ID != nil {
				match := matchResourceGroup.FindStringSubmatch(to.String(cluster.ID))
				if len(match) < 2 || match[1] == "" {
					return nil, http.StatusInternalServerError, fmt.Errorf("could not parse virtual network ID")
				}
				tmpCluster.ResourceGroup = match[1]
			}
			clusters = append(clusters, tmpCluster)
		}

		err = clusterList.NextWithContext(ctx)
		if err != nil {
			return nil, http.StatusInternalServerError, err
		}
	}

	return encodeOutput(clusters)
}

func listVMSizes(ctx context.Context, cap *Capabilities) ([]byte, int, error) {
	if cap.ResourceLocation == "" {
		return nil, http.StatusBadRequest, fmt.Errorf("region is required")
	}

	virtualMachine, err := NewVirtualMachineClient(cap)
	if err != nil {
		return nil, http.StatusInternalServerError, err
	}
	vmMachineSizeList, err := virtualMachine.List(ctx, cap.ResourceLocation)
	if err != nil {
		return nil, http.StatusBadRequest, fmt.Errorf("failed to get VM sizes: %v", err)
	}

	vmSizes := make([]string, 0, len(*vmMachineSizeList.Value))

	for _, virtualMachineSize := range *vmMachineSizeList.Value {
		vmSizes = append(vmSizes, to.String(virtualMachineSize.Name))
	}

	return encodeOutput(vmSizes)
}

type locationsResponseBody struct {
	Name        string `json:"name"`
	DisplayName string `json:"displayName"`
}

func listLocations(ctx context.Context, cap *Capabilities) ([]byte, int, error) {
	clientSubscription, err := NewSubscriptionServiceClient(cap)
	if err != nil {
		return nil, http.StatusInternalServerError, err
	}
	locationList, err := clientSubscription.ListLocations(ctx, cap.SubscriptionID)
	if err != nil {
		return nil, http.StatusBadRequest, fmt.Errorf("failed to get locations: %v", err)
	}

	var locations []locationsResponseBody

	for _, location := range *locationList.Value {
		locations = append(locations, locationsResponseBody{
			Name:        to.String(location.Name),
			DisplayName: to.String(location.DisplayName),
		})
	}

	return encodeOutput(locations)
}

func encodeOutput(result interface{}) ([]byte, int, error) {
	data, err := json.Marshal(&result)
	if err != nil {
		return data, http.StatusInternalServerError, err
	}

	return data, http.StatusOK, err
}

func upgradeableVersionsMap(upgradeProfile containerservice.OrchestratorVersionProfile) map[string]bool {
	rval := make(map[string]bool, 0)

	if upgradeProfile.Upgrades == nil {
		// already on latest version, no upgrades available
		return rval
	}
	for _, profile := range *upgradeProfile.Upgrades {
		rval[*profile.OrchestratorVersion] = true
	}

	return rval
}
<|endoftext|>


<commit_before>package class

import (
	cf "jvmgo/classfile"
)

type ConstantMethodref struct {
	className  string
	name       string
	descriptor string
	argCount   uint
	cp         *ConstantPool
	method     *Method
	vslot      int
}

func newConstantMethodref(cp *ConstantPool, methodrefInfo *cf.ConstantMethodrefInfo) *ConstantMethodref {
	return &ConstantMethodref{
		className:  methodrefInfo.ClassName(),
		name:       methodrefInfo.Name(),
		descriptor: methodrefInfo.Descriptor(),
		argCount:   calcArgCount(methodrefInfo.Descriptor()),
		cp:         cp,
		vslot:      -1,
	}
}

func (self *ConstantMethodref) ArgCount() uint {
	return self.argCount
}

func (self *ConstantMethodref) StaticMethod() *Method {
	if self.method == nil {
		self.resolveStaticMethod()
	}
	return self.method
}
func (self *ConstantMethodref) resolveStaticMethod() {
	method := self.findMethod(true)
	if method != nil {
		self.method = method
	} else {
		// todo
		panic("static method not found!")
	}
}

func (self *ConstantMethodref) SpecialMethod() *Method {
	if self.method == nil {
		self.resolveSpecialMethod()
	}
	return self.method
}
func (self *ConstantMethodref) resolveSpecialMethod() {
	method := self.findMethod(false)
	if method != nil {
		self.method = method
		return
	}

	// todo
	// class := self.cp.class.classLoader.LoadClass(self.className)
	// if class.IsInterface() {
	// 	method = self.findMethodInInterfaces(class)
	// 	if method != nil {
	// 		self.method = method
	// 		return
	// 	}
	// }

	// todo
	println("@@@::" + self.className + "::" + self.name + "::" + self.descriptor)
	panic("special method not found!")
}

func (self *ConstantMethodref) findMethod(isStatic bool) *Method {
	class := self.cp.class.classLoader.LoadClass(self.className)
	return class.getMethod(self.name, self.descriptor, isStatic)
}

// todo
func (self *ConstantMethodref) findMethodInInterfaces(iface *Class) *Method {
	for _, m := range iface.methods {
		if !m.IsAbstract() {
			if m.name == self.name && m.descriptor == self.descriptor {
				return m
			}
		}
	}

	for _, superIface := range iface.interfaces {
		if m := self.findMethodInInterfaces(superIface); m != nil {
			return m
		}
	}

	return nil
}

func (self *ConstantMethodref) GetVirtualMethod(ref *Obj) *Method {
	if self.vslot < 0 {
		self.vslot = getVslot(ref.class, self.name, self.descriptor)
	}
	return ref.class.vtable[self.vslot]
}

// todo
func (self *ConstantMethodref) FindInterfaceMethod(ref *Obj) *Method {
	for class := ref.class; class != nil; class = class.superClass {
		method := class.getMethod(self.name, self.descriptor, false)
		if method != nil {
			return method
		}
	}

	// todo
	panic("virtual method not found!")
}

type ConstantInterfaceMethodref struct {
	ConstantMethodref
}

// todo
func newConstantInterfaceMethodref(cp *ConstantPool, methodrefInfo *cf.ConstantInterfaceMethodrefInfo) *ConstantInterfaceMethodref {
	methodref := &ConstantInterfaceMethodref{}
	methodref.cp = cp
	methodref.className = methodrefInfo.ClassName()
	methodref.name = methodrefInfo.Name()
	methodref.descriptor = methodrefInfo.Descriptor()
	methodref.argCount = calcArgCount(methodref.descriptor)
	return methodref
}
<commit_msg>code refactor<commit_after>package class

import (
	cf "jvmgo/classfile"
)

type ConstantMethodref struct {
	className  string
	name       string
	descriptor string
	argCount   uint
	cp         *ConstantPool
	method     *Method
	vslot      int
}

func newConstantMethodref(cp *ConstantPool, methodrefInfo *cf.ConstantMethodrefInfo) *ConstantMethodref {
	return &ConstantMethodref{
		className:  methodrefInfo.ClassName(),
		name:       methodrefInfo.Name(),
		descriptor: methodrefInfo.Descriptor(),
		argCount:   calcArgCount(methodrefInfo.Descriptor()),
		cp:         cp,
		vslot:      -1,
	}
}

func (self *ConstantMethodref) ArgCount() uint {
	return self.argCount
}

func (self *ConstantMethodref) StaticMethod() *Method {
	if self.method == nil {
		self.resolveStaticMethod()
	}
	return self.method
}
func (self *ConstantMethodref) resolveStaticMethod() {
	method := self.findMethod(true)
	if method != nil {
		self.method = method
	} else {
		// todo
		panic("static method not found!")
	}
}

func (self *ConstantMethodref) SpecialMethod() *Method {
	if self.method == nil {
		self.resolveSpecialMethod()
	}
	return self.method
}
func (self *ConstantMethodref) resolveSpecialMethod() {
	method := self.findMethod(false)
	if method != nil {
		self.method = method
		return
	}

	// todo
	// class := self.cp.class.classLoader.LoadClass(self.className)
	// if class.IsInterface() {
	// 	method = self.findMethodInInterfaces(class)
	// 	if method != nil {
	// 		self.method = method
	// 		return
	// 	}
	// }

	// todo
	println("@@@::" + self.className + "::" + self.name + "::" + self.descriptor)
	panic("special method not found!")
}

func (self *ConstantMethodref) findMethod(isStatic bool) *Method {
	class := self.cp.class.classLoader.LoadClass(self.className)
	return class.getMethod(self.name, self.descriptor, isStatic)
}

// todo
func (self *ConstantMethodref) findMethodInInterfaces(iface *Class) *Method {
	for _, m := range iface.methods {
		if !m.IsAbstract() {
			if m.name == self.name && m.descriptor == self.descriptor {
				return m
			}
		}
	}

	for _, superIface := range iface.interfaces {
		if m := self.findMethodInInterfaces(superIface); m != nil {
			return m
		}
	}

	return nil
}

func (self *ConstantMethodref) GetVirtualMethod(ref *Obj) *Method {
	if self.vslot < 0 {
		self.vslot = getVslot(ref.class, self.name, self.descriptor)
	}
	return ref.class.vtable[self.vslot]
}

// todo
func (self *ConstantMethodref) FindInterfaceMethod(ref *Obj) *Method {
	for class := ref.class; class != nil; class = class.superClass {
		method := class.getMethod(self.name, self.descriptor, false)
		if method != nil {
			return method
		}
	}

	// todo
	panic("virtual method not found!")
}

type ConstantInterfaceMethodref struct {
	ConstantMethodref
}

// todo
func newConstantInterfaceMethodref(cp *ConstantPool, methodrefInfo *cf.ConstantInterfaceMethodrefInfo) *ConstantInterfaceMethodref {
	return &ConstantInterfaceMethodref{
		ConstantMethodref{
			className:  methodrefInfo.ClassName(),
			name:       methodrefInfo.Name(),
			descriptor: methodrefInfo.Descriptor(),
			argCount:   calcArgCount(methodrefInfo.Descriptor()),
			cp:         cp,
			vslot:      -1,
		},
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Knative Authors

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

        https://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 assert

import (
	"fmt"
	"strings"

	cloudevents "github.com/cloudevents/sdk-go/v2"
	cetest "github.com/cloudevents/sdk-go/v2/test"

	pkgeventshub "knative.dev/reconciler-test/pkg/eventshub"
	"knative.dev/reconciler-test/pkg/test_images/eventshub"
)

// Matcher that never fails
func Any() pkgeventshub.EventInfoMatcher {
	return func(ei eventshub.EventInfo) error {
		return nil
	}
}

// Matcher that fails if there is an error in the EventInfo
func NoError() pkgeventshub.EventInfoMatcher {
	return func(ei eventshub.EventInfo) error {
		if ei.Error != "" {
			return fmt.Errorf("not expecting an error in event info: %s", ei.Error)
		}
		return nil
	}
}

// Convert a matcher that checks valid messages to a function
// that checks EventInfo structures, returning an error for any that don't
// contain valid events.
func MatchEvent(evf ...cetest.EventMatcher) pkgeventshub.EventInfoMatcher {
	return func(ei eventshub.EventInfo) error {
		if ei.Event == nil {
			return fmt.Errorf("Saw nil event")
		} else {
			return cetest.AllOf(evf...)(*ei.Event)
		}
	}
}

// Convert a matcher that checks valid messages to a function
// that checks EventInfo structures, returning an error for any that don't
// contain valid events.
func HasAdditionalHeader(key, value string) pkgeventshub.EventInfoMatcher {
	key = strings.ToLower(key)
	return func(ei eventshub.EventInfo) error {
		for k, v := range ei.HTTPHeaders {
			if strings.ToLower(k) == key && v[0] == value {
				return nil
			}
		}
		return fmt.Errorf("cannot find header '%s' = '%s' between the headers", key, value)
	}
}

// Reexport kinds here to simplify the usage
const (
	EventReceived = eventshub.EventReceived
	EventRejected = eventshub.EventRejected

	EventSent     = eventshub.EventSent
	EventResponse = eventshub.EventResponse
)

// MatchKind matches the kind of EventInfo
func MatchKind(kind eventshub.EventKind) pkgeventshub.EventInfoMatcher {
	return func(info eventshub.EventInfo) error {
		if kind != info.Kind {
			return fmt.Errorf("event kind don't match. Expected: '%s', Actual: '%s'", kind, info.Kind)
		}
		return nil
	}
}

// MatchStatusCode matches the status code of EventInfo
func MatchStatusCode(statusCode int) pkgeventshub.EventInfoMatcher {
	return func(info eventshub.EventInfo) error {
		if info.StatusCode != statusCode {
			return fmt.Errorf("event status code don't match. Expected: '%d', Actual: '%d'", statusCode, info.StatusCode)
		}
		return nil
	}
}

// MatchHeartBeatsImageMessage matches that the data field of the event, in the format of the heartbeats image, contains the following msg field
func MatchHeartBeatsImageMessage(expectedMsg string) cetest.EventMatcher {
	return cetest.AllOf(
		cetest.HasDataContentType(cloudevents.ApplicationJSON),
		func(have cloudevents.Event) error {
			var m map[string]interface{}
			err := have.DataAs(&m)
			if err != nil {
				return fmt.Errorf("cannot parse heartbeats message %s", err.Error())
			}
			if m["msg"].(string) != expectedMsg {
				return fmt.Errorf("heartbeats message don't match. Expected: '%s', Actual: '%s'", expectedMsg, m["msg"].(string))
			}
			return nil
		},
	)
}
<commit_msg>Re-exporting EventInfo (#87)<commit_after>/*
Copyright 2020 The Knative Authors

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

        https://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 assert

import (
	"fmt"
	"strings"

	cloudevents "github.com/cloudevents/sdk-go/v2"
	cetest "github.com/cloudevents/sdk-go/v2/test"

	pkgeventshub "knative.dev/reconciler-test/pkg/eventshub"
	"knative.dev/reconciler-test/pkg/test_images/eventshub"
)

// EventInfo is a re-export of eventshub.EventInfo for convenience
type EventInfo = eventshub.EventInfo

// Matcher that never fails
func Any() pkgeventshub.EventInfoMatcher {
	return func(ei eventshub.EventInfo) error {
		return nil
	}
}

// Matcher that fails if there is an error in the EventInfo
func NoError() pkgeventshub.EventInfoMatcher {
	return func(ei eventshub.EventInfo) error {
		if ei.Error != "" {
			return fmt.Errorf("not expecting an error in event info: %s", ei.Error)
		}
		return nil
	}
}

// Convert a matcher that checks valid messages to a function
// that checks EventInfo structures, returning an error for any that don't
// contain valid events.
func MatchEvent(evf ...cetest.EventMatcher) pkgeventshub.EventInfoMatcher {
	return func(ei eventshub.EventInfo) error {
		if ei.Event == nil {
			return fmt.Errorf("Saw nil event")
		} else {
			return cetest.AllOf(evf...)(*ei.Event)
		}
	}
}

// Convert a matcher that checks valid messages to a function
// that checks EventInfo structures, returning an error for any that don't
// contain valid events.
func HasAdditionalHeader(key, value string) pkgeventshub.EventInfoMatcher {
	key = strings.ToLower(key)
	return func(ei eventshub.EventInfo) error {
		for k, v := range ei.HTTPHeaders {
			if strings.ToLower(k) == key && v[0] == value {
				return nil
			}
		}
		return fmt.Errorf("cannot find header '%s' = '%s' between the headers", key, value)
	}
}

// Reexport kinds here to simplify the usage
const (
	EventReceived = eventshub.EventReceived
	EventRejected = eventshub.EventRejected

	EventSent     = eventshub.EventSent
	EventResponse = eventshub.EventResponse
)

// MatchKind matches the kind of EventInfo
func MatchKind(kind eventshub.EventKind) pkgeventshub.EventInfoMatcher {
	return func(info eventshub.EventInfo) error {
		if kind != info.Kind {
			return fmt.Errorf("event kind don't match. Expected: '%s', Actual: '%s'", kind, info.Kind)
		}
		return nil
	}
}

// MatchStatusCode matches the status code of EventInfo
func MatchStatusCode(statusCode int) pkgeventshub.EventInfoMatcher {
	return func(info eventshub.EventInfo) error {
		if info.StatusCode != statusCode {
			return fmt.Errorf("event status code don't match. Expected: '%d', Actual: '%d'", statusCode, info.StatusCode)
		}
		return nil
	}
}

// MatchHeartBeatsImageMessage matches that the data field of the event, in the format of the heartbeats image, contains the following msg field
func MatchHeartBeatsImageMessage(expectedMsg string) cetest.EventMatcher {
	return cetest.AllOf(
		cetest.HasDataContentType(cloudevents.ApplicationJSON),
		func(have cloudevents.Event) error {
			var m map[string]interface{}
			err := have.DataAs(&m)
			if err != nil {
				return fmt.Errorf("cannot parse heartbeats message %s", err.Error())
			}
			if m["msg"].(string) != expectedMsg {
				return fmt.Errorf("heartbeats message don't match. Expected: '%s', Actual: '%s'", expectedMsg, m["msg"].(string))
			}
			return nil
		},
	)
}
<|endoftext|>


<commit_before>// Package dht implements a distributed hash table that satisfies the ipfs routing
// interface. This DHT is modeled after Kademlia with S/Kademlia modifications.
package dht

import (
	"context"
	"crypto/rand"
	"fmt"
	"time"

	u "github.com/ipfs/go-ipfs-util"
	goprocess "github.com/jbenet/goprocess"
	periodicproc "github.com/jbenet/goprocess/periodic"
	peer "github.com/libp2p/go-libp2p-peer"
	routing "github.com/libp2p/go-libp2p-routing"
)

// BootstrapConfig specifies parameters used bootstrapping the DHT.
//
// Note there is a tradeoff between the bootstrap period and the
// number of queries. We could support a higher period with less
// queries.
type BootstrapConfig struct {
	Queries int           // how many queries to run per period
	Period  time.Duration // how often to run periodi cbootstrap.
	Timeout time.Duration // how long to wait for a bootstrao query to run
}

var DefaultBootstrapConfig = BootstrapConfig{
	// For now, this is set to 1 query.
	// We are currently more interested in ensuring we have a properly formed
	// DHT than making sure our dht minimizes traffic. Once we are more certain
	// of our implementation's robustness, we should lower this down to 8 or 4.
	Queries: 1,

	// For now, this is set to 1 minute, which is a medium period. We are
	// We are currently more interested in ensuring we have a properly formed
	// DHT than making sure our dht minimizes traffic.
	Period: time.Duration(5 * time.Minute),

	Timeout: time.Duration(10 * time.Second),
}

// Bootstrap ensures the dht routing table remains healthy as peers come and go.
// it builds up a list of peers by requesting random peer IDs. The Bootstrap
// process will run a number of queries each time, and run every time signal fires.
// These parameters are configurable.
//
// As opposed to BootstrapWithConfig, Bootstrap satisfies the routing interface
func (dht *IpfsDHT) Bootstrap(ctx context.Context) error {
	proc, err := dht.BootstrapWithConfig(DefaultBootstrapConfig)
	if err != nil {
		return err
	}

	// wait till ctx or dht.Context exits.
	// we have to do it this way to satisfy the Routing interface (contexts)
	go func() {
		defer proc.Close()
		select {
		case <-ctx.Done():
		case <-dht.Context().Done():
		}
	}()

	return nil
}

// BootstrapWithConfig ensures the dht routing table remains healthy as peers come and go.
// it builds up a list of peers by requesting random peer IDs. The Bootstrap
// process will run a number of queries each time, and run every time signal fires.
// These parameters are configurable.
//
// BootstrapWithConfig returns a process, so the user can stop it.
func (dht *IpfsDHT) BootstrapWithConfig(cfg BootstrapConfig) (goprocess.Process, error) {
	if cfg.Queries <= 0 {
		return nil, fmt.Errorf("invalid number of queries: %d", cfg.Queries)
	}

	tickch := make(chan struct{}, 1)
	proc := periodicproc.OnSignal(tickch, dht.bootstrapWorker(cfg))
	go func() {
		tickch <- struct{}{}
		for {
			select {
			case <-time.After(cfg.Period):
				select {
				case tickch <- struct{}{}:
				default:
					// Don't queue ticks, like Tickers
					log.Warning("Previous bootstrapping attempt not completed within bootstrapping period")
					continue
				}
			case <-dht.Context().Done():
				return
			}
		}
	}()

	return proc, nil
}

// SignalBootstrap ensures the dht routing table remains healthy as peers come and go.
// it builds up a list of peers by requesting random peer IDs. The Bootstrap
// process will run a number of queries each time, and run every time signal fires.
// These parameters are configurable.
//
// SignalBootstrap returns a process, so the user can stop it.
func (dht *IpfsDHT) BootstrapOnSignal(cfg BootstrapConfig, signal <-chan time.Time) (goprocess.Process, error) {
	if cfg.Queries <= 0 {
		return nil, fmt.Errorf("invalid number of queries: %d", cfg.Queries)
	}

	if signal == nil {
		return nil, fmt.Errorf("invalid signal: %v", signal)
	}

	proc := periodicproc.Ticker(signal, dht.bootstrapWorker(cfg))

	return proc, nil
}

func (dht *IpfsDHT) bootstrapWorker(cfg BootstrapConfig) func(worker goprocess.Process) {
	return func(worker goprocess.Process) {
		// it would be useful to be able to send out signals of when we bootstrap, too...
		// maybe this is a good case for whole module event pub/sub?

		ctx := dht.Context()
		if err := dht.runBootstrap(ctx, cfg); err != nil {
			log.Warning(err)
			// A bootstrapping error is important to notice but not fatal.
		}
	}
}

// runBootstrap builds up list of peers by requesting random peer IDs
func (dht *IpfsDHT) runBootstrap(ctx context.Context, cfg BootstrapConfig) error {
	bslog := func(msg string) {
		log.Debugf("DHT %s dhtRunBootstrap %s -- routing table size: %d", dht.self, msg, dht.routingTable.Size())
	}
	bslog("start")
	defer bslog("end")
	defer log.EventBegin(ctx, "dhtRunBootstrap").Done()

	var merr u.MultiErr

	randomID := func() peer.ID {
		// 16 random bytes is not a valid peer id. it may be fine becuase
		// the dht will rehash to its own keyspace anyway.
		id := make([]byte, 16)
		rand.Read(id)
		id = u.Hash(id)
		return peer.ID(id)
	}

	// bootstrap sequentially, as results will compound
	runQuery := func(ctx context.Context, id peer.ID) {
		ctx, cancel := context.WithTimeout(ctx, cfg.Timeout)
		defer cancel()

		p, err := dht.FindPeer(ctx, id)
		if err == routing.ErrNotFound {
			// this isn't an error. this is precisely what we expect.
		} else if err != nil {
			merr = append(merr, err)
		} else {
			// woah, actually found a peer with that ID? this shouldn't happen normally
			// (as the ID we use is not a real ID). this is an odd error worth logging.
			err := fmt.Errorf("Bootstrap peer error: Actually FOUND peer. (%s, %s)", id, p)
			log.Warningf("%s", err)
			merr = append(merr, err)
		}
	}

	// these should be parallel normally. but can make them sequential for debugging.
	// note that the core/bootstrap context deadline should be extended too for that.
	for i := 0; i < cfg.Queries; i++ {
		id := randomID()
		log.Debugf("Bootstrapping query (%d/%d) to random ID: %s", i+1, cfg.Queries, id)
		runQuery(ctx, id)
	}

	// Find self to distribute peer info to our neighbors.
	// Do this after bootstrapping.
	log.Debugf("Bootstrapping query to self: %s", dht.self)
	runQuery(ctx, dht.self)

	if len(merr) > 0 {
		return merr
	}
	return nil
}
<commit_msg>Remove redundant continue<commit_after>// Package dht implements a distributed hash table that satisfies the ipfs routing
// interface. This DHT is modeled after Kademlia with S/Kademlia modifications.
package dht

import (
	"context"
	"crypto/rand"
	"fmt"
	"time"

	u "github.com/ipfs/go-ipfs-util"
	goprocess "github.com/jbenet/goprocess"
	periodicproc "github.com/jbenet/goprocess/periodic"
	peer "github.com/libp2p/go-libp2p-peer"
	routing "github.com/libp2p/go-libp2p-routing"
)

// BootstrapConfig specifies parameters used bootstrapping the DHT.
//
// Note there is a tradeoff between the bootstrap period and the
// number of queries. We could support a higher period with less
// queries.
type BootstrapConfig struct {
	Queries int           // how many queries to run per period
	Period  time.Duration // how often to run periodi cbootstrap.
	Timeout time.Duration // how long to wait for a bootstrao query to run
}

var DefaultBootstrapConfig = BootstrapConfig{
	// For now, this is set to 1 query.
	// We are currently more interested in ensuring we have a properly formed
	// DHT than making sure our dht minimizes traffic. Once we are more certain
	// of our implementation's robustness, we should lower this down to 8 or 4.
	Queries: 1,

	// For now, this is set to 1 minute, which is a medium period. We are
	// We are currently more interested in ensuring we have a properly formed
	// DHT than making sure our dht minimizes traffic.
	Period: time.Duration(5 * time.Minute),

	Timeout: time.Duration(10 * time.Second),
}

// Bootstrap ensures the dht routing table remains healthy as peers come and go.
// it builds up a list of peers by requesting random peer IDs. The Bootstrap
// process will run a number of queries each time, and run every time signal fires.
// These parameters are configurable.
//
// As opposed to BootstrapWithConfig, Bootstrap satisfies the routing interface
func (dht *IpfsDHT) Bootstrap(ctx context.Context) error {
	proc, err := dht.BootstrapWithConfig(DefaultBootstrapConfig)
	if err != nil {
		return err
	}

	// wait till ctx or dht.Context exits.
	// we have to do it this way to satisfy the Routing interface (contexts)
	go func() {
		defer proc.Close()
		select {
		case <-ctx.Done():
		case <-dht.Context().Done():
		}
	}()

	return nil
}

// BootstrapWithConfig ensures the dht routing table remains healthy as peers come and go.
// it builds up a list of peers by requesting random peer IDs. The Bootstrap
// process will run a number of queries each time, and run every time signal fires.
// These parameters are configurable.
//
// BootstrapWithConfig returns a process, so the user can stop it.
func (dht *IpfsDHT) BootstrapWithConfig(cfg BootstrapConfig) (goprocess.Process, error) {
	if cfg.Queries <= 0 {
		return nil, fmt.Errorf("invalid number of queries: %d", cfg.Queries)
	}

	tickch := make(chan struct{}, 1)
	proc := periodicproc.OnSignal(tickch, dht.bootstrapWorker(cfg))
	go func() {
		tickch <- struct{}{}
		for {
			select {
			case <-time.After(cfg.Period):
				select {
				case tickch <- struct{}{}:
				default:
					// Don't queue ticks, like Tickers
					log.Warning("Previous bootstrapping attempt not completed within bootstrapping period")
				}
			case <-dht.Context().Done():
				return
			}
		}
	}()

	return proc, nil
}

// SignalBootstrap ensures the dht routing table remains healthy as peers come and go.
// it builds up a list of peers by requesting random peer IDs. The Bootstrap
// process will run a number of queries each time, and run every time signal fires.
// These parameters are configurable.
//
// SignalBootstrap returns a process, so the user can stop it.
func (dht *IpfsDHT) BootstrapOnSignal(cfg BootstrapConfig, signal <-chan time.Time) (goprocess.Process, error) {
	if cfg.Queries <= 0 {
		return nil, fmt.Errorf("invalid number of queries: %d", cfg.Queries)
	}

	if signal == nil {
		return nil, fmt.Errorf("invalid signal: %v", signal)
	}

	proc := periodicproc.Ticker(signal, dht.bootstrapWorker(cfg))

	return proc, nil
}

func (dht *IpfsDHT) bootstrapWorker(cfg BootstrapConfig) func(worker goprocess.Process) {
	return func(worker goprocess.Process) {
		// it would be useful to be able to send out signals of when we bootstrap, too...
		// maybe this is a good case for whole module event pub/sub?

		ctx := dht.Context()
		if err := dht.runBootstrap(ctx, cfg); err != nil {
			log.Warning(err)
			// A bootstrapping error is important to notice but not fatal.
		}
	}
}

// runBootstrap builds up list of peers by requesting random peer IDs
func (dht *IpfsDHT) runBootstrap(ctx context.Context, cfg BootstrapConfig) error {
	bslog := func(msg string) {
		log.Debugf("DHT %s dhtRunBootstrap %s -- routing table size: %d", dht.self, msg, dht.routingTable.Size())
	}
	bslog("start")
	defer bslog("end")
	defer log.EventBegin(ctx, "dhtRunBootstrap").Done()

	var merr u.MultiErr

	randomID := func() peer.ID {
		// 16 random bytes is not a valid peer id. it may be fine becuase
		// the dht will rehash to its own keyspace anyway.
		id := make([]byte, 16)
		rand.Read(id)
		id = u.Hash(id)
		return peer.ID(id)
	}

	// bootstrap sequentially, as results will compound
	runQuery := func(ctx context.Context, id peer.ID) {
		ctx, cancel := context.WithTimeout(ctx, cfg.Timeout)
		defer cancel()

		p, err := dht.FindPeer(ctx, id)
		if err == routing.ErrNotFound {
			// this isn't an error. this is precisely what we expect.
		} else if err != nil {
			merr = append(merr, err)
		} else {
			// woah, actually found a peer with that ID? this shouldn't happen normally
			// (as the ID we use is not a real ID). this is an odd error worth logging.
			err := fmt.Errorf("Bootstrap peer error: Actually FOUND peer. (%s, %s)", id, p)
			log.Warningf("%s", err)
			merr = append(merr, err)
		}
	}

	// these should be parallel normally. but can make them sequential for debugging.
	// note that the core/bootstrap context deadline should be extended too for that.
	for i := 0; i < cfg.Queries; i++ {
		id := randomID()
		log.Debugf("Bootstrapping query (%d/%d) to random ID: %s", i+1, cfg.Queries, id)
		runQuery(ctx, id)
	}

	// Find self to distribute peer info to our neighbors.
	// Do this after bootstrapping.
	log.Debugf("Bootstrapping query to self: %s", dht.self)
	runQuery(ctx, dht.self)

	if len(merr) > 0 {
		return merr
	}
	return nil
}
<|endoftext|>


<commit_before>package junos

import (
	"encoding/xml"
	"fmt"
	"regexp"
)

// Devices contains a list of managed devices.
type Devices struct {
	XMLName xml.Name `xml:"devices"`
	Devices []Device `xml:"device"`
}

// A Device contains information about each individual device.
type Device struct {
	ID        int    `xml:"key,attr"`
	Family    string `xml:"deviceFamily"`
	Version   string `xml:"OSVersion"`
	Platform  string `xml:"platform"`
	Serial    string `xml:"serialNumber"`
	IPAddress string `xml:"ipAddr"`
	Name      string `xml:"name"`
}

// addDeviceIPXML is the XML we send (POST) for adding a device by IP address.
var addDeviceIPXML = `
<discover-devices>
    <ipAddressDiscoveryTarget>
        <ipAddress>%s</ipAddress>
    </ipAddressDiscoveryTarget>
    <sshCredential>
        <userName>%s</userName>
        <password>%s</password>
    </sshCredential>
    <manageDiscoveredSystemsFlag>true</manageDiscoveredSystemsFlag>
    <usePing>true</usePing>
</discover-devices>
`

// addDeviceHostXML is the XML we send (POST) for adding a device by hostname.
var addDeviceHostXML = `
<discover-devices>
    <hostNameDiscoveryTarget>
        <hostName>%s</hostName>
    </hostNameDiscoveryTarget>
    <sshCredential>
        <userName>%s</userName>
        <password>%s</password>
    </sshCredential>
    <manageDiscoveredSystemsFlag>true</manageDiscoveredSystemsFlag>
    <usePing>true</usePing>
</discover-devices>
`

// getDeviceID returns the ID of a managed device.
func (s *JunosSpace) getDeviceID(device interface{}, sd bool) (int, error) {
	var err error
	var deviceID int
	var sds *SecurityDevices
	ipRegex := regexp.MustCompile(`(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})`)
	if sd {
		sds, err = s.SecurityDevices()
	}
	devices, err := s.Devices()
	if err != nil {
		return 0, err
	}

	switch device.(type) {
	case int:
		deviceID = device.(int)
	case string:
		if sd {
			if ipRegex.MatchString(device.(string)) {
				for _, d := range sds.Devices {
					if d.IPAddress == device.(string) {
						deviceID = d.ID
					}
				}
			}
			for _, d := range sds.Devices {
				if d.Name == device.(string) {
					deviceID = d.ID
				}
			}
		}
		if ipRegex.MatchString(device.(string)) {
			for _, d := range devices.Devices {
				if d.IPAddress == device.(string) {
					deviceID = d.ID
				}
			}
		}
		for _, d := range devices.Devices {
			if d.Name == device.(string) {
				deviceID = d.ID
			}
		}
	}

	return deviceID, nil
}

// AddDevice adds a new managed device to Junos Space, and returns the Job ID.
func (s *JunosSpace) AddDevice(host, user, password string) (int, error) {
	var job jobID
	var addDevice string
	ipRegex := regexp.MustCompile(`(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})`)

	if ipRegex.MatchString(host) {
		addDevice = addDeviceIPXML
	}

	addDevice = addDeviceHostXML

	req := &APIRequest{
		Method:      "post",
		URL:         "/api/space/device-management/discover-devices",
		Body:        fmt.Sprintf(addDevice, host, user, password),
		ContentType: contentDiscoverDevices,
	}
	data, err := s.APICall(req)
	if err != nil {
		return 0, err
	}

	err = xml.Unmarshal(data, &job)
	if err != nil {
		return 0, err
	}

	return job.ID, nil
}

// Devices queries the Junos Space server and returns all of the information
// about each device that is managed by Space.
func (s *JunosSpace) Devices() (*Devices, error) {
	var devices Devices
	req := &APIRequest{
		Method: "get",
		URL:    "/api/space/device-management/devices",
	}
	data, err := s.APICall(req)
	if err != nil {
		return nil, err
	}

	err = xml.Unmarshal(data, &devices)
	if err != nil {
		return nil, err
	}

	return &devices, nil
}

// RemoveDevice removes a device from Junos Space. You can specify the device ID, name
// or IP address.
func (s *JunosSpace) RemoveDevice(device interface{}) error {
	var err error
	deviceID, err := s.getDeviceID(device, false)
	if err != nil {
		return err
	}

	if deviceID != 0 {
		req := &APIRequest{
			Method: "delete",
			URL:    fmt.Sprintf("/api/space/device-management/devices/%d", deviceID),
		}
		_, err = s.APICall(req)
		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>Removed boolean option from getDeviceID()<commit_after>package junos

import (
	"encoding/xml"
	"fmt"
	"regexp"
)

// Devices contains a list of managed devices.
type Devices struct {
	XMLName xml.Name `xml:"devices"`
	Devices []Device `xml:"device"`
}

// A Device contains information about each individual device.
type Device struct {
	ID        int    `xml:"key,attr"`
	Family    string `xml:"deviceFamily"`
	Version   string `xml:"OSVersion"`
	Platform  string `xml:"platform"`
	Serial    string `xml:"serialNumber"`
	IPAddress string `xml:"ipAddr"`
	Name      string `xml:"name"`
}

// addDeviceIPXML is the XML we send (POST) for adding a device by IP address.
var addDeviceIPXML = `
<discover-devices>
    <ipAddressDiscoveryTarget>
        <ipAddress>%s</ipAddress>
    </ipAddressDiscoveryTarget>
    <sshCredential>
        <userName>%s</userName>
        <password>%s</password>
    </sshCredential>
    <manageDiscoveredSystemsFlag>true</manageDiscoveredSystemsFlag>
    <usePing>true</usePing>
</discover-devices>
`

// addDeviceHostXML is the XML we send (POST) for adding a device by hostname.
var addDeviceHostXML = `
<discover-devices>
    <hostNameDiscoveryTarget>
        <hostName>%s</hostName>
    </hostNameDiscoveryTarget>
    <sshCredential>
        <userName>%s</userName>
        <password>%s</password>
    </sshCredential>
    <manageDiscoveredSystemsFlag>true</manageDiscoveredSystemsFlag>
    <usePing>true</usePing>
</discover-devices>
`

// getDeviceID returns the ID of a managed device.
func (s *JunosSpace) getDeviceID(device interface{}) (int, error) {
	var err error
	var deviceID int
	ipRegex := regexp.MustCompile(`(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})`)
	devices, err := s.Devices()
	if err != nil {
		return 0, err
	}

	switch device.(type) {
	case int:
		deviceID = device.(int)
	case string:
		if ipRegex.MatchString(device.(string)) {
			for _, d := range devices.Devices {
				if d.IPAddress == device.(string) {
					deviceID = d.ID
				}
			}
		}
		for _, d := range devices.Devices {
			if d.Name == device.(string) {
				deviceID = d.ID
			}
		}
	}

	return deviceID, nil
}

// AddDevice adds a new managed device to Junos Space, and returns the Job ID.
func (s *JunosSpace) AddDevice(host, user, password string) (int, error) {
	var job jobID
	var addDevice string
	ipRegex := regexp.MustCompile(`(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})`)

	if ipRegex.MatchString(host) {
		addDevice = addDeviceIPXML
	}

	addDevice = addDeviceHostXML

	req := &APIRequest{
		Method:      "post",
		URL:         "/api/space/device-management/discover-devices",
		Body:        fmt.Sprintf(addDevice, host, user, password),
		ContentType: contentDiscoverDevices,
	}
	data, err := s.APICall(req)
	if err != nil {
		return 0, err
	}

	err = xml.Unmarshal(data, &job)
	if err != nil {
		return 0, err
	}

	return job.ID, nil
}

// Devices queries the Junos Space server and returns all of the information
// about each device that is managed by Space.
func (s *JunosSpace) Devices() (*Devices, error) {
	var devices Devices
	req := &APIRequest{
		Method: "get",
		URL:    "/api/space/device-management/devices",
	}
	data, err := s.APICall(req)
	if err != nil {
		return nil, err
	}

	err = xml.Unmarshal(data, &devices)
	if err != nil {
		return nil, err
	}

	return &devices, nil
}

// RemoveDevice removes a device from Junos Space. You can specify the device ID, name
// or IP address.
func (s *JunosSpace) RemoveDevice(device interface{}) error {
	var err error
	deviceID, err := s.getDeviceID(device)
	if err != nil {
		return err
	}

	if deviceID != 0 {
		req := &APIRequest{
			Method: "delete",
			URL:    fmt.Sprintf("/api/space/device-management/devices/%d", deviceID),
		}
		_, err = s.APICall(req)
		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package goblet

import (
	"net/http"
	"strconv"
	"time"
)

func (c *Context) GetLoginId() (string, bool) {
	return c.GetLoginIdAs(USERCOOKIENAME)
}

func (c *Context) GetLoginIdAs(name string) (string, bool) {
	cookie, err := c.SignedCookie(name + "Id")
	if cookie != nil && err == nil {
		return cookie.Value, true
	}
	return "", false
}

func (c *Context) AddLoginIdAs(id interface{}, name string, timeduration ...time.Duration) {
	if name == "" {
		name = "user"
	}
	var userid string
	switch rid := id.(type) {
	case string:
		userid = rid
	case int64:
		userid = strconv.FormatInt(rid, 10)

	}
	if timeduration == nil {
		c.addLoginAs(name, userid)
	} else {
		c.addLoginAs(name, userid, timeduration[0])
	}

}

func (c *Context) AddLoginId(id interface{}, timeduration ...time.Duration) {
	var userid string
	switch rid := id.(type) {
	case string:
		userid = rid
	case int64:
		userid = strconv.FormatInt(rid, 10)

	}
	if timeduration == nil {
		c.addLoginAs("user", userid)
	} else {
		c.addLoginAs("user", userid, timeduration[0])
	}

}

func (c *Context) addLoginAs(name string, id string, timeduration ...time.Duration) {
	expire := time.Now().AddDate(0, 0, 1)
	if timeduration != nil {
		expire = time.Now().Add(timeduration[0])
	}
	cookie := new(http.Cookie)
	cookie.Name = name + "Id"
	cookie.Value = id
	cookie.Expires = expire
	cookie.Path = "/"
	cookie.RawExpires = expire.Format(time.UnixDate)
	c.AddSignedCookie(cookie)
}

//Delete the login cookie saved
func (c *Context) DelLogin() {
	c.delLoginAs("user")
}

func (c *Context) delLoginAs(name string) {
	cookie, err := c.SignedCookie(name + "Id")
	if cookie != nil && err == nil {
		cookie.MaxAge = -1
		expire := time.Now()
		cookie.Expires = expire
		cookie.RawExpires = expire.Format(time.UnixDate)
		c.AddSignedCookie(cookie)
	}
}
<commit_msg>addlogin support int and int32 DelLogin add param name .<commit_after>package goblet

import (
	"net/http"
	"strconv"
	"time"
)

func (c *Context) GetLoginId() (string, bool) {
	return c.GetLoginIdAs(USERCOOKIENAME)
}

func (c *Context) GetLoginIdAs(name string) (string, bool) {
	cookie, err := c.SignedCookie(name + "Id")
	if cookie != nil && err == nil {
		return cookie.Value, true
	}
	return "", false
}

func (c *Context) AddLoginIdAs(id interface{}, name string, timeduration ...time.Duration) {
	if name == "" {
		name = "user"
	}
	var userid string
	switch rid := id.(type) {
	case string:
		userid = rid
	case int:
		userid = strconv.FormatInt(int64(rid), 10)
	case int32:
		userid = strconv.FormatInt(int64(rid), 10)
	case int64:
		userid = strconv.FormatInt(rid, 10)
	}
	if timeduration == nil {
		c.addLoginAs(name, userid)
	} else {
		c.addLoginAs(name, userid, timeduration[0])
	}

}

func (c *Context) AddLoginId(id interface{}, timeduration ...time.Duration) {
	var userid string
	switch rid := id.(type) {
	case string:
		userid = rid
	case int:
		userid = strconv.FormatInt(int64(rid), 10)
	case int32:
		userid = strconv.FormatInt(int64(rid), 10)
	case int64:
		userid = strconv.FormatInt(rid, 10)
	}
	if timeduration == nil {
		c.addLoginAs("user", userid)
	} else {
		c.addLoginAs("user", userid, timeduration[0])
	}

}

func (c *Context) addLoginAs(name string, id string, timeduration ...time.Duration) {
	expire := time.Now().AddDate(0, 0, 1)
	if timeduration != nil {
		expire = time.Now().Add(timeduration[0])
	}
	cookie := new(http.Cookie)
	cookie.Name = name + "Id"
	cookie.Value = id
	cookie.Expires = expire
	cookie.Path = "/"
	cookie.RawExpires = expire.Format(time.UnixDate)
	c.AddSignedCookie(cookie)
}

//Delete the login cookie saved
func (c *Context) DelLogin(names ...string) {
	if names != nil {
		for _, v := range names {
			c.delLoginAs(v)
		}
	} else {
		c.delLoginAs("user")
	}
}

func (c *Context) delLoginAs(name string) {
	cookie, err := c.SignedCookie(name + "Id")
	if cookie != nil && err == nil {
		cookie.MaxAge = -1
		expire := time.Now()
		cookie.Expires = expire
		cookie.RawExpires = expire.Format(time.UnixDate)
		c.AddSignedCookie(cookie)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"time"

	"github.com/apex/log"
	"github.com/apex/log/handlers/cli"

	"github.com/apex/apex/cmd/apex/root"
	"github.com/apex/apex/stats"

	// commands
	_ "github.com/apex/apex/cmd/apex/autocomplete"
	_ "github.com/apex/apex/cmd/apex/build"
	_ "github.com/apex/apex/cmd/apex/delete"
	_ "github.com/apex/apex/cmd/apex/deploy"
	_ "github.com/apex/apex/cmd/apex/docs"
	_ "github.com/apex/apex/cmd/apex/infra"
	_ "github.com/apex/apex/cmd/apex/init"
	_ "github.com/apex/apex/cmd/apex/invoke"
	_ "github.com/apex/apex/cmd/apex/list"
	_ "github.com/apex/apex/cmd/apex/logs"
	_ "github.com/apex/apex/cmd/apex/metrics"
	_ "github.com/apex/apex/cmd/apex/rollback"
	_ "github.com/apex/apex/cmd/apex/upgrade"
	_ "github.com/apex/apex/cmd/apex/version"

	// plugins
	_ "github.com/apex/apex/plugins/env"
	_ "github.com/apex/apex/plugins/golang"
	_ "github.com/apex/apex/plugins/hooks"
	_ "github.com/apex/apex/plugins/inference"
	_ "github.com/apex/apex/plugins/java"
	_ "github.com/apex/apex/plugins/nodejs"
	_ "github.com/apex/apex/plugins/python"
	_ "github.com/apex/apex/plugins/shim"
)

// Terraform commands.
var tf = []string{
	"apply",
	"destroy",
	"get",
	"graph",
	"init",
	"output",
	"plan",
	"refresh",
	"remote",
	"show",
	"taint",
	"untaint",
	"validate",
	"version",
}

func main() {
	log.SetHandler(cli.Default)

	args := os.Args[1:]

	// Cobra does not (currently) allow us to pass flags for a sub-command
	// as if they were arguments, so we inject -- here after the first TF command.
	// TODO(tj): replace with a real solution and send PR to Cobra #251
	if len(os.Args) > 1 && os.Args[1] == "infra" {
		off := 1

	out:
		for i, a := range args {
			for _, cmd := range tf {
				if a == cmd {
					off = i
					break out
				}
			}
		}

		args = append(args[0:off], append([]string{"--"}, args[off:]...)...)
	}

	root.Command.SetArgs(args)

	if err := root.Command.Execute(); err != nil {
		log.Fatalf("Error: %s", err)
	}

	stats.Client.ConditionalFlush(500, 24*time.Hour)
}
<commit_msg>fix flags before infra sub-command. Closes #421<commit_after>package main

import (
	"os"
	"time"

	"github.com/apex/log"
	"github.com/apex/log/handlers/cli"

	"github.com/apex/apex/cmd/apex/root"
	"github.com/apex/apex/stats"

	// commands
	_ "github.com/apex/apex/cmd/apex/autocomplete"
	_ "github.com/apex/apex/cmd/apex/build"
	_ "github.com/apex/apex/cmd/apex/delete"
	_ "github.com/apex/apex/cmd/apex/deploy"
	_ "github.com/apex/apex/cmd/apex/docs"
	_ "github.com/apex/apex/cmd/apex/infra"
	_ "github.com/apex/apex/cmd/apex/init"
	_ "github.com/apex/apex/cmd/apex/invoke"
	_ "github.com/apex/apex/cmd/apex/list"
	_ "github.com/apex/apex/cmd/apex/logs"
	_ "github.com/apex/apex/cmd/apex/metrics"
	_ "github.com/apex/apex/cmd/apex/rollback"
	_ "github.com/apex/apex/cmd/apex/upgrade"
	_ "github.com/apex/apex/cmd/apex/version"

	// plugins
	_ "github.com/apex/apex/plugins/env"
	_ "github.com/apex/apex/plugins/golang"
	_ "github.com/apex/apex/plugins/hooks"
	_ "github.com/apex/apex/plugins/inference"
	_ "github.com/apex/apex/plugins/java"
	_ "github.com/apex/apex/plugins/nodejs"
	_ "github.com/apex/apex/plugins/python"
	_ "github.com/apex/apex/plugins/shim"
)

// Terraform commands.
var tf = []string{
	"apply",
	"destroy",
	"get",
	"graph",
	"init",
	"output",
	"plan",
	"refresh",
	"remote",
	"show",
	"taint",
	"untaint",
	"validate",
	"version",
}

// TODO(tj): remove this evil hack, necessary for now for cases such as:
//
//   $ apex --env prod infra deploy
//
// instead of:
//
//   $ apex infra --env prod deploy
//
func isInfra(args []string) bool {
	for _, arg := range args {
		if arg == "infra" {
			return true
		}
	}
	return false
}

func main() {
	log.SetHandler(cli.Default)

	args := os.Args[1:]

	// Cobra does not (currently) allow us to pass flags for a sub-command
	// as if they were arguments, so we inject -- here after the first TF command.
	// TODO(tj): replace with a real solution and send PR to Cobra #251
	if len(os.Args) > 1 && isInfra(os.Args) {
		off := 1

	out:
		for i, a := range args {
			for _, cmd := range tf {
				if a == cmd {
					off = i
					break out
				}
			}
		}

		args = append(args[0:off], append([]string{"--"}, args[off:]...)...)
	}

	root.Command.SetArgs(args)

	if err := root.Command.Execute(); err != nil {
		log.Fatalf("Error: %s", err)
	}

	stats.Client.ConditionalFlush(500, 24*time.Hour)
}
<|endoftext|>


<commit_before>//  Copyright (c) 2014 Couchbase, Inc.
//  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 main

import (
	"expvar"
	"flag"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"path"
	"runtime"
	"runtime/pprof"
	"strings"
	"time"

	"github.com/gorilla/mux"

	"github.com/blevesearch/bleve"
	bleveHttp "github.com/blevesearch/bleve/http"
	bleveRegistry "github.com/blevesearch/bleve/registry"

	log "github.com/couchbase/clog"
	"github.com/couchbase/go-couchbase"
	"github.com/couchbaselabs/cbft"
)

var VERSION = "v0.0.2"

var expvars = expvar.NewMap("stats")

func main() {
	flag.Parse()

	if flags.Help {
		flag.Usage()
		os.Exit(2)
	}

	if flags.Version {
		fmt.Printf("%s main: %s, data: %s\n", path.Base(os.Args[0]),
			VERSION, cbft.VERSION)
		os.Exit(0)
	}

	log.Printf("main: %s started (%s/%s)",
		os.Args[0], VERSION, cbft.VERSION)

	rand.Seed(time.Now().UTC().UnixNano())

	go dumpOnSignalForPlatform()

	mr, err := cbft.NewMsgRing(os.Stderr, 1000)
	if err != nil {
		log.Fatalf("main: could not create MsgRing, err: %v", err)
	}
	log.SetOutput(mr)

	MainWelcome(flagAliases)

	// If cfg is down, we error, leaving it to some user-supplied
	// outside watchdog to backoff and restart/retry.
	cfg, err := MainCfg(flags.CfgConnect, flags.DataDir)
	if err != nil {
		log.Fatalf("main: could not start cfg, cfgConnect: %s, err: %v\n"+
			"  Please check that your -cfg/-cfgConnect parameter (%q)\n"+
			"  is correct and/or that your configuration provider\n"+
			"  is available.",
			flags.CfgConnect, err, flags.CfgConnect)
		return
	}

	uuid, err := MainUUID(flags.DataDir)
	if err != nil {
		log.Fatalf(fmt.Sprintf("%v", err))
		return
	}

	var tagsArr []string
	if flags.Tags != "" {
		tagsArr = strings.Split(flags.Tags, ",")
	}

	expvars.Set("indexes", bleveHttp.IndexStats())

	router, err := MainStart(cfg, uuid, tagsArr,
		flags.Container, flags.Weight,
		flags.BindHttp, flags.DataDir,
		flags.StaticDir, flags.StaticETag,
		flags.Server, flags.Register, mr)
	if err != nil {
		log.Fatal(err)
	}

	http.Handle("/", router)

	log.Printf("main: listening on: %v", flags.BindHttp)
	err = http.ListenAndServe(flags.BindHttp, nil)
	if err != nil {
		log.Fatalf("main: listen, err: %v\n"+
			"  Please check that your -bindHttp parameter (%q)\n"+
			"  is correct and available.", err, flags.BindHttp)
	}
}

func MainWelcome(flagAliases map[string][]string) {
	flag.VisitAll(func(f *flag.Flag) {
		if flagAliases[f.Name] != nil {
			log.Printf("  -%s=%q\n", f.Name, f.Value)
		}
	})
	log.Printf("  GOMAXPROCS=%d", runtime.GOMAXPROCS(-1))

	log.Printf("main: registered bleve stores")
	types, instances := bleveRegistry.KVStoreTypesAndInstances()
	for _, s := range types {
		log.Printf("  %s", s)
	}
	for _, s := range instances {
		log.Printf("  %s", s)
	}
}

func MainUUID(dataDir string) (string, error) {
	uuid := cbft.NewUUID()
	uuidPath := dataDir + string(os.PathSeparator) + "cbft.uuid"
	uuidBuf, err := ioutil.ReadFile(uuidPath)
	if err == nil {
		uuid = strings.TrimSpace(string(uuidBuf))
		if uuid == "" {
			return "", fmt.Errorf("error: could not parse uuidPath: %s",
				uuidPath)
		}
		log.Printf("main: manager uuid: %s", uuid)
		log.Printf("main: manager uuid was reloaded")
	} else {
		log.Printf("main: manager uuid: %s", uuid)
		log.Printf("main: manager uuid was generated")
	}
	err = ioutil.WriteFile(uuidPath, []byte(uuid), 0600)
	if err != nil {
		return "", fmt.Errorf("error: could not write uuidPath: %s\n"+
			"  Please check that your -data/-dataDir parameter (%q)\n"+
			"  is to a writable directory where cbft can store\n"+
			"  index data.",
			uuidPath, dataDir)
	}
	return uuid, nil
}

func MainStart(cfg cbft.Cfg, uuid string, tags []string, container string,
	weight int, bindHttp, dataDir, staticDir, staticETag, server string,
	register string, mr *cbft.MsgRing) (
	*mux.Router, error) {
	if server == "" {
		return nil, fmt.Errorf("error: server URL required (-server)")
	}

	if server != "." {
		_, err := couchbase.Connect(server)
		if err != nil {
			return nil, fmt.Errorf("error: could not connect to server,"+
				" URL: %s, err: %v\n"+
				"  Please check that your -server parameter (%q)\n"+
				"  is correct and the server is available.",
				server, err, server)
		}
	}

	mgr := cbft.NewManager(cbft.VERSION, cfg, uuid, tags, container, weight,
		bindHttp, dataDir, server, &MainHandlers{})
	err := mgr.Start(register)
	if err != nil {
		return nil, err
	}

	router, _, err :=
		cbft.NewManagerRESTRouter(VERSION, mgr, staticDir, staticETag, mr)

	return router, err
}

type MainHandlers struct{}

func (meh *MainHandlers) OnRegisterPIndex(pindex *cbft.PIndex) {
	bindex, ok := pindex.Impl.(bleve.Index)
	if ok {
		bleveHttp.RegisterIndexName(pindex.Name, bindex)
	}
}

func (meh *MainHandlers) OnUnregisterPIndex(pindex *cbft.PIndex) {
	bleveHttp.UnregisterIndexByName(pindex.Name)
}

func dumpOnSignal(signals ...os.Signal) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, signals...)
	for _ = range c {
		log.Printf("dump: goroutine...")
		pprof.Lookup("goroutine").WriteTo(os.Stderr, 1)
		log.Printf("dump: heap...")
		pprof.Lookup("heap").WriteTo(os.Stderr, 1)
	}
}
<commit_msg>log admin web UI URL more prominently<commit_after>//  Copyright (c) 2014 Couchbase, Inc.
//  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 main

import (
	"expvar"
	"flag"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"path"
	"runtime"
	"runtime/pprof"
	"strings"
	"time"

	"github.com/gorilla/mux"

	"github.com/blevesearch/bleve"
	bleveHttp "github.com/blevesearch/bleve/http"
	bleveRegistry "github.com/blevesearch/bleve/registry"

	log "github.com/couchbase/clog"
	"github.com/couchbase/go-couchbase"
	"github.com/couchbaselabs/cbft"
)

var VERSION = "v0.0.2"

var expvars = expvar.NewMap("stats")

func main() {
	flag.Parse()

	if flags.Help {
		flag.Usage()
		os.Exit(2)
	}

	if flags.Version {
		fmt.Printf("%s main: %s, data: %s\n", path.Base(os.Args[0]),
			VERSION, cbft.VERSION)
		os.Exit(0)
	}

	log.Printf("main: %s started (%s/%s)",
		os.Args[0], VERSION, cbft.VERSION)

	rand.Seed(time.Now().UTC().UnixNano())

	go dumpOnSignalForPlatform()

	mr, err := cbft.NewMsgRing(os.Stderr, 1000)
	if err != nil {
		log.Fatalf("main: could not create MsgRing, err: %v", err)
	}
	log.SetOutput(mr)

	MainWelcome(flagAliases)

	// If cfg is down, we error, leaving it to some user-supplied
	// outside watchdog to backoff and restart/retry.
	cfg, err := MainCfg(flags.CfgConnect, flags.DataDir)
	if err != nil {
		log.Fatalf("main: could not start cfg, cfgConnect: %s, err: %v\n"+
			"  Please check that your -cfg/-cfgConnect parameter (%q)\n"+
			"  is correct and/or that your configuration provider\n"+
			"  is available.",
			flags.CfgConnect, err, flags.CfgConnect)
		return
	}

	uuid, err := MainUUID(flags.DataDir)
	if err != nil {
		log.Fatalf(fmt.Sprintf("%v", err))
		return
	}

	var tagsArr []string
	if flags.Tags != "" {
		tagsArr = strings.Split(flags.Tags, ",")
	}

	expvars.Set("indexes", bleveHttp.IndexStats())

	router, err := MainStart(cfg, uuid, tagsArr,
		flags.Container, flags.Weight,
		flags.BindHttp, flags.DataDir,
		flags.StaticDir, flags.StaticETag,
		flags.Server, flags.Register, mr)
	if err != nil {
		log.Fatal(err)
	}

	http.Handle("/", router)

	log.Printf("main: listening on: %s", flags.BindHttp)
	u := flags.BindHttp
	if u[0] == ':' {
		u = "localhost" + u
	}
	log.Printf("------------------------------------------------------------")
	log.Printf("admin web UI is available: http://%s", u)
	log.Printf("------------------------------------------------------------")
	err = http.ListenAndServe(flags.BindHttp, nil)
	if err != nil {
		log.Fatalf("main: listen, err: %v\n"+
			"  Please check that your -bindHttp parameter (%q)\n"+
			"  is correct and available.", err, flags.BindHttp)
	}
}

func MainWelcome(flagAliases map[string][]string) {
	flag.VisitAll(func(f *flag.Flag) {
		if flagAliases[f.Name] != nil {
			log.Printf("  -%s=%q\n", f.Name, f.Value)
		}
	})
	log.Printf("  GOMAXPROCS=%d", runtime.GOMAXPROCS(-1))

	log.Printf("main: registered bleve stores")
	types, instances := bleveRegistry.KVStoreTypesAndInstances()
	for _, s := range types {
		log.Printf("  %s", s)
	}
	for _, s := range instances {
		log.Printf("  %s", s)
	}
}

func MainUUID(dataDir string) (string, error) {
	uuid := cbft.NewUUID()
	uuidPath := dataDir + string(os.PathSeparator) + "cbft.uuid"
	uuidBuf, err := ioutil.ReadFile(uuidPath)
	if err == nil {
		uuid = strings.TrimSpace(string(uuidBuf))
		if uuid == "" {
			return "", fmt.Errorf("error: could not parse uuidPath: %s",
				uuidPath)
		}
		log.Printf("main: manager uuid: %s", uuid)
		log.Printf("main: manager uuid was reloaded")
	} else {
		log.Printf("main: manager uuid: %s", uuid)
		log.Printf("main: manager uuid was generated")
	}
	err = ioutil.WriteFile(uuidPath, []byte(uuid), 0600)
	if err != nil {
		return "", fmt.Errorf("error: could not write uuidPath: %s\n"+
			"  Please check that your -data/-dataDir parameter (%q)\n"+
			"  is to a writable directory where cbft can store\n"+
			"  index data.",
			uuidPath, dataDir)
	}
	return uuid, nil
}

func MainStart(cfg cbft.Cfg, uuid string, tags []string, container string,
	weight int, bindHttp, dataDir, staticDir, staticETag, server string,
	register string, mr *cbft.MsgRing) (
	*mux.Router, error) {
	if server == "" {
		return nil, fmt.Errorf("error: server URL required (-server)")
	}

	if server != "." {
		_, err := couchbase.Connect(server)
		if err != nil {
			return nil, fmt.Errorf("error: could not connect to server,"+
				" URL: %s, err: %v\n"+
				"  Please check that your -server parameter (%q)\n"+
				"  is correct and the server is available.",
				server, err, server)
		}
	}

	mgr := cbft.NewManager(cbft.VERSION, cfg, uuid, tags, container, weight,
		bindHttp, dataDir, server, &MainHandlers{})
	err := mgr.Start(register)
	if err != nil {
		return nil, err
	}

	router, _, err :=
		cbft.NewManagerRESTRouter(VERSION, mgr, staticDir, staticETag, mr)

	return router, err
}

type MainHandlers struct{}

func (meh *MainHandlers) OnRegisterPIndex(pindex *cbft.PIndex) {
	bindex, ok := pindex.Impl.(bleve.Index)
	if ok {
		bleveHttp.RegisterIndexName(pindex.Name, bindex)
	}
}

func (meh *MainHandlers) OnUnregisterPIndex(pindex *cbft.PIndex) {
	bleveHttp.UnregisterIndexByName(pindex.Name)
}

func dumpOnSignal(signals ...os.Signal) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, signals...)
	for _ = range c {
		log.Printf("dump: goroutine...")
		pprof.Lookup("goroutine").WriteTo(os.Stderr, 1)
		log.Printf("dump: heap...")
		pprof.Lookup("heap").WriteTo(os.Stderr, 1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"
	"io"
	"os"
	"sync"
	"text/tabwriter"
	"time"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/content"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/images"
	"github.com/containerd/containerd/log"
	"github.com/containerd/containerd/progress"
	"github.com/containerd/containerd/remotes"
	digest "github.com/opencontainers/go-digest"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/urfave/cli"
)

var fetchCommand = cli.Command{
	Name:      "fetch",
	Usage:     "fetch all content for an image into containerd",
	ArgsUsage: "[flags] <remote> <object>",
	Description: `Fetch an image into containerd.

This command ensures that containerd has all the necessary resources to build
an image's rootfs and convert the configuration to a runtime format supported
by containerd.

This command uses the same syntax, of remote and object, as 'ctr fetch-object'.
We may want to make this nicer, but agnostism is preferred for the moment.

Right now, the responsibility of the daemon and the cli aren't quite clear. Do
not use this implementation as a guide. The end goal should be having metadata,
content and snapshots ready for a direct use via the 'ctr run'.

Most of this is experimental and there are few leaps to make this work.`,
	Flags: registryFlags,
	Action: func(clicontext *cli.Context) error {
		var (
			ref = clicontext.Args().First()
		)
		ctx, cancel := appContext(clicontext)
		defer cancel()

		_, err := fetch(ctx, ref, clicontext)
		return err
	},
}

func fetch(ctx context.Context, ref string, clicontext *cli.Context) (containerd.Image, error) {
	client, err := newClient(clicontext)
	if err != nil {
		return nil, err
	}

	resolver, err := getResolver(ctx, clicontext)
	if err != nil {
		return nil, err
	}

	ongoing := newJobs(ref)

	pctx, stopProgress := context.WithCancel(ctx)
	progress := make(chan struct{})

	go func() {
		showProgress(pctx, ongoing, client.ContentStore(), os.Stdout)
		close(progress)
	}()

	h := images.HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
		if desc.MediaType != images.MediaTypeDockerSchema1Manifest {
			ongoing.add(desc)
		}
		return nil, nil
	})

	log.G(pctx).WithField("image", ref).Debug("fetching")

	img, err := client.Pull(pctx, ref, containerd.WithResolver(resolver), containerd.WithImageHandler(h), containerd.WithSchema1Conversion)
	stopProgress()
	if err != nil {
		return nil, err
	}

	<-progress
	return img, nil
}

func showProgress(ctx context.Context, ongoing *jobs, cs content.Store, out io.Writer) {
	var (
		ticker   = time.NewTicker(100 * time.Millisecond)
		fw       = progress.NewWriter(out)
		start    = time.Now()
		statuses = map[string]statusInfo{}
		done     bool
	)
	defer ticker.Stop()

outer:
	for {
		select {
		case <-ticker.C:
			fw.Flush()

			tw := tabwriter.NewWriter(fw, 1, 8, 1, ' ', 0)

			resolved := "resolved"
			if !ongoing.isResolved() {
				resolved = "resolving"
			}
			statuses[ongoing.name] = statusInfo{
				Ref:    ongoing.name,
				Status: resolved,
			}
			keys := []string{ongoing.name}

			activeSeen := map[string]struct{}{}
			if !done {
				active, err := cs.ListStatuses(ctx, "")
				if err != nil {
					log.G(ctx).WithError(err).Error("active check failed")
					continue
				}
				// update status of active entries!
				for _, active := range active {
					statuses[active.Ref] = statusInfo{
						Ref:       active.Ref,
						Status:    "downloading",
						Offset:    active.Offset,
						Total:     active.Total,
						StartedAt: active.StartedAt,
						UpdatedAt: active.UpdatedAt,
					}
					activeSeen[active.Ref] = struct{}{}
				}
			}

			// now, update the items in jobs that are not in active
			for _, j := range ongoing.jobs() {
				key := remotes.MakeRefKey(ctx, j)
				keys = append(keys, key)
				if _, ok := activeSeen[key]; ok {
					continue
				}

				status, ok := statuses[key]
				if !done && (!ok || status.Status == "downloading") {
					info, err := cs.Info(ctx, j.Digest)
					if err != nil {
						if !errdefs.IsNotFound(err) {
							log.G(ctx).WithError(err).Errorf("failed to get content info")
							continue outer
						} else {
							statuses[key] = statusInfo{
								Ref:    key,
								Status: "waiting",
							}
						}
					} else if info.CreatedAt.After(start) {
						statuses[key] = statusInfo{
							Ref:       key,
							Status:    "done",
							Offset:    info.Size,
							Total:     info.Size,
							UpdatedAt: info.CreatedAt,
						}
					} else {
						statuses[key] = statusInfo{
							Ref:    key,
							Status: "exists",
						}
					}
				} else if done {
					if ok {
						if status.Status != "done" && status.Status != "exists" {
							status.Status = "done"
							statuses[key] = status
						}
					} else {
						statuses[key] = statusInfo{
							Ref:    key,
							Status: "done",
						}
					}
				}
			}

			var ordered []statusInfo
			for _, key := range keys {
				ordered = append(ordered, statuses[key])
			}

			display(tw, ordered, start)
			tw.Flush()

			if done {
				fw.Flush()
				return
			}
		case <-ctx.Done():
			done = true // allow ui to update once more
		}
	}
}

// jobs provides a way of identifying the download keys for a particular task
// encountering during the pull walk.
//
// This is very minimal and will probably be replaced with something more
// featured.
type jobs struct {
	name     string
	added    map[digest.Digest]struct{}
	descs    []ocispec.Descriptor
	mu       sync.Mutex
	resolved bool
}

func newJobs(name string) *jobs {
	return &jobs{
		name:  name,
		added: map[digest.Digest]struct{}{},
	}
}

func (j *jobs) add(desc ocispec.Descriptor) {
	j.mu.Lock()
	defer j.mu.Unlock()
	j.resolved = true

	if _, ok := j.added[desc.Digest]; ok {
		return
	}
	j.descs = append(j.descs, desc)
	j.added[desc.Digest] = struct{}{}
}

func (j *jobs) jobs() []ocispec.Descriptor {
	j.mu.Lock()
	defer j.mu.Unlock()

	var descs []ocispec.Descriptor
	return append(descs, j.descs...)
}

func (j *jobs) isResolved() bool {
	j.mu.Lock()
	defer j.mu.Unlock()
	return j.resolved
}

type statusInfo struct {
	Ref       string
	Status    string
	Offset    int64
	Total     int64
	StartedAt time.Time
	UpdatedAt time.Time
}

func display(w io.Writer, statuses []statusInfo, start time.Time) {
	var total int64
	for _, status := range statuses {
		total += status.Offset
		switch status.Status {
		case "downloading", "uploading":
			var bar progress.Bar
			if status.Total > 0.0 {
				bar = progress.Bar(float64(status.Offset) / float64(status.Total))
			}
			fmt.Fprintf(w, "%s:\t%s\t%40r\t%8.8s/%s\t\n",
				status.Ref,
				status.Status,
				bar,
				progress.Bytes(status.Offset), progress.Bytes(status.Total))
		case "resolving", "waiting":
			bar := progress.Bar(0.0)
			fmt.Fprintf(w, "%s:\t%s\t%40r\t\n",
				status.Ref,
				status.Status,
				bar)
		default:
			bar := progress.Bar(1.0)
			fmt.Fprintf(w, "%s:\t%s\t%40r\t\n",
				status.Ref,
				status.Status,
				bar)
		}
	}

	fmt.Fprintf(w, "elapsed: %-4.1fs\ttotal: %7.6v\t(%v)\t\n",
		time.Since(start).Seconds(),
		// TODO(stevvooe): These calculations are actually way off.
		// Need to account for previously downloaded data. These
		// will basically be right for a download the first time
		// but will be skewed if restarting, as it includes the
		// data into the start time before.
		progress.Bytes(total),
		progress.NewBytesPerSecond(total, time.Since(start)))
}
<commit_msg>ctr: dismiss progressbar on --debug<commit_after>package main

import (
	"context"
	"fmt"
	"io"
	"os"
	"sync"
	"text/tabwriter"
	"time"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/content"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/images"
	"github.com/containerd/containerd/log"
	"github.com/containerd/containerd/progress"
	"github.com/containerd/containerd/remotes"
	digest "github.com/opencontainers/go-digest"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/urfave/cli"
)

var fetchCommand = cli.Command{
	Name:      "fetch",
	Usage:     "fetch all content for an image into containerd",
	ArgsUsage: "[flags] <remote> <object>",
	Description: `Fetch an image into containerd.

This command ensures that containerd has all the necessary resources to build
an image's rootfs and convert the configuration to a runtime format supported
by containerd.

This command uses the same syntax, of remote and object, as 'ctr fetch-object'.
We may want to make this nicer, but agnostism is preferred for the moment.

Right now, the responsibility of the daemon and the cli aren't quite clear. Do
not use this implementation as a guide. The end goal should be having metadata,
content and snapshots ready for a direct use via the 'ctr run'.

Most of this is experimental and there are few leaps to make this work.`,
	Flags: registryFlags,
	Action: func(clicontext *cli.Context) error {
		var (
			ref = clicontext.Args().First()
		)
		ctx, cancel := appContext(clicontext)
		defer cancel()

		_, err := fetch(ctx, ref, clicontext)
		return err
	},
}

func fetch(ctx context.Context, ref string, clicontext *cli.Context) (containerd.Image, error) {
	client, err := newClient(clicontext)
	if err != nil {
		return nil, err
	}

	resolver, err := getResolver(ctx, clicontext)
	if err != nil {
		return nil, err
	}

	ongoing := newJobs(ref)

	pctx, stopProgress := context.WithCancel(ctx)
	progress := make(chan struct{})

	go func() {
		if !clicontext.GlobalBool("debug") {
			// no progress bar, because it hides some debug logs
			showProgress(pctx, ongoing, client.ContentStore(), os.Stdout)
		}
		close(progress)
	}()

	h := images.HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
		if desc.MediaType != images.MediaTypeDockerSchema1Manifest {
			ongoing.add(desc)
		}
		return nil, nil
	})

	log.G(pctx).WithField("image", ref).Debug("fetching")

	img, err := client.Pull(pctx, ref, containerd.WithResolver(resolver), containerd.WithImageHandler(h), containerd.WithSchema1Conversion)
	stopProgress()
	if err != nil {
		return nil, err
	}

	<-progress
	return img, nil
}

func showProgress(ctx context.Context, ongoing *jobs, cs content.Store, out io.Writer) {
	var (
		ticker   = time.NewTicker(100 * time.Millisecond)
		fw       = progress.NewWriter(out)
		start    = time.Now()
		statuses = map[string]statusInfo{}
		done     bool
	)
	defer ticker.Stop()

outer:
	for {
		select {
		case <-ticker.C:
			fw.Flush()

			tw := tabwriter.NewWriter(fw, 1, 8, 1, ' ', 0)

			resolved := "resolved"
			if !ongoing.isResolved() {
				resolved = "resolving"
			}
			statuses[ongoing.name] = statusInfo{
				Ref:    ongoing.name,
				Status: resolved,
			}
			keys := []string{ongoing.name}

			activeSeen := map[string]struct{}{}
			if !done {
				active, err := cs.ListStatuses(ctx, "")
				if err != nil {
					log.G(ctx).WithError(err).Error("active check failed")
					continue
				}
				// update status of active entries!
				for _, active := range active {
					statuses[active.Ref] = statusInfo{
						Ref:       active.Ref,
						Status:    "downloading",
						Offset:    active.Offset,
						Total:     active.Total,
						StartedAt: active.StartedAt,
						UpdatedAt: active.UpdatedAt,
					}
					activeSeen[active.Ref] = struct{}{}
				}
			}

			// now, update the items in jobs that are not in active
			for _, j := range ongoing.jobs() {
				key := remotes.MakeRefKey(ctx, j)
				keys = append(keys, key)
				if _, ok := activeSeen[key]; ok {
					continue
				}

				status, ok := statuses[key]
				if !done && (!ok || status.Status == "downloading") {
					info, err := cs.Info(ctx, j.Digest)
					if err != nil {
						if !errdefs.IsNotFound(err) {
							log.G(ctx).WithError(err).Errorf("failed to get content info")
							continue outer
						} else {
							statuses[key] = statusInfo{
								Ref:    key,
								Status: "waiting",
							}
						}
					} else if info.CreatedAt.After(start) {
						statuses[key] = statusInfo{
							Ref:       key,
							Status:    "done",
							Offset:    info.Size,
							Total:     info.Size,
							UpdatedAt: info.CreatedAt,
						}
					} else {
						statuses[key] = statusInfo{
							Ref:    key,
							Status: "exists",
						}
					}
				} else if done {
					if ok {
						if status.Status != "done" && status.Status != "exists" {
							status.Status = "done"
							statuses[key] = status
						}
					} else {
						statuses[key] = statusInfo{
							Ref:    key,
							Status: "done",
						}
					}
				}
			}

			var ordered []statusInfo
			for _, key := range keys {
				ordered = append(ordered, statuses[key])
			}

			display(tw, ordered, start)
			tw.Flush()

			if done {
				fw.Flush()
				return
			}
		case <-ctx.Done():
			done = true // allow ui to update once more
		}
	}
}

// jobs provides a way of identifying the download keys for a particular task
// encountering during the pull walk.
//
// This is very minimal and will probably be replaced with something more
// featured.
type jobs struct {
	name     string
	added    map[digest.Digest]struct{}
	descs    []ocispec.Descriptor
	mu       sync.Mutex
	resolved bool
}

func newJobs(name string) *jobs {
	return &jobs{
		name:  name,
		added: map[digest.Digest]struct{}{},
	}
}

func (j *jobs) add(desc ocispec.Descriptor) {
	j.mu.Lock()
	defer j.mu.Unlock()
	j.resolved = true

	if _, ok := j.added[desc.Digest]; ok {
		return
	}
	j.descs = append(j.descs, desc)
	j.added[desc.Digest] = struct{}{}
}

func (j *jobs) jobs() []ocispec.Descriptor {
	j.mu.Lock()
	defer j.mu.Unlock()

	var descs []ocispec.Descriptor
	return append(descs, j.descs...)
}

func (j *jobs) isResolved() bool {
	j.mu.Lock()
	defer j.mu.Unlock()
	return j.resolved
}

type statusInfo struct {
	Ref       string
	Status    string
	Offset    int64
	Total     int64
	StartedAt time.Time
	UpdatedAt time.Time
}

func display(w io.Writer, statuses []statusInfo, start time.Time) {
	var total int64
	for _, status := range statuses {
		total += status.Offset
		switch status.Status {
		case "downloading", "uploading":
			var bar progress.Bar
			if status.Total > 0.0 {
				bar = progress.Bar(float64(status.Offset) / float64(status.Total))
			}
			fmt.Fprintf(w, "%s:\t%s\t%40r\t%8.8s/%s\t\n",
				status.Ref,
				status.Status,
				bar,
				progress.Bytes(status.Offset), progress.Bytes(status.Total))
		case "resolving", "waiting":
			bar := progress.Bar(0.0)
			fmt.Fprintf(w, "%s:\t%s\t%40r\t\n",
				status.Ref,
				status.Status,
				bar)
		default:
			bar := progress.Bar(1.0)
			fmt.Fprintf(w, "%s:\t%s\t%40r\t\n",
				status.Ref,
				status.Status,
				bar)
		}
	}

	fmt.Fprintf(w, "elapsed: %-4.1fs\ttotal: %7.6v\t(%v)\t\n",
		time.Since(start).Seconds(),
		// TODO(stevvooe): These calculations are actually way off.
		// Need to account for previously downloaded data. These
		// will basically be right for a download the first time
		// but will be skewed if restarting, as it includes the
		// data into the start time before.
		progress.Bytes(total),
		progress.NewBytesPerSecond(total, time.Since(start)))
}
<|endoftext|>