<commit_before>package vangoh

import (
	"bytes"
	"crypto"
	"crypto/hmac"
	_ "crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"errors"
	"hash"
	"net/http"
	"regexp"
	"sort"
	"strings"
	"time"
	"unsafe"
)

// Expected regex format of the Authorization signature.
//
// An authorization signature consists of three parts:
//
// 	Authorization: [ORG] [KEY]:[HMAC_SIGNATURE]
//
// The first component is an organizational tag, which must consist of at least
// one character, and has a valid character set of alphanumeric characters and
// underscores.
//
// This should be followed by a single space, and then the key, which also must
// consist of one or more alphanumeric characters and/or underscores.
//
// The key must be followed by a single colon ':' character, and then the
// signature, encoded in Base64 (valid characters being all alphanumeric, plus
// "+", forward slash "/", and equals sign "=" as padding on the end if
// needed.)
//
// Any leading or trailing whitespace around the header will be trimmed before
// validation.
const AuthRegex = "^[A-Za-z0-9_]+ [A-Za-z0-9_/+]+:" +
	"(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$"

// Newline character, definited in unicode, to avoid platform dependence.
const newline = "\u000A"

// Vangoh is an object that forms the primary point of configuration of the
// middleware HMAC handler. It allows for the configuration of the hashing
// function to use, the headers (specified as regexes) to be included in the
// computed signature, and the mapping between organization tags and the secret
// key providers associated with them.
type Vangoh struct {
	// Indicates if there is one global provider for this HMAC checker. If set to
	// a non-nil pointer to a secretProvider, Vangoh will authenticate all
	// requests against this provider regardless of organization specified in the
	// request's Authorization header.
	singleProvider *secretProvider

	// A map between org tags, as used in the Authentication section, with the
	// secretProvider that provides the identities for that org.
	providersByOrg map[string]secretProvider

	// The hashing function to be used when computing the HMAC hashes.  Common
	// algorithms for HMAC include SHA1, SHA256, and MD5, but any object that
	// implements hash.Hash should work.
	algorithm func() hash.Hash

	// Specifies which headers should be used in computing the HMAC signature for
	// each request.  It is common to have an application-wide prefix for headers
	// to be used, i.e. X-Aur-Meta-User or X-Aur-Locale.  This could be
	// represented with the include header "^X-Aur-".
	includedHeaders map[string]*regexp.Regexp

	// The maximum amount of time that can have passed between the time a request
	// was signed and the time the request was received by the server.
	maxTimeSkew time.Duration

	// When true, Handler() includes specific error details in the response when
	// denying incorrectly-authenticated requests.
	debug bool
}

// Creates a new Vangoh instance with no secret providers.
func New() *Vangoh {
	return &Vangoh{
		singleProvider:  nil,
		providersByOrg:  make(map[string]secretProvider),
		algorithm:       crypto.SHA256.New,
		includedHeaders: make(map[string]*regexp.Regexp),
		maxTimeSkew:     time.Minute * 15,
		debug:           false,
	}
}

// Creates a new Vangoh instance that supports a single
// secretProvider. Attempting to add providers with AddProvider will fail with an
// error.
func NewSingleProvider(provider secretProvider) *Vangoh {
	vg := New()
	vg.singleProvider = &provider
	return vg
}

/*
AddProvider sets the secret provider of a specific organization. If the Vangoh
instance was created to use a single provider for all requests, regardless of
organization tag, calling AddProvider will fail and return an error. If the
organization already has a provider, calling AddProvider will fail and return
an error.

By supporting different providers based on org tags, there is the ability to
configure authentication sources based on user type or purpose. For instance,
if an endpoint is going to be used by both a small set of internal services as
well as external users, you could create a different provider for each, as
demonstrated below.

Example:
	func main() {
		// Create provider for internal services credentials (not included with Vangoh).
		internalProvider := providers.NewInMemoryProvider(...)
		// Create provider for normal user credentials (not included with Vangoh).
		userProvider := providers.NewDatabaseProvider(...)

		vg := vangoh.New()
		_ = vg.AddProvider("INT", internalProvider)
		_ = vg.AddProvider("API", userProvider)

		// ...
	}

In this example, any connections made with the authorization header "INT
[userID]:[signature]" will be authenticated against `internalProvider`, and
connections with the header "API [userID]:[signature]" will be authenticated
against `userProvider`.
*/
func (vg *Vangoh) AddProvider(org string, skp secretProvider) error {
	if vg.singleProvider != nil {
		return errors.New("cannot add a provider when created for a single provider")
	}
	if _, ok := vg.providersByOrg[org]; ok {
		return errors.New("cannot add more than one keyProvider for the same org tag")
	}
	vg.providersByOrg[org] = skp
	return nil
}

func (vg *Vangoh) SetAlgorithm(algorithm func() hash.Hash) {
	vg.algorithm = algorithm
}

func (vg *Vangoh) SetDebug(debug bool) {
	vg.debug = debug
}

/*
IncludeHeader specifies additional headers to include in the construction of
the HMAC signature body for a request.

Given a regex, any non-canonical (e.g. "X-Aur", not "x-aur") headers that match the
regex will be included.

For instance, to match all headers beginning with "X-Aur-", we could include
the header regex "X-Aur-.*".  It is important to note that this funcationality
uses traditional, non-POSIX regular expressions, and will add anchoring to the
provided regex if it is not included.

This means that the regex "X-Aur" will only match headers with key "X-Aur"
exactly.  In order to do prefix matching you must add a wildcard match after,
i.e. "X-Aur.*"
*/
func (vg *Vangoh) IncludeHeader(headerRegex string) error {
	var regexBuf bytes.Buffer
	if !strings.HasPrefix(headerRegex, "^") {
		regexBuf.WriteString("^")
	}
	regexBuf.WriteString(headerRegex)
	if !strings.HasSuffix(headerRegex, "$") {
		regexBuf.WriteString("$")
	}

	regex := regexBuf.String()
	compiled, err := regexp.Compile(regex)
	if err != nil {
		return err
	}
	vg.includedHeaders[regex] = compiled
	return nil
}

/*
SetMaxTimeSkew sets the maximum allowable duration between the date and time specified
in the Date header and the server time when the response is processed.  If the date in
the header exceeds the duration Vangoh will respond to the request with a HTTP status 403 Forbidden.

To match the behavior of AWS (15 minute skew window):
	vg := vangoh.New()
	vg.SetMaxTimeSkew(time.Minute * 15)

When checking the date header, Vangoh follows the precedent of RFC 2616,
accepting dates in any of the following formats:
	ANSIC    = "Mon Jan _2 15:04:05 2006"
	RFC822   = "02 Jan 06 15:04 MST"
	RFC822Z  = "02 Jan 06 15:04 -0700"
	RFC850   = "Monday, 02-Jan-06 15:04:05 MST"
	RFC1123  = "Mon, 02 Jan 2006 15:04:05 MST"
	RFC1123Z = "Mon, 02 Jan 2006 15:04:05 -0700"
*/
func (vg *Vangoh) SetMaxTimeSkew(timeSkew time.Duration) {
	vg.maxTimeSkew = timeSkew
}

// Checks a request for proper authentication details, returning the relevent
// error if the request fails this check or nil if the request passes.
func (vg *Vangoh) AuthenticateRequest(r *http.Request) *AuthenticationError {
	// Parse the ORG, KEY, and SIGNATURE out of the Authorization header.
	authHeader := strings.TrimSpace(r.Header.Get("Authorization"))
	if authHeader == "" {
		return ErrorAuthHeaderMissing
	}
	match, err := regexp.Match(AuthRegex, []byte(authHeader))
	if err != nil || !match {
		return ErrorAuthHeaderMalformed
	}
	orgSplit := strings.Split(authHeader, " ")
	org := orgSplit[0]
	keySplit := strings.Split(orgSplit[1], ":")
	key := []byte(keySplit[0])
	actualSignatureB64 := keySplit[1]

	// Check that the request was made in the acceptable window.
	dateHeader := strings.TrimSpace(r.Header.Get("Date"))
	if dateHeader == "" {
		return ErrorDateHeaderMissing
	}
	date, err := multiFormatDateParse(
		// TODO: break out into const.
		[]string{
			time.RFC822,
			time.RFC822Z,
			time.RFC850,
			time.ANSIC,
			time.RFC1123,
			time.RFC1123Z},
		dateHeader)
	if err != nil {
		return ErrorDateHeaderMalformed
	}
	present := clock.Now()
	if present.Sub(date) > vg.maxTimeSkew || date.Sub(present) > vg.maxTimeSkew {
		return ErrorDateHeaderTooSkewed
	}

	// Load the secret key from the appropriate key provider, given the ID from
	// the Authorization header.
	var provider secretProvider
	if vg.singleProvider != nil {
		provider = *vg.singleProvider
	} else {
		var exists bool
		provider, exists = vg.providersByOrg[org]
		if !exists {
			return ErrorAuthOrgUnknown
		}
	}

	var voidPtr unsafe.Pointer = nil
	var secret []byte

	switch provider := provider.(type) {
	case SecretProviderWithCallback:
		secret, err = provider.GetSecret(key, &voidPtr)
	case SecretProvider:
		secret, err = provider.GetSecret(key)
	}
	if err != nil {
		return ErrorInProviderKeyLookup
	}
	if secret == nil {
		return ErrorSecretNotFound
	}

	// Calculate the b64 signature and compare against the one sent by the client.
	expectedSignature := vg.ConstructSignature(r, secret)
	expectedSignatureB64 := base64.StdEncoding.EncodeToString(expectedSignature)
	if subtle.ConstantTimeCompare([]byte(expectedSignatureB64), []byte(actualSignatureB64)) != 1 {
		return ErrorHMACSignatureMismatch
	}

	switch provider := provider.(type) {
	case SecretProviderWithCallback:
		if voidPtr != nil {
			provider.SuccessCallback(r, &voidPtr)
		} else {
			provider.SuccessCallback(r, nil)
		}
	}
	// If we have made it this far, authentication is successful.
	return nil
}

func (vg *Vangoh) ConstructSignature(r *http.Request, secret []byte) []byte {
	signingString := vg.CreateSigningString(r)
	mac := hmac.New(vg.algorithm, secret)
	mac.Write([]byte(signingString))
	return mac.Sum(nil)
}

func multiFormatDateParse(formats []string, dateStr string) (time.Time, error) {
	for index := range formats {
		if date, err := time.Parse(formats[index], dateStr); err == nil {
			return date, nil
		}
	}
	return time.Time{}, errors.New("Date does not match any valid format")
}

/*
CreateSigningString creates the string used for signature generation, in accordance with
the specifications as laid out in the package documentation. Refer there for more detail,
or to the Amazon Signature V2 documentation: http://docs.aws.amazon.com/AmazonS3/latest/dev/RESTAuthentication.html.
*/
func (vg *Vangoh) CreateSigningString(r *http.Request) string {
	var buffer bytes.Buffer

	buffer.WriteString(r.Method)
	buffer.WriteString(newline)

	buffer.WriteString(r.Header.Get("Content-MD5"))
	buffer.WriteString(newline)

	buffer.WriteString(r.Header.Get("Content-Type"))
	buffer.WriteString(newline)

	buffer.WriteString(r.Header.Get("Date"))
	buffer.WriteString(newline)

	customHeaders := vg.createHeadersString(r)
	buffer.WriteString(customHeaders)

	buffer.WriteString(r.URL.Path)

	return buffer.String()
}

