<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package vfs

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path"
	"strings"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/s3"
	"k8s.io/klog"
	"k8s.io/kops/util/pkg/hashing"
)

type S3Path struct {
	s3Context     *S3Context
	bucket        string
	bucketDetails *S3BucketDetails
	key           string
	etag          *string

	// scheme is configurable in case an S3 compatible custom
	// endpoint is specified
	scheme string
	// sse specifies if server side encryption should be enabled
	sse bool
}

var _ Path = &S3Path{}
var _ HasHash = &S3Path{}

// S3Acl is an ACL implementation for objects on S3
type S3Acl struct {
	RequestACL *string
}

func newS3Path(s3Context *S3Context, scheme string, bucket string, key string, sse bool) *S3Path {
	bucket = strings.TrimSuffix(bucket, "/")
	key = strings.TrimPrefix(key, "/")

	return &S3Path{
		s3Context: s3Context,
		bucket:    bucket,
		key:       key,
		scheme:    scheme,
		sse:       sse,
	}
}

func (p *S3Path) Path() string {
	return p.scheme + "://" + p.bucket + "/" + p.key
}

func (p *S3Path) Bucket() string {
	return p.bucket
}

func (p *S3Path) Key() string {
	return p.key
}

func (p *S3Path) String() string {
	return p.Path()
}

func (p *S3Path) Remove() error {
	client, err := p.client()
	if err != nil {
		return err
	}

	klog.V(8).Infof("removing file %s", p)

	request := &s3.ListObjectVersionsInput{
		Bucket: aws.String(p.bucket),
		Prefix: aws.String(p.key),
	}

	response, err := client.ListObjectVersions(request)
	if err != nil {
		return fmt.Errorf("error listing versions %s: %v", p, err)
	}

	// Sometimes S3 will return paginated results if there are too many versions for an object.
	// This is unlikely with current use cases, but a warning should be triggered in case it ever occurs.
	if aws.BoolValue(response.IsTruncated) {
		klog.Warningf("too many versions for %s", p)
	}

	for _, version := range response.Versions {
		klog.V(8).Infof("removing file %s version %q", p, aws.StringValue(version.VersionId))

		request := &s3.DeleteObjectInput{
			Bucket:    aws.String(p.bucket),
			Key:       aws.String(p.key),
			VersionId: version.VersionId,
		}

		_, err = client.DeleteObject(request)
		if err != nil {
			// TODO: Check for not-exists, return os.NotExist

			return fmt.Errorf("error deleting %s version %q: %v", p, aws.StringValue(version.VersionId), err)
		}
	}

	return nil
}

func (p *S3Path) Join(relativePath ...string) Path {
	args := []string{p.key}
	args = append(args, relativePath...)
	joined := path.Join(args...)
	return &S3Path{
		s3Context: p.s3Context,
		bucket:    p.bucket,
		key:       joined,
		scheme:    p.scheme,
		sse:       p.sse,
	}
}

func (p *S3Path) WriteFile(data io.ReadSeeker, aclObj ACL) error {
	client, err := p.client()
	if err != nil {
		return err
	}

	klog.V(4).Infof("Writing file %q", p)

	request := &s3.PutObjectInput{}
	request.Body = data
	request.Bucket = aws.String(p.bucket)
	request.Key = aws.String(p.key)

	// If we are on an S3 implementation that supports SSE (i.e. not
	// DO), we use server-side-encryption, it doesn't really cost us
	// anything.  But if the bucket has a defaultEncryption policy
	// instead, we honor that - it is likely to be a higher encryption
	// standard.
	sseLog := "-"
	if p.sse {
		defaultEncryption := p.bucketDetails.hasServerSideEncryptionByDefault()
		if defaultEncryption {
			sseLog = "DefaultBucketEncryption"
		} else {
			sseLog = "AES256"
			request.ServerSideEncryption = aws.String("AES256")
		}
	}

	acl := os.Getenv("KOPS_STATE_S3_ACL")
	acl = strings.TrimSpace(acl)
	if acl != "" {
		klog.V(8).Infof("Using KOPS_STATE_S3_ACL=%s", acl)
		request.ACL = aws.String(acl)
	} else if aclObj != nil {
		s3Acl, ok := aclObj.(*S3Acl)
		if !ok {
			return fmt.Errorf("write to %s with ACL of unexpected type %T", p, aclObj)
		}
		request.ACL = s3Acl.RequestACL
	}

	// We don't need Content-MD5: https://github.com/aws/aws-sdk-go/issues/208

	klog.V(8).Infof("Calling S3 PutObject Bucket=%q Key=%q SSE=%q ACL=%q", p.bucket, p.key, sseLog, acl)

	_, err = client.PutObject(request)
	if err != nil {
		if acl != "" {
			return fmt.Errorf("error writing %s (with ACL=%q): %v", p, acl, err)
		}
		return fmt.Errorf("error writing %s: %v", p, err)
	}

	return nil
}

// To prevent concurrent creates on the same file while maintaining atomicity of writes,
// we take a process-wide lock during the operation.
// Not a great approach, but fine for a single process (with low concurrency)
// TODO: should we enable versioning?
var createFileLockS3 sync.Mutex

func (p *S3Path) CreateFile(data io.ReadSeeker, acl ACL) error {
	createFileLockS3.Lock()
	defer createFileLockS3.Unlock()

	// Check if exists
	_, err := p.ReadFile()
	if err == nil {
		return os.ErrExist
	}

	if !os.IsNotExist(err) {
		return err
	}

	return p.WriteFile(data, acl)
}

// ReadFile implements Path::ReadFile
func (p *S3Path) ReadFile() ([]byte, error) {
	var b bytes.Buffer
	_, err := p.WriteTo(&b)
	if err != nil {
		return nil, err
	}
	return b.Bytes(), nil
}

// WriteTo implements io.WriterTo
func (p *S3Path) WriteTo(out io.Writer) (int64, error) {
	client, err := p.client()
	if err != nil {
		return 0, err
	}

	klog.V(4).Infof("Reading file %q", p)

	request := &s3.GetObjectInput{}
	request.Bucket = aws.String(p.bucket)
	request.Key = aws.String(p.key)

	response, err := client.GetObject(request)
	if err != nil {
		if AWSErrorCode(err) == "NoSuchKey" {
			return 0, os.ErrNotExist
		}
		return 0, fmt.Errorf("error fetching %s: %v", p, err)
	}
	defer response.Body.Close()

	n, err := io.Copy(out, response.Body)
	if err != nil {
		return n, fmt.Errorf("error reading %s: %v", p, err)
	}
	return n, nil
}

func (p *S3Path) ReadDir() ([]Path, error) {
	client, err := p.client()
	if err != nil {
		return nil, err
	}

	prefix := p.key
	if prefix != "" && !strings.HasSuffix(prefix, "/") {
		prefix += "/"
	}
	request := &s3.ListObjectsInput{}
	request.Bucket = aws.String(p.bucket)
	request.Prefix = aws.String(prefix)
	request.Delimiter = aws.String("/")

	klog.V(4).Infof("Listing objects in S3 bucket %q with prefix %q", p.bucket, prefix)
	var paths []Path
	err = client.ListObjectsPages(request, func(page *s3.ListObjectsOutput, lastPage bool) bool {
		for _, o := range page.Contents {
			key := aws.StringValue(o.Key)
			if key == prefix {
				// We have reports (#548 and #520) of the directory being returned as a file
				// And this will indeed happen if the directory has been created as a file,
				// which seems to happen if you use some external tools to manipulate the S3 bucket.
				// We need to tolerate that, so skip the parent directory.
				klog.V(4).Infof("Skipping read of directory: %q", key)
				continue
			}
			child := &S3Path{
				s3Context: p.s3Context,
				bucket:    p.bucket,
				key:       key,
				etag:      o.ETag,
				scheme:    p.scheme,
				sse:       p.sse,
			}
			paths = append(paths, child)
		}
		return true
	})
	if err != nil {
		return nil, fmt.Errorf("error listing %s: %v", p, err)
	}
	klog.V(8).Infof("Listed files in %v: %v", p, paths)
	return paths, nil
}

func (p *S3Path) ReadTree() ([]Path, error) {
	client, err := p.client()
	if err != nil {
		return nil, err
	}

	request := &s3.ListObjectsInput{}
	request.Bucket = aws.String(p.bucket)
	prefix := p.key
	if prefix != "" && !strings.HasSuffix(prefix, "/") {
		prefix += "/"
	}
	request.Prefix = aws.String(prefix)
	// No delimiter for recursive search

	var paths []Path
	err = client.ListObjectsPages(request, func(page *s3.ListObjectsOutput, lastPage bool) bool {
		for _, o := range page.Contents {
			key := aws.StringValue(o.Key)
			child := &S3Path{
				s3Context: p.s3Context,
				bucket:    p.bucket,
				key:       key,
				etag:      o.ETag,
				scheme:    p.scheme,
				sse:       p.sse,
			}
			paths = append(paths, child)
		}
		return true
	})
	if err != nil {
		return nil, fmt.Errorf("error listing %s: %v", p, err)
	}
	return paths, nil
}

func (p *S3Path) client() (*s3.S3, error) {
	var err error
	if p.bucketDetails == nil || p.bucketDetails.region == "" {
		bucketDetails, err := p.s3Context.getDetailsForBucket(p.bucket)

		p.bucketDetails = bucketDetails
		if err != nil {
			return nil, err
		}
	}

	client, err := p.s3Context.getClient(p.bucketDetails.region)
	if err != nil {
		return nil, err
	}

	return client, nil
}

func (p *S3Path) Base() string {
	return path.Base(p.key)
}

func (p *S3Path) PreferredHash() (*hashing.Hash, error) {
	return p.Hash(hashing.HashAlgorithmMD5)
}

func (p *S3Path) Hash(a hashing.HashAlgorithm) (*hashing.Hash, error) {
	if a != hashing.HashAlgorithmMD5 {
		return nil, nil
	}

	if p.etag == nil {
		return nil, nil
	}

	md5 := strings.Trim(*p.etag, "\"")

	md5Bytes, err := hex.DecodeString(md5)
	if err != nil {
		return nil, fmt.Errorf("Etag was not a valid MD5 sum: %q", *p.etag)
	}

	return &hashing.Hash{Algorithm: hashing.HashAlgorithmMD5, HashValue: md5Bytes}, nil
}

// AWSErrorCode returns the aws error code, if it is an awserr.Error, otherwise ""
func AWSErrorCode(err error) string {
	if awsError, ok := err.(awserr.Error); ok {
		return awsError.Code()
	}
	return ""
}
<commit_msg>Remove TODO that was not addressed for a long time<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package vfs

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path"
	"strings"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/s3"
	"k8s.io/klog"
	"k8s.io/kops/util/pkg/hashing"
)

type S3Path struct {
	s3Context     *S3Context
	bucket        string
	bucketDetails *S3BucketDetails
	key           string
	etag          *string

	// scheme is configurable in case an S3 compatible custom
	// endpoint is specified
	scheme string
	// sse specifies if server side encryption should be enabled
	sse bool
}

var _ Path = &S3Path{}
var _ HasHash = &S3Path{}

// S3Acl is an ACL implementation for objects on S3
type S3Acl struct {
	RequestACL *string
}

func newS3Path(s3Context *S3Context, scheme string, bucket string, key string, sse bool) *S3Path {
	bucket = strings.TrimSuffix(bucket, "/")
	key = strings.TrimPrefix(key, "/")

	return &S3Path{
		s3Context: s3Context,
		bucket:    bucket,
		key:       key,
		scheme:    scheme,
		sse:       sse,
	}
}

func (p *S3Path) Path() string {
	return p.scheme + "://" + p.bucket + "/" + p.key
}

func (p *S3Path) Bucket() string {
	return p.bucket
}

func (p *S3Path) Key() string {
	return p.key
}

func (p *S3Path) String() string {
	return p.Path()
}

func (p *S3Path) Remove() error {
	client, err := p.client()
	if err != nil {
		return err
	}

	klog.V(8).Infof("removing file %s", p)

	request := &s3.ListObjectVersionsInput{
		Bucket: aws.String(p.bucket),
		Prefix: aws.String(p.key),
	}

	response, err := client.ListObjectVersions(request)
	if err != nil {
		return fmt.Errorf("error listing versions %s: %v", p, err)
	}

	// Sometimes S3 will return paginated results if there are too many versions for an object.
	// This is unlikely with current use cases, but a warning should be triggered in case it ever occurs.
	if aws.BoolValue(response.IsTruncated) {
		klog.Warningf("too many versions for %s", p)
	}

	for _, version := range response.Versions {
		klog.V(8).Infof("removing file %s version %q", p, aws.StringValue(version.VersionId))

		request := &s3.DeleteObjectInput{
			Bucket:    aws.String(p.bucket),
			Key:       aws.String(p.key),
			VersionId: version.VersionId,
		}

		_, err = client.DeleteObject(request)
		if err != nil {
			return fmt.Errorf("error deleting %s version %q: %v", p, aws.StringValue(version.VersionId), err)
		}
	}

	return nil
}

func (p *S3Path) Join(relativePath ...string) Path {
	args := []string{p.key}
	args = append(args, relativePath...)
	joined := path.Join(args...)
	return &S3Path{
		s3Context: p.s3Context,
		bucket:    p.bucket,
		key:       joined,
		scheme:    p.scheme,
		sse:       p.sse,
	}
}

