<commit_before>package main

import "testing"
import "runtime"
import "strconv"
import "time"

func TestMessageInsertAndRetreive(t *testing.T) {
	id := "1"
	say := "'Ello, Mister Polly Parrot!"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	store.Add <- &Message{at, id, say}
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{zero_time, messages_from_store}
	messages := <-messages_from_store
	if len(messages) != 1 {
		t.FailNow()
	}
	if messages[0].Time != at {
		t.Fail()
	}
	if messages[0].ID != id {
		t.Fail()
	}
	if messages[0].Text != say {
		t.Fail()
	}
	close(store.Get)
	close(store.Add)
}

func TestFetchBlocksUntilSpeak(t *testing.T) {
	start_fetch_wait_count := fetch_wait_count.String()
	id := "2"
	say := "I've got a lovely fresh cuttle fish for you"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{zero_time, messages_from_store}
	for start_fetch_wait_count == fetch_wait_count.String() {
		runtime.Gosched()
	}
	store.Add <- &Message{at, id, say}
	messages := <-messages_from_store
	if len(messages) != 1 {
		t.FailNow()
	}
	if messages[0].Time != at {
		t.Fail()
	}
	if messages[0].ID != id {
		t.Fail()
	}
	if messages[0].Text != say {
		t.Fail()
	}
	close(store.Get)
	close(store.Add)
}

func TestMultipleListeners(t *testing.T) {
	id := "3"
	say := "This is your nine o'clock alarm call!"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	const num_clients = 13
	var messages_from_store [num_clients]chan []Message
	for i := 0; i < num_clients; i++ {
		messages_from_store[i] = make(chan []Message, 1)
		store.Get <- &StoreRequest{zero_time, messages_from_store[i]}
	}
	store.Add <- &Message{at, id, say}
	for i := 0; i < num_clients; i++ {
		messages := <-messages_from_store[i]
		if len(messages) != 1 {
			t.FailNow()
		}
		if messages[0].Time != at {
			t.Fail()
		}
		if messages[0].ID != id {
			t.Fail()
		}
		if messages[0].Text != say {
			t.Fail()
		}
	}
	close(store.Get)
	close(store.Add)
}

func parseDuration(s string) time.Duration {
	d, err := time.ParseDuration(s)
	if err != nil {
		panic(err)
	}
	return d
}

func atoi(s string) int {
	i, err := strconv.Atoi(s)
	if err != nil {
		panic(err)
	}
	return i
}

func TestPartialRetreive(t *testing.T) {
	start_speak_count := atoi(speak_count.String())
	id1 := "4"
	id2 := "5"
	id3 := "6"
	say1 := "No, no.....No, 'e's stunned!"
	say2 := "You stunned him, just as he was wakin' up!"
	say3 := "Norwegian Blues stun easily, major."
	base := time.Now()
	at1 := base.Add(parseDuration("-4m"))
	since := base.Add(parseDuration("-3m"))
	at2 := base.Add(parseDuration("-2m"))
	at3 := base.Add(parseDuration("-1m"))
	store := start_store()
	store.Add <- &Message{at1, id1, say1}
	store.Add <- &Message{at2, id2, say2}
	store.Add <- &Message{at3, id3, say3}
	for atoi(speak_count.String()) != start_speak_count+3 {
		runtime.Gosched()
	}
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{since, messages_from_store}
	messages := <-messages_from_store
	if len(messages) != 2 {
		t.FailNow()
	}
	if messages[0].Time != at2 {
		t.Fail()
	}
	if messages[0].ID != id2 {
		t.Fail()
	}
	if messages[0].Text != say2 {
		t.Fail()
	}
	if messages[1].Time != at3 {
		t.Fail()
	}
	if messages[1].ID != id3 {
		t.Fail()
	}
	if messages[1].Text != say3 {
		t.Fail()
	}
	close(store.Get)
	close(store.Add)
}

func TestPrecisePartialRetreive(t *testing.T) {
	start_speak_count := atoi(speak_count.String())
	id1 := "7"
	id2 := "8"
	id3 := "9"
	say1 := "Well, he's...he's, ah...probably pining for the fjords."
	say2 := "PININ' for the FJORDS?!?!?!?"
	say3 := "look, why did he fall flat on his back the moment I got 'im home?"
	base := time.Now()
	at1 := base.Add(parseDuration("-3m"))
	at2 := base.Add(parseDuration("-2m"))
	at3 := base.Add(parseDuration("-1m"))
	since := at2
	store := start_store()
	store.Add <- &Message{at1, id1, say1}
	store.Add <- &Message{at2, id2, say2}
	store.Add <- &Message{at3, id3, say3}
	for atoi(speak_count.String()) != start_speak_count+3 {
		runtime.Gosched()
	}
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{since, messages_from_store}
	messages := <-messages_from_store
	if len(messages) != 1 {
		t.FailNow()
	}
	if messages[0].Time != at3 {
		t.Fail()
	}
	if messages[0].ID != id3 {
		t.Fail()
	}
	if messages[0].Text != say3 {
		t.Fail()
	}
	close(store.Get)
	close(store.Add)
}

func TestTypicalFlow(t *testing.T) {
	id1 := "10"
	id2 := "11"
	say1 := "The Norwegian Blue prefers kippin' on it's back!"
	say2 := "Remarkable bird, innit, squire?  Lovely plumage!"
	store := start_store()

	// A waiting zero-time fetch.
	var zero_time time.Time
	prev_fetch_wait_count := fetch_wait_count.String()
	fetch1 := make(chan []Message, 1)
	store.Get <- &StoreRequest{zero_time, fetch1}
	for prev_fetch_wait_count == fetch_wait_count.String() {
		runtime.Gosched()
	}

	// Someone speaks.  This triggers delivery.
	at1 := time.Now()
	store.Add <- &Message{at1, id1, say1}
	messages1 := <-fetch1
	if len(messages1) != 1 {
		t.FailNow()
	}
	if messages1[0].Time != at1 {
		t.Fail()
	}
	if messages1[0].ID != id1 {
		t.Fail()
	}
	if messages1[0].Text != say1 {
		t.Fail()
	}

	// Upon recipt, client blocks on fetch with since=at1
	prev_fetch_wait_count = fetch_wait_count.String()
	fetch2 := make(chan []Message, 1)
	store.Get <- &StoreRequest{at1, fetch2}
	for prev_fetch_wait_count == fetch_wait_count.String() {
		runtime.Gosched()
	}

	// Someone speaks again.  This triggers another delivery.
	at2 := time.Now()
	if !at2.After(at1) {
		t.Fail()
	}
	store.Add <- &Message{at2, id2, say2}
	messages2 := <-fetch2
	if len(messages2) != 1 {
		t.FailNow()
	}
	if messages2[0].Time != at2 {
		t.Fail()
	}
	if messages2[0].ID != id2 {
		t.Fail()
	}
	if messages2[0].Text != say2 {
		t.Fail()
	}

	close(store.Get)
	close(store.Add)
}
<commit_msg>Factor out copy/pasted message testing<commit_after>package main

import "testing"
import "runtime"
import "strconv"
import "time"

func expectMessage(t *testing.T, m *Message, at time.Time, id, say string) {
	if m.Time != at {
		t.Fail()
	}
	if m.ID != id {
		t.Fail()
	}
	if m.Text != say {
		t.Fail()
	}
}

func TestMessageInsertAndRetreive(t *testing.T) {
	id := "1"
	say := "'Ello, Mister Polly Parrot!"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	store.Add <- &Message{at, id, say}
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{zero_time, messages_from_store}
	messages := <-messages_from_store
	if len(messages) != 1 {
		t.FailNow()
	}
	expectMessage(t, &messages[0], at, id, say)
	close(store.Get)
	close(store.Add)
}

func TestFetchBlocksUntilSpeak(t *testing.T) {
	start_fetch_wait_count := fetch_wait_count.String()
	id := "2"
	say := "I've got a lovely fresh cuttle fish for you"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{zero_time, messages_from_store}
	for start_fetch_wait_count == fetch_wait_count.String() {
		runtime.Gosched()
	}
	store.Add <- &Message{at, id, say}
	messages := <-messages_from_store
	if len(messages) != 1 {
		t.FailNow()
	}
	expectMessage(t, &messages[0], at, id, say)
	close(store.Get)
	close(store.Add)
}

func TestMultipleListeners(t *testing.T) {
	id := "3"
	say := "This is your nine o'clock alarm call!"
	at := time.Now()
	var zero_time time.Time
	store := start_store()
	const num_clients = 13
	var messages_from_store [num_clients]chan []Message
	for i := 0; i < num_clients; i++ {
		messages_from_store[i] = make(chan []Message, 1)
		store.Get <- &StoreRequest{zero_time, messages_from_store[i]}
	}
	store.Add <- &Message{at, id, say}
	for i := 0; i < num_clients; i++ {
		messages := <-messages_from_store[i]
		if len(messages) != 1 {
			t.FailNow()
		}
		expectMessage(t,& messages[0], at, id, say)
	}
	close(store.Get)
	close(store.Add)
}

func parseDuration(s string) time.Duration {
	d, err := time.ParseDuration(s)
	if err != nil {
		panic(err)
	}
	return d
}

func atoi(s string) int {
	i, err := strconv.Atoi(s)
	if err != nil {
		panic(err)
	}
	return i
}

func TestPartialRetreive(t *testing.T) {
	start_speak_count := atoi(speak_count.String())
	id1 := "4"
	id2 := "5"
	id3 := "6"
	say1 := "No, no.....No, 'e's stunned!"
	say2 := "You stunned him, just as he was wakin' up!"
	say3 := "Norwegian Blues stun easily, major."
	base := time.Now()
	at1 := base.Add(parseDuration("-4m"))
	since := base.Add(parseDuration("-3m"))
	at2 := base.Add(parseDuration("-2m"))
	at3 := base.Add(parseDuration("-1m"))
	store := start_store()
	store.Add <- &Message{at1, id1, say1}
	store.Add <- &Message{at2, id2, say2}
	store.Add <- &Message{at3, id3, say3}
	for atoi(speak_count.String()) != start_speak_count+3 {
		runtime.Gosched()
	}
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{since, messages_from_store}
	messages := <-messages_from_store
	if len(messages) != 2 {
		t.FailNow()
	}
	expectMessage(t, &messages[0], at2, id2, say2)
	expectMessage(t, &messages[1], at3, id3, say3)
	close(store.Get)
	close(store.Add)
}

func TestPrecisePartialRetreive(t *testing.T) {
	start_speak_count := atoi(speak_count.String())
	id1 := "7"
	id2 := "8"
	id3 := "9"
	say1 := "Well, he's...he's, ah...probably pining for the fjords."
	say2 := "PININ' for the FJORDS?!?!?!?"
	say3 := "look, why did he fall flat on his back the moment I got 'im home?"
	base := time.Now()
	at1 := base.Add(parseDuration("-3m"))
	at2 := base.Add(parseDuration("-2m"))
	at3 := base.Add(parseDuration("-1m"))
	since := at2
	store := start_store()
	store.Add <- &Message{at1, id1, say1}
	store.Add <- &Message{at2, id2, say2}
	store.Add <- &Message{at3, id3, say3}
	for atoi(speak_count.String()) != start_speak_count+3 {
		runtime.Gosched()
	}
	messages_from_store := make(chan []Message, 1)
	store.Get <- &StoreRequest{since, messages_from_store}
	messages := <-messages_from_store
	if len(messages) != 1 {
		t.FailNow()
	}
	expectMessage(t, &messages[0], at3, id3, say3)
	close(store.Get)
	close(store.Add)
}

func TestTypicalFlow(t *testing.T) {
	id1 := "10"
	id2 := "11"
	say1 := "The Norwegian Blue prefers kippin' on it's back!"
	say2 := "Remarkable bird, innit, squire?  Lovely plumage!"
	store := start_store()

	// A waiting zero-time fetch.
	var zero_time time.Time
	prev_fetch_wait_count := fetch_wait_count.String()
	fetch1 := make(chan []Message, 1)
	store.Get <- &StoreRequest{zero_time, fetch1}
	for prev_fetch_wait_count == fetch_wait_count.String() {
		runtime.Gosched()
	}

	// Someone speaks.  This triggers delivery.
	at1 := time.Now()
	store.Add <- &Message{at1, id1, say1}
	messages1 := <-fetch1
	if len(messages1) != 1 {
		t.FailNow()
	}
	expectMessage(t, &messages1[0], at1, id1, say1)

	// Upon recipt, client blocks on fetch with since=at1
	prev_fetch_wait_count = fetch_wait_count.String()
	fetch2 := make(chan []Message, 1)
	store.Get <- &StoreRequest{at1, fetch2}
	for prev_fetch_wait_count == fetch_wait_count.String() {
		runtime.Gosched()
	}

	// Someone speaks again.  This triggers another delivery.
	at2 := time.Now()
	if !at2.After(at1) {
		t.Fail()
	}
	store.Add <- &Message{at2, id2, say2}
	messages2 := <-fetch2
	if len(messages2) != 1 {
		t.FailNow()
	}
	expectMessage(t, &messages2[0], at2, id2, say2)

	close(store.Get)
	close(store.Add)
}
<|endoftext|>


