<commit_before>package paillier

import (
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"io"
	"math/big"
)

// Public key for a threshold Paillier scheme.
//
// `V` is a generator in  the cyclic group of squares Z_n^2 and is used to
// execute a zero-knowledge proof of a received share decryption.
//
// `Vi` is an array of verification keys for each decryption server `i`.
//
// Key generation, encryption, share decryption and combining for the threshold
// Paillier scheme has been described in [DJN 10], section 5.1.
//
//     [DJN 10]: Ivan Damgard, Mads Jurik, Jesper Buus Nielsen, (2010)
//               A Generalization of Paillier’s Public-Key System
//               with Applications to Electronic Voting
//               Aarhus University, Dept. of Computer Science, BRICS
type ThresholdKey struct {
	PublicKey
	TotalNumberOfDecryptionServers int
	Threshold                      int
	V                              *big.Int
	Vi                             []*big.Int
}

// Returns the value of [(4*delta^2)]^-1  mod n.
// It is a constant value for the given `ThresholdKey` and is used in the last
// step of share combining.
func (this *ThresholdKey) combineSharesConstant() *big.Int {
	tmp := new(big.Int).Mul(FOUR, new(big.Int).Mul(this.delta(), this.delta()))
	return (&big.Int{}).ModInverse(tmp, this.N)
}

// Returns the factorial of the number of `TotalNumberOfDecryptionServers`.
// It is a contant value for the given `ThresholdKey`.
func (this *ThresholdKey) delta() *big.Int {
	return Factorial(this.TotalNumberOfDecryptionServers)
}

// Checks if the number of received, unique shares is less than the
// required threshold.
// This method does not execute ZKP on received shares.
func (this *ThresholdKey) makeVerificationBeforeCombiningPartialDecryptions(shares []*PartialDecryption) error {
	if len(shares) < this.Threshold {
		return errors.New("Threshold not meet")
	}
	tmp := make(map[int]bool)
	for _, share := range shares {
		tmp[share.Id] = true
	}
	if len(tmp) != len(shares) {
		return errors.New("two shares has been created by the same server")
	}
	return nil
}

func (this *ThresholdKey) updateLambda(share1, share2 *PartialDecryption, lambda *big.Int) *big.Int {
	num := new(big.Int).Mul(lambda, big.NewInt(int64(-share2.Id)))
	denom := big.NewInt(int64(share1.Id - share2.Id))
	return new(big.Int).Div(num, denom)
}

// Evaluates lambda parameter for each decrypted share. See second figure in the
// "Share combining" paragraph in [DJK 10], section 5.2.
func (this *ThresholdKey) computeLambda(share *PartialDecryption, shares []*PartialDecryption) *big.Int {
	lambda := this.delta()
	for _, share2 := range shares {
		if share2.Id != share.Id {
			lambda = this.updateLambda(share, share2, lambda)
		}
	}
	return lambda
}

// Used to evaluate c' parameter which combines individual share decryptions.
//
// Modulo division is performed on the computed exponent to avoid creating
// large numbers. This is possible because of the following property of modulo:
// A^B mod C = (A mod C)^B mod C
//
// Modulo division is performed on the computed coefficient because of the
// following property of modulo:
// (AB) mod C = (A mod C * B mod C) mod C
// Note, we need to combine coefficients into single c'.
func (this *ThresholdKey) updateCprime(cprime, lambda *big.Int, share *PartialDecryption) *big.Int {
	twoLambda := new(big.Int).Mul(TWO, lambda)
	ret := this.exp(share.Decryption, twoLambda, this.GetNSquare())
	ret = new(big.Int).Mul(cprime, ret)
	return new(big.Int).Mod(ret, this.GetNSquare())
}

func (this *ThresholdKey) exp(a, b, c *big.Int) *big.Int {
	if b.Cmp(ZERO) == -1 {
		ret := new(big.Int).Exp(a, new(big.Int).Neg(b), c)
		return new(big.Int).ModInverse(ret, c)
	}
	return new(big.Int).Exp(a, b, c)
}

// Executes the last step of message decryption. Takes `cprime` value computed
// from valid shares provided by decryption servers and multiplies this value
// by `combineSharesContant` which is specific to the given public `ThresholdKey`.
func (this *ThresholdKey) computeDecryption(cprime *big.Int) *big.Int {
	l := L(cprime, this.N)
	return new(big.Int).Mod(new(big.Int).Mul(this.combineSharesConstant(), l), this.N)
}

// Combines partial decryptions provided by decryption servers and returns
// decrypted message.
// This function does not verify zero knowledge proofs. Returned message can be
// incorrectly decrypted if an adversary corrupted partial decryption.
func (this *ThresholdKey) CombinePartialDecryptions(shares []*PartialDecryption) (*big.Int, error) {
	if err := this.makeVerificationBeforeCombiningPartialDecryptions(shares); err != nil {
		return nil, err
	}

	cprime := ONE
	for _, share := range shares {
		lambda := this.computeLambda(share, shares)
		cprime = this.updateCprime(cprime, lambda, share)
	}

	return this.computeDecryption(cprime), nil
}

// Combines partial decryptions provided by decription servers and returns
// decrypted message.
// Function verifies zero knowledge proofs and filters out all shares that failed
// verification.
func (this *ThresholdKey) CombinePartialDecryptionsZKP(shares []*PartialDecryptionZKP) (*big.Int, error) {
	ret := make([]*PartialDecryption, 0)
	for _, share := range shares {
		if share.Verify() {
			ret = append(ret, &share.PartialDecryption)
		}
	}
	return this.CombinePartialDecryptions(ret)
}

// Verifies if the decryption of `encryptedMessage` has been done properly.
// It verifies all the zero-knoledge proofs, the value of the decrypted
// and decrypted message. The method returns `nil` if everything is fine.
// Otherwise, it returns an explicative message.
func (this *ThresholdKey) VerifyDecryption(encryptedMessage, decryptedMessage *big.Int, shares []*PartialDecryptionZKP) error {
	for _, share := range shares {
		if share.C.Cmp(encryptedMessage) != 0 {
			return errors.New("The encrypted message is not the same than the one in the shares")
		}
	}
	res, err := this.CombinePartialDecryptionsZKP(shares)
	if err != nil {
		return err
	}
	if res.Cmp(decryptedMessage) != 0 {
		return errors.New("The decrypted message is not the same than the one in the shares")
	}
	return nil
}

// Private key for a threshold Paillier scheme. Holds private information
// for the given decryption server.
// `Id` is the unique identifier of a decryption server and `Share` is a secret
// share generated from hiding polynomial and is used for a partial share decryption.
type ThresholdPrivateKey struct {
	ThresholdKey
	Id    int
	Share *big.Int
}

// Decrypts the cypher text and returns the partial decryption
func (this *ThresholdPrivateKey) Decrypt(c *big.Int) *PartialDecryption {
	ret := new(PartialDecryption)
	ret.Id = this.Id
	exp := new(big.Int).Mul(this.Share, new(big.Int).Mul(TWO, this.delta()))
	ret.Decryption = new(big.Int).Exp(c, exp, this.GetNSquare())

	return ret
}

func (this *ThresholdPrivateKey) copyVi() []*big.Int {
	ret := make([]*big.Int, len(this.Vi))
	for i, vi := range this.Vi {
		ret[i] = new(big.Int).Add(vi, big.NewInt(0))
	}
	return ret
}

func (this *ThresholdPrivateKey) getThresholdKey() *ThresholdKey {
	ret := new(ThresholdKey)
	ret.Threshold = this.Threshold
	ret.TotalNumberOfDecryptionServers = this.TotalNumberOfDecryptionServers
	ret.V = new(big.Int).Add(this.V, big.NewInt(0))
	ret.Vi = this.copyVi()
	ret.N = new(big.Int).Add(this.N, big.NewInt(0))
	return ret
}

func (this *ThresholdPrivateKey) computeZ(r, e *big.Int) *big.Int {
	tmp := new(big.Int).Mul(e, this.delta())
	tmp = new(big.Int).Mul(tmp, this.Share)
	return new(big.Int).Add(r, tmp)
}

func (this *ThresholdPrivateKey) computeHash(a, b, c4, ci2 *big.Int) *big.Int {
	hash := sha256.New()
	hash.Write(a.Bytes())
	hash.Write(b.Bytes())
	hash.Write(c4.Bytes())
	hash.Write(ci2.Bytes())
	return new(big.Int).SetBytes(hash.Sum([]byte{}))
}

func (this *ThresholdPrivateKey) DecryptAndProduceZNP(c *big.Int, random io.Reader) (*PartialDecryptionZKP, error) {
	pd := new(PartialDecryptionZKP)
	pd.Key = this.getThresholdKey()
	pd.C = c
	pd.Id = this.Id
	pd.Decryption = this.Decrypt(c).Decryption

	// choose random number
	r, err := rand.Int(random, this.GetNSquare())
	if err != nil {
		return nil, err
	}
	//  compute a
	c4 := new(big.Int).Exp(c, FOUR, nil)
	a := new(big.Int).Exp(c4, r, this.GetNSquare())

	// compute b
	b := new(big.Int).Exp(this.V, r, this.GetNSquare())

	// compute hash
	ci2 := new(big.Int).Exp(pd.Decryption, big.NewInt(2), nil)

	pd.E = this.computeHash(a, b, c4, ci2)

	pd.Z = this.computeZ(r, pd.E)

	return pd, nil
}

//  Verify if the partial decryption key is well formed.  If well formed,
//  the method return nil else an explicative error is returned.
func (this *ThresholdPrivateKey) Validate(random io.Reader) error {
	m, err := rand.Int(random, this.N)
	if err != nil {
		return err
	}
	c, err := this.Encrypt(m, random)
	if err != nil {
		return err
	}
	proof, err := this.DecryptAndProduceZNP(c.C, random)
	if err != nil {
		return err
	}
	if !proof.Verify() {
		return errors.New("invalid share.")
	}
	return nil
}

type PartialDecryption struct {
	Id         int
	Decryption *big.Int
}

type PartialDecryptionZKP struct {
	PartialDecryption
	Key *ThresholdKey // the public key used to encrypt
	E   *big.Int      // the challenge
	Z   *big.Int      // the value needed to check to verify the decryption
	C   *big.Int      // the input cypher text
}

func (this *PartialDecryptionZKP) verifyPart1() *big.Int {
	c4 := new(big.Int).Exp(this.C, FOUR, nil)
	decryption2 := new(big.Int).Exp(this.Decryption, TWO, nil)

	a1 := new(big.Int).Exp(c4, this.Z, this.Key.GetNSquare())
	a2 := new(big.Int).Exp(decryption2, this.E, this.Key.GetNSquare())
	a2 = new(big.Int).ModInverse(a2, this.Key.GetNSquare())
	a := new(big.Int).Mod(new(big.Int).Mul(a1, a2), this.Key.GetNSquare())
	return a
}

func (this *PartialDecryptionZKP) neg(n *big.Int) *big.Int {
	return new(big.Int).Neg(n)
}

func (this *PartialDecryptionZKP) verifyPart2() *big.Int {
	vi := this.Key.Vi[this.Id-1]
	b1 := new(big.Int).Exp(this.Key.V, this.Z, this.Key.GetNSquare())
	b2 := new(big.Int).Exp(vi, this.E, this.Key.GetNSquare())
	b2 = new(big.Int).ModInverse(b2, this.Key.GetNSquare())
	b := new(big.Int).Mod(new(big.Int).Mul(b1, b2), this.Key.GetNSquare())
	return b
}

func (this *PartialDecryptionZKP) Verify() bool {
	a := this.verifyPart1()
	b := this.verifyPart2()
	hash := sha256.New()
	hash.Write(a.Bytes())
	hash.Write(b.Bytes())
	c4 := new(big.Int).Exp(this.C, FOUR, nil)
	hash.Write(c4.Bytes())
	ci2 := new(big.Int).Exp(this.Decryption, TWO, nil)
	hash.Write(ci2.Bytes())

	expectedE := new(big.Int).SetBytes(hash.Sum([]byte{}))
	return this.E.Cmp(expectedE) == 0
}
<commit_msg>Added documentation to code evaluating ZKP<commit_after>package paillier

