<commit_before>package workflows

import (
	"fmt"
	"strings"

	"github.com/stelligent/mu/common"
)

// NewEnvironmentTerminator create a new workflow for terminating an environment
func NewEnvironmentTerminator(ctx *common.Context, environmentName string) Executor {

	workflow := new(environmentWorkflow)

	return newPipelineExecutor(
		workflow.environmentServiceTerminator(ctx.Config.Namespace, environmentName, ctx.StackManager, ctx.StackManager, ctx.StackManager, ctx.RolesetManager),
		workflow.environmentDbTerminator(ctx.Config.Namespace, environmentName, ctx.StackManager, ctx.StackManager, ctx.StackManager),
		workflow.environmentEcsTerminator(ctx.Config.Namespace, environmentName, ctx.StackManager, ctx.StackManager),
		workflow.environmentRolesetTerminator(ctx.RolesetManager, environmentName),
		workflow.environmentElbTerminator(ctx.Config.Namespace, environmentName, ctx.StackManager, ctx.StackManager),
		workflow.environmentVpcTerminator(ctx.Config.Namespace, environmentName, ctx.StackManager, ctx.StackManager),
	)
}

func (workflow *environmentWorkflow) environmentServiceTerminator(namespace string, environmentName string, stackLister common.StackLister, stackDeleter common.StackDeleter, stackWaiter common.StackWaiter, rolesetDeleter common.RolesetDeleter) Executor {
	return func() error {
		log.Noticef("Terminating Services for environment '%s' ...", environmentName)
		stacks, err := stackLister.ListStacks(common.StackTypeService, namespace)
		if err != nil {
			return err
		}
		for _, stack := range stacks {
			if stack.Tags["environment"] != environmentName {
				continue
			}
			err := stackDeleter.DeleteStack(stack.Name)
			if err != nil {
				return err
			}

			rolesetDeleter.DeleteServiceRoleset(environmentName, stack.Tags["service"])
		}
		for _, stack := range stacks {
			if stack.Tags["environment"] != environmentName {
				continue
			}
			log.Infof("   Undeploying service '%s' from environment '%s'", stack.Tags["service"], environmentName)
			stackWaiter.AwaitFinalStatus(stack.Name)
		}

		return nil
	}
}
func (workflow *environmentWorkflow) environmentDbTerminator(namespace string, environmentName string, stackLister common.StackLister, stackDeleter common.StackDeleter, stackWaiter common.StackWaiter) Executor {
	return func() error {
		log.Noticef("Terminating Databases for environment '%s' ...", environmentName)
		stacks, err := stackLister.ListStacks(common.StackTypeDatabase, namespace)
		if err != nil {
			return err
		}
		for _, stack := range stacks {
			if stack.Tags["environment"] != environmentName {
				continue
			}
			err := stackDeleter.DeleteStack(stack.Name)
			if err != nil {
				return err
			}
		}
		for _, stack := range stacks {
			if stack.Tags["environment"] != environmentName {
				continue
			}
			log.Infof("   Terminating database for service '%s' from environment '%s'", stack.Tags["service"], environmentName)
			stackWaiter.AwaitFinalStatus(stack.Name)
		}

		return nil
	}
}
func (workflow *environmentWorkflow) environmentEcsTerminator(namespace string, environmentName string, stackDeleter common.StackDeleter, stackWaiter common.StackWaiter) Executor {
	return func() error {
		log.Noticef("Terminating environment '%s' ...", environmentName)
		envStackName := common.CreateStackName(namespace, common.StackTypeEnv, environmentName)
		err := stackDeleter.DeleteStack(envStackName)
		if err != nil {
			return err
		}

		stack := stackWaiter.AwaitFinalStatus(envStackName)
		if stack != nil && !strings.HasSuffix(stack.Status, "_COMPLETE") {
			return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
		}

		return nil
	}
}

func (workflow *environmentWorkflow) environmentRolesetTerminator(rolesetDeleter common.RolesetDeleter, environmentName string) Executor {
	return func() error {
		err := rolesetDeleter.DeleteEnvironmentRoleset(environmentName)
		if err != nil {
			return err
		}
		return nil
	}
}

func (workflow *environmentWorkflow) environmentElbTerminator(namespace string, environmentName string, stackDeleter common.StackDeleter, stackWaiter common.StackWaiter) Executor {
	return func() error {
		log.Noticef("Terminating ELB environment '%s' ...", environmentName)
		envStackName := common.CreateStackName(namespace, common.StackTypeLoadBalancer, environmentName)
		err := stackDeleter.DeleteStack(envStackName)
		if err != nil {
			return err
		}

		stack := stackWaiter.AwaitFinalStatus(envStackName)
		if stack != nil && !strings.HasSuffix(stack.Status, "_COMPLETE") {
			return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
		}

		return nil
	}
}
func (workflow *environmentWorkflow) environmentVpcTerminator(namespace string, environmentName string, stackDeleter common.StackDeleter, stackWaiter common.StackWaiter) Executor {
	return func() error {
		log.Noticef("Terminating VPC environment '%s' ...", environmentName)
		vpcStackName := common.CreateStackName(namespace, common.StackTypeVpc, environmentName)
		err := stackDeleter.DeleteStack(vpcStackName)
		if err != nil {
			log.Debugf("Unable to delete VPC, but ignoring error: %v", err)
		}

		stack := stackWaiter.AwaitFinalStatus(vpcStackName)
		if stack != nil && !strings.HasSuffix(stack.Status, "_COMPLETE") {
			return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
		}

		targetStackName := common.CreateStackName(namespace, common.StackTypeTarget, environmentName)
		err = stackDeleter.DeleteStack(targetStackName)
		if err != nil {
			log.Debugf("Unable to delete VPC target, but ignoring error: %v", err)
		}

		stack = stackWaiter.AwaitFinalStatus(targetStackName)
		if stack != nil && !strings.HasSuffix(stack.Status, "_COMPLETE") {
			return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
		}

		return nil
	}
}
<commit_msg>cleanup ingress resources (and ELB) when terminating an EKS environment<commit_after>package workflows

import (
	"fmt"
	"strings"

	"github.com/stelligent/mu/common"
)

// NewEnvironmentTerminator create a new workflow for terminating an environment
func NewEnvironmentTerminator(ctx *common.Context, environmentName string) Executor {

	workflow := new(environmentWorkflow)

	return newPipelineExecutor(
		workflow.environmentLoader(ctx.Config.Namespace, environmentName, ctx.StackManager, &environmentView{}),
		workflow.environmentServiceTerminator(ctx.Config.Namespace, environmentName, ctx.StackManager, ctx.StackManager, ctx.StackManager, ctx.RolesetManager),
		workflow.environmentDbTerminator(ctx.Config.Namespace, environmentName, ctx.StackManager, ctx.StackManager, ctx.StackManager),
		newConditionalExecutor(workflow.isKubernetesProvider(),
			newPipelineExecutor(
				workflow.connectKubernetes(ctx.Config.Namespace, ctx.KubernetesResourceManagerProvider),
				workflow.environmentKubernetesIngressTerminator(environmentName),
				workflow.environmentEcsTerminator(ctx.Config.Namespace, environmentName, ctx.StackManager, ctx.StackManager),
				workflow.environmentRolesetTerminator(ctx.RolesetManager, environmentName),
			),
			newPipelineExecutor(
				workflow.environmentEcsTerminator(ctx.Config.Namespace, environmentName, ctx.StackManager, ctx.StackManager),
				workflow.environmentRolesetTerminator(ctx.RolesetManager, environmentName),
				workflow.environmentElbTerminator(ctx.Config.Namespace, environmentName, ctx.StackManager, ctx.StackManager),
			),
		),
		workflow.environmentVpcTerminator(ctx.Config.Namespace, environmentName, ctx.StackManager, ctx.StackManager),
	)
}

func (workflow *environmentWorkflow) environmentServiceTerminator(namespace string, environmentName string, stackLister common.StackLister, stackDeleter common.StackDeleter, stackWaiter common.StackWaiter, rolesetDeleter common.RolesetDeleter) Executor {
	return func() error {
		log.Noticef("Terminating Services for environment '%s' ...", environmentName)
		stacks, err := stackLister.ListStacks(common.StackTypeService, namespace)
		if err != nil {
			return err
		}
		for _, stack := range stacks {
			if stack.Tags["environment"] != environmentName {
				continue
			}
			err := stackDeleter.DeleteStack(stack.Name)
			if err != nil {
				return err
			}

			rolesetDeleter.DeleteServiceRoleset(environmentName, stack.Tags["service"])
		}
		for _, stack := range stacks {
			if stack.Tags["environment"] != environmentName {
				continue
			}
			log.Infof("   Undeploying service '%s' from environment '%s'", stack.Tags["service"], environmentName)
			stackWaiter.AwaitFinalStatus(stack.Name)
		}

		return nil
	}
}
func (workflow *environmentWorkflow) environmentDbTerminator(namespace string, environmentName string, stackLister common.StackLister, stackDeleter common.StackDeleter, stackWaiter common.StackWaiter) Executor {
	return func() error {
		log.Noticef("Terminating Databases for environment '%s' ...", environmentName)
		stacks, err := stackLister.ListStacks(common.StackTypeDatabase, namespace)
		if err != nil {
			return err
		}
		for _, stack := range stacks {
			if stack.Tags["environment"] != environmentName {
				continue
			}
			err := stackDeleter.DeleteStack(stack.Name)
			if err != nil {
				return err
			}
		}
		for _, stack := range stacks {
			if stack.Tags["environment"] != environmentName {
				continue
			}
			log.Infof("   Terminating database for service '%s' from environment '%s'", stack.Tags["service"], environmentName)
			stackWaiter.AwaitFinalStatus(stack.Name)
		}

		return nil
	}
}
func (workflow *environmentWorkflow) environmentEcsTerminator(namespace string, environmentName string, stackDeleter common.StackDeleter, stackWaiter common.StackWaiter) Executor {
	return func() error {
		log.Noticef("Terminating environment '%s' ...", environmentName)
		envStackName := common.CreateStackName(namespace, common.StackTypeEnv, environmentName)
		err := stackDeleter.DeleteStack(envStackName)
		if err != nil {
			return err
		}

		stack := stackWaiter.AwaitFinalStatus(envStackName)
		if stack != nil && !strings.HasSuffix(stack.Status, "_COMPLETE") {
			return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
		}

		return nil
	}
}

func (workflow *environmentWorkflow) environmentRolesetTerminator(rolesetDeleter common.RolesetDeleter, environmentName string) Executor {
	return func() error {
		err := rolesetDeleter.DeleteEnvironmentRoleset(environmentName)
		if err != nil {
			return err
		}
		return nil
	}
}

func (workflow *environmentWorkflow) environmentKubernetesIngressTerminator(environmentName string) Executor {
	return func() error {
		log.Noticef("Terminating ingress in environment '%s'", environmentName)

		return workflow.kubernetesResourceManager.DeleteResource("v1", "Namespace", "", "mu-ingress")
	}
}

func (workflow *environmentWorkflow) environmentElbTerminator(namespace string, environmentName string, stackDeleter common.StackDeleter, stackWaiter common.StackWaiter) Executor {
	return func() error {
		log.Noticef("Terminating ELB environment '%s' ...", environmentName)
		envStackName := common.CreateStackName(namespace, common.StackTypeLoadBalancer, environmentName)
		err := stackDeleter.DeleteStack(envStackName)
		if err != nil {
			return err
		}

		stack := stackWaiter.AwaitFinalStatus(envStackName)
		if stack != nil && !strings.HasSuffix(stack.Status, "_COMPLETE") {
			return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
		}

		return nil
	}
}
func (workflow *environmentWorkflow) environmentVpcTerminator(namespace string, environmentName string, stackDeleter common.StackDeleter, stackWaiter common.StackWaiter) Executor {
	return func() error {
		log.Noticef("Terminating VPC environment '%s' ...", environmentName)
		vpcStackName := common.CreateStackName(namespace, common.StackTypeVpc, environmentName)
		err := stackDeleter.DeleteStack(vpcStackName)
		if err != nil {
			log.Debugf("Unable to delete VPC, but ignoring error: %v", err)
		}

		stack := stackWaiter.AwaitFinalStatus(vpcStackName)
		if stack != nil && !strings.HasSuffix(stack.Status, "_COMPLETE") {
			return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
		}

		targetStackName := common.CreateStackName(namespace, common.StackTypeTarget, environmentName)
		err = stackDeleter.DeleteStack(targetStackName)
		if err != nil {
			log.Debugf("Unable to delete VPC target, but ignoring error: %v", err)
		}

		stack = stackWaiter.AwaitFinalStatus(targetStackName)
		if stack != nil && !strings.HasSuffix(stack.Status, "_COMPLETE") {
			return fmt.Errorf("Ended in failed status %s %s", stack.Status, stack.StatusReason)
		}

		return nil
	}
}
<|endoftext|>


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

