<commit_before>/*
Copyright 2019 The Kubernetes 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 instancegroups

import (
	"context"
	"fmt"
	"sync"
	"time"

	"k8s.io/client-go/kubernetes"
	"k8s.io/klog"
	api "k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/cloudinstances"
	"k8s.io/kops/pkg/validation"
	"k8s.io/kops/upup/pkg/fi"
)

// RollingUpdateCluster is a struct containing cluster information for a rolling update.
type RollingUpdateCluster struct {
	Cloud fi.Cloud

	// MasterInterval is the amount of time to wait after stopping a master instance
	MasterInterval time.Duration
	// NodeInterval is the amount of time to wait after stopping a non-master instance
	NodeInterval time.Duration
	// BastionInterval is the amount of time to wait after stopping a bastion instance
	BastionInterval time.Duration
	// Interactive prompts user to continue after each instance is updated
	Interactive bool

	Force bool

	// K8sClient is the kubernetes client, used for draining etc
	K8sClient kubernetes.Interface

	// ClusterValidator is used for validating the cluster. Unused if CloudOnly
	ClusterValidator validation.ClusterValidator

	FailOnDrainError bool
	FailOnValidate   bool
	CloudOnly        bool
	ClusterName      string

	// PostDrainDelay is the duration we wait after draining each node
	PostDrainDelay time.Duration

	// ValidationTimeout is the maximum time to wait for the cluster to validate, once we start validation
	ValidationTimeout time.Duration

	// ValidateTickDuration is the amount of time to wait between cluster validation attempts
	ValidateTickDuration time.Duration

	// ValidateSuccessDuration is the amount of time a cluster must continue to validate successfully
	// before updating the next node
	ValidateSuccessDuration time.Duration

	// ValidateCount is the amount of time that a cluster needs to be validated after single node update
	ValidateCount int
}

// AdjustNeedUpdate adjusts the set of instances that need updating, using factors outside those known by the cloud implementation
func (c *RollingUpdateCluster) AdjustNeedUpdate(groups map[string]*cloudinstances.CloudInstanceGroup, cluster *api.Cluster, instanceGroups *api.InstanceGroupList) error {
	for _, group := range groups {
		if group.Ready != nil {
			var newReady []*cloudinstances.CloudInstanceGroupMember
			for _, member := range group.Ready {
				makeNotReady := false
				if member.Node != nil && member.Node.Annotations != nil {
					if _, ok := member.Node.Annotations["kops.k8s.io/needs-update"]; ok {
						makeNotReady = true
					}
				}

				if makeNotReady {
					group.NeedUpdate = append(group.NeedUpdate, member)
				} else {
					newReady = append(newReady, member)
				}
			}
			group.Ready = newReady
		}
	}
	return nil
}

// RollingUpdate performs a rolling update on a K8s Cluster.
func (c *RollingUpdateCluster) RollingUpdate(ctx context.Context, groups map[string]*cloudinstances.CloudInstanceGroup, cluster *api.Cluster, instanceGroups *api.InstanceGroupList) error {
	if len(groups) == 0 {
		klog.Info("Cloud Instance Group length is zero. Not doing a rolling-update.")
		return nil
	}

	var resultsMutex sync.Mutex
	results := make(map[string]error)

	masterGroups := make(map[string]*cloudinstances.CloudInstanceGroup)
	nodeGroups := make(map[string]*cloudinstances.CloudInstanceGroup)
	bastionGroups := make(map[string]*cloudinstances.CloudInstanceGroup)
	for k, group := range groups {
		switch group.InstanceGroup.Spec.Role {
		case api.InstanceGroupRoleNode:
			nodeGroups[k] = group
		case api.InstanceGroupRoleMaster:
			masterGroups[k] = group
		case api.InstanceGroupRoleBastion:
			bastionGroups[k] = group
		default:
			return fmt.Errorf("unknown group type for group %q", group.InstanceGroup.ObjectMeta.Name)
		}
	}

	// Upgrade bastions first; if these go down we can't see anything
	{
		var wg sync.WaitGroup

		for k, bastionGroup := range bastionGroups {
			wg.Add(1)
			go func(k string, group *cloudinstances.CloudInstanceGroup) {
				resultsMutex.Lock()
				results[k] = fmt.Errorf("function panic bastions")
				resultsMutex.Unlock()

				defer wg.Done()

				err := c.rollingUpdateInstanceGroup(ctx, cluster, group, true, c.BastionInterval)

				resultsMutex.Lock()
				results[k] = err
				resultsMutex.Unlock()
			}(k, bastionGroup)
		}

		wg.Wait()
	}

	// Do not continue update if bastion(s) failed
	for _, err := range results {
		if err != nil {
			return fmt.Errorf("bastion not healthy after update, stopping rolling-update: %q", err)
		}
	}

	// Upgrade masters next
	{
		// We run master nodes in series, even if they are in separate instance groups
		// typically they will be in separate instance groups, so we can force the zones,
		// and we don't want to roll all the masters at the same time.  See issue #284

		for _, group := range masterGroups {
			err := c.rollingUpdateInstanceGroup(ctx, cluster, group, false, c.MasterInterval)

			// Do not continue update if master(s) failed, cluster is potentially in an unhealthy state
			if err != nil {
				return fmt.Errorf("master not healthy after update, stopping rolling-update: %q", err)
			}
		}
	}

	// Upgrade nodes
	{
		// We run nodes in series, even if they are in separate instance groups
		// typically they will not being separate instance groups. If you roll the nodes in parallel
		// you can get into a scenario where you can evict multiple statefulset pods from the same
		// statefulset at the same time. Further improvements needs to be made to protect from this as
		// well.

		for k := range nodeGroups {
			results[k] = fmt.Errorf("function panic nodes")
		}

		for k, group := range nodeGroups {
			err := c.rollingUpdateInstanceGroup(ctx, cluster, group, false, c.NodeInterval)

			results[k] = err

			// TODO: Bail on error?
		}
	}

	for _, err := range results {
		if err != nil {
			return err
		}
	}

	klog.Infof("Rolling update completed for cluster %q!", c.ClusterName)
	return nil
}
<commit_msg>Rolling update instance groups in consistent order<commit_after>/*
Copyright 2019 The Kubernetes 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 instancegroups

import (
	"context"
	"fmt"
	"sort"
	"sync"
	"time"

	"k8s.io/client-go/kubernetes"
	"k8s.io/klog"
	api "k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/cloudinstances"
	"k8s.io/kops/pkg/validation"
	"k8s.io/kops/upup/pkg/fi"
)

// RollingUpdateCluster is a struct containing cluster information for a rolling update.
type RollingUpdateCluster struct {
	Cloud fi.Cloud

	// MasterInterval is the amount of time to wait after stopping a master instance
	MasterInterval time.Duration
	// NodeInterval is the amount of time to wait after stopping a non-master instance
	NodeInterval time.Duration
	// BastionInterval is the amount of time to wait after stopping a bastion instance
	BastionInterval time.Duration
	// Interactive prompts user to continue after each instance is updated
	Interactive bool

	Force bool

	// K8sClient is the kubernetes client, used for draining etc
	K8sClient kubernetes.Interface

	// ClusterValidator is used for validating the cluster. Unused if CloudOnly
	ClusterValidator validation.ClusterValidator

	FailOnDrainError bool
	FailOnValidate   bool
	CloudOnly        bool
	ClusterName      string

	// PostDrainDelay is the duration we wait after draining each node
	PostDrainDelay time.Duration

	// ValidationTimeout is the maximum time to wait for the cluster to validate, once we start validation
	ValidationTimeout time.Duration

	// ValidateTickDuration is the amount of time to wait between cluster validation attempts
	ValidateTickDuration time.Duration

	// ValidateSuccessDuration is the amount of time a cluster must continue to validate successfully
	// before updating the next node
	ValidateSuccessDuration time.Duration

	// ValidateCount is the amount of time that a cluster needs to be validated after single node update
	ValidateCount int
}

// AdjustNeedUpdate adjusts the set of instances that need updating, using factors outside those known by the cloud implementation
func (c *RollingUpdateCluster) AdjustNeedUpdate(groups map[string]*cloudinstances.CloudInstanceGroup, cluster *api.Cluster, instanceGroups *api.InstanceGroupList) error {
	for _, group := range groups {
		if group.Ready != nil {
			var newReady []*cloudinstances.CloudInstanceGroupMember
			for _, member := range group.Ready {
				makeNotReady := false
				if member.Node != nil && member.Node.Annotations != nil {
					if _, ok := member.Node.Annotations["kops.k8s.io/needs-update"]; ok {
						makeNotReady = true
					}
				}

				if makeNotReady {
					group.NeedUpdate = append(group.NeedUpdate, member)
				} else {
					newReady = append(newReady, member)
				}
			}
			group.Ready = newReady
		}
	}
	return nil
}

// RollingUpdate performs a rolling update on a K8s Cluster.
func (c *RollingUpdateCluster) RollingUpdate(ctx context.Context, groups map[string]*cloudinstances.CloudInstanceGroup, cluster *api.Cluster, instanceGroups *api.InstanceGroupList) error {
	if len(groups) == 0 {
		klog.Info("Cloud Instance Group length is zero. Not doing a rolling-update.")
		return nil
	}

	var resultsMutex sync.Mutex
	results := make(map[string]error)

	masterGroups := make(map[string]*cloudinstances.CloudInstanceGroup)
	nodeGroups := make(map[string]*cloudinstances.CloudInstanceGroup)
	bastionGroups := make(map[string]*cloudinstances.CloudInstanceGroup)
	for k, group := range groups {
		switch group.InstanceGroup.Spec.Role {
		case api.InstanceGroupRoleNode:
			nodeGroups[k] = group
		case api.InstanceGroupRoleMaster:
			masterGroups[k] = group
		case api.InstanceGroupRoleBastion:
			bastionGroups[k] = group
		default:
			return fmt.Errorf("unknown group type for group %q", group.InstanceGroup.ObjectMeta.Name)
		}
	}

	// Upgrade bastions first; if these go down we can't see anything
	{
		var wg sync.WaitGroup

		for _, k := range sortGroups(bastionGroups) {
			wg.Add(1)
			go func(k string) {
				resultsMutex.Lock()
				results[k] = fmt.Errorf("function panic bastions")
				resultsMutex.Unlock()

				defer wg.Done()

				err := c.rollingUpdateInstanceGroup(ctx, cluster, bastionGroups[k], true, c.BastionInterval)

				resultsMutex.Lock()
				results[k] = err
				resultsMutex.Unlock()
			}(k)
		}

		wg.Wait()
	}

	// Do not continue update if bastion(s) failed
	for _, err := range results {
		if err != nil {
			return fmt.Errorf("bastion not healthy after update, stopping rolling-update: %q", err)
		}
	}

	// Upgrade masters next
	{
		// We run master nodes in series, even if they are in separate instance groups
		// typically they will be in separate instance groups, so we can force the zones,
		// and we don't want to roll all the masters at the same time.  See issue #284

		for _, k := range sortGroups(masterGroups) {
			err := c.rollingUpdateInstanceGroup(ctx, cluster, masterGroups[k], false, c.MasterInterval)

			// Do not continue update if master(s) failed, cluster is potentially in an unhealthy state
			if err != nil {
				return fmt.Errorf("master not healthy after update, stopping rolling-update: %q", err)
			}
		}
	}

	// Upgrade nodes
	{
		// We run nodes in series, even if they are in separate instance groups
		// typically they will not being separate instance groups. If you roll the nodes in parallel
		// you can get into a scenario where you can evict multiple statefulset pods from the same
		// statefulset at the same time. Further improvements needs to be made to protect from this as
		// well.

		for k := range nodeGroups {
			results[k] = fmt.Errorf("function panic nodes")
		}

		for _, k := range sortGroups(nodeGroups) {
			err := c.rollingUpdateInstanceGroup(ctx, cluster, nodeGroups[k], false, c.NodeInterval)

			results[k] = err

			// TODO: Bail on error?
		}
	}

	for _, err := range results {
		if err != nil {
			return err
		}
	}

	klog.Infof("Rolling update completed for cluster %q!", c.ClusterName)
	return nil
}

func sortGroups(groupMap map[string]*cloudinstances.CloudInstanceGroup) []string {
	groups := make([]string, 0, len(groupMap))
	for group := range groupMap {
		groups = append(groups, group)
	}
	sort.Strings(groups)
	return groups
}
<|endoftext|>


<commit_before>package xpi

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

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

// every minute, add an rsa key to the cache. This will block if
// the cache channel is already full, which is what we want anyway
func (s *XPISigner) populateRsaCache(size int) {
	for {
		key, err := rsa.GenerateKey(rand.Reader, size)
		if err != nil {
			log.Fatalf("xpi.populateRsaCache: %v", err)
		}
		s.rsaCache <- key
		time.Sleep(time.Minute)
	}
}

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

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

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

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

	template := s.makeTemplate(cn)

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

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

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

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

// every minute, add an rsa key to the cache. This will block if
// the cache channel is already full, which is what we want anyway
func (s *XPISigner) populateRsaCache(size int) {
	for {
		key, err := rsa.GenerateKey(rand.Reader, size)
		if err != nil {
			log.Fatalf("xpi.populateRsaCache: %v", err)
		}
		s.rsaCache <- key
		time.Sleep(time.Minute)
	}
}

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

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

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

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

	template := s.makeTemplate(cn)

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

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


<commit_before>package models

import (
	"bytes"
	"database/sql"
	"fmt"
	"strconv"
	"strings"
)

// SearchOptions are options that can be passed to SearchSets for filtering
// sets.
type SearchOptions struct {
	// If len is 0, then it should be treated as if all statuses are good.
	Status []int
	Query  string
	// Gamemodes to which limit the results. If len is 0, it means all modes
	// are ok.
	Mode []int

	// Pagination options.
	Offset int
	Amount int
}

func (o SearchOptions) setModes() (total uint8) {
	for _, m := range o.Mode {
		if m < 0 || m >= 4 {
			continue
		}
		total |= 1 << uint8(m)
	}
	return
}

var mysqlStringReplacer = strings.NewReplacer(
	`\`, `\\`,
	`"`, `\"`,
	`'`, `\'`,
	"\x00", `\0`,
	"\n", `\n`,
	"\r", `\r`,
	"\x1a", `\Z`,
)

func sIntCommaSeparated(nums []int) string {
	b := bytes.Buffer{}
	for idx, num := range nums {
		b.WriteString(strconv.Itoa(num))
		if idx != len(nums)-1 {
			b.WriteString(", ")
		}
	}
	return b.String()
}

// SearchSets retrieves sets, filtering them using SearchOptions.
func SearchSets(db, searchDB *sql.DB, opts SearchOptions) ([]Set, error) {
	setIDsQuery := "SELECT id FROM cg WHERE "

	// add filters to query
	// Yes. I know. Prepared statements. But Sphinx doesn't like them, so
	// bummer.
	setIDsQuery += "MATCH('" + mysqlStringReplacer.Replace(opts.Query) + "') "
	if len(opts.Status) != 0 {
		setIDsQuery += "AND ranked_status IN (" + sIntCommaSeparated(opts.Status) + ") "
	}
	if len(opts.Mode) != 0 {
		sm := strconv.Itoa(int(opts.setModes()))
		setIDsQuery += "AND set_modes = " + sm + " "
	}

	// set limit
	setIDsQuery += fmt.Sprintf("ORDER BY WEIGHT() DESC, id DESC LIMIT %d, %d OPTION ranker=sph04", opts.Offset, opts.Amount)

	// fetch rows
	rows, err := searchDB.Query(setIDsQuery)
	if err != nil {
		return nil, err
	}

	// from the rows we will retrieve the IDs of all our sets.
	// we also pre-create the slices containing the sets we will fill later on
	// when we fetch the actual data.
	setIDs := make([]int, 0, opts.Amount)
	sets := make([]Set, 0, opts.Amount)
	// setMap, having an ID, points to a position of a set contained in sets.
	setMap := make(map[int]int, opts.Amount)
	for rows.Next() {
		var id int
		err = rows.Scan(&id)
		if err != nil {
			return nil, err
		}
		setIDs = append(setIDs, id)
		sets = append(sets, Set{})
		setMap[id] = len(sets) - 1
	}

	// short circuit: there are no sets
	if len(sets) == 0 {
		return []Set{}, nil
	}

	setsQuery := "SELECT " + setFields + " FROM sets WHERE id IN (" + inClause(len(setIDs)) + ")"
	args := sIntToSInterface(setIDs)

	rows, err = db.Query(setsQuery, args...)

	if err != nil {
		return nil, err
	}

	// find all beatmaps, but leave children aside for the moment.
	for rows.Next() {
		var s Set
		err = rows.Scan(
			&s.ID, &s.RankedStatus, &s.ApprovedDate, &s.LastUpdate, &s.LastChecked,
			&s.Artist, &s.Title, &s.Creator, &s.Source, &s.Tags, &s.HasVideo, &s.Genre,
			&s.Language, &s.Favourites,
		)
		if err != nil {
			return nil, err
		}
		sets[setMap[s.ID]] = s
	}

	rows, err = db.Query(
		"SELECT "+beatmapFields+" FROM beatmaps WHERE parent_set_id IN ("+
			inClause(len(setIDs))+")",
		sIntToSInterface(setIDs)...,
	)
	if err != nil {
		return nil, err
	}

	for rows.Next() {
		var b Beatmap
		err = rows.Scan(
			&b.ID, &b.ParentSetID, &b.DiffName, &b.FileMD5, &b.Mode, &b.BPM,
			&b.AR, &b.OD, &b.CS, &b.HP, &b.TotalLength, &b.HitLength,
			&b.Playcount, &b.Passcount, &b.MaxCombo, &b.DifficultyRating,
		)
		if err != nil {
			return nil, err
		}
		parentSet, ok := setMap[b.ParentSetID]
		if !ok {
			continue
		}
		sets[parentSet].ChildrenBeatmaps = append(sets[parentSet].ChildrenBeatmaps, b)
	}

	return sets, nil
}
<commit_msg>Fix modes not working<commit_after>package models

import (
	"bytes"
	"database/sql"
	"fmt"
	"strconv"
	"strings"
)

// SearchOptions are options that can be passed to SearchSets for filtering
// sets.
type SearchOptions struct {
	// If len is 0, then it should be treated as if all statuses are good.
	Status []int
	Query  string
	// Gamemodes to which limit the results. If len is 0, it means all modes
	// are ok.
	Mode []int

	// Pagination options.
	Offset int
	Amount int
}

func (o SearchOptions) setModes() (total uint8) {
	for _, m := range o.Mode {
		if m < 0 || m >= 4 {
			continue
		}
		total |= 1 << uint8(m)
	}
	return
}

var mysqlStringReplacer = strings.NewReplacer(
	`\`, `\\`,
	`"`, `\"`,
	`'`, `\'`,
	"\x00", `\0`,
	"\n", `\n`,
	"\r", `\r`,
	"\x1a", `\Z`,
)

func sIntCommaSeparated(nums []int) string {
	b := bytes.Buffer{}
	for idx, num := range nums {
		b.WriteString(strconv.Itoa(num))
		if idx != len(nums)-1 {
			b.WriteString(", ")
		}
	}
	return b.String()
}

// SearchSets retrieves sets, filtering them using SearchOptions.
func SearchSets(db, searchDB *sql.DB, opts SearchOptions) ([]Set, error) {
	sm := strconv.Itoa(int(opts.setModes()))
	setIDsQuery := "SELECT id, set_modes & " + sm + " AS valid_set_modes FROM cg WHERE "

	// add filters to query
	// Yes. I know. Prepared statements. But Sphinx doesn't like them, so
	// bummer.
	setIDsQuery += "MATCH('" + mysqlStringReplacer.Replace(opts.Query) + "') "
	if len(opts.Status) != 0 {
		setIDsQuery += "AND ranked_status IN (" + sIntCommaSeparated(opts.Status) + ") "
	}
	if len(opts.Mode) != 0 {
		// This is a hack. Apparently, Sphinx does not support AND bitwise
		// operations in the WHERE clause, so we're placing that in the SELECT
		// clause and only making sure it's correct in this place.
		setIDsQuery += "AND valid_set_modes = " + sm + " "
	}

	// set limit
	setIDsQuery += fmt.Sprintf("ORDER BY WEIGHT() DESC, id DESC LIMIT %d, %d OPTION ranker=sph04", opts.Offset, opts.Amount)

	// fetch rows
	rows, err := searchDB.Query(setIDsQuery)
	if err != nil {
		return nil, err
	}

	// from the rows we will retrieve the IDs of all our sets.
	// we also pre-create the slices containing the sets we will fill later on
	// when we fetch the actual data.
	setIDs := make([]int, 0, opts.Amount)
	sets := make([]Set, 0, opts.Amount)
	// setMap, having an ID, points to a position of a set contained in sets.
	setMap := make(map[int]int, opts.Amount)
	for rows.Next() {
		var id int
		err = rows.Scan(&id, new(int))
		if err != nil {
			return nil, err
		}
		setIDs = append(setIDs, id)
		sets = append(sets, Set{})
		setMap[id] = len(sets) - 1
	}

	// short circuit: there are no sets
	if len(sets) == 0 {
		return []Set{}, nil
	}

	setsQuery := "SELECT " + setFields + " FROM sets WHERE id IN (" + inClause(len(setIDs)) + ")"
	args := sIntToSInterface(setIDs)

	rows, err = db.Query(setsQuery, args...)

	if err != nil {
		return nil, err
	}

	// find all beatmaps, but leave children aside for the moment.
	for rows.Next() {
		var s Set
		err = rows.Scan(
			&s.ID, &s.RankedStatus, &s.ApprovedDate, &s.LastUpdate, &s.LastChecked,
			&s.Artist, &s.Title, &s.Creator, &s.Source, &s.Tags, &s.HasVideo, &s.Genre,
			&s.Language, &s.Favourites,
		)
		if err != nil {
			return nil, err
		}
		sets[setMap[s.ID]] = s
	}

	rows, err = db.Query(
		"SELECT "+beatmapFields+" FROM beatmaps WHERE parent_set_id IN ("+
			inClause(len(setIDs))+")",
		sIntToSInterface(setIDs)...,
	)
	if err != nil {
		return nil, err
	}

	for rows.Next() {
		var b Beatmap
		err = rows.Scan(
			&b.ID, &b.ParentSetID, &b.DiffName, &b.FileMD5, &b.Mode, &b.BPM,
			&b.AR, &b.OD, &b.CS, &b.HP, &b.TotalLength, &b.HitLength,
			&b.Playcount, &b.Passcount, &b.MaxCombo, &b.DifficultyRating,
		)
		if err != nil {
			return nil, err
		}
		parentSet, ok := setMap[b.ParentSetID]
		if !ok {
			continue
		}
		sets[parentSet].ChildrenBeatmaps = append(sets[parentSet].ChildrenBeatmaps, b)
	}

	return sets, nil
}
<|endoftext|>


<commit_before>/*
   - support for RelayState
   - does not seem to work: incl. Capitalization content-security-policy: referrer no-referrer;
   - redo no-referer - current version does not work !!!
   - MDQ lookup by location also for hub md
   - Trusted proxy
   - wayf:wayf i hub_ops metadata
        - AttributeNameFormat for Krib -> WAYF SPS - ie if none -> repeat wayf error both formats but error
        - schacHomeOrganization
        - schacHomeOrganizationType
   - collect schema errors
   - illegal attributes from IdP - ignore or provoke error
*/

