<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"sort"
	"strings"

	"github.com/spf13/cobra"
)

const DelveMainPackagePath = "github.com/derekparker/delve/cmd/dlv"

var Verbose bool
var TestSet, TestRegex, TestBackend string

func NewMakeCommands() *cobra.Command {
	RootCommand := &cobra.Command{
		Use:   "make.go",
		Short: "make script for delve.",
	}

	RootCommand.AddCommand(&cobra.Command{
		Use:   "check-cert",
		Short: "Check certificate for macOS.",
		Run:   checkCertCmd,
	})

	RootCommand.AddCommand(&cobra.Command{
		Use:   "build",
		Short: "Build delve",
		Run: func(cmd *cobra.Command, args []string) {
			checkCertCmd(nil, nil)
			execute("go", "build", buildFlags(), DelveMainPackagePath)
			if runtime.GOOS == "darwin" && os.Getenv("CERT") != "" {
				codesign("./dlv")
			}
		},
	})

	RootCommand.AddCommand(&cobra.Command{
		Use:   "install",
		Short: "Installs delve",
		Run: func(cmd *cobra.Command, args []string) {
			checkCertCmd(nil, nil)
			execute("go", "install", buildFlags(), DelveMainPackagePath)
			if runtime.GOOS == "darwin" {
				codesign(installedExecutablePath())
			}
		},
	})

	test := &cobra.Command{
		Use:   "test",
		Short: "Tests delve",
		Long: `Tests delve.

Use the flags -s, -r and -b to specify which tests to run. Specifying nothing is equivalent to:

	go run scripts/make.go test -s all -b default
	go run scripts/make.go test -s basic -b lldb
	go run scripts/make.go test -s basic -b rr

with lldb and rr tests only run if the relevant programs are installed.`,
		Run: testCmd,
	}
	test.PersistentFlags().BoolVarP(&Verbose, "verbose", "v", false, "Verbose tests")
	test.PersistentFlags().StringVarP(&TestSet, "test-set", "s", "", `Select the set of tests to run, one of either:
	all		tests all packages
	basic		tests proc, integration and terminal
	integration 	tests github.com/derekparker/delve/service/test
	package-name	test the specified package only
`)
	test.PersistentFlags().StringVarP(&TestRegex, "test-run", "r", "", `Only runs the tests matching the specified regex. This option can only be specified if testset is a single package`)
	test.PersistentFlags().StringVarP(&TestBackend, "test-backend", "b", "", `Runs tests for the specified backend only, one of either:
	default		the default backend
	lldb		lldb backend
	rr		rr backend

This option can only be specified if testset is basic or a single package.`)

	RootCommand.AddCommand(test)

	RootCommand.AddCommand(&cobra.Command{
		Use:   "vendor",
		Short: "vendors dependencies",
		Run: func(cmd *cobra.Command, args []string) {
			execute("glide", "up", "-v")
			execute("glide-vc", "--use-lock-file", "--no-tests", "--only-code")
		},
	})

	return RootCommand
}

func checkCertCmd(cmd *cobra.Command, args []string) {
	// If we're on OSX make sure the proper CERT env var is set.
	if os.Getenv("TRAVIS") == "true" || runtime.GOOS != "darwin" || os.Getenv("CERT") != "" {
		return
	}

	x := exec.Command("scripts/gencert.sh")
	err := x.Run()
	if x.ProcessState != nil && !x.ProcessState.Success() {
		fmt.Printf("An error occurred when generating and installing a new certificate\n")
		os.Exit(1)
	}
	if err != nil {
		log.Fatal(err)
	}
	os.Setenv("CERT", "dlv-cert")
}

func strflatten(v []interface{}) []string {
	r := []string{}
	for _, s := range v {
		switch s := s.(type) {
		case []string:
			r = append(r, s...)
		case string:
			if s != "" {
				r = append(r, s)
			}
		}
	}
	return r
}

func executeq(cmd string, args ...interface{}) {
	x := exec.Command(cmd, strflatten(args)...)
	x.Stdout = os.Stdout
	x.Stderr = os.Stderr
	x.Env = os.Environ()
	err := x.Run()
	if x.ProcessState != nil && !x.ProcessState.Success() {
		os.Exit(1)
	}
	if err != nil {
		log.Fatal(err)
	}
}

func execute(cmd string, args ...interface{}) {
	fmt.Printf("%s %s\n", cmd, strings.Join(quotemaybe(strflatten(args)), " "))
	executeq(cmd, args...)
}

func quotemaybe(args []string) []string {
	for i := range args {
		if strings.Index(args[i], " ") >= 0 {
			args[i] = fmt.Sprintf("%q", args[i])
		}
	}
	return args
}

func getoutput(cmd string, args ...interface{}) string {
	x := exec.Command(cmd, strflatten(args)...)
	x.Env = os.Environ()
	out, err := x.CombinedOutput()
	if err != nil {
		log.Fatal(err)
	}
	if !x.ProcessState.Success() {
		os.Exit(1)
	}
	return string(out)
}

func codesign(path string) {
	execute("codesign", "-s", os.Getenv("CERT"), path)
}

func installedExecutablePath() string {
	if gobin := os.Getenv("GOBIN"); gobin != "" {
		return filepath.Join(gobin, "dlv")
	}
	gopath := strings.Split(getoutput("go", "env", "GOPATH"), ":")
	return filepath.Join(gopath[0], "dlv")
}

func buildFlags() []string {
	buildSHA, err := exec.Command("git", "rev-parse", "HEAD").CombinedOutput()
	if err != nil {
		log.Fatal(err)
	}
	ldFlags := "-X main.Build=" + strings.TrimSpace(string(buildSHA))
	if runtime.GOOS == "darwin" {
		ldFlags = "-s " + ldFlags
	}
	return []string{fmt.Sprintf("-ldflags=%s", ldFlags)}
}

func testFlags() []string {
	wd, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	testFlags := []string{"-count", "1", "-p", "1"}
	if Verbose {
		testFlags = append(testFlags, "-v")
	}
	if runtime.GOOS == "darwin" {
		testFlags = append(testFlags, "-exec="+wd+"/scripts/testsign")
	}
	return testFlags
}

func testCmd(cmd *cobra.Command, args []string) {
	checkCertCmd(nil, nil)

	if os.Getenv("TRAVIS") == "true" && runtime.GOOS == "darwin" {
		os.Setenv("PROCTEST", "lldb")
		executeq("sudo", "-E", "go", "test", testFlags(), allPackages())
		return
	}

	if TestSet == "" && TestBackend == "" {
		if TestRegex != "" {
			fmt.Printf("Can not use --test-run without --test-set\n")
			os.Exit(1)
		}

		fmt.Println("Testing default backend")
		testCmdIntl("all", "", "default")
		if inpath("lldb-server") {
			fmt.Println("\nTesting LLDB backend")
			testCmdIntl("basic", "", "lldb")
		}
		if inpath("rr") {
			fmt.Println("\nTesting RR backend")
			testCmdIntl("basic", "", "rr")
		}
		return
	}

	if TestSet == "" {
		TestSet = "all"
	}

	if TestBackend == "" {
		TestBackend = "default"
	}

	testCmdIntl(TestSet, TestRegex, TestBackend)
}

func testCmdIntl(testSet, testRegex, testBackend string) {
	testPackages := testSetToPackages(testSet)
	if len(testPackages) == 0 {
		fmt.Printf("Unknown test set %q\n", testSet)
		os.Exit(1)
	}

	if testRegex != "" && len(testPackages) != 1 {
		fmt.Printf("Can not use test-run with test set %q\n", testSet)
		os.Exit(1)
	}

	backendFlag := ""
	if testBackend != "" && testBackend != "default" {
		if testSet != "basic" && len(testPackages) != 1 {
			fmt.Printf("Can not use test-backend with test set %q\n", testSet)
			os.Exit(1)
		}
		backendFlag = "-backend=" + testBackend
	}

	if len(testPackages) > 3 {
		execute("go", "test", testFlags(), buildFlags(), testPackages, backendFlag)
	} else if testRegex != "" {
		execute("go", "test", testFlags(), buildFlags(), testPackages, "-run="+testRegex, backendFlag)
	} else {
		execute("go", "test", testFlags(), buildFlags(), testPackages, backendFlag)
	}
}

func testSetToPackages(testSet string) []string {
	switch testSet {
	case "", "all":
		return allPackages()

	case "basic":
		return []string{"github.com/derekparker/delve/pkg/proc", "github.com/derekparker/delve/service/test", "github.com/derekparker/delve/pkg/terminal"}

	case "integration":
		return []string{"github.com/derekparker/delve/service/test"}

	default:
		for _, pkg := range allPackages() {
			if pkg == testSet || strings.HasSuffix(pkg, "/"+testSet) {
				return []string{pkg}
			}
		}
		return nil
	}
}

func inpath(exe string) bool {
	path, _ := exec.LookPath(exe)
	return path != ""
}

func allPackages() []string {
	r := []string{}
	for _, dir := range strings.Split(getoutput("go", "list", "./..."), "\n") {
		dir = strings.TrimSpace(dir)
		if dir == "" || strings.Contains(dir, "/vendor/") || strings.Contains(dir, "/scripts") {
			continue
		}
		r = append(r, dir)
	}
	sort.Strings(r)
	return r
}

func main() {
	NewMakeCommands().Execute()
}
<commit_msg>Fix scripts/make.go for darwin<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"sort"
	"strings"

	"github.com/spf13/cobra"
)

const DelveMainPackagePath = "github.com/derekparker/delve/cmd/dlv"

