<commit_before>// Copyright (C) 2015 Scaleway. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE.md file.

// scw helpers

// Package utils contains helpers
package utils

import (
	"crypto/md5"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"reflect"
	"regexp"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/Sirupsen/logrus"
	log "github.com/Sirupsen/logrus"
	"github.com/moul/gotty-client"
	"github.com/scaleway/scaleway-cli/pkg/sshcommand"
)

// SpawnRedirection is used to redirects the fluxes
type SpawnRedirection struct {
	Stdin  io.Reader
	Stdout io.Writer
	Stderr io.Writer
}

// SSHExec executes a command over SSH and redirects file-descriptors
func SSHExec(publicIPAddress string, privateIPAddress string, command []string, checkConnection bool, gateway string) error {
	gatewayUser := "root"
	gatewayIPAddress := gateway
	if strings.Contains(gateway, "@") {
		parts := strings.Split(gatewayIPAddress, "@")
		gatewayUser = parts[0]
		gatewayIPAddress = parts[1]
		gateway = gatewayUser + "@" + gatewayIPAddress
	}

	if publicIPAddress == "" && gatewayIPAddress == "" {
		return errors.New("server does not have public IP")
	}
	if privateIPAddress == "" && gatewayIPAddress != "" {
		return errors.New("server does not have private IP")
	}

	if checkConnection {
		useGateway := gatewayIPAddress != ""
		if useGateway && !IsTCPPortOpen(fmt.Sprintf("%s:22", gatewayIPAddress)) {
			return errors.New("gateway is not available, try again later")
		}
		if !useGateway && !IsTCPPortOpen(fmt.Sprintf("%s:22", publicIPAddress)) {
			return errors.New("server is not ready, try again later")
		}
	}

	sshCommand := NewSSHExecCmd(publicIPAddress, privateIPAddress, true, command, gateway)

	log.Debugf("Executing: %s", sshCommand)

	spawn := exec.Command("ssh", sshCommand.Slice()[1:]...)
	spawn.Stdout = os.Stdout
	spawn.Stdin = os.Stdin
	spawn.Stderr = os.Stderr
	return spawn.Run()
}

// NewSSHExecCmd computes execve compatible arguments to run a command via ssh
func NewSSHExecCmd(publicIPAddress string, privateIPAddress string, allocateTTY bool, command []string, gatewayIPAddress string) *sshcommand.Command {
	quiet := os.Getenv("DEBUG") != "1"
	secureExec := os.Getenv("SCW_SECURE_EXEC") == "1"
	sshCommand := &sshcommand.Command{
		AllocateTTY:         true,
		Command:             command,
		Host:                publicIPAddress,
		Quiet:               quiet,
		SkipHostKeyChecking: !secureExec,
		User:                "root",
		NoEscapeCommand:     true,
	}
	if gatewayIPAddress != "" {
		sshCommand.Host = privateIPAddress
		sshCommand.Gateway = &sshcommand.Command{
			Host:                gatewayIPAddress,
			SkipHostKeyChecking: !secureExec,
			AllocateTTY:         true,
			Quiet:               quiet,
			User:                "root",
		}
	}

	return sshCommand
}

// GeneratingAnSSHKey generates an SSH key
func GeneratingAnSSHKey(cfg SpawnRedirection, path string, name string) (string, error) {
	args := []string{
		"-t",
		"rsa",
		"-b",
		"4096",
		"-f",
		filepath.Join(path, name),
		"-N",
		"",
		"-C",
		"",
	}
	log.Infof("Executing commands %v", args)
	spawn := exec.Command("ssh-keygen", args...)
	spawn.Stdout = cfg.Stdout
	spawn.Stdin = cfg.Stdin
	spawn.Stderr = cfg.Stderr
	return args[5], spawn.Run()
}

// WaitForTCPPortOpen calls IsTCPPortOpen in a loop
func WaitForTCPPortOpen(dest string) error {
	for {
		if IsTCPPortOpen(dest) {
			break
		}
		time.Sleep(1 * time.Second)
	}
	return nil
}

// IsTCPPortOpen returns true if a TCP communication with "host:port" can be initialized
func IsTCPPortOpen(dest string) bool {
	conn, err := net.DialTimeout("tcp", dest, time.Duration(2000)*time.Millisecond)
	if err == nil {
		defer conn.Close()
	}
	return err == nil
}

// TruncIf ensures the input string does not exceed max size if cond is met
func TruncIf(str string, max int, cond bool) string {
	if cond && len(str) > max {
		return str[:max]
	}
	return str
}

// Wordify convert complex name to a single word without special shell characters
func Wordify(str string) string {
	str = regexp.MustCompile(`[^a-zA-Z0-9-]`).ReplaceAllString(str, "_")
	str = regexp.MustCompile(`__+`).ReplaceAllString(str, "_")
	str = strings.Trim(str, "_")
	return str
}

// PathToTARPathparts returns the two parts of a unix path
func PathToTARPathparts(fullPath string) (string, string) {
	fullPath = strings.TrimRight(fullPath, "/")
	return path.Dir(fullPath), path.Base(fullPath)
}

// RemoveDuplicates transforms an array into a unique array
func RemoveDuplicates(elements []string) []string {
	encountered := map[string]bool{}

	// Create a map of all unique elements.
	for v := range elements {
		encountered[elements[v]] = true
	}

	// Place all keys from the map into a slice.
	result := []string{}
	for key := range encountered {
		result = append(result, key)
	}
	return result
}

// AttachToSerial tries to connect to server serial using 'gotty-client' and fallback with a help message
func AttachToSerial(serverID string, apiToken string) (*gottyclient.Client, chan bool, error) {
	gottyURL := os.Getenv("SCW_GOTTY_URL")
	if gottyURL == "" {
		gottyURL = "https://tty.scaleway.com/v2/"
	}
	URL := fmt.Sprintf("%s?arg=%s&arg=%s", gottyURL, apiToken, serverID)

	logrus.Debug("Connection to ", URL)
	gottycli, err := gottyclient.NewClient(URL)
	if err != nil {
		return nil, nil, err
	}

	if os.Getenv("SCW_TLSVERIFY") == "0" {
		gottycli.SkipTLSVerify = true
	}

	if err = gottycli.Connect(); err != nil {
		return nil, nil, err
	}
	done := make(chan bool)

	fmt.Println("You are connected, type 'Ctrl+q' to quit.")
	go func() {
		gottycli.Loop()
		gottycli.Close()
		done <- true
	}()
	return gottycli, done, nil
}

func rfc4716hex(data []byte) string {
	fingerprint := ""

	for i := 0; i < len(data); i++ {
		fingerprint = fmt.Sprintf("%s%0.2x", fingerprint, data[i])
		if i != len(data)-1 {
			fingerprint = fingerprint + ":"
		}
	}
	return fingerprint
}

// SSHGetFingerprint returns the fingerprint of an SSH key
func SSHGetFingerprint(key []byte) (string, error) {
	publicKey, comment, _, _, err := ssh.ParseAuthorizedKey(key)
	if err != nil {
		return "", err
	}
	switch reflect.TypeOf(publicKey).String() {
	case "*ssh.rsaPublicKey", "*ssh.dsaPublicKey", "*ssh.ecdsaPublicKey":
		md5sum := md5.Sum(publicKey.Marshal())
		return publicKey.Type() + " " + rfc4716hex(md5sum[:]) + " " + comment, nil
	default:
		return "", errors.New("Can't handle this key")
	}
}
<commit_msg>exec: handle piped command<commit_after>// Copyright (C) 2015 Scaleway. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE.md file.

// scw helpers

// Package utils contains helpers
package utils

import (
	"crypto/md5"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"reflect"
	"regexp"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/Sirupsen/logrus"
	log "github.com/Sirupsen/logrus"
	"github.com/mattn/go-isatty"
	"github.com/moul/gotty-client"
	"github.com/scaleway/scaleway-cli/pkg/sshcommand"
)

// SpawnRedirection is used to redirects the fluxes
type SpawnRedirection struct {
	Stdin  io.Reader
	Stdout io.Writer
	Stderr io.Writer
}

// SSHExec executes a command over SSH and redirects file-descriptors
func SSHExec(publicIPAddress string, privateIPAddress string, command []string, checkConnection bool, gateway string) error {
	gatewayUser := "root"
	gatewayIPAddress := gateway
	if strings.Contains(gateway, "@") {
		parts := strings.Split(gatewayIPAddress, "@")
		gatewayUser = parts[0]
		gatewayIPAddress = parts[1]
		gateway = gatewayUser + "@" + gatewayIPAddress
	}

	if publicIPAddress == "" && gatewayIPAddress == "" {
		return errors.New("server does not have public IP")
	}
	if privateIPAddress == "" && gatewayIPAddress != "" {
		return errors.New("server does not have private IP")
	}

	if checkConnection {
		useGateway := gatewayIPAddress != ""
		if useGateway && !IsTCPPortOpen(fmt.Sprintf("%s:22", gatewayIPAddress)) {
			return errors.New("gateway is not available, try again later")
		}
		if !useGateway && !IsTCPPortOpen(fmt.Sprintf("%s:22", publicIPAddress)) {
			return errors.New("server is not ready, try again later")
		}
	}

	sshCommand := NewSSHExecCmd(publicIPAddress, privateIPAddress, isatty.IsTerminal(os.Stdin.Fd()), command, gateway)

	log.Debugf("Executing: %s", sshCommand)

	spawn := exec.Command("ssh", sshCommand.Slice()[1:]...)
	spawn.Stdout = os.Stdout
	spawn.Stdin = os.Stdin
	spawn.Stderr = os.Stderr
	return spawn.Run()
}

// NewSSHExecCmd computes execve compatible arguments to run a command via ssh
func NewSSHExecCmd(publicIPAddress string, privateIPAddress string, allocateTTY bool, command []string, gatewayIPAddress string) *sshcommand.Command {
	quiet := os.Getenv("DEBUG") != "1"
	secureExec := os.Getenv("SCW_SECURE_EXEC") == "1"
	sshCommand := &sshcommand.Command{
		AllocateTTY:         allocateTTY,
		Command:             command,
		Host:                publicIPAddress,
		Quiet:               quiet,
		SkipHostKeyChecking: !secureExec,
		User:                "root",
		NoEscapeCommand:     true,
	}
	if gatewayIPAddress != "" {
		sshCommand.Host = privateIPAddress
		sshCommand.Gateway = &sshcommand.Command{
			Host:                gatewayIPAddress,
			SkipHostKeyChecking: !secureExec,
			AllocateTTY:         allocateTTY,
			Quiet:               quiet,
			User:                "root",
		}
	}

	return sshCommand
}

// GeneratingAnSSHKey generates an SSH key
func GeneratingAnSSHKey(cfg SpawnRedirection, path string, name string) (string, error) {
	args := []string{
		"-t",
		"rsa",
		"-b",
		"4096",
		"-f",
		filepath.Join(path, name),
		"-N",
		"",
		"-C",
		"",
	}
	log.Infof("Executing commands %v", args)
	spawn := exec.Command("ssh-keygen", args...)
	spawn.Stdout = cfg.Stdout
	spawn.Stdin = cfg.Stdin
	spawn.Stderr = cfg.Stderr
	return args[5], spawn.Run()
}

// WaitForTCPPortOpen calls IsTCPPortOpen in a loop
func WaitForTCPPortOpen(dest string) error {
	for {
		if IsTCPPortOpen(dest) {
			break
		}
		time.Sleep(1 * time.Second)
	}
	return nil
}

// IsTCPPortOpen returns true if a TCP communication with "host:port" can be initialized
func IsTCPPortOpen(dest string) bool {
	conn, err := net.DialTimeout("tcp", dest, time.Duration(2000)*time.Millisecond)
	if err == nil {
		defer conn.Close()
	}
	return err == nil
}

// TruncIf ensures the input string does not exceed max size if cond is met
func TruncIf(str string, max int, cond bool) string {
	if cond && len(str) > max {
		return str[:max]
	}
	return str
}

// Wordify convert complex name to a single word without special shell characters
func Wordify(str string) string {
	str = regexp.MustCompile(`[^a-zA-Z0-9-]`).ReplaceAllString(str, "_")
	str = regexp.MustCompile(`__+`).ReplaceAllString(str, "_")
	str = strings.Trim(str, "_")
	return str
}

// PathToTARPathparts returns the two parts of a unix path
func PathToTARPathparts(fullPath string) (string, string) {
	fullPath = strings.TrimRight(fullPath, "/")
	return path.Dir(fullPath), path.Base(fullPath)
}

