<commit_before>package main

import (
	"errors"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strconv"

	"github.com/codegangsta/cli"
	"github.com/twpayne/go-doarama"
)

func newDoaramaClient(c *cli.Context) *doarama.Client {
	return doarama.NewClient(c.GlobalString("apiurl"), c.GlobalString("apiname"), c.GlobalString("apikey"))
}

func newAuthenticatedDoaramaClient(c *cli.Context) (*doarama.Client, error) {
	client := newDoaramaClient(c)
	userId := c.GlobalString("userid")
	userKey := c.GlobalString("userkey")
	switch {
	case userId != "" && userKey == "":
		return client.Anonymous(userId), nil
	case userId == "" && userKey != "":
		return client.Delegate(userKey), nil
	default:
		return nil, errors.New("exactly one of -userid and -userkey must be specified")
	}
}

func newVisualisationURLOptions(c *cli.Context) *doarama.VisualisationURLOptions {
	var vuo doarama.VisualisationURLOptions
	if c.StringSlice("name") != nil {
		vuo.Names = c.StringSlice("name")
	}
	if c.StringSlice("avatar") != nil {
		vuo.Avatars = c.StringSlice("avatar")
	}
	if c.String("avatarbaseurl") != "" {
		vuo.AvatarBaseURL = c.String("avatarbaseurl")
	}
	if c.Bool("fixedaspect") {
		vuo.FixedAspect = c.Bool("fixedaspect")
	}
	if c.Bool("minimalview") {
		vuo.MinimalView = c.Bool("minimalview")
	}
	if c.String("dzml") != "" {
		vuo.DZML = c.String("dzml")
	}
	return &vuo
}

func activityCreateOne(client *doarama.Client, filename string) (*doarama.Activity, error) {
	gpsTrack, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer gpsTrack.Close()
	return client.CreateActivity(filepath.Base(filename), gpsTrack)
}

func activityCreate(c *cli.Context) error {
	client, err := newAuthenticatedDoaramaClient(c)
	if err != nil {
		return err
	}
	typeId := c.Int("typeid")
	for _, arg := range c.Args() {
		a, err := activityCreateOne(client, arg)
		if err != nil {
			log.Print(err)
			continue
		}
		fmt.Printf("ActivityId: %d\n", a.Id)
		if err := a.SetInfo(&doarama.ActivityInfo{
			TypeId: typeId,
		}); err != nil {
			log.Print(err)
			continue
		}
	}
	return nil
}

func activityDelete(c *cli.Context) error {
	client, err := newAuthenticatedDoaramaClient(c)
	if err != nil {
		return err
	}
	var ids []int
	for _, arg := range c.Args() {
		id64, err := strconv.ParseInt(arg, 10, 0)
		if err != nil {
			return err
		}
		ids = append(ids, int(id64))
	}
	for _, id := range ids {
		a := client.Activity(id)
		if err := a.Delete(); err != nil {
			log.Print(err)
			continue
		}
	}
	return nil
}

func create(c *cli.Context) error {
	client, err := newAuthenticatedDoaramaClient(c)
	if err != nil {
		return err
	}
	typeId := c.Int("typeid")
	var as []*doarama.Activity
	for _, arg := range c.Args() {
		a, err := activityCreateOne(client, arg)
		if err != nil {
			break
		}
		err = a.SetInfo(&doarama.ActivityInfo{
			TypeId: typeId,
		})
		if err != nil {
			break
		}
		fmt.Printf("ActivityId: %d\n", a.Id)
		as = append(as, a)
	}
	if err != nil {
		for _, a := range as {
			a.Delete()
		}
		return err
	}
	v, err := client.CreateVisualisation(as)
	if err != nil {
		return err
	}
	fmt.Printf("VisualisationKey: %s\n", v.Key)
	vuo := newVisualisationURLOptions(c)
	fmt.Printf("VisualisationURL: %s\n", v.URL(vuo))
	return nil
}

func queryActivityTypes(c *cli.Context) error {
	client := newDoaramaClient(c)
	ats, err := client.ActivityTypes()
	if err != nil {
		return err
	}
	var names []string
	for name, _ := range ats {
		names = append(names, name)
	}
	sort.Strings(names)
	for _, name := range names {
		fmt.Printf("%s: %d\n", name, ats[name])
	}
	return nil
}

func visualisationCreate(c *cli.Context) error {
	client, err := newAuthenticatedDoaramaClient(c)
	if err != nil {
		return err
	}
	var as []*doarama.Activity
	for _, arg := range c.Args() {
		id64, err := strconv.ParseInt(arg, 10, 0)
		if err != nil {
			return err
		}
		a := client.Activity(int(id64))
		as = append(as, a)
	}
	v, err := client.CreateVisualisation(as)
	if err != nil {
		return err
	}
	fmt.Printf("VisualisationKey: %s\n", v.Key)
	return nil
}

func visualisationDelete(c *cli.Context) error {
	client, err := newAuthenticatedDoaramaClient(c)
	if err != nil {
		return err
	}
	for _, arg := range c.Args() {
		v := client.Visualisation(arg)
		if err := v.Delete(); err != nil {
			log.Print(err)
			continue
		}
	}
	return nil
}

func visualisationURL(c *cli.Context) error {
	client := newDoaramaClient(c)
	vuo := newVisualisationURLOptions(c)
	for _, arg := range c.Args() {
		v := client.Visualisation(arg)
		fmt.Printf("VisualisationURL: %s\n", v.URL(vuo))
	}
	return nil
}

func logError(f func(*cli.Context) error) func(*cli.Context) {
	return func(c *cli.Context) {
		if err := f(c); err != nil {
			log.Fatal(err)
		}
	}
}

func main() {
	app := cli.NewApp()
	app.Name = "doarama"
	app.Usage = "A command line interface to doarama.com"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "apiurl",
			Value:  doarama.API_URL,
			Usage:  "Doarama API URL",
			EnvVar: "DOARAMA_API_URL",
		},
		cli.StringFlag{
			Name:   "apikey",
			Usage:  "Doarama API key",
			EnvVar: "DOARAMA_API_KEY",
		},
		cli.StringFlag{
			Name:   "apiname",
			Usage:  "Doarama API name",
			EnvVar: "DOARAMA_API_NAME",
		},
		cli.StringFlag{
			Name:   "userid",
			Usage:  "Doarama user ID",
			EnvVar: "DOARAMA_USER_ID",
		},
		cli.StringFlag{
			Name:   "userkey",
			Usage:  "Doarama user key",
			EnvVar: "DOARAMA_USER_KEY",
		},
	}
	typeIdFlag := cli.IntFlag{
		Name:  "typeid",
		Usage: "type id",
	}
	nameFlag := cli.StringSliceFlag{
		Name:  "name",
		Usage: "name",
	}
	avatarFlag := cli.StringSliceFlag{
		Name:  "avatar",
		Usage: "avatar",
	}
	avatarBaseUrlFlag := cli.StringFlag{
		Name:  "avatarbaseurl",
		Usage: "avatar base URL",
	}
	fixedAspectFlag := cli.BoolTFlag{
		Name:  "fixedaspect",
		Usage: "fixed aspect",
	}
	minimalViewFlag := cli.BoolFlag{
		Name:  "minimalview",
		Usage: "minimal view",
	}
	dzmlFlag := cli.StringFlag{
		Name:  "dzml",
		Usage: "DZML",
	}
	app.Commands = []cli.Command{
		{
			Name:    "activity",
			Aliases: []string{"a"},
			Usage:   "Manages activities",
			Subcommands: []cli.Command{
				{
					Name:    "create",
					Aliases: []string{"c"},
					Usage:   "Creates an activity from one or more tracklogs",
					Action:  logError(activityCreate),
					Flags: []cli.Flag{
						typeIdFlag,
					},
				},
				{
					Name:    "delete",
					Aliases: []string{"d"},
					Usage:   "Deletes one or more activities by id",
					Action:  logError(activityDelete),
				},
			},
		},
		{
			Name:    "create",
			Aliases: []string{"c"},
			Usage:   "Creates a visualisation URL from one or more tracklogs",
			Action:  logError(create),
			Flags: []cli.Flag{
				typeIdFlag,
				nameFlag,
				avatarFlag,
				avatarBaseUrlFlag,
				fixedAspectFlag,
				minimalViewFlag,
				dzmlFlag,
			},
		},
		{
			Name:    "query-activity-types",
			Aliases: []string{"qat"},
			Usage:   "Queries activity types",
			Action:  logError(queryActivityTypes),
		},
		{
			Name:    "visualisation",
			Aliases: []string{"v"},
			Usage:   "Manages visualisations",
			Subcommands: []cli.Command{
				{
					Name:    "create",
					Aliases: []string{"c"},
					Usage:   "Creates a visualisation from a list of activities",
					Action:  logError(visualisationCreate),
				},
				{
					Name:    "delete",
					Aliases: []string{"d"},
					Usage:   "Deletes one or more visualisations by key",
					Action:  logError(visualisationDelete),
				},
				{
					Name:    "url",
					Aliases: []string{"u"},
					Usage:   "Creates a visualisation URL from a visualisation key",
					Action:  logError(visualisationURL),
					Flags: []cli.Flag{
						nameFlag,
						avatarFlag,
						avatarBaseUrlFlag,
						fixedAspectFlag,
						minimalViewFlag,
						dzmlFlag,
					},
				},
			},
		},
	}
	app.Run(os.Args)
}
<commit_msg>Avoid creating empty visualisations<commit_after>package main

import (
	"errors"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strconv"

	"github.com/codegangsta/cli"
	"github.com/twpayne/go-doarama"
)

func newDoaramaClient(c *cli.Context) *doarama.Client {
	return doarama.NewClient(c.GlobalString("apiurl"), c.GlobalString("apiname"), c.GlobalString("apikey"))
}

func newAuthenticatedDoaramaClient(c *cli.Context) (*doarama.Client, error) {
	client := newDoaramaClient(c)
	userId := c.GlobalString("userid")
	userKey := c.GlobalString("userkey")
	switch {
	case userId != "" && userKey == "":
		return client.Anonymous(userId), nil
	case userId == "" && userKey != "":
		return client.Delegate(userKey), nil
	default:
		return nil, errors.New("exactly one of -userid and -userkey must be specified")
	}
}

func newVisualisationURLOptions(c *cli.Context) *doarama.VisualisationURLOptions {
	var vuo doarama.VisualisationURLOptions
	if c.StringSlice("name") != nil {
		vuo.Names = c.StringSlice("name")
	}
	if c.StringSlice("avatar") != nil {
		vuo.Avatars = c.StringSlice("avatar")
	}
	if c.String("avatarbaseurl") != "" {
		vuo.AvatarBaseURL = c.String("avatarbaseurl")
	}
	if c.Bool("fixedaspect") {
		vuo.FixedAspect = c.Bool("fixedaspect")
	}
	if c.Bool("minimalview") {
		vuo.MinimalView = c.Bool("minimalview")
	}
	if c.String("dzml") != "" {
		vuo.DZML = c.String("dzml")
	}
	return &vuo
}

func activityCreateOne(client *doarama.Client, filename string) (*doarama.Activity, error) {
	gpsTrack, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer gpsTrack.Close()
	return client.CreateActivity(filepath.Base(filename), gpsTrack)
}

func activityCreate(c *cli.Context) error {
	client, err := newAuthenticatedDoaramaClient(c)
	if err != nil {
		return err
	}
	typeId := c.Int("typeid")
	for _, arg := range c.Args() {
		a, err := activityCreateOne(client, arg)
		if err != nil {
			log.Print(err)
			continue
		}
		fmt.Printf("ActivityId: %d\n", a.Id)
		if err := a.SetInfo(&doarama.ActivityInfo{
			TypeId: typeId,
		}); err != nil {
			log.Print(err)
			continue
		}
	}
	return nil
}

func activityDelete(c *cli.Context) error {
	client, err := newAuthenticatedDoaramaClient(c)
	if err != nil {
		return err
	}
	var ids []int
	for _, arg := range c.Args() {
		id64, err := strconv.ParseInt(arg, 10, 0)
		if err != nil {
			return err
		}
		ids = append(ids, int(id64))
	}
	for _, id := range ids {
		a := client.Activity(id)
		if err := a.Delete(); err != nil {
			log.Print(err)
			continue
		}
	}
	return nil
}

