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

package main

import (
	"flag"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/coreos/coreos-cloudinit/config"
	"github.com/coreos/coreos-cloudinit/config/validate"
	"github.com/coreos/coreos-cloudinit/datasource"
	"github.com/coreos/coreos-cloudinit/datasource/configdrive"
	"github.com/coreos/coreos-cloudinit/datasource/file"
	"github.com/coreos/coreos-cloudinit/datasource/metadata/cloudsigma"
	"github.com/coreos/coreos-cloudinit/datasource/metadata/digitalocean"
	"github.com/coreos/coreos-cloudinit/datasource/metadata/ec2"
	"github.com/coreos/coreos-cloudinit/datasource/metadata/packet"
	"github.com/coreos/coreos-cloudinit/datasource/proc_cmdline"
	"github.com/coreos/coreos-cloudinit/datasource/url"
	"github.com/coreos/coreos-cloudinit/datasource/waagent"
	"github.com/coreos/coreos-cloudinit/initialize"
	"github.com/coreos/coreos-cloudinit/network"
	"github.com/coreos/coreos-cloudinit/pkg"
	"github.com/coreos/coreos-cloudinit/system"
)

const (
	version               = "1.4.1+git"
	datasourceInterval    = 100 * time.Millisecond
	datasourceMaxInterval = 30 * time.Second
	datasourceTimeout     = 5 * time.Minute
)

var (
	flags = struct {
		printVersion  bool
		ignoreFailure bool
		sources       struct {
			file                        string
			configDrive                 string
			waagent                     string
			metadataService             bool
			ec2MetadataService          string
			cloudSigmaMetadataService   bool
			digitalOceanMetadataService string
			packetMetadataService       string
			url                         string
			procCmdLine                 bool
		}
		convertNetconf string
		workspace      string
		sshKeyName     string
		oem            string
		validate       bool
	}{}
)

func init() {
	flag.BoolVar(&flags.printVersion, "version", false, "Print the version and exit")
	flag.BoolVar(&flags.ignoreFailure, "ignore-failure", false, "Exits with 0 status in the event of malformed input from user-data")
	flag.StringVar(&flags.sources.file, "from-file", "", "Read user-data from provided file")
	flag.StringVar(&flags.sources.configDrive, "from-configdrive", "", "Read data from provided cloud-drive directory")
	flag.StringVar(&flags.sources.waagent, "from-waagent", "", "Read data from provided waagent directory")
	flag.BoolVar(&flags.sources.metadataService, "from-metadata-service", false, "[DEPRECATED - Use -from-ec2-metadata] Download data from metadata service")
	flag.StringVar(&flags.sources.ec2MetadataService, "from-ec2-metadata", "", "Download EC2 data from the provided url")
	flag.BoolVar(&flags.sources.cloudSigmaMetadataService, "from-cloudsigma-metadata", false, "Download data from CloudSigma server context")
	flag.StringVar(&flags.sources.digitalOceanMetadataService, "from-digitalocean-metadata", "", "Download DigitalOcean data from the provided url")
	flag.StringVar(&flags.sources.packetMetadataService, "from-packet-metadata", "", "Download Packet data from metadata service")
	flag.StringVar(&flags.sources.url, "from-url", "", "Download user-data from provided url")
	flag.BoolVar(&flags.sources.procCmdLine, "from-proc-cmdline", false, fmt.Sprintf("Parse %s for '%s=<url>', using the cloud-config served by an HTTP GET to <url>", proc_cmdline.ProcCmdlineLocation, proc_cmdline.ProcCmdlineCloudConfigFlag))
	flag.StringVar(&flags.oem, "oem", "", "Use the settings specific to the provided OEM")
	flag.StringVar(&flags.convertNetconf, "convert-netconf", "", "Read the network config provided in cloud-drive and translate it from the specified format into networkd unit files")
	flag.StringVar(&flags.workspace, "workspace", "/var/lib/coreos-cloudinit", "Base directory coreos-cloudinit should use to store data")
	flag.StringVar(&flags.sshKeyName, "ssh-key-name", initialize.DefaultSSHKeyName, "Add SSH keys to the system with the given name")
	flag.BoolVar(&flags.validate, "validate", false, "[EXPERIMENTAL] Validate the user-data but do not apply it to the system")
}

type oemConfig map[string]string

var (
	oemConfigs = map[string]oemConfig{
		"digitalocean": oemConfig{
			"from-digitalocean-metadata": "http://169.254.169.254/",
			"convert-netconf":            "digitalocean",
		},
		"ec2-compat": oemConfig{
			"from-ec2-metadata": "http://169.254.169.254/",
			"from-configdrive":  "/media/configdrive",
		},
		"rackspace-onmetal": oemConfig{
			"from-configdrive": "/media/configdrive",
			"convert-netconf":  "debian",
		},
		"azure": oemConfig{
			"from-waagent": "/var/lib/waagent",
		},
		"cloudsigma": oemConfig{
			"from-cloudsigma-metadata": "true",
		},
		"packet": oemConfig{
			"from-packet-metadata": "https://metadata.packet.net/",
			"convert-netconf":      "packet",
		},
	}
)

func main() {
	failure := false

	flag.Parse()

	if c, ok := oemConfigs[flags.oem]; ok {
		for k, v := range c {
			flag.Set(k, v)
		}
	} else if flags.oem != "" {
		oems := make([]string, 0, len(oemConfigs))
		for k := range oemConfigs {
			oems = append(oems, k)
		}
		fmt.Printf("Invalid option to --oem: %q. Supported options: %q\n", flags.oem, oems)
		os.Exit(2)
	}

	if flags.printVersion == true {
		fmt.Printf("coreos-cloudinit version %s\n", version)
		os.Exit(0)
	}

	switch flags.convertNetconf {
	case "":
	case "debian":
	case "digitalocean":
	case "packet":
	default:
		fmt.Printf("Invalid option to -convert-netconf: '%s'. Supported options: 'debian, digitalocean, packet'\n", flags.convertNetconf)
		os.Exit(2)
	}

	dss := getDatasources()
	if len(dss) == 0 {
		fmt.Println("Provide at least one of --from-file, --from-configdrive, --from-ec2-metadata, --from-cloudsigma-metadata, --from-packet-metadata, --from-url or --from-proc-cmdline")
		os.Exit(2)
	}

	ds := selectDatasource(dss)
	if ds == nil {
		fmt.Println("No datasources available in time")
		os.Exit(1)
	}

	fmt.Printf("Fetching user-data from datasource of type %q\n", ds.Type())
	userdataBytes, err := ds.FetchUserdata()
	if err != nil {
		fmt.Printf("Failed fetching user-data from datasource: %v\nContinuing...\n", err)
		failure = true
	}

	if report, err := validate.Validate(userdataBytes); err == nil {
		ret := 0
		for _, e := range report.Entries() {
			fmt.Println(e)
			ret = 1
		}
		if flags.validate {
			os.Exit(ret)
		}
	} else {
		fmt.Printf("Failed while validating user_data (%q)\n", err)
		if flags.validate {
			os.Exit(1)
		}
	}

	fmt.Printf("Fetching meta-data from datasource of type %q\n", ds.Type())
	metadata, err := ds.FetchMetadata()
	if err != nil {
		fmt.Printf("Failed fetching meta-data from datasource: %v\n", err)
		os.Exit(1)
	}

	// Apply environment to user-data
	env := initialize.NewEnvironment("/", ds.ConfigRoot(), flags.workspace, flags.sshKeyName, metadata)
	userdata := env.Apply(string(userdataBytes))

	var ccu *config.CloudConfig
	var script *config.Script
	switch ud, err := initialize.ParseUserData(userdata); err {
	case initialize.ErrIgnitionConfig:
		fmt.Printf("Detected an Ignition config. Exiting...")
		os.Exit(0)
	case nil:
		switch t := ud.(type) {
		case *config.CloudConfig:
			ccu = t
		case *config.Script:
			script = t
		}
	default:
		fmt.Printf("Failed to parse user-data: %v\nContinuing...\n", err)
		failure = true
	}

	fmt.Println("Merging cloud-config from meta-data and user-data")
	cc := mergeConfigs(ccu, metadata)

	var ifaces []network.InterfaceGenerator
	if flags.convertNetconf != "" {
		var err error
		switch flags.convertNetconf {
		case "debian":
			ifaces, err = network.ProcessDebianNetconf(metadata.NetworkConfig)
		case "digitalocean":
			ifaces, err = network.ProcessDigitalOceanNetconf(metadata.NetworkConfig)
		case "packet":
			ifaces, err = network.ProcessPacketNetconf(metadata.NetworkConfig)
		default:
			err = fmt.Errorf("Unsupported network config format %q", flags.convertNetconf)
		}
		if err != nil {
			fmt.Printf("Failed to generate interfaces: %v\n", err)
			os.Exit(1)
		}
	}

	if err = initialize.Apply(cc, ifaces, env); err != nil {
		fmt.Printf("Failed to apply cloud-config: %v\n", err)
		os.Exit(1)
	}

	if script != nil {
		if err = runScript(*script, env); err != nil {
			fmt.Printf("Failed to run script: %v\n", err)
			os.Exit(1)
		}
	}

	if failure && !flags.ignoreFailure {
		os.Exit(1)
	}
}

// mergeConfigs merges certain options from md (meta-data from the datasource)
// onto cc (a CloudConfig derived from user-data), if they are not already set
// on cc (i.e. user-data always takes precedence)
func mergeConfigs(cc *config.CloudConfig, md datasource.Metadata) (out config.CloudConfig) {
	if cc != nil {
		out = *cc
	}

	if md.Hostname != "" {
		if out.Hostname != "" {
			fmt.Printf("Warning: user-data hostname (%s) overrides metadata hostname (%s)\n", out.Hostname, md.Hostname)
		} else {
			out.Hostname = md.Hostname
		}
	}
	for _, key := range md.SSHPublicKeys {
		out.SSHAuthorizedKeys = append(out.SSHAuthorizedKeys, key)
	}
	return
}

// getDatasources creates a slice of possible Datasources for cloudinit based
// on the different source command-line flags.
func getDatasources() []datasource.Datasource {
	dss := make([]datasource.Datasource, 0, 5)
	if flags.sources.file != "" {
		dss = append(dss, file.NewDatasource(flags.sources.file))
	}
	if flags.sources.url != "" {
		dss = append(dss, url.NewDatasource(flags.sources.url))
	}
	if flags.sources.configDrive != "" {
		dss = append(dss, configdrive.NewDatasource(flags.sources.configDrive))
	}
	if flags.sources.metadataService {
		dss = append(dss, ec2.NewDatasource(ec2.DefaultAddress))
	}
	if flags.sources.ec2MetadataService != "" {
		dss = append(dss, ec2.NewDatasource(flags.sources.ec2MetadataService))
	}
	if flags.sources.cloudSigmaMetadataService {
		dss = append(dss, cloudsigma.NewServerContextService())
	}
	if flags.sources.digitalOceanMetadataService != "" {
		dss = append(dss, digitalocean.NewDatasource(flags.sources.digitalOceanMetadataService))
	}
	if flags.sources.waagent != "" {
		dss = append(dss, waagent.NewDatasource(flags.sources.waagent))
	}
	if flags.sources.packetMetadataService != "" {
		dss = append(dss, packet.NewDatasource(flags.sources.packetMetadataService))
	}
	if flags.sources.procCmdLine {
		dss = append(dss, proc_cmdline.NewDatasource())
	}
	return dss
}

// selectDatasource attempts to choose a valid Datasource to use based on its
// current availability. The first Datasource to report to be available is
// returned. Datasources will be retried if possible if they are not
// immediately available. If all Datasources are permanently unavailable or
// datasourceTimeout is reached before one becomes available, nil is returned.
func selectDatasource(sources []datasource.Datasource) datasource.Datasource {
	ds := make(chan datasource.Datasource)
	stop := make(chan struct{})
	var wg sync.WaitGroup

	for _, s := range sources {
		wg.Add(1)
		go func(s datasource.Datasource) {
			defer wg.Done()

			duration := datasourceInterval
			for {
				fmt.Printf("Checking availability of %q\n", s.Type())
				if s.IsAvailable() {
					ds <- s
					return
				} else if !s.AvailabilityChanges() {
					return
				}
				select {
				case <-stop:
					return
				case <-time.After(duration):
					duration = pkg.ExpBackoff(duration, datasourceMaxInterval)
				}
			}
		}(s)
	}

	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	var s datasource.Datasource
	select {
	case s = <-ds:
	case <-done:
	case <-time.After(datasourceTimeout):
	}

	close(stop)
	return s
}