// RemoveDuplicates transforms an array into a unique array
func RemoveDuplicates(elements []string) []string {
	encountered := map[string]bool{}

	// Create a map of all unique elements.
	for v := range elements {
		encountered[elements[v]] = true
	}

	// Place all keys from the map into a slice.
	result := []string{}
	for key := range encountered {
		result = append(result, key)
	}
	return result
}

// AttachToSerial tries to connect to server serial using 'gotty-client' and fallback with a help message
func AttachToSerial(serverID string, apiToken string) (*gottyclient.Client, chan bool, error) {
	gottyURL := os.Getenv("SCW_GOTTY_URL")
	if gottyURL == "" {
		gottyURL = "https://tty.scaleway.com/v2/"
	}
	URL := fmt.Sprintf("%s?arg=%s&arg=%s", gottyURL, apiToken, serverID)

	logrus.Debug("Connection to ", URL)
	gottycli, err := gottyclient.NewClient(URL)
	if err != nil {
		return nil, nil, err
	}

	if os.Getenv("SCW_TLSVERIFY") == "0" {
		gottycli.SkipTLSVerify = true
	}

	if err = gottycli.Connect(); err != nil {
		return nil, nil, err
	}
	done := make(chan bool)

	fmt.Println("You are connected, type 'Ctrl+q' to quit.")
	go func() {
		gottycli.Loop()
		gottycli.Close()
		done <- true
	}()
	return gottycli, done, nil
}

func rfc4716hex(data []byte) string {
	fingerprint := ""

	for i := 0; i < len(data); i++ {
		fingerprint = fmt.Sprintf("%s%0.2x", fingerprint, data[i])
		if i != len(data)-1 {
			fingerprint = fingerprint + ":"
		}
	}
	return fingerprint
}

// SSHGetFingerprint returns the fingerprint of an SSH key
func SSHGetFingerprint(key []byte) (string, error) {
	publicKey, comment, _, _, err := ssh.ParseAuthorizedKey(key)
	if err != nil {
		return "", err
	}
	switch reflect.TypeOf(publicKey).String() {
	case "*ssh.rsaPublicKey", "*ssh.dsaPublicKey", "*ssh.ecdsaPublicKey":
		md5sum := md5.Sum(publicKey.Marshal())
		return publicKey.Type() + " " + rfc4716hex(md5sum[:]) + " " + comment, nil
	default:
		return "", errors.New("Can't handle this key")
	}
}
<|endoftext|>


<commit_before>package web

import "github.com/hellofresh/janus/pkg/config"

// Option represents the available options
type Option func(*Server)

// WithPort sets the server port
func WithPort(port int) Option {
	return func(s *Server) {
		s.Port = port
	}
}

// WithCredentials sets the credentials for the server
func WithCredentials(cred config.Credentials) Option {
	return func(s *Server) {
		s.Credentials = cred
	}
}

// ReadOnly sets the credentials for the server
func ReadOnly(readOnly bool) Option {
	return func(s *Server) {
		s.ReadOnly = readOnly
	}
}

// WithTLS sets the TLS configs for the server
func WithTLS(tls config.TLS) Option {
	return func(s *Server) {
		s.TLS = tls
	}
}
<commit_msg>Added option for the configurations<commit_after>package web

import (
	"github.com/hellofresh/janus/pkg/api"
	"github.com/hellofresh/janus/pkg/config"
)

// Option represents the available options
type Option func(*Server)

// WithConfigurations sets the current configurations in memory
func WithConfigurations(cfgs *api.Configuration) Option {
	return func(s *Server) {
		s.apiHandler.Cfgs = cfgs
	}
}

// WithPort sets the server port
func WithPort(port int) Option {
	return func(s *Server) {
		s.Port = port
	}
}

// WithCredentials sets the credentials for the server
func WithCredentials(cred config.Credentials) Option {
	return func(s *Server) {
		s.Credentials = cred
	}
}

// ReadOnly sets the credentials for the server
func ReadOnly(readOnly bool) Option {
	return func(s *Server) {
		s.ReadOnly = readOnly
	}
}

// WithTLS sets the TLS configs for the server
func WithTLS(tls config.TLS) Option {
	return func(s *Server) {
		s.TLS = tls
	}
}
<|endoftext|>


<commit_before>package stream

import (
	"log"
	"net"

	"fmt"

	"strings"

	"github.com/getgauge/common"
	"github.com/getgauge/gauge/conn"
	"github.com/getgauge/gauge/env"
	"github.com/getgauge/gauge/execution"
	"github.com/getgauge/gauge/execution/event"
	"github.com/getgauge/gauge/execution/rerun"
	"github.com/getgauge/gauge/execution/result"
	"github.com/getgauge/gauge/filter"
	"github.com/getgauge/gauge/gauge"
	gm "github.com/getgauge/gauge/gauge_messages"
	"github.com/getgauge/gauge/logger"
	"github.com/getgauge/gauge/reporter"
	"github.com/getgauge/gauge/util"
	"github.com/getgauge/gauge/validation"
	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc"
)

func Start() {
	port, err := conn.GetPortFromEnvironmentVariable(common.APIV2PortEnvVariableName)
	if err != nil {
		logger.APILog.Error("Failed to start execution API Service. %s \n", err.Error())
		return
	}
	listener, err := net.ListenTCP("tcp", &net.TCPAddr{Port: port})
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	s := grpc.NewServer()
	gm.RegisterExecutionServer(s, &executionServer{})
	go s.Serve(listener)
}

type executionServer struct {
}

func (e *executionServer) Execute(req *gm.ExecutionRequest, stream gm.Execution_ExecuteServer) error {
	errs := setFlags(req)
	if len(errs) > 0 {
		stream.Send(getErrorExecutionResponse(errs...))
		return nil
	}
	execute(req.Specs, stream)
	return nil
}

func execute(specDirs []string, stream gm.Execution_ExecuteServer) {
	res := validation.ValidateSpecs(specDirs)
	if len(res.Errs) > 0 {
		stream.Send(getErrorExecutionResponse(res.Errs...))
		return
	}
	event.InitRegistry()
	listenExecutionEvents(stream)
	rerun.ListenFailedScenarios()
	execution.Execute(res.SpecCollection, res.Runner, nil, res.ErrMap, execution.InParallel, 0)
}

func listenExecutionEvents(stream gm.Execution_ExecuteServer) {
	ch := make(chan event.ExecutionEvent, 0)
	event.Register(ch, event.SuiteStart, event.SpecStart, event.SpecEnd, event.ScenarioStart, event.ScenarioEnd, event.SuiteEnd)
	go func() {
		for {
			e := <-ch
			res := getResponse(e)
			if stream.Send(res) != nil || res.Type == gm.ExecutionResponse_SuiteEnd.Enum() {
				return
			}
		}
	}()
}

func getResponse(e event.ExecutionEvent) *gm.ExecutionResponse {
	switch e.Topic {
	case event.SuiteStart:
		return &gm.ExecutionResponse{Type: gm.ExecutionResponse_SuiteStart.Enum()}
	case event.SpecStart:
		return &gm.ExecutionResponse{
			Type: gm.ExecutionResponse_SpecStart.Enum(),
			ID:   e.ExecutionInfo.CurrentSpec.FileName,
		}
	case event.ScenarioStart:
		return &gm.ExecutionResponse{
			Type: gm.ExecutionResponse_ScenarioStart.Enum(),
			ID:   proto.String(fmt.Sprintf("%s:%d", e.ExecutionInfo.CurrentSpec.GetFileName(), e.Item.(*gauge.Scenario).Heading.LineNo)),
			Result: &gm.Result{
				TableRowNumber: proto.Int64(int64(getDataTableRowNumber(e.Item.(*gauge.Scenario)))),
			},
		}
	case event.ScenarioEnd:
		scn := e.Item.(*gauge.Scenario)
		return &gm.ExecutionResponse{
			Type: gm.ExecutionResponse_ScenarioEnd.Enum(),
			ID:   proto.String(fmt.Sprintf("%s:%d", e.ExecutionInfo.CurrentSpec.GetFileName(), scn.Heading.LineNo)),
			Result: &gm.Result{
				Status:            getStatus(e.Result.(*result.ScenarioResult)),
				ExecutionTime:     proto.Int64(e.Result.ExecTime()),
				Errors:            getErrors(e.Result.(*result.ScenarioResult).ProtoScenario.GetScenarioItems()),
				BeforeHookFailure: getHookFailure(e.Result.GetPreHook()),
				AfterHookFailure:  getHookFailure(e.Result.GetPostHook()),
				TableRowNumber:    proto.Int64(int64(getDataTableRowNumber(scn))),
			},
		}
	case event.SpecEnd:
		return &gm.ExecutionResponse{
			Type: gm.ExecutionResponse_SpecEnd.Enum(),
			ID:   e.ExecutionInfo.CurrentSpec.FileName,
			Result: &gm.Result{
				BeforeHookFailure: getHookFailure(e.Result.GetPreHook()),
				AfterHookFailure:  getHookFailure(e.Result.GetPostHook()),
			},
		}
	case event.SuiteEnd:
		return &gm.ExecutionResponse{
			Type: gm.ExecutionResponse_SuiteEnd.Enum(),
			Result: &gm.Result{
				BeforeHookFailure: getHookFailure(e.Result.GetPreHook()),
				AfterHookFailure:  getHookFailure(e.Result.GetPostHook()),
			},
		}
	}
	return nil
}

func getDataTableRowNumber(scn *gauge.Scenario) int {
	index := scn.DataTableRowIndex
	if scn.DataTableRow.IsInitialized() {
		index++
	}
	return index
}

func getErrorExecutionResponse(errs ...error) *gm.ExecutionResponse {
	var e []*gm.Result_ExecutionError
	for _, err := range errs {
		e = append(e, &gm.Result_ExecutionError{ErrorMessage: proto.String(err.Error())})
	}
	return &gm.ExecutionResponse{
		Type: gm.ExecutionResponse_ErrorResult.Enum(),
		Result: &gm.Result{
			Errors: e,
		},
	}
}

func getHookFailure(hookFailure **gm.ProtoHookFailure) *gm.Result_ExecutionError {
	if hookFailure != nil && *hookFailure != nil {
		return &gm.Result_ExecutionError{
			Screenshot:   (**hookFailure).ScreenShot,
			ErrorMessage: (**hookFailure).ErrorMessage,
			StackTrace:   (**hookFailure).StackTrace,
		}
	}
	return nil
}

func getErrors(items []*gm.ProtoItem) []*gm.Result_ExecutionError {
	var errors []*gm.Result_ExecutionError
	for _, item := range items {
		executionResult := item.GetStep().GetStepExecutionResult()
		res := executionResult.GetExecutionResult()
		switch item.GetItemType() {
		case gm.ProtoItem_Step:
			if executionResult.GetSkipped() {
				errors = append(errors, &gm.Result_ExecutionError{ErrorMessage: executionResult.SkippedReason})
			} else if res.GetFailed() {
				errors = append(errors, &gm.Result_ExecutionError{
					StackTrace:   res.StackTrace,
					ErrorMessage: res.ErrorMessage,
					Screenshot:   res.ScreenShot,
				})
			}
		case gm.ProtoItem_Concept:
			errors = append(errors, getErrors(item.GetConcept().GetSteps())...)
		}
	}
	return errors
}

func getStatus(result *result.ScenarioResult) *gm.Result_Status {
	if result.GetFailed() {
		return gm.Result_FAILED.Enum()
	}
	if result.ProtoScenario.GetSkipped() {
		return gm.Result_SKIPPED.Enum()
	}
	return gm.Result_PASSED.Enum()
}

func setFlags(req *gm.ExecutionRequest) []error {
	resetFlags()
	reporter.IsParallel = req.GetIsParallel()
	execution.InParallel = req.GetIsParallel()
	filter.ExecuteTags = req.GetTags()
	execution.TableRows = req.GetTableRows()
	streams := int(req.GetParallelStreams())
	if streams < 1 {
		streams = util.NumberOfCores()
	}
	execution.NumberOfExecutionStreams = streams
	reporter.NumberOfExecutionStreams = streams
	filter.NumberOfExecutionStreams = streams
	execution.Strategy = strings.ToLower(req.GetStrategy().String())
	filter.DoNotRandomize = req.GetSort()
	reporter.Verbose = true
	logger.Initialize(strings.ToLower(req.GetLogLevel().String()))
	util.SetWorkingDir(req.GetWorkingDir())
	var errs []error
	if req.GetEnv() != "" {
		if err := env.LoadEnv(req.GetEnv()); err != nil {
			errs = append(errs, err)
		}
	}
	return errs
}