<commit_before>package daemon // import "github.com/docker/docker/daemon"

import (
	"sync"

	"github.com/docker/docker/errdefs"
	"github.com/docker/docker/pkg/plugingetter"
	"github.com/docker/docker/pkg/plugins"
	metrics "github.com/docker/go-metrics"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
)

const metricsPluginType = "MetricsCollector"

var (
	containerActions          metrics.LabeledTimer
	networkActions            metrics.LabeledTimer
	hostInfoFunctions         metrics.LabeledTimer
	engineInfo                metrics.LabeledGauge
	engineCpus                metrics.Gauge
	engineMemory              metrics.Gauge
	healthChecksCounter       metrics.Counter
	healthChecksFailedCounter metrics.Counter

	stateCtr *stateCounter
)

func init() {
	ns := metrics.NewNamespace("engine", "daemon", nil)
	containerActions = ns.NewLabeledTimer("container_actions", "The number of seconds it takes to process each container action", "action")
	for _, a := range []string{
		"start",
		"changes",
		"commit",
		"create",
		"delete",
	} {
		containerActions.WithValues(a).Update(0)
	}
	hostInfoFunctions = ns.NewLabeledTimer("host_info_functions", "The number of seconds it takes to call functions gathering info about the host", "function")

	networkActions = ns.NewLabeledTimer("network_actions", "The number of seconds it takes to process each network action", "action")
	engineInfo = ns.NewLabeledGauge("engine", "The information related to the engine and the OS it is running on", metrics.Unit("info"),
		"version",
		"commit",
		"architecture",
		"graphdriver",
		"kernel",
		"os",
		"os_type",
		"os_version",
		"daemon_id", // ID is a randomly generated unique identifier (e.g. UUID4)
	)
	engineCpus = ns.NewGauge("engine_cpus", "The number of cpus that the host system of the engine has", metrics.Unit("cpus"))
	engineMemory = ns.NewGauge("engine_memory", "The number of bytes of memory that the host system of the engine has", metrics.Bytes)
	healthChecksCounter = ns.NewCounter("health_checks", "The total number of health checks")
	healthChecksFailedCounter = ns.NewCounter("health_checks_failed", "The total number of failed health checks")

	stateCtr = newStateCounter(ns.NewDesc("container_states", "The count of containers in various states", metrics.Unit("containers"), "state"))
	ns.Add(stateCtr)

	metrics.Register(ns)
}

type stateCounter struct {
	mu     sync.RWMutex
	states map[string]string
	desc   *prometheus.Desc
}

func newStateCounter(desc *prometheus.Desc) *stateCounter {
	return &stateCounter{
		states: make(map[string]string),
		desc:   desc,
	}
}

func (ctr *stateCounter) get() (running int, paused int, stopped int) {
	ctr.mu.RLock()
	defer ctr.mu.RUnlock()

	states := map[string]int{
		"running": 0,
		"paused":  0,
		"stopped": 0,
	}
	for _, state := range ctr.states {
		states[state]++
	}
	return states["running"], states["paused"], states["stopped"]
}

func (ctr *stateCounter) set(id, label string) {
	ctr.mu.Lock()
	ctr.states[id] = label
	ctr.mu.Unlock()
}

func (ctr *stateCounter) del(id string) {
	ctr.mu.Lock()
	delete(ctr.states, id)
	ctr.mu.Unlock()
}

func (ctr *stateCounter) Describe(ch chan<- *prometheus.Desc) {
	ch <- ctr.desc
}

func (ctr *stateCounter) Collect(ch chan<- prometheus.Metric) {
	running, paused, stopped := ctr.get()
	ch <- prometheus.MustNewConstMetric(ctr.desc, prometheus.GaugeValue, float64(running), "running")
	ch <- prometheus.MustNewConstMetric(ctr.desc, prometheus.GaugeValue, float64(paused), "paused")
	ch <- prometheus.MustNewConstMetric(ctr.desc, prometheus.GaugeValue, float64(stopped), "stopped")
}

func (daemon *Daemon) cleanupMetricsPlugins() {
	ls := daemon.PluginStore.GetAllManagedPluginsByCap(metricsPluginType)
	var wg sync.WaitGroup
	wg.Add(len(ls))

	for _, plugin := range ls {
		p := plugin
		go func() {
			defer wg.Done()

			adapter, err := makePluginAdapter(p)
			if err != nil {
				logrus.WithError(err).WithField("plugin", p.Name()).Error("Error creating metrics plugin adapter")
				return
			}
			if err := adapter.StopMetrics(); err != nil {
				logrus.WithError(err).WithField("plugin", p.Name()).Error("Error stopping plugin metrics collection")
			}
		}()
	}
	wg.Wait()

	if daemon.metricsPluginListener != nil {
		daemon.metricsPluginListener.Close()
	}
}

type metricsPlugin interface {
	StartMetrics() error
	StopMetrics() error
}

func makePluginAdapter(p plugingetter.CompatPlugin) (metricsPlugin, error) {
	if pc, ok := p.(plugingetter.PluginWithV1Client); ok {
		return &metricsPluginAdapter{pc.Client(), p.Name()}, nil
	}

	pa, ok := p.(plugingetter.PluginAddr)
	if !ok {
		return nil, errdefs.System(errors.Errorf("got unknown plugin type %T", p))
	}

	if pa.Protocol() != plugins.ProtocolSchemeHTTPV1 {
		return nil, errors.Errorf("plugin protocol not supported: %s", pa.Protocol())
	}

	addr := pa.Addr()
	client, err := plugins.NewClientWithTimeout(addr.Network()+"://"+addr.String(), nil, pa.Timeout())
	if err != nil {
		return nil, errors.Wrap(err, "error creating metrics plugin client")
	}
	return &metricsPluginAdapter{client, p.Name()}, nil
}

type metricsPluginAdapter struct {
	c    *plugins.Client
	name string
}

func (a *metricsPluginAdapter) StartMetrics() error {
	type metricsPluginResponse struct {
		Err string
	}
	var res metricsPluginResponse
	if err := a.c.Call(metricsPluginType+".StartMetrics", nil, &res); err != nil {
		return errors.Wrap(err, "could not start metrics plugin")
	}
	if res.Err != "" {
		return errors.New(res.Err)
	}
	return nil
}

func (a *metricsPluginAdapter) StopMetrics() error {
	if err := a.c.Call(metricsPluginType+".StopMetrics", nil, nil); err != nil {
		return errors.Wrap(err, "error stopping metrics collector")
	}
	return nil
}
<commit_msg>metrics: DRY metric definitions<commit_after>package daemon // import "github.com/docker/docker/daemon"

import (
	"sync"

	"github.com/docker/docker/errdefs"
	"github.com/docker/docker/pkg/plugingetter"
	"github.com/docker/docker/pkg/plugins"
	metrics "github.com/docker/go-metrics"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
)

const metricsPluginType = "MetricsCollector"

var (
	metricsNS = metrics.NewNamespace("engine", "daemon", nil)

	containerActions  = metricsNS.NewLabeledTimer("container_actions", "The number of seconds it takes to process each container action", "action")
	networkActions    = metricsNS.NewLabeledTimer("network_actions", "The number of seconds it takes to process each network action", "action")
	hostInfoFunctions = metricsNS.NewLabeledTimer("host_info_functions", "The number of seconds it takes to call functions gathering info about the host", "function")

	engineInfo = metricsNS.NewLabeledGauge("engine", "The information related to the engine and the OS it is running on", metrics.Unit("info"),
		"version",
		"commit",
		"architecture",
		"graphdriver",
		"kernel",
		"os",
		"os_type",
		"os_version",
		"daemon_id", // ID is a randomly generated unique identifier (e.g. UUID4)
	)
	engineCpus   = metricsNS.NewGauge("engine_cpus", "The number of cpus that the host system of the engine has", metrics.Unit("cpus"))
	engineMemory = metricsNS.NewGauge("engine_memory", "The number of bytes of memory that the host system of the engine has", metrics.Bytes)

	healthChecksCounter       = metricsNS.NewCounter("health_checks", "The total number of health checks")
	healthChecksFailedCounter = metricsNS.NewCounter("health_checks_failed", "The total number of failed health checks")

	stateCtr = newStateCounter(metricsNS, metricsNS.NewDesc("container_states", "The count of containers in various states", metrics.Unit("containers"), "state"))
)

func init() {
	for _, a := range []string{
		"start",
		"changes",
		"commit",
		"create",
		"delete",
	} {
		containerActions.WithValues(a).Update(0)
	}

	metrics.Register(metricsNS)
}

type stateCounter struct {
	mu     sync.RWMutex
	states map[string]string
	desc   *prometheus.Desc
}

func newStateCounter(ns *metrics.Namespace, desc *prometheus.Desc) *stateCounter {
	c := &stateCounter{
		states: make(map[string]string),
		desc:   desc,
	}
	ns.Add(c)
	return c
}

func (ctr *stateCounter) get() (running int, paused int, stopped int) {
	ctr.mu.RLock()
	defer ctr.mu.RUnlock()

	states := map[string]int{
		"running": 0,
		"paused":  0,
		"stopped": 0,
	}
	for _, state := range ctr.states {
		states[state]++
	}
	return states["running"], states["paused"], states["stopped"]
}

func (ctr *stateCounter) set(id, label string) {
	ctr.mu.Lock()
	ctr.states[id] = label
	ctr.mu.Unlock()
}

func (ctr *stateCounter) del(id string) {
	ctr.mu.Lock()
	delete(ctr.states, id)
	ctr.mu.Unlock()
}

func (ctr *stateCounter) Describe(ch chan<- *prometheus.Desc) {
	ch <- ctr.desc
}

func (ctr *stateCounter) Collect(ch chan<- prometheus.Metric) {
	running, paused, stopped := ctr.get()
	ch <- prometheus.MustNewConstMetric(ctr.desc, prometheus.GaugeValue, float64(running), "running")
	ch <- prometheus.MustNewConstMetric(ctr.desc, prometheus.GaugeValue, float64(paused), "paused")
	ch <- prometheus.MustNewConstMetric(ctr.desc, prometheus.GaugeValue, float64(stopped), "stopped")
}

func (daemon *Daemon) cleanupMetricsPlugins() {
	ls := daemon.PluginStore.GetAllManagedPluginsByCap(metricsPluginType)
	var wg sync.WaitGroup
	wg.Add(len(ls))

	for _, plugin := range ls {
		p := plugin
		go func() {
			defer wg.Done()

			adapter, err := makePluginAdapter(p)
			if err != nil {
				logrus.WithError(err).WithField("plugin", p.Name()).Error("Error creating metrics plugin adapter")
				return
			}
			if err := adapter.StopMetrics(); err != nil {
				logrus.WithError(err).WithField("plugin", p.Name()).Error("Error stopping plugin metrics collection")
			}
		}()
	}
	wg.Wait()

	if daemon.metricsPluginListener != nil {
		daemon.metricsPluginListener.Close()
	}
}

type metricsPlugin interface {
	StartMetrics() error
	StopMetrics() error
}

func makePluginAdapter(p plugingetter.CompatPlugin) (metricsPlugin, error) {
	if pc, ok := p.(plugingetter.PluginWithV1Client); ok {
		return &metricsPluginAdapter{pc.Client(), p.Name()}, nil
	}

	pa, ok := p.(plugingetter.PluginAddr)
	if !ok {
		return nil, errdefs.System(errors.Errorf("got unknown plugin type %T", p))
	}

	if pa.Protocol() != plugins.ProtocolSchemeHTTPV1 {
		return nil, errors.Errorf("plugin protocol not supported: %s", pa.Protocol())
	}

	addr := pa.Addr()
	client, err := plugins.NewClientWithTimeout(addr.Network()+"://"+addr.String(), nil, pa.Timeout())
	if err != nil {
		return nil, errors.Wrap(err, "error creating metrics plugin client")
	}
	return &metricsPluginAdapter{client, p.Name()}, nil
}

type metricsPluginAdapter struct {
	c    *plugins.Client
	name string
}

func (a *metricsPluginAdapter) StartMetrics() error {
	type metricsPluginResponse struct {
		Err string
	}
	var res metricsPluginResponse
	if err := a.c.Call(metricsPluginType+".StartMetrics", nil, &res); err != nil {
		return errors.Wrap(err, "could not start metrics plugin")
	}
	if res.Err != "" {
		return errors.New(res.Err)
	}
	return nil
}

func (a *metricsPluginAdapter) StopMetrics() error {
	if err := a.c.Call(metricsPluginType+".StopMetrics", nil, nil); err != nil {
		return errors.Wrap(err, "error stopping metrics collector")
	}
	return nil
}
<|endoftext|>


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

package status

import (
	"github.com/juju/juju/workload"
)

type listFormatter struct {
	payloads      []workload.Payload
	compatVersion int
}

func newListFormatter(payloads []workload.Payload, compatVersion int) *listFormatter {
	lf := listFormatter{
		payloads:      payloads,
		compatVersion: compatVersion,
	}
	return &lf
}

