<commit_before>// Package trace provies an experimental API for initiating
// traces. Veneur's tracing API also provides
// an opentracing compatibility layer. The Veneur tracing API
// is completely independent of the opentracing
// compatibility layer, with the exception of one convenience
// function.
package trace

import (
	"context"
	"io"
	"math/rand"
	"net"
	"reflect"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/stripe/veneur/ssf"

	"github.com/Sirupsen/logrus"
	"github.com/golang/protobuf/proto"
	opentracing "github.com/opentracing/opentracing-go"
)

// Experimental
const NameKey = "name"

// Experimental
const ResourceKey = "resource"

func init() {
	rand.Seed(time.Now().Unix())
}

// (Experimental)
// If this is set to true,
// traces will be generated but not actually sent.
// This should only be set before any traces are generated
var disabled bool = false

var enabledMtx sync.RWMutex

// Make an unexported `key` type that we use as a String such
// that we don't get lint warnings from using it as a key in
// Context. See https://blog.golang.org/context#TOC_3.2.
type key string

const traceKey key = "trace"

// Service is our service name and should be set exactly once,
// at startup
var Service = ""

const localVeneurAddress = "127.0.0.1:8128"

// For an error to be recorded correctly in DataDog, these three tags
// need to be set
const errorMessageTag = "error.msg"
const errorTypeTag = "error.type"
const errorStackTag = "error.stack"

// Trace is a convenient structural representation
// of a TraceSpan. It is intended to map transparently
// to the more general type SSFSample.
type Trace struct {
	// the ID for the root span
	// which is also the ID for the trace itself
	TraceID int64

	// For the root span, this will be equal
	// to the TraceId
	SpanID int64

	// For the root span, this will be <= 0
	ParentID int64

	// The Resource should be the same for all spans in the same trace
	Resource string

	Start time.Time

	End time.Time

	// If non-zero, the trace will be treated
	// as an error
	Status ssf.SSFSample_Status

	Tags map[string]string

	// Unlike the Resource, this should not contain spaces
	// It should be of the format foo.bar.baz
	Name string

	error bool
}

// Set the end timestamp and finalize Span state
func (t *Trace) finish() {
	t.End = time.Now()
}

// (Experimental)
// Enabled sets tracing to enabled.
func Enable() {
	enabledMtx.Lock()
	defer enabledMtx.Unlock()

	disabled = false
}

// (Experimental)
// Disabled sets tracing to disabled.
func Disable() {
	enabledMtx.Lock()
	defer enabledMtx.Unlock()

	disabled = true
}

func Disabled() bool {
	enabledMtx.RLock()
	defer enabledMtx.RUnlock()

	return disabled
}

// Duration is a convenience function for
// the difference between the Start and End timestamps.
// It assumes the span has already ended.
func (t *Trace) Duration() time.Duration {
	if t.End.IsZero() {
		return -1
	}
	return t.End.Sub(t.Start)
}

// SSFSample converts the Trace to an SSFSpan type.
// It sets the duration, so it assumes the span has already ended.
// (It is safe to call on a span that has not ended, but the duration
// field will be invalid)
func (t *Trace) SSFSpan() *ssf.SSFSpan {
	name := t.Name

	span := &ssf.SSFSpan{
		StartTimestamp: t.Start.UnixNano(),
		Error:          t.error,
		TraceId:        t.TraceID,
		Id:             t.SpanID,
		ParentId:       t.ParentID,
		EndTimestamp:   t.End.UnixNano(),
		Tags:           t.Tags,
		Service:        Service,
	}

	span.Tags[NameKey] = name
	span.Tags[ResourceKey] = t.Resource

	return span
}

// ProtoMarshalTo writes the Trace as a protocol buffer
// in text format to the specified writer.
func (t *Trace) ProtoMarshalTo(w io.Writer) error {
	packet, err := proto.Marshal(t.SSFSpan())
	if err != nil {
		return err
	}
	_, err = w.Write(packet)
	return err
}

// Record sends a trace to the (local) veneur instance,
// which will pass it on to the tracing agent running on the
// global veneur instance.
func (t *Trace) Record(name string, tags map[string]string) error {
	if t.Tags == nil {
		t.Tags = map[string]string{}
	}
	t.finish()

	for k, v := range tags {
		t.Tags[k] = v
	}

	if name == "" {
		name = t.Name
	}

	span := t.SSFSpan()
	span.Tags[NameKey] = name

	err := sendSample(span)
	if err != nil {
		logrus.WithError(err).Error("Error submitting sample")
	}
	return err
}

func (t *Trace) Error(err error) {
	t.Status = ssf.SSFSample_CRITICAL
	t.error = true

	errorType := reflect.TypeOf(err).Name()
	if errorType == "" {
		errorType = "error"
	}

	t.Tags[errorMessageTag] = err.Error()
	t.Tags[errorTypeTag] = errorType
	t.Tags[errorStackTag] = err.Error()
}

// Attach attaches the current trace to the context
// and returns a copy of the context with that trace
// stored under the key "trace".
func (t *Trace) Attach(c context.Context) context.Context {
	return context.WithValue(c, traceKey, t)
}

// SpanFromContext is used to create a child span
// when the parent trace is in the context
func SpanFromContext(c context.Context) *Trace {
	parent, ok := c.Value(traceKey).(*Trace)
	if !ok {
		logrus.WithField("type", reflect.TypeOf(c.Value(traceKey))).Error("expected *Trace from context")
	}
	return StartChildSpan(parent)
}

// StartSpanFromContext is used to create a child span
// when the parent trace is in the context
func StartSpanFromContext(ctx context.Context, name string, opts ...opentracing.StartSpanOption) (s *Span, c context.Context) {
	defer func() {
		if r := recover(); r != nil {
			s = nil
			c = ctx
		}
	}()

	if name == "" {
		pc, _, _, ok := runtime.Caller(1)
		details := runtime.FuncForPC(pc)
		if ok && details != nil {
			name = stripPackageName(details.Name())
			opts = append(opts, NameTag(name))
		}
	}

	sp, c := opentracing.StartSpanFromContext(ctx, name, opts...)

	s = sp.(*Span)
	return s, c
}

// SetParent updates the ParentId, TraceId, and Resource of a trace
// based on the parent's values (SpanId, TraceId, Resource).
func (t *Trace) SetParent(parent *Trace) {
	t.ParentID = parent.SpanID
	t.TraceID = parent.TraceID
	t.Resource = parent.Resource
}

// context returns a spanContext representing the trace
// from the point of view of itself .
// (The parentid for the trace will be set as the parentid for the context)
func (t *Trace) context() *spanContext {

	c := &spanContext{}
	c.Init()
	c.baggageItems["traceid"] = strconv.FormatInt(t.TraceID, 10)
	c.baggageItems["parentid"] = strconv.FormatInt(t.ParentID, 10)
	c.baggageItems["spanid"] = strconv.FormatInt(t.SpanID, 10)
	c.baggageItems[ResourceKey] = t.Resource
	return c
}

// contextAsParent returns a spanContext representing the trace
// from the point of view of its direct children.
// (The SpanId for the trace will be set as the ParentId for the context)
func (t *Trace) contextAsParent() *spanContext {

	c := &spanContext{}
	c.Init()
	c.baggageItems["traceid"] = strconv.FormatInt(t.TraceID, 10)
	c.baggageItems["parentid"] = strconv.FormatInt(t.SpanID, 10)
	c.baggageItems[ResourceKey] = t.Resource
	return c
}

// StartTrace is called by to create the root-level span
// for a trace
func StartTrace(resource string) *Trace {
	traceID := proto.Int64(rand.Int63())

	t := &Trace{
		TraceID:  *traceID,
		SpanID:   *traceID,
		ParentID: 0,
		Resource: resource,
		Tags:     map[string]string{},
	}

	t.Start = time.Now()
	return t
}

// StartChildSpan creates a new Span with the specified parent
func StartChildSpan(parent *Trace) *Trace {
	spanID := proto.Int64(rand.Int63())
	span := &Trace{
		SpanID: *spanID,
	}

	span.SetParent(parent)
	span.Start = time.Now()

	return span
}

// sendSample marshals the sample using protobuf and sends it
// over UDP to the local veneur instance
func sendSample(sample *ssf.SSFSpan) error {
	if Disabled() {
		return nil
	}

	serverAddr, err := net.ResolveUDPAddr("udp", localVeneurAddress)
	if err != nil {
		return err
	}

	conn, err := net.DialUDP("udp", nil, serverAddr)
	if err != nil {
		return err
	}

	defer conn.Close()

	data, err := proto.Marshal(sample)
	if err != nil {
		return err
	}

	_, err = conn.Write(data)
	if err != nil {
		return err
	}
	newSample := &ssf.SSFSpan{}
	err = proto.Unmarshal(data, newSample)
	if err != nil {
		panic(err)
	}

	return nil
}

