<commit_before>/*
Copyright 2020 The cert-manager Authors.

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

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

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

package helper

import (
	"context"
	"crypto/ecdsa"
	"crypto/rsa"
	"crypto/x509"
	"fmt"
	"sort"
	"time"

	corev1 "k8s.io/api/core/v1"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"

	apiutil "github.com/jetstack/cert-manager/pkg/api/util"
	cmapi "github.com/jetstack/cert-manager/pkg/apis/certmanager/v1"
	cmmeta "github.com/jetstack/cert-manager/pkg/apis/meta/v1"
	"github.com/jetstack/cert-manager/pkg/util"
	"github.com/jetstack/cert-manager/pkg/util/pki"
	"github.com/jetstack/cert-manager/test/e2e/framework/log"
)

// WaitForCertificateReady waits for the certificate resource to enter a Ready
// state.
func (h *Helper) WaitForCertificateReady(ns, name string, timeout time.Duration) (*cmapi.Certificate, error) {
	var certificate *cmapi.Certificate
	err := wait.PollImmediate(time.Second, timeout,
		func() (bool, error) {
			var err error
			log.Logf("Waiting for Certificate %v to be ready", name)
			certificate, err = h.CMClient.CertmanagerV1().Certificates(ns).Get(context.TODO(), name, metav1.GetOptions{})
			if err != nil {
				return false, fmt.Errorf("error getting Certificate %v: %v", name, err)
			}
			isReady := apiutil.CertificateHasCondition(certificate, cmapi.CertificateCondition{
				Type:   cmapi.CertificateConditionReady,
				Status: cmmeta.ConditionTrue,
			})
			if !isReady {
				log.Logf("Expected Certificate to have Ready condition 'true' but it has: %v", certificate.Status.Conditions)
				return false, nil
			}
			return true, nil
		},
	)

	// return certificate even when error to use for debugging
	return certificate, err
}

// WaitForCertificateNotReady waits for the certificate resource to enter a
// non-Ready state.
func (h *Helper) WaitForCertificateNotReady(ns, name string, timeout time.Duration) (*cmapi.Certificate, error) {
	var certificate *cmapi.Certificate
	err := wait.PollImmediate(time.Second, timeout,
		func() (bool, error) {
			var err error
			log.Logf("Waiting for Certificate %v to be ready", name)
			certificate, err = h.CMClient.CertmanagerV1().Certificates(ns).Get(context.TODO(), name, metav1.GetOptions{})
			if err != nil {
				return false, fmt.Errorf("error getting Certificate %v: %v", name, err)
			}
			isReady := apiutil.CertificateHasCondition(certificate, cmapi.CertificateCondition{
				Type:   cmapi.CertificateConditionReady,
				Status: cmmeta.ConditionFalse,
			})
			if !isReady {
				log.Logf("Expected Certificate to have Ready condition 'true' but it has: %v", certificate.Status.Conditions)
				return false, nil
			}
			return true, nil
		},
	)

	// return certificate even when error to use for debugging
	return certificate, err
}

// ValidateIssuedCertificate will ensure that the given Certificate has a
// certificate issued for it, and that the details on the x509 certificate are
// correct as defined by the Certificate's spec.
func (h *Helper) ValidateIssuedCertificate(certificate *cmapi.Certificate, rootCAPEM []byte) (*x509.Certificate, error) {
	log.Logf("Getting the TLS certificate Secret resource")
	secret, err := h.KubeClient.CoreV1().Secrets(certificate.Namespace).Get(context.TODO(), certificate.Spec.SecretName, metav1.GetOptions{})
	if err != nil {
		return nil, err
	}
	if !(len(secret.Data) == 2 || len(secret.Data) == 3) {
		return nil, fmt.Errorf("Expected 2 keys in certificate secret, but there was %d", len(secret.Data))
	}

	keyBytes, ok := secret.Data[corev1.TLSPrivateKeyKey]
	if !ok {
		return nil, fmt.Errorf("No private key data found for Certificate %q (secret %q)", certificate.Name, certificate.Spec.SecretName)
	}
	key, err := pki.DecodePrivateKeyBytes(keyBytes)
	if err != nil {
		return nil, err
	}

	// validate private key is of the correct type (rsa or ecdsa)
	privateKey := certificate.Spec.PrivateKey
	if privateKey == nil {
		privateKey = &cmapi.CertificatePrivateKey{}
	}
	switch privateKey.Algorithm {
	case cmapi.PrivateKeyAlgorithm(""),
		cmapi.RSAKeyAlgorithm:
		_, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("Expected private key of type RSA, but it was: %T", key)
		}
	case cmapi.ECDSAKeyAlgorithm:
		_, ok := key.(*ecdsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("Expected private key of type ECDSA, but it was: %T", key)
		}
	default:
		return nil, fmt.Errorf("unrecognised requested private key algorithm %q", certificate.Spec.PrivateKey.Algorithm)
	}

	// TODO: validate private key KeySize

	// check the provided certificate is valid
	expectedOrganization := pki.OrganizationForCertificate(certificate)
	expectedDNSNames := certificate.Spec.DNSNames
	expectedIPAddresses := certificate.Spec.IPAddresses
	uris, err := pki.URIsForCertificate(certificate)
	if err != nil {
		return nil, fmt.Errorf("failed to parse URIs: %s", err)
	}

	expectedURIs := pki.URLsToString(uris)

	certBytes, ok := secret.Data[corev1.TLSCertKey]
	if !ok {
		return nil, fmt.Errorf("No certificate data found for Certificate %q (secret %q)", certificate.Name, certificate.Spec.SecretName)
	}

	cert, err := pki.DecodeX509CertificateBytes(certBytes)
	if err != nil {
		return nil, err
	}

	commonNameCorrect := true
	expectedCN := certificate.Spec.CommonName
	if len(expectedCN) == 0 && len(cert.Subject.CommonName) > 0 {
		// issuers might set an IP or DNSName as CN
		if !util.Contains(cert.DNSNames, cert.Subject.CommonName) && !util.Contains(pki.IPAddressesToString(cert.IPAddresses), cert.Subject.CommonName) {
			commonNameCorrect = false
		}
	} else if expectedCN != cert.Subject.CommonName {
		commonNameCorrect = false
	}

	if !commonNameCorrect || !util.Subset(cert.DNSNames, expectedDNSNames) || !util.EqualUnsorted(pki.URLsToString(cert.URIs), expectedURIs) ||
		!util.Subset(pki.IPAddressesToString(cert.IPAddresses), expectedIPAddresses) ||
		!(len(cert.Subject.Organization) == 0 || util.EqualUnsorted(cert.Subject.Organization, expectedOrganization)) {
		return nil, fmt.Errorf("Expected certificate valid for CN %q, O %v, dnsNames %v, uriSANs %v,but got a certificate valid for CN %q, O %v, dnsNames %v, uriSANs %v, ipAddresses %v",
			expectedCN, expectedOrganization, expectedDNSNames, expectedURIs, cert.Subject.CommonName, cert.Subject.Organization, cert.DNSNames, cert.URIs, cert.IPAddresses)
	}

	if certificate.Status.NotAfter == nil {
		return nil, fmt.Errorf("No certificate expiration found for Certificate %q", certificate.Name)
	}
	if !cert.NotAfter.Equal(certificate.Status.NotAfter.Time) {
		return nil, fmt.Errorf("Expected certificate expiry date to be %v, but got %v", certificate.Status.NotAfter, cert.NotAfter)
	}

	label, ok := secret.Annotations[cmapi.CertificateNameKey]
	if !ok {
		return nil, fmt.Errorf("Expected secret to have certificate-name label, but had none")
	}

	if label != certificate.Name {
		return nil, fmt.Errorf("Expected secret to have certificate-name label with a value of %q, but got %q", certificate.Name, label)
	}

	certificateKeyUsages, certificateExtKeyUsages, err := pki.BuildKeyUsages(certificate.Spec.Usages, certificate.Spec.IsCA)
	if err != nil {
		return nil, fmt.Errorf("failed to build key usages from certificate: %s", err)
	}

	defaultCertKeyUsages, defaultCertExtKeyUsages, err := h.defaultKeyUsagesToAdd(certificate.Namespace, &certificate.Spec.IssuerRef)
	if err != nil {
		return nil, err
	}

	certificateKeyUsages |= defaultCertKeyUsages
	certificateExtKeyUsages = append(certificateExtKeyUsages, defaultCertExtKeyUsages...)

	// If using ECDSA then ignore key encipherment
	if certificate.Spec.PrivateKey != nil && certificate.Spec.PrivateKey.Algorithm == cmapi.ECDSAKeyAlgorithm {
		certificateKeyUsages &^= x509.KeyUsageKeyEncipherment
		cert.KeyUsage &^= x509.KeyUsageKeyEncipherment
	}

	certificateExtKeyUsages = h.deduplicateExtKeyUsages(certificateExtKeyUsages)

	if !h.keyUsagesMatch(cert.KeyUsage, cert.ExtKeyUsage,
		certificateKeyUsages, certificateExtKeyUsages) {
		return nil, fmt.Errorf("key usages and extended key usages do not match: exp=%s got=%s exp=%s got=%s",
			apiutil.KeyUsageStrings(certificateKeyUsages), apiutil.KeyUsageStrings(cert.KeyUsage),
			apiutil.ExtKeyUsageStrings(certificateExtKeyUsages), apiutil.ExtKeyUsageStrings(cert.ExtKeyUsage))
	}

	if !util.EqualUnsorted(cert.EmailAddresses, certificate.Spec.EmailAddresses) {
		return nil, fmt.Errorf("certificate doesn't contain Email SANs: exp=%v got=%v", certificate.Spec.EmailAddresses, cert.EmailAddresses)
	}

	var dnsName string
	if len(expectedDNSNames) > 0 {
		dnsName = expectedDNSNames[0]
	}

	// TODO: move this verification step out of this function
	if rootCAPEM != nil {
		rootCertPool := x509.NewCertPool()
		rootCertPool.AppendCertsFromPEM(rootCAPEM)
		intermediateCertPool := x509.NewCertPool()
		intermediateCertPool.AppendCertsFromPEM(certBytes)
		opts := x509.VerifyOptions{
			DNSName:       dnsName,
			Intermediates: intermediateCertPool,
			Roots:         rootCertPool,
		}

		if _, err := cert.Verify(opts); err != nil {
			return nil, err
		}
	}

	if cond := apiutil.GetCertificateCondition(certificate, cmapi.CertificateConditionReady); cond.Status != cmmeta.ConditionTrue ||
		cond.ObservedGeneration != certificate.Generation {
		return nil, fmt.Errorf(
			"expected Certificate to have ready condition true, with an observedGeneration matching the Certificate generation, got=%+v",
			cond)
	}

	return cert, nil
}

func (h *Helper) deduplicateExtKeyUsages(us []x509.ExtKeyUsage) []x509.ExtKeyUsage {
	extKeyUsagesMap := make(map[x509.ExtKeyUsage]bool)
	for _, e := range us {
		extKeyUsagesMap[e] = true
	}

	us = make([]x509.ExtKeyUsage, 0)
	for e, ok := range extKeyUsagesMap {
		if ok {
			us = append(us, e)
		}
	}

	return us
}

func (h *Helper) WaitCertificateIssued(ns, name string, timeout time.Duration) error {
	certificate, err := h.WaitForCertificateReady(ns, name, timeout)
	if err != nil {
		log.Logf("Error waiting for Certificate to become Ready: %v", err)
		h.Kubectl(ns).DescribeResource("certificate", name)
		h.Kubectl(ns).Describe("order", "challenge")
		h.describeCertificateRequestFromCertificate(ns, certificate)
	}
	return err
}

func (h *Helper) defaultKeyUsagesToAdd(ns string, issuerRef *cmmeta.ObjectReference) (x509.KeyUsage, []x509.ExtKeyUsage, error) {
	var issuerSpec *cmapi.IssuerSpec
	switch issuerRef.Kind {
	case "ClusterIssuer":
		issuerObj, err := h.CMClient.CertmanagerV1().ClusterIssuers().Get(context.TODO(), issuerRef.Name, metav1.GetOptions{})
		if err != nil {
			return 0, nil, fmt.Errorf("failed to find referenced ClusterIssuer %v: %s",
				issuerRef, err)
		}

		issuerSpec = &issuerObj.Spec
	default:
		issuerObj, err := h.CMClient.CertmanagerV1().Issuers(ns).Get(context.TODO(), issuerRef.Name, metav1.GetOptions{})
		if err != nil {
			return 0, nil, fmt.Errorf("failed to find referenced Issuer %v: %s",
				issuerRef, err)
		}

		issuerSpec = &issuerObj.Spec
	}

	var keyUsages x509.KeyUsage
	var extKeyUsages []x509.ExtKeyUsage

	// Vault and ACME issuers will add server auth and client auth extended key
	// usages by default so we need to add them to the list of expected usages
	if issuerSpec.ACME != nil || issuerSpec.Vault != nil {
		extKeyUsages = append(extKeyUsages, x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth)
	}

	// Vault issuers will add key agreement key usage
	if issuerSpec.Vault != nil {
		keyUsages |= x509.KeyUsageKeyAgreement
	}

	// Venafi issue adds server auth key usage
	if issuerSpec.Venafi != nil {
		extKeyUsages = append(extKeyUsages, x509.ExtKeyUsageServerAuth)
	}

	return keyUsages, extKeyUsages, nil
}

func (h *Helper) keyUsagesMatch(aKU x509.KeyUsage, aEKU []x509.ExtKeyUsage,
	bKU x509.KeyUsage, bEKU []x509.ExtKeyUsage) bool {
	if aKU != bKU {
		return false
	}

	if len(aEKU) != len(bEKU) {
		return false
	}

	sort.SliceStable(aEKU, func(i, j int) bool {
		return aEKU[i] < aEKU[j]
	})

	sort.SliceStable(bEKU, func(i, j int) bool {
		return bEKU[i] < bEKU[j]
	})

	for i := range aEKU {
		if aEKU[i] != bEKU[i] {
			return false
		}
	}

	return true
}

func (h *Helper) describeCertificateRequestFromCertificate(ns string, certificate *cmapi.Certificate) {
	if certificate == nil {
		return
	}

	crName, err := apiutil.ComputeName(certificate.Name, certificate.Spec)
	if err != nil {
		log.Logf("Failed to compute CertificateRequest name from certificate: %s", err)
		return
	}
	h.Kubectl(ns).DescribeResource("certificaterequest", crName)
}
<commit_msg>Reorder ready condition if statement to make more clear<commit_after>/*
Copyright 2020 The cert-manager Authors.

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

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

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

package helper

import (
	"context"
	"crypto/ecdsa"
	"crypto/rsa"
	"crypto/x509"
	"fmt"
	"sort"
	"time"

	corev1 "k8s.io/api/core/v1"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"

	apiutil "github.com/jetstack/cert-manager/pkg/api/util"
	cmapi "github.com/jetstack/cert-manager/pkg/apis/certmanager/v1"
	cmmeta "github.com/jetstack/cert-manager/pkg/apis/meta/v1"
	"github.com/jetstack/cert-manager/pkg/util"
	"github.com/jetstack/cert-manager/pkg/util/pki"
	"github.com/jetstack/cert-manager/test/e2e/framework/log"
)

// WaitForCertificateReady waits for the certificate resource to enter a Ready
// state.
func (h *Helper) WaitForCertificateReady(ns, name string, timeout time.Duration) (*cmapi.Certificate, error) {
	var certificate *cmapi.Certificate
	err := wait.PollImmediate(time.Second, timeout,
		func() (bool, error) {
			var err error
			log.Logf("Waiting for Certificate %v to be ready", name)
			certificate, err = h.CMClient.CertmanagerV1().Certificates(ns).Get(context.TODO(), name, metav1.GetOptions{})
			if err != nil {
				return false, fmt.Errorf("error getting Certificate %v: %v", name, err)
			}
			isReady := apiutil.CertificateHasCondition(certificate, cmapi.CertificateCondition{
				Type:   cmapi.CertificateConditionReady,
				Status: cmmeta.ConditionTrue,
			})
			if !isReady {
				log.Logf("Expected Certificate to have Ready condition 'true' but it has: %v", certificate.Status.Conditions)
				return false, nil
			}
			return true, nil
		},
	)

	// return certificate even when error to use for debugging
	return certificate, err
}

// WaitForCertificateNotReady waits for the certificate resource to enter a
// non-Ready state.
func (h *Helper) WaitForCertificateNotReady(ns, name string, timeout time.Duration) (*cmapi.Certificate, error) {
	var certificate *cmapi.Certificate
	err := wait.PollImmediate(time.Second, timeout,
		func() (bool, error) {
			var err error
			log.Logf("Waiting for Certificate %v to be ready", name)
			certificate, err = h.CMClient.CertmanagerV1().Certificates(ns).Get(context.TODO(), name, metav1.GetOptions{})
			if err != nil {
				return false, fmt.Errorf("error getting Certificate %v: %v", name, err)
			}
			isReady := apiutil.CertificateHasCondition(certificate, cmapi.CertificateCondition{
				Type:   cmapi.CertificateConditionReady,
				Status: cmmeta.ConditionFalse,
			})
			if !isReady {
				log.Logf("Expected Certificate to have Ready condition 'true' but it has: %v", certificate.Status.Conditions)
				return false, nil
			}
			return true, nil
		},
	)

	// return certificate even when error to use for debugging
	return certificate, err
}

// ValidateIssuedCertificate will ensure that the given Certificate has a
// certificate issued for it, and that the details on the x509 certificate are
// correct as defined by the Certificate's spec.
func (h *Helper) ValidateIssuedCertificate(certificate *cmapi.Certificate, rootCAPEM []byte) (*x509.Certificate, error) {
	log.Logf("Getting the TLS certificate Secret resource")
	secret, err := h.KubeClient.CoreV1().Secrets(certificate.Namespace).Get(context.TODO(), certificate.Spec.SecretName, metav1.GetOptions{})
	if err != nil {
		return nil, err
	}
	if !(len(secret.Data) == 2 || len(secret.Data) == 3) {
		return nil, fmt.Errorf("Expected 2 keys in certificate secret, but there was %d", len(secret.Data))
	}

	keyBytes, ok := secret.Data[corev1.TLSPrivateKeyKey]
	if !ok {
		return nil, fmt.Errorf("No private key data found for Certificate %q (secret %q)", certificate.Name, certificate.Spec.SecretName)
	}
	key, err := pki.DecodePrivateKeyBytes(keyBytes)
	if err != nil {
		return nil, err
	}

	// validate private key is of the correct type (rsa or ecdsa)
	privateKey := certificate.Spec.PrivateKey
	if privateKey == nil {
		privateKey = &cmapi.CertificatePrivateKey{}
	}
	switch privateKey.Algorithm {
	case cmapi.PrivateKeyAlgorithm(""),
		cmapi.RSAKeyAlgorithm:
		_, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("Expected private key of type RSA, but it was: %T", key)
		}
	case cmapi.ECDSAKeyAlgorithm:
		_, ok := key.(*ecdsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("Expected private key of type ECDSA, but it was: %T", key)
		}
	default:
		return nil, fmt.Errorf("unrecognised requested private key algorithm %q", certificate.Spec.PrivateKey.Algorithm)
	}

	// TODO: validate private key KeySize

	// check the provided certificate is valid
	expectedOrganization := pki.OrganizationForCertificate(certificate)
	expectedDNSNames := certificate.Spec.DNSNames
	expectedIPAddresses := certificate.Spec.IPAddresses
	uris, err := pki.URIsForCertificate(certificate)
	if err != nil {
		return nil, fmt.Errorf("failed to parse URIs: %s", err)
	}

	expectedURIs := pki.URLsToString(uris)

	certBytes, ok := secret.Data[corev1.TLSCertKey]
	if !ok {
		return nil, fmt.Errorf("No certificate data found for Certificate %q (secret %q)", certificate.Name, certificate.Spec.SecretName)
	}

	cert, err := pki.DecodeX509CertificateBytes(certBytes)
	if err != nil {
		return nil, err
	}

	commonNameCorrect := true
	expectedCN := certificate.Spec.CommonName
	if len(expectedCN) == 0 && len(cert.Subject.CommonName) > 0 {
		// issuers might set an IP or DNSName as CN
		if !util.Contains(cert.DNSNames, cert.Subject.CommonName) && !util.Contains(pki.IPAddressesToString(cert.IPAddresses), cert.Subject.CommonName) {
			commonNameCorrect = false
		}
	} else if expectedCN != cert.Subject.CommonName {
		commonNameCorrect = false
	}

	if !commonNameCorrect || !util.Subset(cert.DNSNames, expectedDNSNames) || !util.EqualUnsorted(pki.URLsToString(cert.URIs), expectedURIs) ||
		!util.Subset(pki.IPAddressesToString(cert.IPAddresses), expectedIPAddresses) ||
		!(len(cert.Subject.Organization) == 0 || util.EqualUnsorted(cert.Subject.Organization, expectedOrganization)) {
		return nil, fmt.Errorf("Expected certificate valid for CN %q, O %v, dnsNames %v, uriSANs %v,but got a certificate valid for CN %q, O %v, dnsNames %v, uriSANs %v, ipAddresses %v",
			expectedCN, expectedOrganization, expectedDNSNames, expectedURIs, cert.Subject.CommonName, cert.Subject.Organization, cert.DNSNames, cert.URIs, cert.IPAddresses)
	}

	if certificate.Status.NotAfter == nil {
		return nil, fmt.Errorf("No certificate expiration found for Certificate %q", certificate.Name)
	}
	if !cert.NotAfter.Equal(certificate.Status.NotAfter.Time) {
		return nil, fmt.Errorf("Expected certificate expiry date to be %v, but got %v", certificate.Status.NotAfter, cert.NotAfter)
	}

	label, ok := secret.Annotations[cmapi.CertificateNameKey]
	if !ok {
		return nil, fmt.Errorf("Expected secret to have certificate-name label, but had none")
	}

	if label != certificate.Name {
		return nil, fmt.Errorf("Expected secret to have certificate-name label with a value of %q, but got %q", certificate.Name, label)
	}

	certificateKeyUsages, certificateExtKeyUsages, err := pki.BuildKeyUsages(certificate.Spec.Usages, certificate.Spec.IsCA)
	if err != nil {
		return nil, fmt.Errorf("failed to build key usages from certificate: %s", err)
	}

	defaultCertKeyUsages, defaultCertExtKeyUsages, err := h.defaultKeyUsagesToAdd(certificate.Namespace, &certificate.Spec.IssuerRef)
	if err != nil {
		return nil, err
	}

	certificateKeyUsages |= defaultCertKeyUsages
	certificateExtKeyUsages = append(certificateExtKeyUsages, defaultCertExtKeyUsages...)

	// If using ECDSA then ignore key encipherment
	if certificate.Spec.PrivateKey != nil && certificate.Spec.PrivateKey.Algorithm == cmapi.ECDSAKeyAlgorithm {
		certificateKeyUsages &^= x509.KeyUsageKeyEncipherment
		cert.KeyUsage &^= x509.KeyUsageKeyEncipherment
	}

	certificateExtKeyUsages = h.deduplicateExtKeyUsages(certificateExtKeyUsages)

	if !h.keyUsagesMatch(cert.KeyUsage, cert.ExtKeyUsage,
		certificateKeyUsages, certificateExtKeyUsages) {
		return nil, fmt.Errorf("key usages and extended key usages do not match: exp=%s got=%s exp=%s got=%s",
			apiutil.KeyUsageStrings(certificateKeyUsages), apiutil.KeyUsageStrings(cert.KeyUsage),
			apiutil.ExtKeyUsageStrings(certificateExtKeyUsages), apiutil.ExtKeyUsageStrings(cert.ExtKeyUsage))
	}

	if !util.EqualUnsorted(cert.EmailAddresses, certificate.Spec.EmailAddresses) {
		return nil, fmt.Errorf("certificate doesn't contain Email SANs: exp=%v got=%v", certificate.Spec.EmailAddresses, cert.EmailAddresses)
	}

	var dnsName string
	if len(expectedDNSNames) > 0 {
		dnsName = expectedDNSNames[0]
	}

	// TODO: move this verification step out of this function
	if rootCAPEM != nil {
		rootCertPool := x509.NewCertPool()
		rootCertPool.AppendCertsFromPEM(rootCAPEM)
		intermediateCertPool := x509.NewCertPool()
		intermediateCertPool.AppendCertsFromPEM(certBytes)
		opts := x509.VerifyOptions{
			DNSName:       dnsName,
			Intermediates: intermediateCertPool,
			Roots:         rootCertPool,
		}

		if _, err := cert.Verify(opts); err != nil {
			return nil, err
		}
	}

	readyCond := apiutil.GetCertificateCondition(certificate, cmapi.CertificateConditionReady)
	if readyCond.Status != cmmeta.ConditionTrue || readyCond.ObservedGeneration != certificate.Generation {
		return nil, fmt.Errorf(
			"expected Certificate to have ready condition true, with an observedGeneration matching the Certificate generation, got=%+v",
			readyCond)
	}

	return cert, nil
}

func (h *Helper) deduplicateExtKeyUsages(us []x509.ExtKeyUsage) []x509.ExtKeyUsage {
	extKeyUsagesMap := make(map[x509.ExtKeyUsage]bool)
	for _, e := range us {
		extKeyUsagesMap[e] = true
	}

	us = make([]x509.ExtKeyUsage, 0)
	for e, ok := range extKeyUsagesMap {
		if ok {
			us = append(us, e)
		}
	}

	return us
}

func (h *Helper) WaitCertificateIssued(ns, name string, timeout time.Duration) error {
	certificate, err := h.WaitForCertificateReady(ns, name, timeout)
	if err != nil {
		log.Logf("Error waiting for Certificate to become Ready: %v", err)
		h.Kubectl(ns).DescribeResource("certificate", name)
		h.Kubectl(ns).Describe("order", "challenge")
		h.describeCertificateRequestFromCertificate(ns, certificate)
	}
	return err
}

func (h *Helper) defaultKeyUsagesToAdd(ns string, issuerRef *cmmeta.ObjectReference) (x509.KeyUsage, []x509.ExtKeyUsage, error) {
	var issuerSpec *cmapi.IssuerSpec
	switch issuerRef.Kind {
	case "ClusterIssuer":
		issuerObj, err := h.CMClient.CertmanagerV1().ClusterIssuers().Get(context.TODO(), issuerRef.Name, metav1.GetOptions{})
		if err != nil {
			return 0, nil, fmt.Errorf("failed to find referenced ClusterIssuer %v: %s",
				issuerRef, err)
		}

		issuerSpec = &issuerObj.Spec
	default:
		issuerObj, err := h.CMClient.CertmanagerV1().Issuers(ns).Get(context.TODO(), issuerRef.Name, metav1.GetOptions{})
		if err != nil {
			return 0, nil, fmt.Errorf("failed to find referenced Issuer %v: %s",
				issuerRef, err)
		}

		issuerSpec = &issuerObj.Spec
	}

	var keyUsages x509.KeyUsage
	var extKeyUsages []x509.ExtKeyUsage

	// Vault and ACME issuers will add server auth and client auth extended key
	// usages by default so we need to add them to the list of expected usages
	if issuerSpec.ACME != nil || issuerSpec.Vault != nil {
		extKeyUsages = append(extKeyUsages, x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth)
	}

	// Vault issuers will add key agreement key usage
	if issuerSpec.Vault != nil {
		keyUsages |= x509.KeyUsageKeyAgreement
	}

	// Venafi issue adds server auth key usage
	if issuerSpec.Venafi != nil {
		extKeyUsages = append(extKeyUsages, x509.ExtKeyUsageServerAuth)
	}

	return keyUsages, extKeyUsages, nil
}

func (h *Helper) keyUsagesMatch(aKU x509.KeyUsage, aEKU []x509.ExtKeyUsage,
	bKU x509.KeyUsage, bEKU []x509.ExtKeyUsage) bool {
	if aKU != bKU {
		return false
	}

	if len(aEKU) != len(bEKU) {
		return false
	}

	sort.SliceStable(aEKU, func(i, j int) bool {
		return aEKU[i] < aEKU[j]
	})

	sort.SliceStable(bEKU, func(i, j int) bool {
		return bEKU[i] < bEKU[j]
	})

	for i := range aEKU {
		if aEKU[i] != bEKU[i] {
			return false
		}
	}

	return true
}

func (h *Helper) describeCertificateRequestFromCertificate(ns string, certificate *cmapi.Certificate) {
	if certificate == nil {
		return
	}

	crName, err := apiutil.ComputeName(certificate.Name, certificate.Spec)
	if err != nil {
		log.Logf("Failed to compute CertificateRequest name from certificate: %s", err)
		return
	}
	h.Kubectl(ns).DescribeResource("certificaterequest", crName)
}
<|endoftext|>


<commit_before>package main

/*
* Read data from InfluxDB, and render it
*/