import (
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"io"
	"math/big"
)

// Public key for a threshold Paillier scheme.
//
// `V` is a generator in  the cyclic group of squares Z_n^2 and is used to
// execute a zero-knowledge proof of a received share decryption.
//
// `Vi` is an array of verification keys for each decryption server `i` used to
// execute a zero-knowledge proof of a received share decryption.
//
// Key generation, encryption, share decryption and combining for the threshold
// Paillier scheme has been described in [DJN 10], section 5.1.
//
//     [DJN 10]: Ivan Damgard, Mads Jurik, Jesper Buus Nielsen, (2010)
//               A Generalization of Paillier’s Public-Key System
//               with Applications to Electronic Voting
//               Aarhus University, Dept. of Computer Science, BRICS
type ThresholdKey struct {
	PublicKey
	TotalNumberOfDecryptionServers int
	Threshold                      int
	V                              *big.Int   // needed for ZKP
	Vi                             []*big.Int // needed for ZKP
}

// Returns the value of [(4*delta^2)]^-1  mod n.
// It is a constant value for the given `ThresholdKey` and is used in the last
// step of share combining.
func (this *ThresholdKey) combineSharesConstant() *big.Int {
	tmp := new(big.Int).Mul(FOUR, new(big.Int).Mul(this.delta(), this.delta()))
	return (&big.Int{}).ModInverse(tmp, this.N)
}

// Returns the factorial of the number of `TotalNumberOfDecryptionServers`.
// It is a contant value for the given `ThresholdKey`.
func (this *ThresholdKey) delta() *big.Int {
	return Factorial(this.TotalNumberOfDecryptionServers)
}

// Checks if the number of received, unique shares is less than the
// required threshold.
// This method does not execute ZKP on received shares.
func (this *ThresholdKey) makeVerificationBeforeCombiningPartialDecryptions(shares []*PartialDecryption) error {
	if len(shares) < this.Threshold {
		return errors.New("Threshold not meet")
	}
	tmp := make(map[int]bool)
	for _, share := range shares {
		tmp[share.Id] = true
	}
	if len(tmp) != len(shares) {
		return errors.New("two shares has been created by the same server")
	}
	return nil
}

func (this *ThresholdKey) updateLambda(share1, share2 *PartialDecryption, lambda *big.Int) *big.Int {
	num := new(big.Int).Mul(lambda, big.NewInt(int64(-share2.Id)))
	denom := big.NewInt(int64(share1.Id - share2.Id))
	return new(big.Int).Div(num, denom)
}

// Evaluates lambda parameter for each decrypted share. See second figure in the
// "Share combining" paragraph in [DJK 10], section 5.2.
func (this *ThresholdKey) computeLambda(share *PartialDecryption, shares []*PartialDecryption) *big.Int {
	lambda := this.delta()
	for _, share2 := range shares {
		if share2.Id != share.Id {
			lambda = this.updateLambda(share, share2, lambda)
		}
	}
	return lambda
}

// Used to evaluate c' parameter which combines individual share decryptions.
//
// Modulo division is performed on the computed exponent to avoid creating
// large numbers. This is possible because of the following property of modulo:
// A^B mod C = (A mod C)^B mod C
//
// Modulo division is performed on the computed coefficient because of the
// following property of modulo:
// (AB) mod C = (A mod C * B mod C) mod C
// Note, we need to combine coefficients into single c'.
func (this *ThresholdKey) updateCprime(cprime, lambda *big.Int, share *PartialDecryption) *big.Int {
	twoLambda := new(big.Int).Mul(TWO, lambda)
	ret := this.exp(share.Decryption, twoLambda, this.GetNSquare())
	ret = new(big.Int).Mul(cprime, ret)
	return new(big.Int).Mod(ret, this.GetNSquare())
}

func (this *ThresholdKey) exp(a, b, c *big.Int) *big.Int {
	if b.Cmp(ZERO) == -1 {
		ret := new(big.Int).Exp(a, new(big.Int).Neg(b), c)
		return new(big.Int).ModInverse(ret, c)
	}
	return new(big.Int).Exp(a, b, c)
}

// Executes the last step of message decryption. Takes `cprime` value computed
// from valid shares provided by decryption servers and multiplies this value
// by `combineSharesContant` which is specific to the given public `ThresholdKey`.
func (this *ThresholdKey) computeDecryption(cprime *big.Int) *big.Int {
	l := L(cprime, this.N)
	return new(big.Int).Mod(new(big.Int).Mul(this.combineSharesConstant(), l), this.N)
}

// Combines partial decryptions provided by decryption servers and returns
// decrypted message.
// This function does not verify zero knowledge proofs. Returned message can be
// incorrectly decrypted if an adversary corrupted partial decryption.
func (this *ThresholdKey) CombinePartialDecryptions(shares []*PartialDecryption) (*big.Int, error) {
	if err := this.makeVerificationBeforeCombiningPartialDecryptions(shares); err != nil {
		return nil, err
	}

	cprime := ONE
	for _, share := range shares {
		lambda := this.computeLambda(share, shares)
		cprime = this.updateCprime(cprime, lambda, share)
	}

	return this.computeDecryption(cprime), nil
}

// Combines partial decryptions provided by decription servers and returns
// decrypted message.
// Function verifies zero knowledge proofs and filters out all shares that failed
// verification.
func (this *ThresholdKey) CombinePartialDecryptionsZKP(shares []*PartialDecryptionZKP) (*big.Int, error) {
	ret := make([]*PartialDecryption, 0)
	for _, share := range shares {
		if share.Verify() {
			ret = append(ret, &share.PartialDecryption)
		}
	}
	return this.CombinePartialDecryptions(ret)
}

// Verifies if the decryption of `encryptedMessage` has been done properly.
// It verifies all the zero-knoledge proofs, the value of the encrypted
// and decrypted message. The method returns `nil` if everything is fine.
// Otherwise, it returns an explicative message.
func (this *ThresholdKey) VerifyDecryption(encryptedMessage, decryptedMessage *big.Int, shares []*PartialDecryptionZKP) error {
	for _, share := range shares {
		if share.C.Cmp(encryptedMessage) != 0 {
			return errors.New("The encrypted message is not the same than the one in the shares")
		}
	}
	res, err := this.CombinePartialDecryptionsZKP(shares)
	if err != nil {
		return err
	}
	if res.Cmp(decryptedMessage) != 0 {
		return errors.New("The decrypted message is not the same than the one in the shares")
	}
	return nil
}

// Private key for a threshold Paillier scheme. Holds private information
// for the given decryption server.
// `Id` is the unique identifier of a decryption server and `Share` is a secret
// share generated from hiding polynomial and is used for a partial share decryption.
type ThresholdPrivateKey struct {
	ThresholdKey
	Id    int
	Share *big.Int
}

// Decrypts the cypher text and returns the partial decryption
func (this *ThresholdPrivateKey) Decrypt(c *big.Int) *PartialDecryption {
	ret := new(PartialDecryption)
	ret.Id = this.Id
	exp := new(big.Int).Mul(this.Share, new(big.Int).Mul(TWO, this.delta()))
	ret.Decryption = new(big.Int).Exp(c, exp, this.GetNSquare())

	return ret
}

func (this *ThresholdPrivateKey) copyVi() []*big.Int {
	ret := make([]*big.Int, len(this.Vi))
	for i, vi := range this.Vi {
		ret[i] = new(big.Int).Add(vi, big.NewInt(0))
	}
	return ret
}

func (this *ThresholdPrivateKey) getThresholdKey() *ThresholdKey {
	ret := new(ThresholdKey)
	ret.Threshold = this.Threshold
	ret.TotalNumberOfDecryptionServers = this.TotalNumberOfDecryptionServers
	ret.V = new(big.Int).Add(this.V, big.NewInt(0))
	ret.Vi = this.copyVi()
	ret.N = new(big.Int).Add(this.N, big.NewInt(0))
	return ret
}

func (this *ThresholdPrivateKey) computeZ(r, e *big.Int) *big.Int {
	tmp := new(big.Int).Mul(e, this.delta())
	tmp = new(big.Int).Mul(tmp, this.Share)
	return new(big.Int).Add(r, tmp)
}

func (this *ThresholdPrivateKey) computeHash(a, b, c4, ci2 *big.Int) *big.Int {
	hash := sha256.New()
	hash.Write(a.Bytes())
	hash.Write(b.Bytes())
	hash.Write(c4.Bytes())
	hash.Write(ci2.Bytes())
	return new(big.Int).SetBytes(hash.Sum([]byte{}))
}

func (this *ThresholdPrivateKey) DecryptAndProduceZNP(c *big.Int, random io.Reader) (*PartialDecryptionZKP, error) {
	pd := new(PartialDecryptionZKP)
	pd.Key = this.getThresholdKey()
	pd.C = c
	pd.Id = this.Id
	pd.Decryption = this.Decrypt(c).Decryption

	// choose random number
	r, err := rand.Int(random, this.GetNSquare())
	if err != nil {
		return nil, err
	}
	//  compute a
	c4 := new(big.Int).Exp(c, FOUR, nil)
	a := new(big.Int).Exp(c4, r, this.GetNSquare())

	// compute b
	b := new(big.Int).Exp(this.V, r, this.GetNSquare())

	// compute hash
	ci2 := new(big.Int).Exp(pd.Decryption, big.NewInt(2), nil)

	pd.E = this.computeHash(a, b, c4, ci2)

	pd.Z = this.computeZ(r, pd.E)

	return pd, nil
}

// Verifies if the partial decryption key is well formed.  If well formed,
// the method return nil else an explicative error is returned.
func (this *ThresholdPrivateKey) Validate(random io.Reader) error {
	m, err := rand.Int(random, this.N)
	if err != nil {
		return err
	}
	c, err := this.Encrypt(m, random)
	if err != nil {
		return err
	}
	proof, err := this.DecryptAndProduceZNP(c.C, random)
	if err != nil {
		return err
	}
	if !proof.Verify() {
		return errors.New("invalid share.")
	}
	return nil
}

type PartialDecryption struct {
	Id         int
	Decryption *big.Int
}

// A non-interactive ZKP based on the Fiat–Shamir heuristic. This algorithm
// proves that the decryption server indeed raised secret to his secret exponent
// (`ThresholdPrivateKey.Share`) by comparison with the public verification key
// (`ThresholdKey.Vi`). Recall that v_i = v^(delta s_i).
//
// This is essentialy a protocol for the equality of discrete logs,
// log_{c^4}(c_i^2) = log_v(v_i).
//
// The Fiat-Shamir scheme works as follows in our case:
//
// ZKP construction
//
// - Pick random r
// - Compute E as:
//   E = HASH(a, b, c^4, c_i^2 ), where
//     a = (c^4)^r mod n^2
//     b = V^r mod n^2
//     c is a cyphertext,
//     V is a generator from ThresholdKey.V
//     c_i is a partial decryption for this server
// - Compute Z as:
//    delta * E * s_i + r, where
//      delta is the factorial of the number of decryption servers
//      s_i is a secret share for this server
//
// ZKP verification
//
// - Compute the original a from
//   a = a1 * a2 mod n^2
//   a1 = (c^4)^Z
//   a2 = [ (c_i^2)^E ]^-1 mod n^2
// - Compute the original b from
//   b = b1 * b2 mod n^2
//   b1 = V^Z
//   b2 = [ v_i^E ] -1 mod n^2
// - Rehash H(a, b, c^4, c_i)
// - Compare ZKP hash with the one just computed
type PartialDecryptionZKP struct {
	PartialDecryption
	Key *ThresholdKey // the public key used to encrypt
	E   *big.Int      // the challenge
	Z   *big.Int      // the value needed to check to verify the decryption
	C   *big.Int      // the input cypher text
}

func (this *PartialDecryptionZKP) verifyPart1() *big.Int {
	c4 := new(big.Int).Exp(this.C, FOUR, nil)                  // c^4
	decryption2 := new(big.Int).Exp(this.Decryption, TWO, nil) // c_i^2

	a1 := new(big.Int).Exp(c4, this.Z, this.Key.GetNSquare())          // (c^4)^Z
	a2 := new(big.Int).Exp(decryption2, this.E, this.Key.GetNSquare()) // (c_i^2)^E
	a2 = new(big.Int).ModInverse(a2, this.Key.GetNSquare())
	a := new(big.Int).Mod(new(big.Int).Mul(a1, a2), this.Key.GetNSquare())
	return a
}

