<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package cmd

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"sync"
	"syscall"
	"time"

	"github.com/kelseyhightower/envconfig"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/spf13/afero"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"

	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/lib/consts"
	"github.com/loadimpact/k6/loader"
	"github.com/loadimpact/k6/stats/cloud"
	"github.com/loadimpact/k6/ui"
	"github.com/loadimpact/k6/ui/pb"
)

const (
	cloudFailedToGetProgressErrorCode = 98
	cloudTestRunFailedErrorCode       = 99
)

//nolint:gochecknoglobals
var (
	exitOnRunning = os.Getenv("K6_EXIT_ON_RUNNING") != ""
	showCloudLogs = true
)

//nolint:gochecknoglobals
var cloudCmd = &cobra.Command{
	Use:   "cloud",
	Short: "Run a test on the cloud",
	Long: `Run a test on the cloud.

This will execute the test on the k6 cloud service. Use "k6 login cloud" to authenticate.`,
	Example: `
        k6 cloud script.js`[1:],
	Args: exactArgsWithMsg(1, "arg should either be \"-\", if reading script from stdin, or a path to a script file"),
	RunE: func(cmd *cobra.Command, args []string) error {
		// TODO: don't use the Global logger
		logger := logrus.StandardLogger()
		// we specifically first parse it and return an error if it has bad value and then check if
		// we are going to set it  ... so we always parse it instead of it breaking the command if
		// the cli flag is removed
		if showCloudLogsEnv, ok := os.LookupEnv("K6_SHOW_CLOUD_LOGS"); ok {
			showCloudLogsValue, err := strconv.ParseBool(showCloudLogsEnv)
			if err != nil {
				return fmt.Errorf("parsing K6_SHOW_CLOUD_LOGS returned an error: %w", err)
			}
			if !cmd.Flags().Changed("show-logs") {
				showCloudLogs = showCloudLogsValue
			}

		}
		// TODO: disable in quiet mode?
		_, _ = BannerColor.Fprintf(stdout, "\n%s\n\n", consts.Banner())

		progressBar := pb.New(
			pb.WithConstLeft("Init"),
			pb.WithConstProgress(0, "Parsing script"),
		)
		printBar(progressBar)

		// Runner
		pwd, err := os.Getwd()
		if err != nil {
			return err
		}

		filename := args[0]
		filesystems := loader.CreateFilesystems()
		src, err := loader.ReadSource(logger, filename, pwd, filesystems, os.Stdin)
		if err != nil {
			return err
		}

		runtimeOptions, err := getRuntimeOptions(cmd.Flags(), buildEnvMap(os.Environ()))
		if err != nil {
			return err
		}

		modifyAndPrintBar(progressBar, pb.WithConstProgress(0, "Getting script options"))
		r, err := newRunner(logger, src, runType, filesystems, runtimeOptions)
		if err != nil {
			return err
		}

		modifyAndPrintBar(progressBar, pb.WithConstProgress(0, "Consolidating options"))
		cliOpts, err := getOptions(cmd.Flags())
		if err != nil {
			return err
		}
		conf, err := getConsolidatedConfig(afero.NewOsFs(), Config{Options: cliOpts}, r)
		if err != nil {
			return err
		}

		derivedConf, cerr := deriveAndValidateConfig(conf, r.IsExecutable)
		if cerr != nil {
			return ExitCode{error: cerr, Code: invalidConfigErrorCode}
		}

		// TODO: validate for usage of execution segment
		// TODO: validate for externally controlled executor (i.e. executors that aren't distributable)
		// TODO: move those validations to a separate function and reuse validateConfig()?

		err = r.SetOptions(conf.Options)
		if err != nil {
			return err
		}

		// Cloud config
		cloudConfig := cloud.NewConfig().Apply(derivedConf.Collectors.Cloud)
		if err = envconfig.Process("", &cloudConfig); err != nil {
			return err
		}
		if !cloudConfig.Token.Valid {
			return errors.New("Not logged in, please use `k6 login cloud`.")
		}

		modifyAndPrintBar(progressBar, pb.WithConstProgress(0, "Building the archive"))
		arc := r.MakeArchive()
		// TODO: Fix this
		// We reuse cloud.Config for parsing options.ext.loadimpact, but this probably shouldn't be
		// done as the idea of options.ext is that they are extensible without touching k6. But in
		// order for this to happen we shouldn't actually marshall cloud.Config on top of it because
		// it will be missing some fields that aren't actually mentioned in the struct.
		// So in order for use to copy the fields that we need for loadimpact's api we unmarshal in
		// map[string]interface{} and copy what we need if it isn't set already
		var tmpCloudConfig map[string]interface{}
		if val, ok := arc.Options.External["loadimpact"]; ok {
			dec := json.NewDecoder(bytes.NewReader(val))
			dec.UseNumber() // otherwise float64 are used
			if err := dec.Decode(&tmpCloudConfig); err != nil {
				return err
			}
		}

		if err := cloud.MergeFromExternal(arc.Options.External, &cloudConfig); err != nil {
			return err
		}
		if tmpCloudConfig == nil {
			tmpCloudConfig = make(map[string]interface{}, 3)
		}

		if _, ok := tmpCloudConfig["token"]; !ok && cloudConfig.Token.Valid {
			tmpCloudConfig["token"] = cloudConfig.Token
		}
		if _, ok := tmpCloudConfig["name"]; !ok && cloudConfig.Name.Valid {
			tmpCloudConfig["name"] = cloudConfig.Name
		}
		if _, ok := tmpCloudConfig["projectID"]; !ok && cloudConfig.ProjectID.Valid {
			tmpCloudConfig["projectID"] = cloudConfig.ProjectID
		}

		if arc.Options.External == nil {
			arc.Options.External = make(map[string]json.RawMessage)
		}
		arc.Options.External["loadimpact"], err = json.Marshal(tmpCloudConfig)
		if err != nil {
			return err
		}

		name := cloudConfig.Name.String
		if !cloudConfig.Name.Valid || cloudConfig.Name.String == "" {
			name = filepath.Base(filename)
		}

		// Start cloud test run
		modifyAndPrintBar(progressBar, pb.WithConstProgress(0, "Validating script options"))
		client := cloud.NewClient(logger, cloudConfig.Token.String, cloudConfig.Host.String, consts.Version)
		if err := client.ValidateOptions(arc.Options); err != nil {
			return err
		}

		modifyAndPrintBar(progressBar, pb.WithConstProgress(0, "Uploading archive"))
		refID, err := client.StartCloudTestRun(name, cloudConfig.ProjectID.Int64, arc)
		if err != nil {
			return err
		}

		et, err := lib.NewExecutionTuple(derivedConf.ExecutionSegment, derivedConf.ExecutionSegmentSequence)
		if err != nil {
			return err
		}
		testURL := cloud.URLForResults(refID, cloudConfig)
		executionPlan := derivedConf.Scenarios.GetFullExecutionRequirements(et)
		printExecutionDescription("cloud", filename, testURL, derivedConf, et, executionPlan, nil)

		modifyAndPrintBar(
			progressBar,
			pb.WithConstLeft("Run "),
			pb.WithConstProgress(0, "Initializing the cloud test"),
		)

		progressCtx, progressCancel := context.WithCancel(context.Background())
		progressBarWG := &sync.WaitGroup{}
		progressBarWG.Add(1)
		defer progressBarWG.Wait()
		defer progressCancel()
		go func() {
			showProgress(progressCtx, conf, []*pb.ProgressBar{progressBar}, logger)
			progressBarWG.Done()
		}()

		// The quiet option hides the progress bar and disallow aborting the test
		if quiet {
			return nil
		}

		// Trap Interrupts, SIGINTs and SIGTERMs.
		sigC := make(chan os.Signal, 1)
		signal.Notify(sigC, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
		defer signal.Stop(sigC)

		var (
			startTime   time.Time
			maxDuration time.Duration
		)
		maxDuration, _ = lib.GetEndOffset(executionPlan)

		testProgressLock := &sync.Mutex{}
		var testProgress *cloud.TestProgressResponse
		progressBar.Modify(
			pb.WithProgress(func() (float64, []string) {
				testProgressLock.Lock()
				defer testProgressLock.Unlock()

				if testProgress == nil {
					return 0, []string{"Waiting..."}
				}

				statusText := testProgress.RunStatusText

				if testProgress.RunStatus == lib.RunStatusRunning {
					if startTime.IsZero() {
						startTime = time.Now()
					}
					spent := time.Since(startTime)
					if spent > maxDuration {
						statusText = maxDuration.String()
					} else {
						statusText = fmt.Sprintf("%s/%s", pb.GetFixedLengthDuration(spent, maxDuration), maxDuration)
					}
				}

				return testProgress.Progress, []string{statusText}
			}),
		)

		ticker := time.NewTicker(time.Millisecond * 2000)
		shouldExitLoop := false
		if showCloudLogs {
			go func() {
				logger.Debug("Connecting to cloud logs server...")
				// TODO replace with another context
				if err := cloudConfig.StreamLogsToLogger(context.Background(), logger, refID, 0); err != nil {
					logger.WithError(err).Error("error while tailing cloud logs")
				}
			}()
		}

	runningLoop:
		for {
			select {
			case <-ticker.C:
				newTestProgress, progressErr := client.GetTestProgress(refID)
				if progressErr == nil {
					if (newTestProgress.RunStatus > lib.RunStatusRunning) ||
						(exitOnRunning && newTestProgress.RunStatus == lib.RunStatusRunning) {
						shouldExitLoop = true
					}
					testProgressLock.Lock()
					testProgress = newTestProgress
					testProgressLock.Unlock()
				} else {
					logger.WithError(progressErr).Error("Test progress error")
				}
				if shouldExitLoop {
					break runningLoop
				}
			case sig := <-sigC:
				logger.WithField("sig", sig).Print("Exiting in response to signal...")
				err := client.StopCloudTestRun(refID)
				if err != nil {
					logger.WithError(err).Error("Stop cloud test error")
				}
				shouldExitLoop = true // Exit after the next GetTestProgress call
			}
		}

		if testProgress == nil {
			//nolint:golint
			return ExitCode{error: errors.New("Test progress error"), Code: cloudFailedToGetProgressErrorCode}
		}

		fprintf(stdout, "     test status: %s\n", ui.ValueColor.Sprint(testProgress.RunStatusText))

		if testProgress.ResultStatus == cloud.ResultStatusFailed {
			//nolint:golint
			return ExitCode{error: errors.New("The test has failed"), Code: cloudTestRunFailedErrorCode}
		}

		return nil
	},
}

func cloudCmdFlagSet() *pflag.FlagSet {
	flags := pflag.NewFlagSet("", pflag.ContinueOnError)
	flags.SortFlags = false
	flags.AddFlagSet(optionFlagSet())
	flags.AddFlagSet(runtimeOptionFlagSet(false))

	// TODO: Figure out a better way to handle the CLI flags:
	// - the default value is specified in this way so we don't overwrire whatever
	//   was specified via the environment variable
	// - global variables are not very testable... :/
	flags.BoolVar(&exitOnRunning, "exit-on-running", exitOnRunning, "exits when test reaches the running status")
	// We also need to explicitly set the default value for the usage message here, so setting
	// K6_EXIT_ON_RUNNING=true won't affect the usage message
	flags.Lookup("exit-on-running").DefValue = "false"

	// read the comments above for explanation why this is done this way and what are the problems
	flags.BoolVar(&showCloudLogs, "show-logs", showCloudLogs,
		"enable showing of logs when a test is executed in the cloud")

	return flags
}

func init() {
	RootCmd.AddCommand(cloudCmd)
	cloudCmd.Flags().SortFlags = false
	cloudCmd.Flags().AddFlagSet(cloudCmdFlagSet())
}
<commit_msg>Bump cloud progress to 100% when a test finishes<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package cmd

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"sync"
	"syscall"
	"time"

	"github.com/kelseyhightower/envconfig"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/spf13/afero"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"

	"github.com/loadimpact/k6/lib"
	"github.com/loadimpact/k6/lib/consts"
	"github.com/loadimpact/k6/loader"
	"github.com/loadimpact/k6/stats/cloud"
	"github.com/loadimpact/k6/ui"
	"github.com/loadimpact/k6/ui/pb"
)

const (
	cloudFailedToGetProgressErrorCode = 98
	cloudTestRunFailedErrorCode       = 99
)

//nolint:gochecknoglobals
var (
	exitOnRunning = os.Getenv("K6_EXIT_ON_RUNNING") != ""
	showCloudLogs = true
)

