<commit_before>package secure

import (
	"context"
	"errors"
	"fmt"
	"github.com/Comcast/webpa-common/secure/key"
	"github.com/SermoDigital/jose/jws"
	"github.com/SermoDigital/jose/jwt"
	"regexp"
	"strings"
	"time"
)

var (
	ErrorNoProtectedHeader = errors.New("Missing protected header")
	ErrorNoSigningMethod   = errors.New("Signing method (alg) is missing or unrecognized")
)

// Validator describes the behavior of a type which can validate tokens
type Validator interface {
	// Validate asserts that the given token is valid, most often verifying
	// the credentials in the token.  A separate error is returned to indicate
	// any problems during validation, such as the inability to access a network resource.
	// In general, the contract of this method is that a Token passes validation
	// if and only if it returns BOTH true and a nil error.
	Validate(context.Context, *Token) (bool, error)
}

// ValidatorFunc is a function type that implements Validator
type ValidatorFunc func(context.Context, *Token) (bool, error)

func (v ValidatorFunc) Validate(ctx context.Context, token *Token) (bool, error) {
	return v(ctx, token)
}

// Validators is an aggregate Validator.  A Validators instance considers a token
// valid if any of its validators considers it valid.  An empty Validators rejects
// all tokens.
type Validators []Validator

func (v Validators) Validate(ctx context.Context, token *Token) (valid bool, err error) {
	for _, validator := range v {
		if valid, err = validator.Validate(ctx, token); valid && err == nil {
			return
		}
	}

	return
}

// ExactMatchValidator simply matches a token's value (exluding the prefix, such as "Basic"),
// to a string.
type ExactMatchValidator string

func (v ExactMatchValidator) Validate(ctx context.Context, token *Token) (bool, error) {
	for _, value := range strings.Split(string(v), ",") {
		if value == token.value {
			return true, nil
		}
	}

	return false, nil
}

// JWSValidator provides validation for JWT tokens encoded as JWS.
type JWSValidator struct {
	DefaultKeyId  string
	Resolver      key.Resolver
	Parser        JWSParser
	JWTValidators []*jwt.Validator
}

// capabilityValidation determines if a claim's capability is valid
func capabilityValidation(ctx context.Context, capability string) (valid_capabilities bool) {
	pieces := strings.Split(capability, ":")

	if len(pieces) == 5     &&
	   pieces[0] == "x1"    &&
	   pieces[1] == "webpa" {

		method_value, ok := ctx.Value("method").(string)
		if ok && (pieces[4] == "all" || strings.EqualFold(pieces[4], method_value)) {
			claimPath := fmt.Sprintf("/%s/[^/]+/%s", pieces[2],pieces[3])
			valid_capabilities, _ = regexp.MatchString(claimPath, ctx.Value("path").(string))
		}
	}

	return
}

func (v JWSValidator) Validate(ctx context.Context, token *Token) (valid bool, err error) {
	if token.Type() != Bearer {
		return
	}

	parser := v.Parser
	if parser == nil {
		parser = DefaultJWSParser
	}

	jwsToken, err := parser.ParseJWS(token)
	if err != nil {
		return
	}

	protected := jwsToken.Protected()
	if len(protected) == 0 {
		err = ErrorNoProtectedHeader
		return
	}

	alg, _ := protected.Get("alg").(string)
	signingMethod := jws.GetSigningMethod(alg)
	if signingMethod == nil {
		err = ErrorNoSigningMethod
		return
	}

	keyId, _ := protected.Get("kid").(string)
	if len(keyId) == 0 {
		keyId = v.DefaultKeyId
	}

	pair, err := v.Resolver.ResolveKey(keyId)
	if err != nil {
		return
	}

	// validate the signature
	if len(v.JWTValidators) > 0 {
		// all JWS implementations also implement jwt.JWT
		err = jwsToken.(jwt.JWT).Validate(pair.Public(), signingMethod, v.JWTValidators...)
	} else {
		err = jwsToken.Verify(pair.Public(), signingMethod)
	}

	if nil != err {
		return
	}

	// validate jwt token claims capabilities
	if caps, capOkay := jwsToken.Payload().(jws.Claims).Get("capabilities").([]interface{}); capOkay && len(caps) > 0 {
		for c := 0; c < len(caps); c++ {
			if cap_value, ok := caps[c].(string); ok {
				if valid = capabilityValidation(ctx, cap_value); valid {
					return
				}
			}
		}
	}
/*
	else if caps, capOkay := jwsToken.Payload().(jws.Claims).Get("capabilities").([]string); capOkay && len(caps) > 0 {
		for c := 0; c < len(caps); c++ {
			if valid = capabilityValidation(ctx, caps[c]); valid {
				return
			}
		}
	}
*/
	// This fail
	return
}

// JWTValidatorFactory is a configurable factory for *jwt.Validator instances
type JWTValidatorFactory struct {
	Expected  jwt.Claims `json:"expected"`
	ExpLeeway int        `json:"expLeeway"`
	NbfLeeway int        `json:"nbfLeeway"`
}

func (f *JWTValidatorFactory) expLeeway() time.Duration {
	if f.ExpLeeway > 0 {
		return time.Duration(f.ExpLeeway) * time.Second
	}

	return 0
}

func (f *JWTValidatorFactory) nbfLeeway() time.Duration {
	if f.NbfLeeway > 0 {
		return time.Duration(f.NbfLeeway) * time.Second
	}

	return 0
}

// New returns a jwt.Validator using the configuration expected claims (if any)
// and a validator function that checks the exp and nbf claims.
//
// The SermoDigital library doesn't appear to do anything with the EXP and NBF
// members of jwt.Validator, but this Factory Method populates them anyway.
func (f *JWTValidatorFactory) New(custom ...jwt.ValidateFunc) *jwt.Validator {
	expLeeway := f.expLeeway()
	nbfLeeway := f.nbfLeeway()

	var validateFunc jwt.ValidateFunc
	customCount := len(custom)
	if customCount > 0 {
		validateFunc = func(claims jwt.Claims) (err error) {
			err = claims.Validate(time.Now(), expLeeway, nbfLeeway)
			for index := 0; index < customCount && err == nil; index++ {
				err = custom[index](claims)
			}

			return
		}
	} else {
		// if no custom validate functions were passed, use a simpler function
		validateFunc = func(claims jwt.Claims) error {
			return claims.Validate(time.Now(), expLeeway, nbfLeeway)
		}
	}

	return &jwt.Validator{
		Expected: f.Expected,
		EXP:      expLeeway,
		NBF:      nbfLeeway,
		Fn:       validateFunc,
	}
}
<commit_msg>removing unnecessary old code<commit_after>package secure

import (
	"context"
	"errors"
	"fmt"
	"github.com/Comcast/webpa-common/secure/key"
	"github.com/SermoDigital/jose/jws"
	"github.com/SermoDigital/jose/jwt"
	"regexp"
	"strings"
	"time"
)

var (
	ErrorNoProtectedHeader = errors.New("Missing protected header")
	ErrorNoSigningMethod   = errors.New("Signing method (alg) is missing or unrecognized")
)

// Validator describes the behavior of a type which can validate tokens
type Validator interface {
	// Validate asserts that the given token is valid, most often verifying
	// the credentials in the token.  A separate error is returned to indicate
	// any problems during validation, such as the inability to access a network resource.
	// In general, the contract of this method is that a Token passes validation
	// if and only if it returns BOTH true and a nil error.
	Validate(context.Context, *Token) (bool, error)
}

// ValidatorFunc is a function type that implements Validator
type ValidatorFunc func(context.Context, *Token) (bool, error)

func (v ValidatorFunc) Validate(ctx context.Context, token *Token) (bool, error) {
	return v(ctx, token)
}

// Validators is an aggregate Validator.  A Validators instance considers a token
// valid if any of its validators considers it valid.  An empty Validators rejects
// all tokens.
type Validators []Validator

func (v Validators) Validate(ctx context.Context, token *Token) (valid bool, err error) {
	for _, validator := range v {
		if valid, err = validator.Validate(ctx, token); valid && err == nil {
			return
		}
	}

	return
}

// ExactMatchValidator simply matches a token's value (exluding the prefix, such as "Basic"),
// to a string.
type ExactMatchValidator string

func (v ExactMatchValidator) Validate(ctx context.Context, token *Token) (bool, error) {
	for _, value := range strings.Split(string(v), ",") {
		if value == token.value {
			return true, nil
		}
	}

	return false, nil
}

// JWSValidator provides validation for JWT tokens encoded as JWS.
type JWSValidator struct {
	DefaultKeyId  string
	Resolver      key.Resolver
	Parser        JWSParser
	JWTValidators []*jwt.Validator
}

// capabilityValidation determines if a claim's capability is valid
func capabilityValidation(ctx context.Context, capability string) (valid_capabilities bool) {
	pieces := strings.Split(capability, ":")

	if len(pieces) == 5     &&
	   pieces[0] == "x1"    &&
	   pieces[1] == "webpa" {

		method_value, ok := ctx.Value("method").(string)
		if ok && (pieces[4] == "all" || strings.EqualFold(pieces[4], method_value)) {
			claimPath := fmt.Sprintf("/%s/[^/]+/%s", pieces[2],pieces[3])
			valid_capabilities, _ = regexp.MatchString(claimPath, ctx.Value("path").(string))
		}
	}

	return
}

func (v JWSValidator) Validate(ctx context.Context, token *Token) (valid bool, err error) {
	if token.Type() != Bearer {
		return
	}

	parser := v.Parser
	if parser == nil {
		parser = DefaultJWSParser
	}

	jwsToken, err := parser.ParseJWS(token)
	if err != nil {
		return
	}

	protected := jwsToken.Protected()
	if len(protected) == 0 {
		err = ErrorNoProtectedHeader
		return
	}

	alg, _ := protected.Get("alg").(string)
	signingMethod := jws.GetSigningMethod(alg)
	if signingMethod == nil {
		err = ErrorNoSigningMethod
		return
	}

	keyId, _ := protected.Get("kid").(string)
	if len(keyId) == 0 {
		keyId = v.DefaultKeyId
	}

	pair, err := v.Resolver.ResolveKey(keyId)
	if err != nil {
		return
	}

	// validate the signature
	if len(v.JWTValidators) > 0 {
		// all JWS implementations also implement jwt.JWT
		err = jwsToken.(jwt.JWT).Validate(pair.Public(), signingMethod, v.JWTValidators...)
	} else {
		err = jwsToken.Verify(pair.Public(), signingMethod)
	}

	if nil != err {
		return
	}

	// validate jwt token claims capabilities
	if caps, capOkay := jwsToken.Payload().(jws.Claims).Get("capabilities").([]interface{}); capOkay && len(caps) > 0 {
		for c := 0; c < len(caps); c++ {
			if cap_value, ok := caps[c].(string); ok {
				if valid = capabilityValidation(ctx, cap_value); valid {
					return
				}
			}
		}
	}

	// This fail
	return
}

// JWTValidatorFactory is a configurable factory for *jwt.Validator instances
type JWTValidatorFactory struct {
	Expected  jwt.Claims `json:"expected"`
	ExpLeeway int        `json:"expLeeway"`
	NbfLeeway int        `json:"nbfLeeway"`
}

func (f *JWTValidatorFactory) expLeeway() time.Duration {
	if f.ExpLeeway > 0 {
		return time.Duration(f.ExpLeeway) * time.Second
	}

	return 0
}

func (f *JWTValidatorFactory) nbfLeeway() time.Duration {
	if f.NbfLeeway > 0 {
		return time.Duration(f.NbfLeeway) * time.Second
	}

	return 0
}

