<commit_before>package vm

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/grubby/grubby/ast"
	"github.com/grubby/grubby/interpreter/vm/builtins"
	"github.com/grubby/grubby/parser"
)

type vm struct {
	currentFilename string
	ObjectSpace     map[string]builtins.Value
	Globals         map[string]builtins.Value
	KnownSymbols    map[string]builtins.Value
	Classes         map[string]builtins.Class
}

type VM interface {
	Run(string) (builtins.Value, error)
	Get(string) (builtins.Value, error)
	MustGet(string) builtins.Value

	Set(string, builtins.Value)

	Symbols() map[string]builtins.Value
}

func NewVM(name string) VM {
	vm := &vm{
		currentFilename: name,
		Globals:         make(map[string]builtins.Value),
		ObjectSpace:     make(map[string]builtins.Value),
		KnownSymbols:    make(map[string]builtins.Value),
	}
	vm.registerClasses()

	loadPath := vm.Classes["Array"].New()
	vm.Globals["LOAD_PATH"] = loadPath
	vm.Globals[":"] = loadPath

	objectClass := vm.Classes["Object"]
	vm.ObjectSpace["Object"] = objectClass
	vm.ObjectSpace["Kernel"] = builtins.NewGlobalKernelClass()
	vm.ObjectSpace["File"] = builtins.NewFileClass()
	vm.ObjectSpace["ARGV"] = vm.Classes["Array"].New()

	main := objectClass.New()
	main.AddMethod(builtins.NewMethod("to_s", func(args ...builtins.Value) (builtins.Value, error) {
		return builtins.NewString("main"), nil
	}))
	main.AddMethod(builtins.NewMethod("require", func(args ...builtins.Value) (builtins.Value, error) {
		fileName := args[0].(*builtins.StringValue).String()
		if fileName == "rubygems" {
			// don't "require 'rubygems'"
			return vm.Classes["False"].New(), nil
		}

		for _, pathStr := range loadPath.(*builtins.Array).Members() {
			path := pathStr.(*builtins.StringValue)
			fullPath := filepath.Join(path.String(), fileName+".rb")
			file, err := os.Open(fullPath)
			if err != nil {
				continue
			}

			contents, err := ioutil.ReadAll(file)

			if err == nil {
				originalName := vm.currentFilename
				defer func() {
					vm.currentFilename = originalName
				}()

				vm.currentFilename = file.Name()
				_, rubyErr := vm.Run(string(contents))
				return vm.Classes["True"].New(), rubyErr
			}
		}

		err := fmt.Sprintf("LoadError: cannot load such file -- %s", fileName)
		return nil, builtins.NewLoadError(err)

	}))
	main.AddMethod(builtins.NewMethod("puts", func(args ...builtins.Value) (builtins.Value, error) {
		for _, arg := range args {
			os.Stdout.Write([]byte(arg.String() + "\n"))
		}

		return nil, nil
	}))

	vm.ObjectSpace["main"] = main

	return vm
}

func (vm *vm) registerClasses() {
	vm.Classes = map[string]builtins.Class{
		"Array":   builtins.NewArrayClass().(builtins.Class),
		"Object":  builtins.NewGlobalObjectClass(),
		"Process": builtins.NewProcessClass(),
		"True":    builtins.NewTrueClass(),
		"False":   builtins.NewFalseClass(),
		"Nil":     builtins.NewNilClass(),
	}
}

func (vm *vm) MustGet(key string) builtins.Value {
	val, ok := vm.ObjectSpace[key]
	if ok {
		return val
	}

	val, ok = vm.Globals[key]
	if ok {
		return val
	}

	return nil
}

func (vm *vm) Get(key string) (builtins.Value, error) {
	val, ok := vm.ObjectSpace[key]
	if ok {
		return val, nil
	}

	val, ok = vm.Globals[key]
	if ok {
		return val, nil
	}

	return nil, errors.New(fmt.Sprintf("'%s' is undefined", key))
}

func (vm *vm) Set(key string, value builtins.Value) {
	vm.ObjectSpace[key] = value
}

func (vm *vm) Symbols() map[string]builtins.Value {
	return vm.KnownSymbols
}

type ParseError struct{}

func NewParseError() *ParseError {
	return &ParseError{}
}

func (err *ParseError) Error() string {
	return "parse error"
}

func (vm *vm) Run(input string) (builtins.Value, error) {
	lexer := parser.NewLexer(input)
	result := parser.RubyParse(lexer)
	if result != 0 {
		return nil, NewParseError()
	}

	main := vm.ObjectSpace["main"]
	return vm.executeWithContext(parser.Statements, main)
}

func (vm *vm) executeWithContext(statements []ast.Node, context builtins.Value) (builtins.Value, error) {
	var (
		returnValue builtins.Value
		returnErr   error
	)
	for _, statement := range statements {
		switch statement.(type) {
		case ast.IfBlock:
			truthy := false
			ifBlock := statement.(ast.IfBlock)
			switch ifBlock.Condition.(type) {
			case ast.Boolean:
				truthy = ifBlock.Condition.(ast.Boolean).Value
			case ast.BareReference:
				truthy = ifBlock.Condition.(ast.BareReference).Name == "nil"
			default:
				truthy = true
			}

			if truthy {
				returnValue, returnErr = vm.executeWithContext(ifBlock.Body, context)
			} else {
				returnValue, returnErr = vm.executeWithContext(ifBlock.Else, context)
			}
		case ast.FuncDecl:
			// FIXME: assumes for now this will only ever be at the top level
			// it seems like this should be replaced with context, but that's really context for calling methods, not
			// necessarily for defining new methods
			funcNode := statement.(ast.FuncDecl)
			method := builtins.NewMethod(funcNode.Name.Name, func(args ...builtins.Value) (builtins.Value, error) {
				return nil, nil
			})
			returnValue = method
			vm.ObjectSpace["Kernel"].AddPrivateMethod(method)
		case ast.SimpleString:
			returnValue = builtins.NewString(statement.(ast.SimpleString).Value)
		case ast.InterpolatedString:
			returnValue = builtins.NewString(statement.(ast.InterpolatedString).Value)
		case ast.Boolean:
			if statement.(ast.Boolean).Value {
				returnValue = vm.Classes["True"].New()
			} else {
				returnValue = vm.Classes["False"].New()
			}
		case ast.GlobalVariable:
			returnValue = vm.Globals[statement.(ast.GlobalVariable).Name]
		case ast.ConstantInt:
			returnValue = builtins.NewInt(statement.(ast.ConstantInt).Value)
		case ast.ConstantFloat:
			returnValue = builtins.NewFloat(statement.(ast.ConstantFloat).Value)
		case ast.Symbol:
			name := statement.(ast.Symbol).Name
			maybe, ok := vm.KnownSymbols[name]
			if !ok {
				returnValue = builtins.NewSymbol(name)
				vm.KnownSymbols[name] = returnValue
			} else {
				returnValue = maybe
			}
		case ast.BareReference:
			name := statement.(ast.BareReference).Name
			maybe, ok := vm.ObjectSpace[name]
			if ok {
				returnValue = maybe
			} else {
				returnErr = builtins.NewNameError(name, context.String(), context.Class().String())
			}
		case ast.CallExpression:
			var method builtins.Method
			callExpr := statement.(ast.CallExpression)

			var target builtins.Value
			switch callExpr.Target.(type) {
			case ast.BareReference:
				target = vm.ObjectSpace[callExpr.Target.(ast.BareReference).Name]
			case ast.Class:
				target = vm.Classes[callExpr.Target.(ast.Class).Name]
			case ast.GlobalVariable:
				target = vm.Globals[callExpr.Target.(ast.GlobalVariable).Name]
			case nil:
				target = context
			}

			if target == nil {
				nilValue := vm.Classes["Nil"].New()
				return nil, builtins.NewNoMethodError(callExpr.Func.Name, nilValue.String(), nilValue.Class().String())
			}
			method, err := target.Method(callExpr.Func.Name)

			if err != nil {
				err := builtins.NewNameError(callExpr.Func.Name, target.String(), target.Class().String())
				return nil, err
			}

			args := []builtins.Value{}
			for _, astArgument := range callExpr.Args {
				arg, err := vm.executeWithContext([]ast.Node{astArgument}, context)
				if err != nil {
					return nil, err
				}

				args = append(args, arg)
			}

			returnValue, returnErr = method.Execute(args...)
		case ast.Assignment:
			assignment := statement.(ast.Assignment)
			returnValue, err := vm.executeWithContext([]ast.Node{assignment.RHS}, context)
			if err != nil {
				return nil, err
			}

			switch assignment.LHS.(type) {
			case ast.BareReference:
				ref := assignment.LHS.(ast.BareReference)
				vm.ObjectSpace[ref.Name] = returnValue
			case ast.GlobalVariable:
				globalVar := assignment.LHS.(ast.GlobalVariable)
				vm.Globals[globalVar.Name] = returnValue
			default:
				panic(fmt.Sprintf("unimplemented assignment failure: %#v", assignment.LHS))
			}

		case ast.FileNameConstReference:
			returnValue = builtins.NewString(vm.currentFilename)
		case ast.Begin:
			begin := statement.(ast.Begin)
			_, err := vm.executeWithContext(begin.Body, context)

			if err != nil {
				rubyErr := err.(builtins.Value)
				for _, rescue := range begin.Rescue {
					r := rescue.(ast.Rescue)
					if r.Exception.Class.Name == rubyErr.String() {
						_, err = vm.executeWithContext(r.Body, context)
						if err == nil {
							break
						}
					}
				}
			}

			if err != nil {
				returnErr = err
			}
		default:
			panic(fmt.Sprintf("handled unknown statement type: %T:\n\t\n => %#v\n", statement, statement))
		}
	}

	return returnValue, returnErr
}
<commit_msg>Rename some private fields in the vm<commit_after>package vm

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/grubby/grubby/ast"
	"github.com/grubby/grubby/interpreter/vm/builtins"
	"github.com/grubby/grubby/parser"
)

type vm struct {
	currentFilename string
	ObjectSpace     map[string]builtins.Value
	CurrentGlobals  map[string]builtins.Value
	CurrentSymbols  map[string]builtins.Value
	CurrentClasses  map[string]builtins.Class
}

type VM interface {
	Run(string) (builtins.Value, error)
	Get(string) (builtins.Value, error)
	MustGet(string) builtins.Value

	Set(string, builtins.Value)

	Symbols() map[string]builtins.Value
	Globals() map[string]builtins.Value
	Classes() map[string]builtins.Class
}

func NewVM(name string) VM {
	vm := &vm{
		currentFilename: name,
		CurrentGlobals:  make(map[string]builtins.Value),
		ObjectSpace:     make(map[string]builtins.Value),
		CurrentSymbols:  make(map[string]builtins.Value),
	}
	vm.registerClasses()

	loadPath := vm.CurrentClasses["Array"].New()
	vm.CurrentGlobals["LOAD_PATH"] = loadPath
	vm.CurrentGlobals[":"] = loadPath

	objectClass := vm.CurrentClasses["Object"]
	vm.ObjectSpace["Object"] = objectClass
	vm.ObjectSpace["Kernel"] = builtins.NewGlobalKernelClass()
	vm.ObjectSpace["File"] = builtins.NewFileClass()
	vm.ObjectSpace["ARGV"] = vm.CurrentClasses["Array"].New()

	main := objectClass.New()
	main.AddMethod(builtins.NewMethod("to_s", func(args ...builtins.Value) (builtins.Value, error) {
		return builtins.NewString("main"), nil
	}))
	main.AddMethod(builtins.NewMethod("require", func(args ...builtins.Value) (builtins.Value, error) {
		fileName := args[0].(*builtins.StringValue).String()
		if fileName == "rubygems" {
			// don't "require 'rubygems'"
			return vm.CurrentClasses["False"].New(), nil
		}

		for _, pathStr := range loadPath.(*builtins.Array).Members() {
			path := pathStr.(*builtins.StringValue)
			fullPath := filepath.Join(path.String(), fileName+".rb")
			file, err := os.Open(fullPath)
			if err != nil {
				continue
			}

			contents, err := ioutil.ReadAll(file)

			if err == nil {
				originalName := vm.currentFilename
				defer func() {
					vm.currentFilename = originalName
				}()

				vm.currentFilename = file.Name()
				_, rubyErr := vm.Run(string(contents))
				return vm.CurrentClasses["True"].New(), rubyErr
			}
		}

		err := fmt.Sprintf("LoadError: cannot load such file -- %s", fileName)
		return nil, builtins.NewLoadError(err)
	}))
	main.AddMethod(builtins.NewMethod("puts", func(args ...builtins.Value) (builtins.Value, error) {
		for _, arg := range args {
			os.Stdout.Write([]byte(arg.String() + "\n"))
		}

		return nil, nil
	}))

	vm.ObjectSpace["main"] = main

	return vm
}

func (vm *vm) registerClasses() {
	vm.CurrentClasses = map[string]builtins.Class{
		"Array":   builtins.NewArrayClass().(builtins.Class),
		"Object":  builtins.NewGlobalObjectClass(),
		"Process": builtins.NewProcessClass(),
		"True":    builtins.NewTrueClass(),
		"False":   builtins.NewFalseClass(),
		"Nil":     builtins.NewNilClass(),
	}
}

func (vm *vm) MustGet(key string) builtins.Value {
	val, ok := vm.ObjectSpace[key]
	if ok {
		return val
	}

	val, ok = vm.CurrentGlobals[key]
	if ok {
		return val
	}

	return nil
}

