<commit_before>package linker

import(
	. "tritium/proto"
	proto "goprotobuf.googlecode.com/hg/proto"
	"log"
	"fmt"
	//packager "tritium/packager"
)

type FuncMap map[string]int;

type LinkingContext struct {
	objMap map[string]int
	funList []FuncMap
	textType int
	*Transform
}

type LocalDef map[string]int

func NewObjectLinkingContext(pkg *Package, objs []*ScriptObject) (*LinkingContext) {
	// Setup object lookup map!
	objScriptLookup := make(map[string]int, len(objs))
	for index, obj := range(objs) {
		objScriptLookup[proto.GetString(obj.Name)] = index
	}
	ctx := NewLinkingContext(pkg)
	ctx.objMap = objScriptLookup
	ctx.Objects = objs
	return ctx
}

func NewLinkingContext(pkg *Package) (*LinkingContext){
	// Setup the function map!
	functionLookup := make([]FuncMap, len(pkg.Types))
	for typeId, typeObj := range(pkg.Types) {
		funcMap := make(FuncMap)
		//println("Type:",proto.GetString(typeObj.Name))
		//println("Implements:", proto.GetInt32(typeObj.Implements))
		implements := functionLookup[proto.GetInt32(typeObj.Implements)]
		for index, fun := range(pkg.Functions) {
			stub := fun.Stub()
			funScopeId := proto.GetInt32(fun.ScopeTypeId)
			inherited := false
			if implements != nil {
				_, inherited = implements[stub]
			}
			if (funScopeId == int32(typeId)) || inherited {
				//println(proto.GetString(typeObj.Name), ":", stub)
				funcMap[stub] = index
			}
		}
		functionLookup[typeId] = funcMap
	}

	// Setup the main context object
	ctx := &LinkingContext{
		funList: functionLookup,
		Transform: &Transform{
			Pkg: pkg,
		},
	}

	// Find Text type int -- need its ID to deal with Text literals during processing
	ctx.textType = pkg.GetTypeId("Text")

	return ctx
}

func (ctx *LinkingContext) Link() {
	ctx.link(0, ctx.Pkg.GetTypeId("Text"))
}

func (ctx *LinkingContext) link(objId, scopeType int) {
	obj := ctx.Objects[objId]
	//println(obj.String())
	if proto.GetBool(obj.Linked) == false {
		//println("Linking", proto.GetString(obj.Name))
		obj.ScopeTypeId = proto.Int(scopeType)
		ctx.ProcessInstruction(obj.Root, scopeType)
	} else {
		if scopeType != int(proto.GetInt32(obj.ScopeTypeId)) {
			log.Panic("Imported a script in two different scopes!")
		}
		//println("Already linked", proto.GetString(obj.Name))
	}
}

func (ctx *LinkingContext) ProcessInstruction(ins *Instruction, scopeType int) (returnType int) {
	localScope := make(LocalDef, 0)
	return ctx.ProcessInstructionWithLocalScope(ins, scopeType, localScope)
}

func (ctx *LinkingContext) ProcessInstructionWithLocalScope(ins *Instruction, scopeType int, localScope LocalDef) (returnType int) {
	returnType = -1
	switch *ins.Type {
		case Instruction_IMPORT:
			// set its import_id and blank the value field
			importValue := proto.GetString(ins.Value)
			importId, ok := ctx.objMap[importValue]
			if ok != true {
				log.Panic("Invalid import ", ins.String())
			}
			// Make sure this object is linked with the right scopeType
			ctx.link(importId, scopeType)
			//println("befor", ins.String())
			ins.ObjectId = proto.Int(importId)
			ins.Value = nil
			//println("after", ins.String())
		case Instruction_LOCAL_VAR:
			name := proto.GetString(ins.Value)
			typeId, found := localScope[name]
			if found {
				returnType = typeId
				if len(ins.Arguments) > 0 {
					log.Panic("The local variable %", name, " has been assigned before and cannot be reassigned!")
				}
			} else {
				if len(ins.Arguments) > 0 {
					// We are going to assign something to this variable
					returnType = ctx.ProcessInstructionWithLocalScope(ins.Arguments[0], scopeType, localScope)
					// Duplicate the localScope before we go messing with my parents scope
					parentScope := localScope
					localScope = make(LocalDef, len(parentScope))
					for s, t := range(parentScope) {
						localScope[s] = t
					}
					localScope[proto.GetString(ins.Value)] = returnType
				} else {
					log.Panic("I've never seen the variable %", name, " before! Please assign a value before usage.")
				}
			}
		case Instruction_FUNCTION_CALL:
			stub := proto.GetString(ins.Value)
			if ins.Arguments != nil {
				for _, arg := range(ins.Arguments) {
					//println("arg:", arg)
					//println("scopeType:", scopeType)
					//println("localScope", localScope)
					argReturn := ctx.ProcessInstructionWithLocalScope(arg, scopeType, localScope)
					if argReturn == -1 {
						log.Panic("Invalid argument object", arg.String())
					}
					stub = stub + "," + fmt.Sprintf("%d", argReturn)
				}
			}
			funcId, ok := ctx.funList[scopeType][stub]
			if ok != true {
				println("Available functions...")
				for funcName, _ := range(ctx.funList[scopeType]) {
					println(funcName)
				}
				log.Panic("No such function found....", ins.String(), "with the stub: ",scopeType, stub)
			}
			ins.FunctionId = proto.Int32(int32(funcId))
			fun := ctx.Pkg.Functions[funcId]
			returnType = int(proto.GetInt32(fun.ReturnTypeId))
			opensScopeType := int(proto.GetInt32(fun.OpensTypeId))
			//println("Zomg, found function", fun.String())
			//println("I open a Scope of type ", opensScopeType)

			if ins.Children != nil {
				for _, child := range(ins.Children) {
					ctx.ProcessInstructionWithLocalScope(child, opensScopeType, localScope)
				}
			}
		case Instruction_TEXT:
			returnType = ctx.textType
		case Instruction_BLOCK:
			if ins.Children != nil {
				for _, child := range(ins.Children) {
					returnType = ctx.ProcessInstructionWithLocalScope(child, scopeType, localScope)
				}
			}
	}

	// if function
		// Figure out function signature (name + arg types)
			// have to start at the bottom of the tree (args first) and check types.
		// Is this a real function?
			// aka, text(regexp()) ... have to see that regexp returns Regexp object,
			// which, then, when we go to process text() we notice we don't have a text(Regexp)
			// function, so we need to throw a reasonable error
			// Hrrrm.... need line numbers, huh?
		// Set the function_id if it is real, error otherwise
	return
}
<commit_msg>Add in better type reporting during linking.<commit_after>package linker

import(
	. "tritium/proto"
	proto "goprotobuf.googlecode.com/hg/proto"
	"log"
	"fmt"
	//packager "tritium/packager"
)

type FuncMap map[string]int;

type LinkingContext struct {
	objMap map[string]int
	funList []FuncMap
	textType int
	types []string
	*Transform
}

type LocalDef map[string]int

func NewObjectLinkingContext(pkg *Package, objs []*ScriptObject) (*LinkingContext) {
	// Setup object lookup map!
	objScriptLookup := make(map[string]int, len(objs))
	for index, obj := range(objs) {
		objScriptLookup[proto.GetString(obj.Name)] = index
	}
	ctx := NewLinkingContext(pkg)
	ctx.objMap = objScriptLookup
	ctx.Objects = objs
	return ctx
}

func NewLinkingContext(pkg *Package) (*LinkingContext){
	// Setup the function map!
	functionLookup := make([]FuncMap, len(pkg.Types))
	types := make([]string, len(pkg.Types))

	for typeId, typeObj := range(pkg.Types) {
		funcMap := make(FuncMap)
		types[typeId] = proto.GetString(typeObj.Name)
		//println("Type:",proto.GetString(typeObj.Name))
		//println("Implements:", proto.GetInt32(typeObj.Implements))
		implements := functionLookup[proto.GetInt32(typeObj.Implements)]
		for index, fun := range(pkg.Functions) {
			stub := fun.Stub()
			funScopeId := proto.GetInt32(fun.ScopeTypeId)
			inherited := false
			if implements != nil {
				_, inherited = implements[stub]
			}
			if (funScopeId == int32(typeId)) || inherited {
				//println(proto.GetString(typeObj.Name), ":", stub)
				funcMap[stub] = index
			}
		}
		functionLookup[typeId] = funcMap
	}

	// Setup the main context object
	ctx := &LinkingContext{
		funList: functionLookup,
		types: types,
		Transform: &Transform{
			Pkg: pkg,
		},
	}

	// Find Text type int -- need its ID to deal with Text literals during processing
	ctx.textType = pkg.GetTypeId("Text")

	return ctx
}

func (ctx *LinkingContext) Link() {
	ctx.link(0, ctx.Pkg.GetTypeId("Text"))
}

func (ctx *LinkingContext) link(objId, scopeType int) {
	obj := ctx.Objects[objId]
	//println(obj.String())
	if proto.GetBool(obj.Linked) == false {
		//println("Linking", proto.GetString(obj.Name))
		obj.ScopeTypeId = proto.Int(scopeType)
		ctx.ProcessInstruction(obj.Root, scopeType)
	} else {
		if scopeType != int(proto.GetInt32(obj.ScopeTypeId)) {
			log.Panic("Imported a script in two different scopes!")
		}
		//println("Already linked", proto.GetString(obj.Name))
	}
}

func (ctx *LinkingContext) ProcessInstruction(ins *Instruction, scopeType int) (returnType int) {
	localScope := make(LocalDef, 0)
	return ctx.ProcessInstructionWithLocalScope(ins, scopeType, localScope)
}

func (ctx *LinkingContext) ProcessInstructionWithLocalScope(ins *Instruction, scopeType int, localScope LocalDef) (returnType int) {
	returnType = -1
	switch *ins.Type {
		case Instruction_IMPORT:
			// set its import_id and blank the value field
			importValue := proto.GetString(ins.Value)
			importId, ok := ctx.objMap[importValue]
			if ok != true {
				log.Panic("Invalid import ", ins.String())
			}
			// Make sure this object is linked with the right scopeType
			ctx.link(importId, scopeType)
			//println("befor", ins.String())
			ins.ObjectId = proto.Int(importId)
			ins.Value = nil
			//println("after", ins.String())
		case Instruction_LOCAL_VAR:
			name := proto.GetString(ins.Value)
			typeId, found := localScope[name]
			if found {
				returnType = typeId
				if len(ins.Arguments) > 0 {
					log.Panic("The local variable %", name, " has been assigned before and cannot be reassigned!")
				}
			} else {
				if len(ins.Arguments) > 0 {
					// We are going to assign something to this variable
					returnType = ctx.ProcessInstructionWithLocalScope(ins.Arguments[0], scopeType, localScope)
					// Duplicate the localScope before we go messing with my parents scope
					parentScope := localScope
					localScope = make(LocalDef, len(parentScope))
					for s, t := range(parentScope) {
						localScope[s] = t
					}
					localScope[proto.GetString(ins.Value)] = returnType
				} else {
					log.Panic("I've never seen the variable %", name, " before! Please assign a value before usage.")
				}
			}
		case Instruction_FUNCTION_CALL:
			stub := proto.GetString(ins.Value)
			if ins.Arguments != nil {
				for _, arg := range(ins.Arguments) {
					//println("arg:", arg)
					//println("scopeType:", scopeType)
					//println("localScope", localScope)
					argReturn := ctx.ProcessInstructionWithLocalScope(arg, scopeType, localScope)
					if argReturn == -1 {
						log.Panic("Invalid argument object", arg.String())
					}
					stub = stub + "," + fmt.Sprintf("%d", argReturn)
				}
			}
			funcId, ok := ctx.funList[scopeType][stub]
			if ok != true {
				println("Available functions...")
				for funcName, _ := range(ctx.funList[scopeType]) {
					println(funcName)
				}
				println(scopeType)
				log.Panic("No such function found....", ins.String(), "with the scope: ", ctx.types[scopeType], " and stub ", stub)
			}
			ins.FunctionId = proto.Int32(int32(funcId))
			fun := ctx.Pkg.Functions[funcId]
			returnType = int(proto.GetInt32(fun.ReturnTypeId))
			opensScopeType := int(proto.GetInt32(fun.OpensTypeId))
			if opensScopeType == 0 {
				// If we're a Base scope, don't mess with texas!
				opensScopeType = scopeType
			}
			//println("Zomg, found function", fun.String())
			//println("I open a Scope of type ", opensScopeType)

			if ins.Children != nil {
				for _, child := range(ins.Children) {
					ctx.ProcessInstructionWithLocalScope(child, opensScopeType, localScope)
				}
			}
		case Instruction_TEXT:
			returnType = ctx.textType
		case Instruction_BLOCK:
			if ins.Children != nil {
				for _, child := range(ins.Children) {
					returnType = ctx.ProcessInstructionWithLocalScope(child, scopeType, localScope)
				}
			}
	}

	// if function
		// Figure out function signature (name + arg types)
			// have to start at the bottom of the tree (args first) and check types.
		// Is this a real function?
			// aka, text(regexp()) ... have to see that regexp returns Regexp object,
			// which, then, when we go to process text() we notice we don't have a text(Regexp)
			// function, so we need to throw a reasonable error
			// Hrrrm.... need line numbers, huh?
		// Set the function_id if it is real, error otherwise
	return
}
<|endoftext|>


