<commit_before>package promclient

import (
	"fmt"
	"reflect"
	"sort"

	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/pkg/labels"
)

func IteratorsForValue(v model.Value) []*SeriesIterator {
	switch valueTyped := v.(type) {
	case *model.Scalar:
		return []*SeriesIterator{NewSeriesIterator(v)}
	case *model.String:
		panic("Not implemented")
	case model.Vector:
		iterators := make([]*SeriesIterator, len(valueTyped))
		for i, sample := range valueTyped {
			iterators[i] = NewSeriesIterator(sample)
		}
		return iterators
	case model.Matrix:
		iterators := make([]*SeriesIterator, len(valueTyped))
		for i, stream := range valueTyped {
			iterators[i] = NewSeriesIterator(stream)
		}
		return iterators
	case nil:
		return nil
	default:
		msg := fmt.Sprintf("Unknown type %v", reflect.TypeOf(v))
		panic(msg)
		return nil
	}
}

// NewSeriesIterator return a series iterator for the given value
// TODO: error return if the type is incorrect?
func NewSeriesIterator(v interface{}) *SeriesIterator {
	return &SeriesIterator{V: v, offset: -1}
}

type SeriesIterator struct {
	V      interface{}
	offset int
}

// Seek advances the iterator forward to the value at or after
// the given timestamp.
func (s *SeriesIterator) Seek(t int64) bool {
	switch valueTyped := s.V.(type) {
	case *model.Sample: // From a vector
		return int64(valueTyped.Timestamp) >= t
	case *model.SampleStream: // from a Matrix
		for i := s.offset; i < len(valueTyped.Values); i++ {
			s.offset = i
			if int64(valueTyped.Values[s.offset].Timestamp) >= t {
				return true
			}
		}
		return false
	default:
		msg := fmt.Sprintf("Unknown data type %v", reflect.TypeOf(s.V))
		panic(msg)
	}
	return false
}

// At returns the current timestamp/value pair.
func (s *SeriesIterator) At() (t int64, v float64) {
	switch valueTyped := s.V.(type) {
	case *model.Sample: // From a vector
		return int64(valueTyped.Timestamp), float64(valueTyped.Value)
	case *model.SampleStream: // from a Matrix
		// We assume the list of values is in order, so we'll iterate backwards
		return int64(valueTyped.Values[s.offset].Timestamp), float64(valueTyped.Values[s.offset].Value)
	default:
		msg := fmt.Sprintf("Unknown data type %v", reflect.TypeOf(s.V))
		panic(msg)
	}

	return 0, 0
}

// Next advances the iterator by one.
func (s *SeriesIterator) Next() bool {
	switch valueTyped := s.V.(type) {
	case *model.Sample: // From a vector
		if s.offset < 0 {
			s.offset = 0
			return true
		}
		return false
	case *model.SampleStream: // from a Matrix
		if s.offset < (len(valueTyped.Values) - 1) {
			s.offset++
			return true
		} else {
			return false
		}
	default:
		msg := fmt.Sprintf("Unknown data type %v", reflect.TypeOf(s.V))
		panic(msg)
	}
	return false
}

// Err returns the current error.
func (s *SeriesIterator) Err() error {
	return nil
}

// Returns the metric of the series that the iterator corresponds to.
func (s *SeriesIterator) Labels() labels.Labels {
	switch valueTyped := s.V.(type) {
	case *model.Scalar:
		panic("Unknown metric() scalar?")
	case *model.Sample: // From a vector
		ret := make(labels.Labels, 0, len(valueTyped.Metric))
		for k, v := range valueTyped.Metric {
			ret = append(ret, labels.Label{string(k), string(v)})
		}
		// TODO: move this into prom
		// there is no reason me (the series iterator) should have to sort these
		// if prom needs them sorted sometimes it should be responsible for doing so
		sort.Sort(ret)
		return ret
	case *model.SampleStream:
		ret := make(labels.Labels, 0, len(valueTyped.Metric))
		for k, v := range valueTyped.Metric {
			ret = append(ret, labels.Label{string(k), string(v)})
		}
		// TODO: move this into prom
		// there is no reason me (the series iterator) should have to sort these
		// if prom needs them sorted sometimes it should be responsible for doing so
		sort.Sort(ret)
		return ret
	default:
		panic("Unknown data type!")
	}
}
<commit_msg>Handle case where prometheus calls Seek() on an empty SeriesIterator<commit_after>package promclient

import (
	"fmt"
	"reflect"
	"sort"

	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/pkg/labels"
)

func IteratorsForValue(v model.Value) []*SeriesIterator {
	switch valueTyped := v.(type) {
	case *model.Scalar:
		return []*SeriesIterator{NewSeriesIterator(v)}
	case *model.String:
		panic("Not implemented")
	case model.Vector:
		iterators := make([]*SeriesIterator, len(valueTyped))
		for i, sample := range valueTyped {
			iterators[i] = NewSeriesIterator(sample)
		}
		return iterators
	case model.Matrix:
		iterators := make([]*SeriesIterator, len(valueTyped))
		for i, stream := range valueTyped {
			iterators[i] = NewSeriesIterator(stream)
		}
		return iterators
	case nil:
		return nil
	default:
		msg := fmt.Sprintf("Unknown type %v", reflect.TypeOf(v))
		panic(msg)
		return nil
	}
}

// NewSeriesIterator return a series iterator for the given value
// TODO: error return if the type is incorrect?
func NewSeriesIterator(v interface{}) *SeriesIterator {
	return &SeriesIterator{V: v, offset: -1}
}

type SeriesIterator struct {
	V      interface{}
	offset int
}

// Seek advances the iterator forward to the value at or after
// the given timestamp.
func (s *SeriesIterator) Seek(t int64) bool {
	switch valueTyped := s.V.(type) {
	case *model.Sample: // From a vector
		return int64(valueTyped.Timestamp) >= t
	case *model.SampleStream: // from a Matrix
		// If someone calls Seek() on an empty SampleStream, just return false
		if len(valueTyped.Values) == 0 {
			return false
		}
		for i := s.offset; i < len(valueTyped.Values); i++ {
			s.offset = i
			if int64(valueTyped.Values[s.offset].Timestamp) >= t {
				return true
			}
		}
		return false
	default:
		msg := fmt.Sprintf("Unknown data type %v", reflect.TypeOf(s.V))
		panic(msg)
	}
	return false
}

// At returns the current timestamp/value pair.
func (s *SeriesIterator) At() (t int64, v float64) {
	switch valueTyped := s.V.(type) {
	case *model.Sample: // From a vector
		return int64(valueTyped.Timestamp), float64(valueTyped.Value)
	case *model.SampleStream: // from a Matrix
		// We assume the list of values is in order, so we'll iterate backwards
		return int64(valueTyped.Values[s.offset].Timestamp), float64(valueTyped.Values[s.offset].Value)
	default:
		msg := fmt.Sprintf("Unknown data type %v", reflect.TypeOf(s.V))
		panic(msg)
	}

	return 0, 0
}

// Next advances the iterator by one.
func (s *SeriesIterator) Next() bool {
	switch valueTyped := s.V.(type) {
	case *model.Sample: // From a vector
		if s.offset < 0 {
			s.offset = 0
			return true
		}
		return false
	case *model.SampleStream: // from a Matrix
		if s.offset < (len(valueTyped.Values) - 1) {
			s.offset++
			return true
		} else {
			return false
		}
	default:
		msg := fmt.Sprintf("Unknown data type %v", reflect.TypeOf(s.V))
		panic(msg)
	}
	return false
}

// Err returns the current error.
func (s *SeriesIterator) Err() error {
	return nil
}

// Returns the metric of the series that the iterator corresponds to.
func (s *SeriesIterator) Labels() labels.Labels {
	switch valueTyped := s.V.(type) {
	case *model.Scalar:
		panic("Unknown metric() scalar?")
	case *model.Sample: // From a vector
		ret := make(labels.Labels, 0, len(valueTyped.Metric))
		for k, v := range valueTyped.Metric {
			ret = append(ret, labels.Label{string(k), string(v)})
		}
		// TODO: move this into prom
		// there is no reason me (the series iterator) should have to sort these
		// if prom needs them sorted sometimes it should be responsible for doing so
		sort.Sort(ret)
		return ret
	case *model.SampleStream:
		ret := make(labels.Labels, 0, len(valueTyped.Metric))
		for k, v := range valueTyped.Metric {
			ret = append(ret, labels.Label{string(k), string(v)})
		}
		// TODO: move this into prom
		// there is no reason me (the series iterator) should have to sort these
		// if prom needs them sorted sometimes it should be responsible for doing so
		sort.Sort(ret)
		return ret
	default:
		panic("Unknown data type!")
	}
}
<|endoftext|>


<commit_before>package main

import "github.com/golang/protobuf/protoc-gen-go/descriptor"

func (fg *FileGenerator) GenerateOneofDefinition(prefix string, inMessage *descriptor.DescriptorProto, oneofIndex int) error {
	inOneof := inMessage.GetOneofDecl()[oneofIndex]

	// TODO: Prefix with message name to avoid collisions.
	oneofType := oneofType(inOneof)

	fg.P("")
	fg.P("")
	fg.P("type %s", oneofType)

	fg.In()

	leading := "="
	{
		oneofVariantName := oneofUnspecifiedValue(inOneof)
		fg.P("%s %s", leading, oneofVariantName)
		leading = "|"
	}
	for _, inField := range inMessage.GetField() {
		if inField.OneofIndex != nil && inField.GetOneofIndex() == int32(oneofIndex) {

			oneofVariantName := elmTypeName(inField.GetName())
			oneofArgumentType := fieldElmType(inField)
			fg.P("%s %s %s", leading, oneofVariantName, oneofArgumentType)

			leading = "|"
		}
	}
	fg.Out()

	return nil
}

func (fg *FileGenerator) GenerateOneofDecoder(prefix string, inMessage *descriptor.DescriptorProto, oneofIndex int) error {
	inOneof := inMessage.GetOneofDecl()[oneofIndex]

	// TODO: Prefix with message name to avoid collisions.
	oneofType := oneofType(inOneof)
	decoderName := oneofDecoderName(inOneof)

	fg.P("")
	fg.P("")
	fg.P("%s : JD.Decoder %s", decoderName, oneofType)
	fg.P("%s =", decoderName)

	fg.In()

	fg.P("JD.oneOf")
	fg.In()

	leading := "["
	for _, inField := range inMessage.GetField() {
		if inField.OneofIndex != nil && inField.GetOneofIndex() == int32(oneofIndex) {
			oneofVariantName := elmTypeName(inField.GetName())
			decoderName := fieldDecoderName(inField)
			fg.P("%s JD.map %s (%q := %s)", leading, oneofVariantName, inField.GetJsonName(), decoderName)
			leading = ","
		}
	}
	fg.P("%s JD.succeed %s", leading, oneofUnspecifiedValue(inOneof))
	fg.P("]")
	fg.Out()
	fg.Out()

	return nil
}

func (fg *FileGenerator) GenerateOneofEncoder(prefix string, inMessage *descriptor.DescriptorProto, oneofIndex int) error {
	inOneof := inMessage.GetOneofDecl()[oneofIndex]

	// TODO: Prefix with message name to avoid collisions.
	oneofType := oneofType(inOneof)
	encoderName := oneofEncoderName(inOneof)
	argName := "v"

	fg.P("")
	fg.P("")
	fg.P("%s : %s -> Maybe (String, JE.Value)", encoderName, oneofType)
	fg.P("%s %s =", encoderName, argName)

	fg.In()
	fg.P("case %s of", argName)
	fg.In()

	valueName := "x"
	{
		oneofVariantName := oneofUnspecifiedValue(inOneof)
		fg.P("%s -> Nothing", oneofVariantName)
	}
	for _, inField := range inMessage.GetField() {
		if inField.OneofIndex != nil && inField.GetOneofIndex() == int32(oneofIndex) {
			oneofVariantName := elmTypeName(inField.GetName())
			e := fieldEncoderName(inField)
			fg.P("%s %s -> Just (%q, %s %s)", oneofVariantName, valueName, inField.GetJsonName(), e, valueName)
		}
	}
	fg.Out()
	fg.Out()

	return nil
}

func oneofDecoderName(inOneof *descriptor.OneofDescriptorProto) string {
	typeName := elmTypeName(inOneof.GetName())
	return decoderName(typeName)
}

func oneofEncoderName(inOneof *descriptor.OneofDescriptorProto) string {
	typeName := elmTypeName(inOneof.GetName())
	return encoderName(typeName)
}

func oneofType(inOneof *descriptor.OneofDescriptorProto) string {
	return elmTypeName(inOneof.GetName())
}

func oneofUnspecifiedValue(inOneof *descriptor.OneofDescriptorProto) string {
	return elmTypeName(inOneof.GetName() + "_unspecified")
}
<commit_msg>Add TODO for oneof evaluation order<commit_after>package main

import "github.com/golang/protobuf/protoc-gen-go/descriptor"

