<commit_before>package resource_pool

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

	"github.com/dropbox/godropbox/errors"
)

type ResourceLocationPool struct {
	ResourceLocation string
	Pool             ResourcePool
}

func shuffle(pools []*ResourceLocationPool) {
	for i := len(pools) - 1; i > 0; i-- {
		idx := rand.Intn(i + 1)
		pools[i], pools[idx] = pools[idx], pools[i]
	}
}

// A resource pool implementation which returns handles from the registered
// resource locations in a round robin fashion.
type RoundRobinResourcePool struct {
	options Options

	createPool func(Options) ResourcePool

	rwMutex    sync.RWMutex
	isLameDuck bool
	pools      []*ResourceLocationPool

	counter *int64 // atomic counter
}

// This returns a RoundRobinResourcePool.
func NewRoundRobinResourcePool(
	options Options,
	createPool func(Options) ResourcePool,
	pools ...*ResourceLocationPool) (ResourcePool, error) {

	for _, pool := range pools {
		if pool.ResourceLocation == "" {
			return nil, errors.New("Invalid resource location")
		}
		if pool.Pool == nil {
			return nil, errors.New("Invalid pool")
		}
	}

	if createPool == nil {
		createPool = NewSimpleResourcePool
	}

	counter := new(int64)
	atomic.StoreInt64(counter, 0)

	shuffle(pools)

	return &RoundRobinResourcePool{
		options:    options,
		createPool: createPool,
		rwMutex:    sync.RWMutex{},
		isLameDuck: false,
		pools:      pools,
		counter:    counter,
	}, nil
}

// See ResourcePool for documentation.
func (p *RoundRobinResourcePool) NumActive() int32 {
	p.rwMutex.RLock()
	defer p.rwMutex.RUnlock()

	total := int32(0)
	for _, locPool := range p.pools {
		total += locPool.Pool.NumActive()
	}
	return total
}

// See ResourcePool for documentation.
func (p *RoundRobinResourcePool) NumIdle() int {
	p.rwMutex.RLock()
	defer p.rwMutex.RUnlock()

	total := 0
	for _, locPool := range p.pools {
		total += locPool.Pool.NumIdle()
	}
	return total
}

// See ResourcePool for documentation.
func (p *RoundRobinResourcePool) Register(resourceLocation string) error {
	if resourceLocation == "" {
		return errors.New("Registering invalid resource location")
	}

	p.rwMutex.Lock()
	defer p.rwMutex.Unlock()

	if p.isLameDuck {
		return errors.Newf(
			"Cannot register %s to lame duck resource pool",
			resourceLocation)
	}

	for _, locPool := range p.pools {
		if locPool.ResourceLocation == resourceLocation {
			return nil
		}
	}

	pool := p.createPool(p.options)
	if err := pool.Register(resourceLocation); err != nil {
		return err
	}

	p.pools = append(
		p.pools,
		&ResourceLocationPool{
			ResourceLocation: resourceLocation,
			Pool:             pool,
		})

	shuffle(p.pools)
	return nil
}

// See ResourcePool for documentation.
func (p *RoundRobinResourcePool) Unregister(resourceLocation string) error {
	p.rwMutex.Lock()
	defer p.rwMutex.Unlock()

	idx := -1
	for i, locPool := range p.pools {
		if locPool.ResourceLocation == resourceLocation {
			idx = i
			break
		}
	}

	if idx >= 0 {
		tail := p.pools[idx+1:]
		p.pools = p.pools[0:idx]
		p.pools = append(p.pools, tail...)
		shuffle(p.pools)
	}
	return nil
}

func (p *RoundRobinResourcePool) ListRegistered() []string {
	p.rwMutex.RLock()
	defer p.rwMutex.RUnlock()

	result := make([]string, 0, len(p.pools))
	for _, locPool := range p.pools {
		result = append(result, locPool.ResourceLocation)
	}
	return result
}

// See ResourcePool for documentation.
func (p *RoundRobinResourcePool) Get(key string) (ManagedHandle, error) {

	p.rwMutex.RLock()
	defer p.rwMutex.RUnlock()

	var err error
	var handle ManagedHandle

	for i := 0; i < len(p.pools); i++ {
		next := int(atomic.AddInt64(p.counter, 1) % int64(len(p.pools)))
		pool := p.pools[next].Pool

		handle, err = pool.Get(key)
		if err == nil {
			return handle, nil
		}
	}

	return nil, errors.Wrap(err, "No resource handle available")
}

// See ResourcePool for documentation.
func (p *RoundRobinResourcePool) Release(handle ManagedHandle) error {
	// NOTE: check if the handle belongs to this pool is expensive, so we'll
	// just skip the check.
	return handle.Release()
}

// See ResourcePool for documentation.
func (p *RoundRobinResourcePool) Discard(handle ManagedHandle) error {
	// NOTE: check if the handle belongs to this pool is expensive, so we'll
	// just skip the check.
	return handle.Discard()
}

// See ResourcePool for documentation.
func (p *RoundRobinResourcePool) EnterLameDuckMode() {
	p.rwMutex.RLock()
	defer p.rwMutex.RUnlock()

	p.isLameDuck = true

	for _, locPool := range p.pools {
		locPool.Pool.EnterLameDuckMode()
	}
}
<commit_msg>dedup check<commit_after>package resource_pool

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

	"github.com/dropbox/godropbox/errors"
)

type ResourceLocationPool struct {
	ResourceLocation string
	Pool             ResourcePool
}

func shuffle(pools []*ResourceLocationPool) {
	for i := len(pools) - 1; i > 0; i-- {
		idx := rand.Intn(i + 1)
		pools[i], pools[idx] = pools[idx], pools[i]
	}
}

// A resource pool implementation which returns handles from the registered
// resource locations in a round robin fashion.
type RoundRobinResourcePool struct {
	options Options

	createPool func(Options) ResourcePool

	rwMutex    sync.RWMutex
	isLameDuck bool
	pools      []*ResourceLocationPool

	counter *int64 // atomic counter
}

// This returns a RoundRobinResourcePool.
func NewRoundRobinResourcePool(
	options Options,
	createPool func(Options) ResourcePool,
	pools ...*ResourceLocationPool) (ResourcePool, error) {

	locations := make(map[string]bool)

	for _, pool := range pools {
		if pool.ResourceLocation == "" {
			return nil, errors.New("Invalid resource location")
		}

		if locations[pool.ResourceLocation] {
			return nil, errors.Newf(
				"Duplication resource location %s",
				pool.ResourceLocation)
		}
		locations[pool.ResourceLocation] = true

		if pool.Pool == nil {
			return nil, errors.New("Invalid pool")
		}
	}

	if createPool == nil {
		createPool = NewSimpleResourcePool
	}

	counter := new(int64)
	atomic.StoreInt64(counter, 0)

	shuffle(pools)

	return &RoundRobinResourcePool{
		options:    options,
		createPool: createPool,
		rwMutex:    sync.RWMutex{},
		isLameDuck: false,
		pools:      pools,
		counter:    counter,
	}, nil
}

// See ResourcePool for documentation.
func (p *RoundRobinResourcePool) NumActive() int32 {
	p.rwMutex.RLock()
	defer p.rwMutex.RUnlock()

	total := int32(0)
	for _, locPool := range p.pools {
		total += locPool.Pool.NumActive()
	}
	return total
}

// See ResourcePool for documentation.
func (p *RoundRobinResourcePool) NumIdle() int {
	p.rwMutex.RLock()
	defer p.rwMutex.RUnlock()

	total := 0
	for _, locPool := range p.pools {
		total += locPool.Pool.NumIdle()
	}
	return total
}

// See ResourcePool for documentation.
func (p *RoundRobinResourcePool) Register(resourceLocation string) error {
	if resourceLocation == "" {
		return errors.New("Registering invalid resource location")
	}

	p.rwMutex.Lock()
	defer p.rwMutex.Unlock()

	if p.isLameDuck {
		return errors.Newf(
			"Cannot register %s to lame duck resource pool",
			resourceLocation)
	}

	for _, locPool := range p.pools {
		if locPool.ResourceLocation == resourceLocation {
			return nil
		}
	}

	pool := p.createPool(p.options)
	if err := pool.Register(resourceLocation); err != nil {
		return err
	}

	p.pools = append(
		p.pools,
		&ResourceLocationPool{
			ResourceLocation: resourceLocation,
			Pool:             pool,
		})

	shuffle(p.pools)
	return nil
}

// See ResourcePool for documentation.
func (p *RoundRobinResourcePool) Unregister(resourceLocation string) error {
	p.rwMutex.Lock()
	defer p.rwMutex.Unlock()

	idx := -1
	for i, locPool := range p.pools {
		if locPool.ResourceLocation == resourceLocation {
			idx = i
			break
		}
	}

	if idx >= 0 {
		tail := p.pools[idx+1:]
		p.pools = p.pools[0:idx]
		p.pools = append(p.pools, tail...)
		shuffle(p.pools)
	}
	return nil
}

func (p *RoundRobinResourcePool) ListRegistered() []string {
	p.rwMutex.RLock()
	defer p.rwMutex.RUnlock()

	result := make([]string, 0, len(p.pools))
	for _, locPool := range p.pools {
		result = append(result, locPool.ResourceLocation)
	}
	return result
}

// See ResourcePool for documentation.
func (p *RoundRobinResourcePool) Get(key string) (ManagedHandle, error) {

	p.rwMutex.RLock()
	defer p.rwMutex.RUnlock()

	var err error
	var handle ManagedHandle

	for i := 0; i < len(p.pools); i++ {
		next := int(atomic.AddInt64(p.counter, 1) % int64(len(p.pools)))
		pool := p.pools[next].Pool

		handle, err = pool.Get(key)
		if err == nil {
			return handle, nil
		}
	}

	return nil, errors.Wrap(err, "No resource handle available")
}

// See ResourcePool for documentation.
func (p *RoundRobinResourcePool) Release(handle ManagedHandle) error {
	// NOTE: check if the handle belongs to this pool is expensive, so we'll
	// just skip the check.
	return handle.Release()
}

// See ResourcePool for documentation.
func (p *RoundRobinResourcePool) Discard(handle ManagedHandle) error {
	// NOTE: check if the handle belongs to this pool is expensive, so we'll
	// just skip the check.
	return handle.Discard()
}

// See ResourcePool for documentation.
func (p *RoundRobinResourcePool) EnterLameDuckMode() {
	p.rwMutex.RLock()
	defer p.rwMutex.RUnlock()

	p.isLameDuck = true

	for _, locPool := range p.pools {
		locPool.Pool.EnterLameDuckMode()
	}
}
<|endoftext|>


<commit_before>package agent

import (
	"bytes"
	"fmt"
	"net"
	"os"
	"path"
	"strings"

	"github.com/pkg/errors"

	"github.com/google/uuid"

	"github.com/havoc-io/mutagen"
	"github.com/havoc-io/mutagen/environment"
	"github.com/havoc-io/mutagen/filesystem"
	"github.com/havoc-io/mutagen/ssh"
	"github.com/havoc-io/mutagen/url"
)

var sshAgentPath string

func init() {
	// Compute the agent SSH command.
	// HACK: This assumes that the SSH user's home directory is used as the
	// default working directory for SSH commands. We have to do this because we
	// don't have a portable mechanism to invoke the command relative to the
	// user's home directory (tilde doesn't work on Windows) and we don't want
	// to do a probe of the remote system before invoking the endpoint. This
	// assumption should be fine for 99.9% of cases, but if it becomes a major
	// issue, the only other options I see are probing before invoking (slow) or
	// using the Go SSH library to do this (painful to faithfully emulate
	// OpenSSH's behavior). Perhaps probing could be hidden behind an option?
	// HACK: We're assuming that none of these path components have spaces in
	// them, but since we control all of them, this is probably okay.
	// HACK: When invoking on Windows systems, we can use forward slashes for
	// the path and leave the "exe" suffix off the target name. This saves us a
	// target check.
	sshAgentPath = path.Join(
		filesystem.MutagenDirectoryName,
		agentsDirectoryName,
		mutagen.Version,
		agentBaseName,
	)
}

func probeSSHPOSIX(remote *url.URL, prompter string) (string, string, error) {
	// Try to invoke uname and print kernel and machine name.
	unameSMBytes, err := ssh.Output(prompter, "Probing endpoint", remote, "uname -s -m")
	if err != nil {
		return "", "", errors.Wrap(err, "unable to invoke uname")
	}

	// Parse uname output.
	unameSM := strings.Split(strings.TrimSpace(string(unameSMBytes)), " ")
	if len(unameSM) != 2 {
		return "", "", errors.New("invalid uname output")
	}
	unameS := unameSM[0]
	unameM := unameSM[1]

	// Translate GOOS.
	var goos string
	if unameSIsWindowsPosix(unameS) {
		goos = "windows"
	} else if g, ok := unameSToGOOS[unameS]; ok {
		goos = g
	} else {
		return "", "", errors.New("unknown platform")
	}

	// Translate GOARCH.
	goarch, ok := unameMToGOARCH[unameM]
	if !ok {
		return "", "", errors.New("unknown architecture")
	}

	// Success.
	return goos, goarch, nil
}

