<commit_before>/*
Copyright 2015 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 protobuf

import (
	"fmt"
	"reflect"
	"strings"

	"k8s.io/gengo/generator"
	"k8s.io/gengo/namer"
	"k8s.io/gengo/types"
)

type localNamer struct {
	localPackage types.Name
}

func (n localNamer) Name(t *types.Type) string {
	if t.Key != nil && t.Elem != nil {
		return fmt.Sprintf("map<%s, %s>", n.Name(t.Key), n.Name(t.Elem))
	}
	if len(n.localPackage.Package) != 0 && n.localPackage.Package == t.Name.Package {
		return t.Name.Name
	}
	return t.Name.String()
}

type protobufNamer struct {
	packages       []*protobufPackage
	packagesByPath map[string]*protobufPackage
}

func NewProtobufNamer() *protobufNamer {
	return &protobufNamer{
		packagesByPath: make(map[string]*protobufPackage),
	}
}

func (n *protobufNamer) Name(t *types.Type) string {
	if t.Kind == types.Map {
		return fmt.Sprintf("map<%s, %s>", n.Name(t.Key), n.Name(t.Elem))
	}
	return t.Name.String()
}

func (n *protobufNamer) List() []generator.Package {
	packages := make([]generator.Package, 0, len(n.packages))
	for i := range n.packages {
		packages = append(packages, n.packages[i])
	}
	return packages
}

func (n *protobufNamer) Add(p *protobufPackage) {
	if _, ok := n.packagesByPath[p.PackagePath]; !ok {
		n.packagesByPath[p.PackagePath] = p
		n.packages = append(n.packages, p)
	}
}

func (n *protobufNamer) GoNameToProtoName(name types.Name) types.Name {
	if p, ok := n.packagesByPath[name.Package]; ok {
		return types.Name{
			Name:    name.Name,
			Package: p.PackageName,
			Path:    p.ImportPath(),
		}
	}
	for _, p := range n.packages {
		if _, ok := p.FilterTypes[name]; ok {
			return types.Name{
				Name:    name.Name,
				Package: p.PackageName,
				Path:    p.ImportPath(),
			}
		}
	}
	return types.Name{Name: name.Name}
}

func protoSafePackage(name string) string {
	pkg := strings.Replace(name, "/", ".", -1)
	return strings.Replace(pkg, "-", "_", -1)
}

type typeNameSet map[types.Name]*protobufPackage

// assignGoTypeToProtoPackage looks for Go and Protobuf types that are referenced by a type in
// a package. It will not recurse into protobuf types.
func assignGoTypeToProtoPackage(p *protobufPackage, t *types.Type, local, global typeNameSet, optional map[types.Name]struct{}) {
	newT, isProto := isFundamentalProtoType(t)
	if isProto {
		t = newT
	}
	if otherP, ok := global[t.Name]; ok {
		if _, ok := local[t.Name]; !ok {
			p.Imports.AddType(&types.Type{
				Kind: types.Protobuf,
				Name: otherP.ProtoTypeName(),
			})
		}
		return
	}
	global[t.Name] = p
	if _, ok := local[t.Name]; ok {
		return
	}
	// don't recurse into existing proto types
	if isProto {
		p.Imports.AddType(t)
		return
	}

	local[t.Name] = p
	for _, m := range t.Members {
		if namer.IsPrivateGoName(m.Name) {
			continue
		}
		field := &protoField{}
		tag := reflect.StructTag(m.Tags).Get("protobuf")
		if tag == "-" {
			continue
		}
		if err := protobufTagToField(tag, field, m, t, p.ProtoTypeName()); err == nil && field.Type != nil {
			assignGoTypeToProtoPackage(p, field.Type, local, global, optional)
			continue
		}
		assignGoTypeToProtoPackage(p, m.Type, local, global, optional)
	}
	// TODO: should methods be walked?
	if t.Elem != nil {
		assignGoTypeToProtoPackage(p, t.Elem, local, global, optional)
	}
	if t.Key != nil {
		assignGoTypeToProtoPackage(p, t.Key, local, global, optional)
	}
	if t.Underlying != nil {
		if t.Kind == types.Alias && isOptionalAlias(t) {
			optional[t.Name] = struct{}{}
		}
		assignGoTypeToProtoPackage(p, t.Underlying, local, global, optional)
	}
}

func (n *protobufNamer) AssignTypesToPackages(c *generator.Context) error {
	global := make(typeNameSet)
	for _, p := range n.packages {
		local := make(typeNameSet)
		optional := make(map[types.Name]struct{})
		p.Imports = NewImportTracker(p.ProtoTypeName())
		for _, t := range c.Order {
			if t.Name.Package != p.PackagePath {
				continue
			}
			assignGoTypeToProtoPackage(p, t, local, global, optional)
		}
		p.FilterTypes = make(map[types.Name]struct{})
		p.LocalNames = make(map[string]struct{})
		p.OptionalTypeNames = make(map[string]struct{})
		for k, v := range local {
			if v == p {
				p.FilterTypes[k] = struct{}{}
				p.LocalNames[k.Name] = struct{}{}
				if _, ok := optional[k]; ok {
					p.OptionalTypeNames[k.Name] = struct{}{}
				}
			}
		}
	}
	return nil
}
<commit_msg>[go-to-protobuf] Skip private types and functions<commit_after>/*
Copyright 2015 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 protobuf

import (
	"fmt"
	"reflect"
	"strings"

	"k8s.io/gengo/generator"
	"k8s.io/gengo/namer"
	"k8s.io/gengo/types"
)

type localNamer struct {
	localPackage types.Name
}

func (n localNamer) Name(t *types.Type) string {
	if t.Key != nil && t.Elem != nil {
		return fmt.Sprintf("map<%s, %s>", n.Name(t.Key), n.Name(t.Elem))
	}
	if len(n.localPackage.Package) != 0 && n.localPackage.Package == t.Name.Package {
		return t.Name.Name
	}
	return t.Name.String()
}

type protobufNamer struct {
	packages       []*protobufPackage
	packagesByPath map[string]*protobufPackage
}

func NewProtobufNamer() *protobufNamer {
	return &protobufNamer{
		packagesByPath: make(map[string]*protobufPackage),
	}
}

func (n *protobufNamer) Name(t *types.Type) string {
	if t.Kind == types.Map {
		return fmt.Sprintf("map<%s, %s>", n.Name(t.Key), n.Name(t.Elem))
	}
	return t.Name.String()
}

func (n *protobufNamer) List() []generator.Package {
	packages := make([]generator.Package, 0, len(n.packages))
	for i := range n.packages {
		packages = append(packages, n.packages[i])
	}
	return packages
}

func (n *protobufNamer) Add(p *protobufPackage) {
	if _, ok := n.packagesByPath[p.PackagePath]; !ok {
		n.packagesByPath[p.PackagePath] = p
		n.packages = append(n.packages, p)
	}
}

func (n *protobufNamer) GoNameToProtoName(name types.Name) types.Name {
	if p, ok := n.packagesByPath[name.Package]; ok {
		return types.Name{
			Name:    name.Name,
			Package: p.PackageName,
			Path:    p.ImportPath(),
		}
	}
	for _, p := range n.packages {
		if _, ok := p.FilterTypes[name]; ok {
			return types.Name{
				Name:    name.Name,
				Package: p.PackageName,
				Path:    p.ImportPath(),
			}
		}
	}
	return types.Name{Name: name.Name}
}

func protoSafePackage(name string) string {
	pkg := strings.Replace(name, "/", ".", -1)
	return strings.Replace(pkg, "-", "_", -1)
}

type typeNameSet map[types.Name]*protobufPackage

// assignGoTypeToProtoPackage looks for Go and Protobuf types that are referenced by a type in
// a package. It will not recurse into protobuf types.
func assignGoTypeToProtoPackage(p *protobufPackage, t *types.Type, local, global typeNameSet, optional map[types.Name]struct{}) {
	newT, isProto := isFundamentalProtoType(t)
	if isProto {
		t = newT
	}
	if otherP, ok := global[t.Name]; ok {
		if _, ok := local[t.Name]; !ok {
			p.Imports.AddType(&types.Type{
				Kind: types.Protobuf,
				Name: otherP.ProtoTypeName(),
			})
		}
		return
	}
	global[t.Name] = p
	if _, ok := local[t.Name]; ok {
		return
	}
	// don't recurse into existing proto types
	if isProto {
		p.Imports.AddType(t)
		return
	}

	local[t.Name] = p
	for _, m := range t.Members {
		if namer.IsPrivateGoName(m.Name) {
			continue
		}
		field := &protoField{}
		tag := reflect.StructTag(m.Tags).Get("protobuf")
		if tag == "-" {
			continue
		}
		if err := protobufTagToField(tag, field, m, t, p.ProtoTypeName()); err == nil && field.Type != nil {
			assignGoTypeToProtoPackage(p, field.Type, local, global, optional)
			continue
		}
		assignGoTypeToProtoPackage(p, m.Type, local, global, optional)
	}
	// TODO: should methods be walked?
	if t.Elem != nil {
		assignGoTypeToProtoPackage(p, t.Elem, local, global, optional)
	}
	if t.Key != nil {
		assignGoTypeToProtoPackage(p, t.Key, local, global, optional)
	}
	if t.Underlying != nil {
		if t.Kind == types.Alias && isOptionalAlias(t) {
			optional[t.Name] = struct{}{}
		}
		assignGoTypeToProtoPackage(p, t.Underlying, local, global, optional)
	}
}

// isTypeApplicableToProtobuf checks to see if a type is relevant for protobuf processing.
// Currently, it filters out functions and private types.
func isTypeApplicableToProtobuf(t *types.Type) bool {
	// skip functions -- we don't care about them for protobuf
	if t.Kind == types.Func || (t.Kind == types.DeclarationOf && t.Underlying.Kind == types.Func) {
		return false
	}
	// skip private types
	if namer.IsPrivateGoName(t.Name.Name) {
		return false
	}

	return true
}

func (n *protobufNamer) AssignTypesToPackages(c *generator.Context) error {
	global := make(typeNameSet)
	for _, p := range n.packages {
		local := make(typeNameSet)
		optional := make(map[types.Name]struct{})
		p.Imports = NewImportTracker(p.ProtoTypeName())
		for _, t := range c.Order {
			if t.Name.Package != p.PackagePath {
				continue
			}
			if !isTypeApplicableToProtobuf(t) {
				// skip types that we don't care about, like functions
				continue
			}
			assignGoTypeToProtoPackage(p, t, local, global, optional)
		}
		p.FilterTypes = make(map[types.Name]struct{})
		p.LocalNames = make(map[string]struct{})
		p.OptionalTypeNames = make(map[string]struct{})
		for k, v := range local {
			if v == p {
				p.FilterTypes[k] = struct{}{}
				p.LocalNames[k.Name] = struct{}{}
				if _, ok := optional[k]; ok {
					p.OptionalTypeNames[k.Name] = struct{}{}
				}
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package gengorums

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/printer"
	"go/token"
	"go/types"
	"log"
	"os"
	"path"
	"sort"
	"strings"

	"github.com/google/go-cmp/cmp"
	"golang.org/x/tools/go/packages"
)

const (
	devPkgPath = "./cmd/protoc-gen-gorums/dev"
)

// GenerateBundleFile generates a file with static definitions for Gorums.
func GenerateBundleFile(dst string) {
	pkgIdent, reservedIdents, code := bundle(devPkgPath)
	// escape backticks
	escaped := strings.ReplaceAll(string(code), "`", "`+\"`\"+`")
	src := "// Code generated by protoc-gen-gorums. DO NOT EDIT.\n" +
		"// Source files can be found in: " + devPkgPath + "\n\n" +
		"package gengorums\n\n" +
		generatePkgMap(pkgIdent, reservedIdents) +
		"var staticCode = `" + escaped + "`\n"

	staticContent, err := format.Source([]byte(src))
	if err != nil {
		log.Fatalf("formatting failed: %v", err)
	}
	currentContent, err := os.ReadFile(dst)
	if err != nil {
		log.Fatal(err)
	}
	if diff := cmp.Diff(currentContent, staticContent); diff != "" {
		fmt.Fprintf(os.Stderr, "change detected (-current +new):\n%s", diff)
		fmt.Fprintf(os.Stderr, "\nReview changes above; to revert use:\n")
		fmt.Fprintf(os.Stderr, "mv %s.bak %s\n", dst, dst)
	}
	err = os.WriteFile(dst, []byte(staticContent), 0666)
	if err != nil {
		log.Fatal(err)
	}
}

func generatePkgMap(pkgs map[string]string, reservedIdents []string) string {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, `
	// pkgIdentMap maps from package name to one of the package's identifiers.
	// These identifiers are used by the Gorums protoc plugin to generate
	// appropriate import statements.
	var pkgIdentMap = %#v`+"\n", pkgs)

	fmt.Fprintf(&buf, `
	// reservedIdents holds the set of Gorums reserved identifiers.
	// These identifiers cannot be used to define message types in a proto file.
	var reservedIdents = %#v`+"\n\n", reservedIdents)
	return buf.String()
}

// findIdentifiers examines the given package to find all imported packages,
// and one used identifier in that imported package. These identifiers are
// used by the Gorums protoc plugin to generate appropriate import statements.
func findIdentifiers(fset *token.FileSet, pkgInfo *packages.Package) (map[string]string, []string) {
	pkgIdents := make(map[string][]string)
	for id, obj := range pkgInfo.TypesInfo.Uses {
		pos := fset.Position(id.Pos())
		if ignore(pos.Filename) || !obj.Exported() {
			// ignore identifiers in zorums generated files and unexported identifiers
			continue
		}
		if pkg := obj.Pkg(); pkg != nil {
			switch obj := obj.(type) {
			case *types.Func:
				if typ := obj.Type(); typ != nil {
					if recv := typ.(*types.Signature).Recv(); recv != nil {
						// ignore functions on non-package types
						continue
					}
				}
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())

			case *types.Const, *types.TypeName:
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())

			case *types.Var:
				// no need to import obj's pkg if obj is a field of a struct.
				if obj.IsField() {
					continue
				}
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())
			}
		}
	}
	// Only need to store one identifier for each imported package.
	// However, to ensure stable output, we sort the identifiers
	// and use the first element.
	pkgIdent := make(map[string]string, len(pkgIdents))
	var reservedIdents []string
	for path, idents := range pkgIdents {
		sort.Strings(idents)
		if path == pkgInfo.PkgPath {
			reservedIdents = idents
			continue
		}
		pkgIdent[path] = idents[0]
	}
	return pkgIdent, reservedIdents
}

func addUniqueIdentifier(pkgIdents map[string][]string, path, name string) {
	currentNames := pkgIdents[path]
	for _, known := range currentNames {
		if name == known {
			return
		}
	}
	pkgIdents[path] = append(pkgIdents[path], name)
}

// bundle returns a slice with the code for the given package without imports.
// The returned map contains packages to be imported along with one identifier
// using the relevant import path.
func bundle(pkgPath string) (map[string]string, []string, []byte) {
	cfg := &packages.Config{
		Mode: packages.NeedName | packages.NeedFiles | packages.NeedSyntax | packages.NeedImports | packages.NeedTypes | packages.NeedTypesInfo,
	}
	pkgs, err := packages.Load(cfg, pkgPath)
	if err != nil {
		log.Fatalf("failed to load Gorums dev package: %v", err)
	}
	if packages.PrintErrors(pkgs) > 0 {
		os.Exit(1)
	}
	// Since Load succeeded and pkgPath is a single package, the following is safe
	pkg := pkgs[0]
	var out bytes.Buffer
	printFiles(&out, pkg.Fset, pkg.Syntax)
	pkgIdentMap, reservedIdents := findIdentifiers(pkg.Fset, pkg)
	debug("pkgIdentMap=%v", pkgIdentMap)
	debug("reservedIdents=%v", reservedIdents)
	return pkgIdentMap, reservedIdents, out.Bytes()
}

func debug(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", args...)
}

func printFiles(out *bytes.Buffer, fset *token.FileSet, files []*ast.File) {
	for _, f := range files {
		// filter files in dev package that shouldn't be bundled in template_static.go
		fileName := fset.File(f.Pos()).Name()
		if ignore(fileName) {
			continue
		}

		last := f.Package
		if len(f.Imports) > 0 {
			imp := f.Imports[len(f.Imports)-1]
			last = imp.End()
			if imp.Comment != nil {
				if e := imp.Comment.End(); e > last {
					last = e
				}
			}
		}

		// Pretty-print package-level declarations.
		// but no package or import declarations.
		var buf bytes.Buffer
		for _, decl := range f.Decls {
			if decl, ok := decl.(*ast.GenDecl); ok && decl.Tok == token.IMPORT {
				continue
			}
			beg, end := sourceRange(decl)
			printComments(out, f.Comments, last, beg)

			buf.Reset()
			err := format.Node(&buf, fset, &printer.CommentedNode{Node: decl, Comments: f.Comments})
			if err != nil {
				log.Fatalf("failed to format source AST node: %v", err)
			}
			out.Write(buf.Bytes())
			last = printSameLineComment(out, f.Comments, fset, end)
			out.WriteString("\n\n")
		}
		printLastComments(out, f.Comments, last)
	}
}

// ignore files in dev folder with suffixes that shouldn't be bundled.
func ignore(file string) bool {
	for _, suffix := range []string{".proto", "_test.go"} {
		if strings.HasSuffix(file, suffix) {
			return true
		}
	}
	base := path.Base(file)
	// ignore zorums* files
	return strings.HasPrefix(base, "zorums")
}

// sourceRange returns the [beg, end) interval of source code
// belonging to decl (incl. associated comments).
func sourceRange(decl ast.Decl) (beg, end token.Pos) {
	beg = decl.Pos()
	end = decl.End()

	var doc, com *ast.CommentGroup

	switch d := decl.(type) {
	case *ast.GenDecl:
		doc = d.Doc
		if len(d.Specs) > 0 {
			switch spec := d.Specs[len(d.Specs)-1].(type) {
			case *ast.ValueSpec:
				com = spec.Comment
			case *ast.TypeSpec:
				com = spec.Comment
			}
		}
	case *ast.FuncDecl:
		doc = d.Doc
	}

	if doc != nil {
		beg = doc.Pos()
	}
	if com != nil && com.End() > end {
		end = com.End()
	}

	return beg, end
}

func printComments(out *bytes.Buffer, comments []*ast.CommentGroup, pos, end token.Pos) {
	for _, cg := range comments {
		if pos <= cg.Pos() && cg.Pos() < end {
			for _, c := range cg.List {
				fmt.Fprintln(out, c.Text)
			}
			fmt.Fprintln(out)
		}
	}
}

const infinity = 1 << 30

func printLastComments(out *bytes.Buffer, comments []*ast.CommentGroup, pos token.Pos) {
	printComments(out, comments, pos, infinity)
}

func printSameLineComment(out *bytes.Buffer, comments []*ast.CommentGroup, fset *token.FileSet, pos token.Pos) token.Pos {
	tf := fset.File(pos)
	for _, cg := range comments {
		if pos <= cg.Pos() && tf.Line(cg.Pos()) == tf.Line(pos) {
			for _, c := range cg.List {
				fmt.Fprintln(out, c.Text)
			}
			return cg.End()
		}
	}
	return pos
}
<commit_msg>Simplified src generation to a single fmt.Sprintf<commit_after>package gengorums

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/printer"
	"go/token"
	"go/types"
	"log"
	"os"
	"path"
	"sort"
	"strings"

	"github.com/google/go-cmp/cmp"
	"golang.org/x/tools/go/packages"
)

const (
	devPkgPath = "./cmd/protoc-gen-gorums/dev"
)

// GenerateBundleFile generates a file with static definitions for Gorums.
func GenerateBundleFile(dst string) {
	pkgIdent, reservedIdents, code := bundle(devPkgPath)
	// escape backticks
	escaped := strings.ReplaceAll(string(code), "`", "`+\"`\"+`")
	src := fmt.Sprintf(`// Code generated by protoc-gen-gorums. DO NOT EDIT.
	// Source files can be found in: %s

	package gengorums

	// pkgIdentMap maps from package name to one of the package's identifiers.
	// These identifiers are used by the Gorums protoc plugin to generate
	// appropriate import statements.
	var pkgIdentMap = %#v

	// reservedIdents holds the set of Gorums reserved identifiers.
	// These identifiers cannot be used to define message types in a proto file.
	var reservedIdents = %#v

	var staticCode = `+"`%s`", devPkgPath, pkgIdent, reservedIdents, escaped)

	staticContent, err := format.Source([]byte(src))
	if err != nil {
		log.Fatalf("formatting failed: %v", err)
	}
	currentContent, err := os.ReadFile(dst)
	if err != nil {
		log.Fatal(err)
	}
	if diff := cmp.Diff(currentContent, staticContent); diff != "" {
		fmt.Fprintf(os.Stderr, "change detected (-current +new):\n%s", diff)
		fmt.Fprintf(os.Stderr, "\nReview changes above; to revert use:\n")
		fmt.Fprintf(os.Stderr, "mv %s.bak %s\n", dst, dst)
	}
	err = os.WriteFile(dst, []byte(staticContent), 0666)
	if err != nil {
		log.Fatal(err)
	}
}

// findIdentifiers examines the given package to find all imported packages,
// and one used identifier in that imported package. These identifiers are
// used by the Gorums protoc plugin to generate appropriate import statements.
func findIdentifiers(fset *token.FileSet, pkgInfo *packages.Package) (map[string]string, []string) {
	pkgIdents := make(map[string][]string)
	for id, obj := range pkgInfo.TypesInfo.Uses {
		pos := fset.Position(id.Pos())
		if ignore(pos.Filename) || !obj.Exported() {
			// ignore identifiers in zorums generated files and unexported identifiers
			continue
		}
		if pkg := obj.Pkg(); pkg != nil {
			switch obj := obj.(type) {
			case *types.Func:
				if typ := obj.Type(); typ != nil {
					if recv := typ.(*types.Signature).Recv(); recv != nil {
						// ignore functions on non-package types
						continue
					}
				}
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())

			case *types.Const, *types.TypeName:
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())

			case *types.Var:
				// no need to import obj's pkg if obj is a field of a struct.
				if obj.IsField() {
					continue
				}
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())
			}
		}
	}
	// Only need to store one identifier for each imported package.
	// However, to ensure stable output, we sort the identifiers
	// and use the first element.
	pkgIdent := make(map[string]string, len(pkgIdents))
	var reservedIdents []string
	for path, idents := range pkgIdents {
		sort.Strings(idents)
		if path == pkgInfo.PkgPath {
			reservedIdents = idents
			continue
		}
		pkgIdent[path] = idents[0]
	}
	return pkgIdent, reservedIdents
}

func addUniqueIdentifier(pkgIdents map[string][]string, path, name string) {
	currentNames := pkgIdents[path]
	for _, known := range currentNames {
		if name == known {
			return
		}
	}
	pkgIdents[path] = append(pkgIdents[path], name)
}

// bundle returns a slice with the code for the given package without imports.
// The returned map contains packages to be imported along with one identifier
// using the relevant import path.
func bundle(pkgPath string) (map[string]string, []string, []byte) {
	cfg := &packages.Config{
		Mode: packages.NeedName | packages.NeedFiles | packages.NeedSyntax | packages.NeedImports | packages.NeedTypes | packages.NeedTypesInfo,
	}
	pkgs, err := packages.Load(cfg, pkgPath)
	if err != nil {
		log.Fatalf("failed to load Gorums dev package: %v", err)
	}
	if packages.PrintErrors(pkgs) > 0 {
		os.Exit(1)
	}
	// Since Load succeeded and pkgPath is a single package, the following is safe
	pkg := pkgs[0]
	var out bytes.Buffer
	printFiles(&out, pkg.Fset, pkg.Syntax)
	pkgIdentMap, reservedIdents := findIdentifiers(pkg.Fset, pkg)
	debug("pkgIdentMap=%v", pkgIdentMap)
	debug("reservedIdents=%v", reservedIdents)
	return pkgIdentMap, reservedIdents, out.Bytes()
}

func debug(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", args...)
}

func printFiles(out *bytes.Buffer, fset *token.FileSet, files []*ast.File) {
	for _, f := range files {
		// filter files in dev package that shouldn't be bundled in template_static.go
		fileName := fset.File(f.Pos()).Name()
		if ignore(fileName) {
			continue
		}

		last := f.Package
		if len(f.Imports) > 0 {
			imp := f.Imports[len(f.Imports)-1]
			last = imp.End()
			if imp.Comment != nil {
				if e := imp.Comment.End(); e > last {
					last = e
				}
			}
		}

		// Pretty-print package-level declarations.
		// but no package or import declarations.
		var buf bytes.Buffer
		for _, decl := range f.Decls {
			if decl, ok := decl.(*ast.GenDecl); ok && decl.Tok == token.IMPORT {
				continue
			}
			beg, end := sourceRange(decl)
			printComments(out, f.Comments, last, beg)

			buf.Reset()
			err := format.Node(&buf, fset, &printer.CommentedNode{Node: decl, Comments: f.Comments})
			if err != nil {
				log.Fatalf("failed to format source AST node: %v", err)
			}
			out.Write(buf.Bytes())
			last = printSameLineComment(out, f.Comments, fset, end)
			out.WriteString("\n\n")
		}
		printLastComments(out, f.Comments, last)
	}
}

// ignore files in dev folder with suffixes that shouldn't be bundled.
func ignore(file string) bool {
	for _, suffix := range []string{".proto", "_test.go"} {
		if strings.HasSuffix(file, suffix) {
			return true
		}
	}
	base := path.Base(file)
	// ignore zorums* files
	return strings.HasPrefix(base, "zorums")
}

// sourceRange returns the [beg, end) interval of source code
// belonging to decl (incl. associated comments).
func sourceRange(decl ast.Decl) (beg, end token.Pos) {
	beg = decl.Pos()
	end = decl.End()

	var doc, com *ast.CommentGroup

	switch d := decl.(type) {
	case *ast.GenDecl:
		doc = d.Doc
		if len(d.Specs) > 0 {
			switch spec := d.Specs[len(d.Specs)-1].(type) {
			case *ast.ValueSpec:
				com = spec.Comment
			case *ast.TypeSpec:
				com = spec.Comment
			}
		}
	case *ast.FuncDecl:
		doc = d.Doc
	}

	if doc != nil {
		beg = doc.Pos()
	}
	if com != nil && com.End() > end {
		end = com.End()
	}

	return beg, end
}

func printComments(out *bytes.Buffer, comments []*ast.CommentGroup, pos, end token.Pos) {
	for _, cg := range comments {
		if pos <= cg.Pos() && cg.Pos() < end {
			for _, c := range cg.List {
				fmt.Fprintln(out, c.Text)
			}
			fmt.Fprintln(out)
		}
	}
}

const infinity = 1 << 30

func printLastComments(out *bytes.Buffer, comments []*ast.CommentGroup, pos token.Pos) {
	printComments(out, comments, pos, infinity)
}

func printSameLineComment(out *bytes.Buffer, comments []*ast.CommentGroup, fset *token.FileSet, pos token.Pos) token.Pos {
	tf := fset.File(pos)
	for _, cg := range comments {
		if pos <= cg.Pos() && tf.Line(cg.Pos()) == tf.Line(pos) {
			for _, c := range cg.List {
				fmt.Fprintln(out, c.Text)
			}
			return cg.End()
		}
	}
	return pos
}
<|endoftext|>


<commit_before>package gengorums

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/printer"
	"go/token"
	"go/types"
	"log"
	"os"
	"path"
	"sort"
	"strings"

	"github.com/google/go-cmp/cmp"
	"golang.org/x/tools/go/packages"
)

const (
	devPkgPath = "./cmd/protoc-gen-gorums/dev"
)

// GenerateBundleFile generates a file with static definitions for Gorums.
func GenerateBundleFile(dst string) {
	pkgIdent, reservedIdents, code := bundle(devPkgPath)
	// escape backticks
	escaped := strings.ReplaceAll(string(code), "`", "`+\"`\"+`")
	src := fmt.Sprintf(`// Code generated by protoc-gen-gorums. DO NOT EDIT.
	// Source files can be found in: %s

	package gengorums

	// pkgIdentMap maps from package name to one of the package's identifiers.
	// These identifiers are used by the Gorums protoc plugin to generate
	// appropriate import statements.
	var pkgIdentMap = %#v

	// reservedIdents holds the set of Gorums reserved identifiers.
	// These identifiers cannot be used to define message types in a proto file.
	var reservedIdents = %#v

	var staticCode = `+"`%s`", devPkgPath, pkgIdent, reservedIdents, escaped)

	staticContent, err := format.Source([]byte(src))
	if err != nil {
		log.Fatalf("formatting failed: %v", err)
	}
	currentContent, err := os.ReadFile(dst)
	if err != nil {
		log.Fatal(err)
	}
	if diff := cmp.Diff(currentContent, staticContent); diff != "" {
		fmt.Fprintf(os.Stderr, "change detected (-current +new):\n%s", diff)
		fmt.Fprintf(os.Stderr, "\nReview changes above; to revert use:\n")
		fmt.Fprintf(os.Stderr, "mv %s.bak %s\n", dst, dst)
	}
	err = os.WriteFile(dst, []byte(staticContent), 0666)
	if err != nil {
		log.Fatal(err)
	}
}

// findIdentifiers examines the given package to find all imported packages,
// and one used identifier in that imported package. These identifiers are
// used by the Gorums protoc plugin to generate appropriate import statements.
func findIdentifiers(fset *token.FileSet, pkgInfo *packages.Package) (map[string]string, []string) {
	pkgIdents := make(map[string][]string)
	for id, obj := range pkgInfo.TypesInfo.Uses {
		pos := fset.Position(id.Pos())
		if ignore(pos.Filename) || !obj.Exported() {
			// ignore identifiers in zorums generated files and unexported identifiers
			continue
		}
		if pkg := obj.Pkg(); pkg != nil {
			switch obj := obj.(type) {
			case *types.Func:
				if typ := obj.Type(); typ != nil {
					if recv := typ.(*types.Signature).Recv(); recv != nil {
						// ignore functions on non-package types
						continue
					}
				}
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())

			case *types.Const, *types.TypeName:
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())

			case *types.Var:
				// no need to import obj's pkg if obj is a field of a struct.
				if obj.IsField() {
					continue
				}
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())
			}
		}
	}
	// Only need to store one identifier for each imported package.
	// However, to ensure stable output, we sort the identifiers
	// and use the first element.
	pkgIdent := make(map[string]string, len(pkgIdents))
	var reservedIdents []string
	for path, idents := range pkgIdents {
		sort.Strings(idents)
		if path == pkgInfo.PkgPath {
			reservedIdents = idents
			continue
		}
		pkgIdent[path] = idents[0]
	}
	return pkgIdent, reservedIdents
}

func addUniqueIdentifier(pkgIdents map[string][]string, path, name string) {
	currentNames := pkgIdents[path]
	for _, known := range currentNames {
		if name == known {
			return
		}
	}
	pkgIdents[path] = append(pkgIdents[path], name)
}

// bundle returns a slice with the code for the given package without imports.
// The returned map contains packages to be imported along with one identifier
// using the relevant import path.
func bundle(pkgPath string) (map[string]string, []string, []byte) {
	cfg := &packages.Config{
		Mode: packages.NeedName | packages.NeedFiles | packages.NeedSyntax | packages.NeedImports | packages.NeedTypes | packages.NeedTypesInfo,
	}
	pkgs, err := packages.Load(cfg, pkgPath)
	if err != nil {
		log.Fatalf("failed to load Gorums dev package: %v", err)
	}
	if packages.PrintErrors(pkgs) > 0 {
		os.Exit(1)
	}
	// Since Load succeeded and pkgPath is a single package, the following is safe
	pkg := pkgs[0]
	src := printFiles(pkg)
	pkgIdentMap, reservedIdents := findIdentifiers(pkg.Fset, pkg)
	debug("pkgIdentMap=%v", pkgIdentMap)
	debug("reservedIdents=%v", reservedIdents)
	return pkgIdentMap, reservedIdents, src
}

func debug(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", args...)
}

func printFiles(pkg *packages.Package) []byte {
	out := &bytes.Buffer{}
	for _, f := range pkg.Syntax {
		// filter files in dev package that shouldn't be bundled in template_static.go
		fileName := pkg.Fset.File(f.Pos()).Name()
		if ignore(fileName) {
			continue
		}

		last := f.Package
		if len(f.Imports) > 0 {
			imp := f.Imports[len(f.Imports)-1]
			last = imp.End()
			if imp.Comment != nil {
				if e := imp.Comment.End(); e > last {
					last = e
				}
			}
		}

		// Pretty-print package-level declarations.
		// but no package or import declarations.
		var buf bytes.Buffer
		for _, decl := range f.Decls {
			if decl, ok := decl.(*ast.GenDecl); ok && decl.Tok == token.IMPORT {
				continue
			}
			beg, end := sourceRange(decl)
			printComments(out, f.Comments, last, beg)

			buf.Reset()
			err := format.Node(&buf, pkg.Fset, &printer.CommentedNode{Node: decl, Comments: f.Comments})
			if err != nil {
				log.Fatalf("failed to format source AST node: %v", err)
			}
			out.Write(buf.Bytes())
			last = printSameLineComment(out, f.Comments, pkg.Fset, end)
			out.WriteString("\n\n")
		}
		printLastComments(out, f.Comments, last)
	}
	return out.Bytes()
}

// ignore files in dev folder with suffixes that shouldn't be bundled.
func ignore(file string) bool {
	for _, suffix := range []string{".proto", "_test.go"} {
		if strings.HasSuffix(file, suffix) {
			return true
		}
	}
	base := path.Base(file)
	// ignore zorums* files
	return strings.HasPrefix(base, "zorums")
}

// sourceRange returns the [beg, end) interval of source code
// belonging to decl (incl. associated comments).
func sourceRange(decl ast.Decl) (beg, end token.Pos) {
	beg = decl.Pos()
	end = decl.End()

	var doc, com *ast.CommentGroup

	switch d := decl.(type) {
	case *ast.GenDecl:
		doc = d.Doc
		if len(d.Specs) > 0 {
			switch spec := d.Specs[len(d.Specs)-1].(type) {
			case *ast.ValueSpec:
				com = spec.Comment
			case *ast.TypeSpec:
				com = spec.Comment
			}
		}
	case *ast.FuncDecl:
		doc = d.Doc
	}

	if doc != nil {
		beg = doc.Pos()
	}
	if com != nil && com.End() > end {
		end = com.End()
	}

	return beg, end
}

func printComments(out *bytes.Buffer, comments []*ast.CommentGroup, pos, end token.Pos) {
	for _, cg := range comments {
		if pos <= cg.Pos() && cg.Pos() < end {
			for _, c := range cg.List {
				fmt.Fprintln(out, c.Text)
			}
			fmt.Fprintln(out)
		}
	}
}

const infinity = 1 << 30

func printLastComments(out *bytes.Buffer, comments []*ast.CommentGroup, pos token.Pos) {
	printComments(out, comments, pos, infinity)
}

func printSameLineComment(out *bytes.Buffer, comments []*ast.CommentGroup, fset *token.FileSet, pos token.Pos) token.Pos {
	tf := fset.File(pos)
	for _, cg := range comments {
		if pos <= cg.Pos() && tf.Line(cg.Pos()) == tf.Line(pos) {
			for _, c := range cg.List {
				fmt.Fprintln(out, c.Text)
			}
			return cg.End()
		}
	}
	return pos
}
<commit_msg>Added doc for printFiles<commit_after>package gengorums

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/printer"
	"go/token"
	"go/types"
	"log"
	"os"
	"path"
	"sort"
	"strings"

	"github.com/google/go-cmp/cmp"
	"golang.org/x/tools/go/packages"
)

const (
	devPkgPath = "./cmd/protoc-gen-gorums/dev"
)

// GenerateBundleFile generates a file with static definitions for Gorums.
func GenerateBundleFile(dst string) {
	pkgIdent, reservedIdents, code := bundle(devPkgPath)
	// escape backticks
	escaped := strings.ReplaceAll(string(code), "`", "`+\"`\"+`")
	src := fmt.Sprintf(`// Code generated by protoc-gen-gorums. DO NOT EDIT.
	// Source files can be found in: %s

	package gengorums

	// pkgIdentMap maps from package name to one of the package's identifiers.
	// These identifiers are used by the Gorums protoc plugin to generate
	// appropriate import statements.
	var pkgIdentMap = %#v

	// reservedIdents holds the set of Gorums reserved identifiers.
	// These identifiers cannot be used to define message types in a proto file.
	var reservedIdents = %#v

	var staticCode = `+"`%s`", devPkgPath, pkgIdent, reservedIdents, escaped)

	staticContent, err := format.Source([]byte(src))
	if err != nil {
		log.Fatalf("formatting failed: %v", err)
	}
	currentContent, err := os.ReadFile(dst)
	if err != nil {
		log.Fatal(err)
	}
	if diff := cmp.Diff(currentContent, staticContent); diff != "" {
		fmt.Fprintf(os.Stderr, "change detected (-current +new):\n%s", diff)
		fmt.Fprintf(os.Stderr, "\nReview changes above; to revert use:\n")
		fmt.Fprintf(os.Stderr, "mv %s.bak %s\n", dst, dst)
	}
	err = os.WriteFile(dst, []byte(staticContent), 0666)
	if err != nil {
		log.Fatal(err)
	}
}

// findIdentifiers examines the given package to find all imported packages,
// and one used identifier in that imported package. These identifiers are
// used by the Gorums protoc plugin to generate appropriate import statements.
func findIdentifiers(fset *token.FileSet, pkgInfo *packages.Package) (map[string]string, []string) {
	pkgIdents := make(map[string][]string)
	for id, obj := range pkgInfo.TypesInfo.Uses {
		pos := fset.Position(id.Pos())
		if ignore(pos.Filename) || !obj.Exported() {
			// ignore identifiers in zorums generated files and unexported identifiers
			continue
		}
		if pkg := obj.Pkg(); pkg != nil {
			switch obj := obj.(type) {
			case *types.Func:
				if typ := obj.Type(); typ != nil {
					if recv := typ.(*types.Signature).Recv(); recv != nil {
						// ignore functions on non-package types
						continue
					}
				}
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())

			case *types.Const, *types.TypeName:
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())

			case *types.Var:
				// no need to import obj's pkg if obj is a field of a struct.
				if obj.IsField() {
					continue
				}
				addUniqueIdentifier(pkgIdents, pkg.Path(), obj.Name())
			}
		}
	}
	// Only need to store one identifier for each imported package.
	// However, to ensure stable output, we sort the identifiers
	// and use the first element.
	pkgIdent := make(map[string]string, len(pkgIdents))
	var reservedIdents []string
	for path, idents := range pkgIdents {
		sort.Strings(idents)
		if path == pkgInfo.PkgPath {
			reservedIdents = idents
			continue
		}
		pkgIdent[path] = idents[0]
	}
	return pkgIdent, reservedIdents
}

func addUniqueIdentifier(pkgIdents map[string][]string, path, name string) {
	currentNames := pkgIdents[path]
	for _, known := range currentNames {
		if name == known {
			return
		}
	}
	pkgIdents[path] = append(pkgIdents[path], name)
}

// bundle returns a slice with the code for the given package without imports.
// The returned map contains packages to be imported along with one identifier
// using the relevant import path.
func bundle(pkgPath string) (map[string]string, []string, []byte) {
	cfg := &packages.Config{
		Mode: packages.NeedName | packages.NeedFiles | packages.NeedSyntax | packages.NeedImports | packages.NeedTypes | packages.NeedTypesInfo,
	}
	pkgs, err := packages.Load(cfg, pkgPath)
	if err != nil {
		log.Fatalf("failed to load Gorums dev package: %v", err)
	}
	if packages.PrintErrors(pkgs) > 0 {
		os.Exit(1)
	}
	// Since Load succeeded and pkgPath is a single package, the following is safe
	pkg := pkgs[0]
	src := printFiles(pkg)
	pkgIdentMap, reservedIdents := findIdentifiers(pkg.Fset, pkg)
	debug("pkgIdentMap=%v", pkgIdentMap)
	debug("reservedIdents=%v", reservedIdents)
	return pkgIdentMap, reservedIdents, src
}

func debug(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", args...)
}

// printFiles returns the source files of pkg without import statements.
func printFiles(pkg *packages.Package) []byte {
	out := &bytes.Buffer{}
	for _, f := range pkg.Syntax {
		// filter files in dev package that shouldn't be bundled in template_static.go
		fileName := pkg.Fset.File(f.Pos()).Name()
		if ignore(fileName) {
			continue
		}

		last := f.Package
		if len(f.Imports) > 0 {
			imp := f.Imports[len(f.Imports)-1]
			last = imp.End()
			if imp.Comment != nil {
				if e := imp.Comment.End(); e > last {
					last = e
				}
			}
		}

		// Pretty-print package-level declarations.
		// but no package or import declarations.
		var buf bytes.Buffer
		for _, decl := range f.Decls {
			if decl, ok := decl.(*ast.GenDecl); ok && decl.Tok == token.IMPORT {
				continue
			}
			beg, end := sourceRange(decl)
			printComments(out, f.Comments, last, beg)

			buf.Reset()
			err := format.Node(&buf, pkg.Fset, &printer.CommentedNode{Node: decl, Comments: f.Comments})
			if err != nil {
				log.Fatalf("failed to format source AST node: %v", err)
			}
			out.Write(buf.Bytes())
			last = printSameLineComment(out, f.Comments, pkg.Fset, end)
			out.WriteString("\n\n")
		}
		printLastComments(out, f.Comments, last)
	}
	return out.Bytes()
}

// ignore files in dev folder with suffixes that shouldn't be bundled.
func ignore(file string) bool {
	for _, suffix := range []string{".proto", "_test.go"} {
		if strings.HasSuffix(file, suffix) {
			return true
		}
	}
	base := path.Base(file)
	// ignore zorums* files
	return strings.HasPrefix(base, "zorums")
}

// sourceRange returns the [beg, end) interval of source code
// belonging to decl (incl. associated comments).
func sourceRange(decl ast.Decl) (beg, end token.Pos) {
	beg = decl.Pos()
	end = decl.End()

	var doc, com *ast.CommentGroup

	switch d := decl.(type) {
	case *ast.GenDecl:
		doc = d.Doc
		if len(d.Specs) > 0 {
			switch spec := d.Specs[len(d.Specs)-1].(type) {
			case *ast.ValueSpec:
				com = spec.Comment
			case *ast.TypeSpec:
				com = spec.Comment
			}
		}
	case *ast.FuncDecl:
		doc = d.Doc
	}

	if doc != nil {
		beg = doc.Pos()
	}
	if com != nil && com.End() > end {
		end = com.End()
	}

	return beg, end
}

func printComments(out *bytes.Buffer, comments []*ast.CommentGroup, pos, end token.Pos) {
	for _, cg := range comments {
		if pos <= cg.Pos() && cg.Pos() < end {
			for _, c := range cg.List {
				fmt.Fprintln(out, c.Text)
			}
			fmt.Fprintln(out)
		}
	}
}

const infinity = 1 << 30

func printLastComments(out *bytes.Buffer, comments []*ast.CommentGroup, pos token.Pos) {
	printComments(out, comments, pos, infinity)
}

func printSameLineComment(out *bytes.Buffer, comments []*ast.CommentGroup, fset *token.FileSet, pos token.Pos) token.Pos {
	tf := fset.File(pos)
	for _, cg := range comments {
		if pos <= cg.Pos() && tf.Line(cg.Pos()) == tf.Line(pos) {
			for _, c := range cg.List {
				fmt.Fprintln(out, c.Text)
			}
			return cg.End()
		}
	}
	return pos
}
<|endoftext|>


<commit_before>/*

errors is a package that implements FriendlyError. FriendlyError implements
the error interface, but also has a FriendlyError() method that returns a
message that is reasonable to show to end-users.

Every method that returns a FriendlyError also accepts 0 to n Fields objects
may be provided and they will be combined into a single Fields, with later
values overwriting newer ones.

FriendlyErrors also have With* methods that return a copy of the error with
the given changes. This allows a pattern where a method defines a default
FriendlyError with the friendly error message to use and then when it needs to
return an error uses WithError() to add the specific error message.

*/
package errors