func (p *S3Path) WriteFile(data io.ReadSeeker, aclObj ACL) error {
	client, err := p.client()
	if err != nil {
		return err
	}

	klog.V(4).Infof("Writing file %q", p)

	request := &s3.PutObjectInput{}
	request.Body = data
	request.Bucket = aws.String(p.bucket)
	request.Key = aws.String(p.key)

	// If we are on an S3 implementation that supports SSE (i.e. not
	// DO), we use server-side-encryption, it doesn't really cost us
	// anything.  But if the bucket has a defaultEncryption policy
	// instead, we honor that - it is likely to be a higher encryption
	// standard.
	sseLog := "-"
	if p.sse {
		defaultEncryption := p.bucketDetails.hasServerSideEncryptionByDefault()
		if defaultEncryption {
			sseLog = "DefaultBucketEncryption"
		} else {
			sseLog = "AES256"
			request.ServerSideEncryption = aws.String("AES256")
		}
	}

	acl := os.Getenv("KOPS_STATE_S3_ACL")
	acl = strings.TrimSpace(acl)
	if acl != "" {
		klog.V(8).Infof("Using KOPS_STATE_S3_ACL=%s", acl)
		request.ACL = aws.String(acl)
	} else if aclObj != nil {
		s3Acl, ok := aclObj.(*S3Acl)
		if !ok {
			return fmt.Errorf("write to %s with ACL of unexpected type %T", p, aclObj)
		}
		request.ACL = s3Acl.RequestACL
	}

	// We don't need Content-MD5: https://github.com/aws/aws-sdk-go/issues/208

	klog.V(8).Infof("Calling S3 PutObject Bucket=%q Key=%q SSE=%q ACL=%q", p.bucket, p.key, sseLog, acl)

	_, err = client.PutObject(request)
	if err != nil {
		if acl != "" {
			return fmt.Errorf("error writing %s (with ACL=%q): %v", p, acl, err)
		}
		return fmt.Errorf("error writing %s: %v", p, err)
	}

	return nil
}

// To prevent concurrent creates on the same file while maintaining atomicity of writes,
// we take a process-wide lock during the operation.
// Not a great approach, but fine for a single process (with low concurrency)
// TODO: should we enable versioning?
var createFileLockS3 sync.Mutex

func (p *S3Path) CreateFile(data io.ReadSeeker, acl ACL) error {
	createFileLockS3.Lock()
	defer createFileLockS3.Unlock()

	// Check if exists
	_, err := p.ReadFile()
	if err == nil {
		return os.ErrExist
	}

	if !os.IsNotExist(err) {
		return err
	}

	return p.WriteFile(data, acl)
}

// ReadFile implements Path::ReadFile
func (p *S3Path) ReadFile() ([]byte, error) {
	var b bytes.Buffer
	_, err := p.WriteTo(&b)
	if err != nil {
		return nil, err
	}
	return b.Bytes(), nil
}

// WriteTo implements io.WriterTo
func (p *S3Path) WriteTo(out io.Writer) (int64, error) {
	client, err := p.client()
	if err != nil {
		return 0, err
	}

	klog.V(4).Infof("Reading file %q", p)

	request := &s3.GetObjectInput{}
	request.Bucket = aws.String(p.bucket)
	request.Key = aws.String(p.key)

	response, err := client.GetObject(request)
	if err != nil {
		if AWSErrorCode(err) == "NoSuchKey" {
			return 0, os.ErrNotExist
		}
		return 0, fmt.Errorf("error fetching %s: %v", p, err)
	}
	defer response.Body.Close()

	n, err := io.Copy(out, response.Body)
	if err != nil {
		return n, fmt.Errorf("error reading %s: %v", p, err)
	}
	return n, nil
}

func (p *S3Path) ReadDir() ([]Path, error) {
	client, err := p.client()
	if err != nil {
		return nil, err
	}

	prefix := p.key
	if prefix != "" && !strings.HasSuffix(prefix, "/") {
		prefix += "/"
	}
	request := &s3.ListObjectsInput{}
	request.Bucket = aws.String(p.bucket)
	request.Prefix = aws.String(prefix)
	request.Delimiter = aws.String("/")

	klog.V(4).Infof("Listing objects in S3 bucket %q with prefix %q", p.bucket, prefix)
	var paths []Path
	err = client.ListObjectsPages(request, func(page *s3.ListObjectsOutput, lastPage bool) bool {
		for _, o := range page.Contents {
			key := aws.StringValue(o.Key)
			if key == prefix {
				// We have reports (#548 and #520) of the directory being returned as a file
				// And this will indeed happen if the directory has been created as a file,
				// which seems to happen if you use some external tools to manipulate the S3 bucket.
				// We need to tolerate that, so skip the parent directory.
				klog.V(4).Infof("Skipping read of directory: %q", key)
				continue
			}
			child := &S3Path{
				s3Context: p.s3Context,
				bucket:    p.bucket,
				key:       key,
				etag:      o.ETag,
				scheme:    p.scheme,
				sse:       p.sse,
			}
			paths = append(paths, child)
		}
		return true
	})
	if err != nil {
		return nil, fmt.Errorf("error listing %s: %v", p, err)
	}
	klog.V(8).Infof("Listed files in %v: %v", p, paths)
	return paths, nil
}

func (p *S3Path) ReadTree() ([]Path, error) {
	client, err := p.client()
	if err != nil {
		return nil, err
	}

	request := &s3.ListObjectsInput{}
	request.Bucket = aws.String(p.bucket)
	prefix := p.key
	if prefix != "" && !strings.HasSuffix(prefix, "/") {
		prefix += "/"
	}
	request.Prefix = aws.String(prefix)
	// No delimiter for recursive search

	var paths []Path
	err = client.ListObjectsPages(request, func(page *s3.ListObjectsOutput, lastPage bool) bool {
		for _, o := range page.Contents {
			key := aws.StringValue(o.Key)
			child := &S3Path{
				s3Context: p.s3Context,
				bucket:    p.bucket,
				key:       key,
				etag:      o.ETag,
				scheme:    p.scheme,
				sse:       p.sse,
			}
			paths = append(paths, child)
		}
		return true
	})
	if err != nil {
		return nil, fmt.Errorf("error listing %s: %v", p, err)
	}
	return paths, nil
}

func (p *S3Path) client() (*s3.S3, error) {
	var err error
	if p.bucketDetails == nil || p.bucketDetails.region == "" {
		bucketDetails, err := p.s3Context.getDetailsForBucket(p.bucket)

		p.bucketDetails = bucketDetails
		if err != nil {
			return nil, err
		}
	}

	client, err := p.s3Context.getClient(p.bucketDetails.region)
	if err != nil {
		return nil, err
	}

	return client, nil
}

func (p *S3Path) Base() string {
	return path.Base(p.key)
}

func (p *S3Path) PreferredHash() (*hashing.Hash, error) {
	return p.Hash(hashing.HashAlgorithmMD5)
}

func (p *S3Path) Hash(a hashing.HashAlgorithm) (*hashing.Hash, error) {
	if a != hashing.HashAlgorithmMD5 {
		return nil, nil
	}

	if p.etag == nil {
		return nil, nil
	}

	md5 := strings.Trim(*p.etag, "\"")

	md5Bytes, err := hex.DecodeString(md5)
	if err != nil {
		return nil, fmt.Errorf("Etag was not a valid MD5 sum: %q", *p.etag)
	}

	return &hashing.Hash{Algorithm: hashing.HashAlgorithmMD5, HashValue: md5Bytes}, nil
}

// AWSErrorCode returns the aws error code, if it is an awserr.Error, otherwise ""
func AWSErrorCode(err error) string {
	if awsError, ok := err.(awserr.Error); ok {
		return awsError.Code()
	}
	return ""
}
<|endoftext|>


<commit_before>package util

import (
	"reflect"
	"runtime"
	"testing"
	"time"
)

func TestTimeBoundWithCacheFunc(t *testing.T) {
	tests := []struct {
		name               string
		functionReturns    []int
		executionDurations []time.Duration
		expectedResults    []int
		timeout            time.Duration
		maxAge             time.Duration
	}{
		{
			name:               "immediate",
			functionReturns:    []int{42},
			executionDurations: []time.Duration{0},
			expectedResults:    []int{42},
			timeout:            50 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "immediate_zero_timeout",
			functionReturns:    []int{42},
			executionDurations: []time.Duration{0},
			expectedResults:    []int{42},
			timeout:            0 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "immediate_zero_maxage",
			functionReturns:    []int{42},
			executionDurations: []time.Duration{0},
			expectedResults:    []int{42},
			timeout:            50 * time.Millisecond,
			maxAge:             0 * time.Minute,
		},
		{
			name:               "slow_notimeout",
			functionReturns:    []int{42},
			executionDurations: []time.Duration{100 * time.Millisecond},
			expectedResults:    []int{42},
			timeout:            500 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "slow_timeout",
			functionReturns:    []int{42},
			executionDurations: []time.Duration{100 * time.Millisecond},
			expectedResults:    []int{42},
			timeout:            50 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "fast_then_slow_timeout",
			functionReturns:    []int{42, 88},
			executionDurations: []time.Duration{10 * time.Millisecond, 100 * time.Millisecond},
			expectedResults:    []int{42, 42},
			timeout:            50 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "fast_then_fast",
			functionReturns:    []int{42, 88},
			executionDurations: []time.Duration{10 * time.Millisecond, 20 * time.Millisecond},
			expectedResults:    []int{42, 88},
			timeout:            50 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "slow_timeout_then_slow_timeout",
			functionReturns:    []int{42, 88},
			executionDurations: []time.Duration{100 * time.Millisecond, 100 * time.Millisecond},
			expectedResults:    []int{42, 42},
			timeout:            50 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "fast_then_fast_then_slow_timeout",
			functionReturns:    []int{42, 88, 1},
			executionDurations: []time.Duration{10 * time.Millisecond, 1 * time.Millisecond, 100 * time.Millisecond},
			expectedResults:    []int{42, 88, 88},
			timeout:            50 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "fast_then_slow_then_fast_never_timeout",
			functionReturns:    []int{42, 88, 1},
			executionDurations: []time.Duration{1 * time.Millisecond, 100 * time.Millisecond, 2 * time.Millisecond},
			expectedResults:    []int{42, 88, 1},
			timeout:            500 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "fast_then_slow_timeout_then_fast",
			functionReturns:    []int{42, 88, 1},
			executionDurations: []time.Duration{1 * time.Millisecond, 100 * time.Millisecond, 2 * time.Millisecond},
			expectedResults:    []int{42, 42, 1},
			timeout:            50 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "stale_cache_fast_then_slow_timeout",
			functionReturns:    []int{42, 88},
			executionDurations: []time.Duration{10 * time.Millisecond, 100 * time.Millisecond},
			expectedResults:    []int{42, 88},
			timeout:            50 * time.Millisecond,
			maxAge:             30 * time.Millisecond,
		},
		{
			name:               "stale_cache_slow_timeout_then_slow_timeout",
			functionReturns:    []int{42, 88},
			executionDurations: []time.Duration{100 * time.Millisecond, 100 * time.Millisecond},
			expectedResults:    []int{42, 88},
			timeout:            50 * time.Millisecond,
			maxAge:             30 * time.Millisecond,
		},
		{
			name:               "stale_cache_fast_then_slow_timeout_then_fast_then_slow",
			functionReturns:    []int{42, 88, 1, 33},
			executionDurations: []time.Duration{1 * time.Millisecond, 100 * time.Millisecond, 2 * time.Millisecond, 100 * time.Millisecond},
			expectedResults:    []int{42, 88, 1, 33},
			timeout:            50 * time.Millisecond,
			maxAge:             30 * time.Millisecond,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			numGoRoutineAtRest := runtime.NumGoroutine()
			endTime := time.Now()

			functionReturnChan := make(chan int, 1)
			executionDurationChan := make(chan time.Duration, 1)
			fn := func() interface{} {
				time.Sleep(<-executionDurationChan)
				return <-functionReturnChan
			}
			decoratedFunc := TimeBoundWithCacheFunc(fn, tt.timeout, tt.maxAge)
			var cacheTimestamp time.Time

			for i := range tt.executionDurations {
				// compute end time for all function executions of the test
				finishesAt := time.Now().Add(tt.executionDurations[i])
				if finishesAt.After(endTime) {
					endTime = finishesAt
				}

				functionReturnChan <- tt.functionReturns[i]
				executionDurationChan <- tt.executionDurations[i]
				beforeExecution := time.Now()
				result := decoratedFunc()
				executionDuration := time.Now().Sub(beforeExecution)

				// save last time a result was cached
				if tt.executionDurations[i] <= tt.timeout {
					cacheTimestamp = time.Now()
				}

				// check that function execution took around tt.timeout
				// if the function does not have a cached result or the cache is stale due to maxAge, the timeout is not enforced
				cacheIsStale := time.Now().After(cacheTimestamp.Add(tt.maxAge))
				tolerance := 1 * time.Millisecond
				if !cacheIsStale && executionDuration > tt.timeout+tolerance {
					t.Errorf("iteration %v: decoratedFunc() took too long to execute %v which is greater than timeout (%v)", i, executionDuration, tt.timeout)
				}

				if !reflect.DeepEqual(result, tt.expectedResults[i]) {
					t.Errorf("iteration %v: decoratedFunc() = %v, want %v", i, result, tt.expectedResults[i])
				}
			}

			// detects if any goroutine is leaked after all processing is supposed to be complete
			tolerance := 100 * time.Millisecond
			time.Sleep(time.Until(endTime.Add(tolerance)))
			extraGoRoutines := runtime.NumGoroutine() - numGoRoutineAtRest
			if extraGoRoutines > 0 {
				t.Errorf("too many goroutines left after processing: %d", extraGoRoutines)
			}
		})
	}
}
<commit_msg>Fixed cache stale computation & looser tolerance to account for CircleCI's slowness<commit_after>package util

