<commit_before>package signalfx

import (
	"bufio"
	"bytes"
	"context"
	"crypto/sha1"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/gorilla/mux"
	"github.com/mailru/easyjson"
	"github.com/signalfx/com_signalfx_metrics_protobuf"
	"github.com/signalfx/gateway/logkey"
	"github.com/signalfx/gateway/protocol/signalfx/format"
	"github.com/signalfx/golib/datapoint"
	"github.com/signalfx/golib/datapoint/dpsink"
	"github.com/signalfx/golib/errors"
	"github.com/signalfx/golib/log"
	"github.com/signalfx/golib/sfxclient"
	"github.com/signalfx/golib/web"
)

// JSONDatapointV1 is an alias
type JSONDatapointV1 signalfxformat.JSONDatapointV1

// JSONDatapointV2 is an alias
type JSONDatapointV2 signalfxformat.JSONDatapointV2

// BodySendFormatV2 is an alias
type BodySendFormatV2 signalfxformat.BodySendFormatV2

// ProtobufDecoderV1 creates datapoints out of the V1 protobuf definition
type ProtobufDecoderV1 struct {
	Sink       dpsink.DSink
	TypeGetter MericTypeGetter
	Logger     log.Logger
}

var errInvalidProtobuf = errors.New("invalid protocol buffer sent")
var errProtobufTooLarge = errors.New("protobuf structure too large")
var errInvalidProtobufVarint = errors.New("invalid protobuf varint")

func (decoder *ProtobufDecoderV1) Read(ctx context.Context, req *http.Request) error {
	body := req.Body
	bufferedBody := bufio.NewReaderSize(body, 32768)
	for {
		_, err := bufferedBody.Peek(1)
		if err == io.EOF {
			return nil
		}
		buf, err := bufferedBody.Peek(4) // should be big enough for any varint

		if err != nil {
			decoder.Logger.Log(log.Err, err, "peek error")
			return err
		}
		num, bytesRead := proto.DecodeVarint(buf)
		if bytesRead == 0 {
			// Invalid varint?
			return errInvalidProtobufVarint
		}
		if num > 32768 {
			// Sanity check
			return errProtobufTooLarge
		}
		// Get the varint out
		buf = make([]byte, bytesRead)
		_, err = io.ReadFull(bufferedBody, buf)
		log.IfErr(decoder.Logger, err)

		// Get the structure out
		buf = make([]byte, num)
		_, err = io.ReadFull(bufferedBody, buf)
		if err != nil {
			return fmt.Errorf("unable to fully read protobuf message: %s", err)
		}
		var msg com_signalfx_metrics_protobuf.DataPoint
		err = proto.Unmarshal(buf, &msg)
		if err != nil {
			return err
		}
		if datapointProtobufIsInvalidForV1(&msg) {
			return errInvalidProtobuf
		}
		mt := decoder.TypeGetter.GetMetricTypeFromMap(msg.GetMetric())
		if dp, err := NewProtobufDataPointWithType(&msg, mt); err == nil {
			log.IfErr(decoder.Logger, decoder.Sink.AddDatapoints(ctx, []*datapoint.Datapoint{dp}))
		}
	}
}

func datapointProtobufIsInvalidForV1(msg *com_signalfx_metrics_protobuf.DataPoint) bool {
	return msg.Metric == nil || msg.Value == nil
}

// JSONDecoderV1 creates datapoints out of the v1 JSON definition
type JSONDecoderV1 struct {
	TypeGetter MericTypeGetter
	Sink       dpsink.DSink
	Logger     log.Logger
}

func (decoder *JSONDecoderV1) Read(ctx context.Context, req *http.Request) error {
	dec := json.NewDecoder(req.Body)
	for {
		var d JSONDatapointV1
		if err := dec.Decode(&d); err == io.EOF {
			break
		} else if err != nil {
			return err
		} else {
			if d.Metric == "" {
				continue
			}
			mt := fromMT(decoder.TypeGetter.GetMetricTypeFromMap(d.Metric))
			dp := datapoint.New(d.Metric, map[string]string{"sf_source": d.Source}, datapoint.NewFloatValue(d.Value), mt, time.Now())
			log.IfErr(decoder.Logger, decoder.Sink.AddDatapoints(ctx, []*datapoint.Datapoint{dp}))
		}
	}
	return nil
}

// ProtobufDecoderV2 decodes protocol buffers in signalfx's v2 format and sends them to Sink
type ProtobufDecoderV2 struct {
	Sink   dpsink.Sink
	Logger log.Logger
}

var buffs = sync.Pool{
	New: func() interface{} {
		return new(bytes.Buffer)
	},
}

func (decoder *ProtobufDecoderV2) Read(ctx context.Context, req *http.Request) (err error) {
	jeff := buffs.Get().(*bytes.Buffer)
	defer buffs.Put(jeff)
	jeff.Reset()
	if err = readFromRequest(jeff, req, decoder.Logger); err != nil {
		return err
	}
	var msg com_signalfx_metrics_protobuf.DataPointUploadMessage
	if err = proto.Unmarshal(jeff.Bytes(), &msg); err != nil {
		return err
	}
	dps := make([]*datapoint.Datapoint, 0, len(msg.GetDatapoints()))
	for _, protoDb := range msg.GetDatapoints() {
		if dp, err1 := NewProtobufDataPointWithType(protoDb, com_signalfx_metrics_protobuf.MetricType_GAUGE); err1 == nil {
			dps = append(dps, dp)
		}
	}
	if len(dps) > 0 {
		err = decoder.Sink.AddDatapoints(ctx, dps)
	}
	return err
}

// JSONDecoderV2 decodes v2 json data for signalfx and sends it to Sink
type JSONDecoderV2 struct {
	Sink              dpsink.Sink
	Logger            log.Logger
	unknownMetricType int64
	invalidValue      int64
}

var errInvalidJSONFormat = errors.New("invalid JSON format; please see correct format at https://developers.signalfx.com/docs/datapoint")

// Datapoints returns datapoints for json decoder v2
func (decoder *JSONDecoderV2) Datapoints() []*datapoint.Datapoint {
	return []*datapoint.Datapoint{
		sfxclient.Counter("dropped_points", map[string]string{"protocol": "sfx_json_v2", "reason": "unknown_metric_type"}, atomic.LoadInt64(&decoder.unknownMetricType)),
		sfxclient.Counter("dropped_points", map[string]string{"protocol": "sfx_json_v2", "reason": "invalid_value"}, atomic.LoadInt64(&decoder.invalidValue)),
	}
}

func (decoder *JSONDecoderV2) Read(ctx context.Context, req *http.Request) error {
	var d signalfxformat.JSONDatapointV2
	if err := easyjson.UnmarshalFromReader(req.Body, &d); err != nil {
		return errInvalidJSONFormat
	}
	dps := make([]*datapoint.Datapoint, 0, len(d))
	for metricType, datapoints := range d {
		if len(datapoints) > 0 {
			mt, ok := com_signalfx_metrics_protobuf.MetricType_value[strings.ToUpper(metricType)]
			if !ok {
				message := make([]interface{}, 0, 7)
				message = append(message, logkey.MetricType, metricType, logkey.Struct, datapoints[0])
				message = append(message, getTokenLogFormat(req)...)
				message = append(message, "Unknown metric type")
				decoder.Logger.Log(message...)
				atomic.AddInt64(&decoder.unknownMetricType, int64(len(datapoints)))
				continue
			}
			for _, jsonDatapoint := range datapoints {
				v, err := ValueToValue(jsonDatapoint.Value)
				if err != nil {
					message := make([]interface{}, 0, 7)
					message = append(message, logkey.Struct, jsonDatapoint, log.Err, err)
					message = append(message, getTokenLogFormat(req)...)
					message = append(message, "Unable to get value for datapoint")
					decoder.Logger.Log(message...)
					atomic.AddInt64(&decoder.invalidValue, 1)
					continue
				}
				dp := datapoint.New(jsonDatapoint.Metric, jsonDatapoint.Dimensions, v, fromMT(com_signalfx_metrics_protobuf.MetricType(mt)), fromTs(jsonDatapoint.Timestamp))
				dps = append(dps, dp)
			}
		}
	}
	if len(dps) == 0 {
		return nil
	}
	return decoder.Sink.AddDatapoints(ctx, dps)
}

func getTokenLogFormat(req *http.Request) (ret []interface{}) {
	h := sha1.New()
	head := req.Header.Get(TokenHeaderName)
	if _, err := io.WriteString(h, head); err != nil || head == "" {
		return ret
	}
	ret = append(ret, logkey.SHA1, base64.StdEncoding.EncodeToString(h.Sum(nil)))
	length := len(head) / 2
	return append(ret, logkey.Caller, head[:length])
}

// SetupProtobufV2DatapointPaths tells the router which paths the given handler (which should handle v2 protobufs)
func SetupProtobufV2DatapointPaths(r *mux.Router, handler http.Handler) {
	SetupProtobufV2ByPaths(r, handler, "/v2/datapoint")
}

// SetupProtobufV2ByPaths tells the router which paths the given handler (which should handle v2 protobufs)
func SetupProtobufV2ByPaths(r *mux.Router, handler http.Handler, path string) {
	r.Path(path).Methods("POST").Headers("Content-Type", "application/x-protobuf").Handler(handler)
}

func setupJSONV2(ctx context.Context, r *mux.Router, sink Sink, logger log.Logger, debugContext *web.HeaderCtxFlag, httpChain web.NextConstructor) sfxclient.Collector {
	var additionalConstructors []web.Constructor
	if debugContext != nil {
		additionalConstructors = append(additionalConstructors, debugContext)
	}
	var j2 *JSONDecoderV2
	handler, st := SetupChain(ctx, sink, "json_v2", func(s Sink) ErrorReader {
		j2 = &JSONDecoderV2{Sink: s, Logger: logger}
		return j2
	}, httpChain, logger, additionalConstructors...)
	multi := sfxclient.NewMultiCollector(st, j2)
	SetupJSONV2DatapointPaths(r, handler)
	return multi
}

// SetupJSONV2DatapointPaths tells the router which paths the given handler (which should handle v2 protobufs)
func SetupJSONV2DatapointPaths(r *mux.Router, handler http.Handler) {
	SetupJSONByPaths(r, handler, "/v2/datapoint")
}

func setupProtobufV1(ctx context.Context, r *mux.Router, sink Sink, typeGetter MericTypeGetter, logger log.Logger, httpChain web.NextConstructor) sfxclient.Collector {
	handler, st := SetupChain(ctx, sink, "protobuf_v1", func(s Sink) ErrorReader {
		return &ProtobufDecoderV1{Sink: s, TypeGetter: typeGetter, Logger: logger}
	}, httpChain, logger)

	SetupProtobufV1Paths(r, handler)
	return st
}

// SetupProtobufV1Paths routes to R paths that should handle V1 Protobuf datapoints
func SetupProtobufV1Paths(r *mux.Router, handler http.Handler) {
	r.Path("/datapoint").Methods("POST").Headers("Content-Type", "application/x-protobuf").Handler(handler)
	r.Path("/v1/datapoint").Methods("POST").Headers("Content-Type", "application/x-protobuf").Handler(handler)
}

func setupJSONV1(ctx context.Context, r *mux.Router, sink Sink, typeGetter MericTypeGetter, logger log.Logger, httpChain web.NextConstructor) sfxclient.Collector {
	handler, st := SetupChain(ctx, sink, "json_v1", func(s Sink) ErrorReader {
		return &JSONDecoderV1{Sink: s, TypeGetter: typeGetter, Logger: logger}
	}, httpChain, logger)
	SetupJSONV1Paths(r, handler)

	return st
}

// SetupJSONV1Paths routes to R paths that should handle V1 JSON datapoints
func SetupJSONV1Paths(r *mux.Router, handler http.Handler) {
	SetupJSONByPaths(r, handler, "/datapoint")
	SetupJSONByPaths(r, handler, "/v1/datapoint")
}

func setupProtobufV2(ctx context.Context, r *mux.Router, sink Sink, logger log.Logger, debugContext *web.HeaderCtxFlag, httpChain web.NextConstructor) sfxclient.Collector {
	var additionalConstructors []web.Constructor
	if debugContext != nil {
		additionalConstructors = append(additionalConstructors, debugContext)
	}
	handler, st := SetupChain(ctx, sink, "protobuf_v2", func(s Sink) ErrorReader {
		return &ProtobufDecoderV2{Sink: s, Logger: logger}
	}, httpChain, logger, additionalConstructors...)
	SetupProtobufV2DatapointPaths(r, handler)

	return st
}
<commit_msg>protocol/signalfx: fix URL of datapoint format reference<commit_after>package signalfx

import (
	"bufio"
	"bytes"
	"context"
	"crypto/sha1"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/gorilla/mux"
	"github.com/mailru/easyjson"
	"github.com/signalfx/com_signalfx_metrics_protobuf"
	"github.com/signalfx/gateway/logkey"
	"github.com/signalfx/gateway/protocol/signalfx/format"
	"github.com/signalfx/golib/datapoint"
	"github.com/signalfx/golib/datapoint/dpsink"
	"github.com/signalfx/golib/errors"
	"github.com/signalfx/golib/log"
	"github.com/signalfx/golib/sfxclient"
	"github.com/signalfx/golib/web"
)