func (fg *FileGenerator) GenerateOneofDefinition(prefix string, inMessage *descriptor.DescriptorProto, oneofIndex int) error {
	inOneof := inMessage.GetOneofDecl()[oneofIndex]

	// TODO: Prefix with message name to avoid collisions.
	oneofType := oneofType(inOneof)

	fg.P("")
	fg.P("")
	fg.P("type %s", oneofType)

	fg.In()

	leading := "="
	{
		oneofVariantName := oneofUnspecifiedValue(inOneof)
		fg.P("%s %s", leading, oneofVariantName)
		leading = "|"
	}
	for _, inField := range inMessage.GetField() {
		if inField.OneofIndex != nil && inField.GetOneofIndex() == int32(oneofIndex) {

			oneofVariantName := elmTypeName(inField.GetName())
			oneofArgumentType := fieldElmType(inField)
			fg.P("%s %s %s", leading, oneofVariantName, oneofArgumentType)

			leading = "|"
		}
	}
	fg.Out()

	return nil
}

func (fg *FileGenerator) GenerateOneofDecoder(prefix string, inMessage *descriptor.DescriptorProto, oneofIndex int) error {
	inOneof := inMessage.GetOneofDecl()[oneofIndex]

	// TODO: Prefix with message name to avoid collisions.
	oneofType := oneofType(inOneof)
	decoderName := oneofDecoderName(inOneof)

	fg.P("")
	fg.P("")
	fg.P("%s : JD.Decoder %s", decoderName, oneofType)
	fg.P("%s =", decoderName)

	fg.In()

	fg.P("JD.oneOf")
	fg.In()

	leading := "["
	for _, inField := range inMessage.GetField() {
		if inField.OneofIndex != nil && inField.GetOneofIndex() == int32(oneofIndex) {
			oneofVariantName := elmTypeName(inField.GetName())
			decoderName := fieldDecoderName(inField)
			fg.P("%s JD.map %s (%q := %s)", leading, oneofVariantName, inField.GetJsonName(), decoderName)
			leading = ","
		}
	}
	fg.P("%s JD.succeed %s", leading, oneofUnspecifiedValue(inOneof))
	fg.P("]")
	fg.Out()
	fg.Out()

	return nil
}

func (fg *FileGenerator) GenerateOneofEncoder(prefix string, inMessage *descriptor.DescriptorProto, oneofIndex int) error {
	inOneof := inMessage.GetOneofDecl()[oneofIndex]

	// TODO: Prefix with message name to avoid collisions.
	oneofType := oneofType(inOneof)
	encoderName := oneofEncoderName(inOneof)
	argName := "v"

	fg.P("")
	fg.P("")
	fg.P("%s : %s -> Maybe (String, JE.Value)", encoderName, oneofType)
	fg.P("%s %s =", encoderName, argName)

	fg.In()
	fg.P("case %s of", argName)
	fg.In()

	valueName := "x"
	{
		oneofVariantName := oneofUnspecifiedValue(inOneof)
		fg.P("%s -> Nothing", oneofVariantName)
	}
	// TODO: Evaluate them in reverse order, as per
	// https://developers.google.com/protocol-buffers/docs/proto3#oneof
	for _, inField := range inMessage.GetField() {
		if inField.OneofIndex != nil && inField.GetOneofIndex() == int32(oneofIndex) {
			oneofVariantName := elmTypeName(inField.GetName())
			e := fieldEncoderName(inField)
			fg.P("%s %s -> Just (%q, %s %s)", oneofVariantName, valueName, inField.GetJsonName(), e, valueName)
		}
	}
	fg.Out()
	fg.Out()

	return nil
}

func oneofDecoderName(inOneof *descriptor.OneofDescriptorProto) string {
	typeName := elmTypeName(inOneof.GetName())
	return decoderName(typeName)
}

func oneofEncoderName(inOneof *descriptor.OneofDescriptorProto) string {
	typeName := elmTypeName(inOneof.GetName())
	return encoderName(typeName)
}

func oneofType(inOneof *descriptor.OneofDescriptorProto) string {
	return elmTypeName(inOneof.GetName())
}

func oneofUnspecifiedValue(inOneof *descriptor.OneofDescriptorProto) string {
	return elmTypeName(inOneof.GetName() + "_unspecified")
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"strings"

	libkb "github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	jsonw "github.com/keybase/go-jsonw"
)

// Prove is an engine used for proving ownership of remote accounts,
// like Twitter, GitHub, etc.
type Prove struct {
	arg        *keybase1.StartProofArg
	me         *libkb.User
	st         libkb.ServiceType
	supersede  bool
	proof      *jsonw.Wrapper
	sig        string
	sigID      keybase1.SigID
	postRes    *libkb.PostProofRes
	signingKey libkb.GenericKey
	sigInner   []byte

	remoteNameNormalized string

	libkb.Contextified
}

// NewProve makes a new Prove Engine given an RPC-friendly ProveArg.
func NewProve(g *libkb.GlobalContext, arg *keybase1.StartProofArg) *Prove {
	if arg.SigVersion == nil || libkb.SigVersion(*arg.SigVersion) == libkb.KeybaseNullSigVersion {
		tmp := keybase1.SigVersion(libkb.GetDefaultSigVersion(g))
		arg.SigVersion = &tmp
	}
	return &Prove{
		arg:          arg,
		Contextified: libkb.NewContextified(g),
	}
}

// Name provides the name of this engine for the engine interface contract
func (p *Prove) Name() string {
	return "Prove"
}

// GetPrereqs returns the engine prereqs.
func (p *Prove) Prereqs() Prereqs {
	return Prereqs{Device: true}
}

// RequiredUIs returns the required UIs.
func (p *Prove) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{
		libkb.LogUIKind,
		libkb.ProveUIKind,
		libkb.SecretUIKind,
	}
}

// SubConsumers returns the other UI consumers for this engine.
func (p *Prove) SubConsumers() []libkb.UIConsumer {
	return nil
}

func (p *Prove) loadMe(m libkb.MetaContext) (err error) {
	p.me, err = libkb.LoadMe(libkb.NewLoadUserArgWithMetaContext(m).WithForceReload())
	return err
}

func (p *Prove) checkExists1(m libkb.MetaContext) (err error) {
	proofs := p.me.IDTable().GetActiveProofsFor(p.st)
	if len(proofs) != 0 && !p.arg.Force && p.st.LastWriterWins() {
		lst := proofs[len(proofs)-1]
		var redo bool
		redo, err = m.UIs().ProveUI.PromptOverwrite(m.Ctx(), keybase1.PromptOverwriteArg{
			Account: lst.ToDisplayString(),
			Typ:     keybase1.PromptOverwriteType_SOCIAL,
		})
		if err != nil {
			return err
		}
		if !redo {
			return libkb.NotConfirmedError{}
		}
		p.supersede = true
	}
	return nil
}

func (p *Prove) promptRemoteName(m libkb.MetaContext) error {
	// If the name is already supplied, there's no need to prompt.
	if len(p.arg.Username) > 0 {
		remoteNameNormalized, err := p.st.NormalizeRemoteName(m, p.arg.Username)
		if err == nil {
			p.remoteNameNormalized = remoteNameNormalized
		}
		return err
	}

	// Prompt for the name, retrying if it's invalid.
	var normalizationError error
	for {
		un, err := m.UIs().ProveUI.PromptUsername(m.Ctx(), keybase1.PromptUsernameArg{
			Prompt:    p.st.GetPrompt(),
			PrevError: libkb.ExportErrorAsStatus(m.G(), normalizationError),
		})
		if err != nil {
			// Errors here are conditions like EOF. Return them rather than retrying.
			return err
		}
		var remoteNameNormalized string
		remoteNameNormalized, normalizationError = p.st.NormalizeRemoteName(m, un)
		if normalizationError == nil {
			p.remoteNameNormalized = remoteNameNormalized
			return nil
		}
	}
}

func (p *Prove) checkExists2(m libkb.MetaContext) (err error) {
	defer m.CTrace("Prove#CheckExists2", func() error { return err })()
	if !p.st.LastWriterWins() {
		var found libkb.RemoteProofChainLink
		for _, proof := range p.me.IDTable().GetActiveProofsFor(p.st) {
			_, name := proof.ToKeyValuePair()
			if libkb.Cicmp(name, p.remoteNameNormalized) {
				found = proof
				break
			}
		}
		if found != nil {
			var redo bool
			redo, err = m.UIs().ProveUI.PromptOverwrite(m.Ctx(), keybase1.PromptOverwriteArg{
				Account: found.ToDisplayString(),
				Typ:     keybase1.PromptOverwriteType_SITE,
			})
			if err != nil {
				return err
			}
			if !redo {
				err = libkb.NotConfirmedError{}
				return err
			}
			p.supersede = true
		}
	}
	return nil
}

func (p *Prove) doPrechecks(m libkb.MetaContext) (err error) {
	var w *libkb.Markup
	w, err = p.st.PreProofCheck(m, p.remoteNameNormalized)
	if w != nil {
		if uierr := m.UIs().ProveUI.OutputPrechecks(m.Ctx(), keybase1.OutputPrechecksArg{Text: w.Export()}); uierr != nil {
			m.CWarningf("prove ui OutputPrechecks call error: %s", uierr)
		}
	}
	return err
}

func (p *Prove) doWarnings(m libkb.MetaContext) (err error) {
	if mu := p.st.PreProofWarning(p.remoteNameNormalized); mu != nil {
		var ok bool
		arg := keybase1.PreProofWarningArg{Text: mu.Export()}
		if ok, err = m.UIs().ProveUI.PreProofWarning(m.Ctx(), arg); err == nil && !ok {
			err = libkb.NotConfirmedError{}
			return err
		}
	}
	return nil
}

func (p *Prove) generateProof(m libkb.MetaContext) (err error) {
	ska := libkb.SecretKeyArg{
		Me:      p.me,
		KeyType: libkb.DeviceSigningKeyType,
	}

	p.signingKey, err = m.G().Keyrings.GetSecretKeyWithPrompt(m, m.SecretKeyPromptArg(ska, "tracking signature"))
	if err != nil {
		return err
	}

	sigVersion := libkb.SigVersion(*p.arg.SigVersion)

	if p.proof, err = p.me.ServiceProof(m, p.signingKey, p.st, p.remoteNameNormalized, sigVersion); err != nil {
		return err
	}

	if p.sigInner, err = p.proof.Marshal(); err != nil {
		return err
	}

	p.sig, p.sigID, _, err = libkb.MakeSig(
		m,
		p.signingKey,
		libkb.LinkTypeWebServiceBinding,
		p.sigInner,
		libkb.SigHasRevokes(false),
		keybase1.SeqType_PUBLIC,
		libkb.SigIgnoreIfUnsupported(false),
		p.me,
		sigVersion,
	)

	return err
}

func (p *Prove) postProofToServer(m libkb.MetaContext) (err error) {
	arg := libkb.PostProofArg{
		Sig:               p.sig,
		ProofType:         p.st.GetProofType(),
		RemoteServiceType: p.st.GetTypeName(),
		ID:                p.sigID,
		Supersede:         p.supersede,
		RemoteUsername:    p.remoteNameNormalized,
		RemoteKey:         p.st.GetAPIArgKey(),
		SigningKey:        p.signingKey,
	}
	if libkb.SigVersion(*p.arg.SigVersion) == libkb.KeybaseSignatureV2 {
		arg.SigInner = p.sigInner
	}
	p.postRes, err = libkb.PostProof(m, arg)
	return err
}

func (p *Prove) instructAction(m libkb.MetaContext) (err error) {
	mkp := p.st.PostInstructions(p.remoteNameNormalized)
	var txt string
	if txt, err = p.st.FormatProofText(m, p.postRes, p.arg.Username, p.sigID); err != nil {
		return err
	}
	err = m.UIs().ProveUI.OutputInstructions(m.Ctx(), keybase1.OutputInstructionsArg{
		Instructions: mkp.Export(),
		// If we don't trim newlines here, we'll run into an issue where e.g.
		// Facebook links get corrupted on iOS. See:
		// - https://keybase.atlassian.net/browse/DESKTOP-3335
		// - https://keybase.atlassian.net/browse/CORE-4941
		// All of our proof verifying code (PVL) should already be flexible
		// with surrounding whitespace, because users are pasting proofs by
		// hand anyway.
		Proof: strings.TrimSpace(txt),
	})
	if err != nil {
		return err
	}

	return p.checkAutoPost(m, txt)
}

func (p *Prove) checkAutoPost(m libkb.MetaContext, txt string) error {
	if !p.arg.Auto {
		return nil
	}
	if libkb.RemoteServiceTypes[p.arg.Service] != keybase1.ProofType_ROOTER {
		return nil
	}
	m.CDebugf("making automatic post of proof to rooter")
	apiArg := libkb.APIArg{
		Endpoint:    "rooter",
		SessionType: libkb.APISessionTypeREQUIRED,
		Args: libkb.HTTPArgs{
			"post":     libkb.S{Val: txt},
			"username": libkb.S{Val: p.arg.Username},
		},
		NetContext: m.Ctx(),
	}
	if _, err := m.G().API.Post(apiArg); err != nil {
		m.CDebugf("error posting to rooter: %s", err)
		return err
	}
	return nil
}

