<commit_before>package main

import (
	"testing"
)

func TestMakeMovePawnPush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	to := byte(CartesianToIndex(0, 2))
	from := byte(CartesianToIndex(0, 1))
	move := Move{from, to, MoveQuiet, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) != PAWN {
		t.Fail()
	}
}

func TestMakeMoveWhitePawnDoublePush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	to := byte(CartesianToIndex(0, 3))
	from := byte(CartesianToIndex(0, 1))
	move := Move{from, to, MoveDoublePush, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) !=
		PAWN || board.EnPassant != CartesianToIndex(0, 2) {
		t.Fail()
	}
}

func TestMakeMoveBlackPawnDoublePush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	board.ToMove = BLACK
	to := byte(CartesianToIndex(0, 4))
	from := byte(CartesianToIndex(0, 6))
	move := Move{from, to, MoveDoublePush, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) !=
		PAWN || board.EnPassant != CartesianToIndex(0, 5) {
		t.Fail()
	}
}

func TestMakeMoveWhitePawnCapture(t *testing.T) {
	board, err := Parse("8/8/8/3p4/4P3/8/8/8 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	to := byte(CartesianToIndex(3, 4))
	from := byte(CartesianToIndex(4, 3))
	move := Move{from, to, MoveCapture, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) !=
		PAWN || GetSide(board.Data[to]) != WHITE {
		t.Fail()
	}
}

func TestMakeMoveBlackPawnCapture(t *testing.T) {
	board, err := Parse("8/8/8/3p4/4P3/8/8/8 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	to := byte(CartesianToIndex(4, 3))
	from := byte(CartesianToIndex(3, 4))
	move := Move{from, to, MoveCapture, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) !=
		PAWN || GetSide(board.Data[to]) != BLACK {
		t.Fail()
	}
}

func IsMoveInMoveList(t *testing.T, list []Move, from, to, kind byte) bool {
	for _, m := range list {
		if m.To == to && m.From == from && m.Kind == kind {
			return true
		}
	}
	return false
}

func TestMoveGenWhitePawnPush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(0, 1)), byte(CartesianToIndex(0, 2)), MoveQuiet) {
		t.Fail()
	}
}

func TestMoveGenBlackPawnPush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	board.ToMove = BLACK
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(0, 6)), byte(CartesianToIndex(0, 5)), MoveQuiet) {
		t.Fail()
	}
}

func TestMoveGenWhitePawnDoublePush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(0, 1)), byte(CartesianToIndex(0, 3)), MoveDoublePush) {
		t.Fail()
	}
}

func TestMoveGenBlackPawnDoublePush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	board.ToMove = BLACK
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(0, 6)), byte(CartesianToIndex(0, 4)), MoveDoublePush) {
		t.Fail()
	}
}

func TestMoveGenSlider(t *testing.T) {
	board, err := Parse("4k3/8/8/1r6/8/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(1, 4)), byte(CartesianToIndex(1, 7)), MoveQuiet) {
		t.Fail()
	}
}

func TestMoveGenNonSlider(t *testing.T) {
	board, err := Parse("4k3/8/8/1r6/8/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(4, 7)), byte(CartesianToIndex(4, 6)), MoveQuiet) {
		t.Fail()
	}
	if IsMoveInMoveList(t, moves, byte(CartesianToIndex(4, 7)), byte(CartesianToIndex(4, 3)), MoveQuiet) {
		t.Fail()
	}
}

func TestMoveGenWhitePawnCap(t *testing.T) {
	board, err := Parse("4k3/8/8/3p4/4P3/8/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(4, 3)), byte(CartesianToIndex(3, 4)), MoveCapture) {
		t.Fail()
	}
}

func TestMoveGenBlackPawnCap(t *testing.T) {
	board, err := Parse("4k3/8/8/3p4/4P3/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(3, 4)), byte(CartesianToIndex(4, 3)), MoveCapture) {
		t.Fail()
	}
}

func TestMoveGenSliderCap(t *testing.T) {
	board, err := Parse("4k3/8/8/8/1r1P4/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(1, 3)), byte(CartesianToIndex(3, 3)), MoveCapture) {
		t.Fail()
	}
}

func TestMoveGenNonSliderCap(t *testing.T) {
	board, err := Parse("4k3/8/2n5/8/3P4/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	t.Log(moves)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(2, 5)), byte(CartesianToIndex(3, 3)), MoveCapture) {
		t.Fail()
	}
}

func TestPerft(t *testing.T) {
	tperft(t, 1, 20)
	tperft(t, 2, 400)
	tperft(t, 3, 8902)
}

func tperft(t *testing.T, depth int, expected uint64) {
	result := DoDividePerft(depth)
	t.Log("Result of perft(", depth, "): ", result)
	if result != expected {
		t.FailNow()
	}
}

func TestSquareNotAttacked(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/8/3p4/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if squareattacked(board, CartesianToIndex(4, 4), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByWhitePawn(t *testing.T) {
	board, err := Parse("4k3/3P4/8/8/8/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 7), WHITE) {
		t.FailNow()
	}
}

func TestSquareAttackedByBlackPawn(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/8/3p4/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByQueenFile(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/4q3/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByQueenDiag(t *testing.T) {
	board, err := Parse("4k3/8/8/8/1q6/8/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByRook(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/4r3/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByBishop(t *testing.T) {
	board, err := Parse("4k3/8/8/8/1b6/8/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}
<commit_msg>add test for knight attacking<commit_after>package main

import (
	"testing"
)

func TestMakeMovePawnPush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	to := byte(CartesianToIndex(0, 2))
	from := byte(CartesianToIndex(0, 1))
	move := Move{from, to, MoveQuiet, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) != PAWN {
		t.Fail()
	}
}

func TestMakeMoveWhitePawnDoublePush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	to := byte(CartesianToIndex(0, 3))
	from := byte(CartesianToIndex(0, 1))
	move := Move{from, to, MoveDoublePush, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) !=
		PAWN || board.EnPassant != CartesianToIndex(0, 2) {
		t.Fail()
	}
}

func TestMakeMoveBlackPawnDoublePush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	board.ToMove = BLACK
	to := byte(CartesianToIndex(0, 4))
	from := byte(CartesianToIndex(0, 6))
	move := Move{from, to, MoveDoublePush, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) !=
		PAWN || board.EnPassant != CartesianToIndex(0, 5) {
		t.Fail()
	}
}

func TestMakeMoveWhitePawnCapture(t *testing.T) {
	board, err := Parse("8/8/8/3p4/4P3/8/8/8 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	to := byte(CartesianToIndex(3, 4))
	from := byte(CartesianToIndex(4, 3))
	move := Move{from, to, MoveCapture, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) !=
		PAWN || GetSide(board.Data[to]) != WHITE {
		t.Fail()
	}
}

func TestMakeMoveBlackPawnCapture(t *testing.T) {
	board, err := Parse("8/8/8/3p4/4P3/8/8/8 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	to := byte(CartesianToIndex(4, 3))
	from := byte(CartesianToIndex(3, 4))
	move := Move{from, to, MoveCapture, EMPTY, 0}
	MakeMove(board, &move)
	if GetPiece(board.Data[from]) != EMPTY || GetPiece(board.Data[to]) !=
		PAWN || GetSide(board.Data[to]) != BLACK {
		t.Fail()
	}
}

func IsMoveInMoveList(t *testing.T, list []Move, from, to, kind byte) bool {
	for _, m := range list {
		if m.To == to && m.From == from && m.Kind == kind {
			return true
		}
	}
	return false
}

func TestMoveGenWhitePawnPush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(0, 1)), byte(CartesianToIndex(0, 2)), MoveQuiet) {
		t.Fail()
	}
}

func TestMoveGenBlackPawnPush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	board.ToMove = BLACK
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(0, 6)), byte(CartesianToIndex(0, 5)), MoveQuiet) {
		t.Fail()
	}
}

func TestMoveGenWhitePawnDoublePush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(0, 1)), byte(CartesianToIndex(0, 3)), MoveDoublePush) {
		t.Fail()
	}
}

func TestMoveGenBlackPawnDoublePush(t *testing.T) {
	board, err := Parse(START)
	if err != nil {
		t.FailNow()
	}
	board.ToMove = BLACK
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(0, 6)), byte(CartesianToIndex(0, 4)), MoveDoublePush) {
		t.Fail()
	}
}

func TestMoveGenSlider(t *testing.T) {
	board, err := Parse("4k3/8/8/1r6/8/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(1, 4)), byte(CartesianToIndex(1, 7)), MoveQuiet) {
		t.Fail()
	}
}

func TestMoveGenNonSlider(t *testing.T) {
	board, err := Parse("4k3/8/8/1r6/8/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(4, 7)), byte(CartesianToIndex(4, 6)), MoveQuiet) {
		t.Fail()
	}
	if IsMoveInMoveList(t, moves, byte(CartesianToIndex(4, 7)), byte(CartesianToIndex(4, 3)), MoveQuiet) {
		t.Fail()
	}
}