// JSONDatapointV1 is an alias
type JSONDatapointV1 signalfxformat.JSONDatapointV1

// JSONDatapointV2 is an alias
type JSONDatapointV2 signalfxformat.JSONDatapointV2

// BodySendFormatV2 is an alias
type BodySendFormatV2 signalfxformat.BodySendFormatV2

// ProtobufDecoderV1 creates datapoints out of the V1 protobuf definition
type ProtobufDecoderV1 struct {
	Sink       dpsink.DSink
	TypeGetter MericTypeGetter
	Logger     log.Logger
}

var errInvalidProtobuf = errors.New("invalid protocol buffer sent")
var errProtobufTooLarge = errors.New("protobuf structure too large")
var errInvalidProtobufVarint = errors.New("invalid protobuf varint")

func (decoder *ProtobufDecoderV1) Read(ctx context.Context, req *http.Request) error {
	body := req.Body
	bufferedBody := bufio.NewReaderSize(body, 32768)
	for {
		_, err := bufferedBody.Peek(1)
		if err == io.EOF {
			return nil
		}
		buf, err := bufferedBody.Peek(4) // should be big enough for any varint

		if err != nil {
			decoder.Logger.Log(log.Err, err, "peek error")
			return err
		}
		num, bytesRead := proto.DecodeVarint(buf)
		if bytesRead == 0 {
			// Invalid varint?
			return errInvalidProtobufVarint
		}
		if num > 32768 {
			// Sanity check
			return errProtobufTooLarge
		}
		// Get the varint out
		buf = make([]byte, bytesRead)
		_, err = io.ReadFull(bufferedBody, buf)
		log.IfErr(decoder.Logger, err)

		// Get the structure out
		buf = make([]byte, num)
		_, err = io.ReadFull(bufferedBody, buf)
		if err != nil {
			return fmt.Errorf("unable to fully read protobuf message: %s", err)
		}
		var msg com_signalfx_metrics_protobuf.DataPoint
		err = proto.Unmarshal(buf, &msg)
		if err != nil {
			return err
		}
		if datapointProtobufIsInvalidForV1(&msg) {
			return errInvalidProtobuf
		}
		mt := decoder.TypeGetter.GetMetricTypeFromMap(msg.GetMetric())
		if dp, err := NewProtobufDataPointWithType(&msg, mt); err == nil {
			log.IfErr(decoder.Logger, decoder.Sink.AddDatapoints(ctx, []*datapoint.Datapoint{dp}))
		}
	}
}

func datapointProtobufIsInvalidForV1(msg *com_signalfx_metrics_protobuf.DataPoint) bool {
	return msg.Metric == nil || msg.Value == nil
}

// JSONDecoderV1 creates datapoints out of the v1 JSON definition
type JSONDecoderV1 struct {
	TypeGetter MericTypeGetter
	Sink       dpsink.DSink
	Logger     log.Logger
}

func (decoder *JSONDecoderV1) Read(ctx context.Context, req *http.Request) error {
	dec := json.NewDecoder(req.Body)
	for {
		var d JSONDatapointV1
		if err := dec.Decode(&d); err == io.EOF {
			break
		} else if err != nil {
			return err
		} else {
			if d.Metric == "" {
				continue
			}
			mt := fromMT(decoder.TypeGetter.GetMetricTypeFromMap(d.Metric))
			dp := datapoint.New(d.Metric, map[string]string{"sf_source": d.Source}, datapoint.NewFloatValue(d.Value), mt, time.Now())
			log.IfErr(decoder.Logger, decoder.Sink.AddDatapoints(ctx, []*datapoint.Datapoint{dp}))
		}
	}
	return nil
}

// ProtobufDecoderV2 decodes protocol buffers in signalfx's v2 format and sends them to Sink
type ProtobufDecoderV2 struct {
	Sink   dpsink.Sink
	Logger log.Logger
}

var buffs = sync.Pool{
	New: func() interface{} {
		return new(bytes.Buffer)
	},
}

func (decoder *ProtobufDecoderV2) Read(ctx context.Context, req *http.Request) (err error) {
	jeff := buffs.Get().(*bytes.Buffer)
	defer buffs.Put(jeff)
	jeff.Reset()
	if err = readFromRequest(jeff, req, decoder.Logger); err != nil {
		return err
	}
	var msg com_signalfx_metrics_protobuf.DataPointUploadMessage
	if err = proto.Unmarshal(jeff.Bytes(), &msg); err != nil {
		return err
	}
	dps := make([]*datapoint.Datapoint, 0, len(msg.GetDatapoints()))
	for _, protoDb := range msg.GetDatapoints() {
		if dp, err1 := NewProtobufDataPointWithType(protoDb, com_signalfx_metrics_protobuf.MetricType_GAUGE); err1 == nil {
			dps = append(dps, dp)
		}
	}
	if len(dps) > 0 {
		err = decoder.Sink.AddDatapoints(ctx, dps)
	}
	return err
}

// JSONDecoderV2 decodes v2 json data for signalfx and sends it to Sink
type JSONDecoderV2 struct {
	Sink              dpsink.Sink
	Logger            log.Logger
	unknownMetricType int64
	invalidValue      int64
}

var errInvalidJSONFormat = errors.New("invalid JSON format; please see correct format at https://developers.signalfx.com/ingest_data_reference.html")

// Datapoints returns datapoints for json decoder v2
func (decoder *JSONDecoderV2) Datapoints() []*datapoint.Datapoint {
	return []*datapoint.Datapoint{
		sfxclient.Counter("dropped_points", map[string]string{"protocol": "sfx_json_v2", "reason": "unknown_metric_type"}, atomic.LoadInt64(&decoder.unknownMetricType)),
		sfxclient.Counter("dropped_points", map[string]string{"protocol": "sfx_json_v2", "reason": "invalid_value"}, atomic.LoadInt64(&decoder.invalidValue)),
	}
}

func (decoder *JSONDecoderV2) Read(ctx context.Context, req *http.Request) error {
	var d signalfxformat.JSONDatapointV2
	if err := easyjson.UnmarshalFromReader(req.Body, &d); err != nil {
		return errInvalidJSONFormat
	}
	dps := make([]*datapoint.Datapoint, 0, len(d))
	for metricType, datapoints := range d {
		if len(datapoints) > 0 {
			mt, ok := com_signalfx_metrics_protobuf.MetricType_value[strings.ToUpper(metricType)]
			if !ok {
				message := make([]interface{}, 0, 7)
				message = append(message, logkey.MetricType, metricType, logkey.Struct, datapoints[0])
				message = append(message, getTokenLogFormat(req)...)
				message = append(message, "Unknown metric type")
				decoder.Logger.Log(message...)
				atomic.AddInt64(&decoder.unknownMetricType, int64(len(datapoints)))
				continue
			}
			for _, jsonDatapoint := range datapoints {
				v, err := ValueToValue(jsonDatapoint.Value)
				if err != nil {
					message := make([]interface{}, 0, 7)
					message = append(message, logkey.Struct, jsonDatapoint, log.Err, err)
					message = append(message, getTokenLogFormat(req)...)
					message = append(message, "Unable to get value for datapoint")
					decoder.Logger.Log(message...)
					atomic.AddInt64(&decoder.invalidValue, 1)
					continue
				}
				dp := datapoint.New(jsonDatapoint.Metric, jsonDatapoint.Dimensions, v, fromMT(com_signalfx_metrics_protobuf.MetricType(mt)), fromTs(jsonDatapoint.Timestamp))
				dps = append(dps, dp)
			}
		}
	}
	if len(dps) == 0 {
		return nil
	}
	return decoder.Sink.AddDatapoints(ctx, dps)
}

func getTokenLogFormat(req *http.Request) (ret []interface{}) {
	h := sha1.New()
	head := req.Header.Get(TokenHeaderName)
	if _, err := io.WriteString(h, head); err != nil || head == "" {
		return ret
	}
	ret = append(ret, logkey.SHA1, base64.StdEncoding.EncodeToString(h.Sum(nil)))
	length := len(head) / 2
	return append(ret, logkey.Caller, head[:length])
}

// SetupProtobufV2DatapointPaths tells the router which paths the given handler (which should handle v2 protobufs)
func SetupProtobufV2DatapointPaths(r *mux.Router, handler http.Handler) {
	SetupProtobufV2ByPaths(r, handler, "/v2/datapoint")
}

// SetupProtobufV2ByPaths tells the router which paths the given handler (which should handle v2 protobufs)
func SetupProtobufV2ByPaths(r *mux.Router, handler http.Handler, path string) {
	r.Path(path).Methods("POST").Headers("Content-Type", "application/x-protobuf").Handler(handler)
}

func setupJSONV2(ctx context.Context, r *mux.Router, sink Sink, logger log.Logger, debugContext *web.HeaderCtxFlag, httpChain web.NextConstructor) sfxclient.Collector {
	var additionalConstructors []web.Constructor
	if debugContext != nil {
		additionalConstructors = append(additionalConstructors, debugContext)
	}
	var j2 *JSONDecoderV2
	handler, st := SetupChain(ctx, sink, "json_v2", func(s Sink) ErrorReader {
		j2 = &JSONDecoderV2{Sink: s, Logger: logger}
		return j2
	}, httpChain, logger, additionalConstructors...)
	multi := sfxclient.NewMultiCollector(st, j2)
	SetupJSONV2DatapointPaths(r, handler)
	return multi
}

// SetupJSONV2DatapointPaths tells the router which paths the given handler (which should handle v2 protobufs)
func SetupJSONV2DatapointPaths(r *mux.Router, handler http.Handler) {
	SetupJSONByPaths(r, handler, "/v2/datapoint")
}

func setupProtobufV1(ctx context.Context, r *mux.Router, sink Sink, typeGetter MericTypeGetter, logger log.Logger, httpChain web.NextConstructor) sfxclient.Collector {
	handler, st := SetupChain(ctx, sink, "protobuf_v1", func(s Sink) ErrorReader {
		return &ProtobufDecoderV1{Sink: s, TypeGetter: typeGetter, Logger: logger}
	}, httpChain, logger)

	SetupProtobufV1Paths(r, handler)
	return st
}

// SetupProtobufV1Paths routes to R paths that should handle V1 Protobuf datapoints
func SetupProtobufV1Paths(r *mux.Router, handler http.Handler) {
	r.Path("/datapoint").Methods("POST").Headers("Content-Type", "application/x-protobuf").Handler(handler)
	r.Path("/v1/datapoint").Methods("POST").Headers("Content-Type", "application/x-protobuf").Handler(handler)
}

func setupJSONV1(ctx context.Context, r *mux.Router, sink Sink, typeGetter MericTypeGetter, logger log.Logger, httpChain web.NextConstructor) sfxclient.Collector {
	handler, st := SetupChain(ctx, sink, "json_v1", func(s Sink) ErrorReader {
		return &JSONDecoderV1{Sink: s, TypeGetter: typeGetter, Logger: logger}
	}, httpChain, logger)
	SetupJSONV1Paths(r, handler)

	return st
}

// SetupJSONV1Paths routes to R paths that should handle V1 JSON datapoints
func SetupJSONV1Paths(r *mux.Router, handler http.Handler) {
	SetupJSONByPaths(r, handler, "/datapoint")
	SetupJSONByPaths(r, handler, "/v1/datapoint")
}