import (
	"encoding/json"
	"os"
	"net/http"
	"net/http/httptest"
	"net/url"
	//"strings"
	"time"
	"github.com/Sirupsen/logrus"
	"github.com/influxdata/influxdb/client"

)

const db = "canal-monitoring"

var log = logrus.New()
var command = "select * from ep_stats"

func main() {
	log.Out = os.Stdout
	log.Info("Test New Client to InfluxDB")
	TestNewClient()
	TestClient_Ping()
	TestClient_Query(db, command)
}

func TestNewClient() {
	config := client.Config{}
	_, err := client.NewClient(config)
	if err != nil {
		log.Fatalf("unexpected error.  expected %v, actual %v", nil, err)
	} else {
		log.Info("Connect to InfluxDB successfully")
	}
}

// helper functions
func emptyTestServer() *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(50 * time.Millisecond)
		w.Header().Set("X-Influxdb-Version", "1.2.2")
		return
	}))
}

func TestClient_Ping() {
	log.Info("TestClient_Ping")
	ts := emptyTestServer()
	defer ts.Close()

	u, _ := url.Parse(ts.URL)

	u.Host = "localhost:8086"  // set to InfluxDB default port : 8086

	config := client.Config{URL: *u}

	log.Info("u =", u)

	c, err := client.NewClient(config)
	if err != nil {
		log.Fatalf("unexpected error.  expected %v, actual %v", nil, err)
	}
	d, version, err := c.Ping()
	log.Info("version: ", version)
	log.Info("d: ", d)
	if err != nil {
		log.Fatalf("unexpected error.  expected %v, actual %v", nil, err)
	}
	if d.Nanoseconds() == 0 {
		log.Fatalf("expected a duration greater than zero.  actual %v", d.Nanoseconds())
	}
	if version != "1.2.2" {
		log.Fatalf("unexpected version.  expected %s,  actual %v", "1.2.2", version)
	}
}

