<commit_before>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You under the Apache License, Version 2.0
// (the "License"); you may not use this file except in compliance with
// the License.  You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	pb "github.com/apache/beam/sdks/go/pkg/beam/model/fnexecution_v1"
	"github.com/golang/protobuf/proto"
	"github.com/spf13/cobra"
)

var (
	provisionCmd = &cobra.Command{
		Use:   "provision",
		Short: "Provision commands",
	}

	infoCmd = &cobra.Command{
		Use:   "info",
		Short: "Retrieve provisioning info",
		RunE:  infoFn,
		Args:  cobra.NoArgs,
	}
)

func init() {
	provisionCmd.AddCommand(infoCmd)
}

func infoFn(cmd *cobra.Command, args []string) error {
	ctx, cc, err := dial()
	if err != nil {
		return err
	}
	defer cc.Close()

	client := pb.NewProvisionServiceClient(cc)

	info, err := client.GetProvisionInfo(ctx, &pb.GetProvisionInfoRequest{})
	if err != nil {
		return err
	}

	cmd.Print(proto.MarshalTextString(info.GetInfo()))
	return nil
}
<commit_msg>Update provision.go<commit_after>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You under the Apache License, Version 2.0
// (the "License"); you may not use this file except in compliance with
// the License.  You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	fnpb "github.com/apache/beam/sdks/go/pkg/beam/model/fnexecution_v1"
	"github.com/golang/protobuf/proto"
	"github.com/spf13/cobra"
)

var (
	provisionCmd = &cobra.Command{
		Use:   "provision",
		Short: "Provision commands",
	}

	infoCmd = &cobra.Command{
		Use:   "info",
		Short: "Retrieve provisioning info",
		RunE:  infoFn,
		Args:  cobra.NoArgs,
	}
)

func init() {
	provisionCmd.AddCommand(infoCmd)
}

func infoFn(cmd *cobra.Command, args []string) error {
	ctx, cc, err := dial()
	if err != nil {
		return err
	}
	defer cc.Close()

	client := fnpb.NewProvisionServiceClient(cc)

	info, err := client.GetProvisionInfo(ctx, &fnpb.GetProvisionInfoRequest{})
	if err != nil {
		return err
	}

	cmd.Print(proto.MarshalTextString(info.GetInfo()))
	return nil
}
<|endoftext|>


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

package wallet

import (
	"context"
	"crypto/rand"
	"crypto/tls"
	"net"

	"decred.org/cspp"
	"decred.org/cspp/coinjoin"
	"decred.org/dcrwallet/errors"
	"decred.org/dcrwallet/wallet/txauthor"
	"decred.org/dcrwallet/wallet/txrules"
	"decred.org/dcrwallet/wallet/txsizes"
	"decred.org/dcrwallet/wallet/udb"
	"decred.org/dcrwallet/wallet/walletdb"
	"github.com/decred/dcrd/dcrutil/v3"
	"github.com/decred/dcrd/wire"
	"github.com/decred/go-socks/socks"
	"golang.org/x/sync/errgroup"
)

// must be sorted large to small
var splitPoints = [...]dcrutil.Amount{
	1 << 36, // 687.19476736
	1 << 34, // 171.79869184
	1 << 32, // 042.94967296
	1 << 30, // 010.73741824
	1 << 28, // 002.68435456
	1 << 26, // 000.67108864
	1 << 24, // 000.16777216
	1 << 22, // 000.04194304
	1 << 20, // 000.01048576
	1 << 18, // 000.00262144
}

var splitSems = [len(splitPoints)]chan struct{}{}

func init() {
	for i := range splitSems {
		splitSems[i] = make(chan struct{}, 10)
	}
}

var (
	errNoSplitDenomination = errors.New("no suitable split denomination")
	errThrottledMixRequest = errors.New("throttled mix request for split denomination")
)

// DialFunc provides a method to dial a network connection.
// If the dialed network connection is secured by TLS, TLS
// configuration is provided by the method, not the caller.
type DialFunc func(ctx context.Context, network, addr string) (net.Conn, error)

