<commit_before>package shortener

import (
	"fmt"

	"github.com/davidlu1997/fast-shortener/config"
	"github.com/davidlu1997/fast-shortener/model"
	cache "github.com/patrickmn/go-cache"
)

type CacheShortener struct {
	cache  *cache.Cache
	config *config.Configuration
	Shortener
}

func InitCacheShortener(config *config.Configuration) *CacheShortener {
	cache := cache.New(config.Cache.DefaultDuration, config.Cache.DefaultPurge)
	if cache == nil {
		return nil
	}

	return &CacheShortener{
		cache:  cache,
		config: config,
	}
}

func (c *CacheShortener) Get(key string) *model.Link {
	if link, found := c.cache.Get(key); found {
		return link.(*model.Link)
	}
	return nil
}

func (c *CacheShortener) Put(link *model.Link) error {
	if link == nil {
		return fmt.Errorf("null link")
	}

	link.PreSave()

	if !c.canPut(link) {
		return fmt.Errorf("cannot put link")
	}
	return c.cache.Add(link.Key, link, *link.Duration)
}

func (c *CacheShortener) canPut(link *model.Link) bool {
	return link.IsValid(c.config) && c.cache.ItemCount() < c.config.Cache.MaxSize
}
<commit_msg>Add AfterGet<commit_after>package shortener

import (
	"fmt"

	"github.com/davidlu1997/fast-shortener/config"
	"github.com/davidlu1997/fast-shortener/model"
	cache "github.com/patrickmn/go-cache"
)

type CacheShortener struct {
	cache  *cache.Cache
	config *config.Configuration
	Shortener
}

func InitCacheShortener(config *config.Configuration) *CacheShortener {
	cache := cache.New(config.Cache.DefaultDuration, config.Cache.DefaultPurge)
	if cache == nil {
		return nil
	}

	return &CacheShortener{
		cache:  cache,
		config: config,
	}
}

func (c *CacheShortener) Get(key string) *model.Link {
	if l, found := c.cache.Get(key); found {
		link := l.(*model.Link)
		link.AfterGet()
		return link
	}
	return nil
}

func (c *CacheShortener) Put(link *model.Link) error {
	if link == nil {
		return fmt.Errorf("null link")
	}

	link.PreSave()

	if !c.canPut(link) {
		return fmt.Errorf("cannot put link")
	}
	return c.cache.Add(link.Key, link, *link.Duration)
}

func (c *CacheShortener) canPut(link *model.Link) bool {
	return link.IsValid(c.config) && c.cache.ItemCount() < c.config.Cache.MaxSize
}
<|endoftext|>


<commit_before>// @llr REQ-0-DDLN-SWL-001
package main

import (
	"fmt"
	"html/template"
	"regexp"
	"strings"
	"unicode"

	"github.com/daedaleanai/reqtraq/config"
)

var (
	// REQ, project number, project abbreviation, req type, req number
	// For example: REQ-0-DDLN-SWH-004
	reReqIdStr   = `REQ-(\d+)-(\w+)-(SYS|SWH|SWL|HWH|HWL)-(\d+)`
	ReReqID      = regexp.MustCompile(reReqIdStr)
	ReReqDeleted = regexp.MustCompile(reReqIdStr + ` DELETED`)
	reReqIDBad   = regexp.MustCompile(`(?i)REQ(-(\w+))+`)
	reReqKWD     = regexp.MustCompile(`(?i)(- )?(rationale|parent|parents|safety impact|verification|urgent|important|mode|provenance):`)
)

// ParseReq finds the first REQ-XXX tag and the reserved words and distills a Req from it.
//
// ParseReq parses according to the 'soft' format defined in the SRS:
//            REQ-ID (text)
//            [Rationale:....]
//            [Parent[s]: REQ-ID[, REQ-ID...]]
//            [Safety Impact:...]
//            [Verification:...]
//            [Urgent:...]
//            [Important:...]
//            [Mode:...]
//            [Provenance:...]
//
// ParseReq does NOT validate the values or check if the mandatory attributes are set; use
// the Req.Check() method for that.
//
// Since the parsing is rather 'soft', ParseReq returns verbose errors indicating problems in
// a helpful way, meaning they at least provide enough context for the user to find the text.
// TODO(aroetter): pass in a bool pandoc here
func ParseReq(txt string) (*Req, error) {
	lyx := strings.HasPrefix(txt, "\n")
	head := txt
	if len(head) > 40 {
		head = head[:40]
	}
	defid := ReReqID.FindStringSubmatchIndex(txt)
	if len(defid) == 0 {
		if reReqIDBad.MatchString(head) {
			return nil, fmt.Errorf("malformed requirement: found only malformed ID: %q (doesn't match %q)", head, ReReqID)
		}
		return nil, fmt.Errorf("malformed requirement: missing ID in first 40 characters: %q", head)
	}

	if lyx {
		if defid[0] > 20 {
			return nil, fmt.Errorf("malformed requirement: too much heading garbage before ID: %q", head)
		}
	} else {
		if defid[0] > 0 {
			return nil, fmt.Errorf("malformed requirement: ID must be at the start of the title: %q", head)
		}
	}

	r := &Req{
		ID:         txt[defid[0]:defid[1]],
		Attributes: map[string]string{},
	}

	// chop defining ID and any punctuation
	txt = strings.TrimLeftFunc(txt[defid[1]:], unicode.IsPunct)
	txt = strings.TrimLeftFunc(txt, unicode.IsSpace)

	var attributesStart int
	kwdMatches := reReqKWD.FindAllStringSubmatchIndex(txt, -1)
	if len(kwdMatches) == 0 {
		return nil, fmt.Errorf("requirement %s contains no attributes", r.ID)
	}
	if lyx {
		attributesStart = kwdMatches[0][0]
	} else {
		attributesStart = strings.Index(txt, "\n###### Attributes:\n")
	}
	for i, v := range kwdMatches {
		key := strings.ToUpper(txt[v[4]:v[5]])
		if key == "PARENT" { // make our lives easier, accept both, output only PARENTS
			key = "PARENTS"
		}
		e := len(txt)
		if i < len(kwdMatches)-1 {
			e = kwdMatches[i+1][0]
		}
		if _, ok := r.Attributes[key]; ok {
			return nil, fmt.Errorf("requirement %s contains duplicate attribute: %q", r.ID, key)
		}
		r.Attributes[key] = strings.TrimSpace(txt[v[1]:e])
	}

	// TEXT is anything up to the first keyword we found
	txt = txt[:attributesStart]

	// PARENTS must be punctuation/space separated list of parseable req-ids.
	parents := r.Attributes["PARENTS"]
	parmatch := ReReqID.FindAllStringSubmatchIndex(parents, -1)
	for i, ids := range parmatch {
		val := parents[ids[0]:ids[1]]
		r.ParentIds = append(r.ParentIds, val)
		if i > 0 {
			sep := parents[parmatch[i-1][1]:ids[0]]
			if strings.TrimFunc(sep, func(r rune) bool { return unicode.IsSpace(r) || unicode.IsPunct(r) }) != "" {
				return nil, fmt.Errorf("requirement %s parents: unparseable as list of requirement ids: %q in %q", r.ID, sep, parents)
			}
		}
	}

	level, ok := config.ReqTypeToReqLevel[r.ReqType()]
	if !ok {
		return nil, fmt.Errorf("Invalid request type: %q", r.ReqType())
	}
	r.Level = level

	parts := strings.SplitN(strings.TrimSpace(txt), "\n", 2)
	r.Title = parts[0]
	// TODO(aroetter): Run this body text through pandoc
	r.Body = template.HTML(parts[1])
	return r, nil
}
<commit_msg>remove comment cruft<commit_after>// @llr REQ-0-DDLN-SWL-001
package main

import (
	"fmt"
	"html/template"
	"regexp"
	"strings"
	"unicode"

	"github.com/daedaleanai/reqtraq/config"
)

var (
	// REQ, project number, project abbreviation, req type, req number
	// For example: REQ-0-DDLN-SWH-004
	reReqIdStr   = `REQ-(\d+)-(\w+)-(SYS|SWH|SWL|HWH|HWL)-(\d+)`
	ReReqID      = regexp.MustCompile(reReqIdStr)
	ReReqDeleted = regexp.MustCompile(reReqIdStr + ` DELETED`)
	reReqIDBad   = regexp.MustCompile(`(?i)REQ(-(\w+))+`)
	reReqKWD     = regexp.MustCompile(`(?i)(- )?(rationale|parent|parents|safety impact|verification|urgent|important|mode|provenance):`)
)

// ParseReq finds the first REQ-XXX tag and the reserved words and distills a Req from it.
//
// ParseReq parses according to the 'soft' format defined in the SRS:
//            REQ-ID (text)
//            [Rationale:....]
//            [Parent[s]: REQ-ID[, REQ-ID...]]
//            [Safety Impact:...]
//            [Verification:...]
//            [Urgent:...]
//            [Important:...]
//            [Mode:...]
//            [Provenance:...]
//
// ParseReq does NOT validate the values or check if the mandatory attributes are set; use
// the Req.Check() method for that.
//
// Since the parsing is rather 'soft', ParseReq returns verbose errors indicating problems in
// a helpful way, meaning they at least provide enough context for the user to find the text.
func ParseReq(txt string) (*Req, error) {
	lyx := strings.HasPrefix(txt, "\n")
	head := txt
	if len(head) > 40 {
		head = head[:40]
	}
	defid := ReReqID.FindStringSubmatchIndex(txt)
	if len(defid) == 0 {
		if reReqIDBad.MatchString(head) {
			return nil, fmt.Errorf("malformed requirement: found only malformed ID: %q (doesn't match %q)", head, ReReqID)
		}
		return nil, fmt.Errorf("malformed requirement: missing ID in first 40 characters: %q", head)
	}

	if lyx {
		if defid[0] > 20 {
			return nil, fmt.Errorf("malformed requirement: too much heading garbage before ID: %q", head)
		}
	} else {
		if defid[0] > 0 {
			return nil, fmt.Errorf("malformed requirement: ID must be at the start of the title: %q", head)
		}
	}

	r := &Req{
		ID:         txt[defid[0]:defid[1]],
		Attributes: map[string]string{},
	}

	// chop defining ID and any punctuation
	txt = strings.TrimLeftFunc(txt[defid[1]:], unicode.IsPunct)
	txt = strings.TrimLeftFunc(txt, unicode.IsSpace)

	var attributesStart int
	kwdMatches := reReqKWD.FindAllStringSubmatchIndex(txt, -1)
	if len(kwdMatches) == 0 {
		return nil, fmt.Errorf("requirement %s contains no attributes", r.ID)
	}
	if lyx {
		attributesStart = kwdMatches[0][0]
	} else {
		attributesStart = strings.Index(txt, "\n###### Attributes:\n")
	}
	for i, v := range kwdMatches {
		key := strings.ToUpper(txt[v[4]:v[5]])
		if key == "PARENT" { // make our lives easier, accept both, output only PARENTS
			key = "PARENTS"
		}
		e := len(txt)
		if i < len(kwdMatches)-1 {
			e = kwdMatches[i+1][0]
		}
		if _, ok := r.Attributes[key]; ok {
			return nil, fmt.Errorf("requirement %s contains duplicate attribute: %q", r.ID, key)
		}
		r.Attributes[key] = strings.TrimSpace(txt[v[1]:e])
	}

	// TEXT is anything up to the first keyword we found
	txt = txt[:attributesStart]

	// PARENTS must be punctuation/space separated list of parseable req-ids.
	parents := r.Attributes["PARENTS"]
	parmatch := ReReqID.FindAllStringSubmatchIndex(parents, -1)
	for i, ids := range parmatch {
		val := parents[ids[0]:ids[1]]
		r.ParentIds = append(r.ParentIds, val)
		if i > 0 {
			sep := parents[parmatch[i-1][1]:ids[0]]
			if strings.TrimFunc(sep, func(r rune) bool { return unicode.IsSpace(r) || unicode.IsPunct(r) }) != "" {
				return nil, fmt.Errorf("requirement %s parents: unparseable as list of requirement ids: %q in %q", r.ID, sep, parents)
			}
		}
	}

	level, ok := config.ReqTypeToReqLevel[r.ReqType()]
	if !ok {
		return nil, fmt.Errorf("Invalid request type: %q", r.ReqType())
	}
	r.Level = level

	parts := strings.SplitN(strings.TrimSpace(txt), "\n", 2)
	r.Title = parts[0]
	// TODO(aroetter): Run this body text through pandoc
	r.Body = template.HTML(parts[1])
	return r, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2018-2019 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package txscript

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"testing"

	"github.com/btcsuite/btcd/wire"
)