func TestClient_Query(db string, command string) {
	log.Info("TestClient_Query")
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var data client.Response
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(data)
	}))
	defer ts.Close()

	u, _ := url.Parse(ts.URL)

	u.Host = "localhost:8086"  // set to InfluxDB default port : 8086

	config := client.Config{URL: *u}
	c, err := client.NewClient(config)
	if err != nil {
		log.Fatalf("unexpected error.  expected %v, actual %v", nil, err)
	}

	query := client.Query{}
	query.Database = db
	query.Command = command

	res, err := c.Query(query)
	log.Info("response: ", res)

	if err != nil {
		log.Fatalf("unexpected error.  expected %v, actual %v", nil, err)
	}
}
<commit_msg>update the influxdb measurment, db<commit_after>package main

/*
* Read data from InfluxDB, and render it
*/

import (
	"encoding/json"
	"os"
	"net/http"
	"net/http/httptest"
	"net/url"
	//"strings"
	"time"
	"github.com/Sirupsen/logrus"
	"github.com/influxdata/influxdb/client"

)

const db = "ican"

var log = logrus.New()
var command = "select * from p2prxbyte"

func main() {
	log.Out = os.Stdout
	log.Info("Test New Client to InfluxDB")
	TestNewClient()
	TestClient_Ping()
	TestClient_Query(db, command)
}

func TestNewClient() {
	config := client.Config{}
	_, err := client.NewClient(config)
	if err != nil {
		log.Fatalf("unexpected error.  expected %v, actual %v", nil, err)
	} else {
		log.Info("Connect to InfluxDB successfully")
	}
}

// helper functions
func emptyTestServer() *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(50 * time.Millisecond)
		w.Header().Set("X-Influxdb-Version", "1.2.2")
		return
	}))
}

func TestClient_Ping() {
	log.Info("TestClient_Ping")
	ts := emptyTestServer()
	defer ts.Close()

	u, _ := url.Parse(ts.URL)

	u.Host = "localhost:8086"  // set to InfluxDB default port : 8086

	config := client.Config{URL: *u}

	log.Info("u =", u)

	c, err := client.NewClient(config)
	if err != nil {
		log.Fatalf("unexpected error.  expected %v, actual %v", nil, err)
	}
	d, version, err := c.Ping()
	log.Info("version: ", version)
	log.Info("d: ", d)
	if err != nil {
		log.Fatalf("unexpected error.  expected %v, actual %v", nil, err)
	}
	if d.Nanoseconds() == 0 {
		log.Fatalf("expected a duration greater than zero.  actual %v", d.Nanoseconds())
	}
	if version != "1.2.2" {
		log.Fatalf("unexpected version.  expected %s,  actual %v", "1.2.2", version)
	}
}

func TestClient_Query(db string, command string) {
	log.Info("TestClient_Query")
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var data client.Response
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(data)
	}))
	defer ts.Close()

	u, _ := url.Parse(ts.URL)

	u.Host = "localhost:8086"  // set to InfluxDB default port : 8086

	config := client.Config{URL: *u}
	c, err := client.NewClient(config)
	if err != nil {
		log.Fatalf("unexpected error.  expected %v, actual %v", nil, err)
	}

	query := client.Query{}
	query.Database = db
	query.Command = command

	res, err := c.Query(query)
	log.Info("response: ", res)

	if err != nil {
		log.Fatalf("unexpected error.  expected %v, actual %v", nil, err)
	}
}
<|endoftext|>


<commit_before>package goshare

import (
	"fmt"
	"net/http"
	"runtime"
	"strings"
	"time"

	golhashmap "github.com/abhishekkr/gol/golhashmap"
	"github.com/abhishekkr/gol/goltime"
	"github.com/abhishekkr/goshare/httpd"
)

/* enabling HTTP to formulate DBTasks call from HTTP Request */
func DBRest(httpMethod string, w http.ResponseWriter, req *http.Request) {
	var (
		dbAction       string
		response_bytes []byte
		axn_status     bool
	)

	switch httpMethod {
	case "GET":
		dbAction = "read"

	case "POST", "PUT":
		dbAction = "push"

	case "DELETE":
		dbAction = "delete"

	default:
		// log_this corrupt request
		return
	}

	packet := PacketFromHTTPRequest(dbAction, req)
	if packet.DBAction != "" {
		response_bytes, axn_status = DBTasksOnPacket(packet)
		DBRestResponse(w, req, response_bytes, axn_status)
	}
}

/* send proper response back to client based on success/data/error */
func DBRestResponse(w http.ResponseWriter, req *http.Request, response_bytes []byte, axn_status bool) {
	if !axn_status {
		error_msg := fmt.Sprintf("FATAL Error: (DBTasks) %q", req.Form)
		http.Error(w, error_msg, http.StatusInternalServerError)

	} else if len(response_bytes) == 0 {
		w.Write([]byte("Success"))

	} else {
		w.Write(response_bytes)

	}
}

/*
return Packet identifiable by DBTasksOnAction
*/
func PacketFromHTTPRequest(dbAction string, req *http.Request) Packet {
	packet := Packet{}
	packet.HashMap = make(golhashmap.HashMap)
	packet.DBAction = dbAction

	req.ParseForm()
	task_type := req.FormValue("type")
	if task_type == "" {
		task_type = "default"
	}
	packet.TaskType = task_type
	task_type_tokens := strings.Split(task_type, "-")
	packet.KeyType = task_type_tokens[0]
	if packet.KeyType == "tsds" {
		packet.TimeDot = goltime.TimestampFromHTTPRequest(req)
	}

	if len(task_type_tokens) > 1 {
		packet.ValType = task_type_tokens[1]

		if len(task_type_tokens) == 3 {
			thirdTokenFeatureHTTP(&packet, req)
		}
	}

	dbdata := req.FormValue("dbdata")
	key := req.FormValue("key")
	if key != "" {
		dbdata = fmt.Sprintf("%s\n%s", key, req.FormValue("val"))
	} else if dbdata == "" {
		return Packet{}
	}
	decodeData(&packet, strings.Split(dbdata, "\n"))

	return packet
}

/* third token handler in taskType */
func thirdTokenFeatureHTTP(packet *Packet, req *http.Request) {
	parentNS := req.FormValue("parentNS")
	if parentNS != "" {
		packet.ParentNS = parentNS
	}
}

/* DB Call HTTP Handler */
func DBRestHandler(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	req.ParseForm()

	DBRest(req.Method, w, req)
}

/* HTTP GET DB-GET call handler */
func GetReadKey(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	req.ParseForm()

	DBRest("GET", w, req)
}