func (w *Wallet) MixOutput(ctx context.Context, dialTLS DialFunc, csppserver string, output *wire.OutPoint, changeAccount, mixAccount, mixBranch uint32) error {
	op := errors.Opf("wallet.MixOutput(%v)", output)

	sdiff, err := w.NextStakeDifficulty(ctx)
	if err != nil {
		return errors.E(op, err)
	}

	var updates []func(walletdb.ReadWriteTx) error

	hold, err := w.holdUnlock()
	if err != nil {
		return errors.E(op, err)
	}
	defer hold.release()

	w.lockedOutpointMu.Lock()
	if _, exists := w.lockedOutpoints[*output]; exists {
		w.lockedOutpointMu.Unlock()
		err = errors.Errorf("output %v already locked", output)
		return errors.E(op, err)
	}

	var prevScript []byte
	var amount dcrutil.Amount
	err = walletdb.View(ctx, w.db, func(dbtx walletdb.ReadTx) error {
		txmgrNs := dbtx.ReadBucket(wtxmgrNamespaceKey)
		txDetails, err := w.txStore.TxDetails(txmgrNs, &output.Hash)
		if err != nil {
			return err
		}
		prevScript = txDetails.MsgTx.TxOut[output.Index].PkScript
		amount = dcrutil.Amount(txDetails.MsgTx.TxOut[output.Index].Value)
		return nil
	})
	if err != nil {
		w.lockedOutpointMu.Unlock()
		return errors.E(op, err)
	}
	w.lockedOutpoints[*output] = struct{}{}
	w.lockedOutpointMu.Unlock()

	defer func() {
		w.lockedOutpointMu.Lock()
		delete(w.lockedOutpoints, *output)
		w.lockedOutpointMu.Unlock()
	}()

	var count int
	var mixValue, remValue dcrutil.Amount
	for i, v := range splitPoints {
		// When the sdiff is more than four times this mixed output
		// amount, there is a smaller common mixed amount with more
		// pairing activity (due to CoinShuffle++ participation from
		// ticket buyers).  Skipping this amount and moving to the next
		// smallest common mixed amount will result in quicker pairings,
		// or pairings occurring at all.  The number of mixed outputs is
		// capped to prevent a single mix being overwhelmingly funded by
		// a single output, and to conserve memory resources.
		if i != len(splitPoints)-1 && 4*v >= sdiff {
			continue
		}
		count = int(amount / v)
		if count > 4 {
			count = 4
		}
		if count > 0 {
			remValue = amount - dcrutil.Amount(count)*v
			mixValue = v
			select {
			case <-ctx.Done():
				return errors.E(op, ctx.Err())
			case splitSems[i] <- struct{}{}:
				defer func() { <-splitSems[i] }()
			default:
				return errThrottledMixRequest
			}
			break
		}
	}
	if mixValue == splitPoints[len(splitPoints)-1] {
		remValue = 0
	}
	if mixValue == 0 {
		err := errors.Errorf("output %v (%v): %w", output, amount, errNoSplitDenomination)
		return errors.E(op, err)
	}

	const (
		txVersion = 1
		locktime  = 0
		expiry    = 0
	)
	pairing := coinjoin.EncodeDesc(coinjoin.P2PKHv0, int64(mixValue), txVersion, locktime, expiry)
	ses, err := cspp.NewSession(rand.Reader, debugLog, pairing, count)
	if err != nil {
		return errors.E(op, err)
	}
	var conn net.Conn
	if dialTLS != nil {
		conn, err = dialTLS(ctx, "tcp", csppserver)
	} else {
		conn, err = tls.Dial("tcp", csppserver, nil)
	}
	if err != nil {
		return errors.E(op, err)
	}
	defer conn.Close()
	log.Infof("Dialed CSPPServer %v -> %v", conn.LocalAddr(), conn.RemoteAddr())

	// Create change output from remaining value and contributed fee
	const P2PKHv0Len = 25
	feeRate := w.RelayFee()
	inScriptSizes := []int{txsizes.RedeemP2PKHSigScriptSize}
	outScriptSizes := make([]int, count)
	for i := range outScriptSizes {
		outScriptSizes[i] = P2PKHv0Len
	}
	size := txsizes.EstimateSerializeSizeFromScriptSizes(inScriptSizes, outScriptSizes, P2PKHv0Len)
	changeValue := remValue - txrules.FeeForSerializeSize(feeRate, size)
	var change *wire.TxOut
	if !txrules.IsDustAmount(changeValue, P2PKHv0Len, feeRate) {
		persist := w.deferPersistReturnedChild(ctx, &updates)
		const accountName = "" // not used, so can be faked.
		addr, err := w.nextAddress(ctx, op, persist,
			accountName, changeAccount, udb.InternalBranch, WithGapPolicyIgnore())
		if err != nil {
			return errors.E(op, err)
		}
		changeScript, version, err := addressScript(addr)
		if err != nil {
			return errors.E(op, err)
		}
		change = &wire.TxOut{
			Value:    int64(changeValue),
			PkScript: changeScript,
			Version:  version,
		}
	}

	log.Infof("Mixing output %v (%v)", output, amount)
	cj := w.newCsppJoin(ctx, change, mixValue, mixAccount, mixBranch, count)
	cj.addTxIn(prevScript, &wire.TxIn{
		PreviousOutPoint: *output,
		ValueIn:          int64(amount),
	})
	err = ses.DiceMix(ctx, conn, cj)
	if err != nil {
		return errors.E(op, err)
	}
	cjHash := cj.tx.TxHash()
	log.Infof("Completed CoinShuffle++ mix of output %v in transaction %v", output, &cjHash)

	err = walletdb.Update(ctx, w.db, func(dbtx walletdb.ReadWriteTx) error {
		for _, f := range updates {
			if err := f(dbtx); err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		return errors.E(op, err)
	}

	return nil
}

// MixAccount individually mixes outputs of an account into standard
// denominations, creating newly mixed outputs for a mixed account.
//
// Due to performance concerns of timing out in a CoinShuffle++ run, this
// function may throttle how many of the outputs are mixed each call.
func (w *Wallet) MixAccount(ctx context.Context, dialTLS DialFunc, csppserver string, changeAccount, mixAccount, mixBranch uint32) error {
	const op errors.Op = "wallet.MixAccount"

	hold, err := w.holdUnlock()
	if err != nil {
		return errors.E(op, err)
	}
	defer hold.release()

	_, tipHeight := w.MainChainTip(ctx)
	w.lockedOutpointMu.Lock()
	var credits []Input
	err = walletdb.View(ctx, w.db, func(dbtx walletdb.ReadTx) error {
		var err error
		credits, err = w.findEligibleOutputs(dbtx, changeAccount, 1, tipHeight)
		return err
	})
	w.lockedOutpointMu.Unlock()
	if err != nil {
		return errors.E(op, err)
	}
	unlockedCredits := credits[:0]
	for i := range credits {
		amount := dcrutil.Amount(credits[i].PrevOut.Value)
		if amount <= splitPoints[len(splitPoints)-1] {
			continue
		}
		unlockedCredits = append(unlockedCredits, credits[i])
	}
	credits = unlockedCredits
	shuffle(len(credits), func(i, j int) {
		credits[i], credits[j] = credits[j], credits[i]
	})
	if len(credits) > 32 { // simple throttle
		credits = credits[:32]
	}
	var g errgroup.Group
	for i := range credits {
		op := &credits[i].OutPoint
		g.Go(func() error {
			err := w.MixOutput(ctx, dialTLS, csppserver, op, changeAccount, mixAccount, mixBranch)
			if errors.Is(err, errThrottledMixRequest) {
				return nil
			}
			if errors.Is(err, errNoSplitDenomination) {
				return nil
			}
			if errors.Is(err, socks.ErrPoolMaxConnections) {
				return nil
			}
			return err
		})
	}
	err = g.Wait()
	if err != nil {
		return errors.E(op, err)
	}
	return nil
}

// randomInputSource wraps an InputSource to randomly pick UTXOs.
// This involves reading all UTXOs from the underlying source into memory.
func randomInputSource(source txauthor.InputSource) txauthor.InputSource {
	all, err := source(dcrutil.MaxAmount)
	if err == nil {
		shuffleUTXOs(all)
	}
	var n int
	var tot dcrutil.Amount
	return func(target dcrutil.Amount) (*txauthor.InputDetail, error) {
		if err != nil {
			return nil, err
		}
		if all.Amount <= target {
			return all, nil
		}
		for n < len(all.Inputs) {
			tot += dcrutil.Amount(all.Inputs[n].ValueIn)
			n++
			if tot >= target {
				break
			}
		}
		selected := &txauthor.InputDetail{
			Amount:            tot,
			Inputs:            all.Inputs[:n],
			Scripts:           all.Scripts[:n],
			RedeemScriptSizes: all.RedeemScriptSizes[:n],
		}
		return selected, nil
	}
}
<commit_msg>Fix denom chosen for change mixes of large outputs<commit_after>// Copyright (c) 2019-2020 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wallet

import (
	"context"
	"crypto/rand"
	"crypto/tls"
	"net"

	"decred.org/cspp"
	"decred.org/cspp/coinjoin"
	"decred.org/dcrwallet/errors"
	"decred.org/dcrwallet/wallet/txauthor"
	"decred.org/dcrwallet/wallet/txrules"
	"decred.org/dcrwallet/wallet/txsizes"
	"decred.org/dcrwallet/wallet/udb"
	"decred.org/dcrwallet/wallet/walletdb"
	"github.com/decred/dcrd/dcrutil/v3"
	"github.com/decred/dcrd/wire"
	"github.com/decred/go-socks/socks"
	"golang.org/x/sync/errgroup"
)

// must be sorted large to small
var splitPoints = [...]dcrutil.Amount{
	1 << 36, // 687.19476736
	1 << 34, // 171.79869184
	1 << 32, // 042.94967296
	1 << 30, // 010.73741824
	1 << 28, // 002.68435456
	1 << 26, // 000.67108864
	1 << 24, // 000.16777216
	1 << 22, // 000.04194304
	1 << 20, // 000.01048576
	1 << 18, // 000.00262144
}

var splitSems = [len(splitPoints)]chan struct{}{}

func init() {
	for i := range splitSems {
		splitSems[i] = make(chan struct{}, 10)
	}
}

var (
	errNoSplitDenomination = errors.New("no suitable split denomination")
	errThrottledMixRequest = errors.New("throttled mix request for split denomination")
)

// DialFunc provides a method to dial a network connection.
// If the dialed network connection is secured by TLS, TLS
// configuration is provided by the method, not the caller.
type DialFunc func(ctx context.Context, network, addr string) (net.Conn, error)

func (w *Wallet) MixOutput(ctx context.Context, dialTLS DialFunc, csppserver string, output *wire.OutPoint, changeAccount, mixAccount, mixBranch uint32) error {
	op := errors.Opf("wallet.MixOutput(%v)", output)

	sdiff, err := w.NextStakeDifficulty(ctx)
	if err != nil {
		return errors.E(op, err)
	}

	var updates []func(walletdb.ReadWriteTx) error

	hold, err := w.holdUnlock()
	if err != nil {
		return errors.E(op, err)
	}
	defer hold.release()

	w.lockedOutpointMu.Lock()
	if _, exists := w.lockedOutpoints[*output]; exists {
		w.lockedOutpointMu.Unlock()
		err = errors.Errorf("output %v already locked", output)
		return errors.E(op, err)
	}

	var prevScript []byte
	var amount dcrutil.Amount
	err = walletdb.View(ctx, w.db, func(dbtx walletdb.ReadTx) error {
		txmgrNs := dbtx.ReadBucket(wtxmgrNamespaceKey)
		txDetails, err := w.txStore.TxDetails(txmgrNs, &output.Hash)
		if err != nil {
			return err
		}
		prevScript = txDetails.MsgTx.TxOut[output.Index].PkScript
		amount = dcrutil.Amount(txDetails.MsgTx.TxOut[output.Index].Value)
		return nil
	})
	if err != nil {
		w.lockedOutpointMu.Unlock()
		return errors.E(op, err)
	}
	w.lockedOutpoints[*output] = struct{}{}
	w.lockedOutpointMu.Unlock()

	defer func() {
		w.lockedOutpointMu.Lock()
		delete(w.lockedOutpoints, *output)
		w.lockedOutpointMu.Unlock()
	}()

	var count int
	var mixValue, remValue dcrutil.Amount
	for i, v := range splitPoints {
		// When the sdiff is more than this mixed output amount, there
		// is a smaller common mixed amount with more pairing activity
		// (due to CoinShuffle++ participation from ticket buyers).
		// Skipping this amount and moving to the next smallest common
		// mixed amount will result in quicker pairings, or pairings
		// occurring at all.  The number of mixed outputs is capped to
		// prevent a single mix being overwhelmingly funded by a single
		// output, and to conserve memory resources.
		if i != len(splitPoints)-1 && v >= sdiff {
			continue
		}
		count = int(amount / v)
		if count > 4 {
			count = 4
		}
		if count > 0 {
			remValue = amount - dcrutil.Amount(count)*v
			mixValue = v
			select {
			case <-ctx.Done():
				return errors.E(op, ctx.Err())
			case splitSems[i] <- struct{}{}:
				defer func() { <-splitSems[i] }()
			default:
				return errThrottledMixRequest
			}
			break
		}
	}
	if mixValue == splitPoints[len(splitPoints)-1] {
		remValue = 0
	}
	if mixValue == 0 {
		err := errors.Errorf("output %v (%v): %w", output, amount, errNoSplitDenomination)
		return errors.E(op, err)
	}

	const (
		txVersion = 1
		locktime  = 0
		expiry    = 0
	)
	pairing := coinjoin.EncodeDesc(coinjoin.P2PKHv0, int64(mixValue), txVersion, locktime, expiry)
	ses, err := cspp.NewSession(rand.Reader, debugLog, pairing, count)
	if err != nil {
		return errors.E(op, err)
	}
	var conn net.Conn
	if dialTLS != nil {
		conn, err = dialTLS(ctx, "tcp", csppserver)
	} else {
		conn, err = tls.Dial("tcp", csppserver, nil)
	}
	if err != nil {
		return errors.E(op, err)
	}
	defer conn.Close()
	log.Infof("Dialed CSPPServer %v -> %v", conn.LocalAddr(), conn.RemoteAddr())

	// Create change output from remaining value and contributed fee
	const P2PKHv0Len = 25
	feeRate := w.RelayFee()
	inScriptSizes := []int{txsizes.RedeemP2PKHSigScriptSize}
	outScriptSizes := make([]int, count)
	for i := range outScriptSizes {
		outScriptSizes[i] = P2PKHv0Len
	}
	size := txsizes.EstimateSerializeSizeFromScriptSizes(inScriptSizes, outScriptSizes, P2PKHv0Len)
	changeValue := remValue - txrules.FeeForSerializeSize(feeRate, size)
	var change *wire.TxOut
	if !txrules.IsDustAmount(changeValue, P2PKHv0Len, feeRate) {
		persist := w.deferPersistReturnedChild(ctx, &updates)
		const accountName = "" // not used, so can be faked.
		addr, err := w.nextAddress(ctx, op, persist,
			accountName, changeAccount, udb.InternalBranch, WithGapPolicyIgnore())
		if err != nil {
			return errors.E(op, err)
		}
		changeScript, version, err := addressScript(addr)
		if err != nil {
			return errors.E(op, err)
		}
		change = &wire.TxOut{
			Value:    int64(changeValue),
			PkScript: changeScript,
			Version:  version,
		}
	}

	log.Infof("Mixing output %v (%v)", output, amount)
	cj := w.newCsppJoin(ctx, change, mixValue, mixAccount, mixBranch, count)
	cj.addTxIn(prevScript, &wire.TxIn{
		PreviousOutPoint: *output,
		ValueIn:          int64(amount),
	})
	err = ses.DiceMix(ctx, conn, cj)
	if err != nil {
		return errors.E(op, err)
	}
	cjHash := cj.tx.TxHash()
	log.Infof("Completed CoinShuffle++ mix of output %v in transaction %v", output, &cjHash)

	err = walletdb.Update(ctx, w.db, func(dbtx walletdb.ReadWriteTx) error {
		for _, f := range updates {
			if err := f(dbtx); err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		return errors.E(op, err)
	}

	return nil
}

// MixAccount individually mixes outputs of an account into standard
// denominations, creating newly mixed outputs for a mixed account.
//
// Due to performance concerns of timing out in a CoinShuffle++ run, this
// function may throttle how many of the outputs are mixed each call.
func (w *Wallet) MixAccount(ctx context.Context, dialTLS DialFunc, csppserver string, changeAccount, mixAccount, mixBranch uint32) error {
	const op errors.Op = "wallet.MixAccount"

	hold, err := w.holdUnlock()
	if err != nil {
		return errors.E(op, err)
	}
	defer hold.release()

	_, tipHeight := w.MainChainTip(ctx)
	w.lockedOutpointMu.Lock()
	var credits []Input
	err = walletdb.View(ctx, w.db, func(dbtx walletdb.ReadTx) error {
		var err error
		credits, err = w.findEligibleOutputs(dbtx, changeAccount, 1, tipHeight)
		return err
	})
	w.lockedOutpointMu.Unlock()
	if err != nil {
		return errors.E(op, err)
	}
	unlockedCredits := credits[:0]
	for i := range credits {
		amount := dcrutil.Amount(credits[i].PrevOut.Value)
		if amount <= splitPoints[len(splitPoints)-1] {
			continue
		}
		unlockedCredits = append(unlockedCredits, credits[i])
	}
	credits = unlockedCredits
	shuffle(len(credits), func(i, j int) {
		credits[i], credits[j] = credits[j], credits[i]
	})
	if len(credits) > 32 { // simple throttle
		credits = credits[:32]
	}
	var g errgroup.Group
	for i := range credits {
		op := &credits[i].OutPoint
		g.Go(func() error {
			err := w.MixOutput(ctx, dialTLS, csppserver, op, changeAccount, mixAccount, mixBranch)
			if errors.Is(err, errThrottledMixRequest) {
				return nil
			}
			if errors.Is(err, errNoSplitDenomination) {
				return nil
			}
			if errors.Is(err, socks.ErrPoolMaxConnections) {
				return nil
			}
			return err
		})
	}
	err = g.Wait()
	if err != nil {
		return errors.E(op, err)
	}
	return nil
}

// randomInputSource wraps an InputSource to randomly pick UTXOs.
// This involves reading all UTXOs from the underlying source into memory.
func randomInputSource(source txauthor.InputSource) txauthor.InputSource {
	all, err := source(dcrutil.MaxAmount)
	if err == nil {
		shuffleUTXOs(all)
	}
	var n int
	var tot dcrutil.Amount
	return func(target dcrutil.Amount) (*txauthor.InputDetail, error) {
		if err != nil {
			return nil, err
		}
		if all.Amount <= target {
			return all, nil
		}
		for n < len(all.Inputs) {
			tot += dcrutil.Amount(all.Inputs[n].ValueIn)
			n++
			if tot >= target {
				break
			}
		}
		selected := &txauthor.InputDetail{
			Amount:            tot,
			Inputs:            all.Inputs[:n],
			Scripts:           all.Scripts[:n],
			RedeemScriptSizes: all.RedeemScriptSizes[:n],
		}
		return selected, nil
	}
}
<|endoftext|>


<commit_before>package utils

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
)

type expectedRequest struct {
	URL      string
	Err      error
	Response string
}

// FakeDownloader is like Downloader, but it used in tests
// to stub out results
type FakeDownloader struct {
	expected []expectedRequest
	progress *Progress
}

// Check interface
var (
	_ Downloader = &FakeDownloader{}
)

// NewFakeDownloader creates new expected downloader
func NewFakeDownloader() *FakeDownloader {
	result := &FakeDownloader{}
	result.expected = make([]expectedRequest, 0)
	result.progress = NewProgress()
	result.progress.Start()
	return result
}

// ExpectResponse installs expectation on upcoming download with response
func (f *FakeDownloader) ExpectResponse(url string, response string) *FakeDownloader {
	f.expected = append(f.expected, expectedRequest{URL: url, Response: response})
	return f
}

// ExpectError installs expectation on upcoming download with error
func (f *FakeDownloader) ExpectError(url string, err error) *FakeDownloader {
	f.expected = append(f.expected, expectedRequest{URL: url, Err: err})
	return f
}

// Empty verifies that are planned downloads have happened
func (f *FakeDownloader) Empty() bool {
	return len(f.expected) == 0
}

// DownloadWithChecksum performs fake download by matching against first expectation in the queue, with cheksum verification
func (f *FakeDownloader) DownloadWithChecksum(url string, filename string, result chan<- error, expected ChecksumInfo, ignoreMismatch bool) {
	if len(f.expected) == 0 || f.expected[0].URL != url {
		result <- fmt.Errorf("unexpected request for %s", url)
		return
	}

	expectation := f.expected[0]
	f.expected = f.expected[1:]

	if expectation.Err != nil {
		result <- expectation.Err
		return
	}

	err := os.MkdirAll(filepath.Dir(filename), 0755)
	if err != nil {
		result <- err
		return
	}

	outfile, err := os.Create(filename)
	if err != nil {
		result <- err
		return
	}
	defer outfile.Close()

	cks := NewChecksumWriter()
	w := io.MultiWriter(outfile, cks)

	_, err = w.Write([]byte(expectation.Response))
	if err != nil {
		result <- err
		return
	}

	if expected.Size != -1 {
		if expected.Size != cks.Sum().Size || expected.MD5 != "" && expected.MD5 != cks.Sum().MD5 ||
			expected.SHA1 != "" && expected.SHA1 != cks.Sum().SHA1 || expected.SHA256 != "" && expected.SHA256 != cks.Sum().SHA256 {
			if ignoreMismatch {
				fmt.Printf("WARNING: checksums don't match: %#v != %#v\n", expected, cks.Sum())
			} else {
				result <- fmt.Errorf("checksums don't match: %#v != %#v", expected, cks.Sum())
				return
			}
		}
	}

	result <- nil
	return
}

// Download performs fake download by matching against first expectation in the queue
func (f *FakeDownloader) Download(url string, filename string, result chan<- error) {
	f.DownloadWithChecksum(url, filename, result, ChecksumInfo{Size: -1}, false)
}

// Shutdown does nothing
func (f *FakeDownloader) Shutdown() {
	f.progress.Shutdown()
}

// Pause does nothing
func (f *FakeDownloader) Pause() {
}

// Resume does nothing
func (f *FakeDownloader) Resume() {
}

// GetProgress does nothing
func (f *FakeDownloader) GetProgress() *Progress {
	return f.progress
}
<commit_msg>Add ability to expect responses in any order.<commit_after>package utils

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
)

type expectedRequest struct {
	URL      string
	Err      error
	Response string
}

// FakeDownloader is like Downloader, but it used in tests
// to stub out results
type FakeDownloader struct {
	expected    []expectedRequest
	anyExpected map[string]expectedRequest
	progress    *Progress
}

// Check interface
var (
	_ Downloader = &FakeDownloader{}
)

// NewFakeDownloader creates new expected downloader
func NewFakeDownloader() *FakeDownloader {
	result := &FakeDownloader{}
	result.expected = make([]expectedRequest, 0)
	result.anyExpected = make(map[string]expectedRequest)
	result.progress = NewProgress()
	result.progress.Start()
	return result
}

// ExpectResponse installs expectation on upcoming download with response
func (f *FakeDownloader) ExpectResponse(url string, response string) *FakeDownloader {
	f.expected = append(f.expected, expectedRequest{URL: url, Response: response})
	return f
}

// AnyExpectResponse installs expectation on upcoming download with response in any order (url should be unique)
func (f *FakeDownloader) AnyExpectResponse(url string, response string) *FakeDownloader {
	f.anyExpected[url] = expectedRequest{URL: url, Response: response}
	return f
}

// ExpectError installs expectation on upcoming download with error
func (f *FakeDownloader) ExpectError(url string, err error) *FakeDownloader {
	f.expected = append(f.expected, expectedRequest{URL: url, Err: err})
	return f
}

// Empty verifies that are planned downloads have happened
func (f *FakeDownloader) Empty() bool {
	return len(f.expected) == 0
}

// DownloadWithChecksum performs fake download by matching against first expectation in the queue or any expectation, with cheksum verification
func (f *FakeDownloader) DownloadWithChecksum(url string, filename string, result chan<- error, expected ChecksumInfo, ignoreMismatch bool) {
	var expectation expectedRequest
	if len(f.expected) > 0 && f.expected[0].URL == url {
		expectation, f.expected = f.expected[0], f.expected[1:]
	} else if _, ok := f.anyExpected[url]; ok {
		expectation = f.anyExpected[url]
		delete(f.anyExpected, url)
	} else {
		result <- fmt.Errorf("unexpected request for %s", url)
		return
	}

	if expectation.Err != nil {
		result <- expectation.Err
		return
	}

	err := os.MkdirAll(filepath.Dir(filename), 0755)
	if err != nil {
		result <- err
		return
	}

	outfile, err := os.Create(filename)
	if err != nil {
		result <- err
		return
	}
	defer outfile.Close()

	cks := NewChecksumWriter()
	w := io.MultiWriter(outfile, cks)

	_, err = w.Write([]byte(expectation.Response))
	if err != nil {
		result <- err
		return
	}

	if expected.Size != -1 {
		if expected.Size != cks.Sum().Size || expected.MD5 != "" && expected.MD5 != cks.Sum().MD5 ||
			expected.SHA1 != "" && expected.SHA1 != cks.Sum().SHA1 || expected.SHA256 != "" && expected.SHA256 != cks.Sum().SHA256 {
			if ignoreMismatch {
				fmt.Printf("WARNING: checksums don't match: %#v != %#v for %s\n", expected, cks.Sum(), url)
			} else {
				result <- fmt.Errorf("checksums don't match: %#v != %#v for %s", expected, cks.Sum(), url)
				return
			}
		}
	}

	result <- nil
	return
}

// Download performs fake download by matching against first expectation in the queue
func (f *FakeDownloader) Download(url string, filename string, result chan<- error) {
	f.DownloadWithChecksum(url, filename, result, ChecksumInfo{Size: -1}, false)
}

// Shutdown does nothing
func (f *FakeDownloader) Shutdown() {
	f.progress.Shutdown()
}

// Pause does nothing
func (f *FakeDownloader) Pause() {
}

// Resume does nothing
func (f *FakeDownloader) Resume() {
}

// GetProgress does nothing
func (f *FakeDownloader) GetProgress() *Progress {
	return f.progress
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"os"
	"sync"
	"time"
)

var (
	srv *http.Server
	wg  sync.WaitGroup

	seed sync.Once
)

// Address returns the http://$server:$port/ representation for clients to load
func Address() string {
	if srv == nil {
		return ""
	}
	return fmt.Sprintf("http://%s/", srv.Addr)
}

func Register() {
	if srv != nil {
		return // already initialized
	}

	srv = &http.Server{
		Addr: fmt.Sprintf("127.0.0.1:%d", getPort()),
	}

	http.HandleFunc("/done", func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, "finished\n")
		wg.Done()
	})
}

