<commit_before>// Copyright 2015 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// This file does serialization of programs for executor binary.
// The format aims at simple parsing: binary and irreversible.

// Exec format is an sequence of uint64's which encodes a sequence of calls.
// The sequence is terminated by a speciall call execInstrEOF.
// Each call is (call ID, copyout index, number of arguments, arguments...).
// Each argument is (type, size, value).
// There are 4 types of arguments:
//  - execArgConst: value is const value
//  - execArgResult: value is copyout index we want to reference
//  - execArgData: value is a binary blob (represented as ]size/8[ uint64's)
//  - execArgCsum: runtime checksum calculation
// There are 2 other special calls:
//  - execInstrCopyin: copies its second argument into address specified by first argument
//  - execInstrCopyout: reads value at address specified by first argument (result can be referenced by execArgResult)

package prog

import (
	"fmt"
	"sort"
)

const (
	execInstrEOF = ^uint64(iota)
	execInstrCopyin
	execInstrCopyout
)

const (
	execArgConst = uint64(iota)
	execArgResult
	execArgData
	execArgCsum

	execArgDataReadable = uint64(1 << 63)
)

const (
	ExecArgCsumInet = uint64(iota)
)

const (
	ExecArgCsumChunkData = uint64(iota)
	ExecArgCsumChunkConst
)

const (
	ExecBufferSize = 4 << 20 // keep in sync with kMaxInput in executor.cc
	ExecNoCopyout  = ^uint64(0)
)

// SerializeForExec serializes program p for execution by process pid into the provided buffer.
// Returns number of bytes written to the buffer.
// If the provided buffer is too small for the program an error is returned.
func (p *Prog) SerializeForExec(buffer []byte) (int, error) {
	p.debugValidate()
	w := &execContext{
		target: p.Target,
		buf:    buffer,
		eof:    false,
		args:   make(map[Arg]argInfo),
	}
	for _, c := range p.Calls {
		w.csumMap, w.csumUses = calcChecksumsCall(c)
		w.serializeCall(c)
	}
	w.write(execInstrEOF)
	if w.eof {
		return 0, fmt.Errorf("provided buffer is too small")
	}
	return len(buffer) - len(w.buf), nil
}

func (w *execContext) serializeCall(c *Call) {
	// Calculate arg offsets within structs.
	// Generate copyin instructions that fill in data into pointer arguments.
	w.writeCopyin(c)
	// Generate checksum calculation instructions starting from the last one,
	// since checksum values can depend on values of the latter ones
	w.writeChecksums()
	// Generate the call itself.
	w.write(uint64(c.Meta.ID))
	if c.Ret != nil && len(c.Ret.uses) != 0 {
		if _, ok := w.args[c.Ret]; ok {
			panic("argInfo is already created for return value")
		}
		w.args[c.Ret] = argInfo{Idx: w.copyoutSeq, Ret: true}
		w.write(w.copyoutSeq)
		w.copyoutSeq++
	} else {
		w.write(ExecNoCopyout)
	}
	w.write(uint64(len(c.Args)))
	for _, arg := range c.Args {
		w.writeArg(arg)
	}
	// Generate copyout instructions that persist interesting return values.
	w.writeCopyout(c)
}

type execContext struct {
	target     *Target
	buf        []byte
	eof        bool
	args       map[Arg]argInfo
	copyoutSeq uint64
	// Per-call state cached here to not pass it through all functions.
	csumMap  map[Arg]CsumInfo
	csumUses map[Arg]struct{}
}

type argInfo struct {
	Addr uint64 // physical addr
	Idx  uint64 // copyout instruction index
	Ret  bool
}

func (w *execContext) writeCopyin(c *Call) {
	ForeachArg(c, func(arg Arg, ctx *ArgCtx) {
		if ctx.Base == nil {
			return
		}
		addr := w.target.PhysicalAddr(ctx.Base) + ctx.Offset
		addr -= arg.Type().UnitOffset()
		if w.willBeUsed(arg) {
			w.args[arg] = argInfo{Addr: addr}
		}
		switch arg.(type) {
		case *GroupArg, *UnionArg:
			return
		}
		typ := arg.Type()
		if arg.Dir() == DirOut || IsPad(typ) || (arg.Size() == 0 && !typ.IsBitfield()) {
			return
		}
		w.write(execInstrCopyin)
		w.write(addr)
		w.writeArg(arg)
	})
}

func (w *execContext) willBeUsed(arg Arg) bool {
	if res, ok := arg.(*ResultArg); ok && len(res.uses) != 0 {
		return true
	}
	_, ok1 := w.csumMap[arg]
	_, ok2 := w.csumUses[arg]
	return ok1 || ok2
}

func (w *execContext) writeChecksums() {
	if len(w.csumMap) == 0 {
		return
	}
	csumArgs := make([]Arg, 0, len(w.csumMap))
	for arg := range w.csumMap {
		csumArgs = append(csumArgs, arg)
	}
	sort.Slice(csumArgs, func(i, j int) bool {
		return w.args[csumArgs[i]].Addr < w.args[csumArgs[j]].Addr
	})
	for i := len(csumArgs) - 1; i >= 0; i-- {
		arg := csumArgs[i]
		info := w.csumMap[arg]
		if _, ok := arg.Type().(*CsumType); !ok {
			panic("csum arg is not csum type")
		}
		w.write(execInstrCopyin)
		w.write(w.args[arg].Addr)
		w.write(execArgCsum)
		w.write(arg.Size())
		switch info.Kind {
		case CsumInet:
			w.write(ExecArgCsumInet)
			w.write(uint64(len(info.Chunks)))
			for _, chunk := range info.Chunks {
				switch chunk.Kind {
				case CsumChunkArg:
					w.write(ExecArgCsumChunkData)
					w.write(w.args[chunk.Arg].Addr)
					w.write(chunk.Arg.Size())
				case CsumChunkConst:
					w.write(ExecArgCsumChunkConst)
					w.write(chunk.Value)
					w.write(chunk.Size)
				default:
					panic(fmt.Sprintf("csum chunk has unknown kind %v", chunk.Kind))
				}
			}
		default:
			panic(fmt.Sprintf("csum arg has unknown kind %v", info.Kind))
		}
	}
}

func (w *execContext) writeCopyout(c *Call) {
	ForeachArg(c, func(arg Arg, _ *ArgCtx) {
		if res, ok := arg.(*ResultArg); ok && len(res.uses) != 0 {
			// Create a separate copyout instruction that has own Idx.
			info := w.args[arg]
			if info.Ret {
				return // Idx is already assigned above.
			}
			info.Idx = w.copyoutSeq
			w.copyoutSeq++
			w.args[arg] = info
			w.write(execInstrCopyout)
			w.write(info.Idx)
			w.write(info.Addr)
			w.write(arg.Size())
		}
	})
}

func (w *execContext) write(v uint64) {
	if len(w.buf) < 8 {
		w.eof = true
		return
	}
	w.buf[0] = byte(v >> 0)
	w.buf[1] = byte(v >> 8)
	w.buf[2] = byte(v >> 16)
	w.buf[3] = byte(v >> 24)
	w.buf[4] = byte(v >> 32)
	w.buf[5] = byte(v >> 40)
	w.buf[6] = byte(v >> 48)
	w.buf[7] = byte(v >> 56)
	w.buf = w.buf[8:]
}

func (w *execContext) writeArg(arg Arg) {
	switch a := arg.(type) {
	case *ConstArg:
		val, pidStride := a.Value()
		typ := a.Type()
		w.writeConstArg(typ.UnitSize(), val, typ.BitfieldOffset(), typ.BitfieldLength(), pidStride, typ.Format())
	case *ResultArg:
		if a.Res == nil {
			w.writeConstArg(a.Size(), a.Val, 0, 0, 0, a.Type().Format())
		} else {
			info, ok := w.args[a.Res]
			if !ok {
				panic("no copyout index")
			}
			w.write(execArgResult)
			meta := a.Size() | uint64(a.Type().Format())<<8
			w.write(meta)
			w.write(info.Idx)
			w.write(a.OpDiv)
			w.write(a.OpAdd)
			w.write(a.Type().(*ResourceType).Default())
		}
	case *PointerArg:
		w.writeConstArg(a.Size(), w.target.PhysicalAddr(a), 0, 0, 0, FormatNative)
	case *DataArg:
		data := a.Data()
		if len(data) == 0 {
			return
		}
		w.write(execArgData)
		flags := uint64(len(data))
		if isReadableDataType(a.Type().(*BufferType)) {
			flags |= execArgDataReadable
		}
		w.write(flags)
		padded := len(data)
		if pad := 8 - len(data)%8; pad != 8 {
			padded += pad
		}
		if len(w.buf) < padded {
			w.eof = true
		} else {
			copy(w.buf, data)
			w.buf = w.buf[padded:]
		}
	case *UnionArg:
		w.writeArg(a.Option)
	default:
		panic("unknown arg type")
	}
}

func (w *execContext) writeConstArg(size, val, bfOffset, bfLength, pidStride uint64, bf BinaryFormat) {
	w.write(execArgConst)
	meta := size | uint64(bf)<<8 | bfOffset<<16 | bfLength<<24 | pidStride<<32
	w.write(meta)
	w.write(val)
}
<commit_msg>prog/encodingexec: pad data args with zero bytes<commit_after>// Copyright 2015 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

// This file does serialization of programs for executor binary.
// The format aims at simple parsing: binary and irreversible.

// Exec format is an sequence of uint64's which encodes a sequence of calls.
// The sequence is terminated by a speciall call execInstrEOF.
// Each call is (call ID, copyout index, number of arguments, arguments...).
// Each argument is (type, size, value).
// There are 4 types of arguments:
//  - execArgConst: value is const value
//  - execArgResult: value is copyout index we want to reference
//  - execArgData: value is a binary blob (represented as ]size/8[ uint64's)
//  - execArgCsum: runtime checksum calculation
// There are 2 other special calls:
//  - execInstrCopyin: copies its second argument into address specified by first argument
//  - execInstrCopyout: reads value at address specified by first argument (result can be referenced by execArgResult)

package prog

import (
	"fmt"
	"sort"
)

const (
	execInstrEOF = ^uint64(iota)
	execInstrCopyin
	execInstrCopyout
)

const (
	execArgConst = uint64(iota)
	execArgResult
	execArgData
	execArgCsum

	execArgDataReadable = uint64(1 << 63)
)

const (
	ExecArgCsumInet = uint64(iota)
)

const (
	ExecArgCsumChunkData = uint64(iota)
	ExecArgCsumChunkConst
)

const (
	ExecBufferSize = 4 << 20 // keep in sync with kMaxInput in executor.cc
	ExecNoCopyout  = ^uint64(0)
)

// SerializeForExec serializes program p for execution by process pid into the provided buffer.
// Returns number of bytes written to the buffer.
// If the provided buffer is too small for the program an error is returned.
func (p *Prog) SerializeForExec(buffer []byte) (int, error) {
	p.debugValidate()
	w := &execContext{
		target: p.Target,
		buf:    buffer,
		eof:    false,
		args:   make(map[Arg]argInfo),
	}
	for _, c := range p.Calls {
		w.csumMap, w.csumUses = calcChecksumsCall(c)
		w.serializeCall(c)
	}
	w.write(execInstrEOF)
	if w.eof {
		return 0, fmt.Errorf("provided buffer is too small")
	}
	return len(buffer) - len(w.buf), nil
}

func (w *execContext) serializeCall(c *Call) {
	// Calculate arg offsets within structs.
	// Generate copyin instructions that fill in data into pointer arguments.
	w.writeCopyin(c)
	// Generate checksum calculation instructions starting from the last one,
	// since checksum values can depend on values of the latter ones
	w.writeChecksums()
	// Generate the call itself.
	w.write(uint64(c.Meta.ID))
	if c.Ret != nil && len(c.Ret.uses) != 0 {
		if _, ok := w.args[c.Ret]; ok {
			panic("argInfo is already created for return value")
		}
		w.args[c.Ret] = argInfo{Idx: w.copyoutSeq, Ret: true}
		w.write(w.copyoutSeq)
		w.copyoutSeq++
	} else {
		w.write(ExecNoCopyout)
	}
	w.write(uint64(len(c.Args)))
	for _, arg := range c.Args {
		w.writeArg(arg)
	}
	// Generate copyout instructions that persist interesting return values.
	w.writeCopyout(c)
}

type execContext struct {
	target     *Target
	buf        []byte
	eof        bool
	args       map[Arg]argInfo
	copyoutSeq uint64
	// Per-call state cached here to not pass it through all functions.
	csumMap  map[Arg]CsumInfo
	csumUses map[Arg]struct{}
}

type argInfo struct {
	Addr uint64 // physical addr
	Idx  uint64 // copyout instruction index
	Ret  bool
}

func (w *execContext) writeCopyin(c *Call) {
	ForeachArg(c, func(arg Arg, ctx *ArgCtx) {
		if ctx.Base == nil {
			return
		}
		addr := w.target.PhysicalAddr(ctx.Base) + ctx.Offset
		addr -= arg.Type().UnitOffset()
		if w.willBeUsed(arg) {
			w.args[arg] = argInfo{Addr: addr}
		}
		switch arg.(type) {
		case *GroupArg, *UnionArg:
			return
		}
		typ := arg.Type()
		if arg.Dir() == DirOut || IsPad(typ) || (arg.Size() == 0 && !typ.IsBitfield()) {
			return
		}
		w.write(execInstrCopyin)
		w.write(addr)
		w.writeArg(arg)
	})
}

func (w *execContext) willBeUsed(arg Arg) bool {
	if res, ok := arg.(*ResultArg); ok && len(res.uses) != 0 {
		return true
	}
	_, ok1 := w.csumMap[arg]
	_, ok2 := w.csumUses[arg]
	return ok1 || ok2
}