//nolint:gochecknoglobals
var cloudCmd = &cobra.Command{
	Use:   "cloud",
	Short: "Run a test on the cloud",
	Long: `Run a test on the cloud.

This will execute the test on the k6 cloud service. Use "k6 login cloud" to authenticate.`,
	Example: `
        k6 cloud script.js`[1:],
	Args: exactArgsWithMsg(1, "arg should either be \"-\", if reading script from stdin, or a path to a script file"),
	RunE: func(cmd *cobra.Command, args []string) error {
		// TODO: don't use the Global logger
		logger := logrus.StandardLogger()
		// we specifically first parse it and return an error if it has bad value and then check if
		// we are going to set it  ... so we always parse it instead of it breaking the command if
		// the cli flag is removed
		if showCloudLogsEnv, ok := os.LookupEnv("K6_SHOW_CLOUD_LOGS"); ok {
			showCloudLogsValue, err := strconv.ParseBool(showCloudLogsEnv)
			if err != nil {
				return fmt.Errorf("parsing K6_SHOW_CLOUD_LOGS returned an error: %w", err)
			}
			if !cmd.Flags().Changed("show-logs") {
				showCloudLogs = showCloudLogsValue
			}

		}
		// TODO: disable in quiet mode?
		_, _ = BannerColor.Fprintf(stdout, "\n%s\n\n", consts.Banner())

		progressBar := pb.New(
			pb.WithConstLeft("Init"),
			pb.WithConstProgress(0, "Parsing script"),
		)
		printBar(progressBar)

		// Runner
		pwd, err := os.Getwd()
		if err != nil {
			return err
		}

		filename := args[0]
		filesystems := loader.CreateFilesystems()
		src, err := loader.ReadSource(logger, filename, pwd, filesystems, os.Stdin)
		if err != nil {
			return err
		}

		runtimeOptions, err := getRuntimeOptions(cmd.Flags(), buildEnvMap(os.Environ()))
		if err != nil {
			return err
		}

		modifyAndPrintBar(progressBar, pb.WithConstProgress(0, "Getting script options"))
		r, err := newRunner(logger, src, runType, filesystems, runtimeOptions)
		if err != nil {
			return err
		}

		modifyAndPrintBar(progressBar, pb.WithConstProgress(0, "Consolidating options"))
		cliOpts, err := getOptions(cmd.Flags())
		if err != nil {
			return err
		}
		conf, err := getConsolidatedConfig(afero.NewOsFs(), Config{Options: cliOpts}, r)
		if err != nil {
			return err
		}

		derivedConf, cerr := deriveAndValidateConfig(conf, r.IsExecutable)
		if cerr != nil {
			return ExitCode{error: cerr, Code: invalidConfigErrorCode}
		}

		// TODO: validate for usage of execution segment
		// TODO: validate for externally controlled executor (i.e. executors that aren't distributable)
		// TODO: move those validations to a separate function and reuse validateConfig()?

		err = r.SetOptions(conf.Options)
		if err != nil {
			return err
		}

		// Cloud config
		cloudConfig := cloud.NewConfig().Apply(derivedConf.Collectors.Cloud)
		if err = envconfig.Process("", &cloudConfig); err != nil {
			return err
		}
		if !cloudConfig.Token.Valid {
			return errors.New("Not logged in, please use `k6 login cloud`.")
		}

		modifyAndPrintBar(progressBar, pb.WithConstProgress(0, "Building the archive"))
		arc := r.MakeArchive()
		// TODO: Fix this
		// We reuse cloud.Config for parsing options.ext.loadimpact, but this probably shouldn't be
		// done as the idea of options.ext is that they are extensible without touching k6. But in
		// order for this to happen we shouldn't actually marshall cloud.Config on top of it because
		// it will be missing some fields that aren't actually mentioned in the struct.
		// So in order for use to copy the fields that we need for loadimpact's api we unmarshal in
		// map[string]interface{} and copy what we need if it isn't set already
		var tmpCloudConfig map[string]interface{}
		if val, ok := arc.Options.External["loadimpact"]; ok {
			dec := json.NewDecoder(bytes.NewReader(val))
			dec.UseNumber() // otherwise float64 are used
			if err := dec.Decode(&tmpCloudConfig); err != nil {
				return err
			}
		}

		if err := cloud.MergeFromExternal(arc.Options.External, &cloudConfig); err != nil {
			return err
		}
		if tmpCloudConfig == nil {
			tmpCloudConfig = make(map[string]interface{}, 3)
		}

		if _, ok := tmpCloudConfig["token"]; !ok && cloudConfig.Token.Valid {
			tmpCloudConfig["token"] = cloudConfig.Token
		}
		if _, ok := tmpCloudConfig["name"]; !ok && cloudConfig.Name.Valid {
			tmpCloudConfig["name"] = cloudConfig.Name
		}
		if _, ok := tmpCloudConfig["projectID"]; !ok && cloudConfig.ProjectID.Valid {
			tmpCloudConfig["projectID"] = cloudConfig.ProjectID
		}

		if arc.Options.External == nil {
			arc.Options.External = make(map[string]json.RawMessage)
		}
		arc.Options.External["loadimpact"], err = json.Marshal(tmpCloudConfig)
		if err != nil {
			return err
		}

		name := cloudConfig.Name.String
		if !cloudConfig.Name.Valid || cloudConfig.Name.String == "" {
			name = filepath.Base(filename)
		}

		// Start cloud test run
		modifyAndPrintBar(progressBar, pb.WithConstProgress(0, "Validating script options"))
		client := cloud.NewClient(logger, cloudConfig.Token.String, cloudConfig.Host.String, consts.Version)
		if err := client.ValidateOptions(arc.Options); err != nil {
			return err
		}

		modifyAndPrintBar(progressBar, pb.WithConstProgress(0, "Uploading archive"))
		refID, err := client.StartCloudTestRun(name, cloudConfig.ProjectID.Int64, arc)
		if err != nil {
			return err
		}

		et, err := lib.NewExecutionTuple(derivedConf.ExecutionSegment, derivedConf.ExecutionSegmentSequence)
		if err != nil {
			return err
		}
		testURL := cloud.URLForResults(refID, cloudConfig)
		executionPlan := derivedConf.Scenarios.GetFullExecutionRequirements(et)
		printExecutionDescription("cloud", filename, testURL, derivedConf, et, executionPlan, nil)

		modifyAndPrintBar(
			progressBar,
			pb.WithConstLeft("Run "),
			pb.WithConstProgress(0, "Initializing the cloud test"),
		)

		progressCtx, progressCancel := context.WithCancel(context.Background())
		progressBarWG := &sync.WaitGroup{}
		progressBarWG.Add(1)
		defer progressBarWG.Wait()
		defer progressCancel()
		go func() {
			showProgress(progressCtx, conf, []*pb.ProgressBar{progressBar}, logger)
			progressBarWG.Done()
		}()

		// The quiet option hides the progress bar and disallow aborting the test
		if quiet {
			return nil
		}

		// Trap Interrupts, SIGINTs and SIGTERMs.
		sigC := make(chan os.Signal, 1)
		signal.Notify(sigC, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
		defer signal.Stop(sigC)

		var (
			startTime   time.Time
			maxDuration time.Duration
		)
		maxDuration, _ = lib.GetEndOffset(executionPlan)

		testProgressLock := &sync.Mutex{}
		var testProgress *cloud.TestProgressResponse
		progressBar.Modify(
			pb.WithProgress(func() (float64, []string) {
				testProgressLock.Lock()
				defer testProgressLock.Unlock()

				if testProgress == nil {
					return 0, []string{"Waiting..."}
				}

				statusText := testProgress.RunStatusText

				if testProgress.RunStatus == lib.RunStatusFinished {
					testProgress.Progress = 1
				} else if testProgress.RunStatus == lib.RunStatusRunning {
					if startTime.IsZero() {
						startTime = time.Now()
					}
					spent := time.Since(startTime)
					if spent > maxDuration {
						statusText = maxDuration.String()
					} else {
						statusText = fmt.Sprintf("%s/%s", pb.GetFixedLengthDuration(spent, maxDuration), maxDuration)
					}
				}

				return testProgress.Progress, []string{statusText}
			}),
		)

		ticker := time.NewTicker(time.Millisecond * 2000)
		shouldExitLoop := false
		if showCloudLogs {
			go func() {
				logger.Debug("Connecting to cloud logs server...")
				// TODO replace with another context
				if err := cloudConfig.StreamLogsToLogger(context.Background(), logger, refID, 0); err != nil {
					logger.WithError(err).Error("error while tailing cloud logs")
				}
			}()
		}

	runningLoop:
		for {
			select {
			case <-ticker.C:
				newTestProgress, progressErr := client.GetTestProgress(refID)
				if progressErr == nil {
					if (newTestProgress.RunStatus > lib.RunStatusRunning) ||
						(exitOnRunning && newTestProgress.RunStatus == lib.RunStatusRunning) {
						shouldExitLoop = true
					}
					testProgressLock.Lock()
					testProgress = newTestProgress
					testProgressLock.Unlock()
				} else {
					logger.WithError(progressErr).Error("Test progress error")
				}
				if shouldExitLoop {
					break runningLoop
				}
			case sig := <-sigC:
				logger.WithField("sig", sig).Print("Exiting in response to signal...")
				err := client.StopCloudTestRun(refID)
				if err != nil {
					logger.WithError(err).Error("Stop cloud test error")
				}
				shouldExitLoop = true // Exit after the next GetTestProgress call
			}
		}

		if testProgress == nil {
			//nolint:golint
			return ExitCode{error: errors.New("Test progress error"), Code: cloudFailedToGetProgressErrorCode}
		}

		fprintf(stdout, "     test status: %s\n", ui.ValueColor.Sprint(testProgress.RunStatusText))

		if testProgress.ResultStatus == cloud.ResultStatusFailed {
			//nolint:golint
			return ExitCode{error: errors.New("The test has failed"), Code: cloudTestRunFailedErrorCode}
		}

		return nil
	},
}

func cloudCmdFlagSet() *pflag.FlagSet {
	flags := pflag.NewFlagSet("", pflag.ContinueOnError)
	flags.SortFlags = false
	flags.AddFlagSet(optionFlagSet())
	flags.AddFlagSet(runtimeOptionFlagSet(false))

	// TODO: Figure out a better way to handle the CLI flags:
	// - the default value is specified in this way so we don't overwrire whatever
	//   was specified via the environment variable
	// - global variables are not very testable... :/
	flags.BoolVar(&exitOnRunning, "exit-on-running", exitOnRunning, "exits when test reaches the running status")
	// We also need to explicitly set the default value for the usage message here, so setting
	// K6_EXIT_ON_RUNNING=true won't affect the usage message
	flags.Lookup("exit-on-running").DefValue = "false"

	// read the comments above for explanation why this is done this way and what are the problems
	flags.BoolVar(&showCloudLogs, "show-logs", showCloudLogs,
		"enable showing of logs when a test is executed in the cloud")

	return flags
}

func init() {
	RootCmd.AddCommand(cloudCmd)
	cloudCmd.Flags().SortFlags = false
	cloudCmd.Flags().AddFlagSet(cloudCmdFlagSet())
}
<|endoftext|>


<commit_before>// Copyright 2013 Julien Schmidt. All rights reserved.
// Based on the path package, Copyright 2009 The Go Authors.
// Use of this source code is governed by a BSD-style license that can be found
// at https://github.com/julienschmidt/httprouter/blob/master/LICENSE.

package gin

// CleanPath is the URL version of path.Clean, it returns a canonical URL path
// for p, eliminating . and .. elements.
//
// The following rules are applied iteratively until no further processing can
// be done:
//	1. Replace multiple slashes with a single slash.
//	2. Eliminate each . path name element (the current directory).
//	3. Eliminate each inner .. path name element (the parent directory)
//	   along with the non-.. element that precedes it.
//	4. Eliminate .. elements that begin a rooted path:
//	   that is, replace "/.." by "/" at the beginning of a path.
//
// If the result of this process is an empty string, "/" is returned
func cleanPath(p string) string {
	// Turn empty string into "/"
	if p == "" {
		return "/"
	}

	n := len(p)
	var buf []byte

	// Invariants:
	//      reading from path; r is index of next byte to process.
	//      writing to buf; w is index of next byte to write.

	// path must start with '/'
	r := 1
	w := 1

	if p[0] != '/' {
		r = 0
		buf = make([]byte, n+1)
		buf[0] = '/'
	}

	trailing := n > 2 && p[n-1] == '/'

	// A bit more clunky without a 'lazybuf' like the path package, but the loop
	// gets completely inlined (bufApp). So in contrast to the path package this
	// loop has no expensive function calls (except 1x make)

	for r < n {
		switch {
		case p[r] == '/':
			// empty path element, trailing slash is added after the end
			r++

		case p[r] == '.' && r+1 == n:
			trailing = true
			r++

		case p[r] == '.' && p[r+1] == '/':
			// . element
			r++

		case p[r] == '.' && p[r+1] == '.' && (r+2 == n || p[r+2] == '/'):
			// .. element: remove to last /
			r += 2

			if w > 1 {
				// can backtrack
				w--

				if buf == nil {
					for w > 1 && p[w] != '/' {
						w--
					}
				} else {
					for w > 1 && buf[w] != '/' {
						w--
					}
				}
			}

		default:
			// real path element.
			// add slash if needed
			if w > 1 {
				bufApp(&buf, p, w, '/')
				w++
			}

			// copy element
			for r < n && p[r] != '/' {
				bufApp(&buf, p, w, p[r])
				w++
				r++
			}
		}
	}

	// re-append trailing slash
	if trailing && w > 1 {
		bufApp(&buf, p, w, '/')
		w++
	}

	if buf == nil {
		return p[:w]
	}
	return string(buf[:w])
}

// internal helper to lazily create a buffer if necessary
func bufApp(buf *[]byte, s string, w int, c byte) {
	if *buf == nil {
		if s[w] == c {
			return
		}

		*buf = make([]byte, len(s))
		copy(*buf, s[:w])
	}
	(*buf)[w] = c
}
<commit_msg>fix cleanPath spell (#969)<commit_after>// Copyright 2013 Julien Schmidt. All rights reserved.
// Based on the path package, Copyright 2009 The Go Authors.
// Use of this source code is governed by a BSD-style license that can be found
// at https://github.com/julienschmidt/httprouter/blob/master/LICENSE.

package gin

// cleanPath is the URL version of path.Clean, it returns a canonical URL path
// for p, eliminating . and .. elements.
//
// The following rules are applied iteratively until no further processing can
// be done:
//	1. Replace multiple slashes with a single slash.
//	2. Eliminate each . path name element (the current directory).
//	3. Eliminate each inner .. path name element (the parent directory)
//	   along with the non-.. element that precedes it.
//	4. Eliminate .. elements that begin a rooted path:
//	   that is, replace "/.." by "/" at the beginning of a path.
//
// If the result of this process is an empty string, "/" is returned
func cleanPath(p string) string {
	// Turn empty string into "/"
	if p == "" {
		return "/"
	}

	n := len(p)
	var buf []byte

	// Invariants:
	//      reading from path; r is index of next byte to process.
	//      writing to buf; w is index of next byte to write.

	// path must start with '/'
	r := 1
	w := 1

	if p[0] != '/' {
		r = 0
		buf = make([]byte, n+1)
		buf[0] = '/'
	}

	trailing := n > 2 && p[n-1] == '/'

	// A bit more clunky without a 'lazybuf' like the path package, but the loop
	// gets completely inlined (bufApp). So in contrast to the path package this
	// loop has no expensive function calls (except 1x make)

	for r < n {
		switch {
		case p[r] == '/':
			// empty path element, trailing slash is added after the end
			r++

		case p[r] == '.' && r+1 == n:
			trailing = true
			r++

		case p[r] == '.' && p[r+1] == '/':
			// . element
			r++

		case p[r] == '.' && p[r+1] == '.' && (r+2 == n || p[r+2] == '/'):
			// .. element: remove to last /
			r += 2

			if w > 1 {
				// can backtrack
				w--

				if buf == nil {
					for w > 1 && p[w] != '/' {
						w--
					}
				} else {
					for w > 1 && buf[w] != '/' {
						w--
					}
				}
			}

		default:
			// real path element.
			// add slash if needed
			if w > 1 {
				bufApp(&buf, p, w, '/')
				w++
			}

			// copy element
			for r < n && p[r] != '/' {
				bufApp(&buf, p, w, p[r])
				w++
				r++
			}
		}
	}

	// re-append trailing slash
	if trailing && w > 1 {
		bufApp(&buf, p, w, '/')
		w++
	}

	if buf == nil {
		return p[:w]
	}
	return string(buf[:w])
}

// internal helper to lazily create a buffer if necessary
func bufApp(buf *[]byte, s string, w int, c byte) {
	if *buf == nil {
		if s[w] == c {
			return
		}

		*buf = make([]byte, len(s))
		copy(*buf, s[:w])
	}
	(*buf)[w] = c
}
<|endoftext|>


<commit_before>// Copyright (c) 2014-2015 by Michael Dvorkin. All Rights Reserved.
// Use of this source code is governed by a MIT-style license that can
// be found in the LICENSE file.

// This space is available for rent.
package main

import (
	`github.com/michaeldv/donna`
	`os`
	`runtime`
)

// Ignore previous comment.
func main() {
	// Default engine settings are: 128MB transposition table, 5s per move.
	engine := donna.NewEngine(
		`fancy`, runtime.GOOS == `darwin`,
		`cache`, 128,
		`movetime`, 5000,
		`logfile`, os.Getenv(`DONNA_LOG`),
		`bookfile`, os.Getenv(`DONNA_BOOK`),
	)

	if len(os.Args) > 1 && os.Args[1] == `-i` {
		engine.Repl()
	} else {
		engine.Uci()
	}
}
<commit_msg>Bumped up default cache size to 256MB<commit_after>// Copyright (c) 2014-2015 by Michael Dvorkin. All Rights Reserved.
// Use of this source code is governed by a MIT-style license that can
// be found in the LICENSE file.

// This space is available for rent.
package main

import (
	`github.com/michaeldv/donna`
	`os`
	`runtime`
)

// Ignore previous comment.
func main() {
	// Default engine settings are: 128MB transposition table, 5s per move.
	engine := donna.NewEngine(
		`fancy`, runtime.GOOS == `darwin`,
		`cache`, 256,
		`movetime`, 5000,
		`logfile`, os.Getenv(`DONNA_LOG`),
		`bookfile`, os.Getenv(`DONNA_BOOK`),
	)

	if len(os.Args) > 1 && os.Args[1] == `-i` {
		engine.Repl()
	} else {
		engine.Uci()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
)

const FILE string = "Makefile"
const PROG string = "make"

/* TODO:
add stopping at homedir
configurable filename for aliases?
*/
func main() {
	checkDir, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	for {
		//fmt.Println("Checking:", checkDir)
		if existsAtPath(checkDir) {
			//fmt.Println("FOUND IT in", checkDir)
			os.Exit(runAt(checkDir))
		} else {
			newdir := filepath.Dir(checkDir)
			//fmt.Println("Moving to:", newdir)
			checkDir = newdir
		}
	}
}

func runAt(dir string) int {
	cmd := exec.Command(PROG, os.Args[1:]...)
	cmd.Dir = dir
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err := cmd.Run()
	if err != nil {
		fmt.Printf("exit with errr %v\n", err)
		return 1
	} else {
		return 0
	}
}

func existsAtPath(dir string) bool {
	path := filepath.Join(dir, FILE)
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return false
	}
	//fmt.Println("found:", path)
	return true
}
<commit_msg>Fix infinite loop at root dir if nothing is found<commit_after>package main

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
)

const FILE string = "Makefile"
const PROG string = "make"

/* TODO:
add stopping at homedir
configurable filename for aliases?
*/
func main() {
	checkDir, err := os.Getwd()
	if err != nil {
		fmt.Println("Error getting working directory:", err)
		os.Exit(1)
	}
	for {
		if existsAtPath(checkDir) {
			os.Exit(runAt(checkDir))
		} else if checkDir == "/" {
			fmt.Println("Unable to find", FILE)
			os.Exit(1)
		} else {
			newdir := filepath.Dir(checkDir)
			checkDir = newdir
		}
	}
}

func runAt(dir string) int {
	cmd := exec.Command(PROG, os.Args[1:]...)
	cmd.Dir = dir
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err := cmd.Run()
	if err != nil {
		fmt.Println(PROG, "exited with error", err)
		return 1
	} else {
		return 0
	}
}

func existsAtPath(dir string) bool {
	path := filepath.Join(dir, FILE)
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package cmd

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"github.com/fsouza/go-dockerclient"
	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/spf13/cobra"
)

type Emet struct {
	Id      string
	BaseUrl string
}

func EmetHandler(w http.ResponseWriter, r *http.Request) {

	tmpl, err := template.ParseFiles("emet.tmpl.js")
	if err != nil {
		http.Error(w, err.Error(), 500)
	}

	emet := Emet{BaseUrl: "localhost"}

	err = tmpl.Execute(w, emet)
}

func HomeHandler(w http.ResponseWriter, r *http.Request) {

}
func ListHandler(w http.ResponseWriter, r *http.Request) {

	client, err := docker.NewClientFromEnv()
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	containers, err := client.ListContainers(docker.ListContainersOptions{All: false})
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	golems := []docker.APIContainers{}
	for _, container := range containers {
		if container.Image == "golem" {
			golems = append(golems, container)
		}
	}

	data, err := json.Marshal(golems)
	if err != nil {
		http.Error(w, err.Error(), 500)
	}

	w.Write(data)

	for _, container := range golems {
		fmt.Println(container.ID)
	}
}
func SpawnHandler(w http.ResponseWriter, r *http.Request) {

	// we need id for webstrate

	client, err := docker.NewClientFromEnv()
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}

	// TODO figure out port
	// TODO if container is already running then return

	//w.Write([]byte(vars["id"]))

	// TODO use id

	vars := mux.Vars(r)
	wsid := vars["id"]

	//ctx := context.Background()

	fmt.Println("Pulling image")
	err = client.PullImage(docker.PullImageOptions{
		Repository: "webstrates/golem",
		Tag:        "latest",
	}, docker.AuthConfiguration{})
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	fmt.Println("Pull done")

	// Get current dir
	dir, err := os.Getwd()
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	fmt.Println("dir is " + dir)

	seccomp, err := ioutil.ReadFile(filepath.Join(dir, "chrome.json"))
	if err != nil {
		fmt.Println("Error reading seccomp" + err.Error())
		http.Error(w, err.Error(), 500)
		return
	}

	fmt.Println("Creating container")
	container, err := client.CreateContainer(
		docker.CreateContainerOptions{
			Name: fmt.Sprintf("golem-%s", wsid),
			Config: &docker.Config{
				Image: "webstrates/golem:latest",
				ExposedPorts: map[docker.Port]struct{}{
					"9222/tcp": {},
				},
			},
			HostConfig: &docker.HostConfig{
				Links: []string{"webstrates"},
				PortBindings: map[docker.Port][]docker.PortBinding{
					"9222/tcp": []docker.PortBinding{
						docker.PortBinding{
							HostIP:   "0.0.0.0",
							HostPort: "9222", // TODO make this dynamic
						},
					},
				},
				SecurityOpt: []string{
					fmt.Sprintf("seccomp=%s", string(seccomp)),
				},
			},
		},
	)
	if err != nil {
		fmt.Println("Error creating container" + err.Error())
		http.Error(w, err.Error(), 500)
		return
	}
	fmt.Println("Created container")

	fmt.Println(container.ID)
	fmt.Printf("seccomp=%s", filepath.Join(dir, "chrome.json"))

	fmt.Println("Starting container")
	err = client.StartContainer(container.ID, nil)

	if err != nil {
		fmt.Println("Error starting container" + err.Error())
		http.Error(w, err.Error(), 500)
		return
	}
	fmt.Println("Started container")
	// TODO return json

}
func ResetHandler(w http.ResponseWriter, r *http.Request) {
}
func KillHandler(w http.ResponseWriter, r *http.Request) {
}