func setupProtobufV2(ctx context.Context, r *mux.Router, sink Sink, logger log.Logger, debugContext *web.HeaderCtxFlag, httpChain web.NextConstructor) sfxclient.Collector {
	var additionalConstructors []web.Constructor
	if debugContext != nil {
		additionalConstructors = append(additionalConstructors, debugContext)
	}
	handler, st := SetupChain(ctx, sink, "protobuf_v2", func(s Sink) ErrorReader {
		return &ProtobufDecoderV2{Sink: s, Logger: logger}
	}, httpChain, logger, additionalConstructors...)
	SetupProtobufV2DatapointPaths(r, handler)

	return st
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

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

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

package cacheservice

import (
	"fmt"
	"sync"
	"time"
)

// NewConnFunc is a factory method that creates a CacheService instance
// using given CacheServiceConfig.
type NewConnFunc func(config Config) (CacheService, error)

// services stores all supported cache service.
var services = make(map[string]NewConnFunc)

var mu sync.Mutex

// DefaultCacheService decides the default cache service connection.
var DefaultCacheService string

// Config carries config data for CacheService.
type Config struct {
	Address string
	Timeout time.Duration
}

// Result gives the cached data.
type Result struct {
	Key   string
	Value []byte
	Flags uint16
	Cas   uint64
}

// CacheService defines functions to use a cache service.
type CacheService interface {
	// Get returns cached data for given keys.
	Get(keys ...string) (results []Result, err error)
	// Gets returns cached data for given keys, it is an alternative Get api
	// for using with CAS. Gets returns a CAS identifier with the item. If
	// the item's CAS value has changed since you Gets'ed it, it will not be stored.
	Gets(keys ...string) (results []Result, err error)
	// Set sets the value with specified cache key.
	Set(key string, flags uint16, timeout uint64, value []byte) (stored bool, err error)
	// Add stores the value only if it does not already exist.
	Add(key string, flags uint16, timeout uint64, value []byte) (stored bool, err error)
	// Replace replaces the value, only if the value already exists,
	// for the specified cache key.
	Replace(key string, flags uint16, timeout uint64, value []byte) (stored bool, err error)
	// Append appends the value after the last bytes in an existing item.
	Append(key string, flags uint16, timeout uint64, value []byte) (stored bool, err error)
	// Prepend prepends the value before existing value.
	Prepend(key string, flags uint16, timeout uint64, value []byte) (stored bool, err error)
	// Cas stores the value only if no one else has updated the data since you read it last.
	Cas(key string, flags uint16, timeout uint64, value []byte, cas uint64) (stored bool, err error)
	// Delete deletes the value for the specified cache key.
	Delete(key string) (deleted bool, err error)
	// FlushAll purges the entire cache.
	FlushAll() (err error)
	// Stats returns a list of basic stats.
	Stats(argument string) (result []byte, err error)
	// Close closes the CacheService
	Close()
}

// Register a db connection.
func Register(name string, fn NewConnFunc) {
	mu.Lock()
	defer mu.Unlock()
	if _, ok := services[name]; ok {
		panic(fmt.Sprintf("register a registered key: %s", name))
	}
	services[name] = fn
}

// Connect returns a CacheService using the given config.
func Connect(config Config) (CacheService, error) {
	mu.Lock()
	defer mu.Unlock()
	if DefaultCacheService == "" {
		if len(services) == 1 {
			for _, fn := range services {
				return fn(config)
			}
		}
		panic("there are more than one service connect func " +
			"registered but no default cache service has been specified.")
	}
	fn, ok := services[DefaultCacheService]
	if !ok {
		panic(fmt.Sprintf("service connect function for given default cache service: %s is not found.", DefaultCacheService))
	}
	return fn(config)
}
<commit_msg>update cacheservice Using var block (#4906)<commit_after>/*
Copyright 2017 Google Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

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

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

package cacheservice

import (
	"fmt"
	"sync"
	"time"
)

// NewConnFunc is a factory method that creates a CacheService instance
// using given CacheServiceConfig.
type NewConnFunc func(config Config) (CacheService, error)

var (
	// services stores all supported cache service.
	services = make(map[string]NewConnFunc)

	mu sync.Mutex

	// DefaultCacheService decides the default cache service connection.
	DefaultCacheService string
)

// Config carries config data for CacheService.
type Config struct {
	Address string
	Timeout time.Duration
}

// Result gives the cached data.
type Result struct {
	Key   string
	Value []byte
	Flags uint16
	Cas   uint64
}

// CacheService defines functions to use a cache service.
type CacheService interface {
	// Get returns cached data for given keys.
	Get(keys ...string) (results []Result, err error)
	// Gets returns cached data for given keys, it is an alternative Get api
	// for using with CAS. Gets returns a CAS identifier with the item. If
	// the item's CAS value has changed since you Gets'ed it, it will not be stored.
	Gets(keys ...string) (results []Result, err error)
	// Set sets the value with specified cache key.
	Set(key string, flags uint16, timeout uint64, value []byte) (stored bool, err error)
	// Add stores the value only if it does not already exist.
	Add(key string, flags uint16, timeout uint64, value []byte) (stored bool, err error)
	// Replace replaces the value, only if the value already exists,
	// for the specified cache key.
	Replace(key string, flags uint16, timeout uint64, value []byte) (stored bool, err error)
	// Append appends the value after the last bytes in an existing item.
	Append(key string, flags uint16, timeout uint64, value []byte) (stored bool, err error)
	// Prepend prepends the value before existing value.
	Prepend(key string, flags uint16, timeout uint64, value []byte) (stored bool, err error)
	// Cas stores the value only if no one else has updated the data since you read it last.
	Cas(key string, flags uint16, timeout uint64, value []byte, cas uint64) (stored bool, err error)
	// Delete deletes the value for the specified cache key.
	Delete(key string) (deleted bool, err error)
	// FlushAll purges the entire cache.
	FlushAll() (err error)
	// Stats returns a list of basic stats.
	Stats(argument string) (result []byte, err error)
	// Close closes the CacheService
	Close()
}

// Register a db connection.
func Register(name string, fn NewConnFunc) {
	mu.Lock()
	defer mu.Unlock()
	if _, ok := services[name]; ok {
		panic(fmt.Sprintf("register a registered key: %s", name))
	}
	services[name] = fn
}

// Connect returns a CacheService using the given config.
func Connect(config Config) (CacheService, error) {
	mu.Lock()
	defer mu.Unlock()
	if DefaultCacheService == "" {
		if len(services) == 1 {
			for _, fn := range services {
				return fn(config)
			}
		}
		panic("there are more than one service connect func " +
			"registered but no default cache service has been specified.")
	}
	fn, ok := services[DefaultCacheService]
	if !ok {
		panic(fmt.Sprintf("service connect function for given default cache service: %s is not found.", DefaultCacheService))
	}
	return fn(config)
}
<|endoftext|>


<commit_before><commit_msg>fix use of string formatting log calls<commit_after><|endoftext|>


<commit_before>package dockerinit

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"syscall"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/rancher/os/config"
	"github.com/rancher/os/util"
)

const (
	consoleDone = "/run/console-done"
	dockerConf  = "/var/lib/rancher/conf/docker"
	dockerDone  = "/run/docker-done"
	dockerLog   = "/var/log/docker.log"
)

func Main() {
	if os.Getenv("DOCKER_CONF_SOURCED") == "" {
		if err := sourceDockerConf(os.Args); err != nil {
			log.Warnf("Failed to source %s: %v", dockerConf, err)
		}
	}

	for {
		if _, err := os.Stat(consoleDone); err == nil {
			break
		}
		time.Sleep(200 * time.Millisecond)
	}

	dockerBin := "/usr/bin/docker"
	for _, binPath := range []string{
		"/opt/bin",
		"/usr/local/bin",
		"/var/lib/rancher/docker",
	} {
		if util.ExistsAndExecutable(path.Join(binPath, "dockerd")) {
			dockerBin = path.Join(binPath, "dockerd")
			break
		}
		if util.ExistsAndExecutable(path.Join(binPath, "docker")) {
			dockerBin = path.Join(binPath, "docker")
			break
		}
	}

	if err := syscall.Mount("", "/", "", syscall.MS_SHARED|syscall.MS_REC, ""); err != nil {
		log.Error(err)
	}
	if err := syscall.Mount("", "/run", "", syscall.MS_SHARED|syscall.MS_REC, ""); err != nil {
		log.Error(err)
	}

	mountInfo, err := ioutil.ReadFile("/proc/self/mountinfo")
	if err != nil {
		log.Fatal(err)
	}

	for _, mount := range strings.Split(string(mountInfo), "\n") {
		if strings.Contains(mount, "/var/lib/docker /var/lib/docker") && strings.Contains(mount, "rootfs") {
			os.Setenv("DOCKER_RAMDISK", "1")
		}
	}

	args := []string{
		"dockerlaunch",
		dockerBin,
	}

	if len(os.Args) > 1 {
		args = append(args, os.Args[1:]...)
	}

	if os.Getenv("DOCKER_OPTS") != "" {
		args = append(args, os.Getenv("DOCKER_OPTS"))
	}

	cfg := config.LoadConfig()

	if err := ioutil.WriteFile(dockerDone, []byte(cfg.Rancher.Docker.Engine), 0644); err != nil {
		log.Error(err)
	}

	log.Fatal(syscall.Exec("/usr/bin/dockerlaunch", args, os.Environ()))
}

func sourceDockerConf(args []string) error {
	args = append([]string{
		"bash",
		"-c",
		fmt.Sprintf(`[ -e %s ] && source %s; exec docker-init "$@" >> %s 2>&1`, dockerConf, dockerConf, dockerLog),
	}, args...)
	env := os.Environ()
	env = append(env, "DOCKER_CONF_SOURCED=1")
	return syscall.Exec("/bin/bash", args, env)
}
<commit_msg>Modify how /var/lib/rancher/conf/docker is sourced<commit_after>package dockerinit

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"syscall"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/rancher/os/config"
	"github.com/rancher/os/util"
)

const (
	consoleDone = "/run/console-done"
	dockerConf  = "/var/lib/rancher/conf/docker"
	dockerDone  = "/run/docker-done"
	dockerLog   = "/var/log/docker.log"
)

func Main() {
	for {
		if _, err := os.Stat(consoleDone); err == nil {
			break
		}
		time.Sleep(200 * time.Millisecond)
	}

	dockerBin := "/usr/bin/docker"
	for _, binPath := range []string{
		"/opt/bin",
		"/usr/local/bin",
		"/var/lib/rancher/docker",
	} {
		if util.ExistsAndExecutable(path.Join(binPath, "dockerd")) {
			dockerBin = path.Join(binPath, "dockerd")
			break
		}
		if util.ExistsAndExecutable(path.Join(binPath, "docker")) {
			dockerBin = path.Join(binPath, "docker")
			break
		}
	}

	if err := syscall.Mount("", "/", "", syscall.MS_SHARED|syscall.MS_REC, ""); err != nil {
		log.Error(err)
	}
	if err := syscall.Mount("", "/run", "", syscall.MS_SHARED|syscall.MS_REC, ""); err != nil {
		log.Error(err)
	}

	mountInfo, err := ioutil.ReadFile("/proc/self/mountinfo")
	if err != nil {
		log.Fatal(err)
	}

	for _, mount := range strings.Split(string(mountInfo), "\n") {
		if strings.Contains(mount, "/var/lib/docker /var/lib/docker") && strings.Contains(mount, "rootfs") {
			os.Setenv("DOCKER_RAMDISK", "1")
		}
	}

	args := []string{
		"bash",
		"-c",
		fmt.Sprintf(`[ -e %s ] && source %s; exec /usr/bin/dockerlaunch %s %s $DOCKER_OPTS >> %s 2>&1`, dockerConf, dockerConf, dockerBin, strings.Join(os.Args[1:], " "), dockerLog),
	}

	cfg := config.LoadConfig()

	if err := ioutil.WriteFile(dockerDone, []byte(cfg.Rancher.Docker.Engine), 0644); err != nil {
		log.Error(err)
	}

	log.Fatal(syscall.Exec("/bin/bash", args, os.Environ()))
}
<|endoftext|>


<commit_before>package command

import (
	"flag"
	"fmt"
	"sort"
	"strings"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
	"github.com/ryanuber/columnize"
)

type Lags struct {
	Ui  cli.Ui
	Cmd string

	onlineOnly      bool
	groupPattern    string
	topicPattern    string
	problematicMode bool
	lagThreshold    int
}

