<commit_before>package consul

import (
	"fmt"
	"time"

	metrics "github.com/armon/go-metrics"
	"github.com/hashicorp/consul/acl"
	"github.com/hashicorp/consul/agent/consul/state"
	"github.com/hashicorp/consul/agent/structs"
	memdb "github.com/hashicorp/go-memdb"
	"github.com/mitchellh/copystructure"
)

// The ConfigEntry endpoint is used to query centralized config information
type ConfigEntry struct {
	srv *Server
}

// Apply does an upsert of the given config entry.
func (c *ConfigEntry) Apply(args *structs.ConfigEntryRequest, reply *bool) error {
	if err := c.srv.validateEnterpriseRequest(args.Entry.GetEnterpriseMeta(), true); err != nil {
		return err
	}

	// Ensure that all config entry writes go to the primary datacenter. These will then
	// be replicated to all the other datacenters.
	args.Datacenter = c.srv.config.PrimaryDatacenter

	if done, err := c.srv.forward("ConfigEntry.Apply", args, args, reply); done {
		return err
	}
	defer metrics.MeasureSince([]string{"config_entry", "apply"}, time.Now())

	entMeta := args.Entry.GetEnterpriseMeta()
	authz, err := c.srv.ResolveTokenAndDefaultMeta(args.Token, entMeta, nil)
	if err != nil {
		return err
	}

	// Normalize and validate the incoming config entry.
	if err := args.Entry.Normalize(); err != nil {
		return err
	}
	if err := args.Entry.Validate(); err != nil {
		return err
	}

	if authz != nil && !args.Entry.CanWrite(authz) {
		return acl.ErrPermissionDenied
	}

	if args.Op != structs.ConfigEntryUpsert && args.Op != structs.ConfigEntryUpsertCAS {
		args.Op = structs.ConfigEntryUpsert
	}
	resp, err := c.srv.raftApply(structs.ConfigEntryRequestType, args)
	if err != nil {
		return err
	}
	if respErr, ok := resp.(error); ok {
		return respErr
	}
	if respBool, ok := resp.(bool); ok {
		*reply = respBool
	}

	return nil
}

// Get returns a single config entry by Kind/Name.
func (c *ConfigEntry) Get(args *structs.ConfigEntryQuery, reply *structs.ConfigEntryResponse) error {
	if err := c.srv.validateEnterpriseRequest(&args.EnterpriseMeta, false); err != nil {
		return err
	}

	if done, err := c.srv.forward("ConfigEntry.Get", args, args, reply); done {
		return err
	}
	defer metrics.MeasureSince([]string{"config_entry", "get"}, time.Now())

	authz, err := c.srv.ResolveTokenAndDefaultMeta(args.Token, &args.EnterpriseMeta, nil)
	if err != nil {
		return err
	}

	// Create a dummy config entry to check the ACL permissions.
	lookupEntry, err := structs.MakeConfigEntry(args.Kind, args.Name)
	if err != nil {
		return err
	}
	if authz != nil && !lookupEntry.CanRead(authz) {
		return acl.ErrPermissionDenied
	}

	return c.srv.blockingQuery(
		&args.QueryOptions,
		&reply.QueryMeta,
		func(ws memdb.WatchSet, state *state.Store) error {
			index, entry, err := state.ConfigEntry(ws, args.Kind, args.Name, &args.EnterpriseMeta)
			if err != nil {
				return err
			}

			reply.Index = index
			if entry == nil {
				return nil
			}

			reply.Entry = entry
			return nil
		})
}

// List returns all the config entries of the given kind. If Kind is blank,
// all existing config entries will be returned.
func (c *ConfigEntry) List(args *structs.ConfigEntryQuery, reply *structs.IndexedConfigEntries) error {
	if err := c.srv.validateEnterpriseRequest(&args.EnterpriseMeta, false); err != nil {
		return err
	}

	if done, err := c.srv.forward("ConfigEntry.List", args, args, reply); done {
		return err
	}
	defer metrics.MeasureSince([]string{"config_entry", "list"}, time.Now())

	authz, err := c.srv.ResolveTokenAndDefaultMeta(args.Token, &args.EnterpriseMeta, nil)
	if err != nil {
		return err
	}

	if args.Kind != "" && !structs.ValidateConfigEntryKind(args.Kind) {
		return fmt.Errorf("invalid config entry kind: %s", args.Kind)
	}

	return c.srv.blockingQuery(
		&args.QueryOptions,
		&reply.QueryMeta,
		func(ws memdb.WatchSet, state *state.Store) error {
			index, entries, err := state.ConfigEntriesByKind(ws, args.Kind, &args.EnterpriseMeta)
			if err != nil {
				return err
			}

			// Filter the entries returned by ACL permissions.
			filteredEntries := make([]structs.ConfigEntry, 0, len(entries))
			for _, entry := range entries {
				if authz != nil && !entry.CanRead(authz) {
					continue
				}
				filteredEntries = append(filteredEntries, entry)
			}

			reply.Kind = args.Kind
			reply.Index = index
			reply.Entries = filteredEntries
			return nil
		})
}

// ListAll returns all the known configuration entries
func (c *ConfigEntry) ListAll(args *structs.DCSpecificRequest, reply *structs.IndexedGenericConfigEntries) error {
	if err := c.srv.validateEnterpriseRequest(&args.EnterpriseMeta, false); err != nil {
		return err
	}

	if done, err := c.srv.forward("ConfigEntry.ListAll", args, args, reply); done {
		return err
	}
	defer metrics.MeasureSince([]string{"config_entry", "listAll"}, time.Now())

	authz, err := c.srv.ResolveTokenAndDefaultMeta(args.Token, &args.EnterpriseMeta, nil)
	if err != nil {
		return err
	}

	return c.srv.blockingQuery(
		&args.QueryOptions,
		&reply.QueryMeta,
		func(ws memdb.WatchSet, state *state.Store) error {
			index, entries, err := state.ConfigEntries(ws, &args.EnterpriseMeta)
			if err != nil {
				return err
			}

			// Filter the entries returned by ACL permissions.
			filteredEntries := make([]structs.ConfigEntry, 0, len(entries))
			for _, entry := range entries {
				if authz != nil && !entry.CanRead(authz) {
					continue
				}
				filteredEntries = append(filteredEntries, entry)
			}

			reply.Entries = filteredEntries
			reply.Index = index
			return nil
		})
}

// Delete deletes a config entry.
func (c *ConfigEntry) Delete(args *structs.ConfigEntryRequest, reply *struct{}) error {
	if err := c.srv.validateEnterpriseRequest(args.Entry.GetEnterpriseMeta(), true); err != nil {
		return err
	}

	// Ensure that all config entry writes go to the primary datacenter. These will then
	// be replicated to all the other datacenters.
	args.Datacenter = c.srv.config.PrimaryDatacenter

	if done, err := c.srv.forward("ConfigEntry.Delete", args, args, reply); done {
		return err
	}
	defer metrics.MeasureSince([]string{"config_entry", "delete"}, time.Now())

	authz, err := c.srv.ResolveTokenAndDefaultMeta(args.Token, args.Entry.GetEnterpriseMeta(), nil)
	if err != nil {
		return err
	}

	// Normalize the incoming entry.
	if err := args.Entry.Normalize(); err != nil {
		return err
	}

	if authz != nil && !args.Entry.CanWrite(authz) {
		return acl.ErrPermissionDenied
	}

	args.Op = structs.ConfigEntryDelete
	resp, err := c.srv.raftApply(structs.ConfigEntryRequestType, args)
	if err != nil {
		return err
	}
	if respErr, ok := resp.(error); ok {
		return respErr
	}
	return nil
}