package action

import (
	"errors"
	"github.com/tsuru/tsuru/log"
	"sync"
)

// Result is the value returned by Forward. It is used in the call of the next
// action, and also when rolling back the actions.
type Result interface{}

// Forward is the function called by the pipeline executor in the forward
// phase.  It receives a FWContext instance, that contains the list of
// parameters given to the pipeline executor and the result of the previous
// action in the pipeline (which will be nil for the first action in the
// pipeline).
type Forward func(context FWContext) (Result, error)

// Backward is the function called by the pipeline executor when in the
// backward phase. It receives the context instance, that contains the list of
// parameters given to the pipeline executor and the result of the forward
// phase.
type Backward func(context BWContext)

// FWContext is the context used in calls to Forward functions (forward phase).
type FWContext struct {
	// Result of the previous action.
	Previous Result

	// List of parameters given to the executor.
	Params []interface{}
}

// BWContext is the context used in calls to Backward functions (backward
// phase).
type BWContext struct {
	// Result of the forward phase (for the current action).
	FWResult Result

	// List of parameters given to the executor.
	Params []interface{}
}

// Action defines actions that should be . It is composed of two functions:
// Forward and Backward.
//
// Each action should do only one thing, and do it well. All information that
// is needed to undo the action should be returned by the Forward function.
type Action struct {
	// Name is the action name. Used by the log.
	Name string

	// Function that will be invoked in the forward phase. This value
	// cannot be nil.
	Forward Forward

	// Function that will be invoked in the backward phase. For actions
	// that are not undoable, this attribute should be nil.
	Backward Backward

	// Minimum number of parameters that this action requires to run.
	MinParams int

	// Result of the action. Stored for use in the backward phase.
	result Result

	// mutex for the result
	rMutex sync.Mutex
}

// Pipeline is a list of actions. Each pipeline is atomic: either all actions
// are successfully executed, or none of them are. For that, it's fundamental
// that all actions are really small and atomic.
type Pipeline struct {
	actions []*Action
}

// NewPipeline creates a new pipeline instance with the given list of actions.
func NewPipeline(actions ...*Action) *Pipeline {
	return &Pipeline{actions: actions}
}

func (p *Pipeline) Result() Result {
	action := p.actions[len(p.actions)-1]
	action.rMutex.Lock()
	defer action.rMutex.Unlock()
	return action.result
}

// Execute executes the pipeline.
//
// The execution starts in the forward phase, calling the Forward function of
// all actions. If none of the Forward calls return error, the pipeline
// execution ends in the forward phase and is "committed".
//
// If any of the Forward call fail, the executor switches to the backward phase
// (roll back) and call the Backward function for each action completed. It
// does not call the Backward function of the action that has failed.
//
// After rolling back all completed actions, it returns the original error
// returned by the action that failed.
func (p *Pipeline) Execute(params ...interface{}) error {
	var (
		r   Result
		err error
	)
	if len(p.actions) == 0 {
		return errors.New("No actions to execute.")
	}
	fwCtx := FWContext{Params: params}
	for i, a := range p.actions {
		log.Debugf("[pipeline] running the Forward for the %s action", a.Name)
		if a.Forward == nil {
			err = errors.New("All actions must define the forward function.")
		} else if len(fwCtx.Params) < a.MinParams {
			err = errors.New("Not enough parameters to call Action.Forward.")
		} else {
			r, err = a.Forward(fwCtx)
			a.rMutex.Lock()
			a.result = r
			a.rMutex.Unlock()
			fwCtx.Previous = r
		}
		if err != nil {
			log.Debugf("[pipeline] error running the Forward for the %s action - %s", a.Name, err)
			p.rollback(i-1, params)
			return err
		}
	}
	return nil
}

func (p *Pipeline) rollback(index int, params []interface{}) {
	bwCtx := BWContext{Params: params}
	for i := index; i >= 0; i-- {
		log.Debugf("[pipeline] running Backward for %s action", p.actions[i].Name)
		if p.actions[i].Backward != nil {
			bwCtx.FWResult = p.actions[i].result
			p.actions[i].Backward(bwCtx)
		}
	}
}
<commit_msg>action: added docs for pipeline.Result.<commit_after>// Copyright 2014 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package action

import (
	"errors"
	"github.com/tsuru/tsuru/log"
	"sync"
)

// Result is the value returned by Forward. It is used in the call of the next
// action, and also when rolling back the actions.
type Result interface{}

// Forward is the function called by the pipeline executor in the forward
// phase.  It receives a FWContext instance, that contains the list of
// parameters given to the pipeline executor and the result of the previous
// action in the pipeline (which will be nil for the first action in the
// pipeline).
type Forward func(context FWContext) (Result, error)

// Backward is the function called by the pipeline executor when in the
// backward phase. It receives the context instance, that contains the list of
// parameters given to the pipeline executor and the result of the forward
// phase.
type Backward func(context BWContext)

// FWContext is the context used in calls to Forward functions (forward phase).
type FWContext struct {
	// Result of the previous action.
	Previous Result

	// List of parameters given to the executor.
	Params []interface{}
}

// BWContext is the context used in calls to Backward functions (backward
// phase).
type BWContext struct {
	// Result of the forward phase (for the current action).
	FWResult Result

	// List of parameters given to the executor.
	Params []interface{}
}

// Action defines actions that should be . It is composed of two functions:
// Forward and Backward.
//
// Each action should do only one thing, and do it well. All information that
// is needed to undo the action should be returned by the Forward function.
type Action struct {
	// Name is the action name. Used by the log.
	Name string

	// Function that will be invoked in the forward phase. This value
	// cannot be nil.
	Forward Forward

	// Function that will be invoked in the backward phase. For actions
	// that are not undoable, this attribute should be nil.
	Backward Backward

	// Minimum number of parameters that this action requires to run.
	MinParams int

	// Result of the action. Stored for use in the backward phase.
	result Result

	// mutex for the result
	rMutex sync.Mutex
}

// Pipeline is a list of actions. Each pipeline is atomic: either all actions
// are successfully executed, or none of them are. For that, it's fundamental
// that all actions are really small and atomic.
type Pipeline struct {
	actions []*Action
}

// NewPipeline creates a new pipeline instance with the given list of actions.
func NewPipeline(actions ...*Action) *Pipeline {
	return &Pipeline{actions: actions}

}

// Result returns the result of the last action.
func (p *Pipeline) Result() Result {
	action := p.actions[len(p.actions)-1]
	action.rMutex.Lock()
	defer action.rMutex.Unlock()
	return action.result
}

// Execute executes the pipeline.
//
// The execution starts in the forward phase, calling the Forward function of
// all actions. If none of the Forward calls return error, the pipeline
// execution ends in the forward phase and is "committed".
//
// If any of the Forward call fail, the executor switches to the backward phase
// (roll back) and call the Backward function for each action completed. It
// does not call the Backward function of the action that has failed.
//
// After rolling back all completed actions, it returns the original error
// returned by the action that failed.
func (p *Pipeline) Execute(params ...interface{}) error {
	var (
		r   Result
		err error
	)
	if len(p.actions) == 0 {
		return errors.New("No actions to execute.")
	}
	fwCtx := FWContext{Params: params}
	for i, a := range p.actions {
		log.Debugf("[pipeline] running the Forward for the %s action", a.Name)
		if a.Forward == nil {
			err = errors.New("All actions must define the forward function.")
		} else if len(fwCtx.Params) < a.MinParams {
			err = errors.New("Not enough parameters to call Action.Forward.")
		} else {
			r, err = a.Forward(fwCtx)
			a.rMutex.Lock()
			a.result = r
			a.rMutex.Unlock()
			fwCtx.Previous = r
		}
		if err != nil {
			log.Debugf("[pipeline] error running the Forward for the %s action - %s", a.Name, err)
			p.rollback(i-1, params)
			return err
		}
	}
	return nil
}

func (p *Pipeline) rollback(index int, params []interface{}) {
	bwCtx := BWContext{Params: params}
	for i := index; i >= 0; i-- {
		log.Debugf("[pipeline] running Backward for %s action", p.actions[i].Name)
		if p.actions[i].Backward != nil {
			bwCtx.FWResult = p.actions[i].result
			p.actions[i].Backward(bwCtx)
		}
	}
}
<|endoftext|>


<commit_before>/*
adstop is an ad-blocking transparent HTTP/HTTPS proxy.

It was designed to run on low power, low memory ARM devices and serve a couple
of clients, mostly old smartphones which cannot run adblockers themselves.

Before using it, you have to configure your devices and network to make it
accessible as a transparent proxy. One way to achieve this is to install
a VPN on the server side and redirect all HTTP/HTTPS traffic to the proxy
with routing rules. Then make the client browse through the VPN.

	$ adstop -http localhost:1080 \
		-https localhost:1081     \
		-cache .adstop			  \
		-max-age 24h			  \
		https://easylist-downloads.adblockplus.org/easylist.txt \
		some_local_list.txt

starts the proxy and makes it listen on HTTP on port 1080, HTTPS on port 1081,
fetch and load rules from easylist and a local file, cache easylist in an
.adstop/ directory and refresh it every 24 hours.

Note that HTTPS filtering requires the proxy to intercept the device traffic
and decrypt it. To allow this, you have to add the proxy certificate authority
to your device. By default, adstop uses goproxy ca.pem file but you should
generate your own to avoid opening your device communications to third parties.

TODO: configure the CA on command line to avoid recompiling goproxy with our
own.
*/
package main

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"flag"
	"fmt"
	"log"
	"mime"
	"net"
	"net/http"
	"net/url"
	"sync"
	"time"

	"github.com/elazarl/goproxy"
	"github.com/inconshreveable/go-vhost"
	"github.com/pmezard/adblock/adblock"
)

var (
	httpAddr  = flag.String("http", "localhost:1080", "HTTP handler address")
	httpsAddr = flag.String("https", "localhost:1081", "HTTPS handler address")
	logp      = flag.Bool("log", false, "enable logging")
	cacheDir  = flag.String("cache", ".cache", "cache directory")
	maxAgeArg = flag.String("max-age", "24h", "cached entries max age")
)

type FilteringHandler struct {
	Cache *RuleCache
}

func logRequest(r *http.Request) {
	log.Printf("%s %s %s %s\n", r.Proto, r.Method, r.URL, r.Host)
	buf := &bytes.Buffer{}
	r.Header.Write(buf)
	log.Println(string(buf.Bytes()))
}

func getReferrerDomain(r *http.Request) string {
	ref := r.Header.Get("Referer")
	if len(ref) > 0 {
		u, err := url.Parse(ref)
		if err == nil {
			return u.Host
		}
	}
	return ""
}

type ProxyState struct {
	Duration time.Duration
	URL      string
}

func (h *FilteringHandler) OnRequest(r *http.Request, ctx *goproxy.ProxyCtx) (
	*http.Request, *http.Response) {

	if *logp {
		logRequest(r)
	}

	host := r.URL.Host
	if host == "" {
		host = r.Host
	}
	rq := &adblock.Request{
		URL:          r.URL.String(),
		Domain:       host,
		OriginDomain: getReferrerDomain(r),
	}
	rules := h.Cache.Rules()
	start := time.Now()
	matched, id := rules.Matcher.Match(rq)
	end := time.Now()
	duration := end.Sub(start) / time.Millisecond
	if matched {
		rule := rules.Rules[id]
		log.Printf("rejected in %dms: %s\n", duration, r.URL.String())
		log.Printf("  by %s\n", rule)
		return r, goproxy.NewResponse(r, goproxy.ContentTypeText,
			http.StatusNotFound, "Not Found")
	}
	ctx.UserData = &ProxyState{
		Duration: duration,
		URL:      r.URL.String(),
	}
	return r, nil
}

func (h *FilteringHandler) OnResponse(r *http.Response,
	ctx *goproxy.ProxyCtx) *http.Response {

	if r == nil {
		// Happens if RoundTrip fails
		return r
	}

	state, ok := ctx.UserData.(*ProxyState)
	if !ok {
		// The request was rejected by the previous handler
		return r
	}
	duration2 := time.Duration(0)
	mediaType, _, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
	if err == nil && len(mediaType) > 0 {
		host := ctx.Req.URL.Host
		if host == "" {
			host = ctx.Req.Host
		}
		rq := &adblock.Request{
			URL:          ctx.Req.URL.String(),
			Domain:       host,
			OriginDomain: getReferrerDomain(ctx.Req),
			ContentType:  mediaType,
		}
		// Second level filtering, based on returned content
		rules := h.Cache.Rules()
		start := time.Now()
		matched, id := rules.Matcher.Match(rq)
		end := time.Now()
		duration2 = end.Sub(start) / time.Millisecond
		if matched {
			r.Body.Close()
			rule := rules.Rules[id]
			log.Printf("rejected in %d/%dms: %s\n", state.Duration, duration2,
				state.URL)
			log.Printf("  by %s\n", rule)
			return goproxy.NewResponse(ctx.Req, goproxy.ContentTypeText,
				http.StatusNotFound, "Not Found")
		}
	}
	log.Printf("accepted in %d/%dms: %s\n", state.Duration, duration2, state.URL)
	return r
}