// stripPackageName strips the package name from a function
// name (as formatted by the runtime package)
func stripPackageName(name string) string {
	i := strings.LastIndex(name, "/")
	if i < 0 || i >= len(name)-1 {
		return name
	}

	return name[i+1:]
}
<commit_msg>Remove re-unmarshal in sendSample (#175)<commit_after>// Package trace provies an experimental API for initiating
// traces. Veneur's tracing API also provides
// an opentracing compatibility layer. The Veneur tracing API
// is completely independent of the opentracing
// compatibility layer, with the exception of one convenience
// function.
package trace

import (
	"context"
	"io"
	"math/rand"
	"net"
	"reflect"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/stripe/veneur/ssf"

	"github.com/Sirupsen/logrus"
	"github.com/golang/protobuf/proto"
	opentracing "github.com/opentracing/opentracing-go"
)

// Experimental
const NameKey = "name"

// Experimental
const ResourceKey = "resource"

func init() {
	rand.Seed(time.Now().Unix())
}

// (Experimental)
// If this is set to true,
// traces will be generated but not actually sent.
// This should only be set before any traces are generated
var disabled bool = false

var enabledMtx sync.RWMutex

// Make an unexported `key` type that we use as a String such
// that we don't get lint warnings from using it as a key in
// Context. See https://blog.golang.org/context#TOC_3.2.
type key string

const traceKey key = "trace"

// Service is our service name and should be set exactly once,
// at startup
var Service = ""

const localVeneurAddress = "127.0.0.1:8128"

// For an error to be recorded correctly in DataDog, these three tags
// need to be set
const errorMessageTag = "error.msg"
const errorTypeTag = "error.type"
const errorStackTag = "error.stack"

// Trace is a convenient structural representation
// of a TraceSpan. It is intended to map transparently
// to the more general type SSFSample.
type Trace struct {
	// the ID for the root span
	// which is also the ID for the trace itself
	TraceID int64

	// For the root span, this will be equal
	// to the TraceId
	SpanID int64

	// For the root span, this will be <= 0
	ParentID int64

	// The Resource should be the same for all spans in the same trace
	Resource string

	Start time.Time

	End time.Time

	// If non-zero, the trace will be treated
	// as an error
	Status ssf.SSFSample_Status

	Tags map[string]string

	// Unlike the Resource, this should not contain spaces
	// It should be of the format foo.bar.baz
	Name string

	error bool
}

// Set the end timestamp and finalize Span state
func (t *Trace) finish() {
	t.End = time.Now()
}

// (Experimental)
// Enabled sets tracing to enabled.
func Enable() {
	enabledMtx.Lock()
	defer enabledMtx.Unlock()

	disabled = false
}

// (Experimental)
// Disabled sets tracing to disabled.
func Disable() {
	enabledMtx.Lock()
	defer enabledMtx.Unlock()

	disabled = true
}

func Disabled() bool {
	enabledMtx.RLock()
	defer enabledMtx.RUnlock()

	return disabled
}

// Duration is a convenience function for
// the difference between the Start and End timestamps.
// It assumes the span has already ended.
func (t *Trace) Duration() time.Duration {
	if t.End.IsZero() {
		return -1
	}
	return t.End.Sub(t.Start)
}

// SSFSample converts the Trace to an SSFSpan type.
// It sets the duration, so it assumes the span has already ended.
// (It is safe to call on a span that has not ended, but the duration
// field will be invalid)
func (t *Trace) SSFSpan() *ssf.SSFSpan {
	name := t.Name

	span := &ssf.SSFSpan{
		StartTimestamp: t.Start.UnixNano(),
		Error:          t.error,
		TraceId:        t.TraceID,
		Id:             t.SpanID,
		ParentId:       t.ParentID,
		EndTimestamp:   t.End.UnixNano(),
		Tags:           t.Tags,
		Service:        Service,
	}

	span.Tags[NameKey] = name
	span.Tags[ResourceKey] = t.Resource

	return span
}

// ProtoMarshalTo writes the Trace as a protocol buffer
// in text format to the specified writer.
func (t *Trace) ProtoMarshalTo(w io.Writer) error {
	packet, err := proto.Marshal(t.SSFSpan())
	if err != nil {
		return err
	}
	_, err = w.Write(packet)
	return err
}

// Record sends a trace to the (local) veneur instance,
// which will pass it on to the tracing agent running on the
// global veneur instance.
func (t *Trace) Record(name string, tags map[string]string) error {
	if t.Tags == nil {
		t.Tags = map[string]string{}
	}
	t.finish()

	for k, v := range tags {
		t.Tags[k] = v
	}

	if name == "" {
		name = t.Name
	}

	span := t.SSFSpan()
	span.Tags[NameKey] = name

	err := sendSample(span)
	if err != nil {
		logrus.WithError(err).Error("Error submitting sample")
	}
	return err
}

func (t *Trace) Error(err error) {
	t.Status = ssf.SSFSample_CRITICAL
	t.error = true

	errorType := reflect.TypeOf(err).Name()
	if errorType == "" {
		errorType = "error"
	}

	t.Tags[errorMessageTag] = err.Error()
	t.Tags[errorTypeTag] = errorType
	t.Tags[errorStackTag] = err.Error()
}

// Attach attaches the current trace to the context
// and returns a copy of the context with that trace
// stored under the key "trace".
func (t *Trace) Attach(c context.Context) context.Context {
	return context.WithValue(c, traceKey, t)
}

// SpanFromContext is used to create a child span
// when the parent trace is in the context
func SpanFromContext(c context.Context) *Trace {
	parent, ok := c.Value(traceKey).(*Trace)
	if !ok {
		logrus.WithField("type", reflect.TypeOf(c.Value(traceKey))).Error("expected *Trace from context")
	}
	return StartChildSpan(parent)
}

// StartSpanFromContext is used to create a child span
// when the parent trace is in the context
func StartSpanFromContext(ctx context.Context, name string, opts ...opentracing.StartSpanOption) (s *Span, c context.Context) {
	defer func() {
		if r := recover(); r != nil {
			s = nil
			c = ctx
		}
	}()

	if name == "" {
		pc, _, _, ok := runtime.Caller(1)
		details := runtime.FuncForPC(pc)
		if ok && details != nil {
			name = stripPackageName(details.Name())
			opts = append(opts, NameTag(name))
		}
	}

	sp, c := opentracing.StartSpanFromContext(ctx, name, opts...)

	s = sp.(*Span)
	return s, c
}

// SetParent updates the ParentId, TraceId, and Resource of a trace
// based on the parent's values (SpanId, TraceId, Resource).
func (t *Trace) SetParent(parent *Trace) {
	t.ParentID = parent.SpanID
	t.TraceID = parent.TraceID
	t.Resource = parent.Resource
}

// context returns a spanContext representing the trace
// from the point of view of itself .
// (The parentid for the trace will be set as the parentid for the context)
func (t *Trace) context() *spanContext {

	c := &spanContext{}
	c.Init()
	c.baggageItems["traceid"] = strconv.FormatInt(t.TraceID, 10)
	c.baggageItems["parentid"] = strconv.FormatInt(t.ParentID, 10)
	c.baggageItems["spanid"] = strconv.FormatInt(t.SpanID, 10)
	c.baggageItems[ResourceKey] = t.Resource
	return c
}

// contextAsParent returns a spanContext representing the trace
// from the point of view of its direct children.
// (The SpanId for the trace will be set as the ParentId for the context)
func (t *Trace) contextAsParent() *spanContext {

	c := &spanContext{}
	c.Init()
	c.baggageItems["traceid"] = strconv.FormatInt(t.TraceID, 10)
	c.baggageItems["parentid"] = strconv.FormatInt(t.SpanID, 10)
	c.baggageItems[ResourceKey] = t.Resource
	return c
}

// StartTrace is called by to create the root-level span
// for a trace
func StartTrace(resource string) *Trace {
	traceID := proto.Int64(rand.Int63())

	t := &Trace{
		TraceID:  *traceID,
		SpanID:   *traceID,
		ParentID: 0,
		Resource: resource,
		Tags:     map[string]string{},
	}

	t.Start = time.Now()
	return t
}

// StartChildSpan creates a new Span with the specified parent
func StartChildSpan(parent *Trace) *Trace {
	spanID := proto.Int64(rand.Int63())
	span := &Trace{
		SpanID: *spanID,
	}

	span.SetParent(parent)
	span.Start = time.Now()

	return span
}

// sendSample marshals the sample using protobuf and sends it
// over UDP to the local veneur instance
func sendSample(sample *ssf.SSFSpan) error {
	if Disabled() {
		return nil
	}

	serverAddr, err := net.ResolveUDPAddr("udp", localVeneurAddress)
	if err != nil {
		return err
	}

	conn, err := net.DialUDP("udp", nil, serverAddr)
	if err != nil {
		return err
	}

	defer conn.Close()

	data, err := proto.Marshal(sample)
	if err != nil {
		return err
	}

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

	return nil
}

// stripPackageName strips the package name from a function
// name (as formatted by the runtime package)
func stripPackageName(name string) string {
	i := strings.LastIndex(name, "/")
	if i < 0 || i >= len(name)-1 {
		return name
	}

	return name[i+1:]
}
<|endoftext|>


<commit_before>package trace

import (
	"context"
	"math/rand"
	"net"
	"reflect"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/golang/protobuf/proto"
	"github.com/stripe/veneur/ssf"
)

func init() {
	rand.Seed(time.Now().Unix())
}

// (Experimental)
// If this is set to true,
// traces will be generated but not actually sent.
// This should only be set before any traces are generated
var Disabled bool = false

const traceKey = "trace"

// this should be set exactly once, at startup
var Service = ""

const localVeneurAddress = "127.0.0.1:8128"

type Trace struct {
	// the ID for the root span
	// which is also the ID for the trace itself
	TraceId int64

	// For the root span, this will be equal
	// to the TraceId
	SpanId int64

	// For the root span, this will be <= 0
	ParentId int64

	// The Resource should be the same for all spans in the same trace
	Resource string

	Start time.Time

	// If non-zero, the trace will be treated
	// as an error
	Status ssf.SSFSample_Status

	Tags []*ssf.SSFTag
}

// Record sends a trace to the (local) veneur instance,
// which will pass it on to the tracing agent running on the
// global veneur instance.
func (t *Trace) Record(name string, tags []*ssf.SSFTag) error {
	duration := time.Now().Sub(t.Start).Nanoseconds()

	t.Tags = append(t.Tags, tags...)

	sample := &ssf.SSFSample{
		Metric:    ssf.SSFSample_TRACE,
		Timestamp: t.Start.UnixNano(),
		Status:    t.Status,
		Name:      *proto.String(name),
		Trace: &ssf.SSFTrace{
			TraceId:  t.TraceId,
			Id:       t.SpanId,
			ParentId: t.ParentId,
			Duration: duration,
			Resource: t.Resource,
		},
		SampleRate: *proto.Float32(.10),
		Tags:       t.Tags,
		Service:    Service,
	}

	err := sendSample(sample)
	if err != nil {
		logrus.WithError(err).Error("Error submitting sample")
	}
	return err
}

func (t *Trace) Error(err error) {
	t.Status = ssf.SSFSample_CRITICAL

	errorType := reflect.TypeOf(err).Name()
	if errorType == "" {
		errorType = "error"
	}

	tags := []*ssf.SSFTag{
		{
			Name:  "error.msg",
			Value: err.Error(),
		},
		{
			Name:  "error.type",
			Value: errorType,
		},
		{
			Name:  "error.stack",
			Value: err.Error(),
		},
	}

	t.Tags = append(t.Tags, tags...)
}

// Attach attaches the current trace to the context
// and returns a copy of the context with that trace
// stored under the key "trace".
func (t *Trace) Attach(c context.Context) context.Context {
	return context.WithValue(c, traceKey, t)
}

// SpanFromContext is used to create a child span
// when the parent trace is in the context
func SpanFromContext(c context.Context) *Trace {
	parent, ok := c.Value(traceKey).(*Trace)
	if !ok {
		logrus.WithField("type", reflect.TypeOf(c.Value(traceKey))).Error("expected *Trace from context")
	}

	spanId := proto.Int64(rand.Int63())
	span := &Trace{
		TraceId:  parent.TraceId,
		SpanId:   *spanId,
		ParentId: parent.SpanId,
		Resource: parent.Resource,
		Start:    time.Now(),
	}

	return span
}

// StartTrace is called by to create the root-level span
// for a trace
func StartTrace(resource string) *Trace {
	traceId := proto.Int64(rand.Int63())

	t := &Trace{
		TraceId:  *traceId,
		SpanId:   *traceId,
		ParentId: 0,
		Resource: resource,
	}

	t.Start = time.Now()
	return t
}

// sendSample marshals the sample using protobuf and sends it
// over UDP to the local veneur instance
func sendSample(sample *ssf.SSFSample) error {
	if Disabled {
		return nil
	}

	server_addr, err := net.ResolveUDPAddr("udp", localVeneurAddress)
	if err != nil {
		return err
	}

	conn, err := net.DialUDP("udp", nil, server_addr)
	if err != nil {
		return err
	}

	defer conn.Close()

	data, err := proto.Marshal(sample)
	if err != nil {
		return err
	}

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

	return nil
}
<commit_msg>Start implementing OpenTracing API<commit_after>package trace

import (
	"context"
	"fmt"
	"math/rand"
	"net"
	"reflect"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/golang/protobuf/proto"
	"github.com/opentracing/opentracing-go"
	opentracinglog "github.com/opentracing/opentracing-go/log"
	"github.com/stripe/veneur/ssf"
)

type spanContext struct {
	TraceId      int64
	baggageItems map[string]string
}

// ForeachBaggageItem calls the handler function on each key/val pair in
// the spanContext's baggage items. If the handler function returns false, it
// terminates iteration immediately.
func (c *spanContext) ForeachBaggageItem(handler func(k, v string) bool) {
	for k, v := range c.baggageItems {
		b := handler(k, v)
		if !b {
			return
		}
	}
}

var _ opentracing.SpanContext = &spanContext{}

type Span struct {
	tracer Tracer

	*Trace
	context  *spanContext
	logLines []opentracinglog.Field
}

func (s *Span) Finish() {
	s.FinishWithOptions(opentracing.FinishOptions{
		FinishTime:  time.Now(),
		LogRecords:  nil,
		BulkLogData: nil,
	})

}

// FinishWithOptions finishes the span, but with explicit
// control over timestamps and log data.
// The BulkLogData field is deprecated and ignored.
func (s *Span) FinishWithOptions(opts opentracing.FinishOptions) {
}

func (s *Span) Context() opentracing.SpanContext {
	return s.context
}

func (s *Span) SetOperationName(name string) opentracing.Span {
	s.Trace.Resource = name
	return s
}

// SetTag sets the tags on the underlying span
func (s *Span) SetTag(key string, value interface{}) opentracing.Span {
	tag := ssf.SSFTag{Name: key}
	// TODO mutex
	switch v := value.(type) {
	case fmt.Stringer:
		tag.Value = v.String()
	default:
		// TODO maybe just ban non-strings?
		tag.Value = fmt.Sprintf("%#v", value)
	}
	s.Tags = append(s.Tags, &tag)
	return s
}

// LogFields sets log fields on the underlying span.
// Currently these are ignored, but they can be fun to set anyway!
func (s *Span) LogFields(fields ...opentracinglog.Field) {
	// TODO mutex this
	s.logLines = append(s.logLines, fields...)
}

func (s *Span) LogKV(alternatingKeyValues ...interface{}) {
	// TODO handle error
	fs, _ := opentracinglog.InterleavedKVToFields(alternatingKeyValues...)
	s.LogFields(fs...)
}

func (s *Span) SetBaggageItem(restrictedKey, value string) opentracing.Span {
	s.context.baggageItems[restrictedKey] = value
	return s
}

func (s *Span) BaggageItem(restrictedKey string) string {
	return s.context.baggageItems[restrictedKey]
}

// Tracer returns the tracer that created this Span
func (s *Span) Tracer() opentracing.Tracer {
	return s.tracer
}

// LogEvent is deprecated and unimplemented.
// It is included only to satisfy the opentracing.Span interface.
func (s *Span) LogEvent(event string) {
}

// LogEventWithPayload is deprecated and unimplemented.
// It is included only to satisfy the opentracing.Span interface.
func (s *Span) LogEventWithPayload(event string, payload interface{}) {
}

// Log is deprecated and unimplemented.
// It is included only to satisfy the opentracing.Span interface.
func (s *Span) Log(data opentracing.LogData) {
}

type Tracer struct {
}

func (t Tracer) StartSpan(operationName string, opts ...opentracing.StartSpanOption) opentracing.Span {
	// TODO implement opts

	sso := opentracing.StartSpanOptions{}
	for _, o := range opts {
		o.Apply(&sso)
	}

	trace := StartTrace(operationName)
	return &Span{
		Trace:  trace,
		tracer: t,
	}
}

// Inject injects the provided SpanContext into the carrier for propagation.
// It will return opentracing.ErrUnsupportedFormat if the format is not supported.
// TODO support all the BuiltinFormats
func (t Tracer) Inject(sm opentracing.SpanContext, format interface{}, carrier interface{}) error {
	return opentracing.ErrUnsupportedFormat
}

// Extract returns a SpanContext given the format and the carrier.
// TODO support all the BuiltinFormats
func (t Tracer) Extract(format interface{}, carrier interface{}) (opentracing.SpanContext, error) {
	return nil, opentracing.ErrUnsupportedFormat
}

var _ opentracing.Tracer = &Tracer{}
var _ opentracing.Span = &Span{}

func init() {
	rand.Seed(time.Now().Unix())
}

// (Experimental)
// If this is set to true,
// traces will be generated but not actually sent.
// This should only be set before any traces are generated
var Disabled bool = false

const traceKey = "trace"

// this should be set exactly once, at startup
var Service = ""

const localVeneurAddress = "127.0.0.1:8128"

type Trace struct {
	// the ID for the root span
	// which is also the ID for the trace itself
	TraceId int64

	// For the root span, this will be equal
	// to the TraceId
	SpanId int64

	// For the root span, this will be <= 0
	ParentId int64

	// The Resource should be the same for all spans in the same trace
	Resource string

	Start time.Time

	End time.Time

	// If non-zero, the trace will be treated
	// as an error
	Status ssf.SSFSample_Status

	Tags []*ssf.SSFTag
}

// Set the end timestamp and finalize Span state
func (t *Trace) Finish() {
	t.End = time.Now()
}

// Duration is a convenience function for
// the difference between the Start and End timestamps.
// It assumes the span has already ended.
func (t *Trace) Duration() time.Duration {
	if t.End.IsZero() {
		return -1
	}
	return t.End.Sub(t.Start)
}

// Record sends a trace to the (local) veneur instance,
// which will pass it on to the tracing agent running on the
// global veneur instance.
func (t *Trace) Record(name string, tags []*ssf.SSFTag) error {
	t.Finish()
	duration := t.Duration().Nanoseconds()

	t.Tags = append(t.Tags, tags...)

	sample := &ssf.SSFSample{
		Metric:    ssf.SSFSample_TRACE,
		Timestamp: t.Start.UnixNano(),
		Status:    t.Status,
		Name:      *proto.String(name),
		Trace: &ssf.SSFTrace{
			TraceId:  t.TraceId,
			Id:       t.SpanId,
			ParentId: t.ParentId,
			Duration: duration,
			Resource: t.Resource,
		},
		SampleRate: *proto.Float32(.10),
		Tags:       t.Tags,
		Service:    Service,
	}

	err := sendSample(sample)
	if err != nil {
		logrus.WithError(err).Error("Error submitting sample")
	}
	return err
}

func (t *Trace) Error(err error) {
	t.Status = ssf.SSFSample_CRITICAL

	errorType := reflect.TypeOf(err).Name()
	if errorType == "" {
		errorType = "error"
	}

	tags := []*ssf.SSFTag{
		{
			Name:  "error.msg",
			Value: err.Error(),
		},
		{
			Name:  "error.type",
			Value: errorType,
		},
		{
			Name:  "error.stack",
			Value: err.Error(),
		},
	}

	t.Tags = append(t.Tags, tags...)
}

// Attach attaches the current trace to the context
// and returns a copy of the context with that trace
// stored under the key "trace".
func (t *Trace) Attach(c context.Context) context.Context {
	return context.WithValue(c, traceKey, t)
}

// SpanFromContext is used to create a child span
// when the parent trace is in the context
func SpanFromContext(c context.Context) *Trace {
	parent, ok := c.Value(traceKey).(*Trace)
	if !ok {
		logrus.WithField("type", reflect.TypeOf(c.Value(traceKey))).Error("expected *Trace from context")
	}

	spanId := proto.Int64(rand.Int63())
	span := &Trace{
		TraceId:  parent.TraceId,
		SpanId:   *spanId,
		ParentId: parent.SpanId,
		Resource: parent.Resource,
		Start:    time.Now(),
	}

	return span
}

// StartTrace is called by to create the root-level span
// for a trace
func StartTrace(resource string) *Trace {
	traceId := proto.Int64(rand.Int63())

	t := &Trace{
		TraceId:  *traceId,
		SpanId:   *traceId,
		ParentId: 0,
		Resource: resource,
	}

	t.Start = time.Now()
	return t
}

// sendSample marshals the sample using protobuf and sends it
// over UDP to the local veneur instance
func sendSample(sample *ssf.SSFSample) error {
	if Disabled {
		return nil
	}

	server_addr, err := net.ResolveUDPAddr("udp", localVeneurAddress)
	if err != nil {
		return err
	}

	conn, err := net.DialUDP("udp", nil, server_addr)
	if err != nil {
		return err
	}

	defer conn.Close()

	data, err := proto.Marshal(sample)
	if err != nil {
		return err
	}

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

	return nil
}
<|endoftext|>


<commit_before>package cocoon

import (
	"fmt"

	context "golang.org/x/net/context"

	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"

	"github.com/ellcrys/util"
	"github.com/ncodes/cocoon/core/api/api"
	"github.com/ncodes/cocoon/core/api/api/proto"
	"github.com/ncodes/cocoon/core/client/db"
	"github.com/ncodes/cocoon/core/common"
	"github.com/ncodes/cocoon/core/config"
	"github.com/ncodes/cocoon/core/types"
	logging "github.com/op/go-logging"
)

var log = logging.MustGetLogger("api.client")

// APIAddress is the remote address to the cluster server (TODO: change this to production address)
var APIAddress = util.Env("API_ADDRESS", "127.0.0.1:8005")

func init() {
	log.SetBackend(config.MessageOnlyBackend)
}

// Create a new cocoon
func Create(cocoon *types.Cocoon) error {

	cocoon.ID = util.UUID4()

	err := api.ValidateCocoon(cocoon)
	if err != nil {
		return err
	}

	userSession, err := db.GetUserSessionToken()
	if err != nil {
		return err
	}

	release := types.Release{
		ID:         util.UUID4(),
		CocoonID:   cocoon.ID,
		URL:        cocoon.URL,
		ReleaseTag: cocoon.ReleaseTag,
		Language:   cocoon.Language,
		BuildParam: cocoon.BuildParam,
		Link:       cocoon.Link,
	}

	cocoon.Releases = []string{release.ID}

	conn, err := grpc.Dial(APIAddress, grpc.WithInsecure())
	if err != nil {
		return fmt.Errorf("unable to connect to cluster. please try again")
	}
	defer conn.Close()

	stopSpinner := util.Spinner("Please wait")
	defer stopSpinner()

	client := proto.NewAPIClient(conn)
	md := metadata.Pairs("access_token", userSession.Token)
	ctx := context.Background()
	ctx = metadata.NewContext(ctx, md)
	resp, err := client.CreateCocoon(ctx, &proto.CreateCocoonRequest{
		ID:             cocoon.ID,
		URL:            cocoon.URL,
		Language:       cocoon.Language,
		ReleaseTag:     cocoon.ReleaseTag,
		BuildParam:     cocoon.BuildParam,
		Memory:         cocoon.Memory,
		Link:           cocoon.Link,
		CPUShare:       cocoon.CPUShare,
		Releases:       cocoon.Releases,
		NumSignatories: cocoon.NumSignatories,
		SigThreshold:   cocoon.SigThreshold,
	})

	if err != nil {
		stopSpinner()
		if common.CompareErr(err, types.ErrInvalidOrExpiredToken) == 0 {
			return types.ErrClientNoActiveSession
		}
		return err
	} else if resp.Status != 200 {
		stopSpinner()
		return fmt.Errorf("%s", resp.Body)
	}

	resp, err = client.CreateRelease(context.Background(), &proto.CreateReleaseRequest{
		ID:         release.ID,
		CocoonID:   cocoon.ID,
		URL:        cocoon.URL,
		Link:       cocoon.Link,
		Language:   cocoon.Language,
		ReleaseTag: cocoon.ReleaseTag,
		BuildParam: cocoon.BuildParam,
	})

	if err != nil {
		stopSpinner()
		return err
	} else if resp.Status != 200 {
		stopSpinner()
		return fmt.Errorf("%s", resp.Body)
	}

	stopSpinner()
	log.Info("==> New cocoon created")
	log.Infof("==> Cocoon ID: %s", cocoon.ID)
	log.Infof("==> Release ID: %s", release.ID)

	return nil
}

// Deploy creates and sends a deploy request to the server
func deploy(ctx context.Context, cocoon *types.Cocoon) error {

	conn, err := grpc.Dial(APIAddress, grpc.WithInsecure())
	if err != nil {
		return fmt.Errorf("unable to connect to cluster. please try again")
	}
	defer conn.Close()

	client := proto.NewAPIClient(conn)
	resp, err := client.Deploy(ctx, &proto.DeployRequest{
		ID:         cocoon.ID,
		URL:        cocoon.URL,
		Language:   cocoon.Language,
		ReleaseTag: cocoon.ReleaseTag,
		BuildParam: []byte(cocoon.BuildParam),
		Memory:     cocoon.Memory,
		CPUShare:   cocoon.CPUShare,
	})
	if err != nil {
		return err
	} else if resp.Status != 200 {
		return fmt.Errorf("%s", resp.Body)
	}

	return nil
}

// Start starts a new or stopped cocoon code
func Start(id string) error {

	userSession, err := db.GetUserSessionToken()
	if err != nil {
		return err
	}

	md := metadata.Pairs("access_token", userSession.Token)
	ctx := context.Background()
	ctx = metadata.NewContext(ctx, md)

	conn, err := grpc.Dial(APIAddress, grpc.WithInsecure())
	if err != nil {
		return fmt.Errorf("unable to connect to cluster. please try again")
	}
	defer conn.Close()

	stopSpinner := util.Spinner("Please wait")
	cl := proto.NewAPIClient(conn)
	resp, err := cl.GetCocoon(ctx, &proto.GetCocoonRequest{
		ID:       id,
		Identity: types.NewIdentity(userSession.Email, "").GetHashedEmail(),
	})

	if err != nil {
		stopSpinner()
		if common.CompareErr(err, types.ErrInvalidOrExpiredToken) == 0 {
			return types.ErrClientNoActiveSession
		}
		return err
	} else if resp.Status != 200 {
		stopSpinner()
		return fmt.Errorf("%s", resp.Body)
	}

	var cocoon types.Cocoon
	err = util.FromJSON(resp.Body, &cocoon)

	if err = deploy(ctx, &cocoon); err != nil {
		stopSpinner()
		return err
	}

	stopSpinner()
	log.Info("==> Successfully created a deployment request")
	log.Info("==> ID:", cocoon.ID)

	return nil
}
<commit_msg>Passes link to deployment request<commit_after>package cocoon

import (
	"fmt"

	context "golang.org/x/net/context"

	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"

	"github.com/ellcrys/util"
	"github.com/ncodes/cocoon/core/api/api"
	"github.com/ncodes/cocoon/core/api/api/proto"
	"github.com/ncodes/cocoon/core/client/db"
	"github.com/ncodes/cocoon/core/common"
	"github.com/ncodes/cocoon/core/config"
	"github.com/ncodes/cocoon/core/types"
	logging "github.com/op/go-logging"
)

var log = logging.MustGetLogger("api.client")

// APIAddress is the remote address to the cluster server (TODO: change this to production address)
var APIAddress = util.Env("API_ADDRESS", "127.0.0.1:8005")

func init() {
	log.SetBackend(config.MessageOnlyBackend)
}

// Create a new cocoon
func Create(cocoon *types.Cocoon) error {

	cocoon.ID = util.UUID4()

	err := api.ValidateCocoon(cocoon)
	if err != nil {
		return err
	}

	userSession, err := db.GetUserSessionToken()
	if err != nil {
		return err
	}

	release := types.Release{
		ID:         util.UUID4(),
		CocoonID:   cocoon.ID,
		URL:        cocoon.URL,
		ReleaseTag: cocoon.ReleaseTag,
		Language:   cocoon.Language,
		BuildParam: cocoon.BuildParam,
		Link:       cocoon.Link,
	}

	cocoon.Releases = []string{release.ID}

	conn, err := grpc.Dial(APIAddress, grpc.WithInsecure())
	if err != nil {
		return fmt.Errorf("unable to connect to cluster. please try again")
	}
	defer conn.Close()

	stopSpinner := util.Spinner("Please wait")
	defer stopSpinner()

	client := proto.NewAPIClient(conn)
	md := metadata.Pairs("access_token", userSession.Token)
	ctx := context.Background()
	ctx = metadata.NewContext(ctx, md)
	resp, err := client.CreateCocoon(ctx, &proto.CreateCocoonRequest{
		ID:             cocoon.ID,
		URL:            cocoon.URL,
		Language:       cocoon.Language,
		ReleaseTag:     cocoon.ReleaseTag,
		BuildParam:     cocoon.BuildParam,
		Memory:         cocoon.Memory,
		Link:           cocoon.Link,
		CPUShare:       cocoon.CPUShare,
		Releases:       cocoon.Releases,
		NumSignatories: cocoon.NumSignatories,
		SigThreshold:   cocoon.SigThreshold,
	})

	if err != nil {
		stopSpinner()
		if common.CompareErr(err, types.ErrInvalidOrExpiredToken) == 0 {
			return types.ErrClientNoActiveSession
		}
		return err
	} else if resp.Status != 200 {
		stopSpinner()
		return fmt.Errorf("%s", resp.Body)
	}

	resp, err = client.CreateRelease(context.Background(), &proto.CreateReleaseRequest{
		ID:         release.ID,
		CocoonID:   cocoon.ID,
		URL:        cocoon.URL,
		Link:       cocoon.Link,
		Language:   cocoon.Language,
		ReleaseTag: cocoon.ReleaseTag,
		BuildParam: cocoon.BuildParam,
	})

	if err != nil {
		stopSpinner()
		return err
	} else if resp.Status != 200 {
		stopSpinner()
		return fmt.Errorf("%s", resp.Body)
	}

	stopSpinner()
	log.Info("==> New cocoon created")
	log.Infof("==> Cocoon ID: %s", cocoon.ID)
	log.Infof("==> Release ID: %s", release.ID)

	return nil
}

// Deploy creates and sends a deploy request to the server
func deploy(ctx context.Context, cocoon *types.Cocoon) error {

	conn, err := grpc.Dial(APIAddress, grpc.WithInsecure())
	if err != nil {
		return fmt.Errorf("unable to connect to cluster. please try again")
	}
	defer conn.Close()

	client := proto.NewAPIClient(conn)
	resp, err := client.Deploy(ctx, &proto.DeployRequest{
		ID:         cocoon.ID,
		URL:        cocoon.URL,
		Language:   cocoon.Language,
		ReleaseTag: cocoon.ReleaseTag,
		BuildParam: []byte(cocoon.BuildParam),
		Memory:     cocoon.Memory,
		CPUShare:   cocoon.CPUShare,
		Link:       cocoon.Link,
	})
	if err != nil {
		return err
	} else if resp.Status != 200 {
		return fmt.Errorf("%s", resp.Body)
	}

	return nil
}

// Start starts a new or stopped cocoon code
func Start(id string) error {

	userSession, err := db.GetUserSessionToken()
	if err != nil {
		return err
	}

	md := metadata.Pairs("access_token", userSession.Token)
	ctx := context.Background()
	ctx = metadata.NewContext(ctx, md)

	conn, err := grpc.Dial(APIAddress, grpc.WithInsecure())
	if err != nil {
		return fmt.Errorf("unable to connect to cluster. please try again")
	}
	defer conn.Close()

	stopSpinner := util.Spinner("Please wait")
	cl := proto.NewAPIClient(conn)
	resp, err := cl.GetCocoon(ctx, &proto.GetCocoonRequest{
		ID:       id,
		Identity: types.NewIdentity(userSession.Email, "").GetHashedEmail(),
	})

	if err != nil {
		stopSpinner()
		if common.CompareErr(err, types.ErrInvalidOrExpiredToken) == 0 {
			return types.ErrClientNoActiveSession
		}
		return err
	} else if resp.Status != 200 {
		stopSpinner()
		return fmt.Errorf("%s", resp.Body)
	}

	var cocoon types.Cocoon
	err = util.FromJSON(resp.Body, &cocoon)

	if err = deploy(ctx, &cocoon); err != nil {
		stopSpinner()
		return err
	}

	stopSpinner()
	log.Info("==> Successfully created a deployment request")
	log.Info("==> ID:", cocoon.ID)

	return nil
}
<|endoftext|>


<commit_before>/*
(c) Copyright 2017 Hewlett Packard Enterprise Development LP

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 provisioner

import (
	"fmt"
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/kubernetes"
	storage_v1 "k8s.io/client-go/pkg/apis/storage/v1"
	storage_v1beta1 "k8s.io/client-go/pkg/apis/storage/v1beta1"
	"k8s.io/client-go/tools/cache"
	"nimblestorage/pkg/util"
)

func (p *Provisioner) listAllClasses(options meta_v1.ListOptions) (runtime.Object, error) {
	return p.kubeClient.StorageV1().StorageClasses().List(options)
}
func (p *Provisioner) watchAllClasses(options meta_v1.ListOptions) (watch.Interface, error) {
	return p.kubeClient.StorageV1().StorageClasses().Watch(options)
}

func (p *Provisioner) listBetaAllClasses(options meta_v1.ListOptions) (runtime.Object, error) {
	return p.kubeClient.StorageV1beta1().StorageClasses().List(options)
}
func (p *Provisioner) watchBetaAllClasses(options meta_v1.ListOptions) (watch.Interface, error) {
	return p.kubeClient.StorageV1beta1().StorageClasses().Watch(options)
}

//NewClassReflector provides a controller that watches for PersistentVolumeClasss and takes action on them
func (p *Provisioner) newClassReflector(kubeClient *kubernetes.Clientset) (cache.Store, *cache.Reflector) {
	// In 1.6 and above classes are out of beta
	classListWatch := &cache.ListWatch{
		ListFunc:  p.listAllClasses,
		WatchFunc: p.watchAllClasses,
	}
	// if we're dealing with 1.5, classes are still in beta
	if p.serverVersion.Major == "1" && p.serverVersion.Minor == "5" {
		classListWatch = &cache.ListWatch{
			ListFunc:  p.listBetaAllClasses,
			WatchFunc: p.watchBetaAllClasses,
		}
	}

	classStore := cache.NewStore(cache.DeletionHandlingMetaNamespaceKeyFunc)
	classReflector := cache.NewReflector(classListWatch, &storage_v1.StorageClass{}, classStore, 0)

	return classStore, classReflector
}

func (p *Provisioner) getClass(className string) (*storage_v1.StorageClass, error) {
	classObj, found, err := p.classStore.GetByKey(className)
	if err != nil {
		util.LogError.Printf("error getting class named %s. err=%v", className, err)
		return nil, err
	}
	if !found {
		util.LogError.Printf("unable to find a class named %s", className)
		return nil, fmt.Errorf("unable to find a class named %s", className)
	}
	return getStorageClass(classObj)
}

func getStorageClass(t interface{}) (*storage_v1.StorageClass, error) {
	switch t := t.(type) {
	default:
		return nil, fmt.Errorf("unexpected type %T for %v", t, t)
	case *storage_v1.StorageClass:
		return t, nil
	case *storage_v1beta1.StorageClass:
		return &storage_v1.StorageClass{
			TypeMeta:    t.TypeMeta,
			ObjectMeta:  t.ObjectMeta,
			Provisioner: t.Provisioner,
			Parameters:  t.Parameters,
		}, nil
	case storage_v1beta1.StorageClass:
		return &storage_v1.StorageClass{
			TypeMeta:    t.TypeMeta,
			ObjectMeta:  t.ObjectMeta,
			Provisioner: t.Provisioner,
			Parameters:  t.Parameters,
		}, nil
	case storage_v1.StorageClass:
		return &t, nil
	}
}
<commit_msg>Support 1.5 classes added after we are started<commit_after>/*
(c) Copyright 2017 Hewlett Packard Enterprise Development LP

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 provisioner

import (
	"fmt"
	meta_v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/kubernetes"
	storage_v1 "k8s.io/client-go/pkg/apis/storage/v1"
	storage_v1beta1 "k8s.io/client-go/pkg/apis/storage/v1beta1"
	"k8s.io/client-go/tools/cache"
	"nimblestorage/pkg/util"
)

func (p *Provisioner) listAllClasses(options meta_v1.ListOptions) (runtime.Object, error) {
	return p.kubeClient.StorageV1().StorageClasses().List(options)
}
func (p *Provisioner) watchAllClasses(options meta_v1.ListOptions) (watch.Interface, error) {
	return p.kubeClient.StorageV1().StorageClasses().Watch(options)
}

func (p *Provisioner) listBetaAllClasses(options meta_v1.ListOptions) (runtime.Object, error) {
	return p.kubeClient.StorageV1beta1().StorageClasses().List(options)
}
func (p *Provisioner) watchBetaAllClasses(options meta_v1.ListOptions) (watch.Interface, error) {
	return p.kubeClient.StorageV1beta1().StorageClasses().Watch(options)
}

//NewClassReflector provides a controller that watches for PersistentVolumeClasss and takes action on them
func (p *Provisioner) newClassReflector(kubeClient *kubernetes.Clientset) (cache.Store, *cache.Reflector) {
	classStore := cache.NewStore(cache.DeletionHandlingMetaNamespaceKeyFunc)
	var classReflector *cache.Reflector
	// In 1.6 and above classes are out of beta
	classListWatch := &cache.ListWatch{
		ListFunc:  p.listAllClasses,
		WatchFunc: p.watchAllClasses,
	}
	classReflector = cache.NewReflector(classListWatch, &storage_v1.StorageClass{}, classStore, 0)

	// if we're dealing with 1.5, classes are still in beta
	if p.serverVersion.Major == "1" && p.serverVersion.Minor == "5" {
		classListWatch = &cache.ListWatch{
			ListFunc:  p.listBetaAllClasses,
			WatchFunc: p.watchBetaAllClasses,
		}
		classReflector = cache.NewReflector(classListWatch, &storage_v1beta1.StorageClass{}, classStore, 0)
	}

	return classStore, classReflector
}

func (p *Provisioner) getClass(className string) (*storage_v1.StorageClass, error) {
	classObj, found, err := p.classStore.GetByKey(className)
	if err != nil {
		util.LogError.Printf("error getting class named %s. err=%v", className, err)
		return nil, err
	}
	if !found {
		util.LogError.Printf("unable to find a class named %s", className)
		return nil, fmt.Errorf("unable to find a class named %s", className)
	}
	return getStorageClass(classObj)
}

func getStorageClass(t interface{}) (*storage_v1.StorageClass, error) {
	switch t := t.(type) {
	default:
		return nil, fmt.Errorf("unexpected type %T for %v", t, t)
	case *storage_v1.StorageClass:
		return t, nil
	case *storage_v1beta1.StorageClass:
		return &storage_v1.StorageClass{
			TypeMeta:    t.TypeMeta,
			ObjectMeta:  t.ObjectMeta,
			Provisioner: t.Provisioner,
			Parameters:  t.Parameters,
		}, nil
	case storage_v1beta1.StorageClass:
		return &storage_v1.StorageClass{
			TypeMeta:    t.TypeMeta,
			ObjectMeta:  t.ObjectMeta,
			Provisioner: t.Provisioner,
			Parameters:  t.Parameters,
		}, nil
	case storage_v1.StorageClass:
		return &t, nil
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 Albert Nigmatzianov. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package track

import (
	"strings"

	"github.com/bogem/nehm/util"
)

const (
	clientID = "11a37feb6ccc034d5975f3f803928a32"
)

type Track struct {
	artist string
	title  string

	// Properties needed for JSON unmarshalling
	JArtworkURL string `json:"artwork_url"`
	JAuthor     Author `json:"user"`
	JCreatedAt  string `json:"created_at"`
	JDuration   int    `json:"duration"`
	JID         int    `json:"id"`
	JTitle      string `json:"title"`
	JURL        string `json:"stream_url"`
}

func (t *Track) Artist() string {
	if t.artist == "" {
		t.artist, t.title = t.name()
	}
	return t.artist
}

func (t Track) ArtworkURL() string {
	au := t.JArtworkURL
	if au == "" {
		au = t.JAuthor.AvatarURL
	}
	return strings.Replace(au, "large", "t500x500", 1)
}

func (t Track) Duration() string {
	return util.DurationString(util.ParseDuration(int(t.JDuration)))
}

func (t Track) Filename() string {
	// Replace all filesystem non-friendly runes with the underscore
	toReplace := ":/\\"
	replaceRunes := func(r rune) rune {
		if strings.ContainsRune(toReplace, r) {
			return '_'
		}
		return r
	}

	return strings.Map(replaceRunes, t.Fullname()) + ".mp3"
}

func (t Track) Fullname() string {
	return t.Artist() + " — " + t.Title()
}

func (t Track) ID() int {
	return t.JID
}

// name splits track's title to artist and title if there is one of separators
// in there.
// E.g. if track has title "Michael Jackson - Thriller" then this function will
// return as first string "Michael Jackson" and as second string "Thriller".
func (t Track) name() (string, string) {
	separators := [...]string{" - ", " ~ ", " – "}
	for _, sep := range separators {
		if strings.Contains(t.JTitle, sep) {
			splitted := strings.SplitN(t.JTitle, sep, 2)
			return splitted[0], splitted[1]
		}
	}
	return t.JAuthor.Username, t.JTitle
}

func (t *Track) Title() string {
	if t.title == "" {
		t.artist, t.title = t.name()
	}
	return t.title
}

func (t Track) URL() string {
	return t.JURL + "?client_id=" + clientID
}

func (t Track) Year() string {
	return t.JCreatedAt[0:4]
}
<commit_msg>Make filtering of track filenames stricter on windows<commit_after>// Copyright 2016 Albert Nigmatzianov. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package track

import (
	"runtime"
	"strings"

	"github.com/bogem/nehm/util"
)

const (
	clientID = "11a37feb6ccc034d5975f3f803928a32"
)

type Track struct {
	artist string
	title  string

	// Properties needed for JSON unmarshalling
	JArtworkURL string `json:"artwork_url"`
	JAuthor     Author `json:"user"`
	JCreatedAt  string `json:"created_at"`
	JDuration   int    `json:"duration"`
	JID         int    `json:"id"`
	JTitle      string `json:"title"`
	JURL        string `json:"stream_url"`
}

func (t *Track) Artist() string {
	if t.artist == "" {
		t.artist, t.title = t.name()
	}
	return t.artist
}

func (t Track) ArtworkURL() string {
	au := t.JArtworkURL
	if au == "" {
		au = t.JAuthor.AvatarURL
	}
	return strings.Replace(au, "large", "t500x500", 1)
}

func (t Track) Duration() string {
	return util.DurationString(util.ParseDuration(int(t.JDuration)))
}

func (t Track) Filename() string {
	// Replace all filesystem non-friendly runes with the underscore
	var toReplace string
	if runtime.GOOS == "windows" {
		toReplace = "<>:\"\\/|?*" // https://msdn.microsoft.com/en-us/library/windows/desktop/aa365247(v=vs.85).aspx
	} else {
		toReplace = ":/\\"
	}
	replaceRunes := func(r rune) rune {
		if strings.ContainsRune(toReplace, r) {
			return '_'
		}
		return r
	}

	return strings.Map(replaceRunes, t.Fullname()) + ".mp3"
}

func (t Track) Fullname() string {
	return t.Artist() + " — " + t.Title()
}

func (t Track) ID() int {
	return t.JID
}

// name splits track's title to artist and title if there is one of separators
// in there.
// E.g. if track has title "Michael Jackson - Thriller" then this function will
// return as first string "Michael Jackson" and as second string "Thriller".
func (t Track) name() (string, string) {
	separators := [...]string{" - ", " ~ ", " – "}
	for _, sep := range separators {
		if strings.Contains(t.JTitle, sep) {
			splitted := strings.SplitN(t.JTitle, sep, 2)
			return strings.TrimSpace(splitted[0]), strings.TrimSpace(splitted[1])
		}
	}
	return strings.TrimSpace(t.JAuthor.Username), strings.TrimSpace(t.JTitle)
}

func (t *Track) Title() string {
	if t.title == "" {
		t.artist, t.title = t.name()
	}
	return t.title
}

func (t Track) URL() string {
	return t.JURL + "?client_id=" + clientID
}

func (t Track) Year() string {
	return t.JCreatedAt[0:4]
}
<|endoftext|>


<commit_before>package sentry

import (
	"net/http"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/jianyuan/go-sentry/sentry"
)

func resourceSentryKey() *schema.Resource {
	return &schema.Resource{
		Create: resourceSentryKeyCreate,
		Read:   resourceSentryKeyRead,
		Update: resourceSentryKeyUpdate,
		Delete: resourceSentryKeyDelete,
		Importer: &schema.ResourceImporter{
			State: resourceKeyImport,
		},

		Schema: map[string]*schema.Schema{
			"organization": {
				Type:        schema.TypeString,
				Required:    true,
				Description: "The slug of the organization the key should be created for",
			},
			"project": {
				Type:        schema.TypeString,
				Required:    true,
				Description: "The slug of the project the key should be created for",
			},
			"name": {
				Type:        schema.TypeString,
				Required:    true,
				Description: "The name of the key",
			},
			"public": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"secret": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"project_id": {
				Type:     schema.TypeInt,
				Computed: true,
			},
			"is_active": {
				Type:     schema.TypeBool,
				Computed: true,
			},
			"rate_limit_window": {
				Type:     schema.TypeInt,
				Optional: true,
				Computed: true,
			},
			"rate_limit_count": {
				Type:     schema.TypeInt,
				Optional: true,
				Computed: true,
			},
			"dsn_secret": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"dsn_public": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"dsn_csp": {
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceSentryKeyCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*sentry.Client)

	org := d.Get("organization").(string)
	project := d.Get("project").(string)
	params := &sentry.CreateProjectKeyParams{
		Name: d.Get("name").(string),
	}

	key, _, err := client.ProjectKeys.Create(org, project, params)
	if err != nil {
		return err
	}

	d.SetId(key.ID)
	return resourceSentryKeyRead(d, meta)
}

func resourceSentryKeyRead(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*sentry.Client)

	id := d.Id()
	org := d.Get("organization").(string)
	project := d.Get("project").(string)

	keys, resp, err := client.ProjectKeys.List(org, project)
	if err != nil && resp.StatusCode == http.StatusNotFound {
		d.SetId("")
		return nil
	} else if err != nil {
		return err
	}

	found := false

	for _, key := range keys {
		if key.ID == id {
			d.SetId(key.ID)
			d.Set("name", key.Name)
			d.Set("public", key.Public)
			d.Set("secret", key.Secret)
			d.Set("project_id", key.ProjectID)
			d.Set("is_active", key.IsActive)

			if key.RateLimit != nil {
				d.Set("rate_limit_window", key.RateLimit.Window)
				d.Set("rate_limit_count", key.RateLimit.Count)
			}

			d.Set("dsn_secret", key.DSN.Secret)
			d.Set("dsn_public", key.DSN.Public)
			d.Set("dsn_csp", key.DSN.CSP)

			found = true

			break
		}
	}

	if !found {
		d.SetId("")
	}

	return nil
}

func resourceSentryKeyUpdate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*sentry.Client)

	id := d.Id()
	org := d.Get("organization").(string)
	project := d.Get("project").(string)
	params := &sentry.UpdateProjectKeyParams{
		Name: d.Get("name").(string),
		RateLimit: &sentry.ProjectKeyRateLimit{
			Window: d.Get("rate_limit_window").(int),
			Count:  d.Get("rate_limit_count").(int),
		},
	}

	key, _, err := client.ProjectKeys.Update(org, project, id, params)
	if err != nil {
		return err
	}

	d.SetId(key.ID)
	return resourceSentryKeyRead(d, meta)
}

func resourceSentryKeyDelete(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*sentry.Client)

	id := d.Id()
	org := d.Get("organization").(string)
	project := d.Get("project").(string)

	_, err := client.ProjectKeys.Delete(org, project, id)
	return err
}
<commit_msg>Add rate limit to create operation<commit_after>package sentry

import (
	"net/http"

	"github.com/hashicorp/terraform/helper/schema"
	"github.com/jianyuan/go-sentry/sentry"
)

func resourceSentryKey() *schema.Resource {
	return &schema.Resource{
		Create: resourceSentryKeyCreate,
		Read:   resourceSentryKeyRead,
		Update: resourceSentryKeyUpdate,
		Delete: resourceSentryKeyDelete,
		Importer: &schema.ResourceImporter{
			State: resourceKeyImport,
		},

		Schema: map[string]*schema.Schema{
			"organization": {
				Type:        schema.TypeString,
				Required:    true,
				Description: "The slug of the organization the key should be created for",
			},
			"project": {
				Type:        schema.TypeString,
				Required:    true,
				Description: "The slug of the project the key should be created for",
			},
			"name": {
				Type:        schema.TypeString,
				Required:    true,
				Description: "The name of the key",
			},
			"public": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"secret": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"project_id": {
				Type:     schema.TypeInt,
				Computed: true,
			},
			"is_active": {
				Type:     schema.TypeBool,
				Computed: true,
			},
			"rate_limit_window": {
				Type:     schema.TypeInt,
				Optional: true,
				Computed: true,
			},
			"rate_limit_count": {
				Type:     schema.TypeInt,
				Optional: true,
				Computed: true,
			},
			"dsn_secret": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"dsn_public": {
				Type:     schema.TypeString,
				Computed: true,
			},
			"dsn_csp": {
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceSentryKeyCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*sentry.Client)

	org := d.Get("organization").(string)
	project := d.Get("project").(string)
	params := &sentry.CreateProjectKeyParams{
		Name: d.Get("name").(string),
		RateLimit: &sentry.ProjectKeyRateLimit{
			Window: d.Get("rate_limit_window").(int),
			Count:  d.Get("rate_limit_count").(int),
		},
	}

	key, _, err := client.ProjectKeys.Create(org, project, params)
	if err != nil {
		return err
	}

	d.SetId(key.ID)
	return resourceSentryKeyRead(d, meta)
}

func resourceSentryKeyRead(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*sentry.Client)

	id := d.Id()
	org := d.Get("organization").(string)
	project := d.Get("project").(string)

	keys, resp, err := client.ProjectKeys.List(org, project)
	if err != nil && resp.StatusCode == http.StatusNotFound {
		d.SetId("")
		return nil
	} else if err != nil {
		return err
	}

	found := false

	for _, key := range keys {
		if key.ID == id {
			d.SetId(key.ID)
			d.Set("name", key.Name)
			d.Set("public", key.Public)
			d.Set("secret", key.Secret)
			d.Set("project_id", key.ProjectID)
			d.Set("is_active", key.IsActive)

			if key.RateLimit != nil {
				d.Set("rate_limit_window", key.RateLimit.Window)
				d.Set("rate_limit_count", key.RateLimit.Count)
			}

			d.Set("dsn_secret", key.DSN.Secret)
			d.Set("dsn_public", key.DSN.Public)
			d.Set("dsn_csp", key.DSN.CSP)

			found = true

			break
		}
	}

	if !found {
		d.SetId("")
	}

	return nil
}

func resourceSentryKeyUpdate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*sentry.Client)

	id := d.Id()
	org := d.Get("organization").(string)
	project := d.Get("project").(string)
	params := &sentry.UpdateProjectKeyParams{
		Name: d.Get("name").(string),
		RateLimit: &sentry.ProjectKeyRateLimit{
			Window: d.Get("rate_limit_window").(int),
			Count:  d.Get("rate_limit_count").(int),
		},
	}

	key, _, err := client.ProjectKeys.Update(org, project, id, params)
	if err != nil {
		return err
	}

	d.SetId(key.ID)
	return resourceSentryKeyRead(d, meta)
}

func resourceSentryKeyDelete(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*sentry.Client)

	id := d.Id()
	org := d.Get("organization").(string)
	project := d.Get("project").(string)

	_, err := client.ProjectKeys.Delete(org, project, id)
	return err
}
<|endoftext|>


<commit_before>package workstation

import (
	"fmt"
	"io/ioutil"
	"regexp"
	"runtime"
	"strconv"
	"strings"

	log "github.com/sirupsen/logrus"
	"github.com/xshellinc/iotit/device/config"
	"github.com/xshellinc/tools/dialogs"
	"github.com/xshellinc/tools/lib/help"
	"github.com/xshellinc/tools/lib/sudo"
)

// Linux specific workstation type
type linux struct {
	*workstation
	folder string
}

// Initializes linux workstation with unix type
func newWorkstation(disk string) WorkStation {
	m := new(MountInfo)
	var ms []*MountInfo
	w := &workstation{disk, runtime.GOOS, true, m, ms}
	return &linux{workstation: w, folder: config.MountDir}
}

// Lists available mounts
func (l *linux) ListRemovableDisk() ([]*MountInfo, error) {
	fmt.Println("[+] Listing available disks...")
	regex := regexp.MustCompile(`(sd[a-z])$`)
	regexMmcblk := regexp.MustCompile(`(mmcblk[0-9])$`)
	var (
		devDisks []string
		out      = []*MountInfo{}
	)
	files, _ := ioutil.ReadDir("/dev/")
	for _, f := range files {
		fileName := f.Name()
		if regex.MatchString(fileName) || regexMmcblk.MatchString(fileName) {
			devDisks = append(devDisks, fileName)
		}
	}
	for _, devDisk := range devDisks {
		var p = &MountInfo{}
		diskMap := make(map[string]string)

		r, _ := ioutil.ReadFile("/sys/block/" + devDisk + "/removable")
		removable := strings.Trim(string(r), "\n") == "1"

		sd, _ := ioutil.ReadFile("/sys/block/" + devDisk + "/device/type")
		isSdCard := strings.Trim(string(sd), "\n") == "SD"

		m, _ := ioutil.ReadFile("/sys/block/" + devDisk + "/device/model")
		deviceName := strings.Trim(string(m), "\n")

		// if model is empty, try read from /device/name
		if deviceName == "" {
			n, _ := ioutil.ReadFile("/sys/block/" + devDisk + "/device/name")
			deviceName = strings.Trim(string(n), "\n")
		}

		sizeInSectors, _ := ioutil.ReadFile("/sys/block/" + devDisk + "/size")
		deviceSizeInSectors := strings.Trim(string(sizeInSectors), "\n")
		deviceSizeInSectorsParsed, err := strconv.ParseInt(deviceSizeInSectors, 10, 64)
		if err != nil {
			// unexpected, because there are always integer
			deviceSizeInSectorsParsed = 0
		}

		sectorSize, _ := ioutil.ReadFile("/sys/block/" + devDisk + "/device/erase_size")
		deviceSectorSize := strings.Trim(string(sectorSize), "\n")
		deviceSectorSizeParsed, err := strconv.ParseInt(deviceSectorSize, 10, 64)
		if err != nil {
			// unexpected, because there are always integer
			deviceSectorSizeParsed = 0
		}

		deviceSize := deviceSizeInSectorsParsed * deviceSectorSizeParsed

		diskMap["deviceName"] = deviceName
		diskMap["diskName"] = "/dev/" + devDisk
		diskMap["diskNameRaw"] = "/dev/" + devDisk
		diskMap["deviceSize"] = strconv.FormatInt(deviceSize, 10)

		if removable || isSdCard {
			p.deviceName = diskMap["deviceName"]
			p.diskName = diskMap["diskName"]
			p.diskNameRaw = diskMap["diskNameRaw"]
			p.deviceSize = diskMap["deviceSize"]
			out = append(out, p)
		}
	}

	if !(len(out) > 0) {
		return nil, fmt.Errorf("[-] No mounts found.\n[-] Please insert your SD card and start command again")
	}
	l.workstation.mounts = out
	return out, nil
}

// Unmounts the disk
func (l *linux) Unmount() error {
	if l.workstation.writable != false {
		fmt.Printf("[+] Unmounting disk: %s\n", l.workstation.mount.deviceName)
		stdout, err := help.ExecSudo(sudo.InputMaskedPassword, nil, "umount", l.workstation.mount.diskName)
		if err != nil {
			return fmt.Errorf("Error unmounting disk:%s from %s with error %s, stdout: %s", l.workstation.mount.diskName, l.folder, err.Error(), stdout)
		}
	}
	return nil
}

const diskSelectionTries = 3
const writeAttempts = 3
const cleanTemplate = `n
p
1

w
q
`

// CopyToDisk Notifies user to choose a mount, after that it tries to copy the data
func (l *linux) CopyToDisk(img string) (job *help.BackgroundJob, err error) {
	log.Debug("CopyToDisk")
	_, err = l.ListRemovableDisk()
	if err != nil {
		fmt.Println("[-] SD card is not found, please insert an unlocked SD card")
		return nil, err
	}

	var dev *MountInfo
	if len(l.Disk) == 0 {
		rng := make([]string, len(l.workstation.mounts))
		for i, e := range l.workstation.mounts {
			rng[i] = fmt.Sprintf(dialogs.PrintColored("%s")+" - "+dialogs.PrintColored("%s")+" (%s)", e.deviceName, e.diskName, e.deviceSize)
		}
		num := dialogs.SelectOneDialog("Select disk to format: ", rng)
		dev = l.workstation.mounts[num]
	} else {
		for _, e := range l.workstation.mounts {
			if e.diskName == l.Disk {
				dev = e
				break
			}
		}
		if dev == nil {
			return nil, fmt.Errorf("Disk name not recognised, try to list disks with " + dialogs.PrintColored("disks") + " argument")
		}
	}

	l.workstation.mount = dev
	fmt.Printf("[+] Writing image to %s\n", dev.diskName)
	log.WithField("image", img).WithField("mount", "/media/KERNEL").Debugf("Writing image to %s", dev.diskName)

	if err := l.CleanDisk(dev.diskName); err != nil {
		return nil, err
	}

	job = help.NewBackgroundJob()
	go func() {
		defer job.Close()
		job.Active(true)
		help.ExecCmd("unzip", []string{img, "-d", "/media/KERNEL/"})
		fmt.Println("\r[+] Done writing image to /media/KERNEL")
	}()

	return job, nil
}

// Notifies user to chose a mount, after that it tries to write the data with `diskSelectionTries` number of retries
func (l *linux) WriteToDisk(img string) (job *help.BackgroundJob, err error) {
	for attempt := 0; attempt < diskSelectionTries; attempt++ {
		if attempt > 0 && !dialogs.YesNoDialog("Continue?") {
			break
		}

		_, err = l.ListRemovableDisk()
		if err != nil {
			fmt.Println("[-] SD card is not found, please insert an unlocked SD card")
			continue
		}

		var dev *MountInfo
		if len(l.Disk) == 0 {
			fmt.Println("[+] Available mounts: ")

			rng := make([]string, len(l.workstation.mounts))
			for i, e := range l.workstation.mounts {
				rng[i] = fmt.Sprintf(dialogs.PrintColored("%s")+" - "+dialogs.PrintColored("%s"), e.deviceName, e.diskName)
			}
			num := dialogs.SelectOneDialog("Select mount to format: ", rng)
			dev = l.workstation.mounts[num]
		} else {
			for _, e := range l.workstation.mounts {
				if e.diskName == l.Disk {
					dev = e
					break
				}
			}
			if dev == nil {
				return nil, fmt.Errorf("Disk name not recognised, try to list disks with " + dialogs.PrintColored("disks") + " argument")
			}
		}

		if ok, ferr := help.FileModeMask(dev.diskNameRaw, 0200); !ok || ferr != nil {
			if ferr != nil {
				log.Error(ferr)
				return nil, ferr
			} else {
				fmt.Println("[-] Your card seems locked. Please unlock your SD card")
				err = fmt.Errorf("[-] Your card seems locked.\n[-]  Please unlock your SD card and start command again")
			}
		} else {
			l.workstation.mount = dev
			break
		}
	}

	if err != nil {
		return nil, err
	}

	if len(l.Disk) == 0 && !dialogs.YesNoDialog("Are you sure? ") {
		l.workstation.writable = false
		return nil, nil
	}

	fmt.Printf("[+] Writing %s to %s\n", img, l.workstation.mount.diskName)
	fmt.Println("[+] You may need to enter your user password")

	job = help.NewBackgroundJob()

	go func() {
		defer job.Close()

		args := []string{
			"dd",
			fmt.Sprintf("if=%s", img),
			fmt.Sprintf("of=%s", l.workstation.mount.diskName),
			"bs=4M",
		}

		var err error
		for attempt := 0; attempt < writeAttempts; attempt++ {
			if attempt > 0 && !dialogs.YesNoDialog("Continue?") {
				break
			}
			job.Active(true)

			var out, eut []byte
			if out, eut, err = sudo.Exec(sudo.InputMaskedPassword, job.Progress, args...); err != nil {
				help.LogCmdErrors(string(out), string(eut), err, args...)

				job.Active(false)
				fmt.Println("\r[-] Can't write to disk. Please make sure your password is correct")
			} else {
				job.Active(false)
				fmt.Printf("\r[+] Done writing %s to %s \n", img, l.workstation.mount.diskName)
				break
			}
		}

		if err != nil {
			job.Error(err)
		}
	}()

	l.workstation.writable = true
	return job, nil
}

// Ejects the mounted disk
func (l *linux) Eject() error {
	if l.workstation.writable != false {
		fmt.Printf("[+] Eject your sd card :%s\n", l.workstation.mount.diskName)
		eut, err := help.ExecSudo(sudo.InputMaskedPassword, nil, "eject", l.workstation.mount.diskName)
		if err != nil {
			return fmt.Errorf("eject disk failed: %s\n[-] Cause: %s", l.workstation.mount.diskName, eut)
		}
	}
	return nil
}

// CleanDisk does nothing on linux
func (l *linux) CleanDisk(disk string) error {
	log.WithField("disk", disk).Debug("CleanDisk")
	if disk == "" {
		return fmt.Errorf("No disk to format")
	}
	job := help.NewBackgroundJob()
	go func() {
		defer job.Close()
		job.Active(true)
		help.ExecCmd("umount", []string{disk})

		// fdisk /dev/mmcblk0 < /tmp/fdisk.cmd
		dst := help.GetTempDir() + help.Separator() + "fdisk.cmd"
		help.CreateFile(dst)
		help.WriteFile(dst, cleanTemplate)

		args := []string{
			"fdisk",
			disk,
			"<",
			dst,
		}
		if _, _, err := sudo.Exec(sudo.InputMaskedPassword, job.Progress, args...); err != nil {
			job.Error(err)
		}

		args = []string{
			"mkdosfs",
			"-n",
			"KERNEL",
			disk + "p1",
			"-F",
			"32",
		}
		if _, _, err := sudo.Exec(sudo.InputMaskedPassword, job.Progress, args...); err != nil {
			job.Error(err)
		}
		log.Debug("mkdosfs done")
		help.ExecCmd("mkdir", []string{"/media/KERNEL/"})
		log.Debug("mkdir done")
		help.ExecCmd("mount", []string{disk + "p1", "/media/KERNEL/"})
		log.Debug("mount done")
		l.workstation.mount.deviceName = "/media/KERNEL"
		l.workstation.mount.diskName = disk + "p1"
		l.workstation.writable = true
		job.Active(false)
	}()

	if err := help.WaitJobAndSpin("You need to have "+dialogs.PrintColored("dosfstools")+" package installed. Formatting", job); err != nil {
		return err
	}

	l.workstation.writable = true
	return nil
}

func (l *linux) PrintDisks() {
	l.workstation.printDisks(l)
}
<commit_msg>Fix linux SD cleaning process<commit_after>package workstation

import (
	"fmt"
	"io/ioutil"
	"os"
	"regexp"
	"runtime"
	"strconv"
	"strings"

	log "github.com/sirupsen/logrus"
	"github.com/xshellinc/iotit/device/config"
	"github.com/xshellinc/tools/dialogs"
	"github.com/xshellinc/tools/lib/help"
	"github.com/xshellinc/tools/lib/sudo"
)

// Linux specific workstation type
type linux struct {
	*workstation
	folder string
}

// Initializes linux workstation with unix type
func newWorkstation(disk string) WorkStation {
	m := new(MountInfo)
	var ms []*MountInfo
	w := &workstation{disk, runtime.GOOS, true, m, ms}
	return &linux{workstation: w, folder: config.MountDir}
}

// Lists available mounts
func (l *linux) ListRemovableDisk() ([]*MountInfo, error) {
	fmt.Println("[+] Listing available disks...")
	regex := regexp.MustCompile(`(sd[a-z])$`)
	regexMmcblk := regexp.MustCompile(`(mmcblk[0-9])$`)
	var (
		devDisks []string
		out      = []*MountInfo{}
	)
	files, _ := ioutil.ReadDir("/dev/")
	for _, f := range files {
		fileName := f.Name()
		if regex.MatchString(fileName) || regexMmcblk.MatchString(fileName) {
			devDisks = append(devDisks, fileName)
		}
	}
	for _, devDisk := range devDisks {
		var p = &MountInfo{}
		diskMap := make(map[string]string)

		r, _ := ioutil.ReadFile("/sys/block/" + devDisk + "/removable")
		removable := strings.Trim(string(r), "\n") == "1"

		sd, _ := ioutil.ReadFile("/sys/block/" + devDisk + "/device/type")
		isSdCard := strings.Trim(string(sd), "\n") == "SD"

		m, _ := ioutil.ReadFile("/sys/block/" + devDisk + "/device/model")
		deviceName := strings.Trim(string(m), "\n")

		// if model is empty, try read from /device/name
		if deviceName == "" {
			n, _ := ioutil.ReadFile("/sys/block/" + devDisk + "/device/name")
			deviceName = strings.Trim(string(n), "\n")
		}

		sizeInSectors, _ := ioutil.ReadFile("/sys/block/" + devDisk + "/size")
		deviceSizeInSectors := strings.Trim(string(sizeInSectors), "\n")
		deviceSizeInSectorsParsed, err := strconv.ParseInt(deviceSizeInSectors, 10, 64)
		if err != nil {
			// unexpected, because there are always integer
			deviceSizeInSectorsParsed = 0
		}

		sectorSize, _ := ioutil.ReadFile("/sys/block/" + devDisk + "/device/erase_size")
		deviceSectorSize := strings.Trim(string(sectorSize), "\n")
		deviceSectorSizeParsed, err := strconv.ParseInt(deviceSectorSize, 10, 64)
		if err != nil {
			// unexpected, because there are always integer
			deviceSectorSizeParsed = 0
		}

		deviceSize := deviceSizeInSectorsParsed * deviceSectorSizeParsed

		diskMap["deviceName"] = deviceName
		diskMap["diskName"] = "/dev/" + devDisk
		diskMap["diskNameRaw"] = "/dev/" + devDisk
		diskMap["deviceSize"] = strconv.FormatInt(deviceSize, 10)

		if removable || isSdCard {
			p.deviceName = diskMap["deviceName"]
			p.diskName = diskMap["diskName"]
			p.diskNameRaw = diskMap["diskNameRaw"]
			p.deviceSize = diskMap["deviceSize"]
			out = append(out, p)
		}
	}

	if !(len(out) > 0) {
		return nil, fmt.Errorf("[-] No mounts found.\n[-] Please insert your SD card and start command again")
	}
	l.workstation.mounts = out
	return out, nil
}

// Unmounts the disk
func (l *linux) Unmount() error {
	if l.workstation.writable != false {
		fmt.Printf("[+] Unmounting disk: %s\n", l.workstation.mount.deviceName)
		stdout, err := help.ExecSudo(sudo.InputMaskedPassword, nil, "umount", l.workstation.mount.diskName)
		if err != nil {
			return fmt.Errorf("Error unmounting disk:%s from %s with error %s, stdout: %s", l.workstation.mount.diskName, l.folder, err.Error(), stdout)
		}
	}
	return nil
}

const diskSelectionTries = 3
const writeAttempts = 3
const cleanTemplate = `n
p
1

w
q
`

// CopyToDisk Notifies user to choose a mount, after that it tries to copy the data
func (l *linux) CopyToDisk(img string) (job *help.BackgroundJob, err error) {
	log.Debug("CopyToDisk")
	_, err = l.ListRemovableDisk()
	if err != nil {
		fmt.Println("[-] SD card is not found, please insert an unlocked SD card")
		return nil, err
	}

	var dev *MountInfo
	if len(l.Disk) == 0 {
		rng := make([]string, len(l.workstation.mounts))
		for i, e := range l.workstation.mounts {
			rng[i] = fmt.Sprintf(dialogs.PrintColored("%s")+" - "+dialogs.PrintColored("%s")+" (%s)", e.deviceName, e.diskName, e.deviceSize)
		}
		num := dialogs.SelectOneDialog("Select disk to format: ", rng)
		dev = l.workstation.mounts[num]
	} else {
		for _, e := range l.workstation.mounts {
			if e.diskName == l.Disk {
				dev = e
				break
			}
		}
		if dev == nil {
			return nil, fmt.Errorf("Disk name not recognised, try to list disks with " + dialogs.PrintColored("disks") + " argument")
		}
	}

	l.workstation.mount = dev
	fmt.Printf("[+] Writing image to %s\n", dev.diskName)
	log.WithField("image", img).WithField("mount", "/media/KERNEL").Debugf("Writing image to %s", dev.diskName)

	if err := l.CleanDisk(dev.diskName); err != nil {
		return nil, err
	}

	job = help.NewBackgroundJob()
	go func() {
		defer job.Close()
		job.Active(true)
		help.ExecCmd("unzip", []string{img, "-d", "/media/KERNEL/"})
		fmt.Println("\r[+] Done writing image to /media/KERNEL")
	}()

	return job, nil
}

// Notifies user to chose a mount, after that it tries to write the data with `diskSelectionTries` number of retries
func (l *linux) WriteToDisk(img string) (job *help.BackgroundJob, err error) {
	for attempt := 0; attempt < diskSelectionTries; attempt++ {
		if attempt > 0 && !dialogs.YesNoDialog("Continue?") {
			break
		}

		_, err = l.ListRemovableDisk()
		if err != nil {
			fmt.Println("[-] SD card is not found, please insert an unlocked SD card")
			continue
		}

		var dev *MountInfo
		if len(l.Disk) == 0 {
			fmt.Println("[+] Available mounts: ")

			rng := make([]string, len(l.workstation.mounts))
			for i, e := range l.workstation.mounts {
				rng[i] = fmt.Sprintf(dialogs.PrintColored("%s")+" - "+dialogs.PrintColored("%s"), e.deviceName, e.diskName)
			}
			num := dialogs.SelectOneDialog("Select mount to format: ", rng)
			dev = l.workstation.mounts[num]
		} else {
			for _, e := range l.workstation.mounts {
				if e.diskName == l.Disk {
					dev = e
					break
				}
			}
			if dev == nil {
				return nil, fmt.Errorf("Disk name not recognised, try to list disks with " + dialogs.PrintColored("disks") + " argument")
			}
		}

		if ok, ferr := help.FileModeMask(dev.diskNameRaw, 0200); !ok || ferr != nil {
			if ferr != nil {
				log.Error(ferr)
				return nil, ferr
			} else {
				fmt.Println("[-] Your card seems locked. Please unlock your SD card")
				err = fmt.Errorf("[-] Your card seems locked.\n[-]  Please unlock your SD card and start command again")
			}
		} else {
			l.workstation.mount = dev
			break
		}
	}

	if err != nil {
		return nil, err
	}

	if len(l.Disk) == 0 && !dialogs.YesNoDialog("Are you sure? ") {
		l.workstation.writable = false
		return nil, nil
	}

	fmt.Printf("[+] Writing %s to %s\n", img, l.workstation.mount.diskName)
	fmt.Println("[+] You may need to enter your user password")

	job = help.NewBackgroundJob()

	go func() {
		defer job.Close()

		args := []string{
			"dd",
			fmt.Sprintf("if=%s", img),
			fmt.Sprintf("of=%s", l.workstation.mount.diskName),
			"bs=4M",
		}

		var err error
		for attempt := 0; attempt < writeAttempts; attempt++ {
			if attempt > 0 && !dialogs.YesNoDialog("Continue?") {
				break
			}
			job.Active(true)

			var out, eut []byte
			if out, eut, err = sudo.Exec(sudo.InputMaskedPassword, job.Progress, args...); err != nil {
				help.LogCmdErrors(string(out), string(eut), err, args...)

				job.Active(false)
				fmt.Println("\r[-] Can't write to disk. Please make sure your password is correct")
			} else {
				job.Active(false)
				fmt.Printf("\r[+] Done writing %s to %s \n", img, l.workstation.mount.diskName)
				break
			}
		}

		if err != nil {
			job.Error(err)
		}
	}()

	l.workstation.writable = true
	return job, nil
}

// Ejects the mounted disk
func (l *linux) Eject() error {
	if l.workstation.writable != false {
		fmt.Printf("[+] Eject your sd card :%s\n", l.workstation.mount.diskName)
		eut, err := help.ExecSudo(sudo.InputMaskedPassword, nil, "eject", l.workstation.mount.diskName)
		if err != nil {
			return fmt.Errorf("eject disk failed: %s\n[-] Cause: %s", l.workstation.mount.diskName, eut)
		}
	}
	return nil
}

// CleanDisk does nothing on linux
func (l *linux) CleanDisk(disk string) error {
	log.WithField("disk", disk).Debug("CleanDisk")
	if disk == "" {
		return fmt.Errorf("No disk to format")
	}
	job := help.NewBackgroundJob()
	go func() {
		defer job.Close()
		job.Active(true)
		log.Info("[+] Unmounting")
		cmd := fmt.Sprintf("for n in %s* ; do umount $n ; done", disk)
		if out, err := help.ExecCmd("sh", []string{"-c", cmd}); err == nil {
			log.WithField("cmd", cmd).Debug(strings.TrimSpace(out))
		}

		log.Info("[+] Wiping previous partitions info")
		args := []string{
			"wipefs",
			"-a",
			disk,
		}
		if out, eut, err := sudo.Exec(sudo.InputMaskedPassword, job.Progress, args...); err != nil {
			job.Error(err)
		} else if string(out) != "" || string(eut) != "" {
			log.WithField("out", strings.TrimSpace(string(out))).WithField("eout", strings.TrimSpace(string(eut))).Debug("wipefs output")
		}

		log.Info("[+] Creating new partition table")
		dst := help.GetTempDir() + help.Separator() + "fdisk.in"
		help.CreateFile(dst)
		help.WriteFile(dst, cleanTemplate)

		// fdisk /dev/mmcblk0 < /tmp/fdisk.cmd
		dst = help.GetTempDir() + help.Separator() + "fdisk.cmd"
		help.CreateFile(dst)
		help.WriteFile(dst, fmt.Sprintf("fdisk %s < %s/fdisk.in", disk, help.GetTempDir()))
		os.Chmod(dst, 0755)

		if out, eut, err := sudo.Exec(sudo.InputMaskedPassword, job.Progress, dst); err != nil {
			job.Error(err)
		} else if string(out) != "" || string(eut) != "" {
			log.WithField("out", strings.TrimSpace(string(out))).WithField("eout", strings.TrimSpace(string(eut))).Debug("fdisk output")
		}

		log.Info("[+] Formatting into FAT32 partition")
		partition := disk + "1"
		if strings.Contains(disk, "blk") {
			partition = disk + "p1"
		}
		args = []string{
			"mkdosfs",
			"-n",
			"KERNEL",
			partition,
			"-F",
			"32",
		}
		if out, eut, err := sudo.Exec(sudo.InputMaskedPassword, job.Progress, args...); err != nil {
			job.Error(err)
		} else if string(out) != "" || string(eut) != "" {
			log.WithField("out", strings.TrimSpace(string(out))).WithField("eout", strings.TrimSpace(string(eut))).Debug("mkdosfs output")
		}
		log.Debug("mkdosfs done")
		help.ExecCmd("mkdir", []string{"/media/KERNEL/"})
		log.Debug("mkdir done")
		log.Info("[+] Mounting")
		if out, err := help.ExecCmd("mount", []string{partition, "/media/KERNEL/"}); err != nil {
			log.Debug(strings.TrimSpace(out))
			job.Error(fmt.Errorf("Can't mount %s", partition))
			return
		}
		log.Debug("mount done")
		l.workstation.mount.deviceName = "/media/KERNEL"
		l.workstation.mount.diskName = partition
		l.workstation.writable = true
		job.Active(false)
	}()

	if err := help.WaitJobAndSpin("You need to have "+dialogs.PrintColored("dosfstools")+" package installed. Formatting", job); err != nil {
		return err
	}

	l.workstation.writable = true
	return nil
}

func (l *linux) PrintDisks() {
	l.workstation.printDisks(l)
}
<|endoftext|>


<commit_before>package oort

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"

	"github.com/BurntSushi/toml"
	"github.com/gholt/ring"
	"github.com/pandemicsyn/cmdctrl"
	"github.com/pandemicsyn/syndicate/utils/srvconf"
)

// FExists true if a file or dir exists
func FExists(name string) bool {
	if _, err := os.Stat(name); os.IsNotExist(err) {
		return false
	}
	return true
}

type ccc struct {
	CmdCtrlConfig cmdctrl.ConfigOpts
}

func (o *Server) loadCmdCtrlConfig() error {
	config := ccc{}
	err := o.LoadRingConfig(&config)
	o.CmdCtrlConfig = config.CmdCtrlConfig
	return err
}

func (o *Server) LoadRingConfig(config interface{}) (err error) {
	o.Lock()
	defer o.Unlock()
	log.Println("Using ring version:", o.ring.Version())
	b := bytes.NewReader(o.ring.Conf())
	if b.Len() > 0 {
		_, err = toml.DecodeReader(b, config)
		if err != nil {
			return err
		}
	}
	// Now overlay per node config on top
	n := o.ring.LocalNode()
	if n == nil {
		panic("n is nil")
	}
	b = bytes.NewReader(o.ring.LocalNode().Conf())
	if b.Len() > 0 {
		_, err = toml.DecodeReader(b, config)
		if err != nil {
			return err
		}
	}
	log.Printf("Local Node config is: \n%s", o.ring.LocalNode().Conf())
	log.Printf("Ring config is: \n%s", o.ring.Conf())
	return nil
}

type EnvGetter struct {
	pre string
	sep string
}

func NewEnvGetter(prefix, seperator string) *EnvGetter {
	return &EnvGetter{pre: prefix, sep: seperator}
}

func (e *EnvGetter) Get(varname string) string {
	return os.Getenv(fmt.Sprintf("%s%s%s", e.pre, e.sep, varname))
}

func (o *Server) ObtainConfig() (err error) {
	e := NewEnvGetter(fmt.Sprintf("OORT_%s", strings.ToUpper(o.serviceName)), "_")
	envSkipSRV := e.Get("SKIP_SRV")
	// Check whether we're supposed to skip loading via srv method
	if strings.ToLower(envSkipSRV) != "true" {
		s := &srvconf.SRVLoader{
			SyndicateURL: e.Get("SYNDICATE_OVERRIDE"),
		}
		s.Record, err = genServiceID(o.serviceName, "syndicate", "tcp")
		if err != nil {
			if e.Get("SYNDICATE_OVERRIDE") == "" {
				log.Println(err)
			} else {
				log.Fatalln("No SYNDICATE_OVERRIDE provided and", err)
			}
		}
		if e.Get("SYNDICATE_OVERRIDE") != "" {
			log.Println("Over wrote syndicate url with url from env!", e.Get("SYNDICATE_OVERRIDE"))
		}
		nc, err := s.Load()
		if err != nil {
			return err
		}
		o.ring, err = ring.LoadRing(bytes.NewReader(nc.Ring))
		if err != nil {
			return fmt.Errorf("Error while loading ring for config get via srv lookup: %s", err)
		}
		err = ring.PersistRingOrBuilder(o.ring, nil, fmt.Sprintf("/etc/oort/%s/%d-oort.ring", o.serviceName, o.ring.Version()))
		if err != nil {
			return err
		}
		o.localID = nc.Localid
		o.ring.SetLocalNode(o.localID)
		o.ringFile = fmt.Sprintf("/etc/oort/%s/%d-oort.ring", o.serviceName, o.ring.Version())
		err = o.loadCmdCtrlConfig()
		if err != nil {
			return err
		}
	} else {
		// if you skip the srv load you have to provide all of the info in env vars!
		log.Println("Skipped SRV Config attempting to load from env")
		s, err := strconv.ParseUint(e.Get("LOCALID"), 10, 64)
		if err != nil {
			return fmt.Errorf("Unable to load env specified local id")
		}
		o.localID = s
		o.ringFile = e.Get("RING_FILE")
		o.ring, _, err = ring.RingOrBuilder(o.ringFile)
		if err != nil {
			return fmt.Errorf("Unable to road env specified ring: %s", err)
		}
		o.ring.SetLocalNode(o.localID)
		err = o.loadCmdCtrlConfig()
		if err != nil {
			return err
		}
	}
	return nil
}

//TODO: need to remove the hack to add IAD3 identifier
func genServiceID(service, name, proto string) (string, error) {
	h, _ := os.Hostname()
	d := strings.SplitN(h, ".", 2)
	if len(d) != 2 {
		return "", fmt.Errorf("Unable to determine FQDN, only got short name.")
	}
	return fmt.Sprintf("_%s-%s._%s.%s", service, name, proto, d[1]), nil
}
<commit_msg>Updated to match ring<commit_after>package oort

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"

	"github.com/BurntSushi/toml"
	"github.com/gholt/ring"
	"github.com/pandemicsyn/cmdctrl"
	"github.com/pandemicsyn/syndicate/utils/srvconf"
)

// FExists true if a file or dir exists
func FExists(name string) bool {
	if _, err := os.Stat(name); os.IsNotExist(err) {
		return false
	}
	return true
}

type ccc struct {
	CmdCtrlConfig cmdctrl.ConfigOpts
}

func (o *Server) loadCmdCtrlConfig() error {
	config := ccc{}
	err := o.LoadRingConfig(&config)
	o.CmdCtrlConfig = config.CmdCtrlConfig
	return err
}

func (o *Server) LoadRingConfig(config interface{}) (err error) {
	o.Lock()
	defer o.Unlock()
	log.Println("Using ring version:", o.ring.Version())
	b := bytes.NewReader(o.ring.Config())
	if b.Len() > 0 {
		_, err = toml.DecodeReader(b, config)
		if err != nil {
			return err
		}
	}
	// Now overlay per node config on top
	n := o.ring.LocalNode()
	if n == nil {
		panic("n is nil")
	}
	b = bytes.NewReader(o.ring.LocalNode().Config())
	if b.Len() > 0 {
		_, err = toml.DecodeReader(b, config)
		if err != nil {
			return err
		}
	}
	log.Printf("Local Node config is: \n%s", o.ring.LocalNode().Config())
	log.Printf("Ring config is: \n%s", o.ring.Config())
	return nil
}

type EnvGetter struct {
	pre string
	sep string
}

func NewEnvGetter(prefix, seperator string) *EnvGetter {
	return &EnvGetter{pre: prefix, sep: seperator}
}

func (e *EnvGetter) Get(varname string) string {
	return os.Getenv(fmt.Sprintf("%s%s%s", e.pre, e.sep, varname))
}

func (o *Server) ObtainConfig() (err error) {
	e := NewEnvGetter(fmt.Sprintf("OORT_%s", strings.ToUpper(o.serviceName)), "_")
	envSkipSRV := e.Get("SKIP_SRV")
	// Check whether we're supposed to skip loading via srv method
	if strings.ToLower(envSkipSRV) != "true" {
		s := &srvconf.SRVLoader{
			SyndicateURL: e.Get("SYNDICATE_OVERRIDE"),
		}
		s.Record, err = genServiceID(o.serviceName, "syndicate", "tcp")
		if err != nil {
			if e.Get("SYNDICATE_OVERRIDE") == "" {
				log.Println(err)
			} else {
				log.Fatalln("No SYNDICATE_OVERRIDE provided and", err)
			}
		}
		if e.Get("SYNDICATE_OVERRIDE") != "" {
			log.Println("Over wrote syndicate url with url from env!", e.Get("SYNDICATE_OVERRIDE"))
		}
		nc, err := s.Load()
		if err != nil {
			return err
		}
		o.ring, err = ring.LoadRing(bytes.NewReader(nc.Ring))
		if err != nil {
			return fmt.Errorf("Error while loading ring for config get via srv lookup: %s", err)
		}
		err = ring.PersistRingOrBuilder(o.ring, nil, fmt.Sprintf("/etc/oort/%s/%d-oort.ring", o.serviceName, o.ring.Version()))
		if err != nil {
			return err
		}
		o.localID = nc.Localid
		o.ring.SetLocalNode(o.localID)
		o.ringFile = fmt.Sprintf("/etc/oort/%s/%d-oort.ring", o.serviceName, o.ring.Version())
		err = o.loadCmdCtrlConfig()
		if err != nil {
			return err
		}
	} else {
		// if you skip the srv load you have to provide all of the info in env vars!
		log.Println("Skipped SRV Config attempting to load from env")
		s, err := strconv.ParseUint(e.Get("LOCALID"), 10, 64)
		if err != nil {
			return fmt.Errorf("Unable to load env specified local id")
		}
		o.localID = s
		o.ringFile = e.Get("RING_FILE")
		o.ring, _, err = ring.RingOrBuilder(o.ringFile)
		if err != nil {
			return fmt.Errorf("Unable to road env specified ring: %s", err)
		}
		o.ring.SetLocalNode(o.localID)
		err = o.loadCmdCtrlConfig()
		if err != nil {
			return err
		}
	}
	return nil
}

//TODO: need to remove the hack to add IAD3 identifier
func genServiceID(service, name, proto string) (string, error) {
	h, _ := os.Hostname()
	d := strings.SplitN(h, ".", 2)
	if len(d) != 2 {
		return "", fmt.Errorf("Unable to determine FQDN, only got short name.")
	}
	return fmt.Sprintf("_%s-%s._%s.%s", service, name, proto, d[1]), nil
}
<|endoftext|>