// ResolveServiceConfig
func (c *ConfigEntry) ResolveServiceConfig(args *structs.ServiceConfigRequest, reply *structs.ServiceConfigResponse) error {
	if err := c.srv.validateEnterpriseRequest(&args.EnterpriseMeta, false); err != nil {
		return err
	}

	if done, err := c.srv.forward("ConfigEntry.ResolveServiceConfig", args, args, reply); done {
		return err
	}
	defer metrics.MeasureSince([]string{"config_entry", "resolve_service_config"}, time.Now())

	var authzContext acl.AuthorizerContext
	authz, err := c.srv.ResolveTokenAndDefaultMeta(args.Token, &args.EnterpriseMeta, &authzContext)
	if err != nil {
		return err
	}
	if authz != nil && authz.ServiceRead(args.Name, &authzContext) != acl.Allow {
		return acl.ErrPermissionDenied
	}

	return c.srv.blockingQuery(
		&args.QueryOptions,
		&reply.QueryMeta,
		func(ws memdb.WatchSet, state *state.Store) error {
			reply.Reset()

			reply.MeshGateway.Mode = structs.MeshGatewayModeDefault
			// Pass the WatchSet to both the service and proxy config lookups. If either is updated
			// during the blocking query, this function will be rerun and these state store lookups
			// will both be current.
			index, serviceEntry, err := state.ConfigEntry(ws, structs.ServiceDefaults, args.Name, &args.EnterpriseMeta)
			if err != nil {
				return err
			}
			var serviceConf *structs.ServiceConfigEntry
			var ok bool
			if serviceEntry != nil {
				serviceConf, ok = serviceEntry.(*structs.ServiceConfigEntry)
				if !ok {
					return fmt.Errorf("invalid service config type %T", serviceEntry)
				}
			}

			// Use the default enterprise meta to look up the global proxy defaults. In the future we may allow per-namespace proxy-defaults
			// but not yet.
			_, proxyEntry, err := state.ConfigEntry(ws, structs.ProxyDefaults, structs.ProxyConfigGlobal, structs.DefaultEnterpriseMeta())
			if err != nil {
				return err
			}
			var proxyConf *structs.ProxyConfigEntry
			if proxyEntry != nil {
				proxyConf, ok = proxyEntry.(*structs.ProxyConfigEntry)
				if !ok {
					return fmt.Errorf("invalid proxy config type %T", proxyEntry)
				}
				// Apply the proxy defaults to the sidecar's proxy config
				mapCopy, err := copystructure.Copy(proxyConf.Config)
				if err != nil {
					return fmt.Errorf("failed to copy global proxy-defaults: %v", err)
				}
				reply.ProxyConfig = mapCopy.(map[string]interface{})
				reply.MeshGateway = proxyConf.MeshGateway
				reply.Expose = proxyConf.Expose
			}

			reply.Index = index

			if serviceConf != nil {
				if serviceConf.Expose.Checks {
					reply.Expose.Checks = true
				}
				if len(serviceConf.Expose.Paths) >= 1 {
					reply.Expose.Paths = serviceConf.Expose.Paths
				}
				if serviceConf.MeshGateway.Mode != structs.MeshGatewayModeDefault {
					reply.MeshGateway.Mode = serviceConf.MeshGateway.Mode
				}
				if serviceConf.Protocol != "" {
					if reply.ProxyConfig == nil {
						reply.ProxyConfig = make(map[string]interface{})
					}
					reply.ProxyConfig["protocol"] = serviceConf.Protocol
				}
			}

			// Extract the global protocol from proxyConf for upstream configs.
			var proxyConfGlobalProtocol interface{}
			if proxyConf != nil && proxyConf.Config != nil {
				proxyConfGlobalProtocol = proxyConf.Config["protocol"]
			}

			// map the legacy request structure using only service names
			// to the new ServiceID type.
			upstreamIDs := args.UpstreamIDs
			legacyUpstreams := false

			if len(upstreamIDs) == 0 {
				legacyUpstreams = true

				upstreamIDs = make([]structs.ServiceID, 0)
				for _, upstream := range args.Upstreams {
					upstreamIDs = append(upstreamIDs, structs.NewServiceID(upstream, &args.EnterpriseMeta))
				}
			}

			usConfigs := make(map[structs.ServiceID]map[string]interface{})

			for _, upstream := range upstreamIDs {
				_, upstreamEntry, err := state.ConfigEntry(ws, structs.ServiceDefaults, upstream.ID, &upstream.EnterpriseMeta)
				if err != nil {
					return err
				}
				var upstreamConf *structs.ServiceConfigEntry
				var ok bool
				if upstreamEntry != nil {
					upstreamConf, ok = upstreamEntry.(*structs.ServiceConfigEntry)
					if !ok {
						return fmt.Errorf("invalid service config type %T", upstreamEntry)
					}
				}

				// No upstream found; skip.
				if upstreamConf == nil {
					continue
				}

				// Fallback to proxyConf global protocol.
				protocol := proxyConfGlobalProtocol
				if upstreamConf.Protocol != "" {
					protocol = upstreamConf.Protocol
				}

				// Nothing to configure if a protocol hasn't been set.
				if protocol == nil {
					continue
				}

				usConfigs[upstream] = map[string]interface{}{
					"protocol": protocol,
				}
			}

			// don't allocate the slices just to not fill them
			if len(usConfigs) == 0 {
				return nil
			}

			if legacyUpstreams {
				if reply.UpstreamConfigs == nil {
					reply.UpstreamConfigs = make(map[string]map[string]interface{})
				}
				for us, conf := range usConfigs {
					reply.UpstreamConfigs[us.ID] = conf
				}
			} else {
				if reply.UpstreamIDConfigs == nil {
					reply.UpstreamIDConfigs = make(structs.UpstreamConfigs, 0, len(usConfigs))
				}

				for us, conf := range usConfigs {
					reply.UpstreamIDConfigs = append(reply.UpstreamIDConfigs, structs.UpstreamConfig{Upstream: us, Config: conf})
				}
			}

			return nil
		})
}
<commit_msg>Fix a bug with ACL enforcement of reads on namespaced config entries. (#7239)<commit_after>package consul

import (
	"fmt"
	"time"

	metrics "github.com/armon/go-metrics"
	"github.com/hashicorp/consul/acl"
	"github.com/hashicorp/consul/agent/consul/state"
	"github.com/hashicorp/consul/agent/structs"
	memdb "github.com/hashicorp/go-memdb"
	"github.com/mitchellh/copystructure"
)

// The ConfigEntry endpoint is used to query centralized config information
type ConfigEntry struct {
	srv *Server
}

// Apply does an upsert of the given config entry.
func (c *ConfigEntry) Apply(args *structs.ConfigEntryRequest, reply *bool) error {
	if err := c.srv.validateEnterpriseRequest(args.Entry.GetEnterpriseMeta(), true); err != nil {
		return err
	}

	// Ensure that all config entry writes go to the primary datacenter. These will then
	// be replicated to all the other datacenters.
	args.Datacenter = c.srv.config.PrimaryDatacenter

	if done, err := c.srv.forward("ConfigEntry.Apply", args, args, reply); done {
		return err
	}
	defer metrics.MeasureSince([]string{"config_entry", "apply"}, time.Now())

	entMeta := args.Entry.GetEnterpriseMeta()
	authz, err := c.srv.ResolveTokenAndDefaultMeta(args.Token, entMeta, nil)
	if err != nil {
		return err
	}

	// Normalize and validate the incoming config entry.
	if err := args.Entry.Normalize(); err != nil {
		return err
	}
	if err := args.Entry.Validate(); err != nil {
		return err
	}

	if authz != nil && !args.Entry.CanWrite(authz) {
		return acl.ErrPermissionDenied
	}

	if args.Op != structs.ConfigEntryUpsert && args.Op != structs.ConfigEntryUpsertCAS {
		args.Op = structs.ConfigEntryUpsert
	}
	resp, err := c.srv.raftApply(structs.ConfigEntryRequestType, args)
	if err != nil {
		return err
	}
	if respErr, ok := resp.(error); ok {
		return respErr
	}
	if respBool, ok := resp.(bool); ok {
		*reply = respBool
	}

	return nil
}

// Get returns a single config entry by Kind/Name.
func (c *ConfigEntry) Get(args *structs.ConfigEntryQuery, reply *structs.ConfigEntryResponse) error {
	if err := c.srv.validateEnterpriseRequest(&args.EnterpriseMeta, false); err != nil {
		return err
	}

	if done, err := c.srv.forward("ConfigEntry.Get", args, args, reply); done {
		return err
	}
	defer metrics.MeasureSince([]string{"config_entry", "get"}, time.Now())

	authz, err := c.srv.ResolveTokenAndDefaultMeta(args.Token, &args.EnterpriseMeta, nil)
	if err != nil {
		return err
	}

	// Create a dummy config entry to check the ACL permissions.
	lookupEntry, err := structs.MakeConfigEntry(args.Kind, args.Name)
	if err != nil {
		return err
	}
	lookupEntry.GetEnterpriseMeta().Merge(&args.EnterpriseMeta)

	if authz != nil && !lookupEntry.CanRead(authz) {
		return acl.ErrPermissionDenied
	}

	return c.srv.blockingQuery(
		&args.QueryOptions,
		&reply.QueryMeta,
		func(ws memdb.WatchSet, state *state.Store) error {
			index, entry, err := state.ConfigEntry(ws, args.Kind, args.Name, &args.EnterpriseMeta)
			if err != nil {
				return err
			}

			reply.Index = index
			if entry == nil {
				return nil
			}

			reply.Entry = entry
			return nil
		})
}

// List returns all the config entries of the given kind. If Kind is blank,
// all existing config entries will be returned.
func (c *ConfigEntry) List(args *structs.ConfigEntryQuery, reply *structs.IndexedConfigEntries) error {
	if err := c.srv.validateEnterpriseRequest(&args.EnterpriseMeta, false); err != nil {
		return err
	}

	if done, err := c.srv.forward("ConfigEntry.List", args, args, reply); done {
		return err
	}
	defer metrics.MeasureSince([]string{"config_entry", "list"}, time.Now())

	authz, err := c.srv.ResolveTokenAndDefaultMeta(args.Token, &args.EnterpriseMeta, nil)
	if err != nil {
		return err
	}

	if args.Kind != "" && !structs.ValidateConfigEntryKind(args.Kind) {
		return fmt.Errorf("invalid config entry kind: %s", args.Kind)
	}

	return c.srv.blockingQuery(
		&args.QueryOptions,
		&reply.QueryMeta,
		func(ws memdb.WatchSet, state *state.Store) error {
			index, entries, err := state.ConfigEntriesByKind(ws, args.Kind, &args.EnterpriseMeta)
			if err != nil {
				return err
			}

			// Filter the entries returned by ACL permissions.
			filteredEntries := make([]structs.ConfigEntry, 0, len(entries))
			for _, entry := range entries {
				if authz != nil && !entry.CanRead(authz) {
					continue
				}
				filteredEntries = append(filteredEntries, entry)
			}

			reply.Kind = args.Kind
			reply.Index = index
			reply.Entries = filteredEntries
			return nil
		})
}

// ListAll returns all the known configuration entries
func (c *ConfigEntry) ListAll(args *structs.DCSpecificRequest, reply *structs.IndexedGenericConfigEntries) error {
	if err := c.srv.validateEnterpriseRequest(&args.EnterpriseMeta, false); err != nil {
		return err
	}

	if done, err := c.srv.forward("ConfigEntry.ListAll", args, args, reply); done {
		return err
	}
	defer metrics.MeasureSince([]string{"config_entry", "listAll"}, time.Now())

	authz, err := c.srv.ResolveTokenAndDefaultMeta(args.Token, &args.EnterpriseMeta, nil)
	if err != nil {
		return err
	}

	return c.srv.blockingQuery(
		&args.QueryOptions,
		&reply.QueryMeta,
		func(ws memdb.WatchSet, state *state.Store) error {
			index, entries, err := state.ConfigEntries(ws, &args.EnterpriseMeta)
			if err != nil {
				return err
			}

			// Filter the entries returned by ACL permissions.
			filteredEntries := make([]structs.ConfigEntry, 0, len(entries))
			for _, entry := range entries {
				if authz != nil && !entry.CanRead(authz) {
					continue
				}
				filteredEntries = append(filteredEntries, entry)
			}

			reply.Entries = filteredEntries
			reply.Index = index
			return nil
		})
}

// Delete deletes a config entry.
func (c *ConfigEntry) Delete(args *structs.ConfigEntryRequest, reply *struct{}) error {
	if err := c.srv.validateEnterpriseRequest(args.Entry.GetEnterpriseMeta(), true); err != nil {
		return err
	}

	// Ensure that all config entry writes go to the primary datacenter. These will then
	// be replicated to all the other datacenters.
	args.Datacenter = c.srv.config.PrimaryDatacenter

	if done, err := c.srv.forward("ConfigEntry.Delete", args, args, reply); done {
		return err
	}
	defer metrics.MeasureSince([]string{"config_entry", "delete"}, time.Now())

	authz, err := c.srv.ResolveTokenAndDefaultMeta(args.Token, args.Entry.GetEnterpriseMeta(), nil)
	if err != nil {
		return err
	}

	// Normalize the incoming entry.
	if err := args.Entry.Normalize(); err != nil {
		return err
	}

	if authz != nil && !args.Entry.CanWrite(authz) {
		return acl.ErrPermissionDenied
	}

	args.Op = structs.ConfigEntryDelete
	resp, err := c.srv.raftApply(structs.ConfigEntryRequestType, args)
	if err != nil {
		return err
	}
	if respErr, ok := resp.(error); ok {
		return respErr
	}
	return nil
}

// ResolveServiceConfig
func (c *ConfigEntry) ResolveServiceConfig(args *structs.ServiceConfigRequest, reply *structs.ServiceConfigResponse) error {
	if err := c.srv.validateEnterpriseRequest(&args.EnterpriseMeta, false); err != nil {
		return err
	}

	if done, err := c.srv.forward("ConfigEntry.ResolveServiceConfig", args, args, reply); done {
		return err
	}
	defer metrics.MeasureSince([]string{"config_entry", "resolve_service_config"}, time.Now())

	var authzContext acl.AuthorizerContext
	authz, err := c.srv.ResolveTokenAndDefaultMeta(args.Token, &args.EnterpriseMeta, &authzContext)
	if err != nil {
		return err
	}
	if authz != nil && authz.ServiceRead(args.Name, &authzContext) != acl.Allow {
		return acl.ErrPermissionDenied
	}

	return c.srv.blockingQuery(
		&args.QueryOptions,
		&reply.QueryMeta,
		func(ws memdb.WatchSet, state *state.Store) error {
			reply.Reset()

			reply.MeshGateway.Mode = structs.MeshGatewayModeDefault
			// Pass the WatchSet to both the service and proxy config lookups. If either is updated
			// during the blocking query, this function will be rerun and these state store lookups
			// will both be current.
			index, serviceEntry, err := state.ConfigEntry(ws, structs.ServiceDefaults, args.Name, &args.EnterpriseMeta)
			if err != nil {
				return err
			}
			var serviceConf *structs.ServiceConfigEntry
			var ok bool
			if serviceEntry != nil {
				serviceConf, ok = serviceEntry.(*structs.ServiceConfigEntry)
				if !ok {
					return fmt.Errorf("invalid service config type %T", serviceEntry)
				}
			}

			// Use the default enterprise meta to look up the global proxy defaults. In the future we may allow per-namespace proxy-defaults
			// but not yet.
			_, proxyEntry, err := state.ConfigEntry(ws, structs.ProxyDefaults, structs.ProxyConfigGlobal, structs.DefaultEnterpriseMeta())
			if err != nil {
				return err
			}
			var proxyConf *structs.ProxyConfigEntry
			if proxyEntry != nil {
				proxyConf, ok = proxyEntry.(*structs.ProxyConfigEntry)
				if !ok {
					return fmt.Errorf("invalid proxy config type %T", proxyEntry)
				}
				// Apply the proxy defaults to the sidecar's proxy config
				mapCopy, err := copystructure.Copy(proxyConf.Config)
				if err != nil {
					return fmt.Errorf("failed to copy global proxy-defaults: %v", err)
				}
				reply.ProxyConfig = mapCopy.(map[string]interface{})
				reply.MeshGateway = proxyConf.MeshGateway
				reply.Expose = proxyConf.Expose
			}

			reply.Index = index

			if serviceConf != nil {
				if serviceConf.Expose.Checks {
					reply.Expose.Checks = true
				}
				if len(serviceConf.Expose.Paths) >= 1 {
					reply.Expose.Paths = serviceConf.Expose.Paths
				}
				if serviceConf.MeshGateway.Mode != structs.MeshGatewayModeDefault {
					reply.MeshGateway.Mode = serviceConf.MeshGateway.Mode
				}
				if serviceConf.Protocol != "" {
					if reply.ProxyConfig == nil {
						reply.ProxyConfig = make(map[string]interface{})
					}
					reply.ProxyConfig["protocol"] = serviceConf.Protocol
				}
			}

			// Extract the global protocol from proxyConf for upstream configs.
			var proxyConfGlobalProtocol interface{}
			if proxyConf != nil && proxyConf.Config != nil {
				proxyConfGlobalProtocol = proxyConf.Config["protocol"]
			}

			// map the legacy request structure using only service names
			// to the new ServiceID type.
			upstreamIDs := args.UpstreamIDs
			legacyUpstreams := false

			if len(upstreamIDs) == 0 {
				legacyUpstreams = true

				upstreamIDs = make([]structs.ServiceID, 0)
				for _, upstream := range args.Upstreams {
					upstreamIDs = append(upstreamIDs, structs.NewServiceID(upstream, &args.EnterpriseMeta))
				}
			}

			usConfigs := make(map[structs.ServiceID]map[string]interface{})

			for _, upstream := range upstreamIDs {
				_, upstreamEntry, err := state.ConfigEntry(ws, structs.ServiceDefaults, upstream.ID, &upstream.EnterpriseMeta)
				if err != nil {
					return err
				}
				var upstreamConf *structs.ServiceConfigEntry
				var ok bool
				if upstreamEntry != nil {
					upstreamConf, ok = upstreamEntry.(*structs.ServiceConfigEntry)
					if !ok {
						return fmt.Errorf("invalid service config type %T", upstreamEntry)
					}
				}

				// No upstream found; skip.
				if upstreamConf == nil {
					continue
				}

				// Fallback to proxyConf global protocol.
				protocol := proxyConfGlobalProtocol
				if upstreamConf.Protocol != "" {
					protocol = upstreamConf.Protocol
				}

				// Nothing to configure if a protocol hasn't been set.
				if protocol == nil {
					continue
				}

				usConfigs[upstream] = map[string]interface{}{
					"protocol": protocol,
				}
			}

			// don't allocate the slices just to not fill them
			if len(usConfigs) == 0 {
				return nil
			}

			if legacyUpstreams {
				if reply.UpstreamConfigs == nil {
					reply.UpstreamConfigs = make(map[string]map[string]interface{})
				}
				for us, conf := range usConfigs {
					reply.UpstreamConfigs[us.ID] = conf
				}
			} else {
				if reply.UpstreamIDConfigs == nil {
					reply.UpstreamIDConfigs = make(structs.UpstreamConfigs, 0, len(usConfigs))
				}

				for us, conf := range usConfigs {
					reply.UpstreamIDConfigs = append(reply.UpstreamIDConfigs, structs.UpstreamConfig{Upstream: us, Config: conf})
				}
			}

			return nil
		})
}
<|endoftext|>


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

package oauth2

import (
	"context"
	"time"

	"github.com/pkg/errors"

	"github.com/ory/fosite"
)

type RefreshTokenGrantHandler struct {
	AccessTokenStrategy    AccessTokenStrategy
	RefreshTokenStrategy   RefreshTokenStrategy
	TokenRevocationStorage TokenRevocationStorage

	// AccessTokenLifespan defines the lifetime of an access token.
	AccessTokenLifespan time.Duration

	// RefreshTokenLifespan defines the lifetime of a refresh token.
	RefreshTokenLifespan time.Duration

	ScopeStrategy            fosite.ScopeStrategy
	AudienceMatchingStrategy fosite.AudienceMatchingStrategy
}

// HandleTokenEndpointRequest implements https://tools.ietf.org/html/rfc6749#section-6
func (c *RefreshTokenGrantHandler) HandleTokenEndpointRequest(ctx context.Context, request fosite.AccessRequester) error {
	// grant_type REQUIRED.
	// Value MUST be set to "refresh_token".
	if !request.GetGrantTypes().Exact("refresh_token") {
		return errors.WithStack(fosite.ErrUnknownRequest)
	}

	if !request.GetClient().GetGrantTypes().Has("refresh_token") {
		return errors.WithStack(fosite.ErrInvalidGrant.WithHint("The OAuth 2.0 Client is not allowed to use authorization grant \"refresh_token\"."))
	}

	refresh := request.GetRequestForm().Get("refresh_token")
	signature := c.RefreshTokenStrategy.RefreshTokenSignature(refresh)
	originalRequest, err := c.TokenRevocationStorage.GetRefreshTokenSession(ctx, signature, request.GetSession())
	if errors.Cause(err) == fosite.ErrNotFound {
		return errors.WithStack(fosite.ErrInvalidRequest.WithDebug(err.Error()))
	} else if err != nil {
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	} else if err := c.RefreshTokenStrategy.ValidateRefreshToken(ctx, originalRequest, refresh); err != nil {
		// The authorization server MUST ... validate the refresh token.
		// This needs to happen after store retrieval for the session to be hydrated properly
		return errors.WithStack(fosite.ErrInvalidRequest.WithDebug(err.Error()))
	}

	if !originalRequest.GetGrantedScopes().HasOneOf("offline", "offline_access") {
		return errors.WithStack(fosite.ErrScopeNotGranted.WithHint("The OAuth 2.0 Client was not granted scope \"offline\" or \"offline_access\" and may thus not perform the \"refresh_token\" authorization grant."))

	}

	// The authorization server MUST ... and ensure that the refresh token was issued to the authenticated client
	if originalRequest.GetClient().GetID() != request.GetClient().GetID() {
		return errors.WithStack(fosite.ErrInvalidRequest.WithHint("The OAuth 2.0 Client ID from this request does not match the ID during the initial token issuance."))
	}

	request.SetSession(originalRequest.GetSession().Clone())
	request.SetRequestedScopes(originalRequest.GetRequestedScopes())
	request.SetRequestedAudience(originalRequest.GetRequestedAudience())

	for _, scope := range originalRequest.GetGrantedScopes() {
		if !c.ScopeStrategy(request.GetClient().GetScopes(), scope) {
			return errors.WithStack(fosite.ErrInvalidScope.WithHintf("The OAuth 2.0 Client is not allowed to request scope \"%s\".", scope))
		}
		request.GrantScope(scope)
	}

	if err := c.AudienceMatchingStrategy(request.GetClient().GetAudience(), originalRequest.GetGrantedAudience()); err != nil {
		return err
	}

	for _, audience := range originalRequest.GetGrantedAudience() {
		request.GrantAudience(audience)
	}

	request.GetSession().SetExpiresAt(fosite.AccessToken, time.Now().UTC().Add(c.AccessTokenLifespan).Round(time.Second))
	if c.RefreshTokenLifespan > -1 {
		request.GetSession().SetExpiresAt(fosite.RefreshToken, time.Now().UTC().Add(c.RefreshTokenLifespan).Round(time.Second))
	}

	return nil
}

// PopulateTokenEndpointResponse implements https://tools.ietf.org/html/rfc6749#section-6
func (c *RefreshTokenGrantHandler) PopulateTokenEndpointResponse(ctx context.Context, requester fosite.AccessRequester, responder fosite.AccessResponder) error {
	if !requester.GetGrantTypes().Exact("refresh_token") {
		return errors.WithStack(fosite.ErrUnknownRequest)
	}

	accessToken, accessSignature, err := c.AccessTokenStrategy.GenerateAccessToken(ctx, requester)
	if err != nil {
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	}

	refreshToken, refreshSignature, err := c.RefreshTokenStrategy.GenerateRefreshToken(ctx, requester)
	if err != nil {
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	}

	signature := c.RefreshTokenStrategy.RefreshTokenSignature(requester.GetRequestForm().Get("refresh_token"))
	ts, err := c.TokenRevocationStorage.GetRefreshTokenSession(ctx, signature, nil)
	if err != nil {
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	} else if err := c.TokenRevocationStorage.RevokeAccessToken(ctx, ts.GetID()); err != nil {
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	} else if err := c.TokenRevocationStorage.RevokeRefreshToken(ctx, ts.GetID()); err != nil {
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	}

	storeReq := requester.Sanitize([]string{})
	storeReq.SetID(ts.GetID())
	if err := c.TokenRevocationStorage.CreateAccessTokenSession(ctx, accessSignature, storeReq); err != nil {
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	} else if err := c.TokenRevocationStorage.CreateRefreshTokenSession(ctx, refreshSignature, storeReq); err != nil {
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	}

	responder.SetAccessToken(accessToken)
	responder.SetTokenType("bearer")
	responder.SetExpiresIn(getExpiresIn(requester, fosite.AccessToken, c.AccessTokenLifespan, time.Now().UTC()))
	responder.SetScopes(requester.GetGrantedScopes())
	responder.SetExtra("refresh_token", refreshToken)
	return nil
}
<commit_msg>oauth2: use transactions in the refresh token flow (if the storage implementation implements the `Transactional` interface) to address #309<commit_after>/*
 * Copyright © 2015-2018 Aeneas Rekkas <aeneas+oss@aeneas.io>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * @author		Aeneas Rekkas <aeneas+oss@aeneas.io>
 * @copyright 	2015-2018 Aeneas Rekkas <aeneas+oss@aeneas.io>
 * @license 	Apache-2.0
 *
 */

package oauth2

import (
	"context"
	"time"

	"github.com/ory/fosite/storage"

	"github.com/pkg/errors"

	"github.com/ory/fosite"
)

type RefreshTokenGrantHandler struct {
	AccessTokenStrategy    AccessTokenStrategy
	RefreshTokenStrategy   RefreshTokenStrategy
	TokenRevocationStorage TokenRevocationStorage

	// AccessTokenLifespan defines the lifetime of an access token.
	AccessTokenLifespan time.Duration

	// RefreshTokenLifespan defines the lifetime of a refresh token.
	RefreshTokenLifespan time.Duration

	ScopeStrategy            fosite.ScopeStrategy
	AudienceMatchingStrategy fosite.AudienceMatchingStrategy
}

// HandleTokenEndpointRequest implements https://tools.ietf.org/html/rfc6749#section-6
func (c *RefreshTokenGrantHandler) HandleTokenEndpointRequest(ctx context.Context, request fosite.AccessRequester) error {
	// grant_type REQUIRED.
	// Value MUST be set to "refresh_token".
	if !request.GetGrantTypes().Exact("refresh_token") {
		return errors.WithStack(fosite.ErrUnknownRequest)
	}

	if !request.GetClient().GetGrantTypes().Has("refresh_token") {
		return errors.WithStack(fosite.ErrInvalidGrant.WithHint("The OAuth 2.0 Client is not allowed to use authorization grant \"refresh_token\"."))
	}

	refresh := request.GetRequestForm().Get("refresh_token")
	signature := c.RefreshTokenStrategy.RefreshTokenSignature(refresh)
	originalRequest, err := c.TokenRevocationStorage.GetRefreshTokenSession(ctx, signature, request.GetSession())
	if errors.Cause(err) == fosite.ErrNotFound {
		return errors.WithStack(fosite.ErrInvalidRequest.WithDebug(err.Error()))
	} else if err != nil {
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	} else if err := c.RefreshTokenStrategy.ValidateRefreshToken(ctx, originalRequest, refresh); err != nil {
		// The authorization server MUST ... validate the refresh token.
		// This needs to happen after store retrieval for the session to be hydrated properly
		return errors.WithStack(fosite.ErrInvalidRequest.WithDebug(err.Error()))
	}

	if !originalRequest.GetGrantedScopes().HasOneOf("offline", "offline_access") {
		return errors.WithStack(fosite.ErrScopeNotGranted.WithHint("The OAuth 2.0 Client was not granted scope \"offline\" or \"offline_access\" and may thus not perform the \"refresh_token\" authorization grant."))

	}

	// The authorization server MUST ... and ensure that the refresh token was issued to the authenticated client
	if originalRequest.GetClient().GetID() != request.GetClient().GetID() {
		return errors.WithStack(fosite.ErrInvalidRequest.WithHint("The OAuth 2.0 Client ID from this request does not match the ID during the initial token issuance."))
	}

	request.SetSession(originalRequest.GetSession().Clone())
	request.SetRequestedScopes(originalRequest.GetRequestedScopes())
	request.SetRequestedAudience(originalRequest.GetRequestedAudience())

	for _, scope := range originalRequest.GetGrantedScopes() {
		if !c.ScopeStrategy(request.GetClient().GetScopes(), scope) {
			return errors.WithStack(fosite.ErrInvalidScope.WithHintf("The OAuth 2.0 Client is not allowed to request scope \"%s\".", scope))
		}
		request.GrantScope(scope)
	}

	if err := c.AudienceMatchingStrategy(request.GetClient().GetAudience(), originalRequest.GetGrantedAudience()); err != nil {
		return err
	}

	for _, audience := range originalRequest.GetGrantedAudience() {
		request.GrantAudience(audience)
	}

	request.GetSession().SetExpiresAt(fosite.AccessToken, time.Now().UTC().Add(c.AccessTokenLifespan).Round(time.Second))
	if c.RefreshTokenLifespan > -1 {
		request.GetSession().SetExpiresAt(fosite.RefreshToken, time.Now().UTC().Add(c.RefreshTokenLifespan).Round(time.Second))
	}

	return nil
}

// PopulateTokenEndpointResponse implements https://tools.ietf.org/html/rfc6749#section-6
func (c *RefreshTokenGrantHandler) PopulateTokenEndpointResponse(ctx context.Context, requester fosite.AccessRequester, responder fosite.AccessResponder) error {
	if !requester.GetGrantTypes().Exact("refresh_token") {
		return errors.WithStack(fosite.ErrUnknownRequest)
	}

	accessToken, accessSignature, err := c.AccessTokenStrategy.GenerateAccessToken(ctx, requester)
	if err != nil {
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	}

	refreshToken, refreshSignature, err := c.RefreshTokenStrategy.GenerateRefreshToken(ctx, requester)
	if err != nil {
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	}

	signature := c.RefreshTokenStrategy.RefreshTokenSignature(requester.GetRequestForm().Get("refresh_token"))

	ctx, err = storage.MaybeBeginTx(ctx, c.TokenRevocationStorage)
	if err != nil {
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	}

	ts, err := c.TokenRevocationStorage.GetRefreshTokenSession(ctx, signature, nil)
	if err != nil {
		if rollBackTxnErr := storage.MaybeRollbackTx(ctx, c.TokenRevocationStorage); rollBackTxnErr != nil {
			err = rollBackTxnErr
		}
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	} else if err := c.TokenRevocationStorage.RevokeAccessToken(ctx, ts.GetID()); err != nil {
		if rollBackTxnErr := storage.MaybeRollbackTx(ctx, c.TokenRevocationStorage); rollBackTxnErr != nil {
			err = rollBackTxnErr
		}
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	} else if err := c.TokenRevocationStorage.RevokeRefreshToken(ctx, ts.GetID()); err != nil {
		if rollBackTxnErr := storage.MaybeRollbackTx(ctx, c.TokenRevocationStorage); rollBackTxnErr != nil {
			err = rollBackTxnErr
		}
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	}

	storeReq := requester.Sanitize([]string{})
	storeReq.SetID(ts.GetID())
	if err := c.TokenRevocationStorage.CreateAccessTokenSession(ctx, accessSignature, storeReq); err != nil {
		if rollBackTxnErr := storage.MaybeRollbackTx(ctx, c.TokenRevocationStorage); rollBackTxnErr != nil {
			err = rollBackTxnErr
		}
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	} else if err := c.TokenRevocationStorage.CreateRefreshTokenSession(ctx, refreshSignature, storeReq); err != nil {
		if rollBackTxnErr := storage.MaybeRollbackTx(ctx, c.TokenRevocationStorage); rollBackTxnErr != nil {
			err = rollBackTxnErr
		}
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	}

	responder.SetAccessToken(accessToken)
	responder.SetTokenType("bearer")
	responder.SetExpiresIn(getExpiresIn(requester, fosite.AccessToken, c.AccessTokenLifespan, time.Now().UTC()))
	responder.SetScopes(requester.GetGrantedScopes())
	responder.SetExtra("refresh_token", refreshToken)

	if err := storage.MaybeCommitTx(ctx, c.TokenRevocationStorage); err != nil {
		return errors.WithStack(fosite.ErrServerError.WithDebug(err.Error()))
	}

	return nil
}
<|endoftext|>


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

package importer

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"time"

	daisyUtils "github.com/GoogleCloudPlatform/compute-image-tools/cli_tools/common/utils/daisy"
	"github.com/GoogleCloudPlatform/compute-image-tools/cli_tools/common/utils/storage"
	"github.com/GoogleCloudPlatform/compute-image-tools/cli_tools/daisycommon"
	"github.com/GoogleCloudPlatform/compute-image-tools/daisy"
	daisyCompute "github.com/GoogleCloudPlatform/compute-image-tools/daisy/compute"
	computeBeta "google.golang.org/api/compute/v0.beta"
	"google.golang.org/api/compute/v1"
	"google.golang.org/api/googleapi"
)