func probeSSHWindows(remote *url.URL, prompter string) (string, string, error) {
	// Try to print the remote environment.
	envBytes, err := ssh.Output(prompter, "Probing endpoint", remote, "cmd /c set")
	if err != nil {
		return "", "", errors.Wrap(err, "unable to invoke set")
	}

	// Parse set output.
	env, err := environment.ParseBlock(string(envBytes))
	if err != nil {
		return "", "", errors.Wrap(err, "unable to parse environment")
	}

	// Translate GOOS.
	goos, ok := osEnvToGOOS[env["OS"]]
	if !ok {
		return "", "", errors.New("unknown platform")
	}

	// Translate GOARCH.
	goarch, ok := processorArchitectureEnvToGOARCH[env["PROCESSOR_ARCHITECTURE"]]
	if !ok {
		return "", "", errors.New("unknown architecture")
	}

	// Success.
	return goos, goarch, nil
}

// probeSSHPlatform attempts to identify the properties of the target platform,
// namely GOOS, GOARCH, and whether or not it's a POSIX environment (which it
// might be even on Windows).
func probeSSHPlatform(remote *url.URL, prompter string) (string, string, bool, error) {
	// Attempt to probe for a POSIX platform. This might apply to certain
	// Windows environments as well.
	if goos, goarch, err := probeSSHPOSIX(remote, prompter); err == nil {
		return goos, goarch, true, nil
	}

	// If that fails, attempt a Windows fallback.
	if goos, goarch, err := probeSSHWindows(remote, prompter); err == nil {
		return goos, goarch, false, nil
	}

	// Failure.
	return "", "", false, errors.New("exhausted probing methods")
}

func installSSH(remote *url.URL, prompter string) error {
	// Detect the target platform.
	goos, goarch, posix, err := probeSSHPlatform(remote, prompter)
	if err != nil {
		return errors.Wrap(err, "unable to probe remote platform")
	}

	// Find the appropriate agent binary. Ensure that it's cleaned up when we're
	// done with it.
	agent, err := executableForPlatform(goos, goarch)
	if err != nil {
		return errors.Wrap(err, "unable to get agent for platform")
	}
	defer os.Remove(agent)

	// Copy the agent to the remote. We use a unique identifier for the
	// temporary destination. For Windows remotes, we add a ".exe" suffix, which
	// will automatically make the file executable on the remote (POSIX systems
	// are handled separately below). For POSIX systems, we add a dot prefix to
	// hide the executable a bit.
	// HACK: This assumes that the SSH user's home directory is used as the
	// default destination directory for SCP copies. That should be true in
	// 99.9% of cases, but if it becomes a major issue, we'll need to use the
	// probe information to handle this more carefully.
	randomUUID, err := uuid.NewRandom()
	if err != nil {
		return errors.Wrap(err, "unable to generate UUID for agent copying")
	}
	destination := agentBaseName + randomUUID.String()
	if goos == "windows" {
		destination += ".exe"
	}
	if posix {
		destination = "." + destination
	}
	destinationURL := &url.URL{
		Protocol: remote.Protocol,
		Username: remote.Username,
		Hostname: remote.Hostname,
		Port:     remote.Port,
		Path:     destination,
	}
	if err := ssh.Copy(prompter, "Copying agent", agent, destinationURL); err != nil {
		return errors.Wrap(err, "unable to copy agent binary")
	}

	// Invoke the remote installation. For POSIX remotes, we have to incorporate
	// a "chmod +x" in order for the remote to execute the installer. The POSIX
	// solution is necessary in the event that an installer is sent from a
	// Windows to a POSIX system using SCP, since there's no way to preserve the
	// executability bit (Windows doesn't have one). This will also be applied
	// to Windows POSIX environments, but a "chmod +x" there will have no
	// effect.
	// HACK: This assumes that the SSH user's home directory is used as the
	// default working directory for SSH commands. We have to do this because we
	// don't have a portable mechanism to invoke the command relative to the
	// user's home directory and we don't want to do a probe of the remote
	// system before invoking the endpoint. This assumption should be fine for
	// 99.9% of cases, but if it becomes a major issue, we'll need to use the
	// probe information to handle this more carefully.
	var installCommand string
	if posix {
		installCommand = fmt.Sprintf("chmod +x %s && ./%s install", destination, destination)
	} else {
		installCommand = fmt.Sprintf("%s install", destination)
	}
	if err := ssh.Run(prompter, "Installing agent", remote, installCommand); err != nil {
		return errors.Wrap(err, "unable to invoke agent installation")
	}

	// Success.
	return nil
}

func connectSSH(remote *url.URL, prompter, mode string) (net.Conn, bool, error) {
	// Compute the command to invoke.
	// HACK: We rely on sshAgentPath not having any spaces in it. If we do
	// eventually need to add any, we'll need to fix this up for the shell.
	command := fmt.Sprintf("%s %s", sshAgentPath, mode)

	// Create an SSH process.
	process, err := ssh.Command(prompter, "Connecting to agent", remote, command)
	if err != nil {
		return nil, false, errors.Wrap(err, "unable to create SSH command")
	}

	// Create a connection that wrap's the process' standard input/output.
	connection, err := newAgentConnection(remote, process)
	if err != nil {
		return nil, false, errors.Wrap(err, "unable to create SSH process connection")
	}

	// Redirect the process' standard error output to a buffer so that we can
	// give better feedback in errors. This might be a bit dangerous since this
	// buffer will be attached for the lifetime of the process and we don't know
	// exactly how much output will be received (and thus we could buffer a
	// large amount of it in memory), but generally speaking SSH doens't spit
	// out much error output (unless in debug mode, which we won't be), and the
	// agent doesn't spit out any.
	// TODO: If we do start seeing large allocations in these buffers, a simple
	// size-limited buffer might suffice, at least to get some of the error
	// message.
	// TODO: If we decide we want these errors available outside the agent
	// package, it might be worth moving this buffer into the processStream
	// type, exporting that type, and allowing type assertions that would give
	// access to that buffer. But for now we're mostly just concerned with
	// connection issues.
	errorBuffer := bytes.NewBuffer(nil)
	process.Stderr = errorBuffer

	// Start the process.
	if err = process.Start(); err != nil {
		return nil, false, errors.Wrap(err, "unable to start SSH agent process")
	}

	// Confirm that the process started correctly by performing a version
	// handshake.
	if versionMatch, err := mutagen.ReceiveAndCompareVersion(connection); err != nil {
		// If there's an error, check if SSH exits with a command not found
		// error. We can't really check this until we try to interact with the
		// process and see that it misbehaves. We wouldn't be able to see this
		// returned as an error from the Start method because it just starts the
		// SSH client itself, not the remote command.
		if ssh.IsCommandNotFound(process.Wait()) {
			return nil, true, errors.New("command not found")
		}

		// Otherwise, check if there is any error output that might illuminate
		// what happened. We let this overrule any err value here since that
		// value will probably just be an EOF.
		if errorBuffer.Len() > 0 {
			return nil, false, errors.Errorf(
				"SSH process failed with error output:\n%s",
				strings.TrimSpace(errorBuffer.String()),
			)
		}

		// Otherwise just wrap up whatever error we have.
		return nil, false, errors.Wrap(err, "unable to handshake with SSH agent process")
	} else if !versionMatch {
		return nil, true, errors.New("version mismatch")
	}

	// Done.
	return connection, false, nil
}

func DialSSH(remote *url.URL, prompter, mode string) (net.Conn, error) {
	// Attempt a connection. If this fails, but it's a failure that justfies
	// attempting an install, then continue, otherwise fail.
	if connection, install, err := connectSSH(remote, prompter, mode); err == nil {
		return connection, nil
	} else if !install {
		return nil, errors.Wrap(err, "unable to connect to agent")
	}

	// Attempt to install.
	if err := installSSH(remote, prompter); err != nil {
		return nil, errors.Wrap(err, "unable to install agent")
	}

	// Re-attempt connectivity.
	if connection, _, err := connectSSH(remote, prompter, mode); err != nil {
		return nil, errors.Wrap(err, "unable to connect to agent")
	} else {
		return connection, nil
	}
}
<commit_msg>Decomposed agent install and reduced overhead in some cases.<commit_after>package agent

import (
	"bytes"
	"fmt"
	"net"
	"os"
	"path"
	"runtime"
	"strings"

	"github.com/pkg/errors"

	"github.com/google/uuid"

	"github.com/havoc-io/mutagen"
	"github.com/havoc-io/mutagen/environment"
	"github.com/havoc-io/mutagen/filesystem"
	"github.com/havoc-io/mutagen/ssh"
	"github.com/havoc-io/mutagen/url"
)

var sshAgentPath string

func init() {
	// Compute the agent SSH command.
	// HACK: This assumes that the SSH user's home directory is used as the
	// default working directory for SSH commands. We have to do this because we
	// don't have a portable mechanism to invoke the command relative to the
	// user's home directory (tilde doesn't work on Windows) and we don't want
	// to do a probe of the remote system before invoking the endpoint. This
	// assumption should be fine for 99.9% of cases, but if it becomes a major
	// issue, the only other options I see are probing before invoking (slow) or
	// using the Go SSH library to do this (painful to faithfully emulate
	// OpenSSH's behavior). Perhaps probing could be hidden behind an option?
	// HACK: We're assuming that none of these path components have spaces in
	// them, but since we control all of them, this is probably okay.
	// HACK: When invoking on Windows systems, we can use forward slashes for
	// the path and leave the "exe" suffix off the target name. This saves us a
	// target check.
	sshAgentPath = path.Join(
		filesystem.MutagenDirectoryName,
		agentsDirectoryName,
		mutagen.Version,
		agentBaseName,
	)
}

func probeSSHPOSIX(remote *url.URL, prompter string) (string, string, error) {
	// Try to invoke uname and print kernel and machine name.
	unameSMBytes, err := ssh.Output(prompter, "Probing endpoint", remote, "uname -s -m")
	if err != nil {
		return "", "", errors.Wrap(err, "unable to invoke uname")
	}

	// Parse uname output.
	unameSM := strings.Split(strings.TrimSpace(string(unameSMBytes)), " ")
	if len(unameSM) != 2 {
		return "", "", errors.New("invalid uname output")
	}
	unameS := unameSM[0]
	unameM := unameSM[1]

	// Translate GOOS.
	var goos string
	if unameSIsWindowsPosix(unameS) {
		goos = "windows"
	} else if g, ok := unameSToGOOS[unameS]; ok {
		goos = g
	} else {
		return "", "", errors.New("unknown platform")
	}

	// Translate GOARCH.
	goarch, ok := unameMToGOARCH[unameM]
	if !ok {
		return "", "", errors.New("unknown architecture")
	}

	// Success.
	return goos, goarch, nil
}

func probeSSHWindows(remote *url.URL, prompter string) (string, string, error) {
	// Try to print the remote environment.
	envBytes, err := ssh.Output(prompter, "Probing endpoint", remote, "cmd /c set")
	if err != nil {
		return "", "", errors.Wrap(err, "unable to invoke set")
	}

	// Parse set output.
	env, err := environment.ParseBlock(string(envBytes))
	if err != nil {
		return "", "", errors.Wrap(err, "unable to parse environment")
	}

	// Translate GOOS.
	goos, ok := osEnvToGOOS[env["OS"]]
	if !ok {
		return "", "", errors.New("unknown platform")
	}

	// Translate GOARCH.
	goarch, ok := processorArchitectureEnvToGOARCH[env["PROCESSOR_ARCHITECTURE"]]
	if !ok {
		return "", "", errors.New("unknown architecture")
	}

	// Success.
	return goos, goarch, nil
}

// probeSSHPlatform attempts to identify the properties of the target platform,
// namely GOOS, GOARCH, and whether or not it's a POSIX environment (which it
// might be even on Windows).
func probeSSHPlatform(remote *url.URL, prompter string) (string, string, bool, error) {
	// Attempt to probe for a POSIX platform. This might apply to certain
	// Windows environments as well.
	if goos, goarch, err := probeSSHPOSIX(remote, prompter); err == nil {
		return goos, goarch, true, nil
	}

	// If that fails, attempt a Windows fallback.
	if goos, goarch, err := probeSSHWindows(remote, prompter); err == nil {
		return goos, goarch, false, nil
	}

	// Failure.
	return "", "", false, errors.New("exhausted probing methods")
}

