<commit_before>/*

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 (
	"context"
	"flag"
	"github.com/spf13/pflag"
	"k8s.io/apimachinery/pkg/runtime"
	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp"
	"os"
	elbv2v1alpha1 "sigs.k8s.io/aws-load-balancer-controller/apis/elbv2/v1alpha1"
	elbv2controller "sigs.k8s.io/aws-load-balancer-controller/controllers/elbv2"
	"sigs.k8s.io/aws-load-balancer-controller/controllers/ingress"
	"sigs.k8s.io/aws-load-balancer-controller/controllers/service"
	"sigs.k8s.io/aws-load-balancer-controller/pkg/aws"
	"sigs.k8s.io/aws-load-balancer-controller/pkg/aws/throttle"
	inject "sigs.k8s.io/aws-load-balancer-controller/pkg/inject"
	"sigs.k8s.io/aws-load-balancer-controller/pkg/k8s"
	"sigs.k8s.io/aws-load-balancer-controller/pkg/networking"
	"sigs.k8s.io/aws-load-balancer-controller/pkg/targetgroupbinding"
	corewebhook "sigs.k8s.io/aws-load-balancer-controller/webhooks/core"
	elbv2webhook "sigs.k8s.io/aws-load-balancer-controller/webhooks/elbv2"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/log/zap"
	"sigs.k8s.io/controller-runtime/pkg/metrics"
	// +kubebuilder:scaffold:imports
)

const (
	flagMetricsAddr          = "metrics-addr"
	flagEnableLeaderElection = "enable-leader-election"
	flagK8sClusterName       = "cluster-name"
)

var (
	scheme   = runtime.NewScheme()
	setupLog = ctrl.Log.WithName("setup")
)

func init() {
	_ = clientgoscheme.AddToScheme(scheme)

	_ = elbv2v1alpha1.AddToScheme(scheme)
	// +kubebuilder:scaffold:scheme
}

func main() {
	var metricsAddr string
	var enableLeaderElection bool
	var k8sClusterName string
	awsCloudConfig := aws.CloudConfig{ThrottleConfig: throttle.NewDefaultServiceOperationsThrottleConfig()}
	injectConfig := inject.Config{}
	fs := pflag.NewFlagSet("", pflag.ExitOnError)
	fs.StringVar(&metricsAddr, flagMetricsAddr, ":8080", "The address the metric endpoint binds to.")
	fs.BoolVar(&enableLeaderElection, flagEnableLeaderElection, false,
		"Enable leader election for controller manager. "+
			"Enabling this will ensure there is only one active controller manager.")
	fs.StringVar(&k8sClusterName, flagK8sClusterName, "", "Kubernetes cluster name")
	awsCloudConfig.BindFlags(fs)
	injectConfig.BindFlags(fs)
	fs.AddGoFlagSet(flag.CommandLine)
	if err := fs.Parse(os.Args); err != nil {
		setupLog.Error(err, "invalid flags")
		os.Exit(1)
	}

	ctrl.SetLogger(zap.New(zap.UseDevMode(false)))
	if len(k8sClusterName) == 0 {
		setupLog.Info("Kubernetes cluster name must be specified")
		os.Exit(1)
	}

	cloud, err := aws.NewCloud(awsCloudConfig, metrics.Registry)
	if err != nil {
		setupLog.Error(err, "Unable to initialize AWS cloud")
		os.Exit(1)
	}

	ctrl.SetLogger(zap.New(zap.UseDevMode(true)))

	mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{
		Scheme:             scheme,
		MetricsBindAddress: metricsAddr,
		LeaderElection:     enableLeaderElection,
		Port:               9443,
	})
	if err != nil {
		setupLog.Error(err, "unable to start manager")
		os.Exit(1)
	}

	podENIResolver := networking.NewDefaultPodENIInfoResolver(cloud.EC2(), cloud.VpcID(), ctrl.Log)
	nodeENIResolver := networking.NewDefaultNodeENIInfoResolver(cloud.EC2(), ctrl.Log)
	sgManager := networking.NewDefaultSecurityGroupManager(cloud.EC2(), ctrl.Log)
	sgReconciler := networking.NewDefaultSecurityGroupReconciler(sgManager, ctrl.Log)
	finalizerManager := k8s.NewDefaultFinalizerManager(mgr.GetClient(), ctrl.Log)
	tgbResManager := targetgroupbinding.NewDefaultResourceManager(mgr.GetClient(), cloud.ELBV2(),
		podENIResolver, nodeENIResolver, sgManager, sgReconciler, cloud.VpcID(), k8sClusterName, ctrl.Log)

	subnetResolver := networking.NewSubnetsResolver(cloud.EC2(), cloud.VpcID(), k8sClusterName, ctrl.Log.WithName("subnets-resolver"))
	ingGroupReconciler := ingress.NewGroupReconciler(cloud, mgr.GetClient(), mgr.GetEventRecorderFor("ingress"),
		sgManager, sgReconciler, k8sClusterName, subnetResolver, ctrl.Log)
	tgbReconciler := elbv2controller.NewTargetGroupBindingReconciler(mgr.GetClient(), mgr.GetFieldIndexer(), finalizerManager, tgbResManager,
		ctrl.Log.WithName("controllers").WithName("TargetGroupBinding"))
	svcReconciler := service.NewServiceReconciler(
		cloud,
		mgr.GetClient(),
		sgManager,
		sgReconciler,
		k8sClusterName,
		subnetResolver,
		ctrl.Log.WithName("controllers").WithName("Service"))
	if err = ingGroupReconciler.SetupWithManager(mgr); err != nil {
		setupLog.Error(err, "unable to create controller", "controller", "Ingress")
		os.Exit(1)
	}
	if err := tgbReconciler.SetupWithManager(context.Background(), mgr); err != nil {
		setupLog.Error(err, "unable to create controller", "controller", "TargetGroupBinding")
		os.Exit(1)
	}
	if err = svcReconciler.SetupWithManager(mgr); err != nil {
		setupLog.Error(err, "Unable to create controller", "controller", "Service")
		os.Exit(1)
	}

	podReadinessGateInjector := inject.NewPodReadinessGate(injectConfig, mgr.GetClient(), ctrl.Log.WithName("pod-readiness-gate-injector"))
	corewebhook.NewPodMutator(podReadinessGateInjector).SetupWithManager(mgr)
	elbv2webhook.NewTargetGroupBindingMutator(cloud.ELBV2(), ctrl.Log).SetupWithManager(mgr)
	elbv2webhook.NewTargetGroupBindingValidator(ctrl.Log).SetupWithManager(mgr)
	// +kubebuilder:scaffold:builder

	setupLog.Info("starting manager")
	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
		setupLog.Error(err, "problem running manager")
		os.Exit(1)
	}
}
<commit_msg>enable leader election by default (#1462)<commit_after>/*

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 (
	"context"
	"flag"
	"github.com/spf13/pflag"
	"k8s.io/apimachinery/pkg/runtime"
	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp"
	"os"
	elbv2v1alpha1 "sigs.k8s.io/aws-load-balancer-controller/apis/elbv2/v1alpha1"
	elbv2controller "sigs.k8s.io/aws-load-balancer-controller/controllers/elbv2"
	"sigs.k8s.io/aws-load-balancer-controller/controllers/ingress"
	"sigs.k8s.io/aws-load-balancer-controller/controllers/service"
	"sigs.k8s.io/aws-load-balancer-controller/pkg/aws"
	"sigs.k8s.io/aws-load-balancer-controller/pkg/aws/throttle"
	inject "sigs.k8s.io/aws-load-balancer-controller/pkg/inject"
	"sigs.k8s.io/aws-load-balancer-controller/pkg/k8s"
	"sigs.k8s.io/aws-load-balancer-controller/pkg/networking"
	"sigs.k8s.io/aws-load-balancer-controller/pkg/targetgroupbinding"
	corewebhook "sigs.k8s.io/aws-load-balancer-controller/webhooks/core"
	elbv2webhook "sigs.k8s.io/aws-load-balancer-controller/webhooks/elbv2"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/log/zap"
	"sigs.k8s.io/controller-runtime/pkg/metrics"
	// +kubebuilder:scaffold:imports
)

const (
	flagMetricsAddr          = "metrics-addr"
	flagEnableLeaderElection = "enable-leader-election"
	flagLeaderElectionID     = "leader-election-id"
	flagK8sClusterName       = "cluster-name"
	defaultLeaderElectionID  = "aws-load-balancer-controller-leader"
)

var (
	scheme   = runtime.NewScheme()
	setupLog = ctrl.Log.WithName("setup")
)

func init() {
	_ = clientgoscheme.AddToScheme(scheme)

	_ = elbv2v1alpha1.AddToScheme(scheme)
	// +kubebuilder:scaffold:scheme
}

func main() {
	var metricsAddr string
	var enableLeaderElection bool
	var k8sClusterName string
	var leaderElectionID string
	awsCloudConfig := aws.CloudConfig{ThrottleConfig: throttle.NewDefaultServiceOperationsThrottleConfig()}
	injectConfig := inject.Config{}
	fs := pflag.NewFlagSet("", pflag.ExitOnError)
	fs.StringVar(&metricsAddr, flagMetricsAddr, ":8080", "The address the metric endpoint binds to.")
	fs.BoolVar(&enableLeaderElection, flagEnableLeaderElection, true,
		"Enable leader election for controller manager. "+
			"Enabling this will ensure there is only one active controller manager.")
	fs.StringVar(&leaderElectionID, flagLeaderElectionID, defaultLeaderElectionID,
		"Name of the leader election ID to use for this controller")
	fs.StringVar(&k8sClusterName, flagK8sClusterName, "", "Kubernetes cluster name")
	awsCloudConfig.BindFlags(fs)
	injectConfig.BindFlags(fs)
	fs.AddGoFlagSet(flag.CommandLine)
	if err := fs.Parse(os.Args); err != nil {
		setupLog.Error(err, "invalid flags")
		os.Exit(1)
	}

	ctrl.SetLogger(zap.New(zap.UseDevMode(false)))
	if len(k8sClusterName) == 0 {
		setupLog.Info("Kubernetes cluster name must be specified")
		os.Exit(1)
	}

	cloud, err := aws.NewCloud(awsCloudConfig, metrics.Registry)
	if err != nil {
		setupLog.Error(err, "Unable to initialize AWS cloud")
		os.Exit(1)
	}

	ctrl.SetLogger(zap.New(zap.UseDevMode(true)))

	mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{
		Scheme:             scheme,
		MetricsBindAddress: metricsAddr,
		LeaderElection:     enableLeaderElection,
		LeaderElectionID:   leaderElectionID,
		Port:               9443,
	})
	if err != nil {
		setupLog.Error(err, "unable to start manager")
		os.Exit(1)
	}

	podENIResolver := networking.NewDefaultPodENIInfoResolver(cloud.EC2(), cloud.VpcID(), ctrl.Log)
	nodeENIResolver := networking.NewDefaultNodeENIInfoResolver(cloud.EC2(), ctrl.Log)
	sgManager := networking.NewDefaultSecurityGroupManager(cloud.EC2(), ctrl.Log)
	sgReconciler := networking.NewDefaultSecurityGroupReconciler(sgManager, ctrl.Log)
	finalizerManager := k8s.NewDefaultFinalizerManager(mgr.GetClient(), ctrl.Log)
	tgbResManager := targetgroupbinding.NewDefaultResourceManager(mgr.GetClient(), cloud.ELBV2(),
		podENIResolver, nodeENIResolver, sgManager, sgReconciler, cloud.VpcID(), k8sClusterName, ctrl.Log)

	subnetResolver := networking.NewSubnetsResolver(cloud.EC2(), cloud.VpcID(), k8sClusterName, ctrl.Log.WithName("subnets-resolver"))
	ingGroupReconciler := ingress.NewGroupReconciler(cloud, mgr.GetClient(), mgr.GetEventRecorderFor("ingress"),
		sgManager, sgReconciler, k8sClusterName, subnetResolver, ctrl.Log)
	tgbReconciler := elbv2controller.NewTargetGroupBindingReconciler(mgr.GetClient(), mgr.GetFieldIndexer(), finalizerManager, tgbResManager,
		ctrl.Log.WithName("controllers").WithName("TargetGroupBinding"))
	svcReconciler := service.NewServiceReconciler(
		cloud,
		mgr.GetClient(),
		sgManager,
		sgReconciler,
		k8sClusterName,
		subnetResolver,
		ctrl.Log.WithName("controllers").WithName("Service"))
	if err = ingGroupReconciler.SetupWithManager(mgr); err != nil {
		setupLog.Error(err, "unable to create controller", "controller", "Ingress")
		os.Exit(1)
	}
	if err := tgbReconciler.SetupWithManager(context.Background(), mgr); err != nil {
		setupLog.Error(err, "unable to create controller", "controller", "TargetGroupBinding")
		os.Exit(1)
	}
	if err = svcReconciler.SetupWithManager(mgr); err != nil {
		setupLog.Error(err, "Unable to create controller", "controller", "Service")
		os.Exit(1)
	}

	podReadinessGateInjector := inject.NewPodReadinessGate(injectConfig, mgr.GetClient(), ctrl.Log.WithName("pod-readiness-gate-injector"))
	corewebhook.NewPodMutator(podReadinessGateInjector).SetupWithManager(mgr)
	elbv2webhook.NewTargetGroupBindingMutator(cloud.ELBV2(), ctrl.Log).SetupWithManager(mgr)
	elbv2webhook.NewTargetGroupBindingValidator(ctrl.Log).SetupWithManager(mgr)
	// +kubebuilder:scaffold:builder

	setupLog.Info("starting manager")
	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
		setupLog.Error(err, "problem running manager")
		os.Exit(1)
	}
}
<|endoftext|>


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

package main

import (
	"bytes"
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"os"
	"path"
	"strings"

	"golang.org/x/tools/present"
)

func printf(format string, args ...interface{}) (int, error) {
	return fmt.Fprintf(os.Stderr, format, args...)
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, `%[1]s - generates LaTeX/Beamer slides from present.

Usage of %[1]s:

$ %[1]s [input-file [output.tex]]

Examples:

$ %[1]s input.slide > out.tex
$ %[1]s input.slide out.tex
$ %[1]s < input.slide > out.tex

`,
			os.Args[0],
		)
		flag.PrintDefaults()
	}

	flag.Parse()

	var (
		r      io.Reader
		w      io.Writer
		input  = "stdin"
		output = "stdout"
	)

	switch flag.NArg() {
	case 0:
		r = os.Stdin
		w = os.Stdout
	case 1:
		input = flag.Arg(0)
		f, err := os.Open(input)
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()
		printf("input:  [%s]...\n", input)

		r = f
		w = os.Stdout

	case 2:

		input = flag.Arg(0)
		f, err := os.Open(input)
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()
		printf("input:  [%s]...\n", input)

		tex, err := os.Create(output)
		if err != nil {
			log.Fatalf("could not create output file [%s]: %v\n", output, err)
		}
		defer func() {
			err = tex.Close()
			if err != nil {
				log.Fatalf("could not close output file [%s]: %v\n", output, err)
			}
		}()

		r = f
		w = tex

	default:
		flag.Usage()
		os.Exit(2)
	}

	doc, err := present.Parse(r, input, 0)
	if err != nil {
		log.Fatal(err)
	}

	tmpl, err := initTemplates("templates")
	if err != nil {
		log.Fatal(err)
	}

	buf := new(bytes.Buffer)
	err = doc.Render(buf, tmpl)
	if err != nil {
		log.Fatal(err)
	}

	out := unescapeHTML(buf.Bytes())

	_, err = w.Write(out)
	if err != nil {
		log.Fatalf("could not fill output: %v\n", err)
	}
}

func unescapeHTML(data []byte) []byte {
	out := make([]byte, len(data))
	copy(out, data)
	for _, r := range []struct {
		old string
		new string
	}{
		{
			old: "<",
			new: "<",
		},
		{
			old: ">",
			new: ">",
		},
		{
			old: """,
			new: `"`,
		},
		{
			old: """,
			new: `"`,
		},
		{
			old: "&",
			new: "&",
		},
		{
			old: " ",
			new: " ",
		},
	} {
		out = bytes.Replace(out, []byte(r.old), []byte(r.new), -1)
	}
	return out
}

func initTemplates(base string) (*template.Template, error) {
	fname := path.Join(base, "beamer.tmpl")
	tmpl := template.New("").Funcs(funcs).Delims("<<", ">>")
	_, err := tmpl.ParseFiles(fname)
	if err != nil {
		return nil, err
	}

	return tmpl, err
}

// renderElem implements the elem template function, used to render
// sub-templates.
func renderElem(t *template.Template, e present.Elem) (template.HTML, error) {
	var data interface{} = e
	if s, ok := e.(present.Section); ok {
		data = struct {
			present.Section
			Template *template.Template
		}{s, t}
	}
	return execTemplate(t, e.TemplateName(), data)
}

var (
	funcs = template.FuncMap{}
)

func init() {
	funcs["elem"] = renderElem
	funcs["stringFromBytes"] = func(raw []byte) string { return string(raw) }
	funcs["join"] = func(lines []string) string { return strings.Join(lines, "\n") }
	funcs["nodot"] = func(s string) string {
		if strings.HasPrefix(s, ".") {
			return s[1:]
		}
		return s
	}
}

// execTemplate is a helper to execute a template and return the output as a
// template.HTML value.
func execTemplate(t *template.Template, name string, data interface{}) (template.HTML, error) {
	b := new(bytes.Buffer)
	err := t.ExecuteTemplate(b, name, data)
	if err != nil {
		return "", err
	}
	return template.HTML(b.String()), nil
}
<commit_msg>main: add -base switch to pick up templates<commit_after>// Copyright 2015 The present-tex Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/build"
	"html/template"
	"io"
	"log"
	"os"
	"path"
	"strings"

	"golang.org/x/tools/present"
)

const (
	basePkg         = "github.com/sbinet/present-tex"
	basePathMessage = `
By default, present-tex locates the slide template files and associated
static content by looking for a %q package
in your Go workspaces (GOPATH).
You may use the -base flag to specify an alternate location.
`
)

func printf(format string, args ...interface{}) (int, error) {
	return fmt.Fprintf(os.Stderr, format, args...)
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, `%[1]s - generates LaTeX/Beamer slides from present.

Usage of %[1]s:

$ %[1]s [options] [input-file [output.tex]]

Examples:

$ %[1]s input.slide > out.tex
$ %[1]s input.slide out.tex
$ %[1]s < input.slide > out.tex

Options:
`,
			os.Args[0],
		)
		flag.PrintDefaults()
	}

	tmpldir := ""
	flag.StringVar(&tmpldir, "base", "", "base path for slide templates")

	flag.Parse()

	if tmpldir == "" {
		p, err := build.Default.Import(basePkg, "", build.FindOnly)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't find present-tex files: %v\n", err)
			fmt.Fprintf(os.Stderr, basePathMessage, basePkg)
			os.Exit(1)
		}
		tmpldir = path.Join(p.Dir, "templates")
	}

	var (
		r      io.Reader
		w      io.Writer
		input  = "stdin"
		output = "stdout"
	)

	switch flag.NArg() {
	case 0:
		r = os.Stdin
		w = os.Stdout
	case 1:
		input = flag.Arg(0)
		f, err := os.Open(input)
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()
		printf("input:  [%s]...\n", input)

		r = f
		w = os.Stdout

	case 2:

		input = flag.Arg(0)
		f, err := os.Open(input)
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()
		printf("input:  [%s]...\n", input)

		tex, err := os.Create(output)
		if err != nil {
			log.Fatalf("could not create output file [%s]: %v\n", output, err)
		}
		defer func() {
			err = tex.Close()
			if err != nil {
				log.Fatalf("could not close output file [%s]: %v\n", output, err)
			}
		}()

		r = f
		w = tex

	default:
		flag.Usage()
		os.Exit(2)
	}

	doc, err := present.Parse(r, input, 0)
	if err != nil {
		log.Fatal(err)
	}

	tmpl, err := initTemplates(tmpldir)
	if err != nil {
		log.Fatal(err)
	}

	buf := new(bytes.Buffer)
	err = doc.Render(buf, tmpl)
	if err != nil {
		log.Fatal(err)
	}

	out := unescapeHTML(buf.Bytes())

	_, err = w.Write(out)
	if err != nil {
		log.Fatalf("could not fill output: %v\n", err)
	}
}

func unescapeHTML(data []byte) []byte {
	out := make([]byte, len(data))
	copy(out, data)
	for _, r := range []struct {
		old string
		new string
	}{
		{
			old: "<",
			new: "<",
		},
		{
			old: ">",
			new: ">",
		},
		{
			old: """,
			new: `"`,
		},
		{
			old: """,
			new: `"`,
		},
		{
			old: "&",
			new: "&",
		},
		{
			old: " ",
			new: " ",
		},
	} {
		out = bytes.Replace(out, []byte(r.old), []byte(r.new), -1)
	}
	return out
}

func initTemplates(base string) (*template.Template, error) {
	fname := path.Join(base, "beamer.tmpl")
	tmpl := template.New("").Funcs(funcs).Delims("<<", ">>")
	_, err := tmpl.ParseFiles(fname)
	if err != nil {
		return nil, err
	}

	return tmpl, err
}

// renderElem implements the elem template function, used to render
// sub-templates.
func renderElem(t *template.Template, e present.Elem) (template.HTML, error) {
	var data interface{} = e
	if s, ok := e.(present.Section); ok {
		data = struct {
			present.Section
			Template *template.Template
		}{s, t}
	}
	return execTemplate(t, e.TemplateName(), data)
}

var (
	funcs = template.FuncMap{}
)

func init() {
	funcs["elem"] = renderElem
	funcs["stringFromBytes"] = func(raw []byte) string { return string(raw) }
	funcs["join"] = func(lines []string) string { return strings.Join(lines, "\n") }
	funcs["nodot"] = func(s string) string {
		if strings.HasPrefix(s, ".") {
			return s[1:]
		}
		return s
	}
}

// execTemplate is a helper to execute a template and return the output as a
// template.HTML value.
func execTemplate(t *template.Template, name string, data interface{}) (template.HTML, error) {
	b := new(bytes.Buffer)
	err := t.ExecuteTemplate(b, name, data)
	if err != nil {
		return "", err
	}
	return template.HTML(b.String()), nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	"github.com/codegangsta/cli"
	"github.com/lmas/ApolloStats/src"
)

var (
	g_addr     string
	g_database string
	g_debug    bool
)

func main() {
	app := cli.NewApp()
	app.Version = apollostats.VERSION
	app.Usage = "Run a web server, serving stats for Apollo."
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:        "addr, a",
			Value:       "127.0.0.1:8000",
			Usage:       "serve web pages on this address",
			EnvVar:      "APOLLOSTATS_ADDR",
			Destination: &g_addr,
		},
		cli.StringFlag{
			Name:        "database, d",
			Value:       "user:password@/database",
			Usage:       "database authentication string",
			EnvVar:      "APOLLOSTATS_DBAUTH",
			Destination: &g_database,
		},
		cli.BoolFlag{
			Name:        "debug",
			Usage:       "run in debug mode",
			Destination: &g_debug,
		},
	}
	app.Commands = []cli.Command{
		{Name: "run", Usage: "Run the web server", Action: run_server},
		{Name: "update", Usage: "Update to the latest version", Action: run_update},
	}
	app.Run(os.Args)
}

func run_server(c *cli.Context) {
	db, e := apollostats.OpenDB(g_database)
	if e != nil {
		fmt.Printf("Failed to connect to the database:\n%s\n", e.Error())
		return
	}

	i := apollostats.Instance{
		Debug: g_debug,
		DB:    db,
	}

	i.Init()
	i.Serve(g_addr)
}

func run_update(c *cli.Context) {
}
<commit_msg>Logs db errors in a better way.<commit_after>package main

import (
	"fmt"
	"log"
	"os"

	"github.com/codegangsta/cli"
	"github.com/lmas/ApolloStats/src"
)

var (
	g_addr     string
	g_database string
	g_debug    bool
)

func main() {
	app := cli.NewApp()
	app.Version = apollostats.VERSION
	app.Usage = "Run a web server, serving stats for Apollo."
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:        "addr, a",
			Value:       "127.0.0.1:8000",
			Usage:       "serve web pages on this address",
			EnvVar:      "APOLLOSTATS_ADDR",
			Destination: &g_addr,
		},
		cli.StringFlag{
			Name:        "database, d",
			Value:       "user:password@/database",
			Usage:       "database authentication string",
			EnvVar:      "APOLLOSTATS_DBAUTH",
			Destination: &g_database,
		},
		cli.BoolFlag{
			Name:        "debug",
			Usage:       "run in debug mode",
			Destination: &g_debug,
		},
	}
	app.Commands = []cli.Command{
		{Name: "run", Usage: "Run the web server", Action: run_server},
		{Name: "update", Usage: "Update to the latest version", Action: run_update},
	}
	app.Run(os.Args)
}

func run_server(c *cli.Context) {
	db, e := apollostats.OpenDB(g_database)
	if e != nil {
		log.Fatal(fmt.Sprintf("Failed to connect to the database:\n%s\n", e.Error()))
	}

	i := apollostats.Instance{
		Debug: g_debug,
		DB:    db,
	}

	i.Init()
	i.Serve(g_addr)
}

func run_update(c *cli.Context) {
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"strings"

	flag "github.com/dotcloud/docker/pkg/mflag"

	"github.com/reiki4040/peco"
	myec2 "github.com/reiki4040/rnssh/internal/ec2"
	"github.com/reiki4040/rnssh/internal/rnssh"
)

const (
	Usage = `rnssh

usage:

  rnssh [-f] [-p] [-s] [user@]query strings ...
  rnssh -h
  rnssh -v
  rnssh --init

options:
  -f: reload ec2 instances infomaion. connect to AWS.
      you have to specify this option after you modified instances.
	  (launch, start, stop, modify name, etc...)

  -P: use Public IP address. this is default ssh host type.
      if you have set RNSSH_HOST_TYPE variable and value is private/name,
      you can ssh with Public IP temporarily with this option.

  -p: use Private IP address. for VPN/Direct connect.
      you can ssh with Private IP always with RNSSH_HOST_TYPE environment variable.

  -n: use Name tag.
      this option for ssh config that Host named by ec2 Name tag.

      # example ssh config
      Host EC2-Name-tag
        HostName X.X.X.X (or domain)
        User ec2-user
        IdentityFile path/to/key

      you can ssh with Name tag always with RNSSH_HOST_TYPE environment variable.

  -r: target region. if you set 'AWS_REGION' environment variable, use it.
      if you specify both, use -r options.

  -s: show ssh command string that would be run. (debug)

  --init: start wizard for default setting AWS region and HostType
          and save to config file (~/.rnssh/config)

options for ssh:
  -l: ssh user.
  -i: identity file path.
  --port: ssh port.
  --strict-host-key-checking-no: suppress host key checking.
                                 1: using 'ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'
                                 0: OFF
                                -1: Default(OFF)

options for help:
  -h: show this usage.
  -v: show version.

args:

  query string...: filtering ec2 instances list.

Environment variables:

  RNSSH_HOST_TYPE: this variable can specify default ssh host types.

    public  : use Public IP (default)
    private : use Private IP
    name    : use Name tag

Caution
*caution to option and arguments order*
  the options must be before arguments.

  - working
  rnssh -s query string

  - not working
  rnssh query string -s
`

	ENV_AWS_REGION      = "AWS_REGION"
	ENV_RNSSH_HOST_TYPE = "RNSSH_HOST_TYPE"
)

type CommandOption struct {
	Reload                  bool
	Region                  string
	PrivateIP               bool
	PublicIP                bool
	NameTag                 bool
	SshUser                 string
	IdentityFile            string
	Port                    int
	StrictHostKeyCheckingNo int
}

func (o *CommandOption) Validate() error {
	if err := duplicateHostTypeOption(o.PublicIP, o.PrivateIP, o.NameTag); err != nil {
		return err
	}

	if err := rnssh.IdentityFileCheck(o.IdentityFile); err != nil {
		return err
	}

	if err := rnssh.StrictHostKeyCheckingNoCheck(o.StrictHostKeyCheckingNo); err != nil {
		return err
	}

	return nil
}

func duplicateHostTypeOption(public, private, name bool) error {
	if public && private || private && name || public && name {
		return fmt.Errorf("duplicate specify option -P/-p/-n. please spcify only one.")
	}

	return nil
}

type RnsshOption struct {
	Reload                  bool
	Region                  string
	HostType                string
	SshUser                 string
	IdentityFile            string
	Port                    int
	StrictHostKeyCheckingNo int
}

var (
	version   string
	hash      string
	builddate string

	show_version bool
	show_usage   bool
	initWizard   bool
	showCommand  bool

	// command option
	opt = &CommandOption{}
)

func init() {
	flag.BoolVar(&show_version, []string{"v", "-version"}, false, "show version.")
	flag.BoolVar(&show_usage, []string{"h", "-help"}, false, "show this usage.")
	flag.BoolVar(&initWizard, []string{"-init"}, false, "run initial configuration wizard.")

	flag.BoolVar(&opt.Reload, []string{"f", "-force"}, false, "reload ec2 (force connect to AWS)")
	flag.BoolVar(&opt.PublicIP, []string{"P", "-public-ip"}, false, "ssh with EC2 Public IP")
	flag.BoolVar(&opt.PrivateIP, []string{"p", "-private-ip"}, false, "ssh with EC2 Private IP")
	flag.BoolVar(&opt.NameTag, []string{"n", "-name-tag"}, false, "ssh with EC2 Name tag")
	flag.BoolVar(&showCommand, []string{"s", "-show-command"}, false, "show ssh command that will do (debug)")

	flag.StringVar(&opt.Region, []string{"r", "-region"}, "", "specify region")
	flag.StringVar(&opt.SshUser, []string{"l", "-user"}, "", "specify ssh user")
	flag.StringVar(&opt.IdentityFile, []string{"i", "-identity-file"}, "", "specify ssh identity file")
	flag.IntVar(&opt.Port, []string{"-port"}, 0, "specify ssh port")
	flag.IntVar(&opt.StrictHostKeyCheckingNo, []string{"-strict-host-key-checking-no"}, -1, "suppress host key checking. 1: ON, 0: OFF, -1: default(OFF)")

	flag.Parse()
}

func showVersion() {
	fmt.Printf("%s (%s) built:%s\n", version, hash, builddate)
}

func usage() {
	fmt.Printf("%s\n", Usage)
}

func main() {
	if show_usage {
		usage()
		os.Exit(0)
	}

	if show_version {
		showVersion()
		os.Exit(0)
	}

	if initWizard {
		if err := rnssh.DoConfigWizard(); err != nil {
			fmt.Println(err)
			os.Exit(1)
		} else {
			fmt.Println("saved rnssh config.")
			os.Exit(0)
		}
	}

	err := opt.Validate()
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	err = rnssh.CreateRnsshDir()
	if err != nil {
		fmt.Printf("can not create rnssh dir: %s\n", err.Error())
		os.Exit(1)
	}

	conf, err := rnssh.GetConfig()
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	rOpt := mergeConfig(&conf.Default, *opt)
	if rOpt.Region == "" {
		fmt.Println("region is empty. please specify by region option (-r) or AWS_REGION envirnment variable.")
		os.Exit(1)
	}

	sshArgs, err := chooseAndGenSshArgs(rOpt, flag.Args())
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	if showCommand {
		fmt.Printf("%s %s\n", "ssh", strings.Join(sshArgs, " "))
		os.Exit(0)
	}

	cmd := exec.Command("ssh", sshArgs...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Run()
}

// merge option, config, ENV
// priority [high] option > config > ENV [low]
func mergeConfig(conf *rnssh.RnsshConfig, opt CommandOption) *RnsshOption {

	region := os.Getenv(ENV_AWS_REGION)
	if opt.Region != "" {
		region = opt.Region
	} else {
		if conf.AWSRegion != "" {
			region = conf.AWSRegion
		}
	}

	hostType := os.Getenv(ENV_RNSSH_HOST_TYPE)
	optHostType := getSshTargetType(opt.PublicIP, opt.PrivateIP, opt.NameTag)
	if optHostType != "" {
		hostType = optHostType
	} else {
		if conf.HostType != "" {
			hostType = conf.HostType
		}
	}

	sshUser := conf.SshUser
	if opt.SshUser != "" {
		sshUser = opt.SshUser
	}

	identityFile := conf.SshIdentityFile
	if opt.IdentityFile != "" {
		identityFile = opt.IdentityFile
	}

	port := conf.SshPort
	if opt.Port > 0 {
		port = opt.Port
	}

	strictHostKeyCheckingNo := conf.SshStrictHostKeyCheckingNo
	if opt.StrictHostKeyCheckingNo != -1 {
		strictHostKeyCheckingNo = opt.StrictHostKeyCheckingNo
	}

	return &RnsshOption{
		Region:       region,
		HostType:     hostType,
		SshUser:      sshUser,
		IdentityFile: identityFile,
		Port:         port,
		StrictHostKeyCheckingNo: strictHostKeyCheckingNo,
	}
}

func chooseAndGenSshArgs(rOpt *RnsshOption, cmdArgs []string) ([]string, error) {

	// support user@host format
	sshUser, hostname, err := getSshUserAndHostname(strings.Join(cmdArgs, " "))
	if err != nil {
		return nil, fmt.Errorf("%s\n", err.Error())
	}

	hostType := rnssh.HOST_TYPE_PUBLIC_IP
	if rOpt.HostType != "" {
		hostType = rOpt.HostType
	}

	handler := myec2.DefaultEC2Handler()
	choosableList, err := handler.LoadTargetHost(hostType, rOpt.Region, rOpt.Reload)
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	if len(choosableList) == 0 {
		fmt.Printf("there is no instance. not running %s\n", rOpt.Region)
		os.Exit(1)
	}

	// show ec2 instances and choose intactive
	targetHost, err := chooseTargetHost(choosableList, hostname)
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	sshHost := targetHost.Value()
	sshArgs := genSshArgs(rOpt.SshUser, rOpt.IdentityFile, rOpt.Port, rOpt.StrictHostKeyCheckingNo, sshUser, sshHost)

	return sshArgs, nil
}

func getSshUserAndHostname(sshTarget string) (string, string, error) {
	// support user@host format
	idx := strings.Index(sshTarget, "@")

	// not include(-1) or first char(0)
	if idx > 0 {
		splited := strings.SplitN(sshTarget, "@", 2)
		return splited[0], splited[1], nil
	}

	return "", sshTarget, nil
}

func chooseTargetHost(choices []rnssh.Choosable, defaultQuery string) (rnssh.Choosable, error) {

	if len(choices) == 0 {
		err := fmt.Errorf("there is no running instance.")
		return nil, err
	}

	pecoChoices := make([]peco.Choosable, 0, len(choices))
	for _, c := range choices {
		pecoChoices = append(pecoChoices, c)
	}

	pecoOpt := &peco.PecoOptions{
		OptPrompt: "which do you choose the instance? >",
		OptQuery:  defaultQuery,
	}

	result, err := peco.PecolibWithOptions(pecoChoices, pecoOpt)
	if err != nil || len(result) == 0 {
		err := fmt.Errorf("no select target.")
		return nil, err
	}

	var targetHost rnssh.Choosable
	for _, r := range result {
		if ec2host, ok := r.(rnssh.Choosable); ok {
			targetHost = ec2host
		} else {
			err := fmt.Errorf("this is bug. type is not Choosable: %v", r)
			return nil, err
		}
	}

	return targetHost, nil
}

func getSshTargetType(publicIP, privateIP, nameTag bool) string {

	// overwrite by option
	if publicIP {
		return rnssh.HOST_TYPE_PUBLIC_IP
	}

	if privateIP {
		return rnssh.HOST_TYPE_PRIVATE_IP
	}

	if nameTag {
		return rnssh.HOST_TYPE_NAME_TAG
	}

	return ""
}

func genSshArgs(optSshUser, optIdentityFile string, optPort, optStrictHostKeyCheckingNo int, sshUser, sshHost string) []string {
	args := make([]string, 0)
	if optSshUser != "" {
		args = append(args, "-l"+optSshUser)
	}

	if optIdentityFile != "" {
		args = append(args, "-i"+optIdentityFile)
	}

	if optPort > 0 {
		args = append(args, "-p"+strconv.Itoa(optPort))
	}

	if optStrictHostKeyCheckingNo == 1 {
		args = append(args, "-oStrictHostKeyChecking=no")
		args = append(args, "-oUserKnownHostsFile=/dev/null")
	}

	if sshUser != "" {
		sshHost = sshUser + "@" + sshHost
	}

	args = append(args, sshHost)

	return args
}
<commit_msg>help message to be simple.<commit_after>package main

import (
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"strings"

	flag "github.com/dotcloud/docker/pkg/mflag"

	"github.com/reiki4040/peco"
	myec2 "github.com/reiki4040/rnssh/internal/ec2"
	"github.com/reiki4040/rnssh/internal/rnssh"
)

const (
	Usage = `rnssh - easy ssh login to EC2.

usage:

  rnssh [-f] [-p] [-s] [user@]query strings ...
  rnssh --init

options:
  -f: reload ec2 instances infomaion. connect to AWS.
      you have to specify this option after you modified instances.
	  (launch, start, stop, modify name, etc...)

  -P: use Public IP address. this is default ssh host type.
  -p: use Private IP address. for VPN/Direct connect.
  -n: use Name tag.
      this option for ssh config that Host named by ec2 Name tag.

  -r: target region. you can set default by --init (~/.rnssh/config)

  -s: show ssh command string that would be run. (debug)

  --init: start wizard for default setting AWS region and rnssh host type.
          and save to config file (~/.rnssh/config)

options for ssh:
  -l: ssh user.
  -i: identity file path.
  --port: ssh port.
  --strict-host-key-checking-no: suppress host key checking.
                                 1: using 'ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'
                                 0: OFF
                                -1: Default(OFF)

options for help:
  -h: show this usage.
  -v: show version.

args:
  query string...: filtering ec2 instances list.
`

	ENV_AWS_REGION      = "AWS_REGION"
	ENV_RNSSH_HOST_TYPE = "RNSSH_HOST_TYPE"
)

type CommandOption struct {
	Reload                  bool
	Region                  string
	PrivateIP               bool
	PublicIP                bool
	NameTag                 bool
	SshUser                 string
	IdentityFile            string
	Port                    int
	StrictHostKeyCheckingNo int
}

func (o *CommandOption) Validate() error {
	if err := duplicateHostTypeOption(o.PublicIP, o.PrivateIP, o.NameTag); err != nil {
		return err
	}

	if err := rnssh.IdentityFileCheck(o.IdentityFile); err != nil {
		return err
	}

	if err := rnssh.StrictHostKeyCheckingNoCheck(o.StrictHostKeyCheckingNo); err != nil {
		return err
	}

	return nil
}

func duplicateHostTypeOption(public, private, name bool) error {
	if public && private || private && name || public && name {
		return fmt.Errorf("duplicate specify option -P/-p/-n. please spcify only one.")
	}

	return nil
}

type RnsshOption struct {
	Reload                  bool
	Region                  string
	HostType                string
	SshUser                 string
	IdentityFile            string
	Port                    int
	StrictHostKeyCheckingNo int
}

var (
	version   string
	hash      string
	builddate string

	show_version bool
	show_usage   bool
	initWizard   bool
	showCommand  bool

	// command option
	opt = &CommandOption{}
)

func init() {
	flag.BoolVar(&show_version, []string{"v", "-version"}, false, "show version.")
	flag.BoolVar(&show_usage, []string{"h", "-help"}, false, "show this usage.")
	flag.BoolVar(&initWizard, []string{"-init"}, false, "run initial configuration wizard.")

	flag.BoolVar(&opt.Reload, []string{"f", "-force"}, false, "reload ec2 (force connect to AWS)")
	flag.BoolVar(&opt.PublicIP, []string{"P", "-public-ip"}, false, "ssh with EC2 Public IP")
	flag.BoolVar(&opt.PrivateIP, []string{"p", "-private-ip"}, false, "ssh with EC2 Private IP")
	flag.BoolVar(&opt.NameTag, []string{"n", "-name-tag"}, false, "ssh with EC2 Name tag")
	flag.BoolVar(&showCommand, []string{"s", "-show-command"}, false, "show ssh command that will do (debug)")

	flag.StringVar(&opt.Region, []string{"r", "-region"}, "", "specify region")
	flag.StringVar(&opt.SshUser, []string{"l", "-user"}, "", "specify ssh user")
	flag.StringVar(&opt.IdentityFile, []string{"i", "-identity-file"}, "", "specify ssh identity file")
	flag.IntVar(&opt.Port, []string{"-port"}, 0, "specify ssh port")
	flag.IntVar(&opt.StrictHostKeyCheckingNo, []string{"-strict-host-key-checking-no"}, -1, "suppress host key checking. 1: ON, 0: OFF, -1: default(OFF)")

	flag.Parse()
}

func showVersion() {
	fmt.Printf("%s (%s) built:%s\n", version, hash, builddate)
}

func usage() {
	fmt.Printf("%s\n", Usage)
}

func main() {
	if show_usage {
		usage()
		os.Exit(0)
	}

	if show_version {
		showVersion()
		os.Exit(0)
	}

	if initWizard {
		if err := rnssh.DoConfigWizard(); err != nil {
			fmt.Println(err)
			os.Exit(1)
		} else {
			fmt.Println("saved rnssh config.")
			os.Exit(0)
		}
	}

	err := opt.Validate()
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	err = rnssh.CreateRnsshDir()
	if err != nil {
		fmt.Printf("can not create rnssh dir: %s\n", err.Error())
		os.Exit(1)
	}

	conf, err := rnssh.GetConfig()
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	rOpt := mergeConfig(&conf.Default, *opt)
	if rOpt.Region == "" {
		fmt.Println("region is empty. please specify by region option (-r) or AWS_REGION envirnment variable.")
		os.Exit(1)
	}

	sshArgs, err := chooseAndGenSshArgs(rOpt, flag.Args())
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	if showCommand {
		fmt.Printf("%s %s\n", "ssh", strings.Join(sshArgs, " "))
		os.Exit(0)
	}

	cmd := exec.Command("ssh", sshArgs...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Run()
}

// merge option, config, ENV
// priority [high] option > config > ENV [low]
func mergeConfig(conf *rnssh.RnsshConfig, opt CommandOption) *RnsshOption {

	region := os.Getenv(ENV_AWS_REGION)
	if opt.Region != "" {
		region = opt.Region
	} else {
		if conf.AWSRegion != "" {
			region = conf.AWSRegion
		}
	}

	hostType := os.Getenv(ENV_RNSSH_HOST_TYPE)
	optHostType := getSshTargetType(opt.PublicIP, opt.PrivateIP, opt.NameTag)
	if optHostType != "" {
		hostType = optHostType
	} else {
		if conf.HostType != "" {
			hostType = conf.HostType
		}
	}

	sshUser := conf.SshUser
	if opt.SshUser != "" {
		sshUser = opt.SshUser
	}

	identityFile := conf.SshIdentityFile
	if opt.IdentityFile != "" {
		identityFile = opt.IdentityFile
	}

	port := conf.SshPort
	if opt.Port > 0 {
		port = opt.Port
	}

	strictHostKeyCheckingNo := conf.SshStrictHostKeyCheckingNo
	if opt.StrictHostKeyCheckingNo != -1 {
		strictHostKeyCheckingNo = opt.StrictHostKeyCheckingNo
	}

	return &RnsshOption{
		Region:       region,
		HostType:     hostType,
		SshUser:      sshUser,
		IdentityFile: identityFile,
		Port:         port,
		StrictHostKeyCheckingNo: strictHostKeyCheckingNo,
	}
}

func chooseAndGenSshArgs(rOpt *RnsshOption, cmdArgs []string) ([]string, error) {

	// support user@host format
	sshUser, hostname, err := getSshUserAndHostname(strings.Join(cmdArgs, " "))
	if err != nil {
		return nil, fmt.Errorf("%s\n", err.Error())
	}

	hostType := rnssh.HOST_TYPE_PUBLIC_IP
	if rOpt.HostType != "" {
		hostType = rOpt.HostType
	}

	handler := myec2.DefaultEC2Handler()
	choosableList, err := handler.LoadTargetHost(hostType, rOpt.Region, rOpt.Reload)
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	if len(choosableList) == 0 {
		fmt.Printf("there is no instance. not running %s\n", rOpt.Region)
		os.Exit(1)
	}

	// show ec2 instances and choose intactive
	targetHost, err := chooseTargetHost(choosableList, hostname)
	if err != nil {
		fmt.Printf("%s\n", err.Error())
		os.Exit(1)
	}

	sshHost := targetHost.Value()
	sshArgs := genSshArgs(rOpt.SshUser, rOpt.IdentityFile, rOpt.Port, rOpt.StrictHostKeyCheckingNo, sshUser, sshHost)

	return sshArgs, nil
}

func getSshUserAndHostname(sshTarget string) (string, string, error) {
	// support user@host format
	idx := strings.Index(sshTarget, "@")

	// not include(-1) or first char(0)
	if idx > 0 {
		splited := strings.SplitN(sshTarget, "@", 2)
		return splited[0], splited[1], nil
	}

	return "", sshTarget, nil
}

func chooseTargetHost(choices []rnssh.Choosable, defaultQuery string) (rnssh.Choosable, error) {

	if len(choices) == 0 {
		err := fmt.Errorf("there is no running instance.")
		return nil, err
	}

	pecoChoices := make([]peco.Choosable, 0, len(choices))
	for _, c := range choices {
		pecoChoices = append(pecoChoices, c)
	}

	pecoOpt := &peco.PecoOptions{
		OptPrompt: "which do you choose the instance? >",
		OptQuery:  defaultQuery,
	}

	result, err := peco.PecolibWithOptions(pecoChoices, pecoOpt)
	if err != nil || len(result) == 0 {
		err := fmt.Errorf("no select target.")
		return nil, err
	}

	var targetHost rnssh.Choosable
	for _, r := range result {
		if ec2host, ok := r.(rnssh.Choosable); ok {
			targetHost = ec2host
		} else {
			err := fmt.Errorf("this is bug. type is not Choosable: %v", r)
			return nil, err
		}
	}

	return targetHost, nil
}

func getSshTargetType(publicIP, privateIP, nameTag bool) string {

	// overwrite by option
	if publicIP {
		return rnssh.HOST_TYPE_PUBLIC_IP
	}

	if privateIP {
		return rnssh.HOST_TYPE_PRIVATE_IP
	}

	if nameTag {
		return rnssh.HOST_TYPE_NAME_TAG
	}

	return ""
}

func genSshArgs(optSshUser, optIdentityFile string, optPort, optStrictHostKeyCheckingNo int, sshUser, sshHost string) []string {
	args := make([]string, 0)
	if optSshUser != "" {
		args = append(args, "-l"+optSshUser)
	}

	if optIdentityFile != "" {
		args = append(args, "-i"+optIdentityFile)
	}

	if optPort > 0 {
		args = append(args, "-p"+strconv.Itoa(optPort))
	}

	if optStrictHostKeyCheckingNo == 1 {
		args = append(args, "-oStrictHostKeyChecking=no")
		args = append(args, "-oUserKnownHostsFile=/dev/null")
	}

	if sshUser != "" {
		sshHost = sshUser + "@" + sshHost
	}

	args = append(args, sshHost)

	return args
}
<|endoftext|>


<commit_before>// +build go1.5

package main

import (
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/tools/go/ast/astutil"
	"golang.org/x/tools/go/loader"
)

var (
	verbose   = flag.Bool("debug", false, "verbose")
	logToFile = flag.Bool("log", false, "log to tmp file")
	filename  = flag.String("f", "", "file")
	offset    = flag.Int("o", 0, "offset in file, 1 based")
	stdin     = flag.Bool("i", false, "read file from stdin")
	// 2-pass is the fastest and loader is slowest
	pass = flag.Int("p", 1, "parse algorithm. can be 1, 2 or 3. 2 is the fastest. 3 is slowest but most reliable")
)

var logf *os.File

func lg(format string, arg ...interface{}) {
	if logf != nil {
		fmt.Fprintf(logf, format+"\n", arg...)
	} else if *verbose {
		log.Printf(format, arg...)
	}
}

func parseFile(fn string, src interface{}) *ast.File {
	f, err := parser.ParseFile(fset, fn, src, parser.AllErrors)
	if err != nil {
		// error is expected
		lg("parse file=%s err=%v", fn, err)
	}
	return f
}

func pkgPath(fn string) string {
	dir := filepath.Dir(fn)
	i := strings.LastIndex(dir, "/src/")
	if i < 0 {
		return "main"
	}
	return dir[i+len("/src/"):]
}

func tokenFile(f *ast.File) *token.File {
	return fset.File(f.Package)
}

func printPos(pos token.Pos) string {
	return fset.Position(pos).String()
}

func pkgFiles(p string) (files, imports []string) {
	pkg, err := build.Import(p, "", 0)
	if err != nil {
		lg("import pkg=%s err=%v pkg=%+v", p, err, pkg)
		return nil, nil
	}

	isTest := isTestFile(*filename)
	n := len(pkg.GoFiles)
	m := n
	if isTest {
		n += len(pkg.TestGoFiles)
	}

	out := make([]string, n)
	for i, f := range pkg.GoFiles {
		out[i] = filepath.Join(pkg.Dir, f)
	}
	if isTest {
		for i, f := range pkg.TestGoFiles {
			out[m+i] = filepath.Join(pkg.Dir, f)
		}
	}
	return out, pkg.Imports
}

func isTestFile(fn string) bool {
	fn = strings.TrimSuffix(fn, filepath.Ext(fn))
	return strings.HasSuffix(fn, "_test")
}

func parseMyPkg() (myPkg string, fs []*ast.File, imports []string, chain []ast.Node) {
	myPkg = pkgPath(*filename)
	fns, imports := pkgFiles(myPkg)
	if fns == nil {
		fns = []string{*filename}
	}
	lg("mypkg=%s: files=%v imports=%v", myPkg, fns, imports)

	for _, fn := range fns {
		var f *ast.File
		if fn == *filename && *stdin {
			b, err := ioutil.ReadAll(os.Stdin)
			if err != nil {
				log.Fatalf("read stdin err=%v", err)
			}
			f = parseFile(fn, b)
		} else {
			f = parseFile(fn, nil)
		}

		fs = append(fs, f)

		if fn == *filename {
			tf := tokenFile(f)
			pos := tf.Pos(*offset)
			if pos == token.NoPos {
				lg("no pos found")
				fail()
			}
			lg("pos=%v", printPos(pos))
			chain, _ = astutil.PathEnclosingInterval(f, pos, pos+1)
		}
	}
	return myPkg, fs, imports, chain
}

func findIdent(chain []ast.Node) *ast.Ident {
	for _, nd := range chain {
		if id, ok := nd.(*ast.Ident); ok {
			return id
		}
	}
	return nil
}

type srcImporter struct {
	cfg types.Config
}

func newSrcImporter() srcImporter {
	return srcImporter{
		cfg: types.Config{
			Importer: importer.Default(),
			Error:    func(err error) {},
			DisableUnusedImportCheck: true,
		},
	}
}

func importSrcPkg(cfg *types.Config, path string) (*types.Package, error) {
	fns, _ := pkgFiles(path)
	var fs []*ast.File
	for _, fn := range fns {
		f := parseFile(fn, nil)
		fs = append(fs, f)
	}

	return cfg.Check(path, fset, fs, nil)
}

func (si srcImporter) Import(path string) (*types.Package, error) {
	return importSrcPkg(&si.cfg, path)
}

type hybridImporter struct {
	cfg     types.Config
	imports []string
}

func newHybridImporter(imports ...string) *hybridImporter {
	return &hybridImporter{
		imports: imports,
		cfg: types.Config{
			Importer: importer.Default(),
			Error:    func(err error) {},
			DisableUnusedImportCheck: true,
		},
	}
}

func (hi *hybridImporter) Import(path string) (*types.Package, error) {
	if !hi.contains(path) {
		lg("import pkg=%s using default importer", path)
		if pkg, err := hi.cfg.Importer.Import(path); err == nil {
			return pkg, nil
		} else {
			lg("import pkg=%s using default importer err=%v. fall back to source importer", path, err)
		}
	}

	return importSrcPkg(&hi.cfg, path)
}

func (hi *hybridImporter) contains(path string) bool {
	for _, p := range hi.imports {
		if p == path {
			return true
		}
	}
	return false
}

func onePass(myPkg string, fs []*ast.File, imports []string, target *ast.Ident) {
	cfg := types.Config{
		Importer: newHybridImporter(imports...),
		Error:    func(err error) {},
		DisableUnusedImportCheck: true,
	}
	info := types.Info{
		Uses: make(map[*ast.Ident]types.Object),
	}
	cfg.Check(myPkg, fset, fs, &info)
	obj := info.Uses[target]
	lg("target=%v uses obj=%v", target, obj)
	printTargetObj(obj)
}

// compatible with godef
func fail() {
	fmt.Fprintln(os.Stderr, "godef: no identifier found")
	os.Exit(2)
}

func printTargetObj(obj types.Object) {
	if obj != nil && obj.Pos() != token.NoPos {
		fmt.Println(printPos(obj.Pos()))
	} else {
		fail()
	}
}

func twoPass(myPkg string, fs []*ast.File, target *ast.Ident) {
	// first pass to find out the package of target
	cfg := types.Config{
		Importer: newHybridImporter(""),
		Error:    func(err error) {},
		DisableUnusedImportCheck: true,
	}
	info := types.Info{
		Uses: make(map[*ast.Ident]types.Object),
	}
	cfg.Check(myPkg, fset, fs, &info)

	obj := info.Uses[target]
	if obj == nil {
		lg("object of target=%v not found", target)
		fail()
	}
	// BUG:https://github.com/golang/go/issues/13898
	otherPkg := obj.Pkg().Path()
	lg("obj of target=%v is %v in pkg=%s", target, obj, otherPkg)

	if otherPkg == myPkg {
		lg("found in mypkg")
		printTargetObj(obj)
		return
	}

	// second pass to find out the object of target
	cfg.Importer = newHybridImporter(otherPkg)
	info.Uses = make(map[*ast.Ident]types.Object)
	cfg.Check(myPkg, fset, fs, &info)
	obj = info.Uses[target]
	lg("target=%v in otherpkg obj=%v", target, obj)
	printTargetObj(obj)
}

func parseProgram(myPkg string, fs []*ast.File, target *ast.Ident) {
	cfg := loader.Config{
		Fset:       fset,
		ParserMode: parser.AllErrors,
		TypeChecker: types.Config{
			Error: func(error) {},
			DisableUnusedImportCheck: true,
		},
		TypeCheckFuncBodies: func(path string) bool {
			return path == myPkg
		},
		AllowErrors: true,
	}
	os.Setenv("CGO_ENABLED", "0")

	cfg.CreateFromFiles(myPkg, fs...)
	prog, err := cfg.Load()
	if err != nil {
		lg("load program err=%v", err)
	}

	obj := findIdentObj(prog, target)
	lg("target=%v uses obj=%v", target, obj)
	printTargetObj(obj)
}

func findIdentObj(prog *loader.Program, target *ast.Ident) types.Object {
	pkg := prog.Created[0]
	for id, obj := range pkg.Defs {
		if target == id {
			return obj
		}
	}
	for id, obj := range pkg.Uses {
		if target == id {
			return obj
		}
	}
	return nil
}

var fset = token.NewFileSet()

func main() {
	// flags used by godef
	flag.Bool("A", false, "")
	flag.Bool("a", false, "")
	flag.Bool("acme", false, "")
	flag.Bool("t", false, "")
	flag.Parse()

	if *logToFile {
		logf, _ = ioutil.TempFile(os.TempDir(), "gosym-log.")
	}

	*filename, _ = filepath.Abs(*filename)

	// offset is 1-based, but token.File.Offset is 0-based.
	*offset--
	if *offset < 0 {
		fail()
	}

	lg("args=%v", os.Args)

	myPkg, fs, imports, chain := parseMyPkg()
	target := findIdent(chain)
	if target == nil {
		fail()
	}
	lg("target is %v@%v", target, printPos(target.Pos()))

	switch *pass {
	case 1:
		onePass(myPkg, fs, imports, target)
	case 2:
		twoPass(myPkg, fs, target)
	case 3:
		parseProgram(myPkg, fs, target)
	}
}
<commit_msg>gosym: a parallel mode<commit_after>// +build go1.5

package main

import (
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/tools/go/ast/astutil"
	"golang.org/x/tools/go/loader"
)

var (
	verbose   = flag.Bool("debug", false, "verbose")
	logToFile = flag.Bool("log", false, "log to tmp file")
	filename  = flag.String("f", "", "file")
	offset    = flag.Int("o", 0, "offset in file, 1 based")
	stdin     = flag.Bool("i", false, "read file from stdin")
	pass      = flag.Int("p", 4, "parse algorithm")
)

var logf *os.File

func lg(format string, arg ...interface{}) {
	if logf != nil {
		fmt.Fprintf(logf, format+"\n", arg...)
	} else if *verbose {
		log.Printf(format, arg...)
	}
}

func parseFile(fn string, src interface{}) *ast.File {
	f, err := parser.ParseFile(fset, fn, src, parser.AllErrors)
	if err != nil {
		// error is expected
		lg("parse file=%s err=%v", fn, err)
	}
	return f
}

func pkgPath(fn string) string {
	dir := filepath.Dir(fn)
	i := strings.LastIndex(dir, "/src/")
	if i < 0 {
		return "main"
	}
	return dir[i+len("/src/"):]
}

func tokenFile(f *ast.File) *token.File {
	return fset.File(f.Package)
}

func printPos(pos token.Pos) string {
	return fset.Position(pos).String()
}

func pkgFiles(p string) (files, imports []string) {
	pkg, err := build.Import(p, "", 0)
	if err != nil {
		lg("import pkg=%s err=%v pkg=%+v", p, err, pkg)
		return nil, nil
	}

	isTest := isTestFile(*filename)
	n := len(pkg.GoFiles)
	m := n
	if isTest {
		n += len(pkg.TestGoFiles)
	}

	out := make([]string, n)
	for i, f := range pkg.GoFiles {
		out[i] = filepath.Join(pkg.Dir, f)
	}
	if isTest {
		for i, f := range pkg.TestGoFiles {
			out[m+i] = filepath.Join(pkg.Dir, f)
		}
	}
	return out, pkg.Imports
}

func isTestFile(fn string) bool {
	fn = strings.TrimSuffix(fn, filepath.Ext(fn))
	return strings.HasSuffix(fn, "_test")
}

func parseMyPkg() (myPkg string, fs []*ast.File, imports []string, chain []ast.Node) {
	myPkg = pkgPath(*filename)
	fns, imports := pkgFiles(myPkg)
	if fns == nil {
		fns = []string{*filename}
	}
	lg("mypkg=%s: files=%v imports=%v", myPkg, fns, imports)

	for _, fn := range fns {
		var f *ast.File
		if fn == *filename && *stdin {
			b, err := ioutil.ReadAll(os.Stdin)
			if err != nil {
				log.Fatalf("read stdin err=%v", err)
			}
			f = parseFile(fn, b)
		} else {
			f = parseFile(fn, nil)
		}

		fs = append(fs, f)

		if fn == *filename {
			tf := tokenFile(f)
			pos := tf.Pos(*offset)
			if pos == token.NoPos {
				lg("no pos found")
				fail()
			}
			lg("pos=%v", printPos(pos))
			chain, _ = astutil.PathEnclosingInterval(f, pos, pos+1)
		}
	}
	return myPkg, fs, imports, chain
}

func findIdent(chain []ast.Node) *ast.Ident {
	for _, nd := range chain {
		if id, ok := nd.(*ast.Ident); ok {
			return id
		}
	}
	return nil
}

type srcImporter struct {
	cfg types.Config
}

func newSrcImporter() srcImporter {
	return srcImporter{
		cfg: types.Config{
			Importer: importer.Default(),
			Error:    func(err error) {},
			DisableUnusedImportCheck: true,
		},
	}
}

func importSrcPkg(cfg *types.Config, path string) (*types.Package, error) {
	fns, _ := pkgFiles(path)
	var fs []*ast.File
	for _, fn := range fns {
		f := parseFile(fn, nil)
		fs = append(fs, f)
	}

	return cfg.Check(path, fset, fs, nil)
}

func (si srcImporter) Import(path string) (*types.Package, error) {
	return importSrcPkg(&si.cfg, path)
}

type hybridImporter struct {
	cfg     types.Config
	imports []string
}

func newHybridImporter(imports ...string) *hybridImporter {
	hi := &hybridImporter{
		imports: imports,
		cfg: types.Config{
			Importer: importer.Default(),
			Error:    func(err error) {},
			DisableUnusedImportCheck: true,
		},
	}
	return hi
}

func (hi *hybridImporter) Import(path string) (pkg *types.Package, err error) {
	if !hi.contains(path) {
		lg("import pkg=%s using default importer", path)
		if pkg, err = hi.cfg.Importer.Import(path); err == nil {
			return pkg, nil
		}
		lg("import pkg=%s using default importer err=%v. fall back to source importer", path, err)
	}

	if pkg, err = importSrcPkg(&hi.cfg, path); err == nil {
		return pkg, nil
	}
	lg("import source pkg=%s err=%v pkg=%v", path, err, pkg)
	return pkg, err
}

func (hi *hybridImporter) contains(path string) bool {
	for _, p := range hi.imports {
		if p == path {
			return true
		}
	}
	return false
}

func onePass(myPkg string, fs []*ast.File, imports []string, target *ast.Ident) types.Object {
	cfg := types.Config{
		Importer: newHybridImporter(imports...),
		Error:    func(err error) {},
		DisableUnusedImportCheck: true,
	}
	info := types.Info{
		Uses: make(map[*ast.Ident]types.Object),
	}
	cfg.Check(myPkg, fset, fs, &info)
	return info.Uses[target]
}

// compatible with godef
func fail() {
	fmt.Fprintln(os.Stderr, "godef: no identifier found")
	os.Exit(2)
}

func printTargetObj(obj types.Object) {
	if obj != nil && obj.Pos() != token.NoPos {
		fmt.Println(printPos(obj.Pos()))
	} else {
		fail()
	}
}

func findInMyPkg(myPkg string, fs []*ast.File, target *ast.Ident) (obj types.Object, otherPkg string) {
	cfg := types.Config{
		Importer: importer.Default(),
		Error:    func(err error) {},
		DisableUnusedImportCheck: true,
	}
	info := types.Info{
		Uses: make(map[*ast.Ident]types.Object),
	}
	cfg.Check(myPkg, fset, fs, &info)

	if obj = info.Uses[target]; obj == nil {
		lg("object of target=%v not found", target)
		fail()
	}
	// BUG:https://github.com/golang/go/issues/13898
	otherPkg = obj.Pkg().Path()
	lg("obj of target=%v is %v in pkg=%s", target, obj, otherPkg)

	if otherPkg == myPkg {
		return obj, myPkg
	}
	return nil, otherPkg
}

func twoPass(myPkg string, fs []*ast.File, target *ast.Ident) types.Object {
	// first pass to find out the package of target
	obj, otherPkg := findInMyPkg(myPkg, fs, target)
	if obj != nil {
		return obj
	}

	// second pass to find out the object of target in otherPkg
	cfg := types.Config{
		Importer: newHybridImporter(otherPkg),
		Error:    func(err error) {},
		DisableUnusedImportCheck: true,
	}
	cfg.Importer = newHybridImporter(otherPkg)
	info := types.Info{
		Uses: make(map[*ast.Ident]types.Object),
	}
	cfg.Check(myPkg, fset, fs, &info)
	return info.Uses[target]
}

func parseProgram(myPkg string, fs []*ast.File, target *ast.Ident) types.Object {
	cfg := loader.Config{
		Fset:       fset,
		ParserMode: parser.AllErrors,
		TypeChecker: types.Config{
			Error: func(error) {},
			DisableUnusedImportCheck: true,
		},
		TypeCheckFuncBodies: func(path string) bool {
			return path == myPkg
		},
		AllowErrors: true,
	}
	os.Setenv("CGO_ENABLED", "0")

	cfg.CreateFromFiles(myPkg, fs...)
	prog, err := cfg.Load()
	if err != nil {
		lg("load program err=%v", err)
	}

	return findIdentObj(prog, target)
}

func findIdentObj(prog *loader.Program, target *ast.Ident) types.Object {
	pkg := prog.Created[0]
	for id, obj := range pkg.Uses {
		if target == id {
			return obj
		}
	}
	for id, obj := range pkg.Defs {
		if target == id {
			return obj
		}
	}
	return nil
}

func parallelPass(myPkg string, fs []*ast.File, target *ast.Ident) types.Object {
	out := make(chan types.Object, 2)
	go func() {
		obj, _ := findInMyPkg(myPkg, fs, target)
		out <- obj
	}()
	go func() {
		out <- parseProgram(myPkg, fs, target)
	}()
	for i := 0; i < 2; i++ {
		obj := <-out
		if obj != nil {
			return obj
		}
	}
	return nil
}

var fset = token.NewFileSet()

func main() {
	// flags used by godef
	flag.Bool("A", false, "")
	flag.Bool("a", false, "")
	flag.Bool("acme", false, "")
	flag.Bool("t", false, "")
	flag.Parse()

	if *logToFile {
		logf, _ = ioutil.TempFile(os.TempDir(), "gosym-log.")
	}

	*filename, _ = filepath.Abs(*filename)

	// offset is 1-based, but token.File.Offset is 0-based.
	*offset--
	if *offset < 0 {
		fail()
	}

	lg("args=%v", os.Args)

	myPkg, fs, imports, chain := parseMyPkg()
	target := findIdent(chain)
	if target == nil {
		fail()
	}
	lg("target is %v@%v", target, printPos(target.Pos()))

	// TODO: once issue 13898 is fixed, we can make two-pass work.
	var obj types.Object
	switch *pass {
	case 1:
		obj = onePass(myPkg, fs, imports, target)
	case 2:
		obj = twoPass(myPkg, fs, target)
	case 3:
		obj = parseProgram(myPkg, fs, target)
	case 4:
		obj = parallelPass(myPkg, fs, target)
	}
	lg("target=%v in otherpkg obj=%v", target, obj)
	printTargetObj(obj)
}
<|endoftext|>


<commit_before>// This software is direct fork of https://github.com/knq/chromedp/tree/master/cmd/chromedp-proxy
// with couple of features added
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"path"

	"errors"

	"github.com/Sirupsen/logrus"
)

func main() {
	flag.Parse()

	mux := http.NewServeMux()

	simpleReverseProxy := httputil.NewSingleHostReverseProxy(&url.URL{Scheme: "http", Host: *flagRemote})

	mux.Handle("/json", simpleReverseProxy)
	mux.Handle("/", simpleReverseProxy)

	rootLogger, err := createLogger("connection")
	if err != nil {
		panic(fmt.Sprintf("could not create logger: %s", err))
	}

	logger := rootLogger.WithFields(logrus.Fields{
		fieldLevel: levelConnection,
	})

	mux.HandleFunc("/devtools/page/", func(res http.ResponseWriter, req *http.Request) {

		stream := make(chan *protocolMessage, 1024)
		id := path.Base(req.URL.Path)

		var protocolLogger *logrus.Entry

		if *flagDistributeLogs {
			logger, err := createLogger("inspector-" + id)
			if err != nil {
				panic(fmt.Sprintf("could not create logger: %s", err))
			}

			protocolLogger = logger.WithFields(logrus.Fields{
				fieldLevel:       levelConnection,
				fieldInspectorID: id,
			})

		} else {
			protocolLogger = logger.WithFields(logrus.Fields{
				fieldInspectorID: id,
			})
		}

		go dumpStream(protocolLogger, stream)

		endpoint := "ws://" + *flagRemote + "/devtools/page/" + id

		logger.Infof("---------- connection from %s ----------", req.RemoteAddr)
		logger.Infof("checking protocol versions on: %s", endpoint)

		ver, err := checkVersion()
		if err != nil {
			protocolLogger.Errorf("could not check version: %v", err)
			http.Error(res, "could not check version", 500)
			return
		}

		logger.Infof("protocol version: %s", ver["Protocol-Version"])
		logger.Infof("versions: Chrome(%s), V8(%s), Webkit(%s)", ver["Browser"], ver["V8-Version"], ver["WebKit-Version"])
		logger.Infof("browser user agent: %s", ver["User-Agent"])
		logger.Infof("connecting to %s... ", endpoint)

		// connecting to ws
		out, pres, err := wsDialer.Dial(endpoint, nil)
		if err != nil {
			msg := fmt.Sprintf("could not connect to %s: %v", endpoint, err)
			logger.Error(protocolError(msg))
			http.Error(res, msg, 500)
			return
		}
		defer pres.Body.Close()
		defer out.Close()

		// connect incoming websocket
		logger.Infof("upgrading connection on %s...", req.RemoteAddr)
		in, err := wsUpgrader.Upgrade(res, req, nil)
		if err != nil {
			logger.Errorf("could not upgrade websocket from %s: %v", req.RemoteAddr, err)
			http.Error(res, "could not upgrade websocket connection", 500)
			return
		}
		defer in.Close()

		ctxt, cancel := context.WithCancel(context.Background())
		defer cancel()

		errc := make(chan error, 1)
		go proxyWS(ctxt, stream, in, out, errc)
		go proxyWS(ctxt, stream, out, in, errc)

		<-errc
		logger.Infof("---------- closing %s ----------", req.RemoteAddr)

		if *flagDistributeLogs {
			if closer, ok := protocolLogger.Logger.Out.(io.Closer); ok {
				closer.Close()
			}
		}

		if *flagOnce {
			os.Exit(0)
		}
	})

	log.Fatal(http.ListenAndServe(*flagListen, mux))
}

func dumpStream(logger *logrus.Entry, stream chan *protocolMessage) {
	logger.Printf("Legend: %s, %s, %s, %s, %s, %s", protocolColor("protocol informations"),
		eventsColor("received events"),
		requestColor("sent request frames"),
		requestReplyColor("requests params"),
		responseColor("received responses"),
		errorColor("error response."),
	)

	requests := make(map[uint64]*protocolMessage)
	targetRequests := make(map[uint64]*protocolMessage)

	for {
		select {
		case msg := <-stream:
			if msg.InTarget() {

				var targetLogger *logrus.Entry

				if *flagDistributeLogs {
					logger, err := createLogger(fmt.Sprintf("target-%s", msg.TargetID()))
					if err != nil {
						panic(fmt.Sprintf("could not create logger: %v", err))
					}

					targetLogger = logger.WithFields(logrus.Fields{
						fieldLevel:    levelTarget,
						fieldTargetID: msg.TargetID(),
					})

				} else {
					targetLogger = logger.WithFields(logrus.Fields{
						fieldLevel:    levelTarget,
						fieldTargetID: msg.TargetID(),
					})
				}

				if msg.IsRequest() {
					requests[msg.ID] = nil

					if protocolMessage, err := decodeMessage([]byte(asString(msg.Params["message"]))); err == nil {
						targetRequests[protocolMessage.ID] = protocolMessage

						if *flagShowRequests {
							targetLogger.WithFields(logrus.Fields{
								fieldType:   typeRequest,
								fieldMethod: protocolMessage.Method,
							}).Info(serialize(protocolMessage.Params))
						}

					} else {
						logger.WithFields(logrus.Fields{
							fieldLevel: levelConnection,
						}).Errorf("Could not deserialize message: %+v", err)
					}
				}

				if msg.IsEvent() {
					if protocolMessage, err := decodeMessage([]byte(asString(msg.Params["message"]))); err == nil {
						if protocolMessage.IsEvent() {
							targetLogger.WithFields(logrus.Fields{
								fieldType:   typeEvent,
								fieldMethod: protocolMessage.Method,
							}).Info(serialize(protocolMessage.Params))
						}

						if protocolMessage.IsResponse() {
							var logMessage string
							var logType int
							var logRequest string
							var logMethod string

							if protocolMessage.IsError() {
								logMessage = serialize(protocolMessage.Error)
								logType = typeRequestResponseError
							} else {
								logMessage = serialize(protocolMessage.Result)
								logType = typeRequestResponse
							}

							if request, ok := targetRequests[protocolMessage.ID]; ok && request != nil {
								delete(targetRequests, protocolMessage.ID)
								logRequest = serialize(request.Params)
								logMethod = request.Method

							} else {
								logRequest = errorColor("could not find request with id: %d", protocolMessage.ID)
							}

							targetLogger.WithFields(logrus.Fields{
								fieldType:    logType,
								fieldMethod:  logMethod,
								fieldRequest: logRequest,
							}).Info(logMessage)
						}
					} else {
						logger.WithFields(logrus.Fields{
							fieldLevel: levelConnection,
						}).Errorf("Could not deserialize message: %+v", err)
					}
				}

			} else {
				protocolLogger := logger.WithFields(logrus.Fields{
					fieldLevel:    levelProtocol,
					fieldTargetID: protocolTargetID,
				})

				if msg.IsRequest() {
					requests[msg.ID] = msg

					if *flagShowRequests {
						protocolLogger.WithFields(logrus.Fields{
							fieldType:   typeRequest,
							fieldMethod: msg.Method,
						}).Info(serialize(msg.Params))
					}
				}

				if msg.IsResponse() {

					var logMessage string
					var logType int
					var logRequest string
					var logMethod string

					if msg.IsError() {
						logMessage = serialize(msg.Error)
						logType = typeRequestResponseError
					} else {
						logMessage = serialize(msg.Result)
						logType = typeRequestResponse
					}

					if request, ok := requests[msg.ID]; ok && request != nil {
						logRequest = serialize(request.Params)
						logMethod = request.Method

						delete(requests, msg.ID)

						protocolLogger.WithFields(logrus.Fields{
							fieldType:    logType,
							fieldMethod:  logMethod,
							fieldRequest: logRequest,
						}).Info(logMessage)
					}
				}

				if msg.IsEvent() {
					protocolLogger.WithFields(logrus.Fields{
						fieldType:   typeEvent,
						fieldMethod: msg.Method,
					}).Info(serialize(msg.Params))
				}
			}
		}
	}
}

func checkVersion() (map[string]string, error) {
	cl := &http.Client{}
	req, err := http.NewRequest("GET", "http://"+*flagRemote+"/json/version", nil)
	if err != nil {
		return nil, err
	}

	res, err := cl.Do(req)
	if err != nil {
		return nil, err
	}

	defer res.Body.Close()

	var v map[string]string
	if err := json.NewDecoder(res.Body).Decode(&v); err != nil {
		return nil, errors.New("expected json result")
	}

	return v, nil
}
<commit_msg>#1 Support /devtools/browser/{id} endpoints<commit_after>// This software is direct fork of https://github.com/knq/chromedp/tree/master/cmd/chromedp-proxy
// with couple of features added
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"path"

	"errors"

	"github.com/Sirupsen/logrus"
)

func main() {
	flag.Parse()

	mux := http.NewServeMux()

	simpleReverseProxy := httputil.NewSingleHostReverseProxy(&url.URL{Scheme: "http", Host: *flagRemote})

	mux.Handle("/json", simpleReverseProxy)
	mux.Handle("/", simpleReverseProxy)

	rootLogger, err := createLogger("connection")
	if err != nil {
		panic(fmt.Sprintf("could not create logger: %s", err))
	}

	logger := rootLogger.WithFields(logrus.Fields{
		fieldLevel: levelConnection,
	})

	handlerFunc := func(basePath string) func(http.ResponseWriter, *http.Request) {
		return func(res http.ResponseWriter, req *http.Request) {

			stream := make(chan *protocolMessage, 1024)
			id := path.Base(req.URL.Path)

			var protocolLogger *logrus.Entry

			if *flagDistributeLogs {
				logger, err := createLogger("inspector-" + id)
				if err != nil {
					panic(fmt.Sprintf("could not create logger: %s", err))
				}

				protocolLogger = logger.WithFields(logrus.Fields{
					fieldLevel:       levelConnection,
					fieldInspectorID: id,
				})

			} else {
				protocolLogger = logger.WithFields(logrus.Fields{
					fieldInspectorID: id,
				})
			}

			go dumpStream(protocolLogger, stream)

			endpoint := "ws://" + *flagRemote + "/devtools/" + basePath + "/" + id

			logger.Infof("---------- connection from %s ----------", req.RemoteAddr)
			logger.Infof("checking protocol versions on: %s", endpoint)

			ver, err := checkVersion()
			if err != nil {
				protocolLogger.Errorf("could not check version: %v", err)
				http.Error(res, "could not check version", 500)
				return
			}

			logger.Infof("protocol version: %s", ver["Protocol-Version"])
			logger.Infof("versions: Chrome(%s), V8(%s), Webkit(%s)", ver["Browser"], ver["V8-Version"], ver["WebKit-Version"])
			logger.Infof("browser user agent: %s", ver["User-Agent"])
			logger.Infof("connecting to %s... ", endpoint)

			// connecting to ws
			out, pres, err := wsDialer.Dial(endpoint, nil)
			if err != nil {
				msg := fmt.Sprintf("could not connect to %s: %v", endpoint, err)
				logger.Error(protocolError(msg))
				http.Error(res, msg, 500)
				return
			}
			defer pres.Body.Close()
			defer out.Close()

			// connect incoming websocket
			logger.Infof("upgrading connection on %s...", req.RemoteAddr)
			in, err := wsUpgrader.Upgrade(res, req, nil)
			if err != nil {
				logger.Errorf("could not upgrade websocket from %s: %v", req.RemoteAddr, err)
				http.Error(res, "could not upgrade websocket connection", 500)
				return
			}
			defer in.Close()

			ctxt, cancel := context.WithCancel(context.Background())
			defer cancel()

			errc := make(chan error, 1)
			go proxyWS(ctxt, stream, in, out, errc)
			go proxyWS(ctxt, stream, out, in, errc)

			<-errc
			logger.Infof("---------- closing %s ----------", req.RemoteAddr)

			if *flagDistributeLogs {
				if closer, ok := protocolLogger.Logger.Out.(io.Closer); ok {
					closer.Close()
				}
			}

			if *flagOnce {
				os.Exit(0)
			}
		}
	}

	mux.HandleFunc("/devtools/page/", handlerFunc("page"))
	mux.HandleFunc("/devtools/browser/", handlerFunc("browser"))

	log.Fatal(http.ListenAndServe(*flagListen, mux))
}

func dumpStream(logger *logrus.Entry, stream chan *protocolMessage) {
	logger.Printf("Legend: %s, %s, %s, %s, %s, %s", protocolColor("protocol informations"),
		eventsColor("received events"),
		requestColor("sent request frames"),
		requestReplyColor("requests params"),
		responseColor("received responses"),
		errorColor("error response."),
	)

	requests := make(map[uint64]*protocolMessage)
	targetRequests := make(map[uint64]*protocolMessage)

	for {
		select {
		case msg := <-stream:
			if msg.InTarget() {

				var targetLogger *logrus.Entry

				if *flagDistributeLogs {
					logger, err := createLogger(fmt.Sprintf("target-%s", msg.TargetID()))
					if err != nil {
						panic(fmt.Sprintf("could not create logger: %v", err))
					}

					targetLogger = logger.WithFields(logrus.Fields{
						fieldLevel:    levelTarget,
						fieldTargetID: msg.TargetID(),
					})

				} else {
					targetLogger = logger.WithFields(logrus.Fields{
						fieldLevel:    levelTarget,
						fieldTargetID: msg.TargetID(),
					})
				}

				if msg.IsRequest() {
					requests[msg.ID] = nil

					if protocolMessage, err := decodeMessage([]byte(asString(msg.Params["message"]))); err == nil {
						targetRequests[protocolMessage.ID] = protocolMessage

						if *flagShowRequests {
							targetLogger.WithFields(logrus.Fields{
								fieldType:   typeRequest,
								fieldMethod: protocolMessage.Method,
							}).Info(serialize(protocolMessage.Params))
						}

					} else {
						logger.WithFields(logrus.Fields{
							fieldLevel: levelConnection,
						}).Errorf("Could not deserialize message: %+v", err)
					}
				}

				if msg.IsEvent() {
					if protocolMessage, err := decodeMessage([]byte(asString(msg.Params["message"]))); err == nil {
						if protocolMessage.IsEvent() {
							targetLogger.WithFields(logrus.Fields{
								fieldType:   typeEvent,
								fieldMethod: protocolMessage.Method,
							}).Info(serialize(protocolMessage.Params))
						}

						if protocolMessage.IsResponse() {
							var logMessage string
							var logType int
							var logRequest string
							var logMethod string

							if protocolMessage.IsError() {
								logMessage = serialize(protocolMessage.Error)
								logType = typeRequestResponseError
							} else {
								logMessage = serialize(protocolMessage.Result)
								logType = typeRequestResponse
							}

							if request, ok := targetRequests[protocolMessage.ID]; ok && request != nil {
								delete(targetRequests, protocolMessage.ID)
								logRequest = serialize(request.Params)
								logMethod = request.Method

							} else {
								logRequest = errorColor("could not find request with id: %d", protocolMessage.ID)
							}

							targetLogger.WithFields(logrus.Fields{
								fieldType:    logType,
								fieldMethod:  logMethod,
								fieldRequest: logRequest,
							}).Info(logMessage)
						}
					} else {
						logger.WithFields(logrus.Fields{
							fieldLevel: levelConnection,
						}).Errorf("Could not deserialize message: %+v", err)
					}
				}

			} else {
				protocolLogger := logger.WithFields(logrus.Fields{
					fieldLevel:    levelProtocol,
					fieldTargetID: protocolTargetID,
				})

				if msg.IsRequest() {
					requests[msg.ID] = msg

					if *flagShowRequests {
						protocolLogger.WithFields(logrus.Fields{
							fieldType:   typeRequest,
							fieldMethod: msg.Method,
						}).Info(serialize(msg.Params))
					}
				}

				if msg.IsResponse() {

					var logMessage string
					var logType int
					var logRequest string
					var logMethod string

					if msg.IsError() {
						logMessage = serialize(msg.Error)
						logType = typeRequestResponseError
					} else {
						logMessage = serialize(msg.Result)
						logType = typeRequestResponse
					}

					if request, ok := requests[msg.ID]; ok && request != nil {
						logRequest = serialize(request.Params)
						logMethod = request.Method

						delete(requests, msg.ID)

						protocolLogger.WithFields(logrus.Fields{
							fieldType:    logType,
							fieldMethod:  logMethod,
							fieldRequest: logRequest,
						}).Info(logMessage)
					}
				}

				if msg.IsEvent() {
					protocolLogger.WithFields(logrus.Fields{
						fieldType:   typeEvent,
						fieldMethod: msg.Method,
					}).Info(serialize(msg.Params))
				}
			}
		}
	}
}

func checkVersion() (map[string]string, error) {
	cl := &http.Client{}
	req, err := http.NewRequest("GET", "http://"+*flagRemote+"/json/version", nil)
	if err != nil {
		return nil, err
	}

	res, err := cl.Do(req)
	if err != nil {
		return nil, err
	}

	defer res.Body.Close()

	var v map[string]string
	if err := json.NewDecoder(res.Body).Decode(&v); err != nil {
		return nil, errors.New("expected json result")
	}

	return v, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"

	"github.com/mitchellh/go-homedir"
)

func shortUsage() {
	os.Stderr.WriteString(`
usage: vub [option(s)] <repository-uri>
try 'vub --help' for more information.
`[1:])
}

func usage() {
	os.Stderr.WriteString(`
usage: vub [option(s)] <repository-uri>
install Vim plugin to under the management of vim-unbundle.

options:
  -h, --help                show this help message
  -f, --filetype=TYPE       installing under the ftbundle/TYPE
`[1:])
}

var (
	ShortGitHubURI = regexp.MustCompile(`^[\w\-.]+/[\w\-.]+$`)
)

func ToSourceURI(uri string) (string, error) {
	switch {
	case ShortGitHubURI.MatchString(uri):
		return "https://github.com/" + uri, nil
	default:
		return uri, nil
	}
}

var (
	home, errInit = homedir.Dir()
	dotvim        = filepath.Join(home, ".vim")
)

func ToDestinationPath(uri, filetype string) (string, error) {
	name := filepath.Base(uri)
	if filetype == "" {
		return filepath.Join(dotvim, "bundle", name), nil
	}
	return filepath.Join(dotvim, "ftbundle", filetype, name), nil
}

type Package struct {
	verbose bool
	src     string
	dst     string
}

func NewPackage(uri, filetype string) (*Package, error) {
	src, err := ToSourceURI(uri)
	if err != nil {
		return nil, err
	}
	dst, err := ToDestinationPath(uri, filetype)
	if err != nil {
		return nil, err
	}
	return &Package{
		src: src,
		dst: dst,
	}, nil
}

func (p *Package) toCommand() *exec.Cmd {
	return exec.Command("git", "clone", p.src, p.dst)
}

func (p *Package) Verbose(enable bool) {
	p.verbose = enable
}

func main() {
	var filetype string
	flag.StringVar(&filetype, "f", "", "")
	flag.StringVar(&filetype, "filetype", "", "")

	var isHelp bool
	flag.BoolVar(&isHelp, "h", false, "")
	flag.BoolVar(&isHelp, "help", false, "")
	flag.Usage = usage
	flag.Parse()
	switch {
	case isHelp:
		usage()
		os.Exit(0)
	case flag.NArg() < 1:
		shortUsage()
		os.Exit(2)
	}
	uri := flag.Arg(0)

	p, err := NewPackage(uri, filetype)
	if err != nil {
		fmt.Fprintln(os.Stderr, "vub:", err)
		os.Exit(2)
	}

	c := p.toCommand()
	if err := c.Run(); err != nil {
		fmt.Fprintln(os.Stderr, "vub:", err)
		os.Exit(1)
	}
}
<commit_msg>Implement Install<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/mitchellh/go-homedir"
)

func shortUsage() {
	os.Stderr.WriteString(`
usage: vub [option(s)] <repository-uri>
try 'vub --help' for more information.
`[1:])
}

func usage() {
	os.Stderr.WriteString(`
usage: vub [option(s)] <repository-uri>
install Vim plugin to under the management of vim-unbundle.

options:
  -h, --help                show this help message
  -f, --filetype=TYPE       installing under the ftbundle/TYPE
`[1:])
}

var (
	ShortGitHubURI = regexp.MustCompile(`^[\w\-.]+/[\w\-.]+$`)
)

func ToSourceURI(uri string) (string, error) {
	switch {
	case ShortGitHubURI.MatchString(uri):
		return "https://github.com/" + uri, nil
	default:
		return uri, nil
	}
}

var (
	home, errInit = homedir.Dir()
	dotvim        = filepath.Join(home, ".vim")
)

func ToDestinationPath(uri, filetype string) (string, error) {
	name := filepath.Base(uri)
	if filetype == "" {
		return filepath.Join(dotvim, "bundle", name), nil
	}
	return filepath.Join(dotvim, "ftbundle", filetype, name), nil
}

type Package struct {
	verbose bool
	src     string
	dst     string
}

func NewPackage(uri, filetype string) (*Package, error) {
	src, err := ToSourceURI(uri)
	if err != nil {
		return nil, err
	}
	dst, err := ToDestinationPath(uri, filetype)
	if err != nil {
		return nil, err
	}
	return &Package{
		src: src,
		dst: dst,
	}, nil
}

func (p *Package) toCommand() *exec.Cmd {
	return exec.Command("git", "clone", p.src, p.dst)
}

func (p *Package) Install(w io.Writer) error {
	c := p.toCommand()
	if w != nil && p.verbose {
		c.Stdout = w
		c.Stderr = w
		_, err := io.WriteString(w, strings.Join(c.Args, " "))
		if err != nil {
			return err
		}
	}
	return c.Run()
}

func (p *Package) Verbose(enable bool) {
	p.verbose = enable
}

func main() {
	var filetype string
	flag.StringVar(&filetype, "f", "", "")
	flag.StringVar(&filetype, "filetype", "", "")

	var isHelp bool
	flag.BoolVar(&isHelp, "h", false, "")
	flag.BoolVar(&isHelp, "help", false, "")
	flag.Usage = usage
	flag.Parse()
	switch {
	case isHelp:
		usage()
		os.Exit(0)
	case flag.NArg() < 1:
		shortUsage()
		os.Exit(2)
	}
	uri := flag.Arg(0)

	p, err := NewPackage(uri, filetype)
	if err != nil {
		fmt.Fprintln(os.Stderr, "vub:", err)
		os.Exit(2)
	}

	c := p.toCommand()
	if err := c.Run(); err != nil {
		fmt.Fprintln(os.Stderr, "vub:", err)
		os.Exit(1)
	}
}
<|endoftext|>


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

package main

import (
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"time"

	"gopkg.in/alecthomas/kingpin.v2"

	"github.com/hanwen/go-fuse/fuse"
	"github.com/hanwen/go-fuse/fuse/nodefs"
	"github.com/rcrowley/go-metrics"
	"github.com/square/go-sq-metrics"
	klog "github.com/square/keywhiz-fs/log"
	"golang.org/x/sys/unix"
)

var (
	app = kingpin.New("keywhiz-fs", "A FUSE based file-system client for Keywhiz.")

	certFile      = app.Flag("cert", "PEM-encoded certificate file").PlaceHolder("FILE").Default("").String()
	keyFile       = app.Flag("key", "PEM-encoded private key file").PlaceHolder("FILE").Required().String()
	caFile        = app.Flag("ca", "PEM-encoded CA certificates file").PlaceHolder("FILE").Required().String()
	asuser        = app.Flag("asuser", "Default user to own files").Default("keywhiz").String()
	asgroup       = app.Flag("group", "Default group to own files").Default("keywhiz").String()
	debug         = app.Flag("debug", "Enable debugging output").Default("false").Bool()
	timeout       = app.Flag("timeout", "Timeout for communication with server").Default("20s").Duration()
	cacheTimeout  = app.Flag("cache-timeout", "Timeout for cache eviction. Useful for testing.").Default("1h").Duration()
	metricsURL    = app.Flag("metrics-url", "Collect metrics and POST them periodically to the given URL (via HTTP/JSON).").PlaceHolder("URL").String()
	metricsPrefix = app.Flag("metrics-prefix", "Override the default metrics prefix used for reporting metrics.").PlaceHolder("PREFIX").String()
	syslog        = app.Flag("syslog", "Send logs to syslog instead of stderr.").Default("false").Bool()
	disableMlock  = app.Flag("disable-mlock", "Do not call mlockall on process memory.").Default("false").Bool()
	serverURL     = app.Arg("url", "server url").Required().URL()
	mountpoint    = app.Arg("mountpoint", "mountpoint").Required().String()
	logger        *klog.Logger
)

func main() {
	app.Version(fmt.Sprintf("rev %s-%s on \"%s\"", buildRevision, buildTime, buildMachine))
	kingpin.MustParse(app.Parse(os.Args[1:]))

	logConfig := klog.Config{Debug: *debug, Mountpoint: *mountpoint, Syslog: *syslog}
	logger = klog.New("kwfs_main", logConfig)
	defer logger.Close()

	if *certFile == "" {
		logger.Debugf("Certificate file not specified, assuming certificate also in %s", *keyFile)
		certFile = keyFile
	}

	metricsHandle := setupMetrics(metricsURL, metricsPrefix, *mountpoint)

	if !*disableMlock {
		lockMemory()
	}

	// TODO: move time limit settings to config file?
	// TODO: or at least make it consistent? some are set here, some are set above with app.Flag()
	freshThreshold := *cacheTimeout
	backendDeadline := 5 * time.Second
	maxWait := *timeout + backendDeadline
	delayDeletion := 1 * time.Hour
	timeouts := Timeouts{freshThreshold, backendDeadline, maxWait, delayDeletion}

	client := NewClient(*certFile, *keyFile, *caFile, *serverURL, *timeout, logConfig)

	ownership := NewOwnership(*asuser, *asgroup)
	kwfs, root, err := NewKeywhizFs(&client, ownership, timeouts, metricsHandle, logConfig)
	if err != nil {
		log.Fatalf("KeywhizFs init fail: %v\n", err)
	}

	mountOptions := &fuse.MountOptions{
		AllowOther: true,
		Name:       kwfs.String(),
		Options:    []string{"default_permissions"},
	}

	// Empty Options struct avoids setting a global uid/gid override.
	conn := nodefs.NewFileSystemConnector(root, &nodefs.Options{})
	server, err := fuse.NewServer(conn.RawFS(), *mountpoint, mountOptions)
	if err != nil {
		log.Fatalf("Mount fail: %v\n", err)
	}

	// Catch SIGINT and exit cleanly.
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		for {
			sig := <-c
			logger.Warnf("Got signal %s, unmounting", sig)
			err := server.Unmount()
			if err != nil {
				logger.Warnf("Error while unmounting: %v", err)
			}
		}
	}()

	kwfs.Cache.Warmup()
	server.Serve()
	logger.Infof("Exiting")
}

// Setup metrics
func setupMetrics(metricsURL *string, metricsPrefix *string, mountpoint string) *sqmetrics.SquareMetrics {
	if *metricsURL != "" {
		if !strings.HasPrefix(*metricsURL, "http://") && !strings.HasPrefix(*metricsURL, "https://") {
			log.Fatalf("--metrics-url should start with http:// or https://")
			os.Exit(1)
		}
		log.Printf("metrics enabled; reporting metrics via POST to %s", *metricsURL)
	}

	var prefix string
	if *metricsPrefix != "" {
		prefix = *metricsPrefix
	} else {
		// By default, prefix metrics with escaped mount path. Replace slashes with - for easier aggregation
		prefix = fmt.Sprintf("keywhizfs.%s", strings.Replace(strings.Replace(mountpoint, "-", "--", -1), "/", "-", -1))
	}

	return sqmetrics.NewMetrics(*metricsURL, prefix, (30 * time.Second), metrics.DefaultRegistry)
}

// Locks memory, preventing memory from being written to disk as swap
func lockMemory() {
	err := unix.Mlockall(unix.MCL_FUTURE | unix.MCL_CURRENT)
	switch err {
	case nil:
	case unix.ENOSYS:
		logger.Warnf("mlockall() not implemented on this system")
	case unix.ENOMEM:
		logger.Warnf("mlockall() failed with ENOMEM")
	default:
		log.Fatalf("Could not perform mlockall and prevent swapping memory: %v", err)
	}
}

// Helper function to panic on error
func panicOnError(err error) {
	if err != nil {
		logger.Errorf("panic: %v", err)
		panic(err)
	}
}
<commit_msg>Moved initial Cache.Warmup() to earlier in the startup process, before the filesystem is mounted.<commit_after>// Copyright 2015 Square Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"time"

	"gopkg.in/alecthomas/kingpin.v2"

	"github.com/hanwen/go-fuse/fuse"
	"github.com/hanwen/go-fuse/fuse/nodefs"
	"github.com/rcrowley/go-metrics"
	"github.com/square/go-sq-metrics"
	klog "github.com/square/keywhiz-fs/log"
	"golang.org/x/sys/unix"
)

var (
	app = kingpin.New("keywhiz-fs", "A FUSE based file-system client for Keywhiz.")

	certFile      = app.Flag("cert", "PEM-encoded certificate file").PlaceHolder("FILE").Default("").String()
	keyFile       = app.Flag("key", "PEM-encoded private key file").PlaceHolder("FILE").Required().String()
	caFile        = app.Flag("ca", "PEM-encoded CA certificates file").PlaceHolder("FILE").Required().String()
	asuser        = app.Flag("asuser", "Default user to own files").Default("keywhiz").String()
	asgroup       = app.Flag("group", "Default group to own files").Default("keywhiz").String()
	debug         = app.Flag("debug", "Enable debugging output").Default("false").Bool()
	timeout       = app.Flag("timeout", "Timeout for communication with server").Default("20s").Duration()
	cacheTimeout  = app.Flag("cache-timeout", "Timeout for cache eviction. Useful for testing.").Default("1h").Duration()
	metricsURL    = app.Flag("metrics-url", "Collect metrics and POST them periodically to the given URL (via HTTP/JSON).").PlaceHolder("URL").String()
	metricsPrefix = app.Flag("metrics-prefix", "Override the default metrics prefix used for reporting metrics.").PlaceHolder("PREFIX").String()
	syslog        = app.Flag("syslog", "Send logs to syslog instead of stderr.").Default("false").Bool()
	disableMlock  = app.Flag("disable-mlock", "Do not call mlockall on process memory.").Default("false").Bool()
	serverURL     = app.Arg("url", "server url").Required().URL()
	mountpoint    = app.Arg("mountpoint", "mountpoint").Required().String()
	logger        *klog.Logger
)

func main() {
	app.Version(fmt.Sprintf("rev %s-%s on \"%s\"", buildRevision, buildTime, buildMachine))
	kingpin.MustParse(app.Parse(os.Args[1:]))

	logConfig := klog.Config{Debug: *debug, Mountpoint: *mountpoint, Syslog: *syslog}
	logger = klog.New("kwfs_main", logConfig)
	defer logger.Close()

	if *certFile == "" {
		logger.Debugf("Certificate file not specified, assuming certificate also in %s", *keyFile)
		certFile = keyFile
	}

	metricsHandle := setupMetrics(metricsURL, metricsPrefix, *mountpoint)

	if !*disableMlock {
		lockMemory()
	}

	// TODO: move time limit settings to config file?
	// TODO: or at least make it consistent? some are set here, some are set above with app.Flag()
	freshThreshold := *cacheTimeout
	backendDeadline := 5 * time.Second
	maxWait := *timeout + backendDeadline
	delayDeletion := 1 * time.Hour
	timeouts := Timeouts{freshThreshold, backendDeadline, maxWait, delayDeletion}

	client := NewClient(*certFile, *keyFile, *caFile, *serverURL, *timeout, logConfig)

	ownership := NewOwnership(*asuser, *asgroup)
	kwfs, root, err := NewKeywhizFs(&client, ownership, timeouts, metricsHandle, logConfig)
	if err != nil {
		log.Fatalf("KeywhizFs init fail: %v\n", err)
	}
	kwfs.Cache.Warmup()

	mountOptions := &fuse.MountOptions{
		AllowOther: true,
		Name:       kwfs.String(),
		Options:    []string{"default_permissions"},
	}

	// Empty Options struct avoids setting a global uid/gid override.
	conn := nodefs.NewFileSystemConnector(root, &nodefs.Options{})
	server, err := fuse.NewServer(conn.RawFS(), *mountpoint, mountOptions)
	if err != nil {
		log.Fatalf("Mount fail: %v\n", err)
	}

	// Catch SIGINT and exit cleanly.
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		for {
			sig := <-c
			logger.Warnf("Got signal %s, unmounting", sig)
			err := server.Unmount()
			if err != nil {
				logger.Warnf("Error while unmounting: %v", err)
			}
		}
	}()

	server.Serve()
	logger.Infof("Exiting")
}

// Setup metrics
func setupMetrics(metricsURL *string, metricsPrefix *string, mountpoint string) *sqmetrics.SquareMetrics {
	if *metricsURL != "" {
		if !strings.HasPrefix(*metricsURL, "http://") && !strings.HasPrefix(*metricsURL, "https://") {
			log.Fatalf("--metrics-url should start with http:// or https://")
			os.Exit(1)
		}
		log.Printf("metrics enabled; reporting metrics via POST to %s", *metricsURL)
	}

	var prefix string
	if *metricsPrefix != "" {
		prefix = *metricsPrefix
	} else {
		// By default, prefix metrics with escaped mount path. Replace slashes with - for easier aggregation
		prefix = fmt.Sprintf("keywhizfs.%s", strings.Replace(strings.Replace(mountpoint, "-", "--", -1), "/", "-", -1))
	}

	return sqmetrics.NewMetrics(*metricsURL, prefix, (30 * time.Second), metrics.DefaultRegistry)
}

// Locks memory, preventing memory from being written to disk as swap
func lockMemory() {
	err := unix.Mlockall(unix.MCL_FUTURE | unix.MCL_CURRENT)
	switch err {
	case nil:
	case unix.ENOSYS:
		logger.Warnf("mlockall() not implemented on this system")
	case unix.ENOMEM:
		logger.Warnf("mlockall() failed with ENOMEM")
	default:
		log.Fatalf("Could not perform mlockall and prevent swapping memory: %v", err)
	}
}

// Helper function to panic on error
func panicOnError(err error) {
	if err != nil {
		logger.Errorf("panic: %v", err)
		panic(err)
	}
}
<|endoftext|>