func create(c *cli.Context) error {
	client, err := newAuthenticatedDoaramaClient(c)
	if err != nil {
		return err
	}
	typeId := c.Int("typeid")
	var as []*doarama.Activity
	for _, arg := range c.Args() {
		a, err := activityCreateOne(client, arg)
		if err != nil {
			break
		}
		err = a.SetInfo(&doarama.ActivityInfo{
			TypeId: typeId,
		})
		if err != nil {
			break
		}
		fmt.Printf("ActivityId: %d\n", a.Id)
		as = append(as, a)
	}
	if err != nil {
		for _, a := range as {
			a.Delete()
		}
		return err
	}
	if len(as) == 0 {
		return errors.New("no activitiess specified")
	}
	v, err := client.CreateVisualisation(as)
	if err != nil {
		return err
	}
	fmt.Printf("VisualisationKey: %s\n", v.Key)
	vuo := newVisualisationURLOptions(c)
	fmt.Printf("VisualisationURL: %s\n", v.URL(vuo))
	return nil
}

func queryActivityTypes(c *cli.Context) error {
	client := newDoaramaClient(c)
	ats, err := client.ActivityTypes()
	if err != nil {
		return err
	}
	var names []string
	for name, _ := range ats {
		names = append(names, name)
	}
	sort.Strings(names)
	for _, name := range names {
		fmt.Printf("%s: %d\n", name, ats[name])
	}
	return nil
}

func visualisationCreate(c *cli.Context) error {
	client, err := newAuthenticatedDoaramaClient(c)
	if err != nil {
		return err
	}
	var as []*doarama.Activity
	for _, arg := range c.Args() {
		id64, err := strconv.ParseInt(arg, 10, 0)
		if err != nil {
			return err
		}
		a := client.Activity(int(id64))
		as = append(as, a)
	}
	v, err := client.CreateVisualisation(as)
	if err != nil {
		return err
	}
	fmt.Printf("VisualisationKey: %s\n", v.Key)
	return nil
}

func visualisationDelete(c *cli.Context) error {
	client, err := newAuthenticatedDoaramaClient(c)
	if err != nil {
		return err
	}
	for _, arg := range c.Args() {
		v := client.Visualisation(arg)
		if err := v.Delete(); err != nil {
			log.Print(err)
			continue
		}
	}
	return nil
}

func visualisationURL(c *cli.Context) error {
	client := newDoaramaClient(c)
	vuo := newVisualisationURLOptions(c)
	for _, arg := range c.Args() {
		v := client.Visualisation(arg)
		fmt.Printf("VisualisationURL: %s\n", v.URL(vuo))
	}
	return nil
}

func logError(f func(*cli.Context) error) func(*cli.Context) {
	return func(c *cli.Context) {
		if err := f(c); err != nil {
			log.Fatal(err)
		}
	}
}

func main() {
	app := cli.NewApp()
	app.Name = "doarama"
	app.Usage = "A command line interface to doarama.com"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "apiurl",
			Value:  doarama.API_URL,
			Usage:  "Doarama API URL",
			EnvVar: "DOARAMA_API_URL",
		},
		cli.StringFlag{
			Name:   "apikey",
			Usage:  "Doarama API key",
			EnvVar: "DOARAMA_API_KEY",
		},
		cli.StringFlag{
			Name:   "apiname",
			Usage:  "Doarama API name",
			EnvVar: "DOARAMA_API_NAME",
		},
		cli.StringFlag{
			Name:   "userid",
			Usage:  "Doarama user ID",
			EnvVar: "DOARAMA_USER_ID",
		},
		cli.StringFlag{
			Name:   "userkey",
			Usage:  "Doarama user key",
			EnvVar: "DOARAMA_USER_KEY",
		},
	}
	typeIdFlag := cli.IntFlag{
		Name:  "typeid",
		Usage: "type id",
	}
	nameFlag := cli.StringSliceFlag{
		Name:  "name",
		Usage: "name",
	}
	avatarFlag := cli.StringSliceFlag{
		Name:  "avatar",
		Usage: "avatar",
	}
	avatarBaseUrlFlag := cli.StringFlag{
		Name:  "avatarbaseurl",
		Usage: "avatar base URL",
	}
	fixedAspectFlag := cli.BoolTFlag{
		Name:  "fixedaspect",
		Usage: "fixed aspect",
	}
	minimalViewFlag := cli.BoolFlag{
		Name:  "minimalview",
		Usage: "minimal view",
	}
	dzmlFlag := cli.StringFlag{
		Name:  "dzml",
		Usage: "DZML",
	}
	app.Commands = []cli.Command{
		{
			Name:    "activity",
			Aliases: []string{"a"},
			Usage:   "Manages activities",
			Subcommands: []cli.Command{
				{
					Name:    "create",
					Aliases: []string{"c"},
					Usage:   "Creates an activity from one or more tracklogs",
					Action:  logError(activityCreate),
					Flags: []cli.Flag{
						typeIdFlag,
					},
				},
				{
					Name:    "delete",
					Aliases: []string{"d"},
					Usage:   "Deletes one or more activities by id",
					Action:  logError(activityDelete),
				},
			},
		},
		{
			Name:    "create",
			Aliases: []string{"c"},
			Usage:   "Creates a visualisation URL from one or more tracklogs",
			Action:  logError(create),
			Flags: []cli.Flag{
				typeIdFlag,
				nameFlag,
				avatarFlag,
				avatarBaseUrlFlag,
				fixedAspectFlag,
				minimalViewFlag,
				dzmlFlag,
			},
		},
		{
			Name:    "query-activity-types",
			Aliases: []string{"qat"},
			Usage:   "Queries activity types",
			Action:  logError(queryActivityTypes),
		},
		{
			Name:    "visualisation",
			Aliases: []string{"v"},
			Usage:   "Manages visualisations",
			Subcommands: []cli.Command{
				{
					Name:    "create",
					Aliases: []string{"c"},
					Usage:   "Creates a visualisation from a list of activities",
					Action:  logError(visualisationCreate),
				},
				{
					Name:    "delete",
					Aliases: []string{"d"},
					Usage:   "Deletes one or more visualisations by key",
					Action:  logError(visualisationDelete),
				},
				{
					Name:    "url",
					Aliases: []string{"u"},
					Usage:   "Creates a visualisation URL from a visualisation key",
					Action:  logError(visualisationURL),
					Flags: []cli.Flag{
						nameFlag,
						avatarFlag,
						avatarBaseUrlFlag,
						fixedAspectFlag,
						minimalViewFlag,
						dzmlFlag,
					},
				},
			},
		},
	}
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"html/template"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"strings"

	"github.com/steder/gophernaut"
)

// Event is basically just an enum
type Event int

// Events that can be generated by our child processes
const (
	Start Event = iota
	Shutdown
	PiningForTheFjords
)

// TODO look into "go generate stringer -type Event"
func (e Event) String() string {
	return fmt.Sprintf("Event(%d)", e)
}

var hostname = fmt.Sprintf("http://127.0.0.1:%d", 8080)
var executable = fmt.Sprintf("python -m SimpleHTTPServer %d", 8080)

var hostnames = []string{
	fmt.Sprintf("http://127.0.0.1:%d", 8080),
	fmt.Sprintf("http://127.0.0.1:%d", 8081),
}

var executables = []string{
	fmt.Sprintf("python -m SimpleHTTPServer %d", 8080),
	fmt.Sprintf("python -m SimpleHTTPServer %d", 8081),
}

func copyToLog(dst *log.Logger, src io.Reader) {
	scanner := bufio.NewScanner(src)
	for scanner.Scan() {
		dst.Print(scanner.Text())
	}
}

func startProcess(control <-chan Event, events chan<- Event, executable string) {
	procLog := log.New(os.Stdout, fmt.Sprintf("gopher-worker(%s) ", executable), log.Ldate|log.Ltime)

	commandParts := strings.Split(executable, " ")
	command := exec.Command(commandParts[0], commandParts[1:]...)
	log.Printf("Command: %v\n", command)

	stdout, err := command.StdoutPipe()
	if err != nil {
		procLog.Fatalln("Unable to connect to stdout from command...")
	}
	stderr, err := command.StderrPipe()
	if err != nil {
		procLog.Fatalln("Unable to connect to stderr from command...")
	}

	//go io.Copy(os.Stdout, stdout)
	//go io.Copy(os.Stderr, stderr)
	go copyToLog(procLog, stdout)
	go copyToLog(procLog, stderr)
	command.Start()

	for {
		_, ok := <-control
		if !ok {
			fmt.Println("Killing worker process after receiving close event.")
			command.Process.Kill()
			events <- Shutdown
			break
		}
	}
}

var requestCount = 0

func myHandler(w http.ResponseWriter, myReq *http.Request) {
	requestPath := myReq.URL.Path
	// TODO: multiprocess, pick one of n hostnames based on pool status
	hostname := hostnames[requestCount%2] // TODO get rid of this hard coded 2
	requestCount++
	targetURL, _ := url.Parse(hostname)
	director := func(req *http.Request) {
		targetQuery := targetURL.RawQuery
		req.URL.Scheme = targetURL.Scheme
		// TODO: adjust request host to assign the request to the appropriate child process
		req.URL.Host = targetURL.Host

		// clean up but preserve trailing slash:
		trailing := strings.HasSuffix(req.URL.Path, "/")
		req.URL.Path = path.Join(targetURL.Path, req.URL.Path)
		if trailing && !strings.HasSuffix(req.URL.Path, "/") {
			req.URL.Path += "/"
		}

		// preserve query string:
		if targetQuery == "" || req.URL.RawQuery == "" {
			req.URL.RawQuery = targetQuery + req.URL.RawQuery
		} else {
			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
		}
	}

	proxy := &httputil.ReverseProxy{Director: director}

	staticHandler := http.StripPrefix("/static", http.FileServer(http.Dir("static")))
	adminTemplate := template.Must(template.ParseFiles("templates/admin.html"))
	adminHandler := func(w http.ResponseWriter, req *http.Request) {
		adminTemplate.Execute(w, nil)
	}

	switch {
	case requestPath == "/admin":
		adminHandler(w, myReq)
		return
	case strings.HasPrefix(requestPath, "/static"):
		staticHandler.ServeHTTP(w, myReq)
		return
	}
	proxy.ServeHTTP(w, myReq)
}

func main() {
	log.SetPrefix("gophernaut ")
	log.SetFlags(log.Ldate | log.Ltime)
	c := gophernaut.ReadConfig()
	log.Printf("Host %s and Port %d\n", c.Host, c.Port)

	controlChannel := make(chan Event)
	eventsChannel := make(chan Event)

	// Handle signals to try to do a graceful shutdown:
	receivedSignals := make(chan os.Signal, 1)
	signal.Notify(receivedSignals, os.Interrupt) // , syscall.SIGINT, syscall.SIGTERM)
	go func() {
		for sig := range receivedSignals {
			fmt.Printf("Received signal, %s, shutting down workers...\n", sig)
			break
		}
		close(controlChannel)
		signal.Stop(receivedSignals)
	}()

	// Actually start some processes
	for _, executable := range executables {
		go startProcess(controlChannel, eventsChannel, executable)
	}

	// wait for child processes to exit before shutting down:
	processCount := 2 // TODO get rid of these hard coded 2s!
	stoppedCount := 0
	go func() {
		for event := range eventsChannel {
			if event == Shutdown {
				stoppedCount++
			}
			if processCount == stoppedCount {
				fmt.Printf("%d workers stopped, shutting down.\n", processCount)
				os.Exit(1)
			}
		}
	}()

	log.Printf("Gophernaut is gopher launch!\n")
	// TODO: our own ReverseProxy implementation of at least, ServeHTTP so that we can
	// monitor the response codes to track successes and failures
	log.Fatal(http.ListenAndServe(":8483", http.HandlerFunc(myHandler)))
}
<commit_msg>a little less hardcoding<commit_after>package main

import (
	"bufio"
	"fmt"
	"html/template"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"strings"

	"github.com/steder/gophernaut"
)

// Event is basically just an enum
type Event int

// Events that can be generated by our child processes
const (
	Start Event = iota
	Shutdown
	PiningForTheFjords
)

// TODO look into "go generate stringer -type Event"
func (e Event) String() string {
	return fmt.Sprintf("Event(%d)", e)
}

var hostnames = []string{
	fmt.Sprintf("http://127.0.0.1:%d", 8080),
	fmt.Sprintf("http://127.0.0.1:%d", 8081),
}

var executables = []string{
	fmt.Sprintf("python -m SimpleHTTPServer %d", 8080),
	fmt.Sprintf("python -m SimpleHTTPServer %d", 8081),
}

func copyToLog(dst *log.Logger, src io.Reader) {
	scanner := bufio.NewScanner(src)
	for scanner.Scan() {
		dst.Print(scanner.Text())
	}
}