var Verbose bool
var TestSet, TestRegex, TestBackend string

func NewMakeCommands() *cobra.Command {
	RootCommand := &cobra.Command{
		Use:   "make.go",
		Short: "make script for delve.",
	}

	RootCommand.AddCommand(&cobra.Command{
		Use:   "check-cert",
		Short: "Check certificate for macOS.",
		Run:   checkCertCmd,
	})

	RootCommand.AddCommand(&cobra.Command{
		Use:   "build",
		Short: "Build delve",
		Run: func(cmd *cobra.Command, args []string) {
			checkCertCmd(nil, nil)
			execute("go", "build", buildFlags(), DelveMainPackagePath)
			if runtime.GOOS == "darwin" && os.Getenv("CERT") != "" {
				codesign("./dlv")
			}
		},
	})

	RootCommand.AddCommand(&cobra.Command{
		Use:   "install",
		Short: "Installs delve",
		Run: func(cmd *cobra.Command, args []string) {
			checkCertCmd(nil, nil)
			execute("go", "install", buildFlags(), DelveMainPackagePath)
			if runtime.GOOS == "darwin" {
				codesign(installedExecutablePath())
			}
		},
	})

	test := &cobra.Command{
		Use:   "test",
		Short: "Tests delve",
		Long: `Tests delve.

Use the flags -s, -r and -b to specify which tests to run. Specifying nothing is equivalent to:

	go run scripts/make.go test -s all -b default
	go run scripts/make.go test -s basic -b lldb
	go run scripts/make.go test -s basic -b rr

with lldb and rr tests only run if the relevant programs are installed.`,
		Run: testCmd,
	}
	test.PersistentFlags().BoolVarP(&Verbose, "verbose", "v", false, "Verbose tests")
	test.PersistentFlags().StringVarP(&TestSet, "test-set", "s", "", `Select the set of tests to run, one of either:
	all		tests all packages
	basic		tests proc, integration and terminal
	integration 	tests github.com/derekparker/delve/service/test
	package-name	test the specified package only
`)
	test.PersistentFlags().StringVarP(&TestRegex, "test-run", "r", "", `Only runs the tests matching the specified regex. This option can only be specified if testset is a single package`)
	test.PersistentFlags().StringVarP(&TestBackend, "test-backend", "b", "", `Runs tests for the specified backend only, one of either:
	default		the default backend
	lldb		lldb backend
	rr		rr backend

This option can only be specified if testset is basic or a single package.`)

	RootCommand.AddCommand(test)

	RootCommand.AddCommand(&cobra.Command{
		Use:   "vendor",
		Short: "vendors dependencies",
		Run: func(cmd *cobra.Command, args []string) {
			execute("glide", "up", "-v")
			execute("glide-vc", "--use-lock-file", "--no-tests", "--only-code")
		},
	})

	return RootCommand
}

func checkCertCmd(cmd *cobra.Command, args []string) {
	// If we're on OSX make sure the proper CERT env var is set.
	if os.Getenv("TRAVIS") == "true" || runtime.GOOS != "darwin" || os.Getenv("CERT") != "" {
		return
	}

	x := exec.Command("scripts/gencert.sh")
	err := x.Run()
	if x.ProcessState != nil && !x.ProcessState.Success() {
		fmt.Printf("An error occurred when generating and installing a new certificate\n")
		os.Exit(1)
	}
	if err != nil {
		log.Fatal(err)
	}
	os.Setenv("CERT", "dlv-cert")
}

func strflatten(v []interface{}) []string {
	r := []string{}
	for _, s := range v {
		switch s := s.(type) {
		case []string:
			r = append(r, s...)
		case string:
			if s != "" {
				r = append(r, s)
			}
		}
	}
	return r
}

func executeq(cmd string, args ...interface{}) {
	x := exec.Command(cmd, strflatten(args)...)
	x.Stdout = os.Stdout
	x.Stderr = os.Stderr
	x.Env = os.Environ()
	err := x.Run()
	if x.ProcessState != nil && !x.ProcessState.Success() {
		os.Exit(1)
	}
	if err != nil {
		log.Fatal(err)
	}
}

func execute(cmd string, args ...interface{}) {
	fmt.Printf("%s %s\n", cmd, strings.Join(quotemaybe(strflatten(args)), " "))
	executeq(cmd, args...)
}

func quotemaybe(args []string) []string {
	for i := range args {
		if strings.Index(args[i], " ") >= 0 {
			args[i] = fmt.Sprintf("%q", args[i])
		}
	}
	return args
}

func getoutput(cmd string, args ...interface{}) string {
	x := exec.Command(cmd, strflatten(args)...)
	x.Env = os.Environ()
	out, err := x.CombinedOutput()
	if err != nil {
		log.Fatal(err)
	}
	if !x.ProcessState.Success() {
		os.Exit(1)
	}
	return string(out)
}

func codesign(path string) {
	execute("codesign", "-s", os.Getenv("CERT"), path)
}

func installedExecutablePath() string {
	if gobin := os.Getenv("GOBIN"); gobin != "" {
		return filepath.Join(gobin, "dlv")
	}
	gopath := strings.Split(getoutput("go", "env", "GOPATH"), ":")
	return filepath.Join(strings.TrimSpace(gopath[0]), "bin", "dlv")
}

func buildFlags() []string {
	buildSHA, err := exec.Command("git", "rev-parse", "HEAD").CombinedOutput()
	if err != nil {
		log.Fatal(err)
	}
	ldFlags := "-X main.Build=" + strings.TrimSpace(string(buildSHA))
	if runtime.GOOS == "darwin" {
		ldFlags = "-s " + ldFlags
	}
	return []string{fmt.Sprintf("-ldflags=%s", ldFlags)}
}

func testFlags() []string {
	wd, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	testFlags := []string{"-count", "1", "-p", "1"}
	if Verbose {
		testFlags = append(testFlags, "-v")
	}
	if runtime.GOOS == "darwin" {
		testFlags = append(testFlags, "-exec="+wd+"/scripts/testsign")
	}
	return testFlags
}

func testCmd(cmd *cobra.Command, args []string) {
	checkCertCmd(nil, nil)

	if os.Getenv("TRAVIS") == "true" && runtime.GOOS == "darwin" {
		os.Setenv("PROCTEST", "lldb")
		executeq("sudo", "-E", "go", "test", testFlags(), allPackages())
		return
	}

	if TestSet == "" && TestBackend == "" {
		if TestRegex != "" {
			fmt.Printf("Can not use --test-run without --test-set\n")
			os.Exit(1)
		}

		fmt.Println("Testing default backend")
		testCmdIntl("all", "", "default")
		if inpath("lldb-server") {
			fmt.Println("\nTesting LLDB backend")
			testCmdIntl("basic", "", "lldb")
		}
		if inpath("rr") {
			fmt.Println("\nTesting RR backend")
			testCmdIntl("basic", "", "rr")
		}
		return
	}

	if TestSet == "" {
		TestSet = "all"
	}

	if TestBackend == "" {
		TestBackend = "default"
	}

	testCmdIntl(TestSet, TestRegex, TestBackend)
}

func testCmdIntl(testSet, testRegex, testBackend string) {
	testPackages := testSetToPackages(testSet)
	if len(testPackages) == 0 {
		fmt.Printf("Unknown test set %q\n", testSet)
		os.Exit(1)
	}

	if testRegex != "" && len(testPackages) != 1 {
		fmt.Printf("Can not use test-run with test set %q\n", testSet)
		os.Exit(1)
	}

	backendFlag := ""
	if testBackend != "" && testBackend != "default" {
		if testSet != "basic" && len(testPackages) != 1 {
			fmt.Printf("Can not use test-backend with test set %q\n", testSet)
			os.Exit(1)
		}
		backendFlag = "-backend=" + testBackend
	}

	if len(testPackages) > 3 {
		execute("go", "test", testFlags(), buildFlags(), testPackages, backendFlag)
	} else if testRegex != "" {
		execute("go", "test", testFlags(), buildFlags(), testPackages, "-run="+testRegex, backendFlag)
	} else {
		execute("go", "test", testFlags(), buildFlags(), testPackages, backendFlag)
	}
}

func testSetToPackages(testSet string) []string {
	switch testSet {
	case "", "all":
		return allPackages()

	case "basic":
		return []string{"github.com/derekparker/delve/pkg/proc", "github.com/derekparker/delve/service/test", "github.com/derekparker/delve/pkg/terminal"}

	case "integration":
		return []string{"github.com/derekparker/delve/service/test"}

	default:
		for _, pkg := range allPackages() {
			if pkg == testSet || strings.HasSuffix(pkg, "/"+testSet) {
				return []string{pkg}
			}
		}
		return nil
	}
}

func inpath(exe string) bool {
	path, _ := exec.LookPath(exe)
	return path != ""
}

func allPackages() []string {
	r := []string{}
	for _, dir := range strings.Split(getoutput("go", "list", "./..."), "\n") {
		dir = strings.TrimSpace(dir)
		if dir == "" || strings.Contains(dir, "/vendor/") || strings.Contains(dir, "/scripts") {
			continue
		}
		r = append(r, dir)
	}
	sort.Strings(r)
	return r
}

func main() {
	NewMakeCommands().Execute()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"math"
	"sort"
)