/* HTTP GET DB-POST call handler */
func GetPushKey(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	req.ParseForm()

	DBRest("POST", w, req)
}

/* HTTP GET DB-POST call handler */
func GetDeleteKey(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	req.ParseForm()

	DBRest("DELETE", w, req)
}

/*
GoShare Handler for HTTP Requests
*/
func GoShareHTTP(httpuri string, httpport int) {
	runtime.GOMAXPROCS(runtime.NumCPU())

	http.HandleFunc("/", abkhttpd.F1)
	http.HandleFunc("/help-http", abkhttpd.HelpHTTP)
	http.HandleFunc("/help-zmq", abkhttpd.HelpZMQ)
	http.HandleFunc("/status", abkhttpd.Status)

	http.HandleFunc("/db", DBRestHandler)
	http.HandleFunc("/get", GetReadKey)
	http.HandleFunc("/put", GetPushKey)
	http.HandleFunc("/del", GetDeleteKey)

	srv := &http.Server{
		Addr:        fmt.Sprintf("%s:%d", httpuri, httpport),
		Handler:     http.DefaultServeMux,
		ReadTimeout: time.Duration(5) * time.Second,
	}

	fmt.Printf("access your goshare at http://%s:%d\n", httpuri, httpport)
	err := srv.ListenAndServe()
	fmt.Println("Game Over:", err)
}
<commit_msg>gohttp updated for tsds TimDot creation only in push<commit_after>package goshare

import (
	"fmt"
	"net/http"
	"runtime"
	"strings"
	"time"

	golhashmap "github.com/abhishekkr/gol/golhashmap"
	"github.com/abhishekkr/gol/goltime"
	"github.com/abhishekkr/goshare/httpd"
)

/* enabling HTTP to formulate DBTasks call from HTTP Request */
func DBRest(httpMethod string, w http.ResponseWriter, req *http.Request) {
	var (
		dbAction       string
		response_bytes []byte
		axn_status     bool
	)

	switch httpMethod {
	case "GET":
		dbAction = "read"

	case "POST", "PUT":
		dbAction = "push"

	case "DELETE":
		dbAction = "delete"

	default:
		// log_this corrupt request
		return
	}

	packet := PacketFromHTTPRequest(dbAction, req)
	if packet.DBAction != "" {
		response_bytes, axn_status = DBTasksOnPacket(packet)
		DBRestResponse(w, req, response_bytes, axn_status)
	}
}

/* send proper response back to client based on success/data/error */
func DBRestResponse(w http.ResponseWriter, req *http.Request, response_bytes []byte, axn_status bool) {
	if !axn_status {
		error_msg := fmt.Sprintf("FATAL Error: (DBTasks) %q", req.Form)
		http.Error(w, error_msg, http.StatusInternalServerError)

	} else if len(response_bytes) == 0 {
		w.Write([]byte("Success"))

	} else {
		w.Write(response_bytes)

	}
}

/*
return Packet identifiable by DBTasksOnAction
*/
func PacketFromHTTPRequest(dbAction string, req *http.Request) Packet {
	packet := Packet{}
	packet.HashMap = make(golhashmap.HashMap)
	packet.DBAction = dbAction

	req.ParseForm()
	task_type := req.FormValue("type")
	if task_type == "" {
		task_type = "default"
	}
	packet.TaskType = task_type
	task_type_tokens := strings.Split(task_type, "-")
	packet.KeyType = task_type_tokens[0]
	if packet.KeyType == "tsds" && packet.DBAction == "push" {
		packet.TimeDot = goltime.TimestampFromHTTPRequest(req)
	}

	if len(task_type_tokens) > 1 {
		packet.ValType = task_type_tokens[1]

		if len(task_type_tokens) == 3 {
			thirdTokenFeatureHTTP(&packet, req)
		}
	}

	dbdata := req.FormValue("dbdata")
	key := req.FormValue("key")
	if key != "" {
		dbdata = fmt.Sprintf("%s\n%s", key, req.FormValue("val"))
	} else if dbdata == "" {
		return Packet{}
	}
	decodeData(&packet, strings.Split(dbdata, "\n"))

	return packet
}

/* third token handler in taskType */
func thirdTokenFeatureHTTP(packet *Packet, req *http.Request) {
	parentNS := req.FormValue("parentNS")
	if parentNS != "" {
		packet.ParentNS = parentNS
	}
}

/* DB Call HTTP Handler */
func DBRestHandler(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	req.ParseForm()

	DBRest(req.Method, w, req)
}

/* HTTP GET DB-GET call handler */
func GetReadKey(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	req.ParseForm()

	DBRest("GET", w, req)
}

/* HTTP GET DB-POST call handler */
func GetPushKey(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	req.ParseForm()

	DBRest("POST", w, req)
}

/* HTTP GET DB-POST call handler */
func GetDeleteKey(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	req.ParseForm()

	DBRest("DELETE", w, req)
}

/*
GoShare Handler for HTTP Requests
*/
func GoShareHTTP(httpuri string, httpport int) {
	runtime.GOMAXPROCS(runtime.NumCPU())

	http.HandleFunc("/", abkhttpd.F1)
	http.HandleFunc("/help-http", abkhttpd.HelpHTTP)
	http.HandleFunc("/help-zmq", abkhttpd.HelpZMQ)
	http.HandleFunc("/status", abkhttpd.Status)

	http.HandleFunc("/db", DBRestHandler)
	http.HandleFunc("/get", GetReadKey)
	http.HandleFunc("/put", GetPushKey)
	http.HandleFunc("/del", GetDeleteKey)

	srv := &http.Server{
		Addr:        fmt.Sprintf("%s:%d", httpuri, httpport),
		Handler:     http.DefaultServeMux,
		ReadTimeout: time.Duration(5) * time.Second,
	}

	fmt.Printf("access your goshare at http://%s:%d\n", httpuri, httpport)
	err := srv.ListenAndServe()
	fmt.Println("Game Over:", err)
}
<|endoftext|>


<commit_before>// Package goldie provides test assertions based on golden files. It's typically
// used for testing responses with larger data bodies.
//
// The concept is straight forward. Valid response data is stored in a "golden
// file". The actual response data will be byte compared with the golden file
// and the test will fail if there is a difference.
//
// Updating the golden file can be done by running `go test -update ./...`.
package goldie

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
	"text/template"
)

var (
	// FixtureDir is the folder name for where the fixtures are stored. It's
	// relative to the "go test" path.
	FixtureDir = "fixtures"

	// FileNameSuffix is the suffix appended to the fixtures. Set to empty
	// string to disable file name suffixes.
	FileNameSuffix = ".golden"

	// FlagName is the name of the command line flag for go test.
	FlagName = "update"

	// FilePerms is used to set the permissions on the golden fixture files.
	FilePerms os.FileMode = 0644

	// DirPerms is used to set the permissions on the golden fixture folder.
	DirPerms os.FileMode = 0755

	// update determines if the actual received data should be written to the
	// golden files or not. This should be true when you need to update the
	// data, but false when actually running the tests.
	update = flag.Bool(FlagName, false, "Update golden test file fixture")
)

// Assert compares the actual data received with the expected data in the
// golden files. If the update flag is set, it will also update the golden
// file.
//
// `name` refers to the name of the test and it should typically be unique
// within the package. Also it should be a valid file name (so keeping to
// `a-z0-9\-\_` is a good idea).
func Assert(t *testing.T, name string, actualData []byte) {
	AssertWithTemplate(t, name, nil, actualData)
}

// Assert compares the actual data received with the expected data in the
// golden files after executing it as a template with data parameter.
// If the update flag is set, it will also update the golden file.
// `name` refers to the name of the test and it should typically be unique
// within the package. Also it should be a valid file name (so keeping to
// `a-z0-9\-\_` is a good idea).
func AssertWithTemplate(t *testing.T, name string, data interface{}, actualData []byte) {
	if *update {
		err := Update(name, actualData)
		if err != nil {
			t.Error(err)
			t.FailNow()
		}
	}

	err := compare(name, data, actualData)
	if err != nil {
		switch err.(type) {
		case errFixtureNotFound:
			t.Error(err)
			t.FailNow()
		case errFixtureMismatch:
			t.Error(err)
		default:
			t.Error(err)
		}
	}
}

// Update will update the golden fixtures with the received actual data.
//
// This method does not need to be called from code, but it's exposed so that it
// can be explicitly called if needed. The more common approach would be to
// update using `go test -update ./...`.
func Update(name string, actualData []byte) error {
	if err := ensureDir(filepath.Dir(goldenFileName(name))); err != nil {
		return err
	}

	return ioutil.WriteFile(goldenFileName(name), actualData, FilePerms)
}

// compare is reading the golden fixture file and compare the stored data with
// the actual data.
func compare(name string, data interface{}, actualData []byte) error {
	expectedDataTmpl, err := ioutil.ReadFile(goldenFileName(name))

	if err != nil {
		if os.IsNotExist(err) {
			return newErrFixtureNotFound()
		}

		return fmt.Errorf("Expected %s to be nil", err.Error())
	}

	tmpl, err := template.New("test").Parse(string(expectedDataTmpl))
	if err != nil {
		return fmt.Errorf("Expected %s to be nil", err.Error())
	}

	var expectedData bytes.Buffer
	err = tmpl.Execute(&expectedData, data)
	if err != nil {
		return fmt.Errorf("Expected %s to be nil", err.Error())
	}

	if !bytes.Equal(actualData, expectedData.Bytes()) {
		return newErrFixtureMismatch(
			fmt.Sprintf("Result did not match the golden fixture.\n"+
				"Expected: %s\n"+
				"Got: %s",
				string(expectedData.Bytes()),
				string(actualData)))
	}

	return nil
}

// ensureDir will create the fixture folder if it does not already exist.
func ensureDir(loc string) error {
	s, err := os.Stat(loc)
	switch {
	case err != nil && os.IsNotExist(err):
		// the location does not exist, so make directories to there
		return os.MkdirAll(loc, DirPerms)
	case err == nil && !s.IsDir():
		return newErrFixtureDirectoryIsFile(loc)
	}

	return err
}

// goldenFileName simply returns the file name of the golden file fixture.
func goldenFileName(name string) string {
	return filepath.Join(FixtureDir, fmt.Sprintf("%s%s", name, FileNameSuffix))
}
<commit_msg>Assert wrapper for JSON types<commit_after>// Package goldie provides test assertions based on golden files. It's typically
// used for testing responses with larger data bodies.
//
// The concept is straight forward. Valid response data is stored in a "golden
// file". The actual response data will be byte compared with the golden file
// and the test will fail if there is a difference.
//
// Updating the golden file can be done by running `go test -update ./...`.
package goldie

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
	"text/template"
)

var (
	// FixtureDir is the folder name for where the fixtures are stored. It's
	// relative to the "go test" path.
	FixtureDir = "fixtures"

	// FileNameSuffix is the suffix appended to the fixtures. Set to empty
	// string to disable file name suffixes.
	FileNameSuffix = ".golden"

	// FlagName is the name of the command line flag for go test.
	FlagName = "update"

	// FilePerms is used to set the permissions on the golden fixture files.
	FilePerms os.FileMode = 0644

	// DirPerms is used to set the permissions on the golden fixture folder.
	DirPerms os.FileMode = 0755

	// update determines if the actual received data should be written to the
	// golden files or not. This should be true when you need to update the
	// data, but false when actually running the tests.
	update = flag.Bool(FlagName, false, "Update golden test file fixture")
)

// Assert compares the actual data received with the expected data in the
// golden files. If the update flag is set, it will also update the golden
// file.
//
// `name` refers to the name of the test and it should typically be unique
// within the package. Also it should be a valid file name (so keeping to
// `a-z0-9\-\_` is a good idea).
func Assert(t *testing.T, name string, actualData []byte) {
	AssertWithTemplate(t, name, nil, actualData)
}

// AssertJson compares the actual json data received with expected data in the
// golden files. If the update flag is set, it will also update the golden
// file.
//
// `name` refers to the name of the test and it should typically be unique
// within the package. Also it should be a valid file name (so keeping to
// `a-z0-9\-\_` is a good idea).
func AssertJson(t *testing.T, name string, actualJsonData interface{}) {
	js, err := json.MarshalIndent(actualJsonData, "", "  ")

	if err != nil {
		t.Error(err)
		t.FailNow()
	}

	Assert(t, name, normalizeLF(js))
}