func resetFlags() {
	cores := util.NumberOfCores()
	reporter.IsParallel = false
	execution.InParallel = false
	reporter.Verbose = false
	filter.ExecuteTags = ""
	execution.TableRows = ""
	execution.NumberOfExecutionStreams = cores
	reporter.NumberOfExecutionStreams = cores
	filter.NumberOfExecutionStreams = cores
	execution.Strategy = "lazy"
	filter.DoNotRandomize = false
}
<commit_msg>Setting working dir to user provided value in exection request #165<commit_after>package stream

import (
	"log"
	"net"

	"fmt"

	"strings"

	"github.com/getgauge/common"
	"github.com/getgauge/gauge/conn"
	"github.com/getgauge/gauge/env"
	"github.com/getgauge/gauge/execution"
	"github.com/getgauge/gauge/execution/event"
	"github.com/getgauge/gauge/execution/rerun"
	"github.com/getgauge/gauge/execution/result"
	"github.com/getgauge/gauge/filter"
	"github.com/getgauge/gauge/gauge"
	gm "github.com/getgauge/gauge/gauge_messages"
	"github.com/getgauge/gauge/logger"
	"github.com/getgauge/gauge/reporter"
	"github.com/getgauge/gauge/util"
	"github.com/getgauge/gauge/validation"
	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc"
)

func Start() {
	port, err := conn.GetPortFromEnvironmentVariable(common.APIV2PortEnvVariableName)
	if err != nil {
		logger.APILog.Error("Failed to start execution API Service. %s \n", err.Error())
		return
	}
	listener, err := net.ListenTCP("tcp", &net.TCPAddr{Port: port})
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	s := grpc.NewServer()
	gm.RegisterExecutionServer(s, &executionServer{})
	go s.Serve(listener)
}

type executionServer struct {
}

func (e *executionServer) Execute(req *gm.ExecutionRequest, stream gm.Execution_ExecuteServer) error {
	errs := setFlags(req)
	if len(errs) > 0 {
		stream.Send(getErrorExecutionResponse(errs...))
		return nil
	}
	execute(req.Specs, stream)
	return nil
}

func execute(specDirs []string, stream gm.Execution_ExecuteServer) {
	res := validation.ValidateSpecs(specDirs)
	if len(res.Errs) > 0 {
		stream.Send(getErrorExecutionResponse(res.Errs...))
		return
	}
	event.InitRegistry()
	listenExecutionEvents(stream)
	rerun.ListenFailedScenarios()
	execution.Execute(res.SpecCollection, res.Runner, nil, res.ErrMap, execution.InParallel, 0)
}

func listenExecutionEvents(stream gm.Execution_ExecuteServer) {
	ch := make(chan event.ExecutionEvent, 0)
	event.Register(ch, event.SuiteStart, event.SpecStart, event.SpecEnd, event.ScenarioStart, event.ScenarioEnd, event.SuiteEnd)
	go func() {
		for {
			e := <-ch
			res := getResponse(e)
			if stream.Send(res) != nil || res.Type == gm.ExecutionResponse_SuiteEnd.Enum() {
				return
			}
		}
	}()
}

func getResponse(e event.ExecutionEvent) *gm.ExecutionResponse {
	switch e.Topic {
	case event.SuiteStart:
		return &gm.ExecutionResponse{Type: gm.ExecutionResponse_SuiteStart.Enum()}
	case event.SpecStart:
		return &gm.ExecutionResponse{
			Type: gm.ExecutionResponse_SpecStart.Enum(),
			ID:   e.ExecutionInfo.CurrentSpec.FileName,
		}
	case event.ScenarioStart:
		return &gm.ExecutionResponse{
			Type: gm.ExecutionResponse_ScenarioStart.Enum(),
			ID:   proto.String(fmt.Sprintf("%s:%d", e.ExecutionInfo.CurrentSpec.GetFileName(), e.Item.(*gauge.Scenario).Heading.LineNo)),
			Result: &gm.Result{
				TableRowNumber: proto.Int64(int64(getDataTableRowNumber(e.Item.(*gauge.Scenario)))),
			},
		}
	case event.ScenarioEnd:
		scn := e.Item.(*gauge.Scenario)
		return &gm.ExecutionResponse{
			Type: gm.ExecutionResponse_ScenarioEnd.Enum(),
			ID:   proto.String(fmt.Sprintf("%s:%d", e.ExecutionInfo.CurrentSpec.GetFileName(), scn.Heading.LineNo)),
			Result: &gm.Result{
				Status:            getStatus(e.Result.(*result.ScenarioResult)),
				ExecutionTime:     proto.Int64(e.Result.ExecTime()),
				Errors:            getErrors(e.Result.(*result.ScenarioResult).ProtoScenario.GetScenarioItems()),
				BeforeHookFailure: getHookFailure(e.Result.GetPreHook()),
				AfterHookFailure:  getHookFailure(e.Result.GetPostHook()),
				TableRowNumber:    proto.Int64(int64(getDataTableRowNumber(scn))),
			},
		}
	case event.SpecEnd:
		return &gm.ExecutionResponse{
			Type: gm.ExecutionResponse_SpecEnd.Enum(),
			ID:   e.ExecutionInfo.CurrentSpec.FileName,
			Result: &gm.Result{
				BeforeHookFailure: getHookFailure(e.Result.GetPreHook()),
				AfterHookFailure:  getHookFailure(e.Result.GetPostHook()),
			},
		}
	case event.SuiteEnd:
		return &gm.ExecutionResponse{
			Type: gm.ExecutionResponse_SuiteEnd.Enum(),
			Result: &gm.Result{
				BeforeHookFailure: getHookFailure(e.Result.GetPreHook()),
				AfterHookFailure:  getHookFailure(e.Result.GetPostHook()),
			},
		}
	}
	return nil
}

func getDataTableRowNumber(scn *gauge.Scenario) int {
	index := scn.DataTableRowIndex
	if scn.DataTableRow.IsInitialized() {
		index++
	}
	return index
}

func getErrorExecutionResponse(errs ...error) *gm.ExecutionResponse {
	var e []*gm.Result_ExecutionError
	for _, err := range errs {
		e = append(e, &gm.Result_ExecutionError{ErrorMessage: proto.String(err.Error())})
	}
	return &gm.ExecutionResponse{
		Type: gm.ExecutionResponse_ErrorResult.Enum(),
		Result: &gm.Result{
			Errors: e,
		},
	}
}

func getHookFailure(hookFailure **gm.ProtoHookFailure) *gm.Result_ExecutionError {
	if hookFailure != nil && *hookFailure != nil {
		return &gm.Result_ExecutionError{
			Screenshot:   (**hookFailure).ScreenShot,
			ErrorMessage: (**hookFailure).ErrorMessage,
			StackTrace:   (**hookFailure).StackTrace,
		}
	}
	return nil
}

func getErrors(items []*gm.ProtoItem) []*gm.Result_ExecutionError {
	var errors []*gm.Result_ExecutionError
	for _, item := range items {
		executionResult := item.GetStep().GetStepExecutionResult()
		res := executionResult.GetExecutionResult()
		switch item.GetItemType() {
		case gm.ProtoItem_Step:
			if executionResult.GetSkipped() {
				errors = append(errors, &gm.Result_ExecutionError{ErrorMessage: executionResult.SkippedReason})
			} else if res.GetFailed() {
				errors = append(errors, &gm.Result_ExecutionError{
					StackTrace:   res.StackTrace,
					ErrorMessage: res.ErrorMessage,
					Screenshot:   res.ScreenShot,
				})
			}
		case gm.ProtoItem_Concept:
			errors = append(errors, getErrors(item.GetConcept().GetSteps())...)
		}
	}
	return errors
}

func getStatus(result *result.ScenarioResult) *gm.Result_Status {
	if result.GetFailed() {
		return gm.Result_FAILED.Enum()
	}
	if result.ProtoScenario.GetSkipped() {
		return gm.Result_SKIPPED.Enum()
	}
	return gm.Result_PASSED.Enum()
}

func setFlags(req *gm.ExecutionRequest) []error {
	resetFlags()
	reporter.IsParallel = req.GetIsParallel()
	execution.InParallel = req.GetIsParallel()
	filter.ExecuteTags = req.GetTags()
	execution.TableRows = req.GetTableRows()
	streams := int(req.GetParallelStreams())
	if streams < 1 {
		streams = util.NumberOfCores()
	}
	execution.NumberOfExecutionStreams = streams
	reporter.NumberOfExecutionStreams = streams
	filter.NumberOfExecutionStreams = streams
	execution.Strategy = strings.ToLower(req.GetStrategy().String())
	filter.DoNotRandomize = req.GetSort()
	reporter.Verbose = true
	logger.Initialize(strings.ToLower(req.GetLogLevel().String()))
	if req.GetWorkingDir() != "" {
		util.SetWorkingDir(req.GetWorkingDir())
	}
	var errs []error
	if req.GetEnv() != "" {
		if err := env.LoadEnv(req.GetEnv()); err != nil {
			errs = append(errs, err)
		}
	}
	return errs
}

func resetFlags() {
	cores := util.NumberOfCores()
	reporter.IsParallel = false
	execution.InParallel = false
	reporter.Verbose = false
	filter.ExecuteTags = ""
	execution.TableRows = ""
	execution.NumberOfExecutionStreams = cores
	reporter.NumberOfExecutionStreams = cores
	filter.NumberOfExecutionStreams = cores
	execution.Strategy = "lazy"
	filter.DoNotRandomize = false
}
<|endoftext|>


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

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

package testutils

import (
	"encoding/json"
	"fmt"
	"runtime"
	"strings"
	"testing"
	"time"

	"github.com/uber/tchannel-go"
	"github.com/uber/tchannel-go/atomic"
	"github.com/uber/tchannel-go/testutils/goroutines"

	"github.com/stretchr/testify/assert"
)

// Has a previous test already leaked a goroutine?
var _leakedGoroutine = atomic.NewInt32(0)

// A TestServer encapsulates a TChannel server, a client factory, and functions
// to ensure that we're not leaking resources.
type TestServer struct {
	testing.TB

	// relayIdx is the index of the relay channel, if any, in the channels slice.
	relayIdx int

	// relayHosts is the relayer's SimpleRelayHosts (if any).
	relayHosts *SimpleRelayHosts

	// channels is the list of channels created for this TestServer. The first
	// element is always the initial server.
	channels []*tchannel.Channel

	// channelState the initial runtime state for all channels created
	// as part of the TestServer (including the server).
	channelStates map[*tchannel.Channel]*tchannel.RuntimeState

	verifyOpts *goroutines.VerifyOpts
	postFns    []func()
}

// NewTestServer constructs a TestServer.
func NewTestServer(t testing.TB, opts *ChannelOpts) *TestServer {
	ts := &TestServer{
		TB:            t,
		channelStates: make(map[*tchannel.Channel]*tchannel.RuntimeState),
	}

	ts.NewServer(opts)
	if opts.IncludeRelay {
		ts.addRelay()
	}

	return ts
}

// WithTestServer creates a new TestServer, runs the passed function, and then
// verifies that no resources were leaked.
//
// TODO: run function twice; once with a relay, once without.
func WithTestServer(t testing.TB, chanOpts *ChannelOpts, f func(*TestServer)) {
	withServer(t, chanOpts.Copy(), f)
}

// SetVerifyOpts specifies the options we'll use during teardown to verify that
// no goroutines were leaked.
func (ts *TestServer) SetVerifyOpts(opts *goroutines.VerifyOpts) {
	ts.verifyOpts = opts
}

// Server returns the underlying TChannel for the server (i.e., the channel on
// which we're registering handlers).
//
// To support test cases with relays interposed between clients and servers,
// callers should use the Client(), HostPort(), ServiceName(), and Register()
// methods instead of accessing the server channel explicitly.
func (ts *TestServer) Server() *tchannel.Channel {
	return ts.channels[0]
}

// Relay returns the relay channel, if one is present.
func (ts *TestServer) Relay() *tchannel.Channel {
	if ts.hasRelay() {
		return ts.channels[ts.relayIdx]
	}
	return nil
}

// HostPort returns the host:port for clients to connect to. Note that this may
// not be the same as the host:port of the server channel.
func (ts *TestServer) HostPort() string {
	if ts.hasRelay() {
		return ts.Relay().PeerInfo().HostPort
	}
	return ts.Server().PeerInfo().HostPort
}

// ServiceName returns the service name of the server channel.
func (ts *TestServer) ServiceName() string {
	return ts.Server().PeerInfo().ServiceName
}

// Register registers a handler on the server channel.
func (ts *TestServer) Register(h tchannel.Handler, methodName string) {
	ts.Server().Register(h, methodName)
}

// CloseAndVerify closes all channels verifying each channel as it is closed.
// It then verifies that no goroutines were leaked.
func (ts *TestServer) CloseAndVerify() {
	for i := len(ts.channels) - 1; i >= 0; i-- {
		ch := ts.channels[i]
		ch.Logger().Debugf("TEST: TestServer is closing and verifying channel")
		ts.close(ch)
		ts.verify(ch)
	}
}