// Create the canonicalized header string part of a request's signature body.
func (vg *Vangoh) createHeadersString(r *http.Request) string {
	if len(vg.includedHeaders) == 0 {
		return ""
	}

	// For each defined regex, determine the set of headers that match. Repeat
	// for all regexes, without duplication, to get the final set of custom
	// headers to use.
	var sanitizedHeaders = make(map[string][]string)

	for _, compiledRegex := range vg.includedHeaders {
		for header := range r.Header {
			lowerHeader := strings.ToLower(header)
			if _, found := sanitizedHeaders[lowerHeader]; found {
				continue
			}
			if compiledRegex.MatchString(header) {
				sanitizedHeaders[lowerHeader] = r.Header[header]
			}
		}
	}

	var orderedHeaders []string
	for header := range sanitizedHeaders {
		orderedHeaders = append(orderedHeaders, header)
	}
	sort.Strings(orderedHeaders)

	// At this point sanitized contains all the headers to be included in the
	// hash. Now we need to retrieve their values, and sanitize them
	// appropriately.
	var buffer bytes.Buffer
	for header := range orderedHeaders {
		buffer.WriteString(orderedHeaders[header])
		buffer.WriteString(":")

		var sanitizedValues []string
		for i := range sanitizedHeaders[orderedHeaders[header]] {
			str := sanitizedHeaders[orderedHeaders[header]][i]
			str = strings.TrimSpace(str)
			str = strings.Replace(str, "\n", "", -1)
			sanitizedValues = append(sanitizedValues, str)
		}

		// Note that sanitizedValues are unsorted here - the order that they are
		// specified in the header will affect the hash result. This conforms with
		// the standard set by AWS, though it may be more reliable to add this
		// sorting in at some point.
		for i := range sanitizedValues {
			buffer.WriteString(sanitizedValues[i])
			if i < (len(sanitizedValues) - 1) {
				buffer.WriteString(",")
			}
		}

		buffer.WriteString(newline)
	}

	return buffer.String()
}
<commit_msg>Fix typo.<commit_after>package vangoh

import (
	"bytes"
	"crypto"
	"crypto/hmac"
	_ "crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"errors"
	"hash"
	"net/http"
	"regexp"
	"sort"
	"strings"
	"time"
	"unsafe"
)

// Expected regex format of the Authorization signature.
//
// An authorization signature consists of three parts:
//
// 	Authorization: [ORG] [KEY]:[HMAC_SIGNATURE]
//
// The first component is an organizational tag, which must consist of at least
// one character, and has a valid character set of alphanumeric characters and
// underscores.
//
// This should be followed by a single space, and then the key, which also must
// consist of one or more alphanumeric characters and/or underscores.
//
// The key must be followed by a single colon ':' character, and then the
// signature, encoded in Base64 (valid characters being all alphanumeric, plus
// "+", forward slash "/", and equals sign "=" as padding on the end if
// needed.)
//
// Any leading or trailing whitespace around the header will be trimmed before
// validation.
const AuthRegex = "^[A-Za-z0-9_]+ [A-Za-z0-9_/+]+:" +
	"(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$"

// Newline character, defined in unicode to avoid platform dependence.
const newline = "\u000A"

// Vangoh is an object that forms the primary point of configuration of the
// middleware HMAC handler. It allows for the configuration of the hashing
// function to use, the headers (specified as regexes) to be included in the
// computed signature, and the mapping between organization tags and the secret
// key providers associated with them.
type Vangoh struct {
	// Indicates if there is one global provider for this HMAC checker. If set to
	// a non-nil pointer to a secretProvider, Vangoh will authenticate all
	// requests against this provider regardless of organization specified in the
	// request's Authorization header.
	singleProvider *secretProvider

	// A map between org tags, as used in the Authentication section, with the
	// secretProvider that provides the identities for that org.
	providersByOrg map[string]secretProvider

	// The hashing function to be used when computing the HMAC hashes.  Common
	// algorithms for HMAC include SHA1, SHA256, and MD5, but any object that
	// implements hash.Hash should work.
	algorithm func() hash.Hash

	// Specifies which headers should be used in computing the HMAC signature for
	// each request.  It is common to have an application-wide prefix for headers
	// to be used, i.e. X-Aur-Meta-User or X-Aur-Locale.  This could be
	// represented with the include header "^X-Aur-".
	includedHeaders map[string]*regexp.Regexp

	// The maximum amount of time that can have passed between the time a request
	// was signed and the time the request was received by the server.
	maxTimeSkew time.Duration

	// When true, Handler() includes specific error details in the response when
	// denying incorrectly-authenticated requests.
	debug bool
}

// Creates a new Vangoh instance with no secret providers.
func New() *Vangoh {
	return &Vangoh{
		singleProvider:  nil,
		providersByOrg:  make(map[string]secretProvider),
		algorithm:       crypto.SHA256.New,
		includedHeaders: make(map[string]*regexp.Regexp),
		maxTimeSkew:     time.Minute * 15,
		debug:           false,
	}
}

// Creates a new Vangoh instance that supports a single
// secretProvider. Attempting to add providers with AddProvider will fail with an
// error.
func NewSingleProvider(provider secretProvider) *Vangoh {
	vg := New()
	vg.singleProvider = &provider
	return vg
}

/*
AddProvider sets the secret provider of a specific organization. If the Vangoh
instance was created to use a single provider for all requests, regardless of
organization tag, calling AddProvider will fail and return an error. If the
organization already has a provider, calling AddProvider will fail and return
an error.

By supporting different providers based on org tags, there is the ability to
configure authentication sources based on user type or purpose. For instance,
if an endpoint is going to be used by both a small set of internal services as
well as external users, you could create a different provider for each, as
demonstrated below.

Example:
	func main() {
		// Create provider for internal services credentials (not included with Vangoh).
		internalProvider := providers.NewInMemoryProvider(...)
		// Create provider for normal user credentials (not included with Vangoh).
		userProvider := providers.NewDatabaseProvider(...)

		vg := vangoh.New()
		_ = vg.AddProvider("INT", internalProvider)
		_ = vg.AddProvider("API", userProvider)

		// ...
	}

In this example, any connections made with the authorization header "INT
[userID]:[signature]" will be authenticated against `internalProvider`, and
connections with the header "API [userID]:[signature]" will be authenticated
against `userProvider`.
*/
func (vg *Vangoh) AddProvider(org string, skp secretProvider) error {
	if vg.singleProvider != nil {
		return errors.New("cannot add a provider when created for a single provider")
	}
	if _, ok := vg.providersByOrg[org]; ok {
		return errors.New("cannot add more than one keyProvider for the same org tag")
	}
	vg.providersByOrg[org] = skp
	return nil
}

func (vg *Vangoh) SetAlgorithm(algorithm func() hash.Hash) {
	vg.algorithm = algorithm
}

func (vg *Vangoh) SetDebug(debug bool) {
	vg.debug = debug
}

/*
IncludeHeader specifies additional headers to include in the construction of
the HMAC signature body for a request.

Given a regex, any non-canonical (e.g. "X-Aur", not "x-aur") headers that match the
regex will be included.

For instance, to match all headers beginning with "X-Aur-", we could include
the header regex "X-Aur-.*".  It is important to note that this funcationality
uses traditional, non-POSIX regular expressions, and will add anchoring to the
provided regex if it is not included.

This means that the regex "X-Aur" will only match headers with key "X-Aur"
exactly.  In order to do prefix matching you must add a wildcard match after,
i.e. "X-Aur.*"
*/
func (vg *Vangoh) IncludeHeader(headerRegex string) error {
	var regexBuf bytes.Buffer
	if !strings.HasPrefix(headerRegex, "^") {
		regexBuf.WriteString("^")
	}
	regexBuf.WriteString(headerRegex)
	if !strings.HasSuffix(headerRegex, "$") {
		regexBuf.WriteString("$")
	}

	regex := regexBuf.String()
	compiled, err := regexp.Compile(regex)
	if err != nil {
		return err
	}
	vg.includedHeaders[regex] = compiled
	return nil
}

/*
SetMaxTimeSkew sets the maximum allowable duration between the date and time specified
in the Date header and the server time when the response is processed.  If the date in
the header exceeds the duration Vangoh will respond to the request with a HTTP status 403 Forbidden.

To match the behavior of AWS (15 minute skew window):
	vg := vangoh.New()
	vg.SetMaxTimeSkew(time.Minute * 15)

When checking the date header, Vangoh follows the precedent of RFC 2616,
accepting dates in any of the following formats:
	ANSIC    = "Mon Jan _2 15:04:05 2006"
	RFC822   = "02 Jan 06 15:04 MST"
	RFC822Z  = "02 Jan 06 15:04 -0700"
	RFC850   = "Monday, 02-Jan-06 15:04:05 MST"
	RFC1123  = "Mon, 02 Jan 2006 15:04:05 MST"
	RFC1123Z = "Mon, 02 Jan 2006 15:04:05 -0700"
*/
func (vg *Vangoh) SetMaxTimeSkew(timeSkew time.Duration) {
	vg.maxTimeSkew = timeSkew
}

// Checks a request for proper authentication details, returning the relevent
// error if the request fails this check or nil if the request passes.
func (vg *Vangoh) AuthenticateRequest(r *http.Request) *AuthenticationError {
	// Parse the ORG, KEY, and SIGNATURE out of the Authorization header.
	authHeader := strings.TrimSpace(r.Header.Get("Authorization"))
	if authHeader == "" {
		return ErrorAuthHeaderMissing
	}
	match, err := regexp.Match(AuthRegex, []byte(authHeader))
	if err != nil || !match {
		return ErrorAuthHeaderMalformed
	}
	orgSplit := strings.Split(authHeader, " ")
	org := orgSplit[0]
	keySplit := strings.Split(orgSplit[1], ":")
	key := []byte(keySplit[0])
	actualSignatureB64 := keySplit[1]

	// Check that the request was made in the acceptable window.
	dateHeader := strings.TrimSpace(r.Header.Get("Date"))
	if dateHeader == "" {
		return ErrorDateHeaderMissing
	}
	date, err := multiFormatDateParse(
		// TODO: break out into const.
		[]string{
			time.RFC822,
			time.RFC822Z,
			time.RFC850,
			time.ANSIC,
			time.RFC1123,
			time.RFC1123Z},
		dateHeader)
	if err != nil {
		return ErrorDateHeaderMalformed
	}
	present := clock.Now()
	if present.Sub(date) > vg.maxTimeSkew || date.Sub(present) > vg.maxTimeSkew {
		return ErrorDateHeaderTooSkewed
	}

	// Load the secret key from the appropriate key provider, given the ID from
	// the Authorization header.
	var provider secretProvider
	if vg.singleProvider != nil {
		provider = *vg.singleProvider
	} else {
		var exists bool
		provider, exists = vg.providersByOrg[org]
		if !exists {
			return ErrorAuthOrgUnknown
		}
	}

	var voidPtr unsafe.Pointer = nil
	var secret []byte

	switch provider := provider.(type) {
	case SecretProviderWithCallback:
		secret, err = provider.GetSecret(key, &voidPtr)
	case SecretProvider:
		secret, err = provider.GetSecret(key)
	}
	if err != nil {
		return ErrorInProviderKeyLookup
	}
	if secret == nil {
		return ErrorSecretNotFound
	}

	// Calculate the b64 signature and compare against the one sent by the client.
	expectedSignature := vg.ConstructSignature(r, secret)
	expectedSignatureB64 := base64.StdEncoding.EncodeToString(expectedSignature)
	if subtle.ConstantTimeCompare([]byte(expectedSignatureB64), []byte(actualSignatureB64)) != 1 {
		return ErrorHMACSignatureMismatch
	}

	switch provider := provider.(type) {
	case SecretProviderWithCallback:
		if voidPtr != nil {
			provider.SuccessCallback(r, &voidPtr)
		} else {
			provider.SuccessCallback(r, nil)
		}
	}
	// If we have made it this far, authentication is successful.
	return nil
}

func (vg *Vangoh) ConstructSignature(r *http.Request, secret []byte) []byte {
	signingString := vg.CreateSigningString(r)
	mac := hmac.New(vg.algorithm, secret)
	mac.Write([]byte(signingString))
	return mac.Sum(nil)
}

func multiFormatDateParse(formats []string, dateStr string) (time.Time, error) {
	for index := range formats {
		if date, err := time.Parse(formats[index], dateStr); err == nil {
			return date, nil
		}
	}
	return time.Time{}, errors.New("Date does not match any valid format")
}

