<commit_before>package release

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"os/exec"
	"time"

	"github.com/ghodss/yaml"
	"github.com/go-kit/kit/log"
	"k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/helm/pkg/chartutil"
	k8shelm "k8s.io/helm/pkg/helm"
	hapi_release "k8s.io/helm/pkg/proto/hapi/release"

	"github.com/weaveworks/flux"
	fluxk8s "github.com/weaveworks/flux/cluster/kubernetes"
	flux_v1beta1 "github.com/weaveworks/flux/integrations/apis/flux.weave.works/v1beta1"
)

type Action string

const (
	InstallAction Action = "CREATE"
	UpgradeAction Action = "UPDATE"
)

// Release contains clients needed to provide functionality related to helm releases
type Release struct {
	logger log.Logger

	HelmClient *k8shelm.Client
}

type Releaser interface {
	GetCurrent() (map[string][]DeployInfo, error)
	GetDeployedRelease(name string) (*hapi_release.Release, error)
	Install(dir string, releaseName string, fhr flux_v1beta1.HelmRelease, action Action, opts InstallOptions) (*hapi_release.Release, error)
}

type DeployInfo struct {
	Name string
}

type InstallOptions struct {
	DryRun    bool
	ReuseName bool
}

// New creates a new Release instance.
func New(logger log.Logger, helmClient *k8shelm.Client) *Release {
	r := &Release{
		logger:     logger,
		HelmClient: helmClient,
	}
	return r
}

// GetReleaseName either retrieves the release name from the Custom Resource or constructs a new one
//  in the form : $Namespace-$CustomResourceName
func GetReleaseName(fhr flux_v1beta1.HelmRelease) string {
	namespace := fhr.Namespace
	if namespace == "" {
		namespace = "default"
	}
	releaseName := fhr.Spec.ReleaseName
	if releaseName == "" {
		releaseName = fmt.Sprintf("%s-%s", namespace, fhr.Name)
	}

	return releaseName
}

// GetDeployedRelease returns a release with Deployed status
func (r *Release) GetDeployedRelease(name string) (*hapi_release.Release, error) {
	rls, err := r.HelmClient.ReleaseContent(name)
	if err != nil {
		return nil, err
	}
	if rls.Release.Info.Status.GetCode() == hapi_release.Status_DEPLOYED {
		return rls.GetRelease(), nil
	}
	return nil, nil
}

func (r *Release) canDelete(name string) (bool, error) {
	rls, err := r.HelmClient.ReleaseStatus(name)

	if err != nil {
		r.logger.Log("error", fmt.Sprintf("Error finding status for release (%s): %#v", name, err))
		return false, err
	}
	/*
		"UNKNOWN":          0,
		"DEPLOYED":         1,
		"DELETED":          2,
		"SUPERSEDED":       3,
		"FAILED":           4,
		"DELETING":         5,
		"PENDING_INSTALL":  6,
		"PENDING_UPGRADE":  7,
		"PENDING_ROLLBACK": 8,
	*/
	status := rls.GetInfo().GetStatus()
	r.logger.Log("info", fmt.Sprintf("Release [%s] status: %s", name, status.Code.String()))
	switch status.Code {
	case 1, 4:
		r.logger.Log("info", fmt.Sprintf("Deleting release (%s)", name))
		return true, nil
	case 2:
		r.logger.Log("info", fmt.Sprintf("Release (%s) already deleted", name))
		return false, nil
	default:
		r.logger.Log("info", fmt.Sprintf("Release (%s) with status %s cannot be deleted", name, status.Code.String()))
		return false, fmt.Errorf("Release (%s) with status %s cannot be deleted", name, status.Code.String())
	}
}

// Install performs a Chart release given the directory containing the
// charts, and the HelmRelease specifying the release. Depending
// on the release type, this is either a new release, or an upgrade of
// an existing one.
//
// TODO(michael): cloneDir is only relevant if installing from git;
// either split this procedure into two varieties, or make it more
// general and calculate the path to the chart in the caller.
func (r *Release) Install(chartPath, releaseName string, fhr flux_v1beta1.HelmRelease, action Action, opts InstallOptions, kubeClient *kubernetes.Clientset) (*hapi_release.Release, error) {
	if chartPath == "" {
		return nil, fmt.Errorf("empty path to chart supplied for resource %q", fhr.ResourceID().String())
	}
	_, err := os.Stat(chartPath)
	switch {
	case os.IsNotExist(err):
		return nil, fmt.Errorf("no file or dir at path to chart: %s", chartPath)
	case err != nil:
		return nil, fmt.Errorf("error statting path given for chart %s: %s", chartPath, err.Error())
	}

	r.logger.Log("info", "releaseName", releaseName, "action", action, "options", fmt.Sprintf("%+v", opts))

	// Read values from given valueFile paths (configmaps, etc.)
	mergedValues := chartutil.Values{}
	for _, valueFileSecret := range fhr.Spec.ValueFileSecrets {
		// Read the contents of the secret
		secret, err := kubeClient.CoreV1().Secrets(fhr.Namespace).Get(valueFileSecret.Name, v1.GetOptions{})
		if err != nil {
			r.logger.Log("error", fmt.Sprintf("Cannot get secret %s for Chart release [%s]: %#v", valueFileSecret.Name, releaseName, err))
			return nil, err
		}

		// Load values.yaml file and merge
		var values chartutil.Values
		err = yaml.Unmarshal(secret.Data["values.yaml"], &values)
		if err != nil {
			r.logger.Log("error", fmt.Sprintf("Cannot yaml.Unmashal values.yaml in secret %s for Chart release [%s]: %#v", valueFileSecret.Name, releaseName, err))
			return nil, err
		}
		mergedValues = mergeValues(mergedValues, values)
	}
	// Merge in values after valueFiles
	mergedValues = mergeValues(mergedValues, fhr.Spec.Values)

	strVals, err := mergedValues.YAML()
	if err != nil {
		r.logger.Log("error", fmt.Sprintf("Problem with supplied customizations for Chart release [%s]: %#v", releaseName, err))
		return nil, err
	}
	rawVals := []byte(strVals)

	switch action {
	case InstallAction:
		res, err := r.HelmClient.InstallRelease(
			chartPath,
			fhr.GetNamespace(),
			k8shelm.ValueOverrides(rawVals),
			k8shelm.ReleaseName(releaseName),
			k8shelm.InstallDryRun(opts.DryRun),
			k8shelm.InstallReuseName(opts.ReuseName),
			/*
				helm.InstallReuseName(i.replace),
				helm.InstallDisableHooks(i.disableHooks),
				helm.InstallTimeout(i.timeout),
				helm.InstallWait(i.wait)
			*/
		)

		if err != nil {
			r.logger.Log("error", fmt.Sprintf("Chart release failed: %s: %#v", releaseName, err))
			// if an install fails, purge the release and keep retrying
			r.logger.Log("info", fmt.Sprintf("Deleting failed release: [%s]", releaseName))
			_, err = r.HelmClient.DeleteRelease(releaseName, k8shelm.DeletePurge(true))
			if err != nil {
				r.logger.Log("error", fmt.Sprintf("Release deletion error: %#v", err))
				return nil, err
			}
			return nil, err
		}
		if !opts.DryRun {
			err = r.annotateResources(res.Release, fhr)
		}
		return res.Release, err
	case UpgradeAction:
		res, err := r.HelmClient.UpdateRelease(
			releaseName,
			chartPath,
			k8shelm.UpdateValueOverrides(rawVals),
			k8shelm.UpgradeDryRun(opts.DryRun),
			/*
				helm.UpgradeRecreate(u.recreate),
				helm.UpgradeForce(u.force),
				helm.UpgradeDisableHooks(u.disableHooks),
				helm.UpgradeTimeout(u.timeout),
				helm.ResetValues(u.resetValues),
				helm.ReuseValues(u.reuseValues),
				helm.UpgradeWait(u.wait))
			*/
		)

		if err != nil {
			r.logger.Log("error", fmt.Sprintf("Chart upgrade release failed: %s: %#v", releaseName, err))
			return nil, err
		}
		if !opts.DryRun {
			err = r.annotateResources(res.Release, fhr)
		}
		return res.Release, err
	default:
		err = fmt.Errorf("Valid install options: CREATE, UPDATE. Provided: %s", action)
		r.logger.Log("error", err.Error())
		return nil, err
	}
}

// Delete purges a Chart release
func (r *Release) Delete(name string) error {
	ok, err := r.canDelete(name)
	if !ok {
		if err != nil {
			return err
		}
		return nil
	}

	_, err = r.HelmClient.DeleteRelease(name, k8shelm.DeletePurge(true))
	if err != nil {
		r.logger.Log("error", fmt.Sprintf("Release deletion error: %#v", err))
		return err
	}
	r.logger.Log("info", fmt.Sprintf("Release deleted: [%s]", name))
	return nil
}

// GetCurrent provides Chart releases (stored in tiller ConfigMaps)
//		output:
//						map[namespace][release name] = nil
func (r *Release) GetCurrent() (map[string][]DeployInfo, error) {
	response, err := r.HelmClient.ListReleases()
	if err != nil {
		return nil, r.logger.Log("error", err)
	}
	r.logger.Log("info", fmt.Sprintf("Number of Chart releases: %d\n", response.GetCount()))

	relsM := make(map[string][]DeployInfo)
	var depl []DeployInfo

	for _, r := range response.GetReleases() {
		ns := r.Namespace
		depl = relsM[ns]

		depl = append(depl, DeployInfo{Name: r.Name})
		relsM[ns] = depl
	}
	return relsM, nil
}

// annotateResources annotates each of the resources created (or updated)
// by the release so that we can spot them.
func (r *Release) annotateResources(release *hapi_release.Release, fhr flux_v1beta1.HelmRelease) error {
	args := []string{"annotate", "--overwrite"}
	args = append(args, "--namespace", release.Namespace)
	args = append(args, "-f", "-")
	args = append(args, fluxk8s.AntecedentAnnotation+"="+fhrResourceID(fhr).String())

	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	cmd := exec.CommandContext(ctx, "kubectl", args...)
	cmd.Stdin = bytes.NewBufferString(release.Manifest)

	output, err := cmd.CombinedOutput()
	if err != nil {
		r.logger.Log("output", string(output), "err", err)
	}
	return err
}

// fhrResourceID constructs a flux.ResourceID for a HelmRelease
// resource.
func fhrResourceID(fhr flux_v1beta1.HelmRelease) flux.ResourceID {
	return flux.MakeResourceID(fhr.Namespace, "HelmRelease", fhr.Name)
}

// Merges source and destination `chartutils.Values`, preferring values from the source Values
// This is slightly adapted from https://github.com/helm/helm/blob/master/cmd/helm/install.go#L329
func mergeValues(dest, src chartutil.Values) chartutil.Values {
	for k, v := range src {
		// If the key doesn't exist already, then just set the key to that value
		if _, exists := dest[k]; !exists {
			dest[k] = v
			continue
		}
		nextMap, ok := v.(map[string]interface{})
		// If it isn't another map, overwrite the value
		if !ok {
			dest[k] = v
			continue
		}
		// Edge case: If the key exists in the destination, but isn't a map
		destMap, isMap := dest[k].(map[string]interface{})
		// If the source map has a map for this key, prefer it
		if !isMap {
			dest[k] = v
			continue
		}
		// If we got to this point, it is a map in both, so merge them
		dest[k] = mergeValues(destMap, nextMap)
	}
	return dest
}
<commit_msg>Prevent Helm release deletion of installed charts - purge a Helm release only if there is a single revision and that one failed - prevent Helm release deletion if Kubernetes API connectivity is flaky - fix #1524 #1447<commit_after>package release

import (
	"bytes"
	"context"
	"fmt"
	"os"
	"os/exec"
	"time"

	"github.com/ghodss/yaml"
	"github.com/go-kit/kit/log"
	"k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/helm/pkg/chartutil"
	k8shelm "k8s.io/helm/pkg/helm"
	hapi_release "k8s.io/helm/pkg/proto/hapi/release"

	"github.com/weaveworks/flux"
	fluxk8s "github.com/weaveworks/flux/cluster/kubernetes"
	flux_v1beta1 "github.com/weaveworks/flux/integrations/apis/flux.weave.works/v1beta1"
)

type Action string

const (
	InstallAction Action = "CREATE"
	UpgradeAction Action = "UPDATE"
)

// Release contains clients needed to provide functionality related to helm releases
type Release struct {
	logger log.Logger

	HelmClient *k8shelm.Client
}

type Releaser interface {
	GetCurrent() (map[string][]DeployInfo, error)
	GetDeployedRelease(name string) (*hapi_release.Release, error)
	Install(dir string, releaseName string, fhr flux_v1beta1.HelmRelease, action Action, opts InstallOptions) (*hapi_release.Release, error)
}

type DeployInfo struct {
	Name string
}

type InstallOptions struct {
	DryRun    bool
	ReuseName bool
}

// New creates a new Release instance.
func New(logger log.Logger, helmClient *k8shelm.Client) *Release {
	r := &Release{
		logger:     logger,
		HelmClient: helmClient,
	}
	return r
}

// GetReleaseName either retrieves the release name from the Custom Resource or constructs a new one
//  in the form : $Namespace-$CustomResourceName
func GetReleaseName(fhr flux_v1beta1.HelmRelease) string {
	namespace := fhr.Namespace
	if namespace == "" {
		namespace = "default"
	}
	releaseName := fhr.Spec.ReleaseName
	if releaseName == "" {
		releaseName = fmt.Sprintf("%s-%s", namespace, fhr.Name)
	}

	return releaseName
}

// GetDeployedRelease returns a release with Deployed status
func (r *Release) GetDeployedRelease(name string) (*hapi_release.Release, error) {
	rls, err := r.HelmClient.ReleaseContent(name)
	if err != nil {
		return nil, err
	}
	if rls.Release.Info.Status.GetCode() == hapi_release.Status_DEPLOYED {
		return rls.GetRelease(), nil
	}
	return nil, nil
}