func installSSH(remote *url.URL, prompter string) error {
	// Detect the target platform.
	goos, goarch, posix, err := probeSSHPlatform(remote, prompter)
	if err != nil {
		return errors.Wrap(err, "unable to probe remote platform")
	}

	// Find the appropriate agent binary. Ensure that it's cleaned up when we're
	// done with it.
	agent, err := executableForPlatform(goos, goarch)
	if err != nil {
		return errors.Wrap(err, "unable to get agent for platform")
	}
	defer os.Remove(agent)

	// Copy the agent to the remote. We use a unique identifier for the
	// temporary destination. For Windows remotes, we add a ".exe" suffix, which
	// will automatically make the file executable on the remote (POSIX systems
	// are handled separately below). For POSIX systems, we add a dot prefix to
	// hide the executable a bit.
	// HACK: This assumes that the SSH user's home directory is used as the
	// default destination directory for SCP copies. That should be true in
	// 99.9% of cases, but if it becomes a major issue, we'll need to use the
	// probe information to handle this more carefully.
	randomUUID, err := uuid.NewRandom()
	if err != nil {
		return errors.Wrap(err, "unable to generate UUID for agent copying")
	}
	destination := agentBaseName + randomUUID.String()
	if goos == "windows" {
		destination += ".exe"
	}
	if posix {
		destination = "." + destination
	}
	destinationURL := &url.URL{
		Protocol: remote.Protocol,
		Username: remote.Username,
		Hostname: remote.Hostname,
		Port:     remote.Port,
		Path:     destination,
	}
	if err := ssh.Copy(prompter, "Copying agent", agent, destinationURL); err != nil {
		return errors.Wrap(err, "unable to copy agent binary")
	}

	// For cases where we're copying from a Windows system to a POSIX remote,
	// invoke "chmod +x" to add executability back to the copied binary. This is
	// necessary under the specified circumstances because as soon as the agent
	// binary is extracted from the bundle, it will lose its executability bit
	// since Windows can't preserve this. This will also be applied to Windows
	// POSIX remotes, but a "chmod +x" there will just be a no-op.
	if runtime.GOOS == "windows" && posix {
		executabilityCommand := fmt.Sprintf("chmod +x %s", destination)
		if err := ssh.Run(prompter, "Setting agent executability", remote, executabilityCommand); err != nil {
			return errors.Wrap(err, "unable to set agent executability")
		}
	}

	// Invoke the remote installation.
	// HACK: This assumes that the SSH user's home directory is used as the
	// default working directory for SSH commands. We have to do this because we
	// don't have a portable mechanism to invoke the command relative to the
	// user's home directory and we don't want to do a probe of the remote
	// system before invoking the endpoint. This assumption should be fine for
	// 99.9% of cases, but if it becomes a major issue, we'll need to use the
	// probe information to handle this more carefully.
	var installCommand string
	if posix {
		installCommand = fmt.Sprintf("./%s install", destination, destination)
	} else {
		installCommand = fmt.Sprintf("%s install", destination)
	}
	if err := ssh.Run(prompter, "Installing agent", remote, installCommand); err != nil {
		return errors.Wrap(err, "unable to invoke agent installation")
	}

	// Success.
	return nil
}

func connectSSH(remote *url.URL, prompter, mode string) (net.Conn, bool, error) {
	// Compute the command to invoke.
	// HACK: We rely on sshAgentPath not having any spaces in it. If we do
	// eventually need to add any, we'll need to fix this up for the shell.
	command := fmt.Sprintf("%s %s", sshAgentPath, mode)

	// Create an SSH process.
	process, err := ssh.Command(prompter, "Connecting to agent", remote, command)
	if err != nil {
		return nil, false, errors.Wrap(err, "unable to create SSH command")
	}

	// Create a connection that wrap's the process' standard input/output.
	connection, err := newAgentConnection(remote, process)
	if err != nil {
		return nil, false, errors.Wrap(err, "unable to create SSH process connection")
	}

	// Redirect the process' standard error output to a buffer so that we can
	// give better feedback in errors. This might be a bit dangerous since this
	// buffer will be attached for the lifetime of the process and we don't know
	// exactly how much output will be received (and thus we could buffer a
	// large amount of it in memory), but generally speaking SSH doens't spit
	// out much error output (unless in debug mode, which we won't be), and the
	// agent doesn't spit out any.
	// TODO: If we do start seeing large allocations in these buffers, a simple
	// size-limited buffer might suffice, at least to get some of the error
	// message.
	// TODO: If we decide we want these errors available outside the agent
	// package, it might be worth moving this buffer into the processStream
	// type, exporting that type, and allowing type assertions that would give
	// access to that buffer. But for now we're mostly just concerned with
	// connection issues.
	errorBuffer := bytes.NewBuffer(nil)
	process.Stderr = errorBuffer

	// Start the process.
	if err = process.Start(); err != nil {
		return nil, false, errors.Wrap(err, "unable to start SSH agent process")
	}

	// Confirm that the process started correctly by performing a version
	// handshake.
	if versionMatch, err := mutagen.ReceiveAndCompareVersion(connection); err != nil {
		// If there's an error, check if SSH exits with a command not found
		// error. We can't really check this until we try to interact with the
		// process and see that it misbehaves. We wouldn't be able to see this
		// returned as an error from the Start method because it just starts the
		// SSH client itself, not the remote command.
		if ssh.IsCommandNotFound(process.Wait()) {
			return nil, true, errors.New("command not found")
		}

		// Otherwise, check if there is any error output that might illuminate
		// what happened. We let this overrule any err value here since that
		// value will probably just be an EOF.
		if errorBuffer.Len() > 0 {
			return nil, false, errors.Errorf(
				"SSH process failed with error output:\n%s",
				strings.TrimSpace(errorBuffer.String()),
			)
		}

		// Otherwise just wrap up whatever error we have.
		return nil, false, errors.Wrap(err, "unable to handshake with SSH agent process")
	} else if !versionMatch {
		return nil, true, errors.New("version mismatch")
	}

	// Done.
	return connection, false, nil
}

func DialSSH(remote *url.URL, prompter, mode string) (net.Conn, error) {
	// Attempt a connection. If this fails, but it's a failure that justfies
	// attempting an install, then continue, otherwise fail.
	if connection, install, err := connectSSH(remote, prompter, mode); err == nil {
		return connection, nil
	} else if !install {
		return nil, errors.Wrap(err, "unable to connect to agent")
	}

	// Attempt to install.
	if err := installSSH(remote, prompter); err != nil {
		return nil, errors.Wrap(err, "unable to install agent")
	}

	// Re-attempt connectivity.
	if connection, _, err := connectSSH(remote, prompter, mode); err != nil {
		return nil, errors.Wrap(err, "unable to connect to agent")
	} else {
		return connection, nil
	}
}
<|endoftext|>


<commit_before>package httpexpect

import (
	"bytes"
	"github.com/stretchr/testify/assert"
	"github.com/valyala/fasthttp"
	"io/ioutil"
	"net/http"
	"testing"
)

type mockHandler struct {
	t *testing.T
}

func (c mockHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	body, err := ioutil.ReadAll(req.Body)
	assert.True(c.t, err == nil)

	assert.Equal(c.t, "GET", req.Method)
	assert.Equal(c.t, "http://example.com", req.URL.String())
	assert.Equal(c.t, "body", string(body))

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	_, err = w.Write([]byte(`{"hello":"world"}`))

	assert.True(c.t, err == nil)
}

func TestBinder(t *testing.T) {
	binder := NewBinder(mockHandler{t})

	req, err := http.NewRequest(
		"GET", "http://example.com", bytes.NewReader([]byte("body")))

	if err != nil {
		t.Fatal(err)
	}

	resp, err := binder.Do(req)
	if err != nil {
		t.Fatal(err)
	}

	header := http.Header{
		"Content-Type": {"application/json"},
	}

	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	assert.Equal(t, header, resp.Header)
	assert.Equal(t, `{"hello":"world"}`, string(b))
}

func TestFastBinder(t *testing.T) {
	binder := NewFastBinder(func(ctx *fasthttp.RequestCtx) {
		assert.Equal(t, "POST", string(ctx.Request.Header.Method()))
		assert.Equal(t, "http://example.com", string(ctx.Request.Header.RequestURI()))
		assert.Equal(t, "application/x-www-form-urlencoded",
			string(ctx.Request.Header.ContentType()))

		assert.Equal(t, "bar", string(ctx.FormValue("foo")))
		assert.Equal(t, "foo=bar", string(ctx.Request.Body()))

		ctx.Response.Header.Set("Content-Type", "application/json")
		ctx.Response.SetBody([]byte(`{"hello":"world"}`))
	})

	req, err := http.NewRequest(
		"POST", "http://example.com", bytes.NewReader([]byte("foo=bar")))

	if err != nil {
		t.Fatal(err)
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := binder.Do(req)
	if err != nil {
		t.Fatal(err)
	}

	header := http.Header{
		"Content-Type": {"application/json"},
	}

	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	assert.Equal(t, header, resp.Header)
	assert.Equal(t, `{"hello":"world"}`, string(b))
}
<commit_msg>Improve test coverage<commit_after>package httpexpect

import (
	"bytes"
	"github.com/stretchr/testify/assert"
	"github.com/valyala/fasthttp"
	"io/ioutil"
	"net/http"
	"testing"
)

type mockHandler struct {
	t *testing.T
}

func (c mockHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	body, err := ioutil.ReadAll(req.Body)
	assert.True(c.t, err == nil)

	assert.Equal(c.t, "GET", req.Method)
	assert.Equal(c.t, "http://example.com", req.URL.String())
	assert.Equal(c.t, "body", string(body))

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	_, err = w.Write([]byte(`{"hello":"world"}`))

	assert.True(c.t, err == nil)
}

func TestBinder(t *testing.T) {
	binder := NewBinder(mockHandler{t})

	req, err := http.NewRequest(
		"GET", "http://example.com", bytes.NewReader([]byte("body")))

	if err != nil {
		t.Fatal(err)
	}

	resp, err := binder.Do(req)
	if err != nil {
		t.Fatal(err)
	}

	header := http.Header{
		"Content-Type": {"application/json"},
	}

	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	assert.Equal(t, header, resp.Header)
	assert.Equal(t, `{"hello":"world"}`, string(b))
}

func TestFastBinder(t *testing.T) {
	binder := NewFastBinder(func(ctx *fasthttp.RequestCtx) {
		assert.Equal(t, "POST", string(ctx.Request.Header.Method()))
		assert.Equal(t, "http://example.com", string(ctx.Request.Header.RequestURI()))

		assert.Equal(t, "application/x-www-form-urlencoded",
			string(ctx.Request.Header.ContentType()))

		headers := map[string][]string{}

		ctx.Request.Header.VisitAll(func(k, v []byte) {
			headers[string(k)] = append(headers[string(k)], string(v))
		})

		expected := map[string][]string{
			"Content-Type": {"application/x-www-form-urlencoded"},
			"Some-Header":  {"foo", "bar"},
		}

		assert.Equal(t, expected, headers)

		assert.Equal(t, "bar", string(ctx.FormValue("foo")))
		assert.Equal(t, "foo=bar", string(ctx.Request.Body()))

		ctx.Response.Header.Set("Content-Type", "application/json")
		ctx.Response.SetBody([]byte(`{"hello":"world"}`))
	})

	req, err := http.NewRequest(
		"POST", "http://example.com", bytes.NewReader([]byte("foo=bar")))

	if err != nil {
		t.Fatal(err)
	}

	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.Header.Add("Some-Header", "foo")
	req.Header.Add("Some-Header", "bar")

	resp, err := binder.Do(req)
	if err != nil {
		t.Fatal(err)
	}

	header := http.Header{
		"Content-Type": {"application/json"},
	}

	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		t.Fatal(err)
	}

	assert.Equal(t, header, resp.Header)
	assert.Equal(t, `{"hello":"world"}`, string(b))
}
<|endoftext|>


<commit_before>// Copyright 2017 Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package util

import (
	"io/ioutil"
	"log"
	"os"
	"path"
	"regexp"
	"strings"

	"github.com/google/go-github/github"
)

var (
	ci = NewCIState()
	// SHARegex matches commit SHA's
	SHARegex = regexp.MustCompile("^[a-z0-9]{40}$")
	// ReleaseTagRegex matches release tags
	ReleaseTagRegex = regexp.MustCompile("^[0-9]+.[0-9]+.[0-9]+$")
)