/*
CreateSigningString creates the string used for signature generation, in accordance with
the specifications as laid out in the package documentation. Refer there for more detail,
or to the Amazon Signature V2 documentation: http://docs.aws.amazon.com/AmazonS3/latest/dev/RESTAuthentication.html.
*/
func (vg *Vangoh) CreateSigningString(r *http.Request) string {
	var buffer bytes.Buffer

	buffer.WriteString(r.Method)
	buffer.WriteString(newline)

	buffer.WriteString(r.Header.Get("Content-MD5"))
	buffer.WriteString(newline)

	buffer.WriteString(r.Header.Get("Content-Type"))
	buffer.WriteString(newline)

	buffer.WriteString(r.Header.Get("Date"))
	buffer.WriteString(newline)

	customHeaders := vg.createHeadersString(r)
	buffer.WriteString(customHeaders)

	buffer.WriteString(r.URL.Path)

	return buffer.String()
}

// Create the canonicalized header string part of a request's signature body.
func (vg *Vangoh) createHeadersString(r *http.Request) string {
	if len(vg.includedHeaders) == 0 {
		return ""
	}

	// For each defined regex, determine the set of headers that match. Repeat
	// for all regexes, without duplication, to get the final set of custom
	// headers to use.
	var sanitizedHeaders = make(map[string][]string)

	for _, compiledRegex := range vg.includedHeaders {
		for header := range r.Header {
			lowerHeader := strings.ToLower(header)
			if _, found := sanitizedHeaders[lowerHeader]; found {
				continue
			}
			if compiledRegex.MatchString(header) {
				sanitizedHeaders[lowerHeader] = r.Header[header]
			}
		}
	}

	var orderedHeaders []string
	for header := range sanitizedHeaders {
		orderedHeaders = append(orderedHeaders, header)
	}
	sort.Strings(orderedHeaders)

	// At this point sanitized contains all the headers to be included in the
	// hash. Now we need to retrieve their values, and sanitize them
	// appropriately.
	var buffer bytes.Buffer
	for header := range orderedHeaders {
		buffer.WriteString(orderedHeaders[header])
		buffer.WriteString(":")

		var sanitizedValues []string
		for i := range sanitizedHeaders[orderedHeaders[header]] {
			str := sanitizedHeaders[orderedHeaders[header]][i]
			str = strings.TrimSpace(str)
			str = strings.Replace(str, "\n", "", -1)
			sanitizedValues = append(sanitizedValues, str)
		}

		// Note that sanitizedValues are unsorted here - the order that they are
		// specified in the header will affect the hash result. This conforms with
		// the standard set by AWS, though it may be more reliable to add this
		// sorting in at some point.
		for i := range sanitizedValues {
			buffer.WriteString(sanitizedValues[i])
			if i < (len(sanitizedValues) - 1) {
				buffer.WriteString(",")
			}
		}

		buffer.WriteString(newline)
	}

	return buffer.String()
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/client"
	"github.com/pkg/errors"
)

var stageMutexes []*sync.Mutex
var completedConditions []*sync.Cond
var completedStages []bool
var stageIndex map[string]int

func run(c *client.Client, p *Pipeline, rootpath, filename string) error {
	// generate a version number/name if the pipeline description didn't
	if p.Version == "" {
		p.Version = createName()
	}

	stageMutexes = make([]*sync.Mutex, len(p.Stages))
	completedConditions = make([]*sync.Cond, len(p.Stages))
	completedStages = make([]bool, len(p.Stages))

	for i := range stageMutexes {
		stageMutexes[i] = &sync.Mutex{}
		completedConditions[i] = sync.NewCond(stageMutexes[i])
	}

	stageIndex = make(map[string]int, len(p.Stages))

	// Name to index mapping
	for i, stage := range p.Stages {
		stageIndex[stage.Name] = i
	}

	e := make(chan error, len(p.Stages))

	for i, stage := range p.Stages {
		go func(i int, stage *Stage) {
			mountpath := "/walrus/" + stage.Name
			hostpath := rootpath + "/" + stage.Name

			repo, tag := getRepoAndTag(stage.Image)
			image := repo + ":" + tag
			_, err := c.ImagePull(context.Background(), image, types.ImagePullOptions{})
			if err != nil {
				e <- errors.Wrap(err, "Could not pull image")
				return
			}

			// If the stage has any inputs it waits for these stages to complete
			// before starting
			if len(stage.Inputs) > 0 {
				for _, input := range stage.Inputs {
					index := stageIndex[input]
					cond := completedConditions[index]
					cond.L.Lock()
					for !completedStages[index] {
						cond.Wait()
					}
					cond.L.Unlock()
				}
			}

			// try to open output directory, if it exists then we can serve the
			// "cached"/old results
			_, err = os.Open(hostpath)

			if !stage.Cache || err != nil {

				// first remove any previous container (note that we're ignoring
				// errors)
				c.ContainerRemove(context.Background(), stage.Name, types.ContainerRemoveOptions{})

				// Note the 0777 permission bits. We use such liberal bits since
				// we do not know about the users within the docker containers that
				// are going to be run. We want to fix this later!
				err = os.MkdirAll(hostpath, 0777)
				if err != nil {
					e <- errors.Wrap(err, "Could not create output directory for stage")
					return
				}

				binds := []string{hostpath + ":" + mountpath}
				binds = append(binds, stage.Volumes...)

				resp, err := c.ContainerCreate(context.Background(),
					&container.Config{Image: image,
						Env: stage.Env,
						Cmd: stage.Cmd,
					},
					&container.HostConfig{
						Binds:       binds,
						VolumesFrom: stage.Inputs},
					&network.NetworkingConfig{},
					stage.Name)

				if err != nil {
					e <- errors.Wrap(err, "Could not create container "+stage.Name)
					return
				}
				containerId := resp.ID

				err = c.ContainerStart(context.Background(), containerId,
					types.ContainerStartOptions{})

				if err != nil {
					e <- errors.Wrap(err, "Could not start container")
					return
				}

				_, err = c.ContainerWait(context.Background(), containerId)
				if err != nil {
					e <- errors.Wrap(err, "Failed to wait for container to finish")
					return
				}

			}

			cond := completedConditions[i]
			cond.L.Lock()

			// Notifies waiting stages on completion
			completedStages[i] = true

			cond.Broadcast()
			cond.L.Unlock()

			fmt.Println(stage.Name, "completed successfully.")

			e <- nil
		}(i, stage)
	}

	for i := 0; i < len(p.Stages); i++ {
		err := <-e
		if err != nil {
			return err
		}
	}

	// Restore permission bits to output directory
	// err := filepath.Walk(rootpath, func(name string, info os.FileInfo, err error) error {
	// 	return os.Chmod(name, 0666)
	// })
	// if err != nil {
	// 	return err
	// }

	return nil
}

// Stops any previously run pipeline and deletes the containers. If any, we
// ignore errors from docker.
func stopPreviousRun(c *client.Client, stages []*Stage) {
	for _, stage := range stages {
		c.ContainerKill(context.Background(), stage.Name, "9")
	}
}

// Generate a list of volume mounts on the form
// /hostpath/stagename-containerid/:/walrus/stagename
func getInputVolumes(inputs []string, hostpath string) (volumes []string) {
	for _, input := range inputs {
		volumes = append(volumes, hostpath+"/"+input+":"+"/walrus"+"/"+input)
	}
	return volumes
}

func getRepoAndTag(pipelineImage string) (repo, tag string) {
	repoAndTag := strings.Split(pipelineImage, ":")
	if len(repoAndTag) == 1 {
		tag = "latest"
	} else {
		tag = repoAndTag[1]
	}
	repo = repoAndTag[0]

	return repo, tag
}

// Saves the pipeline configuration (json) to a new .walrus directory in the
// output directory specified by the user. Can be used to determine what
// produced the output in the output directory.
func saveConfiguration(hostpath string, p *Pipeline) error {
	configPath := createConfigPath(hostpath)
	err := os.Mkdir(configPath, 0777)
	if err != nil {
		return errors.Wrap(err, "Could not create directory to save old pipeline results")
	}

	filename := configPath + "/" + "pipeline.json"
	f, err := os.OpenFile(filename, os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return errors.Wrap(err, "Could not open old pipeline configuration")
	}
	b, err := json.Marshal(p)
	if err != nil {
		return errors.Wrap(err, "Could not marshal pipeline configuration")
	}

	_, err = f.Write(b)
	if err != nil {
		return errors.Wrap(err, "Could not write pipeline configuration")
	}

	return nil

}

// Moves the output of the previous runs into new folders for each stage. The
// names are STAGENAME-VERSION.
func savePreviousRun(hostpath string) error {

	// Check if there is any output from the previous runs
	f, err := os.Open(hostpath)
	if err != nil {
		// Output dir does not exist, nothing to back up.
		if os.IsNotExist(err) {
			return nil
		}

		return errors.Wrap(err, "Could not open previous pipeline outputs")
	}
	files, err := f.Readdir(-1)
	if err != nil {
		return errors.Wrap(err, "Could not read output directory")
	}

	if len(files) == 0 {
		return errors.Wrap(err, "No files in output directory, nothing to back up")
	}

	// Read old pipeline description to get it's version (use it for renaming)
	configPath := createConfigPath(hostpath)
	configFilename := configPath + "/" + "pipeline.json"
	p, err := ParseConfig(configFilename)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return errors.Wrap(err, "Could not parse old pipeline configuration")
	}

	absPath, err := filepath.Abs(hostpath)
	if err != nil {
		return errors.Wrap(err, "Could not get the absolute path of the output directory")
	}

	// iterate over all stages and move each output folder to a new directory
	// if the stage is cached the directory is copied. .
	for _, stage := range p.Stages {
		newFilename := absPath + "/" + stage.Name + "-" + p.Version
		oldFilename := absPath + "/" + stage.Name

		if stage.Cache {
			err = copyDir(oldFilename, newFilename)
		} else {
			err = os.Rename(oldFilename, newFilename)
		}
		if err != nil {
			return errors.Wrap(err, "Could not back up pipeline stage directory")
		}
	}

	// back up the walrus directory as well
	walrusDirectory := configPath
	newWalrusDirectory := walrusDirectory + "-" + p.Version
	err = os.Rename(walrusDirectory, newWalrusDirectory)
	if err != nil {
		return err
	}

	return nil
}

// Copies one directory to another. NOTE, only one level, not recursive yet.
func copyDir(src, dest string) error {

	err := os.Mkdir(dest, 0777)
	if err != nil {
		return err
	}
	err = filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
		// nothing to back up
		if err != nil {
			return err
		}

		// if dir then create dir , else a file then copy it
		if info.IsDir() {
		} else {
			srcFile, err := os.Open(path)
			if err != nil {
				return err
			}

			destFilename := dest + "/" + info.Name()
			destFile, err := os.OpenFile(destFilename, os.O_CREATE, 0644)
			if err != nil {
				return err
			}

			_, err = io.Copy(srcFile, destFile)
			if err != nil {
				return err
			}

		}
		return nil
	})
	return err
}

// Returns the full path of the  walrus configuration directory
func createConfigPath(hostpath string) string {
	return hostpath + "/" + ".walrus"
}

func fixMountPaths(stages []*Stage) error {
	for i, stage := range stages {
		updatedVolumes := []string{}
		for _, volume := range stage.Volumes {
			hostClientPath := strings.Split(volume, ":")

			if len(hostClientPath) > 2 {
				return errors.New("Incorrect volume " + volume + " in pipeline description")
			}

			hostPath := hostClientPath[0]
			clientPath := hostClientPath[1]

			if strings.HasPrefix(hostPath, "/") {
				updatedVolumes = append(updatedVolumes, volume)
				continue
			}

			absPath, err := filepath.Abs(hostPath)
			if err != nil {
				return errors.Wrap(err, "Could not get the absolute path of the mount path")
			}
			updatedVolumes = append(updatedVolumes, absPath+":"+clientPath)
		}
		stages[i].Volumes = updatedVolumes
	}

	return nil
}