const DefaultFriendlyError = "An error occurred"

type Fields map[string]interface{}

type Friendly struct {
	msg         string
	friendlyMsg string
	secureMsg   string
	fields      Fields
}

//New creates a new errors.Friendly with the given msg.
func New(msg string, fields ...Fields) *Friendly {
	return &Friendly{
		msg:    msg,
		fields: combineFields(fields...),
	}
}

//NewFriendly returns a new errors.Friendly with the given FriendlyError.
func NewFriendly(friendlyMsg string, fields ...Fields) *Friendly {
	return &Friendly{
		friendlyMsg: friendlyMsg,
		fields:      combineFields(fields...),
	}
}

//NewSecure returns a new errors.Friendly with the given SecureError.
func NewSecure(secureMsg string, fields ...Fields) *Friendly {
	return &Friendly{
		secureMsg: secureMsg,
		fields:    combineFields(fields...),
	}
}

func combineFields(fields ...Fields) Fields {
	var result = make(Fields)
	for _, field := range fields {
		for key, val := range field {
			result[key] = val
		}
	}
	return result
}

//SecureError returns the error message that should only be shown in secure
//contexts, because it may include secret information. If no SecureError
//message has been provided, will return Error()
func (f *Friendly) SecureError() string {
	if f.secureMsg == "" {
		return f.Error()
	}
	return f.secureMsg
}