// CIState defines constants representing possible states of
// continuous integration tests
type CIState struct {
	Success string
	Failure string
	Pending string
	Error   string
}

// NewCIState creates a new CIState
func NewCIState() *CIState {
	return &CIState{
		Success: "success",
		Failure: "failure",
		Pending: "pending",
		Error:   "error",
	}
}

// GetCIState does NOT trust the given combined output but instead walk
// through the CI results, count states, and determine the final state
// as either pending, failure, or success
func GetCIState(combinedStatus *github.CombinedStatus, skipContext func(string) bool) string {
	return GetReqquiredCIState(combinedStatus, nil, skipContext)
}

// GetReqquiredCIState does NOT trust the given combined output but instead walk
// through the CI results, count states, and determine the final state
// as either pending, failure, or success
func GetReqquiredCIState(combinedStatus *github.CombinedStatus,
	requiredChecks *github.RequiredStatusChecks,
	skipContext func(string) bool) string {
	var failures, pending, successes int
	for _, status := range combinedStatus.Statuses {
		if requiredChecks != nil &&
			!IsRequiredCICheck(status.GetContext(), requiredChecks) {
			continue
		}
		if *status.State == ci.Error || *status.State == ci.Failure {
			if skipContext != nil && skipContext(*status.Context) {
				continue
			}
			log.Printf("%s\t failed", status.GetContext())
			failures++
		} else if *status.State == ci.Pending {
			log.Printf("%s\t pending", status.GetContext())
			pending++
		} else if *status.State == ci.Success {
			log.Printf("%s\t passed", status.GetContext())
			successes++
		} else {
			log.Printf("Check Status %s is unknown", *status.State)
		}
	}
	if pending > 0 {
		return ci.Pending
	} else if failures > 0 {
		return ci.Failure
	} else {
		return ci.Success
	}
}

// IsRequiredCICheck returns true if the check is required to pass before a PR
// can be merged.
// statusCxt is the name of the status
func IsRequiredCICheck(statusCxt string,
	requiredChecks *github.RequiredStatusChecks) bool {
	if requiredChecks == nil {
		return false
	}
	for _, requiredCheckCxt := range requiredChecks.Contexts {
		if statusCxt == requiredCheckCxt {
			return true
		}
	}
	return false
}

// GetAPITokenFromFile returns the github api token from tokenFile
func GetAPITokenFromFile(tokenFile string) (string, error) {
	return GetPasswordFromFile(tokenFile)
}

// GetPasswordFromFile get a string usually is a password or token from a local file
func GetPasswordFromFile(file string) (string, error) {
	b, err := ioutil.ReadFile(file)
	if err != nil {
		return "", err
	}
	token := strings.TrimSpace(string(b[:]))
	return token, nil
}

// CloneRepoCheckoutBranch removes previous repo, clone to local machine,
// change directory into the repo, and checkout the given branch.
// Returns the absolute path to repo root
func CloneRepoCheckoutBranch(gclient *GithubClient, repo, baseBranch, newBranch, pathPrefix string) (string, error) {
	if err := os.MkdirAll(pathPrefix, os.FileMode(0755)); err != nil {
		return "", err
	}
	repoPath := path.Join(pathPrefix, repo)
	if err := os.RemoveAll(repoPath); err != nil {
		return "", err
	}
	if err := os.Chdir(pathPrefix); err != nil {
		return "", err
	}
	if _, err := ShellSilent(
		"git clone " + gclient.Remote(repo)); err != nil {
		return "", err
	}
	if err := os.Chdir(repo); err != nil {
		return "", err
	}
	if _, err := Shell("git checkout " + baseBranch); err != nil {
		return "", err
	}
	if newBranch != "" {
		if _, err := Shell("git checkout -b " + newBranch); err != nil {
			return "", err
		}
	}
	return os.Getwd()
}

// RemoveLocalRepo deletes the local git repo just cloned
func RemoveLocalRepo(pathToRepo string) error {
	return os.RemoveAll(pathToRepo)
}

// CreateCommitPushToRemote stages call local changes, create a commit,
// and push to remote tracking branch
func CreateCommitPushToRemote(branch, commitMsg string) error {
	// git commit -am does not work with untracked files
	// track new files first and then create a commit
	if _, err := Shell("git add -A"); err != nil {
		return err
	}
	if _, err := Shell("git commit -m " + commitMsg); err != nil {
		return err
	}
	_, err := Shell("git push -f --set-upstream origin " + branch)
	return err
}
<commit_msg>Fix regression by PR #651 (#670)<commit_after>// Copyright 2017 Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package util

import (
	"io/ioutil"
	"log"
	"os"
	"path"
	"regexp"
	"strings"

	"github.com/google/go-github/github"
)

var (
	ci = NewCIState()
	// SHARegex matches commit SHA's
	SHARegex = regexp.MustCompile("^[a-z0-9]{40}$")
	// ReleaseTagRegex matches release tags
	ReleaseTagRegex = regexp.MustCompile("^[0-9]+.[0-9]+.[0-9]+$")
)

// CIState defines constants representing possible states of
// continuous integration tests
type CIState struct {
	Success string
	Failure string
	Pending string
	Error   string
}

// NewCIState creates a new CIState
func NewCIState() *CIState {
	return &CIState{
		Success: "success",
		Failure: "failure",
		Pending: "pending",
		Error:   "error",
	}
}

// GetCIState does NOT trust the given combined output but instead walk
// through the CI results, count states, and determine the final state
// as either pending, failure, or success
func GetCIState(combinedStatus *github.CombinedStatus, skipContext func(string) bool) string {
	return GetReqquiredCIState(combinedStatus, nil, skipContext)
}

// GetReqquiredCIState does NOT trust the given combined output but instead walk
// through the CI results, count states, and determine the final state
// as either pending, failure, or success
func GetReqquiredCIState(combinedStatus *github.CombinedStatus,
	requiredChecks *github.RequiredStatusChecks,
	skipContext func(string) bool) string {
	var failures, pending, successes int
	for _, status := range combinedStatus.Statuses {
		if requiredChecks != nil &&
			!IsRequiredCICheck(status.GetContext(), requiredChecks) {
			continue
		}
		if *status.State == ci.Error || *status.State == ci.Failure {
			if skipContext != nil && skipContext(*status.Context) {
				continue
			}
			log.Printf("%s\t failed", status.GetContext())
			failures++
		} else if *status.State == ci.Pending {
			log.Printf("%s\t pending", status.GetContext())
			pending++
		} else if *status.State == ci.Success {
			log.Printf("%s\t passed", status.GetContext())
			successes++
		} else {
			log.Printf("Check Status %s is unknown", *status.State)
		}
	}
	if pending > 0 {
		return ci.Pending
	} else if failures > 0 {
		return ci.Failure
	} else {
		return ci.Success
	}
}

// IsRequiredCICheck returns true if the check is required to pass before a PR
// can be merged.
// statusCxt is the name of the status
func IsRequiredCICheck(statusCxt string,
	requiredChecks *github.RequiredStatusChecks) bool {
	if requiredChecks == nil {
		return false
	}
	for _, requiredCheckCxt := range requiredChecks.Contexts {
		if statusCxt == requiredCheckCxt {
			return true
		}
	}
	return false
}

// GetAPITokenFromFile returns the github api token from tokenFile
func GetAPITokenFromFile(tokenFile string) (string, error) {
	return GetPasswordFromFile(tokenFile)
}

// GetPasswordFromFile get a string usually is a password or token from a local file
func GetPasswordFromFile(file string) (string, error) {
	b, err := ioutil.ReadFile(file)
	if err != nil {
		return "", err
	}
	token := strings.TrimSpace(string(b[:]))
	return token, nil
}

// CloneRepoCheckoutBranch removes previous repo, clone to local machine,
// change directory into the repo, and checkout the given branch.
// Returns the absolute path to repo root
func CloneRepoCheckoutBranch(gclient *GithubClient, repo, baseBranch, newBranch, pathPrefix string) (string, error) {
	repoPath := path.Join(pathPrefix, repo)
	if err := os.RemoveAll(repoPath); err != nil {
		return "", err
	}
	if pathPrefix != "" {
		if err := os.MkdirAll(pathPrefix, os.FileMode(0755)); err != nil {
			return "", err
		}
		if err := os.Chdir(pathPrefix); err != nil {
			return "", err
		}
	}
	if _, err := ShellSilent(
		"git clone " + gclient.Remote(repo)); err != nil {
		return "", err
	}
	if err := os.Chdir(repo); err != nil {
		return "", err
	}
	if _, err := Shell("git checkout " + baseBranch); err != nil {
		return "", err
	}
	if newBranch != "" {
		if _, err := Shell("git checkout -b " + newBranch); err != nil {
			return "", err
		}
	}
	return os.Getwd()
}

// RemoveLocalRepo deletes the local git repo just cloned
func RemoveLocalRepo(pathToRepo string) error {
	return os.RemoveAll(pathToRepo)
}

// CreateCommitPushToRemote stages call local changes, create a commit,
// and push to remote tracking branch
func CreateCommitPushToRemote(branch, commitMsg string) error {
	// git commit -am does not work with untracked files
	// track new files first and then create a commit
	if _, err := Shell("git add -A"); err != nil {
		return err
	}
	if _, err := Shell("git commit -m " + commitMsg); err != nil {
		return err
	}
	_, err := Shell("git push -f --set-upstream origin " + branch)
	return err
}
<|endoftext|>


<commit_before>package cors

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"net/http"
	"net/url"
	"strings"
)

var (
	validSites          = map[string]bool{}
	defaultAllowHeaders = []string{"Origin", "Accept", "Content-Type", "Authorization", "ib_id"}
	defaultAllowMethods []string
)

// CORS will set the headers for Cross-origin resource sharing
func CORS() gin.HandlerFunc {
	return func(c *gin.Context) {

		req := c.Request
		origin := req.Header.Get("Origin")

		// Set origin header from sites config
		if isAllowedSite(origin) {
			c.Header("Access-Control-Allow-Origin", origin)
		} else {
			c.Header("Access-Control-Allow-Origin", "")
		}

		c.Header("Vary", "Origin")

		c.Header("Access-Control-Allow-Credentials", "true")

		if req.Method == "OPTIONS" {

			// Add allowed method header
			c.Header("Access-Control-Allow-Methods", strings.Join(defaultAllowMethods, ","))

			// Add allowed headers header
			c.Header("Access-Control-Allow-Headers", strings.Join(defaultAllowHeaders, ","))

			c.Header("Access-Control-Max-Age", "600")

			c.AbortWithStatus(http.StatusOK)

			return

		} else {

			c.Next()

		}

	}
}

func SetDomains(domains, methods []string) {
	// add valid sites to map
	for _, site := range domains {
		validSites[site] = true
	}

	// set methods
	defaultAllowMethods = methods

	fmt.Println(strings.Repeat("*", 60))
	fmt.Printf("%-20v\n\n", "CORS")
	fmt.Printf("%-20v%40v\n", "Domains", strings.Join(domains, ", "))
	fmt.Printf("%-20v%40v\n", "Methods", strings.Join(defaultAllowMethods, ", "))
	fmt.Println(strings.Repeat("*", 60))

	return
}

// Check if origin is allowed
func isAllowedSite(host string) bool {

	// Get the host from the origin
	parsed, err := url.Parse(host)
	if err != nil {
		return false
	}

	if validSites[strings.ToLower(parsed.Host)] {
		return true
	}

	return false

}
<commit_msg>add max-age to cors headers<commit_after>package cors

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"net/http"
	"net/url"
	"strings"
)

var (
	validSites          = map[string]bool{}
	defaultAllowHeaders = []string{"Origin", "Accept", "Content-Type", "Authorization", "ib_id"}
	defaultAllowMethods []string
)

// CORS will set the headers for Cross-origin resource sharing
func CORS() gin.HandlerFunc {
	return func(c *gin.Context) {

		req := c.Request
		origin := req.Header.Get("Origin")

		// Set origin header from sites config
		if isAllowedSite(origin) {
			c.Header("Access-Control-Allow-Origin", origin)
		} else {
			c.Header("Access-Control-Allow-Origin", "")
		}

		c.Header("Vary", "Origin")

		c.Header("Access-Control-Allow-Credentials", "true")

		if req.Method == "OPTIONS" {

			// Add allowed method header
			c.Header("Access-Control-Allow-Methods", strings.Join(defaultAllowMethods, ","))

			// Add allowed headers header
			c.Header("Access-Control-Allow-Headers", strings.Join(defaultAllowHeaders, ","))

			c.Header("Access-Control-Max-Age", "3600")

			c.AbortWithStatus(http.StatusOK)

			return

		} else {

			c.Next()

		}

	}
}