func isCausedByUnsupportedFormat(err error) bool {
	if err == nil {
		return false
	}
	return strings.Contains(err.Error(), "INVALID_IMAGE_FILE")
}

func isCausedByAlphaAPIAccess(err error) bool {
	if err == nil {
		return false
	}
	return strings.Contains(err.Error(), "Required 'Alpha Access' permission")
}

// apiInflater implements `importer.inflater` using the Compute Engine API
type apiInflater struct {
	serialLogs      []string
	args            ImportArguments
	computeClient   daisyCompute.Client
	storageClient   storage.Client
	guestOsFeatures []*computeBeta.GuestOsFeature
	wg              sync.WaitGroup
	cancelChan      chan string
}

func createAPIInflater(args ImportArguments, computeClient daisyCompute.Client, storageClient storage.Client) Inflater {
	inflater := apiInflater{
		serialLogs:    []string{},
		args:          args,
		computeClient: computeClient,
		storageClient: storageClient,
		cancelChan:    make(chan string, 1),
	}
	if args.UefiCompatible {
		inflater.guestOsFeatures = []*computeBeta.GuestOsFeature{{Type: "UEFI_COMPATIBLE"}}
	}
	return &inflater
}

func (inflater *apiInflater) TraceLogs() []string {
	return inflater.serialLogs
}

