<commit_before>package resources

import (
	"errors"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudformation"
	"github.com/aws/aws-sdk-go/service/cloudformation/cloudformationiface"
	"github.com/rebuy-de/aws-nuke/pkg/config"
	"github.com/rebuy-de/aws-nuke/pkg/types"
	"github.com/sirupsen/logrus"
)

const CLOUDFORMATION_MAX_DELETE_ATTEMPT = 3

func init() {
	register("CloudFormationStack", ListCloudFormationStacks)
}

func ListCloudFormationStacks(sess *session.Session) ([]Resource, error) {
	svc := cloudformation.New(sess)

	params := &cloudformation.DescribeStacksInput{}
	resources := make([]Resource, 0)

	for {
		resp, err := svc.DescribeStacks(params)
		if err != nil {
			return nil, err
		}
		for _, stack := range resp.Stacks {
			resources = append(resources, &CloudFormationStack{
				svc:               svc,
				stack:             stack,
				maxDeleteAttempts: CLOUDFORMATION_MAX_DELETE_ATTEMPT,
			})
		}

		if resp.NextToken == nil {
			break
		}

		params.NextToken = resp.NextToken
	}

	return resources, nil
}

type CloudFormationStack struct {
	svc               cloudformationiface.CloudFormationAPI
	stack             *cloudformation.Stack
	maxDeleteAttempts int
	featureFlags      config.FeatureFlags
}

func (cfs *CloudFormationStack) FeatureFlags(ff config.FeatureFlags) {
	cfs.featureFlags = ff
}

func (cfs *CloudFormationStack) Remove() error {
	return cfs.removeWithAttempts(0)
}

func (cfs *CloudFormationStack) removeWithAttempts(attempt int) error {
	if err := cfs.doRemove(); err != nil {
		logrus.Errorf("CloudFormationStack stackName=%s attempt=%d maxAttempts=%d delete failed: %s", *cfs.stack.StackName, attempt, cfs.maxDeleteAttempts, err.Error())
		awsErr, ok := err.(awserr.Error)
		if ok && awsErr.Code() == "ValidationError" &&
			awsErr.Message() == "Stack ["+*cfs.stack.StackName+"] cannot be deleted while TerminationProtection is enabled" {
			if cfs.featureFlags.DisableDeletionProtection.CloudformationStack {
				logrus.Infof("CloudFormationStack stackName=%s attempt=%d maxAttempts=%d updating termination protection", *cfs.stack.StackName, attempt, cfs.maxDeleteAttempts)
				_, err = cfs.svc.UpdateTerminationProtection(&cloudformation.UpdateTerminationProtectionInput{
					EnableTerminationProtection: aws.Bool(false),
					StackName:                   cfs.stack.StackName,
				})
				if err != nil {
					return err
				}
			} else {
				logrus.Warnf("CloudFormationStack stackName=%s attempt=%d maxAttempts=%d set feature flag to disable deletion protection", *cfs.stack.StackName, attempt, cfs.maxDeleteAttempts)
				return err
			}
		}
		if attempt >= cfs.maxDeleteAttempts {
			return errors.New("CFS might not be deleted after this run.")
		} else {
			return cfs.removeWithAttempts(attempt + 1)
		}
	} else {
		return nil
	}
}

func (cfs *CloudFormationStack) doRemove() error {
	o, err := cfs.svc.DescribeStacks(&cloudformation.DescribeStacksInput{
		StackName: cfs.stack.StackName,
	})
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ValidationFailed" && strings.HasSuffix(awsErr.Message(), " does not exist") {
				logrus.Infof("CloudFormationStack stackName=%s no longer exists", *cfs.stack.StackName)
				return nil
			}
		}
		return err
	}
	stack := o.Stacks[0]

	if *stack.StackStatus == cloudformation.StackStatusDeleteComplete {
		//stack already deleted, no need to re-delete
		return nil
	} else if *stack.StackStatus == cloudformation.StackStatusDeleteInProgress {
		logrus.Infof("CloudFormationStack stackName=%s delete in progress. Waiting", *cfs.stack.StackName)
		return cfs.svc.WaitUntilStackDeleteComplete(&cloudformation.DescribeStacksInput{
			StackName: cfs.stack.StackName,
		})
	} else if *stack.StackStatus == cloudformation.StackStatusDeleteFailed {
		logrus.Infof("CloudFormationStack stackName=%s delete failed. Attempting to retain and delete stack", *cfs.stack.StackName)
		// This means the CFS has undeleteable resources.
		// In order to move on with nuking, we retain them in the deletion.
		retainableResources, err := cfs.svc.ListStackResources(&cloudformation.ListStackResourcesInput{
			StackName: cfs.stack.StackName,
		})
		if err != nil {
			return err
		}

		retain := make([]*string, 0)

		for _, r := range retainableResources.StackResourceSummaries {
			if *r.ResourceStatus != cloudformation.ResourceStatusDeleteComplete {
				retain = append(retain, r.LogicalResourceId)
			}
		}

		_, err = cfs.svc.DeleteStack(&cloudformation.DeleteStackInput{
			StackName:       cfs.stack.StackName,
			RetainResources: retain,
		})
		if err != nil {
			return err
		}
		return cfs.svc.WaitUntilStackDeleteComplete(&cloudformation.DescribeStacksInput{
			StackName: cfs.stack.StackName,
		})
	} else {
		if err := cfs.waitForStackToStabilize(*stack.StackStatus); err != nil {
			return err
		} else if _, err := cfs.svc.DeleteStack(&cloudformation.DeleteStackInput{
			StackName: cfs.stack.StackName,
		}); err != nil {
			return err
		} else if err := cfs.svc.WaitUntilStackDeleteComplete(&cloudformation.DescribeStacksInput{
			StackName: cfs.stack.StackName,
		}); err != nil {
			return err
		} else {
			return nil
		}
	}
}
func (cfs *CloudFormationStack) waitForStackToStabilize(currentStatus string) error {
	switch currentStatus {
	case cloudformation.StackStatusUpdateInProgress:
		fallthrough
	case cloudformation.StackStatusUpdateRollbackCompleteCleanupInProgress:
		fallthrough
	case cloudformation.StackStatusUpdateRollbackInProgress:
		logrus.Infof("CloudFormationStack stackName=%s update in progress. Waiting to stabalize", *cfs.stack.StackName)
		return cfs.svc.WaitUntilStackUpdateComplete(&cloudformation.DescribeStacksInput{
			StackName: cfs.stack.StackName,
		})
	case cloudformation.StackStatusCreateInProgress:
		fallthrough
	case cloudformation.StackStatusRollbackInProgress:
		logrus.Infof("CloudFormationStack stackName=%s create in progress. Waiting to stabalize", *cfs.stack.StackName)
		return cfs.svc.WaitUntilStackCreateComplete(&cloudformation.DescribeStacksInput{
			StackName: cfs.stack.StackName,
		})
	default:
		return nil
	}
}

func (cfs *CloudFormationStack) Properties() types.Properties {
	properties := types.NewProperties()
	properties.Set("Name", cfs.stack.StackName)
	properties.Set("CreationTime", cfs.stack.CreationTime.Format(time.RFC3339))
	if cfs.stack.LastUpdatedTime == nil {
		properties.Set("LastUpdatedTime", cfs.stack.CreationTime.Format(time.RFC3339))
	} else {
		properties.Set("LastUpdatedTime", cfs.stack.LastUpdatedTime.Format(time.RFC3339))
	}

	for _, tagValue := range cfs.stack.Tags {
		properties.SetTag(tagValue.Key, tagValue.Value)
	}

	return properties
}

func (cfs *CloudFormationStack) String() string {
	return *cfs.stack.StackName
}
<commit_msg>fix: exclude nested cloudformation stacks (#761)<commit_after>package resources

import (
	"errors"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudformation"
	"github.com/aws/aws-sdk-go/service/cloudformation/cloudformationiface"
	"github.com/rebuy-de/aws-nuke/pkg/config"
	"github.com/rebuy-de/aws-nuke/pkg/types"
	"github.com/sirupsen/logrus"
)

const CLOUDFORMATION_MAX_DELETE_ATTEMPT = 3

func init() {
	register("CloudFormationStack", ListCloudFormationStacks)
}

func ListCloudFormationStacks(sess *session.Session) ([]Resource, error) {
	svc := cloudformation.New(sess)

	params := &cloudformation.DescribeStacksInput{}
	resources := make([]Resource, 0)

	for {
		resp, err := svc.DescribeStacks(params)
		if err != nil {
			return nil, err
		}
		for _, stack := range resp.Stacks {
			if aws.StringValue(stack.ParentId) != "" {
				continue
			}
			resources = append(resources, &CloudFormationStack{
				svc:               svc,
				stack:             stack,
				maxDeleteAttempts: CLOUDFORMATION_MAX_DELETE_ATTEMPT,
			})
		}

		if resp.NextToken == nil {
			break
		}

		params.NextToken = resp.NextToken
	}

	return resources, nil
}

type CloudFormationStack struct {
	svc               cloudformationiface.CloudFormationAPI
	stack             *cloudformation.Stack
	maxDeleteAttempts int
	featureFlags      config.FeatureFlags
}

func (cfs *CloudFormationStack) FeatureFlags(ff config.FeatureFlags) {
	cfs.featureFlags = ff
}

func (cfs *CloudFormationStack) Remove() error {
	return cfs.removeWithAttempts(0)
}

func (cfs *CloudFormationStack) removeWithAttempts(attempt int) error {
	if err := cfs.doRemove(); err != nil {
		logrus.Errorf("CloudFormationStack stackName=%s attempt=%d maxAttempts=%d delete failed: %s", *cfs.stack.StackName, attempt, cfs.maxDeleteAttempts, err.Error())
		awsErr, ok := err.(awserr.Error)
		if ok && awsErr.Code() == "ValidationError" &&
			awsErr.Message() == "Stack ["+*cfs.stack.StackName+"] cannot be deleted while TerminationProtection is enabled" {
			if cfs.featureFlags.DisableDeletionProtection.CloudformationStack {
				logrus.Infof("CloudFormationStack stackName=%s attempt=%d maxAttempts=%d updating termination protection", *cfs.stack.StackName, attempt, cfs.maxDeleteAttempts)
				_, err = cfs.svc.UpdateTerminationProtection(&cloudformation.UpdateTerminationProtectionInput{
					EnableTerminationProtection: aws.Bool(false),
					StackName:                   cfs.stack.StackName,
				})
				if err != nil {
					return err
				}
			} else {
				logrus.Warnf("CloudFormationStack stackName=%s attempt=%d maxAttempts=%d set feature flag to disable deletion protection", *cfs.stack.StackName, attempt, cfs.maxDeleteAttempts)
				return err
			}
		}
		if attempt >= cfs.maxDeleteAttempts {
			return errors.New("CFS might not be deleted after this run.")
		} else {
			return cfs.removeWithAttempts(attempt + 1)
		}
	} else {
		return nil
	}
}

func (cfs *CloudFormationStack) doRemove() error {
	o, err := cfs.svc.DescribeStacks(&cloudformation.DescribeStacksInput{
		StackName: cfs.stack.StackName,
	})
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ValidationFailed" && strings.HasSuffix(awsErr.Message(), " does not exist") {
				logrus.Infof("CloudFormationStack stackName=%s no longer exists", *cfs.stack.StackName)
				return nil
			}
		}
		return err
	}
	stack := o.Stacks[0]

	if *stack.StackStatus == cloudformation.StackStatusDeleteComplete {
		//stack already deleted, no need to re-delete
		return nil
	} else if *stack.StackStatus == cloudformation.StackStatusDeleteInProgress {
		logrus.Infof("CloudFormationStack stackName=%s delete in progress. Waiting", *cfs.stack.StackName)
		return cfs.svc.WaitUntilStackDeleteComplete(&cloudformation.DescribeStacksInput{
			StackName: cfs.stack.StackName,
		})
	} else if *stack.StackStatus == cloudformation.StackStatusDeleteFailed {
		logrus.Infof("CloudFormationStack stackName=%s delete failed. Attempting to retain and delete stack", *cfs.stack.StackName)
		// This means the CFS has undeleteable resources.
		// In order to move on with nuking, we retain them in the deletion.
		retainableResources, err := cfs.svc.ListStackResources(&cloudformation.ListStackResourcesInput{
			StackName: cfs.stack.StackName,
		})
		if err != nil {
			return err
		}

		retain := make([]*string, 0)

		for _, r := range retainableResources.StackResourceSummaries {
			if *r.ResourceStatus != cloudformation.ResourceStatusDeleteComplete {
				retain = append(retain, r.LogicalResourceId)
			}
		}

		_, err = cfs.svc.DeleteStack(&cloudformation.DeleteStackInput{
			StackName:       cfs.stack.StackName,
			RetainResources: retain,
		})
		if err != nil {
			return err
		}
		return cfs.svc.WaitUntilStackDeleteComplete(&cloudformation.DescribeStacksInput{
			StackName: cfs.stack.StackName,
		})
	} else {
		if err := cfs.waitForStackToStabilize(*stack.StackStatus); err != nil {
			return err
		} else if _, err := cfs.svc.DeleteStack(&cloudformation.DeleteStackInput{
			StackName: cfs.stack.StackName,
		}); err != nil {
			return err
		} else if err := cfs.svc.WaitUntilStackDeleteComplete(&cloudformation.DescribeStacksInput{
			StackName: cfs.stack.StackName,
		}); err != nil {
			return err
		} else {
			return nil
		}
	}
}
func (cfs *CloudFormationStack) waitForStackToStabilize(currentStatus string) error {
	switch currentStatus {
	case cloudformation.StackStatusUpdateInProgress:
		fallthrough
	case cloudformation.StackStatusUpdateRollbackCompleteCleanupInProgress:
		fallthrough
	case cloudformation.StackStatusUpdateRollbackInProgress:
		logrus.Infof("CloudFormationStack stackName=%s update in progress. Waiting to stabalize", *cfs.stack.StackName)
		return cfs.svc.WaitUntilStackUpdateComplete(&cloudformation.DescribeStacksInput{
			StackName: cfs.stack.StackName,
		})
	case cloudformation.StackStatusCreateInProgress:
		fallthrough
	case cloudformation.StackStatusRollbackInProgress:
		logrus.Infof("CloudFormationStack stackName=%s create in progress. Waiting to stabalize", *cfs.stack.StackName)
		return cfs.svc.WaitUntilStackCreateComplete(&cloudformation.DescribeStacksInput{
			StackName: cfs.stack.StackName,
		})
	default:
		return nil
	}
}