func (r *Release) canDelete(name string) (bool, error) {
	rls, err := r.HelmClient.ReleaseStatus(name)

	if err != nil {
		r.logger.Log("error", fmt.Sprintf("Error finding status for release (%s): %#v", name, err))
		return false, err
	}
	/*
		"UNKNOWN":          0,
		"DEPLOYED":         1,
		"DELETED":          2,
		"SUPERSEDED":       3,
		"FAILED":           4,
		"DELETING":         5,
		"PENDING_INSTALL":  6,
		"PENDING_UPGRADE":  7,
		"PENDING_ROLLBACK": 8,
	*/
	status := rls.GetInfo().GetStatus()
	r.logger.Log("info", fmt.Sprintf("Release [%s] status: %s", name, status.Code.String()))
	switch status.Code {
	case 1, 4:
		r.logger.Log("info", fmt.Sprintf("Deleting release (%s)", name))
		return true, nil
	case 2:
		r.logger.Log("info", fmt.Sprintf("Release (%s) already deleted", name))
		return false, nil
	default:
		r.logger.Log("info", fmt.Sprintf("Release (%s) with status %s cannot be deleted", name, status.Code.String()))
		return false, fmt.Errorf("Release (%s) with status %s cannot be deleted", name, status.Code.String())
	}
}

// Install performs a Chart release given the directory containing the
// charts, and the HelmRelease specifying the release. Depending
// on the release type, this is either a new release, or an upgrade of
// an existing one.
//
// TODO(michael): cloneDir is only relevant if installing from git;
// either split this procedure into two varieties, or make it more
// general and calculate the path to the chart in the caller.
func (r *Release) Install(chartPath, releaseName string, fhr flux_v1beta1.HelmRelease, action Action, opts InstallOptions, kubeClient *kubernetes.Clientset) (*hapi_release.Release, error) {
	if chartPath == "" {
		return nil, fmt.Errorf("empty path to chart supplied for resource %q", fhr.ResourceID().String())
	}
	_, err := os.Stat(chartPath)
	switch {
	case os.IsNotExist(err):
		return nil, fmt.Errorf("no file or dir at path to chart: %s", chartPath)
	case err != nil:
		return nil, fmt.Errorf("error statting path given for chart %s: %s", chartPath, err.Error())
	}

	r.logger.Log("info", "releaseName", releaseName, "action", action, "options", fmt.Sprintf("%+v", opts))

	// Read values from given valueFile paths (configmaps, etc.)
	mergedValues := chartutil.Values{}
	for _, valueFileSecret := range fhr.Spec.ValueFileSecrets {
		// Read the contents of the secret
		secret, err := kubeClient.CoreV1().Secrets(fhr.Namespace).Get(valueFileSecret.Name, v1.GetOptions{})
		if err != nil {
			r.logger.Log("error", fmt.Sprintf("Cannot get secret %s for Chart release [%s]: %#v", valueFileSecret.Name, releaseName, err))
			return nil, err
		}

		// Load values.yaml file and merge
		var values chartutil.Values
		err = yaml.Unmarshal(secret.Data["values.yaml"], &values)
		if err != nil {
			r.logger.Log("error", fmt.Sprintf("Cannot yaml.Unmashal values.yaml in secret %s for Chart release [%s]: %#v", valueFileSecret.Name, releaseName, err))
			return nil, err
		}
		mergedValues = mergeValues(mergedValues, values)
	}
	// Merge in values after valueFiles
	mergedValues = mergeValues(mergedValues, fhr.Spec.Values)

	strVals, err := mergedValues.YAML()
	if err != nil {
		r.logger.Log("error", fmt.Sprintf("Problem with supplied customizations for Chart release [%s]: %#v", releaseName, err))
		return nil, err
	}
	rawVals := []byte(strVals)

	switch action {
	case InstallAction:
		res, err := r.HelmClient.InstallRelease(
			chartPath,
			fhr.GetNamespace(),
			k8shelm.ValueOverrides(rawVals),
			k8shelm.ReleaseName(releaseName),
			k8shelm.InstallDryRun(opts.DryRun),
			k8shelm.InstallReuseName(opts.ReuseName),
			/*
				helm.InstallReuseName(i.replace),
				helm.InstallDisableHooks(i.disableHooks),
				helm.InstallTimeout(i.timeout),
				helm.InstallWait(i.wait)
			*/
		)

		if err != nil {
			r.logger.Log("error", fmt.Sprintf("Chart release failed: %s: %#v", releaseName, err))
			// purge the release if the install failed but only if this is the first revision
			history, err := r.HelmClient.ReleaseHistory(releaseName, k8shelm.WithMaxHistory(2))
			if err == nil && len(history.Releases) == 1 && history.Releases[0].Info.Status.Code == hapi_release.Status_FAILED {
				r.logger.Log("info", fmt.Sprintf("Deleting failed release: [%s]", releaseName))
				_, err = r.HelmClient.DeleteRelease(releaseName, k8shelm.DeletePurge(true))
				if err != nil {
					r.logger.Log("error", fmt.Sprintf("Release deletion error: %#v", err))
					return nil, err
				}
			}
			return nil, err
		}
		if !opts.DryRun {
			err = r.annotateResources(res.Release, fhr)
		}
		return res.Release, err
	case UpgradeAction:
		res, err := r.HelmClient.UpdateRelease(
			releaseName,
			chartPath,
			k8shelm.UpdateValueOverrides(rawVals),
			k8shelm.UpgradeDryRun(opts.DryRun),
			/*
				helm.UpgradeRecreate(u.recreate),
				helm.UpgradeForce(u.force),
				helm.UpgradeDisableHooks(u.disableHooks),
				helm.UpgradeTimeout(u.timeout),
				helm.ResetValues(u.resetValues),
				helm.ReuseValues(u.reuseValues),
				helm.UpgradeWait(u.wait))
			*/
		)

		if err != nil {
			r.logger.Log("error", fmt.Sprintf("Chart upgrade release failed: %s: %#v", releaseName, err))
			return nil, err
		}
		if !opts.DryRun {
			err = r.annotateResources(res.Release, fhr)
		}
		return res.Release, err
	default:
		err = fmt.Errorf("Valid install options: CREATE, UPDATE. Provided: %s", action)
		r.logger.Log("error", err.Error())
		return nil, err
	}
}

// Delete purges a Chart release
func (r *Release) Delete(name string) error {
	ok, err := r.canDelete(name)
	if !ok {
		if err != nil {
			return err
		}
		return nil
	}

	_, err = r.HelmClient.DeleteRelease(name, k8shelm.DeletePurge(true))
	if err != nil {
		r.logger.Log("error", fmt.Sprintf("Release deletion error: %#v", err))
		return err
	}
	r.logger.Log("info", fmt.Sprintf("Release deleted: [%s]", name))
	return nil
}

// GetCurrent provides Chart releases (stored in tiller ConfigMaps)
//		output:
//						map[namespace][release name] = nil
func (r *Release) GetCurrent() (map[string][]DeployInfo, error) {
	response, err := r.HelmClient.ListReleases()
	if err != nil {
		return nil, r.logger.Log("error", err)
	}
	r.logger.Log("info", fmt.Sprintf("Number of Chart releases: %d\n", response.GetCount()))

	relsM := make(map[string][]DeployInfo)
	var depl []DeployInfo

	for _, r := range response.GetReleases() {
		ns := r.Namespace
		depl = relsM[ns]

		depl = append(depl, DeployInfo{Name: r.Name})
		relsM[ns] = depl
	}
	return relsM, nil
}

// annotateResources annotates each of the resources created (or updated)
// by the release so that we can spot them.
func (r *Release) annotateResources(release *hapi_release.Release, fhr flux_v1beta1.HelmRelease) error {
	args := []string{"annotate", "--overwrite"}
	args = append(args, "--namespace", release.Namespace)
	args = append(args, "-f", "-")
	args = append(args, fluxk8s.AntecedentAnnotation+"="+fhrResourceID(fhr).String())

	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	cmd := exec.CommandContext(ctx, "kubectl", args...)
	cmd.Stdin = bytes.NewBufferString(release.Manifest)

	output, err := cmd.CombinedOutput()
	if err != nil {
		r.logger.Log("output", string(output), "err", err)
	}
	return err
}

// fhrResourceID constructs a flux.ResourceID for a HelmRelease
// resource.
func fhrResourceID(fhr flux_v1beta1.HelmRelease) flux.ResourceID {
	return flux.MakeResourceID(fhr.Namespace, "HelmRelease", fhr.Name)
}

// Merges source and destination `chartutils.Values`, preferring values from the source Values
// This is slightly adapted from https://github.com/helm/helm/blob/master/cmd/helm/install.go#L329
func mergeValues(dest, src chartutil.Values) chartutil.Values {
	for k, v := range src {
		// If the key doesn't exist already, then just set the key to that value
		if _, exists := dest[k]; !exists {
			dest[k] = v
			continue
		}
		nextMap, ok := v.(map[string]interface{})
		// If it isn't another map, overwrite the value
		if !ok {
			dest[k] = v
			continue
		}
		// Edge case: If the key exists in the destination, but isn't a map
		destMap, isMap := dest[k].(map[string]interface{})
		// If the source map has a map for this key, prefer it
		if !isMap {
			dest[k] = v
			continue
		}
		// If we got to this point, it is a map in both, so merge them
		dest[k] = mergeValues(destMap, nextMap)
	}
	return dest
}
<|endoftext|>


<commit_before>package eval

import (
	"github.com/johnny-morrice/godless/api"
	"github.com/johnny-morrice/godless/crdt"
	"github.com/johnny-morrice/godless/internal/crypto"
	"github.com/johnny-morrice/godless/log"
	"github.com/johnny-morrice/godless/query"
	"github.com/pkg/errors"
)

type NamespaceTreeJoin struct {
	query.NoSelectVisitor
	query.NoDebugVisitor
	query.ErrorCollectVisitor
	Namespace   api.NamespaceTree
	tableKey    crdt.TableName
	table       crdt.Table
	privateKeys []crypto.PrivateKey
	keyStore    api.KeyStore
}

func MakeNamespaceTreeJoin(ns api.NamespaceTree, keyStore api.KeyStore) *NamespaceTreeJoin {
	return &NamespaceTreeJoin{
		Namespace: ns,
		keyStore:  keyStore,
	}
}

func (visitor *NamespaceTreeJoin) RunQuery() api.APIResponse {
	fail := api.RESPONSE_FAIL
	fail.Type = api.API_QUERY

	err := visitor.Error()
	if err != nil {
		fail.Err = visitor.Error()
		return fail
	}

	if visitor.tableKey == "" {
		panic("Expected table key")
	}

	err = visitor.Namespace.JoinTable(visitor.tableKey, visitor.table)

	if err != nil {
		fail.Err = errors.Wrap(err, "NamespaceTreeJoin failed")
		return fail
	}

	return api.RESPONSE_QUERY
}

func (visitor *NamespaceTreeJoin) VisitPublicKeyHash(hash crypto.PublicKeyHash) {
	priv, matchErr := visitor.keyStore.GetPrivateKey(hash)

	if matchErr != nil {
		log.Warn("Private key lookup failed with: %v", matchErr)
		visitor.BadPublicKey(hash)
		return
	}

	log.Info("Joining with private key for %v", hash)
	visitor.privateKeys = append(visitor.privateKeys, priv)
}

func (visitor *NamespaceTreeJoin) VisitOpCode(opCode query.QueryOpCode) {
	if opCode != query.JOIN {
		visitor.CollectError(errors.New("Expected JOIN OpCode"))
	}
}

func (visitor *NamespaceTreeJoin) VisitTableKey(tableKey crdt.TableName) {
	if visitor.Error() != nil {
		return
	}

	visitor.tableKey = tableKey
}

func (visitor *NamespaceTreeJoin) VisitJoin(*query.QueryJoin) {
}

func (visitor *NamespaceTreeJoin) LeaveJoin(*query.QueryJoin) {
}

func (visitor *NamespaceTreeJoin) VisitRowJoin(position int, rowJoin *query.QueryRowJoin) {
	if visitor.Error() != nil {
		return
	}

	row := crdt.Row{}

	for k, entryValue := range rowJoin.Entries {
		point, err := visitor.makePoint(entryValue)

		if err != nil {
			visitor.badPrivateKey()
		}

		entry := crdt.MakeEntry([]crdt.Point{point})
		row = row.JoinEntry(k, entry)
	}

	joined := visitor.table.JoinRow(rowJoin.RowKey, row)

	visitor.table = joined
}

func (visitor *NamespaceTreeJoin) makePoint(text crdt.PointText) (crdt.Point, error) {
	return crdt.SignedPoint(text, visitor.privateKeys)
}

func (visitor *NamespaceTreeJoin) badPrivateKey() {
	err := errors.New("Failed to sign Point with bad private key")
	visitor.CollectError(err)
}
<commit_msg>Fix dodgy byte logging<commit_after>package eval

import (
	"github.com/johnny-morrice/godless/api"
	"github.com/johnny-morrice/godless/crdt"
	"github.com/johnny-morrice/godless/internal/crypto"
	"github.com/johnny-morrice/godless/log"
	"github.com/johnny-morrice/godless/query"
	"github.com/pkg/errors"
)

type NamespaceTreeJoin struct {
	query.NoSelectVisitor
	query.NoDebugVisitor
	query.ErrorCollectVisitor
	Namespace   api.NamespaceTree
	tableKey    crdt.TableName
	table       crdt.Table
	privateKeys []crypto.PrivateKey
	keyStore    api.KeyStore
}

func MakeNamespaceTreeJoin(ns api.NamespaceTree, keyStore api.KeyStore) *NamespaceTreeJoin {
	return &NamespaceTreeJoin{
		Namespace: ns,
		keyStore:  keyStore,
	}
}

