<commit_before>// Copyright (c) 2013 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package btcutil

import (
	"bytes"
	"fmt"
	"github.com/conformal/btcwire"
)

// OutOfRangeError describes an error due to accessing an element that is out
// of range.
type OutOfRangeError string

// BlockHeightUnknown is the value returned for a block height that is unknown.
// This is typically because the block has not been inserted into the main chain
// yet.
const BlockHeightUnknown = int64(-1)

// Error satisfies the error interface and prints human-readable errors.
func (e OutOfRangeError) Error() string {
	return string(e)
}

// Block defines a bitcoin block that provides easier and more efficient
// manipulation of raw wire protocol blocks.  It also memoizes hashes for the
// block and its transactions on their first access so subsequent accesses don't
// have to repeat the relatively expensive hashing operations.
type Block struct {
	msgBlock        *btcwire.MsgBlock  // Underlying MsgBlock
	rawBlock        []byte             // Raw wire encoded bytes for the block
	protocolVersion uint32             // Protocol version used to encode rawBlock
	blockSha        *btcwire.ShaHash   // Cached block hash
	blockHeight     int64              // Height in the main block chain
	txShas          []*btcwire.ShaHash // Cached transaction hashes
	txShasGenerated bool               // ALL transaction hashes generated
}

// MsgBlock returns the underlying btcwire.MsgBlock for the Block.
func (b *Block) MsgBlock() *btcwire.MsgBlock {
	// Return the cached block.
	return b.msgBlock
}

// Bytes returns the raw wire protocol encoded bytes for the Block and the
// protocol version used to encode it.  This is equivalent to calling BtcEncode
// on the underlying btcwire.MsgBlock, however it caches the result so
// subsequent calls are more efficient.
func (b *Block) Bytes() ([]byte, uint32, error) {
	// Return the cached raw block bytes and associated protocol version if
	// it has already been generated.
	if len(b.rawBlock) != 0 {
		return b.rawBlock, b.protocolVersion, nil
	}

	// Encode the MsgBlock into raw block bytes.
	var w bytes.Buffer
	err := b.msgBlock.BtcEncode(&w, b.protocolVersion)
	if err != nil {
		return nil, 0, err
	}
	rawBlock := w.Bytes()

	// Cache the encoded bytes and return them.
	b.rawBlock = rawBlock
	return rawBlock, b.protocolVersion, nil
}

// Sha returns the block identifier hash for the Block.  This is equivalent to
// calling BlockSha on the underlying btcwire.MsgBlock, however it caches the
// result so subsequent calls are more efficient.
func (b *Block) Sha() (*btcwire.ShaHash, error) {
	// Return the cached block hash if it has already been generated.
	if b.blockSha != nil {
		return b.blockSha, nil
	}

	// Generate the block hash.  Ignore the error since BlockSha can't
	// currently fail.
	sha, _ := b.msgBlock.BlockSha(b.protocolVersion)

	// Cache the block hash and return it.
	b.blockSha = &sha
	return &sha, nil
}

// TxSha returns the hash for the requested transaction number in the Block.
// The supplied index is 0 based.  That is to say, the first transaction is the
// block is txNum 0.  This is equivalent to calling TxSha on the underlying
// btcwire.MsgTx, however it caches the result so subsequent calls are more
// efficient.
func (b *Block) TxSha(txNum int) (*btcwire.ShaHash, error) {
	// Ensure the requested transaction is in range.
	numTx := b.msgBlock.Header.TxnCount
	if txNum < 0 || uint64(txNum) > numTx {
		str := fmt.Sprintf("transaction index %d is out of range - max %d",
			txNum, numTx-1)
		return nil, OutOfRangeError(str)
	}

	// Generate slice to hold all of the transaction hashes if needed.
	if len(b.txShas) == 0 {
		b.txShas = make([]*btcwire.ShaHash, numTx)
	}

	// Return the cached hash if it has already been generated.
	if b.txShas[txNum] != nil {
		return b.txShas[txNum], nil
	}

	// Generate the hash for the transaction.  Ignore the error since TxSha
	// can't currently fail.
	sha, _ := b.msgBlock.Transactions[txNum].TxSha(b.protocolVersion)

	// Cache the transaction hash and return it.
	b.txShas[txNum] = &sha
	return &sha, nil
}

// TxShas returns a slice of hashes for all transactions in the Block.  This is
// equivalent to calling TxSha on each underlying btcwire.MsgTx, however it
// caches the result so subsequent calls are more efficient.
func (b *Block) TxShas() ([]*btcwire.ShaHash, error) {
	// Return cached hashes if they have ALL already been generated.  This
	// flag is necessary because the transaction hashes are lazily generated
	// in a sparse fashion.
	if b.txShasGenerated {
		return b.txShas, nil
	}

	// Generate slice to hold all of the transaction hashes if needed.
	if len(b.txShas) == 0 {
		b.txShas = make([]*btcwire.ShaHash, b.msgBlock.Header.TxnCount)
	}

	// Generate and cache the transaction hashes for all that haven't already
	// been done.
	for i, hash := range b.txShas {
		if hash == nil {
			// Ignore the error since TxSha can't currently fail.
			sha, _ := b.msgBlock.Transactions[i].TxSha(b.protocolVersion)
			b.txShas[i] = &sha
		}
	}

	b.txShasGenerated = true
	return b.txShas, nil
}

// ProtocolVersion returns the protocol version that was used to create the
// underlying btcwire.MsgBlock.
func (b *Block) ProtocolVersion() uint32 {
	return b.protocolVersion
}

// TxLoc() returns the offsets and lengths of each transaction in a raw block.
// It is used to allow fast indexing into transactions within the raw byte
// stream.
func (b *Block) TxLoc() ([]btcwire.TxLoc, error) {
	rawMsg, pver, err := b.Bytes()
	if err != nil {
		return nil, err
	}
	rbuf := bytes.NewBuffer(rawMsg)

	var mblock btcwire.MsgBlock
	txLocs, err := mblock.BtcDecodeTxLoc(rbuf, pver)
	if err != nil {
		return nil, err
	}
	return txLocs, err
}

// Height returns the saved height of the block in the blockchain.  This value
// will be BlockHeightUnknown if it hasn't already explicitly been set.
func (b *Block) Height() int64 {
	return b.blockHeight
}

// SetHeight sets the height of the block in the blockchain.
func (b *Block) SetHeight(height int64) {
	b.blockHeight = height
}

// NewBlock returns a new instance of a bitcoin block given an underlying
// btcwire.MsgBlock and protocol version.  See Block.
func NewBlock(msgBlock *btcwire.MsgBlock, pver uint32) *Block {
	return &Block{
		msgBlock:        msgBlock,
		protocolVersion: pver,
		blockHeight:     BlockHeightUnknown,
	}
}

// NewBlockFromBytes returns a new instance of a bitcoin block given the
// raw wire encoded bytes and protocol version used to encode those bytes.
// See Block.
func NewBlockFromBytes(rawBlock []byte, pver uint32) (*Block, error) {
	// Decode the raw block bytes into a MsgBlock.
	var msgBlock btcwire.MsgBlock
	br := bytes.NewBuffer(rawBlock)
	err := msgBlock.BtcDecode(br, pver)
	if err != nil {
		return nil, err
	}

	b := Block{
		msgBlock:        &msgBlock,
		rawBlock:        rawBlock,
		protocolVersion: pver,
		blockHeight:     BlockHeightUnknown,
	}
	return &b, nil
}

// NewBlockFromBlockAndBytes returns a new instance of a bitcoin block given
// an underlying btcwire.MsgBlock, protocol version and raw Block.  See Block.
func NewBlockFromBlockAndBytes(msgBlock *btcwire.MsgBlock, rawBlock []byte, pver uint32) *Block {
	return &Block{
		msgBlock:        msgBlock,
		rawBlock:        rawBlock,
		protocolVersion: pver,
		blockHeight:     BlockHeightUnknown,
	}
}
<commit_msg>Make TxLoc comment consistent with others.<commit_after>// Copyright (c) 2013 Conformal Systems LLC.
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package btcutil

import (
	"bytes"
	"fmt"
	"github.com/conformal/btcwire"
)

// OutOfRangeError describes an error due to accessing an element that is out
// of range.
type OutOfRangeError string

// BlockHeightUnknown is the value returned for a block height that is unknown.
// This is typically because the block has not been inserted into the main chain
// yet.
const BlockHeightUnknown = int64(-1)

// Error satisfies the error interface and prints human-readable errors.
func (e OutOfRangeError) Error() string {
	return string(e)
}

// Block defines a bitcoin block that provides easier and more efficient
// manipulation of raw wire protocol blocks.  It also memoizes hashes for the
// block and its transactions on their first access so subsequent accesses don't
// have to repeat the relatively expensive hashing operations.
type Block struct {
	msgBlock        *btcwire.MsgBlock  // Underlying MsgBlock
	rawBlock        []byte             // Raw wire encoded bytes for the block
	protocolVersion uint32             // Protocol version used to encode rawBlock
	blockSha        *btcwire.ShaHash   // Cached block hash
	blockHeight     int64              // Height in the main block chain
	txShas          []*btcwire.ShaHash // Cached transaction hashes
	txShasGenerated bool               // ALL transaction hashes generated
}

// MsgBlock returns the underlying btcwire.MsgBlock for the Block.
func (b *Block) MsgBlock() *btcwire.MsgBlock {
	// Return the cached block.
	return b.msgBlock
}

// Bytes returns the raw wire protocol encoded bytes for the Block and the
// protocol version used to encode it.  This is equivalent to calling BtcEncode
// on the underlying btcwire.MsgBlock, however it caches the result so
// subsequent calls are more efficient.
func (b *Block) Bytes() ([]byte, uint32, error) {
	// Return the cached raw block bytes and associated protocol version if
	// it has already been generated.
	if len(b.rawBlock) != 0 {
		return b.rawBlock, b.protocolVersion, nil
	}

	// Encode the MsgBlock into raw block bytes.
	var w bytes.Buffer
	err := b.msgBlock.BtcEncode(&w, b.protocolVersion)
	if err != nil {
		return nil, 0, err
	}
	rawBlock := w.Bytes()

	// Cache the encoded bytes and return them.
	b.rawBlock = rawBlock
	return rawBlock, b.protocolVersion, nil
}

// Sha returns the block identifier hash for the Block.  This is equivalent to
// calling BlockSha on the underlying btcwire.MsgBlock, however it caches the
// result so subsequent calls are more efficient.
func (b *Block) Sha() (*btcwire.ShaHash, error) {
	// Return the cached block hash if it has already been generated.
	if b.blockSha != nil {
		return b.blockSha, nil
	}

	// Generate the block hash.  Ignore the error since BlockSha can't
	// currently fail.
	sha, _ := b.msgBlock.BlockSha(b.protocolVersion)

	// Cache the block hash and return it.
	b.blockSha = &sha
	return &sha, nil
}

// TxSha returns the hash for the requested transaction number in the Block.
// The supplied index is 0 based.  That is to say, the first transaction is the
// block is txNum 0.  This is equivalent to calling TxSha on the underlying
// btcwire.MsgTx, however it caches the result so subsequent calls are more
// efficient.
func (b *Block) TxSha(txNum int) (*btcwire.ShaHash, error) {
	// Ensure the requested transaction is in range.
	numTx := b.msgBlock.Header.TxnCount
	if txNum < 0 || uint64(txNum) > numTx {
		str := fmt.Sprintf("transaction index %d is out of range - max %d",
			txNum, numTx-1)
		return nil, OutOfRangeError(str)
	}

	// Generate slice to hold all of the transaction hashes if needed.
	if len(b.txShas) == 0 {
		b.txShas = make([]*btcwire.ShaHash, numTx)
	}

	// Return the cached hash if it has already been generated.
	if b.txShas[txNum] != nil {
		return b.txShas[txNum], nil
	}

	// Generate the hash for the transaction.  Ignore the error since TxSha
	// can't currently fail.
	sha, _ := b.msgBlock.Transactions[txNum].TxSha(b.protocolVersion)

	// Cache the transaction hash and return it.
	b.txShas[txNum] = &sha
	return &sha, nil
}