func (inflater *apiInflater) addTraceLog(l string) {
	inflater.serialLogs = append(inflater.serialLogs, l)
}

func (inflater *apiInflater) Inflate() (persistentDisk, shadowTestFields, error) {
	inflater.wg.Add(1)
	defer inflater.wg.Done()

	ctx := context.Background()
	startTime := time.Now()
	diskName := inflater.getShadowDiskName()

	// Create shadow disk
	cd := computeBeta.Disk{
		Name:                diskName,
		SourceStorageObject: inflater.args.SourceFile,
		GuestOsFeatures:     inflater.guestOsFeatures,
	}

	err := inflater.computeClient.CreateDiskBeta(inflater.args.Project, inflater.args.Zone, &cd)
	if err != nil {
		return persistentDisk{}, shadowTestFields{}, daisy.Errf("Failed to create shadow disk: %v", err)
	}

	// Cleanup the shadow disk ignoring error
	defer inflater.computeClient.DeleteDisk(inflater.args.Project, inflater.args.Zone, cd.Name)

	// If received a cancel signal from cancel(), then return early. Otherwise, it will waste
	// 2 min+ on calculateChecksum().
	select {
	case <-inflater.cancelChan:
		return persistentDisk{}, shadowTestFields{}, nil
	default:
	}

	// Prepare return value
	bkt, objPath, err := storage.GetGCSObjectPathElements(inflater.args.SourceFile)
	if err != nil {
		return persistentDisk{}, shadowTestFields{}, err
	}
	sourceFile := inflater.storageClient.GetObject(bkt, objPath).GetObjectHandle()
	attrs, err := sourceFile.Attrs(ctx)
	if err != nil {
		return persistentDisk{}, shadowTestFields{}, daisy.Errf("Failed to get source file attributes: %v", err)
	}
	sourceFileSizeGb := (attrs.Size-1)/1073741824 + 1

	diskURI := fmt.Sprintf("zones/%s/disks/%s", inflater.args.Zone, diskName)
	pd := persistentDisk{
		uri:        diskURI,
		sizeGb:     cd.SizeGb,
		sourceGb:   sourceFileSizeGb,
		sourceType: "vmdk", // only vmdk is supported right now
	}
	ii := shadowTestFields{
		inflationType: "api",
		inflationTime: time.Since(startTime),
	}

	// Calculate checksum by daisy workflow
	inflater.addTraceLog("Started checksum calculation.")
	ii.checksum, err = inflater.calculateChecksum(ctx, diskURI)
	return pd, ii, daisy.Errf("Failed to calculate checksum: %v", err)
}

func (inflater *apiInflater) getShadowDiskName() string {
	return fmt.Sprintf("shadow-disk-%v", inflater.args.ExecutionID)
}

func (inflater *apiInflater) Cancel(reason string) bool {
	// Send cancel signal
	inflater.cancelChan <- reason

	// Wait for inflate() to finish. Otherwise, the whole program might exit
	// before DeleteDisk() was executed.
	inflater.wg.Wait()

	// Expect 404 error to ensure shadow disk has been cleaned up.
	_, err := inflater.computeClient.GetDisk(inflater.args.Project, inflater.args.Zone, inflater.getShadowDiskName())
	if apiErr, ok := err.(*googleapi.Error); !ok || apiErr.Code != 404 {
		if err == nil {
			inflater.addTraceLog(fmt.Sprintf("apiInflater.inflate is canceled, cleanup is failed: %v", reason))
		} else {
			inflater.addTraceLog(fmt.Sprintf("apiInflater.inflate is canceled, cleanup failed to verify: %v", reason))
		}
		return false
	}
	inflater.addTraceLog(fmt.Sprintf("apiInflater.inflate is canceled: %v", reason))
	return true
}

// run a workflow to calculate checksum
func (inflater *apiInflater) calculateChecksum(ctx context.Context, diskURI string) (string, error) {
	w := daisy.New()
	w.Name = "shadow-disk-checksum"
	checksumScript := checksumScriptConst
	w.Steps = map[string]*daisy.Step{
		"create-disks": {
			CreateDisks: &daisy.CreateDisks{
				{
					Disk: compute.Disk{
						Name:        "disk-${NAME}",
						SourceImage: "projects/compute-image-tools/global/images/family/debian-9-worker",
						Type:        "pd-ssd",
					},
					FallbackToPdStandard: true,
				},
			},
		},
		"create-instance": {
			CreateInstances: &daisy.CreateInstances{
				Instances: []*daisy.Instance{
					{
						Instance: compute.Instance{
							Name: "inst-${NAME}",
							Disks: []*compute.AttachedDisk{
								{Source: "disk-${NAME}"},
								{Source: diskURI, Mode: "READ_ONLY"},
							},
							MachineType: "n1-highcpu-4",
							Metadata: &compute.Metadata{
								Items: []*compute.MetadataItems{
									{
										Key:   "startup-script",
										Value: &checksumScript,
									},
								},
							},
							NetworkInterfaces: []*compute.NetworkInterface{
								{
									AccessConfigs: []*compute.AccessConfig{},
								},
							},
						},
					},
				},
			},
		},
		"wait-for-checksum": {
			WaitForInstancesSignal: &daisy.WaitForInstancesSignal{
				{
					Name: "inst-${NAME}",
					SerialOutput: &daisy.SerialOutput{
						Port:         1,
						SuccessMatch: "Checksum calculated.",
						StatusMatch:  "Checksum:",
					},
				},
			},
		},
	}
	w.Dependencies = map[string][]string{
		"create-instance":   {"create-disks"},
		"wait-for-checksum": {"create-instance"},
	}

	// Calculate checksum within 20min.
	daisycommon.SetWorkflowAttributes(w, daisycommon.WorkflowAttributes{
		Project:           inflater.args.Project,
		Zone:              inflater.args.Zone,
		GCSPath:           inflater.args.ScratchBucketGcsPath,
		OAuth:             inflater.args.Oauth,
		Timeout:           "20m",
		ComputeEndpoint:   inflater.args.ComputeEndpoint,
		DisableGCSLogs:    inflater.args.GcsLogsDisabled,
		DisableCloudLogs:  inflater.args.CloudLogsDisabled,
		DisableStdoutLogs: inflater.args.StdoutLogsDisabled,
	})
	err := daisyUtils.RunWorkflowWithCancelSignal(ctx, w)
	if err != nil {
		return "", err
	}
	return w.GetSerialConsoleOutputValue("disk-checksum"), nil
}

// Dup logic in import_image.sh. If change anything here, please change in both places.
const (
	checksumScriptConst = `
		function serialOutputKeyValuePair() {
			echo "<serial-output key:'$1' value:'$2'>"
		}
		CHECK_DEVICE=sdb
		BLOCK_COUNT=$(cat /sys/class/block/$CHECK_DEVICE/size)

		# Check size = 200000*512 = 100MB
		CHECK_COUNT=200000
		CHECKSUM1=$(sudo dd if=/dev/$CHECK_DEVICE ibs=512 skip=0 count=$CHECK_COUNT | md5sum)
		CHECKSUM2=$(sudo dd if=/dev/$CHECK_DEVICE ibs=512 skip=$(( 2000000 - $CHECK_COUNT )) count=$CHECK_COUNT | md5sum)
		CHECKSUM3=$(sudo dd if=/dev/$CHECK_DEVICE ibs=512 skip=$(( 20000000 - $CHECK_COUNT )) count=$CHECK_COUNT | md5sum)
		CHECKSUM4=$(sudo dd if=/dev/$CHECK_DEVICE ibs=512 skip=$(( $BLOCK_COUNT - $CHECK_COUNT )) count=$CHECK_COUNT | md5sum)
		echo "Checksum: $(serialOutputKeyValuePair "disk-checksum" "$CHECKSUM1-$CHECKSUM2-$CHECKSUM3-$CHECKSUM4")"
		echo "Checksum calculated."`
)
<commit_msg>[VMDK shadow test] minor: fix error construction logic (#1435)<commit_after>//  Copyright 2020  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

package importer

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"time"

	daisyUtils "github.com/GoogleCloudPlatform/compute-image-tools/cli_tools/common/utils/daisy"
	"github.com/GoogleCloudPlatform/compute-image-tools/cli_tools/common/utils/storage"
	"github.com/GoogleCloudPlatform/compute-image-tools/cli_tools/daisycommon"
	"github.com/GoogleCloudPlatform/compute-image-tools/daisy"
	daisyCompute "github.com/GoogleCloudPlatform/compute-image-tools/daisy/compute"
	computeBeta "google.golang.org/api/compute/v0.beta"
	"google.golang.org/api/compute/v1"
	"google.golang.org/api/googleapi"
)