var (
	// manyInputsBenchTx is a transaction that contains a lot of inputs which is
	// useful for benchmarking signature hash calculation.
	manyInputsBenchTx wire.MsgTx

	// A mock previous output script to use in the signing benchmark.
	prevOutScript = hexToBytes("a914f5916158e3e2c4551c1796708db8367207ed13bb87")
)

func init() {
	// tx 620f57c92cf05a7f7e7f7d28255d5f7089437bc48e34dcfebf7751d08b7fb8f5
	txHex, err := ioutil.ReadFile("data/many_inputs_tx.hex")
	if err != nil {
		panic(fmt.Sprintf("unable to read benchmark tx file: %v", err))
	}

	txBytes := hexToBytes(string(txHex))
	err = manyInputsBenchTx.Deserialize(bytes.NewReader(txBytes))
	if err != nil {
		panic(err)
	}
}

// BenchmarkCalcSigHash benchmarks how long it takes to calculate the signature
// hashes for all inputs of a transaction with many inputs.
func BenchmarkCalcSigHash(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		for j := 0; j < len(manyInputsBenchTx.TxIn); j++ {
			_, err := CalcSignatureHash(prevOutScript, SigHashAll,
				&manyInputsBenchTx, j)
			if err != nil {
				b.Fatalf("failed to calc signature hash: %v", err)
			}
		}
	}
}

// BenchmarkCalcWitnessSigHash benchmarks how long it takes to calculate the
// witness signature hashes for all inputs of a transaction with many inputs.
func BenchmarkCalcWitnessSigHash(b *testing.B) {
	sigHashes := NewTxSigHashes(&manyInputsBenchTx)

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		for j := 0; j < len(manyInputsBenchTx.TxIn); j++ {
			_, err := CalcWitnessSigHash(
				prevOutScript, sigHashes, SigHashAll,
				&manyInputsBenchTx, j, 5,
			)
			if err != nil {
				b.Fatalf("failed to calc signature hash: %v", err)
			}
		}
	}
}

// genComplexScript returns a script comprised of half as many opcodes as the
// maximum allowed followed by as many max size data pushes fit without
// exceeding the max allowed script size.
func genComplexScript() ([]byte, error) {
	var scriptLen int
	builder := NewScriptBuilder()
	for i := 0; i < MaxOpsPerScript/2; i++ {
		builder.AddOp(OP_TRUE)
		scriptLen++
	}
	maxData := bytes.Repeat([]byte{0x02}, MaxScriptElementSize)
	for i := 0; i < (MaxScriptSize-scriptLen)/(MaxScriptElementSize+3); i++ {
		builder.AddData(maxData)
	}
	return builder.Script()
}

// BenchmarkScriptParsing benchmarks how long it takes to parse a very large
// script.
func BenchmarkScriptParsing(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	const scriptVersion = 0
	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		tokenizer := MakeScriptTokenizer(scriptVersion, script)
		for tokenizer.Next() {
			_ = tokenizer.Opcode()
			_ = tokenizer.Data()
			_ = tokenizer.ByteIndex()
		}
		if err := tokenizer.Err(); err != nil {
			b.Fatalf("failed to parse script: %v", err)
		}
	}
}

// BenchmarkDisasmString benchmarks how long it takes to disassemble a very
// large script.
func BenchmarkDisasmString(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		_, err := DisasmString(script)
		if err != nil {
			b.Fatalf("failed to disasm script: %v", err)
		}
	}
}

// BenchmarkIsPubKeyScript benchmarks how long it takes to analyze a very large
// script to determine if it is a standard pay-to-pubkey script.
func BenchmarkIsPubKeyScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		_ = IsPayToPubKey(script)
	}
}

// BenchmarkIsPubKeyHashScript benchmarks how long it takes to analyze a very
// large script to determine if it is a standard pay-to-pubkey-hash script.
func BenchmarkIsPubKeyHashScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		_ = IsPayToPubKeyHash(script)
	}
}

// BenchmarkIsPayToScriptHash benchmarks how long it takes IsPayToScriptHash to
// analyze a very large script.
func BenchmarkIsPayToScriptHash(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		_ = IsPayToScriptHash(script)
	}
}

// BenchmarkIsMultisigScriptLarge benchmarks how long it takes IsMultisigScript
// to analyze a very large script.
func BenchmarkIsMultisigScriptLarge(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		isMultisig, err := IsMultisigScript(script)
		if err != nil {
			b.Fatalf("unexpected err: %v", err)
		}
		if isMultisig {
			b.Fatalf("script should NOT be reported as mutisig script")
		}
	}
}

// BenchmarkIsMultisigScript benchmarks how long it takes IsMultisigScript to
// analyze a 1-of-2 multisig public key script.
func BenchmarkIsMultisigScript(b *testing.B) {
	multisigShortForm := "1 " +
		"DATA_33 " +
		"0x030478aaaa2be30772f1e69e581610f1840b3cf2fe7228ee0281cd599e5746f81e " +
		"DATA_33 " +
		"0x0284f4d078b236a9ff91661f8ffbe012737cd3507566f30fd97d25f2b23539f3cd " +
		"2 CHECKMULTISIG"
	pkScript := mustParseShortForm(multisigShortForm)

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		isMultisig, err := IsMultisigScript(pkScript)
		if err != nil {
			b.Fatalf("unexpected err: %v", err)
		}
		if !isMultisig {
			b.Fatalf("script should be reported as a mutisig script")
		}
	}
}

// BenchmarkIsMultisigSigScript benchmarks how long it takes IsMultisigSigScript
// to analyze a very large script.
func BenchmarkIsMultisigSigScriptLarge(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		if IsMultisigSigScript(script) {
			b.Fatalf("script should NOT be reported as mutisig sig script")
		}
	}
}

// BenchmarkIsMultisigSigScript benchmarks how long it takes IsMultisigSigScript
// to analyze both a 1-of-2 multisig public key script (which should be false)
// and a signature script comprised of a pay-to-script-hash 1-of-2 multisig
// redeem script (which should be true).
func BenchmarkIsMultisigSigScript(b *testing.B) {
	multisigShortForm := "1 " +
		"DATA_33 " +
		"0x030478aaaa2be30772f1e69e581610f1840b3cf2fe7228ee0281cd599e5746f81e " +
		"DATA_33 " +
		"0x0284f4d078b236a9ff91661f8ffbe012737cd3507566f30fd97d25f2b23539f3cd " +
		"2 CHECKMULTISIG"
	pkScript := mustParseShortForm(multisigShortForm)

	sigHex := "0x304402205795c3ab6ba11331eeac757bf1fc9c34bef0c7e1a9c8bd5eebb8" +
		"82f3b79c5838022001e0ab7b4c7662e4522dc5fa479e4b4133fa88c6a53d895dc1d5" +
		"2eddc7bbcf2801 "
	sigScript := mustParseShortForm("DATA_71 " + sigHex + "DATA_71 " +
		multisigShortForm)

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		if IsMultisigSigScript(pkScript) {
			b.Fatalf("script should NOT be reported as mutisig sig script")
		}
		if !IsMultisigSigScript(sigScript) {
			b.Fatalf("script should be reported as a mutisig sig script")
		}
	}
}
<commit_msg>txscript: Add benchmark for IsPushOnlyScript.<commit_after>// Copyright (c) 2018-2019 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package txscript

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"testing"

	"github.com/btcsuite/btcd/wire"
)

var (
	// manyInputsBenchTx is a transaction that contains a lot of inputs which is
	// useful for benchmarking signature hash calculation.
	manyInputsBenchTx wire.MsgTx

	// A mock previous output script to use in the signing benchmark.
	prevOutScript = hexToBytes("a914f5916158e3e2c4551c1796708db8367207ed13bb87")
)

func init() {
	// tx 620f57c92cf05a7f7e7f7d28255d5f7089437bc48e34dcfebf7751d08b7fb8f5
	txHex, err := ioutil.ReadFile("data/many_inputs_tx.hex")
	if err != nil {
		panic(fmt.Sprintf("unable to read benchmark tx file: %v", err))
	}

	txBytes := hexToBytes(string(txHex))
	err = manyInputsBenchTx.Deserialize(bytes.NewReader(txBytes))
	if err != nil {
		panic(err)
	}
}

// BenchmarkCalcSigHash benchmarks how long it takes to calculate the signature
// hashes for all inputs of a transaction with many inputs.
func BenchmarkCalcSigHash(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		for j := 0; j < len(manyInputsBenchTx.TxIn); j++ {
			_, err := CalcSignatureHash(prevOutScript, SigHashAll,
				&manyInputsBenchTx, j)
			if err != nil {
				b.Fatalf("failed to calc signature hash: %v", err)
			}
		}
	}
}

// BenchmarkCalcWitnessSigHash benchmarks how long it takes to calculate the
// witness signature hashes for all inputs of a transaction with many inputs.
func BenchmarkCalcWitnessSigHash(b *testing.B) {
	sigHashes := NewTxSigHashes(&manyInputsBenchTx)

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		for j := 0; j < len(manyInputsBenchTx.TxIn); j++ {
			_, err := CalcWitnessSigHash(
				prevOutScript, sigHashes, SigHashAll,
				&manyInputsBenchTx, j, 5,
			)
			if err != nil {
				b.Fatalf("failed to calc signature hash: %v", err)
			}
		}
	}
}

// genComplexScript returns a script comprised of half as many opcodes as the
// maximum allowed followed by as many max size data pushes fit without
// exceeding the max allowed script size.
func genComplexScript() ([]byte, error) {
	var scriptLen int
	builder := NewScriptBuilder()
	for i := 0; i < MaxOpsPerScript/2; i++ {
		builder.AddOp(OP_TRUE)
		scriptLen++
	}
	maxData := bytes.Repeat([]byte{0x02}, MaxScriptElementSize)
	for i := 0; i < (MaxScriptSize-scriptLen)/(MaxScriptElementSize+3); i++ {
		builder.AddData(maxData)
	}
	return builder.Script()
}

// BenchmarkScriptParsing benchmarks how long it takes to parse a very large
// script.
func BenchmarkScriptParsing(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	const scriptVersion = 0
	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		tokenizer := MakeScriptTokenizer(scriptVersion, script)
		for tokenizer.Next() {
			_ = tokenizer.Opcode()
			_ = tokenizer.Data()
			_ = tokenizer.ByteIndex()
		}
		if err := tokenizer.Err(); err != nil {
			b.Fatalf("failed to parse script: %v", err)
		}
	}
}

// BenchmarkDisasmString benchmarks how long it takes to disassemble a very
// large script.
func BenchmarkDisasmString(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		_, err := DisasmString(script)
		if err != nil {
			b.Fatalf("failed to disasm script: %v", err)
		}
	}
}

// BenchmarkIsPubKeyScript benchmarks how long it takes to analyze a very large
// script to determine if it is a standard pay-to-pubkey script.
func BenchmarkIsPubKeyScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		_ = IsPayToPubKey(script)
	}
}

// BenchmarkIsPubKeyHashScript benchmarks how long it takes to analyze a very
// large script to determine if it is a standard pay-to-pubkey-hash script.
func BenchmarkIsPubKeyHashScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		_ = IsPayToPubKeyHash(script)
	}
}

// BenchmarkIsPayToScriptHash benchmarks how long it takes IsPayToScriptHash to
// analyze a very large script.
func BenchmarkIsPayToScriptHash(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		_ = IsPayToScriptHash(script)
	}
}