func (p *Prove) promptPostedLoop(m libkb.MetaContext) (err error) {
	found := false
	for i := 0; ; i++ {
		var retry bool
		var status keybase1.ProofStatus
		var warn *libkb.Markup
		retry, err = m.UIs().ProveUI.OkToCheck(m.Ctx(), keybase1.OkToCheckArg{
			Name:    p.st.DisplayName(p.remoteNameNormalized),
			Attempt: i,
		})
		if !retry || err != nil {
			break
		}
		found, status, _, err = libkb.CheckPosted(m, p.postRes.ID)
		if found || err != nil {
			break
		}
		warn, err = p.st.RecheckProofPosting(i, status, p.remoteNameNormalized)
		if warn != nil {
			uierr := m.UIs().ProveUI.DisplayRecheckWarning(m.Ctx(), keybase1.DisplayRecheckWarningArg{
				Text: warn.Export(),
			})
			if uierr != nil {
				m.CWarningf("prove ui DisplayRecheckWarning call error: %s", uierr)
			}
		}
		if err != nil {
			break
		}
	}
	if !found && err == nil {
		err = libkb.ProofNotYetAvailableError{}
	}

	return err
}

func (p *Prove) checkProofText(m libkb.MetaContext) error {
	m.CDebugf("p.postRes.Text: %q", p.postRes.Text)
	m.CDebugf("p.sigID: %q", p.sigID)
	return p.st.CheckProofText(p.postRes.Text, p.sigID, p.sig)
}

func (p *Prove) getServiceType(m libkb.MetaContext) (err error) {
	if p.st = m.G().GetProofServices().GetServiceType(p.arg.Service); p.st == nil {
		err = libkb.BadServiceError{Service: p.arg.Service}
	}
	return err
}

// SigID returns the signature id of the proof posted to the
// server.
func (p *Prove) SigID() keybase1.SigID {
	return p.sigID
}

// Run runs the Prove engine, performing all steps of the proof process.
func (p *Prove) Run(m libkb.MetaContext) (err error) {
	defer m.CTrace("ProofEngine.Run", func() error { return err })()

	stage := func(s string) {
		m.CDebugf("| ProofEngine.Run() %s", s)
	}

	stage("GetServiceType")
	if err = p.getServiceType(m); err != nil {
		return err
	}
	stage("LoadMe")
	if err = p.loadMe(m); err != nil {
		return err
	}
	stage("CheckExists1")
	if err = p.checkExists1(m); err != nil {
		return err
	}
	stage("PromptRemoteName")
	if err = p.promptRemoteName(m); err != nil {
		return err
	}
	stage("CheckExists2")
	if err = p.checkExists2(m); err != nil {
		return err
	}
	stage("DoPrechecks")
	if err = p.doPrechecks(m); err != nil {
		return err
	}
	stage("DoWarnings")
	if err = p.doWarnings(m); err != nil {
		return err
	}
	m.G().LocalSigchainGuard().Set(m.Ctx(), "Prove")
	defer m.G().LocalSigchainGuard().Clear(m.Ctx(), "Prove")
	stage("GenerateProof")
	if err = p.generateProof(m); err != nil {
		return err
	}
	stage("PostProofToServer")
	if err = p.postProofToServer(m); err != nil {
		return err
	}
	m.G().LocalSigchainGuard().Clear(m.Ctx(), "Prove")
	stage("CheckProofText")
	if err = p.checkProofText(m); err != nil {
		return err
	}
	stage("InstructAction")
	if err = p.instructAction(m); err != nil {
		return err
	}

	if !p.arg.PromptPosted {
		m.CDebugf("PromptPosted not set, prove run finished")
		return nil
	}

	stage("PromptPostedLoop")
	if err = p.promptPostedLoop(m); err != nil {
		return err
	}
	m.UIs().LogUI.Notice("Success!")
	return nil
}
<commit_msg>other username (#14088)<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"strings"

	libkb "github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol/keybase1"
	jsonw "github.com/keybase/go-jsonw"
)

// Prove is an engine used for proving ownership of remote accounts,
// like Twitter, GitHub, etc.
type Prove struct {
	arg        *keybase1.StartProofArg
	me         *libkb.User
	st         libkb.ServiceType
	supersede  bool
	proof      *jsonw.Wrapper
	sig        string
	sigID      keybase1.SigID
	postRes    *libkb.PostProofRes
	signingKey libkb.GenericKey
	sigInner   []byte

	remoteNameNormalized string

	libkb.Contextified
}

// NewProve makes a new Prove Engine given an RPC-friendly ProveArg.
func NewProve(g *libkb.GlobalContext, arg *keybase1.StartProofArg) *Prove {
	if arg.SigVersion == nil || libkb.SigVersion(*arg.SigVersion) == libkb.KeybaseNullSigVersion {
		tmp := keybase1.SigVersion(libkb.GetDefaultSigVersion(g))
		arg.SigVersion = &tmp
	}
	return &Prove{
		arg:          arg,
		Contextified: libkb.NewContextified(g),
	}
}

// Name provides the name of this engine for the engine interface contract
func (p *Prove) Name() string {
	return "Prove"
}

// GetPrereqs returns the engine prereqs.
func (p *Prove) Prereqs() Prereqs {
	return Prereqs{Device: true}
}

// RequiredUIs returns the required UIs.
func (p *Prove) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{
		libkb.LogUIKind,
		libkb.ProveUIKind,
		libkb.SecretUIKind,
	}
}

// SubConsumers returns the other UI consumers for this engine.
func (p *Prove) SubConsumers() []libkb.UIConsumer {
	return nil
}

func (p *Prove) loadMe(m libkb.MetaContext) (err error) {
	p.me, err = libkb.LoadMe(libkb.NewLoadUserArgWithMetaContext(m).WithForceReload())
	return err
}

func (p *Prove) checkExists1(m libkb.MetaContext) (err error) {
	proofs := p.me.IDTable().GetActiveProofsFor(p.st)
	if len(proofs) != 0 && !p.arg.Force && p.st.LastWriterWins() {
		lst := proofs[len(proofs)-1]
		var redo bool
		redo, err = m.UIs().ProveUI.PromptOverwrite(m.Ctx(), keybase1.PromptOverwriteArg{
			Account: lst.ToDisplayString(),
			Typ:     keybase1.PromptOverwriteType_SOCIAL,
		})
		if err != nil {
			return err
		}
		if !redo {
			return libkb.NotConfirmedError{}
		}
		p.supersede = true
	}
	return nil
}

func (p *Prove) promptRemoteName(m libkb.MetaContext) error {
	// If the name is already supplied, there's no need to prompt.
	if len(p.arg.Username) > 0 {
		remoteNameNormalized, err := p.st.NormalizeRemoteName(m, p.arg.Username)
		if err == nil {
			p.remoteNameNormalized = remoteNameNormalized
		}
		return err
	}

	// Prompt for the name, retrying if it's invalid.
	var normalizationError error
	for {
		un, err := m.UIs().ProveUI.PromptUsername(m.Ctx(), keybase1.PromptUsernameArg{
			Prompt:    p.st.GetPrompt(),
			PrevError: libkb.ExportErrorAsStatus(m.G(), normalizationError),
		})
		if err != nil {
			// Errors here are conditions like EOF. Return them rather than retrying.
			return err
		}
		var remoteNameNormalized string
		remoteNameNormalized, normalizationError = p.st.NormalizeRemoteName(m, un)
		if normalizationError == nil {
			p.remoteNameNormalized = remoteNameNormalized
			return nil
		}
	}
}

func (p *Prove) checkExists2(m libkb.MetaContext) (err error) {
	defer m.CTrace("Prove#CheckExists2", func() error { return err })()
	if !p.st.LastWriterWins() {
		var found libkb.RemoteProofChainLink
		for _, proof := range p.me.IDTable().GetActiveProofsFor(p.st) {
			_, name := proof.ToKeyValuePair()
			if libkb.Cicmp(name, p.remoteNameNormalized) {
				found = proof
				break
			}
		}
		if found != nil {
			var redo bool
			redo, err = m.UIs().ProveUI.PromptOverwrite(m.Ctx(), keybase1.PromptOverwriteArg{
				Account: found.ToDisplayString(),
				Typ:     keybase1.PromptOverwriteType_SITE,
			})
			if err != nil {
				return err
			}
			if !redo {
				err = libkb.NotConfirmedError{}
				return err
			}
			p.supersede = true
		}
	}
	return nil
}

func (p *Prove) doPrechecks(m libkb.MetaContext) (err error) {
	var w *libkb.Markup
	w, err = p.st.PreProofCheck(m, p.remoteNameNormalized)
	if w != nil {
		if uierr := m.UIs().ProveUI.OutputPrechecks(m.Ctx(), keybase1.OutputPrechecksArg{Text: w.Export()}); uierr != nil {
			m.CWarningf("prove ui OutputPrechecks call error: %s", uierr)
		}
	}
	return err
}

func (p *Prove) doWarnings(m libkb.MetaContext) (err error) {
	if mu := p.st.PreProofWarning(p.remoteNameNormalized); mu != nil {
		var ok bool
		arg := keybase1.PreProofWarningArg{Text: mu.Export()}
		if ok, err = m.UIs().ProveUI.PreProofWarning(m.Ctx(), arg); err == nil && !ok {
			err = libkb.NotConfirmedError{}
			return err
		}
	}
	return nil
}

func (p *Prove) generateProof(m libkb.MetaContext) (err error) {
	ska := libkb.SecretKeyArg{
		Me:      p.me,
		KeyType: libkb.DeviceSigningKeyType,
	}

	p.signingKey, err = m.G().Keyrings.GetSecretKeyWithPrompt(m, m.SecretKeyPromptArg(ska, "tracking signature"))
	if err != nil {
		return err
	}

	sigVersion := libkb.SigVersion(*p.arg.SigVersion)

	if p.proof, err = p.me.ServiceProof(m, p.signingKey, p.st, p.remoteNameNormalized, sigVersion); err != nil {
		return err
	}

	if p.sigInner, err = p.proof.Marshal(); err != nil {
		return err
	}

	p.sig, p.sigID, _, err = libkb.MakeSig(
		m,
		p.signingKey,
		libkb.LinkTypeWebServiceBinding,
		p.sigInner,
		libkb.SigHasRevokes(false),
		keybase1.SeqType_PUBLIC,
		libkb.SigIgnoreIfUnsupported(false),
		p.me,
		sigVersion,
	)

	return err
}

func (p *Prove) postProofToServer(m libkb.MetaContext) (err error) {
	arg := libkb.PostProofArg{
		Sig:               p.sig,
		ProofType:         p.st.GetProofType(),
		RemoteServiceType: p.st.GetTypeName(),
		ID:                p.sigID,
		Supersede:         p.supersede,
		RemoteUsername:    p.remoteNameNormalized,
		RemoteKey:         p.st.GetAPIArgKey(),
		SigningKey:        p.signingKey,
	}
	if libkb.SigVersion(*p.arg.SigVersion) == libkb.KeybaseSignatureV2 {
		arg.SigInner = p.sigInner
	}
	p.postRes, err = libkb.PostProof(m, arg)
	return err
}

func (p *Prove) instructAction(m libkb.MetaContext) (err error) {
	mkp := p.st.PostInstructions(p.remoteNameNormalized)
	var txt string
	if txt, err = p.st.FormatProofText(m, p.postRes, p.me.GetNormalizedName().String(), p.sigID); err != nil {
		return err
	}
	err = m.UIs().ProveUI.OutputInstructions(m.Ctx(), keybase1.OutputInstructionsArg{
		Instructions: mkp.Export(),
		// If we don't trim newlines here, we'll run into an issue where e.g.
		// Facebook links get corrupted on iOS. See:
		// - https://keybase.atlassian.net/browse/DESKTOP-3335
		// - https://keybase.atlassian.net/browse/CORE-4941
		// All of our proof verifying code (PVL) should already be flexible
		// with surrounding whitespace, because users are pasting proofs by
		// hand anyway.
		Proof: strings.TrimSpace(txt),
	})
	if err != nil {
		return err
	}

	return p.checkAutoPost(m, txt)
}

func (p *Prove) checkAutoPost(m libkb.MetaContext, txt string) error {
	if !p.arg.Auto {
		return nil
	}
	if libkb.RemoteServiceTypes[p.arg.Service] != keybase1.ProofType_ROOTER {
		return nil
	}
	m.CDebugf("making automatic post of proof to rooter")
	apiArg := libkb.APIArg{
		Endpoint:    "rooter",
		SessionType: libkb.APISessionTypeREQUIRED,
		Args: libkb.HTTPArgs{
			"post":     libkb.S{Val: txt},
			"username": libkb.S{Val: p.arg.Username},
		},
		NetContext: m.Ctx(),
	}
	if _, err := m.G().API.Post(apiArg); err != nil {
		m.CDebugf("error posting to rooter: %s", err)
		return err
	}
	return nil
}

func (p *Prove) promptPostedLoop(m libkb.MetaContext) (err error) {
	found := false
	for i := 0; ; i++ {
		var retry bool
		var status keybase1.ProofStatus
		var warn *libkb.Markup
		retry, err = m.UIs().ProveUI.OkToCheck(m.Ctx(), keybase1.OkToCheckArg{
			Name:    p.st.DisplayName(p.remoteNameNormalized),
			Attempt: i,
		})
		if !retry || err != nil {
			break
		}
		found, status, _, err = libkb.CheckPosted(m, p.postRes.ID)
		if found || err != nil {
			break
		}
		warn, err = p.st.RecheckProofPosting(i, status, p.remoteNameNormalized)
		if warn != nil {
			uierr := m.UIs().ProveUI.DisplayRecheckWarning(m.Ctx(), keybase1.DisplayRecheckWarningArg{
				Text: warn.Export(),
			})
			if uierr != nil {
				m.CWarningf("prove ui DisplayRecheckWarning call error: %s", uierr)
			}
		}
		if err != nil {
			break
		}
	}
	if !found && err == nil {
		err = libkb.ProofNotYetAvailableError{}
	}

	return err
}