func startProcess(control <-chan Event, events chan<- Event, executable string) {
	procLog := log.New(os.Stdout, fmt.Sprintf("gopher-worker(%s) ", executable), log.Ldate|log.Ltime)

	commandParts := strings.Split(executable, " ")
	command := exec.Command(commandParts[0], commandParts[1:]...)
	log.Printf("Command: %v\n", command)

	stdout, err := command.StdoutPipe()
	if err != nil {
		procLog.Fatalln("Unable to connect to stdout from command...")
	}
	stderr, err := command.StderrPipe()
	if err != nil {
		procLog.Fatalln("Unable to connect to stderr from command...")
	}

	//go io.Copy(os.Stdout, stdout)
	//go io.Copy(os.Stderr, stderr)
	go copyToLog(procLog, stdout)
	go copyToLog(procLog, stderr)
	command.Start()

	for {
		_, ok := <-control
		if !ok {
			fmt.Println("Killing worker process after receiving close event.")
			command.Process.Kill()
			events <- Shutdown
			break
		}
	}
}

var requestCount = 0

func myHandler(w http.ResponseWriter, myReq *http.Request) {
	requestPath := myReq.URL.Path
	// TODO: multiprocess, pick one of n hostnames based on pool status
	hostname := hostnames[requestCount%len(hostnames)]
	requestCount++
	targetURL, _ := url.Parse(hostname)
	director := func(req *http.Request) {
		targetQuery := targetURL.RawQuery
		req.URL.Scheme = targetURL.Scheme
		// TODO: adjust request host to assign the request to the appropriate child process
		req.URL.Host = targetURL.Host

		// clean up but preserve trailing slash:
		trailing := strings.HasSuffix(req.URL.Path, "/")
		req.URL.Path = path.Join(targetURL.Path, req.URL.Path)
		if trailing && !strings.HasSuffix(req.URL.Path, "/") {
			req.URL.Path += "/"
		}

		// preserve query string:
		if targetQuery == "" || req.URL.RawQuery == "" {
			req.URL.RawQuery = targetQuery + req.URL.RawQuery
		} else {
			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
		}
	}

	proxy := &httputil.ReverseProxy{Director: director}

	staticHandler := http.StripPrefix("/static", http.FileServer(http.Dir("static")))
	adminTemplate := template.Must(template.ParseFiles("templates/admin.html"))
	adminHandler := func(w http.ResponseWriter, req *http.Request) {
		adminTemplate.Execute(w, nil)
	}

	switch {
	case requestPath == "/admin":
		adminHandler(w, myReq)
		return
	case strings.HasPrefix(requestPath, "/static"):
		staticHandler.ServeHTTP(w, myReq)
		return
	}
	proxy.ServeHTTP(w, myReq)
}

func main() {
	log.SetPrefix("gophernaut ")
	log.SetFlags(log.Ldate | log.Ltime)
	c := gophernaut.ReadConfig()
	log.Printf("Host %s and Port %d\n", c.Host, c.Port)

	controlChannel := make(chan Event)
	eventsChannel := make(chan Event)

	// Handle signals to try to do a graceful shutdown:
	receivedSignals := make(chan os.Signal, 1)
	signal.Notify(receivedSignals, os.Interrupt) // , syscall.SIGINT, syscall.SIGTERM)
	go func() {
		for sig := range receivedSignals {
			fmt.Printf("Received signal, %s, shutting down workers...\n", sig)
			break
		}
		close(controlChannel)
		signal.Stop(receivedSignals)
	}()

	// Actually start some processes
	for _, executable := range executables {
		go startProcess(controlChannel, eventsChannel, executable)
	}

	// wait for child processes to exit before shutting down:
	processCount := len(executables)
	stoppedCount := 0
	go func() {
		for event := range eventsChannel {
			if event == Shutdown {
				stoppedCount++
			}
			if processCount == stoppedCount {
				fmt.Printf("%d workers stopped, shutting down.\n", processCount)
				os.Exit(1)
			}
		}
	}()

	log.Printf("Gophernaut is gopher launch!\n")
	// TODO: our own ReverseProxy implementation of at least, ServeHTTP so that we can
	// monitor the response codes to track successes and failures
	log.Fatal(http.ListenAndServe(":8483", http.HandlerFunc(myHandler)))
}
<|endoftext|>


<commit_before>/*
Copyright The Helm Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"bytes"
	"fmt"
	"io"

	"github.com/spf13/cobra"
)

const completionDesc = `
Generate autocompletions script for Helm for the specified shell (bash or zsh).

This command can generate shell autocompletions. e.g.

	$ helm completion bash

Can be sourced as such

	$ source <(helm completion bash)
`

var (
	completionShells = map[string]func(out io.Writer, cmd *cobra.Command) error{
		"bash": runCompletionBash,
		"zsh":  runCompletionZsh,
	}
)

func newCompletionCmd(out io.Writer) *cobra.Command {
	shells := []string{}
	for s := range completionShells {
		shells = append(shells, s)
	}

	cmd := &cobra.Command{
		Use:   "completion SHELL",
		Short: "Generate autocompletions script for the specified shell (bash or zsh)",
		Long:  completionDesc,
		RunE: func(cmd *cobra.Command, args []string) error {
			return runCompletion(out, cmd, args)
		},
		ValidArgs: shells,
	}

	return cmd
}

func runCompletion(out io.Writer, cmd *cobra.Command, args []string) error {
	if len(args) == 0 {
		return fmt.Errorf("shell not specified")
	}
	if len(args) > 1 {
		return fmt.Errorf("too many arguments, expected only the shell type")
	}
	run, found := completionShells[args[0]]
	if !found {
		return fmt.Errorf("unsupported shell type %q", args[0])
	}

	return run(out, cmd)
}

func runCompletionBash(out io.Writer, cmd *cobra.Command) error {
	return cmd.Root().GenBashCompletion(out)
}

func runCompletionZsh(out io.Writer, cmd *cobra.Command) error {
	zshInitialization := `
#compdef helm

__helm_bash_source() {
	alias shopt=':'
	alias _expand=_bash_expand
	alias _complete=_bash_comp
	emulate -L sh
	setopt kshglob noshglob braceexpand
	source "$@"
}
__helm_type() {
	# -t is not supported by zsh
	if [ "$1" == "-t" ]; then
		shift
		# fake Bash 4 to disable "complete -o nospace". Instead
		# "compopt +-o nospace" is used in the code to toggle trailing
		# spaces. We don't support that, but leave trailing spaces on
		# all the time
		if [ "$1" = "__helm_compopt" ]; then
			echo builtin
			return 0
		fi
	fi
	type "$@"
}
__helm_compgen() {
	local completions w
	completions=( $(compgen "$@") ) || return $?
	# filter by given word as prefix
	while [[ "$1" = -* && "$1" != -- ]]; do
		shift
		shift
	done
	if [[ "$1" == -- ]]; then
		shift
	fi
	for w in "${completions[@]}"; do
		if [[ "${w}" = "$1"* ]]; then
			echo "${w}"
		fi
	done
}
__helm_compopt() {
	true # don't do anything. Not supported by bashcompinit in zsh
}
__helm_declare() {
	if [ "$1" == "-F" ]; then
		whence -w "$@"
	else
		builtin declare "$@"
	fi
}
__helm_ltrim_colon_completions()
{
	if [[ "$1" == *:* && "$COMP_WORDBREAKS" == *:* ]]; then
		# Remove colon-word prefix from COMPREPLY items
		local colon_word=${1%${1##*:}}
		local i=${#COMPREPLY[*]}
		while [[ $((--i)) -ge 0 ]]; do
			COMPREPLY[$i]=${COMPREPLY[$i]#"$colon_word"}
		done
	fi
}
__helm_get_comp_words_by_ref() {
	cur="${COMP_WORDS[COMP_CWORD]}"
	prev="${COMP_WORDS[${COMP_CWORD}-1]}"
	words=("${COMP_WORDS[@]}")
	cword=("${COMP_CWORD[@]}")
}
__helm_filedir() {
	local RET OLD_IFS w qw
	__debug "_filedir $@ cur=$cur"
	if [[ "$1" = \~* ]]; then
		# somehow does not work. Maybe, zsh does not call this at all
		eval echo "$1"
		return 0
	fi
	OLD_IFS="$IFS"
	IFS=$'\n'
	if [ "$1" = "-d" ]; then
		shift
		RET=( $(compgen -d) )
	else
		RET=( $(compgen -f) )
	fi
	IFS="$OLD_IFS"
	IFS="," __debug "RET=${RET[@]} len=${#RET[@]}"
	for w in ${RET[@]}; do
		if [[ ! "${w}" = "${cur}"* ]]; then
			continue
		fi
		if eval "[[ \"\${w}\" = *.$1 || -d \"\${w}\" ]]"; then
			qw="$(__helm_quote "${w}")"
			if [ -d "${w}" ]; then
				COMPREPLY+=("${qw}/")
			else
				COMPREPLY+=("${qw}")
			fi
		fi
	done
}
__helm_quote() {
	if [[ $1 == \'* || $1 == \"* ]]; then
		# Leave out first character
		printf %q "${1:1}"
	else
		printf %q "$1"
	fi
}
autoload -U +X bashcompinit && bashcompinit
# use word boundary patterns for BSD or GNU sed
LWORD='[[:<:]]'
RWORD='[[:>:]]'
if sed --help 2>&1 | grep -q GNU; then
	LWORD='\<'
	RWORD='\>'
fi
__helm_convert_bash_to_zsh() {
	sed \
	-e 's/declare -F/whence -w/' \
	-e 's/_get_comp_words_by_ref "\$@"/_get_comp_words_by_ref "\$*"/' \
	-e 's/local \([a-zA-Z0-9_]*\)=/local \1; \1=/' \
	-e 's/flags+=("\(--.*\)=")/flags+=("\1"); two_word_flags+=("\1")/' \
	-e 's/must_have_one_flag+=("\(--.*\)=")/must_have_one_flag+=("\1")/' \
	-e "s/${LWORD}_filedir${RWORD}/__helm_filedir/g" \
	-e "s/${LWORD}_get_comp_words_by_ref${RWORD}/__helm_get_comp_words_by_ref/g" \
	-e "s/${LWORD}__ltrim_colon_completions${RWORD}/__helm_ltrim_colon_completions/g" \
	-e "s/${LWORD}compgen${RWORD}/__helm_compgen/g" \
	-e "s/${LWORD}compopt${RWORD}/__helm_compopt/g" \
	-e "s/${LWORD}declare${RWORD}/__helm_declare/g" \
	-e "s/\\\$(type${RWORD}/\$(__helm_type/g" \
	<<'BASH_COMPLETION_EOF'
`
	out.Write([]byte(zshInitialization))

	buf := new(bytes.Buffer)
	cmd.Root().GenBashCompletion(buf)
	out.Write(buf.Bytes())

	zshTail := `
BASH_COMPLETION_EOF
}
__helm_bash_source <(__helm_convert_bash_to_zsh)
`
	out.Write([]byte(zshTail))
	return nil
}
<commit_msg>Remove newline at the start of zsh completion file (#4851)<commit_after>/*
Copyright The Helm Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"bytes"
	"fmt"
	"io"

	"github.com/spf13/cobra"
)

const completionDesc = `
Generate autocompletions script for Helm for the specified shell (bash or zsh).

This command can generate shell autocompletions. e.g.

	$ helm completion bash

Can be sourced as such

	$ source <(helm completion bash)
`

var (
	completionShells = map[string]func(out io.Writer, cmd *cobra.Command) error{
		"bash": runCompletionBash,
		"zsh":  runCompletionZsh,
	}
)

func newCompletionCmd(out io.Writer) *cobra.Command {
	shells := []string{}
	for s := range completionShells {
		shells = append(shells, s)
	}

	cmd := &cobra.Command{
		Use:   "completion SHELL",
		Short: "Generate autocompletions script for the specified shell (bash or zsh)",
		Long:  completionDesc,
		RunE: func(cmd *cobra.Command, args []string) error {
			return runCompletion(out, cmd, args)
		},
		ValidArgs: shells,
	}

	return cmd
}

func runCompletion(out io.Writer, cmd *cobra.Command, args []string) error {
	if len(args) == 0 {
		return fmt.Errorf("shell not specified")
	}
	if len(args) > 1 {
		return fmt.Errorf("too many arguments, expected only the shell type")
	}
	run, found := completionShells[args[0]]
	if !found {
		return fmt.Errorf("unsupported shell type %q", args[0])
	}

	return run(out, cmd)
}

func runCompletionBash(out io.Writer, cmd *cobra.Command) error {
	return cmd.Root().GenBashCompletion(out)
}

func runCompletionZsh(out io.Writer, cmd *cobra.Command) error {
	zshInitialization := `#compdef helm

__helm_bash_source() {
	alias shopt=':'
	alias _expand=_bash_expand
	alias _complete=_bash_comp
	emulate -L sh
	setopt kshglob noshglob braceexpand
	source "$@"
}
__helm_type() {
	# -t is not supported by zsh
	if [ "$1" == "-t" ]; then
		shift
		# fake Bash 4 to disable "complete -o nospace". Instead
		# "compopt +-o nospace" is used in the code to toggle trailing
		# spaces. We don't support that, but leave trailing spaces on
		# all the time
		if [ "$1" = "__helm_compopt" ]; then
			echo builtin
			return 0
		fi
	fi
	type "$@"
}
__helm_compgen() {
	local completions w
	completions=( $(compgen "$@") ) || return $?
	# filter by given word as prefix
	while [[ "$1" = -* && "$1" != -- ]]; do
		shift
		shift
	done
	if [[ "$1" == -- ]]; then
		shift
	fi
	for w in "${completions[@]}"; do
		if [[ "${w}" = "$1"* ]]; then
			echo "${w}"
		fi
	done
}
__helm_compopt() {
	true # don't do anything. Not supported by bashcompinit in zsh
}
__helm_declare() {
	if [ "$1" == "-F" ]; then
		whence -w "$@"
	else
		builtin declare "$@"
	fi
}
__helm_ltrim_colon_completions()
{
	if [[ "$1" == *:* && "$COMP_WORDBREAKS" == *:* ]]; then
		# Remove colon-word prefix from COMPREPLY items
		local colon_word=${1%${1##*:}}
		local i=${#COMPREPLY[*]}
		while [[ $((--i)) -ge 0 ]]; do
			COMPREPLY[$i]=${COMPREPLY[$i]#"$colon_word"}
		done
	fi
}
__helm_get_comp_words_by_ref() {
	cur="${COMP_WORDS[COMP_CWORD]}"
	prev="${COMP_WORDS[${COMP_CWORD}-1]}"
	words=("${COMP_WORDS[@]}")
	cword=("${COMP_CWORD[@]}")
}
__helm_filedir() {
	local RET OLD_IFS w qw
	__debug "_filedir $@ cur=$cur"
	if [[ "$1" = \~* ]]; then
		# somehow does not work. Maybe, zsh does not call this at all
		eval echo "$1"
		return 0
	fi
	OLD_IFS="$IFS"
	IFS=$'\n'
	if [ "$1" = "-d" ]; then
		shift
		RET=( $(compgen -d) )
	else
		RET=( $(compgen -f) )
	fi
	IFS="$OLD_IFS"
	IFS="," __debug "RET=${RET[@]} len=${#RET[@]}"
	for w in ${RET[@]}; do
		if [[ ! "${w}" = "${cur}"* ]]; then
			continue
		fi
		if eval "[[ \"\${w}\" = *.$1 || -d \"\${w}\" ]]"; then
			qw="$(__helm_quote "${w}")"
			if [ -d "${w}" ]; then
				COMPREPLY+=("${qw}/")
			else
				COMPREPLY+=("${qw}")
			fi
		fi
	done
}
__helm_quote() {
	if [[ $1 == \'* || $1 == \"* ]]; then
		# Leave out first character
		printf %q "${1:1}"
	else
		printf %q "$1"
	fi
}
autoload -U +X bashcompinit && bashcompinit
# use word boundary patterns for BSD or GNU sed
LWORD='[[:<:]]'
RWORD='[[:>:]]'
if sed --help 2>&1 | grep -q GNU; then
	LWORD='\<'
	RWORD='\>'
fi
__helm_convert_bash_to_zsh() {
	sed \
	-e 's/declare -F/whence -w/' \
	-e 's/_get_comp_words_by_ref "\$@"/_get_comp_words_by_ref "\$*"/' \
	-e 's/local \([a-zA-Z0-9_]*\)=/local \1; \1=/' \
	-e 's/flags+=("\(--.*\)=")/flags+=("\1"); two_word_flags+=("\1")/' \
	-e 's/must_have_one_flag+=("\(--.*\)=")/must_have_one_flag+=("\1")/' \
	-e "s/${LWORD}_filedir${RWORD}/__helm_filedir/g" \
	-e "s/${LWORD}_get_comp_words_by_ref${RWORD}/__helm_get_comp_words_by_ref/g" \
	-e "s/${LWORD}__ltrim_colon_completions${RWORD}/__helm_ltrim_colon_completions/g" \
	-e "s/${LWORD}compgen${RWORD}/__helm_compgen/g" \
	-e "s/${LWORD}compopt${RWORD}/__helm_compopt/g" \
	-e "s/${LWORD}declare${RWORD}/__helm_declare/g" \
	-e "s/\\\$(type${RWORD}/\$(__helm_type/g" \
	<<'BASH_COMPLETION_EOF'
`
	out.Write([]byte(zshInitialization))

	buf := new(bytes.Buffer)
	cmd.Root().GenBashCompletion(buf)
	out.Write(buf.Bytes())

	zshTail := `
BASH_COMPLETION_EOF
}
__helm_bash_source <(__helm_convert_bash_to_zsh)
`
	out.Write([]byte(zshTail))
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright IBM Corp. All Rights Reserved.

SPDX-License-Identifier: Apache-2.0
*/