// TxShas returns a slice of hashes for all transactions in the Block.  This is
// equivalent to calling TxSha on each underlying btcwire.MsgTx, however it
// caches the result so subsequent calls are more efficient.
func (b *Block) TxShas() ([]*btcwire.ShaHash, error) {
	// Return cached hashes if they have ALL already been generated.  This
	// flag is necessary because the transaction hashes are lazily generated
	// in a sparse fashion.
	if b.txShasGenerated {
		return b.txShas, nil
	}

	// Generate slice to hold all of the transaction hashes if needed.
	if len(b.txShas) == 0 {
		b.txShas = make([]*btcwire.ShaHash, b.msgBlock.Header.TxnCount)
	}

	// Generate and cache the transaction hashes for all that haven't already
	// been done.
	for i, hash := range b.txShas {
		if hash == nil {
			// Ignore the error since TxSha can't currently fail.
			sha, _ := b.msgBlock.Transactions[i].TxSha(b.protocolVersion)
			b.txShas[i] = &sha
		}
	}

	b.txShasGenerated = true
	return b.txShas, nil
}

// ProtocolVersion returns the protocol version that was used to create the
// underlying btcwire.MsgBlock.
func (b *Block) ProtocolVersion() uint32 {
	return b.protocolVersion
}

// TxLoc returns the offsets and lengths of each transaction in a raw block.
// It is used to allow fast indexing into transactions within the raw byte
// stream.
func (b *Block) TxLoc() ([]btcwire.TxLoc, error) {
	rawMsg, pver, err := b.Bytes()
	if err != nil {
		return nil, err
	}
	rbuf := bytes.NewBuffer(rawMsg)

	var mblock btcwire.MsgBlock
	txLocs, err := mblock.BtcDecodeTxLoc(rbuf, pver)
	if err != nil {
		return nil, err
	}
	return txLocs, err
}

// Height returns the saved height of the block in the blockchain.  This value
// will be BlockHeightUnknown if it hasn't already explicitly been set.
func (b *Block) Height() int64 {
	return b.blockHeight
}

// SetHeight sets the height of the block in the blockchain.
func (b *Block) SetHeight(height int64) {
	b.blockHeight = height
}

// NewBlock returns a new instance of a bitcoin block given an underlying
// btcwire.MsgBlock and protocol version.  See Block.
func NewBlock(msgBlock *btcwire.MsgBlock, pver uint32) *Block {
	return &Block{
		msgBlock:        msgBlock,
		protocolVersion: pver,
		blockHeight:     BlockHeightUnknown,
	}
}

// NewBlockFromBytes returns a new instance of a bitcoin block given the
// raw wire encoded bytes and protocol version used to encode those bytes.
// See Block.
func NewBlockFromBytes(rawBlock []byte, pver uint32) (*Block, error) {
	// Decode the raw block bytes into a MsgBlock.
	var msgBlock btcwire.MsgBlock
	br := bytes.NewBuffer(rawBlock)
	err := msgBlock.BtcDecode(br, pver)
	if err != nil {
		return nil, err
	}

	b := Block{
		msgBlock:        &msgBlock,
		rawBlock:        rawBlock,
		protocolVersion: pver,
		blockHeight:     BlockHeightUnknown,
	}
	return &b, nil
}

// NewBlockFromBlockAndBytes returns a new instance of a bitcoin block given
// an underlying btcwire.MsgBlock, protocol version and raw Block.  See Block.
func NewBlockFromBlockAndBytes(msgBlock *btcwire.MsgBlock, rawBlock []byte, pver uint32) *Block {
	return &Block{
		msgBlock:        msgBlock,
		rawBlock:        rawBlock,
		protocolVersion: pver,
		blockHeight:     BlockHeightUnknown,
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The cert-manager 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 http

import (
	"context"
	"fmt"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/selection"
	utilerrors "k8s.io/apimachinery/pkg/util/errors"
	"k8s.io/apimachinery/pkg/util/intstr"

	cmacme "github.com/jetstack/cert-manager/pkg/apis/acme/v1"
	logf "github.com/jetstack/cert-manager/pkg/logs"
)

func (s *Solver) ensureService(ctx context.Context, ch *cmacme.Challenge) (*corev1.Service, error) {
	log := logf.FromContext(ctx).WithName("ensureService")

	log.V(logf.DebugLevel).Info("checking for existing HTTP01 solver services for challenge")
	existingServices, err := s.getServicesForChallenge(ctx, ch)
	if err != nil {
		return nil, err
	}
	if len(existingServices) == 1 {
		logf.WithRelatedResource(log, existingServices[0]).Info("found one existing HTTP01 solver Service for challenge resource")
		return existingServices[0], nil
	}
	if len(existingServices) > 1 {
		log.V(logf.DebugLevel).Info("multiple challenge solver services found for challenge. cleaning up all existing services.")
		err := s.cleanupServices(ctx, ch)
		if err != nil {
			return nil, err
		}
		return nil, fmt.Errorf("multiple existing challenge solver services found and cleaned up. retrying challenge sync")
	}

	log.V(logf.DebugLevel).Info("creating HTTP01 challenge solver service")
	return s.createService(ctx, ch)
}

// getServicesForChallenge returns a list of services that were created to solve
// http challenges for the given domain
func (s *Solver) getServicesForChallenge(ctx context.Context, ch *cmacme.Challenge) ([]*corev1.Service, error) {
	log := logf.FromContext(ctx).WithName("getServicesForChallenge")

	podLabels := podLabels(ch)
	selector := labels.NewSelector()
	for key, val := range podLabels {
		req, err := labels.NewRequirement(key, selection.Equals, []string{val})
		if err != nil {
			return nil, err
		}
		selector = selector.Add(*req)
	}

	serviceList, err := s.serviceLister.Services(ch.Namespace).List(selector)
	if err != nil {
		return nil, err
	}

	var relevantServices []*corev1.Service
	for _, service := range serviceList {
		if !metav1.IsControlledBy(service, ch) {
			logf.WithRelatedResource(log, service).Info("found existing solver pod for this challenge resource, however " +
				"it does not have an appropriate OwnerReference referencing this challenge. Skipping it altogether.")
			continue
		}
		relevantServices = append(relevantServices, service)
	}

	return relevantServices, nil
}

// createService will create the service required to solve this challenge
// in the target API server.
func (s *Solver) createService(ctx context.Context, ch *cmacme.Challenge) (*corev1.Service, error) {
	svc, err := buildService(ch)
	if err != nil {
		return nil, err
	}
	return s.Client.CoreV1().Services(ch.Namespace).Create(ctx, svc, metav1.CreateOptions{})
}

func buildService(ch *cmacme.Challenge) (*corev1.Service, error) {
	podLabels := podLabels(ch)
	service := &corev1.Service{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "cm-acme-http-solver-",
			Namespace:    ch.Namespace,
			Labels:       podLabels,
			Annotations: map[string]string{
				"auth.istio.io/8089": "NONE",
			},
			OwnerReferences: []metav1.OwnerReference{*metav1.NewControllerRef(ch, challengeGvk)},
		},
		Spec: corev1.ServiceSpec{
			Type: corev1.ServiceTypeNodePort,
			Ports: []corev1.ServicePort{
				{
					Name:       "http",
					Port:       acmeSolverListenPort,
					TargetPort: intstr.FromInt(acmeSolverListenPort),
				},
			},
			Selector: podLabels,
		},
	}

	// checking for presence of http01 config and if set serviceType is set, override our default (NodePort)
	serviceType, err := getServiceType(ch)
	if err != nil {
		return nil, err
	}
	service.Spec.Type = serviceType

	return service, nil
}

func (s *Solver) cleanupServices(ctx context.Context, ch *cmacme.Challenge) error {
	log := logf.FromContext(ctx, "cleanupPods")

	services, err := s.getServicesForChallenge(ctx, ch)
	if err != nil {
		return err
	}
	var errs []error
	for _, service := range services {
		log := logf.WithRelatedResource(log, service).V(logf.DebugLevel)
		log.V(logf.DebugLevel).Info("deleting service resource")

		err := s.Client.CoreV1().Services(service.Namespace).Delete(ctx, service.Name, metav1.DeleteOptions{})
		if err != nil {
			log.V(logf.WarnLevel).Info("failed to delete pod resource", "error", err)
			errs = append(errs, err)
			continue
		}
		log.V(logf.DebugLevel).Info("successfully deleted pod resource")
	}
	return utilerrors.NewAggregate(errs)
}
<commit_msg>Fix regression in solver service builder - default service type should be NodePort rather than empty<commit_after>/*
Copyright 2020 The cert-manager 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 http

import (
	"context"
	"fmt"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/selection"
	utilerrors "k8s.io/apimachinery/pkg/util/errors"
	"k8s.io/apimachinery/pkg/util/intstr"

	cmacme "github.com/jetstack/cert-manager/pkg/apis/acme/v1"
	logf "github.com/jetstack/cert-manager/pkg/logs"
)

func (s *Solver) ensureService(ctx context.Context, ch *cmacme.Challenge) (*corev1.Service, error) {
	log := logf.FromContext(ctx).WithName("ensureService")

	log.V(logf.DebugLevel).Info("checking for existing HTTP01 solver services for challenge")
	existingServices, err := s.getServicesForChallenge(ctx, ch)
	if err != nil {
		return nil, err
	}
	if len(existingServices) == 1 {
		logf.WithRelatedResource(log, existingServices[0]).Info("found one existing HTTP01 solver Service for challenge resource")
		return existingServices[0], nil
	}
	if len(existingServices) > 1 {
		log.V(logf.DebugLevel).Info("multiple challenge solver services found for challenge. cleaning up all existing services.")
		err := s.cleanupServices(ctx, ch)
		if err != nil {
			return nil, err
		}
		return nil, fmt.Errorf("multiple existing challenge solver services found and cleaned up. retrying challenge sync")
	}

	log.V(logf.DebugLevel).Info("creating HTTP01 challenge solver service")
	return s.createService(ctx, ch)
}

// getServicesForChallenge returns a list of services that were created to solve
// http challenges for the given domain
func (s *Solver) getServicesForChallenge(ctx context.Context, ch *cmacme.Challenge) ([]*corev1.Service, error) {
	log := logf.FromContext(ctx).WithName("getServicesForChallenge")

	podLabels := podLabels(ch)
	selector := labels.NewSelector()
	for key, val := range podLabels {
		req, err := labels.NewRequirement(key, selection.Equals, []string{val})
		if err != nil {
			return nil, err
		}
		selector = selector.Add(*req)
	}

	serviceList, err := s.serviceLister.Services(ch.Namespace).List(selector)
	if err != nil {
		return nil, err
	}

	var relevantServices []*corev1.Service
	for _, service := range serviceList {
		if !metav1.IsControlledBy(service, ch) {
			logf.WithRelatedResource(log, service).Info("found existing solver pod for this challenge resource, however " +
				"it does not have an appropriate OwnerReference referencing this challenge. Skipping it altogether.")
			continue
		}
		relevantServices = append(relevantServices, service)
	}

	return relevantServices, nil
}

// createService will create the service required to solve this challenge
// in the target API server.
func (s *Solver) createService(ctx context.Context, ch *cmacme.Challenge) (*corev1.Service, error) {
	svc, err := buildService(ch)
	if err != nil {
		return nil, err
	}
	return s.Client.CoreV1().Services(ch.Namespace).Create(ctx, svc, metav1.CreateOptions{})
}

func buildService(ch *cmacme.Challenge) (*corev1.Service, error) {
	podLabels := podLabels(ch)
	service := &corev1.Service{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "cm-acme-http-solver-",
			Namespace:    ch.Namespace,
			Labels:       podLabels,
			Annotations: map[string]string{
				"auth.istio.io/8089": "NONE",
			},
			OwnerReferences: []metav1.OwnerReference{*metav1.NewControllerRef(ch, challengeGvk)},
		},
		Spec: corev1.ServiceSpec{
			Type: corev1.ServiceTypeNodePort,
			Ports: []corev1.ServicePort{
				{
					Name:       "http",
					Port:       acmeSolverListenPort,
					TargetPort: intstr.FromInt(acmeSolverListenPort),
				},
			},
			Selector: podLabels,
		},
	}

	// checking for presence of http01 config and if set serviceType is set, override our default (NodePort)
	serviceType, err := getServiceType(ch)
	if err != nil {
		return nil, err
	}
	if serviceType != "" {
		service.Spec.Type = serviceType
	}

	return service, nil
}

func (s *Solver) cleanupServices(ctx context.Context, ch *cmacme.Challenge) error {
	log := logf.FromContext(ctx, "cleanupPods")

	services, err := s.getServicesForChallenge(ctx, ch)
	if err != nil {
		return err
	}
	var errs []error
	for _, service := range services {
		log := logf.WithRelatedResource(log, service).V(logf.DebugLevel)
		log.V(logf.DebugLevel).Info("deleting service resource")

		err := s.Client.CoreV1().Services(service.Namespace).Delete(ctx, service.Name, metav1.DeleteOptions{})
		if err != nil {
			log.V(logf.WarnLevel).Info("failed to delete pod resource", "error", err)
			errs = append(errs, err)
			continue
		}
		log.V(logf.DebugLevel).Info("successfully deleted pod resource")
	}
	return utilerrors.NewAggregate(errs)
}
<|endoftext|>


<commit_before>/*
Package bloom provides data structures and methods for creating Bloom filters.

A Bloom filter is a representation of a set of _n_ items, where the main
requirement is to make membership queries; _i.e._, whether an item is a
member of a set.

A Bloom filter has two parameters: _m_, a maximum size (typically a reasonably large
multiple of the cardinality of the set to represent) and _k_, the number of hashing
functions on elements of the set. (The actual hashing functions are important, too,
but this is not a parameter for this implementation). A Bloom filter is backed by
a BitSet; a key is represented in the filter by setting the bits at each value of the
hashing functions (modulo _m_). Set membership is done by _testing_ whether the
bits at each value of the hashing functions (again, modulo _m_) are set. If so,
the item is in the set. If the item is actually in the set, a Bloom filter will
never fail (the true positive rate is 1.0); but it is susceptible to false
positives. The art is to choose _k_ and _m_ correctly.

In this implementation, the hashing functions used is a local version of FNV,
a non-cryptographic hashing function, seeded with the index number of
the kth hashing function.

This implementation accepts keys for setting as testing as []byte. Thus, to
add a string item, "Love":

    uint n = 1000
    filter := bloom.New(20*n, 5) // load of 20, 5 keys
    filter.Add([]byte("Love"))

Similarly, to test if "Love" is in bloom:

    if filter.Test([]byte("Love"))

For numeric data, I recommend that you look into the binary/encoding library. But,
for example, to add a uint32 to the filter:

    i := uint32(100)
    n1 := make([]byte,4)
    binary.BigEndian.PutUint32(n1,i)
    f.Add(n1)

Finally, there is a method to estimate the false positive rate of a particular
Bloom filter for a set of size _n_:

    if filter.EstimateFalsePositiveRate(1000) > 0.001

Given the particular hashing scheme, it's best to be empirical about this. Note
that estimating the FP rate will clear the Bloom filter.
*/
package bloom

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"math"

	"github.com/spaolacci/murmur3"
	"github.com/willf/bitset"
)