func (cfs *CloudFormationStack) Properties() types.Properties {
	properties := types.NewProperties()
	properties.Set("Name", cfs.stack.StackName)
	properties.Set("CreationTime", cfs.stack.CreationTime.Format(time.RFC3339))
	if cfs.stack.LastUpdatedTime == nil {
		properties.Set("LastUpdatedTime", cfs.stack.CreationTime.Format(time.RFC3339))
	} else {
		properties.Set("LastUpdatedTime", cfs.stack.LastUpdatedTime.Format(time.RFC3339))
	}

	for _, tagValue := range cfs.stack.Tags {
		properties.SetTag(tagValue.Key, tagValue.Value)
	}

	return properties
}

func (cfs *CloudFormationStack) String() string {
	return *cfs.stack.StackName
}
<|endoftext|>


<commit_before>package opts

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"reflect"
	"strings"
	"testing"
)

func tmpFileWithContent(content string, t *testing.T) string {
	tmpFile, err := ioutil.TempFile("", "envfile-test")
	if err != nil {
		t.Fatal(err)
	}
	defer tmpFile.Close()

	tmpFile.WriteString(content)
	return tmpFile.Name()
}

// Test ParseEnvFile for a file with a few well formatted lines
func TestParseEnvFileGoodFile(t *testing.T) {
	content := `foo=bar
    baz=quux
# comment

_foobar=foobaz
with.dots=working
and_underscore=working too
`
	// Adding a newline + a line with pure whitespace.
	// This is being done like this instead of the block above
	// because it's common for editors to trim trailing whitespace
	// from lines, which becomes annoying since that's the
	// exact thing we need to test.
	content += "\n    \t  "
	tmpFile := tmpFileWithContent(content, t)
	defer os.Remove(tmpFile)

	lines, err := ParseEnvFile(tmpFile)
	if err != nil {
		t.Fatal(err)
	}

	expectedLines := []string{
		"foo=bar",
		"baz=quux",
		"_foobar=foobaz",
		"with.dots=working",
		"and_underscore=working too",
	}

	if !reflect.DeepEqual(lines, expectedLines) {
		t.Fatal("lines not equal to expectedLines")
	}
}

// Test ParseEnvFile for an empty file
func TestParseEnvFileEmptyFile(t *testing.T) {
	tmpFile := tmpFileWithContent("", t)
	defer os.Remove(tmpFile)

	lines, err := ParseEnvFile(tmpFile)
	if err != nil {
		t.Fatal(err)
	}

	if len(lines) != 0 {
		t.Fatal("lines not empty; expected empty")
	}
}

// Test ParseEnvFile for a non existent file
func TestParseEnvFileNonExistentFile(t *testing.T) {
	_, err := ParseEnvFile("foo_bar_baz")
	if err == nil {
		t.Fatal("ParseEnvFile succeeded; expected failure")
	}
	if _, ok := err.(*os.PathError); !ok {
		t.Fatalf("Expected a PathError, got [%v]", err)
	}
}

// Test ParseEnvFile for a badly formatted file
func TestParseEnvFileBadlyFormattedFile(t *testing.T) {
	content := `foo=bar
    f   =quux
`

	tmpFile := tmpFileWithContent(content, t)
	defer os.Remove(tmpFile)

	_, err := ParseEnvFile(tmpFile)
	if err == nil {
		t.Fatalf("Expected an ErrBadKey, got nothing")
	}
	if _, ok := err.(ErrBadKey); !ok {
		t.Fatalf("Expected an ErrBadKey, got [%v]", err)
	}
	expectedMessage := "poorly formatted environment: variable 'f   ' has white spaces"
	if err.Error() != expectedMessage {
		t.Fatalf("Expected [%v], got [%v]", expectedMessage, err.Error())
	}
}

// Test ParseEnvFile for a file with a line exceeding bufio.MaxScanTokenSize
func TestParseEnvFileLineTooLongFile(t *testing.T) {
	content := strings.Repeat("a", bufio.MaxScanTokenSize+42)
	content = fmt.Sprint("foo=", content)

	tmpFile := tmpFileWithContent(content, t)
	defer os.Remove(tmpFile)

	_, err := ParseEnvFile(tmpFile)
	if err == nil {
		t.Fatal("ParseEnvFile succeeded; expected failure")
	}
}

// ParseEnvFile with a random file, pass through
func TestParseEnvFileRandomFile(t *testing.T) {
	content := `first line
another invalid line`
	tmpFile := tmpFileWithContent(content, t)
	defer os.Remove(tmpFile)

	_, err := ParseEnvFile(tmpFile)
	if err == nil {
		t.Fatalf("Expected an ErrBadKey, got nothing")
	}
	if _, ok := err.(ErrBadKey); !ok {
		t.Fatalf("Expected an ErrBadKey, got [%v]", err)
	}
	expectedMessage := "poorly formatted environment: variable 'first line' has white spaces"
	if err.Error() != expectedMessage {
		t.Fatalf("Expected [%v], got [%v]", expectedMessage, err.Error())
	}
}
<commit_msg>add test for undefined variable environment file import<commit_after>package opts

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"os"
	"reflect"
	"strings"
	"testing"
)

func tmpFileWithContent(content string, t *testing.T) string {
	tmpFile, err := ioutil.TempFile("", "envfile-test")
	if err != nil {
		t.Fatal(err)
	}
	defer tmpFile.Close()

	tmpFile.WriteString(content)
	return tmpFile.Name()
}

// Test ParseEnvFile for a file with a few well formatted lines
func TestParseEnvFileGoodFile(t *testing.T) {
	content := `foo=bar
    baz=quux
# comment

_foobar=foobaz
with.dots=working
and_underscore=working too
`
	// Adding a newline + a line with pure whitespace.
	// This is being done like this instead of the block above
	// because it's common for editors to trim trailing whitespace
	// from lines, which becomes annoying since that's the
	// exact thing we need to test.
	content += "\n    \t  "
	tmpFile := tmpFileWithContent(content, t)
	defer os.Remove(tmpFile)

	lines, err := ParseEnvFile(tmpFile)
	if err != nil {
		t.Fatal(err)
	}

	expectedLines := []string{
		"foo=bar",
		"baz=quux",
		"_foobar=foobaz",
		"with.dots=working",
		"and_underscore=working too",
	}

	if !reflect.DeepEqual(lines, expectedLines) {
		t.Fatal("lines not equal to expectedLines")
	}
}

// Test ParseEnvFile for an empty file
func TestParseEnvFileEmptyFile(t *testing.T) {
	tmpFile := tmpFileWithContent("", t)
	defer os.Remove(tmpFile)

	lines, err := ParseEnvFile(tmpFile)
	if err != nil {
		t.Fatal(err)
	}

	if len(lines) != 0 {
		t.Fatal("lines not empty; expected empty")
	}
}

// Test ParseEnvFile for a non existent file
func TestParseEnvFileNonExistentFile(t *testing.T) {
	_, err := ParseEnvFile("foo_bar_baz")
	if err == nil {
		t.Fatal("ParseEnvFile succeeded; expected failure")
	}
	if _, ok := err.(*os.PathError); !ok {
		t.Fatalf("Expected a PathError, got [%v]", err)
	}
}

// Test ParseEnvFile for a badly formatted file
func TestParseEnvFileBadlyFormattedFile(t *testing.T) {
	content := `foo=bar
    f   =quux
`

	tmpFile := tmpFileWithContent(content, t)
	defer os.Remove(tmpFile)

	_, err := ParseEnvFile(tmpFile)
	if err == nil {
		t.Fatalf("Expected an ErrBadKey, got nothing")
	}
	if _, ok := err.(ErrBadKey); !ok {
		t.Fatalf("Expected an ErrBadKey, got [%v]", err)
	}
	expectedMessage := "poorly formatted environment: variable 'f   ' has white spaces"
	if err.Error() != expectedMessage {
		t.Fatalf("Expected [%v], got [%v]", expectedMessage, err.Error())
	}
}

// Test ParseEnvFile for a file with a line exceeding bufio.MaxScanTokenSize
func TestParseEnvFileLineTooLongFile(t *testing.T) {
	content := strings.Repeat("a", bufio.MaxScanTokenSize+42)
	content = fmt.Sprint("foo=", content)

	tmpFile := tmpFileWithContent(content, t)
	defer os.Remove(tmpFile)

	_, err := ParseEnvFile(tmpFile)
	if err == nil {
		t.Fatal("ParseEnvFile succeeded; expected failure")
	}
}

// ParseEnvFile with a random file, pass through
func TestParseEnvFileRandomFile(t *testing.T) {
	content := `first line
another invalid line`
	tmpFile := tmpFileWithContent(content, t)
	defer os.Remove(tmpFile)

	_, err := ParseEnvFile(tmpFile)
	if err == nil {
		t.Fatalf("Expected an ErrBadKey, got nothing")
	}
	if _, ok := err.(ErrBadKey); !ok {
		t.Fatalf("Expected an ErrBadKey, got [%v]", err)
	}
	expectedMessage := "poorly formatted environment: variable 'first line' has white spaces"
	if err.Error() != expectedMessage {
		t.Fatalf("Expected [%v], got [%v]", expectedMessage, err.Error())
	}
}

// ParseEnvFile with environment variable import definitions
func TestParseEnvVariableDefinitionsFile(t *testing.T) {
	content := `# comment=
UNDEFINED_VAR
HOME
`
	tmpFile := tmpFileWithContent(content, t)
	defer os.Remove(tmpFile)

	variables, err := ParseEnvFile(tmpFile)
	if nil != err {
		t.Fatal("There must not be any error")
	}

	if "HOME="+os.Getenv("HOME") != variables[0] {
		t.Fatal("the HOME variable is not properly imported as the first variable (but it is the only one to import)")
	}

	if 1 != len(variables) {
		t.Fatal("exactly one variable is imported (as the other one is not set at all)")
	}
}
<|endoftext|>


<commit_before>package seccomp

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/subgraph/oz"
	"golang.org/x/sys/unix"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"strconv"
	"syscall"
	"unsafe"
)

const (
	STRINGARG = iota + 1
	PTRARG
	INTARG
)

const MAXARGS int = 6
const MAXSTRING int = 100

type SystemCallArgs [MAXARGS]int

type SystemCall struct {
	prefix string
	name   string
	num    int
	args   SystemCallArgs
}

func Tracer() {

	p := new(oz.Profile)
	if err := json.NewDecoder(os.Stdin).Decode(&p); err != nil {
		log.Error("unable to decode profile data: %v", err)
		os.Exit(1)
	}

	var proc_attr syscall.ProcAttr
	var sys_attr syscall.SysProcAttr

	sys_attr.Ptrace = true
	done := false
	proc_attr.Sys = &sys_attr

	cmd := os.Args[1]
	cmdArgs := os.Args[2:]
	log.Info("Tracer running command (%v) arguments (%v)\n", cmd, cmdArgs)
	c := exec.Command(cmd)
	c.SysProcAttr = &syscall.SysProcAttr{Ptrace: true}
	c.Env = os.Environ()
	c.Args = append(c.Args, cmdArgs...)

	pi, err := c.StdinPipe()
	if err != nil {
		fmt.Errorf("error creating stdin pipe for tracer process: %v", err)
		os.Exit(1)
	}
	jdata, err := json.Marshal(p)
	if err != nil {
		fmt.Errorf("Unable to marshal seccomp state: %+v", err)
		os.Exit(1)
	}
	io.Copy(pi, bytes.NewBuffer(jdata))
	log.Info(string(jdata))
	pi.Close()

	children := make(map[int]bool)

	if err := c.Start(); err == nil {
		children[c.Process.Pid] = true
		var s syscall.WaitStatus
		pid, err := syscall.Wait4(-1, &s, syscall.WALL, nil)
		children[pid] = true
		if err != nil {
			log.Error("Error (wait4): %v", err)
		}
		log.Info("Tracing child pid: %v\n", pid)
		for done == false {
			syscall.PtraceSetOptions(pid, unix.PTRACE_O_TRACESECCOMP|unix.PTRACE_O_TRACEFORK|unix.PTRACE_O_TRACEVFORK|unix.PTRACE_O_TRACECLONE|unix.PTRACE_O_TRACEEXIT)
			syscall.PtraceCont(pid, 0)
			pid, err = syscall.Wait4(-1, &s, syscall.WALL, nil)
			if err != nil {
				log.Error("Error (wait4): %v\n", err)
				if len(children) == 0 {
					done = true
				}
				continue
			}
			children[pid] = true
			if s.Exited() == true {
				delete(children, pid)
				log.Info("Child pid %v finished.\n", pid)
				if len(children) == 0 {
					done = true
				}
				continue
			}
			if uint32(s)>>8 == (uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_SECCOMP << 8)) {
				if err != nil {
					log.Error("Error (ptrace): %v", err)
					continue
				}
				var regs syscall.PtraceRegs
				err = syscall.PtraceGetRegs(pid, &regs)
				if err != nil {
					log.Error("Error (ptrace): %v", err)
				}
				systemcall, err := syscallByNum(int(regs.Orig_rax))
				if err != nil {
					log.Error("Error: %v", err)
					continue
				}
				var callrep string = fmt.Sprintf("%s(", systemcall.name)
				var reg uint64 = 0

				for arg := range systemcall.args {

					if systemcall.args[arg] == 0 {
						break
					}

					if arg > 0 {
						callrep += fmt.Sprintf(",")
					}

					switch arg {
					case 0:
						reg = regs.Rdi
					case 1:
						reg = regs.Rsi
					case 2:
						reg = regs.Rdx
					case 3:
						reg = regs.Rcx
					case 4:
						reg = regs.R8
					case 5:
						reg = regs.R9
					}
					if systemcall.args[arg] == STRINGARG {
						str, err := readStringArg(pid, uintptr(reg))
						if err != nil {
							log.Error("Error: %v", err)
						} else {
							callrep += fmt.Sprintf("\"%s\"", str)
						}
					} else if systemcall.args[arg] == INTARG {
						callrep += fmt.Sprintf("%d", uint64(reg))
					} else {
						/* Stringify pointers in writes to stdout/stderr */
						write, err := syscallByName("write")
						if err != nil {
							log.Error("Error: %v", err)
						}
						if systemcall.num == write.num && (regs.Rdi == uint64(syscall.Stdout) || regs.Rdi == uint64(syscall.Stderr)) {
							str, err := readStringArg(pid, uintptr(reg))
							if err != nil {
								log.Error("Error %v", err)
							} else {
								if isPrintableASCII(str) == true {
									callrep += fmt.Sprintf("\"%s\"", str)
								} else {
									callrep += fmt.Sprintf("0x%X", uintptr(reg))
								}
							}
						} else {
							callrep += fmt.Sprintf("0x%X", uintptr(reg))
						}
					}

				}
				callrep += ")"
				log.Info("==============================================\nseccomp hit on sandbox pid %v (%v) syscall %v (%v): \n\n%s\nI ==============================================\n\n", pid, getProcessCmdLine(pid), systemcall.name, regs.Orig_rax, callrep)
			}
		}
	} else {
		log.Error("Error: %v", err)
	}

}