func (w *execContext) writeChecksums() {
	if len(w.csumMap) == 0 {
		return
	}
	csumArgs := make([]Arg, 0, len(w.csumMap))
	for arg := range w.csumMap {
		csumArgs = append(csumArgs, arg)
	}
	sort.Slice(csumArgs, func(i, j int) bool {
		return w.args[csumArgs[i]].Addr < w.args[csumArgs[j]].Addr
	})
	for i := len(csumArgs) - 1; i >= 0; i-- {
		arg := csumArgs[i]
		info := w.csumMap[arg]
		if _, ok := arg.Type().(*CsumType); !ok {
			panic("csum arg is not csum type")
		}
		w.write(execInstrCopyin)
		w.write(w.args[arg].Addr)
		w.write(execArgCsum)
		w.write(arg.Size())
		switch info.Kind {
		case CsumInet:
			w.write(ExecArgCsumInet)
			w.write(uint64(len(info.Chunks)))
			for _, chunk := range info.Chunks {
				switch chunk.Kind {
				case CsumChunkArg:
					w.write(ExecArgCsumChunkData)
					w.write(w.args[chunk.Arg].Addr)
					w.write(chunk.Arg.Size())
				case CsumChunkConst:
					w.write(ExecArgCsumChunkConst)
					w.write(chunk.Value)
					w.write(chunk.Size)
				default:
					panic(fmt.Sprintf("csum chunk has unknown kind %v", chunk.Kind))
				}
			}
		default:
			panic(fmt.Sprintf("csum arg has unknown kind %v", info.Kind))
		}
	}
}

func (w *execContext) writeCopyout(c *Call) {
	ForeachArg(c, func(arg Arg, _ *ArgCtx) {
		if res, ok := arg.(*ResultArg); ok && len(res.uses) != 0 {
			// Create a separate copyout instruction that has own Idx.
			info := w.args[arg]
			if info.Ret {
				return // Idx is already assigned above.
			}
			info.Idx = w.copyoutSeq
			w.copyoutSeq++
			w.args[arg] = info
			w.write(execInstrCopyout)
			w.write(info.Idx)
			w.write(info.Addr)
			w.write(arg.Size())
		}
	})
}

func (w *execContext) write(v uint64) {
	if len(w.buf) < 8 {
		w.eof = true
		return
	}
	w.buf[0] = byte(v >> 0)
	w.buf[1] = byte(v >> 8)
	w.buf[2] = byte(v >> 16)
	w.buf[3] = byte(v >> 24)
	w.buf[4] = byte(v >> 32)
	w.buf[5] = byte(v >> 40)
	w.buf[6] = byte(v >> 48)
	w.buf[7] = byte(v >> 56)
	w.buf = w.buf[8:]
}

func (w *execContext) writeArg(arg Arg) {
	switch a := arg.(type) {
	case *ConstArg:
		val, pidStride := a.Value()
		typ := a.Type()
		w.writeConstArg(typ.UnitSize(), val, typ.BitfieldOffset(), typ.BitfieldLength(), pidStride, typ.Format())
	case *ResultArg:
		if a.Res == nil {
			w.writeConstArg(a.Size(), a.Val, 0, 0, 0, a.Type().Format())
		} else {
			info, ok := w.args[a.Res]
			if !ok {
				panic("no copyout index")
			}
			w.write(execArgResult)
			meta := a.Size() | uint64(a.Type().Format())<<8
			w.write(meta)
			w.write(info.Idx)
			w.write(a.OpDiv)
			w.write(a.OpAdd)
			w.write(a.Type().(*ResourceType).Default())
		}
	case *PointerArg:
		w.writeConstArg(a.Size(), w.target.PhysicalAddr(a), 0, 0, 0, FormatNative)
	case *DataArg:
		data := a.Data()
		if len(data) == 0 {
			return
		}
		w.write(execArgData)
		flags := uint64(len(data))
		if isReadableDataType(a.Type().(*BufferType)) {
			flags |= execArgDataReadable
		}
		w.write(flags)
		padded := len(data)
		if pad := 8 - len(data)%8; pad != 8 {
			padded += pad
		}
		if len(w.buf) < padded {
			w.eof = true
		} else {
			copy(w.buf, data)
			copy(w.buf[len(data):], make([]byte, 8))
			w.buf = w.buf[padded:]
		}
	case *UnionArg:
		w.writeArg(a.Option)
	default:
		panic("unknown arg type")
	}
}

func (w *execContext) writeConstArg(size, val, bfOffset, bfLength, pidStride uint64, bf BinaryFormat) {
	w.write(execArgConst)
	meta := size | uint64(bf)<<8 | bfOffset<<16 | bfLength<<24 | pidStride<<32
	w.write(meta)
	w.write(val)
}
<|endoftext|>


<commit_before>// Copyright 2013-2022 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package musig2

import (
	"encoding/hex"
	"fmt"
	"testing"

	"github.com/btcsuite/btcd/btcec/v2"
	"github.com/btcsuite/btcd/btcec/v2/schnorr"
)

var (
	testPrivBytes = hexToModNScalar("9e0699c91ca1e3b7e3c9ba71eb71c89890872be97576010fe593fbf3fd57e66d")

	testMsg = hexToBytes("c301ba9de5d6053caad9f5eb46523f007702add2c62fa39de03146a36b8026b7")
)

func hexToBytes(s string) []byte {
	b, err := hex.DecodeString(s)
	if err != nil {
		panic("invalid hex in source file: " + s)
	}
	return b
}

func hexToModNScalar(s string) *btcec.ModNScalar {
	b, err := hex.DecodeString(s)
	if err != nil {
		panic("invalid hex in source file: " + s)
	}
	var scalar btcec.ModNScalar
	if overflow := scalar.SetByteSlice(b); overflow {
		panic("hex in source file overflows mod N scalar: " + s)
	}
	return &scalar
}

func genSigner(t *testing.B) signer {
	privKey, err := btcec.NewPrivateKey()
	if err != nil {
		t.Fatalf("unable to gen priv key: %v", err)
	}

	pubKey, err := schnorr.ParsePubKey(
		schnorr.SerializePubKey(privKey.PubKey()),
	)
	if err != nil {
		t.Fatalf("unable to gen key: %v", err)
	}

	nonces, err := GenNonces(WithPublicKey(pubKey))
	if err != nil {
		t.Fatalf("unable to gen nonces: %v", err)
	}

	return signer{
		privKey: privKey,
		pubKey:  pubKey,
		nonces:  nonces,
	}
}

var (
	testSig *PartialSignature
	testErr error
)

// BenchmarkPartialSign benchmarks how long it takes to generate a partial
// signature factoring in if the keys are sorted and also if we're in fast sign
// mode.
func BenchmarkPartialSign(b *testing.B) {
	for _, numSigners := range []int{10, 100} {
		for _, fastSign := range []bool{true, false} {
			for _, sortKeys := range []bool{true, false} {
				name := fmt.Sprintf("num_signers=%v/fast_sign=%v/sort=%v",
					numSigners, fastSign, sortKeys)

				signers := make(signerSet, numSigners)
				for i := 0; i < numSigners; i++ {
					signers[i] = genSigner(b)
				}

				combinedNonce, err := AggregateNonces(signers.pubNonces())
				if err != nil {
					b.Fatalf("unable to generate combined nonce: %v", err)
				}

				var sig *PartialSignature

				var msg [32]byte
				copy(msg[:], testMsg[:])

				keys := signers.keys()

				b.Run(name, func(b *testing.B) {
					var signOpts []SignOption
					if fastSign {
						signOpts = append(signOpts, WithFastSign())
					}
					if sortKeys {
						signOpts = append(signOpts, WithSortedKeys())
					}

					b.ResetTimer()
					b.ReportAllocs()

					for i := 0; i < b.N; i++ {
						sig, err = Sign(
							signers[0].nonces.SecNonce, signers[0].privKey,
							combinedNonce, keys, msg, signOpts...,
						)
						if err != nil {
							b.Fatalf("unable to generate sig: %v", err)
						}
					}

					testSig = sig
					testErr = err
				})
			}
		}
	}
}

// TODO(roasbeef): add impact of sorting ^

var sigOk bool

// BenchmarkPartialVerify benchmarks how long it takes to verify a partial
// signature.
func BenchmarkPartialVerify(b *testing.B) {
	for _, numSigners := range []int{10, 100} {
		for _, sortKeys := range []bool{true, false} {
			name := fmt.Sprintf("sort_keys=%v/num_signers=%v",
				sortKeys, numSigners)

			signers := make(signerSet, numSigners)
			for i := 0; i < numSigners; i++ {
				signers[i] = genSigner(b)
			}

			combinedNonce, err := AggregateNonces(
				signers.pubNonces(),
			)
			if err != nil {
				b.Fatalf("unable to generate combined "+
					"nonce: %v", err)
			}

			var sig *PartialSignature

			var msg [32]byte
			copy(msg[:], testMsg[:])

			b.ReportAllocs()
			b.ResetTimer()

			sig, err = Sign(
				signers[0].nonces.SecNonce, signers[0].privKey,
				combinedNonce, signers.keys(), msg,
			)
			if err != nil {
				b.Fatalf("unable to generate sig: %v", err)
			}

			keys := signers.keys()
			pubKey := signers[0].pubKey

			b.Run(name, func(b *testing.B) {
				var signOpts []SignOption
				if sortKeys {
					signOpts = append(
						signOpts, WithSortedKeys(),
					)
				}

				b.ResetTimer()
				b.ReportAllocs()

				var ok bool
				for i := 0; i < b.N; i++ {
					ok = sig.Verify(
						signers[0].nonces.PubNonce, combinedNonce,
						keys, pubKey, msg, signOpts...,
					)
					if !ok {
						b.Fatalf("generated invalid sig!")
					}
				}
				sigOk = ok
			})

		}
	}
}

var finalSchnorrSig *schnorr.Signature

// BenchmarkCombineSigs benchmarks how long it takes to combine a set amount of
// signatures.
func BenchmarkCombineSigs(b *testing.B) {

	for _, numSigners := range []int{10, 100} {
		signers := make(signerSet, numSigners)
		for i := 0; i < numSigners; i++ {
			signers[i] = genSigner(b)
		}

		combinedNonce, err := AggregateNonces(signers.pubNonces())
		if err != nil {
			b.Fatalf("unable to generate combined nonce: %v", err)
		}

		var msg [32]byte
		copy(msg[:], testMsg[:])

		var finalNonce *btcec.PublicKey
		for i := range signers {
			signer := signers[i]
			partialSig, err := Sign(
				signer.nonces.SecNonce, signer.privKey,
				combinedNonce, signers.keys(), msg,
			)
			if err != nil {
				b.Fatalf("unable to generate partial sig: %v",
					err)
			}

			signers[i].partialSig = partialSig

			if finalNonce == nil {
				finalNonce = partialSig.R
			}
		}

		sigs := signers.partialSigs()

		name := fmt.Sprintf("num_signers=%v", numSigners)
		b.Run(name, func(b *testing.B) {
			b.ResetTimer()
			b.ReportAllocs()

			finalSig := CombineSigs(finalNonce, sigs)

			finalSchnorrSig = finalSig
		})
	}
}

var testNonce [PubNonceSize]byte

// BenchmarkAggregateNonces benchmarks how long it takes to combine nonces.
func BenchmarkAggregateNonces(b *testing.B) {
	for _, numSigners := range []int{10, 100} {
		signers := make(signerSet, numSigners)
		for i := 0; i < numSigners; i++ {
			signers[i] = genSigner(b)
		}

		nonces := signers.pubNonces()

		name := fmt.Sprintf("num_signers=%v", numSigners)
		b.Run(name, func(b *testing.B) {
			b.ResetTimer()
			b.ReportAllocs()

			pubNonce, err := AggregateNonces(nonces)
			if err != nil {
				b.Fatalf("unable to generate nonces: %v", err)
			}

			testNonce = pubNonce
		})
	}
}

var testKey *btcec.PublicKey

// BenchmarkAggregateKeys benchmarks how long it takes to aggregate public
// keys.
func BenchmarkAggregateKeys(b *testing.B) {
	for _, numSigners := range []int{10, 100} {
		for _, sortKeys := range []bool{true, false} {
			signers := make(signerSet, numSigners)
			for i := 0; i < numSigners; i++ {
				signers[i] = genSigner(b)
			}

			signerKeys := signers.keys()

			name := fmt.Sprintf("num_signers=%v/sort_keys=%v",
				numSigners, sortKeys)

			uniqueKeyIndex := secondUniqueKeyIndex(signerKeys, false)

			b.Run(name, func(b *testing.B) {
				b.ResetTimer()
				b.ReportAllocs()

				aggKey, _, _, _ := AggregateKeys(
					signerKeys, sortKeys,
					WithUniqueKeyIndex(uniqueKeyIndex),
				)

				testKey = aggKey.FinalKey
			})
		}
	}
}
<commit_msg>btcec/schnorr/musig2: fix BenchmarkPartialVerify<commit_after>// Copyright 2013-2022 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package musig2

import (
	"encoding/hex"
	"fmt"
	"testing"

	"github.com/btcsuite/btcd/btcec/v2"
	"github.com/btcsuite/btcd/btcec/v2/schnorr"
)

var (
	testPrivBytes = hexToModNScalar("9e0699c91ca1e3b7e3c9ba71eb71c89890872be97576010fe593fbf3fd57e66d")

	testMsg = hexToBytes("c301ba9de5d6053caad9f5eb46523f007702add2c62fa39de03146a36b8026b7")
)

func hexToBytes(s string) []byte {
	b, err := hex.DecodeString(s)
	if err != nil {
		panic("invalid hex in source file: " + s)
	}
	return b
}

func hexToModNScalar(s string) *btcec.ModNScalar {
	b, err := hex.DecodeString(s)
	if err != nil {
		panic("invalid hex in source file: " + s)
	}
	var scalar btcec.ModNScalar
	if overflow := scalar.SetByteSlice(b); overflow {
		panic("hex in source file overflows mod N scalar: " + s)
	}
	return &scalar
}

func genSigner(t *testing.B) signer {
	privKey, err := btcec.NewPrivateKey()
	if err != nil {
		t.Fatalf("unable to gen priv key: %v", err)
	}

	pubKey := privKey.PubKey()

	nonces, err := GenNonces(WithPublicKey(pubKey))
	if err != nil {
		t.Fatalf("unable to gen nonces: %v", err)
	}

	return signer{
		privKey: privKey,
		pubKey:  pubKey,
		nonces:  nonces,
	}
}

var (
	testSig *PartialSignature
	testErr error
)