// TODO(jonboulle): this should probably be refactored and moved into a different module
func runScript(script config.Script, env *initialize.Environment) error {
	err := initialize.PrepWorkspace(env.Workspace())
	if err != nil {
		fmt.Printf("Failed preparing workspace: %v\n", err)
		return err
	}
	path, err := initialize.PersistScriptInWorkspace(script, env.Workspace())
	if err == nil {
		var name string
		name, err = system.ExecuteScript(path)
		initialize.PersistUnitNameInWorkspace(name, env.Workspace())
	}
	return err
}
<commit_msg>coreos-cloudinit: bump to v1.5.0<commit_after>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/coreos/coreos-cloudinit/config"
	"github.com/coreos/coreos-cloudinit/config/validate"
	"github.com/coreos/coreos-cloudinit/datasource"
	"github.com/coreos/coreos-cloudinit/datasource/configdrive"
	"github.com/coreos/coreos-cloudinit/datasource/file"
	"github.com/coreos/coreos-cloudinit/datasource/metadata/cloudsigma"
	"github.com/coreos/coreos-cloudinit/datasource/metadata/digitalocean"
	"github.com/coreos/coreos-cloudinit/datasource/metadata/ec2"
	"github.com/coreos/coreos-cloudinit/datasource/metadata/packet"
	"github.com/coreos/coreos-cloudinit/datasource/proc_cmdline"
	"github.com/coreos/coreos-cloudinit/datasource/url"
	"github.com/coreos/coreos-cloudinit/datasource/waagent"
	"github.com/coreos/coreos-cloudinit/initialize"
	"github.com/coreos/coreos-cloudinit/network"
	"github.com/coreos/coreos-cloudinit/pkg"
	"github.com/coreos/coreos-cloudinit/system"
)

const (
	version               = "1.5.0"
	datasourceInterval    = 100 * time.Millisecond
	datasourceMaxInterval = 30 * time.Second
	datasourceTimeout     = 5 * time.Minute
)

var (
	flags = struct {
		printVersion  bool
		ignoreFailure bool
		sources       struct {
			file                        string
			configDrive                 string
			waagent                     string
			metadataService             bool
			ec2MetadataService          string
			cloudSigmaMetadataService   bool
			digitalOceanMetadataService string
			packetMetadataService       string
			url                         string
			procCmdLine                 bool
		}
		convertNetconf string
		workspace      string
		sshKeyName     string
		oem            string
		validate       bool
	}{}
)

func init() {
	flag.BoolVar(&flags.printVersion, "version", false, "Print the version and exit")
	flag.BoolVar(&flags.ignoreFailure, "ignore-failure", false, "Exits with 0 status in the event of malformed input from user-data")
	flag.StringVar(&flags.sources.file, "from-file", "", "Read user-data from provided file")
	flag.StringVar(&flags.sources.configDrive, "from-configdrive", "", "Read data from provided cloud-drive directory")
	flag.StringVar(&flags.sources.waagent, "from-waagent", "", "Read data from provided waagent directory")
	flag.BoolVar(&flags.sources.metadataService, "from-metadata-service", false, "[DEPRECATED - Use -from-ec2-metadata] Download data from metadata service")
	flag.StringVar(&flags.sources.ec2MetadataService, "from-ec2-metadata", "", "Download EC2 data from the provided url")
	flag.BoolVar(&flags.sources.cloudSigmaMetadataService, "from-cloudsigma-metadata", false, "Download data from CloudSigma server context")
	flag.StringVar(&flags.sources.digitalOceanMetadataService, "from-digitalocean-metadata", "", "Download DigitalOcean data from the provided url")
	flag.StringVar(&flags.sources.packetMetadataService, "from-packet-metadata", "", "Download Packet data from metadata service")
	flag.StringVar(&flags.sources.url, "from-url", "", "Download user-data from provided url")
	flag.BoolVar(&flags.sources.procCmdLine, "from-proc-cmdline", false, fmt.Sprintf("Parse %s for '%s=<url>', using the cloud-config served by an HTTP GET to <url>", proc_cmdline.ProcCmdlineLocation, proc_cmdline.ProcCmdlineCloudConfigFlag))
	flag.StringVar(&flags.oem, "oem", "", "Use the settings specific to the provided OEM")
	flag.StringVar(&flags.convertNetconf, "convert-netconf", "", "Read the network config provided in cloud-drive and translate it from the specified format into networkd unit files")
	flag.StringVar(&flags.workspace, "workspace", "/var/lib/coreos-cloudinit", "Base directory coreos-cloudinit should use to store data")
	flag.StringVar(&flags.sshKeyName, "ssh-key-name", initialize.DefaultSSHKeyName, "Add SSH keys to the system with the given name")
	flag.BoolVar(&flags.validate, "validate", false, "[EXPERIMENTAL] Validate the user-data but do not apply it to the system")
}

type oemConfig map[string]string

var (
	oemConfigs = map[string]oemConfig{
		"digitalocean": oemConfig{
			"from-digitalocean-metadata": "http://169.254.169.254/",
			"convert-netconf":            "digitalocean",
		},
		"ec2-compat": oemConfig{
			"from-ec2-metadata": "http://169.254.169.254/",
			"from-configdrive":  "/media/configdrive",
		},
		"rackspace-onmetal": oemConfig{
			"from-configdrive": "/media/configdrive",
			"convert-netconf":  "debian",
		},
		"azure": oemConfig{
			"from-waagent": "/var/lib/waagent",
		},
		"cloudsigma": oemConfig{
			"from-cloudsigma-metadata": "true",
		},
		"packet": oemConfig{
			"from-packet-metadata": "https://metadata.packet.net/",
			"convert-netconf":      "packet",
		},
	}
)

func main() {
	failure := false

	flag.Parse()

	if c, ok := oemConfigs[flags.oem]; ok {
		for k, v := range c {
			flag.Set(k, v)
		}
	} else if flags.oem != "" {
		oems := make([]string, 0, len(oemConfigs))
		for k := range oemConfigs {
			oems = append(oems, k)
		}
		fmt.Printf("Invalid option to --oem: %q. Supported options: %q\n", flags.oem, oems)
		os.Exit(2)
	}

	if flags.printVersion == true {
		fmt.Printf("coreos-cloudinit version %s\n", version)
		os.Exit(0)
	}

	switch flags.convertNetconf {
	case "":
	case "debian":
	case "digitalocean":
	case "packet":
	default:
		fmt.Printf("Invalid option to -convert-netconf: '%s'. Supported options: 'debian, digitalocean, packet'\n", flags.convertNetconf)
		os.Exit(2)
	}

	dss := getDatasources()
	if len(dss) == 0 {
		fmt.Println("Provide at least one of --from-file, --from-configdrive, --from-ec2-metadata, --from-cloudsigma-metadata, --from-packet-metadata, --from-url or --from-proc-cmdline")
		os.Exit(2)
	}

	ds := selectDatasource(dss)
	if ds == nil {
		fmt.Println("No datasources available in time")
		os.Exit(1)
	}

	fmt.Printf("Fetching user-data from datasource of type %q\n", ds.Type())
	userdataBytes, err := ds.FetchUserdata()
	if err != nil {
		fmt.Printf("Failed fetching user-data from datasource: %v\nContinuing...\n", err)
		failure = true
	}

	if report, err := validate.Validate(userdataBytes); err == nil {
		ret := 0
		for _, e := range report.Entries() {
			fmt.Println(e)
			ret = 1
		}
		if flags.validate {
			os.Exit(ret)
		}
	} else {
		fmt.Printf("Failed while validating user_data (%q)\n", err)
		if flags.validate {
			os.Exit(1)
		}
	}

	fmt.Printf("Fetching meta-data from datasource of type %q\n", ds.Type())
	metadata, err := ds.FetchMetadata()
	if err != nil {
		fmt.Printf("Failed fetching meta-data from datasource: %v\n", err)
		os.Exit(1)
	}

	// Apply environment to user-data
	env := initialize.NewEnvironment("/", ds.ConfigRoot(), flags.workspace, flags.sshKeyName, metadata)
	userdata := env.Apply(string(userdataBytes))

	var ccu *config.CloudConfig
	var script *config.Script
	switch ud, err := initialize.ParseUserData(userdata); err {
	case initialize.ErrIgnitionConfig:
		fmt.Printf("Detected an Ignition config. Exiting...")
		os.Exit(0)
	case nil:
		switch t := ud.(type) {
		case *config.CloudConfig:
			ccu = t
		case *config.Script:
			script = t
		}
	default:
		fmt.Printf("Failed to parse user-data: %v\nContinuing...\n", err)
		failure = true
	}

	fmt.Println("Merging cloud-config from meta-data and user-data")
	cc := mergeConfigs(ccu, metadata)

	var ifaces []network.InterfaceGenerator
	if flags.convertNetconf != "" {
		var err error
		switch flags.convertNetconf {
		case "debian":
			ifaces, err = network.ProcessDebianNetconf(metadata.NetworkConfig)
		case "digitalocean":
			ifaces, err = network.ProcessDigitalOceanNetconf(metadata.NetworkConfig)
		case "packet":
			ifaces, err = network.ProcessPacketNetconf(metadata.NetworkConfig)
		default:
			err = fmt.Errorf("Unsupported network config format %q", flags.convertNetconf)
		}
		if err != nil {
			fmt.Printf("Failed to generate interfaces: %v\n", err)
			os.Exit(1)
		}
	}

	if err = initialize.Apply(cc, ifaces, env); err != nil {
		fmt.Printf("Failed to apply cloud-config: %v\n", err)
		os.Exit(1)
	}

	if script != nil {
		if err = runScript(*script, env); err != nil {
			fmt.Printf("Failed to run script: %v\n", err)
			os.Exit(1)
		}
	}

	if failure && !flags.ignoreFailure {
		os.Exit(1)
	}
}

// mergeConfigs merges certain options from md (meta-data from the datasource)
// onto cc (a CloudConfig derived from user-data), if they are not already set
// on cc (i.e. user-data always takes precedence)
func mergeConfigs(cc *config.CloudConfig, md datasource.Metadata) (out config.CloudConfig) {
	if cc != nil {
		out = *cc
	}

	if md.Hostname != "" {
		if out.Hostname != "" {
			fmt.Printf("Warning: user-data hostname (%s) overrides metadata hostname (%s)\n", out.Hostname, md.Hostname)
		} else {
			out.Hostname = md.Hostname
		}
	}
	for _, key := range md.SSHPublicKeys {
		out.SSHAuthorizedKeys = append(out.SSHAuthorizedKeys, key)
	}
	return
}

// getDatasources creates a slice of possible Datasources for cloudinit based
// on the different source command-line flags.
func getDatasources() []datasource.Datasource {
	dss := make([]datasource.Datasource, 0, 5)
	if flags.sources.file != "" {
		dss = append(dss, file.NewDatasource(flags.sources.file))
	}
	if flags.sources.url != "" {
		dss = append(dss, url.NewDatasource(flags.sources.url))
	}
	if flags.sources.configDrive != "" {
		dss = append(dss, configdrive.NewDatasource(flags.sources.configDrive))
	}
	if flags.sources.metadataService {
		dss = append(dss, ec2.NewDatasource(ec2.DefaultAddress))
	}
	if flags.sources.ec2MetadataService != "" {
		dss = append(dss, ec2.NewDatasource(flags.sources.ec2MetadataService))
	}
	if flags.sources.cloudSigmaMetadataService {
		dss = append(dss, cloudsigma.NewServerContextService())
	}
	if flags.sources.digitalOceanMetadataService != "" {
		dss = append(dss, digitalocean.NewDatasource(flags.sources.digitalOceanMetadataService))
	}
	if flags.sources.waagent != "" {
		dss = append(dss, waagent.NewDatasource(flags.sources.waagent))
	}
	if flags.sources.packetMetadataService != "" {
		dss = append(dss, packet.NewDatasource(flags.sources.packetMetadataService))
	}
	if flags.sources.procCmdLine {
		dss = append(dss, proc_cmdline.NewDatasource())
	}
	return dss
}

// selectDatasource attempts to choose a valid Datasource to use based on its
// current availability. The first Datasource to report to be available is
// returned. Datasources will be retried if possible if they are not
// immediately available. If all Datasources are permanently unavailable or
// datasourceTimeout is reached before one becomes available, nil is returned.
func selectDatasource(sources []datasource.Datasource) datasource.Datasource {
	ds := make(chan datasource.Datasource)
	stop := make(chan struct{})
	var wg sync.WaitGroup

	for _, s := range sources {
		wg.Add(1)
		go func(s datasource.Datasource) {
			defer wg.Done()

			duration := datasourceInterval
			for {
				fmt.Printf("Checking availability of %q\n", s.Type())
				if s.IsAvailable() {
					ds <- s
					return
				} else if !s.AvailabilityChanges() {
					return
				}
				select {
				case <-stop:
					return
				case <-time.After(duration):
					duration = pkg.ExpBackoff(duration, datasourceMaxInterval)
				}
			}
		}(s)
	}

	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	var s datasource.Datasource
	select {
	case s = <-ds:
	case <-done:
	case <-time.After(datasourceTimeout):
	}

	close(stop)
	return s
}