func (this *PartialDecryptionZKP) verifyPart2() *big.Int {
	vi := this.Key.Vi[this.Id-1]                                      // servers are indexed from 1
	b1 := new(big.Int).Exp(this.Key.V, this.Z, this.Key.GetNSquare()) // V^Z
	b2 := new(big.Int).Exp(vi, this.E, this.Key.GetNSquare())         // (v_i)^E
	b2 = new(big.Int).ModInverse(b2, this.Key.GetNSquare())
	b := new(big.Int).Mod(new(big.Int).Mul(b1, b2), this.Key.GetNSquare())
	return b
}

func (this *PartialDecryptionZKP) Verify() bool {
	a := this.verifyPart1()
	b := this.verifyPart2()
	hash := sha256.New()
	hash.Write(a.Bytes())
	hash.Write(b.Bytes())
	c4 := new(big.Int).Exp(this.C, FOUR, nil)
	hash.Write(c4.Bytes())
	ci2 := new(big.Int).Exp(this.Decryption, TWO, nil)
	hash.Write(ci2.Bytes())

	expectedE := new(big.Int).SetBytes(hash.Sum([]byte{}))
	return this.E.Cmp(expectedE) == 0
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

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

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

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

package binlog

import (
	crand "crypto/rand"
	"fmt"
	"math"
	"math/big"
	"sync"

	"golang.org/x/net/context"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/vt/dbconfigs"
	"vitess.io/vitess/go/vt/log"
)

var (
	// ErrBinlogUnavailable is returned by this library when we
	// cannot find a suitable binlog to satisfy the request.
	ErrBinlogUnavailable = fmt.Errorf("cannot find relevant binlogs on this server")
)

// SlaveConnection represents a connection to mysqld that pretends to be a slave
// connecting for replication. Each such connection must identify itself to
// mysqld with a server ID that is unique both among other SlaveConnections and
// among actual slaves in the topology.
type SlaveConnection struct {
	*mysql.Conn
	cp      dbconfigs.Connector
	slaveID uint32
	cancel  context.CancelFunc
	wg      sync.WaitGroup
}

// We use a random slaveid deprecating IDPool which was causing problems with RDS
// either because it was using ids in the same range that we generate or because
// the pool reuses ids.
func getSlaveId() uint32 {
	fmt.Printf("In getSlaveId\n")
	max := big.NewInt(math.MaxInt32)
	id, err := crand.Int(crand.Reader, max)
	if err != nil {
		fmt.Printf("Error in getting slaveid %v\n", err)

		panic(fmt.Sprintf("Could not allocate slave id"))
	}
	fmt.Printf("getSlaveId returning %d\n", uint32(id.Int64()))
	return uint32(id.Int64())
}

// NewSlaveConnection creates a new slave connection to the mysqld instance.
func NewSlaveConnection(cp dbconfigs.Connector) (*SlaveConnection, error) {
	conn, err := connectForReplication(cp)
	if err != nil {
		return nil, err
	}

	sc := &SlaveConnection{
		Conn:    conn,
		cp:      cp,
		slaveID: getSlaveId(),
	}
	log.Infof("new slave connection: slaveID=%d", sc.slaveID)
	return sc, nil
}

// connectForReplication create a MySQL connection ready to use for replication.
func connectForReplication(cp dbconfigs.Connector) (*mysql.Conn, error) {
	ctx := context.Background()
	conn, err := cp.Connect(ctx)
	if err != nil {
		return nil, err
	}
	// Tell the server that we understand the format of events
	// that will be used if binlog_checksum is enabled on the server.
	if _, err := conn.ExecuteFetch("SET @master_binlog_checksum=@@global.binlog_checksum", 0, false); err != nil {
		return nil, fmt.Errorf("failed to set @master_binlog_checksum=@@global.binlog_checksum: %v", err)
	}

	return conn, nil
}

// StartBinlogDumpFromCurrent requests a replication binlog dump from
// the current position.
func (sc *SlaveConnection) StartBinlogDumpFromCurrent(ctx context.Context) (mysql.Position, <-chan mysql.BinlogEvent, error) {
	ctx, sc.cancel = context.WithCancel(ctx)

	masterPosition, err := sc.Conn.MasterPosition()
	if err != nil {
		return mysql.Position{}, nil, fmt.Errorf("failed to get master position: %v", err)
	}

	c, err := sc.StartBinlogDumpFromPosition(ctx, masterPosition)
	return masterPosition, c, err
}

// StartBinlogDumpFromPosition requests a replication binlog dump from
// the master mysqld at the given Position and then sends binlog
// events to the provided channel.
// The stream will continue in the background, waiting for new events if
// necessary, until the connection is closed, either by the master or
// by canceling the context.
//
// Note the context is valid and used until eventChan is closed.
func (sc *SlaveConnection) StartBinlogDumpFromPosition(ctx context.Context, startPos mysql.Position) (<-chan mysql.BinlogEvent, error) {
	ctx, sc.cancel = context.WithCancel(ctx)

	log.Infof("sending binlog dump command: startPos=%v, slaveID=%v", startPos, sc.slaveID)
	if err := sc.SendBinlogDumpCommand(sc.slaveID, startPos); err != nil {
		log.Errorf("couldn't send binlog dump command: %v", err)
		return nil, err
	}

	return sc.streamEvents(ctx), nil
}

// streamEvents returns a channel on which events are streamed.
func (sc *SlaveConnection) streamEvents(ctx context.Context) chan mysql.BinlogEvent {
	// FIXME(alainjobart) I think we can use a buffered channel for better performance.
	eventChan := make(chan mysql.BinlogEvent)

	// Start reading events.
	sc.wg.Add(1)
	go func() {
		defer func() {
			close(eventChan)
			sc.wg.Done()
		}()
		for {
			event, err := sc.Conn.ReadBinlogEvent()
			if err != nil {
				if sqlErr, ok := err.(*mysql.SQLError); ok && sqlErr.Number() == mysql.CRServerLost {
					// CRServerLost = Lost connection to MySQL server during query
					// This is not necessarily an error. It could just be that we closed
					// the connection from outside.
					log.Infof("connection closed during binlog stream (possibly intentional): %v", err)
					return
				}
				log.Errorf("read error while streaming binlog events: %v", err)
				return
			}

			select {
			case eventChan <- event:
			case <-ctx.Done():
				return
			}
		}
	}()
	return eventChan
}

// StartBinlogDumpFromBinlogBeforeTimestamp requests a replication
// binlog dump from the master mysqld starting with a file that has
// timestamps smaller than the provided timestamp, and then sends
// binlog events to the provided channel.
//
// The startup phase will list all the binary logs, and find the one
// that has events starting strictly before the provided timestamp. It
// will then start from there, and stream all events. It is the
// responsibility of the calling site to filter the events more.
//
// MySQL 5.6+ note: we need to do it that way because of the way the
// GTIDSet works. In the previous two streaming functions, we pass in
// the full GTIDSet (that has the list of all transactions seen in
// the replication stream). In this case, we don't know it, all we
// have is the binlog file names. We depend on parsing the first
// PREVIOUS_GTIDS_EVENT event in the logs to get it. So we need the
// caller to parse that event, and it can't be skipped because its
// timestamp is lower. Then, for each subsequent event, the caller
// also needs to add the event GTID to its GTIDSet. Otherwise it won't
// be correct ever. So the caller really needs to build up its GTIDSet
// along the entire file, not just for events whose timestamp is in a
// given range.
//
// The stream will continue in the background, waiting for new events if
// necessary, until the connection is closed, either by the master or
// by canceling the context.
//
// Note the context is valid and used until eventChan is closed.
func (sc *SlaveConnection) StartBinlogDumpFromBinlogBeforeTimestamp(ctx context.Context, timestamp int64) (<-chan mysql.BinlogEvent, error) {
	ctx, sc.cancel = context.WithCancel(ctx)

	filename, err := sc.findFileBeforeTimestamp(ctx, timestamp)
	if err != nil {
		return nil, err
	}

	// Start dumping the logs. The position is '4' to skip the
	// Binlog File Header. See this page for more info:
	// https://dev.mysql.com/doc/internals/en/binlog-file.html
	if err := sc.Conn.WriteComBinlogDump(sc.slaveID, filename, 4, 0); err != nil {
		return nil, fmt.Errorf("failed to send the ComBinlogDump command: %v", err)
	}
	return sc.streamEvents(ctx), nil
}

func (sc *SlaveConnection) findFileBeforeTimestamp(ctx context.Context, timestamp int64) (filename string, err error) {
	// List the binlogs.
	binlogs, err := sc.Conn.ExecuteFetch("SHOW BINARY LOGS", 1000, false)
	if err != nil {
		return "", fmt.Errorf("failed to SHOW BINARY LOGS: %v", err)
	}

	// Start with the most recent binlog file until we find the right event.
	for binlogIndex := len(binlogs.Rows) - 1; binlogIndex >= 0; binlogIndex-- {
		// Exit the loop early if context is canceled.
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		default:
		}

		filename := binlogs.Rows[binlogIndex][0].ToString()
		blTimestamp, err := sc.getBinlogTimeStamp(filename)
		if err != nil {
			return "", err
		}
		if blTimestamp < timestamp {
			// The binlog timestamp is older: we've found a good starting point.
			return filename, nil
		}
	}

	log.Errorf("couldn't find an old enough binlog to match timestamp >= %v (looked at %v files)", timestamp, len(binlogs.Rows))
	return "", ErrBinlogUnavailable
}

func (sc *SlaveConnection) getBinlogTimeStamp(filename string) (blTimestamp int64, err error) {
	conn, err := connectForReplication(sc.cp)
	if err != nil {
		return 0, err
	}
	defer conn.Close()

	if err := conn.WriteComBinlogDump(sc.slaveID, filename, 4, 0); err != nil {
		return 0, fmt.Errorf("failed to send the ComBinlogDump command: %v", err)
	}

	// Get the first event to get its timestamp. We skip
	// events that don't have timestamps (although it seems
	// most do anyway).
	for {
		event, err := conn.ReadBinlogEvent()
		if err != nil {
			return 0, fmt.Errorf("error reading binlog event %v: %v", filename, err)
		}
		if !event.IsValid() {
			return 0, fmt.Errorf("first event from binlog %v is not valid", filename)
		}
		if ts := event.Timestamp(); ts > 0 {
			return int64(ts), nil
		}
	}
}