// BenchmarkPartialSign benchmarks how long it takes to generate a partial
// signature factoring in if the keys are sorted and also if we're in fast sign
// mode.
func BenchmarkPartialSign(b *testing.B) {
	for _, numSigners := range []int{10, 100} {
		for _, fastSign := range []bool{true, false} {
			for _, sortKeys := range []bool{true, false} {
				name := fmt.Sprintf("num_signers=%v/fast_sign=%v/sort=%v",
					numSigners, fastSign, sortKeys)

				signers := make(signerSet, numSigners)
				for i := 0; i < numSigners; i++ {
					signers[i] = genSigner(b)
				}

				combinedNonce, err := AggregateNonces(signers.pubNonces())
				if err != nil {
					b.Fatalf("unable to generate combined nonce: %v", err)
				}

				var sig *PartialSignature

				var msg [32]byte
				copy(msg[:], testMsg[:])

				keys := signers.keys()

				b.Run(name, func(b *testing.B) {
					var signOpts []SignOption
					if fastSign {
						signOpts = append(signOpts, WithFastSign())
					}
					if sortKeys {
						signOpts = append(signOpts, WithSortedKeys())
					}

					b.ResetTimer()
					b.ReportAllocs()

					for i := 0; i < b.N; i++ {
						sig, err = Sign(
							signers[0].nonces.SecNonce, signers[0].privKey,
							combinedNonce, keys, msg, signOpts...,
						)
						if err != nil {
							b.Fatalf("unable to generate sig: %v", err)
						}
					}

					testSig = sig
					testErr = err
				})
			}
		}
	}
}

// TODO(roasbeef): add impact of sorting ^

var sigOk bool

// BenchmarkPartialVerify benchmarks how long it takes to verify a partial
// signature.
func BenchmarkPartialVerify(b *testing.B) {
	for _, numSigners := range []int{10, 100} {
		for _, sortKeys := range []bool{true, false} {
			name := fmt.Sprintf("sort_keys=%v/num_signers=%v",
				sortKeys, numSigners)

			signers := make(signerSet, numSigners)
			for i := 0; i < numSigners; i++ {
				signers[i] = genSigner(b)
			}

			combinedNonce, err := AggregateNonces(
				signers.pubNonces(),
			)
			if err != nil {
				b.Fatalf("unable to generate combined "+
					"nonce: %v", err)
			}

			var sig *PartialSignature

			var msg [32]byte
			copy(msg[:], testMsg[:])

			b.ReportAllocs()
			b.ResetTimer()

			sig, err = Sign(
				signers[0].nonces.SecNonce, signers[0].privKey,
				combinedNonce, signers.keys(), msg,
			)
			if err != nil {
				b.Fatalf("unable to generate sig: %v", err)
			}

			keys := signers.keys()
			pubKey := signers[0].pubKey

			b.Run(name, func(b *testing.B) {
				var signOpts []SignOption
				if sortKeys {
					signOpts = append(
						signOpts, WithSortedKeys(),
					)
				}

				b.ResetTimer()
				b.ReportAllocs()

				var ok bool
				for i := 0; i < b.N; i++ {
					ok = sig.Verify(
						signers[0].nonces.PubNonce, combinedNonce,
						keys, pubKey, msg, signOpts...,
					)
					if !ok {
						b.Fatalf("generated invalid sig!")
					}
				}
				sigOk = ok
			})

		}
	}
}

var finalSchnorrSig *schnorr.Signature

// BenchmarkCombineSigs benchmarks how long it takes to combine a set amount of
// signatures.
func BenchmarkCombineSigs(b *testing.B) {

	for _, numSigners := range []int{10, 100} {
		signers := make(signerSet, numSigners)
		for i := 0; i < numSigners; i++ {
			signers[i] = genSigner(b)
		}

		combinedNonce, err := AggregateNonces(signers.pubNonces())
		if err != nil {
			b.Fatalf("unable to generate combined nonce: %v", err)
		}

		var msg [32]byte
		copy(msg[:], testMsg[:])

		var finalNonce *btcec.PublicKey
		for i := range signers {
			signer := signers[i]
			partialSig, err := Sign(
				signer.nonces.SecNonce, signer.privKey,
				combinedNonce, signers.keys(), msg,
			)
			if err != nil {
				b.Fatalf("unable to generate partial sig: %v",
					err)
			}

			signers[i].partialSig = partialSig

			if finalNonce == nil {
				finalNonce = partialSig.R
			}
		}

		sigs := signers.partialSigs()

		name := fmt.Sprintf("num_signers=%v", numSigners)
		b.Run(name, func(b *testing.B) {
			b.ResetTimer()
			b.ReportAllocs()

			finalSig := CombineSigs(finalNonce, sigs)

			finalSchnorrSig = finalSig
		})
	}
}

var testNonce [PubNonceSize]byte

// BenchmarkAggregateNonces benchmarks how long it takes to combine nonces.
func BenchmarkAggregateNonces(b *testing.B) {
	for _, numSigners := range []int{10, 100} {
		signers := make(signerSet, numSigners)
		for i := 0; i < numSigners; i++ {
			signers[i] = genSigner(b)
		}

		nonces := signers.pubNonces()

		name := fmt.Sprintf("num_signers=%v", numSigners)
		b.Run(name, func(b *testing.B) {
			b.ResetTimer()
			b.ReportAllocs()

			pubNonce, err := AggregateNonces(nonces)
			if err != nil {
				b.Fatalf("unable to generate nonces: %v", err)
			}

			testNonce = pubNonce
		})
	}
}

var testKey *btcec.PublicKey