func readStringArg(pid int, addr uintptr) (s string, err error) {
	buf := make([]byte, unsafe.Sizeof(addr))
	done := false
	err = nil
	for done == false {
		_, err := syscall.PtracePeekText(pid, addr, buf)
		if err != nil {
			fmt.Printf("Error (ptrace): %v\n", err)
		} else {
			for b := range buf {
				if buf[b] == 0 {
					done = true
					break
				} else {
					s += string(buf[b])
					if len(s) > 90 {
						s += "..."
						break
					}
				}
			}
		}
		addr += unsafe.Sizeof(addr)
	}
	return s, nil
}

func readUIntArg(pid int, addr uintptr) (uint64, error) {
	buf := make([]byte, unsafe.Sizeof(addr))

	_, err := syscall.PtracePeekText(pid, addr, buf)
	if err != nil {
		fmt.Printf("Error (ptrace): %v\n", err)
		return 0, err
	} else {
		i := binary.LittleEndian.Uint64(buf)
		return i, nil
	}
	return 0, errors.New("Error.")
}

func readPtrArg(pid int, addr uintptr) (uintptr, error) {

	buf := make([]byte, unsafe.Sizeof(addr))

	_, err := syscall.PtracePeekText(pid, addr, buf)
	if err != nil {
		fmt.Printf("Error (ptrace): %v\n", err)
		return 0, err
	} else {
		i := binary.LittleEndian.Uint64(buf)
		return uintptr(i), nil
	}
	return 0, nil
}

func syscallByNum(num int) (s SystemCall, err error) {
	var q SystemCall = SystemCall{"", "", -1, [6]int{0, 0, 0, 0, 0, 0}}
	for i := range syscalls {
		if syscalls[i].num == num {
			q = syscalls[i]
			return q, nil
		}
	}
	return q, errors.New("System call not found.\n")
}

func syscallByName(name string) (s SystemCall, err error) {
	var q SystemCall = SystemCall{"", "", -1, [6]int{0, 0, 0, 0, 0, 0}}
	for i := range syscalls {
		if syscalls[i].name == name {
			q = syscalls[i]
			return q, nil
		}
	}
	return q, errors.New("System call not found.\n")
}

func isPrintableASCII(s string) bool {
	for _, x := range s {
		if x > 127 {
			return false
		}
	}
	return true
}

func getProcessCmdLine(pid int) string {
	path := "/proc/" + strconv.Itoa(pid) + "/cmdline"
	cmdline, err := ioutil.ReadFile(path)
	for b := range cmdline {
		if b < (len(cmdline) - 1) {
			if cmdline[b] == 0x00 {
				cmdline[b] = 0x20
			}
		}
	}
	if err != nil {
		log.Error("Error (read): %v", err)
		return "unknown"
	}
	return string(cmdline)
}
<commit_msg>Added more cases to handle stopped children<commit_after>package seccomp

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/subgraph/oz"
	"golang.org/x/sys/unix"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"strconv"
	"syscall"
	"unsafe"
)

const (
	STRINGARG = iota + 1
	PTRARG
	INTARG
)

type SystemCallArgs []int

func Tracer() {

	p := new(oz.Profile)
	if err := json.NewDecoder(os.Stdin).Decode(&p); err != nil {
		log.Error("unable to decode profile data: %v", err)
		os.Exit(1)
	}

	var proc_attr syscall.ProcAttr
	var sys_attr syscall.SysProcAttr

	sys_attr.Ptrace = true
	done := false
	proc_attr.Sys = &sys_attr

	cmd := os.Args[1]
	cmdArgs := os.Args[2:]
	log.Info("Tracer running command (%v) arguments (%v)\n", cmd, cmdArgs)
	c := exec.Command(cmd)
	c.SysProcAttr = &syscall.SysProcAttr{Ptrace: true}
	c.Env = os.Environ()
	c.Args = append(c.Args, cmdArgs...)

	pi, err := c.StdinPipe()
	if err != nil {
		fmt.Errorf("error creating stdin pipe for tracer process: %v", err)
		os.Exit(1)
	}
	jdata, err := json.Marshal(p)
	if err != nil {
		fmt.Errorf("Unable to marshal seccomp state: %+v", err)
		os.Exit(1)
	}
	io.Copy(pi, bytes.NewBuffer(jdata))
	log.Info(string(jdata))
	pi.Close()

	children := make(map[int]bool)

	if err := c.Start(); err == nil {
		children[c.Process.Pid] = true
		var s syscall.WaitStatus
		pid, err := syscall.Wait4(-1, &s, syscall.WALL, nil)
		children[pid] = true
		if err != nil {
			log.Error("Error (wait4) here first: %v %i", err, pid)
		}
		log.Info("Tracing child pid: %v\n", pid)
		for done == false {
			syscall.PtraceSetOptions(pid, unix.PTRACE_O_TRACESECCOMP|unix.PTRACE_O_TRACEFORK|unix.PTRACE_O_TRACEVFORK|unix.PTRACE_O_TRACECLONE)
			syscall.PtraceCont(pid, 0)
			pid, err = syscall.Wait4(-1, &s, syscall.WALL, nil)
			if err != nil {
				log.Error("Error (wait4) here: %v %i %v\n", err, pid, children)
				if len(children) == 0 {
					done = true
				}
				continue
			}
			children[pid] = true
			if s.Exited() == true {
				delete(children, pid)
				log.Info("Child pid %v finished.\n", pid)
				if len(children) == 0 {
					done = true
				}
				continue
			}
			if uint32(s)>>8 == (uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_SECCOMP << 8)) {
				if err != nil {
					log.Error("Error (ptrace): %v", err)
					continue
				}
				var regs syscall.PtraceRegs
				err = syscall.PtraceGetRegs(pid, &regs)
				if err != nil {
					log.Error("Error (ptrace): %v", err)
				}
				systemcall, err := syscallByNum(int(regs.Orig_rax))
				if err != nil {
					log.Error("Error: %v", err)
					continue
				}
				var callrep string = fmt.Sprintf("%s(", systemcall.name)
				var reg uint64 = 0

				for arg := range systemcall.args {

					if systemcall.args[arg] == 0 {
						break
					}

					if arg > 0 {
						callrep += fmt.Sprintf(",")
					}

					switch arg {
					case 0:
						reg = regs.Rdi
					case 1:
						reg = regs.Rsi
					case 2:
						reg = regs.Rdx
					case 3:
						reg = regs.Rcx
					case 4:
						reg = regs.R8
					case 5:
						reg = regs.R9
					}
					if systemcall.args[arg] == STRINGARG {
						str, err := readStringArg(pid, uintptr(reg))
						if err != nil {
							log.Error("Error: %v", err)
						} else {
							callrep += fmt.Sprintf("\"%s\"", str)
						}
					} else if systemcall.args[arg] == INTARG {
						callrep += fmt.Sprintf("%d", uint64(reg))
					} else {
						/* Stringify pointers in writes to stdout/stderr */
						write, err := syscallByName("write")
						if err != nil {
							log.Error("Error: %v", err)
						}
						if systemcall.num == write.num && (regs.Rdi == uint64(syscall.Stdout) || regs.Rdi == uint64(syscall.Stderr)) {
							str, err := readStringArg(pid, uintptr(reg))
							if err != nil {
								log.Error("Error %v", err)
							} else {
								if isPrintableASCII(str) == true {
									callrep += fmt.Sprintf("\"%s\"", str)
								} else {
									callrep += fmt.Sprintf("0x%X", uintptr(reg))
								}
							}
						} else {
							callrep += fmt.Sprintf("0x%X", uintptr(reg))
						}
					}

				}
				callrep += ")"
				log.Info("==============================================\nseccomp hit on sandbox pid %v (%v) syscall %v (%v): \n\n%s\nI ==============================================\n\n", pid, getProcessCmdLine(pid), systemcall.name, regs.Orig_rax, callrep)
			} else {
				if s.Signaled() == true {
					log.Error("Other pid signalled %v %v", pid, s)
					delete(children, pid)
					continue
				} else {

					switch uint32(s) >> 8 {

					case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_EXIT << 8):
						log.Error("Ptrace exit event detected pid %v (%s)", pid, getProcessCmdLine(pid))

					case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_CLONE << 8):
						log.Error("Ptrace clone event detected pid %v (%s)", pid, getProcessCmdLine(pid))
						continue
					case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_FORK << 8):
						log.Error("PTrace fork event detected pid %v (%s)", pid, getProcessCmdLine(pid))
						continue
					case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_VFORK << 8):
						log.Error("Ptrace vfork event detected pid %v (%s)", pid, getProcessCmdLine(pid))
						continue
					case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_VFORK_DONE << 8):
						log.Error("Ptrace vfork done event detected pid %v (%s)", pid, getProcessCmdLine(pid))
						continue
					case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_EXEC << 8):
						log.Error("Ptrace exec event detected pid %v (%s)", pid, getProcessCmdLine(pid))
						continue
					case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_STOP << 8):
						log.Error("Ptrace stop event detected pid %v (%s)", pid, getProcessCmdLine(pid))
						continue
					case uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_SECCOMP << 8):
						log.Error("Ptrace seccomp event detected pid %v (%s)", pid, getProcessCmdLine(pid))
						continue
					case uint32(unix.SIGTRAP):
						log.Error("SIGTRAP detected in pid %v (%s)", pid, getProcessCmdLine(pid))
						continue
					case uint32(unix.SIGCHLD):
						log.Error("SIGCHLD detected pid %v (%s)", pid, getProcessCmdLine(pid))
						continue
					case uint32(unix.SIGSTOP):
						log.Error("SIGSTOP detected pid %v (%s)", pid, getProcessCmdLine(pid))
						continue
					default:
						y := s.StopSignal()
						log.Error("Child stopped for unknown reasons pid %v status %v signal %i (%s)", pid, s, y, getProcessCmdLine(pid))
						continue
					}

					/*
						  if uint32(s)>>8 == (uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_EXIT << 8)) {
							log.Error("Ptrace exit event detected. %v %v", pid, s)
						  }
						  if (uint32(s)>>8 == (uint32(unix.SIGTRAP) | (unix.PTRACE_EVENT_CLONE << 8))) || (uint32(s)>>8 == (uint32(unix.SIGTRAP) | unix.PTRACE_EVENT_FORK << 8)) {
							  log.Error("Clone event %v %v", pid, s)
							  continue
						  }
						  if s.Stopped() == true {
							log.Error("Stopped wtf %v %v", pid, s)
						  }
						  if s.CoreDump() == true {
							  log.Error("Core dumped wtf %v %v", pid, s)
						  }
						  log.Error("Not signalled, wtf %v %v", pid, s)
					*/
				}
			}
		}
	} else {
		log.Error("Error: %v", err)
	}

}

func readStringArg(pid int, addr uintptr) (s string, err error) {
	buf := make([]byte, unsafe.Sizeof(addr))
	done := false
	err = nil
	for done == false {
		_, err := syscall.PtracePeekText(pid, addr, buf)
		if err != nil {
			fmt.Printf("Error (ptrace): %v\n", err)
		} else {
			for b := range buf {
				if buf[b] == 0 {
					done = true
					break
				} else {
					s += string(buf[b])
					if len(s) > 90 {
						s += "..."
						break
					}
				}
			}
		}
		addr += unsafe.Sizeof(addr)
	}
	return s, nil
}