package container

import (
	"fmt"
	"io"
	"sync"

	"golang.org/x/net/context"

	"github.com/hyperledger/fabric/common/flogging"
	"github.com/hyperledger/fabric/core/chaincode/platforms"
	"github.com/hyperledger/fabric/core/container/ccintf"
	pb "github.com/hyperledger/fabric/protos/peer"
)

type VMProvider interface {
	NewVM() VM
}

type Builder interface {
	Build() (io.Reader, error)
}

//VM is an abstract virtual image for supporting arbitrary virual machines
type VM interface {
	Start(ctxt context.Context, ccid ccintf.CCID, args []string, env []string, filesToUpload map[string][]byte, builder Builder) error
	Stop(ctxt context.Context, ccid ccintf.CCID, timeout uint, dontkill bool, dontremove bool) error
}

type refCountedLock struct {
	refCount int
	lock     *sync.RWMutex
}

//VMController - manages VMs
//   . abstract construction of different types of VMs (we only care about Docker for now)
//   . manage lifecycle of VM (start with build, start, stop ...
//     eventually probably need fine grained management)
type VMController struct {
	sync.RWMutex
	containerLocks map[string]*refCountedLock
	vmProviders    map[string]VMProvider
}

var vmLogger = flogging.MustGetLogger("container")

// NewVMController creates a new instance of VMController
func NewVMController(vmProviders map[string]VMProvider) *VMController {
	return &VMController{
		containerLocks: make(map[string]*refCountedLock),
		vmProviders:    vmProviders,
	}
}

func (vmc *VMController) newVM(typ string) VM {
	v, ok := vmc.vmProviders[typ]
	if !ok {
		vmLogger.Panicf("Programming error: unsupported VM type: %s", typ)
	}
	return v.NewVM()
}

func (vmc *VMController) lockContainer(id string) {
	//get the container lock under global lock
	vmc.Lock()
	var refLck *refCountedLock
	var ok bool
	if refLck, ok = vmc.containerLocks[id]; !ok {
		refLck = &refCountedLock{refCount: 1, lock: &sync.RWMutex{}}
		vmc.containerLocks[id] = refLck
	} else {
		refLck.refCount++
		vmLogger.Debugf("refcount %d (%s)", refLck.refCount, id)
	}
	vmc.Unlock()
	vmLogger.Debugf("waiting for container(%s) lock", id)
	refLck.lock.Lock()
	vmLogger.Debugf("got container (%s) lock", id)
}

func (vmc *VMController) unlockContainer(id string) {
	vmc.Lock()
	if refLck, ok := vmc.containerLocks[id]; ok {
		if refLck.refCount <= 0 {
			panic("refcnt <= 0")
		}
		refLck.lock.Unlock()
		if refLck.refCount--; refLck.refCount == 0 {
			vmLogger.Debugf("container lock deleted(%s)", id)
			delete(vmc.containerLocks, id)
		}
	} else {
		vmLogger.Debugf("no lock to unlock(%s)!!", id)
	}
	vmc.Unlock()
}

//VMCReq - all requests should implement this interface.
//The context should be passed and tested at each layer till we stop
//note that we'd stop on the first method on the stack that does not
//take context
type VMCReq interface {
	Do(ctxt context.Context, v VM) error
	GetCCID() ccintf.CCID
}

//StartContainerReq - properties for starting a container.
type StartContainerReq struct {
	ccintf.CCID
	Builder       Builder
	Args          []string
	Env           []string
	FilesToUpload map[string][]byte
}

// PlatformBuilder implements the Build interface using
// the platforms package GenerateDockerBuild function.
// XXX This is a pretty awkward spot for the builder, it should
// really probably be pushed into the dockercontroller, as it only
// builds docker images, but, doing so would require contaminating
// the dockercontroller package with the CDS, which is also
// undesirable.
type PlatformBuilder struct {
	DeploymentSpec   *pb.ChaincodeDeploymentSpec
	PlatformRegistry *platforms.Registry
}

// Build a tar stream based on the CDS
func (b *PlatformBuilder) Build() (io.Reader, error) {
	return b.PlatformRegistry.GenerateDockerBuild(b.DeploymentSpec)
}

func (si StartContainerReq) Do(ctxt context.Context, v VM) error {
	return v.Start(ctxt, si.CCID, si.Args, si.Env, si.FilesToUpload, si.Builder)
}

func (si StartContainerReq) GetCCID() ccintf.CCID {
	return si.CCID
}

//StopContainerReq - properties for stopping a container.
type StopContainerReq struct {
	ccintf.CCID
	Timeout uint
	//by default we will kill the container after stopping
	Dontkill bool
	//by default we will remove the container after killing
	Dontremove bool
}

func (si StopContainerReq) Do(ctxt context.Context, v VM) error {
	return v.Stop(ctxt, si.CCID, si.Timeout, si.Dontkill, si.Dontremove)
}

func (si StopContainerReq) GetCCID() ccintf.CCID {
	return si.CCID
}

//Process should be used as follows
//   . construct a context
//   . construct req of the right type (e.g., CreateImageReq)
//   . call it in a go routine
//   . process response in the go routing
//context can be cancelled. VMCProcess will try to cancel calling functions if it can
//For instance docker clients api's such as BuildImage are not cancelable.
//In all cases VMCProcess will wait for the called go routine to return
func (vmc *VMController) Process(ctxt context.Context, vmtype string, req VMCReq) error {
	v := vmc.newVM(vmtype)

	c := make(chan error)
	go func() {
		ccid := req.GetCCID()
		id := ccid.GetName()
		vmc.lockContainer(id)
		err := req.Do(ctxt, v)
		vmc.unlockContainer(id)
		c <- err
	}()

	select {
	case err := <-c:
		return err
	case <-ctxt.Done():
		//TODO cancel req.do ... (needed) ?
		// XXX This logic doesn't make much sense, why return the context error if it's canceled,
		// but still wait for the request to complete, and ignore its error
		<-c
		return ctxt.Err()
	}
}

// GetChaincodePackageBytes creates bytes for docker container generation using the supplied chaincode specification
func GetChaincodePackageBytes(pr *platforms.Registry, spec *pb.ChaincodeSpec) ([]byte, error) {
	if spec == nil || spec.ChaincodeId == nil {
		return nil, fmt.Errorf("invalid chaincode spec")
	}

	return pr.GetDeploymentPayload(spec)
}
<commit_msg>[FAB-10899] ignore context in VMController<commit_after>/*
Copyright IBM Corp. All Rights Reserved.

SPDX-License-Identifier: Apache-2.0
*/

package container

import (
	"fmt"
	"io"
	"sync"

	"golang.org/x/net/context"

	"github.com/hyperledger/fabric/common/flogging"
	"github.com/hyperledger/fabric/core/chaincode/platforms"
	"github.com/hyperledger/fabric/core/container/ccintf"
	pb "github.com/hyperledger/fabric/protos/peer"
)

type VMProvider interface {
	NewVM() VM
}

type Builder interface {
	Build() (io.Reader, error)
}

//VM is an abstract virtual image for supporting arbitrary virual machines
type VM interface {
	Start(ctxt context.Context, ccid ccintf.CCID, args []string, env []string, filesToUpload map[string][]byte, builder Builder) error
	Stop(ctxt context.Context, ccid ccintf.CCID, timeout uint, dontkill bool, dontremove bool) error
}

type refCountedLock struct {
	refCount int
	lock     *sync.RWMutex
}

//VMController - manages VMs
//   . abstract construction of different types of VMs (we only care about Docker for now)
//   . manage lifecycle of VM (start with build, start, stop ...
//     eventually probably need fine grained management)
type VMController struct {
	sync.RWMutex
	containerLocks map[string]*refCountedLock
	vmProviders    map[string]VMProvider
}

var vmLogger = flogging.MustGetLogger("container")