// BenchmarkAggregateKeys benchmarks how long it takes to aggregate public
// keys.
func BenchmarkAggregateKeys(b *testing.B) {
	for _, numSigners := range []int{10, 100} {
		for _, sortKeys := range []bool{true, false} {
			signers := make(signerSet, numSigners)
			for i := 0; i < numSigners; i++ {
				signers[i] = genSigner(b)
			}

			signerKeys := signers.keys()

			name := fmt.Sprintf("num_signers=%v/sort_keys=%v",
				numSigners, sortKeys)

			uniqueKeyIndex := secondUniqueKeyIndex(signerKeys, false)

			b.Run(name, func(b *testing.B) {
				b.ResetTimer()
				b.ReportAllocs()

				aggKey, _, _, _ := AggregateKeys(
					signerKeys, sortKeys,
					WithUniqueKeyIndex(uniqueKeyIndex),
				)

				testKey = aggKey.FinalKey
			})
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
    "encoding/json"
    "math/rand"
    "net/http"

    "github.com/lye/mustache"
)

type CaptchaTask struct {
    Task   string
    Id     string
    Answer string
}

var (
    nextTask int = 0
    CaptchaTasks = []CaptchaTask{{
        Task:   "8 + 4 =",
        Id:     "666",
        Answer: "dvylika",
    }}
)

func GetTask() *CaptchaTask {
    return &CaptchaTasks[nextTask]
}

func GetTaskById(id string) *CaptchaTask {
    // TODO: impl
    return &CaptchaTasks[0]
}

func SetNextTask(task int) {
    if task < 0 {
        task = rand.Int() % len(CaptchaTasks)
    }
    nextTask = task
}

func CheckCaptcha(task *CaptchaTask, input string) bool {
    return input == task.Answer
}

func CaptchaHtml(task *CaptchaTask) string {
    return mustache.RenderFile("tmpl/captcha.mustache", map[string]interface{}{
        "CaptchaTask": task.Task,
    })
}

func WrongCaptchaReply(w http.ResponseWriter, req *http.Request, status string, task *CaptchaTask) {
    var response = map[string]interface{}{
        "status":     status,
        "captcha-id": task.Id,
        "name":       req.FormValue("name"),
        "email":      req.FormValue("email"),
        "website":    req.FormValue("website"),
        "body":       req.FormValue("text"),
    }
    b, err := json.Marshal(response)
    if err != nil {
        logger.Println(err.Error())
        return
    }
    w.Write(b)
}

func RightCaptchaReply(w http.ResponseWriter, redir string) {
    var response = map[string]interface{}{
        "status": "accepted",
        "redir":  redir,
    }
    b, err := json.Marshal(response)
    if err != nil {
        logger.Println(err.Error())
        return
    }
    w.Write(b)
}
<commit_msg>Unhardcode single captcha question<commit_after>package main

import (
    "encoding/json"
    "fmt"
    "math/rand"
    "net/http"

    "github.com/lye/mustache"
)

type CaptchaTask struct {
    Task   string
    Id     string
    Answer string
}

var (
    nextTask     int = 0
    CaptchaTasks []CaptchaTask
)

func init() {
    answers := []string{
        "vienuolika",
        "dvylika",
        "trylika",
        "keturiolika",
        "penkiolika",
        "šešiolika",
        "septyniolika",
        "aštuoniolika",
        "devyniolika",
    }
    CaptchaTasks = make([]CaptchaTask, 0, 0)
    for i := 2; i < 11; i++ {
        task := CaptchaTask{
            Task:   fmt.Sprintf("9 + %d =", i),
            Id:     fmt.Sprintf("%d", 666+i),
            Answer: answers[i-2],
        }
        CaptchaTasks = append(CaptchaTasks, task)
    }
}

func GetTask() *CaptchaTask {
    return &CaptchaTasks[nextTask]
}

func GetTaskById(id string) *CaptchaTask {
    for _, t := range CaptchaTasks {
        if t.Id == id {
            return &t
        }
    }
    return &CaptchaTasks[0]
}

func SetNextTask(task int) {
    if task < 0 {
        task = rand.Int() % len(CaptchaTasks)
    }
    nextTask = task
}

func CheckCaptcha(task *CaptchaTask, input string) bool {
    return input == task.Answer
}

func CaptchaHtml(task *CaptchaTask) string {
    return mustache.RenderFile("tmpl/captcha.mustache", map[string]interface{}{
        "CaptchaTask": task.Task,
    })
}

func WrongCaptchaReply(w http.ResponseWriter, req *http.Request, status string, task *CaptchaTask) {
    var response = map[string]interface{}{
        "status":     status,
        "captcha-id": task.Id,
        "name":       req.FormValue("name"),
        "email":      req.FormValue("email"),
        "website":    req.FormValue("website"),
        "body":       req.FormValue("text"),
    }
    b, err := json.Marshal(response)
    if err != nil {
        logger.Println(err.Error())
        return
    }
    w.Write(b)
}

func RightCaptchaReply(w http.ResponseWriter, redir string) {
    var response = map[string]interface{}{
        "status": "accepted",
        "redir":  redir,
    }
    b, err := json.Marshal(response)
    if err != nil {
        logger.Println(err.Error())
        return
    }
    w.Write(b)
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"math/big"
	"strings"

	"github.com/spf13/cobra"

	"github.com/NebulousLabs/Sia/api"
	"github.com/NebulousLabs/Sia/modules"
)

// coinUnits converts a siacoin amount to base units.
func coinUnits(amount string) (string, error) {
	units := []string{"pS", "nS", "uS", "mS", "SC", "KS", "MS", "GS", "TS"}
	for i, unit := range units {
		if strings.HasSuffix(amount, unit) {
			// scan into big.Rat
			r, ok := new(big.Rat).SetString(strings.TrimSuffix(amount, unit))
			if !ok {
				return "", errors.New("malformed amount")
			}
			// convert units
			exp := 24 + 3*(int64(i)-4)
			mag := new(big.Int).Exp(big.NewInt(10), big.NewInt(exp), nil)
			r.Mul(r, new(big.Rat).SetInt(mag))
			// r must be an integer at this point
			if !r.IsInt() {
				return "", errors.New("non-integer number of hastings")
			}
			return r.RatString(), nil
		}
	}
	return amount, nil // hastings
}

var (
	walletCmd = &cobra.Command{
		Use:   "wallet",
		Short: "Perform wallet actions",
		Long:  "Generate a new address, send coins to another wallet, or view info about the wallet.",
		Run:   wrap(walletstatuscmd),
	}

	walletAddressCmd = &cobra.Command{
		Use:   "address",
		Short: "Get a new wallet address",
		Long:  "Generate a new wallet address.",
		Run:   wrap(walletaddresscmd),
	}

	walletSendCmd = &cobra.Command{
		Use:   "send [amount] [dest]",
		Short: "Send coins to another wallet",
		Long: `Send coins to another wallet. 'dest' must be a 64-byte hexadecimal address.
'amount' can be specified in units, e.g. 1.23KS. Supported units are:
	pS (pico,  10^-12 SC)
	nS (nano,  10^-9 SC)
	uS (micro, 10^-6 SC)
	mS (milli, 10^-3 SC)
	SC
	KS (kilo, 10^3 SC)
	MS (mega, 10^6 SC)
	GS (giga, 10^9 SC)
	TS (tera, 10^12 SC)
If no unit is supplied, hastings (smallest possible unit, 10^-24 SC) will be assumed.`,
		Run: wrap(walletsendcmd),
	}

	walletSiafundsCmd = &cobra.Command{
		Use:   "siafunds",
		Short: "Display siafunds balance",
		Long:  "Display siafunds balance and siacoin claim balance.",
		Run:   wrap(walletsiafundscmd),
	}

	walletSiafundsSendCmd = &cobra.Command{
		Use:   "send [amount] [dest] [keyfiles]",
		Short: "Send siafunds",
		Long: `Send siafunds to an address, and transfer their siacoins to the wallet.
Run 'wallet send --help' to see a list of available units.`,
		Run: walletsiafundssendcmd, // see function docstring
	}

	walletSiafundsTrackCmd = &cobra.Command{
		Use:   "track [keyfile]",
		Short: "Track a siafund address generated by siag",
		Long:  "Track a siafund address generated by siag.",
		Run:   wrap(walletsiafundstrackcmd),
	}

	walletStatusCmd = &cobra.Command{
		Use:   "status",
		Short: "View wallet status",
		Long:  "View wallet status, including the current balance and number of addresses.",
		Run:   wrap(walletstatuscmd),
	}
)

// TODO: this should be defined outside of siac
type walletAddr struct {
	Address string
}

func walletaddresscmd() {
	addr := new(walletAddr)
	err := getAPI("/wallet/address", addr)
	if err != nil {
		fmt.Println("Could not generate new address:", err)
		return
	}
	fmt.Printf("Created new address: %s\n", addr.Address)
}

func walletsendcmd(amount, dest string) {
	adjAmount, err := coinUnits(amount)
	if err != nil {
		fmt.Println("Could not parse amount:", err)
		return
	}
	err = post("/wallet/send", fmt.Sprintf("amount=%s&destination=%s", adjAmount, dest))
	if err != nil {
		fmt.Println("Could not send:", err)
		return
	}
	fmt.Printf("Sent %s to %s\n", adjAmount, dest)
}

func walletsiafundscmd() {
	bal := new(api.WalletSiafundsBalance)
	err := getAPI("/wallet/siafunds/balance", bal)
	if err != nil {
		fmt.Println("Could not get siafunds balance:", err)
		return
	}
	fmt.Printf("Siafunds Balance: %s\nClaim Balance: %s\n", bal.SiafundBalance, bal.SiacoinClaimBalance)
}

// special because list of keyfiles is variadic
func walletsiafundssendcmd(cmd *cobra.Command, args []string) {
	if len(args) < 3 {
		cmd.Usage()
		return
	}
	amount, dest, keyfiles := args[0], args[1], args[2:]
	for i := range keyfiles {
		keyfiles[i] = abs(keyfiles[i])
	}
	qs := fmt.Sprintf("amount=%s&destination=%s&keyfiles=%s", amount, dest, strings.Join(keyfiles, ","))

	err := post("/wallet/siafunds/send", qs)
	if err != nil {
		fmt.Println("Could not track siafunds:", err)
		return
	}
	fmt.Printf("Sent %s siafunds to %s\n", amount, dest)
}

func walletsiafundstrackcmd(keyfile string) {
	err := post("/wallet/siafunds/watchsiagaddress", "keyfile="+abs(keyfile))
	if err != nil {
		fmt.Println("Could not track siafunds:", err)
		return
	}
	fmt.Printf(`Added %s to tracked siafunds.

You must restart siad to update your siafund balance.
Do not delete the original keyfile.
`, keyfile)
}

func walletstatuscmd() {
	status := new(modules.WalletInfo)
	err := getAPI("/wallet/status", status)
	if err != nil {
		fmt.Println("Could not get wallet status:", err)
		return
	}
	fmt.Printf(`Wallet status:
Balance:   %v (confirmed)
           %v (unconfirmed)
Addresses: %d
`, status.Balance, status.FullBalance, status.NumAddresses)
}
<commit_msg>display wallet balance in SC<commit_after>package main

import (
	"errors"
	"fmt"
	"math/big"
	"strings"

	"github.com/spf13/cobra"

	"github.com/NebulousLabs/Sia/api"
	"github.com/NebulousLabs/Sia/modules"
)

// coinUnits converts a siacoin amount to base units.
func coinUnits(amount string) (string, error) {
	units := []string{"pS", "nS", "uS", "mS", "SC", "KS", "MS", "GS", "TS"}
	for i, unit := range units {
		if strings.HasSuffix(amount, unit) {
			// scan into big.Rat
			r, ok := new(big.Rat).SetString(strings.TrimSuffix(amount, unit))
			if !ok {
				return "", errors.New("malformed amount")
			}
			// convert units
			exp := 24 + 3*(int64(i)-4)
			mag := new(big.Int).Exp(big.NewInt(10), big.NewInt(exp), nil)
			r.Mul(r, new(big.Rat).SetInt(mag))
			// r must be an integer at this point
			if !r.IsInt() {
				return "", errors.New("non-integer number of hastings")
			}
			return r.RatString(), nil
		}
	}
	return amount, nil // hastings
}

var (
	walletCmd = &cobra.Command{
		Use:   "wallet",
		Short: "Perform wallet actions",
		Long:  "Generate a new address, send coins to another wallet, or view info about the wallet.",
		Run:   wrap(walletstatuscmd),
	}

	walletAddressCmd = &cobra.Command{
		Use:   "address",
		Short: "Get a new wallet address",
		Long:  "Generate a new wallet address.",
		Run:   wrap(walletaddresscmd),
	}

	walletSendCmd = &cobra.Command{
		Use:   "send [amount] [dest]",
		Short: "Send coins to another wallet",
		Long: `Send coins to another wallet. 'dest' must be a 64-byte hexadecimal address.
'amount' can be specified in units, e.g. 1.23KS. Supported units are:
	pS (pico,  10^-12 SC)
	nS (nano,  10^-9 SC)
	uS (micro, 10^-6 SC)
	mS (milli, 10^-3 SC)
	SC
	KS (kilo, 10^3 SC)
	MS (mega, 10^6 SC)
	GS (giga, 10^9 SC)
	TS (tera, 10^12 SC)
If no unit is supplied, hastings (smallest possible unit, 10^-24 SC) will be assumed.`,
		Run: wrap(walletsendcmd),
	}

	walletSiafundsCmd = &cobra.Command{
		Use:   "siafunds",
		Short: "Display siafunds balance",
		Long:  "Display siafunds balance and siacoin claim balance.",
		Run:   wrap(walletsiafundscmd),
	}

	walletSiafundsSendCmd = &cobra.Command{
		Use:   "send [amount] [dest] [keyfiles]",
		Short: "Send siafunds",
		Long: `Send siafunds to an address, and transfer their siacoins to the wallet.
Run 'wallet send --help' to see a list of available units.`,
		Run: walletsiafundssendcmd, // see function docstring
	}

	walletSiafundsTrackCmd = &cobra.Command{
		Use:   "track [keyfile]",
		Short: "Track a siafund address generated by siag",
		Long:  "Track a siafund address generated by siag.",
		Run:   wrap(walletsiafundstrackcmd),
	}

	walletStatusCmd = &cobra.Command{
		Use:   "status",
		Short: "View wallet status",
		Long:  "View wallet status, including the current balance and number of addresses.",
		Run:   wrap(walletstatuscmd),
	}
)

// TODO: this should be defined outside of siac
type walletAddr struct {
	Address string
}

func walletaddresscmd() {
	addr := new(walletAddr)
	err := getAPI("/wallet/address", addr)
	if err != nil {
		fmt.Println("Could not generate new address:", err)
		return
	}
	fmt.Printf("Created new address: %s\n", addr.Address)
}

func walletsendcmd(amount, dest string) {
	adjAmount, err := coinUnits(amount)
	if err != nil {
		fmt.Println("Could not parse amount:", err)
		return
	}
	err = post("/wallet/send", fmt.Sprintf("amount=%s&destination=%s", adjAmount, dest))
	if err != nil {
		fmt.Println("Could not send:", err)
		return
	}
	fmt.Printf("Sent %s to %s\n", adjAmount, dest)
}

func walletsiafundscmd() {
	bal := new(api.WalletSiafundsBalance)
	err := getAPI("/wallet/siafunds/balance", bal)
	if err != nil {
		fmt.Println("Could not get siafunds balance:", err)
		return
	}
	fmt.Printf("Siafunds Balance: %s\nClaim Balance: %s\n", bal.SiafundBalance, bal.SiacoinClaimBalance)
}

// special because list of keyfiles is variadic
func walletsiafundssendcmd(cmd *cobra.Command, args []string) {
	if len(args) < 3 {
		cmd.Usage()
		return
	}
	amount, dest, keyfiles := args[0], args[1], args[2:]
	for i := range keyfiles {
		keyfiles[i] = abs(keyfiles[i])
	}
	qs := fmt.Sprintf("amount=%s&destination=%s&keyfiles=%s", amount, dest, strings.Join(keyfiles, ","))

	err := post("/wallet/siafunds/send", qs)
	if err != nil {
		fmt.Println("Could not track siafunds:", err)
		return
	}
	fmt.Printf("Sent %s siafunds to %s\n", amount, dest)
}

func walletsiafundstrackcmd(keyfile string) {
	err := post("/wallet/siafunds/watchsiagaddress", "keyfile="+abs(keyfile))
	if err != nil {
		fmt.Println("Could not track siafunds:", err)
		return
	}
	fmt.Printf(`Added %s to tracked siafunds.

You must restart siad to update your siafund balance.
Do not delete the original keyfile.
`, keyfile)
}

func walletstatuscmd() {
	status := new(modules.WalletInfo)
	err := getAPI("/wallet/status", status)
	if err != nil {
		fmt.Println("Could not get wallet status:", err)
		return
	}
	// divide by 1e24 to get SC
	r := new(big.Rat).SetFrac(status.Balance.Big(), new(big.Int).Exp(big.NewInt(10), big.NewInt(24), nil))
	sc, _ := r.Float64()
	fmt.Printf(`Wallet status:
Balance:   %.2f SC
Exact:     %v
Addresses: %d
`, sc, status.Balance, status.NumAddresses)
}
<|endoftext|>


<commit_before>package api

import (
	"crypto/rand"
	"encoding/json"
	"net/http"

	"github.com/docker/notary/signer"
	"github.com/docker/notary/signer/keys"
	"github.com/docker/notary/tuf/signed"
	"github.com/gorilla/mux"

	pb "github.com/docker/notary/proto"
)

// Handlers sets up all the handers for the routes, injecting a specific CryptoService object for them to use
func Handlers(cryptoServices signer.CryptoServiceIndex) *mux.Router {
	r := mux.NewRouter()

	r.Methods("GET").Path("/{ID}").Handler(KeyInfo(cryptoServices))
	r.Methods("POST").Path("/new/{Algorithm}").Handler(CreateKey(cryptoServices))
	r.Methods("POST").Path("/delete").Handler(DeleteKey(cryptoServices))
	r.Methods("POST").Path("/sign").Handler(Sign(cryptoServices))
	return r
}

// getCryptoService handles looking up the correct signing service, given the
// algorithm specified in the HTTP request. If the algorithm isn't specified
// or isn't supported, an error is returned to the client and this function
// returns a nil CryptoService
func getCryptoService(w http.ResponseWriter, algorithm string, cryptoServices signer.CryptoServiceIndex) signed.CryptoService {
	if algorithm == "" {
		http.Error(w, "algorithm not specified", http.StatusBadRequest)
		return nil
	}

	if service, ok := cryptoServices[algorithm]; ok {
		return service
	}

	http.Error(w, "algorithm "+algorithm+" not supported", http.StatusBadRequest)
	return nil
}

// KeyInfo returns a Handler that given a specific Key ID param, returns the public key bits of that key
func KeyInfo(cryptoServices signer.CryptoServiceIndex) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		vars := mux.Vars(r)

		tufKey, _, err := FindKeyByID(cryptoServices, &pb.KeyID{ID: vars["ID"]})
		if err != nil {
			switch err {
			// If we received an ErrInvalidKeyID, the key doesn't exist, return 404
			case keys.ErrInvalidKeyID:
				w.WriteHeader(http.StatusNotFound)
				w.Write([]byte(err.Error()))
				return
			// If we received anything else, it is unexpected, and we return a 500
			default:
				w.WriteHeader(http.StatusInternalServerError)
				w.Write([]byte(err.Error()))
				return
			}
		}
		key := &pb.PublicKey{
			KeyInfo: &pb.KeyInfo{
				KeyID:     &pb.KeyID{ID: tufKey.ID()},
				Algorithm: &pb.Algorithm{Algorithm: tufKey.Algorithm()},
			},
			PublicKey: tufKey.Public(),
		}
		json.NewEncoder(w).Encode(key)
		return
	})
}

// CreateKey returns a handler that generates a new key using the provided
// algorithm. Only the public component of the key is returned.
func CreateKey(cryptoServices signer.CryptoServiceIndex) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		vars := mux.Vars(r)
		cryptoService := getCryptoService(w, vars["Algorithm"], cryptoServices)
		if cryptoService == nil {
			// Error handled inside getCryptoService
			return
		}

		tufKey, err := cryptoService.Create("", vars["Algorithm"])
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte(err.Error()))
			return
		}
		key := &pb.PublicKey{
			KeyInfo: &pb.KeyInfo{
				KeyID:     &pb.KeyID{ID: tufKey.ID()},
				Algorithm: &pb.Algorithm{Algorithm: tufKey.Algorithm()},
			},
			PublicKey: tufKey.Public(),
		}
		json.NewEncoder(w).Encode(key)
		return
	})
}

// DeleteKey returns a handler that delete a specific KeyID
func DeleteKey(cryptoServices signer.CryptoServiceIndex) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var keyID *pb.KeyID
		err := json.NewDecoder(r.Body).Decode(&keyID)
		defer r.Body.Close()
		if err != nil || keyID.ID == "" {
			w.WriteHeader(http.StatusBadRequest)
			jsonErr, _ := json.Marshal("Malformed request")
			w.Write([]byte(jsonErr))
			return
		}

		_, cryptoService, err := FindKeyByID(cryptoServices, keyID)

		if err != nil {
			switch err {
			// If we received an ErrInvalidKeyID, the key doesn't exist, return 404
			case keys.ErrInvalidKeyID:
				w.WriteHeader(http.StatusNotFound)
				w.Write([]byte(err.Error()))
				return
			// If we received anything else, it is unexpected, and we return a 500
			default:
				w.WriteHeader(http.StatusInternalServerError)
				w.Write([]byte(err.Error()))
				return
			}
		}

		if err = cryptoService.RemoveKey(keyID.ID); err != nil {
			switch err {
			// If we received an ErrInvalidKeyID, the key doesn't exist, return 404
			case keys.ErrInvalidKeyID:
				w.WriteHeader(http.StatusNotFound)
				w.Write([]byte(err.Error()))
				return
			// If we received anything else, it is unexpected, and we return a 500
			default:
				w.WriteHeader(http.StatusInternalServerError)
				w.Write([]byte(err.Error()))
				return
			}
		}
		// In case we successfully delete this key, return 200
		return
	})
}