// BenchmarkIsMultisigScriptLarge benchmarks how long it takes IsMultisigScript
// to analyze a very large script.
func BenchmarkIsMultisigScriptLarge(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		isMultisig, err := IsMultisigScript(script)
		if err != nil {
			b.Fatalf("unexpected err: %v", err)
		}
		if isMultisig {
			b.Fatalf("script should NOT be reported as mutisig script")
		}
	}
}

// BenchmarkIsMultisigScript benchmarks how long it takes IsMultisigScript to
// analyze a 1-of-2 multisig public key script.
func BenchmarkIsMultisigScript(b *testing.B) {
	multisigShortForm := "1 " +
		"DATA_33 " +
		"0x030478aaaa2be30772f1e69e581610f1840b3cf2fe7228ee0281cd599e5746f81e " +
		"DATA_33 " +
		"0x0284f4d078b236a9ff91661f8ffbe012737cd3507566f30fd97d25f2b23539f3cd " +
		"2 CHECKMULTISIG"
	pkScript := mustParseShortForm(multisigShortForm)

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		isMultisig, err := IsMultisigScript(pkScript)
		if err != nil {
			b.Fatalf("unexpected err: %v", err)
		}
		if !isMultisig {
			b.Fatalf("script should be reported as a mutisig script")
		}
	}
}

// BenchmarkIsMultisigSigScript benchmarks how long it takes IsMultisigSigScript
// to analyze a very large script.
func BenchmarkIsMultisigSigScriptLarge(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		if IsMultisigSigScript(script) {
			b.Fatalf("script should NOT be reported as mutisig sig script")
		}
	}
}

// BenchmarkIsMultisigSigScript benchmarks how long it takes IsMultisigSigScript
// to analyze both a 1-of-2 multisig public key script (which should be false)
// and a signature script comprised of a pay-to-script-hash 1-of-2 multisig
// redeem script (which should be true).
func BenchmarkIsMultisigSigScript(b *testing.B) {
	multisigShortForm := "1 " +
		"DATA_33 " +
		"0x030478aaaa2be30772f1e69e581610f1840b3cf2fe7228ee0281cd599e5746f81e " +
		"DATA_33 " +
		"0x0284f4d078b236a9ff91661f8ffbe012737cd3507566f30fd97d25f2b23539f3cd " +
		"2 CHECKMULTISIG"
	pkScript := mustParseShortForm(multisigShortForm)

	sigHex := "0x304402205795c3ab6ba11331eeac757bf1fc9c34bef0c7e1a9c8bd5eebb8" +
		"82f3b79c5838022001e0ab7b4c7662e4522dc5fa479e4b4133fa88c6a53d895dc1d5" +
		"2eddc7bbcf2801 "
	sigScript := mustParseShortForm("DATA_71 " + sigHex + "DATA_71 " +
		multisigShortForm)

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		if IsMultisigSigScript(pkScript) {
			b.Fatalf("script should NOT be reported as mutisig sig script")
		}
		if !IsMultisigSigScript(sigScript) {
			b.Fatalf("script should be reported as a mutisig sig script")
		}
	}
}

// BenchmarkIsPushOnlyScript benchmarks how long it takes IsPushOnlyScript to
// analyze a very large script.
func BenchmarkIsPushOnlyScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		_ = IsPushOnlyScript(script)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2018-2019 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package txscript

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"testing"

	"github.com/decred/dcrd/wire"
)

var (
	// manyInputsBenchTx is a transaction that contains a lot of inputs which is
	// useful for benchmarking signature hash calculation.
	manyInputsBenchTx wire.MsgTx

	// A mock previous output script to use in the signing benchmark.
	prevOutScript = hexToBytes("a914f5916158e3e2c4551c1796708db8367207ed13bb87")
)

func init() {
	// tx 620f57c92cf05a7f7e7f7d28255d5f7089437bc48e34dcfebf7751d08b7fb8f5
	txHex, err := ioutil.ReadFile("data/many_inputs_tx.hex")
	if err != nil {
		panic(fmt.Sprintf("unable to read benchmark tx file: %v", err))
	}

	txBytes := hexToBytes(string(txHex))
	err = manyInputsBenchTx.Deserialize(bytes.NewReader(txBytes))
	if err != nil {
		panic(err)
	}
}

// BenchmarkCalcSigHash benchmarks how long it takes to calculate the signature
// hashes for all inputs of a transaction with many inputs.
func BenchmarkCalcSigHash(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for j := 0; j < len(manyInputsBenchTx.TxIn); j++ {
			_, err := CalcSignatureHash(prevOutScript, SigHashAll,
				&manyInputsBenchTx, j, nil)
			if err != nil {
				b.Fatalf("failed to calc signature hash: %v", err)
			}
		}
	}
}

// genComplexScript returns a script comprised of half as many opcodes as the
// maximum allowed followed by as many max size data pushes fit without
// exceeding the max allowed script size.
func genComplexScript() ([]byte, error) {
	var scriptLen int
	builder := NewScriptBuilder()
	for i := 0; i < MaxOpsPerScript/2; i++ {
		builder.AddOp(OP_TRUE)
		scriptLen++
	}
	maxData := bytes.Repeat([]byte{0x02}, MaxScriptElementSize)
	for i := 0; i < (MaxScriptSize-scriptLen)/MaxScriptElementSize; i++ {
		builder.AddData(maxData)
	}
	return builder.Script()
}

// BenchmarkScriptParsing benchmarks how long it takes to parse a very large
// script.
func BenchmarkScriptParsing(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	const scriptVersion = 0
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		tokenizer := MakeScriptTokenizer(scriptVersion, script)
		for tokenizer.Next() {
			_ = tokenizer.Opcode()
			_ = tokenizer.Data()
			_ = tokenizer.ByteIndex()
		}
		if err := tokenizer.Err(); err != nil {
			b.Fatalf("failed to parse script: %v", err)
		}
	}
}

// BenchmarkDisasmString benchmarks how long it takes to disassemble a very
// large script.
func BenchmarkDisasmString(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, err := DisasmString(script)
		if err != nil {
			b.Fatalf("failed to disasm script: %v", err)
		}
	}
}

// BenchmarkIsPayToScriptHash benchmarks how long it takes IsPayToScriptHash to
// analyze a very large script.
func BenchmarkIsPayToScriptHash(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = IsPayToScriptHash(script)
	}
}

// BenchmarkIsMultisigScriptLarge benchmarks how long it takes IsMultisigScript
// to analyze a very large script.
func BenchmarkIsMultisigScriptLarge(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		isMultisig, err := IsMultisigScript(script)
		if err != nil {
			b.Fatalf("unexpected err: %v", err)
		}
		if isMultisig {
			b.Fatalf("script should NOT be reported as mutisig script")
		}
	}
}

// BenchmarkIsMultisigScript benchmarks how long it takes IsMultisigScript to
// analyze a 1-of-2 multisig public key script.
func BenchmarkIsMultisigScript(b *testing.B) {
	multisigShortForm := "1 " +
		"DATA_33 " +
		"0x030478aaaa2be30772f1e69e581610f1840b3cf2fe7228ee0281cd599e5746f81e " +
		"DATA_33 " +
		"0x0284f4d078b236a9ff91661f8ffbe012737cd3507566f30fd97d25f2b23539f3cd " +
		"2 CHECKMULTISIG"
	pkScript := mustParseShortForm(multisigShortForm)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		isMultisig, err := IsMultisigScript(pkScript)
		if err != nil {
			b.Fatalf("unexpected err: %v", err)
		}
		if !isMultisig {
			b.Fatalf("script should be reported as a mutisig script")
		}
	}
}

// BenchmarkIsMultisigSigScript benchmarks how long it takes IsMultisigSigScript
// to analyze a very large script.
func BenchmarkIsMultisigSigScriptLarge(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if IsMultisigSigScript(script) {
			b.Fatalf("script should NOT be reported as mutisig sig script")
		}
	}
}

// BenchmarkIsMultisigSigScript benchmarks how long it takes IsMultisigSigScript
// to analyze both a 1-of-2 multisig public key script (which should be false)
// and a signature script comprised of a pay-to-script-hash 1-of-2 multisig
// redeem script (which should be true).
func BenchmarkIsMultisigSigScript(b *testing.B) {
	multisigShortForm := "1 " +
		"DATA_33 " +
		"0x030478aaaa2be30772f1e69e581610f1840b3cf2fe7228ee0281cd599e5746f81e " +
		"DATA_33 " +
		"0x0284f4d078b236a9ff91661f8ffbe012737cd3507566f30fd97d25f2b23539f3cd " +
		"2 CHECKMULTISIG"
	pkScript := mustParseShortForm(multisigShortForm)

	sigHex := "0x304402205795c3ab6ba11331eeac757bf1fc9c34bef0c7e1a9c8bd5eebb8" +
		"82f3b79c5838022001e0ab7b4c7662e4522dc5fa479e4b4133fa88c6a53d895dc1d5" +
		"2eddc7bbcf2801 "
	sigScript := mustParseShortForm("DATA_71 " + sigHex + "DATA_71 " +
		multisigShortForm)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if IsMultisigSigScript(pkScript) {
			b.Fatalf("script should NOT be reported as mutisig sig script")
		}
		if !IsMultisigSigScript(sigScript) {
			b.Fatalf("script should be reported as a mutisig sig script")
		}
	}
}

// BenchmarkGetSigOpCount benchmarks how long it takes to count the signature
// operations of a very large script.
func BenchmarkGetSigOpCount(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = GetSigOpCount(script)
	}
}

// BenchmarkGetPreciseSigOpCount benchmarks how long it takes to count the
// signature operations of a very large script using the more precise counting
// method.
func BenchmarkGetPreciseSigOpCount(b *testing.B) {
	redeemScript, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	// Create a fake pay-to-script-hash to pass the necessary checks and create
	// the signature script accordingly by pushing the generated "redeem" script
	// as the final data push so the benchmark will cover the p2sh path.
	scriptHash := "0x0000000000000000000000000000000000000001"
	pkScript := mustParseShortForm("HASH160 DATA_20 " + scriptHash + " EQUAL")
	sigScript, err := NewScriptBuilder().AddFullData(redeemScript).Script()
	if err != nil {
		b.Fatalf("failed to create signature script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = GetPreciseSigOpCount(sigScript, pkScript, true)
	}
}

// BenchmarkIsAnyKindOfScriptHash benchmarks how long it takes to
// isAnyKindOfScriptHash to analyze operations of a very large script.
func BenchmarkIsAnyKindOfScriptHash(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = isAnyKindOfScriptHash(script)
	}
}

// BenchmarkIsPushOnlyScript benchmarks how long it takes IsPushOnlyScript to
// analyze a very large script.
func BenchmarkIsPushOnlyScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = IsPushOnlyScript(script)
	}
}

// BenchmarkGetScriptClass benchmarks how long it takes GetScriptClass to
// analyze a very large script.
func BenchmarkGetScriptClass(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	const scriptVersion = 0
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = GetScriptClass(scriptVersion, script)
	}
}

// BenchmarkIsPubKeyScript benchmarks how long it takes to analyze a very large
// script to determine if it is a standard pay-to-pubkey script.
func BenchmarkIsPubKeyScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = isPubKeyScript(script)
	}
}

// BenchmarkIsAltPubKeyScript benchmarks how long it takes to analyze a very
// large script to determine if it is a standard pay-to-alt-pubkey script.
func BenchmarkIsAltPubKeyScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = isPubKeyAltScript(script)
	}
}

// BenchmarkIsPubKeyHashScript benchmarks how long it takes to analyze a very
// large script to determine if it is a standard pay-to-pubkey-hash script.
func BenchmarkIsPubKeyHashScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = isPubKeyHashScript(script)
	}
}

// BenchmarkIsAltPubKeyHashScript benchmarks how long it takes to analyze a very
// large script to determine if it is a standard pay-to-alt-pubkey-hash script.
func BenchmarkIsAltPubKeyHashScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = isPubKeyHashAltScript(script)
	}
}

// BenchmarkIsNullDataScript benchmarks how long it takes to analyze a very
// large script to determine if it is a standard nulldata script.
func BenchmarkIsNullDataScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	const scriptVersion = 0
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = isNullDataScript(scriptVersion, script)
	}
}