func (this *Lags) Run(args []string) (exitCode int) {
	const secondsInMinute = 60
	var (
		cluster string
		zone    string
	)
	cmdFlags := flag.NewFlagSet("lags", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&cluster, "c", "", "")
	cmdFlags.BoolVar(&this.onlineOnly, "l", false, "")
	cmdFlags.BoolVar(&this.problematicMode, "p", false, "")
	cmdFlags.StringVar(&this.groupPattern, "g", "", "")
	cmdFlags.StringVar(&this.topicPattern, "t", "", "")
	cmdFlags.IntVar(&this.lagThreshold, "lag", 5000, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).
		require("-c", "-t").
		invalid(args) {
		return 2
	}

	if this.problematicMode {
		this.onlineOnly = true
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	zkcluster := zkzone.NewCluster(cluster) // panic if invalid cluster
	this.printConsumersLag(zkcluster)
	printSwallowedErrors(this.Ui, zkzone)

	return
}

func (this *Lags) printConsumersLag(zkcluster *zk.ZkCluster) {
	lines := make([]string, 0)
	header := "Group|Topic/Partition|Offet|UpTo|Lag|Refreshed"
	lines = append(lines, header)

	// sort by group name
	consumersByGroup := zkcluster.ConsumersByGroup(this.groupPattern)
	sortedGroups := make([]string, 0, len(consumersByGroup))
	for group, _ := range consumersByGroup {
		sortedGroups = append(sortedGroups, group)
	}
	sort.Strings(sortedGroups)

	for _, group := range sortedGroups {
		sortedTopicAndPartitionIds := make([]string, 0)
		consumers := make(map[string]zk.ConsumerMeta)
		for _, t := range consumersByGroup[group] {
			key := fmt.Sprintf("%s:%s", t.Topic, t.PartitionId)
			sortedTopicAndPartitionIds = append(sortedTopicAndPartitionIds, key)

			consumers[key] = t
		}
		sort.Strings(sortedTopicAndPartitionIds)

		for _, topicAndPartitionId := range sortedTopicAndPartitionIds {
			consumer := consumers[topicAndPartitionId]

			if !patternMatched(consumer.Topic, this.topicPattern) {
				continue
			}
			if !consumer.Online {
				continue
			}

			if consumer.Online {
				lines = append(lines,
					fmt.Sprintf("%s|%s/%s|%s|%s|%s|%s",
						group,
						consumer.Topic, consumer.PartitionId,
						gofmt.Comma(consumer.ProducerOffset),
						gofmt.Comma(consumer.ConsumerOffset),
						gofmt.Comma(consumer.Lag),
						gofmt.PrettySince(consumer.Mtime.Time()),
					))
			}
		}
	}

	this.Ui.Output(columnize.SimpleFormat(lines))
}

func (*Lags) Synopsis() string {
	return "Display online high level consumer group lags for a topic"
}

func (this *Lags) Help() string {
	help := fmt.Sprintf(`
Usage: %s lags [options]

    Display online high level consumer group lags for a topic

    e,g.
    %s -z prod -c trade -t OrderStatusMsg

Options:

    -z zone
      Default %s

    -c cluster

    -g consumer group name pattern

    -t topic name pattern

`, this.Cmd, this.Cmd, ctx.ZkDefaultZone())
	return strings.TrimSpace(help)
}
<commit_msg>tweak of the header title<commit_after>package command

import (
	"flag"
	"fmt"
	"sort"
	"strings"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
	"github.com/ryanuber/columnize"
)

type Lags struct {
	Ui  cli.Ui
	Cmd string

	onlineOnly      bool
	groupPattern    string
	topicPattern    string
	problematicMode bool
	lagThreshold    int
}

func (this *Lags) Run(args []string) (exitCode int) {
	const secondsInMinute = 60
	var (
		cluster string
		zone    string
	)
	cmdFlags := flag.NewFlagSet("lags", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&cluster, "c", "", "")
	cmdFlags.BoolVar(&this.onlineOnly, "l", false, "")
	cmdFlags.BoolVar(&this.problematicMode, "p", false, "")
	cmdFlags.StringVar(&this.groupPattern, "g", "", "")
	cmdFlags.StringVar(&this.topicPattern, "t", "", "")
	cmdFlags.IntVar(&this.lagThreshold, "lag", 5000, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).
		require("-c", "-t").
		invalid(args) {
		return 2
	}

	if this.problematicMode {
		this.onlineOnly = true
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	zkcluster := zkzone.NewCluster(cluster) // panic if invalid cluster
	this.printConsumersLag(zkcluster)
	printSwallowedErrors(this.Ui, zkzone)

	return
}

func (this *Lags) printConsumersLag(zkcluster *zk.ZkCluster) {
	lines := make([]string, 0)
	header := "ConsumerGroup|Topic/Partition|Produced|Consumed|Lag|Refreshed"
	lines = append(lines, header)

	// sort by group name
	consumersByGroup := zkcluster.ConsumersByGroup(this.groupPattern)
	sortedGroups := make([]string, 0, len(consumersByGroup))
	for group, _ := range consumersByGroup {
		sortedGroups = append(sortedGroups, group)
	}
	sort.Strings(sortedGroups)

	for _, group := range sortedGroups {
		sortedTopicAndPartitionIds := make([]string, 0)
		consumers := make(map[string]zk.ConsumerMeta)
		for _, t := range consumersByGroup[group] {
			key := fmt.Sprintf("%s:%s", t.Topic, t.PartitionId)
			sortedTopicAndPartitionIds = append(sortedTopicAndPartitionIds, key)

			consumers[key] = t
		}
		sort.Strings(sortedTopicAndPartitionIds)

		for _, topicAndPartitionId := range sortedTopicAndPartitionIds {
			consumer := consumers[topicAndPartitionId]

			if !patternMatched(consumer.Topic, this.topicPattern) {
				continue
			}
			if !consumer.Online {
				continue
			}

			if consumer.Online {
				lines = append(lines,
					fmt.Sprintf("%s|%s/%s|%s|%s|%s|%s",
						group,
						consumer.Topic, consumer.PartitionId,
						gofmt.Comma(consumer.ProducerOffset),
						gofmt.Comma(consumer.ConsumerOffset),
						gofmt.Comma(consumer.Lag),
						gofmt.PrettySince(consumer.Mtime.Time()),
					))
			}
		}
	}

	this.Ui.Output(columnize.SimpleFormat(lines))
}

func (*Lags) Synopsis() string {
	return "Display online high level consumer group lags for a topic"
}

func (this *Lags) Help() string {
	help := fmt.Sprintf(`
Usage: %s lags [options]

    Display online high level consumer group lags for a topic

    e,g.
    %s -z prod -c trade -t OrderStatusMsg

Options:

    -z zone
      Default %s

    -c cluster

    -g consumer group name pattern

    -t topic name pattern

`, this.Cmd, this.Cmd, ctx.ZkDefaultZone())
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>// gddoexp is a command line tool crated to list eligible packages for
// archiving in GoDoc.org
package main

import (
	"fmt"
	"net/http"
	"os"
	"path"

	"github.com/golang/gddo/database"
	"github.com/gregjones/httpcache"
	"github.com/gregjones/httpcache/diskcache"
	"github.com/rafaeljusto/gddoexp"
)

func main() {
	// add cache to avoid to repeated requests to Github
	gddoexp.HTTPClient = &http.Client{
		Transport: httpcache.NewTransport(
			diskcache.New(path.Join(os.Getenv("HOME"), ".gddoexp")),
		),
	}

	db, err := database.New()
	if err != nil {
		fmt.Println("error connecting to database:", err)
		return
	}

	pkgs, err := db.AllPackages()
	if err != nil {
		fmt.Println("error retrieving all packages:", err)
		return
	}

	for response := range gddoexp.ShouldArchivePackages(pkgs, db) {
		if response.Error != nil {
			fmt.Println(err)
		} else if response.Archive {
			fmt.Printf("package “%s” should be archived\n", response.Path)
		}
	}
}
<commit_msg>Add some initial information in command line<commit_after>// gddoexp is a command line tool crated to list eligible packages for
// archiving in GoDoc.org
package main

import (
	"fmt"
	"net/http"
	"os"
	"path"

	"github.com/golang/gddo/database"
	"github.com/gregjones/httpcache"
	"github.com/gregjones/httpcache/diskcache"
	"github.com/rafaeljusto/gddoexp"
)

func main() {
	// add cache to avoid to repeated requests to Github
	gddoexp.HTTPClient = &http.Client{
		Transport: httpcache.NewTransport(
			diskcache.New(path.Join(os.Getenv("HOME"), ".gddoexp")),
		),
	}

	db, err := database.New()
	if err != nil {
		fmt.Println("error connecting to database:", err)
		return
	}

	pkgs, err := db.AllPackages()
	if err != nil {
		fmt.Println("error retrieving all packages:", err)
		return
	}

	fmt.Printf("%d packages will be analyzed\n", len(pkgs))

	for response := range gddoexp.ShouldArchivePackages(pkgs, db) {
		if response.Error != nil {
			fmt.Println(err)
		} else if response.Archive {
			fmt.Printf("package “%s” should be archived\n", response.Path)
		}
	}
}
<|endoftext|>


<commit_before>// The oplog-sync command performs a maintaince operation on the oplog database to keep it
// in sync with the source data.
//
// The command takes a dump of the source data as input and compares it with the oplog's data.
// For any discrepency, a related oplog event is sent to rectify the oplog's database and all
// its consumers.
//
// The dump must be in a streamable JSON format. Each line is a JSON object with the same schema
// as of the data part of the SEE API response:
//
// 	{"timestamp":"2014-11-06T03:04:39.041-08:00", "parents": ["user/xl2d"], "type":"video", "id":"x34cd"}
// 	{"timestamp":"2014-12-24T02:03:05.167+01:00", "parents": ["user/xkwek"], "type":"video", "id":"x12ab"}
// 	{"timestamp":"2014-12-24T01:03:05.167Z", "parents": ["user/xkwek"], "type":"video", "id":"x54cd"}
//
// The timestamp must represent the last modification date of the object as an RFC 3339 representation.
//
// The oplog-sync command is used with this dump in order to perform the sync. This command will connect
// to the database, do the comparisons and generate the necessary oplog events to fix the deltas. This
// command does not need an oplogd agent to be running.
//
// BE CAREFUL, any object absent of the dump having a timestamp lower than the most recent timestamp
// present in the dump will be deleted from the oplog.
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"os"

	log "github.com/Sirupsen/logrus"
	"github.com/dailymotion/oplog"
)

var (
	debug                = flag.Bool("debug", false, "Show debug log messages.")
	mongoURL             = flag.String("mongo-url", "", "MongoDB URL to connect to.")
	cappedCollectionSize = flag.Int("capped-collection-size", 104857600, "Size of the created MongoDB capped collection size in bytes (default 100MB).")
	maxQueuedEvents      = flag.Uint64("max-queued-events", 100000, "Number of events to queue before starting throwing UDP messages.")
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		flag.PrintDefaults()
		fmt.Print("  <dump file>\n")
	}
	flag.Parse()
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(2)
	}
	file := flag.Arg(0)

	if *debug {
		log.SetLevel(log.DebugLevel)
	}

	ol, err := oplog.New(*mongoURL, *cappedCollectionSize)
	if err != nil {
		log.Fatal(err)
	}

	fh, err := os.Open(file)
	if err != nil {
		log.Fatalf("SYNC cannot open dump file: %s", err)
	}
	defer fh.Close()

	lines, err := lineCounter(fh)
	if err != nil {
		log.Fatalf("SYNC error counting lines: %s", err)
	}
	fh.Seek(0, 0)
	createMap := make(oplog.OperationDataMap, lines)
	updateMap := make(oplog.OperationDataMap)
	deleteMap := make(oplog.OperationDataMap)

	// Load dump in memory
	obd := oplog.OperationData{}
	scanner := bufio.NewScanner(fh)
	line := 0
	for scanner.Scan() {
		line++
		if err := json.Unmarshal(scanner.Bytes(), &obd); err != nil {
			log.Fatalf("SYNC dump unmarshaling error at line %d: %s", line, err)
		}
		if err := obd.Validate(); err != nil {
			log.Fatalf("SYNC invalid operation at line %d: %s", line, err)
		}
		createMap[obd.GetId()] = obd
	}
	if err := scanner.Err(); err != nil {
		log.Fatalf("SYNC dump reading error: %s", err)
	}

	// Scan the oplog db and generate the diff
	if err := ol.Diff(createMap, updateMap, deleteMap); err != nil {
		log.Fatalf("SYNC diff error: %s", err)
	}

	// Generate events to fix the delta
	db := ol.DB()
	defer db.Session.Close()
	op := &oplog.Operation{Event: "create"}
	genEvents := func(opMap oplog.OperationDataMap) {
		for _, obd := range opMap {
			op.Data = &obd
			ol.Append(op, db)
		}
	}
	genEvents(createMap)
	op.Event = "update"
	genEvents(updateMap)
	op.Event = "delete"
	genEvents(deleteMap)
}

func lineCounter(r io.Reader) (int, error) {
	buf := make([]byte, 8196)
	count := 0
	lineSep := []byte{'\n'}

	for {
		c, err := r.Read(buf)
		if err != nil && err != io.EOF {
			return count, err
		}

		count += bytes.Count(buf[:c], lineSep)

		if err == io.EOF {
			break
		}
	}

	return count, nil
}
<commit_msg>Add dry run mode + stats logged before syncing<commit_after>// The oplog-sync command performs a maintaince operation on the oplog database to keep it
// in sync with the source data.
//
// The command takes a dump of the source data as input and compares it with the oplog's data.
// For any discrepency, a related oplog event is sent to rectify the oplog's database and all
// its consumers.
//
// The dump must be in a streamable JSON format. Each line is a JSON object with the same schema
// as of the data part of the SEE API response:
//
// 	{"timestamp":"2014-11-06T03:04:39.041-08:00", "parents": ["user/xl2d"], "type":"video", "id":"x34cd"}
// 	{"timestamp":"2014-12-24T02:03:05.167+01:00", "parents": ["user/xkwek"], "type":"video", "id":"x12ab"}
// 	{"timestamp":"2014-12-24T01:03:05.167Z", "parents": ["user/xkwek"], "type":"video", "id":"x54cd"}
//
// The timestamp must represent the last modification date of the object as an RFC 3339 representation.
//
// The oplog-sync command is used with this dump in order to perform the sync. This command will connect
// to the database, do the comparisons and generate the necessary oplog events to fix the deltas. This
// command does not need an oplogd agent to be running.
//
// BE CAREFUL, any object absent of the dump having a timestamp lower than the most recent timestamp
// present in the dump will be deleted from the oplog.
package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"os"

	log "github.com/Sirupsen/logrus"
	"github.com/dailymotion/oplog"
)