func main() {
	var configFilename = flag.String("f", "pipeline.json", "pipeline description file")
	var outputDir = flag.String("output", "walrus", "where walrus should store output data on the host")

	flag.Parse()

	hostpath, err := filepath.Abs(*outputDir)
	if err != nil {
		fmt.Println("Check hostpath", err)
	}

	flag.Parse()
	client, err := client.NewEnvClient()
	if err != nil {
		fmt.Println(err)
		return
	}

	p, err := ParseConfig(*configFilename)
	if err != nil {
		fmt.Println(err)
		return
	}

	err = fixMountPaths(p.Stages)
	if err != nil {
		fmt.Println(err)
		return
	}

	stopPreviousRun(client, p.Stages)
	err = savePreviousRun(hostpath)
	if err != nil {
		fmt.Println(err)
		return
	}

	err = run(client, p, hostpath, *configFilename)

	if err != nil {
		fmt.Println(err)
		return
	}

	err = saveConfiguration(hostpath, p)
	if err != nil {
		fmt.Println(err)
		return
	} else {
		fmt.Println("All stages completed successfully.",
			"\nOutput written to ", hostpath)
	}

}
<commit_msg>wait until docker image is pulled down<commit_after>package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/client"
	"github.com/pkg/errors"
)

var stageMutexes []*sync.Mutex
var completedConditions []*sync.Cond
var completedStages []bool
var stageIndex map[string]int

func run(c *client.Client, p *Pipeline, rootpath, filename string) error {
	// generate a version number/name if the pipeline description didn't
	if p.Version == "" {
		p.Version = createName()
	}

	stageMutexes = make([]*sync.Mutex, len(p.Stages))
	completedConditions = make([]*sync.Cond, len(p.Stages))
	completedStages = make([]bool, len(p.Stages))

	for i := range stageMutexes {
		stageMutexes[i] = &sync.Mutex{}
		completedConditions[i] = sync.NewCond(stageMutexes[i])
	}

	stageIndex = make(map[string]int, len(p.Stages))

	// Name to index mapping
	for i, stage := range p.Stages {
		stageIndex[stage.Name] = i
	}

	e := make(chan error, len(p.Stages))

	for i, stage := range p.Stages {
		go func(i int, stage *Stage) {
			mountpath := "/walrus/" + stage.Name
			hostpath := rootpath + "/" + stage.Name

			repo, tag := getRepoAndTag(stage.Image)
			image := repo + ":" + tag
			rc, err := c.ImagePull(context.Background(), image,
				types.ImagePullOptions{})
			if err != nil {
				e <- errors.Wrap(err, "Could not pull image")
				return
			}

			defer rc.Close()

			_, err = ioutil.ReadAll(rc)
			if err != nil {
				e <- errors.Wrap(err, "error reading image pull")
			}

			// If the stage has any inputs it waits for these stages to complete
			// before starting
			if len(stage.Inputs) > 0 {
				for _, input := range stage.Inputs {
					index := stageIndex[input]
					cond := completedConditions[index]
					cond.L.Lock()
					for !completedStages[index] {
						cond.Wait()
					}
					cond.L.Unlock()
				}
			}

			// try to open output directory, if it exists then we can serve the
			// "cached"/old results
			_, err = os.Open(hostpath)

			if !stage.Cache || err != nil {

				// first remove any previous container (note that we're ignoring
				// errors)
				c.ContainerRemove(context.Background(), stage.Name, types.ContainerRemoveOptions{})

				// Note the 0777 permission bits. We use such liberal bits since
				// we do not know about the users within the docker containers that
				// are going to be run. We want to fix this later!
				err = os.MkdirAll(hostpath, 0777)
				if err != nil {
					e <- errors.Wrap(err, "Could not create output directory for stage")
					return
				}

				binds := []string{hostpath + ":" + mountpath}
				binds = append(binds, stage.Volumes...)

				resp, err := c.ContainerCreate(context.Background(),
					&container.Config{Image: image,
						Env: stage.Env,
						Cmd: stage.Cmd,
					},
					&container.HostConfig{
						Binds:       binds,
						VolumesFrom: stage.Inputs},
					&network.NetworkingConfig{},
					stage.Name)

				if err != nil {
					e <- errors.Wrap(err, "Could not create container "+stage.Name)
					return
				}
				containerId := resp.ID

				err = c.ContainerStart(context.Background(), containerId,
					types.ContainerStartOptions{})

				if err != nil {
					e <- errors.Wrap(err, "Could not start container")
					return
				}

				_, err = c.ContainerWait(context.Background(), containerId)
				if err != nil {
					e <- errors.Wrap(err, "Failed to wait for container to finish")
					return
				}

			}

			cond := completedConditions[i]
			cond.L.Lock()

			// Notifies waiting stages on completion
			completedStages[i] = true

			cond.Broadcast()
			cond.L.Unlock()

			fmt.Println(stage.Name, "completed successfully.")

			e <- nil
		}(i, stage)
	}

	for i := 0; i < len(p.Stages); i++ {
		err := <-e
		if err != nil {
			return err
		}
	}

	// Restore permission bits to output directory
	// err := filepath.Walk(rootpath, func(name string, info os.FileInfo, err error) error {
	// 	return os.Chmod(name, 0666)
	// })
	// if err != nil {
	// 	return err
	// }

	return nil
}

// Stops any previously run pipeline and deletes the containers. If any, we
// ignore errors from docker.
func stopPreviousRun(c *client.Client, stages []*Stage) {
	for _, stage := range stages {
		c.ContainerKill(context.Background(), stage.Name, "9")
	}
}

// Generate a list of volume mounts on the form
// /hostpath/stagename-containerid/:/walrus/stagename
func getInputVolumes(inputs []string, hostpath string) (volumes []string) {
	for _, input := range inputs {
		volumes = append(volumes, hostpath+"/"+input+":"+"/walrus"+"/"+input)
	}
	return volumes
}

func getRepoAndTag(pipelineImage string) (repo, tag string) {
	repoAndTag := strings.Split(pipelineImage, ":")
	if len(repoAndTag) == 1 {
		tag = "latest"
	} else {
		tag = repoAndTag[1]
	}
	repo = repoAndTag[0]

	return repo, tag
}

// Saves the pipeline configuration (json) to a new .walrus directory in the
// output directory specified by the user. Can be used to determine what
// produced the output in the output directory.
func saveConfiguration(hostpath string, p *Pipeline) error {
	configPath := createConfigPath(hostpath)
	err := os.Mkdir(configPath, 0777)
	if err != nil {
		return errors.Wrap(err, "Could not create directory to save old pipeline results")
	}

	filename := configPath + "/" + "pipeline.json"
	f, err := os.OpenFile(filename, os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return errors.Wrap(err, "Could not open old pipeline configuration")
	}
	b, err := json.Marshal(p)
	if err != nil {
		return errors.Wrap(err, "Could not marshal pipeline configuration")
	}

	_, err = f.Write(b)
	if err != nil {
		return errors.Wrap(err, "Could not write pipeline configuration")
	}

	return nil

}

// Moves the output of the previous runs into new folders for each stage. The
// names are STAGENAME-VERSION.
func savePreviousRun(hostpath string) error {

	// Check if there is any output from the previous runs
	f, err := os.Open(hostpath)
	if err != nil {
		// Output dir does not exist, nothing to back up.
		if os.IsNotExist(err) {
			return nil
		}

		return errors.Wrap(err, "Could not open previous pipeline outputs")
	}
	files, err := f.Readdir(-1)
	if err != nil {
		return errors.Wrap(err, "Could not read output directory")
	}

	if len(files) == 0 {
		return errors.Wrap(err, "No files in output directory, nothing to back up")
	}

	// Read old pipeline description to get it's version (use it for renaming)
	configPath := createConfigPath(hostpath)
	configFilename := configPath + "/" + "pipeline.json"
	p, err := ParseConfig(configFilename)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return errors.Wrap(err, "Could not parse old pipeline configuration")
	}

	absPath, err := filepath.Abs(hostpath)
	if err != nil {
		return errors.Wrap(err, "Could not get the absolute path of the output directory")
	}

	// iterate over all stages and move each output folder to a new directory
	// if the stage is cached the directory is copied. .
	for _, stage := range p.Stages {
		newFilename := absPath + "/" + stage.Name + "-" + p.Version
		oldFilename := absPath + "/" + stage.Name

		if stage.Cache {
			err = copyDir(oldFilename, newFilename)
		} else {
			err = os.Rename(oldFilename, newFilename)
		}
		if err != nil {
			return errors.Wrap(err, "Could not back up pipeline stage directory")
		}
	}

	// back up the walrus directory as well
	walrusDirectory := configPath
	newWalrusDirectory := walrusDirectory + "-" + p.Version
	err = os.Rename(walrusDirectory, newWalrusDirectory)
	if err != nil {
		return err
	}

	return nil
}

// Copies one directory to another. NOTE, only one level, not recursive yet.
func copyDir(src, dest string) error {

	err := os.Mkdir(dest, 0777)
	if err != nil {
		return err
	}
	err = filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
		// nothing to back up
		if err != nil {
			return err
		}

		// if dir then create dir , else a file then copy it
		if info.IsDir() {
		} else {
			srcFile, err := os.Open(path)
			if err != nil {
				return err
			}

			destFilename := dest + "/" + info.Name()
			destFile, err := os.OpenFile(destFilename, os.O_CREATE, 0644)
			if err != nil {
				return err
			}

			_, err = io.Copy(srcFile, destFile)
			if err != nil {
				return err
			}

		}
		return nil
	})
	return err
}

// Returns the full path of the  walrus configuration directory
func createConfigPath(hostpath string) string {
	return hostpath + "/" + ".walrus"
}

func fixMountPaths(stages []*Stage) error {
	for i, stage := range stages {
		updatedVolumes := []string{}
		for _, volume := range stage.Volumes {
			hostClientPath := strings.Split(volume, ":")

			if len(hostClientPath) > 2 {
				return errors.New("Incorrect volume " + volume + " in pipeline description")
			}

			hostPath := hostClientPath[0]
			clientPath := hostClientPath[1]

			if strings.HasPrefix(hostPath, "/") {
				updatedVolumes = append(updatedVolumes, volume)
				continue
			}

			absPath, err := filepath.Abs(hostPath)
			if err != nil {
				return errors.Wrap(err, "Could not get the absolute path of the mount path")
			}
			updatedVolumes = append(updatedVolumes, absPath+":"+clientPath)
		}
		stages[i].Volumes = updatedVolumes
	}

	return nil
}

func main() {
	var configFilename = flag.String("f", "pipeline.json", "pipeline description file")
	var outputDir = flag.String("output", "walrus", "where walrus should store output data on the host")

	flag.Parse()

	hostpath, err := filepath.Abs(*outputDir)
	if err != nil {
		fmt.Println("Check hostpath", err)
	}

	flag.Parse()
	client, err := client.NewEnvClient()
	if err != nil {
		fmt.Println(err)
		return
	}

	p, err := ParseConfig(*configFilename)
	if err != nil {
		fmt.Println(err)
		return
	}

	err = fixMountPaths(p.Stages)
	if err != nil {
		fmt.Println(err)
		return
	}

	stopPreviousRun(client, p.Stages)
	err = savePreviousRun(hostpath)
	if err != nil {
		fmt.Println(err)
		return
	}

	err = run(client, p, hostpath, *configFilename)

	if err != nil {
		fmt.Println(err)
		return
	}

	err = saveConfiguration(hostpath, p)
	if err != nil {
		fmt.Println(err)
		return
	} else {
		fmt.Println("All stages completed successfully.",
			"\nOutput written to ", hostpath)
	}

}
<|endoftext|>