func (p *Prove) checkProofText(m libkb.MetaContext) error {
	m.CDebugf("p.postRes.Text: %q", p.postRes.Text)
	m.CDebugf("p.sigID: %q", p.sigID)
	return p.st.CheckProofText(p.postRes.Text, p.sigID, p.sig)
}

func (p *Prove) getServiceType(m libkb.MetaContext) (err error) {
	if p.st = m.G().GetProofServices().GetServiceType(p.arg.Service); p.st == nil {
		err = libkb.BadServiceError{Service: p.arg.Service}
	}
	return err
}

// SigID returns the signature id of the proof posted to the
// server.
func (p *Prove) SigID() keybase1.SigID {
	return p.sigID
}

// Run runs the Prove engine, performing all steps of the proof process.
func (p *Prove) Run(m libkb.MetaContext) (err error) {
	defer m.CTrace("ProofEngine.Run", func() error { return err })()

	stage := func(s string) {
		m.CDebugf("| ProofEngine.Run() %s", s)
	}

	stage("GetServiceType")
	if err = p.getServiceType(m); err != nil {
		return err
	}
	stage("LoadMe")
	if err = p.loadMe(m); err != nil {
		return err
	}
	stage("CheckExists1")
	if err = p.checkExists1(m); err != nil {
		return err
	}
	stage("PromptRemoteName")
	if err = p.promptRemoteName(m); err != nil {
		return err
	}
	stage("CheckExists2")
	if err = p.checkExists2(m); err != nil {
		return err
	}
	stage("DoPrechecks")
	if err = p.doPrechecks(m); err != nil {
		return err
	}
	stage("DoWarnings")
	if err = p.doWarnings(m); err != nil {
		return err
	}
	m.G().LocalSigchainGuard().Set(m.Ctx(), "Prove")
	defer m.G().LocalSigchainGuard().Clear(m.Ctx(), "Prove")
	stage("GenerateProof")
	if err = p.generateProof(m); err != nil {
		return err
	}
	stage("PostProofToServer")
	if err = p.postProofToServer(m); err != nil {
		return err
	}
	m.G().LocalSigchainGuard().Clear(m.Ctx(), "Prove")
	stage("CheckProofText")
	if err = p.checkProofText(m); err != nil {
		return err
	}
	stage("InstructAction")
	if err = p.instructAction(m); err != nil {
		return err
	}

	if !p.arg.PromptPosted {
		m.CDebugf("PromptPosted not set, prove run finished")
		return nil
	}

	stage("PromptPostedLoop")
	if err = p.promptPostedLoop(m); err != nil {
		return err
	}
	m.UIs().LogUI.Notice("Success!")
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"
	"strings"
	"text/template"

	"github.com/rapidloop/nrpc"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/protoc-gen-go/descriptor"
	"github.com/golang/protobuf/protoc-gen-go/generator"
	plugin "github.com/golang/protobuf/protoc-gen-go/plugin"
)

// baseName returns the last path element of the name, with the last dotted suffix removed.
func baseName(name string) string {
	// First, find the last element
	if i := strings.LastIndex(name, "/"); i >= 0 {
		name = name[i+1:]
	}
	// Now drop the suffix
	if i := strings.LastIndex(name, "."); i >= 0 {
		name = name[0:i]
	}
	return name
}

// goPackageOption interprets the file's go_package option.
// If there is no go_package, it returns ("", "", false).
// If there's a simple name, it returns ("", pkg, true).
// If the option implies an import path, it returns (impPath, pkg, true).
func goPackageOption(d *descriptor.FileDescriptorProto) (impPath, pkg string, ok bool) {
	pkg = d.GetOptions().GetGoPackage()
	if pkg == "" {
		return
	}
	ok = true
	// The presence of a slash implies there's an import path.
	slash := strings.LastIndex(pkg, "/")
	if slash < 0 {
		return
	}
	impPath, pkg = pkg, pkg[slash+1:]
	// A semicolon-delimited suffix overrides the package name.
	sc := strings.IndexByte(impPath, ';')
	if sc < 0 {
		return
	}
	impPath, pkg = impPath[:sc], impPath[sc+1:]
	return
}

// goPackageName returns the Go package name to use in the
// generated Go file.  The result explicit reports whether the name
// came from an option go_package statement.  If explicit is false,
// the name was derived from the protocol buffer's package statement
// or the input file name.
func goPackageName(d *descriptor.FileDescriptorProto) (name string, explicit bool) {
	// Does the file have a "go_package" option?
	if _, pkg, ok := goPackageOption(d); ok {
		return pkg, true
	}

	// Does the file have a package clause?
	if pkg := d.GetPackage(); pkg != "" {
		return pkg, false
	}
	// Use the file base name.
	return baseName(d.GetName()), false
}

// goFileName returns the output name for the generated Go file.
func goFileName(d *descriptor.FileDescriptorProto) string {
	name := *d.Name
	if ext := path.Ext(name); ext == ".proto" || ext == ".protodevel" {
		name = name[:len(name)-len(ext)]
	}
	name += ".nrpc.go"

	// Does the file have a "go_package" option?
	// If it does, it may override the filename.
	if impPath, _, ok := goPackageOption(d); ok && impPath != "" {
		// Replace the existing dirname with the declared import path.
		_, name = path.Split(name)
		name = path.Join(impPath, name)
		return name
	}

	return name
}

func fieldGoType(td *descriptor.FieldDescriptorProto) string {
	// Use protoc-gen-go generator to get the actual go type (for plain types
	// only!)
	t, _ := (*generator.Generator)(nil).GoType(nil, td)
	// We assume proto3, but pass nil to the generator, which will assume proto2.
	// The consequence is a leading star on the type that we need to trim
	return strings.TrimPrefix(t, "*")
}

// splitMessageTypeName split a message type into (package name, type name)
func splitMessageTypeName(name string) (string, string) {
	if len(name) == 0 {
		log.Fatal("Empty message type")
	}
	if name[0] != '.' {
		log.Fatalf("Expect message type name to start with '.', but it is '%s'", name)
	}
	lastDot := strings.LastIndex(name, ".")
	return name[1:lastDot], name[lastDot+1:]
}

func splitTypePath(name string) []string {
	if len(name) == 0 {
		log.Fatal("Empty message type")
	}
	if name[0] != '.' {
		log.Fatalf("Expect message type name to start with '.', but it is '%s'", name)
	}
	return strings.Split(name[1:], ".")
}

func lookupFileDescriptor(name string) *descriptor.FileDescriptorProto {
	for _, fd := range request.GetProtoFile() {
		if fd.GetPackage() == name {
			return fd
		}
	}
	return nil
}

func lookupMessageType(name string) (*descriptor.FileDescriptorProto, *descriptor.DescriptorProto) {
	path := splitTypePath(name)

	pkgpath := path[:len(path)-1]

	var fd *descriptor.FileDescriptorProto
	for {
		pkgname := strings.Join(pkgpath, ".")
		fd = lookupFileDescriptor(pkgname)
		if fd != nil {
			break
		}
		if len(pkgpath) == 1 {
			log.Fatalf("Could not find the .proto file for package '%s' (from message %s)", pkgname, name)
		}
		pkgpath = pkgpath[:len(pkgpath)-1]
	}

	path = path[len(pkgpath):]

	var d *descriptor.DescriptorProto
	for _, mt := range fd.GetMessageType() {
		if mt.GetName() == path[0] {
			d = mt
			break
		}
	}
	if d == nil {
		log.Fatalf("No such type '%s' in package '%s'", path[0], strings.Join(pkgpath, "."))
	}
	for i, token := range path[1:] {
		var found bool
		for _, nd := range d.GetNestedType() {
			if nd.GetName() == token {
				d = nd
				found = true
				break
			}
		}
		if !found {
			log.Fatalf("No such nested type '%s' in '%s.%s'",
				token, strings.Join(pkgpath, "."), strings.Join(path[:i+1], "."))
		}
	}
	return fd, d
}

func getField(d *descriptor.DescriptorProto, name string) *descriptor.FieldDescriptorProto {
	for _, f := range d.GetField() {
		if f.GetName() == name {
			return f
		}
	}
	return nil
}

func getOneofDecl(d *descriptor.DescriptorProto, name string) *descriptor.OneofDescriptorProto {
	for _, of := range d.GetOneofDecl() {
		if of.GetName() == name {
			return of
		}
	}
	return nil
}

func pkgSubject(fd *descriptor.FileDescriptorProto) string {
	e, err := proto.GetExtension(fd.Options, nrpc.E_PackageSubject)
	if err == nil {
		value := e.(*string)
		return *value
	}
	return fd.GetPackage()
}

func getResultType(
	md *descriptor.MethodDescriptorProto,
) string {
	return md.GetOutputType()
}

func getGoType(pbType string) (string, string) {
	if !strings.Contains(pbType, ".") {
		return "", pbType
	}
	fd, _ := lookupMessageType(pbType)
	name := strings.TrimPrefix(pbType, "."+fd.GetPackage()+".")
	name = strings.Replace(name, ".", "_", -1)
	return fd.GetOptions().GetGoPackage(), name
}

func getPkgImportName(goPkg string) string {
	if goPkg == currentFile.GetOptions().GetGoPackage() {
		return ""
	}
	replacer := strings.NewReplacer(".", "_", "/", "_", "-", "_")
	return replacer.Replace(goPkg)
}

var pluginPrometheus bool

var funcMap = template.FuncMap{
	"GoPackageName": func(fd *descriptor.FileDescriptorProto) string {
		p, _ := goPackageName(fd)
		return p
	},
	"GetPkg": func(pkg, s string) string {
		s = strings.TrimPrefix(s, ".")
		s = strings.TrimPrefix(s, pkg)
		s = strings.TrimPrefix(s, ".")
		return s
	},
	"GetExtraImports": func(fd *descriptor.FileDescriptorProto) []string {
		// check all the types used and imports packages from where they come
		var imports = make(map[string]string)
		for _, sd := range fd.GetService() {
			for _, md := range sd.GetMethod() {
				goPkg, _ := getGoType(md.GetInputType())
				pkgImportName := getPkgImportName(goPkg)
				if pkgImportName != "" {
					imports[pkgImportName] = goPkg
				}
				goPkg, _ = getGoType(getResultType(md))
				pkgImportName = getPkgImportName(goPkg)
				if pkgImportName != "" {
					imports[pkgImportName] = goPkg
				}
			}
		}
		var result []string
		for importName, goPkg := range imports {
			result = append(result,
				fmt.Sprintf("%s \"%s\"",
					importName,
					goPkg,
				),
			)
		}
		return result
	},
	"GetPkgSubjectPrefix": func(fd *descriptor.FileDescriptorProto) string {
		if s := pkgSubject(fd); s != "" {
			return s + "."
		}
		return ""
	},
	"GetPkgSubject": pkgSubject,
	"GetPkgSubjectParams": func(fd *descriptor.FileDescriptorProto) []string {
		e, err := proto.GetExtension(fd.Options, nrpc.E_PackageSubjectParams)
		if err == nil {
			value := e.([]string)
			return value
		}
		return nil
	},
	"GetServiceSubject": func(sd *descriptor.ServiceDescriptorProto) string {
		if opts := sd.GetOptions(); opts != nil {
			s, err := proto.GetExtension(opts, nrpc.E_ServiceSubject)
			if err == nil {
				value := s.(*string)
				return *value
			}
		}
		if opts := currentFile.GetOptions(); opts != nil {
			s, err := proto.GetExtension(opts, nrpc.E_ServiceSubjectRule)
			if err == nil {
				switch *(s.(*nrpc.SubjectRule)) {
				case nrpc.SubjectRule_COPY:
					return sd.GetName()
				case nrpc.SubjectRule_TOLOWER:
					return strings.ToLower(sd.GetName())
				}
			}
		}
		return sd.GetName()
	},
	"ServiceNeedsHandler": func(sd *descriptor.ServiceDescriptorProto) bool {
		for _, md := range sd.GetMethod() {
			if md.GetInputType() != ".nrpc.NoRequest" {
				return true
			}
		}
		return false
	},
	"GetMethodSubject": func(md *descriptor.MethodDescriptorProto) string {
		if opts := md.GetOptions(); opts != nil {
			s, err := proto.GetExtension(opts, nrpc.E_MethodSubject)
			if err == nil {
				value := s.(*string)
				return *value
			}
		}
		if opts := currentFile.GetOptions(); opts != nil {
			s, err := proto.GetExtension(opts, nrpc.E_MethodSubjectRule)
			if err == nil {
				switch *(s.(*nrpc.SubjectRule)) {
				case nrpc.SubjectRule_COPY:
					return md.GetName()
				case nrpc.SubjectRule_TOLOWER:
					return strings.ToLower(md.GetName())
				}
			}
		}
		return md.GetName()
	},
	"GetMethodSubjectParams": func(md *descriptor.MethodDescriptorProto) []string {
		if opts := md.GetOptions(); opts != nil {
			if e, err := proto.GetExtension(opts, nrpc.E_MethodSubjectParams); err == nil {
				return e.([]string)
			}
		}
		return []string{}
	},
	"GetServiceSubjectParams": func(sd *descriptor.ServiceDescriptorProto) []string {
		if opts := sd.GetOptions(); opts != nil {
			if e, err := proto.GetExtension(opts, nrpc.E_ServiceSubjectParams); err == nil {
				return e.([]string)
			}
		}
		return []string{}
	},
	"HasStreamedReply": func(md *descriptor.MethodDescriptorProto) bool {
		if opts := md.GetOptions(); opts != nil {
			if e, err := proto.GetExtension(opts, nrpc.E_StreamedReply); err == nil {
				return *(e.(*bool))
			}
		}
		return false
	},
	"Prometheus": func() bool {
		return pluginPrometheus
	},
	"GetResultType": getResultType,
	"GoType": func(pbType string) string {
		goPkg, goType := getGoType(pbType)
		if goPkg != "" {
			importName := getPkgImportName(goPkg)
			if importName != "" {
				goType = importName + "." + goType
			}
		}
		return goType
	},
}