// New returns a jwt.Validator using the configuration expected claims (if any)
// and a validator function that checks the exp and nbf claims.
//
// The SermoDigital library doesn't appear to do anything with the EXP and NBF
// members of jwt.Validator, but this Factory Method populates them anyway.
func (f *JWTValidatorFactory) New(custom ...jwt.ValidateFunc) *jwt.Validator {
	expLeeway := f.expLeeway()
	nbfLeeway := f.nbfLeeway()

	var validateFunc jwt.ValidateFunc
	customCount := len(custom)
	if customCount > 0 {
		validateFunc = func(claims jwt.Claims) (err error) {
			err = claims.Validate(time.Now(), expLeeway, nbfLeeway)
			for index := 0; index < customCount && err == nil; index++ {
				err = custom[index](claims)
			}

			return
		}
	} else {
		// if no custom validate functions were passed, use a simpler function
		validateFunc = func(claims jwt.Claims) error {
			return claims.Validate(time.Now(), expLeeway, nbfLeeway)
		}
	}

	return &jwt.Validator{
		Expected: f.Expected,
		EXP:      expLeeway,
		NBF:      nbfLeeway,
		Fn:       validateFunc,
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package apimachinery

import (
	"context"
	"fmt"
	"reflect"
	"testing"
	"time"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/kubernetes"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"
	watchtools "k8s.io/client-go/tools/watch"
	"k8s.io/kubernetes/test/integration/framework"
)

func noopNormalization(output []string) []string {
	return output
}

func normalizeInformerOutputFunc(initialVal string) func(output []string) []string {
	return func(output []string) []string {
		result := make([]string, 0, len(output))

		// Removes initial value and all of its direct repetitions
		lastVal := initialVal
		for _, v := range output {
			// Make values unique as informer(List+Watch) duplicates some events
			if v == lastVal {
				continue
			}
			result = append(result, v)
			lastVal = v
		}

		return result
	}
}

func noop() {}

func TestWatchRestartsIfTimeoutNotReached(t *testing.T) {
	// Has to be longer than 5 seconds
	timeout := 2 * time.Minute

	// Set up a master
	masterConfig := framework.NewIntegrationTestMasterConfig()
	// Timeout is set random between MinRequestTimeout and 2x
	masterConfig.GenericConfig.MinRequestTimeout = int(timeout.Seconds()) / 4
	_, s, closeFn := framework.RunAMaster(masterConfig)
	defer closeFn()

	config := &restclient.Config{
		Host: s.URL,
	}

	namespaceObject := framework.CreateTestingNamespace("retry-watch", s, t)
	defer framework.DeleteTestingNamespace(namespaceObject, s, t)

	getListFunc := func(c *kubernetes.Clientset, secret *corev1.Secret) func(options metav1.ListOptions) *corev1.SecretList {
		return func(options metav1.ListOptions) *corev1.SecretList {
			options.FieldSelector = fields.OneTermEqualSelector("metadata.name", secret.Name).String()
			res, err := c.CoreV1().Secrets(secret.Namespace).List(context.TODO(), options)
			if err != nil {
				t.Fatalf("Failed to list Secrets: %v", err)
			}
			return res
		}
	}

	getWatchFunc := func(c *kubernetes.Clientset, secret *corev1.Secret) func(options metav1.ListOptions) (watch.Interface, error) {
		return func(options metav1.ListOptions) (watch.Interface, error) {
			options.FieldSelector = fields.OneTermEqualSelector("metadata.name", secret.Name).String()
			res, err := c.CoreV1().Secrets(secret.Namespace).Watch(context.TODO(), options)
			if err != nil {
				t.Fatalf("Failed to create a watcher on Secrets: %v", err)
			}
			return res, err
		}
	}

	generateEvents := func(t *testing.T, c *kubernetes.Clientset, secret *corev1.Secret, referenceOutput *[]string, stopChan chan struct{}, stoppedChan chan struct{}) {
		defer close(stoppedChan)
		counter := 0

		// These 5 seconds are here to protect against a race at the end when we could write something there at the same time as watch.Until ends
		softTimeout := timeout - 5*time.Second
		if softTimeout < 0 {
			panic("Timeout has to be grater than 5 seconds!")
		}
		endChannel := time.After(softTimeout)
		for {
			select {
			// TODO: get this lower once we figure out how to extend ETCD cache
			case <-time.After(1000 * time.Millisecond):
				counter = counter + 1

				patch := fmt.Sprintf(`{"metadata": {"annotations": {"count": "%d"}}}`, counter)
				_, err := c.CoreV1().Secrets(secret.Namespace).Patch(context.TODO(), secret.Name, types.StrategicMergePatchType, []byte(patch), metav1.PatchOptions{})
				if err != nil {
					t.Fatalf("Failed to patch secret: %v", err)
				}

				*referenceOutput = append(*referenceOutput, fmt.Sprintf("%d", counter))
			case <-endChannel:
				return
			case <-stopChan:
				return
			}
		}
	}

	initialCount := "0"
	newTestSecret := func(name string) *corev1.Secret {
		return &corev1.Secret{
			ObjectMeta: metav1.ObjectMeta{
				Name:      name,
				Namespace: namespaceObject.Name,
				Annotations: map[string]string{
					"count": initialCount,
				},
			},
			Data: map[string][]byte{
				"data": []byte("value1\n"),
			},
		}
	}

	tt := []struct {
		name                string
		succeed             bool
		secret              *corev1.Secret
		getWatcher          func(c *kubernetes.Clientset, secret *corev1.Secret) (watch.Interface, error, func())
		normalizeOutputFunc func(referenceOutput []string) []string
	}{
		{
			name:    "regular watcher should fail",
			succeed: false,
			secret:  newTestSecret("secret-01"),
			getWatcher: func(c *kubernetes.Clientset, secret *corev1.Secret) (watch.Interface, error, func()) {
				options := metav1.ListOptions{
					ResourceVersion: secret.ResourceVersion,
				}
				w, err := getWatchFunc(c, secret)(options)
				return w, err, noop
			}, // regular watcher; unfortunately destined to fail
			normalizeOutputFunc: noopNormalization,
		},
		{
			name:    "RetryWatcher survives closed watches",
			succeed: true,
			secret:  newTestSecret("secret-02"),
			getWatcher: func(c *kubernetes.Clientset, secret *corev1.Secret) (watch.Interface, error, func()) {
				lw := &cache.ListWatch{
					WatchFunc: func(options metav1.ListOptions) (watch.Interface, error) {
						return getWatchFunc(c, secret)(options)
					},
				}
				w, err := watchtools.NewRetryWatcher(secret.ResourceVersion, lw)
				return w, err, func() { <-w.Done() }
			},
			normalizeOutputFunc: noopNormalization,
		},
		{
			name:    "InformerWatcher survives closed watches",
			succeed: true,
			secret:  newTestSecret("secret-03"),
			getWatcher: func(c *kubernetes.Clientset, secret *corev1.Secret) (watch.Interface, error, func()) {
				lw := &cache.ListWatch{
					ListFunc: func(options metav1.ListOptions) (runtime.Object, error) {
						return getListFunc(c, secret)(options), nil
					},
					WatchFunc: func(options metav1.ListOptions) (watch.Interface, error) {
						return getWatchFunc(c, secret)(options)
					},
				}
				_, _, w, done := watchtools.NewIndexerInformerWatcher(lw, &corev1.Secret{})
				return w, nil, func() { <-done }
			},
			normalizeOutputFunc: normalizeInformerOutputFunc(initialCount),
		},
	}

	for _, tmptc := range tt {
		tc := tmptc // we need to copy it for parallel runs
		t.Run(tc.name, func(t *testing.T) {
			c, err := kubernetes.NewForConfig(config)
			if err != nil {
				t.Fatalf("Failed to create clientset: %v", err)
			}

			secret, err := c.CoreV1().Secrets(tc.secret.Namespace).Create(context.TODO(), tc.secret, metav1.CreateOptions{})
			if err != nil {
				t.Fatalf("Failed to create testing secret %s/%s: %v", tc.secret.Namespace, tc.secret.Name, err)
			}

			watcher, err, doneFn := tc.getWatcher(c, secret)
			if err != nil {
				t.Fatalf("Failed to create watcher: %v", err)
			}
			defer doneFn()

			var referenceOutput []string
			var output []string
			stopChan := make(chan struct{})
			stoppedChan := make(chan struct{})
			go generateEvents(t, c, secret, &referenceOutput, stopChan, stoppedChan)

			// Record current time to be able to asses if the timeout has been reached
			startTime := time.Now()
			ctx, cancel := watchtools.ContextWithOptionalTimeout(context.Background(), timeout)
			defer cancel()
			_, err = watchtools.UntilWithoutRetry(ctx, watcher, func(event watch.Event) (bool, error) {
				s, ok := event.Object.(*corev1.Secret)
				if !ok {
					t.Fatalf("Received an object that is not a Secret: %#v", event.Object)
				}
				output = append(output, s.Annotations["count"])
				// Watch will never end voluntarily
				return false, nil
			})
			watchDuration := time.Since(startTime)
			close(stopChan)
			<-stoppedChan

			output = tc.normalizeOutputFunc(output)

			t.Logf("Watch duration: %v; timeout: %v", watchDuration, timeout)

			if err == nil && !tc.succeed {
				t.Fatalf("Watch should have timed out but it exited without an error!")
			}

			if err != wait.ErrWaitTimeout && tc.succeed {
				t.Fatalf("Watch exited with error: %v!", err)
			}

			if watchDuration < timeout && tc.succeed {
				t.Fatalf("Watch should have timed out after %v but it timed out prematurely after %v!", timeout, watchDuration)
			}

			if watchDuration >= timeout && !tc.succeed {
				t.Fatalf("Watch should have timed out but it succeeded!")
			}

			if tc.succeed && !reflect.DeepEqual(referenceOutput, output) {
				t.Fatalf("Reference and real output differ! We must have lost some events or read some multiple times!\nRef:  %#v\nReal: %#v", referenceOutput, output)
			}
		})
	}
}
<commit_msg>Shorten watch restart test, run in parallel<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package apimachinery

import (
	"context"
	"fmt"
	"reflect"
	"testing"
	"time"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/kubernetes"
	restclient "k8s.io/client-go/rest"
	"k8s.io/client-go/tools/cache"
	watchtools "k8s.io/client-go/tools/watch"
	"k8s.io/kubernetes/test/integration/framework"
)

func noopNormalization(output []string) []string {
	return output
}

func normalizeInformerOutputFunc(initialVal string) func(output []string) []string {
	return func(output []string) []string {
		result := make([]string, 0, len(output))

		// Removes initial value and all of its direct repetitions
		lastVal := initialVal
		for _, v := range output {
			// Make values unique as informer(List+Watch) duplicates some events
			if v == lastVal {
				continue
			}
			result = append(result, v)
			lastVal = v
		}

		return result
	}
}

func noop() {}

func TestWatchRestartsIfTimeoutNotReached(t *testing.T) {
	// Has to be longer than 5 seconds
	timeout := 30 * time.Second

	// Set up a master
	masterConfig := framework.NewIntegrationTestMasterConfig()
	// Timeout is set random between MinRequestTimeout and 2x
	masterConfig.GenericConfig.MinRequestTimeout = int(timeout.Seconds()) / 4
	_, s, closeFn := framework.RunAMaster(masterConfig)
	defer closeFn()

	config := &restclient.Config{
		Host: s.URL,
	}

	namespaceObject := framework.CreateTestingNamespace("retry-watch", s, t)
	defer framework.DeleteTestingNamespace(namespaceObject, s, t)

	getListFunc := func(c *kubernetes.Clientset, secret *corev1.Secret) func(options metav1.ListOptions) *corev1.SecretList {
		return func(options metav1.ListOptions) *corev1.SecretList {
			options.FieldSelector = fields.OneTermEqualSelector("metadata.name", secret.Name).String()
			res, err := c.CoreV1().Secrets(secret.Namespace).List(context.TODO(), options)
			if err != nil {
				t.Fatalf("Failed to list Secrets: %v", err)
			}
			return res
		}
	}

	getWatchFunc := func(c *kubernetes.Clientset, secret *corev1.Secret) func(options metav1.ListOptions) (watch.Interface, error) {
		return func(options metav1.ListOptions) (watch.Interface, error) {
			options.FieldSelector = fields.OneTermEqualSelector("metadata.name", secret.Name).String()
			res, err := c.CoreV1().Secrets(secret.Namespace).Watch(context.TODO(), options)
			if err != nil {
				t.Fatalf("Failed to create a watcher on Secrets: %v", err)
			}
			return res, err
		}
	}

	generateEvents := func(t *testing.T, c *kubernetes.Clientset, secret *corev1.Secret, referenceOutput *[]string, stopChan chan struct{}, stoppedChan chan struct{}) {
		defer close(stoppedChan)
		counter := 0

		// These 5 seconds are here to protect against a race at the end when we could write something there at the same time as watch.Until ends
		softTimeout := timeout - 5*time.Second
		if softTimeout < 0 {
			panic("Timeout has to be grater than 5 seconds!")
		}
		endChannel := time.After(softTimeout)
		for {
			select {
			// TODO: get this lower once we figure out how to extend ETCD cache
			case <-time.After(1000 * time.Millisecond):
				counter = counter + 1

				patch := fmt.Sprintf(`{"metadata": {"annotations": {"count": "%d"}}}`, counter)
				_, err := c.CoreV1().Secrets(secret.Namespace).Patch(context.TODO(), secret.Name, types.StrategicMergePatchType, []byte(patch), metav1.PatchOptions{})
				if err != nil {
					t.Fatalf("Failed to patch secret: %v", err)
				}

				*referenceOutput = append(*referenceOutput, fmt.Sprintf("%d", counter))
			case <-endChannel:
				return
			case <-stopChan:
				return
			}
		}
	}

	initialCount := "0"
	newTestSecret := func(name string) *corev1.Secret {
		return &corev1.Secret{
			ObjectMeta: metav1.ObjectMeta{
				Name:      name,
				Namespace: namespaceObject.Name,
				Annotations: map[string]string{
					"count": initialCount,
				},
			},
			Data: map[string][]byte{
				"data": []byte("value1\n"),
			},
		}
	}

	tt := []struct {
		name                string
		succeed             bool
		secret              *corev1.Secret
		getWatcher          func(c *kubernetes.Clientset, secret *corev1.Secret) (watch.Interface, error, func())
		normalizeOutputFunc func(referenceOutput []string) []string
	}{
		{
			name:    "regular watcher should fail",
			succeed: false,
			secret:  newTestSecret("secret-01"),
			getWatcher: func(c *kubernetes.Clientset, secret *corev1.Secret) (watch.Interface, error, func()) {
				options := metav1.ListOptions{
					ResourceVersion: secret.ResourceVersion,
				}
				w, err := getWatchFunc(c, secret)(options)
				return w, err, noop
			}, // regular watcher; unfortunately destined to fail
			normalizeOutputFunc: noopNormalization,
		},
		{
			name:    "RetryWatcher survives closed watches",
			succeed: true,
			secret:  newTestSecret("secret-02"),
			getWatcher: func(c *kubernetes.Clientset, secret *corev1.Secret) (watch.Interface, error, func()) {
				lw := &cache.ListWatch{
					WatchFunc: func(options metav1.ListOptions) (watch.Interface, error) {
						return getWatchFunc(c, secret)(options)
					},
				}
				w, err := watchtools.NewRetryWatcher(secret.ResourceVersion, lw)
				return w, err, func() { <-w.Done() }
			},
			normalizeOutputFunc: noopNormalization,
		},
		{
			name:    "InformerWatcher survives closed watches",
			succeed: true,
			secret:  newTestSecret("secret-03"),
			getWatcher: func(c *kubernetes.Clientset, secret *corev1.Secret) (watch.Interface, error, func()) {
				lw := &cache.ListWatch{
					ListFunc: func(options metav1.ListOptions) (runtime.Object, error) {
						return getListFunc(c, secret)(options), nil
					},
					WatchFunc: func(options metav1.ListOptions) (watch.Interface, error) {
						return getWatchFunc(c, secret)(options)
					},
				}
				_, _, w, done := watchtools.NewIndexerInformerWatcher(lw, &corev1.Secret{})
				return w, nil, func() { <-done }
			},
			normalizeOutputFunc: normalizeInformerOutputFunc(initialCount),
		},
	}

	t.Run("group", func(t *testing.T) {
		for _, tmptc := range tt {
			tc := tmptc // we need to copy it for parallel runs
			t.Run(tc.name, func(t *testing.T) {
				t.Parallel()
				c, err := kubernetes.NewForConfig(config)
				if err != nil {
					t.Fatalf("Failed to create clientset: %v", err)
				}

				secret, err := c.CoreV1().Secrets(tc.secret.Namespace).Create(context.TODO(), tc.secret, metav1.CreateOptions{})
				if err != nil {
					t.Fatalf("Failed to create testing secret %s/%s: %v", tc.secret.Namespace, tc.secret.Name, err)
				}

				watcher, err, doneFn := tc.getWatcher(c, secret)
				if err != nil {
					t.Fatalf("Failed to create watcher: %v", err)
				}
				defer doneFn()

				var referenceOutput []string
				var output []string
				stopChan := make(chan struct{})
				stoppedChan := make(chan struct{})
				go generateEvents(t, c, secret, &referenceOutput, stopChan, stoppedChan)

				// Record current time to be able to asses if the timeout has been reached
				startTime := time.Now()
				ctx, cancel := watchtools.ContextWithOptionalTimeout(context.Background(), timeout)
				defer cancel()
				_, err = watchtools.UntilWithoutRetry(ctx, watcher, func(event watch.Event) (bool, error) {
					s, ok := event.Object.(*corev1.Secret)
					if !ok {
						t.Fatalf("Received an object that is not a Secret: %#v", event.Object)
					}
					output = append(output, s.Annotations["count"])
					// Watch will never end voluntarily
					return false, nil
				})
				watchDuration := time.Since(startTime)
				close(stopChan)
				<-stoppedChan

				output = tc.normalizeOutputFunc(output)

				t.Logf("Watch duration: %v; timeout: %v", watchDuration, timeout)

				if err == nil && !tc.succeed {
					t.Fatalf("Watch should have timed out but it exited without an error!")
				}

				if err != wait.ErrWaitTimeout && tc.succeed {
					t.Fatalf("Watch exited with error: %v!", err)
				}

				if watchDuration < timeout && tc.succeed {
					t.Fatalf("Watch should have timed out after %v but it timed out prematurely after %v!", timeout, watchDuration)
				}

				if watchDuration >= timeout && !tc.succeed {
					t.Fatalf("Watch should have timed out but it succeeded!")
				}

				if tc.succeed && !reflect.DeepEqual(referenceOutput, output) {
					t.Fatalf("Reference and real output differ! We must have lost some events or read some multiple times!\nRef:  %#v\nReal: %#v", referenceOutput, output)
				}
			})
		}
	})
}
<|endoftext|>


<commit_before>// Copyright 2015 The Prometheus 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.

// +build !nodevstat

package collector

import (
	"errors"
	"fmt"
	"sync"
	"unsafe"

	"github.com/prometheus/client_golang/prometheus"
)

// #cgo LDFLAGS: -ldevstat -lkvm
// #include "devstat_freebsd.h"
import "C"

const (
	devstatSubsystem = "devstat"
)

type devstatCollector struct {
	mu      sync.Mutex
	devinfo *C.struct_devinfo

	bytes       typedDesc
	bytes_total typedDesc
	transfers   typedDesc
	duration    typedDesc
	busyTime    typedDesc
	blocks      typedDesc
}

func init() {
	Factories["devstat"] = NewDevstatCollector
}

// Takes a prometheus registry and returns a new Collector exposing
// Device stats.
func NewDevstatCollector() (Collector, error) {
	return &devstatCollector{
		devinfo: &C.struct_devinfo{},
		bytes: typedDesc{prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, devstatSubsystem, "bytes_total"),
			"The total number of bytes in transactions.",
			[]string{"device", "type"}, nil,
		), prometheus.CounterValue},
		transfers: typedDesc{prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, devstatSubsystem, "transfers_total"),
			"The total number of transactions.",
			[]string{"device", "type"}, nil,
		), prometheus.CounterValue},
		duration: typedDesc{prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, devstatSubsystem, "duration_seconds_total"),
			"The total duration of transactions in seconds.",
			[]string{"device", "type"}, nil,
		), prometheus.CounterValue},
		busyTime: typedDesc{prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, devstatSubsystem, "busy_time_seconds_total"),
			"Total time the device had one or more transactions outstanding in seconds.",
			[]string{"device"}, nil,
		), prometheus.CounterValue},
		blocks: typedDesc{prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, devstatSubsystem, "blocks_transferred_total"),
			"The total number of blocks transferred.",
			[]string{"device"}, nil,
		), prometheus.CounterValue},
	}, nil
}