func (lf *listFormatter) format() []formattedPayload {
	if lf.payloads == nil {
		return nil
	}

	var formatted []formattedPayload
	for _, payload := range lf.payloads {
		formatted = append(formatted, lf.formatPayload(payload))
	}
	return formatted
}

func (lf *listFormatter) formatPayload(payload workload.Payload) formattedPayload {
	tags := make([]string, len(payload.Tags))
	copy(tags, payload.Tags)
	return formattedPayload{
		Unit:    payload.Unit,
		Machine: payload.Machine,
		ID:      payload.ID,
		Type:    payload.Type,
		Class:   payload.Name,
		Tags:    tags,
		Status:  payload.Status,
	}
}
<commit_msg>Do not format an empty Payload.Tags.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package status

import (
	"github.com/juju/juju/workload"
)

type listFormatter struct {
	payloads      []workload.Payload
	compatVersion int
}

func newListFormatter(payloads []workload.Payload, compatVersion int) *listFormatter {
	lf := listFormatter{
		payloads:      payloads,
		compatVersion: compatVersion,
	}
	return &lf
}

func (lf *listFormatter) format() []formattedPayload {
	if lf.payloads == nil {
		return nil
	}

	var formatted []formattedPayload
	for _, payload := range lf.payloads {
		formatted = append(formatted, lf.formatPayload(payload))
	}
	return formatted
}

func (lf *listFormatter) formatPayload(payload workload.Payload) formattedPayload {
	var tags []string
	if len(payload.Tags) > 0 {
		tags = make([]string, len(payload.Tags))
		copy(tags, payload.Tags)
	}
	return formattedPayload{
		Unit:    payload.Unit,
		Machine: payload.Machine,
		ID:      payload.ID,
		Type:    payload.Type,
		Class:   payload.Name,
		Tags:    tags,
		// TODO(ericsnow) Explicitly convert to a string?
		Status: payload.Status,
	}
}
<|endoftext|>


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

// Package cpufreq provides the current CPU frequency, in MHz, as reported by
// /proc/cpuinfo.
package cpufreq

import (
	"io"
	"strconv"
	"sync"
	"time"

	"github.com/SermoDigital/helpers"
	joe "github.com/mohae/joefriday"
)

const procFile = "/proc/cpuinfo"

// Frequency holds information about the frequency of a system's cpus, in MHz.
// The reported values are the current speeds as reported by /proc/cpuinfo.
type Frequency struct {
	Timestamp int64
	Sockets   uint8
	CPU       []CPU `json:"cpu"`
}

// CPU holds the clock info for a single processor.
type CPU struct {
	Processor       uint16    `json:"processor"`
	CPUMHz          float32  `json:"cpu_mhz"`
	PhysicalID      uint8    `json:"physical_id"`
	CoreID          uint16    `json:"core_id"`
	APICID          uint16    `json:"apicid"`
}

// Profiler is used to process the frequency information.
type Profiler struct {
	joe.Procer
	*joe.Buffer
	Frequency  // this is used too hold the socket/cpu info so that everything doesn't have to be reprocessed.
}

// Returns an initialized Profiler; ready to use.
func NewProfiler() (prof *Profiler, err error) {
	proc, err := joe.NewProc(procFile)
	if err != nil {
		return nil, err
	}
	prof = &Profiler{Procer: proc, Buffer: joe.NewBuffer()}
	err = prof.InitFrequency()
	if err != nil {
		return nil, err
	}
	return prof, nil
}

// Reset resources; after reset the profiler is ready to be used again.
func (prof *Profiler) Reset() error {
	prof.Buffer.Reset()
	return prof.Procer.Reset()
}

// InitFrequency sets the profiler's frequency with the static information so
// that everything doesn't need to be reprocessed every time the frequency is
// requested. This assumes that cpuinfo returns processor information in the
// same order every time.
//
// This shouldn't be used; it's exported for testing reasons.
func (prof *Profiler) InitFrequency() error {
	var (
		err          error
		n            uint64
		pos, cpuCnt  int
		pidFound     bool
		physIDs      []uint8 // tracks unique physical IDs encountered
		cpu          CPU
	)

	prof.Frequency = Frequency{}
	for {
		prof.Line, err = prof.ReadSlice('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			return &joe.ReadError{Err: err}
		}
		prof.Val = prof.Val[:0]
		// First grab the attribute name; everything up to the ':'.  The key may have
		// spaces and has trailing spaces; that gets trimmed.
		for i, v := range prof.Line {
			if v == 0x3A {
				prof.Val = prof.Line[:i]
				pos = i + 1
				break
			}
			//prof.Val = append(prof.Val, v)
		}
		prof.Val = joe.TrimTrailingSpaces(prof.Val[:])
		nameLen := len(prof.Val)
		// if there's no name; skip.
		if nameLen == 0 {
			continue
		}
		// if there's anything left, the value is everything else; trim spaces
		if pos+1 < len(prof.Line) {
			prof.Val = append(prof.Val, joe.TrimTrailingSpaces(prof.Line[pos+1:])...)
		}
		if prof.Val[0] == 'a' {
			if prof.Val[1] == 'p' { // apicid
				n, err = helpers.ParseUint(prof.Val[nameLen:])
				if err != nil {
					return &joe.ParseError{Info: string(prof.Val[:nameLen]), Err: err}
				}
				cpu.APICID = uint16(n)
			}
			continue
		}
		if prof.Val[0] == 'c' {
			if prof.Val[1] == 'o' { // core id
				n, err = helpers.ParseUint(prof.Val[nameLen:])
				if err != nil {
					return &joe.ParseError{Info: string(prof.Val[:nameLen]), Err: err}
				}
				cpu.CoreID = uint16(n)
			}
			continue
		}
		if prof.Val[0] == 'p' {
			if prof.Val[1] == 'h' { // physical id
				n, err = helpers.ParseUint(prof.Val[nameLen:])
				if err != nil {
					return &joe.ParseError{Info: string(prof.Val[:nameLen]), Err: err}
				}
				cpu.PhysicalID = uint8(n)
				for i := range physIDs {
					if physIDs[i] == cpu.PhysicalID {
						pidFound = true
						break
					}
				}
				if pidFound {
					pidFound = false  // reset for next use
				} else {
					// physical id hasn't been encountered yet; add it
					physIDs = append(physIDs, cpu.PhysicalID)
				}
				continue
			}
			// processor starts information about a processor.
			if prof.Val[1] == 'r' { // processor
				if cpuCnt > 0 {
					prof.Frequency.CPU = append(prof.Frequency.CPU, cpu)
				}
				cpuCnt++
				n, err = helpers.ParseUint(prof.Val[nameLen:])
				if err != nil {
					return &joe.ParseError{Info: string(prof.Val[:nameLen]), Err: err}
				}
				cpu = CPU{Processor: uint16(n)}
			}
		}
		continue
	}
	// append the current processor informatin
	prof.Frequency.CPU = append(prof.Frequency.CPU, cpu)
	prof.Frequency.Sockets = uint8(len(physIDs))
	return  nil
}

// returns a copy of the profiler's frequency.
func (prof *Profiler) newFrequency() *Frequency {
	f := &Frequency{Timestamp: time.Now().UTC().UnixNano(), Sockets: prof.Frequency.Sockets, CPU: make([]CPU, len(prof.Frequency.CPU))}
	copy(f.CPU, prof.Frequency.CPU)
	return f
}

// Get returns Frequency information.
func (prof *Profiler) Get() (f *Frequency, err error) {
	f = prof.newFrequency()
	err = prof.Reset()
	if err != nil {
		return nil, err
	}

	var (
		i, pos, nameLen int
		v               byte
		x               float64
	)
	processor := -1  // start at -1 because it'll be incremented before use as it's the first line encountered
	for {
		prof.Line, err = prof.ReadSlice('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			return nil, &joe.ReadError{Err: err}
		}
		prof.Val = prof.Val[:0]
		// First grab the attribute name; everything up to the ':'.  The key may have
		// spaces and has trailing spaces; that gets trimmed.
		for i, v = range prof.Line {
			if v == 0x3A {
				prof.Val = prof.Line[:i]
				pos = i + 1
				break
			}
			//prof.Val = append(prof.Val, v)
		}
		prof.Val = joe.TrimTrailingSpaces(prof.Val[:])
		nameLen = len(prof.Val)
		// if there's no name; skip.
		if nameLen == 0 {
			continue
		}
		// if there's anything left, the value is everything else; trim spaces
		if pos+1 < len(prof.Line) {
			prof.Val = append(prof.Val, joe.TrimTrailingSpaces(prof.Line[pos+1:])...)
		}
		if prof.Val[0] == 'c' {
			if prof.Val[4] == 'M' { // cpu MHz
				x, err = strconv.ParseFloat(string(prof.Val[nameLen:]), 32)
				if err != nil {
					return nil, &joe.ParseError{Info: string(prof.Val[:nameLen]), Err: err}
				}
				f.CPU[processor].CPUMHz = float32(x)
			}
			continue
		}
		if prof.Val[0] == 'p' {
		// processor starts information about a processor.
			if prof.Val[1] == 'r' { // processor
				processor++
			}
		}
	}
	return f, nil
}

var std *Profiler
var stdMu sync.Mutex

// Get returns Frequency using the package's global Profiler.
func Get() (f *Frequency, err error) {
	stdMu.Lock()
	defer stdMu.Unlock()
	if std == nil {
		std, err = NewProfiler()
		if err != nil {
			return nil, err
		}
	}
	return std.Get()
}

// Ticker delivers the CPU Frequencies at intervals.
type Ticker struct {
	*joe.Ticker
	Data chan *Frequency
	*Profiler
	Sockets uint8
}

// NewTicker returns a new Ticker containing a Data channel that delivers the
// data at intervals and an error channel that delivers any errors encountered.
// Stop the ticker to signal the ticker to stop running. Stopping the ticker
// does not close the Data channel; call Close to close both the ticker and the
// data channel.
func NewTicker(d time.Duration) (joe.Tocker, error) {
	p, err := NewProfiler()
	if err != nil {
		return nil, err
	}
	t := Ticker{Ticker: joe.NewTicker(d), Data: make(chan *Frequency), Profiler: p}
	go t.Run()
	return &t, nil
}

// Run runs the ticker.
func (t *Ticker) Run() {
	// ticker
	for {
		select {
		case <-t.Done:
			return
		case <-t.C:
			s, err := t.Get()
			if err != nil {
				t.Errs <- err
				continue
			}
			t.Data <- s
		}
	}
}

// Close closes the ticker resources.
func (t *Ticker) Close() {
	t.Ticker.Close()
	close(t.Data)
}
<commit_msg>add json tags<commit_after>// Copyright 2016 Joel Scoble and The JoeFriday authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package cpufreq provides the current CPU frequency, in MHz, as reported by
// /proc/cpuinfo.
package cpufreq

import (
	"io"
	"strconv"
	"sync"
	"time"

	"github.com/SermoDigital/helpers"
	joe "github.com/mohae/joefriday"
)

const procFile = "/proc/cpuinfo"

// Frequency holds information about the frequency of a system's cpus, in MHz.
// The reported values are the current speeds as reported by /proc/cpuinfo.
type Frequency struct {
	Timestamp int64 `json:"timestamp"`
	Sockets   uint8 `json:"sockets"`
	CPU       []CPU `json:"cpu"`
}

// CPU holds the clock info for a single processor.
type CPU struct {
	Processor  uint16  `json:"processor"`
	CPUMHz     float32 `json:"cpu_mhz"`
	PhysicalID uint8   `json:"physical_id"`
	CoreID     uint16  `json:"core_id"`
	APICID     uint16  `json:"apicid"`
}

// Profiler is used to process the frequency information.
type Profiler struct {
	joe.Procer
	*joe.Buffer
	Frequency // this is used too hold the socket/cpu info so that everything doesn't have to be reprocessed.
}

// Returns an initialized Profiler; ready to use.
func NewProfiler() (prof *Profiler, err error) {
	proc, err := joe.NewProc(procFile)
	if err != nil {
		return nil, err
	}
	prof = &Profiler{Procer: proc, Buffer: joe.NewBuffer()}
	err = prof.InitFrequency()
	if err != nil {
		return nil, err
	}
	return prof, nil
}

// Reset resources; after reset the profiler is ready to be used again.
func (prof *Profiler) Reset() error {
	prof.Buffer.Reset()
	return prof.Procer.Reset()
}