// CachedConfig holds a TLS configuration. It can be in different states:
// - The config is being generated, Config is nil and the Ready channel is set
// - The config is ready, Config is not nil and Ready is closed.
// This mechanism is used to pool concurrent generations of the same certificate.
type CachedConfig struct {
	Config *tls.Config
	Ready  chan struct{}
}

// TLSConfigCache is a goroutine-safe cache of TLS configurations mapped to hosts.
type TLSConfigCache struct {
	cfgBuilder func(string, *goproxy.ProxyCtx) (*tls.Config, error)
	lock       sync.Mutex
	cache      map[string]CachedConfig
	hit        int
	miss       int
}

func NewTLSConfigCache(ca *tls.Certificate) *TLSConfigCache {
	return &TLSConfigCache{
		cfgBuilder: goproxy.TLSConfigFromCA(ca),
		cache:      map[string]CachedConfig{},
	}
}

func (c *TLSConfigCache) GetConfig(host string, ctx *goproxy.ProxyCtx) (*tls.Config, error) {
	c.lock.Lock()
	cached, ok := c.cache[host]
	if !ok {
		// Register a config generation event
		cached = CachedConfig{
			Ready: make(chan struct{}),
		}
		c.cache[host] = cached
	}
	if ok {
		c.hit += 1
	} else {
		c.miss += 1
	}
	hit := c.hit
	miss := c.miss
	c.lock.Unlock()

	ctx.Warnf("signing hit/miss: %d/%d (%.1f%%)", hit, miss,
		100.0*float64(hit)/float64(hit+miss))
	if ok {
		// config is being generated or is ready, grab it
		<-cached.Ready
		cfg := cached.Config
		if cfg == nil {
			return nil, fmt.Errorf("failed to generate TLS config for %s", host)
		}
		return cfg, nil
	}

	// Generate it
	start := time.Now()
	cfg, err := c.cfgBuilder(host, ctx)
	stop := time.Now()
	ctx.Warnf("signing %s in %.0fms", host,
		float64(stop.Sub(start))/float64(time.Millisecond))

	c.lock.Lock()
	if err == nil {
		c.cache[host] = CachedConfig{
			Config: cfg,
			Ready:  cached.Ready,
		}
	} else {
		delete(c.cache, host)
		ctx.Warnf("failed to sign %s: %s", host, err)
	}
	close(cached.Ready)
	c.lock.Unlock()
	return cfg, err
}

// copied/converted from https.go
type dumbResponseWriter struct {
	net.Conn
}

func (dumb dumbResponseWriter) Header() http.Header {
	panic("Header() should not be called on this ResponseWriter")
}

func (dumb dumbResponseWriter) Write(buf []byte) (int, error) {
	if bytes.Equal(buf, []byte("HTTP/1.0 200 OK\r\n\r\n")) {
		// throw away the HTTP OK response from the faux CONNECT request
		return len(buf), nil
	}
	return dumb.Conn.Write(buf)
}

func (dumb dumbResponseWriter) WriteHeader(code int) {
	panic("WriteHeader() should not be called on this ResponseWriter")
}

func (dumb dumbResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	return dumb, bufio.NewReadWriter(bufio.NewReader(dumb), bufio.NewWriter(dumb)), nil
}

func runProxy() error {
	flag.Parse()
	maxAge, err := time.ParseDuration(*maxAgeArg)
	if err != nil {
		return fmt.Errorf("invalid max-age: %s", err)
	}
	if maxAge < 0 {
		return fmt.Errorf("invalid negative max-age")
	}
	log.Printf("loading rules")
	cache, err := NewRuleCache(*cacheDir, flag.Args(), maxAge)
	if err != nil {
		return err
	}
	h := &FilteringHandler{
		Cache: cache,
	}

	log.Printf("starting servers")
	proxy := goproxy.NewProxyHttpServer()
	proxy.NonproxyHandler = http.HandlerFunc(
		func(w http.ResponseWriter, req *http.Request) {
			if req.Host == "" {
				log.Printf("Cannot handle requests without Host header, e.g., HTTP 1.0")
				return
			}
			req.URL.Scheme = "http"
			req.URL.Host = req.Host
			proxy.ServeHTTP(w, req)
		})

	// Cache MITM certificates
	tlsCache := NewTLSConfigCache(&goproxy.GoproxyCa)
	MitmConnect := &goproxy.ConnectAction{
		Action: goproxy.ConnectMitm,
		TLSConfig: func(host string, ctx *goproxy.ProxyCtx) (*tls.Config, error) {
			return tlsCache.GetConfig(host, ctx)
		},
	}
	var AlwaysMitm goproxy.FuncHttpsHandler = func(host string, ctx *goproxy.ProxyCtx) (
		*goproxy.ConnectAction, string) {

		return MitmConnect, host
	}
	proxy.OnRequest().HandleConnect(AlwaysMitm)

	proxy.OnRequest().DoFunc(h.OnRequest)
	proxy.OnResponse().DoFunc(h.OnResponse)
	go func() {
		http.ListenAndServe(*httpAddr, proxy)
	}()

	// listen to the TLS ClientHello but make it a CONNECT request instead
	ln, err := net.Listen("tcp", *httpsAddr)
	if err != nil {
		return err
	}
	for {
		c, err := ln.Accept()
		if err != nil {
			log.Printf("error accepting new connection - %v", err)
			continue
		}
		go func(c net.Conn) {
			tlsConn, err := vhost.TLS(c)
			if err != nil {
				log.Printf("error accepting new connection - %v", err)
			}
			if tlsConn.Host() == "" {
				log.Printf("cannot support non-SNI enabled clients")
				return
			}
			connectReq := &http.Request{
				Method: "CONNECT",
				URL: &url.URL{
					Opaque: tlsConn.Host(),
					Host:   net.JoinHostPort(tlsConn.Host(), "443"),
				},
				Host:   tlsConn.Host(),
				Header: make(http.Header),
			}
			resp := dumbResponseWriter{tlsConn}
			proxy.ServeHTTP(resp, connectReq)
		}(c)
	}
}

func main() {
	err := runProxy()
	if err != nil {
		log.Fatalf("error: %s\n", err)
	}
}
<commit_msg>adstop: add -ca-cert and -ca-key to configure MITM certificate<commit_after>/*
adstop is an ad-blocking transparent HTTP/HTTPS proxy.

It was designed to run on low power, low memory ARM devices and serve a couple
of clients, mostly old smartphones which cannot run adblockers themselves.

Before using it, you have to configure your devices and network to make it
accessible as a transparent proxy. One way to achieve this is to install
a VPN on the server side and redirect all HTTP/HTTPS traffic to the proxy
with routing rules. Then make the client browse through the VPN.

HTTPS filtering requires the proxy to intercept the device traffic and decrypt
it. To allow this, you have to generate a certificate and add it to your
device.

You also need to generate a certificate

	$ adstop -http localhost:1080 \
		-https localhost:1081     \
		-cache .adstop			  \
		-max-age 24h			  \
		-ca-cert /path/to/ca.cert \
		-ca-key /path/to/ca.key   \
		https://easylist-downloads.adblockplus.org/easylist.txt \
		some_local_list.txt

starts the proxy and makes it listen on HTTP on port 1080, HTTPS on port 1081,
fetch and load rules from easylist and a local file, cache easylist in an
.adstop/ directory and refresh it every 24 hours.

*/
package main

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"mime"
	"net"
	"net/http"
	"net/url"
	"sync"
	"time"

	"github.com/elazarl/goproxy"
	"github.com/inconshreveable/go-vhost"
	"github.com/pmezard/adblock/adblock"
)

var (
	httpAddr  = flag.String("http", "localhost:1080", "HTTP handler address")
	httpsAddr = flag.String("https", "localhost:1081", "HTTPS handler address")
	logp      = flag.Bool("log", false, "enable logging")
	cacheDir  = flag.String("cache", ".cache", "cache directory")
	maxAgeArg = flag.String("max-age", "24h", "cached entries max age")
	caCert    = flag.String("ca-cert", "", "path to CA certificate")
	caKey     = flag.String("ca-key", "", "path to CA key")
)

type FilteringHandler struct {
	Cache *RuleCache
}

func logRequest(r *http.Request) {
	log.Printf("%s %s %s %s\n", r.Proto, r.Method, r.URL, r.Host)
	buf := &bytes.Buffer{}
	r.Header.Write(buf)
	log.Println(string(buf.Bytes()))
}

func getReferrerDomain(r *http.Request) string {
	ref := r.Header.Get("Referer")
	if len(ref) > 0 {
		u, err := url.Parse(ref)
		if err == nil {
			return u.Host
		}
	}
	return ""
}

type ProxyState struct {
	Duration time.Duration
	URL      string
}

func (h *FilteringHandler) OnRequest(r *http.Request, ctx *goproxy.ProxyCtx) (
	*http.Request, *http.Response) {

	if *logp {
		logRequest(r)
	}

	host := r.URL.Host
	if host == "" {
		host = r.Host
	}
	rq := &adblock.Request{
		URL:          r.URL.String(),
		Domain:       host,
		OriginDomain: getReferrerDomain(r),
	}
	rules := h.Cache.Rules()
	start := time.Now()
	matched, id := rules.Matcher.Match(rq)
	end := time.Now()
	duration := end.Sub(start) / time.Millisecond
	if matched {
		rule := rules.Rules[id]
		log.Printf("rejected in %dms: %s\n", duration, r.URL.String())
		log.Printf("  by %s\n", rule)
		return r, goproxy.NewResponse(r, goproxy.ContentTypeText,
			http.StatusNotFound, "Not Found")
	}
	ctx.UserData = &ProxyState{
		Duration: duration,
		URL:      r.URL.String(),
	}
	return r, nil
}

func (h *FilteringHandler) OnResponse(r *http.Response,
	ctx *goproxy.ProxyCtx) *http.Response {

	if r == nil {
		// Happens if RoundTrip fails
		return r
	}

	state, ok := ctx.UserData.(*ProxyState)
	if !ok {
		// The request was rejected by the previous handler
		return r
	}
	duration2 := time.Duration(0)
	mediaType, _, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
	if err == nil && len(mediaType) > 0 {
		host := ctx.Req.URL.Host
		if host == "" {
			host = ctx.Req.Host
		}
		rq := &adblock.Request{
			URL:          ctx.Req.URL.String(),
			Domain:       host,
			OriginDomain: getReferrerDomain(ctx.Req),
			ContentType:  mediaType,
		}
		// Second level filtering, based on returned content
		rules := h.Cache.Rules()
		start := time.Now()
		matched, id := rules.Matcher.Match(rq)
		end := time.Now()
		duration2 = end.Sub(start) / time.Millisecond
		if matched {
			r.Body.Close()
			rule := rules.Rules[id]
			log.Printf("rejected in %d/%dms: %s\n", state.Duration, duration2,
				state.URL)
			log.Printf("  by %s\n", rule)
			return goproxy.NewResponse(ctx.Req, goproxy.ContentTypeText,
				http.StatusNotFound, "Not Found")
		}
	}
	log.Printf("accepted in %d/%dms: %s\n", state.Duration, duration2, state.URL)
	return r
}

// CachedConfig holds a TLS configuration. It can be in different states:
// - The config is being generated, Config is nil and the Ready channel is set
// - The config is ready, Config is not nil and Ready is closed.
// This mechanism is used to pool concurrent generations of the same certificate.
type CachedConfig struct {
	Config *tls.Config
	Ready  chan struct{}
}

// TLSConfigCache is a goroutine-safe cache of TLS configurations mapped to hosts.
type TLSConfigCache struct {
	cfgBuilder func(string, *goproxy.ProxyCtx) (*tls.Config, error)
	lock       sync.Mutex
	cache      map[string]CachedConfig
	hit        int
	miss       int
}

func NewTLSConfigCache(ca *tls.Certificate) *TLSConfigCache {
	return &TLSConfigCache{
		cfgBuilder: goproxy.TLSConfigFromCA(ca),
		cache:      map[string]CachedConfig{},
	}
}