// A BloomFilter is a representation of a set of _n_ items, where the main
// requirement is to make membership queries; _i.e._, whether an item is a
// member of a set.
type BloomFilter struct {
	m uint
	k uint
	b *bitset.BitSet
}

// New creates a new Bloom filter with _m_ bits and _k_ hashing functions
func New(m uint, k uint) *BloomFilter {
	return &BloomFilter{m, k, bitset.New(m)}
}

// From creates a new Bloom filter with len(_data_) * 64 bits and _k_ hashing
// functions. The data slice is not going to be reset.
func From(data []uint64, k uint) *BloomFilter {
	m := uint(len(data) * 64)
	return &BloomFilter{m, k, bitset.From(data)}
}

// baseHashes returns the four hash values of data that are used to create k
// hashes
func baseHashes(data []byte) [4]uint64 {
	a1 := []byte{1} // to grab another bit of data
	hasher := murmur3.New128()
	hasher.Write(data)
	v1, v2 := hasher.Sum128()
	hasher.Write(a1)
	v3, v4 := hasher.Sum128()
	return [4]uint64{
		v1, v2, v3, v4,
	}
}

// location returns the ith hashed location using the four base hash values
func (f *BloomFilter) location(h [4]uint64, i uint) (location uint) {
	ii := uint64(i)
	location = uint((h[ii%2] + ii*h[2+(((ii+(ii%2))%4)/2)]) % uint64(f.m))
	return
}

// EstimateParameters estimates requirements for m and k.
// Based on https://bitbucket.org/ww/bloom/src/829aa19d01d9/bloom.go
// used with permission.
func EstimateParameters(n uint, p float64) (m uint, k uint) {
	m = uint(math.Ceil(-1 * float64(n) * math.Log(p) / math.Pow(math.Log(2), 2)))
	k = uint(math.Ceil(math.Log(2) * float64(m) / float64(n)))
	return
}

// NewWithEstimates creates a new Bloom filter for about n items with fp
// false positive rate
func NewWithEstimates(n uint, fp float64) *BloomFilter {
	m, k := EstimateParameters(n, fp)
	return New(m, k)
}

// Cap returns the capacity, _m_, of a Bloom filter
func (f *BloomFilter) Cap() uint {
	return f.m
}

// K returns the number of hash functions used in the BloomFilter
func (f *BloomFilter) K() uint {
	return f.k
}

// Add data to the Bloom Filter. Returns the filter (allows chaining)
func (f *BloomFilter) Add(data []byte) *BloomFilter {
	h := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		f.b.Set(f.location(h, i))
	}
	return f
}

// Merge the data from two Bloom Filters.
func (f *BloomFilter) Merge(g *BloomFilter) error {
	// Make sure the m's and k's are the same, otherwise merging has no real use.
	if f.m != g.m {
		return fmt.Errorf("m's don't match: %d != %d", f.m, g.m)
	}

	if f.k != g.k {
		return fmt.Errorf("k's don't match: %d != %d", f.m, g.m)
	}

	f.b.InPlaceUnion(g.b)
	return nil
}

// Copy creates a copy of a Bloom filter.
func (f *BloomFilter) Copy() *BloomFilter {
	fc := New(f.m, f.k)
	fc.Merge(f)
	return fc
}

// AddString to the Bloom Filter. Returns the filter (allows chaining)
func (f *BloomFilter) AddString(data string) *BloomFilter {
	return f.Add([]byte(data))
}

// Test returns true if the data is in the BloomFilter, false otherwise.
// If true, the result might be a false positive. If false, the data
// is definitely not in the set.
func (f *BloomFilter) Test(data []byte) bool {
	h := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		if !f.b.Test(f.location(h, i)) {
			return false
		}
	}
	return true
}

// TestString returns true if the string is in the BloomFilter, false otherwise.
// If true, the result might be a false positive. If false, the data
// is definitely not in the set.
func (f *BloomFilter) TestString(data string) bool {
	return f.Test([]byte(data))
}

// TestAndAdd is the equivalent to calling Test(data) then Add(data).
// Returns the result of Test.
func (f *BloomFilter) TestAndAdd(data []byte) bool {
	present := true
	h := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		l := f.location(h, i)
		if !f.b.Test(l) {
			present = false
		}
		f.b.Set(l)
	}
	return present
}

// TestAndAddString is the equivalent to calling Test(string) then Add(string).
// Returns the result of Test.
func (f *BloomFilter) TestAndAddString(data string) bool {
	return f.TestAndAdd([]byte(data))
}

// ClearAll clears all the data in a Bloom filter, removing all keys
func (f *BloomFilter) ClearAll() *BloomFilter {
	f.b.ClearAll()
	return f
}

// EstimateFalsePositiveRate returns, for a BloomFilter with a estimate of m bits
// and k hash functions, what the false positive rate will be
// while storing n entries; runs 100,000 tests. This is an empirical
// test using integers as keys. As a side-effect, it clears the BloomFilter.
func (f *BloomFilter) EstimateFalsePositiveRate(n uint) (fpRate float64) {
	rounds := uint32(100000)
	f.ClearAll()
	n1 := make([]byte, 4)
	for i := uint32(0); i < uint32(n); i++ {
		binary.BigEndian.PutUint32(n1, i)
		f.Add(n1)
	}
	fp := 0
	// test for number of rounds
	for i := uint32(0); i < rounds; i++ {
		binary.BigEndian.PutUint32(n1, i+uint32(n)+1)
		if f.Test(n1) {
			//fmt.Printf("%v failed.\n", i+uint32(n)+1)
			fp++
		}
	}
	fpRate = float64(fp) / (float64(rounds))
	f.ClearAll()
	return
}

// bloomFilterJSON is an unexported type for marshaling/unmarshaling BloomFilter struct.
type bloomFilterJSON struct {
	M uint           `json:"m"`
	K uint           `json:"k"`
	B *bitset.BitSet `json:"b"`
}

// MarshalJSON implements json.Marshaler interface.
func (f *BloomFilter) MarshalJSON() ([]byte, error) {
	return json.Marshal(bloomFilterJSON{f.m, f.k, f.b})
}

// UnmarshalJSON implements json.Unmarshaler interface.
func (f *BloomFilter) UnmarshalJSON(data []byte) error {
	var j bloomFilterJSON
	err := json.Unmarshal(data, &j)
	if err != nil {
		return err
	}
	f.m = j.M
	f.k = j.K
	f.b = j.B
	return nil
}

// WriteTo writes a binary representation of the BloomFilter to an i/o stream.
// It returns the number of bytes written.
func (f *BloomFilter) WriteTo(stream io.Writer) (int64, error) {
	err := binary.Write(stream, binary.BigEndian, uint64(f.m))
	if err != nil {
		return 0, err
	}
	err = binary.Write(stream, binary.BigEndian, uint64(f.k))
	if err != nil {
		return 0, err
	}
	numBytes, err := f.b.WriteTo(stream)
	return numBytes + int64(2*binary.Size(uint64(0))), err
}

// ReadFrom reads a binary representation of the BloomFilter (such as might
// have been written by WriteTo()) from an i/o stream. It returns the number
// of bytes read.
func (f *BloomFilter) ReadFrom(stream io.Reader) (int64, error) {
	var m, k uint64
	err := binary.Read(stream, binary.BigEndian, &m)
	if err != nil {
		return 0, err
	}
	err = binary.Read(stream, binary.BigEndian, &k)
	if err != nil {
		return 0, err
	}
	b := &bitset.BitSet{}
	numBytes, err := b.ReadFrom(stream)
	if err != nil {
		return 0, err
	}
	f.m = uint(m)
	f.k = uint(k)
	f.b = b
	return numBytes + int64(2*binary.Size(uint64(0))), nil
}