// BenchmarkIsStakeSubmissionScript benchmarks how long it takes to analyze a
// very large script to determine if it is a standard stake submission script.
func BenchmarkIsStakeSubmissionScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	const scriptVersion = 0
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = isStakeSubmissionScript(scriptVersion, script)
	}
}
<commit_msg>txscript: Add bench for stake generation scripts.<commit_after>// Copyright (c) 2018-2019 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package txscript

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"testing"

	"github.com/decred/dcrd/wire"
)

var (
	// manyInputsBenchTx is a transaction that contains a lot of inputs which is
	// useful for benchmarking signature hash calculation.
	manyInputsBenchTx wire.MsgTx

	// A mock previous output script to use in the signing benchmark.
	prevOutScript = hexToBytes("a914f5916158e3e2c4551c1796708db8367207ed13bb87")
)

func init() {
	// tx 620f57c92cf05a7f7e7f7d28255d5f7089437bc48e34dcfebf7751d08b7fb8f5
	txHex, err := ioutil.ReadFile("data/many_inputs_tx.hex")
	if err != nil {
		panic(fmt.Sprintf("unable to read benchmark tx file: %v", err))
	}

	txBytes := hexToBytes(string(txHex))
	err = manyInputsBenchTx.Deserialize(bytes.NewReader(txBytes))
	if err != nil {
		panic(err)
	}
}

// BenchmarkCalcSigHash benchmarks how long it takes to calculate the signature
// hashes for all inputs of a transaction with many inputs.
func BenchmarkCalcSigHash(b *testing.B) {
	for i := 0; i < b.N; i++ {
		for j := 0; j < len(manyInputsBenchTx.TxIn); j++ {
			_, err := CalcSignatureHash(prevOutScript, SigHashAll,
				&manyInputsBenchTx, j, nil)
			if err != nil {
				b.Fatalf("failed to calc signature hash: %v", err)
			}
		}
	}
}

// genComplexScript returns a script comprised of half as many opcodes as the
// maximum allowed followed by as many max size data pushes fit without
// exceeding the max allowed script size.
func genComplexScript() ([]byte, error) {
	var scriptLen int
	builder := NewScriptBuilder()
	for i := 0; i < MaxOpsPerScript/2; i++ {
		builder.AddOp(OP_TRUE)
		scriptLen++
	}
	maxData := bytes.Repeat([]byte{0x02}, MaxScriptElementSize)
	for i := 0; i < (MaxScriptSize-scriptLen)/MaxScriptElementSize; i++ {
		builder.AddData(maxData)
	}
	return builder.Script()
}

// BenchmarkScriptParsing benchmarks how long it takes to parse a very large
// script.
func BenchmarkScriptParsing(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	const scriptVersion = 0
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		tokenizer := MakeScriptTokenizer(scriptVersion, script)
		for tokenizer.Next() {
			_ = tokenizer.Opcode()
			_ = tokenizer.Data()
			_ = tokenizer.ByteIndex()
		}
		if err := tokenizer.Err(); err != nil {
			b.Fatalf("failed to parse script: %v", err)
		}
	}
}

// BenchmarkDisasmString benchmarks how long it takes to disassemble a very
// large script.
func BenchmarkDisasmString(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, err := DisasmString(script)
		if err != nil {
			b.Fatalf("failed to disasm script: %v", err)
		}
	}
}

// BenchmarkIsPayToScriptHash benchmarks how long it takes IsPayToScriptHash to
// analyze a very large script.
func BenchmarkIsPayToScriptHash(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = IsPayToScriptHash(script)
	}
}

// BenchmarkIsMultisigScriptLarge benchmarks how long it takes IsMultisigScript
// to analyze a very large script.
func BenchmarkIsMultisigScriptLarge(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		isMultisig, err := IsMultisigScript(script)
		if err != nil {
			b.Fatalf("unexpected err: %v", err)
		}
		if isMultisig {
			b.Fatalf("script should NOT be reported as mutisig script")
		}
	}
}

// BenchmarkIsMultisigScript benchmarks how long it takes IsMultisigScript to
// analyze a 1-of-2 multisig public key script.
func BenchmarkIsMultisigScript(b *testing.B) {
	multisigShortForm := "1 " +
		"DATA_33 " +
		"0x030478aaaa2be30772f1e69e581610f1840b3cf2fe7228ee0281cd599e5746f81e " +
		"DATA_33 " +
		"0x0284f4d078b236a9ff91661f8ffbe012737cd3507566f30fd97d25f2b23539f3cd " +
		"2 CHECKMULTISIG"
	pkScript := mustParseShortForm(multisigShortForm)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		isMultisig, err := IsMultisigScript(pkScript)
		if err != nil {
			b.Fatalf("unexpected err: %v", err)
		}
		if !isMultisig {
			b.Fatalf("script should be reported as a mutisig script")
		}
	}
}

// BenchmarkIsMultisigSigScript benchmarks how long it takes IsMultisigSigScript
// to analyze a very large script.
func BenchmarkIsMultisigSigScriptLarge(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if IsMultisigSigScript(script) {
			b.Fatalf("script should NOT be reported as mutisig sig script")
		}
	}
}

// BenchmarkIsMultisigSigScript benchmarks how long it takes IsMultisigSigScript
// to analyze both a 1-of-2 multisig public key script (which should be false)
// and a signature script comprised of a pay-to-script-hash 1-of-2 multisig
// redeem script (which should be true).
func BenchmarkIsMultisigSigScript(b *testing.B) {
	multisigShortForm := "1 " +
		"DATA_33 " +
		"0x030478aaaa2be30772f1e69e581610f1840b3cf2fe7228ee0281cd599e5746f81e " +
		"DATA_33 " +
		"0x0284f4d078b236a9ff91661f8ffbe012737cd3507566f30fd97d25f2b23539f3cd " +
		"2 CHECKMULTISIG"
	pkScript := mustParseShortForm(multisigShortForm)

	sigHex := "0x304402205795c3ab6ba11331eeac757bf1fc9c34bef0c7e1a9c8bd5eebb8" +
		"82f3b79c5838022001e0ab7b4c7662e4522dc5fa479e4b4133fa88c6a53d895dc1d5" +
		"2eddc7bbcf2801 "
	sigScript := mustParseShortForm("DATA_71 " + sigHex + "DATA_71 " +
		multisigShortForm)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if IsMultisigSigScript(pkScript) {
			b.Fatalf("script should NOT be reported as mutisig sig script")
		}
		if !IsMultisigSigScript(sigScript) {
			b.Fatalf("script should be reported as a mutisig sig script")
		}
	}
}

// BenchmarkGetSigOpCount benchmarks how long it takes to count the signature
// operations of a very large script.
func BenchmarkGetSigOpCount(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = GetSigOpCount(script)
	}
}

// BenchmarkGetPreciseSigOpCount benchmarks how long it takes to count the
// signature operations of a very large script using the more precise counting
// method.
func BenchmarkGetPreciseSigOpCount(b *testing.B) {
	redeemScript, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	// Create a fake pay-to-script-hash to pass the necessary checks and create
	// the signature script accordingly by pushing the generated "redeem" script
	// as the final data push so the benchmark will cover the p2sh path.
	scriptHash := "0x0000000000000000000000000000000000000001"
	pkScript := mustParseShortForm("HASH160 DATA_20 " + scriptHash + " EQUAL")
	sigScript, err := NewScriptBuilder().AddFullData(redeemScript).Script()
	if err != nil {
		b.Fatalf("failed to create signature script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = GetPreciseSigOpCount(sigScript, pkScript, true)
	}
}

// BenchmarkIsAnyKindOfScriptHash benchmarks how long it takes to
// isAnyKindOfScriptHash to analyze operations of a very large script.
func BenchmarkIsAnyKindOfScriptHash(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = isAnyKindOfScriptHash(script)
	}
}

// BenchmarkIsPushOnlyScript benchmarks how long it takes IsPushOnlyScript to
// analyze a very large script.
func BenchmarkIsPushOnlyScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = IsPushOnlyScript(script)
	}
}

// BenchmarkGetScriptClass benchmarks how long it takes GetScriptClass to
// analyze a very large script.
func BenchmarkGetScriptClass(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	const scriptVersion = 0
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = GetScriptClass(scriptVersion, script)
	}
}

// BenchmarkIsPubKeyScript benchmarks how long it takes to analyze a very large
// script to determine if it is a standard pay-to-pubkey script.
func BenchmarkIsPubKeyScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = isPubKeyScript(script)
	}
}

// BenchmarkIsAltPubKeyScript benchmarks how long it takes to analyze a very
// large script to determine if it is a standard pay-to-alt-pubkey script.
func BenchmarkIsAltPubKeyScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = isPubKeyAltScript(script)
	}
}

// BenchmarkIsPubKeyHashScript benchmarks how long it takes to analyze a very
// large script to determine if it is a standard pay-to-pubkey-hash script.
func BenchmarkIsPubKeyHashScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = isPubKeyHashScript(script)
	}
}

// BenchmarkIsAltPubKeyHashScript benchmarks how long it takes to analyze a very
// large script to determine if it is a standard pay-to-alt-pubkey-hash script.
func BenchmarkIsAltPubKeyHashScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = isPubKeyHashAltScript(script)
	}
}

// BenchmarkIsNullDataScript benchmarks how long it takes to analyze a very
// large script to determine if it is a standard nulldata script.
func BenchmarkIsNullDataScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	const scriptVersion = 0
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = isNullDataScript(scriptVersion, script)
	}
}

// BenchmarkIsStakeSubmissionScript benchmarks how long it takes to analyze a
// very large script to determine if it is a standard stake submission script.
func BenchmarkIsStakeSubmissionScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	const scriptVersion = 0
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = isStakeSubmissionScript(scriptVersion, script)
	}
}

// BenchmarkIsStakeGenerationScript benchmarks how long it takes to analyze a
// very large script to determine if it is a standard stake generation script.
func BenchmarkIsStakeGenerationScript(b *testing.B) {
	script, err := genComplexScript()
	if err != nil {
		b.Fatalf("failed to create benchmark script: %v", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		pops, _ := parseScript(script)
		_ = isStakeGen(pops)
	}
}
<|endoftext|>


<commit_before>package types

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"crypto/tls"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"hash"
	"strings"
)

//EncryptedKey contains the decryption key data from the saml2 core and xmlenc
//standards.
type EncryptedKey struct {
	// EncryptionMethod string `xml:"EncryptionMethod>Algorithm"`
	X509Data         string `xml:"KeyInfo>X509Data>X509Certificate"`
	CipherValue      string `xml:"CipherData>CipherValue"`
	EncryptionMethod EncryptionMethod
}

//EncryptionMethod specifies the type of encryption that was used.
type EncryptionMethod struct {
	Algorithm    string `xml:",attr"`
	DigestMethod DigestMethod
}

//DigestMethod is a digest type specification
type DigestMethod struct {
	Algorithm string `xml:",attr"`
}

//Well-known public-key encryption methods
const (
	MethodRSAOAEP  = "http://www.w3.org/2001/04/xmlenc#rsa-oaep-mgf1p"
	MethodRSAOAEP2 = "http://www.w3.org/2009/xmlenc11#rsa-oaep"
)

//Well-known private key encryption methods
const (
	MethodAES128GCM = "http://www.w3.org/2009/xmlenc11#aes128-gcm"
	MethodAES128CBC = "http://www.w3.org/2001/04/xmlenc#aes128-cbc"
	MethodAES256CBC = "http://www.w3.org/2001/04/xmlenc#aes256-cbc"
)

//Well-known hash methods
const (
	MethodSHA1   = "http://www.w3.org/2000/09/xmldsig#sha1"
	MethodSHA256 = "http://www.w3.org/2000/09/xmldsig#sha256"
	MethodSHA512 = "http://www.w3.org/2000/09/xmldsig#sha512"
)