func sieveOfEratosthenes(upperBound int) {
	// Use hashmap to prevent duplicates
	primes := map[int]int{}

	upperBoundSquareRootFl := math.Sqrt(float64(upperBound))
	upperBoundSquareRoot := int(upperBoundSquareRootFl)
	isComposite := []bool{}
	for i := 0; i < upperBound; i++ {
		isComposite = append(isComposite, false)
	}
	for m := 2; m <= upperBoundSquareRoot; m++ {
		if !isComposite[m] {
			primes[m] = m
			for k := m * m; k <= int(upperBound); k += m {
				if k <= len(isComposite)-1 {
					isComposite[k] = true
				}
			}
		}
	}
	for m := upperBoundSquareRoot; m <= int(upperBound); m++ {
		if m <= len(isComposite)-1 {
			if !isComposite[m] {
				primes[m] = m
			}
		}
	}

	// Get all keys out of hashmap
	keys := []int{}
	for key := range primes {
		keys = append(keys, key)
	}

	// Sort Keys
	sort.Ints(keys)

	// Print array of all prime numbers up to upperBound
	fmt.Println(keys)
}

func main() {
	sieveOfEratosthenes(25)
}
<commit_msg>Increased upper bound to match visualization<commit_after>package main

import (
	"fmt"
	"math"
	"sort"
)

func sieveOfEratosthenes(upperBound int) {
	// Use hashmap to prevent duplicates
	primes := map[int]int{}

	upperBoundSquareRootFl := math.Sqrt(float64(upperBound))
	upperBoundSquareRoot := int(upperBoundSquareRootFl)
	isComposite := []bool{}
	for i := 0; i < upperBound; i++ {
		isComposite = append(isComposite, false)
	}
	for m := 2; m <= upperBoundSquareRoot; m++ {
		if !isComposite[m] {
			primes[m] = m
			for k := m * m; k <= int(upperBound); k += m {
				if k <= len(isComposite)-1 {
					isComposite[k] = true
				}
			}
		}
	}
	for m := upperBoundSquareRoot; m <= int(upperBound); m++ {
		if m <= len(isComposite)-1 {
			if !isComposite[m] {
				primes[m] = m
			}
		}
	}

	// Get all keys out of hashmap
	keys := []int{}
	for key := range primes {
		keys = append(keys, key)
	}

	// Sort Keys
	sort.Ints(keys)

	// Print array of all prime numbers up to upperBound
	fmt.Println(keys)
}

func main() {
	sieveOfEratosthenes(120)
}
<|endoftext|>


<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/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.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
			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.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",
		},
	}
)

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":
	default:
		fmt.Printf("Invalid option to -convert-netconf: '%s'. Supported options: 'debian, digitalocean'\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-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
	if ud, err := initialize.ParseUserData(userdata); err != nil {
		fmt.Printf("Failed to parse user-data: %v\nContinuing...\n", err)
		failure = true
	} else {
		switch t := ud.(type) {
		case *config.CloudConfig:
			ccu = t
		case *config.Script:
			script = t
		}
	}

	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)
		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.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.4.0+git<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/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.0+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
			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.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",
		},
	}
)

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":
	default:
		fmt.Printf("Invalid option to -convert-netconf: '%s'. Supported options: 'debian, digitalocean'\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-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
	if ud, err := initialize.ParseUserData(userdata); err != nil {
		fmt.Printf("Failed to parse user-data: %v\nContinuing...\n", err)
		failure = true
	} else {
		switch t := ud.(type) {
		case *config.CloudConfig:
			ccu = t
		case *config.Script:
			script = t
		}
	}

	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)
		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.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 main

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

	"github.com/coreos/coreos-cloudinit/datasource"
	"github.com/coreos/coreos-cloudinit/initialize"
	"github.com/coreos/coreos-cloudinit/pkg"
	"github.com/coreos/coreos-cloudinit/system"
)

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

var (
	printVersion  bool
	ignoreFailure bool
	sources       struct {
		file            string
		configDrive     string
		metadataService bool
		url             string
		procCmdLine     bool
	}
	convertNetconf string
	workspace      string
	sshKeyName     string
)

func init() {
	flag.BoolVar(&printVersion, "version", false, "Print the version and exit")
	flag.BoolVar(&ignoreFailure, "ignore-failure", false, "Exits with 0 status in the event of malformed input from user-data")
	flag.StringVar(&sources.file, "from-file", "", "Read user-data from provided file")
	flag.StringVar(&sources.configDrive, "from-configdrive", "", "Read data from provided cloud-drive directory")
	flag.BoolVar(&sources.metadataService, "from-metadata-service", false, "Download data from metadata service")
	flag.StringVar(&sources.url, "from-url", "", "Download user-data from provided url")
	flag.BoolVar(&sources.procCmdLine, "from-proc-cmdline", false, fmt.Sprintf("Parse %s for '%s=<url>', using the cloud-config served by an HTTP GET to <url>", datasource.ProcCmdlineLocation, datasource.ProcCmdlineCloudConfigFlag))
	flag.StringVar(&convertNetconf, "convert-netconf", "", "Read the network config provided in cloud-drive and translate it from the specified format into networkd unit files (requires the -from-configdrive flag)")
	flag.StringVar(&workspace, "workspace", "/var/lib/coreos-cloudinit", "Base directory coreos-cloudinit should use to store data")
	flag.StringVar(&sshKeyName, "ssh-key-name", initialize.DefaultSSHKeyName, "Add SSH keys to the system with the given name")
}

func main() {
	flag.Parse()

	die := func() {
		if ignoreFailure {
			os.Exit(0)
		}
		os.Exit(1)
	}

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

	if convertNetconf != "" && sources.configDrive == "" {
		fmt.Println("-convert-netconf flag requires -from-configdrive")
		os.Exit(1)
	}

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

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

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

	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\n", err)
		die()
	}

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

	// Extract IPv4 addresses from metadata if possible
	var subs map[string]string
	if len(metadataBytes) > 0 {
		subs, err = initialize.ExtractIPsFromMetadata(metadataBytes)
		if err != nil {
			fmt.Printf("Failed extracting IPs from meta-data: %v\n", err)
			die()
		}
	}

	env := initialize.NewEnvironment("/", ds.ConfigRoot(), workspace, convertNetconf, sshKeyName, subs)

	var ccm, ccu *initialize.CloudConfig
	var script *system.Script
	if ccm, err = initialize.ParseMetaData(string(metadataBytes)); err != nil {
		fmt.Printf("Failed to parse meta-data: %v\n", err)
		die()
	}
	if ud, err := initialize.ParseUserData(string(userdataBytes)); err != nil {
		fmt.Printf("Failed to parse user-data: %v\n", err)
		die()
	} else {
		switch t := ud.(type) {
		case *initialize.CloudConfig:
			ccu = t
		case system.Script:
			script = &t
		}
	}

	var cc *initialize.CloudConfig
	if ccm != nil && ccu != nil {
		fmt.Println("Merging cloud-config from meta-data and user-data")
		merged := mergeCloudConfig(*ccu, *ccm)
		cc = &merged
	} else if ccm != nil && ccu == nil {
		fmt.Println("Processing cloud-config from meta-data")
		cc = ccm
	} else if ccm == nil && ccu != nil {
		fmt.Println("Processing cloud-config from user-data")
		cc = ccu
	} else {
		fmt.Println("No cloud-config data to handle.")
	}

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

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

// mergeCloudConfig merges certain options from mdcc (a CloudConfig derived from
// meta-data) onto udcc (a CloudConfig derived from user-data), if they are
// not already set on udcc (i.e. user-data always takes precedence)
// NB: This needs to be kept in sync with ParseMetadata so that it tracks all
// elements of a CloudConfig which that function can populate.
func mergeCloudConfig(mdcc, udcc initialize.CloudConfig) (cc initialize.CloudConfig) {
	if mdcc.Hostname != "" {
		if udcc.Hostname != "" {
			fmt.Printf("Warning: user-data hostname (%s) overrides metadata hostname (%s)", udcc.Hostname, mdcc.Hostname)
		} else {
			udcc.Hostname = mdcc.Hostname
		}

	}
	for _, key := range mdcc.SSHAuthorizedKeys {
		udcc.SSHAuthorizedKeys = append(udcc.SSHAuthorizedKeys, key)
	}
	if mdcc.NetworkConfigPath != "" {
		if udcc.NetworkConfigPath != "" {
			fmt.Printf("Warning: user-data NetworkConfigPath %s overrides metadata NetworkConfigPath %s", udcc.NetworkConfigPath, mdcc.NetworkConfigPath)
		} else {
			udcc.NetworkConfigPath = mdcc.NetworkConfigPath
		}
	}
	return udcc
}

// 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 sources.file != "" {
		dss = append(dss, datasource.NewLocalFile(sources.file))
	}
	if sources.url != "" {
		dss = append(dss, datasource.NewRemoteFile(sources.url))
	}
	if sources.configDrive != "" {
		dss = append(dss, datasource.NewConfigDrive(sources.configDrive))
	}
	if sources.metadataService {
		dss = append(dss, datasource.NewMetadataService())
	}
	if sources.procCmdLine {
		dss = append(dss, datasource.NewProcCmdline())
	}
	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.Tick(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.Tick(datasourceTimeout):
	}

	close(stop)
	return s
}

// TODO(jonboulle): this should probably be refactored and moved into a different module
func runScript(script system.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 0.8.2+git<commit_after>package main

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

	"github.com/coreos/coreos-cloudinit/datasource"
	"github.com/coreos/coreos-cloudinit/initialize"
	"github.com/coreos/coreos-cloudinit/pkg"
	"github.com/coreos/coreos-cloudinit/system"
)

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

var (
	printVersion  bool
	ignoreFailure bool
	sources       struct {
		file            string
		configDrive     string
		metadataService bool
		url             string
		procCmdLine     bool
	}
	convertNetconf string
	workspace      string
	sshKeyName     string
)