func (visitor *NamespaceTreeJoin) RunQuery() api.APIResponse {
	fail := api.RESPONSE_FAIL
	fail.Type = api.API_QUERY

	err := visitor.Error()
	if err != nil {
		fail.Err = visitor.Error()
		return fail
	}

	if visitor.tableKey == "" {
		panic("Expected table key")
	}

	err = visitor.Namespace.JoinTable(visitor.tableKey, visitor.table)

	if err != nil {
		fail.Err = errors.Wrap(err, "NamespaceTreeJoin failed")
		return fail
	}

	return api.RESPONSE_QUERY
}

func (visitor *NamespaceTreeJoin) VisitPublicKeyHash(hash crypto.PublicKeyHash) {
	priv, matchErr := visitor.keyStore.GetPrivateKey(hash)

	if matchErr != nil {
		log.Warn("Private key lookup failed with: %v", matchErr)
		visitor.BadPublicKey(hash)
		return
	}

	log.Info("Joining with private key for %v", string(hash))
	visitor.privateKeys = append(visitor.privateKeys, priv)
}

func (visitor *NamespaceTreeJoin) VisitOpCode(opCode query.QueryOpCode) {
	if opCode != query.JOIN {
		visitor.CollectError(errors.New("Expected JOIN OpCode"))
	}
}

func (visitor *NamespaceTreeJoin) VisitTableKey(tableKey crdt.TableName) {
	if visitor.Error() != nil {
		return
	}

	visitor.tableKey = tableKey
}

func (visitor *NamespaceTreeJoin) VisitJoin(*query.QueryJoin) {
}

func (visitor *NamespaceTreeJoin) LeaveJoin(*query.QueryJoin) {
}

func (visitor *NamespaceTreeJoin) VisitRowJoin(position int, rowJoin *query.QueryRowJoin) {
	if visitor.Error() != nil {
		return
	}

	row := crdt.Row{}

	for k, entryValue := range rowJoin.Entries {
		point, err := visitor.makePoint(entryValue)

		if err != nil {
			visitor.badPrivateKey()
		}

		entry := crdt.MakeEntry([]crdt.Point{point})
		row = row.JoinEntry(k, entry)
	}

	joined := visitor.table.JoinRow(rowJoin.RowKey, row)

	visitor.table = joined
}

func (visitor *NamespaceTreeJoin) makePoint(text crdt.PointText) (crdt.Point, error) {
	return crdt.SignedPoint(text, visitor.privateKeys)
}

func (visitor *NamespaceTreeJoin) badPrivateKey() {
	err := errors.New("Failed to sign Point with bad private key")
	visitor.CollectError(err)
}
<|endoftext|>


<commit_before>package gcsx

import (
	"context"
	"testing"
	"time"

	"github.com/fsouza/fake-gcs-server/fakestorage"
	"github.com/googlecloudplatform/gcsfuse/internal/storage"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcsfake"
	. "github.com/jacobsa/ogletest"
	"github.com/jacobsa/timeutil"
)

func TestFlags(t *testing.T) { RunTests(t) }

const TestBucketName string = "gcsfuse-default-bucket"

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type BucketManagerTest struct {
	fakeStorageServer *fakestorage.Server
	bucketHandle      *storage.BucketHandle
}

var _ SetUpInterface = &BucketManagerTest{}
var _ TearDownInterface = &BucketManagerTest{}

func init() { RegisterTestSuite(&BucketManagerTest{}) }

func (t *BucketManagerTest) SetUp(_ *TestInfo) {
	var err error

	t.fakeStorageServer, err = storage.CreateFakeStorageServer([]fakestorage.Object{storage.GetTestFakeStorageObject()})

	AssertEq(nil, err)

	storageClient := &storage.Storageclient{Client: t.fakeStorageServer.Client()}
	t.bucketHandle, err = storageClient.BucketHandle(TestBucketName)

	AssertEq(nil, err)
	AssertNe(nil, t.bucketHandle)
}

func (t *BucketManagerTest) TearDown() {
	t.fakeStorageServer.Stop()
}

func (t *BucketManagerTest) TestNewBucketManagerMethod() {
	var nilValue *bucketManager = nil
	storageClient := &storage.Storageclient{Client: t.fakeStorageServer.Client()}

	bucketConfig := BucketConfig{
		BillingProject:                     "BillingProject",
		OnlyDir:                            "OnlyDir",
		EgressBandwidthLimitBytesPerSecond: 7,
		OpRateLimitHz:                      11,
		StatCacheCapacity:                  100,
		StatCacheTTL:                       20 * time.Second,
		EnableMonitoring:                   true,
		DebugGcs:                           true,
		AppendThreshold:                    2,
		TmpObjectPrefix:                    "TmpObjectPrefix",
	}

	bm := NewBucketManager(bucketConfig, nil, storageClient)
	ExpectNe(bm, nilValue)
}

func (t *BucketManagerTest) TestSetupGcsBucketWhenEnableStorageClientLibraryIsTrue() {
	var bm bucketManager
	var nilBucket *gcs.Bucket = nil
	bm.storageHandle = &storage.Storageclient{Client: t.fakeStorageServer.Client()}
	bm.config.EnableStorageClientLibrary = true
	bm.config.DebugGcs = true

	bucket, err := bm.SetUpGcsBucket(context.Background(), TestBucketName)

	ExpectNe(&bucket, nilBucket)
	ExpectEq(err, nil)
}

func (t *BucketManagerTest) TestSetupGcsBucketWhenEnableStorageClientLibraryIsFalse() {
	var bm bucketManager
	var nilBucket *gcs.Bucket = nil

	bm.storageHandle = &storage.Storageclient{Client: t.fakeStorageServer.Client()}
	bm.config.EnableStorageClientLibrary = false
	bm.config.BillingProject = "BillingProject"
	bm.conn = &Connection{
		wrapped: gcsfake.NewConn(timeutil.RealClock()),
	}

	bucket, err := bm.SetUpGcsBucket(context.Background(), "fake@bucket")

	ExpectNe(&bucket, nilBucket)
	ExpectEq(err, nil)
}

func (t *BucketManagerTest) TestSetUpBucketMethod() {
	var bm bucketManager
	nilSync := NewSyncerBucket(0, "", nil)

	bucketConfig := BucketConfig{
		BillingProject:                     "BillingProject",
		OnlyDir:                            "OnlyDir",
		EgressBandwidthLimitBytesPerSecond: 7,
		OpRateLimitHz:                      11,
		StatCacheCapacity:                  100,
		StatCacheTTL:                       20 * time.Second,
		EnableMonitoring:                   true,
		DebugGcs:                           true,
		AppendThreshold:                    2,
		TmpObjectPrefix:                    "TmpObjectPrefix",
		EnableStorageClientLibrary:         true,
	}

	ctx := context.Background()

	bm.storageHandle = &storage.Storageclient{Client: t.fakeStorageServer.Client()}
	bm.config = bucketConfig
	bm.gcCtx = ctx
	bm.conn = &Connection{
		wrapped: gcsfake.NewConn(timeutil.RealClock()),
	}

	bucket, err := bm.SetUpBucket(context.Background(), TestBucketName)
	ExpectNe(bucket.Syncer, nilSync.Syncer)
	ExpectEq(err, nil)
}
<commit_msg>Linked changes in main.go file<commit_after>package gcsx

import (
	"context"
	"testing"
	"time"

	"github.com/fsouza/fake-gcs-server/fakestorage"
	"github.com/googlecloudplatform/gcsfuse/internal/storage"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcsfake"
	. "github.com/jacobsa/ogletest"
	"github.com/jacobsa/timeutil"
)

func TestFlags(t *testing.T) { RunTests(t) }

const TestBucketName string = "gcsfuse-default-bucket"

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type BucketManagerTest struct {
	fakeStorageServer *fakestorage.Server
	bucketHandle      *storage.BucketHandle
}

var _ SetUpInterface = &BucketManagerTest{}
var _ TearDownInterface = &BucketManagerTest{}

func init() { RegisterTestSuite(&BucketManagerTest{}) }

func (t *BucketManagerTest) SetUp(_ *TestInfo) {
	var err error

	t.fakeStorageServer, err = storage.CreateFakeStorageServer([]fakestorage.Object{storage.GetTestFakeStorageObject()})

	AssertEq(nil, err)

	storageClient := &storage.Storageclient{Client: t.fakeStorageServer.Client()}
	t.bucketHandle, err = storageClient.BucketHandle(TestBucketName)

	AssertEq(nil, err)
	AssertNe(nil, t.bucketHandle)
}

func (t *BucketManagerTest) TearDown() {
	t.fakeStorageServer.Stop()
}

func (t *BucketManagerTest) TestNewBucketManagerMethod() {
	var nilValue *bucketManager = nil
	storageClient := &storage.Storageclient{Client: t.fakeStorageServer.Client()}

	bucketConfig := BucketConfig{
		BillingProject:                     "BillingProject",
		OnlyDir:                            "OnlyDir",
		EgressBandwidthLimitBytesPerSecond: 7,
		OpRateLimitHz:                      11,
		StatCacheCapacity:                  100,
		StatCacheTTL:                       20 * time.Second,
		EnableMonitoring:                   true,
		DebugGcs:                           true,
		AppendThreshold:                    2,
		TmpObjectPrefix:                    "TmpObjectPrefix",
	}

	bm := NewBucketManager(bucketConfig, nil, storageClient)

	ExpectNe(bm, nilValue)
}

func (t *BucketManagerTest) TestSetupGcsBucketWhenEnableStorageClientLibraryIsTrue() {
	var bm bucketManager
	var nilBucket *gcs.Bucket = nil
	bm.storageHandle = &storage.Storageclient{Client: t.fakeStorageServer.Client()}
	bm.config.EnableStorageClientLibrary = true
	bm.config.DebugGcs = true

	bucket, err := bm.SetUpGcsBucket(context.Background(), TestBucketName)

	ExpectNe(&bucket, nilBucket)
	ExpectEq(err, nil)
}

func (t *BucketManagerTest) TestSetupGcsBucketWhenEnableStorageClientLibraryIsFalse() {
	var bm bucketManager
	var nilBucket *gcs.Bucket = nil

	bm.storageHandle = &storage.Storageclient{Client: t.fakeStorageServer.Client()}
	bm.config.EnableStorageClientLibrary = false
	bm.config.BillingProject = "BillingProject"
	bm.conn = &Connection{
		wrapped: gcsfake.NewConn(timeutil.RealClock()),
	}

	bucket, err := bm.SetUpGcsBucket(context.Background(), "fake@bucket")

	ExpectNe(&bucket, nilBucket)
	ExpectEq(err, nil)
}

func (t *BucketManagerTest) TestSetUpBucketMethod() {
	var bm bucketManager
	nilSync := NewSyncerBucket(0, "", nil)

	bucketConfig := BucketConfig{
		BillingProject:                     "BillingProject",
		OnlyDir:                            "OnlyDir",
		EgressBandwidthLimitBytesPerSecond: 7,
		OpRateLimitHz:                      11,
		StatCacheCapacity:                  100,
		StatCacheTTL:                       20 * time.Second,
		EnableMonitoring:                   true,
		DebugGcs:                           true,
		AppendThreshold:                    2,
		TmpObjectPrefix:                    "TmpObjectPrefix",
		EnableStorageClientLibrary:         true,
	}

	ctx := context.Background()

	bm.storageHandle = &storage.Storageclient{Client: t.fakeStorageServer.Client()}
	bm.config = bucketConfig
	bm.gcCtx = ctx
	bm.conn = &Connection{
		wrapped: gcsfake.NewConn(timeutil.RealClock()),
	}

	bucket, err := bm.SetUpBucket(context.Background(), TestBucketName)

	ExpectNe(bucket.Syncer, nilSync.Syncer)
	ExpectEq(err, nil)
}
<|endoftext|>


<commit_before>package iobufpool_test