func (c *devstatCollector) Update(ch chan<- prometheus.Metric) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	var stats *C.Stats
	n := C._get_stats(c.devinfo, &stats)
	if n == -1 {
		return errors.New("devstat_getdevs failed")
	}

	base := unsafe.Pointer(stats)
	for i := C.int(0); i < n; i++ {
		offset := i * C.int(C.sizeof_Stats)
		stat := (*C.Stats)(unsafe.Pointer(uintptr(base) + uintptr(offset)))

		device := fmt.Sprintf("%s%d", C.GoString(&stat.device[0]), stat.unit)
		ch <- c.bytes.mustNewConstMetric(float64(stat.bytes.read), device, "read")
		ch <- c.bytes.mustNewConstMetric(float64(stat.bytes.write), device, "write")
		ch <- c.transfers.mustNewConstMetric(float64(stat.transfers.other), device, "other")
		ch <- c.transfers.mustNewConstMetric(float64(stat.transfers.read), device, "read")
		ch <- c.transfers.mustNewConstMetric(float64(stat.transfers.write), device, "write")
		ch <- c.duration.mustNewConstMetric(float64(stat.duration.other), device, "other")
		ch <- c.duration.mustNewConstMetric(float64(stat.duration.read), device, "read")
		ch <- c.duration.mustNewConstMetric(float64(stat.duration.write), device, "write")
		ch <- c.busyTime.mustNewConstMetric(float64(stat.busyTime), device)
		ch <- c.blocks.mustNewConstMetric(float64(stat.blocks), device)
	}
	C.free(unsafe.Pointer(stats))
	return nil
}
<commit_msg>Add missing dependency for static FreeBSD build<commit_after>// Copyright 2015 The Prometheus 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.

// +build !nodevstat

package collector

import (
	"errors"
	"fmt"
	"sync"
	"unsafe"

	"github.com/prometheus/client_golang/prometheus"
)

// #cgo LDFLAGS: -ldevstat -lkvm -lelf
// #include "devstat_freebsd.h"
import "C"

const (
	devstatSubsystem = "devstat"
)

type devstatCollector struct {
	mu      sync.Mutex
	devinfo *C.struct_devinfo

	bytes       typedDesc
	bytes_total typedDesc
	transfers   typedDesc
	duration    typedDesc
	busyTime    typedDesc
	blocks      typedDesc
}

func init() {
	Factories["devstat"] = NewDevstatCollector
}

// Takes a prometheus registry and returns a new Collector exposing
// Device stats.
func NewDevstatCollector() (Collector, error) {
	return &devstatCollector{
		devinfo: &C.struct_devinfo{},
		bytes: typedDesc{prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, devstatSubsystem, "bytes_total"),
			"The total number of bytes in transactions.",
			[]string{"device", "type"}, nil,
		), prometheus.CounterValue},
		transfers: typedDesc{prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, devstatSubsystem, "transfers_total"),
			"The total number of transactions.",
			[]string{"device", "type"}, nil,
		), prometheus.CounterValue},
		duration: typedDesc{prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, devstatSubsystem, "duration_seconds_total"),
			"The total duration of transactions in seconds.",
			[]string{"device", "type"}, nil,
		), prometheus.CounterValue},
		busyTime: typedDesc{prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, devstatSubsystem, "busy_time_seconds_total"),
			"Total time the device had one or more transactions outstanding in seconds.",
			[]string{"device"}, nil,
		), prometheus.CounterValue},
		blocks: typedDesc{prometheus.NewDesc(
			prometheus.BuildFQName(Namespace, devstatSubsystem, "blocks_transferred_total"),
			"The total number of blocks transferred.",
			[]string{"device"}, nil,
		), prometheus.CounterValue},
	}, nil
}