// Sign returns a handler that is able to perform signatures on a given blob
func Sign(cryptoServices signer.CryptoServiceIndex) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var sigRequest *pb.SignatureRequest
		err := json.NewDecoder(r.Body).Decode(&sigRequest)
		defer r.Body.Close()
		if err != nil || sigRequest.Content == nil ||
			sigRequest.KeyID == nil {
			w.WriteHeader(http.StatusBadRequest)
			jsonErr, _ := json.Marshal("Malformed request")
			w.Write([]byte(jsonErr))
			return
		}

		tufKey, cryptoService, err := FindKeyByID(cryptoServices, sigRequest.KeyID)
		if err == keys.ErrInvalidKeyID {
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte(err.Error()))
			return
		} else if err != nil {
			// We got an unexpected error
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte(err.Error()))
			return
		}

		privKey, _, err := cryptoService.GetPrivateKey(tufKey.ID())
		if err != nil {
			// We got an unexpected error
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte(err.Error()))
			return
		}
		sig, err := privKey.Sign(rand.Reader, sigRequest.Content, nil)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte(err.Error()))
			return
		}
		signature := &pb.Signature{
			KeyInfo: &pb.KeyInfo{
				KeyID:     &pb.KeyID{ID: tufKey.ID()},
				Algorithm: &pb.Algorithm{Algorithm: tufKey.Algorithm()},
			},
			Algorithm: &pb.Algorithm{Algorithm: privKey.SignatureAlgorithm().String()},
			Content:   sig,
		}

		json.NewEncoder(w).Encode(signature)
		return
	})
}
<commit_msg>Stop injecting to the helper function<commit_after>package api

import (
	"crypto/rand"
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/docker/notary/signer"
	"github.com/docker/notary/signer/keys"
	"github.com/docker/notary/tuf/signed"
	"github.com/gorilla/mux"

	pb "github.com/docker/notary/proto"
)

// Handlers sets up all the handers for the routes, injecting a specific CryptoService object for them to use
func Handlers(cryptoServices signer.CryptoServiceIndex) *mux.Router {
	r := mux.NewRouter()

	r.Methods("GET").Path("/{ID}").Handler(KeyInfo(cryptoServices))
	r.Methods("POST").Path("/new/{Algorithm}").Handler(CreateKey(cryptoServices))
	r.Methods("POST").Path("/delete").Handler(DeleteKey(cryptoServices))
	r.Methods("POST").Path("/sign").Handler(Sign(cryptoServices))
	return r
}

// getCryptoService handles looking up the correct signing service, given the
// algorithm specified in the HTTP request. If the algorithm isn't specified
// or isn't supported, an error is returned to the client and this function
// returns a nil CryptoService
func getCryptoService(algorithm string, cryptoServices signer.CryptoServiceIndex) (signed.CryptoService, error) {
	if algorithm == "" {
		return nil, fmt.Errorf("algorithm not specified")
	}

	if service, ok := cryptoServices[algorithm]; ok {
		return service, nil
	}

	return nil, fmt.Errorf("algorithm " + algorithm + " not supported")
}

// KeyInfo returns a Handler that given a specific Key ID param, returns the public key bits of that key
func KeyInfo(cryptoServices signer.CryptoServiceIndex) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		vars := mux.Vars(r)

		tufKey, _, err := FindKeyByID(cryptoServices, &pb.KeyID{ID: vars["ID"]})
		if err != nil {
			switch err {
			// If we received an ErrInvalidKeyID, the key doesn't exist, return 404
			case keys.ErrInvalidKeyID:
				w.WriteHeader(http.StatusNotFound)
				w.Write([]byte(err.Error()))
				return
			// If we received anything else, it is unexpected, and we return a 500
			default:
				w.WriteHeader(http.StatusInternalServerError)
				w.Write([]byte(err.Error()))
				return
			}
		}
		key := &pb.PublicKey{
			KeyInfo: &pb.KeyInfo{
				KeyID:     &pb.KeyID{ID: tufKey.ID()},
				Algorithm: &pb.Algorithm{Algorithm: tufKey.Algorithm()},
			},
			PublicKey: tufKey.Public(),
		}
		json.NewEncoder(w).Encode(key)
		return
	})
}

// CreateKey returns a handler that generates a new key using the provided
// algorithm. Only the public component of the key is returned.
func CreateKey(cryptoServices signer.CryptoServiceIndex) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		vars := mux.Vars(r)
		cryptoService, err := getCryptoService(vars["Algorithm"], cryptoServices)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		tufKey, err := cryptoService.Create("", vars["Algorithm"])
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte(err.Error()))
			return
		}
		key := &pb.PublicKey{
			KeyInfo: &pb.KeyInfo{
				KeyID:     &pb.KeyID{ID: tufKey.ID()},
				Algorithm: &pb.Algorithm{Algorithm: tufKey.Algorithm()},
			},
			PublicKey: tufKey.Public(),
		}
		json.NewEncoder(w).Encode(key)
		return
	})
}

// DeleteKey returns a handler that delete a specific KeyID
func DeleteKey(cryptoServices signer.CryptoServiceIndex) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var keyID *pb.KeyID
		err := json.NewDecoder(r.Body).Decode(&keyID)
		defer r.Body.Close()
		if err != nil || keyID.ID == "" {
			w.WriteHeader(http.StatusBadRequest)
			jsonErr, _ := json.Marshal("Malformed request")
			w.Write([]byte(jsonErr))
			return
		}

		_, cryptoService, err := FindKeyByID(cryptoServices, keyID)

		if err != nil {
			switch err {
			// If we received an ErrInvalidKeyID, the key doesn't exist, return 404
			case keys.ErrInvalidKeyID:
				w.WriteHeader(http.StatusNotFound)
				w.Write([]byte(err.Error()))
				return
			// If we received anything else, it is unexpected, and we return a 500
			default:
				w.WriteHeader(http.StatusInternalServerError)
				w.Write([]byte(err.Error()))
				return
			}
		}

		if err = cryptoService.RemoveKey(keyID.ID); err != nil {
			switch err {
			// If we received an ErrInvalidKeyID, the key doesn't exist, return 404
			case keys.ErrInvalidKeyID:
				w.WriteHeader(http.StatusNotFound)
				w.Write([]byte(err.Error()))
				return
			// If we received anything else, it is unexpected, and we return a 500
			default:
				w.WriteHeader(http.StatusInternalServerError)
				w.Write([]byte(err.Error()))
				return
			}
		}
		// In case we successfully delete this key, return 200
		return
	})
}

// Sign returns a handler that is able to perform signatures on a given blob
func Sign(cryptoServices signer.CryptoServiceIndex) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var sigRequest *pb.SignatureRequest
		err := json.NewDecoder(r.Body).Decode(&sigRequest)
		defer r.Body.Close()
		if err != nil || sigRequest.Content == nil ||
			sigRequest.KeyID == nil {
			w.WriteHeader(http.StatusBadRequest)
			jsonErr, _ := json.Marshal("Malformed request")
			w.Write([]byte(jsonErr))
			return
		}

		tufKey, cryptoService, err := FindKeyByID(cryptoServices, sigRequest.KeyID)
		if err == keys.ErrInvalidKeyID {
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte(err.Error()))
			return
		} else if err != nil {
			// We got an unexpected error
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte(err.Error()))
			return
		}

		privKey, _, err := cryptoService.GetPrivateKey(tufKey.ID())
		if err != nil {
			// We got an unexpected error
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte(err.Error()))
			return
		}
		sig, err := privKey.Sign(rand.Reader, sigRequest.Content, nil)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte(err.Error()))
			return
		}
		signature := &pb.Signature{
			KeyInfo: &pb.KeyInfo{
				KeyID:     &pb.KeyID{ID: tufKey.ID()},
				Algorithm: &pb.Algorithm{Algorithm: tufKey.Algorithm()},
			},
			Algorithm: &pb.Algorithm{Algorithm: privKey.SignatureAlgorithm().String()},
			Content:   sig,
		}

		json.NewEncoder(w).Encode(signature)
		return
	})
}
<|endoftext|>


<commit_before>// Copyright 2019 The LUCI 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 services

import (
	"context"
	"fmt"
	"math/rand"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes"
	"github.com/golang/protobuf/ptypes/empty"
	"go.chromium.org/gae/service/datastore"
	ds "go.chromium.org/gae/service/datastore"
	"go.chromium.org/gae/service/info"
	"go.chromium.org/gae/service/taskqueue"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/logging"
	log "go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/tsmon"
	"go.chromium.org/luci/common/tsmon/field"
	"go.chromium.org/luci/common/tsmon/metric"
	"go.chromium.org/luci/grpc/grpcutil"
	logdog "go.chromium.org/luci/logdog/api/endpoints/coordinator/services/v1"
	"go.chromium.org/luci/logdog/appengine/coordinator"
)

const archiveQueueName = "archiveTasks"

var (
	leaseTask = metric.NewCounter(
		"logdog/archival/lease_task",
		"Number of leased tasks from the archive queue, as seen by the coordinator",
		nil,
		field.Int("numRetries"))
	deleteTask = metric.NewCounter(
		"logdog/archival/delete_task",
		"Number of delete task request for the archive queue, as seen by the coordinator",
		nil)
)

// TaskArchival tasks an archival of a stream with the given delay at a given experiment rate.
func TaskArchival(c context.Context, state *coordinator.LogStreamState, delay time.Duration, experimentPercent uint32) error {
	// See if we want to task to the new pipeline or now.
	// Choose a number between 0-100.  This is good enough for our purposes.
	randNum := uint32(rand.Uint64() % 100)
	log.Debugf(c, "rolled a die, got %d", randNum)
	if randNum >= experimentPercent {
		log.Debugf(c, "%d >= %d, bypassing new pipeline", randNum, experimentPercent)
		return nil
	}
	log.Debugf(c, "%d < %d, using new pipeline with %s delay", randNum, experimentPercent, delay)

	// Now task the archival.
	state.Updated = clock.Now(c).UTC()
	state.ArchivalKey = []byte{'1'} // Use a fake key just to signal that we've tasked the archival.
	if err := ds.Put(c, state); err != nil {
		log.Fields{
			log.ErrorKey: err,
		}.Errorf(c, "Failed to Put() LogStream.")
		return grpcutil.Internal
	}

	project := string(coordinator.ProjectFromNamespace(state.Parent.Namespace()))
	id := string(state.ID())
	t, err := tqTask(&logdog.ArchiveTask{Project: project, Id: id})
	if err != nil {
		log.WithError(err).Errorf(c, "could not create archival task")
		return grpcutil.Internal
	}
	t.Delay = delay
	if err := taskqueue.Add(c, archiveQueueName, t); err != nil {
		log.WithError(err).Errorf(c, "could not task archival")
		return grpcutil.Internal
	}
	return nil

}

// tqTask returns a taskqueue task for an archive task.
func tqTask(task *logdog.ArchiveTask) (*taskqueue.Task, error) {
	payload, err := proto.Marshal(task)
	return &taskqueue.Task{
		Name:    task.TaskName,
		Payload: payload,
		Method:  "PULL",
	}, err
}

// tqTaskLite returns a taskqueue task for an archive task without the payload.
func tqTaskLite(task *logdog.ArchiveTask) *taskqueue.Task {
	return &taskqueue.Task{
		Name:   task.TaskName,
		Method: "PULL",
	}
}

// archiveTask creates a archiveTask proto from a taskqueue task.
func archiveTask(task *taskqueue.Task) (*logdog.ArchiveTask, error) {
	result := logdog.ArchiveTask{}
	err := proto.Unmarshal(task.Payload, &result)
	result.TaskName = task.Name
	return &result, err
}

// checkArchived is a best-effort check to see if a task is archived.
// If this fails, an error is logged, and it returns false.
func isArchived(c context.Context, task *logdog.ArchiveTask) bool {
	var err error
	defer func() {
		if err != nil {
			logging.WithError(err).Errorf(c, "while checking if %s/%s is archived", task.Project, task.Id)
		}
	}()
	if c, err = info.Namespace(c, "luci."+task.Project); err != nil {
		return false
	}
	state := (&coordinator.LogStream{ID: coordinator.HashID(task.Id)}).State(c)
	if err = datastore.Get(c, state); err != nil {
		return false
	}
	return state.ArchivalState() == coordinator.ArchivedComplete
}

// LeaseArchiveTasks leases archive tasks to the requestor from a pull queue.
func (b *server) LeaseArchiveTasks(c context.Context, req *logdog.LeaseRequest) (*logdog.LeaseResponse, error) {
	duration, err := ptypes.Duration(req.LeaseTime)
	if err != nil {
		return nil, err
	}

	tasks, err := taskqueue.Lease(c, int(req.MaxTasks), archiveQueueName, duration)
	if err != nil {
		logging.WithError(err).Errorf(c, "could not lease %d tasks from queue", req.MaxTasks)
		return nil, err
	}
	archiveTasks := make([]*logdog.ArchiveTask, 0, len(tasks))
	for _, task := range tasks {
		at, err := archiveTask(task)
		if err != nil {
			// Ignore malformed name errors, just log them.
			logging.WithError(err).Errorf(c, "while leasing task")
			continue
		}
		// Optimization: Delete the task if it's already archived.
		if isArchived(c, at) {
			logging.Infof(c, "%s/%s is already archived, deleting.", at.Project, at.Id)
			if err := taskqueue.Delete(c, archiveQueueName, task); err != nil {
				logging.WithError(err).Errorf(c, "failed to delete %s/%s (%s)", at.Project, at.Id, task.Name)
			}
			continue
		}
		logging.Infof(c, "Leasing Project: %s, ID: %s, TaskName: %s", at.Project, at.Id, at.TaskName)
		archiveTasks = append(archiveTasks, at)
		leaseTask.Add(c, 1, task.RetryCount)
	}
	logging.Infof(c, "Leasing %d tasks", len(archiveTasks))
	if err := tsmon.Flush(c); err != nil {
		logging.WithError(err).Errorf(c, "failed to flush tsmon")
	}
	return &logdog.LeaseResponse{Tasks: archiveTasks}, nil
}

// DeleteArchiveTasks deletes archive tasks from the task queue.
// Errors are logged but ignored.
func (b *server) DeleteArchiveTasks(c context.Context, req *logdog.DeleteRequest) (*empty.Empty, error) {
	deleteTask.Add(c, int64(len(req.Tasks)))
	tasks := make([]*taskqueue.Task, 0, len(req.Tasks))
	msg := fmt.Sprintf("Deleting %d tasks", len(req.Tasks))
	for _, at := range req.Tasks {
		msg += fmt.Sprintf("\nProject: %s, ID: %s, TaskName: %s", at.Project, at.Id, at.TaskName)
		tasks = append(tasks, tqTaskLite(at))
	}
	logging.Infof(c, msg)
	err := taskqueue.Delete(c, archiveQueueName, tasks...)
	if err != nil {
		logging.WithError(err).Errorf(c, "while deleting tasks\n%#v", tasks)
	}
	if err := tsmon.Flush(c); err != nil {
		logging.WithError(err).Errorf(c, "failed to flush tsmon")
	}
	return &empty.Empty{}, nil
}
<commit_msg>[logdog] Condense lease task logging<commit_after>// Copyright 2019 The LUCI 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 services