// NewVMController creates a new instance of VMController
func NewVMController(vmProviders map[string]VMProvider) *VMController {
	return &VMController{
		containerLocks: make(map[string]*refCountedLock),
		vmProviders:    vmProviders,
	}
}

func (vmc *VMController) newVM(typ string) VM {
	v, ok := vmc.vmProviders[typ]
	if !ok {
		vmLogger.Panicf("Programming error: unsupported VM type: %s", typ)
	}
	return v.NewVM()
}

func (vmc *VMController) lockContainer(id string) {
	//get the container lock under global lock
	vmc.Lock()
	var refLck *refCountedLock
	var ok bool
	if refLck, ok = vmc.containerLocks[id]; !ok {
		refLck = &refCountedLock{refCount: 1, lock: &sync.RWMutex{}}
		vmc.containerLocks[id] = refLck
	} else {
		refLck.refCount++
		vmLogger.Debugf("refcount %d (%s)", refLck.refCount, id)
	}
	vmc.Unlock()
	vmLogger.Debugf("waiting for container(%s) lock", id)
	refLck.lock.Lock()
	vmLogger.Debugf("got container (%s) lock", id)
}

func (vmc *VMController) unlockContainer(id string) {
	vmc.Lock()
	if refLck, ok := vmc.containerLocks[id]; ok {
		if refLck.refCount <= 0 {
			panic("refcnt <= 0")
		}
		refLck.lock.Unlock()
		if refLck.refCount--; refLck.refCount == 0 {
			vmLogger.Debugf("container lock deleted(%s)", id)
			delete(vmc.containerLocks, id)
		}
	} else {
		vmLogger.Debugf("no lock to unlock(%s)!!", id)
	}
	vmc.Unlock()
}

//VMCReq - all requests should implement this interface.
//The context should be passed and tested at each layer till we stop
//note that we'd stop on the first method on the stack that does not
//take context
type VMCReq interface {
	Do(ctxt context.Context, v VM) error
	GetCCID() ccintf.CCID
}

//StartContainerReq - properties for starting a container.
type StartContainerReq struct {
	ccintf.CCID
	Builder       Builder
	Args          []string
	Env           []string
	FilesToUpload map[string][]byte
}

// PlatformBuilder implements the Build interface using
// the platforms package GenerateDockerBuild function.
// XXX This is a pretty awkward spot for the builder, it should
// really probably be pushed into the dockercontroller, as it only
// builds docker images, but, doing so would require contaminating
// the dockercontroller package with the CDS, which is also
// undesirable.
type PlatformBuilder struct {
	DeploymentSpec   *pb.ChaincodeDeploymentSpec
	PlatformRegistry *platforms.Registry
}

// Build a tar stream based on the CDS
func (b *PlatformBuilder) Build() (io.Reader, error) {
	return b.PlatformRegistry.GenerateDockerBuild(b.DeploymentSpec)
}

func (si StartContainerReq) Do(ctxt context.Context, v VM) error {
	return v.Start(ctxt, si.CCID, si.Args, si.Env, si.FilesToUpload, si.Builder)
}

func (si StartContainerReq) GetCCID() ccintf.CCID {
	return si.CCID
}

//StopContainerReq - properties for stopping a container.
type StopContainerReq struct {
	ccintf.CCID
	Timeout uint
	//by default we will kill the container after stopping
	Dontkill bool
	//by default we will remove the container after killing
	Dontremove bool
}

func (si StopContainerReq) Do(ctxt context.Context, v VM) error {
	return v.Stop(ctxt, si.CCID, si.Timeout, si.Dontkill, si.Dontremove)
}

func (si StopContainerReq) GetCCID() ccintf.CCID {
	return si.CCID
}

func (vmc *VMController) Process(ctxt context.Context, vmtype string, req VMCReq) error {
	v := vmc.newVM(vmtype)
	ccid := req.GetCCID()
	id := ccid.GetName()

	vmc.lockContainer(id)
	defer vmc.unlockContainer(id)
	return req.Do(ctxt, v)
}

// GetChaincodePackageBytes creates bytes for docker container generation using the supplied chaincode specification
func GetChaincodePackageBytes(pr *platforms.Registry, spec *pb.ChaincodeSpec) ([]byte, error) {
	if spec == nil || spec.ChaincodeId == nil {
		return nil, fmt.Errorf("invalid chaincode spec")
	}

	return pr.GetDeploymentPayload(spec)
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"fmt"
	"strings"

	"launchpad.net/gnuflag"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/environs/manual"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/juju"
	"launchpad.net/juju-core/names"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api/params"
)

// sshHostPrefix is the prefix for a machine to be "manually provisioned".
const sshHostPrefix = "ssh:"

var addMachineDoc = `

If no container is specified, a new machine will be
provisioned.  If a container is specified, a new machine will be provisioned
with that container.

To add a container to an existing machine, use the <container>:<machinenumber>
format.

When adding a new machine, you may specify constraints for the machine to be
provisioned.  Constraints cannot be combined with deploying a container to an
existing machine.

Currently, the only supported container type is lxc.

Machines are created in a clean state and ready to have units deployed.

This command also supports manual provisioning of existing machines via SSH. The
target machine must be able to communicate with the API server, and be able to
access the environment storage.

Examples:
   juju add-machine                      (starts a new machine)
   juju add-machine lxc                  (starts a new machine with an lxc container)
   juju add-machine lxc:4                (starts a new lxc container on machine 4)
   juju add-machine --constraints mem=8G (starts a machine with at least 8GB RAM)

See Also:
   juju help constraints
`

// AddMachineCommand starts a new machine and registers it in the environment.
type AddMachineCommand struct {
	cmd.EnvCommandBase
	// If specified, use this series, else use the environment default-series
	Series string
	// If specified, these constraints are merged with those already in the environment.
	Constraints   constraints.Value
	MachineId     string
	ContainerType instance.ContainerType
	SSHHost       string
}

func (c *AddMachineCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "add-machine",
		Args:    "[<container>:machine | <container> | ssh:[user@]host]",
		Purpose: "start a new, empty machine and optionally a container, or add a container to a machine",
		Doc:     addMachineDoc,
	}
}

func (c *AddMachineCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	f.StringVar(&c.Series, "series", "", "the charm series")
	f.Var(constraints.ConstraintsValue{&c.Constraints}, "constraints", "additional machine constraints")
}

func (c *AddMachineCommand) Init(args []string) error {
	if c.Constraints.Container != nil {
		return fmt.Errorf("container constraint %q not allowed when adding a machine", *c.Constraints.Container)
	}
	containerSpec, err := cmd.ZeroOrOneArgs(args)
	if err != nil {
		return err
	}
	if containerSpec == "" {
		return nil
	}
	if strings.HasPrefix(containerSpec, sshHostPrefix) {
		c.SSHHost = containerSpec[len(sshHostPrefix):]
	} else {
		// container arg can either be 'type:machine' or 'type'
		if c.ContainerType, err = instance.ParseContainerType(containerSpec); err != nil {
			if names.IsMachine(containerSpec) || !cmd.IsMachineOrNewContainer(containerSpec) {
				return fmt.Errorf("malformed container argument %q", containerSpec)
			}
			sep := strings.Index(containerSpec, ":")
			c.MachineId = containerSpec[sep+1:]
			c.ContainerType, err = instance.ParseContainerType(containerSpec[:sep])
		}
	}
	return err
}

// addMachine1dot16 runs Client.AddMachines using a direct DB connection to maintain
// compatibility with an API server running 1.16 or older (when AddMachines
// was not available). This fallback can be removed when we no longer maintain
// 1.16 compatibility.
// This was copied directly from the code in AddMachineCommand.Run in 1.16
func (c *AddMachineCommand) addMachine1dot16() (string, error) {
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return "", err
	}
	defer conn.Close()

	series := c.Series
	if series == "" {
		conf, err := conn.State.EnvironConfig()
		if err != nil {
			return "", err
		}
		series = conf.DefaultSeries()
	}
	template := state.MachineTemplate{
		Series:      series,
		Constraints: c.Constraints,
		Jobs:        []state.MachineJob{state.JobHostUnits},
	}
	var m *state.Machine
	switch {
	case c.ContainerType == "":
		m, err = conn.State.AddOneMachine(template)
	case c.MachineId != "":
		m, err = conn.State.AddMachineInsideMachine(template, c.MachineId, c.ContainerType)
	default:
		m, err = conn.State.AddMachineInsideNewMachine(template, template, c.ContainerType)
	}
	if err != nil {
		return "", err
	}
	return m.String(), err
}

func (c *AddMachineCommand) Run(ctx *cmd.Context) error {
	if c.SSHHost != "" {
		args := manual.ProvisionMachineArgs{
			Host:    c.SSHHost,
			EnvName: c.EnvName,
			Stdin:   ctx.Stdin,
			Stdout:  ctx.Stdout,
			Stderr:  ctx.Stderr,
		}
		_, err := manual.ProvisionMachine(args)
		return err
	}

	client, err := juju.NewAPIClientFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer client.Close()

	machineParams := params.AddMachineParams{
		ParentId:      c.MachineId,
		ContainerType: c.ContainerType,
		Series:        c.Series,
		Constraints:   c.Constraints,
		Jobs:          []params.MachineJob{params.JobHostUnits},
	}
	results, err := client.AddMachines([]params.AddMachineParams{machineParams})
	var machineId string
	if params.IsCodeNotImplemented(err) {
		logger.Infof("AddMachines not supported by the API server, " +
			"falling back to 1.16 compatibility mode (direct DB access)")
		machineId, err = c.addMachine1dot16()
	} else if err != nil {
		return err
	} else {
		// Currently, only one machine is added, but in future there may be several added in one call.
		machineInfo := results[0]
		var machineErr *params.Error
		machineId, machineErr = machineInfo.Machine, machineInfo.Error
		if machineErr != nil {
			err = machineErr
		}
	}
	if err != nil {
		return err
	}
	if c.ContainerType == "" {
		logger.Infof("created machine %v", machineId)
	} else {
		logger.Infof("created %q container on machine %v", c.ContainerType, machineId)
	}
	return nil
}
<commit_msg>Added ssh: to add-machine help<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"fmt"
	"strings"

	"launchpad.net/gnuflag"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/environs/manual"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/juju"
	"launchpad.net/juju-core/names"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api/params"
)

// sshHostPrefix is the prefix for a machine to be "manually provisioned".
const sshHostPrefix = "ssh:"

var addMachineDoc = `

If no container is specified, a new machine will be
provisioned.  If a container is specified, a new machine will be provisioned
with that container.

To add a container to an existing machine, use the <container>:<machinenumber>
format.

When adding a new machine, you may specify constraints for the machine to be
provisioned.  Constraints cannot be combined with deploying a container to an
existing machine.

Currently, the only supported container type is lxc.

Machines are created in a clean state and ready to have units deployed.

This command also supports manual provisioning of existing machines via SSH. The
target machine must be able to communicate with the API server, and be able to
access the environment storage.

Examples:
   juju add-machine                      (starts a new machine)
   juju add-machine lxc                  (starts a new machine with an lxc container)
   juju add-machine lxc:4                (starts a new lxc container on machine 4)
   juju add-machine --constraints mem=8G (starts a machine with at least 8GB RAM)
   juju add-machine ssh:user@10.10.0.3   (manually provisions a machine with ssh)

See Also:
   juju help constraints
`

// AddMachineCommand starts a new machine and registers it in the environment.
type AddMachineCommand struct {
	cmd.EnvCommandBase
	// If specified, use this series, else use the environment default-series
	Series string
	// If specified, these constraints are merged with those already in the environment.
	Constraints   constraints.Value
	MachineId     string
	ContainerType instance.ContainerType
	SSHHost       string
}

func (c *AddMachineCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "add-machine",
		Args:    "[<container>:machine | <container> | ssh:[user@]host]",
		Purpose: "start a new, empty machine and optionally a container, or add a container to a machine",
		Doc:     addMachineDoc,
	}
}

func (c *AddMachineCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	f.StringVar(&c.Series, "series", "", "the charm series")
	f.Var(constraints.ConstraintsValue{&c.Constraints}, "constraints", "additional machine constraints")
}

func (c *AddMachineCommand) Init(args []string) error {
	if c.Constraints.Container != nil {
		return fmt.Errorf("container constraint %q not allowed when adding a machine", *c.Constraints.Container)
	}
	containerSpec, err := cmd.ZeroOrOneArgs(args)
	if err != nil {
		return err
	}
	if containerSpec == "" {
		return nil
	}
	if strings.HasPrefix(containerSpec, sshHostPrefix) {
		c.SSHHost = containerSpec[len(sshHostPrefix):]
	} else {
		// container arg can either be 'type:machine' or 'type'
		if c.ContainerType, err = instance.ParseContainerType(containerSpec); err != nil {
			if names.IsMachine(containerSpec) || !cmd.IsMachineOrNewContainer(containerSpec) {
				return fmt.Errorf("malformed container argument %q", containerSpec)
			}
			sep := strings.Index(containerSpec, ":")
			c.MachineId = containerSpec[sep+1:]
			c.ContainerType, err = instance.ParseContainerType(containerSpec[:sep])
		}
	}
	return err
}