// TODO(jonboulle): this should probably be refactored and moved into a different module
func runScript(script config.Script, env *initialize.Environment) error {
	err := initialize.PrepWorkspace(env.Workspace())
	if err != nil {
		fmt.Printf("Failed preparing workspace: %v\n", err)
		return err
	}
	path, err := initialize.PersistScriptInWorkspace(script, env.Workspace())
	if err == nil {
		var name string
		name, err = system.ExecuteScript(path)
		initialize.PersistUnitNameInWorkspace(name, env.Workspace())
	}
	return err
}
<|endoftext|>


<commit_before>package cli

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/op/go-logging"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"
)

var log = logging.MustGetLogger("jira.cli")

type Cli struct {
	endpoint   *url.URL
	opts       map[string]string
	cookieFile string
	ua         *http.Client
}

func New(opts map[string]string) *Cli {
	homedir := os.Getenv("HOME")
	cookieJar, _ := cookiejar.New(nil)
	endpoint, _ := opts["endpoint"]
	url, _ := url.Parse(strings.TrimRight(endpoint, "/"))

	if project, ok := opts["project"]; ok {
		opts["project"] = strings.ToUpper(project)
	}

	cli := &Cli{
		endpoint:   url,
		opts:       opts,
		cookieFile: fmt.Sprintf("%s/.jira.d/cookies.js", homedir),
		ua:         &http.Client{Jar: cookieJar},
	}

	cli.ua.Jar.SetCookies(url, cli.loadCookies())

	return cli
}

func (c *Cli) saveCookies(cookies []*http.Cookie) {
	// expiry in one week from now
	expiry := time.Now().Add(24 * 7 * time.Hour)
	for _, cookie := range cookies {
		cookie.Expires = expiry
	}

	if currentCookies := c.loadCookies(); currentCookies != nil {
		currentCookiesByName := make(map[string]*http.Cookie)
		for _, cookie := range currentCookies {
			currentCookiesByName[cookie.Name] = cookie
		}

		for _, cookie := range cookies {
			currentCookiesByName[cookie.Name] = cookie
		}

		mergedCookies := make([]*http.Cookie, 0, len(currentCookiesByName))
		for _, v := range currentCookiesByName {
			mergedCookies = append(mergedCookies, v)
		}
		jsonWrite(c.cookieFile, mergedCookies)
	} else {
		jsonWrite(c.cookieFile, cookies)
	}
}

func (c *Cli) loadCookies() []*http.Cookie {
	bytes, err := ioutil.ReadFile(c.cookieFile)
	if err != nil && os.IsNotExist(err) {
		// dont load cookies if the file does not exist
		return nil
	}
	if err != nil {
		log.Error("Failed to open %s: %s", c.cookieFile, err)
		os.Exit(1)
	}
	cookies := make([]*http.Cookie, 0)
	err = json.Unmarshal(bytes, &cookies)
	if err != nil {
		log.Error("Failed to parse json from file %s: %s", c.cookieFile, err)
	}
	log.Debug("Loading Cookies: %s", cookies)
	return cookies
}

func (c *Cli) post(uri string, content string) (*http.Response, error) {
	return c.makeRequestWithContent("POST", uri, content)
}

func (c *Cli) put(uri string, content string) (*http.Response, error) {
	return c.makeRequestWithContent("PUT", uri, content)
}

func (c *Cli) makeRequestWithContent(method string, uri string, content string) (*http.Response, error) {
	buffer := bytes.NewBufferString(content)
	req, _ := http.NewRequest(method, uri, buffer)

	log.Info("%s %s", req.Method, req.URL.String())
	if log.IsEnabledFor(logging.DEBUG) {
		logBuffer := bytes.NewBuffer(make([]byte, 0, len(content)))
		req.Write(logBuffer)
		log.Debug("%s", logBuffer)
		// need to recreate the buffer since the offset is now at the end
		// need to be able to rewind the buffer offset, dont know how yet
		req, _ = http.NewRequest(method, uri, bytes.NewBufferString(content))
	}

	if resp, err := c.makeRequest(req); err != nil {
		return nil, err
	} else {
		if resp.StatusCode == 401 {
			if err := c.CmdLogin(); err != nil {
				return nil, err
			}
			req, _ = http.NewRequest(method, uri, bytes.NewBufferString(content))
			return c.makeRequest(req)
		}
		return resp, err
	}
}

func (c *Cli) get(uri string) (*http.Response, error) {
	req, _ := http.NewRequest("GET", uri, nil)
	log.Info("%s %s", req.Method, req.URL.String())
	if log.IsEnabledFor(logging.DEBUG) {
		logBuffer := bytes.NewBuffer(make([]byte, 0))
		req.Write(logBuffer)
		log.Debug("%s", logBuffer)
	}

	if resp, err := c.makeRequest(req); err != nil {
		return nil, err
	} else {
		if resp.StatusCode == 401 {
			if err := c.CmdLogin(); err != nil {
				return nil, err
			}
			return c.makeRequest(req)
		}
		return resp, err
	}
}

func (c *Cli) makeRequest(req *http.Request) (resp *http.Response, err error) {
	req.Header.Set("Content-Type", "application/json")
	if resp, err = c.ua.Do(req); err != nil {
		log.Error("Failed to %s %s: %s", req.Method, req.URL.String(), err)
		return nil, err
	} else {
		if resp.StatusCode < 200 || resp.StatusCode >= 300 && resp.StatusCode != 401 {
			log.Error("response status: %s", resp.Status)
		}

		runtime.SetFinalizer(resp, func(r *http.Response) {
			r.Body.Close()
		})

		if _, ok := resp.Header["Set-Cookie"]; ok {
			c.saveCookies(resp.Cookies())
		}
	}
	return resp, nil
}

func (c *Cli) getTemplate(name string) string {
	if override, ok := c.opts["template"]; ok {
		if _, err := os.Stat(override); err == nil {
			return readFile(override)
		} else {
			if file, err := FindClosestParentPath(fmt.Sprintf(".jira.d/templates/%s", override)); err == nil {
				return readFile(file)
			}
			if dflt, ok := all_templates[override]; ok {
				return dflt
			}
		}
	}
	if file, err := FindClosestParentPath(fmt.Sprintf(".jira.d/templates/%s", name)); err != nil {
		// create-bug etc are special, if we dont find it in the path
		// then just return a generic create template
		if strings.HasPrefix(name, "create-") {
			if file, err := FindClosestParentPath(".jira.d/templates/create"); err != nil {
				return all_templates["create"]
			} else {
				return readFile(file)
			}
		}
		return all_templates[name]
	} else {
		return readFile(file)
	}
}

func (c *Cli) editTemplate(template string, tmpFilePrefix string, templateData map[string]interface{}, templateProcessor func(string) error) error {

	tmpdir := fmt.Sprintf("%s/.jira.d/tmp", os.Getenv("HOME"))
	if err := mkdir(tmpdir); err != nil {
		return err
	}

	fh, err := ioutil.TempFile(tmpdir, tmpFilePrefix)
	if err != nil {
		log.Error("Failed to make temp file in %s: %s", tmpdir, err)
		return err
	}
	defer fh.Close()

	tmpFileName := fmt.Sprintf("%s.yml", fh.Name())
	if err := os.Rename(fh.Name(), tmpFileName); err != nil {
		log.Error("Failed to rename %s to %s: %s", fh.Name(), fmt.Sprintf("%s.yml", fh.Name()), err)
		return err
	}

	err = runTemplate(template, templateData, fh)
	if err != nil {
		return err
	}

	fh.Close()

	editor, ok := c.opts["editor"]
	if !ok {
		editor = os.Getenv("JIRA_EDITOR")
		if editor == "" {
			editor = os.Getenv("EDITOR")
			if editor == "" {
				editor = "vim"
			}
		}
	}

	editing := true
	if val, ok := c.opts["edit"]; ok && val == "false" {
		editing = false
	}

	for true {
		if editing {
			log.Debug("Running: %s %s", editor, tmpFileName)
			cmd := exec.Command(editor, tmpFileName)
			cmd.Stdout, cmd.Stderr, cmd.Stdin = os.Stdout, os.Stderr, os.Stdin
			if err := cmd.Run(); err != nil {
				log.Error("Failed to edit template with %s: %s", editor, err)
				if promptYN("edit again?", true) {
					continue
				}
				return err
			}
		}

		edited := make(map[string]interface{})
		if fh, err := ioutil.ReadFile(tmpFileName); err != nil {
			log.Error("Failed to read tmpfile %s: %s", tmpFileName, err)
			if editing && promptYN("edit again?", true) {
				continue
			}
			return err
		} else {
			if err := yaml.Unmarshal(fh, &edited); err != nil {
				log.Error("Failed to parse YAML: %s", err)
				if editing && promptYN("edit again?", true) {
					continue
				}
				return err
			}
		}

		if fixed, err := yamlFixup(edited); err != nil {
			return err
		} else {
			edited = fixed.(map[string]interface{})
		}

		if _, ok := templateData["meta"]; ok {
			mf := templateData["meta"].(map[string]interface{})["fields"]
			if f, ok := edited["fields"].(map[string]interface{}); ok {
				for k, _ := range f {
					if _, ok := mf.(map[string]interface{})[k]; !ok {
						err := fmt.Errorf("Field %s is not editable", k)
						log.Error("%s", err)
						if editing && promptYN("edit again?", true) {
							continue
						}
						return err
					}
				}
			}
		}

		json, err := jsonEncode(edited)
		if err != nil {
			return err
		}

		if err := templateProcessor(json); err != nil {
			log.Error("%s", err)
			if editing && promptYN("edit again?", true) {
				continue
			}
		}
		return nil
	}
	return nil
}

func (c *Cli) Browse(issue string) error {
	if val, ok := c.opts["browse"]; ok && val == "true" {
		if runtime.GOOS == "darwin" {
			return exec.Command("open", fmt.Sprintf("%s/browse/%s", c.endpoint, issue)).Run()
		} else if runtime.GOOS == "linux" {
			return exec.Command("xdg-open", fmt.Sprintf("%s/browse/%s", c.endpoint, issue)).Run()
		}
	}
	return nil
}
<commit_msg>gofmt<commit_after>package cli

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/op/go-logging"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"
)

var log = logging.MustGetLogger("jira.cli")

type Cli struct {
	endpoint   *url.URL
	opts       map[string]string
	cookieFile string
	ua         *http.Client
}

func New(opts map[string]string) *Cli {
	homedir := os.Getenv("HOME")
	cookieJar, _ := cookiejar.New(nil)
	endpoint, _ := opts["endpoint"]
	url, _ := url.Parse(strings.TrimRight(endpoint, "/"))

	if project, ok := opts["project"]; ok {
		opts["project"] = strings.ToUpper(project)
	}

	cli := &Cli{
		endpoint:   url,
		opts:       opts,
		cookieFile: fmt.Sprintf("%s/.jira.d/cookies.js", homedir),
		ua:         &http.Client{Jar: cookieJar},
	}

	cli.ua.Jar.SetCookies(url, cli.loadCookies())

	return cli
}

func (c *Cli) saveCookies(cookies []*http.Cookie) {
	// expiry in one week from now
	expiry := time.Now().Add(24 * 7 * time.Hour)
	for _, cookie := range cookies {
		cookie.Expires = expiry
	}

	if currentCookies := c.loadCookies(); currentCookies != nil {
		currentCookiesByName := make(map[string]*http.Cookie)
		for _, cookie := range currentCookies {
			currentCookiesByName[cookie.Name] = cookie
		}

		for _, cookie := range cookies {
			currentCookiesByName[cookie.Name] = cookie
		}

		mergedCookies := make([]*http.Cookie, 0, len(currentCookiesByName))
		for _, v := range currentCookiesByName {
			mergedCookies = append(mergedCookies, v)
		}
		jsonWrite(c.cookieFile, mergedCookies)
	} else {
		jsonWrite(c.cookieFile, cookies)
	}
}

func (c *Cli) loadCookies() []*http.Cookie {
	bytes, err := ioutil.ReadFile(c.cookieFile)
	if err != nil && os.IsNotExist(err) {
		// dont load cookies if the file does not exist
		return nil
	}
	if err != nil {
		log.Error("Failed to open %s: %s", c.cookieFile, err)
		os.Exit(1)
	}
	cookies := make([]*http.Cookie, 0)
	err = json.Unmarshal(bytes, &cookies)
	if err != nil {
		log.Error("Failed to parse json from file %s: %s", c.cookieFile, err)
	}
	log.Debug("Loading Cookies: %s", cookies)
	return cookies
}

func (c *Cli) post(uri string, content string) (*http.Response, error) {
	return c.makeRequestWithContent("POST", uri, content)
}

func (c *Cli) put(uri string, content string) (*http.Response, error) {
	return c.makeRequestWithContent("PUT", uri, content)
}