// InitFrequency sets the profiler's frequency with the static information so
// that everything doesn't need to be reprocessed every time the frequency is
// requested. This assumes that cpuinfo returns processor information in the
// same order every time.
//
// This shouldn't be used; it's exported for testing reasons.
func (prof *Profiler) InitFrequency() error {
	var (
		err         error
		n           uint64
		pos, cpuCnt int
		pidFound    bool
		physIDs     []uint8 // tracks unique physical IDs encountered
		cpu         CPU
	)

	prof.Frequency = Frequency{}
	for {
		prof.Line, err = prof.ReadSlice('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			return &joe.ReadError{Err: err}
		}
		prof.Val = prof.Val[:0]
		// First grab the attribute name; everything up to the ':'.  The key may have
		// spaces and has trailing spaces; that gets trimmed.
		for i, v := range prof.Line {
			if v == 0x3A {
				prof.Val = prof.Line[:i]
				pos = i + 1
				break
			}
			//prof.Val = append(prof.Val, v)
		}
		prof.Val = joe.TrimTrailingSpaces(prof.Val[:])
		nameLen := len(prof.Val)
		// if there's no name; skip.
		if nameLen == 0 {
			continue
		}
		// if there's anything left, the value is everything else; trim spaces
		if pos+1 < len(prof.Line) {
			prof.Val = append(prof.Val, joe.TrimTrailingSpaces(prof.Line[pos+1:])...)
		}
		if prof.Val[0] == 'a' {
			if prof.Val[1] == 'p' { // apicid
				n, err = helpers.ParseUint(prof.Val[nameLen:])
				if err != nil {
					return &joe.ParseError{Info: string(prof.Val[:nameLen]), Err: err}
				}
				cpu.APICID = uint16(n)
			}
			continue
		}
		if prof.Val[0] == 'c' {
			if prof.Val[1] == 'o' { // core id
				n, err = helpers.ParseUint(prof.Val[nameLen:])
				if err != nil {
					return &joe.ParseError{Info: string(prof.Val[:nameLen]), Err: err}
				}
				cpu.CoreID = uint16(n)
			}
			continue
		}
		if prof.Val[0] == 'p' {
			if prof.Val[1] == 'h' { // physical id
				n, err = helpers.ParseUint(prof.Val[nameLen:])
				if err != nil {
					return &joe.ParseError{Info: string(prof.Val[:nameLen]), Err: err}
				}
				cpu.PhysicalID = uint8(n)
				for i := range physIDs {
					if physIDs[i] == cpu.PhysicalID {
						pidFound = true
						break
					}
				}
				if pidFound {
					pidFound = false // reset for next use
				} else {
					// physical id hasn't been encountered yet; add it
					physIDs = append(physIDs, cpu.PhysicalID)
				}
				continue
			}
			// processor starts information about a processor.
			if prof.Val[1] == 'r' { // processor
				if cpuCnt > 0 {
					prof.Frequency.CPU = append(prof.Frequency.CPU, cpu)
				}
				cpuCnt++
				n, err = helpers.ParseUint(prof.Val[nameLen:])
				if err != nil {
					return &joe.ParseError{Info: string(prof.Val[:nameLen]), Err: err}
				}
				cpu = CPU{Processor: uint16(n)}
			}
		}
		continue
	}
	// append the current processor informatin
	prof.Frequency.CPU = append(prof.Frequency.CPU, cpu)
	prof.Frequency.Sockets = uint8(len(physIDs))
	return nil
}

// returns a copy of the profiler's frequency.
func (prof *Profiler) newFrequency() *Frequency {
	f := &Frequency{Timestamp: time.Now().UTC().UnixNano(), Sockets: prof.Frequency.Sockets, CPU: make([]CPU, len(prof.Frequency.CPU))}
	copy(f.CPU, prof.Frequency.CPU)
	return f
}

// Get returns Frequency information.
func (prof *Profiler) Get() (f *Frequency, err error) {
	f = prof.newFrequency()
	err = prof.Reset()
	if err != nil {
		return nil, err
	}

	var (
		i, pos, nameLen int
		v               byte
		x               float64
	)
	processor := -1 // start at -1 because it'll be incremented before use as it's the first line encountered
	for {
		prof.Line, err = prof.ReadSlice('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			return nil, &joe.ReadError{Err: err}
		}
		prof.Val = prof.Val[:0]
		// First grab the attribute name; everything up to the ':'.  The key may have
		// spaces and has trailing spaces; that gets trimmed.
		for i, v = range prof.Line {
			if v == 0x3A {
				prof.Val = prof.Line[:i]
				pos = i + 1
				break
			}
			//prof.Val = append(prof.Val, v)
		}
		prof.Val = joe.TrimTrailingSpaces(prof.Val[:])
		nameLen = len(prof.Val)
		// if there's no name; skip.
		if nameLen == 0 {
			continue
		}
		// if there's anything left, the value is everything else; trim spaces
		if pos+1 < len(prof.Line) {
			prof.Val = append(prof.Val, joe.TrimTrailingSpaces(prof.Line[pos+1:])...)
		}
		if prof.Val[0] == 'c' {
			if prof.Val[4] == 'M' { // cpu MHz
				x, err = strconv.ParseFloat(string(prof.Val[nameLen:]), 32)
				if err != nil {
					return nil, &joe.ParseError{Info: string(prof.Val[:nameLen]), Err: err}
				}
				f.CPU[processor].CPUMHz = float32(x)
			}
			continue
		}
		if prof.Val[0] == 'p' {
			// processor starts information about a processor.
			if prof.Val[1] == 'r' { // processor
				processor++
			}
		}
	}
	return f, nil
}

var std *Profiler
var stdMu sync.Mutex

// Get returns Frequency using the package's global Profiler.
func Get() (f *Frequency, err error) {
	stdMu.Lock()
	defer stdMu.Unlock()
	if std == nil {
		std, err = NewProfiler()
		if err != nil {
			return nil, err
		}
	}
	return std.Get()
}

// Ticker delivers the CPU Frequencies at intervals.
type Ticker struct {
	*joe.Ticker
	Data chan *Frequency
	*Profiler
	Sockets uint8
}

// NewTicker returns a new Ticker containing a Data channel that delivers the
// data at intervals and an error channel that delivers any errors encountered.
// Stop the ticker to signal the ticker to stop running. Stopping the ticker
// does not close the Data channel; call Close to close both the ticker and the
// data channel.
func NewTicker(d time.Duration) (joe.Tocker, error) {
	p, err := NewProfiler()
	if err != nil {
		return nil, err
	}
	t := Ticker{Ticker: joe.NewTicker(d), Data: make(chan *Frequency), Profiler: p}
	go t.Run()
	return &t, nil
}

// Run runs the ticker.
func (t *Ticker) Run() {
	// ticker
	for {
		select {
		case <-t.Done:
			return
		case <-t.C:
			s, err := t.Get()
			if err != nil {
				t.Errs <- err
				continue
			}
			t.Data <- s
		}
	}
}

// Close closes the ticker resources.
func (t *Ticker) Close() {
	t.Ticker.Close()
	close(t.Data)
}
<|endoftext|>


<commit_before>package keyman

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
)

var (
	nssDBs = []string{
		filepath.Join(os.Getenv("HOME"), ".pki/nssdb"),
		filepath.Join(os.Getenv("HOME"), "snap/chromium/current/.pki/nssdb"), // Snapcraft
		"/etc/pki/nssdb", // CentOS 7
	}
	FirefoxProfile = os.Getenv("HOME") + "/.mozilla/firefox/*"
)

func DeleteTrustedRootByName(commonName string, prompt string) error {
	return forEachNSSProfile(func(profile string) error {
		cmd := exec.Command("certutil", "-d", profile, "-D", "-n", commonName)
		out, err := cmd.CombinedOutput()
		if err != nil {
			return fmt.Errorf("Unable to run certutil command: %s\n%s", err, out)
		}
		return nil
	})
}

// AddAsTrustedRoot adds the certificate to the user's trust store as a trusted
// root CA. Supports Chrome and Firefox
func (cert *Certificate) AddAsTrustedRoot(prompt string) error {
	tempFileName, err := cert.WriteToTempFile()
	defer os.Remove(tempFileName)
	if err != nil {
		return fmt.Errorf("Unable to create temp file: %s", err)
	}

	return forEachNSSProfile(func(profile string) error {
		// Add it as a trusted cert
		// https://code.google.com/p/chromium/wiki/LinuxCertManagement#Add_a_certificate
		cmd := exec.Command("certutil", "-d", profile, "-A", "-t", "C,,", "-n", cert.X509().Subject.CommonName, "-i", tempFileName)
		out, err := cmd.CombinedOutput()
		if err != nil {
			return fmt.Errorf("Unable to run certutil command: %s\n%s", err, out)
		}
		return nil
	})
}

func pathExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

func forEachNSSProfile(f func(profile string) error) error {
	profiles, _ := filepath.Glob(FirefoxProfile)
	profiles = append(profiles, nssDBs...)
	for _, profile := range profiles {
		if stat, err := os.Stat(profile); err != nil || !stat.IsDir() {
			continue
		}
		if pathExists(filepath.Join(profile, "cert9.db")) {
			if err := f("sql:" + profile); err != nil {
				return err
			}
		} else if pathExists(filepath.Join(profile, "cert8.db")) {
			if err := f("dbm:" + profile); err != nil {
				return err
			}
		}
	}
	return nil
}

// IsInstalled checks whether this certificate is install based purely on looking for a cert
// in the user's nssdb that has the same common name.  This function returns
// true if there are one or more certs in the nssdb whose common name
// matches this cert.
func (cert *Certificate) IsInstalled() (bool, error) {
	found := false
	err := forEachNSSProfile(func(profile string) error {
		cmd := exec.Command("certutil", "-d", profile, "-L", "-n", cert.X509().Subject.CommonName)
		err := cmd.Run()

		if err == nil {
			found = true
		}
		return nil
	})

	return found, err
}
<commit_msg>minor cleanup<commit_after>package keyman

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
)

var (
	nssDBs = []string{
		filepath.Join(os.Getenv("HOME"), ".pki/nssdb"),
		filepath.Join(os.Getenv("HOME"), "snap/chromium/current/.pki/nssdb"), // Snapcraft
		"/etc/pki/nssdb", // CentOS 7
	}
	FirefoxProfile = os.Getenv("HOME") + "/.mozilla/firefox/*"
)

func DeleteTrustedRootByName(commonName string, prompt string) error {
	return forEachNSSProfile(func(profile string) error {
		cmd := exec.Command("certutil", "-d", profile, "-D", "-n", commonName)
		out, err := cmd.CombinedOutput()
		if err != nil {
			return fmt.Errorf("Unable to run certutil command: %w\n%s", err, out)
		}
		return nil
	})
}

// AddAsTrustedRoot adds the certificate to the user's trust store as a trusted
// root CA. Supports Chrome and Firefox
func (cert *Certificate) AddAsTrustedRoot(prompt string) error {
	tempFileName, err := cert.WriteToTempFile()
	defer os.Remove(tempFileName)
	if err != nil {
		return fmt.Errorf("Unable to create temp file: %s", err)
	}

	return forEachNSSProfile(func(profile string) error {
		// Add it as a trusted cert
		// https://code.google.com/p/chromium/wiki/LinuxCertManagement#Add_a_certificate
		cmd := exec.Command("certutil", "-d", profile, "-A", "-t", "C,,", "-n", cert.X509().Subject.CommonName, "-i", tempFileName)
		out, err := cmd.CombinedOutput()
		if err != nil {
			return fmt.Errorf("Unable to run certutil command: %w\n%s", err, out)
		}
		return nil
	})
}

func pathExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

func forEachNSSProfile(f func(profile string) error) error {
	profiles, _ := filepath.Glob(FirefoxProfile)
	profiles = append(profiles, nssDBs...)
	for _, profile := range profiles {
		if stat, err := os.Stat(profile); err != nil || !stat.IsDir() {
			continue
		}
		if pathExists(filepath.Join(profile, "cert9.db")) {
			if err := f("sql:" + profile); err != nil {
				return err
			}
		} else if pathExists(filepath.Join(profile, "cert8.db")) {
			if err := f("dbm:" + profile); err != nil {
				return err
			}
		}
	}
	return nil
}

// IsInstalled checks whether this certificate is install based purely on looking for a cert
// in the user's nssdb that has the same common name.  This function returns
// true if there are one or more certs in the nssdb whose common name
// matches this cert.
func (cert *Certificate) IsInstalled() (bool, error) {
	found := false
	err := forEachNSSProfile(func(profile string) error {
		cmd := exec.Command("certutil", "-d", profile, "-L", "-n", cert.X509().Subject.CommonName)
		if cmd.Run() == nil {
			found = true
		}
		return nil
	})

	return found, err
}
<|endoftext|>


<commit_before>/*
 *  Copyright (c) 2019, Salesforce.com, Inc.
 *  All rights reserved.
 *  Licensed under the BSD 3-Clause license.
 *  For full license text, see LICENSE.txt file in the repo root  or https://opensource.org/licenses/BSD-3-Clause
 */

package main

import (
	"log"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"syscall"
)

func main() {
	jarName := os.Args[0]
	jarArgs := os.Args[1:]

	binary, err := exec.LookPath("java")
	if err != nil {
		log.Fatalf("Java execution error: %v", err)
	}

	args := []string{"-jar", strings.TrimPrefix(jarName, "./")}
	args = append(args, jarArgs...)

	if runtime.GOOS == "windows" {
		cmd := exec.Command(binary, args...)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.Stdin = os.Stdin

		if err := cmd.Run(); err != nil {
			if err, ok := err.(*exec.ExitError); ok {
				if status, ok := err.Sys().(syscall.WaitStatus); ok {
					// Exit with the same code as the Java program
					os.Exit(status.ExitStatus())
				}
			} else {
				log.Fatalf("Bootstrap execution error: %v", err)
			}
		}
	} else {
		err = syscall.Exec(binary, append([]string{"java"}, args...), os.Environ())
		if err != nil {
			log.Fatalf("Bootstrap execution error: %v", os.Environ())
		}
	}
}
<commit_msg>Try finding Java in JAVA_HOME if `java` isn't on PATH<commit_after>/*
 *  Copyright (c) 2019, Salesforce.com, Inc.
 *  All rights reserved.
 *  Licensed under the BSD 3-Clause license.
 *  For full license text, see LICENSE.txt file in the repo root  or https://opensource.org/licenses/BSD-3-Clause
 */