// Close closes the slave connection, which also signals an ongoing dump
// started with StartBinlogDump() to stop and close its BinlogEvent channel.
// The ID for the slave connection is recycled back into the pool.
func (sc *SlaveConnection) Close() {
	if sc.Conn != nil {
		log.Infof("closing slave socket to unblock reads")
		sc.Conn.Close()

		// sc.cancel is set at the beginning of the StartBinlogDump*
		// methods. If we error out before then, it's nil.
		// Note we also may error out before adding 1 to sc.wg,
		// but then the Wait() still works.
		if sc.cancel != nil {
			log.Infof("waiting for slave dump thread to end")
			sc.cancel()
			sc.wg.Wait()
			sc.cancel = nil
		}

		log.Infof("closing slave MySQL client with slaveID %v", sc.slaveID)
		sc.Conn = nil
	}
}
<commit_msg>Use random slaveids instead of using the current idpool<commit_after>/*
Copyright 2019 The Vitess Authors.

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

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

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

package binlog

import (
	crand "crypto/rand"
	"fmt"
	"math"
	"math/big"
	"sync"

	"golang.org/x/net/context"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/vt/dbconfigs"
	"vitess.io/vitess/go/vt/log"
)

var (
	// ErrBinlogUnavailable is returned by this library when we
	// cannot find a suitable binlog to satisfy the request.
	ErrBinlogUnavailable = fmt.Errorf("cannot find relevant binlogs on this server")
)

// SlaveConnection represents a connection to mysqld that pretends to be a slave
// connecting for replication. Each such connection must identify itself to
// mysqld with a server ID that is unique both among other SlaveConnections and
// among actual slaves in the topology.
type SlaveConnection struct {
	*mysql.Conn
	cp      dbconfigs.Connector
	slaveID uint32
	cancel  context.CancelFunc
	wg      sync.WaitGroup
}

// We use a random slaveid deprecating IDPool which was causing problems with RDS
// either because it was using ids in the same range that we generate or because
// the pool reuses ids.
func getSlaveID() uint32 {
	fmt.Printf("In getSlaveId\n")
	max := big.NewInt(math.MaxInt32)
	id, err := crand.Int(crand.Reader, max)
	if err != nil {
		fmt.Printf("Error in getting slaveid %v\n", err)

		panic(fmt.Sprintf("Could not allocate slave id"))
	}
	fmt.Printf("getSlaveId returning %d\n", uint32(id.Int64()))
	return uint32(id.Int64())
}

// NewSlaveConnection creates a new slave connection to the mysqld instance.
func NewSlaveConnection(cp dbconfigs.Connector) (*SlaveConnection, error) {
	conn, err := connectForReplication(cp)
	if err != nil {
		return nil, err
	}

	sc := &SlaveConnection{
		Conn:    conn,
		cp:      cp,
		slaveID: getSlaveID(),
	}
	log.Infof("new slave connection: slaveID=%d", sc.slaveID)
	return sc, nil
}

// connectForReplication create a MySQL connection ready to use for replication.
func connectForReplication(cp dbconfigs.Connector) (*mysql.Conn, error) {
	ctx := context.Background()
	conn, err := cp.Connect(ctx)
	if err != nil {
		return nil, err
	}
	// Tell the server that we understand the format of events
	// that will be used if binlog_checksum is enabled on the server.
	if _, err := conn.ExecuteFetch("SET @master_binlog_checksum=@@global.binlog_checksum", 0, false); err != nil {
		return nil, fmt.Errorf("failed to set @master_binlog_checksum=@@global.binlog_checksum: %v", err)
	}

	return conn, nil
}

// StartBinlogDumpFromCurrent requests a replication binlog dump from
// the current position.
func (sc *SlaveConnection) StartBinlogDumpFromCurrent(ctx context.Context) (mysql.Position, <-chan mysql.BinlogEvent, error) {
	ctx, sc.cancel = context.WithCancel(ctx)

	masterPosition, err := sc.Conn.MasterPosition()
	if err != nil {
		return mysql.Position{}, nil, fmt.Errorf("failed to get master position: %v", err)
	}

	c, err := sc.StartBinlogDumpFromPosition(ctx, masterPosition)
	return masterPosition, c, err
}

// StartBinlogDumpFromPosition requests a replication binlog dump from
// the master mysqld at the given Position and then sends binlog
// events to the provided channel.
// The stream will continue in the background, waiting for new events if
// necessary, until the connection is closed, either by the master or
// by canceling the context.
//
// Note the context is valid and used until eventChan is closed.
func (sc *SlaveConnection) StartBinlogDumpFromPosition(ctx context.Context, startPos mysql.Position) (<-chan mysql.BinlogEvent, error) {
	ctx, sc.cancel = context.WithCancel(ctx)

	log.Infof("sending binlog dump command: startPos=%v, slaveID=%v", startPos, sc.slaveID)
	if err := sc.SendBinlogDumpCommand(sc.slaveID, startPos); err != nil {
		log.Errorf("couldn't send binlog dump command: %v", err)
		return nil, err
	}

	return sc.streamEvents(ctx), nil
}

// streamEvents returns a channel on which events are streamed.
func (sc *SlaveConnection) streamEvents(ctx context.Context) chan mysql.BinlogEvent {
	// FIXME(alainjobart) I think we can use a buffered channel for better performance.
	eventChan := make(chan mysql.BinlogEvent)

	// Start reading events.
	sc.wg.Add(1)
	go func() {
		defer func() {
			close(eventChan)
			sc.wg.Done()
		}()
		for {
			event, err := sc.Conn.ReadBinlogEvent()
			if err != nil {
				if sqlErr, ok := err.(*mysql.SQLError); ok && sqlErr.Number() == mysql.CRServerLost {
					// CRServerLost = Lost connection to MySQL server during query
					// This is not necessarily an error. It could just be that we closed
					// the connection from outside.
					log.Infof("connection closed during binlog stream (possibly intentional): %v", err)
					return
				}
				log.Errorf("read error while streaming binlog events: %v", err)
				return
			}

			select {
			case eventChan <- event:
			case <-ctx.Done():
				return
			}
		}
	}()
	return eventChan
}

// StartBinlogDumpFromBinlogBeforeTimestamp requests a replication
// binlog dump from the master mysqld starting with a file that has
// timestamps smaller than the provided timestamp, and then sends
// binlog events to the provided channel.
//
// The startup phase will list all the binary logs, and find the one
// that has events starting strictly before the provided timestamp. It
// will then start from there, and stream all events. It is the
// responsibility of the calling site to filter the events more.
//
// MySQL 5.6+ note: we need to do it that way because of the way the
// GTIDSet works. In the previous two streaming functions, we pass in
// the full GTIDSet (that has the list of all transactions seen in
// the replication stream). In this case, we don't know it, all we
// have is the binlog file names. We depend on parsing the first
// PREVIOUS_GTIDS_EVENT event in the logs to get it. So we need the
// caller to parse that event, and it can't be skipped because its
// timestamp is lower. Then, for each subsequent event, the caller
// also needs to add the event GTID to its GTIDSet. Otherwise it won't
// be correct ever. So the caller really needs to build up its GTIDSet
// along the entire file, not just for events whose timestamp is in a
// given range.
//
// The stream will continue in the background, waiting for new events if
// necessary, until the connection is closed, either by the master or
// by canceling the context.
//
// Note the context is valid and used until eventChan is closed.
func (sc *SlaveConnection) StartBinlogDumpFromBinlogBeforeTimestamp(ctx context.Context, timestamp int64) (<-chan mysql.BinlogEvent, error) {
	ctx, sc.cancel = context.WithCancel(ctx)

	filename, err := sc.findFileBeforeTimestamp(ctx, timestamp)
	if err != nil {
		return nil, err
	}

	// Start dumping the logs. The position is '4' to skip the
	// Binlog File Header. See this page for more info:
	// https://dev.mysql.com/doc/internals/en/binlog-file.html
	if err := sc.Conn.WriteComBinlogDump(sc.slaveID, filename, 4, 0); err != nil {
		return nil, fmt.Errorf("failed to send the ComBinlogDump command: %v", err)
	}
	return sc.streamEvents(ctx), nil
}

func (sc *SlaveConnection) findFileBeforeTimestamp(ctx context.Context, timestamp int64) (filename string, err error) {
	// List the binlogs.
	binlogs, err := sc.Conn.ExecuteFetch("SHOW BINARY LOGS", 1000, false)
	if err != nil {
		return "", fmt.Errorf("failed to SHOW BINARY LOGS: %v", err)
	}

	// Start with the most recent binlog file until we find the right event.
	for binlogIndex := len(binlogs.Rows) - 1; binlogIndex >= 0; binlogIndex-- {
		// Exit the loop early if context is canceled.
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		default:
		}

		filename := binlogs.Rows[binlogIndex][0].ToString()
		blTimestamp, err := sc.getBinlogTimeStamp(filename)
		if err != nil {
			return "", err
		}
		if blTimestamp < timestamp {
			// The binlog timestamp is older: we've found a good starting point.
			return filename, nil
		}
	}

	log.Errorf("couldn't find an old enough binlog to match timestamp >= %v (looked at %v files)", timestamp, len(binlogs.Rows))
	return "", ErrBinlogUnavailable
}

func (sc *SlaveConnection) getBinlogTimeStamp(filename string) (blTimestamp int64, err error) {
	conn, err := connectForReplication(sc.cp)
	if err != nil {
		return 0, err
	}
	defer conn.Close()

	if err := conn.WriteComBinlogDump(sc.slaveID, filename, 4, 0); err != nil {
		return 0, fmt.Errorf("failed to send the ComBinlogDump command: %v", err)
	}

	// Get the first event to get its timestamp. We skip
	// events that don't have timestamps (although it seems
	// most do anyway).
	for {
		event, err := conn.ReadBinlogEvent()
		if err != nil {
			return 0, fmt.Errorf("error reading binlog event %v: %v", filename, err)
		}
		if !event.IsValid() {
			return 0, fmt.Errorf("first event from binlog %v is not valid", filename)
		}
		if ts := event.Timestamp(); ts > 0 {
			return int64(ts), nil
		}
	}
}

// Close closes the slave connection, which also signals an ongoing dump
// started with StartBinlogDump() to stop and close its BinlogEvent channel.
// The ID for the slave connection is recycled back into the pool.
func (sc *SlaveConnection) Close() {
	if sc.Conn != nil {
		log.Infof("closing slave socket to unblock reads")
		sc.Conn.Close()

		// sc.cancel is set at the beginning of the StartBinlogDump*
		// methods. If we error out before then, it's nil.
		// Note we also may error out before adding 1 to sc.wg,
		// but then the Wait() still works.
		if sc.cancel != nil {
			log.Infof("waiting for slave dump thread to end")
			sc.cancel()
			sc.wg.Wait()
			sc.cancel = nil
		}

		log.Infof("closing slave MySQL client with slaveID %v", sc.slaveID)
		sc.Conn = nil
	}
}
<|endoftext|>


<commit_before>package loadbalancer

import (
	"context"

	"github.com/aws/aws-sdk-go/service/elb"
	"github.com/giantswarm/apiextensions/pkg/apis/provider/v1alpha1"
	"github.com/giantswarm/microerror"

	"github.com/giantswarm/aws-operator/service/awsconfig/v5/key"
)

const (
	cloudProviderClusterTagValue = "owned"
	cloudProviderServiceTagKey   = "kubernetes.io/service-name"
)

func (r *Resource) GetCurrentState(ctx context.Context, obj interface{}) (interface{}, error) {
	customObject, err := key.ToCustomObject(obj)
	if err != nil {
		return nil, microerror.Mask(err)
	}

	currentState, err := r.clusterLoadBalancers(customObject)
	if err != nil {
		return nil, microerror.Mask(err)
	}

	return currentState, nil
}

func (r *Resource) clusterLoadBalancers(customObject v1alpha1.AWSConfig) (*LoadBalancerState, error) {
	lbState := &LoadBalancerState{}

	// We get all load balancers because the API does not allow tag filters.
	output, err := r.clients.ELB.DescribeLoadBalancers(&elb.DescribeLoadBalancersInput{})
	if err != nil {
		return nil, microerror.Mask(err)
	}

	allLBNames := []*string{}
	for _, lb := range output.LoadBalancerDescriptions {
		allLBNames = append(allLBNames, lb.LoadBalancerName)
	}

	// We get tags for all load balancers which needs a separate API call.
	tagsInput := &elb.DescribeTagsInput{
		LoadBalancerNames: allLBNames,
	}
	tagsOutput, err := r.clients.ELB.DescribeTags(tagsInput)
	if err != nil {
		return nil, microerror.Mask(err)
	}

	// We filter based on the AWS cloud provider tags to find load balancers
	// associated with the cluster being processed.
	clusterLBNames := []string{}
	for _, lb := range tagsOutput.TagDescriptions {
		if containsClusterTag(lb.Tags, customObject) && containsServiceTag(lb.Tags) {
			clusterLBNames = append(clusterLBNames, *lb.LoadBalancerName)
		}
	}

	lbState.LoadBalancerNames = clusterLBNames

	return lbState, nil
}

func containsClusterTag(tags []*elb.Tag, customObject v1alpha1.AWSConfig) bool {
	tagKey := key.ClusterCloudProviderTag(customObject)

	for _, tag := range tags {
		if *tag.Key == tagKey && *tag.Value == cloudProviderClusterTagValue {
			return true
		}
	}
	return false
}

func containsServiceTag(tags []*elb.Tag) bool {
	for _, tag := range tags {
		if *tag.Key == cloudProviderServiceTagKey {
			return true
		}
	}

	return false
}
<commit_msg>Only call DescribeTags if there are load balancers (#761)<commit_after>package loadbalancer

import (
	"context"

	"github.com/aws/aws-sdk-go/service/elb"
	"github.com/giantswarm/apiextensions/pkg/apis/provider/v1alpha1"
	"github.com/giantswarm/microerror"

	"github.com/giantswarm/aws-operator/service/awsconfig/v5/key"
)

const (
	cloudProviderClusterTagValue = "owned"
	cloudProviderServiceTagKey   = "kubernetes.io/service-name"
)

func (r *Resource) GetCurrentState(ctx context.Context, obj interface{}) (interface{}, error) {
	customObject, err := key.ToCustomObject(obj)
	if err != nil {
		return nil, microerror.Mask(err)
	}

	currentState, err := r.clusterLoadBalancers(customObject)
	if err != nil {
		return nil, microerror.Mask(err)
	}

	return currentState, nil
}

func (r *Resource) clusterLoadBalancers(customObject v1alpha1.AWSConfig) (*LoadBalancerState, error) {
	lbState := &LoadBalancerState{}
	clusterLBNames := []string{}

	// We get all load balancers because the API does not allow tag filters.
	output, err := r.clients.ELB.DescribeLoadBalancers(&elb.DescribeLoadBalancersInput{})
	if err != nil {
		return nil, microerror.Mask(err)
	}

	allLBNames := []*string{}
	for _, lb := range output.LoadBalancerDescriptions {
		allLBNames = append(allLBNames, lb.LoadBalancerName)
	}

	if len(allLBNames) > 0 {
		// We get tags for all load balancers which needs a separate API call.
		tagsInput := &elb.DescribeTagsInput{
			LoadBalancerNames: allLBNames,
		}
		tagsOutput, err := r.clients.ELB.DescribeTags(tagsInput)
		if err != nil {
			return nil, microerror.Mask(err)
		}

		// We filter based on the AWS cloud provider tags to find load balancers
		// associated with the cluster being processed.
		for _, lb := range tagsOutput.TagDescriptions {
			if containsClusterTag(lb.Tags, customObject) && containsServiceTag(lb.Tags) {
				clusterLBNames = append(clusterLBNames, *lb.LoadBalancerName)
			}
		}
	}

	lbState.LoadBalancerNames = clusterLBNames

	return lbState, nil
}

func containsClusterTag(tags []*elb.Tag, customObject v1alpha1.AWSConfig) bool {
	tagKey := key.ClusterCloudProviderTag(customObject)

	for _, tag := range tags {
		if *tag.Key == tagKey && *tag.Value == cloudProviderClusterTagValue {
			return true
		}
	}
	return false
}

func containsServiceTag(tags []*elb.Tag) bool {
	for _, tag := range tags {
		if *tag.Key == cloudProviderServiceTagKey {
			return true
		}
	}

	return false
}
<|endoftext|>


<commit_before>package controller

import (
  "encoding/json"
  "fmt"
  "time"
  "config"
  "strings"
  "strconv"
  "models"
  "helpers"
  "doggo"
  "io/ioutil"
  "github.com/gocql/gocql"
  request "github.com/DronRathore/goexpress/request"
  response "github.com/DronRathore/goexpress/response"
  express "github.com/DronRathore/goexpress"
)

///////////////// Webhooks ////////////////////////

/*
  Updates Sendgrid status
  Sendgrid is by far the best way to handle emails
  The status varies with granular details which helps in taking
  decisions properly. Sendgrid posts an array of JSON events.
*/
func BatchUpdateSGStatus(req *request.Request, res *response.Response, next func()) {
  username, password, ok := req.GetRaw().BasicAuth()
  if !ok {
    res.Header.SetStatus(403)
    res.JSON(map[string]interface{}{
      "error": "Access denied",
    })
    return
  }
  // Validate Auth headers
  if username != config.Configuration.Sendgrid.Username || password != config.Configuration.Sendgrid.Password {
    res.Header.SetStatus(403)
    res.JSON(map[string]interface{}{
      "error": "Access denied",
    })
    return
  }
  // valid request, lets process the event
  // read and parse the JSON body
  data, err := ioutil.ReadAll(req.GetRaw().Body)
  if err != nil {
    res.Header.SetStatus(400)
    doggo.DoggoEvent("Cannot read sendgrid JSON", err, true)
    return
  }
  var sendgridJson []map[string]interface{}
  err = json.Unmarshal(data, &sendgridJson)
  if err != nil {
    res.Header.SetStatus(400)
    doggo.DoggoEvent("Parse sendgrid JSON failed", err, true)
    return
  }
  // empty events array pushed
  if len(sendgridJson) == 0 {
    res.Header.SetStatus(200)
    res.End()
    return
  }
  // Return the response to SG, and then process the events
  // so as to avoid timeouts
  res.Write("Success").End()

  // loop over all the events that are pushed
  for _, event := range sendgridJson {
    var batchId gocql.UUID
    var logId gocql.UUID
    var emptyId gocql.UUID
    var timestamp time.Time = time.Now()

    eventName := strings.ToLower(event["event"].(string))
    if event["batch_id"] == nil {
      // break, we expect a batch id
    } else {
      batchId, _ = gocql.ParseUUID(event["batch_id"].(string))
    }

    if event["log_id"] != nil {
      logId, _ = gocql.ParseUUID(event["log_id"].(string))
    }

    var log *models.Log = nil
    var shouldRetry bool = false
    var updateObject map[string]interface{}
    switch(eventName){
      case "processed" :
        updateObject = map[string]interface{}{"state": "sent"}
        break;
      case "delivered" :
        updateObject = map[string]interface{}{"state": "delivered"}
      break;
      case "deferred"  :
        updateObject = map[string]interface{}{"state": "deferred"}
        shouldRetry = false
        break;
      case "bounce"    :
        updateObject = map[string]interface{}{"state": "bounced"}
        shouldRetry = true
      break;
      case "dropped"   :
        updateObject = map[string]interface{}{"state": "dropped"}
        shouldRetry = false
      break;
      default: continue
    }

    if event["timestamp"] != nil {
      timestamp = time.Unix(int64(event["timestamp"].(float64)), 0)
    }
    // set the updated timestamp
    updateObject["last_update"] = timestamp

    batch := &models.Batch{Batch_id: batchId}
    success, err := batch.Exists()
    if err != nil {
      // cannot process the event
      // todo: Push to a DB rabbit queue
      doggo.DoggoEvent("Batch Search failed", "BatchId:" + batchId.String() + "\n" + err.Error() , true)
      continue
    }

    if (success == false){
      // batch doesn't exists
      // Ack the request and move on
      continue
    }
    // fetch the email log doc
    if logId != emptyId {
      log = &models.Log{Batch_id: batchId, Log_id: logId}
      success, _ := log.Exists()
      if success == true {
        // this email is done
        // mark sendgrid as been used
        var status = log.GetMap()
        status["sendgrid"] = true
        updateObject["status"] = status
        if shouldRetry == false {
          // Lets update the status and exit
          log.Update(updateObject)
          continue
       }
      } else {
        // drop the packet
        // log_id doesn't exists, dangling packet
        continue
     }
    }
    // this is a retrial/first info event request
    if log == nil {
      // create a log
      fmt.Println("Creating a log")
      log = &models.Log{}
      _, err := log.Create(map[string]interface{}{
        "state"  : eventName,
        "status" : map[string]bool{"sendgrid": true},
        "batch_id": batchId,
        "user_id": batch.User_id,
        "email": event["email"].(string),
        "created_at": time.Now(),
        "last_update": timestamp,
      })
      fmt.Println("Error=>", err)
      if err != nil {
        doggo.DoggoEvent("Log Creation failed", "BatchId:" + batchId.String() + "\n" + err.Error() , true)
      }
    }
    fmt.Println("Log Created")
    // successfully delivered the packet, save the log
    // and continue
    if shouldRetry == false {
      doggo.AddDoggoMetric("sendgrid.success")
      helpers.UpdateStats(batch.User_id, "success")
      _, err := log.Update(updateObject)
      fmt.Println("Log Updated", err)
      continue
    }
    fmt.Println("Trying the next worker")
    // else our mail has been dropped lets see what we can do
    // pick the one which we haven't tried
    doggo.AddDoggoMetric("sendgrid.failed")
    helpers.UpdateStats(batch.User_id, "failed")
    worker := helpers.NextChannelToTry(log.Status)
    if worker != nil {
      // use this worker to retry sending the email
      routingKey := "log." + worker.Namespace + log.Log_id.String()
      fmt.Println("RoutingKey", routingKey)
      // push it into the queue
      worker.Channel.Publish(routingKey, nil)
    } else {
      // exhausted all the trials for this email
      // drop it permanently
      fmt.Println("Le'me just update and exit")
      log.Update(updateObject)
    }
  }
}