// normalizeLF normalizes line feed character set across os (es)
// \r\n (windows) & \r (mac) into \n (unix)
func normalizeLF(d []byte) []byte {
	// replace CR LF \r\n (windows) with LF \n (unix)
	d = bytes.Replace(d, []byte{13, 10}, []byte{10}, -1)
	// replace CF \r (mac) with LF \n (unix)
	d = bytes.Replace(d, []byte{13}, []byte{10}, -1)
	return d
}

// Assert compares the actual data received with the expected data in the
// golden files after executing it as a template with data parameter.
// If the update flag is set, it will also update the golden file.
// `name` refers to the name of the test and it should typically be unique
// within the package. Also it should be a valid file name (so keeping to
// `a-z0-9\-\_` is a good idea).
func AssertWithTemplate(t *testing.T, name string, data interface{}, actualData []byte) {
	if *update {
		err := Update(name, actualData)
		if err != nil {
			t.Error(err)
			t.FailNow()
		}
	}

	err := compare(name, data, actualData)
	if err != nil {
		switch err.(type) {
		case errFixtureNotFound:
			t.Error(err)
			t.FailNow()
		case errFixtureMismatch:
			t.Error(err)
		default:
			t.Error(err)
		}
	}
}

// Update will update the golden fixtures with the received actual data.
//
// This method does not need to be called from code, but it's exposed so that it
// can be explicitly called if needed. The more common approach would be to
// update using `go test -update ./...`.
func Update(name string, actualData []byte) error {
	if err := ensureDir(filepath.Dir(goldenFileName(name))); err != nil {
		return err
	}

	return ioutil.WriteFile(goldenFileName(name), actualData, FilePerms)
}

// compare is reading the golden fixture file and compare the stored data with
// the actual data.
func compare(name string, data interface{}, actualData []byte) error {
	expectedDataTmpl, err := ioutil.ReadFile(goldenFileName(name))

	if err != nil {
		if os.IsNotExist(err) {
			return newErrFixtureNotFound()
		}

		return fmt.Errorf("Expected %s to be nil", err.Error())
	}

	tmpl, err := template.New("test").Parse(string(expectedDataTmpl))
	if err != nil {
		return fmt.Errorf("Expected %s to be nil", err.Error())
	}

	var expectedData bytes.Buffer
	err = tmpl.Execute(&expectedData, data)
	if err != nil {
		return fmt.Errorf("Expected %s to be nil", err.Error())
	}

	if !bytes.Equal(actualData, expectedData.Bytes()) {
		return newErrFixtureMismatch(
			fmt.Sprintf("Result did not match the golden fixture.\n"+
				"Expected: %s\n"+
				"Got: %s",
				string(expectedData.Bytes()),
				string(actualData)))
	}

	return nil
}

// ensureDir will create the fixture folder if it does not already exist.
func ensureDir(loc string) error {
	s, err := os.Stat(loc)
	switch {
	case err != nil && os.IsNotExist(err):
		// the location does not exist, so make directories to there
		return os.MkdirAll(loc, DirPerms)
	case err == nil && !s.IsDir():
		return newErrFixtureDirectoryIsFile(loc)
	}

	return err
}

// goldenFileName simply returns the file name of the golden file fixture.
func goldenFileName(name string) string {
	return filepath.Join(FixtureDir, fmt.Sprintf("%s%s", name, FileNameSuffix))
}
<|endoftext|>


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

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

package gosnmp

import (
	"fmt"
	"math/big"
	"math/rand"
	"net"
	"strconv"
	"time"
)

const (
	// maxOids is the maximum number of oids allowed in a Get()
	maxOids = 60

	// Base OID for MIB-2 defined SNMP variables
	baseOid = ".1.3.6.1.2.1"

	// Java SNMP uses 50, snmp-net uses 10
	defaultMaxRepetitions = 50
)

// LoggingDisabled is set if the Logger is nil, short circuits any 'slog' calls
var LoggingDisabled bool

// GoSNMP represents GoSNMP library state
type GoSNMP struct {

	// Target is an ipv4 address
	Target string

	// Port is a udp port
	Port uint16

	// Version is an SNMP Version
	Version SnmpVersion

	// Community is an SNMP Community string
	Community string

	// MsgFlags is an SNMPV3 MsgFlags
	MsgFlags SnmpV3MsgFlags

	// SecurityModel is an SNMPV3 Security Model
	SecurityModel SnmpV3SecurityModel

	// SecurityParameters is an SNMPV3 User Security Model paramaters struct
	SecurityParameters *UsmSecurityParameters

	// Timeout is the timeout for the SNMP Query
	Timeout time.Duration

	// Set the number of retries to attempt within timeout.
	Retries int

	// Conn is net connection to use, typically establised using GoSNMP.Connect()
	Conn net.Conn

	// Logger is the GoSNMP.Logger to use for debugging. If nil, debugging
	// output will be discarded (/dev/null). For verbose logging to stdout:
	// x.Logger = log.New(os.Stdout, "", 0)
	Logger Logger

	// MaxRepititions sets the GETBULK max-repetitions used by BulkWalk*
	// (default: 50)
	MaxRepetitions int

	// NonRepeaters sets the GETBULK max-repeaters used by BulkWalk*
	// (default: 0 as per RFC 1905)
	NonRepeaters int

	// Internal - used to sync requests to responses
	requestID uint32

	// Internal - used to sync requests to responses - snmpv3
	msgID uint32

	random *rand.Rand
}

// The default connection settings
var Default = &GoSNMP{
	Port:      161,
	Community: "public",
	Version:   Version2c,
	Timeout:   time.Duration(2) * time.Second,
	Retries:   3,
}

// SnmpPDU will be used when doing SNMP Set's
type SnmpPDU struct {

	// Name is an oid in string format eg ".1.3.6.1.4.9.27"
	Name string

	// The type of the value eg Integer
	Type Asn1BER

	// The value to be set by the SNMP set
	Value interface{}
}

// Asn1BER is the type of the SNMP PDU
type Asn1BER byte

// Asn1BER's - http://www.ietf.org/rfc/rfc1442.txt
const (
	EndOfContents     Asn1BER = 0x00
	UnknownType               = 0x00 // TODO these should all be type Asn1BER, however
	Boolean                   = 0x01 // tests fail if implemented. See for example
	Integer                   = 0x02 /// http://stackoverflow.com/questions/5037610/typed-constant-declaration-list.
	BitString                 = 0x03
	OctetString               = 0x04
	Null                      = 0x05
	ObjectIdentifier          = 0x06
	ObjectDescription         = 0x07
	IPAddress                 = 0x40
	Counter32                 = 0x41
	Gauge32                   = 0x42
	TimeTicks                 = 0x43
	Opaque                    = 0x44
	NsapAddress               = 0x45
	Counter64                 = 0x46
	Uinteger32                = 0x47
	NoSuchObject              = 0x80
	NoSuchInstance            = 0x81
	EndOfMibView              = 0x82
)

//
// Public Functions (main interface)
//

// Connect initiates a connection to the target host
func (x *GoSNMP) Connect() error {
	if x.Logger == nil {
		LoggingDisabled = true
	}
	Conn, err := net.DialTimeout("udp", fmt.Sprintf("%s:%d", x.Target, x.Port), x.Timeout)
	if err == nil {
		x.Conn = Conn
	} else {
		return fmt.Errorf("Error establishing connection to host: %s\n", err.Error())
	}
	if x.random == nil {
		x.random = rand.New(rand.NewSource(time.Now().UTC().UnixNano()))
	}
	x.requestID = uint32(x.random.Int31())
	x.msgID = uint32(x.random.Int31())
	return nil
}

func (x *GoSNMP) mkSnmpPacket(pdutype PDUType, nonRepeaters uint8, maxRepetitions uint8) *SnmpPacket {
	return &SnmpPacket{
		Version:            x.Version,
		Community:          x.Community,
		MsgFlags:           x.MsgFlags,
		SecurityModel:      x.SecurityModel,
		SecurityParameters: x.SecurityParameters,
		Error:              0,
		ErrorIndex:         0,
		PDUType:            pdutype,
		NonRepeaters:       nonRepeaters,
		MaxRepetitions:     maxRepetitions,
	}
}

// Get sends an SNMP GET request
func (x *GoSNMP) Get(oids []string) (result *SnmpPacket, err error) {
	oidCount := len(oids)
	if oidCount > maxOids {
		return nil, fmt.Errorf("oid count (%d) is greater than maxOids (%d)",
			oidCount, maxOids)
	}
	// convert oids slice to pdu slice
	var pdus []SnmpPDU
	for _, oid := range oids {
		pdus = append(pdus, SnmpPDU{oid, Null, nil})
	}
	// build up SnmpPacket
	packetOut := x.mkSnmpPacket(GetRequest, 0, 0)
	return x.send(pdus, packetOut)
}

// Set sends an SNMP SET request
func (x *GoSNMP) Set(pdus []SnmpPDU) (result *SnmpPacket, err error) {
	if len(pdus) != 1 {
		return nil, fmt.Errorf("gosnmp currently only supports SNMP SETs for one oid")
	}
	if pdus[0].Type != Integer {
		return nil, fmt.Errorf("gosnmp currently only supports SNMP SETs for Integers")
	}
	// build up SnmpPacket
	packetOut := x.mkSnmpPacket(SetRequest, 0, 0)
	return x.send(pdus, packetOut)
}

// GetNext sends an SNMP GETNEXT request
func (x *GoSNMP) GetNext(oids []string) (result *SnmpPacket, err error) {
	oidCount := len(oids)
	if oidCount > maxOids {
		return nil, fmt.Errorf("oid count (%d) is greater than maxOids (%d)",
			oidCount, maxOids)
	}

	// convert oids slice to pdu slice
	var pdus []SnmpPDU
	for _, oid := range oids {
		pdus = append(pdus, SnmpPDU{oid, Null, nil})
	}

	// Marshal and send the packet
	packetOut := x.mkSnmpPacket(GetNextRequest, 0, 0)

	return x.send(pdus, packetOut)
}

// GetBulk sends an SNMP GETBULK request
func (x *GoSNMP) GetBulk(oids []string, nonRepeaters uint8, maxRepetitions uint8) (result *SnmpPacket, err error) {
	oidCount := len(oids)
	if oidCount > maxOids {
		return nil, fmt.Errorf("oid count (%d) is greater than maxOids (%d)",
			oidCount, maxOids)
	}

	// convert oids slice to pdu slice
	var pdus []SnmpPDU
	for _, oid := range oids {
		pdus = append(pdus, SnmpPDU{oid, Null, nil})
	}

	// Marshal and send the packet
	packetOut := x.mkSnmpPacket(GetBulkRequest, nonRepeaters, maxRepetitions)
	return x.send(pdus, packetOut)
}

//
// SNMP Walk functions - Analogous to net-snmp's snmpwalk commands
//

// WalkFunc is the type of the function called for each data unit visited
// by the Walk function.  If an error is returned processing stops.
type WalkFunc func(dataUnit SnmpPDU) error

// BulkWalk retrieves a subtree of values using GETBULK. As the tree is
// walked walkFn is called for each new value. The function immediately returns
// an error if either there is an underlaying SNMP error (e.g. GetBulk fails),
// or if walkFn returns an error.
func (x *GoSNMP) BulkWalk(rootOid string, walkFn WalkFunc) error {
	return x.walk(GetBulkRequest, rootOid, walkFn)
}

// BulkWalkAll is similar to BulkWalk but returns a filled array of all values
// rather than using a callback function to stream results.
func (x *GoSNMP) BulkWalkAll(rootOid string) (results []SnmpPDU, err error) {
	return x.walkAll(GetBulkRequest, rootOid)
}

// Walk retrieves a subtree of values using GETNEXT - a request is made for each
// value, unlike BulkWalk which does this operation in batches. As the tree is
// walked walkFn is called for each new value. The function immediately returns
// an error if either there is an underlaying SNMP error (e.g. GetNext fails),
// or if walkFn returns an error.
func (x *GoSNMP) Walk(rootOid string, walkFn WalkFunc) error {
	return x.walk(GetNextRequest, rootOid, walkFn)
}

// WalkAll is similar to Walk but returns a filled array of all values rather
// than using a callback function to stream results.
func (x *GoSNMP) WalkAll(rootOid string) (results []SnmpPDU, err error) {
	return x.walkAll(GetNextRequest, rootOid)
}

//
// Public Functions (helpers) - in alphabetical order
//