package main

import (
	"log"
	"os"
	"os/exec"
	"path"
	"runtime"
	"strings"
	"syscall"
)

func main() {
	jarName := os.Args[0]
	jarArgs := os.Args[1:]

	binary, err := exec.LookPath("java")
	if err != nil {
		javaHome, javaHomeSet := os.LookupEnv("JAVA_HOME")
		if !javaHomeSet {
			log.Fatalf("Java not found in PATH and JAVA_HOME not set: %v", err)
		}

		stat, err := os.Stat(javaHome)
		if err != nil {
			log.Fatalf("Java not found in JAVA_HOME: %v", err)
		} else if !stat.IsDir() {
			log.Fatalf("JAVA_HOME is not a directory: %v", err)
		} else {
			binary = path.Join(javaHome, "bin", "java")
		}
	}

	args := []string{"-jar", strings.TrimPrefix(jarName, "./")}
	args = append(args, jarArgs...)

	if runtime.GOOS == "windows" {
		cmd := exec.Command(binary, args...)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.Stdin = os.Stdin

		if err := cmd.Run(); err != nil {
			if err, ok := err.(*exec.ExitError); ok {
				if status, ok := err.Sys().(syscall.WaitStatus); ok {
					// Exit with the same code as the Java program
					os.Exit(status.ExitStatus())
				}
			} else {
				log.Fatalf("Bootstrap execution error: %v", err)
			}
		}
	} else {
		err = syscall.Exec(binary, append([]string{"java"}, args...), os.Environ())
		if err != nil {
			log.Fatalf("Bootstrap execution error: %v", os.Environ())
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
)

func main() {
	fmt.Println("Hello world!")
}
<commit_msg>Read Data<commit_after>package main

import (
	"bufio"
	"fmt"
	"os"
	"strconv"
	"strings"
)

func main() {
	file, _ := os.Open(os.Args[1])

	scanner := bufio.NewScanner(file)
	triangle := make([][]int64, 0)

	for scanner.Scan() {

		line := make([]int64, 0)

		for _, elem := range strings.Split(scanner.Text(), " ") {
			num, _ := strconv.ParseInt(elem, 10, 64)
			line = append(line, num)
		}

		triangle = append(triangle, line)
	}

	fmt.Printf("%v", triangle)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/zenazn/goji/web"
	"gopkg.in/mgo.v2"
)

type SidewinderDirector struct {
	MongoDB string
	session *mgo.Session
}

func (self *SidewinderDirector) Store() *SidewinderStore {
	return &SidewinderStore{self.MongoDB, self.session.Copy()}
}

type SidewinderStore struct {
	mongoDB string
	session *mgo.Session
}

func (self *SidewinderStore) DB() *mgo.Database {
	return self.session.DB(self.mongoDB)
}

func (self *SidewinderStore) Close() {
	self.session.Close()
}

type DeviceDocument struct {
	DeviceId string `_id`
}

func (self *SidewinderStore) AddDevice(deviceId string) error {
	document := DeviceDocument{deviceId}

	err := self.DB().C("devices").Insert(document)
	return err
}

type DatastoreInfo struct {
	BuildInfo     mgo.BuildInfo
	DatabaseNames []string
}

func (self *SidewinderDirector) DatastoreInfo(context web.C, writer http.ResponseWriter, request *http.Request) {
	session := self.Store().session
	buildInfo, err := session.BuildInfo()
	if err != nil {
		fmt.Fprintf(writer, "Could not connect to MongoDB.\n%v", err.Error())
		return
	}
	writer.Header().Set("Content-Type", "application/json")

	databases, err := session.DatabaseNames()

	if err != nil {
		fmt.Fprintf(writer, "Could not retrieve database names.\n%v", err.Error())
		return
	}

	dataStoreInfo := DatastoreInfo{buildInfo, databases}

	err = json.NewEncoder(writer).Encode(&dataStoreInfo)
	if err != nil {
		fmt.Fprintf(writer, "Could not return info from MongoDB.\n%v", err.Error())
		return
	}
}
<commit_msg>Adding a tiny amount of additional info to the store info page.<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/zenazn/goji/web"
	"gopkg.in/mgo.v2"
)

type SidewinderDirector struct {
	MongoDB string
	session *mgo.Session
}

func (self *SidewinderDirector) Store() *SidewinderStore {
	return &SidewinderStore{self.MongoDB, self.session.Copy()}
}

type SidewinderStore struct {
	mongoDB string
	session *mgo.Session
}

func (self *SidewinderStore) DB() *mgo.Database {
	return self.session.DB(self.mongoDB)
}

func (self *SidewinderStore) Close() {
	self.session.Close()
}

type DeviceDocument struct {
	DeviceId string `_id`
}

func (self *SidewinderStore) AddDevice(deviceId string) error {
	document := DeviceDocument{deviceId}

	err := self.DB().C("devices").Insert(document)
	return err
}

type DatastoreInfo struct {
	BuildInfo     mgo.BuildInfo
	LiveServers   []string
	DatabaseNames []string
}

func (self *SidewinderDirector) DatastoreInfo(context web.C, writer http.ResponseWriter, request *http.Request) {
	session := self.Store().session

	buildInfo, err := session.BuildInfo()
	if err != nil {
		fmt.Fprintf(writer, "Could not connect to MongoDB.\n%v", err.Error())
		return
	}
	writer.Header().Set("Content-Type", "application/json")

	databases, err := session.DatabaseNames()

	if err != nil {
		fmt.Fprintf(writer, "Could not retrieve database names.\n%v", err.Error())
		return
	}

	dataStoreInfo := DatastoreInfo{buildInfo, session.LiveServers(), databases}

	err = json.NewEncoder(writer).Encode(&dataStoreInfo)
	if err != nil {
		fmt.Fprintf(writer, "Could not return info from MongoDB.\n%v", err.Error())
		return
	}
}
<|endoftext|>


<commit_before>package nds

import (
	"appengine"
	"appengine/datastore"
	"errors"
	"fmt"
	"reflect"
	"sync"
)

const (
	// multiLimit is the App Engine datastore limit for the number of entities
	// that can be PutMulti or GetMulti in one call.
	multiLimit = 1000
)

var (
	// milMultiError is a convenience slice used to represent a nil error when
	// grouping errors in GetMulti.
	nilMultiError = make(appengine.MultiError, multiLimit)
)

func checkMultiArgs(keys []*datastore.Key, v reflect.Value) error {
	if v.Kind() != reflect.Slice {
		return errors.New("nds: dst is not a slice")
	}

	if len(keys) != v.Len() {
		return errors.New("nds: key and dst slices have different length")
	}

	return nil
}

// GetMulti works just like datastore.GetMulti except it removes the API limit
// of 1000 entities per request by calling datastore.GetMulti as many times as
// required to complete the request.
//
// Increase the datastore timeout if you get datastore_v3: TIMEOUT errors. You
// can do this using
// http://godoc.org/code.google.com/p/appengine-go/appengine#Timeout.
func GetMulti(c appengine.Context,
	keys []*datastore.Key, dst interface{}) error {

	v := reflect.ValueOf(dst)
	if err := checkMultiArgs(keys, v); err != nil {
		return err
	}

	if len(keys) == 0 {
		return nil
	}

	p := len(keys) / multiLimit
	errs := make([]error, p+1)
	wg := sync.WaitGroup{}
	for i := 0; i < p; i++ {
		index := i
		keySlice := keys[i*multiLimit : (i+1)*multiLimit]
		dstSlice := v.Slice(i*multiLimit, (i+1)*multiLimit)

		wg.Add(1)
		go func() {
			errs[index] = datastore.GetMulti(c, keySlice, dstSlice.Interface())
			wg.Done()
		}()
	}

	if len(keys)%multiLimit == 0 {
		errs = errs[:len(errs)-1]
	} else {
		keySlice := keys[p*multiLimit : len(keys)]
		dstSlice := v.Slice(p*multiLimit, len(keys))
		wg.Add(1)
		go func() {
			errs[p] = datastore.GetMulti(c, keySlice, dstSlice.Interface())
			wg.Done()
		}()
	}
	wg.Wait()

	// Quick escape if all errors are nil.
	errsNil := true
	for _, err := range errs {
		if err != nil {
			errsNil = false
		}
	}
	if errsNil {
		return nil
	}

	groupedErrs := make(appengine.MultiError, 0, len(keys))
	for _, err := range errs {
		if err == nil {
			groupedErrs = append(groupedErrs, nilMultiError...)
		} else if me, ok := err.(appengine.MultiError); ok {
			groupedErrs = append(groupedErrs, me...)
		} else {
			return err
		}
	}
	return groupedErrs[:len(keys)]
}

type cacheContext struct {
	appengine.Context
	cache map[string]*datastore.PropertyList
	sync.RWMutex
}

func NewCacheContext(c appengine.Context) appengine.Context {
	return &cacheContext{
		Context: c,
		cache:   map[string]*datastore.PropertyList{},
	}
}

func GetMultiCache(c appengine.Context,
	keys []*datastore.Key, dst interface{}) error {

	v := reflect.ValueOf(dst)
	if err := checkMultiArgs(keys, v); err != nil {
		return err
	}

	if cc, ok := c.(*cacheContext); ok {
		return getMultiCache(cc, keys, v)
	} else {
		return datastore.GetMulti(c, keys, dst)
	}
}

func convertToPropertyLists(
	v reflect.Value) ([]datastore.PropertyList, error) {
	pls := make([]datastore.PropertyList, v.Len())
	for i := range pls {
		pl := datastore.PropertyList{}
		elem := addrValue(v.Index(i))
		if err := SaveStruct(elem.Interface(), &pl); err != nil {
			return nil, err
		}
		pls[i] = pl
	}
	return pls, nil
}

func addrValue(v reflect.Value) reflect.Value {
	if v.Kind() != reflect.Ptr {
		return v.Addr()
	} else {
		return v
	}
}

// getMultiCache gets entities from local cache then the datastore.
// dst argument must be a slice.
func getMultiCache(cc *cacheContext,
	keys []*datastore.Key, dst reflect.Value) error {

	cacheMissIndexes := []int{}
	cacheMissKeys := []*datastore.Key{}
	cacheMissDsts := []datastore.PropertyList{}

	errors := make(appengine.MultiError, dst.Len())
	errsNil := true

	// Load what we can from the local cache.
	cc.RLock()
	for i, key := range keys {
		if pl, ok := cc.cache[key.Encode()]; ok {
			if pl == nil {
				errors[i] = datastore.ErrNoSuchEntity
				errsNil = false
			} else {
				elem := addrValue(dst.Index(i))
				if err := LoadStruct(elem.Interface(), pl); err != nil {
					return err
				}
			}
		} else {
			cacheMissIndexes = append(cacheMissIndexes, i)
			cacheMissKeys = append(cacheMissKeys, key)
			cacheMissDsts = append(cacheMissDsts, datastore.PropertyList{})
		}
	}
	cc.RUnlock()

	// Load from datastore.
	if err := datastore.GetMulti(cc, cacheMissKeys, cacheMissDsts); err == nil {
		// Save to local memory cache.
		putMultiLocalCache(cc, cacheMissKeys, cacheMissDsts)

		// Update the callers slice with values.
		for i, index := range cacheMissIndexes {
			pl := cacheMissDsts[i]
			elem := addrValue(dst.Index(index))
			if err := LoadStruct(elem.Interface(), &pl); err != nil {
				return err
			}
		}
	} else if me, ok := err.(appengine.MultiError); ok {
		for i, err := range me {
			if err == nil {
				putLocalCache(cc, cacheMissKeys[i], cacheMissDsts[i])

				// Update the callers slice with values.
				pl := cacheMissDsts[i]
				index := cacheMissIndexes[i]
				elem := addrValue(dst.Index(index))
				if err := LoadStruct(elem.Interface(), &pl); err != nil {
					return err
				}
			} else if err == datastore.ErrNoSuchEntity {
				putLocalCache(cc, cacheMissKeys[i], nil)
				index := cacheMissIndexes[i]
				errors[index] = datastore.ErrNoSuchEntity
				errsNil = false
				// Possibly should zero the callers slice value here.
			} else {
				return err
			}
		}
	} else {
		return err
	}

	if errsNil {
		return nil
	}
	return errors
}

func PutMultiCache(c appengine.Context,
	keys []*datastore.Key, src interface{}) ([]*datastore.Key, error) {

	v := reflect.ValueOf(src)
	if err := checkMultiArgs(keys, v); err != nil {
		return nil, err
	}

	if cc, ok := c.(*cacheContext); ok {
		if pls, err := convertToPropertyLists(v); err != nil {
			fmt.Println("Convert error", err)
			return nil, err
		} else {
			return putMultiCache(cc, keys, pls)
		}
	} else {
		return datastore.PutMulti(c, keys, src)
	}
}

// putMultiCache puts the entities into the datastore and then its local cache.
func putMultiCache(cc *cacheContext,
	keys []*datastore.Key,
	pls []datastore.PropertyList) ([]*datastore.Key, error) {

	// Save to the datastore.
	completeKeys, err := datastore.PutMulti(cc, keys, pls)
	if err != nil {
		return nil, err
	}

	// Save to local memory cache.
	putMultiLocalCache(cc, completeKeys, pls)

	return completeKeys, nil
}

func putLocalCache(cc *cacheContext,
	key *datastore.Key, pl datastore.PropertyList) {
	cc.Lock()
	cc.cache[key.Encode()] = &pl
	cc.Unlock()
}

func putMultiLocalCache(cc *cacheContext,
	keys []*datastore.Key, pls []datastore.PropertyList) {
	for i, key := range keys {
		putLocalCache(cc, key, pls[i])
	}
}

// SaveStruct saves src to a datastore.PropertyList.
func SaveStruct(src interface{}, pl *datastore.PropertyList) error {
	c, err := make(chan datastore.Property), make(chan error)
	go func() {
		err <- datastore.SaveStruct(src, c)
	}()
	for p := range c {
		*pl = append(*pl, p)
	}
	return <-err
}

// LoadStruct loads a datastore.PropertyList into dst.
func LoadStruct(dst interface{}, pl *datastore.PropertyList) error {
	c := make(chan datastore.Property)
	go func() {
		for _, p := range *pl {
			c <- p
		}
		close(c)
	}()
	return datastore.LoadStruct(dst, c)
}
<commit_msg>Neatened errors variable name.<commit_after>package nds

import (
	"appengine"
	"appengine/datastore"
	"errors"
	"fmt"
	"reflect"
	"sync"
)

const (
	// multiLimit is the App Engine datastore limit for the number of entities
	// that can be PutMulti or GetMulti in one call.
	multiLimit = 1000
)

var (
	// milMultiError is a convenience slice used to represent a nil error when
	// grouping errors in GetMulti.
	nilMultiError = make(appengine.MultiError, multiLimit)
)

func checkMultiArgs(keys []*datastore.Key, v reflect.Value) error {
	if v.Kind() != reflect.Slice {
		return errors.New("nds: dst is not a slice")
	}

	if len(keys) != v.Len() {
		return errors.New("nds: key and dst slices have different length")
	}

	return nil
}

// GetMulti works just like datastore.GetMulti except it removes the API limit
// of 1000 entities per request by calling datastore.GetMulti as many times as
// required to complete the request.
//
// Increase the datastore timeout if you get datastore_v3: TIMEOUT errors. You
// can do this using
// http://godoc.org/code.google.com/p/appengine-go/appengine#Timeout.
func GetMulti(c appengine.Context,
	keys []*datastore.Key, dst interface{}) error {

	v := reflect.ValueOf(dst)
	if err := checkMultiArgs(keys, v); err != nil {
		return err
	}

	if len(keys) == 0 {
		return nil
	}

	p := len(keys) / multiLimit
	errs := make([]error, p+1)
	wg := sync.WaitGroup{}
	for i := 0; i < p; i++ {
		index := i
		keySlice := keys[i*multiLimit : (i+1)*multiLimit]
		dstSlice := v.Slice(i*multiLimit, (i+1)*multiLimit)

		wg.Add(1)
		go func() {
			errs[index] = datastore.GetMulti(c, keySlice, dstSlice.Interface())
			wg.Done()
		}()
	}

	if len(keys)%multiLimit == 0 {
		errs = errs[:len(errs)-1]
	} else {
		keySlice := keys[p*multiLimit : len(keys)]
		dstSlice := v.Slice(p*multiLimit, len(keys))
		wg.Add(1)
		go func() {
			errs[p] = datastore.GetMulti(c, keySlice, dstSlice.Interface())
			wg.Done()
		}()
	}
	wg.Wait()

	// Quick escape if all errors are nil.
	errsNil := true
	for _, err := range errs {
		if err != nil {
			errsNil = false
		}
	}
	if errsNil {
		return nil
	}

	groupedErrs := make(appengine.MultiError, 0, len(keys))
	for _, err := range errs {
		if err == nil {
			groupedErrs = append(groupedErrs, nilMultiError...)
		} else if me, ok := err.(appengine.MultiError); ok {
			groupedErrs = append(groupedErrs, me...)
		} else {
			return err
		}
	}
	return groupedErrs[:len(keys)]
}

type cacheContext struct {
	appengine.Context
	cache map[string]*datastore.PropertyList
	sync.RWMutex
}

func NewCacheContext(c appengine.Context) appengine.Context {
	return &cacheContext{
		Context: c,
		cache:   map[string]*datastore.PropertyList{},
	}
}

func GetMultiCache(c appengine.Context,
	keys []*datastore.Key, dst interface{}) error {

	v := reflect.ValueOf(dst)
	if err := checkMultiArgs(keys, v); err != nil {
		return err
	}

	if cc, ok := c.(*cacheContext); ok {
		return getMultiCache(cc, keys, v)
	} else {
		return datastore.GetMulti(c, keys, dst)
	}
}

func convertToPropertyLists(
	v reflect.Value) ([]datastore.PropertyList, error) {
	pls := make([]datastore.PropertyList, v.Len())
	for i := range pls {
		pl := datastore.PropertyList{}
		elem := addrValue(v.Index(i))
		if err := SaveStruct(elem.Interface(), &pl); err != nil {
			return nil, err
		}
		pls[i] = pl
	}
	return pls, nil
}

func addrValue(v reflect.Value) reflect.Value {
	if v.Kind() != reflect.Ptr {
		return v.Addr()
	} else {
		return v
	}
}

// getMultiCache gets entities from local cache then the datastore.
// dst argument must be a slice.
func getMultiCache(cc *cacheContext,
	keys []*datastore.Key, dst reflect.Value) error {

	cacheMissIndexes := []int{}
	cacheMissKeys := []*datastore.Key{}
	cacheMissDsts := []datastore.PropertyList{}

	errs := make(appengine.MultiError, dst.Len())
	errsNil := true

	// Load what we can from the local cache.
	cc.RLock()
	for i, key := range keys {
		if pl, ok := cc.cache[key.Encode()]; ok {
			if pl == nil {
				errs[i] = datastore.ErrNoSuchEntity
				errsNil = false
			} else {
				elem := addrValue(dst.Index(i))
				if err := LoadStruct(elem.Interface(), pl); err != nil {
					return err
				}
			}
		} else {
			cacheMissIndexes = append(cacheMissIndexes, i)
			cacheMissKeys = append(cacheMissKeys, key)
			cacheMissDsts = append(cacheMissDsts, datastore.PropertyList{})
		}
	}
	cc.RUnlock()

	// Load from datastore.
	if err := datastore.GetMulti(cc, cacheMissKeys, cacheMissDsts); err == nil {
		// Save to local memory cache.
		putMultiLocalCache(cc, cacheMissKeys, cacheMissDsts)

		// Update the callers slice with values.
		for i, index := range cacheMissIndexes {
			pl := cacheMissDsts[i]
			elem := addrValue(dst.Index(index))
			if err := LoadStruct(elem.Interface(), &pl); err != nil {
				return err
			}
		}
	} else if me, ok := err.(appengine.MultiError); ok {
		for i, err := range me {
			if err == nil {
				putLocalCache(cc, cacheMissKeys[i], cacheMissDsts[i])

				// Update the callers slice with values.
				pl := cacheMissDsts[i]
				index := cacheMissIndexes[i]
				elem := addrValue(dst.Index(index))
				if err := LoadStruct(elem.Interface(), &pl); err != nil {
					return err
				}
			} else if err == datastore.ErrNoSuchEntity {
				putLocalCache(cc, cacheMissKeys[i], nil)
				index := cacheMissIndexes[i]
				errs[index] = datastore.ErrNoSuchEntity
				errsNil = false
				// Possibly should zero the callers slice value here.
			} else {
				return err
			}
		}
	} else {
		return err
	}

	if errsNil {
		return nil
	}
	return errs
}

func PutMultiCache(c appengine.Context,
	keys []*datastore.Key, src interface{}) ([]*datastore.Key, error) {

	v := reflect.ValueOf(src)
	if err := checkMultiArgs(keys, v); err != nil {
		return nil, err
	}

	if cc, ok := c.(*cacheContext); ok {
		if pls, err := convertToPropertyLists(v); err != nil {
			fmt.Println("Convert error", err)
			return nil, err
		} else {
			return putMultiCache(cc, keys, pls)
		}
	} else {
		return datastore.PutMulti(c, keys, src)
	}
}

// putMultiCache puts the entities into the datastore and then its local cache.
func putMultiCache(cc *cacheContext,
	keys []*datastore.Key,
	pls []datastore.PropertyList) ([]*datastore.Key, error) {

	// Save to the datastore.
	completeKeys, err := datastore.PutMulti(cc, keys, pls)
	if err != nil {
		return nil, err
	}

	// Save to local memory cache.
	putMultiLocalCache(cc, completeKeys, pls)

	return completeKeys, nil
}

func putLocalCache(cc *cacheContext,
	key *datastore.Key, pl datastore.PropertyList) {
	cc.Lock()
	cc.cache[key.Encode()] = &pl
	cc.Unlock()
}

func putMultiLocalCache(cc *cacheContext,
	keys []*datastore.Key, pls []datastore.PropertyList) {
	for i, key := range keys {
		putLocalCache(cc, key, pls[i])
	}
}

// SaveStruct saves src to a datastore.PropertyList.
func SaveStruct(src interface{}, pl *datastore.PropertyList) error {
	c, err := make(chan datastore.Property), make(chan error)
	go func() {
		err <- datastore.SaveStruct(src, c)
	}()
	for p := range c {
		*pl = append(*pl, p)
	}
	return <-err
}

// LoadStruct loads a datastore.PropertyList into dst.
func LoadStruct(dst interface{}, pl *datastore.PropertyList) error {
	c := make(chan datastore.Property)
	go func() {
		for _, p := range *pl {
			c <- p
		}
		close(c)
	}()
	return datastore.LoadStruct(dst, c)
}
<|endoftext|>


<commit_before>package db

import (
	"database/sql"
	"database/sql/driver"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"sync"
	"time"

	"github.com/lib/pq"

	"golang.org/x/crypto/ssh"
)

// Credentials represent credentials for connecting to a postgres instance via a jumpbox
type Credentials struct {
	DB            string
	Username      string
	Port          string
	Password      string
	Address       string
	CACert        string
	SSHPrivateKey []byte
	SSHPublicIP   string
}

// Runner is function that runs SQL over a jumpbox
type Runner func(sql string) error

func writeTempFile(dir, prefix string, data []byte) (name string, err error) {
	f, err := ioutil.TempFile(dir, prefix)
	if err != nil {
		return "", err
	}
	name = f.Name()
	_, err = f.Write(data)
	if err1 := f.Close(); err == nil {
		err = err1
	}
	if err != nil {
		os.Remove(name)
	}
	return name, err
}

// NewRunner returns a new SQL runner
func NewRunner(creds *Credentials) (Runner, error) {
	key, err := ssh.ParsePrivateKey(creds.SSHPrivateKey)
	if err != nil {
		return nil, err
	}

	sshConfig := &ssh.ClientConfig{
		User: "vcap",
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(key),
		},
	}

	var once sync.Once
	var db *sql.DB
	return func(sqlStr string) error {
		var err error
		once.Do(func() {
			var sshClient *ssh.Client
			sshClient, err = ssh.Dial("tcp", fmt.Sprintf("%s:22", creds.SSHPublicIP), sshConfig)
			if err != nil {
				return
			}

			var caCertFileName string
			caCertFileName, err = writeTempFile("", "concourse-up", []byte(creds.CACert))
			if err != nil {
				return
			}
			defer os.Remove(caCertFileName)
			dialer := &sshDialer{client: sshClient}

			sql.Register("postgres+ssh", dialer)

			db, err = sql.Open("postgres+ssh", postgresArgs(creds, caCertFileName))
			if err != nil {
				return
			}
		})
		if err != nil {
			return err
		}
		rows, err := db.Query(sqlStr)
		if err != nil {
			return err
		}
		if err = rows.Err(); err != nil {
			return err
		}
		return rows.Close()
	}, nil
}