package gohybrid

import (
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"github.com/gorilla/securecookie"
	"github.com/wayf-dk/go-libxml2/types"
	"github.com/wayf-dk/gosaml"
	"github.com/wayf-dk/goxml"
	"html/template"
	"log"
	"net/http"
	"net/url"
	"regexp"
	"sync"
)

type (
	formdata struct {
		Acs          string
		Samlresponse string
	}

	idpsppair struct {
		idp string
		sp  string
	}

	Conf struct {
		DiscoveryService        string
		Domain                  string
		HubEntityID             string
		EptidSalt               string
		HubRequestedAttributes  *goxml.Xp
		Internal, External, Hub gosaml.Md
		SecureCookieHashKey     string
		PostFormTemplate        string
		Basic2uri               map[string]string
		StdTiming               gosaml.IdAndTiming
		ElementsToSign          []string
		AttributeHandler        func(*goxml.Xp, *goxml.Xp, *goxml.Xp, *goxml.Xp) error
	}
)

const (
	idpCertQuery = `./md:IDPSSODescriptor/md:KeyDescriptor[@use="signing" or not(@use)]/ds:KeyInfo/ds:X509Data/ds:X509Certificate`
)

var (
	_     = log.Printf // For debugging; delete when done.
	_     = fmt.Printf
	remap = map[string]idpsppair{
		"https://nemlogin.wayf.dk": idpsppair{"https://saml.nemlog-in.dk", "https://nemlogin.wayf.dk"},
	}

	contextmutex sync.RWMutex
	context      = make(map[*http.Request]map[string]string)
	bify         = regexp.MustCompile("^(https?://)(.*)$")
	debify       = regexp.MustCompile("^(https?://)(?:(?:birk|krib)\\.wayf.dk/(?:birk|krib)\\.php/)(.+)$")

	postform  *template.Template
	hashKey   []byte
	seccookie *securecookie.SecureCookie
	config    = Conf{}
)