func SetDomains(domains, methods []string) {
	// add valid sites to map
	for _, site := range domains {
		validSites[site] = true
	}

	// set methods
	defaultAllowMethods = methods

	fmt.Println(strings.Repeat("*", 60))
	fmt.Printf("%-20v\n\n", "CORS")
	fmt.Printf("%-20v%40v\n", "Domains", strings.Join(domains, ", "))
	fmt.Printf("%-20v%40v\n", "Methods", strings.Join(defaultAllowMethods, ", "))
	fmt.Println(strings.Repeat("*", 60))

	return
}

// Check if origin is allowed
func isAllowedSite(host string) bool {

	// Get the host from the origin
	parsed, err := url.Parse(host)
	if err != nil {
		return false
	}

	if validSites[strings.ToLower(parsed.Host)] {
		return true
	}

	return false

}
<|endoftext|>


<commit_before>// Copyright 2014 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cors

import (
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/codegangsta/negroni"
)

type HttpHeaderGuardRecorder struct {
	*httptest.ResponseRecorder
	savedHeaderMap http.Header
}

func NewRecorder() *HttpHeaderGuardRecorder {
	return &HttpHeaderGuardRecorder{httptest.NewRecorder(), nil}
}

func (gr *HttpHeaderGuardRecorder) WriteHeader(code int) {
	gr.ResponseRecorder.WriteHeader(code)
	gr.savedHeaderMap = gr.ResponseRecorder.Header()
}

func (gr *HttpHeaderGuardRecorder) Header() http.Header {
	if gr.savedHeaderMap != nil {
		// headers were written. clone so we don't get updates
		clone := make(http.Header)
		for k, v := range gr.savedHeaderMap {
			clone[k] = v
		}
		return clone
	} else {
		return gr.ResponseRecorder.Header()
	}
}

func Test_AllowAll(t *testing.T) {
	recorder := httptest.NewRecorder()
	n := negroni.New()
	opts := &Options{
		AllowAllOrigins: true,
	}
	n.Use(negroni.HandlerFunc(opts.Allow))

	r, _ := http.NewRequest("PUT", "foo", nil)
	n.ServeHTTP(recorder, r)

	if recorder.HeaderMap.Get(headerAllowOrigin) != "*" {
		t.Errorf("Allow-Origin header should be *")
	}
}

func Test_AllowRegexMatch(t *testing.T) {
	recorder := httptest.NewRecorder()
	n := negroni.New()
	opts := &Options{
		AllowOrigins: []string{"https://aaa.com", "https://*.foo.com"},
	}
	n.Use(negroni.HandlerFunc(opts.Allow))

	origin := "https://bar.foo.com"
	r, _ := http.NewRequest("PUT", "foo", nil)
	r.Header.Add("Origin", origin)
	n.ServeHTTP(recorder, r)

	headerValue := recorder.HeaderMap.Get(headerAllowOrigin)
	if headerValue != origin {
		t.Errorf("Allow-Origin header should be %v, found %v", origin, headerValue)
	}
}

func Test_AllowRegexNoMatch(t *testing.T) {
	recorder := httptest.NewRecorder()
	n := negroni.New()
	opts := &Options{
		AllowOrigins: []string{"https://*.foo.com"},
	}
	n.Use(negroni.HandlerFunc(opts.Allow))

	origin := "https://ww.foo.com.evil.com"
	r, _ := http.NewRequest("PUT", "foo", nil)
	r.Header.Add("Origin", origin)
	n.ServeHTTP(recorder, r)

	headerValue := recorder.HeaderMap.Get(headerAllowOrigin)
	if headerValue != "" {
		t.Errorf("Allow-Origin header should not exist, found %v", headerValue)
	}
}

func Test_OtherHeaders(t *testing.T) {
	recorder := httptest.NewRecorder()
	n := negroni.New()
	opts := &Options{
		AllowAllOrigins:  true,
		AllowCredentials: true,
		AllowMethods:     []string{"PATCH", "GET"},
		AllowHeaders:     []string{"Origin", "X-whatever"},
		ExposeHeaders:    []string{"Content-Length", "Hello"},
		MaxAge:           5 * time.Minute,
	}
	n.Use(negroni.HandlerFunc(opts.Allow))

	r, _ := http.NewRequest("PUT", "foo", nil)
	n.ServeHTTP(recorder, r)

	credentialsVal := recorder.HeaderMap.Get(headerAllowCredentials)
	methodsVal := recorder.HeaderMap.Get(headerAllowMethods)
	headersVal := recorder.HeaderMap.Get(headerAllowHeaders)
	exposedHeadersVal := recorder.HeaderMap.Get(headerExposeHeaders)
	maxAgeVal := recorder.HeaderMap.Get(headerMaxAge)

	if credentialsVal != "true" {
		t.Errorf("Allow-Credentials is expected to be true, found %v", credentialsVal)
	}

	if methodsVal != "PATCH,GET" {
		t.Errorf("Allow-Methods is expected to be PATCH,GET; found %v", methodsVal)
	}

	if headersVal != "Origin,X-whatever" {
		t.Errorf("Allow-Headers is expected to be Origin,X-whatever; found %v", headersVal)
	}

	if exposedHeadersVal != "Content-Length,Hello" {
		t.Errorf("Expose-Headers are expected to be Content-Length,Hello. Found %v", exposedHeadersVal)
	}

	if maxAgeVal != "300" {
		t.Errorf("Max-Age is expected to be 300, found %v", maxAgeVal)
	}
}

func Test_DefaultAllowHeaders(t *testing.T) {
	recorder := httptest.NewRecorder()
	n := negroni.New()
	opts := &Options{
		AllowAllOrigins: true,
	}
	n.Use(negroni.HandlerFunc(opts.Allow))

	r, _ := http.NewRequest("PUT", "foo", nil)
	n.ServeHTTP(recorder, r)

	headersVal := recorder.HeaderMap.Get(headerAllowHeaders)
	if headersVal != "Origin,Accept,Content-Type,Authorization" {
		t.Errorf("Allow-Headers is expected to be Origin,Accept,Content-Type,Authorization; found %v", headersVal)
	}
}

/*
func Test_Preflight(t *testing.T) {
	recorder := NewRecorder()
	n := negroni.New()
	opts := &Options{
		AllowAllOrigins: true,
		AllowMethods:    []string{"PUT", "PATCH"},
		AllowHeaders:    []string{"Origin", "X-whatever", "X-CaseSensitive"},
	}
	n.Use(negroni.HandlerFunc(opts.Allow))

	// TODO: Rewrite handler in negroni
	m.Options("foo", func(res http.ResponseWriter) {
		res.WriteHeader(500)
	})

	r, _ := http.NewRequest("OPTIONS", "foo", nil)
	r.Header.Add(headerRequestMethod, "PUT")
	r.Header.Add(headerRequestHeaders, "X-whatever, x-casesensitive")
	n.ServeHTTP(recorder, r)

	headers := recorder.Header()
	methodsVal := headers.Get(headerAllowMethods)
	headersVal := headers.Get(headerAllowHeaders)
	originVal := headers.Get(headerAllowOrigin)

	if methodsVal != "PUT,PATCH" {
		t.Errorf("Allow-Methods is expected to be PUT,PATCH, found %v", methodsVal)
	}

	if !strings.Contains(headersVal, "X-whatever") {
		t.Errorf("Allow-Headers is expected to contain X-whatever, found %v", headersVal)
	}

	if !strings.Contains(headersVal, "x-casesensitive") {
		t.Errorf("Allow-Headers is expected to contain x-casesensitive, found %v", headersVal)
	}

	if originVal != "*" {
		t.Errorf("Allow-Origin is expected to be *, found %v", originVal)
	}

	if recorder.Code != http.StatusOK {
		t.Errorf("Status code is expected to be 200, found %d", recorder.Code)
	}
}
*/
func Benchmark_WithoutCORS(b *testing.B) {
	recorder := httptest.NewRecorder()
	n := negroni.New()

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		r, _ := http.NewRequest("PUT", "foo", nil)
		n.ServeHTTP(recorder, r)
	}
}

func Benchmark_WithCORS(b *testing.B) {
	recorder := httptest.NewRecorder()
	n := negroni.New()
	opts := &Options{
		AllowAllOrigins:  true,
		AllowCredentials: true,
		AllowMethods:     []string{"PATCH", "GET"},
		AllowHeaders:     []string{"Origin", "X-whatever"},
		MaxAge:           5 * time.Minute,
	}
	n.Use(negroni.HandlerFunc(opts.Allow))

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		r, _ := http.NewRequest("PUT", "foo", nil)
		n.ServeHTTP(recorder, r)
	}
}
<commit_msg>Update Preflight test to negroni<commit_after>// Copyright 2014 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cors

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/codegangsta/negroni"
)

type HttpHeaderGuardRecorder struct {
	*httptest.ResponseRecorder
	savedHeaderMap http.Header
}

func NewRecorder() *HttpHeaderGuardRecorder {
	return &HttpHeaderGuardRecorder{httptest.NewRecorder(), nil}
}

func (gr *HttpHeaderGuardRecorder) WriteHeader(code int) {
	gr.ResponseRecorder.WriteHeader(code)
	gr.savedHeaderMap = gr.ResponseRecorder.Header()
}

func (gr *HttpHeaderGuardRecorder) Header() http.Header {
	if gr.savedHeaderMap != nil {
		// headers were written. clone so we don't get updates
		clone := make(http.Header)
		for k, v := range gr.savedHeaderMap {
			clone[k] = v
		}
		return clone
	} else {
		return gr.ResponseRecorder.Header()
	}
}

func Test_AllowAll(t *testing.T) {
	recorder := httptest.NewRecorder()
	n := negroni.New()
	opts := &Options{
		AllowAllOrigins: true,
	}
	n.Use(negroni.HandlerFunc(opts.Allow))

	r, _ := http.NewRequest("PUT", "foo", nil)
	n.ServeHTTP(recorder, r)

	if recorder.HeaderMap.Get(headerAllowOrigin) != "*" {
		t.Errorf("Allow-Origin header should be *")
	}
}

func Test_AllowRegexMatch(t *testing.T) {
	recorder := httptest.NewRecorder()
	n := negroni.New()
	opts := &Options{
		AllowOrigins: []string{"https://aaa.com", "https://*.foo.com"},
	}
	n.Use(negroni.HandlerFunc(opts.Allow))

	origin := "https://bar.foo.com"
	r, _ := http.NewRequest("PUT", "foo", nil)
	r.Header.Add("Origin", origin)
	n.ServeHTTP(recorder, r)

	headerValue := recorder.HeaderMap.Get(headerAllowOrigin)
	if headerValue != origin {
		t.Errorf("Allow-Origin header should be %v, found %v", origin, headerValue)
	}
}

func Test_AllowRegexNoMatch(t *testing.T) {
	recorder := httptest.NewRecorder()
	n := negroni.New()
	opts := &Options{
		AllowOrigins: []string{"https://*.foo.com"},
	}
	n.Use(negroni.HandlerFunc(opts.Allow))

	origin := "https://ww.foo.com.evil.com"
	r, _ := http.NewRequest("PUT", "foo", nil)
	r.Header.Add("Origin", origin)
	n.ServeHTTP(recorder, r)

	headerValue := recorder.HeaderMap.Get(headerAllowOrigin)
	if headerValue != "" {
		t.Errorf("Allow-Origin header should not exist, found %v", headerValue)
	}
}

func Test_OtherHeaders(t *testing.T) {
	recorder := httptest.NewRecorder()
	n := negroni.New()
	opts := &Options{
		AllowAllOrigins:  true,
		AllowCredentials: true,
		AllowMethods:     []string{"PATCH", "GET"},
		AllowHeaders:     []string{"Origin", "X-whatever"},
		ExposeHeaders:    []string{"Content-Length", "Hello"},
		MaxAge:           5 * time.Minute,
	}
	n.Use(negroni.HandlerFunc(opts.Allow))

	r, _ := http.NewRequest("PUT", "foo", nil)
	n.ServeHTTP(recorder, r)

	credentialsVal := recorder.HeaderMap.Get(headerAllowCredentials)
	methodsVal := recorder.HeaderMap.Get(headerAllowMethods)
	headersVal := recorder.HeaderMap.Get(headerAllowHeaders)
	exposedHeadersVal := recorder.HeaderMap.Get(headerExposeHeaders)
	maxAgeVal := recorder.HeaderMap.Get(headerMaxAge)

	if credentialsVal != "true" {
		t.Errorf("Allow-Credentials is expected to be true, found %v", credentialsVal)
	}

	if methodsVal != "PATCH,GET" {
		t.Errorf("Allow-Methods is expected to be PATCH,GET; found %v", methodsVal)
	}

	if headersVal != "Origin,X-whatever" {
		t.Errorf("Allow-Headers is expected to be Origin,X-whatever; found %v", headersVal)
	}

	if exposedHeadersVal != "Content-Length,Hello" {
		t.Errorf("Expose-Headers are expected to be Content-Length,Hello. Found %v", exposedHeadersVal)
	}

	if maxAgeVal != "300" {
		t.Errorf("Max-Age is expected to be 300, found %v", maxAgeVal)
	}
}