// NewClient returns a client that with log verification.
// TODO: Verify message exchanges and leaks for client channels as well.
func (ts *TestServer) NewClient(opts *ChannelOpts) *tchannel.Channel {
	return ts.addChannel(newClient, opts.Copy())
}

// NewServer returns a server with log and channel state verification.
func (ts *TestServer) NewServer(opts *ChannelOpts) *tchannel.Channel {
	ch := ts.addChannel(newServer, opts.Copy())
	if ts.hasRelay() {
		ts.relayHosts.Add(ch.ServiceName(), ch.PeerInfo().HostPort)
	}
	return ch
}

// addRelay adds a relay in front of the test server, altering public methods as
// necessary to route traffic through the relay.
func (ts *TestServer) addRelay() {
	ts.relayHosts = NewSimpleRelayHosts(map[string][]string{
		ts.Server().ServiceName(): []string{ts.Server().PeerInfo().HostPort},
	})
	opts := &ChannelOpts{
		ServiceName:    "relay",
		ChannelOptions: tchannel.ChannelOptions{RelayHosts: ts.relayHosts},
	}
	ts.addChannel(NewServer, opts)
	ts.relayIdx = len(ts.channels) - 1
}

func (ts *TestServer) hasRelay() bool {
	return ts.relayIdx > 0
}

func (ts *TestServer) addChannel(createChannel func(t testing.TB, opts *ChannelOpts) *tchannel.Channel, opts *ChannelOpts) *tchannel.Channel {
	ch := createChannel(ts, opts)
	ts.postFns = append(ts.postFns, opts.postFns...)
	ts.channels = append(ts.channels, ch)
	ts.channelStates[ch] = comparableState(ch)
	return ch
}

// close closes all channels in most-recently-created order.
// it waits for the channels to close.
func (ts *TestServer) close(ch *tchannel.Channel) {
	ch.Close()
	ts.waitForChannelClose(ch)
}

func (ts *TestServer) verify(ch *tchannel.Channel) {
	// For the main server channel, we want to ensure there's no goroutine leaks
	// which will wait for all runnable goroutines. We cannot verify goroutines
	// for all channels, as it would detect goroutines in the outer channels.
	if ch == ts.channels[0] {
		ts.verifyNoGoroutinesLeaked()
	}

	ts.verifyRelaysEmpty(ch)
	ts.verifyExchangesCleared(ch)
}

func (ts *TestServer) post() {
	if !ts.Failed() {
		for _, ch := range ts.channels {
			ts.verifyNoStateLeak(ch)
		}
	}
	for _, fn := range ts.postFns {
		fn()
	}
}

func (ts *TestServer) waitForChannelClose(ch *tchannel.Channel) {
	if ts.Failed() {
		return
	}
	started := time.Now()

	var state tchannel.ChannelState
	for i := 0; i < 50; i++ {
		if state = ch.State(); state == tchannel.ChannelClosed {
			return
		}

		runtime.Gosched()
		if i < 5 {
			continue
		}

		sleepFor := time.Duration(i) * 100 * time.Microsecond
		time.Sleep(Timeout(sleepFor))
	}

	// Channel is not closing, fail the test.
	sinceStart := time.Since(started)
	ts.Errorf("Channel did not close after %v, last state: %v", sinceStart, state)
}

func (ts *TestServer) verifyNoStateLeak(ch *tchannel.Channel) {
	initial := ts.channelStates[ch]
	final := comparableState(ch)
	assert.Equal(ts.TB, initial, final, "Runtime state has leaks")
}

func (ts *TestServer) verifyExchangesCleared(ch *tchannel.Channel) {
	if ts.Failed() {
		return
	}
	// Ensure that all the message exchanges are empty.
	serverState := ch.IntrospectState(&tchannel.IntrospectionOptions{
		IncludeExchanges: true,
	})
	if exchangesLeft := describeLeakedExchanges(serverState); exchangesLeft != "" {
		ts.Errorf("Found uncleared message exchanges on server:\n%v", exchangesLeft)
	}
}

func (ts *TestServer) verifyRelaysEmpty(ch *tchannel.Channel) {
	if ts.Failed() {
		return
	}
	var foundErrors bool
	state := ch.IntrospectState(&tchannel.IntrospectionOptions{IncludeExchanges: true})
	for _, peerState := range state.RootPeers {
		var connStates []tchannel.ConnectionRuntimeState
		connStates = append(connStates, peerState.InboundConnections...)
		connStates = append(connStates, peerState.OutboundConnections...)
		for _, connState := range connStates {
			n := connState.Relayer.Count
			if assert.Equal(ts, 0, n, "Found %v left-over items in relayer for %v.", n, connState.LocalHostPort) {
				continue
			}
			foundErrors = true
		}
	}

	if !foundErrors {
		return
	}

	marshalled, err := json.MarshalIndent(state, "", "  ")
	require.NoError(ts, err, "Failed to marshal relayer state")
	// Print out all the exchanges we found.
	ts.Logf("Relayer state:\n%s", marshalled)
}

func (ts *TestServer) verifyNoGoroutinesLeaked() {
	if _leakedGoroutine.Load() == 1 {
		ts.Log("Skipping check for leaked goroutines because of a previous leak.")
		return
	}
	err := goroutines.IdentifyLeaks(ts.verifyOpts)
	if err == nil {
		// No leaks, nothing to do.
		return
	}
	if isFirstLeak := _leakedGoroutine.CAS(0, 1); !isFirstLeak {
		ts.Log("Skipping check for leaked goroutines because of a previous leak.")
		return
	}
	if ts.Failed() {
		// If we've already failed this test, don't pollute the test output with
		// more failures.
		return
	}
	ts.Error(err.Error())
}

func comparableState(ch *tchannel.Channel) *tchannel.RuntimeState {
	s := ch.IntrospectState(&tchannel.IntrospectionOptions{
		IncludeExchanges: true,
	})
	s.OtherChannels = nil
	s.SubChannels = nil
	s.Peers = nil
	return s
}

func describeLeakedExchanges(rs *tchannel.RuntimeState) string {
	var connections []*tchannel.ConnectionRuntimeState
	for _, peer := range rs.RootPeers {
		for _, conn := range peer.InboundConnections {
			connections = append(connections, &conn)
		}
		for _, conn := range peer.OutboundConnections {
			connections = append(connections, &conn)
		}
	}
	return describeLeakedExchangesConns(connections)
}

func describeLeakedExchangesConns(connections []*tchannel.ConnectionRuntimeState) string {
	var exchanges []string
	for _, c := range connections {
		if exch := describeLeakedExchangesSingleConn(c); exch != "" {
			exchanges = append(exchanges, exch)
		}
	}
	return strings.Join(exchanges, "\n")
}

func describeLeakedExchangesSingleConn(cs *tchannel.ConnectionRuntimeState) string {
	var exchanges []string
	checkExchange := func(e tchannel.ExchangeSetRuntimeState) {
		if e.Count > 0 {
			exchanges = append(exchanges, fmt.Sprintf(" %v leftover %v exchanges", e.Name, e.Count))
			for _, v := range e.Exchanges {
				exchanges = append(exchanges, fmt.Sprintf("  exchanges: %+v", v))
			}
		}
	}
	checkExchange(cs.InboundExchange)
	checkExchange(cs.OutboundExchange)
	if len(exchanges) == 0 {
		return ""
	}

	return fmt.Sprintf("Connection %d has leftover exchanges:\n\t%v", cs.ID, strings.Join(exchanges, "\n\t"))
}

func withServer(t testing.TB, chanOpts *ChannelOpts, f func(*TestServer)) {
	ts := NewTestServer(t, chanOpts)
	// Note: We use defer, as we want the postFns to run even if the test
	// goroutine exits (e.g. user calls t.Fatalf).
	defer ts.post()

	f(ts)
	ts.Server().Logger().Debugf("TEST: Test function complete")
	ts.CloseAndVerify()
}
<commit_msg>Run tests twice, once with relay and once without<commit_after>// Copyright (c) 2015 Uber Technologies, Inc.

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

package testutils