var request plugin.CodeGeneratorRequest
var currentFile *descriptor.FileDescriptorProto

func main() {

	log.SetPrefix("protoc-gen-nrpc: ")
	data, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		log.Fatalf("error: reading input: %v", err)
	}

	var response plugin.CodeGeneratorResponse
	if err := proto.Unmarshal(data, &request); err != nil {
		log.Fatalf("error: parsing input proto: %v", err)
	}

	if len(request.GetFileToGenerate()) == 0 {
		log.Fatal("error: no files to generate")
	}

	pluginPrometheus = request.GetParameter() == "plugins=prometheus"

	tmpl, err := template.New(".").Funcs(funcMap).Parse(tFile)
	if err != nil {
		log.Fatal(err)
	}

	for _, name := range request.GetFileToGenerate() {
		var fd *descriptor.FileDescriptorProto
		for _, fd = range request.GetProtoFile() {
			if name == fd.GetName() {
				break
			}
		}
		if fd == nil {
			log.Fatalf("could not find the .proto file for %s", name)
		}

		currentFile = fd

		var buf bytes.Buffer
		if err := tmpl.Execute(&buf, fd); err != nil {
			log.Fatal(err)
		}

		currentFile = nil

		response.File = append(response.File, &plugin.CodeGeneratorResponse_File{
			Name:    proto.String(goFileName(fd)),
			Content: proto.String(buf.String()),
		})
	}

	if data, err = proto.Marshal(&response); err != nil {
		log.Fatalf("error: failed to marshal output proto: %v", err)
	}
	if _, err := os.Stdout.Write(data); err != nil {
		log.Fatalf("error: failed to write output proto: %v", err)
	}
}
<commit_msg>Support go_package with semicolon<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"
	"strings"
	"text/template"

	"github.com/rapidloop/nrpc"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/protoc-gen-go/descriptor"
	"github.com/golang/protobuf/protoc-gen-go/generator"
	plugin "github.com/golang/protobuf/protoc-gen-go/plugin"
)

// baseName returns the last path element of the name, with the last dotted suffix removed.
func baseName(name string) string {
	// First, find the last element
	if i := strings.LastIndex(name, "/"); i >= 0 {
		name = name[i+1:]
	}
	// Now drop the suffix
	if i := strings.LastIndex(name, "."); i >= 0 {
		name = name[0:i]
	}
	return name
}

// getGoPackage returns the file's go_package option.
// If it containts a semicolon, only the part before it is returned.
func getGoPackage(fd *descriptor.FileDescriptorProto) string {
	pkg := fd.GetOptions().GetGoPackage()
	if strings.Contains(pkg, ";") {
		parts := strings.Split(pkg, ";")
		if len(parts) > 2 {
			log.Fatalf(
				"protoc-gen-nrpc: go_package '%s' contains more than 1 ';'",
				pkg)
		}
		pkg = parts[0]
	}

	return pkg
}

// goPackageOption interprets the file's go_package option.
// If there is no go_package, it returns ("", "", false).
// If there's a simple name, it returns ("", pkg, true).
// If the option implies an import path, it returns (impPath, pkg, true).
func goPackageOption(d *descriptor.FileDescriptorProto) (impPath, pkg string, ok bool) {
	pkg = getGoPackage(d)
	if pkg == "" {
		return
	}
	ok = true
	// The presence of a slash implies there's an import path.
	slash := strings.LastIndex(pkg, "/")
	if slash < 0 {
		return
	}
	impPath, pkg = pkg, pkg[slash+1:]
	// A semicolon-delimited suffix overrides the package name.
	sc := strings.IndexByte(impPath, ';')
	if sc < 0 {
		return
	}
	impPath, pkg = impPath[:sc], impPath[sc+1:]
	return
}

// goPackageName returns the Go package name to use in the
// generated Go file.  The result explicit reports whether the name
// came from an option go_package statement.  If explicit is false,
// the name was derived from the protocol buffer's package statement
// or the input file name.
func goPackageName(d *descriptor.FileDescriptorProto) (name string, explicit bool) {
	// Does the file have a "go_package" option?
	if _, pkg, ok := goPackageOption(d); ok {
		return pkg, true
	}

	// Does the file have a package clause?
	if pkg := d.GetPackage(); pkg != "" {
		return pkg, false
	}
	// Use the file base name.
	return baseName(d.GetName()), false
}

// goFileName returns the output name for the generated Go file.
func goFileName(d *descriptor.FileDescriptorProto) string {
	name := *d.Name
	if ext := path.Ext(name); ext == ".proto" || ext == ".protodevel" {
		name = name[:len(name)-len(ext)]
	}
	name += ".nrpc.go"

	// Does the file have a "go_package" option?
	// If it does, it may override the filename.
	if impPath, _, ok := goPackageOption(d); ok && impPath != "" {
		// Replace the existing dirname with the declared import path.
		_, name = path.Split(name)
		name = path.Join(impPath, name)
		return name
	}

	return name
}

func fieldGoType(td *descriptor.FieldDescriptorProto) string {
	// Use protoc-gen-go generator to get the actual go type (for plain types
	// only!)
	t, _ := (*generator.Generator)(nil).GoType(nil, td)
	// We assume proto3, but pass nil to the generator, which will assume proto2.
	// The consequence is a leading star on the type that we need to trim
	return strings.TrimPrefix(t, "*")
}

// splitMessageTypeName split a message type into (package name, type name)
func splitMessageTypeName(name string) (string, string) {
	if len(name) == 0 {
		log.Fatal("Empty message type")
	}
	if name[0] != '.' {
		log.Fatalf("Expect message type name to start with '.', but it is '%s'", name)
	}
	lastDot := strings.LastIndex(name, ".")
	return name[1:lastDot], name[lastDot+1:]
}

func splitTypePath(name string) []string {
	if len(name) == 0 {
		log.Fatal("Empty message type")
	}
	if name[0] != '.' {
		log.Fatalf("Expect message type name to start with '.', but it is '%s'", name)
	}
	return strings.Split(name[1:], ".")
}

func lookupFileDescriptor(name string) *descriptor.FileDescriptorProto {
	for _, fd := range request.GetProtoFile() {
		if fd.GetPackage() == name {
			return fd
		}
	}
	return nil
}

func lookupMessageType(name string) (*descriptor.FileDescriptorProto, *descriptor.DescriptorProto) {
	path := splitTypePath(name)

	pkgpath := path[:len(path)-1]

	var fd *descriptor.FileDescriptorProto
	for {
		pkgname := strings.Join(pkgpath, ".")
		fd = lookupFileDescriptor(pkgname)
		if fd != nil {
			break
		}
		if len(pkgpath) == 1 {
			log.Fatalf("Could not find the .proto file for package '%s' (from message %s)", pkgname, name)
		}
		pkgpath = pkgpath[:len(pkgpath)-1]
	}

	path = path[len(pkgpath):]

	var d *descriptor.DescriptorProto
	for _, mt := range fd.GetMessageType() {
		if mt.GetName() == path[0] {
			d = mt
			break
		}
	}
	if d == nil {
		log.Fatalf("No such type '%s' in package '%s'", path[0], strings.Join(pkgpath, "."))
	}
	for i, token := range path[1:] {
		var found bool
		for _, nd := range d.GetNestedType() {
			if nd.GetName() == token {
				d = nd
				found = true
				break
			}
		}
		if !found {
			log.Fatalf("No such nested type '%s' in '%s.%s'",
				token, strings.Join(pkgpath, "."), strings.Join(path[:i+1], "."))
		}
	}
	return fd, d
}

func getField(d *descriptor.DescriptorProto, name string) *descriptor.FieldDescriptorProto {
	for _, f := range d.GetField() {
		if f.GetName() == name {
			return f
		}
	}
	return nil
}

func getOneofDecl(d *descriptor.DescriptorProto, name string) *descriptor.OneofDescriptorProto {
	for _, of := range d.GetOneofDecl() {
		if of.GetName() == name {
			return of
		}
	}
	return nil
}

func pkgSubject(fd *descriptor.FileDescriptorProto) string {
	e, err := proto.GetExtension(fd.Options, nrpc.E_PackageSubject)
	if err == nil {
		value := e.(*string)
		return *value
	}
	return fd.GetPackage()
}

func getResultType(
	md *descriptor.MethodDescriptorProto,
) string {
	return md.GetOutputType()
}

func getGoType(pbType string) (string, string) {
	if !strings.Contains(pbType, ".") {
		return "", pbType
	}
	fd, _ := lookupMessageType(pbType)
	name := strings.TrimPrefix(pbType, "."+fd.GetPackage()+".")
	name = strings.Replace(name, ".", "_", -1)
	return getGoPackage(fd), name
}

func getPkgImportName(goPkg string) string {
	if goPkg == getGoPackage(currentFile) {
		return ""
	}
	replacer := strings.NewReplacer(".", "_", "/", "_", "-", "_")
	return replacer.Replace(goPkg)
}

var pluginPrometheus bool

var funcMap = template.FuncMap{
	"GoPackageName": func(fd *descriptor.FileDescriptorProto) string {
		p, _ := goPackageName(fd)
		return p
	},
	"GetPkg": func(pkg, s string) string {
		s = strings.TrimPrefix(s, ".")
		s = strings.TrimPrefix(s, pkg)
		s = strings.TrimPrefix(s, ".")
		return s
	},
	"GetExtraImports": func(fd *descriptor.FileDescriptorProto) []string {
		// check all the types used and imports packages from where they come
		var imports = make(map[string]string)
		for _, sd := range fd.GetService() {
			for _, md := range sd.GetMethod() {
				goPkg, _ := getGoType(md.GetInputType())
				pkgImportName := getPkgImportName(goPkg)
				if pkgImportName != "" {
					imports[pkgImportName] = goPkg
				}
				goPkg, _ = getGoType(getResultType(md))
				pkgImportName = getPkgImportName(goPkg)
				if pkgImportName != "" {
					imports[pkgImportName] = goPkg
				}
			}
		}
		var result []string
		for importName, goPkg := range imports {
			result = append(result,
				fmt.Sprintf("%s \"%s\"",
					importName,
					goPkg,
				),
			)
		}
		return result
	},
	"GetPkgSubjectPrefix": func(fd *descriptor.FileDescriptorProto) string {
		if s := pkgSubject(fd); s != "" {
			return s + "."
		}
		return ""
	},
	"GetPkgSubject": pkgSubject,
	"GetPkgSubjectParams": func(fd *descriptor.FileDescriptorProto) []string {
		e, err := proto.GetExtension(fd.Options, nrpc.E_PackageSubjectParams)
		if err == nil {
			value := e.([]string)
			return value
		}
		return nil
	},
	"GetServiceSubject": func(sd *descriptor.ServiceDescriptorProto) string {
		if opts := sd.GetOptions(); opts != nil {
			s, err := proto.GetExtension(opts, nrpc.E_ServiceSubject)
			if err == nil {
				value := s.(*string)
				return *value
			}
		}
		if opts := currentFile.GetOptions(); opts != nil {
			s, err := proto.GetExtension(opts, nrpc.E_ServiceSubjectRule)
			if err == nil {
				switch *(s.(*nrpc.SubjectRule)) {
				case nrpc.SubjectRule_COPY:
					return sd.GetName()
				case nrpc.SubjectRule_TOLOWER:
					return strings.ToLower(sd.GetName())
				}
			}
		}
		return sd.GetName()
	},
	"ServiceNeedsHandler": func(sd *descriptor.ServiceDescriptorProto) bool {
		for _, md := range sd.GetMethod() {
			if md.GetInputType() != ".nrpc.NoRequest" {
				return true
			}
		}
		return false
	},
	"GetMethodSubject": func(md *descriptor.MethodDescriptorProto) string {
		if opts := md.GetOptions(); opts != nil {
			s, err := proto.GetExtension(opts, nrpc.E_MethodSubject)
			if err == nil {
				value := s.(*string)
				return *value
			}
		}
		if opts := currentFile.GetOptions(); opts != nil {
			s, err := proto.GetExtension(opts, nrpc.E_MethodSubjectRule)
			if err == nil {
				switch *(s.(*nrpc.SubjectRule)) {
				case nrpc.SubjectRule_COPY:
					return md.GetName()
				case nrpc.SubjectRule_TOLOWER:
					return strings.ToLower(md.GetName())
				}
			}
		}
		return md.GetName()
	},
	"GetMethodSubjectParams": func(md *descriptor.MethodDescriptorProto) []string {
		if opts := md.GetOptions(); opts != nil {
			if e, err := proto.GetExtension(opts, nrpc.E_MethodSubjectParams); err == nil {
				return e.([]string)
			}
		}
		return []string{}
	},
	"GetServiceSubjectParams": func(sd *descriptor.ServiceDescriptorProto) []string {
		if opts := sd.GetOptions(); opts != nil {
			if e, err := proto.GetExtension(opts, nrpc.E_ServiceSubjectParams); err == nil {
				return e.([]string)
			}
		}
		return []string{}
	},
	"HasStreamedReply": func(md *descriptor.MethodDescriptorProto) bool {
		if opts := md.GetOptions(); opts != nil {
			if e, err := proto.GetExtension(opts, nrpc.E_StreamedReply); err == nil {
				return *(e.(*bool))
			}
		}
		return false
	},
	"Prometheus": func() bool {
		return pluginPrometheus
	},
	"GetResultType": getResultType,
	"GoType": func(pbType string) string {
		goPkg, goType := getGoType(pbType)
		if goPkg != "" {
			importName := getPkgImportName(goPkg)
			if importName != "" {
				goType = importName + "." + goType
			}
		}
		return goType
	},
}

