<commit_before>package alert

import (
	"context"
	"fmt"
	"strings"
	"time"

	g "github.com/onsi/ginkgo"
	o "github.com/onsi/gomega"
	"github.com/openshift/origin/pkg/synthetictests/allowedalerts"
	exutil "github.com/openshift/origin/test/extended/util"
	"github.com/openshift/origin/test/extended/util/disruption"
	helper "github.com/openshift/origin/test/extended/util/prometheus"
	prometheusv1 "github.com/prometheus/client_golang/api/prometheus/v1"
	"github.com/prometheus/common/model"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/kubernetes/test/e2e/framework"
	"k8s.io/kubernetes/test/e2e/upgrades"
)

// UpgradeTest runs verifies invariants regarding what alerts are allowed to fire
// during the upgrade process.
type UpgradeTest struct {
	oc               *exutil.CLI
	prometheusClient prometheusv1.API
}

func (UpgradeTest) Name() string { return "check-for-alerts" }
func (UpgradeTest) DisplayName() string {
	return "[sig-arch] Check if alerts are firing during or after upgrade success"
}

// Setup creates parameters to query Prometheus
func (t *UpgradeTest) Setup(f *framework.Framework) {
	g.By("Setting up upgrade alert test")

	oc := exutil.NewCLIWithFramework(f)
	t.oc = oc
	t.prometheusClient = oc.NewPrometheusClient(context.TODO())
	framework.Logf("Post-upgrade alert test setup complete")
}