func Test_DefaultAllowHeaders(t *testing.T) {
	recorder := httptest.NewRecorder()
	n := negroni.New()
	opts := &Options{
		AllowAllOrigins: true,
	}
	n.Use(negroni.HandlerFunc(opts.Allow))

	r, _ := http.NewRequest("PUT", "foo", nil)
	n.ServeHTTP(recorder, r)

	headersVal := recorder.HeaderMap.Get(headerAllowHeaders)
	if headersVal != "Origin,Accept,Content-Type,Authorization" {
		t.Errorf("Allow-Headers is expected to be Origin,Accept,Content-Type,Authorization; found %v", headersVal)
	}
}

func Test_Preflight(t *testing.T) {
	recorder := NewRecorder()
	n := negroni.New()
	opts := &Options{
		AllowAllOrigins: true,
		AllowMethods:    []string{"PUT", "PATCH"},
		AllowHeaders:    []string{"Origin", "X-whatever", "X-CaseSensitive"},
	}
	n.Use(negroni.HandlerFunc(opts.Allow))
	mux := http.NewServeMux()
	mux.HandleFunc("/foo", func(w http.ResponseWriter, req *http.Request) {
		if req.Method == "OPTIONS" {
			w.WriteHeader(500)
			return
		}
		return
	})
	n.UseHandler(mux)

	r, _ := http.NewRequest("OPTIONS", "/foo", nil)
	r.Header.Add(headerRequestMethod, "PUT")
	r.Header.Add(headerRequestHeaders, "X-whatever, x-casesensitive")
	n.ServeHTTP(recorder, r)

	headers := recorder.Header()
	methodsVal := headers.Get(headerAllowMethods)
	headersVal := headers.Get(headerAllowHeaders)
	originVal := headers.Get(headerAllowOrigin)

	if methodsVal != "PUT,PATCH" {
		t.Errorf("Allow-Methods is expected to be PUT,PATCH, found %v", methodsVal)
	}

	if !strings.Contains(headersVal, "X-whatever") {
		t.Errorf("Allow-Headers is expected to contain X-whatever, found %v", headersVal)
	}

	if !strings.Contains(headersVal, "x-casesensitive") {
		t.Errorf("Allow-Headers is expected to contain x-casesensitive, found %v", headersVal)
	}

	if originVal != "*" {
		t.Errorf("Allow-Origin is expected to be *, found %v", originVal)
	}

	if recorder.Code != http.StatusOK {
		t.Errorf("Status code is expected to be 200, found %d", recorder.Code)
	}
}

func Benchmark_WithoutCORS(b *testing.B) {
	recorder := httptest.NewRecorder()
	n := negroni.New()

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		r, _ := http.NewRequest("PUT", "foo", nil)
		n.ServeHTTP(recorder, r)
	}
}

func Benchmark_WithCORS(b *testing.B) {
	recorder := httptest.NewRecorder()
	n := negroni.New()
	opts := &Options{
		AllowAllOrigins:  true,
		AllowCredentials: true,
		AllowMethods:     []string{"PATCH", "GET"},
		AllowHeaders:     []string{"Origin", "X-whatever"},
		MaxAge:           5 * time.Minute,
	}
	n.Use(negroni.HandlerFunc(opts.Allow))

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		r, _ := http.NewRequest("PUT", "foo", nil)
		n.ServeHTTP(recorder, r)
	}
}
<|endoftext|>


<commit_before>package ldap

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/sirupsen/logrus"
	ldapv2 "gopkg.in/ldap.v2"

	"github.com/pkg/errors"
	"github.com/rancher/norman/httperror"
	"github.com/rancher/types/apis/management.cattle.io/v3"
)

func NewLDAPConn(servers []string, TLS bool, port int64, connectionTimeout int64, caPool *x509.CertPool) (*ldapv2.Conn, error) {
	logrus.Debug("Now creating Ldap connection")
	var lConn *ldapv2.Conn
	var err error
	var tlsConfig *tls.Config
	ldapv2.DefaultTimeout = time.Duration(connectionTimeout) * time.Millisecond
	// TODO implment multi-server support
	if len(servers) != 1 {
		return nil, errors.New("invalid server config. only exactly 1 server is currently supported")
	}
	server := servers[0]
	if TLS {
		tlsConfig = &tls.Config{RootCAs: caPool, InsecureSkipVerify: false, ServerName: server}
		lConn, err = ldapv2.DialTLS("tcp", fmt.Sprintf("%s:%d", server, port), tlsConfig)
		if err != nil {
			return nil, fmt.Errorf("Error creating ssl connection: %v", err)
		}
	} else {
		lConn, err = ldapv2.Dial("tcp", fmt.Sprintf("%s:%d", server, port))
		if err != nil {
			return nil, fmt.Errorf("Error creating connection: %v", err)
		}
	}

	lConn.SetTimeout(time.Duration(connectionTimeout) * time.Millisecond)

	return lConn, nil
}

func GetUserExternalID(username string, loginDomain string) string {
	if strings.Contains(username, "\\") {
		return username
	} else if loginDomain != "" {
		return loginDomain + "\\" + username
	}
	return username
}

func HasPermission(attributes []*ldapv2.EntryAttribute, userObjectClass string, userEnabledAttribute string, userDisabledBitMask int64) bool {
	var permission int64
	if !IsType(attributes, userObjectClass) {
		return true
	}

	if userEnabledAttribute != "" {
		for _, attr := range attributes {
			if attr.Name == userEnabledAttribute {
				if len(attr.Values) > 0 && attr.Values[0] != "" {
					intAttr, err := strconv.ParseInt(attr.Values[0], 10, 64)
					if err != nil {
						logrus.Errorf("Failed to get USER_ENABLED_ATTRIBUTE, error: %v", err)
						return false
					}
					permission = intAttr
				}
			}
		}
	} else {
		return true
	}
	permission = permission & userDisabledBitMask
	return permission != userDisabledBitMask
}

func IsType(search []*ldapv2.EntryAttribute, varType string) bool {
	for _, attrib := range search {
		if attrib.Name == "objectClass" {
			for _, val := range attrib.Values {
				if val == varType {
					return true
				}
			}
		}
	}
	logrus.Debugf("Failed to determine if object is type: %s", varType)
	return false
}

func EscapeLDAPSearchFilter(filter string) string {
	buf := new(bytes.Buffer)
	for i := 0; i < len(filter); i++ {
		currChar := filter[i]
		switch currChar {
		case '\\':
			buf.WriteString("\\5c")
			break
		case '*':
			buf.WriteString("\\2a")
			break
		case '(':
			buf.WriteString("\\28")
			break
		case ')':
			buf.WriteString("\\29")
			break
		case '\u0000':
			buf.WriteString("\\00")
			break
		default:
			buf.WriteString(string(currChar))
		}
	}
	return buf.String()
}

func GetUserSearchAttributes(memberOfAttribute, objectClassAttribute string, config *v3.ActiveDirectoryConfig) []string {
	srchAttributes := strings.Split(config.UserSearchAttribute, "|")
	userSearchAttributes := []string{memberOfAttribute,
		objectClassAttribute,
		config.UserObjectClass,
		config.UserLoginAttribute,
		config.UserNameAttribute,
		config.UserEnabledAttribute}
	userSearchAttributes = append(userSearchAttributes, srchAttributes...)

	return userSearchAttributes
}

func GetGroupSearchAttributes(memberOfAttribute, objectClassAttribute string, config *v3.ActiveDirectoryConfig) []string {
	groupSeachAttributes := []string{memberOfAttribute,
		objectClassAttribute,
		config.GroupObjectClass,
		config.UserLoginAttribute,
		config.GroupNameAttribute,
		config.GroupSearchAttribute}
	return groupSeachAttributes
}

func GetUserSearchAttributesForLDAP(config *v3.LdapConfig) []string {
	userSearchAttributes := []string{"dn", config.UserMemberAttribute,
		"objectClass",
		config.UserObjectClass,
		config.UserLoginAttribute,
		config.UserNameAttribute,
		config.UserEnabledAttribute}
	return userSearchAttributes
}

func GetGroupSearchAttributesForLDAP(config *v3.LdapConfig) []string {
	groupSeachAttributes := []string{config.GroupMemberUserAttribute,
		config.GroupMemberMappingAttribute,
		"objectClass",
		config.GroupObjectClass,
		config.UserLoginAttribute,
		config.GroupNameAttribute,
		config.GroupSearchAttribute}
	return groupSeachAttributes
}

func AuthenticateServiceAccountUser(serviceAccountPassword string, serviceAccountUsername string, defaultLoginDomain string, lConn *ldapv2.Conn) error {
	logrus.Debug("Binding service account username password")
	if serviceAccountPassword == "" {
		return httperror.NewAPIError(httperror.MissingRequired, "service account password not provided")
	}
	sausername := GetUserExternalID(serviceAccountUsername, defaultLoginDomain)
	err := lConn.Bind(sausername, serviceAccountPassword)
	if err != nil {
		if ldapv2.IsErrorWithCode(err, ldapv2.LDAPResultInvalidCredentials) {
			return httperror.WrapAPIError(err, httperror.Unauthorized, "authentication failed")
		}
		return httperror.WrapAPIError(err, httperror.ServerError, "server error while authenticating")
	}

	return nil
}

func Min(a int, b int) int {
	if a < b {
		return a
	}
	return b
}
<commit_msg>delete unnecessary user search attributes for AD<commit_after>package ldap

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/sirupsen/logrus"
	ldapv2 "gopkg.in/ldap.v2"

	"github.com/pkg/errors"
	"github.com/rancher/norman/httperror"
	"github.com/rancher/types/apis/management.cattle.io/v3"
)

func NewLDAPConn(servers []string, TLS bool, port int64, connectionTimeout int64, caPool *x509.CertPool) (*ldapv2.Conn, error) {
	logrus.Debug("Now creating Ldap connection")
	var lConn *ldapv2.Conn
	var err error
	var tlsConfig *tls.Config
	ldapv2.DefaultTimeout = time.Duration(connectionTimeout) * time.Millisecond
	// TODO implment multi-server support
	if len(servers) != 1 {
		return nil, errors.New("invalid server config. only exactly 1 server is currently supported")
	}
	server := servers[0]
	if TLS {
		tlsConfig = &tls.Config{RootCAs: caPool, InsecureSkipVerify: false, ServerName: server}
		lConn, err = ldapv2.DialTLS("tcp", fmt.Sprintf("%s:%d", server, port), tlsConfig)
		if err != nil {
			return nil, fmt.Errorf("Error creating ssl connection: %v", err)
		}
	} else {
		lConn, err = ldapv2.Dial("tcp", fmt.Sprintf("%s:%d", server, port))
		if err != nil {
			return nil, fmt.Errorf("Error creating connection: %v", err)
		}
	}

	lConn.SetTimeout(time.Duration(connectionTimeout) * time.Millisecond)

	return lConn, nil
}

func GetUserExternalID(username string, loginDomain string) string {
	if strings.Contains(username, "\\") {
		return username
	} else if loginDomain != "" {
		return loginDomain + "\\" + username
	}
	return username
}

func HasPermission(attributes []*ldapv2.EntryAttribute, userObjectClass string, userEnabledAttribute string, userDisabledBitMask int64) bool {
	var permission int64
	if !IsType(attributes, userObjectClass) {
		return true
	}

	if userEnabledAttribute != "" {
		for _, attr := range attributes {
			if attr.Name == userEnabledAttribute {
				if len(attr.Values) > 0 && attr.Values[0] != "" {
					intAttr, err := strconv.ParseInt(attr.Values[0], 10, 64)
					if err != nil {
						logrus.Errorf("Failed to get USER_ENABLED_ATTRIBUTE, error: %v", err)
						return false
					}
					permission = intAttr
				}
			}
		}
	} else {
		return true
	}
	permission = permission & userDisabledBitMask
	return permission != userDisabledBitMask
}

func IsType(search []*ldapv2.EntryAttribute, varType string) bool {
	for _, attrib := range search {
		if attrib.Name == "objectClass" {
			for _, val := range attrib.Values {
				if val == varType {
					return true
				}
			}
		}
	}
	logrus.Debugf("Failed to determine if object is type: %s", varType)
	return false
}