// serveCmd represents the serve command
var serveCmd = &cobra.Command{
	Use:   "serve",
	Short: "Start a remote administration interface",
	Run: func(cmd *cobra.Command, args []string) {
		r := mux.NewRouter()

		r.HandleFunc("/", HomeHandler)
		r.HandleFunc("/emet", EmetHandler)
		r.HandleFunc("/ls", ListHandler)
		r.HandleFunc("/spawn/{id}", SpawnHandler)
		r.HandleFunc("/reset/{id}", ResetHandler)
		r.HandleFunc("/kill/{id}", KillHandler)

		srv := &http.Server{
			Handler:   handlers.CORS()(r),
			Addr:      ":8000",
			TLSConfig: &tls.Config{},
			// Good practice: enforce timeouts for servers you create!
			WriteTimeout: 15 * time.Second,
			ReadTimeout:  15 * time.Second,
		}

		log.Fatal(srv.ListenAndServeTLS("server.crt", "server.key"))
	},
}

func init() {
	RootCmd.AddCommand(serveCmd)

	// Here you will define your flags and configuration settings.

	// Cobra supports Persistent Flags which will work for this command
	// and all subcommands, e.g.:
	// serveCmd.PersistentFlags().String("foo", "", "A help for foo")

	// Cobra supports local flags which will only run when this command
	// is called directly, e.g.:
	// serveCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")

}
<commit_msg>Better logging<commit_after>package cmd

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"time"

	log "github.com/sirupsen/logrus"

	"github.com/fsouza/go-dockerclient"
	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/spf13/cobra"
)

type Emet struct {
	Id      string
	BaseUrl string
}

func EmetHandler(w http.ResponseWriter, r *http.Request) {

	tmpl, err := template.ParseFiles("emet.tmpl.js")
	if err != nil {
		http.Error(w, err.Error(), 500)
	}

	emet := Emet{BaseUrl: "localhost"}

	err = tmpl.Execute(w, emet)
}