func getPort() int {
	// prep seed, otherwise math/rand.Intn is deterministic, which means there could be
	// something relying on that port
	seed.Do(func() {
		rand.Seed(time.Now().Unix()) // something
	})

	lowerBound := 1024 // ports below require elevated privileges
	n := rand.Intn(2<<15 - lowerBound)
	return n + lowerBound
}

// Start initializes the http server, binds and accepts connections
func Start() {
	if srv == nil {
		fmt.Fprint(os.Stderr, "ui/server: no http server has been registered")
		return
	}

	wg.Add(1) // force callers to wait

	// spawn off http server
	go func() {
		if err := srv.ListenAndServe(); err != nil {
			fmt.Printf("ERROR: failed creating localhost server err=%s\n", err)
		}
	}()
}

// Stop calls for a shutdown of the http server, if it exists
func Stop() error {
	// hold on until the form has been filled out
	wg.Wait()

	if srv != nil {
		return srv.Shutdown(nil)
	}
	return nil
}
<commit_msg>ui/server: show a nicer message after we're done with certs<commit_after>package server

import (
	"fmt"
	"html/template"
	"io"
	"math/rand"
	"net/http"
	"os"
	"sync"
	"time"
)

var (
	srv *http.Server
	wg  sync.WaitGroup

	seed sync.Once

	doneTpl = `<!doctype html>

  <title>cert-manage</title>

All Done!

You can close this browser window now.

`
)

// Address returns the http://$server:$port/ representation for clients to load
func Address() string {
	if srv == nil {
		return ""
	}
	return fmt.Sprintf("http://%s/", srv.Addr)
}

func Register() {
	if srv != nil {
		return // already initialized
	}

	srv = &http.Server{
		Addr: fmt.Sprintf("127.0.0.1:%d", getPort()),
	}

	http.HandleFunc("/done", func(w http.ResponseWriter, r *http.Request) {
		defer wg.Done()

		t, err := template.New("contents").Parse(doneTpl)
		if err != nil {
			io.WriteString(w, fmt.Sprintf("ERROR: %v", err))
			return
		}
		err = t.Execute(w, nil)
		if err != nil {
			io.WriteString(w, fmt.Sprintf("ERROR: %v", err))
		}
	})
}

func getPort() int {
	// prep seed, otherwise math/rand.Intn is deterministic, which means there could be
	// something relying on that port
	seed.Do(func() {
		rand.Seed(time.Now().Unix()) // something
	})

	lowerBound := 1024 // ports below require elevated privileges
	n := rand.Intn(2<<15 - lowerBound)
	return n + lowerBound
}

// Start initializes the http server, binds and accepts connections
func Start() {
	if srv == nil {
		fmt.Fprint(os.Stderr, "ui/server: no http server has been registered")
		return
	}

	wg.Add(1) // force callers to wait

	// spawn off http server
	go func() {
		if err := srv.ListenAndServe(); err != nil {
			fmt.Printf("ERROR: failed creating localhost server err=%s\n", err)
		}
	}()
}

// Stop calls for a shutdown of the http server, if it exists
func Stop() error {
	// hold on until the form has been filled out
	wg.Wait()

	if srv != nil {
		return srv.Shutdown(nil)
	}
	return nil
}
<|endoftext|>


<commit_before>package wallet

import (
	"encoding/json"
	"fmt"

	log "github.com/sirupsen/logrus"
	"github.com/tendermint/go-wire/data/base58"
	"github.com/tendermint/tmlibs/db"

	"github.com/bytom/account"
	"github.com/bytom/asset"
	"github.com/bytom/blockchain/pseudohsm"
	"github.com/bytom/crypto/ed25519/chainkd"
	"github.com/bytom/crypto/sha3pool"
	"github.com/bytom/protocol"
	"github.com/bytom/protocol/bc"
	"github.com/bytom/protocol/bc/types"
)

//SINGLE single sign
const SINGLE = 1

//RecoveryIndex walletdb recovery cp number
const RecoveryIndex = 5000

var walletKey = []byte("walletInfo")
var dbKeyForimportingPrivateKey = []byte("importingKeysInfo")

//StatusInfo is base valid block info to handle orphan block rollback
type StatusInfo struct {
	WorkHeight uint64
	WorkHash   bc.Hash
	BestHeight uint64
	BestHash   bc.Hash
}

//KeyInfo is key import status
type KeyInfo struct {
	Alias    string       `json:"alias"`
	XPub     chainkd.XPub `json:"xpub"`
	Percent  uint8        `json:"percent"`
	Complete bool         `json:"complete"`
}

//Wallet is related to storing account unspent outputs
type Wallet struct {
	DB                  db.DB
	status              StatusInfo
	AccountMgr          *account.Manager
	AssetReg            *asset.Registry
	Hsm                 *pseudohsm.HSM
	chain               *protocol.Chain
	rescanProgress      chan struct{}
	ImportingPrivateKey bool
	importingKeysInfo   []KeyInfo
}

//NewWallet return a new wallet instance
func NewWallet(walletDB db.DB, account *account.Manager, asset *asset.Registry, hsm *pseudohsm.HSM, chain *protocol.Chain) (*Wallet, error) {
	w := &Wallet{
		DB:                walletDB,
		AccountMgr:        account,
		AssetReg:          asset,
		chain:             chain,
		Hsm:               hsm,
		rescanProgress:    make(chan struct{}, 1),
		importingKeysInfo: make([]KeyInfo, 0),
	}

	if err := w.loadWalletInfo(); err != nil {
		return nil, err
	}

	if err := w.loadKeysInfo(); err != nil {
		return nil, err
	}

	w.ImportingPrivateKey = w.getImportKeyFlag()

	go w.walletUpdater()

	return w, nil
}

//GetWalletInfo return stored wallet info and nil,if error,
//return initial wallet info and err
func (w *Wallet) loadWalletInfo() error {
	if rawWallet := w.DB.Get(walletKey); rawWallet != nil {
		return json.Unmarshal(rawWallet, &w.status)
	}

	block, err := w.chain.GetBlockByHeight(0)
	if err != nil {
		return err
	}
	return w.AttachBlock(block)
}

func (w *Wallet) commitWalletInfo(batch db.Batch) error {
	rawWallet, err := json.Marshal(w.status)
	if err != nil {
		log.WithField("err", err).Error("save wallet info")
		return err
	}

	batch.Set(walletKey, rawWallet)
	batch.Write()
	return nil
}

//GetWalletInfo return stored wallet info and nil,if error,
//return initial wallet info and err
func (w *Wallet) loadKeysInfo() error {
	if rawKeyInfo := w.DB.Get(dbKeyForimportingPrivateKey); rawKeyInfo != nil {
		json.Unmarshal(rawKeyInfo, &w.importingKeysInfo)
		return nil
	}
	return nil
}

func (w *Wallet) commitkeysInfo() error {
	rawKeysInfo, err := json.Marshal(w.importingKeysInfo)
	if err != nil {
		log.WithField("err", err).Error("save wallet info")
		return err
	}
	w.DB.Set(dbKeyForimportingPrivateKey, rawKeysInfo)
	return nil
}

func (w *Wallet) getImportKeyFlag() bool {
	for _, v := range w.importingKeysInfo {
		if v.Complete == false {
			return true
		}
	}
	return false
}

// AttachBlock attach a new block
func (w *Wallet) AttachBlock(block *types.Block) error {
	if block.PreviousBlockHash != w.status.WorkHash {
		log.Warn("wallet skip attachBlock due to status hash not equal to previous hash")
		return nil
	}

	blockHash := block.Hash()
	txStatus, err := w.chain.GetTransactionStatus(&blockHash)
	if err != nil {
		return err
	}

	storeBatch := w.DB.NewBatch()
	w.indexTransactions(storeBatch, block, txStatus)
	w.buildAccountUTXOs(storeBatch, block, txStatus)

	w.status.WorkHeight = block.Height
	w.status.WorkHash = block.Hash()
	if w.status.WorkHeight >= w.status.BestHeight {
		w.status.BestHeight = w.status.WorkHeight
		w.status.BestHash = w.status.WorkHash
	}
	return w.commitWalletInfo(storeBatch)
}

