<commit_before>package referrer

import (
	"github.com/stretchr/testify/assert"
	"strings"
	"testing"
)

var (
	ExampleDirectDomains = []string{"example.com", "sample.com"}
)

func TestSearchNonAscii(t *testing.T) {
	url := "http://ca.search.yahoo.com/search;_ylt=A0geu8fBeW5SqVEAZ2vrFAx.;_ylc=X1MDMjExNDcyMTAwMwRfcgMyBGJjawMwbXFjc3RoOHYybjlkJTI2YiUzRDMlMjZzJTNEYWkEY3NyY3B2aWQDWmxUdFhVZ2V1eVVMYVp6c1VmRmRMUXUyMkxfbjJsSnVlY0VBQlhDWQRmcgN5ZnAtdC03MTUEZnIyA3NiLXRvcARncHJpZANVRFRzSGFBUVF0ZUZHZ2hzZ0N3VDNBBG10ZXN0aWQDbnVsbARuX3JzbHQDMARuX3N1Z2cDMARvcmlnaW4DY2Euc2VhcmNoLnlhaG9vLmNvbQRwb3MDMARwcXN0cgMEcHFzdHJsAwRxc3RybAM0NARxdWVyeQN2aW5kdWVzcHVkc25pbmcgbXlzaG9waWZ5IHJlbmf4cmluZyBta29iZXRpYwR0X3N0bXADMTM4Mjk3MjM1NDIzMwR2dGVzdGlkA01TWUNBQzE-?p=vinduespudsning+myshopify+rengøring+mkobetic&fr2=sb-top&fr=yfp-t-715&rd=r1"
	assert.True(t, strings.Contains(url, "\u00F8"))

	r := NewReferrer(url)
	direct, social, engine, err := r.Parse(ExampleDirectDomains)
	assert.NoError(t, err)
	assert.Nil(t, direct)
	assert.Nil(t, social)

	assert.NotNil(t, engine)
	assert.Equal(t, engine.Label, "Yahoo")
	assert.True(t, strings.Contains(engine.Query, "\u00F8"))
	assert.Equal(t, engine.Query, "vinduespudsning myshopify rengøring mkobetic")
}

func TestReferrerSearchWithExplicitPlus(t *testing.T) {
	url := `http://ca.search.yahoo.com/search;_ylt=A0geu8nVvm5StDIAIxHrFAx.;_ylc=X1MDMjExNDcyMTAwMwRfcgMyBGJjawMwbXFjc3RoOHYybjlkJTI2YiUzRDMlMjZzJTNEYWkEY3NyY3B2aWQDSjNTOW9rZ2V1eVVMYVp6c1VmRmRMUkdDMkxfbjJsSnV2dFVBQmZyWgRmcgN5ZnAtdC03MTUEZnIyA3NiLXRvcARncHJpZANDc01MSGlnTVFOS2k2cDRqcUxERzRBBG10ZXN0aWQDbnVsbARuX3JzbHQDMARuX3N1Z2cDMARvcmlnaW4DY2Euc2VhcmNoLnlhaG9vLmNvbQRwb3MDMARwcXN0cgMEcHFzdHJsAwRxc3RybAM0NARxdWVyeQN2aW5kdWVzcHVkc25pbmcgSk9LQVBPTEFSICIxMSArIDExIiBta29iZXRpYwR0X3N0bXADMTM4Mjk4OTYwMjg3OQR2dGVzdGlkA01TWUNBQzE-?p=vinduespudsning+JOKAPOLAR+"11+%2B+11"+mkobetic&fr2=sb-top&fr=yfp-t-715&rd=r1`

	r := NewReferrer(url)
	direct, social, engine, err := r.Parse(ExampleDirectDomains)
	assert.NoError(t, err)
	assert.Nil(t, direct)
	assert.Nil(t, social)

	assert.NotNil(t, engine)
	assert.Equal(t, engine.Label, "Yahoo")
	assert.True(t, strings.Contains(engine.Query, "11 + 11"))
	assert.Equal(t, engine.Query, `vinduespudsning JOKAPOLAR "11 + 11" mkobetic`)
}

func TestReferrerSearchWithNonAscii(t *testing.T) {
	url := `http://ca.search.yahoo.com/search;_ylt=A0geu8fBeW5SqVEAZ2vrFAx.;_ylc=X1MDMjExNDcyMTAwMwRfcgMyBGJjawMwbXFjc3RoOHYybjlkJTI2YiUzRDMlMjZzJTNEYWkEY3NyY3B2aWQDWmxUdFhVZ2V1eVVMYVp6c1VmRmRMUXUyMkxfbjJsSnVlY0VBQlhDWQRmcgN5ZnAtdC03MTUEZnIyA3NiLXRvcARncHJpZANVRFRzSGFBUVF0ZUZHZ2hzZ0N3VDNBBG10ZXN0aWQDbnVsbARuX3JzbHQDMARuX3N1Z2cDMARvcmlnaW4DY2Euc2VhcmNoLnlhaG9vLmNvbQRwb3MDMARwcXN0cgMEcHFzdHJsAwRxc3RybAM0NARxdWVyeQN2aW5kdWVzcHVkc25pbmcgbXlzaG9waWZ5IHJlbmf4cmluZyBta29iZXRpYwR0X3N0bXADMTM4Mjk3MjM1NDIzMwR2dGVzdGlkA01TWUNBQzE-?p=vinduespudsning+myshopify+rengøring+mkobetic&fr2=sb-top&fr=yfp-t-715&rd=r1`

	r := NewReferrer(url)
	direct, social, engine, err := r.Parse(ExampleDirectDomains)
	assert.NoError(t, err)
	assert.Nil(t, direct)
	assert.Nil(t, social)

	assert.NotNil(t, engine)
	assert.Equal(t, engine.Label, "Yahoo")
	assert.True(t, strings.Contains(engine.Query, "rengøring"))
	assert.Equal(t, engine.Query, `vinduespudsning myshopify rengøring mkobetic`)
}

func TestReferrerSearchWithCyrillics(t *testing.T) {
	url := `http://www.yandex.com/yandsearch?text=%D0%B1%D0%BE%D1%82%D0%B8%D0%BD%D0%BA%D0%B8%20packer-shoes&lr=87&msid=22868.18811.1382712652.60127&noreask=1`

	r := NewReferrer(url)
	direct, social, engine, err := r.Parse(ExampleDirectDomains)
	assert.NoError(t, err)
	assert.Nil(t, direct)
	assert.Nil(t, social)

	assert.NotNil(t, engine)
	assert.Equal(t, engine.Label, "Yandex")
	assert.True(t, strings.Contains(engine.Query, "ботинки"))
	assert.Equal(t, engine.Query, `ботинки packer-shoes`)
}

// func TestReferrerDirect(t *testing.T) {
// 	url := "http://example.com"

// 	r := NewReferrer(url)
// 	direct, social, engine, err := r.Parse(ExampleDirectDomains)
// 	assert.NoError(t, err)
// 	assert.Nil(t, social)
// 	assert.Nil(t, engine)

// 	assert.NotNil(t, direct)
// 	assert.Equal(t, direct.Url, url)
// 	assert.Equal(t, direct.Domain, "example.com")
// }

func TestReferrerInvalid(t *testing.T) {
	url := `buh`

	r := NewReferrer(url)
	direct, social, engine, err := r.Parse(ExampleDirectDomains)

	assert.Error(t, err)
	assert.Nil(t, direct)
	assert.Nil(t, social)
	assert.Nil(t, engine)
}
<commit_msg>Clean up test method names.<commit_after>package referrer

import (
	"github.com/stretchr/testify/assert"
	"strings"
	"testing"
)

var (
	ExampleDirectDomains = []string{"example.com", "sample.com"}
)

func TestInvalidUrl(t *testing.T) {
	url := `buh`

	r := NewReferrer(url)
	direct, social, engine, err := r.Parse(ExampleDirectDomains)

	assert.Error(t, err)
	assert.Nil(t, direct)
	assert.Nil(t, social)
	assert.Nil(t, engine)
}

func TestSearchNonAscii(t *testing.T) {
	url := "http://ca.search.yahoo.com/search;_ylt=A0geu8fBeW5SqVEAZ2vrFAx.;_ylc=X1MDMjExNDcyMTAwMwRfcgMyBGJjawMwbXFjc3RoOHYybjlkJTI2YiUzRDMlMjZzJTNEYWkEY3NyY3B2aWQDWmxUdFhVZ2V1eVVMYVp6c1VmRmRMUXUyMkxfbjJsSnVlY0VBQlhDWQRmcgN5ZnAtdC03MTUEZnIyA3NiLXRvcARncHJpZANVRFRzSGFBUVF0ZUZHZ2hzZ0N3VDNBBG10ZXN0aWQDbnVsbARuX3JzbHQDMARuX3N1Z2cDMARvcmlnaW4DY2Euc2VhcmNoLnlhaG9vLmNvbQRwb3MDMARwcXN0cgMEcHFzdHJsAwRxc3RybAM0NARxdWVyeQN2aW5kdWVzcHVkc25pbmcgbXlzaG9waWZ5IHJlbmf4cmluZyBta29iZXRpYwR0X3N0bXADMTM4Mjk3MjM1NDIzMwR2dGVzdGlkA01TWUNBQzE-?p=vinduespudsning+myshopify+rengøring+mkobetic&fr2=sb-top&fr=yfp-t-715&rd=r1"
	assert.True(t, strings.Contains(url, "\u00F8"))

	r := NewReferrer(url)
	direct, social, engine, err := r.Parse(ExampleDirectDomains)
	assert.NoError(t, err)
	assert.Nil(t, direct)
	assert.Nil(t, social)

	assert.NotNil(t, engine)
	assert.Equal(t, engine.Label, "Yahoo")
	assert.True(t, strings.Contains(engine.Query, "\u00F8"))
	assert.Equal(t, engine.Query, "vinduespudsning myshopify rengøring mkobetic")
}

func TestSearchWithExplicitPlus(t *testing.T) {
	url := `http://ca.search.yahoo.com/search;_ylt=A0geu8nVvm5StDIAIxHrFAx.;_ylc=X1MDMjExNDcyMTAwMwRfcgMyBGJjawMwbXFjc3RoOHYybjlkJTI2YiUzRDMlMjZzJTNEYWkEY3NyY3B2aWQDSjNTOW9rZ2V1eVVMYVp6c1VmRmRMUkdDMkxfbjJsSnV2dFVBQmZyWgRmcgN5ZnAtdC03MTUEZnIyA3NiLXRvcARncHJpZANDc01MSGlnTVFOS2k2cDRqcUxERzRBBG10ZXN0aWQDbnVsbARuX3JzbHQDMARuX3N1Z2cDMARvcmlnaW4DY2Euc2VhcmNoLnlhaG9vLmNvbQRwb3MDMARwcXN0cgMEcHFzdHJsAwRxc3RybAM0NARxdWVyeQN2aW5kdWVzcHVkc25pbmcgSk9LQVBPTEFSICIxMSArIDExIiBta29iZXRpYwR0X3N0bXADMTM4Mjk4OTYwMjg3OQR2dGVzdGlkA01TWUNBQzE-?p=vinduespudsning+JOKAPOLAR+"11+%2B+11"+mkobetic&fr2=sb-top&fr=yfp-t-715&rd=r1`

	r := NewReferrer(url)
	direct, social, engine, err := r.Parse(ExampleDirectDomains)
	assert.NoError(t, err)
	assert.Nil(t, direct)
	assert.Nil(t, social)

	assert.NotNil(t, engine)
	assert.Equal(t, engine.Label, "Yahoo")
	assert.True(t, strings.Contains(engine.Query, "11 + 11"))
	assert.Equal(t, engine.Query, `vinduespudsning JOKAPOLAR "11 + 11" mkobetic`)
}

func TestSearchWithNonAscii(t *testing.T) {
	url := `http://ca.search.yahoo.com/search;_ylt=A0geu8fBeW5SqVEAZ2vrFAx.;_ylc=X1MDMjExNDcyMTAwMwRfcgMyBGJjawMwbXFjc3RoOHYybjlkJTI2YiUzRDMlMjZzJTNEYWkEY3NyY3B2aWQDWmxUdFhVZ2V1eVVMYVp6c1VmRmRMUXUyMkxfbjJsSnVlY0VBQlhDWQRmcgN5ZnAtdC03MTUEZnIyA3NiLXRvcARncHJpZANVRFRzSGFBUVF0ZUZHZ2hzZ0N3VDNBBG10ZXN0aWQDbnVsbARuX3JzbHQDMARuX3N1Z2cDMARvcmlnaW4DY2Euc2VhcmNoLnlhaG9vLmNvbQRwb3MDMARwcXN0cgMEcHFzdHJsAwRxc3RybAM0NARxdWVyeQN2aW5kdWVzcHVkc25pbmcgbXlzaG9waWZ5IHJlbmf4cmluZyBta29iZXRpYwR0X3N0bXADMTM4Mjk3MjM1NDIzMwR2dGVzdGlkA01TWUNBQzE-?p=vinduespudsning+myshopify+rengøring+mkobetic&fr2=sb-top&fr=yfp-t-715&rd=r1`

	r := NewReferrer(url)
	direct, social, engine, err := r.Parse(ExampleDirectDomains)
	assert.NoError(t, err)
	assert.Nil(t, direct)
	assert.Nil(t, social)

	assert.NotNil(t, engine)
	assert.Equal(t, engine.Label, "Yahoo")
	assert.True(t, strings.Contains(engine.Query, "rengøring"))
	assert.Equal(t, engine.Query, `vinduespudsning myshopify rengøring mkobetic`)
}