func (c *Cli) makeRequestWithContent(method string, uri string, content string) (*http.Response, error) {
	buffer := bytes.NewBufferString(content)
	req, _ := http.NewRequest(method, uri, buffer)

	log.Info("%s %s", req.Method, req.URL.String())
	if log.IsEnabledFor(logging.DEBUG) {
		logBuffer := bytes.NewBuffer(make([]byte, 0, len(content)))
		req.Write(logBuffer)
		log.Debug("%s", logBuffer)
		// need to recreate the buffer since the offset is now at the end
		// need to be able to rewind the buffer offset, dont know how yet
		req, _ = http.NewRequest(method, uri, bytes.NewBufferString(content))
	}

	if resp, err := c.makeRequest(req); err != nil {
		return nil, err
	} else {
		if resp.StatusCode == 401 {
			if err := c.CmdLogin(); err != nil {
				return nil, err
			}
			req, _ = http.NewRequest(method, uri, bytes.NewBufferString(content))
			return c.makeRequest(req)
		}
		return resp, err
	}
}

func (c *Cli) get(uri string) (*http.Response, error) {
	req, _ := http.NewRequest("GET", uri, nil)
	log.Info("%s %s", req.Method, req.URL.String())
	if log.IsEnabledFor(logging.DEBUG) {
		logBuffer := bytes.NewBuffer(make([]byte, 0))
		req.Write(logBuffer)
		log.Debug("%s", logBuffer)
	}

	if resp, err := c.makeRequest(req); err != nil {
		return nil, err
	} else {
		if resp.StatusCode == 401 {
			if err := c.CmdLogin(); err != nil {
				return nil, err
			}
			return c.makeRequest(req)
		}
		return resp, err
	}
}

func (c *Cli) makeRequest(req *http.Request) (resp *http.Response, err error) {
	req.Header.Set("Content-Type", "application/json")
	if resp, err = c.ua.Do(req); err != nil {
		log.Error("Failed to %s %s: %s", req.Method, req.URL.String(), err)
		return nil, err
	} else {
		if resp.StatusCode < 200 || resp.StatusCode >= 300 && resp.StatusCode != 401 {
			log.Error("response status: %s", resp.Status)
		}

		runtime.SetFinalizer(resp, func(r *http.Response) {
			r.Body.Close()
		})

		if _, ok := resp.Header["Set-Cookie"]; ok {
			c.saveCookies(resp.Cookies())
		}
	}
	return resp, nil
}

func (c *Cli) getTemplate(name string) string {
	if override, ok := c.opts["template"]; ok {
		if _, err := os.Stat(override); err == nil {
			return readFile(override)
		} else {
			if file, err := FindClosestParentPath(fmt.Sprintf(".jira.d/templates/%s", override)); err == nil {
				return readFile(file)
			}
			if dflt, ok := all_templates[override]; ok {
				return dflt
			}
		}
	}
	if file, err := FindClosestParentPath(fmt.Sprintf(".jira.d/templates/%s", name)); err != nil {
		// create-bug etc are special, if we dont find it in the path
		// then just return a generic create template
		if strings.HasPrefix(name, "create-") {
			if file, err := FindClosestParentPath(".jira.d/templates/create"); err != nil {
				return all_templates["create"]
			} else {
				return readFile(file)
			}
		}
		return all_templates[name]
	} else {
		return readFile(file)
	}
}

func (c *Cli) editTemplate(template string, tmpFilePrefix string, templateData map[string]interface{}, templateProcessor func(string) error) error {

	tmpdir := fmt.Sprintf("%s/.jira.d/tmp", os.Getenv("HOME"))
	if err := mkdir(tmpdir); err != nil {
		return err
	}

	fh, err := ioutil.TempFile(tmpdir, tmpFilePrefix)
	if err != nil {
		log.Error("Failed to make temp file in %s: %s", tmpdir, err)
		return err
	}
	defer fh.Close()

	tmpFileName := fmt.Sprintf("%s.yml", fh.Name())
	if err := os.Rename(fh.Name(), tmpFileName); err != nil {
		log.Error("Failed to rename %s to %s: %s", fh.Name(), fmt.Sprintf("%s.yml", fh.Name()), err)
		return err
	}

	err = runTemplate(template, templateData, fh)
	if err != nil {
		return err
	}

	fh.Close()

	editor, ok := c.opts["editor"]
	if !ok {
		editor = os.Getenv("JIRA_EDITOR")
		if editor == "" {
			editor = os.Getenv("EDITOR")
			if editor == "" {
				editor = "vim"
			}
		}
	}

	editing := true
	if val, ok := c.opts["edit"]; ok && val == "false" {
		editing = false
	}

	for true {
		if editing {
			log.Debug("Running: %s %s", editor, tmpFileName)
			cmd := exec.Command(editor, tmpFileName)
			cmd.Stdout, cmd.Stderr, cmd.Stdin = os.Stdout, os.Stderr, os.Stdin
			if err := cmd.Run(); err != nil {
				log.Error("Failed to edit template with %s: %s", editor, err)
				if promptYN("edit again?", true) {
					continue
				}
				return err
			}
		}

		edited := make(map[string]interface{})
		if fh, err := ioutil.ReadFile(tmpFileName); err != nil {
			log.Error("Failed to read tmpfile %s: %s", tmpFileName, err)
			if editing && promptYN("edit again?", true) {
				continue
			}
			return err
		} else {
			if err := yaml.Unmarshal(fh, &edited); err != nil {
				log.Error("Failed to parse YAML: %s", err)
				if editing && promptYN("edit again?", true) {
					continue
				}
				return err
			}
		}

		if fixed, err := yamlFixup(edited); err != nil {
			return err
		} else {
			edited = fixed.(map[string]interface{})
		}

		if _, ok := templateData["meta"]; ok {
			mf := templateData["meta"].(map[string]interface{})["fields"]
			if f, ok := edited["fields"].(map[string]interface{}); ok {
				for k := range f {
					if _, ok := mf.(map[string]interface{})[k]; !ok {
						err := fmt.Errorf("Field %s is not editable", k)
						log.Error("%s", err)
						if editing && promptYN("edit again?", true) {
							continue
						}
						return err
					}
				}
			}
		}

		json, err := jsonEncode(edited)
		if err != nil {
			return err
		}

		if err := templateProcessor(json); err != nil {
			log.Error("%s", err)
			if editing && promptYN("edit again?", true) {
				continue
			}
		}
		return nil
	}
	return nil
}

func (c *Cli) Browse(issue string) error {
	if val, ok := c.opts["browse"]; ok && val == "true" {
		if runtime.GOOS == "darwin" {
			return exec.Command("open", fmt.Sprintf("%s/browse/%s", c.endpoint, issue)).Run()
		} else if runtime.GOOS == "linux" {
			return exec.Command("xdg-open", fmt.Sprintf("%s/browse/%s", c.endpoint, issue)).Run()
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package jocko

import (
	"context"
	"crypto/tls"
	"math"
	"net"
	"time"
)

const (
	defaultRTT = 1 * time.Second
	maxTimeout = time.Duration(math.MaxInt32) * time.Millisecond
	minTimeout = time.Duration(math.MinInt32) * time.Millisecond
)

// Dialer is like the net.Dialer API but for opening connections to Jocko brokers.
type Dialer struct {
	// Unique ID for client connections established by this Dialer.
	ClientID string
	// Timeout is the max duration a dial will wait for a connect to complete.
	Timeout time.Duration
	// Deadline is the absolute time after which the dial will fail. Zero is no deadline.
	Deadline time.Time
	// LocalAddr is the local address to dial.
	LocalAddr net.Addr
	// RemoteAddr is the remote address to dial.
	RemoteAddr net.Addr
	// KeepAlive is the keep-alive period for a network connection.
	KeepAlive time.Duration
	// FallbackDelay is the duration to wait before spawning a fallback connection. If 0, default duration is 300ms.
	FallbackDelay time.Duration
	// Resolver species an alternative resolver to use.
	Resolver Resolver
	// TLS enables the Dialer to secure connections. If nil, standard net.Conn is used.
	TLS *tls.Config
	// DualStack enables RFC 6555-compliant "happy eyeballs" dialing.
	DualStack bool
}

var (
	defaultDialer = NewDialer("jocko")
)

// NewDialer creates a new dialer.
func NewDialer(clientID string) *Dialer {
	return &Dialer{
		Timeout:   10 * time.Second,
		DualStack: true,
		ClientID:  clientID,
	}
}

// Dial creates a connection to the broker on the given network and address on the default dialer.
func Dial(network, address string) (*Conn, error) {
	return defaultDialer.Dial(network, address)
}

func DialContext(ctx context.Context, network, address string) (*Conn, error) {
	return defaultDialer.DialContext(ctx, network, address)
}

// Dial creates a connection to the broker on the given network and address.
func (d *Dialer) Dial(network, address string) (*Conn, error) {
	return d.DialContext(context.Background(), network, address)
}

func (d *Dialer) DialContext(ctx context.Context, network, address string) (*Conn, error) {
	if d.Timeout != 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, d.Timeout)
		defer cancel()
	}

	if !d.Deadline.IsZero() {
		var cancel context.CancelFunc
		ctx, cancel = context.WithDeadline(ctx, d.Deadline)
		defer cancel()
	}

	c, err := d.dialContext(ctx, network, address)
	if err != nil {
		return nil, err
	}
	return NewConn(c, d.ClientID)
}

func (d *Dialer) dialContext(ctx context.Context, network, address string) (net.Conn, error) {
	if r := d.Resolver; r != nil {
		host, port := splitHostPort(address)
		addrs, err := r.LookupHost(ctx, host)
		if err != nil {
			return nil, err
		}
		if len(addrs) != 0 {
			address = addrs[0]
		}
		if len(port) != 0 {
			address, _ = splitHostPort(address)
			address = net.JoinHostPort(address, port)
		}
	}
	return (&net.Dialer{
		LocalAddr:     d.LocalAddr,
		FallbackDelay: d.FallbackDelay,
		KeepAlive:     d.KeepAlive,
	}).DialContext(ctx, network, address)
}

func splitHostPort(s string) (string, string) {
	host, port, _ := net.SplitHostPort(s)
	if len(host) == 0 && len(port) == 0 {
		host = s
	}
	return host, port
}

// Resolver provides service discovery of the hosts of a kafka cluster.
type Resolver interface {
	// LookupHost looks up the given host using the local resolver.
	LookupHost(ctx context.Context, host string) ([]string, error)
}
<commit_msg>conn: add tls support<commit_after>package jocko

import (
	"context"
	"crypto/tls"
	"math"
	"net"
	"time"
)

const (
	defaultRTT = 1 * time.Second
	maxTimeout = time.Duration(math.MaxInt32) * time.Millisecond
	minTimeout = time.Duration(math.MinInt32) * time.Millisecond
)

// Dialer is like the net.Dialer API but for opening connections to Jocko brokers.
type Dialer struct {
	// Unique ID for client connections established by this Dialer.
	ClientID string
	// Timeout is the max duration a dial will wait for a connect to complete.
	Timeout time.Duration
	// Deadline is the absolute time after which the dial will fail. Zero is no deadline.
	Deadline time.Time
	// LocalAddr is the local address to dial.
	LocalAddr net.Addr
	// RemoteAddr is the remote address to dial.
	RemoteAddr net.Addr
	// KeepAlive is the keep-alive period for a network connection.
	KeepAlive time.Duration
	// FallbackDelay is the duration to wait before spawning a fallback connection. If 0, default duration is 300ms.
	FallbackDelay time.Duration
	// Resolver species an alternative resolver to use.
	Resolver Resolver
	// TLS enables the Dialer to secure connections. If nil, standard net.Conn is used.
	TLS *tls.Config
	// DualStack enables RFC 6555-compliant "happy eyeballs" dialing.
	DualStack bool
}

var (
	defaultDialer = NewDialer("jocko")
)

// NewDialer creates a new dialer.
func NewDialer(clientID string) *Dialer {
	return &Dialer{
		Timeout:   10 * time.Second,
		DualStack: true,
		ClientID:  clientID,
	}
}

// Dial creates a connection to the broker on the given network and address on the default dialer.
func Dial(network, address string) (*Conn, error) {
	return defaultDialer.Dial(network, address)
}

func DialContext(ctx context.Context, network, address string) (*Conn, error) {
	return defaultDialer.DialContext(ctx, network, address)
}

// Dial creates a connection to the broker on the given network and address.
func (d *Dialer) Dial(network, address string) (*Conn, error) {
	return d.DialContext(context.Background(), network, address)
}

func (d *Dialer) DialContext(ctx context.Context, network, address string) (*Conn, error) {
	if d.Timeout != 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, d.Timeout)
		defer cancel()
	}

	if !d.Deadline.IsZero() {
		var cancel context.CancelFunc
		ctx, cancel = context.WithDeadline(ctx, d.Deadline)
		defer cancel()
	}

	c, err := d.dialContext(ctx, network, address)
	if err != nil {
		return nil, err
	}
	return NewConn(c, d.ClientID)
}