func init() {
	flag.BoolVar(&printVersion, "version", false, "Print the version and exit")
	flag.BoolVar(&ignoreFailure, "ignore-failure", false, "Exits with 0 status in the event of malformed input from user-data")
	flag.StringVar(&sources.file, "from-file", "", "Read user-data from provided file")
	flag.StringVar(&sources.configDrive, "from-configdrive", "", "Read data from provided cloud-drive directory")
	flag.BoolVar(&sources.metadataService, "from-metadata-service", false, "Download data from metadata service")
	flag.StringVar(&sources.url, "from-url", "", "Download user-data from provided url")
	flag.BoolVar(&sources.procCmdLine, "from-proc-cmdline", false, fmt.Sprintf("Parse %s for '%s=<url>', using the cloud-config served by an HTTP GET to <url>", datasource.ProcCmdlineLocation, datasource.ProcCmdlineCloudConfigFlag))
	flag.StringVar(&convertNetconf, "convert-netconf", "", "Read the network config provided in cloud-drive and translate it from the specified format into networkd unit files (requires the -from-configdrive flag)")
	flag.StringVar(&workspace, "workspace", "/var/lib/coreos-cloudinit", "Base directory coreos-cloudinit should use to store data")
	flag.StringVar(&sshKeyName, "ssh-key-name", initialize.DefaultSSHKeyName, "Add SSH keys to the system with the given name")
}

func main() {
	flag.Parse()

	die := func() {
		if ignoreFailure {
			os.Exit(0)
		}
		os.Exit(1)
	}

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

	if convertNetconf != "" && sources.configDrive == "" {
		fmt.Println("-convert-netconf flag requires -from-configdrive")
		os.Exit(1)
	}

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

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

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

	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\n", err)
		die()
	}

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

	// Extract IPv4 addresses from metadata if possible
	var subs map[string]string
	if len(metadataBytes) > 0 {
		subs, err = initialize.ExtractIPsFromMetadata(metadataBytes)
		if err != nil {
			fmt.Printf("Failed extracting IPs from meta-data: %v\n", err)
			die()
		}
	}

	env := initialize.NewEnvironment("/", ds.ConfigRoot(), workspace, convertNetconf, sshKeyName, subs)

	var ccm, ccu *initialize.CloudConfig
	var script *system.Script
	if ccm, err = initialize.ParseMetaData(string(metadataBytes)); err != nil {
		fmt.Printf("Failed to parse meta-data: %v\n", err)
		die()
	}
	if ud, err := initialize.ParseUserData(string(userdataBytes)); err != nil {
		fmt.Printf("Failed to parse user-data: %v\n", err)
		die()
	} else {
		switch t := ud.(type) {
		case *initialize.CloudConfig:
			ccu = t
		case system.Script:
			script = &t
		}
	}

	var cc *initialize.CloudConfig
	if ccm != nil && ccu != nil {
		fmt.Println("Merging cloud-config from meta-data and user-data")
		merged := mergeCloudConfig(*ccu, *ccm)
		cc = &merged
	} else if ccm != nil && ccu == nil {
		fmt.Println("Processing cloud-config from meta-data")
		cc = ccm
	} else if ccm == nil && ccu != nil {
		fmt.Println("Processing cloud-config from user-data")
		cc = ccu
	} else {
		fmt.Println("No cloud-config data to handle.")
	}

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

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

// mergeCloudConfig merges certain options from mdcc (a CloudConfig derived from
// meta-data) onto udcc (a CloudConfig derived from user-data), if they are
// not already set on udcc (i.e. user-data always takes precedence)
// NB: This needs to be kept in sync with ParseMetadata so that it tracks all
// elements of a CloudConfig which that function can populate.
func mergeCloudConfig(mdcc, udcc initialize.CloudConfig) (cc initialize.CloudConfig) {
	if mdcc.Hostname != "" {
		if udcc.Hostname != "" {
			fmt.Printf("Warning: user-data hostname (%s) overrides metadata hostname (%s)", udcc.Hostname, mdcc.Hostname)
		} else {
			udcc.Hostname = mdcc.Hostname
		}

	}
	for _, key := range mdcc.SSHAuthorizedKeys {
		udcc.SSHAuthorizedKeys = append(udcc.SSHAuthorizedKeys, key)
	}
	if mdcc.NetworkConfigPath != "" {
		if udcc.NetworkConfigPath != "" {
			fmt.Printf("Warning: user-data NetworkConfigPath %s overrides metadata NetworkConfigPath %s", udcc.NetworkConfigPath, mdcc.NetworkConfigPath)
		} else {
			udcc.NetworkConfigPath = mdcc.NetworkConfigPath
		}
	}
	return udcc
}

// 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 sources.file != "" {
		dss = append(dss, datasource.NewLocalFile(sources.file))
	}
	if sources.url != "" {
		dss = append(dss, datasource.NewRemoteFile(sources.url))
	}
	if sources.configDrive != "" {
		dss = append(dss, datasource.NewConfigDrive(sources.configDrive))
	}
	if sources.metadataService {
		dss = append(dss, datasource.NewMetadataService())
	}
	if sources.procCmdLine {
		dss = append(dss, datasource.NewProcCmdline())
	}
	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.Tick(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.Tick(datasourceTimeout):
	}

	close(stop)
	return s
}