// DetachBlock detach a block and rollback state
func (w *Wallet) DetachBlock(block *types.Block) error {
	blockHash := block.Hash()
	txStatus, err := w.chain.GetTransactionStatus(&blockHash)
	if err != nil {
		return err
	}

	storeBatch := w.DB.NewBatch()
	w.reverseAccountUTXOs(storeBatch, block, txStatus)
	w.deleteTransactions(storeBatch, w.status.BestHeight)

	w.status.BestHeight = block.Height - 1
	w.status.BestHash = block.PreviousBlockHash

	if w.status.WorkHeight > w.status.BestHeight {
		w.status.WorkHeight = w.status.BestHeight
		w.status.WorkHash = w.status.BestHash
	}

	return w.commitWalletInfo(storeBatch)
}

//WalletUpdate process every valid block and reverse every invalid block which need to rollback
func (w *Wallet) walletUpdater() {
	for {
		getRescanNotification(w)
		updateRescanStatus(w)
		for !w.chain.InMainChain(w.status.BestHash) {
			block, err := w.chain.GetBlockByHash(&w.status.BestHash)
			if err != nil {
				log.WithField("err", err).Error("walletUpdater GetBlockByHash")
				return
			}

			if err := w.DetachBlock(block); err != nil {
				log.WithField("err", err).Error("walletUpdater detachBlock")
				return
			}
		}

		block, _ := w.chain.GetBlockByHeight(w.status.WorkHeight + 1)
		if block == nil {
			<-w.chain.BlockWaiter(w.status.WorkHeight + 1)
			continue
		}

		if err := w.AttachBlock(block); err != nil {
			log.WithField("err", err).Error("walletUpdater stop")
			return
		}
	}
}

func getRescanNotification(w *Wallet) {
	select {
	case <-w.rescanProgress:
		w.status.WorkHeight = 0
		block, _ := w.chain.GetBlockByHeight(w.status.WorkHeight)
		w.status.WorkHash = block.Hash()
	default:
		return
	}
}

// ExportAccountPrivKey exports the account private key as a WIF for encoding as a string
// in the Wallet Import Formt.
func (w *Wallet) ExportAccountPrivKey(xpub chainkd.XPub, auth string) (*string, error) {
	xprv, err := w.Hsm.LoadChainKDKey(xpub, auth)
	if err != nil {
		return nil, err
	}
	var hashed [32]byte
	sha3pool.Sum256(hashed[:], xprv[:])

	tmp := append(xprv[:], hashed[:4]...)
	res := base58.Encode(tmp)
	return &res, nil
}

// ImportAccountPrivKey imports the account key in the Wallet Import Formt.
func (w *Wallet) ImportAccountPrivKey(xprv chainkd.XPrv, keyAlias, auth string, index uint64, accountAlias string) (*pseudohsm.XPub, error) {
	if w.Hsm.HasAlias(keyAlias) {
		return nil, pseudohsm.ErrDuplicateKeyAlias
	}
	if w.Hsm.HasKey(xprv) {
		return nil, pseudohsm.ErrDuplicateKey
	}

	if acc, _ := w.AccountMgr.FindByAlias(nil, accountAlias); acc != nil {
		return nil, account.ErrDuplicateAlias
	}

	xpub, _, err := w.Hsm.ImportXPrvKey(auth, keyAlias, xprv)
	if err != nil {
		return nil, err
	}

	newAccount, err := w.AccountMgr.Create(nil, []chainkd.XPub{xpub.XPub}, SINGLE, accountAlias)
	if err != nil {
		w.Hsm.XDelete(xpub.XPub, auth)
		return nil, err
	}
	if err := w.recoveryAccountWalletDB(newAccount, xpub, index, keyAlias); err != nil {
		w.AccountMgr.DeleteAccount(newAccount.ID)
		w.Hsm.XDelete(xpub.XPub, auth)
		return nil, err
	}
	return xpub, nil
}

// ImportAccountXpubKey imports the account key in the Wallet Import Formt.
func (w *Wallet) ImportAccountXpubKey(xpubIndex int, xpub pseudohsm.XPub, cpIndex uint64) error {
	accountAlias := fmt.Sprintf("recovery_%d", xpubIndex)

	if acc, _ := w.AccountMgr.FindByAlias(nil, accountAlias); acc != nil {
		return account.ErrDuplicateAlias
	}

	newAccount, err := w.AccountMgr.Create(nil, []chainkd.XPub{xpub.XPub}, SINGLE, accountAlias)
	if err != nil {
		return err
	}

	return w.recoveryAccountWalletDB(newAccount, &xpub, cpIndex, xpub.Alias)
}

func (w *Wallet) recoveryAccountWalletDB(account *account.Account, XPub *pseudohsm.XPub, index uint64, keyAlias string) error {
	if err := w.createProgram(account, XPub, index); err != nil {
		return err
	}
	w.ImportingPrivateKey = true
	tmp := KeyInfo{
		Alias:    keyAlias,
		XPub:     XPub.XPub,
		Complete: false,
	}
	w.importingKeysInfo = append(w.importingKeysInfo, tmp)
	w.commitkeysInfo()
	w.rescanBlocks()

	return nil
}

func (w *Wallet) createProgram(account *account.Account, XPub *pseudohsm.XPub, index uint64) error {
	for i := uint64(0); i < index; i++ {
		if _, err := w.AccountMgr.CreateAddress(nil, account.ID, false); err != nil {
			return err
		}
	}
	return nil
}

func (w *Wallet) rescanBlocks() {
	select {
	case w.rescanProgress <- struct{}{}:
	default:
		return
	}
}

//GetRescanStatus return key import rescan status
func (w *Wallet) GetRescanStatus() ([]KeyInfo, error) {
	keysInfo := make([]KeyInfo, len(w.importingKeysInfo))

	if rawKeyInfo := w.DB.Get(dbKeyForimportingPrivateKey); rawKeyInfo != nil {
		if err := json.Unmarshal(rawKeyInfo, &keysInfo); err != nil {
			return nil, err
		}
	}

	var status StatusInfo
	if rawWallet := w.DB.Get(walletKey); rawWallet != nil {
		if err := json.Unmarshal(rawWallet, &status); err != nil {
			return nil, err
		}
	}

	for i := range keysInfo {
		if keysInfo[i].Complete == true || status.BestHeight == 0 {
			keysInfo[i].Percent = 100
			continue
		}

		keysInfo[i].Percent = uint8(status.WorkHeight * 100 / status.BestHeight)
		if keysInfo[i].Percent == 100 {
			keysInfo[i].Complete = true
		}
	}
	return keysInfo, nil
}

//updateRescanStatus mark private key import process `Complete` if rescan finished
func updateRescanStatus(w *Wallet) {
	if !w.ImportingPrivateKey || w.status.WorkHeight < w.status.BestHeight {
		return
	}

	w.ImportingPrivateKey = false
	for _, keyInfo := range w.importingKeysInfo {
		keyInfo.Percent = 100
		keyInfo.Complete = true
	}
	w.commitkeysInfo()
}
<commit_msg>improve importing key progress calculation<commit_after>package wallet

import (
	"encoding/json"
	"fmt"

	log "github.com/sirupsen/logrus"
	"github.com/tendermint/go-wire/data/base58"
	"github.com/tendermint/tmlibs/db"

	"github.com/bytom/account"
	"github.com/bytom/asset"
	"github.com/bytom/blockchain/pseudohsm"
	"github.com/bytom/crypto/ed25519/chainkd"
	"github.com/bytom/crypto/sha3pool"
	"github.com/bytom/protocol"
	"github.com/bytom/protocol/bc"
	"github.com/bytom/protocol/bc/types"
)

//SINGLE single sign
const SINGLE = 1

//RecoveryIndex walletdb recovery cp number
const RecoveryIndex = 5000

var walletKey = []byte("walletInfo")
var dbKeyForimportingPrivateKey = []byte("importingKeysInfo")

//StatusInfo is base valid block info to handle orphan block rollback
type StatusInfo struct {
	WorkHeight uint64
	WorkHash   bc.Hash
	BestHeight uint64
	BestHash   bc.Hash
}

//KeyInfo is key import status
type KeyInfo struct {
	Alias    string       `json:"alias"`
	XPub     chainkd.XPub `json:"xpub"`
	Percent  uint8        `json:"percent"`
	Complete bool         `json:"complete"`
}

//Wallet is related to storing account unspent outputs
type Wallet struct {
	DB                  db.DB
	status              StatusInfo
	AccountMgr          *account.Manager
	AssetReg            *asset.Registry
	Hsm                 *pseudohsm.HSM
	chain               *protocol.Chain
	rescanProgress      chan struct{}
	ImportingPrivateKey bool
	importingKeysInfo   []KeyInfo
}

//NewWallet return a new wallet instance
func NewWallet(walletDB db.DB, account *account.Manager, asset *asset.Registry, hsm *pseudohsm.HSM, chain *protocol.Chain) (*Wallet, error) {
	w := &Wallet{
		DB:                walletDB,
		AccountMgr:        account,
		AssetReg:          asset,
		chain:             chain,
		Hsm:               hsm,
		rescanProgress:    make(chan struct{}, 1),
		importingKeysInfo: make([]KeyInfo, 0),
	}

	if err := w.loadWalletInfo(); err != nil {
		return nil, err
	}

	if err := w.loadKeysInfo(); err != nil {
		return nil, err
	}

	w.ImportingPrivateKey = w.getImportKeyFlag()

	go w.walletUpdater()

	return w, nil
}

//GetWalletInfo return stored wallet info and nil,if error,
//return initial wallet info and err
func (w *Wallet) loadWalletInfo() error {
	if rawWallet := w.DB.Get(walletKey); rawWallet != nil {
		return json.Unmarshal(rawWallet, &w.status)
	}

	block, err := w.chain.GetBlockByHeight(0)
	if err != nil {
		return err
	}
	return w.AttachBlock(block)
}

func (w *Wallet) commitWalletInfo(batch db.Batch) error {
	rawWallet, err := json.Marshal(w.status)
	if err != nil {
		log.WithField("err", err).Error("save wallet info")
		return err
	}

	batch.Set(walletKey, rawWallet)
	batch.Write()
	return nil
}

//GetWalletInfo return stored wallet info and nil,if error,
//return initial wallet info and err
func (w *Wallet) loadKeysInfo() error {
	if rawKeyInfo := w.DB.Get(dbKeyForimportingPrivateKey); rawKeyInfo != nil {
		json.Unmarshal(rawKeyInfo, &w.importingKeysInfo)
		return nil
	}
	return nil
}

func (w *Wallet) commitkeysInfo() error {
	rawKeysInfo, err := json.Marshal(w.importingKeysInfo)
	if err != nil {
		log.WithField("err", err).Error("save wallet info")
		return err
	}
	w.DB.Set(dbKeyForimportingPrivateKey, rawKeysInfo)
	return nil
}

func (w *Wallet) getImportKeyFlag() bool {
	for _, v := range w.importingKeysInfo {
		if v.Complete == false {
			return true
		}
	}
	return false
}

// AttachBlock attach a new block
func (w *Wallet) AttachBlock(block *types.Block) error {
	if block.PreviousBlockHash != w.status.WorkHash {
		log.Warn("wallet skip attachBlock due to status hash not equal to previous hash")
		return nil
	}

	blockHash := block.Hash()
	txStatus, err := w.chain.GetTransactionStatus(&blockHash)
	if err != nil {
		return err
	}

	storeBatch := w.DB.NewBatch()
	w.indexTransactions(storeBatch, block, txStatus)
	w.buildAccountUTXOs(storeBatch, block, txStatus)

	w.status.WorkHeight = block.Height
	w.status.WorkHash = block.Hash()
	if w.status.WorkHeight >= w.status.BestHeight {
		w.status.BestHeight = w.status.WorkHeight
		w.status.BestHash = w.status.WorkHash
	}
	return w.commitWalletInfo(storeBatch)
}

// DetachBlock detach a block and rollback state
func (w *Wallet) DetachBlock(block *types.Block) error {
	blockHash := block.Hash()
	txStatus, err := w.chain.GetTransactionStatus(&blockHash)
	if err != nil {
		return err
	}

	storeBatch := w.DB.NewBatch()
	w.reverseAccountUTXOs(storeBatch, block, txStatus)
	w.deleteTransactions(storeBatch, w.status.BestHeight)

	w.status.BestHeight = block.Height - 1
	w.status.BestHash = block.PreviousBlockHash

	if w.status.WorkHeight > w.status.BestHeight {
		w.status.WorkHeight = w.status.BestHeight
		w.status.WorkHash = w.status.BestHash
	}

	return w.commitWalletInfo(storeBatch)
}

//WalletUpdate process every valid block and reverse every invalid block which need to rollback
func (w *Wallet) walletUpdater() {
	for {
		getRescanNotification(w)
		updateRescanStatus(w)
		for !w.chain.InMainChain(w.status.BestHash) {
			block, err := w.chain.GetBlockByHash(&w.status.BestHash)
			if err != nil {
				log.WithField("err", err).Error("walletUpdater GetBlockByHash")
				return
			}

			if err := w.DetachBlock(block); err != nil {
				log.WithField("err", err).Error("walletUpdater detachBlock")
				return
			}
		}

		block, _ := w.chain.GetBlockByHeight(w.status.WorkHeight + 1)
		if block == nil {
			<-w.chain.BlockWaiter(w.status.WorkHeight + 1)
			continue
		}

		if err := w.AttachBlock(block); err != nil {
			log.WithField("err", err).Error("walletUpdater stop")
			return
		}
	}
}