var (
	debug                = flag.Bool("debug", false, "Show debug log messages.")
	dryRun               = flag.Bool("dry-run", false, "Compute diff but do not generate events.")
	mongoURL             = flag.String("mongo-url", "", "MongoDB URL to connect to.")
	cappedCollectionSize = flag.Int("capped-collection-size", 104857600, "Size of the created MongoDB capped collection size in bytes (default 100MB).")
	maxQueuedEvents      = flag.Uint64("max-queued-events", 100000, "Number of events to queue before starting throwing UDP messages.")
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		flag.PrintDefaults()
		fmt.Print("  <dump file>\n")
	}
	flag.Parse()
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(2)
	}
	file := flag.Arg(0)

	if *debug {
		log.SetLevel(log.DebugLevel)
	}

	ol, err := oplog.New(*mongoURL, *cappedCollectionSize)
	if err != nil {
		log.Fatal(err)
	}

	createMap := make(oplog.OperationDataMap)
	updateMap := make(oplog.OperationDataMap)
	deleteMap := make(oplog.OperationDataMap)

	var fh *os.File
	if file == "-" {
		fh = os.Stdin
	} else {
		fh, err = os.Open(file)
		if err != nil {
			log.Fatalf("SYNC cannot open dump file: %s", err)
		}
		defer fh.Close()
	}

	// Load dump in memory
	obd := oplog.OperationData{}
	scanner := bufio.NewScanner(fh)
	line := 0
	for scanner.Scan() {
		line++
		if err := json.Unmarshal(scanner.Bytes(), &obd); err != nil {
			log.Fatalf("SYNC dump unmarshaling error at line %d: %s", line, err)
		}
		if err := obd.Validate(); err != nil {
			log.Fatalf("SYNC invalid operation at line %d: %s", line, err)
		}
		createMap[obd.GetId()] = obd
	}
	if err := scanner.Err(); err != nil {
		log.Fatalf("SYNC dump reading error: %s", err)
	}

	// Scan the oplog db and generate the diff
	if err := ol.Diff(createMap, updateMap, deleteMap); err != nil {
		log.Fatalf("SYNC diff error: %s", err)
	}

	log.Infof("SYNC create: %d, update: %d, delete: %d", len(createMap), len(updateMap), len(deleteMap))

	if *dryRun {
		return
	}

	// Generate events to fix the delta
	db := ol.DB()
	defer db.Session.Close()
	op := &oplog.Operation{Event: "create"}
	genEvents := func(opMap oplog.OperationDataMap) {
		for _, obd := range opMap {
			op.Data = &obd
			ol.Append(op, db)
		}
	}
	log.Debugf("SYNC generating %d create events", len(createMap))
	genEvents(createMap)
	log.Debugf("SYNC generating %d update events", len(updateMap))
	op.Event = "update"
	genEvents(updateMap)
	log.Debugf("SYNC generating %d delete events", len(deleteMap))
	op.Event = "delete"
	genEvents(deleteMap)
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Mirantis

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

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

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

package main

import (
	goflag "flag"
	"fmt"
	"math/rand"
	"os"
	"time"

	"github.com/golang/glog"
	flag "github.com/spf13/pflag"
	"k8s.io/client-go/tools/clientcmd"

	"github.com/Mirantis/virtlet/pkg/api/virtlet.k8s/v1"
	"github.com/Mirantis/virtlet/pkg/cni"
	"github.com/Mirantis/virtlet/pkg/config"
	"github.com/Mirantis/virtlet/pkg/diag"
	"github.com/Mirantis/virtlet/pkg/libvirttools"
	"github.com/Mirantis/virtlet/pkg/manager"
	"github.com/Mirantis/virtlet/pkg/nsfix"
	"github.com/Mirantis/virtlet/pkg/tapmanager"
	"github.com/Mirantis/virtlet/pkg/utils"
	"github.com/Mirantis/virtlet/pkg/version"
)

const (
	wantTapManagerEnv  = "WANT_TAP_MANAGER"
	nodeNameEnv        = "KUBE_NODE_NAME"
	diagSocket         = "/run/virtlet-diag.sock"
	netnsDiagCommand   = `if [ -d /var/run/netns ]; then cd /var/run/netns; for ns in *; do echo "*** ${ns} ***"; ip netns exec "${ns}" ip a; ip netns exec "${ns}" ip r; echo; done; fi`
	criproxyLogCommand = `nsenter -t 1 -m -u -i journalctl -xe -u criproxy -n 20000 --no-pager || true`
	qemuLogDir         = "/var/log/libvirt/qemu"
)

var (
	dumpConfig     = flag.Bool("dump-config", false, "Dump node-specific Virtlet config as a shell script and exit")
	dumpDiag       = flag.Bool("diag", false, "Dump diagnostics as JSON and exit")
	displayVersion = flag.Bool("version", false, "Display version and exit")
	versionFormat  = flag.String("version-format", "text", "Version format to use (text, short, json, yaml)")
)

func configWithDefaults(cfg *v1.VirtletConfig) *v1.VirtletConfig {
	r := config.GetDefaultConfig()
	config.Override(r, cfg)
	return r
}

func runVirtlet(config *v1.VirtletConfig, clientCfg clientcmd.ClientConfig, diagSet *diag.Set) {
	manager := manager.NewVirtletManager(config, nil, clientCfg, diagSet)
	if err := manager.Run(); err != nil {
		glog.Errorf("Error: %v", err)
		os.Exit(1)
	}
}

func runTapManager(config *v1.VirtletConfig) {
	cniClient, err := cni.NewClient(*config.CNIPluginDir, *config.CNIConfigDir)
	if err != nil {
		glog.Errorf("Error initializing CNI client: %v", err)
		os.Exit(1)
	}
	src, err := tapmanager.NewTapFDSource(cniClient, *config.EnableSriov, *config.CalicoSubnetSize)
	if err != nil {
		glog.Errorf("Error creating tap fd source: %v", err)
		os.Exit(1)
	}
	os.Remove(*config.FDServerSocketPath) // FIXME
	s := tapmanager.NewFDServer(*config.FDServerSocketPath, src)
	if err = s.Serve(); err != nil {
		glog.Errorf("FD server returned error: %v", err)
		os.Exit(1)
	}
	if err := libvirttools.ChownForEmulator(*config.FDServerSocketPath, false); err != nil {
		glog.Warningf("Couldn't set tapmanager socket permissions: %v", err)
	}
	for {
		time.Sleep(1000 * time.Hour)
	}
}

func printVersion() {
	out, err := version.Get().ToBytes(*versionFormat)
	if err == nil {
		_, err = os.Stdout.Write(out)
	}
	if err != nil {
		glog.Errorf("Error printing version info: %v", err)
		os.Exit(1)
	}
}

func setLogLevel(config *v1.VirtletConfig) {
	goflag.CommandLine.Parse([]string{
		fmt.Sprintf("-v=%d", *config.LogLevel),
		"-logtostderr=true",
	})
}

func runDiagServer() *diag.Set {
	diagSet := diag.NewDiagSet()
	diagSet.RegisterDiagSource("ip-a", diag.NewCommandSource("txt", []string{"ip", "a"}))
	diagSet.RegisterDiagSource("ip-r", diag.NewCommandSource("txt", []string{"ip", "r"}))
	diagSet.RegisterDiagSource("psaux", diag.NewCommandSource("txt", []string{"ps", "aux"}))
	diagSet.RegisterDiagSource("netns", diag.NewCommandSource("txt", []string{"/bin/bash", "-c", netnsDiagCommand}))
	diagSet.RegisterDiagSource("criproxy", diag.NewCommandSource("log", []string{"/bin/bash", "-c", criproxyLogCommand}))
	diagSet.RegisterDiagSource("libvirt-logs", diag.NewLogDirSource(qemuLogDir))
	diagSet.RegisterDiagSource("stack", diag.StackDumpSource)
	server := diag.NewServer(diagSet)
	go func() {
		err := server.Serve(diagSocket, nil)
		glog.V(1).Infof("Diag server returned: %v", err)
	}()
	return diagSet
}

func doDiag() {
	dr, err := diag.RetrieveDiagnostics(diagSocket)
	if err != nil {
		glog.Errorf("Failed to retrieve diagnostics: %v", err)
		os.Exit(1)
	}
	os.Stdout.Write(dr.ToJSON())
}