func (c *devstatCollector) Update(ch chan<- prometheus.Metric) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	var stats *C.Stats
	n := C._get_stats(c.devinfo, &stats)
	if n == -1 {
		return errors.New("devstat_getdevs failed")
	}

	base := unsafe.Pointer(stats)
	for i := C.int(0); i < n; i++ {
		offset := i * C.int(C.sizeof_Stats)
		stat := (*C.Stats)(unsafe.Pointer(uintptr(base) + uintptr(offset)))

		device := fmt.Sprintf("%s%d", C.GoString(&stat.device[0]), stat.unit)
		ch <- c.bytes.mustNewConstMetric(float64(stat.bytes.read), device, "read")
		ch <- c.bytes.mustNewConstMetric(float64(stat.bytes.write), device, "write")
		ch <- c.transfers.mustNewConstMetric(float64(stat.transfers.other), device, "other")
		ch <- c.transfers.mustNewConstMetric(float64(stat.transfers.read), device, "read")
		ch <- c.transfers.mustNewConstMetric(float64(stat.transfers.write), device, "write")
		ch <- c.duration.mustNewConstMetric(float64(stat.duration.other), device, "other")
		ch <- c.duration.mustNewConstMetric(float64(stat.duration.read), device, "read")
		ch <- c.duration.mustNewConstMetric(float64(stat.duration.write), device, "write")
		ch <- c.busyTime.mustNewConstMetric(float64(stat.busyTime), device)
		ch <- c.blocks.mustNewConstMetric(float64(stat.blocks), device)
	}
	C.free(unsafe.Pointer(stats))
	return nil
}
<|endoftext|>


<commit_before>package tablestruct

import (
	"encoding/json"
	"io"
)

// Map describes a mapping between database tables and Go structs.
type Map []TableMap

// NewMap constructs a new mapping object.
func NewMap(in io.Reader) (*Map, error) {
	var mapper Map
	if err := json.NewDecoder(in).Decode(&mapper); err != nil {
		return nil, err
	}
	return &mapper, nil
}

// Imports generates list of import specs required by generated code.
func (m *Map) Imports() []importSpec {
	return []importSpec{
		{"database/sql", ""},
		{"log", ""},
		{"github.com/lib/pq", "_"},
	}
}
<commit_msg>Remove lib/pq import<commit_after>package tablestruct

import (
	"encoding/json"
	"io"
)

// Map describes a mapping between database tables and Go structs.
type Map []TableMap

// NewMap constructs a new mapping object.
func NewMap(in io.Reader) (*Map, error) {
	var mapper Map
	if err := json.NewDecoder(in).Decode(&mapper); err != nil {
		return nil, err
	}
	return &mapper, nil
}

// Imports generates list of import specs required by generated code.
func (m *Map) Imports() []importSpec {
	return []importSpec{
		{"database/sql", ""},
		{"log", ""},
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"sync"

	"github.com/gorilla/mux"
	"github.com/stts-se/rbg2p"
)

type g2pMutex struct {
	g2ps  map[string]rbg2p.RuleSet
	mutex *sync.RWMutex
}

var g2p = g2pMutex{
	g2ps:  make(map[string]rbg2p.RuleSet),
	mutex: &sync.RWMutex{},
}

func g2pMain(w http.ResponseWriter, r *http.Request) {
	// TODO error if file not found
	http.ServeFile(w, r, "./src/g2p_demo.html")
}

var wSplitRe = regexp.MustCompile(" *, *")

// Word internal struct for json
type Word struct {
	Orth    string   `json:"orth"`
	Transes []string `json:"transes"`
}

func transcribe(w http.ResponseWriter, r *http.Request) {

	lang := r.FormValue("lang")
	if "" == lang {
		msg := "no value for the expected 'lang' parameter"
		log.Println(msg)
		http.Error(w, msg, http.StatusBadRequest)
		return
	}

	words := r.FormValue("words")
	if "" == words {
		msg := "no value for the expected 'words' parameter"
		log.Println(msg)
		http.Error(w, msg, http.StatusBadRequest)
		return
	}
	words = strings.ToLower(words)

	g2p.mutex.RLock()
	defer g2p.mutex.RUnlock()
	ruleSet, ok := g2p.g2ps[lang]
	if !ok {
		msg := "unknown 'lang': " + lang
		langs := listLanguages()
		msg = fmt.Sprintf("%s. Known 'lang' values: %s", msg, strings.Join(langs, ", "))
		log.Println(msg)
		http.Error(w, msg, http.StatusBadRequest)
		return
	}

	res := []Word{}
	for _, orth := range wSplitRe.Split(words, -1) {
		transes, err := ruleSet.Apply(orth)
		if err != nil {
			msg := fmt.Sprintf("couldn't transcribe word : %v", err)
			log.Println(msg)
			http.Error(w, msg, http.StatusInternalServerError)
			return
		}
		tRes := []string{}
		for _, trans := range transes {
			tRes = append(tRes, strings.Join(trans.Phonemes, ruleSet.PhonemeDelimiter))
		}
		wRes := Word{orth, tRes}
		res = append(res, wRes)

	}
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	j, err := json.Marshal(res)
	if err != nil {
		msg := fmt.Sprintf("failed json marshalling : %v", err)
		log.Println(msg)
		http.Error(w, msg, http.StatusInternalServerError)
		return
	}
	fmt.Fprintf(w, string(j))
}

func listLanguages() []string {
	var res []string
	for name := range g2p.g2ps {
		res = append(res, name)
	}
	return res
}

func list(w http.ResponseWriter, r *http.Request) {
	g2p.mutex.RLock()
	res := listLanguages()
	g2p.mutex.RUnlock()

	sort.Strings(res)
	j, err := json.Marshal(res)
	if err != nil {
		msg := fmt.Sprintf("failed json marshalling : %v", err)
		log.Println(msg)
		http.Error(w, msg, http.StatusInternalServerError)
		return
	}
	fmt.Fprintf(w, string(j))
}

// langFromFilePath returns the base file name stripped from any '.g2p' extension
func langFromFilePath(p string) string {
	b := filepath.Base(p)
	if strings.HasSuffix(b, ".g2p") {
		b = b[0 : len(b)-4]
	}
	return b
}

func main() {

	if len(os.Args) != 2 {
		fmt.Fprintf(os.Stderr, "g2pserver <G2P FILES DIR>\n")
		os.Exit(0)
	}

	// g2p file dir. Each file in dir with .g2p extension
	// is treated as a g2p file
	var dir = os.Args[1]

	files, err := ioutil.ReadDir(dir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(0)
	}

	// populate map of g2p rules from files.
	// The base file name minus '.g2p' is the language name.
	var fn string
	for _, f := range files {
		fn = filepath.Join(dir, f.Name())
		if !strings.HasSuffix(fn, ".g2p") {
			fmt.Fprintf(os.Stderr, "g2pserver: skipping file: '%s'\n", fn)
			continue
		}

		ruleSet, err := rbg2p.LoadFile(fn)
		if err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err)
			fmt.Fprintf(os.Stderr, "g2pserver: skipping file: '%s'\n", fn)
			continue
		}

		lang := langFromFilePath(fn)
		g2p.mutex.Lock()
		g2p.g2ps[lang] = ruleSet
		g2p.mutex.Unlock()
		fmt.Fprintf(os.Stderr, "g2p server: loaded file '%s'\n", fn)
	}

	r := mux.NewRouter().StrictSlash(true)

	r.HandleFunc("/rbg2p", g2pMain).Methods("get")
	r.HandleFunc("/rbg2p/transcribe", transcribe).Methods("get", "post")
	r.HandleFunc("/rbg2p/list", list).Methods("get", "post")

	port := ":6771"
	log.Printf("starting g2p server at port %s\n", port)
	err = http.ListenAndServe(port, r)
	if err != nil {

		log.Fatalf("no fun: %v\n", err)
	}

}
<commit_msg>added handler for API call to root (/)<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"sync"

	"github.com/gorilla/mux"
	"github.com/stts-se/rbg2p"
)

type g2pMutex struct {
	g2ps  map[string]rbg2p.RuleSet
	mutex *sync.RWMutex
}

var g2p = g2pMutex{
	g2ps:  make(map[string]rbg2p.RuleSet),
	mutex: &sync.RWMutex{},
}

func g2pMain(w http.ResponseWriter, r *http.Request) {
	// TODO error if file not found
	http.ServeFile(w, r, "./src/g2p_demo.html")
}

var wSplitRe = regexp.MustCompile(" *, *")

// Word internal struct for json
type Word struct {
	Orth    string   `json:"orth"`
	Transes []string `json:"transes"`
}

func transcribe(w http.ResponseWriter, r *http.Request) {

	lang := r.FormValue("lang")
	if "" == lang {
		msg := "no value for the expected 'lang' parameter"
		log.Println(msg)
		http.Error(w, msg, http.StatusBadRequest)
		return
	}

	words := r.FormValue("words")
	if "" == words {
		msg := "no value for the expected 'words' parameter"
		log.Println(msg)
		http.Error(w, msg, http.StatusBadRequest)
		return
	}
	words = strings.ToLower(words)

	g2p.mutex.RLock()
	defer g2p.mutex.RUnlock()
	ruleSet, ok := g2p.g2ps[lang]
	if !ok {
		msg := "unknown 'lang': " + lang
		langs := listLanguages()
		msg = fmt.Sprintf("%s. Known 'lang' values: %s", msg, strings.Join(langs, ", "))
		log.Println(msg)
		http.Error(w, msg, http.StatusBadRequest)
		return
	}

	res := []Word{}
	for _, orth := range wSplitRe.Split(words, -1) {
		transes, err := ruleSet.Apply(orth)
		if err != nil {
			msg := fmt.Sprintf("couldn't transcribe word : %v", err)
			log.Println(msg)
			http.Error(w, msg, http.StatusInternalServerError)
			return
		}
		tRes := []string{}
		for _, trans := range transes {
			tRes = append(tRes, strings.Join(trans.Phonemes, ruleSet.PhonemeDelimiter))
		}
		wRes := Word{orth, tRes}
		res = append(res, wRes)

	}
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	j, err := json.Marshal(res)
	if err != nil {
		msg := fmt.Sprintf("failed json marshalling : %v", err)
		log.Println(msg)
		http.Error(w, msg, http.StatusInternalServerError)
		return
	}
	fmt.Fprintf(w, string(j))
}

func listLanguages() []string {
	var res []string
	for name := range g2p.g2ps {
		res = append(res, name)
	}
	return res
}

func list(w http.ResponseWriter, r *http.Request) {
	g2p.mutex.RLock()
	res := listLanguages()
	g2p.mutex.RUnlock()

	sort.Strings(res)
	j, err := json.Marshal(res)
	if err != nil {
		msg := fmt.Sprintf("failed json marshalling : %v", err)
		log.Println(msg)
		http.Error(w, msg, http.StatusInternalServerError)
		return
	}
	fmt.Fprintf(w, string(j))
}

// langFromFilePath returns the base file name stripped from any '.g2p' extension
func langFromFilePath(p string) string {
	b := filepath.Base(p)
	if strings.HasSuffix(b, ".g2p") {
		b = b[0 : len(b)-4]
	}
	return b
}