func getRescanNotification(w *Wallet) {
	select {
	case <-w.rescanProgress:
		w.status.WorkHeight = 0
		block, _ := w.chain.GetBlockByHeight(w.status.WorkHeight)
		w.status.WorkHash = block.Hash()
	default:
		return
	}
}

// ExportAccountPrivKey exports the account private key as a WIF for encoding as a string
// in the Wallet Import Formt.
func (w *Wallet) ExportAccountPrivKey(xpub chainkd.XPub, auth string) (*string, error) {
	xprv, err := w.Hsm.LoadChainKDKey(xpub, auth)
	if err != nil {
		return nil, err
	}
	var hashed [32]byte
	sha3pool.Sum256(hashed[:], xprv[:])

	tmp := append(xprv[:], hashed[:4]...)
	res := base58.Encode(tmp)
	return &res, nil
}

// ImportAccountPrivKey imports the account key in the Wallet Import Formt.
func (w *Wallet) ImportAccountPrivKey(xprv chainkd.XPrv, keyAlias, auth string, index uint64, accountAlias string) (*pseudohsm.XPub, error) {
	if w.Hsm.HasAlias(keyAlias) {
		return nil, pseudohsm.ErrDuplicateKeyAlias
	}
	if w.Hsm.HasKey(xprv) {
		return nil, pseudohsm.ErrDuplicateKey
	}

	if acc, _ := w.AccountMgr.FindByAlias(nil, accountAlias); acc != nil {
		return nil, account.ErrDuplicateAlias
	}

	xpub, _, err := w.Hsm.ImportXPrvKey(auth, keyAlias, xprv)
	if err != nil {
		return nil, err
	}

	newAccount, err := w.AccountMgr.Create(nil, []chainkd.XPub{xpub.XPub}, SINGLE, accountAlias)
	if err != nil {
		w.Hsm.XDelete(xpub.XPub, auth)
		return nil, err
	}
	if err := w.recoveryAccountWalletDB(newAccount, xpub, index, keyAlias); err != nil {
		w.AccountMgr.DeleteAccount(newAccount.ID)
		w.Hsm.XDelete(xpub.XPub, auth)
		return nil, err
	}
	return xpub, nil
}

// ImportAccountXpubKey imports the account key in the Wallet Import Formt.
func (w *Wallet) ImportAccountXpubKey(xpubIndex int, xpub pseudohsm.XPub, cpIndex uint64) error {
	accountAlias := fmt.Sprintf("recovery_%d", xpubIndex)

	if acc, _ := w.AccountMgr.FindByAlias(nil, accountAlias); acc != nil {
		return account.ErrDuplicateAlias
	}

	newAccount, err := w.AccountMgr.Create(nil, []chainkd.XPub{xpub.XPub}, SINGLE, accountAlias)
	if err != nil {
		return err
	}

	return w.recoveryAccountWalletDB(newAccount, &xpub, cpIndex, xpub.Alias)
}

func (w *Wallet) recoveryAccountWalletDB(account *account.Account, XPub *pseudohsm.XPub, index uint64, keyAlias string) error {
	if err := w.createProgram(account, XPub, index); err != nil {
		return err
	}
	w.ImportingPrivateKey = true
	tmp := KeyInfo{
		Alias:    keyAlias,
		XPub:     XPub.XPub,
		Complete: false,
	}
	w.importingKeysInfo = append(w.importingKeysInfo, tmp)
	w.commitkeysInfo()
	w.rescanBlocks()

	return nil
}

func (w *Wallet) createProgram(account *account.Account, XPub *pseudohsm.XPub, index uint64) error {
	for i := uint64(0); i < index; i++ {
		if _, err := w.AccountMgr.CreateAddress(nil, account.ID, false); err != nil {
			return err
		}
	}
	return nil
}

func (w *Wallet) rescanBlocks() {
	select {
	case w.rescanProgress <- struct{}{}:
	default:
		return
	}
}

//GetRescanStatus return key import rescan status
func (w *Wallet) GetRescanStatus() ([]KeyInfo, error) {
	keysInfo := make([]KeyInfo, len(w.importingKeysInfo))

	if rawKeyInfo := w.DB.Get(dbKeyForimportingPrivateKey); rawKeyInfo != nil {
		if err := json.Unmarshal(rawKeyInfo, &keysInfo); err != nil {
			return nil, err
		}
	}

	return keysInfo, nil
}