func (d *Dialer) dialContext(ctx context.Context, network, address string) (conn net.Conn, err error) {
	if r := d.Resolver; r != nil {
		host, port := splitHostPort(address)
		addrs, err := r.LookupHost(ctx, host)
		if err != nil {
			return nil, err
		}
		if len(addrs) != 0 {
			address = addrs[0]
		}
		if len(port) != 0 {
			address, _ = splitHostPort(address)
			address = net.JoinHostPort(address, port)
		}
	}
	conn, err = (&net.Dialer{
		LocalAddr:     d.LocalAddr,
		FallbackDelay: d.FallbackDelay,
		KeepAlive:     d.KeepAlive,
	}).DialContext(ctx, network, address)
	if err != nil {
		return
	}

	if d.TLS != nil {
		return d.connectTLS(ctx, conn)
	}

	return conn, nil
}

// TODO: add unit tests
func (d *Dialer) connectTLS(ctx context.Context, conn net.Conn) (tlsConn *tls.Conn, err error) {
	tlsConn = tls.Client(conn, d.TLS)
	errc := make(chan error)
	go func() {
		defer close(errc)
		errc <- tlsConn.Handshake()
	}()
	select {
	case <-ctx.Done():
		conn.Close()
		tlsConn.Close()
		<-errc
		err = ctx.Err()
	case err = <-errc:
	}
	return
}

func splitHostPort(s string) (string, string) {
	host, port, _ := net.SplitHostPort(s)
	if len(host) == 0 && len(port) == 0 {
		host = s
	}
	return host, port
}

// Resolver provides service discovery of the hosts of a kafka cluster.
type Resolver interface {
	// LookupHost looks up the given host using the local resolver.
	LookupHost(ctx context.Context, host string) ([]string, error)
}
<|endoftext|>


<commit_before>package workspace

