<commit_before>/*
   Copyright The containerd 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 command

import (
	gocontext "context"
	"io"
	"io/ioutil"
	"net"
	"os"
	"time"

	eventsapi "github.com/containerd/containerd/api/services/events/v1"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/namespaces"
	"github.com/containerd/containerd/pkg/dialer"
	"github.com/gogo/protobuf/types"
	"github.com/pkg/errors"
	"github.com/urfave/cli"
	"google.golang.org/grpc"
)

var publishCommand = cli.Command{
	Name:  "publish",
	Usage: "binary to publish events to containerd",
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "namespace",
			Usage: "namespace to publish to",
		},
		cli.StringFlag{
			Name:  "topic",
			Usage: "topic of the event",
		},
	},
	Action: func(context *cli.Context) error {
		ctx := namespaces.WithNamespace(gocontext.Background(), context.String("namespace"))
		topic := context.String("topic")
		if topic == "" {
			return errors.New("topic required to publish event")
		}
		payload, err := getEventPayload(os.Stdin)
		if err != nil {
			return err
		}
		client, err := connectEvents(context.GlobalString("address"))
		if err != nil {
			return err
		}
		if _, err := client.Publish(ctx, &eventsapi.PublishRequest{
			Topic: topic,
			Event: payload,
		}); err != nil {
			return errdefs.FromGRPC(err)
		}
		return nil
	},
}

func getEventPayload(r io.Reader) (*types.Any, error) {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, err
	}
	var any types.Any
	if err := any.Unmarshal(data); err != nil {
		return nil, err
	}
	return &any, nil
}

func connectEvents(address string) (eventsapi.EventsClient, error) {
	conn, err := connect(address, dialer.Dialer)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to dial %q", address)
	}
	return eventsapi.NewEventsClient(conn), nil
}

func connect(address string, d func(string, time.Duration) (net.Conn, error)) (*grpc.ClientConn, error) {
	gopts := []grpc.DialOption{
		grpc.WithBlock(),
		grpc.WithInsecure(),
		grpc.WithDialer(d),
		grpc.FailOnNonTempDialError(true),
		grpc.WithBackoffMaxDelay(3 * time.Second),
	}
	ctx, cancel := gocontext.WithTimeout(gocontext.Background(), 60*time.Second)
	defer cancel()
	conn, err := grpc.DialContext(ctx, dialer.DialAddress(address), gopts...)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to dial %q", address)
	}
	return conn, nil
}
<commit_msg>Reduce publish connection timeout<commit_after>/*
   Copyright The containerd 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 command

import (
	gocontext "context"
	"io"
	"io/ioutil"
	"net"
	"os"
	"time"

	eventsapi "github.com/containerd/containerd/api/services/events/v1"
	"github.com/containerd/containerd/errdefs"
	"github.com/containerd/containerd/namespaces"
	"github.com/containerd/containerd/pkg/dialer"
	"github.com/gogo/protobuf/types"
	"github.com/pkg/errors"
	"github.com/urfave/cli"
	"google.golang.org/grpc"
)

var publishCommand = cli.Command{
	Name:  "publish",
	Usage: "binary to publish events to containerd",
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "namespace",
			Usage: "namespace to publish to",
		},
		cli.StringFlag{
			Name:  "topic",
			Usage: "topic of the event",
		},
	},
	Action: func(context *cli.Context) error {
		ctx := namespaces.WithNamespace(gocontext.Background(), context.String("namespace"))
		topic := context.String("topic")
		if topic == "" {
			return errors.New("topic required to publish event")
		}
		payload, err := getEventPayload(os.Stdin)
		if err != nil {
			return err
		}
		client, err := connectEvents(context.GlobalString("address"))
		if err != nil {
			return err
		}
		if _, err := client.Publish(ctx, &eventsapi.PublishRequest{
			Topic: topic,
			Event: payload,
		}); err != nil {
			return errdefs.FromGRPC(err)
		}
		return nil
	},
}

func getEventPayload(r io.Reader) (*types.Any, error) {
	data, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, err
	}
	var any types.Any
	if err := any.Unmarshal(data); err != nil {
		return nil, err
	}
	return &any, nil
}

func connectEvents(address string) (eventsapi.EventsClient, error) {
	conn, err := connect(address, dialer.Dialer)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to dial %q", address)
	}
	return eventsapi.NewEventsClient(conn), nil
}

func connect(address string, d func(string, time.Duration) (net.Conn, error)) (*grpc.ClientConn, error) {
	gopts := []grpc.DialOption{
		grpc.WithBlock(),
		grpc.WithInsecure(),
		grpc.WithDialer(d),
		grpc.FailOnNonTempDialError(true),
		grpc.WithBackoffMaxDelay(3 * time.Second),
	}
	ctx, cancel := gocontext.WithTimeout(gocontext.Background(), 2*time.Second)
	defer cancel()
	conn, err := grpc.DialContext(ctx, dialer.DialAddress(address), gopts...)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to dial %q", address)
	}
	return conn, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes 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 cmd

import (
	"io"
	"time"

	"github.com/spf13/cobra"

	"k8s.io/klog/v2"

	kubeadmapiv1 "k8s.io/kubernetes/cmd/kubeadm/app/apis/kubeadm/v1beta3"
	"k8s.io/kubernetes/cmd/kubeadm/app/cmd/options"
	cmdutil "k8s.io/kubernetes/cmd/kubeadm/app/cmd/util"
	kubeadmconstants "k8s.io/kubernetes/cmd/kubeadm/app/constants"
	kubeconfigphase "k8s.io/kubernetes/cmd/kubeadm/app/phases/kubeconfig"
	configutil "k8s.io/kubernetes/cmd/kubeadm/app/util/config"
)

var (
	kubeconfigLongDesc = cmdutil.LongDesc(`
	Kubeconfig file utilities.
	`)

	userKubeconfigLongDesc = cmdutil.LongDesc(`
	Output a kubeconfig file for an additional user.
	`)

	userKubeconfigExample = cmdutil.Examples(`
	# Output a kubeconfig file for an additional user named foo using a kubeadm config file bar
	kubeadm kubeconfig user --client-name=foo --config=bar
	`)
)

// newCmdKubeConfigUtility returns main command for kubeconfig phase
func newCmdKubeConfigUtility(out io.Writer) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "kubeconfig",
		Short: "Kubeconfig file utilities",
		Long:  kubeconfigLongDesc,
	}

	cmd.AddCommand(newCmdUserKubeConfig(out))
	return cmd
}

// newCmdUserKubeConfig returns sub commands for kubeconfig phase
func newCmdUserKubeConfig(out io.Writer) *cobra.Command {

	initCfg := cmdutil.DefaultInitConfiguration()
	clusterCfg := &kubeadmapiv1.ClusterConfiguration{}

	var (
		token, clientName, cfgPath string
		organizations              []string
		validityPeriod             time.Duration
	)

	// Creates the UX Command
	cmd := &cobra.Command{
		Use:     "user",
		Short:   "Output a kubeconfig file for an additional user",
		Long:    userKubeconfigLongDesc,
		Example: userKubeconfigExample,
		RunE: func(cmd *cobra.Command, args []string) error {
			// This call returns the ready-to-use configuration based on the defaults populated by flags
			internalCfg, err := configutil.LoadOrDefaultInitConfiguration(cfgPath, initCfg, clusterCfg)
			if err != nil {
				return err
			}

			if validityPeriod > kubeadmconstants.CertificateValidity {
				klog.Warningf("WARNING: the specified certificate validity period %v is longer than the default duration %v, this may increase security risks.",
					validityPeriod, kubeadmconstants.CertificateValidity)
			}

			notAfter := time.Now().Add(validityPeriod).UTC()

			// if the kubeconfig file for an additional user has to use a token, use it
			if token != "" {
				return kubeconfigphase.WriteKubeConfigWithToken(out, internalCfg, clientName, token, &notAfter)
			}

			// Otherwise, write a kubeconfig file with a generate client cert
			return kubeconfigphase.WriteKubeConfigWithClientCert(out, internalCfg, clientName, organizations, &notAfter)
		},
		Args: cobra.NoArgs,
	}

	options.AddConfigFlag(cmd.Flags(), &cfgPath)

	// Add command specific flags
	cmd.Flags().StringVar(&token, options.TokenStr, token, "The token that should be used as the authentication mechanism for this kubeconfig, instead of client certificates")
	cmd.Flags().StringVar(&clientName, "client-name", clientName, "The name of user. It will be used as the CN if client certificates are created")
	cmd.Flags().StringSliceVar(&organizations, "org", organizations, "The orgnizations of the client certificate. It will be used as the O if client certificates are created")
	cmd.Flags().DurationVar(&validityPeriod, "validity-period", kubeadmconstants.CertificateValidity, "The validity period of the client certificate. It is an offset from the current time.")

	cmd.MarkFlagRequired(options.CfgPath)
	cmd.MarkFlagRequired("client-name")
	return cmd
}
<commit_msg>validate flag cfgPath to make sure it is not empty<commit_after>/*
Copyright 2018 The Kubernetes 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 cmd

import (
	"io"
	"time"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"

	"k8s.io/klog/v2"

	kubeadmapiv1 "k8s.io/kubernetes/cmd/kubeadm/app/apis/kubeadm/v1beta3"
	"k8s.io/kubernetes/cmd/kubeadm/app/cmd/options"
	cmdutil "k8s.io/kubernetes/cmd/kubeadm/app/cmd/util"
	kubeadmconstants "k8s.io/kubernetes/cmd/kubeadm/app/constants"
	kubeconfigphase "k8s.io/kubernetes/cmd/kubeadm/app/phases/kubeconfig"
	configutil "k8s.io/kubernetes/cmd/kubeadm/app/util/config"
)

var (
	kubeconfigLongDesc = cmdutil.LongDesc(`
	Kubeconfig file utilities.
	`)

	userKubeconfigLongDesc = cmdutil.LongDesc(`
	Output a kubeconfig file for an additional user.
	`)

	userKubeconfigExample = cmdutil.Examples(`
	# Output a kubeconfig file for an additional user named foo using a kubeadm config file bar
	kubeadm kubeconfig user --client-name=foo --config=bar
	`)
)

// newCmdKubeConfigUtility returns main command for kubeconfig phase
func newCmdKubeConfigUtility(out io.Writer) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "kubeconfig",
		Short: "Kubeconfig file utilities",
		Long:  kubeconfigLongDesc,
	}

	cmd.AddCommand(newCmdUserKubeConfig(out))
	return cmd
}

// newCmdUserKubeConfig returns sub commands for kubeconfig phase
func newCmdUserKubeConfig(out io.Writer) *cobra.Command {

	initCfg := cmdutil.DefaultInitConfiguration()
	clusterCfg := &kubeadmapiv1.ClusterConfiguration{}

	var (
		token, clientName, cfgPath string
		organizations              []string
		validityPeriod             time.Duration
	)

	// Creates the UX Command
	cmd := &cobra.Command{
		Use:     "user",
		Short:   "Output a kubeconfig file for an additional user",
		Long:    userKubeconfigLongDesc,
		Example: userKubeconfigExample,
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(cfgPath) == 0 {
				return errors.New("the kubeadm configuration path cannot be empty")
			}
			// This call returns the ready-to-use configuration based on the defaults populated by flags
			internalCfg, err := configutil.LoadOrDefaultInitConfiguration(cfgPath, initCfg, clusterCfg)
			if err != nil {
				return err
			}

			if validityPeriod > kubeadmconstants.CertificateValidity {
				klog.Warningf("WARNING: the specified certificate validity period %v is longer than the default duration %v, this may increase security risks.",
					validityPeriod, kubeadmconstants.CertificateValidity)
			}

			notAfter := time.Now().Add(validityPeriod).UTC()

			// if the kubeconfig file for an additional user has to use a token, use it
			if token != "" {
				return kubeconfigphase.WriteKubeConfigWithToken(out, internalCfg, clientName, token, &notAfter)
			}

			// Otherwise, write a kubeconfig file with a generate client cert
			return kubeconfigphase.WriteKubeConfigWithClientCert(out, internalCfg, clientName, organizations, &notAfter)
		},
		Args: cobra.NoArgs,
	}

	options.AddConfigFlag(cmd.Flags(), &cfgPath)

	// Add command specific flags
	cmd.Flags().StringVar(&token, options.TokenStr, token, "The token that should be used as the authentication mechanism for this kubeconfig, instead of client certificates")
	cmd.Flags().StringVar(&clientName, "client-name", clientName, "The name of user. It will be used as the CN if client certificates are created")
	cmd.Flags().StringSliceVar(&organizations, "org", organizations, "The orgnizations of the client certificate. It will be used as the O if client certificates are created")
	cmd.Flags().DurationVar(&validityPeriod, "validity-period", kubeadmconstants.CertificateValidity, "The validity period of the client certificate. It is an offset from the current time.")

	cmd.MarkFlagRequired(options.CfgPath)
	cmd.MarkFlagRequired("client-name")
	return cmd
}
<|endoftext|>


<commit_before>package cloudwatch

import (
	"sync"

	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/metrics"
	"github.com/go-kit/kit/metrics/generic"
)

// CloudWatch ...
type CloudWatch struct {
	mtx      sync.RWMutex
	prefix   string
	svc      *cloudwatch.CloudWatch
	counters map[string]*Counter
	//gauges     map[string]*Gauge
	//histograms map[string]*Histogram
	logger log.Logger
}

func New(prefix string, logger log.Logger, svc *cloudwatch.CloudWatch) *CloudWatch {
	return &CloudWatch{
		prefix:   prefix,
		svc:      svc,
		counters: map[string]*Counter{},
		logger:   logger,
	}
}

func (cw *CloudWatch) NewCounter(name string) *Counter {
	c := NewCounter(cw.prefix, name)
	cw.mtx.Lock()
	cw.counters[cw.prefix+name] = c
	cw.mtx.Unlock()
	return c
}

// Counter is a Graphite counter metric.
type Counter struct {
	namespace string
	c         *generic.Counter
}

// NewCounter returns a new usable counter metric.
func NewCounter(namespace, name string) *Counter {
	return &Counter{
		namespace: namespace,
		c:         generic.NewCounter(name),
	}
}

// With is a no-op.
func (c *Counter) With(labelValues ...string) metrics.Counter {
	return c.c.With(labelValues...)
}

// Add implements counter.
func (c *Counter) Add(delta float64) {
	c.c.Add(delta)
}
<commit_msg>Use WriteLoop approach to reduce api overhead.<commit_after>package cloudwatch

import (
	"sync"

	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/metrics"
	"github.com/go-kit/kit/metrics/generic"
)

// CloudWatch ...
type CloudWatch struct {
	mtx      sync.RWMutex
	prefix   string
	svc      *cloudwatch.CloudWatch
	counters map[string]*Counter
	logger   log.Logger
}

// New ...
func New(prefix string, logger log.Logger, svc *cloudwatch.CloudWatch) *CloudWatch {
	return &CloudWatch{
		prefix:   prefix,
		svc:      svc,
		counters: map[string]*Counter{},
		logger:   logger,
	}
}

func (cw *CloudWatch) NewCounter(name string) *Counter {
	c := NewCounter(name)
	cw.mtx.Lock()
	cw.counters[cw.prefix+name] = c
	cw.mtx.Unlock()
	return c
}

// WriteLoop is a helper method that invokes WriteTo to the passed writer every
// time the passed channel fires. This method blocks until the channel is
// closed, so clients probably want to run it in its own goroutine. For typical
// usage, create a time.Ticker and pass its C channel to this method.
func (cw *CloudWatch) WriteLoop(c <-chan time.Time) {
	for range c {
		cw.mtx.RLock()
		defer cw.mtx.RUnlock()
		now := time.Now()

		for name, c := range cw.counters {
			_, err := cw.svc.PutMetricData(&cloudwatch.PutMetricDataInput{
				Namespace: aws.String(cw.prefix),
				MetricData: []*cloudwatch.MetricDatum{
					{
						MetricName: aws.String(name),
						Dimensions: makeDimensions(c.c.LabelValues()...),
						Value:      aws.Float64(c.c.Value()),
						Timestamp:  aws.Time(now),
					},
				},
			})
			if err != nil {
				cw.logger.Log("during", "WriteLoop", "err", err)
			}
		}
	}
}

// Counter is a CloudWatch counter metric.
type Counter struct {
	c *generic.Counter
}

// NewCounter returns a new usable counter metric.
func NewCounter(name string) *Counter {
	return &Counter{
		c: generic.NewCounter(name),
	}
}

// With is a no-op.
func (c *Counter) With(labelValues ...string) metrics.Counter {
	return c.c.With(labelValues...)
}

// Add implements counter.
func (c *Counter) Add(delta float64) {
	c.c.Add(delta)
}

func makeDimensions(labelValues ...string) []*cloudwatch.Dimension {
	dimensions := make([]*cloudwatch.Dimension, len(labelValues)/2)
	for i, j := 0, 0; i < len(labelValues); i, j := i+2, j+1 {
		dimensions[j] = &cloudwatch.Dimension{
			Name:  aws.String(labelValues[i]),
			Value: aws.String(labelValues[i+1]),
		}
	}
	return dimensions
}
<|endoftext|>


<commit_before>package metrics

import (
	"fmt"
	"time"

	"github.com/fweikert/continuous-integration/metrics/clients"
	"github.com/fweikert/continuous-integration/metrics/data"
	timestamp "github.com/golang/protobuf/ptypes/timestamp"
	metricpb "google.golang.org/genproto/googleapis/api/metric"
	monitoredres "google.golang.org/genproto/googleapis/api/monitoredres"
	monitoringpb "google.golang.org/genproto/googleapis/monitoring/v3"
)

type PlatformLoad struct {
	client  *clients.BuildkiteClient
	columns []Column
	builds  int
}

func (pl *PlatformLoad) Name() string {
	return "platform_load"
}

func (pl *PlatformLoad) Columns() []Column {
	return pl.columns
}

func (pl *PlatformLoad) Collect() (data.DataSet, error) {
	builds, err := pl.client.GetMostRecentBuilds("all", pl.builds)
	if err != nil {
		return nil, fmt.Errorf("Cannot get builds to determine platform load: %v", err)
	}

	result := &loadDataSet{headers: GetColumnNames(pl.columns), ts: time.Now()}
	allPlatforms := make(map[string]bool)
	waiting := make(map[string]int)
	running := make(map[string]int)
	for _, build := range builds {
		for _, job := range build.Jobs {
			// Do not use getPlatform() since it may return "rbe", but here we're only interested in the actual worker OS (which would be "linux" in the rbe case).
			platform := getPlatformFromAgentQueryRules(job.AgentQueryRules)
			if platform == "" || job.CreatedAt == nil || job.FinishedAt != nil {
				continue
			}
			allPlatforms[platform] = true
			switch *job.State {
			case "running":
				running[platform] += 1
			case "scheduled":
				/*
					State "scheduled" = waiting for a worker to become available
					State "waiting" / "waiting_failed" = waiting for another task to finish

					We're only interested in "scheduled" jobs since they may indicate a shortage of workers.
				*/
				waiting[platform] += 1
			}
		}
	}

	result.rows = make([]loadDataRow, 0)
	for platform := range allPlatforms {
		row := loadDataRow{platform: platform, waitingJobs: waiting[platform], runningJobs: running[platform]}
		result.rows = append(result.rows, row)
	}
	return result, nil
}

