<commit_before>/*
Copyright 2019 The Knative Authors

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

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

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

package main

import (
	"context"
	"flag"
	"fmt"
	"log"

	"go.uber.org/zap"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/clientcmd"
	"knative.dev/pkg/configmap"
	"knative.dev/pkg/logging"
	"knative.dev/pkg/signals"
	"knative.dev/pkg/system"
	"knative.dev/pkg/version"
	"knative.dev/pkg/webhook"

	eventsv1alpha1 "github.com/google/knative-gcp/pkg/apis/events/v1alpha1"
	messagingv1alpha1 "github.com/google/knative-gcp/pkg/apis/messaging/v1alpha1"
	pubsubv1alpha1 "github.com/google/knative-gcp/pkg/apis/pubsub/v1alpha1"
)

const (
	component = "webhook"
)

var (
	masterURL  = flag.String("master", "", "The address of the Kubernetes API server. Overrides any value in kubeconfig. Only required if out-of-cluster.")
	kubeconfig = flag.String("kubeconfig", "", "Path to a kubeconfig. Only required if out-of-cluster.")
)

func SharedMain(resourceHandlers map[schema.GroupVersionKind]webhook.GenericCRD) {

	flag.Parse()
	cm, err := configmap.Load("/etc/config-logging")
	if err != nil {
		log.Fatal("Error loading logging configuration:", err)
	}
	config, err := logging.NewConfigFromMap(cm)
	if err != nil {
		log.Fatal("Error parsing logging configuration:", err)
	}
	logger, atomicLevel := logging.NewLoggerFromConfig(config, component)
	defer logger.Sync()
	logger = logger.With(zap.String("cloud.google.com/events", component))

	logger.Info("Starting the Cloud Run Events Webhook")

	// Set up signals so we handle the first shutdown signal gracefully.
	ctx := signals.NewContext()

	clusterConfig, err := clientcmd.BuildConfigFromFlags(*masterURL, *kubeconfig)
	if err != nil {
		logger.Fatalw("Failed to get cluster config", zap.Error(err))
	}

	kubeClient, err := kubernetes.NewForConfig(clusterConfig)
	if err != nil {
		logger.Fatalw("Failed to get the client set", zap.Error(err))
	}

	if err := version.CheckMinimumVersion(kubeClient.Discovery()); err != nil {
		logger.Fatalw("Version check failed", err)
	}

	// Watch the logging config map and dynamically update logging levels.
	configMapWatcher := configmap.NewInformedWatcher(kubeClient, system.Namespace())
	configMapWatcher.Watch(logging.ConfigMapName(), logging.UpdateLevelFromConfigMap(logger, atomicLevel, component))

	// // If you want to control Defaulting or Validation, you can attach config state
	// // to the context by watching the configmap here, and then uncommenting the logic
	// // below.
	// stores := make([]Store, 0, len(factories))
	// for _, sf := range factories {
	// 	store := sf(logger)
	// 	store.WatchConfigs(configMapWatcher)
	// 	stores = append(stores, store)
	// }

	if err = configMapWatcher.Start(ctx.Done()); err != nil {
		logger.Fatalw("Failed to start the ConfigMap watcher", zap.Error(err))
	}

	options := webhook.ControllerOptions{
		ServiceName:                     "webhook",
		DeploymentName:                  "webhook",
		Namespace:                       system.Namespace(),
		Port:                            8443,
		SecretName:                      "webhook-certs",
		ResourceMutatingWebhookName:     fmt.Sprintf("webhook.%s.events.cloud.google.com", system.Namespace()),
		ResourceAdmissionControllerPath: "/",
	}

	// Decorate contexts with the current state of the config.
	ctxFunc := func(ctx context.Context) context.Context {
		// TODO: implement upgrades when eventing needs it:
		//  return v1beta1.WithUpgradeViaDefaulting(store.ToContext(ctx))
		return ctx
	}

	resourceAdmissionController := webhook.NewResourceAdmissionController(resourceHandlers, options, true)
	admissionControllers := map[string]webhook.AdmissionController{
		options.ResourceAdmissionControllerPath: resourceAdmissionController,
	}

	controller, err := webhook.New(kubeClient, options, admissionControllers, logger, ctxFunc)

	if err != nil {
		logger.Fatalw("Failed to create admission controller", zap.Error(err))
	}

	if err = controller.Run(ctx.Done()); err != nil {
		logger.Fatalw("Failed to start the admission controller", zap.Error(err))
	}
}

func main() {
	handlers := map[schema.GroupVersionKind]webhook.GenericCRD{
		messagingv1alpha1.SchemeGroupVersion.WithKind("Channel"):       &messagingv1alpha1.Channel{},
		messagingv1alpha1.SchemeGroupVersion.WithKind("Decorator"):     &messagingv1alpha1.Decorator{},
		eventsv1alpha1.SchemeGroupVersion.WithKind("Storage"):          &eventsv1alpha1.Storage{},
		eventsv1alpha1.SchemeGroupVersion.WithKind("Scheduler"):        &eventsv1alpha1.Scheduler{},
		eventsv1alpha1.SchemeGroupVersion.WithKind("PubSub"):           &eventsv1alpha1.PubSub{},
		pubsubv1alpha1.SchemeGroupVersion.WithKind("PullSubscription"): &pubsubv1alpha1.PullSubscription{},
		pubsubv1alpha1.SchemeGroupVersion.WithKind("Topic"):            &pubsubv1alpha1.Topic{},
	}
	SharedMain(handlers)
}
<commit_msg>Switch to desugared logger. (#335)<commit_after>/*
Copyright 2019 The Knative Authors

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

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

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

package main

import (
	"context"
	"flag"
	"fmt"
	"log"

	"go.uber.org/zap"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/clientcmd"
	"knative.dev/pkg/configmap"
	"knative.dev/pkg/logging"
	"knative.dev/pkg/metrics"
	"knative.dev/pkg/signals"
	"knative.dev/pkg/system"
	"knative.dev/pkg/version"
	"knative.dev/pkg/webhook"

	eventsv1alpha1 "github.com/google/knative-gcp/pkg/apis/events/v1alpha1"
	messagingv1alpha1 "github.com/google/knative-gcp/pkg/apis/messaging/v1alpha1"
	pubsubv1alpha1 "github.com/google/knative-gcp/pkg/apis/pubsub/v1alpha1"
)

const (
	component = "webhook"
)

var (
	masterURL  = flag.String("master", "", "The address of the Kubernetes API server. Overrides any value in kubeconfig. Only required if out-of-cluster.")
	kubeconfig = flag.String("kubeconfig", "", "Path to a kubeconfig. Only required if out-of-cluster.")
)

func SharedMain(resourceHandlers map[schema.GroupVersionKind]webhook.GenericCRD) {

	flag.Parse()
	cm, err := configmap.Load("/etc/config-logging")
	if err != nil {
		log.Fatal("Error loading logging configuration:", err)
	}
	config, err := logging.NewConfigFromMap(cm)
	if err != nil {
		log.Fatal("Error parsing logging configuration:", err)
	}
	sl, atomicLevel := logging.NewLoggerFromConfig(config, component)
	logger := sl.Desugar().With(zap.String("cloud.google.com/events", component))
	defer flush(logger)

	logger.Info("Starting the Cloud Run Events Webhook")

	// Set up signals so we handle the first shutdown signal gracefully.
	ctx := signals.NewContext()

	clusterConfig, err := clientcmd.BuildConfigFromFlags(*masterURL, *kubeconfig)
	if err != nil {
		logger.Fatal("Failed to get cluster config", zap.Error(err))
	}

	kubeClient, err := kubernetes.NewForConfig(clusterConfig)
	if err != nil {
		logger.Fatal("Failed to get the client set", zap.Error(err))
	}

	if err := version.CheckMinimumVersion(kubeClient.Discovery()); err != nil {
		logger.Fatal("Version check failed", zap.Error(err))
	}

	// Watch the logging config map and dynamically update logging levels.
	configMapWatcher := configmap.NewInformedWatcher(kubeClient, system.Namespace())
	configMapWatcher.Watch(logging.ConfigMapName(), logging.UpdateLevelFromConfigMap(logger.Sugar(), atomicLevel, component))

	// // If you want to control Defaulting or Validation, you can attach config state
	// // to the context by watching the configmap here, and then uncommenting the logic
	// // below.
	// stores := make([]Store, 0, len(factories))
	// for _, sf := range factories {
	// 	store := sf(logger)
	// 	store.WatchConfigs(configMapWatcher)
	// 	stores = append(stores, store)
	// }

	if err = configMapWatcher.Start(ctx.Done()); err != nil {
		logger.Fatal("Failed to start the ConfigMap watcher", zap.Error(err))
	}

	options := webhook.ControllerOptions{
		ServiceName:                     "webhook",
		DeploymentName:                  "webhook",
		Namespace:                       system.Namespace(),
		Port:                            8443,
		SecretName:                      "webhook-certs",
		ResourceMutatingWebhookName:     fmt.Sprintf("webhook.%s.events.cloud.google.com", system.Namespace()),
		ResourceAdmissionControllerPath: "/",
	}

	// Decorate contexts with the current state of the config.
	ctxFunc := func(ctx context.Context) context.Context {
		// TODO: implement upgrades when eventing needs it:
		//  return v1beta1.WithUpgradeViaDefaulting(store.ToContext(ctx))
		return ctx
	}

	resourceAdmissionController := webhook.NewResourceAdmissionController(resourceHandlers, options, true)
	admissionControllers := map[string]webhook.AdmissionController{
		options.ResourceAdmissionControllerPath: resourceAdmissionController,
	}

	controller, err := webhook.New(kubeClient, options, admissionControllers, logger.Sugar(), ctxFunc)

	if err != nil {
		logger.Fatal("Failed to create admission controller", zap.Error(err))
	}

	if err = controller.Run(ctx.Done()); err != nil {
		logger.Fatal("Failed to start the admission controller", zap.Error(err))
	}
}

func main() {
	handlers := map[schema.GroupVersionKind]webhook.GenericCRD{
		messagingv1alpha1.SchemeGroupVersion.WithKind("Channel"):       &messagingv1alpha1.Channel{},
		messagingv1alpha1.SchemeGroupVersion.WithKind("Decorator"):     &messagingv1alpha1.Decorator{},
		eventsv1alpha1.SchemeGroupVersion.WithKind("Storage"):          &eventsv1alpha1.Storage{},
		eventsv1alpha1.SchemeGroupVersion.WithKind("Scheduler"):        &eventsv1alpha1.Scheduler{},
		eventsv1alpha1.SchemeGroupVersion.WithKind("PubSub"):           &eventsv1alpha1.PubSub{},
		pubsubv1alpha1.SchemeGroupVersion.WithKind("PullSubscription"): &pubsubv1alpha1.PullSubscription{},
		pubsubv1alpha1.SchemeGroupVersion.WithKind("Topic"):            &pubsubv1alpha1.Topic{},
	}
	SharedMain(handlers)
}

func flush(logger *zap.Logger) {
	_ = logger.Sync()
	metrics.FlushExporter()
}
<|endoftext|>


<commit_before>package commands

import (
	"bufio"
	"bytes"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/EMSSConsulting/Thargo"
	"github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/common"
)

type ArchiveCommand struct {
	Paths     []string `long:"path" description:"Glob based filters used to include files"`
	Untracked bool     `long:"untracked" description:"Add git untracked files"`
	Output    string   `long:"output" description:"The filepath to output file"`
	Silent    bool     `long:"silent" description:"Suppress archiving ouput"`
	List      bool     `long:"list" description:"List files to archive"`

	files map[string]time.Time
	wd    string
}

func (c *ArchiveCommand) isChanged(modTime time.Time) bool {
	for _, fileModTime := range c.files {
		if modTime.Before(fileModTime) {
			return true
		}
	}
	return false
}

func (c *ArchiveCommand) archive() {
	logrus.Infoln("Creating archive", filepath.Base(c.Output), "...")

	// create directories to store archive
	os.MkdirAll(filepath.Dir(c.Output), 0700)

	tempFile, err := ioutil.TempFile(filepath.Dir(c.Output), "archive_")
	if err != nil {
		logrus.Fatalln("Failed to create temporary archive", err)
	}
	tempFile.Close()

	defer os.Remove(tempFile.Name())

	archive, err := thargo.NewArchiveFile(tempFile.Name(), nil)
	if err != nil {
		logrus.Fatalln("Failed to open archive for writing: ", err)
	}

	logrus.Debugln("Temporary file:", tempFile.Name())

	includePredicate := func(entry thargo.Entry) bool {
		header, err := entry.Header()
		if err != nil {
			return false
		}

		// Fix up the header name if the path is outside of the cwd
		if filepath.HasPrefix(header.Name, "../") {
			absPath, err := filepath.Abs(header.Name)
			if err != nil {
				return false
			}

			header.Name = absPath
		}

		// Don't include duplicate files
		if _, exists := c.files[header.Name]; exists {
			return false
		}

		c.files[header.Name] = header.ChangeTime
		return true
	}

	for _, filter := range c.Paths {
		if !c.Silent {
			logrus.Infof("Adding '%s' to archive", filter)
		}

		if err := archive.AddIf(&thargo.FileSystemTarget{
			Path:    c.wd,
			Pattern: filter,
		}, includePredicate); err != nil {
			logrus.Warnf("Failed to add '%s' to archive: %s", filter, err)
		}
	}

	if c.Untracked {
		if err := archive.AddIf(&GitUntrackedFilesTarget{
			Path: c.wd,
		}, includePredicate); err != nil {
			logrus.Warnf("Failed to add git untracked files to archive: %s", err)
		}
	}

	if err := archive.Close(); err != nil {
		logrus.Warningln("Failed to close temp archive: ", err)
	}

	ai, err := os.Stat(c.Output)
	if err != nil && !os.IsNotExist(err) {
		logrus.Fatalln("Failed to verify archive:", c.Output, err)
	}
	if ai != nil && !c.isChanged(ai.ModTime()) {
		logrus.Infoln("Archive is up to date!")
		return
	}

	err = os.Rename(tempFile.Name(), c.Output)
	if err != nil {
		logrus.Warningln("Failed to rename archive:", err)
	}

	logrus.Infoln("Done!")
}

func (c *ArchiveCommand) Execute(context *cli.Context) {
	logrus.SetFormatter(
		&logrus.TextFormatter{
			ForceColors:      true,
			DisableTimestamp: false,
		},
	)

	wd, err := os.Getwd()
	if err != nil {
		logrus.Fatalln("Failed to get current working directory:", err)
	}
	if c.Output == "" && !c.List {
		logrus.Fatalln("Missing archive file name!")
	}

	c.wd = wd
	c.files = make(map[string]time.Time)

	if len(c.Paths) == 0 && !c.Untracked {
		logrus.Infoln("No inclusion filters specified.")
		return
	}

	c.archive()

	if !c.Silent {
		for file := range c.files {
			logrus.Println(" - ", file)
		}
	}
}

func init() {
	common.RegisterCommand2("archive", "find and archive files (internal)", &ArchiveCommand{})
}

// GitUntrackedFilesTarget is a Thargo target which provides a list of untracked
// git files for inclusion in an archive.
type GitUntrackedFilesTarget struct {
	Path string
}

func (t *GitUntrackedFilesTarget) Entries() ([]thargo.Entry, error) {
	entries := []thargo.Entry{}

	var output bytes.Buffer
	cmd := exec.Command("git", "ls-files", "-o")
	cmd.Env = os.Environ()
	cmd.Stdout = &output
	cmd.Stderr = os.Stderr
	logrus.Debugln("Executing command:", strings.Join(cmd.Args, " "))
	err := cmd.Run()
	if err == nil {
		reader := bufio.NewReader(&output)
		for {
			line, _, err := reader.ReadLine()
			if err == io.EOF {
				break
			} else if err != nil {
				logrus.Warningln(err)
				break
			}

			entry, err := t.processFile(string(line))
			if err != nil {
				logrus.Warningln("Failed to include untracked file", line, err)
			} else if entry == nil {
				continue
			}

			entries = append(entries, entry)
		}
	} else {
		logrus.Warningln(err)
	}

	return entries, nil
}

func (t *GitUntrackedFilesTarget) processFile(path string) (thargo.Entry, error) {
	f, err := os.Stat(path)
	if os.IsNotExist(err) {
		return nil, nil
	} else if err != nil {
		return nil, err
	}

	// Skip directories
	if f.IsDir() {
		return nil, nil
	}

	absPath, err := filepath.Abs(path)
	if err != nil {
		return nil, err
	}

	return &thargo.FileEntry{Name: path, Path: absPath, Info: f}, nil
}
<commit_msg>feat: Add support for -list argument to archive command<commit_after>package commands

import (
	"bufio"
	"bytes"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/EMSSConsulting/Thargo"
	"github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/common"
)

type ArchiveCommand struct {
	Paths     []string `long:"path" description:"Glob based filters used to include files"`
	Untracked bool     `long:"untracked" description:"Add git untracked files"`
	Output    string   `long:"output" description:"The filepath to output file"`
	Silent    bool     `long:"silent" description:"Suppress archiving ouput"`
	List      bool     `long:"list" description:"List files to archive"`

	files map[string]time.Time
	wd    string
}

func (c *ArchiveCommand) isChanged(modTime time.Time) bool {
	for _, fileModTime := range c.files {
		if modTime.Before(fileModTime) {
			return true
		}
	}
	return false
}

func (c *ArchiveCommand) archive() {
	logrus.Infoln("Creating archive", filepath.Base(c.Output), "...")

	// create directories to store archive
	os.MkdirAll(filepath.Dir(c.Output), 0700)

	tempFile, err := ioutil.TempFile(filepath.Dir(c.Output), "archive_")
	if err != nil {
		logrus.Fatalln("Failed to create temporary archive", err)
	}
	tempFile.Close()

	defer os.Remove(tempFile.Name())

	archive, err := thargo.NewArchiveFile(tempFile.Name(), nil)
	if err != nil {
		logrus.Fatalln("Failed to open archive for writing: ", err)
	}

	logrus.Debugln("Temporary file:", tempFile.Name())

	includePredicate := func(entry thargo.Entry) bool {
		header, err := entry.Header()
		if err != nil {
			return false
		}

		// Fix up the header name if the path is outside of the cwd
		if filepath.HasPrefix(header.Name, "../") {
			absPath, err := filepath.Abs(header.Name)
			if err != nil {
				return false
			}

			header.Name = absPath
		}

		// Don't include duplicate files
		if _, exists := c.files[header.Name]; exists {
			return false
		}

		if(c.List) {
			println(header.Name)
		}

		c.files[header.Name] = header.ChangeTime
		return true
	}

	for _, filter := range c.Paths {
		if !c.Silent {
			logrus.Infof("Adding '%s' to archive", filter)
		}

		if err := archive.AddIf(&thargo.FileSystemTarget{
			Path:    c.wd,
			Pattern: filter,
		}, includePredicate); err != nil {
			logrus.Warnf("Failed to add '%s' to archive: %s", filter, err)
		}
	}

	if c.Untracked {
		if err := archive.AddIf(&GitUntrackedFilesTarget{
			Path: c.wd,
		}, includePredicate); err != nil {
			logrus.Warnf("Failed to add git untracked files to archive: %s", err)
		}
	}

	if err := archive.Close(); err != nil {
		logrus.Warningln("Failed to close temp archive: ", err)
	}

	ai, err := os.Stat(c.Output)
	if err != nil && !os.IsNotExist(err) {
		logrus.Fatalln("Failed to verify archive:", c.Output, err)
	}
	if ai != nil && !c.isChanged(ai.ModTime()) {
		logrus.Infoln("Archive is up to date!")
		return
	}

	err = os.Rename(tempFile.Name(), c.Output)
	if err != nil {
		logrus.Warningln("Failed to rename archive:", err)
	}

	logrus.Infoln("Done!")
}

func (c *ArchiveCommand) Execute(context *cli.Context) {
	logrus.SetFormatter(
		&logrus.TextFormatter{
			ForceColors:      true,
			DisableTimestamp: false,
		},
	)

	wd, err := os.Getwd()
	if err != nil {
		logrus.Fatalln("Failed to get current working directory:", err)
	}
	if c.Output == "" && !c.List {
		logrus.Fatalln("Missing archive file name!")
	}

	c.wd = wd
	c.files = make(map[string]time.Time)

	if len(c.Paths) == 0 && !c.Untracked {
		logrus.Infoln("No inclusion filters specified.")
		return
	}

	c.archive()

	if !c.Silent {
		for file := range c.files {
			logrus.Println(" - ", file)
		}
	}
}

func init() {
	common.RegisterCommand2("archive", "find and archive files (internal)", &ArchiveCommand{})
}

// GitUntrackedFilesTarget is a Thargo target which provides a list of untracked
// git files for inclusion in an archive.
type GitUntrackedFilesTarget struct {
	Path string
}

func (t *GitUntrackedFilesTarget) Entries() ([]thargo.Entry, error) {
	entries := []thargo.Entry{}

	var output bytes.Buffer
	cmd := exec.Command("git", "ls-files", "-o")
	cmd.Env = os.Environ()
	cmd.Stdout = &output
	cmd.Stderr = os.Stderr
	logrus.Debugln("Executing command:", strings.Join(cmd.Args, " "))
	err := cmd.Run()
	if err == nil {
		reader := bufio.NewReader(&output)
		for {
			line, _, err := reader.ReadLine()
			if err == io.EOF {
				break
			} else if err != nil {
				logrus.Warningln(err)
				break
			}

			entry, err := t.processFile(string(line))
			if err != nil {
				logrus.Warningln("Failed to include untracked file", line, err)
			} else if entry == nil {
				continue
			}

			entries = append(entries, entry)
		}
	} else {
		logrus.Warningln(err)
	}

	return entries, nil
}

func (t *GitUntrackedFilesTarget) processFile(path string) (thargo.Entry, error) {
	f, err := os.Stat(path)
	if os.IsNotExist(err) {
		return nil, nil
	} else if err != nil {
		return nil, err
	}

	// Skip directories
	if f.IsDir() {
		return nil, nil
	}

	absPath, err := filepath.Abs(path)
	if err != nil {
		return nil, err
	}

	return &thargo.FileEntry{Name: path, Path: absPath, Info: f}, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2015 Google Inc. All rights reserved.

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

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

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

package commands

import (
	"errors"
	"flag"
	"fmt"
	"github.com/google/git-appraise/commands/input"
	"github.com/google/git-appraise/repository"
	"github.com/google/git-appraise/review"
	"github.com/google/git-appraise/review/comment"
	"strings"
)

var commentFlagSet = flag.NewFlagSet("comment", flag.ExitOnError)

var (
	commentMessage = commentFlagSet.String("m", "", "Message to attach to the review")
	commentParent  = commentFlagSet.String("p", "", "Parent comment")
	commentFile    = commentFlagSet.String("f", "", "File being commented upon")
	commentLine    = commentFlagSet.Uint("l", 0, "Line being commented upon; requires that the -f flag also be set")
	commentLgtm    = commentFlagSet.Bool("lgtm", false, "'Looks Good To Me'. Set this to express your approval. This cannot be combined with nmw")
	commentNmw     = commentFlagSet.Bool("nmw", false, "'Needs More Work'. Set this to express your disapproval. This cannot be combined with lgtm")
)

// commentHashExists checks if the given comment hash exists in the given comment threads.
func commentHashExists(hashToFind string, threads []review.CommentThread) bool {
	for _, thread := range threads {
		if thread.Hash == hashToFind {
			return true
		}
		if commentHashExists(hashToFind, thread.Children) {
			return true
		}
	}
	return false
}

// checkCommentLocation verifies that the given location exists at the given commit.
func checkCommentLocation(repo repository.Repo, commit, file string, line uint) error {
	contents, err := repo.Show(commit, file)
	if err != nil {
		return err
	}
	lines := strings.Split(contents, "\n")
	if line >= uint(len(lines)) {
		return fmt.Errorf("Line number %d does not exist in file %q", line, file)
	}
	return nil
}

// commentOnReview adds a comment to the current code review.
func commentOnReview(repo repository.Repo, args []string) error {
	commentFlagSet.Parse(args)
	args = commentFlagSet.Args()

	var r *review.Review
	var err error
	if len(args) > 1 {
		return errors.New("Only accepting a single review is supported.")
	}

	if len(args) == 1 {
		r, err = review.Get(repo, args[0])
	} else {
		r, err = review.GetCurrent(repo)
	}

	if err != nil {
		return fmt.Errorf("Failed to load the review: %v\n", err)
	}
	if r == nil {
		return errors.New("There is no matching review.")
	}

	if *commentLgtm && *commentNmw {
		return errors.New("You cannot combine the flags -lgtm and -nmw.")
	}
	if *commentLine != 0 && *commentFile == "" {
		return errors.New("Specifying a line number with the -l flag requires that you also specify a file name with the -f flag.")
	}
	if *commentParent != "" && !commentHashExists(*commentParent, r.Comments) {
		return errors.New("There is no matching parent comment.")
	}

	if *commentMessage == "" {
		*commentMessage, err = input.LaunchEditor(repo, commentFilename)
		if err != nil {
			return err
		}
	}

	commentedUponCommit, err := r.GetHeadCommit()
	if err != nil {
		return err
	}
	location := comment.Location{
		Commit: commentedUponCommit,
	}
	if *commentFile != "" {
		if err := checkCommentLocation(r.Repo, commentedUponCommit, *commentFile, *commentLine); err != nil {
			return fmt.Errorf("Unabled to comment on the given location: %v", err)
		}
		location.Path = *commentFile
		if *commentLine != 0 {
			location.Range = &comment.Range{
				StartLine: uint32(*commentLine),
			}
		}
	}

	userEmail, err := repo.GetUserEmail()
	if err != nil {
		return err
	}
	c := comment.New(userEmail, *commentMessage)
	c.Location = &location
	c.Parent = *commentParent
	if *commentLgtm || *commentNmw {
		resolved := *commentLgtm
		c.Resolved = &resolved
	}
	return r.AddComment(c)
}

// commentCmd defines the "comment" subcommand.
var commentCmd = &Command{
	Usage: func(arg0 string) {
		fmt.Printf("Usage: %s comment [<option>...] [<review-hash>]\n\nOptions:\n", arg0)
		commentFlagSet.PrintDefaults()
	},
	RunMethod: func(repo repository.Repo, args []string) error {
		return commentOnReview(repo, args)
	},
}
<commit_msg>Fix line-number<commit_after>/*
Copyright 2015 Google Inc. All rights reserved.

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

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

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

package commands

import (
	"errors"
	"flag"
	"fmt"
	"github.com/google/git-appraise/commands/input"
	"github.com/google/git-appraise/repository"
	"github.com/google/git-appraise/review"
	"github.com/google/git-appraise/review/comment"
	"strings"
)

var commentFlagSet = flag.NewFlagSet("comment", flag.ExitOnError)

var (
	commentMessage = commentFlagSet.String("m", "", "Message to attach to the review")
	commentParent  = commentFlagSet.String("p", "", "Parent comment")
	commentFile    = commentFlagSet.String("f", "", "File being commented upon")
	commentLine    = commentFlagSet.Uint("l", 0, "Line being commented upon; requires that the -f flag also be set")
	commentLgtm    = commentFlagSet.Bool("lgtm", false, "'Looks Good To Me'. Set this to express your approval. This cannot be combined with nmw")
	commentNmw     = commentFlagSet.Bool("nmw", false, "'Needs More Work'. Set this to express your disapproval. This cannot be combined with lgtm")
)

// commentHashExists checks if the given comment hash exists in the given comment threads.
func commentHashExists(hashToFind string, threads []review.CommentThread) bool {
	for _, thread := range threads {
		if thread.Hash == hashToFind {
			return true
		}
		if commentHashExists(hashToFind, thread.Children) {
			return true
		}
	}
	return false
}

// checkCommentLocation verifies that the given location exists at the given commit.
func checkCommentLocation(repo repository.Repo, commit, file string, line uint) error {
	contents, err := repo.Show(commit, file)
	if err != nil {
		return err
	}
	lines := strings.Split(contents, "\n")
	if line < 1 || line > uint(len(lines)) {
		return fmt.Errorf("Line number %d does not exist in file %q", line, file)
	}
	return nil
}

// commentOnReview adds a comment to the current code review.
func commentOnReview(repo repository.Repo, args []string) error {
	commentFlagSet.Parse(args)
	args = commentFlagSet.Args()

	var r *review.Review
	var err error
	if len(args) > 1 {
		return errors.New("Only accepting a single review is supported.")
	}

	if len(args) == 1 {
		r, err = review.Get(repo, args[0])
	} else {
		r, err = review.GetCurrent(repo)
	}

	if err != nil {
		return fmt.Errorf("Failed to load the review: %v\n", err)
	}
	if r == nil {
		return errors.New("There is no matching review.")
	}

	if *commentLgtm && *commentNmw {
		return errors.New("You cannot combine the flags -lgtm and -nmw.")
	}
	if *commentLine != 0 && *commentFile == "" {
		return errors.New("Specifying a line number with the -l flag requires that you also specify a file name with the -f flag.")
	}
	if *commentParent != "" && !commentHashExists(*commentParent, r.Comments) {
		return errors.New("There is no matching parent comment.")
	}

	if *commentMessage == "" {
		*commentMessage, err = input.LaunchEditor(repo, commentFilename)
		if err != nil {
			return err
		}
	}

	commentedUponCommit, err := r.GetHeadCommit()
	if err != nil {
		return err
	}
	location := comment.Location{
		Commit: commentedUponCommit,
	}
	if *commentFile != "" {
		if err := checkCommentLocation(r.Repo, commentedUponCommit, *commentFile, *commentLine); err != nil {
			return fmt.Errorf("Unabled to comment on the given location: %v", err)
		}
		location.Path = *commentFile
		if *commentLine != 0 {
			location.Range = &comment.Range{
				StartLine: uint32(*commentLine),
			}
		}
	}

	userEmail, err := repo.GetUserEmail()
	if err != nil {
		return err
	}
	c := comment.New(userEmail, *commentMessage)
	c.Location = &location
	c.Parent = *commentParent
	if *commentLgtm || *commentNmw {
		resolved := *commentLgtm
		c.Resolved = &resolved
	}
	return r.AddComment(c)
}

// commentCmd defines the "comment" subcommand.
var commentCmd = &Command{
	Usage: func(arg0 string) {
		fmt.Printf("Usage: %s comment [<option>...] [<review-hash>]\n\nOptions:\n", arg0)
		commentFlagSet.PrintDefaults()
	},
	RunMethod: func(repo repository.Repo, args []string) error {
		return commentOnReview(repo, args)
	},
}
<|endoftext|>


<commit_before>/*
Copyright 2015 Google Inc. All rights reserved.

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

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

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

package commands

import (
	"errors"
	"flag"
	"fmt"

	"github.com/google/git-appraise/commands/input"
	"github.com/google/git-appraise/repository"
	"github.com/google/git-appraise/review"
	"github.com/google/git-appraise/review/comment"
	"github.com/google/git-appraise/review/gpg"
)

var commentFlagSet = flag.NewFlagSet("comment", flag.ExitOnError)
var commentLocation = comment.Range{}

var (
	commentMessageFile = commentFlagSet.String("F", "", "Take the comment from the given file. Use - to read the message from the standard input")
	commentMessage     = commentFlagSet.String("m", "", "Message to attach to the review")
	commentParent      = commentFlagSet.String("p", "", "Parent comment")
	commentFile        = commentFlagSet.String("f", "", "File being commented upon")
	commentLgtm        = commentFlagSet.Bool("lgtm", false, "'Looks Good To Me'. Set this to express your approval. This cannot be combined with nmw")
	commentNmw         = commentFlagSet.Bool("nmw", false, "'Needs More Work'. Set this to express your disapproval. This cannot be combined with lgtm")
	commentSign        = commentFlagSet.Bool("S", false,
		"Sign the contents of the comment")
)

func init() {
	commentFlagSet.Var(&commentLocation, "l", "File location be commented upon; requires that the -f flag also be set")
}

// commentHashExists checks if the given comment hash exists in the given comment threads.
func commentHashExists(hashToFind string, threads []review.CommentThread) bool {
	for _, thread := range threads {
		if thread.Hash == hashToFind {
			return true
		}
		if commentHashExists(hashToFind, thread.Children) {
			return true
		}
	}
	return false
}

// commentOnReview adds a comment to the current code review.
func commentOnReview(repo repository.Repo, args []string) error {
	commentFlagSet.Parse(args)
	args = commentFlagSet.Args()

	var r *review.Review
	var err error
	if len(args) > 1 {
		return errors.New("Only accepting a single review is supported.")
	}

	if len(args) == 1 {
		r, err = review.Get(repo, args[0])
	} else {
		r, err = review.GetCurrent(repo)
	}

	if err != nil {
		return fmt.Errorf("Failed to load the review: %v\n", err)
	}
	if r == nil {
		return errors.New("There is no matching review.")
	}

	if *commentLgtm && *commentNmw {
		return errors.New("You cannot combine the flags -lgtm and -nmw.")
	}
	if commentLocation != (comment.Range{}) && *commentFile == "" {
		return errors.New("Specifying a line number with the -l flag requires that you also specify a file name with the -f flag.")
	}
	if *commentParent != "" && !commentHashExists(*commentParent, r.Comments) {
		return errors.New("There is no matching parent comment.")
	}

	if *commentMessageFile != "" && *commentMessage == "" {
		*commentMessage, err = input.FromFile(*commentMessageFile)
		if err != nil {
			return err
		}
	}
	if *commentMessageFile == "" && *commentMessage == "" {
		*commentMessage, err = input.LaunchEditor(repo, commentFilename)
		if err != nil {
			return err
		}
	}

	commentedUponCommit, err := r.GetHeadCommit()
	if err != nil {
		return err
	}
	location := comment.Location{
		Commit: commentedUponCommit,
	}
	if *commentFile != "" {
		location.Path = *commentFile
		location.Range = &commentLocation
		if err := location.Check(r.Repo); err != nil {
			return fmt.Errorf("Unable to comment on the given location: %v", err)
		}
	}

	userEmail, err := repo.GetUserEmail()
	if err != nil {
		return err
	}
	c := comment.New(userEmail, *commentMessage)
	c.Location = &location
	c.Parent = *commentParent
	if *commentLgtm || *commentNmw {
		resolved := *commentLgtm
		c.Resolved = &resolved
	}

	if *commentSign {
		key, err := repo.GetUserSigningKey()
		if err != nil {
			return err
		}
		err = gpg.Sign(key, &c)
		if err != nil {
			return err
		}
	}

	return r.AddComment(c)
}

// commentCmd defines the "comment" subcommand.
var commentCmd = &Command{
	Usage: func(arg0 string) {
		fmt.Printf("Usage: %s comment [<option>...] [<review-hash>]\n\nOptions:\n", arg0)
		commentFlagSet.PrintDefaults()
	},
	RunMethod: func(repo repository.Repo, args []string) error {
		return commentOnReview(repo, args)
	},
}
<commit_msg>Describe format of comment location specification<commit_after>/*
Copyright 2015 Google Inc. All rights reserved.

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

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

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

package commands

import (
	"errors"
	"flag"
	"fmt"

	"github.com/google/git-appraise/commands/input"
	"github.com/google/git-appraise/repository"
	"github.com/google/git-appraise/review"
	"github.com/google/git-appraise/review/comment"
	"github.com/google/git-appraise/review/gpg"
)

var commentFlagSet = flag.NewFlagSet("comment", flag.ExitOnError)
var commentLocation = comment.Range{}

var (
	commentMessageFile = commentFlagSet.String("F", "", "Take the comment from the given file. Use - to read the message from the standard input")
	commentMessage     = commentFlagSet.String("m", "", "Message to attach to the review")
	commentParent      = commentFlagSet.String("p", "", "Parent comment")
	commentFile        = commentFlagSet.String("f", "", "File being commented upon")
	commentLgtm        = commentFlagSet.Bool("lgtm", false, "'Looks Good To Me'. Set this to express your approval. This cannot be combined with nmw")
	commentNmw         = commentFlagSet.Bool("nmw", false, "'Needs More Work'. Set this to express your disapproval. This cannot be combined with lgtm")
	commentSign        = commentFlagSet.Bool("S", false,
		"Sign the contents of the comment")
)

func init() {
	commentFlagSet.Var(&commentLocation, "l",
		`File location to be commented upon; requires that the -f flag also be set.
Location follows the following format:
    <START LINE>[+<START COLUMN>][:<END LINE>[+<END COLUMN>]]
So, in order to comment starting on the 5th character of the 2nd line until (and
including) the 4th character of the 7th line, use:
    -l 2+5:7+4`)
}

// commentHashExists checks if the given comment hash exists in the given comment threads.
func commentHashExists(hashToFind string, threads []review.CommentThread) bool {
	for _, thread := range threads {
		if thread.Hash == hashToFind {
			return true
		}
		if commentHashExists(hashToFind, thread.Children) {
			return true
		}
	}
	return false
}

// commentOnReview adds a comment to the current code review.
func commentOnReview(repo repository.Repo, args []string) error {
	commentFlagSet.Parse(args)
	args = commentFlagSet.Args()

	var r *review.Review
	var err error
	if len(args) > 1 {
		return errors.New("Only accepting a single review is supported.")
	}

	if len(args) == 1 {
		r, err = review.Get(repo, args[0])
	} else {
		r, err = review.GetCurrent(repo)
	}

	if err != nil {
		return fmt.Errorf("Failed to load the review: %v\n", err)
	}
	if r == nil {
		return errors.New("There is no matching review.")
	}

	if *commentLgtm && *commentNmw {
		return errors.New("You cannot combine the flags -lgtm and -nmw.")
	}
	if commentLocation != (comment.Range{}) && *commentFile == "" {
		return errors.New("Specifying a line number with the -l flag requires that you also specify a file name with the -f flag.")
	}
	if *commentParent != "" && !commentHashExists(*commentParent, r.Comments) {
		return errors.New("There is no matching parent comment.")
	}

	if *commentMessageFile != "" && *commentMessage == "" {
		*commentMessage, err = input.FromFile(*commentMessageFile)
		if err != nil {
			return err
		}
	}
	if *commentMessageFile == "" && *commentMessage == "" {
		*commentMessage, err = input.LaunchEditor(repo, commentFilename)
		if err != nil {
			return err
		}
	}

	commentedUponCommit, err := r.GetHeadCommit()
	if err != nil {
		return err
	}
	location := comment.Location{
		Commit: commentedUponCommit,
	}
	if *commentFile != "" {
		location.Path = *commentFile
		location.Range = &commentLocation
		if err := location.Check(r.Repo); err != nil {
			return fmt.Errorf("Unable to comment on the given location: %v", err)
		}
	}

	userEmail, err := repo.GetUserEmail()
	if err != nil {
		return err
	}
	c := comment.New(userEmail, *commentMessage)
	c.Location = &location
	c.Parent = *commentParent
	if *commentLgtm || *commentNmw {
		resolved := *commentLgtm
		c.Resolved = &resolved
	}

	if *commentSign {
		key, err := repo.GetUserSigningKey()
		if err != nil {
			return err
		}
		err = gpg.Sign(key, &c)
		if err != nil {
			return err
		}
	}

	return r.AddComment(c)
}

// commentCmd defines the "comment" subcommand.
var commentCmd = &Command{
	Usage: func(arg0 string) {
		fmt.Printf("Usage: %s comment [<option>...] [<review-hash>]\n\nOptions:\n", arg0)
		commentFlagSet.PrintDefaults()
	},
	RunMethod: func(repo repository.Repo, args []string) error {
		return commentOnReview(repo, args)
	},
}
<|endoftext|>


<commit_before>// Copyright 2013 Marc Weistroff. All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package gogol

import (
	"bytes"
	"testing"
	"time"
)

func ExampleLogger() {
	handler := NewStdoutHandler(WARNING) // Will log all messages where severity is >= that WARNING
	handler.SetFormatter(NewMinimalLineFormatter())
	logger := NewLogger("example", []Handler{handler}, []Processor{})
	logger.Debug("Debug message that won't be displayed")
	logger.Warning("I sense a disturbance in the force")
	logger.Error("This is an error")
	// Output:
	// example.WARNING: I sense a disturbance in the force
	// example.ERROR: This is an error
}

func TestLogger(t *testing.T) {
	buffer := new(bytes.Buffer)
	handler := NewBufferHandler(buffer, NOTICE)
	handler.SetFormatter(&LineFormatter{"%channel%.%level_name%: %message%"})
	logger := NewLogger("test", []Handler{handler}, []Processor{})

	logger.Debug("foobar")
	if "" != buffer.String() {
		t.Error(buffer.String())
	}
	buffer.Reset()
	logger.Info("foobar")
	if "" != buffer.String() {
		t.Error(buffer.String())
	}
	buffer.Reset()
	logger.Notice("foobar")
	if "test.NOTICE: foobar" != buffer.String() {
		t.Error(buffer.String())
	}
	buffer.Reset()
	logger.Warning("foobar")
	if "test.WARNING: foobar" != buffer.String() {
		t.Error(buffer.String())
	}
	buffer.Reset()
	logger.Error("foobar")
	if "test.ERROR: foobar" != buffer.String() {
		t.Error(buffer.String())
	}
	buffer.Reset()
	logger.Critical("foobar")
	if "test.CRITICAL: foobar" != buffer.String() {
		t.Error(buffer.String())
	}
	buffer.Reset()
	logger.Alert("foobar")
	if "test.ALERT: foobar" != buffer.String() {
		t.Error(buffer.String())
	}
	buffer.Reset()
	logger.Emergency("foobar")
	if "test.EMERGENCY: foobar" != buffer.String() {
		t.Error(buffer.String())
	}
}

func TestLoggerPushPopProcessors(t *testing.T) {
	p1 := NewProcessor(func(r *Record) {}) // Will be called last
	p2 := NewProcessor(func(r *Record) {})

	l := NewLogger("test", []Handler{}, []Processor{})
	l.PushProcessor(p1)
	l.PushProcessor(p2)

	if len(l.processors) != 2 {
		t.Error()
	}
	if l.processors[0] != p2 {
		t.Error()
	}
	if l.processors[1] != p1 {
		t.Error()
	}
	l.PopProcessor()
	if len(l.processors) != 1 {
		t.Error()
	}
	if l.processors[0] != p1 {
		t.Error()
	}
}

func TestLoggerPushPopHandler(t *testing.T) {
	buffer := new(bytes.Buffer)
	h1 := NewBufferHandler(buffer, DEBUG)
	h2 := NewBufferHandler(buffer, DEBUG)
	l := &Logger{}
	l.PushHandler(h1)
	l.PushHandler(h2)

	if len(l.handlers) != 2 {
		t.Error()
	}
	if l.handlers[0] != h2 {
		t.Error()
	}
	if l.handlers[1] != h1 {
		t.Error()
	}
	l.PopHandler()
	if len(l.handlers) != 1 {
		t.Error()
	}
	if l.handlers[0] != h1 {
		t.Error()
	}
}

func TestLoggerWithDefaultProcessor(t *testing.T) {
	buffer := new(bytes.Buffer)
	h1 := NewBufferHandler(buffer, DEBUG)
	h1.SetFormatter(&LineFormatter{"%message% %extra%"})
	l := NewLogger("channel", []Handler{}, []Processor{})
	l.PushHandler(h1)
	p := NewProcessor(func(r *Record) {
		r.Extra["go.date"] = time.Date(2009, 11, 10, 23, 0, 0, 0, time.UTC)
	})
	l.PushProcessor(p)
	l.Debug("foobar")
	if "foobar map[go.date:2009-11-10T23:00:00Z]" != buffer.String() {
		t.Error()
	}
}
<commit_msg>I can't english<commit_after>// Copyright 2013 Marc Weistroff. All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package gogol

import (
	"bytes"
	"testing"
	"time"
)

func ExampleLogger() {
	handler := NewStdoutHandler(WARNING) // Will log all messages where severity is >= WARNING
	handler.SetFormatter(NewMinimalLineFormatter())
	logger := NewLogger("example", []Handler{handler}, []Processor{})
	logger.Debug("Debug message that won't be displayed")
	logger.Warning("I sense a disturbance in the force")
	logger.Error("This is an error")
	// Output:
	// example.WARNING: I sense a disturbance in the force
	// example.ERROR: This is an error
}

func TestLogger(t *testing.T) {
	buffer := new(bytes.Buffer)
	handler := NewBufferHandler(buffer, NOTICE)
	handler.SetFormatter(&LineFormatter{"%channel%.%level_name%: %message%"})
	logger := NewLogger("test", []Handler{handler}, []Processor{})

	logger.Debug("foobar")
	if "" != buffer.String() {
		t.Error(buffer.String())
	}
	buffer.Reset()
	logger.Info("foobar")
	if "" != buffer.String() {
		t.Error(buffer.String())
	}
	buffer.Reset()
	logger.Notice("foobar")
	if "test.NOTICE: foobar" != buffer.String() {
		t.Error(buffer.String())
	}
	buffer.Reset()
	logger.Warning("foobar")
	if "test.WARNING: foobar" != buffer.String() {
		t.Error(buffer.String())
	}
	buffer.Reset()
	logger.Error("foobar")
	if "test.ERROR: foobar" != buffer.String() {
		t.Error(buffer.String())
	}
	buffer.Reset()
	logger.Critical("foobar")
	if "test.CRITICAL: foobar" != buffer.String() {
		t.Error(buffer.String())
	}
	buffer.Reset()
	logger.Alert("foobar")
	if "test.ALERT: foobar" != buffer.String() {
		t.Error(buffer.String())
	}
	buffer.Reset()
	logger.Emergency("foobar")
	if "test.EMERGENCY: foobar" != buffer.String() {
		t.Error(buffer.String())
	}
}

func TestLoggerPushPopProcessors(t *testing.T) {
	p1 := NewProcessor(func(r *Record) {}) // Will be called last
	p2 := NewProcessor(func(r *Record) {})

	l := NewLogger("test", []Handler{}, []Processor{})
	l.PushProcessor(p1)
	l.PushProcessor(p2)

	if len(l.processors) != 2 {
		t.Error()
	}
	if l.processors[0] != p2 {
		t.Error()
	}
	if l.processors[1] != p1 {
		t.Error()
	}
	l.PopProcessor()
	if len(l.processors) != 1 {
		t.Error()
	}
	if l.processors[0] != p1 {
		t.Error()
	}
}

func TestLoggerPushPopHandler(t *testing.T) {
	buffer := new(bytes.Buffer)
	h1 := NewBufferHandler(buffer, DEBUG)
	h2 := NewBufferHandler(buffer, DEBUG)
	l := &Logger{}
	l.PushHandler(h1)
	l.PushHandler(h2)

	if len(l.handlers) != 2 {
		t.Error()
	}
	if l.handlers[0] != h2 {
		t.Error()
	}
	if l.handlers[1] != h1 {
		t.Error()
	}
	l.PopHandler()
	if len(l.handlers) != 1 {
		t.Error()
	}
	if l.handlers[0] != h1 {
		t.Error()
	}
}

func TestLoggerWithDefaultProcessor(t *testing.T) {
	buffer := new(bytes.Buffer)
	h1 := NewBufferHandler(buffer, DEBUG)
	h1.SetFormatter(&LineFormatter{"%message% %extra%"})
	l := NewLogger("channel", []Handler{}, []Processor{})
	l.PushHandler(h1)
	p := NewProcessor(func(r *Record) {
		r.Extra["go.date"] = time.Date(2009, 11, 10, 23, 0, 0, 0, time.UTC)
	})
	l.PushProcessor(p)
	l.Debug("foobar")
	if "foobar map[go.date:2009-11-10T23:00:00Z]" != buffer.String() {
		t.Error()
	}
}
<|endoftext|>


<commit_before>package logging

import (
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/Sirupsen/logrus"
	"gopkg.in/errgo.v1"
)

var closeChan chan<- os.Signal

func SetCloseChan(c chan<- os.Signal) {
	closeChan = c
}

// CheckFatal exits the process if err != nil
func CheckFatal(err error) {
	if err != nil {
		pc, file, line, ok := runtime.Caller(1)
		if !ok {
			file = "?"
			line = 0
		}
		fn := runtime.FuncForPC(pc)
		var fnName string
		if fn == nil {
			fnName = "?()"
		} else {
			dotName := filepath.Ext(fn.Name())
			fnName = strings.TrimLeft(dotName, ".") + "()"
		}
		logrus.Errorf("%s:%d %s", file, line, fnName)
		logrus.Error("Fatal Error:", errgo.Details(err))
		if closeChan != nil {
			logrus.Warn("Sending close message")
			closeChan <- os.Interrupt
		}
		os.Exit(1)
	}
}

// SetupLogging will initialize the logger backend and set the flags.
func SetupLogging(w io.Writer) {
	if w != nil {
		logrus.SetOutput(io.MultiWriter(os.Stderr, w))
	}
	if runtime.GOOS == "windows" { // colored ttys are rare on windows...
		logrus.SetFormatter(&logrus.TextFormatter{DisableColors: true})
	}
	if lvl := os.Getenv("CRYPTIX_LOGLVL"); lvl != "" {
		l, err := logrus.ParseLevel(lvl)
		if err != nil {
			l = logrus.DebugLevel
		}
		logrus.SetLevel(l)
	}
}

// Logger returns a logger where the module field is set to name
// https://github.com/Sirupsen/logrus/issues/144
func Logger(name string) *logrus.Entry {
	if len(name) == 0 {
		logrus.Warnf("missing name parameter")
		name = "undefined"
	}
	return logrus.WithField("module", name)
}
<commit_msg>logging: expose underlying logrus logger for modifications<commit_after>package logging

import (
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/Sirupsen/logrus"
	"gopkg.in/errgo.v1"
)

var closeChan chan<- os.Signal

func SetCloseChan(c chan<- os.Signal) {
	closeChan = c
}

// CheckFatal exits the process if err != nil
func CheckFatal(err error) {
	if err != nil {
		pc, file, line, ok := runtime.Caller(1)
		if !ok {
			file = "?"
			line = 0
		}
		fn := runtime.FuncForPC(pc)
		var fnName string
		if fn == nil {
			fnName = "?()"
		} else {
			dotName := filepath.Ext(fn.Name())
			fnName = strings.TrimLeft(dotName, ".") + "()"
		}
		logrus.Errorf("%s:%d %s", file, line, fnName)
		logrus.Error("Fatal Error:", errgo.Details(err))
		if closeChan != nil {
			logrus.Warn("Sending close message")
			closeChan <- os.Interrupt
		}
		os.Exit(1)
	}
}

var Underlying *logrus.Logger

// SetupLogging will initialize the logger backend and set the flags.
func SetupLogging(w io.Writer) {
	if w != nil {
		logrus.SetOutput(io.MultiWriter(os.Stderr, w))
	}
	if runtime.GOOS == "windows" { // colored ttys are rare on windows...
		logrus.SetFormatter(&logrus.TextFormatter{DisableColors: true})
	}
	if lvl := os.Getenv("CRYPTIX_LOGLVL"); lvl != "" {
		l, err := logrus.ParseLevel(lvl)
		if err != nil {
			logrus.Errorf("logging: could not parse lvl from env, defaulting to debug: %s", err)
			l = logrus.DebugLevel
		}
		logrus.SetLevel(l)
	}
	Underlying = logrus.New()
}

// Logger returns an Entry where the module field is set to name
func Logger(name string) *logrus.Entry {
	if name == "" {
		Underlying.Warnf("missing name parameter")
		name = "undefined"
	}
	return Underlying.WithField("module", name)
}
<|endoftext|>


<commit_before>package stats

import (
	"os"
	"runtime"
	"time"
)

func StartReportingMemoryAndGC(reporter Statter, d time.Duration) {
	ticker := time.Tick(d)
	for {
		select {
		case <-ticker:
			var ms runtime.MemStats
			runtime.ReadMemStats(&ms)

			prefix := "runtime"

			reporter.Measure(prefix, "allocated", int64(ms.Alloc), 1.0)
			reporter.Measure(prefix, "allocated.heap", int64(ms.HeapAlloc), 1.0)
			reporter.Time(prefix, "gc.pause", time.Duration(ms.PauseNs[(ms.NumGC+255)%256]), 1.0)

			// GC CPU percentage.
			reporter.Measure(prefix, "gc.cpufraction", int64(ms.GCCPUFraction*100), 1.0)
		}
	}
}
<commit_msg>fix build<commit_after>package stats

import (
	"runtime"
	"time"
)

func StartReportingMemoryAndGC(reporter Statter, d time.Duration) {
	ticker := time.Tick(d)
	for {
		select {
		case <-ticker:
			var ms runtime.MemStats
			runtime.ReadMemStats(&ms)

			prefix := "runtime"

			reporter.Measure(prefix, "allocated", int64(ms.Alloc), 1.0)
			reporter.Measure(prefix, "allocated.heap", int64(ms.HeapAlloc), 1.0)
			reporter.Time(prefix, "gc.pause", time.Duration(ms.PauseNs[(ms.NumGC+255)%256]), 1.0)

			// GC CPU percentage.
			reporter.Measure(prefix, "gc.cpufraction", int64(ms.GCCPUFraction*100), 1.0)
		}
	}
}
<|endoftext|>


<commit_before>package lzma

import (
	"io"

	"github.com/uli-go/xz/lzbase"
)

type Writer struct {
	lzbase.Writer
	params *Parameters
}

func NewWriter(w io.Writer) (lw *Writer, err error) {
	return NewWriterP(w, Default)
}

func NewWriterP(w io.Writer, p Parameters) (lw *Writer, err error) {
	if w == nil {
		return nil, newError("writer argument w is nil")
	}
	normalizeSizes(&p)
	if err = verifyParameters(&p); err != nil {
		return nil, err
	}
	if p.Size == 0 && !p.SizeInHeader {
		p.EOS = true
	}
	if err = writeHeader(w, &p); err != nil {
		return nil, err
	}
	dict, err := lzbase.NewWriterDict(p.DictSize, p.BufferSize)
	if err != nil {
		return nil, err
	}
	oc := lzbase.NewOpCodec(p.Properties(), dict)
	lw = new(Writer)
	if err = lzbase.InitWriter(&lw.Writer, w, oc,
		lzbase.Parameters{
			SizeInHeader: p.SizeInHeader,
			Size:         p.Size,
			EOS:          p.EOS}); err != nil {
		return nil, err
	}
	return lw, nil
}
<commit_msg>lzma: added documentation comments to Writer<commit_after>package lzma

import (
	"io"

	"github.com/uli-go/xz/lzbase"
)

// Writer supports the LZMA compression of a file.
//
// Using an arithmetic coder it cannot support flushing. A writer must be
// closed.
type Writer struct {
	lzbase.Writer
	params *Parameters
}

// NewWriter creates a new writer. It writes the LZMA header. It will use the
// Default Parameters.
//
// Don't forget to call Close() for the writer after all data has been written.
//
// For high performance use a buffered writer. But be aware that Close will not
// flush it.
func NewWriter(w io.Writer) (lw *Writer, err error) {
	return NewWriterP(w, Default)
}

// NewWriterP creates a new writer with the given Parameters. It writes the
// LZMA header.
//
// Don't forget to call Close() for the writer after all data has been written.
//
// For high performance use a buffered writer. But be aware that Close will not
// flush it.
func NewWriterP(w io.Writer, p Parameters) (lw *Writer, err error) {
	if w == nil {
		return nil, newError("writer argument w is nil")
	}
	normalizeSizes(&p)
	if err = verifyParameters(&p); err != nil {
		return nil, err
	}
	if p.Size == 0 && !p.SizeInHeader {
		p.EOS = true
	}
	if err = writeHeader(w, &p); err != nil {
		return nil, err
	}
	dict, err := lzbase.NewWriterDict(p.DictSize, p.BufferSize)
	if err != nil {
		return nil, err
	}
	oc := lzbase.NewOpCodec(p.Properties(), dict)
	lw = new(Writer)
	if err = lzbase.InitWriter(&lw.Writer, w, oc,
		lzbase.Parameters{
			SizeInHeader: p.SizeInHeader,
			Size:         p.Size,
			EOS:          p.EOS}); err != nil {
		return nil, err
	}
	return lw, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"time"
)

// Console flags
var (
	listen                = flag.String("l", ":8888", "port to accept requests")
	targetProduction      = flag.String("a", "localhost:8080", "where production traffic goes. http://localhost:8080/production")
	altTarget             = flag.String("b", "localhost:8081", "where testing traffic goes. response are skipped. http://localhost:8081/test")
	debug                 = flag.Bool("debug", false, "more logging, showing ignored output")
	productionTimeout     = flag.Int("a.timeout", 2500, "timeout in milliseconds for production traffic")
	alternateTimeout      = flag.Int("b.timeout", 1000, "timeout in milliseconds for alternate site traffic")
	productionHostRewrite = flag.Bool("a.rewrite", false, "rewrite the host header when proxying production traffic")
	alternateHostRewrite  = flag.Bool("b.rewrite", false, "rewrite the host header when proxying alternate site traffic")
	percent               = flag.Float64("p", 100.0, "float64 percentage of traffic to send to testing")
	tlsPrivateKey         = flag.String("key.file", "", "path to the TLS private key file")
	tlsCertificate        = flag.String("cert.file", "", "path to the TLS certificate file")
	forwardClientIP       = flag.Bool("forward-client-ip", false, "enable forwarding of the client IP to the backend using the 'X-Forwarded-For' and 'Forwarded' headers")
	closeConnections      = flag.Bool("close-connections", false, "close connections to the clients and backends")
)

// Sets the request URL.
//
// This turns a inbound request (a request without URL) into an outbound request.
func setRequestTarget(request *http.Request, target *string) {
	URL, err := url.Parse("http://" + *target + request.URL.String())
	if err != nil {
		log.Println(err)
	}
	request.URL = URL
}

// Sends a request and returns the response.
func handleRequest(request *http.Request, timeout time.Duration) *http.Response {
	transport := &http.Transport{
		// NOTE(girone): DialTLS is not needed here, because the teeproxy works
		// as an SSL terminator.
		Dial: (&net.Dialer{ // go1.8 deprecated: Use DialContext instead
			Timeout:   timeout,
			KeepAlive: 10 * timeout,
		}).Dial,
		// Close connections to the production and alternative servers?
		DisableKeepAlives: *closeConnections,
		//IdleConnTimeout: timeout,  // go1.8
		TLSHandshakeTimeout:   timeout,
		ResponseHeaderTimeout: timeout,
		ExpectContinueTimeout: timeout,
	}
	// Do not use http.Client here, because it's higher level and processes
	// redirects internally, which is not what we want.
	//client := &http.Client{
	//	Timeout: timeout,
	//	Transport: transport,
	//}
	//response, err := client.Do(request)
	response, err := transport.RoundTrip(request)
	if err != nil {
		log.Println("Request failed:", err)
	}
	return response
}

// handler contains the address of the main Target and the one for the Alternative target
type handler struct {
	Target      string
	Alternative string
	Randomizer  rand.Rand
}

// ServeHTTP duplicates the incoming request (req) and does the request to the
// Target and the Alternate target discading the Alternate response
func (h handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var productionRequest, alternativeRequest *http.Request
	if *forwardClientIP {
		updateForwardedHeaders(req)
	}
	if *percent == 100.0 || h.Randomizer.Float64()*100 < *percent {
		alternativeRequest, productionRequest = DuplicateRequest(req)
		go func() {
			defer func() {
				if r := recover(); r != nil && *debug {
					log.Println("Recovered in ServeHTTP(alternate request) from:", r)
				}
			}()

			setRequestTarget(alternativeRequest, altTarget)

			if *alternateHostRewrite {
				alternativeRequest.Host = h.Alternative
			}

			timeout := time.Duration(*alternateTimeout) * time.Millisecond
			// This keeps responses from the alternative target away from the outside world.
			alternateResponse := handleRequest(alternativeRequest, timeout)
			if alternateResponse != nil {
				// NOTE(girone): Even though we do not care about the second
				// response, we still need to close the Body reader. Otherwise
				// the connection stays open and we would soon run out of file
				// descriptors.
				alternateResponse.Body.Close()
			}
		}()
	} else {
		productionRequest = req
	}
	defer func() {
		if r := recover(); r != nil && *debug {
			log.Println("Recovered in ServeHTTP(production request) from:", r)
		}
	}()

	setRequestTarget(productionRequest, targetProduction)

	if *productionHostRewrite {
		productionRequest.Host = h.Target
	}

	timeout := time.Duration(*productionTimeout) * time.Millisecond
	resp := handleRequest(productionRequest, timeout)

	if resp != nil {
		defer resp.Body.Close()

		// Forward response headers.
		for k, v := range resp.Header {
			w.Header()[k] = v
		}
		w.WriteHeader(resp.StatusCode)

		// Forward response body.
		io.Copy(w, resp.Body)
	}
}

func main() {
	flag.Parse()

	log.Printf("Starting teeproxy at %s sending to A: %s and B: %s",
		*listen, *targetProduction, *altTarget)

	runtime.GOMAXPROCS(runtime.NumCPU())

	var err error

	var listener net.Listener

	if len(*tlsPrivateKey) > 0 {
		cer, err := tls.LoadX509KeyPair(*tlsCertificate, *tlsPrivateKey)
		if err != nil {
			log.Fatalf("Failed to load certficate: %s and private key: %s", *tlsCertificate, *tlsPrivateKey)
		}

		config := &tls.Config{Certificates: []tls.Certificate{cer}}
		listener, err = tls.Listen("tcp", *listen, config)
		if err != nil {
			log.Fatalf("Failed to listen to %s: %s", *listen, err)
		}
	} else {
		listener, err = net.Listen("tcp", *listen)
		if err != nil {
			log.Fatalf("Failed to listen to %s: %s", *listen, err)
		}
	}

	h := handler{
		Target:      *targetProduction,
		Alternative: *altTarget,
		Randomizer:  *rand.New(rand.NewSource(time.Now().UnixNano())),
	}

	server := &http.Server{
		Handler: h,
	}
	if *closeConnections {
		// Close connections to clients by setting the "Connection": "close" header in the response.
		server.SetKeepAlivesEnabled(false)
	}
	server.Serve(listener)
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() error { return nil }

func DuplicateRequest(request *http.Request) (request1 *http.Request, request2 *http.Request) {
	b1 := new(bytes.Buffer)
	b2 := new(bytes.Buffer)
	w := io.MultiWriter(b1, b2)
	io.Copy(w, request.Body)
	defer request.Body.Close()
	request1 = &http.Request{
		Method:        request.Method,
		URL:           request.URL,
		Proto:         request.Proto,
		ProtoMajor:    request.ProtoMajor,
		ProtoMinor:    request.ProtoMinor,
		Header:        request.Header,
		Body:          nopCloser{b1},
		Host:          request.Host,
		ContentLength: request.ContentLength,
		Close:         true,
	}
	request2 = &http.Request{
		Method:        request.Method,
		URL:           request.URL,
		Proto:         request.Proto,
		ProtoMajor:    request.ProtoMajor,
		ProtoMinor:    request.ProtoMinor,
		Header:        request.Header,
		Body:          nopCloser{b2},
		Host:          request.Host,
		ContentLength: request.ContentLength,
		Close:         true,
	}
	return
}

func updateForwardedHeaders(request *http.Request) {
	positionOfColon := strings.LastIndex(request.RemoteAddr, ":")
	var remoteIP string
	if positionOfColon != -1 {
		remoteIP = request.RemoteAddr[:positionOfColon]
	} else {
		log.Printf("The default format of request.RemoteAddr should be IP:Port but was %s\n", remoteIP)
		remoteIP = request.RemoteAddr
	}
	insertOrExtendForwardedHeader(request, remoteIP)
	insertOrExtendXFFHeader(request, remoteIP)
}

const XFF_HEADER = "X-Forwarded-For"

func insertOrExtendXFFHeader(request *http.Request, remoteIP string) {
	header := request.Header.Get(XFF_HEADER)
	if header != "" {
		// extend
		request.Header.Set(XFF_HEADER, header+", "+remoteIP)
	} else {
		// insert
		request.Header.Set(XFF_HEADER, remoteIP)
	}
}

const FORWARDED_HEADER = "Forwarded"

// Implementation according to rfc7239
func insertOrExtendForwardedHeader(request *http.Request, remoteIP string) {
	extension := "for=" + remoteIP
	header := request.Header.Get(FORWARDED_HEADER)
	if header != "" {
		// extend
		request.Header.Set(FORWARDED_HEADER, header+", "+extension)
	} else {
		// insert
		request.Header.Set(FORWARDED_HEADER, extension)
	}
}
<commit_msg>fix compilation<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"flag"
	"io"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"strings"
	"time"
)

// Console flags
var (
	listen                = flag.String("l", ":8888", "port to accept requests")
	targetProduction      = flag.String("a", "localhost:8080", "where production traffic goes. http://localhost:8080/production")
	altTarget             = flag.String("b", "localhost:8081", "where testing traffic goes. response are skipped. http://localhost:8081/test")
	debug                 = flag.Bool("debug", false, "more logging, showing ignored output")
	productionTimeout     = flag.Int("a.timeout", 2500, "timeout in milliseconds for production traffic")
	alternateTimeout      = flag.Int("b.timeout", 1000, "timeout in milliseconds for alternate site traffic")
	productionHostRewrite = flag.Bool("a.rewrite", false, "rewrite the host header when proxying production traffic")
	alternateHostRewrite  = flag.Bool("b.rewrite", false, "rewrite the host header when proxying alternate site traffic")
	percent               = flag.Float64("p", 100.0, "float64 percentage of traffic to send to testing")
	tlsPrivateKey         = flag.String("key.file", "", "path to the TLS private key file")
	tlsCertificate        = flag.String("cert.file", "", "path to the TLS certificate file")
	forwardClientIP       = flag.Bool("forward-client-ip", false, "enable forwarding of the client IP to the backend using the 'X-Forwarded-For' and 'Forwarded' headers")
	closeConnections      = flag.Bool("close-connections", false, "close connections to the clients and backends")
)

// Sets the request URL.
//
// This turns a inbound request (a request without URL) into an outbound request.
func setRequestTarget(request *http.Request, target *string) {
	URL, err := url.Parse("http://" + *target + request.URL.String())
	if err != nil {
		log.Println(err)
	}
	request.URL = URL
}

// Sends a request and returns the response.
func handleRequest(request *http.Request, timeout time.Duration) *http.Response {
	transport := &http.Transport{
		// NOTE(girone): DialTLS is not needed here, because the teeproxy works
		// as an SSL terminator.
		Dial: (&net.Dialer{ // go1.8 deprecated: Use DialContext instead
			Timeout:   timeout,
			KeepAlive: 10 * timeout,
		}).Dial,
		// Close connections to the production and alternative servers?
		DisableKeepAlives: *closeConnections,
		//IdleConnTimeout: timeout,  // go1.8
		TLSHandshakeTimeout:   timeout,
		ResponseHeaderTimeout: timeout,
		ExpectContinueTimeout: timeout,
	}
	// Do not use http.Client here, because it's higher level and processes
	// redirects internally, which is not what we want.
	//client := &http.Client{
	//	Timeout: timeout,
	//	Transport: transport,
	//}
	//response, err := client.Do(request)
	response, err := transport.RoundTrip(request)
	if err != nil {
		log.Println("Request failed:", err)
	}
	return response
}

// handler contains the address of the main Target and the one for the Alternative target
type handler struct {
	Target      string
	Alternative string
	Randomizer  rand.Rand
}

// ServeHTTP duplicates the incoming request (req) and does the request to the
// Target and the Alternate target discading the Alternate response
func (h handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var productionRequest, alternativeRequest *http.Request
	if *forwardClientIP {
		updateForwardedHeaders(req)
	}
	if *percent == 100.0 || h.Randomizer.Float64()*100 < *percent {
		alternativeRequest, productionRequest = DuplicateRequest(req)
		go func() {
			defer func() {
				if r := recover(); r != nil && *debug {
					log.Println("Recovered in ServeHTTP(alternate request) from:", r)
				}
			}()

			setRequestTarget(alternativeRequest, altTarget)

			if *alternateHostRewrite {
				alternativeRequest.Host = h.Alternative
			}

			timeout := time.Duration(*alternateTimeout) * time.Millisecond
			// This keeps responses from the alternative target away from the outside world.
			alternateResponse := handleRequest(alternativeRequest, timeout)
			if alternateResponse != nil {
				// NOTE(girone): Even though we do not care about the second
				// response, we still need to close the Body reader. Otherwise
				// the connection stays open and we would soon run out of file
				// descriptors.
				alternateResponse.Body.Close()
			}
		}()
	} else {
		productionRequest = req
	}
	defer func() {
		if r := recover(); r != nil && *debug {
			log.Println("Recovered in ServeHTTP(production request) from:", r)
		}
	}()

	setRequestTarget(productionRequest, targetProduction)

	if *productionHostRewrite {
		productionRequest.Host = h.Target
	}

	timeout := time.Duration(*productionTimeout) * time.Millisecond
	resp := handleRequest(productionRequest, timeout)

	if resp != nil {
		defer resp.Body.Close()

		// Forward response headers.
		for k, v := range resp.Header {
			w.Header()[k] = v
		}
		w.WriteHeader(resp.StatusCode)

		// Forward response body.
		io.Copy(w, resp.Body)
	}
}

func main() {
	flag.Parse()

	log.Printf("Starting teeproxy at %s sending to A: %s and B: %s",
		*listen, *targetProduction, *altTarget)

	runtime.GOMAXPROCS(runtime.NumCPU())

	var err error

	var listener net.Listener

	if len(*tlsPrivateKey) > 0 {
		cer, err := tls.LoadX509KeyPair(*tlsCertificate, *tlsPrivateKey)
		if err != nil {
			log.Fatalf("Failed to load certficate: %s and private key: %s", *tlsCertificate, *tlsPrivateKey)
		}

		config := &tls.Config{Certificates: []tls.Certificate{cer}}
		listener, err = tls.Listen("tcp", *listen, config)
		if err != nil {
			log.Fatalf("Failed to listen to %s: %s", *listen, err)
		}
	} else {
		listener, err = net.Listen("tcp", *listen)
		if err != nil {
			log.Fatalf("Failed to listen to %s: %s", *listen, err)
		}
	}

	h := handler{
		Target:      *targetProduction,
		Alternative: *altTarget,
		Randomizer:  *rand.New(rand.NewSource(time.Now().UnixNano())),
	}

	server := &http.Server{
		Handler: h,
	}
	if *closeConnections {
		// Close connections to clients by setting the "Connection": "close" header in the response.
		server.SetKeepAlivesEnabled(false)
	}
	server.Serve(listener)
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() error { return nil }

func DuplicateRequest(request *http.Request) (request1 *http.Request, request2 *http.Request) {
	b1 := new(bytes.Buffer)
	b2 := new(bytes.Buffer)
	w := io.MultiWriter(b1, b2)
	io.Copy(w, request.Body)
	defer request.Body.Close()
	request1 = &http.Request{
		Method:        request.Method,
		URL:           request.URL,
		Proto:         request.Proto,
		ProtoMajor:    request.ProtoMajor,
		ProtoMinor:    request.ProtoMinor,
		Header:        request.Header,
		Body:          nopCloser{b1},
		Host:          request.Host,
		ContentLength: request.ContentLength,
		Close:         true,
	}
	request2 = &http.Request{
		Method:        request.Method,
		URL:           request.URL,
		Proto:         request.Proto,
		ProtoMajor:    request.ProtoMajor,
		ProtoMinor:    request.ProtoMinor,
		Header:        request.Header,
		Body:          nopCloser{b2},
		Host:          request.Host,
		ContentLength: request.ContentLength,
		Close:         true,
	}
	return
}

func updateForwardedHeaders(request *http.Request) {
	positionOfColon := strings.LastIndex(request.RemoteAddr, ":")
	var remoteIP string
	if positionOfColon != -1 {
		remoteIP = request.RemoteAddr[:positionOfColon]
	} else {
		log.Printf("The default format of request.RemoteAddr should be IP:Port but was %s\n", remoteIP)
		remoteIP = request.RemoteAddr
	}
	insertOrExtendForwardedHeader(request, remoteIP)
	insertOrExtendXFFHeader(request, remoteIP)
}

const XFF_HEADER = "X-Forwarded-For"

func insertOrExtendXFFHeader(request *http.Request, remoteIP string) {
	header := request.Header.Get(XFF_HEADER)
	if header != "" {
		// extend
		request.Header.Set(XFF_HEADER, header+", "+remoteIP)
	} else {
		// insert
		request.Header.Set(XFF_HEADER, remoteIP)
	}
}

const FORWARDED_HEADER = "Forwarded"

// Implementation according to rfc7239
func insertOrExtendForwardedHeader(request *http.Request, remoteIP string) {
	extension := "for=" + remoteIP
	header := request.Header.Get(FORWARDED_HEADER)
	if header != "" {
		// extend
		request.Header.Set(FORWARDED_HEADER, header+", "+extension)
	} else {
		// insert
		request.Header.Set(FORWARDED_HEADER, extension)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"math/rand"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/go-telegram-bot-api/telegram-bot-api"
	bt "github.com/ikool-cn/gobeanstalk-connection-pool"
)

var (
	telegramBot *TelegramBot
)

// TelegramBot ...
type TelegramBot struct {
	Name          string
	SelfChatID    int64
	ChannelChatID int64
	ComicPath     string
	DeleteDelay   time.Duration
	Client        *tgbotapi.BotAPI
	Queue         *bt.Pool
	Tube          string
}

// NewTelegramBot ...
func NewTelegramBot(cfg *TelegramConfig, btdAddr string) (t *TelegramBot) {
	bot, err := tgbotapi.NewBotAPI(cfg.Token)
	if err != nil {
		logger.Panicf("tg bot init failed: %+v", err)
	}
	delay, err := time.ParseDuration(cfg.DeleteDelay)
	if err != nil {
		logger.Panicf("delete delay error: %+v", err)
	}

	t = &TelegramBot{
		Name:          bot.Self.UserName,
		SelfChatID:    cfg.SelfChatID,
		ChannelChatID: cfg.ChannelChatID,
		ComicPath:     cfg.ComicPath,
		DeleteDelay:   delay,
		Client:        bot,
		Tube:          "tg",
	}
	t.Queue = &bt.Pool{
		Dial: func() (*bt.Conn, error) {
			return bt.Dial(btdAddr)
		},
		MaxIdle:     10,
		MaxActive:   100,
		IdleTimeout: 60 * time.Second,
		MaxLifetime: 180 * time.Second,
		Wait:        true,
	}
	return
}

func (t *TelegramBot) putQueue(msg []byte) {
	conn, err := t.Queue.Get()
	if err != nil {
		logger.Errorf("%+v: %s", err, string(msg))
		return
	}
	conn.Use(t.Tube)
	_, err = conn.Put(msg, 1, t.DeleteDelay, time.Minute)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
}

func (t *TelegramBot) sendFile(chat int64, file string) (tgbotapi.Message, error) {
	logger.Debugf("[%d]%s", chat, file)
	return t.Client.Send(tgbotapi.NewDocumentUpload(chat, file))
}
func (t *TelegramBot) sendPic(chat int64, file string) (tgbotapi.Message, error) {
	logger.Debugf("[%d]%s", chat, file)
	if strings.HasSuffix(file, ".mp4") {
		return t.Client.Send(tgbotapi.NewVideoUpload(chat, file))
	}
	return t.Client.Send(tgbotapi.NewPhotoUpload(chat, file))
}

func (t *TelegramBot) send(chat int64, msg string) (tgbotapi.Message, error) {
	logger.Debugf("[%d]%s", chat, msg)
	return t.Client.Send(tgbotapi.NewMessage(chat, msg))
}

func (t *TelegramBot) delMessage() {
	for {
		conn, err := t.Queue.Get()
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		conn.Watch(t.Tube)
		job, err := conn.Reserve()
		if err != nil {
			logger.Warningf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		msg := &tgbotapi.Message{}
		err = json.Unmarshal(job.Body, msg)
		if err != nil {
			logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		delMsg := tgbotapi.DeleteMessageConfig{
			ChatID:    msg.Chat.ID,
			MessageID: msg.MessageID,
		}
		logger.Infof(":[%s]{%s}", getMsgTitle(msg), strconv.Quote(msg.Text))

		_, err = t.Client.DeleteMessage(delMsg)
		if err != nil {
			logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		err = conn.Delete(job.ID)
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
		}
		t.Queue.Release(conn, false)
	}
}

func (t *TelegramBot) tgBot() {
	u := tgbotapi.NewUpdate(0)
	u.Timeout = 60
	for {
		updates, err := t.Client.GetUpdatesChan(u)
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		var message *tgbotapi.Message
		for update := range updates {
			if update.Message != nil {
				message = update.Message
			} else if update.EditedMessage != nil {
				message = update.EditedMessage
			} else {
				// unkown msg type
				continue
			}
			if message.Chat.IsGroup() {
				logger.Infof(
					"recv:(%s)[%s]{%s}",
					message.Chat.Title,
					message.From.String(),
					strconv.Quote(message.Text))
			} else {
				logger.Infof(
					"recv:[%s]{%s}",
					message.From.String(),
					strconv.Quote(message.Text),
				)
			}

			if message.IsCommand() {
				switch message.Command() {
				case "start":
					go onStart(t, message)
				case "comic":
					go onComic(t, message)
				case "pic":
					go onPic(t, message)
				default:
					logger.Infof("ignore unkown cmd: %+v", message.Command())
					continue
				}
			} else {
				if message.Text == "" {
					continue
				}
				checkRepeat(t, message)
			}
		}
		logger.Warning("tg bot restarted.")
		time.Sleep(3 * time.Second)
	}
}

func checkRepeat(t *TelegramBot, message *tgbotapi.Message) {
	key := "tg_" + getMsgTitle(message) + "_last"
	flattendMsg := strings.TrimSpace(message.Text)
	defer redisClient.LTrim(key, 0, 10)
	defer redisClient.LPush(key, flattendMsg)

	lastMsgs, err := redisClient.LRange(key, 0, 6).Result()
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	i := 0
	for _, s := range lastMsgs {
		if s == flattendMsg {
			i++
		}
	}
	if i > 1 {
		redisClient.Del(key)
		logger.Infof("repeat: %s", strconv.Quote(message.Text))
		msg := tgbotapi.NewMessage(message.Chat.ID, message.Text)
		t.Client.Send(msg)
	}
}

func onStart(t *TelegramBot, message *tgbotapi.Message) {
	msg := tgbotapi.NewMessage(message.Chat.ID, "呀呀呀")
	msg.ReplyToMessageID = message.MessageID
	t.Client.Send(msg)
}

func onComic(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(t.ComicPath)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	rand.Seed(time.Now().Unix())
	file := files[rand.Intn(len(files))]
	number := strings.Split(strings.Split(file, "@")[1], ".")[0]
	msg := tgbotapi.NewMessage(message.Chat.ID, "https://nhentai.net/g/"+number)

	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))
	msgSent, err := t.Client.Send(msg)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	data, err := json.Marshal(msgSent)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	t.putQueue(data)
}

func onPic(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(twitterBot.ImgPath + "/*")
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	if files == nil {
		logger.Error("find no pics")
	}
	rand.Seed(time.Now().Unix())
	file := files[rand.Intn(len(files))]

	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))
	msgSent, err := t.sendFile(message.Chat.ID, file)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	data, err := json.Marshal(msgSent)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	t.putQueue(data)
}

func getMsgTitle(m *tgbotapi.Message) string {
	if m.Chat.IsGroup() {
		return m.Chat.Title
	}
	return m.From.String()
}
<commit_msg>do not auto delete pictures<commit_after>package main

import (
	"encoding/json"
	"math/rand"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/go-telegram-bot-api/telegram-bot-api"
	bt "github.com/ikool-cn/gobeanstalk-connection-pool"
)

var (
	telegramBot *TelegramBot
)

// TelegramBot ...
type TelegramBot struct {
	Name          string
	SelfChatID    int64
	ChannelChatID int64
	ComicPath     string
	DeleteDelay   time.Duration
	Client        *tgbotapi.BotAPI
	Queue         *bt.Pool
	Tube          string
}

// NewTelegramBot ...
func NewTelegramBot(cfg *TelegramConfig, btdAddr string) (t *TelegramBot) {
	bot, err := tgbotapi.NewBotAPI(cfg.Token)
	if err != nil {
		logger.Panicf("tg bot init failed: %+v", err)
	}
	delay, err := time.ParseDuration(cfg.DeleteDelay)
	if err != nil {
		logger.Panicf("delete delay error: %+v", err)
	}

	t = &TelegramBot{
		Name:          bot.Self.UserName,
		SelfChatID:    cfg.SelfChatID,
		ChannelChatID: cfg.ChannelChatID,
		ComicPath:     cfg.ComicPath,
		DeleteDelay:   delay,
		Client:        bot,
		Tube:          "tg",
	}
	t.Queue = &bt.Pool{
		Dial: func() (*bt.Conn, error) {
			return bt.Dial(btdAddr)
		},
		MaxIdle:     10,
		MaxActive:   100,
		IdleTimeout: 60 * time.Second,
		MaxLifetime: 180 * time.Second,
		Wait:        true,
	}
	return
}

func (t *TelegramBot) putQueue(msg []byte) {
	conn, err := t.Queue.Get()
	if err != nil {
		logger.Errorf("%+v: %s", err, string(msg))
		return
	}
	conn.Use(t.Tube)
	_, err = conn.Put(msg, 1, t.DeleteDelay, time.Minute)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
}

func (t *TelegramBot) sendFile(chat int64, file string) (tgbotapi.Message, error) {
	logger.Debugf("[%d]%s", chat, file)
	return t.Client.Send(tgbotapi.NewDocumentUpload(chat, file))
}
func (t *TelegramBot) sendPic(chat int64, file string) (tgbotapi.Message, error) {
	logger.Debugf("[%d]%s", chat, file)
	if strings.HasSuffix(file, ".mp4") {
		return t.Client.Send(tgbotapi.NewVideoUpload(chat, file))
	}
	return t.Client.Send(tgbotapi.NewPhotoUpload(chat, file))
}

func (t *TelegramBot) send(chat int64, msg string) (tgbotapi.Message, error) {
	logger.Debugf("[%d]%s", chat, msg)
	return t.Client.Send(tgbotapi.NewMessage(chat, msg))
}

func (t *TelegramBot) delMessage() {
	for {
		conn, err := t.Queue.Get()
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		conn.Watch(t.Tube)
		job, err := conn.Reserve()
		if err != nil {
			logger.Warningf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		msg := &tgbotapi.Message{}
		err = json.Unmarshal(job.Body, msg)
		if err != nil {
			logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		delMsg := tgbotapi.DeleteMessageConfig{
			ChatID:    msg.Chat.ID,
			MessageID: msg.MessageID,
		}
		logger.Infof(":[%s]{%s}", getMsgTitle(msg), strconv.Quote(msg.Text))

		_, err = t.Client.DeleteMessage(delMsg)
		if err != nil {
			logger.Errorf("%+v", err)
			err = conn.Bury(job.ID, 0)
			if err != nil {
				logger.Errorf("%+v", err)
			}
			time.Sleep(3 * time.Second)
			continue
		}
		err = conn.Delete(job.ID)
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
		}
		t.Queue.Release(conn, false)
	}
}

func (t *TelegramBot) tgBot() {
	u := tgbotapi.NewUpdate(0)
	u.Timeout = 60
	for {
		updates, err := t.Client.GetUpdatesChan(u)
		if err != nil {
			logger.Errorf("%+v", err)
			time.Sleep(3 * time.Second)
			continue
		}
		var message *tgbotapi.Message
		for update := range updates {
			if update.Message != nil {
				message = update.Message
			} else if update.EditedMessage != nil {
				message = update.EditedMessage
			} else {
				// unkown msg type
				continue
			}
			if message.Chat.IsGroup() {
				logger.Infof(
					"recv:(%s)[%s]{%s}",
					message.Chat.Title,
					message.From.String(),
					strconv.Quote(message.Text))
			} else {
				logger.Infof(
					"recv:[%s]{%s}",
					message.From.String(),
					strconv.Quote(message.Text),
				)
			}

			if message.IsCommand() {
				switch message.Command() {
				case "start":
					go onStart(t, message)
				case "comic":
					go onComic(t, message)
				case "pic":
					go onPic(t, message)
				default:
					logger.Infof("ignore unkown cmd: %+v", message.Command())
					continue
				}
			} else {
				if message.Text == "" {
					continue
				}
				checkRepeat(t, message)
			}
		}
		logger.Warning("tg bot restarted.")
		time.Sleep(3 * time.Second)
	}
}

func checkRepeat(t *TelegramBot, message *tgbotapi.Message) {
	key := "tg_" + getMsgTitle(message) + "_last"
	flattendMsg := strings.TrimSpace(message.Text)
	defer redisClient.LTrim(key, 0, 10)
	defer redisClient.LPush(key, flattendMsg)

	lastMsgs, err := redisClient.LRange(key, 0, 6).Result()
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	i := 0
	for _, s := range lastMsgs {
		if s == flattendMsg {
			i++
		}
	}
	if i > 1 {
		redisClient.Del(key)
		logger.Infof("repeat: %s", strconv.Quote(message.Text))
		msg := tgbotapi.NewMessage(message.Chat.ID, message.Text)
		t.Client.Send(msg)
	}
}

func onStart(t *TelegramBot, message *tgbotapi.Message) {
	msg := tgbotapi.NewMessage(message.Chat.ID, "呀呀呀")
	msg.ReplyToMessageID = message.MessageID
	t.Client.Send(msg)
}

func onComic(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(t.ComicPath)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	rand.Seed(time.Now().Unix())
	file := files[rand.Intn(len(files))]
	number := strings.Split(strings.Split(file, "@")[1], ".")[0]
	msg := tgbotapi.NewMessage(message.Chat.ID, "https://nhentai.net/g/"+number)

	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))
	msgSent, err := t.Client.Send(msg)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	data, err := json.Marshal(msgSent)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	t.putQueue(data)
}

func onPic(t *TelegramBot, message *tgbotapi.Message) {
	files, err := filepath.Glob(twitterBot.ImgPath + "/*")
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	if files == nil {
		logger.Error("find no pics")
	}
	rand.Seed(time.Now().Unix())
	file := files[rand.Intn(len(files))]

	logger.Infof("send:[%s]{%s}", getMsgTitle(message), strconv.Quote(file))
	msgSent, err := t.sendFile(message.Chat.ID, file)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	data, err := json.Marshal(msgSent)
	if err != nil {
		logger.Errorf("%+v", err)
		return
	}
	// t.putQueue(data)
}

func getMsgTitle(m *tgbotapi.Message) string {
	if m.Chat.IsGroup() {
		return m.Chat.Title
	}
	return m.From.String()
}
<|endoftext|>


<commit_before>package multitemplate

import (
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"path/filepath"
	"strings"
	"text/template/parse"
)

// To Register a parser, set it in this map for each extension that would
// correspond to it.
var Parsers = make(map[string]Parser)

// The interface you must have to implement a Parser
type Parser interface {
	ParseTemplate(name, src string, funcs template.FuncMap) (map[string]*parse.Tree, error)
	String() string
}

type Template struct {
	Tmpl  *template.Template
	Base  string
	ctx   *Context
	funcs template.FuncMap
}

func Must(t *Template, err error) *Template {
	if err != nil {
		panic(err)
	}
	return t
}

func New(name string) *Template {
	t := &Template{Tmpl: template.New(name).Funcs(template.FuncMap{}), Base: name}
	t.Funcs(baseFuncMap())
	return t
}
func baseFuncMap() template.FuncMap {
	fm := template.FuncMap{}
	for k, v := range StaticFuncs {
		fm[k] = v
	}
	for k, v := range LoadedFuncs {
		fm[k] = v
	}
	return fm
}

func ParseFiles(filenames ...string) (*Template, error) {
	return (&Template{Tmpl: template.New("root")}).ParseFiles(filenames...)
}

func ParseGlob(pattern string) (*Template, error) {
	return (&Template{Tmpl: template.New("root")}).ParseGlob(pattern)
}

func (t *Template) AddParseTree(name string, tree *parse.Tree) (*Template, error) {
	var e error
	t.Tmpl, e = t.Tmpl.AddParseTree(name, tree)
	return t, e
}

func (t *Template) Clone() (*Template, error) {
	tmpl, err := t.Tmpl.Clone()
	return &Template{tmpl, t.Base, nil, t.funcs}, err
}

func (t *Template) Context(ctx *Context) (*Template, error) {
	tmpl, err := t.Clone()
	if err != nil {
		return nil, err
	}
	tmpl.ctx = ctx
	ctx.tmpl = tmpl
	return tmpl.Funcs(generateFuncs(tmpl)), nil
}

func (t *Template) Execute(w io.Writer, data interface{}) error {
	var tt *Template
	if t.ctx == nil {
		tt, _ = t.Context(NewContext(data))
	}

	e := tt.Tmpl.Execute(tt.ctx.output, data)
	if e == nil {
		return tt.ctx.Close(w)
	}
	return e
}

func (t *Template) ExecuteContext(w io.Writer, ctx *Context) error {
	ctx.executingLayout = false
	tt, e := t.Context(ctx)
	if e != nil {
		return e
	}

	main := ctx.Main
	if ctx.Layout != "" {
		ctx.mainContent, e = tt.ctx.exec(ctx.Main, ctx.Dot)
		if e != nil {
			return e
		}
		main = ctx.Layout
		tt.ctx.executingLayout = true
	}
	return tt.ExecuteTemplate(w, main, ctx.Dot)
}

func (t *Template) ExecuteTemplate(w io.Writer, name string, data interface{}) error {
	tt := t
	if t.ctx == nil {
		tt, _ = t.Context(NewContext(data))
	}

	if e := tt.Tmpl.ExecuteTemplate(tt.ctx.output, name, data); e != nil {
		return e
	}
	return tt.ctx.Close(w)
}

func (t *Template) Funcs(fm template.FuncMap) *Template {
	if t.funcs == nil {
		t.funcs = fm
	} else {
		for k, v := range fm {
			t.funcs[k] = v
		}
	}
	t.Tmpl.Funcs(fm)
	return t
}

func (t *Template) Lookup(name string) *Template {
	tmpl := t.Tmpl.Lookup(name)
	if tmpl != nil {
		return &Template{tmpl, t.Base, nil, t.funcs}
	}
	return nil
}

func (t *Template) Name() string {
	return t.Tmpl.Name()
}

func (t *Template) Parse(name, src, parser string) (*Template, error) {
	p, ok := Parsers[parser]
	if !ok {
		p = &defaultParser{}
	}

	t2, _ := t.Clone()
	trees, err := p.ParseTemplate(name, src, t2.Funcs(generateFuncs(t)).funcs)
	if err != nil {
		return nil, err
	}
	for n, tree := range trees {
		// text/template/parse needs the text of the template to generate errors,
		// but you can't set that without parsing, so make a fake parse run, then swap
		// out the roots while no one's looking. Use 25 delimeters so it all gets parsed
		// as text
		ptt := parse.New(tree.Name)
		ptt.Parse(src, "{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{", "}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}", map[string]*parse.Tree{})
		ptt.Root = tree.Root
		t, err = t.AddParseTree(n, ptt)
		if err != nil {
			return nil, err
		}
	}
	return t, nil
}

func (t *Template) ParseFiles(filenames ...string) (*Template, error) {
	for _, f := range filenames {
		n, p := t.stripBase(f)
		b, e := ioutil.ReadFile(f)
		if e != nil {
			return t, e
		}
		t, e = t.Parse(n, string(b), p)
		if e != nil {
			return t, e
		}
	}
	return t, nil
}

func (t *Template) ParseGlob(pattern string) (*Template, error) {
	filenames, err := filepath.Glob(pattern)
	if err != nil {
		return nil, err
	}
	if len(filenames) == 0 {
		return nil, fmt.Errorf("multitemplate: no files match pattern: %#q", pattern)
	}
	return t.ParseFiles(filenames...)
}

func (t *Template) stripBase(filename string) (name, parser string) {
	name = filename
	if strings.HasPrefix(filename, t.Base) {
		filename = filename[len(t.Base):]
	}
	if filename[0] == '/' || filename[0] == '\\' {
		name = filename[1:]
	}
	base, exts := extensions(name)
	if len(exts) > 0 {
		for _, ext := range exts {
			if _, ok := Parsers[ext]; ok {
				parser = ext
			} else {
				base = base + "." + ext
			}
		}
	}
	name = base
	return
}

func extensions(filename string) (string, []string) {
	name := filepath.Base(filename)
	dirs := filename[:len(filename)-len(name)]
	ext := filepath.Ext(name)
	exts := []string{}
	for ext != "" {
		exts = append(exts, ext[1:])
		name = name[:len(name)-len(ext)]
		ext = filepath.Ext(name)
	}
	return dirs + name, exts
}

func (t *Template) Templates() []*Template {
	tmpls := t.Tmpl.Templates()
	ret := make([]*Template, len(tmpls))
	for i, tmpl := range tmpls {
		ret[i] = &Template{tmpl, t.Base, nil, t.funcs}
	}
	return ret
}
<commit_msg>Don't trust context's to be fully set up<commit_after>package multitemplate

import (
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"path/filepath"
	"strings"
	"text/template/parse"
)

// To Register a parser, set it in this map for each extension that would
// correspond to it.
var Parsers = make(map[string]Parser)

// The interface you must have to implement a Parser
type Parser interface {
	ParseTemplate(name, src string, funcs template.FuncMap) (map[string]*parse.Tree, error)
	String() string
}

type Template struct {
	Tmpl  *template.Template
	Base  string
	ctx   *Context
	funcs template.FuncMap
}

func Must(t *Template, err error) *Template {
	if err != nil {
		panic(err)
	}
	return t
}

func New(name string) *Template {
	t := &Template{Tmpl: template.New(name).Funcs(template.FuncMap{}), Base: name}
	t.Funcs(baseFuncMap())
	return t
}
func baseFuncMap() template.FuncMap {
	fm := template.FuncMap{}
	for k, v := range StaticFuncs {
		fm[k] = v
	}
	for k, v := range LoadedFuncs {
		fm[k] = v
	}
	return fm
}

func ParseFiles(filenames ...string) (*Template, error) {
	return (&Template{Tmpl: template.New("root")}).ParseFiles(filenames...)
}

func ParseGlob(pattern string) (*Template, error) {
	return (&Template{Tmpl: template.New("root")}).ParseGlob(pattern)
}

func (t *Template) AddParseTree(name string, tree *parse.Tree) (*Template, error) {
	var e error
	t.Tmpl, e = t.Tmpl.AddParseTree(name, tree)
	return t, e
}

func (t *Template) Clone() (*Template, error) {
	tmpl, err := t.Tmpl.Clone()
	return &Template{tmpl, t.Base, nil, t.funcs}, err
}

func (t *Template) Context(ctx *Context) (*Template, error) {
	tmpl, err := t.Clone()
	if err != nil {
		return nil, err
	}
	tmpl.ctx = ctx
	ctx.tmpl = tmpl
	if ctx.Yields == nil {
		ctx.Yields = make(map[string]string)
	}
	if ctx.Blocks == nil {
		ctx.Blocks = make(map[string]RenderedBlock)
	}
	if ctx.output == nil {
		ctx.output = newPouchWriter()
	}

	return tmpl.Funcs(generateFuncs(tmpl)), nil
}

func (t *Template) Execute(w io.Writer, data interface{}) error {
	var tt *Template
	if t.ctx == nil {
		tt, _ = t.Context(NewContext(data))
	}

	e := tt.Tmpl.Execute(tt.ctx.output, data)
	if e == nil {
		return tt.ctx.Close(w)
	}
	return e
}

func (t *Template) ExecuteContext(w io.Writer, ctx *Context) error {
	ctx.executingLayout = false
	tt, e := t.Context(ctx)
	if e != nil {
		return e
	}

	main := ctx.Main
	if ctx.Layout != "" {
		ctx.mainContent, e = tt.ctx.exec(ctx.Main, ctx.Dot)
		if e != nil {
			return e
		}
		main = ctx.Layout
		tt.ctx.executingLayout = true
	}
	return tt.ExecuteTemplate(w, main, ctx.Dot)
}

func (t *Template) ExecuteTemplate(w io.Writer, name string, data interface{}) error {
	tt := t
	if t.ctx == nil {
		tt, _ = t.Context(NewContext(data))
	}

	if e := tt.Tmpl.ExecuteTemplate(tt.ctx.output, name, data); e != nil {
		return e
	}
	return tt.ctx.Close(w)
}

func (t *Template) Funcs(fm template.FuncMap) *Template {
	if t.funcs == nil {
		t.funcs = fm
	} else {
		for k, v := range fm {
			t.funcs[k] = v
		}
	}
	t.Tmpl.Funcs(fm)
	return t
}

func (t *Template) Lookup(name string) *Template {
	tmpl := t.Tmpl.Lookup(name)
	if tmpl != nil {
		return &Template{tmpl, t.Base, nil, t.funcs}
	}
	return nil
}

func (t *Template) Name() string {
	return t.Tmpl.Name()
}

func (t *Template) Parse(name, src, parser string) (*Template, error) {
	p, ok := Parsers[parser]
	if !ok {
		p = &defaultParser{}
	}

	t2, _ := t.Clone()
	trees, err := p.ParseTemplate(name, src, t2.Funcs(generateFuncs(t)).funcs)
	if err != nil {
		return nil, err
	}
	for n, tree := range trees {
		// text/template/parse needs the text of the template to generate errors,
		// but you can't set that without parsing, so make a fake parse run, then swap
		// out the roots while no one's looking. Use 25 delimeters so it all gets parsed
		// as text
		ptt := parse.New(tree.Name)
		ptt.Parse(src, "{{{{{{{{{{{{{{{{{{{{{{{{{{{{{{", "}}}}}}}}}}}}}}}}}}}}}}}}}}}}}}", map[string]*parse.Tree{})
		ptt.Root = tree.Root
		t, err = t.AddParseTree(n, ptt)
		if err != nil {
			return nil, err
		}
	}
	return t, nil
}

func (t *Template) ParseFiles(filenames ...string) (*Template, error) {
	for _, f := range filenames {
		n, p := t.stripBase(f)
		b, e := ioutil.ReadFile(f)
		if e != nil {
			return t, e
		}
		t, e = t.Parse(n, string(b), p)
		if e != nil {
			return t, e
		}
	}
	return t, nil
}

func (t *Template) ParseGlob(pattern string) (*Template, error) {
	filenames, err := filepath.Glob(pattern)
	if err != nil {
		return nil, err
	}
	if len(filenames) == 0 {
		return nil, fmt.Errorf("multitemplate: no files match pattern: %#q", pattern)
	}
	return t.ParseFiles(filenames...)
}

func (t *Template) stripBase(filename string) (name, parser string) {
	name = filename
	if strings.HasPrefix(filename, t.Base) {
		filename = filename[len(t.Base):]
	}
	if filename[0] == '/' || filename[0] == '\\' {
		name = filename[1:]
	}
	base, exts := extensions(name)
	if len(exts) > 0 {
		for _, ext := range exts {
			if _, ok := Parsers[ext]; ok {
				parser = ext
			} else {
				base = base + "." + ext
			}
		}
	}
	name = base
	return
}

func extensions(filename string) (string, []string) {
	name := filepath.Base(filename)
	dirs := filename[:len(filename)-len(name)]
	ext := filepath.Ext(name)
	exts := []string{}
	for ext != "" {
		exts = append(exts, ext[1:])
		name = name[:len(name)-len(ext)]
		ext = filepath.Ext(name)
	}
	return dirs + name, exts
}

func (t *Template) Templates() []*Template {
	tmpls := t.Tmpl.Templates()
	ret := make([]*Template, len(tmpls))
	for i, tmpl := range tmpls {
		ret[i] = &Template{tmpl, t.Base, nil, t.funcs}
	}
	return ret
}
<|endoftext|>