import (
	"errors"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

var errMissingMetadata = errors.New("no exercise metadata file found")

// IsMissingMetadata verifies the type of error.
func IsMissingMetadata(err error) bool {
	return err == errMissingMetadata
}

// Workspace represents a user's Exercism workspace.
// It may contain a user's own exercises, and other people's
// exercises that they've downloaded to look at or run locally.
type Workspace struct {
	Dir string
}

// New returns a configured workspace.
func New(dir string) (Workspace, error) {
	_, err := os.Lstat(dir)
	if err != nil {
		return Workspace{}, err
	}
	dir, err = filepath.EvalSymlinks(dir)
	if err != nil {
		return Workspace{}, err
	}
	return Workspace{Dir: dir}, nil
}

// PotentialExercises are a first-level guess at the user's exercises.
// It looks at the workspace structurally, and guesses based on
// the location of the directory. E.g. any top level directory
// within the workspace (except 'users') is assumed to be a
// track, and any directory within there again is assumed to
// be an exercise.
func (ws Workspace) PotentialExercises() ([]Exercise, error) {
	exercises := []Exercise{}

	topInfos, err := ioutil.ReadDir(ws.Dir)
	if err != nil {
		return nil, err
	}
	for _, topInfo := range topInfos {
		if !topInfo.IsDir() {
			continue
		}

		if topInfo.Name() == "users" {
			continue
		}

		subInfos, err := ioutil.ReadDir(filepath.Join(ws.Dir, topInfo.Name()))
		if err != nil {
			return nil, err
		}

		for _, subInfo := range subInfos {
			if !subInfo.IsDir() {
				continue
			}

			exercises = append(exercises, Exercise{Track: topInfo.Name(), Slug: subInfo.Name(), Root: ws.Dir})
		}
	}

	return exercises, nil
}

// Exercises returns the user's exercises within the workspace.
// This doesn't find legacy exercises where the metadata is missing.
func (ws Workspace) Exercises() ([]Exercise, error) {
	candidates, err := ws.PotentialExercises()
	if err != nil {
		return nil, err
	}

	exercises := make([]Exercise, 0, len(candidates))
	for _, candidate := range candidates {
		ok, err := candidate.HasMetadata()
		if err != nil {
			return nil, err
		}
		if ok {
			exercises = append(exercises, candidate)
		}
	}
	return exercises, nil
}

// ExerciseDir determines the root directory of an exercise.
// This is the directory that contains the exercise metadata file.
func (ws Workspace) ExerciseDir(s string) (string, error) {
	if !strings.HasPrefix(s, ws.Dir) {
		return "", errors.New("not in workspace")
	}

	path := s
	for {
		if path == ws.Dir {
			return "", errMissingMetadata
		}
		if _, err := os.Lstat(path); os.IsNotExist(err) {
			return "", err
		}
		if _, err := os.Lstat(filepath.Join(path, metadataFilepath)); err == nil {
			return path, nil
		}
		if _, err := os.Lstat(filepath.Join(path, legacyMetadataFilename)); err == nil {
			return path, nil
		}
		path = filepath.Dir(path)
	}
}
<commit_msg>add teams support to workspace.PotentialExercises<commit_after>package workspace

import (
	"errors"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

var errMissingMetadata = errors.New("no exercise metadata file found")

// IsMissingMetadata verifies the type of error.
func IsMissingMetadata(err error) bool {
	return err == errMissingMetadata
}

// Workspace represents a user's Exercism workspace.
// It may contain a user's own exercises, and other people's
// exercises that they've downloaded to look at or run locally.
type Workspace struct {
	Dir string
}

// New returns a configured workspace.
func New(dir string) (Workspace, error) {
	_, err := os.Lstat(dir)
	if err != nil {
		return Workspace{}, err
	}
	dir, err = filepath.EvalSymlinks(dir)
	if err != nil {
		return Workspace{}, err
	}
	return Workspace{Dir: dir}, nil
}

// PotentialExercises are a first-level guess at the user's exercises.
// It looks at the workspace structurally, and guesses based on
// the location of the directory. E.g. any top level directory
// within the workspace (except 'users') is assumed to be a
// track, and any directory within there again is assumed to
// be an exercise.
func (ws Workspace) PotentialExercises() ([]Exercise, error) {
	exercises := []Exercise{}

	topInfos, err := ioutil.ReadDir(ws.Dir)
	if err != nil {
		return nil, err
	}
	for _, topInfo := range topInfos {
		if !topInfo.IsDir() {
			continue
		}

		if topInfo.Name() == "users" {
			continue
		}

		if topInfo.Name() == "teams" {
			subInfos, err := ioutil.ReadDir(filepath.Join(ws.Dir, "teams"))
			if err != nil {
				return nil, err
			}

			for _, subInfo := range subInfos {
				teamWs, err := New(filepath.Join(ws.Dir, "teams", subInfo.Name()))
				if err != nil {
					return nil, err
				}

				teamExercises, err := teamWs.PotentialExercises()
				if err != nil {
					return nil, err
				}

				exercises = append(exercises, teamExercises...)
			}
			continue
		}

		subInfos, err := ioutil.ReadDir(filepath.Join(ws.Dir, topInfo.Name()))
		if err != nil {
			return nil, err
		}

		for _, subInfo := range subInfos {
			if !subInfo.IsDir() {
				continue
			}

			exercises = append(exercises, Exercise{Track: topInfo.Name(), Slug: subInfo.Name(), Root: ws.Dir})
		}
	}

	return exercises, nil
}

// Exercises returns the user's exercises within the workspace.
// This doesn't find legacy exercises where the metadata is missing.
func (ws Workspace) Exercises() ([]Exercise, error) {
	candidates, err := ws.PotentialExercises()
	if err != nil {
		return nil, err
	}

	exercises := make([]Exercise, 0, len(candidates))
	for _, candidate := range candidates {
		ok, err := candidate.HasMetadata()
		if err != nil {
			return nil, err
		}
		if ok {
			exercises = append(exercises, candidate)
		}
	}
	return exercises, nil
}

// ExerciseDir determines the root directory of an exercise.
// This is the directory that contains the exercise metadata file.
func (ws Workspace) ExerciseDir(s string) (string, error) {
	if !strings.HasPrefix(s, ws.Dir) {
		return "", errors.New("not in workspace")
	}

	path := s
	for {
		if path == ws.Dir {
			return "", errMissingMetadata
		}
		if _, err := os.Lstat(path); os.IsNotExist(err) {
			return "", err
		}
		if _, err := os.Lstat(filepath.Join(path, metadataFilepath)); err == nil {
			return path, nil
		}
		if _, err := os.Lstat(filepath.Join(path, legacyMetadataFilename)); err == nil {
			return path, nil
		}
		path = filepath.Dir(path)
	}
}
<|endoftext|>


<commit_before>package mirango

import (
	"fmt"
	"strings"

	"github.com/mirango/defaults"
	"github.com/mirango/framework"
	"github.com/mirango/validation"
)

type Middleware interface {
	Run(Handler) Handler
}

type MiddlewareFunc func(Handler) Handler

func (f MiddlewareFunc) Run(h Handler) Handler {
	return f(h)
}

func CheckSchemes(schemes []string) MiddlewareFunc {
	return MiddlewareFunc(func(next Handler) Handler {
		return HandlerFunc(func(c *Context) interface{} {

			var schemeAccepted bool
			if c.URL.Scheme == "" {
				c.URL.Scheme = framework.SCHEME_HTTP
			}

			schemeAccepted = containsString(o.GetAllSchemes(), c.URL.Scheme)
			if !schemeAccepted {
				return nil
			}

			return next.ServeHTTP(c)
		})
	})
}

func getEncodingFromAccept(returns []string, r *Request) (string, *Error) {
	var encoding string

	parts := strings.Split(r.Header.Get(framework.HEADER_Accept), ",")

	for _, acceptMime := range parts {
		mime := strings.Trim(strings.Split(acceptMime, ";")[0], " ")
		if 0 == len(mime) || mime == "*/*" {
			if len(returns) == 0 {
				encoding = defaults.MimeType
				break
			} else {
				encoding = returns[0]
				break
			}
		} else {
			if containsString(returns, mime) {
				encoding = mime
				break
			}
		}
	}

	if len(parts) == 0 {
		encoding = defaults.MimeType
	}

	if len(encoding) == 0 {
		// return encoding, NewError(406, "Request", "Encoding requested not valid.")
		encoding = defaults.MimeType
	}

	return encoding, nil
}

func CheckReturns(returns []string, mimeTypeIn paramIn, mimeTypeParam string) MiddlewareFunc {
	return MiddlewareFunc(func(next Handler) Handler {
		return HandlerFunc(func(c *Context) interface{} {
			// if mimeInAccept {
			enc, err := getEncodingFromAccept(returns, c.Request)
			if err != nil {
				return err
			}
			c.encoding = enc
			// } else {
			//
			// }

			return next.ServeHTTP(c)
		})
	})
}

func CheckAccepts(accepts []string) MiddlewareFunc {
	return MiddlewareFunc(func(next Handler) Handler {
		return HandlerFunc(func(c *Context) interface{} {
			// enc, err := getEncodingFromAccept(accepts, c.Request)
			// if err != nil {
			// 	return err
			// }
			// c.encoding = enc

			return next.ServeHTTP(c)
		})
	})
}

func ValidateParams(params *Params) MiddlewareFunc {
	return MiddlewareFunc(func(next Handler) Handler {
		return HandlerFunc(func(c *Context) interface{} {

			var errs *validation.Error

			q := c.URL.Query()
			h := c.Request.Header

			if params.ContainsFiles() {
				c.ParseMultipartForm(defaults.MaxMemory)
			}
			if params.ContainsBodyParams() {
				c.ParseForm()
			}
			for _, p := range params.Get() {
				var pv *validation.Value
				if p.IsIn(IN_QUERY) {
					v, ok := q[p.name]
					if !ok {
						if p.IsRequired() {
							if errs == nil {
								errs = &validation.Error{}
							}
							errs.Append(p.name, fmt.Errorf("param %s is required", p.name))
						}
					} else {
						if !p.IsMultiple() {
							pv = validation.NewValue(p.name, v[0], "query", p.GetAs())
						} else {
							pv = validation.NewMultipleValue(p.name, v, "query", p.GetAs())
						}
					}
				} else if p.IsIn(IN_HEADER) {
					v, ok := h[p.name]
					if !ok {
						if p.IsRequired() {
							if errs == nil {
								errs = &validation.Error{}
							}
							errs.Append(p.name, fmt.Errorf("param %s is required", p.name))
						}
					} else {
						if !p.IsMultiple() {
							pv = validation.NewValue(p.name, v[0], "header", p.GetAs())
						} else {
							pv = validation.NewMultipleValue(p.name, v, "header", p.GetAs())
						}
					}
				} else if p.IsIn(IN_BODY) { // decide what to do when content type is form-encoded
					if p.IsFile() {
						_, ok := c.MultipartForm.File[p.name]
						if !ok {
							if p.IsRequired() {
								if errs == nil {
									errs = &validation.Error{}
								}
								errs.Append(p.name, fmt.Errorf("param %s is required", p.name))
							}
						} else {
							//pv = NewFileParamValue(p.name, v[0], "header")
						}
					} else if !p.IsFile() && params.ContainsFiles() {
						_, ok := c.MultipartForm.Value[p.name]
						if !ok {
							if p.IsRequired() {
								if errs == nil {
									errs = &validation.Error{}
								}
								errs.Append(p.name, fmt.Errorf("param %s is required", p.name))
							}
						} else {
							//pv = NewFileParamValue(p.name, v[0], "header")
						}
					} else {
						v, ok := c.Form[p.name]
						if !ok {
							if p.IsRequired() {
								if errs == nil {
									errs = &validation.Error{}
								}
								errs.Append(p.name, fmt.Errorf("param %s is required", p.name))
							}
						} else {
							if !p.IsMultiple() {
								pv = validation.NewValue(p.name, v[0], "body", p.GetAs())
							} else {
								pv = validation.NewMultipleValue(p.name, v, "body", p.GetAs())
							}
						}
					}
				}
				if pv != nil {
					c.Input[p.name] = pv
				}
			}

			vErrs := params.ValidateAll(c, c.Input)
			if vErrs != nil {
				if errs == nil {
					errs = &validation.Error{}
				}
				errs.UnionAppend(*vErrs)
			}

			if errs != nil {
				return errs
			}

			return next.ServeHTTP(c)
		})
	})
}
<commit_msg>Use Operation as a parameter for checking Milddeware<commit_after>package mirango

import (
	"fmt"
	"strings"

	"github.com/mirango/defaults"
	"github.com/mirango/framework"
	"github.com/mirango/validation"
)

type Middleware interface {
	Run(Handler) Handler
}

type MiddlewareFunc func(Handler) Handler

func (f MiddlewareFunc) Run(h Handler) Handler {
	return f(h)
}

func CheckSchemes(o *Operation) MiddlewareFunc {
	return MiddlewareFunc(func(next Handler) Handler {
		return HandlerFunc(func(c *Context) interface{} {

			var schemeAccepted bool
			if c.URL.Scheme == "" {
				c.URL.Scheme = framework.SCHEME_HTTP
			}

			schemeAccepted = containsString(o.GetAllSchemes(), c.URL.Scheme)
			if !schemeAccepted {
				return nil
			}

			return next.ServeHTTP(c)
		})
	})
}

func getEncodingFromAccept(returns []string, r *Request) (string, *Error) {
	var encoding string

	parts := strings.Split(r.Header.Get(framework.HEADER_Accept), ",")

	for _, acceptMime := range parts {
		mime := strings.Trim(strings.Split(acceptMime, ";")[0], " ")
		if 0 == len(mime) || mime == "*/*" {
			if len(returns) == 0 {
				encoding = defaults.MimeType
				break
			} else {
				encoding = returns[0]
				break
			}
		} else {
			if containsString(returns, mime) {
				encoding = mime
				break
			}
		}
	}

	if len(parts) == 0 {
		encoding = defaults.MimeType
	}

	if len(encoding) == 0 {
		// return encoding, NewError(406, "Request", "Encoding requested not valid.")
		encoding = defaults.MimeType
	}

	return encoding, nil
}

func CheckReturns(o *Operation) MiddlewareFunc {
	return MiddlewareFunc(func(next Handler) Handler {
		return HandlerFunc(func(c *Context) interface{} {
			// if mimeInAccept {
			enc, err := getEncodingFromAccept(o.GetAllReturns(), c.Request)
			if err != nil {
				return err
			}
			c.encoding = enc
			// } else {
			//
			// }

			return next.ServeHTTP(c)
		})
	})
}

func CheckAccepts(o *Operation) MiddlewareFunc {
	return MiddlewareFunc(func(next Handler) Handler {
		return HandlerFunc(func(c *Context) interface{} {
			// enc, err := getEncodingFromAccept(accepts, c.Request)
			// if err != nil {
			// 	return err
			// }
			// c.encoding = enc

			return next.ServeHTTP(c)
		})
	})
}

func CheckParams(o *Operation) MiddlewareFunc {
	return MiddlewareFunc(func(next Handler) Handler {
		return HandlerFunc(func(c *Context) interface{} {
			params := o.GetAllParams()
			var errs *validation.Error

			q := c.URL.Query()
			h := c.Request.Header

			if params.ContainsFiles() {
				c.ParseMultipartForm(defaults.MaxMemory)
			}
			if params.ContainsBodyParams() {
				c.ParseForm()
			}
			for _, p := range params.GetAll() {
				var pv *validation.Value
				if p.IsIn(IN_QUERY) {
					v, ok := q[p.name]
					if !ok {
						if p.IsRequired() {
							if errs == nil {
								errs = &validation.Error{}
							}
							errs.Append(p.name, fmt.Errorf("param %s is required", p.name))
						}
					} else {
						if !p.IsMultiple() {
							pv = validation.NewValue(p.name, v[0], "query", p.GetAs())
						} else {
							pv = validation.NewMultipleValue(p.name, v, "query", p.GetAs())
						}
					}
				} else if p.IsIn(IN_HEADER) {
					v, ok := h[p.name]
					if !ok {
						if p.IsRequired() {
							if errs == nil {
								errs = &validation.Error{}
							}
							errs.Append(p.name, fmt.Errorf("param %s is required", p.name))
						}
					} else {
						if !p.IsMultiple() {
							pv = validation.NewValue(p.name, v[0], "header", p.GetAs())
						} else {
							pv = validation.NewMultipleValue(p.name, v, "header", p.GetAs())
						}
					}
				} else if p.IsIn(IN_BODY) { // decide what to do when content type is form-encoded
					if p.IsFile() {
						_, ok := c.MultipartForm.File[p.name]
						if !ok {
							if p.IsRequired() {
								if errs == nil {
									errs = &validation.Error{}
								}
								errs.Append(p.name, fmt.Errorf("param %s is required", p.name))
							}
						} else {
							//pv = NewFileParamValue(p.name, v[0], "header")
						}
					} else if !p.IsFile() && params.ContainsFiles() {
						_, ok := c.MultipartForm.Value[p.name]
						if !ok {
							if p.IsRequired() {
								if errs == nil {
									errs = &validation.Error{}
								}
								errs.Append(p.name, fmt.Errorf("param %s is required", p.name))
							}
						} else {
							//pv = NewFileParamValue(p.name, v[0], "header")
						}
					} else {
						v, ok := c.Form[p.name]
						if !ok {
							if p.IsRequired() {
								if errs == nil {
									errs = &validation.Error{}
								}
								errs.Append(p.name, fmt.Errorf("param %s is required", p.name))
							}
						} else {
							if !p.IsMultiple() {
								pv = validation.NewValue(p.name, v[0], "body", p.GetAs())
							} else {
								pv = validation.NewMultipleValue(p.name, v, "body", p.GetAs())
							}
						}
					}
				}
				if pv != nil {
					c.Input[p.name] = pv
				}
			}

			vErrs := params.ValidateAll(c, c.Input)
			if vErrs != nil {
				if errs == nil {
					errs = &validation.Error{}
				}
				errs.UnionAppend(*vErrs)
			}

			if errs != nil {
				return errs
			}

			return next.ServeHTTP(c)
		})
	})
}
<|endoftext|>


<commit_before>package buffalo

import (
	"reflect"
	"runtime"
	"strings"
)

// MiddlewareFunc defines the interface for a piece of Buffalo
// Middleware.
/*
	func DoSomething(next Handler) Handler {
		return func(c Context) error {
			// do something before calling the next handler
			err := next(c)
			// do something after call the handler
			return err
		}
	}
*/
type MiddlewareFunc func(Handler) Handler

// Use the specified Middleware for the App.
// When defined on an `*App` the specified middleware will be
// inherited by any `Group` calls that are made on that on
// the App.
func (a *App) Use(mw ...MiddlewareFunc) {
	a.Middleware.Use(mw...)
}

// MiddlewareStack manages the middleware stack for an App/Group.
type MiddlewareStack struct {
	stack []MiddlewareFunc
	skips map[string]bool
}

func (ms *MiddlewareStack) clone() *MiddlewareStack {
	n := newMiddlewareStack()
	for _, s := range ms.stack {
		n.stack = append(n.stack, s)
	}
	for k, v := range ms.skips {
		n.skips[k] = v
	}
	return n
}

// Clear wipes out the current middleware stack for the App/Group,
// any middleware previously defined will be removed leaving an empty
// middleware stack.
func (ms *MiddlewareStack) Clear() {
	ms.stack = []MiddlewareFunc{}
	ms.skips = map[string]bool{}
}

// Use the specified Middleware for the App.
// When defined on an `*App` the specified middleware will be
// inherited by any `Group` calls that are made on that on
// the App.
func (ms *MiddlewareStack) Use(mw ...MiddlewareFunc) {
	ms.stack = append(ms.stack, mw...)
}

// Skip a specified piece of middleware the specified Handlers.
// This is useful for things like wrapping your application in an
// authorization middleare, but skipping it for things the home
// page, the login page, etc...
/*
	a.Middleware.Skip(Authorization, HomeHandler, LoginHandler, RegistrationHandler)
*/
// NOTE: When skipping Resource handlers, you need to first declare your
// resource handler as a type of buffalo.Resource for the Skip function to
// properly recognize and match it.
/*
	// Works:
	var ur Resource
	cr = &carsResource{&buffaloBaseResource{}}
	g = a.Resource("/cars", cr)
	g.Use(SomeMiddleware)
	g.Middleware.Skip(SomeMiddleware, cr.Show)

	// Doesn't Work:
	cr := &carsResource{&buffaloBaseResource{}}
	g = a.Resource("/cars", cr)
	g.Use(SomeMiddleware)
	g.Middleware.Skip(SomeMiddleware, cr.Show)
*/
func (ms *MiddlewareStack) Skip(mw MiddlewareFunc, handlers ...Handler) {
	for _, h := range handlers {
		key := funcKey(mw, h)
		ms.skips[key] = true
	}
}

// Replace a piece of middleware with another piece of middleware. Great for
// testing.
func (ms *MiddlewareStack) Replace(mw1 MiddlewareFunc, mw2 MiddlewareFunc) {
	m1k := funcKey(mw1)
	stack := []MiddlewareFunc{}
	for _, mw := range ms.stack {
		if funcKey(mw) == m1k {
			stack = append(stack, mw2)
		} else {
			stack = append(stack, mw)
		}
	}
	ms.stack = stack
}

func (ms *MiddlewareStack) handler(h Handler) Handler {
	if len(ms.stack) > 0 {
		mh := func(_ Handler) Handler {
			return h
		}

		tstack := []MiddlewareFunc{mh}

		sl := len(ms.stack) - 1
		for i := sl; i >= 0; i-- {
			mw := ms.stack[i]
			key := funcKey(mw, h)
			if !ms.skips[key] {
				tstack = append(tstack, mw)
			}
		}

		for _, mw := range tstack {
			h = mw(h)
		}
		return h
	}
	return h
}

func newMiddlewareStack(mws ...MiddlewareFunc) *MiddlewareStack {
	return &MiddlewareStack{
		stack: mws,
		skips: map[string]bool{},
	}
}

func funcKey(funcs ...interface{}) string {
	names := []string{}
	for _, f := range funcs {
		rv := reflect.ValueOf(f)
		ptr := rv.Pointer()
		if n, ok := keyMap[ptr]; ok {
			names = append(names, n)
			continue
		}
		fnc := runtime.FuncForPC(ptr)
		n := fnc.Name()
		keyMap[ptr] = n
		names = append(names, n)
	}
	return strings.Join(names, "/")
}

func setFuncKey(f interface{}, name string) {
	rv := reflect.ValueOf(f)
	ptr := rv.Pointer()
	keyMap[ptr] = name
}

var keyMap = map[uintptr]string{}
<commit_msg>Fix typos in GoDoc<commit_after>package buffalo

import (
	"reflect"
	"runtime"
	"strings"
)

// MiddlewareFunc defines the interface for a piece of Buffalo
// Middleware.
/*
	func DoSomething(next Handler) Handler {
		return func(c Context) error {
			// do something before calling the next handler
			err := next(c)
			// do something after call the handler
			return err
		}
	}
*/
type MiddlewareFunc func(Handler) Handler

// Use the specified Middleware for the App.
// When defined on an `*App` the specified middleware will be
// inherited by any `Group` calls that are made on that on
// the App.
func (a *App) Use(mw ...MiddlewareFunc) {
	a.Middleware.Use(mw...)
}

// MiddlewareStack manages the middleware stack for an App/Group.
type MiddlewareStack struct {
	stack []MiddlewareFunc
	skips map[string]bool
}

func (ms *MiddlewareStack) clone() *MiddlewareStack {
	n := newMiddlewareStack()
	for _, s := range ms.stack {
		n.stack = append(n.stack, s)
	}
	for k, v := range ms.skips {
		n.skips[k] = v
	}
	return n
}

// Clear wipes out the current middleware stack for the App/Group,
// any middleware previously defined will be removed leaving an empty
// middleware stack.
func (ms *MiddlewareStack) Clear() {
	ms.stack = []MiddlewareFunc{}
	ms.skips = map[string]bool{}
}

// Use the specified Middleware for the App.
// When defined on an `*App` the specified middleware will be
// inherited by any `Group` calls that are made on that on
// the App.
func (ms *MiddlewareStack) Use(mw ...MiddlewareFunc) {
	ms.stack = append(ms.stack, mw...)
}

// Skip a specified piece of middleware the specified Handlers.
// This is useful for things like wrapping your application in an
// authorization middleware, but skipping it for things the home
// page, the login page, etc...
/*
	a.Middleware.Skip(Authorization, HomeHandler, LoginHandler, RegistrationHandler)
*/
// NOTE: When skipping Resource handlers, you need to first declare your
// resource handler as a type of buffalo.Resource for the Skip function to
// properly recognize and match it.
/*
	// Works:
	var cr Resource
	cr = &carsResource{&buffaloBaseResource{}}
	g = a.Resource("/cars", cr)
	g.Use(SomeMiddleware)
	g.Middleware.Skip(SomeMiddleware, cr.Show)

	// Doesn't Work:
	cr := &carsResource{&buffaloBaseResource{}}
	g = a.Resource("/cars", cr)
	g.Use(SomeMiddleware)
	g.Middleware.Skip(SomeMiddleware, cr.Show)
*/
func (ms *MiddlewareStack) Skip(mw MiddlewareFunc, handlers ...Handler) {
	for _, h := range handlers {
		key := funcKey(mw, h)
		ms.skips[key] = true
	}
}

// Replace a piece of middleware with another piece of middleware. Great for
// testing.
func (ms *MiddlewareStack) Replace(mw1 MiddlewareFunc, mw2 MiddlewareFunc) {
	m1k := funcKey(mw1)
	stack := []MiddlewareFunc{}
	for _, mw := range ms.stack {
		if funcKey(mw) == m1k {
			stack = append(stack, mw2)
		} else {
			stack = append(stack, mw)
		}
	}
	ms.stack = stack
}

func (ms *MiddlewareStack) handler(h Handler) Handler {
	if len(ms.stack) > 0 {
		mh := func(_ Handler) Handler {
			return h
		}

		tstack := []MiddlewareFunc{mh}

		sl := len(ms.stack) - 1
		for i := sl; i >= 0; i-- {
			mw := ms.stack[i]
			key := funcKey(mw, h)
			if !ms.skips[key] {
				tstack = append(tstack, mw)
			}
		}

		for _, mw := range tstack {
			h = mw(h)
		}
		return h
	}
	return h
}

func newMiddlewareStack(mws ...MiddlewareFunc) *MiddlewareStack {
	return &MiddlewareStack{
		stack: mws,
		skips: map[string]bool{},
	}
}

func funcKey(funcs ...interface{}) string {
	names := []string{}
	for _, f := range funcs {
		rv := reflect.ValueOf(f)
		ptr := rv.Pointer()
		if n, ok := keyMap[ptr]; ok {
			names = append(names, n)
			continue
		}
		fnc := runtime.FuncForPC(ptr)
		n := fnc.Name()
		keyMap[ptr] = n
		names = append(names, n)
	}
	return strings.Join(names, "/")
}

func setFuncKey(f interface{}, name string) {
	rv := reflect.ValueOf(f)
	ptr := rv.Pointer()
	keyMap[ptr] = name
}

var keyMap = map[uintptr]string{}
<|endoftext|>


<commit_before>package guestmetric

import (
	"bufio"
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
	xenstoreclient "xe-guest-utilities/xenstoreclient"
)

type Collector struct {
	Client xenstoreclient.XenStoreClient
	Ballon bool
	Debug  bool
}

func (c *Collector) CollectOS() (GuestMetric, error) {
	current := make(GuestMetric, 0)
	f, err := os.OpenFile("/var/cache/xe-linux-distribution", os.O_RDONLY, 0666)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, "=") {
			parts := strings.SplitN(line, "=", 2)
			k := strings.TrimSpace(parts[0])
			v := strings.TrimSpace(strings.Trim(strings.TrimSpace(parts[1]), "\""))
			current[k] = v
		}
	}
	return prefixKeys("data/", current), nil
}