<commit_before>package rest

import (
	"bytes"
	"http"
	"io"
	"net"
	"os"
)

type Client struct {
	conn     *http.ClientConn
	resource *http.URL
}

// resource is the url to the base of the resource (i.e., http://127.0.0.1:3000/snips/)
func NewClient(resource string) (*Client, os.Error) {
	var client = new(Client)
	var err os.Error

	// setup host
	if client.resource, err = http.ParseURL(resource); err != nil {
		return nil, err
	}

	// Setup conn
	var tcpConn net.Conn
	if tcpConn, err = net.Dial("tcp", "", client.resource.Host); err != nil {
		return nil, err
	}
	client.conn = http.NewClientConn(tcpConn, nil)

	return client, nil
}

func (client *Client) Close() {
	client.conn.Close()
}

func (client *Client) newRequest(method string, id string) (*http.Request, os.Error) {
	request := new(http.Request)
	var err os.Error

	request.ProtoMajor = 1
	request.ProtoMinor = 1
	request.TransferEncoding = []string{"chunked"}

	request.Method = method

	url := client.resource.String() + id
	if request.URL, err = http.ParseURL(url); err != nil {
		return nil, err
	}

	return request, nil
}

func (client *Client) Request(request *http.Request) (*http.Response, os.Error) {
	var err os.Error
	var response *http.Response

	if err = client.conn.Write(request); err != nil {
		return nil, err
	}

	if response, err = client.conn.Read(request); err != nil {
		return nil, err
	}

	return response, nil
}

// GET /resource/
func (client *Client) Index() (*http.Response, os.Error) {
	var request *http.Request
	var err os.Error

	if request, err = client.newRequest("GET", ""); err != nil {
		return nil, err
	}

	return client.Request(request)
}