func isCausedByUnsupportedFormat(err error) bool {
	if err == nil {
		return false
	}
	return strings.Contains(err.Error(), "INVALID_IMAGE_FILE")
}

func isCausedByAlphaAPIAccess(err error) bool {
	if err == nil {
		return false
	}
	return strings.Contains(err.Error(), "Required 'Alpha Access' permission")
}

// apiInflater implements `importer.inflater` using the Compute Engine API
type apiInflater struct {
	serialLogs      []string
	args            ImportArguments
	computeClient   daisyCompute.Client
	storageClient   storage.Client
	guestOsFeatures []*computeBeta.GuestOsFeature
	wg              sync.WaitGroup
	cancelChan      chan string
}

func createAPIInflater(args ImportArguments, computeClient daisyCompute.Client, storageClient storage.Client) Inflater {
	inflater := apiInflater{
		serialLogs:    []string{},
		args:          args,
		computeClient: computeClient,
		storageClient: storageClient,
		cancelChan:    make(chan string, 1),
	}
	if args.UefiCompatible {
		inflater.guestOsFeatures = []*computeBeta.GuestOsFeature{{Type: "UEFI_COMPATIBLE"}}
	}
	return &inflater
}

func (inflater *apiInflater) TraceLogs() []string {
	return inflater.serialLogs
}

func (inflater *apiInflater) addTraceLog(l string) {
	inflater.serialLogs = append(inflater.serialLogs, l)
}

func (inflater *apiInflater) Inflate() (persistentDisk, shadowTestFields, error) {
	inflater.wg.Add(1)
	defer inflater.wg.Done()

	ctx := context.Background()
	startTime := time.Now()
	diskName := inflater.getShadowDiskName()

	// Create shadow disk
	cd := computeBeta.Disk{
		Name:                diskName,
		SourceStorageObject: inflater.args.SourceFile,
		GuestOsFeatures:     inflater.guestOsFeatures,
	}

	err := inflater.computeClient.CreateDiskBeta(inflater.args.Project, inflater.args.Zone, &cd)
	if err != nil {
		return persistentDisk{}, shadowTestFields{}, daisy.Errf("Failed to create shadow disk: %v", err)
	}

	// Cleanup the shadow disk ignoring error
	defer inflater.computeClient.DeleteDisk(inflater.args.Project, inflater.args.Zone, cd.Name)

	// If received a cancel signal from cancel(), then return early. Otherwise, it will waste
	// 2 min+ on calculateChecksum().
	select {
	case <-inflater.cancelChan:
		return persistentDisk{}, shadowTestFields{}, nil
	default:
	}

	// Prepare return value
	bkt, objPath, err := storage.GetGCSObjectPathElements(inflater.args.SourceFile)
	if err != nil {
		return persistentDisk{}, shadowTestFields{}, err
	}
	sourceFile := inflater.storageClient.GetObject(bkt, objPath).GetObjectHandle()
	attrs, err := sourceFile.Attrs(ctx)
	if err != nil {
		return persistentDisk{}, shadowTestFields{}, daisy.Errf("Failed to get source file attributes: %v", err)
	}
	sourceFileSizeGb := (attrs.Size-1)/1073741824 + 1

	diskURI := fmt.Sprintf("zones/%s/disks/%s", inflater.args.Zone, diskName)
	pd := persistentDisk{
		uri:        diskURI,
		sizeGb:     cd.SizeGb,
		sourceGb:   sourceFileSizeGb,
		sourceType: "vmdk", // only vmdk is supported right now
	}
	ii := shadowTestFields{
		inflationType: "api",
		inflationTime: time.Since(startTime),
	}

	// Calculate checksum by daisy workflow
	inflater.addTraceLog("Started checksum calculation.")
	ii.checksum, err = inflater.calculateChecksum(ctx, diskURI)
	if err != nil {
		err = daisy.Errf("Failed to calculate checksum: %v", err)
	}
	return pd, ii, err
}

func (inflater *apiInflater) getShadowDiskName() string {
	return fmt.Sprintf("shadow-disk-%v", inflater.args.ExecutionID)
}

func (inflater *apiInflater) Cancel(reason string) bool {
	// Send cancel signal
	inflater.cancelChan <- reason

	// Wait for inflate() to finish. Otherwise, the whole program might exit
	// before DeleteDisk() was executed.
	inflater.wg.Wait()

	// Expect 404 error to ensure shadow disk has been cleaned up.
	_, err := inflater.computeClient.GetDisk(inflater.args.Project, inflater.args.Zone, inflater.getShadowDiskName())
	if apiErr, ok := err.(*googleapi.Error); !ok || apiErr.Code != 404 {
		if err == nil {
			inflater.addTraceLog(fmt.Sprintf("apiInflater.inflate is canceled, cleanup is failed: %v", reason))
		} else {
			inflater.addTraceLog(fmt.Sprintf("apiInflater.inflate is canceled, cleanup failed to verify: %v", reason))
		}
		return false
	}
	inflater.addTraceLog(fmt.Sprintf("apiInflater.inflate is canceled: %v", reason))
	return true
}

// run a workflow to calculate checksum
func (inflater *apiInflater) calculateChecksum(ctx context.Context, diskURI string) (string, error) {
	w := daisy.New()
	w.Name = "shadow-disk-checksum"
	checksumScript := checksumScriptConst
	w.Steps = map[string]*daisy.Step{
		"create-disks": {
			CreateDisks: &daisy.CreateDisks{
				{
					Disk: compute.Disk{
						Name:        "disk-${NAME}",
						SourceImage: "projects/compute-image-tools/global/images/family/debian-9-worker",
						Type:        "pd-ssd",
					},
					FallbackToPdStandard: true,
				},
			},
		},
		"create-instance": {
			CreateInstances: &daisy.CreateInstances{
				Instances: []*daisy.Instance{
					{
						Instance: compute.Instance{
							Name: "inst-${NAME}",
							Disks: []*compute.AttachedDisk{
								{Source: "disk-${NAME}"},
								{Source: diskURI, Mode: "READ_ONLY"},
							},
							MachineType: "n1-highcpu-4",
							Metadata: &compute.Metadata{
								Items: []*compute.MetadataItems{
									{
										Key:   "startup-script",
										Value: &checksumScript,
									},
								},
							},
							NetworkInterfaces: []*compute.NetworkInterface{
								{
									AccessConfigs: []*compute.AccessConfig{},
								},
							},
						},
					},
				},
			},
		},
		"wait-for-checksum": {
			WaitForInstancesSignal: &daisy.WaitForInstancesSignal{
				{
					Name: "inst-${NAME}",
					SerialOutput: &daisy.SerialOutput{
						Port:         1,
						SuccessMatch: "Checksum calculated.",
						StatusMatch:  "Checksum:",
					},
				},
			},
		},
	}
	w.Dependencies = map[string][]string{
		"create-instance":   {"create-disks"},
		"wait-for-checksum": {"create-instance"},
	}

	// Calculate checksum within 20min.
	daisycommon.SetWorkflowAttributes(w, daisycommon.WorkflowAttributes{
		Project:           inflater.args.Project,
		Zone:              inflater.args.Zone,
		GCSPath:           inflater.args.ScratchBucketGcsPath,
		OAuth:             inflater.args.Oauth,
		Timeout:           "20m",
		ComputeEndpoint:   inflater.args.ComputeEndpoint,
		DisableGCSLogs:    inflater.args.GcsLogsDisabled,
		DisableCloudLogs:  inflater.args.CloudLogsDisabled,
		DisableStdoutLogs: inflater.args.StdoutLogsDisabled,
	})
	err := daisyUtils.RunWorkflowWithCancelSignal(ctx, w)
	if err != nil {
		return "", err
	}
	return w.GetSerialConsoleOutputValue("disk-checksum"), nil
}

// Dup logic in import_image.sh. If change anything here, please change in both places.
const (
	checksumScriptConst = `
		function serialOutputKeyValuePair() {
			echo "<serial-output key:'$1' value:'$2'>"
		}
		CHECK_DEVICE=sdb
		BLOCK_COUNT=$(cat /sys/class/block/$CHECK_DEVICE/size)

		# Check size = 200000*512 = 100MB
		CHECK_COUNT=200000
		CHECKSUM1=$(sudo dd if=/dev/$CHECK_DEVICE ibs=512 skip=0 count=$CHECK_COUNT | md5sum)
		CHECKSUM2=$(sudo dd if=/dev/$CHECK_DEVICE ibs=512 skip=$(( 2000000 - $CHECK_COUNT )) count=$CHECK_COUNT | md5sum)
		CHECKSUM3=$(sudo dd if=/dev/$CHECK_DEVICE ibs=512 skip=$(( 20000000 - $CHECK_COUNT )) count=$CHECK_COUNT | md5sum)
		CHECKSUM4=$(sudo dd if=/dev/$CHECK_DEVICE ibs=512 skip=$(( $BLOCK_COUNT - $CHECK_COUNT )) count=$CHECK_COUNT | md5sum)
		echo "Checksum: $(serialOutputKeyValuePair "disk-checksum" "$CHECKSUM1-$CHECKSUM2-$CHECKSUM3-$CHECKSUM4")"
		echo "Checksum calculated."`
)
<|endoftext|>


<commit_before>package algoliasearch

const (
	AltCorrection1 string = "altCorrection1"
	AltCorrection2 string = "altCorrection2"
)

type SearchSynonymsRes struct {
	Hits   []Synonym `json:"hits"`
	NbHits int       `json:"nbHits"`
}

type Synonym struct {
	// Common fields
	HighlightedResult Map `json:"_highlightedResult,omitempty"`
	ObjectID          string                 `json:"objectID"`
	Type              string                 `json:"type"`

	// Alternative correction synonym's fields
	Corrections []string `json:"corrections,omitempty"`
	Word        string   `json:"word,omitempty"`

	// One way synonym's fields
	Input    string   `json:"input,omitempty"`
	Synonyms []string `json:"synonyms,omitempty"`

	// Placeholder synonym's fields
	Placeholder  string   `json:"placeholder,omitempty"`
	Replacements []string `json:"replacements,omitempty"`

	// Simple synonym's field (shared with `oneWaySynonym`)
	// Synonyms []string `json:"synonyms"`
}

func NewAltCorrectionSynomym(objectID string, corrections []string, word string, t string) Synonym {
	return Synonym{
		ObjectID:    objectID,
		Type:        t,
		Corrections: corrections,
		Word:        word,
	}
}

func NewOneWaySynomym(objectID string, input string, synonyms []string) Synonym {
	return Synonym{
		ObjectID: objectID,
		Type:     "oneWaySynonym",
		Input:    input,
		Synonyms: synonyms,
	}
}

func NewPlaceholderSynonym(objectID string, placeholder string, replacements []string) Synonym {
	return Synonym{
		ObjectID:     objectID,
		Type:         "placeholder",
		Placeholder:  placeholder,
		Replacements: replacements,
	}
}

func NewSynonym(objectID string, synonyms []string) Synonym {
	return Synonym{
		ObjectID: objectID,
		Type:     "synonym",
		Synonyms: synonyms,
	}
}
<commit_msg>Fix typo in function name (NewOneWaySynomym -> NewOneWaySynonym)<commit_after>package algoliasearch

const (
	AltCorrection1 string = "altCorrection1"
	AltCorrection2 string = "altCorrection2"
)

type SearchSynonymsRes struct {
	Hits   []Synonym `json:"hits"`
	NbHits int       `json:"nbHits"`
}

type Synonym struct {
	// Common fields
	HighlightedResult Map `json:"_highlightedResult,omitempty"`
	ObjectID          string                 `json:"objectID"`
	Type              string                 `json:"type"`

	// Alternative correction synonym's fields
	Corrections []string `json:"corrections,omitempty"`
	Word        string   `json:"word,omitempty"`

	// One way synonym's fields
	Input    string   `json:"input,omitempty"`
	Synonyms []string `json:"synonyms,omitempty"`

	// Placeholder synonym's fields
	Placeholder  string   `json:"placeholder,omitempty"`
	Replacements []string `json:"replacements,omitempty"`

	// Simple synonym's field (shared with `oneWaySynonym`)
	// Synonyms []string `json:"synonyms"`
}

func NewAltCorrectionSynomym(objectID string, corrections []string, word string, t string) Synonym {
	return Synonym{
		ObjectID:    objectID,
		Type:        t,
		Corrections: corrections,
		Word:        word,
	}
}

func NewOneWaySynonym(objectID string, input string, synonyms []string) Synonym {
	return Synonym{
		ObjectID: objectID,
		Type:     "oneWaySynonym",
		Input:    input,
		Synonyms: synonyms,
	}
}