func EscapeLDAPSearchFilter(filter string) string {
	buf := new(bytes.Buffer)
	for i := 0; i < len(filter); i++ {
		currChar := filter[i]
		switch currChar {
		case '\\':
			buf.WriteString("\\5c")
			break
		case '*':
			buf.WriteString("\\2a")
			break
		case '(':
			buf.WriteString("\\28")
			break
		case ')':
			buf.WriteString("\\29")
			break
		case '\u0000':
			buf.WriteString("\\00")
			break
		default:
			buf.WriteString(string(currChar))
		}
	}
	return buf.String()
}

func GetUserSearchAttributes(memberOfAttribute, objectClassAttribute string, config *v3.ActiveDirectoryConfig) []string {
	userSearchAttributes := []string{memberOfAttribute,
		objectClassAttribute,
		config.UserObjectClass,
		config.UserLoginAttribute,
		config.UserNameAttribute,
		config.UserEnabledAttribute}
	return userSearchAttributes
}

func GetGroupSearchAttributes(memberOfAttribute, objectClassAttribute string, config *v3.ActiveDirectoryConfig) []string {
	groupSeachAttributes := []string{memberOfAttribute,
		objectClassAttribute,
		config.GroupObjectClass,
		config.UserLoginAttribute,
		config.GroupNameAttribute,
		config.GroupSearchAttribute}
	return groupSeachAttributes
}

func GetUserSearchAttributesForLDAP(config *v3.LdapConfig) []string {
	userSearchAttributes := []string{"dn", config.UserMemberAttribute,
		"objectClass",
		config.UserObjectClass,
		config.UserLoginAttribute,
		config.UserNameAttribute,
		config.UserEnabledAttribute}
	return userSearchAttributes
}

func GetGroupSearchAttributesForLDAP(config *v3.LdapConfig) []string {
	groupSeachAttributes := []string{config.GroupMemberUserAttribute,
		config.GroupMemberMappingAttribute,
		"objectClass",
		config.GroupObjectClass,
		config.UserLoginAttribute,
		config.GroupNameAttribute,
		config.GroupSearchAttribute}
	return groupSeachAttributes
}

func AuthenticateServiceAccountUser(serviceAccountPassword string, serviceAccountUsername string, defaultLoginDomain string, lConn *ldapv2.Conn) error {
	logrus.Debug("Binding service account username password")
	if serviceAccountPassword == "" {
		return httperror.NewAPIError(httperror.MissingRequired, "service account password not provided")
	}
	sausername := GetUserExternalID(serviceAccountUsername, defaultLoginDomain)
	err := lConn.Bind(sausername, serviceAccountPassword)
	if err != nil {
		if ldapv2.IsErrorWithCode(err, ldapv2.LDAPResultInvalidCredentials) {
			return httperror.WrapAPIError(err, httperror.Unauthorized, "authentication failed")
		}
		return httperror.WrapAPIError(err, httperror.ServerError, "server error while authenticating")
	}

	return nil
}

func Min(a int, b int) int {
	if a < b {
		return a
	}
	return b
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package filters

import (
	"context"
	"fmt"
	"net/http"
	"runtime"
	"sync/atomic"

	flowcontrol "k8s.io/api/flowcontrol/v1beta2"
	apitypes "k8s.io/apimachinery/pkg/types"
	epmetrics "k8s.io/apiserver/pkg/endpoints/metrics"
	apirequest "k8s.io/apiserver/pkg/endpoints/request"
	"k8s.io/apiserver/pkg/server/httplog"
	utilflowcontrol "k8s.io/apiserver/pkg/util/flowcontrol"
	fcmetrics "k8s.io/apiserver/pkg/util/flowcontrol/metrics"
	flowcontrolrequest "k8s.io/apiserver/pkg/util/flowcontrol/request"
	"k8s.io/klog/v2"
)

// PriorityAndFairnessClassification identifies the results of
// classification for API Priority and Fairness
type PriorityAndFairnessClassification struct {
	FlowSchemaName    string
	FlowSchemaUID     apitypes.UID
	PriorityLevelName string
	PriorityLevelUID  apitypes.UID
}

// waitingMark tracks requests waiting rather than being executed
var waitingMark = &requestWatermark{
	phase:            epmetrics.WaitingPhase,
	readOnlyObserver: fcmetrics.ReadWriteConcurrencyObserverPairGenerator.Generate(1, 1, []string{epmetrics.ReadOnlyKind}).RequestsWaiting,
	mutatingObserver: fcmetrics.ReadWriteConcurrencyObserverPairGenerator.Generate(1, 1, []string{epmetrics.MutatingKind}).RequestsWaiting,
}

var atomicMutatingExecuting, atomicReadOnlyExecuting int32
var atomicMutatingWaiting, atomicReadOnlyWaiting int32

// newInitializationSignal is defined for testing purposes.
var newInitializationSignal = utilflowcontrol.NewInitializationSignal

func truncateLogField(s string) string {
	const maxFieldLogLength = 64

	if len(s) > maxFieldLogLength {
		s = s[0:maxFieldLogLength]
	}
	return s
}

// WithPriorityAndFairness limits the number of in-flight
// requests in a fine-grained way.
func WithPriorityAndFairness(
	handler http.Handler,
	longRunningRequestCheck apirequest.LongRunningRequestCheck,
	fcIfc utilflowcontrol.Interface,
	workEstimator flowcontrolrequest.WorkEstimatorFunc,
) http.Handler {
	if fcIfc == nil {
		klog.Warningf("priority and fairness support not found, skipping")
		return handler
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()
		requestInfo, ok := apirequest.RequestInfoFrom(ctx)
		if !ok {
			handleError(w, r, fmt.Errorf("no RequestInfo found in context"))
			return
		}
		user, ok := apirequest.UserFrom(ctx)
		if !ok {
			handleError(w, r, fmt.Errorf("no User found in context"))
			return
		}

		isWatchRequest := watchVerbs.Has(requestInfo.Verb)

		// Skip tracking long running non-watch requests.
		if longRunningRequestCheck != nil && longRunningRequestCheck(r, requestInfo) && !isWatchRequest {
			klog.V(6).Infof("Serving RequestInfo=%#+v, user.Info=%#+v as longrunning\n", requestInfo, user)
			handler.ServeHTTP(w, r)
			return
		}

		var classification *PriorityAndFairnessClassification
		note := func(fs *flowcontrol.FlowSchema, pl *flowcontrol.PriorityLevelConfiguration, flowDistinguisher string) {
			classification = &PriorityAndFairnessClassification{
				FlowSchemaName:    fs.Name,
				FlowSchemaUID:     fs.UID,
				PriorityLevelName: pl.Name,
				PriorityLevelUID:  pl.UID}

			httplog.AddKeyValue(ctx, "apf_pl", truncateLogField(pl.Name))
			httplog.AddKeyValue(ctx, "apf_fs", truncateLogField(fs.Name))
			httplog.AddKeyValue(ctx, "apf_d", truncateLogField(flowDistinguisher))
		}

		var served bool
		isMutatingRequest := !nonMutatingRequestVerbs.Has(requestInfo.Verb)
		noteExecutingDelta := func(delta int32) {
			if isMutatingRequest {
				watermark.recordMutating(int(atomic.AddInt32(&atomicMutatingExecuting, delta)))
			} else {
				watermark.recordReadOnly(int(atomic.AddInt32(&atomicReadOnlyExecuting, delta)))
			}
		}
		noteWaitingDelta := func(delta int32) {
			if isMutatingRequest {
				waitingMark.recordMutating(int(atomic.AddInt32(&atomicMutatingWaiting, delta)))
			} else {
				waitingMark.recordReadOnly(int(atomic.AddInt32(&atomicReadOnlyWaiting, delta)))
			}
		}
		queueNote := func(inQueue bool) {
			if inQueue {
				noteWaitingDelta(1)
			} else {
				noteWaitingDelta(-1)
			}
		}

		// find the estimated amount of work of the request
		// TODO: Estimate cost should also take fcIfc.GetWatchCount(requestInfo) as a parameter.
		workEstimate := workEstimator.EstimateWork(r)
		digest := utilflowcontrol.RequestDigest{
			RequestInfo:  requestInfo,
			User:         user,
			WorkEstimate: workEstimate,
		}

		if isWatchRequest {
			// This channel blocks calling handler.ServeHTTP() until closed, and is closed inside execute().
			// If APF rejects the request, it is never closed.
			shouldStartWatchCh := make(chan struct{})

			watchInitializationSignal := newInitializationSignal()
			// This wraps the request passed to handler.ServeHTTP(),
			// setting a context that plumbs watchInitializationSignal to storage
			var watchReq *http.Request
			// This is set inside execute(), prior to closing shouldStartWatchCh.
			// If the request is rejected by APF it is left nil.
			var forgetWatch utilflowcontrol.ForgetWatchFunc

			defer func() {
				// Protect from the situation when request will not reach storage layer
				// and the initialization signal will not be send.
				if watchInitializationSignal != nil {
					watchInitializationSignal.Signal()
				}
				// Forget the watcher if it was registered.
				//
				// // This is race-free because by this point, one of the following occurred:
				// case <-shouldStartWatchCh: execute() completed the assignment to forgetWatch
				// case <-resultCh: Handle() completed, and Handle() does not return
				//   while execute() is running
				if forgetWatch != nil {
					forgetWatch()
				}
			}()

			execute := func() {
				noteExecutingDelta(1)
				defer noteExecutingDelta(-1)
				served = true
				setResponseHeaders(classification, w)

				forgetWatch = fcIfc.RegisterWatch(requestInfo)

				// Notify the main thread that we're ready to start the watch.
				close(shouldStartWatchCh)

				// Wait until the request is finished from the APF point of view
				// (which is when its initialization is done).
				watchInitializationSignal.Wait()
			}

			// Ensure that an item can be put to resultCh asynchronously.
			resultCh := make(chan interface{}, 1)

			// Call Handle in a separate goroutine.
			// The reason for it is that from APF point of view, the request processing
			// finishes as soon as watch is initialized (which is generally orders of
			// magnitude faster then the watch request itself). This means that Handle()
			// call finishes much faster and for performance reasons we want to reduce
			// the number of running goroutines - so we run the shorter thing in a
			// dedicated goroutine and the actual watch handler in the main one.
			go func() {
				defer func() {
					err := recover()
					// do not wrap the sentinel ErrAbortHandler panic value
					if err != nil && err != http.ErrAbortHandler {
						// Same as stdlib http server code. Manually allocate stack
						// trace buffer size to prevent excessively large logs
						const size = 64 << 10
						buf := make([]byte, size)
						buf = buf[:runtime.Stack(buf, false)]
						err = fmt.Sprintf("%v\n%s", err, buf)
					}

					// Ensure that the result is put into resultCh independently of the panic.
					resultCh <- err
				}()

				// We create handleCtx with explicit cancelation function.
				// The reason for it is that Handle() underneath may start additional goroutine
				// that is blocked on context cancellation. However, from APF point of view,
				// we don't want to wait until the whole watch request is processed (which is
				// when it context is actually cancelled) - we want to unblock the goroutine as
				// soon as the request is processed from the APF point of view.
				//
				// Note that we explicitly do NOT call the actuall handler using that context
				// to avoid cancelling request too early.
				handleCtx, handleCtxCancel := context.WithCancel(ctx)
				defer handleCtxCancel()

				// Note that Handle will return irrespective of whether the request
				// executes or is rejected. In the latter case, the function will return
				// without calling the passed `execute` function.
				fcIfc.Handle(handleCtx, digest, note, queueNote, execute)
			}()

			select {
			case <-shouldStartWatchCh:
				watchCtx := utilflowcontrol.WithInitializationSignal(ctx, watchInitializationSignal)
				watchReq = r.WithContext(watchCtx)
				handler.ServeHTTP(w, watchReq)
				// Protect from the situation when request will not reach storage layer
				// and the initialization signal will not be send.
				// It has to happen before waiting on the resultCh below.
				watchInitializationSignal.Signal()
				// TODO: Consider finishing the request as soon as Handle call panics.
				if err := <-resultCh; err != nil {
					panic(err)
				}
			case err := <-resultCh:
				if err != nil {
					panic(err)
				}
			}
		} else {
			execute := func() {
				noteExecutingDelta(1)
				defer noteExecutingDelta(-1)
				served = true
				setResponseHeaders(classification, w)

				handler.ServeHTTP(w, r)
			}

			fcIfc.Handle(ctx, digest, note, queueNote, execute)
		}

		if !served {
			setResponseHeaders(classification, w)

			if isMutatingRequest {
				epmetrics.DroppedRequests.WithContext(ctx).WithLabelValues(epmetrics.MutatingKind).Inc()
			} else {
				epmetrics.DroppedRequests.WithContext(ctx).WithLabelValues(epmetrics.ReadOnlyKind).Inc()
			}
			epmetrics.RecordRequestTermination(r, requestInfo, epmetrics.APIServerComponent, http.StatusTooManyRequests)
			tooManyRequests(r, w)
		}
	})
}