// GET /resource/id
func (client *Client) Find(id string) (*http.Response, os.Error) {
	var request *http.Request
	var err os.Error

	if request, err = client.newRequest("GET", id); err != nil {
		return nil, err
	}

	return client.Request(request)
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() os.Error {
	return nil
}

// POST /resource
func (client *Client) Create(body string) (*http.Response, os.Error) {
	var request *http.Request
	var err os.Error

	if request, err = client.newRequest("POST", ""); err != nil {
		return nil, err
	}

	request.Body = nopCloser{bytes.NewBufferString(body)}

	return client.Request(request)
}

// PUT /resource/id
func (client *Client) Update(id string, body string) (*http.Response, os.Error) {
	var request *http.Request
	var err os.Error
	if request, err = client.newRequest("PUT", id); err != nil {
		return nil, err
	}

	request.Body = nopCloser{bytes.NewBufferString(body)}

	return client.Request(request)
}

func (client *Client) IdFromURL(urlString string) (string, os.Error) {
	var url *http.URL
	var err os.Error
	if url, err = http.ParseURL(urlString); err != nil {
		return "", err
	}

	return string(url.Path[len(client.resource.Path):]), nil
}

func (client *Client) Delete(id string) (*http.Response, os.Error) {
	var request *http.Request
	var err os.Error
	if request, err = client.newRequest("DELETE", id); err != nil {
		return nil, err
	}

	return client.Request(request)
}
<commit_msg>updated for net.Dial() change<commit_after>package rest

import (
	"bytes"
	"http"
	"io"
	"net"
	"os"
)

type Client struct {
	conn     *http.ClientConn
	resource *http.URL
}

// resource is the url to the base of the resource (i.e., http://127.0.0.1:3000/snips/)
func NewClient(resource string) (*Client, os.Error) {
	var client = new(Client)
	var err os.Error

	// setup host
	if client.resource, err = http.ParseURL(resource); err != nil {
		return nil, err
	}

	// Setup conn
	var tcpConn net.Conn
	if tcpConn, err = net.Dial("tcp", client.resource.Host); err != nil {
		return nil, err
	}
	client.conn = http.NewClientConn(tcpConn, nil)

	return client, nil
}

func (client *Client) Close() {
	client.conn.Close()
}

func (client *Client) newRequest(method string, id string) (*http.Request, os.Error) {
	request := new(http.Request)
	var err os.Error

	request.ProtoMajor = 1
	request.ProtoMinor = 1
	request.TransferEncoding = []string{"chunked"}

	request.Method = method

	url := client.resource.String() + id
	if request.URL, err = http.ParseURL(url); err != nil {
		return nil, err
	}

	return request, nil
}

func (client *Client) Request(request *http.Request) (*http.Response, os.Error) {
	var err os.Error
	var response *http.Response

	if err = client.conn.Write(request); err != nil {
		return nil, err
	}

	if response, err = client.conn.Read(request); err != nil {
		return nil, err
	}

	return response, nil
}

// GET /resource/
func (client *Client) Index() (*http.Response, os.Error) {
	var request *http.Request
	var err os.Error

	if request, err = client.newRequest("GET", ""); err != nil {
		return nil, err
	}

	return client.Request(request)
}

// GET /resource/id
func (client *Client) Find(id string) (*http.Response, os.Error) {
	var request *http.Request
	var err os.Error

	if request, err = client.newRequest("GET", id); err != nil {
		return nil, err
	}

	return client.Request(request)
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() os.Error {
	return nil
}

// POST /resource
func (client *Client) Create(body string) (*http.Response, os.Error) {
	var request *http.Request
	var err os.Error

	if request, err = client.newRequest("POST", ""); err != nil {
		return nil, err
	}

	request.Body = nopCloser{bytes.NewBufferString(body)}

	return client.Request(request)
}

// PUT /resource/id
func (client *Client) Update(id string, body string) (*http.Response, os.Error) {
	var request *http.Request
	var err os.Error
	if request, err = client.newRequest("PUT", id); err != nil {
		return nil, err
	}

	request.Body = nopCloser{bytes.NewBufferString(body)}

	return client.Request(request)
}

func (client *Client) IdFromURL(urlString string) (string, os.Error) {
	var url *http.URL
	var err os.Error
	if url, err = http.ParseURL(urlString); err != nil {
		return "", err
	}

	return string(url.Path[len(client.resource.Path):]), nil
}

func (client *Client) Delete(id string) (*http.Response, os.Error) {
	var request *http.Request
	var err os.Error
	if request, err = client.newRequest("DELETE", id); err != nil {
		return nil, err
	}

	return client.Request(request)
}
<|endoftext|>


<commit_before>package account

import (
	"context"
	"encoding/json"

	"github.com/bytom/blockchain/signers"
	"github.com/bytom/blockchain/txbuilder"
	"github.com/bytom/common"
	"github.com/bytom/consensus"
	"github.com/bytom/crypto/ed25519/chainkd"
	"github.com/bytom/errors"
	"github.com/bytom/protocol/bc"
	"github.com/bytom/protocol/bc/types"
	"github.com/bytom/protocol/vm/vmutil"
)

var (
	//chainTxUtxoNum maximum utxo quantity in a tx
	chainTxUtxoNum = 10
	//chainTxMergeGas chain tx gas
	chainTxMergeGas = uint64(10000000)
)

//DecodeSpendAction unmarshal JSON-encoded data of spend action
func (m *Manager) DecodeSpendAction(data []byte) (txbuilder.Action, error) {
	a := &spendAction{accounts: m}
	return a, json.Unmarshal(data, a)
}

type spendAction struct {
	accounts *Manager
	bc.AssetAmount
	AccountID      string `json:"account_id"`
	UseUnconfirmed bool   `json:"use_unconfirmed"`
}

func (a *spendAction) ActionType() string {
	return "spend_account"
}

// MergeSpendAction merge common assetID and accountID spend action
func MergeSpendAction(actions []txbuilder.Action) []txbuilder.Action {
	resultActions := []txbuilder.Action{}
	spendActionMap := make(map[string]*spendAction)

	for _, act := range actions {
		switch act := act.(type) {
		case *spendAction:
			actionKey := act.AssetId.String() + act.AccountID
			if tmpAct, ok := spendActionMap[actionKey]; ok {
				tmpAct.Amount += act.Amount
				tmpAct.UseUnconfirmed = tmpAct.UseUnconfirmed || act.UseUnconfirmed
			} else {
				spendActionMap[actionKey] = act
				resultActions = append(resultActions, act)
			}
		default:
			resultActions = append(resultActions, act)
		}
	}
	return resultActions
}

//calcMergeGas calculate the gas required that n utxos are merged into one
func calcMergeGas(num int) uint64 {
	gas := uint64(0)
	for num > 1 {
		gas += chainTxMergeGas
		num -= chainTxUtxoNum - 1
	}
	return gas
}

func (m *Manager) reserveBtmUtxoChain(builder *txbuilder.TemplateBuilder, accountID string, amount uint64, useUnconfirmed bool) ([]*UTXO, error) {
	reservedAmount := uint64(0)
	utxos := []*UTXO{}
	for gasAmount := uint64(0); reservedAmount < gasAmount+amount; gasAmount = calcMergeGas(len(utxos)) {
		reserveAmount := amount + gasAmount - reservedAmount
		res, err := m.utxoKeeper.Reserve(accountID, consensus.BTMAssetID, reserveAmount, useUnconfirmed, builder.MaxTime())
		if err != nil {
			return nil, err
		}

		builder.OnRollback(func() { m.utxoKeeper.Cancel(res.id) })
		reservedAmount += reserveAmount + res.change
		utxos = append(utxos, res.utxos[:]...)
	}
	return utxos, nil
}

func (m *Manager) buildBtmTxChain(utxos []*UTXO, signer *signers.Signer) ([]*txbuilder.Template, *UTXO, error) {
	if len(utxos) == 0 {
		return nil, nil, errors.New("mergeSpendActionUTXO utxos num 0")
	}

	tpls := []*txbuilder.Template{}
	if len(utxos) == 1 {
		return tpls, utxos[len(utxos)-1], nil
	}

	acp, err := m.GetLocalCtrlProgramByAddress(utxos[0].Address)
	if err != nil {
		return nil, nil, err
	}

	buildAmount := uint64(0)
	builder := &txbuilder.TemplateBuilder{}
	for index := 0; index < len(utxos); index++ {
		input, sigInst, err := UtxoToInputs(signer, utxos[index])
		if err != nil {
			return nil, nil, err
		}

		if err = builder.AddInput(input, sigInst); err != nil {
			return nil, nil, err
		}

		buildAmount += input.Amount()
		if builder.InputCount() != chainTxUtxoNum && index != len(utxos)-1 {
			continue
		}

		outAmount := buildAmount - chainTxMergeGas
		output := types.NewTxOutput(*consensus.BTMAssetID, outAmount, acp.ControlProgram)
		if err := builder.AddOutput(output); err != nil {
			return nil, nil, err
		}

		tpl, _, err := builder.Build()
		if err != nil {
			return nil, nil, err
		}

		bcOut, err := tpl.Transaction.Output(*tpl.Transaction.ResultIds[0])
		if err != nil {
			return nil, nil, err
		}

		utxos = append(utxos, &UTXO{
			OutputID:            *tpl.Transaction.ResultIds[0],
			AssetID:             *consensus.BTMAssetID,
			Amount:              outAmount,
			ControlProgram:      acp.ControlProgram,
			SourceID:            *bcOut.Source.Ref,
			SourcePos:           bcOut.Source.Position,
			ControlProgramIndex: acp.KeyIndex,
			Address:             acp.Address,
			Change:              acp.Change,
		})

		tpls = append(tpls, tpl)
		buildAmount = 0
		builder = &txbuilder.TemplateBuilder{}
		if index == len(utxos)-2 {
			break
		}
	}
	return tpls, utxos[len(utxos)-1], nil
}

// SpendAccountChain build the spend action with auto merge utxo function
func SpendAccountChain(ctx context.Context, builder *txbuilder.TemplateBuilder, action txbuilder.Action) ([]*txbuilder.Template, error) {
	act, ok := action.(*spendAction)
	if !ok {
		return nil, errors.New("fail to convert the spend action")
	}
	if *act.AssetId != *consensus.BTMAssetID {
		return nil, errors.New("spend chain action only support BTM")
	}

	utxos, err := act.accounts.reserveBtmUtxoChain(builder, act.AccountID, act.Amount, act.UseUnconfirmed)
	if err != nil {
		return nil, err
	}

	acct, err := act.accounts.FindByID(act.AccountID)
	if err != nil {
		return nil, err
	}

	tpls, utxo, err := act.accounts.buildBtmTxChain(utxos, acct.Signer)
	if err != nil {
		return nil, err
	}

	input, sigInst, err := UtxoToInputs(acct.Signer, utxo)
	if err != nil {
		return nil, err
	}

	if err := builder.AddInput(input, sigInst); err != nil {
		return nil, err
	}

	if utxo.Amount > act.Amount {
		if err = builder.AddOutput(types.NewTxOutput(*consensus.BTMAssetID, utxo.Amount-act.Amount, utxo.ControlProgram)); err != nil {
			return nil, errors.Wrap(err, "adding change output")
		}
	}
	return tpls, nil
}

func (a *spendAction) Build(ctx context.Context, b *txbuilder.TemplateBuilder) error {
	var missing []string
	if a.AccountID == "" {
		missing = append(missing, "account_id")
	}
	if a.AssetId.IsZero() {
		missing = append(missing, "asset_id")
	}
	if len(missing) > 0 {
		return txbuilder.MissingFieldsError(missing...)
	}

	acct, err := a.accounts.FindByID(a.AccountID)
	if err != nil {
		return errors.Wrap(err, "get account info")
	}

	res, err := a.accounts.utxoKeeper.Reserve(a.AccountID, a.AssetId, a.Amount, a.UseUnconfirmed, b.MaxTime())
	if err != nil {
		return errors.Wrap(err, "reserving utxos")
	}

	// Cancel the reservation if the build gets rolled back.
	b.OnRollback(func() { a.accounts.utxoKeeper.Cancel(res.id) })
	for _, r := range res.utxos {
		txInput, sigInst, err := UtxoToInputs(acct.Signer, r)
		if err != nil {
			return errors.Wrap(err, "creating inputs")
		}

		if err = b.AddInput(txInput, sigInst); err != nil {
			return errors.Wrap(err, "adding inputs")
		}
	}

	if res.change > 0 {
		acp, err := a.accounts.CreateAddress(a.AccountID, true)
		if err != nil {
			return errors.Wrap(err, "creating control program")
		}

		// Don't insert the control program until callbacks are executed.
		a.accounts.insertControlProgramDelayed(b, acp)
		if err = b.AddOutput(types.NewTxOutput(*a.AssetId, res.change, acp.ControlProgram)); err != nil {
			return errors.Wrap(err, "adding change output")
		}
	}
	return nil
}

//DecodeSpendUTXOAction unmarshal JSON-encoded data of spend utxo action
func (m *Manager) DecodeSpendUTXOAction(data []byte) (txbuilder.Action, error) {
	a := &spendUTXOAction{accounts: m}
	return a, json.Unmarshal(data, a)
}

type spendUTXOAction struct {
	accounts       *Manager
	OutputID       *bc.Hash                     `json:"output_id"`
	UseUnconfirmed bool                         `json:"use_unconfirmed"`
	Arguments      []txbuilder.ContractArgument `json:"arguments"`
}

func (a *spendUTXOAction) ActionType() string {
	return "spend_account_unspent_output"
}

func (a *spendUTXOAction) Build(ctx context.Context, b *txbuilder.TemplateBuilder) error {
	if a.OutputID == nil {
		return txbuilder.MissingFieldsError("output_id")
	}

	res, err := a.accounts.utxoKeeper.ReserveParticular(*a.OutputID, a.UseUnconfirmed, b.MaxTime())
	if err != nil {
		return err
	}

	b.OnRollback(func() { a.accounts.utxoKeeper.Cancel(res.id) })
	var accountSigner *signers.Signer
	if len(res.utxos[0].AccountID) != 0 {
		account, err := a.accounts.FindByID(res.utxos[0].AccountID)
		if err != nil {
			return err
		}

		accountSigner = account.Signer
	}

	txInput, sigInst, err := UtxoToInputs(accountSigner, res.utxos[0])
	if err != nil {
		return err
	}

	if a.Arguments == nil {
		return b.AddInput(txInput, sigInst)
	}

	sigInst = &txbuilder.SigningInstruction{}
	if err := txbuilder.AddContractArgs(sigInst, a.Arguments); err != nil {
		return err
	}

	return b.AddInput(txInput, sigInst)
}

// UtxoToInputs convert an utxo to the txinput
func UtxoToInputs(signer *signers.Signer, u *UTXO) (*types.TxInput, *txbuilder.SigningInstruction, error) {
	txInput := types.NewSpendInput(nil, u.SourceID, u.AssetID, u.Amount, u.SourcePos, u.ControlProgram)
	sigInst := &txbuilder.SigningInstruction{}
	if signer == nil {
		return txInput, sigInst, nil
	}

	path, err := signers.Path(signer, signers.AccountKeySpace, u.Change, u.ControlProgramIndex)
	if err != nil {
		return nil, nil, err
	}
	if u.Address == "" {
		sigInst.AddWitnessKeys(signer.XPubs, path, signer.Quorum)
		return txInput, sigInst, nil
	}

	address, err := common.DecodeAddress(u.Address, &consensus.ActiveNetParams)
	if err != nil {
		return nil, nil, err
	}

	sigInst.AddRawWitnessKeys(signer.XPubs, path, signer.Quorum)
	derivedXPubs := chainkd.DeriveXPubs(signer.XPubs, path)

	switch address.(type) {
	case *common.AddressWitnessPubKeyHash:
		derivedPK := derivedXPubs[0].PublicKey()
		sigInst.WitnessComponents = append(sigInst.WitnessComponents, txbuilder.DataWitness([]byte(derivedPK)))

	case *common.AddressWitnessScriptHash:
		derivedPKs := chainkd.XPubKeys(derivedXPubs)
		script, err := vmutil.P2SPMultiSigProgram(derivedPKs, signer.Quorum)
		if err != nil {
			return nil, nil, err
		}
		sigInst.WitnessComponents = append(sigInst.WitnessComponents, txbuilder.DataWitness(script))

	default:
		return nil, nil, errors.New("unsupport address type")
	}

	return txInput, sigInst, nil
}

// insertControlProgramDelayed takes a template builder and an account
// control program that hasn't been inserted to the database yet. It
// registers callbacks on the TemplateBuilder so that all of the template's
// account control programs are batch inserted if building the rest of
// the template is successful.
func (m *Manager) insertControlProgramDelayed(b *txbuilder.TemplateBuilder, acp *CtrlProgram) {
	m.delayedACPsMu.Lock()
	m.delayedACPs[b] = append(m.delayedACPs[b], acp)
	m.delayedACPsMu.Unlock()

	b.OnRollback(func() {
		m.delayedACPsMu.Lock()
		delete(m.delayedACPs, b)
		m.delayedACPsMu.Unlock()
	})
	b.OnBuild(func() error {
		m.delayedACPsMu.Lock()
		acps := m.delayedACPs[b]
		delete(m.delayedACPs, b)
		m.delayedACPsMu.Unlock()

		// Insert all of the account control programs at once.
		if len(acps) == 0 {
			return nil
		}
		return m.SaveControlPrograms(acps...)
	})
}
<commit_msg>edit the chain utxo number config<commit_after>package account

import (
	"context"
	"encoding/json"

	"github.com/bytom/blockchain/signers"
	"github.com/bytom/blockchain/txbuilder"
	"github.com/bytom/common"
	"github.com/bytom/consensus"
	"github.com/bytom/crypto/ed25519/chainkd"
	"github.com/bytom/errors"
	"github.com/bytom/protocol/bc"
	"github.com/bytom/protocol/bc/types"
	"github.com/bytom/protocol/vm/vmutil"
)

var (
	//chainTxUtxoNum maximum utxo quantity in a tx
	chainTxUtxoNum = 5
	//chainTxMergeGas chain tx gas
	chainTxMergeGas = uint64(10000000)
)

//DecodeSpendAction unmarshal JSON-encoded data of spend action
func (m *Manager) DecodeSpendAction(data []byte) (txbuilder.Action, error) {
	a := &spendAction{accounts: m}
	return a, json.Unmarshal(data, a)
}

type spendAction struct {
	accounts *Manager
	bc.AssetAmount
	AccountID      string `json:"account_id"`
	UseUnconfirmed bool   `json:"use_unconfirmed"`
}

func (a *spendAction) ActionType() string {
	return "spend_account"
}

// MergeSpendAction merge common assetID and accountID spend action
func MergeSpendAction(actions []txbuilder.Action) []txbuilder.Action {
	resultActions := []txbuilder.Action{}
	spendActionMap := make(map[string]*spendAction)

	for _, act := range actions {
		switch act := act.(type) {
		case *spendAction:
			actionKey := act.AssetId.String() + act.AccountID
			if tmpAct, ok := spendActionMap[actionKey]; ok {
				tmpAct.Amount += act.Amount
				tmpAct.UseUnconfirmed = tmpAct.UseUnconfirmed || act.UseUnconfirmed
			} else {
				spendActionMap[actionKey] = act
				resultActions = append(resultActions, act)
			}
		default:
			resultActions = append(resultActions, act)
		}
	}
	return resultActions
}

//calcMergeGas calculate the gas required that n utxos are merged into one
func calcMergeGas(num int) uint64 {
	gas := uint64(0)
	for num > 1 {
		gas += chainTxMergeGas
		num -= chainTxUtxoNum - 1
	}
	return gas
}

func (m *Manager) reserveBtmUtxoChain(builder *txbuilder.TemplateBuilder, accountID string, amount uint64, useUnconfirmed bool) ([]*UTXO, error) {
	reservedAmount := uint64(0)
	utxos := []*UTXO{}
	for gasAmount := uint64(0); reservedAmount < gasAmount+amount; gasAmount = calcMergeGas(len(utxos)) {
		reserveAmount := amount + gasAmount - reservedAmount
		res, err := m.utxoKeeper.Reserve(accountID, consensus.BTMAssetID, reserveAmount, useUnconfirmed, builder.MaxTime())
		if err != nil {
			return nil, err
		}

		builder.OnRollback(func() { m.utxoKeeper.Cancel(res.id) })
		reservedAmount += reserveAmount + res.change
		utxos = append(utxos, res.utxos[:]...)
	}
	return utxos, nil
}

func (m *Manager) buildBtmTxChain(utxos []*UTXO, signer *signers.Signer) ([]*txbuilder.Template, *UTXO, error) {
	if len(utxos) == 0 {
		return nil, nil, errors.New("mergeSpendActionUTXO utxos num 0")
	}

	tpls := []*txbuilder.Template{}
	if len(utxos) == 1 {
		return tpls, utxos[len(utxos)-1], nil
	}

	acp, err := m.GetLocalCtrlProgramByAddress(utxos[0].Address)
	if err != nil {
		return nil, nil, err
	}

	buildAmount := uint64(0)
	builder := &txbuilder.TemplateBuilder{}
	for index := 0; index < len(utxos); index++ {
		input, sigInst, err := UtxoToInputs(signer, utxos[index])
		if err != nil {
			return nil, nil, err
		}

		if err = builder.AddInput(input, sigInst); err != nil {
			return nil, nil, err
		}

		buildAmount += input.Amount()
		if builder.InputCount() != chainTxUtxoNum && index != len(utxos)-1 {
			continue
		}

		outAmount := buildAmount - chainTxMergeGas
		output := types.NewTxOutput(*consensus.BTMAssetID, outAmount, acp.ControlProgram)
		if err := builder.AddOutput(output); err != nil {
			return nil, nil, err
		}

		tpl, _, err := builder.Build()
		if err != nil {
			return nil, nil, err
		}

		bcOut, err := tpl.Transaction.Output(*tpl.Transaction.ResultIds[0])
		if err != nil {
			return nil, nil, err
		}

		utxos = append(utxos, &UTXO{
			OutputID:            *tpl.Transaction.ResultIds[0],
			AssetID:             *consensus.BTMAssetID,
			Amount:              outAmount,
			ControlProgram:      acp.ControlProgram,
			SourceID:            *bcOut.Source.Ref,
			SourcePos:           bcOut.Source.Position,
			ControlProgramIndex: acp.KeyIndex,
			Address:             acp.Address,
			Change:              acp.Change,
		})

		tpls = append(tpls, tpl)
		buildAmount = 0
		builder = &txbuilder.TemplateBuilder{}
		if index == len(utxos)-2 {
			break
		}
	}
	return tpls, utxos[len(utxos)-1], nil
}

// SpendAccountChain build the spend action with auto merge utxo function
func SpendAccountChain(ctx context.Context, builder *txbuilder.TemplateBuilder, action txbuilder.Action) ([]*txbuilder.Template, error) {
	act, ok := action.(*spendAction)
	if !ok {
		return nil, errors.New("fail to convert the spend action")
	}
	if *act.AssetId != *consensus.BTMAssetID {
		return nil, errors.New("spend chain action only support BTM")
	}

	utxos, err := act.accounts.reserveBtmUtxoChain(builder, act.AccountID, act.Amount, act.UseUnconfirmed)
	if err != nil {
		return nil, err
	}

	acct, err := act.accounts.FindByID(act.AccountID)
	if err != nil {
		return nil, err
	}

	tpls, utxo, err := act.accounts.buildBtmTxChain(utxos, acct.Signer)
	if err != nil {
		return nil, err
	}

	input, sigInst, err := UtxoToInputs(acct.Signer, utxo)
	if err != nil {
		return nil, err
	}

	if err := builder.AddInput(input, sigInst); err != nil {
		return nil, err
	}

	if utxo.Amount > act.Amount {
		if err = builder.AddOutput(types.NewTxOutput(*consensus.BTMAssetID, utxo.Amount-act.Amount, utxo.ControlProgram)); err != nil {
			return nil, errors.Wrap(err, "adding change output")
		}
	}
	return tpls, nil
}

func (a *spendAction) Build(ctx context.Context, b *txbuilder.TemplateBuilder) error {
	var missing []string
	if a.AccountID == "" {
		missing = append(missing, "account_id")
	}
	if a.AssetId.IsZero() {
		missing = append(missing, "asset_id")
	}
	if len(missing) > 0 {
		return txbuilder.MissingFieldsError(missing...)
	}

	acct, err := a.accounts.FindByID(a.AccountID)
	if err != nil {
		return errors.Wrap(err, "get account info")
	}

	res, err := a.accounts.utxoKeeper.Reserve(a.AccountID, a.AssetId, a.Amount, a.UseUnconfirmed, b.MaxTime())
	if err != nil {
		return errors.Wrap(err, "reserving utxos")
	}

	// Cancel the reservation if the build gets rolled back.
	b.OnRollback(func() { a.accounts.utxoKeeper.Cancel(res.id) })
	for _, r := range res.utxos {
		txInput, sigInst, err := UtxoToInputs(acct.Signer, r)
		if err != nil {
			return errors.Wrap(err, "creating inputs")
		}

		if err = b.AddInput(txInput, sigInst); err != nil {
			return errors.Wrap(err, "adding inputs")
		}
	}

	if res.change > 0 {
		acp, err := a.accounts.CreateAddress(a.AccountID, true)
		if err != nil {
			return errors.Wrap(err, "creating control program")
		}

		// Don't insert the control program until callbacks are executed.
		a.accounts.insertControlProgramDelayed(b, acp)
		if err = b.AddOutput(types.NewTxOutput(*a.AssetId, res.change, acp.ControlProgram)); err != nil {
			return errors.Wrap(err, "adding change output")
		}
	}
	return nil
}

//DecodeSpendUTXOAction unmarshal JSON-encoded data of spend utxo action
func (m *Manager) DecodeSpendUTXOAction(data []byte) (txbuilder.Action, error) {
	a := &spendUTXOAction{accounts: m}
	return a, json.Unmarshal(data, a)
}

type spendUTXOAction struct {
	accounts       *Manager
	OutputID       *bc.Hash                     `json:"output_id"`
	UseUnconfirmed bool                         `json:"use_unconfirmed"`
	Arguments      []txbuilder.ContractArgument `json:"arguments"`
}

func (a *spendUTXOAction) ActionType() string {
	return "spend_account_unspent_output"
}

func (a *spendUTXOAction) Build(ctx context.Context, b *txbuilder.TemplateBuilder) error {
	if a.OutputID == nil {
		return txbuilder.MissingFieldsError("output_id")
	}

	res, err := a.accounts.utxoKeeper.ReserveParticular(*a.OutputID, a.UseUnconfirmed, b.MaxTime())
	if err != nil {
		return err
	}

	b.OnRollback(func() { a.accounts.utxoKeeper.Cancel(res.id) })
	var accountSigner *signers.Signer
	if len(res.utxos[0].AccountID) != 0 {
		account, err := a.accounts.FindByID(res.utxos[0].AccountID)
		if err != nil {
			return err
		}

		accountSigner = account.Signer
	}

	txInput, sigInst, err := UtxoToInputs(accountSigner, res.utxos[0])
	if err != nil {
		return err
	}

	if a.Arguments == nil {
		return b.AddInput(txInput, sigInst)
	}

	sigInst = &txbuilder.SigningInstruction{}
	if err := txbuilder.AddContractArgs(sigInst, a.Arguments); err != nil {
		return err
	}

	return b.AddInput(txInput, sigInst)
}

// UtxoToInputs convert an utxo to the txinput
func UtxoToInputs(signer *signers.Signer, u *UTXO) (*types.TxInput, *txbuilder.SigningInstruction, error) {
	txInput := types.NewSpendInput(nil, u.SourceID, u.AssetID, u.Amount, u.SourcePos, u.ControlProgram)
	sigInst := &txbuilder.SigningInstruction{}
	if signer == nil {
		return txInput, sigInst, nil
	}

	path, err := signers.Path(signer, signers.AccountKeySpace, u.Change, u.ControlProgramIndex)
	if err != nil {
		return nil, nil, err
	}
	if u.Address == "" {
		sigInst.AddWitnessKeys(signer.XPubs, path, signer.Quorum)
		return txInput, sigInst, nil
	}

	address, err := common.DecodeAddress(u.Address, &consensus.ActiveNetParams)
	if err != nil {
		return nil, nil, err
	}

	sigInst.AddRawWitnessKeys(signer.XPubs, path, signer.Quorum)
	derivedXPubs := chainkd.DeriveXPubs(signer.XPubs, path)

	switch address.(type) {
	case *common.AddressWitnessPubKeyHash:
		derivedPK := derivedXPubs[0].PublicKey()
		sigInst.WitnessComponents = append(sigInst.WitnessComponents, txbuilder.DataWitness([]byte(derivedPK)))

	case *common.AddressWitnessScriptHash:
		derivedPKs := chainkd.XPubKeys(derivedXPubs)
		script, err := vmutil.P2SPMultiSigProgram(derivedPKs, signer.Quorum)
		if err != nil {
			return nil, nil, err
		}
		sigInst.WitnessComponents = append(sigInst.WitnessComponents, txbuilder.DataWitness(script))

	default:
		return nil, nil, errors.New("unsupport address type")
	}

	return txInput, sigInst, nil
}

// insertControlProgramDelayed takes a template builder and an account
// control program that hasn't been inserted to the database yet. It
// registers callbacks on the TemplateBuilder so that all of the template's
// account control programs are batch inserted if building the rest of
// the template is successful.
func (m *Manager) insertControlProgramDelayed(b *txbuilder.TemplateBuilder, acp *CtrlProgram) {
	m.delayedACPsMu.Lock()
	m.delayedACPs[b] = append(m.delayedACPs[b], acp)
	m.delayedACPsMu.Unlock()

	b.OnRollback(func() {
		m.delayedACPsMu.Lock()
		delete(m.delayedACPs, b)
		m.delayedACPsMu.Unlock()
	})
	b.OnBuild(func() error {
		m.delayedACPsMu.Lock()
		acps := m.delayedACPs[b]
		delete(m.delayedACPs, b)
		m.delayedACPsMu.Unlock()

		// Insert all of the account control programs at once.
		if len(acps) == 0 {
			return nil
		}
		return m.SaveControlPrograms(acps...)
	})
}
<|endoftext|>


<commit_before>package nattywad

import (
	"math/rand"
	"net"
	"sync"
	"time"

	"github.com/getlantern/go-natty/natty"
	"github.com/getlantern/waddell"
)

// ServerPeer identifies a server for NAT traversal
type ServerPeer struct {
	// ID: the server's PeerID on waddell (type 4 GUID)
	ID string

	// WaddellAddr: the address of the waddell server on which the server is
	// listening for offers.
	WaddellAddr string

	// Extras: Extra information about the peer (pass-through)
	Extras map[string]interface{}
}

func (p *ServerPeer) CompositeID() string {
	return p.WaddellAddr + "|" + p.ID
}

func (p *ServerPeer) String() string {
	return p.CompositeID()
}

// ClientSuccessCallback is a function that gets invoked when a client NAT
// traversal results in a UDP five tuple.
type ClientSuccessCallback func(info *TraversalInfo)

// ClientFailureCallback is a callback that is invoked if a client NAT traversal
// fails.
type ClientFailureCallback func(info *TraversalInfo)

// TraversalInfo provides information about traversals (successful and failed).
type TraversalInfo struct {
	// Peer: the ServerPeer with which we attempted traversal.
	Peer *ServerPeer

	// ServerRespondedToSignaling: indicates whether nattywad received any
	// signaling messages from the server peer during the traversal.
	ServerRespondedToSignaling bool

	// ServerGotFiveTuple: indicates whether or not the server peer got a
	// FiveTuple of its own (only populated when using nattywad as a client)
	ServerGotFiveTuple bool

	// LocalAddr: on a successful traversal, this contains the local UDP addr of
	// the FiveTuple.
	LocalAddr *net.UDPAddr

	// RemoteAddr: on a successful traversal, this contains the remote UDP addr
	// of the FiveTuple.
	RemoteAddr *net.UDPAddr

	// Duration: the duration of the traversal
	Duration time.Duration
}

// Client is a client that initiates NAT traversals to one or more configured
// servers. When a NAT traversal results in a 5-tuple, the OnFiveTuple callback
// is called.
type Client struct {
	// ClientMgr the ClientMgr to use to obtain Waddell connections
	ClientMgr *waddell.ClientMgr

	// OnSuccess: a callback that's invoked once a five tuple has been
	// obtained. Must be specified in order for Client to work.
	OnSuccess ClientSuccessCallback

	// OnFailure: a optional callback that's invoked if the NAT traversal fails
	// (e.g. times out). If unpopulated, failures aren't reported.
	OnFailure ClientFailureCallback

	// KeepAliveInterval: If specified to a non-zero value, nattywad will send a
	// keepalive message over the waddell channel to keep open the underlying
	// connections.
	KeepAliveInterval time.Duration

	serverPeers  map[string]*ServerPeer
	workers      map[traversalId]*clientWorker
	workersMutex sync.RWMutex
	waddells     map[string]*waddell.Client
	cfgMutex     sync.Mutex
}

// Configure (re)configures this Client to communicate with the given list of
// server peers. Anytime that the list is found to contain a new peer, a NAT
// traversal is attempted to that peer.
func (c *Client) Configure(serverPeers []*ServerPeer) {
	c.cfgMutex.Lock()
	defer c.cfgMutex.Unlock()

	log.Debugf("Configuring nat traversal client with %d server peers", len(serverPeers))

	// Lazily initialize data structures
	if c.serverPeers == nil {
		c.serverPeers = make(map[string]*ServerPeer)
		c.waddells = make(map[string]*waddell.Client)
		c.workers = make(map[traversalId]*clientWorker)
	}

	priorServerPeers := c.serverPeers
	c.serverPeers = make(map[string]*ServerPeer)

	for _, peer := range serverPeers {
		cid := peer.CompositeID()

		if priorServerPeers[cid] == nil {
			// Either we have a new server, or the address changed, try to
			// traverse
			log.Debugf("Attempting traversal to %s", peer.ID)
			peerId, err := waddell.PeerIdFromString(peer.ID)
			if err != nil {
				log.Errorf("Unable to parse PeerID for server peer %s: %s",
					peer.ID, err)
				continue
			}
			c.offer(peer, peerId)
		} else {
			log.Debugf("Already know about %s, not attempting traversal", peer.ID)
		}

		// Keep track of new peer
		c.serverPeers[cid] = peer
	}
}

func (c *Client) offer(serverPeer *ServerPeer, peerId waddell.PeerId) {
	wc := c.waddells[serverPeer.WaddellAddr]
	if wc == nil {
		/* new waddell server--open connection to it */
		var err error
		wc, err = c.ClientMgr.ClientTo(serverPeer.WaddellAddr)
		if err != nil {
			log.Errorf("Unable to connect to waddell: %s", err)
			return
		}
		if c.KeepAliveInterval > 0 {
			// Periodically send a KeepAlive message
			go func() {
				for {
					time.Sleep(c.KeepAliveInterval)
					err := wc.SendKeepAlive()
					if err != nil {
						log.Errorf("Unable to send KeepAlive packet to waddell: %s", err)
						return
					}
				}
			}()
		}
		c.waddells[serverPeer.WaddellAddr] = wc
		go c.receiveMessages(wc.In(NattywadTopic))
	}

	w := &clientWorker{
		out:         wc.Out(NattywadTopic),
		peerId:      peerId,
		onSuccess:   c.OnSuccess,
		onFailure:   c.OnFailure,
		tid:         traversalId(rand.Int31()),
		info:        &TraversalInfo{Peer: serverPeer},
		serverReady: make(chan bool, 10), // make this buffered to prevent deadlocks
	}
	c.addWorker(w)
	go func() {
		w.run()
		c.removeWorker(w)
	}()
}

func (c *Client) receiveMessages(in <-chan *waddell.MessageIn) {
	for wm := range in {
		msg := message(wm.Body)
		log.Tracef("Received %s from %s", msg.getData(), wm.From)
		w := c.getWorker(msg.getTraversalId())
		if w == nil {
			log.Debugf("Got message for unknown traversal %d, skipping", msg.getTraversalId())
			continue
		}
		w.messageReceived(msg)
	}
}

func (c *Client) addWorker(w *clientWorker) {
	c.workersMutex.Lock()
	defer c.workersMutex.Unlock()
	c.workers[w.tid] = w
}

func (c *Client) getWorker(tid traversalId) *clientWorker {
	c.workersMutex.RLock()
	defer c.workersMutex.RUnlock()
	return c.workers[tid]
}

func (c *Client) removeWorker(w *clientWorker) {
	c.workersMutex.Lock()
	defer c.workersMutex.Unlock()
	delete(c.workers, w.tid)
}

// clientWorker encapsulates the work done by the client for a single NAT
// traversal.
type clientWorker struct {
	out         chan<- *waddell.MessageOut
	peerId      waddell.PeerId
	onSuccess   ClientSuccessCallback
	onFailure   ClientFailureCallback
	tid         traversalId
	traversal   *natty.Traversal
	info        *TraversalInfo
	startedAt   time.Time
	serverReady chan bool
}

func (w *clientWorker) run() {
	w.traversal = natty.Offer(Timeout)
	defer func() {
		err := w.traversal.Close()
		if err != nil {
			log.Debugf("Unable to close traversal: %s", err)
		}
	}()

	go w.sendMessages()

	w.startedAt = time.Now()
	ft, err := w.traversal.FiveTuple()
	if err != nil {
		log.Errorf("Traversal to %s failed: %s", w.peerId, err)
		if w.onFailure != nil {
			w.info.Duration = time.Now().Sub(w.startedAt)
			w.onFailure(w.info)
		}
		return
	}
	if <-w.serverReady {
		local, remote, err := ft.UDPAddrs()
		if err != nil {
			log.Errorf("Unable to get UDP addresses for FiveTuple: %s", err)
			return
		}
		w.info.LocalAddr = local
		w.info.RemoteAddr = remote
		w.info.Duration = time.Now().Sub(w.startedAt)
		w.onSuccess(w.info)
	}
}

func (w *clientWorker) sendMessages() {
	for {
		msgOut, done := w.traversal.NextMsgOut()
		if done {
			return
		}
		w.out <- waddell.Message(w.peerId, w.tid.toBytes(), []byte(msgOut))
	}
}

func (w *clientWorker) messageReceived(msg message) {
	msgString := string(msg.getData())

	// Update info
	w.info.ServerRespondedToSignaling = true
	if natty.IsFiveTuple(msgString) {
		w.info.ServerGotFiveTuple = true
	}

	if msgString == ServerReady {
		// Server's ready!
		w.serverReady <- true
	} else {
		w.traversal.MsgIn(msgString)
	}
}
<commit_msg>Fixed comment on ServerGotFiveTuple<commit_after>package nattywad

import (
	"math/rand"
	"net"
	"sync"
	"time"

	"github.com/getlantern/go-natty/natty"
	"github.com/getlantern/waddell"
)

// ServerPeer identifies a server for NAT traversal
type ServerPeer struct {
	// ID: the server's PeerID on waddell (type 4 GUID)
	ID string

	// WaddellAddr: the address of the waddell server on which the server is
	// listening for offers.
	WaddellAddr string

	// Extras: Extra information about the peer (pass-through)
	Extras map[string]interface{}
}

func (p *ServerPeer) CompositeID() string {
	return p.WaddellAddr + "|" + p.ID
}

func (p *ServerPeer) String() string {
	return p.CompositeID()
}

// ClientSuccessCallback is a function that gets invoked when a client NAT
// traversal results in a UDP five tuple.
type ClientSuccessCallback func(info *TraversalInfo)

// ClientFailureCallback is a callback that is invoked if a client NAT traversal
// fails.
type ClientFailureCallback func(info *TraversalInfo)

// TraversalInfo provides information about traversals (successful and failed).
type TraversalInfo struct {
	// Peer: the ServerPeer with which we attempted traversal.
	Peer *ServerPeer

	// ServerRespondedToSignaling: indicates whether nattywad received any
	// signaling messages from the server peer during the traversal.
	ServerRespondedToSignaling bool

	// ServerGotFiveTuple: indicates whether or not the server peer got a
	// FiveTuple of its own.
	ServerGotFiveTuple bool

	// LocalAddr: on a successful traversal, this contains the local UDP addr of
	// the FiveTuple.
	LocalAddr *net.UDPAddr

	// RemoteAddr: on a successful traversal, this contains the remote UDP addr
	// of the FiveTuple.
	RemoteAddr *net.UDPAddr

	// Duration: the duration of the traversal
	Duration time.Duration
}

// Client is a client that initiates NAT traversals to one or more configured
// servers. When a NAT traversal results in a 5-tuple, the OnFiveTuple callback
// is called.
type Client struct {
	// ClientMgr the ClientMgr to use to obtain Waddell connections
	ClientMgr *waddell.ClientMgr

	// OnSuccess: a callback that's invoked once a five tuple has been
	// obtained. Must be specified in order for Client to work.
	OnSuccess ClientSuccessCallback

	// OnFailure: a optional callback that's invoked if the NAT traversal fails
	// (e.g. times out). If unpopulated, failures aren't reported.
	OnFailure ClientFailureCallback

	// KeepAliveInterval: If specified to a non-zero value, nattywad will send a
	// keepalive message over the waddell channel to keep open the underlying
	// connections.
	KeepAliveInterval time.Duration

	serverPeers  map[string]*ServerPeer
	workers      map[traversalId]*clientWorker
	workersMutex sync.RWMutex
	waddells     map[string]*waddell.Client
	cfgMutex     sync.Mutex
}

// Configure (re)configures this Client to communicate with the given list of
// server peers. Anytime that the list is found to contain a new peer, a NAT
// traversal is attempted to that peer.
func (c *Client) Configure(serverPeers []*ServerPeer) {
	c.cfgMutex.Lock()
	defer c.cfgMutex.Unlock()

	log.Debugf("Configuring nat traversal client with %d server peers", len(serverPeers))

	// Lazily initialize data structures
	if c.serverPeers == nil {
		c.serverPeers = make(map[string]*ServerPeer)
		c.waddells = make(map[string]*waddell.Client)
		c.workers = make(map[traversalId]*clientWorker)
	}

	priorServerPeers := c.serverPeers
	c.serverPeers = make(map[string]*ServerPeer)

	for _, peer := range serverPeers {
		cid := peer.CompositeID()

		if priorServerPeers[cid] == nil {
			// Either we have a new server, or the address changed, try to
			// traverse
			log.Debugf("Attempting traversal to %s", peer.ID)
			peerId, err := waddell.PeerIdFromString(peer.ID)
			if err != nil {
				log.Errorf("Unable to parse PeerID for server peer %s: %s",
					peer.ID, err)
				continue
			}
			c.offer(peer, peerId)
		} else {
			log.Debugf("Already know about %s, not attempting traversal", peer.ID)
		}

		// Keep track of new peer
		c.serverPeers[cid] = peer
	}
}

func (c *Client) offer(serverPeer *ServerPeer, peerId waddell.PeerId) {
	wc := c.waddells[serverPeer.WaddellAddr]
	if wc == nil {
		/* new waddell server--open connection to it */
		var err error
		wc, err = c.ClientMgr.ClientTo(serverPeer.WaddellAddr)
		if err != nil {
			log.Errorf("Unable to connect to waddell: %s", err)
			return
		}
		if c.KeepAliveInterval > 0 {
			// Periodically send a KeepAlive message
			go func() {
				for {
					time.Sleep(c.KeepAliveInterval)
					err := wc.SendKeepAlive()
					if err != nil {
						log.Errorf("Unable to send KeepAlive packet to waddell: %s", err)
						return
					}
				}
			}()
		}
		c.waddells[serverPeer.WaddellAddr] = wc
		go c.receiveMessages(wc.In(NattywadTopic))
	}

	w := &clientWorker{
		out:         wc.Out(NattywadTopic),
		peerId:      peerId,
		onSuccess:   c.OnSuccess,
		onFailure:   c.OnFailure,
		tid:         traversalId(rand.Int31()),
		info:        &TraversalInfo{Peer: serverPeer},
		serverReady: make(chan bool, 10), // make this buffered to prevent deadlocks
	}
	c.addWorker(w)
	go func() {
		w.run()
		c.removeWorker(w)
	}()
}

func (c *Client) receiveMessages(in <-chan *waddell.MessageIn) {
	for wm := range in {
		msg := message(wm.Body)
		log.Tracef("Received %s from %s", msg.getData(), wm.From)
		w := c.getWorker(msg.getTraversalId())
		if w == nil {
			log.Debugf("Got message for unknown traversal %d, skipping", msg.getTraversalId())
			continue
		}
		w.messageReceived(msg)
	}
}

func (c *Client) addWorker(w *clientWorker) {
	c.workersMutex.Lock()
	defer c.workersMutex.Unlock()
	c.workers[w.tid] = w
}

func (c *Client) getWorker(tid traversalId) *clientWorker {
	c.workersMutex.RLock()
	defer c.workersMutex.RUnlock()
	return c.workers[tid]
}

func (c *Client) removeWorker(w *clientWorker) {
	c.workersMutex.Lock()
	defer c.workersMutex.Unlock()
	delete(c.workers, w.tid)
}

// clientWorker encapsulates the work done by the client for a single NAT
// traversal.
type clientWorker struct {
	out         chan<- *waddell.MessageOut
	peerId      waddell.PeerId
	onSuccess   ClientSuccessCallback
	onFailure   ClientFailureCallback
	tid         traversalId
	traversal   *natty.Traversal
	info        *TraversalInfo
	startedAt   time.Time
	serverReady chan bool
}

func (w *clientWorker) run() {
	w.traversal = natty.Offer(Timeout)
	defer func() {
		err := w.traversal.Close()
		if err != nil {
			log.Debugf("Unable to close traversal: %s", err)
		}
	}()

	go w.sendMessages()

	w.startedAt = time.Now()
	ft, err := w.traversal.FiveTuple()
	if err != nil {
		log.Errorf("Traversal to %s failed: %s", w.peerId, err)
		if w.onFailure != nil {
			w.info.Duration = time.Now().Sub(w.startedAt)
			w.onFailure(w.info)
		}
		return
	}
	if <-w.serverReady {
		local, remote, err := ft.UDPAddrs()
		if err != nil {
			log.Errorf("Unable to get UDP addresses for FiveTuple: %s", err)
			return
		}
		w.info.LocalAddr = local
		w.info.RemoteAddr = remote
		w.info.Duration = time.Now().Sub(w.startedAt)
		w.onSuccess(w.info)
	}
}

func (w *clientWorker) sendMessages() {
	for {
		msgOut, done := w.traversal.NextMsgOut()
		if done {
			return
		}
		w.out <- waddell.Message(w.peerId, w.tid.toBytes(), []byte(msgOut))
	}
}

func (w *clientWorker) messageReceived(msg message) {
	msgString := string(msg.getData())

	// Update info
	w.info.ServerRespondedToSignaling = true
	if natty.IsFiveTuple(msgString) {
		w.info.ServerGotFiveTuple = true
	}

	if msgString == ServerReady {
		// Server's ready!
		w.serverReady <- true
	} else {
		w.traversal.MsgIn(msgString)
	}
}
<|endoftext|>


<commit_before>package arp

import (
	"bytes"
	"errors"
	"net"
	"syscall"
	"time"

	"github.com/mdlayher/ethernet"
	"github.com/mdlayher/raw"
)

var (
	// errNoIPv4Addr is returned when an interface does not have an IPv4
	// address.
	errNoIPv4Addr = errors.New("no IPv4 address available for interface")
)

// A Client is an ARP client, which can be used to send ARP requests to
// retrieve the MAC address of a machine using its IPv4 address.
type Client struct {
	ifi *net.Interface
	ip  net.IP
	p   net.PacketConn
}

// NewClient creates a new Client using the specified network interface.
// NewClient retrieves the IPv4 address of the interface and binds a raw socket
// to send and receive ARP packets.
func NewClient(ifi *net.Interface) (*Client, error) {
	// Check for a usable IPv4 address for the Client
	addrs, err := ifi.Addrs()
	if err != nil {
		return nil, err
	}
	ip, err := firstIPv4Addr(addrs)
	if err != nil {
		return nil, err
	}

	// Open raw socket to send and receive ARP packets using ethernet frames
	// we build ourselves
	p, err := raw.ListenPacket(ifi, syscall.SOCK_RAW, syscall.ETH_P_ARP)
	if err != nil {
		return nil, err
	}

	return &Client{
		ifi: ifi,
		ip:  ip,
		p:   p,
	}, nil
}

// Close closes the Client's raw socket and stops sending and receiving
// ARP packets.
func (c *Client) Close() error {
	return c.p.Close()
}

// Request performs an ARP request, attempting to retrieve the MAC address
// of a machine using its IPv4 address.
func (c *Client) Request(ip net.IP) (net.HardwareAddr, error) {
	// Create ARP packet addressed to broadcast MAC to attempt to find the
	// hardware address of the input IP address
	arp, err := NewPacket(OperationRequest, c.ifi.HardwareAddr, c.ip, ethernet.Broadcast, ip)
	if err != nil {
		return nil, err
	}
	arpb, err := arp.MarshalBinary()
	if err != nil {
		return nil, err
	}

	// Create thernet frame addressed to broadcast MAC to encapsulate the
	// ARP packet
	eth := &ethernet.Frame{
		DestinationMAC: ethernet.Broadcast,
		SourceMAC:      c.ifi.HardwareAddr,
		EtherType:      ethernet.EtherTypeARP,
		Payload:        arpb,
	}
	ethb, err := eth.MarshalBinary()
	if err != nil {
		return nil, err
	}

	// Write frame to ethernet broadcast address
	_, err = c.p.WriteTo(ethb, &raw.Addr{
		HardwareAddr: ethernet.Broadcast,
	})
	if err != nil {
		return nil, err
	}

	// Loop and wait for replies
	buf := make([]byte, 128)
	for {
		n, _, err := c.p.ReadFrom(buf)
		if err != nil {
			return nil, err
		}

		// Unmarshal ethernet frame and ARP packet
		if err := eth.UnmarshalBinary(buf[:n]); err != nil {
			return nil, err
		}
		if err := arp.UnmarshalBinary(eth.Payload); err != nil {
			return nil, err
		}

		// Check if ARP is in reply to our MAC address
		if bytes.Equal(arp.TargetMAC, c.ifi.HardwareAddr) {
			return arp.SenderMAC, nil
		}
	}
}

// Copyright (c) 2012 The Go Authors. All rights reserved.
// Source code in this file is based on src/net/interface_linux.go,
// from the Go standard library.  The Go license can be found here:
// https://golang.org/LICENSE.

// Documentation taken from net.PacketConn interface.  Thanks:
// http://golang.org/pkg/net/#PacketConn.

// SetDeadline sets the read and write deadlines associated with the
// connection.
func (c *Client) SetDeadline(t time.Time) error {
	return c.p.SetDeadline(t)
}

// SetReadDeadline sets the deadline for future raw socket read calls.
// If the deadline is reached, a raw socket read will fail with a timeout
// (see type net.Error) instead of blocking.
// A zero value for t means a raw socket read will not time out.
func (c *Client) SetReadDeadline(t time.Time) error {
	return c.p.SetReadDeadline(t)
}

// SetWriteDeadline sets the deadline for future raw socket write calls.
// If the deadline is reached, a raw socket write will fail with a timeout
// (see type net.Error) instead of blocking.
// A zero value for t means a raw socket write will not time out.
// Even if a write times out, it may return n > 0, indicating that
// some of the data was successfully written.
func (c *Client) SetWriteDeadline(t time.Time) error {
	return c.p.SetWriteDeadline(t)
}

// firstIPv4Addr attempts to retrieve the first detected IPv4 address from an
// input slice of network addresses.
func firstIPv4Addr(addrs []net.Addr) (net.IP, error) {
	for _, a := range addrs {
		if a.Network() != "ip+net" {
			continue
		}

		ip, _, err := net.ParseCIDR(a.String())
		if err != nil {
			return nil, err
		}

		// "If ip is not an IPv4 address, To4 returns nil."
		// Reference: http://golang.org/pkg/net/#IP.To4
		if ip4 := ip.To4(); ip4 != nil {
			return ip4, nil
		}
	}

	return nil, errNoIPv4Addr
}
<commit_msg>client: fix typo<commit_after>package arp

import (
	"bytes"
	"errors"
	"net"
	"syscall"
	"time"

	"github.com/mdlayher/ethernet"
	"github.com/mdlayher/raw"
)

var (
	// errNoIPv4Addr is returned when an interface does not have an IPv4
	// address.
	errNoIPv4Addr = errors.New("no IPv4 address available for interface")
)

// A Client is an ARP client, which can be used to send ARP requests to
// retrieve the MAC address of a machine using its IPv4 address.
type Client struct {
	ifi *net.Interface
	ip  net.IP
	p   net.PacketConn
}

// NewClient creates a new Client using the specified network interface.
// NewClient retrieves the IPv4 address of the interface and binds a raw socket
// to send and receive ARP packets.
func NewClient(ifi *net.Interface) (*Client, error) {
	// Check for a usable IPv4 address for the Client
	addrs, err := ifi.Addrs()
	if err != nil {
		return nil, err
	}
	ip, err := firstIPv4Addr(addrs)
	if err != nil {
		return nil, err
	}

	// Open raw socket to send and receive ARP packets using ethernet frames
	// we build ourselves
	p, err := raw.ListenPacket(ifi, syscall.SOCK_RAW, syscall.ETH_P_ARP)
	if err != nil {
		return nil, err
	}

	return &Client{
		ifi: ifi,
		ip:  ip,
		p:   p,
	}, nil
}

// Close closes the Client's raw socket and stops sending and receiving
// ARP packets.
func (c *Client) Close() error {
	return c.p.Close()
}

// Request performs an ARP request, attempting to retrieve the MAC address
// of a machine using its IPv4 address.
func (c *Client) Request(ip net.IP) (net.HardwareAddr, error) {
	// Create ARP packet addressed to broadcast MAC to attempt to find the
	// hardware address of the input IP address
	arp, err := NewPacket(OperationRequest, c.ifi.HardwareAddr, c.ip, ethernet.Broadcast, ip)
	if err != nil {
		return nil, err
	}
	arpb, err := arp.MarshalBinary()
	if err != nil {
		return nil, err
	}

	// Create ethernet frame addressed to broadcast MAC to encapsulate the
	// ARP packet
	eth := &ethernet.Frame{
		DestinationMAC: ethernet.Broadcast,
		SourceMAC:      c.ifi.HardwareAddr,
		EtherType:      ethernet.EtherTypeARP,
		Payload:        arpb,
	}
	ethb, err := eth.MarshalBinary()
	if err != nil {
		return nil, err
	}

	// Write frame to ethernet broadcast address
	_, err = c.p.WriteTo(ethb, &raw.Addr{
		HardwareAddr: ethernet.Broadcast,
	})
	if err != nil {
		return nil, err
	}

	// Loop and wait for replies
	buf := make([]byte, 128)
	for {
		n, _, err := c.p.ReadFrom(buf)
		if err != nil {
			return nil, err
		}

		// Unmarshal ethernet frame and ARP packet
		if err := eth.UnmarshalBinary(buf[:n]); err != nil {
			return nil, err
		}
		if err := arp.UnmarshalBinary(eth.Payload); err != nil {
			return nil, err
		}

		// Check if ARP is in reply to our MAC address
		if bytes.Equal(arp.TargetMAC, c.ifi.HardwareAddr) {
			return arp.SenderMAC, nil
		}
	}
}

// Copyright (c) 2012 The Go Authors. All rights reserved.
// Source code in this file is based on src/net/interface_linux.go,
// from the Go standard library.  The Go license can be found here:
// https://golang.org/LICENSE.

// Documentation taken from net.PacketConn interface.  Thanks:
// http://golang.org/pkg/net/#PacketConn.

// SetDeadline sets the read and write deadlines associated with the
// connection.
func (c *Client) SetDeadline(t time.Time) error {
	return c.p.SetDeadline(t)
}

// SetReadDeadline sets the deadline for future raw socket read calls.
// If the deadline is reached, a raw socket read will fail with a timeout
// (see type net.Error) instead of blocking.
// A zero value for t means a raw socket read will not time out.
func (c *Client) SetReadDeadline(t time.Time) error {
	return c.p.SetReadDeadline(t)
}

// SetWriteDeadline sets the deadline for future raw socket write calls.
// If the deadline is reached, a raw socket write will fail with a timeout
// (see type net.Error) instead of blocking.
// A zero value for t means a raw socket write will not time out.
// Even if a write times out, it may return n > 0, indicating that
// some of the data was successfully written.
func (c *Client) SetWriteDeadline(t time.Time) error {
	return c.p.SetWriteDeadline(t)
}

// firstIPv4Addr attempts to retrieve the first detected IPv4 address from an
// input slice of network addresses.
func firstIPv4Addr(addrs []net.Addr) (net.IP, error) {
	for _, a := range addrs {
		if a.Network() != "ip+net" {
			continue
		}

		ip, _, err := net.ParseCIDR(a.String())
		if err != nil {
			return nil, err
		}

		// "If ip is not an IPv4 address, To4 returns nil."
		// Reference: http://golang.org/pkg/net/#IP.To4
		if ip4 := ip.To4(); ip4 != nil {
			return ip4, nil
		}
	}

	return nil, errNoIPv4Addr
}
<|endoftext|>


<commit_before>package main

import (
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/anacrolix/torrent"
	"github.com/anacrolix/torrent/iplist"
	"github.com/dustin/go-humanize"
)

const clearScreen = "\033[H\033[2J"

const torrentBlockListURL = "http://john.bitsurge.net/public/biglist.p2p.gz"

var isHTTP = regexp.MustCompile(`^https?:\/\/`)

// ClientError formats errors coming from the client.
type ClientError struct {
	Type   string
	Origin error
}

func (clientError ClientError) Error() string {
	return fmt.Sprintf("Error %s: %s\n", clientError.Type, clientError.Origin)
}

// Client manages the torrent downloading.
type Client struct {
	Client   *torrent.Client
	Torrent  *torrent.Torrent
	Progress int64
	Uploaded int64
	Port     int
	Seed     bool
}

// NewClient creates a new torrent client based on a magnet or a torrent file.
// If the torrent file is on http, we try downloading it.
func NewClient(torrentPath string, port int, seed bool, tcp bool, maxConnections int) (client Client, err error) {
	var t *torrent.Torrent
	var c *torrent.Client

	client.Port = port

	// Create client.
	c, err = torrent.NewClient(&torrent.Config{
		DataDir:    os.TempDir(),
		NoUpload:   !seed,
		Seed:       seed,
		DisableTCP: !tcp,
	})

	if err != nil {
		return client, ClientError{Type: "creating torrent client", Origin: err}
	}

	client.Client = c
	client.Seed = seed

	// Add torrent.

	// Add as magnet url.
	if strings.HasPrefix(torrentPath, "magnet:") {
		if t, err = c.AddMagnet(torrentPath); err != nil {
			return client, ClientError{Type: "adding torrent", Origin: err}
		}
	} else {
		// Otherwise add as a torrent file.

		// If it's online, we try downloading the file.
		if isHTTP.MatchString(torrentPath) {
			if torrentPath, err = downloadFile(torrentPath); err != nil {
				return client, ClientError{Type: "downloading torrent file", Origin: err}
			}
		}

		// Check if the file exists.
		if _, err = os.Stat(torrentPath); err != nil {
			return client, ClientError{Type: "file not found", Origin: err}
		}

		if t, err = c.AddTorrentFromFile(torrentPath); err != nil {
			return client, ClientError{Type: "adding torrent to the client", Origin: err}
		}
	}

	client.Torrent = t
	client.Torrent.SetMaxEstablishedConns(maxConnections)

	go func() {
		<-t.GotInfo()
		t.DownloadAll()

		// Prioritize first 5% of the file.
		client.getLargestFile().PrioritizeRegion(0, int64(t.NumPieces()/100*5))
	}()

	go client.addBlocklist()

	return
}

// Download and add the blocklist.
func (c *Client) addBlocklist() {
	var err error
	blocklistPath := os.TempDir() + "/go-peerflix-blocklist.gz"

	if _, err = os.Stat(blocklistPath); os.IsNotExist(err) {
		err = downloadBlockList(blocklistPath)
	}

	if err != nil {
		log.Printf("Error downloading blocklist: %s", err)
		return
	}

	// Load blocklist.
	blocklistReader, err := os.Open(blocklistPath)
	if err != nil {
		log.Printf("Error opening blocklist: %s", err)
		return
	}

	// Extract file.
	gzipReader, err := gzip.NewReader(blocklistReader)
	if err != nil {
		log.Printf("Error extracting blocklist: %s", err)
		return
	}

	// Read as iplist.
	blocklist, err := iplist.NewFromReader(gzipReader)
	if err != nil {
		log.Printf("Error reading blocklist: %s", err)
		return
	}

	log.Printf("Setting blocklist.\nFound %d ranges\n", blocklist.NumRanges())
	c.Client.SetIPBlockList(blocklist)
}

func downloadBlockList(blocklistPath string) (err error) {
	log.Printf("Downloading blocklist")
	fileName, err := downloadFile(torrentBlockListURL)
	if err != nil {
		log.Printf("Error downloading blocklist: %s\n", err)
		return
	}

	// Ungzip file.
	in, err := os.Open(fileName)
	if err != nil {
		log.Printf("Error extracting blocklist: %s\n", err)
		return
	}
	defer func() {
		if err = in.Close(); err != nil {
			log.Printf("Error closing the blocklist gzip file: %s", err)
		}
	}()

	// Write to file.
	out, err := os.Create(blocklistPath)
	if err != nil {
		log.Printf("Error writing blocklist: %s\n", err)
		return
	}
	defer func() {
		if err = out.Close(); err != nil {
			log.Printf("Error closing the blocklist file: %s", err)
		}
	}()

	_, err = io.Copy(out, in)
	if err != nil {
		log.Printf("Error writing the blocklist file: %s", err)
		return
	}

	return
}

// Close cleans up the connections.
func (c *Client) Close() {
	c.Torrent.Drop()
	c.Client.Close()
}

// Render outputs the command line interface for the client.
func (c *Client) Render() {
	t := c.Torrent

	if t.Info() == nil {
		return
	}

	currentProgress := t.BytesCompleted()
	downloadSpeed := humanize.Bytes(uint64(currentProgress-c.Progress)) + "/s"
	c.Progress = currentProgress

	complete := humanize.Bytes(uint64(currentProgress))
	size := humanize.Bytes(uint64(t.Info().TotalLength()))

	uploadProgress := t.Stats().DataBytesSent - c.Uploaded
	uploadSpeed := humanize.Bytes(uint64(uploadProgress)) + "/s"
	c.Uploaded = uploadProgress

	print(clearScreen)
	fmt.Println(t.Info().Name)
	fmt.Println(strings.Repeat("=", len(t.Info().Name)))
	if c.ReadyForPlayback() {
		fmt.Printf("Stream: \thttp://localhost:%d\n", c.Port)
	}

	if currentProgress > 0 {
		fmt.Printf("Progress: \t%s / %s  %.2f%%\n", complete, size, c.percentage())
	}
	if currentProgress < t.Info().TotalLength() {
		fmt.Printf("Download speed: %s\n", downloadSpeed)
		if c.Seed {
			fmt.Printf("Upload speed: \t%s\n", uploadSpeed)
		}
	}

}

func (c Client) getLargestFile() *torrent.File {
	var target torrent.File
	var maxSize int64

	for _, file := range c.Torrent.Files() {
		if maxSize < file.Length() {
			maxSize = file.Length()
			target = file
		}
	}

	return &target
}

/*
func (c Client) RenderPieces() (output string) {
	pieces := c.Torrent.PieceStateRuns()
	for i := range pieces {
		piece := pieces[i]

		if piece.Priority == torrent.PiecePriorityReadahead {
			output += "!"
		}

		if piece.Partial {
			output += "P"
		} else if piece.Checking {
			output += "c"
		} else if piece.Complete {
			output += "d"
		} else {
			output += "_"
		}
	}

	return
}
*/

// ReadyForPlayback checks if the torrent is ready for playback or not.
// We wait until 5% of the torrent to start playing.
func (c Client) ReadyForPlayback() bool {
	return c.percentage() > 5
}

// GetFile is an http handler to serve the biggest file managed by the client.
func (c Client) GetFile(w http.ResponseWriter, r *http.Request) {
	target := c.getLargestFile()
	entry, err := NewFileReader(target)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	defer func() {
		if err := entry.Close(); err != nil {
			log.Printf("Error closing file reader: %s\n", err)
		}
	}()

	w.Header().Set("Content-Disposition", "attachment; filename=\""+c.Torrent.Info().Name+"\"")
	http.ServeContent(w, r, target.DisplayPath(), time.Now(), entry)
}

func (c Client) percentage() float64 {
	info := c.Torrent.Info()

	if info == nil {
		return 0
	}

	return float64(c.Torrent.BytesCompleted()) / float64(info.TotalLength()) * 100
}

func downloadFile(URL string) (fileName string, err error) {
	var file *os.File
	if file, err = ioutil.TempFile(os.TempDir(), "go-peerflix"); err != nil {
		return
	}

	defer func() {
		if ferr := file.Close(); ferr != nil {
			log.Printf("Error closing torrent file: %s", ferr)
		}
	}()

	response, err := http.Get(URL)
	if err != nil {
		return
	}

	defer func() {
		if ferr := response.Body.Close(); ferr != nil {
			log.Printf("Error closing torrent file: %s", ferr)
		}
	}()

	_, err = io.Copy(file, response.Body)

	return file.Name(), err
}
<commit_msg>Remove unnecessary check if the torrent file exists<commit_after>package main

import (
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/anacrolix/torrent"
	"github.com/anacrolix/torrent/iplist"
	"github.com/dustin/go-humanize"
)

const clearScreen = "\033[H\033[2J"

const torrentBlockListURL = "http://john.bitsurge.net/public/biglist.p2p.gz"

var isHTTP = regexp.MustCompile(`^https?:\/\/`)

// ClientError formats errors coming from the client.
type ClientError struct {
	Type   string
	Origin error
}

func (clientError ClientError) Error() string {
	return fmt.Sprintf("Error %s: %s\n", clientError.Type, clientError.Origin)
}

// Client manages the torrent downloading.
type Client struct {
	Client   *torrent.Client
	Torrent  *torrent.Torrent
	Progress int64
	Uploaded int64
	Port     int
	Seed     bool
}

// NewClient creates a new torrent client based on a magnet or a torrent file.
// If the torrent file is on http, we try downloading it.
func NewClient(torrentPath string, port int, seed bool, tcp bool, maxConnections int) (client Client, err error) {
	var t *torrent.Torrent
	var c *torrent.Client

	client.Port = port

	// Create client.
	c, err = torrent.NewClient(&torrent.Config{
		DataDir:    os.TempDir(),
		NoUpload:   !seed,
		Seed:       seed,
		DisableTCP: !tcp,
	})

	if err != nil {
		return client, ClientError{Type: "creating torrent client", Origin: err}
	}

	client.Client = c
	client.Seed = seed

	// Add torrent.

	// Add as magnet url.
	if strings.HasPrefix(torrentPath, "magnet:") {
		if t, err = c.AddMagnet(torrentPath); err != nil {
			return client, ClientError{Type: "adding torrent", Origin: err}
		}
	} else {
		// Otherwise add as a torrent file.

		// If it's online, we try downloading the file.
		if isHTTP.MatchString(torrentPath) {
			if torrentPath, err = downloadFile(torrentPath); err != nil {
				return client, ClientError{Type: "downloading torrent file", Origin: err}
			}
		}

		if t, err = c.AddTorrentFromFile(torrentPath); err != nil {
			return client, ClientError{Type: "adding torrent to the client", Origin: err}
		}
	}

	client.Torrent = t
	client.Torrent.SetMaxEstablishedConns(maxConnections)

	go func() {
		<-t.GotInfo()
		t.DownloadAll()

		// Prioritize first 5% of the file.
		client.getLargestFile().PrioritizeRegion(0, int64(t.NumPieces()/100*5))
	}()

	go client.addBlocklist()

	return
}

// Download and add the blocklist.
func (c *Client) addBlocklist() {
	var err error
	blocklistPath := os.TempDir() + "/go-peerflix-blocklist.gz"

	if _, err = os.Stat(blocklistPath); os.IsNotExist(err) {
		err = downloadBlockList(blocklistPath)
	}

	if err != nil {
		log.Printf("Error downloading blocklist: %s", err)
		return
	}

	// Load blocklist.
	blocklistReader, err := os.Open(blocklistPath)
	if err != nil {
		log.Printf("Error opening blocklist: %s", err)
		return
	}

	// Extract file.
	gzipReader, err := gzip.NewReader(blocklistReader)
	if err != nil {
		log.Printf("Error extracting blocklist: %s", err)
		return
	}

	// Read as iplist.
	blocklist, err := iplist.NewFromReader(gzipReader)
	if err != nil {
		log.Printf("Error reading blocklist: %s", err)
		return
	}

	log.Printf("Setting blocklist.\nFound %d ranges\n", blocklist.NumRanges())
	c.Client.SetIPBlockList(blocklist)
}

func downloadBlockList(blocklistPath string) (err error) {
	log.Printf("Downloading blocklist")
	fileName, err := downloadFile(torrentBlockListURL)
	if err != nil {
		log.Printf("Error downloading blocklist: %s\n", err)
		return
	}

	// Ungzip file.
	in, err := os.Open(fileName)
	if err != nil {
		log.Printf("Error extracting blocklist: %s\n", err)
		return
	}
	defer func() {
		if err = in.Close(); err != nil {
			log.Printf("Error closing the blocklist gzip file: %s", err)
		}
	}()

	// Write to file.
	out, err := os.Create(blocklistPath)
	if err != nil {
		log.Printf("Error writing blocklist: %s\n", err)
		return
	}
	defer func() {
		if err = out.Close(); err != nil {
			log.Printf("Error closing the blocklist file: %s", err)
		}
	}()

	_, err = io.Copy(out, in)
	if err != nil {
		log.Printf("Error writing the blocklist file: %s", err)
		return
	}

	return
}

// Close cleans up the connections.
func (c *Client) Close() {
	c.Torrent.Drop()
	c.Client.Close()
}

// Render outputs the command line interface for the client.
func (c *Client) Render() {
	t := c.Torrent

	if t.Info() == nil {
		return
	}

	currentProgress := t.BytesCompleted()
	downloadSpeed := humanize.Bytes(uint64(currentProgress-c.Progress)) + "/s"
	c.Progress = currentProgress

	complete := humanize.Bytes(uint64(currentProgress))
	size := humanize.Bytes(uint64(t.Info().TotalLength()))

	uploadProgress := t.Stats().DataBytesSent - c.Uploaded
	uploadSpeed := humanize.Bytes(uint64(uploadProgress)) + "/s"
	c.Uploaded = uploadProgress

	print(clearScreen)
	fmt.Println(t.Info().Name)
	fmt.Println(strings.Repeat("=", len(t.Info().Name)))
	if c.ReadyForPlayback() {
		fmt.Printf("Stream: \thttp://localhost:%d\n", c.Port)
	}

	if currentProgress > 0 {
		fmt.Printf("Progress: \t%s / %s  %.2f%%\n", complete, size, c.percentage())
	}
	if currentProgress < t.Info().TotalLength() {
		fmt.Printf("Download speed: %s\n", downloadSpeed)
		if c.Seed {
			fmt.Printf("Upload speed: \t%s\n", uploadSpeed)
		}
	}

}

func (c Client) getLargestFile() *torrent.File {
	var target torrent.File
	var maxSize int64

	for _, file := range c.Torrent.Files() {
		if maxSize < file.Length() {
			maxSize = file.Length()
			target = file
		}
	}

	return &target
}

/*
func (c Client) RenderPieces() (output string) {
	pieces := c.Torrent.PieceStateRuns()
	for i := range pieces {
		piece := pieces[i]

		if piece.Priority == torrent.PiecePriorityReadahead {
			output += "!"
		}

		if piece.Partial {
			output += "P"
		} else if piece.Checking {
			output += "c"
		} else if piece.Complete {
			output += "d"
		} else {
			output += "_"
		}
	}

	return
}
*/

// ReadyForPlayback checks if the torrent is ready for playback or not.
// We wait until 5% of the torrent to start playing.
func (c Client) ReadyForPlayback() bool {
	return c.percentage() > 5
}

// GetFile is an http handler to serve the biggest file managed by the client.
func (c Client) GetFile(w http.ResponseWriter, r *http.Request) {
	target := c.getLargestFile()
	entry, err := NewFileReader(target)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	defer func() {
		if err := entry.Close(); err != nil {
			log.Printf("Error closing file reader: %s\n", err)
		}
	}()

	w.Header().Set("Content-Disposition", "attachment; filename=\""+c.Torrent.Info().Name+"\"")
	http.ServeContent(w, r, target.DisplayPath(), time.Now(), entry)
}

func (c Client) percentage() float64 {
	info := c.Torrent.Info()

	if info == nil {
		return 0
	}

	return float64(c.Torrent.BytesCompleted()) / float64(info.TotalLength()) * 100
}

func downloadFile(URL string) (fileName string, err error) {
	var file *os.File
	if file, err = ioutil.TempFile(os.TempDir(), "go-peerflix"); err != nil {
		return
	}

	defer func() {
		if ferr := file.Close(); ferr != nil {
			log.Printf("Error closing torrent file: %s", ferr)
		}
	}()

	response, err := http.Get(URL)
	if err != nil {
		return
	}

	defer func() {
		if ferr := response.Body.Close(); ferr != nil {
			log.Printf("Error closing torrent file: %s", ferr)
		}
	}()

	_, err = io.Copy(file, response.Body)

	return file.Name(), err
}
<|endoftext|>


<commit_before>package gitlab

import (
	"errors"
	"net/http"

	"github.com/nsheridan/cashier/server/auth"
	"github.com/nsheridan/cashier/server/config"

	gitlabapi "github.com/xanzy/go-gitlab"
	"golang.org/x/oauth2"
)

const (
	name = "gitlab"
)

// Config is an implementation of `auth.Provider` for authenticating using a
// Gitlab account.
type Config struct {
	config    *oauth2.Config
	baseurl   string
	group     string
	whitelist map[string]bool
	allusers  bool
}

// New creates a new Gitlab provider from a configuration.
func New(c *config.Auth) (auth.Provider, error) {
	uw := make(map[string]bool)
	for _, u := range c.UsersWhitelist {
		uw[u] = true
	}
	allUsers := false
	if c.ProviderOpts["allusers"] == "true" {
		allUsers = true
	}
	if !allUsers && c.ProviderOpts["group"] == "" && len(uw) == 0 {
		return nil, errors.New("gitlab_opts group and the users whitelist must not be both empty if allusers isn't true")
	}
	siteURL := "https://gitlab.com/"
	if c.ProviderOpts["siteurl"] != "" {
		siteURL = c.ProviderOpts["siteurl"]
		if siteURL[len(siteURL)-1] != '/' {
			return nil, errors.New("gitlab_opts siteurl must end in /")
		}
	} else {
		if allUsers {
			return nil, errors.New("gitlab_opts if allusers is set, siteurl must be set")
		}
	}

	return &Config{
		config: &oauth2.Config{
			ClientID:     c.OauthClientID,
			ClientSecret: c.OauthClientSecret,
			RedirectURL:  c.OauthCallbackURL,
			Endpoint: oauth2.Endpoint{
				AuthURL:  siteURL + "oauth/authorize",
				TokenURL: siteURL + "oauth/token",
			},
			Scopes: []string{
				"api",
			},
		},
		group:     c.ProviderOpts["group"],
		whitelist: uw,
		allusers:  allUsers,
		baseurl:   siteURL + "api/v3/",
	}, nil
}

// A new oauth2 http client.
func (c *Config) newClient(token *oauth2.Token) *http.Client {
	return c.config.Client(oauth2.NoContext, token)
}

// Name returns the name of the provider.
func (c *Config) Name() string {
	return name
}

// Valid validates the oauth token.
func (c *Config) Valid(token *oauth2.Token) bool {
	if !token.Valid() {
		return false
	}
	if c.allusers {
		return true
	}
	if len(c.whitelist) > 0 && !c.whitelist[c.Username(token)] {
		return false
	}
	if c.group == "" {
		// There's no group and token is valid.  Can only reach
		// here if user whitelist is set and user is in whitelist.
		return true
	}
	client := gitlabapi.NewOAuthClient(nil, token.AccessToken)
	client.SetBaseURL(c.baseurl)
	groups, _, err := client.Groups.SearchGroup(c.group)
	if err != nil {
		return false
	}
	for _, g := range groups {
		if g.Path == c.group {
			return true
		}
	}
	return false
}

// Revoke is a no-op revoke method. Gitlab doesn't allow token
// revocation - tokens live for an hour.
// Returns nil to satisfy the Provider interface.
func (c *Config) Revoke(token *oauth2.Token) error {
	return nil
}

// StartSession retrieves an authentication endpoint from Gitlab.
func (c *Config) StartSession(state string) *auth.Session {
	return &auth.Session{
		AuthURL: c.config.AuthCodeURL(state),
	}
}

// Exchange authorizes the session and returns an access token.
func (c *Config) Exchange(code string) (*oauth2.Token, error) {
	return c.config.Exchange(oauth2.NoContext, code)
}

// Username retrieves the username of the Gitlab user.
func (c *Config) Username(token *oauth2.Token) string {
	client := gitlabapi.NewOAuthClient(nil, token.AccessToken)
	client.SetBaseURL(c.baseurl)
	u, _, err := client.Users.CurrentUser()
	if err != nil {
		return ""
	}
	return u.Username
}
<commit_msg>Code cleanup.<commit_after>package gitlab

import (
	"errors"
	"net/http"
	"strconv"

	"github.com/nsheridan/cashier/server/auth"
	"github.com/nsheridan/cashier/server/config"

	gitlabapi "github.com/xanzy/go-gitlab"
	"golang.org/x/oauth2"
)

const (
	name = "gitlab"
)

// Config is an implementation of `auth.Provider` for authenticating using a
// Gitlab account.
type Config struct {
	config    *oauth2.Config
	baseurl   string
	group     string
	whitelist map[string]bool
	allusers  bool
}

// New creates a new Gitlab provider from a configuration.
func New(c *config.Auth) (auth.Provider, error) {
	uw := make(map[string]bool)
	for _, u := range c.UsersWhitelist {
		uw[u] = true
	}
	allUsers, _ := strconv.ParseBool(c.ProviderOpts["allusers"])
	if !allUsers && c.ProviderOpts["group"] == "" && len(uw) == 0 {
		return nil, errors.New("gitlab_opts group and the users whitelist must not be both empty if allusers isn't true")
	}
	siteURL := "https://gitlab.com/"
	if c.ProviderOpts["siteurl"] != "" {
		siteURL = c.ProviderOpts["siteurl"]
		if siteURL[len(siteURL)-1] != '/' {
			return nil, errors.New("gitlab_opts siteurl must end in /")
		}
	} else {
		if allUsers {
			return nil, errors.New("gitlab_opts if allusers is set, siteurl must be set")
		}
	}

	return &Config{
		config: &oauth2.Config{
			ClientID:     c.OauthClientID,
			ClientSecret: c.OauthClientSecret,
			RedirectURL:  c.OauthCallbackURL,
			Endpoint: oauth2.Endpoint{
				AuthURL:  siteURL + "oauth/authorize",
				TokenURL: siteURL + "oauth/token",
			},
			Scopes: []string{
				"api",
			},
		},
		group:     c.ProviderOpts["group"],
		whitelist: uw,
		allusers:  allUsers,
		baseurl:   siteURL + "api/v3/",
	}, nil
}

// A new oauth2 http client.
func (c *Config) newClient(token *oauth2.Token) *http.Client {
	return c.config.Client(oauth2.NoContext, token)
}

// Name returns the name of the provider.
func (c *Config) Name() string {
	return name
}

// Valid validates the oauth token.
func (c *Config) Valid(token *oauth2.Token) bool {
	if !token.Valid() {
		return false
	}
	if c.allusers {
		return true
	}
	if len(c.whitelist) > 0 && !c.whitelist[c.Username(token)] {
		return false
	}
	if c.group == "" {
		// There's no group and token is valid.  Can only reach
		// here if user whitelist is set and user is in whitelist.
		return true
	}
	client := gitlabapi.NewOAuthClient(c.newClient(token), token.AccessToken)
	client.SetBaseURL(c.baseurl)
	groups, _, err := client.Groups.SearchGroup(c.group)
	if err != nil {
		return false
	}
	for _, g := range groups {
		if g.Path == c.group {
			return true
		}
	}
	return false
}

// Revoke is a no-op revoke method. Gitlab doesn't allow token
// revocation - tokens live for an hour.
// Returns nil to satisfy the Provider interface.
func (c *Config) Revoke(token *oauth2.Token) error {
	return nil
}

// StartSession retrieves an authentication endpoint from Gitlab.
func (c *Config) StartSession(state string) *auth.Session {
	return &auth.Session{
		AuthURL: c.config.AuthCodeURL(state),
	}
}

// Exchange authorizes the session and returns an access token.
func (c *Config) Exchange(code string) (*oauth2.Token, error) {
	return c.config.Exchange(oauth2.NoContext, code)
}

// Username retrieves the username of the Gitlab user.
func (c *Config) Username(token *oauth2.Token) string {
	client := gitlabapi.NewOAuthClient(c.newClient(token), token.AccessToken)
	client.SetBaseURL(c.baseurl)
	u, _, err := client.Users.CurrentUser()
	if err != nil {
		return ""
	}
	return u.Username
}
<|endoftext|>


<commit_before>package omise

import (
	"bytes"
	"encoding/json"
	"fmt"
	"go/build"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/omise/omise-go/internal"
)

var _ = fmt.Println

// Client helps you configure and perform HTTP operations against Omise's REST API. It
// should be used with operation structures from the operations subpackage.
type Client struct {
	*http.Client
	debug bool
	pkey  string
	skey  string

	// Overrides
	Endpoints map[internal.Endpoint]string

	// configuration
	APIVersion string
	GoVersion  string
}

// NewClient creates and returns a Client with the given public key and secret key.  Signs
// in to http://omise.co and visit https://dashboard.omise.co/test/dashboard to obtain
// your test (or live) keys.
func NewClient(pkey, skey string) (*Client, error) {
	switch {
	case pkey == "" && skey == "":
		return nil, ErrInvalidKey
	case pkey != "" && !strings.HasPrefix(pkey, "pkey_"):
		return nil, ErrInvalidKey
	case skey != "" && !strings.HasPrefix(skey, "skey_"):
		return nil, ErrInvalidKey
	}

	client := &Client{
		Client: &http.Client{Transport: transport},
		debug:  false,
		pkey:   pkey,
		skey:   skey,

		Endpoints: map[internal.Endpoint]string{},
	}

	if len(build.Default.ReleaseTags) > 0 {
		client.GoVersion = build.Default.ReleaseTags[len(build.Default.ReleaseTags)-1]
	}

	return client, nil
}

// Request creates a new *http.Request that should performs the supplied Operation. Most
// people should use the Do method instead.
func (c *Client) Request(operation internal.Operation) (*http.Request, error) {
	var req *http.Request
	var err error
	req, err = c.buildJSONRequest(operation)
	if err != nil {
		return nil, err
	}

	err = c.setRequestHeaders(req, operation.Describe())
	if err != nil {
		return nil, err
	}

	return req, nil
}

func (c *Client) buildJSONRequest(operation internal.Operation) (*http.Request, error) {
	desc := operation.Describe()

	b, err := json.Marshal(operation)
	if err != nil {
		return nil, err
	}

	body := bytes.NewReader(b)

	endpoint := string(desc.Endpoint)
	if ep, ok := c.Endpoints[desc.Endpoint]; ok {
		endpoint = ep
	}

	return http.NewRequest(desc.Method, endpoint+desc.Path, body)
}

func (c *Client) setRequestHeaders(req *http.Request, desc *internal.Description) error {
	ua := "OmiseGo/2015-11-06"
	if c.GoVersion != "" {
		ua += " Go/" + c.GoVersion
	}

	if desc.ContentType != "" {
		req.Header.Add("Content-Type", desc.ContentType)
	}

	req.Header.Add("User-Agent", ua)
	if c.APIVersion != "" {
		req.Header.Add("Omise-Version", c.APIVersion)
	}

	switch desc.KeyKind() {
	case "public":
		req.SetBasicAuth(c.pkey, "")
	case "secret":
		req.SetBasicAuth(c.skey, "")
	default:
		return ErrInternal("unrecognized endpoint:" + desc.Endpoint)
	}

	return nil
}

// Do performs the supplied operation against Omise's REST API and unmarshal the response
// into the given result parameter. Results are usually basic objects or a list that
// corresponds to the operations being done.
//
// If the operation is successful, result should contains the response data. Otherwise a
// non-nil error should be returned. Error maybe of the omise-go.Error struct type, in
// which case you can further inspect the Code and Message field for more information.
func (c *Client) Do(result interface{}, operation internal.Operation) error {
	req, err := c.Request(operation)
	if err != nil {
		return err
	}

	// response
	resp, err := c.Client.Do(req)
	if resp != nil {
		defer resp.Body.Close()
	}
	if err != nil {
		return err
	}

	buffer, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return &ErrTransport{err, buffer}
	}

	switch {
	case resp.StatusCode != 200:
		err := &Error{StatusCode: resp.StatusCode}
		if err := json.Unmarshal(buffer, err); err != nil {
			return &ErrTransport{err, buffer}
		}

		return err
	} // status == 200 && e == nil

	if c.debug {
		fmt.Println("resp:", resp.StatusCode, string(buffer))
	}

	if result != nil {
		if err := json.Unmarshal(buffer, result); err != nil {
			return &ErrTransport{err, buffer}
		}
	}

	return nil
}
<commit_msg>Remove var _ = fmt.Println<commit_after>package omise

import (
	"bytes"
	"encoding/json"
	"fmt"
	"go/build"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/omise/omise-go/internal"
)

// Client helps you configure and perform HTTP operations against Omise's REST API. It
// should be used with operation structures from the operations subpackage.
type Client struct {
	*http.Client
	debug bool
	pkey  string
	skey  string

	// Overrides
	Endpoints map[internal.Endpoint]string

	// configuration
	APIVersion string
	GoVersion  string
}

// NewClient creates and returns a Client with the given public key and secret key.  Signs
// in to http://omise.co and visit https://dashboard.omise.co/test/dashboard to obtain
// your test (or live) keys.
func NewClient(pkey, skey string) (*Client, error) {
	switch {
	case pkey == "" && skey == "":
		return nil, ErrInvalidKey
	case pkey != "" && !strings.HasPrefix(pkey, "pkey_"):
		return nil, ErrInvalidKey
	case skey != "" && !strings.HasPrefix(skey, "skey_"):
		return nil, ErrInvalidKey
	}

	client := &Client{
		Client: &http.Client{Transport: transport},
		debug:  false,
		pkey:   pkey,
		skey:   skey,

		Endpoints: map[internal.Endpoint]string{},
	}

	if len(build.Default.ReleaseTags) > 0 {
		client.GoVersion = build.Default.ReleaseTags[len(build.Default.ReleaseTags)-1]
	}

	return client, nil
}

// Request creates a new *http.Request that should performs the supplied Operation. Most
// people should use the Do method instead.
func (c *Client) Request(operation internal.Operation) (*http.Request, error) {
	var req *http.Request
	var err error
	req, err = c.buildJSONRequest(operation)
	if err != nil {
		return nil, err
	}

	err = c.setRequestHeaders(req, operation.Describe())
	if err != nil {
		return nil, err
	}

	return req, nil
}

func (c *Client) buildJSONRequest(operation internal.Operation) (*http.Request, error) {
	desc := operation.Describe()

	b, err := json.Marshal(operation)
	if err != nil {
		return nil, err
	}

	body := bytes.NewReader(b)

	endpoint := string(desc.Endpoint)
	if ep, ok := c.Endpoints[desc.Endpoint]; ok {
		endpoint = ep
	}

	return http.NewRequest(desc.Method, endpoint+desc.Path, body)
}

func (c *Client) setRequestHeaders(req *http.Request, desc *internal.Description) error {
	ua := "OmiseGo/2015-11-06"
	if c.GoVersion != "" {
		ua += " Go/" + c.GoVersion
	}

	if desc.ContentType != "" {
		req.Header.Add("Content-Type", desc.ContentType)
	}

	req.Header.Add("User-Agent", ua)
	if c.APIVersion != "" {
		req.Header.Add("Omise-Version", c.APIVersion)
	}

	switch desc.KeyKind() {
	case "public":
		req.SetBasicAuth(c.pkey, "")
	case "secret":
		req.SetBasicAuth(c.skey, "")
	default:
		return ErrInternal("unrecognized endpoint:" + desc.Endpoint)
	}

	return nil
}

// Do performs the supplied operation against Omise's REST API and unmarshal the response
// into the given result parameter. Results are usually basic objects or a list that
// corresponds to the operations being done.
//
// If the operation is successful, result should contains the response data. Otherwise a
// non-nil error should be returned. Error maybe of the omise-go.Error struct type, in
// which case you can further inspect the Code and Message field for more information.
func (c *Client) Do(result interface{}, operation internal.Operation) error {
	req, err := c.Request(operation)
	if err != nil {
		return err
	}

	// response
	resp, err := c.Client.Do(req)
	if resp != nil {
		defer resp.Body.Close()
	}
	if err != nil {
		return err
	}

	buffer, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return &ErrTransport{err, buffer}
	}

	switch {
	case resp.StatusCode != 200:
		err := &Error{StatusCode: resp.StatusCode}
		if err := json.Unmarshal(buffer, err); err != nil {
			return &ErrTransport{err, buffer}
		}

		return err
	} // status == 200 && e == nil

	if c.debug {
		fmt.Println("resp:", resp.StatusCode, string(buffer))
	}

	if result != nil {
		if err := json.Unmarshal(buffer, result); err != nil {
			return &ErrTransport{err, buffer}
		}
	}

	return nil
}
<|endoftext|>