import (
	"encoding/json"
	"fmt"
	"runtime"
	"strings"
	"testing"
	"time"

	"github.com/uber/tchannel-go"
	"github.com/uber/tchannel-go/atomic"
	"github.com/uber/tchannel-go/testutils/goroutines"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// Has a previous test already leaked a goroutine?
var _leakedGoroutine = atomic.NewInt32(0)

// A TestServer encapsulates a TChannel server, a client factory, and functions
// to ensure that we're not leaking resources.
type TestServer struct {
	testing.TB

	// relayIdx is the index of the relay channel, if any, in the channels slice.
	relayIdx int

	// relayHosts is the relayer's SimpleRelayHosts (if any).
	relayHosts *SimpleRelayHosts

	// channels is the list of channels created for this TestServer. The first
	// element is always the initial server.
	channels []*tchannel.Channel

	// channelState the initial runtime state for all channels created
	// as part of the TestServer (including the server).
	channelStates map[*tchannel.Channel]*tchannel.RuntimeState

	verifyOpts *goroutines.VerifyOpts
	postFns    []func()
}

// NewTestServer constructs a TestServer.
func NewTestServer(t testing.TB, opts *ChannelOpts) *TestServer {
	ts := &TestServer{
		TB:            t,
		channelStates: make(map[*tchannel.Channel]*tchannel.RuntimeState),
	}

	ts.NewServer(opts)
	if opts.IncludeRelay {
		ts.addRelay()
	}

	return ts
}

// WithTestServer creates a new TestServer, runs the passed function, and then
// verifies that no resources were leaked.
//
// TODO: run function twice; once with a relay, once without.
func WithTestServer(t testing.TB, chanOpts *ChannelOpts, f func(*TestServer)) {
	withServer(t, chanOpts.Copy(), f)
	if chanOpts != nil && chanOpts.IncludeRelay {
		repeatOpts := chanOpts.Copy()
		repeatOpts.IncludeRelay = false
		withServer(t, repeatOpts, f)
	}
}

// SetVerifyOpts specifies the options we'll use during teardown to verify that
// no goroutines were leaked.
func (ts *TestServer) SetVerifyOpts(opts *goroutines.VerifyOpts) {
	ts.verifyOpts = opts
}

// Server returns the underlying TChannel for the server (i.e., the channel on
// which we're registering handlers).
//
// To support test cases with relays interposed between clients and servers,
// callers should use the Client(), HostPort(), ServiceName(), and Register()
// methods instead of accessing the server channel explicitly.
func (ts *TestServer) Server() *tchannel.Channel {
	return ts.channels[0]
}

// Relay returns the relay channel, if one is present.
func (ts *TestServer) Relay() *tchannel.Channel {
	if ts.hasRelay() {
		return ts.channels[ts.relayIdx]
	}
	return nil
}

// HostPort returns the host:port for clients to connect to. Note that this may
// not be the same as the host:port of the server channel.
func (ts *TestServer) HostPort() string {
	if ts.hasRelay() {
		return ts.Relay().PeerInfo().HostPort
	}
	return ts.Server().PeerInfo().HostPort
}

// ServiceName returns the service name of the server channel.
func (ts *TestServer) ServiceName() string {
	return ts.Server().PeerInfo().ServiceName
}

// Register registers a handler on the server channel.
func (ts *TestServer) Register(h tchannel.Handler, methodName string) {
	ts.Server().Register(h, methodName)
}

// CloseAndVerify closes all channels verifying each channel as it is closed.
// It then verifies that no goroutines were leaked.
func (ts *TestServer) CloseAndVerify() {
	for i := len(ts.channels) - 1; i >= 0; i-- {
		ch := ts.channels[i]
		ch.Logger().Debugf("TEST: TestServer is closing and verifying channel")
		ts.close(ch)
		ts.verify(ch)
	}
}

// NewClient returns a client that with log verification.
// TODO: Verify message exchanges and leaks for client channels as well.
func (ts *TestServer) NewClient(opts *ChannelOpts) *tchannel.Channel {
	return ts.addChannel(newClient, opts.Copy())
}

// NewServer returns a server with log and channel state verification.
func (ts *TestServer) NewServer(opts *ChannelOpts) *tchannel.Channel {
	ch := ts.addChannel(newServer, opts.Copy())
	if ts.hasRelay() {
		ts.relayHosts.Add(ch.ServiceName(), ch.PeerInfo().HostPort)
	}
	return ch
}

// addRelay adds a relay in front of the test server, altering public methods as
// necessary to route traffic through the relay.
func (ts *TestServer) addRelay() {
	ts.relayHosts = NewSimpleRelayHosts(map[string][]string{
		ts.Server().ServiceName(): []string{ts.Server().PeerInfo().HostPort},
	})
	opts := &ChannelOpts{
		ServiceName:    "relay",
		ChannelOptions: tchannel.ChannelOptions{RelayHosts: ts.relayHosts},
	}
	ts.addChannel(NewServer, opts)
	ts.relayIdx = len(ts.channels) - 1
}

func (ts *TestServer) hasRelay() bool {
	return ts.relayIdx > 0
}

func (ts *TestServer) addChannel(createChannel func(t testing.TB, opts *ChannelOpts) *tchannel.Channel, opts *ChannelOpts) *tchannel.Channel {
	ch := createChannel(ts, opts)
	ts.postFns = append(ts.postFns, opts.postFns...)
	ts.channels = append(ts.channels, ch)
	ts.channelStates[ch] = comparableState(ch)
	return ch
}

// close closes all channels in most-recently-created order.
// it waits for the channels to close.
func (ts *TestServer) close(ch *tchannel.Channel) {
	ch.Close()
	ts.waitForChannelClose(ch)
}

func (ts *TestServer) verify(ch *tchannel.Channel) {
	// For the main server channel, we want to ensure there's no goroutine leaks
	// which will wait for all runnable goroutines. We cannot verify goroutines
	// for all channels, as it would detect goroutines in the outer channels.
	if ch == ts.channels[0] {
		ts.verifyNoGoroutinesLeaked()
	}

	ts.verifyRelaysEmpty(ch)
	ts.verifyExchangesCleared(ch)
}

func (ts *TestServer) post() {
	if !ts.Failed() {
		for _, ch := range ts.channels {
			ts.verifyNoStateLeak(ch)
		}
	}
	for _, fn := range ts.postFns {
		fn()
	}
}

func (ts *TestServer) waitForChannelClose(ch *tchannel.Channel) {
	if ts.Failed() {
		return
	}
	started := time.Now()

	var state tchannel.ChannelState
	for i := 0; i < 50; i++ {
		if state = ch.State(); state == tchannel.ChannelClosed {
			return
		}

		runtime.Gosched()
		if i < 5 {
			continue
		}

		sleepFor := time.Duration(i) * 100 * time.Microsecond
		time.Sleep(Timeout(sleepFor))
	}

	// Channel is not closing, fail the test.
	sinceStart := time.Since(started)
	ts.Errorf("Channel did not close after %v, last state: %v", sinceStart, state)
}

func (ts *TestServer) verifyNoStateLeak(ch *tchannel.Channel) {
	initial := ts.channelStates[ch]
	final := comparableState(ch)
	assert.Equal(ts.TB, initial, final, "Runtime state has leaks")
}

func (ts *TestServer) verifyExchangesCleared(ch *tchannel.Channel) {
	if ts.Failed() {
		return
	}
	// Ensure that all the message exchanges are empty.
	serverState := ch.IntrospectState(&tchannel.IntrospectionOptions{
		IncludeExchanges: true,
	})
	if exchangesLeft := describeLeakedExchanges(serverState); exchangesLeft != "" {
		ts.Errorf("Found uncleared message exchanges on server:\n%v", exchangesLeft)
	}
}

func (ts *TestServer) verifyRelaysEmpty(ch *tchannel.Channel) {
	if ts.Failed() {
		return
	}
	var foundErrors bool
	state := ch.IntrospectState(&tchannel.IntrospectionOptions{IncludeExchanges: true})
	for _, peerState := range state.RootPeers {
		var connStates []tchannel.ConnectionRuntimeState
		connStates = append(connStates, peerState.InboundConnections...)
		connStates = append(connStates, peerState.OutboundConnections...)
		for _, connState := range connStates {
			n := connState.Relayer.Count
			if assert.Equal(ts, 0, n, "Found %v left-over items in relayer for %v.", n, connState.LocalHostPort) {
				continue
			}
			foundErrors = true
		}
	}

	if !foundErrors {
		return
	}

	marshalled, err := json.MarshalIndent(state, "", "  ")
	require.NoError(ts, err, "Failed to marshal relayer state")
	// Print out all the exchanges we found.
	ts.Logf("Relayer state:\n%s", marshalled)
}

func (ts *TestServer) verifyNoGoroutinesLeaked() {
	if _leakedGoroutine.Load() == 1 {
		ts.Log("Skipping check for leaked goroutines because of a previous leak.")
		return
	}
	err := goroutines.IdentifyLeaks(ts.verifyOpts)
	if err == nil {
		// No leaks, nothing to do.
		return
	}
	if isFirstLeak := _leakedGoroutine.CAS(0, 1); !isFirstLeak {
		ts.Log("Skipping check for leaked goroutines because of a previous leak.")
		return
	}
	if ts.Failed() {
		// If we've already failed this test, don't pollute the test output with
		// more failures.
		return
	}
	ts.Error(err.Error())
}

func comparableState(ch *tchannel.Channel) *tchannel.RuntimeState {
	s := ch.IntrospectState(&tchannel.IntrospectionOptions{
		IncludeExchanges: true,
	})
	s.OtherChannels = nil
	s.SubChannels = nil
	s.Peers = nil
	return s
}

func describeLeakedExchanges(rs *tchannel.RuntimeState) string {
	var connections []*tchannel.ConnectionRuntimeState
	for _, peer := range rs.RootPeers {
		for _, conn := range peer.InboundConnections {
			connections = append(connections, &conn)
		}
		for _, conn := range peer.OutboundConnections {
			connections = append(connections, &conn)
		}
	}
	return describeLeakedExchangesConns(connections)
}

func describeLeakedExchangesConns(connections []*tchannel.ConnectionRuntimeState) string {
	var exchanges []string
	for _, c := range connections {
		if exch := describeLeakedExchangesSingleConn(c); exch != "" {
			exchanges = append(exchanges, exch)
		}
	}
	return strings.Join(exchanges, "\n")
}

func describeLeakedExchangesSingleConn(cs *tchannel.ConnectionRuntimeState) string {
	var exchanges []string
	checkExchange := func(e tchannel.ExchangeSetRuntimeState) {
		if e.Count > 0 {
			exchanges = append(exchanges, fmt.Sprintf(" %v leftover %v exchanges", e.Name, e.Count))
			for _, v := range e.Exchanges {
				exchanges = append(exchanges, fmt.Sprintf("  exchanges: %+v", v))
			}
		}
	}
	checkExchange(cs.InboundExchange)
	checkExchange(cs.OutboundExchange)
	if len(exchanges) == 0 {
		return ""
	}

	return fmt.Sprintf("Connection %d has leftover exchanges:\n\t%v", cs.ID, strings.Join(exchanges, "\n\t"))
}

func withServer(t testing.TB, chanOpts *ChannelOpts, f func(*TestServer)) {
	ts := NewTestServer(t, chanOpts)
	// Note: We use defer, as we want the postFns to run even if the test
	// goroutine exits (e.g. user calls t.Fatalf).
	defer ts.post()

	f(ts)
	ts.Server().Logger().Debugf("TEST: Test function complete")
	ts.CloseAndVerify()
}
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC
//
// 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 tfgcv

import (
	"context"
	"path/filepath"

	"github.com/GoogleCloudPlatform/terraform-validator/converters/google"
	"github.com/forseti-security/config-validator/pkg/api/validator"
	"github.com/forseti-security/config-validator/pkg/gcv"
	"github.com/pkg/errors"
)

// To be set by Go build tools.
var buildVersion string

// BuildVersion returns the build version of Terraform Validator.
func BuildVersion() string {
	return buildVersion
}

// ValidateAssets instantiates GCV and audits CAI assets using "policies"
// and "lib" folder under policyPath.
func ValidateAssets(assets []google.Asset, policyPath string) (*validator.AuditResponse, error) {
	return ValidateAssetsWithLibrary(assets,
		filepath.Join(policyPath, "policies"),
		filepath.Join(policyPath, "lib"))
}

// ValidateAssetsWithLibrary instantiates GCV and audits CAI assets.
func ValidateAssetsWithLibrary(assets []google.Asset, policyPath, policyLibraryDir string) (*validator.AuditResponse, error) {
	valid, err := gcv.NewValidator(
		gcv.PolicyPath(policyPath),
		gcv.PolicyLibraryDir(policyLibraryDir),
	)
	if err != nil {
		return nil, errors.Wrap(err, "initializing gcv validator")
	}

	pbAssets := make([]*validator.Asset, len(assets))
	for i := range assets {
		pbAssets[i] = &validator.Asset{}
		if err := protoViaJSON(assets[i], pbAssets[i]); err != nil {
			return nil, errors.Wrapf(err, "converting asset %s to proto", assets[i].Name)
		}
	}

	if err := valid.AddData(&validator.AddDataRequest{
		Assets: pbAssets,
	}); err != nil {
		return nil, errors.Wrap(err, "adding data to validator")
	}

	auditResult, err := valid.Audit(context.Background())
	if err != nil {
		return nil, errors.Wrap(err, "auditing")
	}

	return auditResult, nil
}
<commit_msg>rename policyPath to policyRootPath<commit_after>// Copyright 2019 Google LLC
//
// 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 tfgcv

import (
	"context"
	"path/filepath"

	"github.com/GoogleCloudPlatform/terraform-validator/converters/google"
	"github.com/forseti-security/config-validator/pkg/api/validator"
	"github.com/forseti-security/config-validator/pkg/gcv"
	"github.com/pkg/errors"
)

// To be set by Go build tools.
var buildVersion string

// BuildVersion returns the build version of Terraform Validator.
func BuildVersion() string {
	return buildVersion
}

// ValidateAssets instantiates GCV and audits CAI assets using "policies"
// and "lib" folder under policyRootPath.
func ValidateAssets(assets []google.Asset, policyRootPath string) (*validator.AuditResponse, error) {
	return ValidateAssetsWithLibrary(assets,
		filepath.Join(policyRootPath, "policies"),
		filepath.Join(policyRootPath, "lib"))
}

// ValidateAssetsWithLibrary instantiates GCV and audits CAI assets.
func ValidateAssetsWithLibrary(assets []google.Asset, policyPath, policyLibraryDir string) (*validator.AuditResponse, error) {
	valid, err := gcv.NewValidator(
		gcv.PolicyPath(policyPath),
		gcv.PolicyLibraryDir(policyLibraryDir),
	)
	if err != nil {
		return nil, errors.Wrap(err, "initializing gcv validator")
	}

	pbAssets := make([]*validator.Asset, len(assets))
	for i := range assets {
		pbAssets[i] = &validator.Asset{}
		if err := protoViaJSON(assets[i], pbAssets[i]); err != nil {
			return nil, errors.Wrapf(err, "converting asset %s to proto", assets[i].Name)
		}
	}

	if err := valid.AddData(&validator.AddDataRequest{
		Assets: pbAssets,
	}); err != nil {
		return nil, errors.Wrap(err, "adding data to validator")
	}

	auditResult, err := valid.Audit(context.Background())
	if err != nil {
		return nil, errors.Wrap(err, "auditing")
	}

	return auditResult, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"bytes"
	"io/ioutil"
	"strings"
	"testing"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
)

func TestSaltPackSignDeviceRequired(t *testing.T) {
	tc := SetupEngineTest(t, "sign")
	defer tc.Cleanup()

	ctx := &Context{
		SecretUI: &libkb.TestSecretUI{},
	}
	eng := NewSaltPackSign(nil, tc.G)
	err := RunEngine(eng, ctx)
	if err == nil {
		t.Fatal("sign not logged in returned no error")
	}
	if _, ok := err.(libkb.DeviceRequiredError); !ok {
		t.Errorf("error type: %T, expected DeviceRequiredError", err)
	}
}

func TestSaltPackSignVerify(t *testing.T) {
	tc := SetupEngineTest(t, "sign")
	defer tc.Cleanup()

	fu := CreateAndSignupFakeUser(tc, "sign")

	// signTests are defined in pgp_sign_test.  Make sure that saltpack sign can
	// sign/verify the same messages as pgp.
	for _, test := range signTests {
		var sink bytes.Buffer

		sarg := &SaltPackSignArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: ioutil.NopCloser(bytes.NewBufferString(test.input)),
		}

		eng := NewSaltPackSign(sarg, tc.G)
		ctx := &Context{
			IdentifyUI: &FakeIdentifyUI{},
			SecretUI:   fu.NewSecretUI(),
		}

		if err := RunEngine(eng, ctx); err != nil {
			t.Errorf("%s: run error: %s", test.name, err)
			continue
		}

		sig := sink.String()

		if len(sig) == 0 {
			t.Errorf("%s: empty sig", test.name)
		}

		varg := &SaltPackVerifyArg{
			Source: strings.NewReader(sig),
		}
		veng := NewSaltPackVerify(varg, tc.G)

		if err := RunEngine(veng, ctx); err != nil {
			t.Errorf("%s: verify error: %s", test.name, err)
			continue
		}
	}

	// now try the same messages, but generate detached signatures
	for _, test := range signTests {
		var sink bytes.Buffer

		sarg := &SaltPackSignArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: ioutil.NopCloser(bytes.NewBufferString(test.input)),
			Opts: keybase1.SaltPackSignOptions{
				Detached: true,
			},
		}

		eng := NewSaltPackSign(sarg, tc.G)
		ctx := &Context{
			IdentifyUI: &FakeIdentifyUI{},
			SecretUI:   fu.NewSecretUI(),
		}

		if err := RunEngine(eng, ctx); err != nil {
			t.Errorf("(detached) %s: run error: %s", test.name, err)
			continue
		}

		sig := sink.Bytes()

		if len(sig) == 0 {
			t.Errorf("(detached) %s: empty sig", test.name)
		}

		varg := &SaltPackVerifyArg{
			Source: strings.NewReader(test.input),
			Opts: keybase1.SaltPackVerifyOptions{
				Signature: sig,
			},
		}
		veng := NewSaltPackVerify(varg, tc.G)

		if err := RunEngine(veng, ctx); err != nil {
			t.Errorf("(detached) %s: verify error: %s", test.name, err)
			continue
		}
	}
}
<commit_msg>Fix missing sp ui in verify test<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"bytes"
	"io/ioutil"
	"strings"
	"testing"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
)