func readUIntArg(pid int, addr uintptr) (uint64, error) {
	buf := make([]byte, unsafe.Sizeof(addr))

	_, err := syscall.PtracePeekText(pid, addr, buf)
	if err != nil {
		fmt.Printf("Error (ptrace): %v\n", err)
		return 0, err
	} else {
		i := binary.LittleEndian.Uint64(buf)
		return i, nil
	}
	return 0, errors.New("Error.")
}

func readPtrArg(pid int, addr uintptr) (uintptr, error) {

	buf := make([]byte, unsafe.Sizeof(addr))

	_, err := syscall.PtracePeekText(pid, addr, buf)
	if err != nil {
		fmt.Printf("Error (ptrace): %v\n", err)
		return 0, err
	} else {
		i := binary.LittleEndian.Uint64(buf)
		return uintptr(i), nil
	}
	return 0, nil
}

func syscallByNum(num int) (s SystemCall, err error) {
	var q SystemCall = SystemCall{"", "", -1, []int{0, 0, 0, 0, 0, 0}}
	for i := range syscalls {
		if syscalls[i].num == num {
			q = syscalls[i]
			return q, nil
		}
	}
	return q, errors.New("System call not found.\n")
}

func syscallByName(name string) (s SystemCall, err error) {
	var q SystemCall = SystemCall{"", "", -1, []int{0, 0, 0, 0, 0, 0}}
	for i := range syscalls {
		if syscalls[i].name == name {
			q = syscalls[i]
			return q, nil
		}
	}
	return q, errors.New("System call not found.\n")
}

func isPrintableASCII(s string) bool {
	for _, x := range s {
		if x > 127 {
			return false
		}
	}
	return true
}

func getProcessCmdLine(pid int) string {
	path := "/proc/" + strconv.Itoa(pid) + "/cmdline"
	cmdline, err := ioutil.ReadFile(path)
	for b := range cmdline {
		if b <= (len(cmdline) - 1) {
			if cmdline[b] == 0x00 {
				cmdline[b] = 0x20
			}
		}
	}
	if err != nil {
		log.Error("Error (read): %v", err)
		return "unknown"
	}
	return string(cmdline)
}
<|endoftext|>


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

package security

import (
	"reflect"
	"time"

	"v.io/v23"
	"v.io/v23/context"
	"v.io/v23/security"
	"v.io/v23/vdl"
)

// DischargeRefreshFraction determines how early before their expiration
// time we refresh discharges.  A value of 0.5 means we refresh when it
// is only half way to is expiration time.
const DischargeRefreshFraction = 0.5

// If this is attached to the context, we will not fetch discharges.
// We use this to prevent ourselves from fetching discharges in the
// process of fetching discharges, thus creating an infinite loop.
type skipDischargesKey struct{}

type updateResult struct {
	refreshTime time.Time
	discharge   security.Discharge
	done        bool
}

type work struct {
	caveat  security.Caveat
	impetus security.DischargeImpetus
}

// PrepareDischarges retrieves the caveat discharges required for using blessings
// at server. The discharges are either found in the dischargeCache, in the call
// options, or requested from the discharge issuer indicated on the caveat.
// Note that requesting a discharge is an rpc call, so one copy of this
// function must be able to successfully terminate while another is blocked.
// PrepareDischarges also returns a refreshTime, which is the time at which
// PrepareDischarges should be called again (or zero if none of the discharges
// expire).
func PrepareDischarges(
	ctx *context.T,
	blessings security.Blessings,
	serverBlessings []string,
	method string,
	args []interface{}) (map[string]security.Discharge, time.Time) {
	tpCavs := blessings.ThirdPartyCaveats()
	if len(tpCavs) == 0 {
		return nil, time.Time{}
	}
	impetus := mkDischargeImpetus(serverBlessings, method, args)
	// We only want to send the impetus information we really need for each
	// discharge.
	todo := make(map[string]work, len(tpCavs))
	for _, cav := range tpCavs {
		if tp := cav.ThirdPartyDetails(); tp != nil {
			todo[tp.ID()] = work{cav, filteredImpetus(tp.Requirements(), impetus)}
		}
	}
	// Since there may be dependencies in the caveats, we keep retrying
	// until either all discharges can be fetched or no new discharges
	// are fetched.
	var minRefreshTime time.Time
	ch := make(chan *updateResult, len(tpCavs))
	ret := make(map[string]security.Discharge, len(tpCavs))
	for {
		want := len(todo)
		now := time.Now()
		for _, w := range todo {
			updateDischarge(ctx, now, w.impetus, w.caveat, ch)
		}
		got := 0
		for i := 0; i < want; i++ {
			res := <-ch
			id := res.discharge.ID()
			ret[id] = res.discharge
			if res.done {
				minRefreshTime = minTime(minRefreshTime, res.refreshTime)
				delete(todo, id)
				got++
			}
		}
		if got == want {
			return ret, minRefreshTime
		}
		if got == 0 {
			return ret, minTime(minRefreshTime, time.Now().Add(time.Minute))
		}
	}
}

func updateDischarge(
	ctx *context.T,
	now time.Time,
	impetus security.DischargeImpetus,
	caveat security.Caveat,
	out chan<- *updateResult) {
	bstore := v23.GetPrincipal(ctx).BlessingStore()
	dis, ct := bstore.Discharge(caveat, impetus)
	if skip, _ := ctx.Value(skipDischargesKey{}).(bool); skip {
		// We can't fetch discharges while making a call to fetch a discharge.
		// Just go with what we have in the cache.
		out <- &updateResult{done: true, discharge: dis}
		return
	}
	if rt := refreshTime(dis, ct); dis.ID() != "" && (now.Before(rt) || rt.IsZero()) {
		// The cached value is still fresh, just return it.
		out <- &updateResult{done: true, discharge: dis, refreshTime: rt}
		return
	}
	// discharge in blessing store either doesn't exist or may be stale,
	// refresh via RPC.
	go func() {
		tp := caveat.ThirdPartyDetails()
		ctx = context.WithValue(ctx, skipDischargesKey{}, true)
		var newDis security.Discharge
		args, res := []interface{}{caveat, impetus}, []interface{}{&newDis}
		ctx.VI(3).Infof("Fetching discharge for %v", tp)
		if err := v23.GetClient(ctx).Call(ctx, tp.Location(), "Discharge", args, res); err != nil {
			ctx.VI(3).Infof("Discharge fetch for %v failed: %v", tp, err)
			out <- &updateResult{discharge: dis}
			return
		}
		bstore.CacheDischarge(newDis, caveat, impetus)
		out <- &updateResult{done: true, discharge: newDis, refreshTime: refreshTime(newDis, time.Now())}
	}()
}

// filteredImpetus returns a copy of 'before' after removing any values that are not required as per 'r'.
func filteredImpetus(r security.ThirdPartyRequirements, before security.DischargeImpetus) (after security.DischargeImpetus) {
	if r.ReportServer && len(before.Server) > 0 {
		after.Server = make([]security.BlessingPattern, len(before.Server))
		for i := range before.Server {
			after.Server[i] = before.Server[i]
		}
	}
	if r.ReportMethod {
		after.Method = before.Method
	}
	if r.ReportArguments && len(before.Arguments) > 0 {
		after.Arguments = make([]*vdl.Value, len(before.Arguments))
		for i := range before.Arguments {
			after.Arguments[i] = vdl.CopyValue(before.Arguments[i])
		}
	}
	return
}

func refreshTime(dis security.Discharge, cacheTime time.Time) time.Time {
	expiry := dis.Expiry()
	if expiry.IsZero() {
		return time.Time{}
	}
	if cacheTime.IsZero() {
		// If we don't know the cache time, just try to refresh within a minute of
		// the expiry time.
		return expiry.Add(-time.Minute)
	}
	lifetime := expiry.Sub(cacheTime)
	return expiry.Add(-time.Duration(float64(lifetime) * DischargeRefreshFraction))
}

func minTime(a, b time.Time) time.Time {
	if a.IsZero() || (!b.IsZero() && b.Before(a)) {
		return b
	}
	return a
}

func mkDischargeImpetus(serverBlessings []string, method string, args []interface{}) security.DischargeImpetus {
	var impetus security.DischargeImpetus
	if len(serverBlessings) > 0 {
		impetus.Server = make([]security.BlessingPattern, len(serverBlessings))
		for i, b := range serverBlessings {
			impetus.Server[i] = security.BlessingPattern(b)
		}
	}
	impetus.Method = method
	if len(args) > 0 {
		impetus.Arguments = make([]*vdl.Value, len(args))
		for i, a := range args {
			vArg, err := vdl.ValueFromReflect(reflect.ValueOf(a))
			if err != nil {
				continue
			}
			impetus.Arguments[i] = vArg
		}
	}
	return impetus
}
<commit_msg>ref: Construct filtered impetus directly instead of creating impetus and then filtering.<commit_after>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package security

import (
	"reflect"
	"time"

	"v.io/v23"
	"v.io/v23/context"
	"v.io/v23/security"
	"v.io/v23/vdl"
)

// DischargeRefreshFraction determines how early before their expiration
// time we refresh discharges.  A value of 0.5 means we refresh when it
// is only half way to is expiration time.
const DischargeRefreshFraction = 0.5

// If this is attached to the context, we will not fetch discharges.
// We use this to prevent ourselves from fetching discharges in the
// process of fetching discharges, thus creating an infinite loop.
type skipDischargesKey struct{}

type updateResult struct {
	refreshTime time.Time
	discharge   security.Discharge
	done        bool
}

type work struct {
	caveat  security.Caveat
	impetus security.DischargeImpetus
}

// PrepareDischarges retrieves the caveat discharges required for using blessings
// at server. The discharges are either found in the dischargeCache, in the call
// options, or requested from the discharge issuer indicated on the caveat.
// Note that requesting a discharge is an rpc call, so one copy of this
// function must be able to successfully terminate while another is blocked.
// PrepareDischarges also returns a refreshTime, which is the time at which
// PrepareDischarges should be called again (or zero if none of the discharges
// expire).
func PrepareDischarges(
	ctx *context.T,
	blessings security.Blessings,
	serverBlessings []string,
	method string,
	args []interface{}) (map[string]security.Discharge, time.Time) {
	tpCavs := blessings.ThirdPartyCaveats()
	if len(tpCavs) == 0 {
		return nil, time.Time{}
	}
	// We only want to send the impetus information we really need for each
	// discharge.
	todo := make(map[string]work, len(tpCavs))
	for _, cav := range tpCavs {
		if tp := cav.ThirdPartyDetails(); tp != nil {
			todo[tp.ID()] = work{cav, filteredImpetus(tp.Requirements(), serverBlessings, method, args)}
		}
	}
	// Since there may be dependencies in the caveats, we keep retrying
	// until either all discharges can be fetched or no new discharges
	// are fetched.
	var minRefreshTime time.Time
	ch := make(chan *updateResult, len(tpCavs))
	ret := make(map[string]security.Discharge, len(tpCavs))
	for {
		want := len(todo)
		now := time.Now()
		for _, w := range todo {
			updateDischarge(ctx, now, w.impetus, w.caveat, ch)
		}
		got := 0
		for i := 0; i < want; i++ {
			res := <-ch
			id := res.discharge.ID()
			ret[id] = res.discharge
			if res.done {
				minRefreshTime = minTime(minRefreshTime, res.refreshTime)
				delete(todo, id)
				got++
			}
		}
		if got == want {
			return ret, minRefreshTime
		}
		if got == 0 {
			return ret, minTime(minRefreshTime, time.Now().Add(time.Minute))
		}
	}
}

func updateDischarge(
	ctx *context.T,
	now time.Time,
	impetus security.DischargeImpetus,
	caveat security.Caveat,
	out chan<- *updateResult) {
	bstore := v23.GetPrincipal(ctx).BlessingStore()
	dis, ct := bstore.Discharge(caveat, impetus)
	if skip, _ := ctx.Value(skipDischargesKey{}).(bool); skip {
		// We can't fetch discharges while making a call to fetch a discharge.
		// Just go with what we have in the cache.
		out <- &updateResult{done: true, discharge: dis}
		return
	}
	if rt := refreshTime(dis, ct); dis.ID() != "" && (now.Before(rt) || rt.IsZero()) {
		// The cached value is still fresh, just return it.
		out <- &updateResult{done: true, discharge: dis, refreshTime: rt}
		return
	}
	// discharge in blessing store either doesn't exist or may be stale,
	// refresh via RPC.
	go func() {
		tp := caveat.ThirdPartyDetails()
		ctx = context.WithValue(ctx, skipDischargesKey{}, true)
		var newDis security.Discharge
		args, res := []interface{}{caveat, impetus}, []interface{}{&newDis}
		ctx.VI(3).Infof("Fetching discharge for %v", tp)
		if err := v23.GetClient(ctx).Call(ctx, tp.Location(), "Discharge", args, res); err != nil {
			ctx.VI(3).Infof("Discharge fetch for %v failed: %v", tp, err)
			out <- &updateResult{discharge: dis}
			return
		}
		bstore.CacheDischarge(newDis, caveat, impetus)
		out <- &updateResult{done: true, discharge: newDis, refreshTime: refreshTime(newDis, time.Now())}
	}()
}

// filteredImpetus returns a copy of 'before' after removing any values that are not required as per 'r'.
func filteredImpetus(r security.ThirdPartyRequirements, serverBlessings []string, method string, args []interface{}) security.DischargeImpetus {
	if !r.ReportServer {
		serverBlessings = nil
	}
	if !r.ReportMethod {
		method = ""
	}
	if !r.ReportArguments {
		args = nil
	}
	return mkDischargeImpetus(serverBlessings, method, args)
}