<commit_before>package trakt

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
)

type Trakt struct {
	Url         string
	ApiKey      string
	AccessToken string
}

func (t Trakt) Get(path string) (resp *http.Response, err error) {
	url := fmt.Sprintf("%s%s", t.Url, path)
	client := &http.Client{}
	req, err := http.NewRequest("GET", url, nil)
	req.Header.Add("Content-type", "application/json")
	req.Header.Add("trakt-api-version", "2")
	req.Header.Add("trakt-api-key", t.ApiKey)
	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", t.AccessToken))
	return client.Do(req)
}
<commit_msg>Removing unused packages from import<commit_after>package trakt

import (
	"fmt"
	"net/http"
)

type Trakt struct {
	Url         string
	ApiKey      string
	AccessToken string
}

func (t Trakt) Get(path string) (resp *http.Response, err error) {
	url := fmt.Sprintf("%s%s", t.Url, path)
	client := &http.Client{}
	req, err := http.NewRequest("GET", url, nil)
	req.Header.Add("Content-type", "application/json")
	req.Header.Add("trakt-api-version", "2")
	req.Header.Add("trakt-api-key", t.ApiKey)
	req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", t.AccessToken))
	return client.Do(req)
}
<|endoftext|>


<commit_before>package libp2pquic

import (
	ma "github.com/multiformats/go-multiaddr"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Transport", func() {
	var t *QuicTransport

	BeforeEach(func() {
		t = NewQuicTransport(nil)
	})

	It("matches", func() {
		invalidAddr, err := ma.NewMultiaddr("/ip4/127.0.0.1/udp/1234")
		Expect(err).ToNot(HaveOccurred())
		validAddr, err := ma.NewMultiaddr("/ip4/127.0.0.1/udp/1234/quic")
		Expect(err).ToNot(HaveOccurred())
		Expect(t.Matches(invalidAddr)).To(BeFalse())
		Expect(t.Matches(validAddr)).To(BeTrue())
	})
})
<commit_msg>add tests for transport.Listen<commit_after>package libp2pquic

import (
	ma "github.com/multiformats/go-multiaddr"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Transport", func() {
	var t *QuicTransport

	BeforeEach(func() {
		t = NewQuicTransport(nil)
	})

	Context("listening", func() {
		It("creates a new listener", func() {
			maddr, err := ma.NewMultiaddr("/ip4/127.0.0.1/udp/1234/quic")
			Expect(err).ToNot(HaveOccurred())
			ln, err := t.Listen(maddr)
			Expect(err).ToNot(HaveOccurred())
			Expect(ln.Multiaddr()).To(Equal(maddr))
		})

		It("returns an existing listener", func() {
			maddr, err := ma.NewMultiaddr("/ip4/127.0.0.1/udp/1235/quic")
			Expect(err).ToNot(HaveOccurred())
			ln, err := t.Listen(maddr)
			Expect(err).ToNot(HaveOccurred())
			Expect(ln.Multiaddr()).To(Equal(maddr))
			ln2, err := t.Listen(maddr)
			Expect(err).ToNot(HaveOccurred())
			Expect(ln2).To(Equal(ln))
			Expect(t.listeners).To(HaveLen(1))
		})
	})

	It("matches", func() {
		invalidAddr, err := ma.NewMultiaddr("/ip4/127.0.0.1/udp/1234")
		Expect(err).ToNot(HaveOccurred())
		validAddr, err := ma.NewMultiaddr("/ip4/127.0.0.1/udp/1234/quic")
		Expect(err).ToNot(HaveOccurred())
		Expect(t.Matches(invalidAddr)).To(BeFalse())
		Expect(t.Matches(validAddr)).To(BeTrue())
	})
})
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"log"
	"net"
	"time"
)

// Transparently intercept HTTPS connections.

// runTransparentServer transparently intercepts connections, listening at addr.
func runTransparentServer(addr string) error {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	listenerChan <- ln

	localAddresses, err := getLocalAddresses()
	if err != nil {
		log.Println("Error getting list of this server's IP addresses:", err)
		// Continue, but without protection against infinite redirect loops.
	}

	var tempDelay time.Duration

	for {
		conn, err := ln.Accept()
		if err != nil {
			if ne, ok := err.(net.Error); ok && ne.Temporary() {
				if tempDelay == 0 {
					tempDelay = 5 * time.Millisecond
				} else {
					tempDelay *= 2
				}
				if max := 1 * time.Second; tempDelay > max {
					tempDelay = max
				}
				log.Printf("Accept error: %v; retrying in %v", err, tempDelay)
				time.Sleep(tempDelay)
				continue
			}
			return err
		}

		serverAddr, err := realServerAddress(conn)
		if err != nil {
			log.Println("Error getting original address for intercepted connection:", err)
			continue
		}
		user, _, _ := net.SplitHostPort(conn.RemoteAddr().String())

		var server string
		if tcpAddr, ok := serverAddr.(*net.TCPAddr); ok {
			server = tcpAddr.IP.String()
		} else {
			server = serverAddr.String()
		}
		if localAddresses[server] {
			// This is not an intercepted connection; it is a direct connection to
			// our transparent port. If we bump it, we will end up with an infinite
			// loop of redirects.
			logTLS(user, serverAddr.String(), "", errors.New("infinite redirect loop"))
			conn.Close()
			continue
		}

		go SSLBump(conn, serverAddr.String(), user, "")
	}

	panic("unreachable")
}

// getLocalAddresses returns a set of the IP addresses of this machine's
// network interfaces.
func getLocalAddresses() (map[string]bool, error) {
	res := map[string]bool{}

	ifs, err := net.Interfaces()
	if err != nil {
		return nil, err
	}

	for _, iface := range ifs {
		addrs, err := iface.Addrs()
		if err != nil {
			return nil, err
		}
		for _, a := range addrs {
			if ipNet, ok := a.(*net.IPNet); ok {
				a = &net.IPAddr{IP: ipNet.IP}
			}
			res[a.String()] = true
		}
	}

	return res, nil
}
<commit_msg>Make error message for transparent connections more informative.<commit_after>package main

import (
	"errors"
	"log"
	"net"
	"time"
)

// Transparently intercept HTTPS connections.

// runTransparentServer transparently intercepts connections, listening at addr.
func runTransparentServer(addr string) error {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	listenerChan <- ln

	localAddresses, err := getLocalAddresses()
	if err != nil {
		log.Println("Error getting list of this server's IP addresses:", err)
		// Continue, but without protection against infinite redirect loops.
	}

	var tempDelay time.Duration

	for {
		conn, err := ln.Accept()
		if err != nil {
			if ne, ok := err.(net.Error); ok && ne.Temporary() {
				if tempDelay == 0 {
					tempDelay = 5 * time.Millisecond
				} else {
					tempDelay *= 2
				}
				if max := 1 * time.Second; tempDelay > max {
					tempDelay = max
				}
				log.Printf("Accept error: %v; retrying in %v", err, tempDelay)
				time.Sleep(tempDelay)
				continue
			}
			return err
		}

		serverAddr, err := realServerAddress(conn)
		if err != nil {
			log.Printf("Error getting original address for intercepted connection from %v: %v", conn.RemoteAddr(), err)
			continue
		}
		user, _, _ := net.SplitHostPort(conn.RemoteAddr().String())

		var server string
		if tcpAddr, ok := serverAddr.(*net.TCPAddr); ok {
			server = tcpAddr.IP.String()
		} else {
			server = serverAddr.String()
		}
		if localAddresses[server] {
			// This is not an intercepted connection; it is a direct connection to
			// our transparent port. If we bump it, we will end up with an infinite
			// loop of redirects.
			logTLS(user, serverAddr.String(), "", errors.New("infinite redirect loop"))
			conn.Close()
			continue
		}

		go SSLBump(conn, serverAddr.String(), user, "")
	}

	panic("unreachable")
}

// getLocalAddresses returns a set of the IP addresses of this machine's
// network interfaces.
func getLocalAddresses() (map[string]bool, error) {
	res := map[string]bool{}

	ifs, err := net.Interfaces()
	if err != nil {
		return nil, err
	}

	for _, iface := range ifs {
		addrs, err := iface.Addrs()
		if err != nil {
			return nil, err
		}
		for _, a := range addrs {
			if ipNet, ok := a.(*net.IPNet); ok {
				a = &net.IPAddr{IP: ipNet.IP}
			}
			res[a.String()] = true
		}
	}

	return res, nil
}
<|endoftext|>


<commit_before>package widget

import (
	"fmt"
	"html/template"
	"os"
	"path"

	"github.com/jinzhu/gorm"
	"github.com/qor/admin"
	"github.com/qor/qor/resource"
	"github.com/qor/roles"
)

var (
	root, _                = os.Getwd()
	viewPaths              []string
	registeredWidgets      []*Widget
	registeredWidgetsGroup []*WidgetsGroup
)

// Config widget config
type Config struct {
	DB            *gorm.DB
	Admin         *admin.Admin
	PreviewAssets []string
}

func init() {
	if path := os.Getenv("WEB_ROOT"); path != "" {
		root = path
	}
}

// New new widgets container
func New(config *Config) *Widgets {
	widgets := &Widgets{Config: config, funcMaps: template.FuncMap{}, AssetFileSystem: &admin.AssetFileSystem{}}

	if root != "" {
		widgets.RegisterViewPath(path.Join(root, "app/views/widgets"))
	}
	widgets.RegisterViewPath("app/views/widgets")
	return widgets
}

// Widgets widgets container
type Widgets struct {
	funcMaps              template.FuncMap
	Config                *Config
	Resource              *admin.Resource
	AssetFileSystem       admin.AssetFSInterface
	WidgetSettingResource *admin.Resource
}

// SetAssetFS set asset fs for render
func (widgets *Widgets) SetAssetFS(assetFS admin.AssetFSInterface) {
	for _, viewPath := range viewPaths {
		assetFS.RegisterPath(viewPath)
	}

	assetFS.Compile()

	widgets.AssetFileSystem = assetFS
}

// RegisterWidget register a new widget
func (widgets *Widgets) RegisterWidget(w *Widget) {
	registeredWidgets = append(registeredWidgets, w)
}

// RegisterWidgetsGroup register widgets group
func (widgets *Widgets) RegisterWidgetsGroup(group *WidgetsGroup) {
	registeredWidgetsGroup = append(registeredWidgetsGroup, group)
}

// RegisterFuncMap register view funcs, it could be used when render templates
func (widgets *Widgets) RegisterFuncMap(name string, fc interface{}) {
	widgets.funcMaps[name] = fc
}