// addMachine1dot16 runs Client.AddMachines using a direct DB connection to maintain
// compatibility with an API server running 1.16 or older (when AddMachines
// was not available). This fallback can be removed when we no longer maintain
// 1.16 compatibility.
// This was copied directly from the code in AddMachineCommand.Run in 1.16
func (c *AddMachineCommand) addMachine1dot16() (string, error) {
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return "", err
	}
	defer conn.Close()

	series := c.Series
	if series == "" {
		conf, err := conn.State.EnvironConfig()
		if err != nil {
			return "", err
		}
		series = conf.DefaultSeries()
	}
	template := state.MachineTemplate{
		Series:      series,
		Constraints: c.Constraints,
		Jobs:        []state.MachineJob{state.JobHostUnits},
	}
	var m *state.Machine
	switch {
	case c.ContainerType == "":
		m, err = conn.State.AddOneMachine(template)
	case c.MachineId != "":
		m, err = conn.State.AddMachineInsideMachine(template, c.MachineId, c.ContainerType)
	default:
		m, err = conn.State.AddMachineInsideNewMachine(template, template, c.ContainerType)
	}
	if err != nil {
		return "", err
	}
	return m.String(), err
}

func (c *AddMachineCommand) Run(ctx *cmd.Context) error {
	if c.SSHHost != "" {
		args := manual.ProvisionMachineArgs{
			Host:    c.SSHHost,
			EnvName: c.EnvName,
			Stdin:   ctx.Stdin,
			Stdout:  ctx.Stdout,
			Stderr:  ctx.Stderr,
		}
		_, err := manual.ProvisionMachine(args)
		return err
	}

	client, err := juju.NewAPIClientFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer client.Close()

	machineParams := params.AddMachineParams{
		ParentId:      c.MachineId,
		ContainerType: c.ContainerType,
		Series:        c.Series,
		Constraints:   c.Constraints,
		Jobs:          []params.MachineJob{params.JobHostUnits},
	}
	results, err := client.AddMachines([]params.AddMachineParams{machineParams})
	var machineId string
	if params.IsCodeNotImplemented(err) {
		logger.Infof("AddMachines not supported by the API server, " +
			"falling back to 1.16 compatibility mode (direct DB access)")
		machineId, err = c.addMachine1dot16()
	} else if err != nil {
		return err
	} else {
		// Currently, only one machine is added, but in future there may be several added in one call.
		machineInfo := results[0]
		var machineErr *params.Error
		machineId, machineErr = machineInfo.Machine, machineInfo.Error
		if machineErr != nil {
			err = machineErr
		}
	}
	if err != nil {
		return err
	}
	if c.ContainerType == "" {
		logger.Infof("created machine %v", machineId)
	} else {
		logger.Infof("created %q container on machine %v", c.ContainerType, machineId)
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// 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 (
	"fmt"
	"log"

	"go.opentelemetry.io/collector/component"
	"go.opentelemetry.io/collector/service"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"

	"github.com/GoogleCloudPlatform/opentelemetry-operations-collector/internal/env"
	"github.com/GoogleCloudPlatform/opentelemetry-operations-collector/internal/levelchanger"
	"github.com/GoogleCloudPlatform/opentelemetry-operations-collector/internal/version"
)

func main() {
	if err := env.Create(); err != nil {
		log.Fatalf("failed to build environment variables for config: %v", err)
	}

	factories, err := components()
	if err != nil {
		log.Fatalf("failed to build default components: %v", err)
	}

	info := component.BuildInfo{
		Command:     "google-cloud-metrics-agent",
		Description: "Google Cloud Metrics Agent",
		Version:     version.Version,
	}

	params := service.CollectorSettings{
		Factories: factories,
		BuildInfo: info,
		LoggingOptions: []zap.Option{
			levelchanger.NewLevelChangerOption(
				zapcore.ErrorLevel,
				zapcore.DebugLevel,
				// We would like the Error logs from this file to be logged at Debug instead.
				// https://github.com/open-telemetry/opentelemetry-collector/blob/831373ae6c6959f6c9258ac585a2ec0ab19a074f/receiver/scraperhelper/scrapercontroller.go#L198
				levelchanger.FilePathLevelChangeCondition("scrapercontroller.go")),
		},
	}

	if err := run(params); err != nil {
		log.Fatal(err)
	}
}

func runInteractive(params service.CollectorSettings) error {
	cmd := service.NewCommand(params)
	err := cmd.Execute()
	if err != nil {
		return fmt.Errorf("application run finished with error: %w", err)
	}

	return nil
}
<commit_msg>Do not fail otel if useragent can't retrieved (#114)<commit_after>// Copyright 2020 Google LLC
//
// 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 (
	"fmt"
	"log"

	"go.opentelemetry.io/collector/component"
	"go.opentelemetry.io/collector/service"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"

	"github.com/GoogleCloudPlatform/opentelemetry-operations-collector/internal/env"
	"github.com/GoogleCloudPlatform/opentelemetry-operations-collector/internal/levelchanger"
	"github.com/GoogleCloudPlatform/opentelemetry-operations-collector/internal/version"
)

func main() {
	if err := env.Create(); err != nil {
		log.Printf("failed to build environment variables for config: %v", err)
	}

	factories, err := components()
	if err != nil {
		log.Fatalf("failed to build default components: %v", err)
	}

	info := component.BuildInfo{
		Command:     "google-cloud-metrics-agent",
		Description: "Google Cloud Metrics Agent",
		Version:     version.Version,
	}

	params := service.CollectorSettings{
		Factories: factories,
		BuildInfo: info,
		LoggingOptions: []zap.Option{
			levelchanger.NewLevelChangerOption(
				zapcore.ErrorLevel,
				zapcore.DebugLevel,
				// We would like the Error logs from this file to be logged at Debug instead.
				// https://github.com/open-telemetry/opentelemetry-collector/blob/831373ae6c6959f6c9258ac585a2ec0ab19a074f/receiver/scraperhelper/scrapercontroller.go#L198
				levelchanger.FilePathLevelChangeCondition("scrapercontroller.go")),
		},
	}

	if err := run(params); err != nil {
		log.Fatal(err)
	}
}

func runInteractive(params service.CollectorSettings) error {
	cmd := service.NewCommand(params)
	err := cmd.Execute()
	if err != nil {
		return fmt.Errorf("application run finished with error: %w", err)
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/bosun-monitor/scollector/_third_party/github.com/StackExchange/slog"
	"github.com/bosun-monitor/scollector/_third_party/github.com/bosun-monitor/collect"
	"github.com/bosun-monitor/scollector/_third_party/github.com/bosun-monitor/metadata"
	"github.com/bosun-monitor/scollector/_third_party/github.com/bosun-monitor/opentsdb"
	"github.com/bosun-monitor/scollector/_third_party/github.com/bosun-monitor/util"
	"github.com/bosun-monitor/scollector/collectors"
)

// These constants should remain in source control as their zero values.
const (
	// VersionDate should be set at build time as a date: 20140721184001.
	VersionDate uint64 = 0
	// VersionID should be set at build time as the most recent commit hash.
	VersionID string = ""
)

var (
	flagFilter          = flag.String("f", "", "Filters collectors matching this term. Works with all other arguments.")
	flagList            = flag.Bool("l", false, "List available collectors.")
	flagPrint           = flag.Bool("p", false, "Print to screen instead of sending to a host")
	flagHost            = flag.String("h", "", `Bosun or OpenTSDB host. Ex: "http://bosun.example.com:8070".`)
	flagColDir          = flag.String("c", "", `External collectors directory.`)
	flagBatchSize       = flag.Int("b", 0, "OpenTSDB batch size. Used for debugging bad data.")
	flagSNMP            = flag.String("s", "", "SNMP host to poll of the format: \"community@host[,community@host...]\".")
	flagICMP            = flag.String("i", "", "ICMP host to ping of the format: \"host[,host...]\".")
	flagVsphere         = flag.String("v", "", `vSphere host to poll of the format: "user:password@host[,user:password@host...]".`)
	flagFake            = flag.Int("fake", 0, "Generates X fake data points on the test.fake metric per second.")
	flagDebug           = flag.Bool("d", false, "Enables debug output.")
	flagFullHost        = flag.Bool("u", false, `Enables full hostnames: doesn't truncate to first ".".`)
	flagTags            = flag.String("t", "", `Tags to add to every datapoint in the format dc=ny,rack=3. If a collector specifies the same tag key, this one will be overwritten.`)
	flagDisableMetadata = flag.Bool("m", false, "Disable sending of metadata.")
	flagVersion         = flag.Bool("version", false, "Prints the version and exits.")
	flagDisableDefault  = flag.Bool("n", false, "Disable sending of scollector self metrics.")

	procs []*collectors.WatchedProc

	mains []func()
)

func readConf() {
	p, err := exePath()
	if err != nil {
		slog.Error(err)
		return
	}
	dir := filepath.Dir(p)
	p = filepath.Join(dir, "scollector.conf")
	b, err := ioutil.ReadFile(p)
	if err != nil {
		if *flagDebug {
			slog.Error(err)
		}
		return
	}
	for i, line := range strings.Split(string(b), "\n") {
		if strings.TrimSpace(line) == "" {
			continue
		}
		sp := strings.SplitN(line, "=", 2)
		if len(sp) != 2 {
			if *flagDebug {
				slog.Errorf("expected = in %v:%v", p, i+1)
			}
			continue
		}
		k := strings.TrimSpace(sp[0])
		v := strings.TrimSpace(sp[1])
		f := func(s *string) {
			if *s == "" {
				*s = v
			}
		}
		switch k {
		case "host":
			f(flagHost)
		case "filter":
			f(flagFilter)
		case "coldir":
			f(flagColDir)
		case "snmp":
			f(flagSNMP)
		case "icmp":
			f(flagICMP)
		case "tags":
			f(flagTags)
		case "vsphere":
			f(flagVsphere)
		case "process":
			p, err := collectors.NewWatchedProc(v)
			if err != nil {
				slog.Fatal(err)
			}
			procs = append(procs, p)
		default:
			if *flagDebug {
				slog.Errorf("unknown key in %v:%v", p, i+1)
			}
		}
	}
}

func main() {
	flag.Parse()
	if *flagPrint || *flagDebug {
		slog.Set(&slog.StdLog{Log: log.New(os.Stdout, "", log.LstdFlags)})
	}
	if *flagVersion {
		fmt.Printf("scollector version %v (%v)\n", VersionDate, VersionID)
		os.Exit(0)
	}
	if *flagTags != "" {
		var err error
		collectors.AddTags, err = opentsdb.ParseTags(*flagTags)
		if err != nil {
			slog.Fatalf("failed to parse additional tags %v: %v", *flagTags, err)
		}
	}
	for _, m := range mains {
		m()
	}
	readConf()
	util.FullHostname = *flagFullHost
	util.Set()
	if *flagColDir != "" {
		collectors.InitPrograms(*flagColDir)
	}
	if *flagSNMP != "" {
		for _, s := range strings.Split(*flagSNMP, ",") {
			sp := strings.Split(s, "@")
			if len(sp) != 2 {
				slog.Fatal("invalid snmp string:", *flagSNMP)
			}
			collectors.SNMPIfaces(sp[0], sp[1])
			collectors.SNMPCisco(sp[0], sp[1])
		}
	}
	if *flagICMP != "" {
		for _, s := range strings.Split(*flagICMP, ",") {
			collectors.ICMP(s)
		}
	}
	if *flagVsphere != "" {
		for _, s := range strings.Split(*flagVsphere, ",") {
			sp := strings.SplitN(s, ":", 2)
			if len(sp) != 2 {
				slog.Fatal("invalid vsphere string:", *flagVsphere)
			}
			user := sp[0]
			idx := strings.LastIndex(sp[1], "@")
			if idx == -1 {
				slog.Fatal("invalid vsphere string:", *flagVsphere)
			}
			pwd := sp[1][:idx]
			host := sp[1][idx+1:]
			if len(user) == 0 || len(pwd) == 0 || len(host) == 0 {
				slog.Fatal("invalid vsphere string:", *flagVsphere)
			}
			collectors.Vsphere(user, pwd, host)
		}
	}
	if len(procs) > 0 {
		if err := collectors.WatchProcesses(procs); err != nil {
			log.Fatal(err)
		}
	}
	if *flagFake > 0 {
		collectors.InitFake(*flagFake)
	}
	collect.Debug = *flagDebug
	util.Debug = *flagDebug
	if *flagDisableDefault {
		collect.DisableDefaultCollectors = true
	}
	c := collectors.Search(*flagFilter)
	for _, col := range c {
		col.Init()
	}
	u, err := parseHost()
	if *flagList {
		list(c)
		return
	} else if err != nil {
		slog.Fatal("invalid host:", *flagHost)
	}
	if *flagPrint {
		collectors.DefaultFreq = time.Second * 3
		slog.Infoln("Set default frequency to", collectors.DefaultFreq)
		collect.Print = true
	}
	if !*flagDisableMetadata {
		if err := metadata.Init(u, *flagDebug); err != nil {
			slog.Fatal(err)
		}
	}
	cdp := collectors.Run(c)
	if u != nil {
		slog.Infoln("OpenTSDB host:", u)
	}
	if err := collect.InitChan(u, "scollector", cdp); err != nil {
		slog.Fatal(err)
	}
	if VersionDate > 0 {
		if err := collect.Put("version", nil, VersionDate); err != nil {
			slog.Error(err)
		}
	}
	if *flagBatchSize > 0 {
		collect.BatchSize = *flagBatchSize
	}
	go func() {
		const maxMem = 500 * 1024 * 1024 // 500MB
		var m runtime.MemStats
		for _ = range time.Tick(time.Minute) {
			runtime.ReadMemStats(&m)
			if m.Alloc > maxMem {
				panic("memory max reached")
			}
		}
	}()
	select {}
}

func exePath() (string, error) {
	prog := os.Args[0]
	p, err := filepath.Abs(prog)
	if err != nil {
		return "", err
	}
	fi, err := os.Stat(p)
	if err == nil {
		if !fi.Mode().IsDir() {
			return p, nil
		}
		err = fmt.Errorf("%s is directory", p)
	}
	if filepath.Ext(p) == "" {
		p += ".exe"
		fi, err := os.Stat(p)
		if err == nil {
			if !fi.Mode().IsDir() {
				return p, nil
			}
			err = fmt.Errorf("%s is directory", p)
		}
	}
	return "", err
}

func list(cs []collectors.Collector) {
	for _, c := range cs {
		fmt.Println(c.Name())
	}
}

func parseHost() (*url.URL, error) {
	if *flagHost == "" {
		*flagHost = "bosun"
	}
	if !strings.Contains(*flagHost, "//") {
		*flagHost = "http://" + *flagHost
	}
	return url.Parse(*flagHost)
}

func printPut(c chan *opentsdb.DataPoint) {
	for dp := range c {
		b, _ := json.Marshal(dp)
		slog.Info(string(b))
	}
}
<commit_msg>cmd/scollector: Read additional tags *after* reading the config<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/bosun-monitor/scollector/_third_party/github.com/StackExchange/slog"
	"github.com/bosun-monitor/scollector/_third_party/github.com/bosun-monitor/collect"
	"github.com/bosun-monitor/scollector/_third_party/github.com/bosun-monitor/metadata"
	"github.com/bosun-monitor/scollector/_third_party/github.com/bosun-monitor/opentsdb"
	"github.com/bosun-monitor/scollector/_third_party/github.com/bosun-monitor/util"
	"github.com/bosun-monitor/scollector/collectors"
)