func main() {

	if len(os.Args) != 2 {
		fmt.Fprintf(os.Stderr, "g2pserver <G2P FILES DIR>\n")
		os.Exit(0)
	}

	// g2p file dir. Each file in dir with .g2p extension
	// is treated as a g2p file
	var dir = os.Args[1]

	files, err := ioutil.ReadDir(dir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		os.Exit(0)
	}

	// populate map of g2p rules from files.
	// The base file name minus '.g2p' is the language name.
	var fn string
	for _, f := range files {
		fn = filepath.Join(dir, f.Name())
		if !strings.HasSuffix(fn, ".g2p") {
			fmt.Fprintf(os.Stderr, "g2pserver: skipping file: '%s'\n", fn)
			continue
		}

		ruleSet, err := rbg2p.LoadFile(fn)
		if err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err)
			fmt.Fprintf(os.Stderr, "g2pserver: skipping file: '%s'\n", fn)
			continue
		}

		lang := langFromFilePath(fn)
		g2p.mutex.Lock()
		g2p.g2ps[lang] = ruleSet
		g2p.mutex.Unlock()
		fmt.Fprintf(os.Stderr, "g2p server: loaded file '%s'\n", fn)
	}

	r := mux.NewRouter().StrictSlash(true)

	r.HandleFunc("/rbg2p", g2pMain).Methods("get")
	r.HandleFunc("/", g2pMain).Methods("get")
	r.HandleFunc("/rbg2p/transcribe", transcribe).Methods("get", "post")
	r.HandleFunc("/rbg2p/list", list).Methods("get", "post")

	port := ":6771"
	log.Printf("starting g2p server at port %s\n", port)
	err = http.ListenAndServe(port, r)
	if err != nil {

		log.Fatalf("no fun: %v\n", err)
	}

}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2015 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 *
 */

package elasticseach

import (
	"encoding/json"
	"errors"
	"strconv"

	elastigo "github.com/mattbaird/elastigo/lib"

	"github.com/redhat-cip/skydive/config"
	"github.com/redhat-cip/skydive/flow"
	"github.com/redhat-cip/skydive/logging"
	"github.com/redhat-cip/skydive/storage"
)

type OldFlow struct {
	UUID       string
	LayersPath string
	EtherSrc   string
	EtherDst   string
	Ipv4Src    string
	Ipv4Dst    string
	PortSrc    uint32
	PortDst    uint32
	ID         uint64
	Timestamp  uint64

	ProbeGraphPath string

	IfSrcName      string
	IfSrcType      string
	IfSrcGraphPath string
	IfSrcTenantID  string
	IfSrcVNI       uint64

	IfDstName      string
	IfDstType      string
	IfDstGraphPath string
	IfDstTenantID  string
	IfDstVNI       uint64
}

func flow2OldFlow(f *flow.Flow) OldFlow {
	fs := f.GetStatistics()
	eth := fs.Endpoints[flow.FlowEndpointType_ETHERNET.Value()]
	ip := fs.Endpoints[flow.FlowEndpointType_IPV4.Value()]
	port := fs.Endpoints[flow.FlowEndpointType_TCPPORT.Value()]
	if port != nil {
		port = fs.Endpoints[flow.FlowEndpointType_UDPPORT.Value()]
		if port != nil {
			port = fs.Endpoints[flow.FlowEndpointType_SCTPPORT.Value()]
		}
	}

	of := OldFlow{}
	of.UUID = f.UUID
	of.LayersPath = f.LayersPath
	of.EtherSrc = eth.AB.Value
	of.EtherDst = eth.BA.Value
	of.Ipv4Src = ""
	of.Ipv4Dst = ""
	of.PortSrc = 0
	of.PortDst = 0
	if ip != nil {
		of.Ipv4Src = ip.AB.Value
		of.Ipv4Dst = ip.BA.Value
	}
	if port != nil {
		portInt, _ := strconv.Atoi(port.AB.Value)
		of.PortSrc = uint32(portInt)
		portInt, _ = strconv.Atoi(port.BA.Value)
		of.PortDst = uint32(portInt)
	}
	of.ID = 0
	of.Timestamp = uint64(fs.Start)

	of.ProbeGraphPath = f.ProbeGraphPath

	of.IfSrcName = ""
	of.IfSrcType = ""
	of.IfSrcGraphPath = ""
	of.IfSrcTenantID = ""
	of.IfSrcVNI = 0

	of.IfDstName = ""
	of.IfDstType = ""
	of.IfDstGraphPath = ""
	of.IfDstTenantID = ""
	of.IfDstVNI = 0
	return of
}

type ElasticSearchStorage struct {
	connection *elastigo.Conn
}

func (c *ElasticSearchStorage) StoreFlows(flows []*flow.Flow) error {
	/* TODO(safchain) bulk insert */
	for _, flow := range flows {
		j, err := json.Marshal(flow)
		if err == nil {
			logging.GetLogger().Debug("Indexing: %s", string(j))
		}

		_, err = c.connection.Index("skydive", "flow", flow.UUID, nil, flow2OldFlow(flow))
		if err != nil {
			logging.GetLogger().Error("Error while indexing: %s", err.Error())
			continue
		}
	}

	return nil
}

func (c *ElasticSearchStorage) SearchFlows(filters storage.Filters) ([]*flow.Flow, error) {
	query := map[string]interface{}{
		"sort": map[string]interface{}{
			"Timestamp": map[string]string{
				"order": "desc",
			},
		},
		"from": 0,
		"size": 5,
	}
	if len(filters) > 0 {
		query = map[string]interface{}{
			"query": map[string]interface{}{
				"term": filters,
			},
			"sort": map[string]interface{}{
				"Timestamp": map[string]string{
					"order": "desc",
				},
			},
			"from": 0,
			"size": 5,
		}
	}

	q, err := json.Marshal(query)
	if err != nil {
		return nil, err
	}

	out, err := c.connection.Search("skydive", "flow", nil, string(q))
	if err != nil {
		return nil, err
	}

	flows := []*flow.Flow{}

	if out.Hits.Len() > 0 {
		for _, d := range out.Hits.Hits {
			f := new(flow.Flow)
			err := json.Unmarshal([]byte(*d.Source), f)
			if err != nil {
				return nil, err
			}

			flows = append(flows, f)
		}
	}

	return flows, nil
}

func (c *ElasticSearchStorage) initialize() error {
	req, err := c.connection.NewRequest("GET", "/skydive", "")
	if err != nil {
		return err
	}

	var response map[string]interface{}
	code, _, _ := req.Do(&response)
	if code == 200 {
		return nil
	}

	// template to remove the analyzer
	req, err = c.connection.NewRequest("PUT", "/skydive", "")
	if err != nil {
		return err
	}

	body := `{"mappings":{"flow":{"dynamic_templates":[{"notanalyzed":{"match":"*","mapping":{"type":"string","index":"not_analyzed"}}}]}}}`
	req.SetBodyString(body)

	code, _, err = req.Do(&response)
	if err != nil {
		return err
	}

	if code != 200 {
		return errors.New("Unable to create the skydive index: " + strconv.FormatInt(int64(code), 10))
	}

	return nil
}

var ErrBadConfig = errors.New("elasticseach : Config file is misconfigured, check elasticsearch key format")

func New() (*ElasticSearchStorage, error) {
	c := elastigo.NewConn()

	elasticonfig := config.GetConfig().Section("storage").Key("elasticsearch").Strings(":")
	if len(elasticonfig) != 2 {
		return nil, ErrBadConfig
	}
	c.Domain = elasticonfig[0]
	c.Port = elasticonfig[1]

	storage := &ElasticSearchStorage{connection: c}

	err := storage.initialize()
	if err != nil {
		return nil, err
	}

	return storage, nil
}
<commit_msg>[storage][elasticsearch] New FlatFlow format<commit_after>/*
 * Copyright (C) 2015 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 *
 */

package elasticseach

import (
	"encoding/json"
	"errors"
	"strconv"

	elastigo "github.com/mattbaird/elastigo/lib"

	"github.com/redhat-cip/skydive/config"
	"github.com/redhat-cip/skydive/flow"
	"github.com/redhat-cip/skydive/logging"
	"github.com/redhat-cip/skydive/storage"
)

type FlatFlowEndpointStatistics struct {
	Type    flow.FlowEndpointType
	Value   string
	Packets uint64
	Bytes   uint64
}

type FlatFlowEndpointsStatistics struct {
	Type flow.FlowEndpointType
	AB   FlatFlowEndpointStatistics /* A->B */
	BA   FlatFlowEndpointStatistics /* B->A */
}

type FlatFlowStatistics struct {
	Start     int64
	Last      int64
	Endpoints []FlatFlowEndpointsStatistics
}

type FlatFlow struct {
	UUID       string
	LayersPath string
	/* Data Flow info */
	Statistics FlatFlowStatistics
	/* Topology info */
	ProbeGraphPath string
	IfSrcGraphPath string
	IfDstGraphPath string
}

func flow2FlatFlow(f *flow.Flow) FlatFlow {
	ff := FlatFlow{}

	ff.UUID = f.UUID
	ff.LayersPath = f.LayersPath
	ff.ProbeGraphPath = f.ProbeGraphPath
	ff.IfSrcGraphPath = f.IfSrcGraphPath
	ff.IfDstGraphPath = f.IfDstGraphPath

	fs := f.GetStatistics()
	ffs := FlatFlowStatistics{}
	ffs.Start = fs.Start
	ffs.Last = fs.Last
	for _, endp := range fs.Endpoints {
		ffes := FlatFlowEndpointsStatistics{}
		ffes.Type = endp.Type
		ffes.AB.Type = endp.Type
		ffes.AB.Value = endp.AB.Value
		ffes.AB.Packets = endp.AB.Packets
		ffes.AB.Bytes = endp.AB.Bytes
		ffes.BA.Type = endp.Type
		ffes.BA.Value = endp.BA.Value
		ffes.BA.Packets = endp.BA.Packets
		ffes.BA.Bytes = endp.BA.Bytes
		ffs.Endpoints = append(ffs.Endpoints, ffes)
	}
	ff.Statistics = ffs
	return ff
}

type ElasticSearchStorage struct {
	connection *elastigo.Conn
}

func (c *ElasticSearchStorage) StoreFlows(flows []*flow.Flow) error {
	/* TODO(safchain) bulk insert */
	for _, flow := range flows {
		_, err := c.connection.Index("skydive", "flow", flow.UUID, nil, flow2FlatFlow(flow))
		if err != nil {
			logging.GetLogger().Error("Error while indexing: %s", err.Error())
			continue
		}
	}

	return nil
}

func (c *ElasticSearchStorage) SearchFlows(filters storage.Filters) ([]*flow.Flow, error) {
	query := map[string]interface{}{
		"sort": map[string]interface{}{
			"Timestamp": map[string]string{
				"order": "desc",
			},
		},
		"from": 0,
		"size": 5,
	}
	if len(filters) > 0 {
		query = map[string]interface{}{
			"query": map[string]interface{}{
				"term": filters,
			},
			"sort": map[string]interface{}{
				"Timestamp": map[string]string{
					"order": "desc",
				},
			},
			"from": 0,
			"size": 5,
		}
	}

	q, err := json.Marshal(query)
	if err != nil {
		return nil, err
	}

	out, err := c.connection.Search("skydive", "flow", nil, string(q))
	if err != nil {
		return nil, err
	}

	flows := []*flow.Flow{}

	if out.Hits.Len() > 0 {
		for _, d := range out.Hits.Hits {
			f := new(flow.Flow)
			err := json.Unmarshal([]byte(*d.Source), f)
			if err != nil {
				return nil, err
			}

			flows = append(flows, f)
		}
	}

	return flows, nil
}