func HomeHandler(w http.ResponseWriter, r *http.Request) {

}
func ListHandler(w http.ResponseWriter, r *http.Request) {

	client, err := docker.NewClientFromEnv()
	if err != nil {
		log.WithError(err).Error("Error creating docker client")
		http.Error(w, err.Error(), 500)
		return
	}

	containers, err := client.ListContainers(docker.ListContainersOptions{All: false})
	if err != nil {
		log.WithError(err).Error("Error listing containers")
		http.Error(w, err.Error(), 500)
		return
	}

	golems := []docker.APIContainers{}
	for _, container := range containers {
		if container.Image == "golem" {
			golems = append(golems, container)
		}
	}

	data, err := json.Marshal(golems)
	if err != nil {
		http.Error(w, err.Error(), 500)
	}

	w.Write(data)
}
func SpawnHandler(w http.ResponseWriter, r *http.Request) {

	// we need id for webstrate

	client, err := docker.NewClientFromEnv()
	if err != nil {
		log.WithError(err).Error("Could create docker client")
		http.Error(w, err.Error(), 500)
		return
	}

	// TODO figure out port
	// TODO if container is already running then return

	//w.Write([]byte(vars["id"]))

	// TODO use id

	vars := mux.Vars(r)
	wsid := vars["id"]

	//ctx := context.Background()

	repository := "webstrates/golem"
	tag := "latest"

	log.WithFields(log.Fields{"image": fmt.Sprintf("%s:%s", repository, tag)}).Info("Pulling image")

	err = client.PullImage(docker.PullImageOptions{
		Repository: "webstrates/golem",
		Tag:        "latest",
	}, docker.AuthConfiguration{})
	if err != nil {
		http.Error(w, err.Error(), 500)
		return
	}
	log.WithFields(log.Fields{"image": fmt.Sprintf("%s:%s", repository, tag)}).Info("Pull done")

	// Get current dir
	dir, err := os.Getwd()
	if err != nil {
		log.WithError(err).Error("Could not discover current directory")
		http.Error(w, err.Error(), 500)
		return
	}

	seccomp, err := ioutil.ReadFile(filepath.Join(dir, "chrome.json"))
	if err != nil {
		log.WithError(err).Error("Could not read seccomp profile")
		http.Error(w, err.Error(), 500)
		return
	}

	// TODO add environment variables (WEBSTRATEID)
	log.WithFields(log.Fields{"webstrateid": wsid}).Info("Creating container")
	container, err := client.CreateContainer(
		docker.CreateContainerOptions{
			Name: fmt.Sprintf("golem-%s", wsid),
			Config: &docker.Config{
				Image: "webstrates/golem:latest",
				ExposedPorts: map[docker.Port]struct{}{
					"9222/tcp": {},
				},
			},
			HostConfig: &docker.HostConfig{
				Links: []string{"webstrates"},
				PortBindings: map[docker.Port][]docker.PortBinding{
					"9222/tcp": []docker.PortBinding{
						docker.PortBinding{
							HostIP:   "0.0.0.0",
							HostPort: "9222", // TODO make this dynamic
						},
					},
				},
				SecurityOpt: []string{
					fmt.Sprintf("seccomp=%s", string(seccomp)),
				},
			},
		},
	)
	if err != nil {
		log.WithError(err).Error("Error creating container")
		http.Error(w, err.Error(), 500)
		return
	}
	log.WithFields(log.Fields{"webstrateid": wsid, "containerid": container.ID}).Info("Created container, starting ...")

	err = client.StartContainer(container.ID, nil)

	if err != nil {
		log.WithError(err).Error("Error starting container")
		http.Error(w, err.Error(), 500)
		return
	}

	// TODO return something

}
func ResetHandler(w http.ResponseWriter, r *http.Request) {
	// TODO reset handler
}
func KillHandler(w http.ResponseWriter, r *http.Request) {
	// TODO kill handler
}

// serveCmd represents the serve command
var serveCmd = &cobra.Command{
	Use:   "serve",
	Short: "Start a remote administration interface",
	Run: func(cmd *cobra.Command, args []string) {
		r := mux.NewRouter()

		r.HandleFunc("/", HomeHandler)
		r.HandleFunc("/emet", EmetHandler)
		r.HandleFunc("/ls", ListHandler)
		r.HandleFunc("/spawn/{id}", SpawnHandler)
		r.HandleFunc("/reset/{id}", ResetHandler)
		r.HandleFunc("/kill/{id}", KillHandler)

		srv := &http.Server{
			Handler:   handlers.CORS()(r),
			Addr:      ":8000",
			TLSConfig: &tls.Config{},
			// Good practice: enforce timeouts for servers you create!
			WriteTimeout: 15 * time.Second,
			ReadTimeout:  15 * time.Second,
		}

		log.Fatal(srv.ListenAndServeTLS("server.crt", "server.key"))
	},
}

func init() {
	RootCmd.AddCommand(serveCmd)

	// Here you will define your flags and configuration settings.

	// Cobra supports Persistent Flags which will work for this command
	// and all subcommands, e.g.:
	// serveCmd.PersistentFlags().String("foo", "", "A help for foo")

	// Cobra supports local flags which will only run when this command
	// is called directly, e.g.:
	// serveCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")

}
<|endoftext|>


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

package cmd

import (
	"crypto/tls"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"code.gitea.io/git"
	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/base"
	"code.gitea.io/gitea/modules/httplib"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"
	"github.com/Unknwon/com"
	gouuid "github.com/satori/go.uuid"
	"github.com/urfave/cli"
)

const (
	accessDenied = "Repository does not exist or you do not have access"
)

// CmdServ represents the available serv sub-command.
var CmdServ = cli.Command{
	Name:        "serv",
	Usage:       "This command should only be called by SSH shell",
	Description: `Serv provide access auth for repositories`,
	Action:      runServ,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "config, c",
			Value: "custom/conf/app.ini",
			Usage: "Custom configuration file path",
		},
	},
}

func setup(logPath string) {
	setting.NewContext()
	log.NewGitLogger(filepath.Join(setting.LogRootPath, logPath))

	models.LoadConfigs()

	if setting.UseSQLite3 || setting.UseTiDB {
		workDir, _ := setting.WorkDir()
		if err := os.Chdir(workDir); err != nil {
			log.GitLogger.Fatal(4, "Fail to change directory %s: %v", workDir, err)
		}
	}

	models.SetEngine()
}

func parseCmd(cmd string) (string, string) {
	ss := strings.SplitN(cmd, " ", 2)
	if len(ss) != 2 {
		return "", ""
	}
	return ss[0], strings.Replace(ss[1], "'/", "'", 1)
}

var (
	allowedCommands = map[string]models.AccessMode{
		"git-upload-pack":    models.AccessModeRead,
		"git-upload-archive": models.AccessModeRead,
		"git-receive-pack":   models.AccessModeWrite,
	}
)

func fail(userMessage, logMessage string, args ...interface{}) {
	fmt.Fprintln(os.Stderr, "Gogs:", userMessage)

	if len(logMessage) > 0 {
		if !setting.ProdMode {
			fmt.Fprintf(os.Stderr, logMessage+"\n", args...)
		}
		log.GitLogger.Fatal(3, logMessage, args...)
		return
	}

	log.GitLogger.Close()
	os.Exit(1)
}

func handleUpdateTask(uuid string, user, repoUser *models.User, reponame string, isWiki bool) {
	task, err := models.GetUpdateTaskByUUID(uuid)
	if err != nil {
		if models.IsErrUpdateTaskNotExist(err) {
			log.GitLogger.Trace("No update task is presented: %s", uuid)
			return
		}
		log.GitLogger.Fatal(2, "GetUpdateTaskByUUID: %v", err)
	} else if err = models.DeleteUpdateTaskByUUID(uuid); err != nil {
		log.GitLogger.Fatal(2, "DeleteUpdateTaskByUUID: %v", err)
	}

	if isWiki {
		return
	}

	if err = models.PushUpdate(models.PushUpdateOptions{
		RefFullName:  task.RefName,
		OldCommitID:  task.OldCommitID,
		NewCommitID:  task.NewCommitID,
		PusherID:     user.ID,
		PusherName:   user.Name,
		RepoUserName: repoUser.Name,
		RepoName:     reponame,
	}); err != nil {
		log.GitLogger.Error(2, "Update: %v", err)
	}

	// Ask for running deliver hook and test pull request tasks.
	reqURL := setting.LocalURL + repoUser.Name + "/" + reponame + "/tasks/trigger?branch=" +
		strings.TrimPrefix(task.RefName, git.BRANCH_PREFIX) + "&secret=" + base.EncodeMD5(repoUser.Salt) + "&pusher=" + com.ToStr(user.ID)
	log.GitLogger.Trace("Trigger task: %s", reqURL)

	resp, err := httplib.Head(reqURL).SetTLSClientConfig(&tls.Config{
		InsecureSkipVerify: true,
	}).Response()
	if err == nil {
		resp.Body.Close()
		if resp.StatusCode/100 != 2 {
			log.GitLogger.Error(2, "Fail to trigger task: not 2xx response code")
		}
	} else {
		log.GitLogger.Error(2, "Fail to trigger task: %v", err)
	}
}

func runServ(c *cli.Context) error {
	if c.IsSet("config") {
		setting.CustomConf = c.String("config")
	}

	setup("serv.log")

	if setting.SSH.Disabled {
		println("Gogs: SSH has been disabled")
		return nil
	}

	if len(c.Args()) < 1 {
		fail("Not enough arguments", "Not enough arguments")
	}

	cmd := os.Getenv("SSH_ORIGINAL_COMMAND")
	if len(cmd) == 0 {
		println("Hi there, You've successfully authenticated, but Gogs does not provide shell access.")
		println("If this is unexpected, please log in with password and setup Gogs under another user.")
		return nil
	}

	verb, args := parseCmd(cmd)
	repoPath := strings.ToLower(strings.Trim(args, "'"))
	rr := strings.SplitN(repoPath, "/", 2)
	if len(rr) != 2 {
		fail("Invalid repository path", "Invalid repository path: %v", args)
	}
	username := strings.ToLower(rr[0])
	reponame := strings.ToLower(strings.TrimSuffix(rr[1], ".git"))

	isWiki := false
	if strings.HasSuffix(reponame, ".wiki") {
		isWiki = true
		reponame = reponame[:len(reponame)-5]
	}

	repoUser, err := models.GetUserByName(username)
	if err != nil {
		if models.IsErrUserNotExist(err) {
			fail("Repository owner does not exist", "Unregistered owner: %s", username)
		}
		fail("Internal error", "Failed to get repository owner (%s): %v", username, err)
	}

	repo, err := models.GetRepositoryByName(repoUser.ID, reponame)
	if err != nil {
		if models.IsErrRepoNotExist(err) {
			fail(accessDenied, "Repository does not exist: %s/%s", repoUser.Name, reponame)
		}
		fail("Internal error", "Failed to get repository: %v", err)
	}

	requestedMode, has := allowedCommands[verb]
	if !has {
		fail("Unknown git command", "Unknown git command %s", verb)
	}

	// Prohibit push to mirror repositories.
	if requestedMode > models.AccessModeRead && repo.IsMirror {
		fail("mirror repository is read-only", "")
	}

	// Allow anonymous clone for public repositories.
	var (
		keyID int64
		user  *models.User
	)
	if requestedMode == models.AccessModeWrite || repo.IsPrivate {
		keys := strings.Split(c.Args()[0], "-")
		if len(keys) != 2 {
			fail("Key ID format error", "Invalid key argument: %s", c.Args()[0])
		}

		key, err := models.GetPublicKeyByID(com.StrTo(keys[1]).MustInt64())
		if err != nil {
			fail("Invalid key ID", "Invalid key ID[%s]: %v", c.Args()[0], err)
		}
		keyID = key.ID

		// Check deploy key or user key.
		if key.Type == models.KeyTypeDeploy {
			if key.Mode < requestedMode {
				fail("Key permission denied", "Cannot push with deployment key: %d", key.ID)
			}
			// Check if this deploy key belongs to current repository.
			if !models.HasDeployKey(key.ID, repo.ID) {
				fail("Key access denied", "Deploy key access denied: [key_id: %d, repo_id: %d]", key.ID, repo.ID)
			}

			// Update deploy key activity.
			deployKey, err := models.GetDeployKeyByRepo(key.ID, repo.ID)
			if err != nil {
				fail("Internal error", "GetDeployKey: %v", err)
			}

			deployKey.Updated = time.Now()
			if err = models.UpdateDeployKey(deployKey); err != nil {
				fail("Internal error", "UpdateDeployKey: %v", err)
			}
		} else {
			user, err = models.GetUserByKeyID(key.ID)
			if err != nil {
				fail("internal error", "Failed to get user by key ID(%d): %v", keyID, err)
			}

			mode, err := models.AccessLevel(user, repo)
			if err != nil {
				fail("Internal error", "Fail to check access: %v", err)
			} else if mode < requestedMode {
				clientMessage := accessDenied
				if mode >= models.AccessModeRead {
					clientMessage = "You do not have sufficient authorization for this action"
				}
				fail(clientMessage,
					"User %s does not have level %v access to repository %s",
					user.Name, requestedMode, repoPath)
			}
		}
	}

	os.Setenv("GITEA_PUSHER_NAME", user.Name)

	uuid := gouuid.NewV4().String()
	os.Setenv("GITEA_UUID", uuid)
	// Keep the old env variable name for backward compability
	os.Setenv("uuid", uuid)

	// Special handle for Windows.
	if setting.IsWindows {
		verb = strings.Replace(verb, "-", " ", 1)
	}

	var gitcmd *exec.Cmd
	verbs := strings.Split(verb, " ")
	if len(verbs) == 2 {
		gitcmd = exec.Command(verbs[0], verbs[1], repoPath)
	} else {
		gitcmd = exec.Command(verb, repoPath)
	}
	gitcmd.Dir = setting.RepoRootPath
	gitcmd.Stdout = os.Stdout
	gitcmd.Stdin = os.Stdin
	gitcmd.Stderr = os.Stderr
	if err = gitcmd.Run(); err != nil {
		fail("Internal error", "Failed to execute git command: %v", err)
	}

	if requestedMode == models.AccessModeWrite {
		handleUpdateTask(uuid, user, repoUser, reponame, isWiki)
	}

	// Update user key activity.
	if keyID > 0 {
		key, err := models.GetPublicKeyByID(keyID)
		if err != nil {
			fail("Internal error", "GetPublicKeyById: %v", err)
		}

		key.Updated = time.Now()
		if err = models.UpdatePublicKey(key); err != nil {
			fail("Internal error", "UpdatePublicKey: %v", err)
		}
	}

	return nil
}
<commit_msg>Fixes panic when there's no user initialized (#358)<commit_after>// Copyright 2014 The Gogs Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package cmd