// CREATE TABLE platform_load (timestamp DATETIME, platform VARCHAR(255), waiting_jobs INT, running_jobs INT, PRIMARY KEY(timestamp, platform));
func CreatePlatformLoad(client *clients.BuildkiteClient, builds int) *PlatformLoad {
	columns := []Column{Column{"timestamp", true}, Column{"platform", true}, Column{"waiting_jobs", false}, Column{"running_jobs", false}}
	return &PlatformLoad{client: client, columns: columns, builds: builds}
}

type loadDataRow struct {
	platform    string
	waitingJobs int
	runningJobs int
}

type loadDataSet struct {
	headers []string
	ts      time.Time
	rows    []loadDataRow
}

func (lds *loadDataSet) GetData() *data.LegacyDataSet {
	rawSet := data.CreateDataSet(lds.headers)
	for _, row := range lds.rows {
		rawRow := []interface{}{lds.ts, row.platform, row.waitingJobs, row.runningJobs}
		rawSet.Data = append(rawSet.Data, rawRow)
	}
	return rawSet
}

func (lds *loadDataSet) CreateTimeSeriesRequest(projectID string) *monitoringpb.CreateTimeSeriesRequest {
	ts := &timestamp.Timestamp{
		Seconds: lds.ts.Unix(),
	}
	series := make([]*monitoringpb.TimeSeries, len(lds.rows))
	for i, row := range lds.rows {
		series[i] = row.createTimeSeries("required_workers", ts)
	}

	return &monitoringpb.CreateTimeSeriesRequest{
		Name:       "projects/" + projectID,
		TimeSeries: series,
	}
}

func (ldr *loadDataRow) createTimeSeries(metricType string, ts *timestamp.Timestamp) *monitoringpb.TimeSeries {
	return &monitoringpb.TimeSeries{
		Metric: &metricpb.Metric{
			Type: metricType,
			Labels: map[string]string{
				"platform": ldr.platform,
			},
		},
		Resource: &monitoredres.MonitoredResource{
			Type: "global",
			Labels: map[string]string{
				"platform": ldr.platform,
			},
		},
		Points: []*monitoringpb.Point{{
			Interval: &monitoringpb.TimeInterval{
				StartTime: ts,
				EndTime:   ts,
			},
			Value: &monitoringpb.TypedValue{
				Value: &monitoringpb.TypedValue_Int64Value{
					Int64Value: int64(ldr.waitingJobs + ldr.runningJobs),
				},
			},
		}},
	}
}
<commit_msg>Split into multiple Stackdriver metrics with real names.<commit_after>package metrics

import (
	"fmt"
	"time"

	"github.com/fweikert/continuous-integration/metrics/clients"
	"github.com/fweikert/continuous-integration/metrics/data"
	timestamp "github.com/golang/protobuf/ptypes/timestamp"
	metricpb "google.golang.org/genproto/googleapis/api/metric"
	monitoredres "google.golang.org/genproto/googleapis/api/monitoredres"
	monitoringpb "google.golang.org/genproto/googleapis/monitoring/v3"
)

const baseMetricType = "custom.googleapis.com/bazel/ci"

type PlatformLoad struct {
	client  *clients.BuildkiteClient
	columns []Column
	builds  int
}

func (pl *PlatformLoad) Name() string {
	return "platform_load"
}

func (pl *PlatformLoad) Columns() []Column {
	return pl.columns
}

func (pl *PlatformLoad) Collect() (data.DataSet, error) {
	builds, err := pl.client.GetMostRecentBuilds("all", pl.builds)
	if err != nil {
		return nil, fmt.Errorf("Cannot get builds to determine platform load: %v", err)
	}

	result := &loadDataSet{headers: GetColumnNames(pl.columns), ts: time.Now()}
	allPlatforms := make(map[string]bool)
	waiting := make(map[string]int)
	running := make(map[string]int)
	for _, build := range builds {
		for _, job := range build.Jobs {
			// Do not use getPlatform() since it may return "rbe", but here we're only interested in the actual worker OS (which would be "linux" in the rbe case).
			platform := getPlatformFromAgentQueryRules(job.AgentQueryRules)
			if platform == "" || job.CreatedAt == nil || job.FinishedAt != nil {
				continue
			}
			allPlatforms[platform] = true
			switch *job.State {
			case "running":
				running[platform] += 1
			case "scheduled":
				/*
					State "scheduled" = waiting for a worker to become available
					State "waiting" / "waiting_failed" = waiting for another task to finish

					We're only interested in "scheduled" jobs since they may indicate a shortage of workers.
				*/
				waiting[platform] += 1
			}
		}
	}

	result.rows = make([]loadDataRow, 0)
	for platform := range allPlatforms {
		row := loadDataRow{platform: platform, waitingJobs: waiting[platform], runningJobs: running[platform]}
		result.rows = append(result.rows, row)
	}
	return result, nil
}

// CREATE TABLE platform_load (timestamp DATETIME, platform VARCHAR(255), waiting_jobs INT, running_jobs INT, PRIMARY KEY(timestamp, platform));
func CreatePlatformLoad(client *clients.BuildkiteClient, builds int) *PlatformLoad {
	columns := []Column{Column{"timestamp", true}, Column{"platform", true}, Column{"waiting_jobs", false}, Column{"running_jobs", false}}
	return &PlatformLoad{client: client, columns: columns, builds: builds}
}

type loadDataRow struct {
	platform    string
	waitingJobs int
	runningJobs int
}

type loadDataSet struct {
	headers []string
	ts      time.Time
	rows    []loadDataRow
}

func (lds *loadDataSet) GetData() *data.LegacyDataSet {
	rawSet := data.CreateDataSet(lds.headers)
	for _, row := range lds.rows {
		rawRow := []interface{}{lds.ts, row.platform, row.waitingJobs, row.runningJobs}
		rawSet.Data = append(rawSet.Data, rawRow)
	}
	return rawSet
}

func (lds *loadDataSet) CreateTimeSeriesRequest(projectID string) *monitoringpb.CreateTimeSeriesRequest {
	ts := &timestamp.Timestamp{
		Seconds: lds.ts.Unix(),
	}
	series := make([]*monitoringpb.TimeSeries, len(lds.rows)*2)
	for i, row := range lds.rows {
		series[i] = createTimeSeries(ts, row.platform, "waiting_jobs", row.waitingJobs)
		series[i] = createTimeSeries(ts, row.platform, "running_jobs", row.runningJobs)
	}

	return &monitoringpb.CreateTimeSeriesRequest{
		Name:       "projects/" + projectID,
		TimeSeries: series,
	}
}

func createTimeSeries(ts *timestamp.Timestamp, platform, metricSuffix string, value int) *monitoringpb.TimeSeries {
	return &monitoringpb.TimeSeries{
		Metric: &metricpb.Metric{
			Type: fmt.Sprintf("%s/%s/%s", baseMetricType, platform, metricSuffix),
		},
		Resource: &monitoredres.MonitoredResource{
			Type: "global",
		},
		Points: []*monitoringpb.Point{{
			Interval: &monitoringpb.TimeInterval{
				StartTime: ts,
				EndTime:   ts,
			},
			Value: &monitoringpb.TypedValue{
				Value: &monitoringpb.TypedValue_Int64Value{
					Int64Value: int64(value),
				},
			},
		}},
	}
}
<|endoftext|>


