<commit_before>// Copyright 2018 Google LLC
//
// 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
//
//     https://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 transformer invokes the golang HTML parser, executes the
// individual transformers (unless overridden), and prints the output
// to the provided string.
package transformer

import (
	"net/url"
	"regexp"
	"strings"

	"github.com/ampproject/amppackager/transformer/internal/amphtml"
	"github.com/ampproject/amppackager/transformer/internal/htmlnode"
	"github.com/ampproject/amppackager/transformer/printer"
	rpb "github.com/ampproject/amppackager/transformer/request"
	"github.com/ampproject/amppackager/transformer/transformers"
	"github.com/pkg/errors"
	"golang.org/x/net/html/atom"
	"golang.org/x/net/html"
)

// Transformer functions must be added here in order to be passed in from
// the command line or invoked from other languages. Please keep alphabetical.
// Case-insensitive lookup.
//
// NOTE: The string mapping is necessary as a language cross-over to
// allow explicit transformer invocation (via the CUSTOM config).
var transformerFunctionMap = map[string]func(*transformers.Context) error{
	"ampboilerplate":        transformers.AMPBoilerplate,
	"ampruntimecss":         transformers.AMPRuntimeCSS,
	"linktag":               transformers.LinkTag,
	"metatag":               transformers.MetaTag,
	"nodecleanup":           transformers.NodeCleanup,
	"reorderhead":           transformers.ReorderHead,
	"serversiderendering":   transformers.ServerSideRendering,
	"stripjs":               transformers.StripJS,
	"transformedidentifier": transformers.TransformedIdentifier,
	"url":                   transformers.URL,
	"urlrewrite":            transformers.URLRewrite,
}

// The map of config to the list of transformers, in the order in
// which they should be executed.
var configMap = map[rpb.Request_TransformersConfig][]func(*transformers.Context) error{
	rpb.Request_DEFAULT: {
		// NodeCleanup should be first.
		transformers.NodeCleanup,
		transformers.StripJS,
		transformers.MetaTag,
		// TODO(alin04): Reenable LinkTag once validation is done.
		// transformers.LinkTag,
		// end TODO
		transformers.URL,
		transformers.AMPBoilerplate,
		// TODO(alin04): Reenable SSR once validation is done.
		// transformers.ServerSideRendering,
		// end TODO
		// AmpRuntimeCss must run after ServerSideRendering
		// TODO(alin04): Reenable AMPRuntimeCSS and
		// TransformedIdentifier once validation is done
		// transformers.AMPRuntimeCSS,
		// transformers.TransformedIdentifier,
		// end TODO
		// ReorderHead should run after all transformers that modify the
		// <head>, as they may do so without preserving the proper order.
		transformers.ReorderHead,
		transformers.URLRewrite,
	},
	rpb.Request_NONE: {},
	rpb.Request_VALIDATION: {
		// TODO(alin04): Fill this in
		transformers.ReorderHead,
	},
	rpb.Request_CUSTOM: {},
}

// The maximum number of preloads to place in the Link header. This limit
// should be enforced by AMP Caches, to protect any pages that prefetch the SXG
// from an unnecessary number of fetches.
const maxPreloads = 20

// Override for tests.
var runTransformers = func(c *transformers.Context, fns []func(*transformers.Context) error) error {
	// Invoke the configured transformers
	for _, f := range fns {
		if err := f(c); err != nil {
			return errors.WithStack(err)
		}
	}
	return nil
}

// ampAttrRE is a regexp to match html amp attributes. Its group capture should
// be compared against ampFormatSuffixes.
var ampAttrRE = func() *regexp.Regexp {
	r := regexp.MustCompile(`\A(?:⚡|amp)(?:4(.+))?\z`)
	if len(r.SubexpNames()) != 2 {
		panic("ampAttrRE must have 1 subexpression")
	}
	return r
}()

// The allowed AMP formats, and their serialization as an html "amp4" attribute.
var ampFormatSuffixes = map[rpb.Request_HtmlFormat]string{
	rpb.Request_AMP:       "",
	rpb.Request_AMP4ADS:   "ads",
	rpb.Request_AMP4EMAIL: "email",
}

// The keys from ampFormatSuffixes.
var ampFormats = func() []rpb.Request_HtmlFormat {
	ret := []rpb.Request_HtmlFormat{}
	for v := range ampFormatSuffixes {
		ret = append(ret, v)
	}
	return ret
}()

// isAllowed returns true iff the declared format, as parsed from the html amp
// attribute, corresponds to one of the allowed formats. If allowedFormats is
// empty, then any AMP format is allowed.
func isAllowed(declaredFormat string, allowedFormats []rpb.Request_HtmlFormat) bool {
	// Default to all formats.
	if len(allowedFormats) == 0 {
		allowedFormats = ampFormats
	}
	for _, allowedFormat := range allowedFormats {
		suffix, ok := ampFormatSuffixes[allowedFormat]
		if ok && declaredFormat == suffix {
			return true
		}
	}
	return false
}

// setDOM parses the input HTML and sets c.DOM to the parsed DOM struct.
func setDOM(c *transformers.Context, s string) error {
	doc, err := html.Parse(strings.NewReader(s))
	if err != nil {
		return errors.Wrap(err, "Error parsing input HTML")
	}

	dom, err := amphtml.NewDOM(doc)
	if err != nil {
		return err
	}
	c.DOM = dom
	return nil
}

// requireAMPAttribute returns an error if the <html> tag doesn't contain an
// attribute indicating that the document is AMP.
func requireAMPAttribute(dom *amphtml.DOM, allowedFormats []rpb.Request_HtmlFormat) error {
	for _, attr := range dom.HTMLNode.Attr {
		if attr.Namespace == "" {
			if match := ampAttrRE.FindStringSubmatch(attr.Key); match != nil {
				if isAllowed(match[1], allowedFormats) {
					return nil
				}
			}
		}
	}
	return errors.New("html tag is missing an AMP attribute")
}

// extractPreloads returns a list of absolute URLs of the resources to preload,
// in the order to preload them. It depends on transformers.ReorderHead having
// run.
func extractPreloads(dom *amphtml.DOM) []*rpb.Metadata_Preload {
	// If you add additional preloads here, verify that they can not be
	// unintentionally author supplied.
	preloads := []*rpb.Metadata_Preload{}
	for child := dom.HeadNode.FirstChild; child != nil; child = child.NextSibling {
		switch child.DataAtom {
		case atom.Script:
			if src, ok := htmlnode.GetAttributeVal(child, "", "src"); ok {
				preloads = append(preloads, &rpb.Metadata_Preload{Url: src, As: "script"})
			}
		case atom.Link:
			if rel, ok := htmlnode.GetAttributeVal(child, "", "rel"); ok {
				if strings.EqualFold(rel, "stylesheet") {
					if href, ok := htmlnode.GetAttributeVal(child, "", "href"); ok {
						preloads = append(preloads, &rpb.Metadata_Preload{Url: href, As: "style"})
					}
				}
			}
		}
		if len(preloads) == maxPreloads {
			break
		}
	}
	return preloads
}

// setBaseURL derives the absolute base URL, and sets it on c.BaseURL. The value
// is derived using the <base> href in the DOM, if it exists. If the href is
// relative, it is parsed in the context of the document URL.
// This must run after DocumentURL is set on the context.
func setBaseURL(c *transformers.Context) {
	if n, ok := htmlnode.FindNode(c.DOM.HeadNode, atom.Base); ok {
		if v, ok := htmlnode.GetAttributeVal(n, "", "href"); ok {
			if u, err := c.DocumentURL.Parse(v); err == nil {
				c.BaseURL = u
				return
			}
		}
	}
	c.BaseURL = c.DocumentURL
}

// Process will parse the given request, which contains the HTML to
// transform, applying the requested list of transformers, and return the
// transformed HTML and list of resources to preload (absolute URLs), or an
// error.
//
// If the requested list of transformers is empty, apply the default.
func Process(r *rpb.Request) (string, *rpb.Metadata, error) {
	context := &transformers.Context{}
	var err error

	err = setDOM(context, r.Html)
	if err != nil {
		return "", nil, err
	}

	if err = requireAMPAttribute(context.DOM, r.AllowedFormats); err != nil {
		return "", nil, err
	}

	fns := configMap[r.Config]
	if r.Config == rpb.Request_CUSTOM {
		for _, val := range r.Transformers {
			fn, ok := transformerFunctionMap[strings.ToLower(val)]
			if !ok {
				return "", nil, errors.Errorf("transformer doesn't exist: %s", val)
			}
			fns = append(fns, fn)
		}
	}

	context.DocumentURL, err = url.Parse(r.DocumentUrl)
	if err != nil {
		return "", nil, err
	}

	context.Version = r.Version
	if r.Version == 0 {
		context.Version, err = SelectVersion(nil)
		if err != nil {
			return "", nil, err
		}
	}

	// This must run AFTER DocumentURL is parsed.
	setBaseURL(context)

	if err := runTransformers(context, fns); err != nil {
		return "", nil, err
	}
	var o strings.Builder
	err = printer.Print(&o, context.DOM.RootNode)
	if err != nil {
		return "", nil, err
	}
	return o.String(), &rpb.Metadata{Preloads: extractPreloads(context.DOM)}, nil
}
<commit_msg>Now that preconnects are added, move URLRewrite transformer before ReorderHead<commit_after>// Copyright 2018 Google LLC
//
// 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
//
//     https://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 transformer invokes the golang HTML parser, executes the
// individual transformers (unless overridden), and prints the output
// to the provided string.
package transformer

import (
	"net/url"
	"regexp"
	"strings"

	"github.com/ampproject/amppackager/transformer/internal/amphtml"
	"github.com/ampproject/amppackager/transformer/internal/htmlnode"
	"github.com/ampproject/amppackager/transformer/printer"
	rpb "github.com/ampproject/amppackager/transformer/request"
	"github.com/ampproject/amppackager/transformer/transformers"
	"github.com/pkg/errors"
	"golang.org/x/net/html/atom"
	"golang.org/x/net/html"
)

// Transformer functions must be added here in order to be passed in from
// the command line or invoked from other languages. Please keep alphabetical.
// Case-insensitive lookup.
//
// NOTE: The string mapping is necessary as a language cross-over to
// allow explicit transformer invocation (via the CUSTOM config).
var transformerFunctionMap = map[string]func(*transformers.Context) error{
	"ampboilerplate":        transformers.AMPBoilerplate,
	"ampruntimecss":         transformers.AMPRuntimeCSS,
	"linktag":               transformers.LinkTag,
	"metatag":               transformers.MetaTag,
	"nodecleanup":           transformers.NodeCleanup,
	"reorderhead":           transformers.ReorderHead,
	"serversiderendering":   transformers.ServerSideRendering,
	"stripjs":               transformers.StripJS,
	"transformedidentifier": transformers.TransformedIdentifier,
	"url":                   transformers.URL,
	"urlrewrite":            transformers.URLRewrite,
}

// The map of config to the list of transformers, in the order in
// which they should be executed.
var configMap = map[rpb.Request_TransformersConfig][]func(*transformers.Context) error{
	rpb.Request_DEFAULT: {
		// NodeCleanup should be first.
		transformers.NodeCleanup,
		transformers.StripJS,
		transformers.MetaTag,
		// TODO(alin04): Reenable LinkTag once validation is done.
		// transformers.LinkTag,
		// end TODO
		transformers.URL,
		transformers.AMPBoilerplate,
		// TODO(alin04): Reenable SSR once validation is done.
		// transformers.ServerSideRendering,
		// end TODO
		// AmpRuntimeCss must run after ServerSideRendering
		// TODO(alin04): Reenable AMPRuntimeCSS and
		// TransformedIdentifier once validation is done
		// transformers.AMPRuntimeCSS,
		// transformers.TransformedIdentifier,
		// end TODO
		transformers.URLRewrite,
		// ReorderHead should run after all transformers that modify the
		// <head>, as they may do so without preserving the proper order.
		transformers.ReorderHead,
	},
	rpb.Request_NONE: {},
	rpb.Request_VALIDATION: {
		// TODO(alin04): Fill this in
		transformers.ReorderHead,
	},
	rpb.Request_CUSTOM: {},
}

// The maximum number of preloads to place in the Link header. This limit
// should be enforced by AMP Caches, to protect any pages that prefetch the SXG
// from an unnecessary number of fetches.
const maxPreloads = 20

// Override for tests.
var runTransformers = func(c *transformers.Context, fns []func(*transformers.Context) error) error {
	// Invoke the configured transformers
	for _, f := range fns {
		if err := f(c); err != nil {
			return errors.WithStack(err)
		}
	}
	return nil
}

// ampAttrRE is a regexp to match html amp attributes. Its group capture should
// be compared against ampFormatSuffixes.
var ampAttrRE = func() *regexp.Regexp {
	r := regexp.MustCompile(`\A(?:⚡|amp)(?:4(.+))?\z`)
	if len(r.SubexpNames()) != 2 {
		panic("ampAttrRE must have 1 subexpression")
	}
	return r
}()

// The allowed AMP formats, and their serialization as an html "amp4" attribute.
var ampFormatSuffixes = map[rpb.Request_HtmlFormat]string{
	rpb.Request_AMP:       "",
	rpb.Request_AMP4ADS:   "ads",
	rpb.Request_AMP4EMAIL: "email",
}

// The keys from ampFormatSuffixes.
var ampFormats = func() []rpb.Request_HtmlFormat {
	ret := []rpb.Request_HtmlFormat{}
	for v := range ampFormatSuffixes {
		ret = append(ret, v)
	}
	return ret
}()