import (
	"crypto/tls"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"code.gitea.io/git"
	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/base"
	"code.gitea.io/gitea/modules/httplib"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"
	"github.com/Unknwon/com"
	gouuid "github.com/satori/go.uuid"
	"github.com/urfave/cli"
)

const (
	accessDenied = "Repository does not exist or you do not have access"
)

// CmdServ represents the available serv sub-command.
var CmdServ = cli.Command{
	Name:        "serv",
	Usage:       "This command should only be called by SSH shell",
	Description: `Serv provide access auth for repositories`,
	Action:      runServ,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "config, c",
			Value: "custom/conf/app.ini",
			Usage: "Custom configuration file path",
		},
	},
}

func setup(logPath string) {
	setting.NewContext()
	log.NewGitLogger(filepath.Join(setting.LogRootPath, logPath))

	models.LoadConfigs()

	if setting.UseSQLite3 || setting.UseTiDB {
		workDir, _ := setting.WorkDir()
		if err := os.Chdir(workDir); err != nil {
			log.GitLogger.Fatal(4, "Fail to change directory %s: %v", workDir, err)
		}
	}

	models.SetEngine()
}

func parseCmd(cmd string) (string, string) {
	ss := strings.SplitN(cmd, " ", 2)
	if len(ss) != 2 {
		return "", ""
	}
	return ss[0], strings.Replace(ss[1], "'/", "'", 1)
}

var (
	allowedCommands = map[string]models.AccessMode{
		"git-upload-pack":    models.AccessModeRead,
		"git-upload-archive": models.AccessModeRead,
		"git-receive-pack":   models.AccessModeWrite,
	}
)

func fail(userMessage, logMessage string, args ...interface{}) {
	fmt.Fprintln(os.Stderr, "Gogs:", userMessage)

	if len(logMessage) > 0 {
		if !setting.ProdMode {
			fmt.Fprintf(os.Stderr, logMessage+"\n", args...)
		}
		log.GitLogger.Fatal(3, logMessage, args...)
		return
	}

	log.GitLogger.Close()
	os.Exit(1)
}

func handleUpdateTask(uuid string, user, repoUser *models.User, reponame string, isWiki bool) {
	task, err := models.GetUpdateTaskByUUID(uuid)
	if err != nil {
		if models.IsErrUpdateTaskNotExist(err) {
			log.GitLogger.Trace("No update task is presented: %s", uuid)
			return
		}
		log.GitLogger.Fatal(2, "GetUpdateTaskByUUID: %v", err)
	} else if err = models.DeleteUpdateTaskByUUID(uuid); err != nil {
		log.GitLogger.Fatal(2, "DeleteUpdateTaskByUUID: %v", err)
	}

	if isWiki {
		return
	}

	if err = models.PushUpdate(models.PushUpdateOptions{
		RefFullName:  task.RefName,
		OldCommitID:  task.OldCommitID,
		NewCommitID:  task.NewCommitID,
		PusherID:     user.ID,
		PusherName:   user.Name,
		RepoUserName: repoUser.Name,
		RepoName:     reponame,
	}); err != nil {
		log.GitLogger.Error(2, "Update: %v", err)
	}

	// Ask for running deliver hook and test pull request tasks.
	reqURL := setting.LocalURL + repoUser.Name + "/" + reponame + "/tasks/trigger?branch=" +
		strings.TrimPrefix(task.RefName, git.BRANCH_PREFIX) + "&secret=" + base.EncodeMD5(repoUser.Salt) + "&pusher=" + com.ToStr(user.ID)
	log.GitLogger.Trace("Trigger task: %s", reqURL)

	resp, err := httplib.Head(reqURL).SetTLSClientConfig(&tls.Config{
		InsecureSkipVerify: true,
	}).Response()
	if err == nil {
		resp.Body.Close()
		if resp.StatusCode/100 != 2 {
			log.GitLogger.Error(2, "Fail to trigger task: not 2xx response code")
		}
	} else {
		log.GitLogger.Error(2, "Fail to trigger task: %v", err)
	}
}

func runServ(c *cli.Context) error {
	if c.IsSet("config") {
		setting.CustomConf = c.String("config")
	}

	setup("serv.log")

	if setting.SSH.Disabled {
		println("Gogs: SSH has been disabled")
		return nil
	}

	if len(c.Args()) < 1 {
		fail("Not enough arguments", "Not enough arguments")
	}

	cmd := os.Getenv("SSH_ORIGINAL_COMMAND")
	if len(cmd) == 0 {
		println("Hi there, You've successfully authenticated, but Gogs does not provide shell access.")
		println("If this is unexpected, please log in with password and setup Gogs under another user.")
		return nil
	}

	verb, args := parseCmd(cmd)
	repoPath := strings.ToLower(strings.Trim(args, "'"))
	rr := strings.SplitN(repoPath, "/", 2)
	if len(rr) != 2 {
		fail("Invalid repository path", "Invalid repository path: %v", args)
	}
	username := strings.ToLower(rr[0])
	reponame := strings.ToLower(strings.TrimSuffix(rr[1], ".git"))

	isWiki := false
	if strings.HasSuffix(reponame, ".wiki") {
		isWiki = true
		reponame = reponame[:len(reponame)-5]
	}

	repoUser, err := models.GetUserByName(username)
	if err != nil {
		if models.IsErrUserNotExist(err) {
			fail("Repository owner does not exist", "Unregistered owner: %s", username)
		}
		fail("Internal error", "Failed to get repository owner (%s): %v", username, err)
	}

	repo, err := models.GetRepositoryByName(repoUser.ID, reponame)
	if err != nil {
		if models.IsErrRepoNotExist(err) {
			fail(accessDenied, "Repository does not exist: %s/%s", repoUser.Name, reponame)
		}
		fail("Internal error", "Failed to get repository: %v", err)
	}

	requestedMode, has := allowedCommands[verb]
	if !has {
		fail("Unknown git command", "Unknown git command %s", verb)
	}

	// Prohibit push to mirror repositories.
	if requestedMode > models.AccessModeRead && repo.IsMirror {
		fail("mirror repository is read-only", "")
	}

	// Allow anonymous clone for public repositories.
	var (
		keyID int64
		user  *models.User
	)
	if requestedMode == models.AccessModeWrite || repo.IsPrivate {
		keys := strings.Split(c.Args()[0], "-")
		if len(keys) != 2 {
			fail("Key ID format error", "Invalid key argument: %s", c.Args()[0])
		}

		key, err := models.GetPublicKeyByID(com.StrTo(keys[1]).MustInt64())
		if err != nil {
			fail("Invalid key ID", "Invalid key ID[%s]: %v", c.Args()[0], err)
		}
		keyID = key.ID

		// Check deploy key or user key.
		if key.Type == models.KeyTypeDeploy {
			if key.Mode < requestedMode {
				fail("Key permission denied", "Cannot push with deployment key: %d", key.ID)
			}
			// Check if this deploy key belongs to current repository.
			if !models.HasDeployKey(key.ID, repo.ID) {
				fail("Key access denied", "Deploy key access denied: [key_id: %d, repo_id: %d]", key.ID, repo.ID)
			}

			// Update deploy key activity.
			deployKey, err := models.GetDeployKeyByRepo(key.ID, repo.ID)
			if err != nil {
				fail("Internal error", "GetDeployKey: %v", err)
			}

			deployKey.Updated = time.Now()
			if err = models.UpdateDeployKey(deployKey); err != nil {
				fail("Internal error", "UpdateDeployKey: %v", err)
			}
		} else {
			user, err = models.GetUserByKeyID(key.ID)
			if err != nil {
				fail("internal error", "Failed to get user by key ID(%d): %v", keyID, err)
			}

			mode, err := models.AccessLevel(user, repo)
			if err != nil {
				fail("Internal error", "Fail to check access: %v", err)
			} else if mode < requestedMode {
				clientMessage := accessDenied
				if mode >= models.AccessModeRead {
					clientMessage = "You do not have sufficient authorization for this action"
				}
				fail(clientMessage,
					"User %s does not have level %v access to repository %s",
					user.Name, requestedMode, repoPath)
			}

			os.Setenv("GITEA_PUSHER_NAME", user.Name)
		}
	}

	uuid := gouuid.NewV4().String()
	os.Setenv("GITEA_UUID", uuid)
	// Keep the old env variable name for backward compability
	os.Setenv("uuid", uuid)

	// Special handle for Windows.
	if setting.IsWindows {
		verb = strings.Replace(verb, "-", " ", 1)
	}

	var gitcmd *exec.Cmd
	verbs := strings.Split(verb, " ")
	if len(verbs) == 2 {
		gitcmd = exec.Command(verbs[0], verbs[1], repoPath)
	} else {
		gitcmd = exec.Command(verb, repoPath)
	}
	gitcmd.Dir = setting.RepoRootPath
	gitcmd.Stdout = os.Stdout
	gitcmd.Stdin = os.Stdin
	gitcmd.Stderr = os.Stderr
	if err = gitcmd.Run(); err != nil {
		fail("Internal error", "Failed to execute git command: %v", err)
	}

	if requestedMode == models.AccessModeWrite {
		handleUpdateTask(uuid, user, repoUser, reponame, isWiki)
	}

	// Update user key activity.
	if keyID > 0 {
		key, err := models.GetPublicKeyByID(keyID)
		if err != nil {
			fail("Internal error", "GetPublicKeyById: %v", err)
		}

		key.Updated = time.Now()
		if err = models.UpdatePublicKey(key); err != nil {
			fail("Internal error", "UpdatePublicKey: %v", err)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 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 cmd

import (
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"regexp"
	"strconv"
	"strings"
	"syscall"

	"golang.org/x/crypto/ssh/terminal"
)

var httpHeaderRegexp = regexp.MustCompile(`HTTP/.*? (\d+)`)

type ShellToContainerCmd struct {
	GuessingCommand
}

func (c *ShellToContainerCmd) Info() *Info {
	return &Info{
		Name:  "app-shell",
		Usage: "app-shell [container-id] -a/--app <appname>",
		Desc: `Opens a remote shell inside container, using the API server as a proxy. You
can access an app container just giving app name.

Also, you can access a specific container from this app. In this case, you
have to specify part of the container's ID. You can list current container's
IDs using [[tsuru app-info]].

Open a remote shell to the given container, or to one of the containers of the given app.`,
		MinArgs: 0,
	}
}

func (c *ShellToContainerCmd) Run(context *Context, client *Client) error {
	var width, height int
	if stdin, ok := context.Stdin.(*os.File); ok {
		fd := int(stdin.Fd())
		if terminal.IsTerminal(fd) {
			width, height, _ = terminal.GetSize(fd)
			oldState, err := terminal.MakeRaw(fd)
			if err != nil {
				return err
			}
			defer terminal.Restore(fd, oldState)
			sigChan := make(chan os.Signal, 2)
			go func(c <-chan os.Signal) {
				if _, ok := <-c; ok {
					terminal.Restore(fd, oldState)
					os.Exit(1)
				}
			}(sigChan)
			signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL, syscall.SIGQUIT)
		}
	}
	queryString := make(url.Values)
	queryString.Set("width", strconv.Itoa(width))
	queryString.Set("height", strconv.Itoa(height))
	if len(context.Args) > 0 {
		queryString.Set("container", context.Args[0])
	}
	appName, err := c.Guess()
	if err != nil {
		return err
	}
	serverURL, err := GetURL(fmt.Sprintf("/apps/%s/shell?%s", appName, queryString.Encode()))
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", serverURL, nil)
	if err != nil {
		return err
	}
	request.Close = true
	token, err := ReadToken()
	if err == nil {
		request.Header.Set("Authorization", "bearer "+token)
	}
	parsedURL, _ := url.Parse(serverURL)
	conn, err := net.Dial("tcp", parsedURL.Host)
	if err != nil {
		return err
	}
	defer conn.Close()
	request.Write(conn)
	bytesLimit := 12
	var readStr string
	byteBuffer := make([]byte, 1)
	for i := 0; i < bytesLimit && byteBuffer[0] != '\n'; i++ {
		_, err := conn.Read(byteBuffer)
		if err != nil {
			break
		}
		readStr += string(byteBuffer)
	}
	matches := httpHeaderRegexp.FindAllStringSubmatch(readStr, -1)
	if len(matches) > 0 && len(matches[0]) > 1 {
		return errors.New(strings.TrimSpace(readStr))
	} else {
		context.Stdout.Write([]byte(readStr))
	}
	errs := make(chan error, 2)
	quit := make(chan bool)
	go io.Copy(conn, context.Stdin)
	go func() {
		defer close(quit)
		_, err := io.Copy(context.Stdout, conn)
		if err != nil && err != io.EOF {
			errs <- err
		}
	}()
	<-quit
	close(errs)
	return <-errs
}
<commit_msg>cmd/shell: support portless targets<commit_after>// Copyright 2015 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 cmd

import (
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"regexp"
	"strconv"
	"strings"
	"syscall"

	"golang.org/x/crypto/ssh/terminal"
)

var httpHeaderRegexp = regexp.MustCompile(`HTTP/.*? (\d+)`)

type ShellToContainerCmd struct {
	GuessingCommand
}

func (c *ShellToContainerCmd) Info() *Info {
	return &Info{
		Name:  "app-shell",
		Usage: "app-shell [container-id] -a/--app <appname>",
		Desc: `Opens a remote shell inside container, using the API server as a proxy. You
can access an app container just giving app name.

Also, you can access a specific container from this app. In this case, you
have to specify part of the container's ID. You can list current container's
IDs using [[tsuru app-info]].

Open a remote shell to the given container, or to one of the containers of the given app.`,
		MinArgs: 0,
	}
}

func (c *ShellToContainerCmd) Run(context *Context, client *Client) error {
	var width, height int
	if stdin, ok := context.Stdin.(*os.File); ok {
		fd := int(stdin.Fd())
		if terminal.IsTerminal(fd) {
			width, height, _ = terminal.GetSize(fd)
			oldState, err := terminal.MakeRaw(fd)
			if err != nil {
				return err
			}
			defer terminal.Restore(fd, oldState)
			sigChan := make(chan os.Signal, 2)
			go func(c <-chan os.Signal) {
				if _, ok := <-c; ok {
					terminal.Restore(fd, oldState)
					os.Exit(1)
				}
			}(sigChan)
			signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL, syscall.SIGQUIT)
		}
	}
	queryString := make(url.Values)
	queryString.Set("width", strconv.Itoa(width))
	queryString.Set("height", strconv.Itoa(height))
	if len(context.Args) > 0 {
		queryString.Set("container", context.Args[0])
	}
	appName, err := c.Guess()
	if err != nil {
		return err
	}
	serverURL, err := GetURL(fmt.Sprintf("/apps/%s/shell?%s", appName, queryString.Encode()))
	if err != nil {
		return err
	}
	request, err := http.NewRequest("GET", serverURL, nil)
	if err != nil {
		return err
	}
	request.Close = true
	token, err := ReadToken()
	if err == nil {
		request.Header.Set("Authorization", "bearer "+token)
	}
	parsedURL, _ := url.Parse(serverURL)
	host := parsedURL.Host
	if _, _, err := net.SplitHostPort(host); err != nil {
		port := "80"
		if parsedURL.Scheme == "https" {
			port = "443"
		}
		host += ":" + port
	}
	conn, err := net.Dial("tcp", host)
	if err != nil {
		return err
	}
	defer conn.Close()
	request.Write(conn)
	bytesLimit := 12
	var readStr string
	byteBuffer := make([]byte, 1)
	for i := 0; i < bytesLimit && byteBuffer[0] != '\n'; i++ {
		_, err := conn.Read(byteBuffer)
		if err != nil {
			break
		}
		readStr += string(byteBuffer)
	}
	matches := httpHeaderRegexp.FindAllStringSubmatch(readStr, -1)
	if len(matches) > 0 && len(matches[0]) > 1 {
		return errors.New(strings.TrimSpace(readStr))
	} else {
		context.Stdout.Write([]byte(readStr))
	}
	errs := make(chan error, 2)
	quit := make(chan bool)
	go io.Copy(conn, context.Stdin)
	go func() {
		defer close(quit)
		_, err := io.Copy(context.Stdout, conn)
		if err != nil && err != io.EOF {
			errs <- err
		}
	}()
	<-quit
	close(errs)
	return <-errs
}
<|endoftext|>


<commit_before>package peco

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"os/exec"
	"regexp"
	"sort"
	"strings"
	"unicode"
)