//updateRescanStatus mark private key import process `Complete` if rescan finished
func updateRescanStatus(w *Wallet) {
	if !w.ImportingPrivateKey {
		return
	}

	if w.status.WorkHeight < w.status.BestHeight {
		percent := uint8(w.status.WorkHeight * 100 / w.status.BestHeight)
		for _, keyInfo := range w.importingKeysInfo {
			keyInfo.Percent = percent
		}
		return
	}

	w.ImportingPrivateKey = false
	for _, keyInfo := range w.importingKeysInfo {
		keyInfo.Percent = 100
		keyInfo.Complete = true
	}
	w.commitkeysInfo()
	// TODO: delete the generated but not used addresses
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package utils

import (
	"crypto/tls"
	"net/http"
	"sync"
)

var insecureClient = (*http.Client)(nil)
var insecureClientMutex = sync.Mutex{}

func init() {
	// See https://code.google.com/p/go/issues/detail?id=4677
	// We need to force the connection to close each time so that we don't
	// hit the above Go bug.
	defaultTransport := http.DefaultTransport.(*http.Transport)
	defaultTransport.DisableKeepAlives = true
	registerFileProtocol(defaultTransport)
}

// registerFileProtocol registers support for file:// URLs on the
// given transport.
func registerFileProtocol(transport *http.Transport) {
	transport.RegisterProtocol("file", http.NewFileTransport(http.Dir("/")))
}

// SSLHostnameVerification is used as a switch for when a given provider might
// use self-signed credentials and we should not try to verify the hostname on
// the TLS/SSL certificates
type SSLHostnameVerification bool

const (
	// VerifySSLHostnames ensures we verify the hostname on the certificate
	// matches the host we are connecting and is signed
	VerifySSLHostnames = SSLHostnameVerification(true)
	// NoVerifySSLHostnames informs us to skip verifying the hostname
	// matches a valid certificate
	NoVerifySSLHostnames = SSLHostnameVerification(false)
)

// GetHTTPClient returns either a standard http client or
// non validating client depending on the value of verify.
func GetHTTPClient(verify SSLHostnameVerification) *http.Client {
	if verify == VerifySSLHostnames {
		return GetValidatingHTTPClient()
	}
	return GetNonValidatingHTTPClient()
}

// GetValidatingHTTPClient returns a new http.Client that
// verifies the server's certificate chain and hostname.
func GetValidatingHTTPClient() *http.Client {
	logger.Infof("hostname SSL verification enabled")
	return http.DefaultClient
}

// GetNonValidatingHTTPClient returns a new http.Client that
// does not verify the server's certificate chain and hostname.
func GetNonValidatingHTTPClient() *http.Client {
	logger.Infof("hostname SSL verification disabled")
	insecureClientMutex.Lock()
	defer insecureClientMutex.Unlock()
	if insecureClient == nil {
		insecureConfig := &tls.Config{InsecureSkipVerify: true}
		insecureTransport := NewHttpTLSTransport(insecureConfig)
		insecureClient = &http.Client{Transport: insecureTransport}
	}
	return insecureClient
}

// NewHttpTLSTransport returns a new http.Transport constructed with the TLS config
// and the necessary parameters for Juju.
func NewHttpTLSTransport(tlsConfig *tls.Config) *http.Transport {
	// See https://code.google.com/p/go/issues/detail?id=4677
	// We need to force the connection to close each time so that we don't
	// hit the above Go bug.
	transport := &http.Transport{
		TLSClientConfig:   tlsConfig,
		DisableKeepAlives: true,
	}
	registerFileProtocol(transport)
	return transport
}
<commit_msg>Fix comment<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package utils

import (
	"crypto/tls"
	"net/http"
	"sync"
)

var insecureClient = (*http.Client)(nil)
var insecureClientMutex = sync.Mutex{}

func init() {
	// See https://code.google.com/p/go/issues/detail?id=4677
	// We need to force the connection to close each time so that we don't
	// hit the above Go bug.
	defaultTransport := http.DefaultTransport.(*http.Transport)
	defaultTransport.DisableKeepAlives = true
	registerFileProtocol(defaultTransport)
}

// registerFileProtocol registers support for file:// URLs on the given transport.
func registerFileProtocol(transport *http.Transport) {
	transport.RegisterProtocol("file", http.NewFileTransport(http.Dir("/")))
}

// SSLHostnameVerification is used as a switch for when a given provider might
// use self-signed credentials and we should not try to verify the hostname on
// the TLS/SSL certificates
type SSLHostnameVerification bool

const (
	// VerifySSLHostnames ensures we verify the hostname on the certificate
	// matches the host we are connecting and is signed
	VerifySSLHostnames = SSLHostnameVerification(true)
	// NoVerifySSLHostnames informs us to skip verifying the hostname
	// matches a valid certificate
	NoVerifySSLHostnames = SSLHostnameVerification(false)
)

// GetHTTPClient returns either a standard http client or
// non validating client depending on the value of verify.
func GetHTTPClient(verify SSLHostnameVerification) *http.Client {
	if verify == VerifySSLHostnames {
		return GetValidatingHTTPClient()
	}
	return GetNonValidatingHTTPClient()
}

// GetValidatingHTTPClient returns a new http.Client that
// verifies the server's certificate chain and hostname.
func GetValidatingHTTPClient() *http.Client {
	logger.Infof("hostname SSL verification enabled")
	return http.DefaultClient
}

// GetNonValidatingHTTPClient returns a new http.Client that
// does not verify the server's certificate chain and hostname.
func GetNonValidatingHTTPClient() *http.Client {
	logger.Infof("hostname SSL verification disabled")
	insecureClientMutex.Lock()
	defer insecureClientMutex.Unlock()
	if insecureClient == nil {
		insecureConfig := &tls.Config{InsecureSkipVerify: true}
		insecureTransport := NewHttpTLSTransport(insecureConfig)
		insecureClient = &http.Client{Transport: insecureTransport}
	}
	return insecureClient
}

// NewHttpTLSTransport returns a new http.Transport constructed with the TLS config
// and the necessary parameters for Juju.
func NewHttpTLSTransport(tlsConfig *tls.Config) *http.Transport {
	// See https://code.google.com/p/go/issues/detail?id=4677
	// We need to force the connection to close each time so that we don't
	// hit the above Go bug.
	transport := &http.Transport{
		TLSClientConfig:   tlsConfig,
		DisableKeepAlives: true,
	}
	registerFileProtocol(transport)
	return transport
}
<|endoftext|>


<commit_before>package utils

import (
	"errors"

	e "github.com/techjanitor/pram-get/errors"
)

var (
	ErrUserNotConfirmed error = errors.New("Account not confirmed")
	ErrUserBanned       error = errors.New("Account banned")
	ErrUserLocked       error = errors.New("Account locked")
	userdataWorker      *userWorker
)

// struct for database insert worker
type userWorker struct {
	queue chan *User
}

// user struct
type User struct {
	Id              uint   `json:"id"`
	Name            string `json:"name"`
	Email           string `json:"email"`
	Group           uint   `json:"group"`
	IsConfirmed     bool   `json:"-"`
	IsLocked        bool   `json:"-"`
	IsBanned        bool   `json:"-"`
	IsAuthenticated bool   `json:"-"`
}

func init() {
	// make worker channel
	userdataWorker = &userWorker{
		make(chan *User, 64),
	}

	go func() {

		// Get Database handle
		db, err := GetDb()
		if err != nil {
			return
		}

		// prepare query for users table
		ps1, err := db.Prepare("SELECT usergroup_id,user_name,user_email,user_confirmed,user_locked,user_banned FROM users WHERE user_id = ?")
		if err != nil {
			return
		}

		// range through tasks channel
		for u := range userdataWorker.queue {

			// input data
			err = ps1.QueryRow(u.Id).Scan(&u.Group, &u.Name, &u.Email, &u.IsConfirmed, &u.IsLocked, &u.IsBanned)
			if err != nil {
				return
			}

			userdataWorker.queue <- u

		}

	}()

}

// get the user info from id
func (u *User) Info() (err error) {

	// this needs an id
	if u.Id == 0 || u.Id == 1 {
		return e.ErrInvalidParam
	}

	userdataWorker.queue <- u

	// if account is not confirmed
	if !u.IsConfirmed {
		return ErrUserNotConfirmed
	}

	// if locked
	if u.IsLocked {
		return ErrUserLocked
	}

	// if banned
	if u.IsBanned {
		return ErrUserBanned
	}

	// mark authenticated
	u.IsAuthenticated = true

	return

}
<commit_msg>test async user info<commit_after>package utils

import (
	"errors"
	"fmt"

	e "github.com/techjanitor/pram-get/errors"
)

var (
	ErrUserNotConfirmed error = errors.New("Account not confirmed")
	ErrUserBanned       error = errors.New("Account banned")
	ErrUserLocked       error = errors.New("Account locked")
	userdataWorker      *userWorker
)

// struct for database insert worker
type userWorker struct {
	queue chan *User
}

// user struct
type User struct {
	Id              uint   `json:"id"`
	Name            string `json:"name"`
	Email           string `json:"email"`
	Group           uint   `json:"group"`
	IsConfirmed     bool   `json:"-"`
	IsLocked        bool   `json:"-"`
	IsBanned        bool   `json:"-"`
	IsAuthenticated bool   `json:"-"`
}

func init() {
	// make worker channel
	userdataWorker = &userWorker{
		make(chan *User, 64),
	}

	go func() {

		// Get Database handle
		db, err := GetDb()
		if err != nil {
			return
		}

		// prepare query for users table
		ps1, err := db.Prepare("SELECT usergroup_id,user_name,user_email,user_confirmed,user_locked,user_banned FROM users WHERE user_id = ?")
		if err != nil {
			return
		}

		// range through tasks channel
		for u := range userdataWorker.queue {

			// input data
			err = ps1.QueryRow(u.Id).Scan(&u.Group, &u.Name, &u.Email, &u.IsConfirmed, &u.IsLocked, &u.IsBanned)
			if err != nil {
				return
			}

			userdataWorker.queue <- u

		}

	}()

}

// get the user info from id
func (u *User) Info() (err error) {

	// this needs an id
	if u.Id == 0 || u.Id == 1 {
		return e.ErrInvalidParam
	}

	userdataWorker.queue <- u

	fmt.Printf("%s\n", u)

	// if account is not confirmed
	if !u.IsConfirmed {
		return ErrUserNotConfirmed
	}

	// if locked
	if u.IsLocked {
		return ErrUserLocked
	}

	// if banned
	if u.IsBanned {
		return ErrUserBanned
	}

	// mark authenticated
	u.IsAuthenticated = true

	return

}
<|endoftext|>


<commit_before>package taskrunner

import (
	"context"
	"strings"

	log "github.com/hashicorp/go-hclog"
	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
	cconfig "github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/taskenv"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/plugins/drivers"
)

const (
	// TaskDirHookIsDoneDataKey is used to mark whether the hook is done. We
	// do not use the Done response value because we still need to set the
	// environment variables every time a task starts.
	// TODO(0.9.1): Use the resp.Env map and switch to resp.Done. We need to
	// remove usage of the envBuilder
	TaskDirHookIsDoneDataKey = "is_done"
)

type taskDirHook struct {
	runner *TaskRunner
	logger log.Logger
}

func newTaskDirHook(runner *TaskRunner, logger log.Logger) *taskDirHook {
	td := &taskDirHook{
		runner: runner,
	}
	td.logger = logger.Named(td.Name())
	return td
}

func (h *taskDirHook) Name() string {
	// Copied in client/state when upgrading from <0.9 schemas, so if you
	// change it here you also must change it there.
	return "task_dir"
}

func (h *taskDirHook) Prestart(ctx context.Context, req *interfaces.TaskPrestartRequest, resp *interfaces.TaskPrestartResponse) error {
	if v, ok := req.HookData[TaskDirHookIsDoneDataKey]; ok && v == "true" {
		fsi := h.runner.driverCapabilities.FSIsolation
		setEnvvars(h.runner.envBuilder, fsi, h.runner.taskDir, h.runner.clientConfig)
		resp.HookData = map[string]string{
			TaskDirHookIsDoneDataKey: "true",
		}
		return nil
	}

	cc := h.runner.clientConfig
	chroot := cconfig.DefaultChrootEnv
	if len(cc.ChrootEnv) > 0 {
		chroot = cc.ChrootEnv
	}

	// Emit the event that we are going to be building the task directory
	h.runner.EmitEvent(structs.NewTaskEvent(structs.TaskSetup).SetMessage(structs.TaskBuildingTaskDir))

	// Build the task directory structure
	fsi := h.runner.driverCapabilities.FSIsolation
	err := h.runner.taskDir.Build(fsi == drivers.FSIsolationChroot, chroot)
	if err != nil {
		return err
	}

	// Update the environment variables based on the built task directory
	setEnvvars(h.runner.envBuilder, fsi, h.runner.taskDir, h.runner.clientConfig)
	resp.HookData = map[string]string{
		TaskDirHookIsDoneDataKey: "true",
	}
	return nil
}

// setEnvvars sets path and host env vars depending on the FS isolation used.
func setEnvvars(envBuilder *taskenv.Builder, fsi drivers.FSIsolation, taskDir *allocdir.TaskDir, conf *cconfig.Config) {
	// Set driver-specific environment variables
	switch fsi {
	case drivers.FSIsolationNone:
		// Use host paths
		envBuilder.SetAllocDir(taskDir.SharedAllocDir)
		envBuilder.SetTaskLocalDir(taskDir.LocalDir)
		envBuilder.SetSecretsDir(taskDir.SecretsDir)
	default:
		// filesystem isolation; use container paths
		envBuilder.SetAllocDir(allocdir.SharedAllocContainerPath)
		envBuilder.SetTaskLocalDir(allocdir.TaskLocalContainerPath)
		envBuilder.SetSecretsDir(allocdir.TaskSecretsContainerPath)
	}

	// Set the host environment variables for non-image based drivers
	if fsi != drivers.FSIsolationImage {
		filter := strings.Split(conf.ReadDefault("env.blacklist", cconfig.DefaultEnvBlacklist), ",")
		envBuilder.SetHostEnvvars(filter)
	}
}
<commit_msg>Fix usage of fsi variable<commit_after>package taskrunner

import (
	"context"
	"strings"

	log "github.com/hashicorp/go-hclog"
	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
	cconfig "github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/taskenv"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/plugins/drivers"
)

const (
	// TaskDirHookIsDoneDataKey is used to mark whether the hook is done. We
	// do not use the Done response value because we still need to set the
	// environment variables every time a task starts.
	// TODO(0.9.1): Use the resp.Env map and switch to resp.Done. We need to
	// remove usage of the envBuilder
	TaskDirHookIsDoneDataKey = "is_done"
)

type taskDirHook struct {
	runner *TaskRunner
	logger log.Logger
}

func newTaskDirHook(runner *TaskRunner, logger log.Logger) *taskDirHook {
	td := &taskDirHook{
		runner: runner,
	}
	td.logger = logger.Named(td.Name())
	return td
}

func (h *taskDirHook) Name() string {
	// Copied in client/state when upgrading from <0.9 schemas, so if you
	// change it here you also must change it there.
	return "task_dir"
}

func (h *taskDirHook) Prestart(ctx context.Context, req *interfaces.TaskPrestartRequest, resp *interfaces.TaskPrestartResponse) error {
	fsi := h.runner.driverCapabilities.FSIsolation
	if v, ok := req.HookData[TaskDirHookIsDoneDataKey]; ok && v == "true" {
		setEnvvars(h.runner.envBuilder, fsi, h.runner.taskDir, h.runner.clientConfig)
		resp.HookData = map[string]string{
			TaskDirHookIsDoneDataKey: "true",
		}
		return nil
	}

	cc := h.runner.clientConfig
	chroot := cconfig.DefaultChrootEnv
	if len(cc.ChrootEnv) > 0 {
		chroot = cc.ChrootEnv
	}

	// Emit the event that we are going to be building the task directory
	h.runner.EmitEvent(structs.NewTaskEvent(structs.TaskSetup).SetMessage(structs.TaskBuildingTaskDir))

	// Build the task directory structure
	err := h.runner.taskDir.Build(fsi == drivers.FSIsolationChroot, chroot)
	if err != nil {
		return err
	}

	// Update the environment variables based on the built task directory
	setEnvvars(h.runner.envBuilder, fsi, h.runner.taskDir, h.runner.clientConfig)
	resp.HookData = map[string]string{
		TaskDirHookIsDoneDataKey: "true",
	}
	return nil
}

// setEnvvars sets path and host env vars depending on the FS isolation used.
func setEnvvars(envBuilder *taskenv.Builder, fsi drivers.FSIsolation, taskDir *allocdir.TaskDir, conf *cconfig.Config) {
	// Set driver-specific environment variables
	switch fsi {
	case drivers.FSIsolationNone:
		// Use host paths
		envBuilder.SetAllocDir(taskDir.SharedAllocDir)
		envBuilder.SetTaskLocalDir(taskDir.LocalDir)
		envBuilder.SetSecretsDir(taskDir.SecretsDir)
	default:
		// filesystem isolation; use container paths
		envBuilder.SetAllocDir(allocdir.SharedAllocContainerPath)
		envBuilder.SetTaskLocalDir(allocdir.TaskLocalContainerPath)
		envBuilder.SetSecretsDir(allocdir.TaskSecretsContainerPath)
	}

	// Set the host environment variables for non-image based drivers
	if fsi != drivers.FSIsolationImage {
		filter := strings.Split(conf.ReadDefault("env.blacklist", cconfig.DefaultEnvBlacklist), ",")
		envBuilder.SetHostEnvvars(filter)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

// Package webhook implements a generic HTTP webhook plugin.
package webhook

import (
	"context"
	"fmt"
	"time"

	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/runtime/serializer"
	"k8s.io/apimachinery/pkg/util/net"
	utilnet "k8s.io/apimachinery/pkg/util/net"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

// defaultRequestTimeout is set for all webhook request. This is the absolute
// timeout of the HTTP request, including reading the response body.
const defaultRequestTimeout = 30 * time.Second

// GenericWebhook defines a generic client for webhooks with commonly used capabilities,
// such as retry requests.
type GenericWebhook struct {
	RestClient     *rest.RESTClient
	InitialBackoff time.Duration
	ShouldRetry    func(error) bool
}

// DefaultShouldRetry is a default implementation for the GenericWebhook ShouldRetry function property.
// If the error reason is one of: networking (connection reset) or http (InternalServerError (500), GatewayTimeout (504), TooManyRequests (429)),
// or apierrors.SuggestsClientDelay() returns true, then the function advises a retry.
// Otherwise it returns false for an immediate fail.
func DefaultShouldRetry(err error) bool {
	// these errors indicate a transient error that should be retried.
	if net.IsConnectionReset(err) || apierrors.IsInternalError(err) || apierrors.IsTimeout(err) || apierrors.IsTooManyRequests(err) {
		return true
	}
	// if the error sends the Retry-After header, we respect it as an explicit confirmation we should retry.
	if _, shouldRetry := apierrors.SuggestsClientDelay(err); shouldRetry {
		return true
	}
	return false
}

// NewGenericWebhook creates a new GenericWebhook from the provided kubeconfig file.
func NewGenericWebhook(scheme *runtime.Scheme, codecFactory serializer.CodecFactory, kubeConfigFile string, groupVersions []schema.GroupVersion, initialBackoff time.Duration, customDial utilnet.DialFunc) (*GenericWebhook, error) {
	return newGenericWebhook(scheme, codecFactory, kubeConfigFile, groupVersions, initialBackoff, defaultRequestTimeout, customDial)
}

func newGenericWebhook(scheme *runtime.Scheme, codecFactory serializer.CodecFactory, kubeConfigFile string, groupVersions []schema.GroupVersion, initialBackoff, requestTimeout time.Duration, customDial utilnet.DialFunc) (*GenericWebhook, error) {
	for _, groupVersion := range groupVersions {
		if !scheme.IsVersionRegistered(groupVersion) {
			return nil, fmt.Errorf("webhook plugin requires enabling extension resource: %s", groupVersion)
		}
	}

	loadingRules := clientcmd.NewDefaultClientConfigLoadingRules()
	loadingRules.ExplicitPath = kubeConfigFile
	loader := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(loadingRules, &clientcmd.ConfigOverrides{})

	clientConfig, err := loader.ClientConfig()
	if err != nil {
		return nil, err
	}

	// Kubeconfigs can't set a timeout, this can only be set through a command line flag.
	//
	// https://github.com/kubernetes/client-go/blob/master/tools/clientcmd/overrides.go
	//
	// Set this to something reasonable so request to webhooks don't hang forever.
	clientConfig.Timeout = requestTimeout

	// Avoid client-side rate limiting talking to the webhook backend.
	// Rate limiting should happen when deciding how many requests to serve.
	clientConfig.QPS = -1

	codec := codecFactory.LegacyCodec(groupVersions...)
	clientConfig.ContentConfig.NegotiatedSerializer = serializer.NegotiatedSerializerWrapper(runtime.SerializerInfo{Serializer: codec})

	clientConfig.Dial = customDial

	restClient, err := rest.UnversionedRESTClientFor(clientConfig)
	if err != nil {
		return nil, err
	}

	return &GenericWebhook{restClient, initialBackoff, DefaultShouldRetry}, nil
}

// WithExponentialBackoff will retry webhookFn() up to 5 times with exponentially increasing backoff when
// it returns an error for which this GenericWebhook's ShouldRetry function returns true, confirming it to
// be retriable. If no ShouldRetry has been defined for the webhook, then the default one is used (DefaultShouldRetry).
func (g *GenericWebhook) WithExponentialBackoff(ctx context.Context, webhookFn func() rest.Result) rest.Result {
	var result rest.Result
	shouldRetry := g.ShouldRetry
	if shouldRetry == nil {
		shouldRetry = DefaultShouldRetry
	}
	WithExponentialBackoff(ctx, g.InitialBackoff, func() error {
		result = webhookFn()
		return result.Error()
	}, shouldRetry)
	return result
}

// WithExponentialBackoff will retry webhookFn up to 5 times with exponentially increasing backoff when
// it returns an error for which shouldRetry returns true, confirming it to be retriable.
func WithExponentialBackoff(ctx context.Context, initialBackoff time.Duration, webhookFn func() error, shouldRetry func(error) bool) error {
	backoff := wait.Backoff{
		Duration: initialBackoff,
		Factor:   1.5,
		Jitter:   0.2,
		Steps:    5,
	}

	var err error
	wait.ExponentialBackoff(backoff, func() (bool, error) {
		err = webhookFn()
		if ctx.Err() != nil {
			// we timed out or were cancelled, we should not retry
			return true, err
		}
		if shouldRetry(err) {
			return false, nil
		}
		if err != nil {
			return false, err
		}
		return true, nil
	})
	return err
}
<commit_msg>Remove double import of k8s.io/apimachinery/pkg/util/net<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

// Package webhook implements a generic HTTP webhook plugin.
package webhook

import (
	"context"
	"fmt"
	"time"

	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/runtime/serializer"
	utilnet "k8s.io/apimachinery/pkg/util/net"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

// defaultRequestTimeout is set for all webhook request. This is the absolute
// timeout of the HTTP request, including reading the response body.
const defaultRequestTimeout = 30 * time.Second

// GenericWebhook defines a generic client for webhooks with commonly used capabilities,
// such as retry requests.
type GenericWebhook struct {
	RestClient     *rest.RESTClient
	InitialBackoff time.Duration
	ShouldRetry    func(error) bool
}

// DefaultShouldRetry is a default implementation for the GenericWebhook ShouldRetry function property.
// If the error reason is one of: networking (connection reset) or http (InternalServerError (500), GatewayTimeout (504), TooManyRequests (429)),
// or apierrors.SuggestsClientDelay() returns true, then the function advises a retry.
// Otherwise it returns false for an immediate fail.
func DefaultShouldRetry(err error) bool {
	// these errors indicate a transient error that should be retried.
	if utilnet.IsConnectionReset(err) || apierrors.IsInternalError(err) || apierrors.IsTimeout(err) || apierrors.IsTooManyRequests(err) {
		return true
	}
	// if the error sends the Retry-After header, we respect it as an explicit confirmation we should retry.
	if _, shouldRetry := apierrors.SuggestsClientDelay(err); shouldRetry {
		return true
	}
	return false
}

// NewGenericWebhook creates a new GenericWebhook from the provided kubeconfig file.
func NewGenericWebhook(scheme *runtime.Scheme, codecFactory serializer.CodecFactory, kubeConfigFile string, groupVersions []schema.GroupVersion, initialBackoff time.Duration, customDial utilnet.DialFunc) (*GenericWebhook, error) {
	return newGenericWebhook(scheme, codecFactory, kubeConfigFile, groupVersions, initialBackoff, defaultRequestTimeout, customDial)
}

func newGenericWebhook(scheme *runtime.Scheme, codecFactory serializer.CodecFactory, kubeConfigFile string, groupVersions []schema.GroupVersion, initialBackoff, requestTimeout time.Duration, customDial utilnet.DialFunc) (*GenericWebhook, error) {
	for _, groupVersion := range groupVersions {
		if !scheme.IsVersionRegistered(groupVersion) {
			return nil, fmt.Errorf("webhook plugin requires enabling extension resource: %s", groupVersion)
		}
	}

	loadingRules := clientcmd.NewDefaultClientConfigLoadingRules()
	loadingRules.ExplicitPath = kubeConfigFile
	loader := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(loadingRules, &clientcmd.ConfigOverrides{})

	clientConfig, err := loader.ClientConfig()
	if err != nil {
		return nil, err
	}

	// Kubeconfigs can't set a timeout, this can only be set through a command line flag.
	//
	// https://github.com/kubernetes/client-go/blob/master/tools/clientcmd/overrides.go
	//
	// Set this to something reasonable so request to webhooks don't hang forever.
	clientConfig.Timeout = requestTimeout

	// Avoid client-side rate limiting talking to the webhook backend.
	// Rate limiting should happen when deciding how many requests to serve.
	clientConfig.QPS = -1

	codec := codecFactory.LegacyCodec(groupVersions...)
	clientConfig.ContentConfig.NegotiatedSerializer = serializer.NegotiatedSerializerWrapper(runtime.SerializerInfo{Serializer: codec})

	clientConfig.Dial = customDial

	restClient, err := rest.UnversionedRESTClientFor(clientConfig)
	if err != nil {
		return nil, err
	}

	return &GenericWebhook{restClient, initialBackoff, DefaultShouldRetry}, nil
}

// WithExponentialBackoff will retry webhookFn() up to 5 times with exponentially increasing backoff when
// it returns an error for which this GenericWebhook's ShouldRetry function returns true, confirming it to
// be retriable. If no ShouldRetry has been defined for the webhook, then the default one is used (DefaultShouldRetry).
func (g *GenericWebhook) WithExponentialBackoff(ctx context.Context, webhookFn func() rest.Result) rest.Result {
	var result rest.Result
	shouldRetry := g.ShouldRetry
	if shouldRetry == nil {
		shouldRetry = DefaultShouldRetry
	}
	WithExponentialBackoff(ctx, g.InitialBackoff, func() error {
		result = webhookFn()
		return result.Error()
	}, shouldRetry)
	return result
}

// WithExponentialBackoff will retry webhookFn up to 5 times with exponentially increasing backoff when
// it returns an error for which shouldRetry returns true, confirming it to be retriable.
func WithExponentialBackoff(ctx context.Context, initialBackoff time.Duration, webhookFn func() error, shouldRetry func(error) bool) error {
	backoff := wait.Backoff{
		Duration: initialBackoff,
		Factor:   1.5,
		Jitter:   0.2,
		Steps:    5,
	}

	var err error
	wait.ExponentialBackoff(backoff, func() (bool, error) {
		err = webhookFn()
		if ctx.Err() != nil {
			// we timed out or were cancelled, we should not retry
			return true, err
		}
		if shouldRetry(err) {
			return false, nil
		}
		if err != nil {
			return false, err
		}
		return true, nil
	})
	return err
}
<|endoftext|>


<commit_before>package fetcher

import (
	"fmt"
	"testing"
)

func TestGetArticle(t *testing.T) {
	a, err := GetArticle("http://www.yasuhisay.info/entry/20090516/1242480413")
	if err != nil {
		t.Error(err.Error())
	}

	if a.Title == "" {
		t.Error("Title must not be empty")
	}
	if a.Description == "" {
		t.Error("Description must not be empty")
	}
	if a.OgType != "article" {
		t.Error("OgType must be article")
	}
	if a.StatusCode != 200 {
		t.Error("StatusCode must be 200")
	}
}

func TestGetArticleWithInvalidEncoding(t *testing.T) {
	url := "http://www.atmarkit.co.jp/ait/articles/1702/20/news021.html"
	a, err := GetArticle(url)
	if err == nil {
		t.Error(fmt.Sprintf("Error must occur for this url: %s", url))
	}
}
<commit_msg>unused<commit_after>package fetcher

import (
	"fmt"
	"testing"
)

func TestGetArticle(t *testing.T) {
	a, err := GetArticle("http://www.yasuhisay.info/entry/20090516/1242480413")
	if err != nil {
		t.Error(err.Error())
	}

	if a.Title == "" {
		t.Error("Title must not be empty")
	}
	if a.Description == "" {
		t.Error("Description must not be empty")
	}
	if a.OgType != "article" {
		t.Error("OgType must be article")
	}
	if a.StatusCode != 200 {
		t.Error("StatusCode must be 200")
	}
}

func TestGetArticleWithInvalidEncoding(t *testing.T) {
	url := "http://www.atmarkit.co.jp/ait/articles/1702/20/news021.html"
	_, err := GetArticle(url)
	if err == nil {
		t.Error(fmt.Sprintf("Error must occur for this url: %s", url))
	}
}
<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"fmt"
	"net/http"
	"regexp"
	"time"

	"github.com/nicklaw5/helix"
	log "github.com/sirupsen/logrus"
)

type channelPointRedemption struct {
	Subscription struct {
		ID        string `json:"id"`
		Status    string `json:"status"`
		Type      string `json:"type"`
		Version   string `json:"version"`
		Condition struct {
			BroadcasterUserID string `json:"broadcaster_user_id"`
			RewardID          string `json:"reward_id"`
		} `json:"condition"`
		Transport struct {
			Method   string `json:"method"`
			Callback string `json:"callback"`
		} `json:"transport"`
		CreatedAt time.Time `json:"created_at"`
		Cost      int       `json:"cost"`
	} `json:"subscription"`
	Event struct {
		BroadcasterUserID    string    `json:"broadcaster_user_id"`
		BroadcasterUserLogin string    `json:"broadcaster_user_login"`
		BroadcasterUserName  string    `json:"broadcaster_user_name"`
		ID                   string    `json:"id"`
		UserID               string    `json:"user_id"`
		UserLogin            string    `json:"user_login"`
		UserName             string    `json:"user_name"`
		UserInput            string    `json:"user_input"`
		Status               string    `json:"status"`
		RedeemedAt           time.Time `json:"redeemed_at"`
		Reward               struct {
			ID     string `json:"id"`
			Title  string `json:"title"`
			Prompt string `json:"prompt"`
			Cost   int    `json:"cost"`
		} `json:"reward"`
	} `json:"event"`
}

var bttvRegex = regexp.MustCompile(`https?:\/\/betterttv.com\/emotes\/(\w*)`)

func (s *Server) subscribeChannelPoints() {
	// Twitch doesn't need a user token here, always an app token eventhough the user has to authenticate beforehand.
	// Internally they check if the app token has authenticated users
	// s.helixUserClient.Client.SetUserAccessToken(s.store.Client.HGet("accessToken", "77829817").Val())
	response, err := s.helixUserClient.Client.CreateEventSubSubscription(
		&helix.EventSubSubscription{
			Condition: helix.EventSubCondition{BroadcasterUserID: "77829817"},
			Transport: helix.EventSubTransport{Method: "webhook", Callback: s.cfg.HttpBaseUrl + "/api/redemption", Secret: s.cfg.Secret},
			Type:      "channel.channel_points_custom_reward_redemption.add",
			Version:   "1",
		},
	)
	if err != nil {
		log.Error(err)
	}

	log.Info(response)
}

func (s *Server) handleChannelPointsRedemption(w http.ResponseWriter, r *http.Request) {
	// verified := helix.VerifyEventSubNotification(s.cfg.Secret, r.Header, "")
	// if !verified {
	// 	http.Error(w, "failed verfication", http.StatusPreconditionFailed)
	// }

	if r.Header.Get("Twitch-Eventsub-Message-Type") == "webhook_callback_verification" {
		s.handleChallenge(w, r)
		return
	}

	var redemption channelPointRedemption
	err := json.NewDecoder(r.Body).Decode(&redemption)
	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	matches := bttvRegex.FindAllStringSubmatch(redemption.Event.UserInput, -1)
	if len(matches) == 1 && len(matches[0]) == 2 {
		err = s.emotechief.SetEmote(redemption.Event.BroadcasterUserID, matches[0][1])
		if err != nil {
			log.Warn(err)
		}
		fmt.Fprint(w, "success")
		return
	}

	log.Warnf("Could not find emote in message: %s", redemption.Event.UserInput)
	http.Error(w, "Could not find emote", http.StatusBadRequest)
}

func (s *Server) handleChallenge(w http.ResponseWriter, r *http.Request) {
	var event struct{ Challenge string }
	err := json.NewDecoder(r.Body).Decode(&event)
	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	log.Infof("Challenge success: %s", event.Challenge)
	fmt.Fprint(w, event.Challenge)
}
<commit_msg>handle verfication<commit_after>package server

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"regexp"
	"time"

	"github.com/nicklaw5/helix"
	log "github.com/sirupsen/logrus"
)

type channelPointRedemption struct {
	Subscription struct {
		ID        string `json:"id"`
		Status    string `json:"status"`
		Type      string `json:"type"`
		Version   string `json:"version"`
		Condition struct {
			BroadcasterUserID string `json:"broadcaster_user_id"`
			RewardID          string `json:"reward_id"`
		} `json:"condition"`
		Transport struct {
			Method   string `json:"method"`
			Callback string `json:"callback"`
		} `json:"transport"`
		CreatedAt time.Time `json:"created_at"`
		Cost      int       `json:"cost"`
	} `json:"subscription"`
	Event struct {
		BroadcasterUserID    string    `json:"broadcaster_user_id"`
		BroadcasterUserLogin string    `json:"broadcaster_user_login"`
		BroadcasterUserName  string    `json:"broadcaster_user_name"`
		ID                   string    `json:"id"`
		UserID               string    `json:"user_id"`
		UserLogin            string    `json:"user_login"`
		UserName             string    `json:"user_name"`
		UserInput            string    `json:"user_input"`
		Status               string    `json:"status"`
		RedeemedAt           time.Time `json:"redeemed_at"`
		Reward               struct {
			ID     string `json:"id"`
			Title  string `json:"title"`
			Prompt string `json:"prompt"`
			Cost   int    `json:"cost"`
		} `json:"reward"`
	} `json:"event"`
}

var bttvRegex = regexp.MustCompile(`https?:\/\/betterttv.com\/emotes\/(\w*)`)

func (s *Server) subscribeChannelPoints() {
	// Twitch doesn't need a user token here, always an app token eventhough the user has to authenticate beforehand.
	// Internally they check if the app token has authenticated users
	// s.helixUserClient.Client.SetUserAccessToken(s.store.Client.HGet("accessToken", "77829817").Val())
	response, err := s.helixUserClient.Client.CreateEventSubSubscription(
		&helix.EventSubSubscription{
			Condition: helix.EventSubCondition{BroadcasterUserID: "77829817"},
			Transport: helix.EventSubTransport{Method: "webhook", Callback: s.cfg.HttpBaseUrl + "/api/redemption", Secret: s.cfg.Secret},
			Type:      "channel.channel_points_custom_reward_redemption.add",
			Version:   "1",
		},
	)
	if err != nil {
		log.Error(err)
	}

	log.Info(response)
}

func (s *Server) handleChannelPointsRedemption(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		log.Error(err)
		http.Error(w, "failed reading body", http.StatusBadRequest)
	}

	verified := helix.VerifyEventSubNotification(s.cfg.Secret, r.Header, string(body))
	if !verified {
		log.Errorf("Failed verification: %s", r.Header.Get("Twitch-Eventsub-Message-Id"))
		http.Error(w, "failed verfication", http.StatusPreconditionFailed)
	}

	if r.Header.Get("Twitch-Eventsub-Message-Type") == "webhook_callback_verification" {
		s.handleChallenge(w, r)
		return
	}

	var redemption channelPointRedemption
	err = json.Unmarshal(body, &redemption)
	if err != nil {
		log.Error(err)
		http.Error(w, "Failed decoding body: "+err.Error(), http.StatusBadRequest)
		return
	}

	matches := bttvRegex.FindAllStringSubmatch(redemption.Event.UserInput, -1)
	if len(matches) == 1 && len(matches[0]) == 2 {
		err = s.emotechief.SetEmote(redemption.Event.BroadcasterUserID, matches[0][1])
		if err != nil {
			log.Warn(err)
		}
		fmt.Fprint(w, "success")
		return
	}

	log.Warnf("Could not find emote in message: %s", redemption.Event.UserInput)
	http.Error(w, "Could not find emote", http.StatusBadRequest)
}

func (s *Server) handleChallenge(w http.ResponseWriter, r *http.Request) {
	var event struct{ Challenge string }
	err := json.NewDecoder(r.Body).Decode(&event)
	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	log.Infof("Challenge success: %s", event.Challenge)
	fmt.Fprint(w, event.Challenge)
}
<|endoftext|>


<commit_before>// Copyright 2017 The Ebiten Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build js

package mp3

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"

	"github.com/gopherjs/gopherjs/js"
	"github.com/hajimehoshi/ebiten/audio"
)

// TODO: This just uses decodeAudioData, that can treat audio files other than MP3.

type Stream struct {
	leftData   []float32
	rightData  []float32
	posInBytes int
}

var errTryAgain = errors.New("audio/mp3: try again")

func (s *Stream) Read(b []byte) (int, error) {
	l := len(s.leftData)*4 - s.posInBytes
	if l > len(b) {
		l = len(b)
	}
	l = l / 4 * 4
	const max = 1<<15 - 1
	for i := 0; i < l/4; i++ {
		il := int32(s.leftData[s.posInBytes/4+i] * max)
		if il > max {
			il = max
		}
		if il < -max {
			il = -max
		}
		ir := int32(s.rightData[s.posInBytes/4+i] * max)
		if ir > max {
			ir = max
		}
		if ir < -max {
			ir = -max
		}
		b[4*i] = uint8(il)
		b[4*i+1] = uint8(il >> 8)
		b[4*i+2] = uint8(ir)
		b[4*i+3] = uint8(ir >> 8)
	}
	s.posInBytes += l
	if s.posInBytes == len(s.leftData)*4 {
		return l, io.EOF
	}
	return l, nil
}

func (s *Stream) Seek(offset int64, whence int) (int64, error) {
	next := int64(0)
	switch whence {
	case 0:
		next = offset
	case 1:
		next = int64(s.posInBytes) + offset
	case 2:
		next = s.Size() + offset
	}
	s.posInBytes = int(next)
	return next, nil
}

func (s *Stream) Close() error {
	return nil
}

func (s *Stream) Size() int64 {
	return int64(len(s.leftData) * 4)
}

// seekNextFrame seeks the next frame and returns the new buffer with the new position.
// seekNextFrame also returns true when seeking is successful, or false otherwise.
//
// Seeking is necessary when decoding fails. Safari's MP3 decoder can't treat IDs well (#438).
func seekNextFrame(buf []uint8) ([]uint8, bool) {
	// TODO: Need to skip tags explicitly? (hajimehoshi/go-mp3#9)

	if len(buf) < 1 {
		return nil, false
	}
	buf = buf[1:]

	for {
		if buf[0] == 0xff && buf[1]&0xfe == 0xfe {
			break
		}
		buf = buf[1:]
		if len(buf) < 2 {
			return nil, false
		}
	}
	return buf, true
}

func Decode(context *audio.Context, src audio.ReadSeekCloser) (*Stream, error) {
	b, err := ioutil.ReadAll(src)
	if err != nil {
		return nil, err
	}
	if err := src.Close(); err != nil {
		return nil, err
	}
	var s *Stream
	for {
		s, err = decode(context, b)
		if err == errTryAgain {
			buf, ok := seekNextFrame(b)
			if !ok {
				return nil, fmt.Errorf("audio/mp3: Decode failed: invalid format?")
			}
			b = buf
			continue
		}
		if err != nil {
			return nil, err
		}
		break
	}
	return s, nil
}

var offlineAudioContextClass *js.Object

func init() {
	if klass := js.Global.Get("OfflineAudioContext"); klass != js.Undefined {
		offlineAudioContextClass = klass
		return
	}
	if klass := js.Global.Get("webkitOfflineAudioContext"); klass != js.Undefined {
		offlineAudioContextClass = klass
		return
	}
}

func decode(context *audio.Context, buf []uint8) (*Stream, error) {
	if offlineAudioContextClass == nil {
		return nil, errors.New("audio/mp3: OfflineAudioContext is not available")
	}

	s := &Stream{}
	ch := make(chan error)

	// TODO: webkitOfflineAudioContext's constructor might causes 'Syntax' error
	// when the sample rate is not so usual like 22050 on Safari.
	// Should this be handled before calling this?

	// TODO: 1 is a correct second argument?
	oc := offlineAudioContextClass.New(2, 1, context.SampleRate())
	oc.Call("decodeAudioData", js.NewArrayBuffer(buf), func(buf *js.Object) {
		s.leftData = buf.Call("getChannelData", 0).Interface().([]float32)
		switch n := buf.Get("numberOfChannels").Int(); n {
		case 1:
			s.rightData = s.leftData
		case 2:
			s.rightData = buf.Call("getChannelData", 1).Interface().([]float32)
		default:
			ch <- fmt.Errorf("audio/mp3: number of channels must be 1 or 2 but %d", n)
		}
		close(ch)
	}, func(err *js.Object) {
		if err != nil {
			ch <- fmt.Errorf("audio/mp3: decodeAudioData failed: %v", err)
		} else {
			// On Safari, error value might be null and it is needed to retry decoding
			// from the next frame (#438).
			ch <- errTryAgain
		}
		close(ch)
	})
	if err := <-ch; err != nil {
		return nil, err
	}
	return s, nil
}
<commit_msg>audio/mp3: Force context switch when decoding<commit_after>// Copyright 2017 The Ebiten Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build js

package mp3

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"runtime"

	"github.com/gopherjs/gopherjs/js"
	"github.com/hajimehoshi/ebiten/audio"
)