//Error returns the error message, implementing the error interface. If no
//specific message has been provided, will fall back on FriendlyError(). The
//Error() value is OK to show in insecure contexts (i.e. on the client) but it
//just might be confusing to users.
func (f *Friendly) Error() string {
	if f.msg == "" {
		return f.FriendlyError()
	}
	return f.msg
}

//FriendlyError is the error string that is OK to show in insecure contexts to
//end-users. It is generally a much simplified version of the message. If no
//specific FriendlyMessage has been provided, will return DefaultFriendlyError.
func (f *Friendly) FriendlyError() string {
	if f.friendlyMsg == "" {
		return DefaultFriendlyError
	}
	return f.friendlyMsg
}

//Fields returns the Fields object for this error. Fields should only be used
//in secure contexts (e.g. logging on a server) because they may contain
//private information.
func (f *Friendly) Fields() Fields {
	return f.fields
}

//Extend returns a new FriendlyError where the Error() message is prepended
//with this new message and a delimiter. The SecureError and FriendlyError
//message are left untouched.
func (f *Friendly) Extend(msg string, fields ...Fields) *Friendly {
	return &Friendly{
		secureMsg:   f.secureMsg,
		friendlyMsg: f.friendlyMsg,
		msg:         msg + " : " + f.Error(),
		fields:      combineFields(append([]Fields{f.fields}, fields...)...),
	}
}