// These are used as keys in the config file
const (
	IgnoreCaseMatch    = "IgnoreCase"
	CaseSensitiveMatch = "CaseSensitive"
	SmartCaseMatch     = "SmartCase"
	RegexpMatch        = "Regexp"
)

var ignoreCaseFlags = []string{"i"}
var defaultFlags = []string{}

type regexpFlags interface {
	flags(string) []string
}
type regexpFlagList []string

func (r regexpFlagList) flags(_ string) []string {
	return []string(r)
}

type regexpFlagFunc func(string) []string

func (r regexpFlagFunc) flags(s string) []string {
	return r(s)
}

func containsUpper(query string) bool {
	for _, c := range query {
		if unicode.IsUpper(c) {
			return true
		}
	}
	return false
}

func regexpFor(q string, flags []string, quotemeta bool) (*regexp.Regexp, error) {
	reTxt := q
	if quotemeta {
		reTxt = regexp.QuoteMeta(q)
	}

	if flags != nil && len(flags) > 0 {
		reTxt = fmt.Sprintf("(?%s)%s", strings.Join(flags, ""), reTxt)
	}

	re, err := regexp.Compile(reTxt)
	if err != nil {
		return nil, err
	}
	return re, nil
}

func queryToRegexps(flags regexpFlags, quotemeta bool, query string) ([]*regexp.Regexp, error) {
	queries := strings.Split(strings.TrimSpace(query), " ")
	regexps := make([]*regexp.Regexp, 0)

	for _, q := range queries {
		re, err := regexpFor(q, flags.flags(query), quotemeta)
		if err != nil {
			return nil, err
		}
		regexps = append(regexps, re)
	}

	return regexps, nil
}

// sort related stuff
type byMatchStart [][]int

func (m byMatchStart) Len() int {
	return len(m)
}

func (m byMatchStart) Swap(i, j int) {
	m[i], m[j] = m[j], m[i]
}

func (m byMatchStart) Less(i, j int) bool {
	if m[i][0] < m[j][0] {
		return true
	}

	if m[i][0] == m[j][0] {
		return m[i][1]-m[i][0] < m[i][1]-m[i][0]
	}

	return false
}
func matchContains(a []int, b []int) bool {
	return a[0] <= b[0] && a[1] >= b[1]
}

func matchOverlaps(a []int, b []int) bool {
	return a[0] <= b[0] && a[1] >= b[0] ||
		a[0] <= b[1] && a[1] >= b[1]
}

func mergeMatches(a []int, b []int) []int {
	ret := make([]int, 2)

	// Note: In practice this should never happen
	// because we're sorting by N[0] before calling
	// this routine, but for completeness' sake...
	if a[0] < b[0] {
		ret[0] = a[0]
	} else {
		ret[0] = b[0]
	}

	if a[1] < b[1] {
		ret[1] = b[1]
	} else {
		ret[1] = a[1]
	}
	return ret
}

// Filter is responsible for the actual "grep" part of peco
type Filter struct {
	*Ctx
}

// Work is the actual work horse that that does the matching
// in a goroutine of its own. It wraps Matcher.Match().
func (f *Filter) Work(cancel chan struct{}, q HubReq) {
	trace("Filter.Work: START\n")
	defer trace("Filter.Work: END\n")
	defer q.Done()

	query := q.DataString()
	if query == "" {
		trace("Filter.Work: Resetting activingLineBuffer")
		f.ResetActiveLineBuffer()
	} else {
		f.rawLineBuffer.cancelCh = cancel
		f.rawLineBuffer.Replay()

		filter := f.Filter().Clone()
		filter.SetQuery(query)
		trace("Running %#v filter using query '%s'", filter, query)

		filter.Accept(f.rawLineBuffer)
		buf := NewRawLineBuffer()
		buf.onEnd = func() { f.SendStatusMsg("") }
		buf.Accept(filter)

		f.SetActiveLineBuffer(buf)
	}

	if ! f.config.StickySelection {
		f.SelectionClear()
	}
}

// Loop keeps watching for incoming queries, and upon receiving
// a query, spawns a goroutine to do the heavy work. It also
// checks for previously running queries, so we can avoid
// running many goroutines doing the grep at the same time
func (f *Filter) Loop() {
	defer f.ReleaseWaitGroup()

	// previous holds a channel that can cancel the previous
	// query. This is used when multiple queries come in succession
	// and the previous query is discarded anyway
	var previous chan struct{}
	for {
		select {
		case <-f.LoopCh():
			return
		case q := <-f.QueryCh():
			if previous != nil {
				// Tell the previous query to stop
				previous <- struct{}{}
			}
			previous = make(chan struct{}, 1)

			f.SendStatusMsg("Running query...")
			go f.Work(previous, q)
		}
	}
}

type QueryFilterer interface {
	Pipeliner
	Cancel()
	Clone() QueryFilterer
	Accept(Pipeliner)
	SetQuery(string)
	String() string
}

type SelectionFilter struct {
	sel *Selection
}

func (sf SelectionFilter) Name() string {
	return "SelectionFilter"
}

type RegexpFilter struct {
	simplePipeline
	compiledQuery []*regexp.Regexp
	flags         regexpFlags
	quotemeta     bool
	query         string
	name          string
	onEnd         func()
}

func NewRegexpFilter() *RegexpFilter {
	return &RegexpFilter{}
}

func (rf RegexpFilter) Clone() QueryFilterer {
	return &RegexpFilter{
		simplePipeline{},
		nil,
		rf.flags,
		rf.quotemeta,
		rf.query,
		rf.name,
		nil,
	}
}

func (rf *RegexpFilter) Accept(p Pipeliner) {
	cancelCh, incomingCh := p.Pipeline()
	rf.cancelCh = cancelCh
	rf.outputCh = make(chan Line)
	go acceptPipeline(cancelCh, incomingCh, rf.outputCh,
		&pipelineCtx{rf.filter, rf.onEnd})
}

var ErrFilterDidNotMatch = errors.New("error: filter did not match against given line")

func (rf *RegexpFilter) filter(l Line) (Line, error) {
	trace("RegexpFilter.filter: START")
	defer trace("RegexpFilter.filter: END")
	regexps, err := rf.getQueryAsRegexps()
	if err != nil {
		return nil, err
	}
	v := l.DisplayString()
	allMatched := true
	matches := [][]int{}
TryRegexps:
	for _, rx := range regexps {
		trace("RegexpFilter.filter: matching '%s' against '%s'", v, rx)
		match := rx.FindAllStringSubmatchIndex(v, -1)
		if match == nil {
			allMatched = false
			break TryRegexps
		}
		matches = append(matches, match...)
	}

	if !allMatched {
		return nil, ErrFilterDidNotMatch
	}

	trace("RegexpFilter.filter: line matched pattern\n")
	sort.Sort(byMatchStart(matches))

	// We need to "dedupe" the results. For example, if we matched the
	// same region twice, we don't want that to be drawn

	deduped := make([][]int, 0, len(matches))

	for i, m := range matches {
		// Always push the first one
		if i == 0 {
			deduped = append(deduped, m)
			continue
		}

		prev := deduped[len(deduped)-1]
		switch {
		case matchContains(prev, m):
			// If the previous match contains this one, then
			// don't do anything
			continue
		case matchOverlaps(prev, m):
			// If the previous match overlaps with this one,
			// merge the results and make it a bigger one
			deduped[len(deduped)-1] = mergeMatches(prev, m)
		default:
			deduped = append(deduped, m)
		}
	}
	return NewMatchedLine(l, deduped), nil
}

func (rf *RegexpFilter) getQueryAsRegexps() ([]*regexp.Regexp, error) {
	if q := rf.compiledQuery; q != nil {
		return q, nil
	}
	q, err := queryToRegexps(rf.flags, rf.quotemeta, rf.query)
	if err != nil {
		return nil, err
	}

	rf.compiledQuery = q
	return q, nil
}

func (rf *RegexpFilter) SetQuery(q string) {
	rf.query = q
	rf.compiledQuery = nil
}

func (rf RegexpFilter) String() string {
	return rf.name
}

type FilterSet struct {
	filters []QueryFilterer
	current int
}

func (fs *FilterSet) Size() int {
	return len(fs.filters)
}

func (fs *FilterSet) Add(qf QueryFilterer) error {
	fs.filters = append(fs.filters, qf)
	return nil
}

func (fs *FilterSet) Rotate() {
	fs.current++
	if fs.current >= len(fs.filters) {
		fs.current = 0
	}
	trace("FilterSet.Rotate: now filter in effect is %s", fs.filters[fs.current])
}

var ErrFilterNotFound = errors.New("specified filter was not found")

func (fs *FilterSet) SetCurrentByName(name string) error {
	for i, f := range fs.filters {
		if f.String() == name {
			fs.current = i
			return nil
		}
	}
	return ErrFilterNotFound
}

func (fs *FilterSet) GetCurrent() QueryFilterer {
	return fs.filters[fs.current]
}

func NewIgnoreCaseFilter() *RegexpFilter {
	return &RegexpFilter{
		flags:     regexpFlagList(ignoreCaseFlags),
		quotemeta: true,
		name:      "IgnoreCase",
	}
}

func NewCaseSensitiveFilter() *RegexpFilter {
	return &RegexpFilter{
		flags:     regexpFlagList(defaultFlags),
		quotemeta: true,
		name:      "CaseSensitive",
	}
}

// SmartCaseFilter turns ON the ignore-case flag in the regexp
// if the query contains a upper-case character
func NewSmartCaseFilter() *RegexpFilter {
	return &RegexpFilter{
		flags: regexpFlagFunc(func(q string) []string {
			if containsUpper(q) {
				return defaultFlags
			}
			return []string{"i"}
		}),
		quotemeta: true,
		name:      "SmartCase",
	}
}