// TODO: This just uses decodeAudioData, that can treat audio files other than MP3.

type Stream struct {
	leftData   []float32
	rightData  []float32
	posInBytes int
}

var errTryAgain = errors.New("audio/mp3: try again")

func (s *Stream) Read(b []byte) (int, error) {
	l := len(s.leftData)*4 - s.posInBytes
	if l > len(b) {
		l = len(b)
	}
	l = l / 4 * 4
	const max = 1<<15 - 1
	for i := 0; i < l/4; i++ {
		il := int32(s.leftData[s.posInBytes/4+i] * max)
		if il > max {
			il = max
		}
		if il < -max {
			il = -max
		}
		ir := int32(s.rightData[s.posInBytes/4+i] * max)
		if ir > max {
			ir = max
		}
		if ir < -max {
			ir = -max
		}
		b[4*i] = uint8(il)
		b[4*i+1] = uint8(il >> 8)
		b[4*i+2] = uint8(ir)
		b[4*i+3] = uint8(ir >> 8)
	}
	s.posInBytes += l
	if s.posInBytes == len(s.leftData)*4 {
		return l, io.EOF
	}
	return l, nil
}

func (s *Stream) Seek(offset int64, whence int) (int64, error) {
	next := int64(0)
	switch whence {
	case 0:
		next = offset
	case 1:
		next = int64(s.posInBytes) + offset
	case 2:
		next = s.Size() + offset
	}
	s.posInBytes = int(next)
	return next, nil
}