import (
	"reflect"
	"runtime"
	"testing"
	"time"
)

func TestTimeBoundWithCacheFunc(t *testing.T) {
	tests := []struct {
		name               string
		functionReturns    []int
		executionDurations []time.Duration
		expectedResults    []int
		timeout            time.Duration
		maxAge             time.Duration
	}{
		{
			name:               "immediate",
			functionReturns:    []int{42},
			executionDurations: []time.Duration{0},
			expectedResults:    []int{42},
			timeout:            50 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "immediate_zero_timeout",
			functionReturns:    []int{42},
			executionDurations: []time.Duration{0},
			expectedResults:    []int{42},
			timeout:            0 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "immediate_zero_maxage",
			functionReturns:    []int{42},
			executionDurations: []time.Duration{0},
			expectedResults:    []int{42},
			timeout:            50 * time.Millisecond,
			maxAge:             0 * time.Minute,
		},
		{
			name:               "slow_notimeout",
			functionReturns:    []int{42},
			executionDurations: []time.Duration{100 * time.Millisecond},
			expectedResults:    []int{42},
			timeout:            500 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "slow_timeout",
			functionReturns:    []int{42},
			executionDurations: []time.Duration{100 * time.Millisecond},
			expectedResults:    []int{42},
			timeout:            50 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "fast_then_slow_timeout",
			functionReturns:    []int{42, 88},
			executionDurations: []time.Duration{10 * time.Millisecond, 100 * time.Millisecond},
			expectedResults:    []int{42, 42},
			timeout:            50 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "fast_then_fast",
			functionReturns:    []int{42, 88},
			executionDurations: []time.Duration{10 * time.Millisecond, 20 * time.Millisecond},
			expectedResults:    []int{42, 88},
			timeout:            50 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "slow_timeout_then_slow_timeout",
			functionReturns:    []int{42, 88},
			executionDurations: []time.Duration{100 * time.Millisecond, 100 * time.Millisecond},
			expectedResults:    []int{42, 42},
			timeout:            50 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "fast_then_fast_then_slow_timeout",
			functionReturns:    []int{42, 88, 1},
			executionDurations: []time.Duration{10 * time.Millisecond, 1 * time.Millisecond, 100 * time.Millisecond},
			expectedResults:    []int{42, 88, 88},
			timeout:            50 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "fast_then_slow_then_fast_never_timeout",
			functionReturns:    []int{42, 88, 1},
			executionDurations: []time.Duration{1 * time.Millisecond, 100 * time.Millisecond, 2 * time.Millisecond},
			expectedResults:    []int{42, 88, 1},
			timeout:            500 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "fast_then_slow_timeout_then_fast",
			functionReturns:    []int{42, 88, 1},
			executionDurations: []time.Duration{1 * time.Millisecond, 100 * time.Millisecond, 2 * time.Millisecond},
			expectedResults:    []int{42, 42, 1},
			timeout:            50 * time.Millisecond,
			maxAge:             5 * time.Minute,
		},
		{
			name:               "stale_cache_fast_then_slow_timeout",
			functionReturns:    []int{42, 88},
			executionDurations: []time.Duration{10 * time.Millisecond, 100 * time.Millisecond},
			expectedResults:    []int{42, 88},
			timeout:            50 * time.Millisecond,
			maxAge:             30 * time.Millisecond,
		},
		{
			name:               "stale_cache_slow_timeout_then_slow_timeout",
			functionReturns:    []int{42, 88},
			executionDurations: []time.Duration{100 * time.Millisecond, 100 * time.Millisecond},
			expectedResults:    []int{42, 88},
			timeout:            50 * time.Millisecond,
			maxAge:             30 * time.Millisecond,
		},
		{
			name:               "stale_cache_fast_then_slow_timeout_then_fast_then_slow",
			functionReturns:    []int{42, 88, 1, 33},
			executionDurations: []time.Duration{1 * time.Millisecond, 100 * time.Millisecond, 2 * time.Millisecond, 100 * time.Millisecond},
			expectedResults:    []int{42, 88, 1, 33},
			timeout:            50 * time.Millisecond,
			maxAge:             30 * time.Millisecond,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			numGoRoutineAtRest := runtime.NumGoroutine()
			endTime := time.Now()

			functionReturnChan := make(chan int, 1)
			executionDurationChan := make(chan time.Duration, 1)
			fn := func() interface{} {
				time.Sleep(<-executionDurationChan)
				return <-functionReturnChan
			}
			decoratedFunc := TimeBoundWithCacheFunc(fn, tt.timeout, tt.maxAge)
			var cacheTimestamp time.Time

			for i := range tt.executionDurations {
				// compute end time for all function executions of the test
				finishesAt := time.Now().Add(tt.executionDurations[i])
				if finishesAt.After(endTime) {
					endTime = finishesAt
				}

				functionReturnChan <- tt.functionReturns[i]
				executionDurationChan <- tt.executionDurations[i]
				beforeExecution := time.Now()
				result := decoratedFunc()
				executionDuration := time.Now().Sub(beforeExecution)

				// check that function execution took around tt.timeout
				// if the function does not have a cached result or the cache is stale due to maxAge, the timeout is not enforced
				cacheIsStale := time.Now().After(cacheTimestamp.Add(tt.maxAge))
				tolerance := 20 * time.Millisecond
				if !cacheIsStale && executionDuration > tt.timeout+tolerance {
					t.Errorf("iteration %v: decoratedFunc() took too long to execute %v which is greater than timeout (%v)", i, executionDuration, tt.timeout)
				}

				if !reflect.DeepEqual(result, tt.expectedResults[i]) {
					t.Errorf("iteration %v: decoratedFunc() = %v, want %v", i, result, tt.expectedResults[i])
				}

				// save last time a result was cached
				if tt.executionDurations[i] <= tt.timeout {
					cacheTimestamp = time.Now()
				}
			}

			// detects if any goroutine is leaked after all processing is supposed to be complete
			tolerance := 100 * time.Millisecond
			time.Sleep(time.Until(endTime.Add(tolerance)))
			extraGoRoutines := runtime.NumGoroutine() - numGoRoutineAtRest
			if extraGoRoutines > 0 {
				t.Errorf("too many goroutines left after processing: %d", extraGoRoutines)
			}
		})
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 The Serviced 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 utils

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"syscall"

	"github.com/zenoss/glog"
)

// An error type for failed docker exec attempts.
type DockerExecError struct {
	Command string
	ExecErr error
}

func (err DockerExecError) Error() string {
	return fmt.Sprintf("Error running command: %s: %s", err.command, err.err)
}

// ExecDockerExec execs the command using docker exec
func ExecDockerExec(containerID string, bashcmd []string) error {
	command, err := generateDockerExecCommand(containerID, bashcmd, false)
	if err != nil {
		return err
	}
	glog.V(1).Infof("exec command for container:%v command: %v\n", containerID, command)
	return syscall.Exec(command[0], command[0:], os.Environ())
}

// RunDockerExec runs the command using docker exec
func RunDockerExec(containerID string, bashcmd []string) ([]byte, error) {
	oldStdin := os.Stdin
	os.Stdin = nil // temporary stdin=nil https://github.com/docker/docker/pull/9537
	command, err := generateDockerExecCommand(containerID, bashcmd, true)
	os.Stdin = oldStdin
	if err != nil {
		return nil, err
	}
	thecmd := exec.Command(command[0], command[1:]...)
	output, err := thecmd.CombinedOutput()
	if err != nil {
		err = DockerExecError{strings.Join(command, " "), err}
		return output, err
	}
	glog.V(1).Infof("Successfully ran command:'%s' output: %s\n", command, output)
	return output, nil
}

// generateDockerExecCommand returns a slice containing docker exec command to exec
func generateDockerExecCommand(containerID string, bashcmd []string, prependBash bool) ([]string, error) {
	if containerID == "" {
		return []string{}, fmt.Errorf("will not attach to container with empty containerID")
	}

	exeMap, err := exePaths([]string{"docker"})
	if err != nil {
		return []string{}, err
	}

	// TODO: add '-h' hostname to specify the container hostname when that
	// feature becomes available
	attachCmd := []string{exeMap["docker"], "exec"}

	if Isatty(os.Stdin) {
		attachCmd = append(attachCmd, "-t")
	}
	if Isatty(os.Stdout) && Isatty(os.Stdin) {
		attachCmd = append(attachCmd, "-i")
	}
	attachCmd = append(attachCmd, containerID)

	if prependBash {
		attachCmd = append(attachCmd, "/bin/bash", "-c", fmt.Sprintf("%s", strings.Join(bashcmd, " ")))
	} else {
		attachCmd = append(attachCmd, bashcmd...)
	}
	glog.V(1).Infof("attach command for container:%v command: %v\n", containerID, attachCmd)
	return attachCmd, nil
}

// hasFeatureDockerExec returns true if docker exec is supported
func hasFeatureDockerExec() bool {
	command := []string{"docker", "exec"}

	thecmd := exec.Command(command[0], command[1:]...)
	output, err := thecmd.CombinedOutput()
	// when docker exec is supported, we expect above 'docker exec' to fail,
	// but provide usage
	glog.V(1).Infof("Successfully ran command:'%s'  err: %s  output: %s\n", command, err, output)

	return strings.Contains(string(output), "Usage: docker exec")
}

// AttachAndRun attaches to a container and runs the command
func AttachAndRun(containerID string, bashcmd []string) ([]byte, error) {
	return RunDockerExec(containerID, bashcmd)
}

// AttachAndExec attaches to a container and execs the command
func AttachAndExec(containerID string, bashcmd []string) error {
	return ExecDockerExec(containerID, bashcmd)
}

// exePaths returns the full path to the given executables in a map
func exePaths(exes []string) (map[string]string, error) {
	exeMap := map[string]string{}

	for _, exe := range exes {
		path, err := exec.LookPath(exe)
		if err != nil {
			glog.Errorf("exe:'%v' not found error:%v\n", exe, err)
			return nil, err
		}

		exeMap[exe] = path
	}

	return exeMap, nil
}
<commit_msg>Make DockerExecError members public.<commit_after>// Copyright 2014 The Serviced 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 utils

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"syscall"

	"github.com/zenoss/glog"
)

// An error type for failed docker exec attempts.
type DockerExecError struct {
	Command string
	ExecErr error
}

func (err DockerExecError) Error() string {
	return fmt.Sprintf("Error running command: %s: %s", err.Command, err.ExecErr)
}

// ExecDockerExec execs the command using docker exec
func ExecDockerExec(containerID string, bashcmd []string) error {
	command, err := generateDockerExecCommand(containerID, bashcmd, false)
	if err != nil {
		return err
	}
	glog.V(1).Infof("exec command for container:%v command: %v\n", containerID, command)
	return syscall.Exec(command[0], command[0:], os.Environ())
}

// RunDockerExec runs the command using docker exec
func RunDockerExec(containerID string, bashcmd []string) ([]byte, error) {
	oldStdin := os.Stdin
	os.Stdin = nil // temporary stdin=nil https://github.com/docker/docker/pull/9537
	command, err := generateDockerExecCommand(containerID, bashcmd, true)
	os.Stdin = oldStdin
	if err != nil {
		return nil, err
	}
	thecmd := exec.Command(command[0], command[1:]...)
	output, err := thecmd.CombinedOutput()
	if err != nil {
		err = DockerExecError{strings.Join(command, " "), err}
		return output, err
	}
	glog.V(1).Infof("Successfully ran command:'%s' output: %s\n", command, output)
	return output, nil
}

// generateDockerExecCommand returns a slice containing docker exec command to exec
func generateDockerExecCommand(containerID string, bashcmd []string, prependBash bool) ([]string, error) {
	if containerID == "" {
		return []string{}, fmt.Errorf("will not attach to container with empty containerID")
	}

	exeMap, err := exePaths([]string{"docker"})
	if err != nil {
		return []string{}, err
	}

	// TODO: add '-h' hostname to specify the container hostname when that
	// feature becomes available
	attachCmd := []string{exeMap["docker"], "exec"}

	if Isatty(os.Stdin) {
		attachCmd = append(attachCmd, "-t")
	}
	if Isatty(os.Stdout) && Isatty(os.Stdin) {
		attachCmd = append(attachCmd, "-i")
	}
	attachCmd = append(attachCmd, containerID)

	if prependBash {
		attachCmd = append(attachCmd, "/bin/bash", "-c", fmt.Sprintf("%s", strings.Join(bashcmd, " ")))
	} else {
		attachCmd = append(attachCmd, bashcmd...)
	}
	glog.V(1).Infof("attach command for container:%v command: %v\n", containerID, attachCmd)
	return attachCmd, nil
}

// hasFeatureDockerExec returns true if docker exec is supported
func hasFeatureDockerExec() bool {
	command := []string{"docker", "exec"}

	thecmd := exec.Command(command[0], command[1:]...)
	output, err := thecmd.CombinedOutput()
	// when docker exec is supported, we expect above 'docker exec' to fail,
	// but provide usage
	glog.V(1).Infof("Successfully ran command:'%s'  err: %s  output: %s\n", command, err, output)

	return strings.Contains(string(output), "Usage: docker exec")
}