// isAllowed returns true iff the declared format, as parsed from the html amp
// attribute, corresponds to one of the allowed formats. If allowedFormats is
// empty, then any AMP format is allowed.
func isAllowed(declaredFormat string, allowedFormats []rpb.Request_HtmlFormat) bool {
	// Default to all formats.
	if len(allowedFormats) == 0 {
		allowedFormats = ampFormats
	}
	for _, allowedFormat := range allowedFormats {
		suffix, ok := ampFormatSuffixes[allowedFormat]
		if ok && declaredFormat == suffix {
			return true
		}
	}
	return false
}

// setDOM parses the input HTML and sets c.DOM to the parsed DOM struct.
func setDOM(c *transformers.Context, s string) error {
	doc, err := html.Parse(strings.NewReader(s))
	if err != nil {
		return errors.Wrap(err, "Error parsing input HTML")
	}

	dom, err := amphtml.NewDOM(doc)
	if err != nil {
		return err
	}
	c.DOM = dom
	return nil
}

// requireAMPAttribute returns an error if the <html> tag doesn't contain an
// attribute indicating that the document is AMP.
func requireAMPAttribute(dom *amphtml.DOM, allowedFormats []rpb.Request_HtmlFormat) error {
	for _, attr := range dom.HTMLNode.Attr {
		if attr.Namespace == "" {
			if match := ampAttrRE.FindStringSubmatch(attr.Key); match != nil {
				if isAllowed(match[1], allowedFormats) {
					return nil
				}
			}
		}
	}
	return errors.New("html tag is missing an AMP attribute")
}

// extractPreloads returns a list of absolute URLs of the resources to preload,
// in the order to preload them. It depends on transformers.ReorderHead having
// run.
func extractPreloads(dom *amphtml.DOM) []*rpb.Metadata_Preload {
	// If you add additional preloads here, verify that they can not be
	// unintentionally author supplied.
	preloads := []*rpb.Metadata_Preload{}
	for child := dom.HeadNode.FirstChild; child != nil; child = child.NextSibling {
		switch child.DataAtom {
		case atom.Script:
			if src, ok := htmlnode.GetAttributeVal(child, "", "src"); ok {
				preloads = append(preloads, &rpb.Metadata_Preload{Url: src, As: "script"})
			}
		case atom.Link:
			if rel, ok := htmlnode.GetAttributeVal(child, "", "rel"); ok {
				if strings.EqualFold(rel, "stylesheet") {
					if href, ok := htmlnode.GetAttributeVal(child, "", "href"); ok {
						preloads = append(preloads, &rpb.Metadata_Preload{Url: href, As: "style"})
					}
				}
			}
		}
		if len(preloads) == maxPreloads {
			break
		}
	}
	return preloads
}

// setBaseURL derives the absolute base URL, and sets it on c.BaseURL. The value
// is derived using the <base> href in the DOM, if it exists. If the href is
// relative, it is parsed in the context of the document URL.
// This must run after DocumentURL is set on the context.
func setBaseURL(c *transformers.Context) {
	if n, ok := htmlnode.FindNode(c.DOM.HeadNode, atom.Base); ok {
		if v, ok := htmlnode.GetAttributeVal(n, "", "href"); ok {
			if u, err := c.DocumentURL.Parse(v); err == nil {
				c.BaseURL = u
				return
			}
		}
	}
	c.BaseURL = c.DocumentURL
}

// Process will parse the given request, which contains the HTML to
// transform, applying the requested list of transformers, and return the
// transformed HTML and list of resources to preload (absolute URLs), or an
// error.
//
// If the requested list of transformers is empty, apply the default.
func Process(r *rpb.Request) (string, *rpb.Metadata, error) {
	context := &transformers.Context{}
	var err error

	err = setDOM(context, r.Html)
	if err != nil {
		return "", nil, err
	}

	if err = requireAMPAttribute(context.DOM, r.AllowedFormats); err != nil {
		return "", nil, err
	}

	fns := configMap[r.Config]
	if r.Config == rpb.Request_CUSTOM {
		for _, val := range r.Transformers {
			fn, ok := transformerFunctionMap[strings.ToLower(val)]
			if !ok {
				return "", nil, errors.Errorf("transformer doesn't exist: %s", val)
			}
			fns = append(fns, fn)
		}
	}

	context.DocumentURL, err = url.Parse(r.DocumentUrl)
	if err != nil {
		return "", nil, err
	}

	context.Version = r.Version
	if r.Version == 0 {
		context.Version, err = SelectVersion(nil)
		if err != nil {
			return "", nil, err
		}
	}

	// This must run AFTER DocumentURL is parsed.
	setBaseURL(context)

	if err := runTransformers(context, fns); err != nil {
		return "", nil, err
	}
	var o strings.Builder
	err = printer.Print(&o, context.DOM.RootNode)
	if err != nil {
		return "", nil, err
	}
	return o.String(), &rpb.Metadata{Preloads: extractPreloads(context.DOM)}, nil
}
<|endoftext|>


<commit_before>package main

import (
  "fmt"
  "log"
  "strings"
  "errors"
  "github.com/Jeffail/gabs"
)

// Name
// Traits
// Tags/Groups
type Person struct {
  Name string
  Json *gabs.Container
}

func new_person(name_in string) *Person {
  p := new(Person)
  p.Name = name_in
  jsonObj := gabs.New()
  jsonObj.Set(name_in, "person", "name")
  // jsonObj.Set("c", "person", "traits")
  // jsonObj.ArrayAppend("", "person", "tags")
  jsonObj.Array("person", "tags")
  p.Json = jsonObj
  return p
}

// See about comparing input to what already exists
func (p Person) add_trait(trait_path, trait_name, trait_text string) {
  current_json := p.Json

  // Get the full dotted path to use for all operations
  path := create_path(trait_path, trait_name)

  // 2 situations, the trait already exists or it doesnt
  // If it exists array it (there is a check for the array also)
  // If it doesnt exist add it

  // There are 2 additional situations if the trait exists
  // The array already exists (add to it)
  // The array doesn't already exist (make it and add to it)

  // It it exists, else it does not
  if current_json.ExistsP(path) {

    // Get an error if there is no array
    _, err := current_json.Path(path).Children()

    // If there is an error we need to make the array
    if err != nil {
      // Get the original value (that isnt an array)
      current_val := current_json.Path(path).Data()
      // Create an array
      current_json.ArrayP(path)
      // Add the original value to the array
      current_json.ArrayAppendP(current_val, path)
    }

    // Add the value we had originally wanted to add
    current_json.ArrayAppendP(trait_text, path)
  } else {
    current_json.SetP(trait_text, path)
  }

}

// If the name will be unique it will change
func (p Person) change_name(name string) error {
  if acceptable_person_name(name) {
    p.Name = name
    return nil
  } else {
    return errors.New("This name has been previously assinged, try another.")
  }
}

// This will have some content at some point to ensure uniquity among names
func acceptable_person_name(name string) bool {
  if true != false {
    return true
  } else {
    return false
  }
}

// See about comparing input to what already exists
func (p Person) add_tag(tag_name string) error {
  current_json := p.Json
  if acceptable_tag_name(tag_name) {
    current_json.ArrayAppend(tag_name, "person", "tags")
  } else {
    return errors.New("This tag has been previously assinged, try another.")
  }

}

// This will have some content at some point to ensure uniquity among tags
func acceptable_tag_name(name string) bool {
  if true != false {
    return true
  } else {
    return false
  }
}

func get_all_people_with_tag(tag_name string) {

}

func (p Person) get_person_tags() map[string]bool {
  tags := make(map[string]bool)
  children, err := p.Json.S("person", "tags").Children()
  check_err(err)
  for _, child := range children {
  	tags[child.Data().(string)] = true
  }
  return tags
}

func (p Person) get_person_trails() {
  children, err := p.Json.S("tags").ChildrenMap()
  check_err(err)
  for key, child := range children {
  	fmt.Printf("key: %v, value: %v\n", key, child.Data().(string))
  }
}

func (p Person) t_delete_trait(trait_name string) {
  err := p.Json.DeleteP("person.traits.relative")
  check_err(err)
}

func main() {

  traits_path := "person.traits"

  p1 := new_person("ben")
  p1.add_tag("friend")
  p1.add_tag("geog 1000")

  p2 := new_person("steve")
  p2.add_trait(traits_path, "relative.brother", "uno")
  // p2.add_trait(traits_path, "relative.sister", "dos")
  p2.add_trait(traits_path, "relative.brother", "tres")
  p2.add_trait(traits_path, "relative.brother", "quad")
  //p2.add_trait(traits_path, "relative", "tre")

  // fmt.Println(p2.Json.ExistsP("person.traits.relative"))

  // p2.t_delete_trait("bob")

  // fmt.Println(p1.Name, "->", p2.Name)
  // fmt.Println(p1.Json.String())
  fmt.Println(p2.Json.String())
  //
  //
  //
  // p2.get_person_trails()

}

// Support methods

func create_path(path, addition string) string {
  if !strings.HasSuffix(path, ".") {
    path = path + "."
  }
  if strings.HasPrefix(addition, ".") {
    addition = addition[1:]
  }
  return (path + addition)
}

func check_err(err error) {
  if err != nil {
    log.Fatal(err)
  }
}
<commit_msg>tags are now more selective<commit_after>package main

import (
  "fmt"
  "log"
  "strings"
  "errors"
  "github.com/Jeffail/gabs"
)

// Name
// Traits
// Tags/Groups
type Person struct {
  Name string
  Json *gabs.Container
}

func new_person(name_in string) *Person {
  p := new(Person)
  p.Name = name_in
  jsonObj := gabs.New()
  jsonObj.Set(name_in, "person", "name")
  // jsonObj.Set("c", "person", "traits")
  // jsonObj.ArrayAppend("", "person", "tags")
  jsonObj.Array("person", "tags")
  p.Json = jsonObj
  return p
}

// See about comparing input to what already exists
func (p Person) add_trait(trait_path, trait_name, trait_text string) {
  current_json := p.Json

  // Get the full dotted path to use for all operations
  path := create_path(trait_path, trait_name)

  // 2 situations, the trait already exists or it doesnt
  // If it exists array it (there is a check for the array also)
  // If it doesnt exist add it

  // There are 2 additional situations if the trait exists
  // The array already exists (add to it)
  // The array doesn't already exist (make it and add to it)

  // It it exists, else it does not
  if current_json.ExistsP(path) {

    // Get an error if there is no array
    _, err := current_json.Path(path).Children()

    // If there is an error we need to make the array
    if err != nil {
      // Get the original value (that isnt an array)
      current_val := current_json.Path(path).Data()
      // Create an array
      current_json.ArrayP(path)
      // Add the original value to the array
      current_json.ArrayAppendP(current_val, path)
    }

    // Add the value we had originally wanted to add
    current_json.ArrayAppendP(trait_text, path)
  } else {
    current_json.SetP(trait_text, path)
  }
}

// If the name will be unique it will change
func (p Person) change_name(name string) error {
  if acceptable_person_name(name) {
    p.Name = name
    return nil
  } else {
    return errors.New("This name has been previously assinged, try another.")
  }
}

// This will have some content at some point to ensure uniquity among names
func acceptable_person_name(name string) bool {
  if true != false {
    return true
  } else {
    return false
  }
}

// See about comparing input to what already exists
func (p Person) add_tag(tag_name string) error {
  current_json := p.Json
  if p.acceptable_tag_name(tag_name) {
    current_json.ArrayAppend(tag_name, "person", "tags")
    return nil
  } else {
    return errors.New("This tag has been previously assinged, try another.")
  }

}

// This will have some content at some point to ensure uniquity among tags
func (p Person) acceptable_tag_name(name string) bool {

  tag_set := p.get_person_tags()

  // If the tag is currently associated with the person do not add it again
  // If the statement is true then the tag is associated with the person
  if tag_set[name] == false {
    return true
  } else {
    return false
  }
}

func get_all_people_with_tag(tag_name string) {

}

func (p Person) get_person_tags() map[string]bool {
  // _, err := current_json.Path(path).Children()
  current_json := p.Json
  tags := make(map[string]bool)
  children, err := current_json.Path("person.tags").Children()
  check_err(err)
  for _, child := range children {
  	tags[child.Data().(string)] = true
  }
  return tags
}

func (p Person) get_person_trails() {
  children, err := p.Json.S("tags").ChildrenMap()
  check_err(err)
  for key, child := range children {
  	fmt.Printf("key: %v, value: %v\n", key, child.Data().(string))
  }
}

func (p Person) t_delete_trait(trait_name string) {
  err := p.Json.DeleteP("person.traits.relative")
  check_err(err)
}

func main() {

  traits_path := "person.traits"

  p1 := new_person("ben")
  p1.add_tag("friend")
  p1.add_tag("geog 1000")

  p2 := new_person("steve")
  p2.add_trait(traits_path, "relative.brother", "uno")
  // p2.add_trait(traits_path, "relative.sister", "dos")
  p2.add_trait(traits_path, "relative.brother", "tres")
  p2.add_trait(traits_path, "relative.brother", "quad")
  p2.add_trait(traits_path, "location.current", "md")

  p1.add_tag("dog person")
  p1.add_tag("cat person")
  p1.add_tag("cat person")
  //p2.add_trait(traits_path, "relative", "tre")

  // fmt.Println(p2.Json.ExistsP("person.traits.relative"))

  // p2.t_delete_trait("bob")

  // fmt.Println(p1.Name, "->", p2.Name)
  fmt.Println(p1.Json.String())
  fmt.Println(p2.Json.String())
  //
  //
  //
  // p2.get_person_trails()

}

// Support methods

// Import people json objects from a file
func import_people_from_file() {

}