type ExternalCmdFilter struct {
	simplePipeline
	enableSep       bool
	cmd             string
	args            []string
	name            string
	query           string
	thresholdBufsiz int
}

func NewExternalCmdFilter(name, cmd string, args []string, threshold int, enableSep bool) *ExternalCmdFilter {
	trace("name = %s, cmd = %s, args = %#v", name, cmd, args)
	return &ExternalCmdFilter{
		simplePipeline:  simplePipeline{},
		enableSep:       enableSep,
		cmd:             cmd,
		args:            args,
		name:            name,
		thresholdBufsiz: threshold,
	}
}

func (ecf ExternalCmdFilter) Clone() QueryFilterer {
	return &ExternalCmdFilter{
		simplePipeline:  simplePipeline{},
		enableSep:       ecf.enableSep,
		cmd:             ecf.cmd,
		args:            ecf.args,
		name:            ecf.name,
		thresholdBufsiz: ecf.thresholdBufsiz,
	}
}

func (ecf *ExternalCmdFilter) Verify() error {
	if ecf.cmd == "" {
		return fmt.Errorf("no executable specified for custom matcher '%s'", ecf.name)
	}

	if _, err := exec.LookPath(ecf.cmd); err != nil {
		return err
	}
	return nil
}

func (ecf *ExternalCmdFilter) Accept(p Pipeliner) {
	cancelCh, incomingCh := p.Pipeline()
	outputCh := make(chan Line)
	ecf.cancelCh = cancelCh
	ecf.outputCh = outputCh

	go func() {
		defer close(outputCh)

		defer trace("ExternalCmdFilter.Accept: DONE")

		// for every N lines, execute the external command
		buf := []Line{}
		for l := range incomingCh {
			buf = append(buf, l)
			if len(buf) < ecf.thresholdBufsiz {
				continue
			}

			ecf.launchExternalCmd(buf, cancelCh, outputCh)
			buf = []Line{} // drain
		}

		if len(buf) > 0 {
			ecf.launchExternalCmd(buf, cancelCh, outputCh)
		}
	}()
}

func (ecf *ExternalCmdFilter) SetQuery(q string) {
	ecf.query = q
}

func (ecf ExternalCmdFilter) String() string {
	return ecf.name
}

func (ecf *ExternalCmdFilter) launchExternalCmd(buf []Line, cancelCh chan struct{}, outputCh chan Line) {
	defer func() { recover() }() // ignore errors

	trace("ExternalCmdFilter.launchExternalCmd: START")
	defer trace("ExternalCmdFilter.launchExternalCmd: END")

	trace("buf = %v", buf)

	args := append([]string{ecf.query}, ecf.args...)
	cmd := exec.Command(ecf.cmd, args...)

	inbuf := &bytes.Buffer{}
	for _, l := range buf {
		inbuf.WriteString(l.DisplayString() + "\n")
	}

	cmd.Stdin = inbuf
	r, err := cmd.StdoutPipe()
	if err != nil {
		return
	}

	trace("cmd = %#v", cmd)
	err = cmd.Start()
	if err != nil {
		return
	}

	go cmd.Wait()

	cmdCh := make(chan Line)
	go func(cmdCh chan Line, rdr *bufio.Reader) {
		defer func() { recover() }()
		defer close(cmdCh)
		for {
			b, _, err := rdr.ReadLine()
			if len(b) > 0 {
				// TODO: need to redo the spec for custom matchers
				// This is the ONLY location where we need to actually
				// RECREATE a RawLine, and thus the only place where
				// ctx.enableSep is required.
				cmdCh <- NewMatchedLine(NewRawLine(string(b), ecf.enableSep), nil)
			}
			if err != nil {
				break
			}
		}
	}(cmdCh, bufio.NewReader(r))

	defer func() {
		if p := cmd.Process; p != nil {
			p.Kill()
		}
	}()

	defer trace("Done waiting for cancel or line")

	for {
		select {
		case <-cancelCh:
			return
		case l, ok := <-cmdCh:
			if l == nil || !ok {
				return
			}
			trace("Custom: l = %s", l.DisplayString())
			outputCh <- l
		}
	}
}
<commit_msg>Replace $QUERY like how the document says<commit_after>package peco

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"os/exec"
	"regexp"
	"sort"
	"strings"
	"unicode"
)

// These are used as keys in the config file
const (
	IgnoreCaseMatch    = "IgnoreCase"
	CaseSensitiveMatch = "CaseSensitive"
	SmartCaseMatch     = "SmartCase"
	RegexpMatch        = "Regexp"
)

var ignoreCaseFlags = []string{"i"}
var defaultFlags = []string{}

type regexpFlags interface {
	flags(string) []string
}
type regexpFlagList []string

func (r regexpFlagList) flags(_ string) []string {
	return []string(r)
}

type regexpFlagFunc func(string) []string

func (r regexpFlagFunc) flags(s string) []string {
	return r(s)
}

func containsUpper(query string) bool {
	for _, c := range query {
		if unicode.IsUpper(c) {
			return true
		}
	}
	return false
}

func regexpFor(q string, flags []string, quotemeta bool) (*regexp.Regexp, error) {
	reTxt := q
	if quotemeta {
		reTxt = regexp.QuoteMeta(q)
	}

	if flags != nil && len(flags) > 0 {
		reTxt = fmt.Sprintf("(?%s)%s", strings.Join(flags, ""), reTxt)
	}

	re, err := regexp.Compile(reTxt)
	if err != nil {
		return nil, err
	}
	return re, nil
}

func queryToRegexps(flags regexpFlags, quotemeta bool, query string) ([]*regexp.Regexp, error) {
	queries := strings.Split(strings.TrimSpace(query), " ")
	regexps := make([]*regexp.Regexp, 0)

	for _, q := range queries {
		re, err := regexpFor(q, flags.flags(query), quotemeta)
		if err != nil {
			return nil, err
		}
		regexps = append(regexps, re)
	}

	return regexps, nil
}

// sort related stuff
type byMatchStart [][]int

func (m byMatchStart) Len() int {
	return len(m)
}

func (m byMatchStart) Swap(i, j int) {
	m[i], m[j] = m[j], m[i]
}

func (m byMatchStart) Less(i, j int) bool {
	if m[i][0] < m[j][0] {
		return true
	}

	if m[i][0] == m[j][0] {
		return m[i][1]-m[i][0] < m[i][1]-m[i][0]
	}

	return false
}
func matchContains(a []int, b []int) bool {
	return a[0] <= b[0] && a[1] >= b[1]
}

func matchOverlaps(a []int, b []int) bool {
	return a[0] <= b[0] && a[1] >= b[0] ||
		a[0] <= b[1] && a[1] >= b[1]
}

func mergeMatches(a []int, b []int) []int {
	ret := make([]int, 2)

	// Note: In practice this should never happen
	// because we're sorting by N[0] before calling
	// this routine, but for completeness' sake...
	if a[0] < b[0] {
		ret[0] = a[0]
	} else {
		ret[0] = b[0]
	}

	if a[1] < b[1] {
		ret[1] = b[1]
	} else {
		ret[1] = a[1]
	}
	return ret
}

// Filter is responsible for the actual "grep" part of peco
type Filter struct {
	*Ctx
}

// Work is the actual work horse that that does the matching
// in a goroutine of its own. It wraps Matcher.Match().
func (f *Filter) Work(cancel chan struct{}, q HubReq) {
	trace("Filter.Work: START\n")
	defer trace("Filter.Work: END\n")
	defer q.Done()

	query := q.DataString()
	if query == "" {
		trace("Filter.Work: Resetting activingLineBuffer")
		f.ResetActiveLineBuffer()
	} else {
		f.rawLineBuffer.cancelCh = cancel
		f.rawLineBuffer.Replay()

		filter := f.Filter().Clone()
		filter.SetQuery(query)
		trace("Running %#v filter using query '%s'", filter, query)

		filter.Accept(f.rawLineBuffer)
		buf := NewRawLineBuffer()
		buf.onEnd = func() { f.SendStatusMsg("") }
		buf.Accept(filter)

		f.SetActiveLineBuffer(buf)
	}

	if ! f.config.StickySelection {
		f.SelectionClear()
	}
}

// Loop keeps watching for incoming queries, and upon receiving
// a query, spawns a goroutine to do the heavy work. It also
// checks for previously running queries, so we can avoid
// running many goroutines doing the grep at the same time
func (f *Filter) Loop() {
	defer f.ReleaseWaitGroup()

	// previous holds a channel that can cancel the previous
	// query. This is used when multiple queries come in succession
	// and the previous query is discarded anyway
	var previous chan struct{}
	for {
		select {
		case <-f.LoopCh():
			return
		case q := <-f.QueryCh():
			if previous != nil {
				// Tell the previous query to stop
				previous <- struct{}{}
			}
			previous = make(chan struct{}, 1)

			f.SendStatusMsg("Running query...")
			go f.Work(previous, q)
		}
	}
}

type QueryFilterer interface {
	Pipeliner
	Cancel()
	Clone() QueryFilterer
	Accept(Pipeliner)
	SetQuery(string)
	String() string
}

type SelectionFilter struct {
	sel *Selection
}

func (sf SelectionFilter) Name() string {
	return "SelectionFilter"
}

type RegexpFilter struct {
	simplePipeline
	compiledQuery []*regexp.Regexp
	flags         regexpFlags
	quotemeta     bool
	query         string
	name          string
	onEnd         func()
}

func NewRegexpFilter() *RegexpFilter {
	return &RegexpFilter{}
}

func (rf RegexpFilter) Clone() QueryFilterer {
	return &RegexpFilter{
		simplePipeline{},
		nil,
		rf.flags,
		rf.quotemeta,
		rf.query,
		rf.name,
		nil,
	}
}

func (rf *RegexpFilter) Accept(p Pipeliner) {
	cancelCh, incomingCh := p.Pipeline()
	rf.cancelCh = cancelCh
	rf.outputCh = make(chan Line)
	go acceptPipeline(cancelCh, incomingCh, rf.outputCh,
		&pipelineCtx{rf.filter, rf.onEnd})
}

var ErrFilterDidNotMatch = errors.New("error: filter did not match against given line")

func (rf *RegexpFilter) filter(l Line) (Line, error) {
	trace("RegexpFilter.filter: START")
	defer trace("RegexpFilter.filter: END")
	regexps, err := rf.getQueryAsRegexps()
	if err != nil {
		return nil, err
	}
	v := l.DisplayString()
	allMatched := true
	matches := [][]int{}
TryRegexps:
	for _, rx := range regexps {
		trace("RegexpFilter.filter: matching '%s' against '%s'", v, rx)
		match := rx.FindAllStringSubmatchIndex(v, -1)
		if match == nil {
			allMatched = false
			break TryRegexps
		}
		matches = append(matches, match...)
	}

	if !allMatched {
		return nil, ErrFilterDidNotMatch
	}

	trace("RegexpFilter.filter: line matched pattern\n")
	sort.Sort(byMatchStart(matches))

	// We need to "dedupe" the results. For example, if we matched the
	// same region twice, we don't want that to be drawn

	deduped := make([][]int, 0, len(matches))

	for i, m := range matches {
		// Always push the first one
		if i == 0 {
			deduped = append(deduped, m)
			continue
		}

		prev := deduped[len(deduped)-1]
		switch {
		case matchContains(prev, m):
			// If the previous match contains this one, then
			// don't do anything
			continue
		case matchOverlaps(prev, m):
			// If the previous match overlaps with this one,
			// merge the results and make it a bigger one
			deduped[len(deduped)-1] = mergeMatches(prev, m)
		default:
			deduped = append(deduped, m)
		}
	}
	return NewMatchedLine(l, deduped), nil
}

func (rf *RegexpFilter) getQueryAsRegexps() ([]*regexp.Regexp, error) {
	if q := rf.compiledQuery; q != nil {
		return q, nil
	}
	q, err := queryToRegexps(rf.flags, rf.quotemeta, rf.query)
	if err != nil {
		return nil, err
	}

	rf.compiledQuery = q
	return q, nil
}

func (rf *RegexpFilter) SetQuery(q string) {
	rf.query = q
	rf.compiledQuery = nil
}

func (rf RegexpFilter) String() string {
	return rf.name
}

type FilterSet struct {
	filters []QueryFilterer
	current int
}

func (fs *FilterSet) Size() int {
	return len(fs.filters)
}

func (fs *FilterSet) Add(qf QueryFilterer) error {
	fs.filters = append(fs.filters, qf)
	return nil
}

func (fs *FilterSet) Rotate() {
	fs.current++
	if fs.current >= len(fs.filters) {
		fs.current = 0
	}
	trace("FilterSet.Rotate: now filter in effect is %s", fs.filters[fs.current])
}

var ErrFilterNotFound = errors.New("specified filter was not found")

func (fs *FilterSet) SetCurrentByName(name string) error {
	for i, f := range fs.filters {
		if f.String() == name {
			fs.current = i
			return nil
		}
	}
	return ErrFilterNotFound
}

func (fs *FilterSet) GetCurrent() QueryFilterer {
	return fs.filters[fs.current]
}

func NewIgnoreCaseFilter() *RegexpFilter {
	return &RegexpFilter{
		flags:     regexpFlagList(ignoreCaseFlags),
		quotemeta: true,
		name:      "IgnoreCase",
	}
}