// GobEncode implements gob.GobEncoder interface.
func (f *BloomFilter) GobEncode() ([]byte, error) {
	var buf bytes.Buffer
	_, err := f.WriteTo(&buf)
	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

// GobDecode implements gob.GobDecoder interface.
func (f *BloomFilter) GobDecode(data []byte) error {
	buf := bytes.NewBuffer(data)
	_, err := f.ReadFrom(buf)

	return err
}

// Equal tests for the equality of two Bloom filters
func (f *BloomFilter) Equal(g *BloomFilter) bool {
	return f.m == g.m && f.k == g.k && f.b.Equal(g.b)
}
<commit_msg>minor fix in the documentation<commit_after>/*
Package bloom provides data structures and methods for creating Bloom filters.

A Bloom filter is a representation of a set of _n_ items, where the main
requirement is to make membership queries; _i.e._, whether an item is a
member of a set.

A Bloom filter has two parameters: _m_, a maximum size (typically a reasonably large
multiple of the cardinality of the set to represent) and _k_, the number of hashing
functions on elements of the set. (The actual hashing functions are important, too,
but this is not a parameter for this implementation). A Bloom filter is backed by
a BitSet; a key is represented in the filter by setting the bits at each value of the
hashing functions (modulo _m_). Set membership is done by _testing_ whether the
bits at each value of the hashing functions (again, modulo _m_) are set. If so,
the item is in the set. If the item is actually in the set, a Bloom filter will
never fail (the true positive rate is 1.0); but it is susceptible to false
positives. The art is to choose _k_ and _m_ correctly.

In this implementation, the hashing functions used is murmurhash,
a non-cryptographic hashing function.

This implementation accepts keys for setting as testing as []byte. Thus, to
add a string item, "Love":

    uint n = 1000
    filter := bloom.New(20*n, 5) // load of 20, 5 keys
    filter.Add([]byte("Love"))

Similarly, to test if "Love" is in bloom:

    if filter.Test([]byte("Love"))

For numeric data, I recommend that you look into the binary/encoding library. But,
for example, to add a uint32 to the filter:

    i := uint32(100)
    n1 := make([]byte,4)
    binary.BigEndian.PutUint32(n1,i)
    f.Add(n1)

Finally, there is a method to estimate the false positive rate of a particular
Bloom filter for a set of size _n_:

    if filter.EstimateFalsePositiveRate(1000) > 0.001

Given the particular hashing scheme, it's best to be empirical about this. Note
that estimating the FP rate will clear the Bloom filter.
*/
package bloom

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"math"

	"github.com/spaolacci/murmur3"
	"github.com/willf/bitset"
)

// A BloomFilter is a representation of a set of _n_ items, where the main
// requirement is to make membership queries; _i.e._, whether an item is a
// member of a set.
type BloomFilter struct {
	m uint
	k uint
	b *bitset.BitSet
}

// New creates a new Bloom filter with _m_ bits and _k_ hashing functions
func New(m uint, k uint) *BloomFilter {
	return &BloomFilter{m, k, bitset.New(m)}
}

// From creates a new Bloom filter with len(_data_) * 64 bits and _k_ hashing
// functions. The data slice is not going to be reset.
func From(data []uint64, k uint) *BloomFilter {
	m := uint(len(data) * 64)
	return &BloomFilter{m, k, bitset.From(data)}
}

// baseHashes returns the four hash values of data that are used to create k
// hashes
func baseHashes(data []byte) [4]uint64 {
	a1 := []byte{1} // to grab another bit of data
	hasher := murmur3.New128()
	hasher.Write(data)
	v1, v2 := hasher.Sum128()
	hasher.Write(a1)
	v3, v4 := hasher.Sum128()
	return [4]uint64{
		v1, v2, v3, v4,
	}
}

// location returns the ith hashed location using the four base hash values
func (f *BloomFilter) location(h [4]uint64, i uint) (location uint) {
	ii := uint64(i)
	location = uint((h[ii%2] + ii*h[2+(((ii+(ii%2))%4)/2)]) % uint64(f.m))
	return
}

// EstimateParameters estimates requirements for m and k.
// Based on https://bitbucket.org/ww/bloom/src/829aa19d01d9/bloom.go
// used with permission.
func EstimateParameters(n uint, p float64) (m uint, k uint) {
	m = uint(math.Ceil(-1 * float64(n) * math.Log(p) / math.Pow(math.Log(2), 2)))
	k = uint(math.Ceil(math.Log(2) * float64(m) / float64(n)))
	return
}

// NewWithEstimates creates a new Bloom filter for about n items with fp
// false positive rate
func NewWithEstimates(n uint, fp float64) *BloomFilter {
	m, k := EstimateParameters(n, fp)
	return New(m, k)
}

// Cap returns the capacity, _m_, of a Bloom filter
func (f *BloomFilter) Cap() uint {
	return f.m
}

// K returns the number of hash functions used in the BloomFilter
func (f *BloomFilter) K() uint {
	return f.k
}

// Add data to the Bloom Filter. Returns the filter (allows chaining)
func (f *BloomFilter) Add(data []byte) *BloomFilter {
	h := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		f.b.Set(f.location(h, i))
	}
	return f
}

// Merge the data from two Bloom Filters.
func (f *BloomFilter) Merge(g *BloomFilter) error {
	// Make sure the m's and k's are the same, otherwise merging has no real use.
	if f.m != g.m {
		return fmt.Errorf("m's don't match: %d != %d", f.m, g.m)
	}

	if f.k != g.k {
		return fmt.Errorf("k's don't match: %d != %d", f.m, g.m)
	}

	f.b.InPlaceUnion(g.b)
	return nil
}

// Copy creates a copy of a Bloom filter.
func (f *BloomFilter) Copy() *BloomFilter {
	fc := New(f.m, f.k)
	fc.Merge(f)
	return fc
}

// AddString to the Bloom Filter. Returns the filter (allows chaining)
func (f *BloomFilter) AddString(data string) *BloomFilter {
	return f.Add([]byte(data))
}

// Test returns true if the data is in the BloomFilter, false otherwise.
// If true, the result might be a false positive. If false, the data
// is definitely not in the set.
func (f *BloomFilter) Test(data []byte) bool {
	h := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		if !f.b.Test(f.location(h, i)) {
			return false
		}
	}
	return true
}

// TestString returns true if the string is in the BloomFilter, false otherwise.
// If true, the result might be a false positive. If false, the data
// is definitely not in the set.
func (f *BloomFilter) TestString(data string) bool {
	return f.Test([]byte(data))
}

// TestAndAdd is the equivalent to calling Test(data) then Add(data).
// Returns the result of Test.
func (f *BloomFilter) TestAndAdd(data []byte) bool {
	present := true
	h := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		l := f.location(h, i)
		if !f.b.Test(l) {
			present = false
		}
		f.b.Set(l)
	}
	return present
}

// TestAndAddString is the equivalent to calling Test(string) then Add(string).
// Returns the result of Test.
func (f *BloomFilter) TestAndAddString(data string) bool {
	return f.TestAndAdd([]byte(data))
}

// ClearAll clears all the data in a Bloom filter, removing all keys
func (f *BloomFilter) ClearAll() *BloomFilter {
	f.b.ClearAll()
	return f
}

// EstimateFalsePositiveRate returns, for a BloomFilter with a estimate of m bits
// and k hash functions, what the false positive rate will be
// while storing n entries; runs 100,000 tests. This is an empirical
// test using integers as keys. As a side-effect, it clears the BloomFilter.
func (f *BloomFilter) EstimateFalsePositiveRate(n uint) (fpRate float64) {
	rounds := uint32(100000)
	f.ClearAll()
	n1 := make([]byte, 4)
	for i := uint32(0); i < uint32(n); i++ {
		binary.BigEndian.PutUint32(n1, i)
		f.Add(n1)
	}
	fp := 0
	// test for number of rounds
	for i := uint32(0); i < rounds; i++ {
		binary.BigEndian.PutUint32(n1, i+uint32(n)+1)
		if f.Test(n1) {
			//fmt.Printf("%v failed.\n", i+uint32(n)+1)
			fp++
		}
	}
	fpRate = float64(fp) / (float64(rounds))
	f.ClearAll()
	return
}

// bloomFilterJSON is an unexported type for marshaling/unmarshaling BloomFilter struct.
type bloomFilterJSON struct {
	M uint           `json:"m"`
	K uint           `json:"k"`
	B *bitset.BitSet `json:"b"`
}

// MarshalJSON implements json.Marshaler interface.
func (f *BloomFilter) MarshalJSON() ([]byte, error) {
	return json.Marshal(bloomFilterJSON{f.m, f.k, f.b})
}

// UnmarshalJSON implements json.Unmarshaler interface.
func (f *BloomFilter) UnmarshalJSON(data []byte) error {
	var j bloomFilterJSON
	err := json.Unmarshal(data, &j)
	if err != nil {
		return err
	}
	f.m = j.M
	f.k = j.K
	f.b = j.B
	return nil
}

// WriteTo writes a binary representation of the BloomFilter to an i/o stream.
// It returns the number of bytes written.
func (f *BloomFilter) WriteTo(stream io.Writer) (int64, error) {
	err := binary.Write(stream, binary.BigEndian, uint64(f.m))
	if err != nil {
		return 0, err
	}
	err = binary.Write(stream, binary.BigEndian, uint64(f.k))
	if err != nil {
		return 0, err
	}
	numBytes, err := f.b.WriteTo(stream)
	return numBytes + int64(2*binary.Size(uint64(0))), err
}

// ReadFrom reads a binary representation of the BloomFilter (such as might
// have been written by WriteTo()) from an i/o stream. It returns the number
// of bytes read.
func (f *BloomFilter) ReadFrom(stream io.Reader) (int64, error) {
	var m, k uint64
	err := binary.Read(stream, binary.BigEndian, &m)
	if err != nil {
		return 0, err
	}
	err = binary.Read(stream, binary.BigEndian, &k)
	if err != nil {
		return 0, err
	}
	b := &bitset.BitSet{}
	numBytes, err := b.ReadFrom(stream)
	if err != nil {
		return 0, err
	}
	f.m = uint(m)
	f.k = uint(k)
	f.b = b
	return numBytes + int64(2*binary.Size(uint64(0))), nil
}