//WithFriendly returns a copy of err where the friendlyMsg is set to friendlyMsg
func (f *Friendly) WithFriendly(friendlyMsg string, fields ...Fields) *Friendly {
	return &Friendly{
		secureMsg:   f.secureMsg,
		msg:         f.msg,
		friendlyMsg: friendlyMsg,
		fields:      combineFields(append([]Fields{f.fields}, fields...)...),
	}
}

//WithError returns a copy of err where the Error() is set to msg. See a;so
//Extend, which prepends a new message to the front of the existing message.
func (f *Friendly) WithError(msg string, fields ...Fields) *Friendly {
	return &Friendly{
		secureMsg:   f.secureMsg,
		msg:         msg,
		friendlyMsg: f.friendlyMsg,
		fields:      combineFields(append([]Fields{f.fields}, fields...)...),
	}
}

//WithSecure returns a copy of err where the SecureError() is set to secureMsg.
func (f *Friendly) WithSecure(secureMsg string, fields ...Fields) *Friendly {
	return &Friendly{
		secureMsg:   secureMsg,
		msg:         f.msg,
		friendlyMsg: f.friendlyMsg,
		fields:      combineFields(append([]Fields{f.fields}, fields...)...),
	}
}
<commit_msg>Adds NewWrapped convenience method to errors. Part of #243.<commit_after>/*

errors is a package that implements FriendlyError. FriendlyError implements
the error interface, but also has a FriendlyError() method that returns a
message that is reasonable to show to end-users.

Every method that returns a FriendlyError also accepts 0 to n Fields objects
may be provided and they will be combined into a single Fields, with later
values overwriting newer ones.