//SHA-1 is commonly used for certificate fingerprints (openssl -fingerprint and ADFS thumbprint).
//SHA-1 is sufficient for our purposes here (error message).
func debugKeyFp(keyBytes []byte) string {
	hashFunc := sha1.New()
	hashFunc.Write(keyBytes)
	sum := strings.ToLower(hex.EncodeToString(hashFunc.Sum(nil)))
	var ret string
	for idx := 0; idx+1 < len(sum); idx += 2 {
		if idx == 0 {
			ret += sum[idx : idx+2]
		} else {
			ret += ":" + sum[idx:idx+2]
		}
	}
	return ret
}

//DecryptSymmetricKey returns the private key contained in the EncryptedKey document
func (ek *EncryptedKey) DecryptSymmetricKey(cert *tls.Certificate) (cipher.Block, error) {
	encCert, err := base64.StdEncoding.DecodeString(ek.X509Data)
	if err != nil {
		return nil, fmt.Errorf("error getting certificate from encryptedkey: %v", err)
	}

	if len(cert.Certificate) < 1 {
		return nil, fmt.Errorf("decryption tls.Certificate has no public certs attached")
	}

	if !bytes.Equal(cert.Certificate[0], encCert) {
		return nil, fmt.Errorf("key decryption attempted with mismatched cert, SP cert(%.11s), assertion cert(%.11s)",
			debugKeyFp(cert.Certificate[0]), debugKeyFp(encCert))
	}

	cipherText, err := base64.StdEncoding.DecodeString(ek.CipherValue)
	if err != nil {
		return nil, err
	}

	switch pk := cert.PrivateKey.(type) {
	case *rsa.PrivateKey:
		var h hash.Hash

		switch ek.EncryptionMethod.DigestMethod.Algorithm {
		case MethodSHA1:
			h = sha1.New()
		case MethodSHA256:
			h = sha256.New()
		case MethodSHA512:
			h = sha512.New()
		}

		switch ek.EncryptionMethod.Algorithm {
		case MethodRSAOAEP, MethodRSAOAEP2:
			pt, err := rsa.DecryptOAEP(h, rand.Reader, pk, cipherText, nil)
			if err != nil {
				return nil, fmt.Errorf("rsa internal error: %v", err)
			}

			b, err := aes.NewCipher(pt)
			if err != nil {
				return nil, err
			}

			return b, nil
		default:
			return nil, fmt.Errorf("unsupported encryption algorithm: %s", ek.EncryptionMethod.Algorithm)
		}
	}
	return nil, fmt.Errorf("no cipher for decoding symmetric key")
}
<commit_msg>Cert fingerprint is empty when cert is empty.<commit_after>package types

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"crypto/tls"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"hash"
	"strings"
)

//EncryptedKey contains the decryption key data from the saml2 core and xmlenc
//standards.
type EncryptedKey struct {
	// EncryptionMethod string `xml:"EncryptionMethod>Algorithm"`
	X509Data         string `xml:"KeyInfo>X509Data>X509Certificate"`
	CipherValue      string `xml:"CipherData>CipherValue"`
	EncryptionMethod EncryptionMethod
}

//EncryptionMethod specifies the type of encryption that was used.
type EncryptionMethod struct {
	Algorithm    string `xml:",attr"`
	DigestMethod DigestMethod
}

//DigestMethod is a digest type specification
type DigestMethod struct {
	Algorithm string `xml:",attr"`
}

//Well-known public-key encryption methods
const (
	MethodRSAOAEP  = "http://www.w3.org/2001/04/xmlenc#rsa-oaep-mgf1p"
	MethodRSAOAEP2 = "http://www.w3.org/2009/xmlenc11#rsa-oaep"
)

//Well-known private key encryption methods
const (
	MethodAES128GCM = "http://www.w3.org/2009/xmlenc11#aes128-gcm"
	MethodAES128CBC = "http://www.w3.org/2001/04/xmlenc#aes128-cbc"
	MethodAES256CBC = "http://www.w3.org/2001/04/xmlenc#aes256-cbc"
)

//Well-known hash methods
const (
	MethodSHA1   = "http://www.w3.org/2000/09/xmldsig#sha1"
	MethodSHA256 = "http://www.w3.org/2000/09/xmldsig#sha256"
	MethodSHA512 = "http://www.w3.org/2000/09/xmldsig#sha512"
)

//SHA-1 is commonly used for certificate fingerprints (openssl -fingerprint and ADFS thumbprint).
//SHA-1 is sufficient for our purposes here (error message).
func debugKeyFp(keyBytes []byte) string {
	if len(keyBytes) < 1 {
		return ""
	}
	hashFunc := sha1.New()
	hashFunc.Write(keyBytes)
	sum := strings.ToLower(hex.EncodeToString(hashFunc.Sum(nil)))
	var ret string
	for idx := 0; idx+1 < len(sum); idx += 2 {
		if idx == 0 {
			ret += sum[idx : idx+2]
		} else {
			ret += ":" + sum[idx:idx+2]
		}
	}
	return ret
}

//DecryptSymmetricKey returns the private key contained in the EncryptedKey document
func (ek *EncryptedKey) DecryptSymmetricKey(cert *tls.Certificate) (cipher.Block, error) {
	encCert, err := base64.StdEncoding.DecodeString(ek.X509Data)
	if err != nil {
		return nil, fmt.Errorf("error getting certificate from encryptedkey: %v", err)
	}

	if len(cert.Certificate) < 1 {
		return nil, fmt.Errorf("decryption tls.Certificate has no public certs attached")
	}

	if !bytes.Equal(cert.Certificate[0], encCert) {
		return nil, fmt.Errorf("key decryption attempted with mismatched cert, SP cert(%.11s), assertion cert(%.11s)",
			debugKeyFp(cert.Certificate[0]), debugKeyFp(encCert))
	}

	cipherText, err := base64.StdEncoding.DecodeString(ek.CipherValue)
	if err != nil {
		return nil, err
	}

	switch pk := cert.PrivateKey.(type) {
	case *rsa.PrivateKey:
		var h hash.Hash

		switch ek.EncryptionMethod.DigestMethod.Algorithm {
		case MethodSHA1:
			h = sha1.New()
		case MethodSHA256:
			h = sha256.New()
		case MethodSHA512:
			h = sha512.New()
		}

		switch ek.EncryptionMethod.Algorithm {
		case MethodRSAOAEP, MethodRSAOAEP2:
			pt, err := rsa.DecryptOAEP(h, rand.Reader, pk, cipherText, nil)
			if err != nil {
				return nil, fmt.Errorf("rsa internal error: %v", err)
			}

			b, err := aes.NewCipher(pt)
			if err != nil {
				return nil, err
			}

			return b, nil
		default:
			return nil, fmt.Errorf("unsupported encryption algorithm: %s", ek.EncryptionMethod.Algorithm)
		}
	}
	return nil, fmt.Errorf("no cipher for decoding symmetric key")
}
<|endoftext|>


<commit_before>package querying

