<commit_before>// Copyright 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// 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.

// Binary proto_generator generates Protobuf3 code corresponding to an input
// YANG schema. The input set of modules are read, parsed using goyang, and
// handled as input to the ygen package which generates the corresponding
// set of Protobuf3 messages.
package main

import (
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	log "github.com/golang/glog"
	"github.com/openconfig/goyang/pkg/yang"
	"github.com/openconfig/ygot/ygen"
)

var (
	yangPaths           = flag.String("path", "", "Comma separated list of paths to be recursively searched for included modules or submodules within the defined YANG modules.")
	compressPaths       = flag.Bool("compress_paths", false, "If set to true, the schema's paths are compressed, according to OpenConfig YANG module conventions.")
	excludeModules      = flag.String("exclude_modules", "", "Comma separated set of module names that should be excluded from code generation. This can be used to ensure overlapping namespaces can be ignored.")
	packageName         = flag.String("package_name", "openconfig", "The name of the Proto package that generated messages should belong to as their parent.")
	enumPackageName     = flag.String("enum_package_name", "enums", "The name of the package within the generated package that should contain global enum definitions.")
	outputDir           = flag.String("output_dir", "", "The path to which files should be output, hierarchical folders are created for the generated messages.")
	ignoreCircDeps      = flag.Bool("ignore_circdeps", false, "If set to true, circular dependencies between submodules are ignored.")
	baseImportPath      = flag.String("base_import_path", "", "The base import path that should be used for this package, for example a URL to the GitHub repo that the protobuf messages are stored in.")
	ywrapperPath        = flag.String("ywrapper_path", ygen.DefaultYwrapperPath, "The path to the ywrapper.proto file, excluding the file name. Used to import the ywrapper protobuf that specifies the wrapper messages for scalar protobuf types.")
	yextPath            = flag.String("yext_path", ygen.DefaultYextPath, "The path to the yext.proto file, excluding the file name. Used to import the yext protobuf that specifies YANG-specific field options for protobuf.")
	generateFakeRoot    = flag.Bool("generate_fakeroot", false, "If set to true, a fake element at the root of the data tree is generated. The fake root's name can be controlled with the fakeroot_name flag.")
	fakeRootName        = flag.String("fakeroot_name", "Device", "The name of the fake root entity.")
	annotateSchemaPaths = flag.Bool("add_schemapaths", true, "If set to true, the schema path of each YANG entity is added as a protobuf field option")
	annotateEnumNames   = flag.Bool("add_enumnames", true, "If set to true, each value within output enums will be annotated with the label in the original YANG schema.")
	packageHierarchy    = flag.Bool("package_hierarchy", false, "If set to true, an individual protobuf package is output per level of the YANG schema tree.")
	callerName          = flag.String("caller_name", "proto_generator", "The name of the generator binary that should be recorded in output files.")
	excludeState        = flag.Bool("exclude_state", false, "If set to true, state (config false) fields in the YANG schema are not included in the generated Protobuf messages.")
)