func TestSaltPackSignDeviceRequired(t *testing.T) {
	tc := SetupEngineTest(t, "sign")
	defer tc.Cleanup()

	ctx := &Context{
		SecretUI: &libkb.TestSecretUI{},
	}
	eng := NewSaltPackSign(nil, tc.G)
	err := RunEngine(eng, ctx)
	if err == nil {
		t.Fatal("sign not logged in returned no error")
	}
	if _, ok := err.(libkb.DeviceRequiredError); !ok {
		t.Errorf("error type: %T, expected DeviceRequiredError", err)
	}
}

func TestSaltPackSignVerify(t *testing.T) {
	tc := SetupEngineTest(t, "sign")
	defer tc.Cleanup()

	fu := CreateAndSignupFakeUser(tc, "sign")

	// signTests are defined in pgp_sign_test.  Make sure that saltpack sign can
	// sign/verify the same messages as pgp.
	for _, test := range signTests {
		var sink bytes.Buffer

		sarg := &SaltPackSignArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: ioutil.NopCloser(bytes.NewBufferString(test.input)),
		}

		eng := NewSaltPackSign(sarg, tc.G)
		ctx := &Context{
			IdentifyUI: &FakeIdentifyUI{},
			SecretUI:   fu.NewSecretUI(),
		}

		if err := RunEngine(eng, ctx); err != nil {
			t.Errorf("%s: run error: %s", test.name, err)
			continue
		}

		sig := sink.String()

		if len(sig) == 0 {
			t.Errorf("%s: empty sig", test.name)
		}

		varg := &SaltPackVerifyArg{
			Source: strings.NewReader(sig),
		}
		veng := NewSaltPackVerify(varg, tc.G)

		ctx.SaltPackUI = fakeSaltPackUI{}

		if err := RunEngine(veng, ctx); err != nil {
			t.Errorf("%s: verify error: %s", test.name, err)
			continue
		}
	}

	// now try the same messages, but generate detached signatures
	for _, test := range signTests {
		var sink bytes.Buffer

		sarg := &SaltPackSignArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: ioutil.NopCloser(bytes.NewBufferString(test.input)),
			Opts: keybase1.SaltPackSignOptions{
				Detached: true,
			},
		}

		eng := NewSaltPackSign(sarg, tc.G)
		ctx := &Context{
			IdentifyUI: &FakeIdentifyUI{},
			SecretUI:   fu.NewSecretUI(),
		}

		if err := RunEngine(eng, ctx); err != nil {
			t.Errorf("(detached) %s: run error: %s", test.name, err)
			continue
		}

		sig := sink.Bytes()

		if len(sig) == 0 {
			t.Errorf("(detached) %s: empty sig", test.name)
		}

		varg := &SaltPackVerifyArg{
			Source: strings.NewReader(test.input),
			Opts: keybase1.SaltPackVerifyOptions{
				Signature: sig,
			},
		}
		veng := NewSaltPackVerify(varg, tc.G)

		ctx.SaltPackUI = fakeSaltPackUI{}

		if err := RunEngine(veng, ctx); err != nil {
			t.Errorf("(detached) %s: verify error: %s", test.name, err)
			continue
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"bytes"
	"io/ioutil"
	"strings"
	"testing"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
)

func TestSaltpackSignDeviceRequired(t *testing.T) {
	tc := SetupEngineTest(t, "sign")
	defer tc.Cleanup()

	ctx := &Context{
		SecretUI: &libkb.TestSecretUI{},
	}
	eng := NewSaltpackSign(nil, tc.G)
	err := RunEngine(eng, ctx)
	if err == nil {
		t.Fatal("sign not logged in returned no error")
	}
	if _, ok := err.(libkb.DeviceRequiredError); !ok {
		t.Errorf("error type: %T, expected DeviceRequiredError", err)
	}
}

func TestSaltpackSignVerify(t *testing.T) {
	tc := SetupEngineTest(t, "sign")
	defer tc.Cleanup()

	fu := CreateAndSignupFakeUser(tc, "sign")

	// signTests are defined in pgp_sign_test.  Make sure that saltpack sign can
	// sign/verify the same messages as pgp.
	for _, test := range signTests {
		var sink bytes.Buffer

		sarg := &SaltpackSignArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: ioutil.NopCloser(bytes.NewBufferString(test.input)),
		}

		eng := NewSaltpackSign(sarg, tc.G)
		ctx := &Context{
			IdentifyUI: &FakeIdentifyUI{},
			SecretUI:   fu.NewSecretUI(),
		}

		if err := RunEngine(eng, ctx); err != nil {
			t.Errorf("%s: run error: %s", test.name, err)
			continue
		}

		sig := sink.String()

		if len(sig) == 0 {
			t.Errorf("%s: empty sig", test.name)
		}

		varg := &SaltpackVerifyArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: strings.NewReader(sig),
		}
		veng := NewSaltpackVerify(varg, tc.G)

		ctx.SaltpackUI = fakeSaltpackUI{}

		if err := RunEngine(veng, ctx); err != nil {
			t.Errorf("%s: verify error: %s", test.name, err)
			continue
		}
	}

	// now try the same messages, but generate detached signatures
	for _, test := range signTests {
		var sink bytes.Buffer

		sarg := &SaltpackSignArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: ioutil.NopCloser(bytes.NewBufferString(test.input)),
			Opts: keybase1.SaltpackSignOptions{
				Detached: true,
			},
		}

		eng := NewSaltpackSign(sarg, tc.G)
		ctx := &Context{
			IdentifyUI: &FakeIdentifyUI{},
			SecretUI:   fu.NewSecretUI(),
		}

		if err := RunEngine(eng, ctx); err != nil {
			t.Errorf("(detached) %s: run error: %s", test.name, err)
			continue
		}

		sig := sink.Bytes()

		if len(sig) == 0 {
			t.Errorf("(detached) %s: empty sig", test.name)
		}

		varg := &SaltpackVerifyArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: strings.NewReader(test.input),
			Opts: keybase1.SaltpackVerifyOptions{
				Signature: sig,
			},
		}
		veng := NewSaltpackVerify(varg, tc.G)

		ctx.SaltpackUI = fakeSaltpackUI{}

		if err := RunEngine(veng, ctx); err != nil {
			t.Errorf("(detached) %s: verify error: %s", test.name, err)
			continue
		}
	}
}

func TestSaltpackSignVerifyBinary(t *testing.T) {
	tc := SetupEngineTest(t, "sign")
	defer tc.Cleanup()

	fu := CreateAndSignupFakeUser(tc, "sign")

	// signTests are defined in pgp_sign_test.  Make sure that saltpack sign can
	// sign/verify the same messages as pgp.
	for _, test := range signTests {
		var sink bytes.Buffer

		sarg := &SaltpackSignArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: ioutil.NopCloser(bytes.NewBufferString(test.input)),
			Opts: keybase1.SaltpackSignOptions{
				Binary: true,
			},
		}

		eng := NewSaltpackSign(sarg, tc.G)
		ctx := &Context{
			IdentifyUI: &FakeIdentifyUI{},
			SecretUI:   fu.NewSecretUI(),
		}

		if err := RunEngine(eng, ctx); err != nil {
			t.Errorf("%s: run error: %s", test.name, err)
			continue
		}

		sig := sink.String()

		if len(sig) == 0 {
			t.Errorf("%s: empty sig", test.name)
		}

		varg := &SaltpackVerifyArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: strings.NewReader(sig),
		}
		veng := NewSaltpackVerify(varg, tc.G)

		ctx.SaltpackUI = fakeSaltpackUI{}

		if err := RunEngine(veng, ctx); err != nil {
			t.Errorf("%s: verify error: %s", test.name, err)
			continue
		}
	}

	// now try the same messages, but generate detached signatures
	for _, test := range signTests {
		var sink bytes.Buffer

		sarg := &SaltpackSignArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: ioutil.NopCloser(bytes.NewBufferString(test.input)),
			Opts: keybase1.SaltpackSignOptions{
				Binary:   true,
				Detached: true,
			},
		}

		eng := NewSaltpackSign(sarg, tc.G)
		ctx := &Context{
			IdentifyUI: &FakeIdentifyUI{},
			SecretUI:   fu.NewSecretUI(),
		}

		if err := RunEngine(eng, ctx); err != nil {
			t.Errorf("(detached) %s: run error: %s", test.name, err)
			continue
		}

		sig := sink.Bytes()

		if len(sig) == 0 {
			t.Errorf("(detached) %s: empty sig", test.name)
		}

		varg := &SaltpackVerifyArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: strings.NewReader(test.input),
			Opts: keybase1.SaltpackVerifyOptions{
				Signature: sig,
			},
		}
		veng := NewSaltpackVerify(varg, tc.G)

		ctx.SaltpackUI = fakeSaltpackUI{}

		if err := RunEngine(veng, ctx); err != nil {
			t.Errorf("(detached) %s: verify error: %s", test.name, err)
			continue
		}
	}
}
<commit_msg>Add tests for SignedBy option (for self and non-self)<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package engine

import (
	"bytes"
	"io/ioutil"
	"strings"
	"testing"

	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
)

func TestSaltpackSignDeviceRequired(t *testing.T) {
	tc := SetupEngineTest(t, "sign")
	defer tc.Cleanup()

	ctx := &Context{
		SecretUI: &libkb.TestSecretUI{},
	}
	eng := NewSaltpackSign(nil, tc.G)
	err := RunEngine(eng, ctx)
	if err == nil {
		t.Fatal("sign not logged in returned no error")
	}
	if _, ok := err.(libkb.DeviceRequiredError); !ok {
		t.Errorf("error type: %T, expected DeviceRequiredError", err)
	}
}