func (c *ElasticSearchStorage) initialize() error {
	req, err := c.connection.NewRequest("GET", "/skydive", "")
	if err != nil {
		return err
	}

	var response map[string]interface{}
	code, _, _ := req.Do(&response)
	if code == 200 {
		return nil
	}

	// template to remove the analyzer
	req, err = c.connection.NewRequest("PUT", "/skydive", "")
	if err != nil {
		return err
	}

	body := `{"mappings":{"flow":{"dynamic_templates":[{"notanalyzed":{"match":"*","mapping":{"type":"string","index":"not_analyzed"}}}]}}}`
	req.SetBodyString(body)

	code, _, err = req.Do(&response)
	if err != nil {
		return err
	}

	if code != 200 {
		return errors.New("Unable to create the skydive index: " + strconv.FormatInt(int64(code), 10))
	}

	return nil
}

var ErrBadConfig = errors.New("elasticseach : Config file is misconfigured, check elasticsearch key format")

func New() (*ElasticSearchStorage, error) {
	c := elastigo.NewConn()

	elasticonfig := config.GetConfig().Section("storage").Key("elasticsearch").Strings(":")
	if len(elasticonfig) != 2 {
		return nil, ErrBadConfig
	}
	c.Domain = elasticonfig[0]
	c.Port = elasticonfig[1]

	storage := &ElasticSearchStorage{connection: c}

	err := storage.initialize()
	if err != nil {
		return nil, err
	}

	return storage, nil
}
<|endoftext|>


<commit_before>package capn

import (
	"bytes"
	"errors"
	"io"
	"math"
)

var (
	errBufferCall     = errors.New("capn: can't call on a memory buffer")
	ErrInvalidSegment = errors.New("capn: invalid segment id")
	ErrTooMuchData    = errors.New("capn: too much data in stream")
)

type buffer Segment

// NewBuffer creates an expanding single segment buffer. Creating new objects
// will expand the buffer. Data can be nil (or length 0 with some capacity) if
// creating a new session. If parsing an existing segment than data should be
// the segment contents and will not be copied.
func NewBuffer(data []byte) *Segment {
	if uint64(len(data)) > uint64(math.MaxUint32) {
		return nil
	}

	b := &buffer{}
	b.Message = b
	b.Data = data
	return (*Segment)(b)
}

func (b *buffer) NewSegment(minsz int) (*Segment, error) {
	if uint64(len(b.Data)) > uint64(math.MaxUint32)-uint64(minsz) {
		return nil, ErrOverlarge
	}
	b.Data = append(b.Data, make([]byte, minsz)...)
	b.Data = b.Data[:len(b.Data)-minsz]
	return (*Segment)(b), nil
}

func (b *buffer) Lookup(segid uint32) (*Segment, error) {
	if segid == 0 {
		return (*Segment)(b), nil
	} else {
		return nil, ErrInvalidSegment
	}
}

type MultiBuffer struct {
	segments []*Segment
}

// NewmultiBuffer creates a new multi segment message. Creating new objects
// will try and reuse the buffers available, but will create new ones if there
// is insufficient capacity. When parsing an existing message data should be
// the list of segments. The data buffers will not be copied.
func NewMultiBuffer(data [][]byte) *Segment {
	m := &MultiBuffer{make([]*Segment, len(data))}
	for i, d := range data {
		m.segments[i] = &Segment{m, d, uint32(i)}
	}
	if len(data) > 0 {
		return m.segments[0]
	}
	return &Segment{m, nil, 0xFFFFFFFF}
}

var (
	MaxSegmentNumber = 1024
	MaxTotalSize     = 1024 * 1024 * 1024
)

func (m *MultiBuffer) NewSegment(minsz int) (*Segment, error) {
	for _, s := range m.segments {
		if len(s.Data)+minsz <= cap(s.Data) {
			return s, nil
		}
	}

	if minsz < 4096 {
		minsz = 4096
	}
	s := &Segment{m, make([]byte, 0, minsz), uint32(len(m.segments))}
	m.segments = append(m.segments, s)
	return s, nil
}

func (m *MultiBuffer) Lookup(segid uint32) (*Segment, error) {
	if uint(segid) < uint(len(m.segments)) {
		return m.segments[segid], nil
	} else {
		return nil, ErrInvalidSegment
	}
}

// ReadFromStream reads a non-packed serialized stream from r. buf is used to
// buffer the read contents, can be nil, and is provided so that the buffer
// can be reused between messages. The returned segment is the first segment
// read, which contains the root pointer.
func ReadFromStream(r io.Reader, buf *bytes.Buffer) (*Segment, error) {
	if buf == nil {
		buf = new(bytes.Buffer)
	} else {
		buf.Reset()
	}

	if _, err := io.CopyN(buf, r, 4); err != nil {
		return nil, err
	}

	if little32(buf.Bytes()[:]) >= uint32(MaxSegmentNumber) {
		return nil, ErrTooMuchData
	}

	segnum := int(little32(buf.Bytes()[:]) + 1)
	hdrsz := 8*(segnum/2) + 4

	if _, err := io.CopyN(buf, r, int64(hdrsz)); err != nil {
		return nil, err
	}

	total := 0
	for i := 0; i < segnum; i++ {
		sz := little32(buf.Bytes()[4*i+4:])
		if uint64(total)+uint64(sz)*8 > uint64(MaxTotalSize) {
			return nil, ErrTooMuchData
		}
		total += int(sz) * 8
	}

	if _, err := io.CopyN(buf, r, int64(total)); err != nil {
		return nil, err
	}

	hdrv := buf.Bytes()[4 : hdrsz+4]
	datav := buf.Bytes()[hdrsz+4:]

	if segnum == 1 {
		sz := int(little32(hdrv)) * 8
		return NewBuffer(datav[:sz]), nil
	}

	m := &MultiBuffer{make([]*Segment, segnum)}
	for i := 0; i < segnum; i++ {
		sz := int(little32(hdrv[4*i:])) * 8
		m.segments[i] = &Segment{m, datav[:sz], uint32(i)}
		datav = datav[sz:]
	}

	return m.segments[0], nil
}

// ReadFromMemoryZeroCopy: like ReadFromStream, but reads a non-packed
// serialized stream that already resides in memory in the argument data.
// The returned segment is the first segment read, which contains
// the root pointer. The returned bytesRead says how many bytes were
// consumed from data in making seg. The caller should advance the
// data slice by doing data = data[bytesRead:] between successive calls
// to ReadFromMemoryZeroCopy().
func ReadFromMemoryZeroCopy(data []byte) (seg *Segment, bytesRead int64, err error) {

	if len(data) < 4 {
		return nil, 0, io.EOF
	}

	if little32(data[0:4]) >= uint32(MaxSegmentNumber) {
		return nil, 0, ErrTooMuchData
	}

	segnum := int(little32(data[0:4]) + 1)
	hdrsz := 8*(segnum/2) + 4

	b := data[0:(hdrsz + 4)]

	total := 0
	for i := 0; i < segnum; i++ {
		sz := little32(b[4*i+4:])
		if uint64(total)+uint64(sz)*8 > uint64(MaxTotalSize) {
			return nil, 0, ErrTooMuchData
		}
		total += int(sz) * 8
	}
	if total == 0 {
		return nil, 0, io.EOF
	}

	hdrv := data[4:(hdrsz + 4)]
	datav := data[hdrsz+4:]
	m := &MultiBuffer{make([]*Segment, segnum)}
	for i := 0; i < segnum; i++ {
		sz := int(little32(hdrv[4*i:])) * 8
		m.segments[i] = &Segment{m, datav[:sz], uint32(i)}
		datav = datav[sz:]
	}

	return m.segments[0], int64(4 + hdrsz + total), nil
}

// WriteTo writes the message that the segment is part of to the
// provided stream in serialized form.
func (s *Segment) WriteTo(w io.Writer) (int64, error) {
	segnum := uint32(1)
	for {
		if seg, _ := s.Message.Lookup(segnum); seg == nil {
			break
		}
		segnum++
	}

	hdrv := make([]uint8, 8*(segnum/2)+8)
	putLittle32(hdrv, segnum-1)
	for i := uint32(0); i < segnum; i++ {
		seg, _ := s.Message.Lookup(i)
		putLittle32(hdrv[4*i+4:], uint32(len(seg.Data)/8))
	}

	if n, err := w.Write(hdrv); err != nil {
		return int64(n), err
	}
	written := int64(len(hdrv))

	for i := uint32(0); i < segnum; i++ {
		seg, _ := s.Message.Lookup(i)
		if n, err := w.Write(seg.Data); err != nil {
			return written + int64(n), err
		} else {
			written += int64(n)
		}
	}

	return written, nil
}
<commit_msg>refer to github.com/glycerine/go-capnproto in the generated bindings<commit_after>package capn

import (
	"bytes"
	"errors"
	"io"
	"math"
)

var (
	errBufferCall     = errors.New("capn: can't call on a memory buffer")
	ErrInvalidSegment = errors.New("capn: invalid segment id")
	ErrTooMuchData    = errors.New("capn: too much data in stream")
)

type buffer Segment

// NewBuffer creates an expanding single segment buffer. Creating new objects
// will expand the buffer. Data can be nil (or length 0 with some capacity) if
// creating a new session. If parsing an existing segment than data should be
// the segment contents and will not be copied.
func NewBuffer(data []byte) *Segment {
	if uint64(len(data)) > uint64(math.MaxUint32) {
		return nil
	}

	b := &buffer{}
	b.Message = b
	b.Data = data
	return (*Segment)(b)
}

func (b *buffer) NewSegment(minsz int) (*Segment, error) {
	if uint64(len(b.Data)) > uint64(math.MaxUint32)-uint64(minsz) {
		return nil, ErrOverlarge
	}
	b.Data = append(b.Data, make([]byte, minsz)...)
	b.Data = b.Data[:len(b.Data)-minsz]
	return (*Segment)(b), nil
}

func (b *buffer) Lookup(segid uint32) (*Segment, error) {
	if segid == 0 {
		return (*Segment)(b), nil
	} else {
		return nil, ErrInvalidSegment
	}
}

type multiBuffer struct {
	segments []*Segment
}

// NewmultiBuffer creates a new multi segment message. Creating new objects
// will try and reuse the buffers available, but will create new ones if there
// is insufficient capacity. When parsing an existing message data should be
// the list of segments. The data buffers will not be copied.
func NewmultiBuffer(data [][]byte) *Segment {
	m := &multiBuffer{make([]*Segment, len(data))}
	for i, d := range data {
		m.segments[i] = &Segment{m, d, uint32(i)}
	}
	if len(data) > 0 {
		return m.segments[0]
	}
	return &Segment{m, nil, 0xFFFFFFFF}
}

var (
	MaxSegmentNumber = 1024
	MaxTotalSize     = 1024 * 1024 * 1024
)

func (m *multiBuffer) NewSegment(minsz int) (*Segment, error) {
	for _, s := range m.segments {
		if len(s.Data)+minsz <= cap(s.Data) {
			return s, nil
		}
	}

	if minsz < 4096 {
		minsz = 4096
	}
	s := &Segment{m, make([]byte, 0, minsz), uint32(len(m.segments))}
	m.segments = append(m.segments, s)
	return s, nil
}

func (m *multiBuffer) Lookup(segid uint32) (*Segment, error) {
	if uint(segid) < uint(len(m.segments)) {
		return m.segments[segid], nil
	} else {
		return nil, ErrInvalidSegment
	}
}