/*
  MailGun Webhook
  Mailgun service only provides delivered, bounced and hardfail status
  so we can only take actions in case of a bounced email only.
*/
func BatchUpdateMGStatus(req *request.Request, res *response.Response, next func()) {
  var batchId gocql.UUID
  var logId gocql.UUID
  var emptyId gocql.UUID
  var email string
  var updateObject map[string]interface{}
  var err error
  var batch *models.Batch
  var log *models.Log = nil
  var success bool
  var boundary string
  var timestamp time.Time
  // Mailgun has this strange way of sending status updates
  // if it is a failed event than you get a form-data else a simple
  // post form is sent in case of delivered status
  if helpers.IsMultipart(req.Header["content-type"], &boundary) {
    form := helpers.ReadMultiPartForm(boundary, req.GetRaw().Body)
    if form == nil {
      res.Header.SetStatus(400)
      res.Write("Bad Request").End()
      return
    }
    // inject the form-data values
    req.Body = form.Value
  }
  // validate the request, should have valid auth header
  if helpers.IsValidMGRequest(req) == false {
    res.Header.SetStatus(403)
    res.Write("Access Denied").End()
    return
  }
  // simple checks
  if len(req.Body["event"]) == 0 || len(req.Body["recipient"]) == 0 || len(req.Body["batch-id"]) == 0 {
    res.Header.SetStatus(400)
    res.Write("Bad Request").End()
    return
  }
  // Cool, it is a valid webhook request from MG
  // process it.
  batch_id_data := req.Body["batch-id"][0]
  if batch_id_data == "" {
    // batch_id not present, drop the request
    res.Header.SetStatus(400)
    res.Write("Bad Request").End()
    return
  }
  // parse the batch_id
  batchId, err = gocql.ParseUUID(batch_id_data)
  if err != nil {
    // we messed up! Why would you send a bad JSON dron?
    res.Header.SetStatus(400)
    res.Write("Bad Request").End()
    return
  }
  // see if we have a log_id
  if len(req.Body["log-id"]) == 1 {
    logId, err = gocql.ParseUUID(req.Body["log-id"][0])
  }
  email = req.Body["recipient"][0]
  var eventName = req.Body["event"][0]
  var shouldRetry bool = false
  // send success to the client and then process it
  res.Write("success").End()

  switch(eventName){
    case "delivered" :
      updateObject = map[string]interface{}{"state": "delivered"}
    break;
    case "bounced"    :
      updateObject = map[string]interface{}{"state": "bounced"}
      shouldRetry = true
    break;
    case "hardfail"   :
      updateObject = map[string]interface{}{"state": "dropped"}
      shouldRetry = false
    break;
   case "dropped":
      updateObject = map[string]interface{}{"state": "dropped"}
      shouldRetry = false
    default: return // we can add further click/open events but OOS
  }
  // retrieve the timestamp
  if len(req.Body["timestamp"]) != 0 {
    val, _ := strconv.Atoi(req.Body["timestamp"][0])
    timestamp = time.Unix(int64(val), 0)
  }
  // set the updated timestamp
  updateObject["last_update"] = timestamp

  batch = &models.Batch{Batch_id: batchId}
  success, err = batch.Exists()
  if err != nil {
    // cannot process the event
    // todo: Push to a DB rabbit queue
    doggo.DoggoEvent("Batch Search failed", "BatchId:" + batchId.String() + "\n" + err.Error() , true)
    return
  }

  if (success == false){
    // batch doesn't exists
    // Ack the request and move on
    return
  }
  // fetch the email log doc
  if logId != emptyId {
    log = &models.Log{Batch_id: batchId, Log_id: logId}
    success, _ = log.Exists()
    if success == true {
      // this email is done
      // mark mailgun as been used
      var status = log.GetMap()
      status["mailgun"] = true
      updateObject["status"] = status
      if shouldRetry == false {
        // Lets update the status and exit
        log.Update(updateObject)
        return
     }
    } else {
      // drop the packet
      // log_id doesn't exists, dangling packet
      return
   }
  }
  // this is a retrial mail
  if log == nil {
    // create a log
    log = &models.Log{}
    _, err = log.Create(map[string]interface{}{
      "state"  : eventName,
      "status" : map[string]bool{"mailgun": true},
      "batch_id": batchId,
      "user_id": batch.User_id,
      "email": email,
      "created_at": time.Now(),
      "last_update": timestamp,
    })
    if err != nil {
      doggo.DoggoEvent("Log Creation failed", "BatchId:" + batchId.String() + "\n" + err.Error() , true)
    }
  }
  // successfully delivered the packet, save the log
  // and return
  if shouldRetry == false {
    helpers.UpdateStats(batch.User_id, "success")
    doggo.AddDoggoMetric("mailgun.success")
    _, err = log.Update(updateObject)
    return
  }
  fmt.Println("Trying the next worker")
  doggo.AddDoggoMetric("mailgun.failed")
  helpers.UpdateStats(batch.User_id, "failed")
  // else our mail has been dropped lets see what we can do
  // pick the one which we haven't tried
  worker := helpers.NextChannelToTry(log.Status)
  if worker != nil {
    // use this worker to retry sending the email
    routingKey := "log." + worker.Namespace + log.Log_id.String()
    fmt.Println("RoutingKey", routingKey)
    // push it into the queue
    worker.Channel.Publish(routingKey, nil)
  } else {
    // exhausted all the trials for this email
    // drop it permanently
    fmt.Println("Le'me just update and exit")
    log.Update(updateObject)
  }
}