FriendlyErrors also have With* methods that return a copy of the error with
the given changes. This allows a pattern where a method defines a default
FriendlyError with the friendly error message to use and then when it needs to
return an error uses WithError() to add the specific error message.

*/
package errors

const DefaultFriendlyError = "An error occurred"

type Fields map[string]interface{}

type Friendly struct {
	msg         string
	friendlyMsg string
	secureMsg   string
	fields      Fields
}

//New creates a new errors.Friendly with the given msg.
func New(msg string, fields ...Fields) *Friendly {
	return &Friendly{
		msg:    msg,
		fields: combineFields(fields...),
	}
}

//NewFriendly returns a new errors.Friendly with the given FriendlyError.
func NewFriendly(friendlyMsg string, fields ...Fields) *Friendly {
	return &Friendly{
		friendlyMsg: friendlyMsg,
		fields:      combineFields(fields...),
	}
}

//NewSecure returns a new errors.Friendly with the given SecureError.
func NewSecure(secureMsg string, fields ...Fields) *Friendly {
	return &Friendly{
		secureMsg: secureMsg,
		fields:    combineFields(fields...),
	}
}

//NewWrapped takes a generic error. It returns either a new error with that
//message or, if nil, returns nil.
func NewWrapped(err error, fields ...Fields) *Friendly {
	if err == nil {
		return nil
	}
	return New(err.Error(), fields...)
}

func combineFields(fields ...Fields) Fields {
	var result = make(Fields)
	for _, field := range fields {
		for key, val := range field {
			result[key] = val
		}
	}
	return result
}

//SecureError returns the error message that should only be shown in secure
//contexts, because it may include secret information. If no SecureError
//message has been provided, will return Error()
func (f *Friendly) SecureError() string {
	if f.secureMsg == "" {
		return f.Error()
	}
	return f.secureMsg
}

//Error returns the error message, implementing the error interface. If no
//specific message has been provided, will fall back on FriendlyError(). The
//Error() value is OK to show in insecure contexts (i.e. on the client) but it
//just might be confusing to users.
func (f *Friendly) Error() string {
	if f.msg == "" {
		return f.FriendlyError()
	}
	return f.msg
}

//FriendlyError is the error string that is OK to show in insecure contexts to
//end-users. It is generally a much simplified version of the message. If no
//specific FriendlyMessage has been provided, will return DefaultFriendlyError.
func (f *Friendly) FriendlyError() string {
	if f.friendlyMsg == "" {
		return DefaultFriendlyError
	}
	return f.friendlyMsg
}

//Fields returns the Fields object for this error. Fields should only be used
//in secure contexts (e.g. logging on a server) because they may contain
//private information.
func (f *Friendly) Fields() Fields {
	return f.fields
}

//Extend returns a new FriendlyError where the Error() message is prepended
//with this new message and a delimiter. The SecureError and FriendlyError
//message are left untouched.
func (f *Friendly) Extend(msg string, fields ...Fields) *Friendly {
	return &Friendly{
		secureMsg:   f.secureMsg,
		friendlyMsg: f.friendlyMsg,
		msg:         msg + " : " + f.Error(),
		fields:      combineFields(append([]Fields{f.fields}, fields...)...),
	}
}

//WithFriendly returns a copy of err where the friendlyMsg is set to friendlyMsg
func (f *Friendly) WithFriendly(friendlyMsg string, fields ...Fields) *Friendly {
	return &Friendly{
		secureMsg:   f.secureMsg,
		msg:         f.msg,
		friendlyMsg: friendlyMsg,
		fields:      combineFields(append([]Fields{f.fields}, fields...)...),
	}
}

//WithError returns a copy of err where the Error() is set to msg. See a;so
//Extend, which prepends a new message to the front of the existing message.
func (f *Friendly) WithError(msg string, fields ...Fields) *Friendly {
	return &Friendly{
		secureMsg:   f.secureMsg,
		msg:         msg,
		friendlyMsg: f.friendlyMsg,
		fields:      combineFields(append([]Fields{f.fields}, fields...)...),
	}
}

//WithSecure returns a copy of err where the SecureError() is set to secureMsg.
func (f *Friendly) WithSecure(secureMsg string, fields ...Fields) *Friendly {
	return &Friendly{
		secureMsg:   secureMsg,
		msg:         f.msg,
		friendlyMsg: f.friendlyMsg,
		fields:      combineFields(append([]Fields{f.fields}, fields...)...),
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 The etcd 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 raft

import (
	"bytes"
	"fmt"

	pb "go.etcd.io/etcd/raft/raftpb"
)

func (st StateType) MarshalJSON() ([]byte, error) {
	return []byte(fmt.Sprintf("%q", st.String())), nil
}

// uint64Slice implements sort interface
type uint64Slice []uint64

func (p uint64Slice) Len() int           { return len(p) }
func (p uint64Slice) Less(i, j int) bool { return p[i] < p[j] }
func (p uint64Slice) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }

func min(a, b uint64) uint64 {
	if a > b {
		return b
	}
	return a
}

func max(a, b uint64) uint64 {
	if a > b {
		return a
	}
	return b
}

func IsLocalMsg(msgt pb.MessageType) bool {
	return msgt == pb.MsgHup || msgt == pb.MsgBeat || msgt == pb.MsgUnreachable ||
		msgt == pb.MsgSnapStatus || msgt == pb.MsgCheckQuorum
}

func IsResponseMsg(msgt pb.MessageType) bool {
	return msgt == pb.MsgAppResp || msgt == pb.MsgVoteResp || msgt == pb.MsgHeartbeatResp || msgt == pb.MsgUnreachable || msgt == pb.MsgPreVoteResp
}

// voteResponseType maps vote and prevote message types to their corresponding responses.
func voteRespMsgType(msgt pb.MessageType) pb.MessageType {
	switch msgt {
	case pb.MsgVote:
		return pb.MsgVoteResp
	case pb.MsgPreVote:
		return pb.MsgPreVoteResp
	default:
		panic(fmt.Sprintf("not a vote message: %s", msgt))
	}
}

// EntryFormatter can be implemented by the application to provide human-readable formatting
// of entry data. Nil is a valid EntryFormatter and will use a default format.
type EntryFormatter func([]byte) string

// DescribeMessage returns a concise human-readable description of a
// Message for debugging.
func DescribeMessage(m pb.Message, f EntryFormatter) string {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, "%x->%x %v Term:%d Log:%d/%d", m.From, m.To, m.Type, m.Term, m.LogTerm, m.Index)
	if m.Reject {
		fmt.Fprintf(&buf, " Rejected")
		if m.RejectHint != 0 {
			fmt.Fprintf(&buf, "(Hint:%d)", m.RejectHint)
		}
	}
	if m.Commit != 0 {
		fmt.Fprintf(&buf, " Commit:%d", m.Commit)
	}
	if len(m.Entries) > 0 {
		fmt.Fprintf(&buf, " Entries:[")
		for i, e := range m.Entries {
			if i != 0 {
				buf.WriteString(", ")
			}
			buf.WriteString(DescribeEntry(e, f))
		}
		fmt.Fprintf(&buf, "]")
	}
	if !IsEmptySnap(m.Snapshot) {
		fmt.Fprintf(&buf, " Snapshot:%v", m.Snapshot)
	}
	return buf.String()
}

// PayloadSize is the size of the payload of this Entry. Notably, it does not
// depend on its Index or Term.
func PayloadSize(e pb.Entry) int {
	return len(e.Data)
}

// DescribeEntry returns a concise human-readable description of an
// Entry for debugging.
func DescribeEntry(e pb.Entry, f EntryFormatter) string {
	var formatted string
	if e.Type == pb.EntryNormal && f != nil {
		formatted = f(e.Data)
	} else {
		formatted = fmt.Sprintf("%q", e.Data)
	}
	return fmt.Sprintf("%d/%d %s %s", e.Term, e.Index, e.Type, formatted)
}

// DescribeEntries calls DescribeEntry for each Entry, adding a newline to
// each.
func DescribeEntries(ents []pb.Entry, f EntryFormatter) string {
	var buf bytes.Buffer
	for _, e := range ents {
		_, _ = buf.WriteString(DescribeEntry(e, f) + "\n")
	}
	return buf.String()
}

func limitSize(ents []pb.Entry, maxSize uint64) []pb.Entry {
	if len(ents) == 0 {
		return ents
	}
	size := ents[0].Size()
	var limit int
	for limit = 1; limit < len(ents); limit++ {
		size += ents[limit].Size()
		if uint64(size) > maxSize {
			break
		}
	}
	return ents[:limit]
}
<commit_msg>raft: print RejectHint of zero on MsgAppResp<commit_after>// Copyright 2015 The etcd 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 raft

import (
	"bytes"
	"fmt"

	pb "go.etcd.io/etcd/raft/raftpb"
)

func (st StateType) MarshalJSON() ([]byte, error) {
	return []byte(fmt.Sprintf("%q", st.String())), nil
}

// uint64Slice implements sort interface
type uint64Slice []uint64

func (p uint64Slice) Len() int           { return len(p) }
func (p uint64Slice) Less(i, j int) bool { return p[i] < p[j] }
func (p uint64Slice) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }

func min(a, b uint64) uint64 {
	if a > b {
		return b
	}
	return a
}

func max(a, b uint64) uint64 {
	if a > b {
		return a
	}
	return b
}

func IsLocalMsg(msgt pb.MessageType) bool {
	return msgt == pb.MsgHup || msgt == pb.MsgBeat || msgt == pb.MsgUnreachable ||
		msgt == pb.MsgSnapStatus || msgt == pb.MsgCheckQuorum
}

func IsResponseMsg(msgt pb.MessageType) bool {
	return msgt == pb.MsgAppResp || msgt == pb.MsgVoteResp || msgt == pb.MsgHeartbeatResp || msgt == pb.MsgUnreachable || msgt == pb.MsgPreVoteResp
}

// voteResponseType maps vote and prevote message types to their corresponding responses.
func voteRespMsgType(msgt pb.MessageType) pb.MessageType {
	switch msgt {
	case pb.MsgVote:
		return pb.MsgVoteResp
	case pb.MsgPreVote:
		return pb.MsgPreVoteResp
	default:
		panic(fmt.Sprintf("not a vote message: %s", msgt))
	}
}

// EntryFormatter can be implemented by the application to provide human-readable formatting
// of entry data. Nil is a valid EntryFormatter and will use a default format.
type EntryFormatter func([]byte) string

// DescribeMessage returns a concise human-readable description of a
// Message for debugging.
func DescribeMessage(m pb.Message, f EntryFormatter) string {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, "%x->%x %v Term:%d Log:%d/%d", m.From, m.To, m.Type, m.Term, m.LogTerm, m.Index)
	if m.Reject {
		fmt.Fprintf(&buf, " Rejected (Hint: %d)", m.RejectHint)
	}
	if m.Commit != 0 {
		fmt.Fprintf(&buf, " Commit:%d", m.Commit)
	}
	if len(m.Entries) > 0 {
		fmt.Fprintf(&buf, " Entries:[")
		for i, e := range m.Entries {
			if i != 0 {
				buf.WriteString(", ")
			}
			buf.WriteString(DescribeEntry(e, f))
		}
		fmt.Fprintf(&buf, "]")
	}
	if !IsEmptySnap(m.Snapshot) {
		fmt.Fprintf(&buf, " Snapshot:%v", m.Snapshot)
	}
	return buf.String()
}

// PayloadSize is the size of the payload of this Entry. Notably, it does not
// depend on its Index or Term.
func PayloadSize(e pb.Entry) int {
	return len(e.Data)
}

// DescribeEntry returns a concise human-readable description of an
// Entry for debugging.
func DescribeEntry(e pb.Entry, f EntryFormatter) string {
	var formatted string
	if e.Type == pb.EntryNormal && f != nil {
		formatted = f(e.Data)
	} else {
		formatted = fmt.Sprintf("%q", e.Data)
	}
	return fmt.Sprintf("%d/%d %s %s", e.Term, e.Index, e.Type, formatted)
}

// DescribeEntries calls DescribeEntry for each Entry, adding a newline to
// each.
func DescribeEntries(ents []pb.Entry, f EntryFormatter) string {
	var buf bytes.Buffer
	for _, e := range ents {
		_, _ = buf.WriteString(DescribeEntry(e, f) + "\n")
	}
	return buf.String()
}