import (
	"testing"

	"github.com/jackc/pgx/v5/internal/iobufpool"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestGetCap(t *testing.T) {
	tests := []struct {
		requestedLen int
		expectedCap  int
	}{
		{requestedLen: 0, expectedCap: 256},
		{requestedLen: 128, expectedCap: 256},
		{requestedLen: 255, expectedCap: 256},
		{requestedLen: 256, expectedCap: 256},
		{requestedLen: 257, expectedCap: 512},
		{requestedLen: 511, expectedCap: 512},
		{requestedLen: 512, expectedCap: 512},
		{requestedLen: 513, expectedCap: 1024},
		{requestedLen: 1023, expectedCap: 1024},
		{requestedLen: 1024, expectedCap: 1024},
		{requestedLen: 33554431, expectedCap: 33554432},
		{requestedLen: 33554432, expectedCap: 33554432},

		// Above 32 MiB skip the pool and allocate exactly the requested size.
		{requestedLen: 33554433, expectedCap: 33554433},
	}
	for _, tt := range tests {
		buf := iobufpool.Get(tt.requestedLen)
		assert.Equalf(t, tt.requestedLen, len(buf), "bad len for requestedLen: %d", len(buf), tt.requestedLen)
		assert.Equalf(t, tt.expectedCap, cap(buf), "bad cap for requestedLen: %d", tt.requestedLen)
	}
}

func TestPutHandlesWrongSizedBuffers(t *testing.T) {
	for putBufSize := range []int{0, 1, 128, 250, 256, 257, 1023, 1024, 1025, 1 << 28} {
		putBuf := make([]byte, putBufSize)
		iobufpool.Put(putBuf)

		tests := []struct {
			requestedLen int
			expectedCap  int
		}{
			{requestedLen: 0, expectedCap: 256},
			{requestedLen: 128, expectedCap: 256},
			{requestedLen: 255, expectedCap: 256},
			{requestedLen: 256, expectedCap: 256},
			{requestedLen: 257, expectedCap: 512},
			{requestedLen: 511, expectedCap: 512},
			{requestedLen: 512, expectedCap: 512},
			{requestedLen: 513, expectedCap: 1024},
			{requestedLen: 1023, expectedCap: 1024},
			{requestedLen: 1024, expectedCap: 1024},
			{requestedLen: 33554431, expectedCap: 33554432},
			{requestedLen: 33554432, expectedCap: 33554432},

			// Above 32 MiB skip the pool and allocate exactly the requested size.
			{requestedLen: 33554433, expectedCap: 33554433},
		}
		for _, tt := range tests {
			getBuf := iobufpool.Get(tt.requestedLen)
			assert.Equalf(t, tt.requestedLen, len(getBuf), "len(putBuf): %d, requestedLen: %d", len(putBuf), tt.requestedLen)
			assert.Equalf(t, tt.expectedCap, cap(getBuf), "cap(putBuf): %d, requestedLen: %d", cap(putBuf), tt.requestedLen)
		}
	}
}

func TestPutGetBufferReuse(t *testing.T) {
	buf := iobufpool.Get(4)
	buf[0] = 1
	iobufpool.Put(buf)
	buf = iobufpool.Get(4)
	require.Equal(t, byte(1), buf[0])
}
<commit_msg>Fix flickering test in CI<commit_after>package iobufpool_test

import (
	"testing"

	"github.com/jackc/pgx/v5/internal/iobufpool"
	"github.com/stretchr/testify/assert"
)

func TestGetCap(t *testing.T) {
	tests := []struct {
		requestedLen int
		expectedCap  int
	}{
		{requestedLen: 0, expectedCap: 256},
		{requestedLen: 128, expectedCap: 256},
		{requestedLen: 255, expectedCap: 256},
		{requestedLen: 256, expectedCap: 256},
		{requestedLen: 257, expectedCap: 512},
		{requestedLen: 511, expectedCap: 512},
		{requestedLen: 512, expectedCap: 512},
		{requestedLen: 513, expectedCap: 1024},
		{requestedLen: 1023, expectedCap: 1024},
		{requestedLen: 1024, expectedCap: 1024},
		{requestedLen: 33554431, expectedCap: 33554432},
		{requestedLen: 33554432, expectedCap: 33554432},

		// Above 32 MiB skip the pool and allocate exactly the requested size.
		{requestedLen: 33554433, expectedCap: 33554433},
	}
	for _, tt := range tests {
		buf := iobufpool.Get(tt.requestedLen)
		assert.Equalf(t, tt.requestedLen, len(buf), "bad len for requestedLen: %d", len(buf), tt.requestedLen)
		assert.Equalf(t, tt.expectedCap, cap(buf), "bad cap for requestedLen: %d", tt.requestedLen)
	}
}

func TestPutHandlesWrongSizedBuffers(t *testing.T) {
	for putBufSize := range []int{0, 1, 128, 250, 256, 257, 1023, 1024, 1025, 1 << 28} {
		putBuf := make([]byte, putBufSize)
		iobufpool.Put(putBuf)

		tests := []struct {
			requestedLen int
			expectedCap  int
		}{
			{requestedLen: 0, expectedCap: 256},
			{requestedLen: 128, expectedCap: 256},
			{requestedLen: 255, expectedCap: 256},
			{requestedLen: 256, expectedCap: 256},
			{requestedLen: 257, expectedCap: 512},
			{requestedLen: 511, expectedCap: 512},
			{requestedLen: 512, expectedCap: 512},
			{requestedLen: 513, expectedCap: 1024},
			{requestedLen: 1023, expectedCap: 1024},
			{requestedLen: 1024, expectedCap: 1024},
			{requestedLen: 33554431, expectedCap: 33554432},
			{requestedLen: 33554432, expectedCap: 33554432},

			// Above 32 MiB skip the pool and allocate exactly the requested size.
			{requestedLen: 33554433, expectedCap: 33554433},
		}
		for _, tt := range tests {
			getBuf := iobufpool.Get(tt.requestedLen)
			assert.Equalf(t, tt.requestedLen, len(getBuf), "len(putBuf): %d, requestedLen: %d", len(putBuf), tt.requestedLen)
			assert.Equalf(t, tt.expectedCap, cap(getBuf), "cap(putBuf): %d, requestedLen: %d", cap(putBuf), tt.requestedLen)
		}
	}
}

func TestPutGetBufferReuse(t *testing.T) {
	// There is no way to guarantee a buffer will be reused. It should be, but a GC between the Put and the Get will cause
	// it not to be. So try many times.
	for i := 0; i < 100000; i++ {
		buf := iobufpool.Get(4)
		buf[0] = 1
		iobufpool.Put(buf)
		buf = iobufpool.Get(4)
		if buf[0] == 1 {
			return
		}
	}

	t.Error("buffer was never reused")
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: MIT

package search

import (
	"testing"

	"github.com/issue9/assert"

	"github.com/caixw/apidoc/v7/core"
)

func TestTokenBuilder_append(t *testing.T) {
	a := assert.New(t)

	b := &tokenBuilder{}
	b.append(core.Range{
		Start: core.Position{Line: 1, Character: 11},
		End:   core.Position{Line: 1, Character: 12},
	}, 1)
	a.Equal(b.tokens[0], []int{1, 11, 1, 1, 0})

	// 空值，不会添加内容
	b.append(core.Range{}, 1)
	a.Equal(1, len(b.tokens))

	// 长度为 0
	a.Panic(func() {
		b.append(core.Range{
			Start: core.Position{Line: 1, Character: 11},
			End:   core.Position{Line: 1, Character: 11},
		}, 1)
	})

	// 长度为负数
	a.Panic(func() {
		b.append(core.Range{
			Start: core.Position{Line: 1, Character: 11},
			End:   core.Position{Line: 1, Character: 10},
		}, 1)
	})
}

func TestTokenBuilder_build(t *testing.T) {
	a := assert.New(t)

	b := &tokenBuilder{
		tokens: [][]int{
			{1, 2, 5, 0, 0},
			{1, 12, 5, 0, 0},
			{2, 2, 7, 0, 0},
			{2, 2, 5, 0, 0},
			{2, 5, 5, 0, 0},
			{11, 1, 5, 0, 0},
		},
	}
	a.Equal(b.build(), []int{
		1, 2, 5, 0, 0,
		0, 10, 5, 0, 0,
		1, 2, 7, 0, 0,
		0, 0, 5, 0, 0,
		0, 3, 5, 0, 0,
		9, 1, 5, 0, 0,
	})
}

func TestTokenBuilder_sort(t *testing.T) {
	a := assert.New(t)

	b := &tokenBuilder{
		tokens: [][]int{
			{11, 1, 5, 0, 0},
			{1, 2, 5, 0, 0},
			{1, 12, 5, 0, 0},
			{2, 5, 5, 0, 0},
			{2, 2, 7, 0, 0},
			{2, 2, 5, 0, 0},
		},
	}

	b.sort()
	a.Equal(b.tokens, [][]int{
		{1, 2, 5, 0, 0},
		{1, 12, 5, 0, 0},
		{2, 2, 7, 0, 0},
		{2, 2, 5, 0, 0},
		{2, 5, 5, 0, 0},
		{11, 1, 5, 0, 0},
	})
}
<commit_msg>test(internal/lsp/search): 添加针对 Tokens 的测试用例<commit_after>// SPDX-License-Identifier: MIT

package search

import (
	"testing"

	"github.com/issue9/assert"

	"github.com/caixw/apidoc/v7/core"
	"github.com/caixw/apidoc/v7/core/messagetest"
	"github.com/caixw/apidoc/v7/internal/ast"
)

func TestTokenBuilder_append(t *testing.T) {
	a := assert.New(t)

	b := &tokenBuilder{}
	b.append(core.Range{
		Start: core.Position{Line: 1, Character: 11},
		End:   core.Position{Line: 1, Character: 12},
	}, 1)
	a.Equal(b.tokens[0], []int{1, 11, 1, 1, 0})

	// 空值，不会添加内容
	b.append(core.Range{}, 1)
	a.Equal(1, len(b.tokens))

	// 长度为 0
	a.Panic(func() {
		b.append(core.Range{
			Start: core.Position{Line: 1, Character: 11},
			End:   core.Position{Line: 1, Character: 11},
		}, 1)
	})

	// 长度为负数
	a.Panic(func() {
		b.append(core.Range{
			Start: core.Position{Line: 1, Character: 11},
			End:   core.Position{Line: 1, Character: 10},
		}, 1)
	})
}

func TestTokenBuilder_build(t *testing.T) {
	a := assert.New(t)

	b := &tokenBuilder{
		tokens: [][]int{
			{1, 2, 5, 0, 0},
			{1, 12, 5, 0, 0},
			{2, 2, 7, 0, 0},
			{2, 2, 5, 0, 0},
			{2, 5, 5, 0, 0},
			{11, 1, 5, 0, 0},
		},
	}
	a.Equal(b.build(), []int{
		1, 2, 5, 0, 0,
		0, 10, 5, 0, 0,
		1, 2, 7, 0, 0,
		0, 0, 5, 0, 0,
		0, 3, 5, 0, 0,
		9, 1, 5, 0, 0,
	})
}

func TestTokenBuilder_sort(t *testing.T) {
	a := assert.New(t)

	b := &tokenBuilder{
		tokens: [][]int{
			{11, 1, 5, 0, 0},
			{1, 2, 5, 0, 0},
			{1, 12, 5, 0, 0},
			{2, 5, 5, 0, 0},
			{2, 2, 7, 0, 0},
			{2, 2, 5, 0, 0},
		},
	}

	b.sort()
	a.Equal(b.tokens, [][]int{
		{1, 2, 5, 0, 0},
		{1, 12, 5, 0, 0},
		{2, 2, 7, 0, 0},
		{2, 2, 5, 0, 0},
		{2, 5, 5, 0, 0},
		{11, 1, 5, 0, 0},
	})
}

func TestTokens(t *testing.T) {
	a := assert.New(t)

	b := `<apidoc version="1.1.1">
	<title>标题</title>
	<mimetype>xml</mimetype>
	<api method="GET">
		<path path="/users" />
		<response status="200" />
	</api>

	<api method="POST">
		<path path="/users" />
		<response status="200" />
	</api>
</apidoc>`
	blk := core.Block{Data: []byte(b), Location: core.Location{URI: "doc.go"}}
	doc := &ast.APIDoc{}
	rslt := messagetest.NewMessageHandler()
	doc.Parse(rslt.Handler, blk)
	rslt.Handler.Stop()
	a.Empty(rslt.Errors)

	a.Equal(Tokens(doc, "doc.go", 1, 2, 3), []int{
		0, 1, 6, 1, 0, // <apidoc>
		0, 7, 7, 2, 0,
		0, 9, 5, 3, 0,

		1, 2, 5, 1, 0, // <title>
		// {1, 8, 2, 0, 0}, // 元素内容不作解析，直接采用默认的注释颜色
		0, 10, 5, 1, 0, // </title>

		1, 2, 8, 1, 0, // <mimetype>
		// {2, 11, 3, 0, 0}, // 元素内容不作解析，直接采用默认的注释颜色
		0, 14, 8, 1, 0, // </mimetype>

		1, 2, 3, 1, 0, // <api>
		0, 4, 6, 2, 0,
		0, 8, 3, 3, 0,

		1, 3, 4, 1, 0, // path
		0, 5, 4, 2, 0,
		0, 6, 6, 3, 0,

		1, 3, 8, 1, 0, // response
		0, 9, 6, 2, 0,
		0, 8, 3, 3, 0,

		1, 3, 3, 1, 0, // </api>

		2, 2, 3, 1, 0, // <api>
		0, 4, 6, 2, 0,
		0, 8, 4, 3, 0,

		1, 3, 4, 1, 0, // path
		0, 5, 4, 2, 0,
		0, 6, 6, 3, 0,

		1, 3, 8, 1, 0, // response
		0, 9, 6, 2, 0,
		0, 8, 3, 3, 0,

		1, 3, 3, 1, 0, // </api>

		1, 2, 6, 1, 0, // </apidoc>
	})
}
<|endoftext|>


<commit_before>package oval

import "encoding/xml"

// Root : root object
type Root struct {
	XMLName     xml.Name    `xml:"oval_definitions"`
	Generator   Generator   `xml:"generator"`
	Definitions Definitions `xml:"definitions"`
	Tests       Tests       `xml:"tests"`
	Objects     Objects     `xml:"objects"`
	States      States      `xml:"states"`
}

// Generator : >generator
type Generator struct {
	XMLName        xml.Name `xml:"generator"`
	ProductName    string   `xml:"product_name"`
	ProductVersion string   `xml:"product_version"`
	SchemaVersion  string   `xml:"schema_version"`
	Timestamp      string   `xml:"timestamp"`
}

// Definitions : >definitions
type Definitions struct {
	XMLName     xml.Name     `xml:"definitions"`
	Definitions []Definition `xml:"definition"`
}

// Definition : >definitions>definition
type Definition struct {
	XMLName     xml.Name    `xml:"definition"`
	ID          string      `xml:"id,attr"`
	Class       string      `xml:"class,attr"`
	Title       string      `xml:"metadata>title"`
	Affecteds   []Affected  `xml:"metadata>affected"`
	References  []Reference `xml:"metadata>reference"`
	Description string      `xml:"metadata>description"`
	Advisory    Advisory    `xml:"metadata>advisory"` // RedHat, Oracle, Ubuntu
	Debian      Debian      `xml:"metadata>debian"`   // Debian
	Criteria    Criteria    `xml:"criteria"`
}

// Criteria : >definitions>definition>criteria
type Criteria struct {
	XMLName    xml.Name    `xml:"criteria"`
	Operator   string      `xml:"operator,attr"`
	Criterias  []Criteria  `xml:"criteria"`
	Criterions []Criterion `xml:"criterion"`
}

// Criterion : >definitions>definition>criteria>*>criterion
type Criterion struct {
	XMLName xml.Name `xml:"criterion"`
	Negate  bool     `xml:"negate,attr"`
	TestRef string   `xml:"test_ref,attr"`
	Comment string   `xml:"comment,attr"`
}

// Affected : >definitions>definition>metadata>affected
type Affected struct {
	XMLName   xml.Name `xml:"affected"`
	Family    string   `xml:"family,attr"`
	Platforms []string `xml:"platform"`
}

// Reference : >definitions>definition>metadata>reference
type Reference struct {
	XMLName xml.Name `xml:"reference"`
	Source  string   `xml:"source,attr"`
	RefID   string   `xml:"ref_id,attr"`
	RefURL  string   `xml:"ref_url,attr"`
}

// Advisory : >definitions>definition>metadata>advisory
// RedHat and Ubuntu OVAL
type Advisory struct {
	XMLName         xml.Name   `xml:"advisory"`
	Severity        string     `xml:"severity"`
	Cves            []Cve      `xml:"cve"`
	Bugzillas       []Bugzilla `xml:"bugzilla"`
	AffectedCPEList []string   `xml:"affected_cpe_list>cpe"`
	Refs            []Ref      `xml:"ref"` // Ubuntu Only
	Bugs            []Bug      `xml:"bug"` // Ubuntu Only
}

// Ref : >definitions>definition>metadata>advisory>ref
// Ubuntu OVAL
type Ref struct {
	XMLName xml.Name `xml:"ref"`
	URL     string   `xml:",chardata"`
}

// Bug : >definitions>definition>metadata>advisory>bug
// Ubuntu OVAL
type Bug struct {
	XMLName xml.Name `xml:"bug"`
	URL     string   `xml:",chardata"`
}

// Cve : >definitions>definition>metadata>advisory>cve
// RedHat OVAL
type Cve struct {
	XMLName xml.Name `xml:"cve"`
	CveID   string   `xml:",chardata"`
	Cvss2   string   `xml:"cvss2,attr"`
	Cvss3   string   `xml:"cvss3,attr"`
	Cwe     string   `xml:"cwe,attr"`
	Href    string   `xml:"href,attr"`
	Public  string   `xml:"public,attr"`
}

// Bugzilla : >definitions>definition>metadata>advisory>bugzilla
// RedHat OVAL
type Bugzilla struct {
	XMLName xml.Name `xml:"bugzilla"`
	ID      string   `xml:"id,attr"`
	URL     string   `xml:"href,attr"`
	Title   string   `xml:",chardata"`
}

// Debian : >definitions>definition>metadata>debian
type Debian struct {
	XMLName  xml.Name `xml:"debian"`
	MoreInfo string   `xml:"moreinfo"`
	Date     string   `xml:"date"`
}

// Tests : >tests
type Tests struct {
	XMLName        xml.Name        `xml:"tests"`
	LineTests      []LineTest      `xml:"line_test"`
	Version55Tests []Version55Test `xml:"version55_test"`
}

// LineTest : >tests>line_test
type LineTest struct {
	XMLName       xml.Name    `xml:"line_test"`
	ID            string      `xml:"id,attr"`
	StateOperator string      `xml:"state_operator,attr"`
	ObjectRefs    []ObjectRef `xml:"object"`
	StateRefs     []StateRef  `xml:"state"`
	Comment       string      `xml:"comment,attr"`
}

// Version55Test : >tests>version55_test
type Version55Test struct {
	XMLName       xml.Name    `xml:"version55_test"`
	ID            string      `xml:"id,attr"`
	StateOperator string      `xml:"state_operator,attr"`
	ObjectRefs    []ObjectRef `xml:"object"`
	StateRefs     []StateRef  `xml:"state"`
	Comment       string      `xml:"comment,attr"`
}

// ObjectRef : >tests>line_test>object-object_ref
//           : >tests>version55_test>object-object_ref
type ObjectRef struct {
	XMLName   xml.Name `xml:"object"`
	ObjectRef string   `xml:"object_ref,attr"`
}

// StateRef : >tests>line_test>state-state_ref
//          : >tests>version55_test>state-state_ref
type StateRef struct {
	XMLName  xml.Name `xml:"state"`
	StateRef string   `xml:"state_ref,attr"`
}

// Objects : >objects
type Objects struct {
	XMLName          xml.Name          `xml:"objects"`
	LineObjects      []LineObject      `xml:"line_object"`
	Version55Objects []Version55Object `xml:"version55_object"`
}

// LineObject : >objects>line_object
type LineObject struct {
	XMLName         xml.Name `xml:"line_object"`
	ID              string   `xml:"id,attr"`
	ShowSubcommands []string `xml:"show_subcommand"`
}

// Version55Object : >objects>version55_object
type Version55Object struct {
	XMLName xml.Name `xml:"version55_object"`
	ID      string   `xml:"id,attr"`
}

// States : >states
type States struct {
	XMLName         xml.Name         `xml:"states"`
	LineStates      []LineState      `xml:"line_state"`
	Version55States []Version55State `xml:"version55_state"`
}

// LineState : >states>line_state
type LineState struct {
	XMLName        xml.Name   `xml:"line_state"`
	ID             string     `xml:"id,attr"`
	ShowSubcommand string     `xml:"show_subcommand"`
	ConfigLine     ConfigLine `xml:"config_line"`
}

// ConfigLine : >states>line_state>config_line
type ConfigLine struct {
	XMLName   xml.Name `xml:"config_line"`
	Body      string   `xml:",chardata"`
	Operation string   `xml:"operation,attr"`
}

// Version55State : >states>version55_state
type Version55State struct {
	XMLName       xml.Name `xml:"version55_state"`
	ID            string   `xml:"id,attr"`
	VersionString string   `xml:"version_string"`
}
<commit_msg>Add Advisory>Issued,Updated (#8)<commit_after>package oval

import (
	"encoding/xml"
)

// Root : root object
type Root struct {
	XMLName     xml.Name    `xml:"oval_definitions"`
	Generator   Generator   `xml:"generator"`
	Definitions Definitions `xml:"definitions"`
	Tests       Tests       `xml:"tests"`
	Objects     Objects     `xml:"objects"`
	States      States      `xml:"states"`
}

// Generator : >generator
type Generator struct {
	XMLName        xml.Name `xml:"generator"`
	ProductName    string   `xml:"product_name"`
	ProductVersion string   `xml:"product_version"`
	SchemaVersion  string   `xml:"schema_version"`
	Timestamp      string   `xml:"timestamp"`
}

// Definitions : >definitions
type Definitions struct {
	XMLName     xml.Name     `xml:"definitions"`
	Definitions []Definition `xml:"definition"`
}

// Definition : >definitions>definition
type Definition struct {
	XMLName     xml.Name    `xml:"definition"`
	ID          string      `xml:"id,attr"`
	Class       string      `xml:"class,attr"`
	Title       string      `xml:"metadata>title"`
	Affecteds   []Affected  `xml:"metadata>affected"`
	References  []Reference `xml:"metadata>reference"`
	Description string      `xml:"metadata>description"`
	Advisory    Advisory    `xml:"metadata>advisory"` // RedHat, Oracle, Ubuntu
	Debian      Debian      `xml:"metadata>debian"`   // Debian
	Criteria    Criteria    `xml:"criteria"`
}

// Criteria : >definitions>definition>criteria
type Criteria struct {
	XMLName    xml.Name    `xml:"criteria"`
	Operator   string      `xml:"operator,attr"`
	Criterias  []Criteria  `xml:"criteria"`
	Criterions []Criterion `xml:"criterion"`
}

// Criterion : >definitions>definition>criteria>*>criterion
type Criterion struct {
	XMLName xml.Name `xml:"criterion"`
	Negate  bool     `xml:"negate,attr"`
	TestRef string   `xml:"test_ref,attr"`
	Comment string   `xml:"comment,attr"`
}

// Affected : >definitions>definition>metadata>affected
type Affected struct {
	XMLName   xml.Name `xml:"affected"`
	Family    string   `xml:"family,attr"`
	Platforms []string `xml:"platform"`
}

// Reference : >definitions>definition>metadata>reference
type Reference struct {
	XMLName xml.Name `xml:"reference"`
	Source  string   `xml:"source,attr"`
	RefID   string   `xml:"ref_id,attr"`
	RefURL  string   `xml:"ref_url,attr"`
}

// Advisory : >definitions>definition>metadata>advisory
// RedHat and Ubuntu OVAL
type Advisory struct {
	XMLName         xml.Name   `xml:"advisory"`
	Severity        string     `xml:"severity"`
	Cves            []Cve      `xml:"cve"`
	Bugzillas       []Bugzilla `xml:"bugzilla"`
	AffectedCPEList []string   `xml:"affected_cpe_list>cpe"`
	Refs            []Ref      `xml:"ref"` // Ubuntu Only
	Bugs            []Bug      `xml:"bug"` // Ubuntu Only
	Issued          struct {
		Date string `xml:"date,attr"`
	} `xml:"issued"`
	Updated struct {
		Date string `xml:"date,attr"`
	} `xml:"updated"`
}

// Ref : >definitions>definition>metadata>advisory>ref
// Ubuntu OVAL
type Ref struct {
	XMLName xml.Name `xml:"ref"`
	URL     string   `xml:",chardata"`
}

// Bug : >definitions>definition>metadata>advisory>bug
// Ubuntu OVAL
type Bug struct {
	XMLName xml.Name `xml:"bug"`
	URL     string   `xml:",chardata"`
}

// Cve : >definitions>definition>metadata>advisory>cve
// RedHat OVAL
type Cve struct {
	XMLName xml.Name `xml:"cve"`
	CveID   string   `xml:",chardata"`
	Cvss2   string   `xml:"cvss2,attr"`
	Cvss3   string   `xml:"cvss3,attr"`
	Cwe     string   `xml:"cwe,attr"`
	Href    string   `xml:"href,attr"`
	Public  string   `xml:"public,attr"`
}

// Bugzilla : >definitions>definition>metadata>advisory>bugzilla
// RedHat OVAL
type Bugzilla struct {
	XMLName xml.Name `xml:"bugzilla"`
	ID      string   `xml:"id,attr"`
	URL     string   `xml:"href,attr"`
	Title   string   `xml:",chardata"`
}

// Debian : >definitions>definition>metadata>debian
type Debian struct {
	XMLName  xml.Name `xml:"debian"`
	MoreInfo string   `xml:"moreinfo"`
	Date     string   `xml:"date"`
}

// Tests : >tests
type Tests struct {
	XMLName        xml.Name        `xml:"tests"`
	LineTests      []LineTest      `xml:"line_test"`
	Version55Tests []Version55Test `xml:"version55_test"`
}

// LineTest : >tests>line_test
type LineTest struct {
	XMLName       xml.Name    `xml:"line_test"`
	ID            string      `xml:"id,attr"`
	StateOperator string      `xml:"state_operator,attr"`
	ObjectRefs    []ObjectRef `xml:"object"`
	StateRefs     []StateRef  `xml:"state"`
	Comment       string      `xml:"comment,attr"`
}

// Version55Test : >tests>version55_test
type Version55Test struct {
	XMLName       xml.Name    `xml:"version55_test"`
	ID            string      `xml:"id,attr"`
	StateOperator string      `xml:"state_operator,attr"`
	ObjectRefs    []ObjectRef `xml:"object"`
	StateRefs     []StateRef  `xml:"state"`
	Comment       string      `xml:"comment,attr"`
}

// ObjectRef : >tests>line_test>object-object_ref
//           : >tests>version55_test>object-object_ref
type ObjectRef struct {
	XMLName   xml.Name `xml:"object"`
	ObjectRef string   `xml:"object_ref,attr"`
}

// StateRef : >tests>line_test>state-state_ref
//          : >tests>version55_test>state-state_ref
type StateRef struct {
	XMLName  xml.Name `xml:"state"`
	StateRef string   `xml:"state_ref,attr"`
}

// Objects : >objects
type Objects struct {
	XMLName          xml.Name          `xml:"objects"`
	LineObjects      []LineObject      `xml:"line_object"`
	Version55Objects []Version55Object `xml:"version55_object"`
}

// LineObject : >objects>line_object
type LineObject struct {
	XMLName         xml.Name `xml:"line_object"`
	ID              string   `xml:"id,attr"`
	ShowSubcommands []string `xml:"show_subcommand"`
}

// Version55Object : >objects>version55_object
type Version55Object struct {
	XMLName xml.Name `xml:"version55_object"`
	ID      string   `xml:"id,attr"`
}

// States : >states
type States struct {
	XMLName         xml.Name         `xml:"states"`
	LineStates      []LineState      `xml:"line_state"`
	Version55States []Version55State `xml:"version55_state"`
}

// LineState : >states>line_state
type LineState struct {
	XMLName        xml.Name   `xml:"line_state"`
	ID             string     `xml:"id,attr"`
	ShowSubcommand string     `xml:"show_subcommand"`
	ConfigLine     ConfigLine `xml:"config_line"`
}

// ConfigLine : >states>line_state>config_line
type ConfigLine struct {
	XMLName   xml.Name `xml:"config_line"`
	Body      string   `xml:",chardata"`
	Operation string   `xml:"operation,attr"`
}

// Version55State : >states>version55_state
type Version55State struct {
	XMLName       xml.Name `xml:"version55_state"`
	ID            string   `xml:"id,attr"`
	VersionString string   `xml:"version_string"`
}
<|endoftext|>


<commit_before>package helpers

import (
	"github.com/jfrogdev/jfrog-cli-go/utils/ioutils"
	"github.com/jfrogdev/jfrog-cli-go/utils/cliutils"
	"time"
	"net/http"
	"io"
	"io/ioutil"
	"errors"
	"bufio"
	"encoding/json"
	"github.com/jfrogdev/jfrog-cli-go/utils/cliutils/log"
)

const BINTRAY_RECONNECT_HEADER = "X-Bintray-Stream-Reconnect-Id"

type StreamManager struct {
	HttpClientDetails ioutils.HttpClientDetails
	Url               string
	IncludeFilter     map[string]struct{}
	ReconnectId       string
}

func (sm *StreamManager) ReadStream(resp *http.Response, writer io.Writer, lastServerInteraction *time.Time) {
	ioReader := resp.Body
	bodyReader := bufio.NewReader(ioReader)
	sm.handleStream(bodyReader, writer, lastServerInteraction)
}

func (sm *StreamManager) handleStream(ioReader io.Reader, writer io.Writer, lastServerInteraction *time.Time) {
	bodyReader := bufio.NewReader(ioReader)
	pReader, pWriter := io.Pipe()
	defer pWriter.Close()
	go func() {
		defer pReader.Close()
		for {
			line, _, err := bodyReader.ReadLine()
			if err != nil {
				break
			}
			*lastServerInteraction = time.Now()
			_, err = pWriter.Write(line)
			if err != nil {
				break
			}
		}
	}()
	streamDecoder := json.NewDecoder(pReader)
	streamEncoder := json.NewEncoder(writer)
	sm.parseStream(streamDecoder, streamEncoder)
}

func (sm *StreamManager) parseStream(streamDecoder *json.Decoder, streamEncoder *json.Encoder) error {
	for {
		var decodedJson map[string]interface{}
		if e := streamDecoder.Decode(&decodedJson); e != nil {
			return e
		}
		if _, ok := sm.IncludeFilter[decodedJson["type"].(string)]; ok || len(sm.IncludeFilter) == 0 {
			if e := streamEncoder.Encode(&decodedJson); e != nil {
				return e
			}
		}
	}
}

func (sm *StreamManager) isReconnection() bool {
	return len(sm.ReconnectId) > 0
}

func (sm *StreamManager) setReconnectHeader() {
	if sm.HttpClientDetails.Headers == nil {
		sm.HttpClientDetails.Headers = map[string]string{}
	}
	sm.HttpClientDetails.Headers[BINTRAY_RECONNECT_HEADER] = sm.ReconnectId
}

func (sm *StreamManager) Connect() (bool, *http.Response) {
	if sm.isReconnection() {
		sm.setReconnectHeader()
	}
	log.Info("Connecting...")
	resp, _, _, e := ioutils.Stream(sm.Url, sm.HttpClientDetails)
	if e != nil {
		return false, resp
	}
	if resp.StatusCode != 200 {
		cliutils.CheckError(errors.New("response: " + resp.Status))
		msgBody, _ := ioutil.ReadAll(resp.Body)
		resp.Body.Close()
		if resp.StatusCode > 400 && resp.StatusCode < 500 {
			cliutils.Exit(cliutils.ExitCodeError, string(msgBody))
		}
		return false, resp

	}
	sm.ReconnectId = resp.Header.Get(BINTRAY_RECONNECT_HEADER)
	log.Info("Connected.")
	return true, resp
}
<commit_msg>Support Bintray stream API<commit_after>package helpers

import (
	"github.com/jfrogdev/jfrog-cli-go/utils/ioutils"
	"github.com/jfrogdev/jfrog-cli-go/utils/cliutils"
	"time"
	"net/http"
	"io"
	"io/ioutil"
	"errors"
	"bufio"
	"encoding/json"
	"github.com/jfrogdev/jfrog-cli-go/utils/cliutils/log"
)

const BINTRAY_RECONNECT_HEADER = "X-Bintray-Stream-Reconnect-Id"

type StreamManager struct {
	HttpClientDetails ioutils.HttpClientDetails
	Url               string
	IncludeFilter     map[string]struct{}
	ReconnectId       string
}

func (sm *StreamManager) ReadStream(resp *http.Response, writer io.Writer, lastServerInteraction *time.Time) {
	ioReader := resp.Body
	bodyReader := bufio.NewReader(ioReader)
	sm.handleStream(bodyReader, writer, lastServerInteraction)
}

func (sm *StreamManager) handleStream(ioReader io.Reader, writer io.Writer, lastServerInteraction *time.Time) {
	bodyReader := bufio.NewReader(ioReader)
	pReader, pWriter := io.Pipe()
	defer pWriter.Close()
	go func() {
		defer pReader.Close()
		for {
			line, _, err := bodyReader.ReadLine()
			if err != nil {
				log.Debug(err)
				break
			}
			*lastServerInteraction = time.Now()
			_, err = pWriter.Write(line)
			if err != nil {
				log.Debug(err)
				break
			}
		}
	}()
	streamDecoder := json.NewDecoder(pReader)
	streamEncoder := json.NewEncoder(writer)
	sm.parseStream(streamDecoder, streamEncoder)
}

func (sm *StreamManager) parseStream(streamDecoder *json.Decoder, streamEncoder *json.Encoder) error {
	for {
		var decodedJson map[string]interface{}
		if e := streamDecoder.Decode(&decodedJson); e != nil {
			log.Debug(e)
			return e
		}
		if _, ok := sm.IncludeFilter[decodedJson["type"].(string)]; ok || len(sm.IncludeFilter) == 0 {
			if e := streamEncoder.Encode(&decodedJson); e != nil {
				log.Debug(e)
				return e
			}
		}
	}
}

func (sm *StreamManager) isReconnection() bool {
	return len(sm.ReconnectId) > 0
}

func (sm *StreamManager) setReconnectHeader() {
	if sm.HttpClientDetails.Headers == nil {
		sm.HttpClientDetails.Headers = map[string]string{}
	}
	sm.HttpClientDetails.Headers[BINTRAY_RECONNECT_HEADER] = sm.ReconnectId
}

func (sm *StreamManager) Connect() (bool, *http.Response) {
	if sm.isReconnection() {
		sm.setReconnectHeader()
	}
	log.Info("Connecting...")
	resp, _, _, e := ioutils.Stream(sm.Url, sm.HttpClientDetails)
	if e != nil {
		return false, resp
	}
	if resp.StatusCode != 200 {
		cliutils.CheckError(errors.New("response: " + resp.Status))
		msgBody, _ := ioutil.ReadAll(resp.Body)
		resp.Body.Close()
		if resp.StatusCode > 400 && resp.StatusCode < 500 {
			cliutils.Exit(cliutils.ExitCodeError, string(msgBody))
		}
		return false, resp

	}
	sm.ReconnectId = resp.Header.Get(BINTRAY_RECONNECT_HEADER)
	log.Info("Connected.")
	return true, resp
}
<|endoftext|>


<commit_before>package webdav

import (
	"net/http"
	"strings"

	"fmt"
	"time"

	"io/ioutil"

	"github.com/petergtz/bitsgo/config"
	"github.com/petergtz/bitsgo/httputil"
)

func NewWebdavResourceSigner(config config.WebdavBlobstoreConfig) *WebdavResourceSigner {
	return &WebdavResourceSigner{
		webdavPrivateEndpoint: config.PrivateEndpoint,
		webdavPublicEndpoint:  config.PublicEndpoint,
		httpClient:            NewHttpClient(config.CACert(), config.SkipCertVerify),
		webdavUsername:        config.Username,
		webdavPassword:        config.Password,
	}
}

type WebdavResourceSigner struct {
	httpClient            *http.Client
	webdavPrivateEndpoint string
	webdavPublicEndpoint  string
	webdavUsername        string
	webdavPassword        string
}

func (signer *WebdavResourceSigner) Sign(resource string, method string) string {
	var url string
	switch strings.ToLower(method) {
	case "put":
		// TODO why do we need a "/" before the resource?
		url = fmt.Sprintf(signer.webdavPrivateEndpoint+"/sign_for_put?path=/%v&expires=%v", resource, time.Now().Unix()+3600)
	case "get":
		url = fmt.Sprintf(signer.webdavPrivateEndpoint+"/sign?path=/%v&expires=%v", resource, time.Now().Unix()+3600)
	}
	request, _ := http.NewRequest("GET", url, nil)
	request.SetBasicAuth(signer.webdavUsername, signer.webdavPassword)
	response, e := signer.httpClient.Do(request)
	if e != nil {
		return "Error during signing. Error: " + e.Error()
	}
	if response.StatusCode != http.StatusOK {
		return "Error during signing. Error code: " + response.Status
	}
	defer response.Body.Close()
	content, e := ioutil.ReadAll(response.Body)
	if e != nil {
		return "Error reading response body. Error: " + e.Error()
	}

	signedUrl := httputil.MustParse(string(content))

	// TODO Is this really what we want to do?
	signedUrl.Host = httputil.MustParse(signer.webdavPublicEndpoint).Host
	signedUrl.Scheme = httputil.MustParse(signer.webdavPublicEndpoint).Scheme

	return signedUrl.String()
}
<commit_msg>Use Request builder in webdav/signed_urls.go<commit_after>package webdav

import (
	"net/http"
	"strings"

	"fmt"
	"time"

	"io/ioutil"

	"github.com/petergtz/bitsgo/config"
	"github.com/petergtz/bitsgo/httputil"
)

func NewWebdavResourceSigner(config config.WebdavBlobstoreConfig) *WebdavResourceSigner {
	return &WebdavResourceSigner{
		webdavPrivateEndpoint: config.PrivateEndpoint,
		webdavPublicEndpoint:  config.PublicEndpoint,
		httpClient:            NewHttpClient(config.CACert(), config.SkipCertVerify),
		webdavUsername:        config.Username,
		webdavPassword:        config.Password,
	}
}

type WebdavResourceSigner struct {
	httpClient            *http.Client
	webdavPrivateEndpoint string
	webdavPublicEndpoint  string
	webdavUsername        string
	webdavPassword        string
}

func (signer *WebdavResourceSigner) Sign(resource string, method string) string {
	var url string
	switch strings.ToLower(method) {
	case "put":
		// TODO why do we need a "/" before the resource?
		url = fmt.Sprintf(signer.webdavPrivateEndpoint+"/sign_for_put?path=/%v&expires=%v", resource, time.Now().Unix()+3600)
	case "get":
		url = fmt.Sprintf(signer.webdavPrivateEndpoint+"/sign?path=/%v&expires=%v", resource, time.Now().Unix()+3600)
	}
	response, e := signer.httpClient.Do(
		httputil.NewRequest("GET", url, nil).
			WithBasicAuth(signer.webdavUsername, signer.webdavPassword).
			Build())
	if e != nil {
		return "Error during signing. Error: " + e.Error()
	}
	if response.StatusCode != http.StatusOK {
		return "Error during signing. Error code: " + response.Status
	}
	defer response.Body.Close()
	content, e := ioutil.ReadAll(response.Body)
	if e != nil {
		return "Error reading response body. Error: " + e.Error()
	}

	signedUrl := httputil.MustParse(string(content))

	// TODO Is this really what we want to do?
	signedUrl.Host = httputil.MustParse(signer.webdavPublicEndpoint).Host
	signedUrl.Scheme = httputil.MustParse(signer.webdavPublicEndpoint).Scheme

	return signedUrl.String()
}
<|endoftext|>


<commit_before>package queryset

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/types"
	"io"
	"sort"
	"strings"
	"text/template"

	"golang.org/x/tools/go/loader"

	"github.com/jinzhu/gorm"
	"github.com/to6ka/go-queryset/parser"
	"github.com/to6ka/go-queryset/queryset/methods"
)

var qsTmpl = template.Must(
	template.New("generator").
		Funcs(template.FuncMap{
			"lcf":      methods.LowerCaseFirstRune,
			"todbname": gorm.ToDBName,
		}).
		Parse(qsCode),
)

type querySetStructConfig struct {
	StructName string
	Name       string
	Methods    methodsSlice
	Fields     []parser.StructField
}

type methodsSlice []methods.Method

func (s methodsSlice) Len() int { return len(s) }
func (s methodsSlice) Less(i, j int) bool {
	return strings.Compare(s[i].GetMethodName(), s[j].GetMethodName()) < 0
}
func (s methodsSlice) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

type querySetStructConfigSlice []querySetStructConfig

func (s querySetStructConfigSlice) Len() int { return len(s) }
func (s querySetStructConfigSlice) Less(i, j int) bool {
	return strings.Compare(s[i].Name, s[j].Name) < 0
}
func (s querySetStructConfigSlice) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

type baseFieldInfo struct {
	name      string // name of field
	typeName  string // name of type of field
	isStruct  bool
	isNumeric bool
	isTime    bool
}

type fieldInfo struct {
	pointed *baseFieldInfo
	baseFieldInfo
	isPointer bool
}

func (fi fieldInfo) getPointed() fieldInfo {
	return fieldInfo{
		baseFieldInfo: *fi.pointed,
	}
}

func getQuerySetMethodsForField(f fieldInfo, qsTypeName string) []methods.Method {
	fCtx := methods.NewQsFieldContext("", f.name, f.typeName, qsTypeName)
	basicTypeMethods := []methods.Method{
		methods.NewBinaryFilterMethod(fCtx.WithName("eq")),
		methods.NewBinaryFilterMethod(fCtx.WithName("ne")),
	}
	if !f.isTime {
		inMethod := methods.NewInFilterMethod(fCtx)
		notInMethod := methods.NewNotInFilterMethod(fCtx)
		basicTypeMethods = append(basicTypeMethods, inMethod, notInMethod)
	}

	numericMethods := []methods.Method{
		methods.NewBinaryFilterMethod(fCtx.WithName("lt")),
		methods.NewBinaryFilterMethod(fCtx.WithName("gt")),
		methods.NewBinaryFilterMethod(fCtx.WithName("lte")),
		methods.NewBinaryFilterMethod(fCtx.WithName("gte")),
		methods.NewOrderAscByMethod(f.name, qsTypeName),
		methods.NewOrderDescByMethod(f.name, qsTypeName),
	}

	if f.isNumeric {
		return append(basicTypeMethods, numericMethods...)
	}

	if f.isStruct {
		// Association was found (any struct or struct pointer)
		return []methods.Method{methods.NewPreloadMethod(f.name, qsTypeName)}
	}

	if f.isPointer {
		ptrMethods := getQuerySetMethodsForField(f.getPointed(), qsTypeName)
		return append(ptrMethods,
			methods.NewIsNullMethod(f.name, qsTypeName),
			methods.NewIsNotNullMethod(f.name, qsTypeName))
	}

	// it's a string
	return basicTypeMethods
}

func getOriginalTypeName(t *types.Named, pkgInfo *loader.PackageInfo) string {
	if t.Obj().Pkg() == pkgInfo.Pkg {
		// t is from the same package as a struct
		return t.Obj().Name()
	}

	// t is an imported from another package type
	return fmt.Sprintf("%s.%s", t.Obj().Pkg().Name(), t.Obj().Name())
}

func generateFieldInfo(pkgInfo *loader.PackageInfo, name string,
	fieldType fmt.Stringer) *fieldInfo {

	typeName := fieldType.String()

	if typeName == "time.Time" {
		return &fieldInfo{
			baseFieldInfo: baseFieldInfo{
				name:      name,
				typeName:  typeName,
				isNumeric: true,
				isTime:    true,
			},
		}
	}

	switch t := fieldType.(type) {
	case *types.Basic:
		return &fieldInfo{
			baseFieldInfo: baseFieldInfo{
				name:      name,
				typeName:  typeName,
				isNumeric: t.Info()&types.IsNumeric != 0,
			},
		}
	case *types.Named:
		r := generateFieldInfo(pkgInfo, name, t.Underlying())
		if r != nil {
			r.typeName = getOriginalTypeName(t, pkgInfo)
		}
		return r
	case *types.Struct:
		return &fieldInfo{
			baseFieldInfo: baseFieldInfo{
				name:     name,
				typeName: typeName,
				isStruct: true,
			},
		}
	case *types.Pointer:
		pf := generateFieldInfo(pkgInfo, name, t.Elem())
		return &fieldInfo{
			baseFieldInfo: baseFieldInfo{
				name:     name,
				typeName: typeName,
			},
			isPointer: true,
			pointed:   &pf.baseFieldInfo,
		}
	default:
		// no filtering is needed
		return nil
	}
}

func getQuerySetFieldMethods(fields []fieldInfo, qsTypeName string) []methods.Method {
	ret := []methods.Method{}
	for _, f := range fields {
		methods := getQuerySetMethodsForField(f, qsTypeName)
		ret = append(ret, methods...)
	}

	return ret
}

func getUpdaterTypeName(structTypeName string) string {
	return structTypeName + "Updater"
}

func getUpdaterMethods(fields []fieldInfo, structTypeName string) []methods.Method {
	updaterTypeName := getUpdaterTypeName(structTypeName)
	ret := []methods.Method{
		methods.NewUpdaterUpdateMethod(updaterTypeName),
		methods.NewUpdaterUpdateNumMethod(updaterTypeName),
	}
	for _, f := range fields {
		if f.isPointer {
			// TODO
			continue
		}
		dbSchemaTypeName := structTypeName + "DBSchema"
		ret = append(ret,
			methods.NewUpdaterSetMethod(f.name, f.typeName, updaterTypeName,
				dbSchemaTypeName))
	}
	return ret
}

func getMethodsForStruct(structTypeName string, fieldInfos []fieldInfo) []methods.Method {
	qsTypeName := structTypeName + "QuerySet"

	ret := []methods.Method{
		methods.NewLimitMethod(qsTypeName),
		methods.NewAllMethod(structTypeName, qsTypeName),
		methods.NewOneMethod(structTypeName, qsTypeName),
		methods.NewGetUpdaterMethod(qsTypeName, getUpdaterTypeName(structTypeName)),
		methods.NewDeleteMethod(qsTypeName, structTypeName),
		methods.NewStructModifierMethod("Create", structTypeName),
		methods.NewStructModifierMethod("Delete", structTypeName),
		methods.NewCountMethod(qsTypeName),
	}

	fieldMethods := getQuerySetFieldMethods(fieldInfos, qsTypeName)
	ret = append(ret, fieldMethods...)

	ret = append(ret, getUpdaterMethods(fieldInfos, structTypeName)...)

	return ret
}

func doesNeedToGenerateQuerySet(doc *ast.CommentGroup) bool {
	if doc == nil {
		return false
	}

	for _, c := range doc.List {
		parts := strings.Split(strings.TrimSpace(c.Text), ":")
		ok := len(parts) == 2 &&
			strings.TrimSpace(strings.TrimPrefix(parts[0], "//")) == "gen" &&
			strings.TrimSpace(parts[1]) == "qs"
		if ok {
			return true
		}
	}

	return false
}

func generateQuerySetConfigs(pkgInfo *loader.PackageInfo,
	structs parser.ParsedStructs) querySetStructConfigSlice {

	querySetStructConfigs := querySetStructConfigSlice{}

	for structTypeName, ps := range structs {
		if !doesNeedToGenerateQuerySet(ps.Doc) {
			continue
		}

		fieldInfos := []fieldInfo{}
		for _, f := range ps.Fields {
			fi := generateFieldInfo(pkgInfo, f.Name, f.Type)
			if fi == nil {
				continue
			}
			fieldInfos = append(fieldInfos, *fi)
		}

		methods := getMethodsForStruct(structTypeName, fieldInfos)

		qsConfig := querySetStructConfig{
			StructName: structTypeName,
			Name:       structTypeName + "QuerySet",
			Methods:    methods,
			Fields:     ps.Fields,
		}
		sort.Sort(qsConfig.Methods)
		querySetStructConfigs = append(querySetStructConfigs, qsConfig)
	}

	return querySetStructConfigs
}

// GenerateQuerySetsForStructs is an internal method to retrieve querysets
// generated code from parsed structs
func GenerateQuerySetsForStructs(pkgInfo *loader.PackageInfo, structs parser.ParsedStructs) (io.Reader, error) {

	querySetStructConfigs := generateQuerySetConfigs(pkgInfo, structs)
	if len(querySetStructConfigs) == 0 {
		return nil, nil
	}

	sort.Sort(querySetStructConfigs)

	var b bytes.Buffer
	err := qsTmpl.Execute(&b, struct {
		Configs querySetStructConfigSlice
	}{
		Configs: querySetStructConfigs,
	})

	if err != nil {
		return nil, fmt.Errorf("can't generate structs query sets: %s", err)
	}

	return &b, nil
}

const qsCode = `
// ===== BEGIN of all query sets

{{ range .Configs }}
  // ===== BEGIN of query set {{ .Name }}

	// {{ .Name }} is an queryset type for {{ .StructName }}
  type {{ .Name }} struct {
	  db *gorm.DB
  }

  // New{{ .Name }} constructs new {{ .Name }}
  func New{{ .Name }}(db *gorm.DB) {{ .Name }} {
	  return {{ .Name }}{
		  db: db.Model(&{{ .StructName }}{}),
	  }
  }

	func (qs {{ .Name }}) w(db *gorm.DB) {{ .Name }} {
	  return New{{ .Name }}(db)
  }

	{{ range .Methods }}
		{{ .GetDoc .GetMethodName }}
		func ({{ .GetReceiverDeclaration }}) {{ .GetMethodName }}({{ .GetArgsDeclaration }})
		{{- .GetReturnValuesDeclaration }} {
      {{ .GetBody }}
		}
	{{ end }}

  // ===== END of query set {{ .Name }}

	// ===== BEGIN of {{ .StructName }} modifiers

	{{ $ft := printf "%s%s" .StructName "DBSchemaField" | lcf }}
	type {{ $ft }} string

	// {{ .StructName }}DBSchema stores db field names of {{ .StructName }}
	var {{ .StructName }}DBSchema = struct {
		{{ range .Fields }}
			{{ .Name }} {{ $ft }}
		{{- end }}
	}{
		{{ range .Fields }}
			{{ .Name }}: {{ $ft }}("{{ .Name | todbname }}"),
		{{- end }}
	}

	// Update updates {{ .StructName }} fields by primary key
	func (o *{{ .StructName }}) Update(db *gorm.DB, fields ...{{ $ft }}) error {
		dbNameToFieldName := map[string]interface{}{
			{{- range .Fields }}
				"{{ .Name | todbname }}": o.{{ .Name }},
			{{- end }}
		}
		u := map[string]interface{}{}
		for _, f := range fields {
			fs := string(f)
			u[fs] = dbNameToFieldName[fs]
		}
		if err := db.Model(o).Updates(u).Error; err != nil {
			if err == gorm.ErrRecordNotFound {
				return err
			}

			return fmt.Errorf("can't update {{ .StructName }} %v fields %v: %s",
				o, fields, err)
		}

		return nil
	}

	// {{ .StructName }}Updater is an {{ .StructName }} updates manager
	type {{ .StructName }}Updater struct {
		fields map[string]interface{}
		db *gorm.DB
	}

	// New{{ .StructName }}Updater creates new {{ .StructName }} updater
	func New{{ .StructName }}Updater(db *gorm.DB) {{ .StructName }}Updater {
		return {{ .StructName }}Updater{
			fields: map[string]interface{}{},
			db: db.Model(&{{ .StructName }}{}),
		}
	}

	// ===== END of {{ .StructName }} modifiers
{{ end }}

// ===== END of all query sets
`
<commit_msg>identation fix<commit_after>package queryset

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/types"
	"io"
	"sort"
	"strings"
	"text/template"

	"golang.org/x/tools/go/loader"

	"github.com/jinzhu/gorm"
	"github.com/to6ka/go-queryset/parser"
	"github.com/to6ka/go-queryset/queryset/methods"
)

var qsTmpl = template.Must(
	template.New("generator").
		Funcs(template.FuncMap{
			"lcf":      methods.LowerCaseFirstRune,
			"todbname": gorm.ToDBName,
		}).
		Parse(qsCode),
)

type querySetStructConfig struct {
	StructName string
	Name       string
	Methods    methodsSlice
	Fields     []parser.StructField
}

type methodsSlice []methods.Method

func (s methodsSlice) Len() int { return len(s) }
func (s methodsSlice) Less(i, j int) bool {
	return strings.Compare(s[i].GetMethodName(), s[j].GetMethodName()) < 0
}
func (s methodsSlice) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

type querySetStructConfigSlice []querySetStructConfig

func (s querySetStructConfigSlice) Len() int { return len(s) }
func (s querySetStructConfigSlice) Less(i, j int) bool {
	return strings.Compare(s[i].Name, s[j].Name) < 0
}
func (s querySetStructConfigSlice) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

type baseFieldInfo struct {
	name      string // name of field
	typeName  string // name of type of field
	isStruct  bool
	isNumeric bool
	isTime    bool
}

type fieldInfo struct {
	pointed *baseFieldInfo
	baseFieldInfo
	isPointer bool
}

func (fi fieldInfo) getPointed() fieldInfo {
	return fieldInfo{
		baseFieldInfo: *fi.pointed,
	}
}

func getQuerySetMethodsForField(f fieldInfo, qsTypeName string) []methods.Method {
	fCtx := methods.NewQsFieldContext("", f.name, f.typeName, qsTypeName)
	basicTypeMethods := []methods.Method{
		methods.NewBinaryFilterMethod(fCtx.WithName("eq")),
		methods.NewBinaryFilterMethod(fCtx.WithName("ne")),
	}
	if !f.isTime {
		inMethod := methods.NewInFilterMethod(fCtx)
		notInMethod := methods.NewNotInFilterMethod(fCtx)
		basicTypeMethods = append(basicTypeMethods, inMethod, notInMethod)
	}

	numericMethods := []methods.Method{
		methods.NewBinaryFilterMethod(fCtx.WithName("lt")),
		methods.NewBinaryFilterMethod(fCtx.WithName("gt")),
		methods.NewBinaryFilterMethod(fCtx.WithName("lte")),
		methods.NewBinaryFilterMethod(fCtx.WithName("gte")),
		methods.NewOrderAscByMethod(f.name, qsTypeName),
		methods.NewOrderDescByMethod(f.name, qsTypeName),
	}

	if f.isNumeric {
		return append(basicTypeMethods, numericMethods...)
	}

	if f.isStruct {
		// Association was found (any struct or struct pointer)
		return []methods.Method{methods.NewPreloadMethod(f.name, qsTypeName)}
	}

	if f.isPointer {
		ptrMethods := getQuerySetMethodsForField(f.getPointed(), qsTypeName)
		return append(ptrMethods,
			methods.NewIsNullMethod(f.name, qsTypeName),
			methods.NewIsNotNullMethod(f.name, qsTypeName))
	}

	// it's a string
	return basicTypeMethods
}

func getOriginalTypeName(t *types.Named, pkgInfo *loader.PackageInfo) string {
	if t.Obj().Pkg() == pkgInfo.Pkg {
		// t is from the same package as a struct
		return t.Obj().Name()
	}

	// t is an imported from another package type
	return fmt.Sprintf("%s.%s", t.Obj().Pkg().Name(), t.Obj().Name())
}

func generateFieldInfo(pkgInfo *loader.PackageInfo, name string,
	fieldType fmt.Stringer) *fieldInfo {

	typeName := fieldType.String()

	if typeName == "time.Time" {
		return &fieldInfo{
			baseFieldInfo: baseFieldInfo{
				name:      name,
				typeName:  typeName,
				isNumeric: true,
				isTime:    true,
			},
		}
	}

	switch t := fieldType.(type) {
	case *types.Basic:
		return &fieldInfo{
			baseFieldInfo: baseFieldInfo{
				name:      name,
				typeName:  typeName,
				isNumeric: t.Info()&types.IsNumeric != 0,
			},
		}
	case *types.Named:
		r := generateFieldInfo(pkgInfo, name, t.Underlying())
		if r != nil {
			r.typeName = getOriginalTypeName(t, pkgInfo)
		}
		return r
	case *types.Struct:
		return &fieldInfo{
			baseFieldInfo: baseFieldInfo{
				name:     name,
				typeName: typeName,
				isStruct: true,
			},
		}
	case *types.Pointer:
		pf := generateFieldInfo(pkgInfo, name, t.Elem())
		return &fieldInfo{
			baseFieldInfo: baseFieldInfo{
				name:     name,
				typeName: typeName,
			},
			isPointer: true,
			pointed:   &pf.baseFieldInfo,
		}
	default:
		// no filtering is needed
		return nil
	}
}

func getQuerySetFieldMethods(fields []fieldInfo, qsTypeName string) []methods.Method {
	ret := []methods.Method{}
	for _, f := range fields {
		methods := getQuerySetMethodsForField(f, qsTypeName)
		ret = append(ret, methods...)
	}

	return ret
}

func getUpdaterTypeName(structTypeName string) string {
	return structTypeName + "Updater"
}

func getUpdaterMethods(fields []fieldInfo, structTypeName string) []methods.Method {
	updaterTypeName := getUpdaterTypeName(structTypeName)
	ret := []methods.Method{
		methods.NewUpdaterUpdateMethod(updaterTypeName),
		methods.NewUpdaterUpdateNumMethod(updaterTypeName),
	}
	for _, f := range fields {
		if f.isPointer {
			// TODO
			continue
		}
		dbSchemaTypeName := structTypeName + "DBSchema"
		ret = append(ret,
			methods.NewUpdaterSetMethod(f.name, f.typeName, updaterTypeName,
				dbSchemaTypeName))
	}
	return ret
}

func getMethodsForStruct(structTypeName string, fieldInfos []fieldInfo) []methods.Method {
	qsTypeName := structTypeName + "QuerySet"

	ret := []methods.Method{
		methods.NewLimitMethod(qsTypeName),
		methods.NewAllMethod(structTypeName, qsTypeName),
		methods.NewOneMethod(structTypeName, qsTypeName),
		methods.NewGetUpdaterMethod(qsTypeName, getUpdaterTypeName(structTypeName)),
		methods.NewDeleteMethod(qsTypeName, structTypeName),
		methods.NewStructModifierMethod("Create", structTypeName),
		methods.NewStructModifierMethod("Delete", structTypeName),
		methods.NewCountMethod(qsTypeName),
	}

	fieldMethods := getQuerySetFieldMethods(fieldInfos, qsTypeName)
	ret = append(ret, fieldMethods...)

	ret = append(ret, getUpdaterMethods(fieldInfos, structTypeName)...)

	return ret
}

func doesNeedToGenerateQuerySet(doc *ast.CommentGroup) bool {
	if doc == nil {
		return false
	}

	for _, c := range doc.List {
		parts := strings.Split(strings.TrimSpace(c.Text), ":")
		ok := len(parts) == 2 &&
			strings.TrimSpace(strings.TrimPrefix(parts[0], "//")) == "gen" &&
			strings.TrimSpace(parts[1]) == "qs"
		if ok {
			return true
		}
	}

	return false
}

func generateQuerySetConfigs(pkgInfo *loader.PackageInfo,
	structs parser.ParsedStructs) querySetStructConfigSlice {

	querySetStructConfigs := querySetStructConfigSlice{}

	for structTypeName, ps := range structs {
		if !doesNeedToGenerateQuerySet(ps.Doc) {
			continue
		}

		fieldInfos := []fieldInfo{}
		for _, f := range ps.Fields {
			fi := generateFieldInfo(pkgInfo, f.Name, f.Type)
			if fi == nil {
				continue
			}
			fieldInfos = append(fieldInfos, *fi)
		}

		methods := getMethodsForStruct(structTypeName, fieldInfos)

		qsConfig := querySetStructConfig{
			StructName: structTypeName,
			Name:       structTypeName + "QuerySet",
			Methods:    methods,
			Fields:     ps.Fields,
		}
		sort.Sort(qsConfig.Methods)
		querySetStructConfigs = append(querySetStructConfigs, qsConfig)
	}

	return querySetStructConfigs
}

// GenerateQuerySetsForStructs is an internal method to retrieve querysets
// generated code from parsed structs
func GenerateQuerySetsForStructs(pkgInfo *loader.PackageInfo, structs parser.ParsedStructs) (io.Reader, error) {

	querySetStructConfigs := generateQuerySetConfigs(pkgInfo, structs)
	if len(querySetStructConfigs) == 0 {
		return nil, nil
	}

	sort.Sort(querySetStructConfigs)

	var b bytes.Buffer
	err := qsTmpl.Execute(&b, struct {
		Configs querySetStructConfigSlice
	}{
		Configs: querySetStructConfigs,
	})

	if err != nil {
		return nil, fmt.Errorf("can't generate structs query sets: %s", err)
	}

	return &b, nil
}

const qsCode = `
// ===== BEGIN of all query sets

{{ range .Configs }}
	// ===== BEGIN of query set {{ .Name }}

	// {{ .Name }} is an queryset type for {{ .StructName }}
	type {{ .Name }} struct {
		db *gorm.DB
	}

	// New{{ .Name }} constructs new {{ .Name }}
	func New{{ .Name }}(db *gorm.DB) {{ .Name }} {
		return {{ .Name }}{
			db: db.Model(&{{ .StructName }}{}),
		}
	}

	func (qs {{ .Name }}) w(db *gorm.DB) {{ .Name }} {
		return New{{ .Name }}(db)
	}

	{{ range .Methods }}
		{{ .GetDoc .GetMethodName }}
		func ({{ .GetReceiverDeclaration }}) {{ .GetMethodName }}({{ .GetArgsDeclaration }})
		{{- .GetReturnValuesDeclaration }} {
			{{ .GetBody }}
		}
	{{ end }}

	// ===== END of query set {{ .Name }}

	// ===== BEGIN of {{ .StructName }} modifiers

	{{ $ft := printf "%s%s" .StructName "DBSchemaField" | lcf }}
	type {{ $ft }} string

	// {{ .StructName }}DBSchema stores db field names of {{ .StructName }}
	var {{ .StructName }}DBSchema = struct {
		{{ range .Fields }}
			{{ .Name }} {{ $ft }}
		{{- end }}
	}{
		{{ range .Fields }}
			{{ .Name }}: {{ $ft }}("{{ .Name | todbname }}"),
		{{- end }}
	}

	// Update updates {{ .StructName }} fields by primary key
	func (o *{{ .StructName }}) Update(db *gorm.DB, fields ...{{ $ft }}) error {
		dbNameToFieldName := map[string]interface{}{
			{{- range .Fields }}
				"{{ .Name | todbname }}": o.{{ .Name }},
			{{- end }}
		}
		u := map[string]interface{}{}
		for _, f := range fields {
			fs := string(f)
			u[fs] = dbNameToFieldName[fs]
		}
		if err := db.Model(o).Updates(u).Error; err != nil {
			if err == gorm.ErrRecordNotFound {
				return err
			}

			return fmt.Errorf("can't update {{ .StructName }} %v fields %v: %s",
				o, fields, err)
		}

		return nil
	}

	// {{ .StructName }}Updater is an {{ .StructName }} updates manager
	type {{ .StructName }}Updater struct {
		fields map[string]interface{}
		db *gorm.DB
	}

	// New{{ .StructName }}Updater creates new {{ .StructName }} updater
	func New{{ .StructName }}Updater(db *gorm.DB) {{ .StructName }}Updater {
		return {{ .StructName }}Updater{
			fields: map[string]interface{}{},
			db: db.Model(&{{ .StructName }}{}),
		}
	}

	// ===== END of {{ .StructName }} modifiers
{{ end }}

// ===== END of all query sets
`
<|endoftext|>


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

package smtp

import (
	"crypto/hmac"
	"crypto/md5"
	"errors"
	"fmt"
)

// Auth is implemented by an SMTP authentication mechanism.
type Auth interface {
	// Start begins an authentication with a server.
	// It returns the name of the authentication protocol
	// and optionally data to include in the initial AUTH message
	// sent to the server. It can return proto == "" to indicate
	// that the authentication should be skipped.
	// If it returns a non-nil error, the SMTP client aborts
	// the authentication attempt and closes the connection.
	Start(server *ServerInfo) (proto string, toServer []byte, err error)

	// Next continues the authentication. The server has just sent
	// the fromServer data. If more is true, the server expects a
	// response, which Next should return as toServer; otherwise
	// Next should return toServer == nil.
	// If Next returns a non-nil error, the SMTP client aborts
	// the authentication attempt and closes the connection.
	Next(fromServer []byte, more bool) (toServer []byte, err error)
}

// ServerInfo records information about an SMTP server.
type ServerInfo struct {
	Name string   // SMTP server name
	TLS  bool     // using TLS, with valid certificate for Name
	Auth []string // advertised authentication mechanisms
}

type plainAuth struct {
	identity, username, password string
	host                         string
}

// PlainAuth returns an Auth that implements the PLAIN authentication
// mechanism as defined in RFC 4616. The returned Auth uses the given
// username and password to authenticate to host and act as identity.
// Usually identity should be the empty string, to act as username.
//
// PlainAuth will only send the credentials if the connection is using TLS
// or is connected to localhost. Otherwise authentication will fail with an
// error, without sending the credentials.
func PlainAuth(identity, username, password, host string) Auth {
	return &plainAuth{identity, username, password, host}
}

func isLocalhost(name string) bool {
	return name == "localhost" || name == "127.0.0.1" || name == "::1"
}

func (a *plainAuth) Start(server *ServerInfo) (string, []byte, error) {
	// Must have TLS, or else localhost server.
	// Note: If TLS is not true, then we can't trust ANYTHING in ServerInfo.
	// In particular, it doesn't matter if the server advertises PLAIN auth.
	// That might just be the attacker saying
	// "it's ok, you can trust me with your password."
	if !server.TLS && !isLocalhost(server.Name) {
		return "", nil, errors.New("unencrypted connection")
	}
	if server.Name != a.host {
		return "", nil, errors.New("wrong host name")
	}
	resp := []byte(a.identity + "\x00" + a.username + "\x00" + a.password)
	return "PLAIN", resp, nil
}

func (a *plainAuth) Next(fromServer []byte, more bool) ([]byte, error) {
	if more {
		// We've already sent everything.
		return nil, errors.New("unexpected server challenge")
	}
	return nil, nil
}

type cramMD5Auth struct {
	username, secret string
}

// CRAMMD5Auth returns an Auth that implements the CRAM-MD5 authentication
// mechanism as defined in RFC 2195.
// The returned Auth uses the given username and secret to authenticate
// to the server using the challenge-response mechanism.
func CRAMMD5Auth(username, secret string) Auth {
	return &cramMD5Auth{username, secret}
}

func (a *cramMD5Auth) Start(server *ServerInfo) (string, []byte, error) {
	return "CRAM-MD5", nil, nil
}

func (a *cramMD5Auth) Next(fromServer []byte, more bool) ([]byte, error) {
	if more {
		d := hmac.New(md5.New, []byte(a.secret))
		d.Write(fromServer)
		s := make([]byte, 0, d.Size())
		return []byte(fmt.Sprintf("%s %x", a.username, d.Sum(s))), nil
	}
	return nil, nil
}
<commit_msg>net/smtp: match actual behavior for Auth.Start<commit_after>// Copyright 2010 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package smtp

import (
	"crypto/hmac"
	"crypto/md5"
	"errors"
	"fmt"
)

// Auth is implemented by an SMTP authentication mechanism.
type Auth interface {
	// Start begins an authentication with a server.
	// It returns the name of the authentication protocol
	// and optionally data to include in the initial AUTH message
	// sent to the server.
	// If it returns a non-nil error, the SMTP client aborts
	// the authentication attempt and closes the connection.
	Start(server *ServerInfo) (proto string, toServer []byte, err error)

	// Next continues the authentication. The server has just sent
	// the fromServer data. If more is true, the server expects a
	// response, which Next should return as toServer; otherwise
	// Next should return toServer == nil.
	// If Next returns a non-nil error, the SMTP client aborts
	// the authentication attempt and closes the connection.
	Next(fromServer []byte, more bool) (toServer []byte, err error)
}

// ServerInfo records information about an SMTP server.
type ServerInfo struct {
	Name string   // SMTP server name
	TLS  bool     // using TLS, with valid certificate for Name
	Auth []string // advertised authentication mechanisms
}

type plainAuth struct {
	identity, username, password string
	host                         string
}

// PlainAuth returns an Auth that implements the PLAIN authentication
// mechanism as defined in RFC 4616. The returned Auth uses the given
// username and password to authenticate to host and act as identity.
// Usually identity should be the empty string, to act as username.
//
// PlainAuth will only send the credentials if the connection is using TLS
// or is connected to localhost. Otherwise authentication will fail with an
// error, without sending the credentials.
func PlainAuth(identity, username, password, host string) Auth {
	return &plainAuth{identity, username, password, host}
}

func isLocalhost(name string) bool {
	return name == "localhost" || name == "127.0.0.1" || name == "::1"
}

func (a *plainAuth) Start(server *ServerInfo) (string, []byte, error) {
	// Must have TLS, or else localhost server.
	// Note: If TLS is not true, then we can't trust ANYTHING in ServerInfo.
	// In particular, it doesn't matter if the server advertises PLAIN auth.
	// That might just be the attacker saying
	// "it's ok, you can trust me with your password."
	if !server.TLS && !isLocalhost(server.Name) {
		return "", nil, errors.New("unencrypted connection")
	}
	if server.Name != a.host {
		return "", nil, errors.New("wrong host name")
	}
	resp := []byte(a.identity + "\x00" + a.username + "\x00" + a.password)
	return "PLAIN", resp, nil
}

func (a *plainAuth) Next(fromServer []byte, more bool) ([]byte, error) {
	if more {
		// We've already sent everything.
		return nil, errors.New("unexpected server challenge")
	}
	return nil, nil
}

type cramMD5Auth struct {
	username, secret string
}

// CRAMMD5Auth returns an Auth that implements the CRAM-MD5 authentication
// mechanism as defined in RFC 2195.
// The returned Auth uses the given username and secret to authenticate
// to the server using the challenge-response mechanism.
func CRAMMD5Auth(username, secret string) Auth {
	return &cramMD5Auth{username, secret}
}

func (a *cramMD5Auth) Start(server *ServerInfo) (string, []byte, error) {
	return "CRAM-MD5", nil, nil
}

func (a *cramMD5Auth) Next(fromServer []byte, more bool) ([]byte, error) {
	if more {
		d := hmac.New(md5.New, []byte(a.secret))
		d.Write(fromServer)
		s := make([]byte, 0, d.Size())
		return []byte(fmt.Sprintf("%s %x", a.username, d.Sum(s))), nil
	}
	return nil, nil
}
<|endoftext|>