// These constants should remain in source control as their zero values.
const (
	// VersionDate should be set at build time as a date: 20140721184001.
	VersionDate uint64 = 0
	// VersionID should be set at build time as the most recent commit hash.
	VersionID string = ""
)

var (
	flagFilter          = flag.String("f", "", "Filters collectors matching this term. Works with all other arguments.")
	flagList            = flag.Bool("l", false, "List available collectors.")
	flagPrint           = flag.Bool("p", false, "Print to screen instead of sending to a host")
	flagHost            = flag.String("h", "", `Bosun or OpenTSDB host. Ex: "http://bosun.example.com:8070".`)
	flagColDir          = flag.String("c", "", `External collectors directory.`)
	flagBatchSize       = flag.Int("b", 0, "OpenTSDB batch size. Used for debugging bad data.")
	flagSNMP            = flag.String("s", "", "SNMP host to poll of the format: \"community@host[,community@host...]\".")
	flagICMP            = flag.String("i", "", "ICMP host to ping of the format: \"host[,host...]\".")
	flagVsphere         = flag.String("v", "", `vSphere host to poll of the format: "user:password@host[,user:password@host...]".`)
	flagFake            = flag.Int("fake", 0, "Generates X fake data points on the test.fake metric per second.")
	flagDebug           = flag.Bool("d", false, "Enables debug output.")
	flagFullHost        = flag.Bool("u", false, `Enables full hostnames: doesn't truncate to first ".".`)
	flagTags            = flag.String("t", "", `Tags to add to every datapoint in the format dc=ny,rack=3. If a collector specifies the same tag key, this one will be overwritten.`)
	flagDisableMetadata = flag.Bool("m", false, "Disable sending of metadata.")
	flagVersion         = flag.Bool("version", false, "Prints the version and exits.")
	flagDisableDefault  = flag.Bool("n", false, "Disable sending of scollector self metrics.")

	procs []*collectors.WatchedProc

	mains []func()
)

func readConf() {
	p, err := exePath()
	if err != nil {
		slog.Error(err)
		return
	}
	dir := filepath.Dir(p)
	p = filepath.Join(dir, "scollector.conf")
	b, err := ioutil.ReadFile(p)
	if err != nil {
		if *flagDebug {
			slog.Error(err)
		}
		return
	}
	for i, line := range strings.Split(string(b), "\n") {
		if strings.TrimSpace(line) == "" {
			continue
		}
		sp := strings.SplitN(line, "=", 2)
		if len(sp) != 2 {
			if *flagDebug {
				slog.Errorf("expected = in %v:%v", p, i+1)
			}
			continue
		}
		k := strings.TrimSpace(sp[0])
		v := strings.TrimSpace(sp[1])
		f := func(s *string) {
			if *s == "" {
				*s = v
			}
		}
		switch k {
		case "host":
			f(flagHost)
		case "filter":
			f(flagFilter)
		case "coldir":
			f(flagColDir)
		case "snmp":
			f(flagSNMP)
		case "icmp":
			f(flagICMP)
		case "tags":
			f(flagTags)
		case "vsphere":
			f(flagVsphere)
		case "process":
			p, err := collectors.NewWatchedProc(v)
			if err != nil {
				slog.Fatal(err)
			}
			procs = append(procs, p)
		default:
			if *flagDebug {
				slog.Errorf("unknown key in %v:%v", p, i+1)
			}
		}
	}
}

func main() {
	flag.Parse()
	if *flagPrint || *flagDebug {
		slog.Set(&slog.StdLog{Log: log.New(os.Stdout, "", log.LstdFlags)})
	}
	if *flagVersion {
		fmt.Printf("scollector version %v (%v)\n", VersionDate, VersionID)
		os.Exit(0)
	}
	for _, m := range mains {
		m()
	}
	readConf()
	if *flagTags != "" {
		var err error
		collectors.AddTags, err = opentsdb.ParseTags(*flagTags)
		if err != nil {
			slog.Fatalf("failed to parse additional tags %v: %v", *flagTags, err)
		}
	}
	util.FullHostname = *flagFullHost
	util.Set()
	if *flagColDir != "" {
		collectors.InitPrograms(*flagColDir)
	}
	if *flagSNMP != "" {
		for _, s := range strings.Split(*flagSNMP, ",") {
			sp := strings.Split(s, "@")
			if len(sp) != 2 {
				slog.Fatal("invalid snmp string:", *flagSNMP)
			}
			collectors.SNMPIfaces(sp[0], sp[1])
			collectors.SNMPCisco(sp[0], sp[1])
		}
	}
	if *flagICMP != "" {
		for _, s := range strings.Split(*flagICMP, ",") {
			collectors.ICMP(s)
		}
	}
	if *flagVsphere != "" {
		for _, s := range strings.Split(*flagVsphere, ",") {
			sp := strings.SplitN(s, ":", 2)
			if len(sp) != 2 {
				slog.Fatal("invalid vsphere string:", *flagVsphere)
			}
			user := sp[0]
			idx := strings.LastIndex(sp[1], "@")
			if idx == -1 {
				slog.Fatal("invalid vsphere string:", *flagVsphere)
			}
			pwd := sp[1][:idx]
			host := sp[1][idx+1:]
			if len(user) == 0 || len(pwd) == 0 || len(host) == 0 {
				slog.Fatal("invalid vsphere string:", *flagVsphere)
			}
			collectors.Vsphere(user, pwd, host)
		}
	}
	if len(procs) > 0 {
		if err := collectors.WatchProcesses(procs); err != nil {
			log.Fatal(err)
		}
	}
	if *flagFake > 0 {
		collectors.InitFake(*flagFake)
	}
	collect.Debug = *flagDebug
	util.Debug = *flagDebug
	if *flagDisableDefault {
		collect.DisableDefaultCollectors = true
	}
	c := collectors.Search(*flagFilter)
	for _, col := range c {
		col.Init()
	}
	u, err := parseHost()
	if *flagList {
		list(c)
		return
	} else if err != nil {
		slog.Fatal("invalid host:", *flagHost)
	}
	if *flagPrint {
		collectors.DefaultFreq = time.Second * 3
		slog.Infoln("Set default frequency to", collectors.DefaultFreq)
		collect.Print = true
	}
	if !*flagDisableMetadata {
		if err := metadata.Init(u, *flagDebug); err != nil {
			slog.Fatal(err)
		}
	}
	cdp := collectors.Run(c)
	if u != nil {
		slog.Infoln("OpenTSDB host:", u)
	}
	if err := collect.InitChan(u, "scollector", cdp); err != nil {
		slog.Fatal(err)
	}
	if VersionDate > 0 {
		if err := collect.Put("version", nil, VersionDate); err != nil {
			slog.Error(err)
		}
	}
	if *flagBatchSize > 0 {
		collect.BatchSize = *flagBatchSize
	}
	go func() {
		const maxMem = 500 * 1024 * 1024 // 500MB
		var m runtime.MemStats
		for _ = range time.Tick(time.Minute) {
			runtime.ReadMemStats(&m)
			if m.Alloc > maxMem {
				panic("memory max reached")
			}
		}
	}()
	select {}
}

func exePath() (string, error) {
	prog := os.Args[0]
	p, err := filepath.Abs(prog)
	if err != nil {
		return "", err
	}
	fi, err := os.Stat(p)
	if err == nil {
		if !fi.Mode().IsDir() {
			return p, nil
		}
		err = fmt.Errorf("%s is directory", p)
	}
	if filepath.Ext(p) == "" {
		p += ".exe"
		fi, err := os.Stat(p)
		if err == nil {
			if !fi.Mode().IsDir() {
				return p, nil
			}
			err = fmt.Errorf("%s is directory", p)
		}
	}
	return "", err
}

func list(cs []collectors.Collector) {
	for _, c := range cs {
		fmt.Println(c.Name())
	}
}

func parseHost() (*url.URL, error) {
	if *flagHost == "" {
		*flagHost = "bosun"
	}
	if !strings.Contains(*flagHost, "//") {
		*flagHost = "http://" + *flagHost
	}
	return url.Parse(*flagHost)
}

func printPut(c chan *opentsdb.DataPoint) {
	for dp := range c {
		b, _ := json.Marshal(dp)
		slog.Info(string(b))
	}
}
<|endoftext|>