func limitSize(ents []pb.Entry, maxSize uint64) []pb.Entry {
	if len(ents) == 0 {
		return ents
	}
	size := ents[0].Size()
	var limit int
	for limit = 1; limit < len(ents); limit++ {
		size += ents[limit].Size()
		if uint64(size) > maxSize {
			break
		}
	}
	return ents[:limit]
}
<|endoftext|>


<commit_before>// +build darwin

package darwin

import "testing"

func TestCPUUsageGenerator_Generate(t *testing.T) {
	g := &CPUUsageGenerator{}
	values, err := g.Generate()

	if err != nil {
		t.Errorf("error should be nil but got: %s", err)
	}

	metricName := []string{"cpu.user.percentage", "cpu.system.percentage", "cpu.idle.percentage"}
	for _, n := range metricName {
		if _, ok := values[n]; !ok {
			t.Errorf("should have '%s': %v", n, values)
		}
	}

	t.Logf("cpu metrics: %+v", values)
}
<commit_msg>add tests for darwin cpu generator<commit_after>// +build darwin

package darwin

import (
	"reflect"
	"testing"

	"github.com/hatena/mackerel-container-agent/metric"
)

func TestCPUUsageGenerator_Generate(t *testing.T) {
	g := &CPUUsageGenerator{}
	values, err := g.Generate()

	if err != nil {
		t.Errorf("error should be nil but got: %s", err)
	}

	metricName := []string{"cpu.user.percentage", "cpu.system.percentage", "cpu.idle.percentage"}
	for _, n := range metricName {
		if _, ok := values[n]; !ok {
			t.Errorf("should have '%s': %v", n, values)
		}
	}

	t.Logf("cpu metrics: %+v", values)
}

func TestCPUUsageGenerator_parseIostatOutput(t *testing.T) {
	testCases := []struct {
		output string
		values metric.Values
	}{
		{
			output: `      cpu    load average
 us sy id   1m   5m   15m
 19  9 72  2.50 3.04 3.20
 16 12 72  2.50 3.04 3.20
`,
			values: metric.Values{
				"cpu.user.percentage":   16.0,
				"cpu.system.percentage": 12.0,
				"cpu.idle.percentage":   72.0,
			},
		},
	}
	for _, testCase := range testCases {
		got, err := parseIostatOutput(testCase.output)
		if err != nil {
			t.Errorf("error should be nil but got: %s", err)
		}
		if !reflect.DeepEqual(map[string]float64(got), map[string]float64(testCase.values)) {
			t.Errorf("metric values should be %#v but got: %#v", testCase.values, got)
		}
	}
}
<|endoftext|>


<commit_before>package models

import (
	"socialapi/workers/common/runner"
	"testing"

	"github.com/koding/bongo"
	. "github.com/smartystreets/goconvey/convey"
)

func TestInteractiongetAccountId(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while getting account id", t, func() {
		Convey("it should have error if interaction id is not set", func() {
			i := NewInteraction()

			in, err := i.getAccountId()
			So(err, ShouldNotBeNil)
			So(err.Error(), ShouldContainSubstring, "couldnt find accountId from content")
			So(in, ShouldEqual, 0)
		})

		Convey("it should have error if account is not set in db", func() {
			i := NewInteraction()
			i.Id = 4590

			in, err := i.getAccountId()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, bongo.RecordNotFound)
			So(in, ShouldEqual, 0)
		})

		Convey("it should return quickly account id if id is set", func() {
			i := NewInteraction()
			i.AccountId = 1020

			in, err := i.getAccountId()
			So(err, ShouldBeNil)
			So(in, ShouldEqual, i.AccountId)
		})
	})
}
<commit_msg>social/tests: isExempt tests are added for interaction<commit_after>package models

import (
	"socialapi/workers/common/runner"
	"testing"

	"github.com/koding/bongo"
	. "github.com/smartystreets/goconvey/convey"
)

func TestInteractiongetAccountId(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("while getting account id", t, func() {
		Convey("it should have error if interaction id is not set", func() {
			i := NewInteraction()

			in, err := i.getAccountId()
			So(err, ShouldNotBeNil)
			So(err.Error(), ShouldContainSubstring, "couldnt find accountId from content")
			So(in, ShouldEqual, 0)
		})

		Convey("it should have error if account is not set in db", func() {
			i := NewInteraction()
			i.Id = 4590

			in, err := i.getAccountId()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, bongo.RecordNotFound)
			So(in, ShouldEqual, 0)
		})

		Convey("it should return quickly account id if id is set", func() {
			i := NewInteraction()
			i.AccountId = 1020

			in, err := i.getAccountId()
			So(err, ShouldBeNil)
			So(in, ShouldEqual, i.AccountId)
		})
	})
}