func TestSearchWithCyrillics(t *testing.T) {
	url := `http://www.yandex.com/yandsearch?text=%D0%B1%D0%BE%D1%82%D0%B8%D0%BD%D0%BA%D0%B8%20packer-shoes&lr=87&msid=22868.18811.1382712652.60127&noreask=1`

	r := NewReferrer(url)
	direct, social, engine, err := r.Parse(ExampleDirectDomains)
	assert.NoError(t, err)
	assert.Nil(t, direct)
	assert.Nil(t, social)

	assert.NotNil(t, engine)
	assert.Equal(t, engine.Label, "Yandex")
	assert.True(t, strings.Contains(engine.Query, "ботинки"))
	assert.Equal(t, engine.Query, `ботинки packer-shoes`)
}

// func TestReferrerDirect(t *testing.T) {
//  url := "http://example.com"

//  r := NewReferrer(url)
//  direct, social, engine, err := r.Parse(ExampleDirectDomains)
//  assert.NoError(t, err)
//  assert.Nil(t, social)
//  assert.Nil(t, engine)

//  assert.NotNil(t, direct)
//  assert.Equal(t, direct.Url, url)
//  assert.Equal(t, direct.Domain, "example.com")
// }
<|endoftext|>


<commit_before>package invoicesrpc

import (
	"bytes"
	"context"
	"crypto/rand"
	"errors"
	"fmt"
	"math"
	"time"

	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcutil"
	"github.com/davecgh/go-spew/spew"

	"github.com/lightningnetwork/lnd/channeldb"
	"github.com/lightningnetwork/lnd/lntypes"
	"github.com/lightningnetwork/lnd/lnwire"
	"github.com/lightningnetwork/lnd/netann"
	"github.com/lightningnetwork/lnd/zpay32"
)

// AddInvoiceConfig contains dependencies for invoice creation.
type AddInvoiceConfig struct {
	// AddInvoice is called to add the invoice to the registry.
	AddInvoice func(invoice *channeldb.Invoice, paymentHash lntypes.Hash) (
		uint64, error)

	// IsChannelActive is used to generate valid hop hints.
	IsChannelActive func(chanID lnwire.ChannelID) bool

	// ChainParams are required to properly decode invoice payment requests
	// that are marshalled over rpc.
	ChainParams *chaincfg.Params

	// NodeSigner is an implementation of the MessageSigner implementation
	// that's backed by the identity private key of the running lnd node.
	NodeSigner *netann.NodeSigner

	// MaxPaymentMSat is the maximum allowed payment.
	MaxPaymentMSat lnwire.MilliSatoshi

	// DefaultCLTVExpiry is the default invoice expiry if no values is
	// specified.
	DefaultCLTVExpiry uint32

	// ChanDB is a global boltdb instance which is needed to access the
	// channel graph.
	ChanDB *channeldb.DB
}

// AddInvoiceData contains the required data to create a new invoice.
type AddInvoiceData struct {
	// An optional memo to attach along with the invoice. Used for record
	// keeping purposes for the invoice's creator, and will also be set in
	// the description field of the encoded payment request if the
	// description_hash field is not being used.
	Memo string

	// The preimage which will allow settling an incoming HTLC payable to
	// this preimage. If Preimage is set, Hash should be nil. If both
	// Preimage and Hash are nil, a random preimage is generated.
	Preimage *lntypes.Preimage

	// The hash of the preimage. If Hash is set, Preimage should be nil.
	// This condition indicates that we have a 'hold invoice' for which the
	// htlc will be accepted and held until the preimage becomes known.
	Hash *lntypes.Hash

	// The value of this invoice in millisatoshis.
	Value lnwire.MilliSatoshi

	// Hash (SHA-256) of a description of the payment. Used if the
	// description of payment (memo) is too long to naturally fit within the
	// description field of an encoded payment request.
	DescriptionHash []byte

	// Payment request expiry time in seconds. Default is 3600 (1 hour).
	Expiry int64

	// Fallback on-chain address.
	FallbackAddr string

	// Delta to use for the time-lock of the CLTV extended to the final hop.
	CltvExpiry uint64

	// Whether this invoice should include routing hints for private
	// channels.
	Private bool
}

// AddInvoice attempts to add a new invoice to the invoice database. Any
// duplicated invoices are rejected, therefore all invoices *must* have a
// unique payment preimage.
func AddInvoice(ctx context.Context, cfg *AddInvoiceConfig,
	invoice *AddInvoiceData) (*lntypes.Hash, *channeldb.Invoice, error) {

	var (
		paymentPreimage lntypes.Preimage
		paymentHash     lntypes.Hash
	)

	switch {

	// Only either preimage or hash can be set.
	case invoice.Preimage != nil && invoice.Hash != nil:
		return nil, nil,
			errors.New("preimage and hash both set")

	// Prevent the unknown preimage magic value from being used for a
	// regular invoice. This would cause the invoice the be handled as if it
	// was a hold invoice.
	case invoice.Preimage != nil &&
		*invoice.Preimage == channeldb.UnknownPreimage:

		return nil, nil,
			fmt.Errorf("cannot use all zeroes as a preimage")

	// Prevent the hash of the unknown preimage magic value to be used for a
	// hold invoice. This would make it impossible to settle the invoice,
	// because it would still be interpreted as not having a preimage.
	case invoice.Hash != nil &&
		*invoice.Hash == channeldb.UnknownPreimage.Hash():

		return nil, nil,
			fmt.Errorf("cannot use hash of all zeroes preimage")

	// If no hash or preimage is given, generate a random preimage.
	case invoice.Preimage == nil && invoice.Hash == nil:
		if _, err := rand.Read(paymentPreimage[:]); err != nil {
			return nil, nil, err
		}
		paymentHash = paymentPreimage.Hash()

	// If just a hash is given, we create a hold invoice by setting the
	// preimage to unknown.
	case invoice.Preimage == nil && invoice.Hash != nil:
		paymentPreimage = channeldb.UnknownPreimage
		paymentHash = *invoice.Hash

	// A specific preimage was supplied. Use that for the invoice.
	case invoice.Preimage != nil && invoice.Hash == nil:
		paymentPreimage = *invoice.Preimage
		paymentHash = invoice.Preimage.Hash()
	}

	// The size of the memo, receipt and description hash attached must not
	// exceed the maximum values for either of the fields.
	if len(invoice.Memo) > channeldb.MaxMemoSize {
		return nil, nil, fmt.Errorf("memo too large: %v bytes "+
			"(maxsize=%v)", len(invoice.Memo), channeldb.MaxMemoSize)
	}
	if len(invoice.DescriptionHash) > 0 && len(invoice.DescriptionHash) != 32 {
		return nil, nil, fmt.Errorf("description hash is %v bytes, must be 32",
			len(invoice.DescriptionHash))
	}

	// The value of the invoice must not be negative.
	if invoice.Value < 0 {
		return nil, nil, fmt.Errorf("payments of negative value "+
			"are not allowed, value is %v", invoice.Value)
	}

	amtMSat := invoice.Value

	// The value of the invoice must also not exceed the current soft-limit
	// on the largest payment within the network.
	if amtMSat > cfg.MaxPaymentMSat {
		return nil, nil, fmt.Errorf("payment of %v is too large, max "+
			"payment allowed is %v", invoice.Value,
			cfg.MaxPaymentMSat.ToSatoshis(),
		)
	}

	// We also create an encoded payment request which allows the
	// caller to compactly send the invoice to the payer. We'll create a
	// list of options to be added to the encoded payment request. For now
	// we only support the required fields description/description_hash,
	// expiry, fallback address, and the amount field.
	var options []func(*zpay32.Invoice)

	// We only include the amount in the invoice if it is greater than 0.
	// By not including the amount, we enable the creation of invoices that
	// allow the payee to specify the amount of satoshis they wish to send.
	if amtMSat > 0 {
		options = append(options, zpay32.Amount(amtMSat))
	}

	// If specified, add a fallback address to the payment request.
	if len(invoice.FallbackAddr) > 0 {
		addr, err := btcutil.DecodeAddress(invoice.FallbackAddr,
			cfg.ChainParams)
		if err != nil {
			return nil, nil, fmt.Errorf("invalid fallback address: %v",
				err)
		}
		options = append(options, zpay32.FallbackAddr(addr))
	}

	// If expiry is set, specify it. If it is not provided, no expiry time
	// will be explicitly added to this payment request, which will imply
	// the default 3600 seconds.
	if invoice.Expiry > 0 {

		// We'll ensure that the specified expiry is restricted to sane
		// number of seconds. As a result, we'll reject an invoice with
		// an expiry greater than 1 year.
		maxExpiry := time.Hour * 24 * 365
		expSeconds := invoice.Expiry

		if float64(expSeconds) > maxExpiry.Seconds() {
			return nil, nil, fmt.Errorf("expiry of %v seconds "+
				"greater than max expiry of %v seconds",
				float64(expSeconds), maxExpiry.Seconds())
		}

		expiry := time.Duration(invoice.Expiry) * time.Second
		options = append(options, zpay32.Expiry(expiry))
	}

	// If the description hash is set, then we add it do the list of options.
	// If not, use the memo field as the payment request description.
	if len(invoice.DescriptionHash) > 0 {
		var descHash [32]byte
		copy(descHash[:], invoice.DescriptionHash[:])
		options = append(options, zpay32.DescriptionHash(descHash))
	} else {
		// Use the memo field as the description. If this is not set
		// this will just be an empty string.
		options = append(options, zpay32.Description(invoice.Memo))
	}

	// We'll use our current default CLTV value unless one was specified as
	// an option on the command line when creating an invoice.
	switch {
	case invoice.CltvExpiry > math.MaxUint16:
		return nil, nil, fmt.Errorf("CLTV delta of %v is too large, max "+
			"accepted is: %v", invoice.CltvExpiry, math.MaxUint16)
	case invoice.CltvExpiry != 0:
		options = append(options,
			zpay32.CLTVExpiry(invoice.CltvExpiry))
	default:
		// TODO(roasbeef): assumes set delta between versions
		defaultDelta := cfg.DefaultCLTVExpiry
		options = append(options, zpay32.CLTVExpiry(uint64(defaultDelta)))
	}

	// If we were requested to include routing hints in the invoice, then
	// we'll fetch all of our available private channels and create routing
	// hints for them.
	if invoice.Private {
		openChannels, err := cfg.ChanDB.FetchAllChannels()
		if err != nil {
			return nil, nil, fmt.Errorf("could not fetch all channels")
		}

		graph := cfg.ChanDB.ChannelGraph()

		numHints := 0
		for _, channel := range openChannels {
			// We'll restrict the number of individual route hints
			// to 20 to avoid creating overly large invoices.
			if numHints >= 20 {
				break
			}

			// Since we're only interested in our private channels,
			// we'll skip public ones.
			isPublic := channel.ChannelFlags&lnwire.FFAnnounceChannel != 0
			if isPublic {
				continue
			}

			// Make sure the counterparty has enough balance in the
			// channel for our amount. We do this in order to reduce
			// payment errors when attempting to use this channel
			// as a hint.
			chanPoint := lnwire.NewChanIDFromOutPoint(
				&channel.FundingOutpoint,
			)
			if amtMSat >= channel.LocalCommitment.RemoteBalance {
				log.Debugf("Skipping channel %v due to "+
					"not having enough remote balance",
					chanPoint)
				continue
			}

			// Make sure the channel is active.
			if !cfg.IsChannelActive(chanPoint) {
				log.Debugf("Skipping channel %v due to not "+
					"being eligible to forward payments",
					chanPoint)
				continue
			}

			// To ensure we don't leak unadvertised nodes, we'll
			// make sure our counterparty is publicly advertised
			// within the network. Otherwise, we'll end up leaking
			// information about nodes that intend to stay
			// unadvertised, like in the case of a node only having
			// private channels.
			var remotePub [33]byte
			copy(remotePub[:], channel.IdentityPub.SerializeCompressed())
			isRemoteNodePublic, err := graph.IsPublicNode(remotePub)
			if err != nil {
				log.Errorf("Unable to determine if node %x "+
					"is advertised: %v", remotePub, err)
				continue
			}

			if !isRemoteNodePublic {
				log.Debugf("Skipping channel %v due to "+
					"counterparty %x being unadvertised",
					chanPoint, remotePub)
				continue
			}

			// Fetch the policies for each end of the channel.
			chanID := channel.ShortChanID().ToUint64()
			info, p1, p2, err := graph.FetchChannelEdgesByID(chanID)
			if err != nil {
				log.Errorf("Unable to fetch the routing "+
					"policies for the edges of the channel "+
					"%v: %v", chanPoint, err)
				continue
			}

			// Now, we'll need to determine which is the correct
			// policy for HTLCs being sent from the remote node.
			var remotePolicy *channeldb.ChannelEdgePolicy
			if bytes.Equal(remotePub[:], info.NodeKey1Bytes[:]) {
				remotePolicy = p1
			} else {
				remotePolicy = p2
			}

			// If for some reason we don't yet have the edge for
			// the remote party, then we'll just skip adding this
			// channel as a routing hint.
			if remotePolicy == nil {
				continue
			}

			// Finally, create the routing hint for this channel and
			// add it to our list of route hints.
			hint := zpay32.HopHint{
				NodeID:      channel.IdentityPub,
				ChannelID:   chanID,
				FeeBaseMSat: uint32(remotePolicy.FeeBaseMSat),
				FeeProportionalMillionths: uint32(
					remotePolicy.FeeProportionalMillionths,
				),
				CLTVExpiryDelta: remotePolicy.TimeLockDelta,
			}

			// Include the route hint in our set of options that
			// will be used when creating the invoice.
			routeHint := []zpay32.HopHint{hint}
			options = append(options, zpay32.RouteHint(routeHint))

			numHints++
		}

	}

	// Set a blank feature vector, as our invoice generation forbids nil
	// features.
	invoiceFeatures := lnwire.NewFeatureVector(
		lnwire.NewRawFeatureVector(), lnwire.Features,
	)
	options = append(options, zpay32.Features(invoiceFeatures))

	// Create and encode the payment request as a bech32 (zpay32) string.
	creationDate := time.Now()
	payReq, err := zpay32.NewInvoice(
		cfg.ChainParams, paymentHash, creationDate, options...,
	)
	if err != nil {
		return nil, nil, err
	}

	payReqString, err := payReq.Encode(
		zpay32.MessageSigner{
			SignCompact: cfg.NodeSigner.SignDigestCompact,
		},
	)
	if err != nil {
		return nil, nil, err
	}

	newInvoice := &channeldb.Invoice{
		CreationDate:   creationDate,
		Memo:           []byte(invoice.Memo),
		PaymentRequest: []byte(payReqString),
		Terms: channeldb.ContractTerm{
			FinalCltvDelta:  int32(payReq.MinFinalCLTVExpiry()),
			Expiry:          payReq.Expiry(),
			Value:           amtMSat,
			PaymentPreimage: paymentPreimage,
			Features:        invoiceFeatures,
		},
	}

	log.Tracef("[addinvoice] adding new invoice %v",
		newLogClosure(func() string {
			return spew.Sdump(newInvoice)
		}),
	)

	// With all sanity checks passed, write the invoice to the database.
	_, err = cfg.AddInvoice(newInvoice, paymentHash)
	if err != nil {
		return nil, nil, err
	}

	return &paymentHash, newInvoice, nil
}
<commit_msg>rpcserver+invoicesrpc: set payment addr on new invoices<commit_after>package invoicesrpc