// AttachAndRun attaches to a container and runs the command
func AttachAndRun(containerID string, bashcmd []string) ([]byte, error) {
	return RunDockerExec(containerID, bashcmd)
}

// AttachAndExec attaches to a container and execs the command
func AttachAndExec(containerID string, bashcmd []string) error {
	return ExecDockerExec(containerID, bashcmd)
}

// exePaths returns the full path to the given executables in a map
func exePaths(exes []string) (map[string]string, error) {
	exeMap := map[string]string{}

	for _, exe := range exes {
		path, err := exec.LookPath(exe)
		if err != nil {
			glog.Errorf("exe:'%v' not found error:%v\n", exe, err)
			return nil, err
		}

		exeMap[exe] = path
	}

	return exeMap, nil
}
<|endoftext|>


<commit_before>package client

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"time"

	"github.com/jcmturner/gokrb5/v8/iana/errorcode"
	"github.com/jcmturner/gokrb5/v8/messages"
)

// SendToKDC performs network actions to send data to the KDC.
func (cl *Client) sendToKDC(b []byte, realm string) ([]byte, error) {
	var rb []byte
	if cl.Config.LibDefaults.UDPPreferenceLimit == 1 {
		//1 means we should always use TCP
		rb, errtcp := cl.sendKDCTCP(realm, b)
		if errtcp != nil {
			if e, ok := errtcp.(messages.KRBError); ok {
				return rb, e
			}
			return rb, fmt.Errorf("communication error with KDC via TCP: %v", errtcp)
		}
		return rb, nil
	}
	if len(b) <= cl.Config.LibDefaults.UDPPreferenceLimit {
		//Try UDP first, TCP second
		rb, errudp := cl.sendKDCUDP(realm, b)
		if errudp != nil {
			if e, ok := errudp.(messages.KRBError); ok && e.ErrorCode != errorcode.KRB_ERR_RESPONSE_TOO_BIG {
				// Got a KRBError from KDC
				// If this is not a KRB_ERR_RESPONSE_TOO_BIG we will return immediately otherwise will try TCP.
				return rb, e
			}
			// Try TCP
			r, errtcp := cl.sendKDCTCP(realm, b)
			if errtcp != nil {
				if e, ok := errtcp.(messages.KRBError); ok {
					// Got a KRBError
					return r, e
				}
				return r, fmt.Errorf("failed to communicate with KDC. Attempts made with UDP (%v) and then TCP (%v)", errudp, errtcp)
			}
			rb = r
		}
		return rb, nil
	}
	//Try TCP first, UDP second
	rb, errtcp := cl.sendKDCTCP(realm, b)
	if errtcp != nil {
		if e, ok := errtcp.(messages.KRBError); ok {
			// Got a KRBError from KDC so returning and not trying UDP.
			return rb, e
		}
		rb, errudp := cl.sendKDCUDP(realm, b)
		if errudp != nil {
			if e, ok := errudp.(messages.KRBError); ok {
				// Got a KRBError
				return rb, e
			}
			return rb, fmt.Errorf("failed to communicate with KDC. Attempts made with TCP (%v) and then UDP (%v)", errtcp, errudp)
		}
	}
	return rb, nil
}

// sendKDCUDP sends bytes to the KDC via UDP.
func (cl *Client) sendKDCUDP(realm string, b []byte) ([]byte, error) {
	var r []byte
	_, kdcs, err := cl.Config.GetKDCs(realm, false)
	if err != nil {
		return r, err
	}
	r, err = dialSendUDP(kdcs, b)
	if err != nil {
		return r, err
	}
	return checkForKRBError(r)
}

// dialSendUDP establishes a UDP connection to a KDC.
func dialSendUDP(kdcs map[int]string, b []byte) ([]byte, error) {
	var errs []string
	for i := 1; i <= len(kdcs); i++ {
		udpAddr, err := net.ResolveUDPAddr("udp", kdcs[i])
		if err != nil {
			errs = append(errs, fmt.Sprintf("error resolving KDC address: %v", err))
			continue
		}

		conn, err := net.DialTimeout("udp", udpAddr.String(), 5*time.Second)
		if err != nil {
			errs = append(errs, fmt.Sprintf("error setting dial timeout on connection to %s: %v", kdcs[i], err))
			continue
		}
		if err := conn.SetDeadline(time.Now().Add(5 * time.Second)); err != nil {
			errs = append(errs, fmt.Sprintf("error setting deadline on connection to %s: %v", kdcs[i], err))
			continue
		}
		// conn is guaranteed to be a UDPConn
		rb, err := sendUDP(conn.(*net.UDPConn), b)
		if err != nil {
			errs = append(errs, fmt.Sprintf("error sneding to %s: %v", kdcs[i], err))
			continue
		}
		return rb, nil
	}
	return nil, fmt.Errorf("error sending to a KDC: %s", strings.Join(errs, "; "))
}

// sendUDP sends bytes to connection over UDP.
func sendUDP(conn *net.UDPConn, b []byte) ([]byte, error) {
	var r []byte
	defer conn.Close()
	_, err := conn.Write(b)
	if err != nil {
		return r, fmt.Errorf("error sending to (%s): %v", conn.RemoteAddr().String(), err)
	}
	udpbuf := make([]byte, 4096)
	n, _, err := conn.ReadFrom(udpbuf)
	r = udpbuf[:n]
	if err != nil {
		return r, fmt.Errorf("sending over UDP failed to %s: %v", conn.RemoteAddr().String(), err)
	}
	if len(r) < 1 {
		return r, fmt.Errorf("no response data from %s", conn.RemoteAddr().String())
	}
	return r, nil
}

// sendKDCTCP sends bytes to the KDC via TCP.
func (cl *Client) sendKDCTCP(realm string, b []byte) ([]byte, error) {
	var r []byte
	_, kdcs, err := cl.Config.GetKDCs(realm, true)
	if err != nil {
		return r, err
	}
	r, err = dialSendTCP(kdcs, b)
	if err != nil {
		return r, err
	}
	return checkForKRBError(r)
}

// dialKDCTCP establishes a TCP connection to a KDC.
func dialSendTCP(kdcs map[int]string, b []byte) ([]byte, error) {
	var errs []string
	for i := 1; i <= len(kdcs); i++ {
		tcpAddr, err := net.ResolveTCPAddr("tcp", kdcs[i])
		if err != nil {
			errs = append(errs, fmt.Sprintf("error resolving KDC address: %v", err))
			continue
		}

		conn, err := net.DialTimeout("tcp", tcpAddr.String(), 5*time.Second)
		if err != nil {
			errs = append(errs, fmt.Sprintf("error setting dial timeout on connection to %s: %v", kdcs[i], err))
			continue
		}
		if err := conn.SetDeadline(time.Now().Add(5 * time.Second)); err != nil {
			errs = append(errs, fmt.Sprintf("error setting deadline on connection to %s: %v", kdcs[i], err))
			continue
		}
		// conn is guaranteed to be a TCPConn
		rb, err := sendTCP(conn.(*net.TCPConn), b)
		if err != nil {
			errs = append(errs, fmt.Sprintf("error sneding to %s: %v", kdcs[i], err))
			continue
		}
		return rb, nil
	}
	return nil, fmt.Errorf("error sending to a KDC: %s", strings.Join(errs, "; "))
}

// sendTCP sends bytes to connection over TCP.
func sendTCP(conn *net.TCPConn, b []byte) ([]byte, error) {
	defer conn.Close()
	var r []byte
	// RFC 4120 7.2.2 specifies the first 4 bytes indicate the length of the message in big endian order.
	hb := make([]byte, 4, 4)
	binary.BigEndian.PutUint32(hb, uint32(len(b)))
	b = append(hb, b...)

	_, err := conn.Write(b)
	if err != nil {
		return r, fmt.Errorf("error sending to KDC (%s): %v", conn.RemoteAddr().String(), err)
	}

	sh := make([]byte, 4, 4)
	_, err = conn.Read(sh)
	if err != nil {
		return r, fmt.Errorf("error reading response size header: %v", err)
	}
	s := binary.BigEndian.Uint32(sh)

	rb := make([]byte, s, s)
	_, err = io.ReadFull(conn, rb)
	if err != nil {
		return r, fmt.Errorf("error reading response: %v", err)
	}
	if len(rb) < 1 {
		return r, fmt.Errorf("no response data from KDC %s", conn.RemoteAddr().String())
	}
	return rb, nil
}

// checkForKRBError checks if the response bytes from the KDC are a KRBError.
func checkForKRBError(b []byte) ([]byte, error) {
	var KRBErr messages.KRBError
	if err := KRBErr.Unmarshal(b); err == nil {
		return b, KRBErr
	}
	return b, nil
}
<commit_msg>Update network.go<commit_after>package client

import (
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"strings"
	"time"

	"github.com/jcmturner/gokrb5/v8/iana/errorcode"
	"github.com/jcmturner/gokrb5/v8/messages"
)

// SendToKDC performs network actions to send data to the KDC.
func (cl *Client) sendToKDC(b []byte, realm string) ([]byte, error) {
	var rb []byte
	if cl.Config.LibDefaults.UDPPreferenceLimit == 1 {
		//1 means we should always use TCP
		rb, errtcp := cl.sendKDCTCP(realm, b)
		if errtcp != nil {
			if e, ok := errtcp.(messages.KRBError); ok {
				return rb, e
			}
			return rb, fmt.Errorf("communication error with KDC via TCP: %v", errtcp)
		}
		return rb, nil
	}
	if len(b) <= cl.Config.LibDefaults.UDPPreferenceLimit {
		//Try UDP first, TCP second
		rb, errudp := cl.sendKDCUDP(realm, b)
		if errudp != nil {
			if e, ok := errudp.(messages.KRBError); ok && e.ErrorCode != errorcode.KRB_ERR_RESPONSE_TOO_BIG {
				// Got a KRBError from KDC
				// If this is not a KRB_ERR_RESPONSE_TOO_BIG we will return immediately otherwise will try TCP.
				return rb, e
			}
			// Try TCP
			r, errtcp := cl.sendKDCTCP(realm, b)
			if errtcp != nil {
				if e, ok := errtcp.(messages.KRBError); ok {
					// Got a KRBError
					return r, e
				}
				return r, fmt.Errorf("failed to communicate with KDC. Attempts made with UDP (%v) and then TCP (%v)", errudp, errtcp)
			}
			rb = r
		}
		return rb, nil
	}
	//Try TCP first, UDP second
	rb, errtcp := cl.sendKDCTCP(realm, b)
	if errtcp != nil {
		if e, ok := errtcp.(messages.KRBError); ok {
			// Got a KRBError from KDC so returning and not trying UDP.
			return rb, e
		}
		rb, errudp := cl.sendKDCUDP(realm, b)
		if errudp != nil {
			if e, ok := errudp.(messages.KRBError); ok {
				// Got a KRBError
				return rb, e
			}
			return rb, fmt.Errorf("failed to communicate with KDC. Attempts made with TCP (%v) and then UDP (%v)", errtcp, errudp)
		}
	}
	return rb, nil
}

// sendKDCUDP sends bytes to the KDC via UDP.
func (cl *Client) sendKDCUDP(realm string, b []byte) ([]byte, error) {
	var r []byte
	_, kdcs, err := cl.Config.GetKDCs(realm, false)
	if err != nil {
		return r, err
	}
	r, err = dialSendUDP(kdcs, b)
	if err != nil {
		return r, err
	}
	return checkForKRBError(r)
}

// dialSendUDP establishes a UDP connection to a KDC.
func dialSendUDP(kdcs map[int]string, b []byte) ([]byte, error) {
	var errs []string
	for i := 1; i <= len(kdcs); i++ {
		udpAddr, err := net.ResolveUDPAddr("udp", kdcs[i])
		if err != nil {
			errs = append(errs, fmt.Sprintf("error resolving KDC address: %v", err))
			continue
		}

		conn, err := net.DialTimeout("udp", udpAddr.String(), 5*time.Second)
		if err != nil {
			errs = append(errs, fmt.Sprintf("error setting dial timeout on connection to %s: %v", kdcs[i], err))
			continue
		}
		if err := conn.SetDeadline(time.Now().Add(5 * time.Second)); err != nil {
			errs = append(errs, fmt.Sprintf("error setting deadline on connection to %s: %v", kdcs[i], err))
			continue
		}
		// conn is guaranteed to be a UDPConn
		rb, err := sendUDP(conn.(*net.UDPConn), b)
		if err != nil {
			errs = append(errs, fmt.Sprintf("error sneding to %s: %v", kdcs[i], err))
			continue
		}
		return rb, nil
	}
	return nil, fmt.Errorf("error sending to a KDC: %s", strings.Join(errs, "; "))
}

// sendUDP sends bytes to connection over UDP.
func sendUDP(conn *net.UDPConn, b []byte) ([]byte, error) {
	var r []byte
	defer conn.Close()
	_, err := conn.Write(b)
	if err != nil {
		return r, fmt.Errorf("error sending to (%s): %v", conn.RemoteAddr().String(), err)
	}
	udpbuf := make([]byte, 4096)
	n, _, err := conn.ReadFrom(udpbuf)
	r = udpbuf[:n]
	if err != nil {
		return r, fmt.Errorf("sending over UDP failed to %s: %v", conn.RemoteAddr().String(), err)
	}
	if len(r) < 1 {
		return r, fmt.Errorf("no response data from %s", conn.RemoteAddr().String())
	}
	return r, nil
}