func (c *Collector) CollectMisc() (GuestMetric, error) {
	current := make(GuestMetric, 0)
	if c.Ballon {
		current["control/feature-balloon"] = "1"
	} else {
		current["control/feature-balloon"] = "0"
	}
	current["attr/PVAddons/Installed"] = "1"
	current["attr/PVAddons/MajorVersion"] = "@PRODUCT_MAJOR_VERSION@"
	current["attr/PVAddons/MinorVersion"] = "@PRODUCT_MINOR_VERSION@"
	current["attr/PVAddons/MicroVersion"] = "@PRODUCT_MICRO_VERSION@"
	current["attr/PVAddons/BuildVersion"] = "@NUMERIC_BUILD_NUMBER@"

	return current, nil
}

func (c *Collector) CollectMemory() (GuestMetric, error) {
	current := make(GuestMetric, 0)
	f, err := os.Open("/proc/meminfo")
	if err != nil {
		return nil, err
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		parts := regexp.MustCompile(`\w+`).FindAllString(scanner.Text(), -1)
		switch parts[0] {
		case "MemTotal":
			current["meminfo_total"] = parts[1]
		case "MemAvailable":
			current["meminfo_free"] = parts[1]
		}
	}
	return prefixKeys("data/", current), nil
}

func enumNetworkAddresses(iface string) (GuestMetric, error) {
	const (
		IP_RE   string = `(\d{1,3}\.){3}\d{1,3}`
		IPV6_RE string = `[\da-f:]+[\da-f]`
	)

	var (
		IP_IPV4_ADDR_RE       = regexp.MustCompile(`inet\s*(` + IP_RE + `).*\se[a-zA-Z0-9]+[\s\n]`)
		IP_IPV6_ADDR_RE       = regexp.MustCompile(`inet6\s*(` + IPV6_RE + `)`)
		IFCONFIG_IPV4_ADDR_RE = regexp.MustCompile(`inet addr:\s*(` + IP_RE + `)`)
		IFCONFIG_IPV6_ADDR_RE = regexp.MustCompile(`inet6 addr:\s*(` + IPV6_RE + `)`)
	)

	d := make(GuestMetric, 0)

	var v4re, v6re *regexp.Regexp
	var out string
	var err error
	if out, err = runCmd("ip", "addr", "show", iface); err == nil {
		v4re = IP_IPV4_ADDR_RE
		v6re = IP_IPV6_ADDR_RE
	} else if out, err = runCmd("ifconfig", iface); err == nil {
		v4re = IFCONFIG_IPV4_ADDR_RE
		v6re = IFCONFIG_IPV6_ADDR_RE
	} else {
		return nil, fmt.Errorf("Cannot find ip/ifconfig command")
	}

	m := v4re.FindAllStringSubmatch(out, -1)
	if m != nil {
		for i, parts := range m {
			d[fmt.Sprintf("ipv4/%d", i)] = parts[1]
		}
	}
	m = v6re.FindAllStringSubmatch(out, -1)
	if m != nil {
		for i, parts := range m {
			d[fmt.Sprintf("ipv6/%d", i)] = parts[1]
		}
	}

	return d, nil
}

func getPlainVifId(path string) (string, error) {
	nodenamePath := fmt.Sprintf("%s/device/nodename", path)
	strLine, err := readSysfs(nodenamePath)
	if err != nil {
		return "", err
	}
	vifId := ""
	reNodename := regexp.MustCompile(`^device\/vif\/(\d+)$`)
	if matched := reNodename.FindStringSubmatch(strLine); matched != nil {
		vifId = matched[1]
	}
	if vifId == "" {
		return "", fmt.Errorf("Not found string like \"device/vif/[id]\" in file %s", nodenamePath)
	} else {
		return vifId, nil
	}
}

func (c *Collector) getSriovVifId(path string) (string, error) {
	sriovDevicePath := "xenserver/device/net-sriov-vf"
	macAddress, err := readSysfs(path + "/address")
	if err != nil {
		return "", err
	}
	subPaths, err := c.Client.List(sriovDevicePath)
	if err != nil {
		return "", err
	}
	for _, subPath := range subPaths {
		iterMac, err := c.Client.Read(fmt.Sprintf("%s/%s/mac", sriovDevicePath, subPath))
		if err != nil {
			continue
		}
		if iterMac == macAddress {
			return subPath, nil
		}
	}
	return "", fmt.Errorf("Cannot find a MAC address to map with %s", path)
}

// return vif_xenstore_prefix * vif_id * error where
// `vif_xenstore_prefix` could be either `attr/vif` for plain VIF or
// `xenserver/attr/net-sriov-vf` for SR-IOV VIF
func (c *Collector) getTargetXenstorePath(path string) (string, string, error) {
	plainVifPrefix := "attr/vif"
	sriovVifPrefix := "xenserver/attr/net-sriov-vf"
	// try to get `vif_id` from nodename interface, only a plain VIF have the nodename interface.
	vifId, err1 := getPlainVifId(path)
	if vifId != "" {
		return plainVifPrefix, vifId, nil
	}
	// not a plain VIF, it could possible be an SR-IOV VIF, try to get vif_id from MAC address mapping
	vifId, err2 := c.getSriovVifId(path)
	if vifId != "" {
		return sriovVifPrefix, vifId, nil
	}
	return "", "", fmt.Errorf("Failed to get VIF ID, errors: %s | %s", err1.Error(), err2.Error())
}

func (c *Collector) CollectNetworkAddr() (GuestMetric, error) {
	current := make(GuestMetric, 0)

	var paths []string
	vifNamePrefixList := [...]string{"eth", "eno", "ens", "emp", "enx", "enX"}
	for _, prefix := range vifNamePrefixList {
		prefixPaths, err := filepath.Glob(fmt.Sprintf("/sys/class/net/%s*", prefix))
		if err != nil {
			return nil, err
		}
		paths = append(paths, prefixPaths...)
	}
	for _, path := range paths {
		// a path is going to be like "/sys/class/net/eth0"
		prefix, vifId, err := c.getTargetXenstorePath(path)
		if err != nil {
			continue
		}
		iface := filepath.Base(path)
		if addrs, err := enumNetworkAddresses(iface); err == nil {
			for tag, addr := range addrs {
				current[fmt.Sprintf("%s/%s/%s", prefix, vifId, tag)] = addr
			}
		}
	}
	return current, nil
}