// GobEncode implements gob.GobEncoder interface.
func (f *BloomFilter) GobEncode() ([]byte, error) {
	var buf bytes.Buffer
	_, err := f.WriteTo(&buf)
	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

// GobDecode implements gob.GobDecoder interface.
func (f *BloomFilter) GobDecode(data []byte) error {
	buf := bytes.NewBuffer(data)
	_, err := f.ReadFrom(buf)

	return err
}

// Equal tests for the equality of two Bloom filters
func (f *BloomFilter) Equal(g *BloomFilter) bool {
	return f.m == g.m && f.k == g.k && f.b.Equal(g.b)
}
<|endoftext|>


<commit_before>/*
Copyright AppsCode Inc. and Contributors

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 tableconvertor

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
	"time"

	"kmodules.xyz/resource-metadata/pkg/tableconvertor/printers"

	"github.com/Masterminds/sprig/v3"
	prom_op "github.com/prometheus-operator/prometheus-operator/pkg/apis/monitoring/v1"
	"gomodules.xyz/jsonpath"
	core "k8s.io/api/core/v1"
	rbac "k8s.io/api/rbac/v1"
	metatable "k8s.io/apimachinery/pkg/api/meta/table"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/duration"
	kubedb "kubedb.dev/apimachinery/apis/kubedb/v1alpha2"
)

var templateFns = sprig.TxtFuncMap()

func init() {
	templateFns["jp"] = jsonpathFn
	templateFns["k8s_fmt_selector"] = formatLabelSelectorFn
	templateFns["k8s_fmt_label"] = formatLabelsFn
	templateFns["k8s_age"] = ageFn
	templateFns["k8s_svc_ports"] = servicePortsFn
	templateFns["k8s_container_ports"] = containerPortFn
	templateFns["k8s_container_images"] = containerImagesFn
	templateFns["k8s_volumes"] = volumesFn
	templateFns["k8s_volumeMounts"] = volumeMountsFn
	templateFns["k8s_duration"] = durationFn
	templateFns["fmt_list"] = fmtListFn
	templateFns["prom_ns_selector"] = promNamespaceSelectorFn
	templateFns["map_key_count"] = mapKeyCountFn
	templateFns["kubedb_db_mode"] = kubedbDBModeFn
	templateFns["kubedb_db_replicas"] = kubedbDBReplicasFn
	templateFns["kubedb_db_resources"] = kubedbDBResourcesFn
	templateFns["rbac_subjects"] = rbacSubjects
	templateFns["cert_validity"] = certificateValidity
}

func jsonpathFn(expr string, data interface{}, jsonoutput ...bool) (interface{}, error) {
	enableJSONoutput := len(jsonoutput) > 0 && jsonoutput[0]

	jp := jsonpath.New("jp")
	if err := jp.Parse(expr); err != nil {
		return nil, fmt.Errorf("unrecognized column definition %q", expr)
	}
	jp.AllowMissingKeys(true)
	jp.EnableJSONOutput(enableJSONoutput)

	var buf bytes.Buffer
	err := jp.Execute(&buf, data)
	if err != nil {
		return nil, err
	}

	if enableJSONoutput {
		var v []interface{}
		err = json.Unmarshal(buf.Bytes(), &v)
		return v, err
	}
	return buf.String(), err
}

func formatLabelSelectorFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}
	var sel metav1.LabelSelector
	err := json.Unmarshal([]byte(data), &sel)
	if err != nil {
		return "", err
	}
	return metav1.FormatLabelSelector(&sel), nil
}

func formatLabelsFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}
	var label map[string]string
	err := json.Unmarshal([]byte(data), &label)
	if err != nil {
		return "", err
	}
	return labels.FormatLabels(label), nil
}

func ageFn(data string) (string, error) {
	if data == "" {
		return "", nil
	}
	var timestamp metav1.Time
	err := timestamp.UnmarshalQueryParameter(data)
	if err != nil {
		return "", err
	}
	return metatable.ConvertToHumanReadableDateType(timestamp), nil
}

func servicePortsFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}
	var ports []core.ServicePort
	err := json.Unmarshal([]byte(data), &ports)
	if err != nil {
		return "", err
	}
	return printers.MakeServicePortString(ports), nil
}

func containerPortFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}
	var ports []core.ContainerPort
	err := json.Unmarshal([]byte(data), &ports)
	if err != nil {
		return "", err
	}
	pieces := make([]string, len(ports))
	for ix := range ports {
		port := &ports[ix]
		pieces[ix] = fmt.Sprintf("%d/%s", port.ContainerPort, port.Protocol)
		if port.HostPort > 0 {
			pieces[ix] = fmt.Sprintf("%d:%d/%s", port.ContainerPort, port.HostPort, port.Protocol)
		}
	}
	return strings.Join(pieces, ","), nil
}

func volumesFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}
	var volumes []core.Volume
	err := json.Unmarshal([]byte(data), &volumes)
	if err != nil {
		return "", err
	}
	ss := "["
	for i := range volumes {
		ss += describeVolume(volumes[i])
		if i < len(volumes)-1 {
			ss += ","
		}
	}
	ss += "]"
	return ss, nil
}

func volumeMountsFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}
	var mounts []core.VolumeMount
	ss := make([]string, 0)
	err := json.Unmarshal([]byte(data), &mounts)
	if err != nil {
		return "", err
	}

	for i := range mounts {
		mnt := fmt.Sprintf("%s:%s", mounts[i].Name, mounts[i].MountPath)
		if mounts[i].SubPath != "" {
			mnt = fmt.Sprintf("%s:%s:%s", mounts[i].Name, mounts[i].MountPath, mounts[i].SubPath)
		}
		ss = append(ss, mnt)
	}
	return strings.Join(ss, "\n"), nil
}

func fmtListFn(data string) (string, error) {
	// Return empty list if the data is empty. This helps to avoid object parsing error.
	// ref: https://stackoverflow.com/a/18419503
	if len(data) == 0 {
		return "[]", nil
	}
	return data, nil
}

type promNamespaceSelector struct {
	metav1.ObjectMeta `json:"metadata"`
	Spec              promNamespaceSelectorSpec `json:"spec"`
}
type promNamespaceSelectorSpec struct {
	NamespaceSelector *prom_op.NamespaceSelector `json:"namespaceSelector,omitempty"`
}

func promNamespaceSelectorFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}
	var selOpts promNamespaceSelector
	err := json.Unmarshal([]byte(data), &selOpts)
	if err != nil {
		return "", err
	}
	// If selector field is empty, then all namespaces are the targets.
	if selOpts.Spec.NamespaceSelector == nil {
		return "All", nil
	} else if len(selOpts.Spec.NamespaceSelector.MatchNames) != 0 {
		// If an array of namespace is provided, then those namespaces are the target
		return strings.Join(selOpts.Spec.NamespaceSelector.MatchNames, ", "), nil
	} else if !selOpts.Spec.NamespaceSelector.Any {
		// If "any: false" is set in the namespace selector field, only the object namespace is the target.
		return selOpts.Namespace, nil
	}
	return "", nil
}

func containerImagesFn(data string) (string, error) {
	var imagesBuffer bytes.Buffer
	if strings.TrimSpace(data) == "" {
		return "", nil
	}

	var containers []core.Container
	err := json.Unmarshal([]byte(data), &containers)
	if err != nil {
		return "", err
	}
	for i, container := range containers {
		imagesBuffer.WriteString(container.Image)
		if i != len(containers)-1 {
			imagesBuffer.WriteString(",")
		}
	}
	return imagesBuffer.String(), nil
}

func durationFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}

	// make the data a valid json
	ss := strings.Split(strings.TrimSuffix(data, ","), ",")
	jsonData := "["
	for i := range ss {
		jsonData += fmt.Sprintf("%q", ss[i])
		if i < len(ss)-1 {
			jsonData += ","
		}
	}
	jsonData += "]"

	var tt []metav1.Time
	err := json.Unmarshal([]byte(jsonData), &tt)
	if err != nil {
		return "", err
	}
	if len(tt) == 1 {
		// only start time is does exist
		return duration.HumanDuration(time.Since(tt[0].Time)), nil
	} else {
		return duration.HumanDuration(tt[1].Sub(tt[0].Time)), nil
	}
}

func mapKeyCountFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}

	var m map[string]string
	err := json.Unmarshal([]byte(data), &m)
	if err != nil {
		return "", err
	}

	return strconv.Itoa(len(m)), nil
}

func kubedbDBModeFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}
	var obj unstructured.Unstructured
	err := json.Unmarshal([]byte(data), &obj)
	if err != nil {
		return "", err
	}

	switch obj.GetKind() {
	case kubedb.ResourceKindMongoDB:
		db := new(kubedb.MongoDB)
		err := runtime.DefaultUnstructuredConverter.FromUnstructured(obj.UnstructuredContent(), db)
		if err != nil {
			return "", err
		}
		if db.Spec.ShardTopology != nil {
			return "Sharded", nil
		} else if db.Spec.ReplicaSet != nil {
			return "ReplicaSet", nil
		}
		return "Standalone", nil
	}
	return "", fmt.Errorf("failed to detectect database mode. Reason: Unknown database type")
}

func kubedbDBReplicasFn(data string) (string, error) {
	var obj unstructured.Unstructured
	err := json.Unmarshal([]byte(data), &obj)
	if err != nil {
		return "", err
	}

	switch obj.GetKind() {
	case kubedb.ResourceKindMongoDB:
		db := new(kubedb.MongoDB)
		err := runtime.DefaultUnstructuredConverter.FromUnstructured(obj.UnstructuredContent(), db)
		if err != nil {
			return "", err
		}
		if db.Spec.ShardTopology != nil {
			t := db.Spec.ShardTopology
			return fmt.Sprintf("%d, %d, %d", t.Shard.Replicas, t.ConfigServer.Replicas, t.Mongos.Replicas), nil
		} else if db.Spec.ReplicaSet != nil {
			return fmt.Sprintf("%d", *db.Spec.Replicas), nil
		} else {
			return "1", nil
		}
	}
	return "", fmt.Errorf("failed to detect replica number. Reason: Unknown database type")
}

func kubedbDBResourcesFn(data string) (string, error) {
	var obj unstructured.Unstructured
	err := json.Unmarshal([]byte(data), &obj)
	if err != nil {
		return "", err
	}

	switch obj.GetKind() {
	case kubedb.ResourceKindMongoDB:
		db := new(kubedb.MongoDB)
		err := runtime.DefaultUnstructuredConverter.FromUnstructured(obj.UnstructuredContent(), db)
		if err != nil {
			return "", err
		}
		cpu, memory, storage := mongoDBResources(db.Spec)
		return fmt.Sprintf("{%q:%q, %q:%q, %q:%q}", core.ResourceCPU, cpu, core.ResourceMemory, memory, core.ResourceStorage, storage), nil
	}
	return "", fmt.Errorf("failed to extract CPU information. Reason: Unknown database type")
}

func rbacSubjects(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}

	var subjects []rbac.Subject
	err := json.Unmarshal([]byte(data), &subjects)
	if err != nil {
		return "", err
	}
	var ss []string
	for i := range subjects {
		s := fmt.Sprintf("%s %s", subjects[i].Kind, subjects[i].Name)
		if subjects[i].Namespace != "" {
			s = fmt.Sprintf("%s %s/%s", subjects[i].Kind, subjects[i].Namespace, subjects[i].Name)
		}
		ss = append(ss, s)
	}
	return strings.Join(ss, ","), nil
}

func certificateValidity(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}

	certStatus := struct {
		NotBefore metav1.Time `json:"notBefore"`
		NotAfter  metav1.Time `json:"notAfter"`
	}{}
	err := json.Unmarshal([]byte(data), &certStatus)
	if err != nil {
		return "", err
	}

	if certStatus.NotBefore.After(time.Now()) {
		return "Not valid yet", nil
	} else if time.Now().After(certStatus.NotAfter.Time) {
		return "Expired", nil
	}
	return duration.HumanDuration(time.Until(certStatus.NotAfter.Time)), nil
}
<commit_msg>Add TxtFuncMap (#124)<commit_after>/*
Copyright AppsCode Inc. and Contributors

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 tableconvertor

import (
	"bytes"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"
	"text/template"
	"time"

	"kmodules.xyz/resource-metadata/pkg/tableconvertor/printers"

	"github.com/Masterminds/sprig/v3"
	prom_op "github.com/prometheus-operator/prometheus-operator/pkg/apis/monitoring/v1"
	"gomodules.xyz/jsonpath"
	core "k8s.io/api/core/v1"
	rbac "k8s.io/api/rbac/v1"
	metatable "k8s.io/apimachinery/pkg/api/meta/table"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/duration"
	kubedb "kubedb.dev/apimachinery/apis/kubedb/v1alpha2"
)

var templateFns = sprig.TxtFuncMap()

func init() {
	templateFns["jp"] = jsonpathFn
	templateFns["k8s_fmt_selector"] = formatLabelSelectorFn
	templateFns["k8s_fmt_label"] = formatLabelsFn
	templateFns["k8s_age"] = ageFn
	templateFns["k8s_svc_ports"] = servicePortsFn
	templateFns["k8s_container_ports"] = containerPortFn
	templateFns["k8s_container_images"] = containerImagesFn
	templateFns["k8s_volumes"] = volumesFn
	templateFns["k8s_volumeMounts"] = volumeMountsFn
	templateFns["k8s_duration"] = durationFn
	templateFns["fmt_list"] = fmtListFn
	templateFns["prom_ns_selector"] = promNamespaceSelectorFn
	templateFns["map_key_count"] = mapKeyCountFn
	templateFns["kubedb_db_mode"] = kubedbDBModeFn
	templateFns["kubedb_db_replicas"] = kubedbDBReplicasFn
	templateFns["kubedb_db_resources"] = kubedbDBResourcesFn
	templateFns["rbac_subjects"] = rbacSubjects
	templateFns["cert_validity"] = certificateValidity
}

// TxtFuncMap returns a 'text/template'.FuncMap
func TxtFuncMap() template.FuncMap {
	gfm := make(map[string]interface{}, len(templateFns))
	for k, v := range templateFns {
		gfm[k] = v
	}
	return gfm
}

func jsonpathFn(expr string, data interface{}, jsonoutput ...bool) (interface{}, error) {
	enableJSONoutput := len(jsonoutput) > 0 && jsonoutput[0]

	jp := jsonpath.New("jp")
	if err := jp.Parse(expr); err != nil {
		return nil, fmt.Errorf("unrecognized column definition %q", expr)
	}
	jp.AllowMissingKeys(true)
	jp.EnableJSONOutput(enableJSONoutput)

	var buf bytes.Buffer
	err := jp.Execute(&buf, data)
	if err != nil {
		return nil, err
	}

	if enableJSONoutput {
		var v []interface{}
		err = json.Unmarshal(buf.Bytes(), &v)
		return v, err
	}
	return buf.String(), err
}

func formatLabelSelectorFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}
	var sel metav1.LabelSelector
	err := json.Unmarshal([]byte(data), &sel)
	if err != nil {
		return "", err
	}
	return metav1.FormatLabelSelector(&sel), nil
}

func formatLabelsFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}
	var label map[string]string
	err := json.Unmarshal([]byte(data), &label)
	if err != nil {
		return "", err
	}
	return labels.FormatLabels(label), nil
}

func ageFn(data string) (string, error) {
	if data == "" {
		return "", nil
	}
	var timestamp metav1.Time
	err := timestamp.UnmarshalQueryParameter(data)
	if err != nil {
		return "", err
	}
	return metatable.ConvertToHumanReadableDateType(timestamp), nil
}

func servicePortsFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}
	var ports []core.ServicePort
	err := json.Unmarshal([]byte(data), &ports)
	if err != nil {
		return "", err
	}
	return printers.MakeServicePortString(ports), nil
}

func containerPortFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}
	var ports []core.ContainerPort
	err := json.Unmarshal([]byte(data), &ports)
	if err != nil {
		return "", err
	}
	pieces := make([]string, len(ports))
	for ix := range ports {
		port := &ports[ix]
		pieces[ix] = fmt.Sprintf("%d/%s", port.ContainerPort, port.Protocol)
		if port.HostPort > 0 {
			pieces[ix] = fmt.Sprintf("%d:%d/%s", port.ContainerPort, port.HostPort, port.Protocol)
		}
	}
	return strings.Join(pieces, ","), nil
}

func volumesFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}
	var volumes []core.Volume
	err := json.Unmarshal([]byte(data), &volumes)
	if err != nil {
		return "", err
	}
	ss := "["
	for i := range volumes {
		ss += describeVolume(volumes[i])
		if i < len(volumes)-1 {
			ss += ","
		}
	}
	ss += "]"
	return ss, nil
}

func volumeMountsFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}
	var mounts []core.VolumeMount
	ss := make([]string, 0)
	err := json.Unmarshal([]byte(data), &mounts)
	if err != nil {
		return "", err
	}

	for i := range mounts {
		mnt := fmt.Sprintf("%s:%s", mounts[i].Name, mounts[i].MountPath)
		if mounts[i].SubPath != "" {
			mnt = fmt.Sprintf("%s:%s:%s", mounts[i].Name, mounts[i].MountPath, mounts[i].SubPath)
		}
		ss = append(ss, mnt)
	}
	return strings.Join(ss, "\n"), nil
}

func fmtListFn(data string) (string, error) {
	// Return empty list if the data is empty. This helps to avoid object parsing error.
	// ref: https://stackoverflow.com/a/18419503
	if len(data) == 0 {
		return "[]", nil
	}
	return data, nil
}

type promNamespaceSelector struct {
	metav1.ObjectMeta `json:"metadata"`
	Spec              promNamespaceSelectorSpec `json:"spec"`
}
type promNamespaceSelectorSpec struct {
	NamespaceSelector *prom_op.NamespaceSelector `json:"namespaceSelector,omitempty"`
}

func promNamespaceSelectorFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}
	var selOpts promNamespaceSelector
	err := json.Unmarshal([]byte(data), &selOpts)
	if err != nil {
		return "", err
	}
	// If selector field is empty, then all namespaces are the targets.
	if selOpts.Spec.NamespaceSelector == nil {
		return "All", nil
	} else if len(selOpts.Spec.NamespaceSelector.MatchNames) != 0 {
		// If an array of namespace is provided, then those namespaces are the target
		return strings.Join(selOpts.Spec.NamespaceSelector.MatchNames, ", "), nil
	} else if !selOpts.Spec.NamespaceSelector.Any {
		// If "any: false" is set in the namespace selector field, only the object namespace is the target.
		return selOpts.Namespace, nil
	}
	return "", nil
}

func containerImagesFn(data string) (string, error) {
	var imagesBuffer bytes.Buffer
	if strings.TrimSpace(data) == "" {
		return "", nil
	}

	var containers []core.Container
	err := json.Unmarshal([]byte(data), &containers)
	if err != nil {
		return "", err
	}
	for i, container := range containers {
		imagesBuffer.WriteString(container.Image)
		if i != len(containers)-1 {
			imagesBuffer.WriteString(",")
		}
	}
	return imagesBuffer.String(), nil
}

func durationFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}

	// make the data a valid json
	ss := strings.Split(strings.TrimSuffix(data, ","), ",")
	jsonData := "["
	for i := range ss {
		jsonData += fmt.Sprintf("%q", ss[i])
		if i < len(ss)-1 {
			jsonData += ","
		}
	}
	jsonData += "]"

	var tt []metav1.Time
	err := json.Unmarshal([]byte(jsonData), &tt)
	if err != nil {
		return "", err
	}
	if len(tt) == 1 {
		// only start time is does exist
		return duration.HumanDuration(time.Since(tt[0].Time)), nil
	} else {
		return duration.HumanDuration(tt[1].Sub(tt[0].Time)), nil
	}
}

func mapKeyCountFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}

	var m map[string]string
	err := json.Unmarshal([]byte(data), &m)
	if err != nil {
		return "", err
	}

	return strconv.Itoa(len(m)), nil
}

func kubedbDBModeFn(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}
	var obj unstructured.Unstructured
	err := json.Unmarshal([]byte(data), &obj)
	if err != nil {
		return "", err
	}

	switch obj.GetKind() {
	case kubedb.ResourceKindMongoDB:
		db := new(kubedb.MongoDB)
		err := runtime.DefaultUnstructuredConverter.FromUnstructured(obj.UnstructuredContent(), db)
		if err != nil {
			return "", err
		}
		if db.Spec.ShardTopology != nil {
			return "Sharded", nil
		} else if db.Spec.ReplicaSet != nil {
			return "ReplicaSet", nil
		}
		return "Standalone", nil
	}
	return "", fmt.Errorf("failed to detectect database mode. Reason: Unknown database type")
}

func kubedbDBReplicasFn(data string) (string, error) {
	var obj unstructured.Unstructured
	err := json.Unmarshal([]byte(data), &obj)
	if err != nil {
		return "", err
	}

	switch obj.GetKind() {
	case kubedb.ResourceKindMongoDB:
		db := new(kubedb.MongoDB)
		err := runtime.DefaultUnstructuredConverter.FromUnstructured(obj.UnstructuredContent(), db)
		if err != nil {
			return "", err
		}
		if db.Spec.ShardTopology != nil {
			t := db.Spec.ShardTopology
			return fmt.Sprintf("%d, %d, %d", t.Shard.Replicas, t.ConfigServer.Replicas, t.Mongos.Replicas), nil
		} else if db.Spec.ReplicaSet != nil {
			return fmt.Sprintf("%d", *db.Spec.Replicas), nil
		} else {
			return "1", nil
		}
	}
	return "", fmt.Errorf("failed to detect replica number. Reason: Unknown database type")
}

func kubedbDBResourcesFn(data string) (string, error) {
	var obj unstructured.Unstructured
	err := json.Unmarshal([]byte(data), &obj)
	if err != nil {
		return "", err
	}

	switch obj.GetKind() {
	case kubedb.ResourceKindMongoDB:
		db := new(kubedb.MongoDB)
		err := runtime.DefaultUnstructuredConverter.FromUnstructured(obj.UnstructuredContent(), db)
		if err != nil {
			return "", err
		}
		cpu, memory, storage := mongoDBResources(db.Spec)
		return fmt.Sprintf("{%q:%q, %q:%q, %q:%q}", core.ResourceCPU, cpu, core.ResourceMemory, memory, core.ResourceStorage, storage), nil
	}
	return "", fmt.Errorf("failed to extract CPU information. Reason: Unknown database type")
}

func rbacSubjects(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}

	var subjects []rbac.Subject
	err := json.Unmarshal([]byte(data), &subjects)
	if err != nil {
		return "", err
	}
	var ss []string
	for i := range subjects {
		s := fmt.Sprintf("%s %s", subjects[i].Kind, subjects[i].Name)
		if subjects[i].Namespace != "" {
			s = fmt.Sprintf("%s %s/%s", subjects[i].Kind, subjects[i].Namespace, subjects[i].Name)
		}
		ss = append(ss, s)
	}
	return strings.Join(ss, ","), nil
}

func certificateValidity(data string) (string, error) {
	if strings.TrimSpace(data) == "" {
		return "", nil
	}

	certStatus := struct {
		NotBefore metav1.Time `json:"notBefore"`
		NotAfter  metav1.Time `json:"notAfter"`
	}{}
	err := json.Unmarshal([]byte(data), &certStatus)
	if err != nil {
		return "", err
	}

	if certStatus.NotBefore.After(time.Now()) {
		return "Not valid yet", nil
	} else if time.Now().After(certStatus.NotAfter.Time) {
		return "Expired", nil
	}
	return duration.HumanDuration(time.Until(certStatus.NotAfter.Time)), nil
}
<|endoftext|>


<commit_before>package amqp

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"github.com/RichardKnop/machinery/v1/brokers/errs"
	"github.com/RichardKnop/machinery/v1/brokers/iface"
	"github.com/RichardKnop/machinery/v1/common"
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"github.com/pkg/errors"
	"github.com/streadway/amqp"
)

type AMQPConnection struct {
	queueName    string
	connection   *amqp.Connection
	channel      *amqp.Channel
	queue        amqp.Queue
	confirmation <-chan amqp.Confirmation
	errorchan    <-chan *amqp.Error
	cleanup      chan struct{}
}

// Broker represents an AMQP broker
type Broker struct {
	common.Broker
	common.AMQPConnector
	processingWG sync.WaitGroup // use wait group to make sure task processing completes on interrupt signal

	connections      map[string]*AMQPConnection
	connectionsMutex sync.RWMutex
}

// New creates new Broker instance
func New(cnf *config.Config) iface.Broker {
	return &Broker{Broker: common.NewBroker(cnf), AMQPConnector: common.AMQPConnector{}, connections: make(map[string]*AMQPConnection)}
}

// StartConsuming enters a loop and waits for incoming messages
func (b *Broker) StartConsuming(consumerTag string, concurrency int, taskProcessor iface.TaskProcessor) (bool, error) {
	b.Broker.StartConsuming(consumerTag, concurrency, taskProcessor)

	queueName := taskProcessor.CustomQueue()
	if queueName == "" {
		queueName = b.GetConfig().DefaultQueue
	}

	conn, channel, queue, _, amqpCloseChan, err := b.Connect(
		b.GetConfig().Broker,
		b.GetConfig().TLSConfig,
		b.GetConfig().AMQP.Exchange,     // exchange name
		b.GetConfig().AMQP.ExchangeType, // exchange type
		queueName,                       // queue name
		true,                            // queue durable
		false,                           // queue delete when unused
		b.GetConfig().AMQP.BindingKey,   // queue binding key
		nil,                             // exchange declare args
		amqp.Table(b.GetConfig().AMQP.QueueDeclareArgs), // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		b.GetRetryFunc()(b.GetRetryStopChan())
		return b.GetRetry(), err
	}
	defer b.Close(channel, conn)

	if err = channel.Qos(
		b.GetConfig().AMQP.PrefetchCount,
		0,     // prefetch size
		false, // global
	); err != nil {
		return b.GetRetry(), fmt.Errorf("Channel qos error: %s", err)
	}

	deliveries, err := channel.Consume(
		queue.Name,  // queue
		consumerTag, // consumer tag
		false,       // auto-ack
		false,       // exclusive
		false,       // no-local
		false,       // no-wait
		nil,         // arguments
	)
	if err != nil {
		return b.GetRetry(), fmt.Errorf("Queue consume error: %s", err)
	}

	log.INFO.Print("[*] Waiting for messages. To exit press CTRL+C")

	if err := b.consume(deliveries, concurrency, taskProcessor, amqpCloseChan); err != nil {
		return b.GetRetry(), err
	}

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()

	return b.GetRetry(), nil
}

// StopConsuming quits the loop
func (b *Broker) StopConsuming() {
	b.Broker.StopConsuming()

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()
}

// GetOrOpenConnection will return a connection on a particular queue name. Open connections
// are saved to avoid having to reopen connection for multiple queues
func (b *Broker) GetOrOpenConnection(queueName string, queueBindingKey string, exchangeDeclareArgs, queueDeclareArgs, queueBindingArgs amqp.Table) (*AMQPConnection, error) {
	var err error

	b.connectionsMutex.Lock()
	defer b.connectionsMutex.Unlock()

	conn, ok := b.connections[queueName]
	if !ok {
		conn = &AMQPConnection{
			queueName: queueName,
			cleanup:   make(chan struct{}),
		}
		conn.connection, conn.channel, conn.queue, conn.confirmation, conn.errorchan, err = b.Connect(
			b.GetConfig().Broker,
			b.GetConfig().TLSConfig,
			b.GetConfig().AMQP.Exchange,     // exchange name
			b.GetConfig().AMQP.ExchangeType, // exchange type
			queueName,                       // queue name
			true,                            // queue durable
			false,                           // queue delete when unused
			queueBindingKey,                 // queue binding key
			exchangeDeclareArgs,             // exchange declare args
			queueDeclareArgs,                // queue declare args
			queueBindingArgs,                // queue binding args
		)
		if err != nil {
			return nil, errors.Wrapf(err, "Failed to connect to queue %s", queueName)
		}

		// Reconnect to the channel if it disconnects/errors out
		go func() {
			select {
			case err = <-conn.errorchan:
				log.INFO.Printf("Error occured on queue: %s. Reconnecting", queueName)
				b.connectionsMutex.Lock()
				delete(b.connections, queueName)
				b.connectionsMutex.Unlock()
				_, err := b.GetOrOpenConnection(queueName, queueBindingKey, exchangeDeclareArgs, queueDeclareArgs, queueBindingArgs)
				if err != nil {
					log.ERROR.Printf("Failed to reopen queue: %s.", queueName)
				}
			case <-conn.cleanup:
				return
			}
			return
		}()
		b.connections[queueName] = conn
	}
	return conn, nil
}

func (b *Broker) CloseConnections() error {
	b.connectionsMutex.Lock()
	defer b.connectionsMutex.Unlock()

	for key, conn := range b.connections {
		if err := b.Close(conn.channel, conn.connection); err != nil {
			log.ERROR.Print("Failed to close channel")
			return nil
		}
		close(conn.cleanup)
		delete(b.connections, key)
	}
	return nil
}

// Publish places a new message on the default queue
func (b *Broker) Publish(ctx context.Context, signature *tasks.Signature) error {
	// Adjust routing key (this decides which queue the message will be published to)
	b.AdjustRoutingKey(signature)

	msg, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// Check the ETA signature field, if it is set and it is in the future,
	// delay the task
	if signature.ETA != nil {
		now := time.Now().UTC()

		if signature.ETA.After(now) {
			delayMs := int64(signature.ETA.Sub(now) / time.Millisecond)

			return b.delay(signature, delayMs)
		}
	}

	queue := b.GetConfig().DefaultQueue
	bindingKey := b.GetConfig().AMQP.BindingKey // queue binding key
	if b.isDirectExchange() {
		queue = signature.RoutingKey
		bindingKey = signature.RoutingKey
	}

	connection, err := b.GetOrOpenConnection(
		queue,
		bindingKey, // queue binding key
		nil,        // exchange declare args
		amqp.Table(b.GetConfig().AMQP.QueueDeclareArgs), // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		return errors.Wrapf(err, "Failed to get a connection for queue %s", queue)
	}

	channel := connection.channel
	confirmsChan := connection.confirmation

	if err := channel.Publish(
		b.GetConfig().AMQP.Exchange, // exchange name
		signature.RoutingKey,        // routing key
		false,                       // mandatory
		false,                       // immediate
		amqp.Publishing{
			Headers:      amqp.Table(signature.Headers),
			ContentType:  "application/json",
			Body:         msg,
			Priority:     signature.Priority,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return errors.Wrap(err, "Failed to publish task")
	}

	confirmed := <-confirmsChan

	if confirmed.Ack {
		return nil
	}

	return fmt.Errorf("Failed delivery of delivery tag: %v", confirmed.DeliveryTag)
}

// consume takes delivered messages from the channel and manages a worker pool
// to process tasks concurrently
func (b *Broker) consume(deliveries <-chan amqp.Delivery, concurrency int, taskProcessor iface.TaskProcessor, amqpCloseChan <-chan *amqp.Error) error {
	pool := make(chan struct{}, concurrency)

	// initialize worker pool with maxWorkers workers
	go func() {
		for i := 0; i < concurrency; i++ {
			pool <- struct{}{}
		}
	}()

	// make channel with a capacity makes it become a buffered channel so that a worker which wants to
	// push an error to `errorsChan` doesn't need to be blocked while the for-loop is blocked waiting
	// a worker, that is, it avoids a possible deadlock
	errorsChan := make(chan error, 1)

	for {
		select {
		case amqpErr := <-amqpCloseChan:
			return amqpErr
		case err := <-errorsChan:
			return err
		case d := <-deliveries:
			if concurrency > 0 {
				// get worker from pool (blocks until one is available)
				<-pool
			}

			b.processingWG.Add(1)

			// Consume the task inside a gotourine so multiple tasks
			// can be processed concurrently
			go func() {
				if err := b.consumeOne(d, taskProcessor); err != nil {
					errorsChan <- err
				}

				b.processingWG.Done()

				if concurrency > 0 {
					// give worker back to pool
					pool <- struct{}{}
				}
			}()
		case <-b.GetStopChan():
			return nil
		}
	}
}

// consumeOne processes a single message using TaskProcessor
func (b *Broker) consumeOne(delivery amqp.Delivery, taskProcessor iface.TaskProcessor) error {
	if len(delivery.Body) == 0 {
		delivery.Nack(true, false)                     // multiple, requeue
		return errors.New("Received an empty message") // RabbitMQ down?
	}

	var multiple, requeue = false, false

	// Unmarshal message body into signature struct
	signature := new(tasks.Signature)
	decoder := json.NewDecoder(bytes.NewReader(delivery.Body))
	decoder.UseNumber()
	if err := decoder.Decode(signature); err != nil {
		delivery.Nack(multiple, requeue)
		return errs.NewErrCouldNotUnmarshaTaskSignature(delivery.Body, err)
	}

	// If the task is not registered, we nack it and requeue,
	// there might be different workers for processing specific tasks
	if !b.IsTaskRegistered(signature.Name) {
		if !delivery.Redelivered {
			requeue = true
			log.INFO.Printf("Task not registered with this worker. Requeing message: %s", delivery.Body)
		}
		if !signature.IgnoreWhenTaskNotRegistered {
			delivery.Nack(multiple, requeue)
		}
		return nil
	}

	log.INFO.Printf("Received new message: %s", delivery.Body)

	err := taskProcessor.Process(signature)
	delivery.Ack(multiple)
	return err
}

// delay a task by delayDuration miliseconds, the way it works is a new queue
// is created without any consumers, the message is then published to this queue
// with appropriate ttl expiration headers, after the expiration, it is sent to
// the proper queue with consumers
func (b *Broker) delay(signature *tasks.Signature, delayMs int64) error {
	if delayMs <= 0 {
		return errors.New("Cannot delay task by 0ms")
	}

	message, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// It's necessary to redeclare the queue each time (to zero its TTL timer).
	queueName := fmt.Sprintf(
		"delay.%d.%s.%s",
		delayMs, // delay duration in mileseconds
		b.GetConfig().AMQP.Exchange,
		signature.RoutingKey, // routing key
	)
	declareQueueArgs := amqp.Table{
		// Exchange where to send messages after TTL expiration.
		"x-dead-letter-exchange": b.GetConfig().AMQP.Exchange,
		// Routing key which use when resending expired messages.
		"x-dead-letter-routing-key": signature.RoutingKey,
		// Time in milliseconds
		// after that message will expire and be sent to destination.
		"x-message-ttl": delayMs,
		// Time after that the queue will be deleted.
		"x-expires": delayMs * 2,
	}
	conn, channel, _, _, _, err := b.Connect(
		b.GetConfig().Broker,
		b.GetConfig().TLSConfig,
		b.GetConfig().AMQP.Exchange,     // exchange name
		b.GetConfig().AMQP.ExchangeType, // exchange type
		queueName,                       // queue name
		true,                            // queue durable
		b.GetConfig().AMQP.AutoDelete,   // queue delete when unused
		queueName,                       // queue binding key
		nil,                             // exchange declare args
		declareQueueArgs,                // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		return err
	}

	defer b.Close(channel, conn)

	if err := channel.Publish(
		b.GetConfig().AMQP.Exchange, // exchange
		queueName,                   // routing key
		false,                       // mandatory
		false,                       // immediate
		amqp.Publishing{
			Headers:      amqp.Table(signature.Headers),
			ContentType:  "application/json",
			Body:         message,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return err
	}

	return nil
}

func (b *Broker) isDirectExchange() bool {
	return b.GetConfig().AMQP != nil && b.GetConfig().AMQP.ExchangeType == "direct"
}

// AdjustRoutingKey makes sure the routing key is correct.
// If the routing key is an empty string:
// a) set it to binding key for direct exchange type
// b) set it to default queue name
func (b *Broker) AdjustRoutingKey(s *tasks.Signature) {
	if s.RoutingKey != "" {
		return
	}

	if b.isDirectExchange() {
		// The routing algorithm behind a direct exchange is simple - a message goes
		// to the queues whose binding key exactly matches the routing key of the message.
		s.RoutingKey = b.GetConfig().AMQP.BindingKey
		return
	}

	s.RoutingKey = b.GetConfig().DefaultQueue
}
<commit_msg>Fix typo 'occured' to 'occurred'<commit_after>package amqp

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"github.com/RichardKnop/machinery/v1/brokers/errs"
	"github.com/RichardKnop/machinery/v1/brokers/iface"
	"github.com/RichardKnop/machinery/v1/common"
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"github.com/pkg/errors"
	"github.com/streadway/amqp"
)

type AMQPConnection struct {
	queueName    string
	connection   *amqp.Connection
	channel      *amqp.Channel
	queue        amqp.Queue
	confirmation <-chan amqp.Confirmation
	errorchan    <-chan *amqp.Error
	cleanup      chan struct{}
}

// Broker represents an AMQP broker
type Broker struct {
	common.Broker
	common.AMQPConnector
	processingWG sync.WaitGroup // use wait group to make sure task processing completes on interrupt signal

	connections      map[string]*AMQPConnection
	connectionsMutex sync.RWMutex
}

// New creates new Broker instance
func New(cnf *config.Config) iface.Broker {
	return &Broker{Broker: common.NewBroker(cnf), AMQPConnector: common.AMQPConnector{}, connections: make(map[string]*AMQPConnection)}
}

// StartConsuming enters a loop and waits for incoming messages
func (b *Broker) StartConsuming(consumerTag string, concurrency int, taskProcessor iface.TaskProcessor) (bool, error) {
	b.Broker.StartConsuming(consumerTag, concurrency, taskProcessor)

	queueName := taskProcessor.CustomQueue()
	if queueName == "" {
		queueName = b.GetConfig().DefaultQueue
	}

	conn, channel, queue, _, amqpCloseChan, err := b.Connect(
		b.GetConfig().Broker,
		b.GetConfig().TLSConfig,
		b.GetConfig().AMQP.Exchange,     // exchange name
		b.GetConfig().AMQP.ExchangeType, // exchange type
		queueName,                       // queue name
		true,                            // queue durable
		false,                           // queue delete when unused
		b.GetConfig().AMQP.BindingKey,   // queue binding key
		nil,                             // exchange declare args
		amqp.Table(b.GetConfig().AMQP.QueueDeclareArgs), // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		b.GetRetryFunc()(b.GetRetryStopChan())
		return b.GetRetry(), err
	}
	defer b.Close(channel, conn)

	if err = channel.Qos(
		b.GetConfig().AMQP.PrefetchCount,
		0,     // prefetch size
		false, // global
	); err != nil {
		return b.GetRetry(), fmt.Errorf("Channel qos error: %s", err)
	}

	deliveries, err := channel.Consume(
		queue.Name,  // queue
		consumerTag, // consumer tag
		false,       // auto-ack
		false,       // exclusive
		false,       // no-local
		false,       // no-wait
		nil,         // arguments
	)
	if err != nil {
		return b.GetRetry(), fmt.Errorf("Queue consume error: %s", err)
	}

	log.INFO.Print("[*] Waiting for messages. To exit press CTRL+C")

	if err := b.consume(deliveries, concurrency, taskProcessor, amqpCloseChan); err != nil {
		return b.GetRetry(), err
	}

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()

	return b.GetRetry(), nil
}

// StopConsuming quits the loop
func (b *Broker) StopConsuming() {
	b.Broker.StopConsuming()

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()
}

// GetOrOpenConnection will return a connection on a particular queue name. Open connections
// are saved to avoid having to reopen connection for multiple queues
func (b *Broker) GetOrOpenConnection(queueName string, queueBindingKey string, exchangeDeclareArgs, queueDeclareArgs, queueBindingArgs amqp.Table) (*AMQPConnection, error) {
	var err error

	b.connectionsMutex.Lock()
	defer b.connectionsMutex.Unlock()

	conn, ok := b.connections[queueName]
	if !ok {
		conn = &AMQPConnection{
			queueName: queueName,
			cleanup:   make(chan struct{}),
		}
		conn.connection, conn.channel, conn.queue, conn.confirmation, conn.errorchan, err = b.Connect(
			b.GetConfig().Broker,
			b.GetConfig().TLSConfig,
			b.GetConfig().AMQP.Exchange,     // exchange name
			b.GetConfig().AMQP.ExchangeType, // exchange type
			queueName,                       // queue name
			true,                            // queue durable
			false,                           // queue delete when unused
			queueBindingKey,                 // queue binding key
			exchangeDeclareArgs,             // exchange declare args
			queueDeclareArgs,                // queue declare args
			queueBindingArgs,                // queue binding args
		)
		if err != nil {
			return nil, errors.Wrapf(err, "Failed to connect to queue %s", queueName)
		}

		// Reconnect to the channel if it disconnects/errors out
		go func() {
			select {
			case err = <-conn.errorchan:
				log.INFO.Printf("Error occurred on queue: %s. Reconnecting", queueName)
				b.connectionsMutex.Lock()
				delete(b.connections, queueName)
				b.connectionsMutex.Unlock()
				_, err := b.GetOrOpenConnection(queueName, queueBindingKey, exchangeDeclareArgs, queueDeclareArgs, queueBindingArgs)
				if err != nil {
					log.ERROR.Printf("Failed to reopen queue: %s.", queueName)
				}
			case <-conn.cleanup:
				return
			}
			return
		}()
		b.connections[queueName] = conn
	}
	return conn, nil
}

func (b *Broker) CloseConnections() error {
	b.connectionsMutex.Lock()
	defer b.connectionsMutex.Unlock()

	for key, conn := range b.connections {
		if err := b.Close(conn.channel, conn.connection); err != nil {
			log.ERROR.Print("Failed to close channel")
			return nil
		}
		close(conn.cleanup)
		delete(b.connections, key)
	}
	return nil
}

// Publish places a new message on the default queue
func (b *Broker) Publish(ctx context.Context, signature *tasks.Signature) error {
	// Adjust routing key (this decides which queue the message will be published to)
	b.AdjustRoutingKey(signature)

	msg, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// Check the ETA signature field, if it is set and it is in the future,
	// delay the task
	if signature.ETA != nil {
		now := time.Now().UTC()

		if signature.ETA.After(now) {
			delayMs := int64(signature.ETA.Sub(now) / time.Millisecond)

			return b.delay(signature, delayMs)
		}
	}

	queue := b.GetConfig().DefaultQueue
	bindingKey := b.GetConfig().AMQP.BindingKey // queue binding key
	if b.isDirectExchange() {
		queue = signature.RoutingKey
		bindingKey = signature.RoutingKey
	}

	connection, err := b.GetOrOpenConnection(
		queue,
		bindingKey, // queue binding key
		nil,        // exchange declare args
		amqp.Table(b.GetConfig().AMQP.QueueDeclareArgs), // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		return errors.Wrapf(err, "Failed to get a connection for queue %s", queue)
	}

	channel := connection.channel
	confirmsChan := connection.confirmation

	if err := channel.Publish(
		b.GetConfig().AMQP.Exchange, // exchange name
		signature.RoutingKey,        // routing key
		false,                       // mandatory
		false,                       // immediate
		amqp.Publishing{
			Headers:      amqp.Table(signature.Headers),
			ContentType:  "application/json",
			Body:         msg,
			Priority:     signature.Priority,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return errors.Wrap(err, "Failed to publish task")
	}

	confirmed := <-confirmsChan

	if confirmed.Ack {
		return nil
	}

	return fmt.Errorf("Failed delivery of delivery tag: %v", confirmed.DeliveryTag)
}

// consume takes delivered messages from the channel and manages a worker pool
// to process tasks concurrently
func (b *Broker) consume(deliveries <-chan amqp.Delivery, concurrency int, taskProcessor iface.TaskProcessor, amqpCloseChan <-chan *amqp.Error) error {
	pool := make(chan struct{}, concurrency)

	// initialize worker pool with maxWorkers workers
	go func() {
		for i := 0; i < concurrency; i++ {
			pool <- struct{}{}
		}
	}()

	// make channel with a capacity makes it become a buffered channel so that a worker which wants to
	// push an error to `errorsChan` doesn't need to be blocked while the for-loop is blocked waiting
	// a worker, that is, it avoids a possible deadlock
	errorsChan := make(chan error, 1)

	for {
		select {
		case amqpErr := <-amqpCloseChan:
			return amqpErr
		case err := <-errorsChan:
			return err
		case d := <-deliveries:
			if concurrency > 0 {
				// get worker from pool (blocks until one is available)
				<-pool
			}

			b.processingWG.Add(1)

			// Consume the task inside a gotourine so multiple tasks
			// can be processed concurrently
			go func() {
				if err := b.consumeOne(d, taskProcessor); err != nil {
					errorsChan <- err
				}

				b.processingWG.Done()

				if concurrency > 0 {
					// give worker back to pool
					pool <- struct{}{}
				}
			}()
		case <-b.GetStopChan():
			return nil
		}
	}
}

// consumeOne processes a single message using TaskProcessor
func (b *Broker) consumeOne(delivery amqp.Delivery, taskProcessor iface.TaskProcessor) error {
	if len(delivery.Body) == 0 {
		delivery.Nack(true, false)                     // multiple, requeue
		return errors.New("Received an empty message") // RabbitMQ down?
	}

	var multiple, requeue = false, false

	// Unmarshal message body into signature struct
	signature := new(tasks.Signature)
	decoder := json.NewDecoder(bytes.NewReader(delivery.Body))
	decoder.UseNumber()
	if err := decoder.Decode(signature); err != nil {
		delivery.Nack(multiple, requeue)
		return errs.NewErrCouldNotUnmarshaTaskSignature(delivery.Body, err)
	}

	// If the task is not registered, we nack it and requeue,
	// there might be different workers for processing specific tasks
	if !b.IsTaskRegistered(signature.Name) {
		if !delivery.Redelivered {
			requeue = true
			log.INFO.Printf("Task not registered with this worker. Requeing message: %s", delivery.Body)
		}
		if !signature.IgnoreWhenTaskNotRegistered {
			delivery.Nack(multiple, requeue)
		}
		return nil
	}

	log.INFO.Printf("Received new message: %s", delivery.Body)

	err := taskProcessor.Process(signature)
	delivery.Ack(multiple)
	return err
}

// delay a task by delayDuration miliseconds, the way it works is a new queue
// is created without any consumers, the message is then published to this queue
// with appropriate ttl expiration headers, after the expiration, it is sent to
// the proper queue with consumers
func (b *Broker) delay(signature *tasks.Signature, delayMs int64) error {
	if delayMs <= 0 {
		return errors.New("Cannot delay task by 0ms")
	}

	message, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// It's necessary to redeclare the queue each time (to zero its TTL timer).
	queueName := fmt.Sprintf(
		"delay.%d.%s.%s",
		delayMs, // delay duration in mileseconds
		b.GetConfig().AMQP.Exchange,
		signature.RoutingKey, // routing key
	)
	declareQueueArgs := amqp.Table{
		// Exchange where to send messages after TTL expiration.
		"x-dead-letter-exchange": b.GetConfig().AMQP.Exchange,
		// Routing key which use when resending expired messages.
		"x-dead-letter-routing-key": signature.RoutingKey,
		// Time in milliseconds
		// after that message will expire and be sent to destination.
		"x-message-ttl": delayMs,
		// Time after that the queue will be deleted.
		"x-expires": delayMs * 2,
	}
	conn, channel, _, _, _, err := b.Connect(
		b.GetConfig().Broker,
		b.GetConfig().TLSConfig,
		b.GetConfig().AMQP.Exchange,     // exchange name
		b.GetConfig().AMQP.ExchangeType, // exchange type
		queueName,                       // queue name
		true,                            // queue durable
		b.GetConfig().AMQP.AutoDelete,   // queue delete when unused
		queueName,                       // queue binding key
		nil,                             // exchange declare args
		declareQueueArgs,                // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		return err
	}

	defer b.Close(channel, conn)

	if err := channel.Publish(
		b.GetConfig().AMQP.Exchange, // exchange
		queueName,                   // routing key
		false,                       // mandatory
		false,                       // immediate
		amqp.Publishing{
			Headers:      amqp.Table(signature.Headers),
			ContentType:  "application/json",
			Body:         message,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return err
	}

	return nil
}

func (b *Broker) isDirectExchange() bool {
	return b.GetConfig().AMQP != nil && b.GetConfig().AMQP.ExchangeType == "direct"
}

// AdjustRoutingKey makes sure the routing key is correct.
// If the routing key is an empty string:
// a) set it to binding key for direct exchange type
// b) set it to default queue name
func (b *Broker) AdjustRoutingKey(s *tasks.Signature) {
	if s.RoutingKey != "" {
		return
	}

	if b.isDirectExchange() {
		// The routing algorithm behind a direct exchange is simple - a message goes
		// to the queues whose binding key exactly matches the routing key of the message.
		s.RoutingKey = b.GetConfig().AMQP.BindingKey
		return
	}

	s.RoutingKey = b.GetConfig().DefaultQueue
}
<|endoftext|>


<commit_before>package messages

type AliasApprovedMessage struct {
	MessageType uint8
}

func NewAliasApprovedMessage() AliasApprovedMessage {
	return AliasApprovedMessage{
		MessageType: TypeAliasApproved,
	}
}

func NewAliasApprovedMessageFromBytes(messageBytes []byte) AliasApprovedMessage {
	return AliasApprovedMessage{
		MessageType: uint8(messageBytes[0]),
	}
}

func (ms *AliasApprovedMessage) Bytes() []byte {
	b := make([]byte, 1)
	b[0] = byte(ms.MessageType)
	return b
}
<commit_msg>Commented AliasApproved.go<commit_after>package messages

// AliasApprovedMessage is the struct which describes an Alias approved message
type AliasApprovedMessage struct {
	//MessageType is the message type
	MessageType uint8
}

// NewAliasApprovedMessage returns an instance of AliasApprovedMessage
func NewAliasApprovedMessage() AliasApprovedMessage {
	return AliasApprovedMessage{
		MessageType: TypeAliasApproved,
	}
}

// NewAliasApprovedMessageFromBytes returns an instance of AliasApprovedMessage
// Based on a given slice of bytes
func NewAliasApprovedMessageFromBytes(messageBytes []byte) AliasApprovedMessage {
	return AliasApprovedMessage{
		MessageType: uint8(messageBytes[0]),
	}
}

// Bytes returns a slice of bytes representing an AliasApprovedMessage
// which can be sent through a connection
func (ms *AliasApprovedMessage) Bytes() []byte {
	b := make([]byte, 1)
	b[0] = byte(ms.MessageType)
	return b
}
<|endoftext|>