// Partition - returns true when dividing a slice into
// partitionSize lengths, including last partition which may be smaller
// than partitionSize. This is useful when you have a large array of OIDs
// to run Get() on. See the tests for example usage.
//
// For example for a slice of 8 items to be broken into partitions of
// length 3, Partition returns true for the currentPosition having
// the following values:
//
// 0  1  2  3  4  5  6  7
//       T        T     T
//
func Partition(currentPosition, partitionSize, sliceLength int) bool {
	if currentPosition < 0 || currentPosition >= sliceLength {
		return false
	}
	if partitionSize == 1 { // redundant, but an obvious optimisation
		return true
	}
	if currentPosition%partitionSize == partitionSize-1 {
		return true
	}
	if currentPosition == sliceLength-1 {
		return true
	}
	return false
}

// ToBigInt converts SnmpPDU.Value to big.Int, or returns a zero big.Int for
// non int-like types (eg strings).
//
// This is a convenience function to make working with SnmpPDU's easier - it
// reduces the need for type assertions. A big.Int is convenient, as SNMP can
// return int32, uint32, and uint64.
func ToBigInt(value interface{}) *big.Int {
	var val int64
	switch value := value.(type) { // shadow
	case int:
		val = int64(value)
	case int8:
		val = int64(value)
	case int16:
		val = int64(value)
	case int32:
		val = int64(value)
	case int64:
		val = int64(value)
	case uint:
		val = int64(value)
	case uint8:
		val = int64(value)
	case uint16:
		val = int64(value)
	case uint32:
		val = int64(value)
	case uint64:
		return (uint64ToBigInt(value))
	case string:
		// for testing and other apps - numbers may appear as strings
		var err error
		if val, err = strconv.ParseInt(value, 10, 64); err != nil {
			return new(big.Int)
		}
	default:
		return new(big.Int)
	}
	return big.NewInt(val)
}
<commit_msg>Change GoSNMP.SecurityParameters to interface{} so multiple security models can be used.<commit_after>// Copyright 2012-2014 The GoSNMP Authors. All rights reserved.  Use of this
// source code is governed by a BSD-style license that can be found in the
// LICENSE file.

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

package gosnmp

import (
	"fmt"
	"math/big"
	"math/rand"
	"net"
	"strconv"
	"time"
)

const (
	// maxOids is the maximum number of oids allowed in a Get()
	maxOids = 60

	// Base OID for MIB-2 defined SNMP variables
	baseOid = ".1.3.6.1.2.1"

	// Java SNMP uses 50, snmp-net uses 10
	defaultMaxRepetitions = 50
)

// LoggingDisabled is set if the Logger is nil, short circuits any 'slog' calls
var LoggingDisabled bool

// GoSNMP represents GoSNMP library state
type GoSNMP struct {

	// Target is an ipv4 address
	Target string

	// Port is a udp port
	Port uint16

	// Version is an SNMP Version
	Version SnmpVersion

	// Community is an SNMP Community string
	Community string

	// MsgFlags is an SNMPV3 MsgFlags
	MsgFlags SnmpV3MsgFlags

	// SecurityModel is an SNMPV3 Security Model
	SecurityModel SnmpV3SecurityModel

	// SecurityParameters is an SNMPV3 Security Model paramaters struct
	SecurityParameters interface{}

	// Timeout is the timeout for the SNMP Query
	Timeout time.Duration

	// Set the number of retries to attempt within timeout.
	Retries int

	// Conn is net connection to use, typically establised using GoSNMP.Connect()
	Conn net.Conn

	// Logger is the GoSNMP.Logger to use for debugging. If nil, debugging
	// output will be discarded (/dev/null). For verbose logging to stdout:
	// x.Logger = log.New(os.Stdout, "", 0)
	Logger Logger

	// MaxRepititions sets the GETBULK max-repetitions used by BulkWalk*
	// (default: 50)
	MaxRepetitions int

	// NonRepeaters sets the GETBULK max-repeaters used by BulkWalk*
	// (default: 0 as per RFC 1905)
	NonRepeaters int

	// Internal - used to sync requests to responses
	requestID uint32

	// Internal - used to sync requests to responses - snmpv3
	msgID uint32

	random *rand.Rand
}

// The default connection settings
var Default = &GoSNMP{
	Port:      161,
	Community: "public",
	Version:   Version2c,
	Timeout:   time.Duration(2) * time.Second,
	Retries:   3,
}

// SnmpPDU will be used when doing SNMP Set's
type SnmpPDU struct {

	// Name is an oid in string format eg ".1.3.6.1.4.9.27"
	Name string

	// The type of the value eg Integer
	Type Asn1BER

	// The value to be set by the SNMP set
	Value interface{}
}

// Asn1BER is the type of the SNMP PDU
type Asn1BER byte

// Asn1BER's - http://www.ietf.org/rfc/rfc1442.txt
const (
	EndOfContents     Asn1BER = 0x00
	UnknownType               = 0x00 // TODO these should all be type Asn1BER, however
	Boolean                   = 0x01 // tests fail if implemented. See for example
	Integer                   = 0x02 /// http://stackoverflow.com/questions/5037610/typed-constant-declaration-list.
	BitString                 = 0x03
	OctetString               = 0x04
	Null                      = 0x05
	ObjectIdentifier          = 0x06
	ObjectDescription         = 0x07
	IPAddress                 = 0x40
	Counter32                 = 0x41
	Gauge32                   = 0x42
	TimeTicks                 = 0x43
	Opaque                    = 0x44
	NsapAddress               = 0x45
	Counter64                 = 0x46
	Uinteger32                = 0x47
	NoSuchObject              = 0x80
	NoSuchInstance            = 0x81
	EndOfMibView              = 0x82
)

//
// Public Functions (main interface)
//

// Connect initiates a connection to the target host
func (x *GoSNMP) Connect() error {
	if x.Logger == nil {
		LoggingDisabled = true
	}
	Conn, err := net.DialTimeout("udp", fmt.Sprintf("%s:%d", x.Target, x.Port), x.Timeout)
	if err == nil {
		x.Conn = Conn
	} else {
		return fmt.Errorf("Error establishing connection to host: %s\n", err.Error())
	}
	if x.random == nil {
		x.random = rand.New(rand.NewSource(time.Now().UTC().UnixNano()))
	}
	x.requestID = uint32(x.random.Int31())
	x.msgID = uint32(x.random.Int31())
	return nil
}

func (x *GoSNMP) mkSnmpPacket(pdutype PDUType, nonRepeaters uint8, maxRepetitions uint8) *SnmpPacket {
	return &SnmpPacket{
		Version:            x.Version,
		Community:          x.Community,
		MsgFlags:           x.MsgFlags,
		SecurityModel:      x.SecurityModel,
		SecurityParameters: x.SecurityParameters,
		Error:              0,
		ErrorIndex:         0,
		PDUType:            pdutype,
		NonRepeaters:       nonRepeaters,
		MaxRepetitions:     maxRepetitions,
	}
}

// Get sends an SNMP GET request
func (x *GoSNMP) Get(oids []string) (result *SnmpPacket, err error) {
	oidCount := len(oids)
	if oidCount > maxOids {
		return nil, fmt.Errorf("oid count (%d) is greater than maxOids (%d)",
			oidCount, maxOids)
	}
	// convert oids slice to pdu slice
	var pdus []SnmpPDU
	for _, oid := range oids {
		pdus = append(pdus, SnmpPDU{oid, Null, nil})
	}
	// build up SnmpPacket
	packetOut := x.mkSnmpPacket(GetRequest, 0, 0)
	return x.send(pdus, packetOut)
}

// Set sends an SNMP SET request
func (x *GoSNMP) Set(pdus []SnmpPDU) (result *SnmpPacket, err error) {
	if len(pdus) != 1 {
		return nil, fmt.Errorf("gosnmp currently only supports SNMP SETs for one oid")
	}
	if pdus[0].Type != Integer {
		return nil, fmt.Errorf("gosnmp currently only supports SNMP SETs for Integers")
	}
	// build up SnmpPacket
	packetOut := x.mkSnmpPacket(SetRequest, 0, 0)
	return x.send(pdus, packetOut)
}

// GetNext sends an SNMP GETNEXT request
func (x *GoSNMP) GetNext(oids []string) (result *SnmpPacket, err error) {
	oidCount := len(oids)
	if oidCount > maxOids {
		return nil, fmt.Errorf("oid count (%d) is greater than maxOids (%d)",
			oidCount, maxOids)
	}

	// convert oids slice to pdu slice
	var pdus []SnmpPDU
	for _, oid := range oids {
		pdus = append(pdus, SnmpPDU{oid, Null, nil})
	}

	// Marshal and send the packet
	packetOut := x.mkSnmpPacket(GetNextRequest, 0, 0)

	return x.send(pdus, packetOut)
}

// GetBulk sends an SNMP GETBULK request
func (x *GoSNMP) GetBulk(oids []string, nonRepeaters uint8, maxRepetitions uint8) (result *SnmpPacket, err error) {
	oidCount := len(oids)
	if oidCount > maxOids {
		return nil, fmt.Errorf("oid count (%d) is greater than maxOids (%d)",
			oidCount, maxOids)
	}

	// convert oids slice to pdu slice
	var pdus []SnmpPDU
	for _, oid := range oids {
		pdus = append(pdus, SnmpPDU{oid, Null, nil})
	}

	// Marshal and send the packet
	packetOut := x.mkSnmpPacket(GetBulkRequest, nonRepeaters, maxRepetitions)
	return x.send(pdus, packetOut)
}

//
// SNMP Walk functions - Analogous to net-snmp's snmpwalk commands
//

// WalkFunc is the type of the function called for each data unit visited
// by the Walk function.  If an error is returned processing stops.
type WalkFunc func(dataUnit SnmpPDU) error

// BulkWalk retrieves a subtree of values using GETBULK. As the tree is
// walked walkFn is called for each new value. The function immediately returns
// an error if either there is an underlaying SNMP error (e.g. GetBulk fails),
// or if walkFn returns an error.
func (x *GoSNMP) BulkWalk(rootOid string, walkFn WalkFunc) error {
	return x.walk(GetBulkRequest, rootOid, walkFn)
}

// BulkWalkAll is similar to BulkWalk but returns a filled array of all values
// rather than using a callback function to stream results.
func (x *GoSNMP) BulkWalkAll(rootOid string) (results []SnmpPDU, err error) {
	return x.walkAll(GetBulkRequest, rootOid)
}

// Walk retrieves a subtree of values using GETNEXT - a request is made for each
// value, unlike BulkWalk which does this operation in batches. As the tree is
// walked walkFn is called for each new value. The function immediately returns
// an error if either there is an underlaying SNMP error (e.g. GetNext fails),
// or if walkFn returns an error.
func (x *GoSNMP) Walk(rootOid string, walkFn WalkFunc) error {
	return x.walk(GetNextRequest, rootOid, walkFn)
}

// WalkAll is similar to Walk but returns a filled array of all values rather
// than using a callback function to stream results.
func (x *GoSNMP) WalkAll(rootOid string) (results []SnmpPDU, err error) {
	return x.walkAll(GetNextRequest, rootOid)
}

//
// Public Functions (helpers) - in alphabetical order
//

// Partition - returns true when dividing a slice into
// partitionSize lengths, including last partition which may be smaller
// than partitionSize. This is useful when you have a large array of OIDs
// to run Get() on. See the tests for example usage.
//
// For example for a slice of 8 items to be broken into partitions of
// length 3, Partition returns true for the currentPosition having
// the following values:
//
// 0  1  2  3  4  5  6  7
//       T        T     T
//
func Partition(currentPosition, partitionSize, sliceLength int) bool {
	if currentPosition < 0 || currentPosition >= sliceLength {
		return false
	}
	if partitionSize == 1 { // redundant, but an obvious optimisation
		return true
	}
	if currentPosition%partitionSize == partitionSize-1 {
		return true
	}
	if currentPosition == sliceLength-1 {
		return true
	}
	return false
}

// ToBigInt converts SnmpPDU.Value to big.Int, or returns a zero big.Int for
// non int-like types (eg strings).
//
// This is a convenience function to make working with SnmpPDU's easier - it
// reduces the need for type assertions. A big.Int is convenient, as SNMP can
// return int32, uint32, and uint64.
func ToBigInt(value interface{}) *big.Int {
	var val int64
	switch value := value.(type) { // shadow
	case int:
		val = int64(value)
	case int8:
		val = int64(value)
	case int16:
		val = int64(value)
	case int32:
		val = int64(value)
	case int64:
		val = int64(value)
	case uint:
		val = int64(value)
	case uint8:
		val = int64(value)
	case uint16:
		val = int64(value)
	case uint32:
		val = int64(value)
	case uint64:
		return (uint64ToBigInt(value))
	case string:
		// for testing and other apps - numbers may appear as strings
		var err error
		if val, err = strconv.ParseInt(value, 10, 64); err != nil {
			return new(big.Int)
		}
	default:
		return new(big.Int)
	}
	return big.NewInt(val)
}
<|endoftext|>