func NewPlaceholderSynonym(objectID string, placeholder string, replacements []string) Synonym {
	return Synonym{
		ObjectID:     objectID,
		Type:         "placeholder",
		Placeholder:  placeholder,
		Replacements: replacements,
	}
}

func NewSynonym(objectID string, synonyms []string) Synonym {
	return Synonym{
		ObjectID: objectID,
		Type:     "synonym",
		Synonyms: synonyms,
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 Google Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package subcommands implements a simple way for a single command to have many
// subcommands, each of which takes arguments and so forth.
package subcommands

import (
	"context"
	"flag"
	"fmt"
	"io"
	"os"
	"path"
	"sort"
)

// A Command represents a single command.
type Command interface {
	// Name returns the name of the command.
	Name() string

	// Synopsis returns a short string (less than one line) describing the command.
	Synopsis() string

	// Usage returns a long string explaining the command and giving usage
	// information.
	Usage() string

	// SetFlags adds the flags for this command to the specified set.
	SetFlags(*flag.FlagSet)

	// Execute executes the command and returns an ExitStatus.
	Execute(ctx context.Context, f *flag.FlagSet, args ...interface{}) ExitStatus
}

// A Commander represents a set of commands.
type Commander struct {
	commands  []*commandGroup
	topFlags  *flag.FlagSet // top-level flags
	important []string      // important top-level flags
	name      string        // normally path.Base(os.Args[0])

	Output io.Writer // Output specifies where the commander should write its output (default: os.Stdout).
	Error  io.Writer // Error specifies where the commander should write its error (default: os.Stderr).
}

// A commandGroup represents a set of commands about a common topic.
type commandGroup struct {
	name     string
	commands []Command
}

// An ExitStatus represents a Posix exit status that a subcommand
// expects to be returned to the shell.
type ExitStatus int

const (
	ExitSuccess ExitStatus = iota
	ExitFailure
	ExitUsageError
)

// NewCommander returns a new commander with the specified top-level
// flags and command name. The Usage function for the topLevelFlags
// will be set as well.
func NewCommander(topLevelFlags *flag.FlagSet, name string) *Commander {
	cdr := &Commander{
		topFlags: topLevelFlags,
		name:     name,
		Output:   os.Stdout,
		Error:    os.Stderr,
	}
	topLevelFlags.Usage = func() { cdr.explain(cdr.Error) }
	return cdr
}

// Register adds a subcommand to the supported subcommands in the
// specified group. (Help output is sorted and arranged by group name.)
// The empty string is an acceptable group name; such subcommands are
// explained first before named groups.
func (cdr *Commander) Register(cmd Command, group string) {
	for _, g := range cdr.commands {
		if g.name == group {
			g.commands = append(g.commands, cmd)
			return
		}
	}
	cdr.commands = append(cdr.commands, &commandGroup{
		name:     group,
		commands: []Command{cmd},
	})
}

// ImportantFlag marks a top-level flag as important, which means it
// will be printed out as part of the output of an ordinary "help"
// subcommand.  (All flags, important or not, are printed by the
// "flags" subcommand.)
func (cdr *Commander) ImportantFlag(name string) {
	cdr.important = append(cdr.important, name)
}

// Execute should be called once the top-level-flags on a Commander
// have been initialized. It finds the correct subcommand and executes
// it, and returns an ExitStatus with the result. On a usage error, an
// appropriate message is printed to os.Stderr, and ExitUsageError is
// returned. The additional args are provided as-is to the Execute method
// of the selected Command.
func (cdr *Commander) Execute(ctx context.Context, args ...interface{}) ExitStatus {
	if cdr.topFlags.NArg() < 1 {
		cdr.topFlags.Usage()
		return ExitUsageError
	}

	name := cdr.topFlags.Arg(0)

	for _, group := range cdr.commands {
		for _, cmd := range group.commands {
			if name != cmd.Name() {
				continue
			}
			f := flag.NewFlagSet(name, flag.ContinueOnError)
			f.Usage = func() { explain(cdr.Error, cmd) }
			cmd.SetFlags(f)
			if f.Parse(cdr.topFlags.Args()[1:]) != nil {
				return ExitUsageError
			}
			return cmd.Execute(ctx, f, args...)
		}
	}

	// Cannot find this command.
	cdr.topFlags.Usage()
	return ExitUsageError
}

// Sorting of a slice of command groups.
type byGroupName []*commandGroup

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

// explain prints a brief description of all the subcommands and the
// important top-level flags.
func (cdr *Commander) explain(w io.Writer) {
	fmt.Fprintf(w, "Usage: %s <flags> <subcommand> <subcommand args>\n\n", cdr.name)
	sort.Sort(byGroupName(cdr.commands))
	for _, group := range cdr.commands {
		explainGroup(w, group)
	}
	if cdr.topFlags == nil {
		fmt.Fprintln(w, "\nNo top level flags.")
		return
	}
	if len(cdr.important) == 0 {
		fmt.Fprintf(w, "\nUse \"%s flags\" for a list of top-level flags\n", cdr.name)
		return
	}

	fmt.Fprintf(w, "\nTop-level flags (use \"%s flags\" for a full list):\n", cdr.name)
	for _, name := range cdr.important {
		f := cdr.topFlags.Lookup(name)
		if f == nil {
			panic("Important flag is not defined")
		}
		fmt.Fprintf(w, "  -%s=%s: %s\n", f.Name, f.DefValue, f.Usage)
	}
}

// Sorting of the commands within a group.
func (g commandGroup) Len() int           { return len(g.commands) }
func (g commandGroup) Less(i, j int) bool { return g.commands[i].Name() < g.commands[j].Name() }
func (g commandGroup) Swap(i, j int)      { g.commands[i], g.commands[j] = g.commands[j], g.commands[i] }

// explainGroup explains all the subcommands for a particular group.
func explainGroup(w io.Writer, group *commandGroup) {
	if len(group.commands) == 0 {
		return
	}
	if group.name == "" {
		fmt.Fprintf(w, "Subcommands:\n")
	} else {
		fmt.Fprintf(w, "Subcommands for %s:\n", group.name)
	}
	sort.Sort(group)
	for _, cmd := range group.commands {
		fmt.Fprintf(w, "\t%-15s  %s\n", cmd.Name(), cmd.Synopsis())
	}
	fmt.Fprintln(w)
}

// explainCmd prints a brief description of a single command.
func explain(w io.Writer, cmd Command) {
	fmt.Fprintf(w, "%s", cmd.Usage())
	subflags := flag.NewFlagSet(cmd.Name(), flag.PanicOnError)
	subflags.SetOutput(w)
	cmd.SetFlags(subflags)
	subflags.PrintDefaults()
}

// A helper is a Command implementing a "help" command for
// a given Commander.
type helper Commander

func (h *helper) Name() string           { return "help" }
func (h *helper) Synopsis() string       { return "describe subcommands and their syntax" }
func (h *helper) SetFlags(*flag.FlagSet) {}
func (h *helper) Usage() string {
	return `help [<subcommand>]:
	With an argument, prints detailed information on the use of
	the specified subcommand. With no argument, print a list of
	all commands and a brief description of each.
`
}
func (h *helper) Execute(_ context.Context, f *flag.FlagSet, args ...interface{}) ExitStatus {
	switch f.NArg() {
	case 0:
		(*Commander)(h).explain(h.Output)
		return ExitSuccess

	case 1:
		for _, group := range h.commands {
			for _, cmd := range group.commands {
				if f.Arg(0) != cmd.Name() {
					continue
				}
				explain(h.Output, cmd)
				return ExitSuccess
			}
		}
		fmt.Fprintf(h.Error, "Subcommand %s not understood\n", f.Arg(0))
	}

	f.Usage()
	return ExitUsageError
}

// HelpCommand returns a Command which implements a "help" subcommand.
func (cdr *Commander) HelpCommand() Command {
	return (*helper)(cdr)
}

// A flagger is a Command implementing a "flags" command for a given Commander.
type flagger Commander

func (flg *flagger) Name() string           { return "flags" }
func (flg *flagger) Synopsis() string       { return "describe all known top-level flags" }
func (flg *flagger) SetFlags(*flag.FlagSet) {}
func (flg *flagger) Usage() string {
	return `flags [<subcommand>]:
	With an argument, print all flags of <subcommand>. Else,
	print a description of all known top-level flags.  (The basic
	help information only discusses the most generally important
	top-level flags.)
`
}
func (flg *flagger) Execute(_ context.Context, f *flag.FlagSet, _ ...interface{}) ExitStatus {
	if f.NArg() > 1 {
		f.Usage()
		return ExitUsageError
	}

	if f.NArg() == 0 {
		if flg.topFlags == nil {
			fmt.Fprintln(flg.Output, "No top-level flags are defined.")
		} else {
			flg.topFlags.PrintDefaults()
		}
		return ExitSuccess
	}

	for _, group := range flg.commands {
		for _, cmd := range group.commands {
			if f.Arg(0) != cmd.Name() {
				continue
			}
			subflags := flag.NewFlagSet(cmd.Name(), flag.PanicOnError)
			subflags.SetOutput(flg.Output)
			cmd.SetFlags(subflags)
			subflags.PrintDefaults()
			return ExitSuccess
		}
	}
	fmt.Fprintf(flg.Error, "Subcommand %s not understood\n", f.Arg(0))
	return ExitFailure
}

// FlagsCommand returns a Command which implements a "flags" subcommand.
func (cdr *Commander) FlagsCommand() Command {
	return (*flagger)(cdr)
}

// A lister is a Command implementing a "commands" command for a given Commander.
type lister Commander

func (l *lister) Name() string           { return "commands" }
func (l *lister) Synopsis() string       { return "list all command names" }
func (l *lister) SetFlags(*flag.FlagSet) {}
func (l *lister) Usage() string {
	return `commands:
	Print a list of all commands.
`
}
func (l *lister) Execute(_ context.Context, f *flag.FlagSet, _ ...interface{}) ExitStatus {
	if f.NArg() != 0 {
		f.Usage()
		return ExitUsageError
	}

	for _, group := range l.commands {
		for _, cmd := range group.commands {
			fmt.Fprintf(l.Output, "%s\n", cmd.Name())
		}
	}
	return ExitSuccess
}

// CommandsCommand returns Command which implements a "commands" subcommand.
func (cdr *Commander) CommandsCommand() Command {
	return (*lister)(cdr)
}

// DefaultCommander is the default commander using flag.CommandLine for flags
// and os.Args[0] for the command name.
var DefaultCommander *Commander

func init() {
	DefaultCommander = NewCommander(flag.CommandLine, path.Base(os.Args[0]))
}

// Register adds a subcommand to the supported subcommands in the
// specified group. (Help output is sorted and arranged by group
// name.)  The empty string is an acceptable group name; such
// subcommands are explained first before named groups. It is a
// wrapper around DefaultCommander.Register.
func Register(cmd Command, group string) {
	DefaultCommander.Register(cmd, group)
}

// ImportantFlag marks a top-level flag as important, which means it
// will be printed out as part of the output of an ordinary "help"
// subcommand.  (All flags, important or not, are printed by the
// "flags" subcommand.) It is a wrapper around
// DefaultCommander.ImportantFlag.
func ImportantFlag(name string) {
	DefaultCommander.ImportantFlag(name)
}

// Execute should be called once the default flags have been
// initialized by flag.Parse. It finds the correct subcommand and
// executes it, and returns an ExitStatus with the result. On a usage
// error, an appropriate message is printed to os.Stderr, and
// ExitUsageError is returned. The additional args are provided as-is
// to the Execute method of the selected Command. It is a wrapper
// around DefaultCommander.Execute.
func Execute(ctx context.Context, args ...interface{}) ExitStatus {
	return DefaultCommander.Execute(ctx, args...)
}

// HelpCommand returns a Command which implements "help" for the
// DefaultCommander. Use Register(HelpCommand(), <group>) for it to be
// recognized.
func HelpCommand() Command {
	return DefaultCommander.HelpCommand()
}

// FlagsCommand returns a Command which implements "flags" for the
// DefaultCommander. Use Register(FlagsCommand(), <group>) for it to be
// recognized.
func FlagsCommand() Command {
	return DefaultCommander.FlagsCommand()
}

// CommandsCommand returns Command which implements a "commands" subcommand.
func CommandsCommand() Command {
	return DefaultCommander.CommandsCommand()
}
<commit_msg>Update panic message to include the name of the undefined, important flag (#12)<commit_after>/*
Copyright 2016 Google Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Package subcommands implements a simple way for a single command to have many
// subcommands, each of which takes arguments and so forth.
package subcommands

import (
	"context"
	"flag"
	"fmt"
	"io"
	"os"
	"path"
	"sort"
)

// A Command represents a single command.
type Command interface {
	// Name returns the name of the command.
	Name() string

	// Synopsis returns a short string (less than one line) describing the command.
	Synopsis() string

	// Usage returns a long string explaining the command and giving usage
	// information.
	Usage() string

	// SetFlags adds the flags for this command to the specified set.
	SetFlags(*flag.FlagSet)

	// Execute executes the command and returns an ExitStatus.
	Execute(ctx context.Context, f *flag.FlagSet, args ...interface{}) ExitStatus
}

// A Commander represents a set of commands.
type Commander struct {
	commands  []*commandGroup
	topFlags  *flag.FlagSet // top-level flags
	important []string      // important top-level flags
	name      string        // normally path.Base(os.Args[0])

	Output io.Writer // Output specifies where the commander should write its output (default: os.Stdout).
	Error  io.Writer // Error specifies where the commander should write its error (default: os.Stderr).
}

// A commandGroup represents a set of commands about a common topic.
type commandGroup struct {
	name     string
	commands []Command
}

// An ExitStatus represents a Posix exit status that a subcommand
// expects to be returned to the shell.
type ExitStatus int

const (
	ExitSuccess ExitStatus = iota
	ExitFailure
	ExitUsageError
)

// NewCommander returns a new commander with the specified top-level
// flags and command name. The Usage function for the topLevelFlags
// will be set as well.
func NewCommander(topLevelFlags *flag.FlagSet, name string) *Commander {
	cdr := &Commander{
		topFlags: topLevelFlags,
		name:     name,
		Output:   os.Stdout,
		Error:    os.Stderr,
	}
	topLevelFlags.Usage = func() { cdr.explain(cdr.Error) }
	return cdr
}

// Register adds a subcommand to the supported subcommands in the
// specified group. (Help output is sorted and arranged by group name.)
// The empty string is an acceptable group name; such subcommands are
// explained first before named groups.
func (cdr *Commander) Register(cmd Command, group string) {
	for _, g := range cdr.commands {
		if g.name == group {
			g.commands = append(g.commands, cmd)
			return
		}
	}
	cdr.commands = append(cdr.commands, &commandGroup{
		name:     group,
		commands: []Command{cmd},
	})
}

// ImportantFlag marks a top-level flag as important, which means it
// will be printed out as part of the output of an ordinary "help"
// subcommand.  (All flags, important or not, are printed by the
// "flags" subcommand.)
func (cdr *Commander) ImportantFlag(name string) {
	cdr.important = append(cdr.important, name)
}

// Execute should be called once the top-level-flags on a Commander
// have been initialized. It finds the correct subcommand and executes
// it, and returns an ExitStatus with the result. On a usage error, an
// appropriate message is printed to os.Stderr, and ExitUsageError is
// returned. The additional args are provided as-is to the Execute method
// of the selected Command.
func (cdr *Commander) Execute(ctx context.Context, args ...interface{}) ExitStatus {
	if cdr.topFlags.NArg() < 1 {
		cdr.topFlags.Usage()
		return ExitUsageError
	}

	name := cdr.topFlags.Arg(0)

	for _, group := range cdr.commands {
		for _, cmd := range group.commands {
			if name != cmd.Name() {
				continue
			}
			f := flag.NewFlagSet(name, flag.ContinueOnError)
			f.Usage = func() { explain(cdr.Error, cmd) }
			cmd.SetFlags(f)
			if f.Parse(cdr.topFlags.Args()[1:]) != nil {
				return ExitUsageError
			}
			return cmd.Execute(ctx, f, args...)
		}
	}

	// Cannot find this command.
	cdr.topFlags.Usage()
	return ExitUsageError
}

// Sorting of a slice of command groups.
type byGroupName []*commandGroup

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

// explain prints a brief description of all the subcommands and the
// important top-level flags.
func (cdr *Commander) explain(w io.Writer) {
	fmt.Fprintf(w, "Usage: %s <flags> <subcommand> <subcommand args>\n\n", cdr.name)
	sort.Sort(byGroupName(cdr.commands))
	for _, group := range cdr.commands {
		explainGroup(w, group)
	}
	if cdr.topFlags == nil {
		fmt.Fprintln(w, "\nNo top level flags.")
		return
	}
	if len(cdr.important) == 0 {
		fmt.Fprintf(w, "\nUse \"%s flags\" for a list of top-level flags\n", cdr.name)
		return
	}

	fmt.Fprintf(w, "\nTop-level flags (use \"%s flags\" for a full list):\n", cdr.name)
	for _, name := range cdr.important {
		f := cdr.topFlags.Lookup(name)
		if f == nil {
			panic(fmt.Sprintf("Important flag (%s) is not defined", name))
		}
		fmt.Fprintf(w, "  -%s=%s: %s\n", f.Name, f.DefValue, f.Usage)
	}
}

// Sorting of the commands within a group.
func (g commandGroup) Len() int           { return len(g.commands) }
func (g commandGroup) Less(i, j int) bool { return g.commands[i].Name() < g.commands[j].Name() }
func (g commandGroup) Swap(i, j int)      { g.commands[i], g.commands[j] = g.commands[j], g.commands[i] }

// explainGroup explains all the subcommands for a particular group.
func explainGroup(w io.Writer, group *commandGroup) {
	if len(group.commands) == 0 {
		return
	}
	if group.name == "" {
		fmt.Fprintf(w, "Subcommands:\n")
	} else {
		fmt.Fprintf(w, "Subcommands for %s:\n", group.name)
	}
	sort.Sort(group)
	for _, cmd := range group.commands {
		fmt.Fprintf(w, "\t%-15s  %s\n", cmd.Name(), cmd.Synopsis())
	}
	fmt.Fprintln(w)
}

// explainCmd prints a brief description of a single command.
func explain(w io.Writer, cmd Command) {
	fmt.Fprintf(w, "%s", cmd.Usage())
	subflags := flag.NewFlagSet(cmd.Name(), flag.PanicOnError)
	subflags.SetOutput(w)
	cmd.SetFlags(subflags)
	subflags.PrintDefaults()
}

// A helper is a Command implementing a "help" command for
// a given Commander.
type helper Commander

func (h *helper) Name() string           { return "help" }
func (h *helper) Synopsis() string       { return "describe subcommands and their syntax" }
func (h *helper) SetFlags(*flag.FlagSet) {}
func (h *helper) Usage() string {
	return `help [<subcommand>]:
	With an argument, prints detailed information on the use of
	the specified subcommand. With no argument, print a list of
	all commands and a brief description of each.
`
}
func (h *helper) Execute(_ context.Context, f *flag.FlagSet, args ...interface{}) ExitStatus {
	switch f.NArg() {
	case 0:
		(*Commander)(h).explain(h.Output)
		return ExitSuccess

	case 1:
		for _, group := range h.commands {
			for _, cmd := range group.commands {
				if f.Arg(0) != cmd.Name() {
					continue
				}
				explain(h.Output, cmd)
				return ExitSuccess
			}
		}
		fmt.Fprintf(h.Error, "Subcommand %s not understood\n", f.Arg(0))
	}

	f.Usage()
	return ExitUsageError
}

// HelpCommand returns a Command which implements a "help" subcommand.
func (cdr *Commander) HelpCommand() Command {
	return (*helper)(cdr)
}

// A flagger is a Command implementing a "flags" command for a given Commander.
type flagger Commander

func (flg *flagger) Name() string           { return "flags" }
func (flg *flagger) Synopsis() string       { return "describe all known top-level flags" }
func (flg *flagger) SetFlags(*flag.FlagSet) {}
func (flg *flagger) Usage() string {
	return `flags [<subcommand>]:
	With an argument, print all flags of <subcommand>. Else,
	print a description of all known top-level flags.  (The basic
	help information only discusses the most generally important
	top-level flags.)
`
}
func (flg *flagger) Execute(_ context.Context, f *flag.FlagSet, _ ...interface{}) ExitStatus {
	if f.NArg() > 1 {
		f.Usage()
		return ExitUsageError
	}

	if f.NArg() == 0 {
		if flg.topFlags == nil {
			fmt.Fprintln(flg.Output, "No top-level flags are defined.")
		} else {
			flg.topFlags.PrintDefaults()
		}
		return ExitSuccess
	}

	for _, group := range flg.commands {
		for _, cmd := range group.commands {
			if f.Arg(0) != cmd.Name() {
				continue
			}
			subflags := flag.NewFlagSet(cmd.Name(), flag.PanicOnError)
			subflags.SetOutput(flg.Output)
			cmd.SetFlags(subflags)
			subflags.PrintDefaults()
			return ExitSuccess
		}
	}
	fmt.Fprintf(flg.Error, "Subcommand %s not understood\n", f.Arg(0))
	return ExitFailure
}

// FlagsCommand returns a Command which implements a "flags" subcommand.
func (cdr *Commander) FlagsCommand() Command {
	return (*flagger)(cdr)
}

// A lister is a Command implementing a "commands" command for a given Commander.
type lister Commander

func (l *lister) Name() string           { return "commands" }
func (l *lister) Synopsis() string       { return "list all command names" }
func (l *lister) SetFlags(*flag.FlagSet) {}
func (l *lister) Usage() string {
	return `commands:
	Print a list of all commands.
`
}
func (l *lister) Execute(_ context.Context, f *flag.FlagSet, _ ...interface{}) ExitStatus {
	if f.NArg() != 0 {
		f.Usage()
		return ExitUsageError
	}

	for _, group := range l.commands {
		for _, cmd := range group.commands {
			fmt.Fprintf(l.Output, "%s\n", cmd.Name())
		}
	}
	return ExitSuccess
}

// CommandsCommand returns Command which implements a "commands" subcommand.
func (cdr *Commander) CommandsCommand() Command {
	return (*lister)(cdr)
}

// DefaultCommander is the default commander using flag.CommandLine for flags
// and os.Args[0] for the command name.
var DefaultCommander *Commander

func init() {
	DefaultCommander = NewCommander(flag.CommandLine, path.Base(os.Args[0]))
}

// Register adds a subcommand to the supported subcommands in the
// specified group. (Help output is sorted and arranged by group
// name.)  The empty string is an acceptable group name; such
// subcommands are explained first before named groups. It is a
// wrapper around DefaultCommander.Register.
func Register(cmd Command, group string) {
	DefaultCommander.Register(cmd, group)
}

// ImportantFlag marks a top-level flag as important, which means it
// will be printed out as part of the output of an ordinary "help"
// subcommand.  (All flags, important or not, are printed by the
// "flags" subcommand.) It is a wrapper around
// DefaultCommander.ImportantFlag.
func ImportantFlag(name string) {
	DefaultCommander.ImportantFlag(name)
}

// Execute should be called once the default flags have been
// initialized by flag.Parse. It finds the correct subcommand and
// executes it, and returns an ExitStatus with the result. On a usage
// error, an appropriate message is printed to os.Stderr, and
// ExitUsageError is returned. The additional args are provided as-is
// to the Execute method of the selected Command. It is a wrapper
// around DefaultCommander.Execute.
func Execute(ctx context.Context, args ...interface{}) ExitStatus {
	return DefaultCommander.Execute(ctx, args...)
}

// HelpCommand returns a Command which implements "help" for the
// DefaultCommander. Use Register(HelpCommand(), <group>) for it to be
// recognized.
func HelpCommand() Command {
	return DefaultCommander.HelpCommand()
}

// FlagsCommand returns a Command which implements "flags" for the
// DefaultCommander. Use Register(FlagsCommand(), <group>) for it to be
// recognized.
func FlagsCommand() Command {
	return DefaultCommander.FlagsCommand()
}

// CommandsCommand returns Command which implements a "commands" subcommand.
func CommandsCommand() Command {
	return DefaultCommander.CommandsCommand()
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"

	"./tree"
	// "github.com/marcsantiago/GoTree/tree"
)

func main() {
	// Using time go run example.go, takes
	// go run example.go  235.89s user 11.92s system 174% cpu 2:21.75 total
	// While suming with this function is a bit slower then just calling GetTreeTotal()
	// the real bottle neck is adding that many Nodes to the tree. It uses a lot of memory, so be careful when
	// executing this funcion
	t := tree.NewTree()
	log.Println("Creating a tree with 16777216 nodes")
	t.GenerateRandomTreeIteratively(16777216)
	log.Println("Adding Nodes using the slower Sum() function")
	log.Println("Total:", t.Sum())

	t = tree.NewTree()
	t.AddIteratively(10)
	t.AddIteratively(100)
	t.AddIteratively(2)
	t.AddIteratively(3)
	fmt.Println(t.Sum()) // should be 142
	fmt.Println(t.GetTreeTotal())
	fmt.Println("edges", t.CountEdges())
	fmt.Printf("%+v\n", t.TreeToArray())
	t.PrintTree()

	t = tree.NewTree()
	t.AddIteratively(1)
	t.AddIteratively(2)
	t.AddIteratively(3)
	t.AddIteratively(4)
	t.AddIteratively(10)
	t.AddIteratively(9)
	t.AddIteratively(8)
	t.PrintTree()
}
<commit_msg>updating run example comments. GenerateRandomTreeIteratively created a large tree in no time at all, the bottle neck is the Sum() function.. GetTreeTotal(), should be used... or better yet.. Tree.Total since thats an exported value<commit_after>package main

import (
	"fmt"
	"log"

	"./tree"
	// "github.com/marcsantiago/GoTree/tree"
)

func main() {
	// Using time go run example.go, takes
	// go run example.go  81.86s user 11.17s system 195% cpu 47.536 total
	// While suming with this function is a slower then just calling GetTreeTotal()
	// the real bottle neck is adding that many Nodes to the tree. It uses a lot of memory, so be careful when
	// executing this funcion
	t := tree.NewTree()
	log.Println("Creating a tree with 16777216 nodes")
	t.GenerateRandomTreeIteratively(16777216)
	log.Println("Adding Nodes using the slower Sum() function")
	log.Println("Total:", t.Sum())

	t = tree.NewTree()
	t.AddIteratively(10)
	t.AddIteratively(100)
	t.AddIteratively(2)
	t.AddIteratively(3)
	fmt.Println(t.Sum()) // should be 142
	fmt.Println(t.GetTreeTotal())
	fmt.Println("edges", t.CountEdges())
	fmt.Printf("%+v\n", t.TreeToArray())
	t.PrintTree()

	t = tree.NewTree()
	t.AddIteratively(1)
	t.AddIteratively(2)
	t.AddIteratively(3)
	t.AddIteratively(4)
	t.AddIteratively(10)
	t.AddIteratively(9)
	t.AddIteratively(8)
	t.PrintTree()
}
<|endoftext|>


<commit_before>package fenwick

import "errors"

var IndexError = errors.New("Index is out of range")

type Fenwick struct {
	data []int
}

func New(size uint) *Fenwick {
	var f Fenwick
	f.data = make([]int, size)
	return &f
}

func (f *Fenwick) Add(pos int, value int) error {
	if pos < 0 || pos >= len(f.data) {
		return IndexError
	}
	for ; pos < len(f.data); pos = pos | (pos + 1) {
		f.data[pos] += value
	}
	return nil
}

func (f *Fenwick) Sum(begin, end int) int {
	endSum := f.prefixSum(end)
	beginSum := f.prefixSum(begin - 1)
	return endSum - beginSum
}

func (f *Fenwick) prefixSum(pos int) (result int) {
	if pos >= len(f.data) {
		pos = len(f.data) - 1
	}
	for ; pos >= 0; pos = (pos & (pos + 1)) - 1 {
		result += f.data[pos]
	}
	return
}
<commit_msg>add comments<commit_after>/* Fenwick tree (https://en.wikipedia.org/wiki/Fenwick_tree) */
package fenwick

import "errors"

var IndexError = errors.New("Index is out of range")

// Basic structure
type Fenwick struct {
	data []int
}

// New creates a Fenwick tree of fixed size
func New(size uint) *Fenwick {
	var f Fenwick
	f.data = make([]int, size)
	return &f
}

// Increment value at position by value
func (f *Fenwick) Add(pos int, value int) error {
	if pos < 0 || pos >= len(f.data) {
		return IndexError
	}
	for ; pos < len(f.data); pos = pos | (pos + 1) {
		f.data[pos] += value
	}
	return nil
}

// Return sum of value on segment [begin, end] (including both begin and end)
func (f *Fenwick) Sum(begin, end int) int {
	endSum := f.prefixSum(end)
	beginSum := f.prefixSum(begin - 1)
	return endSum - beginSum
}

func (f *Fenwick) prefixSum(pos int) (result int) {
	if pos >= len(f.data) {
		pos = len(f.data) - 1
	}
	for ; pos >= 0; pos = (pos & (pos + 1)) - 1 {
		result += f.data[pos]
	}
	return
}
<|endoftext|>


<commit_before>package stats

import (
	"bytes"
	"io"
	"net"
	"time"

	"github.com/raintank/worldping-api/pkg/log"
)

var (
	queueItems      *Range32
	genDataDuration *Gauge32
	flushDuration   *LatencyHistogram15s32
	messageSize     *Gauge32
	connected       *Bool
)

type GraphiteMetric interface {
	// Report the measurements in graphite format and reset measurements for the next interval if needed
	ReportGraphite(prefix []byte, buf []byte, now time.Time) []byte
}

type Graphite struct {
	prefix []byte
	addr   string

	timeout    time.Duration
	toGraphite chan []byte
}

func NewGraphite(prefix, addr string, interval, bufferSize int, timeout time.Duration) {
	if len(prefix) != 0 && prefix[len(prefix)-1] != '.' {
		prefix = prefix + "."
	}
	NewGauge32("stats.graphite.write_queue.size").Set(bufferSize)
	queueItems = NewRange32("stats.graphite.write_queue.items")
	// metric stats.generate_message is how long it takes to generate the stats
	genDataDuration = NewGauge32("stats.generate_message.duration")
	flushDuration = NewLatencyHistogram15s32("stats.graphite.flush")
	messageSize = NewGauge32("stats.message_size")
	connected = NewBool("stats.graphite.connected")

	g := &Graphite{
		prefix:     []byte(prefix),
		addr:       addr,
		toGraphite: make(chan []byte, bufferSize),
		timeout:    timeout,
	}
	go g.writer()
	go g.reporter(interval)
}

func (g *Graphite) reporter(interval int) {
	ticker := tick(time.Duration(interval) * time.Second)
	for now := range ticker {
		log.Debug("stats flushing for", now, "to graphite")
		queueItems.Value(len(g.toGraphite))
		if cap(g.toGraphite) != 0 && len(g.toGraphite) == cap(g.toGraphite) {
			// no space in buffer, no use in doing any work
			continue
		}

		pre := time.Now()

		buf := make([]byte, 0)

		var fullPrefix bytes.Buffer
		for name, metric := range registry.list() {
			fullPrefix.Reset()
			fullPrefix.Write(g.prefix)
			fullPrefix.WriteString(name)
			fullPrefix.WriteRune('.')
			buf = metric.ReportGraphite(fullPrefix.Bytes(), buf, now)
		}

		genDataDuration.Set(int(time.Since(pre).Nanoseconds()))
		messageSize.Set(len(buf))
		g.toGraphite <- buf
		queueItems.Value(len(g.toGraphite))
	}
}

// writer connects to graphite and submits all pending data to it
func (g *Graphite) writer() {
	var conn net.Conn
	var err error

	assureConn := func() net.Conn {
		connected.Set(conn != nil)
		for conn == nil {
			time.Sleep(time.Second)
			conn, err = net.Dial("tcp", g.addr)
			if err == nil {
				log.Info("stats now connected to %s", g.addr)
				go g.checkEOF(conn)
			} else {
				log.Warn("stats dialing %s failed: %s. will retry", g.addr, err.Error())
			}
			connected.Set(conn != nil)
		}
		return conn
	}

	for buf := range g.toGraphite {
		queueItems.Value(len(g.toGraphite))
		var ok bool
		for !ok {
			conn = assureConn()
			conn.SetDeadline(time.Now().Add(g.timeout))
			pre := time.Now()
			_, err = conn.Write(buf)
			if err == nil {
				ok = true
				flushDuration.Value(time.Since(pre))
			} else {
				log.Warn("stats failed to write to graphite: %s (took %s). will retry...", err, time.Now().Sub(pre))
				conn.Close()
				conn = nil
			}
		}
	}
}

// normally the remote end should never write anything back
// but we know when we get EOF that the other end closed the conn
// if not for this, we can happily write and flush without getting errors (in Go) but getting RST tcp packets back (!)
// props to Tv` for this trick.
func (g *Graphite) checkEOF(conn net.Conn) {
	b := make([]byte, 1024)
	for {
		num, err := conn.Read(b)
		if err == io.EOF {
			log.Info("Graphite.checkEOF: remote closed conn. closing conn")
			conn.Close()
			return
		}

		// just in case i misunderstand something or the remote behaves badly
		if num != 0 {
			log.Warn("Graphite.checkEOF: read unexpected data from peer: %s\n", b[:num])
			continue
		}

		if err != io.EOF {
			log.Warn("Graphite.checkEOF: closing conn: %s\n", err)
			conn.Close()
			return
		}
	}
}
<commit_msg>only set write deadline<commit_after>package stats

import (
	"bytes"
	"io"
	"net"
	"time"

	"github.com/raintank/worldping-api/pkg/log"
)

var (
	queueItems      *Range32
	genDataDuration *Gauge32
	flushDuration   *LatencyHistogram15s32
	messageSize     *Gauge32
	connected       *Bool
)

type GraphiteMetric interface {
	// Report the measurements in graphite format and reset measurements for the next interval if needed
	ReportGraphite(prefix []byte, buf []byte, now time.Time) []byte
}

type Graphite struct {
	prefix []byte
	addr   string

	timeout    time.Duration
	toGraphite chan []byte
}

func NewGraphite(prefix, addr string, interval, bufferSize int, timeout time.Duration) {
	if len(prefix) != 0 && prefix[len(prefix)-1] != '.' {
		prefix = prefix + "."
	}
	NewGauge32("stats.graphite.write_queue.size").Set(bufferSize)
	queueItems = NewRange32("stats.graphite.write_queue.items")
	// metric stats.generate_message is how long it takes to generate the stats
	genDataDuration = NewGauge32("stats.generate_message.duration")
	flushDuration = NewLatencyHistogram15s32("stats.graphite.flush")
	messageSize = NewGauge32("stats.message_size")
	connected = NewBool("stats.graphite.connected")

	g := &Graphite{
		prefix:     []byte(prefix),
		addr:       addr,
		toGraphite: make(chan []byte, bufferSize),
		timeout:    timeout,
	}
	go g.writer()
	go g.reporter(interval)
}

func (g *Graphite) reporter(interval int) {
	ticker := tick(time.Duration(interval) * time.Second)
	for now := range ticker {
		log.Debug("stats flushing for", now, "to graphite")
		queueItems.Value(len(g.toGraphite))
		if cap(g.toGraphite) != 0 && len(g.toGraphite) == cap(g.toGraphite) {
			// no space in buffer, no use in doing any work
			continue
		}

		pre := time.Now()

		buf := make([]byte, 0)

		var fullPrefix bytes.Buffer
		for name, metric := range registry.list() {
			fullPrefix.Reset()
			fullPrefix.Write(g.prefix)
			fullPrefix.WriteString(name)
			fullPrefix.WriteRune('.')
			buf = metric.ReportGraphite(fullPrefix.Bytes(), buf, now)
		}

		genDataDuration.Set(int(time.Since(pre).Nanoseconds()))
		messageSize.Set(len(buf))
		g.toGraphite <- buf
		queueItems.Value(len(g.toGraphite))
	}
}

// writer connects to graphite and submits all pending data to it
func (g *Graphite) writer() {
	var conn net.Conn
	var err error

	assureConn := func() net.Conn {
		connected.Set(conn != nil)
		for conn == nil {
			time.Sleep(time.Second)
			conn, err = net.Dial("tcp", g.addr)
			if err == nil {
				log.Info("stats now connected to %s", g.addr)
				go g.checkEOF(conn)
			} else {
				log.Warn("stats dialing %s failed: %s. will retry", g.addr, err.Error())
			}
			connected.Set(conn != nil)
		}
		return conn
	}

	for buf := range g.toGraphite {
		queueItems.Value(len(g.toGraphite))
		var ok bool
		for !ok {
			conn = assureConn()
			conn.SetWriteDeadline(time.Now().Add(g.timeout))
			pre := time.Now()
			_, err = conn.Write(buf)
			if err == nil {
				ok = true
				flushDuration.Value(time.Since(pre))
			} else {
				log.Warn("stats failed to write to graphite: %s (took %s). will retry...", err, time.Now().Sub(pre))
				conn.Close()
				conn = nil
			}
		}
	}
}

// normally the remote end should never write anything back
// but we know when we get EOF that the other end closed the conn
// if not for this, we can happily write and flush without getting errors (in Go) but getting RST tcp packets back (!)
// props to Tv` for this trick.
func (g *Graphite) checkEOF(conn net.Conn) {
	b := make([]byte, 1024)
	for {
		num, err := conn.Read(b)
		if err == io.EOF {
			log.Info("Graphite.checkEOF: remote closed conn. closing conn")
			conn.Close()
			return
		}

		// just in case i misunderstand something or the remote behaves badly
		if num != 0 {
			log.Warn("Graphite.checkEOF: read unexpected data from peer: %s\n", b[:num])
			continue
		}

		if err != io.EOF {
			log.Warn("Graphite.checkEOF: closing conn: %s\n", err)
			conn.Close()
			return
		}
	}
}
<|endoftext|>