// TODO(jonboulle): this should probably be refactored and moved into a different module
func runScript(script system.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 main

import (
	"github.com/AutoRoute/l2"
	"github.com/AutoRoute/node"

	"code.google.com/p/tuntap"

	"flag"
	"fmt"
	"log"
	"net"
	"runtime"
	"strings"
)

var listen = flag.String("listen", "127.0.0.1:34321",
	"The address to listen to incoming connections on")
var nolisten = flag.Bool("nolisten", false, "Disables listening")
var connect = flag.String("connect", "",
	"Comma separated list of addresses to connect to")
var autodiscover = flag.Bool("auto", false,
	"Whether we should try and find neighboring routers")
var dev_name = flag.String("interface", "", "Interface to discover on")
var tcptun = flag.String("tcptun", "",
	"Address to try and tcp tunnel to")
var keyfile = flag.String("keyfile", "",
	"The keyfile we should check for a key and write our current key to")

func GetLinkLocalAddr(dev net.Interface) (*net.IPAddr, error) {
	dev_addrs, err := dev.Addrs()
	if err != nil {
		return nil, err
	}

	var ll_addr net.IP

	for _, dev_addr := range dev_addrs {
		addr, _, err := net.ParseCIDR(dev_addr.String())
		if err != nil {
			return nil, err
		}

		if addr.IsLinkLocalUnicast() {
			ll_addr = addr
			break
		}
	}

	if ll_addr == nil {
		log.Fatalf("Couldn't find interface link-local addresses %v", dev.Name)
	}

	ll_addr_zone := fmt.Sprintf("%s%%%s", ll_addr.String(), dev.Name)
	resolved_ll_addr, err := net.ResolveIPAddr("ip6", ll_addr_zone)
	if err != nil {
		return nil, err
	}
	return resolved_ll_addr, nil
}

func FindNeighbors(dev net.Interface, ll_addr *net.IPAddr, key node.PublicKey) <-chan *node.FrameData {
	conn, err := l2.ConnectExistingDevice(dev.Name)
	if err != nil {
		log.Fatal(err)
	}
	nf := node.NewNeighborData(key, ll_addr)
	channel, err := nf.Find(dev.HardwareAddr, conn)
	if err != nil {
		log.Fatal(err)
	}
	return channel
}

func Probe(key node.PrivateKey, n *node.Server, devs []net.Interface) {
	public_key := key.PublicKey()

	for _, dev := range devs {
		if dev.Name == "lo" {
			continue
		}

		ll_addr, err := GetLinkLocalAddr(dev)
		if err != nil {
			log.Fatal(err)
		}

		neighbors := FindNeighbors(dev, ll_addr, public_key)
		go func() {
			for neighbor := range neighbors {
				log.Printf("Neighbour Found %v", neighbor.NodeAddr)
				err := n.Connect(neighbor.LLAddrStr)
				if err != nil {
					log.Printf("Error connecting: %v", err)
					continue
				}
				log.Printf("Connection established to %v", neighbor.NodeAddr)
			}
		}()
	}
}

func main() {

	runtime.GOMAXPROCS(runtime.NumCPU())

	log.Print("Starting")

	flag.Parse()
	quit := make(chan bool)

	var key node.PrivateKey
	var err error

	if len(*keyfile) > 0 {
		key, err = node.LoadKey(*keyfile)
	} else {
		log.Print("Generating key")
		key, err = node.NewECDSAKey()
		if err != nil {
			log.Fatal(err)
		}
	}
	log.Printf("Key is %x", key.PublicKey().Hash())

	n := node.NewServer(key)

	if *autodiscover {
		var devs []net.Interface

		if *dev_name == "" {
			devs, err = net.Interfaces()
			if err != nil {
				log.Fatal(err)
			}
		} else {
			dev, err := net.InterfaceByName(*dev_name)
			if err != nil {
				log.Fatal(err)
			}
			devs = []net.Interface{*dev}
		}

		go Probe(key, n, devs)
	}

	if !*nolisten {
		log.Print("Starting Listening")
		err := n.Listen(*listen)
		if err != nil {
			log.Printf("Error listening: %v", err)
		}
	}

	for _, ip := range strings.Split(*connect, ",") {
		if len(ip) == 0 {
			continue
		}
		err := n.Connect(ip)
		if err != nil {
			log.Printf("Error connecting to %s: %v", ip, err)
		}
	}

	if len(*tcptun) > 0 {
		log.Printf("Establishing tcp tunnel to %v", *tcptun)
		i, err := tuntap.Open("tun%d", tuntap.DevTun)
		if err != nil {
			log.Fatal(err)
		}
		dest := ""
		_, err = fmt.Sscanf(*tcptun, "%x", &dest)
		if err != nil {
			log.Fatal(err)
		}
		t := node.NewTCPTunnel(i, n.Node(), node.NodeAddress(dest), 1)
		t = t
		<-quit
	}

	<-quit
}
<commit_msg>Better error handling for getting a link local address. Shouldn't fatal if nothing is connected to an interface.<commit_after>package main

import (
	"github.com/AutoRoute/l2"
	"github.com/AutoRoute/node"

	"code.google.com/p/tuntap"

	"flag"
	"fmt"
	"log"
	"net"
	"runtime"
	"strings"
)

var listen = flag.String("listen", "127.0.0.1:34321",
	"The address to listen to incoming connections on")
var nolisten = flag.Bool("nolisten", false, "Disables listening")
var connect = flag.String("connect", "",
	"Comma separated list of addresses to connect to")
var autodiscover = flag.Bool("auto", false,
	"Whether we should try and find neighboring routers")
var dev_name = flag.String("interface", "", "Interface to discover on")
var tcptun = flag.String("tcptun", "",
	"Address to try and tcp tunnel to")
var keyfile = flag.String("keyfile", "",
	"The keyfile we should check for a key and write our current key to")

type LinkLocalError struct {
	msg   string
	fatal bool
}

func (e LinkLocalError) Error() string { return e.msg }
func (e LinkLocalError) IsFatal() bool { return e.fatal }

func GetLinkLocalAddr(dev net.Interface) (*net.IPAddr, error) {
	dev_addrs, err := dev.Addrs()
	if err != nil {
		return nil, err
	}

	var ll_addr net.IP

	for _, dev_addr := range dev_addrs {
		addr, _, err := net.ParseCIDR(dev_addr.String())
		if err != nil {
			return nil, err
		}

		if addr.IsLinkLocalUnicast() {
			ll_addr = addr
			break
		}
	}

	if ll_addr == nil {
		msg := fmt.Sprintf("Couldn't find link-local address on interface %v", dev.Name)
		e := LinkLocalError{msg, false}
		return nil, error(e)
	}

	ll_addr_zone := fmt.Sprintf("%s%%%s", ll_addr.String(), dev.Name)
	resolved_ll_addr, err := net.ResolveIPAddr("ip6", ll_addr_zone)
	if err != nil {
		e := LinkLocalError{err.Error(), true}
		return nil, error(e)
	}
	return resolved_ll_addr, nil
}

func FindNeighbors(dev net.Interface, ll_addr *net.IPAddr, key node.PublicKey) <-chan *node.FrameData {
	conn, err := l2.ConnectExistingDevice(dev.Name)
	if err != nil {
		log.Fatal(err)
	}
	nf := node.NewNeighborData(key, ll_addr)
	channel, err := nf.Find(dev.HardwareAddr, conn)
	if err != nil {
		log.Fatal(err)
	}
	return channel
}

func Probe(key node.PrivateKey, n *node.Server, devs []net.Interface) {
	public_key := key.PublicKey()

	// find neighbors of each interface
	for _, dev := range devs {
		if dev.Name == "lo" {
			continue
		}

		ll_addr, err := GetLinkLocalAddr(dev)
		if err != nil {
			e, _ := err.(LinkLocalError)
			if e.IsFatal() {
				log.Fatal(e)
			} else {
				log.Print(e)
			}
		}

		neighbors := FindNeighbors(dev, ll_addr, public_key)
		go func() {
			for neighbor := range neighbors {
				log.Printf("Neighbour Found %v", neighbor.NodeAddr)
				err := n.Connect(neighbor.LLAddrStr)
				if err != nil {
					log.Printf("Error connecting: %v", err)
					continue
				}
				log.Printf("Connection established to %v", neighbor.NodeAddr)
			}
		}()
	}
}

func main() {

	runtime.GOMAXPROCS(runtime.NumCPU())

	log.Print("Starting")

	flag.Parse()
	quit := make(chan bool)

	var key node.PrivateKey
	var err error

	if len(*keyfile) > 0 {
		key, err = node.LoadKey(*keyfile)
	} else {
		log.Print("Generating key")
		key, err = node.NewECDSAKey()
		if err != nil {
			log.Fatal(err)
		}
	}
	log.Printf("Key is %x", key.PublicKey().Hash())

	n := node.NewServer(key)

	if *autodiscover {
		var devs []net.Interface

		if *dev_name == "" {
			devs, err = net.Interfaces()
			if err != nil {
				log.Fatal(err)
			}
		} else {
			dev, err := net.InterfaceByName(*dev_name)
			if err != nil {
				log.Fatal(err)
			}
			devs = []net.Interface{*dev}
		}

		log.Printf("Starting to probe on interfaces: ")
		for _, dev := range devs {
			log.Printf("%v", dev.Name)
		}
		go Probe(key, n, devs)
	}

	if !*nolisten {
		log.Print("Starting to listen")
		err := n.Listen(*listen)
		if err != nil {
			log.Printf("Error listening: %v", err)
		}
	}

	for _, ip := range strings.Split(*connect, ",") {
		if len(ip) == 0 {
			continue
		}
		err := n.Connect(ip)
		if err != nil {
			log.Printf("Error connecting to %s: %v", ip, err)
		}
	}

	if len(*tcptun) > 0 {
		log.Printf("Establishing tcp tunnel to %v", *tcptun)
		i, err := tuntap.Open("tun%d", tuntap.DevTun)
		if err != nil {
			log.Fatal(err)
		}
		dest := ""
		_, err = fmt.Sscanf(*tcptun, "%x", &dest)
		if err != nil {
			log.Fatal(err)
		}
		t := node.NewTCPTunnel(i, n.Node(), node.NodeAddress(dest), 1)
		t = t
		<-quit
	}

	<-quit
}
<|endoftext|>


<commit_before>// This is an NSQ client that reads the specified topic/channel
// and re-publishes the messages to destination nsqd via TCP

package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strconv"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/bitly/go-hostpool"
	"github.com/bitly/timer_metrics"
	"github.com/nsqio/go-nsq"
	"github.com/nsqio/nsq/internal/app"
	"github.com/nsqio/nsq/internal/protocol"
	"github.com/nsqio/nsq/internal/version"
)

const (
	ModeRoundRobin = iota
	ModeHostPool
)

var (
	showVersion = flag.Bool("version", false, "print version string")
	channel     = flag.String("channel", "nsq_to_nsq", "nsq channel")
	destTopic   = flag.String("destination-topic", "", "activate single destination nsq topic Note: default destination topic is the topic itself")
	maxInFlight = flag.Int("max-in-flight", 200, "max number of messages to allow in flight")

	statusEvery = flag.Int("status-every", 250, "the # of requests between logging status (per destination), 0 disables")
	mode        = flag.String("mode", "hostpool", "the upstream request mode options: round-robin, hostpool (default), epsilon-greedy")

	nsqdTCPAddrs        = app.StringArray{}
	lookupdHTTPAddrs    = app.StringArray{}
	destNsqdTCPAddrs    = app.StringArray{}
	whitelistJSONFields = app.StringArray{}
	topics              = app.StringArray{}

	requireJSONField = flag.String("require-json-field", "", "for JSON messages: only pass messages that contain this field")
	requireJSONValue = flag.String("require-json-value", "", "for JSON messages: only pass messages in which the required field has this value")
)

func init() {
	flag.Var(&nsqdTCPAddrs, "nsqd-tcp-address", "nsqd TCP address (may be given multiple times)")
	flag.Var(&destNsqdTCPAddrs, "destination-nsqd-tcp-address", "destination nsqd TCP address (may be given multiple times)")
	flag.Var(&lookupdHTTPAddrs, "lookupd-http-address", "lookupd HTTP address (may be given multiple times)")
	flag.Var(&topics, "topic", "nsq topic (may be given multiple times)")
	flag.Var(&whitelistJSONFields, "whitelist-json-field", "for JSON messages: pass this field (may be given multiple times)")
}

type PublishHandler struct {
	// 64bit atomic vars need to be first for proper alignment on 32bit platforms
	counter uint64

	addresses app.StringArray
	producers map[string]*nsq.Producer
	mode      int
	hostPool  hostpool.HostPool
	respChan  chan *nsq.ProducerTransaction

	requireJSONValueParsed   bool
	requireJSONValueIsNumber bool
	requireJSONNumber        float64

	perAddressStatus map[string]*timer_metrics.TimerMetrics
	timermetrics     *timer_metrics.TimerMetrics
}

type TopicHandler struct {
	publishHandler   *PublishHandler
	destinationTopic string
}

func (ph *PublishHandler) responder() {
	var msg *nsq.Message
	var startTime time.Time
	var address string
	var hostPoolResponse hostpool.HostPoolResponse

	for t := range ph.respChan {
		switch ph.mode {
		case ModeRoundRobin:
			msg = t.Args[0].(*nsq.Message)
			startTime = t.Args[1].(time.Time)
			hostPoolResponse = nil
			address = t.Args[2].(string)
		case ModeHostPool:
			msg = t.Args[0].(*nsq.Message)
			startTime = t.Args[1].(time.Time)
			hostPoolResponse = t.Args[2].(hostpool.HostPoolResponse)
			address = hostPoolResponse.Host()
		}

		success := t.Error == nil

		if hostPoolResponse != nil {
			if !success {
				hostPoolResponse.Mark(errors.New("failed"))
			} else {
				hostPoolResponse.Mark(nil)
			}
		}

		if success {
			msg.Finish()
		} else {
			msg.Requeue(-1)
		}

		ph.perAddressStatus[address].Status(startTime)
		ph.timermetrics.Status(startTime)
	}
}

