<commit_before>/*
Copyright 2021 The Skaffold 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 configlocations

import (
	"context"
	"path"
	"reflect"
	"strconv"
	"strings"

	kyaml "sigs.k8s.io/kustomize/kyaml/yaml"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/output/log"
	sErrors "github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/errors"
	latestV1 "github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/latest/v1"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
)

type YAMLInfo struct {
	RNode      *kyaml.RNode
	SourceFile string
}

type Location struct {
	SourceFile  string
	StartLine   int
	StartColumn int
	EndLine     int
	EndColumn   int
}

type YAMLInfos struct {
	yamlInfos               map[uintptr]map[string]YAMLInfo
	FieldsOverrodeByProfile map[string]YAMLOverrideInfo // map of schema path -> profile name -- ex: /artifacts/0/image -> "overwrite-artifacte-image-profile"
}

func (m *YAMLInfos) GetYamlInfosCopy() map[uintptr]map[string]YAMLInfo {
	yamlInfos := map[uintptr]map[string]YAMLInfo{}
	for ptr, mp := range m.yamlInfos {
		tmpmp := map[string]YAMLInfo{}
		for k, v := range mp {
			tmpmp[k] = YAMLInfo{
				RNode:      v.RNode.Copy(),
				SourceFile: v.SourceFile,
			}
		}
		yamlInfos[ptr] = tmpmp
	}
	return yamlInfos
}

func MissingLocation() *Location {
	return &Location{
		SourceFile:  "",
		StartLine:   -1,
		StartColumn: -1,
		EndLine:     -1,
		EndColumn:   -1,
	}
}

func NewYAMLInfos() *YAMLInfos {
	return &YAMLInfos{
		yamlInfos: map[uintptr]map[string]YAMLInfo{},
	}
}

type YAMLOverrideInfo struct {
	ProfileName    string
	PatchIndex     int
	PatchOperation string
	PatchCopyFrom  string
}

// Parse parses a skaffold config entry collecting file location information for each schema config object
func Parse(sourceFile string, config *latestV1.SkaffoldConfig, fieldsOverrodeByProfile map[string]YAMLOverrideInfo) (*YAMLInfos, error) {
	yamlInfos, err := buildMapOfSchemaObjPointerToYAMLInfos(sourceFile, config, map[uintptr]map[string]YAMLInfo{}, fieldsOverrodeByProfile)
	return &YAMLInfos{
			yamlInfos:               yamlInfos,
			FieldsOverrodeByProfile: fieldsOverrodeByProfile,
		},
		err
}

// Locate gets the location for a skaffold schema struct pointer
func (m *YAMLInfos) Locate(obj interface{}) *Location {
	return m.locate(obj, "")
}

// Locate gets the location for a skaffold schema struct pointer
func (m *YAMLInfos) LocateElement(obj interface{}, idx int) *Location {
	return m.locate(obj, strconv.Itoa(idx))
}

// Locate gets the location for a skaffold schema struct pointer
func (m *YAMLInfos) LocateField(obj interface{}, fieldName string) *Location {
	return m.locate(obj, fieldName)
}

// Locate gets the location for a skaffold schema struct pointer
func (m *YAMLInfos) LocateByPointer(ptr uintptr) *Location {
	if m == nil {
		log.Entry(context.TODO()).Infof("YamlInfos is nil, unable to complete call to LocateByPointer for pointer: %d", ptr)
		return MissingLocation()
	}
	if _, ok := m.yamlInfos[ptr]; !ok {
		log.Entry(context.TODO()).Infof("no map entry found when attempting LocateByPointer for pointer: %d", ptr)
		return MissingLocation()
	}
	node, ok := m.yamlInfos[ptr][""]
	if !ok {
		log.Entry(context.TODO()).Infof("no map entry found when attempting LocateByPointer for pointer: %d", ptr)
		return MissingLocation()
	}
	// iterate over kyaml.RNode text to get endline and endcolumn information
	nodeText, err := node.RNode.String()
	if err != nil {
		return MissingLocation()
	}
	log.Entry(context.TODO()).Infof("map entry found when executing LocateByPointer for pointer: %d", ptr)
	lines, cols := getLinesAndColsOfString(nodeText)

	// TODO(aaron-prindle) all line & col values seem 1 greater than expected in actual use, will need to check to see how it works with IDE
	return &Location{
		SourceFile:  node.SourceFile,
		StartLine:   node.RNode.Document().Line,
		StartColumn: node.RNode.Document().Column,
		EndLine:     node.RNode.Document().Line + lines,
		EndColumn:   cols,
	}
}

func (m *YAMLInfos) locate(obj interface{}, key string) *Location {
	if m == nil {
		log.Entry(context.TODO()).Infof("YamlInfos is nil, unable to complete call to locate with params: %v of type %T", obj, obj)
		return MissingLocation()
	}
	v := reflect.ValueOf(obj)
	if v.Kind() != reflect.Ptr {
		log.Entry(context.TODO()).Infof("non pointer object passed to locate: %v of type %T", obj, obj)
		return MissingLocation()
	}
	if _, ok := m.yamlInfos[v.Pointer()]; !ok {
		log.Entry(context.TODO()).Infof("no map entry found when attempting locate for %v of type %T and pointer: %d", obj, obj, v.Pointer())
		return MissingLocation()
	}
	node, ok := m.yamlInfos[v.Pointer()][key]
	if !ok {
		log.Entry(context.TODO()).Infof("no map entry found when attempting locate for %v of type %T and pointer: %d", obj, obj, v.Pointer())
		return MissingLocation()
	}
	// iterate over kyaml.RNode text to get endline and endcolumn information
	nodeText, err := node.RNode.String()
	if err != nil {
		return MissingLocation()
	}
	log.Entry(context.TODO()).Infof("map entry found when executing locate for %v of type %T and pointer: %d", obj, obj, v.Pointer())
	lines, cols := getLinesAndColsOfString(nodeText)

	// TODO(aaron-prindle) all line & col values seem 1 greater than expected in actual use, will need to check to see how it works with IDE
	return &Location{
		SourceFile:  node.SourceFile,
		StartLine:   node.RNode.Document().Line,
		StartColumn: node.RNode.Document().Column,
		EndLine:     node.RNode.Document().Line + lines,
		EndColumn:   cols,
	}
}

func getLinesAndColsOfString(str string) (int, int) {
	line := 0
	col := 0
	for i := range str {
		col++
		if str[i] == '\n' {
			line++
			col = 0
		}
	}
	return line, col
}

func buildMapOfSchemaObjPointerToYAMLInfos(sourceFile string, config *latestV1.SkaffoldConfig, yamlInfos map[uintptr]map[string]YAMLInfo,
	fieldsOverrodeByProfile map[string]YAMLOverrideInfo) (map[uintptr]map[string]YAMLInfo, error) {
	skaffoldConfigText, err := util.ReadConfiguration(sourceFile)
	if err != nil {
		return nil, sErrors.ConfigParsingError(err)
	}
	root, err := kyaml.Parse(string(skaffoldConfigText))
	if err != nil {
		return nil, err
	}
	// TODO(aaron-prindle) perhaps add some defensive logic to recover from panic when using reflection and instead return error?
	return generateObjPointerToYAMLNodeMap(sourceFile, reflect.ValueOf(config), reflect.ValueOf(nil), "", "", []string{},
		root, root, -1, fieldsOverrodeByProfile, map[interface{}]bool{}, yamlInfos, false)
}

// generateObjPointerToYAMLNodeMap recursively walks through a structs fields (taking into account profile and patch profile overrides)
// and collects the corresponding yaml node for each field
func generateObjPointerToYAMLNodeMap(sourceFile string, v reflect.Value, parentV reflect.Value, fieldName, yamlTag string, schemaPath []string,
	rootRNode *kyaml.RNode, rNode *kyaml.RNode, containerIdx int, fieldPathsOverrodeByProfiles map[string]YAMLOverrideInfo,
	visited map[interface{}]bool, yamlInfos map[uintptr]map[string]YAMLInfo, isPatchProfileElemOverride bool) (map[uintptr]map[string]YAMLInfo, error) {
	// TODO(aaron-prindle) need to verify if generateObjPointerToYAMLNodeMap adds entries for 'map' types, luckily the skaffold schema
	// only has map[string]string and they are leaf nodes as well which this should work fine for doing the recursion for the time being
	var err error

	// add current obj/field to schema path if criteria met
	switch {
	case containerIdx >= 0:
		schemaPath = append(schemaPath, strconv.Itoa(containerIdx))
	case yamlTag != "":
		schemaPath = append(schemaPath, yamlTag)
	}
	// check if current obj/field was overridden by a profile
	if yamlOverrideInfo, ok := fieldPathsOverrodeByProfiles["/"+path.Join(schemaPath...)]; ok {
		// reset yaml node path from root path to given profile path ("/" -> "/profile/name=profileName/etc...")
		rNode, err = rootRNode.Pipe(kyaml.Lookup("profiles"), kyaml.MatchElementList([]string{"name"}, []string{yamlOverrideInfo.ProfileName}))
		if err != nil {
			return nil, err
		}
		switch {
		case yamlOverrideInfo.PatchIndex < 0: // this schema obj/field has a profile override (NOT a patch profile override)
			// moves parent node path from being rooted at default yaml '/' to being rooted at '/profile/name=profileName/...'
			for i := 0; i < len(schemaPath)-1; i++ {
				rNode, err = rNode.Pipe(kyaml.Lookup(schemaPath[i]))
				if err != nil {
					return nil, err
				}
			}
		default: // this schema obj/field has a patch profile override
			// NOTE: 'remove' patch operations are not included in fieldPathsOverrodeByProfiles as there
			//  is no work to be done on them (they were already removed from the schema)

			// TODO(aaron-prindle) verify UX makes sense to use the "FROM" copy node to get yaml information from
			if yamlOverrideInfo.PatchOperation == "copy" {
				fromPath := strings.Split(yamlOverrideInfo.PatchCopyFrom, "/")
				var kf kyaml.Filter
				for i := 0; i < len(fromPath)-1; i++ {
					if pathNum, err := strconv.Atoi(fromPath[i]); err == nil {
						// this path element is a number
						kf = kyaml.ElementIndexer{Index: pathNum}
					} else {
						// this path element isn't a number
						kf = kyaml.Lookup(fromPath[i])
					}
					rNode, err = rNode.Pipe(kf)
					if err != nil {
						return nil, err
					}
				}
			} else {
				rNode, err = rNode.Pipe(kyaml.Lookup("patches"), kyaml.ElementIndexer{Index: yamlOverrideInfo.PatchIndex})
				if err != nil {
					return nil, err
				}
				yamlTag = "value"
			}
			isPatchProfileElemOverride = true
		}
	}
	if rNode == nil {
		return yamlInfos, nil
	}

	// drill down through pointers and interfaces to get a value we can use
	for v.Kind() == reflect.Ptr || v.Kind() == reflect.Interface {
		if v.Kind() == reflect.Ptr {
			// Check for recursive data
			if visited[v.Interface()] {
				return yamlInfos, nil
			}
			visited[v.Interface()] = true
		}
		v = v.Elem()
	}

	if yamlTag != "" { // check that struct is not `yaml:",inline"`
		// traverse kyaml node tree to current obj/field location
		var kf kyaml.Filter
		switch {
		case rNode.YNode().Kind == kyaml.SequenceNode:
			kf = kyaml.ElementIndexer{Index: containerIdx}
		default:
			kf = kyaml.Lookup(yamlTag)
		}
		rNode, err = rNode.Pipe(kf)
		if err != nil {
			return nil, err
		}
		if rNode == nil {
			return yamlInfos, nil
		}

		// this case is so that the line #'s of primitive values can be "located" as they are not addressable but we can
		// map the parent address and put the child field in second map
		if parentV.CanAddr() {
			if _, ok := yamlInfos[parentV.Addr().Pointer()]; !ok {
				yamlInfos[parentV.Addr().Pointer()] = map[string]YAMLInfo{}
			}
			// add parent relationship entry to yaml info map
			if containerIdx >= 0 {
				yamlInfos[parentV.Addr().Pointer()][strconv.Itoa(containerIdx)] = YAMLInfo{
					RNode:      rNode,
					SourceFile: sourceFile,
				}
			} else {
				yamlInfos[parentV.Addr().Pointer()][fieldName] = YAMLInfo{
					RNode:      rNode,
					SourceFile: sourceFile,
				}
			}
		}
	}

	if v.CanAddr() {
		if _, ok := yamlInfos[v.Addr().Pointer()]; !ok {
			yamlInfos[v.Addr().Pointer()] = map[string]YAMLInfo{}
		}
		// add current node entry to yaml info map
		yamlInfos[v.Addr().Pointer()][""] = YAMLInfo{
			RNode:      rNode,
			SourceFile: sourceFile,
		}
	}

	switch v.Kind() {
	// TODO(aaron-prindle) add reflect.Map support here as well, currently no struct fields have nested struct in map field so ok for now
	case reflect.Slice, reflect.Array:
		for i := 0; i < v.Len(); i++ {
			generateObjPointerToYAMLNodeMap(sourceFile, v.Index(i), v, fieldName+"["+strconv.Itoa(i)+"]", yamlTag+"["+strconv.Itoa(i)+"]", schemaPath,
				rootRNode, rNode, i, fieldPathsOverrodeByProfiles, visited, yamlInfos, isPatchProfileElemOverride)
		}
	case reflect.Struct:
		t := v.Type() // use type to get number and names of fields
		for i := 0; i < t.NumField(); i++ {
			field := t.Field(i)
			// TODO(aaron-prindle) verify this value works for structs that are `yaml:",inline"`
			newYamlTag := field.Name
			if yamlTagToken := field.Tag.Get("yaml"); yamlTagToken != "" && yamlTagToken != "-" {
				// check for possible comma as in "...,omitempty"
				var commaIdx int
				if commaIdx = strings.Index(yamlTagToken, ","); commaIdx < 0 {
					commaIdx = len(yamlTagToken)
				}
				newYamlTag = yamlTagToken[:commaIdx]
			}
			generateObjPointerToYAMLNodeMap(sourceFile, v.Field(i), v, field.Name, newYamlTag, schemaPath, rootRNode, rNode, -1,
				fieldPathsOverrodeByProfiles, visited, yamlInfos, isPatchProfileElemOverride)
		}
	}
	return yamlInfos, nil
}
<commit_msg>fix: add panic fix and recovery logic to reflection for yaml line number info (#7276)<commit_after>/*
Copyright 2021 The Skaffold 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 configlocations

import (
	"context"
	"path"
	"reflect"
	"strconv"
	"strings"

	kyaml "sigs.k8s.io/kustomize/kyaml/yaml"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/output/log"
	sErrors "github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/errors"
	latestV1 "github.com/GoogleContainerTools/skaffold/pkg/skaffold/schema/latest/v1"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/util"
)

type YAMLInfo struct {
	RNode      *kyaml.RNode
	SourceFile string
}

type Location struct {
	SourceFile  string
	StartLine   int
	StartColumn int
	EndLine     int
	EndColumn   int
}

type YAMLInfos struct {
	yamlInfos               map[uintptr]map[string]YAMLInfo
	FieldsOverrodeByProfile map[string]YAMLOverrideInfo // map of schema path -> profile name -- ex: /artifacts/0/image -> "overwrite-artifacte-image-profile"
}

func (m *YAMLInfos) GetYamlInfosCopy() map[uintptr]map[string]YAMLInfo {
	yamlInfos := map[uintptr]map[string]YAMLInfo{}
	for ptr, mp := range m.yamlInfos {
		tmpmp := map[string]YAMLInfo{}
		for k, v := range mp {
			tmpmp[k] = YAMLInfo{
				RNode:      v.RNode.Copy(),
				SourceFile: v.SourceFile,
			}
		}
		yamlInfos[ptr] = tmpmp
	}
	return yamlInfos
}

func MissingLocation() *Location {
	return &Location{
		SourceFile:  "",
		StartLine:   -1,
		StartColumn: -1,
		EndLine:     -1,
		EndColumn:   -1,
	}
}

func NewYAMLInfos() *YAMLInfos {
	return &YAMLInfos{
		yamlInfos: map[uintptr]map[string]YAMLInfo{},
	}
}

type YAMLOverrideInfo struct {
	ProfileName    string
	PatchIndex     int
	PatchOperation string
	PatchCopyFrom  string
}

// Parse parses a skaffold config entry collecting file location information for each schema config object
func Parse(sourceFile string, config *latestV1.SkaffoldConfig, fieldsOverrodeByProfile map[string]YAMLOverrideInfo) (*YAMLInfos, error) {
	yamlInfos, err := buildMapOfSchemaObjPointerToYAMLInfos(sourceFile, config, map[uintptr]map[string]YAMLInfo{}, fieldsOverrodeByProfile)
	return &YAMLInfos{
			yamlInfos:               yamlInfos,
			FieldsOverrodeByProfile: fieldsOverrodeByProfile,
		},
		err
}

// Locate gets the location for a skaffold schema struct pointer
func (m *YAMLInfos) Locate(obj interface{}) *Location {
	return m.locate(obj, "")
}

// Locate gets the location for a skaffold schema struct pointer
func (m *YAMLInfos) LocateElement(obj interface{}, idx int) *Location {
	return m.locate(obj, strconv.Itoa(idx))
}

// Locate gets the location for a skaffold schema struct pointer
func (m *YAMLInfos) LocateField(obj interface{}, fieldName string) *Location {
	return m.locate(obj, fieldName)
}

// Locate gets the location for a skaffold schema struct pointer
func (m *YAMLInfos) LocateByPointer(ptr uintptr) *Location {
	if m == nil {
		log.Entry(context.TODO()).Infof("YamlInfos is nil, unable to complete call to LocateByPointer for pointer: %d", ptr)
		return MissingLocation()
	}
	if _, ok := m.yamlInfos[ptr]; !ok {
		log.Entry(context.TODO()).Infof("no map entry found when attempting LocateByPointer for pointer: %d", ptr)
		return MissingLocation()
	}
	node, ok := m.yamlInfos[ptr][""]
	if !ok {
		log.Entry(context.TODO()).Infof("no map entry found when attempting LocateByPointer for pointer: %d", ptr)
		return MissingLocation()
	}
	// iterate over kyaml.RNode text to get endline and endcolumn information
	nodeText, err := node.RNode.String()
	if err != nil {
		return MissingLocation()
	}
	log.Entry(context.TODO()).Infof("map entry found when executing LocateByPointer for pointer: %d", ptr)
	lines, cols := getLinesAndColsOfString(nodeText)

	// TODO(aaron-prindle) all line & col values seem 1 greater than expected in actual use, will need to check to see how it works with IDE
	return &Location{
		SourceFile:  node.SourceFile,
		StartLine:   node.RNode.Document().Line,
		StartColumn: node.RNode.Document().Column,
		EndLine:     node.RNode.Document().Line + lines,
		EndColumn:   cols,
	}
}

func (m *YAMLInfos) locate(obj interface{}, key string) *Location {
	if m == nil {
		log.Entry(context.TODO()).Infof("YamlInfos is nil, unable to complete call to locate with params: %v of type %T", obj, obj)
		return MissingLocation()
	}
	v := reflect.ValueOf(obj)
	if v.Kind() != reflect.Ptr {
		log.Entry(context.TODO()).Infof("non pointer object passed to locate: %v of type %T", obj, obj)
		return MissingLocation()
	}
	if _, ok := m.yamlInfos[v.Pointer()]; !ok {
		log.Entry(context.TODO()).Infof("no map entry found when attempting locate for %v of type %T and pointer: %d", obj, obj, v.Pointer())
		return MissingLocation()
	}
	node, ok := m.yamlInfos[v.Pointer()][key]
	if !ok {
		log.Entry(context.TODO()).Infof("no map entry found when attempting locate for %v of type %T and pointer: %d", obj, obj, v.Pointer())
		return MissingLocation()
	}
	// iterate over kyaml.RNode text to get endline and endcolumn information
	nodeText, err := node.RNode.String()
	if err != nil {
		return MissingLocation()
	}
	log.Entry(context.TODO()).Infof("map entry found when executing locate for %v of type %T and pointer: %d", obj, obj, v.Pointer())
	lines, cols := getLinesAndColsOfString(nodeText)

	// TODO(aaron-prindle) all line & col values seem 1 greater than expected in actual use, will need to check to see how it works with IDE
	return &Location{
		SourceFile:  node.SourceFile,
		StartLine:   node.RNode.Document().Line,
		StartColumn: node.RNode.Document().Column,
		EndLine:     node.RNode.Document().Line + lines,
		EndColumn:   cols,
	}
}

func getLinesAndColsOfString(str string) (int, int) {
	line := 0
	col := 0
	for i := range str {
		col++
		if str[i] == '\n' {
			line++
			col = 0
		}
	}
	return line, col
}

func buildMapOfSchemaObjPointerToYAMLInfos(sourceFile string, config *latestV1.SkaffoldConfig, yamlInfos map[uintptr]map[string]YAMLInfo,
	fieldsOverrodeByProfile map[string]YAMLOverrideInfo) (map[uintptr]map[string]YAMLInfo, error) {
	defer func() {
		if err := recover(); err != nil {
			log.Entry(context.TODO()).Errorf(
				"panic occurred during schema reflection for yaml line number information: %v", err)
		}
	}()

	skaffoldConfigText, err := util.ReadConfiguration(sourceFile)
	if err != nil {
		return nil, sErrors.ConfigParsingError(err)
	}
	root, err := kyaml.Parse(string(skaffoldConfigText))
	if err != nil {
		return nil, err
	}

	return generateObjPointerToYAMLNodeMap(sourceFile, reflect.ValueOf(config), reflect.ValueOf(nil), "", "", []string{},
		root, root, -1, fieldsOverrodeByProfile, yamlInfos, false)
}

// generateObjPointerToYAMLNodeMap recursively walks through a structs fields (taking into account profile and patch profile overrides)
// and collects the corresponding yaml node for each field
func generateObjPointerToYAMLNodeMap(sourceFile string, v reflect.Value, parentV reflect.Value, fieldName, yamlTag string, schemaPath []string,
	rootRNode *kyaml.RNode, rNode *kyaml.RNode, containerIdx int, fieldPathsOverrodeByProfiles map[string]YAMLOverrideInfo, yamlInfos map[uintptr]map[string]YAMLInfo, isPatchProfileElemOverride bool) (map[uintptr]map[string]YAMLInfo, error) {
	// TODO(aaron-prindle) need to verify if generateObjPointerToYAMLNodeMap adds entries for 'map' types, luckily the skaffold schema
	// only has map[string]string and they are leaf nodes as well which this should work fine for doing the recursion for the time being
	var err error

	// add current obj/field to schema path if criteria met
	switch {
	case containerIdx >= 0:
		schemaPath = append(schemaPath, strconv.Itoa(containerIdx))
	case yamlTag != "":
		schemaPath = append(schemaPath, yamlTag)
	}
	// check if current obj/field was overridden by a profile
	if yamlOverrideInfo, ok := fieldPathsOverrodeByProfiles["/"+path.Join(schemaPath...)]; ok {
		// reset yaml node path from root path to given profile path ("/" -> "/profile/name=profileName/etc...")
		rNode, err = rootRNode.Pipe(kyaml.Lookup("profiles"), kyaml.MatchElementList([]string{"name"}, []string{yamlOverrideInfo.ProfileName}))
		if err != nil {
			return nil, err
		}
		switch {
		case yamlOverrideInfo.PatchIndex < 0: // this schema obj/field has a profile override (NOT a patch profile override)
			// moves parent node path from being rooted at default yaml '/' to being rooted at '/profile/name=profileName/...'
			for i := 0; i < len(schemaPath)-1; i++ {
				rNode, err = rNode.Pipe(kyaml.Lookup(schemaPath[i]))
				if err != nil {
					return nil, err
				}
			}
		default: // this schema obj/field has a patch profile override
			// NOTE: 'remove' patch operations are not included in fieldPathsOverrodeByProfiles as there
			//  is no work to be done on them (they were already removed from the schema)

			// TODO(aaron-prindle) verify UX makes sense to use the "FROM" copy node to get yaml information from
			if yamlOverrideInfo.PatchOperation == "copy" {
				fromPath := strings.Split(yamlOverrideInfo.PatchCopyFrom, "/")
				var kf kyaml.Filter
				for i := 0; i < len(fromPath)-1; i++ {
					if pathNum, err := strconv.Atoi(fromPath[i]); err == nil {
						// this path element is a number
						kf = kyaml.ElementIndexer{Index: pathNum}
					} else {
						// this path element isn't a number
						kf = kyaml.Lookup(fromPath[i])
					}
					rNode, err = rNode.Pipe(kf)
					if err != nil {
						return nil, err
					}
				}
			} else {
				rNode, err = rNode.Pipe(kyaml.Lookup("patches"), kyaml.ElementIndexer{Index: yamlOverrideInfo.PatchIndex})
				if err != nil {
					return nil, err
				}
				yamlTag = "value"
			}
			isPatchProfileElemOverride = true
		}
	}
	if rNode == nil {
		return yamlInfos, nil
	}

	// drill down through pointers and interfaces to get a value we can use
	for v.Kind() == reflect.Ptr || v.Kind() == reflect.Interface {
		v = v.Elem()
	}

	if yamlTag != "" { // check that struct is not `yaml:",inline"`
		// traverse kyaml node tree to current obj/field location
		var kf kyaml.Filter
		switch {
		case rNode.YNode().Kind == kyaml.SequenceNode:
			kf = kyaml.ElementIndexer{Index: containerIdx}
		default:
			kf = kyaml.Lookup(yamlTag)
		}
		rNode, err = rNode.Pipe(kf)
		if err != nil {
			return nil, err
		}
		if rNode == nil {
			return yamlInfos, nil
		}

		// this case is so that the line #'s of primitive values can be "located" as they are not addressable but we can
		// map the parent address and put the child field in second map
		if parentV.CanAddr() {
			if _, ok := yamlInfos[parentV.Addr().Pointer()]; !ok {
				yamlInfos[parentV.Addr().Pointer()] = map[string]YAMLInfo{}
			}
			// add parent relationship entry to yaml info map
			if containerIdx >= 0 {
				yamlInfos[parentV.Addr().Pointer()][strconv.Itoa(containerIdx)] = YAMLInfo{
					RNode:      rNode,
					SourceFile: sourceFile,
				}
			} else {
				yamlInfos[parentV.Addr().Pointer()][fieldName] = YAMLInfo{
					RNode:      rNode,
					SourceFile: sourceFile,
				}
			}
		}
	}

	if v.CanAddr() {
		if _, ok := yamlInfos[v.Addr().Pointer()]; !ok {
			yamlInfos[v.Addr().Pointer()] = map[string]YAMLInfo{}
		}
		// add current node entry to yaml info map
		yamlInfos[v.Addr().Pointer()][""] = YAMLInfo{
			RNode:      rNode,
			SourceFile: sourceFile,
		}
	}

	switch v.Kind() {
	// TODO(aaron-prindle) add reflect.Map support here as well, currently no struct fields have nested struct in map field so ok for now
	case reflect.Slice, reflect.Array:
		for i := 0; i < v.Len(); i++ {
			generateObjPointerToYAMLNodeMap(sourceFile, v.Index(i), v, fieldName+"["+strconv.Itoa(i)+"]", yamlTag+"["+strconv.Itoa(i)+"]", schemaPath,
				rootRNode, rNode, i, fieldPathsOverrodeByProfiles, yamlInfos, isPatchProfileElemOverride)
		}
	case reflect.Struct:
		t := v.Type() // use type to get number and names of fields
		for i := 0; i < t.NumField(); i++ {
			field := t.Field(i)
			// TODO(aaron-prindle) verify this value works for structs that are `yaml:",inline"`
			newYamlTag := field.Name
			if yamlTagToken := field.Tag.Get("yaml"); yamlTagToken != "" && yamlTagToken != "-" {
				// check for possible comma as in "...,omitempty"
				var commaIdx int
				if commaIdx = strings.Index(yamlTagToken, ","); commaIdx < 0 {
					commaIdx = len(yamlTagToken)
				}
				newYamlTag = yamlTagToken[:commaIdx]
			}
			generateObjPointerToYAMLNodeMap(sourceFile, v.Field(i), v, field.Name, newYamlTag, schemaPath, rootRNode, rNode, -1,
				fieldPathsOverrodeByProfiles, yamlInfos, isPatchProfileElemOverride)
		}
	}
	return yamlInfos, nil
}
<|endoftext|>


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

// Unix cryptographically secure pseudorandom number
// generator.

package rand

import (
	"crypto/aes"
	"io"
	"os"
	"sync"
	"time"
)

// Easy implementation: read from /dev/urandom.
// This is sufficient on Linux, OS X, and FreeBSD.

func init() { Reader = &devReader{name: "/dev/urandom"} }

// A devReader satisfies reads by reading the file named name.
type devReader struct {
	name string
	f    *os.File
	mu   sync.Mutex
}

func (r *devReader) Read(b []byte) (n int, err os.Error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.f == nil {
		f, err := os.Open(r.name, os.O_RDONLY, 0)
		if f == nil {
			return 0, err
		}
		r.f = f
	}
	return r.f.Read(b)
}

// Alternate pseudo-random implementation for use on
// systems without a reliable /dev/urandom.  So far we
// haven't needed it.

// newReader returns a new pseudorandom generator that
// seeds itself by reading from entropy.  If entropy == nil,
// the generator seeds itself by reading from the system's
// random number generator, typically /dev/random.
// The Read method on the returned reader always returns
// the full amount asked for, or else it returns an error.
//
// The generator uses the X9.31 algorithm with AES-128,
// reseeding after every 1 MB of generated data.
func newReader(entropy io.Reader) io.Reader {
	if entropy == nil {
		entropy = &devReader{name: "/dev/random"}
	}
	return &reader{entropy: entropy}
}

type reader struct {
	mu                   sync.Mutex
	budget               int // number of bytes that can be generated
	cipher               *aes.Cipher
	entropy              io.Reader
	time, seed, dst, key [aes.BlockSize]byte
}

func (r *reader) Read(b []byte) (n int, err os.Error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	n = len(b)

	for len(b) > 0 {
		if r.budget == 0 {
			_, err := io.ReadFull(r.entropy, r.seed[0:])
			if err != nil {
				return n - len(b), err
			}
			_, err = io.ReadFull(r.entropy, r.key[0:])
			if err != nil {
				return n - len(b), err
			}
			r.cipher, err = aes.NewCipher(r.key[0:])
			if err != nil {
				return n - len(b), err
			}
			r.budget = 1 << 20 // reseed after generating 1MB
		}
		r.budget -= aes.BlockSize

		// ANSI X9.31 (== X9.17) algorithm, but using AES in place of 3DES.
		//
		// single block:
		// t = encrypt(time)
		// dst = encrypt(t^seed)
		// seed = encrypt(t^dst)
		ns := time.Nanoseconds()
		r.time[0] = byte(ns >> 56)
		r.time[1] = byte(ns >> 48)
		r.time[2] = byte(ns >> 40)
		r.time[3] = byte(ns >> 32)
		r.time[4] = byte(ns >> 24)
		r.time[5] = byte(ns >> 16)
		r.time[6] = byte(ns >> 8)
		r.time[7] = byte(ns)
		r.cipher.Encrypt(r.time[0:], r.time[0:])
		for i := 0; i < aes.BlockSize; i++ {
			r.dst[i] = r.time[i] ^ r.seed[i]
		}
		r.cipher.Encrypt(r.dst[0:], r.dst[0:])
		for i := 0; i < aes.BlockSize; i++ {
			r.seed[i] = r.time[i] ^ r.dst[i]
		}
		r.cipher.Encrypt(r.seed[0:], r.seed[0:])

		m := copy(b, r.dst[0:])
		b = b[m:]
	}

	return n, nil
}
<commit_msg>crypto/rand: Added read buffer to speed up requests for small amounts of bytes.<commit_after>// Copyright 2010 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Unix cryptographically secure pseudorandom number
// generator.

package rand

import (
	"bufio"
	"crypto/aes"
	"io"
	"os"
	"sync"
	"time"
)

// Easy implementation: read from /dev/urandom.
// This is sufficient on Linux, OS X, and FreeBSD.

func init() { Reader = &devReader{name: "/dev/urandom"} }

// A devReader satisfies reads by reading the file named name.
type devReader struct {
	name string
	f    io.Reader
	mu   sync.Mutex
}

func (r *devReader) Read(b []byte) (n int, err os.Error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.f == nil {
		f, err := os.Open(r.name, os.O_RDONLY, 0)
		if f == nil {
			return 0, err
		}
		r.f = bufio.NewReader(f)
	}
	return r.f.Read(b)
}

// Alternate pseudo-random implementation for use on
// systems without a reliable /dev/urandom.  So far we
// haven't needed it.

// newReader returns a new pseudorandom generator that
// seeds itself by reading from entropy.  If entropy == nil,
// the generator seeds itself by reading from the system's
// random number generator, typically /dev/random.
// The Read method on the returned reader always returns
// the full amount asked for, or else it returns an error.
//
// The generator uses the X9.31 algorithm with AES-128,
// reseeding after every 1 MB of generated data.
func newReader(entropy io.Reader) io.Reader {
	if entropy == nil {
		entropy = &devReader{name: "/dev/random"}
	}
	return &reader{entropy: entropy}
}

type reader struct {
	mu                   sync.Mutex
	budget               int // number of bytes that can be generated
	cipher               *aes.Cipher
	entropy              io.Reader
	time, seed, dst, key [aes.BlockSize]byte
}

func (r *reader) Read(b []byte) (n int, err os.Error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	n = len(b)

	for len(b) > 0 {
		if r.budget == 0 {
			_, err := io.ReadFull(r.entropy, r.seed[0:])
			if err != nil {
				return n - len(b), err
			}
			_, err = io.ReadFull(r.entropy, r.key[0:])
			if err != nil {
				return n - len(b), err
			}
			r.cipher, err = aes.NewCipher(r.key[0:])
			if err != nil {
				return n - len(b), err
			}
			r.budget = 1 << 20 // reseed after generating 1MB
		}
		r.budget -= aes.BlockSize

		// ANSI X9.31 (== X9.17) algorithm, but using AES in place of 3DES.
		//
		// single block:
		// t = encrypt(time)
		// dst = encrypt(t^seed)
		// seed = encrypt(t^dst)
		ns := time.Nanoseconds()
		r.time[0] = byte(ns >> 56)
		r.time[1] = byte(ns >> 48)
		r.time[2] = byte(ns >> 40)
		r.time[3] = byte(ns >> 32)
		r.time[4] = byte(ns >> 24)
		r.time[5] = byte(ns >> 16)
		r.time[6] = byte(ns >> 8)
		r.time[7] = byte(ns)
		r.cipher.Encrypt(r.time[0:], r.time[0:])
		for i := 0; i < aes.BlockSize; i++ {
			r.dst[i] = r.time[i] ^ r.seed[i]
		}
		r.cipher.Encrypt(r.dst[0:], r.dst[0:])
		for i := 0; i < aes.BlockSize; i++ {
			r.seed[i] = r.time[i] ^ r.dst[i]
		}
		r.cipher.Encrypt(r.seed[0:], r.seed[0:])

		m := copy(b, r.dst[0:])
		b = b[m:]
	}

	return n, nil
}
<|endoftext|>


<commit_before>package discovery

import (
	"strings"
	"sync"

	"github.com/coreos/go-etcd/etcd"
)

type etcdClient struct {
	client *etcd.Client
}

func newEtcdClient(addresses ...string) *etcdClient {
	return &etcdClient{etcd.NewClient(addresses)}
}

func (c *etcdClient) Close() error {
	c.client.Close()
	return nil
}

func (c *etcdClient) Get(key string) (string, bool, error) {
	response, err := c.client.Get(key, false, false)
	if err != nil {
		if strings.HasPrefix(err.Error(), "100: Key not found") {
			return "", false, nil
		}
		return "", false, err
	}
	return response.Node.Value, true, nil
}

func (c *etcdClient) GetAll(key string) (map[string]string, error) {
	response, err := c.client.Get(key, false, true)
	result := make(map[string]string, 0)
	if err != nil {
		if strings.HasPrefix(err.Error(), "100: Key not found") {
			return result, nil
		}
		return nil, err
	}
	nodeToMap(response.Node, result)
	return result, nil
}

func (c *etcdClient) Watch(key string, cancel chan bool, callBack func(string) error) (retErr error) {
	if err := callBack(""); err != nil {
		return err
	}
	localCancel := make(chan bool)
	var once sync.Once
	var errOnce sync.Once
	go func() {
		select {
		case <-cancel:
			once.Do(func() { close(localCancel) })
		case <-localCancel:
		}
	}()
	receiver := make(chan *etcd.Response)
	defer close(receiver)
	go func() {
		for response := range receiver {
			if err := callBack(response.Node.Value); err != nil {
				errOnce.Do(func() { retErr = err })
				once.Do(func() { close(localCancel) })
			}
		}
	}()
	if _, err := c.client.Watch(key, 0, false, receiver, localCancel); err != nil {
		errOnce.Do(func() { retErr = err })
	}
	return
}

func (c *etcdClient) WatchAll(key string, cancel chan bool, callBack func(map[string]string) error) (retErr error) {
	if err := callBack(nil); err != nil {
		return err
	}
	localCancel := make(chan bool)
	var errOnce sync.Once
	var once sync.Once
	go func() {
		select {
		case <-cancel:
			once.Do(func() { close(localCancel) })
		case <-localCancel:
		}
	}()
	receiver := make(chan *etcd.Response)
	defer close(receiver)
	go func() {
		for response := range receiver {
			value := make(map[string]string)
			nodeToMap(response.Node, value)
			if err := callBack(value); err != nil {
				errOnce.Do(func() { retErr = err })
				once.Do(func() { close(localCancel) })
			}
		}
	}()
	if _, err := c.client.Watch(key, 0, true, receiver, localCancel); err != nil {
		errOnce.Do(func() { retErr = err })
	}
	return
}

func (c *etcdClient) Set(key string, value string, ttl uint64) error {
	_, err := c.client.Set(key, value, ttl)
	return err
}

func (c *etcdClient) Create(key string, value string, ttl uint64) error {
	_, err := c.client.Create(key, value, ttl)
	return err
}

func (c *etcdClient) CreateInDir(dir string, value string, ttl uint64) error {
	_, err := c.client.CreateInOrder(dir, value, ttl)
	return err
}

func (c *etcdClient) Delete(key string) error {
	_, err := c.client.Delete(key, false)
	return err
}

func (c *etcdClient) CheckAndSet(key string, value string, ttl uint64, oldValue string) error {
	_, err := c.client.CompareAndSwap(key, value, ttl, oldValue, 0)
	return err
}

func nodeToMap(node *etcd.Node, out map[string]string) {
	if !node.Dir {
		out[strings.TrimPrefix(node.Key, "/")] = node.Value
	} else {
		for _, node := range node.Nodes {
			nodeToMap(node, out)
		}
	}
}
<commit_msg>Adds a todokjk<commit_after>package discovery

import (
	"strings"
	"sync"

	"github.com/coreos/go-etcd/etcd"
)

type etcdClient struct {
	client *etcd.Client
}

func newEtcdClient(addresses ...string) *etcdClient {
	return &etcdClient{etcd.NewClient(addresses)}
}

func (c *etcdClient) Close() error {
	c.client.Close()
	return nil
}

func (c *etcdClient) Get(key string) (string, bool, error) {
	response, err := c.client.Get(key, false, false)
	if err != nil {
		if strings.HasPrefix(err.Error(), "100: Key not found") {
			return "", false, nil
		}
		return "", false, err
	}
	return response.Node.Value, true, nil
}

func (c *etcdClient) GetAll(key string) (map[string]string, error) {
	response, err := c.client.Get(key, false, true)
	result := make(map[string]string, 0)
	if err != nil {
		if strings.HasPrefix(err.Error(), "100: Key not found") {
			return result, nil
		}
		return nil, err
	}
	nodeToMap(response.Node, result)
	return result, nil
}

//TODO jdoliner we're doing some acrobatics here to make the cancel chan work
//we should think about if there's a better way to do this or maybe abstract
//this functionality into a class
func (c *etcdClient) Watch(key string, cancel chan bool, callBack func(string) error) (retErr error) {
	if err := callBack(""); err != nil {
		return err
	}
	localCancel := make(chan bool)
	var once sync.Once
	var errOnce sync.Once
	go func() {
		select {
		case <-cancel:
			once.Do(func() { close(localCancel) })
		case <-localCancel:
		}
	}()
	receiver := make(chan *etcd.Response)
	defer close(receiver)
	go func() {
		for response := range receiver {
			if err := callBack(response.Node.Value); err != nil {
				errOnce.Do(func() { retErr = err })
				once.Do(func() { close(localCancel) })
			}
		}
	}()
	if _, err := c.client.Watch(key, 0, false, receiver, localCancel); err != nil {
		errOnce.Do(func() { retErr = err })
	}
	return
}

func (c *etcdClient) WatchAll(key string, cancel chan bool, callBack func(map[string]string) error) (retErr error) {
	if err := callBack(nil); err != nil {
		return err
	}
	localCancel := make(chan bool)
	var errOnce sync.Once
	var once sync.Once
	go func() {
		select {
		case <-cancel:
			once.Do(func() { close(localCancel) })
		case <-localCancel:
		}
	}()
	receiver := make(chan *etcd.Response)
	defer close(receiver)
	go func() {
		for response := range receiver {
			value := make(map[string]string)
			nodeToMap(response.Node, value)
			if err := callBack(value); err != nil {
				errOnce.Do(func() { retErr = err })
				once.Do(func() { close(localCancel) })
			}
		}
	}()
	if _, err := c.client.Watch(key, 0, true, receiver, localCancel); err != nil {
		errOnce.Do(func() { retErr = err })
	}
	return
}

func (c *etcdClient) Set(key string, value string, ttl uint64) error {
	_, err := c.client.Set(key, value, ttl)
	return err
}

func (c *etcdClient) Create(key string, value string, ttl uint64) error {
	_, err := c.client.Create(key, value, ttl)
	return err
}

func (c *etcdClient) CreateInDir(dir string, value string, ttl uint64) error {
	_, err := c.client.CreateInOrder(dir, value, ttl)
	return err
}

func (c *etcdClient) Delete(key string) error {
	_, err := c.client.Delete(key, false)
	return err
}

func (c *etcdClient) CheckAndSet(key string, value string, ttl uint64, oldValue string) error {
	_, err := c.client.CompareAndSwap(key, value, ttl, oldValue, 0)
	return err
}

func nodeToMap(node *etcd.Node, out map[string]string) {
	if !node.Dir {
		out[strings.TrimPrefix(node.Key, "/")] = node.Value
	} else {
		for _, node := range node.Nodes {
			nodeToMap(node, out)
		}
	}
}
<|endoftext|>


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

package png

import (
	"bufio"
	"fmt"
	"image"
	"image/color"
	"io"
	"os"
	"strings"
	"testing"
)

var filenames = []string{
	"basn0g01",
	"basn0g01-30",
	"basn0g02",
	"basn0g02-29",
	"basn0g04",
	"basn0g04-31",
	"basn0g08",
	"basn0g16",
	"basn2c08",
	"basn2c16",
	"basn3p01",
	"basn3p02",
	"basn3p04",
	"basn3p08",
	"basn3p08-trns",
	"basn4a08",
	"basn4a16",
	"basn6a08",
	"basn6a16",
}

var filenamesShort = []string{
	"basn0g01",
	"basn0g04-31",
	"basn6a16",
}

func readPNG(filename string) (image.Image, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return Decode(f)
}

// An approximation of the sng command-line tool.
func sng(w io.WriteCloser, filename string, png image.Image) {
	defer w.Close()
	bounds := png.Bounds()
	cm := png.ColorModel()
	var bitdepth int
	switch cm {
	case color.RGBAModel, color.NRGBAModel, color.AlphaModel, color.GrayModel:
		bitdepth = 8
	default:
		bitdepth = 16
	}
	cpm, _ := cm.(color.Palette)
	var paletted *image.Paletted
	if cpm != nil {
		switch {
		case len(cpm) <= 2:
			bitdepth = 1
		case len(cpm) <= 4:
			bitdepth = 2
		case len(cpm) <= 16:
			bitdepth = 4
		default:
			bitdepth = 8
		}
		paletted = png.(*image.Paletted)
	}

	// Write the filename and IHDR.
	io.WriteString(w, "#SNG: from "+filename+".png\nIHDR {\n")
	fmt.Fprintf(w, "    width: %d; height: %d; bitdepth: %d;\n", bounds.Dx(), bounds.Dy(), bitdepth)
	switch {
	case cm == color.RGBAModel, cm == color.RGBA64Model:
		io.WriteString(w, "    using color;\n")
	case cm == color.NRGBAModel, cm == color.NRGBA64Model:
		io.WriteString(w, "    using color alpha;\n")
	case cm == color.GrayModel, cm == color.Gray16Model:
		io.WriteString(w, "    using grayscale;\n")
	case cpm != nil:
		io.WriteString(w, "    using color palette;\n")
	default:
		io.WriteString(w, "unknown PNG decoder color model\n")
	}
	io.WriteString(w, "}\n")

	// We fake a gAMA output. The test files have a gAMA chunk but the go PNG parser ignores it
	// (the PNG spec section 11.3 says "Ancillary chunks may be ignored by a decoder").
	io.WriteString(w, "gAMA {1.0000}\n")

	// Write the PLTE and tRNS (if applicable).
	if cpm != nil {
		lastAlpha := -1
		io.WriteString(w, "PLTE {\n")
		for i, c := range cpm {
			r, g, b, a := c.RGBA()
			if a != 0xffff {
				lastAlpha = i
			}
			r >>= 8
			g >>= 8
			b >>= 8
			fmt.Fprintf(w, "    (%3d,%3d,%3d)     # rgb = (0x%02x,0x%02x,0x%02x)\n", r, g, b, r, g, b)
		}
		io.WriteString(w, "}\n")
		if lastAlpha != -1 {
			io.WriteString(w, "tRNS {\n")
			for i := 0; i <= lastAlpha; i++ {
				_, _, _, a := cpm[i].RGBA()
				a >>= 8
				fmt.Fprintf(w, " %d", a)
			}
			io.WriteString(w, "}\n")
		}
	}

	// Write the IMAGE.
	io.WriteString(w, "IMAGE {\n    pixels hex\n")
	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		switch {
		case cm == color.GrayModel:
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				gray := png.At(x, y).(color.Gray)
				fmt.Fprintf(w, "%02x", gray.Y)
			}
		case cm == color.Gray16Model:
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				gray16 := png.At(x, y).(color.Gray16)
				fmt.Fprintf(w, "%04x ", gray16.Y)
			}
		case cm == color.RGBAModel:
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				rgba := png.At(x, y).(color.RGBA)
				fmt.Fprintf(w, "%02x%02x%02x ", rgba.R, rgba.G, rgba.B)
			}
		case cm == color.RGBA64Model:
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				rgba64 := png.At(x, y).(color.RGBA64)
				fmt.Fprintf(w, "%04x%04x%04x ", rgba64.R, rgba64.G, rgba64.B)
			}
		case cm == color.NRGBAModel:
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				nrgba := png.At(x, y).(color.NRGBA)
				fmt.Fprintf(w, "%02x%02x%02x%02x ", nrgba.R, nrgba.G, nrgba.B, nrgba.A)
			}
		case cm == color.NRGBA64Model:
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				nrgba64 := png.At(x, y).(color.NRGBA64)
				fmt.Fprintf(w, "%04x%04x%04x%04x ", nrgba64.R, nrgba64.G, nrgba64.B, nrgba64.A)
			}
		case cpm != nil:
			var b, c int
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				b = b<<uint(bitdepth) | int(paletted.ColorIndexAt(x, y))
				c++
				if c == 8/bitdepth {
					fmt.Fprintf(w, "%02x", b)
					b = 0
					c = 0
				}
			}
		}
		io.WriteString(w, "\n")
	}
	io.WriteString(w, "}\n")
}

func TestReader(t *testing.T) {
	names := filenames
	if testing.Short() {
		names = filenamesShort
	}
	for _, fn := range names {
		// Read the .png file.
		img, err := readPNG("testdata/pngsuite/" + fn + ".png")
		if err != nil {
			t.Error(fn, err)
			continue
		}

		if fn == "basn4a16" {
			// basn4a16.sng is gray + alpha but sng() will produce true color + alpha
			// so we just check a single random pixel.
			c := img.At(2, 1).(color.NRGBA64)
			if c.R != 0x11a7 || c.G != 0x11a7 || c.B != 0x11a7 || c.A != 0x1085 {
				t.Error(fn, fmt.Errorf("wrong pixel value at (2, 1): %x", c))
			}
			continue
		}

		piper, pipew := io.Pipe()
		pb := bufio.NewReader(piper)
		go sng(pipew, fn, img)
		defer piper.Close()

		// Read the .sng file.
		sf, err := os.Open("testdata/pngsuite/" + fn + ".sng")
		if err != nil {
			t.Error(fn, err)
			continue
		}
		defer sf.Close()
		sb := bufio.NewReader(sf)
		if err != nil {
			t.Error(fn, err)
			continue
		}

		// Compare the two, in SNG format, line by line.
		for {
			ps, perr := pb.ReadString('\n')
			ss, serr := sb.ReadString('\n')
			if perr == io.EOF && serr == io.EOF {
				break
			}
			if perr != nil {
				t.Error(fn, perr)
				break
			}
			if serr != nil {
				t.Error(fn, serr)
				break
			}
			if ps != ss {
				t.Errorf("%s: Mismatch\n%sversus\n%s\n", fn, ps, ss)
				break
			}
		}
	}
}

var readerErrors = []struct {
	file string
	err  string
}{
	{"invalid-zlib.png", "zlib checksum error"},
	{"invalid-crc32.png", "invalid checksum"},
	{"invalid-noend.png", "unexpected EOF"},
	{"invalid-trunc.png", "unexpected EOF"},
}

func TestReaderError(t *testing.T) {
	for _, tt := range readerErrors {
		img, err := readPNG("testdata/" + tt.file)
		if err == nil {
			t.Errorf("decoding %s: missing error", tt.file)
			continue
		}
		if !strings.Contains(err.Error(), tt.err) {
			t.Errorf("decoding %s: %s, want %s", tt.file, err, tt.err)
		}
		if img != nil {
			t.Errorf("decoding %s: have image + error", tt.file)
		}
	}
}
<commit_msg>fix build: update image/png test<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package png

import (
	"bufio"
	"fmt"
	"image"
	"image/color"
	"io"
	"os"
	"strings"
	"testing"
)

var filenames = []string{
	"basn0g01",
	"basn0g01-30",
	"basn0g02",
	"basn0g02-29",
	"basn0g04",
	"basn0g04-31",
	"basn0g08",
	"basn0g16",
	"basn2c08",
	"basn2c16",
	"basn3p01",
	"basn3p02",
	"basn3p04",
	"basn3p08",
	"basn3p08-trns",
	"basn4a08",
	"basn4a16",
	"basn6a08",
	"basn6a16",
}

var filenamesShort = []string{
	"basn0g01",
	"basn0g04-31",
	"basn6a16",
}

func readPNG(filename string) (image.Image, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return Decode(f)
}

// An approximation of the sng command-line tool.
func sng(w io.WriteCloser, filename string, png image.Image) {
	defer w.Close()
	bounds := png.Bounds()
	cm := png.ColorModel()
	var bitdepth int
	switch cm {
	case color.RGBAModel, color.NRGBAModel, color.AlphaModel, color.GrayModel:
		bitdepth = 8
	default:
		bitdepth = 16
	}
	cpm, _ := cm.(color.Palette)
	var paletted *image.Paletted
	if cpm != nil {
		switch {
		case len(cpm) <= 2:
			bitdepth = 1
		case len(cpm) <= 4:
			bitdepth = 2
		case len(cpm) <= 16:
			bitdepth = 4
		default:
			bitdepth = 8
		}
		paletted = png.(*image.Paletted)
	}

	// Write the filename and IHDR.
	io.WriteString(w, "#SNG: from "+filename+".png\nIHDR {\n")
	fmt.Fprintf(w, "    width: %d; height: %d; bitdepth: %d;\n", bounds.Dx(), bounds.Dy(), bitdepth)
	switch {
	case cm == color.RGBAModel, cm == color.RGBA64Model:
		io.WriteString(w, "    using color;\n")
	case cm == color.NRGBAModel, cm == color.NRGBA64Model:
		io.WriteString(w, "    using color alpha;\n")
	case cm == color.GrayModel, cm == color.Gray16Model:
		io.WriteString(w, "    using grayscale;\n")
	case cpm != nil:
		io.WriteString(w, "    using color palette;\n")
	default:
		io.WriteString(w, "unknown PNG decoder color model\n")
	}
	io.WriteString(w, "}\n")

	// We fake a gAMA output. The test files have a gAMA chunk but the go PNG parser ignores it
	// (the PNG spec section 11.3 says "Ancillary chunks may be ignored by a decoder").
	io.WriteString(w, "gAMA {1.0000}\n")

	// Write the PLTE and tRNS (if applicable).
	if cpm != nil {
		lastAlpha := -1
		io.WriteString(w, "PLTE {\n")
		for i, c := range cpm {
			r, g, b, a := c.RGBA()
			if a != 0xffff {
				lastAlpha = i
			}
			r >>= 8
			g >>= 8
			b >>= 8
			fmt.Fprintf(w, "    (%3d,%3d,%3d)     # rgb = (0x%02x,0x%02x,0x%02x)\n", r, g, b, r, g, b)
		}
		io.WriteString(w, "}\n")
		if lastAlpha != -1 {
			io.WriteString(w, "tRNS {\n")
			for i := 0; i <= lastAlpha; i++ {
				_, _, _, a := cpm[i].RGBA()
				a >>= 8
				fmt.Fprintf(w, " %d", a)
			}
			io.WriteString(w, "}\n")
		}
	}

	// Write the IMAGE.
	io.WriteString(w, "IMAGE {\n    pixels hex\n")
	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		switch {
		case cm == color.GrayModel:
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				gray := png.At(x, y).(color.Gray)
				fmt.Fprintf(w, "%02x", gray.Y)
			}
		case cm == color.Gray16Model:
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				gray16 := png.At(x, y).(color.Gray16)
				fmt.Fprintf(w, "%04x ", gray16.Y)
			}
		case cm == color.RGBAModel:
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				rgba := png.At(x, y).(color.RGBA)
				fmt.Fprintf(w, "%02x%02x%02x ", rgba.R, rgba.G, rgba.B)
			}
		case cm == color.RGBA64Model:
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				rgba64 := png.At(x, y).(color.RGBA64)
				fmt.Fprintf(w, "%04x%04x%04x ", rgba64.R, rgba64.G, rgba64.B)
			}
		case cm == color.NRGBAModel:
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				nrgba := png.At(x, y).(color.NRGBA)
				fmt.Fprintf(w, "%02x%02x%02x%02x ", nrgba.R, nrgba.G, nrgba.B, nrgba.A)
			}
		case cm == color.NRGBA64Model:
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				nrgba64 := png.At(x, y).(color.NRGBA64)
				fmt.Fprintf(w, "%04x%04x%04x%04x ", nrgba64.R, nrgba64.G, nrgba64.B, nrgba64.A)
			}
		case cpm != nil:
			var b, c int
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				b = b<<uint(bitdepth) | int(paletted.ColorIndexAt(x, y))
				c++
				if c == 8/bitdepth {
					fmt.Fprintf(w, "%02x", b)
					b = 0
					c = 0
				}
			}
		}
		io.WriteString(w, "\n")
	}
	io.WriteString(w, "}\n")
}

func TestReader(t *testing.T) {
	names := filenames
	if testing.Short() {
		names = filenamesShort
	}
	for _, fn := range names {
		// Read the .png file.
		img, err := readPNG("testdata/pngsuite/" + fn + ".png")
		if err != nil {
			t.Error(fn, err)
			continue
		}

		if fn == "basn4a16" {
			// basn4a16.sng is gray + alpha but sng() will produce true color + alpha
			// so we just check a single random pixel.
			c := img.At(2, 1).(color.NRGBA64)
			if c.R != 0x11a7 || c.G != 0x11a7 || c.B != 0x11a7 || c.A != 0x1085 {
				t.Error(fn, fmt.Errorf("wrong pixel value at (2, 1): %x", c))
			}
			continue
		}

		piper, pipew := io.Pipe()
		pb := bufio.NewReader(piper)
		go sng(pipew, fn, img)
		defer piper.Close()

		// Read the .sng file.
		sf, err := os.Open("testdata/pngsuite/" + fn + ".sng")
		if err != nil {
			t.Error(fn, err)
			continue
		}
		defer sf.Close()
		sb := bufio.NewReader(sf)
		if err != nil {
			t.Error(fn, err)
			continue
		}

		// Compare the two, in SNG format, line by line.
		for {
			ps, perr := pb.ReadString('\n')
			ss, serr := sb.ReadString('\n')
			if perr == io.EOF && serr == io.EOF {
				break
			}
			if perr != nil {
				t.Error(fn, perr)
				break
			}
			if serr != nil {
				t.Error(fn, serr)
				break
			}
			if ps != ss {
				t.Errorf("%s: Mismatch\n%sversus\n%s\n", fn, ps, ss)
				break
			}
		}
	}
}

var readerErrors = []struct {
	file string
	err  string
}{
	{"invalid-zlib.png", "zlib: invalid checksum"},
	{"invalid-crc32.png", "invalid checksum"},
	{"invalid-noend.png", "unexpected EOF"},
	{"invalid-trunc.png", "unexpected EOF"},
}

func TestReaderError(t *testing.T) {
	for _, tt := range readerErrors {
		img, err := readPNG("testdata/" + tt.file)
		if err == nil {
			t.Errorf("decoding %s: missing error", tt.file)
			continue
		}
		if !strings.Contains(err.Error(), tt.err) {
			t.Errorf("decoding %s: %s, want %s", tt.file, err, tt.err)
		}
		if img != nil {
			t.Errorf("decoding %s: have image + error", tt.file)
		}
	}
}
<|endoftext|>


<commit_before>package sso

import (
	"encoding/json"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"github.com/skygeario/skygear-server/pkg/core/config"
	"github.com/skygeario/skygear-server/pkg/core/errors"
)

type AccessTokenResp map[string]interface{}

func NewBearerAccessTokenResp(accessToken string) AccessTokenResp {
	return AccessTokenResp{
		"token_type":   "Bearer",
		"access_token": accessToken,
	}
}

func (r AccessTokenResp) IDToken() string {
	idToken, ok := r["id_token"].(string)
	if ok {
		return idToken
	}
	return ""
}

func (r AccessTokenResp) AccessToken() string {
	accessToken, ok := r["access_token"].(string)
	if ok {
		return accessToken
	}
	return ""
}

func (r AccessTokenResp) ExpiresIn() int {
	expires, hasExpires := r["expires"]
	expiresIn, hasExpiresIn := r["expires_in"]

	// Facebook use "expires" instead of "expires_in"
	if hasExpires && !hasExpiresIn {
		expiresIn = expires
	}

	switch v := expiresIn.(type) {
	// Azure AD v2 uses string instead of number
	case string:
		i, err := strconv.Atoi(v)
		if err != nil {
			return 0
		}
		return i
	case float64:
		return int(v)
	default:
		return 0
	}
}

func (r AccessTokenResp) TokenType() string {
	tokenType, ok := r["token_type"].(string)
	if !ok {
		return ""
	}
	tokenType = strings.ToLower(tokenType)
	switch tokenType {
	case "basic":
		return "Basic"
	case "digest":
		return "Digest"
	case "bearer":
		return "Bearer"
	// We do not care about other less common schemes.
	default:
		return tokenType
	}
}

func (r AccessTokenResp) Validate() error {
	if r.AccessToken() == "" {
		err := NewSSOFailed(SSOUnauthorized, "unexpected authorization response")
		return err
	}

	return nil
}

func fetchAccessTokenResp(
	code string,
	accessTokenURL string,
	redirectURL string,
	oauthConfig *config.OAuthConfiguration,
	providerConfig config.OAuthProviderConfiguration,
) (r AccessTokenResp, err error) {
	v := url.Values{}
	v.Set("grant_type", "authorization_code")
	v.Add("code", code)
	v.Add("redirect_uri", redirectURL)
	v.Add("client_id", providerConfig.ClientID)
	v.Add("client_secret", providerConfig.ClientSecret)

	// nolint: gosec
	resp, err := http.PostForm(accessTokenURL, v)
	if resp != nil {
		defer resp.Body.Close()
	}
	if err != nil {
		err = errors.WithSecondaryError(
			NewSSOFailed(NetworkFailed, "failed to connect authorization server"),
			err,
		)
		return
	}

	if resp.StatusCode == 200 {
		err = json.NewDecoder(resp.Body).Decode(&r)
		if err != nil {
			return
		}
	} else { // normally 400 Bad Request
		var errResp oauthErrorResp
		err = json.NewDecoder(resp.Body).Decode(&errResp)
		if err != nil {
			return
		}
		err = errResp.AsError()
	}

	return
}
<commit_msg>Assume token_type is bearer if absent<commit_after>package sso

import (
	"encoding/json"
	"net/http"
	"net/url"
	"strconv"
	"strings"

	"github.com/skygeario/skygear-server/pkg/core/config"
	"github.com/skygeario/skygear-server/pkg/core/errors"
)

type AccessTokenResp map[string]interface{}

func NewBearerAccessTokenResp(accessToken string) AccessTokenResp {
	return AccessTokenResp{
		"token_type":   "Bearer",
		"access_token": accessToken,
	}
}

func (r AccessTokenResp) IDToken() string {
	idToken, ok := r["id_token"].(string)
	if ok {
		return idToken
	}
	return ""
}

func (r AccessTokenResp) AccessToken() string {
	accessToken, ok := r["access_token"].(string)
	if ok {
		return accessToken
	}
	return ""
}

func (r AccessTokenResp) ExpiresIn() int {
	expires, hasExpires := r["expires"]
	expiresIn, hasExpiresIn := r["expires_in"]

	// Facebook use "expires" instead of "expires_in"
	if hasExpires && !hasExpiresIn {
		expiresIn = expires
	}

	switch v := expiresIn.(type) {
	// Azure AD v2 uses string instead of number
	case string:
		i, err := strconv.Atoi(v)
		if err != nil {
			return 0
		}
		return i
	case float64:
		return int(v)
	default:
		return 0
	}
}

func (r AccessTokenResp) TokenType() string {
	tokenType, ok := r["token_type"].(string)
	if !ok {
		// LinkedIn does not include token_type in the response.
		return "Bearer"
	}
	tokenType = strings.ToLower(tokenType)
	switch tokenType {
	case "basic":
		return "Basic"
	case "digest":
		return "Digest"
	case "bearer":
		return "Bearer"
	// We do not care about other less common schemes.
	default:
		return tokenType
	}
}

func (r AccessTokenResp) Validate() error {
	if r.AccessToken() == "" {
		err := NewSSOFailed(SSOUnauthorized, "unexpected authorization response")
		return err
	}

	return nil
}

func fetchAccessTokenResp(
	code string,
	accessTokenURL string,
	redirectURL string,
	oauthConfig *config.OAuthConfiguration,
	providerConfig config.OAuthProviderConfiguration,
) (r AccessTokenResp, err error) {
	v := url.Values{}
	v.Set("grant_type", "authorization_code")
	v.Add("code", code)
	v.Add("redirect_uri", redirectURL)
	v.Add("client_id", providerConfig.ClientID)
	v.Add("client_secret", providerConfig.ClientSecret)

	// nolint: gosec
	resp, err := http.PostForm(accessTokenURL, v)
	if resp != nil {
		defer resp.Body.Close()
	}
	if err != nil {
		err = errors.WithSecondaryError(
			NewSSOFailed(NetworkFailed, "failed to connect authorization server"),
			err,
		)
		return
	}

	if resp.StatusCode == 200 {
		err = json.NewDecoder(resp.Body).Decode(&r)
		if err != nil {
			return
		}
	} else { // normally 400 Bad Request
		var errResp oauthErrorResp
		err = json.NewDecoder(resp.Body).Decode(&errResp)
		if err != nil {
			return
		}
		err = errResp.AsError()
	}

	return
}
<|endoftext|>


<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 genericapiserver

import (
	"fmt"
	"mime"
	"strings"

	"k8s.io/kubernetes/pkg/api/unversioned"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/runtime/serializer/recognizer"
	"k8s.io/kubernetes/pkg/storage/storagebackend"
	"k8s.io/kubernetes/pkg/util/sets"

	"github.com/golang/glog"
)

// StorageFactory is the interface to locate the storage for a given GroupResource
type StorageFactory interface {
	// New finds the storage destination for the given group and resource. It will
	// return an error if the group has no storage destination configured.
	NewConfig(groupResource unversioned.GroupResource) (*storagebackend.Config, error)

	// ResourcePrefix returns the overridden resource prefix for the GroupResource
	// This allows for cohabitation of resources with different native types and provides
	// centralized control over the shape of etcd directories
	ResourcePrefix(groupResource unversioned.GroupResource) string

	// Backends gets all backends for all registered storage destinations.
	// Used for getting all instances for health validations.
	Backends() []string
}

// DefaultStorageFactory takes a GroupResource and returns back its storage interface.  This result includes:
// 1. Merged etcd config, including: auth, server locations, prefixes
// 2. Resource encodings for storage: group,version,kind to store as
// 3. Cohabitating default: some resources like hpa are exposed through multiple APIs.  They must agree on 1 and 2
type DefaultStorageFactory struct {
	// StorageConfig describes how to create a storage backend in general.
	// Its authentication information will be used for every storage.Interface returned.
	StorageConfig storagebackend.Config

	Overrides map[unversioned.GroupResource]groupResourceOverrides

	// DefaultMediaType is the media type used to store resources. If it is not set, "application/json" is used.
	DefaultMediaType string

	// DefaultSerializer is used to create encoders and decoders for the storage.Interface.
	DefaultSerializer runtime.StorageSerializer

	// ResourceEncodingConfig describes how to encode a particular GroupVersionResource
	ResourceEncodingConfig ResourceEncodingConfig

	// APIResourceConfigSource indicates whether the *storage* is enabled, NOT the API
	// This is discrete from resource enablement because those are separate concerns.  How this source is configured
	// is left to the caller.
	APIResourceConfigSource APIResourceConfigSource

	// newStorageCodecFn exists to be overwritten for unit testing.
	newStorageCodecFn func(storageMediaType string, ns runtime.StorageSerializer, storageVersion, memoryVersion unversioned.GroupVersion, config storagebackend.Config) (codec runtime.Codec, err error)
}

type groupResourceOverrides struct {
	// etcdLocation contains the list of "special" locations that are used for particular GroupResources
	// These are merged on top of the StorageConfig when requesting the storage.Interface for a given GroupResource
	etcdLocation []string
	// etcdPrefix is the base location for a GroupResource.
	etcdPrefix string
	// etcdResourcePrefix is the location to use to store a particular type under the `etcdPrefix` location
	// If empty, the default mapping is used.  If the default mapping doesn't contain an entry, it will use
	// the ToLowered name of the resource, not including the group.
	etcdResourcePrefix string
	// mediaType is the desired serializer to choose. If empty, the default is chosen.
	mediaType string
	// serializer contains the list of "special" serializers for a GroupResource.  Resource=* means for the entire group
	serializer runtime.StorageSerializer
	// cohabitatingResources keeps track of which resources must be stored together.  This happens when we have multiple ways
	// of exposing one set of concepts.  autoscaling.HPA and extensions.HPA as a for instance
	// The order of the slice matters!  It is the priority order of lookup for finding a storage location
	cohabitatingResources []unversioned.GroupResource
}

var _ StorageFactory = &DefaultStorageFactory{}

const AllResources = "*"

func NewDefaultStorageFactory(config storagebackend.Config, defaultMediaType string, defaultSerializer runtime.StorageSerializer, resourceEncodingConfig ResourceEncodingConfig, resourceConfig APIResourceConfigSource) *DefaultStorageFactory {
	if len(defaultMediaType) == 0 {
		defaultMediaType = runtime.ContentTypeJSON
	}
	return &DefaultStorageFactory{
		StorageConfig:           config,
		Overrides:               map[unversioned.GroupResource]groupResourceOverrides{},
		DefaultMediaType:        defaultMediaType,
		DefaultSerializer:       defaultSerializer,
		ResourceEncodingConfig:  resourceEncodingConfig,
		APIResourceConfigSource: resourceConfig,

		newStorageCodecFn: NewStorageCodec,
	}
}

func (s *DefaultStorageFactory) SetEtcdLocation(groupResource unversioned.GroupResource, location []string) {
	overrides := s.Overrides[groupResource]
	overrides.etcdLocation = location
	s.Overrides[groupResource] = overrides
}

func (s *DefaultStorageFactory) SetEtcdPrefix(groupResource unversioned.GroupResource, prefix string) {
	overrides := s.Overrides[groupResource]
	overrides.etcdPrefix = prefix
	s.Overrides[groupResource] = overrides
}

// SetResourceEtcdPrefix sets the prefix for a resource, but not the base-dir.  You'll end up in `etcdPrefix/resourceEtcdPrefix`.
func (s *DefaultStorageFactory) SetResourceEtcdPrefix(groupResource unversioned.GroupResource, prefix string) {
	overrides := s.Overrides[groupResource]
	overrides.etcdResourcePrefix = prefix
	s.Overrides[groupResource] = overrides
}

func (s *DefaultStorageFactory) SetSerializer(groupResource unversioned.GroupResource, mediaType string, serializer runtime.StorageSerializer) {
	overrides := s.Overrides[groupResource]
	overrides.mediaType = mediaType
	overrides.serializer = serializer
	s.Overrides[groupResource] = overrides
}

// AddCohabitatingResources links resources together the order of the slice matters!  its the priority order of lookup for finding a storage location
func (s *DefaultStorageFactory) AddCohabitatingResources(groupResources ...unversioned.GroupResource) {
	for _, groupResource := range groupResources {
		overrides := s.Overrides[groupResource]
		overrides.cohabitatingResources = groupResources
		s.Overrides[groupResource] = overrides
	}
}

func getAllResourcesAlias(resource unversioned.GroupResource) unversioned.GroupResource {
	return unversioned.GroupResource{Group: resource.Group, Resource: AllResources}
}

func (s *DefaultStorageFactory) getStorageGroupResource(groupResource unversioned.GroupResource) unversioned.GroupResource {
	for _, potentialStorageResource := range s.Overrides[groupResource].cohabitatingResources {
		if s.APIResourceConfigSource.AnyVersionOfResourceEnabled(potentialStorageResource) {
			return potentialStorageResource
		}
	}

	return groupResource
}

// New finds the storage destination for the given group and resource. It will
// return an error if the group has no storage destination configured.
func (s *DefaultStorageFactory) NewConfig(groupResource unversioned.GroupResource) (*storagebackend.Config, error) {
	chosenStorageResource := s.getStorageGroupResource(groupResource)

	groupOverride := s.Overrides[getAllResourcesAlias(chosenStorageResource)]
	exactResourceOverride := s.Overrides[chosenStorageResource]

	overriddenEtcdLocations := []string{}
	if len(groupOverride.etcdLocation) > 0 {
		overriddenEtcdLocations = groupOverride.etcdLocation
	}
	if len(exactResourceOverride.etcdLocation) > 0 {
		overriddenEtcdLocations = exactResourceOverride.etcdLocation
	}

	etcdPrefix := s.StorageConfig.Prefix
	if len(groupOverride.etcdPrefix) > 0 {
		etcdPrefix = groupOverride.etcdPrefix
	}
	if len(exactResourceOverride.etcdPrefix) > 0 {
		etcdPrefix = exactResourceOverride.etcdPrefix
	}

	etcdMediaType := s.DefaultMediaType
	if len(groupOverride.mediaType) != 0 {
		etcdMediaType = groupOverride.mediaType
	}
	if len(exactResourceOverride.mediaType) != 0 {
		etcdMediaType = exactResourceOverride.mediaType
	}

	etcdSerializer := s.DefaultSerializer
	if groupOverride.serializer != nil {
		etcdSerializer = groupOverride.serializer
	}
	if exactResourceOverride.serializer != nil {
		etcdSerializer = exactResourceOverride.serializer
	}
	// operate on copy
	config := s.StorageConfig
	config.Prefix = etcdPrefix
	if len(overriddenEtcdLocations) > 0 {
		config.ServerList = overriddenEtcdLocations
	}

	storageEncodingVersion, err := s.ResourceEncodingConfig.StorageEncodingFor(chosenStorageResource)
	if err != nil {
		return nil, err
	}
	internalVersion, err := s.ResourceEncodingConfig.InMemoryEncodingFor(groupResource)
	if err != nil {
		return nil, err
	}

	codec, err := s.newStorageCodecFn(etcdMediaType, etcdSerializer, storageEncodingVersion, internalVersion, config)
	if err != nil {
		return nil, err
	}

	glog.V(3).Infof("storing %v in %v, reading as %v from %v", groupResource, storageEncodingVersion, internalVersion, config)
	config.Codec = codec
	return &config, nil
}

// Get all backends for all registered storage destinations.
// Used for getting all instances for health validations.
func (s *DefaultStorageFactory) Backends() []string {
	backends := sets.NewString(s.StorageConfig.ServerList...)

	for _, overrides := range s.Overrides {
		backends.Insert(overrides.etcdLocation...)
	}
	return backends.List()
}

// NewStorageCodec assembles a storage codec for the provided storage media type, the provided serializer, and the requested
// storage and memory versions.
func NewStorageCodec(storageMediaType string, ns runtime.StorageSerializer, storageVersion, memoryVersion unversioned.GroupVersion, config storagebackend.Config) (runtime.Codec, error) {
	mediaType, _, err := mime.ParseMediaType(storageMediaType)
	if err != nil {
		return nil, fmt.Errorf("%q is not a valid mime-type", storageMediaType)
	}
	serializer, ok := runtime.SerializerInfoForMediaType(ns.SupportedMediaTypes(), mediaType)
	if !ok {
		return nil, fmt.Errorf("unable to find serializer for %q", storageMediaType)
	}

	s := serializer.Serializer

	// etcd2 only supports string data - we must wrap any result before returning
	// TODO: storagebackend should return a boolean indicating whether it supports binary data
	if !serializer.EncodesAsText && (config.Type == storagebackend.StorageTypeUnset || config.Type == storagebackend.StorageTypeETCD2) {
		glog.V(4).Infof("Wrapping the underlying binary storage serializer with a base64 encoding for etcd2")
		s = runtime.NewBase64Serializer(s)
	}

	encoder := ns.EncoderForVersion(
		s,
		runtime.NewMultiGroupVersioner(
			storageVersion,
			unversioned.GroupKind{Group: storageVersion.Group},
			unversioned.GroupKind{Group: memoryVersion.Group},
		),
	)

	ds := recognizer.NewDecoder(s, ns.UniversalDeserializer())
	decoder := ns.DecoderToVersion(
		ds,
		runtime.NewMultiGroupVersioner(
			memoryVersion,
			unversioned.GroupKind{Group: memoryVersion.Group},
			unversioned.GroupKind{Group: storageVersion.Group},
		),
	)

	return runtime.NewCodec(encoder, decoder), nil
}

var specialDefaultResourcePrefixes = map[unversioned.GroupResource]string{
	unversioned.GroupResource{Group: "", Resource: "replicationControllers"}: "controllers",
	unversioned.GroupResource{Group: "", Resource: "replicationcontrollers"}: "controllers",
	unversioned.GroupResource{Group: "", Resource: "endpoints"}:              "services/endpoints",
	unversioned.GroupResource{Group: "", Resource: "nodes"}:                  "minions",
	unversioned.GroupResource{Group: "", Resource: "services"}:               "services/specs",
	unversioned.GroupResource{Group: "extensions", Resource: "ingresses"}:    "ingress",
}

func (s *DefaultStorageFactory) ResourcePrefix(groupResource unversioned.GroupResource) string {
	chosenStorageResource := s.getStorageGroupResource(groupResource)
	groupOverride := s.Overrides[getAllResourcesAlias(chosenStorageResource)]
	exactResourceOverride := s.Overrides[chosenStorageResource]

	etcdResourcePrefix := specialDefaultResourcePrefixes[chosenStorageResource]
	if len(groupOverride.etcdResourcePrefix) > 0 {
		etcdResourcePrefix = groupOverride.etcdResourcePrefix
	}
	if len(exactResourceOverride.etcdResourcePrefix) > 0 {
		etcdResourcePrefix = exactResourceOverride.etcdResourcePrefix
	}
	if len(etcdResourcePrefix) == 0 {
		etcdResourcePrefix = strings.ToLower(chosenStorageResource.Resource)
	}

	return etcdResourcePrefix
}
<commit_msg>Storage factory should not hardcode special resources<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 genericapiserver

import (
	"fmt"
	"mime"
	"strings"

	"k8s.io/kubernetes/pkg/api/unversioned"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/runtime/serializer/recognizer"
	"k8s.io/kubernetes/pkg/storage/storagebackend"
	"k8s.io/kubernetes/pkg/util/sets"

	"github.com/golang/glog"
)

// StorageFactory is the interface to locate the storage for a given GroupResource
type StorageFactory interface {
	// New finds the storage destination for the given group and resource. It will
	// return an error if the group has no storage destination configured.
	NewConfig(groupResource unversioned.GroupResource) (*storagebackend.Config, error)

	// ResourcePrefix returns the overridden resource prefix for the GroupResource
	// This allows for cohabitation of resources with different native types and provides
	// centralized control over the shape of etcd directories
	ResourcePrefix(groupResource unversioned.GroupResource) string

	// Backends gets all backends for all registered storage destinations.
	// Used for getting all instances for health validations.
	Backends() []string
}

// DefaultStorageFactory takes a GroupResource and returns back its storage interface.  This result includes:
// 1. Merged etcd config, including: auth, server locations, prefixes
// 2. Resource encodings for storage: group,version,kind to store as
// 3. Cohabitating default: some resources like hpa are exposed through multiple APIs.  They must agree on 1 and 2
type DefaultStorageFactory struct {
	// StorageConfig describes how to create a storage backend in general.
	// Its authentication information will be used for every storage.Interface returned.
	StorageConfig storagebackend.Config

	Overrides map[unversioned.GroupResource]groupResourceOverrides

	DefaultResourcePrefixes map[unversioned.GroupResource]string

	// DefaultMediaType is the media type used to store resources. If it is not set, "application/json" is used.
	DefaultMediaType string

	// DefaultSerializer is used to create encoders and decoders for the storage.Interface.
	DefaultSerializer runtime.StorageSerializer

	// ResourceEncodingConfig describes how to encode a particular GroupVersionResource
	ResourceEncodingConfig ResourceEncodingConfig

	// APIResourceConfigSource indicates whether the *storage* is enabled, NOT the API
	// This is discrete from resource enablement because those are separate concerns.  How this source is configured
	// is left to the caller.
	APIResourceConfigSource APIResourceConfigSource

	// newStorageCodecFn exists to be overwritten for unit testing.
	newStorageCodecFn func(storageMediaType string, ns runtime.StorageSerializer, storageVersion, memoryVersion unversioned.GroupVersion, config storagebackend.Config) (codec runtime.Codec, err error)
}

type groupResourceOverrides struct {
	// etcdLocation contains the list of "special" locations that are used for particular GroupResources
	// These are merged on top of the StorageConfig when requesting the storage.Interface for a given GroupResource
	etcdLocation []string
	// etcdPrefix is the base location for a GroupResource.
	etcdPrefix string
	// etcdResourcePrefix is the location to use to store a particular type under the `etcdPrefix` location
	// If empty, the default mapping is used.  If the default mapping doesn't contain an entry, it will use
	// the ToLowered name of the resource, not including the group.
	etcdResourcePrefix string
	// mediaType is the desired serializer to choose. If empty, the default is chosen.
	mediaType string
	// serializer contains the list of "special" serializers for a GroupResource.  Resource=* means for the entire group
	serializer runtime.StorageSerializer
	// cohabitatingResources keeps track of which resources must be stored together.  This happens when we have multiple ways
	// of exposing one set of concepts.  autoscaling.HPA and extensions.HPA as a for instance
	// The order of the slice matters!  It is the priority order of lookup for finding a storage location
	cohabitatingResources []unversioned.GroupResource
}

var _ StorageFactory = &DefaultStorageFactory{}

const AllResources = "*"

// specialDefaultResourcePrefixes are prefixes compiled into Kubernetes.
// TODO: move out of this package, it is not generic
var specialDefaultResourcePrefixes = map[unversioned.GroupResource]string{
	unversioned.GroupResource{Group: "", Resource: "replicationControllers"}: "controllers",
	unversioned.GroupResource{Group: "", Resource: "replicationcontrollers"}: "controllers",
	unversioned.GroupResource{Group: "", Resource: "endpoints"}:              "services/endpoints",
	unversioned.GroupResource{Group: "", Resource: "nodes"}:                  "minions",
	unversioned.GroupResource{Group: "", Resource: "services"}:               "services/specs",
	unversioned.GroupResource{Group: "extensions", Resource: "ingresses"}:    "ingress",
}

func NewDefaultStorageFactory(config storagebackend.Config, defaultMediaType string, defaultSerializer runtime.StorageSerializer, resourceEncodingConfig ResourceEncodingConfig, resourceConfig APIResourceConfigSource) *DefaultStorageFactory {
	if len(defaultMediaType) == 0 {
		defaultMediaType = runtime.ContentTypeJSON
	}
	return &DefaultStorageFactory{
		StorageConfig:           config,
		Overrides:               map[unversioned.GroupResource]groupResourceOverrides{},
		DefaultMediaType:        defaultMediaType,
		DefaultSerializer:       defaultSerializer,
		ResourceEncodingConfig:  resourceEncodingConfig,
		APIResourceConfigSource: resourceConfig,
		DefaultResourcePrefixes: specialDefaultResourcePrefixes,

		newStorageCodecFn: NewStorageCodec,
	}
}

func (s *DefaultStorageFactory) SetEtcdLocation(groupResource unversioned.GroupResource, location []string) {
	overrides := s.Overrides[groupResource]
	overrides.etcdLocation = location
	s.Overrides[groupResource] = overrides
}

func (s *DefaultStorageFactory) SetEtcdPrefix(groupResource unversioned.GroupResource, prefix string) {
	overrides := s.Overrides[groupResource]
	overrides.etcdPrefix = prefix
	s.Overrides[groupResource] = overrides
}

// SetResourceEtcdPrefix sets the prefix for a resource, but not the base-dir.  You'll end up in `etcdPrefix/resourceEtcdPrefix`.
func (s *DefaultStorageFactory) SetResourceEtcdPrefix(groupResource unversioned.GroupResource, prefix string) {
	overrides := s.Overrides[groupResource]
	overrides.etcdResourcePrefix = prefix
	s.Overrides[groupResource] = overrides
}

func (s *DefaultStorageFactory) SetSerializer(groupResource unversioned.GroupResource, mediaType string, serializer runtime.StorageSerializer) {
	overrides := s.Overrides[groupResource]
	overrides.mediaType = mediaType
	overrides.serializer = serializer
	s.Overrides[groupResource] = overrides
}

// AddCohabitatingResources links resources together the order of the slice matters!  its the priority order of lookup for finding a storage location
func (s *DefaultStorageFactory) AddCohabitatingResources(groupResources ...unversioned.GroupResource) {
	for _, groupResource := range groupResources {
		overrides := s.Overrides[groupResource]
		overrides.cohabitatingResources = groupResources
		s.Overrides[groupResource] = overrides
	}
}

func getAllResourcesAlias(resource unversioned.GroupResource) unversioned.GroupResource {
	return unversioned.GroupResource{Group: resource.Group, Resource: AllResources}
}

func (s *DefaultStorageFactory) getStorageGroupResource(groupResource unversioned.GroupResource) unversioned.GroupResource {
	for _, potentialStorageResource := range s.Overrides[groupResource].cohabitatingResources {
		if s.APIResourceConfigSource.AnyVersionOfResourceEnabled(potentialStorageResource) {
			return potentialStorageResource
		}
	}

	return groupResource
}

// New finds the storage destination for the given group and resource. It will
// return an error if the group has no storage destination configured.
func (s *DefaultStorageFactory) NewConfig(groupResource unversioned.GroupResource) (*storagebackend.Config, error) {
	chosenStorageResource := s.getStorageGroupResource(groupResource)

	groupOverride := s.Overrides[getAllResourcesAlias(chosenStorageResource)]
	exactResourceOverride := s.Overrides[chosenStorageResource]

	overriddenEtcdLocations := []string{}
	if len(groupOverride.etcdLocation) > 0 {
		overriddenEtcdLocations = groupOverride.etcdLocation
	}
	if len(exactResourceOverride.etcdLocation) > 0 {
		overriddenEtcdLocations = exactResourceOverride.etcdLocation
	}

	etcdPrefix := s.StorageConfig.Prefix
	if len(groupOverride.etcdPrefix) > 0 {
		etcdPrefix = groupOverride.etcdPrefix
	}
	if len(exactResourceOverride.etcdPrefix) > 0 {
		etcdPrefix = exactResourceOverride.etcdPrefix
	}

	etcdMediaType := s.DefaultMediaType
	if len(groupOverride.mediaType) != 0 {
		etcdMediaType = groupOverride.mediaType
	}
	if len(exactResourceOverride.mediaType) != 0 {
		etcdMediaType = exactResourceOverride.mediaType
	}

	etcdSerializer := s.DefaultSerializer
	if groupOverride.serializer != nil {
		etcdSerializer = groupOverride.serializer
	}
	if exactResourceOverride.serializer != nil {
		etcdSerializer = exactResourceOverride.serializer
	}
	// operate on copy
	config := s.StorageConfig
	config.Prefix = etcdPrefix
	if len(overriddenEtcdLocations) > 0 {
		config.ServerList = overriddenEtcdLocations
	}

	storageEncodingVersion, err := s.ResourceEncodingConfig.StorageEncodingFor(chosenStorageResource)
	if err != nil {
		return nil, err
	}
	internalVersion, err := s.ResourceEncodingConfig.InMemoryEncodingFor(groupResource)
	if err != nil {
		return nil, err
	}

	codec, err := s.newStorageCodecFn(etcdMediaType, etcdSerializer, storageEncodingVersion, internalVersion, config)
	if err != nil {
		return nil, err
	}

	glog.V(3).Infof("storing %v in %v, reading as %v from %v", groupResource, storageEncodingVersion, internalVersion, config)
	config.Codec = codec
	return &config, nil
}

// Get all backends for all registered storage destinations.
// Used for getting all instances for health validations.
func (s *DefaultStorageFactory) Backends() []string {
	backends := sets.NewString(s.StorageConfig.ServerList...)

	for _, overrides := range s.Overrides {
		backends.Insert(overrides.etcdLocation...)
	}
	return backends.List()
}

// NewStorageCodec assembles a storage codec for the provided storage media type, the provided serializer, and the requested
// storage and memory versions.
func NewStorageCodec(storageMediaType string, ns runtime.StorageSerializer, storageVersion, memoryVersion unversioned.GroupVersion, config storagebackend.Config) (runtime.Codec, error) {
	mediaType, _, err := mime.ParseMediaType(storageMediaType)
	if err != nil {
		return nil, fmt.Errorf("%q is not a valid mime-type", storageMediaType)
	}
	serializer, ok := runtime.SerializerInfoForMediaType(ns.SupportedMediaTypes(), mediaType)
	if !ok {
		return nil, fmt.Errorf("unable to find serializer for %q", storageMediaType)
	}

	s := serializer.Serializer

	// etcd2 only supports string data - we must wrap any result before returning
	// TODO: storagebackend should return a boolean indicating whether it supports binary data
	if !serializer.EncodesAsText && (config.Type == storagebackend.StorageTypeUnset || config.Type == storagebackend.StorageTypeETCD2) {
		glog.V(4).Infof("Wrapping the underlying binary storage serializer with a base64 encoding for etcd2")
		s = runtime.NewBase64Serializer(s)
	}

	encoder := ns.EncoderForVersion(
		s,
		runtime.NewMultiGroupVersioner(
			storageVersion,
			unversioned.GroupKind{Group: storageVersion.Group},
			unversioned.GroupKind{Group: memoryVersion.Group},
		),
	)

	ds := recognizer.NewDecoder(s, ns.UniversalDeserializer())
	decoder := ns.DecoderToVersion(
		ds,
		runtime.NewMultiGroupVersioner(
			memoryVersion,
			unversioned.GroupKind{Group: memoryVersion.Group},
			unversioned.GroupKind{Group: storageVersion.Group},
		),
	)

	return runtime.NewCodec(encoder, decoder), nil
}

func (s *DefaultStorageFactory) ResourcePrefix(groupResource unversioned.GroupResource) string {
	chosenStorageResource := s.getStorageGroupResource(groupResource)
	groupOverride := s.Overrides[getAllResourcesAlias(chosenStorageResource)]
	exactResourceOverride := s.Overrides[chosenStorageResource]

	etcdResourcePrefix := s.DefaultResourcePrefixes[chosenStorageResource]
	if len(groupOverride.etcdResourcePrefix) > 0 {
		etcdResourcePrefix = groupOverride.etcdResourcePrefix
	}
	if len(exactResourceOverride.etcdResourcePrefix) > 0 {
		etcdResourcePrefix = exactResourceOverride.etcdResourcePrefix
	}
	if len(etcdResourcePrefix) == 0 {
		etcdResourcePrefix = strings.ToLower(chosenStorageResource.Resource)
	}

	return etcdResourcePrefix
}
<|endoftext|>


<commit_before>package util

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"runtime/debug"
	"strings"
	"sync"
	"unsafe"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	gzip "github.com/klauspost/pgzip"
	"go.etcd.io/bbolt"

	"github.com/grafana/loki/pkg/storage/chunk"
	"github.com/grafana/loki/pkg/storage/chunk/local"
	util_log "github.com/grafana/loki/pkg/util/log"
)

const (
	delimiter = "/"
	sep       = "\xff"
)

var (
	gzipReader = sync.Pool{}
	gzipWriter = sync.Pool{}
)

// getGzipReader gets or creates a new CompressionReader and reset it to read from src
func getGzipReader(src io.Reader) io.Reader {
	if r := gzipReader.Get(); r != nil {
		reader := r.(*gzip.Reader)
		err := reader.Reset(src)
		if err != nil {
			panic(err)
		}
		return reader
	}
	reader, err := gzip.NewReader(src)
	if err != nil {
		panic(err)
	}
	return reader
}

// putGzipReader places back in the pool a CompressionReader
func putGzipReader(reader io.Reader) {
	gzipReader.Put(reader)
}

// getGzipWriter gets or creates a new CompressionWriter and reset it to write to dst
func getGzipWriter(dst io.Writer) io.WriteCloser {
	if w := gzipWriter.Get(); w != nil {
		writer := w.(*gzip.Writer)
		writer.Reset(dst)
		return writer
	}
	return gzip.NewWriter(dst)
}

// PutWriter places back in the pool a CompressionWriter
func putGzipWriter(writer io.WriteCloser) {
	gzipWriter.Put(writer)
}

type IndexStorageClient interface {
	GetFile(ctx context.Context, tableName, fileName string) (io.ReadCloser, error)
	GetUserFile(ctx context.Context, tableName, userID, fileName string) (io.ReadCloser, error)
}

type GetFileFunc func() (io.ReadCloser, error)

// DownloadFileFromStorage downloads a file from storage to given location.
func DownloadFileFromStorage(destination string, decompressFile bool, sync bool, logger log.Logger, getFileFunc GetFileFunc) error {
	readCloser, err := getFileFunc()
	if err != nil {
		return err
	}

	defer func() {
		if err := readCloser.Close(); err != nil {
			level.Error(logger).Log("msg", "failed to close read closer", "err", err)
		}
	}()

	f, err := os.Create(destination)
	if err != nil {
		return err
	}

	defer func() {
		if err := f.Close(); err != nil {
			level.Warn(logger).Log("msg", "failed to close file", "file", destination)
		}
	}()
	var objectReader io.Reader = readCloser
	if decompressFile {
		decompressedReader := getGzipReader(readCloser)
		defer putGzipReader(decompressedReader)

		objectReader = decompressedReader
	}

	_, err = io.Copy(f, objectReader)
	if err != nil {
		return err
	}

	level.Info(logger).Log("msg", "downloaded file")
	if sync {
		return f.Sync()
	}
	return nil
}

func BuildIndexFileName(tableName, uploader, dbName string) string {
	// Files are stored with <uploader>-<db-name>
	objectKey := fmt.Sprintf("%s-%s", uploader, dbName)

	// if the file is a migrated one then don't add its name to the object key otherwise we would re-upload them again here with a different name.
	if tableName == dbName {
		objectKey = uploader
	}

	return objectKey
}

func CompressFile(src, dest string, sync bool) error {
	level.Info(util_log.Logger).Log("msg", "compressing the file", "src", src, "dest", dest)
	uncompressedFile, err := os.Open(src)
	if err != nil {
		return err
	}

	defer func() {
		if err := uncompressedFile.Close(); err != nil {
			level.Error(util_log.Logger).Log("msg", "failed to close uncompressed file", "path", src, "err", err)
		}
	}()

	compressedFile, err := os.Create(dest)
	if err != nil {
		return err
	}

	defer func() {
		if err := compressedFile.Close(); err != nil {
			level.Error(util_log.Logger).Log("msg", "failed to close compressed file", "path", dest, "err", err)
		}
	}()

	compressedWriter := getGzipWriter(compressedFile)
	defer putGzipWriter(compressedWriter)

	_, err = io.Copy(compressedWriter, uncompressedFile)
	if err != nil {
		return err
	}

	err = compressedWriter.Close()
	if err == nil {
		return err
	}
	if sync {
		return compressedFile.Sync()
	}
	return nil
}

type result struct {
	boltdb *bbolt.DB
	err    error
}

// SafeOpenBoltdbFile will recover from a panic opening a DB file, and return the panic message in the err return object.
func SafeOpenBoltdbFile(path string) (*bbolt.DB, error) {
	result := make(chan *result)
	// Open the file in a separate goroutine because we want to change
	// the behavior of a Fault for just this operation and not for the
	// calling goroutine
	go safeOpenBoltDbFile(path, result)
	res := <-result
	return res.boltdb, res.err
}

func safeOpenBoltDbFile(path string, ret chan *result) {
	// boltdb can throw faults which are not caught by recover unless we turn them into panics
	debug.SetPanicOnFault(true)
	res := &result{}

	defer func() {
		if r := recover(); r != nil {
			res.err = fmt.Errorf("recovered from panic opening boltdb file: %v", r)
		}

		// Return the result object on the channel to unblock the calling thread
		ret <- res
	}()

	b, err := local.OpenBoltdbFile(path)
	res.boltdb = b
	res.err = err
}

func ValidateSharedStoreKeyPrefix(prefix string) error {
	if prefix == "" {
		return errors.New("shared store key prefix must be set")
	} else if strings.Contains(prefix, "\\") {
		// When using windows filesystem as object store the implementation of ObjectClient in Cortex takes care of conversion of separator.
		// We just need to always use `/` as a path separator.
		return fmt.Errorf("shared store key prefix should only have '%s' as a path separator", delimiter)
	} else if strings.HasPrefix(prefix, delimiter) {
		return errors.New("shared store key prefix should never start with a path separator i.e '/'")
	} else if !strings.HasSuffix(prefix, delimiter) {
		return errors.New("shared store key prefix should end with a path separator i.e '/'")
	}

	return nil
}

func QueryKey(q chunk.IndexQuery) string {
	ret := q.TableName + sep + q.HashValue

	if len(q.RangeValuePrefix) != 0 {
		ret += sep + string(q.RangeValuePrefix)
	}

	if len(q.RangeValueStart) != 0 {
		ret += sep + string(q.RangeValueStart)
	}

	if len(q.ValueEqual) != 0 {
		ret += sep + string(q.ValueEqual)
	}

	return ret
}

func IsCompressedFile(filename string) bool {
	return strings.HasSuffix(filename, ".gz")
}

func LoggerWithFilename(logger log.Logger, filename string) log.Logger {
	return log.With(logger, "file-name", filename)
}

func GetUnsafeBytes(s string) []byte {
	return *((*[]byte)(unsafe.Pointer(&s)))
}

func GetUnsafeString(buf []byte) string {
	return *((*string)(unsafe.Pointer(&buf)))
}
<commit_msg>add some timing to the logs for downloading and opening index tables (#5042)<commit_after>package util

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"runtime/debug"
	"strings"
	"sync"
	"time"
	"unsafe"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	gzip "github.com/klauspost/pgzip"
	"go.etcd.io/bbolt"

	"github.com/grafana/loki/pkg/storage/chunk"
	"github.com/grafana/loki/pkg/storage/chunk/local"
	util_log "github.com/grafana/loki/pkg/util/log"
)

const (
	delimiter = "/"
	sep       = "\xff"
)

var (
	gzipReader = sync.Pool{}
	gzipWriter = sync.Pool{}
)

// getGzipReader gets or creates a new CompressionReader and reset it to read from src
func getGzipReader(src io.Reader) io.Reader {
	if r := gzipReader.Get(); r != nil {
		reader := r.(*gzip.Reader)
		err := reader.Reset(src)
		if err != nil {
			panic(err)
		}
		return reader
	}
	reader, err := gzip.NewReader(src)
	if err != nil {
		panic(err)
	}
	return reader
}

// putGzipReader places back in the pool a CompressionReader
func putGzipReader(reader io.Reader) {
	gzipReader.Put(reader)
}

// getGzipWriter gets or creates a new CompressionWriter and reset it to write to dst
func getGzipWriter(dst io.Writer) io.WriteCloser {
	if w := gzipWriter.Get(); w != nil {
		writer := w.(*gzip.Writer)
		writer.Reset(dst)
		return writer
	}
	return gzip.NewWriter(dst)
}

// PutWriter places back in the pool a CompressionWriter
func putGzipWriter(writer io.WriteCloser) {
	gzipWriter.Put(writer)
}

type IndexStorageClient interface {
	GetFile(ctx context.Context, tableName, fileName string) (io.ReadCloser, error)
	GetUserFile(ctx context.Context, tableName, userID, fileName string) (io.ReadCloser, error)
}

type GetFileFunc func() (io.ReadCloser, error)

// DownloadFileFromStorage downloads a file from storage to given location.
func DownloadFileFromStorage(destination string, decompressFile bool, sync bool, logger log.Logger, getFileFunc GetFileFunc) error {
	start := time.Now()
	readCloser, err := getFileFunc()
	if err != nil {
		return err
	}

	defer func() {
		if err := readCloser.Close(); err != nil {
			level.Error(logger).Log("msg", "failed to close read closer", "err", err)
		}
	}()

	f, err := os.Create(destination)
	if err != nil {
		return err
	}

	defer func() {
		if err := f.Close(); err != nil {
			level.Warn(logger).Log("msg", "failed to close file", "file", destination)
		}
	}()
	var objectReader io.Reader = readCloser
	if decompressFile {
		decompressedReader := getGzipReader(readCloser)
		defer putGzipReader(decompressedReader)

		objectReader = decompressedReader
	}

	_, err = io.Copy(f, objectReader)
	if err != nil {
		return err
	}

	level.Info(logger).Log("msg", "downloaded file", "total_time", time.Since(start))
	if sync {
		return f.Sync()
	}
	return nil
}

func BuildIndexFileName(tableName, uploader, dbName string) string {
	// Files are stored with <uploader>-<db-name>
	objectKey := fmt.Sprintf("%s-%s", uploader, dbName)

	// if the file is a migrated one then don't add its name to the object key otherwise we would re-upload them again here with a different name.
	if tableName == dbName {
		objectKey = uploader
	}

	return objectKey
}

func CompressFile(src, dest string, sync bool) error {
	level.Info(util_log.Logger).Log("msg", "compressing the file", "src", src, "dest", dest)
	uncompressedFile, err := os.Open(src)
	if err != nil {
		return err
	}

	defer func() {
		if err := uncompressedFile.Close(); err != nil {
			level.Error(util_log.Logger).Log("msg", "failed to close uncompressed file", "path", src, "err", err)
		}
	}()

	compressedFile, err := os.Create(dest)
	if err != nil {
		return err
	}

	defer func() {
		if err := compressedFile.Close(); err != nil {
			level.Error(util_log.Logger).Log("msg", "failed to close compressed file", "path", dest, "err", err)
		}
	}()

	compressedWriter := getGzipWriter(compressedFile)
	defer putGzipWriter(compressedWriter)

	_, err = io.Copy(compressedWriter, uncompressedFile)
	if err != nil {
		return err
	}

	err = compressedWriter.Close()
	if err == nil {
		return err
	}
	if sync {
		return compressedFile.Sync()
	}
	return nil
}

type result struct {
	boltdb *bbolt.DB
	err    error
}

// SafeOpenBoltdbFile will recover from a panic opening a DB file, and return the panic message in the err return object.
func SafeOpenBoltdbFile(path string) (*bbolt.DB, error) {
	result := make(chan *result)
	// Open the file in a separate goroutine because we want to change
	// the behavior of a Fault for just this operation and not for the
	// calling goroutine
	go safeOpenBoltDbFile(path, result)
	res := <-result
	return res.boltdb, res.err
}

func safeOpenBoltDbFile(path string, ret chan *result) {
	// boltdb can throw faults which are not caught by recover unless we turn them into panics
	debug.SetPanicOnFault(true)
	res := &result{}

	defer func() {
		if r := recover(); r != nil {
			res.err = fmt.Errorf("recovered from panic opening boltdb file: %v", r)
		}

		// Return the result object on the channel to unblock the calling thread
		ret <- res
	}()

	b, err := local.OpenBoltdbFile(path)
	res.boltdb = b
	res.err = err
}

func ValidateSharedStoreKeyPrefix(prefix string) error {
	if prefix == "" {
		return errors.New("shared store key prefix must be set")
	} else if strings.Contains(prefix, "\\") {
		// When using windows filesystem as object store the implementation of ObjectClient in Cortex takes care of conversion of separator.
		// We just need to always use `/` as a path separator.
		return fmt.Errorf("shared store key prefix should only have '%s' as a path separator", delimiter)
	} else if strings.HasPrefix(prefix, delimiter) {
		return errors.New("shared store key prefix should never start with a path separator i.e '/'")
	} else if !strings.HasSuffix(prefix, delimiter) {
		return errors.New("shared store key prefix should end with a path separator i.e '/'")
	}

	return nil
}

func QueryKey(q chunk.IndexQuery) string {
	ret := q.TableName + sep + q.HashValue

	if len(q.RangeValuePrefix) != 0 {
		ret += sep + string(q.RangeValuePrefix)
	}

	if len(q.RangeValueStart) != 0 {
		ret += sep + string(q.RangeValueStart)
	}

	if len(q.ValueEqual) != 0 {
		ret += sep + string(q.ValueEqual)
	}

	return ret
}

func IsCompressedFile(filename string) bool {
	return strings.HasSuffix(filename, ".gz")
}

func LoggerWithFilename(logger log.Logger, filename string) log.Logger {
	return log.With(logger, "file-name", filename)
}

func GetUnsafeBytes(s string) []byte {
	return *((*[]byte)(unsafe.Pointer(&s)))
}

func GetUnsafeString(buf []byte) string {
	return *((*string)(unsafe.Pointer(&buf)))
}
<|endoftext|>