<commit_before>/*
Copyright 2016 The Kubernetes 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 templates

import (
	"bytes"
	"fmt"
	"strings"
	"text/template"
	"unicode"

	"k8s.io/kubernetes/pkg/kubectl/util/term"

	"github.com/spf13/cobra"
	flag "github.com/spf13/pflag"
)

type FlagExposer interface {
	ExposeFlags(cmd *cobra.Command, flags ...string) FlagExposer
}

func ActsAsRootCommand(cmd *cobra.Command, filters []string, groups ...CommandGroup) FlagExposer {
	if cmd == nil {
		panic("nil root command")
	}
	templater := &templater{
		RootCmd:       cmd,
		UsageTemplate: MainUsageTemplate(),
		HelpTemplate:  MainHelpTemplate(),
		CommandGroups: groups,
		Filtered:      filters,
	}
	cmd.SetUsageFunc(templater.UsageFunc())
	cmd.SetHelpFunc(templater.HelpFunc())
	return templater
}

func UseOptionsTemplates(cmd *cobra.Command) {
	templater := &templater{
		UsageTemplate: OptionsUsageTemplate(),
		HelpTemplate:  OptionsHelpTemplate(),
	}
	cmd.SetUsageFunc(templater.UsageFunc())
	cmd.SetHelpFunc(templater.HelpFunc())
}

type templater struct {
	UsageTemplate string
	HelpTemplate  string
	RootCmd       *cobra.Command
	CommandGroups
	Filtered []string
}

func (templater *templater) ExposeFlags(cmd *cobra.Command, flags ...string) FlagExposer {
	cmd.SetUsageFunc(templater.UsageFunc(flags...))
	return templater
}

func (templater *templater) HelpFunc() func(*cobra.Command, []string) {
	return func(c *cobra.Command, s []string) {
		t := template.New("help")
		t.Funcs(templater.templateFuncs())
		template.Must(t.Parse(templater.HelpTemplate))
		out := term.NewResponsiveWriter(c.OutOrStdout())
		err := t.Execute(out, c)
		if err != nil {
			c.Println(err)
		}
	}
}

func (templater *templater) UsageFunc(exposedFlags ...string) func(*cobra.Command) error {
	return func(c *cobra.Command) error {
		t := template.New("usage")
		t.Funcs(templater.templateFuncs(exposedFlags...))
		template.Must(t.Parse(templater.UsageTemplate))
		out := term.NewResponsiveWriter(c.OutOrStderr())
		return t.Execute(out, c)
	}
}

func (templater *templater) templateFuncs(exposedFlags ...string) template.FuncMap {
	return template.FuncMap{
		"trim":                strings.TrimSpace,
		"trimRight":           func(s string) string { return strings.TrimRightFunc(s, unicode.IsSpace) },
		"trimLeft":            func(s string) string { return strings.TrimLeftFunc(s, unicode.IsSpace) },
		"gt":                  cobra.Gt,
		"eq":                  cobra.Eq,
		"rpad":                rpad,
		"appendIfNotPresent":  appendIfNotPresent,
		"flagsNotIntersected": flagsNotIntersected,
		"visibleFlags":        visibleFlags,
		"flagsUsages":         flagsUsages,
		"cmdGroups":           templater.cmdGroups,
		"cmdGroupsString":     templater.cmdGroupsString,
		"rootCmd":             templater.rootCmdName,
		"isRootCmd":           templater.isRootCmd,
		"optionsCmdFor":       templater.optionsCmdFor,
		"usageLine":           templater.usageLine,
		"exposed": func(c *cobra.Command) *flag.FlagSet {
			exposed := flag.NewFlagSet("exposed", flag.ContinueOnError)
			if len(exposedFlags) > 0 {
				for _, name := range exposedFlags {
					if flag := c.Flags().Lookup(name); flag != nil {
						exposed.AddFlag(flag)
					}
				}
			}
			return exposed
		},
	}
}

func (templater *templater) cmdGroups(c *cobra.Command, all []*cobra.Command) []CommandGroup {
	if len(templater.CommandGroups) > 0 && c == templater.RootCmd {
		all = filter(all, templater.Filtered...)
		return AddAdditionalCommands(templater.CommandGroups, "Other Commands:", all)
	}
	all = filter(all, "options")
	return []CommandGroup{
		{
			Message:  "Available Commands:",
			Commands: all,
		},
	}
}

func (t *templater) cmdGroupsString(c *cobra.Command) string {
	groups := []string{}
	for _, cmdGroup := range t.cmdGroups(c, c.Commands()) {
		cmds := []string{cmdGroup.Message}
		for _, cmd := range cmdGroup.Commands {
			if cmd.Runnable() {
				cmds = append(cmds, "  "+rpad(cmd.Name(), cmd.NamePadding())+" "+cmd.Short)
			}
		}
		groups = append(groups, strings.Join(cmds, "\n"))
	}
	return strings.Join(groups, "\n\n")
}

func (t *templater) rootCmdName(c *cobra.Command) string {
	return t.rootCmd(c).CommandPath()
}

func (t *templater) isRootCmd(c *cobra.Command) bool {
	return t.rootCmd(c) == c
}

func (t *templater) parents(c *cobra.Command) []*cobra.Command {
	parents := []*cobra.Command{c}
	for current := c; !t.isRootCmd(current) && current.HasParent(); {
		current = current.Parent()
		parents = append(parents, current)
	}
	return parents
}

func (t *templater) rootCmd(c *cobra.Command) *cobra.Command {
	if c != nil && !c.HasParent() {
		return c
	}
	if t.RootCmd == nil {
		panic("nil root cmd")
	}
	return t.RootCmd
}

func (t *templater) optionsCmdFor(c *cobra.Command) string {
	if !c.Runnable() {
		return ""
	}
	rootCmdStructure := t.parents(c)
	for i := len(rootCmdStructure) - 1; i >= 0; i-- {
		cmd := rootCmdStructure[i]
		if _, _, err := cmd.Find([]string{"options"}); err == nil {
			return cmd.CommandPath() + " options"
		}
	}
	return ""
}

func (t *templater) usageLine(c *cobra.Command) string {
	usage := c.UseLine()
	suffix := "[options]"
	if c.HasFlags() && !strings.Contains(usage, suffix) {
		usage += " " + suffix
	}
	return usage
}

func flagsUsages(f *flag.FlagSet) string {
	x := new(bytes.Buffer)

	f.VisitAll(func(flag *flag.Flag) {
		if flag.Hidden {
			return
		}
		format := "--%s=%s: %s\n"

		if flag.Value.Type() == "string" {
			format = "--%s='%s': %s\n"
		}

		if len(flag.Shorthand) > 0 {
			format = "  -%s, " + format
		} else {
			format = "   %s   " + format
		}

		fmt.Fprintf(x, format, flag.Shorthand, flag.Name, flag.DefValue, flag.Usage)
	})

	return x.String()
}

func rpad(s string, padding int) string {
	template := fmt.Sprintf("%%-%ds", padding)
	return fmt.Sprintf(template, s)
}

func indentLines(s string, indentation int) string {
	r := []string{}
	for _, line := range strings.Split(s, "\n") {
		indented := strings.Repeat(" ", indentation) + line
		r = append(r, indented)
	}
	return strings.Join(r, "\n")
}

func appendIfNotPresent(s, stringToAppend string) string {
	if strings.Contains(s, stringToAppend) {
		return s
	}
	return s + " " + stringToAppend
}

func flagsNotIntersected(l *flag.FlagSet, r *flag.FlagSet) *flag.FlagSet {
	f := flag.NewFlagSet("notIntersected", flag.ContinueOnError)
	l.VisitAll(func(flag *flag.Flag) {
		if r.Lookup(flag.Name) == nil {
			f.AddFlag(flag)
		}
	})
	return f
}

func visibleFlags(l *flag.FlagSet) *flag.FlagSet {
	hidden := "help"
	f := flag.NewFlagSet("visible", flag.ContinueOnError)
	l.VisitAll(func(flag *flag.Flag) {
		if flag.Name != hidden {
			f.AddFlag(flag)
		}
	})
	return f
}

func filter(cmds []*cobra.Command, names ...string) []*cobra.Command {
	out := []*cobra.Command{}
	for _, c := range cmds {
		if c.Hidden {
			continue
		}
		skip := false
		for _, name := range names {
			if name == c.Name() {
				skip = true
				break
			}
		}
		if skip {
			continue
		}
		out = append(out, c)
	}
	return out
}
<commit_msg>Filter unavailable commands in help<commit_after>/*
Copyright 2016 The Kubernetes 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 templates

import (
	"bytes"
	"fmt"
	"strings"
	"text/template"
	"unicode"

	"k8s.io/kubernetes/pkg/kubectl/util/term"

	"github.com/spf13/cobra"
	flag "github.com/spf13/pflag"
)

type FlagExposer interface {
	ExposeFlags(cmd *cobra.Command, flags ...string) FlagExposer
}

func ActsAsRootCommand(cmd *cobra.Command, filters []string, groups ...CommandGroup) FlagExposer {
	if cmd == nil {
		panic("nil root command")
	}
	templater := &templater{
		RootCmd:       cmd,
		UsageTemplate: MainUsageTemplate(),
		HelpTemplate:  MainHelpTemplate(),
		CommandGroups: groups,
		Filtered:      filters,
	}
	cmd.SetUsageFunc(templater.UsageFunc())
	cmd.SetHelpFunc(templater.HelpFunc())
	return templater
}

func UseOptionsTemplates(cmd *cobra.Command) {
	templater := &templater{
		UsageTemplate: OptionsUsageTemplate(),
		HelpTemplate:  OptionsHelpTemplate(),
	}
	cmd.SetUsageFunc(templater.UsageFunc())
	cmd.SetHelpFunc(templater.HelpFunc())
}

type templater struct {
	UsageTemplate string
	HelpTemplate  string
	RootCmd       *cobra.Command
	CommandGroups
	Filtered []string
}

func (templater *templater) ExposeFlags(cmd *cobra.Command, flags ...string) FlagExposer {
	cmd.SetUsageFunc(templater.UsageFunc(flags...))
	return templater
}

func (templater *templater) HelpFunc() func(*cobra.Command, []string) {
	return func(c *cobra.Command, s []string) {
		t := template.New("help")
		t.Funcs(templater.templateFuncs())
		template.Must(t.Parse(templater.HelpTemplate))
		out := term.NewResponsiveWriter(c.OutOrStdout())
		err := t.Execute(out, c)
		if err != nil {
			c.Println(err)
		}
	}
}

func (templater *templater) UsageFunc(exposedFlags ...string) func(*cobra.Command) error {
	return func(c *cobra.Command) error {
		t := template.New("usage")
		t.Funcs(templater.templateFuncs(exposedFlags...))
		template.Must(t.Parse(templater.UsageTemplate))
		out := term.NewResponsiveWriter(c.OutOrStderr())
		return t.Execute(out, c)
	}
}

func (templater *templater) templateFuncs(exposedFlags ...string) template.FuncMap {
	return template.FuncMap{
		"trim":                strings.TrimSpace,
		"trimRight":           func(s string) string { return strings.TrimRightFunc(s, unicode.IsSpace) },
		"trimLeft":            func(s string) string { return strings.TrimLeftFunc(s, unicode.IsSpace) },
		"gt":                  cobra.Gt,
		"eq":                  cobra.Eq,
		"rpad":                rpad,
		"appendIfNotPresent":  appendIfNotPresent,
		"flagsNotIntersected": flagsNotIntersected,
		"visibleFlags":        visibleFlags,
		"flagsUsages":         flagsUsages,
		"cmdGroups":           templater.cmdGroups,
		"cmdGroupsString":     templater.cmdGroupsString,
		"rootCmd":             templater.rootCmdName,
		"isRootCmd":           templater.isRootCmd,
		"optionsCmdFor":       templater.optionsCmdFor,
		"usageLine":           templater.usageLine,
		"exposed": func(c *cobra.Command) *flag.FlagSet {
			exposed := flag.NewFlagSet("exposed", flag.ContinueOnError)
			if len(exposedFlags) > 0 {
				for _, name := range exposedFlags {
					if flag := c.Flags().Lookup(name); flag != nil {
						exposed.AddFlag(flag)
					}
				}
			}
			return exposed
		},
	}
}

func (templater *templater) cmdGroups(c *cobra.Command, all []*cobra.Command) []CommandGroup {
	if len(templater.CommandGroups) > 0 && c == templater.RootCmd {
		all = filter(all, templater.Filtered...)
		return AddAdditionalCommands(templater.CommandGroups, "Other Commands:", all)
	}
	all = filter(all, "options")
	return []CommandGroup{
		{
			Message:  "Available Commands:",
			Commands: all,
		},
	}
}

func (t *templater) cmdGroupsString(c *cobra.Command) string {
	groups := []string{}
	for _, cmdGroup := range t.cmdGroups(c, c.Commands()) {
		cmds := []string{cmdGroup.Message}
		for _, cmd := range cmdGroup.Commands {
			if cmd.IsAvailableCommand() {
				cmds = append(cmds, "  "+rpad(cmd.Name(), cmd.NamePadding())+" "+cmd.Short)
			}
		}
		groups = append(groups, strings.Join(cmds, "\n"))
	}
	return strings.Join(groups, "\n\n")
}

func (t *templater) rootCmdName(c *cobra.Command) string {
	return t.rootCmd(c).CommandPath()
}

func (t *templater) isRootCmd(c *cobra.Command) bool {
	return t.rootCmd(c) == c
}

func (t *templater) parents(c *cobra.Command) []*cobra.Command {
	parents := []*cobra.Command{c}
	for current := c; !t.isRootCmd(current) && current.HasParent(); {
		current = current.Parent()
		parents = append(parents, current)
	}
	return parents
}

func (t *templater) rootCmd(c *cobra.Command) *cobra.Command {
	if c != nil && !c.HasParent() {
		return c
	}
	if t.RootCmd == nil {
		panic("nil root cmd")
	}
	return t.RootCmd
}

func (t *templater) optionsCmdFor(c *cobra.Command) string {
	if !c.Runnable() {
		return ""
	}
	rootCmdStructure := t.parents(c)
	for i := len(rootCmdStructure) - 1; i >= 0; i-- {
		cmd := rootCmdStructure[i]
		if _, _, err := cmd.Find([]string{"options"}); err == nil {
			return cmd.CommandPath() + " options"
		}
	}
	return ""
}

func (t *templater) usageLine(c *cobra.Command) string {
	usage := c.UseLine()
	suffix := "[options]"
	if c.HasFlags() && !strings.Contains(usage, suffix) {
		usage += " " + suffix
	}
	return usage
}

func flagsUsages(f *flag.FlagSet) string {
	x := new(bytes.Buffer)

	f.VisitAll(func(flag *flag.Flag) {
		if flag.Hidden {
			return
		}
		format := "--%s=%s: %s\n"

		if flag.Value.Type() == "string" {
			format = "--%s='%s': %s\n"
		}

		if len(flag.Shorthand) > 0 {
			format = "  -%s, " + format
		} else {
			format = "   %s   " + format
		}

		fmt.Fprintf(x, format, flag.Shorthand, flag.Name, flag.DefValue, flag.Usage)
	})

	return x.String()
}

func rpad(s string, padding int) string {
	template := fmt.Sprintf("%%-%ds", padding)
	return fmt.Sprintf(template, s)
}

func indentLines(s string, indentation int) string {
	r := []string{}
	for _, line := range strings.Split(s, "\n") {
		indented := strings.Repeat(" ", indentation) + line
		r = append(r, indented)
	}
	return strings.Join(r, "\n")
}

func appendIfNotPresent(s, stringToAppend string) string {
	if strings.Contains(s, stringToAppend) {
		return s
	}
	return s + " " + stringToAppend
}

func flagsNotIntersected(l *flag.FlagSet, r *flag.FlagSet) *flag.FlagSet {
	f := flag.NewFlagSet("notIntersected", flag.ContinueOnError)
	l.VisitAll(func(flag *flag.Flag) {
		if r.Lookup(flag.Name) == nil {
			f.AddFlag(flag)
		}
	})
	return f
}

func visibleFlags(l *flag.FlagSet) *flag.FlagSet {
	hidden := "help"
	f := flag.NewFlagSet("visible", flag.ContinueOnError)
	l.VisitAll(func(flag *flag.Flag) {
		if flag.Name != hidden {
			f.AddFlag(flag)
		}
	})
	return f
}

func filter(cmds []*cobra.Command, names ...string) []*cobra.Command {
	out := []*cobra.Command{}
	for _, c := range cmds {
		if c.Hidden {
			continue
		}
		skip := false
		for _, name := range names {
			if name == c.Name() {
				skip = true
				break
			}
		}
		if skip {
			continue
		}
		out = append(out, c)
	}
	return out
}
<|endoftext|>


<commit_before>package expression

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

const (
	ResTrue         = iota
	ResFalse        = iota
	ResCompileError = iota
	ResEvalError    = iota
)

func evaluate(t *testing.T, expressionStr string, params *Parameters, expectedResult int) {
	// Check for compilation
	expr, err := NewExpression(expressionStr)
	if !assert.Equal(t, expectedResult != ResCompileError, err == nil, "Expression compilation (success vs. error): "+expressionStr) || expectedResult == ResCompileError {
		return
	}

	// Check for evaluation
	result, err := expr.EvaluateAsBool(params)
	if !assert.Equal(t, expectedResult != ResEvalError, err == nil, "Expression evaluation (success vs. error): "+expressionStr) || expectedResult == ResEvalError {
		return
	}

	// Check for result
	assert.Equal(t, expectedResult == ResTrue, result, "Expression evaluation result: "+expressionStr)
}

func evaluateWithCache(t *testing.T, expressionStr string, params *Parameters, expectedResult int, cache *Cache) {
	// Check for compilation & evaluation
	for i := 0; i < 10; i++ {
		result, err := cache.EvaluateAsBool(expressionStr, params)
		if !assert.Equal(t, expectedResult != ResCompileError && expectedResult != ResEvalError, err == nil, "[Cache] Expression compilation && evaluation (success vs. error): "+expressionStr) || expectedResult == ResCompileError || expectedResult == ResEvalError {
			return
		}

		// Check for result
		assert.Equal(t, expectedResult == ResTrue, result, "[Cache] Expression evaluation result: "+expressionStr)
	}
}

func TestExpressions(t *testing.T) {
	params := NewParams(
		map[string]string{
			"foo":         "10",
			"unusedLabel": "3",
			"a":           "valueOfA",
			"bar":         "true",
			"anotherbar":  "t",
		},

		map[string]interface{}{
			"service": struct {
				Name   string
				Labels map[string]string
			}{
				"myservicename",
				map[string]string{
					"Name": "Value",
				},
			},
		},
	)

	tests := []struct {
		expression string
		result     int
	}{
		// true (checking bool, int, string)
		{"anotherbar == true", ResTrue},
		{"anotherbar", ResTrue},
		{"bar == true", ResTrue},
		{"bar", ResTrue},
		{"foo > 5", ResTrue},
		{"a == 'valueOfA'", ResTrue},
		{"foo > 5 && a == 'valueOfA'", ResTrue},

		// true (IN function)
		{"in(a, 'valueOfC', 'valueOfB', 'valueOfA')", ResTrue},
		{"in(foo, 10, 20, 30)", ResTrue},
		{"in(a, 'valueOfX', 'valueOfY', 'valueOfZ')", ResFalse},

		// false
		{"anotherbar == 'p'", ResFalse},
		{"'A' + 'B' == 5", ResFalse},

		// check when expression involves a missing label
		{"foo > 5 && missingLabel == 'requiredValue'", ResFalse},
		{"foo > 5 && missingLabel == 239", ResFalse},

		// we are explicitly converting all integer-like params to integers, so this should always be false (expected behavior)
		{"foo == '10'", ResFalse},

		// we are explicitly converting all bool-like params to bool, so this should always be false (expected behavior)
		{"anotherbar == 't'", ResFalse},

		// check that struct expressions work
		{"service.Name == 'myservicename'", ResTrue},
		{"service.Name == 'incorrectservicename'", ResFalse},
		{"service.Labels.Name + 'B' == 'ValueB'", ResTrue},
		{"serviceMissing.LabelsMissing.Name + 'B' == 'ValueB'", ResFalse},

		// evaluation error
		{"foo + 10 + 'test' > 0", ResEvalError},

		// not a boolean
		{"'a' + 'b' + bar", ResEvalError},

		// compilation error
		{"(5 + 10 > 9", ResCompileError},
	}

	// Evaluate without cache
	for _, test := range tests {
		evaluate(t, test.expression, params, test.result)
	}

	cache := NewCache()
	for _, test := range tests {
		evaluateWithCache(t, test.expression, params, test.result, cache)
	}
}
<commit_msg>added extra unit tests for IN() function in expression engine<commit_after>package expression

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

const (
	ResTrue         = iota
	ResFalse        = iota
	ResCompileError = iota
	ResEvalError    = iota
)

func evaluate(t *testing.T, expressionStr string, params *Parameters, expectedResult int) {
	// Check for compilation
	expr, err := NewExpression(expressionStr)
	if !assert.Equal(t, expectedResult != ResCompileError, err == nil, "Expression compilation (success vs. error): "+expressionStr) || expectedResult == ResCompileError {
		return
	}

	// Check for evaluation
	result, err := expr.EvaluateAsBool(params)
	if !assert.Equal(t, expectedResult != ResEvalError, err == nil, "Expression evaluation (success vs. error): "+expressionStr) || expectedResult == ResEvalError {
		return
	}

	// Check for result
	assert.Equal(t, expectedResult == ResTrue, result, "Expression evaluation result: "+expressionStr)
}

func evaluateWithCache(t *testing.T, expressionStr string, params *Parameters, expectedResult int, cache *Cache) {
	// Check for compilation & evaluation
	for i := 0; i < 10; i++ {
		result, err := cache.EvaluateAsBool(expressionStr, params)
		if !assert.Equal(t, expectedResult != ResCompileError && expectedResult != ResEvalError, err == nil, "[Cache] Expression compilation && evaluation (success vs. error): "+expressionStr) || expectedResult == ResCompileError || expectedResult == ResEvalError {
			return
		}

		// Check for result
		assert.Equal(t, expectedResult == ResTrue, result, "[Cache] Expression evaluation result: "+expressionStr)
	}
}

func TestExpressions(t *testing.T) {
	params := NewParams(
		map[string]string{
			"foo":         "10",
			"unusedLabel": "3",
			"a":           "valueOfA",
			"bar":         "true",
			"anotherbar":  "t",
		},

		map[string]interface{}{
			"service": struct {
				Name   string
				Labels map[string]string
			}{
				"myservicename",
				map[string]string{
					"Name": "Value",
				},
			},
		},
	)

	tests := []struct {
		expression string
		result     int
	}{
		// true (checking bool, int, string)
		{"anotherbar == true", ResTrue},
		{"anotherbar", ResTrue},
		{"bar == true", ResTrue},
		{"bar", ResTrue},
		{"foo > 5", ResTrue},
		{"a == 'valueOfA'", ResTrue},
		{"foo > 5 && a == 'valueOfA'", ResTrue},

		// false
		{"anotherbar == 'p'", ResFalse},
		{"'A' + 'B' == 5", ResFalse},

		// IN function
		{"in(a, 'valueOfC', 'valueOfB', 'valueOfA')", ResTrue},
		{"in(foo, 10, 20, 30)", ResTrue},
		{"in(a, 'valueOfX', 'valueOfY', 'valueOfZ')", ResFalse},
		{"in(a, )", ResCompileError},
		{"in()", ResEvalError},
		{"in(5)", ResFalse},

		// check when expression involves a missing label
		{"foo > 5 && missingLabel == 'requiredValue'", ResFalse},
		{"foo > 5 && missingLabel == 239", ResFalse},

		// we are explicitly converting all integer-like params to integers, so this should always be false (expected behavior)
		{"foo == '10'", ResFalse},

		// we are explicitly converting all bool-like params to bool, so this should always be false (expected behavior)
		{"anotherbar == 't'", ResFalse},

		// check that struct expressions work
		{"service.Name == 'myservicename'", ResTrue},
		{"service.Name == 'incorrectservicename'", ResFalse},
		{"service.Labels.Name + 'B' == 'ValueB'", ResTrue},
		{"serviceMissing.LabelsMissing.Name + 'B' == 'ValueB'", ResFalse},

		// evaluation error
		{"foo + 10 + 'test' > 0", ResEvalError},

		// not a boolean
		{"'a' + 'b' + bar", ResEvalError},

		// compilation error
		{"(5 + 10 > 9", ResCompileError},
	}

	// Evaluate without cache
	for _, test := range tests {
		evaluate(t, test.expression, params, test.result)
	}

	cache := NewCache()
	for _, test := range tests {
		evaluateWithCache(t, test.expression, params, test.result, cache)
	}
}
<|endoftext|>


<commit_before>package integration_test

import (
	"testing"

	"github.com/feedhenry/mcp-standalone/pkg/data"
	"github.com/feedhenry/mcp-standalone/pkg/mobile"
	"github.com/feedhenry/mcp-standalone/pkg/mobile/integration"
	"io"
	"io/ioutil"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/kubernetes/fake"
	v1 "k8s.io/client-go/pkg/api/v1"
	ktesting "k8s.io/client-go/testing"
	"net/http"
	"strings"
)

type mockAuthChecker struct {
	checkBoolRes bool
	checkErrRes  error
}

func (mac *mockAuthChecker) Check(resource, namespace string, client mobile.ExternalHTTPRequester) (bool, error) {
	return mac.checkBoolRes, mac.checkErrRes
}

type mockExternalHTTPRequester struct {
	ResponseBody  io.ReadCloser
	ResponseError error
	ResponseCode  int
}

func (mehr *mockExternalHTTPRequester) Do(req *http.Request) (*http.Response, error) {
	res := &http.Response{Body: mehr.ResponseBody, StatusCode: mehr.ResponseCode}
	return res, mehr.ResponseError
}

func (mehr *mockExternalHTTPRequester) Get(url string) (*http.Response, error) {
	res := &http.Response{Body: mehr.ResponseBody, StatusCode: mehr.ResponseCode}
	return res, mehr.ResponseError
}

func newMockExternalHTTPRequester(body io.ReadCloser, status int, err error) mobile.ExternalHTTPRequester {
	return &mockExternalHTTPRequester{ResponseBody: body, ResponseError: err, ResponseCode: status}
}

func TestMobileServiceDiscovery(t *testing.T) {
	cases := []struct {
		Name          string
		ExpectError   bool
		ServiceCruder func() mobile.ServiceCruder
		authChecker   func() mobile.AuthChecker
		client        mobile.ExternalHTTPRequester
		Validate      func(svs []*mobile.Service, t *testing.T)
	}{
		{
			Name:   "test discover mobile services ok",
			client: newMockExternalHTTPRequester(ioutil.NopCloser(strings.NewReader("")), http.StatusOK, nil),
			authChecker: func() mobile.AuthChecker {
				return &mockAuthChecker{}
			},
			ServiceCruder: func() mobile.ServiceCruder {
				client := &fake.Clientset{}
				client.AddReactor("list", "secrets", func(a ktesting.Action) (bool, runtime.Object, error) {
					list := v1.SecretList{
						Items: []v1.Secret{
							{
								ObjectMeta: metav1.ObjectMeta{
									Labels: map[string]string{"group": "mobileapp"},
								},
								Data: map[string][]byte{
									"uri":  []byte("http://fh-sync.com"),
									"name": []byte("fh-sync-server"),
									"type": []byte("fh-sync-server"),
								},
							},
							{
								Data: map[string][]byte{
									"uri": []byte("http://somehost.com"),
								},
							},
						},
					}
					return true, &list, nil
				})
				return data.NewMobileServiceRepo(client.CoreV1().Secrets("test"))
			},
			Validate: func(svs []*mobile.Service, t *testing.T) {
				if svs == nil {
					t.Fatalf("expected a list of services but got none")
				}
				if len(svs) != 1 {
					t.Fatalf("expected 1 service to be returned but got %v", len(svs))
				}
				svc := svs[0]
				if svc.Host != "http://fh-sync.com" {
					t.Fatalf("unexpected host %s", svc.Host)
				}
				if svc.Name != "fh-sync-server" {
					t.Fatalf("expected svc name to be fh-sync-server but got %s", svc.Name)
				}
			},
		},
	}

	for _, tc := range cases {
		t.Run(tc.Name, func(t *testing.T) {
			is := integration.MobileService{}
			svcs, err := is.DiscoverMobileServices(tc.ServiceCruder(), tc.authChecker(), tc.client)
			if tc.ExpectError && err == nil {
				t.Fatalf("expected an err but got none!")
			}
			if !tc.ExpectError && err != nil {
				t.Fatalf("did not expect and err but got one %v", err)
			}
			tc.Validate(svcs, t)
		})
	}
}
<commit_msg>refactor mobile.User and re-use mocks<commit_after>package integration_test

import (
	"testing"

	"github.com/feedhenry/mcp-standalone/pkg/data"
	"github.com/feedhenry/mcp-standalone/pkg/mobile"
	"github.com/feedhenry/mcp-standalone/pkg/mobile/integration"
	"github.com/feedhenry/mcp-standalone/pkg/mock"
	"io/ioutil"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/kubernetes/fake"
	v1 "k8s.io/client-go/pkg/api/v1"
	ktesting "k8s.io/client-go/testing"
	"net/http"
	"strings"
)

type mockAuthChecker struct {
	checkBoolRes bool
	checkErrRes  error
}

func (mac *mockAuthChecker) Check(resource, namespace string, client mobile.ExternalHTTPRequester) (bool, error) {
	return mac.checkBoolRes, mac.checkErrRes
}

func TestMobileServiceDiscovery(t *testing.T) {
	cases := []struct {
		Name          string
		ExpectError   bool
		ServiceCruder func() mobile.ServiceCruder
		authChecker   func() mobile.AuthChecker
		client        mobile.ExternalHTTPRequester
		Validate      func(svs []*mobile.Service, t *testing.T)
	}{
		{
			Name: "test discover mobile services ok",
			client: &mock.Requester{
				Responder: func(host string, path string, method string, t *testing.T) (*http.Response, error) {
					return &http.Response{
						StatusCode: http.StatusOK,
						Body:       ioutil.NopCloser(strings.NewReader("")),
					}, nil
				},
			},
			authChecker: func() mobile.AuthChecker {
				return &mockAuthChecker{}
			},
			ServiceCruder: func() mobile.ServiceCruder {
				client := &fake.Clientset{}
				client.AddReactor("list", "secrets", func(a ktesting.Action) (bool, runtime.Object, error) {
					list := v1.SecretList{
						Items: []v1.Secret{
							{
								ObjectMeta: metav1.ObjectMeta{
									Labels: map[string]string{"group": "mobileapp"},
								},
								Data: map[string][]byte{
									"uri":  []byte("http://fh-sync.com"),
									"name": []byte("fh-sync-server"),
									"type": []byte("fh-sync-server"),
								},
							},
							{
								Data: map[string][]byte{
									"uri": []byte("http://somehost.com"),
								},
							},
						},
					}
					return true, &list, nil
				})
				return data.NewMobileServiceRepo(client.CoreV1().Secrets("test"))
			},
			Validate: func(svs []*mobile.Service, t *testing.T) {
				if svs == nil {
					t.Fatalf("expected a list of services but got none")
				}
				if len(svs) != 1 {
					t.Fatalf("expected 1 service to be returned but got %v", len(svs))
				}
				svc := svs[0]
				if svc.Host != "http://fh-sync.com" {
					t.Fatalf("unexpected host %s", svc.Host)
				}
				if svc.Name != "fh-sync-server" {
					t.Fatalf("expected svc name to be fh-sync-server but got %s", svc.Name)
				}
			},
		},
	}

	for _, tc := range cases {
		t.Run(tc.Name, func(t *testing.T) {
			is := integration.MobileService{}
			svcs, err := is.DiscoverMobileServices(tc.ServiceCruder(), tc.authChecker(), tc.client)
			if tc.ExpectError && err == nil {
				t.Fatalf("expected an err but got none!")
			}
			if !tc.ExpectError && err != nil {
				t.Fatalf("did not expect and err but got one %v", err)
			}
			tc.Validate(svcs, t)
		})
	}
}
<|endoftext|>


<commit_before>package status_test

import (
	"testing"

	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/status"
	"github.com/stretchr/testify/assert"
)

func TestLiqFromRaw(t *testing.T) {
	testCases := map[string]struct {
		data    []interface{}
		err     func(*testing.T, error)
		success func(*testing.T, interface{})
	}{
		"empty slice": {
			data: []interface{}{},
			err: func(t *testing.T, got error) {
				assert.Error(t, got)
			},
			success: func(t *testing.T, got interface{}) {
				assert.Nil(t, got)
			},
		},
		"liq raw invalid pld": {
			data: []interface{}{1},
			err: func(t *testing.T, got error) {
				assert.Error(t, got)
			},
			success: func(t *testing.T, got interface{}) {
				assert.Nil(t, got)
			},
		},
		"liquidation raw valid pld": {
			data: []interface{}{
				"pos", 145400868, 1609144352338, nil, "tETHF0:USTF0",
				-1.67288094, 730.96, nil, 1, 1, nil, 736.13,
			},
			err: func(t *testing.T, got error) {
				assert.Nil(t, got)
			},
			success: func(t *testing.T, got interface{}) {
				assert.Equal(t, got, &status.Liquidation{
					Symbol:        "tETHF0:USTF0",
					PositionID:    145400868,
					MTS:           1609144352338,
					Amount:        -1.67288094,
					BasePrice:     730.96,
					IsMatch:       1,
					IsMarketSold:  1,
					PriceAcquired: 736.13,
				})
			},
		},
	}

	for k, v := range testCases {
		t.Run(k, func(t *testing.T) {
			got, err := status.LiqFromRaw(v.data)
			v.err(t, err)
			v.success(t, got)
		})
	}
}
<commit_msg>status.LiqSnapshotFromRaw test coverage<commit_after>package status_test

import (
	"testing"

	"github.com/bitfinexcom/bitfinex-api-go/pkg/models/status"
	"github.com/stretchr/testify/assert"
)

func TestLiqFromRaw(t *testing.T) {
	testCases := map[string]struct {
		data    []interface{}
		err     func(*testing.T, error)
		success func(*testing.T, interface{})
	}{
		"empty pld": {
			data: []interface{}{},
			err: func(t *testing.T, got error) {
				assert.Error(t, got)
			},
			success: func(t *testing.T, got interface{}) {
				assert.Nil(t, got)
			},
		},
		"invalid pld": {
			data: []interface{}{1},
			err: func(t *testing.T, got error) {
				assert.Error(t, got)
			},
			success: func(t *testing.T, got interface{}) {
				assert.Nil(t, got)
			},
		},
		"valid pld": {
			data: []interface{}{
				"pos", 145400868, 1609144352338, nil, "tETHF0:USTF0",
				-1.67288094, 730.96, nil, 1, 1, nil, 736.13,
			},
			err: func(t *testing.T, got error) {
				assert.Nil(t, got)
			},
			success: func(t *testing.T, got interface{}) {
				assert.Equal(t, got, &status.Liquidation{
					Symbol:        "tETHF0:USTF0",
					PositionID:    145400868,
					MTS:           1609144352338,
					Amount:        -1.67288094,
					BasePrice:     730.96,
					IsMatch:       1,
					IsMarketSold:  1,
					PriceAcquired: 736.13,
				})
			},
		},
	}

	for k, v := range testCases {
		t.Run(k, func(t *testing.T) {
			got, err := status.LiqFromRaw(v.data)
			v.err(t, err)
			v.success(t, got)
		})
	}
}

func TestLiqSnapshotFromRaw(t *testing.T) {
	testCases := map[string]struct {
		data    [][]interface{}
		err     func(*testing.T, error)
		success func(*testing.T, interface{})
	}{
		"empty slice": {
			data: [][]interface{}{},
			err: func(t *testing.T, got error) {
				assert.Error(t, got)
			},
			success: func(t *testing.T, got interface{}) {
				assert.Nil(t, got)
			},
		},
		"invalid pld": {
			data: [][]interface{}{{1}},
			err: func(t *testing.T, got error) {
				assert.Error(t, got)
			},
			success: func(t *testing.T, got interface{}) {
				assert.Nil(t, got)
			},
		},
		"valid pld": {
			data: [][]interface{}{
				{
					"pos", 145400868, 1609144352338, nil, "tETHF0:USTF0",
					-1.67288094, 730.96, nil, 1, 1, nil, 736.13,
				},
				{
					"pos", 145400869, 1609144352338, nil, "tETHF0:USTF0",
					-1.67288094, 730.96, nil, 1, 1, nil, 736.13, nil, 1000,
				},
			},
			err: func(t *testing.T, got error) {
				assert.Nil(t, got)
			},
			success: func(t *testing.T, got interface{}) {
				assert.Equal(t, got, &status.LiquidationsSnapshot{
					Snapshot: []*status.Liquidation{
						{
							Symbol:        "tETHF0:USTF0",
							PositionID:    145400868,
							MTS:           1609144352338,
							Amount:        -1.67288094,
							BasePrice:     730.96,
							IsMatch:       1,
							IsMarketSold:  1,
							PriceAcquired: 736.13,
						},
						{
							Symbol:        "tETHF0:USTF0",
							PositionID:    145400869,
							MTS:           1609144352338,
							Amount:        -1.67288094,
							BasePrice:     730.96,
							IsMatch:       1,
							IsMarketSold:  1,
							PriceAcquired: 736.13,
						},
					},
				})
			},
		},
	}

	for k, v := range testCases {
		t.Run(k, func(t *testing.T) {
			got, err := status.LiqSnapshotFromRaw(v.data)
			v.err(t, err)
			v.success(t, got)
		})
	}
}
<|endoftext|>


<commit_before>// Copyright 2021 The gVisor 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 msgqueue implements System V message queues.
package msgqueue

import (
	"gvisor.dev/gvisor/pkg/abi/linux"
	"gvisor.dev/gvisor/pkg/context"
	"gvisor.dev/gvisor/pkg/errors/linuxerr"
	"gvisor.dev/gvisor/pkg/sentry/fs"
	"gvisor.dev/gvisor/pkg/sentry/kernel/auth"
	"gvisor.dev/gvisor/pkg/sentry/kernel/ipc"
	ktime "gvisor.dev/gvisor/pkg/sentry/kernel/time"
	"gvisor.dev/gvisor/pkg/sync"
	"gvisor.dev/gvisor/pkg/waiter"
)

const (
	// System-wide limit for maximum number of queues.
	maxQueues = linux.MSGMNI

	// Maximum size of a queue in bytes.
	maxQueueBytes = linux.MSGMNB

	// Maximum size of a message in bytes.
	maxMessageBytes = linux.MSGMAX
)

// Registry contains a set of message queues that can be referenced using keys
// or IDs.
//
// +stateify savable
type Registry struct {
	// mu protects all the fields below.
	mu sync.Mutex `state:"nosave"`

	// reg defines basic fields and operations needed for all SysV registries.
	reg *ipc.Registry
}

// NewRegistry returns a new Registry ready to be used.
func NewRegistry(userNS *auth.UserNamespace) *Registry {
	return &Registry{
		reg: ipc.NewRegistry(userNS),
	}
}

// Queue represents a SysV message queue, described by sysvipc(7).
//
// +stateify savable
type Queue struct {
	// registry is the registry owning this queue. Immutable.
	registry *Registry

	// mu protects all the fields below.
	mu sync.Mutex `state:"nosave"`

	// dead is set to true when a queue is removed from the registry and should
	// not be used. Operations on the queue should check dead, and return
	// EIDRM if set to true.
	dead bool

	// obj defines basic fields that should be included in all SysV IPC objects.
	obj *ipc.Object

	// senders holds a queue of blocked message senders. Senders are notified
	// when enough space is available in the queue to insert their message.
	senders waiter.Queue

	// receivers holds a queue of blocked receivers. Receivers are notified
	// when a new message is inserted into the queue and can be received.
	receivers waiter.Queue

	// messages is a list of sent messages.
	messages msgList

	// sendTime is the last time a msgsnd was perfomed.
	sendTime ktime.Time

	// receiveTime is the last time a msgrcv was performed.
	receiveTime ktime.Time

	// changeTime is the last time the queue was modified using msgctl.
	changeTime ktime.Time

	// byteCount is the current number of message bytes in the queue.
	byteCount uint64

	// messageCount is the current number of messages in the queue.
	messageCount uint64

	// maxBytes is the maximum allowed number of bytes in the queue, and is also
	// used as a limit for the number of total possible messages.
	maxBytes uint64

	// sendPID is the PID of the process that performed the last msgsnd.
	sendPID int32

	// receivePID is the PID of the process that performed the last msgrcv.
	receivePID int32
}

// Message represents a message exchanged through a Queue via msgsnd(2) and
// msgrcv(2).
//
// +stateify savable
type Message struct {
	msgEntry

	// Type is an integer representing the type of the sent message.
	Type int64

	// Text is an untyped block of memory.
	Text []byte

	// Size is the size of Text.
	Size uint64
}

// Blocker is used for blocking Queue.Send, and Queue.Receive calls that serves
// as an abstracted version of kernel.Task. kernel.Task is not directly used to
// prevent circular dependencies.
type Blocker interface {
	Block(C <-chan struct{}) error
}

// FindOrCreate creates a new message queue or returns an existing one. See
// msgget(2).
func (r *Registry) FindOrCreate(ctx context.Context, key ipc.Key, mode linux.FileMode, private, create, exclusive bool) (*Queue, error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	if !private {
		queue, err := r.reg.Find(ctx, key, mode, create, exclusive)
		if err != nil {
			return nil, err
		}

		if queue != nil {
			return queue.(*Queue), nil
		}
	}

	// Check system-wide limits.
	if r.reg.ObjectCount() >= maxQueues {
		return nil, linuxerr.ENOSPC
	}

	return r.newQueueLocked(ctx, key, fs.FileOwnerFromContext(ctx), fs.FilePermsFromMode(mode))
}

// newQueueLocked creates a new queue using the given fields. An error is
// returned if there're no more available identifiers.
//
// Precondition: r.mu must be held.
func (r *Registry) newQueueLocked(ctx context.Context, key ipc.Key, creator fs.FileOwner, perms fs.FilePermissions) (*Queue, error) {
	q := &Queue{
		registry:    r,
		obj:         ipc.NewObject(r.reg.UserNS, key, creator, creator, perms),
		sendTime:    ktime.ZeroTime,
		receiveTime: ktime.ZeroTime,
		changeTime:  ktime.NowFromContext(ctx),
		maxBytes:    maxQueueBytes,
	}

	err := r.reg.Register(q)
	if err != nil {
		return nil, err
	}
	return q, nil
}

// Remove removes the queue with specified ID. All waiters (readers and
// writers) and writers will be awakened and fail. Remove will return an error
// if the ID is invalid, or the the user doesn't have privileges.
func (r *Registry) Remove(id ipc.ID, creds *auth.Credentials) error {
	r.mu.Lock()
	defer r.mu.Unlock()

	r.reg.Remove(id, creds)
	return nil
}

// FindByID returns the queue with the specified ID and an error if the ID
// doesn't exist.
func (r *Registry) FindByID(id ipc.ID) (*Queue, error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	mech := r.reg.FindByID(id)
	if mech == nil {
		return nil, linuxerr.EINVAL
	}
	return mech.(*Queue), nil
}

// Send appends a message to the message queue, and returns an error if sending
// fails. See msgsnd(2).
func (q *Queue) Send(ctx context.Context, m Message, b Blocker, wait bool, pid int32) (err error) {
	// Try to perform a non-blocking send using queue.append. If EWOULDBLOCK
	// is returned, start the blocking procedure. Otherwise, return normally.
	creds := auth.CredentialsFromContext(ctx)
	if err := q.append(ctx, m, creds, pid); err != linuxerr.EWOULDBLOCK {
		return err
	}

	if !wait {
		return linuxerr.EAGAIN
	}

	e, ch := waiter.NewChannelEntry(nil)
	q.senders.EventRegister(&e, waiter.EventOut)

	for {
		if err = q.append(ctx, m, creds, pid); err != linuxerr.EWOULDBLOCK {
			break
		}
		b.Block(ch)
	}

	q.senders.EventUnregister(&e)
	return err
}

// append appends a message to the queue's message list and notifies waiting
// receivers that a message has been inserted. It returns an error if adding
// the message would cause the queue to exceed its maximum capacity, which can
// be used as a signal to block the task. Other errors should be returned as is.
func (q *Queue) append(ctx context.Context, m Message, creds *auth.Credentials, pid int32) error {
	if m.Type <= 0 {
		return linuxerr.EINVAL
	}

	q.mu.Lock()
	defer q.mu.Unlock()

	if !q.obj.CheckPermissions(creds, fs.PermMask{Write: true}) {
		// The calling process does not have write permission on the message
		// queue, and does not have the CAP_IPC_OWNER capability in the user
		// namespace that governs its IPC namespace.
		return linuxerr.EACCES
	}

	// Queue was removed while the process was waiting.
	if q.dead {
		return linuxerr.EIDRM
	}

	// Check if sufficient space is available (the queue isn't full.) From
	// the man pages:
	//
	// "A message queue is considered to be full if either of the following
	// conditions is true:
	//
	//  • Adding a new message to the queue would cause the total number
	//    of bytes in the queue to exceed the queue's maximum size (the
	//    msg_qbytes field).
	//
	//  • Adding another message to the queue would cause the total
	//    number of messages in the queue to exceed the queue's maximum
	//    size (the msg_qbytes field).  This check is necessary to
	//    prevent an unlimited number of zero-length messages being
	//    placed on the queue.  Although such messages contain no data,
	//    they nevertheless consume (locked) kernel memory."
	//
	// The msg_qbytes field in our implementation is q.maxBytes.
	if m.Size+q.byteCount > q.maxBytes || q.messageCount+1 > q.maxBytes {
		return linuxerr.EWOULDBLOCK
	}

	// Copy the message into the queue.
	q.messages.PushBack(&m)

	q.byteCount += m.Size
	q.messageCount++
	q.sendPID = pid
	q.sendTime = ktime.NowFromContext(ctx)

	// Notify receivers about the new message.
	q.receivers.Notify(waiter.EventIn)

	return nil
}

// Lock implements ipc.Mechanism.Lock.
func (q *Queue) Lock() {
	q.mu.Lock()
}

// Unlock implements ipc.mechanism.Unlock.
//
// +checklocksignore
func (q *Queue) Unlock() {
	q.mu.Unlock()
}

// Object implements ipc.Mechanism.Object.
func (q *Queue) Object() *ipc.Object {
	return q.obj
}

// Destroy implements ipc.Mechanism.Destroy.
func (q *Queue) Destroy() {
	q.dead = true

	// Notify waiters. Senders and receivers will try to run, and return an
	// error (EIDRM). Waiters should remove themselves from the queue after
	// waking up.
	q.senders.Notify(waiter.EventOut)
	q.receivers.Notify(waiter.EventIn)
}

// ID returns queue's ID.
func (q *Queue) ID() ipc.ID {
	return q.obj.ID
}
<commit_msg>Implement Queue.Receive.<commit_after>// Copyright 2021 The gVisor 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 msgqueue implements System V message queues.
package msgqueue

import (
	"gvisor.dev/gvisor/pkg/abi/linux"
	"gvisor.dev/gvisor/pkg/context"
	"gvisor.dev/gvisor/pkg/errors/linuxerr"
	"gvisor.dev/gvisor/pkg/sentry/fs"
	"gvisor.dev/gvisor/pkg/sentry/kernel/auth"
	"gvisor.dev/gvisor/pkg/sentry/kernel/ipc"
	ktime "gvisor.dev/gvisor/pkg/sentry/kernel/time"
	"gvisor.dev/gvisor/pkg/sync"
	"gvisor.dev/gvisor/pkg/waiter"
)

const (
	// System-wide limit for maximum number of queues.
	maxQueues = linux.MSGMNI

	// Maximum size of a queue in bytes.
	maxQueueBytes = linux.MSGMNB

	// Maximum size of a message in bytes.
	maxMessageBytes = linux.MSGMAX
)

// Registry contains a set of message queues that can be referenced using keys
// or IDs.
//
// +stateify savable
type Registry struct {
	// mu protects all the fields below.
	mu sync.Mutex `state:"nosave"`

	// reg defines basic fields and operations needed for all SysV registries.
	reg *ipc.Registry
}

// NewRegistry returns a new Registry ready to be used.
func NewRegistry(userNS *auth.UserNamespace) *Registry {
	return &Registry{
		reg: ipc.NewRegistry(userNS),
	}
}

// Queue represents a SysV message queue, described by sysvipc(7).
//
// +stateify savable
type Queue struct {
	// registry is the registry owning this queue. Immutable.
	registry *Registry

	// mu protects all the fields below.
	mu sync.Mutex `state:"nosave"`

	// dead is set to true when a queue is removed from the registry and should
	// not be used. Operations on the queue should check dead, and return
	// EIDRM if set to true.
	dead bool

	// obj defines basic fields that should be included in all SysV IPC objects.
	obj *ipc.Object

	// senders holds a queue of blocked message senders. Senders are notified
	// when enough space is available in the queue to insert their message.
	senders waiter.Queue

	// receivers holds a queue of blocked receivers. Receivers are notified
	// when a new message is inserted into the queue and can be received.
	receivers waiter.Queue

	// messages is a list of sent messages.
	messages msgList

	// sendTime is the last time a msgsnd was perfomed.
	sendTime ktime.Time

	// receiveTime is the last time a msgrcv was performed.
	receiveTime ktime.Time

	// changeTime is the last time the queue was modified using msgctl.
	changeTime ktime.Time

	// byteCount is the current number of message bytes in the queue.
	byteCount uint64

	// messageCount is the current number of messages in the queue.
	messageCount uint64

	// maxBytes is the maximum allowed number of bytes in the queue, and is also
	// used as a limit for the number of total possible messages.
	maxBytes uint64

	// sendPID is the PID of the process that performed the last msgsnd.
	sendPID int32

	// receivePID is the PID of the process that performed the last msgrcv.
	receivePID int32
}

// Message represents a message exchanged through a Queue via msgsnd(2) and
// msgrcv(2).
//
// +stateify savable
type Message struct {
	msgEntry

	// Type is an integer representing the type of the sent message.
	Type int64

	// Text is an untyped block of memory.
	Text []byte

	// Size is the size of Text.
	Size uint64
}

// Blocker is used for blocking Queue.Send, and Queue.Receive calls that serves
// as an abstracted version of kernel.Task. kernel.Task is not directly used to
// prevent circular dependencies.
type Blocker interface {
	Block(C <-chan struct{}) error
}

// FindOrCreate creates a new message queue or returns an existing one. See
// msgget(2).
func (r *Registry) FindOrCreate(ctx context.Context, key ipc.Key, mode linux.FileMode, private, create, exclusive bool) (*Queue, error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	if !private {
		queue, err := r.reg.Find(ctx, key, mode, create, exclusive)
		if err != nil {
			return nil, err
		}

		if queue != nil {
			return queue.(*Queue), nil
		}
	}

	// Check system-wide limits.
	if r.reg.ObjectCount() >= maxQueues {
		return nil, linuxerr.ENOSPC
	}

	return r.newQueueLocked(ctx, key, fs.FileOwnerFromContext(ctx), fs.FilePermsFromMode(mode))
}

// newQueueLocked creates a new queue using the given fields. An error is
// returned if there're no more available identifiers.
//
// Precondition: r.mu must be held.
func (r *Registry) newQueueLocked(ctx context.Context, key ipc.Key, creator fs.FileOwner, perms fs.FilePermissions) (*Queue, error) {
	q := &Queue{
		registry:    r,
		obj:         ipc.NewObject(r.reg.UserNS, key, creator, creator, perms),
		sendTime:    ktime.ZeroTime,
		receiveTime: ktime.ZeroTime,
		changeTime:  ktime.NowFromContext(ctx),
		maxBytes:    maxQueueBytes,
	}

	err := r.reg.Register(q)
	if err != nil {
		return nil, err
	}
	return q, nil
}

// Remove removes the queue with specified ID. All waiters (readers and
// writers) and writers will be awakened and fail. Remove will return an error
// if the ID is invalid, or the the user doesn't have privileges.
func (r *Registry) Remove(id ipc.ID, creds *auth.Credentials) error {
	r.mu.Lock()
	defer r.mu.Unlock()

	r.reg.Remove(id, creds)
	return nil
}

// FindByID returns the queue with the specified ID and an error if the ID
// doesn't exist.
func (r *Registry) FindByID(id ipc.ID) (*Queue, error) {
	r.mu.Lock()
	defer r.mu.Unlock()

	mech := r.reg.FindByID(id)
	if mech == nil {
		return nil, linuxerr.EINVAL
	}
	return mech.(*Queue), nil
}

// Send appends a message to the message queue, and returns an error if sending
// fails. See msgsnd(2).
func (q *Queue) Send(ctx context.Context, m Message, b Blocker, wait bool, pid int32) (err error) {
	// Try to perform a non-blocking send using queue.append. If EWOULDBLOCK
	// is returned, start the blocking procedure. Otherwise, return normally.
	creds := auth.CredentialsFromContext(ctx)
	if err := q.append(ctx, m, creds, pid); err != linuxerr.EWOULDBLOCK {
		return err
	}

	if !wait {
		return linuxerr.EAGAIN
	}

	e, ch := waiter.NewChannelEntry(nil)
	q.senders.EventRegister(&e, waiter.EventOut)

	for {
		if err = q.append(ctx, m, creds, pid); err != linuxerr.EWOULDBLOCK {
			break
		}
		b.Block(ch)
	}

	q.senders.EventUnregister(&e)
	return err
}

// append appends a message to the queue's message list and notifies waiting
// receivers that a message has been inserted. It returns an error if adding
// the message would cause the queue to exceed its maximum capacity, which can
// be used as a signal to block the task. Other errors should be returned as is.
func (q *Queue) append(ctx context.Context, m Message, creds *auth.Credentials, pid int32) error {
	if m.Type <= 0 {
		return linuxerr.EINVAL
	}

	q.mu.Lock()
	defer q.mu.Unlock()

	if !q.obj.CheckPermissions(creds, fs.PermMask{Write: true}) {
		// The calling process does not have write permission on the message
		// queue, and does not have the CAP_IPC_OWNER capability in the user
		// namespace that governs its IPC namespace.
		return linuxerr.EACCES
	}

	// Queue was removed while the process was waiting.
	if q.dead {
		return linuxerr.EIDRM
	}

	// Check if sufficient space is available (the queue isn't full.) From
	// the man pages:
	//
	// "A message queue is considered to be full if either of the following
	// conditions is true:
	//
	//  • Adding a new message to the queue would cause the total number
	//    of bytes in the queue to exceed the queue's maximum size (the
	//    msg_qbytes field).
	//
	//  • Adding another message to the queue would cause the total
	//    number of messages in the queue to exceed the queue's maximum
	//    size (the msg_qbytes field).  This check is necessary to
	//    prevent an unlimited number of zero-length messages being
	//    placed on the queue.  Although such messages contain no data,
	//    they nevertheless consume (locked) kernel memory."
	//
	// The msg_qbytes field in our implementation is q.maxBytes.
	if m.Size+q.byteCount > q.maxBytes || q.messageCount+1 > q.maxBytes {
		return linuxerr.EWOULDBLOCK
	}

	// Copy the message into the queue.
	q.messages.PushBack(&m)

	q.byteCount += m.Size
	q.messageCount++
	q.sendPID = pid
	q.sendTime = ktime.NowFromContext(ctx)

	// Notify receivers about the new message.
	q.receivers.Notify(waiter.EventIn)

	return nil
}

// Receive removes a message from the queue and returns it. See msgrcv(2).
func (q *Queue) Receive(ctx context.Context, b Blocker, mType int64, maxSize int64, wait, truncate, except bool, pid int32) (msg *Message, err error) {
	if maxSize < 0 || maxSize > maxMessageBytes {
		return nil, linuxerr.EINVAL
	}
	max := uint64(maxSize)

	// Try to perform a non-blocking receive using queue.pop. If EWOULDBLOCK
	// is returned, start the blocking procedure. Otherwise, return normally.
	creds := auth.CredentialsFromContext(ctx)
	if msg, err := q.pop(ctx, creds, mType, max, truncate, except, pid); err != linuxerr.EWOULDBLOCK {
		return msg, err
	}

	if !wait {
		return nil, linuxerr.ENOMSG
	}

	e, ch := waiter.NewChannelEntry(nil)
	q.receivers.EventRegister(&e, waiter.EventIn)

	for {
		if msg, err = q.pop(ctx, creds, mType, max, truncate, except, pid); err != linuxerr.EWOULDBLOCK {
			break
		}
		b.Block(ch)
	}
	q.receivers.EventUnregister(&e)
	return msg, err
}

// pop pops the first message from the queue that matches the given type. It
// returns an error for all the cases specified in msgrcv(2). If the queue is
// empty or no message of the specified type is available, a EWOULDBLOCK error
// is returned, which can then be used as a signal to block the process or fail.
func (q *Queue) pop(ctx context.Context, creds *auth.Credentials, mType int64, maxSize uint64, truncate, except bool, pid int32) (msg *Message, _ error) {
	q.mu.Lock()
	defer q.mu.Unlock()

	if !q.obj.CheckPermissions(creds, fs.PermMask{Read: true}) {
		// The calling process does not have read permission on the message
		// queue, and does not have the CAP_IPC_OWNER capability in the user
		// namespace that governs its IPC namespace.
		return nil, linuxerr.EACCES
	}

	// Queue was removed while the process was waiting.
	if q.dead {
		return nil, linuxerr.EIDRM
	}

	if q.messages.Empty() {
		return nil, linuxerr.EWOULDBLOCK
	}

	// Get a message from the queue.
	switch {
	case mType == 0:
		msg = q.messages.Front()
	case mType > 0:
		msg = q.msgOfType(mType, except)
	case mType < 0:
		msg = q.msgOfTypeLessThan(-1 * mType)
	}

	// If no message exists, return a blocking singal.
	if msg == nil {
		return nil, linuxerr.EWOULDBLOCK
	}

	// Check message's size is acceptable.
	if maxSize < msg.Size {
		if !truncate {
			return nil, linuxerr.E2BIG
		}
		msg.Size = maxSize
		msg.Text = msg.Text[:maxSize+1]
	}

	q.messages.Remove(msg)

	q.byteCount -= msg.Size
	q.messageCount--
	q.receivePID = pid
	q.receiveTime = ktime.NowFromContext(ctx)

	// Notify senders about available space.
	q.senders.Notify(waiter.EventOut)

	return msg, nil
}

// msgOfType returns the first message with the specified type, nil if no
// message is found. If except is true, the first message of a type not equal
// to mType will be returned.
//
// Precondition: caller must hold q.mu.
func (q *Queue) msgOfType(mType int64, except bool) *Message {
	if except {
		for msg := q.messages.Front(); msg != nil; msg = msg.Next() {
			if msg.Type != mType {
				return msg
			}
		}
		return nil
	}

	for msg := q.messages.Front(); msg != nil; msg = msg.Next() {
		if msg.Type == mType {
			return msg
		}
	}
	return nil
}

// msgOfTypeLessThan return the the first message with the lowest type less
// than or equal to mType, nil if no such message exists.
//
// Precondition: caller must hold q.mu.
func (q *Queue) msgOfTypeLessThan(mType int64) (m *Message) {
	min := mType
	for msg := q.messages.Front(); msg != nil; msg = msg.Next() {
		if msg.Type <= mType && msg.Type < min {
			m = msg
			min = msg.Type
		}
	}
	return m
}

// Lock implements ipc.Mechanism.Lock.
func (q *Queue) Lock() {
	q.mu.Lock()
}

// Unlock implements ipc.mechanism.Unlock.
//
// +checklocksignore
func (q *Queue) Unlock() {
	q.mu.Unlock()
}

// Object implements ipc.Mechanism.Object.
func (q *Queue) Object() *ipc.Object {
	return q.obj
}

// Destroy implements ipc.Mechanism.Destroy.
func (q *Queue) Destroy() {
	q.dead = true

	// Notify waiters. Senders and receivers will try to run, and return an
	// error (EIDRM). Waiters should remove themselves from the queue after
	// waking up.
	q.senders.Notify(waiter.EventOut)
	q.receivers.Notify(waiter.EventIn)
}

// ID returns queue's ID.
func (q *Queue) ID() ipc.ID {
	return q.obj.ID
}
<|endoftext|>


<commit_before>package v1

import (
	"encoding/json"
	"net/http"
)

const (
	StatusImplementationError = "implementationError"
	StatusUnauthorized        = "unauthorized"
	StatusError               = "error"
	StatusSuccess             = "success"
)

// ServiceResponse represents a general response container for (payment-related) API
// requests
type ServiceResponse struct {
	HttpStatus int `json:"-"`
	Status     string
	Info       string
	Response   interface{}
	Error      interface{}
}

// default service responses
var (
	ErrReadJson = ServiceResponse{
		http.StatusBadRequest,
		StatusImplementationError,
		"could not read request",
		nil,
		"JSON decoding error",
	}
	ErrUnauthorized = ServiceResponse{
		http.StatusUnauthorized,
		StatusUnauthorized,
		"unauthorized",
		nil,
		"unauthorized",
	}
	ErrDatabase = ServiceResponse{
		http.StatusInternalServerError,
		StatusError,
		"database error",
		nil,
		"database error",
	}
	ErrSystem = ServiceResponse{
		http.StatusInternalServerError,
		StatusError,
		"internal error",
		nil,
		"internal error",
	}
	ErrInval = ServiceResponse{
		http.StatusBadRequest,
		StatusImplementationError,
		"invalid value",
		nil,
		"invalid value",
	}
	ErrNotFound = ServiceResponse{
		http.StatusNotFound,
		StatusError,
		"resource not found",
		nil,
		"resource not found",
	}
)

func (sr *ServiceResponse) Write(w http.ResponseWriter) error {

	// set default http states
	if sr.HttpStatus == 0 && sr.Status == StatusSuccess && sr.Error == nil {
		sr.HttpStatus = http.StatusOK
	} else if sr.HttpStatus == 0 {
		sr.HttpStatus = http.StatusInternalServerError
	}

	w.WriteHeader(sr.HttpStatus)

	// json encode response struct
	je := json.NewEncoder(w)
	err := je.Encode(sr)
	if err != nil {
		return err
	}

	return err
}
<commit_msg>include service version in service response: service version 1.1<commit_after>package v1

import (
	"encoding/json"
	"net/http"
)

const (
	StatusImplementationError = "implementationError"
	StatusUnauthorized        = "unauthorized"
	StatusError               = "error"
	StatusSuccess             = "success"
)

const (
	ServiceVersion = "1.1"
)

// ServiceResponse represents a general response container for (payment-related) API
// requests
type ServiceResponse struct {
	HttpStatus int `json:"-"`
	Version    string
	Status     string
	Info       string
	Response   interface{}
	Error      interface{}
}

// default service responses
var (
	ErrReadJson = ServiceResponse{
		http.StatusBadRequest,
		ServiceVersion,
		StatusImplementationError,
		"could not read request",
		nil,
		"JSON decoding error",
	}
	ErrUnauthorized = ServiceResponse{
		http.StatusUnauthorized,
		ServiceVersion,
		StatusUnauthorized,
		"unauthorized",
		nil,
		"unauthorized",
	}
	ErrDatabase = ServiceResponse{
		http.StatusInternalServerError,
		ServiceVersion,
		StatusError,
		"database error",
		nil,
		"database error",
	}
	ErrSystem = ServiceResponse{
		http.StatusInternalServerError,
		ServiceVersion,
		StatusError,
		"internal error",
		nil,
		"internal error",
	}
	ErrInval = ServiceResponse{
		http.StatusBadRequest,
		ServiceVersion,
		StatusImplementationError,
		"invalid value",
		nil,
		"invalid value",
	}
	ErrNotFound = ServiceResponse{
		http.StatusNotFound,
		ServiceVersion,
		StatusError,
		"resource not found",
		nil,
		"resource not found",
	}
)

func (sr *ServiceResponse) Write(w http.ResponseWriter) error {

	// set default http states
	if sr.HttpStatus == 0 && sr.Status == StatusSuccess && sr.Error == nil {
		sr.HttpStatus = http.StatusOK
	} else if sr.HttpStatus == 0 {
		sr.HttpStatus = http.StatusInternalServerError
	}

	w.WriteHeader(sr.HttpStatus)

	// json encode response struct
	je := json.NewEncoder(w)
	err := je.Encode(sr)
	if err != nil {
		return err
	}

	return err
}
<|endoftext|>


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

package zap

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"math"
	"reflect"
	"sort"

	"github.com/blevesearch/bleve/index"
	"github.com/blevesearch/bleve/size"
	"github.com/golang/snappy"
)

var reflectStaticSizedocValueReader int

func init() {
	var dvi docValueReader
	reflectStaticSizedocValueReader = int(reflect.TypeOf(dvi).Size())
}

type docNumTermsVisitor func(docNum uint64, terms []byte) error

type docValueReader struct {
	field          string
	curChunkNum    uint64
	chunkOffsets   []uint64
	dvDataLoc      uint64
	curChunkHeader []MetaData
	curChunkData   []byte // compressed data cache
	uncompressed   []byte // temp buf for snappy decompression
}

func (di *docValueReader) size() int {
	return reflectStaticSizedocValueReader + size.SizeOfPtr +
		len(di.field) +
		len(di.chunkOffsets)*size.SizeOfUint64 +
		len(di.curChunkHeader)*reflectStaticSizeMetaData +
		len(di.curChunkData)
}

func (di *docValueReader) cloneInto(rv *docValueReader) *docValueReader {
	if rv == nil {
		rv = &docValueReader{}
	}

	rv.field = di.field
	rv.curChunkNum = math.MaxUint64
	rv.chunkOffsets = di.chunkOffsets // immutable, so it's sharable
	rv.dvDataLoc = di.dvDataLoc
	rv.curChunkHeader = nil
	rv.curChunkData = nil
	rv.uncompressed = nil

	return rv
}

func (di *docValueReader) fieldName() string {
	return di.field
}

func (di *docValueReader) curChunkNumber() uint64 {
	return di.curChunkNum
}

func (s *SegmentBase) loadFieldDocValueReader(field string,
	fieldDvLocStart, fieldDvLocEnd uint64) (*docValueReader, error) {
	// get the docValue offset for the given fields
	if fieldDvLocStart == fieldNotUninverted {
		return nil, fmt.Errorf("loadFieldDocValueReader: "+
			"no docValues found for field: %s", field)
	}

	// read the number of chunks, and chunk offsets position
	var numChunks, chunkOffsetsPosition uint64

	if fieldDvLocEnd-fieldDvLocStart > 16 {
		numChunks = binary.BigEndian.Uint64(s.mem[fieldDvLocEnd-8 : fieldDvLocEnd])
		// read the length of chunk offsets
		chunkOffsetsLen := binary.BigEndian.Uint64(s.mem[fieldDvLocEnd-16 : fieldDvLocEnd-8])
		// acquire position of chunk offsets
		chunkOffsetsPosition = (fieldDvLocEnd - 16) - chunkOffsetsLen
	}

	fdvIter := &docValueReader{
		curChunkNum:  math.MaxUint64,
		field:        field,
		chunkOffsets: make([]uint64, int(numChunks)),
	}

	// read the chunk offsets
	var offset uint64
	for i := 0; i < int(numChunks); i++ {
		loc, read := binary.Uvarint(s.mem[chunkOffsetsPosition+offset : chunkOffsetsPosition+offset+binary.MaxVarintLen64])
		if read <= 0 {
			return nil, fmt.Errorf("corrupted chunk offset during segment load")
		}
		fdvIter.chunkOffsets[i] = loc
		offset += uint64(read)
	}

	// set the data offset
	fdvIter.dvDataLoc = fieldDvLocStart

	return fdvIter, nil
}

func (di *docValueReader) loadDvChunk(chunkNumber uint64, s *SegmentBase) error {
	// advance to the chunk where the docValues
	// reside for the given docNum
	destChunkDataLoc, curChunkEnd := di.dvDataLoc, di.dvDataLoc
	start, end := readChunkBoundary(int(chunkNumber), di.chunkOffsets)
	destChunkDataLoc += start
	curChunkEnd += end

	// read the number of docs reside in the chunk
	numDocs, read := binary.Uvarint(s.mem[destChunkDataLoc : destChunkDataLoc+binary.MaxVarintLen64])
	if read <= 0 {
		return fmt.Errorf("failed to read the chunk")
	}
	chunkMetaLoc := destChunkDataLoc + uint64(read)

	offset := uint64(0)
	di.curChunkHeader = make([]MetaData, int(numDocs))
	for i := 0; i < int(numDocs); i++ {
		di.curChunkHeader[i].DocNum, read = binary.Uvarint(s.mem[chunkMetaLoc+offset : chunkMetaLoc+offset+binary.MaxVarintLen64])
		offset += uint64(read)
		di.curChunkHeader[i].DocDvOffset, read = binary.Uvarint(s.mem[chunkMetaLoc+offset : chunkMetaLoc+offset+binary.MaxVarintLen64])
		offset += uint64(read)
	}

	compressedDataLoc := chunkMetaLoc + offset
	dataLength := curChunkEnd - compressedDataLoc
	di.curChunkData = s.mem[compressedDataLoc : compressedDataLoc+dataLength]
	di.curChunkNum = chunkNumber
	di.uncompressed = di.uncompressed[:0]
	return nil
}

func (di *docValueReader) iterateAllDocValues(s *SegmentBase, visitor docNumTermsVisitor) error {
	for i := 0; i < len(di.chunkOffsets); i++ {
		err := di.loadDvChunk(uint64(i), s)
		if err != nil {
			return err
		}

		// uncompress the already loaded data
		uncompressed, err := snappy.Decode(di.uncompressed[:cap(di.uncompressed)], di.curChunkData)
		if err != nil {
			return err
		}
		di.uncompressed = uncompressed

		start := uint64(0)
		for _, entry := range di.curChunkHeader {
			err = visitor(entry.DocNum, uncompressed[start:entry.DocDvOffset])
			if err != nil {
				return err
			}

			start = entry.DocDvOffset
		}
	}

	return nil
}

func (di *docValueReader) visitDocValues(docNum uint64,
	visitor index.DocumentFieldTermVisitor) error {
	// binary search the term locations for the docNum
	start, end := di.getDocValueLocs(docNum)
	if start == math.MaxUint64 || end == math.MaxUint64 {
		return nil
	}

	var uncompressed []byte
	var err error
	// use the uncompressed copy if available
	if len(di.uncompressed) > 0 {
		uncompressed = di.uncompressed
	} else {
		// uncompress the already loaded data
		uncompressed, err = snappy.Decode(di.uncompressed[:cap(di.uncompressed)], di.curChunkData)
		if err != nil {
			return err
		}
		di.uncompressed = uncompressed
	}

	// pick the terms for the given docNum
	uncompressed = uncompressed[start:end]
	for {
		i := bytes.Index(uncompressed, termSeparatorSplitSlice)
		if i < 0 {
			break
		}

		visitor(di.field, uncompressed[0:i])
		uncompressed = uncompressed[i+1:]
	}

	return nil
}

func (di *docValueReader) getDocValueLocs(docNum uint64) (uint64, uint64) {
	i := sort.Search(len(di.curChunkHeader), func(i int) bool {
		return di.curChunkHeader[i].DocNum >= docNum
	})
	if i < len(di.curChunkHeader) && di.curChunkHeader[i].DocNum == docNum {
		return ReadDocValueBoundary(i, di.curChunkHeader)
	}
	return math.MaxUint64, math.MaxUint64
}

// VisitDocumentFieldTerms is an implementation of the
// DocumentFieldTermVisitable interface
func (s *SegmentBase) VisitDocumentFieldTerms(localDocNum uint64, fields []string,
	visitor index.DocumentFieldTermVisitor) error {
	var dvIterClone *docValueReader
	fieldIDPlus1 := uint16(0)
	ok := true
	for _, field := range fields {
		if fieldIDPlus1, ok = s.fieldsMap[field]; !ok {
			continue
		}
		// find the chunkNumber where the docValues are stored
		docInChunk := localDocNum / uint64(s.chunkFactor)

		if dvIter, exists := s.fieldDvReaders[fieldIDPlus1-1]; exists &&
			dvIter != nil {
			dvIterClone = dvIter.cloneInto(dvIterClone)

			// check if the chunk is already loaded
			if docInChunk != dvIterClone.curChunkNumber() {
				err := dvIterClone.loadDvChunk(docInChunk, s)
				if err != nil {
					continue
				}
			}

			_ = dvIterClone.visitDocValues(localDocNum, visitor)
		}
	}
	return nil
}

// VisitableDocValueFields returns the list of fields with
// persisted doc value terms ready to be visitable using the
// VisitDocumentFieldTerms method.
func (s *Segment) VisitableDocValueFields() ([]string, error) {
	rv := make([]string, 0, len(s.fieldDvReaders))
	for fieldID, field := range s.fieldsInv {
		if dvIter, ok := s.fieldDvReaders[uint16(fieldID)]; ok &&
			dvIter != nil {
			rv = append(rv, field)
		}
	}
	return rv, nil
}
<commit_msg>fixes #906 - docvalues handles empty chunk<commit_after>//  Copyright (c) 2017 Couchbase, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 		http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package zap

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"math"
	"reflect"
	"sort"

	"github.com/blevesearch/bleve/index"
	"github.com/blevesearch/bleve/size"
	"github.com/golang/snappy"
)

var reflectStaticSizedocValueReader int

func init() {
	var dvi docValueReader
	reflectStaticSizedocValueReader = int(reflect.TypeOf(dvi).Size())
}

type docNumTermsVisitor func(docNum uint64, terms []byte) error

type docValueReader struct {
	field          string
	curChunkNum    uint64
	chunkOffsets   []uint64
	dvDataLoc      uint64
	curChunkHeader []MetaData
	curChunkData   []byte // compressed data cache
	uncompressed   []byte // temp buf for snappy decompression
}

func (di *docValueReader) size() int {
	return reflectStaticSizedocValueReader + size.SizeOfPtr +
		len(di.field) +
		len(di.chunkOffsets)*size.SizeOfUint64 +
		len(di.curChunkHeader)*reflectStaticSizeMetaData +
		len(di.curChunkData)
}

func (di *docValueReader) cloneInto(rv *docValueReader) *docValueReader {
	if rv == nil {
		rv = &docValueReader{}
	}

	rv.field = di.field
	rv.curChunkNum = math.MaxUint64
	rv.chunkOffsets = di.chunkOffsets // immutable, so it's sharable
	rv.dvDataLoc = di.dvDataLoc
	rv.curChunkHeader = nil
	rv.curChunkData = nil
	rv.uncompressed = nil

	return rv
}

func (di *docValueReader) fieldName() string {
	return di.field
}

func (di *docValueReader) curChunkNumber() uint64 {
	return di.curChunkNum
}

func (s *SegmentBase) loadFieldDocValueReader(field string,
	fieldDvLocStart, fieldDvLocEnd uint64) (*docValueReader, error) {
	// get the docValue offset for the given fields
	if fieldDvLocStart == fieldNotUninverted {
		return nil, fmt.Errorf("loadFieldDocValueReader: "+
			"no docValues found for field: %s", field)
	}

	// read the number of chunks, and chunk offsets position
	var numChunks, chunkOffsetsPosition uint64

	if fieldDvLocEnd-fieldDvLocStart > 16 {
		numChunks = binary.BigEndian.Uint64(s.mem[fieldDvLocEnd-8 : fieldDvLocEnd])
		// read the length of chunk offsets
		chunkOffsetsLen := binary.BigEndian.Uint64(s.mem[fieldDvLocEnd-16 : fieldDvLocEnd-8])
		// acquire position of chunk offsets
		chunkOffsetsPosition = (fieldDvLocEnd - 16) - chunkOffsetsLen
	}

	fdvIter := &docValueReader{
		curChunkNum:  math.MaxUint64,
		field:        field,
		chunkOffsets: make([]uint64, int(numChunks)),
	}

	// read the chunk offsets
	var offset uint64
	for i := 0; i < int(numChunks); i++ {
		loc, read := binary.Uvarint(s.mem[chunkOffsetsPosition+offset : chunkOffsetsPosition+offset+binary.MaxVarintLen64])
		if read <= 0 {
			return nil, fmt.Errorf("corrupted chunk offset during segment load")
		}
		fdvIter.chunkOffsets[i] = loc
		offset += uint64(read)
	}

	// set the data offset
	fdvIter.dvDataLoc = fieldDvLocStart

	return fdvIter, nil
}

func (di *docValueReader) loadDvChunk(chunkNumber uint64, s *SegmentBase) error {
	// advance to the chunk where the docValues
	// reside for the given docNum
	destChunkDataLoc, curChunkEnd := di.dvDataLoc, di.dvDataLoc
	start, end := readChunkBoundary(int(chunkNumber), di.chunkOffsets)
    if start >= end {
        di.curChunkHeader = di.curChunkHeader[:0]
        di.curChunkData = nil
        di.curChunkNum = chunkNumber
        di.uncompressed = di.uncompressed[:0]
        return nil
    }

	destChunkDataLoc += start
	curChunkEnd += end

	// read the number of docs reside in the chunk
	numDocs, read := binary.Uvarint(s.mem[destChunkDataLoc : destChunkDataLoc+binary.MaxVarintLen64])
	if read <= 0 {
		return fmt.Errorf("failed to read the chunk")
	}
	chunkMetaLoc := destChunkDataLoc + uint64(read)

	offset := uint64(0)
	di.curChunkHeader = make([]MetaData, int(numDocs))
	for i := 0; i < int(numDocs); i++ {
		di.curChunkHeader[i].DocNum, read = binary.Uvarint(s.mem[chunkMetaLoc+offset : chunkMetaLoc+offset+binary.MaxVarintLen64])
		offset += uint64(read)
		di.curChunkHeader[i].DocDvOffset, read = binary.Uvarint(s.mem[chunkMetaLoc+offset : chunkMetaLoc+offset+binary.MaxVarintLen64])
		offset += uint64(read)
	}

	compressedDataLoc := chunkMetaLoc + offset
	dataLength := curChunkEnd - compressedDataLoc
	di.curChunkData = s.mem[compressedDataLoc : compressedDataLoc+dataLength]
	di.curChunkNum = chunkNumber
	di.uncompressed = di.uncompressed[:0]
	return nil
}

func (di *docValueReader) iterateAllDocValues(s *SegmentBase, visitor docNumTermsVisitor) error {
	for i := 0; i < len(di.chunkOffsets); i++ {
		err := di.loadDvChunk(uint64(i), s)
		if err != nil {
			return err
		}
        if di.curChunkData == nil || len(di.curChunkHeader) <= 0 {
            continue
        }

		// uncompress the already loaded data
		uncompressed, err := snappy.Decode(di.uncompressed[:cap(di.uncompressed)], di.curChunkData)
		if err != nil {
			return err
		}
		di.uncompressed = uncompressed

		start := uint64(0)
		for _, entry := range di.curChunkHeader {
			err = visitor(entry.DocNum, uncompressed[start:entry.DocDvOffset])
			if err != nil {
				return err
			}

			start = entry.DocDvOffset
		}
	}

	return nil
}

func (di *docValueReader) visitDocValues(docNum uint64,
	visitor index.DocumentFieldTermVisitor) error {
	// binary search the term locations for the docNum
	start, end := di.getDocValueLocs(docNum)
    if start == math.MaxUint64 || end == math.MaxUint64 || start == end {
		return nil
	}

	var uncompressed []byte
	var err error
	// use the uncompressed copy if available
	if len(di.uncompressed) > 0 {
		uncompressed = di.uncompressed
	} else {
		// uncompress the already loaded data
		uncompressed, err = snappy.Decode(di.uncompressed[:cap(di.uncompressed)], di.curChunkData)
		if err != nil {
			return err
		}
		di.uncompressed = uncompressed
	}

	// pick the terms for the given docNum
	uncompressed = uncompressed[start:end]
	for {
		i := bytes.Index(uncompressed, termSeparatorSplitSlice)
		if i < 0 {
			break
		}

		visitor(di.field, uncompressed[0:i])
		uncompressed = uncompressed[i+1:]
	}

	return nil
}

func (di *docValueReader) getDocValueLocs(docNum uint64) (uint64, uint64) {
	i := sort.Search(len(di.curChunkHeader), func(i int) bool {
		return di.curChunkHeader[i].DocNum >= docNum
	})
	if i < len(di.curChunkHeader) && di.curChunkHeader[i].DocNum == docNum {
		return ReadDocValueBoundary(i, di.curChunkHeader)
	}
	return math.MaxUint64, math.MaxUint64
}

// VisitDocumentFieldTerms is an implementation of the
// DocumentFieldTermVisitable interface
func (s *SegmentBase) VisitDocumentFieldTerms(localDocNum uint64, fields []string,
	visitor index.DocumentFieldTermVisitor) error {
	var dvIterClone *docValueReader
	fieldIDPlus1 := uint16(0)
	ok := true
	for _, field := range fields {
		if fieldIDPlus1, ok = s.fieldsMap[field]; !ok {
			continue
		}
		// find the chunkNumber where the docValues are stored
		docInChunk := localDocNum / uint64(s.chunkFactor)

		if dvIter, exists := s.fieldDvReaders[fieldIDPlus1-1]; exists &&
			dvIter != nil {
			dvIterClone = dvIter.cloneInto(dvIterClone)

			// check if the chunk is already loaded
			if docInChunk != dvIterClone.curChunkNumber() {
				err := dvIterClone.loadDvChunk(docInChunk, s)
				if err != nil {
					continue
				}
			}

			_ = dvIterClone.visitDocValues(localDocNum, visitor)
		}
	}
	return nil
}

// VisitableDocValueFields returns the list of fields with
// persisted doc value terms ready to be visitable using the
// VisitDocumentFieldTerms method.
func (s *Segment) VisitableDocValueFields() ([]string, error) {
	rv := make([]string, 0, len(s.fieldDvReaders))
	for fieldID, field := range s.fieldsInv {
		if dvIter, ok := s.fieldDvReaders[uint16(fieldID)]; ok &&
			dvIter != nil {
			rv = append(rv, field)
		}
	}
	return rv, nil
}
<|endoftext|>


<commit_before>package glacier

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/rdwilliamson/aws"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"time"
)

type jobRequest struct {
	Type        string
	ArchiveId   string `json:",omitempty"` // required for archive retrieval
	Description string `json:",omitempty"`
	Format      string `json:",omitempty"`
	SNSTopic    string `json:",omitempty"`
}

type Archive struct {
	ArchiveId          string
	ArchiveDescription string
	CreationDate       time.Time
	Size               uint64
	SHA256TreeHash     string
}

type Inventory struct {
	VaultARN      string
	InventoryDate time.Time
	ArchiveList   []Archive
}

// TODO write unmarshaler instead of using these structs
type archive struct {
	ArchiveId          string
	ArchiveDescription string
	CreationDate       string
	Size               uint64
	SHA256TreeHash     string
}
type inventory struct {
	VaultARN      string
	InventoryDate string
	ArchiveList   []archive
}

type Job struct {
	Action               string
	ArchiveId            string
	ArchiveSizeInBytes   uint64
	Completed            bool
	CompletionDate       time.Time
	CreationDate         time.Time
	InventorySizeInBytes uint
	JobDescription       string
	JobId                string
	SHA256TreeHash       string
	SNSTopic             string
	StatusCode           string
	StatusMessage        string
	VaultARN             string
}

// TODO write unmarshaler instead of using these structs
type job struct {
	Action               string
	ArchiveId            *string
	ArchiveSizeInBytes   *uint64
	Completed            bool
	CompletionDate       *string
	CreationDate         string
	InventorySizeInBytes *uint
	JobDescription       *string
	JobId                string
	SHA256TreeHash       *string
	SNSTopic             *string
	StatusCode           string
	StatusMessage        *string
	VaultARN             string
}
type jobList struct {
	Marker  *string
	JobList []job
}

func (c *Connection) InitiateRetrievalJob(vault, archive, topic,
	description string) (string, error) {
	j := jobRequest{Type: "archive-retrieval", ArchiveId: archive,
		Description: description, SNSTopic: topic}
	rawBody, _ := json.Marshal(j)
	body := bytes.NewReader(rawBody)

	request, err := http.NewRequest("POST",
		"https://"+c.Signature.Region.Glacier+"/-/vaults/"+vault+"/jobs", body)
	if err != nil {
		return "", err
	}
	request.Header.Add("x-amz-glacier-version", "2012-06-01")

	err = c.Signature.Sign(request, body, nil)
	if err != nil {
		return "", err
	}

	response, err := c.Client.Do(request)
	if err != nil {
		return "", err
	}

	if response.StatusCode != 202 {
		body, err := ioutil.ReadAll(response.Body)
		if err != nil {
			return "", err
		}
		err = response.Body.Close()
		if err != nil {
			return "", err
		}
		var e aws.Error
		err = json.Unmarshal(body, &e)
		if err != nil {
			return "", err
		}
		return "", e
	}

	return response.Header.Get("x-amz-job-id"), nil
}

func (c *Connection) InitiateInventoryJob(vault, description,
	topic string) (string, error) {
	j := jobRequest{Type: "inventory-retrieval", Description: description,
		SNSTopic: topic}
	rawBody, err := json.Marshal(j)
	if err != nil {
		return "", err
	}
	body := bytes.NewReader(rawBody)

	request, err := http.NewRequest("POST",
		"https://"+c.Signature.Region.Glacier+"/-/vaults/"+vault+"/jobs", body)
	if err != nil {
		return "", err
	}
	request.Header.Add("x-amz-glacier-version", "2012-06-01")

	err = c.Signature.Sign(request, body, nil)
	if err != nil {
		return "", err
	}

	response, err := c.Client.Do(request)
	if err != nil {
		return "", err
	}

	if response.StatusCode != 202 {
		body, err := ioutil.ReadAll(response.Body)
		if err != nil {
			return "", err
		}
		err = response.Body.Close()
		if err != nil {
			return "", err
		}
		var e aws.Error
		err = json.Unmarshal(body, &e)
		if err != nil {
			return "", err
		}
		return "", e
	}

	return response.Header.Get("x-amz-job-id"), nil
}

func (c *Connection) DescribeJob() error {
	return nil
}

func (c *Connection) GetRetrievalJob(vault, job string, start,
	end uint) (io.ReadCloser, error) {
	request, err := http.NewRequest("GET", "https://"+
		c.Signature.Region.Glacier+"/-/vaults/"+vault+"/jobs/"+job+"/output",
		nil)
	if err != nil {
		return nil, err
	}
	request.Header.Add("x-amz-glacier-version", "2012-06-01")

	if end > 0 {
		request.Header.Add("Range", fmt.Sprintf("bytes %d-%d/*", start, end))
	}

	response, err := c.Client.Do(request)
	if err != nil {
		return nil, err
	}

	if response.StatusCode != 200 {
		body, err := ioutil.ReadAll(response.Body)
		if err != nil {
			return nil, err
		}
		err = response.Body.Close()
		if err != nil {
			return nil, err
		}
		var e aws.Error
		err = json.Unmarshal(body, &e)
		if err != nil {
			return nil, err
		}
		return nil, e
	}

	// TODO return content range and x-amz-sha256-tree-hash
	return nil, err
}

func (c *Connection) GetInventoryJob(vault, job string) (Inventory, error) {
	request, err := http.NewRequest("GET", "https://"+
		c.Signature.Region.Glacier+"/-/vaults/"+vault+"/jobs/"+job+"/output",
		nil)
	if err != nil {
		return Inventory{}, err
	}
	request.Header.Add("x-amz-glacier-version", "2012-06-01")

	err = c.Signature.Sign(request, nil, nil)
	if err != nil {
		return Inventory{}, err
	}

	response, err := c.Client.Do(request)
	if err != nil {
		return Inventory{}, err
	}

	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return Inventory{}, err
	}
	err = response.Body.Close()
	if err != nil {
		return Inventory{}, err
	}

	if response.StatusCode != 200 {
		var e aws.Error
		err = json.Unmarshal(body, &e)
		if err != nil {
			return Inventory{}, err
		}
		return Inventory{}, e
	}

	var i inventory
	err = json.Unmarshal(body, &i)
	if err != nil {
		return Inventory{}, err
	}

	var result Inventory
	result.VaultARN = i.VaultARN
	result.InventoryDate, err = time.Parse(time.RFC3339, i.InventoryDate)
	if err != nil {
		return Inventory{}, err
	}
	result.ArchiveList = make([]Archive, len(i.ArchiveList))
	for j, v := range i.ArchiveList {
		result.ArchiveList[j].ArchiveId = v.ArchiveId
		result.ArchiveList[j].ArchiveDescription = v.ArchiveDescription
		result.ArchiveList[j].CreationDate, err = time.Parse(time.RFC3339,
			v.CreationDate)
		if err != nil {
			return Inventory{}, err
		}
		result.ArchiveList[j].Size = v.Size
		result.ArchiveList[j].SHA256TreeHash = v.SHA256TreeHash
	}

	return result, nil
}

func (c *Connection) ListJobs(vault, completed, limit, marker,
	statusCode string) ([]Job, string, error) {
	get, err := url.Parse("https://" + c.Signature.Region.Glacier +
		"/-/vaults/" + vault + "/jobs")
	if err != nil {
		return nil, "", err
	}

	query := get.Query()
	if completed != "" {
		// TODO validate, true or false
		query.Add("completed", completed)
	}
	if limit != "" {
		// TODO validate, 1 - 1000
		query.Add("limit", limit)
	}
	if marker != "" {
		// TODO validate
		query.Add("marker", marker)
	}
	if statusCode != "" {
		// TODO validate, InProgress, Succeeded, or Failed
		query.Add("statuscode", statusCode)
	}
	get.RawQuery = query.Encode()

	request, err := http.NewRequest("GET", get.String(), nil)
	if err != nil {
		return nil, "", err
	}
	request.Header.Add("x-amz-glacier-version", "2012-06-01")

	c.Signature.Sign(request, nil, nil)

	response, err := c.Client.Do(request)
	if err != nil {
		return nil, "", err
	}

	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, "", err
	}
	err = response.Body.Close()
	if err != nil {
		return nil, "", err
	}

	if response.StatusCode != 200 {
		var e aws.Error
		err = json.Unmarshal(body, &e)
		if err != nil {
			return nil, "", err
		}
		return nil, "", e
	}

	var jl jobList
	err = json.Unmarshal(body, &jl)
	if err != nil {
		return nil, "", err
	}

	var err1 error
	jobs := make([]Job, len(jl.JobList))
	for i, v := range jl.JobList {
		jobs[i].Action = v.Action
		if v.ArchiveId != nil {
			jobs[i].ArchiveId = *v.ArchiveId
		}
		if v.ArchiveSizeInBytes != nil {
			jobs[i].ArchiveSizeInBytes = *v.ArchiveSizeInBytes
		}
		jobs[i].Completed = v.Completed
		if v.CompletionDate != nil {
			jobs[i].CompletionDate, err = time.Parse(time.RFC3339, *v.CompletionDate)
		}
		if err != nil {
			err1 = err
		}
		jobs[i].CreationDate, err = time.Parse(time.RFC3339, v.CreationDate)
		if err != nil {
			err1 = err
		}
		if v.InventorySizeInBytes != nil {
			jobs[i].InventorySizeInBytes = *v.InventorySizeInBytes
		}
		if v.JobDescription != nil {
			jobs[i].JobDescription = *v.JobDescription
		}
		jobs[i].JobId = v.JobId
		if v.SHA256TreeHash != nil {
			jobs[i].SHA256TreeHash = *v.SHA256TreeHash
		}
		if v.SNSTopic != nil {
			jobs[i].SNSTopic = *v.SNSTopic
		}
		jobs[i].StatusCode = v.StatusCode
		if v.StatusMessage != nil {
			jobs[i].StatusMessage = *v.StatusMessage
		}
		jobs[i].VaultARN = v.VaultARN
	}

	var m string
	if jl.Marker != nil {
		m = *jl.Marker
	}

	return jobs, m, err1
}
<commit_msg>Added retrieval job.<commit_after>package glacier

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/rdwilliamson/aws"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"time"
)

type jobRequest struct {
	Type        string
	ArchiveId   string `json:",omitempty"` // required for archive retrieval
	Description string `json:",omitempty"`
	Format      string `json:",omitempty"`
	SNSTopic    string `json:",omitempty"`
}

type Archive struct {
	ArchiveId          string
	ArchiveDescription string
	CreationDate       time.Time
	Size               uint64
	SHA256TreeHash     string
}

type Inventory struct {
	VaultARN      string
	InventoryDate time.Time
	ArchiveList   []Archive
}

// TODO write unmarshaler instead of using these structs
type archive struct {
	ArchiveId          string
	ArchiveDescription string
	CreationDate       string
	Size               uint64
	SHA256TreeHash     string
}
type inventory struct {
	VaultARN      string
	InventoryDate string
	ArchiveList   []archive
}

type Job struct {
	Action               string
	ArchiveId            string
	ArchiveSizeInBytes   uint64
	Completed            bool
	CompletionDate       time.Time
	CreationDate         time.Time
	InventorySizeInBytes uint
	JobDescription       string
	JobId                string
	SHA256TreeHash       string
	SNSTopic             string
	StatusCode           string
	StatusMessage        string
	VaultARN             string
}

// TODO write unmarshaler instead of using these structs
type job struct {
	Action               string
	ArchiveId            *string
	ArchiveSizeInBytes   *uint64
	Completed            bool
	CompletionDate       *string
	CreationDate         string
	InventorySizeInBytes *uint
	JobDescription       *string
	JobId                string
	SHA256TreeHash       *string
	SNSTopic             *string
	StatusCode           string
	StatusMessage        *string
	VaultARN             string
}
type jobList struct {
	Marker  *string
	JobList []job
}

func (c *Connection) InitiateRetrievalJob(vault, archive, topic,
	description string) (string, error) {
	j := jobRequest{Type: "archive-retrieval", ArchiveId: archive,
		Description: description, SNSTopic: topic}
	rawBody, _ := json.Marshal(j)
	body := bytes.NewReader(rawBody)

	request, err := http.NewRequest("POST",
		"https://"+c.Signature.Region.Glacier+"/-/vaults/"+vault+"/jobs", body)
	if err != nil {
		return "", err
	}
	request.Header.Add("x-amz-glacier-version", "2012-06-01")

	err = c.Signature.Sign(request, body, nil)
	if err != nil {
		return "", err
	}

	response, err := c.Client.Do(request)
	if err != nil {
		return "", err
	}

	if response.StatusCode != 202 {
		body, err := ioutil.ReadAll(response.Body)
		if err != nil {
			return "", err
		}
		err = response.Body.Close()
		if err != nil {
			return "", err
		}
		var e aws.Error
		err = json.Unmarshal(body, &e)
		if err != nil {
			return "", err
		}
		return "", e
	}

	return response.Header.Get("x-amz-job-id"), nil
}

func (c *Connection) InitiateInventoryJob(vault, description,
	topic string) (string, error) {
	j := jobRequest{Type: "inventory-retrieval", Description: description,
		SNSTopic: topic}
	rawBody, err := json.Marshal(j)
	if err != nil {
		return "", err
	}
	body := bytes.NewReader(rawBody)

	request, err := http.NewRequest("POST",
		"https://"+c.Signature.Region.Glacier+"/-/vaults/"+vault+"/jobs", body)
	if err != nil {
		return "", err
	}
	request.Header.Add("x-amz-glacier-version", "2012-06-01")

	err = c.Signature.Sign(request, body, nil)
	if err != nil {
		return "", err
	}

	response, err := c.Client.Do(request)
	if err != nil {
		return "", err
	}

	if response.StatusCode != 202 {
		body, err := ioutil.ReadAll(response.Body)
		if err != nil {
			return "", err
		}
		err = response.Body.Close()
		if err != nil {
			return "", err
		}
		var e aws.Error
		err = json.Unmarshal(body, &e)
		if err != nil {
			return "", err
		}
		return "", e
	}

	return response.Header.Get("x-amz-job-id"), nil
}

func (c *Connection) DescribeJob() error {
	return nil
}

func (c *Connection) GetRetrievalJob(vault, job string, start, end uint) (io.ReadCloser, error) {
	request, err := http.NewRequest("GET", "https://"+
		c.Signature.Region.Glacier+"/-/vaults/"+vault+"/jobs/"+job+"/output",
		nil)
	if err != nil {
		return nil, err
	}
	request.Header.Add("x-amz-glacier-version", "2012-06-01")
	if end > 0 {
		request.Header.Add("Range", fmt.Sprintf("bytes %d-%d/*", start, end))
	}

	c.Signature.Sign(request, nil, nil)

	response, err := c.Client.Do(request)
	if err != nil {
		return nil, err
	}

	if response.StatusCode != 200 {
		body, err := ioutil.ReadAll(response.Body)
		if err != nil {
			return nil, err
		}
		response.Body.Close()

		var e aws.Error
		err = json.Unmarshal(body, &e)
		if err != nil {
			return nil, err
		}
		return nil, e
	}

	// TODO return content range and x-amz-sha256-tree-hash
	return response.Body, nil
}

func (c *Connection) GetInventoryJob(vault, job string) (Inventory, error) {
	request, err := http.NewRequest("GET", "https://"+
		c.Signature.Region.Glacier+"/-/vaults/"+vault+"/jobs/"+job+"/output",
		nil)
	if err != nil {
		return Inventory{}, err
	}
	request.Header.Add("x-amz-glacier-version", "2012-06-01")

	err = c.Signature.Sign(request, nil, nil)
	if err != nil {
		return Inventory{}, err
	}

	response, err := c.Client.Do(request)
	if err != nil {
		return Inventory{}, err
	}

	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return Inventory{}, err
	}
	err = response.Body.Close()
	if err != nil {
		return Inventory{}, err
	}

	if response.StatusCode != 200 {
		var e aws.Error
		err = json.Unmarshal(body, &e)
		if err != nil {
			return Inventory{}, err
		}
		return Inventory{}, e
	}

	var i inventory
	err = json.Unmarshal(body, &i)
	if err != nil {
		return Inventory{}, err
	}

	var result Inventory
	result.VaultARN = i.VaultARN
	result.InventoryDate, err = time.Parse(time.RFC3339, i.InventoryDate)
	if err != nil {
		return Inventory{}, err
	}
	result.ArchiveList = make([]Archive, len(i.ArchiveList))
	for j, v := range i.ArchiveList {
		result.ArchiveList[j].ArchiveId = v.ArchiveId
		result.ArchiveList[j].ArchiveDescription = v.ArchiveDescription
		result.ArchiveList[j].CreationDate, err = time.Parse(time.RFC3339,
			v.CreationDate)
		if err != nil {
			return Inventory{}, err
		}
		result.ArchiveList[j].Size = v.Size
		result.ArchiveList[j].SHA256TreeHash = v.SHA256TreeHash
	}

	return result, nil
}

func (c *Connection) ListJobs(vault, completed, limit, marker,
	statusCode string) ([]Job, string, error) {
	get, err := url.Parse("https://" + c.Signature.Region.Glacier +
		"/-/vaults/" + vault + "/jobs")
	if err != nil {
		return nil, "", err
	}

	query := get.Query()
	if completed != "" {
		// TODO validate, true or false
		query.Add("completed", completed)
	}
	if limit != "" {
		// TODO validate, 1 - 1000
		query.Add("limit", limit)
	}
	if marker != "" {
		// TODO validate
		query.Add("marker", marker)
	}
	if statusCode != "" {
		// TODO validate, InProgress, Succeeded, or Failed
		query.Add("statuscode", statusCode)
	}
	get.RawQuery = query.Encode()

	request, err := http.NewRequest("GET", get.String(), nil)
	if err != nil {
		return nil, "", err
	}
	request.Header.Add("x-amz-glacier-version", "2012-06-01")

	c.Signature.Sign(request, nil, nil)

	response, err := c.Client.Do(request)
	if err != nil {
		return nil, "", err
	}

	body, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, "", err
	}
	err = response.Body.Close()
	if err != nil {
		return nil, "", err
	}

	if response.StatusCode != 200 {
		var e aws.Error
		err = json.Unmarshal(body, &e)
		if err != nil {
			return nil, "", err
		}
		return nil, "", e
	}

	var jl jobList
	err = json.Unmarshal(body, &jl)
	if err != nil {
		return nil, "", err
	}

	var err1 error
	jobs := make([]Job, len(jl.JobList))
	for i, v := range jl.JobList {
		jobs[i].Action = v.Action
		if v.ArchiveId != nil {
			jobs[i].ArchiveId = *v.ArchiveId
		}
		if v.ArchiveSizeInBytes != nil {
			jobs[i].ArchiveSizeInBytes = *v.ArchiveSizeInBytes
		}
		jobs[i].Completed = v.Completed
		if v.CompletionDate != nil {
			jobs[i].CompletionDate, err = time.Parse(time.RFC3339, *v.CompletionDate)
		}
		if err != nil {
			err1 = err
		}
		jobs[i].CreationDate, err = time.Parse(time.RFC3339, v.CreationDate)
		if err != nil {
			err1 = err
		}
		if v.InventorySizeInBytes != nil {
			jobs[i].InventorySizeInBytes = *v.InventorySizeInBytes
		}
		if v.JobDescription != nil {
			jobs[i].JobDescription = *v.JobDescription
		}
		jobs[i].JobId = v.JobId
		if v.SHA256TreeHash != nil {
			jobs[i].SHA256TreeHash = *v.SHA256TreeHash
		}
		if v.SNSTopic != nil {
			jobs[i].SNSTopic = *v.SNSTopic
		}
		jobs[i].StatusCode = v.StatusCode
		if v.StatusMessage != nil {
			jobs[i].StatusMessage = *v.StatusMessage
		}
		jobs[i].VaultARN = v.VaultARN
	}

	var m string
	if jl.Marker != nil {
		m = *jl.Marker
	}

	return jobs, m, err1
}
<|endoftext|>