func Config(configuration Conf) {
	config = configuration
	hashKey, _ := hex.DecodeString(config.SecureCookieHashKey)
	seccookie = securecookie.New(hashKey, nil)
	postform = template.Must(template.New("post").Parse(config.PostFormTemplate))
}

func SsoService(w http.ResponseWriter, r *http.Request) (err error) {
	defer r.Body.Close()
	// handle non ok urls gracefully
	// var err error
	// check issuer and acs in md
	// receiveRequest -> request, issuer md, receiver md
	//     check for IDPList 1st in md, then in request then in query
	//     sanitize idp from query or request
	request, spmd, _, err := gosaml.ReceiveSAMLRequest(r, config.Internal, config.Hub)
	if err != nil {
		return
	}
	entityID := spmd.Query1(nil, "@entityID")
	idp := spmd.Query1(nil, "//IDPList/ProviderID") // Need to find a place for IDPList
	if idp == "" {
		idp = request.Query1(nil, "IDPList/ProviderID")
	}
	if idp == "" {
		idp = r.URL.Query().Get("idpentityid")
	}
	if idp == "" {
		data := url.Values{}
		data.Set("return", "https://"+r.Host+r.RequestURI)
		data.Set("returnIDParam", "idpentityid")
		data.Set("entityID", entityID)
		http.Redirect(w, r, config.DiscoveryService+data.Encode(), http.StatusFound)
	} else {
		var idpmd *goxml.Xp
		/**/
		// check overlap btw ad-hoc feds for the idp and the sp
		kribID := bify.ReplaceAllString(entityID, "${1}krib.wayf.dk/krib.php/$2")
		if kribID == entityID {
			kribID = "urn:oid:1.3.6.1.4.1.39153:42:" + entityID
		}

		request.QueryDashP(nil, "/saml:Issuer", kribID, nil)
		acs := request.Query1(nil, "@AssertionConsumerServiceURL")
		acsurl := bify.ReplaceAllString(acs, "${1}krib.wayf.dk/krib.php/$2")
		request.QueryDashP(nil, "@AssertionConsumerServiceURL", acsurl, nil)
		/**/
		idpmd, err = config.External.MDQ(idp)
		if err != nil {
			return
		}
		const ssoquery = "./md:IDPSSODescriptor/md:SingleSignOnService[@Binding='urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect']/@Location"
		ssoservice := idpmd.Query1(nil, ssoquery)
		if ssoservice == "" {

		}
		request.QueryDashP(nil, "@Destination", ssoservice, nil)
		u, _ := gosaml.SAMLRequest2Url(request, "", "", "")
		log.Println(request.Doc.Dump(true))
		http.Redirect(w, r, u.String(), http.StatusFound)
	}
	return
}

func BirkService(w http.ResponseWriter, r *http.Request) (err error) {
	// use incoming request for crafting the new one
	// remember to add the Scoping element to inform the IdP of requesterID - if stated in metadata for the IdP
	// check ad-hoc feds overlab
	defer r.Body.Close()
	// get the sp as well to check for allowed acs
	request, _, mdbirkidp, err := gosaml.ReceiveSAMLRequest(r, config.External, config.External)
	if err != nil {
		return
	}
	// Save the issuer and destination in a cookie for when the response comes back

	cookievalue, err := seccookie.Encode("BIRK", gosaml.Deflate(request.Doc.Dump(true)))
	http.SetCookie(w, &http.Cookie{Name: "BIRK", Value: cookievalue, Domain: config.Domain, Path: "/", Secure: true, HttpOnly: true})

	idp := debify.ReplaceAllString(mdbirkidp.Query1(nil, "@entityID"), "$1$2")

	var mdhub, mdidp *goxml.Xp
	// are we remapping - for now only use case is https://nemlogin.wayf.dk -> https://saml.nemlog-in.dk
	if rm, ok := remap[idp]; ok {
		mdidp, err = config.Internal.MDQ(rm.idp)
		mdhub, err = config.Internal.MDQ(rm.sp)
	} else {
		mdidp, err = config.Internal.MDQ(idp)
		mdhub, err = config.Hub.MDQ(config.HubEntityID)
	}
	// use a std request - we take care of NameID etc in acsService below
	newrequest := gosaml.NewAuthnRequest(config.StdTiming.Refresh(), mdhub, mdidp)
	u, _ := gosaml.SAMLRequest2Url(newrequest, "", "", "") // not signed so blank key, pw and algo
	http.Redirect(w, r, u.String(), http.StatusFound)
	return
}

func AcsService(w http.ResponseWriter, r *http.Request) (err error) {
	defer r.Body.Close()
	birk, err := r.Cookie("BIRK")
	if err != nil {
		return err
	}

	value := []byte{}
	if err = seccookie.Decode("BIRK", birk.Value, &value); err != nil {
		return
	}

	// we checked the request when we received in birkService - we can use it without fear ie. we just parse it
	log.Println("cookie", string(gosaml.Inflate(value)))
	request := goxml.NewXp(string(gosaml.Inflate(value)))

	http.SetCookie(w, &http.Cookie{Name: "BIRK", Value: "", Domain: config.Domain, Path: "/", Secure: true, HttpOnly: true, MaxAge: -1})
	sp_md, err := config.External.MDQ(request.Query1(nil, "/samlp:AuthnRequest/saml:Issuer"))
	if err != nil {
		return
	}

	response, idp_md, _, err := gosaml.ReceiveSAMLResponse(r, config.Internal, config.Hub)
	if err != nil {
		return
	}

	err = config.AttributeHandler(idp_md, config.HubRequestedAttributes, sp_md, response)
	if err != nil {
		return
	}

	birkmd, err := config.External.MDQ(request.Query1(nil, "/samlp:AuthnRequest/@Destination"))
	if err != nil {
		return
	}
	nameid := response.Query(nil, "./saml:Assertion/saml:Subject/saml:NameID")[0]
	// respect nameID in req, give persistent id + all computed attributes + nameformat conversion
	nameidformat := sp_md.Query1(nil, "./md:SPSSODescriptor/md:NameIDFormat")
	if nameidformat == gosaml.Persistent {
		response.QueryDashP(nameid, "@Format", gosaml.Persistent, nil)
		eptid := response.Query1(nil, `./saml:Assertion/saml:AttributeStatement/saml:Attribute[@FriendlyName="eduPersonTargetedID"]/saml:AttributeValue`)
		response.QueryDashP(nameid, ".", eptid, nil)
	} else if nameidformat == gosaml.Transient {
		response.QueryDashP(nameid, ".", gosaml.Id(), nil)
	}

	newresponse := gosaml.NewResponse(config.StdTiming.Refresh(), birkmd, sp_md, request, response)

	for _, q := range config.ElementsToSign {
		err = gosaml.SignResponse(newresponse, q, birkmd)
		if err != nil {
			return
		}
	}

	// when consent as a service is ready - we will post to that
	acs := newresponse.Query1(nil, "@Destination")

	data := formdata{Acs: acs, Samlresponse: base64.StdEncoding.EncodeToString([]byte(newresponse.Doc.Dump(false)))}
	postform.Execute(w, data)
	return
}

func KribService(w http.ResponseWriter, r *http.Request) (err error) {
	// check ad-hoc feds overlap
	defer r.Body.Close()

	response, _, _, err := gosaml.ReceiveSAMLResponse(r, config.External, config.External)
	if err != nil {
		return
	}
	destination := debify.ReplaceAllString(response.Query1(nil, "@Destination"), "$1$2")
	response.QueryDashP(nil, "@Destination", destination, nil)
	response.QueryDashP(nil, "./saml:Assertion/saml:Subject/saml:SubjectConfirmation/saml:SubjectConfirmationData/@Recipient", destination, nil)
	issuer := config.HubEntityID
	response.QueryDashP(nil, "./saml:Issuer", issuer, nil)
	response.QueryDashP(nil, "./saml:Assertion/saml:Issuer", issuer, nil)
	// Krib always receives attributes with nameformat=urn. Before sending to the real SP we need to look into
	// the metadata for SP to determine the actual nameformat - as WAYF supports both for internal SPs.
	mdsp, err := config.Internal.MDQ(destination)
	if err != nil {
		return
	}
	requestedattributes := mdsp.Query(nil, "./md:SPSSODescriptor/md:AttributeConsumingService/md:RequestedAttribute")
	attributestatement := response.Query(nil, "./saml:Assertion/saml:AttributeStatement")[0]
	for _, attr := range requestedattributes {
		nameFormat, _ := attr.(types.Element).GetAttribute("NameFormat")
		if nameFormat.NodeValue() == gosaml.Basic {
			basicname, _ := attr.(types.Element).GetAttribute("Name")
			uriname := config.Basic2uri[basicname.NodeValue()]
			responseattribute := response.Query(attributestatement, "saml:Attribute[@Name='"+uriname+"']")
			if len(responseattribute) > 0 {
				responseattribute[0].(types.Element).SetAttribute("Name", basicname.NodeValue())
				responseattribute[0].(types.Element).SetAttribute("NameFormat", gosaml.Basic)
			}
		}
	}

	mdhub, err := config.Hub.MDQ(config.HubEntityID)
	if err != nil {
		return
	}

	for _, q := range config.ElementsToSign {
		err = gosaml.SignResponse(response, q, mdhub)
		if err != nil {
			return
		}
	}

	data := formdata{Acs: destination, Samlresponse: base64.StdEncoding.EncodeToString([]byte(response.Doc.Dump(false)))}
	postform.Execute(w, data)
	return
}
<commit_msg>Added support for RelqyState<commit_after>/*
   - support for RelayState
   - does not seem to work: incl. Capitalization content-security-policy: referrer no-referrer;
   - redo no-referer - current version does not work !!!
   - MDQ lookup by location also for hub md
   - Trusted proxy
   - wayf:wayf i hub_ops metadata
        - AttributeNameFormat for Krib -> WAYF SPS - ie if none -> repeat wayf error both formats but error
        - schacHomeOrganization
        - schacHomeOrganizationType
   - collect schema errors
   - illegal attributes from IdP - ignore or provoke error
*/