func (c *TLSConfigCache) GetConfig(host string, ctx *goproxy.ProxyCtx) (*tls.Config, error) {
	c.lock.Lock()
	cached, ok := c.cache[host]
	if !ok {
		// Register a config generation event
		cached = CachedConfig{
			Ready: make(chan struct{}),
		}
		c.cache[host] = cached
	}
	if ok {
		c.hit += 1
	} else {
		c.miss += 1
	}
	hit := c.hit
	miss := c.miss
	c.lock.Unlock()

	ctx.Warnf("signing hit/miss: %d/%d (%.1f%%)", hit, miss,
		100.0*float64(hit)/float64(hit+miss))
	if ok {
		// config is being generated or is ready, grab it
		<-cached.Ready
		cfg := cached.Config
		if cfg == nil {
			return nil, fmt.Errorf("failed to generate TLS config for %s", host)
		}
		return cfg, nil
	}

	// Generate it
	start := time.Now()
	cfg, err := c.cfgBuilder(host, ctx)
	stop := time.Now()
	ctx.Warnf("signing %s in %.0fms", host,
		float64(stop.Sub(start))/float64(time.Millisecond))

	c.lock.Lock()
	if err == nil {
		c.cache[host] = CachedConfig{
			Config: cfg,
			Ready:  cached.Ready,
		}
	} else {
		delete(c.cache, host)
		ctx.Warnf("failed to sign %s: %s", host, err)
	}
	close(cached.Ready)
	c.lock.Unlock()
	return cfg, err
}

// copied/converted from https.go
type dumbResponseWriter struct {
	net.Conn
}

func (dumb dumbResponseWriter) Header() http.Header {
	panic("Header() should not be called on this ResponseWriter")
}

func (dumb dumbResponseWriter) Write(buf []byte) (int, error) {
	if bytes.Equal(buf, []byte("HTTP/1.0 200 OK\r\n\r\n")) {
		// throw away the HTTP OK response from the faux CONNECT request
		return len(buf), nil
	}
	return dumb.Conn.Write(buf)
}

func (dumb dumbResponseWriter) WriteHeader(code int) {
	panic("WriteHeader() should not be called on this ResponseWriter")
}

func (dumb dumbResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	return dumb, bufio.NewReadWriter(bufio.NewReader(dumb), bufio.NewWriter(dumb)), nil
}

func makeCertificate(certPath, keyPath string) (*tls.Certificate, error) {
	cert, err := ioutil.ReadFile(certPath)
	if err != nil {
		return nil, fmt.Errorf("cannot load CA certificate: %s", err)
	}
	key, err := ioutil.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("cannot load CA key: %s", err)
	}
	ca, err := tls.X509KeyPair(cert, key)
	return &ca, err
}

func runProxy() error {
	flag.Parse()
	if *caCert == "" || *caKey == "" {
		return fmt.Errorf("CA certificate and key must be specified")
	}
	ca, err := makeCertificate(*caCert, *caKey)
	if err != nil {
		return err
	}

	maxAge, err := time.ParseDuration(*maxAgeArg)
	if err != nil {
		return fmt.Errorf("invalid max-age: %s", err)
	}
	if maxAge < 0 {
		return fmt.Errorf("invalid negative max-age")
	}
	log.Printf("loading rules")
	cache, err := NewRuleCache(*cacheDir, flag.Args(), maxAge)
	if err != nil {
		return err
	}
	h := &FilteringHandler{
		Cache: cache,
	}

	log.Printf("starting servers")
	proxy := goproxy.NewProxyHttpServer()
	proxy.NonproxyHandler = http.HandlerFunc(
		func(w http.ResponseWriter, req *http.Request) {
			if req.Host == "" {
				log.Printf("Cannot handle requests without Host header, e.g., HTTP 1.0")
				return
			}
			req.URL.Scheme = "http"
			req.URL.Host = req.Host
			proxy.ServeHTTP(w, req)
		})

	// Cache MITM certificates
	tlsCache := NewTLSConfigCache(ca)
	MitmConnect := &goproxy.ConnectAction{
		Action: goproxy.ConnectMitm,
		TLSConfig: func(host string, ctx *goproxy.ProxyCtx) (*tls.Config, error) {
			return tlsCache.GetConfig(host, ctx)
		},
	}
	var AlwaysMitm goproxy.FuncHttpsHandler = func(host string, ctx *goproxy.ProxyCtx) (
		*goproxy.ConnectAction, string) {

		return MitmConnect, host
	}
	proxy.OnRequest().HandleConnect(AlwaysMitm)

	proxy.OnRequest().DoFunc(h.OnRequest)
	proxy.OnResponse().DoFunc(h.OnResponse)
	go func() {
		http.ListenAndServe(*httpAddr, proxy)
	}()

	// listen to the TLS ClientHello but make it a CONNECT request instead
	ln, err := net.Listen("tcp", *httpsAddr)
	if err != nil {
		return err
	}
	for {
		c, err := ln.Accept()
		if err != nil {
			log.Printf("error accepting new connection - %v", err)
			continue
		}
		go func(c net.Conn) {
			tlsConn, err := vhost.TLS(c)
			if err != nil {
				log.Printf("error accepting new connection - %v", err)
			}
			if tlsConn.Host() == "" {
				log.Printf("cannot support non-SNI enabled clients")
				return
			}
			connectReq := &http.Request{
				Method: "CONNECT",
				URL: &url.URL{
					Opaque: tlsConn.Host(),
					Host:   net.JoinHostPort(tlsConn.Host(), "443"),
				},
				Host:   tlsConn.Host(),
				Header: make(http.Header),
			}
			resp := dumbResponseWriter{tlsConn}
			proxy.ServeHTTP(resp, connectReq)
		}(c)
	}
}

func main() {
	err := runProxy()
	if err != nil {
		log.Fatalf("error: %s\n", err)
	}
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/terraform/flatmap"
	"github.com/hashicorp/terraform/helper/config"
	"github.com/hashicorp/terraform/helper/diff"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
	"github.com/mitchellh/goamz/ec2"
)

func resource_aws_security_group_create(
	s *terraform.ResourceState,
	d *terraform.ResourceDiff,
	meta interface{}) (*terraform.ResourceState, error) {
	p := meta.(*ResourceProvider)
	ec2conn := p.ec2conn

	// Merge the diff into the state so that we have all the attributes
	// properly.
	rs := s.MergeDiff(d)

	securityGroupOpts := ec2.SecurityGroup{
		Name: rs.Attributes["name"],
	}

	if rs.Attributes["vpc_id"] != "" {
		securityGroupOpts.VpcId = rs.Attributes["vpc_id"]
	}

	if rs.Attributes["description"] != "" {
		securityGroupOpts.Description = rs.Attributes["description"]
	}

	log.Printf("[DEBUG] Security Group create configuration: %#v", securityGroupOpts)
	createResp, err := ec2conn.CreateSecurityGroup(securityGroupOpts)
	if err != nil {
		return nil, fmt.Errorf("Error creating Security Group: %s", err)
	}

	rs.ID = createResp.Id
	group := createResp.SecurityGroup

	log.Printf("[INFO] Security Group ID: %s", rs.ID)

	// Wait for the security group to truly exist
	log.Printf(
		"[DEBUG] Waiting for SG (%s) to exist",
		s.ID)
	stateConf := &resource.StateChangeConf{
		Pending: []string{""},
		Target:  "exists",
		Refresh: SGStateRefreshFunc(ec2conn, rs.ID),
		Timeout: 1 * time.Minute,
	}
	if _, err := stateConf.WaitForState(); err != nil {
		return s, fmt.Errorf(
			"Error waiting for SG (%s) to become available: %s",
			rs.ID, err)
	}

	// Expand the "ingress" array to goamz compat []ec2.IPPerm
	ingressRules := []ec2.IPPerm{}
	v, ok := flatmap.Expand(rs.Attributes, "ingress").([]interface{})
	if ok {
		ingressRules = expandIPPerms(v)
	}

	if len(ingressRules) > 0 {
		_, err = ec2conn.AuthorizeSecurityGroup(group, ingressRules)
		if err != nil {
			return rs, fmt.Errorf("Error authorizing security group ingress rules: %s", err)
		}
	}

	return resource_aws_security_group_refresh(rs, meta)
}

func resource_aws_security_group_update(
	s *terraform.ResourceState,
	d *terraform.ResourceDiff,
	meta interface{}) (*terraform.ResourceState, error) {

	rs := s.MergeDiff(d)
	log.Printf("ResourceDiff: %s", d)
	log.Printf("ResourceState: %s", s)
	log.Printf("Merged: %s", rs)

	return nil, fmt.Errorf("Did not update")
}

func resource_aws_security_group_destroy(
	s *terraform.ResourceState,
	meta interface{}) error {
	p := meta.(*ResourceProvider)
	ec2conn := p.ec2conn

	log.Printf("[DEBUG] Security Group destroy: %v", s.ID)

	_, err := ec2conn.DeleteSecurityGroup(ec2.SecurityGroup{Id: s.ID})
	if err != nil {
		ec2err, ok := err.(*ec2.Error)
		if ok && ec2err.Code == "InvalidGroup.NotFound" {
			return nil
		}
	}

	return err
}

func resource_aws_security_group_refresh(
	s *terraform.ResourceState,
	meta interface{}) (*terraform.ResourceState, error) {
	p := meta.(*ResourceProvider)
	ec2conn := p.ec2conn

	sgRaw, _, err := SGStateRefreshFunc(ec2conn, s.ID)()
	if err != nil {
		return s, err
	}
	if sgRaw == nil {
		return nil, nil
	}

	return resource_aws_security_group_update_state(
		s, sgRaw.(*ec2.SecurityGroupInfo))
}

func resource_aws_security_group_diff(
	s *terraform.ResourceState,
	c *terraform.ResourceConfig,
	meta interface{}) (*terraform.ResourceDiff, error) {

	b := &diff.ResourceBuilder{
		Attrs: map[string]diff.AttrType{
			"name":        diff.AttrTypeCreate,
			"description": diff.AttrTypeUpdate,
			"ingress":     diff.AttrTypeUpdate,
			"vpc_id":      diff.AttrTypeCreate,
		},

		ComputedAttrs: []string{
			"owner_id",
			"vpc_id",
		},
	}

	return b.Diff(s, c)
}

func resource_aws_security_group_update_state(
	s *terraform.ResourceState,
	sg *ec2.SecurityGroupInfo) (*terraform.ResourceState, error) {

	s.Attributes["description"] = sg.Description
	s.Attributes["name"] = sg.Name
	s.Attributes["vpc_id"] = sg.VpcId
	s.Attributes["owner_id"] = sg.OwnerId

	// Flatten our ingress values
	toFlatten := make(map[string]interface{})
	toFlatten["ingress"] = flattenIPPerms(sg.IPPerms)

	for k, v := range flatmap.Flatten(toFlatten) {
		s.Attributes[k] = v
	}

	s.Dependencies = nil
	if s.Attributes["vpc_id"] != "" {
		s.Dependencies = append(s.Dependencies,
			terraform.ResourceDependency{ID: s.Attributes["vpc_id"]},
		)
	}

	return s, nil
}

func resource_aws_security_group_validation() *config.Validator {
	return &config.Validator{
		Required: []string{
			"name",
			"ingress.*",
			"ingress.*.from_port",
			"ingress.*.to_port",
			"ingress.*.protocol",
		},
		Optional: []string{
			"description",
			"vpc_id",
			"owner_id",
			"ingress.*.cidr_blocks.*",
			"ingress.*.security_groups.*",
		},
	}
}