var WebhookController = func() interface{} {
  var Router = express.Router()
  Router.Post("/api/webhook/sendgrid", BatchUpdateSGStatus)
  Router.Post("/api/webhook/mailgun", BatchUpdateMGStatus)
  return *Router
}()
<commit_msg>fix for webhook<commit_after>package controller

import (
  "encoding/json"
  "fmt"
  "time"
  "config"
  "strings"
  "strconv"
  "models"
  "helpers"
  "doggo"
  "io/ioutil"
  "github.com/gocql/gocql"
  request "github.com/DronRathore/goexpress/request"
  response "github.com/DronRathore/goexpress/response"
  express "github.com/DronRathore/goexpress"
)

///////////////// Webhooks ////////////////////////

/*
  Updates Sendgrid status
  Sendgrid is by far the best way to handle emails
  The status varies with granular details which helps in taking
  decisions properly. Sendgrid posts an array of JSON events.
*/
func BatchUpdateSGStatus(req *request.Request, res *response.Response, next func()) {
  username, password, ok := req.GetRaw().BasicAuth()
  if !ok {
    res.Header.SetStatus(403)
    res.JSON(map[string]interface{}{
      "error": "Access denied",
    })
    return
  }
  // Validate Auth headers
  if username != config.Configuration.Sendgrid.Username || password != config.Configuration.Sendgrid.Password {
    res.Header.SetStatus(403)
    res.JSON(map[string]interface{}{
      "error": "Access denied",
    })
    return
  }
  // valid request, lets process the event
  // read and parse the JSON body
  data, err := ioutil.ReadAll(req.GetRaw().Body)
  if err != nil {
    res.Header.SetStatus(400)
    doggo.DoggoEvent("Cannot read sendgrid JSON", err, true)
    return
  }
  var sendgridJson []map[string]interface{}
  err = json.Unmarshal(data, &sendgridJson)
  if err != nil {
    res.Header.SetStatus(400)
    doggo.DoggoEvent("Parse sendgrid JSON failed", err, true)
    return
  }
  // empty events array pushed
  if len(sendgridJson) == 0 {
    res.Header.SetStatus(200)
    res.End()
    return
  }
  // Return the response to SG, and then process the events
  // so as to avoid timeouts
  res.Write("Success").End()

  // loop over all the events that are pushed
  for _, event := range sendgridJson {
    var batchId gocql.UUID
    var logId gocql.UUID
    var emptyId gocql.UUID
    var timestamp time.Time = time.Now()

    eventName := strings.ToLower(event["event"].(string))
    if event["batch_id"] == nil {
      // break, we expect a batch id
    } else {
      batchId, _ = gocql.ParseUUID(event["batch_id"].(string))
    }

    if event["log_id"] != nil {
      logId, _ = gocql.ParseUUID(event["log_id"].(string))
    }

    var log *models.Log = nil
    var shouldRetry bool = false
    var updateObject map[string]interface{}
    switch(eventName){
      case "processed" :
        updateObject = map[string]interface{}{"state": "sent"}
        break;
      case "delivered" :
        updateObject = map[string]interface{}{"state": "delivered"}
      break;
      case "deferred"  :
        updateObject = map[string]interface{}{"state": "deferred"}
        shouldRetry = false
        break;
      case "bounce"    :
        updateObject = map[string]interface{}{"state": "bounced"}
        shouldRetry = true
      break;
      case "dropped"   :
        updateObject = map[string]interface{}{"state": "dropped"}
        shouldRetry = false
      break;
      default: continue
    }

    if event["timestamp"] != nil {
      timestamp = time.Unix(int64(event["timestamp"].(float64)), 0)
    }
    // set the updated timestamp
    updateObject["last_update"] = timestamp

    batch := &models.Batch{Batch_id: batchId}
    success, err := batch.Exists()
    if err != nil {
      // cannot process the event
      // todo: Push to a DB rabbit queue
      doggo.DoggoEvent("Batch Search failed", "BatchId:" + batchId.String() + "\n" + err.Error() , true)
      continue
    }

    if (success == false){
      // batch doesn't exists
      // Ack the request and move on
      continue
    }
    // fetch the email log doc
    if logId != emptyId {
      log = &models.Log{Batch_id: batchId, Log_id: logId}
      success, _ := log.Exists()
      if success == true {
        // this email is done
        // mark sendgrid as been used
        var status = log.GetMap()
        status["sendgrid"] = true
        updateObject["status"] = status
        if shouldRetry == false {
          // Lets update the status and exit
          log.Update(updateObject)
          continue
       }
      } else {
        // drop the packet
        // log_id doesn't exists, dangling packet
        continue
     }
    }
    // this is a retrial/first info event request
    if log == nil {
      // create a log
      fmt.Println("Creating a log")
      log = &models.Log{}
      _, err := log.Create(map[string]interface{}{
        "state"  : eventName,
        "status" : map[string]bool{"sendgrid": true},
        "batch_id": batchId,
        "user_id": batch.User_id,
        "email": event["email"].(string),
        "created_at": time.Now(),
        "last_update": timestamp,
      })
      fmt.Println("Error=>", err)
      if err != nil {
        doggo.DoggoEvent("Log Creation failed", "BatchId:" + batchId.String() + "\n" + err.Error() , true)
      }
    }
    fmt.Println("Log Created")
    // successfully delivered the packet, save the log
    // and continue
    if shouldRetry == false {
      doggo.AddDoggoMetric("sendgrid.success")
      helpers.UpdateStats(batch.User_id, "success")
      _, err := log.Update(updateObject)
      fmt.Println("Log Updated", err)
      continue
    }
    fmt.Println("Trying the next worker")
    // else our mail has been dropped lets see what we can do
    // pick the one which we haven't tried
    doggo.AddDoggoMetric("sendgrid.failed")
    helpers.UpdateStats(batch.User_id, "failed")
    worker := helpers.NextChannelToTry(log.Status)
    if worker != nil {
      // use this worker to retry sending the email
      routingKey := "log." + worker.Namespace + log.Log_id.String()
      fmt.Println("RoutingKey", routingKey)
      // push it into the queue
      worker.Channel.Publish(routingKey, nil)
    } else {
      // exhausted all the trials for this email
      // drop it permanently
      fmt.Println("Le'me just update and exit")
      updateObject["state"] = "failed"
      log.Update(updateObject)
    }
  }
}

/*
  MailGun Webhook
  Mailgun service only provides delivered, bounced and hardfail status
  so we can only take actions in case of a bounced email only.
*/
func BatchUpdateMGStatus(req *request.Request, res *response.Response, next func()) {
  var batchId gocql.UUID
  var logId gocql.UUID
  var emptyId gocql.UUID
  var email string
  var updateObject map[string]interface{}
  var err error
  var batch *models.Batch
  var log *models.Log = nil
  var success bool
  var boundary string
  var timestamp time.Time
  // Mailgun has this strange way of sending status updates
  // if it is a failed event than you get a form-data else a simple
  // post form is sent in case of delivered status
  if helpers.IsMultipart(req.Header["content-type"], &boundary) {
    form := helpers.ReadMultiPartForm(boundary, req.GetRaw().Body)
    if form == nil {
      res.Header.SetStatus(400)
      res.Write("Bad Request").End()
      return
    }
    // inject the form-data values
    req.Body = form.Value
  }
  // validate the request, should have valid auth header
  if helpers.IsValidMGRequest(req) == false {
    res.Header.SetStatus(403)
    res.Write("Access Denied").End()
    return
  }
  // simple checks
  if len(req.Body["event"]) == 0 || len(req.Body["recipient"]) == 0 || len(req.Body["batch-id"]) == 0 {
    res.Header.SetStatus(400)
    res.Write("Bad Request").End()
    return
  }
  // Cool, it is a valid webhook request from MG
  // process it.
  batch_id_data := req.Body["batch-id"][0]
  if batch_id_data == "" {
    // batch_id not present, drop the request
    res.Header.SetStatus(400)
    res.Write("Bad Request").End()
    return
  }
  // parse the batch_id
  batchId, err = gocql.ParseUUID(batch_id_data)
  if err != nil {
    // we messed up! Why would you send a bad JSON dron?
    res.Header.SetStatus(400)
    res.Write("Bad Request").End()
    return
  }
  // see if we have a log_id
  if len(req.Body["log-id"]) == 1 {
    logId, err = gocql.ParseUUID(req.Body["log-id"][0])
  }
  email = req.Body["recipient"][0]
  var eventName = req.Body["event"][0]
  var shouldRetry bool = false
  // send success to the client and then process it
  res.Write("success").End()

  switch(eventName){
    case "delivered" :
      updateObject = map[string]interface{}{"state": "delivered"}
    break;
    case "bounced"    :
      updateObject = map[string]interface{}{"state": "bounced"}
      shouldRetry = true
    break;
    case "hardfail"   :
      updateObject = map[string]interface{}{"state": "dropped"}
      shouldRetry = false
    break;
   case "dropped":
      updateObject = map[string]interface{}{"state": "dropped"}
      shouldRetry = false
    default: return // we can add further click/open events but OOS
  }
  // retrieve the timestamp
  if len(req.Body["timestamp"]) != 0 {
    val, _ := strconv.Atoi(req.Body["timestamp"][0])
    timestamp = time.Unix(int64(val), 0)
  }
  // set the updated timestamp
  updateObject["last_update"] = timestamp

  batch = &models.Batch{Batch_id: batchId}
  success, err = batch.Exists()
  if err != nil {
    // cannot process the event
    // todo: Push to a DB rabbit queue
    doggo.DoggoEvent("Batch Search failed", "BatchId:" + batchId.String() + "\n" + err.Error() , true)
    return
  }

  if (success == false){
    // batch doesn't exists
    // Ack the request and move on
    return
  }
  // fetch the email log doc
  if logId != emptyId {
    log = &models.Log{Batch_id: batchId, Log_id: logId}
    success, _ = log.Exists()
    if success == true {
      // this email is done
      // mark mailgun as been used
      var status = log.GetMap()
      status["mailgun"] = true
      updateObject["status"] = status
      if shouldRetry == false {
        // Lets update the status and exit
        log.Update(updateObject)
        return
     }
    } else {
      // drop the packet
      // log_id doesn't exists, dangling packet
      return
   }
  }
  // this is a retrial mail
  if log == nil {
    // create a log
    log = &models.Log{}
    _, err = log.Create(map[string]interface{}{
      "state"  : eventName,
      "status" : map[string]bool{"mailgun": true},
      "batch_id": batchId,
      "user_id": batch.User_id,
      "email": email,
      "created_at": time.Now(),
      "last_update": timestamp,
    })
    if err != nil {
      doggo.DoggoEvent("Log Creation failed", "BatchId:" + batchId.String() + "\n" + err.Error() , true)
    }
  }
  // successfully delivered the packet, save the log
  // and return
  if shouldRetry == false {
    helpers.UpdateStats(batch.User_id, "success")
    doggo.AddDoggoMetric("mailgun.success")
    _, err = log.Update(updateObject)
    return
  }
  fmt.Println("Trying the next worker")
  doggo.AddDoggoMetric("mailgun.failed")
  helpers.UpdateStats(batch.User_id, "failed")
  // else our mail has been dropped lets see what we can do
  // pick the one which we haven't tried
  worker := helpers.NextChannelToTry(log.Status)
  if worker != nil {
    // use this worker to retry sending the email
    routingKey := "log." + worker.Namespace + log.Log_id.String()
    fmt.Println("RoutingKey", routingKey)
    // push it into the queue
    worker.Channel.Publish(routingKey, nil)
  } else {
    // exhausted all the trials for this email
    // drop it permanently
    updateObject["state"] = "failed"
    fmt.Println("Le'me just update and exit")
    log.Update(updateObject)
  }
}