package gohybrid

import (
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"github.com/gorilla/securecookie"
	"github.com/wayf-dk/go-libxml2/types"
	"github.com/wayf-dk/gosaml"
	"github.com/wayf-dk/goxml"
	"html/template"
	"log"
	"net/http"
	"net/url"
	"regexp"
	"sync"
)

type (
	formdata struct {
		Acs          string
		Samlresponse string
		RelayState   string
	}

	idpsppair struct {
		idp string
		sp  string
	}

	Conf struct {
		DiscoveryService        string
		Domain                  string
		HubEntityID             string
		EptidSalt               string
		HubRequestedAttributes  *goxml.Xp
		Internal, External, Hub gosaml.Md
		SecureCookieHashKey     string
		PostFormTemplate        string
		Basic2uri               map[string]string
		StdTiming               gosaml.IdAndTiming
		ElementsToSign          []string
		AttributeHandler        func(*goxml.Xp, *goxml.Xp, *goxml.Xp, *goxml.Xp) error
	}
)

const (
	idpCertQuery = `./md:IDPSSODescriptor/md:KeyDescriptor[@use="signing" or not(@use)]/ds:KeyInfo/ds:X509Data/ds:X509Certificate`
)

var (
	_     = log.Printf // For debugging; delete when done.
	_     = fmt.Printf
	remap = map[string]idpsppair{
		"https://nemlogin.wayf.dk": idpsppair{"https://saml.nemlog-in.dk", "https://nemlogin.wayf.dk"},
	}

	contextmutex sync.RWMutex
	context      = make(map[*http.Request]map[string]string)
	bify         = regexp.MustCompile("^(https?://)(.*)$")
	debify       = regexp.MustCompile("^(https?://)(?:(?:birk|krib)\\.wayf.dk/(?:birk|krib)\\.php/)(.+)$")

	postform  *template.Template
	hashKey   []byte
	seccookie *securecookie.SecureCookie
	config    = Conf{}
)

func Config(configuration Conf) {
	config = configuration
	hashKey, _ := hex.DecodeString(config.SecureCookieHashKey)
	seccookie = securecookie.New(hashKey, nil)
	postform = template.Must(template.New("post").Parse(config.PostFormTemplate))
}

func SsoService(w http.ResponseWriter, r *http.Request) (err error) {
	defer r.Body.Close()
	// handle non ok urls gracefully
	// var err error
	// check issuer and acs in md
	// receiveRequest -> request, issuer md, receiver md
	//     check for IDPList 1st in md, then in request then in query
	//     sanitize idp from query or request
	request, spmd, _, relayState, err := gosaml.ReceiveSAMLRequest(r, config.Internal, config.Hub)
	if err != nil {
		return
	}
	entityID := spmd.Query1(nil, "@entityID")
	idp := spmd.Query1(nil, "//IDPList/ProviderID") // Need to find a place for IDPList
	if idp == "" {
		idp = request.Query1(nil, "IDPList/ProviderID")
	}
	if idp == "" {
		idp = r.URL.Query().Get("idpentityid")
	}
	if idp == "" {
		data := url.Values{}
		data.Set("return", "https://"+r.Host+r.RequestURI)
		data.Set("returnIDParam", "idpentityid")
		data.Set("entityID", entityID)
		http.Redirect(w, r, config.DiscoveryService+data.Encode(), http.StatusFound)
	} else {
		var idpmd *goxml.Xp
		/**/
		// check overlap btw ad-hoc feds for the idp and the sp
		kribID := bify.ReplaceAllString(entityID, "${1}krib.wayf.dk/krib.php/$2")
		if kribID == entityID {
			kribID = "urn:oid:1.3.6.1.4.1.39153:42:" + entityID
		}

		request.QueryDashP(nil, "/saml:Issuer", kribID, nil)
		acs := request.Query1(nil, "@AssertionConsumerServiceURL")
		acsurl := bify.ReplaceAllString(acs, "${1}krib.wayf.dk/krib.php/$2")
		request.QueryDashP(nil, "@AssertionConsumerServiceURL", acsurl, nil)
		/**/
		idpmd, err = config.External.MDQ(idp)
		if err != nil {
			return
		}
		const ssoquery = "./md:IDPSSODescriptor/md:SingleSignOnService[@Binding='urn:oasis:names:tc:SAML:2.0:bindings:HTTP-Redirect']/@Location"
		ssoservice := idpmd.Query1(nil, ssoquery)
		if ssoservice == "" {

		}
		request.QueryDashP(nil, "@Destination", ssoservice, nil)
		u, _ := gosaml.SAMLRequest2Url(request, relayState, "", "", "")
		log.Println(request.Doc.Dump(true))
		http.Redirect(w, r, u.String(), http.StatusFound)
	}
	return
}

func BirkService(w http.ResponseWriter, r *http.Request) (err error) {
	// use incoming request for crafting the new one
	// remember to add the Scoping element to inform the IdP of requesterID - if stated in metadata for the IdP
	// check ad-hoc feds overlab
	defer r.Body.Close()
	// get the sp as well to check for allowed acs
	request, _, mdbirkidp, relayState, err := gosaml.ReceiveSAMLRequest(r, config.External, config.External)
	if err != nil {
		return
	}
	// Save the issuer and destination in a cookie for when the response comes back

	cookievalue, err := seccookie.Encode("BIRK", gosaml.Deflate(request.Doc.Dump(true)))
	http.SetCookie(w, &http.Cookie{Name: "BIRK", Value: cookievalue, Domain: config.Domain, Path: "/", Secure: true, HttpOnly: true})

	idp := debify.ReplaceAllString(mdbirkidp.Query1(nil, "@entityID"), "$1$2")

	var mdhub, mdidp *goxml.Xp
	// are we remapping - for now only use case is https://nemlogin.wayf.dk -> https://saml.nemlog-in.dk
	if rm, ok := remap[idp]; ok {
		mdidp, err = config.Internal.MDQ(rm.idp)
		mdhub, err = config.Internal.MDQ(rm.sp)
	} else {
		mdidp, err = config.Internal.MDQ(idp)
		mdhub, err = config.Hub.MDQ(config.HubEntityID)
	}
	// use a std request - we take care of NameID etc in acsService below
	newrequest := gosaml.NewAuthnRequest(config.StdTiming.Refresh(), mdhub, mdidp)
	u, _ := gosaml.SAMLRequest2Url(newrequest, relayState, "", "", "") // not signed so blank key, pw and algo
	http.Redirect(w, r, u.String(), http.StatusFound)
	return
}

func AcsService(w http.ResponseWriter, r *http.Request) (err error) {
	defer r.Body.Close()
	birk, err := r.Cookie("BIRK")
	if err != nil {
		return err
	}

	value := []byte{}
	if err = seccookie.Decode("BIRK", birk.Value, &value); err != nil {
		return
	}

	// we checked the request when we received in birkService - we can use it without fear ie. we just parse it
	log.Println("cookie", string(gosaml.Inflate(value)))
	request := goxml.NewXp(string(gosaml.Inflate(value)))

	http.SetCookie(w, &http.Cookie{Name: "BIRK", Value: "", Domain: config.Domain, Path: "/", Secure: true, HttpOnly: true, MaxAge: -1})
	sp_md, err := config.External.MDQ(request.Query1(nil, "/samlp:AuthnRequest/saml:Issuer"))
	if err != nil {
		return
	}

	response, idp_md, _, relayState, err := gosaml.ReceiveSAMLResponse(r, config.Internal, config.Hub)
	if err != nil {
		return
	}

	err = config.AttributeHandler(idp_md, config.HubRequestedAttributes, sp_md, response)
	if err != nil {
		return
	}

	birkmd, err := config.External.MDQ(request.Query1(nil, "/samlp:AuthnRequest/@Destination"))
	if err != nil {
		return
	}
	nameid := response.Query(nil, "./saml:Assertion/saml:Subject/saml:NameID")[0]
	// respect nameID in req, give persistent id + all computed attributes + nameformat conversion
	nameidformat := sp_md.Query1(nil, "./md:SPSSODescriptor/md:NameIDFormat")
	if nameidformat == gosaml.Persistent {
		response.QueryDashP(nameid, "@Format", gosaml.Persistent, nil)
		eptid := response.Query1(nil, `./saml:Assertion/saml:AttributeStatement/saml:Attribute[@FriendlyName="eduPersonTargetedID"]/saml:AttributeValue`)
		response.QueryDashP(nameid, ".", eptid, nil)
	} else if nameidformat == gosaml.Transient {
		response.QueryDashP(nameid, ".", gosaml.Id(), nil)
	}

	newresponse := gosaml.NewResponse(config.StdTiming.Refresh(), birkmd, sp_md, request, response)

	for _, q := range config.ElementsToSign {
		err = gosaml.SignResponse(newresponse, q, birkmd)
		if err != nil {
			return
		}
	}

	// when consent as a service is ready - we will post to that
	acs := newresponse.Query1(nil, "@Destination")

	data := formdata{Acs: acs, Samlresponse: base64.StdEncoding.EncodeToString([]byte(newresponse.Doc.Dump(false))), RelayState: relayState}
	postform.Execute(w, data)
	return
}