// sendKDCTCP sends bytes to the KDC via TCP.
func (cl *Client) sendKDCTCP(realm string, b []byte) ([]byte, error) {
	var r []byte
	_, kdcs, err := cl.Config.GetKDCs(realm, true)
	if err != nil {
		return r, err
	}
	r, err = dialSendTCP(kdcs, b)
	if err != nil {
		return r, err
	}
	return checkForKRBError(r)
}

// dialKDCTCP establishes a TCP connection to a KDC.
func dialSendTCP(kdcs map[int]string, b []byte) ([]byte, error) {
	var errs []string
	for i := 1; i <= len(kdcs); i++ {
		tcpAddr, err := net.ResolveTCPAddr("tcp", kdcs[i])
		if err != nil {
			errs = append(errs, fmt.Sprintf("error resolving KDC address: %v", err))
			continue
		}

		conn, err := net.DialTimeout("tcp", tcpAddr.String(), 5*time.Second)
		if err != nil {
			errs = append(errs, fmt.Sprintf("error setting dial timeout on connection to %s: %v", kdcs[i], err))
			continue
		}
		if err := conn.SetDeadline(time.Now().Add(5 * time.Second)); err != nil {
			errs = append(errs, fmt.Sprintf("error setting deadline on connection to %s: %v", kdcs[i], err))
			continue
		}
		// conn is guaranteed to be a TCPConn
		rb, err := sendTCP(conn.(*net.TCPConn), b)
		if err != nil {
			errs = append(errs, fmt.Sprintf("error sneding to %s: %v", kdcs[i], err))
			continue
		}
		return rb, nil
	}
	return nil, fmt.Errorf("error sending to a KDC: %s", strings.Join(errs, "; "))
}

// sendTCP sends bytes to connection over TCP.
func sendTCP(conn *net.TCPConn, b []byte) ([]byte, error) {
	defer conn.Close()
	var r []byte
	// RFC 4120 7.2.2 specifies the first 4 bytes indicate the length of the message in big endian order.
	hb := make([]byte, 4, 4)
	binary.BigEndian.PutUint32(hb, uint32(len(b)))
	b = append(hb, b...)

	_, err := conn.Write(b)
	if err != nil {
		return r, fmt.Errorf("error sending to KDC (%s): %v", conn.RemoteAddr().String(), err)
	}

	sh := make([]byte, 4, 4)
	_, err = conn.Read(sh)
	if err != nil {
		return r, fmt.Errorf("error reading response size header: %v", err)
	}
	s := binary.BigEndian.Uint32(sh)

	rb := make([]byte, s, s)
	_, err = io.ReadFull(conn, rb)
	if err != nil {
		return r, fmt.Errorf("error reading response: %v", err)
	}
	if len(rb) < 1 {
		return r, fmt.Errorf("no response data from KDC %s", conn.RemoteAddr().String())
	}
	return rb, nil
}