func (ph *PublishHandler) shouldPassMessage(js map[string]interface{}) (bool, bool) {
	pass := true
	backoff := false

	if *requireJSONField == "" {
		return pass, backoff
	}

	if *requireJSONValue != "" && !ph.requireJSONValueParsed {
		// cache conversion in case needed while filtering json
		var err error
		ph.requireJSONNumber, err = strconv.ParseFloat(*requireJSONValue, 64)
		ph.requireJSONValueIsNumber = (err == nil)
		ph.requireJSONValueParsed = true
	}

	v, ok := js[*requireJSONField]
	if !ok {
		pass = false
		if *requireJSONValue != "" {
			log.Printf("ERROR: missing field to check required value")
			backoff = true
		}
	} else if *requireJSONValue != "" {
		// if command-line argument can't convert to float, then it can't match a number
		// if it can, also integers (up to 2^53 or so) can be compared as float64
		if s, ok := v.(string); ok {
			if s != *requireJSONValue {
				pass = false
			}
		} else if ph.requireJSONValueIsNumber {
			f, ok := v.(float64)
			if !ok || f != ph.requireJSONNumber {
				pass = false
			}
		} else {
			// json value wasn't a plain string, and argument wasn't a number
			// give up on comparisons of other types
			pass = false
		}
	}

	return pass, backoff
}

func filterMessage(js map[string]interface{}, rawMsg []byte) ([]byte, error) {
	if len(whitelistJSONFields) == 0 {
		// no change
		return rawMsg, nil
	}

	newMsg := make(map[string]interface{}, len(whitelistJSONFields))

	for _, key := range whitelistJSONFields {
		value, ok := js[key]
		if ok {
			// avoid printing int as float (go 1.0)
			switch tvalue := value.(type) {
			case float64:
				ivalue := int64(tvalue)
				if float64(ivalue) == tvalue {
					newMsg[key] = ivalue
				} else {
					newMsg[key] = tvalue
				}
			default:
				newMsg[key] = value
			}
		}
	}

	newRawMsg, err := json.Marshal(newMsg)
	if err != nil {
		return nil, fmt.Errorf("unable to marshal filtered message %v", newMsg)
	}
	return newRawMsg, nil
}

func (t *TopicHandler) HandleMessage(m *nsq.Message) error {
	return t.publishHandler.HandleMessage(m, t.destinationTopic)
}

func (ph *PublishHandler) HandleMessage(m *nsq.Message, destinationTopic string) error {
	var err error
	msgBody := m.Body

	if *requireJSONField != "" || len(whitelistJSONFields) > 0 {
		var js map[string]interface{}
		err = json.Unmarshal(msgBody, &js)
		if err != nil {
			log.Printf("ERROR: Unable to decode json: %s", msgBody)
			return nil
		}

		if pass, backoff := ph.shouldPassMessage(js); !pass {
			if backoff {
				return errors.New("backoff")
			}
			return nil
		}

		msgBody, err = filterMessage(js, msgBody)

		if err != nil {
			log.Printf("ERROR: filterMessage() failed: %s", err)
			return err
		}
	}

	startTime := time.Now()

	switch ph.mode {
	case ModeRoundRobin:
		counter := atomic.AddUint64(&ph.counter, 1)
		idx := counter % uint64(len(ph.addresses))
		addr := ph.addresses[idx]
		p := ph.producers[addr]
		err = p.PublishAsync(destinationTopic, msgBody, ph.respChan, m, startTime, addr)
	case ModeHostPool:
		hostPoolResponse := ph.hostPool.Get()
		p := ph.producers[hostPoolResponse.Host()]
		err = p.PublishAsync(destinationTopic, msgBody, ph.respChan, m, startTime, hostPoolResponse)
		if err != nil {
			hostPoolResponse.Mark(err)
		}
	}

	if err != nil {
		return err
	}
	m.DisableAutoResponse()
	return nil
}

func hasArg(s string) bool {
	argExist := false
	flag.Visit(func(f *flag.Flag) {
		if f.Name == s {
			argExist = true
		}
	})
	return argExist
}

func main() {
	var selectedMode int

	cCfg := nsq.NewConfig()
	pCfg := nsq.NewConfig()

	flag.Var(&nsq.ConfigFlag{cCfg}, "consumer-opt", "option to passthrough to nsq.Consumer (may be given multiple times, see http://godoc.org/github.com/nsqio/go-nsq#Config)")
	flag.Var(&nsq.ConfigFlag{pCfg}, "producer-opt", "option to passthrough to nsq.Producer (may be given multiple times, see http://godoc.org/github.com/nsqio/go-nsq#Config)")

	flag.Parse()

	if *showVersion {
		fmt.Printf("nsq_to_nsq v%s\n", version.Binary)
		return
	}

	if len(topics) == 0 || *channel == "" {
		log.Fatal("--topic and --channel are required")
	}

	for _, topic := range topics {
		if !protocol.IsValidTopicName(topic) {
			log.Fatal("--topic is invalid")
		}
	}

	if *destTopic != "" && !protocol.IsValidTopicName(*destTopic) {
		log.Fatal("--destination-topic is invalid")
	}

	if !protocol.IsValidChannelName(*channel) {
		log.Fatal("--channel is invalid")
	}

	if len(nsqdTCPAddrs) == 0 && len(lookupdHTTPAddrs) == 0 {
		log.Fatal("--nsqd-tcp-address or --lookupd-http-address required")
	}
	if len(nsqdTCPAddrs) > 0 && len(lookupdHTTPAddrs) > 0 {
		log.Fatal("use --nsqd-tcp-address or --lookupd-http-address not both")
	}

	if len(destNsqdTCPAddrs) == 0 {
		log.Fatal("--destination-nsqd-tcp-address required")
	}

	switch *mode {
	case "round-robin":
		selectedMode = ModeRoundRobin
	case "hostpool", "epsilon-greedy":
		selectedMode = ModeHostPool
	}

	termChan := make(chan os.Signal, 1)
	signal.Notify(termChan, syscall.SIGINT, syscall.SIGTERM)

	defaultUA := fmt.Sprintf("nsq_to_nsq/%s go-nsq/%s", version.Binary, nsq.VERSION)

	cCfg.UserAgent = defaultUA
	cCfg.MaxInFlight = *maxInFlight
	pCfg.UserAgent = defaultUA

	producers := make(map[string]*nsq.Producer)
	for _, addr := range destNsqdTCPAddrs {
		producer, err := nsq.NewProducer(addr, pCfg)
		if err != nil {
			log.Fatalf("failed creating producer %s", err)
		}
		producers[addr] = producer
	}

	perAddressStatus := make(map[string]*timer_metrics.TimerMetrics)
	if len(destNsqdTCPAddrs) == 1 {
		// disable since there is only one address
		perAddressStatus[destNsqdTCPAddrs[0]] = timer_metrics.NewTimerMetrics(0, "")
	} else {
		for _, a := range destNsqdTCPAddrs {
			perAddressStatus[a] = timer_metrics.NewTimerMetrics(*statusEvery,
				fmt.Sprintf("[%s]:", a))
		}
	}

	hostPool := hostpool.New(destNsqdTCPAddrs)
	if *mode == "epsilon-greedy" {
		hostPool = hostpool.NewEpsilonGreedy(destNsqdTCPAddrs, 0, &hostpool.LinearEpsilonValueCalculator{})
	}

	var consumerList []*nsq.Consumer
	var singleDestinationTopicHandler *TopicHandler

	publisherHandlerRef := &PublishHandler{
		addresses:        destNsqdTCPAddrs,
		producers:        producers,
		mode:             selectedMode,
		hostPool:         hostPool,
		respChan:         make(chan *nsq.ProducerTransaction, len(destNsqdTCPAddrs)),
		perAddressStatus: perAddressStatus,
		timermetrics:     timer_metrics.NewTimerMetrics(*statusEvery, "[aggregate]:"),
	}

	if *destTopic != "" {
		singleDestinationTopicHandler = &TopicHandler{
			publishHandler:   publisherHandlerRef,
			destinationTopic: *destTopic,
		}
	}

	for _, topic := range topics {
		consumer, err := nsq.NewConsumer(topic, *channel, cCfg)
		consumerList = append(consumerList, consumer)
		if err != nil {
			log.Fatal(err)
		}
		if singleDestinationTopicHandler != nil {
			consumer.AddConcurrentHandlers(singleDestinationTopicHandler, len(destNsqdTCPAddrs))
		} else {
			topicHandler := &TopicHandler{
				publishHandler:   publisherHandlerRef,
				destinationTopic: topic,
			}

			consumer.AddConcurrentHandlers(topicHandler, len(destNsqdTCPAddrs))
		}
	}
	for i := 0; i < len(destNsqdTCPAddrs); i++ {
		go publisherHandlerRef.responder()
	}

	for _, consumer := range consumerList {
		err := consumer.ConnectToNSQDs(nsqdTCPAddrs)
		if err != nil {
			log.Fatal(err)
		}
	}

	for _, consumer := range consumerList {
		err := consumer.ConnectToNSQLookupds(lookupdHTTPAddrs)
		if err != nil {
			log.Fatal(err)
		}
	}

	<-termChan // wait for signal

	for _, consumer := range consumerList {
		consumer.Stop()
	}
	for _, consumer := range consumerList {
		<-consumer.StopChan
	}
}
<commit_msg>nsq_to_nsq: simplify destination-topic handling<commit_after>// This is an NSQ client that reads the specified topic/channel
// and re-publishes the messages to destination nsqd via TCP