// Export people json objects to a file
func export_people_to_file() {

}

// Create a path to a particular location within the json
func create_path(path, addition string) string {
  if !strings.HasSuffix(path, ".") {
    path = path + "."
  }
  if strings.HasPrefix(addition, ".") {
    addition = addition[1:]
  }
  return (path + addition)
}

// Break if there is an error passed in
func check_err(err error) {
  if err != nil {
    log.Fatal(err)
  }
}
<|endoftext|>


<commit_before>package git

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"

	. "github.com/polydawn/go-errcat"
	"gopkg.in/src-d/go-git.v4/plumbing/filemode"
	"gopkg.in/src-d/go-git.v4/plumbing/object"

	"go.polydawn.net/go-timeless-api"
	"go.polydawn.net/go-timeless-api/rio"
	"go.polydawn.net/go-timeless-api/util"
	"go.polydawn.net/rio/config"
	"go.polydawn.net/rio/fs"
	"go.polydawn.net/rio/fs/osfs"
	"go.polydawn.net/rio/fsOp"
	"go.polydawn.net/rio/transmat/mixins/cache"
	"go.polydawn.net/rio/transmat/mixins/filters"
	"go.polydawn.net/rio/transmat/mixins/fshash"
)

var (
	_ rio.UnpackFunc = Unpack
)

func Unpack(
	ctx context.Context, // Long-running call.  Cancellable.
	wareID api.WareID, // What wareID to fetch for unpacking.
	path string, // Where to unpack the fileset (absolute path).
	filt api.FilesetFilters, // Optionally: filters we should apply while unpacking.
	placementMode rio.PlacementMode, // Optionally: a placement mode (default is "copy").
	warehouses []api.WarehouseAddr, // Warehouses we can try to fetch from.
	mon rio.Monitor, // Optionally: callbacks for progress monitoring.
) (_ api.WareID, err error) {
	if mon.Chan != nil {
		defer close(mon.Chan)
	}
	defer RequireErrorHasCategory(&err, rio.ErrorCategory(""))

	// Sanitize arguments.
	if wareID.Type != PackType {
		return api.WareID{}, Errorf(rio.ErrUsage, "this transmat implementation only supports packtype %q (not %q)", PackType, wareID.Type)
	}
	if placementMode == "" {
		placementMode = rio.Placement_Copy
	}
	// Wrap the direct unpack func with cache behavior; call that.
	return cache.Lrn2Cache(
		osfs.New(config.GetCacheBasePath()),
		unpack,
	)(ctx, wareID, path, filt, placementMode, warehouses, mon)
}

func unpack(
	ctx context.Context,
	wareID api.WareID,
	path string,
	filt api.FilesetFilters,
	placementMode rio.PlacementMode,
	warehouses []api.WarehouseAddr,
	mon rio.Monitor,
) (_ api.WareID, err error) {
	defer RequireErrorHasCategory(&err, rio.ErrorCategory(""))

	// Sanitize arguments.
	path2 := fs.MustAbsolutePath(path)
	filt2, err := apiutil.ProcessFilters(filt, apiutil.FilterPurposeUnpack)
	if err != nil {
		return api.WareID{}, Errorf(rio.ErrUsage, "invalid filter specification: %s", err)
	}

	// Pick a warehouse and get a reader.
	//  This is a *very* expensive operation for git.  It's less
	//  of "pick a warehouse" and more "download the whole thing and hope we
	//  get what we wanted" (which is very ironic for a system that has
	//  a CAS system on its inside, yes).
	whCtrl, err := pick(ctx,
		wareID,
		warehouses,
		osfs.New(config.GetCacheBasePath().Join(fs.MustRelPath("git/objs"))),
		mon,
	)
	if err != nil {
		return api.WareID{}, err
	}
	tr, err := whCtrl.GetTree(wareID.Hash)
	if err != nil {
		panic(err)
	}
	tw := object.NewTreeWalker(tr, true, nil)

	// Construct filesystem wrapper to use for all our ops.
	afs := osfs.New(path2)

	// Make the root dir.  Git doesn't have metadata for the tree root.
	conjuredFmeta := fshash.DefaultDirMetadata()
	filters.Apply(filt2, &conjuredFmeta)
	if err := fsOp.PlaceFile(afs, conjuredFmeta, nil, filt2.SkipChown); err != nil {
		return api.WareID{}, Errorf(rio.ErrInoperablePath, "error while unpacking: %s", err)
	}

	// Extract.
	// Iterate over each entry, mutating filesystem as we go.
	dirs := make([]fs.RelPath, 1, 200) // Keep for dir time repair at end.
	dirs[0] = fs.RelPath{}
	for {
		fmeta := fs.Metadata{}
		name, te, err := tw.Next()

		// Check for done.
		if err == io.EOF {
			break // sucess!  end of archive.
		}
		if err != nil {
			return api.WareID{}, Errorf(rio.ErrWareCorrupt, "corrupt git tree: %s", err)
		}
		if ctx.Err() != nil {
			return api.WareID{}, Errorf(rio.ErrCancelled, "cancelled")
		}
		//fmt.Fprintf(os.Stderr, "walking git tree %s -- %#v\n", name, te)

		// Reshuffle metainfo to our default format.
		fmeta.Name = fs.MustRelPath(name)
		switch te.Mode {
		case filemode.Dir:
			fmeta.Type = fs.Type_Dir
			fmeta.Perms = 0755
			dirs = append(dirs, fmeta.Name)
		case filemode.Regular:
			fmeta.Type = fs.Type_File
			fmeta.Perms = 0644
		case filemode.Executable:
			fmeta.Type = fs.Type_File
			fmeta.Perms = 0755
		case filemode.Symlink:
			fmeta.Type = fs.Type_Symlink
			fmeta.Perms = 0644
			// Hang on, extracting a symlink is actually rough.
			tf, err := tr.TreeEntryFile(&te)
			if err != nil {
				return api.WareID{}, Errorf(rio.ErrWareCorrupt, "corrupt git tree: %s", err)
			}
			reader, err := tf.Blob.Reader()
			if err != nil {
				return api.WareID{}, Errorf(rio.ErrWareCorrupt, "corrupt git tree: %s", err)
			}
			blob, err := ioutil.ReadAll(reader)
			if err != nil {
				return api.WareID{}, Errorf(rio.ErrWareCorrupt, "corrupt git tree: %s", err)
			}
			fmeta.Linkname = string(blob)
		case filemode.Submodule:
			// TODO
			continue
		case filemode.Empty:
			fallthrough
		case filemode.Deprecated:
			fallthrough
		default:
			panic(fmt.Errorf("unknown git filemode %#v", te.Mode))
		}
		fmeta.Mtime = apiutil.DefaultMtime // git doesn't have time info

		// Apply filters.
		filters.Apply(filt2, &fmeta)

		// Place the file.
		switch fmeta.Type {
		case fs.Type_File:
			tf, err := tr.TreeEntryFile(&te)
			if err != nil {
				return api.WareID{}, Errorf(rio.ErrWareCorrupt, "corrupt git tree: %s", err)
			}
			reader, err := tf.Blob.Reader()
			if err != nil {
				return api.WareID{}, Errorf(rio.ErrWareCorrupt, "corrupt git tree: %s", err)
			}
			if err := fsOp.PlaceFile(afs, fmeta, reader, filt2.SkipChown); err != nil {
				return api.WareID{}, Errorf(rio.ErrInoperablePath, "error while unpacking: %s", err)
			}
			reader.Close()
		default:
			if err := fsOp.PlaceFile(afs, fmeta, nil, filt2.SkipChown); err != nil {
				return api.WareID{}, Errorf(rio.ErrInoperablePath, "error while unpacking: %s", err)
			}
		}
	}

	// Cleanup dir times with a post-order traversal over the bucket.
	//  Files and dirs placed inside dirs cause the parent's mtime to update, so we have to re-pave them.
	for i := len(dirs) - 1; i >= 0; i-- {
		if err := afs.SetTimesNano(dirs[i], fs.DefaultAtime, fs.DefaultAtime); err != nil {
			return api.WareID{}, Errorf(rio.ErrInoperablePath, "error while unpacking: %s", err)
		}
	}

	// That's it.  Checkout should have already checked the hash, so we just return it.
	return wareID, nil
}
<commit_msg>git: extract an unpackOneRepo method, preparing to recurse on it for submodule unpack support.<commit_after>package git

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"

	. "github.com/polydawn/go-errcat"
	"gopkg.in/src-d/go-git.v4/plumbing/filemode"
	"gopkg.in/src-d/go-git.v4/plumbing/object"

	"go.polydawn.net/go-timeless-api"
	"go.polydawn.net/go-timeless-api/rio"
	"go.polydawn.net/go-timeless-api/util"
	"go.polydawn.net/rio/config"
	"go.polydawn.net/rio/fs"
	"go.polydawn.net/rio/fs/osfs"
	"go.polydawn.net/rio/fsOp"
	"go.polydawn.net/rio/transmat/mixins/cache"
	"go.polydawn.net/rio/transmat/mixins/filters"
	"go.polydawn.net/rio/transmat/mixins/fshash"
	gitWarehouse "go.polydawn.net/rio/warehouse/impl/git"
)

var (
	_ rio.UnpackFunc = Unpack
)

func Unpack(
	ctx context.Context, // Long-running call.  Cancellable.
	wareID api.WareID, // What wareID to fetch for unpacking.
	path string, // Where to unpack the fileset (absolute path).
	filt api.FilesetFilters, // Optionally: filters we should apply while unpacking.
	placementMode rio.PlacementMode, // Optionally: a placement mode (default is "copy").
	warehouses []api.WarehouseAddr, // Warehouses we can try to fetch from.
	mon rio.Monitor, // Optionally: callbacks for progress monitoring.
) (_ api.WareID, err error) {
	if mon.Chan != nil {
		defer close(mon.Chan)
	}
	defer RequireErrorHasCategory(&err, rio.ErrorCategory(""))

	// Sanitize arguments.
	if wareID.Type != PackType {
		return api.WareID{}, Errorf(rio.ErrUsage, "this transmat implementation only supports packtype %q (not %q)", PackType, wareID.Type)
	}
	if placementMode == "" {
		placementMode = rio.Placement_Copy
	}
	// Wrap the direct unpack func with cache behavior; call that.
	return cache.Lrn2Cache(
		osfs.New(config.GetCacheBasePath()),
		unpack,
	)(ctx, wareID, path, filt, placementMode, warehouses, mon)
}

func unpack(
	ctx context.Context,
	wareID api.WareID,
	path string,
	filt api.FilesetFilters,
	placementMode rio.PlacementMode,
	warehouses []api.WarehouseAddr,
	mon rio.Monitor,
) (_ api.WareID, err error) {
	defer RequireErrorHasCategory(&err, rio.ErrorCategory(""))

	// Sanitize arguments.
	path2 := fs.MustAbsolutePath(path)
	filt2, err := apiutil.ProcessFilters(filt, apiutil.FilterPurposeUnpack)
	if err != nil {
		return api.WareID{}, Errorf(rio.ErrUsage, "invalid filter specification: %s", err)
	}

	// Pick a warehouse and get a reader.
	//  This is a *very* expensive operation for git.  It's less
	//  of "pick a warehouse" and more "download the whole thing and hope we
	//  get what we wanted" (which is very ironic for a system that has
	//  a CAS system on its inside, yes).
	whCtrl, err := pick(ctx,
		wareID,
		warehouses,
		osfs.New(config.GetCacheBasePath().Join(fs.MustRelPath("git/objs"))),
		mon,
	)
	if err != nil {
		return api.WareID{}, err
	}

	// Open a tree to walk in the main repo.
	//  We'll do submodule checkouts somewhere deep in the middle of this.
	tr, err := whCtrl.GetTree(wareID.Hash)
	if err != nil {
		panic(err)
	}

	// Construct filesystem wrapper to use for all our ops.
	afs := osfs.New(path2)

	// Walk.
	if err := unpackOneRepo(ctx, tr, afs, true, filt2, nil, mon); err != nil {
		return api.WareID{}, err
	}

	// That's it.  Checkout should have already checked the hash, so we just return it.
	return wareID, nil
}

