<commit_before>// Copyright 2016 Russell Haering et al.
//
// 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
//
//     https://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 saml2

import (
	"bytes"
	"compress/flate"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"fmt"
	"io/ioutil"

	"encoding/xml"

	"github.com/beevik/etree"
	"github.com/russellhaering/gosaml2/types"
	dsig "github.com/russellhaering/goxmldsig"
	"github.com/russellhaering/goxmldsig/etreeutils"
	rtvalidator "github.com/mattermost/xml-roundtrip-validator"
)

func (sp *SAMLServiceProvider) validationContext() *dsig.ValidationContext {
	ctx := dsig.NewDefaultValidationContext(sp.IDPCertificateStore)
	ctx.Clock = sp.Clock
	return ctx
}

// validateResponseAttributes validates a SAML Response's tag and attributes. It does
// not inspect child elements of the Response at all.
func (sp *SAMLServiceProvider) validateResponseAttributes(response *types.Response) error {
	if response.Destination != "" && response.Destination != sp.AssertionConsumerServiceURL {
		return ErrInvalidValue{
			Key:      DestinationAttr,
			Expected: sp.AssertionConsumerServiceURL,
			Actual:   response.Destination,
		}
	}

	if response.Version != "2.0" {
		return ErrInvalidValue{
			Reason:   ReasonUnsupported,
			Key:      "SAML version",
			Expected: "2.0",
			Actual:   response.Version,
		}
	}

	return nil
}

// validateLogoutResponseAttributes validates a SAML Response's tag and attributes. It does
// not inspect child elements of the Response at all.
func (sp *SAMLServiceProvider) validateLogoutResponseAttributes(response *types.LogoutResponse) error {
	if response.Destination != "" && response.Destination != sp.ServiceProviderSLOURL {
		return ErrInvalidValue{
			Key:      DestinationAttr,
			Expected: sp.ServiceProviderSLOURL,
			Actual:   response.Destination,
		}
	}

	if response.Version != "2.0" {
		return ErrInvalidValue{
			Reason:   ReasonUnsupported,
			Key:      "SAML version",
			Expected: "2.0",
			Actual:   response.Version,
		}
	}

	return nil
}

func xmlUnmarshalElement(el *etree.Element, obj interface{}) error {
	doc := etree.NewDocument()
	doc.SetRoot(el)
	data, err := doc.WriteToBytes()
	if err != nil {
		return err
	}

	err = xml.Unmarshal(data, obj)
	if err != nil {
		return err
	}
	return nil
}

func (sp *SAMLServiceProvider) getDecryptCert() (*tls.Certificate, error) {
	if sp.SPKeyStore == nil {
		return nil, fmt.Errorf("no decryption certs available")
	}

	//This is the tls.Certificate we'll use to decrypt any encrypted assertions
	var decryptCert tls.Certificate

	switch crt := sp.SPKeyStore.(type) {
	case dsig.TLSCertKeyStore:
		// Get the tls.Certificate directly if possible
		decryptCert = tls.Certificate(crt)

	default:

		//Otherwise, construct one from the results of GetKeyPair
		pk, cert, err := sp.SPKeyStore.GetKeyPair()
		if err != nil {
			return nil, fmt.Errorf("error getting keypair: %v", err)
		}

		decryptCert = tls.Certificate{
			Certificate: [][]byte{cert},
			PrivateKey:  pk,
		}
	}

	if sp.ValidateEncryptionCert {
		// Check Validity period of certificate
		if len(decryptCert.Certificate) < 1 || len(decryptCert.Certificate[0]) < 1 {
			return nil, fmt.Errorf("empty decryption cert")
		} else if cert, err := x509.ParseCertificate(decryptCert.Certificate[0]); err != nil {
			return nil, fmt.Errorf("invalid x509 decryption cert: %v", err)
		} else {
			now := sp.Clock.Now()
			if now.Before(cert.NotBefore) || now.After(cert.NotAfter) {
				return nil, fmt.Errorf("decryption cert is not valid at this time")
			}
		}
	}

	return &decryptCert, nil
}

func (sp *SAMLServiceProvider) decryptAssertions(el *etree.Element) error {
	var decryptCert *tls.Certificate

	decryptAssertion := func(ctx etreeutils.NSContext, encryptedElement *etree.Element) error {
		if encryptedElement.Parent() != el {
			return fmt.Errorf("found encrypted assertion with unexpected parent element: %s", encryptedElement.Parent().Tag)
		}

		detached, err := etreeutils.NSDetatch(ctx, encryptedElement) // make a detached copy
		if err != nil {
			return fmt.Errorf("unable to detach encrypted assertion: %v", err)
		}

		encryptedAssertion := &types.EncryptedAssertion{}
		err = xmlUnmarshalElement(detached, encryptedAssertion)
		if err != nil {
			return fmt.Errorf("unable to unmarshal encrypted assertion: %v", err)
		}

		if decryptCert == nil {
			decryptCert, err = sp.getDecryptCert()
			if err != nil {
				return fmt.Errorf("unable to get decryption certificate: %v", err)
			}
		}

		raw, derr := encryptedAssertion.DecryptBytes(decryptCert)
		if derr != nil {
			return fmt.Errorf("unable to decrypt encrypted assertion: %v", derr)
		}

		doc, _, err := parseResponse(raw)
		if err != nil {
			return fmt.Errorf("unable to create element from decrypted assertion bytes: %v", derr)
		}

		// Replace the original encrypted assertion with the decrypted one.
		if el.RemoveChild(encryptedElement) == nil {
			// Out of an abundance of caution, make sure removed worked
			panic("unable to remove encrypted assertion")
		}

		el.AddChild(doc.Root())
		return nil
	}

	if err := etreeutils.NSFindIterate(el, SAMLAssertionNamespace, EncryptedAssertionTag, decryptAssertion); err != nil {
		return err
	} else {
		return nil
	}
}

func (sp *SAMLServiceProvider) validateElementSignature(el *etree.Element) (*etree.Element, error) {
	return sp.validationContext().Validate(el)
}

func (sp *SAMLServiceProvider) validateAssertionSignatures(el *etree.Element) error {
	signedAssertions := 0
	unsignedAssertions := 0
	validateAssertion := func(ctx etreeutils.NSContext, unverifiedAssertion *etree.Element) error {
		if unverifiedAssertion.Parent() != el {
			return fmt.Errorf("found assertion with unexpected parent element: %s", unverifiedAssertion.Parent().Tag)
		}

		detached, err := etreeutils.NSDetatch(ctx, unverifiedAssertion) // make a detached copy
		if err != nil {
			return fmt.Errorf("unable to detach unverified assertion: %v", err)
		}

		assertion, err := sp.validationContext().Validate(detached)
		if err == dsig.ErrMissingSignature {
			unsignedAssertions++
			return nil
		} else if err != nil {
			return err
		}

		// Replace the original unverified Assertion with the verified one. Note that
		// if the Response is not signed, only signed Assertions (and not the parent Response) can be trusted.
		if el.RemoveChild(unverifiedAssertion) == nil {
			// Out of an abundance of caution, check to make sure an Assertion was actually
			// removed. If it wasn't a programming error has occurred.
			panic("unable to remove assertion")
		}

		el.AddChild(assertion)
		signedAssertions++

		return nil
	}

	if err := etreeutils.NSFindIterate(el, SAMLAssertionNamespace, AssertionTag, validateAssertion); err != nil {
		return err
	} else if signedAssertions > 0 && unsignedAssertions > 0 {
		return fmt.Errorf("invalid to have both signed and unsigned assertions")
	} else if signedAssertions < 1 {
		return dsig.ErrMissingSignature
	} else {
		return nil
	}
}

//ValidateEncodedResponse both decodes and validates, based on SP
//configuration, an encoded, signed response. It will also appropriately
//decrypt a response if the assertion was encrypted
func (sp *SAMLServiceProvider) ValidateEncodedResponse(encodedResponse string) (*types.Response, error) {
	raw, err := base64.StdEncoding.DecodeString(encodedResponse)
	if err != nil {
		return nil, err
	}

	// Parse the raw response
	doc, el, err := parseResponse(raw)
	if err != nil {
		return nil, err
	}

	var responseSignatureValidated bool
	if !sp.SkipSignatureValidation {
		el, err = sp.validateElementSignature(el)
		if err == dsig.ErrMissingSignature {
			// Unfortunately we just blew away our Response
			el = doc.Root()
		} else if err != nil {
			return nil, err
		} else if el == nil {
			return nil, fmt.Errorf("missing transformed response")
		} else {
			responseSignatureValidated = true
		}
	}

	err = sp.decryptAssertions(el)
	if err != nil {
		return nil, err
	}

	var assertionSignaturesValidated bool
	if !sp.SkipSignatureValidation {
		err = sp.validateAssertionSignatures(el)
		if err == dsig.ErrMissingSignature {
			if !responseSignatureValidated {
				return nil, fmt.Errorf("response and/or assertions must be signed")
			}
		} else if err != nil {
			return nil, err
		} else {
			assertionSignaturesValidated = true
		}
	}

	decodedResponse := &types.Response{}
	err = xmlUnmarshalElement(el, decodedResponse)
	if err != nil {
		return nil, fmt.Errorf("unable to unmarshal response: %v", err)
	}
	decodedResponse.SignatureValidated = responseSignatureValidated
	if assertionSignaturesValidated {
		for idx := 0; idx < len(decodedResponse.Assertions); idx++ {
			decodedResponse.Assertions[idx].SignatureValidated = true
		}
	}

	err = sp.Validate(decodedResponse)
	if err != nil {
		return nil, err
	}

	return decodedResponse, nil
}

// DecodeUnverifiedBaseResponse decodes several attributes from a SAML response for the purpose
// of determining how to validate the response. This is useful for Service Providers which
// expose a single Assertion Consumer Service URL but consume Responses from many IdPs.
func DecodeUnverifiedBaseResponse(encodedResponse string) (*types.UnverifiedBaseResponse, error) {
	raw, err := base64.StdEncoding.DecodeString(encodedResponse)
	if err != nil {
		return nil, err
	}

	var response *types.UnverifiedBaseResponse

	err = maybeDeflate(raw, func(maybeXML []byte) error {
		response = &types.UnverifiedBaseResponse{}
		return xml.Unmarshal(maybeXML, response)
	})
	if err != nil {
		return nil, err
	}

	return response, nil
}

// maybeDeflate invokes the passed decoder over the passed data. If an error is
// returned, it then attempts to deflate the passed data before re-invoking
// the decoder over the deflated data.
func maybeDeflate(data []byte, decoder func([]byte) error) error {
	err := decoder(data)
	if err == nil {
		return nil
	}

	deflated, err := ioutil.ReadAll(flate.NewReader(bytes.NewReader(data)))
	if err != nil {
		return err
	}

	return decoder(deflated)
}

// parseResponse is a helper function that was refactored out so that the XML parsing behavior can be isolated and unit tested
func parseResponse(xml []byte) (*etree.Document, *etree.Element, error) {
	var doc *etree.Document
	var rawXML []byte

	err := maybeDeflate(xml, func(xml []byte) error {
		doc = etree.NewDocument()
		rawXML = xml
		return doc.ReadFromBytes(xml)
	})
	if err != nil {
		return nil, nil, err
	}

	el := doc.Root()
	if el == nil {
		return nil, nil, fmt.Errorf("unable to parse response")
	}

	// Examine the response for attempts to exploit weaknesses in Go's encoding/xml
	err = rtvalidator.Validate(bytes.NewReader(rawXML))
	if err != nil {
		return nil, nil, err
	}

	return doc, el, nil
}