func refreshTime(dis security.Discharge, cacheTime time.Time) time.Time {
	expiry := dis.Expiry()
	if expiry.IsZero() {
		return time.Time{}
	}
	if cacheTime.IsZero() {
		// If we don't know the cache time, just try to refresh within a minute of
		// the expiry time.
		return expiry.Add(-time.Minute)
	}
	lifetime := expiry.Sub(cacheTime)
	return expiry.Add(-time.Duration(float64(lifetime) * DischargeRefreshFraction))
}

func minTime(a, b time.Time) time.Time {
	if a.IsZero() || (!b.IsZero() && b.Before(a)) {
		return b
	}
	return a
}

func mkDischargeImpetus(serverBlessings []string, method string, args []interface{}) security.DischargeImpetus {
	var impetus security.DischargeImpetus
	if len(serverBlessings) > 0 {
		impetus.Server = make([]security.BlessingPattern, len(serverBlessings))
		for i, b := range serverBlessings {
			impetus.Server[i] = security.BlessingPattern(b)
		}
	}
	impetus.Method = method
	if len(args) > 0 {
		impetus.Arguments = make([]*vdl.Value, len(args))
		for i, a := range args {
			vArg, err := vdl.ValueFromReflect(reflect.ValueOf(a))
			if err != nil {
				continue
			}
			impetus.Arguments[i] = vArg
		}
	}
	return impetus
}
<|endoftext|>


<commit_before>package mailgun

type Complaint struct {
	Count     int
	CreatedAt string
	Address   string
}

func (m *mailgunImpl) GetComplaints(limit, skip int) (int, []interface{}, error) {
	// TODO - this is NOT complete!
	return -1, nil, nil
}
<commit_msg>Added more to the complaint api.<commit_after>package mailgun

import (
	"github.com/mbanzon/simplehttp"
)

const (
	complaintsEndpoint = "complaints"
)

type Complaint struct {
	Count     int    `json:"count"`
	CreatedAt string `json:"created_at"`
	Address   string `json:"address"`
}

type complaintContainer struct {
	TotalCount int         `json:"total_count"`
	Items      []Complaint `json:"items"`
}

func (m *mailgunImpl) GetComplaints(limit, skip int) (int, []interface{}, error) {
	simplehttp.NewGetRequest(generateApiUrl(m, complaintsEndpoint))
	// TODO - this is NOT complete!
	return -1, nil, nil
}
<|endoftext|>


<commit_before>package suggest

/*
 * inspired by
 * http://www.aaai.org/ocs/index.php/AAAI/AAAI10/paper/viewFile/1939/2234
 * http://nlp.stanford.edu/IR-book/html/htmledition/k-gram-indexes-for-wildcard-queries-1.html
 * http://bazhenov.me/blog/2012/08/04/autocomplete.html
 * http://www.aclweb.org/anthology/C10-1096
 */

import (
	"container/heap"
)

type invertedListsT map[int][]int

type NGramIndex struct {
	k          int
	alphabet   Alphabet
	clean      *cleaner
	indices    []invertedListsT
	dictionary []string
	index      int
	config     *conf
}

type conf struct {
	threshold   float64 // 0 - 1
	measureName MeasureT
	pad         string
	wrap        string
}

var defaultConf *conf

func init() {
	defaultConf = &conf{0.5, COSINE, "$", "$"}
}

func NewNGramIndex(k int) *NGramIndex {
	if k < 2 || k > 4 {
		panic("k should be in [2, 4]")
	}

	// TODO declare as constructor argument
	alphabet := NewCompositeAlphabet([]Alphabet{
		NewEnglishAlphabet(),
		NewNumberAlphabet(),
		NewRussianAlphabet(),
		// TODO use config.pad here
		NewSimpleAlphabet([]rune{'$'}),
	})

	// TODO use in constructor
	clean := newCleaner(alphabet.Chars(), defaultConf.pad, defaultConf.wrap)

	return &NGramIndex{
		k, alphabet, clean, make([]invertedListsT, 0), make([]string, 0), 0, defaultConf,
	}
}

// Add given word to invertedList
func (self *NGramIndex) AddWord(word string) {
	prepared := self.prepareString(word)
	set := self.getNGramSet(prepared)
	cardinality := len(set)

	if len(self.indices) <= cardinality {
		tmp := make([]invertedListsT, cardinality, cardinality*2)
		copy(tmp, self.indices)
		self.indices = tmp
	}

	invertedLists := self.indices[cardinality-1]
	if invertedLists == nil {
		invertedLists = make(invertedListsT)
		self.indices[cardinality-1] = invertedLists
	}

	for _, index := range set {
		invertedLists[index] = append(invertedLists[index], self.index)
	}

	self.dictionary = append(self.dictionary, word)
	self.index++
}

// Return top-k similar strings
func (self *NGramIndex) Suggest(word string, topK int) []string {
	result := make([]string, 0, topK)
	preparedWord := self.prepareString(word)
	if len(preparedWord) < self.k {
		return result
	}

	candidates := self.search(preparedWord, topK)
	for candidates.Len() > 0 {
		r := heap.Pop(candidates).(*rank)
		result = append([]string{self.dictionary[r.id]}, result...)
	}

	return result
}

func (self *NGramIndex) search(word string, topK int) *heapImpl {
	set := self.getNGramSet(word)
	sizeA := len(set)

	h := &heapImpl{}
	mm := getMeasure(self.config.measureName)
	alpha := self.config.threshold
	bMin, bMax := mm.minY(alpha, sizeA), mm.maxY(alpha, sizeA)
	rid := make([][]int, 0, sizeA)
	for sizeB := bMax; sizeB >= bMin; sizeB-- {
		if len(self.indices) <= sizeB {
			continue
		}

		rid = rid[:0]
		// find max word id for memory optimize
		invertedLists := self.indices[sizeB]
		for _, index := range set {
			list := invertedLists[index]
			if len(list) > 0 {
				rid = append(rid, list)
			}
		}

		t := mm.threshold(alpha, sizeA, sizeB)
		if len(rid) < t || t == 0 {
			continue
		}

		counts := divideSkip(rid, t)
		// use heap search for finding top k items in a list efficiently
		// see http://stevehanov.ca/blog/index.php?id=122
		for inter := len(counts) - 1; inter >= 0; inter-- {
			if len(counts[inter]) == 0 {
				continue
			}

			list := counts[inter]
			for _, id := range list {
				distance := mm.distance(inter, sizeA, sizeB)
				if h.Len() < topK || h.Top().(*rank).distance > distance {
					if h.Len() == topK {
						heap.Pop(h)
					}

					heap.Push(h, &rank{id, distance})
				}
			}
		}
	}

	return h
}

// Return unique ngrams
func (self *NGramIndex) getNGramSet(word string) []int {
	ngrams := SplitIntoNGrams(word, self.k)
	set := make(map[int]struct{}, len(ngrams))
	list := make([]int, 0, len(ngrams))
	for _, ngram := range ngrams {
		index := self.ngramToIndex(ngram)
		_, found := set[index]
		set[index] = struct{}{}
		if !found {
			list = append(list, index)
		}
	}

	return list
}

//
func (self *NGramIndex) ngramToIndex(ngram string) int {
	index := 0
	for _, char := range ngram {
		i := self.alphabet.MapChar(char)
		if index == INVALID_CHAR {
			panic("Invalid char was detected")
		}

		index = index*self.alphabet.Size() + i
	}

	return index
}

// Prepare string for indexing
func (self *NGramIndex) prepareString(word string) string {
	return self.clean.cleanAndWrap(word)
}
<commit_msg>cleanup code<commit_after>package suggest

/*
 * inspired by
 * http://www.aaai.org/ocs/index.php/AAAI/AAAI10/paper/viewFile/1939/2234
 * http://nlp.stanford.edu/IR-book/html/htmledition/k-gram-indexes-for-wildcard-queries-1.html
 * http://bazhenov.me/blog/2012/08/04/autocomplete.html
 * http://www.aclweb.org/anthology/C10-1096
 */

import (
	"container/heap"
)

type invertedListsT map[int][]int

type NGramIndex struct {
	k          int
	alphabet   Alphabet
	clean      *cleaner
	indices    []invertedListsT
	dictionary []string
	index      int
	config     *conf
}

type conf struct {
	threshold   float64 // 0 - 1
	measureName MeasureT
	pad         string
	wrap        string
}

var defaultConf *conf

func init() {
	defaultConf = &conf{0.5, COSINE, "$", "$"}
}

func NewNGramIndex(k int) *NGramIndex {
	if k < 2 || k > 4 {
		panic("k should be in [2, 4]")
	}

	// TODO declare as constructor argument
	alphabet := NewCompositeAlphabet([]Alphabet{
		NewEnglishAlphabet(),
		NewNumberAlphabet(),
		NewRussianAlphabet(),
		// TODO use config.pad here
		NewSimpleAlphabet([]rune{'$'}),
	})

	// TODO use in constructor
	clean := newCleaner(alphabet.Chars(), defaultConf.pad, defaultConf.wrap)

	return &NGramIndex{
		k, alphabet, clean, make([]invertedListsT, 0), make([]string, 0), 0, defaultConf,
	}
}

// Add given word to invertedList
func (self *NGramIndex) AddWord(word string) {
	prepared := self.prepareString(word)
	set := self.getNGramSet(prepared)
	cardinality := len(set)

	if len(self.indices) <= cardinality {
		tmp := make([]invertedListsT, cardinality, cardinality*2)
		copy(tmp, self.indices)
		self.indices = tmp
	}

	invertedLists := self.indices[cardinality-1]
	if invertedLists == nil {
		invertedLists = make(invertedListsT)
		self.indices[cardinality-1] = invertedLists
	}

	for _, index := range set {
		invertedLists[index] = append(invertedLists[index], self.index)
	}

	self.dictionary = append(self.dictionary, word)
	self.index++
}

// Return top-k similar strings
func (self *NGramIndex) Suggest(word string, topK int) []string {
	result := make([]string, 0, topK)
	preparedWord := self.prepareString(word)
	if len(preparedWord) < self.k {
		return result
	}

	candidates := self.search(preparedWord, topK)
	for candidates.Len() > 0 {
		r := heap.Pop(candidates).(*rank)
		result = append([]string{self.dictionary[r.id]}, result...)
	}

	return result
}

func (self *NGramIndex) search(word string, topK int) *heapImpl {
	set := self.getNGramSet(word)
	sizeA := len(set)

	h := &heapImpl{}
	mm := getMeasure(self.config.measureName)
	alpha := self.config.threshold
	bMin, bMax := mm.minY(alpha, sizeA), mm.maxY(alpha, sizeA)
	rid := make([][]int, 0, sizeA)
	lenIndices := len(self.indices)
	for sizeB := bMax; sizeB >= bMin; sizeB-- {
		if lenIndices <= sizeB {
			continue
		}

		threshold := mm.threshold(alpha, sizeA, sizeB)
		if threshold == 0 {
			continue
		}

		rid = rid[:0]
		invertedLists := self.indices[sizeB]
		for _, index := range set {
			list := invertedLists[index]
			if len(list) > 0 {
				rid = append(rid, list)
			}
		}

		if len(rid) < threshold {
			continue
		}

		counts := divideSkip(rid, threshold)
		// use heap search for finding top k items in a list efficiently
		// see http://stevehanov.ca/blog/index.php?id=122
		for inter := len(counts) - 1; inter >= threshold; inter-- {
			for _, id := range counts[inter] {
				distance := mm.distance(inter, sizeA, sizeB)
				if h.Len() < topK || h.Top().(*rank).distance > distance {
					if h.Len() == topK {
						heap.Pop(h)
					}

					heap.Push(h, &rank{id, distance})
				}
			}
		}
	}

	return h
}

// Return unique ngrams
func (self *NGramIndex) getNGramSet(word string) []int {
	ngrams := SplitIntoNGrams(word, self.k)
	set := make(map[int]struct{}, len(ngrams))
	list := make([]int, 0, len(ngrams))
	for _, ngram := range ngrams {
		index := self.ngramToIndex(ngram)
		_, found := set[index]
		set[index] = struct{}{}
		if !found {
			list = append(list, index)
		}
	}

	return list
}

//
func (self *NGramIndex) ngramToIndex(ngram string) int {
	index := 0
	for _, char := range ngram {
		i := self.alphabet.MapChar(char)
		if index == INVALID_CHAR {
			panic("Invalid char was detected")
		}

		index = index*self.alphabet.Size() + i
	}

	return index
}

// Prepare string for indexing
func (self *NGramIndex) prepareString(word string) string {
	return self.clean.cleanAndWrap(word)
}
<|endoftext|>


<commit_before>package wrappa_test