import (
	"context"
	"fmt"
	"math/rand"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes"
	"github.com/golang/protobuf/ptypes/empty"
	"go.chromium.org/gae/service/datastore"
	ds "go.chromium.org/gae/service/datastore"
	"go.chromium.org/gae/service/info"
	"go.chromium.org/gae/service/taskqueue"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/logging"
	log "go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/tsmon"
	"go.chromium.org/luci/common/tsmon/field"
	"go.chromium.org/luci/common/tsmon/metric"
	"go.chromium.org/luci/grpc/grpcutil"
	logdog "go.chromium.org/luci/logdog/api/endpoints/coordinator/services/v1"
	"go.chromium.org/luci/logdog/appengine/coordinator"
)

const archiveQueueName = "archiveTasks"

var (
	leaseTask = metric.NewCounter(
		"logdog/archival/lease_task",
		"Number of leased tasks from the archive queue, as seen by the coordinator",
		nil,
		field.Int("numRetries"))
	deleteTask = metric.NewCounter(
		"logdog/archival/delete_task",
		"Number of delete task request for the archive queue, as seen by the coordinator",
		nil)
)

// TaskArchival tasks an archival of a stream with the given delay at a given experiment rate.
func TaskArchival(c context.Context, state *coordinator.LogStreamState, delay time.Duration, experimentPercent uint32) error {
	// See if we want to task to the new pipeline or now.
	// Choose a number between 0-100.  This is good enough for our purposes.
	randNum := uint32(rand.Uint64() % 100)
	log.Debugf(c, "rolled a die, got %d", randNum)
	if randNum >= experimentPercent {
		log.Debugf(c, "%d >= %d, bypassing new pipeline", randNum, experimentPercent)
		return nil
	}
	log.Debugf(c, "%d < %d, using new pipeline with %s delay", randNum, experimentPercent, delay)

	// Now task the archival.
	state.Updated = clock.Now(c).UTC()
	state.ArchivalKey = []byte{'1'} // Use a fake key just to signal that we've tasked the archival.
	if err := ds.Put(c, state); err != nil {
		log.Fields{
			log.ErrorKey: err,
		}.Errorf(c, "Failed to Put() LogStream.")
		return grpcutil.Internal
	}

	project := string(coordinator.ProjectFromNamespace(state.Parent.Namespace()))
	id := string(state.ID())
	t, err := tqTask(&logdog.ArchiveTask{Project: project, Id: id})
	if err != nil {
		log.WithError(err).Errorf(c, "could not create archival task")
		return grpcutil.Internal
	}
	t.Delay = delay
	if err := taskqueue.Add(c, archiveQueueName, t); err != nil {
		log.WithError(err).Errorf(c, "could not task archival")
		return grpcutil.Internal
	}
	return nil

}

// tqTask returns a taskqueue task for an archive task.
func tqTask(task *logdog.ArchiveTask) (*taskqueue.Task, error) {
	payload, err := proto.Marshal(task)
	return &taskqueue.Task{
		Name:    task.TaskName,
		Payload: payload,
		Method:  "PULL",
	}, err
}

// tqTaskLite returns a taskqueue task for an archive task without the payload.
func tqTaskLite(task *logdog.ArchiveTask) *taskqueue.Task {
	return &taskqueue.Task{
		Name:   task.TaskName,
		Method: "PULL",
	}
}

// archiveTask creates a archiveTask proto from a taskqueue task.
func archiveTask(task *taskqueue.Task) (*logdog.ArchiveTask, error) {
	result := logdog.ArchiveTask{}
	err := proto.Unmarshal(task.Payload, &result)
	result.TaskName = task.Name
	return &result, err
}

// checkArchived is a best-effort check to see if a task is archived.
// If this fails, an error is logged, and it returns false.
func isArchived(c context.Context, task *logdog.ArchiveTask) bool {
	var err error
	defer func() {
		if err != nil {
			logging.WithError(err).Errorf(c, "while checking if %s/%s is archived", task.Project, task.Id)
		}
	}()
	if c, err = info.Namespace(c, "luci."+task.Project); err != nil {
		return false
	}
	state := (&coordinator.LogStream{ID: coordinator.HashID(task.Id)}).State(c)
	if err = datastore.Get(c, state); err != nil {
		return false
	}
	return state.ArchivalState() == coordinator.ArchivedComplete
}

// LeaseArchiveTasks leases archive tasks to the requestor from a pull queue.
func (b *server) LeaseArchiveTasks(c context.Context, req *logdog.LeaseRequest) (*logdog.LeaseResponse, error) {
	duration, err := ptypes.Duration(req.LeaseTime)
	if err != nil {
		return nil, err
	}

	tasks, err := taskqueue.Lease(c, int(req.MaxTasks), archiveQueueName, duration)
	if err != nil {
		logging.WithError(err).Errorf(c, "could not lease %d tasks from queue", req.MaxTasks)
		return nil, err
	}
	archiveTasks := make([]*logdog.ArchiveTask, 0, len(tasks))
	leasedTaskMessage := "Leasing (project/id/task_name):\n"
	for _, task := range tasks {
		at, err := archiveTask(task)
		if err != nil {
			// Ignore malformed name errors, just log them.
			logging.WithError(err).Errorf(c, "while leasing task")
			continue
		}
		// Optimization: Delete the task if it's already archived.
		if isArchived(c, at) {
			logging.Infof(c, "%s/%s is already archived, deleting.", at.Project, at.Id)
			if err := taskqueue.Delete(c, archiveQueueName, task); err != nil {
				logging.WithError(err).Errorf(c, "failed to delete %s/%s (%s)", at.Project, at.Id, task.Name)
			}
			continue
		}
		leasedTaskMessage += fmt.Sprintf("%s/%s/%s\n", at.Project, at.Id, at.TaskName)
		archiveTasks = append(archiveTasks, at)
		leaseTask.Add(c, 1, task.RetryCount)
	}
	logging.Infof(c, "Leasing %d tasks", len(archiveTasks))
	logging.Debugf(c, leasedTaskMessage)
	if err := tsmon.Flush(c); err != nil {
		logging.WithError(err).Errorf(c, "failed to flush tsmon")
	}
	return &logdog.LeaseResponse{Tasks: archiveTasks}, nil
}

// DeleteArchiveTasks deletes archive tasks from the task queue.
// Errors are logged but ignored.
func (b *server) DeleteArchiveTasks(c context.Context, req *logdog.DeleteRequest) (*empty.Empty, error) {
	deleteTask.Add(c, int64(len(req.Tasks)))
	tasks := make([]*taskqueue.Task, 0, len(req.Tasks))
	msg := fmt.Sprintf("Deleting %d tasks", len(req.Tasks))
	for _, at := range req.Tasks {
		msg += fmt.Sprintf("\nProject: %s, ID: %s, TaskName: %s", at.Project, at.Id, at.TaskName)
		tasks = append(tasks, tqTaskLite(at))
	}
	logging.Infof(c, msg)
	err := taskqueue.Delete(c, archiveQueueName, tasks...)
	if err != nil {
		logging.WithError(err).Errorf(c, "while deleting tasks\n%#v", tasks)
	}
	if err := tsmon.Flush(c); err != nil {
		logging.WithError(err).Errorf(c, "failed to flush tsmon")
	}
	return &empty.Empty{}, nil
}
<|endoftext|>


<commit_before>// Copyright 2014 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"os"
	"strings"
	"testing"
)

func TestSubmitErrors(t *testing.T) {
	gt := newGitTest(t)
	defer gt.done()

	srv := newGerritServer(t)
	defer srv.done()

	t.Logf("> no commit")
	testMainDied(t, "submit")
	testPrintedStderr(t, "cannot submit: no changes pending")
	write(t, gt.client+"/file1", "")
	trun(t, gt.client, "git", "add", "file1")
	trun(t, gt.client, "git", "commit", "-m", "msg\n\nChange-Id: I123456789\n")

	t.Logf("> staged changes")
	write(t, gt.client+"/file1", "asdf")
	trun(t, gt.client, "git", "add", "file1")
	testMainDied(t, "submit")
	testPrintedStderr(t, "cannot submit: staged changes exist",
		"git status", "!git stash", "!git add", "git-codereview change")
	testNoStdout(t)

	t.Logf("> unstaged changes")
	write(t, gt.client+"/file1", "actual content")
	testMainDied(t, "submit")
	testPrintedStderr(t, "cannot submit: unstaged changes exist",
		"git status", "git stash", "git add", "git-codereview change")
	testNoStdout(t)
	testRan(t)
	trun(t, gt.client, "git", "add", "file1")
	trun(t, gt.client, "git", "commit", "--amend", "--no-edit")

	t.Logf("> not found")
	testMainDied(t, "submit")
	testPrintedStderr(t, "change not found on Gerrit server")

	const id = "I123456789"

	t.Logf("> malformed json")
	srv.setJSON(id, "XXX")
	testMainDied(t, "submit")
	testRan(t) // nothing
	testPrintedStderr(t, "malformed json response")

	t.Logf("> unexpected change status")
	srv.setJSON(id, `{"status": "UNEXPECTED"}`)
	testMainDied(t, "submit")
	testRan(t) // nothing
	testPrintedStderr(t, "cannot submit: unexpected Gerrit change status \"UNEXPECTED\"")

	t.Logf("> already merged")
	srv.setJSON(id, `{"status": "MERGED"}`)
	testMainDied(t, "submit")
	testRan(t) // nothing
	testPrintedStderr(t, "cannot submit: change already submitted, run 'git sync'")

	t.Logf("> abandoned")
	srv.setJSON(id, `{"status": "ABANDONED"}`)
	testMainDied(t, "submit")
	testRan(t) // nothing
	testPrintedStderr(t, "cannot submit: change abandoned")

	t.Logf("> missing approval")
	srv.setJSON(id, `{"status": "NEW", "labels": {"Code-Review": {}}}`)
	testMainDied(t, "submit")
	testRan(t) // nothing
	testPrintedStderr(t, "cannot submit: change missing Code-Review approval")

	t.Logf("> rejection")
	srv.setJSON(id, `{"status": "NEW", "labels": {"Code-Review": {"rejected": {}}}}`)
	testMainDied(t, "submit")
	testRan(t) // nothing
	testPrintedStderr(t, "cannot submit: change has Code-Review rejection")

	t.Logf("> unmergeable")
	srv.setJSON(id, `{"status": "NEW", "mergeable": false, "labels": {"Code-Review": {"approved": {}}}}`)
	testMainDied(t, "submit")
	testRan(t, "git push -q origin HEAD:refs/for/master")
	testPrintedStderr(t, "cannot submit: conflicting changes submitted, run 'git sync'")

	t.Logf("> submit with unexpected status")
	const newJSON = `{"status": "NEW", "mergeable": true, "labels": {"Code-Review": {"approved": {}}}}`
	srv.setJSON(id, newJSON)
	srv.setReply("/a/changes/proj~master~I123456789/submit", gerritReply{body: ")]}'\n" + newJSON})
	testMainDied(t, "submit")
	testRan(t, "git push -q origin HEAD:refs/for/master")
	testPrintedStderr(t, "submit error: unexpected post-submit Gerrit change status \"NEW\"")
}

func TestSubmitTimeout(t *testing.T) {
	gt := newGitTest(t)
	defer gt.done()
	srv := newGerritServer(t)
	defer srv.done()

	gt.work(t)

	setJSON := func(json string) {
		srv.setReply("/a/changes/proj~master~I123456789", gerritReply{body: ")]}'\n" + json})
	}

	t.Log("> submit with timeout")
	const submittedJSON = `{"status": "SUBMITTED", "mergeable": true, "labels": {"Code-Review": {"approved": {}}}}`
	setJSON(submittedJSON)
	srv.setReply("/a/changes/proj~master~I123456789/submit", gerritReply{body: ")]}'\n" + submittedJSON})
	testMainDied(t, "submit")
	testRan(t, "git push -q origin HEAD:refs/for/master")
	testPrintedStderr(t, "cannot submit: timed out waiting for change to be submitted by Gerrit")
}

func TestSubmit(t *testing.T) {
	gt := newGitTest(t)
	defer gt.done()
	srv := newGerritServer(t)
	defer srv.done()

	gt.work(t)
	trun(t, gt.client, "git", "tag", "-f", "work.mailed")
	clientHead := strings.TrimSpace(trun(t, gt.client, "git", "log", "-n", "1", "--format=format:%H"))

	write(t, gt.server+"/file", "another change")
	trun(t, gt.server, "git", "add", "file")
	trun(t, gt.server, "git", "commit", "-m", "conflict")
	serverHead := strings.TrimSpace(trun(t, gt.server, "git", "log", "-n", "1", "--format=format:%H"))

	t.Log("> submit")
	var (
		newJSON       = `{"status": "NEW", "mergeable": true, "current_revision": "` + clientHead + `", "labels": {"Code-Review": {"approved": {}}}}`
		submittedJSON = `{"status": "SUBMITTED", "mergeable": true, "current_revision": "` + clientHead + `", "labels": {"Code-Review": {"approved": {}}}}`
		mergedJSON    = `{"status": "MERGED", "mergeable": true, "current_revision": "` + serverHead + `", "labels": {"Code-Review": {"approved": {}}}}`
	)
	submitted := false
	npoll := 0
	srv.setReply("/a/changes/proj~master~I123456789", gerritReply{f: func() gerritReply {
		if !submitted {
			return gerritReply{body: ")]}'\n" + newJSON}
		}
		if npoll++; npoll <= 2 {
			return gerritReply{body: ")]}'\n" + submittedJSON}
		}
		return gerritReply{body: ")]}'\n" + mergedJSON}
	}})
	srv.setReply("/a/changes/proj~master~I123456789/submit", gerritReply{f: func() gerritReply {
		if submitted {
			return gerritReply{status: 409}
		}
		submitted = true
		return gerritReply{body: ")]}'\n" + submittedJSON}
	}})
	testMain(t, "submit")
	testRan(t,
		"git fetch -q",
		"git checkout -q -B work "+serverHead+" --")
}