import (
	"bytes"
	"context"
	"crypto/rand"
	"errors"
	"fmt"
	"math"
	"time"

	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcutil"
	"github.com/davecgh/go-spew/spew"

	"github.com/lightningnetwork/lnd/channeldb"
	"github.com/lightningnetwork/lnd/lntypes"
	"github.com/lightningnetwork/lnd/lnwire"
	"github.com/lightningnetwork/lnd/netann"
	"github.com/lightningnetwork/lnd/zpay32"
)

// AddInvoiceConfig contains dependencies for invoice creation.
type AddInvoiceConfig struct {
	// AddInvoice is called to add the invoice to the registry.
	AddInvoice func(invoice *channeldb.Invoice, paymentHash lntypes.Hash) (
		uint64, error)

	// IsChannelActive is used to generate valid hop hints.
	IsChannelActive func(chanID lnwire.ChannelID) bool

	// ChainParams are required to properly decode invoice payment requests
	// that are marshalled over rpc.
	ChainParams *chaincfg.Params

	// NodeSigner is an implementation of the MessageSigner implementation
	// that's backed by the identity private key of the running lnd node.
	NodeSigner *netann.NodeSigner

	// MaxPaymentMSat is the maximum allowed payment.
	MaxPaymentMSat lnwire.MilliSatoshi

	// DefaultCLTVExpiry is the default invoice expiry if no values is
	// specified.
	DefaultCLTVExpiry uint32

	// ChanDB is a global boltdb instance which is needed to access the
	// channel graph.
	ChanDB *channeldb.DB
}

// AddInvoiceData contains the required data to create a new invoice.
type AddInvoiceData struct {
	// An optional memo to attach along with the invoice. Used for record
	// keeping purposes for the invoice's creator, and will also be set in
	// the description field of the encoded payment request if the
	// description_hash field is not being used.
	Memo string

	// The preimage which will allow settling an incoming HTLC payable to
	// this preimage. If Preimage is set, Hash should be nil. If both
	// Preimage and Hash are nil, a random preimage is generated.
	Preimage *lntypes.Preimage

	// The hash of the preimage. If Hash is set, Preimage should be nil.
	// This condition indicates that we have a 'hold invoice' for which the
	// htlc will be accepted and held until the preimage becomes known.
	Hash *lntypes.Hash

	// The value of this invoice in millisatoshis.
	Value lnwire.MilliSatoshi

	// Hash (SHA-256) of a description of the payment. Used if the
	// description of payment (memo) is too long to naturally fit within the
	// description field of an encoded payment request.
	DescriptionHash []byte

	// Payment request expiry time in seconds. Default is 3600 (1 hour).
	Expiry int64

	// Fallback on-chain address.
	FallbackAddr string

	// Delta to use for the time-lock of the CLTV extended to the final hop.
	CltvExpiry uint64

	// Whether this invoice should include routing hints for private
	// channels.
	Private bool
}

// AddInvoice attempts to add a new invoice to the invoice database. Any
// duplicated invoices are rejected, therefore all invoices *must* have a
// unique payment preimage.
func AddInvoice(ctx context.Context, cfg *AddInvoiceConfig,
	invoice *AddInvoiceData) (*lntypes.Hash, *channeldb.Invoice, error) {

	var (
		paymentPreimage lntypes.Preimage
		paymentHash     lntypes.Hash
	)

	switch {

	// Only either preimage or hash can be set.
	case invoice.Preimage != nil && invoice.Hash != nil:
		return nil, nil,
			errors.New("preimage and hash both set")

	// Prevent the unknown preimage magic value from being used for a
	// regular invoice. This would cause the invoice the be handled as if it
	// was a hold invoice.
	case invoice.Preimage != nil &&
		*invoice.Preimage == channeldb.UnknownPreimage:

		return nil, nil,
			fmt.Errorf("cannot use all zeroes as a preimage")

	// Prevent the hash of the unknown preimage magic value to be used for a
	// hold invoice. This would make it impossible to settle the invoice,
	// because it would still be interpreted as not having a preimage.
	case invoice.Hash != nil &&
		*invoice.Hash == channeldb.UnknownPreimage.Hash():

		return nil, nil,
			fmt.Errorf("cannot use hash of all zeroes preimage")

	// If no hash or preimage is given, generate a random preimage.
	case invoice.Preimage == nil && invoice.Hash == nil:
		if _, err := rand.Read(paymentPreimage[:]); err != nil {
			return nil, nil, err
		}
		paymentHash = paymentPreimage.Hash()

	// If just a hash is given, we create a hold invoice by setting the
	// preimage to unknown.
	case invoice.Preimage == nil && invoice.Hash != nil:
		paymentPreimage = channeldb.UnknownPreimage
		paymentHash = *invoice.Hash

	// A specific preimage was supplied. Use that for the invoice.
	case invoice.Preimage != nil && invoice.Hash == nil:
		paymentPreimage = *invoice.Preimage
		paymentHash = invoice.Preimage.Hash()
	}

	// The size of the memo, receipt and description hash attached must not
	// exceed the maximum values for either of the fields.
	if len(invoice.Memo) > channeldb.MaxMemoSize {
		return nil, nil, fmt.Errorf("memo too large: %v bytes "+
			"(maxsize=%v)", len(invoice.Memo), channeldb.MaxMemoSize)
	}
	if len(invoice.DescriptionHash) > 0 && len(invoice.DescriptionHash) != 32 {
		return nil, nil, fmt.Errorf("description hash is %v bytes, must be 32",
			len(invoice.DescriptionHash))
	}

	// The value of the invoice must not be negative.
	if invoice.Value < 0 {
		return nil, nil, fmt.Errorf("payments of negative value "+
			"are not allowed, value is %v", invoice.Value)
	}

	amtMSat := invoice.Value

	// The value of the invoice must also not exceed the current soft-limit
	// on the largest payment within the network.
	if amtMSat > cfg.MaxPaymentMSat {
		return nil, nil, fmt.Errorf("payment of %v is too large, max "+
			"payment allowed is %v", invoice.Value,
			cfg.MaxPaymentMSat.ToSatoshis(),
		)
	}

	// We also create an encoded payment request which allows the
	// caller to compactly send the invoice to the payer. We'll create a
	// list of options to be added to the encoded payment request. For now
	// we only support the required fields description/description_hash,
	// expiry, fallback address, and the amount field.
	var options []func(*zpay32.Invoice)

	// We only include the amount in the invoice if it is greater than 0.
	// By not including the amount, we enable the creation of invoices that
	// allow the payee to specify the amount of satoshis they wish to send.
	if amtMSat > 0 {
		options = append(options, zpay32.Amount(amtMSat))
	}

	// If specified, add a fallback address to the payment request.
	if len(invoice.FallbackAddr) > 0 {
		addr, err := btcutil.DecodeAddress(invoice.FallbackAddr,
			cfg.ChainParams)
		if err != nil {
			return nil, nil, fmt.Errorf("invalid fallback address: %v",
				err)
		}
		options = append(options, zpay32.FallbackAddr(addr))
	}

	// If expiry is set, specify it. If it is not provided, no expiry time
	// will be explicitly added to this payment request, which will imply
	// the default 3600 seconds.
	if invoice.Expiry > 0 {

		// We'll ensure that the specified expiry is restricted to sane
		// number of seconds. As a result, we'll reject an invoice with
		// an expiry greater than 1 year.
		maxExpiry := time.Hour * 24 * 365
		expSeconds := invoice.Expiry

		if float64(expSeconds) > maxExpiry.Seconds() {
			return nil, nil, fmt.Errorf("expiry of %v seconds "+
				"greater than max expiry of %v seconds",
				float64(expSeconds), maxExpiry.Seconds())
		}

		expiry := time.Duration(invoice.Expiry) * time.Second
		options = append(options, zpay32.Expiry(expiry))
	}

	// If the description hash is set, then we add it do the list of options.
	// If not, use the memo field as the payment request description.
	if len(invoice.DescriptionHash) > 0 {
		var descHash [32]byte
		copy(descHash[:], invoice.DescriptionHash[:])
		options = append(options, zpay32.DescriptionHash(descHash))
	} else {
		// Use the memo field as the description. If this is not set
		// this will just be an empty string.
		options = append(options, zpay32.Description(invoice.Memo))
	}

	// We'll use our current default CLTV value unless one was specified as
	// an option on the command line when creating an invoice.
	switch {
	case invoice.CltvExpiry > math.MaxUint16:
		return nil, nil, fmt.Errorf("CLTV delta of %v is too large, max "+
			"accepted is: %v", invoice.CltvExpiry, math.MaxUint16)
	case invoice.CltvExpiry != 0:
		options = append(options,
			zpay32.CLTVExpiry(invoice.CltvExpiry))
	default:
		// TODO(roasbeef): assumes set delta between versions
		defaultDelta := cfg.DefaultCLTVExpiry
		options = append(options, zpay32.CLTVExpiry(uint64(defaultDelta)))
	}

	// If we were requested to include routing hints in the invoice, then
	// we'll fetch all of our available private channels and create routing
	// hints for them.
	if invoice.Private {
		openChannels, err := cfg.ChanDB.FetchAllChannels()
		if err != nil {
			return nil, nil, fmt.Errorf("could not fetch all channels")
		}

		graph := cfg.ChanDB.ChannelGraph()

		numHints := 0
		for _, channel := range openChannels {
			// We'll restrict the number of individual route hints
			// to 20 to avoid creating overly large invoices.
			if numHints >= 20 {
				break
			}

			// Since we're only interested in our private channels,
			// we'll skip public ones.
			isPublic := channel.ChannelFlags&lnwire.FFAnnounceChannel != 0
			if isPublic {
				continue
			}

			// Make sure the counterparty has enough balance in the
			// channel for our amount. We do this in order to reduce
			// payment errors when attempting to use this channel
			// as a hint.
			chanPoint := lnwire.NewChanIDFromOutPoint(
				&channel.FundingOutpoint,
			)
			if amtMSat >= channel.LocalCommitment.RemoteBalance {
				log.Debugf("Skipping channel %v due to "+
					"not having enough remote balance",
					chanPoint)
				continue
			}

			// Make sure the channel is active.
			if !cfg.IsChannelActive(chanPoint) {
				log.Debugf("Skipping channel %v due to not "+
					"being eligible to forward payments",
					chanPoint)
				continue
			}

			// To ensure we don't leak unadvertised nodes, we'll
			// make sure our counterparty is publicly advertised
			// within the network. Otherwise, we'll end up leaking
			// information about nodes that intend to stay
			// unadvertised, like in the case of a node only having
			// private channels.
			var remotePub [33]byte
			copy(remotePub[:], channel.IdentityPub.SerializeCompressed())
			isRemoteNodePublic, err := graph.IsPublicNode(remotePub)
			if err != nil {
				log.Errorf("Unable to determine if node %x "+
					"is advertised: %v", remotePub, err)
				continue
			}

			if !isRemoteNodePublic {
				log.Debugf("Skipping channel %v due to "+
					"counterparty %x being unadvertised",
					chanPoint, remotePub)
				continue
			}

			// Fetch the policies for each end of the channel.
			chanID := channel.ShortChanID().ToUint64()
			info, p1, p2, err := graph.FetchChannelEdgesByID(chanID)
			if err != nil {
				log.Errorf("Unable to fetch the routing "+
					"policies for the edges of the channel "+
					"%v: %v", chanPoint, err)
				continue
			}

			// Now, we'll need to determine which is the correct
			// policy for HTLCs being sent from the remote node.
			var remotePolicy *channeldb.ChannelEdgePolicy
			if bytes.Equal(remotePub[:], info.NodeKey1Bytes[:]) {
				remotePolicy = p1
			} else {
				remotePolicy = p2
			}

			// If for some reason we don't yet have the edge for
			// the remote party, then we'll just skip adding this
			// channel as a routing hint.
			if remotePolicy == nil {
				continue
			}

			// Finally, create the routing hint for this channel and
			// add it to our list of route hints.
			hint := zpay32.HopHint{
				NodeID:      channel.IdentityPub,
				ChannelID:   chanID,
				FeeBaseMSat: uint32(remotePolicy.FeeBaseMSat),
				FeeProportionalMillionths: uint32(
					remotePolicy.FeeProportionalMillionths,
				),
				CLTVExpiryDelta: remotePolicy.TimeLockDelta,
			}

			// Include the route hint in our set of options that
			// will be used when creating the invoice.
			routeHint := []zpay32.HopHint{hint}
			options = append(options, zpay32.RouteHint(routeHint))

			numHints++
		}

	}

	// Set a blank feature vector, as our invoice generation forbids nil
	// features.
	invoiceFeatures := lnwire.NewFeatureVector(
		lnwire.NewRawFeatureVector(), lnwire.Features,
	)
	options = append(options, zpay32.Features(invoiceFeatures))

	// Generate and set a random payment address for this invoice. If the
	// sender understands payment addresses, this can be used to avoid
	// intermediaries probing the receiver.
	var paymentAddr [32]byte
	if _, err := rand.Read(paymentAddr[:]); err != nil {
		return nil, nil, err
	}
	options = append(options, zpay32.PaymentAddr(paymentAddr))

	// Create and encode the payment request as a bech32 (zpay32) string.
	creationDate := time.Now()
	payReq, err := zpay32.NewInvoice(
		cfg.ChainParams, paymentHash, creationDate, options...,
	)
	if err != nil {
		return nil, nil, err
	}

	payReqString, err := payReq.Encode(
		zpay32.MessageSigner{
			SignCompact: cfg.NodeSigner.SignDigestCompact,
		},
	)
	if err != nil {
		return nil, nil, err
	}

	newInvoice := &channeldb.Invoice{
		CreationDate:   creationDate,
		Memo:           []byte(invoice.Memo),
		PaymentRequest: []byte(payReqString),
		Terms: channeldb.ContractTerm{
			FinalCltvDelta:  int32(payReq.MinFinalCLTVExpiry()),
			Expiry:          payReq.Expiry(),
			Value:           amtMSat,
			PaymentPreimage: paymentPreimage,
			PaymentAddr:     paymentAddr,
			Features:        invoiceFeatures,
		},
	}

	log.Tracef("[addinvoice] adding new invoice %v",
		newLogClosure(func() string {
			return spew.Sdump(newInvoice)
		}),
	)

	// With all sanity checks passed, write the invoice to the database.
	_, err = cfg.AddInvoice(newInvoice, paymentHash)
	if err != nil {
		return nil, nil, err
	}

	return &paymentHash, newInvoice, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"net/http"
	_ "net/http/pprof"
	"regexp"
	"strconv"
	"sync"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/log"

	consul_api "github.com/hashicorp/consul/api"
	consul "github.com/hashicorp/consul/consul/structs"
)