import (
	"github.com/concourse/concourse/atc"
	"github.com/concourse/concourse/atc/db/dbfakes"
	"github.com/concourse/concourse/atc/wrappa"
	"github.com/tedsuo/rata"

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

var _ = Describe("RejectArchivedWrappa", func() {
	var (
		raWrappa         *wrappa.RejectArchivedWrappa
		raHandlerFactory wrappa.RejectArchivedHandlerFactory
	)

	BeforeEach(func() {
		fakeTeamFactory := new(dbfakes.FakeTeamFactory)
		raHandlerFactory = wrappa.NewRejectArchivedHandlerFactory(fakeTeamFactory)
		raWrappa = wrappa.NewRejectArchivedWrappa(raHandlerFactory)
	})

	It("wraps endpoints", func() {
		inputHandlers := rata.Handlers{}

		for _, route := range atc.Routes {
			inputHandlers[route.Name] = &stupidHandler{}
		}

		expectedHandlers := rata.Handlers{

			atc.PausePipeline:           raHandlerFactory.RejectArchived(inputHandlers[atc.PausePipeline]),
			atc.CreateJobBuild:          raHandlerFactory.RejectArchived(inputHandlers[atc.CreateJobBuild]),
			atc.ScheduleJob:             raHandlerFactory.RejectArchived(inputHandlers[atc.ScheduleJob]),
			atc.CheckResource:           raHandlerFactory.RejectArchived(inputHandlers[atc.CheckResource]),
			atc.CheckResourceType:       raHandlerFactory.RejectArchived(inputHandlers[atc.CheckResourceType]),
			atc.DisableResourceVersion:  raHandlerFactory.RejectArchived(inputHandlers[atc.DisableResourceVersion]),
			atc.EnableResourceVersion:   raHandlerFactory.RejectArchived(inputHandlers[atc.EnableResourceVersion]),
			atc.PinResourceVersion:      raHandlerFactory.RejectArchived(inputHandlers[atc.PinResourceVersion]),
			atc.UnpinResource:           raHandlerFactory.RejectArchived(inputHandlers[atc.UnpinResource]),
			atc.SetPinCommentOnResource: raHandlerFactory.RejectArchived(inputHandlers[atc.SetPinCommentOnResource]),
			atc.RerunJobBuild:           raHandlerFactory.RejectArchived(inputHandlers[atc.RerunJobBuild]),

			atc.GetConfig:                     inputHandlers[atc.GetConfig],
			atc.GetBuild:                      inputHandlers[atc.GetBuild],
			atc.BuildResources:                inputHandlers[atc.BuildResources],
			atc.BuildEvents:                   inputHandlers[atc.BuildEvents],
			atc.ListBuildArtifacts:            inputHandlers[atc.ListBuildArtifacts],
			atc.GetBuildPreparation:           inputHandlers[atc.GetBuildPreparation],
			atc.GetBuildPlan:                  inputHandlers[atc.GetBuildPlan],
			atc.AbortBuild:                    inputHandlers[atc.AbortBuild],
			atc.PruneWorker:                   inputHandlers[atc.PruneWorker],
			atc.LandWorker:                    inputHandlers[atc.LandWorker],
			atc.ReportWorkerContainers:        inputHandlers[atc.ReportWorkerContainers],
			atc.ReportWorkerVolumes:           inputHandlers[atc.ReportWorkerVolumes],
			atc.RetireWorker:                  inputHandlers[atc.RetireWorker],
			atc.ListDestroyingContainers:      inputHandlers[atc.ListDestroyingContainers],
			atc.ListDestroyingVolumes:         inputHandlers[atc.ListDestroyingVolumes],
			atc.GetPipeline:                   inputHandlers[atc.GetPipeline],
			atc.GetJobBuild:                   inputHandlers[atc.GetJobBuild],
			atc.PipelineBadge:                 inputHandlers[atc.PipelineBadge],
			atc.JobBadge:                      inputHandlers[atc.JobBadge],
			atc.ListJobs:                      inputHandlers[atc.ListJobs],
			atc.GetJob:                        inputHandlers[atc.GetJob],
			atc.ListJobBuilds:                 inputHandlers[atc.ListJobBuilds],
			atc.ListPipelineBuilds:            inputHandlers[atc.ListPipelineBuilds],
			atc.GetResource:                   inputHandlers[atc.GetResource],
			atc.ListBuildsWithVersionAsInput:  inputHandlers[atc.ListBuildsWithVersionAsInput],
			atc.ListBuildsWithVersionAsOutput: inputHandlers[atc.ListBuildsWithVersionAsOutput],
			atc.ListResources:                 inputHandlers[atc.ListResources],
			atc.ListResourceTypes:             inputHandlers[atc.ListResourceTypes],
			atc.ListResourceVersions:          inputHandlers[atc.ListResourceVersions],
			atc.GetResourceCausality:          inputHandlers[atc.GetResourceCausality],
			atc.GetResourceVersion:            inputHandlers[atc.GetResourceVersion],
			atc.CreateBuild:                   inputHandlers[atc.CreateBuild],
			atc.GetContainer:                  inputHandlers[atc.GetContainer],
			atc.HijackContainer:               inputHandlers[atc.HijackContainer],
			atc.ListContainers:                inputHandlers[atc.ListContainers],
			atc.ListVolumes:                   inputHandlers[atc.ListVolumes],
			atc.ListTeamBuilds:                inputHandlers[atc.ListTeamBuilds],
			atc.ListWorkers:                   inputHandlers[atc.ListWorkers],
			atc.RegisterWorker:                inputHandlers[atc.RegisterWorker],
			atc.HeartbeatWorker:               inputHandlers[atc.HeartbeatWorker],
			atc.DeleteWorker:                  inputHandlers[atc.DeleteWorker],
			atc.GetTeam:                       inputHandlers[atc.GetTeam],
			atc.SetTeam:                       inputHandlers[atc.SetTeam],
			atc.RenameTeam:                    inputHandlers[atc.RenameTeam],
			atc.DestroyTeam:                   inputHandlers[atc.DestroyTeam],
			atc.GetUser:                       inputHandlers[atc.GetUser],
			atc.GetInfo:                       inputHandlers[atc.GetInfo],
			atc.GetCheck:                      inputHandlers[atc.GetCheck],
			atc.DownloadCLI:                   inputHandlers[atc.DownloadCLI],
			atc.CheckResourceWebHook:          inputHandlers[atc.CheckResourceWebHook],
			atc.ListAllPipelines:              inputHandlers[atc.ListAllPipelines],
			atc.ListBuilds:                    inputHandlers[atc.ListBuilds],
			atc.ListPipelines:                 inputHandlers[atc.ListPipelines],
			atc.ListAllJobs:                   inputHandlers[atc.ListAllJobs],
			atc.ListAllResources:              inputHandlers[atc.ListAllResources],
			atc.ListTeams:                     inputHandlers[atc.ListTeams],
			atc.MainJobBadge:                  inputHandlers[atc.MainJobBadge],
			atc.GetWall:                       inputHandlers[atc.GetWall],
			atc.GetLogLevel:                   inputHandlers[atc.GetLogLevel],
			atc.SetLogLevel:                   inputHandlers[atc.SetLogLevel],
			atc.GetInfoCreds:                  inputHandlers[atc.GetInfoCreds],
			atc.ListActiveUsersSince:          inputHandlers[atc.ListActiveUsersSince],
			atc.SetWall:                       inputHandlers[atc.SetWall],
			atc.ClearWall:                     inputHandlers[atc.ClearWall],
			atc.DeletePipeline:                inputHandlers[atc.DeletePipeline],
			atc.GetCC:                         inputHandlers[atc.GetCC],
			atc.GetVersionsDB:                 inputHandlers[atc.GetVersionsDB],
			atc.ListJobInputs:                 inputHandlers[atc.ListJobInputs],
			atc.OrderPipelines:                inputHandlers[atc.OrderPipelines],
			atc.PauseJob:                      inputHandlers[atc.PauseJob],
			atc.ArchivePipeline:               inputHandlers[atc.ArchivePipeline],
			atc.RenamePipeline:                inputHandlers[atc.RenamePipeline],
			atc.SaveConfig:                    inputHandlers[atc.SaveConfig],
			atc.UnpauseJob:                    inputHandlers[atc.UnpauseJob],
			atc.UnpausePipeline:               inputHandlers[atc.UnpausePipeline],
			atc.ExposePipeline:                inputHandlers[atc.ExposePipeline],
			atc.HidePipeline:                  inputHandlers[atc.HidePipeline],
			atc.CreatePipelineBuild:           inputHandlers[atc.CreatePipelineBuild],
			atc.ClearTaskCache:                inputHandlers[atc.ClearTaskCache],
			atc.CreateArtifact:                inputHandlers[atc.CreateArtifact],
			atc.GetArtifact:                   inputHandlers[atc.GetArtifact],
		}

		wrappedHandlers := raWrappa.Wrap(inputHandlers)

		for name, _ := range inputHandlers {
			Expect(wrappedHandlers[name]).To(BeIdenticalTo(expectedHandlers[name]), "handler is "+name)
		}
	})
})
<commit_msg>wrappa: Reduce test size of RejectArchivedWrappa<commit_after>package wrappa_test

import (
	"github.com/concourse/concourse/atc"
	"github.com/concourse/concourse/atc/db/dbfakes"
	"github.com/concourse/concourse/atc/wrappa"
	"github.com/tedsuo/rata"

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

var _ = Describe("RejectArchivedWrappa", func() {
	var (
		raWrappa         *wrappa.RejectArchivedWrappa
		raHandlerFactory wrappa.RejectArchivedHandlerFactory
	)

	BeforeEach(func() {
		fakeTeamFactory := new(dbfakes.FakeTeamFactory)
		raHandlerFactory = wrappa.NewRejectArchivedHandlerFactory(fakeTeamFactory)
		raWrappa = wrappa.NewRejectArchivedWrappa(raHandlerFactory)
	})

	It("wraps endpoints", func() {
		inputHandlers := rata.Handlers{}

		for _, route := range atc.Routes {
			inputHandlers[route.Name] = &stupidHandler{}
		}

		rejectArchivedRoutes := []string{
			atc.PausePipeline,
			atc.CreateJobBuild,
			atc.ScheduleJob,
			atc.CheckResource,
			atc.CheckResourceType,
			atc.DisableResourceVersion,
			atc.EnableResourceVersion,
			atc.PinResourceVersion,
			atc.UnpinResource,
			atc.SetPinCommentOnResource,
			atc.RerunJobBuild,
		}

		rejectArchivedLookup := make(map[string]bool)
		for _, name := range rejectArchivedRoutes {
			rejectArchivedLookup[name] = true
		}

		wrappedHandlers := raWrappa.Wrap(inputHandlers)

		for name, handler := range inputHandlers {
			expectedHandler := handler
			if rejectArchivedLookup[name] {
				expectedHandler = raHandlerFactory.RejectArchived(expectedHandler)
			}
			Expect(wrappedHandlers[name]).To(BeIdenticalTo(expectedHandler), "handler is "+name)
		}
	})
})
<|endoftext|>


<commit_before>package main

import "github.com/rogeralsing/gam/actor"
import "github.com/rogeralsing/gam/remoting"
import "fmt"
import "bufio"
import "os"
import "github.com/rogeralsing/gam/examples/remoting/messages"

type MyActor struct{
	count int
}

func (state *MyActor) Receive(context actor.Context) {
	switch msg := context.Message().(type) {
	case string:
		fmt.Println(msg)
	case *messages.Response:
		state.count++
		if state.count % 1000 == 0 {
			fmt.Println(state.count)
		}
	}
}

func main() {
	remoting.StartServer("localhost:8090")

	pid := actor.SpawnTemplate(&MyActor{})
	message := &messages.Echo{Message: "hej", Sender: pid}
	remote := actor.NewPID("localhost:8091", "foo")
	for i := 0; i < 100000; i++ {
		remote.Tell(message)
	}

	bufio.NewReader(os.Stdin).ReadString('\n')
}
<commit_msg>replaced fmt in example<commit_after>package main

import "github.com/rogeralsing/gam/actor"
import "github.com/rogeralsing/gam/remoting"
import "fmt"
import "log"
import "bufio"
import "os"
import "github.com/rogeralsing/gam/examples/remoting/messages"

type MyActor struct{
	count int
}

func (state *MyActor) Receive(context actor.Context) {
	switch msg := context.Message().(type) {
	case string:
		fmt.Println(msg)
	case *messages.Response:
		state.count++
		if state.count % 1000 == 0 {
			log.Println(state.count)
		}
	}
}

func main() {
	remoting.StartServer("localhost:8090")

	pid := actor.SpawnTemplate(&MyActor{})
	message := &messages.Echo{Message: "hej", Sender: pid}
	remote := actor.NewPID("localhost:8091", "foo")
	for i := 0; i < 100000; i++ {
		remote.Tell(message)
	}

	bufio.NewReader(os.Stdin).ReadString('\n')
}
<|endoftext|>


<commit_before>package gocql

import (
	"fmt"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"
)

type nodeState int32

func (n nodeState) String() string {
	if n == NodeUp {
		return "UP"
	} else if n == NodeDown {
		return "DOWN"
	}
	return fmt.Sprintf("UNKNOWN_%d", n)
}

const (
	NodeUp nodeState = iota
	NodeDown
)

type cassVersion struct {
	Major, Minor, Patch int
}

func (c *cassVersion) UnmarshalCQL(info TypeInfo, data []byte) error {
	version := strings.TrimSuffix(string(data), "-SNAPSHOT")
	v := strings.Split(version, ".")
	if len(v) != 3 {
		return fmt.Errorf("invalid schema_version: %v", string(data))
	}

	var err error
	c.Major, err = strconv.Atoi(v[0])
	if err != nil {
		return fmt.Errorf("invalid major version %v: %v", v[0], err)
	}

	c.Minor, err = strconv.Atoi(v[1])
	if err != nil {
		return fmt.Errorf("invalid minor version %v: %v", v[1], err)
	}

	c.Patch, err = strconv.Atoi(v[2])
	if err != nil {
		return fmt.Errorf("invalid patch version %v: %v", v[2], err)
	}

	return nil
}

func (c cassVersion) String() string {
	return fmt.Sprintf("v%d.%d.%d", c.Major, c.Minor, c.Patch)
}

func (c cassVersion) nodeUpDelay() time.Duration {
	if c.Major >= 2 && c.Minor >= 2 {
		// CASSANDRA-8236
		return 0
	}

	return 10 * time.Second
}

type HostInfo struct {
	// TODO(zariel): reduce locking maybe, not all values will change, but to ensure
	// that we are thread safe use a mutex to access all fields.
	mu         sync.RWMutex
	peer       string
	port       int
	dataCenter string
	rack       string
	hostId     string
	version    cassVersion
	state      nodeState
	tokens     []string
}

func (h *HostInfo) Equal(host *HostInfo) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	host.mu.RLock()
	defer host.mu.RUnlock()

	return h.peer == host.peer && h.hostId == host.hostId
}

func (h *HostInfo) Peer() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.peer
}