func postgresArgs(creds *Credentials, caCertPath string) string {
	return fmt.Sprintf(
		"user=%s password=%s dbname=%s host=%s port=%s sslmode=verify-full sslrootcert=%s",
		creds.Username,
		creds.Password,
		creds.DB,
		creds.Address,
		creds.Port,
		caCertPath,
	)
}

type sshDialer struct {
	client *ssh.Client
}

func (dialer *sshDialer) Open(s string) (_ driver.Conn, err error) {
	return pq.DialOpen(dialer, s)
}

func (dialer *sshDialer) Dial(network, address string) (net.Conn, error) {
	return dialer.client.Dial(network, address)
}

func (dialer *sshDialer) DialTimeout(network, address string, timeout time.Duration) (net.Conn, error) {
	return dialer.client.Dial(network, address)
}
<commit_msg>use InsecureIgnoreHostKey<commit_after>package db

import (
	"database/sql"
	"database/sql/driver"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"sync"
	"time"

	"github.com/lib/pq"

	"golang.org/x/crypto/ssh"
)

// Credentials represent credentials for connecting to a postgres instance via a jumpbox
type Credentials struct {
	DB            string
	Username      string
	Port          string
	Password      string
	Address       string
	CACert        string
	SSHPrivateKey []byte
	SSHPublicIP   string
}