func TestInteractionisExempt(t *testing.T) {
	r := runner.New("test")
	if err := r.Init(); err != nil {
		t.Fatalf("couldnt start bongo %s", err.Error())
	}
	defer r.Close()

	Convey("While testing interaction is exempt or not", t, func() {
		Convey("it should have error while getting account id from db when channel id is not set", func() {
			i := NewInteraction()

			ie, err := i.isExempt()
			So(err, ShouldNotBeNil)
			So(err.Error(), ShouldContainSubstring, "couldnt find accountId from content")
			So(ie, ShouldEqual, false)
		})

		Convey("it should have error if account id is not set", func() {
			i := NewInteraction()
			i.Id = 1098

			ie, err := i.isExempt()
			So(err, ShouldNotBeNil)
			So(err, ShouldEqual, bongo.RecordNotFound)
			So(ie, ShouldEqual, false)
		})

	})
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Vitess 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 vindexes

import (
	"bytes"
	"fmt"

	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/key"
)

var (
	_ SingleColumn = (*UnicodeLooseXXHash)(nil)
)

// UnicodeLooseXXHash is a vindex that normalizes and hashes unicode strings
// to a keyspace id. It conservatively converts the string to its base
// characters before hashing. This is also known as UCA level 1.
// Ref: http://www.unicode.org/reports/tr10/#Multi_Level_Comparison.
// This is compatible with MySQL's utf8_unicode_ci collation.
type UnicodeLooseXXHash struct {
	name string
}

// NewUnicodeLooseXXHash creates a new UnicodeLooseXXHash struct.
func NewUnicodeLooseXXHash(name string, _ map[string]string) (Vindex, error) {
	return &UnicodeLooseXXHash{name: name}, nil
}

// String returns the name of the vindex.
func (vind *UnicodeLooseXXHash) String() string {
	return vind.name
}

// Cost returns the cost as 1.
func (vind *UnicodeLooseXXHash) Cost() int {
	return 1
}

// IsUnique returns true since the Vindex is unique.
func (vind *UnicodeLooseXXHash) IsUnique() bool {
	return true
}

// NeedsVCursor satisfies the Vindex interface.
func (vind *UnicodeLooseXXHash) NeedsVCursor() bool {
	return false
}

// Verify returns true if ids maps to ksids.
func (vind *UnicodeLooseXXHash) Verify(_ VCursor, ids []sqltypes.Value, ksids [][]byte) ([]bool, error) {
	out := make([]bool, len(ids))
	for i := range ids {
		data, err := unicodeHash(vXXHash, ids[i])
		if err != nil {
			return nil, fmt.Errorf("UnicodeLooseXXHash.Verify: %v", err)
		}
		out[i] = bytes.Equal(data, ksids[i])
	}
	return out, nil
}

// Map can map ids to key.Destination objects.
func (vind *UnicodeLooseXXHash) Map(cursor VCursor, ids []sqltypes.Value) ([]key.Destination, error) {
	out := make([]key.Destination, 0, len(ids))
	for _, id := range ids {
		data, err := unicodeHash(vXXHash, id)
		if err != nil {
			return nil, fmt.Errorf("UnicodeLooseXXHash.Map: %v", err)
		}
		out = append(out, key.DestinationKeyspaceID(data))
	}
	return out, nil
}

func init() {
	Register("unicode_loose_xxhash", NewUnicodeLooseXXHash)
}
<commit_msg>feat: unicodeloosexxhash vindex implemented hashing interface<commit_after>/*
Copyright 2020 The Vitess 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 vindexes

import (
	"bytes"
	"fmt"

	"vitess.io/vitess/go/sqltypes"
	"vitess.io/vitess/go/vt/key"
)

var (
	_ SingleColumn = (*UnicodeLooseXXHash)(nil)
	_ Hashing      = (*UnicodeLooseXXHash)(nil)
)

// UnicodeLooseXXHash is a vindex that normalizes and hashes unicode strings
// to a keyspace id. It conservatively converts the string to its base
// characters before hashing. This is also known as UCA level 1.
// Ref: http://www.unicode.org/reports/tr10/#Multi_Level_Comparison.
// This is compatible with MySQL's utf8_unicode_ci collation.
type UnicodeLooseXXHash struct {
	name string
}

// NewUnicodeLooseXXHash creates a new UnicodeLooseXXHash struct.
func NewUnicodeLooseXXHash(name string, _ map[string]string) (Vindex, error) {
	return &UnicodeLooseXXHash{name: name}, nil
}

// String returns the name of the vindex.
func (vind *UnicodeLooseXXHash) String() string {
	return vind.name
}

// Cost returns the cost as 1.
func (vind *UnicodeLooseXXHash) Cost() int {
	return 1
}

// IsUnique returns true since the Vindex is unique.
func (vind *UnicodeLooseXXHash) IsUnique() bool {
	return true
}

// NeedsVCursor satisfies the Vindex interface.
func (vind *UnicodeLooseXXHash) NeedsVCursor() bool {
	return false
}

// Verify returns true if ids maps to ksids.
func (vind *UnicodeLooseXXHash) Verify(_ VCursor, ids []sqltypes.Value, ksids [][]byte) ([]bool, error) {
	out := make([]bool, 0, len(ids))
	for i, id := range ids {
		data, err := vind.Hash(id)
		if err != nil {
			return nil, fmt.Errorf("UnicodeLooseXXHash.Verify: %v", err)
		}
		out = append(out, bytes.Equal(data, ksids[i]))
	}
	return out, nil
}

// Map can map ids to key.Destination objects.
func (vind *UnicodeLooseXXHash) Map(cursor VCursor, ids []sqltypes.Value) ([]key.Destination, error) {
	out := make([]key.Destination, 0, len(ids))
	for _, id := range ids {
		data, err := vind.Hash(id)
		if err != nil {
			return nil, fmt.Errorf("UnicodeLooseXXHash.Map: %v", err)
		}
		out = append(out, key.DestinationKeyspaceID(data))
	}
	return out, nil
}

func (vind *UnicodeLooseXXHash) Hash(id sqltypes.Value) ([]byte, error) {
	return unicodeHash(vXXHash, id)
}

func init() {
	Register("unicode_loose_xxhash", NewUnicodeLooseXXHash)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"

	"github.com/teh-cmc/seq/rr_seq"
)

// -----------------------------------------------------------------------------

var s *rrs.RRServer

func usage() string { return "./server addr peer1_addr [peerN_addr...]" }

func main() {
	flag.Parse()
	if flag.NArg() < 1 {
		fmt.Println("Usage:", usage())
		os.Exit(1)
	}

	f, err := ioutil.TempFile("", "server_example")
	if err != nil {
		log.Fatal(err)
	}
	path := f.Name()
	_ = f.Close()

	s, err = rrs.NewRRServer(flag.Arg(0), path, flag.Args()[1:]...)
	log.Println("ready")

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	<-c

	_ = s.Close()
}
<commit_msg>RRSeq serv example: added mandatory path param<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"

	"github.com/teh-cmc/seq/rr_seq"
)

// -----------------------------------------------------------------------------

var s *rrs.RRServer

func usage() string { return "./server addr path peer1_addr [peerN_addr...]" }

func main() {
	flag.Parse()
	if flag.NArg() < 3 {
		fmt.Println("Usage:", usage())
		os.Exit(1)
	}

	s, err := rrs.NewRRServer(flag.Arg(0), flag.Arg(1), flag.Args()[2:]...)
	log.Println("ready")

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	<-c

	_ = s.Close()
}
<|endoftext|>


<commit_before>package rancher

import (
	"fmt"

	"golang.org/x/net/context"

	"github.com/Sirupsen/logrus"
	"github.com/rancher/go-rancher/v2"
	"github.com/rancher/rancher-compose-executor/config"
	"github.com/rancher/rancher-compose-executor/project"
)

type RancherVolumesFactory struct {
	Context *Context
}

func (f *RancherVolumesFactory) Create(projectName string, volumeConfigs map[string]*config.VolumeConfig, serviceConfigs *config.ServiceConfigs) (project.Volumes, error) {
	volumes := make([]*Volume, 0, len(volumeConfigs))
	for name, config := range volumeConfigs {
		volume := NewVolume(projectName, name, config, f.Context)
		volumes = append(volumes, volume)
	}
	return &Volumes{
		volumes: volumes,
		Context: f.Context,
	}, nil
}

type Volumes struct {
	volumes []*Volume
	Context *Context
}

func (v *Volumes) Initialize(ctx context.Context) error {
	for _, volume := range v.volumes {
		if err := volume.EnsureItExists(ctx); err != nil {
			return err
		}
	}

	return nil
}

func (v *Volumes) Remove(ctx context.Context) error {
	for _, volume := range v.volumes {
		if err := volume.Remove(ctx); err != nil {
			return err
		}
	}
	return nil
}

type Volume struct {
	context       *Context
	name          string
	projectName   string
	driver        string
	driverOptions map[string]string
	external      bool
	perContainer  bool
}

// InspectTemplate looks up a volume template
func (v *Volume) InspectTemplate(ctx context.Context) (*client.VolumeTemplate, error) {
	volumes, err := v.context.Client.VolumeTemplate.List(&client.ListOpts{
		Filters: map[string]interface{}{
			"name":    v.name,
			"stackId": v.context.Stack.Id,
		},
	})
	if err != nil {
		return nil, err
	}

	if len(volumes.Data) > 0 {
		return &volumes.Data[0], nil
	}

	return nil, nil
}

// InspectExternal looks up a volume
func (v *Volume) InspectExternal(ctx context.Context) (*client.Volume, error) {
	volumes, err := v.context.Client.Volume.List(&client.ListOpts{
		Filters: map[string]interface{}{
			"name": v.name,
		},
	})
	if err != nil {
		return nil, err
	}

	if len(volumes.Data) > 0 {
		return &volumes.Data[0], nil
	}

	return nil, nil
}

func (v *Volume) Remove(ctx context.Context) error {
	if v.external {
		return nil
	}

	volumeResource, err := v.InspectTemplate(ctx)
	if err != nil {
		return err
	}
	return v.context.Client.VolumeTemplate.Delete(volumeResource)
}

func (v *Volume) EnsureItExists(ctx context.Context) error {
	if v.external {
		volumeResource, err := v.InspectExternal(ctx)
		if err != nil {
			return err
		}

		if volumeResource == nil {
			return fmt.Errorf("Volume %s declared as external, but could not be found. Please create the volume manually and try again.", v.name)
		}

		return nil
	}

	volumeResource, err := v.InspectTemplate(ctx)
	if err != nil {
		return err
	}

	if volumeResource == nil {
		logrus.Infof("Creating volume template %s", v.name)
		return v.create(ctx)
	} else {
		logrus.Infof("Existing volume template found for %s", v.name)
	}

	if v.driver != "" && volumeResource.Driver != v.driver {
		return fmt.Errorf("Volume %q needs to be recreated - driver has changed", v.name)
	}
	return nil
}

func (v *Volume) create(ctx context.Context) error {
	driverOptions := map[string]interface{}{}
	for k, v := range v.driverOptions {
		driverOptions[k] = v
	}
	_, err := v.context.Client.VolumeTemplate.Create(&client.VolumeTemplate{
		Name:         v.name,
		Driver:       v.driver,
		DriverOpts:   driverOptions,
		StackId:      v.context.Stack.Id,
		PerContainer: v.perContainer,
	})
	return err
}

func NewVolume(projectName, name string, config *config.VolumeConfig, context *Context) *Volume {
	return &Volume{
		context:       context,
		name:          name,
		projectName:   projectName,
		driver:        config.Driver,
		driverOptions: config.DriverOpts,
		external:      config.External.External,
		perContainer:  config.PerContainer,
	}
}
<commit_msg>External volumes should also be volume templates<commit_after>package rancher

import (
	"fmt"

	"golang.org/x/net/context"

	"github.com/Sirupsen/logrus"
	"github.com/rancher/go-rancher/v2"
	"github.com/rancher/rancher-compose-executor/config"
	"github.com/rancher/rancher-compose-executor/project"
)

type RancherVolumesFactory struct {
	Context *Context
}

func (f *RancherVolumesFactory) Create(projectName string, volumeConfigs map[string]*config.VolumeConfig, serviceConfigs *config.ServiceConfigs) (project.Volumes, error) {
	volumes := make([]*Volume, 0, len(volumeConfigs))
	for name, config := range volumeConfigs {
		volume := NewVolume(projectName, name, config, f.Context)
		volumes = append(volumes, volume)
	}
	return &Volumes{
		volumes: volumes,
		Context: f.Context,
	}, nil
}

type Volumes struct {
	volumes []*Volume
	Context *Context
}

func (v *Volumes) Initialize(ctx context.Context) error {
	for _, volume := range v.volumes {
		if err := volume.EnsureItExists(ctx); err != nil {
			return err
		}
	}

	return nil
}

func (v *Volumes) Remove(ctx context.Context) error {
	for _, volume := range v.volumes {
		if err := volume.Remove(ctx); err != nil {
			return err
		}
	}
	return nil
}

type Volume struct {
	context       *Context
	name          string
	projectName   string
	driver        string
	driverOptions map[string]string
	external      bool
	perContainer  bool
}

// Inspect looks up a volume template
func (v *Volume) Inspect(ctx context.Context) (*client.VolumeTemplate, error) {
	filters := map[string]interface{}{
		"name": v.name,
	}
	if !v.external {
		filters["stackId"] = v.context.Stack.Id
	}

	volumes, err := v.context.Client.VolumeTemplate.List(&client.ListOpts{
		Filters: filters,
	})
	if err != nil {
		return nil, err
	}

	if len(volumes.Data) > 0 {
		return &volumes.Data[0], nil
	}

	return nil, nil
}

func (v *Volume) Remove(ctx context.Context) error {
	if v.external {
		return nil
	}

	volumeResource, err := v.Inspect(ctx)
	if err != nil {
		return err
	}
	return v.context.Client.VolumeTemplate.Delete(volumeResource)
}

func (v *Volume) EnsureItExists(ctx context.Context) error {
	volumeResource, err := v.Inspect(ctx)
	if err != nil {
		return err
	}

	if volumeResource == nil {
		if v.external {
			return fmt.Errorf("Volume %s declared as external, but could not be found. Please create the volume manually and try again.", v.name)
		}
		logrus.Infof("Creating volume template %s", v.name)
		return v.create(ctx)
	} else {
		logrus.Infof("Existing volume template found for %s", v.name)
	}

	if v.driver != "" && volumeResource.Driver != v.driver {
		return fmt.Errorf("Volume %q needs to be recreated - driver has changed", v.name)
	}
	return nil
}

func (v *Volume) create(ctx context.Context) error {
	driverOptions := map[string]interface{}{}
	for k, v := range v.driverOptions {
		driverOptions[k] = v
	}
	_, err := v.context.Client.VolumeTemplate.Create(&client.VolumeTemplate{
		Name:         v.name,
		Driver:       v.driver,
		DriverOpts:   driverOptions,
		StackId:      v.context.Stack.Id,
		PerContainer: v.perContainer,
	})
	return err
}

func NewVolume(projectName, name string, config *config.VolumeConfig, context *Context) *Volume {
	return &Volume{
		context:       context,
		name:          name,
		projectName:   projectName,
		driver:        config.Driver,
		driverOptions: config.DriverOpts,
		external:      config.External.External,
		perContainer:  config.PerContainer,
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2018 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 ghcache implements an HTTP cache optimized for caching responses
// from the GitHub API (https://api.github.com).
//
// Specifically, it enforces a cache policy that revalidates every cache hit
// with a conditional request to upstream regardless of cache entry freshness
// because conditional requests for unchanged resources don't cost any API
// tokens!!! See: https://developer.github.com/v3/#conditional-requests
//
// It also provides request coalescing and prometheus instrumentation.
package ghcache

import (
	"context"
	"crypto/sha256"
	"net/http"
	"path"
	"strings"
	"time"

	"github.com/gomodule/redigo/redis"
	"github.com/gregjones/httpcache"
	"github.com/gregjones/httpcache/diskcache"
	rediscache "github.com/gregjones/httpcache/redis"
	"github.com/peterbourgon/diskv"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
	"golang.org/x/sync/semaphore"
	"k8s.io/test-infra/ghproxy/ghmetrics"
)

type CacheResponseMode string

// Cache response modes describe how ghcache fulfilled a request.
const (
	CacheModeHeader = "X-Cache-Mode"

	ModeError   CacheResponseMode = "ERROR"    // internal error handling request
	ModeNoStore CacheResponseMode = "NO-STORE" // response not cacheable
	ModeMiss    CacheResponseMode = "MISS"     // not in cache, request proxied and response cached.
	ModeChanged CacheResponseMode = "CHANGED"  // cache value invalid: resource changed, cache updated
	// The modes below are the happy cases in which the request is fulfilled for
	// free (no API tokens used).
	ModeCoalesced   CacheResponseMode = "COALESCED"   // coalesced request, this is a copied response
	ModeRevalidated CacheResponseMode = "REVALIDATED" // cached value revalidated and returned
)

func CacheModeIsFree(mode CacheResponseMode) bool {
	switch mode {
	case ModeCoalesced:
		return true
	case ModeRevalidated:
		return true
	case ModeError:
		// In this case we did not successfully communicate with the GH API, so no
		// token is used, but we also don't return a response, so ModeError won't
		// ever be returned as a value of CacheModeHeader.
		return true
	}
	return false
}

// cacheCounter provides the 'ghcache_responses' counter vec that is indexed
// by the cache response mode.
var cacheCounter = prometheus.NewCounterVec(
	prometheus.CounterOpts{
		Name: "ghcache_responses",
		Help: "How many cache responses of each cache response mode there are.",
	},
	[]string{"mode"},
)

// outboundConcurrencyGauge provides the 'concurrent_outbound_requests' gauge that
// is global to the proxy.
var outboundConcurrencyGauge = prometheus.NewGauge(prometheus.GaugeOpts{
	Name: "concurrent_outbound_requests",
	Help: "How many concurrent requests are in flight to GitHub servers.",
})

// pendingOutboundConnectionsGauge provides the 'pending_outbound_requests' gauge that
// is global to the proxy.
var pendingOutboundConnectionsGauge = prometheus.NewGauge(prometheus.GaugeOpts{
	Name: "pending_outbound_requests",
	Help: "How many pending requests are waiting to be sent to GitHub servers.",
})

func init() {
	prometheus.MustRegister(cacheCounter)
	prometheus.MustRegister(outboundConcurrencyGauge)
	prometheus.MustRegister(pendingOutboundConnectionsGauge)
}

func cacheResponseMode(headers http.Header) CacheResponseMode {
	if strings.Contains(headers.Get("Cache-Control"), "no-store") {
		return ModeNoStore
	}
	if strings.Contains(headers.Get("Status"), "304 Not Modified") {
		return ModeRevalidated
	}
	if headers.Get("X-Conditional-Request") != "" {
		return ModeChanged
	}
	return ModeMiss
}

func newThrottlingTransport(maxConcurrency int, delegate http.RoundTripper) http.RoundTripper {
	return &throttlingTransport{sem: semaphore.NewWeighted(int64(maxConcurrency)), delegate: delegate}
}

// throttlingTransport throttles outbound concurrency from the proxy
type throttlingTransport struct {
	sem      *semaphore.Weighted
	delegate http.RoundTripper
}

func (c *throttlingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	pendingOutboundConnectionsGauge.Inc()
	if err := c.sem.Acquire(context.Background(), 1); err != nil {
		logrus.WithField("cache-key", req.URL.String()).WithError(err).Error("Internal error acquiring semaphore.")
		return nil, err
	}
	defer c.sem.Release(1)
	pendingOutboundConnectionsGauge.Dec()
	outboundConcurrencyGauge.Inc()
	defer outboundConcurrencyGauge.Dec()
	return c.delegate.RoundTrip(req)
}

// upstreamTransport changes response headers from upstream before they
// reach the cache layer in order to force the caching policy we require.
//
// By default github responds to PR requests with:
//    Cache-Control: private, max-age=60, s-maxage=60
// Which means the httpcache would not consider anything stale for 60 seconds.
// However, we want to always revalidate cache entries using ETags and last
// modified times so this RoundTripper overrides response headers to:
//    Cache-Control: no-cache
// This instructs the cache to store the response, but always consider it stale.
type upstreamTransport struct {
	delegate http.RoundTripper
}

func (u upstreamTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	etag := req.Header.Get("if-none-match")

	// get authorization header to convert to sha256
	authHeader := req.Header.Get("Authorization")
	hasher := sha256.New()
	hasher.Write([]byte(authHeader))
	authHeaderHash := string(hasher.Sum(nil))
	if authHeader == "" {
		logrus.Warnf("Couldn't retrieve 'Authorization' header, %s is the hash of an empty string", authHeaderHash)
	}

	reqStartTime := time.Now()
	// Don't modify request, just pass to delegate.
	resp, err := u.delegate.RoundTrip(req)
	if err != nil {
		logrus.WithField("cache-key", req.URL.String()).WithError(err).Error("Error from upstream (GitHub).")
		return nil, err
	}
	responseTime := time.Now()

	if resp.StatusCode >= 400 {
		// Don't store errors. They can't be revalidated to save API tokens.
		resp.Header.Set("Cache-Control", "no-store")
	} else {
		resp.Header.Set("Cache-Control", "no-cache")
	}
	if etag != "" {
		resp.Header.Set("X-Conditional-Request", etag)
	}

	apiVersion := "v3"
	if strings.HasPrefix(req.URL.Path, "search") || strings.HasPrefix(req.URL.Path, "/search") {
		apiVersion = "v4"
	}

	ghmetrics.CollectGitHubTokenMetrics(authHeaderHash, apiVersion, resp.Header, reqStartTime, responseTime)

	return resp, nil
}

// NewDiskCache creates a GitHub cache RoundTripper that is backed by a disk
// cache.
func NewDiskCache(delegate http.RoundTripper, cacheDir string, cacheSizeGB, maxConcurrency int) http.RoundTripper {
	return NewFromCache(delegate, diskcache.NewWithDiskv(
		diskv.New(diskv.Options{
			BasePath:     path.Join(cacheDir, "data"),
			TempDir:      path.Join(cacheDir, "temp"),
			CacheSizeMax: uint64(cacheSizeGB) * uint64(1000000000), // convert G to B
		})),
		maxConcurrency,
	)
}

// NewMemCache creates a GitHub cache RoundTripper that is backed by a memory
// cache.
func NewMemCache(delegate http.RoundTripper, maxConcurrency int) http.RoundTripper {
	return NewFromCache(delegate, httpcache.NewMemoryCache(), maxConcurrency)
}

// NewFromCache creates a GitHub cache RoundTripper that is backed by the
// specified httpcache.Cache implementation.
func NewFromCache(delegate http.RoundTripper, cache httpcache.Cache, maxConcurrency int) http.RoundTripper {
	cacheTransport := httpcache.NewTransport(cache)
	cacheTransport.Transport = newThrottlingTransport(maxConcurrency, upstreamTransport{delegate: delegate})
	return &requestCoalescer{
		keys:     make(map[string]*responseWaiter),
		delegate: cacheTransport,
	}
}

// NewRedisCache creates a GitHub cache RoundTripper that is backed by a Redis
// cache.
func NewRedisCache(delegate http.RoundTripper, redisAddress string, maxConcurrency int) http.RoundTripper {
	conn, err := redis.Dial("tcp", redisAddress)
	if err != nil {
		logrus.WithError(err).Fatal("Error connecting to Redis")
	}
	return NewFromCache(delegate, rediscache.NewWithClient(conn), maxConcurrency)
}
<commit_msg>Be more careful when manipulating auth header in metrics<commit_after>/*
Copyright 2018 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 ghcache implements an HTTP cache optimized for caching responses
// from the GitHub API (https://api.github.com).
//
// Specifically, it enforces a cache policy that revalidates every cache hit
// with a conditional request to upstream regardless of cache entry freshness
// because conditional requests for unchanged resources don't cost any API
// tokens!!! See: https://developer.github.com/v3/#conditional-requests
//
// It also provides request coalescing and prometheus instrumentation.
package ghcache

import (
	"context"
	"crypto/sha256"
	"fmt"
	"net/http"
	"path"
	"strings"
	"time"

	"github.com/gomodule/redigo/redis"
	"github.com/gregjones/httpcache"
	"github.com/gregjones/httpcache/diskcache"
	rediscache "github.com/gregjones/httpcache/redis"
	"github.com/peterbourgon/diskv"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
	"golang.org/x/sync/semaphore"
	"k8s.io/test-infra/ghproxy/ghmetrics"
)

type CacheResponseMode string

// Cache response modes describe how ghcache fulfilled a request.
const (
	CacheModeHeader = "X-Cache-Mode"

	ModeError   CacheResponseMode = "ERROR"    // internal error handling request
	ModeNoStore CacheResponseMode = "NO-STORE" // response not cacheable
	ModeMiss    CacheResponseMode = "MISS"     // not in cache, request proxied and response cached.
	ModeChanged CacheResponseMode = "CHANGED"  // cache value invalid: resource changed, cache updated
	// The modes below are the happy cases in which the request is fulfilled for
	// free (no API tokens used).
	ModeCoalesced   CacheResponseMode = "COALESCED"   // coalesced request, this is a copied response
	ModeRevalidated CacheResponseMode = "REVALIDATED" // cached value revalidated and returned
)

func CacheModeIsFree(mode CacheResponseMode) bool {
	switch mode {
	case ModeCoalesced:
		return true
	case ModeRevalidated:
		return true
	case ModeError:
		// In this case we did not successfully communicate with the GH API, so no
		// token is used, but we also don't return a response, so ModeError won't
		// ever be returned as a value of CacheModeHeader.
		return true
	}
	return false
}

// cacheCounter provides the 'ghcache_responses' counter vec that is indexed
// by the cache response mode.
var cacheCounter = prometheus.NewCounterVec(
	prometheus.CounterOpts{
		Name: "ghcache_responses",
		Help: "How many cache responses of each cache response mode there are.",
	},
	[]string{"mode"},
)

// outboundConcurrencyGauge provides the 'concurrent_outbound_requests' gauge that
// is global to the proxy.
var outboundConcurrencyGauge = prometheus.NewGauge(prometheus.GaugeOpts{
	Name: "concurrent_outbound_requests",
	Help: "How many concurrent requests are in flight to GitHub servers.",
})

// pendingOutboundConnectionsGauge provides the 'pending_outbound_requests' gauge that
// is global to the proxy.
var pendingOutboundConnectionsGauge = prometheus.NewGauge(prometheus.GaugeOpts{
	Name: "pending_outbound_requests",
	Help: "How many pending requests are waiting to be sent to GitHub servers.",
})

func init() {
	prometheus.MustRegister(cacheCounter)
	prometheus.MustRegister(outboundConcurrencyGauge)
	prometheus.MustRegister(pendingOutboundConnectionsGauge)
}

func cacheResponseMode(headers http.Header) CacheResponseMode {
	if strings.Contains(headers.Get("Cache-Control"), "no-store") {
		return ModeNoStore
	}
	if strings.Contains(headers.Get("Status"), "304 Not Modified") {
		return ModeRevalidated
	}
	if headers.Get("X-Conditional-Request") != "" {
		return ModeChanged
	}
	return ModeMiss
}

func newThrottlingTransport(maxConcurrency int, delegate http.RoundTripper) http.RoundTripper {
	return &throttlingTransport{sem: semaphore.NewWeighted(int64(maxConcurrency)), delegate: delegate}
}

// throttlingTransport throttles outbound concurrency from the proxy
type throttlingTransport struct {
	sem      *semaphore.Weighted
	delegate http.RoundTripper
}

func (c *throttlingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	pendingOutboundConnectionsGauge.Inc()
	if err := c.sem.Acquire(context.Background(), 1); err != nil {
		logrus.WithField("cache-key", req.URL.String()).WithError(err).Error("Internal error acquiring semaphore.")
		return nil, err
	}
	defer c.sem.Release(1)
	pendingOutboundConnectionsGauge.Dec()
	outboundConcurrencyGauge.Inc()
	defer outboundConcurrencyGauge.Dec()
	return c.delegate.RoundTrip(req)
}

// upstreamTransport changes response headers from upstream before they
// reach the cache layer in order to force the caching policy we require.
//
// By default github responds to PR requests with:
//    Cache-Control: private, max-age=60, s-maxage=60
// Which means the httpcache would not consider anything stale for 60 seconds.
// However, we want to always revalidate cache entries using ETags and last
// modified times so this RoundTripper overrides response headers to:
//    Cache-Control: no-cache
// This instructs the cache to store the response, but always consider it stale.
type upstreamTransport struct {
	delegate http.RoundTripper
}

func (u upstreamTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	etag := req.Header.Get("if-none-match")

	// get authorization header to convert to sha256
	authHeader := req.Header.Get("Authorization")
	if authHeader == "" {
		logrus.Warn("Couldn't retrieve 'Authorization' header, adding to unknown bucket")
		authHeader = "unknown"
	}
	hasher := sha256.New()
	hasher.Write([]byte(authHeader))
	authHeaderHash := fmt.Sprintf("%x", hasher.Sum(nil)) // use %x to make this a utf-8 string for use as a label

	reqStartTime := time.Now()
	// Don't modify request, just pass to delegate.
	resp, err := u.delegate.RoundTrip(req)
	if err != nil {
		logrus.WithField("cache-key", req.URL.String()).WithError(err).Error("Error from upstream (GitHub).")
		return nil, err
	}
	responseTime := time.Now()

	if resp.StatusCode >= 400 {
		// Don't store errors. They can't be revalidated to save API tokens.
		resp.Header.Set("Cache-Control", "no-store")
	} else {
		resp.Header.Set("Cache-Control", "no-cache")
	}
	if etag != "" {
		resp.Header.Set("X-Conditional-Request", etag)
	}

	apiVersion := "v3"
	if strings.HasPrefix(req.URL.Path, "search") || strings.HasPrefix(req.URL.Path, "/search") {
		apiVersion = "v4"
	}

	ghmetrics.CollectGitHubTokenMetrics(authHeaderHash, apiVersion, resp.Header, reqStartTime, responseTime)

	return resp, nil
}

// NewDiskCache creates a GitHub cache RoundTripper that is backed by a disk
// cache.
func NewDiskCache(delegate http.RoundTripper, cacheDir string, cacheSizeGB, maxConcurrency int) http.RoundTripper {
	return NewFromCache(delegate, diskcache.NewWithDiskv(
		diskv.New(diskv.Options{
			BasePath:     path.Join(cacheDir, "data"),
			TempDir:      path.Join(cacheDir, "temp"),
			CacheSizeMax: uint64(cacheSizeGB) * uint64(1000000000), // convert G to B
		})),
		maxConcurrency,
	)
}

// NewMemCache creates a GitHub cache RoundTripper that is backed by a memory
// cache.
func NewMemCache(delegate http.RoundTripper, maxConcurrency int) http.RoundTripper {
	return NewFromCache(delegate, httpcache.NewMemoryCache(), maxConcurrency)
}

// NewFromCache creates a GitHub cache RoundTripper that is backed by the
// specified httpcache.Cache implementation.
func NewFromCache(delegate http.RoundTripper, cache httpcache.Cache, maxConcurrency int) http.RoundTripper {
	cacheTransport := httpcache.NewTransport(cache)
	cacheTransport.Transport = newThrottlingTransport(maxConcurrency, upstreamTransport{delegate: delegate})
	return &requestCoalescer{
		keys:     make(map[string]*responseWaiter),
		delegate: cacheTransport,
	}
}

// NewRedisCache creates a GitHub cache RoundTripper that is backed by a Redis
// cache.
func NewRedisCache(delegate http.RoundTripper, redisAddress string, maxConcurrency int) http.RoundTripper {
	conn, err := redis.Dial("tcp", redisAddress)
	if err != nil {
		logrus.WithError(err).Fatal("Error connecting to Redis")
	}
	return NewFromCache(delegate, rediscache.NewWithClient(conn), maxConcurrency)
}
<|endoftext|>