func TestMoveGenWhitePawnCap(t *testing.T) {
	board, err := Parse("4k3/8/8/3p4/4P3/8/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(4, 3)), byte(CartesianToIndex(3, 4)), MoveCapture) {
		t.Fail()
	}
}

func TestMoveGenBlackPawnCap(t *testing.T) {
	board, err := Parse("4k3/8/8/3p4/4P3/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(3, 4)), byte(CartesianToIndex(4, 3)), MoveCapture) {
		t.Fail()
	}
}

func TestMoveGenSliderCap(t *testing.T) {
	board, err := Parse("4k3/8/8/8/1r1P4/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(1, 3)), byte(CartesianToIndex(3, 3)), MoveCapture) {
		t.Fail()
	}
}

func TestMoveGenNonSliderCap(t *testing.T) {
	board, err := Parse("4k3/8/2n5/8/3P4/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	moves := MoveGen(board)
	t.Log(moves)
	if !IsMoveInMoveList(t, moves, byte(CartesianToIndex(2, 5)), byte(CartesianToIndex(3, 3)), MoveCapture) {
		t.Fail()
	}
}

func TestPerft(t *testing.T) {
	tperft(t, 1, 20)
	tperft(t, 2, 400)
	tperft(t, 3, 8902)
}

func tperft(t *testing.T, depth int, expected uint64) {
	result := DoDividePerft(depth)
	t.Log("Result of perft(", depth, "): ", result)
	if result != expected {
		t.FailNow()
	}
}

func TestSquareNotAttacked(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/8/3p4/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if squareattacked(board, CartesianToIndex(4, 4), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByWhitePawn(t *testing.T) {
	board, err := Parse("4k3/3P4/8/8/8/8/8/4K3 b - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 7), WHITE) {
		t.FailNow()
	}
}

func TestSquareAttackedByBlackPawn(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/8/3p4/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByQueenFile(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/4q3/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByQueenDiag(t *testing.T) {
	board, err := Parse("4k3/8/8/8/1q6/8/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByRook(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/4r3/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByBishop(t *testing.T) {
	board, err := Parse("4k3/8/8/8/1b6/8/8/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}

func TestSquareAttackedByKnight(t *testing.T) {
	board, err := Parse("4k3/8/8/8/8/8/2n5/4K3 w - - 0 1")
	if err != nil {
		t.FailNow()
	}
	if !squareattacked(board, CartesianToIndex(4, 0), BLACK) {
		t.FailNow()
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2019 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"
	"github.com/decred/dcrd/dcrutil/v2"
	"github.com/decred/dcrd/wire"
	"github.com/decred/dcrwallet/errors/v2"
	"github.com/decred/dcrwallet/wallet/v3/internal/txsizes"
	"github.com/decred/dcrwallet/wallet/v3/txauthor"
	"github.com/decred/dcrwallet/wallet/v3/txrules"
	"github.com/decred/dcrwallet/wallet/v3/udb"
	"github.com/decred/dcrwallet/wallet/v3/walletdb"
	"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")

// 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)

	var updates []func(walletdb.ReadWriteTx) error

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

	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 {
		return errors.E(op, err)
	}

	w.lockedOutpointMu.Lock()
	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 {
		count = int(amount / v)
		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] }()
			}
			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, infoLog, 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)
	}
	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)
		addr, err := w.nextAddress(ctx, op, persist, 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, int(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 mixes all possible outputs of a mixing change account into
// standard denominations, creating newly mixed outputs for a mixed account.
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)
	credits, err := w.FindEligibleOutputs(ctx, changeAccount, 1, tipHeight)
	if err != nil {
		return errors.E(op, err)
	}

	var g errgroup.Group
	for i := range credits {
		if credits[i].Amount <= splitPoints[len(splitPoints)-1] {
			continue
		}
		op := &credits[i].OutPoint
		g.Go(func() error {
			err := w.MixOutput(ctx, dialTLS, csppserver, op, changeAccount, mixAccount, mixBranch)
			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>Throttle goroutines spawned by MixAccount<commit_after>// Copyright (c) 2019 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"
	"github.com/decred/dcrd/dcrutil/v2"
	"github.com/decred/dcrd/wire"
	"github.com/decred/dcrwallet/errors/v2"
	"github.com/decred/dcrwallet/wallet/v3/internal/txsizes"
	"github.com/decred/dcrwallet/wallet/v3/txauthor"
	"github.com/decred/dcrwallet/wallet/v3/txrules"
	"github.com/decred/dcrwallet/wallet/v3/udb"
	"github.com/decred/dcrwallet/wallet/v3/walletdb"
	"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")

// 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)

	var updates []func(walletdb.ReadWriteTx) error

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

	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 {
		return errors.E(op, err)
	}

	w.lockedOutpointMu.Lock()
	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 {
		count = int(amount / v)
		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] }()
			}
			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, infoLog, 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)
	}
	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)
		addr, err := w.nextAddress(ctx, op, persist, 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, int(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)
	credits, err := w.FindEligibleOutputs(ctx, changeAccount, 1, tipHeight)
	if err != nil {
		return errors.E(op, err)
	}
	unlockedCredits := credits[:0]
	for i := range credits {
		if credits[i].Amount <= splitPoints[len(splitPoints)-1] {
			continue
		}
		w.lockedOutpointMu.Lock()
		_, locked := w.lockedOutpoints[credits[i].OutPoint]
		w.lockedOutpointMu.Unlock()
		if !locked {
			unlockedCredits = append(unlockedCredits, credits[i])
		}
	}
	credits = unlockedCredits
	shuffle(len(credits), func(i, j int) {
		credits[i], credits[j] = credits[j], credits[i]
	})
	credits = credits[:32] // simple throttle
	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, 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>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package core

import (
	"context"
	"strings"
	"sync"
	"time"

	"github.com/loadimpact/k6/core/local"
	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/lib/metrics"
	"github.com/loadimpact/k6/stats"
	log "github.com/sirupsen/logrus"
	"gopkg.in/guregu/null.v3"
)

const (
	TickRate        = 1 * time.Millisecond
	MetricsRate     = 1 * time.Second
	CollectRate     = 10 * time.Millisecond
	ThresholdsRate  = 2 * time.Second
	ShutdownTimeout = 10 * time.Second

	BackoffAmount = 50 * time.Millisecond
	BackoffMax    = 10 * time.Second
)

// The Engine is the beating heart of K6.
type Engine struct {
	runLock sync.Mutex

	Executor     lib.Executor
	Options      lib.Options
	Collector    lib.Collector
	NoThresholds bool

	logger *log.Logger

	Metrics     map[string]*stats.Metric
	MetricsLock sync.RWMutex

	// Assigned to metrics upon first received sample.
	thresholds map[string]stats.Thresholds
	submetrics map[string][]*stats.Submetric

	// Are thresholds tainted?
	thresholdsTainted bool
}

func NewEngine(ex lib.Executor, o lib.Options) (*Engine, error) {
	if ex == nil {
		ex = local.New(nil)
	}

	e := &Engine{
		Executor: ex,
		Options:  o,
		Metrics:  make(map[string]*stats.Metric),
	}
	e.SetLogger(log.StandardLogger())

	if err := ex.SetVUsMax(o.VUsMax.Int64); err != nil {
		return nil, err
	}
	if err := ex.SetVUs(o.VUs.Int64); err != nil {
		return nil, err
	}
	ex.SetPaused(o.Paused.Bool)
	ex.SetStages(o.Stages)
	ex.SetEndTime(o.Duration)
	ex.SetEndIterations(o.Iterations)

	e.thresholds = o.Thresholds
	e.submetrics = make(map[string][]*stats.Submetric)
	for name := range e.thresholds {
		if !strings.Contains(name, "{") {
			continue
		}

		parent, sm := stats.NewSubmetric(name)
		e.submetrics[parent] = append(e.submetrics[parent], sm)
	}

	return e, nil
}

func (e *Engine) Run(ctx context.Context) error {
	e.runLock.Lock()
	defer e.runLock.Unlock()

	e.logger.Debug("Engine: Starting with parameters...")
	for i, st := range e.Executor.GetStages() {
		fields := make(log.Fields)
		if st.Target.Valid {
			fields["tgt"] = st.Target.Int64
		}
		if st.Duration.Valid {
			fields["d"] = st.Duration.Duration
		}
		e.logger.WithFields(fields).Debugf(" - stage #%d", i)
	}

	fields := make(log.Fields)
	if endTime := e.Executor.GetEndTime(); endTime.Valid {
		fields["time"] = endTime.Duration
	}
	if endIter := e.Executor.GetEndIterations(); endIter.Valid {
		fields["iter"] = endIter.Int64
	}
	e.logger.WithFields(fields).Debug(" - end conditions (if any)")

	collectorwg := sync.WaitGroup{}
	collectorctx, collectorcancel := context.WithCancel(context.Background())
	if e.Collector != nil {
		collectorwg.Add(1)
		go func() {
			e.Collector.Run(collectorctx)
			collectorwg.Done()
		}()
	}

	subctx, subcancel := context.WithCancel(context.Background())
	subwg := sync.WaitGroup{}

	// Run metrics emission.
	subwg.Add(1)
	go func() {
		e.runMetricsEmission(subctx)
		e.logger.Debug("Engine: Emission terminated")
		subwg.Done()
	}()

	// Run thresholds.
	if !e.NoThresholds {
		subwg.Add(1)
		go func() {
			e.runThresholds(subctx, subcancel)
			e.logger.Debug("Engine: Thresholds terminated")
			subwg.Done()
		}()
	}

	// Run the executor.
	out := make(chan []stats.Sample)
	errC := make(chan error)
	subwg.Add(1)
	go func() {
		errC <- e.Executor.Run(subctx, out)
		e.logger.Debug("Engine: Executor terminated")
		subwg.Done()
	}()

	defer func() {
		// Shut down subsystems.
		subcancel()

		// Process samples until the subsystems have shut down.
		// Filter out samples produced past the end of a test.
		go func() {
			if errC != nil {
				<-errC
				errC = nil
			}
			subwg.Wait()
			close(out)
		}()
		for samples := range out {
			e.processSamples(samples...)
		}

		// Emit final metrics.
		e.emitMetrics()

		// Process final thresholds.
		if !e.NoThresholds {
			e.processThresholds(nil)
		}

		// Finally, shut down collector.
		collectorcancel()
		collectorwg.Wait()
	}()

	for {
		select {
		case samples := <-out:
			e.processSamples(samples...)
		case err := <-errC:
			errC = nil
			if err != nil {
				e.logger.WithError(err).Debug("run: executor returned an error")
				return err
			}
			e.logger.Debug("run: executor terminated")
			return nil
		case <-ctx.Done():
			e.logger.Debug("run: context expired; exiting...")
			return nil
		}
	}
}

func (e *Engine) IsTainted() bool {
	return e.thresholdsTainted
}

func (e *Engine) SetLogger(l *log.Logger) {
	e.logger = l
	e.Executor.SetLogger(l)
}

func (e *Engine) GetLogger() *log.Logger {
	return e.logger
}

func (e *Engine) runMetricsEmission(ctx context.Context) {
	ticker := time.NewTicker(MetricsRate)
	for {
		select {
		case <-ticker.C:
			e.emitMetrics()
		case <-ctx.Done():
			return
		}
	}
}

func (e *Engine) emitMetrics() {
	t := time.Now()
	e.processSamples(
		stats.Sample{
			Time:   t,
			Metric: metrics.VUs,
			Value:  float64(e.Executor.GetVUs()),
		},
		stats.Sample{
			Time:   t,
			Metric: metrics.VUsMax,
			Value:  float64(e.Executor.GetVUsMax()),
		},
	)
}

func (e *Engine) runThresholds(ctx context.Context, abort func()) {
	ticker := time.NewTicker(ThresholdsRate)
	for {
		select {
		case <-ticker.C:
			e.processThresholds(abort)
		case <-ctx.Done():
			return
		}
	}
}

func (e *Engine) processThresholds(abort func()) {
	e.MetricsLock.Lock()
	defer e.MetricsLock.Unlock()

	t := e.Executor.GetTime()
	abortOnFail := false

	e.thresholdsTainted = false
	for _, m := range e.Metrics {
		if len(m.Thresholds.Thresholds) == 0 {
			continue
		}
		m.Tainted = null.BoolFrom(false)

		e.logger.WithField("m", m.Name).Debug("running thresholds")
		succ, err := m.Thresholds.Run(m.Sink, t)
		if err != nil {
			e.logger.WithField("m", m.Name).WithError(err).Error("Threshold error")
			continue
		}
		if !succ {
			e.logger.WithField("m", m.Name).Debug("Thresholds failed")
			m.Tainted = null.BoolFrom(true)
			e.thresholdsTainted = true
			if !abortOnFail && m.Thresholds.Abort {
				abortOnFail = true
			}
		}
	}

	if abortOnFail && abort != nil {
		abort()
	}
}

func (e *Engine) processSamples(samples ...stats.Sample) {
	if len(samples) == 0 {
		return
	}

	e.MetricsLock.Lock()
	defer e.MetricsLock.Unlock()

	for i, sample := range samples {
		m, ok := e.Metrics[sample.Metric.Name]
		if !ok {
			m = sample.Metric
			m.Thresholds = e.thresholds[m.Name]
			m.Submetrics = e.submetrics[m.Name]
			e.Metrics[m.Name] = m
		}
		m.Sink.Add(sample)

		for _, sm := range m.Submetrics {
			passing := true
			for k, v := range sm.Tags {
				if sample.Tags[k] != v {
					passing = false
					break
				}
			}
			if !passing {
				continue
			}

			if sm.Metric == nil {
				sm.Metric = stats.New(sm.Name, sample.Metric.Type, sample.Metric.Contains)
				sm.Metric.Sub = *sm
				sm.Metric.Thresholds = e.thresholds[sm.Name]
				e.Metrics[sm.Name] = sm.Metric
			}
			sm.Metric.Sink.Add(sample)
		}

		if e.Options.RunTags == nil {
			continue
		}

		if samples[i].Tags == nil {
			samples[i].Tags = e.Options.RunTags
			continue
		}

		for k, v := range e.Options.RunTags {
			samples[i].Tags[k] = v
		}
	}

	if e.Collector != nil {
		e.Collector.Collect(samples)
	}
}
<commit_msg>Adding check to prevent overwritting existing tags on samples<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package core

import (
	"context"
	"strings"
	"sync"
	"time"

	"github.com/loadimpact/k6/core/local"
	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/lib/metrics"
	"github.com/loadimpact/k6/stats"
	log "github.com/sirupsen/logrus"
	"gopkg.in/guregu/null.v3"
)

const (
	TickRate        = 1 * time.Millisecond
	MetricsRate     = 1 * time.Second
	CollectRate     = 10 * time.Millisecond
	ThresholdsRate  = 2 * time.Second
	ShutdownTimeout = 10 * time.Second

	BackoffAmount = 50 * time.Millisecond
	BackoffMax    = 10 * time.Second
)

// The Engine is the beating heart of K6.
type Engine struct {
	runLock sync.Mutex

	Executor     lib.Executor
	Options      lib.Options
	Collector    lib.Collector
	NoThresholds bool

	logger *log.Logger

	Metrics     map[string]*stats.Metric
	MetricsLock sync.RWMutex

	// Assigned to metrics upon first received sample.
	thresholds map[string]stats.Thresholds
	submetrics map[string][]*stats.Submetric

	// Are thresholds tainted?
	thresholdsTainted bool
}

func NewEngine(ex lib.Executor, o lib.Options) (*Engine, error) {
	if ex == nil {
		ex = local.New(nil)
	}

	e := &Engine{
		Executor: ex,
		Options:  o,
		Metrics:  make(map[string]*stats.Metric),
	}
	e.SetLogger(log.StandardLogger())

	if err := ex.SetVUsMax(o.VUsMax.Int64); err != nil {
		return nil, err
	}
	if err := ex.SetVUs(o.VUs.Int64); err != nil {
		return nil, err
	}
	ex.SetPaused(o.Paused.Bool)
	ex.SetStages(o.Stages)
	ex.SetEndTime(o.Duration)
	ex.SetEndIterations(o.Iterations)

	e.thresholds = o.Thresholds
	e.submetrics = make(map[string][]*stats.Submetric)
	for name := range e.thresholds {
		if !strings.Contains(name, "{") {
			continue
		}

		parent, sm := stats.NewSubmetric(name)
		e.submetrics[parent] = append(e.submetrics[parent], sm)
	}

	return e, nil
}

func (e *Engine) Run(ctx context.Context) error {
	e.runLock.Lock()
	defer e.runLock.Unlock()

	e.logger.Debug("Engine: Starting with parameters...")
	for i, st := range e.Executor.GetStages() {
		fields := make(log.Fields)
		if st.Target.Valid {
			fields["tgt"] = st.Target.Int64
		}
		if st.Duration.Valid {
			fields["d"] = st.Duration.Duration
		}
		e.logger.WithFields(fields).Debugf(" - stage #%d", i)
	}

	fields := make(log.Fields)
	if endTime := e.Executor.GetEndTime(); endTime.Valid {
		fields["time"] = endTime.Duration
	}
	if endIter := e.Executor.GetEndIterations(); endIter.Valid {
		fields["iter"] = endIter.Int64
	}
	e.logger.WithFields(fields).Debug(" - end conditions (if any)")

	collectorwg := sync.WaitGroup{}
	collectorctx, collectorcancel := context.WithCancel(context.Background())
	if e.Collector != nil {
		collectorwg.Add(1)
		go func() {
			e.Collector.Run(collectorctx)
			collectorwg.Done()
		}()
	}

	subctx, subcancel := context.WithCancel(context.Background())
	subwg := sync.WaitGroup{}

	// Run metrics emission.
	subwg.Add(1)
	go func() {
		e.runMetricsEmission(subctx)
		e.logger.Debug("Engine: Emission terminated")
		subwg.Done()
	}()

	// Run thresholds.
	if !e.NoThresholds {
		subwg.Add(1)
		go func() {
			e.runThresholds(subctx, subcancel)
			e.logger.Debug("Engine: Thresholds terminated")
			subwg.Done()
		}()
	}

	// Run the executor.
	out := make(chan []stats.Sample)
	errC := make(chan error)
	subwg.Add(1)
	go func() {
		errC <- e.Executor.Run(subctx, out)
		e.logger.Debug("Engine: Executor terminated")
		subwg.Done()
	}()

	defer func() {
		// Shut down subsystems.
		subcancel()

		// Process samples until the subsystems have shut down.
		// Filter out samples produced past the end of a test.
		go func() {
			if errC != nil {
				<-errC
				errC = nil
			}
			subwg.Wait()
			close(out)
		}()
		for samples := range out {
			e.processSamples(samples...)
		}

		// Emit final metrics.
		e.emitMetrics()

		// Process final thresholds.
		if !e.NoThresholds {
			e.processThresholds(nil)
		}

		// Finally, shut down collector.
		collectorcancel()
		collectorwg.Wait()
	}()

	for {
		select {
		case samples := <-out:
			e.processSamples(samples...)
		case err := <-errC:
			errC = nil
			if err != nil {
				e.logger.WithError(err).Debug("run: executor returned an error")
				return err
			}
			e.logger.Debug("run: executor terminated")
			return nil
		case <-ctx.Done():
			e.logger.Debug("run: context expired; exiting...")
			return nil
		}
	}
}

func (e *Engine) IsTainted() bool {
	return e.thresholdsTainted
}

func (e *Engine) SetLogger(l *log.Logger) {
	e.logger = l
	e.Executor.SetLogger(l)
}

func (e *Engine) GetLogger() *log.Logger {
	return e.logger
}

func (e *Engine) runMetricsEmission(ctx context.Context) {
	ticker := time.NewTicker(MetricsRate)
	for {
		select {
		case <-ticker.C:
			e.emitMetrics()
		case <-ctx.Done():
			return
		}
	}
}

func (e *Engine) emitMetrics() {
	t := time.Now()
	e.processSamples(
		stats.Sample{
			Time:   t,
			Metric: metrics.VUs,
			Value:  float64(e.Executor.GetVUs()),
		},
		stats.Sample{
			Time:   t,
			Metric: metrics.VUsMax,
			Value:  float64(e.Executor.GetVUsMax()),
		},
	)
}

func (e *Engine) runThresholds(ctx context.Context, abort func()) {
	ticker := time.NewTicker(ThresholdsRate)
	for {
		select {
		case <-ticker.C:
			e.processThresholds(abort)
		case <-ctx.Done():
			return
		}
	}
}

func (e *Engine) processThresholds(abort func()) {
	e.MetricsLock.Lock()
	defer e.MetricsLock.Unlock()

	t := e.Executor.GetTime()
	abortOnFail := false

	e.thresholdsTainted = false
	for _, m := range e.Metrics {
		if len(m.Thresholds.Thresholds) == 0 {
			continue
		}
		m.Tainted = null.BoolFrom(false)

		e.logger.WithField("m", m.Name).Debug("running thresholds")
		succ, err := m.Thresholds.Run(m.Sink, t)
		if err != nil {
			e.logger.WithField("m", m.Name).WithError(err).Error("Threshold error")
			continue
		}
		if !succ {
			e.logger.WithField("m", m.Name).Debug("Thresholds failed")
			m.Tainted = null.BoolFrom(true)
			e.thresholdsTainted = true
			if !abortOnFail && m.Thresholds.Abort {
				abortOnFail = true
			}
		}
	}

	if abortOnFail && abort != nil {
		abort()
	}
}

func (e *Engine) processSamples(samples ...stats.Sample) {
	if len(samples) == 0 {
		return
	}

	e.MetricsLock.Lock()
	defer e.MetricsLock.Unlock()

	for i, sample := range samples {
		m, ok := e.Metrics[sample.Metric.Name]
		if !ok {
			m = sample.Metric
			m.Thresholds = e.thresholds[m.Name]
			m.Submetrics = e.submetrics[m.Name]
			e.Metrics[m.Name] = m
		}
		m.Sink.Add(sample)

		for _, sm := range m.Submetrics {
			passing := true
			for k, v := range sm.Tags {
				if sample.Tags[k] != v {
					passing = false
					break
				}
			}
			if !passing {
				continue
			}

			if sm.Metric == nil {
				sm.Metric = stats.New(sm.Name, sample.Metric.Type, sample.Metric.Contains)
				sm.Metric.Sub = *sm
				sm.Metric.Thresholds = e.thresholds[sm.Name]
				e.Metrics[sm.Name] = sm.Metric
			}
			sm.Metric.Sink.Add(sample)
		}

		if e.Options.RunTags == nil {
			continue
		}

		if samples[i].Tags == nil {
			samples[i].Tags = e.Options.RunTags
			continue
		}

		for k, v := range e.Options.RunTags {
			//only set tags that haven't been already set on the sample
			if _, ok := samples[i].Tags[k]; !ok {
				samples[i].Tags[k] = v
			}
		}
	}

	if e.Collector != nil {
		e.Collector.Collect(samples)
	}
}
<|endoftext|>


<commit_before>package sharedaction

import (
	"archive/zip"
	"crypto/sha1"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"code.cloudfoundry.org/cli/actor/actionerror"
	"code.cloudfoundry.org/ykk"
	ignore "github.com/sabhiram/go-gitignore"
	log "github.com/sirupsen/logrus"
)

const (
	DefaultFolderPermissions      = 0755
	DefaultArchiveFilePermissions = 0744
	MaxResourceMatchChunkSize     = 1000
)

var DefaultIgnoreLines = []string{
	".cfignore",
	".DS_Store",
	".git",
	".gitignore",
	".hg",
	".svn",
	"_darcs",
	"manifest.yaml",
	"manifest.yml",
}

type Resource struct {
	Filename string      `json:"fn"`
	Mode     os.FileMode `json:"mode"`
	SHA1     string      `json:"sha1"`
	Size     int64       `json:"size"`
}

// GatherArchiveResources returns a list of resources for an archive.
func (actor Actor) GatherArchiveResources(archivePath string) ([]Resource, error) {
	var resources []Resource

	archive, err := os.Open(archivePath)
	if err != nil {
		return nil, err
	}
	defer archive.Close()

	reader, err := actor.newArchiveReader(archive)
	if err != nil {
		return nil, err
	}

	gitIgnore, err := actor.generateArchiveCFIgnoreMatcher(reader.File)
	if err != nil {
		log.Errorln("reading .cfignore file:", err)
		return nil, err
	}

	for _, archivedFile := range reader.File {
		filename := filepath.ToSlash(archivedFile.Name)
		if gitIgnore.MatchesPath(filename) {
			continue
		}

		resource := Resource{Filename: filename}
		info := archivedFile.FileInfo()

		switch {
		case info.IsDir():
			resource.Mode = DefaultFolderPermissions
		case info.Mode()&os.ModeSymlink == os.ModeSymlink:
			resource.Mode = info.Mode()
		default:
			fileReader, err := archivedFile.Open()
			if err != nil {
				return nil, err
			}
			defer fileReader.Close()

			hash := sha1.New()

			_, err = io.Copy(hash, fileReader)
			if err != nil {
				return nil, err
			}

			resource.Mode = DefaultArchiveFilePermissions
			resource.SHA1 = fmt.Sprintf("%x", hash.Sum(nil))
			resource.Size = archivedFile.FileInfo().Size()
		}

		resources = append(resources, resource)
	}
	return resources, nil
}

// GatherDirectoryResources returns a list of resources for a directory.
func (actor Actor) GatherDirectoryResources(sourceDir string) ([]Resource, error) {
	var (
		resources []Resource
		gitIgnore *ignore.GitIgnore
	)

	gitIgnore, err := actor.generateDirectoryCFIgnoreMatcher(sourceDir)
	if err != nil {
		log.Errorln("reading .cfignore file:", err)
		return nil, err
	}

	evalDir, err := filepath.EvalSymlinks(sourceDir)
	if err != nil {
		log.Errorln("evaluating symlink:", err)
		return nil, err
	}

	walkErr := filepath.Walk(evalDir, func(fullPath string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		relPath, err := filepath.Rel(evalDir, fullPath)
		if err != nil {
			return err
		}

		// if file ignored contine to the next file
		if gitIgnore.MatchesPath(relPath) {
			return nil
		}

		if relPath == "." {
			return nil
		}

		resource := Resource{
			Filename: filepath.ToSlash(relPath),
		}

		switch {
		case info.IsDir():
			// If the file is a directory
			resource.Mode = DefaultFolderPermissions
		case info.Mode()&os.ModeSymlink == os.ModeSymlink:
			// If the file is a Symlink we just set the mode of the file
			// We won't be using any sha information since we don't do
			// any resource matching on symlinks.
			resource.Mode = fixMode(info.Mode())
		default:
			// If the file is regular we want to open
			// and calculate the sha of the file
			file, err := os.Open(fullPath)
			if err != nil {
				return err
			}
			defer file.Close()

			sum := sha1.New()
			_, err = io.Copy(sum, file)
			if err != nil {
				return err
			}

			resource.Mode = fixMode(info.Mode())
			resource.SHA1 = fmt.Sprintf("%x", sum.Sum(nil))
			resource.Size = info.Size()
		}

		resources = append(resources, resource)
		return nil
	})

	if len(resources) == 0 {
		return nil, actionerror.EmptyDirectoryError{Path: sourceDir}
	}

	return resources, walkErr
}

// ZipArchiveResources zips an archive and a sorted (based on full
// path/filename) list of resources and returns the location. On Windows, the
// filemode for user is forced to be readable and executable.
func (actor Actor) ZipArchiveResources(sourceArchivePath string, filesToInclude []Resource) (string, error) {
	log.WithField("sourceArchive", sourceArchivePath).Info("zipping source files from archive")
	zipFile, err := ioutil.TempFile("", "cf-cli-")
	if err != nil {
		return "", err
	}
	defer zipFile.Close()
	zipPath := zipFile.Name()

	writer := zip.NewWriter(zipFile)
	defer writer.Close()

	source, err := os.Open(sourceArchivePath)
	if err != nil {
		return zipPath, err
	}
	defer source.Close()

	reader, err := actor.newArchiveReader(source)
	if err != nil {
		return zipPath, err
	}

	for _, archiveFile := range reader.File {
		resource, ok := actor.findInResources(archiveFile.Name, filesToInclude)
		if !ok {
			log.WithField("archiveFileName", archiveFile.Name).Debug("skipping file")
			continue
		}

		log.WithField("archiveFileName", archiveFile.Name).Debug("zipping file")
		// archiveFile.Open opens the symlink file, not the file it points too
		reader, openErr := archiveFile.Open()
		if openErr != nil {
			log.WithField("archiveFile", archiveFile.Name).Errorln("opening path in dir:", openErr)
			return zipPath, openErr
		}
		defer reader.Close()

		err = actor.addFileToZipFromFileSystem(
			resource.Filename, reader, archiveFile.FileInfo(),
			resource, writer,
		)
		if err != nil {
			log.WithField("archiveFileName", archiveFile.Name).Errorln("zipping file:", err)
			return zipPath, err
		}
		reader.Close()
	}

	log.WithFields(log.Fields{
		"zip_file_location": zipFile.Name(),
		"zipped_file_count": len(filesToInclude),
	}).Info("zip file created")
	return zipPath, nil
}

// ZipDirectoryResources zips a directory and a sorted (based on full
// path/filename) list of resources and returns the location. On Windows, the
// filemode for user is forced to be readable and executable.
func (actor Actor) ZipDirectoryResources(sourceDir string, filesToInclude []Resource) (string, error) {
	log.WithField("sourceDir", sourceDir).Info("zipping source files from directory")
	zipFile, err := ioutil.TempFile("", "cf-cli-")
	if err != nil {
		return "", err
	}
	defer zipFile.Close()
	zipPath := zipFile.Name()

	writer := zip.NewWriter(zipFile)
	defer writer.Close()

	for _, resource := range filesToInclude {
		fullPath := filepath.Join(sourceDir, resource.Filename)
		log.WithField("fullPath", fullPath).Debug("zipping file")

		fileInfo, err := os.Lstat(fullPath)
		if err != nil {
			log.WithField("fullPath", fullPath).Errorln("stat error in dir:", err)
			return zipPath, err
		}

		log.WithField("file-mode", fileInfo.Mode().String()).Debug("resource file info")
		if fileInfo.Mode()&os.ModeSymlink == os.ModeSymlink {
			// we need to user os.Readlink to read a symlink file from a directory
			err = actor.addLinkToZipFromFileSystem(fullPath, fileInfo, resource, writer)
			if err != nil {
				log.WithField("fullPath", fullPath).Errorln("zipping file:", err)
				return zipPath, err
			}
		} else {
			srcFile, err := os.Open(fullPath)
			defer srcFile.Close()
			if err != nil {
				log.WithField("fullPath", fullPath).Errorln("opening path in dir:", err)
				return zipPath, err
			}

			err = actor.addFileToZipFromFileSystem(
				fullPath, srcFile, fileInfo,
				resource, writer,
			)
			srcFile.Close()
			if err != nil {
				log.WithField("fullPath", fullPath).Errorln("zipping file:", err)
				return zipPath, err
			}
		}
	}

	log.WithFields(log.Fields{
		"zip_file_location": zipFile.Name(),
		"zipped_file_count": len(filesToInclude),
	}).Info("zip file created")
	return zipPath, nil
}

func (Actor) addLinkToZipFromFileSystem(srcPath string,
	fileInfo os.FileInfo, resource Resource,
	zipFile *zip.Writer,
) error {
	header, err := zip.FileInfoHeader(fileInfo)
	if err != nil {
		log.WithField("srcPath", srcPath).Errorln("getting file info in dir:", err)
		return err
	}

	header.Name = resource.Filename
	header.Method = zip.Deflate

	log.WithFields(log.Fields{
		"srcPath":  srcPath,
		"destPath": header.Name,
		"mode":     header.Mode().String(),
	}).Debug("setting mode for file")

	destFileWriter, err := zipFile.CreateHeader(header)
	if err != nil {
		log.Errorln("creating header:", err)
		return err
	}

	pathInSymlink, err := os.Readlink(srcPath)
	if err != nil {
		return err
	}
	log.WithField("path", pathInSymlink).Debug("resolving symlink")
	symLinkContents := strings.NewReader(pathInSymlink)
	if _, err := io.Copy(destFileWriter, symLinkContents); err != nil {
		log.WithField("srcPath", srcPath).Errorln("copying data in dir:", err)
		return err
	}

	return nil
}

func (Actor) addFileToZipFromFileSystem(srcPath string,
	srcFile io.Reader, fileInfo os.FileInfo, resource Resource,
	zipFile *zip.Writer,
) error {
	header, err := zip.FileInfoHeader(fileInfo)
	if err != nil {
		log.WithField("srcPath", srcPath).Errorln("getting file info in dir:", err)
		return err
	}

	header.Name = resource.Filename

	// An extra '/' indicates that this file is a directory
	if fileInfo.IsDir() && !strings.HasSuffix(resource.Filename, "/") {
		header.Name += "/"
	}
	header.Method = zip.Deflate
	header.SetMode(resource.Mode)

	log.WithFields(log.Fields{
		"srcPath":  srcPath,
		"destPath": header.Name,
		"mode":     header.Mode().String(),
	}).Debug("setting mode for file")

	destFileWriter, err := zipFile.CreateHeader(header)
	if err != nil {
		log.Errorln("creating header:", err)
		return err
	}

	if fileInfo.Mode().IsRegular() {
		sum := sha1.New()
		multi := io.MultiWriter(sum, destFileWriter)

		if _, err := io.Copy(multi, srcFile); err != nil {
			log.WithField("srcPath", srcPath).Errorln("copying data in dir:", err)
			return err
		}

		if currentSum := fmt.Sprintf("%x", sum.Sum(nil)); resource.SHA1 != currentSum {
			log.WithFields(log.Fields{
				"expected":   resource.SHA1,
				"currentSum": currentSum,
			}).Error("setting mode for file")
			return actionerror.FileChangedError{Filename: srcPath}
		}
	} else if fileInfo.Mode()&os.ModeSymlink == os.ModeSymlink {
		io.Copy(destFileWriter, srcFile)
	}

	return nil
}

func (Actor) generateArchiveCFIgnoreMatcher(files []*zip.File) (*ignore.GitIgnore, error) {
	for _, item := range files {
		if strings.HasSuffix(item.Name, ".cfignore") {
			fileReader, err := item.Open()
			if err != nil {
				return nil, err
			}
			defer fileReader.Close()

			raw, err := ioutil.ReadAll(fileReader)
			if err != nil {
				return nil, err
			}
			s := append(DefaultIgnoreLines, strings.Split(string(raw), "\n")...)
			return ignore.CompileIgnoreLines(s...)
		}
	}
	return ignore.CompileIgnoreLines(DefaultIgnoreLines...)
}

func (actor Actor) generateDirectoryCFIgnoreMatcher(sourceDir string) (*ignore.GitIgnore, error) {
	pathToCFIgnore := filepath.Join(sourceDir, ".cfignore")

	additionalIgnoreLines := DefaultIgnoreLines

	// If verbose logging has files in the current dir, ignore them
	_, traceFiles := actor.Config.Verbose()
	for _, traceFilePath := range traceFiles {
		if relPath, err := filepath.Rel(sourceDir, traceFilePath); err == nil {
			additionalIgnoreLines = append(additionalIgnoreLines, relPath)
		}
	}

	if _, err := os.Stat(pathToCFIgnore); !os.IsNotExist(err) {
		return ignore.CompileIgnoreFileAndLines(pathToCFIgnore, additionalIgnoreLines...)
	} else {
		return ignore.CompileIgnoreLines(additionalIgnoreLines...)
	}
}

func (Actor) findInResources(path string, filesToInclude []Resource) (Resource, bool) {
	for _, resource := range filesToInclude {
		if resource.Filename == filepath.ToSlash(path) {
			log.WithField("resource", resource.Filename).Debug("found resource in files to include")
			return resource, true
		}
	}

	log.WithField("path", path).Debug("did not find resource in files to include")
	return Resource{}, false
}

func (Actor) newArchiveReader(archive *os.File) (*zip.Reader, error) {
	info, err := archive.Stat()
	if err != nil {
		return nil, err
	}

	return ykk.NewReader(archive, info.Size())
}

func (actor Actor) CreateArchive(bitsPath string, resources []Resource) (io.ReadCloser, int64, error) {
	archivePath, err := actor.ZipDirectoryResources(bitsPath, resources)
	_ = err

	return actor.ReadArchive(archivePath)
}

func (Actor) ReadArchive(archivePath string) (io.ReadCloser, int64, error) {
	archive, err := os.Open(archivePath)
	if err != nil {
		log.WithField("archivePath", archivePath).Errorln("opening temp archive:", err)
		return nil, -1, err
	}

	archiveInfo, err := archive.Stat()
	if err != nil {
		archive.Close()
		log.WithField("archivePath", archivePath).Errorln("stat temp archive:", err)
		return nil, -1, err
	}

	return archive, archiveInfo.Size(), nil
}
<commit_msg>add additional debug logging for CF ignore on push<commit_after>package sharedaction

import (
	"archive/zip"
	"crypto/sha1"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"code.cloudfoundry.org/cli/actor/actionerror"
	"code.cloudfoundry.org/ykk"
	ignore "github.com/sabhiram/go-gitignore"
	log "github.com/sirupsen/logrus"
)

const (
	DefaultFolderPermissions      = 0755
	DefaultArchiveFilePermissions = 0744
	MaxResourceMatchChunkSize     = 1000
)

var DefaultIgnoreLines = []string{
	".cfignore",
	".DS_Store",
	".git",
	".gitignore",
	".hg",
	".svn",
	"_darcs",
	"manifest.yaml",
	"manifest.yml",
}

type Resource struct {
	Filename string      `json:"fn"`
	Mode     os.FileMode `json:"mode"`
	SHA1     string      `json:"sha1"`
	Size     int64       `json:"size"`
}

// GatherArchiveResources returns a list of resources for an archive.
func (actor Actor) GatherArchiveResources(archivePath string) ([]Resource, error) {
	var resources []Resource

	archive, err := os.Open(archivePath)
	if err != nil {
		return nil, err
	}
	defer archive.Close()

	reader, err := actor.newArchiveReader(archive)
	if err != nil {
		return nil, err
	}

	gitIgnore, err := actor.generateArchiveCFIgnoreMatcher(reader.File)
	if err != nil {
		log.Errorln("reading .cfignore file:", err)
		return nil, err
	}

	for _, archivedFile := range reader.File {
		filename := filepath.ToSlash(archivedFile.Name)
		if gitIgnore.MatchesPath(filename) {
			continue
		}

		resource := Resource{Filename: filename}
		info := archivedFile.FileInfo()

		switch {
		case info.IsDir():
			resource.Mode = DefaultFolderPermissions
		case info.Mode()&os.ModeSymlink == os.ModeSymlink:
			resource.Mode = info.Mode()
		default:
			fileReader, err := archivedFile.Open()
			if err != nil {
				return nil, err
			}
			defer fileReader.Close()

			hash := sha1.New()

			_, err = io.Copy(hash, fileReader)
			if err != nil {
				return nil, err
			}

			resource.Mode = DefaultArchiveFilePermissions
			resource.SHA1 = fmt.Sprintf("%x", hash.Sum(nil))
			resource.Size = archivedFile.FileInfo().Size()
		}

		resources = append(resources, resource)
	}
	return resources, nil
}

// GatherDirectoryResources returns a list of resources for a directory.
func (actor Actor) GatherDirectoryResources(sourceDir string) ([]Resource, error) {
	var (
		resources []Resource
		gitIgnore *ignore.GitIgnore
	)

	gitIgnore, err := actor.generateDirectoryCFIgnoreMatcher(sourceDir)
	if err != nil {
		log.Errorln("reading .cfignore file:", err)
		return nil, err
	}

	evalDir, err := filepath.EvalSymlinks(sourceDir)
	if err != nil {
		log.Errorln("evaluating symlink:", err)
		return nil, err
	}

	walkErr := filepath.Walk(evalDir, func(fullPath string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		relPath, err := filepath.Rel(evalDir, fullPath)
		if err != nil {
			return err
		}

		// if file ignored contine to the next file
		if gitIgnore.MatchesPath(relPath) {
			return nil
		}

		if relPath == "." {
			return nil
		}

		resource := Resource{
			Filename: filepath.ToSlash(relPath),
		}

		switch {
		case info.IsDir():
			// If the file is a directory
			resource.Mode = DefaultFolderPermissions
		case info.Mode()&os.ModeSymlink == os.ModeSymlink:
			// If the file is a Symlink we just set the mode of the file
			// We won't be using any sha information since we don't do
			// any resource matching on symlinks.
			resource.Mode = fixMode(info.Mode())
		default:
			// If the file is regular we want to open
			// and calculate the sha of the file
			file, err := os.Open(fullPath)
			if err != nil {
				return err
			}
			defer file.Close()

			sum := sha1.New()
			_, err = io.Copy(sum, file)
			if err != nil {
				return err
			}

			resource.Mode = fixMode(info.Mode())
			resource.SHA1 = fmt.Sprintf("%x", sum.Sum(nil))
			resource.Size = info.Size()
		}

		resources = append(resources, resource)
		return nil
	})

	if len(resources) == 0 {
		return nil, actionerror.EmptyDirectoryError{Path: sourceDir}
	}

	return resources, walkErr
}

// ZipArchiveResources zips an archive and a sorted (based on full
// path/filename) list of resources and returns the location. On Windows, the
// filemode for user is forced to be readable and executable.
func (actor Actor) ZipArchiveResources(sourceArchivePath string, filesToInclude []Resource) (string, error) {
	log.WithField("sourceArchive", sourceArchivePath).Info("zipping source files from archive")
	zipFile, err := ioutil.TempFile("", "cf-cli-")
	if err != nil {
		return "", err
	}
	defer zipFile.Close()
	zipPath := zipFile.Name()

	writer := zip.NewWriter(zipFile)
	defer writer.Close()

	source, err := os.Open(sourceArchivePath)
	if err != nil {
		return zipPath, err
	}
	defer source.Close()

	reader, err := actor.newArchiveReader(source)
	if err != nil {
		return zipPath, err
	}

	for _, archiveFile := range reader.File {
		resource, ok := actor.findInResources(archiveFile.Name, filesToInclude)
		if !ok {
			log.WithField("archiveFileName", archiveFile.Name).Debug("skipping file")
			continue
		}

		log.WithField("archiveFileName", archiveFile.Name).Debug("zipping file")
		// archiveFile.Open opens the symlink file, not the file it points too
		reader, openErr := archiveFile.Open()
		if openErr != nil {
			log.WithField("archiveFile", archiveFile.Name).Errorln("opening path in dir:", openErr)
			return zipPath, openErr
		}
		defer reader.Close()

		err = actor.addFileToZipFromFileSystem(
			resource.Filename, reader, archiveFile.FileInfo(),
			resource, writer,
		)
		if err != nil {
			log.WithField("archiveFileName", archiveFile.Name).Errorln("zipping file:", err)
			return zipPath, err
		}
		reader.Close()
	}

	log.WithFields(log.Fields{
		"zip_file_location": zipFile.Name(),
		"zipped_file_count": len(filesToInclude),
	}).Info("zip file created")
	return zipPath, nil
}

// ZipDirectoryResources zips a directory and a sorted (based on full
// path/filename) list of resources and returns the location. On Windows, the
// filemode for user is forced to be readable and executable.
func (actor Actor) ZipDirectoryResources(sourceDir string, filesToInclude []Resource) (string, error) {
	log.WithField("sourceDir", sourceDir).Info("zipping source files from directory")
	zipFile, err := ioutil.TempFile("", "cf-cli-")
	if err != nil {
		return "", err
	}
	defer zipFile.Close()
	zipPath := zipFile.Name()

	writer := zip.NewWriter(zipFile)
	defer writer.Close()

	for _, resource := range filesToInclude {
		fullPath := filepath.Join(sourceDir, resource.Filename)
		log.WithField("fullPath", fullPath).Debug("zipping file")

		fileInfo, err := os.Lstat(fullPath)
		if err != nil {
			log.WithField("fullPath", fullPath).Errorln("stat error in dir:", err)
			return zipPath, err
		}

		log.WithField("file-mode", fileInfo.Mode().String()).Debug("resource file info")
		if fileInfo.Mode()&os.ModeSymlink == os.ModeSymlink {
			// we need to user os.Readlink to read a symlink file from a directory
			err = actor.addLinkToZipFromFileSystem(fullPath, fileInfo, resource, writer)
			if err != nil {
				log.WithField("fullPath", fullPath).Errorln("zipping file:", err)
				return zipPath, err
			}
		} else {
			srcFile, err := os.Open(fullPath)
			defer srcFile.Close()
			if err != nil {
				log.WithField("fullPath", fullPath).Errorln("opening path in dir:", err)
				return zipPath, err
			}

			err = actor.addFileToZipFromFileSystem(
				fullPath, srcFile, fileInfo,
				resource, writer,
			)
			srcFile.Close()
			if err != nil {
				log.WithField("fullPath", fullPath).Errorln("zipping file:", err)
				return zipPath, err
			}
		}
	}

	log.WithFields(log.Fields{
		"zip_file_location": zipFile.Name(),
		"zipped_file_count": len(filesToInclude),
	}).Info("zip file created")
	return zipPath, nil
}

func (Actor) addLinkToZipFromFileSystem(srcPath string,
	fileInfo os.FileInfo, resource Resource,
	zipFile *zip.Writer,
) error {
	header, err := zip.FileInfoHeader(fileInfo)
	if err != nil {
		log.WithField("srcPath", srcPath).Errorln("getting file info in dir:", err)
		return err
	}

	header.Name = resource.Filename
	header.Method = zip.Deflate

	log.WithFields(log.Fields{
		"srcPath":  srcPath,
		"destPath": header.Name,
		"mode":     header.Mode().String(),
	}).Debug("setting mode for file")

	destFileWriter, err := zipFile.CreateHeader(header)
	if err != nil {
		log.Errorln("creating header:", err)
		return err
	}

	pathInSymlink, err := os.Readlink(srcPath)
	if err != nil {
		return err
	}
	log.WithField("path", pathInSymlink).Debug("resolving symlink")
	symLinkContents := strings.NewReader(pathInSymlink)
	if _, err := io.Copy(destFileWriter, symLinkContents); err != nil {
		log.WithField("srcPath", srcPath).Errorln("copying data in dir:", err)
		return err
	}

	return nil
}

func (Actor) addFileToZipFromFileSystem(srcPath string,
	srcFile io.Reader, fileInfo os.FileInfo, resource Resource,
	zipFile *zip.Writer,
) error {
	header, err := zip.FileInfoHeader(fileInfo)
	if err != nil {
		log.WithField("srcPath", srcPath).Errorln("getting file info in dir:", err)
		return err
	}

	header.Name = resource.Filename

	// An extra '/' indicates that this file is a directory
	if fileInfo.IsDir() && !strings.HasSuffix(resource.Filename, "/") {
		header.Name += "/"
	}
	header.Method = zip.Deflate
	header.SetMode(resource.Mode)

	log.WithFields(log.Fields{
		"srcPath":  srcPath,
		"destPath": header.Name,
		"mode":     header.Mode().String(),
	}).Debug("setting mode for file")

	destFileWriter, err := zipFile.CreateHeader(header)
	if err != nil {
		log.Errorln("creating header:", err)
		return err
	}

	if fileInfo.Mode().IsRegular() {
		sum := sha1.New()
		multi := io.MultiWriter(sum, destFileWriter)

		if _, err := io.Copy(multi, srcFile); err != nil {
			log.WithField("srcPath", srcPath).Errorln("copying data in dir:", err)
			return err
		}

		if currentSum := fmt.Sprintf("%x", sum.Sum(nil)); resource.SHA1 != currentSum {
			log.WithFields(log.Fields{
				"expected":   resource.SHA1,
				"currentSum": currentSum,
			}).Error("setting mode for file")
			return actionerror.FileChangedError{Filename: srcPath}
		}
	} else if fileInfo.Mode()&os.ModeSymlink == os.ModeSymlink {
		io.Copy(destFileWriter, srcFile)
	}

	return nil
}

func (Actor) generateArchiveCFIgnoreMatcher(files []*zip.File) (*ignore.GitIgnore, error) {
	for _, item := range files {
		if strings.HasSuffix(item.Name, ".cfignore") {
			fileReader, err := item.Open()
			if err != nil {
				return nil, err
			}
			defer fileReader.Close()

			raw, err := ioutil.ReadAll(fileReader)
			if err != nil {
				return nil, err
			}
			s := append(DefaultIgnoreLines, strings.Split(string(raw), "\n")...)
			return ignore.CompileIgnoreLines(s...)
		}
	}
	return ignore.CompileIgnoreLines(DefaultIgnoreLines...)
}

func (actor Actor) generateDirectoryCFIgnoreMatcher(sourceDir string) (*ignore.GitIgnore, error) {
	pathToCFIgnore := filepath.Join(sourceDir, ".cfignore")
	log.WithField("pathToCFIgnore", pathToCFIgnore).Debug("using ignore file")

	additionalIgnoreLines := DefaultIgnoreLines

	// If verbose logging has files in the current dir, ignore them
	_, traceFiles := actor.Config.Verbose()
	for _, traceFilePath := range traceFiles {
		if relPath, err := filepath.Rel(sourceDir, traceFilePath); err == nil {
			additionalIgnoreLines = append(additionalIgnoreLines, relPath)
		}
	}

	if _, err := os.Stat(pathToCFIgnore); !os.IsNotExist(err) {
		return ignore.CompileIgnoreFileAndLines(pathToCFIgnore, additionalIgnoreLines...)
	} else {
		return ignore.CompileIgnoreLines(additionalIgnoreLines...)
	}
}

func (Actor) findInResources(path string, filesToInclude []Resource) (Resource, bool) {
	for _, resource := range filesToInclude {
		if resource.Filename == filepath.ToSlash(path) {
			log.WithField("resource", resource.Filename).Debug("found resource in files to include")
			return resource, true
		}
	}

	log.WithField("path", path).Debug("did not find resource in files to include")
	return Resource{}, false
}

func (Actor) newArchiveReader(archive *os.File) (*zip.Reader, error) {
	info, err := archive.Stat()
	if err != nil {
		return nil, err
	}

	return ykk.NewReader(archive, info.Size())
}

func (actor Actor) CreateArchive(bitsPath string, resources []Resource) (io.ReadCloser, int64, error) {
	archivePath, err := actor.ZipDirectoryResources(bitsPath, resources)
	_ = err

	return actor.ReadArchive(archivePath)
}

func (Actor) ReadArchive(archivePath string) (io.ReadCloser, int64, error) {
	archive, err := os.Open(archivePath)
	if err != nil {
		log.WithField("archivePath", archivePath).Errorln("opening temp archive:", err)
		return nil, -1, err
	}

	archiveInfo, err := archive.Stat()
	if err != nil {
		archive.Close()
		log.WithField("archivePath", archivePath).Errorln("stat temp archive:", err)
		return nil, -1, err
	}

	return archive, archiveInfo.Size(), nil
}
<|endoftext|>


<commit_before>package dependency

import (
	"encoding/gob"
	"fmt"
	"log"
	"net/url"
	"regexp"
	"sort"

	"github.com/pkg/errors"
)

var (
	// Ensure implements
	_ Dependency = (*CatalogNodeQuery)(nil)

	// CatalogNodeQueryRe is the regular expression to use.
	CatalogNodeQueryRe = regexp.MustCompile(`\A` + nameRe + dcRe + `\z`)
)

func init() {
	gob.Register([]*CatalogNode{})
	gob.Register([]*CatalogNodeService{})
}

// CatalogNodeQuery represents a single node from the Consul catalog.
type CatalogNodeQuery struct {
	stopCh chan struct{}

	dc   string
	name string
}

// CatalogNode is a wrapper around the node and its services.
type CatalogNode struct {
	Node     *Node
	Services []*CatalogNodeService
}

// CatalogNodeService is a service on a single node.
type CatalogNodeService struct {
	ID                string
	Service           string
	Tags              ServiceTags
	Port              int
	Address           string
	EnableTagOverride bool
}

// NewCatalogNodeQuery parses the given string into a dependency. If the name is
// empty then the name of the local agent is used.
func NewCatalogNodeQuery(s string) (*CatalogNodeQuery, error) {
	if s != "" && !CatalogNodeQueryRe.MatchString(s) {
		return nil, fmt.Errorf("catalog.node: invalid format: %q", s)
	}

	m := regexpMatch(CatalogNodeQueryRe, s)
	return &CatalogNodeQuery{
		dc:     m["dc"],
		name:   m["name"],
		stopCh: make(chan struct{}, 1),
	}, nil
}

// Fetch queries the Consul API defined by the given client and returns a
// of CatalogNode object.
func (d *CatalogNodeQuery) Fetch(clients *ClientSet, opts *QueryOptions) (interface{}, *ResponseMetadata, error) {
	select {
	case <-d.stopCh:
		return nil, nil, ErrStopped
	default:
	}

	opts = opts.Merge(&QueryOptions{
		Datacenter: d.dc,
	})

	if d.name == "" {
		log.Printf("[TRACE] %s: getting local agent name", d)
		name, err := clients.Consul().Agent().NodeName()
		if err != nil {
			return nil, nil, errors.Wrapf(err, d.String())
		}
		d.name = name
	}

	log.Printf("[TRACE] %s: GET %s", d, &url.URL{
		Path:     "/v1/catalog/node/" + d.name,
		RawQuery: opts.String(),
	})
	node, qm, err := clients.Consul().Catalog().Node(d.name, opts.ToConsulOpts())
	if err != nil {
		return nil, nil, errors.Wrap(err, d.String())
	}

	log.Printf("[TRACE] %s: returned response", d)

	rm := &ResponseMetadata{
		LastIndex:   qm.LastIndex,
		LastContact: qm.LastContact,
	}

	if node == nil {
		log.Printf("[WARN] %s: no node exists with the name %q", d, d.name)
		var node CatalogNode
		return &node, rm, nil
	}

	services := make([]*CatalogNodeService, 0, len(node.Services))
	for _, v := range node.Services {
		services = append(services, &CatalogNodeService{
			ID:                v.ID,
			Service:           v.Service,
			Tags:              ServiceTags(deepCopyAndSortTags(v.Tags)),
			Port:              v.Port,
			Address:           v.Address,
			EnableTagOverride: v.EnableTagOverride,
		})
	}
	sort.Stable(ByService(services))

	detail := &CatalogNode{
		Node: &Node{
			Node:            node.Node.Node,
			Address:         node.Node.Address,
			TaggedAddresses: node.Node.TaggedAddresses,
		},
		Services: services,
	}

	return detail, rm, nil
}

// CanShare returns a boolean if this dependency is shareable.
func (d *CatalogNodeQuery) CanShare() bool {
	return false
}

// String returns the human-friendly version of this dependency.
func (d *CatalogNodeQuery) String() string {
	name := d.name
	if d.dc != "" {
		name = name + "@" + d.dc
	}

	if name == "" {
		return "catalog.node"
	}
	return fmt.Sprintf("catalog.node(%s)", name)
}

// Stop halts the dependency's fetch function.
func (d *CatalogNodeQuery) Stop() {
	close(d.stopCh)
}

// Type returns the type of this dependency.
func (d *CatalogNodeQuery) Type() Type {
	return TypeConsul
}

// ByService is a sorter of node services by their service name and then ID.
type ByService []*CatalogNodeService

func (s ByService) Len() int      { return len(s) }
func (s ByService) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s ByService) Less(i, j int) bool {
	if s[i].Service == s[j].Service {
		return s[i].ID <= s[j].ID
	}
	return s[i].Service <= s[j].Service
}
<commit_msg>Do not alter dependency after it's created<commit_after>package dependency

import (
	"encoding/gob"
	"fmt"
	"log"
	"net/url"
	"regexp"
	"sort"

	"github.com/pkg/errors"
)

var (
	// Ensure implements
	_ Dependency = (*CatalogNodeQuery)(nil)

	// CatalogNodeQueryRe is the regular expression to use.
	CatalogNodeQueryRe = regexp.MustCompile(`\A` + nameRe + dcRe + `\z`)
)

func init() {
	gob.Register([]*CatalogNode{})
	gob.Register([]*CatalogNodeService{})
}

// CatalogNodeQuery represents a single node from the Consul catalog.
type CatalogNodeQuery struct {
	stopCh chan struct{}

	dc   string
	name string
}

// CatalogNode is a wrapper around the node and its services.
type CatalogNode struct {
	Node     *Node
	Services []*CatalogNodeService
}

// CatalogNodeService is a service on a single node.
type CatalogNodeService struct {
	ID                string
	Service           string
	Tags              ServiceTags
	Port              int
	Address           string
	EnableTagOverride bool
}

// NewCatalogNodeQuery parses the given string into a dependency. If the name is
// empty then the name of the local agent is used.
func NewCatalogNodeQuery(s string) (*CatalogNodeQuery, error) {
	if s != "" && !CatalogNodeQueryRe.MatchString(s) {
		return nil, fmt.Errorf("catalog.node: invalid format: %q", s)
	}

	m := regexpMatch(CatalogNodeQueryRe, s)
	return &CatalogNodeQuery{
		dc:     m["dc"],
		name:   m["name"],
		stopCh: make(chan struct{}, 1),
	}, nil
}

// Fetch queries the Consul API defined by the given client and returns a
// of CatalogNode object.
func (d *CatalogNodeQuery) Fetch(clients *ClientSet, opts *QueryOptions) (interface{}, *ResponseMetadata, error) {
	select {
	case <-d.stopCh:
		return nil, nil, ErrStopped
	default:
	}

	opts = opts.Merge(&QueryOptions{
		Datacenter: d.dc,
	})

	// Grab the name
	name := d.name

	if name == "" {
		log.Printf("[TRACE] %s: getting local agent name", d)
		var err error
		name, err = clients.Consul().Agent().NodeName()
		if err != nil {
			return nil, nil, errors.Wrapf(err, d.String())
		}
	}

	log.Printf("[TRACE] %s: GET %s", d, &url.URL{
		Path:     "/v1/catalog/node/" + name,
		RawQuery: opts.String(),
	})
	node, qm, err := clients.Consul().Catalog().Node(name, opts.ToConsulOpts())
	if err != nil {
		return nil, nil, errors.Wrap(err, d.String())
	}

	log.Printf("[TRACE] %s: returned response", d)

	rm := &ResponseMetadata{
		LastIndex:   qm.LastIndex,
		LastContact: qm.LastContact,
	}

	if node == nil {
		log.Printf("[WARN] %s: no node exists with the name %q", d, name)
		var node CatalogNode
		return &node, rm, nil
	}

	services := make([]*CatalogNodeService, 0, len(node.Services))
	for _, v := range node.Services {
		services = append(services, &CatalogNodeService{
			ID:                v.ID,
			Service:           v.Service,
			Tags:              ServiceTags(deepCopyAndSortTags(v.Tags)),
			Port:              v.Port,
			Address:           v.Address,
			EnableTagOverride: v.EnableTagOverride,
		})
	}
	sort.Stable(ByService(services))

	detail := &CatalogNode{
		Node: &Node{
			Node:            node.Node.Node,
			Address:         node.Node.Address,
			TaggedAddresses: node.Node.TaggedAddresses,
		},
		Services: services,
	}

	return detail, rm, nil
}

// CanShare returns a boolean if this dependency is shareable.
func (d *CatalogNodeQuery) CanShare() bool {
	return false
}

// String returns the human-friendly version of this dependency.
func (d *CatalogNodeQuery) String() string {
	name := d.name
	if d.dc != "" {
		name = name + "@" + d.dc
	}

	if name == "" {
		return "catalog.node"
	}
	return fmt.Sprintf("catalog.node(%s)", name)
}

// Stop halts the dependency's fetch function.
func (d *CatalogNodeQuery) Stop() {
	close(d.stopCh)
}

// Type returns the type of this dependency.
func (d *CatalogNodeQuery) Type() Type {
	return TypeConsul
}

// ByService is a sorter of node services by their service name and then ID.
type ByService []*CatalogNodeService

func (s ByService) Len() int      { return len(s) }
func (s ByService) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s ByService) Less(i, j int) bool {
	if s[i].Service == s[j].Service {
		return s[i].ID <= s[j].ID
	}
	return s[i].Service <= s[j].Service
}
<|endoftext|>