// checkForKRBError checks if the response bytes from the KDC are a KRBError.
func checkForKRBError(b []byte) ([]byte, error) {
	var KRBErr messages.KRBError
	if err := KRBErr.Unmarshal(b); err == nil {
		return b, KRBErr
	}
	return b, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Knative Authors

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

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

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

package metrics

import (
	"errors"
	"fmt"
	"sort"
	"strings"
	"sync"
	"time"

	"go.opencensus.io/metric/metricdata"
	"go.opencensus.io/resource"
	"go.opencensus.io/stats"
	"go.opencensus.io/stats/view"
	"go.opencensus.io/tag"
)

type storedViews struct {
	views []*view.View
	lock  sync.Mutex
}

type meterExporter struct {
	m view.Meter    // NOTE: DO NOT RETURN THIS DIRECTLY; the view.Meter will not work for the empty Resource
	o stats.Options // Cache the option to reduce allocations
	e view.Exporter
}

// ResourceExporterFactory provides a hook for producing separate view.Exporters
// for each observed Resource. This is needed because OpenCensus support for
// Resources is a bit tacked-on rather than being a first-class component like
// Tags are.
type ResourceExporterFactory func(*resource.Resource) (view.Exporter, error)
type meters struct {
	meters  map[string]*meterExporter
	factory ResourceExporterFactory
	lock    sync.Mutex
}

// Lock regime: lock allMeters before resourceViews. The critical path is in
// optionForResource, which must lock allMeters, but only needs to lock
// resourceViews if a new meter needs to be created.
var resourceViews = storedViews{}
var allMeters = meters{
	meters: map[string]*meterExporter{"": &defaultMeter},
}

// RegisterResourceView is similar to view.Register(), except that it will
// register the view across all Resources tracked by the system, rather than
// simply the default view.
func RegisterResourceView(views ...*view.View) error {
	var err error
	allMeters.lock.Lock()
	defer allMeters.lock.Unlock()
	resourceViews.lock.Lock()
	defer resourceViews.lock.Unlock()
	for _, meter := range allMeters.meters {
		// make a copy of views to avoid data races
		viewCopy := make([]*view.View, 0, len(views))
		for _, v := range views {
			c := *v
			viewCopy = append(viewCopy, &c)
		}
		if e := meter.m.Register(viewCopy...); e != nil {
			err = e
		}
	}
	if err != nil {
		return err
	}
	resourceViews.views = append(resourceViews.views, views...)
	return nil
}

// UnregisterResourceView is similar to view.Unregiste(), except that it will
// unregister the view across all Resources tracked byt he system, rather than
// simply the default view.
func UnregisterResourceView(views ...*view.View) {
	allMeters.lock.Lock()
	defer allMeters.lock.Unlock()

	resourceViews.lock.Lock()
	defer resourceViews.lock.Unlock()

	for _, meter := range allMeters.meters {
		// Since we make a defensive copy of all views in RegisterResourceView,
		// the caller might not have the same view pointer that was registered.
		// Use Meter.Find() to find the view with a matching name.
		for _, v := range views {
			name := v.Name
			if v.Name == "" {
				name = v.Measure.Name()
			}
			if v := meter.m.Find(name); v != nil {
				meter.m.Unregister(v)
			}
		}
	}

	j := 0
	for _, view := range resourceViews.views {
		toRemove := false
		for _, viewToRemove := range views {
			if view == viewToRemove {
				toRemove = true
				break
			}
		}
		if !toRemove {
			resourceViews.views[j] = view
			j++
		}
	}
	resourceViews.views = resourceViews.views[:j]
}

func setFactory(f ResourceExporterFactory) error {
	if f == nil {
		return errors.New("do not setFactory(nil)")
	}

	allMeters.lock.Lock()
	defer allMeters.lock.Unlock()

	allMeters.factory = f

	var retErr error

	for r, meter := range allMeters.meters {
		e, err := f(resourceFromKey(r))
		if err != nil {
			retErr = err
			continue // Keep trying to clean up remaining Meters.
		}
		meter.m.UnregisterExporter(meter.e)
		meter.m.RegisterExporter(e)
		meter.e = e
	}
	return retErr
}

func setReportingPeriod(mc *metricsConfig) {
	allMeters.lock.Lock()
	defer allMeters.lock.Unlock()

	rp := time.Duration(0)
	if mc != nil {
		rp = mc.reportingPeriod
	}
	for _, meter := range allMeters.meters {
		meter.m.SetReportingPeriod(rp)
	}
}

func flushResourceExporters() {
	allMeters.lock.Lock()
	defer allMeters.lock.Unlock()

	for _, meter := range allMeters.meters {
		flushGivenExporter(meter.e)
	}
}

// ClearMetersForTest clears the internal set of metrics being exported,
// including cleaning up background threads.
func ClearMetersForTest() {
	allMeters.lock.Lock()
	defer allMeters.lock.Unlock()

	for k, meter := range allMeters.meters {
		if k == "" {
			continue
		}
		meter.m.Stop()
		delete(allMeters.meters, k)
	}
}

func meterExporterForResource(r *resource.Resource) *meterExporter {
	key := resourceToKey(r)
	mE := allMeters.meters[key]
	if mE == nil {
		mE = &meterExporter{}
		allMeters.meters[key] = mE
	}
	if mE.o != nil {
		return mE
	}
	mE.m = view.NewMeter()
	mE.m.SetResource(r)
	mE.m.Start()

	mc := getCurMetricsConfig()
	if mc != nil {
		mE.m.SetReportingPeriod(mc.reportingPeriod)
	}
	resourceViews.lock.Lock()
	defer resourceViews.lock.Unlock()
	// make a copy of views to avoid data races
	viewsCopy := make([]*view.View, 0, len(resourceViews.views))
	for _, v := range resourceViews.views {
		c := *v
		viewsCopy = append(viewsCopy, &c)
	}
	mE.m.Register(viewsCopy...)
	mE.o = stats.WithRecorder(mE.m)
	allMeters.meters[key] = mE
	return mE
}

// optionForResource finds or creates a stats exporter for the resource, and
// returns a stats.Option indicating which meter to record to.
func optionForResource(r *resource.Resource) (stats.Options, error) {
	allMeters.lock.Lock()
	defer allMeters.lock.Unlock()

	mE := meterExporterForResource(r)
	if mE == nil {
		return nil, fmt.Errorf("failed lookup for resource %v", r)
	}

	if mE.e != nil {
		// Assume the exporter is already started.
		return mE.o, nil
	}

	if allMeters.factory == nil {
		if mE.o != nil {
			// If we can't create exporters but we have a Meter, return that.
			return mE.o, nil
		}
		return nil, fmt.Errorf("whoops, allMeters.factory is nil")
	}
	exporter, err := allMeters.factory(r)
	if err != nil {
		return nil, err
	}

	mE.m.RegisterExporter(exporter)
	mE.e = exporter
	return mE.o, nil
}

func resourceToKey(r *resource.Resource) string {
	if r == nil {
		return ""
	}
	var s strings.Builder
	l := len(r.Type)
	kvs := make([]string, 0, len(r.Labels))
	for k, v := range r.Labels {
		l += len(k) + len(v) + 2
		// We use byte values 1 and 2 to avoid colliding with valid resource labels
		// and to make unpacking easy
		kvs = append(kvs, fmt.Sprintf("\x01%s\x02%s", k, v))
	}
	s.Grow(l)
	s.WriteString(r.Type)

	sort.Strings(kvs) // Go maps are unsorted, so sort by key to produce stable output.
	for _, kv := range kvs {
		s.WriteString(kv)
	}

	return s.String()
}

func resourceFromKey(s string) *resource.Resource {
	if s == "" {
		return nil
	}
	r := &resource.Resource{Labels: map[string]string{}}
	parts := strings.Split(s, "\x01")
	r.Type = parts[0]
	for _, label := range parts[1:] {
		keyValue := strings.SplitN(label, "\x02", 2)
		r.Labels[keyValue[0]] = keyValue[1]
	}
	return r
}

// defaultMeterImpl is a pass-through to the default worker in OpenCensus.
type defaultMeterImpl struct{}

var _ view.Meter = (*defaultMeterImpl)(nil)

// defaultMeter is a custom meterExporter which passes through to the default
// worker in OpenCensus. This allows legacy code that uses OpenCensus and does
// not store a Resource in the context to continue to interoperate.
var defaultMeter = meterExporter{
	m: &defaultMeterImpl{},
	o: stats.WithRecorder(nil),
}

func (*defaultMeterImpl) Record(*tag.Map, interface{}, map[string]interface{}) {
	// using an empty option prevents this from being called
}

// Find calls view.Find
func (*defaultMeterImpl) Find(name string) *view.View {
	return view.Find(name)
}

// Register calls view.Register
func (*defaultMeterImpl) Register(views ...*view.View) error {
	return view.Register(views...)
}
func (*defaultMeterImpl) Unregister(views ...*view.View) {
	view.Unregister(views...)
}
func (*defaultMeterImpl) SetReportingPeriod(t time.Duration) {
	view.SetReportingPeriod(t)
}
func (*defaultMeterImpl) RegisterExporter(e view.Exporter) {
	view.RegisterExporter(e)
}
func (*defaultMeterImpl) UnregisterExporter(e view.Exporter) {
	view.UnregisterExporter(e)
}
func (*defaultMeterImpl) Start() {}
func (*defaultMeterImpl) Stop()  {}
func (*defaultMeterImpl) RetrieveData(viewName string) ([]*view.Row, error) {
	return view.RetrieveData(viewName)
}
func (*defaultMeterImpl) SetResource(*resource.Resource) {
}

// Read is implemented to support casting defaultMeterImpl to a metricproducer.Producer,
// but returns no values because the prometheus exporter (which is the only consumer)
// already has a built in path which collects these metrics via metricexport, which calls
// concat(x.Read() for x in metricproducer.GlobalManager.GetAll()).
func (*defaultMeterImpl) Read() []*metricdata.Metric {
	return []*metricdata.Metric{}
}
<commit_msg>Deep copy views to avoid races (#1720)<commit_after>/*
Copyright 2020 The Knative Authors

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

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

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

package metrics

import (
	"errors"
	"fmt"
	"sort"
	"strings"
	"sync"
	"time"

	"go.opencensus.io/metric/metricdata"
	"go.opencensus.io/resource"
	"go.opencensus.io/stats"
	"go.opencensus.io/stats/view"
	"go.opencensus.io/tag"
)

type storedViews struct {
	views []*view.View
	lock  sync.Mutex
}

type meterExporter struct {
	m view.Meter    // NOTE: DO NOT RETURN THIS DIRECTLY; the view.Meter will not work for the empty Resource
	o stats.Options // Cache the option to reduce allocations
	e view.Exporter
}

// ResourceExporterFactory provides a hook for producing separate view.Exporters
// for each observed Resource. This is needed because OpenCensus support for
// Resources is a bit tacked-on rather than being a first-class component like
// Tags are.
type ResourceExporterFactory func(*resource.Resource) (view.Exporter, error)
type meters struct {
	meters  map[string]*meterExporter
	factory ResourceExporterFactory
	lock    sync.Mutex
}

// Lock regime: lock allMeters before resourceViews. The critical path is in
// optionForResource, which must lock allMeters, but only needs to lock
// resourceViews if a new meter needs to be created.
var resourceViews = storedViews{}
var allMeters = meters{
	meters: map[string]*meterExporter{"": &defaultMeter},
}

// RegisterResourceView is similar to view.Register(), except that it will
// register the view across all Resources tracked by the system, rather than
// simply the default view.
func RegisterResourceView(views ...*view.View) error {
	var err error
	allMeters.lock.Lock()
	defer allMeters.lock.Unlock()
	resourceViews.lock.Lock()
	defer resourceViews.lock.Unlock()
	for _, meter := range allMeters.meters {
		// make a copy of views to avoid data races
		viewCopy := copyViews(views)
		if e := meter.m.Register(viewCopy...); e != nil {
			err = e
		}
	}
	if err != nil {
		return err
	}
	resourceViews.views = append(resourceViews.views, views...)
	return nil
}

// UnregisterResourceView is similar to view.Unregiste(), except that it will
// unregister the view across all Resources tracked byt he system, rather than
// simply the default view.
func UnregisterResourceView(views ...*view.View) {
	allMeters.lock.Lock()
	defer allMeters.lock.Unlock()

	resourceViews.lock.Lock()
	defer resourceViews.lock.Unlock()

	for _, meter := range allMeters.meters {
		// Since we make a defensive copy of all views in RegisterResourceView,
		// the caller might not have the same view pointer that was registered.
		// Use Meter.Find() to find the view with a matching name.
		for _, v := range views {
			name := v.Name
			if v.Name == "" {
				name = v.Measure.Name()
			}
			if v := meter.m.Find(name); v != nil {
				meter.m.Unregister(v)
			}
		}
	}

	j := 0
	for _, view := range resourceViews.views {
		toRemove := false
		for _, viewToRemove := range views {
			if view == viewToRemove {
				toRemove = true
				break
			}
		}
		if !toRemove {
			resourceViews.views[j] = view
			j++
		}
	}
	resourceViews.views = resourceViews.views[:j]
}

func setFactory(f ResourceExporterFactory) error {
	if f == nil {
		return errors.New("do not setFactory(nil)")
	}

	allMeters.lock.Lock()
	defer allMeters.lock.Unlock()

	allMeters.factory = f

	var retErr error

	for r, meter := range allMeters.meters {
		e, err := f(resourceFromKey(r))
		if err != nil {
			retErr = err
			continue // Keep trying to clean up remaining Meters.
		}
		meter.m.UnregisterExporter(meter.e)
		meter.m.RegisterExporter(e)
		meter.e = e
	}
	return retErr
}

func setReportingPeriod(mc *metricsConfig) {
	allMeters.lock.Lock()
	defer allMeters.lock.Unlock()

	rp := time.Duration(0)
	if mc != nil {
		rp = mc.reportingPeriod
	}
	for _, meter := range allMeters.meters {
		meter.m.SetReportingPeriod(rp)
	}
}

func flushResourceExporters() {
	allMeters.lock.Lock()
	defer allMeters.lock.Unlock()

	for _, meter := range allMeters.meters {
		flushGivenExporter(meter.e)
	}
}

// ClearMetersForTest clears the internal set of metrics being exported,
// including cleaning up background threads.
func ClearMetersForTest() {
	allMeters.lock.Lock()
	defer allMeters.lock.Unlock()

	for k, meter := range allMeters.meters {
		if k == "" {
			continue
		}
		meter.m.Stop()
		delete(allMeters.meters, k)
	}
}

func meterExporterForResource(r *resource.Resource) *meterExporter {
	key := resourceToKey(r)
	mE := allMeters.meters[key]
	if mE == nil {
		mE = &meterExporter{}
		allMeters.meters[key] = mE
	}
	if mE.o != nil {
		return mE
	}
	mE.m = view.NewMeter()
	mE.m.SetResource(r)
	mE.m.Start()

	mc := getCurMetricsConfig()
	if mc != nil {
		mE.m.SetReportingPeriod(mc.reportingPeriod)
	}
	resourceViews.lock.Lock()
	defer resourceViews.lock.Unlock()
	// make a copy of views to avoid data races
	viewsCopy := copyViews(resourceViews.views)
	mE.m.Register(viewsCopy...)
	mE.o = stats.WithRecorder(mE.m)
	allMeters.meters[key] = mE
	return mE
}

// optionForResource finds or creates a stats exporter for the resource, and
// returns a stats.Option indicating which meter to record to.
func optionForResource(r *resource.Resource) (stats.Options, error) {
	allMeters.lock.Lock()
	defer allMeters.lock.Unlock()

	mE := meterExporterForResource(r)
	if mE == nil {
		return nil, fmt.Errorf("failed lookup for resource %v", r)
	}

	if mE.e != nil {
		// Assume the exporter is already started.
		return mE.o, nil
	}

	if allMeters.factory == nil {
		if mE.o != nil {
			// If we can't create exporters but we have a Meter, return that.
			return mE.o, nil
		}
		return nil, fmt.Errorf("whoops, allMeters.factory is nil")
	}
	exporter, err := allMeters.factory(r)
	if err != nil {
		return nil, err
	}

	mE.m.RegisterExporter(exporter)
	mE.e = exporter
	return mE.o, nil
}

func resourceToKey(r *resource.Resource) string {
	if r == nil {
		return ""
	}
	var s strings.Builder
	l := len(r.Type)
	kvs := make([]string, 0, len(r.Labels))
	for k, v := range r.Labels {
		l += len(k) + len(v) + 2
		// We use byte values 1 and 2 to avoid colliding with valid resource labels
		// and to make unpacking easy
		kvs = append(kvs, fmt.Sprintf("\x01%s\x02%s", k, v))
	}
	s.Grow(l)
	s.WriteString(r.Type)

	sort.Strings(kvs) // Go maps are unsorted, so sort by key to produce stable output.
	for _, kv := range kvs {
		s.WriteString(kv)
	}

	return s.String()
}

func resourceFromKey(s string) *resource.Resource {
	if s == "" {
		return nil
	}
	r := &resource.Resource{Labels: map[string]string{}}
	parts := strings.Split(s, "\x01")
	r.Type = parts[0]
	for _, label := range parts[1:] {
		keyValue := strings.SplitN(label, "\x02", 2)
		r.Labels[keyValue[0]] = keyValue[1]
	}
	return r
}

func copyViews(views []*view.View) []*view.View {
	viewsCopy := make([]*view.View, 0, len(resourceViews.views))
	for _, v := range views {
		c := *v
		c.TagKeys = make([]tag.Key, len(v.TagKeys))
		copy(c.TagKeys, v.TagKeys)
		agg := *v.Aggregation
		c.Aggregation = &agg
		c.Aggregation.Buckets = make([]float64, len(v.Aggregation.Buckets))
		copy(c.Aggregation.Buckets, v.Aggregation.Buckets)
		viewsCopy = append(viewsCopy, &c)
	}
	return viewsCopy
}

// defaultMeterImpl is a pass-through to the default worker in OpenCensus.
type defaultMeterImpl struct{}

var _ view.Meter = (*defaultMeterImpl)(nil)

// defaultMeter is a custom meterExporter which passes through to the default
// worker in OpenCensus. This allows legacy code that uses OpenCensus and does
// not store a Resource in the context to continue to interoperate.
var defaultMeter = meterExporter{
	m: &defaultMeterImpl{},
	o: stats.WithRecorder(nil),
}

func (*defaultMeterImpl) Record(*tag.Map, interface{}, map[string]interface{}) {
	// using an empty option prevents this from being called
}

// Find calls view.Find
func (*defaultMeterImpl) Find(name string) *view.View {
	return view.Find(name)
}

// Register calls view.Register
func (*defaultMeterImpl) Register(views ...*view.View) error {
	return view.Register(views...)
}
func (*defaultMeterImpl) Unregister(views ...*view.View) {
	view.Unregister(views...)
}
func (*defaultMeterImpl) SetReportingPeriod(t time.Duration) {
	view.SetReportingPeriod(t)
}
func (*defaultMeterImpl) RegisterExporter(e view.Exporter) {
	view.RegisterExporter(e)
}
func (*defaultMeterImpl) UnregisterExporter(e view.Exporter) {
	view.UnregisterExporter(e)
}
func (*defaultMeterImpl) Start() {}
func (*defaultMeterImpl) Stop()  {}
func (*defaultMeterImpl) RetrieveData(viewName string) ([]*view.Row, error) {
	return view.RetrieveData(viewName)
}
func (*defaultMeterImpl) SetResource(*resource.Resource) {
}

// Read is implemented to support casting defaultMeterImpl to a metricproducer.Producer,
// but returns no values because the prometheus exporter (which is the only consumer)
// already has a built in path which collects these metrics via metricexport, which calls
// concat(x.Read() for x in metricproducer.GlobalManager.GetAll()).
func (*defaultMeterImpl) Read() []*metricdata.Metric {
	return []*metricdata.Metric{}
}
<|endoftext|>


<commit_before>package libnetwork

import (
	"fmt"
	"math/rand"
	"net"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/libnetwork/iptables"
	"github.com/docker/libnetwork/netutils"
	"github.com/miekg/dns"
)

// Resolver represents the embedded DNS server in Docker. It operates
// by listening on container's loopback interface for DNS queries.
type Resolver interface {
	// Start starts the name server for the container
	Start() error
	// Stop stops the name server for the container. Stopped resolver
	// can be reused after running the SetupFunc again.
	Stop()
	// SetupFunc() provides the setup function that should be run
	// in the container's network namespace.
	SetupFunc() func()
	// NameServer() returns the IP of the DNS resolver for the
	// containers.
	NameServer() string
	// SetExtServers configures the external nameservers the resolver
	// should use to forward queries
	SetExtServers([]string)
	// FlushExtServers clears the cached UDP connections to external
	// nameservers
	FlushExtServers()
	// ResolverOptions returns resolv.conf options that should be set
	ResolverOptions() []string
}

const (
	resolverIP      = "127.0.0.11"
	dnsPort         = "53"
	ptrIPv4domain   = ".in-addr.arpa."
	ptrIPv6domain   = ".ip6.arpa."
	respTTL         = 600
	maxExtDNS       = 3 //max number of external servers to try
	extIOTimeout    = 3 * time.Second
	defaultRespSize = 512
)

type extDNSEntry struct {
	ipStr   string
	extConn net.Conn
	extOnce sync.Once
}

// resolver implements the Resolver interface
type resolver struct {
	sb         *sandbox
	extDNSList [maxExtDNS]extDNSEntry
	server     *dns.Server
	conn       *net.UDPConn
	tcpServer  *dns.Server
	tcpListen  *net.TCPListener
	err        error
}

func init() {
	rand.Seed(time.Now().Unix())
}

// NewResolver creates a new instance of the Resolver
func NewResolver(sb *sandbox) Resolver {
	return &resolver{
		sb:  sb,
		err: fmt.Errorf("setup not done yet"),
	}
}

func (r *resolver) SetupFunc() func() {
	return (func() {
		var err error

		// DNS operates primarily on UDP
		addr := &net.UDPAddr{
			IP: net.ParseIP(resolverIP),
		}

		r.conn, err = net.ListenUDP("udp", addr)
		if err != nil {
			r.err = fmt.Errorf("error in opening name server socket %v", err)
			return
		}
		laddr := r.conn.LocalAddr()
		_, ipPort, _ := net.SplitHostPort(laddr.String())

		// Listen on a TCP as well
		tcpaddr := &net.TCPAddr{
			IP: net.ParseIP(resolverIP),
		}

		r.tcpListen, err = net.ListenTCP("tcp", tcpaddr)
		if err != nil {
			r.err = fmt.Errorf("error in opening name TCP server socket %v", err)
			return
		}
		ltcpaddr := r.tcpListen.Addr()
		_, tcpPort, _ := net.SplitHostPort(ltcpaddr.String())
		rules := [][]string{
			{"-t", "nat", "-A", "OUTPUT", "-d", resolverIP, "-p", "udp", "--dport", dnsPort, "-j", "DNAT", "--to-destination", laddr.String()},
			{"-t", "nat", "-A", "POSTROUTING", "-s", resolverIP, "-p", "udp", "--sport", ipPort, "-j", "SNAT", "--to-source", ":" + dnsPort},
			{"-t", "nat", "-A", "OUTPUT", "-d", resolverIP, "-p", "tcp", "--dport", dnsPort, "-j", "DNAT", "--to-destination", ltcpaddr.String()},
			{"-t", "nat", "-A", "POSTROUTING", "-s", resolverIP, "-p", "tcp", "--sport", tcpPort, "-j", "SNAT", "--to-source", ":" + dnsPort},
		}

		for _, rule := range rules {
			r.err = iptables.RawCombinedOutputNative(rule...)
			if r.err != nil {
				return
			}
		}
		r.err = nil
	})
}

func (r *resolver) Start() error {
	// make sure the resolver has been setup before starting
	if r.err != nil {
		return r.err
	}
	s := &dns.Server{Handler: r, PacketConn: r.conn}
	r.server = s
	go func() {
		s.ActivateAndServe()
	}()

	tcpServer := &dns.Server{Handler: r, Listener: r.tcpListen}
	r.tcpServer = tcpServer
	go func() {
		tcpServer.ActivateAndServe()
	}()
	return nil
}

func (r *resolver) FlushExtServers() {
	for i := 0; i < maxExtDNS; i++ {
		r.extDNSList[i].extConn = nil
		r.extDNSList[i].extOnce = sync.Once{}
	}
}

func (r *resolver) Stop() {
	r.FlushExtServers()

	if r.server != nil {
		r.server.Shutdown()
	}
	if r.tcpServer != nil {
		r.tcpServer.Shutdown()
	}
	r.conn = nil
	r.tcpServer = nil
	r.err = fmt.Errorf("setup not done yet")
}

func (r *resolver) SetExtServers(dns []string) {
	l := len(dns)
	if l > maxExtDNS {
		l = maxExtDNS
	}
	for i := 0; i < l; i++ {
		r.extDNSList[i].ipStr = dns[i]
	}
}

func (r *resolver) NameServer() string {
	return resolverIP
}

func (r *resolver) ResolverOptions() []string {
	return []string{"ndots:0"}
}

func setCommonFlags(msg *dns.Msg) {
	msg.RecursionAvailable = true
}

func shuffleAddr(addr []net.IP) []net.IP {
	for i := len(addr) - 1; i > 0; i-- {
		r := rand.Intn(i + 1)
		addr[i], addr[r] = addr[r], addr[i]
	}
	return addr
}

func createRespMsg(query *dns.Msg) *dns.Msg {
	resp := new(dns.Msg)
	resp.SetReply(query)
	setCommonFlags(resp)

	return resp
}

func (r *resolver) handleIPQuery(name string, query *dns.Msg, ipType int) (*dns.Msg, error) {
	addr, ipv6Miss := r.sb.ResolveName(name, ipType)
	if addr == nil && ipv6Miss {
		// Send a reply without any Answer sections
		log.Debugf("Lookup name %s present without IPv6 address", name)
		resp := createRespMsg(query)
		return resp, nil
	}
	if addr == nil {
		return nil, nil
	}

	log.Debugf("Lookup for %s: IP %v", name, addr)

	resp := createRespMsg(query)
	if len(addr) > 1 {
		addr = shuffleAddr(addr)
	}
	if ipType == netutils.IPv4 {
		for _, ip := range addr {
			rr := new(dns.A)
			rr.Hdr = dns.RR_Header{Name: name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: respTTL}
			rr.A = ip
			resp.Answer = append(resp.Answer, rr)
		}
	} else {
		for _, ip := range addr {
			rr := new(dns.AAAA)
			rr.Hdr = dns.RR_Header{Name: name, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: respTTL}
			rr.AAAA = ip
			resp.Answer = append(resp.Answer, rr)
		}
	}
	return resp, nil
}

func (r *resolver) handlePTRQuery(ptr string, query *dns.Msg) (*dns.Msg, error) {
	parts := []string{}

	if strings.HasSuffix(ptr, ptrIPv4domain) {
		parts = strings.Split(ptr, ptrIPv4domain)
	} else if strings.HasSuffix(ptr, ptrIPv6domain) {
		parts = strings.Split(ptr, ptrIPv6domain)
	} else {
		return nil, fmt.Errorf("invalid PTR query, %v", ptr)
	}

	host := r.sb.ResolveIP(parts[0])
	if len(host) == 0 {
		return nil, nil
	}

	log.Debugf("Lookup for IP %s: name %s", parts[0], host)
	fqdn := dns.Fqdn(host)

	resp := new(dns.Msg)
	resp.SetReply(query)
	setCommonFlags(resp)

	rr := new(dns.PTR)
	rr.Hdr = dns.RR_Header{Name: ptr, Rrtype: dns.TypePTR, Class: dns.ClassINET, Ttl: respTTL}
	rr.Ptr = fqdn
	resp.Answer = append(resp.Answer, rr)
	return resp, nil
}

func truncateResp(resp *dns.Msg, maxSize int, isTCP bool) {
	if !isTCP {
		resp.Truncated = true
	}

	// trim the Answer RRs one by one till the whole message fits
	// within the reply size
	for resp.Len() > maxSize {
		resp.Answer = resp.Answer[:len(resp.Answer)-1]
	}
}

func (r *resolver) ServeDNS(w dns.ResponseWriter, query *dns.Msg) {
	var (
		extConn net.Conn
		resp    *dns.Msg
		err     error
	)

	if query == nil || len(query.Question) == 0 {
		return
	}
	name := query.Question[0].Name
	if query.Question[0].Qtype == dns.TypeA {
		resp, err = r.handleIPQuery(name, query, netutils.IPv4)
	} else if query.Question[0].Qtype == dns.TypeAAAA {
		resp, err = r.handleIPQuery(name, query, netutils.IPv6)
	} else if query.Question[0].Qtype == dns.TypePTR {
		resp, err = r.handlePTRQuery(name, query)
	}

	if err != nil {
		log.Error(err)
		return
	}

	proto := w.LocalAddr().Network()
	maxSize := 0
	if proto == "tcp" {
		maxSize = dns.MaxMsgSize - 1
	} else if proto == "udp" {
		optRR := query.IsEdns0()
		if optRR != nil {
			maxSize = int(optRR.UDPSize())
		}
		if maxSize < defaultRespSize {
			maxSize = defaultRespSize
		}
	}

	if resp != nil {
		if resp.Len() > maxSize {
			truncateResp(resp, maxSize, proto == "tcp")
		}
	} else {
		for i := 0; i < maxExtDNS; i++ {
			extDNS := &r.extDNSList[i]
			if extDNS.ipStr == "" {
				break
			}
			log.Debugf("Querying ext dns %s:%s for %s[%d]", proto, extDNS.ipStr, name, query.Question[0].Qtype)

			extConnect := func() {
				addr := fmt.Sprintf("%s:%d", extDNS.ipStr, 53)
				extConn, err = net.DialTimeout(proto, addr, extIOTimeout)
			}

			// For udp clients connection is persisted to reuse for further queries.
			// Accessing extDNS.extConn be a race here between go rouines. Hence the
			// connection setup is done in a Once block and fetch the extConn again
			extConn = extDNS.extConn
			if extConn == nil || proto == "tcp" {
				if proto == "udp" {
					extDNS.extOnce.Do(func() {
						r.sb.execFunc(extConnect)
						extDNS.extConn = extConn
					})
					extConn = extDNS.extConn
				} else {
					r.sb.execFunc(extConnect)
				}
				if err != nil {
					log.Debugf("Connect failed, %s", err)
					continue
				}
			}
			// If two go routines are executing in parralel one will
			// block on the Once.Do and in case of error connecting
			// to the external server it will end up with a nil err
			// but extConn also being nil.
			if extConn == nil {
				continue
			}

			// Timeout has to be set for every IO operation.
			extConn.SetDeadline(time.Now().Add(extIOTimeout))
			co := &dns.Conn{Conn: extConn}

			defer func() {
				if proto == "tcp" {
					co.Close()
				}
			}()
			err = co.WriteMsg(query)
			if err != nil {
				log.Debugf("Send to DNS server failed, %s", err)
				continue
			}

			resp, err = co.ReadMsg()
			if err != nil {
				log.Debugf("Read from DNS server failed, %s", err)
				continue
			}

			resp.Compress = true
			break
		}

		if resp == nil {
			return
		}
	}

	err = w.WriteMsg(resp)
	if err != nil {
		log.Errorf("error writing resolver resp, %s", err)
	}
}
<commit_msg>Limit number of concurrent DNS queries<commit_after>package libnetwork

import (
	"fmt"
	"math/rand"
	"net"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/libnetwork/iptables"
	"github.com/docker/libnetwork/netutils"
	"github.com/miekg/dns"
)

// Resolver represents the embedded DNS server in Docker. It operates
// by listening on container's loopback interface for DNS queries.
type Resolver interface {
	// Start starts the name server for the container
	Start() error
	// Stop stops the name server for the container. Stopped resolver
	// can be reused after running the SetupFunc again.
	Stop()
	// SetupFunc() provides the setup function that should be run
	// in the container's network namespace.
	SetupFunc() func()
	// NameServer() returns the IP of the DNS resolver for the
	// containers.
	NameServer() string
	// SetExtServers configures the external nameservers the resolver
	// should use to forward queries
	SetExtServers([]string)
	// FlushExtServers clears the cached UDP connections to external
	// nameservers
	FlushExtServers()
	// ResolverOptions returns resolv.conf options that should be set
	ResolverOptions() []string
}

const (
	resolverIP      = "127.0.0.11"
	dnsPort         = "53"
	ptrIPv4domain   = ".in-addr.arpa."
	ptrIPv6domain   = ".ip6.arpa."
	respTTL         = 600
	maxExtDNS       = 3 //max number of external servers to try
	extIOTimeout    = 3 * time.Second
	defaultRespSize = 512
	maxConcurrent   = 50
	logInterval     = 2 * time.Second
)

type extDNSEntry struct {
	ipStr   string
	extConn net.Conn
	extOnce sync.Once
}

// resolver implements the Resolver interface
type resolver struct {
	sb         *sandbox
	extDNSList [maxExtDNS]extDNSEntry
	server     *dns.Server
	conn       *net.UDPConn
	tcpServer  *dns.Server
	tcpListen  *net.TCPListener
	err        error
	count      int32
	tStamp     time.Time
	queryLock  sync.Mutex
}

func init() {
	rand.Seed(time.Now().Unix())
}

// NewResolver creates a new instance of the Resolver
func NewResolver(sb *sandbox) Resolver {
	return &resolver{
		sb:  sb,
		err: fmt.Errorf("setup not done yet"),
	}
}

func (r *resolver) SetupFunc() func() {
	return (func() {
		var err error

		// DNS operates primarily on UDP
		addr := &net.UDPAddr{
			IP: net.ParseIP(resolverIP),
		}

		r.conn, err = net.ListenUDP("udp", addr)
		if err != nil {
			r.err = fmt.Errorf("error in opening name server socket %v", err)
			return
		}
		laddr := r.conn.LocalAddr()
		_, ipPort, _ := net.SplitHostPort(laddr.String())

		// Listen on a TCP as well
		tcpaddr := &net.TCPAddr{
			IP: net.ParseIP(resolverIP),
		}

		r.tcpListen, err = net.ListenTCP("tcp", tcpaddr)
		if err != nil {
			r.err = fmt.Errorf("error in opening name TCP server socket %v", err)
			return
		}
		ltcpaddr := r.tcpListen.Addr()
		_, tcpPort, _ := net.SplitHostPort(ltcpaddr.String())
		rules := [][]string{
			{"-t", "nat", "-A", "OUTPUT", "-d", resolverIP, "-p", "udp", "--dport", dnsPort, "-j", "DNAT", "--to-destination", laddr.String()},
			{"-t", "nat", "-A", "POSTROUTING", "-s", resolverIP, "-p", "udp", "--sport", ipPort, "-j", "SNAT", "--to-source", ":" + dnsPort},
			{"-t", "nat", "-A", "OUTPUT", "-d", resolverIP, "-p", "tcp", "--dport", dnsPort, "-j", "DNAT", "--to-destination", ltcpaddr.String()},
			{"-t", "nat", "-A", "POSTROUTING", "-s", resolverIP, "-p", "tcp", "--sport", tcpPort, "-j", "SNAT", "--to-source", ":" + dnsPort},
		}

		for _, rule := range rules {
			r.err = iptables.RawCombinedOutputNative(rule...)
			if r.err != nil {
				return
			}
		}
		r.err = nil
	})
}

func (r *resolver) Start() error {
	// make sure the resolver has been setup before starting
	if r.err != nil {
		return r.err
	}
	s := &dns.Server{Handler: r, PacketConn: r.conn}
	r.server = s
	go func() {
		s.ActivateAndServe()
	}()

	tcpServer := &dns.Server{Handler: r, Listener: r.tcpListen}
	r.tcpServer = tcpServer
	go func() {
		tcpServer.ActivateAndServe()
	}()
	return nil
}

func (r *resolver) FlushExtServers() {
	for i := 0; i < maxExtDNS; i++ {
		r.extDNSList[i].extConn = nil
		r.extDNSList[i].extOnce = sync.Once{}
	}
}

func (r *resolver) Stop() {
	r.FlushExtServers()

	if r.server != nil {
		r.server.Shutdown()
	}
	if r.tcpServer != nil {
		r.tcpServer.Shutdown()
	}
	r.conn = nil
	r.tcpServer = nil
	r.err = fmt.Errorf("setup not done yet")
	r.tStamp = time.Time{}
	r.count = 0
	r.queryLock = sync.Mutex{}
}

func (r *resolver) SetExtServers(dns []string) {
	l := len(dns)
	if l > maxExtDNS {
		l = maxExtDNS
	}
	for i := 0; i < l; i++ {
		r.extDNSList[i].ipStr = dns[i]
	}
}

func (r *resolver) NameServer() string {
	return resolverIP
}

func (r *resolver) ResolverOptions() []string {
	return []string{"ndots:0"}
}

func setCommonFlags(msg *dns.Msg) {
	msg.RecursionAvailable = true
}

func shuffleAddr(addr []net.IP) []net.IP {
	for i := len(addr) - 1; i > 0; i-- {
		r := rand.Intn(i + 1)
		addr[i], addr[r] = addr[r], addr[i]
	}
	return addr
}

func createRespMsg(query *dns.Msg) *dns.Msg {
	resp := new(dns.Msg)
	resp.SetReply(query)
	setCommonFlags(resp)

	return resp
}

func (r *resolver) handleIPQuery(name string, query *dns.Msg, ipType int) (*dns.Msg, error) {
	addr, ipv6Miss := r.sb.ResolveName(name, ipType)
	if addr == nil && ipv6Miss {
		// Send a reply without any Answer sections
		log.Debugf("Lookup name %s present without IPv6 address", name)
		resp := createRespMsg(query)
		return resp, nil
	}
	if addr == nil {
		return nil, nil
	}

	log.Debugf("Lookup for %s: IP %v", name, addr)

	resp := createRespMsg(query)
	if len(addr) > 1 {
		addr = shuffleAddr(addr)
	}
	if ipType == netutils.IPv4 {
		for _, ip := range addr {
			rr := new(dns.A)
			rr.Hdr = dns.RR_Header{Name: name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: respTTL}
			rr.A = ip
			resp.Answer = append(resp.Answer, rr)
		}
	} else {
		for _, ip := range addr {
			rr := new(dns.AAAA)
			rr.Hdr = dns.RR_Header{Name: name, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: respTTL}
			rr.AAAA = ip
			resp.Answer = append(resp.Answer, rr)
		}
	}
	return resp, nil
}

func (r *resolver) handlePTRQuery(ptr string, query *dns.Msg) (*dns.Msg, error) {
	parts := []string{}

	if strings.HasSuffix(ptr, ptrIPv4domain) {
		parts = strings.Split(ptr, ptrIPv4domain)
	} else if strings.HasSuffix(ptr, ptrIPv6domain) {
		parts = strings.Split(ptr, ptrIPv6domain)
	} else {
		return nil, fmt.Errorf("invalid PTR query, %v", ptr)
	}

	host := r.sb.ResolveIP(parts[0])
	if len(host) == 0 {
		return nil, nil
	}

	log.Debugf("Lookup for IP %s: name %s", parts[0], host)
	fqdn := dns.Fqdn(host)

	resp := new(dns.Msg)
	resp.SetReply(query)
	setCommonFlags(resp)

	rr := new(dns.PTR)
	rr.Hdr = dns.RR_Header{Name: ptr, Rrtype: dns.TypePTR, Class: dns.ClassINET, Ttl: respTTL}
	rr.Ptr = fqdn
	resp.Answer = append(resp.Answer, rr)
	return resp, nil
}

func truncateResp(resp *dns.Msg, maxSize int, isTCP bool) {
	if !isTCP {
		resp.Truncated = true
	}

	// trim the Answer RRs one by one till the whole message fits
	// within the reply size
	for resp.Len() > maxSize {
		resp.Answer = resp.Answer[:len(resp.Answer)-1]
	}
}

func (r *resolver) ServeDNS(w dns.ResponseWriter, query *dns.Msg) {
	var (
		extConn net.Conn
		resp    *dns.Msg
		err     error
	)

	if query == nil || len(query.Question) == 0 {
		return
	}
	name := query.Question[0].Name
	if query.Question[0].Qtype == dns.TypeA {
		resp, err = r.handleIPQuery(name, query, netutils.IPv4)
	} else if query.Question[0].Qtype == dns.TypeAAAA {
		resp, err = r.handleIPQuery(name, query, netutils.IPv6)
	} else if query.Question[0].Qtype == dns.TypePTR {
		resp, err = r.handlePTRQuery(name, query)
	}

	if err != nil {
		log.Error(err)
		return
	}

	proto := w.LocalAddr().Network()
	maxSize := 0
	if proto == "tcp" {
		maxSize = dns.MaxMsgSize - 1
	} else if proto == "udp" {
		optRR := query.IsEdns0()
		if optRR != nil {
			maxSize = int(optRR.UDPSize())
		}
		if maxSize < defaultRespSize {
			maxSize = defaultRespSize
		}
	}

	if resp != nil {
		if resp.Len() > maxSize {
			truncateResp(resp, maxSize, proto == "tcp")
		}
	} else {
		for i := 0; i < maxExtDNS; i++ {
			extDNS := &r.extDNSList[i]
			if extDNS.ipStr == "" {
				break
			}
			log.Debugf("Query %s[%d] from %s, forwarding to %s:%s", name, query.Question[0].Qtype,
				w.LocalAddr().String(), proto, extDNS.ipStr)

			extConnect := func() {
				addr := fmt.Sprintf("%s:%d", extDNS.ipStr, 53)
				extConn, err = net.DialTimeout(proto, addr, extIOTimeout)
			}

			// For udp clients connection is persisted to reuse for further queries.
			// Accessing extDNS.extConn be a race here between go rouines. Hence the
			// connection setup is done in a Once block and fetch the extConn again
			extConn = extDNS.extConn
			if extConn == nil || proto == "tcp" {
				if proto == "udp" {
					extDNS.extOnce.Do(func() {
						r.sb.execFunc(extConnect)
						extDNS.extConn = extConn
					})
					extConn = extDNS.extConn
				} else {
					r.sb.execFunc(extConnect)
				}
				if err != nil {
					log.Debugf("Connect failed, %s", err)
					continue
				}
			}
			// If two go routines are executing in parralel one will
			// block on the Once.Do and in case of error connecting
			// to the external server it will end up with a nil err
			// but extConn also being nil.
			if extConn == nil {
				continue
			}

			// Timeout has to be set for every IO operation.
			extConn.SetDeadline(time.Now().Add(extIOTimeout))
			co := &dns.Conn{Conn: extConn}

			if r.concurrentQueryInc() == false {
				old := r.tStamp
				r.tStamp = time.Now()
				if r.tStamp.Sub(old) > logInterval {
					log.Errorf("More than %v concurrent queries from %s", maxConcurrent, w.LocalAddr().String())
				}
				continue
			}

			defer func() {
				if proto == "tcp" {
					co.Close()
				}
			}()
			err = co.WriteMsg(query)
			if err != nil {
				r.concurrentQueryDec()
				log.Debugf("Send to DNS server failed, %s", err)
				continue
			}

			resp, err = co.ReadMsg()
			r.concurrentQueryDec()
			if err != nil {
				log.Debugf("Read from DNS server failed, %s", err)
				continue
			}

			resp.Compress = true
			break
		}

		if resp == nil {
			return
		}
	}

	err = w.WriteMsg(resp)
	if err != nil {
		log.Errorf("error writing resolver resp, %s", err)
	}
}

func (r *resolver) concurrentQueryInc() bool {
	r.queryLock.Lock()
	defer r.queryLock.Unlock()
	if r.count == maxConcurrent {
		return false
	}
	r.count++
	return true
}

func (r *resolver) concurrentQueryDec() bool {
	r.queryLock.Lock()
	defer r.queryLock.Unlock()
	if r.count == 0 {
		return false
	}
	r.count--
	return true
}
<|endoftext|>


<commit_before>package websocket_test

import (
	"net/http"
	"net/url"
	"strings"
	"testing"

	gorilla "github.com/gorilla/websocket"
	"github.com/spring1843/chat-server/src/chat"
	"github.com/spring1843/chat-server/src/config"
	"github.com/spring1843/chat-server/src/drivers/websocket"
)

func TestCantStartAndConnect(t *testing.T) {
	config := config.Config{
		WebAddress: "127.0.0.1:4008",
	}

	chatServer := chat.NewServer()
	chatServer.Listen()
	websocket.SetWebSocket(chatServer)

	http.HandleFunc("/ws", websocket.Handler)

	go func() {
		if err := http.ListenAndServe(config.WebAddress, nil); err != nil {
			t.Fatalf("Failed listening to Websocet on %s. Error: %s", config.WebAddress, err)
		}
	}()

	nickName := "User1"
	conn := connectUsers(t, nickName, config)
	disconnectUsers(t, conn, chatServer)
}

func connectUsers(t *testing.T, nickname string, config config.Config) *gorilla.Conn {
	url := url.URL{Scheme: "ws", Host: config.WebAddress, Path: "/ws"}

	conn, _, err := gorilla.DefaultDialer.Dial(url.String(), nil)
	if err != nil {
		t.Fatalf("Websocket Dial error: %s", err.Error())
	}

	_, message, err := conn.ReadMessage()
	if err != nil {
		t.Fatalf("Error while reading connection %s", err.Error())
	}

	if !strings.Contains(string(message), "Welcome") {
		t.Error("Could not receive welcome message")
	}

	if err := conn.WriteMessage(1, []byte(nickname)); err != nil {
		t.Fatalf("Error writing to connection. Error %s", err)
	}

	_, message, err = conn.ReadMessage()
	if err != nil {
		t.Fatalf("Error while reading connection. Error %s", err.Error())
	}

	expect := "Welcome " + nickname
	if !strings.Contains(string(message), expect) {
		t.Fatalf("Could not set user %s, expected 'Thanks User1' got %s", nickname, expect)
	}

	return conn
}

func disconnectUsers(t *testing.T, conn *gorilla.Conn, chatServer *chat.Server) {
	if err := conn.WriteMessage(1, []byte(`/quit`)); err != nil {
		t.Fatalf("Error writing to connection. Error %s", err)
	}

	_, message, err := conn.ReadMessage()
	if err != nil {
		t.Fatalf("Failed reading from WebSocket connection. Error %s", err)
	}
	if !strings.Contains(string(message), "Good Bye") {
		t.Fatalf("Could not quit from server. Expected 'Good Bye' got %s", string(message))
	}

	if chatServer.IsUserConnected("User1") {
		t.Fatal("User is still connected to server after quiting")
	}
}
<commit_msg>Change websocket test to includes more users<commit_after>package websocket_test

import (
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"testing"

	gorilla "github.com/gorilla/websocket"
	"github.com/spring1843/chat-server/src/chat"
	"github.com/spring1843/chat-server/src/config"
	"github.com/spring1843/chat-server/src/drivers/websocket"
)

func TestCantStartAndConnect(t *testing.T) {
	config := config.Config{
		WebAddress: "127.0.0.1:4008",
	}

	chatServer := chat.NewServer()
	chatServer.Listen()
	websocket.SetWebSocket(chatServer)

	http.HandleFunc("/ws", websocket.Handler)

	go func() {
		if err := http.ListenAndServe(config.WebAddress, nil); err != nil {
			t.Fatalf("Failed listening to Websocet on %s. Error: %s", config.WebAddress, err)
		}
	}()

	tryouts := 10
	conns := make([]*gorilla.Conn, tryouts, tryouts)
	i := 0
	for i < tryouts {
		nickName := fmt.Sprintf("user%d", i)
		conns[i] = connectUser(t, nickName, config)
		i++
	}

	if chatServer.ConnectedUsersCount() != tryouts {
		t.Fatalf("Expected user count to be %d after disconnecting users, got %d", tryouts, chatServer.ConnectedUsersCount())
	}

	i = 0
	for i < tryouts {
		disconnectUser(t, conns[i], chatServer)
		i++
	}

	if chatServer.ConnectedUsersCount() != 0 {
		t.Fatalf("Expected user count to be %d after disconnecting users, got %d", 0, chatServer.ConnectedUsersCount())
	}
}

func connectUser(t *testing.T, nickname string, config config.Config) *gorilla.Conn {
	url := url.URL{Scheme: "ws", Host: config.WebAddress, Path: "/ws"}

	conn, _, err := gorilla.DefaultDialer.Dial(url.String(), nil)
	if err != nil {
		t.Fatalf("Websocket Dial error: %s", err.Error())
	}

	_, message, err := conn.ReadMessage()
	if err != nil {
		t.Fatalf("Error while reading connection %s", err.Error())
	}

	if !strings.Contains(string(message), "Welcome") {
		t.Error("Could not receive welcome message")
	}

	if err := conn.WriteMessage(1, []byte(nickname)); err != nil {
		t.Fatalf("Error writing to connection. Error %s", err)
	}

	_, message, err = conn.ReadMessage()
	if err != nil {
		t.Fatalf("Error while reading connection. Error %s", err.Error())
	}

	expect := "Welcome " + nickname
	if !strings.Contains(string(message), expect) {
		t.Fatalf("Could not set user %s, expected 'Thanks User1' got %s", nickname, expect)
	}

	return conn
}

func disconnectUser(t *testing.T, conn *gorilla.Conn, chatServer *chat.Server) {
	if err := conn.WriteMessage(1, []byte(`/quit`)); err != nil {
		t.Fatalf("Error writing to connection. Error %s", err)
	}

	_, message, err := conn.ReadMessage()
	if err != nil {
		t.Fatalf("Failed reading from WebSocket connection. Error %s", err)
	}
	if !strings.Contains(string(message), "Good Bye") {
		t.Fatalf("Could not quit from server. Expected 'Good Bye' got %s", string(message))
	}

	if chatServer.IsUserConnected("User1") {
		t.Fatal("User is still connected to server after quiting")
	}
}
<|endoftext|>