const (
	namespace = "consul"
)

var (
	serviceLabelNames = []string{"service", "node"}
	memberLabelNames  = []string{"member"}
)

// Exporter collects Consul stats from the given server and exports them using
// the prometheus metrics package.
type Exporter struct {
	URI   string
	mutex sync.RWMutex

	up, clusterServers                                prometheus.Gauge
	nodeCount, serviceCount                           prometheus.Counter
	serviceNodesTotal, serviceNodesHealthy, keyValues *prometheus.GaugeVec
	client                                            *consul_api.Client
	kvPrefix					  string
	kvFilter					  *regexp.Regexp
}

// NewExporter returns an initialized Exporter.
func NewExporter(uri string, kvPrefix string, kvFilter string) *Exporter {
	// Set up our Consul client connection.
	consul_client, _ := consul_api.NewClient(&consul_api.Config{
		Address: uri,
	})

	// Init our exporter.
	return &Exporter{
		URI: uri,
		up: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "up",
			Help:      "Was the last query of Consul successful.",
		}),

		clusterServers: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "raft_peers",
			Help:      "How many peers (servers) are in the Raft cluster.",
		}),

		nodeCount: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "serf_lan_members",
			Help:      "How many members are in the cluster.",
		}),

		serviceCount: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "catalog_services",
			Help:      "How many services are in the cluster.",
		}),

		serviceNodesTotal: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "catalog_service_nodes",
				Help:      "Number of nodes currently registered for this service.",
			},
			[]string{"service"},
		),

		serviceNodesHealthy: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "catalog_service_node_healthy",
				Help:      "Is this service healthy on this node?",
			},
			[]string{"service", "node"},
		),

		keyValues: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "catalog_kv",
				Help:      "The values for selected keys in Consul's key/value catalog. Keys with non-numeric values are omitted.",
			},
			[]string{"key"},
		),

		client: consul_client,
		kvPrefix: kvPrefix,
		kvFilter: regexp.MustCompile(kvFilter),
	}
}

// Describe describes all the metrics ever exported by the Consul exporter. It
// implements prometheus.Collector.
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	ch <- e.up.Desc()
	ch <- e.nodeCount.Desc()
	ch <- e.serviceCount.Desc()
	ch <- e.clusterServers.Desc()

	e.serviceNodesTotal.Describe(ch)
	e.serviceNodesHealthy.Describe(ch)
	e.keyValues.Describe(ch)
}

// Collect fetches the stats from configured Consul location and delivers them
// as Prometheus metrics. It implements prometheus.Collector.
func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	services := make(chan []*consul_api.ServiceEntry)

	go e.queryClient(services)

	e.mutex.Lock() // To protect metrics from concurrent collects.
	defer e.mutex.Unlock()

	// Reset metrics.
	e.serviceNodesTotal.Reset()
	e.serviceNodesHealthy.Reset()

	e.setMetrics(services)

	ch <- e.up
	ch <- e.clusterServers
	ch <- e.nodeCount
	ch <- e.serviceCount

	e.serviceNodesTotal.Collect(ch)
	e.serviceNodesHealthy.Collect(ch)

	e.keyValues.Reset()
	e.setKeyValues()
	e.keyValues.Collect(ch)
}

func (e *Exporter) queryClient(services chan<- []*consul_api.ServiceEntry) {
	defer close(services)

	// How many peers are in the Consul cluster?
	peers, err := e.client.Status().Peers()

	if err != nil {
		e.up.Set(0)
		log.Errorf("Query error is %v", err)
		return
	}

	// We'll use peers to decide that we're up.
	e.up.Set(1)
	e.clusterServers.Set(float64(len(peers)))

	// How many nodes are registered?
	nodes, _, err := e.client.Catalog().Nodes(&consul_api.QueryOptions{})

	if err != nil {
		// FIXME: How should we handle a partial failure like this?
	} else {
		e.nodeCount.Set(float64(len(nodes)))
	}

	// Query for the full list of services.
	serviceNames, _, err := e.client.Catalog().Services(&consul_api.QueryOptions{})
	e.serviceCount.Set(float64(len(serviceNames)))

	if err != nil {
		// FIXME: How should we handle a partial failure like this?
		return
	}

	e.serviceCount.Set(float64(len(serviceNames)))

	for s := range serviceNames {
		s_entries, _, err := e.client.Health().Service(s, "", false, &consul_api.QueryOptions{})

		if err != nil {
			log.Errorf("Failed to query service health: %v", err)
			continue
		}

		services <- s_entries
	}
}

func (e *Exporter) setMetrics(services <-chan []*consul_api.ServiceEntry) {
	// Each service will be an array of ServiceEntry structs.
	for service := range services {
		if len(service) == 0 {
			// Not sure this should ever happen, but catch it just in case...
			continue
		}

		// We should have one ServiceEntry per node, so use that for total nodes.
		e.serviceNodesTotal.WithLabelValues(service[0].Service.Service).Set(float64(len(service)))

		for _, entry := range service {
			// We have a Node, a Service, and one or more Checks. Our
			// service-node combo is passing if all checks have a `status`
			// of "passing."

			passing := 1

			for _, hc := range entry.Checks {
				if hc.Status != consul.HealthPassing {
					passing = 0
					break
				}
			}

			log.Infof("%v/%v status is %v", entry.Service.Service, entry.Node.Node, passing)

			e.serviceNodesHealthy.WithLabelValues(entry.Service.Service, entry.Node.Node).Set(float64(passing))
		}
	}
}

func (e *Exporter) setKeyValues() {
	kv := e.client.KV()

	pairs, _, err := kv.List(e.kvPrefix, &consul_api.QueryOptions{})
	if err != nil {
		log.Errorf("Error fetching key/values: %s", err)
	} else {
		for _, pair := range pairs {
			if e.kvFilter.MatchString(pair.Key) {
				val, err := strconv.ParseFloat(string(pair.Value), 64)
				if err == nil {
					e.keyValues.WithLabelValues(pair.Key).Set(val)
				}
			}
		}
	}
}

func main() {
	var (
		listenAddress = flag.String("web.listen-address", ":9107", "Address to listen on for web interface and telemetry.")
		metricsPath   = flag.String("web.telemetry-path", "/metrics", "Path under which to expose metrics.")
		consulServer  = flag.String("consul.server", "localhost:8500", "HTTP API address of a Consul server or agent.")
		kvPrefix      = flag.String("kv.prefix", "none", "Prefix from which to expose key/value pairs.")
		kvFilter      = flag.String("kv.filter", ".*", "Regex that determines which keys to expose.")
	)
	flag.Parse()

	exporter := NewExporter(*consulServer, *kvPrefix, *kvFilter)
	prometheus.MustRegister(exporter)

	log.Infof("Starting Server: %s", *listenAddress)
	http.Handle(*metricsPath, prometheus.Handler())
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`
             <title>Consul Exporter</title>

Consul Exporter

<a href='` + *metricsPath + `'>Metrics

             `))
	})
	log.Fatal(http.ListenAndServe(*listenAddress, nil))
}
<commit_msg>Simplify control structure by returning on error<commit_after>package main

import (
	"flag"
	"net/http"
	_ "net/http/pprof"
	"regexp"
	"strconv"
	"sync"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/log"

	consul_api "github.com/hashicorp/consul/api"
	consul "github.com/hashicorp/consul/consul/structs"
)

const (
	namespace = "consul"
)

var (
	serviceLabelNames = []string{"service", "node"}
	memberLabelNames  = []string{"member"}
)

// Exporter collects Consul stats from the given server and exports them using
// the prometheus metrics package.
type Exporter struct {
	URI   string
	mutex sync.RWMutex

	up, clusterServers                                prometheus.Gauge
	nodeCount, serviceCount                           prometheus.Counter
	serviceNodesTotal, serviceNodesHealthy, keyValues *prometheus.GaugeVec
	client                                            *consul_api.Client
	kvPrefix					  string
	kvFilter					  *regexp.Regexp
}

// NewExporter returns an initialized Exporter.
func NewExporter(uri string, kvPrefix string, kvFilter string) *Exporter {
	// Set up our Consul client connection.
	consul_client, _ := consul_api.NewClient(&consul_api.Config{
		Address: uri,
	})

	// Init our exporter.
	return &Exporter{
		URI: uri,
		up: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "up",
			Help:      "Was the last query of Consul successful.",
		}),

		clusterServers: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "raft_peers",
			Help:      "How many peers (servers) are in the Raft cluster.",
		}),

		nodeCount: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "serf_lan_members",
			Help:      "How many members are in the cluster.",
		}),

		serviceCount: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: namespace,
			Name:      "catalog_services",
			Help:      "How many services are in the cluster.",
		}),

		serviceNodesTotal: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "catalog_service_nodes",
				Help:      "Number of nodes currently registered for this service.",
			},
			[]string{"service"},
		),

		serviceNodesHealthy: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "catalog_service_node_healthy",
				Help:      "Is this service healthy on this node?",
			},
			[]string{"service", "node"},
		),

		keyValues: prometheus.NewGaugeVec(
			prometheus.GaugeOpts{
				Namespace: namespace,
				Name:      "catalog_kv",
				Help:      "The values for selected keys in Consul's key/value catalog. Keys with non-numeric values are omitted.",
			},
			[]string{"key"},
		),

		client: consul_client,
		kvPrefix: kvPrefix,
		kvFilter: regexp.MustCompile(kvFilter),
	}
}

// Describe describes all the metrics ever exported by the Consul exporter. It
// implements prometheus.Collector.
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	ch <- e.up.Desc()
	ch <- e.nodeCount.Desc()
	ch <- e.serviceCount.Desc()
	ch <- e.clusterServers.Desc()

	e.serviceNodesTotal.Describe(ch)
	e.serviceNodesHealthy.Describe(ch)
	e.keyValues.Describe(ch)
}

// Collect fetches the stats from configured Consul location and delivers them
// as Prometheus metrics. It implements prometheus.Collector.
func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	services := make(chan []*consul_api.ServiceEntry)

	go e.queryClient(services)

	e.mutex.Lock() // To protect metrics from concurrent collects.
	defer e.mutex.Unlock()

	// Reset metrics.
	e.serviceNodesTotal.Reset()
	e.serviceNodesHealthy.Reset()

	e.setMetrics(services)

	ch <- e.up
	ch <- e.clusterServers
	ch <- e.nodeCount
	ch <- e.serviceCount

	e.serviceNodesTotal.Collect(ch)
	e.serviceNodesHealthy.Collect(ch)

	e.keyValues.Reset()
	e.setKeyValues()
	e.keyValues.Collect(ch)
}

func (e *Exporter) queryClient(services chan<- []*consul_api.ServiceEntry) {
	defer close(services)

	// How many peers are in the Consul cluster?
	peers, err := e.client.Status().Peers()

	if err != nil {
		e.up.Set(0)
		log.Errorf("Query error is %v", err)
		return
	}

	// We'll use peers to decide that we're up.
	e.up.Set(1)
	e.clusterServers.Set(float64(len(peers)))

	// How many nodes are registered?
	nodes, _, err := e.client.Catalog().Nodes(&consul_api.QueryOptions{})

	if err != nil {
		// FIXME: How should we handle a partial failure like this?
	} else {
		e.nodeCount.Set(float64(len(nodes)))
	}

	// Query for the full list of services.
	serviceNames, _, err := e.client.Catalog().Services(&consul_api.QueryOptions{})
	e.serviceCount.Set(float64(len(serviceNames)))

	if err != nil {
		// FIXME: How should we handle a partial failure like this?
		return
	}

	e.serviceCount.Set(float64(len(serviceNames)))

	for s := range serviceNames {
		s_entries, _, err := e.client.Health().Service(s, "", false, &consul_api.QueryOptions{})

		if err != nil {
			log.Errorf("Failed to query service health: %v", err)
			continue
		}

		services <- s_entries
	}
}