package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strconv"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/bitly/go-hostpool"
	"github.com/bitly/timer_metrics"
	"github.com/nsqio/go-nsq"
	"github.com/nsqio/nsq/internal/app"
	"github.com/nsqio/nsq/internal/protocol"
	"github.com/nsqio/nsq/internal/version"
)

const (
	ModeRoundRobin = iota
	ModeHostPool
)

var (
	showVersion = flag.Bool("version", false, "print version string")
	channel     = flag.String("channel", "nsq_to_nsq", "nsq channel")
	destTopic   = flag.String("destination-topic", "", "use this destination topic for all consumed topics (default is consumed topic name)")
	maxInFlight = flag.Int("max-in-flight", 200, "max number of messages to allow in flight")

	statusEvery = flag.Int("status-every", 250, "the # of requests between logging status (per destination), 0 disables")
	mode        = flag.String("mode", "hostpool", "the upstream request mode options: round-robin, hostpool (default), epsilon-greedy")

	nsqdTCPAddrs        = app.StringArray{}
	lookupdHTTPAddrs    = app.StringArray{}
	destNsqdTCPAddrs    = app.StringArray{}
	whitelistJSONFields = app.StringArray{}
	topics              = app.StringArray{}

	requireJSONField = flag.String("require-json-field", "", "for JSON messages: only pass messages that contain this field")
	requireJSONValue = flag.String("require-json-value", "", "for JSON messages: only pass messages in which the required field has this value")
)

func init() {
	flag.Var(&nsqdTCPAddrs, "nsqd-tcp-address", "nsqd TCP address (may be given multiple times)")
	flag.Var(&destNsqdTCPAddrs, "destination-nsqd-tcp-address", "destination nsqd TCP address (may be given multiple times)")
	flag.Var(&lookupdHTTPAddrs, "lookupd-http-address", "lookupd HTTP address (may be given multiple times)")
	flag.Var(&topics, "topic", "nsq topic (may be given multiple times)")
	flag.Var(&whitelistJSONFields, "whitelist-json-field", "for JSON messages: pass this field (may be given multiple times)")
}

type PublishHandler struct {
	// 64bit atomic vars need to be first for proper alignment on 32bit platforms
	counter uint64

	addresses app.StringArray
	producers map[string]*nsq.Producer
	mode      int
	hostPool  hostpool.HostPool
	respChan  chan *nsq.ProducerTransaction

	requireJSONValueParsed   bool
	requireJSONValueIsNumber bool
	requireJSONNumber        float64

	perAddressStatus map[string]*timer_metrics.TimerMetrics
	timermetrics     *timer_metrics.TimerMetrics
}

type TopicHandler struct {
	publishHandler   *PublishHandler
	destinationTopic string
}

func (ph *PublishHandler) responder() {
	var msg *nsq.Message
	var startTime time.Time
	var address string
	var hostPoolResponse hostpool.HostPoolResponse

	for t := range ph.respChan {
		switch ph.mode {
		case ModeRoundRobin:
			msg = t.Args[0].(*nsq.Message)
			startTime = t.Args[1].(time.Time)
			hostPoolResponse = nil
			address = t.Args[2].(string)
		case ModeHostPool:
			msg = t.Args[0].(*nsq.Message)
			startTime = t.Args[1].(time.Time)
			hostPoolResponse = t.Args[2].(hostpool.HostPoolResponse)
			address = hostPoolResponse.Host()
		}

		success := t.Error == nil

		if hostPoolResponse != nil {
			if !success {
				hostPoolResponse.Mark(errors.New("failed"))
			} else {
				hostPoolResponse.Mark(nil)
			}
		}

		if success {
			msg.Finish()
		} else {
			msg.Requeue(-1)
		}

		ph.perAddressStatus[address].Status(startTime)
		ph.timermetrics.Status(startTime)
	}
}

func (ph *PublishHandler) shouldPassMessage(js map[string]interface{}) (bool, bool) {
	pass := true
	backoff := false

	if *requireJSONField == "" {
		return pass, backoff
	}

	if *requireJSONValue != "" && !ph.requireJSONValueParsed {
		// cache conversion in case needed while filtering json
		var err error
		ph.requireJSONNumber, err = strconv.ParseFloat(*requireJSONValue, 64)
		ph.requireJSONValueIsNumber = (err == nil)
		ph.requireJSONValueParsed = true
	}

	v, ok := js[*requireJSONField]
	if !ok {
		pass = false
		if *requireJSONValue != "" {
			log.Printf("ERROR: missing field to check required value")
			backoff = true
		}
	} else if *requireJSONValue != "" {
		// if command-line argument can't convert to float, then it can't match a number
		// if it can, also integers (up to 2^53 or so) can be compared as float64
		if s, ok := v.(string); ok {
			if s != *requireJSONValue {
				pass = false
			}
		} else if ph.requireJSONValueIsNumber {
			f, ok := v.(float64)
			if !ok || f != ph.requireJSONNumber {
				pass = false
			}
		} else {
			// json value wasn't a plain string, and argument wasn't a number
			// give up on comparisons of other types
			pass = false
		}
	}

	return pass, backoff
}

func filterMessage(js map[string]interface{}, rawMsg []byte) ([]byte, error) {
	if len(whitelistJSONFields) == 0 {
		// no change
		return rawMsg, nil
	}

	newMsg := make(map[string]interface{}, len(whitelistJSONFields))

	for _, key := range whitelistJSONFields {
		value, ok := js[key]
		if ok {
			// avoid printing int as float (go 1.0)
			switch tvalue := value.(type) {
			case float64:
				ivalue := int64(tvalue)
				if float64(ivalue) == tvalue {
					newMsg[key] = ivalue
				} else {
					newMsg[key] = tvalue
				}
			default:
				newMsg[key] = value
			}
		}
	}

	newRawMsg, err := json.Marshal(newMsg)
	if err != nil {
		return nil, fmt.Errorf("unable to marshal filtered message %v", newMsg)
	}
	return newRawMsg, nil
}

func (t *TopicHandler) HandleMessage(m *nsq.Message) error {
	return t.publishHandler.HandleMessage(m, t.destinationTopic)
}

func (ph *PublishHandler) HandleMessage(m *nsq.Message, destinationTopic string) error {
	var err error
	msgBody := m.Body

	if *requireJSONField != "" || len(whitelistJSONFields) > 0 {
		var js map[string]interface{}
		err = json.Unmarshal(msgBody, &js)
		if err != nil {
			log.Printf("ERROR: Unable to decode json: %s", msgBody)
			return nil
		}

		if pass, backoff := ph.shouldPassMessage(js); !pass {
			if backoff {
				return errors.New("backoff")
			}
			return nil
		}

		msgBody, err = filterMessage(js, msgBody)

		if err != nil {
			log.Printf("ERROR: filterMessage() failed: %s", err)
			return err
		}
	}

	startTime := time.Now()

	switch ph.mode {
	case ModeRoundRobin:
		counter := atomic.AddUint64(&ph.counter, 1)
		idx := counter % uint64(len(ph.addresses))
		addr := ph.addresses[idx]
		p := ph.producers[addr]
		err = p.PublishAsync(destinationTopic, msgBody, ph.respChan, m, startTime, addr)
	case ModeHostPool:
		hostPoolResponse := ph.hostPool.Get()
		p := ph.producers[hostPoolResponse.Host()]
		err = p.PublishAsync(destinationTopic, msgBody, ph.respChan, m, startTime, hostPoolResponse)
		if err != nil {
			hostPoolResponse.Mark(err)
		}
	}

	if err != nil {
		return err
	}
	m.DisableAutoResponse()
	return nil
}

func hasArg(s string) bool {
	argExist := false
	flag.Visit(func(f *flag.Flag) {
		if f.Name == s {
			argExist = true
		}
	})
	return argExist
}