// SGStateRefreshFunc returns a resource.StateRefreshFunc that is used to watch
// a security group.
func SGStateRefreshFunc(conn *ec2.EC2, id string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		sgs := []ec2.SecurityGroup{ec2.SecurityGroup{Id: id}}
		resp, err := conn.SecurityGroups(sgs, nil)
		if err != nil {
			if ec2err, ok := err.(*ec2.Error); ok &&
				ec2err.Code == "InvalidSecurityGroupID.NotFound" {
				resp = nil
			} else {
				log.Printf("Error on SGStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil {
			return nil, "", nil
		}

		group := &resp.Groups[0]
		return group, "exists", nil
	}
}
<commit_msg>providers/aws/aws_security_group: we need to check for one other error<commit_after>package aws

import (
	"fmt"
	"log"
	"time"

	"github.com/hashicorp/terraform/flatmap"
	"github.com/hashicorp/terraform/helper/config"
	"github.com/hashicorp/terraform/helper/diff"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
	"github.com/mitchellh/goamz/ec2"
)

func resource_aws_security_group_create(
	s *terraform.ResourceState,
	d *terraform.ResourceDiff,
	meta interface{}) (*terraform.ResourceState, error) {
	p := meta.(*ResourceProvider)
	ec2conn := p.ec2conn

	// Merge the diff into the state so that we have all the attributes
	// properly.
	rs := s.MergeDiff(d)

	securityGroupOpts := ec2.SecurityGroup{
		Name: rs.Attributes["name"],
	}

	if rs.Attributes["vpc_id"] != "" {
		securityGroupOpts.VpcId = rs.Attributes["vpc_id"]
	}

	if rs.Attributes["description"] != "" {
		securityGroupOpts.Description = rs.Attributes["description"]
	}

	log.Printf("[DEBUG] Security Group create configuration: %#v", securityGroupOpts)
	createResp, err := ec2conn.CreateSecurityGroup(securityGroupOpts)
	if err != nil {
		return nil, fmt.Errorf("Error creating Security Group: %s", err)
	}

	rs.ID = createResp.Id
	group := createResp.SecurityGroup

	log.Printf("[INFO] Security Group ID: %s", rs.ID)

	// Wait for the security group to truly exist
	log.Printf(
		"[DEBUG] Waiting for SG (%s) to exist",
		s.ID)
	stateConf := &resource.StateChangeConf{
		Pending: []string{""},
		Target:  "exists",
		Refresh: SGStateRefreshFunc(ec2conn, rs.ID),
		Timeout: 1 * time.Minute,
	}
	if _, err := stateConf.WaitForState(); err != nil {
		return s, fmt.Errorf(
			"Error waiting for SG (%s) to become available: %s",
			rs.ID, err)
	}

	// Expand the "ingress" array to goamz compat []ec2.IPPerm
	ingressRules := []ec2.IPPerm{}
	v, ok := flatmap.Expand(rs.Attributes, "ingress").([]interface{})
	if ok {
		ingressRules = expandIPPerms(v)
	}

	if len(ingressRules) > 0 {
		_, err = ec2conn.AuthorizeSecurityGroup(group, ingressRules)
		if err != nil {
			return rs, fmt.Errorf("Error authorizing security group ingress rules: %s", err)
		}
	}

	return resource_aws_security_group_refresh(rs, meta)
}

func resource_aws_security_group_update(
	s *terraform.ResourceState,
	d *terraform.ResourceDiff,
	meta interface{}) (*terraform.ResourceState, error) {

	rs := s.MergeDiff(d)
	log.Printf("ResourceDiff: %s", d)
	log.Printf("ResourceState: %s", s)
	log.Printf("Merged: %s", rs)

	return nil, fmt.Errorf("Did not update")
}

func resource_aws_security_group_destroy(
	s *terraform.ResourceState,
	meta interface{}) error {
	p := meta.(*ResourceProvider)
	ec2conn := p.ec2conn

	log.Printf("[DEBUG] Security Group destroy: %v", s.ID)

	_, err := ec2conn.DeleteSecurityGroup(ec2.SecurityGroup{Id: s.ID})
	if err != nil {
		ec2err, ok := err.(*ec2.Error)
		if ok && ec2err.Code == "InvalidGroup.NotFound" {
			return nil
		}
	}

	return err
}

func resource_aws_security_group_refresh(
	s *terraform.ResourceState,
	meta interface{}) (*terraform.ResourceState, error) {
	p := meta.(*ResourceProvider)
	ec2conn := p.ec2conn

	sgRaw, _, err := SGStateRefreshFunc(ec2conn, s.ID)()
	if err != nil {
		return s, err
	}
	if sgRaw == nil {
		return nil, nil
	}

	return resource_aws_security_group_update_state(
		s, sgRaw.(*ec2.SecurityGroupInfo))
}

func resource_aws_security_group_diff(
	s *terraform.ResourceState,
	c *terraform.ResourceConfig,
	meta interface{}) (*terraform.ResourceDiff, error) {

	b := &diff.ResourceBuilder{
		Attrs: map[string]diff.AttrType{
			"name":        diff.AttrTypeCreate,
			"description": diff.AttrTypeUpdate,
			"ingress":     diff.AttrTypeUpdate,
			"vpc_id":      diff.AttrTypeCreate,
		},

		ComputedAttrs: []string{
			"owner_id",
			"vpc_id",
		},
	}

	return b.Diff(s, c)
}

func resource_aws_security_group_update_state(
	s *terraform.ResourceState,
	sg *ec2.SecurityGroupInfo) (*terraform.ResourceState, error) {

	s.Attributes["description"] = sg.Description
	s.Attributes["name"] = sg.Name
	s.Attributes["vpc_id"] = sg.VpcId
	s.Attributes["owner_id"] = sg.OwnerId

	// Flatten our ingress values
	toFlatten := make(map[string]interface{})
	toFlatten["ingress"] = flattenIPPerms(sg.IPPerms)

	for k, v := range flatmap.Flatten(toFlatten) {
		s.Attributes[k] = v
	}

	s.Dependencies = nil
	if s.Attributes["vpc_id"] != "" {
		s.Dependencies = append(s.Dependencies,
			terraform.ResourceDependency{ID: s.Attributes["vpc_id"]},
		)
	}

	return s, nil
}

func resource_aws_security_group_validation() *config.Validator {
	return &config.Validator{
		Required: []string{
			"name",
			"ingress.*",
			"ingress.*.from_port",
			"ingress.*.to_port",
			"ingress.*.protocol",
		},
		Optional: []string{
			"description",
			"vpc_id",
			"owner_id",
			"ingress.*.cidr_blocks.*",
			"ingress.*.security_groups.*",
		},
	}
}

// SGStateRefreshFunc returns a resource.StateRefreshFunc that is used to watch
// a security group.
func SGStateRefreshFunc(conn *ec2.EC2, id string) resource.StateRefreshFunc {
	return func() (interface{}, string, error) {
		sgs := []ec2.SecurityGroup{ec2.SecurityGroup{Id: id}}
		resp, err := conn.SecurityGroups(sgs, nil)
		if err != nil {
			if ec2err, ok := err.(*ec2.Error); ok {
				if ec2err.Code == "InvalidSecurityGroupID.NotFound" ||
					ec2err.Code == "InvalidGroup.NotFound" {
					resp = nil
					err = nil
				}
			}

			if err != nil {
				log.Printf("Error on SGStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil {
			return nil, "", nil
		}

		group := &resp.Groups[0]
		return group, "exists", nil
	}
}
<|endoftext|>


<commit_before>/*--------------------------------------------------------*\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: https://hprose.com                     |
|                                                          |
| rpc/core/service.go                                      |
|                                                          |
| LastModified: Feb 17, 2021                               |
| Author: Ma Bingyao <andot@hprose.com>                    |
|                                                          |
\*________________________________________________________*/

package core

import (
	"context"
	"reflect"
	"sync"
)

// Server is a generic interface used to represent any server.
type Server interface{}

// Handler is an interface used to bind service to any server.
type Handler interface {
	Bind(server Server) error
}

// HandlerFactory is a constructor for Handler.
type HandlerFactory func(service *Service) Handler

var handlerFactories sync.Map
var serverTypes sync.Map

// RegisterHandler for Service.
func RegisterHandler(name string, handlerFactory HandlerFactory, serverType reflect.Type) {
	handlerFactories.Store(name, handlerFactory)
	serverTypes.Store(serverType, name)
}

// Service for RPC.
type Service struct {
	Codec            ServiceCodec
	MaxRequestLength int
	Options          Dict
	invokeManager    PluginManager
	ioManager        PluginManager
	handlers         map[string]Handler
	methodManager
}

// NewService returns an instance of Service.
func NewService() *Service {
	service := &Service{
		Codec:            serviceCodec{},
		MaxRequestLength: 0x7FFFFFFF,
		Options:          NewSafeDict(),
	}
	handlerFactories.Range(func(key, value interface{}) bool {
		handler := value.(HandlerFactory)(service)
		service.handlers[key.(string)] = handler
		return true
	})
	service.ioManager = NewIOManager(service.Process)
	service.invokeManager = NewInvokeManager(service.Execute)
	service.AddFunction(service.methodManager.Names, "~")
	return service
}

// Bind to server.
func (s *Service) Bind(server Server) error {
	serverType := reflect.TypeOf(server)
	if name, ok := serverTypes.Load(serverType); ok {
		handler := s.handlers[name.(string)]
		return handler.Bind(server)
	}
	return UnsupportedServerTypeError{serverType}
}

// Handler returns the handler by the specified name.
func (s *Service) Handler(name string) Handler {
	return s.handlers[name]
}

// Handle the reqeust and returns the response.
func (s *Service) Handle(context context.Context, request []byte) ([]byte, error) {
	return s.ioManager.Handler().(NextIOHandler)(context, request)
}

// Process the reqeust and returns the response.
func (s *Service) Process(context context.Context, request []byte) ([]byte, error) {
	serviceContext := FromContext(context).(*ServiceContext)
	name, args, err := s.Codec.Decode(request, serviceContext)
	if err != nil {
		return s.Codec.Encode(err, serviceContext)
	}
	var result interface{}
	func() {
		defer func() {
			if p := recover(); p != nil {
				result = NewPanicError(p)
			}
		}()
		results, err := s.invokeManager.Handler().(NextInvokeHandler)(context, name, args)
		if err != nil {
			result = err
			return
		}
		switch len(results) {
		case 0:
			result = nil
		case 1:
			result = results[0]
		default:
			result = results
		}
	}()
	return s.Codec.Encode(result, serviceContext)
}

// Execute the method and returns the results.
func (s *Service) Execute(context context.Context, name string, args []interface{}) (result []interface{}, err error) {
	serviceContext := FromContext(context).(*ServiceContext)
	method := serviceContext.Method
	if method.Missing() {
		if method.PassContext() {
			return method.(contextMissingMethod)(context, name, args)
		}
		return method.(missingMethod)(name, args)
	}
	n := len(args)
	var in []reflect.Value
	if method.PassContext() {
		in = make([]reflect.Value, n+1)
		in[0] = reflect.ValueOf(context)
		for i := 0; i < n; i++ {
			in[i+1] = reflect.ValueOf(args[i])
		}
	} else {
		in = make([]reflect.Value, n)
		for i := 0; i < n; i++ {
			in[i] = reflect.ValueOf(args[i])
		}
	}
	f := method.Func()
	out := f.Call(in)
	n = len(out)
	if f.Type().Out(n - 1).Implements(errorType) {
		if !out[n-1].IsNil() {
			err = out[n-1].Interface().(error)
		}
		out = out[:n-1]
		n--
	}
	for i := 0; i < n; i++ {
		result = append(result, out[i].Interface())
	}
	return
}

func splitPluginHandlers(handlers []PluginHandler) (invokeHandlers []PluginHandler, ioHandler []PluginHandler) {
	for _, handler := range handlers {
		switch handler.(type) {
		case InvokeHandler:
			invokeHandlers = append(invokeHandlers, handler)
		case IOHandler:
			ioHandler = append(ioHandler, handler)
		}
	}
	return
}

// Use plugin handlers.
func (s *Service) Use(handler ...PluginHandler) *Service {
	invokeHandlers, ioHandler := splitPluginHandlers(handler)
	if len(invokeHandlers) > 0 {
		s.invokeManager.Use(invokeHandlers...)
	}
	if len(ioHandler) > 0 {
		s.ioManager.Use(ioHandler...)
	}
	return s
}

// Unuse plugin handlers.
func (s *Service) Unuse(handler ...PluginHandler) *Service {
	invokeHandlers, ioHandler := splitPluginHandlers(handler)
	if len(invokeHandlers) > 0 {
		s.invokeManager.Unuse(invokeHandlers...)
	}
	if len(ioHandler) > 0 {
		s.ioManager.Unuse(ioHandler...)
	}
	return s
}

// Remove is used for unpublishing method by the specified name.
func (s *Service) Remove(name string) *Service {
	s.methodManager.Remove(name)
	return s
}

// Add is used for publishing the method.
func (s *Service) Add(method Method) *Service {
	s.methodManager.Add(method)
	return s
}

// AddFunction is used for publishing function f with name.
func (s *Service) AddFunction(f interface{}, name string) *Service {
	s.methodManager.AddFunction(f, name)
	return s
}

// AddMethod is used for publishing method named name on target with alias.
func (s *Service) AddMethod(name string, target interface{}, alias ...string) *Service {
	s.methodManager.AddMethod(name, target, alias...)
	return s
}

// AddMethods is used for publishing methods named names on target with namespace.
func (s *Service) AddMethods(names []string, target interface{}, namespace ...string) *Service {
	s.methodManager.AddMethods(names, target, namespace...)
	return s
}

// AddInstanceMethods is used for publishing all the public methods and func fields with namespace.
func (s *Service) AddInstanceMethods(target interface{}, namespace ...string) *Service {
	s.methodManager.AddInstanceMethods(target, namespace...)
	return s
}

// AddAllMethods will publish all methods and non-nil function fields on the
// obj self and on its anonymous or non-anonymous struct fields (or pointer to
// pointer ... to pointer struct fields). This is a recursive operation.
// So it's a pit, if you do not know what you are doing, do not step on.
func (s *Service) AddAllMethods(target interface{}, namespace ...string) *Service {
	s.methodManager.AddAllMethods(target, namespace...)
	return s
}

// AddMissingMethod is used for publishing a method,
// all methods not explicitly published will be redirected to this method.
func (s *Service) AddMissingMethod(f interface{}) *Service {
	s.methodManager.AddMissingMethod(f)
	return s
}

// AddNetRPCMethods is used for publishing methods defined for net/rpc.
func (s *Service) AddNetRPCMethods(rcvr interface{}, namespace ...string) *Service {
	s.methodManager.AddNetRPCMethods(rcvr, namespace...)
	return s
}
<commit_msg>Improved HandlerFactory<commit_after>/*--------------------------------------------------------*\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: https://hprose.com                     |
|                                                          |
| rpc/core/service.go                                      |
|                                                          |
| LastModified: Feb 18, 2021                               |
| Author: Ma Bingyao <andot@hprose.com>                    |
|                                                          |
\*________________________________________________________*/

package core

import (
	"context"
	"reflect"
	"sync"
)

// Server is a generic interface used to represent any server.
type Server interface{}

// Handler is an interface used to bind service to any server.
type Handler interface {
	Bind(server Server)
}

// HandlerFactory is a constructor for Handler.
type HandlerFactory interface {
	ServerTypes() []reflect.Type
	New(service *Service) Handler
}

var handlerFactories sync.Map
var serverTypes sync.Map

// RegisterHandler for Service.
func RegisterHandler(name string, handlerFactory HandlerFactory) {
	handlerFactories.Store(name, handlerFactory)
	for _, serverType := range handlerFactory.ServerTypes() {
		if value, loaded := serverTypes.LoadOrStore(serverType, []string{name}); loaded {
			names := value.([]string)
			names = append(names, name)
			serverTypes.Store(serverType, names)
		}
	}
}

// Service for RPC.
type Service struct {
	Codec            ServiceCodec
	MaxRequestLength int
	Options          Dict
	invokeManager    PluginManager
	ioManager        PluginManager
	handlers         map[string]Handler
	methodManager
}

// NewService returns an instance of Service.
func NewService() *Service {
	service := &Service{
		Codec:            serviceCodec{},
		MaxRequestLength: 0x7FFFFFFF,
		Options:          NewSafeDict(),
	}
	handlerFactories.Range(func(key, value interface{}) bool {
		handler := value.(HandlerFactory).New(service)
		service.handlers[key.(string)] = handler
		return true
	})
	service.ioManager = NewIOManager(service.Process)
	service.invokeManager = NewInvokeManager(service.Execute)
	service.AddFunction(service.methodManager.Names, "~")
	return service
}

// Bind to server.
func (s *Service) Bind(server Server) error {
	serverType := reflect.TypeOf(server)
	if value, ok := serverTypes.Load(serverType); ok {
		names := value.([]string)
		for _, name := range names {
			s.handlers[name].Bind(server)
		}
	}
	return UnsupportedServerTypeError{serverType}
}

// Handler returns the handler by the specified name.
func (s *Service) Handler(name string) Handler {
	return s.handlers[name]
}

// Handle the reqeust and returns the response.
func (s *Service) Handle(context context.Context, request []byte) ([]byte, error) {
	return s.ioManager.Handler().(NextIOHandler)(context, request)
}

// Process the reqeust and returns the response.
func (s *Service) Process(context context.Context, request []byte) ([]byte, error) {
	serviceContext := FromContext(context).(*ServiceContext)
	name, args, err := s.Codec.Decode(request, serviceContext)
	if err != nil {
		return s.Codec.Encode(err, serviceContext)
	}
	var result interface{}
	func() {
		defer func() {
			if p := recover(); p != nil {
				result = NewPanicError(p)
			}
		}()
		results, err := s.invokeManager.Handler().(NextInvokeHandler)(context, name, args)
		if err != nil {
			result = err
			return
		}
		switch len(results) {
		case 0:
			result = nil
		case 1:
			result = results[0]
		default:
			result = results
		}
	}()
	return s.Codec.Encode(result, serviceContext)
}

// Execute the method and returns the results.
func (s *Service) Execute(context context.Context, name string, args []interface{}) (result []interface{}, err error) {
	serviceContext := FromContext(context).(*ServiceContext)
	method := serviceContext.Method
	if method.Missing() {
		if method.PassContext() {
			return method.(contextMissingMethod)(context, name, args)
		}
		return method.(missingMethod)(name, args)
	}
	n := len(args)
	var in []reflect.Value
	if method.PassContext() {
		in = make([]reflect.Value, n+1)
		in[0] = reflect.ValueOf(context)
		for i := 0; i < n; i++ {
			in[i+1] = reflect.ValueOf(args[i])
		}
	} else {
		in = make([]reflect.Value, n)
		for i := 0; i < n; i++ {
			in[i] = reflect.ValueOf(args[i])
		}
	}
	f := method.Func()
	out := f.Call(in)
	n = len(out)
	if f.Type().Out(n - 1).Implements(errorType) {
		if !out[n-1].IsNil() {
			err = out[n-1].Interface().(error)
		}
		out = out[:n-1]
		n--
	}
	for i := 0; i < n; i++ {
		result = append(result, out[i].Interface())
	}
	return
}

func splitPluginHandlers(handlers []PluginHandler) (invokeHandlers []PluginHandler, ioHandler []PluginHandler) {
	for _, handler := range handlers {
		switch handler.(type) {
		case InvokeHandler:
			invokeHandlers = append(invokeHandlers, handler)
		case IOHandler:
			ioHandler = append(ioHandler, handler)
		}
	}
	return
}

// Use plugin handlers.
func (s *Service) Use(handler ...PluginHandler) *Service {
	invokeHandlers, ioHandler := splitPluginHandlers(handler)
	if len(invokeHandlers) > 0 {
		s.invokeManager.Use(invokeHandlers...)
	}
	if len(ioHandler) > 0 {
		s.ioManager.Use(ioHandler...)
	}
	return s
}

// Unuse plugin handlers.
func (s *Service) Unuse(handler ...PluginHandler) *Service {
	invokeHandlers, ioHandler := splitPluginHandlers(handler)
	if len(invokeHandlers) > 0 {
		s.invokeManager.Unuse(invokeHandlers...)
	}
	if len(ioHandler) > 0 {
		s.ioManager.Unuse(ioHandler...)
	}
	return s
}

// Remove is used for unpublishing method by the specified name.
func (s *Service) Remove(name string) *Service {
	s.methodManager.Remove(name)
	return s
}

// Add is used for publishing the method.
func (s *Service) Add(method Method) *Service {
	s.methodManager.Add(method)
	return s
}

// AddFunction is used for publishing function f with name.
func (s *Service) AddFunction(f interface{}, name string) *Service {
	s.methodManager.AddFunction(f, name)
	return s
}

// AddMethod is used for publishing method named name on target with alias.
func (s *Service) AddMethod(name string, target interface{}, alias ...string) *Service {
	s.methodManager.AddMethod(name, target, alias...)
	return s
}

// AddMethods is used for publishing methods named names on target with namespace.
func (s *Service) AddMethods(names []string, target interface{}, namespace ...string) *Service {
	s.methodManager.AddMethods(names, target, namespace...)
	return s
}

// AddInstanceMethods is used for publishing all the public methods and func fields with namespace.
func (s *Service) AddInstanceMethods(target interface{}, namespace ...string) *Service {
	s.methodManager.AddInstanceMethods(target, namespace...)
	return s
}

// AddAllMethods will publish all methods and non-nil function fields on the
// obj self and on its anonymous or non-anonymous struct fields (or pointer to
// pointer ... to pointer struct fields). This is a recursive operation.
// So it's a pit, if you do not know what you are doing, do not step on.
func (s *Service) AddAllMethods(target interface{}, namespace ...string) *Service {
	s.methodManager.AddAllMethods(target, namespace...)
	return s
}

// AddMissingMethod is used for publishing a method,
// all methods not explicitly published will be redirected to this method.
func (s *Service) AddMissingMethod(f interface{}) *Service {
	s.methodManager.AddMissingMethod(f)
	return s
}

// AddNetRPCMethods is used for publishing methods defined for net/rpc.
func (s *Service) AddNetRPCMethods(rcvr interface{}, namespace ...string) *Service {
	s.methodManager.AddNetRPCMethods(rcvr, namespace...)
	return s
}
<|endoftext|>


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

// Copyright 2015 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package rsets

import (
	"github.com/ngaut/log"
	"github.com/pingcap/tidb/context"
	"github.com/pingcap/tidb/expression"
	"github.com/pingcap/tidb/parser/opcode"
	"github.com/pingcap/tidb/plan"
	"github.com/pingcap/tidb/plan/plans"
	"github.com/pingcap/tidb/util/types"
)

var (
	_ plan.Planner = (*WhereRset)(nil)
)

// WhereRset is record set for where filter.
type WhereRset struct {
	Expr expression.Expression
	Src  plan.Plan
}

func (r *WhereRset) planBinOp(ctx context.Context, x *expression.BinaryOperation) (plan.Plan, error) {
	var err error
	var p2 plan.Plan
	var filtered bool

	p := r.Src
	switch x.Op {
	case opcode.EQ, opcode.GE, opcode.GT, opcode.LE, opcode.LT, opcode.NE:
		if p2, filtered, err = p.Filter(ctx, x); err != nil {
			return nil, err
		}
		if filtered {
			return p2, nil
		}
	case opcode.AndAnd:
		var in []expression.Expression
		var f func(expression.Expression)
		f = func(e expression.Expression) {
			b, ok := e.(*expression.BinaryOperation)
			if !ok || b.Op != opcode.AndAnd {
				in = append(in, e)
				return
			}

			f(b.L)
			f(b.R)
		}
		f(x)
		out := []expression.Expression{}
		p := r.Src
		isNewPlan := false
		for _, e := range in {
			p2, filtered, err = p.Filter(ctx, e)
			if err != nil {
				return nil, err
			}

			if !filtered {
				out = append(out, e)
				continue
			}

			p = p2
			isNewPlan = true
		}

		if !isNewPlan {
			break
		}

		if len(out) == 0 {
			return p, nil
		}

		for len(out) > 1 {
			n := len(out)
			e := expression.NewBinaryOperation(opcode.AndAnd, out[n-2], out[n-1])

			out = out[:n-1]
			out[n-2] = e
		}

		return &plans.FilterDefaultPlan{Plan: p, Expr: out[0]}, nil
	default:
		// TODO: better plan for `OR`.
		log.Warn("TODO: better plan for", x.Op)
	}

	return &plans.FilterDefaultPlan{Plan: p, Expr: x}, nil
}

func (r *WhereRset) planIdent(ctx context.Context, x *expression.Ident) (plan.Plan, error) {
	p := r.Src
	p2, filtered, err := p.Filter(ctx, x)
	if err != nil {
		return nil, err
	}

	if filtered {
		return p2, nil
	}

	return &plans.FilterDefaultPlan{Plan: p, Expr: x}, nil
}

func (r *WhereRset) planIsNull(ctx context.Context, x *expression.IsNull) (plan.Plan, error) {
	p := r.Src

	cns := expression.MentionedColumns(x.Expr)
	if len(cns) == 0 {
		return &plans.FilterDefaultPlan{Plan: p, Expr: x}, nil
	}

	p2, filtered, err := p.Filter(ctx, x)
	if err != nil {
		return nil, err
	}

	if filtered {
		return p2, nil
	}

	return &plans.FilterDefaultPlan{Plan: p, Expr: x}, nil
}

func (r *WhereRset) planUnaryOp(ctx context.Context, x *expression.UnaryOperation) (plan.Plan, error) {
	p := r.Src
	p2, filtered, err := p.Filter(ctx, x)
	if err != nil {
		return nil, err
	}

	if filtered {
		return p2, nil
	}

	return &plans.FilterDefaultPlan{Plan: p, Expr: x}, nil
}

func (r *WhereRset) planStatic(ctx context.Context, e expression.Expression) (plan.Plan, error) {
	val, err := e.Eval(nil, nil)
	if err != nil {
		return nil, err
	}

	if val == nil {
		// like `select * from t where null`.
		return &plans.NullPlan{Fields: r.Src.GetFields()}, nil
	}

	n, err := types.ToBool(val)
	if err != nil {
		return nil, err
	}

	if n == 0 {
		// like `select * from t where 0`.
		return &plans.NullPlan{Fields: r.Src.GetFields()}, nil
	}

	return &plans.FilterDefaultPlan{Plan: r.Src, Expr: e}, nil
}

// Plan gets NullPlan/FilterDefaultPlan.
func (r *WhereRset) Plan(ctx context.Context) (plan.Plan, error) {
	expr := r.Expr.Clone()
	if expr.IsStatic() {
		// IsStaic means we have a const value for where condition, and we don't need any index.
		return r.planStatic(ctx, expr)
	}

	switch x := expr.(type) {
	case *expression.BinaryOperation:
		return r.planBinOp(ctx, x)
	case *expression.Ident:
		return r.planIdent(ctx, x)
	case *expression.IsNull:
		return r.planIsNull(ctx, x)
	case *expression.PatternIn:
		// TODO: optimize
		// TODO: show plan
	case *expression.PatternLike:
		// TODO: optimize
	case *expression.PatternRegexp:
		// TODO: optimize
	case *expression.UnaryOperation:
		return r.planUnaryOp(ctx, x)
	default:
		log.Warnf("%v not supported in where rset now", r.Expr)
	}

	return &plans.FilterDefaultPlan{Plan: r.Src, Expr: expr}, nil
}
<commit_msg>rsets: should check where  condition even using index plan.<commit_after>// Copyright 2014 The ql Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSES/QL-LICENSE file.

// Copyright 2015 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package rsets

import (
	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/tidb/context"
	"github.com/pingcap/tidb/expression"
	"github.com/pingcap/tidb/parser/opcode"
	"github.com/pingcap/tidb/plan"
	"github.com/pingcap/tidb/plan/plans"
	"github.com/pingcap/tidb/util/types"
)

var (
	_ plan.Planner = (*WhereRset)(nil)
)

// WhereRset is record set for where filter.
type WhereRset struct {
	Expr expression.Expression
	Src  plan.Plan
}

func (r *WhereRset) planBinOp(ctx context.Context, x *expression.BinaryOperation) (plan.Plan, error) {
	var err error
	var p2 plan.Plan
	var filtered bool

	p := r.Src
	switch x.Op {
	case opcode.EQ, opcode.GE, opcode.GT, opcode.LE, opcode.LT, opcode.NE:
		if p2, filtered, err = p.Filter(ctx, x); err != nil {
			return nil, err
		}
		if filtered {
			return p2, nil
		}
	case opcode.AndAnd:
		var in []expression.Expression
		var f func(expression.Expression)
		f = func(e expression.Expression) {
			b, ok := e.(*expression.BinaryOperation)
			if !ok || b.Op != opcode.AndAnd {
				in = append(in, e)
				return
			}

			f(b.L)
			f(b.R)
		}
		f(x)
		out := []expression.Expression{}
		p := r.Src
		isNewPlan := false
		for _, e := range in {
			p2, filtered, err = p.Filter(ctx, e)
			if err != nil {
				return nil, err
			}

			if !filtered {
				out = append(out, e)
				continue
			}

			p = p2
			isNewPlan = true
		}

		if !isNewPlan {
			break
		}

		if len(out) == 0 {
			return p, nil
		}

		for len(out) > 1 {
			n := len(out)
			e := expression.NewBinaryOperation(opcode.AndAnd, out[n-2], out[n-1])

			out = out[:n-1]
			out[n-2] = e
		}

		return &plans.FilterDefaultPlan{Plan: p, Expr: out[0]}, nil
	default:
		// TODO: better plan for `OR`.
		log.Warn("TODO: better plan for", x.Op)
	}

	return &plans.FilterDefaultPlan{Plan: p, Expr: x}, nil
}

func (r *WhereRset) planIdent(ctx context.Context, x *expression.Ident) (plan.Plan, error) {
	p := r.Src
	p2, filtered, err := p.Filter(ctx, x)
	if err != nil {
		return nil, err
	}

	if filtered {
		return p2, nil
	}

	return &plans.FilterDefaultPlan{Plan: p, Expr: x}, nil
}

func (r *WhereRset) planIsNull(ctx context.Context, x *expression.IsNull) (plan.Plan, error) {
	p := r.Src

	cns := expression.MentionedColumns(x.Expr)
	if len(cns) == 0 {
		return &plans.FilterDefaultPlan{Plan: p, Expr: x}, nil
	}

	p2, filtered, err := p.Filter(ctx, x)
	if err != nil {
		return nil, err
	}

	if filtered {
		return p2, nil
	}

	return &plans.FilterDefaultPlan{Plan: p, Expr: x}, nil
}

func (r *WhereRset) planUnaryOp(ctx context.Context, x *expression.UnaryOperation) (plan.Plan, error) {
	p := r.Src
	p2, filtered, err := p.Filter(ctx, x)
	if err != nil {
		return nil, err
	}

	if filtered {
		return p2, nil
	}

	return &plans.FilterDefaultPlan{Plan: p, Expr: x}, nil
}

func (r *WhereRset) planStatic(ctx context.Context, e expression.Expression) (plan.Plan, error) {
	val, err := e.Eval(nil, nil)
	if err != nil {
		return nil, err
	}

	if val == nil {
		// like `select * from t where null`.
		return &plans.NullPlan{Fields: r.Src.GetFields()}, nil
	}

	n, err := types.ToBool(val)
	if err != nil {
		return nil, err
	}

	if n == 0 {
		// like `select * from t where 0`.
		return &plans.NullPlan{Fields: r.Src.GetFields()}, nil
	}

	return &plans.FilterDefaultPlan{Plan: r.Src, Expr: e}, nil
}

// Plan gets NullPlan/FilterDefaultPlan.
func (r *WhereRset) Plan(ctx context.Context) (plan.Plan, error) {
	expr := r.Expr.Clone()
	if expr.IsStatic() {
		// IsStaic means we have a const value for where condition, and we don't need any index.
		return r.planStatic(ctx, expr)
	}

	var (
		src = r.Src
		err error
	)

	switch x := expr.(type) {
	case *expression.BinaryOperation:
		src, err = r.planBinOp(ctx, x)
	case *expression.Ident:
		src, err = r.planIdent(ctx, x)
	case *expression.IsNull:
		src, err = r.planIsNull(ctx, x)
	case *expression.PatternIn:
		// TODO: optimize
		// TODO: show plan
	case *expression.PatternLike:
		// TODO: optimize
	case *expression.PatternRegexp:
		// TODO: optimize
	case *expression.UnaryOperation:
		src, err = r.planUnaryOp(ctx, x)
	default:
		log.Warnf("%v not supported in where rset now", r.Expr)
	}

	if err != nil {
		return nil, errors.Trace(err)
	}

	if _, ok := src.(*plans.FilterDefaultPlan); ok {
		return src, nil
	}

	// We must use a FilterDefaultPlan here to wrap filtered plan.
	// Alghough we can check where condition using index plan, we still need
	// to check again after the FROM phase if the FROM phase contains outer join.
	// TODO: if FROM phase doesn't contain outer join, we can return filtered plan directly.
	return &plans.FilterDefaultPlan{Plan: src, Expr: expr}, nil
}
<|endoftext|>


<commit_before>// Copyright 2020 The LUCI 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 main is the main entry point for the app.
package main

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"go.chromium.org/luci/common/data/rand/mathrand"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/proto/access"
	"go.chromium.org/luci/config/server/cfgmodule"
	"go.chromium.org/luci/grpc/prpc"
	"go.chromium.org/luci/server"
	"go.chromium.org/luci/server/cron"
	"go.chromium.org/luci/server/gaeemulation"
	"go.chromium.org/luci/server/module"
	"go.chromium.org/luci/server/router"
	"go.chromium.org/luci/server/tq"

	// Enable datastore transactional tasks support.
	_ "go.chromium.org/luci/server/tq/txn/datastore"

	"go.chromium.org/luci/buildbucket/appengine/internal/config"
	"go.chromium.org/luci/buildbucket/appengine/rpc"
	pb "go.chromium.org/luci/buildbucket/proto"
)

// isBeefy returns whether the request was intended for the beefy service.
func isBeefy(req *http.Request) bool {
	return strings.Contains(req.Host, "beefy")
}

// isDev returns whether the request was intended for the dev instance.
func isDev(req *http.Request) bool {
	return strings.HasSuffix(req.Host, "-dev.appspot.com")
}

func main() {
	mods := []module.Module{
		cfgmodule.NewModuleFromFlags(),
		cron.NewModuleFromFlags(),
		gaeemulation.NewModuleFromFlags(),
		tq.NewModuleFromFlags(),
	}

	server.Main(nil, mods, func(srv *server.Server) error {
		// Proxy buildbucket.v2.Builds pRPC requests back to the Python
		// service in order to achieve a programmatic traffic split.
		// Because of the way dispatch routes work, requests are proxied
		// to a copy of the Python service hosted at a different path.
		// TODO(crbug/1042991): Remove the proxy once the go service handles all traffic.
		pythonURL, err := url.Parse(fmt.Sprintf("https://default-dot-%s.appspot.com/python", srv.Options.CloudProject))
		if err != nil {
			panic(err)
		}
		beefyURL, err := url.Parse(fmt.Sprintf("https://beefy-dot-%s.appspot.com/python", srv.Options.CloudProject))
		if err != nil {
			panic(err)
		}
		prx := httputil.NewSingleHostReverseProxy(pythonURL)
		prx.Director = func(req *http.Request) {
			target := pythonURL
			if isBeefy(req) {
				target = beefyURL
			}
			// According to net.Request documentation, setting Host is unnecessary
			// because URL.Host is supposed to be used for outbound requests.
			// However, on GAE, it seems that req.Host is incorrectly used.
			req.Host = target.Host
			req.URL.Scheme = target.Scheme
			req.URL.Host = target.Host
			req.URL.Path = fmt.Sprintf("%s%s", target.Path, req.URL.Path)
		}
		// makeOverride returns a prpc.Override which allows the given percentage of requests
		// through to this service, proxying the remainder to Python.
		makeOverride := func(prodPct, devPct int) func(*router.Context) bool {
			return func(ctx *router.Context) bool {
				// TODO(crbug/1090540): remove env k-v
				ctx.Context = context.WithValue(ctx.Context, "env", "Prod")
				pct := prodPct
				if isDev(ctx.Request) {
					pct = devPct
					// TODO(crbug/1090540): remove env k-v
					ctx.Context = context.WithValue(ctx.Context, "env", "Dev")
				}
				switch val := ctx.Request.Header.Get("Should-Proxy"); val {
				case "true":
					pct = 0
					logging.Debugf(ctx.Context, "request demanded to be proxied")
				case "false":
					pct = 100
					logging.Debugf(ctx.Context, "request demanded not to be proxied")
				}
				ctx.Context = rpc.WithTrafficSplit(ctx.Context, pct)
				if mathrand.Intn(ctx.Context, 100) < pct {
					return false
				}
				target := pythonURL
				if isBeefy(ctx.Request) {
					target = beefyURL
				}
				logging.Debugf(ctx.Context, "proxying request to %s", target)
				prx.ServeHTTP(ctx.Writer, ctx.Request)
				return true
			}
		}

		srv.PRPC.AccessControl = prpc.AllowOriginAll
		access.RegisterAccessServer(srv.PRPC, &access.UnimplementedAccessServer{})
		pb.RegisterBuildsServer(srv.PRPC, rpc.NewBuilds())
		pb.RegisterBuildersServer(srv.PRPC, rpc.NewBuilders())
		// TODO(crbug/1082369): Remove this workaround once field masks can be decoded.
		srv.PRPC.HackFixFieldMasksForJSON = true

		// makeOverride(prod % -> Go, dev % -> Go).
		srv.PRPC.RegisterOverride("buildbucket.v2.Builds", "ScheduleBuild", makeOverride(10, 100))

		cron.RegisterHandler("update_config", config.UpdateSettingsCfg)
		return nil
	})
}
<commit_msg>[buildbucket] Set prod ScheduleBuild experiment to 0%<commit_after>// Copyright 2020 The LUCI 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 main is the main entry point for the app.
package main

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"go.chromium.org/luci/common/data/rand/mathrand"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/proto/access"
	"go.chromium.org/luci/config/server/cfgmodule"
	"go.chromium.org/luci/grpc/prpc"
	"go.chromium.org/luci/server"
	"go.chromium.org/luci/server/cron"
	"go.chromium.org/luci/server/gaeemulation"
	"go.chromium.org/luci/server/module"
	"go.chromium.org/luci/server/router"
	"go.chromium.org/luci/server/tq"

	// Enable datastore transactional tasks support.
	_ "go.chromium.org/luci/server/tq/txn/datastore"

	"go.chromium.org/luci/buildbucket/appengine/internal/config"
	"go.chromium.org/luci/buildbucket/appengine/rpc"
	pb "go.chromium.org/luci/buildbucket/proto"
)

// isBeefy returns whether the request was intended for the beefy service.
func isBeefy(req *http.Request) bool {
	return strings.Contains(req.Host, "beefy")
}

// isDev returns whether the request was intended for the dev instance.
func isDev(req *http.Request) bool {
	return strings.HasSuffix(req.Host, "-dev.appspot.com")
}

func main() {
	mods := []module.Module{
		cfgmodule.NewModuleFromFlags(),
		cron.NewModuleFromFlags(),
		gaeemulation.NewModuleFromFlags(),
		tq.NewModuleFromFlags(),
	}

	server.Main(nil, mods, func(srv *server.Server) error {
		// Proxy buildbucket.v2.Builds pRPC requests back to the Python
		// service in order to achieve a programmatic traffic split.
		// Because of the way dispatch routes work, requests are proxied
		// to a copy of the Python service hosted at a different path.
		// TODO(crbug/1042991): Remove the proxy once the go service handles all traffic.
		pythonURL, err := url.Parse(fmt.Sprintf("https://default-dot-%s.appspot.com/python", srv.Options.CloudProject))
		if err != nil {
			panic(err)
		}
		beefyURL, err := url.Parse(fmt.Sprintf("https://beefy-dot-%s.appspot.com/python", srv.Options.CloudProject))
		if err != nil {
			panic(err)
		}
		prx := httputil.NewSingleHostReverseProxy(pythonURL)
		prx.Director = func(req *http.Request) {
			target := pythonURL
			if isBeefy(req) {
				target = beefyURL
			}
			// According to net.Request documentation, setting Host is unnecessary
			// because URL.Host is supposed to be used for outbound requests.
			// However, on GAE, it seems that req.Host is incorrectly used.
			req.Host = target.Host
			req.URL.Scheme = target.Scheme
			req.URL.Host = target.Host
			req.URL.Path = fmt.Sprintf("%s%s", target.Path, req.URL.Path)
		}
		// makeOverride returns a prpc.Override which allows the given percentage of requests
		// through to this service, proxying the remainder to Python.
		makeOverride := func(prodPct, devPct int) func(*router.Context) bool {
			return func(ctx *router.Context) bool {
				// TODO(crbug/1090540): remove env k-v
				ctx.Context = context.WithValue(ctx.Context, "env", "Prod")
				pct := prodPct
				if isDev(ctx.Request) {
					pct = devPct
					// TODO(crbug/1090540): remove env k-v
					ctx.Context = context.WithValue(ctx.Context, "env", "Dev")
				}
				switch val := ctx.Request.Header.Get("Should-Proxy"); val {
				case "true":
					pct = 0
					logging.Debugf(ctx.Context, "request demanded to be proxied")
				case "false":
					pct = 100
					logging.Debugf(ctx.Context, "request demanded not to be proxied")
				}
				ctx.Context = rpc.WithTrafficSplit(ctx.Context, pct)
				if mathrand.Intn(ctx.Context, 100) < pct {
					return false
				}
				target := pythonURL
				if isBeefy(ctx.Request) {
					target = beefyURL
				}
				logging.Debugf(ctx.Context, "proxying request to %s", target)
				prx.ServeHTTP(ctx.Writer, ctx.Request)
				return true
			}
		}

		srv.PRPC.AccessControl = prpc.AllowOriginAll
		access.RegisterAccessServer(srv.PRPC, &access.UnimplementedAccessServer{})
		pb.RegisterBuildsServer(srv.PRPC, rpc.NewBuilds())
		pb.RegisterBuildersServer(srv.PRPC, rpc.NewBuilders())
		// TODO(crbug/1082369): Remove this workaround once field masks can be decoded.
		srv.PRPC.HackFixFieldMasksForJSON = true

		// makeOverride(prod % -> Go, dev % -> Go).
		srv.PRPC.RegisterOverride("buildbucket.v2.Builds", "ScheduleBuild", makeOverride(0, 100))

		cron.RegisterHandler("update_config", config.UpdateSettingsCfg)
		return nil
	})
}
<|endoftext|>


<commit_before>// Copyright 2020 The LUCI 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 main is the main entry point for the app.
package main

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"go.chromium.org/luci/common/data/rand/mathrand"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/proto/access"
	"go.chromium.org/luci/grpc/prpc"
	"go.chromium.org/luci/server"
	"go.chromium.org/luci/server/gaeemulation"
	"go.chromium.org/luci/server/module"
	"go.chromium.org/luci/server/router"
	"go.chromium.org/luci/server/tq"

	// Enable datastore transactional tasks support.
	_ "go.chromium.org/luci/server/tq/txn/datastore"

	"go.chromium.org/luci/buildbucket/appengine/rpc"
	pb "go.chromium.org/luci/buildbucket/proto"
)