<commit_before>/*
 * Minio Cloud Storage, (C) 2015, 2016 Minio, 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 cmd

import (
	"bufio"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"math/big"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"strconv"
	"sync"
	"testing"
	"time"
)

func TestClose(t *testing.T) {
	// Create ServerMux
	m := NewMuxServer("", nil)

	if err := m.Close(); err != nil {
		t.Error("Server errored while trying to Close", err)
	}
}

func TestMuxServer(t *testing.T) {
	ts := httptest.NewUnstartedServer(nil)
	defer ts.Close()

	// Create ServerMux
	m := NewMuxServer("", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "hello")
	}))

	// Set the test server config to the mux
	ts.Config = &m.Server
	ts.Start()

	// Create a MuxListener
	ml, err := NewMuxListener(ts.Listener, m.WaitGroup, "", "")
	if err != nil {
		t.Fatal(err)
	}
	m.listener = ml

	client := http.Client{}
	res, err := client.Get(ts.URL)
	if err != nil {
		t.Fatal(err)
	}

	got, err := ioutil.ReadAll(res.Body)
	if err != nil {
		t.Fatal(err)
	}

	if string(got) != "hello" {
		t.Errorf("got %q, want hello", string(got))
	}

	// Make sure there is only 1 connection
	m.mu.Lock()
	if len(m.conns) < 1 {
		t.Fatal("Should have 1 connections")
	}
	m.mu.Unlock()

	// Close the server
	m.Close()

	// Make sure there are zero connections
	m.mu.Lock()
	if len(m.conns) > 0 {
		t.Fatal("Should have 0 connections")
	}
	m.mu.Unlock()
}

func TestServerCloseBlocking(t *testing.T) {
	ts := httptest.NewUnstartedServer(nil)
	defer ts.Close()

	// Create ServerMux
	m := NewMuxServer("", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "hello")
	}))

	// Set the test server config to the mux
	ts.Config = &m.Server
	ts.Start()

	// Create a MuxListener
	// var err error
	ml, err := NewMuxListener(ts.Listener, m.WaitGroup, "", "")
	if err != nil {
		t.Fatal(err)
	}
	m.listener = ml

	dial := func() net.Conn {
		c, cerr := net.Dial("tcp", ts.Listener.Addr().String())
		if cerr != nil {
			t.Fatal(err)
		}
		return c
	}

	// Dial to open a StateNew but don't send anything
	cnew := dial()
	defer cnew.Close()

	// Dial another connection but idle after a request to have StateIdle
	cidle := dial()
	defer cidle.Close()
	cidle.Write([]byte("HEAD / HTTP/1.1\r\nHost: foo\r\n\r\n"))
	_, err = http.ReadResponse(bufio.NewReader(cidle), nil)
	if err != nil {
		t.Fatal(err)
	}

	// Make sure we don't block forever.
	m.Close()

	// Make sure there are zero connections
	m.mu.Lock()
	if len(m.conns) > 0 {
		t.Fatal("Should have 0 connections")
	}
	m.mu.Unlock()
}

func TestListenAndServe(t *testing.T) {
	wait := make(chan struct{})
	addr := "127.0.0.1:" + strconv.Itoa(getFreePort())
	errc := make(chan error)
	once := &sync.Once{}

	// Create ServerMux and when we receive a request we stop waiting
	m := NewMuxServer(addr, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "hello")
		once.Do(func() { close(wait) })
	}))

	// ListenAndServe in a goroutine, but we don't know when it's ready
	go func() { errc <- m.ListenAndServe() }()

	// Make sure we don't block by closing wait after a timeout
	tf := time.AfterFunc(time.Millisecond*500, func() { errc <- errors.New("Unable to connect to server") })

	// Keep trying the server until it's accepting connections
	go func() {
		client := http.Client{Timeout: time.Millisecond * 10}
		ok := false
		for !ok {
			res, _ := client.Get("http://" + addr)
			if res != nil && res.StatusCode == http.StatusOK {
				ok = true
			}
		}

		tf.Stop() // Cancel the timeout since we made a successful request
	}()

	// Block until we get an error or wait closed
	select {
	case err := <-errc:
		if err != nil {
			t.Fatal(err)
		}
	case <-wait:
		m.Close() // Shutdown the ServerMux
		return
	}
}

func TestListenAndServeTLS(t *testing.T) {
	wait := make(chan struct{})
	addr := "127.0.0.1:" + strconv.Itoa(getFreePort())
	errc := make(chan error)
	once := &sync.Once{}

	// Create ServerMux and when we receive a request we stop waiting
	m := NewMuxServer(addr, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "hello")
		once.Do(func() { close(wait) })
	}))

	// Create a cert
	err := createCertsPath()
	if err != nil {
		t.Fatal(err)
	}
	certFile := mustGetCertFile()
	keyFile := mustGetKeyFile()
	defer os.RemoveAll(certFile)
	defer os.RemoveAll(keyFile)

	err = generateTestCert(addr)
	if err != nil {
		t.Error(err)
		return
	}

	// ListenAndServe in a goroutine, but we don't know when it's ready
	go func() { errc <- m.ListenAndServeTLS(certFile, keyFile) }()

	// Make sure we don't block by closing wait after a timeout
	tf := time.AfterFunc(time.Millisecond*500, func() { errc <- errors.New("Unable to connect to server") })

	// Keep trying the server until it's accepting connections
	go func() {
		tr := &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}
		client := http.Client{
			Timeout:   time.Millisecond * 10,
			Transport: tr,
		}
		ok := false
		for !ok {
			res, _ := client.Get("https://" + addr)
			if res != nil && res.StatusCode == http.StatusOK {
				ok = true
			}
		}

		tf.Stop() // Cancel the timeout since we made a successful request
	}()

	// Block until we get an error or wait closed
	select {
	case err := <-errc:
		if err != nil {
			t.Error(err)
			return
		}
	case <-wait:
		m.Close() // Shutdown the ServerMux
		return
	}
}

// generateTestCert creates a cert and a key used for testing only
func generateTestCert(host string) error {
	certPath := mustGetCertFile()
	keyPath := mustGetKeyFile()
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return err
	}

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return err
	}

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: []string{"Minio Test Cert"},
		},
		NotBefore: time.Now(),
		NotAfter:  time.Now().Add(time.Minute * 1),

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	if ip := net.ParseIP(host); ip != nil {
		template.IPAddresses = append(template.IPAddresses, ip)
	}

	template.IsCA = true
	template.KeyUsage |= x509.KeyUsageCertSign

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
	if err != nil {
		return err
	}

	certOut, err := os.Create(certPath)
	if err != nil {
		return err
	}
	pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	certOut.Close()

	keyOut, err := os.OpenFile(keyPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return err
	}
	pem.Encode(keyOut, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})
	keyOut.Close()
	return nil
}
<commit_msg>Fix rare 'go test -race' failure in ListenServe{Plain,TLS} (#2588)<commit_after>/*
 * Minio Cloud Storage, (C) 2015, 2016 Minio, 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 cmd

import (
	"bufio"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"math/big"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"strconv"
	"sync"
	"testing"
	"time"
)

func TestClose(t *testing.T) {
	// Create ServerMux
	m := NewMuxServer("", nil)

	if err := m.Close(); err != nil {
		t.Error("Server errored while trying to Close", err)
	}
}

func TestMuxServer(t *testing.T) {
	ts := httptest.NewUnstartedServer(nil)
	defer ts.Close()

	// Create ServerMux
	m := NewMuxServer("", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "hello")
	}))

	// Set the test server config to the mux
	ts.Config = &m.Server
	ts.Start()

	// Create a MuxListener
	ml, err := NewMuxListener(ts.Listener, m.WaitGroup, "", "")
	if err != nil {
		t.Fatal(err)
	}
	m.listener = ml

	client := http.Client{}
	res, err := client.Get(ts.URL)
	if err != nil {
		t.Fatal(err)
	}

	got, err := ioutil.ReadAll(res.Body)
	if err != nil {
		t.Fatal(err)
	}

	if string(got) != "hello" {
		t.Errorf("got %q, want hello", string(got))
	}

	// Make sure there is only 1 connection
	m.mu.Lock()
	if len(m.conns) < 1 {
		t.Fatal("Should have 1 connections")
	}
	m.mu.Unlock()

	// Close the server
	m.Close()

	// Make sure there are zero connections
	m.mu.Lock()
	if len(m.conns) > 0 {
		t.Fatal("Should have 0 connections")
	}
	m.mu.Unlock()
}

func TestServerCloseBlocking(t *testing.T) {
	ts := httptest.NewUnstartedServer(nil)
	defer ts.Close()

	// Create ServerMux
	m := NewMuxServer("", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "hello")
	}))

	// Set the test server config to the mux
	ts.Config = &m.Server
	ts.Start()

	// Create a MuxListener
	// var err error
	ml, err := NewMuxListener(ts.Listener, m.WaitGroup, "", "")
	if err != nil {
		t.Fatal(err)
	}
	m.listener = ml

	dial := func() net.Conn {
		c, cerr := net.Dial("tcp", ts.Listener.Addr().String())
		if cerr != nil {
			t.Fatal(err)
		}
		return c
	}

	// Dial to open a StateNew but don't send anything
	cnew := dial()
	defer cnew.Close()

	// Dial another connection but idle after a request to have StateIdle
	cidle := dial()
	defer cidle.Close()
	cidle.Write([]byte("HEAD / HTTP/1.1\r\nHost: foo\r\n\r\n"))
	_, err = http.ReadResponse(bufio.NewReader(cidle), nil)
	if err != nil {
		t.Fatal(err)
	}

	// Make sure we don't block forever.
	m.Close()

	// Make sure there are zero connections
	m.mu.Lock()
	if len(m.conns) > 0 {
		t.Fatal("Should have 0 connections")
	}
	m.mu.Unlock()
}

func TestListenAndServePlain(t *testing.T) {
	wait := make(chan struct{})
	addr := "127.0.0.1:" + strconv.Itoa(getFreePort())
	errc := make(chan error)
	once := &sync.Once{}

	// Create ServerMux and when we receive a request we stop waiting
	m := NewMuxServer(addr, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "hello")
		once.Do(func() { close(wait) })
	}))

	// ListenAndServe in a goroutine, but we don't know when it's ready
	go func() { errc <- m.ListenAndServe() }()

	// Make sure we don't block by closing wait after a timeout
	tf := time.AfterFunc(time.Millisecond*500, func() { errc <- errors.New("Unable to connect to server") })

	wg := &sync.WaitGroup{}
	wg.Add(1)
	// Keep trying the server until it's accepting connections
	go func() {
		client := http.Client{Timeout: time.Millisecond * 10}
		ok := false
		for !ok {
			res, _ := client.Get("http://" + addr)
			if res != nil && res.StatusCode == http.StatusOK {
				ok = true
			}
		}

		wg.Done()
		tf.Stop() // Cancel the timeout since we made a successful request
	}()

	wg.Wait()

	// Block until we get an error or wait closed
	select {
	case err := <-errc:
		if err != nil {
			t.Fatal(err)
		}
	case <-wait:
		m.Close() // Shutdown the ServerMux
		return
	}
}

func TestListenAndServeTLS(t *testing.T) {
	wait := make(chan struct{})
	addr := "127.0.0.1:" + strconv.Itoa(getFreePort())
	errc := make(chan error)
	once := &sync.Once{}

	// Create ServerMux and when we receive a request we stop waiting
	m := NewMuxServer(addr, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, "hello")
		once.Do(func() { close(wait) })
	}))

	// Create a cert
	err := createCertsPath()
	if err != nil {
		t.Fatal(err)
	}
	certFile := mustGetCertFile()
	keyFile := mustGetKeyFile()
	defer os.RemoveAll(certFile)
	defer os.RemoveAll(keyFile)

	err = generateTestCert(addr)
	if err != nil {
		t.Error(err)
		return
	}

	// ListenAndServe in a goroutine, but we don't know when it's ready
	go func() { errc <- m.ListenAndServeTLS(certFile, keyFile) }()

	// Make sure we don't block by closing wait after a timeout
	tf := time.AfterFunc(time.Millisecond*500, func() { errc <- errors.New("Unable to connect to server") })
	wg := &sync.WaitGroup{}
	wg.Add(1)

	// Keep trying the server until it's accepting connections
	go func() {
		tr := &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}
		client := http.Client{
			Timeout:   time.Millisecond * 10,
			Transport: tr,
		}
		ok := false
		for !ok {
			res, _ := client.Get("https://" + addr)
			if res != nil && res.StatusCode == http.StatusOK {
				ok = true
			}
		}

		wg.Done()
		tf.Stop() // Cancel the timeout since we made a successful request
	}()

	wg.Wait()

	// Block until we get an error or wait closed
	select {
	case err := <-errc:
		if err != nil {
			t.Error(err)
			return
		}
	case <-wait:
		m.Close() // Shutdown the ServerMux
		return
	}
}

// generateTestCert creates a cert and a key used for testing only
func generateTestCert(host string) error {
	certPath := mustGetCertFile()
	keyPath := mustGetKeyFile()
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return err
	}

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return err
	}

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: []string{"Minio Test Cert"},
		},
		NotBefore: time.Now(),
		NotAfter:  time.Now().Add(time.Minute * 1),

		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	if ip := net.ParseIP(host); ip != nil {
		template.IPAddresses = append(template.IPAddresses, ip)
	}

	template.IsCA = true
	template.KeyUsage |= x509.KeyUsageCertSign

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
	if err != nil {
		return err
	}

	certOut, err := os.Create(certPath)
	if err != nil {
		return err
	}
	pem.Encode(certOut, &pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	certOut.Close()

	keyOut, err := os.OpenFile(keyPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		return err
	}
	pem.Encode(keyOut, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})
	keyOut.Close()
	return nil
}
<|endoftext|>