// main parses command-line flags to determine the set of YANG modules for
// which code generation should be performed, and calls the codegen library
// to generate Go code corresponding to their schema. The output is written
// to the specified file.
func main() {
	flag.Parse()
	// Extract the set of modules that code is to be generated for,
	// throwing an error if the set is empty.
	generateModules := flag.Args()
	if len(generateModules) == 0 {
		log.Exitln("Error: no input modules specified")
	}

	if *outputDir == "" {
		log.Exitln("Error: an output directory must be specified")
	}

	// Determine the set of paths that should be searched for included
	// modules. This is supplied by the user as a set of comma-separated
	// paths, so we split the string. Additionally, for each path
	// specified, we append "..." to ensure that the directory is
	// recursively searched.
	includePaths := []string{}
	if len(*yangPaths) > 0 {
		pathParts := strings.Split(*yangPaths, ",")
		for _, path := range pathParts {
			includePaths = append(includePaths, filepath.Join(path, "..."))
		}
	}

	// Determine which modules the user has requested to be excluded from
	// code generation.
	modsExcluded := []string{}
	if len(*excludeModules) > 0 {
		modParts := strings.Split(*excludeModules, ",")
		for _, mod := range modParts {
			modsExcluded = append(modsExcluded, mod)
		}
	}

	// Perform the code generation.
	cg := ygen.NewYANGCodeGenerator(&ygen.GeneratorConfig{
		CompressOCPaths:  *compressPaths,
		ExcludeModules:   modsExcluded,
		PackageName:      *packageName,
		GenerateFakeRoot: *generateFakeRoot,
		FakeRootName:     *fakeRootName,
		Caller:           *callerName,
		YANGParseOptions: yang.Options{
			IgnoreSubmoduleCircularDependencies: *ignoreCircDeps,
		},
		ProtoOptions: ygen.ProtoOpts{
			BaseImportPath:      *baseImportPath,
			YwrapperPath:        *ywrapperPath,
			YextPath:            *yextPath,
			AnnotateSchemaPaths: *annotateSchemaPaths,
			AnnotateEnumNames:   *annotateEnumNames,
			NestedMessages:      !*packageHierarchy,
		},
		ExcludeState: *excludeState,
	})

	generatedProtoCode, err := cg.GenerateProto3(generateModules, includePaths)
	if err != nil {
		log.Exitf("%v\n", err)
	}

	for _, p := range generatedProtoCode.Packages {
		fp := filepath.Join(append([]string{*outputDir}, p.FilePath[:len(p.FilePath)-1]...)...)
		if err := os.MkdirAll(fp, 0755); err != nil {
			log.Exitf("could not create directory %v, got error: %v", fp, err)
		}

		f, err := os.Create(filepath.Join(fp, p.FilePath[len(p.FilePath)-1]))
		if err != nil {
			log.Exitf("could not create file %v, got error: %v", fp, err)
		}
		defer f.Close()

		f.WriteString(p.Header)
		for _, m := range p.Messages {
			f.WriteString(fmt.Sprintf("%s\n", m))
		}
		for _, e := range p.Enums {
			f.WriteString(e)
		}
		f.Sync()
	}
}
<commit_msg>Honour enum package name. (#272)<commit_after>// Copyright 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// 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.

// Binary proto_generator generates Protobuf3 code corresponding to an input
// YANG schema. The input set of modules are read, parsed using goyang, and
// handled as input to the ygen package which generates the corresponding
// set of Protobuf3 messages.
package main

import (
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	log "github.com/golang/glog"
	"github.com/openconfig/goyang/pkg/yang"
	"github.com/openconfig/ygot/ygen"
)

var (
	yangPaths           = flag.String("path", "", "Comma separated list of paths to be recursively searched for included modules or submodules within the defined YANG modules.")
	compressPaths       = flag.Bool("compress_paths", false, "If set to true, the schema's paths are compressed, according to OpenConfig YANG module conventions.")
	excludeModules      = flag.String("exclude_modules", "", "Comma separated set of module names that should be excluded from code generation. This can be used to ensure overlapping namespaces can be ignored.")
	packageName         = flag.String("package_name", "openconfig", "The name of the Proto package that generated messages should belong to as their parent.")
	enumPackageName     = flag.String("enum_package_name", "enums", "The name of the package within the generated package that should contain global enum definitions.")
	outputDir           = flag.String("output_dir", "", "The path to which files should be output, hierarchical folders are created for the generated messages.")
	ignoreCircDeps      = flag.Bool("ignore_circdeps", false, "If set to true, circular dependencies between submodules are ignored.")
	baseImportPath      = flag.String("base_import_path", "", "The base import path that should be used for this package, for example a URL to the GitHub repo that the protobuf messages are stored in.")
	ywrapperPath        = flag.String("ywrapper_path", ygen.DefaultYwrapperPath, "The path to the ywrapper.proto file, excluding the file name. Used to import the ywrapper protobuf that specifies the wrapper messages for scalar protobuf types.")
	yextPath            = flag.String("yext_path", ygen.DefaultYextPath, "The path to the yext.proto file, excluding the file name. Used to import the yext protobuf that specifies YANG-specific field options for protobuf.")
	generateFakeRoot    = flag.Bool("generate_fakeroot", false, "If set to true, a fake element at the root of the data tree is generated. The fake root's name can be controlled with the fakeroot_name flag.")
	fakeRootName        = flag.String("fakeroot_name", "Device", "The name of the fake root entity.")
	annotateSchemaPaths = flag.Bool("add_schemapaths", true, "If set to true, the schema path of each YANG entity is added as a protobuf field option")
	annotateEnumNames   = flag.Bool("add_enumnames", true, "If set to true, each value within output enums will be annotated with the label in the original YANG schema.")
	packageHierarchy    = flag.Bool("package_hierarchy", false, "If set to true, an individual protobuf package is output per level of the YANG schema tree.")
	callerName          = flag.String("caller_name", "proto_generator", "The name of the generator binary that should be recorded in output files.")
	excludeState        = flag.Bool("exclude_state", false, "If set to true, state (config false) fields in the YANG schema are not included in the generated Protobuf messages.")
)

// main parses command-line flags to determine the set of YANG modules for
// which code generation should be performed, and calls the codegen library
// to generate Go code corresponding to their schema. The output is written
// to the specified file.
func main() {
	flag.Parse()
	// Extract the set of modules that code is to be generated for,
	// throwing an error if the set is empty.
	generateModules := flag.Args()
	if len(generateModules) == 0 {
		log.Exitln("Error: no input modules specified")
	}

	if *outputDir == "" {
		log.Exitln("Error: an output directory must be specified")
	}

	// Determine the set of paths that should be searched for included
	// modules. This is supplied by the user as a set of comma-separated
	// paths, so we split the string. Additionally, for each path
	// specified, we append "..." to ensure that the directory is
	// recursively searched.
	includePaths := []string{}
	if len(*yangPaths) > 0 {
		pathParts := strings.Split(*yangPaths, ",")
		for _, path := range pathParts {
			includePaths = append(includePaths, filepath.Join(path, "..."))
		}
	}

	// Determine which modules the user has requested to be excluded from
	// code generation.
	modsExcluded := []string{}
	if len(*excludeModules) > 0 {
		modParts := strings.Split(*excludeModules, ",")
		for _, mod := range modParts {
			modsExcluded = append(modsExcluded, mod)
		}
	}

	// Perform the code generation.
	cg := ygen.NewYANGCodeGenerator(&ygen.GeneratorConfig{
		CompressOCPaths:  *compressPaths,
		ExcludeModules:   modsExcluded,
		PackageName:      *packageName,
		GenerateFakeRoot: *generateFakeRoot,
		FakeRootName:     *fakeRootName,
		Caller:           *callerName,
		YANGParseOptions: yang.Options{
			IgnoreSubmoduleCircularDependencies: *ignoreCircDeps,
		},
		ProtoOptions: ygen.ProtoOpts{
			BaseImportPath:      *baseImportPath,
			YwrapperPath:        *ywrapperPath,
			YextPath:            *yextPath,
			AnnotateSchemaPaths: *annotateSchemaPaths,
			AnnotateEnumNames:   *annotateEnumNames,
			NestedMessages:      !*packageHierarchy,
			EnumPackageName:     *enumPackageName,
		},
		ExcludeState: *excludeState,
	})

	generatedProtoCode, err := cg.GenerateProto3(generateModules, includePaths)
	if err != nil {
		log.Exitf("%v\n", err)
	}

	for _, p := range generatedProtoCode.Packages {
		fp := filepath.Join(append([]string{*outputDir}, p.FilePath[:len(p.FilePath)-1]...)...)
		if err := os.MkdirAll(fp, 0755); err != nil {
			log.Exitf("could not create directory %v, got error: %v", fp, err)
		}

		f, err := os.Create(filepath.Join(fp, p.FilePath[len(p.FilePath)-1]))
		if err != nil {
			log.Exitf("could not create file %v, got error: %v", fp, err)
		}
		defer f.Close()

		f.WriteString(p.Header)
		for _, m := range p.Messages {
			f.WriteString(fmt.Sprintf("%s\n", m))
		}
		for _, e := range p.Enums {
			f.WriteString(e)
		}
		f.Sync()
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package common_test

import (
	"fmt"
	"os"
	"time"

	"github.com/juju/loggo"
	"github.com/juju/testing"
	jc "github.com/juju/testing/checkers"
	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/cloudinit"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/environs/network"
	"launchpad.net/juju-core/environs/storage"
	envtesting "launchpad.net/juju-core/environs/testing"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/provider/common"
	coretesting "launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/tools"
	"launchpad.net/juju-core/utils/ssh"
)

type BootstrapSuite struct {
	coretesting.FakeHomeSuite
	envtesting.ToolsFixture
}

var _ = gc.Suite(&BootstrapSuite{})

type cleaner interface {
	AddCleanup(testing.CleanupFunc)
}

func (s *BootstrapSuite) SetUpTest(c *gc.C) {
	s.FakeHomeSuite.SetUpTest(c)
	s.ToolsFixture.SetUpTest(c)
	s.PatchValue(common.ConnectSSH, func(_ ssh.Client, host, checkHostScript string) error {
		return fmt.Errorf("mock connection failure to %s", host)
	})
}

func (s *BootstrapSuite) TearDownTest(c *gc.C) {
	s.ToolsFixture.TearDownTest(c)
	s.LoggingSuite.TearDownTest(c)
}

func newStorage(suite cleaner, c *gc.C) storage.Storage {
	closer, stor, _ := envtesting.CreateLocalTestStorage(c)
	suite.AddCleanup(func(*gc.C) { closer.Close() })
	envtesting.UploadFakeTools(c, stor)
	return stor
}

func minimalConfig(c *gc.C) *config.Config {
	attrs := map[string]interface{}{
		"name":           "whatever",
		"type":           "anything, really",
		"ca-cert":        coretesting.CACert,
		"ca-private-key": coretesting.CAKey,
	}
	cfg, err := config.New(config.UseDefaults, attrs)
	c.Assert(err, gc.IsNil)
	return cfg
}

func configGetter(c *gc.C) configFunc {
	cfg := minimalConfig(c)
	return func() *config.Config { return cfg }
}

func (s *BootstrapSuite) TestCannotStartInstance(c *gc.C) {
	checkPlacement := "directive"
	checkCons := constraints.MustParse("mem=8G")

	startInstance := func(
		placement string, cons constraints.Value, _, _ []string, possibleTools tools.List, mcfg *cloudinit.MachineConfig,
	) (
		instance.Instance, *instance.HardwareCharacteristics, []network.Info, error,
	) {
		c.Assert(placement, gc.DeepEquals, checkPlacement)
		c.Assert(cons, gc.DeepEquals, checkCons)
		c.Assert(mcfg, gc.DeepEquals, environs.NewBootstrapMachineConfig(mcfg.SystemPrivateSSHKey))
		return nil, nil, nil, fmt.Errorf("meh, not started")
	}

	env := &mockEnviron{
		storage:       newStorage(s, c),
		startInstance: startInstance,
		config:        configGetter(c),
	}

	ctx := coretesting.Context(c)
	err := common.Bootstrap(ctx, env, environs.BootstrapParams{
		Constraints: checkCons,
		Placement:   checkPlacement,
	})
	c.Assert(err, gc.ErrorMatches, "cannot start bootstrap instance: meh, not started")
}

func (s *BootstrapSuite) TestCannotRecordStartedInstance(c *gc.C) {
	innerStorage := newStorage(s, c)
	stor := &mockStorage{Storage: innerStorage}

	startInstance := func(
		_ string, _ constraints.Value, _, _ []string, _ tools.List, _ *cloudinit.MachineConfig,
	) (
		instance.Instance, *instance.HardwareCharacteristics, []network.Info, error,
	) {
		stor.putErr = fmt.Errorf("suddenly a wild blah")
		return &mockInstance{id: "i-blah"}, nil, nil, nil
	}

	var stopped []instance.Instance
	stopInstances := func(instances []instance.Instance) error {
		stopped = append(stopped, instances...)
		return nil
	}

	env := &mockEnviron{
		storage:       stor,
		startInstance: startInstance,
		stopInstances: stopInstances,
		config:        configGetter(c),
	}

	ctx := coretesting.Context(c)
	err := common.Bootstrap(ctx, env, environs.BootstrapParams{})
	c.Assert(err, gc.ErrorMatches, "cannot save state: suddenly a wild blah")
	c.Assert(stopped, gc.HasLen, 1)
	c.Assert(stopped[0].Id(), gc.Equals, instance.Id("i-blah"))
}

func (s *BootstrapSuite) TestCannotRecordThenCannotStop(c *gc.C) {
	innerStorage := newStorage(s, c)
	stor := &mockStorage{Storage: innerStorage}

	startInstance := func(
		_ string, _ constraints.Value, _, _ []string, _ tools.List, _ *cloudinit.MachineConfig,
	) (
		instance.Instance, *instance.HardwareCharacteristics, []network.Info, error,
	) {
		stor.putErr = fmt.Errorf("suddenly a wild blah")
		return &mockInstance{id: "i-blah"}, nil, nil, nil
	}

	var stopped []instance.Instance
	stopInstances := func(instances []instance.Instance) error {
		stopped = append(stopped, instances...)
		return fmt.Errorf("bork bork borken")
	}

	tw := &loggo.TestWriter{}
	c.Assert(loggo.RegisterWriter("bootstrap-tester", tw, loggo.DEBUG), gc.IsNil)
	defer loggo.RemoveWriter("bootstrap-tester")

	env := &mockEnviron{
		storage:       stor,
		startInstance: startInstance,
		stopInstances: stopInstances,
		config:        configGetter(c),
	}

	ctx := coretesting.Context(c)
	err := common.Bootstrap(ctx, env, environs.BootstrapParams{})
	c.Assert(err, gc.ErrorMatches, "cannot save state: suddenly a wild blah")
	c.Assert(stopped, gc.HasLen, 1)
	c.Assert(stopped[0].Id(), gc.Equals, instance.Id("i-blah"))
	c.Assert(tw.Log, jc.LogMatches, []jc.SimpleMessage{{
		loggo.ERROR, `cannot stop failed bootstrap instance "i-blah": bork bork borken`,
	}})
}

func (s *BootstrapSuite) TestSuccess(c *gc.C) {
	stor := newStorage(s, c)
	checkInstanceId := "i-success"
	checkHardware := instance.MustParseHardware("mem=2T")

	startInstance := func(
		_ string, _ constraints.Value, _, _ []string, _ tools.List, mcfg *cloudinit.MachineConfig,
	) (
		instance.Instance, *instance.HardwareCharacteristics, []network.Info, error,
	) {
		return &mockInstance{id: checkInstanceId}, &checkHardware, nil, nil
	}
	var mocksConfig = minimalConfig(c)
	var getConfigCalled int
	getConfig := func() *config.Config {
		getConfigCalled++
		return mocksConfig
	}
	setConfig := func(c *config.Config) error {
		mocksConfig = c
		return nil
	}

	restore := envtesting.DisableFinishBootstrap()
	defer restore()

	env := &mockEnviron{
		storage:       stor,
		startInstance: startInstance,
		config:        getConfig,
		setConfig:     setConfig,
	}
	originalAuthKeys := env.Config().AuthorizedKeys()
	ctx := coretesting.Context(c)
	err := common.Bootstrap(ctx, env, environs.BootstrapParams{})
	c.Assert(err, gc.IsNil)

	authKeys := env.Config().AuthorizedKeys()
	c.Assert(authKeys, gc.Not(gc.Equals), originalAuthKeys)
	c.Assert(authKeys, jc.HasSuffix, "juju-system-key\n")
}

type neverRefreshes struct {
}

func (neverRefreshes) Refresh() error {
	return nil
}

type neverAddresses struct {
	neverRefreshes
}

func (neverAddresses) Addresses() ([]instance.Address, error) {
	return nil, nil
}

var testSSHTimeout = config.SSHTimeoutOpts{
	Timeout:        coretesting.ShortWait,
	RetryDelay:     1 * time.Millisecond,
	AddressesDelay: 1 * time.Millisecond,
}

func (s *BootstrapSuite) TestWaitSSHTimesOutWaitingForAddresses(c *gc.C) {
	ctx := coretesting.Context(c)
	_, err := common.WaitSSH(ctx, nil, ssh.DefaultClient, "/bin/true", neverAddresses{}, testSSHTimeout)
	c.Check(err, gc.ErrorMatches, `waited for `+testSSHTimeout.Timeout.String()+` without getting any addresses`)
	c.Check(coretesting.Stderr(ctx), gc.Matches, "Waiting for address\n")
}

func (s *BootstrapSuite) TestWaitSSHKilledWaitingForAddresses(c *gc.C) {
	ctx := coretesting.Context(c)
	interrupted := make(chan os.Signal, 1)
	go func() {
		<-time.After(2 * time.Millisecond)
		interrupted <- os.Interrupt
	}()
	_, err := common.WaitSSH(ctx, interrupted, ssh.DefaultClient, "/bin/true", neverAddresses{}, testSSHTimeout)
	c.Check(err, gc.ErrorMatches, "interrupted")
	c.Check(coretesting.Stderr(ctx), gc.Matches, "Waiting for address\n")
}

type brokenAddresses struct {
	neverRefreshes
}

func (brokenAddresses) Addresses() ([]instance.Address, error) {
	return nil, fmt.Errorf("Addresses will never work")
}

func (s *BootstrapSuite) TestWaitSSHStopsOnBadError(c *gc.C) {
	ctx := coretesting.Context(c)
	_, err := common.WaitSSH(ctx, nil, ssh.DefaultClient, "/bin/true", brokenAddresses{}, testSSHTimeout)
	c.Check(err, gc.ErrorMatches, "getting addresses: Addresses will never work")
	c.Check(coretesting.Stderr(ctx), gc.Equals, "Waiting for address\n")
}

type neverOpensPort struct {
	neverRefreshes
	addr string
}

func (n *neverOpensPort) Addresses() ([]instance.Address, error) {
	return instance.NewAddresses(n.addr), nil
}

func (s *BootstrapSuite) TestWaitSSHTimesOutWaitingForDial(c *gc.C) {
	ctx := coretesting.Context(c)
	// 0.x.y.z addresses are always invalid
	_, err := common.WaitSSH(ctx, nil, ssh.DefaultClient, "/bin/true", &neverOpensPort{addr: "0.1.2.3"}, testSSHTimeout)
	c.Check(err, gc.ErrorMatches,
		`waited for `+testSSHTimeout.Timeout.String()+` without being able to connect: mock connection failure to 0.1.2.3`)
	c.Check(coretesting.Stderr(ctx), gc.Matches,
		"Waiting for address\n"+
			"(Attempting to connect to 0.1.2.3:22\n)+")
}

type interruptOnDial struct {
	neverRefreshes
	name        string
	interrupted chan os.Signal
	returned    bool
}

func (i *interruptOnDial) Addresses() ([]instance.Address, error) {
	// kill the tomb the second time Addresses is called
	if !i.returned {
		i.returned = true
	} else {
		i.interrupted <- os.Interrupt
	}
	return []instance.Address{instance.NewAddress(i.name, instance.NetworkUnknown)}, nil
}

func (s *BootstrapSuite) TestWaitSSHKilledWaitingForDial(c *gc.C) {
	ctx := coretesting.Context(c)
	timeout := testSSHTimeout
	timeout.Timeout = 1 * time.Minute
	interrupted := make(chan os.Signal, 1)
	_, err := common.WaitSSH(ctx, interrupted, ssh.DefaultClient, "", &interruptOnDial{name: "0.1.2.3", interrupted: interrupted}, timeout)
	c.Check(err, gc.ErrorMatches, "interrupted")
	// Exact timing is imprecise but it should have tried a few times before being killed
	c.Check(coretesting.Stderr(ctx), gc.Matches,
		"Waiting for address\n"+
			"(Attempting to connect to 0.1.2.3:22\n)+")
}

type addressesChange struct {
	addrs [][]string
}

func (ac *addressesChange) Refresh() error {
	if len(ac.addrs) > 1 {
		ac.addrs = ac.addrs[1:]
	}
	return nil
}

func (ac *addressesChange) Addresses() ([]instance.Address, error) {
	var addrs []instance.Address
	for _, addr := range ac.addrs[0] {
		addrs = append(addrs, instance.NewAddress(addr, instance.NetworkUnknown))
	}
	return addrs, nil
}

func (s *BootstrapSuite) TestWaitSSHRefreshAddresses(c *gc.C) {
	ctx := coretesting.Context(c)
	_, err := common.WaitSSH(ctx, nil, ssh.DefaultClient, "", &addressesChange{addrs: [][]string{
		nil,
		nil,
		[]string{"0.1.2.3"},
		[]string{"0.1.2.3"},
		nil,
		[]string{"0.1.2.4"},
	}}, testSSHTimeout)
	// Not necessarily the last one in the list, due to scheduling.
	c.Check(err, gc.ErrorMatches,
		`waited for `+testSSHTimeout.Timeout.String()+` without being able to connect: mock connection failure to 0.1.2.[34]`)
	stderr := coretesting.Stderr(ctx)
	c.Check(stderr, gc.Matches,
		"Waiting for address\n"+
			"(.|\n)*(Attempting to connect to 0.1.2.3:22\n)+(.|\n)*")
	c.Check(stderr, gc.Matches,
		"Waiting for address\n"+
			"(.|\n)*(Attempting to connect to 0.1.2.4:22\n)+(.|\n)*")
}
<commit_msg>provider/common: fix time-sensitive test<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package common_test

import (
	"fmt"
	"os"
	"time"

	"github.com/juju/loggo"
	"github.com/juju/testing"
	jc "github.com/juju/testing/checkers"
	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/cloudinit"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/environs/network"
	"launchpad.net/juju-core/environs/storage"
	envtesting "launchpad.net/juju-core/environs/testing"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/provider/common"
	coretesting "launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/tools"
	"launchpad.net/juju-core/utils/ssh"
)

type BootstrapSuite struct {
	coretesting.FakeHomeSuite
	envtesting.ToolsFixture
}

var _ = gc.Suite(&BootstrapSuite{})

type cleaner interface {
	AddCleanup(testing.CleanupFunc)
}

func (s *BootstrapSuite) SetUpTest(c *gc.C) {
	s.FakeHomeSuite.SetUpTest(c)
	s.ToolsFixture.SetUpTest(c)
	s.PatchValue(common.ConnectSSH, func(_ ssh.Client, host, checkHostScript string) error {
		return fmt.Errorf("mock connection failure to %s", host)
	})
}

func (s *BootstrapSuite) TearDownTest(c *gc.C) {
	s.ToolsFixture.TearDownTest(c)
	s.LoggingSuite.TearDownTest(c)
}

func newStorage(suite cleaner, c *gc.C) storage.Storage {
	closer, stor, _ := envtesting.CreateLocalTestStorage(c)
	suite.AddCleanup(func(*gc.C) { closer.Close() })
	envtesting.UploadFakeTools(c, stor)
	return stor
}

func minimalConfig(c *gc.C) *config.Config {
	attrs := map[string]interface{}{
		"name":           "whatever",
		"type":           "anything, really",
		"ca-cert":        coretesting.CACert,
		"ca-private-key": coretesting.CAKey,
	}
	cfg, err := config.New(config.UseDefaults, attrs)
	c.Assert(err, gc.IsNil)
	return cfg
}

func configGetter(c *gc.C) configFunc {
	cfg := minimalConfig(c)
	return func() *config.Config { return cfg }
}

func (s *BootstrapSuite) TestCannotStartInstance(c *gc.C) {
	checkPlacement := "directive"
	checkCons := constraints.MustParse("mem=8G")

	startInstance := func(
		placement string, cons constraints.Value, _, _ []string, possibleTools tools.List, mcfg *cloudinit.MachineConfig,
	) (
		instance.Instance, *instance.HardwareCharacteristics, []network.Info, error,
	) {
		c.Assert(placement, gc.DeepEquals, checkPlacement)
		c.Assert(cons, gc.DeepEquals, checkCons)
		c.Assert(mcfg, gc.DeepEquals, environs.NewBootstrapMachineConfig(mcfg.SystemPrivateSSHKey))
		return nil, nil, nil, fmt.Errorf("meh, not started")
	}

	env := &mockEnviron{
		storage:       newStorage(s, c),
		startInstance: startInstance,
		config:        configGetter(c),
	}

	ctx := coretesting.Context(c)
	err := common.Bootstrap(ctx, env, environs.BootstrapParams{
		Constraints: checkCons,
		Placement:   checkPlacement,
	})
	c.Assert(err, gc.ErrorMatches, "cannot start bootstrap instance: meh, not started")
}

func (s *BootstrapSuite) TestCannotRecordStartedInstance(c *gc.C) {
	innerStorage := newStorage(s, c)
	stor := &mockStorage{Storage: innerStorage}

	startInstance := func(
		_ string, _ constraints.Value, _, _ []string, _ tools.List, _ *cloudinit.MachineConfig,
	) (
		instance.Instance, *instance.HardwareCharacteristics, []network.Info, error,
	) {
		stor.putErr = fmt.Errorf("suddenly a wild blah")
		return &mockInstance{id: "i-blah"}, nil, nil, nil
	}

	var stopped []instance.Instance
	stopInstances := func(instances []instance.Instance) error {
		stopped = append(stopped, instances...)
		return nil
	}

	env := &mockEnviron{
		storage:       stor,
		startInstance: startInstance,
		stopInstances: stopInstances,
		config:        configGetter(c),
	}

	ctx := coretesting.Context(c)
	err := common.Bootstrap(ctx, env, environs.BootstrapParams{})
	c.Assert(err, gc.ErrorMatches, "cannot save state: suddenly a wild blah")
	c.Assert(stopped, gc.HasLen, 1)
	c.Assert(stopped[0].Id(), gc.Equals, instance.Id("i-blah"))
}

func (s *BootstrapSuite) TestCannotRecordThenCannotStop(c *gc.C) {
	innerStorage := newStorage(s, c)
	stor := &mockStorage{Storage: innerStorage}

	startInstance := func(
		_ string, _ constraints.Value, _, _ []string, _ tools.List, _ *cloudinit.MachineConfig,
	) (
		instance.Instance, *instance.HardwareCharacteristics, []network.Info, error,
	) {
		stor.putErr = fmt.Errorf("suddenly a wild blah")
		return &mockInstance{id: "i-blah"}, nil, nil, nil
	}

	var stopped []instance.Instance
	stopInstances := func(instances []instance.Instance) error {
		stopped = append(stopped, instances...)
		return fmt.Errorf("bork bork borken")
	}

	tw := &loggo.TestWriter{}
	c.Assert(loggo.RegisterWriter("bootstrap-tester", tw, loggo.DEBUG), gc.IsNil)
	defer loggo.RemoveWriter("bootstrap-tester")

	env := &mockEnviron{
		storage:       stor,
		startInstance: startInstance,
		stopInstances: stopInstances,
		config:        configGetter(c),
	}

	ctx := coretesting.Context(c)
	err := common.Bootstrap(ctx, env, environs.BootstrapParams{})
	c.Assert(err, gc.ErrorMatches, "cannot save state: suddenly a wild blah")
	c.Assert(stopped, gc.HasLen, 1)
	c.Assert(stopped[0].Id(), gc.Equals, instance.Id("i-blah"))
	c.Assert(tw.Log, jc.LogMatches, []jc.SimpleMessage{{
		loggo.ERROR, `cannot stop failed bootstrap instance "i-blah": bork bork borken`,
	}})
}

func (s *BootstrapSuite) TestSuccess(c *gc.C) {
	stor := newStorage(s, c)
	checkInstanceId := "i-success"
	checkHardware := instance.MustParseHardware("mem=2T")

	startInstance := func(
		_ string, _ constraints.Value, _, _ []string, _ tools.List, mcfg *cloudinit.MachineConfig,
	) (
		instance.Instance, *instance.HardwareCharacteristics, []network.Info, error,
	) {
		return &mockInstance{id: checkInstanceId}, &checkHardware, nil, nil
	}
	var mocksConfig = minimalConfig(c)
	var getConfigCalled int
	getConfig := func() *config.Config {
		getConfigCalled++
		return mocksConfig
	}
	setConfig := func(c *config.Config) error {
		mocksConfig = c
		return nil
	}

	restore := envtesting.DisableFinishBootstrap()
	defer restore()

	env := &mockEnviron{
		storage:       stor,
		startInstance: startInstance,
		config:        getConfig,
		setConfig:     setConfig,
	}
	originalAuthKeys := env.Config().AuthorizedKeys()
	ctx := coretesting.Context(c)
	err := common.Bootstrap(ctx, env, environs.BootstrapParams{})
	c.Assert(err, gc.IsNil)

	authKeys := env.Config().AuthorizedKeys()
	c.Assert(authKeys, gc.Not(gc.Equals), originalAuthKeys)
	c.Assert(authKeys, jc.HasSuffix, "juju-system-key\n")
}

type neverRefreshes struct {
}

func (neverRefreshes) Refresh() error {
	return nil
}

type neverAddresses struct {
	neverRefreshes
}

func (neverAddresses) Addresses() ([]instance.Address, error) {
	return nil, nil
}

var testSSHTimeout = config.SSHTimeoutOpts{
	Timeout:        coretesting.ShortWait,
	RetryDelay:     1 * time.Millisecond,
	AddressesDelay: 1 * time.Millisecond,
}

func (s *BootstrapSuite) TestWaitSSHTimesOutWaitingForAddresses(c *gc.C) {
	ctx := coretesting.Context(c)
	_, err := common.WaitSSH(ctx, nil, ssh.DefaultClient, "/bin/true", neverAddresses{}, testSSHTimeout)
	c.Check(err, gc.ErrorMatches, `waited for `+testSSHTimeout.Timeout.String()+` without getting any addresses`)
	c.Check(coretesting.Stderr(ctx), gc.Matches, "Waiting for address\n")
}

func (s *BootstrapSuite) TestWaitSSHKilledWaitingForAddresses(c *gc.C) {
	ctx := coretesting.Context(c)
	interrupted := make(chan os.Signal, 1)
	interrupted <- os.Interrupt
	_, err := common.WaitSSH(ctx, interrupted, ssh.DefaultClient, "/bin/true", neverAddresses{}, testSSHTimeout)
	c.Check(err, gc.ErrorMatches, "interrupted")
	c.Check(coretesting.Stderr(ctx), gc.Matches, "Waiting for address\n")
}

type brokenAddresses struct {
	neverRefreshes
}

func (brokenAddresses) Addresses() ([]instance.Address, error) {
	return nil, fmt.Errorf("Addresses will never work")
}

func (s *BootstrapSuite) TestWaitSSHStopsOnBadError(c *gc.C) {
	ctx := coretesting.Context(c)
	_, err := common.WaitSSH(ctx, nil, ssh.DefaultClient, "/bin/true", brokenAddresses{}, testSSHTimeout)
	c.Check(err, gc.ErrorMatches, "getting addresses: Addresses will never work")
	c.Check(coretesting.Stderr(ctx), gc.Equals, "Waiting for address\n")
}

type neverOpensPort struct {
	neverRefreshes
	addr string
}

func (n *neverOpensPort) Addresses() ([]instance.Address, error) {
	return instance.NewAddresses(n.addr), nil
}

func (s *BootstrapSuite) TestWaitSSHTimesOutWaitingForDial(c *gc.C) {
	ctx := coretesting.Context(c)
	// 0.x.y.z addresses are always invalid
	_, err := common.WaitSSH(ctx, nil, ssh.DefaultClient, "/bin/true", &neverOpensPort{addr: "0.1.2.3"}, testSSHTimeout)
	c.Check(err, gc.ErrorMatches,
		`waited for `+testSSHTimeout.Timeout.String()+` without being able to connect: mock connection failure to 0.1.2.3`)
	c.Check(coretesting.Stderr(ctx), gc.Matches,
		"Waiting for address\n"+
			"(Attempting to connect to 0.1.2.3:22\n)+")
}

type interruptOnDial struct {
	neverRefreshes
	name        string
	interrupted chan os.Signal
	returned    bool
}

func (i *interruptOnDial) Addresses() ([]instance.Address, error) {
	// kill the tomb the second time Addresses is called
	if !i.returned {
		i.returned = true
	} else {
		i.interrupted <- os.Interrupt
	}
	return []instance.Address{instance.NewAddress(i.name, instance.NetworkUnknown)}, nil
}

func (s *BootstrapSuite) TestWaitSSHKilledWaitingForDial(c *gc.C) {
	ctx := coretesting.Context(c)
	timeout := testSSHTimeout
	timeout.Timeout = 1 * time.Minute
	interrupted := make(chan os.Signal, 1)
	_, err := common.WaitSSH(ctx, interrupted, ssh.DefaultClient, "", &interruptOnDial{name: "0.1.2.3", interrupted: interrupted}, timeout)
	c.Check(err, gc.ErrorMatches, "interrupted")
	// Exact timing is imprecise but it should have tried a few times before being killed
	c.Check(coretesting.Stderr(ctx), gc.Matches,
		"Waiting for address\n"+
			"(Attempting to connect to 0.1.2.3:22\n)+")
}

type addressesChange struct {
	addrs [][]string
}

func (ac *addressesChange) Refresh() error {
	if len(ac.addrs) > 1 {
		ac.addrs = ac.addrs[1:]
	}
	return nil
}

func (ac *addressesChange) Addresses() ([]instance.Address, error) {
	var addrs []instance.Address
	for _, addr := range ac.addrs[0] {
		addrs = append(addrs, instance.NewAddress(addr, instance.NetworkUnknown))
	}
	return addrs, nil
}

func (s *BootstrapSuite) TestWaitSSHRefreshAddresses(c *gc.C) {
	ctx := coretesting.Context(c)
	_, err := common.WaitSSH(ctx, nil, ssh.DefaultClient, "", &addressesChange{addrs: [][]string{
		nil,
		nil,
		[]string{"0.1.2.3"},
		[]string{"0.1.2.3"},
		nil,
		[]string{"0.1.2.4"},
	}}, testSSHTimeout)
	// Not necessarily the last one in the list, due to scheduling.
	c.Check(err, gc.ErrorMatches,
		`waited for `+testSSHTimeout.Timeout.String()+` without being able to connect: mock connection failure to 0.1.2.[34]`)
	stderr := coretesting.Stderr(ctx)
	c.Check(stderr, gc.Matches,
		"Waiting for address\n"+
			"(.|\n)*(Attempting to connect to 0.1.2.3:22\n)+(.|\n)*")
	c.Check(stderr, gc.Matches,
		"Waiting for address\n"+
			"(.|\n)*(Attempting to connect to 0.1.2.4:22\n)+(.|\n)*")
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"							//used for unifrom database access
	"fmt"									//print statements
	"github.com/go-martini/martini"			//extra frame work build on net/http
	_ "github.com/lib/pq"					//go sql driver
	"net/http"								//framework
	"strings"
	"os"
)

func SetupDB() *sql.DB {
	//POSTGRESDB_PORT_5432_TCP_PORT  10.254.76.103
	db, err := sql.Open("postgres", "host=os.Getenv(POSTGRESDB_PORT_5432_TCP_PORT) user=postgres dbname=postgres sslmode=disable") 		//my only lib/pq usage? login into postgres database
	PanicIf(err)

	return db
}

func PanicIf(err error) {
	if err != nil {
		panic(err)
	}
}

func ShowDB(db *sql.DB, r *http.Request, rw http.ResponseWriter) { //localhost:3000/?search=name
		search := "%" + r.URL.Query().Get("search") + "%"
		rows, err := db.Query(`SELECT fname, LName, cost, city
                           FROM custauth
                           WHERE FName ILIKE $1
                           OR LName ILIKE $1
                           OR city ILIKE $1`, search)
		PanicIf(err)
		defer rows.Close()

		var FirstName, LastName, cost, city string
		for rows.Next() {
			err := rows.Scan(&FirstName, &LastName, &cost, &city)
			PanicIf(err)
			fmt.Fprintf(rw, "First Name: %s\nLast Name: %s\nCost: %s\nCity: %s\n\n", FirstName, LastName, cost, city)
		}
	}

func InsertPur(r *http.Request, db *sql.DB){
//	_, err := db.Query("INSERT INTO custauth (custid, fname, lname, city,state, country,email,cost,errorflag) VALUES ($1, $2, $3, $4, $5, $6, $7, $8,$9)",
	_, err := db.Exec("INSERT INTO custauth (custid, fname, lname, city,state, country,email,cost,errorflag) VALUES ($1, $2, $3, $4, $5, $6, $7, $8,$9)",
	r.FormValue("custid"),
		r.FormValue("fname"),
		r.FormValue("lname"),
		r.FormValue("city"),
		r.FormValue("state"),
		r.FormValue("country"),
		r.FormValue("email"),
		r.FormValue("cost"),
		r.FormValue("errorflag"))

	PanicIf(err)
}

func main() {
	m := martini.Classic()
	m.Map(SetupDB())
	m.Get("/", func() string {return  "Hello to GoSQL database v3"})
  	m.Get("/var", func() string {
	for _, e := range os.Environ() {
        pair := strings.Split(e, "=")
		 fmt.Println(pair[0])}
		 return "yo"
	})
	m.Get("/show", ShowDB)
	m.Post("/add", InsertPur)
	m.Run()
}
<commit_msg>Testing KUBERNETES_SERVICE_HOST<commit_after>package main

import (
	"database/sql"							//used for unifrom database access
	"fmt"									//print statements
	"github.com/go-martini/martini"			//extra frame work build on net/http
	_ "github.com/lib/pq"					//go sql driver
	"net/http"								//framework
	"strings"
	"os"
)

func SetupDB() *sql.DB {
	//  10.254.76.103
	db, err := sql.Open("postgres", "host=os.Getenv(KUBERNETES_SERVICE_HOST) user=postgres dbname=postgres sslmode=disable") 		//my only lib/pq usage? login into postgres database
	PanicIf(err)

	return db
}

func PanicIf(err error) {
	if err != nil {
		panic(err)
	}
}

func ShowDB(db *sql.DB, r *http.Request, rw http.ResponseWriter) { //localhost:3000/?search=name
		search := "%" + r.URL.Query().Get("search") + "%"
		rows, err := db.Query(`SELECT fname, LName, cost, city
                           FROM custauth
                           WHERE FName ILIKE $1
                           OR LName ILIKE $1
                           OR city ILIKE $1`, search)
		PanicIf(err)
		defer rows.Close()

		var FirstName, LastName, cost, city string
		for rows.Next() {
			err := rows.Scan(&FirstName, &LastName, &cost, &city)
			PanicIf(err)
			fmt.Fprintf(rw, "First Name: %s\nLast Name: %s\nCost: %s\nCity: %s\n\n", FirstName, LastName, cost, city)
		}
	}

func InsertPur(r *http.Request, db *sql.DB){
//	_, err := db.Query("INSERT INTO custauth (custid, fname, lname, city,state, country,email,cost,errorflag) VALUES ($1, $2, $3, $4, $5, $6, $7, $8,$9)",
	_, err := db.Exec("INSERT INTO custauth (custid, fname, lname, city,state, country,email,cost,errorflag) VALUES ($1, $2, $3, $4, $5, $6, $7, $8,$9)",
	r.FormValue("custid"),
		r.FormValue("fname"),
		r.FormValue("lname"),
		r.FormValue("city"),
		r.FormValue("state"),
		r.FormValue("country"),
		r.FormValue("email"),
		r.FormValue("cost"),
		r.FormValue("errorflag"))

	PanicIf(err)
}

func main() {
	m := martini.Classic()
	m.Map(SetupDB())
	m.Get("/", func() string {return  "Hello to GoSQL database v1.4"})
  	m.Get("/var", func() string {
	for _, e := range os.Environ() {
        pair := strings.Split(e, "=")
		 fmt.Println(pair[0])}
		 return "yo"
	})
	m.Get("/show", ShowDB)
	m.Post("/add", InsertPur)
	m.Run()
}
<|endoftext|>


<commit_before>/*

Copyright (c) 2010 Andrea Fazzi

Permission is hereby granted, free of charge, to any person obtaining
a copy of this software and associated documentation files (the
"Software"), to deal in the Software without restriction, including
without limitation the rights to use, copy, modify, merge, publish,
distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so, subject to
the following conditions:

The above copyright notice and this permission notice shall be
included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

*/

package spectrum

import (
	"testing"
)

var (
	src = []uint32{ 1,  2,  3,  4,
                        5,  6,  7,  8,
	                9,  10, 11, 12,
	                13, 14, 15, 16 }

	expected = []uint32{ 1,   1,   2,   2,  3,  3,  4,  4,
	                     1,   1,   2,   2,  3,  3,  4,  4,
	                     5,   5,   6,   6,  7,  7,  8,  8,
	                     5,   5,   6,   6,  7,  7,  8,  8,
	                     9,   9,   10,  10, 11, 11, 12, 12,
	                     9,   9,   10,  10, 11, 11, 12, 12,
	                     13,  13,  14,  14, 15, 15, 16, 16,
	                     13,  13,  14,  14, 15, 15, 16, 16 }

	dst [16*4]uint32
)

type testSurface struct {
	data []uint32
	w, h uint
}

func (s *testSurface) Width() uint {
	return s.w
}

func (s *testSurface) Height() uint {
	return s.h
}

func (s *testSurface) Bpp() uint {
	return 1
}

func (s *testSurface) SizeInBytes() uint {
	return s.Width() * s.Height() * s.Bpp()
}

func (s *testSurface) Size() uint {
	return s.Width() * s.Height()
}

func (s *testSurface) getValueAt(id uint) uint32 {
	return s.data[id]
}

func (s *testSurface) setPixelValue(x, y uint, value uint32) {
	s.data[s.Width() * y + x] = value
}

func (s *testSurface) setValueAt(id uint, value uint32) { }
func (s *testSurface) setPixel(x, y uint, color [3]byte) { }
func (s *testSurface) setPixelAt(id uint, color [3]byte) { }

func TestScale2x(t *testing.T) {
	Scale2x(&testSurface{ data: src, w: 4, h: 4}, &testSurface{ data: &dst, w: 8, h: 8 })

	for i, val := range expected {
		if val != dst[i] {
			t.Errorf("Expected %d at %d but got %d", val, i, dst[i])
		}
	}
}

<commit_msg>Remove filter_test.go<commit_after><|endoftext|>


<commit_before>// Package camo provides an HTTP proxy server with content type
// restrictions as well as regex host allow list support.
package camo

import (
	"errors"
	"io"
	"net"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"syscall"
	"time"

	"github.com/cactus/go-camo/camo/encoding"
	"github.com/cactus/gologit"
	"github.com/gorilla/mux"
	httpclient "github.com/mreiferson/go-httpclient"
)

// Config holds configuration data used when creating a Proxy with New.
type Config struct {
	// HMACKey is a byte slice to be used as the hmac key
	HMACKey []byte
	// AllowList is a list of string represenstations of regex (not compiled
	// regex) that are used as a whitelist filter. If an AllowList is present,
	// then anything not matching is dropped. If no AllowList is present,
	// no Allow filtering is done.
	AllowList []string
	// MaxSize is the maximum valid image size response (in bytes).
	MaxSize int64
	// MaxRedirects is the maximum number of redirects to follow.
	MaxRedirects int
	// Request timeout is a timeout for fetching upstream data.
	RequestTimeout time.Duration
	// Server name used in Headers and Via checks
	ServerName string
	// Default headers to add to each response
	AddHeaders map[string]string
	// Keepalive enable/disable
	DisableKeepAlivesFE bool
	DisableKeepAlivesBE bool
}

// ProxyMetrics interface for Proxy to use for stats/metrics.
// This must be goroutine safe, as AddBytes and AddServed will be called from
// many goroutines.
type ProxyMetrics interface {
	AddBytes(bc int64)
	AddServed()
}

// A Proxy is a Camo like HTTP proxy, that provides content type
// restrictions as well as regex host allow list support.
type Proxy struct {
	client *http.Client
	config *Config
	// compiled allow list regex
	allowList []*regexp.Regexp
	metrics   ProxyMetrics
}

// NotFound handles client requests that did not match a valid
// mux route.
func (p *Proxy) NotFoundHandler() http.Handler {
	handler := func(w http.ResponseWriter, req *http.Request) {
		h := w.Header()
		for k, v := range p.config.AddHeaders {
			h.Set(k, v)
		}
		h.Set("Server", p.config.ServerName)
		h.Set("Date", formattedDate.String())
		http.Error(w, "404 Not Found", 404)
	}
	return http.HandlerFunc(handler)
}

// ServerHTTP handles the client request, validates the request is validly
// HMAC signed, filters based on the Allow list, and then proxies
// valid requests to the desired endpoint. Responses are filtered for
// proper image content types.
func (p *Proxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	gologit.Debugln("Request:", req.URL)
	if p.metrics != nil {
		go p.metrics.AddServed()
	}

	// add default headers
	wh := w.Header()
	for k, v := range p.config.AddHeaders {
		wh.Set(k, v)
	}
	wh.Set("Server", p.config.ServerName)
	wh.Set("Date", formattedDate.String())
	if p.config.DisableKeepAlivesFE {
		wh.Set("Connection", "close")
	}

	if req.Header.Get("Via") == p.config.ServerName {
		http.Error(w, "Request loop failure", http.StatusNotFound)
		return
	}

	vars := mux.Vars(req)
	sURL, ok := encoding.DecodeURL(p.config.HMACKey, vars["sigHash"], vars["encodedURL"])
	if !ok {
		http.Error(w, "Bad Signature", http.StatusForbidden)
		return
	}
	gologit.Debugln("URL:", sURL)
	gologit.Debugln("Client request:", req)

	u, err := url.Parse(sURL)
	if err != nil {
		gologit.Debugln("url parse error:", err)
		http.Error(w, "Bad url", http.StatusBadRequest)
		return
	}

	u.Host = strings.ToLower(u.Host)
	if u.Host == "" || localhostRegex.MatchString(u.Host) {
		http.Error(w, "Bad url host", http.StatusNotFound)
		return
	}

	// if allowList is set, require match
	matchFound := true
	if len(p.allowList) > 0 {
		matchFound = false
		for _, rgx := range p.allowList {
			if rgx.MatchString(u.Host) {
				matchFound = true
			}
		}
	}
	if !matchFound {
		http.Error(w, "Allowlist host failure", http.StatusNotFound)
		return
	}

	// filter out rfc1918 hosts
	ip := net.ParseIP(u.Host)
	if ip != nil {
		if addr1918PrefixRegex.MatchString(ip.String()) {
			http.Error(w, "Denylist host failure", http.StatusNotFound)
			return
		}
	}

	nreq, err := http.NewRequest(req.Method, sURL, nil)
	if err != nil {
		gologit.Debugln("Could not create NewRequest", err)
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	}

	// filter headers
	p.copyHeader(&nreq.Header, &req.Header, &ValidReqHeaders)
	if req.Header.Get("X-Forwarded-For") == "" {
		host, _, err := net.SplitHostPort(req.RemoteAddr)
		if err == nil && !addr1918PrefixRegex.MatchString(host) {
			nreq.Header.Add("X-Forwarded-For", host)
		}
	}

	// add an accept header if the client didn't send one
	if nreq.Header.Get("Accept") == "" {
		nreq.Header.Add("Accept", "image/*")
	}

	nreq.Header.Add("user-agent", p.config.ServerName)
	nreq.Header.Add("via", p.config.ServerName)

	gologit.Debugln("Built outgoing request:", nreq)

	resp, err := p.client.Do(nreq)
	if err != nil {
		gologit.Debugln("Could not connect to endpoint", err)
		// this is a bit janky, but better than peeling off the
		// 3 layers of wrapped errors and trying to get to net.OpErr and
		// still having to rely on string comparison to find out if it is
		// a net.errClosing or not.
		errString := err.Error()
		if strings.Contains(errString, "timeout") {
			http.Error(w, "Error Fetching Resource", http.StatusGatewayTimeout)
		} else if strings.Contains(errString, "use of closed") {
			http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		} else {
			// some other error. call it a not found (camo compliant)
			http.Error(w, "Error Fetching Resource", http.StatusNotFound)
		}
		return
	}
	defer resp.Body.Close()
	gologit.Debugln("Response from upstream:", resp)

	// check for too large a response
	if resp.ContentLength > p.config.MaxSize {
		gologit.Debugln("Content length exceeded", sURL)
		http.Error(w, "Content length exceeded", http.StatusNotFound)
		return
	}

	switch resp.StatusCode {
	case 200:
		// check content type
		if !strings.HasPrefix(resp.Header.Get("Content-Type"), "image/") {
			gologit.Debugln("Non-Image content-type returned", u)
			http.Error(w, "Non-Image content-type returned",
				http.StatusBadRequest)
			return
		}
	case 300:
		gologit.Debugln("Multiple choices not supported")
		http.Error(w, "Multiple choices not supported", http.StatusNotFound)
		return
	case 301, 302, 303, 307:
		// if we get a redirect here, we either disabled following,
		// or followed until max depth and still got one (redirect loop)
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 304:
		h := w.Header()
		p.copyHeader(&h, &resp.Header, &ValidRespHeaders)
		w.WriteHeader(304)
		return
	case 404:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 500, 502, 503, 504:
		// upstream errors should probably just 502. client can try later.
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	default:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}

	h := w.Header()
	p.copyHeader(&h, &resp.Header, &ValidRespHeaders)
	w.WriteHeader(resp.StatusCode)

	// since this uses io.Copy from the respBody, it is streaming
	// from the request to the response. This means it will nearly
	// always end up with a chunked response.
	bW, err := io.Copy(w, resp.Body)
	if err != nil {
		if opErr, ok := err.(*net.OpError); ok {
			switch opErr.Err {
			case syscall.EPIPE, syscall.ECONNRESET:
				// broken pipe - endpoint terminated the conn
				// connection reset by peer - endpoint terminated the conn
				// log as debug only.
				gologit.Debugln("OpError writing response:", err)
			default:
				// log anything else normally
				gologit.Println("OpError writing response:", err)
			}
		} else {
			// unknown error and not an OpError.
			gologit.Println("Error writing response:", err)
		}
		return
	}

	if p.metrics != nil {
		go p.metrics.AddBytes(bW)
	}
	gologit.Debugln("Response to client:", w)
}

// copy headers from src into dst
// empty filter map will result in no filtering being done
func (p *Proxy) copyHeader(dst, src *http.Header, filter *map[string]bool) {
	f := *filter
	filtering := false
	if len(f) > 0 {
		filtering = true
	}

	for k, vv := range *src {
		if x, ok := f[k]; filtering && (!ok || !x) {
			continue
		}
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// SetMetricsCollector sets a proxy metrics (ProxyMetrics interface) for
// the proxy
func (p *Proxy) SetMetricsCollector(pm ProxyMetrics) {
	p.metrics = pm
}

// New returns a new Proxy. An error is returned if there was a failure
// to parse the regex from the passed Config.
func New(pc Config) (*Proxy, error) {
	tr := &httpclient.Transport{
		MaxIdleConnsPerHost: 8,
		ConnectTimeout:      2 * time.Second,
		RequestTimeout:      pc.RequestTimeout,
		DisableKeepAlives:   pc.DisableKeepAlivesBE,
		// no need for compression with images
		// some xml/svg can be compressed, but apparently some clients can
		// exhibit weird behavior when those are compressed
		DisableCompression:  true,
	}

	// spawn an idle conn trimmer
	go func() {
		// prunes every 5 minutes. this is just a guess at an
		// initial value. very busy severs may want to lower this...
		for {
			time.Sleep(5 * time.Minute)
			tr.CloseIdleConnections()
		}
	}()

	client := &http.Client{Transport: tr}
	client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
		if len(via) >= pc.MaxRedirects {
			return errors.New("Too many redirects")
		}
		return nil
	}

	var allow []*regexp.Regexp
	var c *regexp.Regexp
	var err error
	// compile allow list
	for _, v := range pc.AllowList {
		c, err = regexp.Compile(strings.TrimSpace(v))
		if err != nil {
			return nil, err
		}
		allow = append(allow, c)
	}

	return &Proxy{
		client:    client,
		config:    &pc,
		allowList: allow}, nil
}
<commit_msg>make path param parsing router agnostic<commit_after>// Package camo provides an HTTP proxy server with content type
// restrictions as well as regex host allow list support.
package camo

import (
	"errors"
	"io"
	"net"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"syscall"
	"time"

	"github.com/cactus/go-camo/camo/encoding"
	"github.com/cactus/gologit"
	httpclient "github.com/mreiferson/go-httpclient"
)

// Config holds configuration data used when creating a Proxy with New.
type Config struct {
	// HMACKey is a byte slice to be used as the hmac key
	HMACKey []byte
	// AllowList is a list of string represenstations of regex (not compiled
	// regex) that are used as a whitelist filter. If an AllowList is present,
	// then anything not matching is dropped. If no AllowList is present,
	// no Allow filtering is done.
	AllowList []string
	// MaxSize is the maximum valid image size response (in bytes).
	MaxSize int64
	// MaxRedirects is the maximum number of redirects to follow.
	MaxRedirects int
	// Request timeout is a timeout for fetching upstream data.
	RequestTimeout time.Duration
	// Server name used in Headers and Via checks
	ServerName string
	// Default headers to add to each response
	AddHeaders map[string]string
	// Keepalive enable/disable
	DisableKeepAlivesFE bool
	DisableKeepAlivesBE bool
}

// ProxyMetrics interface for Proxy to use for stats/metrics.
// This must be goroutine safe, as AddBytes and AddServed will be called from
// many goroutines.
type ProxyMetrics interface {
	AddBytes(bc int64)
	AddServed()
}

// A Proxy is a Camo like HTTP proxy, that provides content type
// restrictions as well as regex host allow list support.
type Proxy struct {
	client *http.Client
	config *Config
	// compiled allow list regex
	allowList []*regexp.Regexp
	metrics   ProxyMetrics
}

// NotFound handles client requests that did not match a valid
// mux route.
func (p *Proxy) NotFoundHandler() http.Handler {
	handler := func(w http.ResponseWriter, req *http.Request) {
		h := w.Header()
		for k, v := range p.config.AddHeaders {
			h.Set(k, v)
		}
		h.Set("Server", p.config.ServerName)
		h.Set("Date", formattedDate.String())
		http.Error(w, "404 Not Found", 404)
	}
	return http.HandlerFunc(handler)
}

// ServerHTTP handles the client request, validates the request is validly
// HMAC signed, filters based on the Allow list, and then proxies
// valid requests to the desired endpoint. Responses are filtered for
// proper image content types.
func (p *Proxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	gologit.Debugln("Request:", req.URL)
	if p.metrics != nil {
		go p.metrics.AddServed()
	}

	// add default headers
	wh := w.Header()
	for k, v := range p.config.AddHeaders {
		wh.Set(k, v)
	}
	wh.Set("Server", p.config.ServerName)
	wh.Set("Date", formattedDate.String())
	if p.config.DisableKeepAlivesFE {
		wh.Set("Connection", "close")
	}

	if req.Header.Get("Via") == p.config.ServerName {
		http.Error(w, "Request loop failure", http.StatusNotFound)
		return
	}

	// split path and get components
	components := strings.Split(req.URL.Path, "/")
	if len(components) < 3 {
		http.Error(w, "Malformed request path", http.StatusNotFound)
		return
	}
	sigHash, encodedURL := components[1], components[2]

	sURL, ok := encoding.DecodeURL(p.config.HMACKey, sigHash, encodedURL)
	if !ok {
		http.Error(w, "Bad Signature", http.StatusForbidden)
		return
	}
	gologit.Debugln("URL:", sURL)
	gologit.Debugln("Client request:", req)

	u, err := url.Parse(sURL)
	if err != nil {
		gologit.Debugln("url parse error:", err)
		http.Error(w, "Bad url", http.StatusBadRequest)
		return
	}

	u.Host = strings.ToLower(u.Host)
	if u.Host == "" || localhostRegex.MatchString(u.Host) {
		http.Error(w, "Bad url host", http.StatusNotFound)
		return
	}

	// if allowList is set, require match
	matchFound := true
	if len(p.allowList) > 0 {
		matchFound = false
		for _, rgx := range p.allowList {
			if rgx.MatchString(u.Host) {
				matchFound = true
			}
		}
	}
	if !matchFound {
		http.Error(w, "Allowlist host failure", http.StatusNotFound)
		return
	}

	// filter out rfc1918 hosts
	ip := net.ParseIP(u.Host)
	if ip != nil {
		if addr1918PrefixRegex.MatchString(ip.String()) {
			http.Error(w, "Denylist host failure", http.StatusNotFound)
			return
		}
	}

	nreq, err := http.NewRequest(req.Method, sURL, nil)
	if err != nil {
		gologit.Debugln("Could not create NewRequest", err)
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	}

	// filter headers
	p.copyHeader(&nreq.Header, &req.Header, &ValidReqHeaders)
	if req.Header.Get("X-Forwarded-For") == "" {
		host, _, err := net.SplitHostPort(req.RemoteAddr)
		if err == nil && !addr1918PrefixRegex.MatchString(host) {
			nreq.Header.Add("X-Forwarded-For", host)
		}
	}

	// add an accept header if the client didn't send one
	if nreq.Header.Get("Accept") == "" {
		nreq.Header.Add("Accept", "image/*")
	}

	nreq.Header.Add("user-agent", p.config.ServerName)
	nreq.Header.Add("via", p.config.ServerName)

	gologit.Debugln("Built outgoing request:", nreq)

	resp, err := p.client.Do(nreq)
	if err != nil {
		gologit.Debugln("Could not connect to endpoint", err)
		// this is a bit janky, but better than peeling off the
		// 3 layers of wrapped errors and trying to get to net.OpErr and
		// still having to rely on string comparison to find out if it is
		// a net.errClosing or not.
		errString := err.Error()
		if strings.Contains(errString, "timeout") {
			http.Error(w, "Error Fetching Resource", http.StatusGatewayTimeout)
		} else if strings.Contains(errString, "use of closed") {
			http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		} else {
			// some other error. call it a not found (camo compliant)
			http.Error(w, "Error Fetching Resource", http.StatusNotFound)
		}
		return
	}
	defer resp.Body.Close()
	gologit.Debugln("Response from upstream:", resp)

	// check for too large a response
	if resp.ContentLength > p.config.MaxSize {
		gologit.Debugln("Content length exceeded", sURL)
		http.Error(w, "Content length exceeded", http.StatusNotFound)
		return
	}

	switch resp.StatusCode {
	case 200:
		// check content type
		if !strings.HasPrefix(resp.Header.Get("Content-Type"), "image/") {
			gologit.Debugln("Non-Image content-type returned", u)
			http.Error(w, "Non-Image content-type returned",
				http.StatusBadRequest)
			return
		}
	case 300:
		gologit.Debugln("Multiple choices not supported")
		http.Error(w, "Multiple choices not supported", http.StatusNotFound)
		return
	case 301, 302, 303, 307:
		// if we get a redirect here, we either disabled following,
		// or followed until max depth and still got one (redirect loop)
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 304:
		h := w.Header()
		p.copyHeader(&h, &resp.Header, &ValidRespHeaders)
		w.WriteHeader(304)
		return
	case 404:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	case 500, 502, 503, 504:
		// upstream errors should probably just 502. client can try later.
		http.Error(w, "Error Fetching Resource", http.StatusBadGateway)
		return
	default:
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}

	h := w.Header()
	p.copyHeader(&h, &resp.Header, &ValidRespHeaders)
	w.WriteHeader(resp.StatusCode)

	// since this uses io.Copy from the respBody, it is streaming
	// from the request to the response. This means it will nearly
	// always end up with a chunked response.
	bW, err := io.Copy(w, resp.Body)
	if err != nil {
		if opErr, ok := err.(*net.OpError); ok {
			switch opErr.Err {
			case syscall.EPIPE, syscall.ECONNRESET:
				// broken pipe - endpoint terminated the conn
				// connection reset by peer - endpoint terminated the conn
				// log as debug only.
				gologit.Debugln("OpError writing response:", err)
			default:
				// log anything else normally
				gologit.Println("OpError writing response:", err)
			}
		} else {
			// unknown error and not an OpError.
			gologit.Println("Error writing response:", err)
		}
		return
	}

	if p.metrics != nil {
		go p.metrics.AddBytes(bW)
	}
	gologit.Debugln("Response to client:", w)
}

// copy headers from src into dst
// empty filter map will result in no filtering being done
func (p *Proxy) copyHeader(dst, src *http.Header, filter *map[string]bool) {
	f := *filter
	filtering := false
	if len(f) > 0 {
		filtering = true
	}

	for k, vv := range *src {
		if x, ok := f[k]; filtering && (!ok || !x) {
			continue
		}
		for _, v := range vv {
			dst.Add(k, v)
		}
	}
}

// SetMetricsCollector sets a proxy metrics (ProxyMetrics interface) for
// the proxy
func (p *Proxy) SetMetricsCollector(pm ProxyMetrics) {
	p.metrics = pm
}

// New returns a new Proxy. An error is returned if there was a failure
// to parse the regex from the passed Config.
func New(pc Config) (*Proxy, error) {
	tr := &httpclient.Transport{
		MaxIdleConnsPerHost: 8,
		ConnectTimeout:      2 * time.Second,
		RequestTimeout:      pc.RequestTimeout,
		DisableKeepAlives:   pc.DisableKeepAlivesBE,
		// no need for compression with images
		// some xml/svg can be compressed, but apparently some clients can
		// exhibit weird behavior when those are compressed
		DisableCompression:  true,
	}

	// spawn an idle conn trimmer
	go func() {
		// prunes every 5 minutes. this is just a guess at an
		// initial value. very busy severs may want to lower this...
		for {
			time.Sleep(5 * time.Minute)
			tr.CloseIdleConnections()
		}
	}()

	client := &http.Client{Transport: tr}
	client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
		if len(via) >= pc.MaxRedirects {
			return errors.New("Too many redirects")
		}
		return nil
	}

	var allow []*regexp.Regexp
	var c *regexp.Regexp
	var err error
	// compile allow list
	for _, v := range pc.AllowList {
		c, err = regexp.Compile(strings.TrimSpace(v))
		if err != nil {
			return nil, err
		}
		allow = append(allow, c)
	}

	return &Proxy{
		client:    client,
		config:    &pc,
		allowList: allow}, nil
}
<|endoftext|>


<commit_before>package api

import (
	. "gopkg.in/check.v1"
)

type RepositoriesClientTest struct {
	BaseClientTest
}

var _ = Suite(&RepositoriesClientTest{newBaseClientTest()})

func (t *RepositoriesClientTest) SetUpTest(c *C) {
	t.BaseClientTest.SetUpTest(c)
	t.client.SetAdminSecret(adminSecret)
}

func (t *RepositoriesClientTest) TestRegister(c *C) {
	err := t.client.Register(t.pubKey)
	c.Assert(err, IsNil)
}

func (t *RepositoriesClientTest) TestConnError(c *C) {
	t.server.Close()
	err := t.client.Register(t.pubKey)
	c.Assert(err, NotNil)
}
<commit_msg>api: added additional admin secret missing test to the repositories_client_test.go file.<commit_after>package api

import (
	. "gopkg.in/check.v1"
)

type RepositoriesClientTest struct {
	BaseClientTest
}

var _ = Suite(&RepositoriesClientTest{newBaseClientTest()})

func (t *RepositoriesClientTest) SetUpTest(c *C) {
	t.BaseClientTest.SetUpTest(c)
	t.client.SetAdminSecret(adminSecret)
}

func (t *RepositoriesClientTest) TestRegister(c *C) {
	err := t.client.Register(t.pubKey)
	c.Assert(err, IsNil)
}

func (t *RepositoriesClientTest) TestConnError(c *C) {
	t.server.Close()
	err := t.client.Register(t.pubKey)
	c.Assert(err, NotNil)
}

func (t *RepositoriesClientTest) TestAdminSecretError(c *C) {
	t.client.adminSecret = []byte{}
	err := t.client.Register(t.pubKey)
	c.Assert(err, NotNil)
}
<|endoftext|>


<commit_before>package logging

import (
	"bytes"
	"fmt"
	"sort"
	"strings"
	"time"

	log "github.com/sirupsen/logrus"
)

const quoteCharacter = "\""

type textFormatter struct {
}

func (f *textFormatter) Format(entry *log.Entry) ([]byte, error) {
	var b *bytes.Buffer
	keys := make([]string, 0, len(entry.Data))
	for k := range entry.Data {
		keys = append(keys, k)
	}

	sort.Strings(keys)
	if entry.Buffer != nil {
		b = entry.Buffer
	} else {
		b = &bytes.Buffer{}
	}

	prefixFieldClashes(entry.Data)

	timestampFormat := time.RFC3339
	b.WriteString(entry.Time.Format(timestampFormat))
	b.WriteByte(' ')
	f.appendValue(b, entry.Level.String())
	b.WriteByte(' ')

	tid := entry.Data["tid"]
	if tid != nil {
		b.WriteByte('<')
		b.WriteString(tid.(string))
		b.WriteString("> ")
	}

	if entry.Message != "" {
		f.appendValue(b, entry.Message)
	}
	b.WriteByte(' ')
	for _, key := range keys {
		if key == "tid" {
			continue
		}
		f.appendKeyValue(b, key, entry.Data[key])
	}
	b.WriteByte('\n')
	return b.Bytes(), nil
}

func (f *textFormatter) needsQuoting(text string) bool {
	for _, ch := range text {
		if !((ch >= 'a' && ch <= 'z') ||
			(ch >= 'A' && ch <= 'Z') ||
			(ch >= '0' && ch <= '9') ||
			ch == '-' || ch == '.') {
			return true
		}
	}
	return false
}

func (f *textFormatter) appendKeyValue(b *bytes.Buffer, key string, value interface{}) {
	b.WriteString(key)
	b.WriteByte('=')
	f.appendValue(b, value)
	b.WriteByte(' ')
}

func (f *textFormatter) appendValue(b *bytes.Buffer, value interface{}) {
	switch value := value.(type) {
	case string:
		if !f.needsQuoting(value) {
			b.WriteString(value)
		} else {
			b.WriteString(f.quoteString(value))
		}
	case error:
		errmsg := value.Error()
		if !f.needsQuoting(errmsg) {
			b.WriteString(errmsg)
		} else {
			b.WriteString(f.quoteString(errmsg))
		}
	default:
		fmt.Fprint(b, value)
	}
}

func (f *textFormatter) quoteString(v string) string {
	escapedQuote := fmt.Sprintf("\\%s", quoteCharacter)
	escapedValue := strings.Replace(v, quoteCharacter, escapedQuote, -1)

	return fmt.Sprintf("%s%v%s", quoteCharacter, escapedValue, quoteCharacter)
}

// This is to not silently overwrite `time`, `msg` and `level` fields when
// dumping it. If this code wasn't there doing:
//
//  logrus.WithField("level", 1).Info("hello")
//
// Would just silently drop the user provided level. Instead with this code
// it'll logged as:
//
//  {"level": "info", "fields.level": 1, "msg": "hello", "time": "..."}
//
// It's not exported because it's still using Data in an opinionated way. It's to
// avoid code duplication between the two default formatters.
func prefixFieldClashes(data log.Fields) {
	if t, ok := data["time"]; ok {
		data["fields.time"] = t
	}

	if m, ok := data["msg"]; ok {
		data["fields.msg"] = m
	}

	if l, ok := data["level"]; ok {
		data["fields.level"] = l
	}
}
<commit_msg>Increase timestamp precision in console log formatter<commit_after>package logging

import (
	"bytes"
	"fmt"
	"sort"
	"strings"

	log "github.com/sirupsen/logrus"
)

const quoteCharacter = "\""

type textFormatter struct {
}

func (f *textFormatter) Format(entry *log.Entry) ([]byte, error) {
	var b *bytes.Buffer
	keys := make([]string, 0, len(entry.Data))
	for k := range entry.Data {
		keys = append(keys, k)
	}

	sort.Strings(keys)
	if entry.Buffer != nil {
		b = entry.Buffer
	} else {
		b = &bytes.Buffer{}
	}

	prefixFieldClashes(entry.Data)

	timestampFormat := "2006-01-02 15:04:05.000000 Z07"
	b.WriteString(entry.Time.Format(timestampFormat))
	b.WriteByte(' ')
	f.appendValue(b, entry.Level.String())
	b.WriteByte(' ')

	tid := entry.Data["tid"]
	if tid != nil {
		b.WriteByte('<')
		b.WriteString(tid.(string))
		b.WriteString("> ")
	}

	if entry.Message != "" {
		f.appendValue(b, entry.Message)
	}
	b.WriteByte(' ')
	for _, key := range keys {
		if key == "tid" {
			continue
		}
		f.appendKeyValue(b, key, entry.Data[key])
	}
	b.WriteByte('\n')
	return b.Bytes(), nil
}

func (f *textFormatter) needsQuoting(text string) bool {
	for _, ch := range text {
		if !((ch >= 'a' && ch <= 'z') ||
			(ch >= 'A' && ch <= 'Z') ||
			(ch >= '0' && ch <= '9') ||
			ch == '-' || ch == '.') {
			return true
		}
	}
	return false
}

func (f *textFormatter) appendKeyValue(b *bytes.Buffer, key string, value interface{}) {
	b.WriteString(key)
	b.WriteByte('=')
	f.appendValue(b, value)
	b.WriteByte(' ')
}

func (f *textFormatter) appendValue(b *bytes.Buffer, value interface{}) {
	switch value := value.(type) {
	case string:
		if !f.needsQuoting(value) {
			b.WriteString(value)
		} else {
			b.WriteString(f.quoteString(value))
		}
	case error:
		errmsg := value.Error()
		if !f.needsQuoting(errmsg) {
			b.WriteString(errmsg)
		} else {
			b.WriteString(f.quoteString(errmsg))
		}
	default:
		fmt.Fprint(b, value)
	}
}

func (f *textFormatter) quoteString(v string) string {
	escapedQuote := fmt.Sprintf("\\%s", quoteCharacter)
	escapedValue := strings.Replace(v, quoteCharacter, escapedQuote, -1)

	return fmt.Sprintf("%s%v%s", quoteCharacter, escapedValue, quoteCharacter)
}

// This is to not silently overwrite `time`, `msg` and `level` fields when
// dumping it. If this code wasn't there doing:
//
//  logrus.WithField("level", 1).Info("hello")
//
// Would just silently drop the user provided level. Instead with this code
// it'll logged as:
//
//  {"level": "info", "fields.level": 1, "msg": "hello", "time": "..."}
//
// It's not exported because it's still using Data in an opinionated way. It's to
// avoid code duplication between the two default formatters.
func prefixFieldClashes(data log.Fields) {
	if t, ok := data["time"]; ok {
		data["fields.time"] = t
	}

	if m, ok := data["msg"]; ok {
		data["fields.msg"] = m
	}

	if l, ok := data["level"]; ok {
		data["fields.level"] = l
	}
}
<|endoftext|>


<commit_before>package carbon

import (
	"fmt"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"

	"go.uber.org/zap"

	"github.com/lomik/carbon-clickhouse/helper/RowBinary"
	"github.com/lomik/carbon-clickhouse/receiver"
	"github.com/lomik/carbon-clickhouse/uploader"
	"github.com/lomik/carbon-clickhouse/writer"
	"github.com/lomik/zapwriter"
)

type App struct {
	sync.RWMutex
	Config           *Config
	Writer           *writer.Writer
	Uploaders        map[string]uploader.Uploader
	UDP              receiver.Receiver
	TCP              receiver.Receiver
	Pickle           receiver.Receiver
	Grpc             receiver.Receiver
	Prometheus       receiver.Receiver
	TelegrafHttpJson receiver.Receiver
	Collector        *Collector // (!!!) Should be re-created on every change config/modules
	writeChan        chan *RowBinary.WriteBuffer
	exit             chan bool
	ConfigFilename   string
}

// New App instance
func New(configFilename string) *App {
	app := &App{
		exit:           make(chan bool),
		ConfigFilename: configFilename,
	}

	return app
}

// configure loads config from config file, schemas.conf, aggregation.conf
func (app *App) configure() error {
	cfg, err := ReadConfig(app.ConfigFilename)
	if err != nil {
		return err
	}

	// carbon-cache prefix
	if hostname, err := os.Hostname(); err == nil {
		hostname = strings.Replace(hostname, ".", "_", -1)
		cfg.Common.MetricPrefix = strings.Replace(cfg.Common.MetricPrefix, "{host}", hostname, -1)
	} else {
		cfg.Common.MetricPrefix = strings.Replace(cfg.Common.MetricPrefix, "{host}", "localhost", -1)
	}

	if cfg.Common.MetricEndpoint == "" {
		cfg.Common.MetricEndpoint = MetricEndpointLocal
	}

	if cfg.Common.MetricEndpoint != MetricEndpointLocal {
		u, err := url.Parse(cfg.Common.MetricEndpoint)

		if err != nil {
			return fmt.Errorf("common.metric-endpoint parse error: %s", err.Error())
		}

		if u.Scheme != "tcp" && u.Scheme != "udp" {
			return fmt.Errorf("common.metric-endpoint supports only tcp and udp protocols. %#v is unsupported", u.Scheme)
		}
	}

	err = cfg.TagDesc.Configure()
	if err != nil {
		return err
	}

	app.Config = cfg

	return nil
}

// ParseConfig loads config from config file
func (app *App) ParseConfig() error {
	app.Lock()
	defer app.Unlock()

	return app.configure()
}

// Stop all socket listeners
func (app *App) stopListeners() {
	logger := zapwriter.Logger("app")

	if app.TCP != nil {
		app.TCP.Stop()
		app.TCP = nil
		logger.Debug("finished", zap.String("module", "tcp"))
	}

	if app.Pickle != nil {
		app.Pickle.Stop()
		app.Pickle = nil
		logger.Debug("finished", zap.String("module", "pickle"))
	}

	if app.UDP != nil {
		app.UDP.Stop()
		app.UDP = nil
		logger.Debug("finished", zap.String("module", "udp"))
	}

	if app.Grpc != nil {
		app.Grpc.Stop()
		app.Grpc = nil
		logger.Debug("finished", zap.String("module", "grpc"))
	}

	if app.Prometheus != nil {
		app.Prometheus.Stop()
		app.Prometheus = nil
		logger.Debug("finished", zap.String("module", "prometheus"))
	}

	if app.TelegrafHttpJson != nil {
		app.TelegrafHttpJson.Stop()
		app.TelegrafHttpJson = nil
		logger.Debug("finished", zap.String("module", "telegraf_http_json"))
	}
}

func (app *App) stopAll() {
	logger := zapwriter.Logger("app")

	app.stopListeners()

	if app.Collector != nil {
		app.Collector.Stop()
		app.Collector = nil
		logger.Debug("finished", zap.String("module", "collector"))
	}

	if app.Writer != nil {
		app.Writer.Stop()
		app.Writer = nil
		logger.Debug("finished", zap.String("module", "writer"))
	}

	if app.Uploaders != nil {
		for n, u := range app.Uploaders {
			u.Stop()
			logger.Debug("finished", zap.String("module", "uploader"), zap.String("name", n))
		}
		app.Uploaders = nil
	}

	if app.exit != nil {
		close(app.exit)
		app.exit = nil
		logger.Debug("close(app.exit)", zap.String("module", "app"))
	}
}

// Stop force stop all components
func (app *App) Stop() {
	app.Lock()
	defer app.Unlock()
	app.stopAll()
}

// Start starts
func (app *App) Start() (err error) {
	app.Lock()
	defer app.Unlock()

	defer func() {
		if err != nil {
			app.stopAll()
		}
	}()

	conf := app.Config

	runtime.GOMAXPROCS(conf.Common.MaxCPU)

	app.writeChan = make(chan *RowBinary.WriteBuffer)

	/* WRITER start */
	uploaders := make([]string, 0, len(conf.Upload))
	for t := range conf.Upload {
		uploaders = append(uploaders, t)
	}

	conf.Data.AutoInterval.SetDefault(conf.Data.FileInterval.Value())

	app.Writer = writer.New(
		app.writeChan,
		conf.Data.Path,
		conf.Data.AutoInterval,
		conf.Data.CompAlgo.CompAlgo,
		conf.Data.CompLevel,
		uploaders,
		nil,
	)
	app.Writer.Start()
	/* WRITER end */

	/* UPLOADER start */
	app.Uploaders = make(map[string]uploader.Uploader)
	for uploaderName, uploaderConfig := range conf.Upload {
		uploaderDir := filepath.Join(conf.Data.Path, uploaderName)
		if err := os.MkdirAll(uploaderDir, 0755); err != nil {
			return err
		}
		up, err := uploader.New(uploaderDir, uploaderName, uploaderConfig)
		if err != nil {
			return err
		}
		app.Uploaders[uploaderName] = up

		// debug cache dump
		if dumper, ok := up.(uploader.DebugCacheDumper); ok {
			func(uploaderName string, d uploader.DebugCacheDumper) {
				http.HandleFunc(fmt.Sprintf("/debug/upload/%s/cache/", uploaderName), func(w http.ResponseWriter, r *http.Request) {
					d.CacheDump(w)
				})
			}(uploaderName, dumper)
		}
	}

	for _, uploader := range app.Uploaders {
		uploader.Start()
	}
	/* UPLOADER end */

	/* RECEIVER start */
	if conf.Tcp.Enabled {
		app.TCP, err = receiver.New(
			"tcp://"+conf.Tcp.Listen,
			app.Config.TagDesc,
			receiver.ParseThreads(runtime.GOMAXPROCS(-1)*2),
			receiver.WriteChan(app.writeChan),
			receiver.DropFuture(uint32(conf.Tcp.DropFuture.Value().Seconds())),
			receiver.DropPast(uint32(conf.Tcp.DropPast.Value().Seconds())),
		)

		if err != nil {
			return
		}

		http.HandleFunc("/debug/receive/tcp/dropped/", app.TCP.DroppedHandler)
	}

	if conf.Udp.Enabled {
		app.UDP, err = receiver.New(
			"udp://"+conf.Udp.Listen,
			app.Config.TagDesc,
			receiver.ParseThreads(runtime.GOMAXPROCS(-1)*2),
			receiver.WriteChan(app.writeChan),
			receiver.DropFuture(uint32(conf.Udp.DropFuture.Value().Seconds())),
			receiver.DropPast(uint32(conf.Udp.DropPast.Value().Seconds())),
		)

		if err != nil {
			return
		}

		http.HandleFunc("/debug/receive/udp/dropped/", app.UDP.DroppedHandler)
	}

	if conf.Pickle.Enabled {
		app.Pickle, err = receiver.New(
			"pickle://"+conf.Pickle.Listen,
			app.Config.TagDesc,
			receiver.ParseThreads(runtime.GOMAXPROCS(-1)*2),
			receiver.WriteChan(app.writeChan),
			receiver.DropFuture(uint32(conf.Pickle.DropFuture.Value().Seconds())),
			receiver.DropPast(uint32(conf.Pickle.DropPast.Value().Seconds())),
		)

		if err != nil {
			return
		}

		http.HandleFunc("/debug/receive/pickle/dropped/", app.Pickle.DroppedHandler)
	}

	if conf.Grpc.Enabled {
		app.Grpc, err = receiver.New(
			"grpc://"+conf.Grpc.Listen,
			app.Config.TagDesc,
			receiver.WriteChan(app.writeChan),
			receiver.DropFuture(uint32(conf.Grpc.DropFuture.Value().Seconds())),
			receiver.DropPast(uint32(conf.Grpc.DropPast.Value().Seconds())),
		)

		if err != nil {
			return
		}

		http.HandleFunc("/debug/receive/grpc/dropped/", app.Grpc.DroppedHandler)
	}

	if conf.Prometheus.Enabled {
		app.Prometheus, err = receiver.New(
			"prometheus://"+conf.Prometheus.Listen,
			app.Config.TagDesc,
			receiver.WriteChan(app.writeChan),
			receiver.DropFuture(uint32(conf.Prometheus.DropFuture.Value().Seconds())),
			receiver.DropPast(uint32(conf.Prometheus.DropPast.Value().Seconds())),
		)

		if err != nil {
			return
		}

		http.HandleFunc("/debug/receive/prometheus/dropped/", app.Prometheus.DroppedHandler)
	}

	if conf.TelegrafHttpJson.Enabled {
		app.TelegrafHttpJson, err = receiver.New(
			"telegraf+http+json://"+conf.TelegrafHttpJson.Listen,
			app.Config.TagDesc,
			receiver.WriteChan(app.writeChan),
			receiver.DropFuture(uint32(conf.TelegrafHttpJson.DropFuture.Value().Seconds())),
			receiver.DropPast(uint32(conf.TelegrafHttpJson.DropPast.Value().Seconds())),
		)

		if err != nil {
			return
		}

		http.HandleFunc("/debug/receive/telegraf_http_json/dropped/", app.TelegrafHttpJson.DroppedHandler)
	}
	/* RECEIVER end */

	/* COLLECTOR start */
	app.Collector = NewCollector(app)
	/* COLLECTOR end */

	return
}

// Reset cache in uploaders
func (app *App) Reset() {
	logger := zapwriter.Logger("app")
	logger.Info("HUP received")

	for n, u := range app.Uploaders {
		if v, ok := u.(uploader.UploaderWithReset); ok {
			logger.Info("reset cache", zap.String("module", "uploader"), zap.String("name", n))
			go v.Reset()
		}
	}
}

// Loop ...
func (app *App) Loop() {
	app.RLock()
	exitChan := app.exit
	app.RUnlock()

	if exitChan != nil {
		<-app.exit
	}
}
<commit_msg>create main data directory<commit_after>package carbon

import (
	"fmt"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"

	"go.uber.org/zap"

	"github.com/lomik/carbon-clickhouse/helper/RowBinary"
	"github.com/lomik/carbon-clickhouse/receiver"
	"github.com/lomik/carbon-clickhouse/uploader"
	"github.com/lomik/carbon-clickhouse/writer"
	"github.com/lomik/zapwriter"
)

type App struct {
	sync.RWMutex
	Config           *Config
	Writer           *writer.Writer
	Uploaders        map[string]uploader.Uploader
	UDP              receiver.Receiver
	TCP              receiver.Receiver
	Pickle           receiver.Receiver
	Grpc             receiver.Receiver
	Prometheus       receiver.Receiver
	TelegrafHttpJson receiver.Receiver
	Collector        *Collector // (!!!) Should be re-created on every change config/modules
	writeChan        chan *RowBinary.WriteBuffer
	exit             chan bool
	ConfigFilename   string
}

// New App instance
func New(configFilename string) *App {
	app := &App{
		exit:           make(chan bool),
		ConfigFilename: configFilename,
	}

	return app
}

// configure loads config from config file, schemas.conf, aggregation.conf
func (app *App) configure() error {
	cfg, err := ReadConfig(app.ConfigFilename)
	if err != nil {
		return err
	}

	// carbon-cache prefix
	if hostname, err := os.Hostname(); err == nil {
		hostname = strings.Replace(hostname, ".", "_", -1)
		cfg.Common.MetricPrefix = strings.Replace(cfg.Common.MetricPrefix, "{host}", hostname, -1)
	} else {
		cfg.Common.MetricPrefix = strings.Replace(cfg.Common.MetricPrefix, "{host}", "localhost", -1)
	}

	if cfg.Common.MetricEndpoint == "" {
		cfg.Common.MetricEndpoint = MetricEndpointLocal
	}

	if cfg.Common.MetricEndpoint != MetricEndpointLocal {
		u, err := url.Parse(cfg.Common.MetricEndpoint)

		if err != nil {
			return fmt.Errorf("common.metric-endpoint parse error: %s", err.Error())
		}

		if u.Scheme != "tcp" && u.Scheme != "udp" {
			return fmt.Errorf("common.metric-endpoint supports only tcp and udp protocols. %#v is unsupported", u.Scheme)
		}
	}

	err = cfg.TagDesc.Configure()
	if err != nil {
		return err
	}

	app.Config = cfg

	return nil
}

// ParseConfig loads config from config file
func (app *App) ParseConfig() error {
	app.Lock()
	defer app.Unlock()

	return app.configure()
}

// Stop all socket listeners
func (app *App) stopListeners() {
	logger := zapwriter.Logger("app")

	if app.TCP != nil {
		app.TCP.Stop()
		app.TCP = nil
		logger.Debug("finished", zap.String("module", "tcp"))
	}

	if app.Pickle != nil {
		app.Pickle.Stop()
		app.Pickle = nil
		logger.Debug("finished", zap.String("module", "pickle"))
	}

	if app.UDP != nil {
		app.UDP.Stop()
		app.UDP = nil
		logger.Debug("finished", zap.String("module", "udp"))
	}

	if app.Grpc != nil {
		app.Grpc.Stop()
		app.Grpc = nil
		logger.Debug("finished", zap.String("module", "grpc"))
	}

	if app.Prometheus != nil {
		app.Prometheus.Stop()
		app.Prometheus = nil
		logger.Debug("finished", zap.String("module", "prometheus"))
	}

	if app.TelegrafHttpJson != nil {
		app.TelegrafHttpJson.Stop()
		app.TelegrafHttpJson = nil
		logger.Debug("finished", zap.String("module", "telegraf_http_json"))
	}
}

func (app *App) stopAll() {
	logger := zapwriter.Logger("app")

	app.stopListeners()

	if app.Collector != nil {
		app.Collector.Stop()
		app.Collector = nil
		logger.Debug("finished", zap.String("module", "collector"))
	}

	if app.Writer != nil {
		app.Writer.Stop()
		app.Writer = nil
		logger.Debug("finished", zap.String("module", "writer"))
	}

	if app.Uploaders != nil {
		for n, u := range app.Uploaders {
			u.Stop()
			logger.Debug("finished", zap.String("module", "uploader"), zap.String("name", n))
		}
		app.Uploaders = nil
	}

	if app.exit != nil {
		close(app.exit)
		app.exit = nil
		logger.Debug("close(app.exit)", zap.String("module", "app"))
	}
}

// Stop force stop all components
func (app *App) Stop() {
	app.Lock()
	defer app.Unlock()
	app.stopAll()
}

// Start starts
func (app *App) Start() (err error) {
	app.Lock()
	defer app.Unlock()

	defer func() {
		if err != nil {
			app.stopAll()
		}
	}()

	conf := app.Config

	runtime.GOMAXPROCS(conf.Common.MaxCPU)

	app.writeChan = make(chan *RowBinary.WriteBuffer)

	/* WRITER start */
	uploaders := make([]string, 0, len(conf.Upload))
	for t := range conf.Upload {
		uploaders = append(uploaders, t)
	}

	conf.Data.AutoInterval.SetDefault(conf.Data.FileInterval.Value())

	if err := os.MkdirAll(conf.Data.Path, 0755); err != nil {
		return err
	}

	app.Writer = writer.New(
		app.writeChan,
		conf.Data.Path,
		conf.Data.AutoInterval,
		conf.Data.CompAlgo.CompAlgo,
		conf.Data.CompLevel,
		uploaders,
		nil,
	)
	app.Writer.Start()
	/* WRITER end */

	/* UPLOADER start */
	app.Uploaders = make(map[string]uploader.Uploader)
	for uploaderName, uploaderConfig := range conf.Upload {
		uploaderDir := filepath.Join(conf.Data.Path, uploaderName)
		if err := os.MkdirAll(uploaderDir, 0755); err != nil {
			return err
		}
		up, err := uploader.New(uploaderDir, uploaderName, uploaderConfig)
		if err != nil {
			return err
		}
		app.Uploaders[uploaderName] = up

		// debug cache dump
		if dumper, ok := up.(uploader.DebugCacheDumper); ok {
			func(uploaderName string, d uploader.DebugCacheDumper) {
				http.HandleFunc(fmt.Sprintf("/debug/upload/%s/cache/", uploaderName), func(w http.ResponseWriter, r *http.Request) {
					d.CacheDump(w)
				})
			}(uploaderName, dumper)
		}
	}

	for _, uploader := range app.Uploaders {
		uploader.Start()
	}
	/* UPLOADER end */

	/* RECEIVER start */
	if conf.Tcp.Enabled {
		app.TCP, err = receiver.New(
			"tcp://"+conf.Tcp.Listen,
			app.Config.TagDesc,
			receiver.ParseThreads(runtime.GOMAXPROCS(-1)*2),
			receiver.WriteChan(app.writeChan),
			receiver.DropFuture(uint32(conf.Tcp.DropFuture.Value().Seconds())),
			receiver.DropPast(uint32(conf.Tcp.DropPast.Value().Seconds())),
		)

		if err != nil {
			return
		}

		http.HandleFunc("/debug/receive/tcp/dropped/", app.TCP.DroppedHandler)
	}

	if conf.Udp.Enabled {
		app.UDP, err = receiver.New(
			"udp://"+conf.Udp.Listen,
			app.Config.TagDesc,
			receiver.ParseThreads(runtime.GOMAXPROCS(-1)*2),
			receiver.WriteChan(app.writeChan),
			receiver.DropFuture(uint32(conf.Udp.DropFuture.Value().Seconds())),
			receiver.DropPast(uint32(conf.Udp.DropPast.Value().Seconds())),
		)

		if err != nil {
			return
		}

		http.HandleFunc("/debug/receive/udp/dropped/", app.UDP.DroppedHandler)
	}

	if conf.Pickle.Enabled {
		app.Pickle, err = receiver.New(
			"pickle://"+conf.Pickle.Listen,
			app.Config.TagDesc,
			receiver.ParseThreads(runtime.GOMAXPROCS(-1)*2),
			receiver.WriteChan(app.writeChan),
			receiver.DropFuture(uint32(conf.Pickle.DropFuture.Value().Seconds())),
			receiver.DropPast(uint32(conf.Pickle.DropPast.Value().Seconds())),
		)

		if err != nil {
			return
		}

		http.HandleFunc("/debug/receive/pickle/dropped/", app.Pickle.DroppedHandler)
	}

	if conf.Grpc.Enabled {
		app.Grpc, err = receiver.New(
			"grpc://"+conf.Grpc.Listen,
			app.Config.TagDesc,
			receiver.WriteChan(app.writeChan),
			receiver.DropFuture(uint32(conf.Grpc.DropFuture.Value().Seconds())),
			receiver.DropPast(uint32(conf.Grpc.DropPast.Value().Seconds())),
		)

		if err != nil {
			return
		}

		http.HandleFunc("/debug/receive/grpc/dropped/", app.Grpc.DroppedHandler)
	}

	if conf.Prometheus.Enabled {
		app.Prometheus, err = receiver.New(
			"prometheus://"+conf.Prometheus.Listen,
			app.Config.TagDesc,
			receiver.WriteChan(app.writeChan),
			receiver.DropFuture(uint32(conf.Prometheus.DropFuture.Value().Seconds())),
			receiver.DropPast(uint32(conf.Prometheus.DropPast.Value().Seconds())),
		)

		if err != nil {
			return
		}

		http.HandleFunc("/debug/receive/prometheus/dropped/", app.Prometheus.DroppedHandler)
	}

	if conf.TelegrafHttpJson.Enabled {
		app.TelegrafHttpJson, err = receiver.New(
			"telegraf+http+json://"+conf.TelegrafHttpJson.Listen,
			app.Config.TagDesc,
			receiver.WriteChan(app.writeChan),
			receiver.DropFuture(uint32(conf.TelegrafHttpJson.DropFuture.Value().Seconds())),
			receiver.DropPast(uint32(conf.TelegrafHttpJson.DropPast.Value().Seconds())),
		)

		if err != nil {
			return
		}

		http.HandleFunc("/debug/receive/telegraf_http_json/dropped/", app.TelegrafHttpJson.DroppedHandler)
	}
	/* RECEIVER end */

	/* COLLECTOR start */
	app.Collector = NewCollector(app)
	/* COLLECTOR end */

	return
}

// Reset cache in uploaders
func (app *App) Reset() {
	logger := zapwriter.Logger("app")
	logger.Info("HUP received")

	for n, u := range app.Uploaders {
		if v, ok := u.(uploader.UploaderWithReset); ok {
			logger.Info("reset cache", zap.String("module", "uploader"), zap.String("name", n))
			go v.Reset()
		}
	}
}

// Loop ...
func (app *App) Loop() {
	app.RLock()
	exitChan := app.exit
	app.RUnlock()

	if exitChan != nil {
		<-app.exit
	}
}
<|endoftext|>


<commit_before>//go:build linux && cgo && !agent
// +build linux,cgo,!agent

package db

import (
	"fmt"

	"github.com/pkg/errors"

	"github.com/lxc/lxd/lxd/db/query"
)

//go:generate -command mapper lxd-generate db mapper -t operations.mapper.go
//go:generate mapper reset
//go:generate mapper stmt -p db -e operation objects
//go:generate mapper stmt -p db -e operation objects-by-NodeID
//go:generate mapper stmt -p db -e operation objects-by-ID
//go:generate mapper stmt -p db -e operation objects-by-UUID
//go:generate mapper stmt -p db -e operation create-or-replace struct=Operation
//go:generate mapper stmt -p db -e operation delete-by-UUID
//go:generate mapper stmt -p db -e operation delete-by-NodeID

//go:generate mapper method -p db -e operation List
//go:generate mapper method -p db -e operation CreateOrReplace struct=Operation
//go:generate mapper method -p db -e operation DeleteOne
//go:generate mapper method -p db -e operation DeleteMany

// Operation holds information about a single LXD operation running on a node
// in the cluster.
type Operation struct {
	ID          int64         `db:"primary=yes"`                               // Stable database identifier
	UUID        string        `db:"primary=yes"`                               // User-visible identifier
	NodeAddress string        `db:"join=nodes.address&omit=create-or-replace"` // Address of the node the operation is running on
	ProjectID   *int64        `db:"omit=objects"`                              // ID of the project for the operation.
	NodeID      int64         // ID of the node the operation is running on
	Type        OperationType // Type of the operation
}

// OperationFilter specifies potential query parameter fields.
type OperationFilter struct {
	ID     int64
	NodeID int64
	UUID   string
}

// GetLocalOperations returns all operations associated with this node.
func (c *ClusterTx) GetLocalOperations() ([]Operation, error) {
	return c.operations("node_id=?", c.nodeID)
}

// GetLocalOperationsUUIDs returns the UUIDs of all operations associated with this
// node.
func (c *ClusterTx) GetLocalOperationsUUIDs() ([]string, error) {
	stmt := "SELECT uuid FROM operations WHERE node_id=?"
	return query.SelectStrings(c.tx, stmt, c.nodeID)
}

// GetNodesWithRunningOperations returns a list of nodes that have running operations
func (c *ClusterTx) GetNodesWithRunningOperations(project string) ([]string, error) {
	stmt := `
SELECT DISTINCT nodes.address
  FROM operations
  LEFT OUTER JOIN projects ON projects.id = operations.project_id
  JOIN nodes ON nodes.id = operations.node_id
 WHERE projects.name = ? OR operations.project_id IS NULL
`
	return query.SelectStrings(c.tx, stmt, project)
}

// GetOperationsOfType returns a list operations that belong to the specified project and have the desired type.
func (c *ClusterTx) GetOperationsOfType(projectName string, opType OperationType) ([]Operation, error) {
	var ops []Operation

	stmt := `
SELECT operations.id, operations.uuid, operations.type, nodes.address
  FROM operations
  LEFT JOIN projects on projects.id = operations.project_id
  JOIN nodes on nodes.id = operations.node_id
WHERE (projects.name = ? OR operations.project_id IS NULL) and operations.type = ?
`
	rows, err := c.tx.Query(stmt, projectName, opType)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		var op Operation
		err := rows.Scan(&op.ID, &op.UUID, &op.Type, &op.NodeAddress)
		if err != nil {
			return nil, err
		}

		ops = append(ops, op)
	}
	if rows.Err() != nil {
		return nil, err
	}

	return ops, nil
}

// GetOperationWithID returns the operation with the given ID.
func (c *ClusterTx) GetOperationWithID(opID int) (Operation, error) {
	null := Operation{}
	operations, err := c.operations("id=?", opID)
	if err != nil {
		return null, err
	}
	switch len(operations) {
	case 0:
		return null, ErrNoSuchObject
	case 1:
		return operations[0], nil
	default:
		return null, fmt.Errorf("More than one operation matches")
	}
}

// GetOperationByUUID returns the operation with the given UUID.
func (c *ClusterTx) GetOperationByUUID(uuid string) (Operation, error) {
	null := Operation{}
	operations, err := c.operations("uuid=?", uuid)
	if err != nil {
		return null, err
	}
	switch len(operations) {
	case 0:
		return null, ErrNoSuchObject
	case 1:
		return operations[0], nil
	default:
		return null, fmt.Errorf("More than one operation matches")
	}
}

// CreateOperation adds a new operations to the table.
func (c *ClusterTx) CreateOperation(project, uuid string, typ OperationType) (int64, error) {
	var projectID interface{}

	if project != "" {
		var err error
		projectID, err = c.GetProjectID(project)
		if err != nil {
			return -1, errors.Wrap(err, "Fetch project ID")
		}
	} else {
		projectID = nil
	}

	columns := []string{"uuid", "node_id", "type", "project_id"}
	values := []interface{}{uuid, c.nodeID, typ, projectID}
	return query.UpsertObject(c.tx, "operations", columns, values)
}

// RemoveOperation removes the operation with the given UUID.
func (c *ClusterTx) RemoveOperation(uuid string) error {
	result, err := c.tx.Exec("DELETE FROM operations WHERE uuid=?", uuid)
	if err != nil {
		return err
	}
	n, err := result.RowsAffected()
	if err != nil {
		return err
	}
	if n != 1 {
		return fmt.Errorf("query deleted %d rows instead of 1", n)
	}
	return nil
}

// Remove all operations for the given node.
func (c *ClusterTx) removeNodeOperations(nodeID int64) error {
	_, err := c.tx.Exec("DELETE FROM operations WHERE node_id=?", nodeID)
	if err != nil {
		return err
	}

	return nil
}

// Operations returns all operations in the cluster, filtered by the given clause.
func (c *ClusterTx) operations(where string, args ...interface{}) ([]Operation, error) {
	operations := []Operation{}
	dest := func(i int) []interface{} {
		operations = append(operations, Operation{})
		return []interface{}{
			&operations[i].ID,
			&operations[i].UUID,
			&operations[i].NodeAddress,
			&operations[i].Type,
		}
	}
	sql := `
SELECT operations.id, uuid, nodes.address, type FROM operations JOIN nodes ON nodes.id = node_id `
	if where != "" {
		sql += fmt.Sprintf("WHERE %s ", where)
	}
	sql += "ORDER BY operations.id"
	stmt, err := c.tx.Prepare(sql)
	if err != nil {
		return nil, err
	}
	defer stmt.Close()
	err = query.SelectObjects(stmt, dest, args...)
	if err != nil {
		return nil, errors.Wrap(err, "Failed to fetch operations")
	}
	return operations, nil
}
<commit_msg>lxd/db/operations: remove hard-coded functions<commit_after>//go:build linux && cgo && !agent
// +build linux,cgo,!agent

package db

import (
	"github.com/lxc/lxd/lxd/db/query"
)

//go:generate -command mapper lxd-generate db mapper -t operations.mapper.go
//go:generate mapper reset
//go:generate mapper stmt -p db -e operation objects
//go:generate mapper stmt -p db -e operation objects-by-NodeID
//go:generate mapper stmt -p db -e operation objects-by-ID
//go:generate mapper stmt -p db -e operation objects-by-UUID
//go:generate mapper stmt -p db -e operation create-or-replace struct=Operation
//go:generate mapper stmt -p db -e operation delete-by-UUID
//go:generate mapper stmt -p db -e operation delete-by-NodeID

//go:generate mapper method -p db -e operation List
//go:generate mapper method -p db -e operation CreateOrReplace struct=Operation
//go:generate mapper method -p db -e operation DeleteOne
//go:generate mapper method -p db -e operation DeleteMany

// Operation holds information about a single LXD operation running on a node
// in the cluster.
type Operation struct {
	ID          int64         `db:"primary=yes"`                               // Stable database identifier
	UUID        string        `db:"primary=yes"`                               // User-visible identifier
	NodeAddress string        `db:"join=nodes.address&omit=create-or-replace"` // Address of the node the operation is running on
	ProjectID   *int64        `db:"omit=objects"`                              // ID of the project for the operation.
	NodeID      int64         // ID of the node the operation is running on
	Type        OperationType // Type of the operation
}

// OperationFilter specifies potential query parameter fields.
type OperationFilter struct {
	ID     int64
	NodeID int64
	UUID   string
}

// GetNodesWithRunningOperations returns a list of nodes that have running operations
func (c *ClusterTx) GetNodesWithRunningOperations(project string) ([]string, error) {
	stmt := `
SELECT DISTINCT nodes.address
  FROM operations
  LEFT OUTER JOIN projects ON projects.id = operations.project_id
  JOIN nodes ON nodes.id = operations.node_id
 WHERE projects.name = ? OR operations.project_id IS NULL
`
	return query.SelectStrings(c.tx, stmt, project)
}

// GetOperationsOfType returns a list operations that belong to the specified project and have the desired type.
func (c *ClusterTx) GetOperationsOfType(projectName string, opType OperationType) ([]Operation, error) {
	var ops []Operation

	stmt := `
SELECT operations.id, operations.uuid, operations.type, nodes.address
  FROM operations
  LEFT JOIN projects on projects.id = operations.project_id
  JOIN nodes on nodes.id = operations.node_id
WHERE (projects.name = ? OR operations.project_id IS NULL) and operations.type = ?
`
	rows, err := c.tx.Query(stmt, projectName, opType)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	for rows.Next() {
		var op Operation
		err := rows.Scan(&op.ID, &op.UUID, &op.Type, &op.NodeAddress)
		if err != nil {
			return nil, err
		}

		ops = append(ops, op)
	}
	if rows.Err() != nil {
		return nil, err
	}

	return ops, nil
}
<|endoftext|>