func main() {
	var selectedMode int

	cCfg := nsq.NewConfig()
	pCfg := nsq.NewConfig()

	flag.Var(&nsq.ConfigFlag{cCfg}, "consumer-opt", "option to passthrough to nsq.Consumer (may be given multiple times, see http://godoc.org/github.com/nsqio/go-nsq#Config)")
	flag.Var(&nsq.ConfigFlag{pCfg}, "producer-opt", "option to passthrough to nsq.Producer (may be given multiple times, see http://godoc.org/github.com/nsqio/go-nsq#Config)")

	flag.Parse()

	if *showVersion {
		fmt.Printf("nsq_to_nsq v%s\n", version.Binary)
		return
	}

	if len(topics) == 0 || *channel == "" {
		log.Fatal("--topic and --channel are required")
	}

	for _, topic := range topics {
		if !protocol.IsValidTopicName(topic) {
			log.Fatal("--topic is invalid")
		}
	}

	if *destTopic != "" && !protocol.IsValidTopicName(*destTopic) {
		log.Fatal("--destination-topic is invalid")
	}

	if !protocol.IsValidChannelName(*channel) {
		log.Fatal("--channel is invalid")
	}

	if len(nsqdTCPAddrs) == 0 && len(lookupdHTTPAddrs) == 0 {
		log.Fatal("--nsqd-tcp-address or --lookupd-http-address required")
	}
	if len(nsqdTCPAddrs) > 0 && len(lookupdHTTPAddrs) > 0 {
		log.Fatal("use --nsqd-tcp-address or --lookupd-http-address not both")
	}

	if len(destNsqdTCPAddrs) == 0 {
		log.Fatal("--destination-nsqd-tcp-address required")
	}

	switch *mode {
	case "round-robin":
		selectedMode = ModeRoundRobin
	case "hostpool", "epsilon-greedy":
		selectedMode = ModeHostPool
	}

	termChan := make(chan os.Signal, 1)
	signal.Notify(termChan, syscall.SIGINT, syscall.SIGTERM)

	defaultUA := fmt.Sprintf("nsq_to_nsq/%s go-nsq/%s", version.Binary, nsq.VERSION)

	cCfg.UserAgent = defaultUA
	cCfg.MaxInFlight = *maxInFlight
	pCfg.UserAgent = defaultUA

	producers := make(map[string]*nsq.Producer)
	for _, addr := range destNsqdTCPAddrs {
		producer, err := nsq.NewProducer(addr, pCfg)
		if err != nil {
			log.Fatalf("failed creating producer %s", err)
		}
		producers[addr] = producer
	}

	perAddressStatus := make(map[string]*timer_metrics.TimerMetrics)
	if len(destNsqdTCPAddrs) == 1 {
		// disable since there is only one address
		perAddressStatus[destNsqdTCPAddrs[0]] = timer_metrics.NewTimerMetrics(0, "")
	} else {
		for _, a := range destNsqdTCPAddrs {
			perAddressStatus[a] = timer_metrics.NewTimerMetrics(*statusEvery,
				fmt.Sprintf("[%s]:", a))
		}
	}

	hostPool := hostpool.New(destNsqdTCPAddrs)
	if *mode == "epsilon-greedy" {
		hostPool = hostpool.NewEpsilonGreedy(destNsqdTCPAddrs, 0, &hostpool.LinearEpsilonValueCalculator{})
	}

	var consumerList []*nsq.Consumer

	publisher := &PublishHandler{
		addresses:        destNsqdTCPAddrs,
		producers:        producers,
		mode:             selectedMode,
		hostPool:         hostPool,
		respChan:         make(chan *nsq.ProducerTransaction, len(destNsqdTCPAddrs)),
		perAddressStatus: perAddressStatus,
		timermetrics:     timer_metrics.NewTimerMetrics(*statusEvery, "[aggregate]:"),
	}

	for _, topic := range topics {
		consumer, err := nsq.NewConsumer(topic, *channel, cCfg)
		consumerList = append(consumerList, consumer)
		if err != nil {
			log.Fatal(err)
		}

		publishTopic := topic
		if *destTopic != "" {
			publishTopic = *destTopic
		}
		topicHandler := &TopicHandler{
			publishHandler:   publisher,
			destinationTopic: publishTopic,
		}
		consumer.AddConcurrentHandlers(topicHandler, len(destNsqdTCPAddrs))
	}
	for i := 0; i < len(destNsqdTCPAddrs); i++ {
		go publisher.responder()
	}

	for _, consumer := range consumerList {
		err := consumer.ConnectToNSQDs(nsqdTCPAddrs)
		if err != nil {
			log.Fatal(err)
		}
	}

	for _, consumer := range consumerList {
		err := consumer.ConnectToNSQLookupds(lookupdHTTPAddrs)
		if err != nil {
			log.Fatal(err)
		}
	}

	<-termChan // wait for signal

	for _, consumer := range consumerList {
		consumer.Stop()
	}
	for _, consumer := range consumerList {
		<-consumer.StopChan
	}
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net/http"

	"gopkg.in/mgo.v2"
)

// Period hold the information of an activity
type Period struct {
	Type string // pomodoro or rest
}

type Periods []Period

// createPeriod handler lets you create a period obect.
//
func createPeriod(w http.ResponseWriter, r *http.Request) {

	var uri string
	uri = getMongoURI()

	log.Println("uri found: ", uri)

	var err error
	var session *mgo.Session

	log.Println("start dial")
	if session, err = mgo.Dial(uri); err != nil {
		panic(err)
	}
	log.Println("dial done")

	defer session.Close()
	log.Println("setting session mode")
	session.SetMode(mgo.Monotonic, true)

	log.Println("trying to get pomotime database")
	collection := session.DB("").C("period")

	log.Println("start insert to period entity")
	var p Period
	p = Period{"pomodoro"}

	if err = collection.Insert(&p); err != nil {
		log.Fatal(err)
	}

	data := struct {
		Status           int    `json:"status"`
		DeveloperMessage string `json:"developerMessage"`
		UserMessage      string `json:"userMessage"`
	}{
		Status:           200,
		DeveloperMessage: "app correctly stored period in database.",
		UserMessage:      "your data was stored correctly.",
	}

	log.Println("sending response")
	if err := renderJson(w, data); err != nil {
		log.Println(err)
	}
}

// getPeriods handler returns a list of periods.
//
func getPeriods(w http.ResponseWriter, r *http.Request) {

	var uri string
	uri = getMongoURI()

	log.Println("uri found: ", uri)

	var err error
	var session *mgo.Session

	log.Println("start dial")
	if session, err = mgo.Dial(uri); err != nil {
		panic(err)
	}
	log.Println("dial done")

	defer session.Close()
	log.Println("setting session mode")
	session.SetMode(mgo.Monotonic, true)

	log.Println("trying database")
	collection := session.DB("").C("period")

	var results []Period

	log.Println("start search for all period objects")
	if err = collection.Find(nil).All(&results); err != nil {
		log.Fatal(err)
	}

	data := struct {
		Status           int     `json:"status"`
		DeveloperMessage string  `json:"developerMessage"`
		UserMessage      string  `json:"userMessage"`
		Periods          Periods `json:"periods"`
	}{
		Status:           200,
		DeveloperMessage: "returning all periods from database.",
		UserMessage:      "",
		Periods:          results,
	}

	log.Println("sending response")
	if err := renderJson(w, data); err != nil {
		log.Println(err)
	}
}
<commit_msg>pass param value to store period type handle errors<commit_after>package main

import (
	"log"
	"net/http"

	"gopkg.in/mgo.v2"
)

// Period hold the information of an activity
type Period struct {
	Type string `json:"type"` // pomodoro or rest
}

type Periods []Period

// createPeriod handler lets you create a period obect.
//
func createPeriod(w http.ResponseWriter, r *http.Request) {

	var strType string

	if strType = r.FormValue("type"); len(strType) == 0 {
		log.Println("empty type")
		data := struct {
			Status           int    `json:"status"`
			DeveloperMessage string `json:"developerMessage"`
			UserMessage      string `json:"userMessage"`
		}{
			Status:           400,
			DeveloperMessage: "createPeriod had empty type parameter, unable to create period.",
			UserMessage:      "Oops, something went wrong, we are unable to save your data right now.",
		}

		log.Println("sending response")
		if err := renderJson(w, data); err != nil {
			log.Println(err)
		}
		return
	} else {
		if strType != "pomodoro" && strType != "rest" {
			log.Println("empty type")
			data := struct {
				Status           int    `json:"status"`
				DeveloperMessage string `json:"developerMessage"`
				UserMessage      string `json:"userMessage"`
			}{
				Status:           400,
				DeveloperMessage: "createPeriod has wrong type, unable to create period.",
				UserMessage:      "Oops, something went wrong, we are unable to save your data right now.",
			}

			log.Println("sending response")
			if err := renderJson(w, data); err != nil {
				log.Println(err)
			}
			return
		}
	}

	var uri string
	uri = getMongoURI()

	log.Println("uri found: ", uri)

	var err error
	var session *mgo.Session

	log.Println("start dial")
	if session, err = mgo.Dial(uri); err != nil {
		panic(err)
	}
	log.Println("dial done")

	defer session.Close()
	log.Println("setting session mode")
	session.SetMode(mgo.Monotonic, true)

	log.Println("trying to get pomotime database")
	collection := session.DB("").C("period")

	log.Println("start insert to period entity")
	var p Period
	p = Period{strType}

	if err = collection.Insert(&p); err != nil {
		log.Fatal(err)
	}

	data := struct {
		Status           int    `json:"status"`
		DeveloperMessage string `json:"developerMessage"`
		UserMessage      string `json:"userMessage"`
	}{
		Status:           200,
		DeveloperMessage: "app correctly stored period in database.",
		UserMessage:      "your data was stored correctly.",
	}

	log.Println("sending response")
	if err := renderJson(w, data); err != nil {
		log.Println(err)
	}
}

// getPeriods handler returns a list of periods.
//
func getPeriods(w http.ResponseWriter, r *http.Request) {

	var uri string
	uri = getMongoURI()

	log.Println("uri found: ", uri)

	var err error
	var session *mgo.Session

	log.Println("start dial")
	if session, err = mgo.Dial(uri); err != nil {
		panic(err)
	}
	log.Println("dial done")

	defer session.Close()
	log.Println("setting session mode")
	session.SetMode(mgo.Monotonic, true)

	log.Println("trying database")
	collection := session.DB("").C("period")

	var results []Period

	log.Println("start search for all period objects")
	if err = collection.Find(nil).All(&results); err != nil {
		log.Fatal(err)
	}

	data := struct {
		Status           int     `json:"status"`
		DeveloperMessage string  `json:"developerMessage"`
		UserMessage      string  `json:"userMessage"`
		Periods          Periods `json:"periods"`
	}{
		Status:           200,
		DeveloperMessage: "returning all periods from database.",
		UserMessage:      "",
		Periods:          results,
	}

	log.Println("sending response")
	if err := renderJson(w, data); err != nil {
		log.Println(err)
	}
}
<|endoftext|>