func (s *Stream) Close() error {
	return nil
}

func (s *Stream) Size() int64 {
	return int64(len(s.leftData) * 4)
}

// seekNextFrame seeks the next frame and returns the new buffer with the new position.
// seekNextFrame also returns true when seeking is successful, or false otherwise.
//
// Seeking is necessary when decoding fails. Safari's MP3 decoder can't treat IDs well (#438).
func seekNextFrame(buf []byte) ([]byte, bool) {
	// TODO: Need to skip tags explicitly? (hajimehoshi/go-mp3#9)

	if len(buf) < 1 {
		return nil, false
	}
	buf = buf[1:]

	for {
		if buf[0] == 0xff && buf[1]&0xfe == 0xfe {
			break
		}
		buf = buf[1:]
		if len(buf) < 2 {
			return nil, false
		}
	}
	return buf, true
}

func Decode(context *audio.Context, src audio.ReadSeekCloser) (*Stream, error) {
	b, err := ioutil.ReadAll(src)
	if err != nil {
		return nil, err
	}
	if err := src.Close(); err != nil {
		return nil, err
	}
	var s *Stream
	for {
		s, err = decode(context, b)
		if err == errTryAgain {
			buf, ok := seekNextFrame(b)
			if !ok {
				return nil, fmt.Errorf("audio/mp3: Decode failed: invalid format?")
			}
			b = buf
			continue
		}
		if err != nil {
			return nil, err
		}
		break
	}
	return s, nil
}

var offlineAudioContextClass *js.Object

func init() {
	if klass := js.Global.Get("OfflineAudioContext"); klass != js.Undefined {
		offlineAudioContextClass = klass
		return
	}
	if klass := js.Global.Get("webkitOfflineAudioContext"); klass != js.Undefined {
		offlineAudioContextClass = klass
		return
	}
}

func decode(context *audio.Context, buf []byte) (*Stream, error) {
	if offlineAudioContextClass == nil {
		return nil, errors.New("audio/mp3: OfflineAudioContext is not available")
	}

	s := &Stream{}
	ch := make(chan error)

	// TODO: webkitOfflineAudioContext's constructor might causes 'Syntax' error
	// when the sample rate is not so usual like 22050 on Safari.
	// Should this be handled before calling this?

	// TODO: 1 is a correct second argument?
	oc := offlineAudioContextClass.New(2, 1, context.SampleRate())
	oc.Call("decodeAudioData", js.NewArrayBuffer(buf), func(buf *js.Object) {
		s.leftData = buf.Call("getChannelData", 0).Interface().([]float32)
		switch n := buf.Get("numberOfChannels").Int(); n {
		case 1:
			s.rightData = s.leftData
		case 2:
			s.rightData = buf.Call("getChannelData", 1).Interface().([]float32)
		default:
			ch <- fmt.Errorf("audio/mp3: number of channels must be 1 or 2 but %d", n)
		}
		close(ch)
	}, func(err *js.Object) {
		if err != nil {
			ch <- fmt.Errorf("audio/mp3: decodeAudioData failed: %v", err)
		} else {
			// On Safari, error value might be null and it is needed to retry decoding
			// from the next frame (#438).
			ch <- errTryAgain
		}
		close(ch)
	})

	// GopherJS's bug? Without Gosched(), receiving might block forever.
	runtime.Gosched()

	if err := <-ch; err != nil {
		return nil, err
	}
	return s, nil
}
<|endoftext|>