<commit_before>package handler

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

	"github.com/gilcrest/errs"
	"github.com/gilcrest/go-api-basic/controller/movieController"
	"github.com/gorilla/mux"
)

// AddMovie handles POST requests for the /movies endpoint
// and creates a movie in the database
func (ah *AppHandler) AddMovie(w http.ResponseWriter, r *http.Request) {
	const op errs.Op = "handler/AppHandler.AddMovie"

	// Declare rqst as an instance of movieController.MovieRequest
	rqst := new(movieController.MovieRequest)
	// Decode JSON HTTP request body into a Decoder type
	// and unmarshal that into rqst
	err := json.NewDecoder(r.Body).Decode(&rqst)
	defer r.Body.Close()
	if err != nil {
		err = errs.RE(http.StatusBadRequest, errs.InvalidRequest, errs.E(op, err))
		errs.HTTPError(w, err)
		return
	}

	// retrieve the context from the http.Request
	ctx := r.Context()

	// Initialize the MovieController
	mc := movieController.NewMovieController(ah.App, ah.StandardResponseFields)

	// Send the request context and request struct to the controller
	// Receive a response or error in return
	resp, err := mc.Add(ctx, rqst)
	if err != nil {
		err = errs.RE(http.StatusBadRequest, errs.E(op, err))
		errs.HTTPError(w, err)
		return
	}

	// Encode response struct to JSON for the response body
	err = json.NewEncoder(w).Encode(*resp)
	if err != nil {
		err = errs.RE(http.StatusInternalServerError, errs.E(op, errs.Internal))
		errs.HTTPError(w, err)
		return
	}
}

// FindByID handles GET requests for the /movies/{id} endpoint
// and finds a movie by it's ID
func (ah *AppHandler) FindByID(w http.ResponseWriter, r *http.Request) {
	const op errs.Op = "handler/AppHandler.FindByID"

	vars := mux.Vars(r)
	id := vars["id"]

	// retrieve the context from the http.Request
	ctx := r.Context()

	// Initialize the MovieController
	mc := movieController.NewMovieController(ah.App, ah.StandardResponseFields)

	// Send the request context and request struct to the controller
	// Receive a response or error in return
	resp, err := mc.FindByID(ctx, id)
	if err != nil {
		err = errs.RE(http.StatusBadRequest, errs.InvalidRequest, errs.E(op, err))
		errs.HTTPError(w, err)
		return
	}

	// Encode response struct to JSON for the response body
	err = json.NewEncoder(w).Encode(resp)
	if err != nil {
		err = errs.RE(http.StatusInternalServerError, errs.E(op, errs.Internal))
		errs.HTTPError(w, err)
		return
	}
}

// FindAll handles GET requests for the /movies endpoint
// and finds all movies
func (ah *AppHandler) FindAll(w http.ResponseWriter, r *http.Request) {
	const op errs.Op = "handler/AppHandler.FindAll"

	// retrieve the context from the http.Request
	ctx := r.Context()

	// Initialize the MovieController
	mc := movieController.NewMovieController(ah.App, ah.StandardResponseFields)

	// Send the request context and request struct to the controller
	// Receive a response or error in return
	resp, err := mc.FindAll(ctx)
	if err != nil {
		err = errs.RE(http.StatusBadRequest, errs.InvalidRequest, errs.E(op, err))
		errs.HTTPError(w, err)
		return
	}

	// Encode response struct to JSON for the response body
	err = json.NewEncoder(w).Encode(resp)
	if err != nil {
		err = errs.RE(http.StatusInternalServerError, errs.E(op, errs.Internal))
		errs.HTTPError(w, err)
		return
	}
}

// Update handles PUT requests for the /movies/{id} endpoint
// and updates the given movie
func (ah *AppHandler) Update(w http.ResponseWriter, r *http.Request) {
	const op errs.Op = "handler/AppHandler.Update"

	vars := mux.Vars(r)
	id := vars["id"]

	// Declare rqst as an instance of movieController.AddMovieRequest
	rqst := new(movieController.MovieRequest)
	// Decode JSON HTTP request body into a Decoder type
	// and unmarshal that into rqst
	err := json.NewDecoder(r.Body).Decode(&rqst)
	defer r.Body.Close()
	if err != nil {
		err = errs.RE(http.StatusBadRequest, errs.InvalidRequest, errs.E(op, err))
		errs.HTTPError(w, err)
		return
	}

	// retrieve the context from the http.Request
	ctx := r.Context()

	// Initialize the MovieController
	mc := movieController.NewMovieController(ah.App, ah.StandardResponseFields)

	// Send the request context and request struct to the controller
	// Receive a response or error in return
	resp, err := mc.Update(ctx, id, rqst)
	if err != nil {
		err = errs.RE(http.StatusBadRequest, errs.InvalidRequest, errs.E(op, err))
		errs.HTTPError(w, err)
		return
	}

	// Encode response struct to JSON for the response body
	err = json.NewEncoder(w).Encode(resp)
	if err != nil {
		err = errs.RE(http.StatusInternalServerError, errs.E(op, errs.Internal))
		errs.HTTPError(w, err)
		return
	}
}

// Delete handles DELETE requests for the /movies/{id} endpoint
// and updates the given movie
func (ah *AppHandler) Delete(w http.ResponseWriter, r *http.Request) {
	const op errs.Op = "handler/AppHandler.Delete"

	vars := mux.Vars(r)
	id := vars["id"]

	// retrieve the context from the http.Request
	ctx := r.Context()

	// Initialize the MovieController
	mc := movieController.NewMovieController(ah.App, ah.StandardResponseFields)

	// Send the request context and request struct to the controller
	// Receive a response or error in return
	resp, err := mc.Delete(ctx, id)
	if err != nil {
		err = errs.RE(http.StatusBadRequest, errs.InvalidRequest, errs.E(op, err))
		errs.HTTPError(w, err)
		return
	}

	// Encode response struct to JSON for the response body
	err = json.NewEncoder(w).Encode(resp)
	if err != nil {
		err = errs.RE(http.StatusInternalServerError, errs.E(op, errs.Internal))
		errs.HTTPError(w, err)
		return
	}
}
<commit_msg>Add nil body check<commit_after>package handler

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

	"github.com/gilcrest/errs"
	"github.com/gilcrest/go-api-basic/controller/movieController"
	"github.com/gorilla/mux"
)

// AddMovie handles POST requests for the /movies endpoint
// and creates a movie in the database
func (ah *AppHandler) AddMovie(w http.ResponseWriter, r *http.Request) {
	const op errs.Op = "handler/AppHandler.AddMovie"

	// If the request body is nil, json.NewDecoder will panic
	// avoid that by checking for nil in advance
	if r.Body == nil {
		err := errs.RE(http.StatusBadRequest, errs.InvalidRequest, errs.E(op, "Request Body cannot be empty"))
		errs.HTTPError(w, err)
		return
	}
	// Declare rqst as an instance of movieController.MovieRequest
	rqst := new(movieController.MovieRequest)
	// Decode JSON HTTP request body into a Decoder type
	// and unmarshal that into rqst
	err := json.NewDecoder(r.Body).Decode(&rqst)
	defer r.Body.Close()
	if err != nil {
		err = errs.RE(http.StatusBadRequest, errs.InvalidRequest, errs.E(op, err))
		errs.HTTPError(w, err)
		return
	}

	// retrieve the context from the http.Request
	ctx := r.Context()

	// Initialize the MovieController
	mc := movieController.NewMovieController(ah.App, ah.StandardResponseFields)

	// Send the request context and request struct to the controller
	// Receive a response or error in return
	resp, err := mc.Add(ctx, rqst)
	if err != nil {
		err = errs.RE(http.StatusBadRequest, errs.E(op, err))
		errs.HTTPError(w, err)
		return
	}

	// Encode response struct to JSON for the response body
	err = json.NewEncoder(w).Encode(*resp)
	if err != nil {
		err = errs.RE(http.StatusInternalServerError, errs.E(op, errs.Internal))
		errs.HTTPError(w, err)
		return
	}
}

// FindByID handles GET requests for the /movies/{id} endpoint
// and finds a movie by it's ID
func (ah *AppHandler) FindByID(w http.ResponseWriter, r *http.Request) {
	const op errs.Op = "handler/AppHandler.FindByID"

	vars := mux.Vars(r)
	id := vars["id"]

	// retrieve the context from the http.Request
	ctx := r.Context()

	// Initialize the MovieController
	mc := movieController.NewMovieController(ah.App, ah.StandardResponseFields)

	// Send the request context and request struct to the controller
	// Receive a response or error in return
	resp, err := mc.FindByID(ctx, id)
	if err != nil {
		err = errs.RE(http.StatusBadRequest, errs.InvalidRequest, errs.E(op, err))
		errs.HTTPError(w, err)
		return
	}

	// Encode response struct to JSON for the response body
	err = json.NewEncoder(w).Encode(resp)
	if err != nil {
		err = errs.RE(http.StatusInternalServerError, errs.E(op, errs.Internal))
		errs.HTTPError(w, err)
		return
	}
}

// FindAll handles GET requests for the /movies endpoint
// and finds all movies
func (ah *AppHandler) FindAll(w http.ResponseWriter, r *http.Request) {
	const op errs.Op = "handler/AppHandler.FindAll"

	// retrieve the context from the http.Request
	ctx := r.Context()

	// Initialize the MovieController
	mc := movieController.NewMovieController(ah.App, ah.StandardResponseFields)

	// Send the request context and request struct to the controller
	// Receive a response or error in return
	resp, err := mc.FindAll(ctx)
	if err != nil {
		err = errs.RE(http.StatusBadRequest, errs.InvalidRequest, errs.E(op, err))
		errs.HTTPError(w, err)
		return
	}

	// Encode response struct to JSON for the response body
	err = json.NewEncoder(w).Encode(resp)
	if err != nil {
		err = errs.RE(http.StatusInternalServerError, errs.E(op, errs.Internal))
		errs.HTTPError(w, err)
		return
	}
}

// Update handles PUT requests for the /movies/{id} endpoint
// and updates the given movie
func (ah *AppHandler) Update(w http.ResponseWriter, r *http.Request) {
	const op errs.Op = "handler/AppHandler.Update"

	// If the request body is nil, json.NewDecoder will panic
	// avoid that by checking for nil in advance
	if r.Body == nil {
		err := errs.RE(http.StatusBadRequest, errs.InvalidRequest, errs.E(op, "Request Body cannot be empty"))
		errs.HTTPError(w, err)
		return
	}

	vars := mux.Vars(r)
	id := vars["id"]

	// Declare rqst as an instance of movieController.AddMovieRequest
	rqst := new(movieController.MovieRequest)
	// Decode JSON HTTP request body into a Decoder type
	// and unmarshal that into rqst
	err := json.NewDecoder(r.Body).Decode(&rqst)
	defer r.Body.Close()
	if err != nil {
		err = errs.RE(http.StatusBadRequest, errs.InvalidRequest, errs.E(op, err))
		errs.HTTPError(w, err)
		return
	}

	// retrieve the context from the http.Request
	ctx := r.Context()

	// Initialize the MovieController
	mc := movieController.NewMovieController(ah.App, ah.StandardResponseFields)

	// Send the request context and request struct to the controller
	// Receive a response or error in return
	resp, err := mc.Update(ctx, id, rqst)
	if err != nil {
		err = errs.RE(http.StatusBadRequest, errs.InvalidRequest, errs.E(op, err))
		errs.HTTPError(w, err)
		return
	}

	// Encode response struct to JSON for the response body
	err = json.NewEncoder(w).Encode(resp)
	if err != nil {
		err = errs.RE(http.StatusInternalServerError, errs.E(op, errs.Internal))
		errs.HTTPError(w, err)
		return
	}
}

// Delete handles DELETE requests for the /movies/{id} endpoint
// and updates the given movie
func (ah *AppHandler) Delete(w http.ResponseWriter, r *http.Request) {
	const op errs.Op = "handler/AppHandler.Delete"

	vars := mux.Vars(r)
	id := vars["id"]

	// retrieve the context from the http.Request
	ctx := r.Context()

	// Initialize the MovieController
	mc := movieController.NewMovieController(ah.App, ah.StandardResponseFields)

	// Send the request context and request struct to the controller
	// Receive a response or error in return
	resp, err := mc.Delete(ctx, id)
	if err != nil {
		err = errs.RE(http.StatusBadRequest, errs.InvalidRequest, errs.E(op, err))
		errs.HTTPError(w, err)
		return
	}

	// Encode response struct to JSON for the response body
	err = json.NewEncoder(w).Encode(resp)
	if err != nil {
		err = errs.RE(http.StatusInternalServerError, errs.E(op, errs.Internal))
		errs.HTTPError(w, err)
		return
	}
}
<|endoftext|>


<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package mqtt

import (
	"encoding/json"
	"fmt"
	"time"

	"github.com/TheThingsNetwork/ttn/utils/random"
	"github.com/apex/log"
	MQTT "github.com/eclipse/paho.mqtt.golang"
)

const QoS = 0x00