func (h *HostInfo) setPeer(peer string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.peer = peer
	return h
}

func (h *HostInfo) DataCenter() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.dataCenter
}

func (h *HostInfo) setDataCenter(dataCenter string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.dataCenter = dataCenter
	return h
}

func (h *HostInfo) Rack() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.rack
}

func (h *HostInfo) setRack(rack string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.rack = rack
	return h
}

func (h *HostInfo) HostID() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.hostId
}

func (h *HostInfo) setHostID(hostID string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.hostId = hostID
	return h
}

func (h *HostInfo) Version() cassVersion {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.version
}

func (h *HostInfo) setVersion(major, minor, patch int) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.version = cassVersion{major, minor, patch}
	return h
}

func (h *HostInfo) State() nodeState {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.state
}

func (h *HostInfo) setState(state nodeState) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.state = state
	return h
}

func (h *HostInfo) Tokens() []string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.tokens
}

func (h *HostInfo) setTokens(tokens []string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.tokens = tokens
	return h
}

func (h *HostInfo) Port() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.port
}

func (h *HostInfo) setPort(port int) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.port = port
	return h
}

func (h *HostInfo) update(from *HostInfo) {
	h.mu.Lock()
	defer h.mu.Unlock()

	h.tokens = from.tokens
	h.version = from.version
	h.hostId = from.hostId
	h.dataCenter = from.dataCenter
}

func (h *HostInfo) IsUp() bool {
	return h.State() == NodeUp
}

func (h *HostInfo) String() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return fmt.Sprintf("[hostinfo peer=%q port=%d data_centre=%q rack=%q host_id=%q version=%q state=%s num_tokens=%d]", h.peer, h.port, h.dataCenter, h.rack, h.hostId, h.version, h.state, len(h.tokens))
}

// Polls system.peers at a specific interval to find new hosts
type ringDescriber struct {
	dcFilter   string
	rackFilter string
	session    *Session
	closeChan  chan bool
	// indicates that we can use system.local to get the connections remote address
	localHasRpcAddr bool

	mu              sync.Mutex
	prevHosts       []*HostInfo
	prevPartitioner string
}

func checkSystemLocal(control *controlConn) (bool, error) {
	iter := control.query("SELECT broadcast_address FROM system.local")
	if err := iter.err; err != nil {
		if errf, ok := err.(*errorFrame); ok {
			if errf.code == errSyntax {
				return false, nil
			}
		}

		return false, err
	}

	return true, nil
}

func (r *ringDescriber) GetHosts() (hosts []*HostInfo, partitioner string, err error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	// we need conn to be the same because we need to query system.peers and system.local
	// on the same node to get the whole cluster

	const (
		legacyLocalQuery = "SELECT data_center, rack, host_id, tokens, partitioner, release_version FROM system.local"
		// only supported in 2.2.0, 2.1.6, 2.0.16
		localQuery = "SELECT broadcast_address, data_center, rack, host_id, tokens, partitioner, release_version FROM system.local"
	)

	localHost := &HostInfo{}
	if r.localHasRpcAddr {
		iter := r.session.control.query(localQuery)
		if iter == nil {
			return r.prevHosts, r.prevPartitioner, nil
		}

		iter.Scan(&localHost.peer, &localHost.dataCenter, &localHost.rack,
			&localHost.hostId, &localHost.tokens, &partitioner, &localHost.version)

		if err = iter.Close(); err != nil {
			return nil, "", err
		}
	} else {
		iter := r.session.control.query(legacyLocalQuery)
		if iter == nil {
			return r.prevHosts, r.prevPartitioner, nil
		}

		iter.Scan(&localHost.dataCenter, &localHost.rack, &localHost.hostId, &localHost.tokens, &partitioner, &localHost.version)

		if err = iter.Close(); err != nil {
			return nil, "", err
		}

		addr, _, err := net.SplitHostPort(r.session.control.addr())
		if err != nil {
			// this should not happen, ever, as this is the address that was dialed by conn, here
			// a panic makes sense, please report a bug if it occurs.
			panic(err)
		}

		localHost.peer = addr
	}

	localHost.port = r.session.cfg.Port

	hosts = []*HostInfo{localHost}

	iter := r.session.control.query("SELECT rpc_address, data_center, rack, host_id, tokens, release_version FROM system.peers")
	if iter == nil {
		return r.prevHosts, r.prevPartitioner, nil
	}

	host := &HostInfo{port: r.session.cfg.Port}
	for iter.Scan(&host.peer, &host.dataCenter, &host.rack, &host.hostId, &host.tokens, &host.version) {
		if r.matchFilter(host) {
			hosts = append(hosts, host)
		}
		host = &HostInfo{
			port: r.session.cfg.Port,
		}
	}

	if err = iter.Close(); err != nil {
		return nil, "", err
	}

	r.prevHosts = hosts
	r.prevPartitioner = partitioner

	return hosts, partitioner, nil
}

func (r *ringDescriber) matchFilter(host *HostInfo) bool {
	if r.dcFilter != "" && r.dcFilter != host.DataCenter() {
		return false
	}

	if r.rackFilter != "" && r.rackFilter != host.Rack() {
		return false
	}

	return true
}

func (r *ringDescriber) refreshRing() error {
	// if we have 0 hosts this will return the previous list of hosts to
	// attempt to reconnect to the cluster otherwise we would never find
	// downed hosts again, could possibly have an optimisation to only
	// try to add new hosts if GetHosts didnt error and the hosts didnt change.
	hosts, partitioner, err := r.GetHosts()
	if err != nil {
		return err
	}

	// TODO: move this to session
	// TODO: handle removing hosts here
	for _, h := range hosts {
		if host, ok := r.session.ring.addHostIfMissing(h); !ok {
			r.session.pool.addHost(h)
		} else {
			host.update(h)
		}
	}

	r.session.pool.SetPartitioner(partitioner)
	return nil
}
<commit_msg>DSE reports >3 elements in the version string<commit_after>package gocql

import (
	"fmt"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"
)

type nodeState int32

func (n nodeState) String() string {
	if n == NodeUp {
		return "UP"
	} else if n == NodeDown {
		return "DOWN"
	}
	return fmt.Sprintf("UNKNOWN_%d", n)
}

const (
	NodeUp nodeState = iota
	NodeDown
)

type cassVersion struct {
	Major, Minor, Patch int
}

func (c *cassVersion) UnmarshalCQL(info TypeInfo, data []byte) error {
	version := strings.TrimSuffix(string(data), "-SNAPSHOT")
	v := strings.Split(version, ".")
	if len(v) < 3 {
		return fmt.Errorf("invalid schema_version: %v", string(data))
	}

	var err error
	c.Major, err = strconv.Atoi(v[0])
	if err != nil {
		return fmt.Errorf("invalid major version %v: %v", v[0], err)
	}

	c.Minor, err = strconv.Atoi(v[1])
	if err != nil {
		return fmt.Errorf("invalid minor version %v: %v", v[1], err)
	}

	c.Patch, err = strconv.Atoi(v[2])
	if err != nil {
		return fmt.Errorf("invalid patch version %v: %v", v[2], err)
	}

	return nil
}

func (c cassVersion) String() string {
	return fmt.Sprintf("v%d.%d.%d", c.Major, c.Minor, c.Patch)
}

func (c cassVersion) nodeUpDelay() time.Duration {
	if c.Major >= 2 && c.Minor >= 2 {
		// CASSANDRA-8236
		return 0
	}

	return 10 * time.Second
}

type HostInfo struct {
	// TODO(zariel): reduce locking maybe, not all values will change, but to ensure
	// that we are thread safe use a mutex to access all fields.
	mu         sync.RWMutex
	peer       string
	port       int
	dataCenter string
	rack       string
	hostId     string
	version    cassVersion
	state      nodeState
	tokens     []string
}

func (h *HostInfo) Equal(host *HostInfo) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	host.mu.RLock()
	defer host.mu.RUnlock()

	return h.peer == host.peer && h.hostId == host.hostId
}

func (h *HostInfo) Peer() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.peer
}

func (h *HostInfo) setPeer(peer string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.peer = peer
	return h
}

func (h *HostInfo) DataCenter() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.dataCenter
}

func (h *HostInfo) setDataCenter(dataCenter string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.dataCenter = dataCenter
	return h
}

func (h *HostInfo) Rack() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.rack
}

func (h *HostInfo) setRack(rack string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.rack = rack
	return h
}

func (h *HostInfo) HostID() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.hostId
}

func (h *HostInfo) setHostID(hostID string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.hostId = hostID
	return h
}

func (h *HostInfo) Version() cassVersion {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.version
}

func (h *HostInfo) setVersion(major, minor, patch int) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.version = cassVersion{major, minor, patch}
	return h
}

func (h *HostInfo) State() nodeState {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.state
}

func (h *HostInfo) setState(state nodeState) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.state = state
	return h
}

func (h *HostInfo) Tokens() []string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.tokens
}

func (h *HostInfo) setTokens(tokens []string) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.tokens = tokens
	return h
}

func (h *HostInfo) Port() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.port
}

func (h *HostInfo) setPort(port int) *HostInfo {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.port = port
	return h
}

func (h *HostInfo) update(from *HostInfo) {
	h.mu.Lock()
	defer h.mu.Unlock()

	h.tokens = from.tokens
	h.version = from.version
	h.hostId = from.hostId
	h.dataCenter = from.dataCenter
}

func (h *HostInfo) IsUp() bool {
	return h.State() == NodeUp
}

func (h *HostInfo) String() string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return fmt.Sprintf("[hostinfo peer=%q port=%d data_centre=%q rack=%q host_id=%q version=%q state=%s num_tokens=%d]", h.peer, h.port, h.dataCenter, h.rack, h.hostId, h.version, h.state, len(h.tokens))
}

// Polls system.peers at a specific interval to find new hosts
type ringDescriber struct {
	dcFilter   string
	rackFilter string
	session    *Session
	closeChan  chan bool
	// indicates that we can use system.local to get the connections remote address
	localHasRpcAddr bool

	mu              sync.Mutex
	prevHosts       []*HostInfo
	prevPartitioner string
}

func checkSystemLocal(control *controlConn) (bool, error) {
	iter := control.query("SELECT broadcast_address FROM system.local")
	if err := iter.err; err != nil {
		if errf, ok := err.(*errorFrame); ok {
			if errf.code == errSyntax {
				return false, nil
			}
		}

		return false, err
	}

	return true, nil
}

func (r *ringDescriber) GetHosts() (hosts []*HostInfo, partitioner string, err error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	// we need conn to be the same because we need to query system.peers and system.local
	// on the same node to get the whole cluster

	const (
		legacyLocalQuery = "SELECT data_center, rack, host_id, tokens, partitioner, release_version FROM system.local"
		// only supported in 2.2.0, 2.1.6, 2.0.16
		localQuery = "SELECT broadcast_address, data_center, rack, host_id, tokens, partitioner, release_version FROM system.local"
	)

	localHost := &HostInfo{}
	if r.localHasRpcAddr {
		iter := r.session.control.query(localQuery)
		if iter == nil {
			return r.prevHosts, r.prevPartitioner, nil
		}

		iter.Scan(&localHost.peer, &localHost.dataCenter, &localHost.rack,
			&localHost.hostId, &localHost.tokens, &partitioner, &localHost.version)

		if err = iter.Close(); err != nil {
			return nil, "", err
		}
	} else {
		iter := r.session.control.query(legacyLocalQuery)
		if iter == nil {
			return r.prevHosts, r.prevPartitioner, nil
		}

		iter.Scan(&localHost.dataCenter, &localHost.rack, &localHost.hostId, &localHost.tokens, &partitioner, &localHost.version)

		if err = iter.Close(); err != nil {
			return nil, "", err
		}

		addr, _, err := net.SplitHostPort(r.session.control.addr())
		if err != nil {
			// this should not happen, ever, as this is the address that was dialed by conn, here
			// a panic makes sense, please report a bug if it occurs.
			panic(err)
		}

		localHost.peer = addr
	}

	localHost.port = r.session.cfg.Port

	hosts = []*HostInfo{localHost}

	iter := r.session.control.query("SELECT rpc_address, data_center, rack, host_id, tokens, release_version FROM system.peers")
	if iter == nil {
		return r.prevHosts, r.prevPartitioner, nil
	}

	host := &HostInfo{port: r.session.cfg.Port}
	for iter.Scan(&host.peer, &host.dataCenter, &host.rack, &host.hostId, &host.tokens, &host.version) {
		if r.matchFilter(host) {
			hosts = append(hosts, host)
		}
		host = &HostInfo{
			port: r.session.cfg.Port,
		}
	}

	if err = iter.Close(); err != nil {
		return nil, "", err
	}

	r.prevHosts = hosts
	r.prevPartitioner = partitioner

	return hosts, partitioner, nil
}

func (r *ringDescriber) matchFilter(host *HostInfo) bool {
	if r.dcFilter != "" && r.dcFilter != host.DataCenter() {
		return false
	}

	if r.rackFilter != "" && r.rackFilter != host.Rack() {
		return false
	}

	return true
}