func main() {
	nsfix.HandleReexec()
	clientCfg := utils.BindFlags(flag.CommandLine)
	var cb *config.Binder
	cb = config.NewBinder(flag.CommandLine)
	flag.Parse()
	localConfig := cb.GetConfig()

	rand.Seed(time.Now().UnixNano())
	setLogLevel(configWithDefaults(localConfig))
	switch {
	case *displayVersion:
		printVersion()
	case *dumpConfig:
		nodeConfig := config.NewNodeConfig(clientCfg)
		nodeName := os.Getenv(nodeNameEnv)
		cfg, err := nodeConfig.LoadConfig(localConfig, nodeName)
		if err != nil {
			glog.Warningf("Failed to load per-node configs, using local config only: %v", err)
			cfg = localConfig
		}
		if _, err := os.Stdout.Write([]byte(config.DumpEnv(cfg))); err != nil {
			glog.Errorf("Error writing config: %v", err)
			os.Exit(1)
		}
	case *dumpDiag:
		doDiag()
	default:
		localConfig = configWithDefaults(localConfig)
		go runTapManager(localConfig)
		diagSet := runDiagServer()
		runVirtlet(localConfig, clientCfg, diagSet)
	}
}
<commit_msg>delete other changes except the deletion of constant "wantTapManagerEnv" in virtlet.go<commit_after>/*
Copyright 2017 Mirantis

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

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

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

package main

import (
	goflag "flag"
	"fmt"
	"math/rand"
	"os"
	"time"

	"github.com/golang/glog"
	flag "github.com/spf13/pflag"
	"k8s.io/client-go/tools/clientcmd"

	"github.com/Mirantis/virtlet/pkg/api/virtlet.k8s/v1"
	"github.com/Mirantis/virtlet/pkg/cni"
	"github.com/Mirantis/virtlet/pkg/config"
	"github.com/Mirantis/virtlet/pkg/diag"
	"github.com/Mirantis/virtlet/pkg/libvirttools"
	"github.com/Mirantis/virtlet/pkg/manager"
	"github.com/Mirantis/virtlet/pkg/nsfix"
	"github.com/Mirantis/virtlet/pkg/tapmanager"
	"github.com/Mirantis/virtlet/pkg/utils"
	"github.com/Mirantis/virtlet/pkg/version"
)

const (
	nodeNameEnv        = "KUBE_NODE_NAME"
	diagSocket         = "/run/virtlet-diag.sock"
	netnsDiagCommand   = `if [ -d /var/run/netns ]; then cd /var/run/netns; for ns in *; do echo "*** ${ns} ***"; ip netns exec "${ns}" ip a; ip netns exec "${ns}" ip r; echo; done; fi`
	criproxyLogCommand = `nsenter -t 1 -m -u -i journalctl -xe -u criproxy -n 20000 --no-pager || true`
	qemuLogDir         = "/var/log/libvirt/qemu"
)

var (
	dumpConfig     = flag.Bool("dump-config", false, "Dump node-specific Virtlet config as a shell script and exit")
	dumpDiag       = flag.Bool("diag", false, "Dump diagnostics as JSON and exit")
	displayVersion = flag.Bool("version", false, "Display version and exit")
	versionFormat  = flag.String("version-format", "text", "Version format to use (text, short, json, yaml)")
)

func configWithDefaults(cfg *v1.VirtletConfig) *v1.VirtletConfig {
	r := config.GetDefaultConfig()
	config.Override(r, cfg)
	return r
}

func runVirtlet(config *v1.VirtletConfig, clientCfg clientcmd.ClientConfig, diagSet *diag.Set) {
	manager := manager.NewVirtletManager(config, nil, clientCfg, diagSet)
	if err := manager.Run(); err != nil {
		glog.Errorf("Error: %v", err)
		os.Exit(1)
	}
}

func runTapManager(config *v1.VirtletConfig) {
	cniClient, err := cni.NewClient(*config.CNIPluginDir, *config.CNIConfigDir)
	if err != nil {
		glog.Errorf("Error initializing CNI client: %v", err)
		os.Exit(1)
	}
	src, err := tapmanager.NewTapFDSource(cniClient, *config.EnableSriov, *config.CalicoSubnetSize)
	if err != nil {
		glog.Errorf("Error creating tap fd source: %v", err)
		os.Exit(1)
	}
	os.Remove(*config.FDServerSocketPath) // FIXME
	s := tapmanager.NewFDServer(*config.FDServerSocketPath, src)
	if err = s.Serve(); err != nil {
		glog.Errorf("FD server returned error: %v", err)
		os.Exit(1)
	}
	if err := libvirttools.ChownForEmulator(*config.FDServerSocketPath, false); err != nil {
		glog.Warningf("Couldn't set tapmanager socket permissions: %v", err)
	}
	for {
		time.Sleep(1000 * time.Hour)
	}
}

func printVersion() {
	out, err := version.Get().ToBytes(*versionFormat)
	if err == nil {
		_, err = os.Stdout.Write(out)
	}
	if err != nil {
		glog.Errorf("Error printing version info: %v", err)
		os.Exit(1)
	}
}

func setLogLevel(config *v1.VirtletConfig) {
	goflag.CommandLine.Parse([]string{
		fmt.Sprintf("-v=%d", *config.LogLevel),
		"-logtostderr=true",
	})
}

func runDiagServer() *diag.Set {
	diagSet := diag.NewDiagSet()
	diagSet.RegisterDiagSource("ip-a", diag.NewCommandSource("txt", []string{"ip", "a"}))
	diagSet.RegisterDiagSource("ip-r", diag.NewCommandSource("txt", []string{"ip", "r"}))
	diagSet.RegisterDiagSource("psaux", diag.NewCommandSource("txt", []string{"ps", "aux"}))
	diagSet.RegisterDiagSource("netns", diag.NewCommandSource("txt", []string{"/bin/bash", "-c", netnsDiagCommand}))
	diagSet.RegisterDiagSource("criproxy", diag.NewCommandSource("log", []string{"/bin/bash", "-c", criproxyLogCommand}))
	diagSet.RegisterDiagSource("libvirt-logs", diag.NewLogDirSource(qemuLogDir))
	diagSet.RegisterDiagSource("stack", diag.StackDumpSource)
	server := diag.NewServer(diagSet)
	go func() {
		err := server.Serve(diagSocket, nil)
		glog.V(1).Infof("Diag server returned: %v", err)
	}()
	return diagSet
}

func doDiag() {
	dr, err := diag.RetrieveDiagnostics(diagSocket)
	if err != nil {
		glog.Errorf("Failed to retrieve diagnostics: %v", err)
		os.Exit(1)
	}
	os.Stdout.Write(dr.ToJSON())
}

func main() {
	nsfix.HandleReexec()
	clientCfg := utils.BindFlags(flag.CommandLine)
	var cb *config.Binder
	cb = config.NewBinder(flag.CommandLine)
	flag.Parse()
	localConfig := cb.GetConfig()

	rand.Seed(time.Now().UnixNano())
	setLogLevel(configWithDefaults(localConfig))
	switch {
	case *displayVersion:
		printVersion()
	case *dumpConfig:
		nodeConfig := config.NewNodeConfig(clientCfg)
		nodeName := os.Getenv(nodeNameEnv)
		cfg, err := nodeConfig.LoadConfig(localConfig, nodeName)
		if err != nil {
			glog.Warningf("Failed to load per-node configs, using local config only: %v", err)
			cfg = localConfig
		}
		if _, err := os.Stdout.Write([]byte(config.DumpEnv(cfg))); err != nil {
			glog.Errorf("Error writing config: %v", err)
			os.Exit(1)
		}
	case *dumpDiag:
		doDiag()
	default:
		localConfig = configWithDefaults(localConfig)
		go runTapManager(localConfig)
		diagSet := runDiagServer()
		runVirtlet(localConfig, clientCfg, diagSet)
	}
}
<|endoftext|>


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

package flow

import (
	"reflect"
	"sort"
	"testing"

	"gonum.org/v1/gonum/graph"
	"gonum.org/v1/gonum/graph/internal/ordered"
	"gonum.org/v1/gonum/graph/simple"
)

var dominatorsTests = []struct {
	n     graph.Node
	edges []simple.Edge

	want DominatorTree
}{
	{ // Example from Lengauer and Tarjan http://www.dtic.mil/dtic/tr/fulltext/u2/a054144.pdf fig 1.
		n: char('R'),
		edges: []simple.Edge{
			{F: char('A'), T: char('D')},
			{F: char('B'), T: char('A')},
			{F: char('B'), T: char('D')},
			{F: char('B'), T: char('E')},
			{F: char('C'), T: char('F')},
			{F: char('C'), T: char('G')},
			{F: char('D'), T: char('L')},
			{F: char('E'), T: char('H')},
			{F: char('F'), T: char('I')},
			{F: char('G'), T: char('I')},
			{F: char('G'), T: char('J')},
			{F: char('H'), T: char('E')},
			{F: char('H'), T: char('K')},
			{F: char('I'), T: char('K')},
			{F: char('J'), T: char('I')},
			{F: char('K'), T: char('I')},
			{F: char('K'), T: char('R')},
			{F: char('L'), T: char('H')},
			{F: char('R'), T: char('A')},
			{F: char('R'), T: char('B')},
			{F: char('R'), T: char('C')},
		},

		want: DominatorTree{
			root: char('R'),
			dominatorOf: map[int64]graph.Node{
				'A': char('R'),
				'B': char('R'),
				'C': char('R'),
				'D': char('R'),
				'E': char('R'),
				'F': char('C'),
				'G': char('C'),
				'H': char('R'),
				'I': char('R'),
				'J': char('G'),
				'K': char('R'),
				'L': char('D'),
			},
			dominatedBy: map[int64][]graph.Node{
				'C': {char('F'), char('G')},
				'D': {char('L')},
				'G': {char('J')},
				'R': {char('A'), char('B'), char('C'), char('D'), char('E'), char('H'), char('I'), char('K')},
			},
		},
	},
	{ // WP example: https://en.wikipedia.org/w/index.php?title=Dominator_(graph_theory)&oldid=758099236.
		n: simple.Node(1),
		edges: []simple.Edge{
			{F: simple.Node(1), T: simple.Node(2)},
			{F: simple.Node(2), T: simple.Node(3)},
			{F: simple.Node(2), T: simple.Node(4)},
			{F: simple.Node(2), T: simple.Node(6)},
			{F: simple.Node(3), T: simple.Node(5)},
			{F: simple.Node(4), T: simple.Node(5)},
			{F: simple.Node(5), T: simple.Node(2)},
		},

		want: DominatorTree{
			root: simple.Node(1),
			dominatorOf: map[int64]graph.Node{
				2: simple.Node(1),
				3: simple.Node(2),
				4: simple.Node(2),
				5: simple.Node(2),
				6: simple.Node(2),
			},
			dominatedBy: map[int64][]graph.Node{
				1: {simple.Node(2)},
				2: {simple.Node(3), simple.Node(4), simple.Node(5), simple.Node(6)},
			},
		},
	},
	{ // WP example with node IDs decremented by 1.
		n: simple.Node(0),
		edges: []simple.Edge{
			{F: simple.Node(0), T: simple.Node(1)},
			{F: simple.Node(1), T: simple.Node(2)},
			{F: simple.Node(1), T: simple.Node(3)},
			{F: simple.Node(1), T: simple.Node(5)},
			{F: simple.Node(2), T: simple.Node(4)},
			{F: simple.Node(3), T: simple.Node(4)},
			{F: simple.Node(4), T: simple.Node(1)},
		},

		want: DominatorTree{
			root: simple.Node(0),
			dominatorOf: map[int64]graph.Node{
				1: simple.Node(0),
				2: simple.Node(1),
				3: simple.Node(1),
				4: simple.Node(1),
				5: simple.Node(1),
			},
			dominatedBy: map[int64][]graph.Node{
				0: {simple.Node(1)},
				1: {simple.Node(2), simple.Node(3), simple.Node(4), simple.Node(5)},
			},
		},
	},
}

type char rune

func (n char) ID() int64      { return int64(n) }
func (n char) String() string { return string(n) }

func TestDominators(t *testing.T) {
	for _, test := range dominatorsTests {
		g := simple.NewDirectedGraph()
		for _, e := range test.edges {
			g.SetEdge(e)
		}

		for _, alg := range []struct {
			name string
			fn   func(graph.Node, graph.Directed) DominatorTree
		}{
			{"Dominators", Dominators},
			{"DominatorsSLT", DominatorsSLT},
		} {
			got := alg.fn(test.n, g)

			if !reflect.DeepEqual(got.Root(), test.want.root) {
				t.Errorf("unexpected dominator tree root from %s: got:%v want:%v",
					alg.name, got.root, test.want.root)
			}

			if !reflect.DeepEqual(got.dominatorOf, test.want.dominatorOf) {
				t.Errorf("unexpected dominator tree from %s: got:%v want:%v",
					alg.name, got.dominatorOf, test.want.dominatorOf)
			}

			for q, want := range test.want.dominatorOf {
				node := got.DominatorOf(q)
				if node != want {
					t.Errorf("unexpected dominator tree result from %s dominated of %v: got:%v want:%v",
						alg.name, q, node, want)
				}
			}

			for _, nodes := range got.dominatedBy {
				sort.Sort(ordered.ByID(nodes))
			}

			if !reflect.DeepEqual(got.dominatedBy, test.want.dominatedBy) {
				t.Errorf("unexpected dominator tree from %s: got:%v want:%v",
					alg.name, got.dominatedBy, test.want.dominatedBy)
			}

			for q, want := range test.want.dominatedBy {
				nodes := got.DominatedBy(q)
				if !reflect.DeepEqual(nodes, want) {
					t.Errorf("unexpected dominator tree result from %s dominated by %v: got:%v want:%v",
						alg.name, q, nodes, want)
				}
			}
		}
	}
}
<commit_msg>graph/flow: improve consistency of coverage reports<commit_after>// Copyright ©2017 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package flow

import (
	"reflect"
	"sort"
	"testing"

	"gonum.org/v1/gonum/graph"
	"gonum.org/v1/gonum/graph/internal/ordered"
	"gonum.org/v1/gonum/graph/simple"
)

var dominatorsTests = []struct {
	n     graph.Node
	edges []simple.Edge

	want DominatorTree
}{
	{ // Example from Lengauer and Tarjan http://www.dtic.mil/dtic/tr/fulltext/u2/a054144.pdf fig 1.
		n: char('R'),
		edges: []simple.Edge{
			{F: char('A'), T: char('D')},
			{F: char('B'), T: char('A')},
			{F: char('B'), T: char('D')},
			{F: char('B'), T: char('E')},
			{F: char('C'), T: char('F')},
			{F: char('C'), T: char('G')},
			{F: char('D'), T: char('L')},
			{F: char('E'), T: char('H')},
			{F: char('F'), T: char('I')},
			{F: char('G'), T: char('I')},
			{F: char('G'), T: char('J')},
			{F: char('H'), T: char('E')},
			{F: char('H'), T: char('K')},
			{F: char('I'), T: char('K')},
			{F: char('J'), T: char('I')},
			{F: char('K'), T: char('I')},
			{F: char('K'), T: char('R')},
			{F: char('L'), T: char('H')},
			{F: char('R'), T: char('A')},
			{F: char('R'), T: char('B')},
			{F: char('R'), T: char('C')},
		},

		want: DominatorTree{
			root: char('R'),
			dominatorOf: map[int64]graph.Node{
				'A': char('R'),
				'B': char('R'),
				'C': char('R'),
				'D': char('R'),
				'E': char('R'),
				'F': char('C'),
				'G': char('C'),
				'H': char('R'),
				'I': char('R'),
				'J': char('G'),
				'K': char('R'),
				'L': char('D'),
			},
			dominatedBy: map[int64][]graph.Node{
				'C': {char('F'), char('G')},
				'D': {char('L')},
				'G': {char('J')},
				'R': {char('A'), char('B'), char('C'), char('D'), char('E'), char('H'), char('I'), char('K')},
			},
		},
	},
	{ // WP example: https://en.wikipedia.org/w/index.php?title=Dominator_(graph_theory)&oldid=758099236.
		n: simple.Node(1),
		edges: []simple.Edge{
			{F: simple.Node(1), T: simple.Node(2)},
			{F: simple.Node(2), T: simple.Node(3)},
			{F: simple.Node(2), T: simple.Node(4)},
			{F: simple.Node(2), T: simple.Node(6)},
			{F: simple.Node(3), T: simple.Node(5)},
			{F: simple.Node(4), T: simple.Node(5)},
			{F: simple.Node(5), T: simple.Node(2)},
		},

		want: DominatorTree{
			root: simple.Node(1),
			dominatorOf: map[int64]graph.Node{
				2: simple.Node(1),
				3: simple.Node(2),
				4: simple.Node(2),
				5: simple.Node(2),
				6: simple.Node(2),
			},
			dominatedBy: map[int64][]graph.Node{
				1: {simple.Node(2)},
				2: {simple.Node(3), simple.Node(4), simple.Node(5), simple.Node(6)},
			},
		},
	},
	{ // WP example with node IDs decremented by 1.
		n: simple.Node(0),
		edges: []simple.Edge{
			{F: simple.Node(0), T: simple.Node(1)},
			{F: simple.Node(1), T: simple.Node(2)},
			{F: simple.Node(1), T: simple.Node(3)},
			{F: simple.Node(1), T: simple.Node(5)},
			{F: simple.Node(2), T: simple.Node(4)},
			{F: simple.Node(3), T: simple.Node(4)},
			{F: simple.Node(4), T: simple.Node(1)},
		},

		want: DominatorTree{
			root: simple.Node(0),
			dominatorOf: map[int64]graph.Node{
				1: simple.Node(0),
				2: simple.Node(1),
				3: simple.Node(1),
				4: simple.Node(1),
				5: simple.Node(1),
			},
			dominatedBy: map[int64][]graph.Node{
				0: {simple.Node(1)},
				1: {simple.Node(2), simple.Node(3), simple.Node(4), simple.Node(5)},
			},
		},
	},
}

type char rune

func (n char) ID() int64      { return int64(n) }
func (n char) String() string { return string(n) }

func TestDominators(t *testing.T) {
	// The dominator functions are non-deterministic due
	// to map iteration ordering, so repeat the tests to
	// ensure consistent coverage. The value of 100 was
	// chosen empirically to have no observed reduction
	// in coverage in several hundred runs of go test -cover.
	for i := 0; i < 100; i++ {
		for _, test := range dominatorsTests {
			g := simple.NewDirectedGraph()
			for _, e := range test.edges {
				g.SetEdge(e)
			}

			for _, alg := range []struct {
				name string
				fn   func(graph.Node, graph.Directed) DominatorTree
			}{
				{"Dominators", Dominators},
				{"DominatorsSLT", DominatorsSLT},
			} {
				got := alg.fn(test.n, g)

				if !reflect.DeepEqual(got.Root(), test.want.root) {
					t.Errorf("unexpected dominator tree root from %s: got:%v want:%v",
						alg.name, got.root, test.want.root)
				}

				if !reflect.DeepEqual(got.dominatorOf, test.want.dominatorOf) {
					t.Errorf("unexpected dominator tree from %s: got:%v want:%v",
						alg.name, got.dominatorOf, test.want.dominatorOf)
				}

				for q, want := range test.want.dominatorOf {
					node := got.DominatorOf(q)
					if node != want {
						t.Errorf("unexpected dominator tree result from %s dominated of %v: got:%v want:%v",
							alg.name, q, node, want)
					}
				}

				for _, nodes := range got.dominatedBy {
					sort.Sort(ordered.ByID(nodes))
				}

				if !reflect.DeepEqual(got.dominatedBy, test.want.dominatedBy) {
					t.Errorf("unexpected dominator tree from %s: got:%v want:%v",
						alg.name, got.dominatedBy, test.want.dominatedBy)
				}

				for q, want := range test.want.dominatedBy {
					nodes := got.DominatedBy(q)
					if !reflect.DeepEqual(nodes, want) {
						t.Errorf("unexpected dominator tree result from %s dominated by %v: got:%v want:%v",
							alg.name, q, nodes, want)
					}
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>package schemabuilder

import (
	"context"
	"fmt"
	"reflect"
	"sort"

	"github.com/samsarahq/thunder/graphql"
)

func (sb *schemaBuilder) buildStruct(typ reflect.Type) error {
	if sb.types[typ] != nil {
		return nil
	}

	if typ == unionType {
		return fmt.Errorf("schemabuilder.Union can only be used as an embedded anonymous non-pointer struct")
	}

	if hasUnionMarkerEmbedded(typ) {
		return sb.buildUnionStruct(typ)
	}

	var name string
	var description string
	var methods Methods
	var objectKey string
	if object, ok := sb.objects[typ]; ok {
		name = object.Name
		description = object.Description
		methods = object.Methods
		objectKey = object.key
	}

	if name == "" {
		name = typ.Name()
		if name == "" {
			return fmt.Errorf("bad type %s: should have a name", typ)
		}
	}

	object := &graphql.Object{
		Name:        name,
		Description: description,
		Fields:      make(map[string]*graphql.Field),
	}
	sb.types[typ] = object

	for i := 0; i < typ.NumField(); i++ {
		field := typ.Field(i)
		fieldInfo, err := parseGraphQLFieldInfo(field)
		if err != nil {
			return fmt.Errorf("bad type %s: %s", typ, fieldInfo.Name)
		}
		if fieldInfo.Skipped {
			continue
		}

		if _, ok := object.Fields[fieldInfo.Name]; ok {
			return fmt.Errorf("bad type %s: two fields named %s", typ, fieldInfo.Name)
		}

		built, err := sb.buildField(field)
		if err != nil {
			return fmt.Errorf("bad field %s on type %s: %s", fieldInfo.Name, typ, err)
		}
		object.Fields[fieldInfo.Name] = built
		if fieldInfo.KeyField {
			if object.Key != nil {
				return fmt.Errorf("bad type %s: multiple key fields", typ)
			}
			if !isScalarType(built.Type) {
				return fmt.Errorf("bad type %s: key type must be scalar, got %T", typ, built.Type)
			}
			object.Key = built.Resolve
		}
	}

	var names []string
	for name := range methods {
		names = append(names, name)
	}
	sort.Strings(names)

	for _, name := range names {
		method := methods[name]

		if method.Paginated {
			typedField, err := sb.buildPaginatedField(typ, method)
			if err != nil {
				return err
			}
			object.Fields[name] = typedField
			continue
		}

		built, err := sb.buildFunction(typ, method)
		if err != nil {
			return fmt.Errorf("bad method %s on type %s: %s", name, typ, err)
		}
		object.Fields[name] = built
	}

	if objectKey != "" {
		keyPtr, ok := object.Fields[objectKey]
		if !ok {
			return fmt.Errorf("key field doesn't exist on object")
		}

		if !isScalarType(keyPtr.Type) {
			return fmt.Errorf("bad type %s: key type must be scalar, got %s", typ, keyPtr.Type.String())
		}
		object.Key = keyPtr.Resolve
	}

	return nil
}

func hasUnionMarkerEmbedded(typ reflect.Type) bool {
	for i := 0; i < typ.NumField(); i++ {
		field := typ.Field(i)
		if field.Anonymous && field.Type == unionType {
			return true
		}
	}
	return false
}

func (sb *schemaBuilder) buildUnionStruct(typ reflect.Type) error {
	var name string
	var description string

	if name == "" {
		name = typ.Name()
		if name == "" {
			return fmt.Errorf("bad type %s: should have a name", typ)
		}
	}

	union := &graphql.Union{
		Name:        name,
		Description: description,
		Types:       make(map[string]*graphql.Object),
	}
	sb.types[typ] = union

	for i := 0; i < typ.NumField(); i++ {
		field := typ.Field(i)
		if field.PkgPath != "" || (field.Anonymous && field.Type == unionType) {
			continue
		}

		if !field.Anonymous {
			return fmt.Errorf("bad type %s: union type member types must be anonymous", name)
		}

		typ, err := sb.getType(field.Type)
		if err != nil {
			return err
		}

		obj, ok := typ.(*graphql.Object)
		if !ok {
			return fmt.Errorf("bad type %s: union type member must be a pointer to a struct, received %s", name, typ.String())
		}

		if union.Types[obj.Name] != nil {
			return fmt.Errorf("bad type %s: union type member may only appear once", name)
		}

		union.Types[obj.Name] = obj
	}
	return nil
}

func isScalarType(typ graphql.Type) bool {
	if nonNull, ok := typ.(*graphql.NonNull); ok {
		typ = nonNull.Type
	}
	if _, ok := typ.(*graphql.Scalar); !ok {
		return false
	}
	return true
}

func (sb *schemaBuilder) buildField(field reflect.StructField) (*graphql.Field, error) {
	retType, err := sb.getType(field.Type)
	if err != nil {
		return nil, err
	}

	return &graphql.Field{
		Resolve: func(ctx context.Context, source, args interface{}, selectionSet *graphql.SelectionSet) (interface{}, error) {
			value := reflect.ValueOf(source)
			if value.Kind() == reflect.Ptr {
				value = value.Elem()
			}
			return value.FieldByIndex(field.Index).Interface(), nil
		},
		Type:           retType,
		ParseArguments: nilParseArguments,
	}, nil
}
<commit_msg>go schemabuilder/output: Add comments for field output functions<commit_after>package schemabuilder

import (
	"context"
	"fmt"
	"reflect"
	"sort"

	"github.com/samsarahq/thunder/graphql"
)

// buildStruct is a function for building the graphql.Type for a passed in
// struct type.  This function reads the "Object" information and Fields of the
// passed in struct to create a "graphql.Object" type that represents this type
// and can be used to resolve graphql requests.
func (sb *schemaBuilder) buildStruct(typ reflect.Type) error {
	if sb.types[typ] != nil {
		return nil
	}

	if typ == unionType {
		return fmt.Errorf("schemabuilder.Union can only be used as an embedded anonymous non-pointer struct")
	}

	if hasUnionMarkerEmbedded(typ) {
		return sb.buildUnionStruct(typ)
	}

	var name string
	var description string
	var methods Methods
	var objectKey string
	if object, ok := sb.objects[typ]; ok {
		name = object.Name
		description = object.Description
		methods = object.Methods
		objectKey = object.key
	}

	if name == "" {
		name = typ.Name()
		if name == "" {
			return fmt.Errorf("bad type %s: should have a name", typ)
		}
	}

	object := &graphql.Object{
		Name:        name,
		Description: description,
		Fields:      make(map[string]*graphql.Field),
	}
	sb.types[typ] = object

	for i := 0; i < typ.NumField(); i++ {
		field := typ.Field(i)
		fieldInfo, err := parseGraphQLFieldInfo(field)
		if err != nil {
			return fmt.Errorf("bad type %s: %s", typ, fieldInfo.Name)
		}
		if fieldInfo.Skipped {
			continue
		}

		if _, ok := object.Fields[fieldInfo.Name]; ok {
			return fmt.Errorf("bad type %s: two fields named %s", typ, fieldInfo.Name)
		}

		built, err := sb.buildField(field)
		if err != nil {
			return fmt.Errorf("bad field %s on type %s: %s", fieldInfo.Name, typ, err)
		}
		object.Fields[fieldInfo.Name] = built
		if fieldInfo.KeyField {
			if object.Key != nil {
				return fmt.Errorf("bad type %s: multiple key fields", typ)
			}
			if !isScalarType(built.Type) {
				return fmt.Errorf("bad type %s: key type must be scalar, got %T", typ, built.Type)
			}
			object.Key = built.Resolve
		}
	}

	var names []string
	for name := range methods {
		names = append(names, name)
	}
	sort.Strings(names)

	for _, name := range names {
		method := methods[name]

		if method.Paginated {
			typedField, err := sb.buildPaginatedField(typ, method)
			if err != nil {
				return err
			}
			object.Fields[name] = typedField
			continue
		}

		built, err := sb.buildFunction(typ, method)
		if err != nil {
			return fmt.Errorf("bad method %s on type %s: %s", name, typ, err)
		}
		object.Fields[name] = built
	}

	if objectKey != "" {
		keyPtr, ok := object.Fields[objectKey]
		if !ok {
			return fmt.Errorf("key field doesn't exist on object")
		}

		if !isScalarType(keyPtr.Type) {
			return fmt.Errorf("bad type %s: key type must be scalar, got %s", typ, keyPtr.Type.String())
		}
		object.Key = keyPtr.Resolve
	}

	return nil
}

// hasUnionMarkerEmbedded determines if a struct has an embedded schemabuilder.Union
// field embedded on the type.
func hasUnionMarkerEmbedded(typ reflect.Type) bool {
	for i := 0; i < typ.NumField(); i++ {
		field := typ.Field(i)
		if field.Anonymous && field.Type == unionType {
			return true
		}
	}
	return false
}

// buildUnionStruct builds a graphql.Union type that has the ability to be one
// of many different graphql types.
func (sb *schemaBuilder) buildUnionStruct(typ reflect.Type) error {
	var name string
	var description string

	if name == "" {
		name = typ.Name()
		if name == "" {
			return fmt.Errorf("bad type %s: should have a name", typ)
		}
	}

	union := &graphql.Union{
		Name:        name,
		Description: description,
		Types:       make(map[string]*graphql.Object),
	}
	sb.types[typ] = union

	for i := 0; i < typ.NumField(); i++ {
		field := typ.Field(i)
		if field.PkgPath != "" || (field.Anonymous && field.Type == unionType) {
			continue
		}

		if !field.Anonymous {
			return fmt.Errorf("bad type %s: union type member types must be anonymous", name)
		}

		typ, err := sb.getType(field.Type)
		if err != nil {
			return err
		}

		obj, ok := typ.(*graphql.Object)
		if !ok {
			return fmt.Errorf("bad type %s: union type member must be a pointer to a struct, received %s", name, typ.String())
		}

		if union.Types[obj.Name] != nil {
			return fmt.Errorf("bad type %s: union type member may only appear once", name)
		}

		union.Types[obj.Name] = obj
	}
	return nil
}

// isScalarType returns whether a graphql.Type is a scalar type (or a non-null
// wrapped scalar type).
func isScalarType(typ graphql.Type) bool {
	if nonNull, ok := typ.(*graphql.NonNull); ok {
		typ = nonNull.Type
	}
	if _, ok := typ.(*graphql.Scalar); !ok {
		return false
	}
	return true
}

// buildField generates a graphQL field for a struct's field.  This field can be
// used to "resolve" a response for a graphql request.
func (sb *schemaBuilder) buildField(field reflect.StructField) (*graphql.Field, error) {
	retType, err := sb.getType(field.Type)
	if err != nil {
		return nil, err
	}

	return &graphql.Field{
		Resolve: func(ctx context.Context, source, args interface{}, selectionSet *graphql.SelectionSet) (interface{}, error) {
			value := reflect.ValueOf(source)
			if value.Kind() == reflect.Ptr {
				value = value.Elem()
			}
			return value.FieldByIndex(field.Index).Interface(), nil
		},
		Type:           retType,
		ParseArguments: nilParseArguments,
	}, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2015 Arista Networks, Inc.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package test

import (
	"fmt"
	"runtime"
	"testing"
)

// ShouldPanic will test is a function is panicking
func ShouldPanic(t *testing.T, fn func()) {
	defer func() {
		if r := recover(); r == nil {
			t.Errorf("%sThe function %p should have panicked",
				getCallerInfo(), fn)
		}
	}()

	fn()
}

// ShouldPanicWith will test is a function is panicking with a specific message
func ShouldPanicWith(t *testing.T, msg interface{}, fn func()) {
	defer func() {
		if r := recover(); r == nil {
			t.Errorf("%sThe function %p should have panicked",
				getCallerInfo(), fn)
		} else if d := Diff(msg, r); len(d) != 0 {
			t.Errorf("%sThe function %p panicked with the wrong message.\n"+
				"Expected: %#v\nReceived: %#v\nDiff:%s",
				getCallerInfo(), fn, msg, r, d)
		}
	}()

	fn()
}

func getCallerInfo() string {
	_, file, line, ok := runtime.Caller(4)
	if !ok {
		return ""
	}
	return fmt.Sprintf("%s:%d\n", file, line)
}
<commit_msg>test/panic: use t.Helper()<commit_after>// Copyright (c) 2015 Arista Networks, Inc.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package test

import (
	"fmt"
	"runtime"
	"testing"
)

// ShouldPanic will test is a function is panicking
func ShouldPanic(t *testing.T, fn func()) {
	t.Helper()
	defer func() {
		t.Helper()
		if r := recover(); r == nil {
			t.Errorf("%sThe function %p should have panicked",
				getCallerInfo(), fn)
		}
	}()

	fn()
}

// ShouldPanicWith will test is a function is panicking with a specific message
func ShouldPanicWith(t *testing.T, msg interface{}, fn func()) {
	t.Helper()
	defer func() {
		t.Helper()
		if r := recover(); r == nil {
			t.Errorf("%sThe function %p should have panicked",
				getCallerInfo(), fn)
		} else if d := Diff(msg, r); len(d) != 0 {
			t.Errorf("%sThe function %p panicked with the wrong message.\n"+
				"Expected: %#v\nReceived: %#v\nDiff:%s",
				getCallerInfo(), fn, msg, r, d)
		}
	}()

	fn()
}

func getCallerInfo() string {
	_, file, line, ok := runtime.Caller(4)
	if !ok {
		return ""
	}
	return fmt.Sprintf("%s:%d\n", file, line)
}
<|endoftext|>


<commit_before>package engine

import (
	"fmt"

	"github.com/keybase/client/go/libkb"
	triplesec "github.com/keybase/go-triplesec"
)

type SignupEngine struct {
	pwsalt     []byte
	tspkey     libkb.PassphraseStream
	uid        libkb.UID
	me         *libkb.User
	signingKey libkb.GenericKey
	arg        *SignupEngineRunArg
	lks        *libkb.LKSec
	libkb.Contextified
}

type SignupEngineRunArg struct {
	Username   string
	Email      string
	InviteCode string
	Passphrase string
	DeviceName string
	SkipGPG    bool
	SkipMail   bool
}

func NewSignupEngine(arg *SignupEngineRunArg, g *libkb.GlobalContext) *SignupEngine {
	return &SignupEngine{
		arg:          arg,
		Contextified: libkb.NewContextified(g),
	}
}

func (s *SignupEngine) Name() string {
	return "Signup"
}

func (s *SignupEngine) RequiredUIs() []libkb.UIKind {
	return nil
}

func (s *SignupEngine) GetPrereqs() EnginePrereqs { return EnginePrereqs{} }

func (s *SignupEngine) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{
		&DetKeyEngine{},
		&GPGImportKeyEngine{},
		&DeviceWrap{},
	}
}

func (s *SignupEngine) Init() error {
	return nil
}

func (s *SignupEngine) SetArg(arg *SignupEngineRunArg) {
	s.arg = arg
}

func (s *SignupEngine) CheckRegistered() (err error) {
	s.G().Log.Debug("+ libkb.SignupJoinEngine::CheckRegistered")
	if cr := s.G().Env.GetConfig(); cr == nil {
		err = fmt.Errorf("No configuration file available")
	} else if u := cr.GetUID(); u != nil {
		err = libkb.AlreadyRegisteredError{Uid: *u}
	}
	s.G().Log.Debug("- libkb.SignupJoinEngine::CheckRegistered -> %s", libkb.ErrToOk(err))
	return err
}

func (s *SignupEngine) PostInviteRequest(arg libkb.InviteRequestArg) error {
	return libkb.PostInviteRequest(arg)
}

func (s *SignupEngine) GetMe() *libkb.User {
	return s.me
}

func (s *SignupEngine) Run(ctx *Context) error {
	// make sure we're starting with a clear login state:
	s.G().Logout()

	f := func() error {
		if err := s.genTSPassKey(s.arg.Passphrase); err != nil {
			return err
		}

		if err := s.join(s.arg.Username, s.arg.Email, s.arg.InviteCode, s.arg.SkipMail); err != nil {
			return err
		}

		if err := s.registerDevice(ctx, s.arg.DeviceName); err != nil {
			return err
		}

		if err := s.genDetKeys(ctx); err != nil {
			return fmt.Errorf("detkeys error: %s", err)
		}

		if s.arg.SkipGPG {
			return nil
		}

		if wantsGPG, err := s.checkGPG(ctx); err != nil {
			return err
		} else if wantsGPG {
			if err := s.addGPG(ctx, true); err != nil {
				return fmt.Errorf("addGPG error: %s", err)
			}
		}

		return nil
	}
	return s.G().LoginState().ExternalFunc(f, "SignupEngine - Run")
}

func (s *SignupEngine) genTSPassKey(passphrase string) error {
	salt, err := libkb.RandBytes(triplesec.SaltLen)
	if err != nil {
		return err
	}
	s.pwsalt = salt
	var tsec *triplesec.Cipher
	tsec, s.tspkey, err = libkb.StretchPassphrase(passphrase, salt)
	s.G().LoginState().SetSignupStreamCache(tsec, s.tspkey)
	return err
}

func (s *SignupEngine) join(username, email, inviteCode string, skipMail bool) error {
	joinEngine := NewSignupJoinEngine(s.G())

	arg := SignupJoinEngineRunArg{
		Username:   username,
		Email:      email,
		InviteCode: inviteCode,
		PWHash:     s.tspkey.PWHash(),
		PWSalt:     s.pwsalt,
		SkipMail:   skipMail,
	}
	res := joinEngine.Run(arg)
	if res.Err != nil {
		return res
	}

	s.uid = *res.UID
	user, err := libkb.LoadUser(libkb.LoadUserArg{Uid: res.UID, PublicKeyOptional: true})
	if err != nil {
		return err
	}
	s.me = user
	return nil
}

func (s *SignupEngine) registerDevice(ctx *Context, deviceName string) error {
	s.lks = libkb.NewLKSec(s.tspkey.LksClientHalf(), s.G())
	args := &DeviceWrapArgs{
		Me:         s.me,
		DeviceName: deviceName,
		Lks:        s.lks,
		IsEldest:   true,
	}
	eng := NewDeviceWrap(args, s.G())
	if err := RunEngine(eng, ctx); err != nil {
		return err
	}
	s.signingKey = eng.SigningKey()

	return nil
}

func (s *SignupEngine) genDetKeys(ctx *Context) error {
	arg := &DetKeyArgs{
		Tsp:         s.tspkey,
		Me:          s.me,
		SigningKey:  s.signingKey,
		EldestKeyID: s.signingKey.GetKid(),
	}
	eng := NewDetKeyEngine(arg, s.G())
	return RunEngine(eng, ctx)
}

func (s *SignupEngine) checkGPG(ctx *Context) (bool, error) {
	eng := NewGPGImportKeyEngine(nil, s.G())
	return eng.WantsGPG(ctx)
}

func (s *SignupEngine) addGPG(ctx *Context, allowMulti bool) error {
	s.G().Log.Debug("SignupEngine.addGPG.  signingKey: %v\n", s.signingKey)
	arg := GPGImportKeyArg{Signer: s.signingKey, AllowMulti: allowMulti, Me: s.me, Lks: s.lks}
	eng := NewGPGImportKeyEngine(&arg, s.G())
	if err := RunEngine(eng, ctx); err != nil {
		return err
	}

	if s.signingKey == nil {
		s.signingKey = eng.LastKey()
	}
	return nil
}
<commit_msg>passing Self: true to LoadUser<commit_after>package engine

import (
	"fmt"

	"github.com/keybase/client/go/libkb"
	triplesec "github.com/keybase/go-triplesec"
)

type SignupEngine struct {
	pwsalt     []byte
	tspkey     libkb.PassphraseStream
	uid        libkb.UID
	me         *libkb.User
	signingKey libkb.GenericKey
	arg        *SignupEngineRunArg
	lks        *libkb.LKSec
	libkb.Contextified
}

type SignupEngineRunArg struct {
	Username   string
	Email      string
	InviteCode string
	Passphrase string
	DeviceName string
	SkipGPG    bool
	SkipMail   bool
}

func NewSignupEngine(arg *SignupEngineRunArg, g *libkb.GlobalContext) *SignupEngine {
	return &SignupEngine{
		arg:          arg,
		Contextified: libkb.NewContextified(g),
	}
}

func (s *SignupEngine) Name() string {
	return "Signup"
}

func (s *SignupEngine) RequiredUIs() []libkb.UIKind {
	return nil
}

func (s *SignupEngine) GetPrereqs() EnginePrereqs { return EnginePrereqs{} }

func (s *SignupEngine) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{
		&DetKeyEngine{},
		&GPGImportKeyEngine{},
		&DeviceWrap{},
	}
}

func (s *SignupEngine) Init() error {
	return nil
}

func (s *SignupEngine) SetArg(arg *SignupEngineRunArg) {
	s.arg = arg
}

func (s *SignupEngine) CheckRegistered() (err error) {
	s.G().Log.Debug("+ libkb.SignupJoinEngine::CheckRegistered")
	if cr := s.G().Env.GetConfig(); cr == nil {
		err = fmt.Errorf("No configuration file available")
	} else if u := cr.GetUID(); u != nil {
		err = libkb.AlreadyRegisteredError{Uid: *u}
	}
	s.G().Log.Debug("- libkb.SignupJoinEngine::CheckRegistered -> %s", libkb.ErrToOk(err))
	return err
}

func (s *SignupEngine) PostInviteRequest(arg libkb.InviteRequestArg) error {
	return libkb.PostInviteRequest(arg)
}

func (s *SignupEngine) GetMe() *libkb.User {
	return s.me
}

func (s *SignupEngine) Run(ctx *Context) error {
	// make sure we're starting with a clear login state:
	s.G().Logout()

	f := func() error {
		if err := s.genTSPassKey(s.arg.Passphrase); err != nil {
			return err
		}

		if err := s.join(s.arg.Username, s.arg.Email, s.arg.InviteCode, s.arg.SkipMail); err != nil {
			return err
		}

		if err := s.registerDevice(ctx, s.arg.DeviceName); err != nil {
			return err
		}

		if err := s.genDetKeys(ctx); err != nil {
			return fmt.Errorf("detkeys error: %s", err)
		}

		if s.arg.SkipGPG {
			return nil
		}

		if wantsGPG, err := s.checkGPG(ctx); err != nil {
			return err
		} else if wantsGPG {
			if err := s.addGPG(ctx, true); err != nil {
				return fmt.Errorf("addGPG error: %s", err)
			}
		}

		return nil
	}
	return s.G().LoginState().ExternalFunc(f, "SignupEngine - Run")
}

func (s *SignupEngine) genTSPassKey(passphrase string) error {
	salt, err := libkb.RandBytes(triplesec.SaltLen)
	if err != nil {
		return err
	}
	s.pwsalt = salt
	var tsec *triplesec.Cipher
	tsec, s.tspkey, err = libkb.StretchPassphrase(passphrase, salt)
	s.G().LoginState().SetSignupStreamCache(tsec, s.tspkey)
	return err
}

func (s *SignupEngine) join(username, email, inviteCode string, skipMail bool) error {
	joinEngine := NewSignupJoinEngine(s.G())

	arg := SignupJoinEngineRunArg{
		Username:   username,
		Email:      email,
		InviteCode: inviteCode,
		PWHash:     s.tspkey.PWHash(),
		PWSalt:     s.pwsalt,
		SkipMail:   skipMail,
	}
	res := joinEngine.Run(arg)
	if res.Err != nil {
		return res
	}

	s.uid = *res.UID
	user, err := libkb.LoadUser(libkb.LoadUserArg{Self: true, Uid: res.UID, PublicKeyOptional: true})
	if err != nil {
		return err
	}
	s.me = user
	return nil
}

func (s *SignupEngine) registerDevice(ctx *Context, deviceName string) error {
	s.lks = libkb.NewLKSec(s.tspkey.LksClientHalf(), s.G())
	args := &DeviceWrapArgs{
		Me:         s.me,
		DeviceName: deviceName,
		Lks:        s.lks,
		IsEldest:   true,
	}
	eng := NewDeviceWrap(args, s.G())
	if err := RunEngine(eng, ctx); err != nil {
		return err
	}
	s.signingKey = eng.SigningKey()

	return nil
}

func (s *SignupEngine) genDetKeys(ctx *Context) error {
	arg := &DetKeyArgs{
		Tsp:         s.tspkey,
		Me:          s.me,
		SigningKey:  s.signingKey,
		EldestKeyID: s.signingKey.GetKid(),
	}
	eng := NewDetKeyEngine(arg, s.G())
	return RunEngine(eng, ctx)
}

func (s *SignupEngine) checkGPG(ctx *Context) (bool, error) {
	eng := NewGPGImportKeyEngine(nil, s.G())
	return eng.WantsGPG(ctx)
}

func (s *SignupEngine) addGPG(ctx *Context, allowMulti bool) error {
	s.G().Log.Debug("SignupEngine.addGPG.  signingKey: %v\n", s.signingKey)
	arg := GPGImportKeyArg{Signer: s.signingKey, AllowMulti: allowMulti, Me: s.me, Lks: s.lks}
	eng := NewGPGImportKeyEngine(&arg, s.G())
	if err := RunEngine(eng, ctx); err != nil {
		return err
	}

	if s.signingKey == nil {
		s.signingKey = eng.LastKey()
	}
	return nil
}
<|endoftext|>