// Client connects to the MQTT server and can publish/subscribe on uplink, downlink and activations from devices
type Client interface {
	Connect() error
	Disconnect()

	IsConnected() bool

	// Uplink pub/sub
	PublishUplink(payload UplinkMessage) Token
	SubscribeDeviceUplink(appID string, devID string, handler UplinkHandler) Token
	SubscribeAppUplink(appID string, handler UplinkHandler) Token
	SubscribeUplink(handler UplinkHandler) Token
	UnsubscribeDeviceUplink(appID string, devID string) Token
	UnsubscribeAppUplink(appID string) Token
	UnsubscribeUplink() Token

	// Downlink pub/sub
	PublishDownlink(payload DownlinkMessage) Token
	SubscribeDeviceDownlink(appID string, devID string, handler DownlinkHandler) Token
	SubscribeAppDownlink(appID string, handler DownlinkHandler) Token
	SubscribeDownlink(handler DownlinkHandler) Token
	UnsubscribeDeviceDownlink(appID string, devID string) Token
	UnsubscribeAppDownlink(appID string) Token
	UnsubscribeDownlink() Token

	// Activation pub/sub
	PublishActivation(payload Activation) Token
	SubscribeDeviceActivations(appID string, devID string, handler ActivationHandler) Token
	SubscribeAppActivations(appID string, handler ActivationHandler) Token
	SubscribeActivations(handler ActivationHandler) Token
	UnsubscribeDeviceActivations(appID string, devID string) Token
	UnsubscribeAppActivations(appID string) Token
	UnsubscribeActivations() Token
}

// Token is returned on asyncronous functions
type Token interface {
	// Wait for the function to finish
	Wait() bool
	// Wait for the function to finish or return false after a certain time
	WaitTimeout(time.Duration) bool
	// The error associated with the result of the function (nil if everything okay)
	Error() error
}

type simpleToken struct {
	err error
}

// Wait always returns true
func (t *simpleToken) Wait() bool {
	return true
}

// WaitTimeout always returns true
func (t *simpleToken) WaitTimeout(_ time.Duration) bool {
	return true
}

// Error contains the error if present
func (t *simpleToken) Error() error {
	return t.err
}

// UplinkHandler is called for uplink messages
type UplinkHandler func(client Client, appID string, devID string, req UplinkMessage)

// DownlinkHandler is called for downlink messages
type DownlinkHandler func(client Client, appID string, devID string, req DownlinkMessage)

// ActivationHandler is called for activations
type ActivationHandler func(client Client, appID string, devID string, req Activation)

// DefaultClient is the default MQTT client for The Things Network
type DefaultClient struct {
	mqtt          MQTT.Client
	ctx           log.Interface
	subscriptions map[string]MQTT.MessageHandler
}

// NewClient creates a new DefaultClient
func NewClient(ctx log.Interface, id, username, password string, brokers ...string) Client {
	mqttOpts := MQTT.NewClientOptions()

	for _, broker := range brokers {
		mqttOpts.AddBroker(broker)
	}

	mqttOpts.SetClientID(fmt.Sprintf("%s-%s", id, random.String(16)))
	mqttOpts.SetUsername(username)
	mqttOpts.SetPassword(password)

	// TODO: Some tuning of these values probably won't hurt:
	mqttOpts.SetKeepAlive(30 * time.Second)
	mqttOpts.SetPingTimeout(10 * time.Second)

	mqttOpts.SetCleanSession(true)

	mqttOpts.SetDefaultPublishHandler(func(client MQTT.Client, msg MQTT.Message) {
		ctx.WithField("message", msg).Warn("Received unhandled message")
	})

	var reconnecting bool

	mqttOpts.SetConnectionLostHandler(func(client MQTT.Client, err error) {
		ctx.WithError(err).Warn("Disconnected, reconnecting...")
		reconnecting = true
	})

	ttnClient := &DefaultClient{
		ctx:           ctx,
		subscriptions: make(map[string]MQTT.MessageHandler),
	}

	mqttOpts.SetOnConnectHandler(func(client MQTT.Client) {
		ctx.Info("Connected to MQTT")
		if reconnecting {
			for topic, handler := range ttnClient.subscriptions {
				ctx.Infof("Re-subscribing to %s", topic)
				ttnClient.subscribe(topic, handler)
			}
			reconnecting = false
		}
	})

	ttnClient.mqtt = MQTT.NewClient(mqttOpts)

	return ttnClient
}

var (
	// ConnectRetries says how many times the client should retry a failed connection
	ConnectRetries = 10
	// ConnectRetryDelay says how long the client should wait between retries
	ConnectRetryDelay = time.Second
)

// Connect to the MQTT broker. It will retry for ConnectRetries times with a delay of ConnectRetryDelay between retries
func (c *DefaultClient) Connect() error {
	if c.mqtt.IsConnected() {
		return nil
	}
	var err error
	for retries := 0; retries < ConnectRetries; retries++ {
		token := c.mqtt.Connect()
		finished := token.WaitTimeout(1 * time.Second)
		err = token.Error()
		if finished && err == nil {
			break
		}
		c.ctx.WithError(err).Warn("Could not connect to MQTT Broker. Retrying...")
		<-time.After(ConnectRetryDelay)
	}
	if err != nil {
		return fmt.Errorf("Could not connect to MQTT Broker: %s", err)
	}
	return nil
}

func (c *DefaultClient) subscribe(topic string, handler MQTT.MessageHandler) Token {
	c.subscriptions[topic] = handler
	return c.mqtt.Subscribe(topic, QoS, handler)
}

func (c *DefaultClient) unsubscribe(topic string) Token {
	delete(c.subscriptions, topic)
	return c.mqtt.Unsubscribe(topic)
}

// Disconnect from the MQTT broker
func (c *DefaultClient) Disconnect() {
	if !c.mqtt.IsConnected() {
		return
	}
	c.ctx.Debug("Disconnecting from MQTT")
	c.mqtt.Disconnect(25)
}

// IsConnected returns true if there is a connection to the MQTT broker
func (c *DefaultClient) IsConnected() bool {
	return c.mqtt.IsConnected()
}

// PublishUplink publishes an uplink message to the MQTT broker
func (c *DefaultClient) PublishUplink(dataUp UplinkMessage) Token {
	topic := DeviceTopic{dataUp.AppID, dataUp.DevID, Uplink}
	dataUp.AppID = ""
	dataUp.DevID = ""
	msg, err := json.Marshal(dataUp)
	if err != nil {
		return &simpleToken{fmt.Errorf("Unable to marshal the message payload")}
	}
	return c.mqtt.Publish(topic.String(), QoS, false, msg)
}

// SubscribeDeviceUplink subscribes to all uplink messages for the given application and device
func (c *DefaultClient) SubscribeDeviceUplink(appID string, devID string, handler UplinkHandler) Token {
	topic := DeviceTopic{appID, devID, Uplink}
	return c.subscribe(topic.String(), func(mqtt MQTT.Client, msg MQTT.Message) {
		// Determine the actual topic
		topic, err := ParseDeviceTopic(msg.Topic())
		if err != nil {
			c.ctx.WithField("topic", msg.Topic()).WithError(err).Warn("Received message on invalid uplink topic")
			return
		}

		// Unmarshal the payload
		dataUp := &UplinkMessage{}
		err = json.Unmarshal(msg.Payload(), dataUp)
		dataUp.AppID = topic.AppID
		dataUp.DevID = topic.DevID

		if err != nil {
			c.ctx.WithError(err).Warn("Could not unmarshal uplink")
			return
		}

		// Call the uplink handler
		handler(c, topic.AppID, topic.DevID, *dataUp)
	})
}

// SubscribeAppUplink subscribes to all uplink messages for the given application
func (c *DefaultClient) SubscribeAppUplink(appID string, handler UplinkHandler) Token {
	return c.SubscribeDeviceUplink(appID, "", handler)
}

// SubscribeUplink subscribes to all uplink messages that the current user has access to
func (c *DefaultClient) SubscribeUplink(handler UplinkHandler) Token {
	return c.SubscribeDeviceUplink("", "", handler)
}

// UnsubscribeDeviceUplink unsubscribes from the uplink messages for the given application and device
func (c *DefaultClient) UnsubscribeDeviceUplink(appID string, devID string) Token {
	topic := DeviceTopic{appID, devID, Uplink}
	return c.unsubscribe(topic.String())
}

// UnsubscribeAppUplink unsubscribes from the uplink messages for the given application
func (c *DefaultClient) UnsubscribeAppUplink(appID string) Token {
	return c.UnsubscribeDeviceUplink(appID, "")
}

// UnsubscribeUplink unsubscribes from the uplink messages that the current user has access to
func (c *DefaultClient) UnsubscribeUplink() Token {
	return c.UnsubscribeDeviceUplink("", "")
}

// PublishDownlink publishes a downlink message
func (c *DefaultClient) PublishDownlink(dataDown DownlinkMessage) Token {
	topic := DeviceTopic{dataDown.AppID, dataDown.DevID, Downlink}
	dataDown.AppID = ""
	dataDown.DevID = ""
	msg, err := json.Marshal(dataDown)
	if err != nil {
		return &simpleToken{fmt.Errorf("Unable to marshal the message payload")}
	}
	return c.mqtt.Publish(topic.String(), QoS, false, msg)
}

// SubscribeDeviceDownlink subscribes to all downlink messages for the given application and device
func (c *DefaultClient) SubscribeDeviceDownlink(appID string, devID string, handler DownlinkHandler) Token {
	topic := DeviceTopic{appID, devID, Downlink}
	return c.subscribe(topic.String(), func(mqtt MQTT.Client, msg MQTT.Message) {
		// Determine the actual topic
		topic, err := ParseDeviceTopic(msg.Topic())
		if err != nil {
			c.ctx.WithField("topic", msg.Topic()).WithError(err).Warn("Received message on invalid Downlink topic")
			return
		}

		// Unmarshal the payload
		dataDown := &DownlinkMessage{}
		err = json.Unmarshal(msg.Payload(), dataDown)
		if err != nil {
			c.ctx.WithError(err).Warn("Could not unmarshal Downlink")
			return
		}
		dataDown.AppID = topic.AppID
		dataDown.DevID = topic.DevID

		// Call the Downlink handler
		handler(c, topic.AppID, topic.DevID, *dataDown)
	})
}

// SubscribeAppDownlink subscribes to all downlink messages for the given application
func (c *DefaultClient) SubscribeAppDownlink(appID string, handler DownlinkHandler) Token {
	return c.SubscribeDeviceDownlink(appID, "", handler)
}

// SubscribeDownlink subscribes to all downlink messages that the current user has access to
func (c *DefaultClient) SubscribeDownlink(handler DownlinkHandler) Token {
	return c.SubscribeDeviceDownlink("", "", handler)
}

// UnsubscribeDeviceDownlink unsubscribes from the downlink messages for the given application and device
func (c *DefaultClient) UnsubscribeDeviceDownlink(appID string, devID string) Token {
	topic := DeviceTopic{appID, devID, Downlink}
	return c.unsubscribe(topic.String())
}

// UnsubscribeAppDownlink unsubscribes from the downlink messages for the given application
func (c *DefaultClient) UnsubscribeAppDownlink(appID string) Token {
	return c.UnsubscribeDeviceDownlink(appID, "")
}

// UnsubscribeDownlink unsubscribes from the downlink messages that the current user has access to
func (c *DefaultClient) UnsubscribeDownlink() Token {
	return c.UnsubscribeDeviceDownlink("", "")
}

// PublishActivation publishes an activation
func (c *DefaultClient) PublishActivation(activation Activation) Token {
	topic := DeviceTopic{activation.AppID, activation.DevID, Activations}
	activation.AppID = ""
	activation.DevID = ""
	msg, err := json.Marshal(activation)
	if err != nil {
		return &simpleToken{fmt.Errorf("Unable to marshal the message payload")}
	}
	return c.mqtt.Publish(topic.String(), QoS, false, msg)
}

// SubscribeDeviceActivations subscribes to all activations for the given application and device
func (c *DefaultClient) SubscribeDeviceActivations(appID string, devID string, handler ActivationHandler) Token {
	topic := DeviceTopic{appID, devID, Activations}
	return c.subscribe(topic.String(), func(mqtt MQTT.Client, msg MQTT.Message) {
		// Determine the actual topic
		topic, err := ParseDeviceTopic(msg.Topic())
		if err != nil {
			c.ctx.WithField("topic", msg.Topic()).WithError(err).Warn("Received message on invalid Activations topic")
			return
		}

		// Unmarshal the payload
		activation := &Activation{}
		err = json.Unmarshal(msg.Payload(), activation)
		if err != nil {
			c.ctx.WithError(err).Warn("Could not unmarshal Activation")
			return
		}
		activation.AppID = topic.AppID
		activation.DevID = topic.DevID

		// Call the Activation handler
		handler(c, topic.AppID, topic.DevID, *activation)
	})
}

// SubscribeAppActivations subscribes to all activations for the given application
func (c *DefaultClient) SubscribeAppActivations(appID string, handler ActivationHandler) Token {
	return c.SubscribeDeviceActivations(appID, "", handler)
}

// SubscribeActivations subscribes to all activations that the current user has access to
func (c *DefaultClient) SubscribeActivations(handler ActivationHandler) Token {
	return c.SubscribeDeviceActivations("", "", handler)
}