// ConfigureQorResourceBeforeInitialize a method used to config Widget for qor admin
func (widgets *Widgets) ConfigureQorResourceBeforeInitialize(res resource.Resourcer) {
	if res, ok := res.(*admin.Resource); ok {
		// register view paths
		res.GetAdmin().RegisterViewPath("github.com/qor/widget/views")

		// set resources
		widgets.Resource = res

		// set setting resource
		if widgets.WidgetSettingResource == nil {
			widgets.WidgetSettingResource = res.GetAdmin().NewResource(&QorWidgetSetting{}, &admin.Config{Name: res.Name})
		}

		res.Name = widgets.WidgetSettingResource.Name

		for funcName, fc := range funcMap {
			res.GetAdmin().RegisterFuncMap(funcName, fc)
		}

		// configure routes
		controller := widgetController{Widgets: widgets}
		router := res.GetAdmin().GetRouter()
		router.Get(widgets.WidgetSettingResource.ToParam(), controller.Index, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
		router.Get(fmt.Sprintf("%v/new", widgets.WidgetSettingResource.ToParam()), controller.New, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
		router.Get(fmt.Sprintf("%v/!setting", widgets.WidgetSettingResource.ToParam()), controller.Setting, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
		router.Get(fmt.Sprintf("%v/%v", widgets.WidgetSettingResource.ToParam(), widgets.WidgetSettingResource.ParamIDName()), controller.Edit, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
		router.Get(fmt.Sprintf("%v/%v/!preview", widgets.WidgetSettingResource.ToParam(), widgets.WidgetSettingResource.ParamIDName()), controller.Preview, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
		router.Get(fmt.Sprintf("%v/%v/edit", widgets.WidgetSettingResource.ToParam(), widgets.WidgetSettingResource.ParamIDName()), controller.Edit, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
		router.Put(fmt.Sprintf("%v/%v", widgets.WidgetSettingResource.ToParam(), widgets.WidgetSettingResource.ParamIDName()), controller.Update, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
		router.Post(widgets.WidgetSettingResource.ToParam(), controller.Update, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
		router.Get(fmt.Sprintf("%v/inline-edit", res.ToParam()), controller.InlineEdit, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
	}
}

// Widget widget struct
type Widget struct {
	Name        string
	PreviewIcon string
	Group       string
	Templates   []string
	Setting     *admin.Resource
	Permission  *roles.Permission
	Context     func(context *Context, setting interface{}) *Context
}

// WidgetsGroup widgets Group
type WidgetsGroup struct {
	Name    string
	Widgets []string
}

// GetWidget get widget by name
func GetWidget(name string) *Widget {
	for _, w := range registeredWidgets {
		if w.Name == name {
			return w
		}
	}

	for _, g := range registeredWidgetsGroup {
		if g.Name == name {
			for _, widgetName := range g.Widgets {
				return GetWidget(widgetName)
			}
		}
	}
	return nil
}
<commit_msg>Don't compile when set assetfs<commit_after>package widget

import (
	"fmt"
	"html/template"
	"os"
	"path"

	"github.com/jinzhu/gorm"
	"github.com/qor/admin"
	"github.com/qor/qor/resource"
	"github.com/qor/roles"
)

var (
	root, _                = os.Getwd()
	viewPaths              []string
	registeredWidgets      []*Widget
	registeredWidgetsGroup []*WidgetsGroup
)

// Config widget config
type Config struct {
	DB            *gorm.DB
	Admin         *admin.Admin
	PreviewAssets []string
}

func init() {
	if path := os.Getenv("WEB_ROOT"); path != "" {
		root = path
	}
}

// New new widgets container
func New(config *Config) *Widgets {
	widgets := &Widgets{Config: config, funcMaps: template.FuncMap{}, AssetFileSystem: &admin.AssetFileSystem{}}

	if root != "" {
		widgets.RegisterViewPath(path.Join(root, "app/views/widgets"))
	}
	widgets.RegisterViewPath("app/views/widgets")
	return widgets
}

// Widgets widgets container
type Widgets struct {
	funcMaps              template.FuncMap
	Config                *Config
	Resource              *admin.Resource
	AssetFileSystem       admin.AssetFSInterface
	WidgetSettingResource *admin.Resource
}

// SetAssetFS set asset fs for render
func (widgets *Widgets) SetAssetFS(assetFS admin.AssetFSInterface) {
	for _, viewPath := range viewPaths {
		assetFS.RegisterPath(viewPath)
	}

	widgets.AssetFileSystem = assetFS
}

// RegisterWidget register a new widget
func (widgets *Widgets) RegisterWidget(w *Widget) {
	registeredWidgets = append(registeredWidgets, w)
}

// RegisterWidgetsGroup register widgets group
func (widgets *Widgets) RegisterWidgetsGroup(group *WidgetsGroup) {
	registeredWidgetsGroup = append(registeredWidgetsGroup, group)
}

// RegisterFuncMap register view funcs, it could be used when render templates
func (widgets *Widgets) RegisterFuncMap(name string, fc interface{}) {
	widgets.funcMaps[name] = fc
}

// ConfigureQorResourceBeforeInitialize a method used to config Widget for qor admin
func (widgets *Widgets) ConfigureQorResourceBeforeInitialize(res resource.Resourcer) {
	if res, ok := res.(*admin.Resource); ok {
		// register view paths
		res.GetAdmin().RegisterViewPath("github.com/qor/widget/views")

		// set resources
		widgets.Resource = res

		// set setting resource
		if widgets.WidgetSettingResource == nil {
			widgets.WidgetSettingResource = res.GetAdmin().NewResource(&QorWidgetSetting{}, &admin.Config{Name: res.Name})
		}

		res.Name = widgets.WidgetSettingResource.Name

		for funcName, fc := range funcMap {
			res.GetAdmin().RegisterFuncMap(funcName, fc)
		}

		// configure routes
		controller := widgetController{Widgets: widgets}
		router := res.GetAdmin().GetRouter()
		router.Get(widgets.WidgetSettingResource.ToParam(), controller.Index, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
		router.Get(fmt.Sprintf("%v/new", widgets.WidgetSettingResource.ToParam()), controller.New, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
		router.Get(fmt.Sprintf("%v/!setting", widgets.WidgetSettingResource.ToParam()), controller.Setting, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
		router.Get(fmt.Sprintf("%v/%v", widgets.WidgetSettingResource.ToParam(), widgets.WidgetSettingResource.ParamIDName()), controller.Edit, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
		router.Get(fmt.Sprintf("%v/%v/!preview", widgets.WidgetSettingResource.ToParam(), widgets.WidgetSettingResource.ParamIDName()), controller.Preview, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
		router.Get(fmt.Sprintf("%v/%v/edit", widgets.WidgetSettingResource.ToParam(), widgets.WidgetSettingResource.ParamIDName()), controller.Edit, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
		router.Put(fmt.Sprintf("%v/%v", widgets.WidgetSettingResource.ToParam(), widgets.WidgetSettingResource.ParamIDName()), controller.Update, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
		router.Post(widgets.WidgetSettingResource.ToParam(), controller.Update, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
		router.Get(fmt.Sprintf("%v/inline-edit", res.ToParam()), controller.InlineEdit, &admin.RouteConfig{Resource: widgets.WidgetSettingResource})
	}
}

// Widget widget struct
type Widget struct {
	Name        string
	PreviewIcon string
	Group       string
	Templates   []string
	Setting     *admin.Resource
	Permission  *roles.Permission
	Context     func(context *Context, setting interface{}) *Context
}

// WidgetsGroup widgets Group
type WidgetsGroup struct {
	Name    string
	Widgets []string
}

// GetWidget get widget by name
func GetWidget(name string) *Widget {
	for _, w := range registeredWidgets {
		if w.Name == name {
			return w
		}
	}

	for _, g := range registeredWidgetsGroup {
		if g.Name == name {
			for _, widgetName := range g.Widgets {
				return GetWidget(widgetName)
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package statisticalaccumulator

import "math/big"

type StatisticalAccumulator struct {
	n              *big.Rat
	sigmaXI        *big.Rat
	sigmaXISquared *big.Rat

	one *big.Rat
}

func NewStatisticalAccumulator() (*StatisticalAccumulator) {

	n := new(big.Rat)
	n.SetInt64(0)

	sigmaXI := new(big.Rat)
	sigmaXI.SetInt64(0)

	sigmaXISquared := new(big.Rat)
	sigmaXISquared.SetInt64(0)

	one := new(big.Rat)
	one.SetInt64(1)

	me := StatisticalAccumulator{
		n:n,
		sigmaXI:sigmaXI,
		sigmaXISquared:sigmaXISquared,
		one:one,
	}

	return &me
}

func (me *StatisticalAccumulator) PutString(x string) {

	xx := new(big.Rat)
	xx.SetString(x)

	me.PutRat(xx)
}

func (me *StatisticalAccumulator) PutRat(x *big.Rat) {

	me.n.Add(me.n, me.one)

	me.sigmaXI.Add(me.sigmaXI, x)

	xSquared := new(big.Rat)
	xSquared.Mul(x, x)

	me.sigmaXISquared.Add(me.sigmaXISquared, xSquared)
}

func (me *StatisticalAccumulator) N() (*big.Rat) {
	n := new(big.Rat)
	n.Set(me.n)

	return n
}

func (me *StatisticalAccumulator) Mean() (*big.Rat) {
	mean := new(big.Rat)

	mean.Inv(me.n)

	mean.Mul(mean, me.sigmaXI)

	return mean
}

func (me *StatisticalAccumulator) Variance() (*big.Rat) {
	variance := new(big.Rat)

	variance.Inv(me.n)

	variance.Mul(variance, me.sigmaXISquared)

	temp := new(big.Rat)

	temp.Mul(me.n, me.n)
	temp.Inv(temp)

	temp.Mul(temp, me.sigmaXI)
	temp.Mul(temp, me.sigmaXI)

	variance.Sub(variance, temp)

	return variance
}
<commit_msg>NewStatisticalAccumulator() -> New()<commit_after>package statisticalaccumulator

import "math/big"

type StatisticalAccumulator struct {
	n              *big.Rat
	sigmaXI        *big.Rat
	sigmaXISquared *big.Rat

	one *big.Rat
}

func New() (*StatisticalAccumulator) {

	n := new(big.Rat)
	n.SetInt64(0)

	sigmaXI := new(big.Rat)
	sigmaXI.SetInt64(0)

	sigmaXISquared := new(big.Rat)
	sigmaXISquared.SetInt64(0)

	one := new(big.Rat)
	one.SetInt64(1)

	me := StatisticalAccumulator{
		n:n,
		sigmaXI:sigmaXI,
		sigmaXISquared:sigmaXISquared,
		one:one,
	}

	return &me
}

func (me *StatisticalAccumulator) PutString(x string) {

	xx := new(big.Rat)
	xx.SetString(x)

	me.PutRat(xx)
}

func (me *StatisticalAccumulator) PutRat(x *big.Rat) {

	me.n.Add(me.n, me.one)

	me.sigmaXI.Add(me.sigmaXI, x)

	xSquared := new(big.Rat)
	xSquared.Mul(x, x)

	me.sigmaXISquared.Add(me.sigmaXISquared, xSquared)
}

func (me *StatisticalAccumulator) N() (*big.Rat) {
	n := new(big.Rat)
	n.Set(me.n)

	return n
}

func (me *StatisticalAccumulator) Mean() (*big.Rat) {
	mean := new(big.Rat)

	mean.Inv(me.n)

	mean.Mul(mean, me.sigmaXI)

	return mean
}

func (me *StatisticalAccumulator) Variance() (*big.Rat) {
	variance := new(big.Rat)

	variance.Inv(me.n)

	variance.Mul(variance, me.sigmaXISquared)

	temp := new(big.Rat)

	temp.Mul(me.n, me.n)
	temp.Inv(temp)

	temp.Mul(temp, me.sigmaXI)
	temp.Mul(temp, me.sigmaXI)

	variance.Sub(variance, temp)

	return variance
}
<|endoftext|>


<commit_before>// Copyright 2015 The Cockroach 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.
//
// Author: Bram Gruneir (bram+code@cockroachlabs.com)

package storage

import (
	"time"

	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/config"
	"github.com/cockroachdb/cockroach/gossip"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/log"
	"github.com/coreos/etcd/raft"
)

const (
	// raftLogQueueMaxSize is the max size of the queue.
	raftLogQueueMaxSize = 100
	// RaftLogQueueTimerDuration is the duration between checking the
	// raft logs.
	RaftLogQueueTimerDuration = time.Second
	// RaftLogQueueStaleThreshold is the minimum threshold for stale raft log
	// entries. A stale entry is one which all replicas of the range have
	// progressed past and thus is no longer needed and can be pruned.
	RaftLogQueueStaleThreshold = 1
)

// raftLogQueue manages a queue of replicas slated to have their raft logs
// truncated by removing unneeded entries.
type raftLogQueue struct {
	baseQueue
	db *client.DB
}

// newRaftLogQueue returns a new instance of raftLogQueue.
func newRaftLogQueue(db *client.DB, gossip *gossip.Gossip) *raftLogQueue {
	rlq := &raftLogQueue{
		db: db,
	}
	rlq.baseQueue = makeBaseQueue("raftlog", rlq, gossip, raftLogQueueMaxSize)
	return rlq
}

func (*raftLogQueue) needsLeaderLease() bool {
	return false
}

func (*raftLogQueue) acceptsUnsplitRanges() bool {
	return true
}

// getTruncatableIndexes returns the total number of stale raft log entries that
// can be truncated and the oldest index that cannot be pruned.
func getTruncatableIndexes(r *Replica) (uint64, uint64, error) {
	rangeID := r.RangeID
	raftStatus := r.store.RaftStatus(rangeID)
	if raftStatus == nil {
		if log.V(1) {
			log.Infof("the raft group doesn't exist for range %d", rangeID)
		}
		return 0, 0, nil
	}

	// Is this the raft leader?
	if raftStatus.RaftState != raft.StateLeader {
		return 0, 0, nil
	}

	// Find the oldest index still in use by the range.
	oldestIndex := raftStatus.Applied
	for _, progress := range raftStatus.Progress {
		if progress.Match < oldestIndex {
			oldestIndex = progress.Match
		}
	}

	r.mu.Lock()
	defer r.mu.Unlock()
	firstIndex, err := r.FirstIndex()
	if err != nil {
		return 0, 0, util.Errorf("error retrieving first index for range %d: %s", rangeID, err)
	}

	if oldestIndex < firstIndex {
		return 0, 0, util.Errorf("raft log's oldest index is less than the first index for range %d", rangeID)
	}

	// Return the number of truncatable indexes.
	return oldestIndex - firstIndex, oldestIndex, nil
}

// shouldQueue determines whether a range should be queued for truncating. This
// is true only if the replica is the raft leader and if the total number of
// the range's raft log's stale entries exceeds RaftLogQueueStaleThreshold.
func (*raftLogQueue) shouldQueue(now roachpb.Timestamp, r *Replica, _ config.SystemConfig) (shouldQ bool,
	priority float64) {

	truncatableIndexes, _, err := getTruncatableIndexes(r)
	if err != nil {
		log.Warning(err)
		return false, 0
	}

	return truncatableIndexes > RaftLogQueueStaleThreshold, float64(truncatableIndexes)
}

// process truncates the raft log of the range if the replica is the raft
// leader and if the total number of the range's raft log's stale entries
// exceeds RaftLogQueueStaleThreshold.
func (rlq *raftLogQueue) process(now roachpb.Timestamp, r *Replica, _ config.SystemConfig) error {

	truncatableIndexes, oldestIndex, err := getTruncatableIndexes(r)
	if err != nil {
		return err
	}

	// Can and should the raft logs be truncated?
	if truncatableIndexes > RaftLogQueueStaleThreshold {
		if log.V(1) {
			log.Infof("truncating the raft log of range %d to %d", r.RangeID, oldestIndex)
		}
		b := &client.Batch{}
		b.InternalAddRequest(&roachpb.TruncateLogRequest{
			Span:    roachpb.Span{Key: r.Desc().StartKey.AsRawKey()},
			Index:   oldestIndex,
			RangeID: r.RangeID,
		})
		return rlq.db.Run(b).GoError()
	}
	return nil
}

// timer returns interval between processing successive queued truncations.
func (*raftLogQueue) timer() time.Duration {
	return RaftLogQueueTimerDuration
}

// purgatoryChan returns nil.
func (*raftLogQueue) purgatoryChan() <-chan struct{} {
	return nil
}
<commit_msg>storage: add more info to a raft truncation error message<commit_after>// Copyright 2015 The Cockroach 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.
//
// Author: Bram Gruneir (bram+code@cockroachlabs.com)

package storage

import (
	"time"

	"github.com/cockroachdb/cockroach/client"
	"github.com/cockroachdb/cockroach/config"
	"github.com/cockroachdb/cockroach/gossip"
	"github.com/cockroachdb/cockroach/roachpb"
	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/log"
	"github.com/coreos/etcd/raft"
)

const (
	// raftLogQueueMaxSize is the max size of the queue.
	raftLogQueueMaxSize = 100
	// RaftLogQueueTimerDuration is the duration between checking the
	// raft logs.
	RaftLogQueueTimerDuration = time.Second
	// RaftLogQueueStaleThreshold is the minimum threshold for stale raft log
	// entries. A stale entry is one which all replicas of the range have
	// progressed past and thus is no longer needed and can be pruned.
	RaftLogQueueStaleThreshold = 1
)

// raftLogQueue manages a queue of replicas slated to have their raft logs
// truncated by removing unneeded entries.
type raftLogQueue struct {
	baseQueue
	db *client.DB
}

// newRaftLogQueue returns a new instance of raftLogQueue.
func newRaftLogQueue(db *client.DB, gossip *gossip.Gossip) *raftLogQueue {
	rlq := &raftLogQueue{
		db: db,
	}
	rlq.baseQueue = makeBaseQueue("raftlog", rlq, gossip, raftLogQueueMaxSize)
	return rlq
}

func (*raftLogQueue) needsLeaderLease() bool {
	return false
}

func (*raftLogQueue) acceptsUnsplitRanges() bool {
	return true
}

// getTruncatableIndexes returns the total number of stale raft log entries that
// can be truncated and the oldest index that cannot be pruned.
func getTruncatableIndexes(r *Replica) (uint64, uint64, error) {
	rangeID := r.RangeID
	raftStatus := r.store.RaftStatus(rangeID)
	if raftStatus == nil {
		if log.V(1) {
			log.Infof("the raft group doesn't exist for range %d", rangeID)
		}
		return 0, 0, nil
	}

	// Is this the raft leader?
	if raftStatus.RaftState != raft.StateLeader {
		return 0, 0, nil
	}

	// Find the oldest index still in use by the range.
	oldestIndex := raftStatus.Applied
	for _, progress := range raftStatus.Progress {
		if progress.Match < oldestIndex {
			oldestIndex = progress.Match
		}
	}

	r.mu.Lock()
	defer r.mu.Unlock()
	firstIndex, err := r.FirstIndex()
	if err != nil {
		return 0, 0, util.Errorf("error retrieving first index for range %d: %s", rangeID, err)
	}

	if oldestIndex < firstIndex {
		return 0, 0, util.Errorf("raft log's oldest index (%d) is less than the first index (%d) for range %d",
			oldestIndex, firstIndex, rangeID)
	}

	// Return the number of truncatable indexes.
	return oldestIndex - firstIndex, oldestIndex, nil
}

// shouldQueue determines whether a range should be queued for truncating. This
// is true only if the replica is the raft leader and if the total number of
// the range's raft log's stale entries exceeds RaftLogQueueStaleThreshold.
func (*raftLogQueue) shouldQueue(now roachpb.Timestamp, r *Replica, _ config.SystemConfig) (shouldQ bool,
	priority float64) {

	truncatableIndexes, _, err := getTruncatableIndexes(r)
	if err != nil {
		log.Warning(err)
		return false, 0
	}

	return truncatableIndexes > RaftLogQueueStaleThreshold, float64(truncatableIndexes)
}

// process truncates the raft log of the range if the replica is the raft
// leader and if the total number of the range's raft log's stale entries
// exceeds RaftLogQueueStaleThreshold.
func (rlq *raftLogQueue) process(now roachpb.Timestamp, r *Replica, _ config.SystemConfig) error {

	truncatableIndexes, oldestIndex, err := getTruncatableIndexes(r)
	if err != nil {
		return err
	}

	// Can and should the raft logs be truncated?
	if truncatableIndexes > RaftLogQueueStaleThreshold {
		if log.V(1) {
			log.Infof("truncating the raft log of range %d to %d", r.RangeID, oldestIndex)
		}
		b := &client.Batch{}
		b.InternalAddRequest(&roachpb.TruncateLogRequest{
			Span:    roachpb.Span{Key: r.Desc().StartKey.AsRawKey()},
			Index:   oldestIndex,
			RangeID: r.RangeID,
		})
		return rlq.db.Run(b).GoError()
	}
	return nil
}

// timer returns interval between processing successive queued truncations.
func (*raftLogQueue) timer() time.Duration {
	return RaftLogQueueTimerDuration
}

// purgatoryChan returns nil.
func (*raftLogQueue) purgatoryChan() <-chan struct{} {
	return nil
}
<|endoftext|>


<commit_before>// +build linux

package storage

import (
	"io/ioutil"
	"os"
	"testing"
	"math/rand"
	"github.com/golang/glog"
	"github.com/stretchr/testify/assert"
)

func TestWAL(t *testing.T) {
	assert := assert.New(t)

	//Create temp directory
	dir, err := ioutil.TempDir("", "IosWALTests")
	if err != nil {
		glog.Fatal(err)
	}
	defer os.RemoveAll(dir) // clean up

	//create file
  testFile := dir + "/test.temp"
	wal := openWriteAheadFile(testFile, "fsync")
	actualBytes, err := ioutil.ReadFile(testFile)
	assert.Equal(64*1000*1000,len(actualBytes), "File is expected size")

	//verfiy that write ahead logging works
  expectedBytes := make([]byte, 100)
	rand.Read(expectedBytes)
  wal.writeAhead(expectedBytes)
  actualBytes, err = ioutil.ReadFile(testFile)
  assert.Nil(err)
  //assert.Equal(1001,len(actualBytes), "Number of bytes read is not same as bytes written")
  assert.Equal(expectedBytes,actualBytes[len(actualBytes)-100:], "Bytes read are not same as written")

}
<commit_msg>shifting write ahead test bytes by 1<commit_after>// +build linux

package storage

import (
	"io/ioutil"
	"os"
	"testing"
	"math/rand"
	"github.com/golang/glog"
	"github.com/stretchr/testify/assert"
)

func TestWAL(t *testing.T) {
	assert := assert.New(t)

	//Create temp directory
	dir, err := ioutil.TempDir("", "IosWALTests")
	if err != nil {
		glog.Fatal(err)
	}
	defer os.RemoveAll(dir) // clean up

	//create file
  testFile := dir + "/test.temp"
	wal := openWriteAheadFile(testFile, "fsync")
	actualBytes, err := ioutil.ReadFile(testFile)
	assert.Equal(64*1000*1000,len(actualBytes), "File is expected size")

	//verfiy that write ahead logging works
  expectedBytes := make([]byte, 100)
	rand.Read(expectedBytes)
  wal.writeAhead(expectedBytes)
  actualBytes, err = ioutil.ReadFile(testFile)
  assert.Nil(err)
  //assert.Equal(1001,len(actualBytes), "Number of bytes read is not same as bytes written")
  assert.Equal(expectedBytes,actualBytes[len(actualBytes)-100+1:], "Bytes read are not same as written")

}
<|endoftext|>


<commit_before>// Package revolver provides a revolving file writer.
package revolver

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"
)

type revWriter struct {
	dir      string
	prefix   string
	suffix   string
	middle   func() string
	maxBytes int
	maxFiles int
	size     int
	file     *os.File
	lock     *sync.Mutex // synchronizes file operations
}

// Must wraps the call to NewWriter and returns a io.WriteCloser or panics
func Must(w io.WriteCloser, err error) io.WriteCloser {
	if err != nil {
		panic(fmt.Errorf("could not create revolving log writer, %v", err))
	}
	return w
}

// New  returns a io.WriteCloser that writes revolving files as specified by the given conf.
// Calling New will always create a new file even if there is space left in other files.
// If the configured directory doesn't exist it will be created.
func New(conf Conf) (io.WriteCloser, error) {
	if err := ValidConf(conf); err != nil {
		return nil, err
	}
	conf = clean(conf)
	return NewQuick(conf.Dir, conf.Prefix, conf.Suffix, conf.Middle, conf.MaxBytes, conf.MaxFiles)
}

// NewQuick returns a io.WriteCloser that writes revolving files.
// Calling New will always create a new file even if there is space left in other files.
// If the configured directory doesn't exist it will be created.
func NewQuick(dir, prefix, suffix string, middle func() string, maxBytes, maxFiles int) (io.WriteCloser, error) {
	if prefix == "" {
		return nil, fmt.Errorf("revolver, prefix can not be empty")
	}
	if middle == nil {
		middle = func() string { return "" }
	}
	if maxBytes < 1 {
		return nil, fmt.Errorf("revolver, maxBytes must be > 0")
	}
	if maxFiles < 1 {
		return nil, fmt.Errorf("revolver, maxFiles must be > 0")
	}

	if err := setupDirs(dir); err != nil {
		return nil, fmt.Errorf("revolver setup, %v", err)
	}
	if err := countAndRemoveFiles(dir, prefix, maxFiles); err != nil {
		return nil, fmt.Errorf("revolver, remove, %v", err)
	}

	file, err := createFile(dir, prefix, suffix, middle)
	if err != nil {
		return nil, fmt.Errorf("revolver, create, %v", err)
	}

	return &revWriter{
		dir:      filepath.Clean(dir),
		prefix:   filepath.Clean(prefix),
		suffix:   suffix,
		middle:   middle,
		maxBytes: maxBytes,
		maxFiles: maxFiles,
		file:     file,
		lock:     &sync.Mutex{},
	}, nil
}

// Write the given bytes into the log file specified by the given conf.
// If there is not enough file space left,surplus files will be deleted and a new file will be created.
func (l *revWriter) Write(p []byte) (n int, err error) {
	l.lock.Lock()
	defer l.lock.Unlock()

	size := len(p)
	if size > l.maxBytes {
		return 0, fmt.Errorf("revolver, bytes to write %d over max file size %d", size, l.maxBytes)
	}
	if l.file == nil || l.size+size > l.maxBytes {
		if err := l.close(); err != nil {
			return 0, fmt.Errorf("revolver, close, %v", err)
		}

		if err := countAndRemoveFiles(l.dir, l.prefix, l.maxFiles); err != nil {
			return 0, fmt.Errorf("revolver, remove, %v", err)
		}

		file, err := createFile(l.dir, l.prefix, l.suffix, l.middle)
		if err != nil {
			return 0, fmt.Errorf("revolver, create, %v", err)

		}
		l.file = file
		l.size = 0
	}

	l.size += size
	return l.file.Write(p)

}

// Close closes the current log file and sets the writer reference to nil.
// If the file reference is nil, the returned err is always be nil.
// Writing to a nil referencing writer cleans up surplus files and creates a new file.
func (l *revWriter) Close() error {
	l.lock.Lock()
	defer l.lock.Unlock()
	return l.close()
}

func (l *revWriter) close() error {
	if l.file == nil {
		return nil
	}
	err := l.file.Close()
	l.file = nil
	return err

}
<commit_msg>Modified comments.<commit_after>// Package revolver provides a revolving file writer.
package revolver

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"
)

type revWriter struct {
	dir      string
	prefix   string
	suffix   string
	middle   func() string
	maxBytes int
	maxFiles int
	size     int
	file     *os.File
	lock     *sync.Mutex // synchronizes file operations
}

// Must wraps the call to NewWriter and returns a io.WriteCloser or panics
func Must(w io.WriteCloser, err error) io.WriteCloser {
	if err != nil {
		panic(fmt.Errorf("could not create revolving log writer, %v", err))
	}
	return w
}

// New  returns a io.WriteCloser that writes revolving files as specified by the given conf.
// Calling New will always create a new file even if there is space left in other files.
// If the configured directory doesn't exist it will be created.
func New(conf Conf) (io.WriteCloser, error) {
	if err := ValidConf(conf); err != nil {
		return nil, err
	}
	conf = clean(conf)
	return NewQuick(conf.Dir, conf.Prefix, conf.Suffix, conf.Middle, conf.MaxBytes, conf.MaxFiles)
}

// NewQuick is like New with the difference that no Conf struct is needed.
// Calling New will always create a new file even if there is space left in other files.
// If the configured directory doesn't exist it will be created.
func NewQuick(dir, prefix, suffix string, middle func() string, maxBytes, maxFiles int) (io.WriteCloser, error) {
	if prefix == "" {
		return nil, fmt.Errorf("revolver, prefix can not be empty")
	}
	if middle == nil {
		middle = func() string { return "" }
	}
	if maxBytes < 1 {
		return nil, fmt.Errorf("revolver, maxBytes must be > 0")
	}
	if maxFiles < 1 {
		return nil, fmt.Errorf("revolver, maxFiles must be > 0")
	}

	if err := setupDirs(dir); err != nil {
		return nil, fmt.Errorf("revolver setup, %v", err)
	}
	if err := countAndRemoveFiles(dir, prefix, maxFiles); err != nil {
		return nil, fmt.Errorf("revolver, remove, %v", err)
	}

	file, err := createFile(dir, prefix, suffix, middle)
	if err != nil {
		return nil, fmt.Errorf("revolver, create, %v", err)
	}

	return &revWriter{
		dir:      filepath.Clean(dir),
		prefix:   filepath.Clean(prefix),
		suffix:   suffix,
		middle:   middle,
		maxBytes: maxBytes,
		maxFiles: maxFiles,
		file:     file,
		lock:     &sync.Mutex{},
	}, nil
}

// Write writes the given bytes into the current file. The specifics of the file are specified on writer creation.
// If there is not enough file space left,surplus files will be deleted and a new file will be created.
func (l *revWriter) Write(p []byte) (n int, err error) {
	l.lock.Lock()
	defer l.lock.Unlock()

	size := len(p)
	if size > l.maxBytes {
		return 0, fmt.Errorf("revolver, bytes to write %d over max file size %d", size, l.maxBytes)
	}
	if l.file == nil || l.size+size > l.maxBytes {
		if err := l.close(); err != nil {
			return 0, fmt.Errorf("revolver, close, %v", err)
		}

		if err := countAndRemoveFiles(l.dir, l.prefix, l.maxFiles); err != nil {
			return 0, fmt.Errorf("revolver, remove, %v", err)
		}

		file, err := createFile(l.dir, l.prefix, l.suffix, l.middle)
		if err != nil {
			return 0, fmt.Errorf("revolver, create, %v", err)

		}
		l.file = file
		l.size = 0
	}

	l.size += size
	return l.file.Write(p)

}

// Close closes the current log file and sets the writer reference to nil.
// If the file reference is nil, the returned err is always be nil.
// Writing to a nil referencing writer cleans up surplus files and creates a new file.
func (l *revWriter) Close() error {
	l.lock.Lock()
	defer l.lock.Unlock()
	return l.close()
}

func (l *revWriter) close() error {
	if l.file == nil {
		return nil
	}
	err := l.file.Close()
	l.file = nil
	return err

}
<|endoftext|>


<commit_before>package zerolog

import (
	"io"
	"sync"
)

// LevelWriter defines as interface a writer may implement in order
// to receive level information with payload.
type LevelWriter interface {
	io.Writer
	WriteLevel(level Level, p []byte) (n int, err error)
}

type levelWriterAdapter struct {
	io.Writer
}

func (lw levelWriterAdapter) WriteLevel(l Level, p []byte) (n int, err error) {
	return lw.Write(p)
}

type syncWriter struct {
	mu sync.Mutex
	lw LevelWriter
}

// SyncWriter wraps w so that each call to Write is synchronized with a mutex.
// This syncer can be the call to writer's Write method is not thread safe.
// Note that os.File Write operation is using write() syscall which is supposed
// to be thread-safe on POSIX systems. So there is no need to use this with
// os.File on such systems as zerolog guaranties to issue a single Write call
// per log event.
func SyncWriter(w io.Writer) io.Writer {
	if lw, ok := w.(LevelWriter); ok {
		return &syncWriter{lw: lw}
	}
	return &syncWriter{lw: levelWriterAdapter{w}}
}

// Write implements the io.Writer interface.
func (s *syncWriter) Write(p []byte) (n int, err error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.lw.Write(p)
}

// WriteLevel implements the LevelWriter interface.
func (s *syncWriter) WriteLevel(l Level, p []byte) (n int, err error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.lw.WriteLevel(l, p)
}

type multiLevelWriter struct {
	writers []LevelWriter
}

func (t multiLevelWriter) Write(p []byte) (n int, err error) {
	for _, w := range t.writers {
		n, err = w.Write(p)
		if err != nil {
			return
		}
		if n != len(p) {
			err = io.ErrShortWrite
			return
		}
	}
	return len(p), nil
}

func (t multiLevelWriter) WriteLevel(l Level, p []byte) (n int, err error) {
	for _, w := range t.writers {
		n, err = w.WriteLevel(l, p)
		if err != nil {
			return
		}
		if n != len(p) {
			err = io.ErrShortWrite
			return
		}
	}
	return len(p), nil
}

// MultiLevelWriter creates a writer that duplicates its writes to all the
// provided writers, similar to the Unix tee(1) command. If some writers
// implement LevelWriter, their WriteLevel method will be used instead of Write.
func MultiLevelWriter(writers ...io.Writer) LevelWriter {
	lwriters := make([]LevelWriter, 0, len(writers))
	for _, w := range writers {
		if lw, ok := w.(LevelWriter); ok {
			lwriters = append(lwriters, lw)
		} else {
			lwriters = append(lwriters, levelWriterAdapter{w})
		}
	}
	return multiLevelWriter{lwriters}
}
<commit_msg>Typo fix (#233)<commit_after>package zerolog

import (
	"io"
	"sync"
)

// LevelWriter defines as interface a writer may implement in order
// to receive level information with payload.
type LevelWriter interface {
	io.Writer
	WriteLevel(level Level, p []byte) (n int, err error)
}

type levelWriterAdapter struct {
	io.Writer
}

func (lw levelWriterAdapter) WriteLevel(l Level, p []byte) (n int, err error) {
	return lw.Write(p)
}

type syncWriter struct {
	mu sync.Mutex
	lw LevelWriter
}

// SyncWriter wraps w so that each call to Write is synchronized with a mutex.
// This syncer can be the call to writer's Write method is not thread safe.
// Note that os.File Write operation is using write() syscall which is supposed
// to be thread-safe on POSIX systems. So there is no need to use this with
// os.File on such systems as zerolog guarantees to issue a single Write call
// per log event.
func SyncWriter(w io.Writer) io.Writer {
	if lw, ok := w.(LevelWriter); ok {
		return &syncWriter{lw: lw}
	}
	return &syncWriter{lw: levelWriterAdapter{w}}
}

// Write implements the io.Writer interface.
func (s *syncWriter) Write(p []byte) (n int, err error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.lw.Write(p)
}

// WriteLevel implements the LevelWriter interface.
func (s *syncWriter) WriteLevel(l Level, p []byte) (n int, err error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.lw.WriteLevel(l, p)
}

type multiLevelWriter struct {
	writers []LevelWriter
}

func (t multiLevelWriter) Write(p []byte) (n int, err error) {
	for _, w := range t.writers {
		n, err = w.Write(p)
		if err != nil {
			return
		}
		if n != len(p) {
			err = io.ErrShortWrite
			return
		}
	}
	return len(p), nil
}

func (t multiLevelWriter) WriteLevel(l Level, p []byte) (n int, err error) {
	for _, w := range t.writers {
		n, err = w.WriteLevel(l, p)
		if err != nil {
			return
		}
		if n != len(p) {
			err = io.ErrShortWrite
			return
		}
	}
	return len(p), nil
}

// MultiLevelWriter creates a writer that duplicates its writes to all the
// provided writers, similar to the Unix tee(1) command. If some writers
// implement LevelWriter, their WriteLevel method will be used instead of Write.
func MultiLevelWriter(writers ...io.Writer) LevelWriter {
	lwriters := make([]LevelWriter, 0, len(writers))
	for _, w := range writers {
		if lw, ok := w.(LevelWriter); ok {
			lwriters = append(lwriters, lw)
		} else {
			lwriters = append(lwriters, levelWriterAdapter{w})
		}
	}
	return multiLevelWriter{lwriters}
}
<|endoftext|>