func readSysfs(filename string) (string, error) {
	f, err := os.OpenFile(filename, os.O_RDONLY, 0666)
	if err != nil {
		return "", err
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	scanner.Scan()
	return scanner.Text(), nil
}

func (c *Collector) CollectDisk() (GuestMetric, error) {
	pi := make(GuestMetric, 0)

	disks := make([]string, 0)
	paths, err := filepath.Glob("/sys/block/*/device")
	if err != nil {
		return nil, err
	}
	for _, path := range paths {
		disk := filepath.Base(strings.TrimSuffix(filepath.Dir(path), "/"))
		disks = append(disks, disk)
	}

	var sortedDisks sort.StringSlice = disks
	sortedDisks.Sort()

	part_idx := 0
	for _, disk := range sortedDisks[:] {
		paths, err = filepath.Glob(fmt.Sprintf("/dev/%s?*", disk))
		if err != nil {
			return nil, err
		}
		for _, path := range paths {
			p := filepath.Base(path)
			line, err := readSysfs(fmt.Sprintf("/sys/block/%s/%s/size", disk, p))
			if err != nil {
				return nil, err
			}
			size, err := strconv.ParseInt(line, 10, 64)
			if err != nil {
				return nil, err
			}
			blocksize := 512
			if bs, err := readSysfs(fmt.Sprintf("/sys/block/%s/queue/physical_block_size", p)); err == nil {
				if bs1, err := strconv.Atoi(bs); err == nil {
					blocksize = bs1
				}
			}
			real_dev := ""
			if c.Client != nil {
				nodename, err := readSysfs(fmt.Sprintf("/sys/block/%s/device/nodename", disk))
				if err == nil {
					backend, err := c.Client.Read(fmt.Sprintf("%s/backend", nodename))
					if err != nil {
						return nil, err
					}
					real_dev, err = c.Client.Read(fmt.Sprintf("%s/dev", backend))
					if err != nil {
						return nil, err
					}
				}
			}
			name := path
			blkid, err := runCmd("blkid", "-s", "UUID", path)
			if err != nil {
				// ignore blkid errors
				blkid = ""
			}
			if strings.Contains(blkid, "=") {
				parts := strings.SplitN(strings.TrimSpace(blkid), "=", 2)
				name = fmt.Sprintf("%s(%s)", name, strings.Trim(parts[1], "\""))
			}
			i := map[string]string{
				"extents/0": real_dev,
				"name":      name,
				"size":      strconv.FormatInt(size*int64(blocksize), 10),
			}
			output, err := runCmd("pvs", "--noheadings", "--units", "b", path)
			if err == nil && output != "" {
				parts := regexp.MustCompile(`\s+`).Split(output, -1)[1:]
				i["free"] = strings.TrimSpace(parts[5])[:len(parts[5])-1]
				i["filesystem"] = strings.TrimSpace(parts[2])
				i["mount_points/0"] = "[LVM]"
			} else {
				output, err = runCmd("mount")
				if err == nil {
					m := regexp.MustCompile(`(?m)^(\S+) on (\S+) type (\S+)`).FindAllStringSubmatch(output, -1)
					if m != nil {
						for _, parts := range m {
							if parts[1] == path {
								i["mount_points/0"] = parts[2]
								i["filesystem"] = parts[3]
								break
							}
						}
					}
				}
				output, err = runCmd("df", path)
				if err == nil {
					scanner := bufio.NewScanner(bytes.NewReader([]byte(output)))
					scanner.Scan()
					scanner.Scan()
					parts := regexp.MustCompile(`\s+`).Split(scanner.Text(), -1)
					free, err := strconv.ParseInt(parts[3], 10, 64)
					if err == nil {
						i["free"] = strconv.FormatInt(free*1024, 10)
					}
				}
			}
			for k, v := range i {
				pi[fmt.Sprintf("data/volumes/%d/%s", part_idx, k)] = v
			}
			part_idx += 1
		}
	}
	return pi, nil
}
<commit_msg>Fix identification of blockdev block size<commit_after>package guestmetric

import (
	"bufio"
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
	xenstoreclient "xe-guest-utilities/xenstoreclient"
)

type Collector struct {
	Client xenstoreclient.XenStoreClient
	Ballon bool
	Debug  bool
}

func (c *Collector) CollectOS() (GuestMetric, error) {
	current := make(GuestMetric, 0)
	f, err := os.OpenFile("/var/cache/xe-linux-distribution", os.O_RDONLY, 0666)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, "=") {
			parts := strings.SplitN(line, "=", 2)
			k := strings.TrimSpace(parts[0])
			v := strings.TrimSpace(strings.Trim(strings.TrimSpace(parts[1]), "\""))
			current[k] = v
		}
	}
	return prefixKeys("data/", current), nil
}

func (c *Collector) CollectMisc() (GuestMetric, error) {
	current := make(GuestMetric, 0)
	if c.Ballon {
		current["control/feature-balloon"] = "1"
	} else {
		current["control/feature-balloon"] = "0"
	}
	current["attr/PVAddons/Installed"] = "1"
	current["attr/PVAddons/MajorVersion"] = "@PRODUCT_MAJOR_VERSION@"
	current["attr/PVAddons/MinorVersion"] = "@PRODUCT_MINOR_VERSION@"
	current["attr/PVAddons/MicroVersion"] = "@PRODUCT_MICRO_VERSION@"
	current["attr/PVAddons/BuildVersion"] = "@NUMERIC_BUILD_NUMBER@"

	return current, nil
}

func (c *Collector) CollectMemory() (GuestMetric, error) {
	current := make(GuestMetric, 0)
	f, err := os.Open("/proc/meminfo")
	if err != nil {
		return nil, err
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		parts := regexp.MustCompile(`\w+`).FindAllString(scanner.Text(), -1)
		switch parts[0] {
		case "MemTotal":
			current["meminfo_total"] = parts[1]
		case "MemAvailable":
			current["meminfo_free"] = parts[1]
		}
	}
	return prefixKeys("data/", current), nil
}

func enumNetworkAddresses(iface string) (GuestMetric, error) {
	const (
		IP_RE   string = `(\d{1,3}\.){3}\d{1,3}`
		IPV6_RE string = `[\da-f:]+[\da-f]`
	)

	var (
		IP_IPV4_ADDR_RE       = regexp.MustCompile(`inet\s*(` + IP_RE + `).*\se[a-zA-Z0-9]+[\s\n]`)
		IP_IPV6_ADDR_RE       = regexp.MustCompile(`inet6\s*(` + IPV6_RE + `)`)
		IFCONFIG_IPV4_ADDR_RE = regexp.MustCompile(`inet addr:\s*(` + IP_RE + `)`)
		IFCONFIG_IPV6_ADDR_RE = regexp.MustCompile(`inet6 addr:\s*(` + IPV6_RE + `)`)
	)

	d := make(GuestMetric, 0)

	var v4re, v6re *regexp.Regexp
	var out string
	var err error
	if out, err = runCmd("ip", "addr", "show", iface); err == nil {
		v4re = IP_IPV4_ADDR_RE
		v6re = IP_IPV6_ADDR_RE
	} else if out, err = runCmd("ifconfig", iface); err == nil {
		v4re = IFCONFIG_IPV4_ADDR_RE
		v6re = IFCONFIG_IPV6_ADDR_RE
	} else {
		return nil, fmt.Errorf("Cannot find ip/ifconfig command")
	}

	m := v4re.FindAllStringSubmatch(out, -1)
	if m != nil {
		for i, parts := range m {
			d[fmt.Sprintf("ipv4/%d", i)] = parts[1]
		}
	}
	m = v6re.FindAllStringSubmatch(out, -1)
	if m != nil {
		for i, parts := range m {
			d[fmt.Sprintf("ipv6/%d", i)] = parts[1]
		}
	}

	return d, nil
}

func getPlainVifId(path string) (string, error) {
	nodenamePath := fmt.Sprintf("%s/device/nodename", path)
	strLine, err := readSysfs(nodenamePath)
	if err != nil {
		return "", err
	}
	vifId := ""
	reNodename := regexp.MustCompile(`^device\/vif\/(\d+)$`)
	if matched := reNodename.FindStringSubmatch(strLine); matched != nil {
		vifId = matched[1]
	}
	if vifId == "" {
		return "", fmt.Errorf("Not found string like \"device/vif/[id]\" in file %s", nodenamePath)
	} else {
		return vifId, nil
	}
}

func (c *Collector) getSriovVifId(path string) (string, error) {
	sriovDevicePath := "xenserver/device/net-sriov-vf"
	macAddress, err := readSysfs(path + "/address")
	if err != nil {
		return "", err
	}
	subPaths, err := c.Client.List(sriovDevicePath)
	if err != nil {
		return "", err
	}
	for _, subPath := range subPaths {
		iterMac, err := c.Client.Read(fmt.Sprintf("%s/%s/mac", sriovDevicePath, subPath))
		if err != nil {
			continue
		}
		if iterMac == macAddress {
			return subPath, nil
		}
	}
	return "", fmt.Errorf("Cannot find a MAC address to map with %s", path)
}

// return vif_xenstore_prefix * vif_id * error where
// `vif_xenstore_prefix` could be either `attr/vif` for plain VIF or
// `xenserver/attr/net-sriov-vf` for SR-IOV VIF
func (c *Collector) getTargetXenstorePath(path string) (string, string, error) {
	plainVifPrefix := "attr/vif"
	sriovVifPrefix := "xenserver/attr/net-sriov-vf"
	// try to get `vif_id` from nodename interface, only a plain VIF have the nodename interface.
	vifId, err1 := getPlainVifId(path)
	if vifId != "" {
		return plainVifPrefix, vifId, nil
	}
	// not a plain VIF, it could possible be an SR-IOV VIF, try to get vif_id from MAC address mapping
	vifId, err2 := c.getSriovVifId(path)
	if vifId != "" {
		return sriovVifPrefix, vifId, nil
	}
	return "", "", fmt.Errorf("Failed to get VIF ID, errors: %s | %s", err1.Error(), err2.Error())
}

func (c *Collector) CollectNetworkAddr() (GuestMetric, error) {
	current := make(GuestMetric, 0)

	var paths []string
	vifNamePrefixList := [...]string{"eth", "eno", "ens", "emp", "enx", "enX"}
	for _, prefix := range vifNamePrefixList {
		prefixPaths, err := filepath.Glob(fmt.Sprintf("/sys/class/net/%s*", prefix))
		if err != nil {
			return nil, err
		}
		paths = append(paths, prefixPaths...)
	}
	for _, path := range paths {
		// a path is going to be like "/sys/class/net/eth0"
		prefix, vifId, err := c.getTargetXenstorePath(path)
		if err != nil {
			continue
		}
		iface := filepath.Base(path)
		if addrs, err := enumNetworkAddresses(iface); err == nil {
			for tag, addr := range addrs {
				current[fmt.Sprintf("%s/%s/%s", prefix, vifId, tag)] = addr
			}
		}
	}
	return current, nil
}

func readSysfs(filename string) (string, error) {
	f, err := os.OpenFile(filename, os.O_RDONLY, 0666)
	if err != nil {
		return "", err
	}
	defer f.Close()
	scanner := bufio.NewScanner(f)
	scanner.Scan()
	return scanner.Text(), nil
}

func (c *Collector) CollectDisk() (GuestMetric, error) {
	pi := make(GuestMetric, 0)

	disks := make([]string, 0)
	paths, err := filepath.Glob("/sys/block/*/device")
	if err != nil {
		return nil, err
	}
	for _, path := range paths {
		disk := filepath.Base(strings.TrimSuffix(filepath.Dir(path), "/"))
		disks = append(disks, disk)
	}

	var sortedDisks sort.StringSlice = disks
	sortedDisks.Sort()

	part_idx := 0
	for _, disk := range sortedDisks[:] {
		paths, err = filepath.Glob(fmt.Sprintf("/dev/%s?*", disk))
		if err != nil {
			return nil, err
		}
		for _, path := range paths {
			p := filepath.Base(path)
			line, err := readSysfs(fmt.Sprintf("/sys/block/%s/%s/size", disk, p))
			if err != nil {
				return nil, err
			}
			size, err := strconv.ParseInt(line, 10, 64)
			if err != nil {
				return nil, err
			}
			blocksize := 512
			if bs, err := readSysfs(fmt.Sprintf("/sys/block/%s/queue/physical_block_size", disk)); err == nil {
				if bs1, err := strconv.Atoi(bs); err == nil {
					blocksize = bs1
				}
			}
			real_dev := ""
			if c.Client != nil {
				nodename, err := readSysfs(fmt.Sprintf("/sys/block/%s/device/nodename", disk))
				if err == nil {
					backend, err := c.Client.Read(fmt.Sprintf("%s/backend", nodename))
					if err != nil {
						return nil, err
					}
					real_dev, err = c.Client.Read(fmt.Sprintf("%s/dev", backend))
					if err != nil {
						return nil, err
					}
				}
			}
			name := path
			blkid, err := runCmd("blkid", "-s", "UUID", path)
			if err != nil {
				// ignore blkid errors
				blkid = ""
			}
			if strings.Contains(blkid, "=") {
				parts := strings.SplitN(strings.TrimSpace(blkid), "=", 2)
				name = fmt.Sprintf("%s(%s)", name, strings.Trim(parts[1], "\""))
			}
			i := map[string]string{
				"extents/0": real_dev,
				"name":      name,
				"size":      strconv.FormatInt(size*int64(blocksize), 10),
			}
			output, err := runCmd("pvs", "--noheadings", "--units", "b", path)
			if err == nil && output != "" {
				parts := regexp.MustCompile(`\s+`).Split(output, -1)[1:]
				i["free"] = strings.TrimSpace(parts[5])[:len(parts[5])-1]
				i["filesystem"] = strings.TrimSpace(parts[2])
				i["mount_points/0"] = "[LVM]"
			} else {
				output, err = runCmd("mount")
				if err == nil {
					m := regexp.MustCompile(`(?m)^(\S+) on (\S+) type (\S+)`).FindAllStringSubmatch(output, -1)
					if m != nil {
						for _, parts := range m {
							if parts[1] == path {
								i["mount_points/0"] = parts[2]
								i["filesystem"] = parts[3]
								break
							}
						}
					}
				}
				output, err = runCmd("df", path)
				if err == nil {
					scanner := bufio.NewScanner(bytes.NewReader([]byte(output)))
					scanner.Scan()
					scanner.Scan()
					parts := regexp.MustCompile(`\s+`).Split(scanner.Text(), -1)
					free, err := strconv.ParseInt(parts[3], 10, 64)
					if err == nil {
						i["free"] = strconv.FormatInt(free*1024, 10)
					}
				}
			}
			for k, v := range i {
				pi[fmt.Sprintf("data/volumes/%d/%s", part_idx, k)] = v
			}
			part_idx += 1
		}
	}
	return pi, nil
}
<|endoftext|>