func TestSubmitMultiple(t *testing.T) {
	gt := newGitTest(t)
	defer gt.done()

	srv := newGerritServer(t)
	defer srv.done()

	cl1, cl2 := testSubmitMultiple(t, gt, srv)
	testMain(t, "submit", cl1.CurrentRevision, cl2.CurrentRevision)
}

func TestSubmitInteractive(t *testing.T) {
	gt := newGitTest(t)
	defer gt.done()

	srv := newGerritServer(t)
	defer srv.done()

	cl1, cl2 := testSubmitMultiple(t, gt, srv)
	os.Setenv("GIT_EDITOR", "echo "+cl1.CurrentRevision+" > ")
	testMain(t, "submit", "-i")
	if cl1.Status != "MERGED" {
		t.Fatalf("want cl1.Status == MERGED; got %v", cl1.Status)
	}
	if cl2.Status != "NEW" {
		t.Fatalf("want cl2.Status == NEW; got %v", cl1.Status)
	}
}

func testSubmitMultiple(t *testing.T, gt *gitTest, srv *gerritServer) (*GerritChange, *GerritChange) {
	write(t, gt.client+"/file1", "")
	trun(t, gt.client, "git", "add", "file1")
	trun(t, gt.client, "git", "commit", "-m", "msg\n\nChange-Id: I0000001\n")
	hash1 := strings.TrimSpace(trun(t, gt.client, "git", "log", "-n", "1", "--format=format:%H"))

	write(t, gt.client+"/file2", "")
	trun(t, gt.client, "git", "add", "file2")
	trun(t, gt.client, "git", "commit", "-m", "msg\n\nChange-Id: I0000002\n")
	hash2 := strings.TrimSpace(trun(t, gt.client, "git", "log", "-n", "1", "--format=format:%H"))

	testMainDied(t, "submit")
	testPrintedStderr(t, "cannot submit: multiple changes pending")

	cl1 := GerritChange{
		Status:          "NEW",
		Mergeable:       true,
		CurrentRevision: hash1,
		Labels:          map[string]*GerritLabel{"Code-Review": &GerritLabel{Approved: new(GerritAccount)}},
	}
	cl2 := GerritChange{
		Status:          "NEW",
		Mergeable:       false,
		CurrentRevision: hash2,
		Labels:          map[string]*GerritLabel{"Code-Review": &GerritLabel{Approved: new(GerritAccount)}},
	}

	srv.setReply("/a/changes/proj~master~I0000001", gerritReply{f: func() gerritReply {
		return gerritReply{json: cl1}
	}})
	srv.setReply("/a/changes/proj~master~I0000001/submit", gerritReply{f: func() gerritReply {
		if cl1.Status != "NEW" {
			return gerritReply{status: 409}
		}
		cl1.Status = "MERGED"
		cl2.Mergeable = true
		return gerritReply{json: cl1}
	}})
	srv.setReply("/a/changes/proj~master~I0000002", gerritReply{f: func() gerritReply {
		return gerritReply{json: cl2}
	}})
	srv.setReply("/a/changes/proj~master~I0000002/submit", gerritReply{f: func() gerritReply {
		if cl2.Status != "NEW" || !cl2.Mergeable {
			return gerritReply{status: 409}
		}
		cl2.Status = "MERGED"
		return gerritReply{json: cl2}
	}})
	return &cl1, &cl2
}
<commit_msg>git-codereview: skip TestSubmitInteractive on windows (fixes build)<commit_after>// Copyright 2014 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"os"
	"runtime"
	"strings"
	"testing"
)

func TestSubmitErrors(t *testing.T) {
	gt := newGitTest(t)
	defer gt.done()

	srv := newGerritServer(t)
	defer srv.done()

	t.Logf("> no commit")
	testMainDied(t, "submit")
	testPrintedStderr(t, "cannot submit: no changes pending")
	write(t, gt.client+"/file1", "")
	trun(t, gt.client, "git", "add", "file1")
	trun(t, gt.client, "git", "commit", "-m", "msg\n\nChange-Id: I123456789\n")

	t.Logf("> staged changes")
	write(t, gt.client+"/file1", "asdf")
	trun(t, gt.client, "git", "add", "file1")
	testMainDied(t, "submit")
	testPrintedStderr(t, "cannot submit: staged changes exist",
		"git status", "!git stash", "!git add", "git-codereview change")
	testNoStdout(t)

	t.Logf("> unstaged changes")
	write(t, gt.client+"/file1", "actual content")
	testMainDied(t, "submit")
	testPrintedStderr(t, "cannot submit: unstaged changes exist",
		"git status", "git stash", "git add", "git-codereview change")
	testNoStdout(t)
	testRan(t)
	trun(t, gt.client, "git", "add", "file1")
	trun(t, gt.client, "git", "commit", "--amend", "--no-edit")

	t.Logf("> not found")
	testMainDied(t, "submit")
	testPrintedStderr(t, "change not found on Gerrit server")

	const id = "I123456789"

	t.Logf("> malformed json")
	srv.setJSON(id, "XXX")
	testMainDied(t, "submit")
	testRan(t) // nothing
	testPrintedStderr(t, "malformed json response")

	t.Logf("> unexpected change status")
	srv.setJSON(id, `{"status": "UNEXPECTED"}`)
	testMainDied(t, "submit")
	testRan(t) // nothing
	testPrintedStderr(t, "cannot submit: unexpected Gerrit change status \"UNEXPECTED\"")

	t.Logf("> already merged")
	srv.setJSON(id, `{"status": "MERGED"}`)
	testMainDied(t, "submit")
	testRan(t) // nothing
	testPrintedStderr(t, "cannot submit: change already submitted, run 'git sync'")

	t.Logf("> abandoned")
	srv.setJSON(id, `{"status": "ABANDONED"}`)
	testMainDied(t, "submit")
	testRan(t) // nothing
	testPrintedStderr(t, "cannot submit: change abandoned")

	t.Logf("> missing approval")
	srv.setJSON(id, `{"status": "NEW", "labels": {"Code-Review": {}}}`)
	testMainDied(t, "submit")
	testRan(t) // nothing
	testPrintedStderr(t, "cannot submit: change missing Code-Review approval")

	t.Logf("> rejection")
	srv.setJSON(id, `{"status": "NEW", "labels": {"Code-Review": {"rejected": {}}}}`)
	testMainDied(t, "submit")
	testRan(t) // nothing
	testPrintedStderr(t, "cannot submit: change has Code-Review rejection")

	t.Logf("> unmergeable")
	srv.setJSON(id, `{"status": "NEW", "mergeable": false, "labels": {"Code-Review": {"approved": {}}}}`)
	testMainDied(t, "submit")
	testRan(t, "git push -q origin HEAD:refs/for/master")
	testPrintedStderr(t, "cannot submit: conflicting changes submitted, run 'git sync'")

	t.Logf("> submit with unexpected status")
	const newJSON = `{"status": "NEW", "mergeable": true, "labels": {"Code-Review": {"approved": {}}}}`
	srv.setJSON(id, newJSON)
	srv.setReply("/a/changes/proj~master~I123456789/submit", gerritReply{body: ")]}'\n" + newJSON})
	testMainDied(t, "submit")
	testRan(t, "git push -q origin HEAD:refs/for/master")
	testPrintedStderr(t, "submit error: unexpected post-submit Gerrit change status \"NEW\"")
}

func TestSubmitTimeout(t *testing.T) {
	gt := newGitTest(t)
	defer gt.done()
	srv := newGerritServer(t)
	defer srv.done()

	gt.work(t)

	setJSON := func(json string) {
		srv.setReply("/a/changes/proj~master~I123456789", gerritReply{body: ")]}'\n" + json})
	}

	t.Log("> submit with timeout")
	const submittedJSON = `{"status": "SUBMITTED", "mergeable": true, "labels": {"Code-Review": {"approved": {}}}}`
	setJSON(submittedJSON)
	srv.setReply("/a/changes/proj~master~I123456789/submit", gerritReply{body: ")]}'\n" + submittedJSON})
	testMainDied(t, "submit")
	testRan(t, "git push -q origin HEAD:refs/for/master")
	testPrintedStderr(t, "cannot submit: timed out waiting for change to be submitted by Gerrit")
}

func TestSubmit(t *testing.T) {
	gt := newGitTest(t)
	defer gt.done()
	srv := newGerritServer(t)
	defer srv.done()

	gt.work(t)
	trun(t, gt.client, "git", "tag", "-f", "work.mailed")
	clientHead := strings.TrimSpace(trun(t, gt.client, "git", "log", "-n", "1", "--format=format:%H"))

	write(t, gt.server+"/file", "another change")
	trun(t, gt.server, "git", "add", "file")
	trun(t, gt.server, "git", "commit", "-m", "conflict")
	serverHead := strings.TrimSpace(trun(t, gt.server, "git", "log", "-n", "1", "--format=format:%H"))

	t.Log("> submit")
	var (
		newJSON       = `{"status": "NEW", "mergeable": true, "current_revision": "` + clientHead + `", "labels": {"Code-Review": {"approved": {}}}}`
		submittedJSON = `{"status": "SUBMITTED", "mergeable": true, "current_revision": "` + clientHead + `", "labels": {"Code-Review": {"approved": {}}}}`
		mergedJSON    = `{"status": "MERGED", "mergeable": true, "current_revision": "` + serverHead + `", "labels": {"Code-Review": {"approved": {}}}}`
	)
	submitted := false
	npoll := 0
	srv.setReply("/a/changes/proj~master~I123456789", gerritReply{f: func() gerritReply {
		if !submitted {
			return gerritReply{body: ")]}'\n" + newJSON}
		}
		if npoll++; npoll <= 2 {
			return gerritReply{body: ")]}'\n" + submittedJSON}
		}
		return gerritReply{body: ")]}'\n" + mergedJSON}
	}})
	srv.setReply("/a/changes/proj~master~I123456789/submit", gerritReply{f: func() gerritReply {
		if submitted {
			return gerritReply{status: 409}
		}
		submitted = true
		return gerritReply{body: ")]}'\n" + submittedJSON}
	}})
	testMain(t, "submit")
	testRan(t,
		"git fetch -q",
		"git checkout -q -B work "+serverHead+" --")
}

func TestSubmitMultiple(t *testing.T) {
	gt := newGitTest(t)
	defer gt.done()

	srv := newGerritServer(t)
	defer srv.done()

	cl1, cl2 := testSubmitMultiple(t, gt, srv)
	testMain(t, "submit", cl1.CurrentRevision, cl2.CurrentRevision)
}

func TestSubmitInteractive(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Skip("see golang.org/issue/13406")
	}

	gt := newGitTest(t)
	defer gt.done()

	srv := newGerritServer(t)
	defer srv.done()

	cl1, cl2 := testSubmitMultiple(t, gt, srv)
	os.Setenv("GIT_EDITOR", "echo "+cl1.CurrentRevision+" > ")
	testMain(t, "submit", "-i")
	if cl1.Status != "MERGED" {
		t.Fatalf("want cl1.Status == MERGED; got %v", cl1.Status)
	}
	if cl2.Status != "NEW" {
		t.Fatalf("want cl2.Status == NEW; got %v", cl1.Status)
	}
}

func testSubmitMultiple(t *testing.T, gt *gitTest, srv *gerritServer) (*GerritChange, *GerritChange) {
	write(t, gt.client+"/file1", "")
	trun(t, gt.client, "git", "add", "file1")
	trun(t, gt.client, "git", "commit", "-m", "msg\n\nChange-Id: I0000001\n")
	hash1 := strings.TrimSpace(trun(t, gt.client, "git", "log", "-n", "1", "--format=format:%H"))

	write(t, gt.client+"/file2", "")
	trun(t, gt.client, "git", "add", "file2")
	trun(t, gt.client, "git", "commit", "-m", "msg\n\nChange-Id: I0000002\n")
	hash2 := strings.TrimSpace(trun(t, gt.client, "git", "log", "-n", "1", "--format=format:%H"))

	testMainDied(t, "submit")
	testPrintedStderr(t, "cannot submit: multiple changes pending")

	cl1 := GerritChange{
		Status:          "NEW",
		Mergeable:       true,
		CurrentRevision: hash1,
		Labels:          map[string]*GerritLabel{"Code-Review": &GerritLabel{Approved: new(GerritAccount)}},
	}
	cl2 := GerritChange{
		Status:          "NEW",
		Mergeable:       false,
		CurrentRevision: hash2,
		Labels:          map[string]*GerritLabel{"Code-Review": &GerritLabel{Approved: new(GerritAccount)}},
	}

	srv.setReply("/a/changes/proj~master~I0000001", gerritReply{f: func() gerritReply {
		return gerritReply{json: cl1}
	}})
	srv.setReply("/a/changes/proj~master~I0000001/submit", gerritReply{f: func() gerritReply {
		if cl1.Status != "NEW" {
			return gerritReply{status: 409}
		}
		cl1.Status = "MERGED"
		cl2.Mergeable = true
		return gerritReply{json: cl1}
	}})
	srv.setReply("/a/changes/proj~master~I0000002", gerritReply{f: func() gerritReply {
		return gerritReply{json: cl2}
	}})
	srv.setReply("/a/changes/proj~master~I0000002/submit", gerritReply{f: func() gerritReply {
		if cl2.Status != "NEW" || !cl2.Mergeable {
			return gerritReply{status: 409}
		}
		cl2.Status = "MERGED"
		return gerritReply{json: cl2}
	}})
	return &cl1, &cl2
}
<|endoftext|>


<commit_before>package zygo

import (
	"errors"
	"fmt"
	"time"
)

var UtcTz *time.Location
var NYC *time.Location

func init() {
	var err error
	UtcTz, err = time.LoadLocation("UTC")
	panicOn(err)
	NYC, err = time.LoadLocation("America/New_York")
	panicOn(err)
}

type SexpTime struct {
	Tm time.Time
}

func (r *SexpTime) Type() *RegisteredType {
	return nil // TODO what should this be?
}