// UnsubscribeDeviceActivations unsubscribes from the activations for the given application and device
func (c *DefaultClient) UnsubscribeDeviceActivations(appID string, devID string) Token {
	topic := DeviceTopic{appID, devID, Activations}
	return c.unsubscribe(topic.String())
}

// UnsubscribeAppActivations unsubscribes from the activations for the given application
func (c *DefaultClient) UnsubscribeAppActivations(appID string) Token {
	return c.UnsubscribeDeviceActivations(appID, "")
}

// UnsubscribeActivations unsubscribes from the activations that the current user has access to
func (c *DefaultClient) UnsubscribeActivations() Token {
	return c.UnsubscribeDeviceActivations("", "")
}
<commit_msg>Fix nil pointer when token expired<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package mqtt

import (
	"encoding/json"
	"fmt"
	"time"

	"github.com/TheThingsNetwork/ttn/utils/random"
	"github.com/apex/log"
	MQTT "github.com/eclipse/paho.mqtt.golang"
)

const QoS = 0x00

// Client connects to the MQTT server and can publish/subscribe on uplink, downlink and activations from devices
type Client interface {
	Connect() error
	Disconnect()

	IsConnected() bool

	// Uplink pub/sub
	PublishUplink(payload UplinkMessage) Token
	SubscribeDeviceUplink(appID string, devID string, handler UplinkHandler) Token
	SubscribeAppUplink(appID string, handler UplinkHandler) Token
	SubscribeUplink(handler UplinkHandler) Token
	UnsubscribeDeviceUplink(appID string, devID string) Token
	UnsubscribeAppUplink(appID string) Token
	UnsubscribeUplink() Token

	// Downlink pub/sub
	PublishDownlink(payload DownlinkMessage) Token
	SubscribeDeviceDownlink(appID string, devID string, handler DownlinkHandler) Token
	SubscribeAppDownlink(appID string, handler DownlinkHandler) Token
	SubscribeDownlink(handler DownlinkHandler) Token
	UnsubscribeDeviceDownlink(appID string, devID string) Token
	UnsubscribeAppDownlink(appID string) Token
	UnsubscribeDownlink() Token

	// Activation pub/sub
	PublishActivation(payload Activation) Token
	SubscribeDeviceActivations(appID string, devID string, handler ActivationHandler) Token
	SubscribeAppActivations(appID string, handler ActivationHandler) Token
	SubscribeActivations(handler ActivationHandler) Token
	UnsubscribeDeviceActivations(appID string, devID string) Token
	UnsubscribeAppActivations(appID string) Token
	UnsubscribeActivations() Token
}

// Token is returned on asyncronous functions
type Token interface {
	// Wait for the function to finish
	Wait() bool
	// Wait for the function to finish or return false after a certain time
	WaitTimeout(time.Duration) bool
	// The error associated with the result of the function (nil if everything okay)
	Error() error
}

type simpleToken struct {
	err error
}

// Wait always returns true
func (t *simpleToken) Wait() bool {
	return true
}

// WaitTimeout always returns true
func (t *simpleToken) WaitTimeout(_ time.Duration) bool {
	return true
}

// Error contains the error if present
func (t *simpleToken) Error() error {
	return t.err
}

// UplinkHandler is called for uplink messages
type UplinkHandler func(client Client, appID string, devID string, req UplinkMessage)

// DownlinkHandler is called for downlink messages
type DownlinkHandler func(client Client, appID string, devID string, req DownlinkMessage)

// ActivationHandler is called for activations
type ActivationHandler func(client Client, appID string, devID string, req Activation)

// DefaultClient is the default MQTT client for The Things Network
type DefaultClient struct {
	mqtt          MQTT.Client
	ctx           log.Interface
	subscriptions map[string]MQTT.MessageHandler
}

// NewClient creates a new DefaultClient
func NewClient(ctx log.Interface, id, username, password string, brokers ...string) Client {
	mqttOpts := MQTT.NewClientOptions()

	for _, broker := range brokers {
		mqttOpts.AddBroker(broker)
	}

	mqttOpts.SetClientID(fmt.Sprintf("%s-%s", id, random.String(16)))
	mqttOpts.SetUsername(username)
	mqttOpts.SetPassword(password)

	// TODO: Some tuning of these values probably won't hurt:
	mqttOpts.SetKeepAlive(30 * time.Second)
	mqttOpts.SetPingTimeout(10 * time.Second)

	mqttOpts.SetCleanSession(true)

	mqttOpts.SetDefaultPublishHandler(func(client MQTT.Client, msg MQTT.Message) {
		ctx.WithField("message", msg).Warn("Received unhandled message")
	})

	var reconnecting bool

	mqttOpts.SetConnectionLostHandler(func(client MQTT.Client, err error) {
		ctx.WithError(err).Warn("Disconnected, reconnecting...")
		reconnecting = true
	})

	ttnClient := &DefaultClient{
		ctx:           ctx,
		subscriptions: make(map[string]MQTT.MessageHandler),
	}

	mqttOpts.SetOnConnectHandler(func(client MQTT.Client) {
		ctx.Info("Connected to MQTT")
		if reconnecting {
			for topic, handler := range ttnClient.subscriptions {
				ctx.Infof("Re-subscribing to %s", topic)
				ttnClient.subscribe(topic, handler)
			}
			reconnecting = false
		}
	})

	ttnClient.mqtt = MQTT.NewClient(mqttOpts)

	return ttnClient
}

var (
	// ConnectRetries says how many times the client should retry a failed connection
	ConnectRetries = 10
	// ConnectRetryDelay says how long the client should wait between retries
	ConnectRetryDelay = time.Second
)

// Connect to the MQTT broker. It will retry for ConnectRetries times with a delay of ConnectRetryDelay between retries
func (c *DefaultClient) Connect() error {
	if c.mqtt.IsConnected() {
		return nil
	}
	var err error
	for retries := 0; retries < ConnectRetries; retries++ {
		token := c.mqtt.Connect()
		finished := token.WaitTimeout(1 * time.Second)
		if !finished {
			c.ctx.Warn("MQTT connection took longer than expected...")
			token.Wait()
		}
		err = token.Error()
		if err == nil {
			break
		}
		c.ctx.WithError(err).Warn("Could not connect to MQTT Broker. Retrying...")
		<-time.After(ConnectRetryDelay)
	}
	if err != nil {
		return fmt.Errorf("Could not connect to MQTT Broker: %s", err)
	}
	return nil
}

func (c *DefaultClient) subscribe(topic string, handler MQTT.MessageHandler) Token {
	c.subscriptions[topic] = handler
	return c.mqtt.Subscribe(topic, QoS, handler)
}

func (c *DefaultClient) unsubscribe(topic string) Token {
	delete(c.subscriptions, topic)
	return c.mqtt.Unsubscribe(topic)
}

// Disconnect from the MQTT broker
func (c *DefaultClient) Disconnect() {
	if !c.mqtt.IsConnected() {
		return
	}
	c.ctx.Debug("Disconnecting from MQTT")
	c.mqtt.Disconnect(25)
}

// IsConnected returns true if there is a connection to the MQTT broker
func (c *DefaultClient) IsConnected() bool {
	return c.mqtt.IsConnected()
}

// PublishUplink publishes an uplink message to the MQTT broker
func (c *DefaultClient) PublishUplink(dataUp UplinkMessage) Token {
	topic := DeviceTopic{dataUp.AppID, dataUp.DevID, Uplink}
	dataUp.AppID = ""
	dataUp.DevID = ""
	msg, err := json.Marshal(dataUp)
	if err != nil {
		return &simpleToken{fmt.Errorf("Unable to marshal the message payload")}
	}
	return c.mqtt.Publish(topic.String(), QoS, false, msg)
}

// SubscribeDeviceUplink subscribes to all uplink messages for the given application and device
func (c *DefaultClient) SubscribeDeviceUplink(appID string, devID string, handler UplinkHandler) Token {
	topic := DeviceTopic{appID, devID, Uplink}
	return c.subscribe(topic.String(), func(mqtt MQTT.Client, msg MQTT.Message) {
		// Determine the actual topic
		topic, err := ParseDeviceTopic(msg.Topic())
		if err != nil {
			c.ctx.WithField("topic", msg.Topic()).WithError(err).Warn("Received message on invalid uplink topic")
			return
		}

		// Unmarshal the payload
		dataUp := &UplinkMessage{}
		err = json.Unmarshal(msg.Payload(), dataUp)
		dataUp.AppID = topic.AppID
		dataUp.DevID = topic.DevID

		if err != nil {
			c.ctx.WithError(err).Warn("Could not unmarshal uplink")
			return
		}

		// Call the uplink handler
		handler(c, topic.AppID, topic.DevID, *dataUp)
	})
}

// SubscribeAppUplink subscribes to all uplink messages for the given application
func (c *DefaultClient) SubscribeAppUplink(appID string, handler UplinkHandler) Token {
	return c.SubscribeDeviceUplink(appID, "", handler)
}

// SubscribeUplink subscribes to all uplink messages that the current user has access to
func (c *DefaultClient) SubscribeUplink(handler UplinkHandler) Token {
	return c.SubscribeDeviceUplink("", "", handler)
}

// UnsubscribeDeviceUplink unsubscribes from the uplink messages for the given application and device
func (c *DefaultClient) UnsubscribeDeviceUplink(appID string, devID string) Token {
	topic := DeviceTopic{appID, devID, Uplink}
	return c.unsubscribe(topic.String())
}

// UnsubscribeAppUplink unsubscribes from the uplink messages for the given application
func (c *DefaultClient) UnsubscribeAppUplink(appID string) Token {
	return c.UnsubscribeDeviceUplink(appID, "")
}

// UnsubscribeUplink unsubscribes from the uplink messages that the current user has access to
func (c *DefaultClient) UnsubscribeUplink() Token {
	return c.UnsubscribeDeviceUplink("", "")
}

// PublishDownlink publishes a downlink message
func (c *DefaultClient) PublishDownlink(dataDown DownlinkMessage) Token {
	topic := DeviceTopic{dataDown.AppID, dataDown.DevID, Downlink}
	dataDown.AppID = ""
	dataDown.DevID = ""
	msg, err := json.Marshal(dataDown)
	if err != nil {
		return &simpleToken{fmt.Errorf("Unable to marshal the message payload")}
	}
	return c.mqtt.Publish(topic.String(), QoS, false, msg)
}

// SubscribeDeviceDownlink subscribes to all downlink messages for the given application and device
func (c *DefaultClient) SubscribeDeviceDownlink(appID string, devID string, handler DownlinkHandler) Token {
	topic := DeviceTopic{appID, devID, Downlink}
	return c.subscribe(topic.String(), func(mqtt MQTT.Client, msg MQTT.Message) {
		// Determine the actual topic
		topic, err := ParseDeviceTopic(msg.Topic())
		if err != nil {
			c.ctx.WithField("topic", msg.Topic()).WithError(err).Warn("Received message on invalid Downlink topic")
			return
		}

		// Unmarshal the payload
		dataDown := &DownlinkMessage{}
		err = json.Unmarshal(msg.Payload(), dataDown)
		if err != nil {
			c.ctx.WithError(err).Warn("Could not unmarshal Downlink")
			return
		}
		dataDown.AppID = topic.AppID
		dataDown.DevID = topic.DevID

		// Call the Downlink handler
		handler(c, topic.AppID, topic.DevID, *dataDown)
	})
}

// SubscribeAppDownlink subscribes to all downlink messages for the given application
func (c *DefaultClient) SubscribeAppDownlink(appID string, handler DownlinkHandler) Token {
	return c.SubscribeDeviceDownlink(appID, "", handler)
}

// SubscribeDownlink subscribes to all downlink messages that the current user has access to
func (c *DefaultClient) SubscribeDownlink(handler DownlinkHandler) Token {
	return c.SubscribeDeviceDownlink("", "", handler)
}

// UnsubscribeDeviceDownlink unsubscribes from the downlink messages for the given application and device
func (c *DefaultClient) UnsubscribeDeviceDownlink(appID string, devID string) Token {
	topic := DeviceTopic{appID, devID, Downlink}
	return c.unsubscribe(topic.String())
}

// UnsubscribeAppDownlink unsubscribes from the downlink messages for the given application
func (c *DefaultClient) UnsubscribeAppDownlink(appID string) Token {
	return c.UnsubscribeDeviceDownlink(appID, "")
}

// UnsubscribeDownlink unsubscribes from the downlink messages that the current user has access to
func (c *DefaultClient) UnsubscribeDownlink() Token {
	return c.UnsubscribeDeviceDownlink("", "")
}

// PublishActivation publishes an activation
func (c *DefaultClient) PublishActivation(activation Activation) Token {
	topic := DeviceTopic{activation.AppID, activation.DevID, Activations}
	activation.AppID = ""
	activation.DevID = ""
	msg, err := json.Marshal(activation)
	if err != nil {
		return &simpleToken{fmt.Errorf("Unable to marshal the message payload")}
	}
	return c.mqtt.Publish(topic.String(), QoS, false, msg)
}