var WebhookController = func() interface{} {
  var Router = express.Router()
  Router.Post("/api/webhook/sendgrid", BatchUpdateSGStatus)
  Router.Post("/api/webhook/mailgun", BatchUpdateMGStatus)
  return *Router
}()
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"math/rand"
	"os"
	"strings"

	"time"

	"github.com/grafana/grafana/pkg/log"
	"github.com/nsqio/go-nsq"
	"github.com/raintank/met"
	"github.com/raintank/met/helper"
	"github.com/raintank/raintank-metric/dur"
	"github.com/raintank/raintank-metric/msg"
	"github.com/raintank/raintank-metric/schema"
)

const maxMpubSize = 5 * 1024 * 1024 // nsq errors if more. not sure if can be changed
const maxMetricPerMsg = 1000        // emperically found through benchmarks (should result in 64~128k messages)
var (
	producer          *nsq.Producer
	metricsPublished  met.Count
	messagesPublished met.Count
	messagesSize      met.Meter
	metricsPerMessage met.Meter
	publishDuration   met.Timer
)

var (
	showVersion = flag.Bool("version", false, "print version string")

	topic = flag.String("topic", "metrics", "NSQ topic")

	//producerOpts     = flag.String("producer-opt", "", "option to passthrough to nsq.Producer (may be given multiple times as comma-separated list, see http://godoc.org/github.com/nsqio/go-nsq#Config)")
	nsqdTCPAddr  = flag.String("nsqd-tcp-address", "localhost:4150", "nsqd TCP address")
	logLevel     = flag.Int("log-level", 2, "log level. 0=TRACE|1=DEBUG|2=INFO|3=WARN|4=ERROR|5=CRITICAL|6=FATAL")
	orgs         = flag.Int("orgs", 2000, "how many orgs to simulate")
	keysPerOrg   = flag.Int("keys-per-org", 100, "how many metrics per orgs to simulate")
	speedup      = flag.Int("speedup", 1, "for each advancement of real time, how many advancements of fake data to simulate")
	metricPeriod = flag.Int("metricPeriod", 1, "period in seconds between metric points")
	flushPeriod  = flag.Int("flushPeriod", 100, "period in ms between flushes. metricPeriod must be cleanly divisible by flushPeriod. does not affect volume/throughput per se. the message is adjusted as to keep the volume/throughput constant")
	offset       = flag.String("offset", "0", "offset duration expression. (how far back in time to start. e.g. 1month, 6h, etc")
	stopAtNow    = flag.Bool("stop-at-now", false, "stop program instead of starting to write data with future timestamps")
	statsdAddr   = flag.String("statsd-addr", "localhost:8125", "statsd address")
	statsdType   = flag.String("statsd-type", "standard", "statsd type: standard or datadog")
)

func main() {
	flag.Parse()

	log.NewLogger(0, "console", fmt.Sprintf(`{"level": %d, "formatting":true}`, *logLevel))

	if *showVersion {
		fmt.Println("fake_metrics_to_nsq")
		return
	}

	if *topic == "" {
		log.Fatal(4, "--topic is required")
	}

	hostname, err := os.Hostname()
	if err != nil {
		log.Fatal(4, "failed to lookup hostname. %s", err)
	}
	stats, err := helper.New(true, *statsdAddr, *statsdType, "fake_metrics_to_nsq", strings.Replace(hostname, ".", "_", -1))
	if err != nil {
		log.Fatal(4, "failed to initialize statsd. %s", err)
	}

	cfg := nsq.NewConfig()
	cfg.UserAgent = "fake_metrics_to_nsq"
	producer, err = nsq.NewProducer(*nsqdTCPAddr, cfg)
	if err != nil {
		log.Fatal(0, "failed to initialize nsq producer.", err)
	}
	err = producer.Ping()
	if err != nil {
		log.Fatal(0, "can't connect to nsqd: %s", err)
	}
	metricsPublished = stats.NewCount("metricpublisher.metrics-published")
	messagesPublished = stats.NewCount("metricpublisher.messages-published")
	messagesSize = stats.NewMeter("metricpublisher.message_size", 0)
	metricsPerMessage = stats.NewMeter("metricpublisher.metrics_per_message", 0)
	publishDuration = stats.NewTimer("metricpublisher.publish_duration", 0)
	if (1000**metricPeriod)%*flushPeriod != 0 {
		panic("metricPeriod must be cleanly divisible by flushPeriod")
	}
	off := int(dur.MustParseUsec("offset", *offset))
	run(*orgs, *keysPerOrg, *metricPeriod, *flushPeriod, off, *speedup)
}

func Reslice(in []*schema.MetricData, size int) [][]*schema.MetricData {
	numSubSlices := len(in) / size
	if len(in)%size > 0 {
		numSubSlices += 1
	}
	out := make([][]*schema.MetricData, numSubSlices)
	for i := 0; i < numSubSlices; i++ {
		start := i * size
		end := (i + 1) * size
		if end > len(in) {
			out[i] = in[start:]
		} else {
			out[i] = in[start:end]
		}
	}
	return out
}

func Publish(metrics []*schema.MetricData) error {
	if len(metrics) == 0 {
		return nil
	}
	// typical metrics seem to be around 300B
	// nsqd allows <= 10MiB messages.
	// we ideally have 64kB ~ 1MiB messages (see benchmark https://gist.github.com/Dieterbe/604232d35494eae73f15)
	// at 300B, about 3500 msg fit in 1MiB
	// in worst case, this allows messages up to 2871B
	// this could be made more robust of course

	// real world findings in dev-stack with env-load:
	// 159569B msg /795  metrics per msg = 200B per msg
	// so peak message size is about 3500*200 = 700k (seen 711k)

	subslices := Reslice(metrics, 3500)

	for _, subslice := range subslices {
		id := time.Now().UnixNano()
		data, err := msg.CreateMsg(subslice, id, msg.FormatMetricDataArrayMsgp)
		if err != nil {
			log.Fatal(0, "Fatal error creating metric message: %s", err)
		}
		metricsPublished.Inc(int64(len(subslice)))
		messagesPublished.Inc(1)
		messagesSize.Value(int64(len(data)))
		metricsPerMessage.Value(int64(len(subslice)))
		pre := time.Now()
		err = producer.Publish(*topic, data)
		publishDuration.Value(time.Since(pre))
		if err != nil {
			log.Fatal(0, "can't publish to nsqd: %s", err)
		}
		//log.Info("published metrics %d size=%d", id, len(data))
	}
	return nil
}

func run(orgs, keysPerOrg, metricPeriod, flushPeriod, offset, speedup int) {
	// examples:
	// flushperiod - metricperiod - speedup -> data to flush each flush
	// 100           1              1          1/10 -> flush 1 of 10 fractions
	// 100           1              2          1/5 -> flush 1 of 5 fractions
	// 1000          1              2          2x -> flush ratio 2
	// 2000          1              1          2x -> flush ratio 2
	// 2000          1              2          4x -> flush ratio 4

	if flushPeriod*speedup >= 1000*metricPeriod {
		runMultiplied(orgs, keysPerOrg, metricPeriod, flushPeriod, offset, speedup)
	} else {
		runDivided(orgs, keysPerOrg, metricPeriod, flushPeriod, offset, speedup)
	}
}

// we need to send multiples of the data at every flush
func runMultiplied(orgs, keysPerOrg, metricPeriod, flushPeriod, offset, speedup int) {
	if flushPeriod*speedup%(1000*metricPeriod) != 0 {
		panic(fmt.Sprintf("not a good fit. metricPeriod*1000 should fit in flushPeriod*speedup"))
	}
	ratio := flushPeriod * speedup / 1000 / metricPeriod
	uniqueKeys := orgs * keysPerOrg
	totalKeys := uniqueKeys * ratio
	tickDur := time.Duration(flushPeriod) * time.Millisecond
	tick := time.NewTicker(tickDur)
	fmt.Printf("each %s, sending %d metrics. (sending the %d metrics %d times per flush, each time advancing them by %d s)\n", tickDur, totalKeys, uniqueKeys, ratio, metricPeriod)

	metrics := make([]*schema.MetricData, totalKeys)
	for r := 1; r <= ratio; r++ {
		for o := 1; o <= orgs; o++ {
			for k := 1; k <= keysPerOrg; k++ {
				i := (r-1)*orgs*keysPerOrg + (o-1)*keysPerOrg + k - 1
				metrics[i] = &schema.MetricData{
					Name:       fmt.Sprintf("some.id.of.a.metric.%d", k),
					Metric:     "some.id.of.a.metric",
					OrgId:      o,
					Interval:   metricPeriod,
					Value:      0,
					Unit:       "ms",
					TargetType: "gauge",
					Tags:       []string{"some_tag", "ok"},
				}
				metrics[i].SetId()
			}
		}
	}
	mp := int64(metricPeriod)
	ts := time.Now().Unix() - int64(offset) - mp

	for range tick.C {
		now := time.Now().Unix()
		for i := 0; i < len(metrics); i++ {
			if i%uniqueKeys == 0 {
				ts += mp
				if ts > now && *stopAtNow {
					return
				}

			}
			metrics[i].Time = ts
			metrics[i].Value = rand.Float64() * float64(i+1)
		}
		Publish(metrics)
	}
}

// we need to send fractions of the data at every flush
func runDivided(orgs, keysPerOrg, metricPeriod, flushPeriod, offset, speedup int) {
	totalKeys := orgs * keysPerOrg
	fractions := 1000 * metricPeriod / flushPeriod / speedup
	if totalKeys%fractions != 0 {
		panic(fmt.Sprintf("not a good fit: (%d orgs * %d keysPerOrg)=%d metrics but 1000*%d metricPeriod / %d flushPeriod / %d speedup = %d fractions)", orgs, keysPerOrg, totalKeys, metricPeriod, flushPeriod, speedup, fractions))
	}
	metricsPerFrac := totalKeys / fractions
	tickDur := time.Duration(flushPeriod) * time.Millisecond
	tick := time.NewTicker(tickDur)
	fmt.Printf("each %s, sending %d out of %d metrics. timestamps increment by %d s every %d flushes)\n", tickDur, metricsPerFrac, totalKeys, metricPeriod, fractions)

	metrics := make([]*schema.MetricData, totalKeys)
	for o := 1; o <= orgs; o++ {
		for k := 1; k <= keysPerOrg; k++ {
			i := (o-1)*keysPerOrg + k - 1
			metrics[i] = &schema.MetricData{
				Name:       fmt.Sprintf("some.id.of.a.metric.%d", k),
				Metric:     "some.id.of.a.metric",
				OrgId:      o,
				Interval:   metricPeriod,
				Value:      0,
				Unit:       "ms",
				TargetType: "gauge",
				Tags:       []string{"some_tag", "ok"},
			}
			metrics[i].SetId()
		}
	}

	mp := int64(metricPeriod)
	ts := time.Now().Unix() - int64(offset) - mp
	endIndex, startIndex := 0, 0
	for range tick.C {
		now := time.Now().Unix()
		startIndex = endIndex
		if startIndex == len(metrics) {
			startIndex = 0
		}
		endIndex = startIndex + metricsPerFrac
		if startIndex == 0 {
			ts += mp
			if ts > now && *stopAtNow {
				return
			}
		}
		for i := startIndex; i < endIndex; i++ {
			metrics[i].Time = ts
			metrics[i].Value = rand.Float64() * float64(i+1)
		}
		Publish(metrics[startIndex:endIndex])
	}
}
<commit_msg>explain better<commit_after>package main