func KribService(w http.ResponseWriter, r *http.Request) (err error) {
	// check ad-hoc feds overlap
	defer r.Body.Close()

	response, _, _, relayState, err := gosaml.ReceiveSAMLResponse(r, config.External, config.External)
	if err != nil {
		return
	}
	destination := debify.ReplaceAllString(response.Query1(nil, "@Destination"), "$1$2")
	response.QueryDashP(nil, "@Destination", destination, nil)
	response.QueryDashP(nil, "./saml:Assertion/saml:Subject/saml:SubjectConfirmation/saml:SubjectConfirmationData/@Recipient", destination, nil)
	issuer := config.HubEntityID
	response.QueryDashP(nil, "./saml:Issuer", issuer, nil)
	response.QueryDashP(nil, "./saml:Assertion/saml:Issuer", issuer, nil)
	// Krib always receives attributes with nameformat=urn. Before sending to the real SP we need to look into
	// the metadata for SP to determine the actual nameformat - as WAYF supports both for internal SPs.
	mdsp, err := config.Internal.MDQ(destination)
	if err != nil {
		return
	}
	requestedattributes := mdsp.Query(nil, "./md:SPSSODescriptor/md:AttributeConsumingService/md:RequestedAttribute")
	attributestatement := response.Query(nil, "./saml:Assertion/saml:AttributeStatement")[0]
	for _, attr := range requestedattributes {
		nameFormat, _ := attr.(types.Element).GetAttribute("NameFormat")
		if nameFormat.NodeValue() == gosaml.Basic {
			basicname, _ := attr.(types.Element).GetAttribute("Name")
			uriname := config.Basic2uri[basicname.NodeValue()]
			responseattribute := response.Query(attributestatement, "saml:Attribute[@Name='"+uriname+"']")
			if len(responseattribute) > 0 {
				responseattribute[0].(types.Element).SetAttribute("Name", basicname.NodeValue())
				responseattribute[0].(types.Element).SetAttribute("NameFormat", gosaml.Basic)
			}
		}
	}

	mdhub, err := config.Hub.MDQ(config.HubEntityID)
	if err != nil {
		return
	}

	for _, q := range config.ElementsToSign {
		err = gosaml.SignResponse(response, q, mdhub)
		if err != nil {
			return
		}
	}

	data := formdata{Acs: destination, Samlresponse: base64.StdEncoding.EncodeToString([]byte(response.Doc.Dump(false))), RelayState: relayState}
	postform.Execute(w, data)
	return
}
<|endoftext|>


<commit_before>package logr

import (
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/tcotav/base/config"
	"os"
	"runtime"
)

var stackTrace = false
var configFile = "log.cfg"
var log = logrus.New()

const Linfo = "info"
const Lfatal = "fatal"
const Lwarn = "lwarn"
const Ldebug = "debug"
const Lpanic = "panic"
const Lerror = "error"

// SetConfig sets the path to the log config file we want to use AND resets the
// module to use it.
func SetConfig(path string) {
	configFile = path
	log = logrus.New()
}

func setLogLevel(lvl string) {
	switch lvl {
	case Linfo:
		logrus.SetLevel(logrus.InfoLevel)
	case Lerror:
		logrus.SetLevel(logrus.ErrorLevel)
	case Lfatal:
		logrus.SetLevel(logrus.FatalLevel)
	case Lwarn:
		logrus.SetLevel(logrus.WarnLevel)
	case Ldebug:
		logrus.SetLevel(logrus.DebugLevel)
	case Lpanic:
		logrus.SetLevel(logrus.PanicLevel)
	}
}

func init() {
	logcfg, err := config.ParseConfig(configFile)
	if err != nil {
		log.Printf("%s logging config not found")
		return
	}

	// do we want to dump stack traces?
	s, _ := logcfg["stacktrace"]

	if s == "true" {
		stackTrace = true
	}

	// put overrides here
	logLevel := logcfg["loglevel"]
	if logLevel != "" {
		setLogLevel(logLevel)
	}
	return
}

func DumpStackTrace(lvl string, tagsrc string, msg string) {
	if stackTrace {
		//stack trace
		var stack [4096]byte
		runtime.Stack(stack[:], false)
		logLine(lvl, tagsrc, fmt.Sprintf("%s", stack[:]))
	}
}

func LogLine(lvl string, tagsrc string, msg string) {
	logLine(lvl, tagsrc, msg)
	// additional actions
	switch lvl {
	case Lerror:
		DumpStackTrace(lvl, tagsrc, msg)
	case Lfatal:
		DumpStackTrace(lvl, tagsrc, msg)
		os.Exit(3)
	}
}

// logLine only handles the logging to target
func logLine(lvl string, tagsrc string, msg string) {
	l := log.WithFields(logrus.Fields{
		"src": tagsrc,
	})
	switch lvl {
	case Linfo:
		l.Info(msg)
	case Lerror:
		l.Error(msg)
	case Lfatal:
		l.Fatal(msg)
	case Lwarn:
		l.Warn(msg)
	case Ldebug:
		l.Debug(msg)
	case Lpanic:
		l.Panic(msg)
	default:
		l.Info(msg)
	}
}
<commit_msg>no config handled a bit better<commit_after>package logr

import (
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/tcotav/base/config"
	"os"
	"runtime"
)

var stackTrace = false
var configFile = "log.cfg"
var log = logrus.New()

const Linfo = "info"
const Lfatal = "fatal"
const Lwarn = "lwarn"
const Ldebug = "debug"
const Lpanic = "panic"
const Lerror = "error"

// SetConfig sets the path to the log config file we want to use AND resets the
// module to use it.
func SetConfig(path string) {
	configFile = path
	log = logrus.New()
}

func setLogLevel(lvl string) {
	switch lvl {
	case Linfo:
		logrus.SetLevel(logrus.InfoLevel)
	case Lerror:
		logrus.SetLevel(logrus.ErrorLevel)
	case Lfatal:
		logrus.SetLevel(logrus.FatalLevel)
	case Lwarn:
		logrus.SetLevel(logrus.WarnLevel)
	case Ldebug:
		logrus.SetLevel(logrus.DebugLevel)
	case Lpanic:
		logrus.SetLevel(logrus.PanicLevel)
	}
}

func init() {
	logcfg, err := config.ParseConfig(configFile)
	if err != nil {
		log.Printf("No logging config found -- using defaults")
		setLogLevel("info")
		return
	}

	// do we want to dump stack traces?
	s, _ := logcfg["stacktrace"]

	if s == "true" {
		stackTrace = true
	}

	// put overrides here
	logLevel := logcfg["loglevel"]
	if logLevel != "" {
		setLogLevel(logLevel)
	}
	return
}

func DumpStackTrace(lvl string, tagsrc string, msg string) {
	if stackTrace {
		//stack trace
		var stack [4096]byte
		runtime.Stack(stack[:], false)
		logLine(lvl, tagsrc, fmt.Sprintf("%s", stack[:]))
	}
}

func LogLine(lvl string, tagsrc string, msg string) {
	logLine(lvl, tagsrc, msg)
	// additional actions
	switch lvl {
	case Lerror:
		DumpStackTrace(lvl, tagsrc, msg)
	case Lfatal:
		DumpStackTrace(lvl, tagsrc, msg)
		os.Exit(3)
	}
}

// logLine only handles the logging to target
func logLine(lvl string, tagsrc string, msg string) {
	l := log.WithFields(logrus.Fields{
		"src": tagsrc,
	})
	switch lvl {
	case Linfo:
		l.Info(msg)
	case Lerror:
		l.Error(msg)
	case Lfatal:
		l.Fatal(msg)
	case Lwarn:
		l.Warn(msg)
	case Ldebug:
		l.Debug(msg)
	case Lpanic:
		l.Panic(msg)
	default:
		l.Info(msg)
	}
}
<|endoftext|>


<commit_before>//  Copyright (c) 2014 Couchbase, Inc.
//  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 searchers

import (
	"math"

	"github.com/blevesearch/bleve/index"
	"github.com/blevesearch/bleve/search"
	"github.com/blevesearch/bleve/search/scorers"
)

type BooleanSearcher struct {
	indexReader     index.IndexReader
	mustSearcher    search.Searcher
	shouldSearcher  search.Searcher
	mustNotSearcher search.Searcher
	queryNorm       float64
	currMust        *search.DocumentMatch
	currShould      *search.DocumentMatch
	currMustNot     *search.DocumentMatch
	currentID       index.IndexInternalID
	min             uint64
	scorer          *scorers.ConjunctionQueryScorer
	matches         []*search.DocumentMatch
	initialized     bool
}

func NewBooleanSearcher(indexReader index.IndexReader, mustSearcher search.Searcher, shouldSearcher search.Searcher, mustNotSearcher search.Searcher, explain bool) (*BooleanSearcher, error) {
	// build our searcher
	rv := BooleanSearcher{
		indexReader:     indexReader,
		mustSearcher:    mustSearcher,
		shouldSearcher:  shouldSearcher,
		mustNotSearcher: mustNotSearcher,
		scorer:          scorers.NewConjunctionQueryScorer(explain),
		matches:         make([]*search.DocumentMatch, 2),
	}
	rv.computeQueryNorm()
	return &rv, nil
}

func (s *BooleanSearcher) computeQueryNorm() {
	// first calculate sum of squared weights
	sumOfSquaredWeights := 0.0
	if s.mustSearcher != nil {
		sumOfSquaredWeights += s.mustSearcher.Weight()
	}
	if s.shouldSearcher != nil {
		sumOfSquaredWeights += s.shouldSearcher.Weight()
	}

	// now compute query norm from this
	s.queryNorm = 1.0 / math.Sqrt(sumOfSquaredWeights)
	// finally tell all the downstream searchers the norm
	if s.mustSearcher != nil {
		s.mustSearcher.SetQueryNorm(s.queryNorm)
	}
	if s.shouldSearcher != nil {
		s.shouldSearcher.SetQueryNorm(s.queryNorm)
	}
}

func (s *BooleanSearcher) initSearchers(ctx *search.SearchContext) error {
	var err error
	// get all searchers pointing at their first match
	if s.mustSearcher != nil {
		if s.currMust != nil {
			ctx.DocumentMatchPool.Put(s.currMust)
		}
		s.currMust, err = s.mustSearcher.Next(ctx)
		if err != nil {
			return err
		}
	}

	if s.shouldSearcher != nil {
		if s.currShould != nil {
			ctx.DocumentMatchPool.Put(s.currShould)
		}
		s.currShould, err = s.shouldSearcher.Next(ctx)
		if err != nil {
			return err
		}
	}

	if s.mustNotSearcher != nil {
		if s.currMustNot != nil {
			ctx.DocumentMatchPool.Put(s.currMustNot)
		}
		s.currMustNot, err = s.mustNotSearcher.Next(ctx)
		if err != nil {
			return err
		}
	}

	if s.mustSearcher != nil && s.currMust != nil {
		s.currentID = s.currMust.IndexInternalID
	} else if s.mustSearcher == nil && s.currShould != nil {
		s.currentID = s.currShould.IndexInternalID
	} else {
		s.currentID = nil
	}

	s.initialized = true
	return nil
}

func (s *BooleanSearcher) advanceNextMust(ctx *search.SearchContext, skipReturn *search.DocumentMatch) error {
	var err error

	if s.mustSearcher != nil {
		if s.currMust != skipReturn {
			ctx.DocumentMatchPool.Put(s.currMust)
		}
		s.currMust, err = s.mustSearcher.Next(ctx)
		if err != nil {
			return err
		}
	} else if s.mustSearcher == nil {
		if s.currShould != skipReturn {
			ctx.DocumentMatchPool.Put(s.currShould)
		}
		s.currShould, err = s.shouldSearcher.Next(ctx)
		if err != nil {
			return err
		}
	}

	if s.mustSearcher != nil && s.currMust != nil {
		s.currentID = s.currMust.IndexInternalID
	} else if s.mustSearcher == nil && s.currShould != nil {
		s.currentID = s.currShould.IndexInternalID
	} else {
		s.currentID = nil
	}
	return nil
}

func (s *BooleanSearcher) Weight() float64 {
	var rv float64
	if s.mustSearcher != nil {
		rv += s.mustSearcher.Weight()
	}
	if s.shouldSearcher != nil {
		rv += s.shouldSearcher.Weight()
	}

	return rv
}

func (s *BooleanSearcher) SetQueryNorm(qnorm float64) {
	if s.mustSearcher != nil {
		s.mustSearcher.SetQueryNorm(qnorm)
	}
	if s.shouldSearcher != nil {
		s.shouldSearcher.SetQueryNorm(qnorm)
	}
}