<commit_before>package simulationrep

import (
	"net/http"
	"sync"
	"time"

	"code.cloudfoundry.org/bbs/models"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/rep"
)

type SimulationRep struct {
	stack                  string
	zone                   string
	totalResources         rep.Resources
	lrps                   map[string]rep.LRP
	tasks                  map[string]rep.Task
	startingContainerCount int
	volumeDrivers          []string

	lock *sync.Mutex
}

func New(stack string, zone string, totalResources rep.Resources, volumeDrivers []string) rep.SimClient {
	return &SimulationRep{
		stack:          stack,
		totalResources: totalResources,
		lrps:           map[string]rep.LRP{},
		tasks:          map[string]rep.Task{},
		startingContainerCount: 0,
		zone:          zone,
		volumeDrivers: volumeDrivers,

		lock: &sync.Mutex{},
	}
}

func (r *SimulationRep) State(_ lager.Logger) (rep.CellState, error) {
	r.lock.Lock()
	defer r.lock.Unlock()

	lrps := []rep.LRP{}
	for _, lrp := range r.lrps {
		lrps = append(lrps, lrp)
	}

	tasks := []rep.Task{}
	for _, task := range r.tasks {
		tasks = append(tasks, task)
	}

	availableResources := r.availableResources()

	// util.RandomSleep(800, 900)

	return rep.CellState{
		RootFSProviders: rep.RootFSProviders{
			models.PreloadedRootFSScheme: rep.NewFixedSetRootFSProvider(r.stack),
		},
		AvailableResources: availableResources,
		TotalResources:     r.totalResources,
		LRPs:               lrps,
		Tasks:              tasks,
		StartingContainerCount: r.startingContainerCount,
		Zone:          r.zone,
		VolumeDrivers: r.volumeDrivers,
	}, nil
}

func (r *SimulationRep) Perform(_ lager.Logger, work rep.Work) (rep.Work, error) {
	r.lock.Lock()
	defer r.lock.Unlock()

	failedWork := rep.Work{}

	availableResources := r.availableResources()

	for _, start := range work.LRPs {
		hasRoom := availableResources.Containers >= 0
		hasRoom = hasRoom && availableResources.MemoryMB >= start.MemoryMB
		hasRoom = hasRoom && availableResources.DiskMB >= start.DiskMB

		if hasRoom {
			r.lrps[start.Identifier()] = start

			availableResources.Containers -= 1
			if start.Domain == "auction" {
				r.startingContainerCount++
			}
			availableResources.MemoryMB -= start.MemoryMB
			availableResources.DiskMB -= start.DiskMB
		} else {
			failedWork.LRPs = append(failedWork.LRPs, start)
		}
	}

	for _, task := range work.Tasks {
		hasRoom := availableResources.Containers >= 0
		hasRoom = hasRoom && availableResources.MemoryMB >= task.MemoryMB
		hasRoom = hasRoom && availableResources.DiskMB >= task.DiskMB

		if hasRoom {
			r.tasks[task.TaskGuid] = task

			availableResources.Containers -= 1
			if task.Domain == "auction" {
				r.startingContainerCount++
			}
			availableResources.MemoryMB -= task.MemoryMB
			availableResources.DiskMB -= task.DiskMB
		} else {
			failedWork.Tasks = append(failedWork.Tasks, task)
		}
	}

	return failedWork, nil
}

//simulation only

func (r *SimulationRep) Reset() error {
	r.lock.Lock()
	defer r.lock.Unlock()

	r.lrps = map[string]rep.LRP{}
	r.tasks = map[string]rep.Task{}
	r.startingContainerCount = 0
	return nil
}

//these are rep client methods the auction does not use

func (rep *SimulationRep) StopLRPInstance(models.ActualLRPKey, models.ActualLRPInstanceKey) error {
	panic("UNIMPLEMENTED METHOD")
}

func (rep *SimulationRep) CancelTask(string) error {
	panic("UNIMPLEMENTED METHOD")
}

func (rep *SimulationRep) SetStateClient(client *http.Client) {
	panic("UNIMPLEMENTED METHOD")
}

func (rep *SimulationRep) StateClientTimeout() time.Duration {
	panic("UNIMPLEMENTED METHOD")
}

//internal -- no locks here

func (rep *SimulationRep) availableResources() rep.Resources {
	resources := rep.totalResources
	for _, lrp := range rep.lrps {
		resources.MemoryMB -= lrp.MemoryMB
		resources.DiskMB -= lrp.DiskMB
		resources.Containers -= 1
	}
	for _, task := range rep.tasks {
		resources.MemoryMB -= task.MemoryMB
		resources.DiskMB -= task.DiskMB
		resources.Containers -= 1
	}
	return resources
}
<commit_msg>add logger to CancelTask and StopLRPInstance in rep Client<commit_after>package simulationrep

import (
	"net/http"
	"sync"
	"time"

	"code.cloudfoundry.org/bbs/models"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/rep"
)

type SimulationRep struct {
	stack                  string
	zone                   string
	totalResources         rep.Resources
	lrps                   map[string]rep.LRP
	tasks                  map[string]rep.Task
	startingContainerCount int
	volumeDrivers          []string

	lock *sync.Mutex
}

func New(stack string, zone string, totalResources rep.Resources, volumeDrivers []string) rep.SimClient {
	return &SimulationRep{
		stack:          stack,
		totalResources: totalResources,
		lrps:           map[string]rep.LRP{},
		tasks:          map[string]rep.Task{},
		startingContainerCount: 0,
		zone:          zone,
		volumeDrivers: volumeDrivers,

		lock: &sync.Mutex{},
	}
}

func (r *SimulationRep) State(_ lager.Logger) (rep.CellState, error) {
	r.lock.Lock()
	defer r.lock.Unlock()

	lrps := []rep.LRP{}
	for _, lrp := range r.lrps {
		lrps = append(lrps, lrp)
	}

	tasks := []rep.Task{}
	for _, task := range r.tasks {
		tasks = append(tasks, task)
	}

	availableResources := r.availableResources()

	// util.RandomSleep(800, 900)

	return rep.CellState{
		RootFSProviders: rep.RootFSProviders{
			models.PreloadedRootFSScheme: rep.NewFixedSetRootFSProvider(r.stack),
		},
		AvailableResources: availableResources,
		TotalResources:     r.totalResources,
		LRPs:               lrps,
		Tasks:              tasks,
		StartingContainerCount: r.startingContainerCount,
		Zone:          r.zone,
		VolumeDrivers: r.volumeDrivers,
	}, nil
}

func (r *SimulationRep) Perform(_ lager.Logger, work rep.Work) (rep.Work, error) {
	r.lock.Lock()
	defer r.lock.Unlock()

	failedWork := rep.Work{}

	availableResources := r.availableResources()

	for _, start := range work.LRPs {
		hasRoom := availableResources.Containers >= 0
		hasRoom = hasRoom && availableResources.MemoryMB >= start.MemoryMB
		hasRoom = hasRoom && availableResources.DiskMB >= start.DiskMB

		if hasRoom {
			r.lrps[start.Identifier()] = start

			availableResources.Containers -= 1
			if start.Domain == "auction" {
				r.startingContainerCount++
			}
			availableResources.MemoryMB -= start.MemoryMB
			availableResources.DiskMB -= start.DiskMB
		} else {
			failedWork.LRPs = append(failedWork.LRPs, start)
		}
	}

	for _, task := range work.Tasks {
		hasRoom := availableResources.Containers >= 0
		hasRoom = hasRoom && availableResources.MemoryMB >= task.MemoryMB
		hasRoom = hasRoom && availableResources.DiskMB >= task.DiskMB

		if hasRoom {
			r.tasks[task.TaskGuid] = task

			availableResources.Containers -= 1
			if task.Domain == "auction" {
				r.startingContainerCount++
			}
			availableResources.MemoryMB -= task.MemoryMB
			availableResources.DiskMB -= task.DiskMB
		} else {
			failedWork.Tasks = append(failedWork.Tasks, task)
		}
	}

	return failedWork, nil
}

//simulation only

func (r *SimulationRep) Reset() error {
	r.lock.Lock()
	defer r.lock.Unlock()

	r.lrps = map[string]rep.LRP{}
	r.tasks = map[string]rep.Task{}
	r.startingContainerCount = 0
	return nil
}

//these are rep client methods the auction does not use

func (rep *SimulationRep) StopLRPInstance(lager.Logger, models.ActualLRPKey, models.ActualLRPInstanceKey) error {
	panic("UNIMPLEMENTED METHOD")
}

func (rep *SimulationRep) CancelTask(lager.Logger, string) error {
	panic("UNIMPLEMENTED METHOD")
}

func (rep *SimulationRep) SetStateClient(client *http.Client) {
	panic("UNIMPLEMENTED METHOD")
}

func (rep *SimulationRep) StateClientTimeout() time.Duration {
	panic("UNIMPLEMENTED METHOD")
}

//internal -- no locks here

func (rep *SimulationRep) availableResources() rep.Resources {
	resources := rep.totalResources
	for _, lrp := range rep.lrps {
		resources.MemoryMB -= lrp.MemoryMB
		resources.DiskMB -= lrp.DiskMB
		resources.Containers -= 1
	}
	for _, task := range rep.tasks {
		resources.MemoryMB -= task.MemoryMB
		resources.DiskMB -= task.DiskMB
		resources.Containers -= 1
	}
	return resources
}
<|endoftext|>


<commit_before>package agent

import (
	"context"
	"fmt"
	"log"
	"time"

	"github.com/coreos/go-systemd/login1"
	"k8s.io/client-go/1.5/kubernetes"
	v1core "k8s.io/client-go/1.5/kubernetes/typed/core/v1"
	"k8s.io/client-go/1.5/pkg/api"
	v1api "k8s.io/client-go/1.5/pkg/api/v1"
	"k8s.io/client-go/1.5/pkg/fields"
	"k8s.io/client-go/1.5/pkg/watch"

	"github.com/coreos-inc/container-linux-update-operator/internal/constants"
	"github.com/coreos-inc/container-linux-update-operator/internal/drain"
	"github.com/coreos-inc/container-linux-update-operator/internal/k8sutil"
	"github.com/coreos-inc/container-linux-update-operator/internal/updateengine"
)

type Klocksmith struct {
	node string
	kc   *kubernetes.Clientset
	nc   v1core.NodeInterface
	ue   *updateengine.Client
	lc   *login1.Conn
}

func New(node string) (*Klocksmith, error) {
	// set up kubernetes in-cluster client
	kc, err := k8sutil.InClusterClient()
	if err != nil {
		return nil, fmt.Errorf("error creating kubernetes client: %v", err)
	}

	// node interface
	nc := kc.Nodes()

	// set up update_engine client
	ue, err := updateengine.New()
	if err != nil {
		return nil, fmt.Errorf("error establishing connection to update_engine dbus: %v", err)
	}

	// set up login1 client for our eventual reboot
	lc, err := login1.New()
	if err != nil {
		return nil, fmt.Errorf("error establishing connection to logind dbus: %v", err)
	}

	return &Klocksmith{node, kc, nc, ue, lc}, nil
}