import (
	"errors"

	. "github.com/herald-it/goncord/models"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

func FindUserID(obj *User, c *mgo.Collection) (*User, error) {
	var results []User

	err := c.Find(
		bson.M{"$or": []bson.M{
			bson.M{"_id": obj.Id},
		}}).All(&results)
	if err != nil {
		return nil, err
	}

	if len(results) > 1 {
		return nil, errors.New("Find user returned more 1 record.")
	}

	if len(results) == 0 {
		return nil, nil
	}

	return &results[0], nil
}

func FindUser(obj *User, c *mgo.Collection) (*User, error) {
	var results []User

	err := c.Find(
		bson.M{"$or": []bson.M{
			bson.M{"login": obj.Login},
			bson.M{"email": obj.Email},
		}}).All(&results)
	if err != nil {
		return nil, err
	}

	if len(results) > 1 {
		return nil, errors.New("Find user returned more 1 record.")
	}

	if len(results) == 0 {
		return nil, nil
	}

	return &results[0], nil
}

func FindDumpToken(obj *DumpToken, c *mgo.Collection) (*DumpToken, error) {
	var results []DumpToken

	err := c.Find(obj).All(&results)
	if err != nil {
		return nil, err
	}

	if len(results) > 1 {
		return nil, errors.New("Find user returned more 1 record.")
	}

	if len(results) == 0 {
		return nil, nil
	}

	return &results[0], nil
}
<commit_msg>Was found std method find by id<commit_after>package querying

import (
	"errors"

	. "github.com/herald-it/goncord/models"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

func FindUserID(obj *User, c *mgo.Collection) (*User, error) {
	var results []User

	err := c.FindId(obj.Id).All(&results)
	if err != nil {
		return nil, err
	}

	if len(results) > 1 {
		return nil, errors.New("Find user returned more 1 record.")
	}

	if len(results) == 0 {
		return nil, nil
	}

	return &results[0], nil
}

func FindUser(obj *User, c *mgo.Collection) (*User, error) {
	var results []User

	err := c.Find(
		bson.M{"$or": []bson.M{
			bson.M{"login": obj.Login},
			bson.M{"email": obj.Email},
		}}).All(&results)
	if err != nil {
		return nil, err
	}

	if len(results) > 1 {
		return nil, errors.New("Find user returned more 1 record.")
	}

	if len(results) == 0 {
		return nil, nil
	}

	return &results[0], nil
}

func FindDumpToken(obj *DumpToken, c *mgo.Collection) (*DumpToken, error) {
	var results []DumpToken

	err := c.Find(obj).All(&results)
	if err != nil {
		return nil, err
	}

	if len(results) > 1 {
		return nil, errors.New("Find user returned more 1 record.")
	}

	if len(results) == 0 {
		return nil, nil
	}

	return &results[0], nil
}
<|endoftext|>


<commit_before>package ripper

import (
	"github.com/goby-lang/goby/compiler"
	"github.com/goby-lang/goby/compiler/bytecode"
	"github.com/goby-lang/goby/compiler/lexer"
	"github.com/goby-lang/goby/compiler/parser"
	"github.com/goby-lang/goby/compiler/token"
	"github.com/goby-lang/goby/vm"
	"github.com/goby-lang/goby/vm/classes"
	"github.com/goby-lang/goby/vm/errors"
	"strings"
)

// Ripper is a loadable library and has abilities to obtain parsed/lexed/tokenized Goby codes from String.
// The library would be convenient for validating Goby codes when building lint tools,
// as well as the tests for Goby's compiler.
// For now, Ripper is a class and has only class methods, but I think this should finally be a 'newable' module with more sophisticated instance methods.

// Imported objects from vm
type Object = vm.Object
type VM = vm.VM
type Thread = vm.Thread
type Method = vm.Method
type StringObject = vm.StringObject
type HashObject = vm.HashObject
type ArrayObject = vm.ArrayObject

// Class methods --------------------------------------------------------

// Returns the list of instruction code generated from Goby code.
// Returns an error when the code is invalid.
// The return value is a "tuple" style nested array:
// - `Array`: contains an instruction set
//   - `arg_types:` (none if `nil`)
// 		 - `names:` array of names (string)
//     - `types:` array of types (integer)
//   - `instructions:` array of instructions
//     - `action:` string
//     - `anchor:` integer
//     - `line:` integer
//     - `params:` array of parameters (string)
//     - `source_line:` integer
//     - `arg_set:` (none if `nil`)
//   		 - `names:` array of names (string)
//       - `types:` array of types (integer)
//
// ```ruby
// require 'ripper'; Ripper.instruction "10.times do |i| puts i end"
// #=>
//
// require 'ripper'; Ripper.instruction "10.times do |i| puts i" # the code is invalid
// #=> TypeError: InternalError%!(EXTRA string=String, string=Invalid Goby code)
// ```
//
// @param Goby code [String]
// @return [Array]
func instruction(receiver Object, sourceLine int, t *Thread, args []Object) Object {
	if len(args) != 1 {
		return t.VM().InitErrorObject(errors.ArgumentError, sourceLine, "Expect 1 argument. got=%d", len(args))
	}

	arg, ok := args[0].(*StringObject)
	if !ok {
		return t.VM().InitErrorObject(errors.TypeError, sourceLine, errors.WrongArgumentTypeFormat, classes.StringClass, args[0].Class().Name)
	}

	i, _ := compiler.CompileToInstructions(arg.Value().(string), parser.NormalMode)

	return convertToTuple(i, t.VM())
}

// Returns a nested array that contains the line #, type of the tokenize, and the literal of the tokenize.
// Note that the class method does not return any errors even though the provided Goby code is invalid.
//
// ```ruby
// require 'ripper'; Ripper.lex "10.times do |i| puts i end"
// #=> [[0, "on_int", "10"], [0, "on_dot", "."], [0, "on_ident", "times"], [0, "on_do", "do"], [0, "on_bar", "|"], [0, "on_ident", "i"], [0, "on_bar", "|"], [0, "on_ident", "puts"], [0, "on_ident", "i"], [0, "on_end", "end"], [0, "on_eof", ""]]
//
// require 'ripper'; Ripper.lex "10.times do |i| puts i" # the code is invalid
// #=> [[0, "on_int", "10"], [0, "on_dot", "."], [0, "on_ident", "times"], [0, "on_do", "do"], [0, "on_bar", "|"], [0, "on_ident", "i"], [0, "on_bar", "|"], [0, "on_ident", "puts"], [0, "on_ident", "i"], [0, "on_eof", ""]]
// ```
//
// @param Goby code [String]
// @return [Array]
func lex(receiver Object, sourceLine int, t *Thread, args []Object) Object  {
	if len(args) != 1 {
		return t.VM().InitErrorObject(errors.ArgumentError, sourceLine, "Expect 1 argument. got=%d", len(args))
	}

	arg, ok := args[0].(*StringObject)
	if !ok {
		return t.VM().InitErrorObject(errors.TypeError, sourceLine, errors.WrongArgumentTypeFormat, classes.StringClass, args[0].Class().Name)
	}

	l := lexer.New(arg.Value().(string))
	array := t.VM().InitArrayObject([]Object{})
	var elements []Object
	var nextToken token.Token
	for i := 0; ; i++ {
		nextToken = l.NextToken()
		elements = append(elements, t.VM().InitIntegerObject(nextToken.Line))
		elements = append(elements, t.VM().InitStringObject(convertLex(nextToken.Type)))
		elements = append(elements, t.VM().InitStringObject(nextToken.Literal))
		array.Elements = append(array.Elements, t.VM().InitArrayObject(elements))
		if nextToken.Type == token.EOF {
			break
		}
		elements = nil
	}
	return array
}

// Just to disable creating instances.
func new(receiver Object, sourceLine int, t *Thread, args []Object) Object {
	return t.VM().InitUnsupportedMethodError(sourceLine, "#new", receiver)
}

// Returns the parsed Goby codes as a String object.
// Returns an error when the code is invalid.
//
// ```ruby
// require 'ripper'; Ripper.parse "10.times do |i| puts i end"
// #=> "10.times() do |i|
// #=> self.puts(i)
// #=> end"
//
// require 'ripper'; Ripper.parse "10.times do |i| puts i" # the code is invalid
// #=> TypeError: InternalError%!(EXTRA string=String, string=Invalid Goby code)
// ```
//
// @param Goby code [String]
// @return [String]
func parse(receiver Object, sourceLine int, t *Thread, args []Object) Object {
	if len(args) != 1 {
		return t.VM().InitErrorObject(errors.ArgumentError, sourceLine, "Expect 1 argument. got=%d", len(args))
	}

	arg, ok := args[0].(*StringObject)
	if !ok {
		return t.VM().InitErrorObject(errors.TypeError, sourceLine, errors.WrongArgumentTypeFormat, classes.StringClass, args[0].Class().Name)
	}

	l := lexer.New(arg.Value().(string))
	p := parser.New(l)
	program, _ := p.ParseProgram()

	return t.VM().InitStringObject(program.String())
}

// Returns a tokenized Goby codes as an Array object.
// Note that this does not return any errors even though the provided code is invalid.
//
// ```ruby
// require 'ripper'; Ripper.tokenize "10.times do |i| puts i end"
// #=> ["10", ".", "times", "do", "|", "i", "|", "puts", "i", "end", "EOF"]
//
// require 'ripper'; Ripper.parse "10.times do |i| puts i" # the code is invalid
// #=> ["10", ".", "times", "do", "|", "i", "|", "puts", "i", "EOF"]
// ```
//
// @param Goby code [String]
// @return [String]
func tokenize(receiver Object, sourceLine int, t *Thread, args []Object) Object {
	if len(args) != 1 {
		return t.VM().InitErrorObject(errors.ArgumentError, sourceLine, "Expect 1 argument. got=%d", len(args))
	}

	arg, ok := args[0].(*StringObject)
	if !ok {
		return t.VM().InitErrorObject(errors.TypeError, sourceLine, errors.WrongArgumentTypeFormat, classes.StringClass, args[0].Class().Name)
	}

	l := lexer.New(arg.Value().(string))
	el := []Object{}
	var nt token.Token
	for i := 0; ; i++ {
		nt = l.NextToken()
		if nt.Type == token.EOF {
			el = append(el, t.VM().InitStringObject("EOF"))
			break
		}
		el = append(el, t.VM().InitStringObject(nt.Literal))
	}
	return t.VM().InitArrayObject(el)
}

// Internal functions ===================================================
func init() {
	vm.RegisterExternalClass("ripper", vm.ExternalClass("Ripper", "ripper.gb",
		// class methods
		map[string]vm.Method{
			"instruction": instruction,
			"lex":         lex,
			"new": 				 new,
			"parse":       parse,
			"tokenize":    tokenize,
		},
		// instance methods
		map[string]vm.Method{},
	))
}

// Other helper functions ----------------------------------------------

func convertToTuple(instSet []*bytecode.InstructionSet, v *VM) *ArrayObject {
	ary := []Object{}
	for _, instruction := range instSet {
		hashInstLevel1 := make(map[string]Object)
		hashInstLevel1["name"] = v.InitStringObject(instruction.Name())
		hashInstLevel1["type"] = v.InitStringObject(instruction.Type())
		if instruction.ArgTypes() != nil {
			hashInstLevel1["arg_types"] = getArgNameType(instruction.ArgTypes(), v)
		}
		ary = append(ary, v.InitHashObject(hashInstLevel1))

		arrayInst := []Object{}
		for _, ins := range instruction.Instructions {
			hashInstLevel2 := make(map[string]Object)
			hashInstLevel2["action"] = v.InitStringObject(ins.Action)
			hashInstLevel2["line"] = v.InitIntegerObject(ins.Line())
			hashInstLevel2["source_line"] = v.InitIntegerObject(ins.SourceLine())
			anchor, _ := ins.AnchorLine()
			hashInstLevel2["anchor"] = v.InitIntegerObject(anchor)

			arrayParams := []Object{}
			for _, param := range ins.Params {
				arrayParams = append(arrayParams, v.InitStringObject(param))
			}
			hashInstLevel2["params"] = v.InitArrayObject(arrayParams)

			if ins.ArgSet != nil {
				hashInstLevel1["arg_set"] = getArgNameType(ins.ArgSet, v)
			}

			arrayInst = append(arrayInst, v.InitHashObject(hashInstLevel2))
		}

		hashInstLevel1["instructions"] = v.InitArrayObject(arrayInst)
		ary = append(ary, v.InitHashObject(hashInstLevel1))
	}
	return v.InitArrayObject(ary)
}

func getArgNameType(argSet *bytecode.ArgSet, v *VM) *HashObject {
	h := make(map[string]Object)

	aName := []Object{}
	for _, argname := range argSet.Names() {
		aName = append(aName, v.InitStringObject(argname))
	}
	h["names"] = v.InitArrayObject(aName)

	aType := []Object{}
	for _, argtype := range argSet.Types() {
		aType = append(aType, v.InitIntegerObject(argtype))
	}

	h["types"] = v.InitArrayObject(aType)
	return v.InitHashObject(h)
}

// TODO: This should finally be auto-generated from tokenize.go
func convertLex(t token.Type) string {
	var s string
	switch t {
	case token.Asterisk:
		s = "asterisk"
	case token.And:
		s = "and"
	case token.Assign:
		s = "assign"
	case token.Bang:
		s = "bang"
	case token.Bar:
		s = "bar"
	case token.Colon:
		s = "colon"
	case token.Comma:
		s = "comma"
	case token.COMP:
		s = "comp"
	case token.Dot:
		s = "dot"
	case token.Eq:
		s = "eq"
	case token.GT:
		s = "gt"
	case token.GTE:
		s = "gte"
	case token.LBrace:
		s = "lbrace"
	case token.LBracket:
		s = "lbracket"
	case token.LParen:
		s = "lparen"
	case token.LT:
		s = "lt"
	case token.LTE:
		s = "lte"
	case token.Match:
		s = "match"
	case token.Minus:
		s = "minus"
	case token.MinusEq:
		s = "minuseq"
	case token.Modulo:
		s = "modulo"
	case token.NotEq:
		s = "noteq"
	case token.Or:
		s = "or"
	case token.OrEq:
		s = "oreq"
	case token.Plus:
		s = "plus"
	case token.PlusEq:
		s = "pluseq"
	case token.Pow:
		s = "pow"
	case token.Range:
		s = "range"
	case token.RBrace:
		s = "rbrace"
	case token.RBracket:
		s = "rbracket"
	case token.ResolutionOperator:
		s = "resolutionoperator"
	case token.RParen:
		s = "rparen"
	case token.Semicolon:
		s = "semicolon"
	case token.Slash:
		s = "slash"
	default:
		s = strings.ToLower(string(t))
	}

	return "on_" + s
}
<commit_msg>Fix apidoc in Ripper#instruction<commit_after>package ripper

import (
	"github.com/goby-lang/goby/compiler"
	"github.com/goby-lang/goby/compiler/bytecode"
	"github.com/goby-lang/goby/compiler/lexer"
	"github.com/goby-lang/goby/compiler/parser"
	"github.com/goby-lang/goby/compiler/token"
	"github.com/goby-lang/goby/vm"
	"github.com/goby-lang/goby/vm/classes"
	"github.com/goby-lang/goby/vm/errors"
	"strings"
)

// Ripper is a loadable library and has abilities to parse/lex/tokenize/get instructions of Goby codes from String.
// The library would be convenient for validating Goby codes when building lint tools,
// as well as the tests for Goby's compiler.
// For now, Ripper is a class and has only class methods, but I think this should finally be a 'newable' module with more sophisticated instance methods.

// Imported objects from vm
type Object = vm.Object
type VM = vm.VM
type Thread = vm.Thread
type Method = vm.Method
type StringObject = vm.StringObject
type HashObject = vm.HashObject
type ArrayObject = vm.ArrayObject

// Class methods --------------------------------------------------------

// Returns the list of instruction code generated from Goby code.
// Returns `[]` when the Goby code is invalid.
// The return value is a "tuple" style nested array:
// - `Array`: contains an instruction set
//   - `arg_types:` (none if `nil`)
// 		 - `names:` array of names (string)
//     - `types:` array of types (integer)
//   - `instructions:` array of instructions
//     - `action:` string
//     - `anchor:` integer
//     - `line:` integer
//     - `params:` array of parameters (string)
//     - `source_line:` integer
//     - `arg_set:` (none if `nil`)
//   		 - `names:` array of names (string)
//       - `types:` array of types (integer)
//
// ```ruby
// require 'ripper'; Ripper.instruction "10.times do |i| puts i end"
// #=>
//
// require 'ripper'; Ripper.instruction "10.times do |i| puts i" # the code is invalid
// #=> []
// ```
//
// @param Goby code [String]
// @return [Array]
func instruction(receiver Object, sourceLine int, t *Thread, args []Object) Object {
	if len(args) != 1 {
		return t.VM().InitErrorObject(errors.ArgumentError, sourceLine, "Expect 1 argument. got=%d", len(args))
	}

	arg, ok := args[0].(*StringObject)
	if !ok {
		return t.VM().InitErrorObject(errors.TypeError, sourceLine, errors.WrongArgumentTypeFormat, classes.StringClass, args[0].Class().Name)
	}

	i, _ := compiler.CompileToInstructions(arg.Value().(string), parser.NormalMode)

	return convertToTuple(i, t.VM())
}

// Returns a nested array that contains the line #, type of the tokenize, and the literal of the tokenize.
// Note that the class method does not return any errors even though the provided Goby code is invalid.
//
// ```ruby
// require 'ripper'; Ripper.lex "10.times do |i| puts i end"
// #=> [[0, "on_int", "10"], [0, "on_dot", "."], [0, "on_ident", "times"], [0, "on_do", "do"], [0, "on_bar", "|"], [0, "on_ident", "i"], [0, "on_bar", "|"], [0, "on_ident", "puts"], [0, "on_ident", "i"], [0, "on_end", "end"], [0, "on_eof", ""]]
//
// require 'ripper'; Ripper.lex "10.times do |i| puts i" # the code is invalid
// #=> [[0, "on_int", "10"], [0, "on_dot", "."], [0, "on_ident", "times"], [0, "on_do", "do"], [0, "on_bar", "|"], [0, "on_ident", "i"], [0, "on_bar", "|"], [0, "on_ident", "puts"], [0, "on_ident", "i"], [0, "on_eof", ""]]
// ```
//
// @param Goby code [String]
// @return [Array]
func lex(receiver Object, sourceLine int, t *Thread, args []Object) Object  {
	if len(args) != 1 {
		return t.VM().InitErrorObject(errors.ArgumentError, sourceLine, "Expect 1 argument. got=%d", len(args))
	}

	arg, ok := args[0].(*StringObject)
	if !ok {
		return t.VM().InitErrorObject(errors.TypeError, sourceLine, errors.WrongArgumentTypeFormat, classes.StringClass, args[0].Class().Name)
	}

	l := lexer.New(arg.Value().(string))
	array := t.VM().InitArrayObject([]Object{})
	var elements []Object
	var nextToken token.Token
	for i := 0; ; i++ {
		nextToken = l.NextToken()
		elements = append(elements, t.VM().InitIntegerObject(nextToken.Line))
		elements = append(elements, t.VM().InitStringObject(convertLex(nextToken.Type)))
		elements = append(elements, t.VM().InitStringObject(nextToken.Literal))
		array.Elements = append(array.Elements, t.VM().InitArrayObject(elements))
		if nextToken.Type == token.EOF {
			break
		}
		elements = nil
	}
	return array
}

// Just to disable creating instances.
func new(receiver Object, sourceLine int, t *Thread, args []Object) Object {
	return t.VM().InitUnsupportedMethodError(sourceLine, "#new", receiver)
}

// Returns the parsed Goby codes as a String object.
// Returns an error when the code is invalid.
//
// ```ruby
// require 'ripper'; Ripper.parse "10.times do |i| puts i end"
// #=> "10.times() do |i|
// #=> self.puts(i)
// #=> end"
//
// require 'ripper'; Ripper.parse "10.times do |i| puts i" # the code is invalid
// #=> TypeError: InternalError%!(EXTRA string=String, string=Invalid Goby code)
// ```
//
// @param Goby code [String]
// @return [String]
func parse(receiver Object, sourceLine int, t *Thread, args []Object) Object {
	if len(args) != 1 {
		return t.VM().InitErrorObject(errors.ArgumentError, sourceLine, "Expect 1 argument. got=%d", len(args))
	}

	arg, ok := args[0].(*StringObject)
	if !ok {
		return t.VM().InitErrorObject(errors.TypeError, sourceLine, errors.WrongArgumentTypeFormat, classes.StringClass, args[0].Class().Name)
	}

	l := lexer.New(arg.Value().(string))
	p := parser.New(l)
	program, _ := p.ParseProgram()

	return t.VM().InitStringObject(program.String())
}

// Returns a tokenized Goby codes as an Array object.
// Note that this does not return any errors even though the provided code is invalid.
//
// ```ruby
// require 'ripper'; Ripper.tokenize "10.times do |i| puts i end"
// #=> ["10", ".", "times", "do", "|", "i", "|", "puts", "i", "end", "EOF"]
//
// require 'ripper'; Ripper.parse "10.times do |i| puts i" # the code is invalid
// #=> ["10", ".", "times", "do", "|", "i", "|", "puts", "i", "EOF"]
// ```
//
// @param Goby code [String]
// @return [String]
func tokenize(receiver Object, sourceLine int, t *Thread, args []Object) Object {
	if len(args) != 1 {
		return t.VM().InitErrorObject(errors.ArgumentError, sourceLine, "Expect 1 argument. got=%d", len(args))
	}

	arg, ok := args[0].(*StringObject)
	if !ok {
		return t.VM().InitErrorObject(errors.TypeError, sourceLine, errors.WrongArgumentTypeFormat, classes.StringClass, args[0].Class().Name)
	}

	l := lexer.New(arg.Value().(string))
	el := []Object{}
	var nt token.Token
	for i := 0; ; i++ {
		nt = l.NextToken()
		if nt.Type == token.EOF {
			el = append(el, t.VM().InitStringObject("EOF"))
			break
		}
		el = append(el, t.VM().InitStringObject(nt.Literal))
	}
	return t.VM().InitArrayObject(el)
}

// Internal functions ===================================================
func init() {
	vm.RegisterExternalClass("ripper", vm.ExternalClass("Ripper", "ripper.gb",
		// class methods
		map[string]vm.Method{
			"instruction": instruction,
			"lex":         lex,
			"new": 				 new,
			"parse":       parse,
			"tokenize":    tokenize,
		},
		// instance methods
		map[string]vm.Method{},
	))
}

// Other helper functions ----------------------------------------------

func convertToTuple(instSet []*bytecode.InstructionSet, v *VM) *ArrayObject {
	ary := []Object{}
	for _, instruction := range instSet {
		hashInstLevel1 := make(map[string]Object)
		hashInstLevel1["name"] = v.InitStringObject(instruction.Name())
		hashInstLevel1["type"] = v.InitStringObject(instruction.Type())
		if instruction.ArgTypes() != nil {
			hashInstLevel1["arg_types"] = getArgNameType(instruction.ArgTypes(), v)
		}
		ary = append(ary, v.InitHashObject(hashInstLevel1))

		arrayInst := []Object{}
		for _, ins := range instruction.Instructions {
			hashInstLevel2 := make(map[string]Object)
			hashInstLevel2["action"] = v.InitStringObject(ins.Action)
			hashInstLevel2["line"] = v.InitIntegerObject(ins.Line())
			hashInstLevel2["source_line"] = v.InitIntegerObject(ins.SourceLine())
			anchor, _ := ins.AnchorLine()
			hashInstLevel2["anchor"] = v.InitIntegerObject(anchor)

			arrayParams := []Object{}
			for _, param := range ins.Params {
				arrayParams = append(arrayParams, v.InitStringObject(param))
			}
			hashInstLevel2["params"] = v.InitArrayObject(arrayParams)

			if ins.ArgSet != nil {
				hashInstLevel1["arg_set"] = getArgNameType(ins.ArgSet, v)
			}

			arrayInst = append(arrayInst, v.InitHashObject(hashInstLevel2))
		}

		hashInstLevel1["instructions"] = v.InitArrayObject(arrayInst)
		ary = append(ary, v.InitHashObject(hashInstLevel1))
	}
	return v.InitArrayObject(ary)
}

func getArgNameType(argSet *bytecode.ArgSet, v *VM) *HashObject {
	h := make(map[string]Object)

	aName := []Object{}
	for _, argname := range argSet.Names() {
		aName = append(aName, v.InitStringObject(argname))
	}
	h["names"] = v.InitArrayObject(aName)

	aType := []Object{}
	for _, argtype := range argSet.Types() {
		aType = append(aType, v.InitIntegerObject(argtype))
	}

	h["types"] = v.InitArrayObject(aType)
	return v.InitHashObject(h)
}

// TODO: This should finally be auto-generated from tokenize.go
func convertLex(t token.Type) string {
	var s string
	switch t {
	case token.Asterisk:
		s = "asterisk"
	case token.And:
		s = "and"
	case token.Assign:
		s = "assign"
	case token.Bang:
		s = "bang"
	case token.Bar:
		s = "bar"
	case token.Colon:
		s = "colon"
	case token.Comma:
		s = "comma"
	case token.COMP:
		s = "comp"
	case token.Dot:
		s = "dot"
	case token.Eq:
		s = "eq"
	case token.GT:
		s = "gt"
	case token.GTE:
		s = "gte"
	case token.LBrace:
		s = "lbrace"
	case token.LBracket:
		s = "lbracket"
	case token.LParen:
		s = "lparen"
	case token.LT:
		s = "lt"
	case token.LTE:
		s = "lte"
	case token.Match:
		s = "match"
	case token.Minus:
		s = "minus"
	case token.MinusEq:
		s = "minuseq"
	case token.Modulo:
		s = "modulo"
	case token.NotEq:
		s = "noteq"
	case token.Or:
		s = "or"
	case token.OrEq:
		s = "oreq"
	case token.Plus:
		s = "plus"
	case token.PlusEq:
		s = "pluseq"
	case token.Pow:
		s = "pow"
	case token.Range:
		s = "range"
	case token.RBrace:
		s = "rbrace"
	case token.RBracket:
		s = "rbracket"
	case token.ResolutionOperator:
		s = "resolutionoperator"
	case token.RParen:
		s = "rparen"
	case token.Semicolon:
		s = "semicolon"
	case token.Slash:
		s = "slash"
	default:
		s = strings.ToLower(string(t))
	}

	return "on_" + s
}
<|endoftext|>


<commit_before>package sockjsclient

// http://sockjs.github.io/sockjs-protocol/sockjs-protocol-0.3.3.html

import (
	crand "crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/gorilla/websocket"
)

func ConnectWebsocketSession(baseURL string) (*WebsocketSession, error) {
	dialURL, err := url.Parse(baseURL)
	if err != nil {
		return nil, err
	}

	// will be used to set the origin header
	originalScheme := dialURL.Scheme

	if err := replaceSchemeWithWS(dialURL); err != nil {
		return nil, err
	}

	if err := addMissingPortAndSlash(dialURL); err != nil {
		return nil, err
	}

	serverID := threeDigits()
	sessionID := randomStringLength(20)

	// Add server_id and session_id to the path.
	dialURL.Path += serverID + "/" + sessionID + "/websocket"

	requestHeader := http.Header{}
	requestHeader.Add("Origin", originalScheme+"://"+dialURL.Host)

	conn, _, err := websocket.DefaultDialer.Dial(dialURL.String(), requestHeader)
	if err != nil {
		return nil, err
	}

	session := NewWebsocketSession(conn)
	session.id = sessionID
	return session, nil
}

type WebsocketSession struct {
	conn     *websocket.Conn
	id       string
	messages []string
}

func NewWebsocketSession(conn *websocket.Conn) *WebsocketSession {
	return &WebsocketSession{
		conn: conn,
	}
}

func (w *WebsocketSession) RemoteAddr() string {
	return w.conn.RemoteAddr().String()
}

// ID returns a session id
func (w *WebsocketSession) ID() string {
	return w.id
}

// Recv reads one text frame from session
func (w *WebsocketSession) Recv() (string, error) {
	// Return previously received messages if there is any.
	if len(w.messages) > 0 {
		msg := w.messages[0]
		w.messages = w.messages[1:]
		return msg, nil
	}

read_frame:
	// Read one SockJS frame.
	_, buf, err := w.conn.ReadMessage()
	if err != nil {
		return "", err
	}

	if len(buf) == 0 {
		return "", errors.New("unexpected empty message")
	}

	frameType := buf[0]
	data := buf[1:]

	switch frameType {
	case 'o':
		// TODO handle open
		goto read_frame
	case 'a':
		var messages []string
		err = json.Unmarshal(data, &messages)
		if err != nil {
			return "", err
		}
		w.messages = append(w.messages, messages...)
	case 'm':
		var message string
		err = json.Unmarshal(data, &message)
		if err != nil {
			return "", err
		}
		w.messages = append(w.messages, message)
	case 'c':
		return "", errors.New("session closed")
	case 'h':
		// TODO handle heartbeat
		goto read_frame
	default:
		return "", errors.New("invalid frame type")
	}

	// Return first message in slice.
	if len(w.messages) == 0 {
		return "", errors.New("no message")
	}
	msg := w.messages[0]
	w.messages = w.messages[1:]
	return msg, nil
}

// Send sends one text frame to session
func (w *WebsocketSession) Send(str string) error {
	b, _ := json.Marshal([]string{str})
	return w.conn.WriteMessage(websocket.TextMessage, b)
}

// Close closes the session with provided code and reason.
func (w *WebsocketSession) Close(status uint32, reason string) error {
	return w.conn.Close()
}

var r = rand.New(rand.NewSource(time.Now().UnixNano()))

// threeDigits is used to generate a server_id.
func threeDigits() string {
	var i uint64
	i = uint64(r.Int31())
	if i < 100 {
		i += 100
	}
	return strconv.FormatUint(i, 10)[:3]
}

// randomStringLength is used to generate a session_id.
func randomStringLength(length int) string {
	size := (length * 6 / 8) + 1
	r := make([]byte, size)
	crand.Read(r)
	return base64.URLEncoding.EncodeToString(r)[:length]
}

func replaceSchemeWithWS(u *url.URL) error {
	switch u.Scheme {
	case "http":
		u.Scheme = "ws"
	case "https":
		u.Scheme = "wss"
	default:
		return fmt.Errorf("invalid scheme in url: %s", u.Scheme)
	}
	return nil
}

// addMissingPortAndSlash appends 80 or 443 depending on the scheme
// if there is no port number in the URL.
// Also it adds "/" to the end of path if path does not ends with "/".
func addMissingPortAndSlash(u *url.URL) error {
	_, _, err := net.SplitHostPort(u.Host)
	if err != nil {
		if missingPortErr, ok := err.(*net.AddrError); ok && missingPortErr.Err == "missing port in address" {
			var port string
			switch u.Scheme {
			case "ws":
				port = "80"
			case "wss":
				port = "443"
			default:
				return fmt.Errorf("unknown scheme: %s", u.Scheme)
			}
			u.Host = net.JoinHostPort(strings.TrimRight(missingPortErr.Addr, ":"), port)
		} else {
			return err
		}
	}

	if u.Path == "" || u.Path[len(u.Path)-1:] != "/" {
		u.Path += "/"
	}

	return nil
}
<commit_msg>socksclient: check marshal errors<commit_after>package sockjsclient

// http://sockjs.github.io/sockjs-protocol/sockjs-protocol-0.3.3.html

import (
	crand "crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/gorilla/websocket"
)

func ConnectWebsocketSession(baseURL string) (*WebsocketSession, error) {
	dialURL, err := url.Parse(baseURL)
	if err != nil {
		return nil, err
	}

	// will be used to set the origin header
	originalScheme := dialURL.Scheme

	if err := replaceSchemeWithWS(dialURL); err != nil {
		return nil, err
	}

	if err := addMissingPortAndSlash(dialURL); err != nil {
		return nil, err
	}

	serverID := threeDigits()
	sessionID := randomStringLength(20)

	// Add server_id and session_id to the path.
	dialURL.Path += serverID + "/" + sessionID + "/websocket"

	requestHeader := http.Header{}
	requestHeader.Add("Origin", originalScheme+"://"+dialURL.Host)

	conn, _, err := websocket.DefaultDialer.Dial(dialURL.String(), requestHeader)
	if err != nil {
		return nil, err
	}

	session := NewWebsocketSession(conn)
	session.id = sessionID
	return session, nil
}

type WebsocketSession struct {
	conn     *websocket.Conn
	id       string
	messages []string
}

func NewWebsocketSession(conn *websocket.Conn) *WebsocketSession {
	return &WebsocketSession{
		conn: conn,
	}
}

func (w *WebsocketSession) RemoteAddr() string {
	return w.conn.RemoteAddr().String()
}

// ID returns a session id
func (w *WebsocketSession) ID() string {
	return w.id
}

// Recv reads one text frame from session
func (w *WebsocketSession) Recv() (string, error) {
	// Return previously received messages if there is any.
	if len(w.messages) > 0 {
		msg := w.messages[0]
		w.messages = w.messages[1:]
		return msg, nil
	}

read_frame:
	// Read one SockJS frame.
	_, buf, err := w.conn.ReadMessage()
	if err != nil {
		return "", err
	}

	if len(buf) == 0 {
		return "", errors.New("unexpected empty message")
	}

	frameType := buf[0]
	data := buf[1:]

	switch frameType {
	case 'o':
		// TODO handle open
		goto read_frame
	case 'a':
		var messages []string
		err = json.Unmarshal(data, &messages)
		if err != nil {
			return "", err
		}
		w.messages = append(w.messages, messages...)
	case 'm':
		var message string
		err = json.Unmarshal(data, &message)
		if err != nil {
			return "", err
		}
		w.messages = append(w.messages, message)
	case 'c':
		return "", errors.New("session closed")
	case 'h':
		// TODO handle heartbeat
		goto read_frame
	default:
		return "", errors.New("invalid frame type")
	}

	// Return first message in slice.
	if len(w.messages) == 0 {
		return "", errors.New("no message")
	}
	msg := w.messages[0]
	w.messages = w.messages[1:]
	return msg, nil
}

// Send sends one text frame to session
func (w *WebsocketSession) Send(str string) error {
	b, err := json.Marshal([]string{str})
	if err != nil {
		return err
	}

	return w.conn.WriteMessage(websocket.TextMessage, b)
}

// Close closes the session with provided code and reason.
func (w *WebsocketSession) Close(status uint32, reason string) error {
	return w.conn.Close()
}

var r = rand.New(rand.NewSource(time.Now().UnixNano()))

// threeDigits is used to generate a server_id.
func threeDigits() string {
	var i uint64
	i = uint64(r.Int31())
	if i < 100 {
		i += 100
	}
	return strconv.FormatUint(i, 10)[:3]
}

// randomStringLength is used to generate a session_id.
func randomStringLength(length int) string {
	size := (length * 6 / 8) + 1
	r := make([]byte, size)
	crand.Read(r)
	return base64.URLEncoding.EncodeToString(r)[:length]
}

func replaceSchemeWithWS(u *url.URL) error {
	switch u.Scheme {
	case "http":
		u.Scheme = "ws"
	case "https":
		u.Scheme = "wss"
	default:
		return fmt.Errorf("invalid scheme in url: %s", u.Scheme)
	}
	return nil
}

// addMissingPortAndSlash appends 80 or 443 depending on the scheme
// if there is no port number in the URL.
// Also it adds "/" to the end of path if path does not ends with "/".
func addMissingPortAndSlash(u *url.URL) error {
	_, _, err := net.SplitHostPort(u.Host)
	if err != nil {
		if missingPortErr, ok := err.(*net.AddrError); ok && missingPortErr.Err == "missing port in address" {
			var port string
			switch u.Scheme {
			case "ws":
				port = "80"
			case "wss":
				port = "443"
			default:
				return fmt.Errorf("unknown scheme: %s", u.Scheme)
			}
			u.Host = net.JoinHostPort(strings.TrimRight(missingPortErr.Addr, ":"), port)
		} else {
			return err
		}
	}

	if u.Path == "" || u.Path[len(u.Path)-1:] != "/" {
		u.Path += "/"
	}

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package legacyregistry

import (
	"net/http"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"

	"k8s.io/component-base/metrics"
)

var (
	defaultRegistry = metrics.NewKubeRegistry()
	// DefaultGatherer exposes the global registry gatherer
	DefaultGatherer metrics.Gatherer = defaultRegistry
	// Reset calls reset on the global registry
	Reset = defaultRegistry.Reset
	// MustRegister registers registerable metrics but uses the global registry.
	MustRegister = defaultRegistry.MustRegister
	// RawMustRegister registers prometheus collectors but uses the global registry, this
	// bypasses the metric stability framework
	//
	// Deprecated
	RawMustRegister = defaultRegistry.RawMustRegister

	// Register registers a collectable metric but uses the global registry
	Register = defaultRegistry.Register
)

func init() {
	//nolint:staticcheck // SA1019 - replacement function still calls prometheus.NewProcessCollector().
	RawMustRegister(prometheus.NewProcessCollector(prometheus.ProcessCollectorOpts{}))
	//nolint:staticcheck // SA1019 - replacement function still calls prometheus.NewGoCollector().
	RawMustRegister(prometheus.NewGoCollector())
}

// Handler returns an HTTP handler for the DefaultGatherer. It is
// already instrumented with InstrumentHandler (using "prometheus" as handler
// name).
func Handler() http.Handler {
	return promhttp.InstrumentMetricHandler(prometheus.DefaultRegisterer, promhttp.HandlerFor(defaultRegistry, promhttp.HandlerOpts{}))
}

// HandlerWithReset returns an HTTP handler for the DefaultGatherer but invokes
// registry reset if the http method is DELETE.
func HandlerWithReset() http.Handler {
	return promhttp.InstrumentMetricHandler(
		prometheus.DefaultRegisterer,
		metrics.HandlerWithReset(defaultRegistry, metrics.HandlerOpts{}))
}

// CustomRegister registers a custom collector but uses the global registry.
func CustomRegister(c metrics.StableCollector) error {
	err := defaultRegistry.CustomRegister(c)

	//TODO(RainbowMango): Maybe we can wrap this error by error wrapping.(Golang 1.13)
	_ = prometheus.Register(c)

	return err
}

// CustomMustRegister registers custom collectors but uses the global registry.
func CustomMustRegister(cs ...metrics.StableCollector) {
	defaultRegistry.CustomMustRegister(cs...)

	for _, c := range cs {
		prometheus.MustRegister(c)
	}
}
<commit_msg>Enable new set of Go metrics in Kubernetes<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package legacyregistry

import (
	"net/http"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/collectors"
	"github.com/prometheus/client_golang/prometheus/promhttp"

	"k8s.io/component-base/metrics"
)

var (
	defaultRegistry = metrics.NewKubeRegistry()
	// DefaultGatherer exposes the global registry gatherer
	DefaultGatherer metrics.Gatherer = defaultRegistry
	// Reset calls reset on the global registry
	Reset = defaultRegistry.Reset
	// MustRegister registers registerable metrics but uses the global registry.
	MustRegister = defaultRegistry.MustRegister
	// RawMustRegister registers prometheus collectors but uses the global registry, this
	// bypasses the metric stability framework
	//
	// Deprecated
	RawMustRegister = defaultRegistry.RawMustRegister

	// Register registers a collectable metric but uses the global registry
	Register = defaultRegistry.Register
)

func init() {
	RawMustRegister(collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}))
	RawMustRegister(collectors.NewGoCollector(collectors.WithGoCollectorRuntimeMetrics(collectors.MetricsAll)))
}

// Handler returns an HTTP handler for the DefaultGatherer. It is
// already instrumented with InstrumentHandler (using "prometheus" as handler
// name).
func Handler() http.Handler {
	return promhttp.InstrumentMetricHandler(prometheus.DefaultRegisterer, promhttp.HandlerFor(defaultRegistry, promhttp.HandlerOpts{}))
}

// HandlerWithReset returns an HTTP handler for the DefaultGatherer but invokes
// registry reset if the http method is DELETE.
func HandlerWithReset() http.Handler {
	return promhttp.InstrumentMetricHandler(
		prometheus.DefaultRegisterer,
		metrics.HandlerWithReset(defaultRegistry, metrics.HandlerOpts{}))
}

// CustomRegister registers a custom collector but uses the global registry.
func CustomRegister(c metrics.StableCollector) error {
	err := defaultRegistry.CustomRegister(c)

	//TODO(RainbowMango): Maybe we can wrap this error by error wrapping.(Golang 1.13)
	_ = prometheus.Register(c)

	return err
}

// CustomMustRegister registers custom collectors but uses the global registry.
func CustomMustRegister(cs ...metrics.StableCollector) {
	defaultRegistry.CustomMustRegister(cs...)

	for _, c := range cs {
		prometheus.MustRegister(c)
	}
}
<|endoftext|>