func (s *BooleanSearcher) Next(ctx *search.SearchContext) (*search.DocumentMatch, error) {

	if !s.initialized {
		err := s.initSearchers(ctx)
		if err != nil {
			return nil, err
		}
	}

	var err error
	var rv *search.DocumentMatch

	for s.currentID != nil {
		if s.currMustNot != nil {
			cmp := s.currMustNot.IndexInternalID.Compare(s.currentID)
			if cmp < 0 {
				ctx.DocumentMatchPool.Put(s.currMustNot)
				// advance must not searcher to our candidate entry
				s.currMustNot, err = s.mustNotSearcher.Advance(ctx, s.currentID)
				if err != nil {
					return nil, err
				}
				if s.currMustNot != nil && s.currMustNot.IndexInternalID.Equals(s.currentID) {
					// the candidate is excluded
					err = s.advanceNextMust(ctx, nil)
					if err != nil {
						return nil, err
					}
					continue
				}
			} else if cmp == 0 {
				// the candidate is excluded
				err = s.advanceNextMust(ctx, nil)
				if err != nil {
					return nil, err
				}
				continue
			}
		}

		shouldCmpOrNil := 1 // NOTE: shouldCmp will also be 1 when currShould == nil.
		if s.currShould != nil {
			shouldCmpOrNil = s.currShould.IndexInternalID.Compare(s.currentID)
		}

		if shouldCmpOrNil < 0 {
			ctx.DocumentMatchPool.Put(s.currShould)
			// advance should searcher to our candidate entry
			s.currShould, err = s.shouldSearcher.Advance(ctx, s.currentID)
			if err != nil {
				return nil, err
			}
			if s.currShould != nil && s.currShould.IndexInternalID.Equals(s.currentID) {
				// score bonus matches should
				var cons []*search.DocumentMatch
				if s.currMust != nil {
					cons = s.matches
					cons[0] = s.currMust
					cons[1] = s.currShould
				} else {
					cons = s.matches[0:1]
					cons[0] = s.currShould
				}
				rv = s.scorer.Score(ctx, cons)
				err = s.advanceNextMust(ctx, rv)
				if err != nil {
					return nil, err
				}
				break
			} else if s.shouldSearcher.Min() == 0 {
				// match is OK anyway
				cons := s.matches[0:1]
				cons[0] = s.currMust
				rv = s.scorer.Score(ctx, cons)
				err = s.advanceNextMust(ctx, rv)
				if err != nil {
					return nil, err
				}
				break
			}
		} else if shouldCmpOrNil == 0 {
			// score bonus matches should
			var cons []*search.DocumentMatch
			if s.currMust != nil {
				cons = s.matches
				cons[0] = s.currMust
				cons[1] = s.currShould
			} else {
				cons = s.matches[0:1]
				cons[0] = s.currShould
			}
			rv = s.scorer.Score(ctx, cons)
			err = s.advanceNextMust(ctx, rv)
			if err != nil {
				return nil, err
			}
			break
		} else if s.shouldSearcher == nil || s.shouldSearcher.Min() == 0 {
			// match is OK anyway
			cons := s.matches[0:1]
			cons[0] = s.currMust
			rv = s.scorer.Score(ctx, cons)
			err = s.advanceNextMust(ctx, rv)
			if err != nil {
				return nil, err
			}
			break
		}

		err = s.advanceNextMust(ctx, nil)
		if err != nil {
			return nil, err
		}
	}
	return rv, nil
}

func (s *BooleanSearcher) Advance(ctx *search.SearchContext, ID index.IndexInternalID) (*search.DocumentMatch, error) {

	if !s.initialized {
		err := s.initSearchers(ctx)
		if err != nil {
			return nil, err
		}
	}

	var err error
	if s.mustSearcher != nil {
		if s.currMust != nil {
			ctx.DocumentMatchPool.Put(s.currMust)
		}
		s.currMust, err = s.mustSearcher.Advance(ctx, ID)
		if err != nil {
			return nil, err
		}
	}
	if s.shouldSearcher != nil {
		if s.currShould != nil {
			ctx.DocumentMatchPool.Put(s.currShould)
		}
		s.currShould, err = s.shouldSearcher.Advance(ctx, ID)
		if err != nil {
			return nil, err
		}
	}
	if s.mustNotSearcher != nil {
		if s.currMustNot != nil {
			ctx.DocumentMatchPool.Put(s.currMustNot)
		}
		s.currMustNot, err = s.mustNotSearcher.Advance(ctx, ID)
		if err != nil {
			return nil, err
		}
	}

	if s.mustSearcher != nil && s.currMust != nil {
		s.currentID = s.currMust.IndexInternalID
	} else if s.mustSearcher == nil && s.currShould != nil {
		s.currentID = s.currShould.IndexInternalID
	} else {
		s.currentID = nil
	}

	return s.Next(ctx)
}

func (s *BooleanSearcher) Count() uint64 {

	// for now return a worst case
	var sum uint64
	if s.mustSearcher != nil {
		sum += s.mustSearcher.Count()
	}
	if s.shouldSearcher != nil {
		sum += s.shouldSearcher.Count()
	}
	return sum
}