func (e *Exporter) setMetrics(services <-chan []*consul_api.ServiceEntry) {
	// Each service will be an array of ServiceEntry structs.
	for service := range services {
		if len(service) == 0 {
			// Not sure this should ever happen, but catch it just in case...
			continue
		}

		// We should have one ServiceEntry per node, so use that for total nodes.
		e.serviceNodesTotal.WithLabelValues(service[0].Service.Service).Set(float64(len(service)))

		for _, entry := range service {
			// We have a Node, a Service, and one or more Checks. Our
			// service-node combo is passing if all checks have a `status`
			// of "passing."

			passing := 1

			for _, hc := range entry.Checks {
				if hc.Status != consul.HealthPassing {
					passing = 0
					break
				}
			}

			log.Infof("%v/%v status is %v", entry.Service.Service, entry.Node.Node, passing)

			e.serviceNodesHealthy.WithLabelValues(entry.Service.Service, entry.Node.Node).Set(float64(passing))
		}
	}
}

func (e *Exporter) setKeyValues() {
	kv := e.client.KV()

	pairs, _, err := kv.List(e.kvPrefix, &consul_api.QueryOptions{})
	if err != nil {
		log.Errorf("Error fetching key/values: %s", err)
		return
	}

	for _, pair := range pairs {
		if e.kvFilter.MatchString(pair.Key) {
			val, err := strconv.ParseFloat(string(pair.Value), 64)
			if err == nil {
				e.keyValues.WithLabelValues(pair.Key).Set(val)
			}
		}
	}
}

func main() {
	var (
		listenAddress = flag.String("web.listen-address", ":9107", "Address to listen on for web interface and telemetry.")
		metricsPath   = flag.String("web.telemetry-path", "/metrics", "Path under which to expose metrics.")
		consulServer  = flag.String("consul.server", "localhost:8500", "HTTP API address of a Consul server or agent.")
		kvPrefix      = flag.String("kv.prefix", "none", "Prefix from which to expose key/value pairs.")
		kvFilter      = flag.String("kv.filter", ".*", "Regex that determines which keys to expose.")
	)
	flag.Parse()

	exporter := NewExporter(*consulServer, *kvPrefix, *kvFilter)
	prometheus.MustRegister(exporter)

	log.Infof("Starting Server: %s", *listenAddress)
	http.Handle(*metricsPath, prometheus.Handler())
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`
             <title>Consul Exporter</title>

Consul Exporter

<a href='` + *metricsPath + `'>Metrics

             `))
	})
	log.Fatal(http.ListenAndServe(*listenAddress, nil))
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"gopkg.in/telegram-bot-api.v4"
	"time"
	//"os/user"
)

func regpi(msg *tgbotapi.Message, update tgbotapi.Update) {
	var reply tgbotapi.MessageConfig
	var user User

	userId := msg.From.ID
	groupId := msg.Chat.ID

	gdb.Where("userId = ? AND groupId = ?", userId, groupId).First(&user)

	if len(msg.From.UserName) == 0 {
		reply = tgbotapi.NewMessage(msg.Chat.ID, "A girl has no name.")
	} else if user.Id == 0 {
		user.Username = "@" + msg.From.UserName
		user.UserId = msg.From.ID
		user.GroupID = int(groupId)
		user.Score = 0
		user.Usernick = fmt.Sprintf("%s %s", msg.From.FirstName, msg.From.LastName)
		user.Quota = 6
		gdb.Create(&user)
		gdb.Model(&user).Update(User{Quota: 6})
		reply = tgbotapi.NewMessage(msg.Chat.ID, fmt.Sprintf("Ты регнулся, [%s](tg://user?id=%d)\n", user.Usernick, user.UserId))
		reply.ParseMode = tgbotapi.ModeMarkdown
	} else {
		reply = tgbotapi.NewMessage(msg.Chat.ID, fmt.Sprint("Эй, ты уже в игре!"))
	}

	reply.ReplyToMessageID = update.Message.MessageID
	bot.Send(reply)
}

func showpid(msg *tgbotapi.Message) {
	var users []User

	groupId := msg.Chat.ID

	gdb.Where("groupId = ?", groupId).Find(&users)

	if len(users) != 0 {
		output := "Кандидаты в пидоры дня:\n"
		for _, i := range users {
			if len(i.Usernick) > 0 {
				output += fmt.Sprintf("[%s](tg://user?id=%d)\n", i.Usernick, i.UserId)
			} else {
				output += fmt.Sprintf("[%s](tg://user?id=%d)\n", i.Username[1:], i.UserId)
			}
		}
		output += "Хочешь себя увидеть тут?\nЖми /regpi"
		reply := tgbotapi.NewMessage(msg.Chat.ID, output)
		reply.ParseMode = tgbotapi.ModeMarkdown
		bot.Send(reply)
	} else {
		output := "Пидоров нет! Будь первым! Жми /regpi"
		bot.Send(tgbotapi.NewMessage(msg.Chat.ID, output))
	}
}

func pidorStat(msg *tgbotapi.Message) {
	titles := map[int]string{
		1 : "Пидоратор",
		2 : "Пидороль",
		3 : "Герцопидор",
		4 : "Пиркиз",
		5 : "Пидорон",
	}

	var users []User
	var reply tgbotapi.MessageConfig
	var flag bool
	var currentUserName string

	groupId := msg.Chat.ID
	counter := 0
	var titlesCounter int

	gdb.Where("groupId = ?", groupId).Order("score desc").Find(&users)

	output := "Статистика (первые 5):\n"
	titlesCounter = 1
	for _, i := range users {
		if i.Score != 0 {
			if len(i.Usernick) > 0 {
				currentUserName = i.Usernick
			} else {
				currentUserName = i.Username[1:]
			}

			output += fmt.Sprintf("[%s](tg://user?id=%d) - %d (%s)\n", currentUserName, i.UserId, i.Score, titles[titlesCounter])
			titlesCounter++
			flag = true

			if counter > 5 {
				break
			} else {
				counter++
			}
		}
	}

	if flag {
		reply = tgbotapi.NewMessage(msg.Chat.ID, output)
		reply.ParseMode = tgbotapi.ModeMarkdown
	} else {
		reply = tgbotapi.NewMessage(msg.Chat.ID, "Пидор дня еще ни разу не был выбран! Жми /pidor")
	}

	bot.Send(reply)
}

func startQuiz(msg *tgbotapi.Message) {
	firstPhrases := []string {
		"Инициализирую поиск пидора дня...",
		"Внимание, ищу пидора!",
		"Ну-ка дай-ка...",
		"Такс, кто тут у нас мало каши ел?",
		"Инициализация.Поиск.",
	}

	secondPhrases := []string {
		"Кажется я что-то вижу!",
		"Не может быть!",
		"Пожалуй препроверю...",
		"Найден!",
		"Прям по Бабичу!",
		"Как предсказал Великий Мейстер...",
	}

	var reply tgbotapi.MessageConfig
	var theUser User
	var users []User
	var randomUser int
	var currentUserName string
	var winner User
	var winnerScore int
	var available Available

	groupId := msg.Chat.ID

	gdb.Where("groupId = ?", groupId).Find(&users)
	gdb.Where("groupId = ?", groupId).First(&available)

	rowsCounted := len(users)
	if rowsCounted == 0 {
		reply = tgbotapi.NewMessage(msg.Chat.ID, "Нет участников! Жми /regpi")
		bot.Send(reply)
	} else {
		if available.Flag {
			lenOfCurrentUsers := len(users)
			if lenOfCurrentUsers == 1 {
				randomUser = 0
			} else {
				randomUser = random(0, lenOfCurrentUsers - 1)
			}

			gdb.Where("id = ?", users[randomUser].Id).First(&winner)

			reply = tgbotapi.NewMessage(msg.Chat.ID, firstPhrases[random(0, len(firstPhrases) - 1)])
			bot.Send(reply)
			time.Sleep(time.Second * 2)
			reply = tgbotapi.NewMessage(msg.Chat.ID, secondPhrases[random(0, len(secondPhrases) - 1)])
			bot.Send(reply)
			gdb.Where("id = ? and groupId = ?", theUser, groupId).First(&winner)
			winnerScore = winner.Score + 1
			gdb.Model(&users).Where("id = ?", winner.Id).UpdateColumn("score", winnerScore)
			time.Sleep(time.Second * 2)
			if len(winner.Usernick) > 0 {
				currentUserName = winner.Usernick
			} else {
				currentUserName = winner.Username
			}
			reply = tgbotapi.NewMessage(msg.Chat.ID, fmt.Sprintf("Ага! 🎉🎉🎉 Сегодня пидор - [%s](tg://user?id=%d)", currentUserName, winner.UserId))
			reply.ParseMode = tgbotapi.ModeMarkdown
			bot.Send(reply)
			gdb.Model(&available).Where("groupId = ?", groupId).Update("flag", false)
			gdb.Model(&available).Where("groupId = ?", groupId).Update("userId", winner.Id)
		} else {
			var currentUser User
			gdb.Where("id = ?", available.UserId).First(&currentUser)
			if len(currentUser.Usernick) > 0 {
				currentUserName = currentUser.Usernick
			} else {
				currentUserName = currentUser.Username[1:]
			}
			reply = tgbotapi.NewMessage(msg.Chat.ID, fmt.Sprintf("🎉Сегодня у нас уже есть победитель - [%s](tg://user?id=%d)🎉", currentUserName, currentUser.UserId))
			reply.ParseMode = tgbotapi.ModeMarkdown
			bot.Send(reply)
		}
	}
}

func kekogen(msg *tgbotapi.Message) {
	var reply tgbotapi.MessageConfig
	var user User

	userId := msg.From.ID
	groupId := msg.Chat.ID

	gdb.Where("userId = ? and groupId = ?", userId, groupId).First(&user)
	currentQuota := user.Quota

	if user.Id > 0 {
		if currentQuota > 1 {
			vowels := []string {
				"а", "о", "и", "е", "у",  "я",
			}
			consonants := []string {
				"в", "д", "к", "л", "м", "н", "п", "р", "с", "т", "ф", "х", "ш",
			}
			result := "кек"

			for x:= 0; x < 5; x++ {
				if x % 2 == 0 {
					result += vowels[random(0, len(vowels) - 1)]
				} else {
					result += consonants[random(0, len(consonants) - 1)]
				}
			}

			gdb.Model(&user).Update(User{Quota: currentQuota - 1})
			reply = tgbotapi.NewMessage(msg.Chat.ID, fmt.Sprintf(result))
		} else {
			reply = tgbotapi.NewMessage(msg.Chat.ID, fmt.Sprintf("Твои кеки на сегодня кончились!"))
		}
	} else {
		reply = tgbotapi.NewMessage(msg.Chat.ID, fmt.Sprintf("Жми /regpi чтобы зарегистрироваться и получить свои кеки!"))
	}

	bot.Send(reply)
}<commit_msg>fix stat<commit_after>package main

import (
	"fmt"
	"gopkg.in/telegram-bot-api.v4"
	"time"
	//"os/user"
)

func regpi(msg *tgbotapi.Message, update tgbotapi.Update) {
	var reply tgbotapi.MessageConfig
	var user User

	userId := msg.From.ID
	groupId := msg.Chat.ID

	gdb.Where("userId = ? AND groupId = ?", userId, groupId).First(&user)

	if len(msg.From.UserName) == 0 {
		reply = tgbotapi.NewMessage(msg.Chat.ID, "A girl has no name.")
	} else if user.Id == 0 {
		user.Username = "@" + msg.From.UserName
		user.UserId = msg.From.ID
		user.GroupID = int(groupId)
		user.Score = 0
		user.Usernick = fmt.Sprintf("%s %s", msg.From.FirstName, msg.From.LastName)
		user.Quota = 6
		gdb.Create(&user)
		gdb.Model(&user).Update(User{Quota: 6})
		reply = tgbotapi.NewMessage(msg.Chat.ID, fmt.Sprintf("Ты регнулся, [%s](tg://user?id=%d)\n", user.Usernick, user.UserId))
		reply.ParseMode = tgbotapi.ModeMarkdown
	} else {
		reply = tgbotapi.NewMessage(msg.Chat.ID, fmt.Sprint("Эй, ты уже в игре!"))
	}

	reply.ReplyToMessageID = update.Message.MessageID
	bot.Send(reply)
}

func showpid(msg *tgbotapi.Message) {
	var users []User

	groupId := msg.Chat.ID

	gdb.Where("groupId = ?", groupId).Find(&users)

	if len(users) != 0 {
		output := "Кандидаты в пидоры дня:\n"
		for _, i := range users {
			if len(i.Usernick) > 0 {
				output += fmt.Sprintf("[%s](tg://user?id=%d)\n", i.Usernick, i.UserId)
			} else {
				output += fmt.Sprintf("[%s](tg://user?id=%d)\n", i.Username[1:], i.UserId)
			}
		}
		output += "Хочешь себя увидеть тут?\nЖми /regpi"
		reply := tgbotapi.NewMessage(msg.Chat.ID, output)
		reply.ParseMode = tgbotapi.ModeMarkdown
		bot.Send(reply)
	} else {
		output := "Пидоров нет! Будь первым! Жми /regpi"
		bot.Send(tgbotapi.NewMessage(msg.Chat.ID, output))
	}
}