// ReadFromStream reads a non-packed serialized stream from r. buf is used to
// buffer the read contents, can be nil, and is provided so that the buffer
// can be reused between messages. The returned segment is the first segment
// read, which contains the root pointer.
func ReadFromStream(r io.Reader, buf *bytes.Buffer) (*Segment, error) {
	if buf == nil {
		buf = new(bytes.Buffer)
	} else {
		buf.Reset()
	}

	if _, err := io.CopyN(buf, r, 4); err != nil {
		return nil, err
	}

	if little32(buf.Bytes()[:]) >= uint32(MaxSegmentNumber) {
		return nil, ErrTooMuchData
	}

	segnum := int(little32(buf.Bytes()[:]) + 1)
	hdrsz := 8*(segnum/2) + 4

	if _, err := io.CopyN(buf, r, int64(hdrsz)); err != nil {
		return nil, err
	}

	total := 0
	for i := 0; i < segnum; i++ {
		sz := little32(buf.Bytes()[4*i+4:])
		if uint64(total)+uint64(sz)*8 > uint64(MaxTotalSize) {
			return nil, ErrTooMuchData
		}
		total += int(sz) * 8
	}

	if _, err := io.CopyN(buf, r, int64(total)); err != nil {
		return nil, err
	}

	hdrv := buf.Bytes()[4 : hdrsz+4]
	datav := buf.Bytes()[hdrsz+4:]

	if segnum == 1 {
		sz := int(little32(hdrv)) * 8
		return NewBuffer(datav[:sz]), nil
	}

	m := &MultiBuffer{make([]*Segment, segnum)}
	for i := 0; i < segnum; i++ {
		sz := int(little32(hdrv[4*i:])) * 8
		m.segments[i] = &Segment{m, datav[:sz], uint32(i)}
		datav = datav[sz:]
	}

	return m.segments[0], nil
}

// ReadFromMemoryZeroCopy: like ReadFromStream, but reads a non-packed
// serialized stream that already resides in memory in the argument data.
// The returned segment is the first segment read, which contains
// the root pointer. The returned bytesRead says how many bytes were
// consumed from data in making seg. The caller should advance the
// data slice by doing data = data[bytesRead:] between successive calls
// to ReadFromMemoryZeroCopy().
func ReadFromMemoryZeroCopy(data []byte) (seg *Segment, bytesRead int64, err error) {

	if len(data) < 4 {
		return nil, 0, io.EOF
	}

	if little32(data[0:4]) >= uint32(MaxSegmentNumber) {
		return nil, 0, ErrTooMuchData
	}

	segnum := int(little32(data[0:4]) + 1)
	hdrsz := 8*(segnum/2) + 4

	b := data[0:(hdrsz + 4)]

	total := 0
	for i := 0; i < segnum; i++ {
		sz := little32(b[4*i+4:])
		if uint64(total)+uint64(sz)*8 > uint64(MaxTotalSize) {
			return nil, 0, ErrTooMuchData
		}
		total += int(sz) * 8
	}
	if total == 0 {
		return nil, 0, io.EOF
	}

	hdrv := data[4:(hdrsz + 4)]
	datav := data[hdrsz+4:]
	m := &multiBuffer{make([]*Segment, segnum)}
	for i := 0; i < segnum; i++ {
		sz := int(little32(hdrv[4*i:])) * 8
		m.segments[i] = &Segment{m, datav[:sz], uint32(i)}
		datav = datav[sz:]
	}

	return m.segments[0], int64(4 + hdrsz + total), nil
}

// WriteTo writes the message that the segment is part of to the
// provided stream in serialized form.
func (s *Segment) WriteTo(w io.Writer) (int64, error) {
	segnum := uint32(1)
	for {
		if seg, _ := s.Message.Lookup(segnum); seg == nil {
			break
		}
		segnum++
	}

	hdrv := make([]uint8, 8*(segnum/2)+8)
	putLittle32(hdrv, segnum-1)
	for i := uint32(0); i < segnum; i++ {
		seg, _ := s.Message.Lookup(i)
		putLittle32(hdrv[4*i+4:], uint32(len(seg.Data)/8))
	}

	if n, err := w.Write(hdrv); err != nil {
		return int64(n), err
	}
	written := int64(len(hdrv))

	for i := uint32(0); i < segnum; i++ {
		seg, _ := s.Message.Lookup(i)
		if n, err := w.Write(seg.Data); err != nil {
			return written + int64(n), err
		} else {
			written += int64(n)
		}
	}

	return written, nil
}
<|endoftext|>


<commit_before>package isolated

import (
	helpers "code.cloudfoundry.org/cli/integration/helpers"
	"code.cloudfoundry.org/cli/util/configv3"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Config", func() {
	Describe("Version Management", func() {
		var oldTarget string
		var oldVersion int

		BeforeEach(func() {
			config := helpers.GetConfig()
			oldTarget = config.Target()
			oldVersion = config.ConfigFile.ConfigVersion
		})
		AfterEach(func() {
			helpers.SetConfig(func(config *configv3.Config) {
				config.ConfigFile.ConfigVersion = oldVersion
				config.ConfigFile.Target = oldTarget
			})
		})

		It("reset config to default if version mismatch", func() {
			helpers.SetConfig(func(config *configv3.Config) {
				config.ConfigFile.ConfigVersion = configv3.CurrentConfigVersion - 1
				config.ConfigFile.Target = "api.my-target"
			})
			helpers.LoginCF()
			config := helpers.GetConfig()
			Expect(config.ConfigFile.ConfigVersion).To(Equal(configv3.CurrentConfigVersion))
			Expect(config.ConfigFile.Target).To(Equal(""))
		})
	})
})
<commit_msg>The ConfigVersion change saga continues<commit_after>package isolated

import (
	"code.cloudfoundry.org/cli/integration/helpers"
	"code.cloudfoundry.org/cli/util/configv3"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Config", func() {
	Describe("Version Management", func() {
		var oldTarget string
		var oldVersion int
		var oldSkipSSLValidation bool

		BeforeEach(func() {
			config := helpers.GetConfig()
			oldTarget = config.Target()
			oldVersion = config.ConfigFile.ConfigVersion
			oldSkipSSLValidation = config.ConfigFile.SkipSSLValidation

		})

		It("reset config to default if version mismatch", func() {
			helpers.SetConfig(func(config *configv3.Config) {
				config.ConfigFile.ConfigVersion = configv3.CurrentConfigVersion - 1
				config.ConfigFile.Target = "api.my-target"
			})
			helpers.LoginCF()
			config := helpers.GetConfig()
			Expect(config.ConfigFile.ConfigVersion).To(Equal(configv3.CurrentConfigVersion))
			Expect(config.ConfigFile.Target).To(Equal(""))
			helpers.SetConfig(func(config *configv3.Config) {
				config.ConfigFile.ConfigVersion = oldVersion
				config.ConfigFile.Target = oldTarget
				config.ConfigFile.SkipSSLValidation = oldSkipSSLValidation
			})
		})
	})
})
<|endoftext|>


<commit_before>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel

import (
	"errors"
	"sync"

	"github.com/uber/tchannel-go/typed"
	"golang.org/x/net/context"
)

var (
	errDuplicateMex        = errors.New("multiple attempts to use the message id")
	errMexChannelFull      = NewSystemError(ErrCodeBusy, "cannot send frame to message exchange channel")
	errUnexpectedFrameType = errors.New("unexpected frame received")
)

const (
	messageExchangeSetInbound  = "inbound"
	messageExchangeSetOutbound = "outbound"

	// mexChannelBufferSize is the size of the message exchange channel buffer.
	mexChannelBufferSize = 2
)

// A messageExchange tracks this Connections's side of a message exchange with a
// peer.  Each message exchange has a channel that can be used to receive
// frames from the peer, and a Context that can controls when the exchange has
// timed out or been cancelled.
type messageExchange struct {
	recvCh    chan *Frame
	ctx       context.Context
	msgID     uint32
	msgType   messageType
	mexset    *messageExchangeSet
	framePool FramePool
}

// forwardPeerFrame forwards a frame from a peer to the message exchange, where
// it can be pulled by whatever application thread is handling the exchange
func (mex *messageExchange) forwardPeerFrame(frame *Frame) error {
	if err := mex.ctx.Err(); err != nil {
		return GetContextError(err)
	}
	select {
	case mex.recvCh <- frame:
		return nil
	case <-mex.ctx.Done():
		// Note: One slow reader processing a large request could stall the connection.
		// If we see this, we need to increase the recvCh buffer size.
		return GetContextError(mex.ctx.Err())
	}
}

// recvPeerFrame waits for a new frame from the peer, or until the context
// expires or is cancelled
func (mex *messageExchange) recvPeerFrame() (*Frame, error) {
	if err := mex.ctx.Err(); err != nil {
		return nil, GetContextError(err)
	}

	select {
	case frame := <-mex.recvCh:
		if frame.Header.ID != mex.msgID {
			mex.mexset.log.WithFields(
				LogField{"msgId", mex.msgID},
				LogField{"header", frame.Header},
			).Error("recvPeerFrame received msg with unexpected ID.")
			return nil, errUnexpectedFrameType
		}
		return frame, nil
	case <-mex.ctx.Done():
		return nil, GetContextError(mex.ctx.Err())
	}
}

// recvPeerFrameOfType waits for a new frame of a given type from the peer, failing
// if the next frame received is not of that type.
// If an error frame is returned, then the errorMessage is returned as the error.
func (mex *messageExchange) recvPeerFrameOfType(msgType messageType) (*Frame, error) {
	frame, err := mex.recvPeerFrame()
	if err != nil {
		return nil, err
	}

	switch frame.Header.messageType {
	case msgType:
		return frame, nil

	case messageTypeError:
		// If we read an error frame, we can release it once we deserialize it.
		defer mex.framePool.Release(frame)

		errMsg := errorMessage{
			id: frame.Header.ID,
		}
		var rbuf typed.ReadBuffer
		rbuf.Wrap(frame.SizedPayload())
		if err := errMsg.read(&rbuf); err != nil {
			return nil, err
		}
		return nil, errMsg

	default:
		// TODO(mmihic): Should be treated as a protocol error
		mex.mexset.log.WithFields(
			LogField{"header", frame.Header},
			LogField{"expectedType", msgType},
			LogField{"expectedID", mex.msgID},
		).Warn("Received unexpected frame.")
		return nil, errUnexpectedFrameType
	}
}

// shutdown shuts down the message exchange, removing it from the message
// exchange set so  that it cannot receive more messages from the peer.  The
// receive channel remains open, however, in case there are concurrent
// goroutines sending to it.
func (mex *messageExchange) shutdown() {
	mex.mexset.removeExchange(mex.msgID)
}

// inboundTimeout is called when an exchange times out, but a handler may still be
// running in the background. Since the handler may still write to the exchange, we
// cannot shutdown the exchange, but we should remove it from the connection's
// exchange list.
func (mex *messageExchange) inboundTimeout() {
	mex.mexset.timeoutExchange(mex.msgID)
}

// A messageExchangeSet manages a set of active message exchanges.  It is
// mainly used to route frames from a peer to the appropriate messageExchange,
// or to cancel or mark a messageExchange as being in error.  Each Connection
// maintains two messageExchangeSets, one to manage exchanges that it has
// initiated (outbound), and another to manage exchanges that the peer has
// initiated (inbound).  The message-type specific handlers are responsible for
// ensuring that their message exchanges are properly registered and removed
// from the corresponding exchange set.
type messageExchangeSet struct {
	sync.RWMutex

	log        Logger
	name       string
	onRemoved  func()
	onAdded    func()
	sendChRefs sync.WaitGroup

	// exchanges is mutable, and is protected by the mutex.
	exchanges map[uint32]*messageExchange
}

// newMessageExchangeSet creates a new messageExchangeSet with a given name.
func newMessageExchangeSet(log Logger, name string) *messageExchangeSet {
	return &messageExchangeSet{
		name:      name,
		log:       log.WithFields(LogField{"exchange", name}),
		exchanges: make(map[uint32]*messageExchange),
	}
}