// isBeefy returns whether the request was intended for the beefy service.
func isBeefy(req *http.Request) bool {
	return strings.Contains(req.Host, "beefy")
}

// isDev returns whether the request was intended for the dev instance.
func isDev(req *http.Request) bool {
	return strings.HasSuffix(req.Host, "-dev.appspot.com")
}

func main() {
	mods := []module.Module{
		gaeemulation.NewModuleFromFlags(),
		tq.NewModuleFromFlags(),
	}

	server.Main(nil, mods, func(srv *server.Server) error {
		// Proxy buildbucket.v2.Builds pRPC requests back to the Python
		// service in order to achieve a programmatic traffic split.
		// Because of the way dispatch routes work, requests are proxied
		// to a copy of the Python service hosted at a different path.
		// TODO(crbug/1042991): Remove the proxy once the go service handles all traffic.
		pythonURL, err := url.Parse(fmt.Sprintf("https://default-dot-%s.appspot.com/python", srv.Options.CloudProject))
		if err != nil {
			panic(err)
		}
		beefyURL, err := url.Parse(fmt.Sprintf("https://beefy-dot-%s.appspot.com/python", srv.Options.CloudProject))
		if err != nil {
			panic(err)
		}
		prx := httputil.NewSingleHostReverseProxy(pythonURL)
		prx.Director = func(req *http.Request) {
			target := pythonURL
			if isBeefy(req) {
				target = beefyURL
			}
			// According to net.Request documentation, setting Host is unnecessary
			// because URL.Host is supposed to be used for outbound requests.
			// However, on GAE, it seems that req.Host is incorrectly used.
			req.Host = target.Host
			req.URL.Scheme = target.Scheme
			req.URL.Host = target.Host
			req.URL.Path = fmt.Sprintf("%s%s", target.Path, req.URL.Path)
		}
		// makeOverride returns a prpc.Override which allows the given percentage of requests
		// through to this service, proxying the remainder to Python.
		makeOverride := func(prodPct, devPct int) func(*router.Context) bool {
			return func(ctx *router.Context) bool {
				// TODO(crbug/1090540): remove env k-v
				ctx.Context = context.WithValue(ctx.Context, "env", "Prod")
				pct := prodPct
				if isDev(ctx.Request) {
					pct = devPct
					// TODO(crbug/1090540): remove env k-v
					ctx.Context = context.WithValue(ctx.Context, "env", "Dev")
				}
				switch val := ctx.Request.Header.Get("Should-Proxy"); val {
				case "true":
					pct = 0
					logging.Debugf(ctx.Context, "request demanded to be proxied")
				case "false":
					pct = 100
					logging.Debugf(ctx.Context, "request demanded not to be proxied")
				}
				if mathrand.Intn(ctx.Context, 100) < pct {
					return false
				}
				target := pythonURL
				if isBeefy(ctx.Request) {
					target = beefyURL
				}
				logging.Debugf(ctx.Context, "proxying request to %s", target)
				prx.ServeHTTP(ctx.Writer, ctx.Request)
				return true
			}
		}

		srv.PRPC.AccessControl = prpc.AllowOriginAll
		access.RegisterAccessServer(srv.PRPC, &access.UnimplementedAccessServer{})
		pb.RegisterBuildsServer(srv.PRPC, rpc.NewBuilds())
		pb.RegisterBuildersServer(srv.PRPC, rpc.NewBuilders())
		// TODO(crbug/1082369): Remove this workaround once field masks can be decoded.
		srv.PRPC.HackFixFieldMasksForJSON = true

		// makeOverride(prod % -> Go, dev % -> Go).
		srv.PRPC.RegisterOverride("buildbucket.v2.Builds", "Batch", makeOverride(0, 0))
		srv.PRPC.RegisterOverride("buildbucket.v2.Builds", "CancelBuild", makeOverride(0, 0))
		srv.PRPC.RegisterOverride("buildbucket.v2.Builds", "SearchBuilds", makeOverride(0, 10))
		srv.PRPC.RegisterOverride("buildbucket.v2.Builds", "ScheduleBuild", makeOverride(0, 0))
		srv.PRPC.RegisterOverride("buildbucket.v2.Builds", "UpdateBuild", makeOverride(0, 0))
		return nil
	})
}
<commit_msg>[buildbucket] Randomly sample 1% of prod traffic for search comparison<commit_after>// Copyright 2020 The LUCI 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 main is the main entry point for the app.
package main

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strings"

	"go.chromium.org/luci/common/data/rand/mathrand"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/proto/access"
	"go.chromium.org/luci/grpc/prpc"
	"go.chromium.org/luci/server"
	"go.chromium.org/luci/server/gaeemulation"
	"go.chromium.org/luci/server/module"
	"go.chromium.org/luci/server/router"
	"go.chromium.org/luci/server/tq"

	// Enable datastore transactional tasks support.
	_ "go.chromium.org/luci/server/tq/txn/datastore"

	"go.chromium.org/luci/buildbucket/appengine/rpc"
	pb "go.chromium.org/luci/buildbucket/proto"
)