func NewCaseSensitiveFilter() *RegexpFilter {
	return &RegexpFilter{
		flags:     regexpFlagList(defaultFlags),
		quotemeta: true,
		name:      "CaseSensitive",
	}
}

// SmartCaseFilter turns ON the ignore-case flag in the regexp
// if the query contains a upper-case character
func NewSmartCaseFilter() *RegexpFilter {
	return &RegexpFilter{
		flags: regexpFlagFunc(func(q string) []string {
			if containsUpper(q) {
				return defaultFlags
			}
			return []string{"i"}
		}),
		quotemeta: true,
		name:      "SmartCase",
	}
}

type ExternalCmdFilter struct {
	simplePipeline
	enableSep       bool
	cmd             string
	args            []string
	name            string
	query           string
	thresholdBufsiz int
}

func NewExternalCmdFilter(name, cmd string, args []string, threshold int, enableSep bool) *ExternalCmdFilter {
	trace("name = %s, cmd = %s, args = %#v", name, cmd, args)
	if len(args) == 0 {
		args = []string{ "$QUERY" }
	}

	return &ExternalCmdFilter{
		simplePipeline:  simplePipeline{},
		enableSep:       enableSep,
		cmd:             cmd,
		args:            args,
		name:            name,
		thresholdBufsiz: threshold,
	}
}

func (ecf ExternalCmdFilter) Clone() QueryFilterer {
	return &ExternalCmdFilter{
		simplePipeline:  simplePipeline{},
		enableSep:       ecf.enableSep,
		cmd:             ecf.cmd,
		args:            ecf.args,
		name:            ecf.name,
		thresholdBufsiz: ecf.thresholdBufsiz,
	}
}

func (ecf *ExternalCmdFilter) Verify() error {
	if ecf.cmd == "" {
		return fmt.Errorf("no executable specified for custom matcher '%s'", ecf.name)
	}

	if _, err := exec.LookPath(ecf.cmd); err != nil {
		return err
	}
	return nil
}

func (ecf *ExternalCmdFilter) Accept(p Pipeliner) {
	cancelCh, incomingCh := p.Pipeline()
	outputCh := make(chan Line)
	ecf.cancelCh = cancelCh
	ecf.outputCh = outputCh

	go func() {
		defer close(outputCh)

		defer trace("ExternalCmdFilter.Accept: DONE")

		// for every N lines, execute the external command
		buf := []Line{}
		for l := range incomingCh {
			buf = append(buf, l)
			if len(buf) < ecf.thresholdBufsiz {
				continue
			}

			ecf.launchExternalCmd(buf, cancelCh, outputCh)
			buf = []Line{} // drain
		}

		if len(buf) > 0 {
			ecf.launchExternalCmd(buf, cancelCh, outputCh)
		}
	}()
}

func (ecf *ExternalCmdFilter) SetQuery(q string) {
	ecf.query = q
}

func (ecf ExternalCmdFilter) String() string {
	return ecf.name
}

func (ecf *ExternalCmdFilter) launchExternalCmd(buf []Line, cancelCh chan struct{}, outputCh chan Line) {
	defer func() { recover() }() // ignore errors

	trace("ExternalCmdFilter.launchExternalCmd: START")
	defer trace("ExternalCmdFilter.launchExternalCmd: END")

	trace("buf = %v", buf)

	args := append([]string(nil), ecf.args...)
	for i, v := range args {
		if v == "$QUERY" {
			args[i] = ecf.query
		}
	}
	cmd := exec.Command(ecf.cmd, args...)

	inbuf := &bytes.Buffer{}
	for _, l := range buf {
		inbuf.WriteString(l.DisplayString() + "\n")
	}

	cmd.Stdin = inbuf
	r, err := cmd.StdoutPipe()
	if err != nil {
		return
	}

	trace("cmd = %#v", cmd)
	err = cmd.Start()
	if err != nil {
		return
	}

	go cmd.Wait()

	cmdCh := make(chan Line)
	go func(cmdCh chan Line, rdr *bufio.Reader) {
		defer func() { recover() }()
		defer close(cmdCh)
		for {
			b, _, err := rdr.ReadLine()
			if len(b) > 0 {
				// TODO: need to redo the spec for custom matchers
				// This is the ONLY location where we need to actually
				// RECREATE a RawLine, and thus the only place where
				// ctx.enableSep is required.
				cmdCh <- NewMatchedLine(NewRawLine(string(b), ecf.enableSep), nil)
			}
			if err != nil {
				break
			}
		}
	}(cmdCh, bufio.NewReader(r))

	defer func() {
		if p := cmd.Process; p != nil {
			p.Kill()
		}
	}()

	defer trace("Done waiting for cancel or line")

	for {
		select {
		case <-cancelCh:
			return
		case l, ok := <-cmdCh:
			if l == nil || !ok {
				return
			}
			trace("Custom: l = %s", l.DisplayString())
			outputCh <- l
		}
	}
}
<|endoftext|>


<commit_before>package scenario

import (
	"io"
	"math/rand"

	"fmt"
	"io/ioutil"
	"math"
	"strconv"

	"github.com/PuerkitoBio/goquery"
	"github.com/catatsuy/isucon6-final/bench/action"
	"github.com/catatsuy/isucon6-final/bench/fails"
	"github.com/catatsuy/isucon6-final/bench/seed"
	"github.com/catatsuy/isucon6-final/bench/session"
	"github.com/catatsuy/isucon6-final/bench/svg"
)

func newSession(origins []string) *session.Session {
	return session.New(origins[rand.Intn(len(origins))])
}

func fetchCSRFToken(s *session.Session, path string) (string, bool) {
	var token string

	ok := action.Get(s, path, func(body io.Reader, l *fails.Logger) bool {
		doc, ok := makeDocument(body, l)
		if !ok {
			return false
		}

		token, ok = extractCsrfToken(doc, l)

		return ok
	})
	if !ok {
		return "", false
	}
	return token, true
}

func makeDocument(body io.Reader, l *fails.Logger) (*goquery.Document, bool) {
	doc, err := goquery.NewDocumentFromReader(body)
	if err != nil {
		l.Add("ページのHTMLがパースできませんでした", err)
		return nil, false
	}
	return doc, true
}

func extractCsrfToken(doc *goquery.Document, l *fails.Logger) (string, bool) {
	token := ""

	doc.Find("html").Each(func(_ int, selection *goquery.Selection) {
		if t, ok := selection.Attr("data-csrf-token"); ok {
			token = t
		}
	})

	ok := token != ""
	if !ok {
		l.Add("トークンが取得できませんでした", nil)
	}

	return token, ok
}

func extractImages(doc *goquery.Document) []string {
	imageUrls := []string{}

	doc.Find("img").Each(func(_ int, selection *goquery.Selection) {
		if url, ok := selection.Attr("src"); ok {
			imageUrls = append(imageUrls, url)
		}
	})

	return imageUrls
}

// 描いた線がsvgに反映されるか
func checkStrokeReflectedToSVG(s *session.Session, roomID int64, strokeID int64, stroke seed.Stroke) bool {
	imageURL := "/img/" + strconv.FormatInt(roomID, 10)

	return action.Get(s, imageURL, func(body io.Reader, l *fails.Logger) bool {
		b, err := ioutil.ReadAll(body)
		if err != nil {
			l.Critical("内容が読み込めませんでした", err)
			return false
		}
		data, err := svg.Parse(b)
		if err != nil {
			l.Critical("SVGがパースできませんでした", err)
			return false
		}
		for i, polyLine := range data.PolyLines {
			if data.PolyLines[i].ID == strconv.FormatInt(strokeID, 10) {
				if len(stroke.Points) != len(polyLine.Points) {
					l.Critical("投稿が反映されていません（pointが足りません）", err)
					return false
				}
				for j, p := range polyLine.Points {
					if math.Abs(float64(stroke.Points[j].X)-float64(p.X)) > 0.1 || math.Abs(float64(stroke.Points[j].Y)-float64(p.Y)) > 0.1 {
						fmt.Println(stroke.Points[j].X, p.X, stroke.Points[j].Y, p.Y)
						l.Critical("投稿が反映されていません（x,yの値が改変されています）", err)
						return false
					}
				}
				return true
			}
		}
		// ここに来るのは、IDがstroke.IDと同じpolylineが一つも無かったとき
		l.Critical("投稿が反映されていません", err)
		return false
	})
}

// TODO: ステータスコード以外にもチェックしたい
func loadImages(s *session.Session, images []string) bool {
	status := true
	for _, image := range images {
		ok := action.Get(s, image, func(body io.Reader, l *fails.Logger) bool {
			return false
		})
		status = status && ok
	}
	return status

	// TODO: 画像を並列リクエストするようにしてみたが、 connection reset by peer というエラーが出るので直列に戻した
	// もしかすると s.Transport.MaxIdleConnsPerHost ずつ処理するといけるのかも
	//errs := make(chan error, len(images))
	//for _, image := range images {
	//	go func(image string) {
	//		err := s.Get(image, func(status int, body io.Reader) error {
	//			if status != 200 {
	//				return errors.New("ステータスが200ではありません: " + strconv.Itoa(status))
	//			}
	//			return nil
	//		})
	//		errs <- err
	//	}(image)
	//}
	//var lastErr error
	//for i := 0; i < len(images); i++ {
	//	err := <-errs
	//	if err != nil {
	//		lastErr = err
	//	}
	//}
	//return lastErr
}
<commit_msg>loadImages makes MaxIdleConnsPerHost requests at a time<commit_after>package scenario

import (
	"io"
	"math/rand"

	"fmt"
	"io/ioutil"
	"math"
	"strconv"

	"github.com/PuerkitoBio/goquery"
	"github.com/catatsuy/isucon6-final/bench/action"
	"github.com/catatsuy/isucon6-final/bench/fails"
	"github.com/catatsuy/isucon6-final/bench/seed"
	"github.com/catatsuy/isucon6-final/bench/session"
	"github.com/catatsuy/isucon6-final/bench/svg"
)

func newSession(origins []string) *session.Session {
	return session.New(origins[rand.Intn(len(origins))])
}

func fetchCSRFToken(s *session.Session, path string) (string, bool) {
	var token string

	ok := action.Get(s, path, func(body io.Reader, l *fails.Logger) bool {
		doc, ok := makeDocument(body, l)
		if !ok {
			return false
		}

		token, ok = extractCsrfToken(doc, l)

		return ok
	})
	if !ok {
		return "", false
	}
	return token, true
}

func makeDocument(body io.Reader, l *fails.Logger) (*goquery.Document, bool) {
	doc, err := goquery.NewDocumentFromReader(body)
	if err != nil {
		l.Add("ページのHTMLがパースできませんでした", err)
		return nil, false
	}
	return doc, true
}

func extractCsrfToken(doc *goquery.Document, l *fails.Logger) (string, bool) {
	token := ""

	doc.Find("html").Each(func(_ int, selection *goquery.Selection) {
		if t, ok := selection.Attr("data-csrf-token"); ok {
			token = t
		}
	})

	ok := token != ""
	if !ok {
		l.Add("トークンが取得できませんでした", nil)
	}

	return token, ok
}

func extractImages(doc *goquery.Document) []string {
	imageUrls := []string{}

	doc.Find("img").Each(func(_ int, selection *goquery.Selection) {
		if url, ok := selection.Attr("src"); ok {
			imageUrls = append(imageUrls, url)
		}
	})

	return imageUrls
}

// 描いた線がsvgに反映されるか
func checkStrokeReflectedToSVG(s *session.Session, roomID int64, strokeID int64, stroke seed.Stroke) bool {
	imageURL := "/img/" + strconv.FormatInt(roomID, 10)

	return action.Get(s, imageURL, func(body io.Reader, l *fails.Logger) bool {
		b, err := ioutil.ReadAll(body)
		if err != nil {
			l.Critical("内容が読み込めませんでした", err)
			return false
		}
		data, err := svg.Parse(b)
		if err != nil {
			l.Critical("SVGがパースできませんでした", err)
			return false
		}
		for i, polyLine := range data.PolyLines {
			if data.PolyLines[i].ID == strconv.FormatInt(strokeID, 10) {
				if len(stroke.Points) != len(polyLine.Points) {
					l.Critical("投稿が反映されていません（pointが足りません）", err)
					return false
				}
				for j, p := range polyLine.Points {
					if math.Abs(float64(stroke.Points[j].X)-float64(p.X)) > 0.1 || math.Abs(float64(stroke.Points[j].Y)-float64(p.Y)) > 0.1 {
						fmt.Println(stroke.Points[j].X, p.X, stroke.Points[j].Y, p.Y)
						l.Critical("投稿が反映されていません（x,yの値が改変されています）", err)
						return false
					}
				}
				return true
			}
		}
		// ここに来るのは、IDがstroke.IDと同じpolylineが一つも無かったとき
		l.Critical("投稿が反映されていません", err)
		return false
	})
}

func loadImages(s *session.Session, images []string) bool {
	ch := make(chan struct{}, session.MaxIdleConnsPerHost)
	OK := true
	for _, image := range images {
		ch <- struct{}{}
		go func(image string) {
			ok := action.Get(s, image, func(body io.Reader, l *fails.Logger) bool {
				return true
			})
			if !ok {
				OK = false // ture -> false になるだけなのでmutexは不要と思われ
			}
			<-ch
		}(image)
	}
	return OK
}
<|endoftext|>