func TestSaltpackSignVerify(t *testing.T) {
	tc := SetupEngineTest(t, "sign")
	defer tc.Cleanup()

	fu := CreateAndSignupFakeUser(tc, "sign")

	// signTests are defined in pgp_sign_test.  Make sure that saltpack sign can
	// sign/verify the same messages as pgp.
	for _, test := range signTests {
		var sink bytes.Buffer

		sarg := &SaltpackSignArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: ioutil.NopCloser(bytes.NewBufferString(test.input)),
		}

		eng := NewSaltpackSign(sarg, tc.G)
		ctx := &Context{
			IdentifyUI: &FakeIdentifyUI{},
			SecretUI:   fu.NewSecretUI(),
		}

		if err := RunEngine(eng, ctx); err != nil {
			t.Errorf("%s: run error: %s", test.name, err)
			continue
		}

		sig := sink.String()

		if len(sig) == 0 {
			t.Errorf("%s: empty sig", test.name)
		}

		varg := &SaltpackVerifyArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: strings.NewReader(sig),
		}
		veng := NewSaltpackVerify(varg, tc.G)

		ctx.SaltpackUI = fakeSaltpackUI{}

		if err := RunEngine(veng, ctx); err != nil {
			t.Errorf("%s: verify error: %s", test.name, err)
			continue
		}

		// test SignedBy option:
		varg = &SaltpackVerifyArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: strings.NewReader(sig),
			Opts: keybase1.SaltpackVerifyOptions{
				SignedBy: fu.Username,
			},
		}
		veng = NewSaltpackVerify(varg, tc.G)
		if err := RunEngine(veng, ctx); err != nil {
			t.Errorf("%s: verify w/ SignedBy error: %s", test.name, err)
			continue
		}

		varg = &SaltpackVerifyArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: strings.NewReader(sig),
			Opts: keybase1.SaltpackVerifyOptions{
				SignedBy: "unknown",
			},
		}
		veng = NewSaltpackVerify(varg, tc.G)
		if err := RunEngine(veng, ctx); err == nil {
			t.Errorf("%s: verify w/ SignedBy=unknown worked, should have failed", test.name)
			continue
		}
	}

	// now try the same messages, but generate detached signatures
	for _, test := range signTests {
		var sink bytes.Buffer

		sarg := &SaltpackSignArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: ioutil.NopCloser(bytes.NewBufferString(test.input)),
			Opts: keybase1.SaltpackSignOptions{
				Detached: true,
			},
		}

		eng := NewSaltpackSign(sarg, tc.G)
		ctx := &Context{
			IdentifyUI: &FakeIdentifyUI{},
			SecretUI:   fu.NewSecretUI(),
		}

		if err := RunEngine(eng, ctx); err != nil {
			t.Errorf("(detached) %s: run error: %s", test.name, err)
			continue
		}

		sig := sink.Bytes()

		if len(sig) == 0 {
			t.Errorf("(detached) %s: empty sig", test.name)
		}

		varg := &SaltpackVerifyArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: strings.NewReader(test.input),
			Opts: keybase1.SaltpackVerifyOptions{
				Signature: sig,
			},
		}
		veng := NewSaltpackVerify(varg, tc.G)

		ctx.SaltpackUI = fakeSaltpackUI{}

		if err := RunEngine(veng, ctx); err != nil {
			t.Errorf("(detached) %s: verify error: %s", test.name, err)
			continue
		}
	}
}

func TestSaltpackSignVerifyBinary(t *testing.T) {
	tc := SetupEngineTest(t, "sign")
	defer tc.Cleanup()

	fu := CreateAndSignupFakeUser(tc, "sign")

	// signTests are defined in pgp_sign_test.  Make sure that saltpack sign can
	// sign/verify the same messages as pgp.
	for _, test := range signTests {
		var sink bytes.Buffer

		sarg := &SaltpackSignArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: ioutil.NopCloser(bytes.NewBufferString(test.input)),
			Opts: keybase1.SaltpackSignOptions{
				Binary: true,
			},
		}

		eng := NewSaltpackSign(sarg, tc.G)
		ctx := &Context{
			IdentifyUI: &FakeIdentifyUI{},
			SecretUI:   fu.NewSecretUI(),
		}

		if err := RunEngine(eng, ctx); err != nil {
			t.Errorf("%s: run error: %s", test.name, err)
			continue
		}

		sig := sink.String()

		if len(sig) == 0 {
			t.Errorf("%s: empty sig", test.name)
		}

		varg := &SaltpackVerifyArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: strings.NewReader(sig),
		}
		veng := NewSaltpackVerify(varg, tc.G)

		ctx.SaltpackUI = fakeSaltpackUI{}

		if err := RunEngine(veng, ctx); err != nil {
			t.Errorf("%s: verify error: %s", test.name, err)
			continue
		}
	}

	// now try the same messages, but generate detached signatures
	for _, test := range signTests {
		var sink bytes.Buffer

		sarg := &SaltpackSignArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: ioutil.NopCloser(bytes.NewBufferString(test.input)),
			Opts: keybase1.SaltpackSignOptions{
				Binary:   true,
				Detached: true,
			},
		}

		eng := NewSaltpackSign(sarg, tc.G)
		ctx := &Context{
			IdentifyUI: &FakeIdentifyUI{},
			SecretUI:   fu.NewSecretUI(),
		}

		if err := RunEngine(eng, ctx); err != nil {
			t.Errorf("(detached) %s: run error: %s", test.name, err)
			continue
		}

		sig := sink.Bytes()

		if len(sig) == 0 {
			t.Errorf("(detached) %s: empty sig", test.name)
		}

		varg := &SaltpackVerifyArg{
			Sink:   libkb.NopWriteCloser{W: &sink},
			Source: strings.NewReader(test.input),
			Opts: keybase1.SaltpackVerifyOptions{
				Signature: sig,
			},
		}
		veng := NewSaltpackVerify(varg, tc.G)

		ctx.SaltpackUI = fakeSaltpackUI{}

		if err := RunEngine(veng, ctx); err != nil {
			t.Errorf("(detached) %s: verify error: %s", test.name, err)
			continue
		}
	}
}

func TestSaltpackSignVerifyNotSelf(t *testing.T) {
	tc := SetupEngineTest(t, "sign")
	defer tc.Cleanup()

	signer := CreateAndSignupFakeUser(tc, "sign")

	var sink bytes.Buffer

	sarg := &SaltpackSignArg{
		Sink:   libkb.NopWriteCloser{W: &sink},
		Source: ioutil.NopCloser(bytes.NewBufferString("this is from me")),
	}

	eng := NewSaltpackSign(sarg, tc.G)
	ctx := &Context{
		IdentifyUI: &FakeIdentifyUI{},
		SecretUI:   signer.NewSecretUI(),
	}

	if err := RunEngine(eng, ctx); err != nil {
		t.Fatal(err)
	}

	sig := sink.String()

	if len(sig) == 0 {
		t.Fatal("empty sig")
	}

	Logout(tc)

	_ = CreateAndSignupFakeUser(tc, "sign")

	// no user assertion
	varg := &SaltpackVerifyArg{
		Sink:   libkb.NopWriteCloser{W: &sink},
		Source: strings.NewReader(sig),
	}
	veng := NewSaltpackVerify(varg, tc.G)

	ctx.SaltpackUI = fakeSaltpackUI{}

	if err := RunEngine(veng, ctx); err != nil {
		t.Fatalf("verify error: %s", err)
	}

	// valid user assertion
	varg = &SaltpackVerifyArg{
		Sink:   libkb.NopWriteCloser{W: &sink},
		Source: strings.NewReader(sig),
		Opts: keybase1.SaltpackVerifyOptions{
			SignedBy: signer.Username,
		},
	}
	veng = NewSaltpackVerify(varg, tc.G)
	if err := RunEngine(veng, ctx); err != nil {
		t.Fatalf("verify w/ SignedBy error: %s", err)
	}

	// invalid user assertion
	varg = &SaltpackVerifyArg{
		Sink:   libkb.NopWriteCloser{W: &sink},
		Source: strings.NewReader(sig),
		Opts: keybase1.SaltpackVerifyOptions{
			SignedBy: "unknown",
		},
	}
	veng = NewSaltpackVerify(varg, tc.G)
	if err := RunEngine(veng, ctx); err == nil {
		t.Errorf("verify w/ SignedBy unknown didn't fail")
	}
}
<|endoftext|>


<commit_before>package forward

import (
	"log"

	"github.com/moby/vpnkit/go/pkg/libproxy"
	"github.com/moby/vpnkit/go/pkg/vpnkit"
)

// Listen on stream sockets and forward to a remote multiplexer.

type network interface {
	listen(vpnkit.Port) (listener, error)
}

type listener interface {
	accept() (libproxy.Conn, error)
	close() error
}

func makeStream(c common, n network) (*stream, error) {
	l, err := n.listen(c.port)
	if err != nil {
		return nil, err
	}

	return &stream{
		c,
		l,
	}, nil
}

type stream struct {
	common
	l listener
}

func (s *stream) Run() {
	for {
		src, err := s.l.accept()
		if err != nil {
			log.Printf("Stopping accepting connections on %s", s.port.String())
			return
		}
		mux := s.ctrl.Mux()
		dest, err := mux.Dial(*s.dest)
		if err != nil {
			log.Printf("unable to connect on %s: %s", s.port.String(), err)
			if err := src.Close(); err != nil {
				log.Printf("unable to Close on %s: %s", s.port.String(), err)
			}
			return
		}
		go func() {
			if err := libproxy.ProxyStream(src, dest, s.quit); err != nil {
				log.Printf("unable to proxy on %s: %s", s.port.String(), err)
			}
			if err := src.Close(); err != nil {
				log.Printf("unable to Close on %s: %s", s.port.String(), err)
			}
		}()

	}
}

func (s *stream) Stop() {
	log.Printf("Removing %s", s.port.String())
	s.l.close()
	close(s.quit)
}
<commit_msg>stream: don't stop forwarding if the multiplexer terminates<commit_after>package forward

import (
	"log"

	"github.com/moby/vpnkit/go/pkg/libproxy"
	"github.com/moby/vpnkit/go/pkg/vpnkit"
)

// Listen on stream sockets and forward to a remote multiplexer.

type network interface {
	listen(vpnkit.Port) (listener, error)
}

type listener interface {
	accept() (libproxy.Conn, error)
	close() error
}

func makeStream(c common, n network) (*stream, error) {
	l, err := n.listen(c.port)
	if err != nil {
		return nil, err
	}

	return &stream{
		c,
		l,
	}, nil
}

type stream struct {
	common
	l listener
}

func (s *stream) Run() {
	for {
		src, err := s.l.accept()
		if err != nil {
			log.Printf("Stopping accepting connections on %s", s.port.String())
			return
		}
		mux := s.ctrl.Mux()
		dest, err := mux.Dial(*s.dest)
		if err != nil {
			log.Printf("unable to connect on %s: %s", s.port.String(), err)
			if err := src.Close(); err != nil {
				log.Printf("unable to Close on %s: %s", s.port.String(), err)
			}
			continue // Multiplexer could be disconnected
		}
		go func() {
			if err := libproxy.ProxyStream(src, dest, s.quit); err != nil {
				log.Printf("unable to proxy on %s: %s", s.port.String(), err)
			}
			if err := src.Close(); err != nil {
				log.Printf("unable to Close on %s: %s", s.port.String(), err)
			}
		}()

	}
}

func (s *stream) Stop() {
	log.Printf("Removing %s", s.port.String())
	s.l.close()
	close(s.quit)
}
<|endoftext|>


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

package vtgate

import (
	"flag"
	"sync"
	"time"

	log "github.com/golang/glog"

	"github.com/youtube/vitess/go/stats"
	"github.com/youtube/vitess/go/vt/topo"
)

var (
	srvTopoCacheTTL = flag.Duration("srv_topo_cache_ttl", 1*time.Second, "how long to use cached entries for topology")
)

// SrvTopoServer is a subset of topo.Server that only contains the serving
// graph read-only calls used by clients to resolve serving addresses.
type SrvTopoServer interface {
	GetSrvKeyspaceNames(cell string) ([]string, error)

	GetSrvKeyspace(cell, keyspace string) (*topo.SrvKeyspace, error)

	GetEndPoints(cell, keyspace, shard string, tabletType topo.TabletType) (*topo.EndPoints, error)
}

// ResilientSrvTopoServer is an implementation of SrvTopoServer based
// on another SrvTopoServer that uses a cache for two purposes:
// - limit the QPS to the underlying SrvTopoServer
// - return the last known value of the data if there is an error
type ResilientSrvTopoServer struct {
	topoServer SrvTopoServer

	// stats for the cache layer
	queryCount  *stats.Counters
	cachedCount *stats.Counters
	errorCount  *stats.Counters

	// mu protects the cache map itself, not the individual values
	// in the cache.
	mutex                 sync.Mutex
	srvKeyspaceNamesCache map[string]*srvKeyspaceNamesEntry
	srvKeyspaceCache      map[string]*srvKeyspaceEntry
	endPointsCache        map[string]*endPointsEntry
}

type srvKeyspaceNamesEntry struct {
	// the mutex protects any access to this structure (read or write)
	mutex sync.Mutex

	insertionTime time.Time
	value         []string
}

type srvKeyspaceEntry struct {
	// the mutex protects any access to this structure (read or write)
	mutex sync.Mutex

	insertionTime time.Time
	value         *topo.SrvKeyspace
}

type endPointsEntry struct {
	// the mutex protects any access to this structure (read or write)
	mutex sync.Mutex

	insertionTime time.Time
	value         *topo.EndPoints
}

// NewResilientSrvTopoServer creates a new ResilientSrvTopoServer
// based on the provided SrvTopoServer.
func NewResilientSrvTopoServer(base SrvTopoServer) *ResilientSrvTopoServer {
	return &ResilientSrvTopoServer{
		topoServer: base,

		queryCount:  stats.NewCounters("ResilientSrvTopoServerQueryCount"),
		cachedCount: stats.NewCounters("ResilientSrvTopoServerCachedCount"),
		errorCount:  stats.NewCounters("ResilientSrvTopoServerErrorCount"),

		srvKeyspaceNamesCache: make(map[string]*srvKeyspaceNamesEntry),
		srvKeyspaceCache:      make(map[string]*srvKeyspaceEntry),
		endPointsCache:        make(map[string]*endPointsEntry),
	}
}