// isBeefy returns whether the request was intended for the beefy service.
func isBeefy(req *http.Request) bool {
	return strings.Contains(req.Host, "beefy")
}

// isDev returns whether the request was intended for the dev instance.
func isDev(req *http.Request) bool {
	return strings.HasSuffix(req.Host, "-dev.appspot.com")
}

func main() {
	mods := []module.Module{
		gaeemulation.NewModuleFromFlags(),
		tq.NewModuleFromFlags(),
	}

	server.Main(nil, mods, func(srv *server.Server) error {
		// Proxy buildbucket.v2.Builds pRPC requests back to the Python
		// service in order to achieve a programmatic traffic split.
		// Because of the way dispatch routes work, requests are proxied
		// to a copy of the Python service hosted at a different path.
		// TODO(crbug/1042991): Remove the proxy once the go service handles all traffic.
		pythonURL, err := url.Parse(fmt.Sprintf("https://default-dot-%s.appspot.com/python", srv.Options.CloudProject))
		if err != nil {
			panic(err)
		}
		beefyURL, err := url.Parse(fmt.Sprintf("https://beefy-dot-%s.appspot.com/python", srv.Options.CloudProject))
		if err != nil {
			panic(err)
		}
		prx := httputil.NewSingleHostReverseProxy(pythonURL)
		prx.Director = func(req *http.Request) {
			target := pythonURL
			if isBeefy(req) {
				target = beefyURL
			}
			// According to net.Request documentation, setting Host is unnecessary
			// because URL.Host is supposed to be used for outbound requests.
			// However, on GAE, it seems that req.Host is incorrectly used.
			req.Host = target.Host
			req.URL.Scheme = target.Scheme
			req.URL.Host = target.Host
			req.URL.Path = fmt.Sprintf("%s%s", target.Path, req.URL.Path)
		}
		// makeOverride returns a prpc.Override which allows the given percentage of requests
		// through to this service, proxying the remainder to Python.
		makeOverride := func(prodPct, devPct int) func(*router.Context) bool {
			return func(ctx *router.Context) bool {
				// TODO(crbug/1090540): remove env k-v
				ctx.Context = context.WithValue(ctx.Context, "env", "Prod")
				pct := prodPct
				if isDev(ctx.Request) {
					pct = devPct
					// TODO(crbug/1090540): remove env k-v
					ctx.Context = context.WithValue(ctx.Context, "env", "Dev")
				}
				switch val := ctx.Request.Header.Get("Should-Proxy"); val {
				case "true":
					pct = 0
					logging.Debugf(ctx.Context, "request demanded to be proxied")
				case "false":
					pct = 100
					logging.Debugf(ctx.Context, "request demanded not to be proxied")
				}
				if mathrand.Intn(ctx.Context, 100) < pct {
					return false
				}
				target := pythonURL
				if isBeefy(ctx.Request) {
					target = beefyURL
				}
				logging.Debugf(ctx.Context, "proxying request to %s", target)
				prx.ServeHTTP(ctx.Writer, ctx.Request)
				return true
			}
		}

		srv.PRPC.AccessControl = prpc.AllowOriginAll
		access.RegisterAccessServer(srv.PRPC, &access.UnimplementedAccessServer{})
		pb.RegisterBuildsServer(srv.PRPC, rpc.NewBuilds())
		pb.RegisterBuildersServer(srv.PRPC, rpc.NewBuilders())
		// TODO(crbug/1082369): Remove this workaround once field masks can be decoded.
		srv.PRPC.HackFixFieldMasksForJSON = true

		// makeOverride(prod % -> Go, dev % -> Go).
		srv.PRPC.RegisterOverride("buildbucket.v2.Builds", "Batch", makeOverride(0, 0))
		srv.PRPC.RegisterOverride("buildbucket.v2.Builds", "CancelBuild", makeOverride(0, 0))
		srv.PRPC.RegisterOverride("buildbucket.v2.Builds", "SearchBuilds", makeOverride(1, 10))
		srv.PRPC.RegisterOverride("buildbucket.v2.Builds", "ScheduleBuild", makeOverride(0, 0))
		srv.PRPC.RegisterOverride("buildbucket.v2.Builds", "UpdateBuild", makeOverride(0, 0))
		return nil
	})
}
<|endoftext|>