func (vm *vm) Get(key string) (builtins.Value, error) {
	val, ok := vm.ObjectSpace[key]
	if ok {
		return val, nil
	}

	val, ok = vm.CurrentGlobals[key]
	if ok {
		return val, nil
	}

	return nil, errors.New(fmt.Sprintf("'%s' is undefined", key))
}

func (vm *vm) Set(key string, value builtins.Value) {
	vm.ObjectSpace[key] = value
}

func (vm *vm) Symbols() map[string]builtins.Value {
	return vm.CurrentSymbols
}

func (vm *vm) Globals() map[string]builtins.Value {
	return vm.CurrentGlobals
}

func (vm *vm) Classes() map[string]builtins.Class {
	return vm.CurrentClasses
}

type ParseError struct{}

func NewParseError() *ParseError {
	return &ParseError{}
}

func (err *ParseError) Error() string {
	return "parse error"
}

func (vm *vm) Run(input string) (builtins.Value, error) {
	lexer := parser.NewLexer(input)
	result := parser.RubyParse(lexer)
	if result != 0 {
		return nil, NewParseError()
	}

	main := vm.ObjectSpace["main"]
	return vm.executeWithContext(parser.Statements, main)
}

func (vm *vm) executeWithContext(statements []ast.Node, context builtins.Value) (builtins.Value, error) {
	var (
		returnValue builtins.Value
		returnErr   error
	)
	for _, statement := range statements {
		switch statement.(type) {
		case ast.IfBlock:
			truthy := false
			ifBlock := statement.(ast.IfBlock)
			switch ifBlock.Condition.(type) {
			case ast.Boolean:
				truthy = ifBlock.Condition.(ast.Boolean).Value
			case ast.BareReference:
				truthy = ifBlock.Condition.(ast.BareReference).Name == "nil"
			default:
				truthy = true
			}

			if truthy {
				returnValue, returnErr = vm.executeWithContext(ifBlock.Body, context)
			} else {
				returnValue, returnErr = vm.executeWithContext(ifBlock.Else, context)
			}
		case ast.FuncDecl:
			// FIXME: assumes for now this will only ever be at the top level
			// it seems like this should be replaced with context, but that's really context for calling methods, not
			// necessarily for defining new methods
			funcNode := statement.(ast.FuncDecl)
			method := builtins.NewMethod(funcNode.Name.Name, func(args ...builtins.Value) (builtins.Value, error) {
				return nil, nil
			})
			returnValue = method
			vm.ObjectSpace["Kernel"].AddPrivateMethod(method)
		case ast.SimpleString:
			returnValue = builtins.NewString(statement.(ast.SimpleString).Value)
		case ast.InterpolatedString:
			returnValue = builtins.NewString(statement.(ast.InterpolatedString).Value)
		case ast.Boolean:
			if statement.(ast.Boolean).Value {
				returnValue = vm.CurrentClasses["True"].New()
			} else {
				returnValue = vm.CurrentClasses["False"].New()
			}
		case ast.GlobalVariable:
			returnValue = vm.CurrentGlobals[statement.(ast.GlobalVariable).Name]
		case ast.ConstantInt:
			returnValue = builtins.NewInt(statement.(ast.ConstantInt).Value)
		case ast.ConstantFloat:
			returnValue = builtins.NewFloat(statement.(ast.ConstantFloat).Value)
		case ast.Symbol:
			name := statement.(ast.Symbol).Name
			maybe, ok := vm.CurrentSymbols[name]
			if !ok {
				returnValue = builtins.NewSymbol(name)
				vm.CurrentSymbols[name] = returnValue
			} else {
				returnValue = maybe
			}
		case ast.BareReference:
			name := statement.(ast.BareReference).Name
			maybe, ok := vm.ObjectSpace[name]
			if ok {
				returnValue = maybe
			} else {
				returnErr = builtins.NewNameError(name, context.String(), context.Class().String())
			}
		case ast.CallExpression:
			var method builtins.Method
			callExpr := statement.(ast.CallExpression)

			var target builtins.Value
			switch callExpr.Target.(type) {
			case ast.BareReference:
				target = vm.ObjectSpace[callExpr.Target.(ast.BareReference).Name]
			case ast.Class:
				target = vm.CurrentClasses[callExpr.Target.(ast.Class).Name]
			case ast.GlobalVariable:
				target = vm.CurrentGlobals[callExpr.Target.(ast.GlobalVariable).Name]
			case nil:
				target = context
			}

			if target == nil {
				nilValue := vm.CurrentClasses["Nil"].New()
				return nil, builtins.NewNoMethodError(callExpr.Func.Name, nilValue.String(), nilValue.Class().String())
			}
			method, err := target.Method(callExpr.Func.Name)

			if err != nil {
				err := builtins.NewNameError(callExpr.Func.Name, target.String(), target.Class().String())
				return nil, err
			}

			args := []builtins.Value{}
			for _, astArgument := range callExpr.Args {
				arg, err := vm.executeWithContext([]ast.Node{astArgument}, context)
				if err != nil {
					return nil, err
				}

				args = append(args, arg)
			}

			returnValue, returnErr = method.Execute(args...)
		case ast.Assignment:
			assignment := statement.(ast.Assignment)
			returnValue, err := vm.executeWithContext([]ast.Node{assignment.RHS}, context)
			if err != nil {
				return nil, err
			}

			switch assignment.LHS.(type) {
			case ast.BareReference:
				ref := assignment.LHS.(ast.BareReference)
				vm.ObjectSpace[ref.Name] = returnValue
			case ast.GlobalVariable:
				globalVar := assignment.LHS.(ast.GlobalVariable)
				vm.CurrentGlobals[globalVar.Name] = returnValue
			default:
				panic(fmt.Sprintf("unimplemented assignment failure: %#v", assignment.LHS))
			}

		case ast.FileNameConstReference:
			returnValue = builtins.NewString(vm.currentFilename)
		case ast.Begin:
			begin := statement.(ast.Begin)
			_, err := vm.executeWithContext(begin.Body, context)

			if err != nil {
				rubyErr := err.(builtins.Value)
				for _, rescue := range begin.Rescue {
					r := rescue.(ast.Rescue)
					if r.Exception.Class.Name == rubyErr.String() {
						_, err = vm.executeWithContext(r.Body, context)
						if err == nil {
							break
						}
					}
				}
			}

			if err != nil {
				returnErr = err
			}
		default:
			panic(fmt.Sprintf("handled unknown statement type: %T:\n\t\n => %#v\n", statement, statement))
		}
	}

	return returnValue, returnErr
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The OpenEBS Authors

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

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

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

package cstorpoolit

import (
	"testing"

	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/openebs/CITF"
	"github.com/openebs/CITF/citf_options"
	apis "github.com/openebs/CITF/pkg/apis/openebs.io/v1alpha1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// TestIntegrationCstorPool function instantiate the cstor pool test suite.
func TestIntegrationCstorPool(t *testing.T) {
	// RegisterFailHandler is used to register failed test cases and produce readable output.
	RegisterFailHandler(Fail)
	// RunSpecs runs all the test cases in the suite.
	RunSpecs(t, "Cstor pool integration test suite")
}

// Create an instance of CITF to use the inbuilt functions that will help
// communicating with the kube-apiserver.
var citfInstance, err = citf.NewCITF(&citfoptions.CreateOptions{
	// K8SInclude is true to get the kube-config from the machine where the suite is running.
	// Kube-config is a config file that establishes communication to the k8s cluster.
	K8SInclude: true,
})

// ToDo: Set up cluster environment before runninng all test cases ( i.e. BeforeSuite)
// The environment set up by BeforeSuite is going to persist for all
// the test cases under run

//var _ = BeforeSuite(func() {
//	//var err error
//	//
//	//Expect(err).NotTo(HaveOccurred())
//})

// ToDo: Set up tear down of cluster environment ( i.e Aftersuite)

// ToDo: Set up cluster environment before every test cases that will be run (i.e. preRunHook)
// ToDo: Reset cluster environment after every test cases that will be run  ( i.e postRunHook)
var _ = Describe("Integration Test", func() {
	// Test Case #1 (sparse-striped-auto-spc). Type : Positive
	When("We apply sparse-striped-auto spc yaml with maxPool count equal to 3 on a k8s cluster having at least 3 capable node", func() {
		It("pool resources count should be 3 with no error and online status", func() {
			// TODO: Create a generic util function in utils.go to convert yaml into go object.
			// ToDo: More POC regarding this util converter function.
			// Functions generic to both cstor-pool and cstor-vol should go inside common directory

			// 1.Read SPC yaml form a file.
			// 2.Convert SPC yaml to json.
			// 3.Marshall json to SPC go object.

			// Create a storage pool claim object
			spcObject := &apis.StoragePoolClaim{
				ObjectMeta: metav1.ObjectMeta{
					Name: "disk-claim-auto",
				},
				Spec: apis.StoragePoolClaimSpec{
					Name:     "sparse-claim-auto",
					Type:     "sparse",
					MaxPools: 3,
					PoolSpec: apis.CStorPoolAttr{
						PoolType: "striped",
					},
				},
			}
			// Call CITF to create StoragePoolClaim in k8s.
			spcGot, err := citfInstance.K8S.CreateStoragePoolClaim(spcObject)
			Expect(err).To(BeNil())
			// We expect nil error.

			// We expect 3 cstorPool objects.
			var maxRetry int
			var cspCount int
			maxRetry = 10
			for i := 0; i < maxRetry; i++ {
				cspCount, err = getCstorPoolCount(spcGot.Name, citfInstance)
				if err != nil {
					break
				}
				if cspCount == 3 {
					break
				}
				time.Sleep(time.Second * 5)
			}
			Expect(cspCount).To(Equal(3))
			// We expect 3 pool deployments.
			var deployCount int
			maxRetry = 10
			for i := 0; i < maxRetry; i++ {
				deployCount, err = getPoolDeployCount(spcGot.Name, citfInstance)
				if err != nil {
					break
				}
				if deployCount == 3 {
					break
				}
				time.Sleep(time.Second * 5)
			}
			Expect(cspCount).To(Equal(3))
			// We expect 3 storagePool objects.
			var spCount int
			maxRetry = 10
			for i := 0; i < maxRetry; i++ {
				spCount, err = getStoragePoolCount(spcGot.Name, citfInstance)
				if err != nil {
					break
				}
				if spCount == 3 {
					break
				}
				time.Sleep(time.Second * 5)
			}
			Expect(spCount).To(Equal(3))

			// We expect 'online' status on all the three cstorPool objects(i.e. 3 online counts)
			var onlineCspCount int
			maxRetry = 10
			for i := 0; i < maxRetry; i++ {
				onlineCspCount, err = getCstorPoolStatus(spcGot.Name, citfInstance)
				if err != nil {
					break
				}
				if onlineCspCount == 3 {
					break
				}
				time.Sleep(time.Second * 5)
			}
			Expect(onlineCspCount).To(Equal(3))

		})
	})
	// Test Case #2 (sparse-mirrored-auto-spc) Type:positive
	/*When("We apply sparse-mirrored-auto spc yaml with maxPool count equal to 3 on a k8s cluster having at least 3 capable node", func() {
		It("pool resources count should be 3 with no error and online status", func() {
			// Create a storage pool claim object
			// Call CITF to create StoragePoolClaim in k8s.
			// We expect nil error.
			// We expect 3 cstorPool objects.
			// We expect 3 pool deployments.
			// We expect 3 storagePool objects.
			// We expect 'online' status on all the three cstorPool objects(i.e. 3 online counts)

		})
	})*/

	// TODo: Add more test cases. Refer to following design doc
	// https://docs.google.com/document/d/1QAYK-Bsehc7v66kscXCiMJ7_pTIjzNmwyl43tF92gWA/edit
})
<commit_msg>Adds IT for sparse-mirrored-auto(dynamic) spc with maxpool count 0 (2190) (#732)<commit_after>/*
Copyright 2018 The OpenEBS Authors

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

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

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

package cstorpoolit

import (
	"testing"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	citf "github.com/openebs/CITF"
	citfoptions "github.com/openebs/CITF/citf_options"
	apis "github.com/openebs/CITF/pkg/apis/openebs.io/v1alpha1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// TestIntegrationCstorPool function instantiate the cstor pool test suite.
func TestIntegrationCstorPool(t *testing.T) {
	// RegisterFailHandler is used to register failed test cases and produce readable output.
	RegisterFailHandler(Fail)
	// RunSpecs runs all the test cases in the suite.
	RunSpecs(t, "Cstor pool integration test suite")
}

// Create an instance of CITF to use the inbuilt functions that will help
// communicating with the kube-apiserver.
var citfInstance, err = citf.NewCITF(&citfoptions.CreateOptions{
	// K8SInclude is true to get the kube-config from the machine where the suite is running.
	// Kube-config is a config file that establishes communication to the k8s cluster.
	K8SInclude: true,
})

// ToDo: Set up cluster environment before runninng all test cases ( i.e. BeforeSuite)
// The environment set up by BeforeSuite is going to persist for all
// the test cases under run

//var _ = BeforeSuite(func() {
//	//var err error
//	//
//	//Expect(err).NotTo(HaveOccurred())
//})

// ToDo: Set up tear down of cluster environment ( i.e Aftersuite)

// ToDo: Set up cluster environment before every test cases that will be run (i.e. preRunHook)
// ToDo: Reset cluster environment after every test cases that will be run  ( i.e postRunHook)
var _ = Describe("Integration Test", func() {
	// Test Case #1 (sparse-striped-auto-spc). Type : Positive
	When("We apply sparse-striped-auto spc yaml with maxPool count equal to 3 on a k8s cluster having at least 3 capable node", func() {
		It("pool resources count should be 3 with no error and online status", func() {
			// TODO: Create a generic util function in utils.go to convert yaml into go object.
			// ToDo: More POC regarding this util converter function.
			// Functions generic to both cstor-pool and cstor-vol should go inside common directory

			// 1.Read SPC yaml form a file.
			// 2.Convert SPC yaml to json.
			// 3.Marshall json to SPC go object.

			// Create a storage pool claim object
			spcObject := &apis.StoragePoolClaim{
				ObjectMeta: metav1.ObjectMeta{
					Name: "disk-claim-auto",
				},
				Spec: apis.StoragePoolClaimSpec{
					Name:     "sparse-claim-auto",
					Type:     "sparse",
					MaxPools: 3,
					PoolSpec: apis.CStorPoolAttr{
						PoolType: "striped",
					},
				},
			}
			// Call CITF to create StoragePoolClaim in k8s.
			spcGot, err := citfInstance.K8S.CreateStoragePoolClaim(spcObject)
			Expect(err).To(BeNil())
			// We expect nil error.

			// We expect 3 cstorPool objects.
			var maxRetry int
			var cspCount int
			maxRetry = 10
			for i := 0; i < maxRetry; i++ {
				cspCount, err = getCstorPoolCount(spcGot.Name, citfInstance)
				if err != nil {
					break
				}
				if cspCount == 3 {
					break
				}
				time.Sleep(time.Second * 5)
			}
			Expect(cspCount).To(Equal(3))
			// We expect 3 pool deployments.
			var deployCount int
			maxRetry = 10
			for i := 0; i < maxRetry; i++ {
				deployCount, err = getPoolDeployCount(spcGot.Name, citfInstance)
				if err != nil {
					break
				}
				if deployCount == 3 {
					break
				}
				time.Sleep(time.Second * 5)
			}
			Expect(cspCount).To(Equal(3))
			// We expect 3 storagePool objects.
			var spCount int
			maxRetry = 10
			for i := 0; i < maxRetry; i++ {
				spCount, err = getStoragePoolCount(spcGot.Name, citfInstance)
				if err != nil {
					break
				}
				if spCount == 3 {
					break
				}
				time.Sleep(time.Second * 5)
			}
			Expect(spCount).To(Equal(3))

			// We expect 'online' status on all the three cstorPool objects(i.e. 3 online counts)
			var onlineCspCount int
			maxRetry = 10
			for i := 0; i < maxRetry; i++ {
				onlineCspCount, err = getCstorPoolStatus(spcGot.Name, citfInstance)
				if err != nil {
					break
				}
				if onlineCspCount == 3 {
					break
				}
				time.Sleep(time.Second * 5)
			}
			Expect(onlineCspCount).To(Equal(3))

		})
	})

	// Test Case #2 (sparse-mirrored-auto-spc). Type : Negative
	When("We apply sparse-mirrored-auto spc yaml with maxPool count equal to 0 on a k8s cluster", func() {
		It("pool resources count should be 0 with no error and online status", func() {
			// TODO: Create a generic util function in utils.go to convert yaml into go object.
			// ToDo: More POC regarding this util converter function.
			// Functions generic to both cstor-pool and cstor-vol should go inside common directory

			// 1.Read SPC yaml form a file.
			// 2.Convert SPC yaml to json.
			// 3.Marshall json to SPC go object.

			// Create a storage pool claim object
			spcObject := &apis.StoragePoolClaim{
				ObjectMeta: metav1.ObjectMeta{
					Name: "disk-claim-auto",
				},
				Spec: apis.StoragePoolClaimSpec{
					Name:     "sparse-claim-auto",
					Type:     "sparse",
					MaxPools: 0,
					PoolSpec: apis.CStorPoolAttr{
						PoolType: "mirrored",
					},
				},
			}
			// Call CITF to create StoragePoolClaim in k8s.
			spcGot, err := citfInstance.K8S.CreateStoragePoolClaim(spcObject)
			Expect(err).To(BeNil())
			// We expect nil error.

			// We expect 0 cstorPool objects.
			var maxRetry int
			var cspCount int
			maxRetry = 10
			for i := 0; i < maxRetry; i++ {
				cspCount, err = getCstorPoolCount(spcGot.Name, citfInstance)
				if err != nil {
					break
				}
				if cspCount == 0 {
					break
				}
				time.Sleep(time.Second * 5)
			}
			Expect(cspCount).To(Equal(0))
			// We expect 0 pool deployments.
			var deployCount int
			maxRetry = 10
			for i := 0; i < maxRetry; i++ {
				deployCount, err = getPoolDeployCount(spcGot.Name, citfInstance)
				if err != nil {
					break
				}
				if deployCount == 0 {
					break
				}
				time.Sleep(time.Second * 5)
			}
			Expect(cspCount).To(Equal(0))
			// We expect 0 storagePool objects.
			var spCount int
			maxRetry = 10
			for i := 0; i < maxRetry; i++ {
				spCount, err = getStoragePoolCount(spcGot.Name, citfInstance)
				if err != nil {
					break
				}
				if spCount == 0 {
					break
				}
				time.Sleep(time.Second * 5)
			}
			Expect(spCount).To(Equal(0))

			// We don't expect 'online' status on any of the cstorPool objects.
			var onlineCspCount int
			maxRetry = 10
			for i := 0; i < maxRetry; i++ {
				onlineCspCount, err = getCstorPoolStatus(spcGot.Name, citfInstance)
				if err != nil {
					break
				}
				if onlineCspCount == 0 {
					break
				}
				time.Sleep(time.Second * 5)
			}
			Expect(onlineCspCount).To(Equal(0))

		})
	})

	// TODo: Add more test cases. Refer to following design doc
	// https://docs.google.com/document/d/1QAYK-Bsehc7v66kscXCiMJ7_pTIjzNmwyl43tF92gWA/edit
})
<|endoftext|>


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

package http

import (
	"bytes"
	"encoding/base64"
	"encoding/binary"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/ThomsonReutersEikon/go-ntlm/ntlm"
	"github.com/pkg/errors"
)

// NTLMNegotiator is a http.Roundtripper that automatically converts basic authentication to
// NTLM authentication.
type NTLMNegotiator struct{ http.RoundTripper }

func (n NTLMNegotiator) RoundTrip(req *http.Request) (res *http.Response, err error) {
	// Use default round tripper if not provided
	rt := n.RoundTripper
	if rt == nil {
		rt = http.DefaultTransport
	}

	username, password, err := getCredentialsFromHeader(req.Header.Get("Authorization"))
	if err != nil {
		return nil, errors.Wrap(err, "get basic credentials from header failed")
	}

	// Save request body
	body := bytes.Buffer{}
	if req.Body != nil {
		_, err = body.ReadFrom(req.Body)
		if err != nil {
			return nil, err
		}

		if err := req.Body.Close(); err != nil {
			return nil, err
		}
		req.Body = ioutil.NopCloser(bytes.NewReader(body.Bytes()))
	}

	// Sending first request to get challenge data
	req.Header.Set("Authorization", "NTLM TlRMTVNTUAABAAAABoIIAAAAAAAAAAAAAAAAAAAAAAA=")
	res, err = rt.RoundTrip(req)
	if err != nil {
		return nil, err
	}
	if res.StatusCode != http.StatusUnauthorized {
		return res, err
	}

	if _, err := io.Copy(ioutil.Discard, res.Body); err != nil {
		return nil, err
	}
	if err := res.Body.Close(); err != nil {
		return nil, err
	}
	req.Body = ioutil.NopCloser(bytes.NewReader(body.Bytes()))

	// retrieve Www-Authenticate header from response
	ntlmChallenge := res.Header.Get("WWW-Authenticate")
	if ntlmChallenge == "" {
		return nil, errors.New("Empty WWW-Authenticate header")
	}

	challengeBytes, err := base64.StdEncoding.DecodeString(strings.Replace(ntlmChallenge, "NTLM ", "", -1))
	if err != nil {
		return nil, err
	}

	session, err := ntlm.CreateClientSession(ntlm.Version2, ntlm.ConnectionlessMode)
	if err != nil {
		return nil, err
	}

	session.SetUserInfo(username, password, "")

	// parse NTLM challenge
	challenge, err := ntlm.ParseChallengeMessage(challengeBytes)
	if err != nil {
		return nil, err
	}

	err = session.ProcessChallengeMessage(challenge)
	if err != nil {
		return nil, err
	}

	// authenticate user
	authenticate, err := session.GenerateAuthenticateMessage()
	if err != nil {
		return nil, err
	}

	// set NTLM Authorization header
	req.Header.Set("Authorization", "NTLM "+base64.StdEncoding.EncodeToString(authenticate.Bytes()))

	return rt.RoundTrip(req)
}

func getCredentialsFromHeader(header string) (string, string, error) {
	credBytes, err := base64.StdEncoding.DecodeString(strings.Replace(header, "Basic ", "", -1))
	if err != nil {
		return "", "", err
	}

	parts := strings.SplitN(string(credBytes), ":", 2)
	return parts[0], parts[1], nil
}

func ntlmHandler(username, password string) func(w http.ResponseWriter, r *http.Request) {
	challenges := make(map[string]*ntlm.ChallengeMessage)
	return func(w http.ResponseWriter, r *http.Request) {
		// Make sure there is some kind of authentication
		if r.Header.Get("Authorization") == "" {
			w.Header().Set("WWW-Authenticate", "NTLM")
			w.WriteHeader(401)
			return
		}

		// Parse the proxy authorization header
		auth := r.Header.Get("Authorization")
		parts := strings.SplitN(auth, " ", 2)
		authType := parts[0]
		authPayload := parts[1]

		// Filter out unsupported authentication methods
		if authType != "NTLM" {
			w.Header().Set("WWW-Authenticate", "NTLM")
			w.WriteHeader(401)
			return
		}

		// Decode base64 auth data and get NTLM message type
		rawAuthPayload, _ := base64.StdEncoding.DecodeString(authPayload)
		ntlmMessageType := binary.LittleEndian.Uint32(rawAuthPayload[8:12])

		// Handle NTLM negotiate message
		if ntlmMessageType == 1 {
			session, err := ntlm.CreateServerSession(ntlm.Version2, ntlm.ConnectionOrientedMode)
			if err != nil {
				return
			}

			session.SetUserInfo(username, password, "")

			challenge, err := session.GenerateChallengeMessage()
			if err != nil {
				return
			}

			challenges[r.RemoteAddr] = challenge

			authPayload := base64.StdEncoding.EncodeToString(challenge.Bytes())

			w.Header().Set("WWW-Authenticate", "NTLM "+authPayload)
			w.WriteHeader(401)

			return
		}

		if ntlmMessageType == 3 {
			challenge := challenges[r.RemoteAddr]
			if challenge == nil {
				w.Header().Set("WWW-Authenticate", "NTLM")
				w.WriteHeader(401)
				return
			}

			msg, err := ntlm.ParseAuthenticateMessage(rawAuthPayload, 2)
			if err != nil {
				msg2, err := ntlm.ParseAuthenticateMessage(rawAuthPayload, 1)

				if err != nil {
					return
				}

				session, err := ntlm.CreateServerSession(ntlm.Version1, ntlm.ConnectionOrientedMode)
				if err != nil {
					return
				}

				session.SetServerChallenge(challenge.ServerChallenge)
				session.SetUserInfo(username, password, "")

				err = session.ProcessAuthenticateMessage(msg2)
				if err != nil {
					w.Header().Set("WWW-Authenticate", "NTLM")
					w.WriteHeader(401)
					return
				}
			} else {
				session, err := ntlm.CreateServerSession(ntlm.Version2, ntlm.ConnectionOrientedMode)
				if err != nil {
					return
				}

				session.SetServerChallenge(challenge.ServerChallenge)
				session.SetUserInfo(username, password, "")

				err = session.ProcessAuthenticateMessage(msg)
				if err != nil {
					w.Header().Set("WWW-Authenticate", "NTLM")
					w.WriteHeader(401)
					return
				}
			}
		}

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

package http

import (
	"bytes"
	"encoding/base64"
	"encoding/binary"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/ThomsonReutersEikon/go-ntlm/ntlm"
	"github.com/pkg/errors"
)

// NTLMNegotiator is a http.Roundtripper that automatically converts basic authentication to
// NTLM authentication.
type NTLMNegotiator struct{ http.RoundTripper }

func (n NTLMNegotiator) RoundTrip(req *http.Request) (res *http.Response, err error) {
	// Use default round tripper if not provided
	rt := n.RoundTripper
	if rt == nil {
		rt = http.DefaultTransport
	}

	username, password, err := getCredentialsFromHeader(req.Header.Get("Authorization"))
	if err != nil {
		return nil, errors.Wrap(err, "get basic credentials from header failed")
	}

	// Save request body
	body := bytes.Buffer{}
	if req.Body != nil {
		_, err = body.ReadFrom(req.Body)
		if err != nil {
			return nil, err
		}

		if err := req.Body.Close(); err != nil {
			return nil, err
		}
		req.Body = ioutil.NopCloser(bytes.NewReader(body.Bytes()))
	}

	// Sending first request to get challenge data
	req.Header.Set("Authorization", "NTLM TlRMTVNTUAABAAAABoIIAAAAAAAAAAAAAAAAAAAAAAA=")
	res, err = rt.RoundTrip(req)
	if err != nil {
		return nil, err
	}
	if res.StatusCode != http.StatusUnauthorized {
		return res, err
	}

	if _, err := io.Copy(ioutil.Discard, res.Body); err != nil {
		return nil, err
	}
	if err := res.Body.Close(); err != nil {
		return nil, err
	}
	req.Body = ioutil.NopCloser(bytes.NewReader(body.Bytes()))

	// retrieve Www-Authenticate header from response
	ntlmChallenge := res.Header.Get("WWW-Authenticate")
	if ntlmChallenge == "" {
		return nil, errors.New("Empty WWW-Authenticate header")
	}

	challengeBytes, err := base64.StdEncoding.DecodeString(strings.Replace(ntlmChallenge, "NTLM ", "", -1))
	if err != nil {
		return nil, err
	}

	session, err := ntlm.CreateClientSession(ntlm.Version2, ntlm.ConnectionlessMode)
	if err != nil {
		return nil, err
	}

	session.SetUserInfo(username, password, "")

	// parse NTLM challenge
	challenge, err := ntlm.ParseChallengeMessage(challengeBytes)
	if err != nil {
		return nil, err
	}

	err = session.ProcessChallengeMessage(challenge)
	if err != nil {
		return nil, err
	}

	// authenticate user
	authenticate, err := session.GenerateAuthenticateMessage()
	if err != nil {
		return nil, err
	}

	// set NTLM Authorization header
	req.Header.Set("Authorization", "NTLM "+base64.StdEncoding.EncodeToString(authenticate.Bytes()))

	return rt.RoundTrip(req)
}

func getCredentialsFromHeader(header string) (string, string, error) {
	credBytes, err := base64.StdEncoding.DecodeString(strings.Replace(header, "Basic ", "", -1))
	if err != nil {
		return "", "", err
	}

	parts := strings.SplitN(string(credBytes), ":", 2)
	return parts[0], parts[1], nil
}

func ntlmHandler(username, password string) func(w http.ResponseWriter, r *http.Request) {
	challenges := make(map[string]*ntlm.ChallengeMessage)
	return func(w http.ResponseWriter, r *http.Request) {
		// Make sure there is some kind of authentication
		if r.Header.Get("Authorization") == "" {
			w.Header().Set("WWW-Authenticate", "NTLM")
			w.WriteHeader(401)
			return
		}

		// Parse the proxy authorization header
		auth := r.Header.Get("Authorization")
		parts := strings.SplitN(auth, " ", 2)
		authType := parts[0]
		authPayload := parts[1]

		// Filter out unsupported authentication methods
		if authType != "NTLM" {
			w.Header().Set("WWW-Authenticate", "NTLM")
			w.WriteHeader(401)
			return
		}

		// Decode base64 auth data and get NTLM message type
		rawAuthPayload, _ := base64.StdEncoding.DecodeString(authPayload)
		ntlmMessageType := binary.LittleEndian.Uint32(rawAuthPayload[8:12])

		// Handle NTLM negotiate message
		if ntlmMessageType == 1 {
			session, err := ntlm.CreateServerSession(ntlm.Version2, ntlm.ConnectionOrientedMode)
			if err != nil {
				return
			}

			session.SetUserInfo(username, password, "")

			challenge, err := session.GenerateChallengeMessage()
			if err != nil {
				return
			}

			challenges[r.RemoteAddr] = challenge

			authPayload := base64.StdEncoding.EncodeToString(challenge.Bytes())

			w.Header().Set("WWW-Authenticate", "NTLM "+authPayload)
			w.WriteHeader(401)

			return
		}

		if ntlmMessageType == 3 {
			challenge := challenges[r.RemoteAddr]
			if challenge == nil {
				w.Header().Set("WWW-Authenticate", "NTLM")
				w.WriteHeader(401)
				return
			}

			msg, err := ntlm.ParseAuthenticateMessage(rawAuthPayload, 2)
			if err != nil {
				msg2, err := ntlm.ParseAuthenticateMessage(rawAuthPayload, 1)

				if err != nil {
					return
				}

				session, err := ntlm.CreateServerSession(ntlm.Version1, ntlm.ConnectionOrientedMode)
				if err != nil {
					return
				}

				session.SetServerChallenge(challenge.ServerChallenge)
				session.SetUserInfo(username, password, "")

				err = session.ProcessAuthenticateMessage(msg2)
				if err != nil {
					w.Header().Set("WWW-Authenticate", "NTLM")
					w.WriteHeader(401)
					return
				}
			} else {
				session, err := ntlm.CreateServerSession(ntlm.Version2, ntlm.ConnectionOrientedMode)
				if err != nil {
					return
				}

				session.SetServerChallenge(challenge.ServerChallenge)
				session.SetUserInfo(username, password, "")

				err = session.ProcessAuthenticateMessage(msg)
				if err != nil {
					w.Header().Set("WWW-Authenticate", "NTLM")
					w.WriteHeader(401)
					return
				}
			}
		}

		data := "authenticated"
		w.Header().Set("Content-Length", fmt.Sprint(len(data)))
		fmt.Fprint(w, data)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"golang.org/x/crypto/openpgp"
	"io/ioutil"
	"log"
	"os"
	"os/user"
	"path/filepath"
	"strings"
)

const jaegerTemplateExtension = ".jgrt"
const jaegerDBExtension = ".jgrdb"
const jaegerDBDescription = "JaegerDB - Jaeger database management program\n\nJaeger is a JSON encoded GPG encrypted key value store. It is useful for separating development with operations and keeping configuration files secure."
const jaegerQuote = "\"Stacker Pentecost: Haven't you heard Mr. Beckett? The world is coming to an end. So where would you rather die? Here? Or in a Jaeger!\" - Pacific Rim"
const jaegerDBRecommendedUsage = "RECOMMENDED:\n    jaegerdb -j file.txt.jgrdb -a \"Field1\" -v \"Secret value\"\n\nThis will run JaegerDB with the default options and assume the following:\n    Keyring file: ~/.gnupg/jaeger_pubring.gpg"

var debug debugging = false

type debugging bool

func (d debugging) Printf(format string, args ...interface{}) {
	// From: https://groups.google.com/forum/#!msg/golang-nuts/gU7oQGoCkmg/BNIl-TqB-4wJ
	if d {
		log.Printf(format, args...)
	}
}

type Data struct {
	Properties []Property
}

type Property struct {
	Name           string `json:"Name"`
	EncryptedValue string `json:"EncryptedValue"`
}

func main() {
	// Define flags
	// TODO: View individual property and unencrypted value. 'get'
	var (
		addKey         = flag.String("a", "", "Add property")
		changeKey      = flag.String("c", "", "Change property")
		debugFlag      = flag.Bool("d", false, "Enable Debug")
		deleteKey      = flag.String("delete", "", "Delete property")
		initializeFlag = flag.Bool("init", false, "Create an initial blank JSON GPG database file")
		jsonGPGDB      = flag.String("j", "", "JSON GPG database file. eg. file.txt.jgrdb")
		keyringFile    = flag.String("k", "", "Keyring file. Public key in ASCII armored format. eg. pubring.asc")
		value          = flag.String("v", "", "Value for property to use")
	)

	flag.Usage = func() {
		fmt.Printf("%s\n%s\n\n%s\n\n", jaegerDBDescription, jaegerQuote, jaegerDBRecommendedUsage)
		fmt.Fprintf(os.Stderr, "OPTIONS:\n")
		flag.PrintDefaults()
	}

	flag.Parse()

	if *debugFlag {
		debug = true
	}

	if *jsonGPGDB == "" {
		assumedJaegerDB, err := checkExistsJaegerDB()
		if err != nil {
			flag.Usage()
			log.Fatalf("\n\nError: %s", err)
		}
		*jsonGPGDB = assumedJaegerDB
	}

	if *initializeFlag {
		err := initializeJSONGPGDB(jsonGPGDB)
		if err != nil {
			log.Fatal(err)
		} else {
			fmt.Println("Initialized JSON GPG database and wrote to file:", *jsonGPGDB)
			os.Exit(0)
		}
	}

	if *deleteKey != "" {
		err := deleteKeyJaegerDB(deleteKey, jsonGPGDB)
		if err != nil {
			log.Fatal(err)
		} else {
			fmt.Println("Deleted property and wrote to file:", *jsonGPGDB)
			os.Exit(0)
		}
	}

	var entitylist openpgp.EntityList

	if *keyringFile == "" {
		_, entitylist = processPublicKeyRing()
	} else {
		_, entitylist = processArmoredKeyRingFile(keyringFile)
	}

	if *addKey != "" {
		if *value == "" {
			flag.Usage()
			log.Fatalf("\n\nError: No value for add key operation specified")
		}
		err := addKeyJaegerDB(addKey, value, jsonGPGDB, entitylist)
		if err != nil {
			log.Fatal(err)
		} else {
			fmt.Println("Added property and wrote to file:", *jsonGPGDB)
			os.Exit(0)
		}
	}

	if *changeKey != "" {
		if *value == "" {
			flag.Usage()
			log.Fatalf("\n\nError: No value for change key operation specified")
		}
		err := changeKeyJaegerDB(changeKey, value, jsonGPGDB, entitylist)
		if err != nil {
			log.Fatal(err)
		} else {
			fmt.Println("Changed property and wrote to file:", *jsonGPGDB)
			os.Exit(0)
		}
	}

	if *deleteKey == "" && *addKey == "" && *changeKey == "" {
		log.Fatalf("\n\nError: No JSON GPG database operations specified")
	}

}

func checkExistsJaegerT() (string, error) {
	// Check that one template file is in the current directory and use that file
	files, err := filepath.Glob("*.jgrt")
	if err != nil {
		return "", err
	}
	if len(files) == 1 {
		return files[0], nil
	}
	return "", fmt.Errorf("No input template file specified")
}

func checkExistsJaegerDB() (string, error) {
	// If no JSON GPG database file is explicitly specified, check that one JSON GPG database file is in the
	// current directory and use that file
	files, err := filepath.Glob("*.jgrdb")
	if err != nil {
		return "", err
	}
	if len(files) == 1 {
		return files[0], nil
	}

	assumedTemplate, err := checkExistsJaegerT()

	if assumedTemplate != "" {
		basefilename := strings.TrimSuffix(assumedTemplate, jaegerTemplateExtension)
		jsonGPGDB := fmt.Sprintf("%v%v", basefilename, jaegerDBExtension)
		return jsonGPGDB, nil
	}

	return "", fmt.Errorf("Please specify a JSON GPG database file")
}

func initializeJSONGPGDB(jsonGPGDB *string) error {
	if _, err := os.Stat(*jsonGPGDB); err == nil {
		return fmt.Errorf("ERR: File already exists: %v", *jsonGPGDB)
	}

	var newP []Property

	newData := Data{newP}

	bytes, err := json.MarshalIndent(newData, "", "    ")
	if err != nil {
		return fmt.Errorf("error:", err)
	}

	debug.Printf("b: %v", string(bytes))

	// Writing file
	// To handle large files, use a file buffer: http://stackoverflow.com/a/9739903/603745
	if err := ioutil.WriteFile(*jsonGPGDB, bytes, 0644); err != nil {
		return fmt.Errorf("error:", err)
	}

	return nil
}

func encodeBase64EncryptedMessage(s string, entitylist openpgp.EntityList) string {
	// Encrypt message using public key and then encode with base64
	debug.Printf("entitylist: #%v", entitylist)
	buf := new(bytes.Buffer)
	w, err := openpgp.Encrypt(buf, entitylist, nil, nil, nil)
	if err != nil {
		log.Fatalln("ERR: Error encrypting message - ", err)
	}

	_, err = w.Write([]byte(s))
	if err != nil {
	}
	err = w.Close()
	if err != nil {
	}

	// Output as base64 encoded string
	bytes, err := ioutil.ReadAll(buf)
	str := base64.StdEncoding.EncodeToString(bytes)

	debug.Printf("Public key encrypted message (base64 encoded): %v", str)

	return str
}

func processPublicKeyRing() (entity *openpgp.Entity, entitylist openpgp.EntityList) {
	// TODO: Handle a specified recipient
	// Get default public keyring location
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}

	jaegerPublicKeyRing := fmt.Sprintf("%v/.gnupg/jaeger_pubring.gpg", usr.HomeDir)
	publicKeyRing := ""

	if _, err := os.Stat(jaegerPublicKeyRing); err == nil {
		publicKeyRing = jaegerPublicKeyRing
	} else {
		publicKeyRing = fmt.Sprintf("%v/.gnupg/pubring.gpg", usr.HomeDir)
	}

	debug.Printf("publicKeyRing file:", publicKeyRing)
	publicKeyRingBuffer, err := os.Open(publicKeyRing)
	if err != nil {
		panic(err)
	}
	entitylist, err = openpgp.ReadKeyRing(publicKeyRingBuffer)
	if err != nil {
		log.Fatal(err)
	}

	entity = entitylist[0]
	debug.Printf("Public key default keyring:", entity.Identities)

	return entity, entitylist
}

func processArmoredKeyRingFile(keyringFile *string) (entity *openpgp.Entity, entitylist openpgp.EntityList) {
	keyringFileBuffer, err := os.Open(*keyringFile)
	if err != nil {
		log.Fatalln("ERROR: Unable to read keyring file")
	}
	entitylist, err = openpgp.ReadArmoredKeyRing(keyringFileBuffer)
	if err != nil {
		log.Fatal(err)
	}
	entity = entitylist[0]
	debug.Printf("Public key from ASCII armored string:", entity.Identities)

	return entity, entitylist
}

func addKeyJaegerDB(key *string, value *string, jsonGPGDB *string, entitylist openpgp.EntityList) error {
	// json handling
	jsonGPGDBBuffer, err := ioutil.ReadFile(*jsonGPGDB)
	if err != nil {
		return fmt.Errorf("ERROR: Unable to read JSON GPG DB file")
	}

	var j Data
	if err := json.Unmarshal(jsonGPGDBBuffer, &j); err != nil {
		return fmt.Errorf("error:", err)
	}
	debug.Printf("json unmarshal: %v", j)

	found := false

	var newP []Property

	p := Property{Name: *key, EncryptedValue: encodeBase64EncryptedMessage(*value, entitylist)}

	// Search
	for i, _ := range j.Properties {
		property := &j.Properties[i]
		debug.Printf("i: %v, Name: %#v, EncryptedValue: %#v\n", i, property.Name, property.EncryptedValue)
		if property.Name == *key {
			j.Properties[i] = p
			found = true
			break
		}
	}

	if found {
		return fmt.Errorf("\n\nError: Property '%s' already exists.", *key)
	}

	newP = append(j.Properties, p)

	debug.Printf("new properties: %v", newP)

	newData := Data{newP}

	bytes, err := json.MarshalIndent(newData, "", "    ")
	if err != nil {
		return fmt.Errorf("error:", err)
	}

	debug.Printf("b: %v", string(bytes))

	// Writing file
	// To handle large files, use a file buffer: http://stackoverflow.com/a/9739903/603745
	if err := ioutil.WriteFile(*jsonGPGDB, bytes, 0644); err != nil {
		return fmt.Errorf("error:", err)
	}

	return nil
}

func changeKeyJaegerDB(key *string, value *string, jsonGPGDB *string, entitylist openpgp.EntityList) error {
	// json handling
	jsonGPGDBBuffer, err := ioutil.ReadFile(*jsonGPGDB)
	if err != nil {
		return fmt.Errorf("ERROR: Unable to read JSON GPG DB file")
	}

	var j Data
	if err := json.Unmarshal(jsonGPGDBBuffer, &j); err != nil {
		return fmt.Errorf("error:", err)
	}
	debug.Printf("json unmarshal: %v", j)

	found := false

	// New property to replace the old
	p := Property{Name: *key, EncryptedValue: encodeBase64EncryptedMessage(*value, entitylist)}

	// Search and replace
	for i, _ := range j.Properties {
		property := &j.Properties[i]
		debug.Printf("i: %v, Name: %#v, EncryptedValue: %#v\n", i, property.Name, property.EncryptedValue)
		if property.Name == *key {
			j.Properties[i] = p
			found = true
			break
		}
	}

	if !found {
		return fmt.Errorf("\n\nError: Property '%s' not found.", *key)
	}

	bytes, err := json.MarshalIndent(j, "", "    ")
	if err != nil {
		return fmt.Errorf("error:", err)
	}

	debug.Printf("b: %v", string(bytes))

	// Writing file
	// To handle large files, use a file buffer: http://stackoverflow.com/a/9739903/603745
	if err := ioutil.WriteFile(*jsonGPGDB, bytes, 0644); err != nil {
		return fmt.Errorf("error:", err)
	}

	return nil
}

func deleteKeyJaegerDB(key *string, jsonGPGDB *string) error {
	debug.Printf("deleteKeyJaegerDB key: %v", *key)

	// json handling
	jsonGPGDBBuffer, err := ioutil.ReadFile(*jsonGPGDB)
	if err != nil {
		return fmt.Errorf("ERROR: Unable to read JSON GPG DB file")
	}

	var j Data
	if err := json.Unmarshal(jsonGPGDBBuffer, &j); err != nil {
		return fmt.Errorf("error:", err)
	}
	debug.Printf("json unmarshal: %v", j)

	var newP []Property
	found := false

	for i, _ := range j.Properties {
		property := &j.Properties[i]
		debug.Printf("i: %v, Name: %#v, EncryptedValue: %#v\n", i, property.Name, property.EncryptedValue)
		if property.Name == *key {
			// https://code.google.com/p/go-wiki/wiki/SliceTricks
			newP = j.Properties[:i+copy(j.Properties[i:], j.Properties[i+1:])]
			found = true
			break
		}
	}

	if !found {
		return fmt.Errorf("\n\nError: Property '%s' not found.", *key)
	}

	debug.Printf("new properties: %v", newP)

	newData := Data{newP}

	bytes, err := json.MarshalIndent(newData, "", "    ")
	if err != nil {
		return fmt.Errorf("error:", err)
	}

	debug.Printf("b: %v", string(bytes))

	// Writing file
	// To handle large files, use a file buffer: http://stackoverflow.com/a/9739903/603745
	if err := ioutil.WriteFile(*jsonGPGDB, bytes, 0644); err != nil {
		return fmt.Errorf("error:", err)
	}

	return nil
}
<commit_msg>Minor code cleanups<commit_after>package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"golang.org/x/crypto/openpgp"
	"io/ioutil"
	"log"
	"os"
	"os/user"
	"path/filepath"
	"strings"
)

const jaegerTemplateExtension = ".jgrt"
const jaegerDBExtension = ".jgrdb"
const jaegerDBDescription = "JaegerDB - Jaeger database management program\n\nJaeger is a JSON encoded GPG encrypted key value store. It is useful for separating development with operations and keeping configuration files secure."
const jaegerQuote = "\"Stacker Pentecost: Haven't you heard Mr. Beckett? The world is coming to an end. So where would you rather die? Here? Or in a Jaeger!\" - Pacific Rim"
const jaegerDBRecommendedUsage = "RECOMMENDED:\n    jaegerdb -j file.txt.jgrdb -a \"Field1\" -v \"Secret value\"\n\nThis will run JaegerDB with the default options and assume the following:\n    Keyring file: ~/.gnupg/jaeger_pubring.gpg"

var debug debugging = false

type debugging bool

func (d debugging) Printf(format string, args ...interface{}) {
	// From: https://groups.google.com/forum/#!msg/golang-nuts/gU7oQGoCkmg/BNIl-TqB-4wJ
	if d {
		log.Printf(format, args...)
	}
}

type Data struct {
	Properties []Property
}

type Property struct {
	Name           string `json:"Name"`
	EncryptedValue string `json:"EncryptedValue"`
}

func main() {
	// Define flags
	// TODO: View individual property and unencrypted value. 'get'
	var (
		addKey         = flag.String("a", "", "Add property")
		changeKey      = flag.String("c", "", "Change property")
		debugFlag      = flag.Bool("d", false, "Enable Debug")
		deleteKey      = flag.String("delete", "", "Delete property")
		initializeFlag = flag.Bool("init", false, "Create an initial blank JSON GPG database file")
		jsonGPGDB      = flag.String("j", "", "JSON GPG database file. eg. file.txt.jgrdb")
		keyringFile    = flag.String("k", "", "Keyring file. Public key in ASCII armored format. eg. pubring.asc")
		value          = flag.String("v", "", "Value for property to use")
	)

	flag.Usage = func() {
		fmt.Printf("%s\n%s\n\n%s\n\n", jaegerDBDescription, jaegerQuote, jaegerDBRecommendedUsage)
		fmt.Fprintf(os.Stderr, "OPTIONS:\n")
		flag.PrintDefaults()
	}

	flag.Parse()

	if *debugFlag {
		debug = true
	}

	if *jsonGPGDB == "" {
		assumedJaegerDB, err := checkExistsJaegerDB()
		if err != nil {
			flag.Usage()
			log.Fatalf("\n\nError: %s", err)
		}
		*jsonGPGDB = assumedJaegerDB
	}

	if *initializeFlag {
		err := initializeJSONGPGDB(jsonGPGDB)
		if err != nil {
			log.Fatal(err)
		} else {
			fmt.Println("Initialized JSON GPG database and wrote to file:", *jsonGPGDB)
			os.Exit(0)
		}
	}

	if *deleteKey != "" {
		err := deleteKeyJaegerDB(deleteKey, jsonGPGDB)
		if err != nil {
			log.Fatal(err)
		} else {
			fmt.Println("Deleted property and wrote to file:", *jsonGPGDB)
			os.Exit(0)
		}
	}

	var entitylist openpgp.EntityList

	if *keyringFile == "" {
		_, entitylist = processPublicKeyRing()
	} else {
		_, entitylist = processArmoredKeyRingFile(keyringFile)
	}

	if *addKey != "" {
		if *value == "" {
			flag.Usage()
			log.Fatalf("\n\nError: No value for add key operation specified")
		}
		err := addKeyJaegerDB(addKey, value, jsonGPGDB, entitylist)
		if err != nil {
			log.Fatal(err)
		} else {
			fmt.Println("Added property and wrote to file:", *jsonGPGDB)
			os.Exit(0)
		}
	}

	if *changeKey != "" {
		if *value == "" {
			flag.Usage()
			log.Fatalf("\n\nError: No value for change key operation specified")
		}
		err := changeKeyJaegerDB(changeKey, value, jsonGPGDB, entitylist)
		if err != nil {
			log.Fatal(err)
		} else {
			fmt.Println("Changed property and wrote to file:", *jsonGPGDB)
			os.Exit(0)
		}
	}

	if *deleteKey == "" && *addKey == "" && *changeKey == "" {
		log.Fatalf("\n\nError: No JSON GPG database operations specified")
	}

}

func checkExistsJaegerT() (string, error) {
	// Check that one template file is in the current directory and use that file
	files, err := filepath.Glob("*.jgrt")
	if err != nil {
		return "", err
	}
	if len(files) == 1 {
		return files[0], nil
	}
	return "", fmt.Errorf("No input template file specified")
}

func checkExistsJaegerDB() (string, error) {
	// If no JSON GPG database file is explicitly specified, check that one JSON GPG database file is in the
	// current directory and use that file
	files, err := filepath.Glob("*.jgrdb")
	if err != nil {
		return "", err
	}
	if len(files) == 1 {
		return files[0], nil
	}

	assumedTemplate, err := checkExistsJaegerT()

	if assumedTemplate != "" {
		basefilename := strings.TrimSuffix(assumedTemplate, jaegerTemplateExtension)
		jsonGPGDB := fmt.Sprintf("%v%v", basefilename, jaegerDBExtension)
		return jsonGPGDB, nil
	}

	return "", fmt.Errorf("Please specify a JSON GPG database file")
}

func initializeJSONGPGDB(jsonGPGDB *string) error {
	if _, err := os.Stat(*jsonGPGDB); err == nil {
		return fmt.Errorf("ERR: File already exists: %v", *jsonGPGDB)
	}

	var newP []Property

	newData := Data{newP}

	bytes, err := json.MarshalIndent(newData, "", "    ")
	if err != nil {
		return fmt.Errorf("error:", err)
	}

	debug.Printf("b: %v", string(bytes))

	// Writing file
	// To handle large files, use a file buffer: http://stackoverflow.com/a/9739903/603745
	if err := ioutil.WriteFile(*jsonGPGDB, bytes, 0644); err != nil {
		return fmt.Errorf("error:", err)
	}

	return nil
}

func encodeBase64EncryptedMessage(s string, entitylist openpgp.EntityList) string {
	// Encrypt message using public key and then encode with base64
	debug.Printf("entitylist: #%v", entitylist)
	buf := new(bytes.Buffer)
	w, err := openpgp.Encrypt(buf, entitylist, nil, nil, nil)
	if err != nil {
		log.Fatalln("ERR: Error encrypting message - ", err)
	}

	_, err = w.Write([]byte(s))
	err = w.Close()

	// Output as base64 encoded string
	bytes, err := ioutil.ReadAll(buf)
	str := base64.StdEncoding.EncodeToString(bytes)

	debug.Printf("Public key encrypted message (base64 encoded): %v", str)

	return str
}

func processPublicKeyRing() (entity *openpgp.Entity, entitylist openpgp.EntityList) {
	// TODO: Handle a specified recipient
	// Get default public keyring location
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}

	jaegerPublicKeyRing := fmt.Sprintf("%v/.gnupg/jaeger_pubring.gpg", usr.HomeDir)
	publicKeyRing := ""

	if _, err := os.Stat(jaegerPublicKeyRing); err == nil {
		publicKeyRing = jaegerPublicKeyRing
	} else {
		publicKeyRing = fmt.Sprintf("%v/.gnupg/pubring.gpg", usr.HomeDir)
	}

	debug.Printf("publicKeyRing file:", publicKeyRing)
	publicKeyRingBuffer, err := os.Open(publicKeyRing)
	if err != nil {
		panic(err)
	}
	entitylist, err = openpgp.ReadKeyRing(publicKeyRingBuffer)
	if err != nil {
		log.Fatal(err)
	}

	entity = entitylist[0]
	debug.Printf("Public key default keyring:", entity.Identities)

	return entity, entitylist
}

func processArmoredKeyRingFile(keyringFile *string) (entity *openpgp.Entity, entitylist openpgp.EntityList) {
	keyringFileBuffer, err := os.Open(*keyringFile)
	if err != nil {
		log.Fatalln("ERROR: Unable to read keyring file")
	}
	entitylist, err = openpgp.ReadArmoredKeyRing(keyringFileBuffer)
	if err != nil {
		log.Fatal(err)
	}
	entity = entitylist[0]
	debug.Printf("Public key from ASCII armored string:", entity.Identities)

	return entity, entitylist
}

func addKeyJaegerDB(key *string, value *string, jsonGPGDB *string, entitylist openpgp.EntityList) error {
	// json handling
	jsonGPGDBBuffer, err := ioutil.ReadFile(*jsonGPGDB)
	if err != nil {
		return fmt.Errorf("ERROR: Unable to read JSON GPG DB file")
	}

	var j Data
	if err := json.Unmarshal(jsonGPGDBBuffer, &j); err != nil {
		return fmt.Errorf("error:", err)
	}
	debug.Printf("json unmarshal: %v", j)

	found := false

	var newP []Property

	p := Property{Name: *key, EncryptedValue: encodeBase64EncryptedMessage(*value, entitylist)}

	// Search
	for i := range j.Properties {
		property := &j.Properties[i]
		debug.Printf("i: %v, Name: %#v, EncryptedValue: %#v\n", i, property.Name, property.EncryptedValue)
		if property.Name == *key {
			j.Properties[i] = p
			found = true
			break
		}
	}

	if found {
		return fmt.Errorf("\n\nError: Property '%s' already exists.", *key)
	}

	newP = append(j.Properties, p)

	debug.Printf("new properties: %v", newP)

	newData := Data{newP}

	bytes, err := json.MarshalIndent(newData, "", "    ")
	if err != nil {
		return fmt.Errorf("error:", err)
	}

	debug.Printf("b: %v", string(bytes))

	// Writing file
	// To handle large files, use a file buffer: http://stackoverflow.com/a/9739903/603745
	if err := ioutil.WriteFile(*jsonGPGDB, bytes, 0644); err != nil {
		return fmt.Errorf("error:", err)
	}

	return nil
}

func changeKeyJaegerDB(key *string, value *string, jsonGPGDB *string, entitylist openpgp.EntityList) error {
	// json handling
	jsonGPGDBBuffer, err := ioutil.ReadFile(*jsonGPGDB)
	if err != nil {
		return fmt.Errorf("ERROR: Unable to read JSON GPG DB file")
	}

	var j Data
	if err := json.Unmarshal(jsonGPGDBBuffer, &j); err != nil {
		return fmt.Errorf("error:", err)
	}
	debug.Printf("json unmarshal: %v", j)

	found := false

	// New property to replace the old
	p := Property{Name: *key, EncryptedValue: encodeBase64EncryptedMessage(*value, entitylist)}

	// Search and replace
	for i := range j.Properties {
		property := &j.Properties[i]
		debug.Printf("i: %v, Name: %#v, EncryptedValue: %#v\n", i, property.Name, property.EncryptedValue)
		if property.Name == *key {
			j.Properties[i] = p
			found = true
			break
		}
	}

	if !found {
		return fmt.Errorf("\n\nError: Property '%s' not found.", *key)
	}

	bytes, err := json.MarshalIndent(j, "", "    ")
	if err != nil {
		return fmt.Errorf("error:", err)
	}

	debug.Printf("b: %v", string(bytes))

	// Writing file
	// To handle large files, use a file buffer: http://stackoverflow.com/a/9739903/603745
	if err := ioutil.WriteFile(*jsonGPGDB, bytes, 0644); err != nil {
		return fmt.Errorf("error:", err)
	}

	return nil
}

func deleteKeyJaegerDB(key *string, jsonGPGDB *string) error {
	debug.Printf("deleteKeyJaegerDB key: %v", *key)

	// json handling
	jsonGPGDBBuffer, err := ioutil.ReadFile(*jsonGPGDB)
	if err != nil {
		return fmt.Errorf("ERROR: Unable to read JSON GPG DB file")
	}

	var j Data
	if err := json.Unmarshal(jsonGPGDBBuffer, &j); err != nil {
		return fmt.Errorf("error:", err)
	}
	debug.Printf("json unmarshal: %v", j)

	var newP []Property
	found := false

	for i := range j.Properties {
		property := &j.Properties[i]
		debug.Printf("i: %v, Name: %#v, EncryptedValue: %#v\n", i, property.Name, property.EncryptedValue)
		if property.Name == *key {
			// https://code.google.com/p/go-wiki/wiki/SliceTricks
			newP = j.Properties[:i+copy(j.Properties[i:], j.Properties[i+1:])]
			found = true
			break
		}
	}

	if !found {
		return fmt.Errorf("\n\nError: Property '%s' not found.", *key)
	}

	debug.Printf("new properties: %v", newP)

	newData := Data{newP}

	bytes, err := json.MarshalIndent(newData, "", "    ")
	if err != nil {
		return fmt.Errorf("error:", err)
	}

	debug.Printf("b: %v", string(bytes))

	// Writing file
	// To handle large files, use a file buffer: http://stackoverflow.com/a/9739903/603745
	if err := ioutil.WriteFile(*jsonGPGDB, bytes, 0644); err != nil {
		return fmt.Errorf("error:", err)
	}

	return nil
}
<|endoftext|>


<commit_before>// This file was auto-generated via go generate.
// DO NOT UPDATE MANUALLY
package vlog_test

import "fmt"
import "testing"
import "os"

import "v.io/core/veyron/lib/modules"
import "v.io/core/veyron/lib/testutil"

func init() {
	modules.RegisterChild("child", ``, child)
}

func TestMain(m *testing.M) {
	testutil.Init()
	if modules.IsModulesProcess() {
		if err := modules.Dispatch(); err != nil {
			fmt.Fprintf(os.Stderr, "modules.Dispatch failed: %v\n", err)
			os.Exit(1)
		}
		return
	}
	os.Exit(m.Run())
}
<commit_msg>TBR: Copyright header added. MultiPart: 2/5<commit_after>// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// This file was auto-generated via go generate.
// DO NOT UPDATE MANUALLY
package vlog_test

import "fmt"
import "testing"
import "os"

import "v.io/core/veyron/lib/modules"
import "v.io/core/veyron/lib/testutil"

func init() {
	modules.RegisterChild("child", ``, child)
}

func TestMain(m *testing.M) {
	testutil.Init()
	if modules.IsModulesProcess() {
		if err := modules.Dispatch(); err != nil {
			fmt.Fprintf(os.Stderr, "modules.Dispatch failed: %v\n", err)
			os.Exit(1)
		}
		return
	}
	os.Exit(m.Run())
}
<|endoftext|>


<commit_before>/**
 * (C) Copyright 2013 Deft Labs
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at:
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package deftlabsutil

import (
	"io"
	"time"
	"os/exec"
)

// TimedCmdExec allows you to execute commands for a "max" period of time. After
// the specified amount of time, the command will be killed. If you
// want access to the stdout/err then you need to pass in a Writer to receive the data.
// This does not support sending data to the command executing. This method will panic if
// maxTimeMs is <= 0. If the cmdName length is zero, it will also panic. This returns true
// if the process was killed.
func CmdExecWithMaxTime(cmdName string, maxTimeMs int64, stdOutPipe io.Writer, stdErrPipe io.Writer, args ...string) (bool, error) {

	if maxTimeMs <= 0 {
		panic("You must set maxTimeMs to greater than zero")
	}

	if len(cmdName) == 0 {
		panic("You must set cmdName to a non-mepty string")
	}

	cmd := exec.Command(cmdName, args...)
	done := make(chan error)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return false, err
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return false, err
	}

	if err = cmd.Start(); err != nil {
		return false, err
	}

	if stdErrPipe != nil {
		go io.Copy(stdErrPipe, stderr)
	}

	if stdOutPipe != nil {
		go io.Copy(stdOutPipe, stdout)
	}

	go func() { done <- cmd.Wait() }()

	select {
		case <-time.After(time.Duration(maxTimeMs) * time.Millisecond): {
			err = cmd.Process.Kill()
			<-done
			return true, err
		}

		case err = <-done: return false, err
	}
}

<commit_msg>moved int64 to int<commit_after>/**
 * (C) Copyright 2013 Deft Labs
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at:
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package deftlabsutil

import (
	"io"
	"time"
	"os/exec"
)

// TimedCmdExec allows you to execute commands for a "max" period of time. After
// the specified amount of time, the command will be killed. If you
// want access to the stdout/err then you need to pass in a Writer to receive the data.
// This does not support sending data to the command executing. This method will panic if
// maxTimeMs is <= 0. If the cmdName length is zero, it will also panic. This returns true
// if the process was killed.
func CmdExecWithMaxTime(cmdName string, maxTimeMs int, stdOutPipe io.Writer, stdErrPipe io.Writer, args ...string) (bool, error) {

	if maxTimeMs <= 0 {
		panic("You must set maxTimeMs to greater than zero")
	}

	if len(cmdName) == 0 {
		panic("You must set cmdName to a non-mepty string")
	}

	cmd := exec.Command(cmdName, args...)
	done := make(chan error)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return false, err
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return false, err
	}

	if err = cmd.Start(); err != nil {
		return false, err
	}

	if stdErrPipe != nil {
		go io.Copy(stdErrPipe, stderr)
	}

	if stdOutPipe != nil {
		go io.Copy(stdOutPipe, stdout)
	}

	go func() { done <- cmd.Wait() }()

	select {
		case <-time.After(time.Duration(maxTimeMs) * time.Millisecond): {
			err = cmd.Process.Kill()
			<-done
			return true, err
		}

		case err = <-done: return false, err
	}
}

<|endoftext|>


<commit_before>package s3

import (
	"io"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/petergtz/bitsgo"
	"github.com/petergtz/bitsgo/blobstores/validate"
	"github.com/petergtz/bitsgo/config"
	"github.com/petergtz/bitsgo/logger"
	"github.com/pkg/errors"
)

type Blobstore struct {
	s3Client *s3.S3
	bucket   string
}

func NewBlobstore(config config.S3BlobstoreConfig) *Blobstore {
	validate.NotEmpty(config.AccessKeyID)
	validate.NotEmpty(config.Bucket)
	validate.NotEmpty(config.Region)
	validate.NotEmpty(config.SecretAccessKey)

	if config.Host != "" {
		validate.NotEmpty(config.Host)
	}

	return &Blobstore{
		s3Client: newS3Client(config.Region, config.AccessKeyID, config.SecretAccessKey, config.Host),
		bucket:   config.Bucket,
	}
}

func (blobstore *Blobstore) Exists(path string) (bool, error) {
	_, e := blobstore.s3Client.HeadObject(&s3.HeadObjectInput{
		Bucket: &blobstore.bucket,
		Key:    &path,
	})
	if e != nil {
		if isS3NotFoundError(e) {
			return false, nil
		}
		return false, errors.Wrapf(e, "Failed to check for %v/%v", blobstore.bucket, path)
	}
	return true, nil
}

func (blobstore *Blobstore) HeadOrRedirectAsGet(path string) (redirectLocation string, err error) {
	request, _ := blobstore.s3Client.GetObjectRequest(&s3.GetObjectInput{
		Bucket: &blobstore.bucket,
		Key:    &path,
	})
	return signedURLFrom(request, blobstore.bucket, path)
}

func (blobstore *Blobstore) Get(path string) (body io.ReadCloser, err error) {
	logger.Log.Debugw("Get from S3", "bucket", blobstore.bucket, "path", path)
	output, e := blobstore.s3Client.GetObject(&s3.GetObjectInput{
		Bucket: &blobstore.bucket,
		Key:    &path,
	})
	if e != nil {
		if isS3NotFoundError(e) {
			return nil, bitsgo.NewNotFoundError()
		}
		return nil, errors.Wrapf(e, "Path %v", path)
	}
	return output.Body, nil
}

func (blobstore *Blobstore) GetOrRedirect(path string) (body io.ReadCloser, redirectLocation string, err error) {
	request, _ := blobstore.s3Client.GetObjectRequest(&s3.GetObjectInput{
		Bucket: &blobstore.bucket,
		Key:    &path,
	})
	signedUrl, e := signedURLFrom(request, blobstore.bucket, path)
	return nil, signedUrl, e
}

func (blobstore *Blobstore) Put(path string, src io.ReadSeeker) error {
	logger.Log.Debugw("Put to S3", "bucket", blobstore.bucket, "path", path)
	_, e := blobstore.s3Client.PutObject(&s3.PutObjectInput{
		Bucket: &blobstore.bucket,
		Key:    &path,
		Body:   src,
	})
	if e != nil {
		return errors.Wrapf(e, "Path %v", path)
	}
	return nil
}

func signedURLFrom(req *request.Request, bucket, path string) (string, error) {
	signedURL, e := req.Presign(time.Hour)
	if e != nil {
		return "", errors.Wrapf(e, "Bucket/Path %v/%v", bucket, path)
	}
	return signedURL, nil

}

func (blobstore *Blobstore) Copy(src, dest string) error {
	logger.Log.Debugw("Copy in S3", "bucket", blobstore.bucket, "src", src, "dest", dest)
	_, e := blobstore.s3Client.CopyObject(&s3.CopyObjectInput{
		Key:        &dest,
		CopySource: aws.String(blobstore.bucket + "/" + src),
		Bucket:     &blobstore.bucket,
	})
	if e != nil {
		if isS3NotFoundError(e) {
			return bitsgo.NewNotFoundError()
		}
		return errors.Wrapf(e, "Error while trying to copy src %v to dest %v in bucket %v", src, dest, blobstore.bucket)
	}
	return nil
}

func (blobstore *Blobstore) Delete(path string) error {
	_, e := blobstore.s3Client.DeleteObject(&s3.DeleteObjectInput{
		Bucket: &blobstore.bucket,
		Key:    &path,
	})
	if e != nil {
		if isS3NotFoundError(e) {
			return bitsgo.NewNotFoundError()
		}
		return errors.Wrapf(e, "Path %v", path)
	}
	return nil
}

func (blobstore *Blobstore) DeleteDir(prefix string) error {
	deletionErrs := []error{}
	e := blobstore.s3Client.ListObjectsPages(
		&s3.ListObjectsInput{
			Bucket: &blobstore.bucket,
			Prefix: &prefix,
		},
		func(p *s3.ListObjectsOutput, lastPage bool) (shouldContinue bool) {
			for _, object := range p.Contents {
				e := blobstore.Delete(*object.Key)
				if e != nil {
					if _, isNotFoundError := e.(*bitsgo.NotFoundError); !isNotFoundError {
						deletionErrs = append(deletionErrs, e)
					}
				}
			}
			return true
		})
	if e != nil {
		return errors.Wrapf(e, "Prefix %v, errors from deleting: %v", prefix, deletionErrs)
	}
	if len(deletionErrs) != 0 {
		return errors.Errorf("Prefix %v, errors from deleting: %v", prefix, deletionErrs)
	}
	return nil
}

func (signer *Blobstore) Sign(resource string, method string, expirationTime time.Time) (signedURL string) {
	var request *request.Request
	switch strings.ToLower(method) {
	case "put":
		request, _ = signer.s3Client.PutObjectRequest(&s3.PutObjectInput{
			Bucket: aws.String(signer.bucket),
			Key:    aws.String(resource),
		})
	case "get":
		request, _ = signer.s3Client.GetObjectRequest(&s3.GetObjectInput{
			Bucket: aws.String(signer.bucket),
			Key:    aws.String(resource),
		})
	default:
		panic("The only supported methods are 'put' and 'get'. But got '" + method + "'")
	}
	// TODO use clock
	signedURL, e := request.Presign(expirationTime.Sub(time.Now()))
	if e != nil {
		panic(e)
	}
	logger.Log.Debugw("Signed URL", "verb", method, "signed-url", signedURL)
	return
}
<commit_msg>Remove check for optional host field<commit_after>package s3

import (
	"io"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/petergtz/bitsgo"
	"github.com/petergtz/bitsgo/blobstores/validate"
	"github.com/petergtz/bitsgo/config"
	"github.com/petergtz/bitsgo/logger"
	"github.com/pkg/errors"
)

type Blobstore struct {
	s3Client *s3.S3
	bucket   string
}

func NewBlobstore(config config.S3BlobstoreConfig) *Blobstore {
	validate.NotEmpty(config.AccessKeyID)
	validate.NotEmpty(config.Bucket)
	validate.NotEmpty(config.Region)
	validate.NotEmpty(config.SecretAccessKey)

	return &Blobstore{
		s3Client: newS3Client(config.Region, config.AccessKeyID, config.SecretAccessKey, config.Host),
		bucket:   config.Bucket,
	}
}

func (blobstore *Blobstore) Exists(path string) (bool, error) {
	_, e := blobstore.s3Client.HeadObject(&s3.HeadObjectInput{
		Bucket: &blobstore.bucket,
		Key:    &path,
	})
	if e != nil {
		if isS3NotFoundError(e) {
			return false, nil
		}
		return false, errors.Wrapf(e, "Failed to check for %v/%v", blobstore.bucket, path)
	}
	return true, nil
}

func (blobstore *Blobstore) HeadOrRedirectAsGet(path string) (redirectLocation string, err error) {
	request, _ := blobstore.s3Client.GetObjectRequest(&s3.GetObjectInput{
		Bucket: &blobstore.bucket,
		Key:    &path,
	})
	return signedURLFrom(request, blobstore.bucket, path)
}

func (blobstore *Blobstore) Get(path string) (body io.ReadCloser, err error) {
	logger.Log.Debugw("Get from S3", "bucket", blobstore.bucket, "path", path)
	output, e := blobstore.s3Client.GetObject(&s3.GetObjectInput{
		Bucket: &blobstore.bucket,
		Key:    &path,
	})
	if e != nil {
		if isS3NotFoundError(e) {
			return nil, bitsgo.NewNotFoundError()
		}
		return nil, errors.Wrapf(e, "Path %v", path)
	}
	return output.Body, nil
}

func (blobstore *Blobstore) GetOrRedirect(path string) (body io.ReadCloser, redirectLocation string, err error) {
	request, _ := blobstore.s3Client.GetObjectRequest(&s3.GetObjectInput{
		Bucket: &blobstore.bucket,
		Key:    &path,
	})
	signedUrl, e := signedURLFrom(request, blobstore.bucket, path)
	return nil, signedUrl, e
}

func (blobstore *Blobstore) Put(path string, src io.ReadSeeker) error {
	logger.Log.Debugw("Put to S3", "bucket", blobstore.bucket, "path", path)
	_, e := blobstore.s3Client.PutObject(&s3.PutObjectInput{
		Bucket: &blobstore.bucket,
		Key:    &path,
		Body:   src,
	})
	if e != nil {
		return errors.Wrapf(e, "Path %v", path)
	}
	return nil
}

func signedURLFrom(req *request.Request, bucket, path string) (string, error) {
	signedURL, e := req.Presign(time.Hour)
	if e != nil {
		return "", errors.Wrapf(e, "Bucket/Path %v/%v", bucket, path)
	}
	return signedURL, nil

}

func (blobstore *Blobstore) Copy(src, dest string) error {
	logger.Log.Debugw("Copy in S3", "bucket", blobstore.bucket, "src", src, "dest", dest)
	_, e := blobstore.s3Client.CopyObject(&s3.CopyObjectInput{
		Key:        &dest,
		CopySource: aws.String(blobstore.bucket + "/" + src),
		Bucket:     &blobstore.bucket,
	})
	if e != nil {
		if isS3NotFoundError(e) {
			return bitsgo.NewNotFoundError()
		}
		return errors.Wrapf(e, "Error while trying to copy src %v to dest %v in bucket %v", src, dest, blobstore.bucket)
	}
	return nil
}

func (blobstore *Blobstore) Delete(path string) error {
	_, e := blobstore.s3Client.DeleteObject(&s3.DeleteObjectInput{
		Bucket: &blobstore.bucket,
		Key:    &path,
	})
	if e != nil {
		if isS3NotFoundError(e) {
			return bitsgo.NewNotFoundError()
		}
		return errors.Wrapf(e, "Path %v", path)
	}
	return nil
}

func (blobstore *Blobstore) DeleteDir(prefix string) error {
	deletionErrs := []error{}
	e := blobstore.s3Client.ListObjectsPages(
		&s3.ListObjectsInput{
			Bucket: &blobstore.bucket,
			Prefix: &prefix,
		},
		func(p *s3.ListObjectsOutput, lastPage bool) (shouldContinue bool) {
			for _, object := range p.Contents {
				e := blobstore.Delete(*object.Key)
				if e != nil {
					if _, isNotFoundError := e.(*bitsgo.NotFoundError); !isNotFoundError {
						deletionErrs = append(deletionErrs, e)
					}
				}
			}
			return true
		})
	if e != nil {
		return errors.Wrapf(e, "Prefix %v, errors from deleting: %v", prefix, deletionErrs)
	}
	if len(deletionErrs) != 0 {
		return errors.Errorf("Prefix %v, errors from deleting: %v", prefix, deletionErrs)
	}
	return nil
}

func (signer *Blobstore) Sign(resource string, method string, expirationTime time.Time) (signedURL string) {
	var request *request.Request
	switch strings.ToLower(method) {
	case "put":
		request, _ = signer.s3Client.PutObjectRequest(&s3.PutObjectInput{
			Bucket: aws.String(signer.bucket),
			Key:    aws.String(resource),
		})
	case "get":
		request, _ = signer.s3Client.GetObjectRequest(&s3.GetObjectInput{
			Bucket: aws.String(signer.bucket),
			Key:    aws.String(resource),
		})
	default:
		panic("The only supported methods are 'put' and 'get'. But got '" + method + "'")
	}
	// TODO use clock
	signedURL, e := request.Presign(expirationTime.Sub(time.Now()))
	if e != nil {
		panic(e)
	}
	logger.Log.Debugw("Signed URL", "verb", method, "signed-url", signedURL)
	return
}
<|endoftext|>


<commit_before>package isolated

import (
	"io/ioutil"

	. "code.cloudfoundry.org/cli/integration/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = Describe("create-buildpack command", func() {
	Context("when the wrong data type is provided as the position argument", func() {
		It("outputs an error message to the user, provides help text, and exits 1", func() {
			// Avoid file does not exist error.
			err := ioutil.WriteFile("some-file", []byte{}, 0400)
			Expect(err).NotTo(HaveOccurred())

			session := CF("create-buildpack", "some-buildpack", "some-file", "not-an-integer")
			Eventually(session).Should(Exit(1))
			Expect(session.Err).To(Say("Incorrect usage: Value for POSITION must be integer"))
			Expect(session.Out).To(Say("cf create-buildpack BUILDPACK PATH POSITION")) // help
		})
	})
})
<commit_msg>create-buildpack integration test cleans after itself<commit_after>package isolated

import (
	"io/ioutil"
	"os"

	. "code.cloudfoundry.org/cli/integration/helpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"
)

var _ = FDescribe("create-buildpack command", func() {
	Context("when the wrong data type is provided as the position argument", func() {
		var filename string

		BeforeEach(func() {
			// The args take a filepath. Creating a real file will avoid the file
			// does not exist error, and trigger the correct error case we are
			// testing.
			filename = "some-file"
			err := ioutil.WriteFile(filename, []byte{}, 0400)
			Expect(err).NotTo(HaveOccurred())
		})

		AfterEach(func() {
			err := os.Remove(filename)
			Expect(err).NotTo(HaveOccurred())
		})

		It("outputs an error message to the user, provides help text, and exits 1", func() {
			session := CF("create-buildpack", "some-buildpack", "some-file", "not-an-integer")
			Eventually(session).Should(Exit(1))
			Expect(session.Err).To(Say("Incorrect usage: Value for POSITION must be integer"))
			Expect(session.Out).To(Say("cf create-buildpack BUILDPACK PATH POSITION")) // help
		})
	})
})
<|endoftext|>


<commit_before>package backend

import (
	"bufio"
	"bytes"
	"fmt"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/gogo/protobuf/proto"
	"github.com/lomik/graphite-clickhouse/carbonzipperpb"
	"github.com/uber-go/zap"
)

type Points []Point

func (s Points) Len() int      { return len(s) }
func (s Points) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

type ByKey struct{ Points }

func (s ByKey) Less(i, j int) bool {
	c := strings.Compare(s.Points[i].Metric, s.Points[j].Metric)

	switch c {
	case -1:
		return true
	case 1:
		return false
	case 0:
		return s.Points[i].Time < s.Points[j].Time
	}

	return false
}

type RenderHandler struct {
	config *Config
}

func (h *RenderHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	logger := Logger(r.Context())
	target := r.URL.Query().Get("target")

	if strings.IndexByte(target, '\'') > -1 { // sql injection dumb fix
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	var prefix string
	var err error

	if h.config.ClickHouse.ExtraPrefix != "" {
		prefix, target, err = RemoveExtraPrefix(h.config.ClickHouse.ExtraPrefix, target)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		if target == "" {
			h.Reply(w, r, make([]Point, 0), 0, 0, "")
			return
		}
	}

	fromTimestamp, err := strconv.ParseInt(r.URL.Query().Get("from"), 10, 32)
	if err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	untilTimestamp, err := strconv.ParseInt(r.URL.Query().Get("until"), 10, 32)
	if err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	dateWhere := fmt.Sprintf(
		"(Date >='%s' AND Date <= '%s' AND Time >= %d AND Time <= %d)",
		time.Unix(fromTimestamp, 0).Format("2006-01-02"),
		time.Unix(untilTimestamp, 0).Format("2006-01-02"),
		fromTimestamp,
		untilTimestamp,
	)

	var pathWhere string

	if hasWildcard(target) {
		// Search in small index table first
		treeWhere := makeWhere(target, true)
		if treeWhere == "" {
			http.Error(w, "Bad or unsupported query", http.StatusBadRequest)
			return
		}

		treeData, err := Query(
			r.Context(),
			h.config.ClickHouse.Url,
			fmt.Sprintf("SELECT Path FROM %s PREWHERE %s GROUP BY Path", h.config.ClickHouse.TreeTable, treeWhere),
			h.config.ClickHouse.TreeTimeout.Value(),
		)

		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		listBuf := bytes.NewBuffer(nil)
		first := true
		for _, p := range strings.Split(string(treeData), "\n") {
			if p == "" {
				continue
			}

			if !first {
				listBuf.Write([]byte{','})
			}
			first = false

			listBuf.WriteString("'" + p + "'") // SQL-Injection
		}

		if listBuf.Len() == 0 {
			h.Reply(w, r, make([]Point, 0), 0, 0, "")
			return
		}

		pathWhere = fmt.Sprintf(
			"Path IN (%s)",
			string(listBuf.Bytes()),
		)

		// pathWhere = fmt.Sprintf(
		// 	"Path IN (SELECT Path FROM %s WHERE %s)",
		// 	h.config.ClickHouse.DataTable,
		// )
		// pathWhere = makeWhere(target, false)
	} else {
		pathWhere = fmt.Sprintf("Path = '%s'", target)
	}

	// @TODO: change format to RowBinary
	query := fmt.Sprintf(
		`
		SELECT
			Path, Time, Value, Timestamp
		FROM %s
		PREWHERE (%s) AND (%s)
		FORMAT RowBinary
		`,
		h.config.ClickHouse.DataTable,
		pathWhere,
		dateWhere,
	)

	body, err := Query(
		r.Context(),
		h.config.ClickHouse.Url,
		query,
		h.config.ClickHouse.DataTimeout.Value(),
	)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	parseStart := time.Now()

	data, err := DataParse(body)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	logger.Debug("parse", zap.Duration("time_ns", time.Since(parseStart)))

	sortStart := time.Now()
	sort.Sort(data)
	logger.Debug("sort", zap.Duration("time_ns", time.Since(sortStart)))

	data.Points = PointsUniq(data.Points)

	// pp.Println(points)
	h.Reply(w, r, data.Points, int32(fromTimestamp), int32(untilTimestamp), prefix)
}

func (h *RenderHandler) Reply(w http.ResponseWriter, r *http.Request, points []Point, from, until int32, prefix string) {
	start := time.Now()
	switch r.URL.Query().Get("format") {
	case "pickle":
		h.ReplyPickle(w, r, points, from, until, prefix)
	case "protobuf":
		h.ReplyProtobuf(w, r, points, from, until, prefix)
	}
	Logger(r.Context()).Debug("reply", zap.Duration("time_ns", time.Since(start)))
}

func (h *RenderHandler) ReplyPickle(w http.ResponseWriter, r *http.Request, points []Point, from, until int32, prefix string) {
	var rollupTime time.Duration
	var pickleTime time.Duration

	defer func() {
		Logger(r.Context()).Debug("rollup", zap.Duration("time_ns", rollupTime))
		Logger(r.Context()).Debug("pickle", zap.Duration("time_ns", pickleTime))
	}()

	if len(points) == 0 {
		w.Write(PickleEmptyList)
		return
	}

	writer := bufio.NewWriterSize(w, 1024*1024)
	p := NewPickler(writer)
	defer writer.Flush()

	p.List()

	writeMetric := func(points []Point) {
		rollupStart := time.Now()
		points, step := h.config.Rollup.RollupMetric(points)
		rollupTime += time.Since(rollupStart)

		pickleStart := time.Now()
		p.Dict()

		p.String("name")
		if prefix != "" {
			p.String(prefix + "." + points[0].Metric)
		} else {
			p.String(points[0].Metric)
		}
		p.SetItem()

		p.String("step")
		p.Uint32(uint32(step))
		p.SetItem()

		start := from - (from % step)
		if start < from {
			start += step
		}
		end := until - (until % step)
		last := start - step

		p.String("values")
		p.List()
		for _, point := range points {
			if point.Time < start || point.Time > end {
				continue
			}

			if point.Time > last+step {
				p.AppendNulls(int(((point.Time - last) / step) - 1))
			}

			p.AppendFloat64(point.Value)

			last = point.Time
		}

		if end > last {
			p.AppendNulls(int((end - last) / step))
		}
		p.SetItem()

		p.String("start")
		p.Uint32(uint32(start))
		p.SetItem()

		p.String("end")
		p.Uint32(uint32(end))
		p.SetItem()

		p.Append()
		pickleTime += time.Since(pickleStart)
	}

	// group by Metric
	var i, n int
	// i - current position of iterator
	// n - position of the first record with current metric
	l := len(points)

	for i = 1; i < l; i++ {
		if points[i].Metric != points[n].Metric {
			writeMetric(points[n:i])
			n = i
			continue
		}
	}
	writeMetric(points[n:i])

	p.Stop()
}

func (h *RenderHandler) ReplyProtobuf(w http.ResponseWriter, r *http.Request, points []Point, from, until int32, prefix string) {
	if len(points) == 0 {
		return
	}

	var multiResponse carbonzipperpb.MultiFetchResponse

	writeMetric := func(points []Point) {
		points, step := h.config.Rollup.RollupMetric(points)

		var name string

		if prefix != "" {
			name = prefix + "." + points[0].Metric
		} else {
			name = points[0].Metric
		}

		start := from - (from % step)
		if start < from {
			start += step
		}
		stop := until - (until % step)
		count := ((stop - start) / step) + 1

		response := carbonzipperpb.FetchResponse{
			Name:      proto.String(name),
			StartTime: &start,
			StopTime:  &stop,
			StepTime:  &step,
			Values:    make([]float64, count),
			IsAbsent:  make([]bool, count),
		}

		var index int32
		// skip points before start
		for index = 0; points[index].Time < start; index++ {
		}

		for i := int32(0); i < count; i++ {
			if index < int32(len(points)) && points[index].Time == start+step*i {
				response.Values[i] = points[index].Value
				response.IsAbsent[i] = false
				index++
			} else {
				response.Values[i] = 0
				response.IsAbsent[i] = true
			}
		}

		multiResponse.Metrics = append(multiResponse.Metrics, &response)
	}

	// group by Metric
	var i, n int
	// i - current position of iterator
	// n - position of the first record with current metric
	l := len(points)

	for i = 1; i < l; i++ {
		if points[i].Metric != points[n].Metric {
			writeMetric(points[n:i])
			n = i
			continue
		}
	}
	writeMetric(points[n:i])

	body, _ := proto.Marshal(&multiResponse)
	w.Write(body)
}

func NewRenderHandler(config *Config) *RenderHandler {
	return &RenderHandler{
		config: config,
	}
}
<commit_msg>remove unused sort helper<commit_after>package backend

import (
	"bufio"
	"bytes"
	"fmt"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/gogo/protobuf/proto"
	"github.com/lomik/graphite-clickhouse/carbonzipperpb"
	"github.com/uber-go/zap"
)

type RenderHandler struct {
	config *Config
}

func (h *RenderHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	logger := Logger(r.Context())
	target := r.URL.Query().Get("target")

	if strings.IndexByte(target, '\'') > -1 { // sql injection dumb fix
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	var prefix string
	var err error

	if h.config.ClickHouse.ExtraPrefix != "" {
		prefix, target, err = RemoveExtraPrefix(h.config.ClickHouse.ExtraPrefix, target)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		if target == "" {
			h.Reply(w, r, make([]Point, 0), 0, 0, "")
			return
		}
	}

	fromTimestamp, err := strconv.ParseInt(r.URL.Query().Get("from"), 10, 32)
	if err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	untilTimestamp, err := strconv.ParseInt(r.URL.Query().Get("until"), 10, 32)
	if err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	dateWhere := fmt.Sprintf(
		"(Date >='%s' AND Date <= '%s' AND Time >= %d AND Time <= %d)",
		time.Unix(fromTimestamp, 0).Format("2006-01-02"),
		time.Unix(untilTimestamp, 0).Format("2006-01-02"),
		fromTimestamp,
		untilTimestamp,
	)

	var pathWhere string

	if hasWildcard(target) {
		// Search in small index table first
		treeWhere := makeWhere(target, true)
		if treeWhere == "" {
			http.Error(w, "Bad or unsupported query", http.StatusBadRequest)
			return
		}

		treeData, err := Query(
			r.Context(),
			h.config.ClickHouse.Url,
			fmt.Sprintf("SELECT Path FROM %s PREWHERE %s GROUP BY Path", h.config.ClickHouse.TreeTable, treeWhere),
			h.config.ClickHouse.TreeTimeout.Value(),
		)

		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		listBuf := bytes.NewBuffer(nil)
		first := true
		for _, p := range strings.Split(string(treeData), "\n") {
			if p == "" {
				continue
			}

			if !first {
				listBuf.Write([]byte{','})
			}
			first = false

			listBuf.WriteString("'" + p + "'") // SQL-Injection
		}

		if listBuf.Len() == 0 {
			h.Reply(w, r, make([]Point, 0), 0, 0, "")
			return
		}

		pathWhere = fmt.Sprintf(
			"Path IN (%s)",
			string(listBuf.Bytes()),
		)

		// pathWhere = fmt.Sprintf(
		// 	"Path IN (SELECT Path FROM %s WHERE %s)",
		// 	h.config.ClickHouse.DataTable,
		// )
		// pathWhere = makeWhere(target, false)
	} else {
		pathWhere = fmt.Sprintf("Path = '%s'", target)
	}

	// @TODO: change format to RowBinary
	query := fmt.Sprintf(
		`
		SELECT
			Path, Time, Value, Timestamp
		FROM %s
		PREWHERE (%s) AND (%s)
		FORMAT RowBinary
		`,
		h.config.ClickHouse.DataTable,
		pathWhere,
		dateWhere,
	)

	body, err := Query(
		r.Context(),
		h.config.ClickHouse.Url,
		query,
		h.config.ClickHouse.DataTimeout.Value(),
	)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	parseStart := time.Now()

	data, err := DataParse(body)

	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	logger.Debug("parse", zap.Duration("time_ns", time.Since(parseStart)))

	sortStart := time.Now()
	sort.Sort(data)
	logger.Debug("sort", zap.Duration("time_ns", time.Since(sortStart)))

	data.Points = PointsUniq(data.Points)

	// pp.Println(points)
	h.Reply(w, r, data.Points, int32(fromTimestamp), int32(untilTimestamp), prefix)
}

func (h *RenderHandler) Reply(w http.ResponseWriter, r *http.Request, points []Point, from, until int32, prefix string) {
	start := time.Now()
	switch r.URL.Query().Get("format") {
	case "pickle":
		h.ReplyPickle(w, r, points, from, until, prefix)
	case "protobuf":
		h.ReplyProtobuf(w, r, points, from, until, prefix)
	}
	Logger(r.Context()).Debug("reply", zap.Duration("time_ns", time.Since(start)))
}

func (h *RenderHandler) ReplyPickle(w http.ResponseWriter, r *http.Request, points []Point, from, until int32, prefix string) {
	var rollupTime time.Duration
	var pickleTime time.Duration

	defer func() {
		Logger(r.Context()).Debug("rollup", zap.Duration("time_ns", rollupTime))
		Logger(r.Context()).Debug("pickle", zap.Duration("time_ns", pickleTime))
	}()

	if len(points) == 0 {
		w.Write(PickleEmptyList)
		return
	}

	writer := bufio.NewWriterSize(w, 1024*1024)
	p := NewPickler(writer)
	defer writer.Flush()

	p.List()

	writeMetric := func(points []Point) {
		rollupStart := time.Now()
		points, step := h.config.Rollup.RollupMetric(points)
		rollupTime += time.Since(rollupStart)

		pickleStart := time.Now()
		p.Dict()

		p.String("name")
		if prefix != "" {
			p.String(prefix + "." + points[0].Metric)
		} else {
			p.String(points[0].Metric)
		}
		p.SetItem()

		p.String("step")
		p.Uint32(uint32(step))
		p.SetItem()

		start := from - (from % step)
		if start < from {
			start += step
		}
		end := until - (until % step)
		last := start - step

		p.String("values")
		p.List()
		for _, point := range points {
			if point.Time < start || point.Time > end {
				continue
			}

			if point.Time > last+step {
				p.AppendNulls(int(((point.Time - last) / step) - 1))
			}

			p.AppendFloat64(point.Value)

			last = point.Time
		}

		if end > last {
			p.AppendNulls(int((end - last) / step))
		}
		p.SetItem()

		p.String("start")
		p.Uint32(uint32(start))
		p.SetItem()

		p.String("end")
		p.Uint32(uint32(end))
		p.SetItem()

		p.Append()
		pickleTime += time.Since(pickleStart)
	}

	// group by Metric
	var i, n int
	// i - current position of iterator
	// n - position of the first record with current metric
	l := len(points)

	for i = 1; i < l; i++ {
		if points[i].Metric != points[n].Metric {
			writeMetric(points[n:i])
			n = i
			continue
		}
	}
	writeMetric(points[n:i])

	p.Stop()
}

func (h *RenderHandler) ReplyProtobuf(w http.ResponseWriter, r *http.Request, points []Point, from, until int32, prefix string) {
	if len(points) == 0 {
		return
	}

	var multiResponse carbonzipperpb.MultiFetchResponse

	writeMetric := func(points []Point) {
		points, step := h.config.Rollup.RollupMetric(points)

		var name string

		if prefix != "" {
			name = prefix + "." + points[0].Metric
		} else {
			name = points[0].Metric
		}

		start := from - (from % step)
		if start < from {
			start += step
		}
		stop := until - (until % step)
		count := ((stop - start) / step) + 1

		response := carbonzipperpb.FetchResponse{
			Name:      proto.String(name),
			StartTime: &start,
			StopTime:  &stop,
			StepTime:  &step,
			Values:    make([]float64, count),
			IsAbsent:  make([]bool, count),
		}

		var index int32
		// skip points before start
		for index = 0; points[index].Time < start; index++ {
		}

		for i := int32(0); i < count; i++ {
			if index < int32(len(points)) && points[index].Time == start+step*i {
				response.Values[i] = points[index].Value
				response.IsAbsent[i] = false
				index++
			} else {
				response.Values[i] = 0
				response.IsAbsent[i] = true
			}
		}

		multiResponse.Metrics = append(multiResponse.Metrics, &response)
	}

	// group by Metric
	var i, n int
	// i - current position of iterator
	// n - position of the first record with current metric
	l := len(points)

	for i = 1; i < l; i++ {
		if points[i].Metric != points[n].Metric {
			writeMetric(points[n:i])
			n = i
			continue
		}
	}
	writeMetric(points[n:i])

	body, _ := proto.Marshal(&multiResponse)
	w.Write(body)
}

func NewRenderHandler(config *Config) *RenderHandler {
	return &RenderHandler{
		config: config,
	}
}
<|endoftext|>