func pidorStat(msg *tgbotapi.Message) {
	titles := map[int]string{
		1 : "Пидоратор",
		2 : "Пидороль",
		3 : "Герцопидор",
		4 : "Пиркиз",
		5 : "Пидорон",
	}

	var users []User
	var reply tgbotapi.MessageConfig
	var flag bool
	var currentUserName string

	groupId := msg.Chat.ID
	counter := 0
	var titlesCounter int

	gdb.Where("groupId = ?", groupId).Order("score desc").Find(&users)

	output := "Статистика (первые 5):\n"
	titlesCounter = 1
	for _, i := range users {
		if i.Score != 0 {
			if len(i.Usernick) > 0 {
				currentUserName = i.Usernick
			} else {
				currentUserName = i.Username[1:]
			}

			output += fmt.Sprintf("[%s](tg://user?id=%d) - %d (%s)\n", currentUserName, i.UserId, i.Score, titles[titlesCounter])
			titlesCounter++
			flag = true

			if counter < 5 {
				break
			} else {
				counter++
			}
		}
	}

	if flag {
		reply = tgbotapi.NewMessage(msg.Chat.ID, output)
		reply.ParseMode = tgbotapi.ModeMarkdown
	} else {
		reply = tgbotapi.NewMessage(msg.Chat.ID, "Пидор дня еще ни разу не был выбран! Жми /pidor")
	}

	bot.Send(reply)
}

func startQuiz(msg *tgbotapi.Message) {
	firstPhrases := []string {
		"Инициализирую поиск пидора дня...",
		"Внимание, ищу пидора!",
		"Ну-ка дай-ка...",
		"Такс, кто тут у нас мало каши ел?",
		"Инициализация.Поиск.",
	}

	secondPhrases := []string {
		"Кажется я что-то вижу!",
		"Не может быть!",
		"Пожалуй препроверю...",
		"Найден!",
		"Прям по Бабичу!",
		"Как предсказал Великий Мейстер...",
	}

	var reply tgbotapi.MessageConfig
	var theUser User
	var users []User
	var randomUser int
	var currentUserName string
	var winner User
	var winnerScore int
	var available Available

	groupId := msg.Chat.ID

	gdb.Where("groupId = ?", groupId).Find(&users)
	gdb.Where("groupId = ?", groupId).First(&available)

	rowsCounted := len(users)
	if rowsCounted == 0 {
		reply = tgbotapi.NewMessage(msg.Chat.ID, "Нет участников! Жми /regpi")
		bot.Send(reply)
	} else {
		if available.Flag {
			lenOfCurrentUsers := len(users)
			if lenOfCurrentUsers == 1 {
				randomUser = 0
			} else {
				randomUser = random(0, lenOfCurrentUsers - 1)
			}

			gdb.Where("id = ?", users[randomUser].Id).First(&winner)

			reply = tgbotapi.NewMessage(msg.Chat.ID, firstPhrases[random(0, len(firstPhrases) - 1)])
			bot.Send(reply)
			time.Sleep(time.Second * 2)
			reply = tgbotapi.NewMessage(msg.Chat.ID, secondPhrases[random(0, len(secondPhrases) - 1)])
			bot.Send(reply)
			gdb.Where("id = ? and groupId = ?", theUser, groupId).First(&winner)
			winnerScore = winner.Score + 1
			gdb.Model(&users).Where("id = ?", winner.Id).UpdateColumn("score", winnerScore)
			time.Sleep(time.Second * 2)
			if len(winner.Usernick) > 0 {
				currentUserName = winner.Usernick
			} else {
				currentUserName = winner.Username
			}
			reply = tgbotapi.NewMessage(msg.Chat.ID, fmt.Sprintf("Ага! 🎉🎉🎉 Сегодня пидор - [%s](tg://user?id=%d)", currentUserName, winner.UserId))
			reply.ParseMode = tgbotapi.ModeMarkdown
			bot.Send(reply)
			gdb.Model(&available).Where("groupId = ?", groupId).Update("flag", false)
			gdb.Model(&available).Where("groupId = ?", groupId).Update("userId", winner.Id)
		} else {
			var currentUser User
			gdb.Where("id = ?", available.UserId).First(&currentUser)
			if len(currentUser.Usernick) > 0 {
				currentUserName = currentUser.Usernick
			} else {
				currentUserName = currentUser.Username[1:]
			}
			reply = tgbotapi.NewMessage(msg.Chat.ID, fmt.Sprintf("🎉Сегодня у нас уже есть победитель - [%s](tg://user?id=%d)🎉", currentUserName, currentUser.UserId))
			reply.ParseMode = tgbotapi.ModeMarkdown
			bot.Send(reply)
		}
	}
}

func kekogen(msg *tgbotapi.Message) {
	var reply tgbotapi.MessageConfig
	var user User

	userId := msg.From.ID
	groupId := msg.Chat.ID

	gdb.Where("userId = ? and groupId = ?", userId, groupId).First(&user)
	currentQuota := user.Quota

	if user.Id > 0 {
		if currentQuota > 1 {
			vowels := []string {
				"а", "о", "и", "е", "у",  "я",
			}
			consonants := []string {
				"в", "д", "к", "л", "м", "н", "п", "р", "с", "т", "ф", "х", "ш",
			}
			result := "кек"

			for x:= 0; x < 5; x++ {
				if x % 2 == 0 {
					result += vowels[random(0, len(vowels) - 1)]
				} else {
					result += consonants[random(0, len(consonants) - 1)]
				}
			}

			gdb.Model(&user).Update(User{Quota: currentQuota - 1})
			reply = tgbotapi.NewMessage(msg.Chat.ID, fmt.Sprintf(result))
		} else {
			reply = tgbotapi.NewMessage(msg.Chat.ID, fmt.Sprintf("Твои кеки на сегодня кончились!"))
		}
	} else {
		reply = tgbotapi.NewMessage(msg.Chat.ID, fmt.Sprintf("Жми /regpi чтобы зарегистрироваться и получить свои кеки!"))
	}

	bot.Send(reply)
}<|endoftext|>


<commit_before>package report

import(
	"strings"

	"golang.org/x/net/context"
	"google.golang.org/appengine/user"

	"github.com/skypies/flightdb/metar"
)

// TODO: just embed / extract this data from the regular context
type ReportingContext struct {
	context.Context
	metar.Archive
	UserEmail       string
}

var(
	// Lower case everything
	ACLFOIA = []string{
		"adam@worrall.cc",
		"raymonde.guindon@gmail.com",
		"rguindon@alumni.stanford.edu",
		"meekgee@gmail.com",
		"nancyjordan650@gmail.com",
		"borg@smwlaw.com",
		"matt@classalameda.com",
		"jnelson@wiai.com",
		"robert.holbrook@gmail.com",
		"jtsunnyvaleair1@gmail.com",
		"mthurd2003@gmail.com",
		"fabrice.beretta@gmail.com",
	}
)

func (r *Report)setupReportingContext(ctx context.Context) error {
	r.ReportingContext.Context = ctx

	metar,err := metar.LookupArchive(ctx, "KSFO",
		r.Options.Start.AddDate(0,0,-1), r.Options.End.AddDate(0,0,1))
	if err != nil {
		return err
	}
	r.ReportingContext.Archive = *metar

	user := user.Current(ctx)
	if user != nil {
		r.ReportingContext.UserEmail = user.Email
	}

	r.AddACLs()

	//airframes := ref.NewAirframeCache(ctx)

	return nil
}

func (r *Report)AddACLs() {
	email := r.ReportingContext.UserEmail
	if userInList(email, ACLFOIA) {
		r.Options.CanSeeFOIA = true
	}
}

func userInList(user string, acl []string) bool {
	for _,e := range acl {
		if strings.ToLower(user) == strings.ToLower(e) { return true }
	}
	return false
}
<commit_msg>Add a new person to the FOIA ACL<commit_after>package report

import(
	"strings"

	"golang.org/x/net/context"
	"google.golang.org/appengine/user"

	"github.com/skypies/flightdb/metar"
)

// TODO: just embed / extract this data from the regular context
type ReportingContext struct {
	context.Context
	metar.Archive
	UserEmail       string
}

var(
	// Lower case everything
	ACLFOIA = []string{
		"adam@worrall.cc",
		"raymonde.guindon@gmail.com",
		"rguindon@alumni.stanford.edu",
		"meekgee@gmail.com",
		"nancyjordan650@gmail.com",
		"borg@smwlaw.com",
		"matt@classalameda.com",
		"jnelson@wiai.com",
		"robert.holbrook@gmail.com",
		"jtsunnyvaleair1@gmail.com",
		"mthurd2003@gmail.com",
		"fabrice.beretta@gmail.com",
		"kaneshirokimberly@gmail.com",
	}
)

func (r *Report)setupReportingContext(ctx context.Context) error {
	r.ReportingContext.Context = ctx

	metar,err := metar.LookupArchive(ctx, "KSFO",
		r.Options.Start.AddDate(0,0,-1), r.Options.End.AddDate(0,0,1))
	if err != nil {
		return err
	}
	r.ReportingContext.Archive = *metar

	user := user.Current(ctx)
	if user != nil {
		r.ReportingContext.UserEmail = user.Email
	}

	r.AddACLs()

	//airframes := ref.NewAirframeCache(ctx)

	return nil
}

func (r *Report)AddACLs() {
	email := r.ReportingContext.UserEmail
	if userInList(email, ACLFOIA) {
		r.Options.CanSeeFOIA = true
	}
}

func userInList(user string, acl []string) bool {
	for _,e := range acl {
		if strings.ToLower(user) == strings.ToLower(e) { return true }
	}
	return false
}
<|endoftext|>


<commit_before>package objx

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

var getFixtures = []struct {
	data   string
	get    interface{}
	output interface{}
}{
	// get
	{
		data:   `{"name": "Mat"}`,
		get:    "name",
		output: "Mat",
	},
	// dot notation
	{
		data:   `{"address": {"city": "Boulder"}}`,
		get:    "address.city",
		output: "Boulder",
	},
	{
		data:   `{"one": {"two": {"three": {"four": "hello"}}}}`,
		get:    "one.two.three.four",
		output: "hello",
	},
	// no data
	{
		data:   `{"one": {"two": {"three": {"four": "hello"}}}}`,
		get:    "one.ten",
		output: nil,
	},
	// array [] notation
	{
		data:   `{"tags": ["one", "two", "three"]}`,
		get:    "tags[1]",
		output: "two",
	},
	// array [] notation in map
	{
		data:   `{"types": { "tags": ["one", "two", "three"]}}`,
		get:    "types.tags[1]",
		output: "two",
	},
	{
		data:   `{"tags": [{"name":"one"}, {"name":"two"}, {"name":"three"}]}`,
		get:    "tags[1].name",
		output: "two",
	},
	// array number
	{
		data:   `["one", "two", "three"]`,
		get:    1,
		output: "two",
	},
}

func TestGetFixtures(t *testing.T) {

	for index, fixture := range getFixtures {
		o := MustJson(fixture.data)
		newO := o.Get(fixture.get)
		assert.Equal(t, fixture.output, newO.Obj, "Get fixture %d failed: %v", index+1, fixture)
	}

}
<commit_msg>Add complex test and test name<commit_after>package objx

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

var getFixtures = []struct {
	name   string
	data   string
	get    interface{}
	output interface{}
}{
	// get
	{
		name:   "Simple get",
		data:   `{"name": "Mat"}`,
		get:    "name",
		output: "Mat",
	},
	// dot notation
	{
		name:   "Get with dot notation",
		data:   `{"address": {"city": "Boulder"}}`,
		get:    "address.city",
		output: "Boulder",
	},
	{
		name:   "Deep get with dot notation",
		data:   `{"one": {"two": {"three": {"four": "hello"}}}}`,
		get:    "one.two.three.four",
		output: "hello",
	},
	// no data
	{
		name:   "Get missing with dot notation",
		data:   `{"one": {"two": {"three": {"four": "hello"}}}}`,
		get:    "one.ten",
		output: nil,
	},
	// array [] notation
	{
		name:   "Get with array notation",
		data:   `{"tags": ["one", "two", "three"]}`,
		get:    "tags[1]",
		output: "two",
	},
	// array [] notation in map
	{
		name:   "Get with array and dot notation",
		data:   `{"types": { "tags": ["one", "two", "three"]}}`,
		get:    "types.tags[1]",
		output: "two",
	},
	{
		name:   "Get with array and dot notation - field after array",
		data:   `{"tags": [{"name":"one"}, {"name":"two"}, {"name":"three"}]}`,
		get:    "tags[1].name",
		output: "two",
	},
	{
		name:   "Complex get with array and dot notation",
		data:   `{"tags": [{"list": [{"one":"pizza"}]}]}`,
		get:    "tags[0].list[0].one",
		output: "pizza",
	},
	// array number
	{
		name:   "Get with integer argument",
		data:   `["one", "two", "three"]`,
		get:    1,
		output: "two",
	},
}

func TestGetFixtures(t *testing.T) {

	for index, fixture := range getFixtures {
		o := MustJson(fixture.data)
		newO := o.Get(fixture.get)
		assert.Equal(t, fixture.output, newO.Obj, "Get fixture %d (%s) failed: %v", index+1, fixture.name, fixture)
	}

}
<|endoftext|>


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

package featuretests

import (
	"fmt"
	"reflect"
	"strings"
	"time"

	"github.com/juju/cmd"
	"github.com/juju/errors"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/api/modelmanager"
	undertakerapi "github.com/juju/juju/api/undertaker"
	"github.com/juju/juju/cmd/juju/commands"
	"github.com/juju/juju/environs/configstore"
	"github.com/juju/juju/juju"
	jujutesting "github.com/juju/juju/juju/testing"
	"github.com/juju/juju/provider/dummy"
	"github.com/juju/juju/state"
	"github.com/juju/juju/testing"
	"github.com/juju/juju/testing/factory"
	"github.com/juju/juju/worker/undertaker"
)

type cmdControllerSuite struct {
	jujutesting.JujuConnSuite
}