// SubscribeDeviceActivations subscribes to all activations for the given application and device
func (c *DefaultClient) SubscribeDeviceActivations(appID string, devID string, handler ActivationHandler) Token {
	topic := DeviceTopic{appID, devID, Activations}
	return c.subscribe(topic.String(), func(mqtt MQTT.Client, msg MQTT.Message) {
		// Determine the actual topic
		topic, err := ParseDeviceTopic(msg.Topic())
		if err != nil {
			c.ctx.WithField("topic", msg.Topic()).WithError(err).Warn("Received message on invalid Activations topic")
			return
		}

		// Unmarshal the payload
		activation := &Activation{}
		err = json.Unmarshal(msg.Payload(), activation)
		if err != nil {
			c.ctx.WithError(err).Warn("Could not unmarshal Activation")
			return
		}
		activation.AppID = topic.AppID
		activation.DevID = topic.DevID

		// Call the Activation handler
		handler(c, topic.AppID, topic.DevID, *activation)
	})
}

// SubscribeAppActivations subscribes to all activations for the given application
func (c *DefaultClient) SubscribeAppActivations(appID string, handler ActivationHandler) Token {
	return c.SubscribeDeviceActivations(appID, "", handler)
}

// SubscribeActivations subscribes to all activations that the current user has access to
func (c *DefaultClient) SubscribeActivations(handler ActivationHandler) Token {
	return c.SubscribeDeviceActivations("", "", handler)
}

// UnsubscribeDeviceActivations unsubscribes from the activations for the given application and device
func (c *DefaultClient) UnsubscribeDeviceActivations(appID string, devID string) Token {
	topic := DeviceTopic{appID, devID, Activations}
	return c.unsubscribe(topic.String())
}

// UnsubscribeAppActivations unsubscribes from the activations for the given application
func (c *DefaultClient) UnsubscribeAppActivations(appID string) Token {
	return c.UnsubscribeDeviceActivations(appID, "")
}

// UnsubscribeActivations unsubscribes from the activations that the current user has access to
func (c *DefaultClient) UnsubscribeActivations() Token {
	return c.UnsubscribeDeviceActivations("", "")
}
<|endoftext|>


<commit_before>package urlsign

import (
	"crypto/hmac"
	"crypto/sha512"
	"encoding/base64"
	"log"
	"net/http"
	"net/url"

	"github.com/divoxx/stackerr"
)

var Log *log.Logger

var (
	ErrOpaqueURL         = stackerr.New("can't sign or verify opaque URL")
	ErrSignatureMismatch = stackerr.New("signature mismatch")
)

type Signer struct {
	Key       []byte
	ParamName string
}

func (s Signer) SignURL(u *url.URL) error {
	if u.Opaque != "" {
		return ErrOpaqueURL
	}

	chunks := []string{u.Scheme, u.Host, u.Path, u.RawQuery}
	if u.User != nil {
		chunks = append(chunks, u.String())
	}

	sig := computeSignature(s.Key, chunks)

	q := u.Query()
	q.Set(s.ParamName, sig)
	u.RawQuery = q.Encode()

	return nil
}

func (s Signer) VerifyURL(u *url.URL) error {
	if u.Opaque != "" {
		return ErrOpaqueURL
	}

	q := u.Query()
	sig := q.Get(s.ParamName)
	q.Del(s.ParamName)

	chunks := []string{u.Scheme, u.Host, u.Path, q.Encode()}
	if u.User != nil {
		chunks = append(chunks, u.String())
	}

	expSig := computeSignature(s.Key, chunks)

	if sig != expSig {
		Log.Printf("signature mismatch, expected %s but got %s", expSig, sig)
		return ErrSignatureMismatch
	}

	return nil
}

func (s Signer) VerifyReq(req *http.Request) error {
	u := url.URL{Host: req.Host}

	if p := req.Header.Get("X-Forwarded-Proto"); p != "" {
		u.Scheme = p
	} else if req.TLS != nil {
		u.Scheme = "https"
	} else {
		u.Scheme = "http"
	}

	return s.VerifyURL(u.ResolveReference(req.URL))
}

func computeSignature(key []byte, chunks []string) string {
	h := hmac.New(sha512.New, key)
	for _, c := range chunks {
		h.Write([]byte(c))
	}

	return base64.URLEncoding.EncodeToString(h.Sum(nil))
}
<commit_msg>Fix logger<commit_after>package urlsign

import (
	"crypto/hmac"
	"crypto/sha512"
	"encoding/base64"
	"log"
	"net/http"
	"net/url"

	"github.com/divoxx/stackerr"
)

var (
	ErrOpaqueURL         = stackerr.New("can't sign or verify opaque URL")
	ErrSignatureMismatch = stackerr.New("signature mismatch")
)

type Signer struct {
	Key       []byte
	ParamName string
}

func (s Signer) SignURL(u *url.URL) error {
	if u.Opaque != "" {
		return ErrOpaqueURL
	}

	chunks := []string{u.Scheme, u.Host, u.Path, u.RawQuery}
	if u.User != nil {
		chunks = append(chunks, u.String())
	}

	sig := computeSignature(s.Key, chunks)

	q := u.Query()
	q.Set(s.ParamName, sig)
	u.RawQuery = q.Encode()

	return nil
}

func (s Signer) VerifyURL(u *url.URL) error {
	if u.Opaque != "" {
		return ErrOpaqueURL
	}

	q := u.Query()
	sig := q.Get(s.ParamName)
	q.Del(s.ParamName)

	chunks := []string{u.Scheme, u.Host, u.Path, q.Encode()}
	if u.User != nil {
		chunks = append(chunks, u.String())
	}

	expSig := computeSignature(s.Key, chunks)

	if sig != expSig {
		log.Printf("signature mismatch, expected %s but got %s", expSig, sig)
		return ErrSignatureMismatch
	}

	return nil
}

func (s Signer) VerifyReq(req *http.Request) error {
	u := url.URL{Host: req.Host}

	if p := req.Header.Get("X-Forwarded-Proto"); p != "" {
		u.Scheme = p
	} else if req.TLS != nil {
		u.Scheme = "https"
	} else {
		u.Scheme = "http"
	}

	return s.VerifyURL(u.ResolveReference(req.URL))
}

func computeSignature(key []byte, chunks []string) string {
	h := hmac.New(sha512.New, key)
	for _, c := range chunks {
		h.Write([]byte(c))
	}

	return base64.URLEncoding.EncodeToString(h.Sum(nil))
}
<|endoftext|>


<commit_before>package mstate

import (
	"errors"
	"fmt"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"labix.org/v2/mgo/txn"
	"launchpad.net/juju-core/trivial"
)

// ResolvedMode describes the way state transition errors
// are resolved.
type ResolvedMode int

const (
	ResolvedNone ResolvedMode = iota
	ResolvedRetryHooks
	ResolvedNoHooks
	nResolvedModes
)

// AssignmentPolicy controls what machine a unit will be assigned to.
type AssignmentPolicy string

const (
	// AssignLocal indicates that all service units should be assigned
	// to machine 0.
	AssignLocal AssignmentPolicy = "local"
	// AssignUnused indicates that every service unit should be assigned
	// to a dedicated machine, and that new machines should be launched
	// if required.
	AssignUnused AssignmentPolicy = "unused"
)

// NeedsUpgrade describes if a unit needs an
// upgrade and if this is forced.
type NeedsUpgrade struct {
	Upgrade bool
	Force   bool
}

// Port identifies a network port number for a particular protocol.
type Port struct {
	Protocol string `yaml:"proto"`
	Number   int    `yaml:"port"`
}

// UnitSettings holds information about a service unit's settings within a
// relation.
type UnitSettings struct {
	Version  int
	Settings map[string]interface{}
}

// unitDoc represents the internal state of a unit in MongoDB.
type unitDoc struct {
	Name           string `bson:"_id"`
	Service        string
	Principal      string
	PublicAddress  string
	PrivateAddress string
	MachineId      *int
	Resolved       ResolvedMode
	NeedsUpgrade   *NeedsUpgrade
	Life           Life
}

// Unit represents the state of a service unit.
type Unit struct {
	st  *State
	doc unitDoc
}

func newUnit(st *State, udoc *unitDoc) *Unit {
	return &Unit{
		st:  st,
		doc: *udoc,
	}
}

// ServiceName returns the service name.
func (u *Unit) ServiceName() string {
	return u.doc.Service
}

// String returns the unit as string.
func (u *Unit) String() string {
	return u.doc.Name
}

// Name returns the unit name.
func (u *Unit) Name() string {
	return u.doc.Name
}

// Resolved returns the resolved mode for the unit.
func (u *Unit) Resolved() (mode ResolvedMode, err error) {
	return u.doc.Resolved, nil
}

// IsPrincipal returns whether the unit is deployed in its own container,
// and can therefore have subordinate services deployed alongside it.
func (u *Unit) IsPrincipal() bool {
	return u.doc.Principal == ""
}

// PublicAddress returns the public address of the unit.
func (u *Unit) PublicAddress() (string, error) {
	if u.doc.PublicAddress == "" {
		return "", fmt.Errorf("public address of unit %q not found", u)
	}
	return u.doc.PublicAddress, nil
}

// PrivateAddress returns the public address of the unit.
func (u *Unit) PrivateAddress() (string, error) {
	if u.doc.PrivateAddress == "" {
		return "", fmt.Errorf("private address of unit %q not found", u)
	}
	return u.doc.PrivateAddress, nil
}

func (u *Unit) Refresh() error {
	err := u.st.units.FindId(u.doc.Name).One(&u.doc)
	if err != nil {
		return fmt.Errorf("cannot refresh unit %q: %v", u, err)
	}
	return nil
}

// AssignedMachineId returns the id of the assigned machine.
func (u *Unit) AssignedMachineId() (id int, err error) {
	defer trivial.ErrorContextf(&err, "cannot get machine id of unit %q", u)
	if u.IsPrincipal() {
		if u.doc.MachineId == nil {
			return 0, errors.New("unit not assigned to machine")
		}
		return *u.doc.MachineId, nil
	}
	pudoc := unitDoc{}
	sel := bson.D{{"_id", u.doc.Principal}, {"life", Alive}}
	err = u.st.units.Find(sel).One(&pudoc)
	if err != nil {
		return 0, err
	}
	if pudoc.MachineId == nil {
		return 0, errors.New("unit not assigned to machine")
	}
	return *pudoc.MachineId, nil
}

// AssignToMachine assigns this unit to a given machine.
func (u *Unit) AssignToMachine(m *Machine) (err error) {
	sel := bson.D{
		{"_id", u.doc.Name},
		{"$or", []bson.D{
			bson.D{{"machineid", nil}},
			bson.D{{"machineid", m.Id()}},
		}},
		{"life", Alive},
	}
	op := []txn.Operation{{
		Collection: u.st.units.Name,
		DocId:      u.doc.Name,
		Assert:     sel,
		Change:     bson.D{{"$set", bson.D{{"machineid", m.Id()}}}},
	}}
	err = u.st.runner.Run(op, "", nil)
	if err != nil {
		return fmt.Errorf("cannot assign unit %q to machine %s: %v", u, m, err)
	}
	u.doc.MachineId = &m.doc.Id
	return nil
}

// UnassignFromMachine removes the assignment between this unit and the
// machine it's assigned to.
func (u *Unit) UnassignFromMachine() (err error) {
	change := bson.D{{"$set", bson.D{{"machineid", nil}}}}
	sel := bson.D{{"_id", u.doc.Name}}
	err = u.st.units.Update(sel, change)
	if err != nil {
		return fmt.Errorf("cannot unassign unit %q from machine: %v", u, err)
	}
	u.doc.MachineId = nil
	return nil
}

// SetPublicAddress sets the public address of the unit.
func (u *Unit) SetPublicAddress(address string) error {
	change := bson.D{{"$set", bson.D{{"publicaddress", address}}}}
	sel := bson.D{{"_id", u.doc.Name}}
	err := u.st.units.Update(sel, change)
	if err != nil {
		return fmt.Errorf("cannot set public address of unit %q: %v", u, err)
	}
	u.doc.PublicAddress = address
	return nil
}

// SetPrivateAddress sets the public address of the unit.
func (u *Unit) SetPrivateAddress(address string) error {
	change := bson.D{{"$set", bson.D{{"privateaddress", address}}}}
	sel := bson.D{{"_id", u.doc.Name}}
	err := u.st.units.Update(sel, change)
	if err != nil {
		return fmt.Errorf("cannot set private address of unit %q: %v", u, err)
	}
	u.doc.PrivateAddress = address
	return nil
}

// SetResolved marks the unit as having had any previous state transition
// problems resolved, and informs the unit that it may attempt to
// reestablish normal workflow. The resolved mode parameter informs
// whether to attempt to reexecute previous failed hooks or to continue
// as if they had succeeded before.
func (u *Unit) SetResolved(mode ResolvedMode) (err error) {
	defer trivial.ErrorContextf(&err, "cannot set resolved mode for unit %q", u)
	if !(0 <= mode && mode < nResolvedModes) {
		return fmt.Errorf("invalid error resolution mode: %v", mode)
	}
	change := bson.D{{"$set", bson.D{{"resolved", mode}}}}
	sel := bson.D{
		{"_id", u.doc.Name},
		{"resolved", ResolvedNone},
	}
	err = u.st.units.Update(sel, change)
	if err == mgo.ErrNotFound {
		return errors.New("flag already set")
	}
	if err != nil {
		return err
	}
	u.doc.Resolved = mode
	return nil
}