// StartPriorityAndFairnessWatermarkMaintenance starts the goroutines to observe and maintain watermarks for
// priority-and-fairness requests.
func StartPriorityAndFairnessWatermarkMaintenance(stopCh <-chan struct{}) {
	startWatermarkMaintenance(watermark, stopCh)
	startWatermarkMaintenance(waitingMark, stopCh)
}

func setResponseHeaders(classification *PriorityAndFairnessClassification, w http.ResponseWriter) {
	if classification == nil {
		return
	}

	// We intentionally set the UID of the flow-schema and priority-level instead of name. This is so that
	// the names that cluster-admins choose for categorization and priority levels are not exposed, also
	// the names might make it obvious to the users that they are rejected due to classification with low priority.
	w.Header().Set(flowcontrol.ResponseHeaderMatchedPriorityLevelConfigurationUID, string(classification.PriorityLevelUID))
	w.Header().Set(flowcontrol.ResponseHeaderMatchedFlowSchemaUID, string(classification.FlowSchemaUID))
}
<commit_msg>Rename httplog entry from "apf_d" to "apf_fd"<commit_after>/*
Copyright 2019 The Kubernetes Authors.

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

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

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

package filters

import (
	"context"
	"fmt"
	"net/http"
	"runtime"
	"sync/atomic"

	flowcontrol "k8s.io/api/flowcontrol/v1beta2"
	apitypes "k8s.io/apimachinery/pkg/types"
	epmetrics "k8s.io/apiserver/pkg/endpoints/metrics"
	apirequest "k8s.io/apiserver/pkg/endpoints/request"
	"k8s.io/apiserver/pkg/server/httplog"
	utilflowcontrol "k8s.io/apiserver/pkg/util/flowcontrol"
	fcmetrics "k8s.io/apiserver/pkg/util/flowcontrol/metrics"
	flowcontrolrequest "k8s.io/apiserver/pkg/util/flowcontrol/request"
	"k8s.io/klog/v2"
)

// PriorityAndFairnessClassification identifies the results of
// classification for API Priority and Fairness
type PriorityAndFairnessClassification struct {
	FlowSchemaName    string
	FlowSchemaUID     apitypes.UID
	PriorityLevelName string
	PriorityLevelUID  apitypes.UID
}

// waitingMark tracks requests waiting rather than being executed
var waitingMark = &requestWatermark{
	phase:            epmetrics.WaitingPhase,
	readOnlyObserver: fcmetrics.ReadWriteConcurrencyObserverPairGenerator.Generate(1, 1, []string{epmetrics.ReadOnlyKind}).RequestsWaiting,
	mutatingObserver: fcmetrics.ReadWriteConcurrencyObserverPairGenerator.Generate(1, 1, []string{epmetrics.MutatingKind}).RequestsWaiting,
}

var atomicMutatingExecuting, atomicReadOnlyExecuting int32
var atomicMutatingWaiting, atomicReadOnlyWaiting int32

// newInitializationSignal is defined for testing purposes.
var newInitializationSignal = utilflowcontrol.NewInitializationSignal

func truncateLogField(s string) string {
	const maxFieldLogLength = 64

	if len(s) > maxFieldLogLength {
		s = s[0:maxFieldLogLength]
	}
	return s
}

// WithPriorityAndFairness limits the number of in-flight
// requests in a fine-grained way.
func WithPriorityAndFairness(
	handler http.Handler,
	longRunningRequestCheck apirequest.LongRunningRequestCheck,
	fcIfc utilflowcontrol.Interface,
	workEstimator flowcontrolrequest.WorkEstimatorFunc,
) http.Handler {
	if fcIfc == nil {
		klog.Warningf("priority and fairness support not found, skipping")
		return handler
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()
		requestInfo, ok := apirequest.RequestInfoFrom(ctx)
		if !ok {
			handleError(w, r, fmt.Errorf("no RequestInfo found in context"))
			return
		}
		user, ok := apirequest.UserFrom(ctx)
		if !ok {
			handleError(w, r, fmt.Errorf("no User found in context"))
			return
		}

		isWatchRequest := watchVerbs.Has(requestInfo.Verb)

		// Skip tracking long running non-watch requests.
		if longRunningRequestCheck != nil && longRunningRequestCheck(r, requestInfo) && !isWatchRequest {
			klog.V(6).Infof("Serving RequestInfo=%#+v, user.Info=%#+v as longrunning\n", requestInfo, user)
			handler.ServeHTTP(w, r)
			return
		}

		var classification *PriorityAndFairnessClassification
		note := func(fs *flowcontrol.FlowSchema, pl *flowcontrol.PriorityLevelConfiguration, flowDistinguisher string) {
			classification = &PriorityAndFairnessClassification{
				FlowSchemaName:    fs.Name,
				FlowSchemaUID:     fs.UID,
				PriorityLevelName: pl.Name,
				PriorityLevelUID:  pl.UID}

			httplog.AddKeyValue(ctx, "apf_pl", truncateLogField(pl.Name))
			httplog.AddKeyValue(ctx, "apf_fs", truncateLogField(fs.Name))
			httplog.AddKeyValue(ctx, "apf_fd", truncateLogField(flowDistinguisher))
		}

		var served bool
		isMutatingRequest := !nonMutatingRequestVerbs.Has(requestInfo.Verb)
		noteExecutingDelta := func(delta int32) {
			if isMutatingRequest {
				watermark.recordMutating(int(atomic.AddInt32(&atomicMutatingExecuting, delta)))
			} else {
				watermark.recordReadOnly(int(atomic.AddInt32(&atomicReadOnlyExecuting, delta)))
			}
		}
		noteWaitingDelta := func(delta int32) {
			if isMutatingRequest {
				waitingMark.recordMutating(int(atomic.AddInt32(&atomicMutatingWaiting, delta)))
			} else {
				waitingMark.recordReadOnly(int(atomic.AddInt32(&atomicReadOnlyWaiting, delta)))
			}
		}
		queueNote := func(inQueue bool) {
			if inQueue {
				noteWaitingDelta(1)
			} else {
				noteWaitingDelta(-1)
			}
		}

		// find the estimated amount of work of the request
		// TODO: Estimate cost should also take fcIfc.GetWatchCount(requestInfo) as a parameter.
		workEstimate := workEstimator.EstimateWork(r)
		digest := utilflowcontrol.RequestDigest{
			RequestInfo:  requestInfo,
			User:         user,
			WorkEstimate: workEstimate,
		}

		if isWatchRequest {
			// This channel blocks calling handler.ServeHTTP() until closed, and is closed inside execute().
			// If APF rejects the request, it is never closed.
			shouldStartWatchCh := make(chan struct{})

			watchInitializationSignal := newInitializationSignal()
			// This wraps the request passed to handler.ServeHTTP(),
			// setting a context that plumbs watchInitializationSignal to storage
			var watchReq *http.Request
			// This is set inside execute(), prior to closing shouldStartWatchCh.
			// If the request is rejected by APF it is left nil.
			var forgetWatch utilflowcontrol.ForgetWatchFunc

			defer func() {
				// Protect from the situation when request will not reach storage layer
				// and the initialization signal will not be send.
				if watchInitializationSignal != nil {
					watchInitializationSignal.Signal()
				}
				// Forget the watcher if it was registered.
				//
				// // This is race-free because by this point, one of the following occurred:
				// case <-shouldStartWatchCh: execute() completed the assignment to forgetWatch
				// case <-resultCh: Handle() completed, and Handle() does not return
				//   while execute() is running
				if forgetWatch != nil {
					forgetWatch()
				}
			}()

			execute := func() {
				noteExecutingDelta(1)
				defer noteExecutingDelta(-1)
				served = true
				setResponseHeaders(classification, w)

				forgetWatch = fcIfc.RegisterWatch(requestInfo)

				// Notify the main thread that we're ready to start the watch.
				close(shouldStartWatchCh)

				// Wait until the request is finished from the APF point of view
				// (which is when its initialization is done).
				watchInitializationSignal.Wait()
			}

			// Ensure that an item can be put to resultCh asynchronously.
			resultCh := make(chan interface{}, 1)

			// Call Handle in a separate goroutine.
			// The reason for it is that from APF point of view, the request processing
			// finishes as soon as watch is initialized (which is generally orders of
			// magnitude faster then the watch request itself). This means that Handle()
			// call finishes much faster and for performance reasons we want to reduce
			// the number of running goroutines - so we run the shorter thing in a
			// dedicated goroutine and the actual watch handler in the main one.
			go func() {
				defer func() {
					err := recover()
					// do not wrap the sentinel ErrAbortHandler panic value
					if err != nil && err != http.ErrAbortHandler {
						// Same as stdlib http server code. Manually allocate stack
						// trace buffer size to prevent excessively large logs
						const size = 64 << 10
						buf := make([]byte, size)
						buf = buf[:runtime.Stack(buf, false)]
						err = fmt.Sprintf("%v\n%s", err, buf)
					}

					// Ensure that the result is put into resultCh independently of the panic.
					resultCh <- err
				}()

				// We create handleCtx with explicit cancelation function.
				// The reason for it is that Handle() underneath may start additional goroutine
				// that is blocked on context cancellation. However, from APF point of view,
				// we don't want to wait until the whole watch request is processed (which is
				// when it context is actually cancelled) - we want to unblock the goroutine as
				// soon as the request is processed from the APF point of view.
				//
				// Note that we explicitly do NOT call the actuall handler using that context
				// to avoid cancelling request too early.
				handleCtx, handleCtxCancel := context.WithCancel(ctx)
				defer handleCtxCancel()

				// Note that Handle will return irrespective of whether the request
				// executes or is rejected. In the latter case, the function will return
				// without calling the passed `execute` function.
				fcIfc.Handle(handleCtx, digest, note, queueNote, execute)
			}()

			select {
			case <-shouldStartWatchCh:
				watchCtx := utilflowcontrol.WithInitializationSignal(ctx, watchInitializationSignal)
				watchReq = r.WithContext(watchCtx)
				handler.ServeHTTP(w, watchReq)
				// Protect from the situation when request will not reach storage layer
				// and the initialization signal will not be send.
				// It has to happen before waiting on the resultCh below.
				watchInitializationSignal.Signal()
				// TODO: Consider finishing the request as soon as Handle call panics.
				if err := <-resultCh; err != nil {
					panic(err)
				}
			case err := <-resultCh:
				if err != nil {
					panic(err)
				}
			}
		} else {
			execute := func() {
				noteExecutingDelta(1)
				defer noteExecutingDelta(-1)
				served = true
				setResponseHeaders(classification, w)

				handler.ServeHTTP(w, r)
			}

			fcIfc.Handle(ctx, digest, note, queueNote, execute)
		}

		if !served {
			setResponseHeaders(classification, w)

			if isMutatingRequest {
				epmetrics.DroppedRequests.WithContext(ctx).WithLabelValues(epmetrics.MutatingKind).Inc()
			} else {
				epmetrics.DroppedRequests.WithContext(ctx).WithLabelValues(epmetrics.ReadOnlyKind).Inc()
			}
			epmetrics.RecordRequestTermination(r, requestInfo, epmetrics.APIServerComponent, http.StatusTooManyRequests)
			tooManyRequests(r, w)
		}
	})
}

// StartPriorityAndFairnessWatermarkMaintenance starts the goroutines to observe and maintain watermarks for
// priority-and-fairness requests.
func StartPriorityAndFairnessWatermarkMaintenance(stopCh <-chan struct{}) {
	startWatermarkMaintenance(watermark, stopCh)
	startWatermarkMaintenance(waitingMark, stopCh)
}

func setResponseHeaders(classification *PriorityAndFairnessClassification, w http.ResponseWriter) {
	if classification == nil {
		return
	}

	// We intentionally set the UID of the flow-schema and priority-level instead of name. This is so that
	// the names that cluster-admins choose for categorization and priority levels are not exposed, also
	// the names might make it obvious to the users that they are rejected due to classification with low priority.
	w.Header().Set(flowcontrol.ResponseHeaderMatchedPriorityLevelConfigurationUID, string(classification.PriorityLevelUID))
	w.Header().Set(flowcontrol.ResponseHeaderMatchedFlowSchemaUID, string(classification.FlowSchemaUID))
}
<|endoftext|>