func (s *cmdControllerSuite) run(c *gc.C, args ...string) *cmd.Context {
	context := testing.Context(c)
	command := commands.NewJujuCommand(context)
	c.Assert(testing.InitCommand(command, args), jc.ErrorIsNil)
	c.Assert(command.Run(context), jc.ErrorIsNil)
	return context
}

func (s *cmdControllerSuite) createEnv(c *gc.C, envname string, isServer bool) {
	modelManager := modelmanager.NewClient(s.APIState)
	_, err := modelManager.CreateModel(s.AdminUserTag(c).Id(), nil, map[string]interface{}{
		"name":            envname,
		"authorized-keys": "ssh-key",
		"controller":      isServer,
	})
	c.Assert(err, jc.ErrorIsNil)
}

func (s *cmdControllerSuite) TestControllerListCommand(c *gc.C) {
	context := s.run(c, "list-controllers")
	expectedOutput := `
CONTROLLER  MODEL  USER         SERVER
kontroll*   admin  admin@local

`[1:]
	c.Assert(testing.Stdout(context), gc.Equals, expectedOutput)
}

func (s *cmdControllerSuite) TestControllerModelsCommand(c *gc.C) {
	s.createEnv(c, "new-model", false)
	context := s.run(c, "list-models")
	c.Assert(testing.Stdout(context), gc.Equals, ""+
		"NAME       OWNER        LAST CONNECTION\n"+
		"admin      admin@local  just now\n"+
		"new-model  admin@local  never connected\n"+
		"\n")
}

func (s *cmdControllerSuite) TestCreateModel(c *gc.C) {
	// The JujuConnSuite doesn't set up an ssh key in the fake home dir,
	// so fake one on the command line.  The dummy provider also expects
	// a config value for 'controller'.
	context := s.run(c, "create-model", "new-model", "authorized-keys=fake-key", "controller=false")
	c.Check(testing.Stdout(context), gc.Equals, "")
	c.Check(testing.Stderr(context), gc.Equals, "created model \"new-model\"\n")

	// Make sure that the saved server details are sufficient to connect
	// to the api server.
	api, err := juju.NewAPIConnection(s.ControllerStore, "kontroll", "admin@local", "new-model", nil)
	c.Assert(err, jc.ErrorIsNil)
	api.Close()
}

func (s *cmdControllerSuite) TestControllerDestroy(c *gc.C) {
	st := s.Factory.MakeModel(c, &factory.ModelParams{
		Name:        "just-a-controller",
		ConfigAttrs: testing.Attrs{"controller": true},
	})
	defer st.Close()

	stop := make(chan struct{})
	done := make(chan struct{})
	// In order for the destroy controller command to complete we need to run
	// the code that the cleaner and undertaker workers would be running in
	// the agent in order to progress the lifecycle of the hosted model,
	// and cleanup the documents.
	go func() {
		defer close(done)
		a := testing.LongAttempt.Start()
		for a.Next() {
			err := s.State.Cleanup()
			c.Check(err, jc.ErrorIsNil)
			err = st.ProcessDyingModel()
			if errors.Cause(err) != state.ErrModelNotDying {
				c.Check(err, jc.ErrorIsNil)
				if err == nil {
					// success!
					return
				}
			}
			select {
			case <-stop:
				return
			default:
				// retry
			}
		}
	}()

	s.run(c, "destroy-controller", "kontroll", "-y", "--destroy-all-models", "--debug")
	close(stop)
	<-done

	store, err := configstore.Default()
	_, err = store.ReadInfo("kontroll:admin")
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
}

func (s *cmdControllerSuite) TestRemoveBlocks(c *gc.C) {
	s.State.SwitchBlockOn(state.DestroyBlock, "TestBlockDestroyModel")
	s.State.SwitchBlockOn(state.ChangeBlock, "TestChangeBlock")

	s.run(c, "remove-all-blocks")

	blocks, err := s.State.AllBlocksForController()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(blocks, gc.HasLen, 0)
}

func (s *cmdControllerSuite) TestControllerKill(c *gc.C) {
	st := s.Factory.MakeModel(c, &factory.ModelParams{
		Name: "foo",
	})

	st.SwitchBlockOn(state.DestroyBlock, "TestBlockDestroyModel")
	st.Close()

	s.run(c, "kill-controller", "kontroll", "-y")

	store, err := configstore.Default()
	_, err = store.ReadInfo("kontroll:admin")
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
}

func (s *cmdControllerSuite) TestListBlocks(c *gc.C) {
	s.State.SwitchBlockOn(state.DestroyBlock, "TestBlockDestroyModel")
	s.State.SwitchBlockOn(state.ChangeBlock, "TestChangeBlock")

	ctx := s.run(c, "list-all-blocks", "--format", "json")
	expected := fmt.Sprintf(`[{"name":"admin","model-uuid":"%s","owner-tag":"%s","blocks":["BlockDestroy","BlockChange"]}]`,
		s.State.ModelUUID(), s.AdminUserTag(c).String())

	strippedOut := strings.Replace(testing.Stdout(ctx), "\n", "", -1)
	c.Check(strippedOut, gc.Equals, expected)
}

func (s *cmdControllerSuite) TestSystemKillCallsEnvironDestroyOnHostedEnviron(c *gc.C) {
	st := s.Factory.MakeModel(c, &factory.ModelParams{
		Name: "foo",
	})
	defer st.Close()

	st.SwitchBlockOn(state.DestroyBlock, "TestBlockDestroyModel")
	st.Close()

	opc := make(chan dummy.Operation, 200)
	dummy.Listen(opc)

	client := undertakerapi.NewClient(s.APIState)

	startTime := time.Date(2015, time.September, 1, 17, 2, 1, 0, time.UTC)
	mClock := testing.NewClock(startTime)
	undertaker.NewUndertaker(client, mClock)

	store, err := configstore.Default()
	_, err = store.ReadInfo("kontroll:admin")
	c.Assert(err, jc.ErrorIsNil)

	s.run(c, "kill-controller", "kontroll", "-y")

	// Ensure that Destroy was called on the hosted model ...
	opRecvTimeout(c, st, opc, dummy.OpDestroy{})

	// ... and that the configstore was removed.
	_, err = store.ReadInfo("kontroll:admin")
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
}

// opRecvTimeout waits for any of the given kinds of operation to
// be received from ops, and times out if not.
func opRecvTimeout(c *gc.C, st *state.State, opc <-chan dummy.Operation, kinds ...dummy.Operation) dummy.Operation {
	st.StartSync()
	for {
		select {
		case op := <-opc:
			for _, k := range kinds {
				if reflect.TypeOf(op) == reflect.TypeOf(k) {
					return op
				}
			}
			c.Logf("discarding unknown event %#v", op)
		case <-time.After(testing.LongWait):
			c.Fatalf("time out wating for operation")
		}
	}
}
<commit_msg>Fix test<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package featuretests

import (
	"fmt"
	"reflect"
	"strings"
	"time"

	"github.com/juju/cmd"
	"github.com/juju/errors"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/api/modelmanager"
	undertakerapi "github.com/juju/juju/api/undertaker"
	"github.com/juju/juju/cmd/juju/commands"
	"github.com/juju/juju/environs/configstore"
	"github.com/juju/juju/juju"
	jujutesting "github.com/juju/juju/juju/testing"
	"github.com/juju/juju/provider/dummy"
	"github.com/juju/juju/state"
	"github.com/juju/juju/testing"
	"github.com/juju/juju/testing/factory"
	"github.com/juju/juju/worker/undertaker"
)

type cmdControllerSuite struct {
	jujutesting.JujuConnSuite
}

func (s *cmdControllerSuite) run(c *gc.C, args ...string) *cmd.Context {
	context := testing.Context(c)
	command := commands.NewJujuCommand(context)
	c.Assert(testing.InitCommand(command, args), jc.ErrorIsNil)
	c.Assert(command.Run(context), jc.ErrorIsNil)
	return context
}

func (s *cmdControllerSuite) createEnv(c *gc.C, envname string, isServer bool) {
	modelManager := modelmanager.NewClient(s.APIState)
	_, err := modelManager.CreateModel(s.AdminUserTag(c).Id(), nil, map[string]interface{}{
		"name":            envname,
		"authorized-keys": "ssh-key",
		"controller":      isServer,
	})
	c.Assert(err, jc.ErrorIsNil)
}

func (s *cmdControllerSuite) TestControllerListCommand(c *gc.C) {
	context := s.run(c, "list-controllers")
	expectedOutput := `
CONTROLLER  MODEL  USER         SERVER
kontroll*   admin  admin@local

`[1:]
	c.Assert(testing.Stdout(context), gc.Equals, expectedOutput)
}

func (s *cmdControllerSuite) TestControllerModelsCommand(c *gc.C) {
	s.createEnv(c, "new-model", false)
	context := s.run(c, "list-models")
	c.Assert(testing.Stdout(context), gc.Equals, ""+
		"NAME       OWNER        LAST CONNECTION\n"+
		"admin*     admin@local  just now\n"+
		"new-model  admin@local  never connected\n"+
		"\n")
}

func (s *cmdControllerSuite) TestCreateModel(c *gc.C) {
	// The JujuConnSuite doesn't set up an ssh key in the fake home dir,
	// so fake one on the command line.  The dummy provider also expects
	// a config value for 'controller'.
	context := s.run(c, "create-model", "new-model", "authorized-keys=fake-key", "controller=false")
	c.Check(testing.Stdout(context), gc.Equals, "")
	c.Check(testing.Stderr(context), gc.Equals, "created model \"new-model\"\n")

	// Make sure that the saved server details are sufficient to connect
	// to the api server.
	api, err := juju.NewAPIConnection(s.ControllerStore, "kontroll", "admin@local", "new-model", nil)
	c.Assert(err, jc.ErrorIsNil)
	api.Close()
}

func (s *cmdControllerSuite) TestControllerDestroy(c *gc.C) {
	st := s.Factory.MakeModel(c, &factory.ModelParams{
		Name:        "just-a-controller",
		ConfigAttrs: testing.Attrs{"controller": true},
	})
	defer st.Close()

	stop := make(chan struct{})
	done := make(chan struct{})
	// In order for the destroy controller command to complete we need to run
	// the code that the cleaner and undertaker workers would be running in
	// the agent in order to progress the lifecycle of the hosted model,
	// and cleanup the documents.
	go func() {
		defer close(done)
		a := testing.LongAttempt.Start()
		for a.Next() {
			err := s.State.Cleanup()
			c.Check(err, jc.ErrorIsNil)
			err = st.ProcessDyingModel()
			if errors.Cause(err) != state.ErrModelNotDying {
				c.Check(err, jc.ErrorIsNil)
				if err == nil {
					// success!
					return
				}
			}
			select {
			case <-stop:
				return
			default:
				// retry
			}
		}
	}()

	s.run(c, "destroy-controller", "kontroll", "-y", "--destroy-all-models", "--debug")
	close(stop)
	<-done

	store, err := configstore.Default()
	_, err = store.ReadInfo("kontroll:admin")
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
}

func (s *cmdControllerSuite) TestRemoveBlocks(c *gc.C) {
	s.State.SwitchBlockOn(state.DestroyBlock, "TestBlockDestroyModel")
	s.State.SwitchBlockOn(state.ChangeBlock, "TestChangeBlock")

	s.run(c, "remove-all-blocks")

	blocks, err := s.State.AllBlocksForController()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(blocks, gc.HasLen, 0)
}

func (s *cmdControllerSuite) TestControllerKill(c *gc.C) {
	st := s.Factory.MakeModel(c, &factory.ModelParams{
		Name: "foo",
	})

	st.SwitchBlockOn(state.DestroyBlock, "TestBlockDestroyModel")
	st.Close()

	s.run(c, "kill-controller", "kontroll", "-y")

	store, err := configstore.Default()
	_, err = store.ReadInfo("kontroll:admin")
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
}

func (s *cmdControllerSuite) TestListBlocks(c *gc.C) {
	s.State.SwitchBlockOn(state.DestroyBlock, "TestBlockDestroyModel")
	s.State.SwitchBlockOn(state.ChangeBlock, "TestChangeBlock")

	ctx := s.run(c, "list-all-blocks", "--format", "json")
	expected := fmt.Sprintf(`[{"name":"admin","model-uuid":"%s","owner-tag":"%s","blocks":["BlockDestroy","BlockChange"]}]`,
		s.State.ModelUUID(), s.AdminUserTag(c).String())

	strippedOut := strings.Replace(testing.Stdout(ctx), "\n", "", -1)
	c.Check(strippedOut, gc.Equals, expected)
}

func (s *cmdControllerSuite) TestSystemKillCallsEnvironDestroyOnHostedEnviron(c *gc.C) {
	st := s.Factory.MakeModel(c, &factory.ModelParams{
		Name: "foo",
	})
	defer st.Close()

	st.SwitchBlockOn(state.DestroyBlock, "TestBlockDestroyModel")
	st.Close()

	opc := make(chan dummy.Operation, 200)
	dummy.Listen(opc)

	client := undertakerapi.NewClient(s.APIState)

	startTime := time.Date(2015, time.September, 1, 17, 2, 1, 0, time.UTC)
	mClock := testing.NewClock(startTime)
	undertaker.NewUndertaker(client, mClock)

	store, err := configstore.Default()
	_, err = store.ReadInfo("kontroll:admin")
	c.Assert(err, jc.ErrorIsNil)

	s.run(c, "kill-controller", "kontroll", "-y")

	// Ensure that Destroy was called on the hosted model ...
	opRecvTimeout(c, st, opc, dummy.OpDestroy{})

	// ... and that the configstore was removed.
	_, err = store.ReadInfo("kontroll:admin")
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
}