var request plugin.CodeGeneratorRequest
var currentFile *descriptor.FileDescriptorProto

func main() {

	log.SetPrefix("protoc-gen-nrpc: ")
	data, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		log.Fatalf("error: reading input: %v", err)
	}

	var response plugin.CodeGeneratorResponse
	if err := proto.Unmarshal(data, &request); err != nil {
		log.Fatalf("error: parsing input proto: %v", err)
	}

	if len(request.GetFileToGenerate()) == 0 {
		log.Fatal("error: no files to generate")
	}

	pluginPrometheus = request.GetParameter() == "plugins=prometheus"

	tmpl, err := template.New(".").Funcs(funcMap).Parse(tFile)
	if err != nil {
		log.Fatal(err)
	}

	for _, name := range request.GetFileToGenerate() {
		var fd *descriptor.FileDescriptorProto
		for _, fd = range request.GetProtoFile() {
			if name == fd.GetName() {
				break
			}
		}
		if fd == nil {
			log.Fatalf("could not find the .proto file for %s", name)
		}

		currentFile = fd

		var buf bytes.Buffer
		if err := tmpl.Execute(&buf, fd); err != nil {
			log.Fatal(err)
		}

		currentFile = nil

		response.File = append(response.File, &plugin.CodeGeneratorResponse_File{
			Name:    proto.String(goFileName(fd)),
			Content: proto.String(buf.String()),
		})
	}

	if data, err = proto.Marshal(&response); err != nil {
		log.Fatalf("error: failed to marshal output proto: %v", err)
	}
	if _, err := os.Stdout.Write(data); err != nil {
		log.Fatalf("error: failed to write output proto: %v", err)
	}
}
<|endoftext|>


<commit_before>package engine

import (
	"github.com/keybase/client/go/libkb"
	jsonw "github.com/keybase/go-jsonw"
)

type TrackOptions struct {
	TrackLocalOnly bool // true: only track locally, false: track locally and remotely
	TrackApprove   bool // true: don't ask for confirmation, false: ask for confirmation
}

type TrackEngineArg struct {
	TheirName string
	Me        *libkb.User
	Options   TrackOptions
}

type TrackEngine struct {
	arg                 *TrackEngineArg
	res                 *IdRes
	them                *libkb.User
	signingKeyPub       libkb.GenericKey
	signingKeyPriv      libkb.GenericKey
	trackStatementBytes []byte
	trackStatement      *jsonw.Wrapper
	sig                 string
	sigid               *libkb.SigId
	lockedKey           *libkb.SKB
}

// NewTrackEngine creates a default TrackEngine for tracking theirName.
func NewTrackEngine(arg *TrackEngineArg) *TrackEngine {
	return &TrackEngine{arg: arg}
}

func (s *TrackEngine) Name() string {
	return "Track"
}

func (e *TrackEngine) GetPrereqs() EnginePrereqs {
	return EnginePrereqs{
		Session: true,
	}
}

func (k *TrackEngine) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{
		libkb.SecretUIKind,
	}
}

func (s *TrackEngine) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{
		NewIdentify(nil),
	}
}

func (e *TrackEngine) Run(ctx *Context) error {
	if err := e.loadMe(); err != nil {
		return err
	}

	iarg := NewIdentifyTrackArg(e.arg.TheirName, true, e.arg.Options)
	ieng := NewIdentify(iarg)
	if err := RunEngine(ieng, ctx); err != nil {
		return err
	}
	ti := ieng.TrackInstructions()
	e.them = ieng.User()

	e.res = &IdRes{Outcome: ieng.Outcome(), User: e.them}

	var err error
	e.signingKeyPub, err = e.arg.Me.SigningKeyPub()
	if err != nil {
		return err
	}

	if e.trackStatement, err = e.arg.Me.TrackingProofFor(e.signingKeyPub, e.them); err != nil {
		return err
	}

	if e.trackStatementBytes, err = e.trackStatement.Marshal(); err != nil {
		return err
	}

	G.Log.Debug("| Tracking statement: %s", string(e.trackStatementBytes))

	if ti.Remote {
		err = e.storeRemoteTrack(ctx)
	} else if ti.Local {
		err = e.storeLocalTrack()
	}
	return err
}

func (e *TrackEngine) User() *libkb.User {
	return e.them
}

/*
func (e *TrackEngine) loadThem() error {
	if e.arg.Them == nil && len(e.arg.TheirName) == 0 {
		return fmt.Errorf("No 'them' passed to TrackEngine")
	}
	if e.arg.Them == nil {
		if u, err := libkb.LoadUser(libkb.LoadUserArg{
			Name:        e.arg.TheirName,
			Self:        false,
			ForceReload: false,
		}); err != nil {
			return err
		} else {
			e.arg.Them = u
		}
	}
	return nil
}
*/

func (e *TrackEngine) loadMe() error {
	if e.arg.Me != nil {
		return nil
	}

	me, err := libkb.LoadMe(libkb.LoadUserArg{})
	if err != nil {
		return err
	}
	e.arg.Me = me
	return nil
}

func (e *TrackEngine) storeLocalTrack() error {
	return libkb.StoreLocalTrack(e.arg.Me.GetUid(), e.them.GetUid(), e.trackStatement)
}

func (e *TrackEngine) storeRemoteTrack(ctx *Context) (err error) {
	G.Log.Debug("+ StoreRemoteTrack")
	defer G.Log.Debug("- StoreRemoteTrack -> %s", libkb.ErrToOk(err))

	arg := libkb.SecretKeyArg{Reason: "tracking signature", Ui: ctx.SecretUI, Me: e.arg.Me, All: true}
	if e.signingKeyPriv, err = G.Keyrings.GetSecretKey(arg); err != nil {
		return
	} else if e.signingKeyPriv == nil {
		err = libkb.NoSecretKeyError{}
		return
	}

	if e.sig, e.sigid, err = e.signingKeyPriv.SignToString(e.trackStatementBytes); err != nil {
		return
	}

	_, err = G.API.Post(libkb.ApiArg{
		Endpoint:    "follow",
		NeedSession: true,
		Args: libkb.HttpArgs{
			"sig_id_base":  libkb.S{Val: e.sigid.ToString(false)},
			"sig_id_short": libkb.S{Val: e.sigid.ToShortId()},
			"sig":          libkb.S{Val: e.sig},
			"uid":          e.them.GetUid(),
			"type":         libkb.S{Val: "track"},
			"signing_kid":  e.signingKeyPub.GetKid(),
		},
	})

	return
}
<commit_msg>cleanup<commit_after>package engine

import (
	"github.com/keybase/client/go/libkb"
	jsonw "github.com/keybase/go-jsonw"
)

type TrackOptions struct {
	TrackLocalOnly bool // true: only track locally, false: track locally and remotely
	TrackApprove   bool // true: don't ask for confirmation, false: ask for confirmation
}

type TrackEngineArg struct {
	TheirName string
	Me        *libkb.User
	Options   TrackOptions
}

type TrackEngine struct {
	arg                 *TrackEngineArg
	res                 *IdRes
	them                *libkb.User
	signingKeyPub       libkb.GenericKey
	signingKeyPriv      libkb.GenericKey
	trackStatementBytes []byte
	trackStatement      *jsonw.Wrapper
	sig                 string
	sigid               *libkb.SigId
	lockedKey           *libkb.SKB
}

// NewTrackEngine creates a default TrackEngine for tracking theirName.
func NewTrackEngine(arg *TrackEngineArg) *TrackEngine {
	return &TrackEngine{arg: arg}
}

func (s *TrackEngine) Name() string {
	return "Track"
}

func (e *TrackEngine) GetPrereqs() EnginePrereqs {
	return EnginePrereqs{
		Session: true,
	}
}

func (k *TrackEngine) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{
		libkb.SecretUIKind,
	}
}

func (s *TrackEngine) SubConsumers() []libkb.UIConsumer {
	return []libkb.UIConsumer{
		NewIdentify(nil),
	}
}

func (e *TrackEngine) Run(ctx *Context) error {
	if err := e.loadMe(); err != nil {
		return err
	}

	iarg := NewIdentifyTrackArg(e.arg.TheirName, true, e.arg.Options)
	ieng := NewIdentify(iarg)
	if err := RunEngine(ieng, ctx); err != nil {
		return err
	}
	ti := ieng.TrackInstructions()
	e.them = ieng.User()

	e.res = &IdRes{Outcome: ieng.Outcome(), User: e.them}

	var err error
	e.signingKeyPub, err = e.arg.Me.SigningKeyPub()
	if err != nil {
		return err
	}

	if e.trackStatement, err = e.arg.Me.TrackingProofFor(e.signingKeyPub, e.them); err != nil {
		return err
	}

	if e.trackStatementBytes, err = e.trackStatement.Marshal(); err != nil {
		return err
	}

	G.Log.Debug("| Tracking statement: %s", string(e.trackStatementBytes))

	if ti.Remote {
		err = e.storeRemoteTrack(ctx)
	} else if ti.Local {
		err = e.storeLocalTrack()
	}
	return err
}

func (e *TrackEngine) User() *libkb.User {
	return e.them
}

func (e *TrackEngine) loadMe() error {
	if e.arg.Me != nil {
		return nil
	}

	me, err := libkb.LoadMe(libkb.LoadUserArg{})
	if err != nil {
		return err
	}
	e.arg.Me = me
	return nil
}

func (e *TrackEngine) storeLocalTrack() error {
	return libkb.StoreLocalTrack(e.arg.Me.GetUid(), e.them.GetUid(), e.trackStatement)
}

func (e *TrackEngine) storeRemoteTrack(ctx *Context) (err error) {
	G.Log.Debug("+ StoreRemoteTrack")
	defer G.Log.Debug("- StoreRemoteTrack -> %s", libkb.ErrToOk(err))

	arg := libkb.SecretKeyArg{Reason: "tracking signature", Ui: ctx.SecretUI, Me: e.arg.Me, All: true}
	if e.signingKeyPriv, err = G.Keyrings.GetSecretKey(arg); err != nil {
		return
	} else if e.signingKeyPriv == nil {
		err = libkb.NoSecretKeyError{}
		return
	}

	if e.sig, e.sigid, err = e.signingKeyPriv.SignToString(e.trackStatementBytes); err != nil {
		return
	}

	_, err = G.API.Post(libkb.ApiArg{
		Endpoint:    "follow",
		NeedSession: true,
		Args: libkb.HttpArgs{
			"sig_id_base":  libkb.S{Val: e.sigid.ToString(false)},
			"sig_id_short": libkb.S{Val: e.sigid.ToShortId()},
			"sig":          libkb.S{Val: e.sig},
			"uid":          e.them.GetUid(),
			"type":         libkb.S{Val: "track"},
			"signing_kid":  e.signingKeyPub.GetKid(),
		},
	})

	return
}
<|endoftext|>


<commit_before>// Copyright 2016 Marc-Antoine Ruel. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

package modules

import (
	"errors"
	"sync"
	"time"

	mqtt "github.com/eclipse/paho.mqtt.golang"
)

// MQTT main purpose is to hide the horror that paho.mqtt.golang is.
//
// It intentionally has a much simpler surface.
type MQTT struct {
	client mqtt.Client

	mu          sync.Mutex
	subscribers []*subscription
}