func (r *ringDescriber) refreshRing() error {
	// if we have 0 hosts this will return the previous list of hosts to
	// attempt to reconnect to the cluster otherwise we would never find
	// downed hosts again, could possibly have an optimisation to only
	// try to add new hosts if GetHosts didnt error and the hosts didnt change.
	hosts, partitioner, err := r.GetHosts()
	if err != nil {
		return err
	}

	// TODO: move this to session
	// TODO: handle removing hosts here
	for _, h := range hosts {
		if host, ok := r.session.ring.addHostIfMissing(h); !ok {
			r.session.pool.addHost(h)
		} else {
			host.update(h)
		}
	}

	r.session.pool.SetPartitioner(partitioner)
	return nil
}
<|endoftext|>


<commit_before>// Package oauth provides an HTTP handler to handle OAuth2
// redirect and callback requests for the bebop web app.
package oauth

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

	"github.com/go-chi/chi"
	"github.com/satori/go.uuid"
	"golang.org/x/oauth2"

	"github.com/disintegration/bebop/jwt"
	"github.com/disintegration/bebop/store"
)

const (
	stateCookie   = "bebop_oauth_state"
	resultCookie  = "bebop_oauth_result"
	clientTimeout = 10 * time.Second
)

// Config is a configuration of an OAuth handler.
type Config struct {
	Logger     *log.Logger
	UserStore  store.UserStore
	JWTService jwt.Service
	MountURL   string
	CookiePath string
}

// Handler handles oauth2 authentication requests.
type Handler struct {
	*Config
	providers map[string]*provider
	router    chi.Router
}

// New creates a new handler based on the given config.
func New(config *Config) *Handler {
	h := &Handler{Config: config}

	h.providers = make(map[string]*provider)

	h.router = chi.NewRouter()
	h.router.Get("/begin/:provider", h.handleBegin)
	h.router.Get("/end/:provider", h.handleEnd)

	return h
}

// AddProvider adds a new provider to oauth handler.
func (h *Handler) AddProvider(name, id, secret string) error {
	pc, ok := providerConfigs[name]
	if !ok {
		return fmt.Errorf("oauth: unknown provider: %q", name)
	}

	if id == "" {
		return fmt.Errorf("oauth: empty client id of provider %q", name)
	}

	if secret == "" {
		return fmt.Errorf("oauth: empty client secret of provider %q", name)
	}

	if h.providers == nil {
		h.providers = make(map[string]*provider)
	}

	h.providers[name] = &provider{
		config: &oauth2.Config{
			ClientID:     id,
			ClientSecret: secret,
			RedirectURL:  h.MountURL + "/end/" + name,
			Endpoint:     pc.endpoint,
			Scopes:       pc.scopes,
		},
		getUser: pc.getUser,
	}

	return nil
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h.router.ServeHTTP(w, r)
}

func (h *Handler) handleBegin(w http.ResponseWriter, r *http.Request) {
	providerName := chi.URLParam(r, "provider")
	provider, ok := h.providers[providerName]
	if !ok {
		http.NotFound(w, r)
		return
	}

	state := h.genState()
	redirectURL := provider.config.AuthCodeURL(state)

	http.SetCookie(w, &http.Cookie{
		Name:     stateCookie,
		Value:    state,
		Path:     h.CookiePath,
		HttpOnly: true,
		Secure:   strings.HasPrefix(h.MountURL, "https"),
		MaxAge:   1 * 60 * 60,
	})

	http.Redirect(w, r, redirectURL, http.StatusFound)
}

func (h *Handler) handleEnd(w http.ResponseWriter, r *http.Request) {
	providerName := chi.URLParam(r, "provider")
	provider, ok := h.providers[providerName]
	if !ok {
		http.NotFound(w, r)
		return
	}

	cookie, err := r.Cookie(stateCookie)
	if err != nil {
		h.handleError(w, "failed to get oauth state cookie: %s", err)
		return
	}

	http.SetCookie(w, &http.Cookie{
		Name:     stateCookie,
		Value:    "",
		Path:     h.CookiePath,
		HttpOnly: true,
		MaxAge:   -1,
	})

	state := cookie.Value
	if state == "" {
		h.handleError(w, "empty oauth state cookie")
		return
	}

	queryState := r.URL.Query().Get("state")
	if queryState != state {
		h.handleError(w, "bad state value")
		return
	}

	queryCode := r.URL.Query().Get("code")
	if queryCode == "" {
		h.handleError(w, "empty code value")
		return
	}

	ctx, cancel := context.WithTimeout(r.Context(), clientTimeout)
	defer cancel()

	token, err := provider.config.Exchange(ctx, queryCode)
	if err != nil {
		h.handleError(w, "exchange failed: %s", err)
		return
	}
	if !token.Valid() {
		h.handleError(w, "invalid token")
		return
	}

	u, err := provider.getUser(provider.config.Client(ctx, token))
	if err != nil {
		h.handleError(w, "get provider user: %s", err)
		return
	}

	if u.id == "" {
		h.handleError(w, "provider user id is empty")
		return
	}

	var authToken string

	user, err := h.UserStore.GetByAuth(providerName, u.id)
	switch err {
	case nil:
		if user.Blocked {
			h.renderOAuthResult(w, "error:UserBlocked")
			return
		}

		authToken, err = h.JWTService.Create(user.ID)
		if err != nil {
			h.handleError(w, "failed to create auth token: %s", err)
			return
		}

	case store.ErrNotFound:
		userID, err := h.UserStore.New(providerName, u.id)
		if err != nil {
			h.handleError(w, "failed to create user: %s", err)
			return
		}

		authToken, err = h.JWTService.Create(userID)
		if err != nil {
			h.handleError(w, "failed to create auth token: %s", err)
			return
		}

	default:
		h.handleError(w, "failed to get user by auth: %s", err)
		return
	}

	h.renderOAuthResult(w, "success:"+authToken)
	return
}

func (h *Handler) renderOAuthResult(w http.ResponseWriter, message string) {
	http.SetCookie(w, &http.Cookie{
		Name:   resultCookie,
		Value:  message,
		Path:   h.CookiePath,
		Secure: strings.HasPrefix(h.MountURL, "https"),
		MaxAge: 10 * 60,
	})
	fmt.Fprint(w, `<!doctype html><title>OAuth</title><script>try {opener.bebopOAuthEnd()} finally {window.close()}</script>`)
}

func (h *Handler) handleError(w http.ResponseWriter, format string, a ...interface{}) {
	pc, _, _, _ := runtime.Caller(1)
	callerNameSplit := strings.Split(runtime.FuncForPC(pc).Name(), ".")
	funcName := callerNameSplit[len(callerNameSplit)-1]
	h.Logger.Printf("ERROR: %s: %s", funcName, fmt.Sprintf(format, a...))
	h.renderOAuthResult(w, "error:Other")
}

func (h *Handler) genState() string {
	return uuid.NewV4().String()
}
<commit_msg>oauth: fix routes for chi v3<commit_after>// Package oauth provides an HTTP handler to handle OAuth2
// redirect and callback requests for the bebop web app.
package oauth

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

	"github.com/go-chi/chi"
	"github.com/satori/go.uuid"
	"golang.org/x/oauth2"

	"github.com/disintegration/bebop/jwt"
	"github.com/disintegration/bebop/store"
)

const (
	stateCookie   = "bebop_oauth_state"
	resultCookie  = "bebop_oauth_result"
	clientTimeout = 10 * time.Second
)

// Config is a configuration of an OAuth handler.
type Config struct {
	Logger     *log.Logger
	UserStore  store.UserStore
	JWTService jwt.Service
	MountURL   string
	CookiePath string
}

// Handler handles oauth2 authentication requests.
type Handler struct {
	*Config
	providers map[string]*provider
	router    chi.Router
}

// New creates a new handler based on the given config.
func New(config *Config) *Handler {
	h := &Handler{Config: config}

	h.providers = make(map[string]*provider)

	h.router = chi.NewRouter()
	h.router.Get("/begin/{provider}", h.handleBegin)
	h.router.Get("/end/{provider}", h.handleEnd)

	return h
}

// AddProvider adds a new provider to oauth handler.
func (h *Handler) AddProvider(name, id, secret string) error {
	pc, ok := providerConfigs[name]
	if !ok {
		return fmt.Errorf("oauth: unknown provider: %q", name)
	}

	if id == "" {
		return fmt.Errorf("oauth: empty client id of provider %q", name)
	}

	if secret == "" {
		return fmt.Errorf("oauth: empty client secret of provider %q", name)
	}

	if h.providers == nil {
		h.providers = make(map[string]*provider)
	}

	h.providers[name] = &provider{
		config: &oauth2.Config{
			ClientID:     id,
			ClientSecret: secret,
			RedirectURL:  h.MountURL + "/end/" + name,
			Endpoint:     pc.endpoint,
			Scopes:       pc.scopes,
		},
		getUser: pc.getUser,
	}

	return nil
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h.router.ServeHTTP(w, r)
}

func (h *Handler) handleBegin(w http.ResponseWriter, r *http.Request) {
	providerName := chi.URLParam(r, "provider")
	provider, ok := h.providers[providerName]
	if !ok {
		http.NotFound(w, r)
		return
	}

	state := h.genState()
	redirectURL := provider.config.AuthCodeURL(state)

	http.SetCookie(w, &http.Cookie{
		Name:     stateCookie,
		Value:    state,
		Path:     h.CookiePath,
		HttpOnly: true,
		Secure:   strings.HasPrefix(h.MountURL, "https"),
		MaxAge:   1 * 60 * 60,
	})

	http.Redirect(w, r, redirectURL, http.StatusFound)
}

func (h *Handler) handleEnd(w http.ResponseWriter, r *http.Request) {
	providerName := chi.URLParam(r, "provider")
	provider, ok := h.providers[providerName]
	if !ok {
		http.NotFound(w, r)
		return
	}

	cookie, err := r.Cookie(stateCookie)
	if err != nil {
		h.handleError(w, "failed to get oauth state cookie: %s", err)
		return
	}

	http.SetCookie(w, &http.Cookie{
		Name:     stateCookie,
		Value:    "",
		Path:     h.CookiePath,
		HttpOnly: true,
		MaxAge:   -1,
	})

	state := cookie.Value
	if state == "" {
		h.handleError(w, "empty oauth state cookie")
		return
	}

	queryState := r.URL.Query().Get("state")
	if queryState != state {
		h.handleError(w, "bad state value")
		return
	}

	queryCode := r.URL.Query().Get("code")
	if queryCode == "" {
		h.handleError(w, "empty code value")
		return
	}

	ctx, cancel := context.WithTimeout(r.Context(), clientTimeout)
	defer cancel()

	token, err := provider.config.Exchange(ctx, queryCode)
	if err != nil {
		h.handleError(w, "exchange failed: %s", err)
		return
	}
	if !token.Valid() {
		h.handleError(w, "invalid token")
		return
	}

	u, err := provider.getUser(provider.config.Client(ctx, token))
	if err != nil {
		h.handleError(w, "get provider user: %s", err)
		return
	}

	if u.id == "" {
		h.handleError(w, "provider user id is empty")
		return
	}

	var authToken string

	user, err := h.UserStore.GetByAuth(providerName, u.id)
	switch err {
	case nil:
		if user.Blocked {
			h.renderOAuthResult(w, "error:UserBlocked")
			return
		}

		authToken, err = h.JWTService.Create(user.ID)
		if err != nil {
			h.handleError(w, "failed to create auth token: %s", err)
			return
		}

	case store.ErrNotFound:
		userID, err := h.UserStore.New(providerName, u.id)
		if err != nil {
			h.handleError(w, "failed to create user: %s", err)
			return
		}

		authToken, err = h.JWTService.Create(userID)
		if err != nil {
			h.handleError(w, "failed to create auth token: %s", err)
			return
		}

	default:
		h.handleError(w, "failed to get user by auth: %s", err)
		return
	}

	h.renderOAuthResult(w, "success:"+authToken)
	return
}

func (h *Handler) renderOAuthResult(w http.ResponseWriter, message string) {
	http.SetCookie(w, &http.Cookie{
		Name:   resultCookie,
		Value:  message,
		Path:   h.CookiePath,
		Secure: strings.HasPrefix(h.MountURL, "https"),
		MaxAge: 10 * 60,
	})
	fmt.Fprint(w, `<!doctype html><title>OAuth</title><script>try {opener.bebopOAuthEnd()} finally {window.close()}</script>`)
}

func (h *Handler) handleError(w http.ResponseWriter, format string, a ...interface{}) {
	pc, _, _, _ := runtime.Caller(1)
	callerNameSplit := strings.Split(runtime.FuncForPC(pc).Name(), ".")
	funcName := callerNameSplit[len(callerNameSplit)-1]
	h.Logger.Printf("ERROR: %s: %s", funcName, fmt.Sprintf(format, a...))
	h.renderOAuthResult(w, "error:Other")
}

func (h *Handler) genState() string {
	return uuid.NewV4().String()
}
<|endoftext|>


<commit_before>// Copyright (c) 2017, 2020, Oracle and/or its affiliates. All rights reserved.
// Licensed under the Mozilla Public License v2.0

package oci

import (
	"log"
)

const Version = "3.90.1"

func PrintVersion() {
	log.Printf("[INFO] terraform-provider-oci %s\n", Version)
}
<commit_msg>Finalize changelog and release for version v3.91.0<commit_after>// Copyright (c) 2017, 2020, Oracle and/or its affiliates. All rights reserved.
// Licensed under the Mozilla Public License v2.0

package oci

import (
	"log"
)

const Version = "3.91.0"

func PrintVersion() {
	log.Printf("[INFO] terraform-provider-oci %s\n", Version)
}
<|endoftext|>