// DecodeUnverifiedLogoutResponse decodes several attributes from a SAML Logout response, without doing any verifications.
func DecodeUnverifiedLogoutResponse(encodedResponse string) (*types.LogoutResponse, error) {
	raw, err := base64.StdEncoding.DecodeString(encodedResponse)
	if err != nil {
		return nil, err
	}

	var response *types.LogoutResponse

	err = maybeDeflate(raw, func(maybeXML []byte) error {
		response = &types.LogoutResponse{}
		return xml.Unmarshal(maybeXML, response)
	})
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (sp *SAMLServiceProvider) ValidateEncodedLogoutResponsePOST(encodedResponse string) (*types.LogoutResponse, error) {
	raw, err := base64.StdEncoding.DecodeString(encodedResponse)
	if err != nil {
		return nil, err
	}

	// Parse the raw response
	doc, el, err := parseResponse(raw)
	if err != nil {
		return nil, err
	}

	var responseSignatureValidated bool
	if !sp.SkipSignatureValidation {
		el, err = sp.validateElementSignature(el)
		if err == dsig.ErrMissingSignature {
			// Unfortunately we just blew away our Response
			el = doc.Root()
		} else if err != nil {
			return nil, err
		} else if el == nil {
			return nil, fmt.Errorf("missing transformed logout response")
		} else {
			responseSignatureValidated = true
		}
	}

	decodedResponse := &types.LogoutResponse{}
	err = xmlUnmarshalElement(el, decodedResponse)
	if err != nil {
		return nil, fmt.Errorf("unable to unmarshal logout response: %v", err)
	}
	decodedResponse.SignatureValidated = responseSignatureValidated

	err = sp.ValidateDecodedLogoutResponse(decodedResponse)
	if err != nil {
		return nil, err
	}

	return decodedResponse, nil
}<commit_msg>fixes panic on malformed input<commit_after>// Copyright 2016 Russell Haering et al.
//
// 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
//
//     https://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 saml2

import (
	"bytes"
	"compress/flate"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"fmt"
	"io/ioutil"

	"encoding/xml"

	"github.com/beevik/etree"
	"github.com/russellhaering/gosaml2/types"
	dsig "github.com/russellhaering/goxmldsig"
	"github.com/russellhaering/goxmldsig/etreeutils"
	rtvalidator "github.com/mattermost/xml-roundtrip-validator"
)

func (sp *SAMLServiceProvider) validationContext() *dsig.ValidationContext {
	ctx := dsig.NewDefaultValidationContext(sp.IDPCertificateStore)
	ctx.Clock = sp.Clock
	return ctx
}

// validateResponseAttributes validates a SAML Response's tag and attributes. It does
// not inspect child elements of the Response at all.
func (sp *SAMLServiceProvider) validateResponseAttributes(response *types.Response) error {
	if response.Destination != "" && response.Destination != sp.AssertionConsumerServiceURL {
		return ErrInvalidValue{
			Key:      DestinationAttr,
			Expected: sp.AssertionConsumerServiceURL,
			Actual:   response.Destination,
		}
	}

	if response.Version != "2.0" {
		return ErrInvalidValue{
			Reason:   ReasonUnsupported,
			Key:      "SAML version",
			Expected: "2.0",
			Actual:   response.Version,
		}
	}

	return nil
}

// validateLogoutResponseAttributes validates a SAML Response's tag and attributes. It does
// not inspect child elements of the Response at all.
func (sp *SAMLServiceProvider) validateLogoutResponseAttributes(response *types.LogoutResponse) error {
	if response.Destination != "" && response.Destination != sp.ServiceProviderSLOURL {
		return ErrInvalidValue{
			Key:      DestinationAttr,
			Expected: sp.ServiceProviderSLOURL,
			Actual:   response.Destination,
		}
	}

	if response.Version != "2.0" {
		return ErrInvalidValue{
			Reason:   ReasonUnsupported,
			Key:      "SAML version",
			Expected: "2.0",
			Actual:   response.Version,
		}
	}

	return nil
}

func xmlUnmarshalElement(el *etree.Element, obj interface{}) error {
	doc := etree.NewDocument()
	doc.SetRoot(el)
	data, err := doc.WriteToBytes()
	if err != nil {
		return err
	}

	err = xml.Unmarshal(data, obj)
	if err != nil {
		return err
	}
	return nil
}

func (sp *SAMLServiceProvider) getDecryptCert() (*tls.Certificate, error) {
	if sp.SPKeyStore == nil {
		return nil, fmt.Errorf("no decryption certs available")
	}

	//This is the tls.Certificate we'll use to decrypt any encrypted assertions
	var decryptCert tls.Certificate

	switch crt := sp.SPKeyStore.(type) {
	case dsig.TLSCertKeyStore:
		// Get the tls.Certificate directly if possible
		decryptCert = tls.Certificate(crt)

	default:

		//Otherwise, construct one from the results of GetKeyPair
		pk, cert, err := sp.SPKeyStore.GetKeyPair()
		if err != nil {
			return nil, fmt.Errorf("error getting keypair: %v", err)
		}

		decryptCert = tls.Certificate{
			Certificate: [][]byte{cert},
			PrivateKey:  pk,
		}
	}

	if sp.ValidateEncryptionCert {
		// Check Validity period of certificate
		if len(decryptCert.Certificate) < 1 || len(decryptCert.Certificate[0]) < 1 {
			return nil, fmt.Errorf("empty decryption cert")
		} else if cert, err := x509.ParseCertificate(decryptCert.Certificate[0]); err != nil {
			return nil, fmt.Errorf("invalid x509 decryption cert: %v", err)
		} else {
			now := sp.Clock.Now()
			if now.Before(cert.NotBefore) || now.After(cert.NotAfter) {
				return nil, fmt.Errorf("decryption cert is not valid at this time")
			}
		}
	}

	return &decryptCert, nil
}

func (sp *SAMLServiceProvider) decryptAssertions(el *etree.Element) error {
	var decryptCert *tls.Certificate

	decryptAssertion := func(ctx etreeutils.NSContext, encryptedElement *etree.Element) error {
		if encryptedElement.Parent() != el {
			return fmt.Errorf("found encrypted assertion with unexpected parent element: %s", encryptedElement.Parent().Tag)
		}

		detached, err := etreeutils.NSDetatch(ctx, encryptedElement) // make a detached copy
		if err != nil {
			return fmt.Errorf("unable to detach encrypted assertion: %v", err)
		}

		encryptedAssertion := &types.EncryptedAssertion{}
		err = xmlUnmarshalElement(detached, encryptedAssertion)
		if err != nil {
			return fmt.Errorf("unable to unmarshal encrypted assertion: %v", err)
		}

		if decryptCert == nil {
			decryptCert, err = sp.getDecryptCert()
			if err != nil {
				return fmt.Errorf("unable to get decryption certificate: %v", err)
			}
		}

		raw, derr := encryptedAssertion.DecryptBytes(decryptCert)
		if derr != nil {
			return fmt.Errorf("unable to decrypt encrypted assertion: %v", derr)
		}

		doc, _, err := parseResponse(raw)
		if err != nil {
			return fmt.Errorf("unable to create element from decrypted assertion bytes: %v", derr)
		}

		// Replace the original encrypted assertion with the decrypted one.
		if el.RemoveChild(encryptedElement) == nil {
			// Out of an abundance of caution, make sure removed worked
			panic("unable to remove encrypted assertion")
		}

		el.AddChild(doc.Root())
		return nil
	}

	if err := etreeutils.NSFindIterate(el, SAMLAssertionNamespace, EncryptedAssertionTag, decryptAssertion); err != nil {
		return err
	} else {
		return nil
	}
}

func (sp *SAMLServiceProvider) validateElementSignature(el *etree.Element) (*etree.Element, error) {
	return sp.validationContext().Validate(el)
}

func (sp *SAMLServiceProvider) validateAssertionSignatures(el *etree.Element) error {
	signedAssertions := 0
	unsignedAssertions := 0
	validateAssertion := func(ctx etreeutils.NSContext, unverifiedAssertion *etree.Element) error {
		parent := unverifiedAssertion.Parent()
		if parent == nil {
			return fmt.Errorf("parent is nil")
		}
		if parent != el {
			return fmt.Errorf("found assertion with unexpected parent element: %s", unverifiedAssertion.Parent().Tag)
		}

		detached, err := etreeutils.NSDetatch(ctx, unverifiedAssertion) // make a detached copy
		if err != nil {
			return fmt.Errorf("unable to detach unverified assertion: %v", err)
		}

		assertion, err := sp.validationContext().Validate(detached)
		if err == dsig.ErrMissingSignature {
			unsignedAssertions++
			return nil
		} else if err != nil {
			return err
		}

		// Replace the original unverified Assertion with the verified one. Note that
		// if the Response is not signed, only signed Assertions (and not the parent Response) can be trusted.
		if el.RemoveChild(unverifiedAssertion) == nil {
			// Out of an abundance of caution, check to make sure an Assertion was actually
			// removed. If it wasn't a programming error has occurred.
			panic("unable to remove assertion")
		}

		el.AddChild(assertion)
		signedAssertions++

		return nil
	}

	if err := etreeutils.NSFindIterate(el, SAMLAssertionNamespace, AssertionTag, validateAssertion); err != nil {
		return err
	} else if signedAssertions > 0 && unsignedAssertions > 0 {
		return fmt.Errorf("invalid to have both signed and unsigned assertions")
	} else if signedAssertions < 1 {
		return dsig.ErrMissingSignature
	} else {
		return nil
	}
}

//ValidateEncodedResponse both decodes and validates, based on SP
//configuration, an encoded, signed response. It will also appropriately
//decrypt a response if the assertion was encrypted
func (sp *SAMLServiceProvider) ValidateEncodedResponse(encodedResponse string) (*types.Response, error) {
	raw, err := base64.StdEncoding.DecodeString(encodedResponse)
	if err != nil {
		return nil, err
	}

	// Parse the raw response
	doc, el, err := parseResponse(raw)
	if err != nil {
		return nil, err
	}

	var responseSignatureValidated bool
	if !sp.SkipSignatureValidation {
		el, err = sp.validateElementSignature(el)
		if err == dsig.ErrMissingSignature {
			// Unfortunately we just blew away our Response
			el = doc.Root()
		} else if err != nil {
			return nil, err
		} else if el == nil {
			return nil, fmt.Errorf("missing transformed response")
		} else {
			responseSignatureValidated = true
		}
	}

	err = sp.decryptAssertions(el)
	if err != nil {
		return nil, err
	}

	var assertionSignaturesValidated bool
	if !sp.SkipSignatureValidation {
		err = sp.validateAssertionSignatures(el)
		if err == dsig.ErrMissingSignature {
			if !responseSignatureValidated {
				return nil, fmt.Errorf("response and/or assertions must be signed")
			}
		} else if err != nil {
			return nil, err
		} else {
			assertionSignaturesValidated = true
		}
	}

	decodedResponse := &types.Response{}
	err = xmlUnmarshalElement(el, decodedResponse)
	if err != nil {
		return nil, fmt.Errorf("unable to unmarshal response: %v", err)
	}
	decodedResponse.SignatureValidated = responseSignatureValidated
	if assertionSignaturesValidated {
		for idx := 0; idx < len(decodedResponse.Assertions); idx++ {
			decodedResponse.Assertions[idx].SignatureValidated = true
		}
	}

	err = sp.Validate(decodedResponse)
	if err != nil {
		return nil, err
	}

	return decodedResponse, nil
}

// DecodeUnverifiedBaseResponse decodes several attributes from a SAML response for the purpose
// of determining how to validate the response. This is useful for Service Providers which
// expose a single Assertion Consumer Service URL but consume Responses from many IdPs.
func DecodeUnverifiedBaseResponse(encodedResponse string) (*types.UnverifiedBaseResponse, error) {
	raw, err := base64.StdEncoding.DecodeString(encodedResponse)
	if err != nil {
		return nil, err
	}

	var response *types.UnverifiedBaseResponse

	err = maybeDeflate(raw, func(maybeXML []byte) error {
		response = &types.UnverifiedBaseResponse{}
		return xml.Unmarshal(maybeXML, response)
	})
	if err != nil {
		return nil, err
	}

	return response, nil
}

// maybeDeflate invokes the passed decoder over the passed data. If an error is
// returned, it then attempts to deflate the passed data before re-invoking
// the decoder over the deflated data.
func maybeDeflate(data []byte, decoder func([]byte) error) error {
	err := decoder(data)
	if err == nil {
		return nil
	}

	deflated, err := ioutil.ReadAll(flate.NewReader(bytes.NewReader(data)))
	if err != nil {
		return err
	}

	return decoder(deflated)
}

// parseResponse is a helper function that was refactored out so that the XML parsing behavior can be isolated and unit tested
func parseResponse(xml []byte) (*etree.Document, *etree.Element, error) {
	var doc *etree.Document
	var rawXML []byte

	err := maybeDeflate(xml, func(xml []byte) error {
		doc = etree.NewDocument()
		rawXML = xml
		return doc.ReadFromBytes(xml)
	})
	if err != nil {
		return nil, nil, err
	}

	el := doc.Root()
	if el == nil {
		return nil, nil, fmt.Errorf("unable to parse response")
	}

	// Examine the response for attempts to exploit weaknesses in Go's encoding/xml
	err = rtvalidator.Validate(bytes.NewReader(rawXML))
	if err != nil {
		return nil, nil, err
	}

	return doc, el, nil
}

// DecodeUnverifiedLogoutResponse decodes several attributes from a SAML Logout response, without doing any verifications.
func DecodeUnverifiedLogoutResponse(encodedResponse string) (*types.LogoutResponse, error) {
	raw, err := base64.StdEncoding.DecodeString(encodedResponse)
	if err != nil {
		return nil, err
	}

	var response *types.LogoutResponse

	err = maybeDeflate(raw, func(maybeXML []byte) error {
		response = &types.LogoutResponse{}
		return xml.Unmarshal(maybeXML, response)
	})
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (sp *SAMLServiceProvider) ValidateEncodedLogoutResponsePOST(encodedResponse string) (*types.LogoutResponse, error) {
	raw, err := base64.StdEncoding.DecodeString(encodedResponse)
	if err != nil {
		return nil, err
	}

	// Parse the raw response
	doc, el, err := parseResponse(raw)
	if err != nil {
		return nil, err
	}

	var responseSignatureValidated bool
	if !sp.SkipSignatureValidation {
		el, err = sp.validateElementSignature(el)
		if err == dsig.ErrMissingSignature {
			// Unfortunately we just blew away our Response
			el = doc.Root()
		} else if err != nil {
			return nil, err
		} else if el == nil {
			return nil, fmt.Errorf("missing transformed logout response")
		} else {
			responseSignatureValidated = true
		}
	}

	decodedResponse := &types.LogoutResponse{}
	err = xmlUnmarshalElement(el, decodedResponse)
	if err != nil {
		return nil, fmt.Errorf("unable to unmarshal logout response: %v", err)
	}
	decodedResponse.SignatureValidated = responseSignatureValidated

	err = sp.ValidateDecodedLogoutResponse(decodedResponse)
	if err != nil {
		return nil, err
	}

	return decodedResponse, nil
}<|endoftext|>


<commit_before>package profile_test

import (
	"fmt"
	"strings"
	"testing"
	"time"

	. "github.com/PhilipRasmussen/minecraft/profile"
)

/************
* TEST DATA *
************/

// 1. known account of author
const (
	nergName = "Nergalic"
	nergID   = "087cc153c3434ff7ac497de1569affa1"
	// Model is expected to be Steve
	// Cape is not expected
)

var nergHist = []pastName{{"GeneralSezuan", time.Unix(1423047705, 0)}}

// 2. known account of author
const (
	breeName = "BreeSakana"
	breeID   = "d9a5b542ce88442aaab38ec13e6c7773"
)

var breeHist = []pastName{}

// akronman1, the holder of the 1.000.000th-Minecraft-copy cape
// Profile may in theory be deleted at any moment...
//
// TODO: Substitute for a cape profile under author's control
const capeID = "d90b68bc81724329a047f1186dcd4336"

/******************
* TEST STRUCTURES *
******************/

// Username --> profile info
// Must maintain consistency with the other test structures
var loadTestUsers = map[string]user{
	nergName:                  {nergID, nergName},
	strings.ToUpper(nergName): {nergID, nergName}, // Casing should not matter
	breeName:                  {breeID, breeName},
}

// Should match loadTestUsers without duplicates
var loadManyTestRes = []string{breeName, nergName}

// Keys should match entries in loadTestUsers
var pastNames = map[string][]pastName{
	breeName: breeHist,
	nergName: nergHist,

	// TODO: Add profiles with longer name history to verify order
}

/*************
* TEST TYPES *
*************/

type user struct {
	id   string
	name string
}

type pastName struct {
	name  string
	until time.Time
}

type propertySet struct {

	// Do not check their values, simply check their presence
	skinURL, capeURL *bool // expected skin and cape (if any)

	model *Model // expected model (if any)

	user string // expect profile to match loadTestUsers[user]
}

/**********
* HELPERS *
**********/

var oneSec time.Duration = time.Unix(1, 0).Sub(time.Unix(0, 0))

// Verifies the basics of a loaded profile, i.e. that no errors happened and that
// the ID and name attributes are as expected.
func assertBasicInfo(t *testing.T, fn string, err error, got *Profile, expect user) {

	checkForErr(t, err, fn)

	if id := got.ID(); id != expect.id {

		t.Errorf("%s.ID() = %q; want %q", fn, id, expect.id)
	}
	if name := got.Name(); name != expect.name {

		t.Errorf("%s.Name() = %q; want %q", fn, name, expect.name)
	}
}

// Verifies that no error occurred, otherwise reports it.
// err is the error, fn is the function invocation which returned the error.
func checkForErr(t *testing.T, err error, fn string) {

	if err != nil {

		t.Logf("%s returned error: %s", fn, err)

		if _, tmr := err.(ErrTooManyRequests); tmr {

			t.Log("Likely problem: Test has been executed too frequently. Wait 1-10 minutes before retrying.")
		}

		t.FailNow()
	}
}

/*************
* TEST CASES *
*************/

// Test that the correct ID and case-corrected name gets loaded
func TestLoad(t *testing.T) {

	for n, expect := range loadTestUsers {

		got, err := Load(n)
		assertBasicInfo(t, fmt.Sprintf("Load(%q)", n), err, got, expect)
	}
}

// Test that the correct ID and case-corrected name gets loaded for past names
func TestLoadAtTime(t *testing.T) {

	for n, hist := range pastNames {

		expect := loadTestUsers[n]

		for _, p := range hist {

			tmfmt := p.until.Format(time.RFC3339)

			got, err := LoadAtTime(p.name, p.until)
			assertBasicInfo(t, fmt.Sprintf("LoadAtTime(%q, %s)", p.name, tmfmt), err, got, expect)
		}
	}
}

// No profile should be found for these usernames at these time instants.
func TestLoadAtTimeFailure(t *testing.T) {

	for n, hist := range pastNames {

		for _, p := range hist {

			p.until = p.until.Add(oneSec)
			tmfmt := p.until.Format(time.RFC3339)

			got, err := LoadAtTime(p.name, p.until)

			if _, nsp := err.(ErrNoSuchUser); !nsp {

				checkForErr(t, err, fmt.Sprintf("LoadAtTime(%q, %s)", p.name, tmfmt))

				t.Errorf("LoadAtTime(%q, %s) = %s; want ErrNoSuchUser error", n, tmfmt, got)
			}
		}
	}
}

// Verify that name histories are reported with expected stats and in ascending order.
// Also verify that LoadWithNameHistory actually preloads the profile's name history
//
// TODO: Supply test samples to verify ascending order
func TestLoadWithNameHistory(t *testing.T) {

	for n, hist := range pastNames {

		expect := loadTestUsers[n]
		fn := fmt.Sprintf("LoadWithNameHistory(%q)", expect.id)

		got, err := LoadWithNameHistory(expect.id)

		if got.NameHistory() == nil {

			t.Errorf("%s.NameHistory() = nil; should already be loaded", fn)
		}

		assertBasicInfo(t, fn, err, got, expect)

		// Test name history for correctness and ascending order
		h, err := got.LoadNameHistory()

		// Never ought to happen
		checkForErr(t, err, fn)

		var nameMatch, timeMatch bool
		for i, p := range hist {

			nameMatch = h[i].Name() == p.name
			if !nameMatch {

				t.Errorf("%s.LoadNameHistory()[%d].Name() = %q; want %q", fn, i, h[i].Name(), p.name)
				continue
			}

			timeMatch = h[i].Until() == p.until
			if !timeMatch {

				t.Errorf("%s.LoadNameHistory()[%d].Until() = %s; want %s",
					fn, i, h[i].Until().Format(time.RFC3339), p.until.Format(time.RFC3339))
			}
		}
	}
}

// Verify that:
// 1) LoadWithProperties successfully loads a profile with correct name and ID
// 2) LoadWithProperties preloads properties
// 3) Skins are handled successfully, a) present or b) not
// 4) Model are handled successfully, a) present or b) not (Alex/Steve)
// 5) Capes are handled successfully, b) present or b) not
//
// TODO: 3b, 4a
func TestLoadWithProperties(t *testing.T) {

	///////////////////////////
	// TEST 1, 2, 3a, 4b, 5b //
	///////////////////////////

	u1 := loadTestUsers[nergName]

	fn := fmt.Sprintf("LoadWithProperties(%q)", u1.id)
	p1, err := LoadWithProperties(u1.id)

	// 1) Correct name and ID, no errors
	assertBasicInfo(t, fn, err, p1, u1)

	// 2) Properties preloaded
	if p1.Properties() == nil {

		t.Errorf("%s.Properties() = nil; want preloaded", fn)
	}

	pp1, err := p1.LoadProperties()

	// Never ought to happen
	checkForErr(t, err, fn)

	// 3a) Skin set
	if _, ok := pp1.SkinURL(); !ok {

		t.Errorf("%s.Properties().SkinURL() = \"\"; want URL", fn)
	}

	// 4b) Model is Steve
	if m := pp1.Model(); m != Steve {

		t.Errorf("%s.Properties().Model() = %s; want %s", fn, m, Steve)
	}

	// 5b) No cape
	if c, ok := pp1.CapeURL(); ok {

		t.Errorf("%s.Properties().CapeURL() = %q; want %q", fn, c, "")
	}

	////////////////
	// TEST 2, 5a //
	////////////////

	fn = fmt.Sprintf("LoadWithProperties(%q)", capeID)

	p2, err := LoadWithProperties(capeID)
	checkForErr(t, err, fn)

	// 2 + 5a) Check for cape and blow up if not preloaded
	if p2.Properties().CapeURL == nil {

		t.Errorf("%s.Properties().CapeURL() = \"\"; want URL", fn)
	}
}

// Test that ErrTooManyRequests is returned if we exceed the rate limit.
// Allow the operation to succeed in case Mojang changes the rate limit.
func TestLoadWithPropertiesFailure(t *testing.T) {

	id := loadTestUsers[breeName].id

	for i := 0; i < 3; i++ {

		_, err := LoadWithProperties(id)
		if _, nsp := err.(ErrTooManyRequests); !nsp && err != nil {

			t.Fatalf("LoadWithProperties(%q) returned non-ErrTooManyRequests error: %s", id, err)
		}
	}
}

// Test that multiple users may be loaded (correctly) at once.
// Also test that empty and non-existing usernames are ignored
// and duplicates only are returned once.
func TestLoadMany(t *testing.T) {

	testUsers := []string{"", "I_DONT_ËXIST_ÆØÅ39"}

	for n, _ := range loadTestUsers {

		testUsers = append(testUsers, n)
	}

	fn := fmt.Sprintf("LoadMany(%#v)", testUsers)

	ps, err := LoadMany(testUsers...)
	checkForErr(t, err, fn)

	expect := loadManyTestRes

	// Verify that the correct number of profiles were returned
	if len(ps) != len(expect) {

		t.Errorf("len(%s) = %d; want %d", fn, len(ps), len(expect))
		t.Errorf("%s = %s; want %s", fn, ps, expect)

		t.FailNow()
	}

	// Verify that the loaded data was correct.
	for i := range ps {

		got := ps[i]
		assertBasicInfo(t, fn, nil, got, loadTestUsers[expect[i]])
	}
}

// Test that LoadMany actually succeeds at requesting LoadManyMaxSize profiles.
func TestLoadManyMax(t *testing.T) {

	var genUsers []string
	for i := 0; i < LoadManyMaxSize; i++ {

		genUsers = append(genUsers, "user"+string(i))
	}

	_, err := LoadMany(genUsers...)

	// Verify that no error occurred
	checkForErr(t, err, "LoadMany(<LoadManyMaxSize USERNAMES>)")
}
<commit_msg>Changed a Logf method call into a more semantically correct Errorf call.<commit_after>package profile_test

import (
	"fmt"
	"strings"
	"testing"
	"time"

	. "github.com/PhilipRasmussen/minecraft/profile"
)

/************
* TEST DATA *
************/

// 1. known account of author
const (
	nergName = "Nergalic"
	nergID   = "087cc153c3434ff7ac497de1569affa1"
	// Model is expected to be Steve
	// Cape is not expected
)

var nergHist = []pastName{{"GeneralSezuan", time.Unix(1423047705, 0)}}

// 2. known account of author
const (
	breeName = "BreeSakana"
	breeID   = "d9a5b542ce88442aaab38ec13e6c7773"
)

var breeHist = []pastName{}

// akronman1, the holder of the 1.000.000th-Minecraft-copy cape
// Profile may in theory be deleted at any moment...
//
// TODO: Substitute for a cape profile under author's control
const capeID = "d90b68bc81724329a047f1186dcd4336"

/******************
* TEST STRUCTURES *
******************/

// Username --> profile info
// Must maintain consistency with the other test structures
var loadTestUsers = map[string]user{
	nergName:                  {nergID, nergName},
	strings.ToUpper(nergName): {nergID, nergName}, // Casing should not matter
	breeName:                  {breeID, breeName},
}

// Should match loadTestUsers without duplicates
var loadManyTestRes = []string{breeName, nergName}

// Keys should match entries in loadTestUsers
var pastNames = map[string][]pastName{
	breeName: breeHist,
	nergName: nergHist,

	// TODO: Add profiles with longer name history to verify order
}

/*************
* TEST TYPES *
*************/

type user struct {
	id   string
	name string
}

type pastName struct {
	name  string
	until time.Time
}

type propertySet struct {

	// Do not check their values, simply check their presence
	skinURL, capeURL *bool // expected skin and cape (if any)

	model *Model // expected model (if any)

	user string // expect profile to match loadTestUsers[user]
}

/**********
* HELPERS *
**********/

var oneSec time.Duration = time.Unix(1, 0).Sub(time.Unix(0, 0))

// Verifies the basics of a loaded profile, i.e. that no errors happened and that
// the ID and name attributes are as expected.
func assertBasicInfo(t *testing.T, fn string, err error, got *Profile, expect user) {

	checkForErr(t, err, fn)

	if id := got.ID(); id != expect.id {

		t.Errorf("%s.ID() = %q; want %q", fn, id, expect.id)
	}
	if name := got.Name(); name != expect.name {

		t.Errorf("%s.Name() = %q; want %q", fn, name, expect.name)
	}
}

// Verifies that no error occurred, otherwise reports it.
// err is the error, fn is the function invocation which returned the error.
func checkForErr(t *testing.T, err error, fn string) {

	if err != nil {

		t.Errorf("%s returned error: %s", fn, err)

		if _, tmr := err.(ErrTooManyRequests); tmr {

			t.Log("Likely problem: Test has been executed too frequently. Wait 1-10 minutes before retrying.")
		}

		t.FailNow()
	}
}

/*************
* TEST CASES *
*************/

// Test that the correct ID and case-corrected name gets loaded
func TestLoad(t *testing.T) {

	for n, expect := range loadTestUsers {

		got, err := Load(n)
		assertBasicInfo(t, fmt.Sprintf("Load(%q)", n), err, got, expect)
	}
}

// Test that the correct ID and case-corrected name gets loaded for past names
func TestLoadAtTime(t *testing.T) {

	for n, hist := range pastNames {

		expect := loadTestUsers[n]

		for _, p := range hist {

			tmfmt := p.until.Format(time.RFC3339)

			got, err := LoadAtTime(p.name, p.until)
			assertBasicInfo(t, fmt.Sprintf("LoadAtTime(%q, %s)", p.name, tmfmt), err, got, expect)
		}
	}
}

// No profile should be found for these usernames at these time instants.
func TestLoadAtTimeFailure(t *testing.T) {

	for n, hist := range pastNames {

		for _, p := range hist {

			p.until = p.until.Add(oneSec)
			tmfmt := p.until.Format(time.RFC3339)

			got, err := LoadAtTime(p.name, p.until)

			if _, nsp := err.(ErrNoSuchUser); !nsp {

				checkForErr(t, err, fmt.Sprintf("LoadAtTime(%q, %s)", p.name, tmfmt))

				t.Errorf("LoadAtTime(%q, %s) = %s; want ErrNoSuchUser error", n, tmfmt, got)
			}
		}
	}
}

// Verify that name histories are reported with expected stats and in ascending order.
// Also verify that LoadWithNameHistory actually preloads the profile's name history
//
// TODO: Supply test samples to verify ascending order
func TestLoadWithNameHistory(t *testing.T) {

	for n, hist := range pastNames {

		expect := loadTestUsers[n]
		fn := fmt.Sprintf("LoadWithNameHistory(%q)", expect.id)

		got, err := LoadWithNameHistory(expect.id)

		if got.NameHistory() == nil {

			t.Errorf("%s.NameHistory() = nil; should already be loaded", fn)
		}

		assertBasicInfo(t, fn, err, got, expect)

		// Test name history for correctness and ascending order
		h, err := got.LoadNameHistory()

		// Never ought to happen
		checkForErr(t, err, fn)

		var nameMatch, timeMatch bool
		for i, p := range hist {

			nameMatch = h[i].Name() == p.name
			if !nameMatch {

				t.Errorf("%s.LoadNameHistory()[%d].Name() = %q; want %q", fn, i, h[i].Name(), p.name)
				continue
			}

			timeMatch = h[i].Until() == p.until
			if !timeMatch {

				t.Errorf("%s.LoadNameHistory()[%d].Until() = %s; want %s",
					fn, i, h[i].Until().Format(time.RFC3339), p.until.Format(time.RFC3339))
			}
		}
	}
}

// Verify that:
// 1) LoadWithProperties successfully loads a profile with correct name and ID
// 2) LoadWithProperties preloads properties
// 3) Skins are handled successfully, a) present or b) not
// 4) Model are handled successfully, a) present or b) not (Alex/Steve)
// 5) Capes are handled successfully, b) present or b) not
//
// TODO: 3b, 4a
func TestLoadWithProperties(t *testing.T) {

	///////////////////////////
	// TEST 1, 2, 3a, 4b, 5b //
	///////////////////////////

	u1 := loadTestUsers[nergName]

	fn := fmt.Sprintf("LoadWithProperties(%q)", u1.id)
	p1, err := LoadWithProperties(u1.id)

	// 1) Correct name and ID, no errors
	assertBasicInfo(t, fn, err, p1, u1)

	// 2) Properties preloaded
	if p1.Properties() == nil {

		t.Errorf("%s.Properties() = nil; want preloaded", fn)
	}

	pp1, err := p1.LoadProperties()

	// Never ought to happen
	checkForErr(t, err, fn)

	// 3a) Skin set
	if _, ok := pp1.SkinURL(); !ok {

		t.Errorf("%s.Properties().SkinURL() = \"\"; want URL", fn)
	}

	// 4b) Model is Steve
	if m := pp1.Model(); m != Steve {

		t.Errorf("%s.Properties().Model() = %s; want %s", fn, m, Steve)
	}

	// 5b) No cape
	if c, ok := pp1.CapeURL(); ok {

		t.Errorf("%s.Properties().CapeURL() = %q; want %q", fn, c, "")
	}

	////////////////
	// TEST 2, 5a //
	////////////////

	fn = fmt.Sprintf("LoadWithProperties(%q)", capeID)

	p2, err := LoadWithProperties(capeID)
	checkForErr(t, err, fn)

	// 2 + 5a) Check for cape and blow up if not preloaded
	if p2.Properties().CapeURL == nil {

		t.Errorf("%s.Properties().CapeURL() = \"\"; want URL", fn)
	}
}

// Test that ErrTooManyRequests is returned if we exceed the rate limit.
// Allow the operation to succeed in case Mojang changes the rate limit.
func TestLoadWithPropertiesFailure(t *testing.T) {

	id := loadTestUsers[breeName].id

	for i := 0; i < 3; i++ {

		_, err := LoadWithProperties(id)
		if _, nsp := err.(ErrTooManyRequests); !nsp && err != nil {

			t.Fatalf("LoadWithProperties(%q) returned non-ErrTooManyRequests error: %s", id, err)
		}
	}
}

// Test that multiple users may be loaded (correctly) at once.
// Also test that empty and non-existing usernames are ignored
// and duplicates only are returned once.
func TestLoadMany(t *testing.T) {

	testUsers := []string{"", "I_DONT_ËXIST_ÆØÅ39"}

	for n, _ := range loadTestUsers {

		testUsers = append(testUsers, n)
	}

	fn := fmt.Sprintf("LoadMany(%#v)", testUsers)

	ps, err := LoadMany(testUsers...)
	checkForErr(t, err, fn)

	expect := loadManyTestRes

	// Verify that the correct number of profiles were returned
	if len(ps) != len(expect) {

		t.Errorf("len(%s) = %d; want %d", fn, len(ps), len(expect))
		t.Errorf("%s = %s; want %s", fn, ps, expect)

		t.FailNow()
	}

	// Verify that the loaded data was correct.
	for i := range ps {

		got := ps[i]
		assertBasicInfo(t, fn, nil, got, loadTestUsers[expect[i]])
	}
}

// Test that LoadMany actually succeeds at requesting LoadManyMaxSize profiles.
func TestLoadManyMax(t *testing.T) {

	var genUsers []string
	for i := 0; i < LoadManyMaxSize; i++ {

		genUsers = append(genUsers, "user"+string(i))
	}

	_, err := LoadMany(genUsers...)

	// Verify that no error occurred
	checkForErr(t, err, "LoadMany(<LoadManyMaxSize USERNAMES>)")
}
<|endoftext|>


<commit_before>/*
 *    Copyright (C) 2015 Stefan Luecke
 *
 *    This program is free software: you can redistribute it and/or modify
 *    it under the terms of the GNU Affero General Public License as published
 *    by the Free Software Foundation, either version 3 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU Affero General Public License for more details.
 *
 *    You should have received a copy of the GNU Affero General Public License
 *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 *    Authors: Stefan Luecke <glaxx@glaxx.net>
 */

package backend

import (
	"crypto/rand"
	log "github.com/Sirupsen/logrus"
	"github.com/emicklei/go-restful"
	"gopkg.in/mgo.v2"
	mrand "math/rand"
	"net/http"
	"os"
)

const (
	ERROR_INVALID_ID    = "Error: Invalid ID"
	ERROR_STMT_PREPARE  = "Error: Statement prepare failed"
	ERROR_INVALID_INPUT = "Error: Invalid Input"
	ERROR_INTERNAL      = "Error: Internal Server Error"
	ERROR_INSERT        = "Error: DB Insert failed"
	ERROR_QUERY         = "Error: DB Query failed"
	PEPPER_SIZE         = 64
)

var db *mgo.Session

var uCol *mgo.Collection
var iCol *mgo.Collection
var ihCol *mgo.Collection
var pCol *mgo.Collection
var phCol *mgo.Collection

var pepper []byte

var idgen *idgenerator

func RegisterDatabase(s *mgo.Session, dbname string) {
	db = s
	uCol = s.DB(dbname).C("user")
	iCol = s.DB(dbname).C("item")
	ihCol = s.DB(dbname).C("item_history")
	pCol = s.DB(dbname).C("policy")
	phCol = s.DB(dbname).C("policy_history")
	idgen = NewIDGenerator(s.DB(dbname).C("counters"))
}

func ReadPepper(path string) {
	f, er := os.Open(path)
	if er != nil {
		err := er.(*os.PathError)
		log.WithFields(log.Fields{"Path": err.Path, "Op": err.Op}).Debug(err.Err)
		if err.Err.Error() == "no such file or directory" {
			log.Warn("Pepper file not found - creating ...")
			pepper = createPepper(path)
			return
		}
		log.Fatal(err)
	}
	defer f.Close()
	fi, er := f.Stat()
	if er != nil {
		log.Fatal(er)
	}
	if fi.Size() != PEPPER_SIZE {
		log.WithFields(log.Fields{"File Size": fi.Size(), "Expected Size": PEPPER_SIZE}).Fatal("Invalid pepper length - your file may be corrupt. Check your disk for errors.")
	}
	pepper = make([]byte, PEPPER_SIZE)
	bytes, er := f.Read(pepper)
	if er != nil || bytes != PEPPER_SIZE {
		log.WithFields(log.Fields{"Read": bytes, "Expected": PEPPER_SIZE}).Fatal(er)
	}
}

func createPepper(path string) []byte {
	res := make([]byte, PEPPER_SIZE)
	b, err := rand.Read(res)
	if err != nil || b != PEPPER_SIZE {
		log.WithFields(log.Fields{"Read": b, "Expected": PEPPER_SIZE}).Fatal(err)
	}
	f, err := os.Create(path)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	err = f.Chmod(0600)
	if err != nil {
		log.Fatal(err)
	}
	b, err = f.Write(res)
	if err != nil || b != PEPPER_SIZE {
		log.Fatal(err)
	}
	return res
}

func DebugLoggingFilter(rq *restful.Request, rs *restful.Response, ch *restful.FilterChain) {
	id := uint32(mrand.Int31())
	log.WithFields(log.Fields{
		"ID": id, "Path": rq.SelectedRoutePath()}).
		Debug("Got Request")

	log.WithFields(log.Fields{
		"ID": id, "PathParameters": rq.PathParameters()}).
		Debug()

	log.WithFields(log.Fields{
		"ID": id, "Method": rq.Request.Method}).
		Debug()

	log.WithFields(log.Fields{
		"ID": id, "Protocol": rq.Request.Proto}).
		Debug()

	log.WithFields(log.Fields{
		"ID": id, "Host": rq.Request.Header.Get("Host")}).
		Debug()

	log.WithFields(log.Fields{
		"ID": id, "Upgrade": rq.Request.Header.Get("Upgrade")}).
		Debug()

	log.WithFields(log.Fields{
		"ID": id, "User-Agent": rq.Request.Header.Get("User-Agent")}).
		Debug()

	log.WithFields(log.Fields{
		"ID": id, "Content-Length": rq.Request.ContentLength}).
		Debug()
}

func CloseIDGen() {
	idgen.StopIDGenerator()
}

func returnsInternalServerError(b *restful.RouteBuilder) {
	b.Returns(http.StatusInternalServerError, ERROR_INTERNAL, nil)
}

func returnsNotFound(b *restful.RouteBuilder) {
	b.Returns(http.StatusNotFound, ERROR_INVALID_ID, nil)
}

func returnsUpdateSuccessful(b *restful.RouteBuilder) {
	b.Returns(http.StatusOK, "Update successful", nil)
}

func returnsDeleteSuccessful(b *restful.RouteBuilder) {
	b.Returns(http.StatusOK, "Delete successful", nil)
}

func returnsBadRequest(b *restful.RouteBuilder) {
	b.Returns(http.StatusBadRequest, "Failed to parse input", nil)
}
<commit_msg>Process chain after logging<commit_after>/*
 *    Copyright (C) 2015 Stefan Luecke
 *
 *    This program is free software: you can redistribute it and/or modify
 *    it under the terms of the GNU Affero General Public License as published
 *    by the Free Software Foundation, either version 3 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU Affero General Public License for more details.
 *
 *    You should have received a copy of the GNU Affero General Public License
 *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 *    Authors: Stefan Luecke <glaxx@glaxx.net>
 */

package backend

import (
	"crypto/rand"
	log "github.com/Sirupsen/logrus"
	"github.com/emicklei/go-restful"
	"gopkg.in/mgo.v2"
	mrand "math/rand"
	"net/http"
	"os"
)

const (
	ERROR_INVALID_ID    = "Error: Invalid ID"
	ERROR_STMT_PREPARE  = "Error: Statement prepare failed"
	ERROR_INVALID_INPUT = "Error: Invalid Input"
	ERROR_INTERNAL      = "Error: Internal Server Error"
	ERROR_INSERT        = "Error: DB Insert failed"
	ERROR_QUERY         = "Error: DB Query failed"
	PEPPER_SIZE         = 64
)

var db *mgo.Session

var uCol *mgo.Collection
var iCol *mgo.Collection
var ihCol *mgo.Collection
var pCol *mgo.Collection
var phCol *mgo.Collection

var pepper []byte

var idgen *idgenerator

func RegisterDatabase(s *mgo.Session, dbname string) {
	db = s
	uCol = s.DB(dbname).C("user")
	iCol = s.DB(dbname).C("item")
	ihCol = s.DB(dbname).C("item_history")
	pCol = s.DB(dbname).C("policy")
	phCol = s.DB(dbname).C("policy_history")
	idgen = NewIDGenerator(s.DB(dbname).C("counters"))
}

func ReadPepper(path string) {
	f, er := os.Open(path)
	if er != nil {
		err := er.(*os.PathError)
		log.WithFields(log.Fields{"Path": err.Path, "Op": err.Op}).Debug(err.Err)
		if err.Err.Error() == "no such file or directory" {
			log.Warn("Pepper file not found - creating ...")
			pepper = createPepper(path)
			return
		}
		log.Fatal(err)
	}
	defer f.Close()
	fi, er := f.Stat()
	if er != nil {
		log.Fatal(er)
	}
	if fi.Size() != PEPPER_SIZE {
		log.WithFields(log.Fields{"File Size": fi.Size(), "Expected Size": PEPPER_SIZE}).Fatal("Invalid pepper length - your file may be corrupt. Check your disk for errors.")
	}
	pepper = make([]byte, PEPPER_SIZE)
	bytes, er := f.Read(pepper)
	if er != nil || bytes != PEPPER_SIZE {
		log.WithFields(log.Fields{"Read": bytes, "Expected": PEPPER_SIZE}).Fatal(er)
	}
}

func createPepper(path string) []byte {
	res := make([]byte, PEPPER_SIZE)
	b, err := rand.Read(res)
	if err != nil || b != PEPPER_SIZE {
		log.WithFields(log.Fields{"Read": b, "Expected": PEPPER_SIZE}).Fatal(err)
	}
	f, err := os.Create(path)
	if err != nil {
		log.Fatal(err)
	}
	defer f.Close()
	err = f.Chmod(0600)
	if err != nil {
		log.Fatal(err)
	}
	b, err = f.Write(res)
	if err != nil || b != PEPPER_SIZE {
		log.Fatal(err)
	}
	return res
}

func DebugLoggingFilter(rq *restful.Request, rs *restful.Response, ch *restful.FilterChain) {
	id := uint32(mrand.Int31())
	log.WithFields(log.Fields{
		"ID": id, "Path": rq.SelectedRoutePath()}).
		Debug("Got Request")

	log.WithFields(log.Fields{
		"ID": id, "PathParameters": rq.PathParameters()}).
		Debug()

	log.WithFields(log.Fields{
		"ID": id, "Method": rq.Request.Method}).
		Debug()

	log.WithFields(log.Fields{
		"ID": id, "Protocol": rq.Request.Proto}).
		Debug()

	log.WithFields(log.Fields{
		"ID": id, "Host": rq.Request.Header.Get("Host")}).
		Debug()

	log.WithFields(log.Fields{
		"ID": id, "Upgrade": rq.Request.Header.Get("Upgrade")}).
		Debug()

	log.WithFields(log.Fields{
		"ID": id, "User-Agent": rq.Request.Header.Get("User-Agent")}).
		Debug()

	log.WithFields(log.Fields{
		"ID": id, "Content-Length": rq.Request.ContentLength}).
		Debug()

	ch.ProcessFilter(rq, rs)
}

func CloseIDGen() {
	idgen.StopIDGenerator()
}

func returnsInternalServerError(b *restful.RouteBuilder) {
	b.Returns(http.StatusInternalServerError, ERROR_INTERNAL, nil)
}

func returnsNotFound(b *restful.RouteBuilder) {
	b.Returns(http.StatusNotFound, ERROR_INVALID_ID, nil)
}

func returnsUpdateSuccessful(b *restful.RouteBuilder) {
	b.Returns(http.StatusOK, "Update successful", nil)
}

func returnsDeleteSuccessful(b *restful.RouteBuilder) {
	b.Returns(http.StatusOK, "Delete successful", nil)
}

func returnsBadRequest(b *restful.RouteBuilder) {
	b.Returns(http.StatusBadRequest, "Failed to parse input", nil)
}
<|endoftext|>


<commit_before>package guard

import (
	"math"
	"math/rand"
	"sync/atomic"
	"time"
)

// BackoffStrategy is a backoff strategy.
type BackoffStrategy interface {
	// NextInterval returns the next interval.
	NextInterval() time.Duration

	// Reset creates the clone of the current strategy with an initialized state.
	Reset() BackoffStrategy
}

// ConstantBackoff creates BackoffStrategy with a constant interval.
// NextInterval() always returns given parameter d.
func ConstantBackoff(d time.Duration) BackoffStrategy {
	return &constantBackoff{d}
}

type constantBackoff struct {
	Interval time.Duration
}

func (c *constantBackoff) NextInterval() time.Duration {
	return c.Interval
}

func (c *constantBackoff) Reset() BackoffStrategy {
	return c
}

// NoBackoff creates BackoffStrategy without an interval.
// NextInterval() always returns 0.
func NoBackoff() BackoffStrategy {
	return noBackoff{}
}

type noBackoff struct{}

func (n noBackoff) NextInterval() time.Duration {
	return 0
}

func (n noBackoff) Reset() BackoffStrategy {
	return n
}

// ExponentialBackoff creates BackoffStrategy with an exponential backoff.
//
// Let N be a retry count of the process, the value of NextInterval(N) is calculated by following formula.
//
//  NextInterval(N) = BaseInterval(N) * [1-RandomizationFactor, 1+RandomizationFactor)
//  BaseInterval(N) = min(BaseInterval(N-1) * Multiplier, MaxInterval)
//  BaseInterval(1) = min(InitialInterval, MaxInterval)
//
// The default parameters.
//
//  InitialInterval:     200 (ms)
//  MaxInterval:         1 (min)
//  Multiplier:          2
//  RandomizationFactor: 0.2
//  Randomizer:          rand.New(rand.NewSource(time.Now().Unix()))
//
// Example intervals.
//
//  +----+----------------------+----------------------+
//  | N  | BaseInterval(N) (ms) | NextInterval(N) (ms) |
//  +----+----------------------+----------------------+
//  |  1 |                  200 | [160, 240)           |
//  |  2 |                  400 | [320, 480)           |
//  |  3 |                  800 | [640, 960)           |
//  |  4 |                 1600 | [1280, 1920)         |
//  |  5 |                 3200 | [2560, 3840)         |
//  |  6 |                 6400 | [5120, 7680)         |
//  |  7 |                12800 | [10240, 15360)       |
//  |  8 |                25600 | [20480, 30720)       |
//  |  9 |                51200 | [40960, 61440)       |
//  | 10 |                60000 | [48000, 72000)       |
//  | 11 |                60000 | [48000, 72000)       |
//  +----+----------------------+----------------------+
//
// Note: MaxInterval effects only the base interval.
// The actual interval may exceed MaxInterval depending on RandomizationFactor.
func ExponentialBackoff(options ...ExponentialBackoffOption) BackoffStrategy {
	e := &exponentialBackoff{
		initialInterval:     float64(200 * time.Millisecond),
		maxInterval:         float64(time.Minute),
		multiplier:          2,
		randomizationFactor: 0.2,
	}

	for _, o := range options {
		o(e)
	}

	if e.randomizer == nil {
		e.randomizer = rand.New(rand.NewSource(time.Now().Unix()))
	}
	e.baseInterval = math.Float64bits(e.initialInterval)

	return e
}

type exponentialBackoff struct {
	initialInterval     float64
	maxInterval         float64
	multiplier          float64
	randomizationFactor float64
	randomizer          Randomizer

	baseInterval uint64 // baseInterval actually represents float64. use uint64 for CompareAndSwap.
}

func (e *exponentialBackoff) NextInterval() time.Duration {
	var baseInterval float64
	for {
		old := atomic.LoadUint64(&e.baseInterval)
		baseInterval = math.Float64frombits(old)

		if baseInterval > e.maxInterval {
			baseInterval = e.maxInterval
		}
		if atomic.CompareAndSwapUint64(&e.baseInterval, old, math.Float64bits(baseInterval*e.multiplier)) {
			break
		}
	}

	rnd := (1 - e.randomizationFactor) + (2 * e.randomizationFactor * e.randomizer.Float64())
	nextBackoff := time.Duration(baseInterval * rnd)

	return nextBackoff
}

func (e *exponentialBackoff) Reset() BackoffStrategy {
	clone := *e
	clone.baseInterval = math.Float64bits(clone.initialInterval)
	return &clone
}

// ExponentialBackoffOption is the optional parameter for ExponentialBackoff.
type ExponentialBackoffOption func(*exponentialBackoff)

// WithInitialInterval set the initial interval of ExponentialBackoff.
func WithInitialInterval(d time.Duration) ExponentialBackoffOption {
	return ExponentialBackoffOption(func(e *exponentialBackoff) {
		e.initialInterval = float64(d)
	})
}

// WithMaxInterval set the maximum interval of ExponentialBackoff.
func WithMaxInterval(d time.Duration) ExponentialBackoffOption {
	return ExponentialBackoffOption(func(e *exponentialBackoff) {
		e.maxInterval = float64(d)
	})
}

// WithMultiplier set the multiplier of ExponentialBackoff.
func WithMultiplier(f float64) ExponentialBackoffOption {
	return ExponentialBackoffOption(func(e *exponentialBackoff) {
		e.multiplier = f
	})
}

// WithRandomizationFactor set the randomization factor of ExponentialBackoff.
func WithRandomizationFactor(f float64) ExponentialBackoffOption {
	return ExponentialBackoffOption(func(e *exponentialBackoff) {
		e.randomizationFactor = f
	})
}

// WithRandomizer set the randomizer of ExponentialBackoff.
func WithRandomizer(r Randomizer) ExponentialBackoffOption {
	return ExponentialBackoffOption(func(e *exponentialBackoff) {
		e.randomizer = r
	})
}
<commit_msg>Refactor NextInterval of exponential backoff<commit_after>package guard

import (
	"math"
	"math/rand"
	"sync/atomic"
	"time"
)

// BackoffStrategy is a backoff strategy.
type BackoffStrategy interface {
	// NextInterval returns the next interval.
	NextInterval() time.Duration

	// Reset creates the clone of the current strategy with an initialized state.
	Reset() BackoffStrategy
}

// ConstantBackoff creates BackoffStrategy with a constant interval.
// NextInterval() always returns given parameter d.
func ConstantBackoff(d time.Duration) BackoffStrategy {
	return &constantBackoff{d}
}

type constantBackoff struct {
	Interval time.Duration
}

func (c *constantBackoff) NextInterval() time.Duration {
	return c.Interval
}

func (c *constantBackoff) Reset() BackoffStrategy {
	return c
}

// NoBackoff creates BackoffStrategy without an interval.
// NextInterval() always returns 0.
func NoBackoff() BackoffStrategy {
	return noBackoff{}
}

type noBackoff struct{}

func (n noBackoff) NextInterval() time.Duration {
	return 0
}

func (n noBackoff) Reset() BackoffStrategy {
	return n
}

// ExponentialBackoff creates BackoffStrategy with an exponential backoff.
//
// Let N be a retry count of the process, the value of NextInterval(N) is calculated by following formula.
//
//  NextInterval(N) = BaseInterval(N) * [1-RandomizationFactor, 1+RandomizationFactor)
//  BaseInterval(N) = min(BaseInterval(N-1) * Multiplier, MaxInterval)
//  BaseInterval(1) = min(InitialInterval, MaxInterval)
//
// The default parameters.
//
//  InitialInterval:     200 (ms)
//  MaxInterval:         1 (min)
//  Multiplier:          2
//  RandomizationFactor: 0.2
//  Randomizer:          rand.New(rand.NewSource(time.Now().Unix()))
//
// Example intervals.
//
//  +----+----------------------+----------------------+
//  | N  | BaseInterval(N) (ms) | NextInterval(N) (ms) |
//  +----+----------------------+----------------------+
//  |  1 |                  200 | [160, 240)           |
//  |  2 |                  400 | [320, 480)           |
//  |  3 |                  800 | [640, 960)           |
//  |  4 |                 1600 | [1280, 1920)         |
//  |  5 |                 3200 | [2560, 3840)         |
//  |  6 |                 6400 | [5120, 7680)         |
//  |  7 |                12800 | [10240, 15360)       |
//  |  8 |                25600 | [20480, 30720)       |
//  |  9 |                51200 | [40960, 61440)       |
//  | 10 |                60000 | [48000, 72000)       |
//  | 11 |                60000 | [48000, 72000)       |
//  +----+----------------------+----------------------+
//
// Note: MaxInterval effects only the base interval.
// The actual interval may exceed MaxInterval depending on RandomizationFactor.
func ExponentialBackoff(options ...ExponentialBackoffOption) BackoffStrategy {
	e := &exponentialBackoff{
		initialInterval:     float64(200 * time.Millisecond),
		maxInterval:         float64(time.Minute),
		multiplier:          2,
		randomizationFactor: 0.2,
	}

	for _, o := range options {
		o(e)
	}

	if e.randomizer == nil {
		e.randomizer = rand.New(rand.NewSource(time.Now().Unix()))
	}
	e.baseInterval = math.Float64bits(e.initialInterval)

	return e
}

type exponentialBackoff struct {
	initialInterval     float64
	maxInterval         float64
	multiplier          float64
	randomizationFactor float64
	randomizer          Randomizer

	baseInterval uint64 // baseInterval actually represents float64. use uint64 for CompareAndSwap.
}

func (e *exponentialBackoff) NextInterval() time.Duration {
	baseInterval := e.BaseInterval()

	rnd := (1 - e.randomizationFactor) + (2 * e.randomizationFactor * e.randomizer.Float64())
	nextBackoff := time.Duration(baseInterval * rnd)

	return nextBackoff
}

func (e *exponentialBackoff) BaseInterval() float64 {
	for {
		old := atomic.LoadUint64(&e.baseInterval)
		baseInterval := math.Float64frombits(old)

		if baseInterval > e.maxInterval {
			baseInterval = e.maxInterval
		}
		if atomic.CompareAndSwapUint64(&e.baseInterval, old, math.Float64bits(baseInterval*e.multiplier)) {
			return baseInterval
		}
	}
}

func (e *exponentialBackoff) Reset() BackoffStrategy {
	clone := *e
	clone.baseInterval = math.Float64bits(clone.initialInterval)
	return &clone
}

// ExponentialBackoffOption is the optional parameter for ExponentialBackoff.
type ExponentialBackoffOption func(*exponentialBackoff)

// WithInitialInterval set the initial interval of ExponentialBackoff.
func WithInitialInterval(d time.Duration) ExponentialBackoffOption {
	return ExponentialBackoffOption(func(e *exponentialBackoff) {
		e.initialInterval = float64(d)
	})
}

// WithMaxInterval set the maximum interval of ExponentialBackoff.
func WithMaxInterval(d time.Duration) ExponentialBackoffOption {
	return ExponentialBackoffOption(func(e *exponentialBackoff) {
		e.maxInterval = float64(d)
	})
}

// WithMultiplier set the multiplier of ExponentialBackoff.
func WithMultiplier(f float64) ExponentialBackoffOption {
	return ExponentialBackoffOption(func(e *exponentialBackoff) {
		e.multiplier = f
	})
}

// WithRandomizationFactor set the randomization factor of ExponentialBackoff.
func WithRandomizationFactor(f float64) ExponentialBackoffOption {
	return ExponentialBackoffOption(func(e *exponentialBackoff) {
		e.randomizationFactor = f
	})
}

// WithRandomizer set the randomizer of ExponentialBackoff.
func WithRandomizer(r Randomizer) ExponentialBackoffOption {
	return ExponentialBackoffOption(func(e *exponentialBackoff) {
		e.randomizer = r
	})
}
<|endoftext|>


<commit_before>package model

import (
	"github.com/olivere/elastic"
)

// ConcurrentElastic is ElasticDB implementation of ConcurrentStorage.
type ConcurrentElastic struct {
	DB            *ElasticDB
	actionsCached map[string][]string
}

// Count returns number of Pageviews matching the filter defined by PageviewOptions.
func (pDB *ConcurrentElastic) Count(options AggregateOptions) (CountRowCollection, bool, error) {
	extras := make(map[string]elastic.Aggregation)

	search := pDB.DB.Client.Search().
		Index(TableConcurrents).
		Type("_doc").
		Size(0) // return no specific results

	search, err := pDB.DB.addSearchFilters(search, TableConcurrents, options)
	if err != nil {
		return nil, false, err
	}

	search, err = pDB.DB.addGroupBy(search, TableConcurrents, options, extras, nil)
	if err != nil {
		return nil, false, err
	}

	// get results
	result, err := search.Do(pDB.DB.Context)
	if err != nil {
		return nil, false, err
	}

	if len(options.GroupBy) == 0 {
		// extract simplified results (no aggregation)
		return CountRowCollection{
			CountRow{
				Count: int(result.Hits.TotalHits),
			},
		}, true, nil
	}

	// extract results
	return pDB.DB.countRowCollectionFromAggregations(result, options)
}
<commit_msg>Comment fixed<commit_after>package model

import (
	"github.com/olivere/elastic"
)

// ConcurrentElastic is ElasticDB implementation of ConcurrentStorage.
type ConcurrentElastic struct {
	DB            *ElasticDB
	actionsCached map[string][]string
}

// Count returns number of Concurrents matching the filter defined by AggregateOptions.
func (pDB *ConcurrentElastic) Count(options AggregateOptions) (CountRowCollection, bool, error) {
	extras := make(map[string]elastic.Aggregation)

	search := pDB.DB.Client.Search().
		Index(TableConcurrents).
		Type("_doc").
		Size(0) // return no specific results

	search, err := pDB.DB.addSearchFilters(search, TableConcurrents, options)
	if err != nil {
		return nil, false, err
	}

	search, err = pDB.DB.addGroupBy(search, TableConcurrents, options, extras, nil)
	if err != nil {
		return nil, false, err
	}

	// get results
	result, err := search.Do(pDB.DB.Context)
	if err != nil {
		return nil, false, err
	}

	if len(options.GroupBy) == 0 {
		// extract simplified results (no aggregation)
		return CountRowCollection{
			CountRow{
				Count: int(result.Hits.TotalHits),
			},
		}, true, nil
	}

	// extract results
	return pDB.DB.countRowCollectionFromAggregations(result, options)
}
<|endoftext|>


<commit_before>// Copyright 2015 The Chromium 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 internal contains code used internally by common/auth.
package internal

import (
	"encoding/json"
	"net/http"
	"strings"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/oauth2"

	"github.com/luci/luci-go/common/clock"
	"github.com/luci/luci-go/common/errors"
)

var (
	// ErrInsufficientAccess is returned by MintToken() if token can't be minted
	// for given OAuth scopes. For example, if GCE instance wasn't granted access
	// to requested scopes when it was created.
	ErrInsufficientAccess = errors.New("can't get access token for given scopes")

	// ErrBadRefreshToken is returned by RefreshToken if refresh token was revoked
	// or otherwise invalid. It means MintToken must be used to get a new refresh
	// token.
	ErrBadRefreshToken = errors.New("refresh_token is not valid")

	// ErrBadCredentials is returned by MintToken or RefreshToken if provided
	// offline credentials (like service account key) are invalid.
	ErrBadCredentials = errors.New("invalid service account credentials")
)

// TokenProvider knows how to mint new tokens, refresh existing ones, marshal
// and unmarshal tokens to byte buffers.
type TokenProvider interface {
	// RequiresInteraction is true if provider may start user interaction
	// in MintToken.
	RequiresInteraction() bool

	// CacheSeed is an optional byte string to use when constructing cache entry
	// name for access token cache. Different seeds will result in different
	// cache entries.
	CacheSeed() []byte

	// MintToken launches authentication flow (possibly interactive) and returns
	// a new refreshable token (or error). It must never return (nil, nil).
	MintToken() (*oauth2.Token, error)

	// RefreshToken takes existing token (probably expired, but not necessarily)
	// and returns a new refreshed token. It should never do any user interaction.
	// If a user interaction is required, a error should be returned instead.
	RefreshToken(*oauth2.Token) (*oauth2.Token, error)
}

// TransportFromContext returns http.RoundTripper buried inside the given
// context.
func TransportFromContext(ctx context.Context) http.RoundTripper {
	// When nil is passed to NewClient it skips all OAuth stuff and returns
	// client extracted from the context or http.DefaultClient.
	c := oauth2.NewClient(ctx, nil)
	if c == http.DefaultClient {
		return http.DefaultTransport
	}
	return c.Transport
}

// MarshalToken converts a token into byte buffer.
func MarshalToken(tok *oauth2.Token) ([]byte, error) {
	return json.MarshalIndent(&tokenOnDisk{
		AccessToken:  tok.AccessToken,
		TokenType:    tok.Type(),
		RefreshToken: tok.RefreshToken,
		ExpiresAtSec: tok.Expiry.Unix(),
	}, "", "\t")
}

// UnmarshalToken takes byte buffer produced by MarshalToken and returns
// original token.
func UnmarshalToken(data []byte) (*oauth2.Token, error) {
	onDisk := tokenOnDisk{}
	if err := json.Unmarshal(data, &onDisk); err != nil {
		return nil, err
	}
	return &oauth2.Token{
		AccessToken:  onDisk.AccessToken,
		TokenType:    onDisk.TokenType,
		RefreshToken: onDisk.RefreshToken,
		Expiry:       time.Unix(onDisk.ExpiresAtSec, 0),
	}, nil
}

// TokenExpiresIn returns True if the token is not valid or expires within given
// duration.
func TokenExpiresIn(ctx context.Context, t *oauth2.Token, lifetime time.Duration) bool {
	if t == nil || t.AccessToken == "" {
		return true
	}
	if t.Expiry.IsZero() {
		return false
	}
	expiry := t.Expiry.Add(-lifetime)
	return expiry.Before(clock.Now(ctx))
}

// EqualTokens returns true if both token object have same access token.
//
// 'nil' token corresponds to an empty access token.
func EqualTokens(a, b *oauth2.Token) bool {
	if a == b {
		return true
	}
	aTok := ""
	if a != nil {
		aTok = a.AccessToken
	}
	bTok := ""
	if b != nil {
		bTok = b.AccessToken
	}
	return aTok == bTok
}

// tokenOnDisk describes JSON produced by MarshalToken.
type tokenOnDisk struct {
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token,omitempty"`
	TokenType    string `json:"token_type,omitempty"`
	ExpiresAtSec int64  `json:"expires_at,omitempty"`
}

// isBadTokenError sniffs out HTTP 400 from token source errors.
func isBadTokenError(err error) bool {
	// See https://github.com/golang/oauth2/blob/master/internal/token.go.
	// Unfortunately, fmt.Errorf is used there, so there's no other way to
	// differentiate between bad tokens and transient errors.
	return err != nil && strings.Contains(err.Error(), "400 Bad Request")
}

// grabToken uses token source to create a new token.
//
// It recognizes transient errors.
func grabToken(src oauth2.TokenSource) (*oauth2.Token, error) {
	switch tok, err := src.Token(); {
	case isBadTokenError(err):
		return nil, err
	case err != nil:
		// More often than not errors here are transient (network connectivity
		// errors, HTTP 500 responses, etc). It is difficult to categorize them,
		// since oauth2 library uses fmt.Errorf(...) for errors. Retrying a fatal
		// error a bunch of times is not very bad, so pick safer approach and assume
		// any error is transient. Revoked refresh token or bad credentials (most
		// common source of fatal errors) is already handled above.
		return nil, errors.WrapTransient(err)
	default:
		return tok, nil
	}
}
<commit_msg>Add back 'version' and 'flavor' fields to token cache file.<commit_after>// Copyright 2015 The Chromium 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 internal contains code used internally by common/auth.
package internal

import (
	"encoding/json"
	"net/http"
	"strings"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/oauth2"

	"github.com/luci/luci-go/common/clock"
	"github.com/luci/luci-go/common/errors"
)

var (
	// ErrInsufficientAccess is returned by MintToken() if token can't be minted
	// for given OAuth scopes. For example, if GCE instance wasn't granted access
	// to requested scopes when it was created.
	ErrInsufficientAccess = errors.New("can't get access token for given scopes")

	// ErrBadRefreshToken is returned by RefreshToken if refresh token was revoked
	// or otherwise invalid. It means MintToken must be used to get a new refresh
	// token.
	ErrBadRefreshToken = errors.New("refresh_token is not valid")

	// ErrBadCredentials is returned by MintToken or RefreshToken if provided
	// offline credentials (like service account key) are invalid.
	ErrBadCredentials = errors.New("invalid service account credentials")
)

// TokenProvider knows how to mint new tokens, refresh existing ones, marshal
// and unmarshal tokens to byte buffers.
type TokenProvider interface {
	// RequiresInteraction is true if provider may start user interaction
	// in MintToken.
	RequiresInteraction() bool

	// CacheSeed is an optional byte string to use when constructing cache entry
	// name for access token cache. Different seeds will result in different
	// cache entries.
	CacheSeed() []byte

	// MintToken launches authentication flow (possibly interactive) and returns
	// a new refreshable token (or error). It must never return (nil, nil).
	MintToken() (*oauth2.Token, error)

	// RefreshToken takes existing token (probably expired, but not necessarily)
	// and returns a new refreshed token. It should never do any user interaction.
	// If a user interaction is required, a error should be returned instead.
	RefreshToken(*oauth2.Token) (*oauth2.Token, error)
}

// TransportFromContext returns http.RoundTripper buried inside the given
// context.
func TransportFromContext(ctx context.Context) http.RoundTripper {
	// When nil is passed to NewClient it skips all OAuth stuff and returns
	// client extracted from the context or http.DefaultClient.
	c := oauth2.NewClient(ctx, nil)
	if c == http.DefaultClient {
		return http.DefaultTransport
	}
	return c.Transport
}

// MarshalToken converts a token into byte buffer.
func MarshalToken(tok *oauth2.Token) ([]byte, error) {
	// TODO(vadimsh): Remove 'flavor' and 'version' when new code that doesn't use
	// them is deployed everywhere.
	flavor := "service_account"
	if tok.RefreshToken != "" {
		flavor = "user"
	}
	return json.MarshalIndent(&tokenOnDisk{
		Version:      "1",    // not actually used, exists for backward compatibility
		Flavor:       flavor, // same
		AccessToken:  tok.AccessToken,
		TokenType:    tok.Type(),
		RefreshToken: tok.RefreshToken,
		ExpiresAtSec: tok.Expiry.Unix(),
	}, "", "\t")
}

// UnmarshalToken takes byte buffer produced by MarshalToken and returns
// original token.
func UnmarshalToken(data []byte) (*oauth2.Token, error) {
	onDisk := tokenOnDisk{}
	if err := json.Unmarshal(data, &onDisk); err != nil {
		return nil, err
	}
	return &oauth2.Token{
		AccessToken:  onDisk.AccessToken,
		TokenType:    onDisk.TokenType,
		RefreshToken: onDisk.RefreshToken,
		Expiry:       time.Unix(onDisk.ExpiresAtSec, 0),
	}, nil
}

// TokenExpiresIn returns True if the token is not valid or expires within given
// duration.
func TokenExpiresIn(ctx context.Context, t *oauth2.Token, lifetime time.Duration) bool {
	if t == nil || t.AccessToken == "" {
		return true
	}
	if t.Expiry.IsZero() {
		return false
	}
	expiry := t.Expiry.Add(-lifetime)
	return expiry.Before(clock.Now(ctx))
}

// EqualTokens returns true if both token object have same access token.
//
// 'nil' token corresponds to an empty access token.
func EqualTokens(a, b *oauth2.Token) bool {
	if a == b {
		return true
	}
	aTok := ""
	if a != nil {
		aTok = a.AccessToken
	}
	bTok := ""
	if b != nil {
		bTok = b.AccessToken
	}
	return aTok == bTok
}

// tokenOnDisk describes JSON produced by MarshalToken.
type tokenOnDisk struct {
	Version      string `json:"version,omitempty"` // not actually used, exists for backward compatibility
	Flavor       string `json:"flavor,omitempty"`  // same
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token,omitempty"`
	TokenType    string `json:"token_type,omitempty"`
	ExpiresAtSec int64  `json:"expires_at,omitempty"`
}

// isBadTokenError sniffs out HTTP 400 from token source errors.
func isBadTokenError(err error) bool {
	// See https://github.com/golang/oauth2/blob/master/internal/token.go.
	// Unfortunately, fmt.Errorf is used there, so there's no other way to
	// differentiate between bad tokens and transient errors.
	return err != nil && strings.Contains(err.Error(), "400 Bad Request")
}

// grabToken uses token source to create a new token.
//
// It recognizes transient errors.
func grabToken(src oauth2.TokenSource) (*oauth2.Token, error) {
	switch tok, err := src.Token(); {
	case isBadTokenError(err):
		return nil, err
	case err != nil:
		// More often than not errors here are transient (network connectivity
		// errors, HTTP 500 responses, etc). It is difficult to categorize them,
		// since oauth2 library uses fmt.Errorf(...) for errors. Retrying a fatal
		// error a bunch of times is not very bad, so pick safer approach and assume
		// any error is transient. Revoked refresh token or bad credentials (most
		// common source of fatal errors) is already handled above.
		return nil, errors.WrapTransient(err)
	default:
		return tok, nil
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 The rkt 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.

// Disabled on kvm due to https://github.com/coreos/rkt/issues/3382
// +build !fly,!kvm

package main

import (
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/coreos/rkt/tests/testutils"
)

// TestAppSandboxSmoke is a basic smoke test for `rkt app` sandbox
// and related commands.
func TestAppSandboxSmoke(t *testing.T) {
	actionTimeout := 30 * time.Second
	imageName := "coreos.com/rkt-inspect/hello"
	appName := "hello-app"
	msg := "HelloFromAppInSandbox"

	aciHello := patchTestACI("rkt-inspect-hello.aci", "--name="+imageName, "--exec=/inspect --print-msg="+msg)
	defer os.Remove(aciHello)

	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	if err := os.Setenv("RKT_EXPERIMENT_APP", "true"); err != nil {
		panic(err)
	}
	defer os.Unsetenv("RKT_EXPERIMENT_APP")

	fetch := ctx.ExecCmd("fetch", "--insecure-options=image", aciHello)
	fetch.Env = append(fetch.Env, "RKT_EXPERIMENT_APP=true")
	t.Log("Running", fetch.Args)
	if out, err := fetch.CombinedOutput(); err != nil {
		t.Fatal(err, "output", out)
	}

	tmpDir := createTempDirOrPanic("rkt-test-cri-")
	uuidFile := filepath.Join(tmpDir, "uuid")
	defer os.RemoveAll(tmpDir)

	rkt := ctx.Cmd() + " app sandbox --uuid-file-save=" + uuidFile
	child := spawnOrFail(t, rkt)

	// wait for the sandbox to start
	podUUID, err := waitPodReady(ctx, t, uuidFile, actionTimeout)
	if err != nil {
		t.Fatal(err)
	}

	add := ctx.ExecCmd("app", "add", "--debug", podUUID, imageName, "--name="+appName)
	t.Log("Running", add.Args)
	if out, err := add.CombinedOutput(); err != nil {
		t.Fatal(err, "output", out)
	}

	start := ctx.ExecCmd("app", "start", "--debug", podUUID, "--app="+appName)
	t.Log("Running", start.Args)
	if out, err := start.CombinedOutput(); err != nil {
		t.Fatal(err, "output", out)
	}

	if err := expectTimeoutWithOutput(child, msg, actionTimeout); err != nil {
		t.Fatalf("Expected %q but not found: %v", msg, err)
	}

	remove := ctx.ExecCmd("app", "rm", "--debug", podUUID, "--app="+appName)
	t.Log("Running", remove.Args)
	if out, err := remove.CombinedOutput(); err != nil {
		t.Fatal(err, "output", out)
	}

	stop := ctx.ExecCmd("stop", podUUID)
	t.Log("Running", stop.Args)
	if out, err := stop.CombinedOutput(); err != nil {
		t.Fatal(err, "output", out)
	}

	waitOrFail(t, child, 0)
}
<commit_msg>tests/sandbox: add tests for start/remove and multiple apps<commit_after>// Copyright 2016 The rkt 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.

// Disabled on kvm due to https://github.com/coreos/rkt/issues/3382
// +build !fly,!kvm

package main

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/coreos/gexpect"
	"github.com/coreos/rkt/tests/testutils"
)

// TestAppSandboxOneApp is a basic test for `rkt app` sandbox.
// It starts the sandbox, adds one app, starts it, and removes it.
func TestAppSandboxAddStartRemove(t *testing.T) {
	testSandbox(t, func(ctx *testutils.RktRunCtx, child *gexpect.ExpectSubprocess, podUUID string) {
		actionTimeout := 30 * time.Second
		imageName := "coreos.com/rkt-inspect/hello"
		appName := "hello-app"
		msg := "HelloFromAppInSandbox"

		aciHello := patchTestACI("rkt-inspect-hello.aci", "--name="+imageName, "--exec=/inspect --print-msg="+msg)
		defer os.Remove(aciHello)

		testCmd{ctx.ExecCmd("fetch", "--insecure-options=image", aciHello)}.CombinedOutput(t)
		testCmd{ctx.ExecCmd("app", "add", "--debug", podUUID, imageName, "--name="+appName)}.CombinedOutput(t)
		testCmd{ctx.ExecCmd("app", "start", "--debug", podUUID, "--app="+appName)}.CombinedOutput(t)

		if err := expectTimeoutWithOutput(child, msg, actionTimeout); err != nil {
			t.Fatalf("Expected %q but not found: %v", msg, err)
		}

		testCmd{ctx.ExecCmd("app", "rm", "--debug", podUUID, "--app="+appName)}.CombinedOutput(t)

		out := testCmd{ctx.ExecCmd("app", "list", "--no-legend", podUUID)}.CombinedOutput(t)
		if out != "\n" {
			t.Errorf("unexpected output %q", out)
			return
		}
	})
}

// TestAppSandboxMultipleApps tests multiple apps in a sandbox:
// one that exits successfully, one that exits with an error, and one that keeps running.
func TestAppSandboxMultipleApps(t *testing.T) {
	testSandbox(t, func(ctx *testutils.RktRunCtx, child *gexpect.ExpectSubprocess, podUUID string) {
		actionTimeout := 60 * time.Second

		type app struct {
			name, image, exec, aci string
		}

		apps := []app{
			{
				name:  "winner",
				image: "coreos.com/rkt-inspect/success",
				exec:  "/inspect -print-msg=SUCCESS",
				aci:   "rkt-inspect-success.aci",
			},
			{
				name:  "loser",
				image: "coreos.com/rkt-inspect/fail",
				exec:  "/inspect -print-msg=FAILED -exit-code=12",
				aci:   "rkt-inspect-loser.aci",
			},
			{
				name:  "sleeper",
				image: "coreos.com/rkt-inspect/sleep",
				exec:  "/inspect -print-msg=SLEEP -sleep=120",
				aci:   "rkt-inspect-sleep.aci",
			},
		}

		// create, fetch, add, and start all apps in the sandbox
		for _, app := range apps {
			aci := patchTestACI(app.aci, "--name="+app.image, "--exec="+app.exec)
			defer os.Remove(aci)

			testCmd{ctx.ExecCmd("fetch", "--insecure-options=image", aci)}.CombinedOutput(t)
			testCmd{ctx.ExecCmd("app", "add", "--debug", podUUID, app.image, "--name="+app.name)}.CombinedOutput(t)
			testCmd{ctx.ExecCmd("app", "start", "--debug", podUUID, "--app="+app.name)}.CombinedOutput(t)
		}

		// check for app output messages
		for _, msg := range []string{
			"SUCCESS",
			"FAILED",
			"SLEEP",
		} {
			if err := expectTimeoutWithOutput(child, msg, actionTimeout); err != nil {
				t.Fatalf("Expected %q but not found: %v", msg, err)
			}
		}

		// total retry timeout: 10s
		r := retry{
			n: 20,
			t: 500 * time.Millisecond,
		}

		// assert `rkt app list` for the apps
		if err := r.Retry(func() error {
			got := testCmd{ctx.ExecCmd("app", "list", "--no-legend", podUUID)}.CombinedOutput(t)

			if strings.Contains(got, "winner\texited") &&
				strings.Contains(got, "loser\texited") &&
				strings.Contains(got, "sleeper\trunning") {
				return nil
			}

			return fmt.Errorf("unexpected result, got %q", got)
		}); err != nil {
			t.Error(err)
			return
		}

		// assert `rkt app status` for the apps
		for _, app := range []struct {
			name          string
			checkExitCode bool
			exitCode      int
			state         string
		}{
			{
				name:          "winner",
				checkExitCode: true,
				exitCode:      0,
				state:         "exited",
			},
			{
				name:          "loser",
				checkExitCode: true,
				exitCode:      12,
				state:         "exited",
			},
			{
				name:  "sleeper",
				state: "running",
			},
		} {
			if err := r.Retry(func() error {
				got := testCmd{ctx.ExecCmd("app", "status", podUUID, "--app="+app.name)}.CombinedOutput(t)
				ok := true

				if app.checkExitCode {
					ok = ok && strings.Contains(got, "exit_code="+strconv.Itoa(app.exitCode))
				}

				ok = ok && strings.Contains(got, "state="+app.state)

				if !ok {
					return fmt.Errorf("unexpected result, got %q", got)
				}

				return nil
			}); err != nil {
				t.Error(err)
				return
			}
		}

		// remove all apps
		for _, app := range apps {
			testCmd{ctx.ExecCmd("app", "rm", "--debug", podUUID, "--app="+app.name)}.CombinedOutput(t)
		}

		// assert empty `rkt app list`, no need for retrying,
		// as after removal no leftovers are expected to be present
		got := testCmd{ctx.ExecCmd("app", "list", "--no-legend", podUUID)}.CombinedOutput(t)
		if got != "\n" {
			t.Errorf("unexpected result, got %q", got)
			return
		}
	})
}

// TestAppSandboxRestart tests multiple apps in a sandbox and restarts one of them.
func TestAppSandboxRestart(t *testing.T) {
	testSandbox(t, func(ctx *testutils.RktRunCtx, child *gexpect.ExpectSubprocess, podUUID string) {
		type app struct {
			name, image, exec, aci string
		}

		apps := []app{
			{
				name:  "app1",
				image: "coreos.com/rkt-inspect/app1",
				exec:  "/inspect -sleep=120",
				aci:   "rkt-inspect-app1.aci",
			},
			{
				name:  "app2",
				image: "coreos.com/rkt-inspect/app1",
				exec:  "/inspect -sleep=120",
				aci:   "rkt-inspect-app1.aci",
			},
		}

		// create, fetch, add, and start all apps in the sandbox
		for _, app := range apps {
			aci := patchTestACI(app.aci, "--name="+app.image, "--exec="+app.exec)
			defer os.Remove(aci)

			testCmd{ctx.ExecCmd("fetch", "--insecure-options=image", aci)}.CombinedOutput(t)
			testCmd{ctx.ExecCmd("app", "add", "--debug", podUUID, app.image, "--name="+app.name)}.CombinedOutput(t)
			testCmd{ctx.ExecCmd("app", "start", "--debug", podUUID, "--app="+app.name)}.CombinedOutput(t)
		}

		// total retry timeout: 10s
		r := retry{
			n: 20,
			t: 500 * time.Millisecond,
		}

		// assert `rkt app list` for the apps
		if err := r.Retry(func() error {
			got := testCmd{ctx.ExecCmd("app", "list", "--no-legend", podUUID)}.CombinedOutput(t)

			if strings.Contains(got, "app1\trunning") &&
				strings.Contains(got, "app2\trunning") {
				return nil
			}

			return fmt.Errorf("unexpected result, got %q", got)
		}); err != nil {
			t.Error(err)
			return
		}

		assertStatus := func(name, status string) error {
			return r.Retry(func() error {
				got := testCmd{ctx.ExecCmd("app", "status", podUUID, "--app="+name)}.CombinedOutput(t)

				if !strings.Contains(got, status) {
					return fmt.Errorf("unexpected result, got %q", got)
				}

				return nil
			})
		}

		// assert `rkt app status` for the apps
		for _, app := range apps {
			if err := assertStatus(app.name, "state=running"); err != nil {
				t.Error(err)
				return
			}
		}

		// stop app1
		testCmd{ctx.ExecCmd("app", "stop", podUUID, "--app=app1")}.CombinedOutput(t)

		// assert `rkt app status` for the apps
		for _, app := range []struct {
			name   string
			status string
		}{
			{
				name:   "app1",
				status: "state=exited",
			},
			{
				name:   "app2",
				status: "state=running",
			},
		} {
			if err := assertStatus(app.name, app.status); err != nil {
				t.Error(err)
				return
			}
		}

		// assert `rkt app list` for the apps
		if err := r.Retry(func() error {
			got := testCmd{ctx.ExecCmd("app", "list", "--no-legend", podUUID)}.CombinedOutput(t)

			if strings.Contains(got, "app1\texited") &&
				strings.Contains(got, "app2\trunning") {
				return nil
			}

			return fmt.Errorf("unexpected result, got %q", got)
		}); err != nil {
			t.Error(err)
			return
		}

		// start app1
		testCmd{ctx.ExecCmd("app", "start", podUUID, "--app=app1")}.CombinedOutput(t)

		// assert `rkt app status` for the apps
		for _, app := range []struct {
			name   string
			status string
		}{
			{
				name:   "app1",
				status: "state=running",
			},
			{
				name:   "app2",
				status: "state=running",
			},
		} {
			if err := assertStatus(app.name, app.status); err != nil {
				t.Error(err)
				return
			}
		}

		// assert `rkt app list` for the apps
		if err := r.Retry(func() error {
			got := testCmd{ctx.ExecCmd("app", "list", "--no-legend", podUUID)}.CombinedOutput(t)

			if strings.Contains(got, "app1\trunning") &&
				strings.Contains(got, "app2\trunning") {
				return nil
			}

			return fmt.Errorf("unexpected result, got %q", got)
		}); err != nil {
			t.Error(err)
			return
		}
	})
}

func testSandbox(t *testing.T, testFunc func(*testutils.RktRunCtx, *gexpect.ExpectSubprocess, string)) {
	if err := os.Setenv("RKT_EXPERIMENT_APP", "true"); err != nil {
		panic(err)
	}
	defer os.Unsetenv("RKT_EXPERIMENT_APP")

	tmpDir := createTempDirOrPanic("rkt-test-cri-")
	uuidFile := filepath.Join(tmpDir, "uuid")
	defer os.RemoveAll(tmpDir)

	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	rkt := ctx.Cmd() + " app sandbox --uuid-file-save=" + uuidFile
	child := spawnOrFail(t, rkt)

	// wait for the sandbox to start
	podUUID, err := waitPodReady(ctx, t, uuidFile, 30*time.Second)
	if err != nil {
		t.Fatal(err)
	}

	testFunc(ctx, child, podUUID)

	// assert that the pod is still running
	got := testCmd{ctx.ExecCmd("status", podUUID)}.CombinedOutput(t)
	if !strings.Contains(got, "state=running") {
		t.Errorf("unexpected result, got %q", got)
		return
	}

	testCmd{ctx.ExecCmd("stop", podUUID)}.CombinedOutput(t)

	waitOrFail(t, child, 0)
}

// retry is the struct that represents retrying function calls.
type retry struct {
	n int
	t time.Duration
}

// Retry retries the given function f n times with a delay t between invocations
// until no error is retrurned from f or n is exceeded.
// The last occured error is returned.
func (r retry) Retry(f func() error) error {
	var err error

	for i := 0; i < r.n; i++ {
		err = f()
		if err == nil {
			return nil
		}
		time.Sleep(r.t)
	}

	return err
}

type testCmd struct {
	*exec.Cmd
}

func (c testCmd) CombinedOutput(t *testing.T) string {
	t.Log("Running", c.Args)
	out, err := c.Cmd.CombinedOutput()

	if err != nil {
		t.Fatal(err, "output", out)
	}

	return string(out)
}
<|endoftext|>