// Runner is function that runs SQL over a jumpbox
type Runner func(sql string) error

func writeTempFile(dir, prefix string, data []byte) (name string, err error) {
	f, err := ioutil.TempFile(dir, prefix)
	if err != nil {
		return "", err
	}
	name = f.Name()
	_, err = f.Write(data)
	if err1 := f.Close(); err == nil {
		err = err1
	}
	if err != nil {
		os.Remove(name)
	}
	return name, err
}

// NewRunner returns a new SQL runner
func NewRunner(creds *Credentials) (Runner, error) {
	key, err := ssh.ParsePrivateKey(creds.SSHPrivateKey)
	if err != nil {
		return nil, err
	}

	sshConfig := &ssh.ClientConfig{
		User: "vcap",
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(key),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), //TODO: probably don't do this
	}

	var once sync.Once
	var db *sql.DB
	return func(sqlStr string) error {
		var err error
		once.Do(func() {
			var sshClient *ssh.Client
			sshClient, err = ssh.Dial("tcp", fmt.Sprintf("%s:22", creds.SSHPublicIP), sshConfig)
			if err != nil {
				return
			}

			var caCertFileName string
			caCertFileName, err = writeTempFile("", "concourse-up", []byte(creds.CACert))
			if err != nil {
				return
			}
			defer os.Remove(caCertFileName)
			dialer := &sshDialer{client: sshClient}

			sql.Register("postgres+ssh", dialer)

			db, err = sql.Open("postgres+ssh", postgresArgs(creds, caCertFileName))
			if err != nil {
				return
			}
		})
		if err != nil {
			return err
		}
		rows, err := db.Query(sqlStr)
		if err != nil {
			return err
		}
		if err = rows.Err(); err != nil {
			return err
		}
		return rows.Close()
	}, nil
}

func postgresArgs(creds *Credentials, caCertPath string) string {
	return fmt.Sprintf(
		"user=%s password=%s dbname=%s host=%s port=%s sslmode=verify-full sslrootcert=%s",
		creds.Username,
		creds.Password,
		creds.DB,
		creds.Address,
		creds.Port,
		caCertPath,
	)
}

type sshDialer struct {
	client *ssh.Client
}

func (dialer *sshDialer) Open(s string) (_ driver.Conn, err error) {
	return pq.DialOpen(dialer, s)
}

func (dialer *sshDialer) Dial(network, address string) (net.Conn, error) {
	return dialer.client.Dial(network, address)
}

func (dialer *sshDialer) DialTimeout(network, address string, timeout time.Duration) (net.Conn, error) {
	return dialer.client.Dial(network, address)
}
<|endoftext|>


<commit_before>package db

import (
	"bytes"
	"encoding/binary"
	"fmt"

	"bazil.org/bazil/fs/clock"
	wirepeer "bazil.org/bazil/peer/wire"
	"github.com/boltdb/bolt"
	"github.com/golang/protobuf/proto"
)

// VolumeConflicts tracks the alternate versions of directory
// entries.
type VolumeConflicts struct {
	b *bolt.Bucket
}

func (VolumeConflicts) pathToKey(parentInode uint64, name string, clock []byte) []byte {
	buf := make([]byte, 8, 8+len(name)+1+len(clock))
	binary.BigEndian.PutUint64(buf, parentInode)
	buf = append(buf, name...)
	buf = append(buf, '\x00')
	buf = append(buf, clock...)
	return buf
}

func (VolumeConflicts) dirToKeyPrefix(parentInode uint64) []byte {
	buf := make([]byte, 8)
	binary.BigEndian.PutUint64(buf, parentInode)
	return buf
}

func (vc *VolumeConflicts) Add(parentInode uint64, c *clock.Clock, de *wirepeer.Dirent) error {
	clockBuf, err := c.MarshalBinary()
	if err != nil {
		return fmt.Errorf("error marshaling clock: %v", err)
	}
	key := vc.pathToKey(parentInode, de.Name, clockBuf)

	tmp := *de
	tmp.Name = ""
	tmp.Clock = nil
	buf, err := proto.Marshal(&tmp)
	if err != nil {
		return fmt.Errorf("error marshaling dirent: %v", err)
	}

	if err := vc.b.Put(key, buf); err != nil {
		return err
	}
	return nil
}

func (vc *VolumeConflicts) List(parentInode uint64, name string) *VolumeConflictsCursor {
	c := vc.b.Cursor()
	prefix := vc.pathToKey(parentInode, name, nil)
	return &VolumeConflictsCursor{
		prefix: prefix,
		c:      c,
	}
}

// ListAll iterates over all of the conflict entries for this directory.
func (vc *VolumeConflicts) ListAll(parentInode uint64) *VolumeConflictsCursor {
	c := vc.b.Cursor()
	prefix := vc.dirToKeyPrefix(parentInode)
	return &VolumeConflictsCursor{
		prefix: prefix,
		c:      c,
	}
}

func (vc *VolumeConflicts) Get(parentInode uint64, name string, clockBuf []byte) *VolumeConflictsItem {
	k := vc.pathToKey(parentInode, name, clockBuf)
	v := vc.b.Get(k)
	if v == nil {
		return nil
	}
	return &VolumeConflictsItem{
		name:  k[8 : 8+len(name)],
		clock: k[8+len(name)+1:],
		data:  v,
	}
}

type VolumeConflictsCursor struct {
	prefix []byte
	c      *bolt.Cursor
}

func (c *VolumeConflictsCursor) First() *VolumeConflictsItem {
	return c.item(c.c.Seek(c.prefix))
}

func (c *VolumeConflictsCursor) Next() *VolumeConflictsItem {
	return c.item(c.c.Next())
}

// Delete the current item.
func (c *VolumeConflictsCursor) Delete() error {
	return c.c.Delete()
}

func (c *VolumeConflictsCursor) item(k, v []byte) *VolumeConflictsItem {
	if !bytes.HasPrefix(k, c.prefix) {
		// past the end of the dirent for List, or dir for ListAll
		return nil
	}
	name := k[8:]
	idx := bytes.IndexByte(name, '\x00')
	if idx == -1 {
		// corrupt entry?
		return nil
	}
	name = name[:idx]
	clock := k[8+idx+1:]
	return &VolumeConflictsItem{
		name:  name,
		clock: clock,
		data:  v,
	}
}

type VolumeConflictsItem struct {
	name  []byte
	clock []byte
	data  []byte
}

// Name returns the file name for this item.
//
// This is mostly useful when used with ListAll.
//
// Returned value is valid after the transaction.
func (item *VolumeConflictsItem) Name() string {
	return string(item.name)
}

// Clock returns the clock for this item.
//
// Returned value is valid after the transaction.
func (item *VolumeConflictsItem) Clock() (*clock.Clock, error) {
	var c clock.Clock
	if err := c.UnmarshalBinary(item.clock); err != nil {
		return nil, fmt.Errorf("error unmarshaling clock: %v", err)
	}
	return &c, nil
}

// Dirent returns the directory entry for this item.
//
// out is valid after the transaction.
func (item *VolumeConflictsItem) Dirent(out *wirepeer.Dirent) error {
	if err := proto.Unmarshal(item.data, out); err != nil {
		return fmt.Errorf("error unmarshaling dirent: %v", err)
	}
	return nil
}
<commit_msg>db: Support deleting conflict entries<commit_after>package db

import (
	"bytes"
	"encoding/binary"
	"fmt"

	"bazil.org/bazil/fs/clock"
	wirepeer "bazil.org/bazil/peer/wire"
	"bazil.org/fuse"
	"github.com/boltdb/bolt"
	"github.com/golang/protobuf/proto"
)

// VolumeConflicts tracks the alternate versions of directory
// entries.
type VolumeConflicts struct {
	b *bolt.Bucket
}

func (VolumeConflicts) pathToKey(parentInode uint64, name string, clock []byte) []byte {
	buf := make([]byte, 8, 8+len(name)+1+len(clock))
	binary.BigEndian.PutUint64(buf, parentInode)
	buf = append(buf, name...)
	buf = append(buf, '\x00')
	buf = append(buf, clock...)
	return buf
}

func (VolumeConflicts) dirToKeyPrefix(parentInode uint64) []byte {
	buf := make([]byte, 8)
	binary.BigEndian.PutUint64(buf, parentInode)
	return buf
}

func (vc *VolumeConflicts) Add(parentInode uint64, c *clock.Clock, de *wirepeer.Dirent) error {
	clockBuf, err := c.MarshalBinary()
	if err != nil {
		return fmt.Errorf("error marshaling clock: %v", err)
	}
	key := vc.pathToKey(parentInode, de.Name, clockBuf)

	tmp := *de
	tmp.Name = ""
	tmp.Clock = nil
	buf, err := proto.Marshal(&tmp)
	if err != nil {
		return fmt.Errorf("error marshaling dirent: %v", err)
	}

	if err := vc.b.Put(key, buf); err != nil {
		return err
	}
	return nil
}

func (vc *VolumeConflicts) List(parentInode uint64, name string) *VolumeConflictsCursor {
	c := vc.b.Cursor()
	prefix := vc.pathToKey(parentInode, name, nil)
	return &VolumeConflictsCursor{
		prefix: prefix,
		c:      c,
	}
}

// ListAll iterates over all of the conflict entries for this directory.
func (vc *VolumeConflicts) ListAll(parentInode uint64) *VolumeConflictsCursor {
	c := vc.b.Cursor()
	prefix := vc.dirToKeyPrefix(parentInode)
	return &VolumeConflictsCursor{
		prefix: prefix,
		c:      c,
	}
}

func (vc *VolumeConflicts) Get(parentInode uint64, name string, clockBuf []byte) *VolumeConflictsItem {
	k := vc.pathToKey(parentInode, name, clockBuf)
	v := vc.b.Get(k)
	if v == nil {
		return nil
	}
	return &VolumeConflictsItem{
		name:  k[8 : 8+len(name)],
		clock: k[8+len(name)+1:],
		data:  v,
	}
}

func (vc *VolumeConflicts) Delete(parentInode uint64, name string, clockBuf []byte) error {
	k := vc.pathToKey(parentInode, name, clockBuf)
	v := vc.b.Get(k)
	if v == nil {
		return fuse.ENOENT
	}
	if err := vc.b.Delete(k); err != nil {
		return err
	}
	return nil
}

type VolumeConflictsCursor struct {
	prefix []byte
	c      *bolt.Cursor
}

func (c *VolumeConflictsCursor) First() *VolumeConflictsItem {
	return c.item(c.c.Seek(c.prefix))
}

func (c *VolumeConflictsCursor) Next() *VolumeConflictsItem {
	return c.item(c.c.Next())
}

// Delete the current item.
func (c *VolumeConflictsCursor) Delete() error {
	return c.c.Delete()
}

func (c *VolumeConflictsCursor) item(k, v []byte) *VolumeConflictsItem {
	if !bytes.HasPrefix(k, c.prefix) {
		// past the end of the dirent for List, or dir for ListAll
		return nil
	}
	name := k[8:]
	idx := bytes.IndexByte(name, '\x00')
	if idx == -1 {
		// corrupt entry?
		return nil
	}
	name = name[:idx]
	clock := k[8+idx+1:]
	return &VolumeConflictsItem{
		name:  name,
		clock: clock,
		data:  v,
	}
}

type VolumeConflictsItem struct {
	name  []byte
	clock []byte
	data  []byte
}

// Name returns the file name for this item.
//
// This is mostly useful when used with ListAll.
//
// Returned value is valid after the transaction.
func (item *VolumeConflictsItem) Name() string {
	return string(item.name)
}

// Clock returns the clock for this item.
//
// Returned value is valid after the transaction.
func (item *VolumeConflictsItem) Clock() (*clock.Clock, error) {
	var c clock.Clock
	if err := c.UnmarshalBinary(item.clock); err != nil {
		return nil, fmt.Errorf("error unmarshaling clock: %v", err)
	}
	return &c, nil
}

