<commit_before>/*
Copyright 2016 The Kubernetes 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

    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 azure

import (
	"fmt"
	"io"
	"io/ioutil"

	"k8s.io/kubernetes/pkg/cloudprovider"
	"k8s.io/kubernetes/pkg/version"

	"github.com/Azure/azure-sdk-for-go/arm/compute"
	"github.com/Azure/azure-sdk-for-go/arm/network"
	"github.com/Azure/azure-sdk-for-go/arm/storage"
	"github.com/Azure/go-autorest/autorest"
	"github.com/Azure/go-autorest/autorest/azure"
	"github.com/ghodss/yaml"
	"time"
)

// CloudProviderName is the value used for the --cloud-provider flag
const CloudProviderName = "azure"

// Config holds the configuration parsed from the --cloud-config flag
// All fields are required unless otherwise specified
type Config struct {
	// The cloud environment identifier. Takes values from https://github.com/Azure/go-autorest/blob/ec5f4903f77ed9927ac95b19ab8e44ada64c1356/autorest/azure/environments.go#L13
	Cloud string `json:"cloud" yaml:"cloud"`
	// The AAD Tenant ID for the Subscription that the cluster is deployed in
	TenantID string `json:"tenantId" yaml:"tenantId"`
	// The ID of the Azure Subscription that the cluster is deployed in
	SubscriptionID string `json:"subscriptionId" yaml:"subscriptionId"`
	// The name of the resource group that the cluster is deployed in
	ResourceGroup string `json:"resourceGroup" yaml:"resourceGroup"`
	// The location of the resource group that the cluster is deployed in
	Location string `json:"location" yaml:"location"`
	// The name of the VNet that the cluster is deployed in
	VnetName string `json:"vnetName" yaml:"vnetName"`
	// The name of the subnet that the cluster is deployed in
	SubnetName string `json:"subnetName" yaml:"subnetName"`
	// The name of the security group attached to the cluster's subnet
	SecurityGroupName string `json:"securityGroupName" yaml:"securityGroupName"`
	// (Optional in 1.6) The name of the route table attached to the subnet that the cluster is deployed in
	RouteTableName string `json:"routeTableName" yaml:"routeTableName"`
	// (Optional) The name of the availability set that should be used as the load balancer backend
	// If this is set, the Azure cloudprovider will only add nodes from that availability set to the load
	// balancer backend pool. If this is not set, and multiple agent pools (availability sets) are used, then
	// the cloudprovider will try to add all nodes to a single backend pool which is forbidden.
	// In other words, if you use multiple agent pools (availability sets), you MUST set this field.
	PrimaryAvailabilitySetName string `json:"primaryAvailabilitySetName" yaml:"primaryAvailabilitySetName"`

	// The ClientID for an AAD application with RBAC access to talk to Azure RM APIs
	AADClientID string `json:"aadClientId" yaml:"aadClientId"`
	// The ClientSecret for an AAD application with RBAC access to talk to Azure RM APIs
	AADClientSecret string `json:"aadClientSecret" yaml:"aadClientSecret"`
}

// Cloud holds the config and clients
type Cloud struct {
	Config
	Environment             azure.Environment
	RoutesClient            network.RoutesClient
	SubnetsClient           network.SubnetsClient
	InterfacesClient        network.InterfacesClient
	RouteTablesClient       network.RouteTablesClient
	LoadBalancerClient      network.LoadBalancersClient
	PublicIPAddressesClient network.PublicIPAddressesClient
	SecurityGroupsClient    network.SecurityGroupsClient
	VirtualMachinesClient   compute.VirtualMachinesClient
	StorageAccountClient    storage.AccountsClient
}

func init() {
	cloudprovider.RegisterCloudProvider(CloudProviderName, NewCloud)
}

// NewCloud returns a Cloud with initialized clients
func NewCloud(configReader io.Reader) (cloudprovider.Interface, error) {
	var az Cloud

	configContents, err := ioutil.ReadAll(configReader)
	if err != nil {
		return nil, err
	}
	err = yaml.Unmarshal(configContents, &az)
	if err != nil {
		return nil, err
	}

	if az.Cloud == "" {
		az.Environment = azure.PublicCloud
	} else {
		az.Environment, err = azure.EnvironmentFromName(az.Cloud)
		if err != nil {
			return nil, err
		}
	}

	oauthConfig, err := az.Environment.OAuthConfigForTenant(az.TenantID)
	if err != nil {
		return nil, err
	}

	servicePrincipalToken, err := azure.NewServicePrincipalToken(
		*oauthConfig,
		az.AADClientID,
		az.AADClientSecret,
		az.Environment.ServiceManagementEndpoint)
	if err != nil {
		return nil, err
	}

	az.SubnetsClient = network.NewSubnetsClient(az.SubscriptionID)
	az.SubnetsClient.BaseURI = az.Environment.ResourceManagerEndpoint
	az.SubnetsClient.Authorizer = servicePrincipalToken
	configureUserAgent(&az.SubnetsClient.Client)

	az.RouteTablesClient = network.NewRouteTablesClient(az.SubscriptionID)
	az.RouteTablesClient.BaseURI = az.Environment.ResourceManagerEndpoint
	az.RouteTablesClient.Authorizer = servicePrincipalToken
	configureUserAgent(&az.RouteTablesClient.Client)

	az.RoutesClient = network.NewRoutesClient(az.SubscriptionID)
	az.RoutesClient.BaseURI = az.Environment.ResourceManagerEndpoint
	az.RoutesClient.Authorizer = servicePrincipalToken
	configureUserAgent(&az.RoutesClient.Client)

	az.InterfacesClient = network.NewInterfacesClient(az.SubscriptionID)
	az.InterfacesClient.BaseURI = az.Environment.ResourceManagerEndpoint
	az.InterfacesClient.Authorizer = servicePrincipalToken
	configureUserAgent(&az.InterfacesClient.Client)

	az.LoadBalancerClient = network.NewLoadBalancersClient(az.SubscriptionID)
	az.LoadBalancerClient.BaseURI = az.Environment.ResourceManagerEndpoint
	az.LoadBalancerClient.Authorizer = servicePrincipalToken
	configureUserAgent(&az.LoadBalancerClient.Client)

	az.VirtualMachinesClient = compute.NewVirtualMachinesClient(az.SubscriptionID)
	az.VirtualMachinesClient.BaseURI = az.Environment.ResourceManagerEndpoint
	az.VirtualMachinesClient.Authorizer = servicePrincipalToken
	az.VirtualMachinesClient.PollingDelay = 5 * time.Second
	configureUserAgent(&az.VirtualMachinesClient.Client)

	az.PublicIPAddressesClient = network.NewPublicIPAddressesClient(az.SubscriptionID)
	az.PublicIPAddressesClient.BaseURI = az.Environment.ResourceManagerEndpoint
	az.PublicIPAddressesClient.Authorizer = servicePrincipalToken
	configureUserAgent(&az.PublicIPAddressesClient.Client)

	az.SecurityGroupsClient = network.NewSecurityGroupsClient(az.SubscriptionID)
	az.SecurityGroupsClient.BaseURI = az.Environment.ResourceManagerEndpoint
	az.SecurityGroupsClient.Authorizer = servicePrincipalToken
	configureUserAgent(&az.SecurityGroupsClient.Client)

	az.StorageAccountClient = storage.NewAccountsClientWithBaseURI(az.Environment.ResourceManagerEndpoint, az.SubscriptionID)
	az.StorageAccountClient.Authorizer = servicePrincipalToken

	return &az, nil
}

// LoadBalancer returns a balancer interface. Also returns true if the interface is supported, false otherwise.
func (az *Cloud) LoadBalancer() (cloudprovider.LoadBalancer, bool) {
	return az, true
}

// Instances returns an instances interface. Also returns true if the interface is supported, false otherwise.
func (az *Cloud) Instances() (cloudprovider.Instances, bool) {
	return az, true
}

// Zones returns a zones interface. Also returns true if the interface is supported, false otherwise.
func (az *Cloud) Zones() (cloudprovider.Zones, bool) {
	return az, true
}

// Clusters returns a clusters interface.  Also returns true if the interface is supported, false otherwise.
func (az *Cloud) Clusters() (cloudprovider.Clusters, bool) {
	return nil, false
}

// Routes returns a routes interface along with whether the interface is supported.
func (az *Cloud) Routes() (cloudprovider.Routes, bool) {
	return az, true
}

// ScrubDNS provides an opportunity for cloud-provider-specific code to process DNS settings for pods.
func (az *Cloud) ScrubDNS(nameservers, searches []string) (nsOut, srchOut []string) {
	return nameservers, searches
}

// ProviderName returns the cloud provider ID.
func (az *Cloud) ProviderName() string {
	return CloudProviderName
}

func configureUserAgent(client *autorest.Client) {
	k8sVersion := version.Get().GitVersion
	client.UserAgent = fmt.Sprintf("%s; %s", client.UserAgent, k8sVersion)
}
<commit_msg>azure: reduce poll delay for all clients to 5 sec<commit_after>/*
Copyright 2016 The Kubernetes 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

    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 azure

import (
	"fmt"
	"io"
	"io/ioutil"

	"k8s.io/kubernetes/pkg/cloudprovider"
	"k8s.io/kubernetes/pkg/version"

	"github.com/Azure/azure-sdk-for-go/arm/compute"
	"github.com/Azure/azure-sdk-for-go/arm/network"
	"github.com/Azure/azure-sdk-for-go/arm/storage"
	"github.com/Azure/go-autorest/autorest"
	"github.com/Azure/go-autorest/autorest/azure"
	"github.com/ghodss/yaml"
	"time"
)

// CloudProviderName is the value used for the --cloud-provider flag
const CloudProviderName = "azure"

// Config holds the configuration parsed from the --cloud-config flag
// All fields are required unless otherwise specified
type Config struct {
	// The cloud environment identifier. Takes values from https://github.com/Azure/go-autorest/blob/ec5f4903f77ed9927ac95b19ab8e44ada64c1356/autorest/azure/environments.go#L13
	Cloud string `json:"cloud" yaml:"cloud"`
	// The AAD Tenant ID for the Subscription that the cluster is deployed in
	TenantID string `json:"tenantId" yaml:"tenantId"`
	// The ID of the Azure Subscription that the cluster is deployed in
	SubscriptionID string `json:"subscriptionId" yaml:"subscriptionId"`
	// The name of the resource group that the cluster is deployed in
	ResourceGroup string `json:"resourceGroup" yaml:"resourceGroup"`
	// The location of the resource group that the cluster is deployed in
	Location string `json:"location" yaml:"location"`
	// The name of the VNet that the cluster is deployed in
	VnetName string `json:"vnetName" yaml:"vnetName"`
	// The name of the subnet that the cluster is deployed in
	SubnetName string `json:"subnetName" yaml:"subnetName"`
	// The name of the security group attached to the cluster's subnet
	SecurityGroupName string `json:"securityGroupName" yaml:"securityGroupName"`
	// (Optional in 1.6) The name of the route table attached to the subnet that the cluster is deployed in
	RouteTableName string `json:"routeTableName" yaml:"routeTableName"`
	// (Optional) The name of the availability set that should be used as the load balancer backend
	// If this is set, the Azure cloudprovider will only add nodes from that availability set to the load
	// balancer backend pool. If this is not set, and multiple agent pools (availability sets) are used, then
	// the cloudprovider will try to add all nodes to a single backend pool which is forbidden.
	// In other words, if you use multiple agent pools (availability sets), you MUST set this field.
	PrimaryAvailabilitySetName string `json:"primaryAvailabilitySetName" yaml:"primaryAvailabilitySetName"`

	// The ClientID for an AAD application with RBAC access to talk to Azure RM APIs
	AADClientID string `json:"aadClientId" yaml:"aadClientId"`
	// The ClientSecret for an AAD application with RBAC access to talk to Azure RM APIs
	AADClientSecret string `json:"aadClientSecret" yaml:"aadClientSecret"`
}

// Cloud holds the config and clients
type Cloud struct {
	Config
	Environment             azure.Environment
	RoutesClient            network.RoutesClient
	SubnetsClient           network.SubnetsClient
	InterfacesClient        network.InterfacesClient
	RouteTablesClient       network.RouteTablesClient
	LoadBalancerClient      network.LoadBalancersClient
	PublicIPAddressesClient network.PublicIPAddressesClient
	SecurityGroupsClient    network.SecurityGroupsClient
	VirtualMachinesClient   compute.VirtualMachinesClient
	StorageAccountClient    storage.AccountsClient
}

func init() {
	cloudprovider.RegisterCloudProvider(CloudProviderName, NewCloud)
}

// NewCloud returns a Cloud with initialized clients
func NewCloud(configReader io.Reader) (cloudprovider.Interface, error) {
	var az Cloud

	configContents, err := ioutil.ReadAll(configReader)
	if err != nil {
		return nil, err
	}
	err = yaml.Unmarshal(configContents, &az)
	if err != nil {
		return nil, err
	}

	if az.Cloud == "" {
		az.Environment = azure.PublicCloud
	} else {
		az.Environment, err = azure.EnvironmentFromName(az.Cloud)
		if err != nil {
			return nil, err
		}
	}

	oauthConfig, err := az.Environment.OAuthConfigForTenant(az.TenantID)
	if err != nil {
		return nil, err
	}

	servicePrincipalToken, err := azure.NewServicePrincipalToken(
		*oauthConfig,
		az.AADClientID,
		az.AADClientSecret,
		az.Environment.ServiceManagementEndpoint)
	if err != nil {
		return nil, err
	}

	az.SubnetsClient = network.NewSubnetsClient(az.SubscriptionID)
	az.SubnetsClient.BaseURI = az.Environment.ResourceManagerEndpoint
	az.SubnetsClient.Authorizer = servicePrincipalToken
	az.SubnetsClient.PollingDelay = 5 * time.Second
	configureUserAgent(&az.SubnetsClient.Client)

	az.RouteTablesClient = network.NewRouteTablesClient(az.SubscriptionID)
	az.RouteTablesClient.BaseURI = az.Environment.ResourceManagerEndpoint
	az.RouteTablesClient.Authorizer = servicePrincipalToken
	az.RouteTablesClient.PollingDelay = 5 * time.Second
	configureUserAgent(&az.RouteTablesClient.Client)

	az.RoutesClient = network.NewRoutesClient(az.SubscriptionID)
	az.RoutesClient.BaseURI = az.Environment.ResourceManagerEndpoint
	az.RoutesClient.Authorizer = servicePrincipalToken
	az.RoutesClient.PollingDelay = 5 * time.Second
	configureUserAgent(&az.RoutesClient.Client)

	az.InterfacesClient = network.NewInterfacesClient(az.SubscriptionID)
	az.InterfacesClient.BaseURI = az.Environment.ResourceManagerEndpoint
	az.InterfacesClient.Authorizer = servicePrincipalToken
	az.InterfacesClient.PollingDelay = 5 * time.Second
	configureUserAgent(&az.InterfacesClient.Client)

	az.LoadBalancerClient = network.NewLoadBalancersClient(az.SubscriptionID)
	az.LoadBalancerClient.BaseURI = az.Environment.ResourceManagerEndpoint
	az.LoadBalancerClient.Authorizer = servicePrincipalToken
	az.LoadBalancerClient.PollingDelay = 5 * time.Second
	configureUserAgent(&az.LoadBalancerClient.Client)

	az.VirtualMachinesClient = compute.NewVirtualMachinesClient(az.SubscriptionID)
	az.VirtualMachinesClient.BaseURI = az.Environment.ResourceManagerEndpoint
	az.VirtualMachinesClient.Authorizer = servicePrincipalToken
	az.VirtualMachinesClient.PollingDelay = 5 * time.Second
	configureUserAgent(&az.VirtualMachinesClient.Client)

	az.PublicIPAddressesClient = network.NewPublicIPAddressesClient(az.SubscriptionID)
	az.PublicIPAddressesClient.BaseURI = az.Environment.ResourceManagerEndpoint
	az.PublicIPAddressesClient.Authorizer = servicePrincipalToken
	az.PublicIPAddressesClient.PollingDelay = 5 * time.Second
	configureUserAgent(&az.PublicIPAddressesClient.Client)

	az.SecurityGroupsClient = network.NewSecurityGroupsClient(az.SubscriptionID)
	az.SecurityGroupsClient.BaseURI = az.Environment.ResourceManagerEndpoint
	az.SecurityGroupsClient.Authorizer = servicePrincipalToken
	az.SecurityGroupsClient.PollingDelay = 5 * time.Second
	configureUserAgent(&az.SecurityGroupsClient.Client)

	az.StorageAccountClient = storage.NewAccountsClientWithBaseURI(az.Environment.ResourceManagerEndpoint, az.SubscriptionID)
	az.StorageAccountClient.Authorizer = servicePrincipalToken

	return &az, nil
}

// LoadBalancer returns a balancer interface. Also returns true if the interface is supported, false otherwise.
func (az *Cloud) LoadBalancer() (cloudprovider.LoadBalancer, bool) {
	return az, true
}

// Instances returns an instances interface. Also returns true if the interface is supported, false otherwise.
func (az *Cloud) Instances() (cloudprovider.Instances, bool) {
	return az, true
}

// Zones returns a zones interface. Also returns true if the interface is supported, false otherwise.
func (az *Cloud) Zones() (cloudprovider.Zones, bool) {
	return az, true
}

// Clusters returns a clusters interface.  Also returns true if the interface is supported, false otherwise.
func (az *Cloud) Clusters() (cloudprovider.Clusters, bool) {
	return nil, false
}

// Routes returns a routes interface along with whether the interface is supported.
func (az *Cloud) Routes() (cloudprovider.Routes, bool) {
	return az, true
}

// ScrubDNS provides an opportunity for cloud-provider-specific code to process DNS settings for pods.
func (az *Cloud) ScrubDNS(nameservers, searches []string) (nsOut, srchOut []string) {
	return nameservers, searches
}

// ProviderName returns the cloud provider ID.
func (az *Cloud) ProviderName() string {
	return CloudProviderName
}

func configureUserAgent(client *autorest.Client) {
	k8sVersion := version.Get().GitVersion
	client.UserAgent = fmt.Sprintf("%s; %s", client.UserAgent, k8sVersion)
}
<|endoftext|>


<commit_before>package minion

import (
	"encoding/gob"
	"encoding/json"
	"errors"
	"fmt"
	"html/template"
	"log"
	"net/http"
	"os"
	"strings"

	"github.com/gorilla/sessions"
)

func init() {
	gob.Register(Principal{})
}

// ErrorFormat defines as which content type an error should be serialized
type ErrorFormat string

const (
	// ErrorAsHTML formats an error using the HTML template `error.html`.
	ErrorAsHTML ErrorFormat = "html"
	// ErrorAsJSON formats the error as JSON object.
	ErrorAsJSON ErrorFormat = "json"
)

// Minion implements basic building blocks that most http servers require
type Minion struct {
	Debug       bool
	Sessions    sessions.Store
	SessionName string
	Templates   *template.Template
	LoginURL    string
	ErrorFormat ErrorFormat
}

// NewMinion creates a new minion instance.
func NewMinion(sessionName string, sessionKey []byte) Minion {
	return Minion{
		Debug:       os.Getenv("DEBUG") == "true",
		Sessions:    sessions.NewCookieStore(sessionKey),
		SessionName: sessionName,
		LoginURL:    "/login",
		ErrorFormat: ErrorAsHTML,
	}
}

// GetSessionValue retrieves a value from the active user session.
func (m Minion) GetSessionValue(w http.ResponseWriter, r *http.Request, name string) (interface{}, error) {
	session, err := m.Sessions.Get(r, m.SessionName)
	if err != nil {
		return nil, err
	}

	return session.Values[name], nil
}

// GetPrincipal returns the currently logged in principal
func (m Minion) GetPrincipal(w http.ResponseWriter, r *http.Request) Principal {
	obj, err := m.GetSessionValue(w, r, "principal")
	if err != nil {
		return Principal{}
	}

	principal, _ := obj.(Principal)
	return principal
}

// StorePrincipal stores a principal in the current browser session.
func (m Minion) StorePrincipal(w http.ResponseWriter, r *http.Request, principal Principal) error {
	session, err := m.Sessions.Get(r, m.SessionName)
	if err != nil {
		return err
	}

	session.Values["principal"] = principal
	return session.Save(r, w)
}

// ClearPrincipal removes the principal from the current browser session.
func (m Minion) ClearPrincipal(w http.ResponseWriter, r *http.Request) error {
	session, err := m.Sessions.Get(r, m.SessionName)
	if err != nil {
		return err
	}

	delete(session.Values, "principal")
	return session.Save(r, w)
}

// Secured requires that the user has at least one of the provided roles before
// the request is forwarded to the secured handler.
func (m Minion) Secured(fn http.HandlerFunc, roles ...string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		principal := m.GetPrincipal(w, r)
		if !principal.Authenticated {
			session, err := m.Sessions.Get(r, m.SessionName)
			if err != nil {
				m.Error(w, r, http.StatusBadRequest, err)
				return
			}

			session.Values["redirect"] = r.URL.String()
			err = session.Save(r, w)
			if err != nil {
				m.Error(w, r, http.StatusInternalServerError, err)
				return
			}

			http.Redirect(w, r, m.LoginURL, http.StatusSeeOther)
			return
		}

		if !principal.HasAnyRole(roles...) {
			m.HTML(w, r, http.StatusForbidden, "403.html", V{})
			return
		}

		fn(w, r)
	}
}

// Error outputs an error using the default error format (HTML with template
// "error.html" or JSON).
func (m Minion) Error(w http.ResponseWriter, r *http.Request, code int, err error) {
	log.Printf("error: %v", err)
	switch m.ErrorFormat {
	case ErrorAsHTML:
		m.HTML(w, r, code, "error.html", V{
			"code":  code,
			"error": err.Error(),
		})

	case ErrorAsJSON:
		m.JSON(w, r, code, V{
			"code":  code,
			"error": err.Error(),
		})
	}
}

// JSON outputs the data encoded as JSON.
func (m Minion) JSON(w http.ResponseWriter, r *http.Request, code int, data interface{}) {
	w.Header().Add("content-type", "application/json; charset=utf-8")
	w.WriteHeader(code)

	err := json.NewEncoder(w).Encode(data)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, "failed to encode json: %v", err)
		log.Printf("failed to encode json: %v", err)
	}
}

// HTML outputs a rendered HTML template to the client. This function also includes
// some default variables into the template scope.
func (m *Minion) HTML(w http.ResponseWriter, r *http.Request, code int, name string, data V) {
	// reload templates in debug mode
	if m.Templates == nil || m.Debug {
		fm := template.FuncMap{
			"div": func(dividend, divisor int) float64 {
				return float64(dividend) / float64(divisor)
			},
			"json": func(v interface{}) template.JS {
				b, _ := json.MarshalIndent(v, "", "  ")
				return template.JS(b)
			},
			"dict": func(values ...interface{}) (map[string]interface{}, error) {
				if len(values)%2 != 0 {
					return nil, errors.New("invalid dict call")
				}
				dict := make(map[string]interface{}, len(values)/2)
				for i := 0; i < len(values); i += 2 {
					key, ok := values[i].(string)
					if !ok {
						return nil, errors.New("dict keys must be strings")
					}
					dict[key] = values[i+1]
				}
				return dict, nil
			},
		}

		var err error
		m.Templates, err = template.New("").Funcs(fm).ParseGlob("templates/*")
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "failed to parse templates: %v", err)
			log.Printf("failed to parse templates: %v", err)
			return
		}
	}

	session, err := m.Sessions.Get(r, m.SessionName)
	if err == nil {
		data["flashes"] = session.Flashes()
		session.Save(r, w)
	}

	principal := m.GetPrincipal(w, r)
	data["principal"] = principal

	w.Header().Add("content-type", "text/html; charset=utf-8")
	err = m.Templates.ExecuteTemplate(w, name, data)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, "failed to execute template %q: %v", name, err)
		log.Printf("failed to execute template %q: %v", name, err)
		return
	}
}

// Principal is an entity that is authenticated and verified.
type Principal struct {
	Authenticated bool
	ID            string
	Login         string
	Roles         string
}

// HasAnyRole checks whether the principal has any of the given roles. Use '*'
// as a wildcard role to match any.
func (u Principal) HasAnyRole(roles ...string) bool {
	if !u.Authenticated {
		return false
	}

	dedup := make(map[string]struct{})
	for _, role := range strings.Split(u.Roles, " ") {
		dedup[role] = struct{}{}
	}

	for _, role := range roles {
		if _, ok := dedup[role]; ok || role == "*" {
			return true
		}
	}

	return false
}

// BindingResult holds validation errors of the binding process from a HTML
// form to a Go struct.
type BindingResult map[string]string

// Valid returns whether the binding was successfull or not.
func (br BindingResult) Valid() bool {
	return len(br) == 0
}

// Fail marks the binding as failed and stores an error for the given field
// that caused the form binding to fail.
func (br BindingResult) Fail(field, err string) {
	br[field] = err
}

// Include copies all errors and state of a binding result
func (br BindingResult) Include(other BindingResult) {
	for field, err := range other {
		br.Fail(field, err)
	}
}

// V is a helper type to quickly build variable maps for templates.
type V map[string]interface{}

// MarshalJSON implements the json.Marshaler interface.
func (v V) MarshalJSON() ([]byte, error) {
	return json.Marshal(map[string]interface{}(v))
}
<commit_msg>simplified interface<commit_after>package minion

import (
	"encoding/gob"
	"encoding/json"
	"errors"
	"fmt"
	"html/template"
	"log"
	"net/http"
	"os"
	"strings"

	"github.com/gorilla/sessions"
)

func init() {
	gob.Register(Principal{})
}

// PrincipalKey is the key used for the principal in the user session.
const PrincipalKey = "__principal__"

// ErrorFormat defines as which content type an error should be serialized
type ErrorFormat string

const (
	// ErrorAsHTML formats an error using the HTML template `error.html`.
	ErrorAsHTML ErrorFormat = "html"
	// ErrorAsJSON formats the error as JSON object.
	ErrorAsJSON ErrorFormat = "json"
)

// Minion implements basic building blocks that most http servers require
type Minion struct {
	Debug       bool
	LoginURL    string
	ErrorFormat ErrorFormat

	sessions    sessions.Store
	sessionName string
	templates   *template.Template
}

// NewMinion creates a new minion instance.
func NewMinion(sessionName string, sessionKey []byte) Minion {
	return Minion{
		Debug:       os.Getenv("DEBUG") == "true",
		LoginURL:    "/login",
		ErrorFormat: ErrorAsHTML,

		sessions:    sessions.NewCookieStore(sessionKey),
		sessionName: sessionName,
	}
}

// Get retrieves a value from the active session. If the value does not
// exist in the session, a provided default is returned
func (m Minion) Get(w http.ResponseWriter, r *http.Request, name string, def interface{}) interface{} {
	session, err := m.sessions.Get(r, m.sessionName)
	if err != nil {
		return def
	}
	value, ok := session.Values[name]
	if !ok {
		return def
	}
	return value
}

// Set stores a value in the active session.
func (m Minion) Set(w http.ResponseWriter, r *http.Request, name string, value interface{}) {
	session, err := m.sessions.Get(r, m.sessionName)
	if err != nil {
		return
	}

	session.Values[name] = value
	session.Save(r, w)
}

// Delete removes a value from the active session.
func (m Minion) Delete(w http.ResponseWriter, r *http.Request, name string) error {
	session, err := m.sessions.Get(r, m.sessionName)
	if err != nil {
		return err
	}

	delete(session.Values, name)
	return session.Save(r, w)
}

// Secured requires that the user has at least one of the provided roles before
// the request is forwarded to the secured handler.
func (m Minion) Secured(fn http.HandlerFunc, roles ...string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		principal := m.Get(w, r, PrincipalKey, Principal{}).(Principal)
		if !principal.Authenticated {
			session, err := m.sessions.Get(r, m.sessionName)
			if err != nil {
				m.Error(w, r, http.StatusBadRequest, err)
				return
			}

			session.Values["redirect"] = r.URL.String()
			err = session.Save(r, w)
			if err != nil {
				m.Error(w, r, http.StatusInternalServerError, err)
				return
			}

			http.Redirect(w, r, m.LoginURL, http.StatusSeeOther)
			return
		}

		if !principal.HasAnyRole(roles...) {
			m.HTML(w, r, http.StatusForbidden, "403.html", V{})
			return
		}

		fn(w, r)
	}
}

// Error outputs an error using the default error format (HTML with template
// "error.html" or JSON).
func (m Minion) Error(w http.ResponseWriter, r *http.Request, code int, err error) {
	log.Printf("error: %v", err)
	switch m.ErrorFormat {
	case ErrorAsHTML:
		m.HTML(w, r, code, "error.html", V{
			"code":  code,
			"error": err.Error(),
		})

	case ErrorAsJSON:
		m.JSON(w, r, code, V{
			"code":  code,
			"error": err.Error(),
		})
	}
}

// JSON outputs the data encoded as JSON.
func (m Minion) JSON(w http.ResponseWriter, r *http.Request, code int, data interface{}) {
	w.Header().Add("content-type", "application/json; charset=utf-8")
	w.WriteHeader(code)

	err := json.NewEncoder(w).Encode(data)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, "failed to encode json: %v", err)
		log.Printf("failed to encode json: %v", err)
	}
}

// HTML outputs a rendered HTML template to the client. This function also includes
// some default variables into the template scope.
func (m *Minion) HTML(w http.ResponseWriter, r *http.Request, code int, name string, data V) {
	// reload templates in debug mode
	if m.templates == nil || m.Debug {
		fm := template.FuncMap{
			"div": func(dividend, divisor int) float64 {
				return float64(dividend) / float64(divisor)
			},
			"json": func(v interface{}) template.JS {
				b, _ := json.MarshalIndent(v, "", "  ")
				return template.JS(b)
			},
			"dict": func(values ...interface{}) (map[string]interface{}, error) {
				if len(values)%2 != 0 {
					return nil, errors.New("invalid dict call")
				}
				dict := make(map[string]interface{}, len(values)/2)
				for i := 0; i < len(values); i += 2 {
					key, ok := values[i].(string)
					if !ok {
						return nil, errors.New("dict keys must be strings")
					}
					dict[key] = values[i+1]
				}
				return dict, nil
			},
		}

		var err error
		m.templates, err = template.New("").Funcs(fm).ParseGlob("templates/*")
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "failed to parse templates: %v", err)
			log.Printf("failed to parse templates: %v", err)
			return
		}
	}

	session, err := m.sessions.Get(r, m.sessionName)
	if err == nil {
		data["flashes"] = session.Flashes()
		session.Save(r, w)
	}

	principal := m.Get(w, r, PrincipalKey, Principal{}).(Principal)
	data["principal"] = principal

	w.Header().Add("content-type", "text/html; charset=utf-8")
	err = m.templates.ExecuteTemplate(w, name, data)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(w, "failed to execute template %q: %v", name, err)
		log.Printf("failed to execute template %q: %v", name, err)
		return
	}
}

// Principal is an entity that is authenticated and verified.
type Principal struct {
	Authenticated bool
	ID            string
	Login         string
	Roles         string
}

// HasAnyRole checks whether the principal has any of the given roles. Use '*'
// as a wildcard role to match any.
func (u Principal) HasAnyRole(roles ...string) bool {
	if !u.Authenticated {
		return false
	}

	dedup := make(map[string]struct{})
	for _, role := range strings.Split(u.Roles, " ") {
		dedup[role] = struct{}{}
	}

	for _, role := range roles {
		if _, ok := dedup[role]; ok || role == "*" {
			return true
		}
	}

	return false
}

// BindingResult holds validation errors of the binding process from a HTML
// form to a Go struct.
type BindingResult map[string]string

// Valid returns whether the binding was successfull or not.
func (br BindingResult) Valid() bool {
	return len(br) == 0
}

// Fail marks the binding as failed and stores an error for the given field
// that caused the form binding to fail.
func (br BindingResult) Fail(field, err string) {
	br[field] = err
}

// Include copies all errors and state of a binding result
func (br BindingResult) Include(other BindingResult) {
	for field, err := range other {
		br.Fail(field, err)
	}
}

// V is a helper type to quickly build variable maps for templates.
type V map[string]interface{}

// MarshalJSON implements the json.Marshaler interface.
func (v V) MarshalJSON() ([]byte, error) {
	return json.Marshal(map[string]interface{}(v))
}
<|endoftext|>


<commit_before>// Copyright 2016 The Netstack Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package sniffer provides the implementation of data-link layer endpoints that
// wrap another endpoint and logs inbound and outbound packets.
//
// Sniffer endpoints can be used in the networking stack by calling New(eID) to
// create a new endpoint, where eID is the ID of the endpoint being wrapped,
// and then passing it as an argument to Stack.CreateNIC().
package sniffer

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"sync/atomic"
	"time"

	"github.com/google/netstack/tcpip"
	"github.com/google/netstack/tcpip/buffer"
	"github.com/google/netstack/tcpip/header"
	"github.com/google/netstack/tcpip/link/rawfile"
	"github.com/google/netstack/tcpip/stack"
	"log"
)

// LogPackets is a flag used to enable or disable packet logging via the log
// package. Valid values are 0 or 1.
//
// LogPackets must be accessed atomically.
var LogPackets uint32 = 1

// LogPacketsToFile is a flag used to enable or disable logging packets to a
// pcap file. Valid values are 0 or 1. A file must have been specified when the
// sniffer was created for this flag to have effect.
//
// LogPacketsToFile must be accessed atomically.
var LogPacketsToFile uint32 = 1

type endpoint struct {
	dispatcher stack.NetworkDispatcher
	lower      stack.LinkEndpoint
	file       *os.File
	maxPCAPLen uint32
}

// New creates a new sniffer link-layer endpoint. It wraps around another
// endpoint and logs packets and they traverse the endpoint.
func New(lower tcpip.LinkEndpointID) tcpip.LinkEndpointID {
	return stack.RegisterLinkEndpoint(&endpoint{
		lower: stack.FindLinkEndpoint(lower),
	})
}

func zoneOffset() (int32, error) {
	loc, err := time.LoadLocation("Local")
	if err != nil {
		return 0, err
	}
	date := time.Date(0, 0, 0, 0, 0, 0, 0, loc)
	_, offset := date.Zone()
	return int32(offset), nil
}

func writePCAPHeader(w io.Writer, maxLen uint32) error {
	offset, err := zoneOffset()
	if err != nil {
		return err
	}
	return binary.Write(w, binary.BigEndian, pcapHeader{
		// From https://wiki.wireshark.org/Development/LibpcapFileFormat
		MagicNumber: 0xa1b2c3d4,

		VersionMajor: 2,
		VersionMinor: 4,
		Thiszone:     offset,
		Sigfigs:      0,
		Snaplen:      maxLen,
		Network:      101, // LINKTYPE_RAW
	})
}

// NewWithFile creates a new sniffer link-layer endpoint. It wraps around
// another endpoint and logs packets and they traverse the endpoint.
//
// Packets can be logged to file in the pcap format in addition to the standard
// human-readable logs.
//
// snapLen is the maximum amount of a packet to be saved. Packets with a length
// less than or equal too snapLen will be saved in their entirety. Longer
// packets will be truncated to snapLen.
func NewWithFile(lower tcpip.LinkEndpointID, file *os.File, snapLen uint32) (tcpip.LinkEndpointID, error) {
	if err := writePCAPHeader(file, snapLen); err != nil {
		return 0, err
	}
	return stack.RegisterLinkEndpoint(&endpoint{
		lower:      stack.FindLinkEndpoint(lower),
		file:       file,
		maxPCAPLen: snapLen,
	}), nil
}

// DeliverNetworkPacket implements the stack.NetworkDispatcher interface. It is
// called by the link-layer endpoint being wrapped when a packet arrives, and
// logs the packet before forwarding to the actual dispatcher.
func (e *endpoint) DeliverNetworkPacket(linkEP stack.LinkEndpoint, remoteLinkAddr tcpip.LinkAddress, protocol tcpip.NetworkProtocolNumber, vv *buffer.VectorisedView) {
	if atomic.LoadUint32(&LogPackets) == 1 {
		LogPacket("recv", protocol, vv.First(), nil)
	}
	if e.file != nil && atomic.LoadUint32(&LogPacketsToFile) == 1 {
		vs := vv.Views()
		bs := make([][]byte, 1, 1+len(vs))
		var length int
		for _, v := range vs {
			if length+len(v) > int(e.maxPCAPLen) {
				l := int(e.maxPCAPLen) - length
				bs = append(bs, []byte(v)[:l])
				length += l
				break
			}
			bs = append(bs, []byte(v))
			length += len(v)
		}
		buf := bytes.NewBuffer(make([]byte, 0, pcapPacketHeaderLen))
		binary.Write(buf, binary.BigEndian, newPCAPPacketHeader(uint32(length), uint32(vv.Size())))
		bs[0] = buf.Bytes()
		if err := rawfile.NonBlockingWriteN(int(e.file.Fd()), bs...); err != nil {
			panic(err)
		}
	}
	e.dispatcher.DeliverNetworkPacket(e, remoteLinkAddr, protocol, vv)
}

// Attach implements the stack.LinkEndpoint interface. It saves the dispatcher
// and registers with the lower endpoint as its dispatcher so that "e" is called
// for inbound packets.
func (e *endpoint) Attach(dispatcher stack.NetworkDispatcher) {
	e.dispatcher = dispatcher
	e.lower.Attach(e)
}

// MTU implements stack.LinkEndpoint.MTU. It just forwards the request to the
// lower endpoint.
func (e *endpoint) MTU() uint32 {
	return e.lower.MTU()
}

// Capabilities implements stack.LinkEndpoint.Capabilities. It just forwards the
// request to the lower endpoint.
func (e *endpoint) Capabilities() stack.LinkEndpointCapabilities {
	return e.lower.Capabilities()
}

// MaxHeaderLength implements the stack.LinkEndpoint interface. It just forwards
// the request to the lower endpoint.
func (e *endpoint) MaxHeaderLength() uint16 {
	return e.lower.MaxHeaderLength()
}

func (e *endpoint) LinkAddress() tcpip.LinkAddress {
	return e.lower.LinkAddress()
}

// WritePacket implements the stack.LinkEndpoint interface. It is called by
// higher-level protocols to write packets; it just logs the packet and forwards
// the request to the lower endpoint.
func (e *endpoint) WritePacket(r *stack.Route, hdr *buffer.Prependable, payload buffer.View, protocol tcpip.NetworkProtocolNumber) *tcpip.Error {
	if atomic.LoadUint32(&LogPackets) == 1 {
		LogPacket("send", protocol, hdr.UsedBytes(), payload)
	}
	if e.file != nil && atomic.LoadUint32(&LogPacketsToFile) == 1 {
		bs := [][]byte{nil, hdr.UsedBytes(), payload}
		var length int

		for i, b := range bs[1:] {
			if rem := int(e.maxPCAPLen) - length; len(b) > rem {
				b = b[:rem]
			}
			bs[i+1] = b
			length += len(b)
		}

		buf := bytes.NewBuffer(make([]byte, 0, pcapPacketHeaderLen))
		binary.Write(buf, binary.BigEndian, newPCAPPacketHeader(uint32(length), uint32(hdr.UsedLength()+len(payload))))
		bs[0] = buf.Bytes()
		if err := rawfile.NonBlockingWriteN(int(e.file.Fd()), bs...); err != nil {
			panic(err)
		}
	}
	return e.lower.WritePacket(r, hdr, payload, protocol)
}

// LogPacket logs the given packet.
func LogPacket(prefix string, protocol tcpip.NetworkProtocolNumber, b, plb []byte) {
	// Figure out the network layer info.
	var transProto uint8
	src := tcpip.Address("unknown")
	dst := tcpip.Address("unknown")
	id := 0
	size := uint16(0)
	switch protocol {
	case header.IPv4ProtocolNumber:
		ipv4 := header.IPv4(b)
		src = ipv4.SourceAddress()
		dst = ipv4.DestinationAddress()
		transProto = ipv4.Protocol()
		size = ipv4.TotalLength() - uint16(ipv4.HeaderLength())
		b = b[ipv4.HeaderLength():]
		id = int(ipv4.ID())

	case header.IPv6ProtocolNumber:
		ipv6 := header.IPv6(b)
		src = ipv6.SourceAddress()
		dst = ipv6.DestinationAddress()
		transProto = ipv6.NextHeader()
		size = ipv6.PayloadLength()
		b = b[header.IPv6MinimumSize:]

	case header.ARPProtocolNumber:
		arp := header.ARP(b)
		log.Printf(
			"%s arp %v (%v) -> %v (%v) valid:%v",
			prefix,
			tcpip.Address(arp.ProtocolAddressSender()), tcpip.LinkAddress(arp.HardwareAddressSender()),
			tcpip.Address(arp.ProtocolAddressTarget()), tcpip.LinkAddress(arp.HardwareAddressTarget()),
			arp.IsValid(),
		)
		return
	default:
		log.Printf("%s unknown network protocol", prefix)
		return
	}

	// Figure out the transport layer info.
	transName := "unknown"
	srcPort := uint16(0)
	dstPort := uint16(0)
	details := ""
	switch tcpip.TransportProtocolNumber(transProto) {
	case header.ICMPv4ProtocolNumber:
		transName = "icmp"
		icmp := header.ICMPv4(b)
		icmpType := "unknown"
		switch icmp.Type() {
		case header.ICMPv4EchoReply:
			icmpType = "echo reply"
		case header.ICMPv4DstUnreachable:
			icmpType = "destination unreachable"
		case header.ICMPv4SrcQuench:
			icmpType = "source quench"
		case header.ICMPv4Redirect:
			icmpType = "redirect"
		case header.ICMPv4Echo:
			icmpType = "echo"
		case header.ICMPv4TimeExceeded:
			icmpType = "time exceeded"
		case header.ICMPv4ParamProblem:
			icmpType = "param problem"
		case header.ICMPv4Timestamp:
			icmpType = "timestamp"
		case header.ICMPv4TimestampReply:
			icmpType = "timestamp reply"
		case header.ICMPv4InfoRequest:
			icmpType = "info request"
		case header.ICMPv4InfoReply:
			icmpType = "info reply"
		}
		log.Printf("%s %s %v -> %v %s len:%d id:%04x code:%d", prefix, transName, src, dst, icmpType, size, id, icmp.Code())
		return

	case header.UDPProtocolNumber:
		transName = "udp"
		udp := header.UDP(b)
		srcPort = udp.SourcePort()
		dstPort = udp.DestinationPort()
		size -= header.UDPMinimumSize

		details = fmt.Sprintf("xsum: 0x%x", udp.Checksum())

	case header.TCPProtocolNumber:
		transName = "tcp"
		tcp := header.TCP(b)
		srcPort = tcp.SourcePort()
		dstPort = tcp.DestinationPort()
		size -= uint16(tcp.DataOffset())

		// Initialize the TCP flags.
		flags := tcp.Flags()
		flagsStr := []byte("FSRPAU")
		for i := range flagsStr {
			if flags&(1<<uint(i)) == 0 {
				flagsStr[i] = ' '
			}
		}
		details = fmt.Sprintf("flags:0x%02x (%v) seqnum: %v ack: %v win: %v xsum:0x%x", flags, string(flagsStr), tcp.SequenceNumber(), tcp.AckNumber(), tcp.WindowSize(), tcp.Checksum())
		if flags&header.TCPFlagSyn != 0 {
			details += fmt.Sprintf(" options: %+v", header.ParseSynOptions(tcp.Options(), flags&header.TCPFlagAck != 0))
		} else {
			details += fmt.Sprintf(" options: %+v", tcp.ParsedOptions())
		}
	default:
		log.Printf("%s %v -> %v unknown transport protocol: %d", prefix, src, dst, transProto)
		return
	}

	log.Printf("%s %s %v:%v -> %v:%v len:%d id:%04x %s", prefix, transName, src, srcPort, dst, dstPort, size, id, details)
}
<commit_msg>Write either packet logs or pcap file. But not both.<commit_after>// Copyright 2016 The Netstack Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package sniffer provides the implementation of data-link layer endpoints that
// wrap another endpoint and logs inbound and outbound packets.
//
// Sniffer endpoints can be used in the networking stack by calling New(eID) to
// create a new endpoint, where eID is the ID of the endpoint being wrapped,
// and then passing it as an argument to Stack.CreateNIC().
package sniffer

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"sync/atomic"
	"time"

	"github.com/google/netstack/tcpip"
	"github.com/google/netstack/tcpip/buffer"
	"github.com/google/netstack/tcpip/header"
	"github.com/google/netstack/tcpip/link/rawfile"
	"github.com/google/netstack/tcpip/stack"
	"log"
)

// LogPackets is a flag used to enable or disable packet logging via the log
// package. Valid values are 0 or 1.
//
// LogPackets must be accessed atomically.
var LogPackets uint32 = 1

// LogPacketsToFile is a flag used to enable or disable logging packets to a
// pcap file. Valid values are 0 or 1. A file must have been specified when the
// sniffer was created for this flag to have effect.
//
// LogPacketsToFile must be accessed atomically.
var LogPacketsToFile uint32 = 1

type endpoint struct {
	dispatcher stack.NetworkDispatcher
	lower      stack.LinkEndpoint
	file       *os.File
	maxPCAPLen uint32
}

// New creates a new sniffer link-layer endpoint. It wraps around another
// endpoint and logs packets and they traverse the endpoint.
func New(lower tcpip.LinkEndpointID) tcpip.LinkEndpointID {
	return stack.RegisterLinkEndpoint(&endpoint{
		lower: stack.FindLinkEndpoint(lower),
	})
}

func zoneOffset() (int32, error) {
	loc, err := time.LoadLocation("Local")
	if err != nil {
		return 0, err
	}
	date := time.Date(0, 0, 0, 0, 0, 0, 0, loc)
	_, offset := date.Zone()
	return int32(offset), nil
}

func writePCAPHeader(w io.Writer, maxLen uint32) error {
	offset, err := zoneOffset()
	if err != nil {
		return err
	}
	return binary.Write(w, binary.BigEndian, pcapHeader{
		// From https://wiki.wireshark.org/Development/LibpcapFileFormat
		MagicNumber: 0xa1b2c3d4,

		VersionMajor: 2,
		VersionMinor: 4,
		Thiszone:     offset,
		Sigfigs:      0,
		Snaplen:      maxLen,
		Network:      101, // LINKTYPE_RAW
	})
}

// NewWithFile creates a new sniffer link-layer endpoint. It wraps around
// another endpoint and logs packets and they traverse the endpoint.
//
// Packets can be logged to file in the pcap format. A sniffer created
// with this function will not emit packets using the standard log
// package.
//
// snapLen is the maximum amount of a packet to be saved. Packets with a length
// less than or equal too snapLen will be saved in their entirety. Longer
// packets will be truncated to snapLen.
func NewWithFile(lower tcpip.LinkEndpointID, file *os.File, snapLen uint32) (tcpip.LinkEndpointID, error) {
	if err := writePCAPHeader(file, snapLen); err != nil {
		return 0, err
	}
	return stack.RegisterLinkEndpoint(&endpoint{
		lower:      stack.FindLinkEndpoint(lower),
		file:       file,
		maxPCAPLen: snapLen,
	}), nil
}

// DeliverNetworkPacket implements the stack.NetworkDispatcher interface. It is
// called by the link-layer endpoint being wrapped when a packet arrives, and
// logs the packet before forwarding to the actual dispatcher.
func (e *endpoint) DeliverNetworkPacket(linkEP stack.LinkEndpoint, remoteLinkAddr tcpip.LinkAddress, protocol tcpip.NetworkProtocolNumber, vv *buffer.VectorisedView) {
	if atomic.LoadUint32(&LogPackets) == 1 && e.file == nil {
		LogPacket("recv", protocol, vv.First(), nil)
	}
	if e.file != nil && atomic.LoadUint32(&LogPacketsToFile) == 1 {
		vs := vv.Views()
		bs := make([][]byte, 1, 1+len(vs))
		var length int
		for _, v := range vs {
			if length+len(v) > int(e.maxPCAPLen) {
				l := int(e.maxPCAPLen) - length
				bs = append(bs, []byte(v)[:l])
				length += l
				break
			}
			bs = append(bs, []byte(v))
			length += len(v)
		}
		buf := bytes.NewBuffer(make([]byte, 0, pcapPacketHeaderLen))
		binary.Write(buf, binary.BigEndian, newPCAPPacketHeader(uint32(length), uint32(vv.Size())))
		bs[0] = buf.Bytes()
		if err := rawfile.NonBlockingWriteN(int(e.file.Fd()), bs...); err != nil {
			panic(err)
		}
	}
	e.dispatcher.DeliverNetworkPacket(e, remoteLinkAddr, protocol, vv)
}

// Attach implements the stack.LinkEndpoint interface. It saves the dispatcher
// and registers with the lower endpoint as its dispatcher so that "e" is called
// for inbound packets.
func (e *endpoint) Attach(dispatcher stack.NetworkDispatcher) {
	e.dispatcher = dispatcher
	e.lower.Attach(e)
}

// MTU implements stack.LinkEndpoint.MTU. It just forwards the request to the
// lower endpoint.
func (e *endpoint) MTU() uint32 {
	return e.lower.MTU()
}

// Capabilities implements stack.LinkEndpoint.Capabilities. It just forwards the
// request to the lower endpoint.
func (e *endpoint) Capabilities() stack.LinkEndpointCapabilities {
	return e.lower.Capabilities()
}

// MaxHeaderLength implements the stack.LinkEndpoint interface. It just forwards
// the request to the lower endpoint.
func (e *endpoint) MaxHeaderLength() uint16 {
	return e.lower.MaxHeaderLength()
}

func (e *endpoint) LinkAddress() tcpip.LinkAddress {
	return e.lower.LinkAddress()
}

// WritePacket implements the stack.LinkEndpoint interface. It is called by
// higher-level protocols to write packets; it just logs the packet and forwards
// the request to the lower endpoint.
func (e *endpoint) WritePacket(r *stack.Route, hdr *buffer.Prependable, payload buffer.View, protocol tcpip.NetworkProtocolNumber) *tcpip.Error {
	if atomic.LoadUint32(&LogPackets) == 1 && e.file == nil {
		LogPacket("send", protocol, hdr.UsedBytes(), payload)
	}
	if e.file != nil && atomic.LoadUint32(&LogPacketsToFile) == 1 {
		bs := [][]byte{nil, hdr.UsedBytes(), payload}
		var length int

		for i, b := range bs[1:] {
			if rem := int(e.maxPCAPLen) - length; len(b) > rem {
				b = b[:rem]
			}
			bs[i+1] = b
			length += len(b)
		}

		buf := bytes.NewBuffer(make([]byte, 0, pcapPacketHeaderLen))
		binary.Write(buf, binary.BigEndian, newPCAPPacketHeader(uint32(length), uint32(hdr.UsedLength()+len(payload))))
		bs[0] = buf.Bytes()
		if err := rawfile.NonBlockingWriteN(int(e.file.Fd()), bs...); err != nil {
			panic(err)
		}
	}
	return e.lower.WritePacket(r, hdr, payload, protocol)
}

// LogPacket logs the given packet.
func LogPacket(prefix string, protocol tcpip.NetworkProtocolNumber, b, plb []byte) {
	// Figure out the network layer info.
	var transProto uint8
	src := tcpip.Address("unknown")
	dst := tcpip.Address("unknown")
	id := 0
	size := uint16(0)
	switch protocol {
	case header.IPv4ProtocolNumber:
		ipv4 := header.IPv4(b)
		src = ipv4.SourceAddress()
		dst = ipv4.DestinationAddress()
		transProto = ipv4.Protocol()
		size = ipv4.TotalLength() - uint16(ipv4.HeaderLength())
		b = b[ipv4.HeaderLength():]
		id = int(ipv4.ID())

	case header.IPv6ProtocolNumber:
		ipv6 := header.IPv6(b)
		src = ipv6.SourceAddress()
		dst = ipv6.DestinationAddress()
		transProto = ipv6.NextHeader()
		size = ipv6.PayloadLength()
		b = b[header.IPv6MinimumSize:]

	case header.ARPProtocolNumber:
		arp := header.ARP(b)
		log.Printf(
			"%s arp %v (%v) -> %v (%v) valid:%v",
			prefix,
			tcpip.Address(arp.ProtocolAddressSender()), tcpip.LinkAddress(arp.HardwareAddressSender()),
			tcpip.Address(arp.ProtocolAddressTarget()), tcpip.LinkAddress(arp.HardwareAddressTarget()),
			arp.IsValid(),
		)
		return
	default:
		log.Printf("%s unknown network protocol", prefix)
		return
	}

	// Figure out the transport layer info.
	transName := "unknown"
	srcPort := uint16(0)
	dstPort := uint16(0)
	details := ""
	switch tcpip.TransportProtocolNumber(transProto) {
	case header.ICMPv4ProtocolNumber:
		transName = "icmp"
		icmp := header.ICMPv4(b)
		icmpType := "unknown"
		switch icmp.Type() {
		case header.ICMPv4EchoReply:
			icmpType = "echo reply"
		case header.ICMPv4DstUnreachable:
			icmpType = "destination unreachable"
		case header.ICMPv4SrcQuench:
			icmpType = "source quench"
		case header.ICMPv4Redirect:
			icmpType = "redirect"
		case header.ICMPv4Echo:
			icmpType = "echo"
		case header.ICMPv4TimeExceeded:
			icmpType = "time exceeded"
		case header.ICMPv4ParamProblem:
			icmpType = "param problem"
		case header.ICMPv4Timestamp:
			icmpType = "timestamp"
		case header.ICMPv4TimestampReply:
			icmpType = "timestamp reply"
		case header.ICMPv4InfoRequest:
			icmpType = "info request"
		case header.ICMPv4InfoReply:
			icmpType = "info reply"
		}
		log.Printf("%s %s %v -> %v %s len:%d id:%04x code:%d", prefix, transName, src, dst, icmpType, size, id, icmp.Code())
		return

	case header.UDPProtocolNumber:
		transName = "udp"
		udp := header.UDP(b)
		srcPort = udp.SourcePort()
		dstPort = udp.DestinationPort()
		size -= header.UDPMinimumSize

		details = fmt.Sprintf("xsum: 0x%x", udp.Checksum())

	case header.TCPProtocolNumber:
		transName = "tcp"
		tcp := header.TCP(b)
		srcPort = tcp.SourcePort()
		dstPort = tcp.DestinationPort()
		size -= uint16(tcp.DataOffset())

		// Initialize the TCP flags.
		flags := tcp.Flags()
		flagsStr := []byte("FSRPAU")
		for i := range flagsStr {
			if flags&(1<<uint(i)) == 0 {
				flagsStr[i] = ' '
			}
		}
		details = fmt.Sprintf("flags:0x%02x (%v) seqnum: %v ack: %v win: %v xsum:0x%x", flags, string(flagsStr), tcp.SequenceNumber(), tcp.AckNumber(), tcp.WindowSize(), tcp.Checksum())
		if flags&header.TCPFlagSyn != 0 {
			details += fmt.Sprintf(" options: %+v", header.ParseSynOptions(tcp.Options(), flags&header.TCPFlagAck != 0))
		} else {
			details += fmt.Sprintf(" options: %+v", tcp.ParsedOptions())
		}
	default:
		log.Printf("%s %v -> %v unknown transport protocol: %d", prefix, src, dst, transProto)
		return
	}

	log.Printf("%s %s %v:%v -> %v:%v len:%d id:%04x %s", prefix, transName, src, srcPort, dst, dstPort, size, id, details)
}
<|endoftext|>


<commit_before>package rpc

import (
	"fmt"
	"net/url"
	"os"
	"sort"
	"strconv"
	"strings"
	"xd/lib/bittorrent/swarm"
	"xd/lib/config"
	"xd/lib/log"
	"xd/lib/rpc"
	t "xd/lib/translate"
	"xd/lib/util"
	"xd/lib/version"
)

func formatRate(r float64) string {
	str := util.FormatRate(r)
	for len(str) < 12 {
		str += " "
	}
	return str
}

// Run runs xd-cli main function
func Run() {
	var args []string
	cmd := "help"
	fname := "torrents.ini"
	if len(os.Args) > 1 {
		cmd = os.Args[1]
		args = os.Args[2:]
	}
	cfg := new(config.Config)
	err := cfg.Load(fname)
	if err != nil {
		log.Errorf("error: %s", err)
		return
	}
	log.SetLevel(cfg.Log.Level)
	var rpcURL string
	if strings.HasPrefix(cfg.RPC.Bind, "unix:") {
		rpcURL = cfg.RPC.Bind
	} else {
		u := url.URL{
			Scheme: "http",
			Host:   cfg.RPC.Bind,
			Path:   rpc.RPCPath,
		}
		rpcURL = u.String()
	}
	swarms := cfg.Bittorrent.Swarms
	count := 0
	switch strings.ToLower(cmd) {
	case "list":
		for count < swarms {
			c := rpc.NewClient(rpcURL, count)
			listTorrents(c)
			count++
		}
	case "add":
		for count < swarms {
			c := rpc.NewClient(rpcURL, count)
			addTorrents(c, args...)
			count++
		}
	case "start":
		for count < swarms {
			c := rpc.NewClient(rpcURL, count)
			startTorrents(c, args...)
			count++
		}
	case "stop":
		for count < swarms {
			c := rpc.NewClient(rpcURL, count)
			stopTorrents(c, args...)
			count++
		}
	case "remove":
		for count < swarms {
			c := rpc.NewClient(rpcURL, count)
			removeTorrents(c, args...)
			count++
		}
	case "delete":
		for count < swarms {
			c := rpc.NewClient(rpcURL, count)
			deleteTorrents(c, args...)
			count++
		}
	case "set-piece-window":
		for count < swarms {
			c := rpc.NewClient(rpcURL, count)
			setPieceWindow(c, args[0])
			count++
		}
	case "version":
		fmt.Println(version.Version())
	case "help":
		printHelp(os.Args[0])
	}
}

func printHelp(cmd string) {
	fmt.Println(t.T("usage: %s [help|version|list|add http://somesite.i2p/some.torrent|set-piece-window n|remove infohash|delete infohash|stop infohash|start infohash]", cmd))
}

func setPieceWindow(c *rpc.Client, str string) {
	n, err := strconv.Atoi(str)
	if err != nil {
		log.Fatalf("error: %s", err.Error())
	}
	c.SetPieceWindow(n)
}

func addTorrents(c *rpc.Client, urls ...string) {
	for idx := range urls {
		fmt.Println(t.T("fetch %s ... ", urls[idx]))
		err := c.AddTorrent(urls[idx])
		if err == nil {
			fmt.Println(t.T("OK"))
		} else {
			fmt.Println(t.E(err))
		}
	}
}

func startTorrents(c *rpc.Client, ih ...string) {
	for idx := range ih {
		fmt.Println(t.T("start %s ... ", ih[idx]))
		err := c.AddTorrent(ih[idx])
		if err == nil {
			fmt.Println(t.T("OK"))
		} else {
			fmt.Println(t.E(err))
		}
	}
}

func stopTorrents(c *rpc.Client, ih ...string) {
	for idx := range ih {
		fmt.Println(t.T("stop %s ... ", ih[idx]))
		err := c.StopTorrent(ih[idx])
		if err == nil {
			fmt.Println(t.T("OK"))
		} else {
			fmt.Println(t.E(err))
		}
	}
}

func removeTorrents(c *rpc.Client, ih ...string) {
	for idx := range ih {
		fmt.Println(t.T("remove %s ... ", ih[idx]))
		err := c.RemoveTorrent(ih[idx])
		if err == nil {
			fmt.Println(t.T("OK"))
		} else {
			fmt.Println(t.E(err))
		}
	}
}

func deleteTorrents(c *rpc.Client, ih ...string) {
	for idx := range ih {
		fmt.Println(t.T("delete %s ... ", ih[idx]))
		err := c.DeleteTorrent(ih[idx])
		if err == nil {
			fmt.Println(t.T("OK"))
		} else {
			fmt.Println(t.E(err))
		}
	}
}

func listTorrents(c *rpc.Client) {
	var err error
	var st swarm.SwarmStatus
	st, err = c.GetSwarmStatus()
	if err != nil {
		log.Errorf("rpc error: %s", err)
		return
	}

	var torrents swarm.TorrentStatusList
	for _, status := range st {
		torrents = append(torrents, status)
	}
	sort.Stable(&torrents)
	for _, status := range torrents {
		fmt.Printf("%s [%s] %s %.2f\n", status.Name, status.Infohash, t.T("progress:"), status.Progress*100)
		fmt.Println(t.T("peers:"))
		sort.Stable(&status.Peers)
		for _, peer := range status.Peers {
			pad := peer.ID

			for len(pad) < 65 {
				pad += " "
			}
			fmt.Printf("\t%stx=%s rx=%s\n", pad, formatRate(peer.TX), formatRate(peer.RX))
		}
		fmt.Printf("%s tx=%s rx=%s (%s: %.2f)\n", status.State, formatRate(status.Peers.TX()), formatRate(status.Peers.RX()), t.T("ratio"), status.Ratio())
		fmt.Println(t.T("files:"))
		for idx, f := range status.Files {
			fmt.Printf("\t[%d] %s (%s: %.2f)\n", idx, f.FileInfo.Path.FilePath(""), t.T("progress:"), f.Progress)
		}
		fmt.Println()
	}
	fmt.Println()
	tx, rx := st.TotalSpeed()
	fmt.Printf("%s: tx=%s rx=%s (%.2f ratio)\n", t.TN("%d torrent", "%d torrents", torrents.Len()), formatRate(tx), formatRate(rx), st.Ratio())
	fmt.Println()
	fmt.Println()
}
<commit_msg>add integer into rpc TN<commit_after>package rpc

import (
	"fmt"
	"net/url"
	"os"
	"sort"
	"strconv"
	"strings"
	"xd/lib/bittorrent/swarm"
	"xd/lib/config"
	"xd/lib/log"
	"xd/lib/rpc"
	t "xd/lib/translate"
	"xd/lib/util"
	"xd/lib/version"
)

func formatRate(r float64) string {
	str := util.FormatRate(r)
	for len(str) < 12 {
		str += " "
	}
	return str
}

// Run runs xd-cli main function
func Run() {
	var args []string
	cmd := "help"
	fname := "torrents.ini"
	if len(os.Args) > 1 {
		cmd = os.Args[1]
		args = os.Args[2:]
	}
	cfg := new(config.Config)
	err := cfg.Load(fname)
	if err != nil {
		log.Errorf("error: %s", err)
		return
	}
	log.SetLevel(cfg.Log.Level)
	var rpcURL string
	if strings.HasPrefix(cfg.RPC.Bind, "unix:") {
		rpcURL = cfg.RPC.Bind
	} else {
		u := url.URL{
			Scheme: "http",
			Host:   cfg.RPC.Bind,
			Path:   rpc.RPCPath,
		}
		rpcURL = u.String()
	}
	swarms := cfg.Bittorrent.Swarms
	count := 0
	switch strings.ToLower(cmd) {
	case "list":
		for count < swarms {
			c := rpc.NewClient(rpcURL, count)
			listTorrents(c)
			count++
		}
	case "add":
		for count < swarms {
			c := rpc.NewClient(rpcURL, count)
			addTorrents(c, args...)
			count++
		}
	case "start":
		for count < swarms {
			c := rpc.NewClient(rpcURL, count)
			startTorrents(c, args...)
			count++
		}
	case "stop":
		for count < swarms {
			c := rpc.NewClient(rpcURL, count)
			stopTorrents(c, args...)
			count++
		}
	case "remove":
		for count < swarms {
			c := rpc.NewClient(rpcURL, count)
			removeTorrents(c, args...)
			count++
		}
	case "delete":
		for count < swarms {
			c := rpc.NewClient(rpcURL, count)
			deleteTorrents(c, args...)
			count++
		}
	case "set-piece-window":
		for count < swarms {
			c := rpc.NewClient(rpcURL, count)
			setPieceWindow(c, args[0])
			count++
		}
	case "version":
		fmt.Println(version.Version())
	case "help":
		printHelp(os.Args[0])
	}
}

func printHelp(cmd string) {
	fmt.Println(t.T("usage: %s [help|version|list|add http://somesite.i2p/some.torrent|set-piece-window n|remove infohash|delete infohash|stop infohash|start infohash]", cmd))
}

func setPieceWindow(c *rpc.Client, str string) {
	n, err := strconv.Atoi(str)
	if err != nil {
		log.Fatalf("error: %s", err.Error())
	}
	c.SetPieceWindow(n)
}

func addTorrents(c *rpc.Client, urls ...string) {
	for idx := range urls {
		fmt.Println(t.T("fetch %s ... ", urls[idx]))
		err := c.AddTorrent(urls[idx])
		if err == nil {
			fmt.Println(t.T("OK"))
		} else {
			fmt.Println(t.E(err))
		}
	}
}

func startTorrents(c *rpc.Client, ih ...string) {
	for idx := range ih {
		fmt.Println(t.T("start %s ... ", ih[idx]))
		err := c.AddTorrent(ih[idx])
		if err == nil {
			fmt.Println(t.T("OK"))
		} else {
			fmt.Println(t.E(err))
		}
	}
}

func stopTorrents(c *rpc.Client, ih ...string) {
	for idx := range ih {
		fmt.Println(t.T("stop %s ... ", ih[idx]))
		err := c.StopTorrent(ih[idx])
		if err == nil {
			fmt.Println(t.T("OK"))
		} else {
			fmt.Println(t.E(err))
		}
	}
}

func removeTorrents(c *rpc.Client, ih ...string) {
	for idx := range ih {
		fmt.Println(t.T("remove %s ... ", ih[idx]))
		err := c.RemoveTorrent(ih[idx])
		if err == nil {
			fmt.Println(t.T("OK"))
		} else {
			fmt.Println(t.E(err))
		}
	}
}

func deleteTorrents(c *rpc.Client, ih ...string) {
	for idx := range ih {
		fmt.Println(t.T("delete %s ... ", ih[idx]))
		err := c.DeleteTorrent(ih[idx])
		if err == nil {
			fmt.Println(t.T("OK"))
		} else {
			fmt.Println(t.E(err))
		}
	}
}

func listTorrents(c *rpc.Client) {
	var err error
	var st swarm.SwarmStatus
	st, err = c.GetSwarmStatus()
	if err != nil {
		log.Errorf("rpc error: %s", err)
		return
	}

	var torrents swarm.TorrentStatusList
	for _, status := range st {
		torrents = append(torrents, status)
	}
	sort.Stable(&torrents)
	for _, status := range torrents {
		fmt.Printf("%s [%s] %s %.2f\n", status.Name, status.Infohash, t.T("progress:"), status.Progress*100)
		fmt.Println(t.T("peers:"))
		sort.Stable(&status.Peers)
		for _, peer := range status.Peers {
			pad := peer.ID

			for len(pad) < 65 {
				pad += " "
			}
			fmt.Printf("\t%stx=%s rx=%s\n", pad, formatRate(peer.TX), formatRate(peer.RX))
		}
		fmt.Printf("%s tx=%s rx=%s (%s: %.2f)\n", status.State, formatRate(status.Peers.TX()), formatRate(status.Peers.RX()), t.T("ratio"), status.Ratio())
		fmt.Println(t.T("files:"))
		for idx, f := range status.Files {
			fmt.Printf("\t[%d] %s (%s: %.2f)\n", idx, f.FileInfo.Path.FilePath(""), t.T("progress:"), f.Progress)
		}
		fmt.Println()
	}
	fmt.Println()
	tx, rx := st.TotalSpeed()
	fmt.Printf("%s: tx=%s rx=%s (%.2f ratio)\n", t.TN("%d torrent", "%d torrents", torrents.Len(), torrents.Len()), formatRate(tx), formatRate(rx), st.Ratio())
	fmt.Println()
	fmt.Println()
}
<|endoftext|>


<commit_before>package main

import (
	"reflect"
	"testing"
)

type t2mInput struct {
	set     noteMap
	txt     []string
	figures map[string]figure
}

type t2mOutput struct {
	channels matrix
	notes    matrix
	mfigures map[int][]midiFigure
}

type t2mTestPair struct {
	in  t2mInput
	out t2mOutput
}

var t2mTestPairs = []t2mTestPair{
	{
		in: t2mInput{
			set: noteMap{
				"ab": midiNote{
					Channel: 1,
					Note:    60,
				},
			},
			txt: []string{
				"ab +f1 -- --",
				"--  ab -- ab",
			},
			figures: map[string]figure{
				"f1": figure{"ab", 111, "x.x"},
			},
		},
		out: t2mOutput{
			channels: matrix{
				row{1, 0, 0, 0},
				row{0, 1, 0, 1},
			},
			notes: matrix{
				row{60, 0, 0, 0},
				row{0, 60, 0, 60},
			},
			mfigures: map[int][]midiFigure{
				1: []midiFigure{
					midiFigure{
						midiEvent{1, 60, 111},
						midiEvent{},
						midiEvent{1, 60, 111},
					},
				},
			},
		},
	},
	{
		in: t2mInput{
			set: noteMap{
				"ab": midiNote{
					Channel: 2,
					Note:    70,
				},
			},
			txt: []string{
				"ab ab . . .",
				". ab +f1 ab",
			},
			figures: map[string]figure{
				"f1": figure{"ab", 109, ".x."},
			},
		},
		out: t2mOutput{
			channels: matrix{
				row{2, 2, 0, 0, 0},
				row{0, 2, 0, 2},
			},
			notes: matrix{
				row{70, 70, 0, 0, 0},
				row{0, 70, 0, 70},
			},
			mfigures: map[int][]midiFigure{
				2: []midiFigure{
					midiFigure{
						midiEvent{},
						midiEvent{2, 70, 109},
						midiEvent{},
					},
				},
			},
		},
	},
}

func TestT2m(t *testing.T) {
	for _, pair := range t2mTestPairs {
		gotC, gotN, gotF := text2matrix(pair.in.set, pair.in.figures, pair.in.txt)
		if !gotC.eq(pair.out.channels) {
			t.Errorf("got %v, wanted %v", gotC, pair.out.channels)
		}
		if !gotN.eq(pair.out.notes) {
			t.Errorf("got %v, wanted %v", gotN, pair.out.notes)
		}
		if !reflect.DeepEqual(gotF, pair.out.mfigures) {
			t.Errorf("got %v, wanted %v", gotF, pair.out.mfigures)
		}
	}
}
<commit_msg>Test for translateFigure()<commit_after>package main

import (
	"reflect"
	"testing"
)

type t2mInput struct {
	set     noteMap
	txt     []string
	figures map[string]figure
}

type t2mOutput struct {
	channels matrix
	notes    matrix
	mfigures map[int][]midiFigure
}

type t2mTestPair struct {
	in  t2mInput
	out t2mOutput
}

var t2mTestPairs = []t2mTestPair{
	{
		in: t2mInput{
			set: noteMap{
				"ab": midiNote{
					Channel: 1,
					Note:    60,
				},
			},
			txt: []string{
				"ab +f1 -- --",
				"--  ab -- ab",
			},
			figures: map[string]figure{
				"f1": figure{"ab", 111, "x.x"},
			},
		},
		out: t2mOutput{
			channels: matrix{
				row{1, 0, 0, 0},
				row{0, 1, 0, 1},
			},
			notes: matrix{
				row{60, 0, 0, 0},
				row{0, 60, 0, 60},
			},
			mfigures: map[int][]midiFigure{
				1: []midiFigure{
					midiFigure{
						midiEvent{1, 60, 111},
						midiEvent{},
						midiEvent{1, 60, 111},
					},
				},
			},
		},
	},
	{
		in: t2mInput{
			set: noteMap{
				"ab": midiNote{
					Channel: 2,
					Note:    70,
				},
			},
			txt: []string{
				"ab ab . . .",
				". ab +f1 ab",
			},
			figures: map[string]figure{
				"f1": figure{"ab", 109, ".x."},
			},
		},
		out: t2mOutput{
			channels: matrix{
				row{2, 2, 0, 0, 0},
				row{0, 2, 0, 2},
			},
			notes: matrix{
				row{70, 70, 0, 0, 0},
				row{0, 70, 0, 70},
			},
			mfigures: map[int][]midiFigure{
				2: []midiFigure{
					midiFigure{
						midiEvent{},
						midiEvent{2, 70, 109},
						midiEvent{},
					},
				},
			},
		},
	},
}

func TestT2m(t *testing.T) {
	for _, pair := range t2mTestPairs {
		gotC, gotN, gotF := text2matrix(pair.in.set, pair.in.figures, pair.in.txt)
		if !gotC.eq(pair.out.channels) {
			t.Errorf("got %v, wanted %v", gotC, pair.out.channels)
		}
		if !gotN.eq(pair.out.notes) {
			t.Errorf("got %v, wanted %v", gotN, pair.out.notes)
		}
		if !reflect.DeepEqual(gotF, pair.out.mfigures) {
			t.Errorf("got %v, wanted %v", gotF, pair.out.mfigures)
		}
	}
}

type tlfTestPair struct {
	in  string
	out midiFigure
}

var tlfTestPairs = []tlfTestPair{
	{
		in:  "f1",
		out: midiFigure{{0, 0, 0}, {7, 65, 65}, {7, 65, 65}},
	},
	{
		in:  "f2",
		out: midiFigure{{0, 0, 0}, {0, 0, 0}, {5, 50, 100}},
	},
	{
		in:  "h1",
		out: midiFigure{{0, 0, 0}, {0, 0, 0}, {5, 70, 80}},
	},
	{
		in:  "qu",
		out: midiFigure{{5, 58, 95}, {5, 58, 95}, {5, 58, 95}},
	},
}

func TestFigures(t *testing.T) {
	drums := new(drums)
	drums.loadFromFile("../testfiles/beat8.yml")
	sets := drums.getSets()
	figures := drums.getFigures()
	for _, fig := range tlfTestPairs {
		got := translateFigure(sets[defaultSet], figures, fig.in)
		if !reflect.DeepEqual(got, fig.out) {
			t.Errorf("%s: got %v, wanted %v", fig.in, got, fig.out)
		}
	}
}
<|endoftext|>


<commit_before>package dashboard

import (
	"fmt"
	"github.com/Cepave/fe/g"
	"github.com/astaxie/beego/orm"
	"regexp"
)

func QueryEndpointidbyNames(endpoints []string, limit int) (enp []Endpoint, err error) {
	q := orm.NewOrm()
	q.Using("graph")
	q.QueryTable("endpoint")
	qb, _ := orm.NewQueryBuilder("mysql")
	qt := qb.Select("*").From("endpoint").Where("endpoint").In(endpoints...).Limit(limit)
	_, err = q.Raw(qt.String()).QueryRows(&enp)
	return
}

func QueryCounterByEndpoints(endpoints []string, limit int) (counters []string, err error) {
	config := g.Config()
	if limit == 0 || limit > config.GraphDB.Limit {
		limit = config.GraphDB.Limit
	}
	enp, aerr := QueryEndpointidbyNames(endpoints, limit)
	if aerr != nil {
		err = aerr
		return
	}
	q := orm.NewOrm()
	q.Using("graph")
	q.QueryTable("endpoint_counter")
	var endpoint_ids = ""
	for _, v := range enp {
		endpoint_ids += fmt.Sprintf("%d,", v.Id)
	}

	pattn, _ := regexp.Compile("\\s*,\\s*$")
	queryperfix := fmt.Sprintf("select distinct(counter) from endpoint_counter where endpoint_id IN(%s) limit %d", pattn.ReplaceAllString(endpoint_ids, ""), limit)
	var enpc []EndpointCounter
	_, err = q.Raw(queryperfix).QueryRows(&enpc)
	for _, v := range enpc {
		counters = append(counters, v.Counter)
	}
	return
}
<commit_msg>[OWL-688] fix invalid MySQL query if the endpoint not found.<commit_after>package dashboard

import (
	"errors"
	"fmt"
	"github.com/Cepave/fe/g"
	"github.com/astaxie/beego/orm"
	"regexp"
)

func QueryEndpointidbyNames(endpoints []string, limit int) (enp []Endpoint, err error) {
	q := orm.NewOrm()
	q.Using("graph")
	q.QueryTable("endpoint")
	qb, _ := orm.NewQueryBuilder("mysql")
	qt := qb.Select("*").From("endpoint").Where("endpoint").In(endpoints...).Limit(limit)
	_, err = q.Raw(qt.String()).QueryRows(&enp)
	return
}

func QueryCounterByEndpoints(endpoints []string, limit int) (counters []string, err error) {
	config := g.Config()
	if limit == 0 || limit > config.GraphDB.Limit {
		limit = config.GraphDB.Limit
	}
	enp, aerr := QueryEndpointidbyNames(endpoints, limit)
	if aerr != nil {
		err = aerr
		return
	}
	q := orm.NewOrm()
	q.Using("graph")
	q.QueryTable("endpoint_counter")
	var endpoint_ids = ""
	for _, v := range enp {
		endpoint_ids += fmt.Sprintf("%d,", v.Id)
	}

	pattn, _ := regexp.Compile("\\s*,\\s*$")
	queryperfix := fmt.Sprintf("select distinct(counter) from endpoint_counter where endpoint_id IN(%s) limit %d", pattn.ReplaceAllString(endpoint_ids, ""), limit)
	var enpc []EndpointCounter
	if len(enp) != 0 {
		_, err = q.Raw(queryperfix).QueryRows(&enpc)
		for _, v := range enpc {
			counters = append(counters, v.Counter)
		}
	} else {
		err = errors.New("The endpoints doesn't exist.")
	}

	return
}
<|endoftext|>


<commit_before>package bind

import (
	"github.com/timeredbull/tsuru/api/unit"
)

// AppContainer provides methdos for a container of apps.
//
// The container stores only the names of the apps.
type AppContainer interface {
	// Adds an app to the container.
	AddApp(string) error

	// Finds an app in the container, returning an index a value >= 0 if it is
	// present, and -1 if not present.
	FindApp(string) int

	// Removes an app form the container.
	RemoveApp(name string) error
}

type EnvVar struct {
	Name         string
	Value        string
	Public       bool
	InstanceName string
}

type App interface {
	GetUnits() []unit.Unit
	SetEnvs([]EnvVar, bool) error
	UnsetEnvs([]string, bool) error
}

type Binder interface {
	Bind(AppContainer, App) error
	Unbind(AppContainer, App) error
}
<commit_msg>bind: changed Binder interface<commit_after>package bind

import (
	"github.com/timeredbull/tsuru/api/unit"
)

// AppContainer provides methdos for a container of apps.
//
// The container stores only the names of the apps.
type AppContainer interface {
	// Adds an app to the container.
	AddApp(string) error

	// Finds an app in the container, returning an index a value >= 0 if it is
	// present, and -1 if not present.
	FindApp(string) int

	// Removes an app form the container.
	RemoveApp(name string) error
}

type EnvVar struct {
	Name         string
	Value        string
	Public       bool
	InstanceName string
}

type App interface {
	GetUnits() []unit.Unit
	SetEnvs([]EnvVar, bool) error
	UnsetEnvs([]string, bool) error
}

type Binder interface {
	AppContainer
	Bind(App) error
	Unbind(App) error
}
<|endoftext|>


<commit_before>package resourcecollector

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/heptio/ark/pkg/discovery"
	"github.com/heptio/ark/pkg/util/collections"
	"github.com/libopenstorage/stork/drivers/volume"
	"github.com/portworx/sched-ops/k8s"
	"github.com/sirupsen/logrus"
	rbacv1 "k8s.io/api/rbac/v1"
	apiextensionsclient "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/rest"
	"k8s.io/kubernetes/pkg/registry/core/service/portallocator"
)

const (
	// Annotation to use when the resource shouldn't be collected
	skipResourceAnnotation = "stork.libopenstorage.ord/skipresource"
)

// ResourceCollector is used to collect and process unstructured objects in namespaces and using label selectors
type ResourceCollector struct {
	Driver           volume.Driver
	discoveryHelper  discovery.Helper
	dynamicInterface dynamic.Interface
}

// Init initializes the resource collector
func (r *ResourceCollector) Init() error {
	config, err := rest.InClusterConfig()
	if err != nil {
		return fmt.Errorf("error getting cluster config: %v", err)
	}

	aeclient, err := apiextensionsclient.NewForConfig(config)
	if err != nil {
		return fmt.Errorf("error getting apiextension client, %v", err)
	}

	discoveryClient := aeclient.Discovery()
	r.discoveryHelper, err = discovery.NewHelper(discoveryClient, logrus.New())
	if err != nil {
		return err
	}
	err = r.discoveryHelper.Refresh()
	if err != nil {
		return err
	}
	r.dynamicInterface, err = dynamic.NewForConfig(config)
	if err != nil {
		return err
	}
	return nil
}

func resourceToBeCollected(resource metav1.APIResource) bool {
	switch resource.Kind {
	case "PersistentVolumeClaim",
		"PersistentVolume",
		"Deployment",
		"DeploymentConfig",
		"StatefulSet",
		"ConfigMap",
		"Service",
		"Secret",
		"DaemonSet",
		"ServiceAccount",
		"Role",
		"RoleBinding",
		"ClusterRole",
		"ClusterRoleBinding",
		"ImageStream",
		"Ingress",
		"Route":
		return true
	default:
		return false
	}
}

// GetResources gets all the resources in the given list of namespaces which match the labelSelectors
func (r *ResourceCollector) GetResources(namespaces []string, labelSelectors map[string]string) ([]runtime.Unstructured, error) {
	err := r.discoveryHelper.Refresh()
	if err != nil {
		return nil, err
	}
	allObjects := make([]runtime.Unstructured, 0)

	// Map to prevent collection of duplicate objects
	resourceMap := make(map[types.UID]bool)

	crbs, err := k8s.Instance().ListClusterRoleBindings()
	if err != nil {
		return nil, err
	}

	for _, group := range r.discoveryHelper.Resources() {
		groupVersion, err := schema.ParseGroupVersion(group.GroupVersion)
		if err != nil {
			return nil, err
		}

		for _, resource := range group.APIResources {
			if !resourceToBeCollected(resource) {
				continue
			}

			for _, ns := range namespaces {
				var dynamicClient dynamic.ResourceInterface
				if !resource.Namespaced {
					dynamicClient = r.dynamicInterface.Resource(groupVersion.WithResource(resource.Name))
				} else {
					dynamicClient = r.dynamicInterface.Resource(groupVersion.WithResource(resource.Name)).Namespace(ns)
				}

				var selectors string
				// PVs don't get the labels from their PVCs, so don't use the label selector
				// Also skip for some other resources that aren't necessarily tied to an application
				switch resource.Kind {
				case "PersistentVolume",
					"ClusterRoleBinding",
					"ClusterRole",
					"ServiceAccount":
				default:
					selectors = labels.Set(labelSelectors).String()
				}
				objectsList, err := dynamicClient.List(metav1.ListOptions{
					LabelSelector: selectors,
				})
				if err != nil {
					return nil, err
				}
				objects, err := meta.ExtractList(objectsList)
				if err != nil {
					return nil, err
				}
				for _, o := range objects {
					runtimeObject, ok := o.(runtime.Unstructured)
					if !ok {
						return nil, fmt.Errorf("error casting object: %v", o)
					}

					collect, err := r.objectToBeCollected(labelSelectors, resourceMap, runtimeObject, crbs, ns)
					if err != nil {
						return nil, fmt.Errorf("error processing object %v: %v", runtimeObject, err)
					}
					if !collect {
						continue
					}
					metadata, err := meta.Accessor(runtimeObject)
					if err != nil {
						return nil, err
					}
					allObjects = append(allObjects, runtimeObject)
					resourceMap[metadata.GetUID()] = true
				}
			}
		}
	}

	err = r.prepareResourcesForCollection(allObjects, namespaces)
	if err != nil {
		return nil, err
	}
	return allObjects, nil
}

// Returns whether an object should be collected or not for the requested
// namespace
func (r *ResourceCollector) objectToBeCollected(
	labelSelectors map[string]string,
	resourceMap map[types.UID]bool,
	object runtime.Unstructured,
	crbs *rbacv1.ClusterRoleBindingList,
	namespace string,
) (bool, error) {
	metadata, err := meta.Accessor(object)
	if err != nil {
		return false, err
	}

	if value, present := metadata.GetAnnotations()[skipResourceAnnotation]; present {
		if skip, err := strconv.ParseBool(value); err == nil && skip {
			return false, err
		}
	}

	// Skip if we've already processed this object
	if _, ok := resourceMap[metadata.GetUID()]; ok {
		return false, nil
	}

	objectType, err := meta.TypeAccessor(object)
	if err != nil {
		return false, err
	}

	switch objectType.GetKind() {
	case "Service":
		return r.serviceToBeCollected(object)
	case "PersistentVolumeClaim":
		return r.pvcToBeCollected(object, namespace)
	case "PersistentVolume":
		return r.pvToBeCollected(labelSelectors, object, namespace)
	case "ClusterRoleBinding":
		return r.clusterRoleBindingToBeCollected(labelSelectors, object, namespace)
	case "ClusterRole":
		return r.clusterRoleToBeCollected(labelSelectors, object, crbs, namespace)
	case "ServiceAccount":
		return r.serviceAccountToBeCollected(object)
	case "Secret":
		return r.secretToBeCollected(object)
	}

	return true, nil
}

func (r *ResourceCollector) prepareResourcesForCollection(
	objects []runtime.Unstructured,
	namespaces []string,
) error {
	for _, o := range objects {
		metadata, err := meta.Accessor(o)
		if err != nil {
			return err
		}

		switch o.GetObjectKind().GroupVersionKind().Kind {
		case "PersistentVolume":
			err := r.preparePVResourceForCollection(o)
			if err != nil {
				return fmt.Errorf("error preparing PV resource %v: %v", metadata.GetName(), err)
			}
		case "Service":
			err := r.prepareServiceResourceForCollection(o)
			if err != nil {
				return fmt.Errorf("error preparing Service resource %v/%v: %v", metadata.GetNamespace(), metadata.GetName(), err)
			}
		case "ClusterRoleBinding":
			err := r.prepareClusterRoleBindingForCollection(o, namespaces)
			if err != nil {
				return fmt.Errorf("error preparing ClusterRoleBindings resource %v: %v", metadata.GetName(), err)
			}
		}

		content := o.UnstructuredContent()
		// Status shouldn't be retained when collecting resources
		delete(content, "status")
		metadataMap, err := collections.GetMap(content, "metadata")
		if err != nil {
			return fmt.Errorf("error getting metadata for resource %v: %v", metadata.GetName(), err)
		}
		// Remove all metadata except some well-known ones
		for key := range metadataMap {
			switch key {
			case "name", "namespace", "labels", "annotations":
			default:
				delete(metadataMap, key)
			}
		}
	}
	return nil
}

func (r *ResourceCollector) prepareResourceForApply(
	object runtime.Unstructured,
	namespaceMappings map[string]string,
	pvNameMappings map[string]string,
) error {
	objectType, err := meta.TypeAccessor(object)
	if err != nil {
		return err
	}

	metadata, err := meta.Accessor(object)
	if err != nil {
		return err
	}
	if metadata.GetNamespace() != "" {
		// Update the namepsace of the object, will be no-op for clustered resources
		metadata.SetNamespace(namespaceMappings[metadata.GetNamespace()])
	}

	switch objectType.GetKind() {
	case "PersistentVolume":
		return r.preparePVResourceForApply(object, pvNameMappings)
	case "PersistentVolumeClaim":
		return r.preparePVCResourceForApply(object, pvNameMappings)
	case "ClusterRoleBinding":
		err := r.prepareClusterRoleBindingForApply(object, namespaceMappings)
		if err != nil {
			return err
		}

	}
	return nil
}

func (r *ResourceCollector) mergeSupportedForResource(
	object runtime.Unstructured,
) bool {
	objectType, err := meta.TypeAccessor(object)
	if err != nil {
		return false
	}
	switch objectType.GetKind() {
	case "ClusterRoleBinding":
		return true
	}
	return false
}

func (r *ResourceCollector) mergeAndUpdateResource(
	object runtime.Unstructured,
) error {
	objectType, err := meta.TypeAccessor(object)
	if err != nil {
		return err
	}

	switch objectType.GetKind() {
	case "ClusterRoleBinding":
		return r.mergeAndUpdateClusterRoleBinding(object)
	}
	return nil
}

// ApplyResource applies a given resource using the provided client interface
func (r *ResourceCollector) ApplyResource(
	dynamicInterface dynamic.Interface,
	object runtime.Unstructured,
	pvNameMappings map[string]string,
	namespaceMappings map[string]string,
	deleteIfPresent bool,
) error {
	metadata, err := meta.Accessor(object)
	if err != nil {
		return err
	}
	objectType, err := meta.TypeAccessor(object)
	if err != nil {
		return err
	}
	resource := &metav1.APIResource{
		Name:       strings.ToLower(objectType.GetKind()) + "s",
		Namespaced: len(metadata.GetNamespace()) > 0,
	}

	destNamespace := ""
	if resource.Namespaced {
		destNamespace = namespaceMappings[metadata.GetNamespace()]
	}
	dynamicClient := dynamicInterface.Resource(
		object.GetObjectKind().GroupVersionKind().GroupVersion().WithResource(resource.Name)).Namespace(destNamespace)

	err = r.prepareResourceForApply(object, namespaceMappings, pvNameMappings)
	if err != nil {
		return err
	}

	_, err = dynamicClient.Create(object.(*unstructured.Unstructured))
	if err != nil {
		if apierrors.IsAlreadyExists(err) || strings.Contains(err.Error(), portallocator.ErrAllocated.Error()) {
			if r.mergeSupportedForResource(object) {
				return r.mergeAndUpdateResource(object)
			} else if strings.Contains(err.Error(), portallocator.ErrAllocated.Error()) {
				err = r.updateService(object)
				if err != nil {
					return err
				}
			} else if deleteIfPresent {
				// Delete the resource if it already exists on the destination
				// cluster and try creating again
				switch objectType.GetKind() {
				case "PersistentVolumeClaim", "PersistentVolume":
					err = nil
				default:
					err = dynamicClient.Delete(metadata.GetName(), &metav1.DeleteOptions{})
					if err != nil && apierrors.IsNotFound(err) {
						return err
					}
				}
			} else {
				return err
			}
			_, err = dynamicClient.Create(object.(*unstructured.Unstructured))
			return err
		}
	}

	return err
}
<commit_msg>Fix type for annotation to not collect resources<commit_after>package resourcecollector

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/heptio/ark/pkg/discovery"
	"github.com/heptio/ark/pkg/util/collections"
	"github.com/libopenstorage/stork/drivers/volume"
	"github.com/portworx/sched-ops/k8s"
	"github.com/sirupsen/logrus"
	rbacv1 "k8s.io/api/rbac/v1"
	apiextensionsclient "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/rest"
	"k8s.io/kubernetes/pkg/registry/core/service/portallocator"
)

const (
	// Annotation to use when the resource shouldn't be collected
	skipResourceAnnotation = "stork.libopenstorage.org/skipresource"
)

// ResourceCollector is used to collect and process unstructured objects in namespaces and using label selectors
type ResourceCollector struct {
	Driver           volume.Driver
	discoveryHelper  discovery.Helper
	dynamicInterface dynamic.Interface
}

// Init initializes the resource collector
func (r *ResourceCollector) Init() error {
	config, err := rest.InClusterConfig()
	if err != nil {
		return fmt.Errorf("error getting cluster config: %v", err)
	}

	aeclient, err := apiextensionsclient.NewForConfig(config)
	if err != nil {
		return fmt.Errorf("error getting apiextension client, %v", err)
	}

	discoveryClient := aeclient.Discovery()
	r.discoveryHelper, err = discovery.NewHelper(discoveryClient, logrus.New())
	if err != nil {
		return err
	}
	err = r.discoveryHelper.Refresh()
	if err != nil {
		return err
	}
	r.dynamicInterface, err = dynamic.NewForConfig(config)
	if err != nil {
		return err
	}
	return nil
}

func resourceToBeCollected(resource metav1.APIResource) bool {
	switch resource.Kind {
	case "PersistentVolumeClaim",
		"PersistentVolume",
		"Deployment",
		"DeploymentConfig",
		"StatefulSet",
		"ConfigMap",
		"Service",
		"Secret",
		"DaemonSet",
		"ServiceAccount",
		"Role",
		"RoleBinding",
		"ClusterRole",
		"ClusterRoleBinding",
		"ImageStream",
		"Ingress",
		"Route":
		return true
	default:
		return false
	}
}

// GetResources gets all the resources in the given list of namespaces which match the labelSelectors
func (r *ResourceCollector) GetResources(namespaces []string, labelSelectors map[string]string) ([]runtime.Unstructured, error) {
	err := r.discoveryHelper.Refresh()
	if err != nil {
		return nil, err
	}
	allObjects := make([]runtime.Unstructured, 0)

	// Map to prevent collection of duplicate objects
	resourceMap := make(map[types.UID]bool)

	crbs, err := k8s.Instance().ListClusterRoleBindings()
	if err != nil {
		return nil, err
	}

	for _, group := range r.discoveryHelper.Resources() {
		groupVersion, err := schema.ParseGroupVersion(group.GroupVersion)
		if err != nil {
			return nil, err
		}

		for _, resource := range group.APIResources {
			if !resourceToBeCollected(resource) {
				continue
			}

			for _, ns := range namespaces {
				var dynamicClient dynamic.ResourceInterface
				if !resource.Namespaced {
					dynamicClient = r.dynamicInterface.Resource(groupVersion.WithResource(resource.Name))
				} else {
					dynamicClient = r.dynamicInterface.Resource(groupVersion.WithResource(resource.Name)).Namespace(ns)
				}

				var selectors string
				// PVs don't get the labels from their PVCs, so don't use the label selector
				// Also skip for some other resources that aren't necessarily tied to an application
				switch resource.Kind {
				case "PersistentVolume",
					"ClusterRoleBinding",
					"ClusterRole",
					"ServiceAccount":
				default:
					selectors = labels.Set(labelSelectors).String()
				}
				objectsList, err := dynamicClient.List(metav1.ListOptions{
					LabelSelector: selectors,
				})
				if err != nil {
					return nil, err
				}
				objects, err := meta.ExtractList(objectsList)
				if err != nil {
					return nil, err
				}
				for _, o := range objects {
					runtimeObject, ok := o.(runtime.Unstructured)
					if !ok {
						return nil, fmt.Errorf("error casting object: %v", o)
					}

					collect, err := r.objectToBeCollected(labelSelectors, resourceMap, runtimeObject, crbs, ns)
					if err != nil {
						return nil, fmt.Errorf("error processing object %v: %v", runtimeObject, err)
					}
					if !collect {
						continue
					}
					metadata, err := meta.Accessor(runtimeObject)
					if err != nil {
						return nil, err
					}
					allObjects = append(allObjects, runtimeObject)
					resourceMap[metadata.GetUID()] = true
				}
			}
		}
	}

	err = r.prepareResourcesForCollection(allObjects, namespaces)
	if err != nil {
		return nil, err
	}
	return allObjects, nil
}

// Returns whether an object should be collected or not for the requested
// namespace
func (r *ResourceCollector) objectToBeCollected(
	labelSelectors map[string]string,
	resourceMap map[types.UID]bool,
	object runtime.Unstructured,
	crbs *rbacv1.ClusterRoleBindingList,
	namespace string,
) (bool, error) {
	metadata, err := meta.Accessor(object)
	if err != nil {
		return false, err
	}

	if value, present := metadata.GetAnnotations()[skipResourceAnnotation]; present {
		if skip, err := strconv.ParseBool(value); err == nil && skip {
			return false, err
		}
	}

	// Skip if we've already processed this object
	if _, ok := resourceMap[metadata.GetUID()]; ok {
		return false, nil
	}

	objectType, err := meta.TypeAccessor(object)
	if err != nil {
		return false, err
	}

	switch objectType.GetKind() {
	case "Service":
		return r.serviceToBeCollected(object)
	case "PersistentVolumeClaim":
		return r.pvcToBeCollected(object, namespace)
	case "PersistentVolume":
		return r.pvToBeCollected(labelSelectors, object, namespace)
	case "ClusterRoleBinding":
		return r.clusterRoleBindingToBeCollected(labelSelectors, object, namespace)
	case "ClusterRole":
		return r.clusterRoleToBeCollected(labelSelectors, object, crbs, namespace)
	case "ServiceAccount":
		return r.serviceAccountToBeCollected(object)
	case "Secret":
		return r.secretToBeCollected(object)
	}

	return true, nil
}

func (r *ResourceCollector) prepareResourcesForCollection(
	objects []runtime.Unstructured,
	namespaces []string,
) error {
	for _, o := range objects {
		metadata, err := meta.Accessor(o)
		if err != nil {
			return err
		}

		switch o.GetObjectKind().GroupVersionKind().Kind {
		case "PersistentVolume":
			err := r.preparePVResourceForCollection(o)
			if err != nil {
				return fmt.Errorf("error preparing PV resource %v: %v", metadata.GetName(), err)
			}
		case "Service":
			err := r.prepareServiceResourceForCollection(o)
			if err != nil {
				return fmt.Errorf("error preparing Service resource %v/%v: %v", metadata.GetNamespace(), metadata.GetName(), err)
			}
		case "ClusterRoleBinding":
			err := r.prepareClusterRoleBindingForCollection(o, namespaces)
			if err != nil {
				return fmt.Errorf("error preparing ClusterRoleBindings resource %v: %v", metadata.GetName(), err)
			}
		}

		content := o.UnstructuredContent()
		// Status shouldn't be retained when collecting resources
		delete(content, "status")
		metadataMap, err := collections.GetMap(content, "metadata")
		if err != nil {
			return fmt.Errorf("error getting metadata for resource %v: %v", metadata.GetName(), err)
		}
		// Remove all metadata except some well-known ones
		for key := range metadataMap {
			switch key {
			case "name", "namespace", "labels", "annotations":
			default:
				delete(metadataMap, key)
			}
		}
	}
	return nil
}

func (r *ResourceCollector) prepareResourceForApply(
	object runtime.Unstructured,
	namespaceMappings map[string]string,
	pvNameMappings map[string]string,
) error {
	objectType, err := meta.TypeAccessor(object)
	if err != nil {
		return err
	}

	metadata, err := meta.Accessor(object)
	if err != nil {
		return err
	}
	if metadata.GetNamespace() != "" {
		// Update the namepsace of the object, will be no-op for clustered resources
		metadata.SetNamespace(namespaceMappings[metadata.GetNamespace()])
	}

	switch objectType.GetKind() {
	case "PersistentVolume":
		return r.preparePVResourceForApply(object, pvNameMappings)
	case "PersistentVolumeClaim":
		return r.preparePVCResourceForApply(object, pvNameMappings)
	case "ClusterRoleBinding":
		err := r.prepareClusterRoleBindingForApply(object, namespaceMappings)
		if err != nil {
			return err
		}

	}
	return nil
}

func (r *ResourceCollector) mergeSupportedForResource(
	object runtime.Unstructured,
) bool {
	objectType, err := meta.TypeAccessor(object)
	if err != nil {
		return false
	}
	switch objectType.GetKind() {
	case "ClusterRoleBinding":
		return true
	}
	return false
}

func (r *ResourceCollector) mergeAndUpdateResource(
	object runtime.Unstructured,
) error {
	objectType, err := meta.TypeAccessor(object)
	if err != nil {
		return err
	}

	switch objectType.GetKind() {
	case "ClusterRoleBinding":
		return r.mergeAndUpdateClusterRoleBinding(object)
	}
	return nil
}

// ApplyResource applies a given resource using the provided client interface
func (r *ResourceCollector) ApplyResource(
	dynamicInterface dynamic.Interface,
	object runtime.Unstructured,
	pvNameMappings map[string]string,
	namespaceMappings map[string]string,
	deleteIfPresent bool,
) error {
	metadata, err := meta.Accessor(object)
	if err != nil {
		return err
	}
	objectType, err := meta.TypeAccessor(object)
	if err != nil {
		return err
	}
	resource := &metav1.APIResource{
		Name:       strings.ToLower(objectType.GetKind()) + "s",
		Namespaced: len(metadata.GetNamespace()) > 0,
	}

	destNamespace := ""
	if resource.Namespaced {
		destNamespace = namespaceMappings[metadata.GetNamespace()]
	}
	dynamicClient := dynamicInterface.Resource(
		object.GetObjectKind().GroupVersionKind().GroupVersion().WithResource(resource.Name)).Namespace(destNamespace)

	err = r.prepareResourceForApply(object, namespaceMappings, pvNameMappings)
	if err != nil {
		return err
	}

	_, err = dynamicClient.Create(object.(*unstructured.Unstructured))
	if err != nil {
		if apierrors.IsAlreadyExists(err) || strings.Contains(err.Error(), portallocator.ErrAllocated.Error()) {
			if r.mergeSupportedForResource(object) {
				return r.mergeAndUpdateResource(object)
			} else if strings.Contains(err.Error(), portallocator.ErrAllocated.Error()) {
				err = r.updateService(object)
				if err != nil {
					return err
				}
			} else if deleteIfPresent {
				// Delete the resource if it already exists on the destination
				// cluster and try creating again
				switch objectType.GetKind() {
				case "PersistentVolumeClaim", "PersistentVolume":
					err = nil
				default:
					err = dynamicClient.Delete(metadata.GetName(), &metav1.DeleteOptions{})
					if err != nil && apierrors.IsNotFound(err) {
						return err
					}
				}
			} else {
				return err
			}
			_, err = dynamicClient.Create(object.(*unstructured.Unstructured))
			return err
		}
	}

	return err
}
<|endoftext|>