// ClearResolved removes any resolved setting on the unit.
func (u *Unit) ClearResolved() error {
	change := bson.D{{"$set", bson.D{{"resolved", ResolvedNone}}}}
	sel := bson.D{{"_id", u.doc.Name}}
	err := u.st.units.Update(sel, change)
	if err != nil {
		return fmt.Errorf("cannot clear resolved mode for unit %q: %v", u, err)
	}
	u.doc.Resolved = ResolvedNone
	return nil
}

// NeedsUpgrade returns whether the unit needs an upgrade
// and if it does, if this is forced.
func (u *Unit) NeedsUpgrade() (*NeedsUpgrade, error) {
	if u.doc.NeedsUpgrade == nil {
		return &NeedsUpgrade{Upgrade: false, Force: false}, nil
	}
	return u.doc.NeedsUpgrade, nil
}

// SetNeedsUpgrade informs the unit that it should perform
// a regular or forced upgrade.
func (u *Unit) SetNeedsUpgrade(force bool) (err error) {
	defer trivial.ErrorContextf(&err, "cannot inform unit %q about upgrade", u)
	nu := &NeedsUpgrade{Upgrade: true, Force: force}
	change := bson.D{{"$set", bson.D{{"needsupgrade", nu}}}}
	sel := bson.D{
		{"_id", u.doc.Name},
		{"$or", []bson.D{
			bson.D{{"needsupgrade", nil}},
			bson.D{{"needsupgrade", nu}},
		}},
	}
	err = u.st.units.Update(sel, change)
	if err == mgo.ErrNotFound {
		return errors.New("upgrade already enabled")
	}
	if err != nil {
		return err
	}
	u.doc.NeedsUpgrade = nu
	return nil
}

// ClearNeedsUpgrade resets the upgrade notification. It is typically
// done by the unit agent before beginning the upgrade.
func (u *Unit) ClearNeedsUpgrade() error {
	change := bson.D{{"$set", bson.D{{"needsupgrade", nil}}}}
	sel := bson.D{{"_id", u.doc.Name}}
	err := u.st.units.Update(sel, change)
	if err != nil {
		return fmt.Errorf("upgrade notification for unit %q cannot be reset: %v", u, err)
	}
	u.doc.NeedsUpgrade = nil
	return nil
}
<commit_msg>mstate: use txn for UnassignFromMachine<commit_after>package mstate

import (
	"errors"
	"fmt"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"labix.org/v2/mgo/txn"
	"launchpad.net/juju-core/trivial"
)

// ResolvedMode describes the way state transition errors
// are resolved.
type ResolvedMode int

const (
	ResolvedNone ResolvedMode = iota
	ResolvedRetryHooks
	ResolvedNoHooks
	nResolvedModes
)

// AssignmentPolicy controls what machine a unit will be assigned to.
type AssignmentPolicy string

const (
	// AssignLocal indicates that all service units should be assigned
	// to machine 0.
	AssignLocal AssignmentPolicy = "local"
	// AssignUnused indicates that every service unit should be assigned
	// to a dedicated machine, and that new machines should be launched
	// if required.
	AssignUnused AssignmentPolicy = "unused"
)

// NeedsUpgrade describes if a unit needs an
// upgrade and if this is forced.
type NeedsUpgrade struct {
	Upgrade bool
	Force   bool
}

// Port identifies a network port number for a particular protocol.
type Port struct {
	Protocol string `yaml:"proto"`
	Number   int    `yaml:"port"`
}

// UnitSettings holds information about a service unit's settings within a
// relation.
type UnitSettings struct {
	Version  int
	Settings map[string]interface{}
}

// unitDoc represents the internal state of a unit in MongoDB.
type unitDoc struct {
	Name           string `bson:"_id"`
	Service        string
	Principal      string
	PublicAddress  string
	PrivateAddress string
	MachineId      *int
	Resolved       ResolvedMode
	NeedsUpgrade   *NeedsUpgrade
	Life           Life
}

// Unit represents the state of a service unit.
type Unit struct {
	st  *State
	doc unitDoc
}

func newUnit(st *State, udoc *unitDoc) *Unit {
	return &Unit{
		st:  st,
		doc: *udoc,
	}
}

// ServiceName returns the service name.
func (u *Unit) ServiceName() string {
	return u.doc.Service
}

// String returns the unit as string.
func (u *Unit) String() string {
	return u.doc.Name
}

// Name returns the unit name.
func (u *Unit) Name() string {
	return u.doc.Name
}

// Resolved returns the resolved mode for the unit.
func (u *Unit) Resolved() (mode ResolvedMode, err error) {
	return u.doc.Resolved, nil
}

// IsPrincipal returns whether the unit is deployed in its own container,
// and can therefore have subordinate services deployed alongside it.
func (u *Unit) IsPrincipal() bool {
	return u.doc.Principal == ""
}

// PublicAddress returns the public address of the unit.
func (u *Unit) PublicAddress() (string, error) {
	if u.doc.PublicAddress == "" {
		return "", fmt.Errorf("public address of unit %q not found", u)
	}
	return u.doc.PublicAddress, nil
}

// PrivateAddress returns the public address of the unit.
func (u *Unit) PrivateAddress() (string, error) {
	if u.doc.PrivateAddress == "" {
		return "", fmt.Errorf("private address of unit %q not found", u)
	}
	return u.doc.PrivateAddress, nil
}

func (u *Unit) Refresh() error {
	err := u.st.units.FindId(u.doc.Name).One(&u.doc)
	if err != nil {
		return fmt.Errorf("cannot refresh unit %q: %v", u, err)
	}
	return nil
}

// AssignedMachineId returns the id of the assigned machine.
func (u *Unit) AssignedMachineId() (id int, err error) {
	defer trivial.ErrorContextf(&err, "cannot get machine id of unit %q", u)
	if u.IsPrincipal() {
		if u.doc.MachineId == nil {
			return 0, errors.New("unit not assigned to machine")
		}
		return *u.doc.MachineId, nil
	}
	pudoc := unitDoc{}
	sel := bson.D{{"_id", u.doc.Principal}, {"life", Alive}}
	err = u.st.units.Find(sel).One(&pudoc)
	if err != nil {
		return 0, err
	}
	if pudoc.MachineId == nil {
		return 0, errors.New("unit not assigned to machine")
	}
	return *pudoc.MachineId, nil
}

// AssignToMachine assigns this unit to a given machine.
func (u *Unit) AssignToMachine(m *Machine) (err error) {
	sel := bson.D{
		{"_id", u.doc.Name},
		{"$or", []bson.D{
			bson.D{{"machineid", nil}},
			bson.D{{"machineid", m.Id()}},
		}},
		{"life", Alive},
	}
	op := []txn.Operation{{
		Collection: u.st.units.Name,
		DocId:      u.doc.Name,
		Assert:     sel,
		Change:     bson.D{{"$set", bson.D{{"machineid", m.Id()}}}},
	}}
	err = u.st.runner.Run(op, "", nil)
	if err != nil {
		return fmt.Errorf("cannot assign unit %q to machine %s: %v", u, m, err)
	}
	u.doc.MachineId = &m.doc.Id
	return nil
}

// UnassignFromMachine removes the assignment between this unit and the
// machine it's assigned to.
func (u *Unit) UnassignFromMachine() (err error) {
	op := []txn.Operation{{
		Collection: u.st.units.Name,
		DocId:      u.doc.Name,
		Assert:     txn.DocExists,
		Change:     bson.D{{"$set", bson.D{{"machineid", nil}}}},
	}}
	err = u.st.runner.Run(op, "", nil)
	if err != nil {
		return fmt.Errorf("cannot unassign unit %q from machine: %v", u, err)
	}
	u.doc.MachineId = nil
	return nil
}

// SetPublicAddress sets the public address of the unit.
func (u *Unit) SetPublicAddress(address string) error {
	change := bson.D{{"$set", bson.D{{"publicaddress", address}}}}
	sel := bson.D{{"_id", u.doc.Name}}
	err := u.st.units.Update(sel, change)
	if err != nil {
		return fmt.Errorf("cannot set public address of unit %q: %v", u, err)
	}
	u.doc.PublicAddress = address
	return nil
}

// SetPrivateAddress sets the public address of the unit.
func (u *Unit) SetPrivateAddress(address string) error {
	change := bson.D{{"$set", bson.D{{"privateaddress", address}}}}
	sel := bson.D{{"_id", u.doc.Name}}
	err := u.st.units.Update(sel, change)
	if err != nil {
		return fmt.Errorf("cannot set private address of unit %q: %v", u, err)
	}
	u.doc.PrivateAddress = address
	return nil
}

// SetResolved marks the unit as having had any previous state transition
// problems resolved, and informs the unit that it may attempt to
// reestablish normal workflow. The resolved mode parameter informs
// whether to attempt to reexecute previous failed hooks or to continue
// as if they had succeeded before.
func (u *Unit) SetResolved(mode ResolvedMode) (err error) {
	defer trivial.ErrorContextf(&err, "cannot set resolved mode for unit %q", u)
	if !(0 <= mode && mode < nResolvedModes) {
		return fmt.Errorf("invalid error resolution mode: %v", mode)
	}
	change := bson.D{{"$set", bson.D{{"resolved", mode}}}}
	sel := bson.D{
		{"_id", u.doc.Name},
		{"resolved", ResolvedNone},
	}
	err = u.st.units.Update(sel, change)
	if err == mgo.ErrNotFound {
		return errors.New("flag already set")
	}
	if err != nil {
		return err
	}
	u.doc.Resolved = mode
	return nil
}

// ClearResolved removes any resolved setting on the unit.
func (u *Unit) ClearResolved() error {
	change := bson.D{{"$set", bson.D{{"resolved", ResolvedNone}}}}
	sel := bson.D{{"_id", u.doc.Name}}
	err := u.st.units.Update(sel, change)
	if err != nil {
		return fmt.Errorf("cannot clear resolved mode for unit %q: %v", u, err)
	}
	u.doc.Resolved = ResolvedNone
	return nil
}

// NeedsUpgrade returns whether the unit needs an upgrade
// and if it does, if this is forced.
func (u *Unit) NeedsUpgrade() (*NeedsUpgrade, error) {
	if u.doc.NeedsUpgrade == nil {
		return &NeedsUpgrade{Upgrade: false, Force: false}, nil
	}
	return u.doc.NeedsUpgrade, nil
}

// SetNeedsUpgrade informs the unit that it should perform
// a regular or forced upgrade.
func (u *Unit) SetNeedsUpgrade(force bool) (err error) {
	defer trivial.ErrorContextf(&err, "cannot inform unit %q about upgrade", u)
	nu := &NeedsUpgrade{Upgrade: true, Force: force}
	change := bson.D{{"$set", bson.D{{"needsupgrade", nu}}}}
	sel := bson.D{
		{"_id", u.doc.Name},
		{"$or", []bson.D{
			bson.D{{"needsupgrade", nil}},
			bson.D{{"needsupgrade", nu}},
		}},
	}
	err = u.st.units.Update(sel, change)
	if err == mgo.ErrNotFound {
		return errors.New("upgrade already enabled")
	}
	if err != nil {
		return err
	}
	u.doc.NeedsUpgrade = nu
	return nil
}

// ClearNeedsUpgrade resets the upgrade notification. It is typically
// done by the unit agent before beginning the upgrade.
func (u *Unit) ClearNeedsUpgrade() error {
	change := bson.D{{"$set", bson.D{{"needsupgrade", nil}}}}
	sel := bson.D{{"_id", u.doc.Name}}
	err := u.st.units.Update(sel, change)
	if err != nil {
		return fmt.Errorf("upgrade notification for unit %q cannot be reset: %v", u, err)
	}
	u.doc.NeedsUpgrade = nil
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"time"

	"github.com/jaracil/nxcli/nxcore"

	"github.com/gopherjs/gopherjs/js"
	"github.com/jaracil/ei"
)

func ret(r interface{}, e error, cb []*js.Object) {
	switch len(cb) {
	case 1:
		cb[0].Invoke(r, WrapError(e))

	case 2:
		if e == nil {
			cb[0].Invoke(r)
		} else {
			cb[1].Invoke(WrapError(e))
		}
	}
}

func WrapError(e error) *js.Object {
	if e == nil {
		return nil
	}
	jserr := js.Global.Get("Object").New()
	if err, ok := e.(*nxcore.JsonRpcErr); ok {
		jserr.Set("code", err.Cod)
		jserr.Set("message", err.Mess)
		jserr.Set("data", err.Dat)
	} else {
		jserr.Set("code", 0)
		jserr.Set("message", e.Error())
		jserr.Set("data", nil)
	}
	return jserr
}

func WrapTask(task *nxcore.Task) *js.Object {
	if task == nil {
		return nil
	}
	jstask := js.Global.Get("Object").New()
	jstask.Set("sendResult", func(res interface{}, cb ...*js.Object) {
		go func() {
			r, e := task.SendResult(res)
			ret(r, e, cb)
		}()
	})
	jstask.Set("sendError", func(code int, msg string, data interface{}, cb ...*js.Object) {
		go func() {
			r, e := task.SendError(code, msg, data)
			ret(r, e, cb)
		}()
	})
	jstask.Set("path", task.Path)
	jstask.Set("method", task.Method)
	jstask.Set("params", task.Params)
	jstask.Set("tags", task.Tags)
	return jstask
}