// Dirent returns the directory entry for this item.
//
// out is valid after the transaction.
func (item *VolumeConflictsItem) Dirent(out *wirepeer.Dirent) error {
	if err := proto.Unmarshal(item.data, out); err != nil {
		return fmt.Errorf("error unmarshaling dirent: %v", err)
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Current limitations:

	- GSS-API authentication is not supported
	- only SOCKS version 5 is supported
	- TCP bind and UDP not yet supported

Example http client over SOCKS5:

	proxy := &socks.Proxy{"127.0.0.1:1080"}
	tr := &http.Transport{
		Dial: func(net, addr string) (net.Conn, error) {
			return proxy.Dial(net, addr)
		},
	}
	client := &http.Client{Transport: tr}
	resp, err := client.Get("https://example.com")
*/
package socks

import (
	"errors"
	"io"
	"net"
	"strconv"
)

const (
	protocolVersion = 5

	defaultPort = 1080

	authNone             = 0
	authGssApi           = 1
	authUsernamePassword = 2
	authUnavailable      = 0xff

	commandTcpConnect   = 1
	commandTcpBind      = 2
	commandUdpAssociate = 3

	addressTypeIPv4   = 1
	addressTypeDomain = 3
	addressTypeIPv6   = 4

	statusRequestGranted          = 0
	statusGeneralFailure          = 1
	statusConnectionNotAllowed    = 2
	statusNetworkUnreachable      = 3
	statusHostUnreachable         = 4
	statusConnectionRefused       = 5
	statusTtlExpired              = 6
	statusCommandNotSupport       = 7
	statusAddressTypeNotSupported = 8
)

var (
	ErrAuthFailed             = errors.New("authentication failed")
	ErrInvalidProxyResponse   = errors.New("invalid proxy response")
	ErrNoAcceptableAuthMethod = errors.New("no acceptable authentication method")

	statusErrors = map[byte]error{
		statusGeneralFailure:          errors.New("general failure"),
		statusConnectionNotAllowed:    errors.New("connection not allowed by ruleset"),
		statusNetworkUnreachable:      errors.New("network unreachable"),
		statusHostUnreachable:         errors.New("host unreachable"),
		statusConnectionRefused:       errors.New("connection refused by destination host"),
		statusTtlExpired:              errors.New("TTL expired"),
		statusCommandNotSupport:       errors.New("command not supported / protocol error"),
		statusAddressTypeNotSupported: errors.New("address type not supported"),
	}
)

type Proxy struct {
	Addr     string
	Username string
	Password string
}

func (p *Proxy) Dial(network, addr string) (net.Conn, error) {
	host, strPort, err := net.SplitHostPort(addr)
	if err != nil {
		return nil, err
	}
	port, err := strconv.Atoi(strPort)
	if err != nil {
		return nil, err
	}

	conn, err := net.Dial("tcp", p.Addr)
	if err != nil {
		return nil, err
	}

	buf := make([]byte, 16+len(p.Username)+len(p.Password))

	// Initial greeting

	buf[0] = protocolVersion
	if p.Username != "" {
		buf = buf[:4]
		buf[1] = 2 // num auth methods
		buf[2] = authNone
		buf[3] = authUsernamePassword
	} else {
		buf = buf[:3]
		buf[1] = 1 // num auth methods
		buf[2] = authNone
	}

	_, err = conn.Write(buf)
	if err != nil {
		conn.Close()
		return nil, err
	}

	// Server's auth choice

	if _, err := io.ReadFull(conn, buf[:2]); err != nil {
		conn.Close()
		return nil, err
	}
	if buf[0] != protocolVersion {
		conn.Close()
		return nil, ErrInvalidProxyResponse
	}
	err = nil
	switch buf[1] {
	default:
		err = ErrInvalidProxyResponse
	case authUnavailable:
		err = ErrNoAcceptableAuthMethod
	case authGssApi:
		err = ErrNoAcceptableAuthMethod
	case authUsernamePassword:
		buf = buf[:3+len(p.Username)+len(p.Password)]
		buf[0] = 1 // version
		buf[1] = byte(len(p.Username))
		copy(buf[2:], p.Username)
		buf[2+len(p.Username)] = byte(len(p.Password))
		copy(buf[3+len(p.Username):], p.Password)
		if _, err = conn.Write(buf); err != nil {
			conn.Close()
			return nil, err
		}
		if _, err = io.ReadFull(conn, buf[:2]); err != nil {
			conn.Close()
			return nil, err
		}
		if buf[0] != 1 { // version
			err = ErrInvalidProxyResponse
		} else if buf[1] != 0 { // 0 = succes, else auth failed
			err = ErrAuthFailed
		}
	case authNone:
		// Do nothing
	}
	if err != nil {
		conn.Close()
		return nil, err
	}

	// Command / connection request

	buf = buf[:7+len(host)]
	buf[0] = protocolVersion
	buf[1] = commandTcpConnect
	buf[2] = 0 // reserved
	buf[3] = addressTypeDomain
	buf[4] = byte(len(host))
	copy(buf[5:], host)
	buf[5+len(host)] = byte(port >> 8)
	buf[6+len(host)] = byte(port & 0xff)
	if _, err := conn.Write(buf); err != nil {
		conn.Close()
		return nil, err
	}

	// Server response

	if _, err := io.ReadFull(conn, buf[:4]); err != nil {
		conn.Close()
		return nil, err
	}

	if buf[0] != protocolVersion {
		conn.Close()
		return nil, ErrInvalidProxyResponse
	}

	if buf[1] != statusRequestGranted {
		conn.Close()
		err := statusErrors[buf[1]]
		if err == nil {
			err = ErrInvalidProxyResponse
		}
		return nil, err
	}

	paddr := &proxiedAddr{net: network}

	switch buf[3] {
	default:
		conn.Close()
		return nil, ErrInvalidProxyResponse
	case addressTypeIPv4:
		if _, err := io.ReadFull(conn, buf[:4]); err != nil {
			conn.Close()
			return nil, err
		}
		paddr.host = net.IP(buf).String()
	case addressTypeIPv6:
		if _, err := io.ReadFull(conn, buf[:16]); err != nil {
			conn.Close()
			return nil, err
		}
		paddr.host = net.IP(buf).String()
	case addressTypeDomain:
		if _, err := io.ReadFull(conn, buf[:1]); err != nil {
			conn.Close()
			return nil, err
		}
		domainLen := buf[0]
		if _, err := io.ReadFull(conn, buf[:domainLen]); err != nil {
			conn.Close()
			return nil, err
		}
		paddr.host = string(buf[:domainLen])
	}

	if _, err := io.ReadFull(conn, buf[:2]); err != nil {
		conn.Close()
		return nil, err
	}
	paddr.port = int(buf[0])<<8 | int(buf[1])

	return &proxiedConn{
		conn:       conn,
		boundAddr:  paddr,
		remoteAddr: &proxiedAddr{network, host, port},
	}, nil
}
<commit_msg>Bug fix in Dial<commit_after>/*
Current limitations:

	- GSS-API authentication is not supported
	- only SOCKS version 5 is supported
	- TCP bind and UDP not yet supported

Example http client over SOCKS5:

	proxy := &socks.Proxy{"127.0.0.1:1080"}
	tr := &http.Transport{
		Dial: func(net, addr string) (net.Conn, error) {
			return proxy.Dial(net, addr)
		},
	}
	client := &http.Client{Transport: tr}
	resp, err := client.Get("https://example.com")
*/
package socks

import (
	"errors"
	"io"
	"net"
	"strconv"
)

const (
	protocolVersion = 5

	defaultPort = 1080

	authNone             = 0
	authGssApi           = 1
	authUsernamePassword = 2
	authUnavailable      = 0xff

	commandTcpConnect   = 1
	commandTcpBind      = 2
	commandUdpAssociate = 3

	addressTypeIPv4   = 1
	addressTypeDomain = 3
	addressTypeIPv6   = 4

	statusRequestGranted          = 0
	statusGeneralFailure          = 1
	statusConnectionNotAllowed    = 2
	statusNetworkUnreachable      = 3
	statusHostUnreachable         = 4
	statusConnectionRefused       = 5
	statusTtlExpired              = 6
	statusCommandNotSupport       = 7
	statusAddressTypeNotSupported = 8
)

var (
	ErrAuthFailed             = errors.New("authentication failed")
	ErrInvalidProxyResponse   = errors.New("invalid proxy response")
	ErrNoAcceptableAuthMethod = errors.New("no acceptable authentication method")

	statusErrors = map[byte]error{
		statusGeneralFailure:          errors.New("general failure"),
		statusConnectionNotAllowed:    errors.New("connection not allowed by ruleset"),
		statusNetworkUnreachable:      errors.New("network unreachable"),
		statusHostUnreachable:         errors.New("host unreachable"),
		statusConnectionRefused:       errors.New("connection refused by destination host"),
		statusTtlExpired:              errors.New("TTL expired"),
		statusCommandNotSupport:       errors.New("command not supported / protocol error"),
		statusAddressTypeNotSupported: errors.New("address type not supported"),
	}
)

type Proxy struct {
	Addr     string
	Username string
	Password string
}

func (p *Proxy) Dial(network, addr string) (net.Conn, error) {
	host, strPort, err := net.SplitHostPort(addr)
	if err != nil {
		return nil, err
	}
	port, err := strconv.Atoi(strPort)
	if err != nil {
		return nil, err
	}

	conn, err := net.Dial("tcp", p.Addr)
	if err != nil {
		return nil, err
	}

	buf := make([]byte, 32+len(host)+len(p.Username)+len(p.Password))

	// Initial greeting

	buf[0] = protocolVersion
	if p.Username != "" {
		buf = buf[:4]
		buf[1] = 2 // num auth methods
		buf[2] = authNone
		buf[3] = authUsernamePassword
	} else {
		buf = buf[:3]
		buf[1] = 1 // num auth methods
		buf[2] = authNone
	}

	_, err = conn.Write(buf)
	if err != nil {
		conn.Close()
		return nil, err
	}

	// Server's auth choice

	if _, err := io.ReadFull(conn, buf[:2]); err != nil {
		conn.Close()
		return nil, err
	}
	if buf[0] != protocolVersion {
		conn.Close()
		return nil, ErrInvalidProxyResponse
	}
	err = nil
	switch buf[1] {
	default:
		err = ErrInvalidProxyResponse
	case authUnavailable:
		err = ErrNoAcceptableAuthMethod
	case authGssApi:
		err = ErrNoAcceptableAuthMethod
	case authUsernamePassword:
		buf = buf[:3+len(p.Username)+len(p.Password)]
		buf[0] = 1 // version
		buf[1] = byte(len(p.Username))
		copy(buf[2:], p.Username)
		buf[2+len(p.Username)] = byte(len(p.Password))
		copy(buf[3+len(p.Username):], p.Password)
		if _, err = conn.Write(buf); err != nil {
			conn.Close()
			return nil, err
		}
		if _, err = io.ReadFull(conn, buf[:2]); err != nil {
			conn.Close()
			return nil, err
		}
		if buf[0] != 1 { // version
			err = ErrInvalidProxyResponse
		} else if buf[1] != 0 { // 0 = succes, else auth failed
			err = ErrAuthFailed
		}
	case authNone:
		// Do nothing
	}
	if err != nil {
		conn.Close()
		return nil, err
	}

	// Command / connection request

	buf = buf[:7+len(host)]
	buf[0] = protocolVersion
	buf[1] = commandTcpConnect
	buf[2] = 0 // reserved
	buf[3] = addressTypeDomain
	buf[4] = byte(len(host))
	copy(buf[5:], host)
	buf[5+len(host)] = byte(port >> 8)
	buf[6+len(host)] = byte(port & 0xff)
	if _, err := conn.Write(buf); err != nil {
		conn.Close()
		return nil, err
	}

	// Server response

	if _, err := io.ReadFull(conn, buf[:4]); err != nil {
		conn.Close()
		return nil, err
	}

	if buf[0] != protocolVersion {
		conn.Close()
		return nil, ErrInvalidProxyResponse
	}

	if buf[1] != statusRequestGranted {
		conn.Close()
		err := statusErrors[buf[1]]
		if err == nil {
			err = ErrInvalidProxyResponse
		}
		return nil, err
	}

	paddr := &proxiedAddr{net: network}

	switch buf[3] {
	default:
		conn.Close()
		return nil, ErrInvalidProxyResponse
	case addressTypeIPv4:
		if _, err := io.ReadFull(conn, buf[:4]); err != nil {
			conn.Close()
			return nil, err
		}
		paddr.host = net.IP(buf).String()
	case addressTypeIPv6:
		if _, err := io.ReadFull(conn, buf[:16]); err != nil {
			conn.Close()
			return nil, err
		}
		paddr.host = net.IP(buf).String()
	case addressTypeDomain:
		if _, err := io.ReadFull(conn, buf[:1]); err != nil {
			conn.Close()
			return nil, err
		}
		domainLen := buf[0]
		if _, err := io.ReadFull(conn, buf[:domainLen]); err != nil {
			conn.Close()
			return nil, err
		}
		paddr.host = string(buf[:domainLen])
	}

	if _, err := io.ReadFull(conn, buf[:2]); err != nil {
		conn.Close()
		return nil, err
	}
	paddr.port = int(buf[0])<<8 | int(buf[1])

	return &proxiedConn{
		conn:       conn,
		boundAddr:  paddr,
		remoteAddr: &proxiedAddr{network, host, port},
	}, nil
}
<|endoftext|>