func (s *BooleanSearcher) Close() error {
	if s.mustSearcher != nil {
		err := s.mustSearcher.Close()
		if err != nil {
			return err
		}
	}
	if s.shouldSearcher != nil {
		err := s.shouldSearcher.Close()
		if err != nil {
			return err
		}
	}
	if s.mustNotSearcher != nil {
		err := s.mustNotSearcher.Close()
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *BooleanSearcher) Min() int {
	return 0
}

func (s *BooleanSearcher) DocumentMatchPoolSize() int {
	rv := 3
	if s.mustSearcher != nil {
		rv += s.mustSearcher.DocumentMatchPoolSize()
	}
	if s.shouldSearcher != nil {
		rv += s.shouldSearcher.DocumentMatchPoolSize()
	}
	if s.mustNotSearcher != nil {
		rv += s.mustNotSearcher.DocumentMatchPoolSize()
	}
	return rv
}
<commit_msg>simplify BooleanSearcher mustSearcher else logic<commit_after>//  Copyright (c) 2014 Couchbase, Inc.
//  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 searchers

import (
	"math"

	"github.com/blevesearch/bleve/index"
	"github.com/blevesearch/bleve/search"
	"github.com/blevesearch/bleve/search/scorers"
)

type BooleanSearcher struct {
	indexReader     index.IndexReader
	mustSearcher    search.Searcher
	shouldSearcher  search.Searcher
	mustNotSearcher search.Searcher
	queryNorm       float64
	currMust        *search.DocumentMatch
	currShould      *search.DocumentMatch
	currMustNot     *search.DocumentMatch
	currentID       index.IndexInternalID
	min             uint64
	scorer          *scorers.ConjunctionQueryScorer
	matches         []*search.DocumentMatch
	initialized     bool
}

func NewBooleanSearcher(indexReader index.IndexReader, mustSearcher search.Searcher, shouldSearcher search.Searcher, mustNotSearcher search.Searcher, explain bool) (*BooleanSearcher, error) {
	// build our searcher
	rv := BooleanSearcher{
		indexReader:     indexReader,
		mustSearcher:    mustSearcher,
		shouldSearcher:  shouldSearcher,
		mustNotSearcher: mustNotSearcher,
		scorer:          scorers.NewConjunctionQueryScorer(explain),
		matches:         make([]*search.DocumentMatch, 2),
	}
	rv.computeQueryNorm()
	return &rv, nil
}

func (s *BooleanSearcher) computeQueryNorm() {
	// first calculate sum of squared weights
	sumOfSquaredWeights := 0.0
	if s.mustSearcher != nil {
		sumOfSquaredWeights += s.mustSearcher.Weight()
	}
	if s.shouldSearcher != nil {
		sumOfSquaredWeights += s.shouldSearcher.Weight()
	}

	// now compute query norm from this
	s.queryNorm = 1.0 / math.Sqrt(sumOfSquaredWeights)
	// finally tell all the downstream searchers the norm
	if s.mustSearcher != nil {
		s.mustSearcher.SetQueryNorm(s.queryNorm)
	}
	if s.shouldSearcher != nil {
		s.shouldSearcher.SetQueryNorm(s.queryNorm)
	}
}

func (s *BooleanSearcher) initSearchers(ctx *search.SearchContext) error {
	var err error
	// get all searchers pointing at their first match
	if s.mustSearcher != nil {
		if s.currMust != nil {
			ctx.DocumentMatchPool.Put(s.currMust)
		}
		s.currMust, err = s.mustSearcher.Next(ctx)
		if err != nil {
			return err
		}
	}

	if s.shouldSearcher != nil {
		if s.currShould != nil {
			ctx.DocumentMatchPool.Put(s.currShould)
		}
		s.currShould, err = s.shouldSearcher.Next(ctx)
		if err != nil {
			return err
		}
	}

	if s.mustNotSearcher != nil {
		if s.currMustNot != nil {
			ctx.DocumentMatchPool.Put(s.currMustNot)
		}
		s.currMustNot, err = s.mustNotSearcher.Next(ctx)
		if err != nil {
			return err
		}
	}

	if s.mustSearcher != nil && s.currMust != nil {
		s.currentID = s.currMust.IndexInternalID
	} else if s.mustSearcher == nil && s.currShould != nil {
		s.currentID = s.currShould.IndexInternalID
	} else {
		s.currentID = nil
	}

	s.initialized = true
	return nil
}

func (s *BooleanSearcher) advanceNextMust(ctx *search.SearchContext, skipReturn *search.DocumentMatch) error {
	var err error

	if s.mustSearcher != nil {
		if s.currMust != skipReturn {
			ctx.DocumentMatchPool.Put(s.currMust)
		}
		s.currMust, err = s.mustSearcher.Next(ctx)
		if err != nil {
			return err
		}
	} else {
		if s.currShould != skipReturn {
			ctx.DocumentMatchPool.Put(s.currShould)
		}
		s.currShould, err = s.shouldSearcher.Next(ctx)
		if err != nil {
			return err
		}
	}

	if s.mustSearcher != nil && s.currMust != nil {
		s.currentID = s.currMust.IndexInternalID
	} else if s.mustSearcher == nil && s.currShould != nil {
		s.currentID = s.currShould.IndexInternalID
	} else {
		s.currentID = nil
	}
	return nil
}

func (s *BooleanSearcher) Weight() float64 {
	var rv float64
	if s.mustSearcher != nil {
		rv += s.mustSearcher.Weight()
	}
	if s.shouldSearcher != nil {
		rv += s.shouldSearcher.Weight()
	}

	return rv
}

func (s *BooleanSearcher) SetQueryNorm(qnorm float64) {
	if s.mustSearcher != nil {
		s.mustSearcher.SetQueryNorm(qnorm)
	}
	if s.shouldSearcher != nil {
		s.shouldSearcher.SetQueryNorm(qnorm)
	}
}

func (s *BooleanSearcher) Next(ctx *search.SearchContext) (*search.DocumentMatch, error) {

	if !s.initialized {
		err := s.initSearchers(ctx)
		if err != nil {
			return nil, err
		}
	}

	var err error
	var rv *search.DocumentMatch

	for s.currentID != nil {
		if s.currMustNot != nil {
			cmp := s.currMustNot.IndexInternalID.Compare(s.currentID)
			if cmp < 0 {
				ctx.DocumentMatchPool.Put(s.currMustNot)
				// advance must not searcher to our candidate entry
				s.currMustNot, err = s.mustNotSearcher.Advance(ctx, s.currentID)
				if err != nil {
					return nil, err
				}
				if s.currMustNot != nil && s.currMustNot.IndexInternalID.Equals(s.currentID) {
					// the candidate is excluded
					err = s.advanceNextMust(ctx, nil)
					if err != nil {
						return nil, err
					}
					continue
				}
			} else if cmp == 0 {
				// the candidate is excluded
				err = s.advanceNextMust(ctx, nil)
				if err != nil {
					return nil, err
				}
				continue
			}
		}

		shouldCmpOrNil := 1 // NOTE: shouldCmp will also be 1 when currShould == nil.
		if s.currShould != nil {
			shouldCmpOrNil = s.currShould.IndexInternalID.Compare(s.currentID)
		}

		if shouldCmpOrNil < 0 {
			ctx.DocumentMatchPool.Put(s.currShould)
			// advance should searcher to our candidate entry
			s.currShould, err = s.shouldSearcher.Advance(ctx, s.currentID)
			if err != nil {
				return nil, err
			}
			if s.currShould != nil && s.currShould.IndexInternalID.Equals(s.currentID) {
				// score bonus matches should
				var cons []*search.DocumentMatch
				if s.currMust != nil {
					cons = s.matches
					cons[0] = s.currMust
					cons[1] = s.currShould
				} else {
					cons = s.matches[0:1]
					cons[0] = s.currShould
				}
				rv = s.scorer.Score(ctx, cons)
				err = s.advanceNextMust(ctx, rv)
				if err != nil {
					return nil, err
				}
				break
			} else if s.shouldSearcher.Min() == 0 {
				// match is OK anyway
				cons := s.matches[0:1]
				cons[0] = s.currMust
				rv = s.scorer.Score(ctx, cons)
				err = s.advanceNextMust(ctx, rv)
				if err != nil {
					return nil, err
				}
				break
			}
		} else if shouldCmpOrNil == 0 {
			// score bonus matches should
			var cons []*search.DocumentMatch
			if s.currMust != nil {
				cons = s.matches
				cons[0] = s.currMust
				cons[1] = s.currShould
			} else {
				cons = s.matches[0:1]
				cons[0] = s.currShould
			}
			rv = s.scorer.Score(ctx, cons)
			err = s.advanceNextMust(ctx, rv)
			if err != nil {
				return nil, err
			}
			break
		} else if s.shouldSearcher == nil || s.shouldSearcher.Min() == 0 {
			// match is OK anyway
			cons := s.matches[0:1]
			cons[0] = s.currMust
			rv = s.scorer.Score(ctx, cons)
			err = s.advanceNextMust(ctx, rv)
			if err != nil {
				return nil, err
			}
			break
		}

		err = s.advanceNextMust(ctx, nil)
		if err != nil {
			return nil, err
		}
	}
	return rv, nil
}

func (s *BooleanSearcher) Advance(ctx *search.SearchContext, ID index.IndexInternalID) (*search.DocumentMatch, error) {

	if !s.initialized {
		err := s.initSearchers(ctx)
		if err != nil {
			return nil, err
		}
	}

	var err error
	if s.mustSearcher != nil {
		if s.currMust != nil {
			ctx.DocumentMatchPool.Put(s.currMust)
		}
		s.currMust, err = s.mustSearcher.Advance(ctx, ID)
		if err != nil {
			return nil, err
		}
	}
	if s.shouldSearcher != nil {
		if s.currShould != nil {
			ctx.DocumentMatchPool.Put(s.currShould)
		}
		s.currShould, err = s.shouldSearcher.Advance(ctx, ID)
		if err != nil {
			return nil, err
		}
	}
	if s.mustNotSearcher != nil {
		if s.currMustNot != nil {
			ctx.DocumentMatchPool.Put(s.currMustNot)
		}
		s.currMustNot, err = s.mustNotSearcher.Advance(ctx, ID)
		if err != nil {
			return nil, err
		}
	}

	if s.mustSearcher != nil && s.currMust != nil {
		s.currentID = s.currMust.IndexInternalID
	} else if s.mustSearcher == nil && s.currShould != nil {
		s.currentID = s.currShould.IndexInternalID
	} else {
		s.currentID = nil
	}

	return s.Next(ctx)
}

func (s *BooleanSearcher) Count() uint64 {

	// for now return a worst case
	var sum uint64
	if s.mustSearcher != nil {
		sum += s.mustSearcher.Count()
	}
	if s.shouldSearcher != nil {
		sum += s.shouldSearcher.Count()
	}
	return sum
}

func (s *BooleanSearcher) Close() error {
	if s.mustSearcher != nil {
		err := s.mustSearcher.Close()
		if err != nil {
			return err
		}
	}
	if s.shouldSearcher != nil {
		err := s.shouldSearcher.Close()
		if err != nil {
			return err
		}
	}
	if s.mustNotSearcher != nil {
		err := s.mustNotSearcher.Close()
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *BooleanSearcher) Min() int {
	return 0
}

func (s *BooleanSearcher) DocumentMatchPoolSize() int {
	rv := 3
	if s.mustSearcher != nil {
		rv += s.mustSearcher.DocumentMatchPoolSize()
	}
	if s.shouldSearcher != nil {
		rv += s.shouldSearcher.DocumentMatchPoolSize()
	}
	if s.mustNotSearcher != nil {
		rv += s.mustNotSearcher.DocumentMatchPoolSize()
	}
	return rv
}
<|endoftext|>


<commit_before>package platform

import (
	"context"
	"fmt"
	"log"
	"strings"

	installer "github.com/tektoncd/operator/pkg/reconciler/shared/tektoninstallerset"
	"knative.dev/pkg/injection"
	"knative.dev/pkg/injection/sharedmain"
	"knative.dev/pkg/signals"
)

// validateControllerNamesOrDie ensures that the list of controller names to be enabled
// are supported by a platform. This function exits on error
func validateControllerNamesOrDie(p Platform) {
	if err := validateControllerNames(p); err != nil {
		log.Fatalf("error validating provided controller names: %v", err)
	}

}

// validateControllerNames ensures that the list of controller names to be enabled
// are supported by a platform
func validateControllerNames(p Platform) error {
	pParams := p.PlatformParams()
	supportedCtrls := p.AllSupportedControllers()
	invalidNamesStr := invalidNames(supportedCtrls, pParams.ControllerNames)
	if len(invalidNamesStr) == 0 {
		return nil
	}
	return ErrorControllerNames(invalidNamesStr, supportedCtrls.ControllerNames())
}

// invalidNames checks if whether there are any names in []CotrollerNames which are
// not present in (supported by) given ControllerMap
func invalidNames(supportedCtrls ControllerMap, cNames []ControllerName) string {
	invalidNames := strings.Builder{}
	for _, cName := range cNames {
		if _, ok := supportedCtrls[cName]; !ok {
			invalidNames.WriteString(string(cName))
			invalidNames.WriteString(",")
		}
	}

	return strings.TrimSuffix(invalidNames.String(), ",")
}

// ErrorControllerNames is a error message format helper
func ErrorControllerNames(invalidNames string, validNames []string) error {
	return fmt.Errorf("un-identified controller names: %s, supported names: %v", invalidNames, validNames)
}

// activeControllers returns a map of the controllers that should be run
// the returned map is a subset of the platform specific map which stores all-supported-controllers
func activeControllers(p Platform) ControllerMap {
	pParams := p.PlatformParams()
	result := ControllerMap{}
	for _, name := range pParams.ControllerNames {
		if namedCtrl, ok := p.AllSupportedControllers()[name]; ok {
			result[name] = namedCtrl
		}
	}
	return result
}

// disabledControllers returns a map of the controllers that should not be run
// the result of disabledControllers is the set of controllers excluded by activeControllers function
// in other words, disabledControllers returns a map which has controllers "not" specified in the controlelrNames input to a platform
// the returned map is a subset of the platform specific map which stores all-supported-controllers
/*
Copyright 2022 The Tekton 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.
*/

func disabledControllers(p Platform) ControllerMap {
	pParams := p.PlatformParams()
	result := p.AllSupportedControllers()
	for _, name := range pParams.ControllerNames {
		delete(result, name)
	}
	return result
}

// contextWithPlatformName  adds platform name to a given context
func contextWithPlatformName(ctx context.Context, pName string) context.Context {
	ctx = context.WithValue(ctx, PlatformNameKey{}, pName)
	return ctx
}

// startMain starts a knative/pkg sharedMain with a context that stores platform name
// and a list of controllers which should be enabled for the given platform
func startMain(p Platform, ctrls ControllerMap) {
	pParams := p.PlatformParams()
	cfg := injection.ParseAndGetRESTConfigOrDie()
	ctx, _ := injection.EnableInjectionOrDie(signals.NewContext(), cfg)
	ctx = contextWithPlatformName(ctx, pParams.Name)
	installer.InitTektonInstallerSetClient(ctx)
	sharedmain.MainWithConfig(ctx,
		pParams.SharedMainName,
		cfg,
		ctrls.ControllerConstructors()...,
	)
}

// StartMainWithAllControllers calls startMain with all controllers
// supported by a platform
func StartMainWithAllControllers(p Platform) {
	startMain(p, p.AllSupportedControllers())
}

// StartMainWithAllControllers calls startMain with a subset of controllers
// specified in the platformConfig of a platform
func StartMainWithSelectedControllers(p Platform) {
	validateControllerNamesOrDie(p)
	selectedCtrls := activeControllers(p)
	startMain(p, selectedCtrls)
}
<commit_msg>Increasing QPS limit to 50 to improve installation<commit_after>package platform

import (
	"context"
	"fmt"
	"log"
	"strings"

	installer "github.com/tektoncd/operator/pkg/reconciler/shared/tektoninstallerset"
	"knative.dev/pkg/injection"
	"knative.dev/pkg/injection/sharedmain"
	"knative.dev/pkg/signals"
)

// validateControllerNamesOrDie ensures that the list of controller names to be enabled
// are supported by a platform. This function exits on error
func validateControllerNamesOrDie(p Platform) {
	if err := validateControllerNames(p); err != nil {
		log.Fatalf("error validating provided controller names: %v", err)
	}

}

// validateControllerNames ensures that the list of controller names to be enabled
// are supported by a platform
func validateControllerNames(p Platform) error {
	pParams := p.PlatformParams()
	supportedCtrls := p.AllSupportedControllers()
	invalidNamesStr := invalidNames(supportedCtrls, pParams.ControllerNames)
	if len(invalidNamesStr) == 0 {
		return nil
	}
	return ErrorControllerNames(invalidNamesStr, supportedCtrls.ControllerNames())
}

// invalidNames checks if whether there are any names in []CotrollerNames which are
// not present in (supported by) given ControllerMap
func invalidNames(supportedCtrls ControllerMap, cNames []ControllerName) string {
	invalidNames := strings.Builder{}
	for _, cName := range cNames {
		if _, ok := supportedCtrls[cName]; !ok {
			invalidNames.WriteString(string(cName))
			invalidNames.WriteString(",")
		}
	}

	return strings.TrimSuffix(invalidNames.String(), ",")
}

// ErrorControllerNames is a error message format helper
func ErrorControllerNames(invalidNames string, validNames []string) error {
	return fmt.Errorf("un-identified controller names: %s, supported names: %v", invalidNames, validNames)
}

// activeControllers returns a map of the controllers that should be run
// the returned map is a subset of the platform specific map which stores all-supported-controllers
func activeControllers(p Platform) ControllerMap {
	pParams := p.PlatformParams()
	result := ControllerMap{}
	for _, name := range pParams.ControllerNames {
		if namedCtrl, ok := p.AllSupportedControllers()[name]; ok {
			result[name] = namedCtrl
		}
	}
	return result
}

// disabledControllers returns a map of the controllers that should not be run
// the result of disabledControllers is the set of controllers excluded by activeControllers function
// in other words, disabledControllers returns a map which has controllers "not" specified in the controlelrNames input to a platform
// the returned map is a subset of the platform specific map which stores all-supported-controllers
/*
Copyright 2022 The Tekton 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.
*/

func disabledControllers(p Platform) ControllerMap {
	pParams := p.PlatformParams()
	result := p.AllSupportedControllers()
	for _, name := range pParams.ControllerNames {
		delete(result, name)
	}
	return result
}

// contextWithPlatformName  adds platform name to a given context
func contextWithPlatformName(ctx context.Context, pName string) context.Context {
	ctx = context.WithValue(ctx, PlatformNameKey{}, pName)
	return ctx
}

// startMain starts a knative/pkg sharedMain with a context that stores platform name
// and a list of controllers which should be enabled for the given platform
func startMain(p Platform, ctrls ControllerMap) {
	pParams := p.PlatformParams()
	cfg := injection.ParseAndGetRESTConfigOrDie()
	cfg.QPS = 50
	ctx, _ := injection.EnableInjectionOrDie(signals.NewContext(), cfg)
	ctx = contextWithPlatformName(ctx, pParams.Name)
	installer.InitTektonInstallerSetClient(ctx)
	sharedmain.MainWithConfig(ctx,
		pParams.SharedMainName,
		cfg,
		ctrls.ControllerConstructors()...,
	)
}

// StartMainWithAllControllers calls startMain with all controllers
// supported by a platform
func StartMainWithAllControllers(p Platform) {
	startMain(p, p.AllSupportedControllers())
}

// StartMainWithAllControllers calls startMain with a subset of controllers
// specified in the platformConfig of a platform
func StartMainWithSelectedControllers(p Platform) {
	validateControllerNamesOrDie(p)
	selectedCtrls := activeControllers(p)
	startMain(p, selectedCtrls)
}
<|endoftext|>


<commit_before>package network

import (
	"net"

	"github.com/lxc/lxd/lxd/cluster"
	"github.com/lxc/lxd/lxd/state"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
)

// Network represents a LXD network.
type Network interface {
	// Load.
	init(state *state.State, id int64, name string, netType string, description string, config map[string]string, status string)
	fillConfig(config map[string]string) error

	// Config.
	ValidateName(name string) error
	Validate(config map[string]string) error
	ID() int64
	Name() string
	Type() string
	Status() string
	Config() map[string]string
	IsUsed() (bool, error)
	DHCPv4Subnet() *net.IPNet
	DHCPv6Subnet() *net.IPNet
	DHCPv4Ranges() []shared.IPRange
	DHCPv6Ranges() []shared.IPRange

	// Actions.
	Create(clusterNotification bool) error
	Start() error
	Stop() error
	Rename(name string) error
	Update(newNetwork api.NetworkPut, targetNode string, clusterNotification bool) error
	HandleHeartbeat(heartbeatData *cluster.APIHeartbeat) error
	Delete(clusterNotification bool) error
}
<commit_msg>lxd/network/network/interfaces: Replaces clusterNotification bool with cluster.ClientType<commit_after>package network

import (
	"net"

	"github.com/lxc/lxd/lxd/cluster"
	"github.com/lxc/lxd/lxd/state"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
)

// Network represents a LXD network.
type Network interface {
	// Load.
	init(state *state.State, id int64, name string, netType string, description string, config map[string]string, status string)
	fillConfig(config map[string]string) error

	// Config.
	ValidateName(name string) error
	Validate(config map[string]string) error
	ID() int64
	Name() string
	Type() string
	Status() string
	Config() map[string]string
	IsUsed() (bool, error)
	DHCPv4Subnet() *net.IPNet
	DHCPv6Subnet() *net.IPNet
	DHCPv4Ranges() []shared.IPRange
	DHCPv6Ranges() []shared.IPRange

	// Actions.
	Create(clientType cluster.ClientType) error
	Start() error
	Stop() error
	Rename(name string) error
	Update(newNetwork api.NetworkPut, targetNode string, clientType cluster.ClientType) error
	HandleHeartbeat(heartbeatData *cluster.APIHeartbeat) error
	Delete(clientType cluster.ClientType) error
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Knative 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 source

import (
	"context"
	"time"

	"go.uber.org/zap"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/tools/cache"

	sourceinformer "knative.dev/pkg/client/injection/ducks/duck/v1/source"
	kubeclient "knative.dev/pkg/client/injection/kube/client"
	"knative.dev/pkg/configmap"
	"knative.dev/pkg/controller"
	"knative.dev/pkg/injection"
	"knative.dev/pkg/injection/clients/dynamicclient"
	"knative.dev/pkg/logging"
	pkgreconciler "knative.dev/pkg/reconciler"

	kedaclient "knative.dev/eventing-autoscaler-keda/pkg/client/injection/keda/client"
	//scaledobjectinformer "knative.dev/eventing-autoscaler-keda/pkg/client/injection/keda/informers/keda/v1alpha1/scaledobject"
	kedaresources "knative.dev/eventing-autoscaler-keda/pkg/reconciler/keda"
)

const (
	// ReconcilerName is the name of the reconciler.
	ReconcilerName = "KEDASourceDucks"
)

// NewController returns a function that initializes the controller and
// Registers event handlers to enqueue events
func NewController(crd string, gvr schema.GroupVersionResource, gvk schema.GroupVersionKind) injection.ControllerConstructor {
	return func(ctx context.Context,
		cmw configmap.Watcher,
	) *controller.Impl {
		logger := logging.FromContext(ctx)
		sourceduckInformer := sourceinformer.Get(ctx)

		var sourceInformer cache.SharedIndexInformer
		var sourceLister cache.GenericLister

		var err error
		for i := 0; i < 10; i++ {
			sourceInformer, sourceLister, err = sourceduckInformer.Get(ctx, gvr)
			if err == nil {
				break
			} else if apierrors.IsNotFound(err) {
				logger.Debug("SourceDuckInformer not found -> waiting", zap.String("GVR", gvr.String()), zap.Error(err))
				time.Sleep(1 * time.Second)
			} else {
				logger.Errorw("Error getting source informer", zap.String("GVR", gvr.String()), zap.Error(err))
				return nil
			}
		}

		//	scaledobjectInformer := scaledobjectinformer.Get(ctx)

		r := &Reconciler{
			kubeClient:      kubeclient.Get(ctx),
			kedaClient:      kedaclient.Get(ctx),
			sourceInterface: dynamicclient.Get(ctx).Resource(gvr),
			sourceLister:    sourceLister,
			gvk:             gvk,
			gvr:             gvr,
		}
		impl := controller.NewImpl(r, logger, ReconcilerName)

		logger.Info("Setting up event handlers")
		sourceInformer.AddEventHandler(cache.FilteringResourceEventHandler{
			FilterFunc: pkgreconciler.AnnotationFilterFunc(kedaresources.AutoscalingClassAnnotation, kedaresources.KEDA, false),
			Handler:    controller.HandleAll(impl.Enqueue),
		})

		// FIXME don't handle updates on ScaledObject.Status field
		// scaledobjectInformer.Informer().AddEventHandler(cache.FilteringResourceEventHandler{
		// 	FilterFunc: controller.FilterControllerGVK(gvk),
		// 	Handler:    controller.HandleAll(impl.EnqueueControllerOf),
		// })

		return impl
	}
}
<commit_msg>don't reconcile on updates on ScaledObject.Status (#120)<commit_after>/*
Copyright 2020 The Knative 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 source

import (
	"context"
	"time"

	"go.uber.org/zap"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/tools/cache"

	sourceinformer "knative.dev/pkg/client/injection/ducks/duck/v1/source"
	kubeclient "knative.dev/pkg/client/injection/kube/client"
	"knative.dev/pkg/configmap"
	"knative.dev/pkg/controller"
	"knative.dev/pkg/injection"
	"knative.dev/pkg/injection/clients/dynamicclient"
	"knative.dev/pkg/logging"
	pkgreconciler "knative.dev/pkg/reconciler"

	kedaclient "knative.dev/eventing-autoscaler-keda/pkg/client/injection/keda/client"
	scaledobjectinformer "knative.dev/eventing-autoscaler-keda/pkg/client/injection/keda/informers/keda/v1alpha1/scaledobject"
	kedaresources "knative.dev/eventing-autoscaler-keda/pkg/reconciler/keda"
)

const (
	// ReconcilerName is the name of the reconciler.
	ReconcilerName = "KEDASourceDucks"
)

// NewController returns a function that initializes the controller and
// Registers event handlers to enqueue events
func NewController(crd string, gvr schema.GroupVersionResource, gvk schema.GroupVersionKind) injection.ControllerConstructor {
	return func(ctx context.Context,
		cmw configmap.Watcher,
	) *controller.Impl {
		logger := logging.FromContext(ctx)
		sourceduckInformer := sourceinformer.Get(ctx)

		var sourceInformer cache.SharedIndexInformer
		var sourceLister cache.GenericLister

		var err error
		for i := 0; i < 10; i++ {
			sourceInformer, sourceLister, err = sourceduckInformer.Get(ctx, gvr)
			if err == nil {
				break
			} else if apierrors.IsNotFound(err) {
				logger.Debug("SourceDuckInformer not found -> waiting", zap.String("GVR", gvr.String()), zap.Error(err))
				time.Sleep(1 * time.Second)
			} else {
				logger.Errorw("Error getting source informer", zap.String("GVR", gvr.String()), zap.Error(err))
				return nil
			}
		}

		scaledobjectInformer := scaledobjectinformer.Get(ctx)

		r := &Reconciler{
			kubeClient:      kubeclient.Get(ctx),
			kedaClient:      kedaclient.Get(ctx),
			sourceInterface: dynamicclient.Get(ctx).Resource(gvr),
			sourceLister:    sourceLister,
			gvk:             gvk,
			gvr:             gvr,
		}
		impl := controller.NewImpl(r, logger, ReconcilerName)

		logger.Info("Setting up event handlers")
		sourceInformer.AddEventHandler(cache.FilteringResourceEventHandler{
			FilterFunc: pkgreconciler.AnnotationFilterFunc(kedaresources.AutoscalingClassAnnotation, kedaresources.KEDA, false),
			Handler:    controller.HandleAll(impl.Enqueue),
		})

		// don't handle updates on ScaledObject.Status field
		scaledobjectInformer.Informer().AddEventHandler(cache.FilteringResourceEventHandler{
			FilterFunc: controller.FilterControllerGVK(gvk),
			Handler: cache.ResourceEventHandlerFuncs{
				AddFunc: impl.EnqueueControllerOf,
				UpdateFunc: func(old, new interface{}) {
					if mOld, ok := old.(metav1.Object); ok {
						if mNew, ok := new.(metav1.Object); ok {
							if mNew.GetGeneration() != mOld.GetGeneration() {
								impl.EnqueueControllerOf(new)
							}
						}
					}
				},
				DeleteFunc: impl.EnqueueControllerOf,
			},
		})

		return impl
	}
}
<|endoftext|>