func (t *SexpTime) SexpString(ps *PrintState) string {
	return t.Tm.String()
}

func NowFunction(env *Zlisp, name string,
	args []Sexp) (Sexp, error) {
	return &SexpTime{Tm: time.Now()}, nil
}

// string -> time.Time
func AsTmFunction(env *Zlisp, name string,
	args []Sexp) (Sexp, error) {

	if len(args) != 1 {
		return SexpNull, WrongNargs
	}

	var str *SexpStr
	switch t := args[0].(type) {
	case *SexpStr:
		str = t
	default:
		return SexpNull,
			errors.New("argument of astm should be a string RFC3999Nano timestamp that we want to convert to time.Time")
	}

	tm, err := time.ParseInLocation(time.RFC3339Nano, str.S, NYC)
	if err != nil {
		return SexpNull, err
	}
	return &SexpTime{Tm: tm.In(NYC)}, nil
}

func TimeitFunction(env *Zlisp, name string,
	args []Sexp) (Sexp, error) {
	if len(args) != 1 {
		return SexpNull, WrongNargs
	}

	var fun *SexpFunction
	switch t := args[0].(type) {
	case *SexpFunction:
		fun = t
	default:
		return SexpNull,
			errors.New("argument of timeit should be function")
	}

	starttime := time.Now()
	elapsed := time.Since(starttime)
	maxseconds := 10.0
	var iterations int

	for iterations = 0; iterations < 10000; iterations++ {
		_, err := env.Apply(fun, []Sexp{})
		if err != nil {
			return SexpNull, err
		}
		elapsed = time.Since(starttime)
		if elapsed.Seconds() > maxseconds {
			break
		}
	}

	fmt.Printf("ran %d iterations in %f seconds\n",
		iterations, elapsed.Seconds())
	fmt.Printf("average %f seconds per run\n",
		elapsed.Seconds()/float64(iterations))

	return SexpNull, nil
}

func (env *Zlisp) ImportTime() {
	env.AddFunction("now", NowFunction)
	env.AddFunction("timeit", TimeitFunction)
	env.AddFunction("astm", AsTmFunction)
}
<commit_msg>New function "millis" that returns milliseconds since Unix Epoch<commit_after>package zygo

import (
	"errors"
	"fmt"
	"time"
)

var UtcTz *time.Location
var NYC *time.Location

func init() {
	var err error
	UtcTz, err = time.LoadLocation("UTC")
	panicOn(err)
	NYC, err = time.LoadLocation("America/New_York")
	panicOn(err)
}

type SexpTime struct {
	Tm time.Time
}

func (r *SexpTime) Type() *RegisteredType {
	return nil // TODO what should this be?
}

func (t *SexpTime) SexpString(ps *PrintState) string {
	return t.Tm.String()
}

func NowFunction(env *Zlisp, name string,
	args []Sexp) (Sexp, error) {
	return &SexpTime{Tm: time.Now()}, nil
}

// string -> time.Time
func AsTmFunction(env *Zlisp, name string,
	args []Sexp) (Sexp, error) {

	if len(args) != 1 {
		return SexpNull, WrongNargs
	}

	var str *SexpStr
	switch t := args[0].(type) {
	case *SexpStr:
		str = t
	default:
		return SexpNull,
			errors.New("argument of astm should be a string RFC3999Nano timestamp that we want to convert to time.Time")
	}

	tm, err := time.ParseInLocation(time.RFC3339Nano, str.S, NYC)
	if err != nil {
		return SexpNull, err
	}
	return &SexpTime{Tm: tm.In(NYC)}, nil
}

func TimeitFunction(env *Zlisp, name string,
	args []Sexp) (Sexp, error) {
	if len(args) != 1 {
		return SexpNull, WrongNargs
	}

	var fun *SexpFunction
	switch t := args[0].(type) {
	case *SexpFunction:
		fun = t
	default:
		return SexpNull,
			errors.New("argument of timeit should be function")
	}

	starttime := time.Now()
	elapsed := time.Since(starttime)
	maxseconds := 10.0
	var iterations int

	for iterations = 0; iterations < 10000; iterations++ {
		_, err := env.Apply(fun, []Sexp{})
		if err != nil {
			return SexpNull, err
		}
		elapsed = time.Since(starttime)
		if elapsed.Seconds() > maxseconds {
			break
		}
	}

	fmt.Printf("ran %d iterations in %f seconds\n",
		iterations, elapsed.Seconds())
	fmt.Printf("average %f seconds per run\n",
		elapsed.Seconds()/float64(iterations))

	return SexpNull, nil
}

func MillisFunction(env *Zlisp, name string,
	args []Sexp) (Sexp, error) {
	millis := time.Now().UnixNano() / 1000000
	return &SexpInt{Val: int64(millis)}, nil
}

func (env *Zlisp) ImportTime() {
	env.AddFunction("now", NowFunction)
	env.AddFunction("timeit", TimeitFunction)
	env.AddFunction("astm", AsTmFunction)
	env.AddFunction("millis", MillisFunction)
}
<|endoftext|>


<commit_before>package instance

import (
	"io"
	"os"
	"time"

	liblxc "gopkg.in/lxc/go-lxc.v2"

	"github.com/lxc/lxd/lxd/backup"
	"github.com/lxc/lxd/lxd/db"
	deviceConfig "github.com/lxc/lxd/lxd/device/config"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/operations"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/idmap"
)

// HookStart hook used when instance has started.
const HookStart = "onstart"

// HookStopNS hook used when instance has stopped but before namespaces have been destroyed.
const HookStopNS = "onstopns"

// HookStop hook used when instance has stopped.
const HookStop = "onstop"

// Possible values for the protocol argument of the Instance.Console() method.
const (
	ConsoleTypeConsole = "console"
	ConsoleTypeVGA     = "vga"
)

// ConfigReader is used to read instance config.
type ConfigReader interface {
	Project() string
	Type() instancetype.Type
	ExpandedConfig() map[string]string
	ExpandedDevices() deviceConfig.Devices
	LocalConfig() map[string]string
	LocalDevices() deviceConfig.Devices
}

// Instance interface.
type Instance interface {
	ConfigReader

	// Instance actions.
	Freeze() error
	Shutdown(timeout time.Duration) error
	Start(stateful bool) error
	Stop(stateful bool) error
	Unfreeze() error
	RegisterDevices()
	SaveConfigFile() error

	IsPrivileged() bool

	// Snapshots & migration & backups.
	Restore(source Instance, stateful bool) error
	Snapshots() ([]Instance, error)
	Backups() ([]backup.Backup, error)
	UpdateBackupFile() error

	// Config handling.
	Rename(newName string) error
	Update(newConfig db.InstanceArgs, userRequested bool) error

	Delete() error
	Export(w io.Writer, properties map[string]string) (api.ImageMetadata, error)

	// Used for security.
	DevPaths() []string

	// Live configuration.
	CGroupSet(key string, value string) error
	VolatileSet(changes map[string]string) error

	// File handling.
	FileExists(path string) error
	FilePull(srcpath string, dstpath string) (int64, int64, os.FileMode, string, []string, error)
	FilePush(fileType string, srcpath string, dstpath string, uid int64, gid int64, mode int, write string) error
	FileRemove(path string) error

	// Console - Allocate and run a console tty or a spice Unix socket.
	Console(protocol string) (*os.File, chan error, error)
	Exec(req api.InstanceExecPost, stdin *os.File, stdout *os.File, stderr *os.File) (Cmd, error)

	// Status
	Render(options ...func(response interface{}) error) (interface{}, interface{}, error)
	RenderFull() (*api.InstanceFull, interface{}, error)
	RenderState() (*api.InstanceState, error)
	IsRunning() bool
	IsFrozen() bool
	IsEphemeral() bool
	IsSnapshot() bool
	IsStateful() bool

	// Hooks.
	DeviceEventHandler(*deviceConfig.RunConfig) error
	OnHook(hookName string, args map[string]string) error

	// Properties.
	ID() int
	Location() string
	Name() string
	Description() string
	Architecture() int
	CreationDate() time.Time
	LastUsedDate() time.Time

	Profiles() []string
	InitPID() int
	State() string
	ExpiryDate() time.Time
	FillNetworkDevice(name string, m deviceConfig.Device) (deviceConfig.Device, error)

	// Paths.
	Path() string
	RootfsPath() string
	TemplatesPath() string
	StatePath() string
	LogFilePath() string
	ConsoleBufferLogPath() string
	LogPath() string
	DevicesPath() string

	// Storage.
	StoragePool() (string, error)

	// Migration.
	Migrate(args *CriuMigrationArgs) error

	// Progress reporting.
	SetOperation(op *operations.Operation)

	// FIXME: Those should be internal functions.
	// Needed for migration for now.
	StorageStart() (bool, error)
	StorageStop() (bool, error)
	DeferTemplateApply(trigger string) error
}

// Container interface is for container specific functions.
type Container interface {
	Instance

	CurrentIdmap() (*idmap.IdmapSet, error)
	DiskIdmap() (*idmap.IdmapSet, error)
	NextIdmap() (*idmap.IdmapSet, error)
	ConsoleLog(opts liblxc.ConsoleLogOptions) (string, error)
	InsertSeccompUnixDevice(prefix string, m deviceConfig.Device, pid int) error
	DevptsFd() (*os.File, error)
}

// CriuMigrationArgs arguments for CRIU migration.
type CriuMigrationArgs struct {
	Cmd          uint
	StateDir     string
	Function     string
	Stop         bool
	ActionScript bool
	DumpDir      string
	PreDumpDir   string
	Features     liblxc.CriuFeatures
}
<commit_msg>lxd/instance/instance/interface: backup.InstanceBackup usage<commit_after>package instance

import (
	"io"
	"os"
	"time"

	liblxc "gopkg.in/lxc/go-lxc.v2"

	"github.com/lxc/lxd/lxd/backup"
	"github.com/lxc/lxd/lxd/db"
	deviceConfig "github.com/lxc/lxd/lxd/device/config"
	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/operations"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/idmap"
)

// HookStart hook used when instance has started.
const HookStart = "onstart"

// HookStopNS hook used when instance has stopped but before namespaces have been destroyed.
const HookStopNS = "onstopns"

// HookStop hook used when instance has stopped.
const HookStop = "onstop"

// Possible values for the protocol argument of the Instance.Console() method.
const (
	ConsoleTypeConsole = "console"
	ConsoleTypeVGA     = "vga"
)

// ConfigReader is used to read instance config.
type ConfigReader interface {
	Project() string
	Type() instancetype.Type
	ExpandedConfig() map[string]string
	ExpandedDevices() deviceConfig.Devices
	LocalConfig() map[string]string
	LocalDevices() deviceConfig.Devices
}

// Instance interface.
type Instance interface {
	ConfigReader

	// Instance actions.
	Freeze() error
	Shutdown(timeout time.Duration) error
	Start(stateful bool) error
	Stop(stateful bool) error
	Unfreeze() error
	RegisterDevices()
	SaveConfigFile() error

	IsPrivileged() bool

	// Snapshots & migration & backups.
	Restore(source Instance, stateful bool) error
	Snapshots() ([]Instance, error)
	Backups() ([]backup.InstanceBackup, error)
	UpdateBackupFile() error

	// Config handling.
	Rename(newName string) error
	Update(newConfig db.InstanceArgs, userRequested bool) error

	Delete() error
	Export(w io.Writer, properties map[string]string) (api.ImageMetadata, error)

	// Used for security.
	DevPaths() []string

	// Live configuration.
	CGroupSet(key string, value string) error
	VolatileSet(changes map[string]string) error

	// File handling.
	FileExists(path string) error
	FilePull(srcpath string, dstpath string) (int64, int64, os.FileMode, string, []string, error)
	FilePush(fileType string, srcpath string, dstpath string, uid int64, gid int64, mode int, write string) error
	FileRemove(path string) error

	// Console - Allocate and run a console tty or a spice Unix socket.
	Console(protocol string) (*os.File, chan error, error)
	Exec(req api.InstanceExecPost, stdin *os.File, stdout *os.File, stderr *os.File) (Cmd, error)

	// Status
	Render(options ...func(response interface{}) error) (interface{}, interface{}, error)
	RenderFull() (*api.InstanceFull, interface{}, error)
	RenderState() (*api.InstanceState, error)
	IsRunning() bool
	IsFrozen() bool
	IsEphemeral() bool
	IsSnapshot() bool
	IsStateful() bool

	// Hooks.
	DeviceEventHandler(*deviceConfig.RunConfig) error
	OnHook(hookName string, args map[string]string) error

	// Properties.
	ID() int
	Location() string
	Name() string
	Description() string
	Architecture() int
	CreationDate() time.Time
	LastUsedDate() time.Time

	Profiles() []string
	InitPID() int
	State() string
	ExpiryDate() time.Time
	FillNetworkDevice(name string, m deviceConfig.Device) (deviceConfig.Device, error)

	// Paths.
	Path() string
	RootfsPath() string
	TemplatesPath() string
	StatePath() string
	LogFilePath() string
	ConsoleBufferLogPath() string
	LogPath() string
	DevicesPath() string

	// Storage.
	StoragePool() (string, error)

	// Migration.
	Migrate(args *CriuMigrationArgs) error

	// Progress reporting.
	SetOperation(op *operations.Operation)

	// FIXME: Those should be internal functions.
	// Needed for migration for now.
	StorageStart() (bool, error)
	StorageStop() (bool, error)
	DeferTemplateApply(trigger string) error
}

// Container interface is for container specific functions.
type Container interface {
	Instance

	CurrentIdmap() (*idmap.IdmapSet, error)
	DiskIdmap() (*idmap.IdmapSet, error)
	NextIdmap() (*idmap.IdmapSet, error)
	ConsoleLog(opts liblxc.ConsoleLogOptions) (string, error)
	InsertSeccompUnixDevice(prefix string, m deviceConfig.Device, pid int) error
	DevptsFd() (*os.File, error)
}

// CriuMigrationArgs arguments for CRIU migration.
type CriuMigrationArgs struct {
	Cmd          uint
	StateDir     string
	Function     string
	Stop         bool
	ActionScript bool
	DumpDir      string
	PreDumpDir   string
	Features     liblxc.CriuFeatures
}
<|endoftext|>