// Run runs klocksmithd, reacting to the update_engine reboot signal by
// draining pods on this kubernetes node and rebooting.
//
// TODO: try to be more resilient against transient failures
func (k *Klocksmith) Run() error {
	// set annotations with helpful info about the coreos node
	vi, err := k8sutil.GetVersionInfo()
	if err != nil {
		return fmt.Errorf("failed to get version info: %v", err)
	}

	labels := map[string]string{
		constants.LabelID:      vi.ID,
		constants.LabelGroup:   vi.Group,
		constants.LabelVersion: vi.Version,
	}

	log.Printf("Setting labels: %#v", labels)
	if err := k8sutil.SetNodeLabels(k.nc, k.node, labels); err != nil {
		return err
	}

	// set coreos.com/update1/reboot-in-progress=false and
	// coreos.com/update1/reboot-needed=false
	anno := map[string]string{
		constants.AnnotationRebootInProgress: "false",
		constants.AnnotationRebootNeeded:     "false",
	}
	log.Printf("Setting annotations %#v", anno)
	if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
		return err
	}

	// we are schedulable now.
	log.Print("Marking node as schedulable")
	if err := k8sutil.Unschedulable(k.nc, k.node, false); err != nil {
		return err
	}

	// block until update_engine says to reboot, and set label.
	log.Printf("Waiting for reboot signal...")
	if err := k.waitForRebootSignal(); err != nil {
		return err
	}

	// indicate we need a reboot
	anno = map[string]string{
		constants.AnnotationRebootNeeded: "true",
	}
	log.Printf("Setting annotations %#v", anno)
	if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
		return err
	}

	// block until constants.AnnotationOkToReboot is set
	log.Printf("Waiting for ok-to-reboot from controller...")
	if err := k.waitForOkToReboot(); err != nil {
		return err
	}

	// set constants.AnnotationRebootInProgress and drain self
	anno = map[string]string{
		constants.AnnotationRebootInProgress: "true",
	}
	log.Printf("Setting annotations %#v", anno)
	if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
		return err
	}

	// drain self equates to:
	// 1. set Unschedulable
	// 2. delete all pods
	// unlike `kubectl drain`, we do not care about emptyDir or orphan pods
	// ('any pods that are neither mirror pods nor managed by
	// ReplicationController, ReplicaSet, DaemonSet or Job')

	log.Print("Marking node as unschedulable")
	if err := k8sutil.Unschedulable(k.nc, k.node, true); err != nil {
		return err
	}

	log.Print("Getting pod list for deletion")
	pods, err := k.getPodsForDeletion()
	if err != nil {
		return err
	}

	// delete the pods.
	// TODO: explicitly don't terminate self? we'll probably just be a
	// mirror pod or daemonset anyway..
	log.Printf("Deleting %d pods", len(pods))
	deleteOptions := api.NewDeleteOptions(30)
	for _, pod := range pods {
		log.Printf("Terminating pod %q...", pod.Name)
		if err := k.kc.Pods(pod.Namespace).Delete(pod.Name, deleteOptions); err != nil {
			return fmt.Errorf("failed terminating pod %q: %v", pod.Name, err)
		}
	}

	log.Print("Node drained, rebooting")

	// reboot
	k.lc.Reboot(false)

	// cross fingers
	time.Sleep(24 * 7 * time.Hour)

	return nil
}

// waitForRebootSignal watches update_engine and waits for
// updateengine.UpdateStatusUpdatedNeedReboot
func (k *Klocksmith) waitForRebootSignal() error {
	// first, check current status.
	status, err := k.ue.GetStatus()
	if err != nil {
		return fmt.Errorf("failed to get update_engine status: %v", err)
	}

	// if we're not in UpdateStatusUpdatedNeedReboot, then wait for it.
	if status.CurrentOperation != updateengine.UpdateStatusUpdatedNeedReboot {
		ctx, cancel := context.WithCancel(context.Background())
		stc := make(chan updateengine.Status)
		go k.ue.RebootNeededSignal(stc, ctx.Done())
		<-stc
		cancel()
	}

	// we are now in UpdateStatusUpdatedNeedReboot state.
	return nil
}

func (k *Klocksmith) waitForOkToReboot() error {
	n, err := k.nc.Get(k.node)
	if err != nil {
		return fmt.Errorf("failed to get self node (%q): %v", k.node, err)
	}

	// XXX: set timeout > 0?
	watcher, err := k.nc.Watch(api.ListOptions{
		FieldSelector:   fields.OneTermEqualSelector("metadata.name", n.Name),
		ResourceVersion: n.ResourceVersion,
	})
	if err != nil {
		return fmt.Errorf("failed to watch self node (%q): %v", k.node, err)
	}

	// hopefully 24 hours is enough time between indicating we need a
	// reboot and the controller telling us to do it
	ev, err := watch.Until(time.Hour*24, watcher, k8sutil.NodeAnnotationCondition(constants.AnnotationOkToReboot, "true"))
	if err != nil {
		return fmt.Errorf("waiting for annotation %q failed: %v", constants.AnnotationOkToReboot, err)
	}

	// sanity check
	no, ok := ev.Object.(*v1api.Node)
	if !ok {
		panic("event contains a non-*api.Node object")
	}

	if no.Annotations[constants.AnnotationOkToReboot] != "true" {
		panic("event did not contain annotation expected")
	}

	return nil
}

func (k *Klocksmith) getPodsForDeletion() ([]v1api.Pod, error) {
	pods, err := drain.GetPodsForDeletion(k.kc, k.node)
	if err != nil {
		return nil, fmt.Errorf("failed to get list of pods for deletion: %v", err)
	}

	// XXX: ignoring kube-system is a simple way to avoid eviciting
	// critical components such as kube-scheduler and
	// kube-controller-manager.

	pods = k8sutil.FilterPods(pods, func(p *v1api.Pod) bool {
		if p.Namespace == "kube-system" {
			return false
		}

		return true
	})

	return pods, nil
}
<commit_msg>internal/agent: constify annotation values<commit_after>package agent

import (
	"context"
	"fmt"
	"log"
	"time"

	"github.com/coreos/go-systemd/login1"
	"k8s.io/client-go/1.5/kubernetes"
	v1core "k8s.io/client-go/1.5/kubernetes/typed/core/v1"
	"k8s.io/client-go/1.5/pkg/api"
	v1api "k8s.io/client-go/1.5/pkg/api/v1"
	"k8s.io/client-go/1.5/pkg/fields"
	"k8s.io/client-go/1.5/pkg/watch"

	"github.com/coreos-inc/container-linux-update-operator/internal/constants"
	"github.com/coreos-inc/container-linux-update-operator/internal/drain"
	"github.com/coreos-inc/container-linux-update-operator/internal/k8sutil"
	"github.com/coreos-inc/container-linux-update-operator/internal/updateengine"
)

type Klocksmith struct {
	node string
	kc   *kubernetes.Clientset
	nc   v1core.NodeInterface
	ue   *updateengine.Client
	lc   *login1.Conn
}

func New(node string) (*Klocksmith, error) {
	// set up kubernetes in-cluster client
	kc, err := k8sutil.InClusterClient()
	if err != nil {
		return nil, fmt.Errorf("error creating kubernetes client: %v", err)
	}

	// node interface
	nc := kc.Nodes()

	// set up update_engine client
	ue, err := updateengine.New()
	if err != nil {
		return nil, fmt.Errorf("error establishing connection to update_engine dbus: %v", err)
	}

	// set up login1 client for our eventual reboot
	lc, err := login1.New()
	if err != nil {
		return nil, fmt.Errorf("error establishing connection to logind dbus: %v", err)
	}

	return &Klocksmith{node, kc, nc, ue, lc}, nil
}

// Run runs klocksmithd, reacting to the update_engine reboot signal by
// draining pods on this kubernetes node and rebooting.
//
// TODO: try to be more resilient against transient failures
func (k *Klocksmith) Run() error {
	// set annotations with helpful info about the coreos node
	vi, err := k8sutil.GetVersionInfo()
	if err != nil {
		return fmt.Errorf("failed to get version info: %v", err)
	}

	labels := map[string]string{
		constants.LabelID:      vi.ID,
		constants.LabelGroup:   vi.Group,
		constants.LabelVersion: vi.Version,
	}

	log.Printf("Setting labels: %#v", labels)
	if err := k8sutil.SetNodeLabels(k.nc, k.node, labels); err != nil {
		return err
	}

	// set coreos.com/update1/reboot-in-progress=false and
	// coreos.com/update1/reboot-needed=false
	anno := map[string]string{
		constants.AnnotationRebootInProgress: constants.False,
		constants.AnnotationRebootNeeded:     constants.False,
	}
	log.Printf("Setting annotations %#v", anno)
	if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
		return err
	}

	// we are schedulable now.
	log.Print("Marking node as schedulable")
	if err := k8sutil.Unschedulable(k.nc, k.node, false); err != nil {
		return err
	}

	// block until update_engine says to reboot, and set label.
	log.Printf("Waiting for reboot signal...")
	if err := k.waitForRebootSignal(); err != nil {
		return err
	}

	// indicate we need a reboot
	anno = map[string]string{
		constants.AnnotationRebootNeeded: constants.True,
	}
	log.Printf("Setting annotations %#v", anno)
	if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
		return err
	}

	// block until constants.AnnotationOkToReboot is set
	log.Printf("Waiting for ok-to-reboot from controller...")
	if err := k.waitForOkToReboot(); err != nil {
		return err
	}

	// set constants.AnnotationRebootInProgress and drain self
	anno = map[string]string{
		constants.AnnotationRebootInProgress: constants.True,
	}
	log.Printf("Setting annotations %#v", anno)
	if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
		return err
	}

	// drain self equates to:
	// 1. set Unschedulable
	// 2. delete all pods
	// unlike `kubectl drain`, we do not care about emptyDir or orphan pods
	// ('any pods that are neither mirror pods nor managed by
	// ReplicationController, ReplicaSet, DaemonSet or Job')

	log.Print("Marking node as unschedulable")
	if err := k8sutil.Unschedulable(k.nc, k.node, true); err != nil {
		return err
	}

	log.Print("Getting pod list for deletion")
	pods, err := k.getPodsForDeletion()
	if err != nil {
		return err
	}

	// delete the pods.
	// TODO: explicitly don't terminate self? we'll probably just be a
	// mirror pod or daemonset anyway..
	log.Printf("Deleting %d pods", len(pods))
	deleteOptions := api.NewDeleteOptions(30)
	for _, pod := range pods {
		log.Printf("Terminating pod %q...", pod.Name)
		if err := k.kc.Pods(pod.Namespace).Delete(pod.Name, deleteOptions); err != nil {
			return fmt.Errorf("failed terminating pod %q: %v", pod.Name, err)
		}
	}

	log.Print("Node drained, rebooting")

	// reboot
	k.lc.Reboot(false)

	// cross fingers
	time.Sleep(24 * 7 * time.Hour)

	return nil
}

// waitForRebootSignal watches update_engine and waits for
// updateengine.UpdateStatusUpdatedNeedReboot
func (k *Klocksmith) waitForRebootSignal() error {
	// first, check current status.
	status, err := k.ue.GetStatus()
	if err != nil {
		return fmt.Errorf("failed to get update_engine status: %v", err)
	}

	// if we're not in UpdateStatusUpdatedNeedReboot, then wait for it.
	if status.CurrentOperation != updateengine.UpdateStatusUpdatedNeedReboot {
		ctx, cancel := context.WithCancel(context.Background())
		stc := make(chan updateengine.Status)
		go k.ue.RebootNeededSignal(stc, ctx.Done())
		<-stc
		cancel()
	}

	// we are now in UpdateStatusUpdatedNeedReboot state.
	return nil
}

func (k *Klocksmith) waitForOkToReboot() error {
	n, err := k.nc.Get(k.node)
	if err != nil {
		return fmt.Errorf("failed to get self node (%q): %v", k.node, err)
	}

	// XXX: set timeout > 0?
	watcher, err := k.nc.Watch(api.ListOptions{
		FieldSelector:   fields.OneTermEqualSelector("metadata.name", n.Name),
		ResourceVersion: n.ResourceVersion,
	})
	if err != nil {
		return fmt.Errorf("failed to watch self node (%q): %v", k.node, err)
	}

	// hopefully 24 hours is enough time between indicating we need a
	// reboot and the controller telling us to do it
	ev, err := watch.Until(time.Hour*24, watcher, k8sutil.NodeAnnotationCondition(constants.AnnotationOkToReboot, constants.True))
	if err != nil {
		return fmt.Errorf("waiting for annotation %q failed: %v", constants.AnnotationOkToReboot, err)
	}

	// sanity check
	no, ok := ev.Object.(*v1api.Node)
	if !ok {
		panic("event contains a non-*api.Node object")
	}

	if no.Annotations[constants.AnnotationOkToReboot] != constants.True {
		panic("event did not contain annotation expected")
	}

	return nil
}

func (k *Klocksmith) getPodsForDeletion() ([]v1api.Pod, error) {
	pods, err := drain.GetPodsForDeletion(k.kc, k.node)
	if err != nil {
		return nil, fmt.Errorf("failed to get list of pods for deletion: %v", err)
	}

	// XXX: ignoring kube-system is a simple way to avoid eviciting
	// critical components such as kube-scheduler and
	// kube-controller-manager.

	pods = k8sutil.FilterPods(pods, func(p *v1api.Pod) bool {
		if p.Namespace == "kube-system" {
			return false
		}

		return true
	})

	return pods, nil
}
<|endoftext|>