func New(server, clientID, user, password string) (*MQTT, error) {
	opts := mqtt.NewClientOptions().AddBroker(server).SetClientID(clientID)
	// Use lower timeouts than the defaults since they are high and the current
	// assumption is local network.
	opts.SetConnectTimeout(5 * time.Second)
	opts.SetKeepAlive(4 * time.Second)
	opts.SetPingTimeout(2 * time.Second)
	if len(user) != 0 {
		opts.SetUsername(user)
	}
	if len(password) != 0 {
		opts.SetPassword(password)
	}
	// TODO(maruel): opts.SetTLSConfig()
	// https://github.com/eclipse/paho.mqtt.golang/blob/master/samples/ssl.go
	// TODO(maruel): opts.SetBinaryWill()
	client := mqtt.NewClient(opts)
	if token := client.Connect(); token.Wait() && token.Error() != nil {
		return nil, token.Error()
	}
	m := &MQTT{client: client}

	// TODO(maruel): Temporary.
	token := m.client.Subscribe("#", byte(BestEffort), func(client mqtt.Client, msgQ mqtt.Message) {
		msg := Message{msgQ.Topic(), msgQ.Payload()}
		m.mu.Lock()
		defer m.mu.Unlock()
		for i := range m.subscribers {
			if m.subscribers[i].topic.match(msg.Topic) {
				m.subscribers[i].publish(msg)
			}
		}

	})
	token.Wait()
	if err := token.Error(); err != nil {
		return nil, err
	}

	return m, nil
}

func (m *MQTT) Close() error {
	m.client.Disconnect(500)
	m.client = nil
	return nil
}

func (m *MQTT) Publish(msg Message, qos QOS, retained bool) error {
	// Make it back synchronous.
	token := m.client.Publish(msg.Topic, byte(qos), retained, msg.Payload)
	token.Wait()
	return token.Error()
}

func (m *MQTT) Subscribe(topic string, qos QOS) (<-chan Message, error) {
	/*
		c := make(chan Message)
		token := m.client.Subscribe(topic, byte(qos), func(client mqtt.Client, msg mqtt.Message) {
			c <- Message{msg.Topic(), msg.Payload()}
		})
		token.Wait()
		return c, token.Error()
	*/
	p := parseTopic(topic)
	if p == nil {
		return nil, errors.New("invalid topic")
	}
	c := make(chan Message)
	m.mu.Lock()
	defer m.mu.Unlock()
	m.subscribers = append(m.subscribers, &subscription{topic: p, channel: c})
	return c, nil

}