// opRecvTimeout waits for any of the given kinds of operation to
// be received from ops, and times out if not.
func opRecvTimeout(c *gc.C, st *state.State, opc <-chan dummy.Operation, kinds ...dummy.Operation) dummy.Operation {
	st.StartSync()
	for {
		select {
		case op := <-opc:
			for _, k := range kinds {
				if reflect.TypeOf(op) == reflect.TypeOf(k) {
					return op
				}
			}
			c.Logf("discarding unknown event %#v", op)
		case <-time.After(testing.LongWait):
			c.Fatalf("time out wating for operation")
		}
	}
}
<|endoftext|>


<commit_before>package xpi

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"crypto/x509/pkix"
	"fmt"
	"math/big"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/pkg/errors"
	"go.mozilla.org/cose"
)

// populateRsaCache adds an rsa key to the cache every
// XPISigner.rsaCacheSleepDuration, blocks when the cache channel is
// full, and should be run as a goroutine
func (s *XPISigner) populateRsaCache(size int) {
	var (
		err error
		key *rsa.PrivateKey
		start time.Time
	)
	for {
		start = time.Now()
		key, err = rsa.GenerateKey(rand.Reader, size)
		if err != nil {
			log.Fatalf("xpi: error generating RSA key for cache: %s", err)
		}

		s.SendHistogram("xpi.rsa_cache.gen_key_dur", time.Since(start))
		s.rsaCache <- key
		time.Sleep(s.rsaCacheGeneratorSleepDuration)
	}
}

// monitorRsaCacheSize sends the number of cached keys and cache size
// to datadog. It should be run as a goroutine
func (s *XPISigner) monitorRsaCacheSize() {
	for {
		s.SendGauge("xpi.rsa_cache.chan_len", len(s.rsaCache))

		// chan capacity should be constant but is useful for
		// knowing % cache filled across deploys
		s.SendGauge("xpi.rsa_cache.chan_cap", cap(s.rsaCache))

		time.Sleep(s.rsaCacheSizeSampleRate)
	}
}

// retrieve a key from the cache or generate one if it takes too long
// or if the size is wrong
func (s *XPISigner) getRsaKey(size int) (*rsa.PrivateKey, error) {
	var (
		err error
		key *rsa.PrivateKey
		start time.Time
	)
	start = time.Now()
	select {
	case key = <-s.rsaCache:
		if key.N.BitLen() != size {
			// it's theoritically impossible for this to happen
			// because the end entity has the same key size has
			// the signer, but we're paranoid so handling it
			log.Warnf("WARNING: xpi rsa cache returned a key of size %d when %d was requested", key.N.BitLen(), size)
			key, err = rsa.GenerateKey(rand.Reader, size)
		}
	case <-time.After(s.rsaCacheFetchTimeout):
		// generate a key if none available
		key, err = rsa.GenerateKey(rand.Reader, size)
	}
	s.SendHistogram("xpi.rsa_cache.get_key", time.Since(start))
	return key, err
}

// makeTemplate returns a pointer to a template for an x509.Certificate EE
func (s *XPISigner) makeTemplate(cn string) *x509.Certificate {
	cndigest := sha256.Sum256([]byte(cn))
	return &x509.Certificate{
		// The maximum length of a serial number per rfc 5280 is 20 bytes / 160 bits
		// https://tools.ietf.org/html/rfc5280#section-4.1.2.2
		// Setting it to nanoseconds guarantees we'll never have two conflicting serials
		SerialNumber: big.NewInt(time.Now().UnixNano()),
		// PKIX requires EE's to have a valid DNS Names when the intermediate has
		// a constraint, so we hash the CN into an fqdn to get something unique enough
		DNSNames: []string{fmt.Sprintf("%x.%x.addons.mozilla.org", cndigest[:16], cndigest[16:])},
		Subject: pkix.Name{
			CommonName:         cn,
			Organization:       []string{"Addons"},
			OrganizationalUnit: []string{s.OU},
			Country:            []string{"US"},
			Province:           []string{"CA"},
			Locality:           []string{"Mountain View"},
		},
		NotBefore:          time.Now(),
		NotAfter:           time.Now().Add(8760 * time.Hour), // one year
		SignatureAlgorithm: s.issuerCert.SignatureAlgorithm,
		KeyUsage:           x509.KeyUsageDigitalSignature,
	}
}

// generateIssuerEEKeyPair returns a public and private key pair
// matching the issuer XPISigner issuerKey size and type
func (s *XPISigner) generateIssuerEEKeyPair() (eeKey crypto.PrivateKey, eePublicKey crypto.PublicKey, err error) {
	switch s.issuerKey.(type) {
	case *rsa.PrivateKey:
		size := s.issuerKey.(*rsa.PrivateKey).N.BitLen()
		eeKey, err = s.getRsaKey(size)
		if err != nil {
			err = errors.Wrapf(err, "failed to generate rsa private key of size %d", size)
			return
		}
		eePublicKey = eeKey.(*rsa.PrivateKey).Public()
	case *ecdsa.PrivateKey:
		curve := s.issuerKey.(*ecdsa.PrivateKey).Curve
		eeKey, err = ecdsa.GenerateKey(curve, rand.Reader)
		if err != nil {
			err = errors.Wrapf(err, "failed to generate ecdsa private key on curve %s", curve.Params().Name)
			return
		}
		eePublicKey = eeKey.(*ecdsa.PrivateKey).Public()
	}
	return
}

// MakeEndEntity generates a private key and certificate ready to sign a given XPI.
//
// The subject CN of the certificate is taken from the `cn` string argument.
//
// The key type is identical to the key type of the signer that issues
// the certificate when the optional `coseAlg` argument is nil. For
// example, if the signer uses an RSA 2048 key, so will the
// end-entity. When `coseAlg` is not nil, a key type of the COSE
// algorithm is generated.
//
// The signature expiration date is copied over from the issuer.
//
// The signed x509 certificate and private key are returned.
func (s *XPISigner) MakeEndEntity(cn string, coseAlg *cose.Algorithm) (eeCert *x509.Certificate, eeKey crypto.PrivateKey, err error) {
	var (
		eePublicKey crypto.PublicKey
		derCert     []byte
	)

	template := s.makeTemplate(cn)

	if coseAlg == nil {
		eeKey, eePublicKey, err = s.generateIssuerEEKeyPair()
		if err != nil {
			err = errors.Wrapf(err, "xpi.MakeEndEntity: error generating key matching issuer")
			return
		}
	} else {
		eeKey, eePublicKey, err = s.generateCOSEKeyPair(coseAlg)
		if err != nil {
			err = errors.Wrapf(err, "xpi.MakeEndEntity: error generating key matching COSE Algorithm type %s", coseAlg.Name)
			return
		}
	}

	derCert, err = x509.CreateCertificate(rand.Reader, template, s.issuerCert, eePublicKey, s.issuerKey)
	if err != nil {
		err = errors.Wrapf(err, "xpi.MakeEndEntity: failed to create certificate")
		return
	}
	if len(derCert) == 0 {
		err = errors.Errorf("xpi.MakeEndEntity: certificate creation failed for an unknown reason")
		return
	}
	eeCert, err = x509.ParseCertificate(derCert)
	if err != nil {
		err = errors.Wrapf(err, "xpi.MakeEndEntity: certificate parsing failed")
	}
	return
}
<commit_msg>xpi: check for nil key in generator<commit_after>package xpi

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"crypto/x509/pkix"
	"fmt"
	"math/big"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/pkg/errors"
	"go.mozilla.org/cose"
)

// populateRsaCache adds an rsa key to the cache every
// XPISigner.rsaCacheSleepDuration, blocks when the cache channel is
// full, and should be run as a goroutine
func (s *XPISigner) populateRsaCache(size int) {
	var (
		err error
		key *rsa.PrivateKey
		start time.Time
	)
	for {
		start = time.Now()
		key, err = rsa.GenerateKey(rand.Reader, size)
		if err != nil {
			log.Fatalf("xpi: error generating RSA key for cache: %s", err)
		}
		if key == nil {
			log.Fatal("xpi: error generated nil RSA key for cache")
		}

		s.SendHistogram("xpi.rsa_cache.gen_key_dur", time.Since(start))
		s.rsaCache <- key
		time.Sleep(s.rsaCacheGeneratorSleepDuration)
	}
}

// monitorRsaCacheSize sends the number of cached keys and cache size
// to datadog. It should be run as a goroutine
func (s *XPISigner) monitorRsaCacheSize() {
	for {
		s.SendGauge("xpi.rsa_cache.chan_len", len(s.rsaCache))

		// chan capacity should be constant but is useful for
		// knowing % cache filled across deploys
		s.SendGauge("xpi.rsa_cache.chan_cap", cap(s.rsaCache))

		time.Sleep(s.rsaCacheSizeSampleRate)
	}
}

// retrieve a key from the cache or generate one if it takes too long
// or if the size is wrong
func (s *XPISigner) getRsaKey(size int) (*rsa.PrivateKey, error) {
	var (
		err error
		key *rsa.PrivateKey
		start time.Time
	)
	start = time.Now()
	select {
	case key = <-s.rsaCache:
		if key.N.BitLen() != size {
			// it's theoritically impossible for this to happen
			// because the end entity has the same key size has
			// the signer, but we're paranoid so handling it
			log.Warnf("WARNING: xpi rsa cache returned a key of size %d when %d was requested", key.N.BitLen(), size)
			key, err = rsa.GenerateKey(rand.Reader, size)
		}
	case <-time.After(s.rsaCacheFetchTimeout):
		// generate a key if none available
		key, err = rsa.GenerateKey(rand.Reader, size)
	}
	s.SendHistogram("xpi.rsa_cache.get_key", time.Since(start))
	return key, err
}

// makeTemplate returns a pointer to a template for an x509.Certificate EE
func (s *XPISigner) makeTemplate(cn string) *x509.Certificate {
	cndigest := sha256.Sum256([]byte(cn))
	return &x509.Certificate{
		// The maximum length of a serial number per rfc 5280 is 20 bytes / 160 bits
		// https://tools.ietf.org/html/rfc5280#section-4.1.2.2
		// Setting it to nanoseconds guarantees we'll never have two conflicting serials
		SerialNumber: big.NewInt(time.Now().UnixNano()),
		// PKIX requires EE's to have a valid DNS Names when the intermediate has
		// a constraint, so we hash the CN into an fqdn to get something unique enough
		DNSNames: []string{fmt.Sprintf("%x.%x.addons.mozilla.org", cndigest[:16], cndigest[16:])},
		Subject: pkix.Name{
			CommonName:         cn,
			Organization:       []string{"Addons"},
			OrganizationalUnit: []string{s.OU},
			Country:            []string{"US"},
			Province:           []string{"CA"},
			Locality:           []string{"Mountain View"},
		},
		NotBefore:          time.Now(),
		NotAfter:           time.Now().Add(8760 * time.Hour), // one year
		SignatureAlgorithm: s.issuerCert.SignatureAlgorithm,
		KeyUsage:           x509.KeyUsageDigitalSignature,
	}
}

// generateIssuerEEKeyPair returns a public and private key pair
// matching the issuer XPISigner issuerKey size and type
func (s *XPISigner) generateIssuerEEKeyPair() (eeKey crypto.PrivateKey, eePublicKey crypto.PublicKey, err error) {
	switch s.issuerKey.(type) {
	case *rsa.PrivateKey:
		size := s.issuerKey.(*rsa.PrivateKey).N.BitLen()
		eeKey, err = s.getRsaKey(size)
		if err != nil {
			err = errors.Wrapf(err, "failed to generate rsa private key of size %d", size)
			return
		}
		eePublicKey = eeKey.(*rsa.PrivateKey).Public()
	case *ecdsa.PrivateKey:
		curve := s.issuerKey.(*ecdsa.PrivateKey).Curve
		eeKey, err = ecdsa.GenerateKey(curve, rand.Reader)
		if err != nil {
			err = errors.Wrapf(err, "failed to generate ecdsa private key on curve %s", curve.Params().Name)
			return
		}
		eePublicKey = eeKey.(*ecdsa.PrivateKey).Public()
	}
	return
}

// MakeEndEntity generates a private key and certificate ready to sign a given XPI.
//
// The subject CN of the certificate is taken from the `cn` string argument.
//
// The key type is identical to the key type of the signer that issues
// the certificate when the optional `coseAlg` argument is nil. For
// example, if the signer uses an RSA 2048 key, so will the
// end-entity. When `coseAlg` is not nil, a key type of the COSE
// algorithm is generated.
//
// The signature expiration date is copied over from the issuer.
//
// The signed x509 certificate and private key are returned.
func (s *XPISigner) MakeEndEntity(cn string, coseAlg *cose.Algorithm) (eeCert *x509.Certificate, eeKey crypto.PrivateKey, err error) {
	var (
		eePublicKey crypto.PublicKey
		derCert     []byte
	)

	template := s.makeTemplate(cn)

	if coseAlg == nil {
		eeKey, eePublicKey, err = s.generateIssuerEEKeyPair()
		if err != nil {
			err = errors.Wrapf(err, "xpi.MakeEndEntity: error generating key matching issuer")
			return
		}
	} else {
		eeKey, eePublicKey, err = s.generateCOSEKeyPair(coseAlg)
		if err != nil {
			err = errors.Wrapf(err, "xpi.MakeEndEntity: error generating key matching COSE Algorithm type %s", coseAlg.Name)
			return
		}
	}

	derCert, err = x509.CreateCertificate(rand.Reader, template, s.issuerCert, eePublicKey, s.issuerKey)
	if err != nil {
		err = errors.Wrapf(err, "xpi.MakeEndEntity: failed to create certificate")
		return
	}
	if len(derCert) == 0 {
		err = errors.Errorf("xpi.MakeEndEntity: certificate creation failed for an unknown reason")
		return
	}
	eeCert, err = x509.ParseCertificate(derCert)
	if err != nil {
		err = errors.Wrapf(err, "xpi.MakeEndEntity: certificate parsing failed")
	}
	return
}
<|endoftext|>