// Test checks if alerts are firing at various points during upgrade.
// An alert firing during an upgrade is a high severity bug - it either points to a real issue in
// a dependency, or a failure of the component, and therefore must be fixed.
func (t *UpgradeTest) Test(f *framework.Framework, done <-chan struct{}, upgrade upgrades.UpgradeType) {
	tolerateDuringSkew := exutil.TolerateVersionSkewInTests()
	firingAlertsWithBugs := helper.MetricConditions{
		{
			Selector: map[string]string{"alertname": "KubePodNotReady", "namespace": "openshift-kube-apiserver-operator"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "KubePodNotReady", "namespace": "openshift-kube-apiserver-operator"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "ClusterOperatorDegraded", "name": "openshift-apiserver"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "ClusterOperatorDown", "name": "authentication"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "ClusterOperatorDown", "name": "machine-config"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1955300",
		},
		{
			Selector: map[string]string{"alertname": "ClusterOperatorDegraded", "name": "authentication"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "KubeDaemonSetRolloutStuck"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1943667",
		},
		{
			Selector: map[string]string{"alertname": "KubeAPIErrorBudgetBurn"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1953798",
			Matches: func(_ *model.Sample) bool {
				return framework.ProviderIs("gce")
			},
		},
		{
			Selector: map[string]string{"alertname": "AggregatedAPIDown", "namespace": "default", "name": "v1beta1.metrics.k8s.io"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1970624",
			Matches: func(_ *model.Sample) bool {
				return framework.ProviderIs("gce")
			},
		},
		{
			// Should be removed one release after the attached bugzilla is fixed.
			Selector: map[string]string{"alertname": "HighlyAvailableWorkloadIncorrectlySpread", "namespace": "openshift-monitoring", "workload": "prometheus-k8s"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1949262",
		},
		{
			// Should be removed one release after the attached bugzilla is fixed.
			Selector: map[string]string{"alertname": "HighlyAvailableWorkloadIncorrectlySpread", "namespace": "openshift-monitoring", "workload": "alertmanager-main"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1955489",
		},
		{
			// Should be removed one release after the attached bugzilla is fixed, or after that bug is fixed in a backport to the previous minor.
			Selector: map[string]string{"alertname": "ExtremelyHighIndividualControlPlaneCPU"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1985073",
			Matches: func(_ *model.Sample) bool {
				return framework.ProviderIs("gce")
			},
		},
		{
			Selector: map[string]string{"alertname": "KubeJobFailed", "namespace": "openshift-multus"}, // not sure how to do a job_name prefix
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=2054426",
		},
	}
	allowedFiringAlerts := helper.MetricConditions{
		{
			Selector: map[string]string{"alertname": "TargetDown", "namespace": "openshift-e2e-loki"},
			Text:     "Loki is nice to have, but we can allow it to be down",
		},
		{
			Selector: map[string]string{"alertname": "KubePodNotReady", "namespace": "openshift-e2e-loki"},
			Text:     "Loki is nice to have, but we can allow it to be down",
		},
		{
			Selector: map[string]string{"alertname": "KubeDeploymentReplicasMismatch", "namespace": "openshift-e2e-loki"},
			Text:     "Loki is nice to have, but we can allow it to be down",
		},
	}

	pendingAlertsWithBugs := helper.MetricConditions{
		{
			Selector: map[string]string{"alertname": "ClusterMonitoringOperatorReconciliationErrors"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1932624",
		},
		{
			Selector: map[string]string{"alertname": "KubeClientErrors"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1925698",
		},
		{
			Selector: map[string]string{"alertname": "NetworkPodsCrashLooping"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=2009078",
		},
	}
	allowedPendingAlerts := helper.MetricConditions{
		{
			Selector: map[string]string{"alertname": "etcdMemberCommunicationSlow"},
			Text:     "Excluded because it triggers during upgrade (detects ~5m of high latency immediately preceeding the end of the test), and we don't want to change the alert because it is correct",
		},
	}

	// we exclude alerts that have their own separate tests.
	for _, alertTest := range allowedalerts.AllAlertTests() {
		switch alertTest.AlertState() {
		case allowedalerts.AlertPending:
			// a pending test covers pending and everything above (firing)
			allowedPendingAlerts = append(allowedPendingAlerts,
				helper.MetricCondition{
					Selector: map[string]string{"alertname": alertTest.AlertName()},
					Text:     "has a separate e2e test",
				},
			)
			allowedFiringAlerts = append(allowedFiringAlerts,
				helper.MetricCondition{
					Selector: map[string]string{"alertname": alertTest.AlertName()},
					Text:     "has a separate e2e test",
				},
			)
		case allowedalerts.AlertInfo:
			// an info test covers all firing
			allowedFiringAlerts = append(allowedFiringAlerts,
				helper.MetricCondition{
					Selector: map[string]string{"alertname": alertTest.AlertName()},
					Text:     "has a separate e2e test",
				},
			)
		}
	}

	knownViolations := sets.NewString()
	unexpectedViolations := sets.NewString()
	unexpectedViolationsAsFlakes := sets.NewString()
	debug := sets.NewString()

	g.By("Checking for alerts")

	start := time.Now()

	// Block until upgrade is done
	g.By("Waiting for upgrade to finish before checking for alerts")
	<-done

	// Additonal delay after upgrade completion to allow pending alerts to settle
	g.By("Waiting before checking for alerts")
	time.Sleep(1 * time.Minute)

	testDuration := time.Now().Sub(start).Round(time.Second)

	// Invariant: No non-info level alerts should have fired during the upgrade
	firingAlertQuery := fmt.Sprintf(`
sort_desc(
count_over_time(ALERTS{alertstate="firing",severity!="info",alertname!~"Watchdog|AlertmanagerReceiversNotConfigured"}[%[1]s:1s])
) > 0
`, testDuration)
	result, err := helper.RunQuery(context.TODO(), t.prometheusClient, firingAlertQuery)
	o.Expect(err).NotTo(o.HaveOccurred(), "unable to check firing alerts during upgrade")
	for _, series := range result.Data.Result {
		labels := helper.StripLabels(series.Metric, "alertname", "alertstate", "prometheus")
		violation := fmt.Sprintf("alert %s fired for %s seconds with labels: %s", series.Metric["alertname"], series.Value, helper.LabelsAsSelector(labels))
		if cause := allowedFiringAlerts.Matches(series); cause != nil {
			debug.Insert(fmt.Sprintf("%s (allowed: %s)", violation, cause.Text))
			continue
		}
		if cause := firingAlertsWithBugs.Matches(series); cause != nil {
			knownViolations.Insert(fmt.Sprintf("%s (open bug: %s)", violation, cause.Text))
		} else {
			unexpectedViolations.Insert(violation)
		}
	}

	// Invariant: There should be no pending alerts 1m after the upgrade completes
	pendingAlertQuery := fmt.Sprintf(`
sort_desc(
  time() * ALERTS + 1
  -
  last_over_time((
    time() * ALERTS{alertname!~"Watchdog|AlertmanagerReceiversNotConfigured",alertstate="pending",severity!="info"}
    unless
    ALERTS offset 1s
  )[%[1]s:1s])
)
`, testDuration)
	result, err = helper.RunQuery(context.TODO(), t.prometheusClient, pendingAlertQuery)
	o.Expect(err).NotTo(o.HaveOccurred(), "unable to retrieve pending alerts after upgrade")
	for _, series := range result.Data.Result {
		labels := helper.StripLabels(series.Metric, "alertname", "alertstate", "prometheus")
		violation := fmt.Sprintf("alert %s pending for %s seconds with labels: %s", series.Metric["alertname"], series.Value, helper.LabelsAsSelector(labels))
		if cause := allowedPendingAlerts.Matches(series); cause != nil {
			debug.Insert(fmt.Sprintf("%s (allowed: %s)", violation, cause.Text))
			continue
		}
		if cause := pendingAlertsWithBugs.Matches(series); cause != nil {
			knownViolations.Insert(fmt.Sprintf("%s (open bug: %s)", violation, cause.Text))
		} else {
			// treat pending errors as a flake right now because we are still trying to determine the scope
			// TODO: move this to unexpectedViolations later
			unexpectedViolationsAsFlakes.Insert(violation)
		}
	}

	if len(debug) > 0 {
		framework.Logf("Alerts were detected during upgrade which are allowed:\n\n%s", strings.Join(debug.List(), "\n"))
	}
	if len(unexpectedViolations) > 0 {
		if !tolerateDuringSkew {
			framework.Failf("Unexpected alerts fired or pending during the upgrade:\n\n%s", strings.Join(unexpectedViolations.List(), "\n"))
		}
	}
	if flakes := sets.NewString().Union(knownViolations).Union(unexpectedViolations).Union(unexpectedViolationsAsFlakes); len(flakes) > 0 {
		disruption.FrameworkFlakef(f, "Unexpected alert behavior during upgrade:\n\n%s", strings.Join(flakes.List(), "\n"))
	}
	framework.Logf("No alerts fired during upgrade")
}

// Teardown cleans up any remaining resources.
func (t *UpgradeTest) Teardown(f *framework.Framework) {
	// rely on the namespace deletion to clean up everything
}
<commit_msg>Improve ignore for ip-reconciler KubeJobFailed alert.<commit_after>package alert

import (
	"context"
	"fmt"
	"strings"
	"time"

	g "github.com/onsi/ginkgo"
	o "github.com/onsi/gomega"
	"github.com/openshift/origin/pkg/synthetictests/allowedalerts"
	exutil "github.com/openshift/origin/test/extended/util"
	"github.com/openshift/origin/test/extended/util/disruption"
	helper "github.com/openshift/origin/test/extended/util/prometheus"
	prometheusv1 "github.com/prometheus/client_golang/api/prometheus/v1"
	"github.com/prometheus/common/model"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/kubernetes/test/e2e/framework"
	"k8s.io/kubernetes/test/e2e/upgrades"
)

// UpgradeTest runs verifies invariants regarding what alerts are allowed to fire
// during the upgrade process.
type UpgradeTest struct {
	oc               *exutil.CLI
	prometheusClient prometheusv1.API
}

func (UpgradeTest) Name() string { return "check-for-alerts" }
func (UpgradeTest) DisplayName() string {
	return "[sig-arch] Check if alerts are firing during or after upgrade success"
}

// Setup creates parameters to query Prometheus
func (t *UpgradeTest) Setup(f *framework.Framework) {
	g.By("Setting up upgrade alert test")

	oc := exutil.NewCLIWithFramework(f)
	t.oc = oc
	t.prometheusClient = oc.NewPrometheusClient(context.TODO())
	framework.Logf("Post-upgrade alert test setup complete")
}

// Test checks if alerts are firing at various points during upgrade.
// An alert firing during an upgrade is a high severity bug - it either points to a real issue in
// a dependency, or a failure of the component, and therefore must be fixed.
func (t *UpgradeTest) Test(f *framework.Framework, done <-chan struct{}, upgrade upgrades.UpgradeType) {
	tolerateDuringSkew := exutil.TolerateVersionSkewInTests()
	firingAlertsWithBugs := helper.MetricConditions{
		{
			Selector: map[string]string{"alertname": "KubePodNotReady", "namespace": "openshift-kube-apiserver-operator"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "KubePodNotReady", "namespace": "openshift-kube-apiserver-operator"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "ClusterOperatorDegraded", "name": "openshift-apiserver"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "ClusterOperatorDown", "name": "authentication"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "ClusterOperatorDown", "name": "machine-config"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1955300",
		},
		{
			Selector: map[string]string{"alertname": "ClusterOperatorDegraded", "name": "authentication"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1939580",
		},
		{
			Selector: map[string]string{"alertname": "KubeDaemonSetRolloutStuck"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1943667",
		},
		{
			Selector: map[string]string{"alertname": "KubeAPIErrorBudgetBurn"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1953798",
			Matches: func(_ *model.Sample) bool {
				return framework.ProviderIs("gce")
			},
		},
		{
			Selector: map[string]string{"alertname": "AggregatedAPIDown", "namespace": "default", "name": "v1beta1.metrics.k8s.io"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1970624",
			Matches: func(_ *model.Sample) bool {
				return framework.ProviderIs("gce")
			},
		},
		{
			// Should be removed one release after the attached bugzilla is fixed.
			Selector: map[string]string{"alertname": "HighlyAvailableWorkloadIncorrectlySpread", "namespace": "openshift-monitoring", "workload": "prometheus-k8s"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1949262",
		},
		{
			// Should be removed one release after the attached bugzilla is fixed.
			Selector: map[string]string{"alertname": "HighlyAvailableWorkloadIncorrectlySpread", "namespace": "openshift-monitoring", "workload": "alertmanager-main"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1955489",
		},
		{
			// Should be removed one release after the attached bugzilla is fixed, or after that bug is fixed in a backport to the previous minor.
			Selector: map[string]string{"alertname": "ExtremelyHighIndividualControlPlaneCPU"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1985073",
			Matches: func(_ *model.Sample) bool {
				return framework.ProviderIs("gce")
			},
		},
		{
			Selector: map[string]string{"alertname": "KubeJobFailed", "namespace": "openshift-multus"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=2054426",
			Matches: func(sample *model.Sample) bool {
				// Only match if the job_name label starts with ip-reconciler:
				if strings.HasPrefix(string(sample.Metric[model.LabelName("job_name")]), "ip-reconciler-") {
					return true
				}
				return false
			},
		},
	}
	allowedFiringAlerts := helper.MetricConditions{
		{
			Selector: map[string]string{"alertname": "TargetDown", "namespace": "openshift-e2e-loki"},
			Text:     "Loki is nice to have, but we can allow it to be down",
		},
		{
			Selector: map[string]string{"alertname": "KubePodNotReady", "namespace": "openshift-e2e-loki"},
			Text:     "Loki is nice to have, but we can allow it to be down",
		},
		{
			Selector: map[string]string{"alertname": "KubeDeploymentReplicasMismatch", "namespace": "openshift-e2e-loki"},
			Text:     "Loki is nice to have, but we can allow it to be down",
		},
	}

	pendingAlertsWithBugs := helper.MetricConditions{
		{
			Selector: map[string]string{"alertname": "ClusterMonitoringOperatorReconciliationErrors"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1932624",
		},
		{
			Selector: map[string]string{"alertname": "KubeClientErrors"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=1925698",
		},
		{
			Selector: map[string]string{"alertname": "NetworkPodsCrashLooping"},
			Text:     "https://bugzilla.redhat.com/show_bug.cgi?id=2009078",
		},
	}
	allowedPendingAlerts := helper.MetricConditions{
		{
			Selector: map[string]string{"alertname": "etcdMemberCommunicationSlow"},
			Text:     "Excluded because it triggers during upgrade (detects ~5m of high latency immediately preceeding the end of the test), and we don't want to change the alert because it is correct",
		},
	}

	// we exclude alerts that have their own separate tests.
	for _, alertTest := range allowedalerts.AllAlertTests() {
		switch alertTest.AlertState() {
		case allowedalerts.AlertPending:
			// a pending test covers pending and everything above (firing)
			allowedPendingAlerts = append(allowedPendingAlerts,
				helper.MetricCondition{
					Selector: map[string]string{"alertname": alertTest.AlertName()},
					Text:     "has a separate e2e test",
				},
			)
			allowedFiringAlerts = append(allowedFiringAlerts,
				helper.MetricCondition{
					Selector: map[string]string{"alertname": alertTest.AlertName()},
					Text:     "has a separate e2e test",
				},
			)
		case allowedalerts.AlertInfo:
			// an info test covers all firing
			allowedFiringAlerts = append(allowedFiringAlerts,
				helper.MetricCondition{
					Selector: map[string]string{"alertname": alertTest.AlertName()},
					Text:     "has a separate e2e test",
				},
			)
		}
	}

	knownViolations := sets.NewString()
	unexpectedViolations := sets.NewString()
	unexpectedViolationsAsFlakes := sets.NewString()
	debug := sets.NewString()

	g.By("Checking for alerts")

	start := time.Now()

	// Block until upgrade is done
	g.By("Waiting for upgrade to finish before checking for alerts")
	<-done

	// Additonal delay after upgrade completion to allow pending alerts to settle
	g.By("Waiting before checking for alerts")
	time.Sleep(1 * time.Minute)

	testDuration := time.Now().Sub(start).Round(time.Second)

	// Invariant: No non-info level alerts should have fired during the upgrade
	firingAlertQuery := fmt.Sprintf(`
sort_desc(
count_over_time(ALERTS{alertstate="firing",severity!="info",alertname!~"Watchdog|AlertmanagerReceiversNotConfigured"}[%[1]s:1s])
) > 0
`, testDuration)
	result, err := helper.RunQuery(context.TODO(), t.prometheusClient, firingAlertQuery)
	o.Expect(err).NotTo(o.HaveOccurred(), "unable to check firing alerts during upgrade")
	for _, series := range result.Data.Result {
		labels := helper.StripLabels(series.Metric, "alertname", "alertstate", "prometheus")
		violation := fmt.Sprintf("alert %s fired for %s seconds with labels: %s", series.Metric["alertname"], series.Value, helper.LabelsAsSelector(labels))
		if cause := allowedFiringAlerts.Matches(series); cause != nil {
			debug.Insert(fmt.Sprintf("%s (allowed: %s)", violation, cause.Text))
			continue
		}
		if cause := firingAlertsWithBugs.Matches(series); cause != nil {
			knownViolations.Insert(fmt.Sprintf("%s (open bug: %s)", violation, cause.Text))
		} else {
			unexpectedViolations.Insert(violation)
		}
	}

	// Invariant: There should be no pending alerts 1m after the upgrade completes
	pendingAlertQuery := fmt.Sprintf(`
sort_desc(
  time() * ALERTS + 1
  -
  last_over_time((
    time() * ALERTS{alertname!~"Watchdog|AlertmanagerReceiversNotConfigured",alertstate="pending",severity!="info"}
    unless
    ALERTS offset 1s
  )[%[1]s:1s])
)
`, testDuration)
	result, err = helper.RunQuery(context.TODO(), t.prometheusClient, pendingAlertQuery)
	o.Expect(err).NotTo(o.HaveOccurred(), "unable to retrieve pending alerts after upgrade")
	for _, series := range result.Data.Result {
		labels := helper.StripLabels(series.Metric, "alertname", "alertstate", "prometheus")
		violation := fmt.Sprintf("alert %s pending for %s seconds with labels: %s", series.Metric["alertname"], series.Value, helper.LabelsAsSelector(labels))
		if cause := allowedPendingAlerts.Matches(series); cause != nil {
			debug.Insert(fmt.Sprintf("%s (allowed: %s)", violation, cause.Text))
			continue
		}
		if cause := pendingAlertsWithBugs.Matches(series); cause != nil {
			knownViolations.Insert(fmt.Sprintf("%s (open bug: %s)", violation, cause.Text))
		} else {
			// treat pending errors as a flake right now because we are still trying to determine the scope
			// TODO: move this to unexpectedViolations later
			unexpectedViolationsAsFlakes.Insert(violation)
		}
	}

	if len(debug) > 0 {
		framework.Logf("Alerts were detected during upgrade which are allowed:\n\n%s", strings.Join(debug.List(), "\n"))
	}
	if len(unexpectedViolations) > 0 {
		if !tolerateDuringSkew {
			framework.Failf("Unexpected alerts fired or pending during the upgrade:\n\n%s", strings.Join(unexpectedViolations.List(), "\n"))
		}
	}
	if flakes := sets.NewString().Union(knownViolations).Union(unexpectedViolations).Union(unexpectedViolationsAsFlakes); len(flakes) > 0 {
		disruption.FrameworkFlakef(f, "Unexpected alert behavior during upgrade:\n\n%s", strings.Join(flakes.List(), "\n"))
	}
	framework.Logf("No alerts fired during upgrade")
}

// Teardown cleans up any remaining resources.
func (t *UpgradeTest) Teardown(f *framework.Framework) {
	// rely on the namespace deletion to clean up everything
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"log"
	"os"

	"github.com/gorilla/mux"
	"github.com/rs/cors"
	// postgres driver for migrate
	_ "github.com/mattes/migrate/driver/postgres"
	"github.com/nuveo/prest/config"
	cfgMiddleware "github.com/nuveo/prest/config/middlewares"
	"github.com/nuveo/prest/config/router"
	"github.com/nuveo/prest/controllers"
	"github.com/nuveo/prest/middlewares"
	"github.com/spf13/cobra"
	"github.com/urfave/negroni"
)

var cfgFile string
var prestConfig config.Prest

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "prest",
	Short: "Serve a RESTful API from any PostgreSQL database",
	Long:  `Serve a RESTful API from any PostgreSQL database, start HTTP server`,
	Run: func(cmd *cobra.Command, args []string) {
		app()
	},
}

// Execute adds all child commands to the root command sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := RootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(-1)
	}
}

func init() {
	config.InitConf()
}

func app() {
	cfg := config.Prest{}

	err := config.Parse(&cfg)
	if err != nil {
		log.Fatalf("Error parsing conf: %s", err)
	}

	n := cfgMiddleware.GetApp()

	r := router.Get()

	r.HandleFunc("/databases", controllers.GetDatabases).Methods("GET")
	r.HandleFunc("/schemas", controllers.GetSchemas).Methods("GET")
	r.HandleFunc("/tables", controllers.GetTables).Methods("GET")
	r.HandleFunc("/_QUERIES/{queriesLocation}/{script}", controllers.ExecuteFromScripts)
	r.HandleFunc("/{database}/{schema}", controllers.GetTablesByDatabaseAndSchema).Methods("GET")

	crudRoutes := mux.NewRouter().PathPrefix("/").Subrouter().StrictSlash(true)

	crudRoutes.HandleFunc("/{database}/{schema}/{table}", controllers.SelectFromTables).Methods("GET")
	crudRoutes.HandleFunc("/{database}/{schema}/{table}", controllers.InsertInTables).Methods("POST")
	crudRoutes.HandleFunc("/{database}/{schema}/{table}", controllers.DeleteFromTable).Methods("DELETE")
	crudRoutes.HandleFunc("/{database}/{schema}/{table}", controllers.UpdateTable).Methods("PUT", "PATCH")

	r.PathPrefix("/").Handler(negroni.New(
		middlewares.AccessControl(),
		negroni.Wrap(crudRoutes),
	))

	if cfg.CORSAllowOrigin != nil {
		c := cors.New(cors.Options{
			AllowedOrigins: cfg.CORSAllowOrigin,
		})
		n.Use(c)
	}

	n.UseHandler(r)
	n.Run(fmt.Sprintf(":%v", cfg.HTTPPort))
}
<commit_msg>Refactoring to prevent multiple configuration loads (#131)<commit_after>package cmd

import (
	"fmt"
	"os"

	"github.com/gorilla/mux"
	"github.com/rs/cors"
	// postgres driver for migrate
	_ "github.com/mattes/migrate/driver/postgres"
	"github.com/nuveo/prest/config"
	cfgMiddleware "github.com/nuveo/prest/config/middlewares"
	"github.com/nuveo/prest/config/router"
	"github.com/nuveo/prest/controllers"
	"github.com/nuveo/prest/middlewares"
	"github.com/spf13/cobra"
	"github.com/urfave/negroni"
)

var cfgFile string
var prestConfig config.Prest

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "prest",
	Short: "Serve a RESTful API from any PostgreSQL database",
	Long:  `Serve a RESTful API from any PostgreSQL database, start HTTP server`,
	Run: func(cmd *cobra.Command, args []string) {
		app()
	},
}

// Execute adds all child commands to the root command sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := RootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(-1)
	}
}

func app() {
	n := cfgMiddleware.GetApp()
	r := router.Get()

	r.HandleFunc("/databases", controllers.GetDatabases).Methods("GET")
	r.HandleFunc("/schemas", controllers.GetSchemas).Methods("GET")
	r.HandleFunc("/tables", controllers.GetTables).Methods("GET")
	r.HandleFunc("/_QUERIES/{queriesLocation}/{script}", controllers.ExecuteFromScripts)
	r.HandleFunc("/{database}/{schema}", controllers.GetTablesByDatabaseAndSchema).Methods("GET")

	crudRoutes := mux.NewRouter().PathPrefix("/").Subrouter().StrictSlash(true)

	crudRoutes.HandleFunc("/{database}/{schema}/{table}", controllers.SelectFromTables).Methods("GET")
	crudRoutes.HandleFunc("/{database}/{schema}/{table}", controllers.InsertInTables).Methods("POST")
	crudRoutes.HandleFunc("/{database}/{schema}/{table}", controllers.DeleteFromTable).Methods("DELETE")
	crudRoutes.HandleFunc("/{database}/{schema}/{table}", controllers.UpdateTable).Methods("PUT", "PATCH")

	r.PathPrefix("/").Handler(negroni.New(
		middlewares.AccessControl(),
		negroni.Wrap(crudRoutes),
	))

	if config.PrestConf.CORSAllowOrigin != nil {
		c := cors.New(cors.Options{
			AllowedOrigins: config.PrestConf.CORSAllowOrigin,
		})
		n.Use(c)
	}

	n.UseHandler(r)
	n.Run(fmt.Sprintf(":%v", config.PrestConf.HTTPPort))
}
<|endoftext|>


<commit_before>// Copyright 2015 The Xorm 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 xorm

import (
	"database/sql"
	"fmt"
	"reflect"

	"github.com/go-xorm/core"
)

// Rows rows wrapper a rows to
type Rows struct {
	NoTypeCheck bool

	session     *Session
	stmt        *core.Stmt
	rows        *core.Rows
	fields      []string
	fieldsCount int
	beanType    reflect.Type
	lastError   error
}

func newRows(session *Session, bean interface{}) (*Rows, error) {
	rows := new(Rows)
	rows.session = session
	rows.beanType = reflect.Indirect(reflect.ValueOf(bean)).Type()

	defer rows.session.resetStatement()

	var sqlStr string
	var args []interface{}

	rows.session.Statement.setRefValue(rValue(bean))
	if len(session.Statement.TableName()) <= 0 {
		return nil, ErrTableNotFound
	}

	if rows.session.Statement.RawSQL == "" {
		sqlStr, args = rows.session.Statement.genGetSQL(bean)
	} else {
		sqlStr = rows.session.Statement.RawSQL
		args = rows.session.Statement.RawParams
	}

	for _, filter := range rows.session.Engine.dialect.Filters() {
		sqlStr = filter.Do(sqlStr, session.Engine.dialect, rows.session.Statement.RefTable)
	}

	rows.session.saveLastSQL(sqlStr, args...)
	var err error
	if rows.session.prepareStmt {
		rows.stmt, err = rows.session.DB().Prepare(sqlStr)
		if err != nil {
			rows.lastError = err
			rows.Close()
			return nil, err
		}

		rows.rows, err = rows.stmt.Query(args...)
		if err != nil {
			rows.lastError = err
			rows.Close()
			return nil, err
		}
	} else {
		rows.rows, err = rows.session.DB().Query(sqlStr, args...)
		if err != nil {
			rows.lastError = err
			rows.Close()
			return nil, err
		}
	}

	rows.fields, err = rows.rows.Columns()
	if err != nil {
		rows.lastError = err
		rows.Close()
		return nil, err
	}
	rows.fieldsCount = len(rows.fields)

	return rows, nil
}

// Next move cursor to next record, return false if end has reached
func (rows *Rows) Next() bool {
	if rows.lastError == nil && rows.rows != nil {
		hasNext := rows.rows.Next()
		if !hasNext {
			rows.lastError = sql.ErrNoRows
		}
		return hasNext
	}
	return false
}

// Err returns the error, if any, that was encountered during iteration. Err may be called after an explicit or implicit Close.
func (rows *Rows) Err() error {
	return rows.lastError
}

// Scan row record to bean properties
func (rows *Rows) Scan(bean interface{}) error {
	if rows.lastError != nil {
		return rows.lastError
	}

	if !rows.NoTypeCheck && reflect.Indirect(reflect.ValueOf(bean)).Type() != rows.beanType {
		return fmt.Errorf("scan arg is incompatible type to [%v]", rows.beanType)
	}

	_, err := rows.session.row2Bean(rows.rows, rows.fields, rows.fieldsCount, bean)
	return err
}

// Close session if session.IsAutoClose is true, and claimed any opened resources
func (rows *Rows) Close() error {
	if rows.session.IsAutoClose {
		defer rows.session.Close()
	}

	if rows.lastError == nil {
		if rows.rows != nil {
			rows.lastError = rows.rows.Close()
			if rows.lastError != nil {
				defer rows.stmt.Close()
				return rows.lastError
			}
		}
		if rows.stmt != nil {
			rows.lastError = rows.stmt.Close()
		}
	} else {
		if rows.stmt != nil {
			defer rows.stmt.Close()
		}
		if rows.rows != nil {
			defer rows.rows.Close()
		}
	}
	return rows.lastError
}
<commit_msg>remove unused field of rows<commit_after>// Copyright 2015 The Xorm 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 xorm

import (
	"database/sql"
	"fmt"
	"reflect"

	"github.com/go-xorm/core"
)

// Rows rows wrapper a rows to
type Rows struct {
	NoTypeCheck bool

	session   *Session
	stmt      *core.Stmt
	rows      *core.Rows
	fields    []string
	beanType  reflect.Type
	lastError error
}

func newRows(session *Session, bean interface{}) (*Rows, error) {
	rows := new(Rows)
	rows.session = session
	rows.beanType = reflect.Indirect(reflect.ValueOf(bean)).Type()

	defer rows.session.resetStatement()

	var sqlStr string
	var args []interface{}

	rows.session.Statement.setRefValue(rValue(bean))
	if len(session.Statement.TableName()) <= 0 {
		return nil, ErrTableNotFound
	}

	if rows.session.Statement.RawSQL == "" {
		sqlStr, args = rows.session.Statement.genGetSQL(bean)
	} else {
		sqlStr = rows.session.Statement.RawSQL
		args = rows.session.Statement.RawParams
	}

	for _, filter := range rows.session.Engine.dialect.Filters() {
		sqlStr = filter.Do(sqlStr, session.Engine.dialect, rows.session.Statement.RefTable)
	}

	rows.session.saveLastSQL(sqlStr, args...)
	var err error
	if rows.session.prepareStmt {
		rows.stmt, err = rows.session.DB().Prepare(sqlStr)
		if err != nil {
			rows.lastError = err
			rows.Close()
			return nil, err
		}

		rows.rows, err = rows.stmt.Query(args...)
		if err != nil {
			rows.lastError = err
			rows.Close()
			return nil, err
		}
	} else {
		rows.rows, err = rows.session.DB().Query(sqlStr, args...)
		if err != nil {
			rows.lastError = err
			rows.Close()
			return nil, err
		}
	}

	rows.fields, err = rows.rows.Columns()
	if err != nil {
		rows.lastError = err
		rows.Close()
		return nil, err
	}

	return rows, nil
}

// Next move cursor to next record, return false if end has reached
func (rows *Rows) Next() bool {
	if rows.lastError == nil && rows.rows != nil {
		hasNext := rows.rows.Next()
		if !hasNext {
			rows.lastError = sql.ErrNoRows
		}
		return hasNext
	}
	return false
}

// Err returns the error, if any, that was encountered during iteration. Err may be called after an explicit or implicit Close.
func (rows *Rows) Err() error {
	return rows.lastError
}

// Scan row record to bean properties
func (rows *Rows) Scan(bean interface{}) error {
	if rows.lastError != nil {
		return rows.lastError
	}

	if !rows.NoTypeCheck && reflect.Indirect(reflect.ValueOf(bean)).Type() != rows.beanType {
		return fmt.Errorf("scan arg is incompatible type to [%v]", rows.beanType)
	}

	_, err := rows.session.row2Bean(rows.rows, rows.fields, len(rows.fields), bean)
	return err
}

// Close session if session.IsAutoClose is true, and claimed any opened resources
func (rows *Rows) Close() error {
	if rows.session.IsAutoClose {
		defer rows.session.Close()
	}

	if rows.lastError == nil {
		if rows.rows != nil {
			rows.lastError = rows.rows.Close()
			if rows.lastError != nil {
				defer rows.stmt.Close()
				return rows.lastError
			}
		}
		if rows.stmt != nil {
			rows.lastError = rows.stmt.Close()
		}
	} else {
		if rows.stmt != nil {
			defer rows.stmt.Close()
		}
		if rows.rows != nil {
			defer rows.rows.Close()
		}
	}
	return rows.lastError
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"flag"
	"fmt"
	"html/template"
	"io"
	"net/http"
	"strings"
	"time"

	"golang.org/x/crypto/acme/autocert"

	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
	"github.com/labstack/gommon/color"
	_ "github.com/mattn/go-sqlite3"
)

var authDisabled bool
var dbFile = "scan.db"

type port struct {
	Port    int    `json:"port"`
	Proto   string `json:"proto"`
	Status  string `json:"status"`
	Service struct {
		Name   string `json:"name"`
		Banner string `json:"banner"`
	} `json:"service"`
}

// Results posted from masscan
type result struct {
	IP    string `json:"ip"`
	Ports []port `json:"ports"`
}

// Data retrieved from the database for display
type ipInfo struct {
	IP        string
	Port      int
	Proto     string
	FirstSeen string
	LastSeen  string
	New       bool
}

// Load all data for displaying in the browser
func load(s, fs, ls string) ([]ipInfo, error) {
	db, err := sql.Open("sqlite3", dbFile)
	if err != nil {
		return []ipInfo{}, err
	}
	defer db.Close()

	var where string
	var cond []string
	var params []interface{}
	if s != "" {
		cond = append(cond, `ip LIKE ?`)
		params = append(params, fmt.Sprintf("%%%s%%", s))
	}
	if fs != "" {
		cond = append(cond, `firstseen= ?`)
		params = append(params, fs)
	}
	if ls != "" {
		cond = append(cond, `lastseen= ?`)
		params = append(params, ls)
	}
	if len(cond) > 0 {
		where = fmt.Sprintf("WHERE %s", strings.Join(cond, " AND "))
	}

	qry := fmt.Sprintf(`SELECT ip, port, proto, firstseen, lastseen FROM scan %s ORDER BY port, proto, ip, lastseen`, where)
	rows, err := db.Query(qry, params...)
	if err != nil {
		return []ipInfo{}, err
	}

	defer rows.Close()

	var data []ipInfo
	var ip, proto, firstseen, lastseen string
	var port int
	var latest time.Time

	for rows.Next() {
		err := rows.Scan(&ip, &port, &proto, &firstseen, &lastseen)
		if err != nil {
			return []ipInfo{}, err
		}
		last, _ := time.Parse("2006-01-02 15:04", lastseen)
		if last.After(latest) {
			latest = last
		}
		data = append(data, ipInfo{ip, port, proto, firstseen, lastseen, false})
	}

	for i := range data {
		f, _ := time.Parse("2006-01-02 15:04", data[i].FirstSeen)
		l, _ := time.Parse("2006-01-02 15:04", data[i].LastSeen)
		if f.Equal(l) && l == latest {
			data[i].New = true
		}
	}

	return data, nil
}

// Save the results posted
func save(results []result) error {
	db, err := sql.Open("sqlite3", dbFile)
	if err != nil {
		return err
	}
	defer db.Close()

	txn, err := db.Begin()
	if err != nil {
		return err
	}

	insert, err := txn.Prepare(`INSERT INTO scan (ip, port, proto, firstseen, lastseen) VALUES (?, ?, ?, ?, ?)`)
	if err != nil {
		txn.Rollback()
		return err
	}
	qry, err := db.Prepare(`SELECT 1 FROM scan WHERE ip=? AND port=? AND proto=?`)
	if err != nil {
		txn.Rollback()
		return err
	}
	update, err := txn.Prepare(`UPDATE scan SET lastseen=? WHERE ip=? AND port=? AND proto=?`)
	if err != nil {
		txn.Rollback()
		return err
	}

	now := time.Now()
	nowString := now.Format("2006-01-02 15:04")

	for _, r := range results {
		// Although it's an array, only one port is in each
		port := r.Ports[0]

		// Search for the IP/port/proto combo
		// If it exists, update `lastseen`, else insert a new record

		// Because we have to scan into something
		var x int
		err := qry.QueryRow(r.IP, port.Port, port.Proto).Scan(&x)
		switch {
		case err == sql.ErrNoRows:
			_, err = insert.Exec(r.IP, port.Port, port.Proto, nowString, nowString)
			if err != nil {
				txn.Rollback()
				return err
			}
			continue
		case err != nil:
			txn.Rollback()
			return err
		}

		_, err = update.Exec(nowString, r.IP, port.Port, port.Proto)
		if err != nil {
			txn.Rollback()
			return err
		}
	}

	txn.Commit()
	return nil
}

// Template is a template
type Template struct {
	templates *template.Template
}

// Render renders template
func (t *Template) Render(w io.Writer, name string, data interface{}, c echo.Context) error {
	return t.templates.ExecuteTemplate(w, name, data)
}

type indexData struct {
	Authenticated bool
	User          string
	Total         int
	Latest        int
	New           int
	LastSeen      string
	Results       []ipInfo
}

// Handler for GET /
func index(c echo.Context) error {
	var user string
	if !authDisabled {
		session, err := store.Get(c.Request(), "user")
		if err != nil {
			return c.String(http.StatusInternalServerError, err.Error())
		}
		if _, ok := session.Values["user"]; !ok {
			data := indexData{}
			return c.Render(http.StatusOK, "index", data)
		}
		user = session.Values["user"].(string)
	}

	ip := c.QueryParam("ip")
	firstSeen := c.QueryParam("firstseen")
	lastSeen := c.QueryParam("lastseen")
	results, err := load(ip, firstSeen, lastSeen)
	if err != nil {
		return c.String(http.StatusInternalServerError, err.Error())
	}

	data := indexData{Authenticated: true, User: user, Results: results, Total: len(results)}

	timeFmt := "2006-01-02 15:04"

	// Find all the latest results and store the number in the struct
	var latest time.Time
	for _, r := range results {
		last, _ := time.Parse(timeFmt, r.LastSeen)
		if last.After(latest) {
			latest = last
		}
	}
	for _, r := range results {
		last, _ := time.Parse(timeFmt, r.LastSeen)
		if last.Equal(latest) {
			data.Latest++
		}
		if r.New {
			data.New++
		}
	}
	data.LastSeen = latest.Format(timeFmt)

	return c.Render(http.StatusOK, "index", data)
}

// Handler for GET /ips.json
// This is used as the prefetch for Typeahead.js
func ips(c echo.Context) error {
	data, err := load("", "", "")
	if err != nil {
		c.String(http.StatusInternalServerError, err.Error())
	}
	var ips []string
	for _, r := range data {
		ips = append(ips, r.IP)
	}
	return c.JSON(http.StatusOK, ips)
}

// Handler for POST /results
func recvResults(c echo.Context) error {
	res := new([]result)
	err := c.Bind(res)
	if err != nil {
		return c.String(http.StatusInternalServerError, err.Error())
	}

	err = save(*res)
	if err != nil {
		return c.String(http.StatusInternalServerError, err.Error())
	}

	return c.NoContent(http.StatusOK)
}

func main() {
	flag.BoolVar(&authDisabled, "no-auth", false, "Disable authentication")
	httpAddr := flag.String("http.addr", ":80", "HTTP address:port")
	httpsAddr := flag.String("https.addr", ":443", "HTTPS address:port")
	tls := flag.Bool("tls", false, "Enable AutoTLS")
	tlsHostname := flag.String("tls.hostname", "", "(Optional) Hostname to restrict AutoTLS")
	flag.Parse()

	t := &Template{
		templates: template.Must(template.ParseGlob("views/*.html")),
	}

	e := echo.New()

	if authDisabled {
		color.Println(color.Red("Authentication Disabled"))
	}

	if *tls {
		if *tlsHostname != "" {
			e.AutoTLSManager.HostPolicy = autocert.HostWhitelist(*tlsHostname)
		}
		e.AutoTLSManager.Cache = autocert.DirCache(".cache")
		e.Pre(middleware.HTTPSRedirect())
	}

	e.Renderer = t
	e.Use(middleware.Logger())
	e.GET("/", index)
	e.GET("/auth", authHandler)
	e.GET("/login", loginHandler)
	e.GET("/ips.json", ips)
	e.POST("/results", recvResults)
	e.Static("/static", "static")

	if *tls {
		go func() { e.Logger.Fatal(e.Start(*httpAddr)) }()
		e.Logger.Fatal(e.StartAutoTLS(*httpsAddr))
	}
	e.Logger.Fatal(e.Start(*httpAddr))
}
<commit_msg>Split results parsing from index handler<commit_after>package main

import (
	"database/sql"
	"flag"
	"fmt"
	"html/template"
	"io"
	"net/http"
	"strings"
	"time"

	"golang.org/x/crypto/acme/autocert"

	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
	"github.com/labstack/gommon/color"
	_ "github.com/mattn/go-sqlite3"
)

var authDisabled bool
var dbFile = "scan.db"

type port struct {
	Port    int    `json:"port"`
	Proto   string `json:"proto"`
	Status  string `json:"status"`
	Service struct {
		Name   string `json:"name"`
		Banner string `json:"banner"`
	} `json:"service"`
}

// Results posted from masscan
type result struct {
	IP    string `json:"ip"`
	Ports []port `json:"ports"`
}

// Data retrieved from the database for display
type ipInfo struct {
	IP        string
	Port      int
	Proto     string
	FirstSeen string
	LastSeen  string
	New       bool
}

// Load all data for displaying in the browser
func load(s, fs, ls string) ([]ipInfo, error) {
	db, err := sql.Open("sqlite3", dbFile)
	if err != nil {
		return []ipInfo{}, err
	}
	defer db.Close()

	var where string
	var cond []string
	var params []interface{}
	if s != "" {
		cond = append(cond, `ip LIKE ?`)
		params = append(params, fmt.Sprintf("%%%s%%", s))
	}
	if fs != "" {
		cond = append(cond, `firstseen= ?`)
		params = append(params, fs)
	}
	if ls != "" {
		cond = append(cond, `lastseen= ?`)
		params = append(params, ls)
	}
	if len(cond) > 0 {
		where = fmt.Sprintf("WHERE %s", strings.Join(cond, " AND "))
	}

	qry := fmt.Sprintf(`SELECT ip, port, proto, firstseen, lastseen FROM scan %s ORDER BY port, proto, ip, lastseen`, where)
	rows, err := db.Query(qry, params...)
	if err != nil {
		return []ipInfo{}, err
	}

	defer rows.Close()

	var data []ipInfo
	var ip, proto, firstseen, lastseen string
	var port int
	var latest time.Time

	for rows.Next() {
		err := rows.Scan(&ip, &port, &proto, &firstseen, &lastseen)
		if err != nil {
			return []ipInfo{}, err
		}
		last, _ := time.Parse("2006-01-02 15:04", lastseen)
		if last.After(latest) {
			latest = last
		}
		data = append(data, ipInfo{ip, port, proto, firstseen, lastseen, false})
	}

	for i := range data {
		f, _ := time.Parse("2006-01-02 15:04", data[i].FirstSeen)
		l, _ := time.Parse("2006-01-02 15:04", data[i].LastSeen)
		if f.Equal(l) && l == latest {
			data[i].New = true
		}
	}

	return data, nil
}

// Save the results posted
func save(results []result) error {
	db, err := sql.Open("sqlite3", dbFile)
	if err != nil {
		return err
	}
	defer db.Close()

	txn, err := db.Begin()
	if err != nil {
		return err
	}

	insert, err := txn.Prepare(`INSERT INTO scan (ip, port, proto, firstseen, lastseen) VALUES (?, ?, ?, ?, ?)`)
	if err != nil {
		txn.Rollback()
		return err
	}
	qry, err := db.Prepare(`SELECT 1 FROM scan WHERE ip=? AND port=? AND proto=?`)
	if err != nil {
		txn.Rollback()
		return err
	}
	update, err := txn.Prepare(`UPDATE scan SET lastseen=? WHERE ip=? AND port=? AND proto=?`)
	if err != nil {
		txn.Rollback()
		return err
	}

	now := time.Now()
	nowString := now.Format("2006-01-02 15:04")

	for _, r := range results {
		// Although it's an array, only one port is in each
		port := r.Ports[0]

		// Search for the IP/port/proto combo
		// If it exists, update `lastseen`, else insert a new record

		// Because we have to scan into something
		var x int
		err := qry.QueryRow(r.IP, port.Port, port.Proto).Scan(&x)
		switch {
		case err == sql.ErrNoRows:
			_, err = insert.Exec(r.IP, port.Port, port.Proto, nowString, nowString)
			if err != nil {
				txn.Rollback()
				return err
			}
			continue
		case err != nil:
			txn.Rollback()
			return err
		}

		_, err = update.Exec(nowString, r.IP, port.Port, port.Proto)
		if err != nil {
			txn.Rollback()
			return err
		}
	}

	txn.Commit()
	return nil
}

// Template is a template
type Template struct {
	templates *template.Template
}

// Render renders template
func (t *Template) Render(w io.Writer, name string, data interface{}, c echo.Context) error {
	return t.templates.ExecuteTemplate(w, name, data)
}

type indexData struct {
	Authenticated bool
	User          string
	scanData
}

type scanData struct {
	Total    int
	Latest   int
	New      int
	LastSeen string
	Results  []ipInfo
}

func resultData(ip, fs, ls string) (scanData, error) {
	results, err := load(ip, fs, ls)
	if err != nil {
		return scanData{}, err
	}

	timeFmt := "2006-01-02 15:04"
	data := scanData{
		Results: results,
		Total:   len(results),
	}

	// Find all the latest results and store the number in the struct
	var latest time.Time
	for _, r := range results {
		last, _ := time.Parse(timeFmt, r.LastSeen)
		if last.After(latest) {
			latest = last
		}
	}
	for _, r := range results {
		last, _ := time.Parse(timeFmt, r.LastSeen)
		if last.Equal(latest) {
			data.Latest++
		}
		if r.New {
			data.New++
		}
	}
	data.LastSeen = latest.Format(timeFmt)

	return data, nil
}

// Handler for GET /
func index(c echo.Context) error {
	var user string
	if !authDisabled {
		session, err := store.Get(c.Request(), "user")
		if err != nil {
			return c.String(http.StatusInternalServerError, err.Error())
		}
		if _, ok := session.Values["user"]; !ok {
			data := indexData{}
			return c.Render(http.StatusOK, "index", data)
		}
		user = session.Values["user"].(string)
	}

	ip := c.QueryParam("ip")
	firstSeen := c.QueryParam("firstseen")
	lastSeen := c.QueryParam("lastseen")

	results, err := resultData(ip, firstSeen, lastSeen)
	if err != nil {
		return c.String(http.StatusInternalServerError, err.Error())
	}

	data := indexData{Authenticated: true, User: user, scanData: results}

	return c.Render(http.StatusOK, "index", data)
}

// Handler for GET /ips.json
// This is used as the prefetch for Typeahead.js
func ips(c echo.Context) error {
	data, err := load("", "", "")
	if err != nil {
		c.String(http.StatusInternalServerError, err.Error())
	}
	var ips []string
	for _, r := range data {
		ips = append(ips, r.IP)
	}
	return c.JSON(http.StatusOK, ips)
}

// Handler for POST /results
func recvResults(c echo.Context) error {
	res := new([]result)
	err := c.Bind(res)
	if err != nil {
		return c.String(http.StatusInternalServerError, err.Error())
	}

	err = save(*res)
	if err != nil {
		return c.String(http.StatusInternalServerError, err.Error())
	}

	return c.NoContent(http.StatusOK)
}

func main() {
	flag.BoolVar(&authDisabled, "no-auth", false, "Disable authentication")
	httpAddr := flag.String("http.addr", ":80", "HTTP address:port")
	httpsAddr := flag.String("https.addr", ":443", "HTTPS address:port")
	tls := flag.Bool("tls", false, "Enable AutoTLS")
	tlsHostname := flag.String("tls.hostname", "", "(Optional) Hostname to restrict AutoTLS")
	flag.Parse()

	t := &Template{
		templates: template.Must(template.ParseGlob("views/*.html")),
	}

	e := echo.New()

	if authDisabled {
		color.Println(color.Red("Authentication Disabled"))
	}

	if *tls {
		if *tlsHostname != "" {
			e.AutoTLSManager.HostPolicy = autocert.HostWhitelist(*tlsHostname)
		}
		e.AutoTLSManager.Cache = autocert.DirCache(".cache")
		e.Pre(middleware.HTTPSRedirect())
	}

	e.Renderer = t
	e.Use(middleware.Logger())
	e.GET("/", index)
	e.GET("/auth", authHandler)
	e.GET("/login", loginHandler)
	e.GET("/ips.json", ips)
	e.POST("/results", recvResults)
	e.Static("/static", "static")

	if *tls {
		go func() { e.Logger.Fatal(e.Start(*httpAddr)) }()
		e.Logger.Fatal(e.StartAutoTLS(*httpsAddr))
	}
	e.Logger.Fatal(e.Start(*httpAddr))
}
<|endoftext|>


<commit_before>package runner

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strconv"
	"syscall"
	"time"

	"code.cloudfoundry.org/garden/client"
	"code.cloudfoundry.org/garden/client/connection"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagertest"
	"github.com/eapache/go-resiliency/retrier"
	"github.com/onsi/ginkgo"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
)

const MNT_DETACH = 0x2

var RootFSPath = os.Getenv("GARDEN_TEST_ROOTFS")
var DataDir = os.Getenv("GARDEN_TEST_GRAPHPATH")
var TarBin = os.Getenv("GARDEN_TAR_PATH")

type RunningGarden struct {
	client.Client

	runner  *ginkgomon.Runner
	process ifrit.Process

	debugIP   string
	debugPort int

	Pid int

	Tmpdir string

	DepotDir  string
	DataDir   string
	GraphPath string

	logger lager.Logger
}

func Start(bin, initBin, nstarBin, dadooBin string, supplyDefaultRootfs bool, argv ...string) *RunningGarden {
	network := "unix"
	addr := fmt.Sprintf("/tmp/garden_%d.sock", GinkgoParallelNode())
	tmpDir := filepath.Join(
		os.TempDir(),
		fmt.Sprintf("test-garden-%d", ginkgo.GinkgoParallelNode()),
	)

	if DataDir == "" {
		// This must be set outside of the Ginkgo node directory (tmpDir) because
		// otherwise the Concourse worker may run into one of the AUFS kernel
		// module bugs that cause the VM to become unresponsive.
		DataDir = "/tmp/aufs_mount"
	}

	graphPath := filepath.Join(DataDir, fmt.Sprintf("node-%d", ginkgo.GinkgoParallelNode()))
	depotDir := filepath.Join(tmpDir, "containers")

	MustMountTmpfs(graphPath)

	r := &RunningGarden{
		DepotDir: depotDir,

		DataDir:   DataDir,
		GraphPath: graphPath,
		Tmpdir:    tmpDir,
		logger:    lagertest.NewTestLogger("garden-runner"),

		Client: client.New(connection.New(network, addr)),
	}

	c := cmd(tmpDir, depotDir, graphPath, network, addr, bin, initBin, nstarBin, dadooBin, TarBin, supplyDefaultRootfs, argv...)
	c.Env = append(os.Environ(), fmt.Sprintf("TMPDIR=%s", tmpDir))
	r.runner = ginkgomon.New(ginkgomon.Config{
		Name:              "guardian",
		Command:           c,
		AnsiColorCode:     "31m",
		StartCheck:        "guardian.started",
		StartCheckTimeout: 30 * time.Second,
	})
	r.process = ifrit.Invoke(r.runner)
	r.Pid = c.Process.Pid

	for i, arg := range c.Args {
		if arg == "--debug-bind-ip" {
			r.debugIP = c.Args[i+1]
		}
		if arg == "--debug-bind-port" {
			r.debugPort, _ = strconv.Atoi(c.Args[i+1])
		}
	}

	return r
}

func (r *RunningGarden) Kill() error {
	r.process.Signal(syscall.SIGKILL)
	select {
	case err := <-r.process.Wait():
		return err
	case <-time.After(time.Second * 10):
		r.process.Signal(syscall.SIGKILL)
		return errors.New("timed out waiting for garden to shutdown after 10 seconds")
	}
}

func (r *RunningGarden) DestroyAndStop() error {
	if err := r.DestroyContainers(); err != nil {
		return err
	}

	if err := r.Stop(); err != nil {
		return err
	}

	return nil
}

func (r *RunningGarden) Stop() error {
	r.process.Signal(syscall.SIGTERM)

	var err error
	for i := 0; i < 5; i++ {
		select {
		case err := <-r.process.Wait():
			return err
		case <-time.After(time.Second * 5):
			r.process.Signal(syscall.SIGTERM)
			err = errors.New("timed out waiting for garden to shutdown after 5 seconds")
		}
	}

	r.process.Signal(syscall.SIGKILL)
	return err
}

func cmd(tmpdir, depotDir, graphPath, network, addr, bin, initBin, nstarBin, dadooBin, tarBin string, supplyDefaultRootfs bool, argv ...string) *exec.Cmd {
	Expect(os.MkdirAll(tmpdir, 0755)).To(Succeed())

	snapshotsPath := filepath.Join(tmpdir, "snapshots")

	Expect(os.MkdirAll(depotDir, 0755)).To(Succeed())

	Expect(os.MkdirAll(snapshotsPath, 0755)).To(Succeed())

	appendDefaultFlag := func(ar []string, key, value string) []string {
		for _, a := range argv {
			if a == key {
				return ar
			}
		}

		if value != "" {
			return append(ar, key, value)
		} else {
			return append(ar, key)
		}
	}

	gardenArgs := make([]string, len(argv))
	copy(gardenArgs, argv)

	switch network {
	case "tcp":
		gardenArgs = appendDefaultFlag(gardenArgs, "--bind-ip", addr)
	case "unix":
		gardenArgs = appendDefaultFlag(gardenArgs, "--bind-socket", addr)
	}

	if supplyDefaultRootfs {
		gardenArgs = appendDefaultFlag(gardenArgs, "--default-rootfs", RootFSPath)
	}

	gardenArgs = appendDefaultFlag(gardenArgs, "--depot", depotDir)
	gardenArgs = appendDefaultFlag(gardenArgs, "--graph", graphPath)
	gardenArgs = appendDefaultFlag(gardenArgs, "--tag", fmt.Sprintf("%d", GinkgoParallelNode()))
	gardenArgs = appendDefaultFlag(gardenArgs, "--network-pool", fmt.Sprintf("10.254.%d.0/22", 4*GinkgoParallelNode()))
	gardenArgs = appendDefaultFlag(gardenArgs, "--init-bin", initBin)
	gardenArgs = appendDefaultFlag(gardenArgs, "--dadoo-bin", dadooBin)
	gardenArgs = appendDefaultFlag(gardenArgs, "--nstar-bin", nstarBin)
	gardenArgs = appendDefaultFlag(gardenArgs, "--tar-bin", tarBin)
	gardenArgs = appendDefaultFlag(gardenArgs, "--port-pool-start", fmt.Sprintf("%d", GinkgoParallelNode()*10000))
	gardenArgs = appendDefaultFlag(gardenArgs, "--debug-bind-ip", "0.0.0.0")
	gardenArgs = appendDefaultFlag(gardenArgs, "--debug-bind-port", fmt.Sprintf("%d", 8080+ginkgo.GinkgoParallelNode()))

	return exec.Command(bin, gardenArgs...)
}

func (r *RunningGarden) Cleanup() {
	// unmount aufs since the docker graph driver leaves this around,
	// otherwise the following commands might fail
	retry := retrier.New(retrier.ConstantBackoff(200, 500*time.Millisecond), nil)

	err := retry.Run(func() error {
		if err := os.RemoveAll(path.Join(r.GraphPath, "aufs")); err == nil {
			return nil // if we can remove it, it's already unmounted
		}

		if err := syscall.Unmount(path.Join(r.GraphPath, "aufs"), MNT_DETACH); err != nil {
			r.logger.Error("failed-unmount-attempt", err)
			return err
		}

		return nil
	})

	if err != nil {
		r.logger.Error("failed-to-unmount", err)
	}

	MustUnmountTmpfs(r.GraphPath)

	// In the kernel version 3.19.0-51-generic the code bellow results in
	// hanging the running VM. We are not deleting the node-X directories. They
	// are empty and the next test will re-use them. We will stick with that
	// workaround until we can test on a newer kernel that will hopefully not
	// have this bug.
	//
	// if err := os.RemoveAll(r.GraphPath); err != nil {
	// 	r.logger.Error("remove-graph", err)
	// }

	r.logger.Info("cleanup-tempdirs")
	if err := os.RemoveAll(r.Tmpdir); err != nil {
		r.logger.Error("cleanup-tempdirs-failed", err, lager.Data{"tmpdir": r.Tmpdir})
	} else {
		r.logger.Info("tempdirs-removed")
	}
}

func (r *RunningGarden) DestroyContainers() error {
	containers, err := r.Containers(nil)
	if err != nil {
		return err
	}

	for _, container := range containers {
		r.Destroy(container.Handle())
	}

	return nil
}

type debugVars struct {
	NumGoRoutines int `json:"numGoRoutines"`
}

func (r *RunningGarden) NumGoroutines() (int, error) {
	debugURL := fmt.Sprintf("http://%s:%d/debug/vars", r.debugIP, r.debugPort)
	res, err := http.Get(debugURL)
	defer res.Body.Close()
	if err != nil {
		return 0, err
	}

	decoder := json.NewDecoder(res.Body)
	var debugVarsData debugVars
	err = decoder.Decode(&debugVarsData)
	if err != nil {
		return 0, err
	}

	return debugVarsData.NumGoRoutines, nil
}

func (r *RunningGarden) Buffer() *gbytes.Buffer {
	return r.runner.Buffer()
}

func (r *RunningGarden) ExitCode() int {
	return r.runner.ExitCode()
}
<commit_msg>Expose an ifrit.Runner that can be used to run guardian<commit_after>package runner

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strconv"
	"syscall"
	"time"

	"code.cloudfoundry.org/garden/client"
	"code.cloudfoundry.org/garden/client/connection"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/lager/lagertest"
	"github.com/eapache/go-resiliency/retrier"
	"github.com/onsi/ginkgo"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
)

const MNT_DETACH = 0x2

var RootFSPath = os.Getenv("GARDEN_TEST_ROOTFS")
var DataDir string
var TarBin = os.Getenv("GARDEN_TAR_PATH")

type RunningGarden struct {
	client.Client

	runner  GardenRunner
	process ifrit.Process

	debugIP   string
	debugPort int

	Pid int

	Tmpdir string

	DepotDir  string
	DataDir   string
	GraphPath string

	logger lager.Logger
}

type GardenRunner struct {
	*ginkgomon.Runner
	Cmd           *exec.Cmd
	TmpDir        string
	GraphPath     string
	DepotDir      string
	DebugIp       string
	DebugPort     int
	Network, Addr string
}

func init() {
	DataDir = os.Getenv("GARDEN_TEST_GRAPHPATH")
	if DataDir == "" {
		// This must be set outside of the Ginkgo node directory (tmpDir) because
		// otherwise the Concourse worker may run into one of the AUFS kernel
		// module bugs that cause the VM to become unresponsive.
		DataDir = "/tmp/aufs_mount"
	}
}

func NewGardenRunner(bin, initBin, nstarBin, dadooBin string, supplyDefaultRootfs bool, argv ...string) GardenRunner {
	r := GardenRunner{}

	r.Network = "unix"
	r.Addr = fmt.Sprintf("/tmp/garden_%d.sock", GinkgoParallelNode())

	r.TmpDir = filepath.Join(
		os.TempDir(),
		fmt.Sprintf("test-garden-%d", ginkgo.GinkgoParallelNode()),
	)

	r.GraphPath = filepath.Join(DataDir, fmt.Sprintf("node-%d", ginkgo.GinkgoParallelNode()))
	r.DepotDir = filepath.Join(r.TmpDir, "containers")

	MustMountTmpfs(r.GraphPath)

	r.Cmd = cmd(r.TmpDir, r.DepotDir, r.GraphPath, r.Network, r.Addr, bin, initBin, nstarBin, dadooBin, TarBin, supplyDefaultRootfs, argv...)
	r.Cmd.Env = append(os.Environ(), fmt.Sprintf("TMPDIR=%s", r.TmpDir))

	for i, arg := range r.Cmd.Args {
		if arg == "--debug-bind-ip" {
			r.DebugIp = r.Cmd.Args[i+1]
		}
		if arg == "--debug-bind-port" {
			r.DebugPort, _ = strconv.Atoi(r.Cmd.Args[i+1])
		}
	}

	r.Runner = ginkgomon.New(ginkgomon.Config{
		Name:              "guardian",
		Command:           r.Cmd,
		AnsiColorCode:     "31m",
		StartCheck:        "guardian.started",
		StartCheckTimeout: 30 * time.Second,
	})

	return r
}

func Start(bin, initBin, nstarBin, dadooBin string, supplyDefaultRootfs bool, argv ...string) *RunningGarden {
	runner := NewGardenRunner(bin, initBin, nstarBin, dadooBin, supplyDefaultRootfs, argv...)

	r := &RunningGarden{
		runner:   runner,
		DepotDir: runner.DepotDir,

		DataDir:   DataDir,
		GraphPath: runner.GraphPath,
		Tmpdir:    runner.TmpDir,
		logger:    lagertest.NewTestLogger("garden-runner"),

		debugIP:   runner.DebugIp,
		debugPort: runner.DebugPort,

		Client: client.New(connection.New(runner.Network, runner.Addr)),
	}

	r.process = ifrit.Invoke(r.runner)
	r.Pid = runner.Cmd.Process.Pid

	return r
}

func (r *RunningGarden) Kill() error {
	r.process.Signal(syscall.SIGKILL)
	select {
	case err := <-r.process.Wait():
		return err
	case <-time.After(time.Second * 10):
		r.process.Signal(syscall.SIGKILL)
		return errors.New("timed out waiting for garden to shutdown after 10 seconds")
	}
}

func (r *RunningGarden) DestroyAndStop() error {
	if err := r.DestroyContainers(); err != nil {
		return err
	}

	if err := r.Stop(); err != nil {
		return err
	}

	return nil
}

func (r *RunningGarden) Stop() error {
	r.process.Signal(syscall.SIGTERM)

	var err error
	for i := 0; i < 5; i++ {
		select {
		case err := <-r.process.Wait():
			return err
		case <-time.After(time.Second * 5):
			r.process.Signal(syscall.SIGTERM)
			err = errors.New("timed out waiting for garden to shutdown after 5 seconds")
		}
	}

	r.process.Signal(syscall.SIGKILL)
	return err
}

func cmd(tmpdir, depotDir, graphPath, network, addr, bin, initBin, nstarBin, dadooBin, tarBin string, supplyDefaultRootfs bool, argv ...string) *exec.Cmd {
	Expect(os.MkdirAll(tmpdir, 0755)).To(Succeed())

	snapshotsPath := filepath.Join(tmpdir, "snapshots")

	Expect(os.MkdirAll(depotDir, 0755)).To(Succeed())

	Expect(os.MkdirAll(snapshotsPath, 0755)).To(Succeed())

	appendDefaultFlag := func(ar []string, key, value string) []string {
		for _, a := range argv {
			if a == key {
				return ar
			}
		}

		if value != "" {
			return append(ar, key, value)
		} else {
			return append(ar, key)
		}
	}

	gardenArgs := make([]string, len(argv))
	copy(gardenArgs, argv)

	switch network {
	case "tcp":
		gardenArgs = appendDefaultFlag(gardenArgs, "--bind-ip", addr)
	case "unix":
		gardenArgs = appendDefaultFlag(gardenArgs, "--bind-socket", addr)
	}

	if supplyDefaultRootfs {
		gardenArgs = appendDefaultFlag(gardenArgs, "--default-rootfs", RootFSPath)
	}

	gardenArgs = appendDefaultFlag(gardenArgs, "--depot", depotDir)
	gardenArgs = appendDefaultFlag(gardenArgs, "--graph", graphPath)
	gardenArgs = appendDefaultFlag(gardenArgs, "--tag", fmt.Sprintf("%d", GinkgoParallelNode()))
	gardenArgs = appendDefaultFlag(gardenArgs, "--network-pool", fmt.Sprintf("10.254.%d.0/22", 4*GinkgoParallelNode()))
	gardenArgs = appendDefaultFlag(gardenArgs, "--init-bin", initBin)
	gardenArgs = appendDefaultFlag(gardenArgs, "--dadoo-bin", dadooBin)
	gardenArgs = appendDefaultFlag(gardenArgs, "--nstar-bin", nstarBin)
	gardenArgs = appendDefaultFlag(gardenArgs, "--tar-bin", tarBin)
	gardenArgs = appendDefaultFlag(gardenArgs, "--port-pool-start", fmt.Sprintf("%d", GinkgoParallelNode()*10000))
	gardenArgs = appendDefaultFlag(gardenArgs, "--debug-bind-ip", "0.0.0.0")
	gardenArgs = appendDefaultFlag(gardenArgs, "--debug-bind-port", fmt.Sprintf("%d", 8080+ginkgo.GinkgoParallelNode()))

	return exec.Command(bin, gardenArgs...)
}

func (r *RunningGarden) Cleanup() {
	// unmount aufs since the docker graph driver leaves this around,
	// otherwise the following commands might fail
	retry := retrier.New(retrier.ConstantBackoff(200, 500*time.Millisecond), nil)

	err := retry.Run(func() error {
		if err := os.RemoveAll(path.Join(r.GraphPath, "aufs")); err == nil {
			return nil // if we can remove it, it's already unmounted
		}

		if err := syscall.Unmount(path.Join(r.GraphPath, "aufs"), MNT_DETACH); err != nil {
			r.logger.Error("failed-unmount-attempt", err)
			return err
		}

		return nil
	})

	if err != nil {
		r.logger.Error("failed-to-unmount", err)
	}

	MustUnmountTmpfs(r.GraphPath)

	// In the kernel version 3.19.0-51-generic the code bellow results in
	// hanging the running VM. We are not deleting the node-X directories. They
	// are empty and the next test will re-use them. We will stick with that
	// workaround until we can test on a newer kernel that will hopefully not
	// have this bug.
	//
	// if err := os.RemoveAll(r.GraphPath); err != nil {
	// 	r.logger.Error("remove-graph", err)
	// }

	r.logger.Info("cleanup-tempdirs")
	if err := os.RemoveAll(r.Tmpdir); err != nil {
		r.logger.Error("cleanup-tempdirs-failed", err, lager.Data{"tmpdir": r.Tmpdir})
	} else {
		r.logger.Info("tempdirs-removed")
	}
}

func (r *RunningGarden) DestroyContainers() error {
	containers, err := r.Containers(nil)
	if err != nil {
		return err
	}

	for _, container := range containers {
		r.Destroy(container.Handle())
	}

	return nil
}

type debugVars struct {
	NumGoRoutines int `json:"numGoRoutines"`
}

func (r *RunningGarden) NumGoroutines() (int, error) {
	debugURL := fmt.Sprintf("http://%s:%d/debug/vars", r.debugIP, r.debugPort)
	res, err := http.Get(debugURL)
	if err != nil {
		return 0, err
	}
	defer res.Body.Close()

	decoder := json.NewDecoder(res.Body)
	var debugVarsData debugVars
	err = decoder.Decode(&debugVarsData)
	if err != nil {
		return 0, err
	}

	return debugVarsData.NumGoRoutines, nil
}

func (r *RunningGarden) Buffer() *gbytes.Buffer {
	return r.runner.Buffer()
}

func (r *RunningGarden) ExitCode() int {
	return r.runner.ExitCode()
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"log"
	"net/http"
)

type Handler struct {
	Pattern, Root string
}

func registerHandlers(mux *http.ServeMux, paths map[string]string) {
	for pattern, root := range paths {
		log.Printf("Registering handler with pattern: %s, root path: %s",
			pattern, root)
		mux.Handle(pattern, http.StripPrefix(pattern, http.FileServer(http.Dir(root))))
	}
}

func parsePaths(paths string) (err error, pathMap map[string]string) {
	err = json.Unmarshal([]byte(paths), &pathMap)
	return
}

func main() {
	var listen = flag.String("listen", ":8080",
		"Interface/port to listen on. eg. :8080 or 127.0.0.1:8080")
	var pathsRaw = flag.String("paths", `{"/": "."}`,
		"Paths to serve. A json object with the keys as the url pattern, and "+
			"the value as the root. Default serves current folder.")
	flag.Parse()

	err, paths := parsePaths(*pathsRaw)

	if err != nil {
		log.Fatal(err)
	}

	mux := http.NewServeMux()
	registerHandlers(mux, paths)
	log.Println("Listening on: ", *listen)

	log.Fatal(http.ListenAndServe(*listen, mux))
}
<commit_msg>Code cleanup<commit_after>package main

import (
	"encoding/json"
	"flag"
	"log"
	"net/http"
)

func registerHandlers(mux *http.ServeMux, paths map[string]string) {
	for pattern, root := range paths {
		log.Printf("Registering handler with pattern: %s, root path: %s",
			pattern, root)
		mux.Handle(pattern, http.StripPrefix(pattern, http.FileServer(http.Dir(root))))
	}
}

func parsePaths(paths string) (pathMap map[string]string, err error) {
	err = json.Unmarshal([]byte(paths), &pathMap)
	return
}

func main() {
	var listen = flag.String("listen", ":8080",
		"Interface/port to listen on. eg. :8080 or 127.0.0.1:8080")
	var pathsRaw = flag.String("paths", `{"/": "."}`,
		"Paths to serve. A json object with the keys as the url pattern, and "+
			"the value as the root. Default serves current folder.")
	flag.Parse()

	paths, err := parsePaths(*pathsRaw)

	if err != nil {
		log.Fatal(err)
	}

	mux := http.NewServeMux()
	registerHandlers(mux, paths)
	log.Println("Listening on: ", *listen)

	log.Fatal(http.ListenAndServe(*listen, mux))
}
<|endoftext|>


<commit_before>package goemail

import "time"

var pools []*Pool

func SetConfig(cs ...*Config) {
	if cs == nil || len(cs) == 0 {
		panic(ERR_CONFIG)
	}
	pools = make([]*Pool, 0, len(cs))
	for _, c := range cs {
		if c == nil || !c.Valid() {
			continue
		}
		pools = append(pools, NewPool(c))
	}
}

func StartService() error {
	if pools == nil || len(pools) == 0 {
		return ERR_POOL
	}
	for _, p := range pools {
		if p != nil {
			p.Start()
		}
	}
	return nil
}

func StopService() error {
	if pools == nil || len(pools) == 0 {
		return ERR_POOL
	}
	for _, p := range pools {
		if p != nil {
			p.Stop()
		}
	}
	return nil
}

var sendIdx int = 0

func SendEmail(e *Email) error {
	if e == nil || !e.Valid() {
		return ERR_EMAIL
	} else if pools == nil || len(pools) == 0 {
		return ERR_POOL
	}
	if sendIdx > 0 && sendIdx >= len(pools) {
		sendIdx = 0
	}
	if p := pools[sendIdx]; p != nil {
		sendIdx++
		if e.TryCount >= 5 {
			return ERR_EMAIL_TOO_MUCH
		} else if err := p.Send(e.SetFrom(p.config.User).Increase()); err != nil {
			return SendEmail(e)
		}
	}
	return nil
}

func SendEmailDelay(e *Email) {
	go func() {
		time.Sleep(time.Second * 3)
		SendEmail(e)
	}()
}
<commit_msg>fix: some bugs<commit_after>package goemail

import (
	"strconv"
	"time"
)

var pools []*Pool

func SetConfig(cs ...*Config) {
	if cs == nil || len(cs) == 0 {
		panic(ERR_CONFIG)
	}
	pools = make([]*Pool, 0, len(cs))
	for _, c := range cs {
		if c == nil || !c.Valid() {
			continue
		}
		pools = append(pools, NewPool(c))
	}
}

func StartService() error {
	if pools == nil || len(pools) == 0 {
		return ERR_POOL
	}
	for _, p := range pools {
		if p != nil {
			p.Start()
		}
	}
	return nil
}

func StopService() error {
	if pools == nil || len(pools) == 0 {
		return ERR_POOL
	}
	for _, p := range pools {
		if p != nil {
			p.Stop()
		}
	}
	return nil
}

var sendIdx int = 0

func SendEmail(e *Email) error {
	if e == nil || !e.Valid() {
		return ERR_EMAIL
	} else if pools == nil || len(pools) == 0 {
		return ERR_POOL
	}
	if sendIdx > 0 && sendIdx >= len(pools) {
		sendIdx = 0
	}
	if p := pools[sendIdx]; p != nil {
		sendIdx++
		if e.TryCount >= 5 {
			return ERR_EMAIL_TOO_MUCH
		} else if err := p.Send(e.
			AddTag(strconv.Itoa(sendIdx)).
			SetFrom(p.config.User).
			Increase()); err != nil {
			return SendEmail(e)
		}
	}
	return nil
}

func SendEmailDelay(e *Email) {
	go func() {
		time.Sleep(time.Second * 3)
		SendEmail(e)
	}()
}
<|endoftext|>


<commit_before>// Copyright 2017 HyperHQ Inc.
//
// SPDX-License-Identifier: Apache-2.0
//

package main

import (
	"errors"
	"io"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"github.com/moby/moby/pkg/term"
	context "golang.org/x/net/context"

	pb "github.com/kata-containers/agent/protocols/grpc"
)

const sigChanSize = 2048

var sigIgnored = map[syscall.Signal]bool{
	syscall.SIGCHLD:  true,
	syscall.SIGPIPE:  true,
	syscall.SIGWINCH: true,
}

type shim struct {
	containerID string
	execID      string

	ctx   context.Context
	agent *shimAgent
}

func newShim(addr, containerID, execID string) (*shim, error) {
	agent, err := newShimAgent(addr)
	if err != nil {
		return nil, err
	}

	return &shim{containerID: containerID,
		execID: execID,
		ctx:    context.Background(),
		agent:  agent}, nil
}

func (s *shim) proxyStdio(wg *sync.WaitGroup, terminal bool) {
	// don't wait the copying of the stdin, because `io.Copy(inPipe, os.Stdin)`
	// can't terminate when no input. todo: find a better way.
	wg.Add(1)
	if !terminal {
		// In case it's not a terminal, we also need to get the output
		// from stderr.
		wg.Add(1)
	}

	inPipe, outPipe, errPipe := shimStdioPipe(s.ctx, s.agent, s.containerID, s.execID)
	go func() {
		_, err1 := io.Copy(inPipe, os.Stdin)
		_, err2 := s.agent.CloseStdin(s.ctx, &pb.CloseStdinRequest{
			ContainerId: s.containerID,
			ExecId:      s.execID})
		if err1 != nil {
			logger().WithError(err1).Warn("copy stdin failed")
		}
		if err2 != nil {
			logger().WithError(err2).Warn("close stdin failed")
		}
	}()

	go func() {
		_, err := io.Copy(os.Stdout, outPipe)
		logger().WithError(err).Info("copy stdout failed")
		wg.Done()
	}()

	if !terminal {
		go func() {
			_, err := io.Copy(os.Stderr, errPipe)
			logger().WithError(err).Info("copy stderr failed")
			wg.Done()
		}()
	}
}

func (s *shim) forwardAllSignals() chan os.Signal {
	sigc := make(chan os.Signal, sigChanSize)
	// handle all signals for the process.
	signal.Notify(sigc)
	signal.Ignore(syscall.SIGCHLD, syscall.SIGPIPE)

	go func() {
		for sig := range sigc {
			sysSig, ok := sig.(syscall.Signal)
			if !ok {
				err := errors.New("unknown signal")
				logger().WithError(err).WithField("signal", sig.String()).Error("")
				continue
			}

			if fatalSignal(sysSig) {
				logger().WithField("signal", sig).Error("received fatal signal")
				die()
			}

			if sigIgnored[sysSig] {
				//ignore these
				continue
			}
			// forward this signal to container
			_, err := s.agent.SignalProcess(s.ctx, &pb.SignalProcessRequest{
				ContainerId: s.containerID,
				ExecId:      s.execID,
				Signal:      uint32(sysSig)})
			if err != nil {
				logger().WithError(err).WithField("signal", sig.String()).Error("forward signal failed")
			}
		}
	}()
	return sigc
}

func (s *shim) resizeTty(fromTty *os.File) error {
	ws, err := term.GetWinsize(fromTty.Fd())
	if err != nil {
		logger().WithError(err).Info("Error getting size")
		return nil
	}

	_, err = s.agent.TtyWinResize(s.ctx, &pb.TtyWinResizeRequest{
		ContainerId: s.containerID,
		ExecId:      s.execID,
		Row:         uint32(ws.Height),
		Column:      uint32(ws.Width)})
	if err != nil {
		logger().WithError(err).Error("set winsize failed")
	}

	return err
}

func (s *shim) monitorTtySize(tty *os.File) {
	s.resizeTty(tty)
	sigchan := make(chan os.Signal, 1)
	signal.Notify(sigchan, syscall.SIGWINCH)
	go func() {
		for range sigchan {
			s.resizeTty(tty)
		}
	}()
}

func (s *shim) wait() (int32, error) {
	resp, err := s.agent.WaitProcess(s.ctx, &pb.WaitProcessRequest{
		ContainerId: s.containerID,
		ExecId:      s.execID})
	if err != nil {
		return 0, err
	}

	return resp.Status, nil
}
<commit_msg>logging: Improve window resize messages<commit_after>// Copyright 2017 HyperHQ Inc.
//
// SPDX-License-Identifier: Apache-2.0
//

package main

import (
	"errors"
	"io"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"github.com/moby/moby/pkg/term"
	context "golang.org/x/net/context"

	pb "github.com/kata-containers/agent/protocols/grpc"
)

const sigChanSize = 2048

var sigIgnored = map[syscall.Signal]bool{
	syscall.SIGCHLD:  true,
	syscall.SIGPIPE:  true,
	syscall.SIGWINCH: true,
}

type shim struct {
	containerID string
	execID      string

	ctx   context.Context
	agent *shimAgent
}

func newShim(addr, containerID, execID string) (*shim, error) {
	agent, err := newShimAgent(addr)
	if err != nil {
		return nil, err
	}

	return &shim{containerID: containerID,
		execID: execID,
		ctx:    context.Background(),
		agent:  agent}, nil
}

func (s *shim) proxyStdio(wg *sync.WaitGroup, terminal bool) {
	// don't wait the copying of the stdin, because `io.Copy(inPipe, os.Stdin)`
	// can't terminate when no input. todo: find a better way.
	wg.Add(1)
	if !terminal {
		// In case it's not a terminal, we also need to get the output
		// from stderr.
		wg.Add(1)
	}

	inPipe, outPipe, errPipe := shimStdioPipe(s.ctx, s.agent, s.containerID, s.execID)
	go func() {
		_, err1 := io.Copy(inPipe, os.Stdin)
		_, err2 := s.agent.CloseStdin(s.ctx, &pb.CloseStdinRequest{
			ContainerId: s.containerID,
			ExecId:      s.execID})
		if err1 != nil {
			logger().WithError(err1).Warn("copy stdin failed")
		}
		if err2 != nil {
			logger().WithError(err2).Warn("close stdin failed")
		}
	}()

	go func() {
		_, err := io.Copy(os.Stdout, outPipe)
		logger().WithError(err).Info("copy stdout failed")
		wg.Done()
	}()

	if !terminal {
		go func() {
			_, err := io.Copy(os.Stderr, errPipe)
			logger().WithError(err).Info("copy stderr failed")
			wg.Done()
		}()
	}
}

func (s *shim) forwardAllSignals() chan os.Signal {
	sigc := make(chan os.Signal, sigChanSize)
	// handle all signals for the process.
	signal.Notify(sigc)
	signal.Ignore(syscall.SIGCHLD, syscall.SIGPIPE)

	go func() {
		for sig := range sigc {
			sysSig, ok := sig.(syscall.Signal)
			if !ok {
				err := errors.New("unknown signal")
				logger().WithError(err).WithField("signal", sig.String()).Error("")
				continue
			}

			if fatalSignal(sysSig) {
				logger().WithField("signal", sig).Error("received fatal signal")
				die()
			}

			if sigIgnored[sysSig] {
				//ignore these
				continue
			}
			// forward this signal to container
			_, err := s.agent.SignalProcess(s.ctx, &pb.SignalProcessRequest{
				ContainerId: s.containerID,
				ExecId:      s.execID,
				Signal:      uint32(sysSig)})
			if err != nil {
				logger().WithError(err).WithField("signal", sig.String()).Error("forward signal failed")
			}
		}
	}()
	return sigc
}

func (s *shim) resizeTty(fromTty *os.File) error {
	ws, err := term.GetWinsize(fromTty.Fd())
	if err != nil {
		logger().WithError(err).Info("Error getting window size")
		return nil
	}

	_, err = s.agent.TtyWinResize(s.ctx, &pb.TtyWinResizeRequest{
		ContainerId: s.containerID,
		ExecId:      s.execID,
		Row:         uint32(ws.Height),
		Column:      uint32(ws.Width)})
	if err != nil {
		logger().WithError(err).Error("set window size failed")
	}

	return err
}

func (s *shim) monitorTtySize(tty *os.File) {
	s.resizeTty(tty)
	sigchan := make(chan os.Signal, 1)
	signal.Notify(sigchan, syscall.SIGWINCH)
	go func() {
		for range sigchan {
			s.resizeTty(tty)
		}
	}()
}

func (s *shim) wait() (int32, error) {
	resp, err := s.agent.WaitProcess(s.ctx, &pb.WaitProcessRequest{
		ContainerId: s.containerID,
		ExecId:      s.execID})
	if err != nil {
		return 0, err
	}

	return resp.Status, nil
}
<|endoftext|>


<commit_before>package vmwarevsphere

import (
	"fmt"
	"path/filepath"
	"strings"

	"github.com/rancher/machine/libmachine/log"
	"github.com/vmware/govmomi/find"
	"github.com/vmware/govmomi/object"
	"github.com/vmware/govmomi/vapi/library"
	vapifinder "github.com/vmware/govmomi/vapi/library/finder"
	"github.com/vmware/govmomi/vapi/vcenter"
	"github.com/vmware/govmomi/vim25/soap"
	"github.com/vmware/govmomi/vim25/types"
)

func (d *Driver) preCreate() error {
	c, err := d.getSoapClient()
	if err != nil {
		return err
	}

	d.finder = find.NewFinder(c.Client, true)

	d.datacenter, err = d.finder.DatacenterOrDefault(d.getCtx(), d.Datacenter)
	if err != nil {
		return err
	}

	d.finder.SetDatacenter(d.datacenter)
	for _, netName := range d.Networks {
		net, err := d.finder.NetworkOrDefault(d.getCtx(), netName)
		if err != nil {
			return err
		}
		d.networks[netName] = net
	}

	if d.HostSystem != "" {
		var err error
		d.hostsystem, err = d.finder.HostSystemOrDefault(d.getCtx(), d.HostSystem)
		if err != nil {
			return err
		}
	}

	if d.Pool != "" {
		// Find specified Resource Pool
		d.resourcepool, err = d.finder.ResourcePool(d.getCtx(), d.Pool)
		if err != nil {
			return err
		}
	} else if d.HostSystem != "" {
		// Pick default Resource Pool for Host System
		d.resourcepool, err = d.hostsystem.ResourcePool(d.getCtx())
		if err != nil {
			return err
		}
	} else {
		// Pick the default Resource Pool for the Datacenter.
		d.resourcepool, err = d.finder.DefaultResourcePool(d.getCtx())
		if err != nil {
			return err
		}
	}

	return nil
}

func (d *Driver) postCreate(vm *object.VirtualMachine) error {
	if err := d.addConfigParams(vm); err != nil {
		return err
	}

	if err := d.cloudInit(vm); err != nil {
		return err
	}

	if err := d.addTags(vm); err != nil {
		return err
	}

	if err := d.addCustomAttributes(vm); err != nil {
		return err
	}

	return d.Start()
}

func (d *Driver) createLegacy() error {
	c, err := d.getSoapClient()
	if err != nil {
		return err
	}

	spec := types.VirtualMachineConfigSpec{
		Name:     d.MachineName,
		GuestId:  "otherLinux64Guest",
		NumCPUs:  int32(d.CPU),
		MemoryMB: int64(d.Memory),
	}

	scsi, err := object.SCSIControllerTypes().CreateSCSIController("pvscsi")
	if err != nil {
		return err
	}

	spec.DeviceChange = append(spec.DeviceChange, &types.VirtualDeviceConfigSpec{
		Operation: types.VirtualDeviceConfigSpecOperationAdd,
		Device:    scsi,
	})

	if d.VAppTransport == "com.vmware.guestInfo" ||
		d.VAppTransport == "iso" {

		vApp := types.VmConfigSpec{
			OvfEnvironmentTransport: []string{d.VAppTransport},
		}

		if d.VAppIpAllocationPolicy == "dhcp" ||
			d.VAppIpAllocationPolicy == "fixed" ||
			d.VAppIpAllocationPolicy == "transient" ||
			d.VAppIpAllocationPolicy == "fixedAllocated" {

			if d.VAppIpProtocol != "IPv4" &&
				d.VAppIpProtocol != "IPv6" {
				d.VAppIpProtocol = "IPv4"
			}

			supportedAllocationScheme := "ovfenv"
			if d.VAppIpAllocationPolicy == "dhcp" {
				supportedAllocationScheme = "dhcp"
			}

			vApp.IpAssignment = &types.VAppIPAssignmentInfo{
				SupportedIpProtocol:       []string{d.VAppIpProtocol},
				SupportedAllocationScheme: []string{supportedAllocationScheme},
				IpProtocol:                d.VAppIpProtocol,
				IpAllocationPolicy:        d.VAppIpAllocationPolicy + "Policy",
			}
		}

		for i, prop := range d.VAppProperties {
			v := strings.SplitN(prop, "=", 2)
			key := v[0]
			typ := "string"
			value := ""
			if len(v) > 1 {
				value = v[1]
			}
			if strings.HasPrefix(value, "ip:") {
				typ = value
				value = ""
			} else if strings.HasPrefix(value, "${") &&
				strings.HasSuffix(value, "}") {
				typ = "expression"
			}
			vApp.Property = append(vApp.Property, types.VAppPropertySpec{
				ArrayUpdateSpec: types.ArrayUpdateSpec{
					Operation: types.ArrayUpdateOperationAdd,
				},
				Info: &types.VAppPropertyInfo{
					Key:          int32(i),
					Id:           key,
					Type:         typ,
					DefaultValue: value,
				},
			})
		}

		spec.VAppConfig = &vApp
	}

	folder, err := d.findFolder()
	if err != nil {
		return err
	}

	ds, err := d.getDatastore(&spec)
	if err != nil {
		return err
	}

	spec.Files = &types.VirtualMachineFileInfo{
		VmPathName: fmt.Sprintf("[%s]", ds.Name()),
	}

	task, err := folder.CreateVM(d.getCtx(), spec, d.resourcepool, d.hostsystem)
	if err != nil {
		return err
	}

	info, err := task.WaitForResult(d.getCtx(), nil)
	if err != nil {
		return err
	}

	log.Infof("Uploading Boot2docker ISO ...")
	vm := object.NewVirtualMachine(c.Client, info.Result.(types.ManagedObjectReference))
	vmPath, err := d.getVmFolder(vm)
	if err != nil {
		return err
	}

	dsurl := ds.NewURL(filepath.Join(vmPath, isoFilename))
	p := soap.DefaultUpload
	if err = c.Client.UploadFile(d.getCtx(), d.ISO, dsurl, &p); err != nil {
		return err
	}

	devices, err := vm.Device(d.getCtx())
	if err != nil {
		return err
	}

	var add []types.BaseVirtualDevice
	controller, err := devices.FindDiskController("scsi")
	if err != nil {
		return err
	}

	disk := devices.CreateDisk(controller, ds.Reference(),
		ds.Path(fmt.Sprintf("%s/%s.vmdk", vmPath, d.MachineName)))

	// Convert MB to KB
	disk.CapacityInKB = int64(d.DiskSize) * 1024
	add = append(add, disk)
	ide, err := devices.FindIDEController("")
	if err != nil {
		return err
	}

	cdrom, err := devices.CreateCdrom(ide)
	if err != nil {
		return err
	}

	add = append(add, devices.InsertIso(cdrom, ds.Path(fmt.Sprintf("%s/%s", vmPath, isoFilename))))
	if err := vm.AddDevice(d.getCtx(), add...); err != nil {
		return err
	}

	if err := d.addNetworks(vm, d.networks); err != nil {
		return err
	}

	err = d.postCreate(vm)
	if err != nil {
		return err
	}

	if err := d.provisionVm(vm); err != nil {
		return err
	}

	return nil
}

func (d *Driver) createFromVmName() error {
	c, err := d.getSoapClient()
	if err != nil {
		return err
	}

	var info *types.TaskInfo
	poolref := d.resourcepool.Reference()
	hostref := d.hostsystem.Reference()

	spec := types.VirtualMachineCloneSpec{
		Location: types.VirtualMachineRelocateSpec{
			Pool: &poolref,
			Host: &hostref,
		},
		Config: &types.VirtualMachineConfigSpec{
			GuestId:  "otherLinux64Guest",
			NumCPUs:  int32(d.CPU),
			MemoryMB: int64(d.Memory),
		},
	}

	ds, err := d.getDatastore(spec.Config)
	if err != nil {
		return err
	}

	spec.Config.Files = &types.VirtualMachineFileInfo{
		VmPathName: fmt.Sprintf("[%s]", ds.Name()),
	}

	dsref := ds.Reference()
	spec.Location.Datastore = &dsref

	vm2Clone, err := d.fetchVM(d.CloneFrom)
	if err != nil {
		return err
	}

	folder, err := d.findFolder()
	if err != nil {
		return err
	}

	task, err := vm2Clone.Clone(d.getCtx(), folder, d.MachineName, spec)
	if err != nil {
		return err
	}

	info, err = task.WaitForResult(d.getCtx(), nil)
	if err != nil {
		return err
	}

	// Retrieve the new VM
	vm := object.NewVirtualMachine(c.Client, info.Result.(types.ManagedObjectReference))
	if err := d.addNetworks(vm, d.networks); err != nil {
		return err
	}

	if err := d.resizeDisk(vm); err != nil {
		return err
	}

	return d.postCreate(vm)
}

func (d *Driver) createFromLibraryName() error {
	c, err := d.getSoapClient()
	if err != nil {
		return err
	}

	folders, err := d.datacenter.Folders(d.getCtx())
	if err != nil {
		return err
	}

	libManager := library.NewManager(d.getRestLogin(c.Client))
	if err := libManager.Login(d.getCtx(), d.getUserInfo()); err != nil {
		return err
	}

	query := fmt.Sprintf("/%s/%s", d.ContentLibrary, d.CloneFrom)
	results, err := vapifinder.NewFinder(libManager).Find(d.getCtx(), query)
	if err != nil {
		return err
	}

	if len(results) < 1 {
		return fmt.Errorf("No results found in content library: %s", d.CloneFrom)
	}

	if len(results) > 1 {
		return fmt.Errorf("More than one result returned from finder query: %s", d.CloneFrom)
	}

	item, ok := results[0].GetResult().(library.Item)
	if !ok {
		return fmt.Errorf("Content Library item is not a template: %q is a %T", d.CloneFrom, item)
	}

	var nets []vcenter.NetworkMapping
	for k, n := range d.networks {
		nets = append(nets, vcenter.NetworkMapping{
			Key:   k,
			Value: n.Reference().Value,
		})
	}

	hostId := ""
	if d.hostsystem != nil {
		hostId = d.hostsystem.Reference().Value
	}

	ds, err := d.getDatastore(&types.VirtualMachineConfigSpec{})
	if err != nil {
		return err
	}

	deploy := vcenter.Deploy{
		DeploymentSpec: vcenter.DeploymentSpec{
			Name:                d.MachineName,
			DefaultDatastoreID:  ds.Reference().Value,
			AcceptAllEULA:       true,
			NetworkMappings:     nets,
			StorageProvisioning: "thin",
		},
		Target: vcenter.Target{
			ResourcePoolID: d.resourcepool.Reference().Value,
			HostID:         hostId,
			FolderID:       folders.VmFolder.Reference().Value,
		},
	}

	m := vcenter.NewManager(libManager.Client)

	ref, err := m.DeployLibraryItem(d.getCtx(), item.ID, deploy)
	if err != nil {
		return err
	}

	obj, err := d.finder.ObjectReference(d.getCtx(), *ref)
	if err != nil {
		return err
	}

	vm := obj.(*object.VirtualMachine)
	if err := d.resizeDisk(vm); err != nil {
		return err
	}

	return d.postCreate(vm)
}
<commit_msg>safer pointer checks on host/pool<commit_after>package vmwarevsphere

import (
	"fmt"
	"path/filepath"
	"strings"

	"github.com/rancher/machine/libmachine/log"
	"github.com/vmware/govmomi/find"
	"github.com/vmware/govmomi/object"
	"github.com/vmware/govmomi/vapi/library"
	vapifinder "github.com/vmware/govmomi/vapi/library/finder"
	"github.com/vmware/govmomi/vapi/vcenter"
	"github.com/vmware/govmomi/vim25/soap"
	"github.com/vmware/govmomi/vim25/types"
)

func (d *Driver) preCreate() error {
	c, err := d.getSoapClient()
	if err != nil {
		return err
	}

	d.finder = find.NewFinder(c.Client, true)

	d.datacenter, err = d.finder.DatacenterOrDefault(d.getCtx(), d.Datacenter)
	if err != nil {
		return err
	}

	d.finder.SetDatacenter(d.datacenter)
	for _, netName := range d.Networks {
		net, err := d.finder.NetworkOrDefault(d.getCtx(), netName)
		if err != nil {
			return err
		}
		d.networks[netName] = net
	}

	if d.HostSystem != "" {
		var err error
		d.hostsystem, err = d.finder.HostSystemOrDefault(d.getCtx(), d.HostSystem)
		if err != nil {
			return err
		}
	}

	if d.Pool != "" {
		// Find specified Resource Pool
		d.resourcepool, err = d.finder.ResourcePool(d.getCtx(), d.Pool)
		if err != nil {
			return err
		}
	} else if d.HostSystem != "" {
		// Pick default Resource Pool for Host System
		d.resourcepool, err = d.hostsystem.ResourcePool(d.getCtx())
		if err != nil {
			return err
		}
	} else {
		// Pick the default Resource Pool for the Datacenter.
		d.resourcepool, err = d.finder.DefaultResourcePool(d.getCtx())
		if err != nil {
			return err
		}
	}

	return nil
}

func (d *Driver) postCreate(vm *object.VirtualMachine) error {
	if err := d.addConfigParams(vm); err != nil {
		return err
	}

	if err := d.cloudInit(vm); err != nil {
		return err
	}

	if err := d.addTags(vm); err != nil {
		return err
	}

	if err := d.addCustomAttributes(vm); err != nil {
		return err
	}

	return d.Start()
}

func (d *Driver) createLegacy() error {
	c, err := d.getSoapClient()
	if err != nil {
		return err
	}

	spec := types.VirtualMachineConfigSpec{
		Name:     d.MachineName,
		GuestId:  "otherLinux64Guest",
		NumCPUs:  int32(d.CPU),
		MemoryMB: int64(d.Memory),
	}

	scsi, err := object.SCSIControllerTypes().CreateSCSIController("pvscsi")
	if err != nil {
		return err
	}

	spec.DeviceChange = append(spec.DeviceChange, &types.VirtualDeviceConfigSpec{
		Operation: types.VirtualDeviceConfigSpecOperationAdd,
		Device:    scsi,
	})

	if d.VAppTransport == "com.vmware.guestInfo" ||
		d.VAppTransport == "iso" {

		vApp := types.VmConfigSpec{
			OvfEnvironmentTransport: []string{d.VAppTransport},
		}

		if d.VAppIpAllocationPolicy == "dhcp" ||
			d.VAppIpAllocationPolicy == "fixed" ||
			d.VAppIpAllocationPolicy == "transient" ||
			d.VAppIpAllocationPolicy == "fixedAllocated" {

			if d.VAppIpProtocol != "IPv4" &&
				d.VAppIpProtocol != "IPv6" {
				d.VAppIpProtocol = "IPv4"
			}

			supportedAllocationScheme := "ovfenv"
			if d.VAppIpAllocationPolicy == "dhcp" {
				supportedAllocationScheme = "dhcp"
			}

			vApp.IpAssignment = &types.VAppIPAssignmentInfo{
				SupportedIpProtocol:       []string{d.VAppIpProtocol},
				SupportedAllocationScheme: []string{supportedAllocationScheme},
				IpProtocol:                d.VAppIpProtocol,
				IpAllocationPolicy:        d.VAppIpAllocationPolicy + "Policy",
			}
		}

		for i, prop := range d.VAppProperties {
			v := strings.SplitN(prop, "=", 2)
			key := v[0]
			typ := "string"
			value := ""
			if len(v) > 1 {
				value = v[1]
			}
			if strings.HasPrefix(value, "ip:") {
				typ = value
				value = ""
			} else if strings.HasPrefix(value, "${") &&
				strings.HasSuffix(value, "}") {
				typ = "expression"
			}
			vApp.Property = append(vApp.Property, types.VAppPropertySpec{
				ArrayUpdateSpec: types.ArrayUpdateSpec{
					Operation: types.ArrayUpdateOperationAdd,
				},
				Info: &types.VAppPropertyInfo{
					Key:          int32(i),
					Id:           key,
					Type:         typ,
					DefaultValue: value,
				},
			})
		}

		spec.VAppConfig = &vApp
	}

	folder, err := d.findFolder()
	if err != nil {
		return err
	}

	ds, err := d.getDatastore(&spec)
	if err != nil {
		return err
	}

	spec.Files = &types.VirtualMachineFileInfo{
		VmPathName: fmt.Sprintf("[%s]", ds.Name()),
	}

	task, err := folder.CreateVM(d.getCtx(), spec, d.resourcepool, d.hostsystem)
	if err != nil {
		return err
	}

	info, err := task.WaitForResult(d.getCtx(), nil)
	if err != nil {
		return err
	}

	log.Infof("Uploading Boot2docker ISO ...")
	vm := object.NewVirtualMachine(c.Client, info.Result.(types.ManagedObjectReference))
	vmPath, err := d.getVmFolder(vm)
	if err != nil {
		return err
	}

	dsurl := ds.NewURL(filepath.Join(vmPath, isoFilename))
	p := soap.DefaultUpload
	if err = c.Client.UploadFile(d.getCtx(), d.ISO, dsurl, &p); err != nil {
		return err
	}

	devices, err := vm.Device(d.getCtx())
	if err != nil {
		return err
	}

	var add []types.BaseVirtualDevice
	controller, err := devices.FindDiskController("scsi")
	if err != nil {
		return err
	}

	disk := devices.CreateDisk(controller, ds.Reference(),
		ds.Path(fmt.Sprintf("%s/%s.vmdk", vmPath, d.MachineName)))

	// Convert MB to KB
	disk.CapacityInKB = int64(d.DiskSize) * 1024
	add = append(add, disk)
	ide, err := devices.FindIDEController("")
	if err != nil {
		return err
	}

	cdrom, err := devices.CreateCdrom(ide)
	if err != nil {
		return err
	}

	add = append(add, devices.InsertIso(cdrom, ds.Path(fmt.Sprintf("%s/%s", vmPath, isoFilename))))
	if err := vm.AddDevice(d.getCtx(), add...); err != nil {
		return err
	}

	if err := d.addNetworks(vm, d.networks); err != nil {
		return err
	}

	err = d.postCreate(vm)
	if err != nil {
		return err
	}

	if err := d.provisionVm(vm); err != nil {
		return err
	}

	return nil
}

func (d *Driver) createFromVmName() error {
	c, err := d.getSoapClient()
	if err != nil {
		return err
	}

	var info *types.TaskInfo
	var loc types.VirtualMachineRelocateSpec

	if d.resourcepool != nil {
		pool := d.resourcepool.Reference()
		loc.Pool = &pool
	}

	if d.hostsystem != nil {
		host := d.hostsystem.Reference()
		loc.Host = &host
	}

	spec := types.VirtualMachineCloneSpec{
		Location: loc,
		Config: &types.VirtualMachineConfigSpec{
			GuestId:  "otherLinux64Guest",
			NumCPUs:  int32(d.CPU),
			MemoryMB: int64(d.Memory),
		},
	}

	ds, err := d.getDatastore(spec.Config)
	if err != nil {
		return err
	}

	spec.Config.Files = &types.VirtualMachineFileInfo{
		VmPathName: fmt.Sprintf("[%s]", ds.Name()),
	}

	dsref := ds.Reference()
	spec.Location.Datastore = &dsref

	vm2Clone, err := d.fetchVM(d.CloneFrom)
	if err != nil {
		return err
	}

	folder, err := d.findFolder()
	if err != nil {
		return err
	}

	task, err := vm2Clone.Clone(d.getCtx(), folder, d.MachineName, spec)
	if err != nil {
		return err
	}

	info, err = task.WaitForResult(d.getCtx(), nil)
	if err != nil {
		return err
	}

	// Retrieve the new VM
	vm := object.NewVirtualMachine(c.Client, info.Result.(types.ManagedObjectReference))
	if err := d.addNetworks(vm, d.networks); err != nil {
		return err
	}

	if err := d.resizeDisk(vm); err != nil {
		return err
	}

	return d.postCreate(vm)
}

func (d *Driver) createFromLibraryName() error {
	c, err := d.getSoapClient()
	if err != nil {
		return err
	}

	folders, err := d.datacenter.Folders(d.getCtx())
	if err != nil {
		return err
	}

	libManager := library.NewManager(d.getRestLogin(c.Client))
	if err := libManager.Login(d.getCtx(), d.getUserInfo()); err != nil {
		return err
	}

	query := fmt.Sprintf("/%s/%s", d.ContentLibrary, d.CloneFrom)
	results, err := vapifinder.NewFinder(libManager).Find(d.getCtx(), query)
	if err != nil {
		return err
	}

	if len(results) < 1 {
		return fmt.Errorf("No results found in content library: %s", d.CloneFrom)
	}

	if len(results) > 1 {
		return fmt.Errorf("More than one result returned from finder query: %s", d.CloneFrom)
	}

	item, ok := results[0].GetResult().(library.Item)
	if !ok {
		return fmt.Errorf("Content Library item is not a template: %q is a %T", d.CloneFrom, item)
	}

	var nets []vcenter.NetworkMapping
	for k, n := range d.networks {
		nets = append(nets, vcenter.NetworkMapping{
			Key:   k,
			Value: n.Reference().Value,
		})
	}

	hostId := ""
	if d.hostsystem != nil {
		hostId = d.hostsystem.Reference().Value
	}

	ds, err := d.getDatastore(&types.VirtualMachineConfigSpec{})
	if err != nil {
		return err
	}

	deploy := vcenter.Deploy{
		DeploymentSpec: vcenter.DeploymentSpec{
			Name:                d.MachineName,
			DefaultDatastoreID:  ds.Reference().Value,
			AcceptAllEULA:       true,
			NetworkMappings:     nets,
			StorageProvisioning: "thin",
		},
		Target: vcenter.Target{
			ResourcePoolID: d.resourcepool.Reference().Value,
			HostID:         hostId,
			FolderID:       folders.VmFolder.Reference().Value,
		},
	}

	m := vcenter.NewManager(libManager.Client)

	ref, err := m.DeployLibraryItem(d.getCtx(), item.ID, deploy)
	if err != nil {
		return err
	}

	obj, err := d.finder.ObjectReference(d.getCtx(), *ref)
	if err != nil {
		return err
	}

	vm := obj.(*object.VirtualMachine)
	if err := d.resizeDisk(vm); err != nil {
		return err
	}

	return d.postCreate(vm)
}
<|endoftext|>


<commit_before>// Package gcfg reads "gitconfig-like" text-based configuration files with
// "name=value" pairs grouped into sections (gcfg files).
// Support for writing gcfg files may be added later.
//
// See ReadInto and the examples to get an idea of how to use it.
//
// This package is still a work in progress, and both the supported syntax and
// the API is subject to change.
//
// The syntax is based on that used by git config:
// http://git-scm.com/docs/git-config#_syntax .
// There are some (planned) differences compared to the git config format:
//  - improve data portability:
//    - must be encoded in UTF-8 (for now) and must not contain the 0 byte
//    - include and "path" type is not supported
//      (path type may be implementable as a user-defined type)
//  - disallow potentially ambiguous or misleading definitions:
//    - `[sec.sub]` format is not allowed (deprecated in gitconfig)
//    - `[sec ""]` is not allowed
//      - use `[sec]` for section name "sec" and empty subsection name
//    - (planned) within a single file, definitions must be contiguous for each:
//      - section: '[secA]' -> '[secB]' -> '[secA]' is an error
//      - subsection: '[sec "A"]' -> '[sec "B"]' -> '[sec "A"]' is an error
//      - multivalued variable: 'multi=a' -> 'other=x' -> 'multi=b' is an error
//
// The package may be usable for handling some of the various "INI file" formats
// used by some programs and libraries, but achieving or maintaining
// compatibility with any of those is not a primary concern.
//
// TODO:
//  - format
//    - define valid section and variable names
//    - reconsider valid escape sequences
//      (gitconfig doesn't support \r in value, \t in subsection name, etc.)
//    - define handling of default value
//    - complete syntax documentation
//  - reading
//    - define internal representation structure
//    - support multi-value variables
//    - support multiple inputs (readers, strings, files)
//    - support declaring encoding (?)
//    - support pointer fields
//    - support varying fields sets for subsections (?)
//  - scanEnum
//    - should use longest match (?)
//    - support matching on unique prefix (?)
//  - writing gcfg files
//  - error handling
//    - make error context accessible programmatically?
//    - limit input size?
//  - move TODOs to issue tracker (eventually)
//
package gcfg
<commit_msg>add and clarify TODOs<commit_after>// Package gcfg reads "gitconfig-like" text-based configuration files with
// "name=value" pairs grouped into sections (gcfg files).
// Support for writing gcfg files may be added later.
//
// See ReadInto and the examples to get an idea of how to use it.
//
// This package is still a work in progress, and both the supported syntax and
// the API is subject to change.
//
// The syntax is based on that used by git config:
// http://git-scm.com/docs/git-config#_syntax .
// There are some (planned) differences compared to the git config format:
//  - improve data portability:
//    - must be encoded in UTF-8 (for now) and must not contain the 0 byte
//    - include and "path" type is not supported
//      (path type may be implementable as a user-defined type)
//  - disallow potentially ambiguous or misleading definitions:
//    - `[sec.sub]` format is not allowed (deprecated in gitconfig)
//    - `[sec ""]` is not allowed
//      - use `[sec]` for section name "sec" and empty subsection name
//    - (planned) within a single file, definitions must be contiguous for each:
//      - section: '[secA]' -> '[secB]' -> '[secA]' is an error
//      - subsection: '[sec "A"]' -> '[sec "B"]' -> '[sec "A"]' is an error
//      - multivalued variable: 'multi=a' -> 'other=x' -> 'multi=b' is an error
//
// The package may be usable for handling some of the various "INI file" formats
// used by some programs and libraries, but achieving or maintaining
// compatibility with any of those is not a primary concern.
//
// TODO:
//  - format
//    - define valid section and variable names
//    - reconsider valid escape sequences
//      (gitconfig doesn't support \r in value, \t in subsection name, etc.)
//    - define handling of "implicit value" for types other than bool
//    - consider handling of numeric values (decimal only by default?)
//    - complete syntax documentation
//  - reading
//    - define internal representation structure
//    - support multi-value variables
//    - support multiple inputs (readers, strings, files)
//    - support declaring encoding (?)
//    - support automatic derefereining of pointer fields (?)
//    - support varying fields sets for subsections (?)
//  - scanEnum
//    - should use longest match (?)
//    - support matching on unique prefix (?)
//  - writing gcfg files
//  - error handling
//    - make error context accessible programmatically?
//    - limit input size?
//  - move TODOs to issue tracker (eventually)
//
package gcfg
<|endoftext|>