import (
	"flag"
	"fmt"
	"math/rand"
	"os"
	"strings"

	"time"

	"github.com/grafana/grafana/pkg/log"
	"github.com/nsqio/go-nsq"
	"github.com/raintank/met"
	"github.com/raintank/met/helper"
	"github.com/raintank/raintank-metric/dur"
	"github.com/raintank/raintank-metric/msg"
	"github.com/raintank/raintank-metric/schema"
)

const maxMpubSize = 5 * 1024 * 1024 // nsq errors if more. not sure if can be changed
const maxMetricPerMsg = 1000        // emperically found through benchmarks (should result in 64~128k messages)
var (
	producer          *nsq.Producer
	metricsPublished  met.Count
	messagesPublished met.Count
	messagesSize      met.Meter
	metricsPerMessage met.Meter
	publishDuration   met.Timer
)

var (
	showVersion = flag.Bool("version", false, "print version string")

	topic = flag.String("topic", "metrics", "NSQ topic")

	//producerOpts     = flag.String("producer-opt", "", "option to passthrough to nsq.Producer (may be given multiple times as comma-separated list, see http://godoc.org/github.com/nsqio/go-nsq#Config)")
	nsqdTCPAddr  = flag.String("nsqd-tcp-address", "localhost:4150", "nsqd TCP address")
	logLevel     = flag.Int("log-level", 2, "log level. 0=TRACE|1=DEBUG|2=INFO|3=WARN|4=ERROR|5=CRITICAL|6=FATAL")
	orgs         = flag.Int("orgs", 2000, "how many orgs to simulate")
	keysPerOrg   = flag.Int("keys-per-org", 100, "how many metrics per orgs to simulate")
	speedup      = flag.Int("speedup", 1, "for each advancement of real time, how many advancements of fake data to simulate")
	metricPeriod = flag.Int("metricPeriod", 1, "period in seconds between metric points")
	flushPeriod  = flag.Int("flushPeriod", 100, "period in ms between flushes. metricPeriod must be cleanly divisible by flushPeriod. does not affect volume/throughput per se. the message is adjusted as to keep the volume/throughput constant")
	offset       = flag.String("offset", "0", "offset duration expression. (how far back in time to start. e.g. 1month, 6h, etc")
	stopAtNow    = flag.Bool("stop-at-now", false, "stop program instead of starting to write data with future timestamps")
	statsdAddr   = flag.String("statsd-addr", "localhost:8125", "statsd address")
	statsdType   = flag.String("statsd-type", "standard", "statsd type: standard or datadog")
)

func main() {
	flag.Parse()

	log.NewLogger(0, "console", fmt.Sprintf(`{"level": %d, "formatting":true}`, *logLevel))

	if *showVersion {
		fmt.Println("fake_metrics_to_nsq")
		return
	}

	if *topic == "" {
		log.Fatal(4, "--topic is required")
	}

	hostname, err := os.Hostname()
	if err != nil {
		log.Fatal(4, "failed to lookup hostname. %s", err)
	}
	stats, err := helper.New(true, *statsdAddr, *statsdType, "fake_metrics_to_nsq", strings.Replace(hostname, ".", "_", -1))
	if err != nil {
		log.Fatal(4, "failed to initialize statsd. %s", err)
	}

	cfg := nsq.NewConfig()
	cfg.UserAgent = "fake_metrics_to_nsq"
	producer, err = nsq.NewProducer(*nsqdTCPAddr, cfg)
	if err != nil {
		log.Fatal(0, "failed to initialize nsq producer.", err)
	}
	err = producer.Ping()
	if err != nil {
		log.Fatal(0, "can't connect to nsqd: %s", err)
	}
	metricsPublished = stats.NewCount("metricpublisher.metrics-published")
	messagesPublished = stats.NewCount("metricpublisher.messages-published")
	messagesSize = stats.NewMeter("metricpublisher.message_size", 0)
	metricsPerMessage = stats.NewMeter("metricpublisher.metrics_per_message", 0)
	publishDuration = stats.NewTimer("metricpublisher.publish_duration", 0)
	if (1000**metricPeriod)%*flushPeriod != 0 {
		panic("metricPeriod must be cleanly divisible by flushPeriod")
	}
	off := int(dur.MustParseUsec("offset", *offset))
	run(*orgs, *keysPerOrg, *metricPeriod, *flushPeriod, off, *speedup)
}

func Reslice(in []*schema.MetricData, size int) [][]*schema.MetricData {
	numSubSlices := len(in) / size
	if len(in)%size > 0 {
		numSubSlices += 1
	}
	out := make([][]*schema.MetricData, numSubSlices)
	for i := 0; i < numSubSlices; i++ {
		start := i * size
		end := (i + 1) * size
		if end > len(in) {
			out[i] = in[start:]
		} else {
			out[i] = in[start:end]
		}
	}
	return out
}

func Publish(metrics []*schema.MetricData) error {
	if len(metrics) == 0 {
		return nil
	}
	// typical metrics seem to be around 300B
	// nsqd allows <= 10MiB messages.
	// we ideally have 64kB ~ 1MiB messages (see benchmark https://gist.github.com/Dieterbe/604232d35494eae73f15)
	// at 300B, about 3500 msg fit in 1MiB
	// in worst case, this allows messages up to 2871B
	// this could be made more robust of course

	// real world findings in dev-stack with env-load:
	// 159569B msg /795  metrics per msg = 200B per msg
	// so peak message size is about 3500*200 = 700k (seen 711k)

	subslices := Reslice(metrics, 3500)

	for _, subslice := range subslices {
		id := time.Now().UnixNano()
		data, err := msg.CreateMsg(subslice, id, msg.FormatMetricDataArrayMsgp)
		if err != nil {
			log.Fatal(0, "Fatal error creating metric message: %s", err)
		}
		metricsPublished.Inc(int64(len(subslice)))
		messagesPublished.Inc(1)
		messagesSize.Value(int64(len(data)))
		metricsPerMessage.Value(int64(len(subslice)))
		pre := time.Now()
		err = producer.Publish(*topic, data)
		publishDuration.Value(time.Since(pre))
		if err != nil {
			log.Fatal(0, "can't publish to nsqd: %s", err)
		}
		//log.Info("published metrics %d size=%d", id, len(data))
	}
	return nil
}

func run(orgs, keysPerOrg, metricPeriod, flushPeriod, offset, speedup int) {
	// examples:
	// flushperiod - metricperiod - speedup -> data to flush each flush
	// 100     ms    1 s            1          1/10 -> flush 1 of 10 fractions runDivided <-- default
	// 100     ms    1 s            2          1/5 -> flush 1 of 5 fractions   runDivided
	// 100     ms    1 s            100        1/5 -> flush 1x                 runMultiplied
	// 1000    ms    1 s            2          2x -> flush 2x                  runMultiplied
	// 2000    ms    1 s            1          2x -> flush 2x                  runMultiplied
	// 2000    ms    1 s            2          4x -> flush 4x                  runMultiplied

	if flushPeriod*speedup >= 1000*metricPeriod {
		runMultiplied(orgs, keysPerOrg, metricPeriod, flushPeriod, offset, speedup)
	} else {
		runDivided(orgs, keysPerOrg, metricPeriod, flushPeriod, offset, speedup)
	}
}

// we need to send multiples of the data at every flush
func runMultiplied(orgs, keysPerOrg, metricPeriod, flushPeriod, offset, speedup int) {
	if flushPeriod*speedup%(1000*metricPeriod) != 0 {
		panic(fmt.Sprintf("not a good fit. metricPeriod*1000 should fit in flushPeriod*speedup"))
	}
	ratio := flushPeriod * speedup / 1000 / metricPeriod
	uniqueKeys := orgs * keysPerOrg
	totalKeys := uniqueKeys * ratio
	tickDur := time.Duration(flushPeriod) * time.Millisecond
	tick := time.NewTicker(tickDur)
	fmt.Printf("each %s, sending %d metrics. (sending the %d metrics %d times per flush, each time advancing them by %d s)\n", tickDur, totalKeys, uniqueKeys, ratio, metricPeriod)

	metrics := make([]*schema.MetricData, totalKeys)
	for r := 1; r <= ratio; r++ {
		for o := 1; o <= orgs; o++ {
			for k := 1; k <= keysPerOrg; k++ {
				i := (r-1)*orgs*keysPerOrg + (o-1)*keysPerOrg + k - 1
				metrics[i] = &schema.MetricData{
					Name:       fmt.Sprintf("some.id.of.a.metric.%d", k),
					Metric:     "some.id.of.a.metric",
					OrgId:      o,
					Interval:   metricPeriod,
					Value:      0,
					Unit:       "ms",
					TargetType: "gauge",
					Tags:       []string{"some_tag", "ok"},
				}
				metrics[i].SetId()
			}
		}
	}
	mp := int64(metricPeriod)
	ts := time.Now().Unix() - int64(offset) - mp

	for range tick.C {
		now := time.Now().Unix()
		for i := 0; i < len(metrics); i++ {
			if i%uniqueKeys == 0 {
				ts += mp
				if ts > now && *stopAtNow {
					return
				}

			}
			metrics[i].Time = ts
			metrics[i].Value = rand.Float64() * float64(i+1)
		}
		Publish(metrics)
	}
}

// we need to send fractions of the data at every flush
func runDivided(orgs, keysPerOrg, metricPeriod, flushPeriod, offset, speedup int) {
	totalKeys := orgs * keysPerOrg
	fractions := 1000 * metricPeriod / flushPeriod / speedup
	if totalKeys%fractions != 0 {
		panic(fmt.Sprintf("not a good fit: (%d orgs * %d keysPerOrg)=%d metrics but 1000*%d metricPeriod / %d flushPeriod / %d speedup = %d fractions)", orgs, keysPerOrg, totalKeys, metricPeriod, flushPeriod, speedup, fractions))
	}
	metricsPerFrac := totalKeys / fractions
	tickDur := time.Duration(flushPeriod) * time.Millisecond
	tick := time.NewTicker(tickDur)
	fmt.Printf("each %s, sending %d out of %d metrics. timestamps increment by %d s every %d flushes)\n", tickDur, metricsPerFrac, totalKeys, metricPeriod, fractions)

	metrics := make([]*schema.MetricData, totalKeys)
	for o := 1; o <= orgs; o++ {
		for k := 1; k <= keysPerOrg; k++ {
			i := (o-1)*keysPerOrg + k - 1
			metrics[i] = &schema.MetricData{
				Name:       fmt.Sprintf("some.id.of.a.metric.%d", k),
				Metric:     "some.id.of.a.metric",
				OrgId:      o,
				Interval:   metricPeriod,
				Value:      0,
				Unit:       "ms",
				TargetType: "gauge",
				Tags:       []string{"some_tag", "ok"},
			}
			metrics[i].SetId()
		}
	}

	mp := int64(metricPeriod)
	ts := time.Now().Unix() - int64(offset) - mp
	endIndex, startIndex := 0, 0
	for range tick.C {
		now := time.Now().Unix()
		startIndex = endIndex
		if startIndex == len(metrics) {
			startIndex = 0
		}
		endIndex = startIndex + metricsPerFrac
		if startIndex == 0 {
			ts += mp
			if ts > now && *stopAtNow {
				return
			}
		}
		for i := startIndex; i < endIndex; i++ {
			metrics[i].Time = ts
			metrics[i].Value = rand.Float64() * float64(i+1)
		}
		Publish(metrics[startIndex:endIndex])
	}
}
<|endoftext|>