// newExchange creates and adds a new message exchange to this set
func (mexset *messageExchangeSet) newExchange(ctx context.Context, framePool FramePool,
	msgType messageType, msgID uint32, bufferSize int) (*messageExchange, error) {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("Creating new %s message exchange for [%v:%d]", mexset.name, msgType, msgID)
	}

	mex := &messageExchange{
		msgType:   msgType,
		msgID:     msgID,
		ctx:       ctx,
		recvCh:    make(chan *Frame, bufferSize),
		mexset:    mexset,
		framePool: framePool,
	}

	mexset.Lock()
	if existingMex := mexset.exchanges[mex.msgID]; existingMex != nil {
		if existingMex == mex {
			mexset.log.WithFields(
				LogField{"name", mexset.name},
				LogField{"msgType", mex.msgType},
				LogField{"msgID", mex.msgID},
			).Warn("mex registered multiple times.")
		} else {
			mexset.log.WithFields(
				LogField{"msgID", mex.msgID},
				LogField{"existingType", existingMex.msgType},
				LogField{"newType", mex.msgType},
			).Warn("Duplicate msg ID for active and new mex.")
		}

		mexset.Unlock()
		return nil, errDuplicateMex
	}

	mexset.exchanges[mex.msgID] = mex
	mexset.sendChRefs.Add(1)
	mexset.Unlock()

	mexset.onAdded()

	// TODO(mmihic): Put into a deadline ordered heap so we can garbage collected expired exchanges
	return mex, nil
}

// removeExchange removes a message exchange from the set, if it exists.
// It decrements the sendChRefs wait group, signalling that this exchange no longer has
// any active goroutines that will try to send to sendCh.
func (mexset *messageExchangeSet) removeExchange(msgID uint32) {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("Removing %s message exchange %d", mexset.name, msgID)
	}

	mexset.Lock()
	delete(mexset.exchanges, msgID)
	mexset.Unlock()

	mexset.sendChRefs.Done()
	mexset.onRemoved()
}

// timeoutExchange is similar to removeExchange, however it does not decrement
// the sendChRefs wait group.
func (mexset *messageExchangeSet) timeoutExchange(msgID uint32) {
	mexset.log.Debugf("Removing %s message exchange %d due to timeout", mexset.name, msgID)

	mexset.Lock()
	delete(mexset.exchanges, msgID)
	mexset.Unlock()

	mexset.onRemoved()
}

// waitForSendCh waits for all goroutines with references to sendCh to complete.
func (mexset *messageExchangeSet) waitForSendCh() {
	mexset.sendChRefs.Wait()
}

func (mexset *messageExchangeSet) count() int {
	mexset.RLock()
	count := len(mexset.exchanges)
	mexset.RUnlock()

	return count
}

// forwardPeerFrame forwards a frame from the peer to the appropriate message
// exchange
func (mexset *messageExchangeSet) forwardPeerFrame(frame *Frame) error {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("forwarding %s %s", mexset.name, frame.Header)
	}

	mexset.RLock()
	mex := mexset.exchanges[frame.Header.ID]
	mexset.RUnlock()

	if mex == nil {
		// This is ok since the exchange might have expired or been cancelled
		mexset.log.Infof("received frame %s for %s message exchange that no longer exists",
			frame.Header, mexset.name)
		return nil
	}

	if err := mex.forwardPeerFrame(frame); err != nil {
		mexset.log.Infof("Unable to forward frame %v length %v to %s: %v",
			frame.Header, frame.Header.FrameSize(), mexset.name, err)
		return err
	}

	return nil
}
<commit_msg>Make removeExchange safe to call multiple times<commit_after>// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel

import (
	"errors"
	"sync"

	"github.com/uber/tchannel-go/typed"
	"golang.org/x/net/context"
)

var (
	errDuplicateMex        = errors.New("multiple attempts to use the message id")
	errMexChannelFull      = NewSystemError(ErrCodeBusy, "cannot send frame to message exchange channel")
	errUnexpectedFrameType = errors.New("unexpected frame received")
)

const (
	messageExchangeSetInbound  = "inbound"
	messageExchangeSetOutbound = "outbound"

	// mexChannelBufferSize is the size of the message exchange channel buffer.
	mexChannelBufferSize = 2
)

// A messageExchange tracks this Connections's side of a message exchange with a
// peer.  Each message exchange has a channel that can be used to receive
// frames from the peer, and a Context that can controls when the exchange has
// timed out or been cancelled.
type messageExchange struct {
	recvCh    chan *Frame
	ctx       context.Context
	msgID     uint32
	msgType   messageType
	mexset    *messageExchangeSet
	framePool FramePool
}

// forwardPeerFrame forwards a frame from a peer to the message exchange, where
// it can be pulled by whatever application thread is handling the exchange
func (mex *messageExchange) forwardPeerFrame(frame *Frame) error {
	if err := mex.ctx.Err(); err != nil {
		return GetContextError(err)
	}
	select {
	case mex.recvCh <- frame:
		return nil
	case <-mex.ctx.Done():
		// Note: One slow reader processing a large request could stall the connection.
		// If we see this, we need to increase the recvCh buffer size.
		return GetContextError(mex.ctx.Err())
	}
}

// recvPeerFrame waits for a new frame from the peer, or until the context
// expires or is cancelled
func (mex *messageExchange) recvPeerFrame() (*Frame, error) {
	if err := mex.ctx.Err(); err != nil {
		return nil, GetContextError(err)
	}

	select {
	case frame := <-mex.recvCh:
		if frame.Header.ID != mex.msgID {
			mex.mexset.log.WithFields(
				LogField{"msgId", mex.msgID},
				LogField{"header", frame.Header},
			).Error("recvPeerFrame received msg with unexpected ID.")
			return nil, errUnexpectedFrameType
		}
		return frame, nil
	case <-mex.ctx.Done():
		return nil, GetContextError(mex.ctx.Err())
	}
}

// recvPeerFrameOfType waits for a new frame of a given type from the peer, failing
// if the next frame received is not of that type.
// If an error frame is returned, then the errorMessage is returned as the error.
func (mex *messageExchange) recvPeerFrameOfType(msgType messageType) (*Frame, error) {
	frame, err := mex.recvPeerFrame()
	if err != nil {
		return nil, err
	}

	switch frame.Header.messageType {
	case msgType:
		return frame, nil

	case messageTypeError:
		// If we read an error frame, we can release it once we deserialize it.
		defer mex.framePool.Release(frame)

		errMsg := errorMessage{
			id: frame.Header.ID,
		}
		var rbuf typed.ReadBuffer
		rbuf.Wrap(frame.SizedPayload())
		if err := errMsg.read(&rbuf); err != nil {
			return nil, err
		}
		return nil, errMsg

	default:
		// TODO(mmihic): Should be treated as a protocol error
		mex.mexset.log.WithFields(
			LogField{"header", frame.Header},
			LogField{"expectedType", msgType},
			LogField{"expectedID", mex.msgID},
		).Warn("Received unexpected frame.")
		return nil, errUnexpectedFrameType
	}
}

// shutdown shuts down the message exchange, removing it from the message
// exchange set so  that it cannot receive more messages from the peer.  The
// receive channel remains open, however, in case there are concurrent
// goroutines sending to it.
func (mex *messageExchange) shutdown() {
	mex.mexset.removeExchange(mex.msgID)
}

// inboundTimeout is called when an exchange times out, but a handler may still be
// running in the background. Since the handler may still write to the exchange, we
// cannot shutdown the exchange, but we should remove it from the connection's
// exchange list.
func (mex *messageExchange) inboundTimeout() {
	mex.mexset.timeoutExchange(mex.msgID)
}

// A messageExchangeSet manages a set of active message exchanges.  It is
// mainly used to route frames from a peer to the appropriate messageExchange,
// or to cancel or mark a messageExchange as being in error.  Each Connection
// maintains two messageExchangeSets, one to manage exchanges that it has
// initiated (outbound), and another to manage exchanges that the peer has
// initiated (inbound).  The message-type specific handlers are responsible for
// ensuring that their message exchanges are properly registered and removed
// from the corresponding exchange set.
type messageExchangeSet struct {
	sync.RWMutex

	log        Logger
	name       string
	onRemoved  func()
	onAdded    func()
	sendChRefs sync.WaitGroup

	// exchanges is mutable, and is protected by the mutex.
	exchanges map[uint32]*messageExchange
}

// newMessageExchangeSet creates a new messageExchangeSet with a given name.
func newMessageExchangeSet(log Logger, name string) *messageExchangeSet {
	return &messageExchangeSet{
		name:      name,
		log:       log.WithFields(LogField{"exchange", name}),
		exchanges: make(map[uint32]*messageExchange),
	}
}

// newExchange creates and adds a new message exchange to this set
func (mexset *messageExchangeSet) newExchange(ctx context.Context, framePool FramePool,
	msgType messageType, msgID uint32, bufferSize int) (*messageExchange, error) {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("Creating new %s message exchange for [%v:%d]", mexset.name, msgType, msgID)
	}

	mex := &messageExchange{
		msgType:   msgType,
		msgID:     msgID,
		ctx:       ctx,
		recvCh:    make(chan *Frame, bufferSize),
		mexset:    mexset,
		framePool: framePool,
	}

	mexset.Lock()
	if existingMex := mexset.exchanges[mex.msgID]; existingMex != nil {
		if existingMex == mex {
			mexset.log.WithFields(
				LogField{"name", mexset.name},
				LogField{"msgType", mex.msgType},
				LogField{"msgID", mex.msgID},
			).Warn("mex registered multiple times.")
		} else {
			mexset.log.WithFields(
				LogField{"msgID", mex.msgID},
				LogField{"existingType", existingMex.msgType},
				LogField{"newType", mex.msgType},
			).Warn("Duplicate msg ID for active and new mex.")
		}

		mexset.Unlock()
		return nil, errDuplicateMex
	}

	mexset.exchanges[mex.msgID] = mex
	mexset.sendChRefs.Add(1)
	mexset.Unlock()

	mexset.onAdded()

	// TODO(mmihic): Put into a deadline ordered heap so we can garbage collected expired exchanges
	return mex, nil
}

// removeExchange removes a message exchange from the set, if it exists.
// It decrements the sendChRefs wait group, signalling that this exchange no longer has
// any active goroutines that will try to send to sendCh.
func (mexset *messageExchangeSet) removeExchange(msgID uint32) {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("Removing %s message exchange %d", mexset.name, msgID)
	}

	mexset.Lock()
	_, found := mexset.exchanges[msgID]
	if found {
		delete(mexset.exchanges, msgID)
	}
	mexset.Unlock()

	if !found {
		mexset.log.WithFields(
			LogField{"msgID", msgID},
		).Error("Tried to remove exchange multiple times")
		return
	}

	// If the message exchange was found, then we perform clean up actions.
	// These clean up actions can only be run once per exchange.
	mexset.sendChRefs.Done()
	mexset.onRemoved()
}

// timeoutExchange is similar to removeExchange, however it does not decrement
// the sendChRefs wait group.
func (mexset *messageExchangeSet) timeoutExchange(msgID uint32) {
	mexset.log.Debugf("Removing %s message exchange %d due to timeout", mexset.name, msgID)

	mexset.Lock()
	delete(mexset.exchanges, msgID)
	mexset.Unlock()

	mexset.onRemoved()
}

// waitForSendCh waits for all goroutines with references to sendCh to complete.
func (mexset *messageExchangeSet) waitForSendCh() {
	mexset.sendChRefs.Wait()
}

func (mexset *messageExchangeSet) count() int {
	mexset.RLock()
	count := len(mexset.exchanges)
	mexset.RUnlock()

	return count
}

// forwardPeerFrame forwards a frame from the peer to the appropriate message
// exchange
func (mexset *messageExchangeSet) forwardPeerFrame(frame *Frame) error {
	if mexset.log.Enabled(LogLevelDebug) {
		mexset.log.Debugf("forwarding %s %s", mexset.name, frame.Header)
	}

	mexset.RLock()
	mex := mexset.exchanges[frame.Header.ID]
	mexset.RUnlock()

	if mex == nil {
		// This is ok since the exchange might have expired or been cancelled
		mexset.log.Infof("received frame %s for %s message exchange that no longer exists",
			frame.Header, mexset.name)
		return nil
	}

	if err := mex.forwardPeerFrame(frame); err != nil {
		mexset.log.Infof("Unable to forward frame %v length %v to %s: %v",
			frame.Header, frame.Header.FrameSize(), mexset.name, err)
		return err
	}

	return nil
}
<|endoftext|>