func (m *MQTT) Unsubscribe(topic string) error {
	/*
		token := m.client.Unsubscribe(topic)
		token.Wait()
		return token.Error()
	*/
	p := parseTopic(topic)
	if p == nil {
		return errors.New("invalid topic")
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	for i := range m.subscribers {
		if m.subscribers[i].topic.isEqual(p) {
			m.subscribers[i].Close()
			copy(m.subscribers[i:], m.subscribers[i+1:])
			m.subscribers = m.subscribers[:len(m.subscribers)-1]
			return nil
		}
	}
	return errors.New("subscription not found")
}

func (m *MQTT) Get(topic string, qos QOS) ([]Message, error) {
	// TODO(maruel): It looks it needs to do a quick Subscribe + poll every
	// messages until one with !msg.Retained() or a timeout then Unsubscribe.
	return nil, errors.New("implement me")
}

var _ Bus = &MQTT{}
<commit_msg>Disable waiting when publishing.<commit_after>// Copyright 2016 Marc-Antoine Ruel. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

package modules

import (
	"errors"
	"sync"
	"time"

	mqtt "github.com/eclipse/paho.mqtt.golang"
)

// MQTT main purpose is to hide the horror that paho.mqtt.golang is.
//
// It intentionally has a much simpler surface.
type MQTT struct {
	client mqtt.Client

	mu          sync.Mutex
	subscribers []*subscription
}

func New(server, clientID, user, password string) (*MQTT, error) {
	opts := mqtt.NewClientOptions().AddBroker(server).SetClientID(clientID)
	// Use lower timeouts than the defaults since they are high and the current
	// assumption is local network.
	opts.SetConnectTimeout(5 * time.Second)
	opts.SetKeepAlive(4 * time.Second)
	opts.SetPingTimeout(2 * time.Second)
	if len(user) != 0 {
		opts.SetUsername(user)
	}
	if len(password) != 0 {
		opts.SetPassword(password)
	}
	// TODO(maruel): opts.SetTLSConfig()
	// https://github.com/eclipse/paho.mqtt.golang/blob/master/samples/ssl.go
	// TODO(maruel): opts.SetBinaryWill()
	client := mqtt.NewClient(opts)
	if token := client.Connect(); token.Wait() && token.Error() != nil {
		return nil, token.Error()
	}
	m := &MQTT{client: client}

	// TODO(maruel): Temporary.
	token := m.client.Subscribe("#", byte(BestEffort), func(client mqtt.Client, msgQ mqtt.Message) {
		msg := Message{msgQ.Topic(), msgQ.Payload()}
		m.mu.Lock()
		defer m.mu.Unlock()
		for i := range m.subscribers {
			if m.subscribers[i].topic.match(msg.Topic) {
				m.subscribers[i].publish(msg)
			}
		}

	})
	token.Wait()
	if err := token.Error(); err != nil {
		return nil, err
	}

	return m, nil
}

func (m *MQTT) Close() error {
	m.client.Disconnect(500)
	m.client = nil
	return nil
}

func (m *MQTT) Publish(msg Message, qos QOS, retained bool) error {
	// TODO(maruel): Make it back synchronous.
	m.client.Publish(msg.Topic, byte(qos), retained, msg.Payload)
	return nil
	// TODO(maruel): I think it may induce multi-second delay. Horrible hack.
	//token := m.client.Publish(msg.Topic, byte(qos), retained, msg.Payload)
	//token.Wait()
	//return token.Error()
}

func (m *MQTT) Subscribe(topic string, qos QOS) (<-chan Message, error) {
	/*
		c := make(chan Message)
		token := m.client.Subscribe(topic, byte(qos), func(client mqtt.Client, msg mqtt.Message) {
			c <- Message{msg.Topic(), msg.Payload()}
		})
		token.Wait()
		return c, token.Error()
	*/
	p := parseTopic(topic)
	if p == nil {
		return nil, errors.New("invalid topic")
	}
	c := make(chan Message)
	m.mu.Lock()
	defer m.mu.Unlock()
	m.subscribers = append(m.subscribers, &subscription{topic: p, channel: c})
	return c, nil

}

func (m *MQTT) Unsubscribe(topic string) error {
	/*
		token := m.client.Unsubscribe(topic)
		token.Wait()
		return token.Error()
	*/
	p := parseTopic(topic)
	if p == nil {
		return errors.New("invalid topic")
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	for i := range m.subscribers {
		if m.subscribers[i].topic.isEqual(p) {
			m.subscribers[i].Close()
			copy(m.subscribers[i:], m.subscribers[i+1:])
			m.subscribers = m.subscribers[:len(m.subscribers)-1]
			return nil
		}
	}
	return errors.New("subscription not found")
}

func (m *MQTT) Get(topic string, qos QOS) ([]Message, error) {
	// TODO(maruel): It looks it needs to do a quick Subscribe + poll every
	// messages until one with !msg.Retained() or a timeout then Unsubscribe.
	return nil, errors.New("implement me")
}

var _ Bus = &MQTT{}
<|endoftext|>


<commit_before>package protocol

import (
	"bufio"
	"crypto/md5"
	"errors"
	"fmt"
	"github.com/s-rah/onionscan/config"
	"github.com/s-rah/onionscan/report"
	"github.com/s-rah/onionscan/utils"
	"golang.org/x/crypto/ssh"
	"net"
)

type SSHProtocolScanner struct {
}

func (sps *SSHProtocolScanner) ScanProtocol(hiddenService string, osc *config.OnionScanConfig, report *report.OnionScanReport) {
	// SSH
	osc.LogInfo(fmt.Sprintf("Checking %s ssh(22)\n", hiddenService))
	conn, err := utils.GetNetworkConnection(hiddenService, 22, osc.TorProxyAddress, osc.Timeout)
	if err != nil {
		osc.LogInfo("Failed to connect to service on port 22\n")
		report.SSHDetected = false
		if conn != nil {
			conn.Close()
		}
	} else {
		// TODO SSH Checking
		report.SSHDetected = true

		config := &ssh.ClientConfig{
			HostKeyCallback: func(hostname string, addr net.Addr, key ssh.PublicKey) error {
				h := md5.New()
				h.Write(key.Marshal())

				fBytes := h.Sum(nil)
				fingerprint := string("")
				for i := 0; i < len(fBytes); i++ {
					if i+1 != len(fBytes) {
						fingerprint = fmt.Sprintf("%s%0.2x:", fingerprint, fBytes[i])
					} else {
						fingerprint = fmt.Sprintf("%s%0.2x", fingerprint, fBytes[i])
					}
				}
				report.SSHKey = fingerprint
				osc.LogInfo(fmt.Sprintf("Found SSH Key %s\n", fingerprint))
				// We don't want to continue
				return errors.New("error")
			},
		}
		ssh.NewClientConn(conn, hiddenService+":22", config)
		if conn != nil {
			conn.Close()
		}
		conn, err = utils.GetNetworkConnection(hiddenService, 22, osc.TorProxyAddress, osc.Timeout)
		if err == nil {
			reader := bufio.NewReader(conn)
			banner, err := reader.ReadString('\n')
			if err == nil {
				report.SSHBanner = banner
				osc.LogInfo(fmt.Sprintf("Found SSH Banner: %s", banner))
			}
		}
		if conn != nil {
			conn.Close()
		}
	}
	osc.LogInfo(fmt.Sprintf("Checking %s ssh(2222)\n", hiddenService))
	conn, err := utils.GetNetworkConnection(hiddenService, 2222, osc.TorProxyAddress, osc.Timeout)
	if err != nil {
		osc.LogInfo("Failed to connect to service on port 2222\n")
		report.SSHDetected = false
		if conn != nil {
			conn.Close()
		}
	} else {
		// TODO SSH Checking
		report.SSHDetected = true

		config := &ssh.ClientConfig{
			HostKeyCallback: func(hostname string, addr net.Addr, key ssh.PublicKey) error {
				h := md5.New()
				h.Write(key.Marshal())

				fBytes := h.Sum(nil)
				fingerprint := string("")
				for i := 0; i < len(fBytes); i++ {
					if i+1 != len(fBytes) {
						fingerprint = fmt.Sprintf("%s%0.2x:", fingerprint, fBytes[i])
					} else {
						fingerprint = fmt.Sprintf("%s%0.2x", fingerprint, fBytes[i])
					}
				}
				report.SSHKey = fingerprint
				osc.LogInfo(fmt.Sprintf("Found SSH Key %s\n", fingerprint))
				// We don't want to continue
				return errors.New("error")
			},
		}
		ssh.NewClientConn(conn, hiddenService+":2222", config)
		if conn != nil {
			conn.Close()
		}
		conn, err = utils.GetNetworkConnection(hiddenService, 2222, osc.TorProxyAddress, osc.Timeout)
		if err == nil {
			reader := bufio.NewReader(conn)
			banner, err := reader.ReadString('\n')
			if err == nil {
				report.SSHBanner = banner
				osc.LogInfo(fmt.Sprintf("Found SSH Banner: %s", banner))
			}
		}
		if conn != nil {
			conn.Close()
		}
	}
}
<commit_msg>Secure Shell<commit_after>package protocol

import (
	"bufio"
	"crypto/md5"
	"errors"
	"fmt"
	"github.com/s-rah/onionscan/config"
	"github.com/s-rah/onionscan/report"
	"github.com/s-rah/onionscan/utils"
	"golang.org/x/crypto/ssh"
	"net"
)

type SSHProtocolScanner struct {
}

func (sps *SSHProtocolScanner) ScanProtocol(hiddenService string, osc *config.OnionScanConfig, report *report.OnionScanReport) {
	// SSH
	osc.LogInfo(fmt.Sprintf("Checking %s ssh(22)\n", hiddenService))
	conn, err := utils.GetNetworkConnection(hiddenService, 22, osc.TorProxyAddress, osc.Timeout)
	if err != nil {
		osc.LogInfo("Failed to connect to service on port 22\n")
		report.SSHDetected = false
		if conn != nil {
			conn.Close()
		}
	} else {
		// TODO SSH Checking
		report.SSHDetected = true

		config := &ssh.ClientConfig{
			HostKeyCallback: func(hostname string, addr net.Addr, key ssh.PublicKey) error {
				h := md5.New()
				h.Write(key.Marshal())

				fBytes := h.Sum(nil)
				fingerprint := string("")
				for i := 0; i < len(fBytes); i++ {
					if i+1 != len(fBytes) {
						fingerprint = fmt.Sprintf("%s%0.2x:", fingerprint, fBytes[i])
					} else {
						fingerprint = fmt.Sprintf("%s%0.2x", fingerprint, fBytes[i])
					}
				}
				report.SSHKey = fingerprint
				osc.LogInfo(fmt.Sprintf("Found SSH Key %s\n", fingerprint))
				// We don't want to continue
				return errors.New("error")
			},
		}
		ssh.NewClientConn(conn, hiddenService+":22", config)
		if conn != nil {
			conn.Close()
		}
		conn, err = utils.GetNetworkConnection(hiddenService, 22, osc.TorProxyAddress, osc.Timeout)
		if err == nil {
			reader := bufio.NewReader(conn)
			banner, err := reader.ReadString('\n')
			if err == nil {
				report.SSHBanner = banner
				osc.LogInfo(fmt.Sprintf("Found SSH Banner: %s", banner))
			}
		}
		if conn != nil {
			conn.Close()
		}
	}
}
<|endoftext|>


<commit_before>package vm

import (
	"fmt"
	"testing"

	"github.com/bytom/errors"
	"github.com/bytom/math/checked"
	"github.com/bytom/testutil"
)

func TestParseOp(t *testing.T) {
	cases := []struct {
		prog    []byte
		pc      uint32
		want    Instruction
		wantErr error
	}{{
		prog: []byte{byte(OP_ADD)},
		want: Instruction{Op: OP_ADD, Len: 1},
	}, {
		prog: []byte{byte(OP_16)},
		want: Instruction{Op: OP_16, Data: []byte{16}, Len: 1},
	}, {
		prog: []byte{byte(OP_DATA_5), 1, 1, 1, 1, 1},
		want: Instruction{Op: OP_DATA_5, Data: []byte{1, 1, 1, 1, 1}, Len: 6},
	}, {
		prog: []byte{byte(OP_DATA_5), 1, 1, 1, 1, 1, 255},
		want: Instruction{Op: OP_DATA_5, Data: []byte{1, 1, 1, 1, 1}, Len: 6},
	}, {
		prog: []byte{byte(OP_PUSHDATA1), 1, 1},
		want: Instruction{Op: OP_PUSHDATA1, Data: []byte{1}, Len: 3},
	}, {
		prog: []byte{byte(OP_PUSHDATA1), 1, 1, 255},
		want: Instruction{Op: OP_PUSHDATA1, Data: []byte{1}, Len: 3},
	}, {
		prog: []byte{byte(OP_PUSHDATA2), 1, 0, 1},
		want: Instruction{Op: OP_PUSHDATA2, Data: []byte{1}, Len: 4},
	}, {
		prog: []byte{byte(OP_PUSHDATA2), 1, 0, 1, 255},
		want: Instruction{Op: OP_PUSHDATA2, Data: []byte{1}, Len: 4},
	}, {
		prog: []byte{byte(OP_PUSHDATA4), 1, 0, 0, 0, 1},
		want: Instruction{Op: OP_PUSHDATA4, Data: []byte{1}, Len: 6},
	}, {
		prog: []byte{byte(OP_PUSHDATA4), 1, 0, 0, 0, 1, 255},
		want: Instruction{Op: OP_PUSHDATA4, Data: []byte{1}, Len: 6},
	}, {
		prog:    []byte{},
		wantErr: ErrShortProgram,
	}, {
		prog:    []byte{byte(OP_0)},
		pc:      1,
		wantErr: ErrShortProgram,
	}, {
		prog:    []byte{byte(OP_DATA_1)},
		wantErr: ErrShortProgram,
	}, {
		prog:    []byte{byte(OP_PUSHDATA1)},
		wantErr: ErrShortProgram,
	}, {
		prog:    []byte{byte(OP_PUSHDATA1), 1},
		wantErr: ErrShortProgram,
	}, {
		prog:    []byte{byte(OP_PUSHDATA2)},
		wantErr: ErrShortProgram,
	}, {
		prog:    []byte{byte(OP_PUSHDATA2), 1, 0},
		wantErr: ErrShortProgram,
	}, {
		prog:    []byte{byte(OP_PUSHDATA4)},
		wantErr: ErrShortProgram,
	}, {
		prog:    []byte{byte(OP_PUSHDATA4), 1, 0, 0, 0},
		wantErr: ErrShortProgram,
	}, {
		pc:      71,
		prog:    []byte{0x6d, 0x6b, 0xaa, 0x20, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x20, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x0, 0x0, 0x4e, 0xff, 0xff, 0xff, 0xff, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30},
		wantErr: checked.ErrOverflow,
	}}

	for _, c := range cases {
		t.Run(fmt.Sprintf("%d: %x", c.pc, c.prog), func(t *testing.T) {
			got, gotErr := ParseOp(c.prog, c.pc)

			if errors.Root(gotErr) != c.wantErr {
				t.Errorf("ParseOp(%x, %d) error = %v want %v", c.prog, c.pc, gotErr, c.wantErr)
			}

			if c.wantErr != nil {
				return
			}

			if !testutil.DeepEqual(got, c.want) {
				t.Errorf("ParseOp(%x, %d) = %+v want %+v", c.prog, c.pc, got, c.want)
			}
		})
	}
}

func TestParseProgram(t *testing.T) {
	cases := []struct {
		prog    []byte
		want    []Instruction
		wantErr error
	}{
		{
			prog: []byte{byte(OP_2), byte(OP_3), byte(OP_ADD), byte(OP_5), byte(OP_NUMEQUAL)},
			want: []Instruction{
				{Op: OP_2, Data: []byte{0x02}, Len: 1},
				{Op: OP_3, Data: []byte{0x03}, Len: 1},
				{Op: OP_ADD, Len: 1},
				{Op: OP_5, Data: []byte{0x05}, Len: 1},
				{Op: OP_NUMEQUAL, Len: 1},
			},
		},
		{
			prog: []byte{255},
			want: []Instruction{
				{Op: 255, Len: 1},
			},
		},
	}

	for _, c := range cases {
		got, gotErr := ParseProgram(c.prog)

		if errors.Root(gotErr) != c.wantErr {
			t.Errorf("ParseProgram(%x) error = %v want %v", c.prog, gotErr, c.wantErr)
		}

		if c.wantErr != nil {
			continue
		}

		if !testutil.DeepEqual(got, c.want) {
			t.Errorf("ParseProgram(%x) = %+v want %+v", c.prog, got, c.want)
		}
	}
}

func TestIsPushData(t *testing.T) {
	cases := []struct {
		want    Instruction
		wantErr error
	}{
		{
			want: Instruction{Op: OP_16, Data: []byte{16}, Len: 1},
		},
		{
			want: Instruction{Op: OP_DATA_32, Data: []byte{16}, Len: 1},
		},
		{
			want: Instruction{Op: OP_FALSE, Data: []byte{}, Len: 1},
		},
		{
			want: Instruction{Op: OP_TRUE, Data: []byte{1}, Len: 1},
		},
		{
			want:    Instruction{Op: OP_JUMP, Data: []byte{0x00000000}, Len: 1},
			wantErr: ErrShortProgram,
		},
		{
			want:    Instruction{Op: OP_ADD, Data: []byte{0x12, 0x56}, Len: 2},
			wantErr: ErrShortProgram,
		},
	}

	for _, c := range cases {
		if c.want.IsPushdata() {
			t.Logf("check success")
		} else if c.wantErr != nil {
			t.Logf("check err success")
		} else {
			t.Errorf("check false: %v -- %v", ops[OP_1].fn, ops[c.want.Op].fn)
		}
	}
}
<commit_msg>fix unit test for ispushdata (#1132)<commit_after>package vm

import (
	"fmt"
	"reflect"
	"testing"

	"github.com/bytom/errors"
	"github.com/bytom/math/checked"
	"github.com/bytom/testutil"
)

func TestParseOp(t *testing.T) {
	cases := []struct {
		prog    []byte
		pc      uint32
		want    Instruction
		wantErr error
	}{{
		prog: []byte{byte(OP_ADD)},
		want: Instruction{Op: OP_ADD, Len: 1},
	}, {
		prog: []byte{byte(OP_16)},
		want: Instruction{Op: OP_16, Data: []byte{16}, Len: 1},
	}, {
		prog: []byte{byte(OP_DATA_5), 1, 1, 1, 1, 1},
		want: Instruction{Op: OP_DATA_5, Data: []byte{1, 1, 1, 1, 1}, Len: 6},
	}, {
		prog: []byte{byte(OP_DATA_5), 1, 1, 1, 1, 1, 255},
		want: Instruction{Op: OP_DATA_5, Data: []byte{1, 1, 1, 1, 1}, Len: 6},
	}, {
		prog: []byte{byte(OP_PUSHDATA1), 1, 1},
		want: Instruction{Op: OP_PUSHDATA1, Data: []byte{1}, Len: 3},
	}, {
		prog: []byte{byte(OP_PUSHDATA1), 1, 1, 255},
		want: Instruction{Op: OP_PUSHDATA1, Data: []byte{1}, Len: 3},
	}, {
		prog: []byte{byte(OP_PUSHDATA2), 1, 0, 1},
		want: Instruction{Op: OP_PUSHDATA2, Data: []byte{1}, Len: 4},
	}, {
		prog: []byte{byte(OP_PUSHDATA2), 1, 0, 1, 255},
		want: Instruction{Op: OP_PUSHDATA2, Data: []byte{1}, Len: 4},
	}, {
		prog: []byte{byte(OP_PUSHDATA4), 1, 0, 0, 0, 1},
		want: Instruction{Op: OP_PUSHDATA4, Data: []byte{1}, Len: 6},
	}, {
		prog: []byte{byte(OP_PUSHDATA4), 1, 0, 0, 0, 1, 255},
		want: Instruction{Op: OP_PUSHDATA4, Data: []byte{1}, Len: 6},
	}, {
		prog:    []byte{},
		wantErr: ErrShortProgram,
	}, {
		prog:    []byte{byte(OP_0)},
		pc:      1,
		wantErr: ErrShortProgram,
	}, {
		prog:    []byte{byte(OP_DATA_1)},
		wantErr: ErrShortProgram,
	}, {
		prog:    []byte{byte(OP_PUSHDATA1)},
		wantErr: ErrShortProgram,
	}, {
		prog:    []byte{byte(OP_PUSHDATA1), 1},
		wantErr: ErrShortProgram,
	}, {
		prog:    []byte{byte(OP_PUSHDATA2)},
		wantErr: ErrShortProgram,
	}, {
		prog:    []byte{byte(OP_PUSHDATA2), 1, 0},
		wantErr: ErrShortProgram,
	}, {
		prog:    []byte{byte(OP_PUSHDATA4)},
		wantErr: ErrShortProgram,
	}, {
		prog:    []byte{byte(OP_PUSHDATA4), 1, 0, 0, 0},
		wantErr: ErrShortProgram,
	}, {
		pc:      71,
		prog:    []byte{0x6d, 0x6b, 0xaa, 0x20, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x20, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x0, 0x0, 0x4e, 0xff, 0xff, 0xff, 0xff, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30},
		wantErr: checked.ErrOverflow,
	}}

	for _, c := range cases {
		t.Run(fmt.Sprintf("%d: %x", c.pc, c.prog), func(t *testing.T) {
			got, gotErr := ParseOp(c.prog, c.pc)

			if errors.Root(gotErr) != c.wantErr {
				t.Errorf("ParseOp(%x, %d) error = %v want %v", c.prog, c.pc, gotErr, c.wantErr)
			}

			if c.wantErr != nil {
				return
			}

			if !testutil.DeepEqual(got, c.want) {
				t.Errorf("ParseOp(%x, %d) = %+v want %+v", c.prog, c.pc, got, c.want)
			}
		})
	}
}

func TestParseProgram(t *testing.T) {
	cases := []struct {
		prog    []byte
		want    []Instruction
		wantErr error
	}{
		{
			prog: []byte{byte(OP_2), byte(OP_3), byte(OP_ADD), byte(OP_5), byte(OP_NUMEQUAL)},
			want: []Instruction{
				{Op: OP_2, Data: []byte{0x02}, Len: 1},
				{Op: OP_3, Data: []byte{0x03}, Len: 1},
				{Op: OP_ADD, Len: 1},
				{Op: OP_5, Data: []byte{0x05}, Len: 1},
				{Op: OP_NUMEQUAL, Len: 1},
			},
		},
		{
			prog: []byte{255},
			want: []Instruction{
				{Op: 255, Len: 1},
			},
		},
	}

	for _, c := range cases {
		got, gotErr := ParseProgram(c.prog)

		if errors.Root(gotErr) != c.wantErr {
			t.Errorf("ParseProgram(%x) error = %v want %v", c.prog, gotErr, c.wantErr)
		}

		if c.wantErr != nil {
			continue
		}

		if !testutil.DeepEqual(got, c.want) {
			t.Errorf("ParseProgram(%x) = %+v want %+v", c.prog, got, c.want)
		}
	}
}

func TestIsPushData(t *testing.T) {
	cases := []struct {
		want    Instruction
		wantErr error
	}{
		{
			want: Instruction{Op: OP_16, Data: []byte{16}, Len: 1},
		},
		{
			want: Instruction{Op: OP_DATA_32, Data: []byte{16}, Len: 1},
		},
		{
			want: Instruction{Op: OP_FALSE, Data: []byte{}, Len: 1},
		},
		{
			want: Instruction{Op: OP_TRUE, Data: []byte{1}, Len: 1},
		},
		{
			want:    Instruction{Op: OP_JUMP, Data: []byte{0x00000000}, Len: 1},
			wantErr: ErrShortProgram,
		},
		{
			want:    Instruction{Op: OP_ADD, Data: []byte{0x12, 0x56}, Len: 2},
			wantErr: ErrShortProgram,
		},
	}

	for _, c := range cases {
		if c.want.IsPushdata() {
			t.Logf("check success")
		} else if c.wantErr != nil {
			t.Logf("check err success")
		} else {
			t.Errorf("check false: %v -- %v", reflect.ValueOf(ops[OP_1].fn), reflect.ValueOf(ops[c.want.Op].fn))
		}
	}
}
<|endoftext|>