func (server *ResilientSrvTopoServer) GetSrvKeyspaceNames(cell string) ([]string, error) {
	// find the entry in the cache, add it if not there
	key := cell
	server.queryCount.Add(key, 1)
	server.mutex.Lock()
	entry, ok := server.srvKeyspaceNamesCache[key]
	if !ok {
		entry = &srvKeyspaceNamesEntry{}
		server.srvKeyspaceNamesCache[key] = entry
	}
	server.mutex.Unlock()

	// Lock the entry, and do everything holding the lock.  This
	// means two concurrent requests will only issue one
	// underlying query.
	entry.mutex.Lock()
	defer entry.mutex.Unlock()

	// If the entry is fresh enough, return it
	if time.Now().Sub(entry.insertionTime) < *srvTopoCacheTTL {
		return entry.value, nil
	}

	// not in cache or too old, get the real value
	result, err := server.topoServer.GetSrvKeyspaceNames(cell)
	if err != nil {
		if entry.insertionTime.IsZero() {
			server.errorCount.Add(key, 1)
			log.Errorf("GetSrvKeyspaceNames(%v) failed: %v (no cached value, returning error)", cell, err)
			return nil, err
		} else {
			server.cachedCount.Add(key, 1)
			log.Warningf("GetSrvKeyspaceNames(%v) failed: %v (returning cached value)", cell, err)
			return entry.value, nil
		}
	}

	// save the value we got and the current time in the cache
	entry.insertionTime = time.Now()
	entry.value = result
	return result, nil
}

func (server *ResilientSrvTopoServer) GetSrvKeyspace(cell, keyspace string) (*topo.SrvKeyspace, error) {
	// find the entry in the cache, add it if not there
	key := cell + ":" + keyspace
	server.queryCount.Add(key, 1)
	server.mutex.Lock()
	entry, ok := server.srvKeyspaceCache[key]
	if !ok {
		entry = &srvKeyspaceEntry{}
		server.srvKeyspaceCache[key] = entry
	}
	server.mutex.Unlock()

	// Lock the entry, and do everything holding the lock.  This
	// means two concurrent requests will only issue one
	// underlying query.
	entry.mutex.Lock()
	defer entry.mutex.Unlock()

	// If the entry is fresh enough, return it
	if time.Now().Sub(entry.insertionTime) < *srvTopoCacheTTL {
		return entry.value, nil
	}

	// not in cache or too old, get the real value
	result, err := server.topoServer.GetSrvKeyspace(cell, keyspace)
	if err != nil {
		if entry.insertionTime.IsZero() {
			server.errorCount.Add(key, 1)
			log.Errorf("GetSrvKeyspace(%v, %v) failed: %v (no cached value, returning error)", cell, keyspace, err)
			return nil, err
		} else {
			server.cachedCount.Add(key, 1)
			log.Warningf("GetSrvKeyspace(%v, %v) failed: %v (returning cached value)", cell, keyspace, err)
			return entry.value, nil
		}
	}

	// save the value we got and the current time in the cache
	entry.insertionTime = time.Now()
	entry.value = result
	return result, nil
}

func (server *ResilientSrvTopoServer) GetEndPoints(cell, keyspace, shard string, tabletType topo.TabletType) (*topo.EndPoints, error) {
	// find the entry in the cache, add it if not there
	key := cell + ":" + keyspace + ":" + shard + ":" + string(tabletType)
	server.queryCount.Add(key, 1)
	server.mutex.Lock()
	entry, ok := server.endPointsCache[key]
	if !ok {
		entry = &endPointsEntry{}
		server.endPointsCache[key] = entry
	}
	server.mutex.Unlock()

	// Lock the entry, and do everything holding the lock.  This
	// means two concurrent requests will only issue one
	// underlying query.
	entry.mutex.Lock()
	defer entry.mutex.Unlock()

	// If the entry is fresh enough, return it
	if time.Now().Sub(entry.insertionTime) < *srvTopoCacheTTL {
		return entry.value, nil
	}

	// not in cache or too old, get the real value
	result, err := server.topoServer.GetEndPoints(cell, keyspace, shard, tabletType)
	if err != nil {
		if entry.insertionTime.IsZero() {
			server.errorCount.Add(key, 1)
			log.Errorf("GetEndPoints(%v, %v, %v, %v) failed: %v (no cached value, returning error)", cell, keyspace, shard, tabletType, err)
			return nil, err
		} else {
			server.cachedCount.Add(key, 1)
			log.Warningf("GetEndPoints(%v, %v%, v, %v) failed: %v (returning cached value)", cell, keyspace, shard, tabletType, err)
			return entry.value, nil
		}
	}

	// save the value we got and the current time in the cache
	entry.insertionTime = time.Now()
	entry.value = result
	return result, nil
}
<commit_msg>New stats toplevel counters.<commit_after>// Copyright 2013, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package vtgate

import (
	"flag"
	"sync"
	"time"

	log "github.com/golang/glog"

	"github.com/youtube/vitess/go/stats"
	"github.com/youtube/vitess/go/vt/topo"
)

var (
	srvTopoCacheTTL = flag.Duration("srv_topo_cache_ttl", 1*time.Second, "how long to use cached entries for topology")
)

// SrvTopoServer is a subset of topo.Server that only contains the serving
// graph read-only calls used by clients to resolve serving addresses.
type SrvTopoServer interface {
	GetSrvKeyspaceNames(cell string) ([]string, error)

	GetSrvKeyspace(cell, keyspace string) (*topo.SrvKeyspace, error)

	GetEndPoints(cell, keyspace, shard string, tabletType topo.TabletType) (*topo.EndPoints, error)
}

// ResilientSrvTopoServer is an implementation of SrvTopoServer based
// on another SrvTopoServer that uses a cache for two purposes:
// - limit the QPS to the underlying SrvTopoServer
// - return the last known value of the data if there is an error
type ResilientSrvTopoServer struct {
	topoServer SrvTopoServer

	// stats for the cache layer
	queryCount  *stats.Counters
	cachedCount *stats.Counters
	errorCount  *stats.Counters

	// mu protects the cache map itself, not the individual values
	// in the cache.
	mutex                 sync.Mutex
	srvKeyspaceNamesCache map[string]*srvKeyspaceNamesEntry
	srvKeyspaceCache      map[string]*srvKeyspaceEntry
	endPointsCache        map[string]*endPointsEntry
}

type srvKeyspaceNamesEntry struct {
	// the mutex protects any access to this structure (read or write)
	mutex sync.Mutex

	insertionTime time.Time
	value         []string
}

type srvKeyspaceEntry struct {
	// the mutex protects any access to this structure (read or write)
	mutex sync.Mutex

	insertionTime time.Time
	value         *topo.SrvKeyspace
}

type endPointsEntry struct {
	// the mutex protects any access to this structure (read or write)
	mutex sync.Mutex

	insertionTime time.Time
	value         *topo.EndPoints
}

// NewResilientSrvTopoServer creates a new ResilientSrvTopoServer
// based on the provided SrvTopoServer.
func NewResilientSrvTopoServer(base SrvTopoServer) *ResilientSrvTopoServer {
	return &ResilientSrvTopoServer{
		topoServer: base,

		queryCount:  stats.NewCounters("ResilientSrvTopoServerQueryCount"),
		cachedCount: stats.NewCounters("ResilientSrvTopoServerCachedCount"),
		errorCount:  stats.NewCounters("ResilientSrvTopoServerErrorCount"),

		srvKeyspaceNamesCache: make(map[string]*srvKeyspaceNamesEntry),
		srvKeyspaceCache:      make(map[string]*srvKeyspaceEntry),
		endPointsCache:        make(map[string]*endPointsEntry),
	}
}

func (server *ResilientSrvTopoServer) GetSrvKeyspaceNames(cell string) ([]string, error) {
	server.queryCount.Add(cell, 1)

	// find the entry in the cache, add it if not there
	key := cell
	server.mutex.Lock()
	entry, ok := server.srvKeyspaceNamesCache[key]
	if !ok {
		entry = &srvKeyspaceNamesEntry{}
		server.srvKeyspaceNamesCache[key] = entry
	}
	server.mutex.Unlock()

	// Lock the entry, and do everything holding the lock.  This
	// means two concurrent requests will only issue one
	// underlying query.
	entry.mutex.Lock()
	defer entry.mutex.Unlock()

	// If the entry is fresh enough, return it
	if time.Now().Sub(entry.insertionTime) < *srvTopoCacheTTL {
		return entry.value, nil
	}

	// not in cache or too old, get the real value
	result, err := server.topoServer.GetSrvKeyspaceNames(cell)
	if err != nil {
		if entry.insertionTime.IsZero() {
			server.errorCount.Add(cell, 1)
			log.Errorf("GetSrvKeyspaceNames(%v) failed: %v (no cached value, returning error)", cell, err)
			return nil, err
		} else {
			server.cachedCount.Add(cell, 1)
			log.Warningf("GetSrvKeyspaceNames(%v) failed: %v (returning cached value)", cell, err)
			return entry.value, nil
		}
	}

	// save the value we got and the current time in the cache
	entry.insertionTime = time.Now()
	entry.value = result
	return result, nil
}

func (server *ResilientSrvTopoServer) GetSrvKeyspace(cell, keyspace string) (*topo.SrvKeyspace, error) {
	server.queryCount.Add(cell, 1)

	// find the entry in the cache, add it if not there
	key := cell + ":" + keyspace
	server.mutex.Lock()
	entry, ok := server.srvKeyspaceCache[key]
	if !ok {
		entry = &srvKeyspaceEntry{}
		server.srvKeyspaceCache[key] = entry
	}
	server.mutex.Unlock()

	// Lock the entry, and do everything holding the lock.  This
	// means two concurrent requests will only issue one
	// underlying query.
	entry.mutex.Lock()
	defer entry.mutex.Unlock()

	// If the entry is fresh enough, return it
	if time.Now().Sub(entry.insertionTime) < *srvTopoCacheTTL {
		return entry.value, nil
	}

	// not in cache or too old, get the real value
	result, err := server.topoServer.GetSrvKeyspace(cell, keyspace)
	if err != nil {
		if entry.insertionTime.IsZero() {
			server.errorCount.Add(cell, 1)
			log.Errorf("GetSrvKeyspace(%v, %v) failed: %v (no cached value, returning error)", cell, keyspace, err)
			return nil, err
		} else {
			server.cachedCount.Add(cell, 1)
			log.Warningf("GetSrvKeyspace(%v, %v) failed: %v (returning cached value)", cell, keyspace, err)
			return entry.value, nil
		}
	}

	// save the value we got and the current time in the cache
	entry.insertionTime = time.Now()
	entry.value = result
	return result, nil
}

func (server *ResilientSrvTopoServer) GetEndPoints(cell, keyspace, shard string, tabletType topo.TabletType) (*topo.EndPoints, error) {
	server.queryCount.Add(cell, 1)

	// find the entry in the cache, add it if not there
	key := cell + ":" + keyspace + ":" + shard + ":" + string(tabletType)
	server.mutex.Lock()
	entry, ok := server.endPointsCache[key]
	if !ok {
		entry = &endPointsEntry{}
		server.endPointsCache[key] = entry
	}
	server.mutex.Unlock()

	// Lock the entry, and do everything holding the lock.  This
	// means two concurrent requests will only issue one
	// underlying query.
	entry.mutex.Lock()
	defer entry.mutex.Unlock()

	// If the entry is fresh enough, return it
	if time.Now().Sub(entry.insertionTime) < *srvTopoCacheTTL {
		return entry.value, nil
	}

	// not in cache or too old, get the real value
	result, err := server.topoServer.GetEndPoints(cell, keyspace, shard, tabletType)
	if err != nil {
		if entry.insertionTime.IsZero() {
			server.errorCount.Add(cell, 1)
			log.Errorf("GetEndPoints(%v, %v, %v, %v) failed: %v (no cached value, returning error)", cell, keyspace, shard, tabletType, err)
			return nil, err
		} else {
			server.cachedCount.Add(cell, 1)
			log.Warningf("GetEndPoints(%v, %v%, v, %v) failed: %v (returning cached value)", cell, keyspace, shard, tabletType, err)
			return entry.value, nil
		}
	}

	// save the value we got and the current time in the cache
	entry.insertionTime = time.Now()
	entry.value = result
	return result, nil
}
<|endoftext|>