func WrapPipe(pipe *nxcore.Pipe) *js.Object {
	if pipe == nil {
		return nil
	}
	jspipe := js.Global.Get("Object").New()
	jspipe.Set("close", func(cb ...*js.Object) {
		go func() {
			r, e := pipe.Close()
			ret(r, e, cb)
		}()
	})
	jspipe.Set("read", func(max int, timeout float64, cb ...*js.Object) {
		go func() {
			r, e := pipe.Read(max, time.Duration(timeout*float64(time.Second)))
			if e != nil {
				ret(nil, e, cb)
				return
			}
			msgs := make([]ei.M, 0)
			for _, msg := range r.Msgs {
				msgs = append(msgs, ei.M{"count": msg.Count, "msg": msg.Msg})
			}
			result := ei.M{"msgs": msgs, "waiting": r.Waiting, "drops": r.Drops}
			ret(result, e, cb)
		}()
	})
	jspipe.Set("write", func(msg interface{}, cb ...*js.Object) {
		go func() {
			r, e := pipe.Write(msg)
			ret(r, e, cb)
		}()
	})
	jspipe.Set("id", pipe.Id())
	return jspipe
}

func WrapNexusConn(nc *nxcore.NexusConn) *js.Object {
	if nc == nil {
		return nil
	}
	jsnc := js.Global.Get("Object").New()
	jsnc.Set("login", func(user string, pass string, cb ...*js.Object) {
		go func() {
			r, e := nc.Login(user, pass)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("taskPush", func(method string, params interface{}, timeout float64, cb ...*js.Object) {
		go func() {
			r, e := nc.TaskPush(method, params, time.Duration(timeout*float64(time.Second)))
			ret(r, e, cb)
		}()
	})
	jsnc.Set("taskPull", func(prefix string, timeout float64, cb ...*js.Object) {
		go func() {
			r, e := nc.TaskPull(prefix, time.Duration(timeout*float64(time.Second)))
			ret(WrapTask(r), e, cb)
		}()
	})
	jsnc.Set("userCreate", func(user string, pass string, cb ...*js.Object) {
		go func() {
			r, e := nc.UserCreate(user, pass)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userDelete", func(user string, cb ...*js.Object) {
		go func() {
			r, e := nc.UserDelete(user)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userSetTags", func(user string, prefix string, tags map[string]interface{}, cb ...*js.Object) {
		go func() {
			r, e := nc.UserSetTags(user, prefix, tags)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userDelTags", func(user string, prefix string, tags []string, cb ...*js.Object) {
		go func() {
			r, e := nc.UserDelTags(user, prefix, tags)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userSetPass", func(user string, pass string, cb ...*js.Object) {
		go func() {
			r, e := nc.UserSetPass(user, pass)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userList", func(prefix string, limit int, skip int, cb ...*js.Object) {
		go func() {
			r, e := nc.UserList(prefix, limit, skip)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userAddTemplate", func(user string, template string, cb ...*js.Object) {
		go func() {
			r, e := nc.UserAddTemplate(user, template)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userDelTemplate", func(user string, template string, cb ...*js.Object) {
		go func() {
			r, e := nc.UserDelTemplate(user, template)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userListTemplate", func(user string, cb ...*js.Object) {
		go func() {
			r, e := nc.UserListTemplate(user)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("sessionList", func(prefix string, limit int, skip int, cb ...*js.Object) {
		go func() {
			r, e := nc.SessionList(prefix, limit, skip)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("sessionKick", func(connId string, cb ...*js.Object) {
		go func() {
			r, e := nc.SessionKick(connId)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("sessionReload", func(connId string, cb ...*js.Object) {
		go func() {
			r, e := nc.SessionReload(connId)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("nodeList", func(limit int, skip int, cb ...*js.Object) {
		go func() {
			r, e := nc.NodeList(limit, skip)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("pipeCreate", func(jopts ei.M, cb ...*js.Object) {
		go func() {
			opts := &nxcore.PipeOpts{
				Length: ei.N(jopts).M("length").IntZ(),
			}
			r, e := nc.PipeCreate(opts)
			ret(WrapPipe(r), e, cb)
		}()
	})
	jsnc.Set("pipeOpen", func(id string, cb ...*js.Object) {
		go func() {
			r, e := nc.PipeOpen(id)
			ret(WrapPipe(r), e, cb)
		}()
	})
	jsnc.Set("topicPublish", func(topic string, msg interface{}, cb ...*js.Object) {
		go func() {
			r, e := nc.TopicPublish(topic, msg)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("topicSubscribe", func(jspipe *js.Object, topic string, cb ...*js.Object) { // !!!Warning, Don't work!!!
		go func() {
			par := ei.M{
				"pipeid": jspipe.Get("id").String(),
				"topic":  topic,
			}
			r, e := nc.Exec("topic.sub", par)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("topicUnsubscribe", func(jspipe *js.Object, topic string, cb ...*js.Object) {
		go func() {
			par := ei.M{
				"pipeid": jspipe.Get("id").String(),
				"topic":  topic,
			}
			r, e := nc.Exec("topic.unsub", par)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("exec", func(method string, params interface{}, cb ...*js.Object) {
		go func() {
			r, e := nc.Exec(method, params)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("close", func(cb ...*js.Object) {
		go func() {
			nc.Close()
			ret(nil, nil, cb)
		}()
	})
	jsnc.Set("closed", nc.Closed)
	jsnc.Set("ping", func(timeout float64, cb ...*js.Object) {
		go func() {
			e := nc.Ping(time.Duration(timeout * float64(time.Second)))
			ret(nil, e, cb)
		}()
	})
	return jsnc
}
<commit_msg>Missing TaskList<commit_after>package main

import (
	"time"

	"github.com/jaracil/nxcli/nxcore"

	"github.com/gopherjs/gopherjs/js"
	"github.com/jaracil/ei"
)

func ret(r interface{}, e error, cb []*js.Object) {
	switch len(cb) {
	case 1:
		cb[0].Invoke(r, WrapError(e))

	case 2:
		if e == nil {
			cb[0].Invoke(r)
		} else {
			cb[1].Invoke(WrapError(e))
		}
	}
}

func WrapError(e error) *js.Object {
	if e == nil {
		return nil
	}
	jserr := js.Global.Get("Object").New()
	if err, ok := e.(*nxcore.JsonRpcErr); ok {
		jserr.Set("code", err.Cod)
		jserr.Set("message", err.Mess)
		jserr.Set("data", err.Dat)
	} else {
		jserr.Set("code", 0)
		jserr.Set("message", e.Error())
		jserr.Set("data", nil)
	}
	return jserr
}

func WrapTask(task *nxcore.Task) *js.Object {
	if task == nil {
		return nil
	}
	jstask := js.Global.Get("Object").New()
	jstask.Set("sendResult", func(res interface{}, cb ...*js.Object) {
		go func() {
			r, e := task.SendResult(res)
			ret(r, e, cb)
		}()
	})
	jstask.Set("sendError", func(code int, msg string, data interface{}, cb ...*js.Object) {
		go func() {
			r, e := task.SendError(code, msg, data)
			ret(r, e, cb)
		}()
	})
	jstask.Set("path", task.Path)
	jstask.Set("method", task.Method)
	jstask.Set("params", task.Params)
	jstask.Set("tags", task.Tags)
	return jstask
}

func WrapPipe(pipe *nxcore.Pipe) *js.Object {
	if pipe == nil {
		return nil
	}
	jspipe := js.Global.Get("Object").New()
	jspipe.Set("close", func(cb ...*js.Object) {
		go func() {
			r, e := pipe.Close()
			ret(r, e, cb)
		}()
	})
	jspipe.Set("read", func(max int, timeout float64, cb ...*js.Object) {
		go func() {
			r, e := pipe.Read(max, time.Duration(timeout*float64(time.Second)))
			if e != nil {
				ret(nil, e, cb)
				return
			}
			msgs := make([]ei.M, 0)
			for _, msg := range r.Msgs {
				msgs = append(msgs, ei.M{"count": msg.Count, "msg": msg.Msg})
			}
			result := ei.M{"msgs": msgs, "waiting": r.Waiting, "drops": r.Drops}
			ret(result, e, cb)
		}()
	})
	jspipe.Set("write", func(msg interface{}, cb ...*js.Object) {
		go func() {
			r, e := pipe.Write(msg)
			ret(r, e, cb)
		}()
	})
	jspipe.Set("id", pipe.Id())
	return jspipe
}

func WrapNexusConn(nc *nxcore.NexusConn) *js.Object {
	if nc == nil {
		return nil
	}
	jsnc := js.Global.Get("Object").New()
	jsnc.Set("login", func(user string, pass string, cb ...*js.Object) {
		go func() {
			r, e := nc.Login(user, pass)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("taskPush", func(method string, params interface{}, timeout float64, cb ...*js.Object) {
		go func() {
			r, e := nc.TaskPush(method, params, time.Duration(timeout*float64(time.Second)))
			ret(r, e, cb)
		}()
	})
	jsnc.Set("taskPull", func(prefix string, timeout float64, cb ...*js.Object) {
		go func() {
			r, e := nc.TaskPull(prefix, time.Duration(timeout*float64(time.Second)))
			ret(WrapTask(r), e, cb)
		}()
	})
	jsnc.Set("taskList", func(prefix string, limit int, skip int, cb ...*js.Object) {
		go func() {
			r, e := nc.TaskList(prefix, limit, skip)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userCreate", func(user string, pass string, cb ...*js.Object) {
		go func() {
			r, e := nc.UserCreate(user, pass)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userDelete", func(user string, cb ...*js.Object) {
		go func() {
			r, e := nc.UserDelete(user)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userSetTags", func(user string, prefix string, tags map[string]interface{}, cb ...*js.Object) {
		go func() {
			r, e := nc.UserSetTags(user, prefix, tags)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userDelTags", func(user string, prefix string, tags []string, cb ...*js.Object) {
		go func() {
			r, e := nc.UserDelTags(user, prefix, tags)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userSetPass", func(user string, pass string, cb ...*js.Object) {
		go func() {
			r, e := nc.UserSetPass(user, pass)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userList", func(prefix string, limit int, skip int, cb ...*js.Object) {
		go func() {
			r, e := nc.UserList(prefix, limit, skip)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userAddTemplate", func(user string, template string, cb ...*js.Object) {
		go func() {
			r, e := nc.UserAddTemplate(user, template)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userDelTemplate", func(user string, template string, cb ...*js.Object) {
		go func() {
			r, e := nc.UserDelTemplate(user, template)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("userListTemplate", func(user string, cb ...*js.Object) {
		go func() {
			r, e := nc.UserListTemplate(user)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("sessionList", func(prefix string, limit int, skip int, cb ...*js.Object) {
		go func() {
			r, e := nc.SessionList(prefix, limit, skip)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("sessionKick", func(connId string, cb ...*js.Object) {
		go func() {
			r, e := nc.SessionKick(connId)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("sessionReload", func(connId string, cb ...*js.Object) {
		go func() {
			r, e := nc.SessionReload(connId)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("nodeList", func(limit int, skip int, cb ...*js.Object) {
		go func() {
			r, e := nc.NodeList(limit, skip)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("pipeCreate", func(jopts ei.M, cb ...*js.Object) {
		go func() {
			opts := &nxcore.PipeOpts{
				Length: ei.N(jopts).M("length").IntZ(),
			}
			r, e := nc.PipeCreate(opts)
			ret(WrapPipe(r), e, cb)
		}()
	})
	jsnc.Set("pipeOpen", func(id string, cb ...*js.Object) {
		go func() {
			r, e := nc.PipeOpen(id)
			ret(WrapPipe(r), e, cb)
		}()
	})
	jsnc.Set("topicPublish", func(topic string, msg interface{}, cb ...*js.Object) {
		go func() {
			r, e := nc.TopicPublish(topic, msg)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("topicSubscribe", func(jspipe *js.Object, topic string, cb ...*js.Object) { // !!!Warning, Don't work!!!
		go func() {
			par := ei.M{
				"pipeid": jspipe.Get("id").String(),
				"topic":  topic,
			}
			r, e := nc.Exec("topic.sub", par)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("topicUnsubscribe", func(jspipe *js.Object, topic string, cb ...*js.Object) {
		go func() {
			par := ei.M{
				"pipeid": jspipe.Get("id").String(),
				"topic":  topic,
			}
			r, e := nc.Exec("topic.unsub", par)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("exec", func(method string, params interface{}, cb ...*js.Object) {
		go func() {
			r, e := nc.Exec(method, params)
			ret(r, e, cb)
		}()
	})
	jsnc.Set("close", func(cb ...*js.Object) {
		go func() {
			nc.Close()
			ret(nil, nil, cb)
		}()
	})
	jsnc.Set("closed", nc.Closed)
	jsnc.Set("ping", func(timeout float64, cb ...*js.Object) {
		go func() {
			e := nc.Ping(time.Duration(timeout * float64(time.Second)))
			ret(nil, e, cb)
		}()
	})
	return jsnc
}
<|endoftext|>