func unpackOneRepo(
	ctx context.Context,
	tr *object.Tree,
	afs fs.FS,
	isRoot bool, // if true, will recurse for submodules (with this set to false).
	filt apiutil.FilesetFilters,
	submoduleCtrls map[string]*gitWarehouse.Controller,
	mon rio.Monitor,
) (err error) {
	tw := object.NewTreeWalker(tr, true, nil)

	// Make the root dir.  Git doesn't have metadata for the tree root.
	conjuredFmeta := fshash.DefaultDirMetadata()
	filters.Apply(filt, &conjuredFmeta)
	if err := fsOp.PlaceFile(afs, conjuredFmeta, nil, filt.SkipChown); err != nil {
		return Errorf(rio.ErrInoperablePath, "error while unpacking: %s", err)
	}

	// Extract.
	// Iterate over each entry, mutating filesystem as we go.
	dirs := make([]fs.RelPath, 1, 200) // Keep for dir time repair at end.
	dirs[0] = fs.RelPath{}
	for {
		fmeta := fs.Metadata{}
		name, te, err := tw.Next()

		// Check for done.
		if err == io.EOF {
			break // sucess!  end of archive.
		}
		if err != nil {
			return Errorf(rio.ErrWareCorrupt, "corrupt git tree: %s", err)
		}
		if ctx.Err() != nil {
			return Errorf(rio.ErrCancelled, "cancelled")
		}
		//fmt.Fprintf(os.Stderr, "walking git tree %s -- %#v\n", name, te)

		// Reshuffle metainfo to our default format.
		fmeta.Name = fs.MustRelPath(name)
		switch te.Mode {
		case filemode.Dir:
			fmeta.Type = fs.Type_Dir
			fmeta.Perms = 0755
			dirs = append(dirs, fmeta.Name)
		case filemode.Regular:
			fmeta.Type = fs.Type_File
			fmeta.Perms = 0644
		case filemode.Executable:
			fmeta.Type = fs.Type_File
			fmeta.Perms = 0755
		case filemode.Symlink:
			fmeta.Type = fs.Type_Symlink
			fmeta.Perms = 0644
			// Hang on, extracting a symlink is actually rough.
			tf, err := tr.TreeEntryFile(&te)
			if err != nil {
				return Errorf(rio.ErrWareCorrupt, "corrupt git tree: %s", err)
			}
			reader, err := tf.Blob.Reader()
			if err != nil {
				return Errorf(rio.ErrWareCorrupt, "corrupt git tree: %s", err)
			}
			blob, err := ioutil.ReadAll(reader)
			if err != nil {
				return Errorf(rio.ErrWareCorrupt, "corrupt git tree: %s", err)
			}
			fmeta.Linkname = string(blob)
		case filemode.Submodule:
			// TODO
			continue
		case filemode.Empty:
			fallthrough
		case filemode.Deprecated:
			fallthrough
		default:
			panic(fmt.Errorf("unknown git filemode %#v", te.Mode))
		}
		fmeta.Mtime = apiutil.DefaultMtime // git doesn't have time info

		// Apply filters.
		filters.Apply(filt, &fmeta)

		// Place the file.
		switch fmeta.Type {
		case fs.Type_File:
			tf, err := tr.TreeEntryFile(&te)
			if err != nil {
				return Errorf(rio.ErrWareCorrupt, "corrupt git tree: %s", err)
			}
			reader, err := tf.Blob.Reader()
			if err != nil {
				return Errorf(rio.ErrWareCorrupt, "corrupt git tree: %s", err)
			}
			if err := fsOp.PlaceFile(afs, fmeta, reader, filt.SkipChown); err != nil {
				return Errorf(rio.ErrInoperablePath, "error while unpacking: %s", err)
			}
			reader.Close()
		default:
			if err := fsOp.PlaceFile(afs, fmeta, nil, filt.SkipChown); err != nil {
				return Errorf(rio.ErrInoperablePath, "error while unpacking: %s", err)
			}
		}
	}

	// Cleanup dir times with a post-order traversal over the bucket.
	//  Files and dirs placed inside dirs cause the parent's mtime to update, so we have to re-pave them.
	for i := len(dirs) - 1; i >= 0; i-- {
		if err := afs.SetTimesNano(dirs[i], fs.DefaultAtime, fs.DefaultAtime); err != nil {
			return Errorf(rio.ErrInoperablePath, "error while unpacking: %s", err)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package cache

import (
	"context"
	"errors"
	"fmt"
	"log"
	"sync/atomic"
	"time"

	"github.com/cespare/xxhash/v2"
	"github.com/go-redis/redis/v8"
	"github.com/klauspost/compress/s2"
	"github.com/vmihailenco/bufpool"
	"github.com/vmihailenco/go-tinylfu"
	"github.com/vmihailenco/msgpack/v5"
	"golang.org/x/sync/singleflight"
)

const (
	compressionThreshold = 64
	timeLen              = 4
)

const (
	noCompression = 0x0
	s2Compression = 0x1
)

var (
	ErrCacheMiss          = errors.New("cache: key is missing")
	errRedisLocalCacheNil = errors.New("cache: both Redis and LocalCache are nil")
)

type rediser interface {
	Set(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.StatusCmd
	SetXX(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.BoolCmd
	SetNX(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.BoolCmd

	Get(ctx context.Context, key string) *redis.StringCmd
	Del(ctx context.Context, keys ...string) *redis.IntCmd
}

type Item struct {
	Ctx context.Context

	Key   string
	Value interface{}

	// TTL is the cache expiration time.
	// Default TTL is 1 hour.
	TTL time.Duration

	// Do returns value to be cached.
	Do func(*Item) (interface{}, error)

	// IfExists only sets the key if it already exist.
	IfExists bool

	// IfNotExists only sets the key if it does not already exist.
	IfNotExists bool

	// SkipLocalCache skips local cache as if it is not set.
	SkipLocalCache bool
}

func (item *Item) Context() context.Context {
	if item.Ctx == nil {
		return context.Background()
	}
	return item.Ctx
}

func (item *Item) value() (interface{}, error) {
	if item.Do != nil {
		return item.Do(item)
	}
	if item.Value != nil {
		return item.Value, nil
	}
	return nil, nil
}

func (item *Item) ttl() time.Duration {
	if item.TTL < 0 {
		return 0
	}
	if item.TTL != 0 && item.TTL < time.Second {
		log.Printf("too short TTL for key=%q: %s", item.Key, item.TTL)
		return time.Hour
	}
	return item.TTL
}

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

type Options struct {
	Redis rediser

	LocalCache    *tinylfu.SyncT
	LocalCacheTTL time.Duration

	StatsEnabled bool
}

func (opt *Options) init() {
	switch opt.LocalCacheTTL {
	case -1:
		opt.LocalCacheTTL = 0
	case 0:
		opt.LocalCacheTTL = time.Minute
	}
}

type Cache struct {
	opt *Options

	group   singleflight.Group
	bufpool bufpool.Pool

	hits   uint64
	misses uint64
}

func New(opt *Options) *Cache {
	opt.init()
	return &Cache{
		opt: opt,
	}
}

// Set caches the item.
func (cd *Cache) Set(item *Item) error {
	_, _, err := cd.set(item)
	return err
}

func (cd *Cache) set(item *Item) ([]byte, bool, error) {
	value, err := item.value()
	if err != nil {
		return nil, false, err
	}

	b, err := cd.Marshal(value)
	if err != nil {
		return nil, false, err
	}

	if cd.opt.LocalCache != nil {
		cd.localSet(item.Key, b)
	}

	if cd.opt.Redis == nil {
		if cd.opt.LocalCache == nil {
			return b, true, errRedisLocalCacheNil
		}
		return b, true, nil
	}

	if item.IfExists {
		return b, true, cd.opt.Redis.SetXX(item.Context(), item.Key, b, item.ttl()).Err()
	}

	if item.IfNotExists {
		return b, true, cd.opt.Redis.SetNX(item.Context(), item.Key, b, item.ttl()).Err()
	}

	return b, true, cd.opt.Redis.Set(item.Context(), item.Key, b, item.ttl()).Err()
}

// Exists reports whether value for the given key exists.
func (cd *Cache) Exists(ctx context.Context, key string) bool {
	return cd.Get(ctx, key, nil) == nil
}

// Get gets the value for the given key.
func (cd *Cache) Get(ctx context.Context, key string, value interface{}) error {
	return cd.get(ctx, key, value, false)
}

// Get gets the value for the given key skipping local cache.
func (cd *Cache) GetSkippingLocalCache(
	ctx context.Context, key string, value interface{},
) error {
	return cd.get(ctx, key, value, true)
}

func (cd *Cache) get(
	ctx context.Context,
	key string,
	value interface{},
	skipLocalCache bool,
) error {
	b, err := cd.getBytes(ctx, key, skipLocalCache)
	if err != nil {
		return err
	}
	return cd.Unmarshal(b, value)
}

func (cd *Cache) getBytes(ctx context.Context, key string, skipLocalCache bool) ([]byte, error) {
	if !skipLocalCache && cd.opt.LocalCache != nil {
		b, ok := cd.localGet(key)
		if ok {
			return b, nil
		}
	}

	if cd.opt.Redis == nil {
		if cd.opt.LocalCache == nil {
			return nil, errRedisLocalCacheNil
		}
		return nil, ErrCacheMiss
	}

	b, err := cd.opt.Redis.Get(ctx, key).Bytes()
	if err != nil {
		if cd.opt.StatsEnabled {
			atomic.AddUint64(&cd.misses, 1)
		}
		if err == redis.Nil {
			return nil, ErrCacheMiss
		}
		return nil, err
	}

	if cd.opt.StatsEnabled {
		atomic.AddUint64(&cd.hits, 1)
	}

	if !skipLocalCache && cd.opt.LocalCache != nil {
		cd.localSet(key, b)
	}
	return b, nil
}

// Once gets the item.Value for the given item.Key from the cache or
// executes, caches, and returns the results of the given item.Func,
// making sure that only one execution is in-flight for a given item.Key
// at a time. If a duplicate comes in, the duplicate caller waits for the
// original to complete and receives the same results.
func (cd *Cache) Once(item *Item) error {
	b, cached, err := cd.getSetItemBytesOnce(item)
	if err != nil {
		return err
	}

	if item.Value == nil || len(b) == 0 {
		return nil
	}

	if err := cd.Unmarshal(b, item.Value); err != nil {
		if cached {
			_ = cd.Delete(item.Context(), item.Key)
			return cd.Once(item)
		}
		return err
	}

	return nil
}

func (cd *Cache) getSetItemBytesOnce(item *Item) (b []byte, cached bool, err error) {
	if cd.opt.LocalCache != nil {
		b, ok := cd.localGet(item.Key)
		if ok {
			return b, true, nil
		}
	}

	v, err, _ := cd.group.Do(item.Key, func() (interface{}, error) {
		b, err := cd.getBytes(item.Context(), item.Key, item.SkipLocalCache)
		if err == nil {
			cached = true
			return b, nil
		}

		b, ok, err := cd.set(item)
		if ok {
			return b, nil
		}
		return nil, err
	})
	if err != nil {
		return nil, false, err
	}
	return v.([]byte), cached, nil
}

func (cd *Cache) Delete(ctx context.Context, key string) error {
	if cd.opt.LocalCache != nil {
		cd.opt.LocalCache.Del(xxhash.Sum64String(key))
	}

	if cd.opt.Redis == nil {
		if cd.opt.LocalCache == nil {
			return errRedisLocalCacheNil
		}
		return nil
	}

	_, err := cd.opt.Redis.Del(ctx, key).Result()
	return err
}

func (cd *Cache) localSet(key string, b []byte) {
	cd.opt.LocalCache.Set(&tinylfu.Item{
		Key:      xxhash.Sum64String(key),
		Value:    b,
		ExpireAt: time.Now().Add(cd.opt.LocalCacheTTL),
	})
}

func (cd *Cache) localGet(key string) ([]byte, bool) {
	val, ok := cd.opt.LocalCache.Get(xxhash.Sum64String(key))
	if !ok {
		return nil, false
	}

	b := val.([]byte)
	return b, true
}

func (cd *Cache) Marshal(value interface{}) ([]byte, error) {
	switch value := value.(type) {
	case nil:
		return nil, nil
	case []byte:
		return value, nil
	case string:
		return []byte(value), nil
	}

	buf := cd.bufpool.Get()
	defer cd.bufpool.Put(buf)

	enc := msgpack.GetEncoder()
	enc.Reset(buf)
	enc.UseCompactInts(true)

	err := enc.Encode(value)

	msgpack.PutEncoder(enc)

	if err != nil {
		return nil, err
	}

	return compress(buf.Bytes()), nil
}

func compress(data []byte) []byte {
	if len(data) < compressionThreshold {
		n := len(data) + 1
		b := make([]byte, n, n+timeLen)
		copy(b, data)
		b[len(b)-1] = noCompression
		return b
	}

	n := s2.MaxEncodedLen(len(data)) + 1
	b := make([]byte, n, n+timeLen)
	b = s2.Encode(b, data)
	b = append(b, s2Compression)
	return b
}

func (cd *Cache) Unmarshal(b []byte, value interface{}) error {
	if len(b) == 0 {
		return nil
	}

	switch value := value.(type) {
	case nil:
		return nil
	case *[]byte:
		clone := make([]byte, len(b))
		copy(clone, b)
		*value = clone
		return nil
	case *string:
		*value = string(b)
		return nil
	}

	switch c := b[len(b)-1]; c {
	case noCompression:
		b = b[:len(b)-1]
	case s2Compression:
		b = b[:len(b)-1]

		n, err := s2.DecodedLen(b)
		if err != nil {
			return err
		}

		buf := bufpool.Get(n)
		defer bufpool.Put(buf)

		b, err = s2.Decode(buf.Bytes(), b)
		if err != nil {
			return err
		}
	default:
		return fmt.Errorf("unknown compression method: %x", c)
	}

	return unmarshal(b, value)
}

// modified from github.com/vmihailenco/msgpack/v5@v5.0.0-alpha.2/decode.go:Unmarshal
//
// add dec.UseJSONTag(true)
func unmarshal(data []byte, v interface{}) error {
	dec := msgpack.GetDecoder()

	dec.ResetBytes(data)
	dec.UseJSONTag(true)
	err := dec.Decode(v)

	msgpack.PutDecoder(dec)

	return err
}

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

type Stats struct {
	Hits   uint64
	Misses uint64
}

// Stats returns cache statistics.
func (cd *Cache) Stats() *Stats {
	if !cd.opt.StatsEnabled {
		return nil
	}
	return &Stats{
		Hits:   atomic.LoadUint64(&cd.hits),
		Misses: atomic.LoadUint64(&cd.misses),
	}
}
<commit_msg>Fully remove UseJSONTag<commit_after>package cache

import (
	"context"
	"errors"
	"fmt"
	"log"
	"sync/atomic"
	"time"

	"github.com/cespare/xxhash/v2"
	"github.com/go-redis/redis/v8"
	"github.com/klauspost/compress/s2"
	"github.com/vmihailenco/bufpool"
	"github.com/vmihailenco/go-tinylfu"
	"github.com/vmihailenco/msgpack/v5"
	"golang.org/x/sync/singleflight"
)

const (
	compressionThreshold = 64
	timeLen              = 4
)

const (
	noCompression = 0x0
	s2Compression = 0x1
)

var (
	ErrCacheMiss          = errors.New("cache: key is missing")
	errRedisLocalCacheNil = errors.New("cache: both Redis and LocalCache are nil")
)

type rediser interface {
	Set(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.StatusCmd
	SetXX(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.BoolCmd
	SetNX(ctx context.Context, key string, value interface{}, expiration time.Duration) *redis.BoolCmd

	Get(ctx context.Context, key string) *redis.StringCmd
	Del(ctx context.Context, keys ...string) *redis.IntCmd
}

type Item struct {
	Ctx context.Context

	Key   string
	Value interface{}

	// TTL is the cache expiration time.
	// Default TTL is 1 hour.
	TTL time.Duration

	// Do returns value to be cached.
	Do func(*Item) (interface{}, error)

	// IfExists only sets the key if it already exist.
	IfExists bool

	// IfNotExists only sets the key if it does not already exist.
	IfNotExists bool

	// SkipLocalCache skips local cache as if it is not set.
	SkipLocalCache bool
}

func (item *Item) Context() context.Context {
	if item.Ctx == nil {
		return context.Background()
	}
	return item.Ctx
}

func (item *Item) value() (interface{}, error) {
	if item.Do != nil {
		return item.Do(item)
	}
	if item.Value != nil {
		return item.Value, nil
	}
	return nil, nil
}

func (item *Item) ttl() time.Duration {
	if item.TTL < 0 {
		return 0
	}
	if item.TTL != 0 && item.TTL < time.Second {
		log.Printf("too short TTL for key=%q: %s", item.Key, item.TTL)
		return time.Hour
	}
	return item.TTL
}

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

type Options struct {
	Redis rediser

	LocalCache    *tinylfu.SyncT
	LocalCacheTTL time.Duration

	StatsEnabled bool
}

func (opt *Options) init() {
	switch opt.LocalCacheTTL {
	case -1:
		opt.LocalCacheTTL = 0
	case 0:
		opt.LocalCacheTTL = time.Minute
	}
}

type Cache struct {
	opt *Options

	group   singleflight.Group
	bufpool bufpool.Pool

	hits   uint64
	misses uint64
}

func New(opt *Options) *Cache {
	opt.init()
	return &Cache{
		opt: opt,
	}
}

// Set caches the item.
func (cd *Cache) Set(item *Item) error {
	_, _, err := cd.set(item)
	return err
}

func (cd *Cache) set(item *Item) ([]byte, bool, error) {
	value, err := item.value()
	if err != nil {
		return nil, false, err
	}

	b, err := cd.Marshal(value)
	if err != nil {
		return nil, false, err
	}

	if cd.opt.LocalCache != nil {
		cd.localSet(item.Key, b)
	}

	if cd.opt.Redis == nil {
		if cd.opt.LocalCache == nil {
			return b, true, errRedisLocalCacheNil
		}
		return b, true, nil
	}

	if item.IfExists {
		return b, true, cd.opt.Redis.SetXX(item.Context(), item.Key, b, item.ttl()).Err()
	}

	if item.IfNotExists {
		return b, true, cd.opt.Redis.SetNX(item.Context(), item.Key, b, item.ttl()).Err()
	}

	return b, true, cd.opt.Redis.Set(item.Context(), item.Key, b, item.ttl()).Err()
}

// Exists reports whether value for the given key exists.
func (cd *Cache) Exists(ctx context.Context, key string) bool {
	return cd.Get(ctx, key, nil) == nil
}

// Get gets the value for the given key.
func (cd *Cache) Get(ctx context.Context, key string, value interface{}) error {
	return cd.get(ctx, key, value, false)
}

// Get gets the value for the given key skipping local cache.
func (cd *Cache) GetSkippingLocalCache(
	ctx context.Context, key string, value interface{},
) error {
	return cd.get(ctx, key, value, true)
}

func (cd *Cache) get(
	ctx context.Context,
	key string,
	value interface{},
	skipLocalCache bool,
) error {
	b, err := cd.getBytes(ctx, key, skipLocalCache)
	if err != nil {
		return err
	}
	return cd.Unmarshal(b, value)
}

func (cd *Cache) getBytes(ctx context.Context, key string, skipLocalCache bool) ([]byte, error) {
	if !skipLocalCache && cd.opt.LocalCache != nil {
		b, ok := cd.localGet(key)
		if ok {
			return b, nil
		}
	}

	if cd.opt.Redis == nil {
		if cd.opt.LocalCache == nil {
			return nil, errRedisLocalCacheNil
		}
		return nil, ErrCacheMiss
	}

	b, err := cd.opt.Redis.Get(ctx, key).Bytes()
	if err != nil {
		if cd.opt.StatsEnabled {
			atomic.AddUint64(&cd.misses, 1)
		}
		if err == redis.Nil {
			return nil, ErrCacheMiss
		}
		return nil, err
	}

	if cd.opt.StatsEnabled {
		atomic.AddUint64(&cd.hits, 1)
	}

	if !skipLocalCache && cd.opt.LocalCache != nil {
		cd.localSet(key, b)
	}
	return b, nil
}

// Once gets the item.Value for the given item.Key from the cache or
// executes, caches, and returns the results of the given item.Func,
// making sure that only one execution is in-flight for a given item.Key
// at a time. If a duplicate comes in, the duplicate caller waits for the
// original to complete and receives the same results.
func (cd *Cache) Once(item *Item) error {
	b, cached, err := cd.getSetItemBytesOnce(item)
	if err != nil {
		return err
	}

	if item.Value == nil || len(b) == 0 {
		return nil
	}

	if err := cd.Unmarshal(b, item.Value); err != nil {
		if cached {
			_ = cd.Delete(item.Context(), item.Key)
			return cd.Once(item)
		}
		return err
	}

	return nil
}

func (cd *Cache) getSetItemBytesOnce(item *Item) (b []byte, cached bool, err error) {
	if cd.opt.LocalCache != nil {
		b, ok := cd.localGet(item.Key)
		if ok {
			return b, true, nil
		}
	}

	v, err, _ := cd.group.Do(item.Key, func() (interface{}, error) {
		b, err := cd.getBytes(item.Context(), item.Key, item.SkipLocalCache)
		if err == nil {
			cached = true
			return b, nil
		}

		b, ok, err := cd.set(item)
		if ok {
			return b, nil
		}
		return nil, err
	})
	if err != nil {
		return nil, false, err
	}
	return v.([]byte), cached, nil
}

func (cd *Cache) Delete(ctx context.Context, key string) error {
	if cd.opt.LocalCache != nil {
		cd.opt.LocalCache.Del(xxhash.Sum64String(key))
	}

	if cd.opt.Redis == nil {
		if cd.opt.LocalCache == nil {
			return errRedisLocalCacheNil
		}
		return nil
	}

	_, err := cd.opt.Redis.Del(ctx, key).Result()
	return err
}

func (cd *Cache) localSet(key string, b []byte) {
	cd.opt.LocalCache.Set(&tinylfu.Item{
		Key:      xxhash.Sum64String(key),
		Value:    b,
		ExpireAt: time.Now().Add(cd.opt.LocalCacheTTL),
	})
}

func (cd *Cache) localGet(key string) ([]byte, bool) {
	val, ok := cd.opt.LocalCache.Get(xxhash.Sum64String(key))
	if !ok {
		return nil, false
	}

	b := val.([]byte)
	return b, true
}

func (cd *Cache) Marshal(value interface{}) ([]byte, error) {
	switch value := value.(type) {
	case nil:
		return nil, nil
	case []byte:
		return value, nil
	case string:
		return []byte(value), nil
	}

	buf := cd.bufpool.Get()
	defer cd.bufpool.Put(buf)

	enc := msgpack.GetEncoder()
	enc.Reset(buf)
	enc.UseCompactInts(true)

	err := enc.Encode(value)

	msgpack.PutEncoder(enc)

	if err != nil {
		return nil, err
	}

	return compress(buf.Bytes()), nil
}

func compress(data []byte) []byte {
	if len(data) < compressionThreshold {
		n := len(data) + 1
		b := make([]byte, n, n+timeLen)
		copy(b, data)
		b[len(b)-1] = noCompression
		return b
	}

	n := s2.MaxEncodedLen(len(data)) + 1
	b := make([]byte, n, n+timeLen)
	b = s2.Encode(b, data)
	b = append(b, s2Compression)
	return b
}

func (cd *Cache) Unmarshal(b []byte, value interface{}) error {
	if len(b) == 0 {
		return nil
	}

	switch value := value.(type) {
	case nil:
		return nil
	case *[]byte:
		clone := make([]byte, len(b))
		copy(clone, b)
		*value = clone
		return nil
	case *string:
		*value = string(b)
		return nil
	}

	switch c := b[len(b)-1]; c {
	case noCompression:
		b = b[:len(b)-1]
	case s2Compression:
		b = b[:len(b)-1]

		n, err := s2.DecodedLen(b)
		if err != nil {
			return err
		}

		buf := bufpool.Get(n)
		defer bufpool.Put(buf)

		b, err = s2.Decode(buf.Bytes(), b)
		if err != nil {
			return err
		}
	default:
		return fmt.Errorf("unknown compression method: %x", c)
	}

	dec := msgpack.GetDecoder()

	dec.ResetBytes(b)
	err := dec.Decode(value)

	msgpack.PutDecoder(dec)

	return err
}

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

type Stats struct {
	Hits   uint64
	Misses uint64
}

// Stats returns cache statistics.
func (cd *Cache) Stats() *Stats {
	if !cd.opt.StatsEnabled {
		return nil
	}
	return &Stats{
		Hits:   atomic.LoadUint64(&cd.hits),
		Misses: atomic.LoadUint64(&cd.misses),
	}
}
<|endoftext|>


<commit_before>package main

import "time"

var psCache map[string]FullParameters
var usersCache map[string][]string
var userPositionCache map[string]UserPositionJSON
var isLearning map[string]bool

func init() {
	go clearCache()
}

func clearCache() {
	for {
		Debug.Println("Clearing cache.")
		isLearning = make(map[string]bool)
		psCache = make(map[string]FullParameters)
		usersCache = make(map[string][]string)
		userPositionCache = make(map[string]UserPositionJSON)
		time.Sleep(time.Minute * 10)
	}
}
<commit_msg>This line prevents binaries from working?<commit_after>package main

import "time"

var psCache map[string]FullParameters
var usersCache map[string][]string
var userPositionCache map[string]UserPositionJSON
var isLearning map[string]bool

func init() {
	go clearCache()
}

func clearCache() {
	for {
		isLearning = make(map[string]bool)
		psCache = make(map[string]FullParameters)
		usersCache = make(map[string][]string)
		userPositionCache = make(map[string]UserPositionJSON)
		time.Sleep(time.Minute * 10)
	}
}
<|endoftext|>


<commit_before>package main

import (
    "fmt"
    "os"
)

type command_fun func(args []string)

var cmd_map = map[string]command_fun{
    "update":  cmd_update,
    "search":  cmd_search,
    "show":    cmd_show,
    "install": cmd_install,
    "create":  cmd_create,
    "publish": cmd_publish,
}

func main() {
    args := os.Args

    // invoked without command
    if len(args) == 1 {
        print_usage()
        return
    }

    // invoked with 'help' command
    if args[1] == "help" || args[1] == "-h" {
        if len(args) == 2 {
            print_usage()
            return
        } else {
            // convert 'gopm help <cmd>' to 'gopm <cmd> -h'
            args = append(args[0:1], args[2], "-h")
        }
    }

    cmd := args[1]
    fun, cmd_is_valid := cmd_map[cmd]

    // invoked with a invalid command
    if !cmd_is_valid {
        fmt.Printf("\n'%v' is not a valid command, see 'gopm help'\n\n", cmd)
        return
    }

    // call the command's function
    fun(args[2:]) // drop the 'gopm' and <command> argument
}

func print_usage() {
    fmt.Print(`
Usage: 'gopm <command> <options>'

where <command> is one of:

    update      update package index to the latest
    search      search for packages
    install     install a package
    create      create your own package
    publish     publish your package to the central index database

'gopm help <command>':
    show command specific options

'gopm help' or 'gopm -h':
    show this help info

`)
}
<commit_msg>add 'gopm show' to the main help<commit_after>package main

import (
    "fmt"
    "os"
)

type command_fun func(args []string)

var cmd_map = map[string]command_fun{
    "update":  cmd_update,
    "search":  cmd_search,
    "show":    cmd_show,
    "install": cmd_install,
    "create":  cmd_create,
    "publish": cmd_publish,
}

func main() {
    args := os.Args

    // invoked without command
    if len(args) == 1 {
        print_usage()
        return
    }

    // invoked with 'help' command
    if args[1] == "help" || args[1] == "-h" {
        if len(args) == 2 {
            print_usage()
            return
        } else {
            // convert 'gopm help <cmd>' to 'gopm <cmd> -h'
            args = append(args[0:1], args[2], "-h")
        }
    }

    cmd := args[1]
    fun, cmd_is_valid := cmd_map[cmd]

    // invoked with a invalid command
    if !cmd_is_valid {
        fmt.Printf("\n'%v' is not a valid command, see 'gopm help'\n\n", cmd)
        return
    }

    // call the command's function
    fun(args[2:]) // drop the 'gopm' and <command> argument
}

func print_usage() {
    fmt.Print(`
Usage: 'gopm <command> <options>'

where <command> is one of:

    update      update your local package index
    search      search for packages
    show        show detailed information for packages
    install     install a package
    create      create your own package
    publish     publish your package to the central index database

'gopm help <command>':
    show command specific options

'gopm help' or 'gopm -h':
    show this help info

`)
}
<|endoftext|>


<commit_before>package libkb

import (
	"bufio"
	"fmt"
	"io"
	"strconv"
	"strings"
)

//=============================================================================

type BucketDict struct {
	d map[string][]*GpgPrimaryKey
}

func NewBuckDict() *BucketDict {
	return &BucketDict{
		d: make(map[string][]*GpgPrimaryKey),
	}
}

func (bd *BucketDict) Add(k string, v *GpgPrimaryKey) {
	k = strings.ToLower(k)
	var bucket []*GpgPrimaryKey
	var found bool
	if bucket, found = bd.d[k]; !found {
		bucket = make([]*GpgPrimaryKey, 0, 1)
	}
	bucket = append(bucket, v)
	bd.d[k] = bucket
}

func (bd BucketDict) Get(k string) []*GpgPrimaryKey {
	k = strings.ToLower(k)
	ret, found := bd.d[k]
	if !found {
		ret = nil
	}
	return ret
}

func (bd BucketDict) Get0Or1(k string) (ret *GpgPrimaryKey, err error) {
	v := bd.Get(k)
	if len(v) > 1 {
		err = GpgError{fmt.Sprintf("Wanted a unique lookup but got %d objects for key %s", len(v), k)}
	} else if len(v) == 1 {
		ret = v[0]
	}
	return
}

//=============================================================================

func Uniquify(inp []string) []string {
	m := make(map[string]bool)
	for _, s := range inp {
		m[strings.ToLower(s)] = true
	}
	ret := make([]string, 0, len(inp))
	for k, _ := range m {
		ret = append(ret, k)
	}
	return ret
}

//=============================================================================

type GpgBaseKey struct {
	Type        string
	Trust       string
	Bits        int
	Algo        int
	Id64        string
	Created     int
	Expires     int
	fingerprint *PgpFingerprint
}

func (k *GpgBaseKey) ParseBase(line *GpgIndexLine) (err error) {
	if line.Len() < 12 {
		err = GpgIndexError{line.lineno, "Not enough fields (need 12)"}
		return
	}

	k.Type = line.At(0)
	k.Trust = line.At(1)
	k.Id64 = line.At(4)

	flexiAtoi := func(s string) (int, error) {
		if len(s) == 0 {
			return 0, nil
		} else {
			return strconv.Atoi(s)
		}
	}

	if k.Bits, err = strconv.Atoi(line.At(2)); err != nil {
	} else if k.Algo, err = strconv.Atoi(line.At(3)); err != nil {
	} else if k.Created, err = strconv.Atoi(line.At(5)); err != nil {
	} else if k.Expires, err = flexiAtoi(line.At(6)); err != nil {
	}

	return
}

//=============================================================================

type GpgFingerprinter interface {
	SetFingerprint(pgp *PgpFingerprint)
}

type GpgPrimaryKey struct {
	GpgBaseKey
	subkeys    []*GpgSubKey
	identities []*Identity
	top        GpgFingerprinter
}

func (k *GpgBaseKey) SetFingerprint(pgp *PgpFingerprint) {
	k.fingerprint = pgp
}

func (k *GpgPrimaryKey) Parse(l *GpgIndexLine) (err error) {
	if err = k.ParseBase(l); err != nil {
	} else if err = k.AddUid(l); err != nil {
	}
	return
}

func NewGpgPrimaryKey() *GpgPrimaryKey {
	ret := &GpgPrimaryKey{}
	ret.top = ret
	return ret
}

func ParseGpgPrimaryKey(l *GpgIndexLine) (key *GpgPrimaryKey, err error) {
	key = NewGpgPrimaryKey()
	err = key.Parse(l)
	return
}

func (k *GpgPrimaryKey) AddUid(l *GpgIndexLine) (err error) {
	var id *Identity
	if f := l.At(9); len(f) == 0 {
	} else if id, err = ParseIdentity(f); err != nil {
	} else {
		k.identities = append(k.identities, id)
	}
	if err != nil {
		err = ErrorToGpgIndexError(l.lineno, err)
	}
	return
}

func (k *GpgPrimaryKey) AddFingerprint(l *GpgIndexLine) (err error) {
	var fp *PgpFingerprint
	if f := l.At(9); len(f) == 0 {
		err = fmt.Errorf("no fingerprint given")
	} else if fp, err = PgpFingerprintFromHex(f); err == nil {
		k.top.SetFingerprint(fp)
	}
	if err != nil {
		err = ErrorToGpgIndexError(l.lineno, err)
	}
	return
}

func (k *GpgPrimaryKey) GetFingerprint() *PgpFingerprint {
	return k.fingerprint
}

func (k *GpgPrimaryKey) GetEmails() []string {
	ret := make([]string, 0, len(k.identities))
	for _, i := range k.identities {
		ret = append(ret, i.Email)
	}
	return ret
}

func (k *GpgPrimaryKey) GetAllId64s() []string {
	var ret []string
	add := func(fp *PgpFingerprint) {
		if fp != nil {
			ret = append(ret, fp.ToKeyId())
		}
	}
	add(k.GetFingerprint())
	for _, sk := range k.subkeys {
		add(sk.fingerprint)
	}
	return ret
}

func (g *GpgPrimaryKey) AddSubkey(l *GpgIndexLine) (err error) {
	var sk *GpgSubKey
	if sk, err = ParseGpgSubKey(l); err == nil {
		g.subkeys = append(g.subkeys, sk)
		g.top = sk
	}
	return
}

func (g *GpgPrimaryKey) ToKey() *GpgPrimaryKey { return g }

func (p *GpgPrimaryKey) AddLine(l *GpgIndexLine) (err error) {
	if l.Len() < 2 {
		err = GpgIndexError{l.lineno, "too few fields"}
	} else {
		f := l.At(0)
		switch f {
		case "fpr":
			err = p.AddFingerprint(l)
		case "uid":
			err = p.AddUid(l)
		case "uat":
		case "sub", "ssb":
			err = p.AddSubkey(l)
		default:
			err = GpgIndexError{l.lineno, fmt.Sprintf("Unknown subfield: %s", f)}
		}

	}
	return err
}

//=============================================================================

type GpgSubKey struct {
	GpgBaseKey
}

func ParseGpgSubKey(l *GpgIndexLine) (sk *GpgSubKey, err error) {
	sk = &GpgSubKey{}
	err = sk.ParseBase(l)
	return
}

//=============================================================================

type GpgIndexElement interface {
	ToKey() *GpgPrimaryKey
}

type GpgKeyIndex struct {
	Keys                        []*GpgPrimaryKey
	Emails, Fingerprints, Id64s *BucketDict
}

func (ki GpgKeyIndex) Len() int {
	return len(ki.Keys)
}
func (ki GpgKeyIndex) Swap(i, j int) {
	ki.Keys[i], ki.Keys[j] = ki.Keys[j], ki.Keys[i]
}
func (ki GpgKeyIndex) Less(i, j int) bool {
	a, b := ki.Keys[i], ki.Keys[j]

	d := len(a.identities) - len(b.identities)
	if d > 0 {
		return true
	} else if d < 0 {
		return false
	}
}

func NewGpgKeyIndex() *GpgKeyIndex {
	return &GpgKeyIndex{
		Keys:         make([]*GpgPrimaryKey, 0, 1),
		Emails:       NewBuckDict(),
		Fingerprints: NewBuckDict(),
		Id64s:        NewBuckDict(),
	}
}

func (ki *GpgKeyIndex) IndexKey(k *GpgPrimaryKey) {
	ki.Keys = append(ki.Keys, k)
	if fp := k.GetFingerprint(); fp != nil {
		ki.Fingerprints.Add(fp.ToString(), k)
	}
	for _, e := range Uniquify(k.GetEmails()) {
		ki.Emails.Add(e, k)
	}
	for _, i := range Uniquify(k.GetAllId64s()) {
		ki.Id64s.Add(i, k)
	}
}

func (k *GpgKeyIndex) PushElement(e GpgIndexElement) {
	if key := e.ToKey(); key != nil {
		k.IndexKey(key)
	}
}

func (ki *GpgKeyIndex) AllFingerprints() []PgpFingerprint {
	ret := make([]PgpFingerprint, 0, 1)
	for _, k := range ki.Keys {
		if fp := k.GetFingerprint(); fp != nil {
			ret = append(ret, *fp)
		}
	}
	return ret
}

//=============================================================================

type GpgIndexLine struct {
	v      []string
	lineno int
}

func (g GpgIndexLine) Len() int        { return len(g.v) }
func (g GpgIndexLine) At(i int) string { return g.v[i] }

func ParseLine(s string, i int) (ret *GpgIndexLine, err error) {
	s = strings.TrimSpace(s)
	v := strings.Split(s, ":")
	if v == nil {
		err = GpgError{fmt.Sprintf("%d: Bad line; split failed", i)}
	} else {
		ret = &GpgIndexLine{v, i}
	}
	return
}

func (l GpgIndexLine) IsNewKey() bool {
	return len(l.v) > 0 && (l.v[0] == "sec" || l.v[0] == "pub")
}

//=============================================================================

type GpgIndexParser struct {
	warnings Warnings
	putback  *GpgIndexLine
	src      *bufio.Reader
	eof      bool
	lineno   int
}

func NewGpgIndexParser() *GpgIndexParser {
	return &GpgIndexParser{
		eof:     false,
		lineno:  0,
		putback: nil,
	}
}

func (p *GpgIndexParser) Warn(w Warning) {
	p.warnings.Push(w)
}

func (p *GpgIndexParser) ParseElement() (ret GpgIndexElement, err error) {
	var line *GpgIndexLine
	line, err = p.GetLine()
	if err != nil || line == nil {
	} else if line.IsNewKey() {
		ret, err = p.ParseKey(line)
	}
	return
}

func (p *GpgIndexParser) ParseKey(l *GpgIndexLine) (ret *GpgPrimaryKey, err error) {
	var line *GpgIndexLine
	ret, err = ParseGpgPrimaryKey(l)
	done := false
	for !done && err == nil && !p.isEof() {
		if line, err = p.GetLine(); line == nil || err != nil {
		} else if line.IsNewKey() {
			p.PutbackLine(line)
			done = true
		} else if e2 := ret.AddLine(line); e2 == nil {
		} else {
			p.warnings.Push(ErrorToWarning(e2))
		}
	}
	return
}

func (p *GpgIndexParser) GetLine() (ret *GpgIndexLine, err error) {
	if p.putback != nil {
		ret = p.putback
		p.putback = nil
	} else if p.isEof() {
	} else if s, e2 := p.src.ReadString(byte('\n')); e2 == nil {
		p.lineno++
		ret, err = ParseLine(s, p.lineno)
	} else if e2 == io.EOF {
		p.eof = true
	} else {
		err = e2
	}
	return
}

func (p *GpgIndexParser) PutbackLine(line *GpgIndexLine) {
	p.putback = line
}

func (p GpgIndexParser) isEof() bool { return p.eof }

func (p *GpgIndexParser) Parse(stream io.Reader) (ki *GpgKeyIndex, err error) {
	p.src = bufio.NewReader(stream)
	ki = NewGpgKeyIndex()
	for err == nil && !p.isEof() {
		var el GpgIndexElement
		if el, err = p.ParseElement(); err == nil && el != nil {
			ki.PushElement(el)
		}
	}
	return
}

//=============================================================================

func ParseGpgIndexStream(stream io.Reader) (ki *GpgKeyIndex, err error, w Warnings) {
	eng := NewGpgIndexParser()
	ki, err = eng.Parse(stream)
	w = eng.warnings
	return
}

//=============================================================================

func (g *GpgCLI) Index(secret bool, query string) (ki *GpgKeyIndex, err error, w Warnings) {
	var k string
	if secret {
		k = "-K"
	} else {
		k = "-k"
	}
	args := []string{"--with-colons", "--fingerprint", k}
	if len(query) > 0 {
		args = append(args, query)
	}
	garg := RunGpg2Arg{
		Arguments: args,
		Stdout:    true,
	}
	if res := g.Run2(garg); res.Err != nil {
		err = res.Err
	} else if ki, err, w = ParseGpgIndexStream(res.Stdout); err != nil {
	} else {
		err = res.Wait()
	}
	return
}

//=============================================================================
<commit_msg>fixed the index<commit_after>package libkb

import (
	"bufio"
	"fmt"
	"golang.org/x/crypto/openpgp/packet"
	"io"
	"sort"
	"strconv"
	"strings"
	"time"
)

//=============================================================================

type BucketDict struct {
	d map[string][]*GpgPrimaryKey
}

func NewBuckDict() *BucketDict {
	return &BucketDict{
		d: make(map[string][]*GpgPrimaryKey),
	}
}

func (bd *BucketDict) Add(k string, v *GpgPrimaryKey) {
	k = strings.ToLower(k)
	var bucket []*GpgPrimaryKey
	var found bool
	if bucket, found = bd.d[k]; !found {
		bucket = make([]*GpgPrimaryKey, 0, 1)
	}
	bucket = append(bucket, v)
	bd.d[k] = bucket
}

func (bd BucketDict) Get(k string) []*GpgPrimaryKey {
	k = strings.ToLower(k)
	ret, found := bd.d[k]
	if !found {
		ret = nil
	}
	return ret
}

func (bd BucketDict) Get0Or1(k string) (ret *GpgPrimaryKey, err error) {
	v := bd.Get(k)
	if len(v) > 1 {
		err = GpgError{fmt.Sprintf("Wanted a unique lookup but got %d objects for key %s", len(v), k)}
	} else if len(v) == 1 {
		ret = v[0]
	}
	return
}

//=============================================================================

func Uniquify(inp []string) []string {
	m := make(map[string]bool)
	for _, s := range inp {
		m[strings.ToLower(s)] = true
	}
	ret := make([]string, 0, len(inp))
	for k, _ := range m {
		ret = append(ret, k)
	}
	return ret
}

//=============================================================================

type GpgBaseKey struct {
	Type        string
	Trust       string
	Bits        int
	Algo        int
	Id64        string
	Created     int
	Expires     int
	fingerprint *PgpFingerprint
}

func (k GpgBaseKey) AlgoString() string {
	switch packet.PublicKeyAlgorithm(k.Algo) {
	case packet.PubKeyAlgoDSA:
		return "D"
	case packet.PubKeyAlgoRSA:
		return "R"
	case packet.PubKeyAlgoECDSA:
		return "E"
	default:
		return "?"
	}
}

func (k GpgBaseKey) ExpirationString() string {
	if k.Expires == 0 {
		return "never"
	} else {
		layout := "2006-01-02"
		return time.Unix(int64(k.Expires), 0).Format(layout)
	}
}

func (k *GpgBaseKey) ParseBase(line *GpgIndexLine) (err error) {
	if line.Len() < 12 {
		err = GpgIndexError{line.lineno, "Not enough fields (need 12)"}
		return
	}

	k.Type = line.At(0)
	k.Trust = line.At(1)
	k.Id64 = line.At(4)

	flexiAtoi := func(s string) (int, error) {
		if len(s) == 0 {
			return 0, nil
		} else {
			return strconv.Atoi(s)
		}
	}

	if k.Bits, err = strconv.Atoi(line.At(2)); err != nil {
	} else if k.Algo, err = strconv.Atoi(line.At(3)); err != nil {
	} else if k.Created, err = strconv.Atoi(line.At(5)); err != nil {
	} else if k.Expires, err = flexiAtoi(line.At(6)); err != nil {
	}

	return
}

//=============================================================================

type GpgFingerprinter interface {
	SetFingerprint(pgp *PgpFingerprint)
}

type GpgPrimaryKey struct {
	GpgBaseKey
	subkeys    []*GpgSubKey
	identities []*Identity
	top        GpgFingerprinter
}

func (k *GpgPrimaryKey) ToRow(i int) []string {
	v := []string{
		fmt.Sprintf("(%d)", i),
		fmt.Sprintf("%d%s", k.Bits, k.AlgoString()),
		k.fingerprint.ToKeyId(),
		k.ExpirationString(),
	}
	for _, i := range k.identities {
		v = append(v, i.Email)
	}
	return v
}

func (k *GpgBaseKey) SetFingerprint(pgp *PgpFingerprint) {
	k.fingerprint = pgp
}

func (k *GpgPrimaryKey) Parse(l *GpgIndexLine) (err error) {
	if err = k.ParseBase(l); err != nil {
	} else if err = k.AddUid(l); err != nil {
	}
	return
}

func NewGpgPrimaryKey() *GpgPrimaryKey {
	ret := &GpgPrimaryKey{}
	ret.top = ret
	return ret
}

func ParseGpgPrimaryKey(l *GpgIndexLine) (key *GpgPrimaryKey, err error) {
	key = NewGpgPrimaryKey()
	err = key.Parse(l)
	return
}

func (k *GpgPrimaryKey) AddUid(l *GpgIndexLine) (err error) {
	var id *Identity
	if f := l.At(9); len(f) == 0 {
	} else if id, err = ParseIdentity(f); err != nil {
	} else {
		k.identities = append(k.identities, id)
	}
	if err != nil {
		err = ErrorToGpgIndexError(l.lineno, err)
	}
	return
}

func (k *GpgPrimaryKey) AddFingerprint(l *GpgIndexLine) (err error) {
	var fp *PgpFingerprint
	if f := l.At(9); len(f) == 0 {
		err = fmt.Errorf("no fingerprint given")
	} else if fp, err = PgpFingerprintFromHex(f); err == nil {
		k.top.SetFingerprint(fp)
	}
	if err != nil {
		err = ErrorToGpgIndexError(l.lineno, err)
	}
	return
}

func (k *GpgPrimaryKey) GetFingerprint() *PgpFingerprint {
	return k.fingerprint
}

func (k *GpgPrimaryKey) GetEmails() []string {
	ret := make([]string, 0, len(k.identities))
	for _, i := range k.identities {
		ret = append(ret, i.Email)
	}
	return ret
}

func (k *GpgPrimaryKey) GetAllId64s() []string {
	var ret []string
	add := func(fp *PgpFingerprint) {
		if fp != nil {
			ret = append(ret, fp.ToKeyId())
		}
	}
	add(k.GetFingerprint())
	for _, sk := range k.subkeys {
		add(sk.fingerprint)
	}
	return ret
}

func (g *GpgPrimaryKey) AddSubkey(l *GpgIndexLine) (err error) {
	var sk *GpgSubKey
	if sk, err = ParseGpgSubKey(l); err == nil {
		g.subkeys = append(g.subkeys, sk)
		g.top = sk
	}
	return
}

func (g *GpgPrimaryKey) ToKey() *GpgPrimaryKey { return g }

func (p *GpgPrimaryKey) AddLine(l *GpgIndexLine) (err error) {
	if l.Len() < 2 {
		err = GpgIndexError{l.lineno, "too few fields"}
	} else {
		f := l.At(0)
		switch f {
		case "fpr":
			err = p.AddFingerprint(l)
		case "uid":
			err = p.AddUid(l)
		case "uat":
		case "sub", "ssb":
			err = p.AddSubkey(l)
		default:
			err = GpgIndexError{l.lineno, fmt.Sprintf("Unknown subfield: %s", f)}
		}

	}
	return err
}

//=============================================================================

type GpgSubKey struct {
	GpgBaseKey
}

func ParseGpgSubKey(l *GpgIndexLine) (sk *GpgSubKey, err error) {
	sk = &GpgSubKey{}
	err = sk.ParseBase(l)
	return
}

//=============================================================================

type GpgIndexElement interface {
	ToKey() *GpgPrimaryKey
}

type GpgKeyIndex struct {
	Keys                        []*GpgPrimaryKey
	Emails, Fingerprints, Id64s *BucketDict
}

func (ki *GpgKeyIndex) Len() int {
	return len(ki.Keys)
}
func (ki *GpgKeyIndex) Swap(i, j int) {
	ki.Keys[i], ki.Keys[j] = ki.Keys[j], ki.Keys[i]
}
func (ki *GpgKeyIndex) Less(i, j int) bool {
	a, b := ki.Keys[i], ki.Keys[j]
	if len(a.identities) > len(b.identities) {
		return true
	} else if len(a.identities) < len(b.identities) {
		return false
	} else if a.Expires == 0 {
		return true
	} else if b.Expires == 0 {
		return false
	} else if a.Expires > b.Expires {
		return true
	} else {
		return false
	}
}

func (p *GpgKeyIndex) GetRowFunc() func() []string {
	i := 0
	return func() []string {
		if i >= len(p.Keys) {
			return nil
		} else {
			ret := p.Keys[i].ToRow(i + 1)
			i++
			return ret
		}
	}
}

func (ki *GpgKeyIndex) Sort() {
	sort.Sort(ki)
}

func NewGpgKeyIndex() *GpgKeyIndex {
	return &GpgKeyIndex{
		Keys:         make([]*GpgPrimaryKey, 0, 1),
		Emails:       NewBuckDict(),
		Fingerprints: NewBuckDict(),
		Id64s:        NewBuckDict(),
	}
}

func (ki *GpgKeyIndex) IndexKey(k *GpgPrimaryKey) {
	ki.Keys = append(ki.Keys, k)
	if fp := k.GetFingerprint(); fp != nil {
		ki.Fingerprints.Add(fp.ToString(), k)
	}
	for _, e := range Uniquify(k.GetEmails()) {
		ki.Emails.Add(e, k)
	}
	for _, i := range Uniquify(k.GetAllId64s()) {
		ki.Id64s.Add(i, k)
	}
}

func (k *GpgKeyIndex) PushElement(e GpgIndexElement) {
	if key := e.ToKey(); key != nil {
		k.IndexKey(key)
	}
}

func (ki *GpgKeyIndex) AllFingerprints() []PgpFingerprint {
	ret := make([]PgpFingerprint, 0, 1)
	for _, k := range ki.Keys {
		if fp := k.GetFingerprint(); fp != nil {
			ret = append(ret, *fp)
		}
	}
	return ret
}

//=============================================================================

type GpgIndexLine struct {
	v      []string
	lineno int
}

func (g GpgIndexLine) Len() int        { return len(g.v) }
func (g GpgIndexLine) At(i int) string { return g.v[i] }

func ParseLine(s string, i int) (ret *GpgIndexLine, err error) {
	s = strings.TrimSpace(s)
	v := strings.Split(s, ":")
	if v == nil {
		err = GpgError{fmt.Sprintf("%d: Bad line; split failed", i)}
	} else {
		ret = &GpgIndexLine{v, i}
	}
	return
}

func (l GpgIndexLine) IsNewKey() bool {
	return len(l.v) > 0 && (l.v[0] == "sec" || l.v[0] == "pub")
}

//=============================================================================

type GpgIndexParser struct {
	warnings Warnings
	putback  *GpgIndexLine
	src      *bufio.Reader
	eof      bool
	lineno   int
}

func NewGpgIndexParser() *GpgIndexParser {
	return &GpgIndexParser{
		eof:     false,
		lineno:  0,
		putback: nil,
	}
}

func (p *GpgIndexParser) Warn(w Warning) {
	p.warnings.Push(w)
}

func (p *GpgIndexParser) ParseElement() (ret GpgIndexElement, err error) {
	var line *GpgIndexLine
	line, err = p.GetLine()
	if err != nil || line == nil {
	} else if line.IsNewKey() {
		ret, err = p.ParseKey(line)
	}
	return
}

func (p *GpgIndexParser) ParseKey(l *GpgIndexLine) (ret *GpgPrimaryKey, err error) {
	var line *GpgIndexLine
	ret, err = ParseGpgPrimaryKey(l)
	done := false
	for !done && err == nil && !p.isEof() {
		if line, err = p.GetLine(); line == nil || err != nil {
		} else if line.IsNewKey() {
			p.PutbackLine(line)
			done = true
		} else if e2 := ret.AddLine(line); e2 == nil {
		} else {
			p.warnings.Push(ErrorToWarning(e2))
		}
	}
	return
}

func (p *GpgIndexParser) GetLine() (ret *GpgIndexLine, err error) {
	if p.putback != nil {
		ret = p.putback
		p.putback = nil
	} else if p.isEof() {
	} else if s, e2 := p.src.ReadString(byte('\n')); e2 == nil {
		p.lineno++
		ret, err = ParseLine(s, p.lineno)
	} else if e2 == io.EOF {
		p.eof = true
	} else {
		err = e2
	}
	return
}

func (p *GpgIndexParser) PutbackLine(line *GpgIndexLine) {
	p.putback = line
}

func (p GpgIndexParser) isEof() bool { return p.eof }

func (p *GpgIndexParser) Parse(stream io.Reader) (ki *GpgKeyIndex, err error) {
	p.src = bufio.NewReader(stream)
	ki = NewGpgKeyIndex()
	for err == nil && !p.isEof() {
		var el GpgIndexElement
		if el, err = p.ParseElement(); err == nil && el != nil {
			ki.PushElement(el)
		}
	}
	ki.Sort()
	return
}

//=============================================================================

func ParseGpgIndexStream(stream io.Reader) (ki *GpgKeyIndex, err error, w Warnings) {
	eng := NewGpgIndexParser()
	ki, err = eng.Parse(stream)
	w = eng.warnings
	return
}

//=============================================================================

func (g *GpgCLI) Index(secret bool, query string) (ki *GpgKeyIndex, err error, w Warnings) {
	var k string
	if secret {
		k = "-K"
	} else {
		k = "-k"
	}
	args := []string{"--with-colons", "--fingerprint", k}
	if len(query) > 0 {
		args = append(args, query)
	}
	garg := RunGpg2Arg{
		Arguments: args,
		Stdout:    true,
	}
	if res := g.Run2(garg); res.Err != nil {
		err = res.Err
	} else if ki, err, w = ParseGpgIndexStream(res.Stdout); err != nil {
	} else {
		err = res.Wait()
	}
	return
}

//=============================================================================
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"go/format"
	"io"
	"path"
	"strings"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/protoc-gen-go/descriptor"
	plugin "github.com/golang/protobuf/protoc-gen-go/plugin"
)

type Generator struct{}

func New() *Generator {
	return &Generator{}
}

func p(w io.Writer, format string, args ...interface{}) {
	if w == nil {
		return
	}
	if _, err := fmt.Fprintf(w, format, args...); err != nil {
		panic(err)
	}
	if _, err := fmt.Fprintln(w); err != nil {
		panic(err)
	}
}

func (g *Generator) Generate(req *plugin.CodeGeneratorRequest) (*plugin.CodeGeneratorResponse, error) {
	// Filter files with Service.
	filterdFiles := make([]*descriptor.FileDescriptorProto, 0)
	for _, f := range req.GetProtoFile() {
		if len(f.GetService()) == 0 {
			continue
		}
		filterdFiles = append(filterdFiles, f)
	}

	// Generate response files from proto files.
	respFiles := make([]*plugin.CodeGeneratorResponse_File, 0)
	for _, f := range filterdFiles {
		respFiles = append(respFiles, g.genRespFile(f))
	}

	// Format response files content.
	for _, f := range respFiles {
		content, err := format.Source([]byte(f.GetContent()))
		if err != nil {
			return nil, err
		}
		f.Content = proto.String(string(content))
	}

	return &plugin.CodeGeneratorResponse{
		File: respFiles,
	}, nil
}

func (g *Generator) genRespFile(file *descriptor.FileDescriptorProto) *plugin.CodeGeneratorResponse_File {
	buf := &bytes.Buffer{}
	g.writePackage(buf, file)
	g.writeImports(buf)
	for _, service := range file.GetService() {
		g.writeService(buf, service)
		for _, method := range service.GetMethod() {
			if method.GetServerStreaming() {
				continue
			}
			g.writeMethod(buf, service, method)
		}
	}
	return &plugin.CodeGeneratorResponse_File{
		Name:    proto.String(basename(file.GetName()) + ".http.go"),
		Content: proto.String(buf.String()),
	}
}

func (g *Generator) writePackage(w io.Writer, f *descriptor.FileDescriptorProto) {
	p(w, "package %s", f.Options.GetGoPackage())
	p(w, "")
}

func (g *Generator) writeImports(w io.Writer) {
	p(w, "import (")
	p(w, "	\"bytes\"")
	p(w, "	\"context\"")
	p(w, "	\"encoding/json\"")
	p(w, "	\"fmt\"")
	p(w, "	\"io\"")
	p(w, "	\"io/ioutil\"")
	p(w, "	\"net/http\"")
	p(w, "")
	p(w, "	\"github.com/golang/protobuf/jsonpb\"")
	p(w, "	\"github.com/golang/protobuf/proto\"")
	p(w, ")")
	p(w, "")
}

func (g *Generator) writeService(w io.Writer, s *descriptor.ServiceDescriptorProto) {
	p(w, "type %sHandler struct{", s.GetName())
	p(w, "	srv %sServer", s.GetName())
	p(w, "}")
	p(w, "")
	p(w, "func New%sHandler(srv %sServer) *%sHandler {", s.GetName(), s.GetName(), s.GetName())
	p(w, "	return &%sHandler{", s.GetName())
	p(w, "		srv: srv,")
	p(w, "	}")
	p(w, "}")
	p(w, "")
}

func (g *Generator) writeMethod(w io.Writer, s *descriptor.ServiceDescriptorProto, m *descriptor.MethodDescriptorProto) {
	p(w, "func (h *%sHandler) %s(cb func(ctx context.Context, w http.ResponseWriter, r *http.Request, arg, ret proto.Message, err error)) http.HandlerFunc {", s.GetName(), m.GetName())
	p(w, "	if cb == nil {")
	p(w, "		cb = func(ctx context.Context, w http.ResponseWriter, r *http.Request, arg, ret proto.Message, err error) {")
	p(w, "			if err != nil {")
	p(w, "				w.WriteHeader(http.StatusInternalServerError)")
	p(w, "				fmt.Fprintf(w, \"%%v: arg = %%v: ret = %%v\", err, arg, ret)")
	p(w, "			}")
	p(w, "		}")
	p(w, "	}")
	p(w, "	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {")
	p(w, "		ctx := r.Context()")
	p(w, "")
	p(w, "		body, err := ioutil.ReadAll(r.Body)")
	p(w, "		if err != nil {")
	p(w, "			cb(ctx, w, r, nil, nil, err)")
	p(w, "			return")
	p(w, "		}")
	p(w, "")
	p(w, "		var arg *%s", ioname(m.GetInputType()))
	p(w, "")
	p(w, "		contentType := r.Header.Get(\"Content-Type\")")
	p(w, "		switch contentType {")
	p(w, "		case \"application/protobuf\", \"application/x-protobuf\":")
	p(w, "			if err := proto.Unmarshal(body, arg); err != nil {")
	p(w, "				cb(ctx, w, r, nil, nil, err)")
	p(w, "				return")
	p(w, "			}")
	p(w, "		case \"application/json\":")
	p(w, "			if err := jsonpb.Unmarshal(bytes.NewBuffer(body), arg); err != nil {")
	p(w, "				cb(ctx, w, r, nil, nil, err)")
	p(w, "				return")
	p(w, "			}")
	p(w, "		default:")
	p(w, "			w.WriteHeader(http.StatusUnsupportedMediaType)")
	p(w, "			if _, err := fmt.Fprintf(w, \"Unsupported Content-Type: %%s\", contentType); err != nil {")
	p(w, "				cb(ctx, w, r, nil, nil, err)")
	p(w, "			}")
	p(w, "			return")
	p(w, "		}")
	p(w, "")
	p(w, "		ret, err := h.srv.%s(ctx, arg)", m.GetName())
	p(w, "		if err != nil {")
	p(w, "			cb(ctx, w, r, arg, ret, err)")
	p(w, "			return")
	p(w, "		}")
	p(w, "")
	p(w, "		switch contentType {")
	p(w, "		case \"application/protobuf\", \"application/x-protobuf\":")
	p(w, "			buf, err := proto.Marshal(ret)")
	p(w, "			if err != nil {")
	p(w, "				cb(ctx, w, r, arg, ret, err)")
	p(w, "				return")
	p(w, "			}")
	p(w, "			if _, err := io.Copy(w, bytes.NewBuffer(buf)); err != nil {")
	p(w, "				cb(ctx, w, r, arg, ret, err)")
	p(w, "				return")
	p(w, "			}")
	p(w, "		case \"application/json\":")
	p(w, "			if err := json.NewEncoder(w).Encode(ret); err != nil {")
	p(w, "				cb(ctx, w, r, arg, ret, err)")
	p(w, "				return")
	p(w, "			}")
	p(w, "		default:")
	p(w, "			w.WriteHeader(http.StatusUnsupportedMediaType)")
	p(w, "			if _, err := fmt.Fprintf(w, \"Unsupported Content-Type: %%s\", contentType); err != nil {")
	p(w, "				cb(ctx, w, r, nil, nil, err)")
	p(w, "			}")
	p(w, "			return")
	p(w, "		}")
	p(w, "")
	p(w, "		cb(ctx, w, r, arg, ret, err)")
	p(w, "	})")
	p(w, "}")
	p(w, "")
}

func basename(name string) string {
	if ext := path.Ext(name); ext == ".proto" || ext == ".protodevel" {
		name = name[:len(name)-len(ext)]
	}
	return name
}

func ioname(name string) string {
	s := strings.Split(name, ".")
	return s[len(s)-1]
}
<commit_msg>Change generate code for testdata<commit_after>package main

import (
	"bytes"
	"fmt"
	"go/format"
	"io"
	"path"
	"strings"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/protoc-gen-go/descriptor"
	plugin "github.com/golang/protobuf/protoc-gen-go/plugin"
)

type Generator struct{}

func New() *Generator {
	return &Generator{}
}

func p(w io.Writer, format string, args ...interface{}) {
	if w == nil {
		return
	}
	if _, err := fmt.Fprintf(w, format, args...); err != nil {
		panic(err)
	}
	if _, err := fmt.Fprintln(w); err != nil {
		panic(err)
	}
}

func (g *Generator) Generate(req *plugin.CodeGeneratorRequest) (*plugin.CodeGeneratorResponse, error) {
	// Filter files with Service.
	filterdFiles := make([]*descriptor.FileDescriptorProto, 0)
	for _, f := range req.GetProtoFile() {
		if len(f.GetService()) == 0 {
			continue
		}
		filterdFiles = append(filterdFiles, f)
	}

	// Generate response files from proto files.
	respFiles := make([]*plugin.CodeGeneratorResponse_File, 0)
	for _, f := range filterdFiles {
		respFiles = append(respFiles, g.genRespFile(f))
	}

	// Format response files content.
	for _, f := range respFiles {
		content, err := format.Source([]byte(f.GetContent()))
		if err != nil {
			return nil, err
		}
		f.Content = proto.String(string(content))
	}

	return &plugin.CodeGeneratorResponse{
		File: respFiles,
	}, nil
}

func (g *Generator) genRespFile(file *descriptor.FileDescriptorProto) *plugin.CodeGeneratorResponse_File {
	buf := &bytes.Buffer{}
	g.writePackage(buf, file)
	g.writeImports(buf)
	for _, service := range file.GetService() {
		g.writeService(buf, service)
		for _, method := range service.GetMethod() {
			if method.GetServerStreaming() {
				continue
			}
			g.writeMethod(buf, service, method)
		}
	}
	return &plugin.CodeGeneratorResponse_File{
		Name:    proto.String(basename(file.GetName()) + ".http.go"),
		Content: proto.String(buf.String()),
	}
}

func (g *Generator) writePackage(w io.Writer, f *descriptor.FileDescriptorProto) {
	p(w, "package %s", f.Options.GetGoPackage())
	p(w, "")
}

func (g *Generator) writeImports(w io.Writer) {
	p(w, "import (")
	p(w, "	\"bytes\"")
	p(w, "	\"context\"")
	p(w, "	\"encoding/json\"")
	p(w, "	\"fmt\"")
	p(w, "	\"io\"")
	p(w, "	\"io/ioutil\"")
	p(w, "	\"net/http\"")
	p(w, "")
	p(w, "	\"github.com/golang/protobuf/jsonpb\"")
	p(w, "	\"github.com/golang/protobuf/proto\"")
	p(w, ")")
	p(w, "")
}

func (g *Generator) writeService(w io.Writer, s *descriptor.ServiceDescriptorProto) {
	p(w, "type %sHandler struct{", s.GetName())
	p(w, "	srv %sServer", s.GetName())
	p(w, "}")
	p(w, "")
	p(w, "func New%sHandler(srv %sServer) *%sHandler {", s.GetName(), s.GetName(), s.GetName())
	p(w, "	return &%sHandler{", s.GetName())
	p(w, "		srv: srv,")
	p(w, "	}")
	p(w, "}")
	p(w, "")
}

func (g *Generator) writeMethod(w io.Writer, s *descriptor.ServiceDescriptorProto, m *descriptor.MethodDescriptorProto) {
	p(w, "func (h *%sHandler) %s(cb func(ctx context.Context, w http.ResponseWriter, r *http.Request, arg, ret proto.Message, err error)) http.HandlerFunc {", s.GetName(), m.GetName())
	p(w, "	if cb == nil {")
	p(w, "		cb = func(ctx context.Context, w http.ResponseWriter, r *http.Request, arg, ret proto.Message, err error) {")
	p(w, "			if err != nil {")
	p(w, "				w.WriteHeader(http.StatusInternalServerError)")
	p(w, "				fmt.Fprintf(w, \"%%v: arg = %%v: ret = %%v\", err, arg, ret)")
	p(w, "			}")
	p(w, "		}")
	p(w, "	}")
	p(w, "	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {")
	p(w, "		var (")
	p(w, "			ctx = r.Context()")
	p(w, "			arg *HelloRequest")
	p(w, "			ret *HelloReply")
	p(w, "			err error")
	p(w, "		)")
	p(w, "		defer cb(ctx, w, r, arg, ret, err)")
	p(w, "")
	p(w, "		body, err := ioutil.ReadAll(r.Body)")
	p(w, "		if err != nil {")
	p(w, "			return")
	p(w, "		}")
	p(w, "")
	p(w, "		contentType := r.Header.Get(\"Content-Type\")")
	p(w, "		switch contentType {")
	p(w, "		case \"application/protobuf\", \"application/x-protobuf\":")
	p(w, "			if err = proto.Unmarshal(body, arg); err != nil {")
	p(w, "				return")
	p(w, "			}")
	p(w, "		case \"application/json\":")
	p(w, "			if err = jsonpb.Unmarshal(bytes.NewBuffer(body), arg); err != nil {")
	p(w, "				return")
	p(w, "			}")
	p(w, "		default:")
	p(w, "			w.WriteHeader(http.StatusUnsupportedMediaType)")
	p(w, "			_, err = fmt.Fprintf(w, \"Unsupported Content-Type: %%s\", contentType)")
	p(w, "			return")
	p(w, "		}")
	p(w, "")
	p(w, "		ret, err = h.srv.%s(ctx, arg)", m.GetName())
	p(w, "		if err != nil {")
	p(w, "			return")
	p(w, "		}")
	p(w, "")
	p(w, "		switch contentType {")
	p(w, "		case \"application/protobuf\", \"application/x-protobuf\":")
	p(w, "			buf, err := proto.Marshal(ret)")
	p(w, "			if err != nil {")
	p(w, "				return")
	p(w, "			}")
	p(w, "			if _, err = io.Copy(w, bytes.NewBuffer(buf)); err != nil {")
	p(w, "				return")
	p(w, "			}")
	p(w, "		case \"application/json\":")
	p(w, "			if err = json.NewEncoder(w).Encode(ret); err != nil {")
	p(w, "				return")
	p(w, "			}")
	p(w, "		default:")
	p(w, "			w.WriteHeader(http.StatusUnsupportedMediaType)")
	p(w, "			_, err = fmt.Fprintf(w, \"Unsupported Content-Type: %%s\", contentType)")
	p(w, "			return")
	p(w, "		}")
	p(w, "	})")
	p(w, "}")
	p(w, "")
}

func basename(name string) string {
	if ext := path.Ext(name); ext == ".proto" || ext == ".protodevel" {
		name = name[:len(name)-len(ext)]
	}
	return name
}

func ioname(name string) string {
	s := strings.Split(name, ".")
	return s[len(s)-1]
}
<|endoftext|>