<commit_before>package aws

import (
	"fmt"
	"log"
	"math/rand"
	"strconv"
	"sync"
	"time"

	awssdk "github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/elbv2"

	"github.com/gliderlabs/registrator/bridge"
	fargo "github.com/hudl/fargo"
)

// LBInfo represents a ELBv2 endpoint
type LBInfo struct {
	DNSName string
	Port    int64
}

type lookupValues struct {
	InstanceID string
	Port       int64
}

type cacheEntry struct {
	*sync.Mutex
	lb *LBInfo
}

type lbCache struct {
	*sync.Mutex
	m map[string]cacheEntry
}

var cache = lbCache{&sync.Mutex{}, make(map[string]cacheEntry)}

type fn func(lookupValues) (*LBInfo, error)

//
// Return a *LBInfo cache entry if it exists, or run the provided function to return data to add to cache
// The locking complexity is purely to make the cache thread safe.
//
func getOrAddCacheEntry(key string, f fn, i lookupValues) (*LBInfo, error) {
	cache.Lock()
	if _, ok := cache.m[key]; !ok {
		cache.m[key] = cacheEntry{&sync.Mutex{}, &LBInfo{Port: 0, DNSName: "should-have-been-set"}}
	}
	cache.m[key].Lock()
	defer cache.m[key].Unlock()
	cache.Unlock()
	if cache.m[key].lb.Port == 0 {
		o, err := f(i)
		if err != nil {
			log.Print("An error occurred trying to add data to the cache:", err)
			return nil, err
		}
		tmp := cache.m[key]
		tmp.lb = o
		cache.m[key] = tmp
	}
	return cache.m[key].lb, nil
}

// RemoveLBCache : Delete any cache of load balancer for this containerID
func RemoveLBCache(key string) {
	cache.Lock()
	delete(cache.m, key)
	cache.Unlock()
}

var registrations = make(map[string]bool)

// Helper function to retrieve all target groups
func getAllTargetGroups(svc *elbv2.ELBV2) ([]*elbv2.DescribeTargetGroupsOutput, error) {
	var tgs []*elbv2.DescribeTargetGroupsOutput
	var e error
	var mark *string

	// Get first page of groups
	tg, e := getTargetGroupsPage(svc, mark)

	if e != nil {
		return nil, e
	}
	tgs = append(tgs, tg)
	mark = tg.NextMarker

	// Page through all remaining target groups generating a slice of DescribeTargetGroupOutputs
	for mark != nil {
		tg, e = getTargetGroupsPage(svc, mark)
		tgs = append(tgs, tg)
		mark = tg.NextMarker
		if e != nil {
			return nil, e
		}
	}
	return tgs, e
}

// Helper function to get a page of target groups
func getTargetGroupsPage(svc *elbv2.ELBV2, marker *string) (*elbv2.DescribeTargetGroupsOutput, error) {
	params := &elbv2.DescribeTargetGroupsInput{
		PageSize: awssdk.Int64(400),
		Marker:   marker,
	}

	// Random wait to try to avoid getting throttled by AWS API
	seed := rand.NewSource(time.Now().UnixNano())
	r2 := rand.New(seed)
	random := r2.Intn(5000)
	period := time.Millisecond * time.Duration(random)
	time.Sleep(period)

	tg, e := svc.DescribeTargetGroups(params)

	if e != nil {
		log.Printf("An error occurred using DescribeTargetGroups: %s \n", e.Error())
		return nil, e
	}
	return tg, nil
}

// GetELBV2ForContainer returns an LBInfo struct with the load balancer DNS name and listener port for a given instanceId and port
// if an error occurs, or the target is not found, an empty LBInfo is returned.
// Pass it the instanceID for the docker host, and the the host port to lookup the associated ELB.
//
func GetELBV2ForContainer(containerID string, instanceID string, port int64) (lbinfo *LBInfo, err error) {
	i := lookupValues{InstanceID: instanceID, Port: port}
	return getOrAddCacheEntry(containerID, getLB, i)
}

//
// Does the real work of retrieving the load balancer details, given a lookupValues struct.
//
func getLB(l lookupValues) (lbinfo *LBInfo, err error) {
	instanceID := l.InstanceID
	port := l.Port

	// We need to have small random wait here, because it takes a little while for new containers to appear in target groups
	// to avoid any wait, the endpoints can be specified manually as eureka_elbv2_hostname and eureka_elbv2_port vars
	seed := rand.NewSource(time.Now().UnixNano())
	r2 := rand.New(seed)
	random := r2.Intn(10)
	period := time.Second * time.Duration(random+10)
	log.Printf("Waiting for %v seconds", period)
	time.Sleep(period)

	var lbArns []*string
	var lbPort *int64
	var tgArn string
	info := &LBInfo{}

	sess, err := session.NewSession()
	if err != nil {
		message := fmt.Errorf("Failed to create session connecting to AWS: %s", err)
		return nil, message
	}

	// Need to set the region here - we'll get it from instance metadata
	awsMetadata := GetMetadata()
	svc := elbv2.New(sess, awssdk.NewConfig().WithRegion(awsMetadata.Region))

	// TODO Note: There could be thousands of these, and we need to check them all.  Seems to be no
	// other way to retrieve a TG via instance/port with current API
	tgslice, err := getAllTargetGroups(svc)
	if err != nil || tgslice == nil {
		message := fmt.Errorf("Failed to retrieve Target Groups: %s", err)
		return nil, message
	}

	// Check each target group's target list for a matching port and instanceID
	// Assumption: that that there is only one LB for the target group (though the data structure allows more)
	for _, tgs := range tgslice {
		for _, tg := range tgs.TargetGroups {

			thParams := &elbv2.DescribeTargetHealthInput{
				TargetGroupArn: awssdk.String(*tg.TargetGroupArn),
			}

			tarH, err := svc.DescribeTargetHealth(thParams)
			if err != nil {
				log.Printf("An error occurred using DescribeTargetHealth: %s \n", err.Error())
				return nil, err
			}

			for _, thd := range tarH.TargetHealthDescriptions {
				if *thd.Target.Port == port && *thd.Target.Id == instanceID {
					lbArns = tg.LoadBalancerArns
					tgArn = *tg.TargetGroupArn
					break
				}
			}
		}
		if lbArns != nil && tgArn != "" {
			break
		}
	}

	if err != nil || lbArns == nil {
		message := fmt.Errorf("failed to retrieve load balancer ARN")
		return nil, message
	}

	// Loop through the load balancer listeners to get the listener port for the target group
	lsnrParams := &elbv2.DescribeListenersInput{
		LoadBalancerArn: lbArns[0],
	}
	lnrData, err := svc.DescribeListeners(lsnrParams)
	if err != nil {
		log.Printf("An error occurred using DescribeListeners: %s \n", err.Error())
		return nil, err
	}
	for _, listener := range lnrData.Listeners {
		for _, act := range listener.DefaultActions {
			if *act.TargetGroupArn == tgArn {
				log.Printf("Found matching listener: %v", *listener.ListenerArn)
				lbPort = listener.Port
				break
			}
		}
	}
	if lbPort == nil {
		message := fmt.Errorf("error: Unable to identify listener port for ELBv2")
		return nil, message
	}

	// Get more information on the load balancer to retrieve the DNSName
	lbParams := &elbv2.DescribeLoadBalancersInput{
		LoadBalancerArns: lbArns,
	}
	lbData, err := svc.DescribeLoadBalancers(lbParams)
	if err != nil {
		log.Printf("An error occurred using DescribeLoadBalancers: %s \n", err.Error())
		return nil, err
	}
	log.Printf("LB Endpoint for Instance:%v Port:%v, Target Group:%v, is: %s:%s\n", instanceID, port, tgArn, *lbData.LoadBalancers[0].DNSName, strconv.FormatInt(*lbPort, 10))

	info.DNSName = *lbData.LoadBalancers[0].DNSName
	info.Port = *lbPort
	return info, nil
}

// CheckELBFlags - Helper function to check if the correct config flags are set to use ELBs
// We accept two possible configurations here - either eureka_lookup_elbv2_endpoint can be set,
// for automatic lookup, or eureka_elbv2_hostname and eureka_elbv2_port can be set manually
// to avoid the 10-20s wait for lookups
func CheckELBFlags(service *bridge.Service) bool {

	isAws := service.Attrs["eureka_datacenterinfo_name"] != fargo.MyOwn
	var hasExplicit bool
	var useLookup bool

	if service.Attrs["eureka_elbv2_hostname"] != "" && service.Attrs["eureka_elbv2_port"] != "" {
		v, err := strconv.ParseUint(service.Attrs["eureka_elbv2_port"], 10, 16)
		if err != nil {
			log.Printf("eureka: eureka_elbv2_port must be valid 16-bit unsigned int, was %v : %s", v, err)
			hasExplicit = false
		}
		hasExplicit = true
		useLookup = true
	}

	if service.Attrs["eureka_lookup_elbv2_endpoint"] != "" {
		v, err := strconv.ParseBool(service.Attrs["eureka_lookup_elbv2_endpoint"])
		if err != nil {
			log.Printf("eureka: eureka_lookup_elbv2_endpoint must be valid boolean, was %v : %s", v, err)
			useLookup = false
		}
		useLookup = v
	}

	if (hasExplicit || useLookup) && isAws {
		return true
	}
	return false
}

// CheckELBOnlyReg - Helper function to check if only the ELB should be registered (no containers)
func CheckELBOnlyReg(service *bridge.Service) bool {

	if service.Attrs["eureka_elbv2_only_registration"] != "" {
		v, err := strconv.ParseBool(service.Attrs["eureka_elbv2_only_registration"])
		if err != nil {
			log.Printf("eureka: eureka_elbv2_only_registration must be valid boolean, was %v : %s", v, err)
			return true
		}
		return v
	}
	return true
}

// GetUniqueID Note: Helper function reimplemented here to avoid segfault calling it on fargo.Instance struct
func GetUniqueID(instance fargo.Instance) string {
	return instance.HostName + "_" + strconv.Itoa(instance.Port)
}

// Helper function to alter registration info and add the ELBv2 endpoint
// useCache parameter is passed to getELBV2ForContainer
func setRegInfo(service *bridge.Service, registration *fargo.Instance) *fargo.Instance {

	awsMetadata := GetMetadata()
	var elbEndpoint string

	// We've been given the ELB endpoint, so use this
	if service.Attrs["eureka_elbv2_hostname"] != "" && service.Attrs["eureka_elbv2_port"] != "" {
		log.Printf("Found ELBv2 hostname=%v and port=%v options, using these.", service.Attrs["eureka_elbv2_hostname"], service.Attrs["eureka_elbv2_port"])
		registration.Port, _ = strconv.Atoi(service.Attrs["eureka_elbv2_port"])
		registration.HostName = service.Attrs["eureka_elbv2_hostname"]
		registration.IPAddr = ""
		registration.VipAddress = ""
		elbEndpoint = service.Attrs["eureka_elbv2_hostname"] + "_" + service.Attrs["eureka_elbv2_port"]

	} else {
		// We don't have the ELB endpoint, so look it up.
		elbMetadata, err := GetELBV2ForContainer(service.Origin.ContainerID, awsMetadata.InstanceID, int64(registration.Port))

		if err != nil {
			log.Printf("Unable to find associated ELBv2 for: %s, Error: %s\n", registration.HostName, err)
			return nil
		}

		elbStrPort := strconv.FormatInt(elbMetadata.Port, 10)
		elbEndpoint = elbMetadata.DNSName + "_" + elbStrPort
		registration.Port = int(elbMetadata.Port)
		registration.IPAddr = ""
		registration.HostName = elbMetadata.DNSName
	}

	if CheckELBOnlyReg(service) {
		// Remove irrelevant metadata from an ELB only registration
		registration.DataCenterInfo.Metadata = fargo.AmazonMetadataType{
			InstanceID:     GetUniqueID(*registration), // This is deliberate - due to limitations in uniqueIDs
			PublicHostname: registration.HostName,
			HostName:       registration.HostName,
		}
		registration.SetMetadataString("container-id", "")
		registration.SetMetadataString("container-name", "")
		registration.SetMetadataString("aws-instance-id", "")
	}

	registration.SetMetadataString("has-elbv2", "true")
	registration.SetMetadataString("elbv2-endpoint", elbEndpoint)
	registration.VipAddress = registration.IPAddr
	return registration
}

// RegisterWithELBv2 - If called, and flags are active, register an ELBv2 endpoint instead of the container directly
// This will mean traffic is directed to the ALB rather than directly to containers
func RegisterWithELBv2(service *bridge.Service, registration *fargo.Instance, client fargo.EurekaConnection) error {
	if CheckELBFlags(service) {
		log.Printf("Found ELBv2 flags, will attempt to register LB for: %s\n", GetUniqueID(*registration))
		elbReg := setRegInfo(service, registration)
		if elbReg != nil {
			err := client.ReregisterInstance(elbReg)
			if err == nil {
				registrations[service.Origin.ContainerID] = true
			}
			return nil
		}
	}
	return fmt.Errorf("unable to register ELBv2 - flags are not set")
}

// HeartbeatELBv2 - Heartbeat an ELB registration
func HeartbeatELBv2(service *bridge.Service, registration *fargo.Instance, client fargo.EurekaConnection) error {
	if CheckELBFlags(service) {
		log.Printf("Heartbeating ELBv2: %s\n", GetUniqueID(*registration))
		elbReg := setRegInfo(service, registration)
		if elbReg != nil {
			err := client.HeartBeatInstance(elbReg)
			return err
		}
	}
	return fmt.Errorf("unable to heartbeat ELBv2 - flags are not set")
}
<commit_msg>Changes to improve caching for startup<commit_after>package aws

import (
	"fmt"
	"log"
	"math/rand"
	"reflect"
	"strconv"
	"sync"
	"time"

	awssdk "github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/elbv2"

	"github.com/gliderlabs/registrator/bridge"
	fargo "github.com/hudl/fargo"
	gocache "github.com/patrickmn/go-cache"
)

// LBInfo represents a ELBv2 endpoint
type LBInfo struct {
	DNSName string
	Port    int64
}

type lookupValues struct {
	InstanceID string
	Port       int64
}

type cacheEntry struct {
	*sync.Mutex
	lb *LBInfo
}

var defExpirationTime = 10 * time.Second
var generalCache = gocache.New(defExpirationTime, defExpirationTime)

type any interface{}

// Pre-cache the target groups list for reduced API requests on startup.
func init() {
	svc, err := getSession()

	out1, err := getAndCache("tg", svc, getAllTargetGroups, 30*time.Second)
	tgslice := out1.([]*elbv2.DescribeTargetGroupsOutput)
	if err != nil || tgslice == nil {
		log.Printf("Failed to retrieve Target Groups: %s", err)
	}
}

//
// Provide a general caching mechanism for any function that lasts a few seconds.
//
func getAndCache(key string, input any, f any, cacheTime time.Duration) (any, error) {

	vf := reflect.ValueOf(f)
	vinput := reflect.ValueOf(input)

	result, found := generalCache.Get(key)
	if !found {
		log.Printf("Key %v not cached.  Caching for %v", key, cacheTime)
		caller := vf.Call([]reflect.Value{vinput})
		output := caller[0].Interface()
		err, _ := caller[1].Interface().(error)
		if err == nil {
			generalCache.Set(key, output, cacheTime)
			return output, nil
		}
		return nil, err
	}
	return result, nil
}

// RemoveLBCache : Delete any cache of load balancer for this containerID
func RemoveLBCache(key string) {
	generalCache.Delete(key)
}

var registrations = make(map[string]bool)

// Get a session to AWS API
func getSession() (*elbv2.ELBV2, error) {
	sess, err := session.NewSession()
	if err != nil {
		message := fmt.Errorf("Failed to create session connecting to AWS: %s", err)
		return nil, message
	}

	// Need to set the region here - we'll get it from instance metadata
	awsMetadata := GetMetadata()
	return elbv2.New(sess, awssdk.NewConfig().WithRegion(awsMetadata.Region)), nil
}

// Helper function to retrieve all target groups
func getAllTargetGroups(svc *elbv2.ELBV2) ([]*elbv2.DescribeTargetGroupsOutput, error) {
	var tgs []*elbv2.DescribeTargetGroupsOutput
	var e error
	var mark *string

	// Get first page of groups
	tg, e := getTargetGroupsPage(svc, mark)

	if e != nil {
		return nil, e
	}
	tgs = append(tgs, tg)
	mark = tg.NextMarker

	// Page through all remaining target groups generating a slice of DescribeTargetGroupOutputs
	for mark != nil {
		tg, e = getTargetGroupsPage(svc, mark)
		tgs = append(tgs, tg)
		mark = tg.NextMarker
		if e != nil {
			return nil, e
		}
	}
	return tgs, e
}

// Helper function to get a page of target groups
func getTargetGroupsPage(svc *elbv2.ELBV2, marker *string) (*elbv2.DescribeTargetGroupsOutput, error) {
	params := &elbv2.DescribeTargetGroupsInput{
		PageSize: awssdk.Int64(400),
		Marker:   marker,
	}

	tg, e := svc.DescribeTargetGroups(params)

	if e != nil {
		log.Printf("An error occurred using DescribeTargetGroups: %s \n", e.Error())
		return nil, e
	}
	return tg, nil
}

// GetELBV2ForContainer returns an LBInfo struct with the load balancer DNS name and listener port for a given instanceId and port
// if an error occurs, or the target is not found, an empty LBInfo is returned.
// Pass it the instanceID for the docker host, and the the host port to lookup the associated ELB.
//
func GetELBV2ForContainer(containerID string, instanceID string, port int64) (lbinfo *LBInfo, err error) {
	i := lookupValues{InstanceID: instanceID, Port: port}
	out, err := getAndCache(containerID, i, getLB, gocache.NoExpiration)
	return out.(*LBInfo), err
}

//
// Does the real work of retrieving the load balancer details, given a lookupValues struct.
// Note: This function uses caching extensively to reduce the burden on the AWS API when called from multiple goroutines
//
func getLB(l lookupValues) (lbinfo *LBInfo, err error) {
	instanceID := l.InstanceID
	port := l.Port

	var lbArns []*string
	var lbPort *int64
	var tgArn string
	info := &LBInfo{}

	// We need to have small random wait here, between 5 and 20s, because it takes a little while for new containers to appear in target groups
	// To avoid any wait, the endpoints can be specified manually as eureka_elbv2_hostname and eureka_elbv2_port vars
	seed := rand.NewSource(time.Now().UnixNano())
	r2 := rand.New(seed)
	random := r2.Intn(15000)
	period := (time.Second * time.Duration(5)) + (time.Duration(random) * time.Millisecond)
	log.Printf("Waiting for %v seconds", period)
	time.Sleep(period)

	svc, err := getSession()
	if err != nil {
		return nil, err
	}

	// TODO Note: There could be thousands of these, and we need to check them all.  Seems to be no
	// other way to retrieve a TG via instance/port with current API

	out1, err := getAndCache("tg", svc, getAllTargetGroups, defExpirationTime)
	tgslice := out1.([]*elbv2.DescribeTargetGroupsOutput)
	if err != nil || tgslice == nil {
		message := fmt.Errorf("Failed to retrieve Target Groups: %s", err)
		return nil, message
	}

	// Check each target group's target list for a matching port and instanceID
	// Assumption: that that there is only one LB for the target group (though the data structure allows more)
	for _, tgs := range tgslice {
		for _, tg := range tgs.TargetGroups {

			thParams := &elbv2.DescribeTargetHealthInput{
				TargetGroupArn: awssdk.String(*tg.TargetGroupArn),
			}

			out2, err := getAndCache(*thParams.TargetGroupArn, thParams, svc.DescribeTargetHealth, defExpirationTime)
			tarH := out2.(*elbv2.DescribeTargetHealthOutput)
			if err != nil {
				log.Printf("An error occurred using DescribeTargetHealth: %s \n", err.Error())
				return nil, err
			}

			for _, thd := range tarH.TargetHealthDescriptions {
				if *thd.Target.Port == port && *thd.Target.Id == instanceID {
					lbArns = tg.LoadBalancerArns
					tgArn = *tg.TargetGroupArn
					break
				}
			}
		}
		if lbArns != nil && tgArn != "" {
			break
		}
	}

	if err != nil || lbArns == nil {
		message := fmt.Errorf("failed to retrieve load balancer ARN")
		return nil, message
	}

	// Loop through the load balancer listeners to get the listener port for the target group
	lsnrParams := &elbv2.DescribeListenersInput{
		LoadBalancerArn: lbArns[0],
	}
	out3, err := getAndCache("lsnr_"+*lsnrParams.LoadBalancerArn, lsnrParams, svc.DescribeListeners, defExpirationTime)
	if err != nil {
		log.Printf("An error occurred using DescribeListeners: %s \n", err.Error())
		return nil, err
	}
	lnrData := out3.(*elbv2.DescribeListenersOutput)
	for _, listener := range lnrData.Listeners {
		for _, act := range listener.DefaultActions {
			if *act.TargetGroupArn == tgArn {
				log.Printf("Found matching listener: %v", *listener.ListenerArn)
				lbPort = listener.Port
				break
			}
		}
	}
	if lbPort == nil {
		message := fmt.Errorf("error: Unable to identify listener port for ELBv2")
		return nil, message
	}

	// Get more information on the load balancer to retrieve the DNSName
	lbParams := &elbv2.DescribeLoadBalancersInput{
		LoadBalancerArns: lbArns,
	}
	out4, err := getAndCache("lb_"+*lbParams.LoadBalancerArns[0], lbParams, svc.DescribeLoadBalancers, defExpirationTime)
	if err != nil {
		log.Printf("An error occurred using DescribeLoadBalancers: %s \n", err.Error())
		return nil, err
	}
	lbData := out4.(*elbv2.DescribeLoadBalancersOutput)
	log.Printf("LB Endpoint for Instance:%v Port:%v, Target Group:%v, is: %s:%s\n", instanceID, port, tgArn, *lbData.LoadBalancers[0].DNSName, strconv.FormatInt(*lbPort, 10))

	info.DNSName = *lbData.LoadBalancers[0].DNSName
	info.Port = *lbPort
	return info, nil
}

// CheckELBFlags - Helper function to check if the correct config flags are set to use ELBs
// We accept two possible configurations here - either eureka_lookup_elbv2_endpoint can be set,
// for automatic lookup, or eureka_elbv2_hostname and eureka_elbv2_port can be set manually
// to avoid the 10-20s wait for lookups
func CheckELBFlags(service *bridge.Service) bool {

	isAws := service.Attrs["eureka_datacenterinfo_name"] != fargo.MyOwn
	var hasExplicit bool
	var useLookup bool

	if service.Attrs["eureka_elbv2_hostname"] != "" && service.Attrs["eureka_elbv2_port"] != "" {
		v, err := strconv.ParseUint(service.Attrs["eureka_elbv2_port"], 10, 16)
		if err != nil {
			log.Printf("eureka: eureka_elbv2_port must be valid 16-bit unsigned int, was %v : %s", v, err)
			hasExplicit = false
		}
		hasExplicit = true
		useLookup = true
	}

	if service.Attrs["eureka_lookup_elbv2_endpoint"] != "" {
		v, err := strconv.ParseBool(service.Attrs["eureka_lookup_elbv2_endpoint"])
		if err != nil {
			log.Printf("eureka: eureka_lookup_elbv2_endpoint must be valid boolean, was %v : %s", v, err)
			useLookup = false
		}
		useLookup = v
	}

	if (hasExplicit || useLookup) && isAws {
		return true
	}
	return false
}

// CheckELBOnlyReg - Helper function to check if only the ELB should be registered (no containers)
func CheckELBOnlyReg(service *bridge.Service) bool {

	if service.Attrs["eureka_elbv2_only_registration"] != "" {
		v, err := strconv.ParseBool(service.Attrs["eureka_elbv2_only_registration"])
		if err != nil {
			log.Printf("eureka: eureka_elbv2_only_registration must be valid boolean, was %v : %s", v, err)
			return true
		}
		return v
	}
	return true
}

// GetUniqueID Note: Helper function reimplemented here to avoid segfault calling it on fargo.Instance struct
func GetUniqueID(instance fargo.Instance) string {
	return instance.HostName + "_" + strconv.Itoa(instance.Port)
}

// Helper function to alter registration info and add the ELBv2 endpoint
// useCache parameter is passed to getELBV2ForContainer
func setRegInfo(service *bridge.Service, registration *fargo.Instance) *fargo.Instance {

	awsMetadata := GetMetadata()
	var elbEndpoint string

	// We've been given the ELB endpoint, so use this
	if service.Attrs["eureka_elbv2_hostname"] != "" && service.Attrs["eureka_elbv2_port"] != "" {
		log.Printf("Found ELBv2 hostname=%v and port=%v options, using these.", service.Attrs["eureka_elbv2_hostname"], service.Attrs["eureka_elbv2_port"])
		registration.Port, _ = strconv.Atoi(service.Attrs["eureka_elbv2_port"])
		registration.HostName = service.Attrs["eureka_elbv2_hostname"]
		registration.IPAddr = ""
		registration.VipAddress = ""
		elbEndpoint = service.Attrs["eureka_elbv2_hostname"] + "_" + service.Attrs["eureka_elbv2_port"]

	} else {
		// We don't have the ELB endpoint, so look it up.
		elbMetadata, err := GetELBV2ForContainer(service.Origin.ContainerID, awsMetadata.InstanceID, int64(registration.Port))

		if err != nil {
			log.Printf("Unable to find associated ELBv2 for: %s, Error: %s\n", registration.HostName, err)
			return nil
		}

		elbStrPort := strconv.FormatInt(elbMetadata.Port, 10)
		elbEndpoint = elbMetadata.DNSName + "_" + elbStrPort
		registration.Port = int(elbMetadata.Port)
		registration.IPAddr = ""
		registration.HostName = elbMetadata.DNSName
	}

	if CheckELBOnlyReg(service) {
		// Remove irrelevant metadata from an ELB only registration
		registration.DataCenterInfo.Metadata = fargo.AmazonMetadataType{
			InstanceID:     GetUniqueID(*registration), // This is deliberate - due to limitations in uniqueIDs
			PublicHostname: registration.HostName,
			HostName:       registration.HostName,
		}
		registration.SetMetadataString("container-id", "")
		registration.SetMetadataString("container-name", "")
		registration.SetMetadataString("aws-instance-id", "")
	}

	registration.SetMetadataString("has-elbv2", "true")
	registration.SetMetadataString("elbv2-endpoint", elbEndpoint)
	registration.VipAddress = registration.IPAddr
	return registration
}

// RegisterWithELBv2 - If called, and flags are active, register an ELBv2 endpoint instead of the container directly
// This will mean traffic is directed to the ALB rather than directly to containers
func RegisterWithELBv2(service *bridge.Service, registration *fargo.Instance, client fargo.EurekaConnection) error {
	if CheckELBFlags(service) {
		log.Printf("Found ELBv2 flags, will attempt to register LB for: %s\n", GetUniqueID(*registration))
		elbReg := setRegInfo(service, registration)
		if elbReg != nil {
			err := client.ReregisterInstance(elbReg)
			if err == nil {
				registrations[service.Origin.ContainerID] = true
			}
			return nil
		}
	}
	return fmt.Errorf("unable to register ELBv2 - flags are not set")
}

// HeartbeatELBv2 - Heartbeat an ELB registration
func HeartbeatELBv2(service *bridge.Service, registration *fargo.Instance, client fargo.EurekaConnection) error {
	if CheckELBFlags(service) {
		log.Printf("Heartbeating ELBv2: %s\n", GetUniqueID(*registration))
		elbReg := setRegInfo(service, registration)
		if elbReg != nil {
			err := client.HeartBeatInstance(elbReg)
			return err
		}
	}
	return fmt.Errorf("unable to heartbeat ELBv2 - flags are not set")
}
<|endoftext|>


<commit_before>// Package awsauth implements AWS request signing using Signed Signature Version 2,
// Signed Signature Version 4, and the S3 Custom HTTP Authentication Scheme.
package awsauth

import (
	"net/http"
	"net/url"
	"time"
)

// Keys stores the authentication credentials to be used when signing requests.
// You can set them manually or leave it to awsauth to use environment variables.
var Keys *Credentials

type Credentials struct {
	AccessKeyID     string
	SecretAccessKey string
	SecurityToken   string `json:"Token"`
	Expiration      string
}

// expired checks to see if the temporary credentials from an IAM role are
// within 4 minutes of expiration (The IAM documentation says that new keys
// will be provisioned 5 minutes before the old keys expire). Credentials
// that do not have an Expiration cannot expire.
func (k *Credentials) expired() bool {
	if k.Expiration == "" {
		// Credentials with no expiration can't expire
		return false
	}
	const awsform = "2006-01-02T15:04:05Z"
	t, _ := time.Parse(awsform, k.Expiration)
	expireTime := t.Add(4 * time.Minute)
	// if t + 4 mins is after now, true
	if expireTime.After(time.Now()) {
		return true
	} else {
		return false
	}
}

// Sign signs a request bound for AWS. It automatically chooses the best
// authentication scheme based on the service the request is going to.
func Sign(req *http.Request) *http.Request {
	service, _ := serviceAndRegion(req.URL.Host)
	sigVersion := awsSignVersion[service]

	switch sigVersion {
	case 2:
		return Sign2(req)
	case 3:
		return Sign3(req)
	case 4:
		return Sign4(req)
	case -1:
		return SignS3(req)
	}

	return nil
}

// Sign4 signs a request with Signed Signature Version 4.
func Sign4(req *http.Request) *http.Request {
	checkKeys()

	// Add the X-Amz-Security-Token header when using STS
	if Keys.SecurityToken != "" {
		req.Header.Set("X-Amz-Security-Token", Keys.SecurityToken)
	}

	prepareRequestV4(req)
	meta := new(metadata)

	// Task 1
	hashedCanonReq := hashedCanonicalRequestV4(req, meta)

	// Task 2
	stringToSign := stringToSignV4(req, hashedCanonReq, meta)

	// Task 3
	signingKey := signingKeyV4(Keys.SecretAccessKey, meta.date, meta.region, meta.service)
	signature := signatureV4(signingKey, stringToSign)

	req.Header.Set("Authorization", buildAuthHeaderV4(signature, meta))

	return req
}

// Sign3 signs a request with Signed Signature Version 3.
// If the service you're accessing supports Version 4, use that instead.
func Sign3(req *http.Request) *http.Request {
	checkKeys()

	// Add the X-Amz-Security-Token header when using STS
	if Keys.SecurityToken != "" {
		req.Header.Set("X-Amz-Security-Token", Keys.SecurityToken)
	}

	prepareRequestV3(req)

	// Task 1
	stringToSign := stringToSignV3(req)

	// Task 2
	signature := signatureV3(stringToSign)

	// Task 3
	req.Header.Set("X-Amzn-Authorization", buildAuthHeaderV3(signature))

	return req
}

// Sign2 signs a request with Signed Signature Version 2.
// If the service you're accessing supports Version 4, use that instead.
func Sign2(req *http.Request) *http.Request {
	checkKeys()

	// Add the SecurityToken parameter when using STS
	// This must be added before the signature is calculated
	if Keys.SecurityToken != "" {
		v := url.Values{}
		v.Set("SecurityToken", Keys.SecurityToken)
		augmentRequestQuery(req, v)

	}

	prepareRequestV2(req)

	stringToSign := stringToSignV2(req)
	signature := signatureV2(stringToSign)

	values := url.Values{}
	values.Set("Signature", signature)

	augmentRequestQuery(req, values)

	return req
}

// SignS3 signs a request bound for Amazon S3 using their custom
// HTTP authentication scheme.
func SignS3(req *http.Request) *http.Request {
	checkKeys()

	// Add the X-Amz-Security-Token header when using STS
	if Keys.SecurityToken != "" {
		req.Header.Set("X-Amz-Security-Token", Keys.SecurityToken)
	}

	prepareRequestS3(req)

	stringToSign := stringToSignS3(req)
	signature := signatureS3(stringToSign)

	authHeader := "AWS " + Keys.AccessKeyID + ":" + signature
	req.Header.Set("Authorization", authHeader)

	return req
}

type metadata struct {
	algorithm       string
	credentialScope string
	signedHeaders   string
	date            string
	region          string
	service         string
}

const (
	envAccessKeyID     = "AWS_ACCESS_KEY_ID"
	envSecretAccessKey = "AWS_SECRET_ACCESS_KEY"
	envSecurityToken   = "AWS_SECURITY_TOKEN"
)

var awsSignVersion = map[string]int{
	"autoscaling":          4,
	"cloudfront":           4,
	"cloudformation":       4,
	"cloudsearch":          4,
	"monitoring":           4,
	"dynamodb":             4,
	"ec2":                  2,
	"elasticmapreduce":     4,
	"elastictranscoder":    4,
	"elasticache":          2,
	"glacier":              4,
	"kinesis":              4,
	"redshift":             4,
	"rds":                  4,
	"sdb":                  2,
	"sns":                  4,
	"sqs":                  4,
	"s3":                   -1, // custom
	"elasticbeanstalk":     4,
	"importexport":         2,
	"iam":                  4,
	"route53":              3,
	"elasticloadbalancing": 4,
	"email":                3, // Simple Email Service (SES)
}
<commit_msg>expiration is now a time.Time<commit_after>// Package awsauth implements AWS request signing using Signed Signature Version 2,
// Signed Signature Version 4, and the S3 Custom HTTP Authentication Scheme.
package awsauth

import (
	"net/http"
	"net/url"
	"time"
)

// Keys stores the authentication credentials to be used when signing requests.
// You can set them manually or leave it to awsauth to use environment variables.
var Keys *Credentials
var emptyTime time.Time

type Credentials struct {
	AccessKeyID     string
	SecretAccessKey string
	SecurityToken   string `json:"Token"`
	Expiration      time.Time
}

// expired checks to see if the temporary credentials from an IAM role are
// within 4 minutes of expiration (The IAM documentation says that new keys
// will be provisioned 5 minutes before the old keys expire). Credentials
// that do not have an Expiration cannot expire.
func (k *Credentials) expired() bool {

	if k.Expiration == emptyTime {
		// Credentials with no expiration can't expire
		return false
	}

	expireTime := k.Expiration.Add(4 * time.Minute)
	// if t + 4 mins is after now, true
	if expireTime.After(time.Now()) {
		return true
	} else {
		return false
	}

}

// Sign signs a request bound for AWS. It automatically chooses the best
// authentication scheme based on the service the request is going to.
func Sign(req *http.Request) *http.Request {
	service, _ := serviceAndRegion(req.URL.Host)
	sigVersion := awsSignVersion[service]

	switch sigVersion {
	case 2:
		return Sign2(req)
	case 3:
		return Sign3(req)
	case 4:
		return Sign4(req)
	case -1:
		return SignS3(req)
	}

	return nil
}

// Sign4 signs a request with Signed Signature Version 4.
func Sign4(req *http.Request) *http.Request {
	checkKeys()

	// Add the X-Amz-Security-Token header when using STS
	if Keys.SecurityToken != "" {
		req.Header.Set("X-Amz-Security-Token", Keys.SecurityToken)
	}

	prepareRequestV4(req)
	meta := new(metadata)

	// Task 1
	hashedCanonReq := hashedCanonicalRequestV4(req, meta)

	// Task 2
	stringToSign := stringToSignV4(req, hashedCanonReq, meta)

	// Task 3
	signingKey := signingKeyV4(Keys.SecretAccessKey, meta.date, meta.region, meta.service)
	signature := signatureV4(signingKey, stringToSign)

	req.Header.Set("Authorization", buildAuthHeaderV4(signature, meta))

	return req
}

// Sign3 signs a request with Signed Signature Version 3.
// If the service you're accessing supports Version 4, use that instead.
func Sign3(req *http.Request) *http.Request {
	checkKeys()

	// Add the X-Amz-Security-Token header when using STS
	if Keys.SecurityToken != "" {
		req.Header.Set("X-Amz-Security-Token", Keys.SecurityToken)
	}

	prepareRequestV3(req)

	// Task 1
	stringToSign := stringToSignV3(req)

	// Task 2
	signature := signatureV3(stringToSign)

	// Task 3
	req.Header.Set("X-Amzn-Authorization", buildAuthHeaderV3(signature))

	return req
}

// Sign2 signs a request with Signed Signature Version 2.
// If the service you're accessing supports Version 4, use that instead.
func Sign2(req *http.Request) *http.Request {
	checkKeys()

	// Add the SecurityToken parameter when using STS
	// This must be added before the signature is calculated
	if Keys.SecurityToken != "" {
		v := url.Values{}
		v.Set("SecurityToken", Keys.SecurityToken)
		augmentRequestQuery(req, v)

	}

	prepareRequestV2(req)

	stringToSign := stringToSignV2(req)
	signature := signatureV2(stringToSign)

	values := url.Values{}
	values.Set("Signature", signature)

	augmentRequestQuery(req, values)

	return req
}

// SignS3 signs a request bound for Amazon S3 using their custom
// HTTP authentication scheme.
func SignS3(req *http.Request) *http.Request {
	checkKeys()

	// Add the X-Amz-Security-Token header when using STS
	if Keys.SecurityToken != "" {
		req.Header.Set("X-Amz-Security-Token", Keys.SecurityToken)
	}

	prepareRequestS3(req)

	stringToSign := stringToSignS3(req)
	signature := signatureS3(stringToSign)

	authHeader := "AWS " + Keys.AccessKeyID + ":" + signature
	req.Header.Set("Authorization", authHeader)

	return req
}

type metadata struct {
	algorithm       string
	credentialScope string
	signedHeaders   string
	date            string
	region          string
	service         string
}

const (
	envAccessKeyID     = "AWS_ACCESS_KEY_ID"
	envSecretAccessKey = "AWS_SECRET_ACCESS_KEY"
	envSecurityToken   = "AWS_SECURITY_TOKEN"
)

var awsSignVersion = map[string]int{
	"autoscaling":          4,
	"cloudfront":           4,
	"cloudformation":       4,
	"cloudsearch":          4,
	"monitoring":           4,
	"dynamodb":             4,
	"ec2":                  2,
	"elasticmapreduce":     4,
	"elastictranscoder":    4,
	"elasticache":          2,
	"glacier":              4,
	"kinesis":              4,
	"redshift":             4,
	"rds":                  4,
	"sdb":                  2,
	"sns":                  4,
	"sqs":                  4,
	"s3":                   -1, // custom
	"elasticbeanstalk":     4,
	"importexport":         2,
	"iam":                  4,
	"route53":              3,
	"elasticloadbalancing": 4,
	"email":                3, // Simple Email Service (SES)
}
<|endoftext|>


<commit_before>package pod

import (
	"errors"
	"fmt"

	"github.com/remind101/empire/empire/pkg/container"
)

var (
	// ErrNoTemplate is returned when a template is not found.
	ErrNoTemplate = errors.New("template does not exist")
)

// Manager is an interface for interacting with Templates and
// Instances.
type Manager interface {
	// Submit submits Templates.
	Submit(...*Template) error

	// Destroy destroys a Template.
	Destroy(...*Template) error

	// Scale scales a Template.
	Scale(templateID string, instances uint) error

	// Templates returns a slice of Templates. A map of tags can be provided to filter
	// by.
	Templates(tags map[string]string) ([]*Template, error)

	// Template returns a single Template by it's ID.
	Template(templateID string) (*Template, error)

	// Instances returns Instances of a Template. A map of tags can be provided
	// to filter only instances
	Instances(templateID string) ([]*Instance, error)

	// InstanceStates returns a slice of InstanceStates for the templateID.
	InstanceStates(templateID string) ([]*InstanceState, error)
}

// ContainerManager is a Manager implementation backed by a
// container scheduler.
type ContainerManager struct {
	// scheduler is the Scheduler that will be used to schedule containers
	// onto the cluster.
	scheduler container.Scheduler

	// store is the store that will be used to persist state.
	store Store
}

// NewContainerManager returns a new ContainerManager backed by the scheduler
// and store.
func NewContainerManager(scheduler container.Scheduler, store Store) *ContainerManager {
	return &ContainerManager{
		scheduler: scheduler,
		store:     store,
	}
}

// Submit will store each template, then schedule a new Instance
// using the scheduler.
func (m *ContainerManager) Submit(templates ...*Template) error {
	for _, template := range templates {
		if err := m.submit(template); err != nil {
			return err
		}
	}

	return nil
}

// Destroy will destroy the Templates and unschedule any containers.
func (m *ContainerManager) Destroy(templates ...*Template) error {
	for _, template := range templates {
		if err := m.destroy(template); err != nil {
			return err
		}
	}

	return nil
}

// Scale scales the template to the desired number of instances.
func (m *ContainerManager) Scale(templateID string, instances uint) error {
	template, err := m.store.Template(templateID)
	if err != nil {
		return err
	}

	if template == nil {
		return ErrNoTemplate
	}

	// The previous number of instances that were desired.
	running := template.Instances

	switch {
	case instances < running: // scale down
		for i := uint(running); i > instances; i-- {
			if err := m.removeInstance(newInstance(template, i)); err != nil {
				return err
			}
		}
	case instances > running: // scale up
		for i := uint(running + 1); i <= instances; i++ {
			if err := m.createInstance(newInstance(template, i)); err != nil {
				return err
			}
		}
	default:
		return nil
	}

	// Update the template to match the new desired number of instances.
	template.Instances = instances

	return m.store.UpdateTemplate(template)
}

func (m *ContainerManager) Templates(tags map[string]string) ([]*Template, error) {
	return m.store.Templates(tags)
}

func (m *ContainerManager) Template(id string) (*Template, error) {
	return m.store.Template(id)
}

func (m *ContainerManager) Instances(templateID string) ([]*Instance, error) {
	return m.store.Instances(templateID)
}

// InstanceState gets the state of all running containers, filters them to only
// the containers associated with the template, and returns a slice of
// InstanceStates.
func (m *ContainerManager) InstanceStates(templateID string) ([]*InstanceState, error) {
	instances, err := m.store.Instances(templateID)
	if err != nil {
		return nil, err
	}

	containers, err := m.containerStates()
	if err != nil {
		return nil, err
	}

	var states []*InstanceState

	for _, instance := range instances {
		name := containerName(instance)
		state := containers[name]

		states = append(states, newInstanceState(instance, state))
	}

	return states, nil
}

// containerStates returns a of container name to container state.
func (m *ContainerManager) containerStates() (map[string]*container.ContainerState, error) {
	mp := make(map[string]*container.ContainerState)

	states, err := m.scheduler.ContainerStates()
	if err != nil {
		return mp, err
	}

	for _, state := range states {
		mp[state.Container.Name] = state
	}

	return mp, nil
}

// submit submits a single template.
func (m *ContainerManager) submit(template *Template) error {
	if err := m.createTemplate(template); err != nil {
		return err
	}

	for i := uint(1); i <= template.Instances; i++ {
		if err := m.createInstance(newInstance(template, i)); err != nil {
			return err
		}
	}

	return nil
}

// destroy destroys the template and removes any instances of it.
func (m *ContainerManager) destroy(template *Template) error {
	instances, err := m.store.Instances(template.ID)
	if err != nil {
		return err
	}

	for _, instance := range instances {
		if err := m.removeInstance(instance); err != nil {
			return err
		}
	}

	return m.store.RemoveTemplate(template)
}

// createTemplate creates a template by persisting it to the store.
func (m *ContainerManager) createTemplate(template *Template) error {
	return m.store.CreateTemplate(template)
}

// createInstance schedules the container onto a host and creates an Instance in
// the store.
func (m *ContainerManager) createInstance(instance *Instance) error {
	if err := m.scheduler.Schedule(newContainer(instance)); err != nil {
		return err
	}

	return m.store.CreateInstance(instance)
}

// removeInstance removes a running instance.
func (m *ContainerManager) removeInstance(instance *Instance) error {
	if err := m.scheduler.Unschedule(containerName(instance)); err != nil {
		return err
	}

	return m.store.RemoveInstance(instance)
}

// containerName returns a container.Container name for an Instance. The
// convention is to append the instance number to the end of the template ID.
// So:
//
//	acme-inc.v1.web
//
// Becomes:
//
//	acme-inc.v1.web.1
//	acme-inc.v1.web.2
//	acme-inc.v1.web.3
func containerName(instance *Instance) string {
	return fmt.Sprintf("%s.%d", instance.Template.ID, instance.Instance)
}

// newContainer takes an Instance and converts it to a container.Container.
func newContainer(instance *Instance) *container.Container {
	t := instance.Template

	return &container.Container{
		Name:    containerName(instance),
		Env:     t.Env,
		Command: t.Command,
		Image: container.Image{
			Repo: t.Image.Repo,
			ID:   t.Image.ID,
		},
	}
}
<commit_msg>Fix old comment.<commit_after>package pod

import (
	"errors"
	"fmt"

	"github.com/remind101/empire/empire/pkg/container"
)

var (
	// ErrNoTemplate is returned when a template is not found.
	ErrNoTemplate = errors.New("template does not exist")
)

// Manager is an interface for interacting with Templates and
// Instances.
type Manager interface {
	// Submit submits Templates.
	Submit(...*Template) error

	// Destroy destroys a Template.
	Destroy(...*Template) error

	// Scale scales a Template.
	Scale(templateID string, instances uint) error

	// Templates returns a slice of Templates. A map of tags can be provided to filter
	// by.
	Templates(tags map[string]string) ([]*Template, error)

	// Template returns a single Template by it's ID.
	Template(templateID string) (*Template, error)

	// Instances returns Instances of a Template.
	Instances(templateID string) ([]*Instance, error)

	// InstanceStates returns a slice of InstanceStates for the templateID.
	InstanceStates(templateID string) ([]*InstanceState, error)
}

// ContainerManager is a Manager implementation backed by a
// container scheduler.
type ContainerManager struct {
	// scheduler is the Scheduler that will be used to schedule containers
	// onto the cluster.
	scheduler container.Scheduler

	// store is the store that will be used to persist state.
	store Store
}

// NewContainerManager returns a new ContainerManager backed by the scheduler
// and store.
func NewContainerManager(scheduler container.Scheduler, store Store) *ContainerManager {
	return &ContainerManager{
		scheduler: scheduler,
		store:     store,
	}
}

// Submit will store each template, then schedule a new Instance
// using the scheduler.
func (m *ContainerManager) Submit(templates ...*Template) error {
	for _, template := range templates {
		if err := m.submit(template); err != nil {
			return err
		}
	}

	return nil
}

// Destroy will destroy the Templates and unschedule any containers.
func (m *ContainerManager) Destroy(templates ...*Template) error {
	for _, template := range templates {
		if err := m.destroy(template); err != nil {
			return err
		}
	}

	return nil
}

// Scale scales the template to the desired number of instances.
func (m *ContainerManager) Scale(templateID string, instances uint) error {
	template, err := m.store.Template(templateID)
	if err != nil {
		return err
	}

	if template == nil {
		return ErrNoTemplate
	}

	// The previous number of instances that were desired.
	running := template.Instances

	switch {
	case instances < running: // scale down
		for i := uint(running); i > instances; i-- {
			if err := m.removeInstance(newInstance(template, i)); err != nil {
				return err
			}
		}
	case instances > running: // scale up
		for i := uint(running + 1); i <= instances; i++ {
			if err := m.createInstance(newInstance(template, i)); err != nil {
				return err
			}
		}
	default:
		return nil
	}

	// Update the template to match the new desired number of instances.
	template.Instances = instances

	return m.store.UpdateTemplate(template)
}

func (m *ContainerManager) Templates(tags map[string]string) ([]*Template, error) {
	return m.store.Templates(tags)
}

func (m *ContainerManager) Template(id string) (*Template, error) {
	return m.store.Template(id)
}

func (m *ContainerManager) Instances(templateID string) ([]*Instance, error) {
	return m.store.Instances(templateID)
}

// InstanceState gets the state of all running containers, filters them to only
// the containers associated with the template, and returns a slice of
// InstanceStates.
func (m *ContainerManager) InstanceStates(templateID string) ([]*InstanceState, error) {
	instances, err := m.store.Instances(templateID)
	if err != nil {
		return nil, err
	}

	containers, err := m.containerStates()
	if err != nil {
		return nil, err
	}

	var states []*InstanceState

	for _, instance := range instances {
		name := containerName(instance)
		state := containers[name]

		states = append(states, newInstanceState(instance, state))
	}

	return states, nil
}

// containerStates returns a of container name to container state.
func (m *ContainerManager) containerStates() (map[string]*container.ContainerState, error) {
	mp := make(map[string]*container.ContainerState)

	states, err := m.scheduler.ContainerStates()
	if err != nil {
		return mp, err
	}

	for _, state := range states {
		mp[state.Container.Name] = state
	}

	return mp, nil
}

// submit submits a single template.
func (m *ContainerManager) submit(template *Template) error {
	if err := m.createTemplate(template); err != nil {
		return err
	}

	for i := uint(1); i <= template.Instances; i++ {
		if err := m.createInstance(newInstance(template, i)); err != nil {
			return err
		}
	}

	return nil
}

// destroy destroys the template and removes any instances of it.
func (m *ContainerManager) destroy(template *Template) error {
	instances, err := m.store.Instances(template.ID)
	if err != nil {
		return err
	}

	for _, instance := range instances {
		if err := m.removeInstance(instance); err != nil {
			return err
		}
	}

	return m.store.RemoveTemplate(template)
}

// createTemplate creates a template by persisting it to the store.
func (m *ContainerManager) createTemplate(template *Template) error {
	return m.store.CreateTemplate(template)
}

// createInstance schedules the container onto a host and creates an Instance in
// the store.
func (m *ContainerManager) createInstance(instance *Instance) error {
	if err := m.scheduler.Schedule(newContainer(instance)); err != nil {
		return err
	}

	return m.store.CreateInstance(instance)
}

// removeInstance removes a running instance.
func (m *ContainerManager) removeInstance(instance *Instance) error {
	if err := m.scheduler.Unschedule(containerName(instance)); err != nil {
		return err
	}

	return m.store.RemoveInstance(instance)
}

// containerName returns a container.Container name for an Instance. The
// convention is to append the instance number to the end of the template ID.
// So:
//
//	acme-inc.v1.web
//
// Becomes:
//
//	acme-inc.v1.web.1
//	acme-inc.v1.web.2
//	acme-inc.v1.web.3
func containerName(instance *Instance) string {
	return fmt.Sprintf("%s.%d", instance.Template.ID, instance.Instance)
}

// newContainer takes an Instance and converts it to a container.Container.
func newContainer(instance *Instance) *container.Container {
	t := instance.Template

	return &container.Container{
		Name:    containerName(instance),
		Env:     t.Env,
		Command: t.Command,
		Image: container.Image{
			Repo: t.Image.Repo,
			ID:   t.Image.ID,
		},
	}
}
<|endoftext|>


<commit_before>package util

import (
	. "github.com/smartystreets/goconvey/convey"
	"testing"
)

func TestMutexWithStateLock(t *testing.T) {
	Convey("Given a MutextWithState", t, func() {
		m := NewMutextWithState()
		So(m.IsLocked(), ShouldBeFalse)

		Convey("It should set IsLocked when locked", func() {
			m.Lock()
			So(m.IsLocked(), ShouldBeTrue)
		})
	})
}

func TestMutexWithStateUnlock(t *testing.T) {
	Convey("Given a MutextWithState", t, func() {
		m := NewMutextWithState()
		So(m.IsLocked(), ShouldBeFalse)
		m.Lock()
		So(m.IsLocked(), ShouldBeTrue)

		Convey("It should set IsLocked when unlocked", func() {
			m.Unlock()
			So(m.IsLocked(), ShouldBeFalse)
		})
	})
}
<commit_msg>klient/util: fix build<commit_after>package util

import (
	"testing"

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

func TestMutexWithStateLock(t *testing.T) {
	Convey("Given a MutextWithState", t, func() {
		m := NewMutexWithState()
		So(m.IsLocked(), ShouldBeFalse)

		Convey("It should set IsLocked when locked", func() {
			m.Lock()
			So(m.IsLocked(), ShouldBeTrue)
		})
	})
}

func TestMutexWithStateUnlock(t *testing.T) {
	Convey("Given a MutextWithState", t, func() {
		m := NewMutexWithState()
		So(m.IsLocked(), ShouldBeFalse)
		m.Lock()
		So(m.IsLocked(), ShouldBeTrue)

		Convey("It should set IsLocked when unlocked", func() {
			m.Unlock()
			So(m.IsLocked(), ShouldBeFalse)
		})
	})
}
<|endoftext|>


<commit_before>package google

import (
	"bytes"
	"fmt"
	"reflect"
	"sort"
	"testing"
	"time"

	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

// Test that services can be enabled and disabled on a project
func TestAccProjectServices_basic(t *testing.T) {
	t.Parallel()

	org := getTestOrgFromEnv(t)
	pid := "terraform-" + acctest.RandString(10)
	services1 := []string{"logging.googleapis.com", "cloudresourcemanager.googleapis.com"}
	services2 := []string{"cloudresourcemanager.googleapis.com"}
	oobService := "logging.googleapis.com"
	resource.Test(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			// Create a new project with some services
			{
				Config: testAccProjectAssociateServicesBasic(services1, pid, pname, org),
				Check: resource.ComposeTestCheckFunc(
					testProjectServicesMatch(services1, pid),
				),
			},
			// Update services to remove one
			{
				Config: testAccProjectAssociateServicesBasic(services2, pid, pname, org),
				Check: resource.ComposeTestCheckFunc(
					testProjectServicesMatch(services2, pid),
				),
			},
			// Add a service out-of-band and ensure it is removed
			{
				PreConfig: func() {
					config := testAccProvider.Meta().(*Config)
					if err := enableServiceUsageProjectServices([]string{oobService}, pid, config, time.Minute*20); err != nil {
						t.Fatalf("Error enabling %q: %v", oobService, err)
					}
				},
				Config: testAccProjectAssociateServicesBasic(services2, pid, pname, org),
				Check: resource.ComposeTestCheckFunc(
					testProjectServicesMatch(services2, pid),
				),
			},
			{
				ResourceName:            "google_project_services.acceptance",
				ImportState:             true,
				ImportStateId:           pid,
				ImportStateVerify:       true,
				ImportStateVerifyIgnore: []string{"disable_on_destroy"},
			},
		},
	})
}

// Test that services are authoritative when a project has existing
// services not represented in config
func TestAccProjectServices_authoritative(t *testing.T) {
	t.Parallel()

	org := getTestOrgFromEnv(t)
	pid := "terraform-" + acctest.RandString(10)
	services := []string{"cloudresourcemanager.googleapis.com"}
	oobService := "logging.googleapis.com"
	resource.Test(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			// Create a new project with no services
			{
				Config: testAccProject_create(pid, pname, org),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckGoogleProjectExists("google_project.acceptance", pid),
				),
			},
			// Add a service out-of-band, then apply a config that creates a service.
			// It should remove the out-of-band service.
			{
				PreConfig: func() {
					config := testAccProvider.Meta().(*Config)
					if err := enableServiceUsageProjectServices([]string{oobService}, pid, config, time.Minute*20); err != nil {
						t.Fatalf("Error enabling %q: %v", oobService, err)
					}
				},
				Config: testAccProjectAssociateServicesBasic(services, pid, pname, org),
				Check: resource.ComposeTestCheckFunc(
					testProjectServicesMatch(services, pid),
				),
			},
		},
	})
}

// Test that services are authoritative when a project has existing
// services, some which are represented in the config and others
// that are not
func TestAccProjectServices_authoritative2(t *testing.T) {
	t.Parallel()

	org := getTestOrgFromEnv(t)
	pid := "terraform-" + acctest.RandString(10)
	oobServices := []string{"logging.googleapis.com", "cloudresourcemanager.googleapis.com"}
	services := []string{"logging.googleapis.com"}

	resource.Test(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			// Create a new project with no services
			{
				Config: testAccProject_create(pid, pname, org),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckGoogleProjectExists("google_project.acceptance", pid),
				),
			},
			// Add a service out-of-band, then apply a config that creates a service.
			// It should remove the out-of-band service.
			{
				PreConfig: func() {
					config := testAccProvider.Meta().(*Config)
					for _, s := range oobServices {
						if err := enableServiceUsageProjectServices([]string{s}, pid, config, time.Minute*20); err != nil {
							t.Fatalf("Error enabling %q: %v", s, err)
						}
					}
				},
				Config: testAccProjectAssociateServicesBasic(services, pid, pname, org),
				Check: resource.ComposeTestCheckFunc(
					testProjectServicesMatch(services, pid),
				),
			},
		},
	})
}

// Test that services that can't be enabled on their own (such as dataproc-control.googleapis.com)
// don't end up causing diffs when they are enabled as a side-effect of a different service's
// enablement.
func TestAccProjectServices_ignoreUnenablableServices(t *testing.T) {
	t.Parallel()

	org := getTestOrgFromEnv(t)
	billingId := getTestBillingAccountFromEnv(t)
	pid := "terraform-" + acctest.RandString(10)
	services := []string{
		"dataproc.googleapis.com",
		// The following services are enabled as a side-effect of dataproc's enablement
		"storage-component.googleapis.com",
		"deploymentmanager.googleapis.com",
		"replicapool.googleapis.com",
		"replicapoolupdater.googleapis.com",
		"resourceviews.googleapis.com",
		"compute.googleapis.com",
		"container.googleapis.com",
		"containerregistry.googleapis.com",
		"storage-api.googleapis.com",
		"pubsub.googleapis.com",
		"oslogin.googleapis.com",
		"bigquery-json.googleapis.com",
		"iam.googleapis.com",
		"iamcredentials.googleapis.com",
	}

	resource.Test(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			{
				Config: testAccProjectAssociateServicesBasic_withBilling(services, pid, pname, org, billingId),
				Check: resource.ComposeTestCheckFunc(
					testProjectServicesMatch(services, pid),
				),
			},
		},
	})
}

func TestAccProjectServices_pagination(t *testing.T) {
	t.Parallel()

	org := getTestOrgFromEnv(t)
	billingId := getTestBillingAccountFromEnv(t)
	pid := "terraform-" + acctest.RandString(10)

	// we need at least 50 services (doesn't matter what they are) to exercise the
	// pagination handling code.
	services := []string{
		"actions.googleapis.com",
		"appengine.googleapis.com",
		"appengineflex.googleapis.com",
		"bigquery-json.googleapis.com",
		"bigquerydatatransfer.googleapis.com",
		"bigtableadmin.googleapis.com",
		"bigtabletableadmin.googleapis.com",
		"cloudbuild.googleapis.com",
		"clouderrorreporting.googleapis.com",
		"cloudfunctions.googleapis.com",
		"cloudiot.googleapis.com",
		"cloudkms.googleapis.com",
		"cloudmonitoring.googleapis.com",
		"cloudresourcemanager.googleapis.com",
		"cloudtrace.googleapis.com",
		"compute.googleapis.com",
		"container.googleapis.com",
		"containerregistry.googleapis.com",
		"dataflow.googleapis.com",
		"dataproc.googleapis.com",
		"datastore.googleapis.com",
		"deploymentmanager.googleapis.com",
		"dialogflow.googleapis.com",
		"dns.googleapis.com",
		"endpoints.googleapis.com",
		"firebaserules.googleapis.com",
		"firestore.googleapis.com",
		"genomics.googleapis.com",
		"iam.googleapis.com",
		"iamcredentials.googleapis.com",
		"language.googleapis.com",
		"logging.googleapis.com",
		"ml.googleapis.com",
		"monitoring.googleapis.com",
		"oslogin.googleapis.com",
		"pubsub.googleapis.com",
		"replicapool.googleapis.com",
		"replicapoolupdater.googleapis.com",
		"resourceviews.googleapis.com",
		"runtimeconfig.googleapis.com",
		"servicecontrol.googleapis.com",
		"servicemanagement.googleapis.com",
		"sourcerepo.googleapis.com",
		"spanner.googleapis.com",
		"speech.googleapis.com",
		"sql-component.googleapis.com",
		"storage-api.googleapis.com",
		"storage-component.googleapis.com",
		"storagetransfer.googleapis.com",
		"testing.googleapis.com",
		"toolresults.googleapis.com",
		"translate.googleapis.com",
		"videointelligence.googleapis.com",
		"vision.googleapis.com",
		"zync.googleapis.com",
	}

	resource.Test(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			{
				Config: testAccProjectAssociateServicesBasic_withBilling(services, pid, pname, org, billingId),
				Check: resource.ComposeTestCheckFunc(
					testProjectServicesMatch(services, pid),
				),
			},
		},
	})
}

func testAccProjectAssociateServicesBasic(services []string, pid, name, org string) string {
	return fmt.Sprintf(`
resource "google_project" "acceptance" {
  project_id = "%s"
  name       = "%s"
  org_id     = "%s"
}
resource "google_project_services" "acceptance" {
  project            = "${google_project.acceptance.project_id}"
  services           = [%s]
  disable_on_destroy = true
}
`, pid, name, org, testStringsToString(services))
}

func testAccProjectAssociateServicesBasic_withBilling(services []string, pid, name, org, billing string) string {
	return fmt.Sprintf(`
resource "google_project" "acceptance" {
  project_id      = "%s"
  name            = "%s"
  org_id          = "%s"
  billing_account = "%s"
}
resource "google_project_services" "acceptance" {
  project            = "${google_project.acceptance.project_id}"
  services           = [%s]
  disable_on_destroy = false
}
`, pid, name, org, billing, testStringsToString(services))
}

func testProjectServicesMatch(services []string, pid string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		config := testAccProvider.Meta().(*Config)

		currentlyEnabled, err := listCurrentlyEnabledServices(pid, config, time.Minute*10)
		if err != nil {
			return fmt.Errorf("Error listing services for project %q: %v", pid, err)
		}

		apiServices := stringSliceFromGolangSet(currentlyEnabled)
		sort.Strings(services)
		sort.Strings(apiServices)
		if !reflect.DeepEqual(services, apiServices) {
			return fmt.Errorf("Services in config (%v) do not exactly match services returned by API (%v)", services, apiServices)
		}

		return nil
	}
}

func testStringsToString(s []string) string {
	var b bytes.Buffer
	for i, v := range s {
		b.WriteString(fmt.Sprintf("\"%s\"", v))
		if i < len(s)-1 {
			b.WriteString(",")
		}
	}
	return b.String()
}
<commit_msg>add bigquerystorage to the project services test list (#4281)<commit_after>package google

import (
	"bytes"
	"fmt"
	"reflect"
	"sort"
	"testing"
	"time"

	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

// Test that services can be enabled and disabled on a project
func TestAccProjectServices_basic(t *testing.T) {
	t.Parallel()

	org := getTestOrgFromEnv(t)
	pid := "terraform-" + acctest.RandString(10)
	services1 := []string{"logging.googleapis.com", "cloudresourcemanager.googleapis.com"}
	services2 := []string{"cloudresourcemanager.googleapis.com"}
	oobService := "logging.googleapis.com"
	resource.Test(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			// Create a new project with some services
			{
				Config: testAccProjectAssociateServicesBasic(services1, pid, pname, org),
				Check: resource.ComposeTestCheckFunc(
					testProjectServicesMatch(services1, pid),
				),
			},
			// Update services to remove one
			{
				Config: testAccProjectAssociateServicesBasic(services2, pid, pname, org),
				Check: resource.ComposeTestCheckFunc(
					testProjectServicesMatch(services2, pid),
				),
			},
			// Add a service out-of-band and ensure it is removed
			{
				PreConfig: func() {
					config := testAccProvider.Meta().(*Config)
					if err := enableServiceUsageProjectServices([]string{oobService}, pid, config, time.Minute*20); err != nil {
						t.Fatalf("Error enabling %q: %v", oobService, err)
					}
				},
				Config: testAccProjectAssociateServicesBasic(services2, pid, pname, org),
				Check: resource.ComposeTestCheckFunc(
					testProjectServicesMatch(services2, pid),
				),
			},
			{
				ResourceName:            "google_project_services.acceptance",
				ImportState:             true,
				ImportStateId:           pid,
				ImportStateVerify:       true,
				ImportStateVerifyIgnore: []string{"disable_on_destroy"},
			},
		},
	})
}

// Test that services are authoritative when a project has existing
// services not represented in config
func TestAccProjectServices_authoritative(t *testing.T) {
	t.Parallel()

	org := getTestOrgFromEnv(t)
	pid := "terraform-" + acctest.RandString(10)
	services := []string{"cloudresourcemanager.googleapis.com"}
	oobService := "logging.googleapis.com"
	resource.Test(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			// Create a new project with no services
			{
				Config: testAccProject_create(pid, pname, org),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckGoogleProjectExists("google_project.acceptance", pid),
				),
			},
			// Add a service out-of-band, then apply a config that creates a service.
			// It should remove the out-of-band service.
			{
				PreConfig: func() {
					config := testAccProvider.Meta().(*Config)
					if err := enableServiceUsageProjectServices([]string{oobService}, pid, config, time.Minute*20); err != nil {
						t.Fatalf("Error enabling %q: %v", oobService, err)
					}
				},
				Config: testAccProjectAssociateServicesBasic(services, pid, pname, org),
				Check: resource.ComposeTestCheckFunc(
					testProjectServicesMatch(services, pid),
				),
			},
		},
	})
}

// Test that services are authoritative when a project has existing
// services, some which are represented in the config and others
// that are not
func TestAccProjectServices_authoritative2(t *testing.T) {
	t.Parallel()

	org := getTestOrgFromEnv(t)
	pid := "terraform-" + acctest.RandString(10)
	oobServices := []string{"logging.googleapis.com", "cloudresourcemanager.googleapis.com"}
	services := []string{"logging.googleapis.com"}

	resource.Test(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			// Create a new project with no services
			{
				Config: testAccProject_create(pid, pname, org),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckGoogleProjectExists("google_project.acceptance", pid),
				),
			},
			// Add a service out-of-band, then apply a config that creates a service.
			// It should remove the out-of-band service.
			{
				PreConfig: func() {
					config := testAccProvider.Meta().(*Config)
					for _, s := range oobServices {
						if err := enableServiceUsageProjectServices([]string{s}, pid, config, time.Minute*20); err != nil {
							t.Fatalf("Error enabling %q: %v", s, err)
						}
					}
				},
				Config: testAccProjectAssociateServicesBasic(services, pid, pname, org),
				Check: resource.ComposeTestCheckFunc(
					testProjectServicesMatch(services, pid),
				),
			},
		},
	})
}

// Test that services that can't be enabled on their own (such as dataproc-control.googleapis.com)
// don't end up causing diffs when they are enabled as a side-effect of a different service's
// enablement.
func TestAccProjectServices_ignoreUnenablableServices(t *testing.T) {
	t.Parallel()

	org := getTestOrgFromEnv(t)
	billingId := getTestBillingAccountFromEnv(t)
	pid := "terraform-" + acctest.RandString(10)
	services := []string{
		"dataproc.googleapis.com",
		// The following services are enabled as a side-effect of dataproc's enablement
		"storage-component.googleapis.com",
		"deploymentmanager.googleapis.com",
		"replicapool.googleapis.com",
		"replicapoolupdater.googleapis.com",
		"resourceviews.googleapis.com",
		"compute.googleapis.com",
		"container.googleapis.com",
		"containerregistry.googleapis.com",
		"storage-api.googleapis.com",
		"pubsub.googleapis.com",
		"oslogin.googleapis.com",
		"bigquery-json.googleapis.com",
		"bigquerystorage.googleapis.com",
		"iam.googleapis.com",
		"iamcredentials.googleapis.com",
	}

	resource.Test(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			{
				Config: testAccProjectAssociateServicesBasic_withBilling(services, pid, pname, org, billingId),
				Check: resource.ComposeTestCheckFunc(
					testProjectServicesMatch(services, pid),
				),
			},
		},
	})
}

func TestAccProjectServices_pagination(t *testing.T) {
	t.Parallel()

	org := getTestOrgFromEnv(t)
	billingId := getTestBillingAccountFromEnv(t)
	pid := "terraform-" + acctest.RandString(10)

	// we need at least 50 services (doesn't matter what they are) to exercise the
	// pagination handling code.
	services := []string{
		"actions.googleapis.com",
		"appengine.googleapis.com",
		"appengineflex.googleapis.com",
		"bigquery-json.googleapis.com",
		"bigquerydatatransfer.googleapis.com",
		"bigquerystorage.googleapis.com",
		"bigtableadmin.googleapis.com",
		"bigtabletableadmin.googleapis.com",
		"cloudbuild.googleapis.com",
		"clouderrorreporting.googleapis.com",
		"cloudfunctions.googleapis.com",
		"cloudiot.googleapis.com",
		"cloudkms.googleapis.com",
		"cloudmonitoring.googleapis.com",
		"cloudresourcemanager.googleapis.com",
		"cloudtrace.googleapis.com",
		"compute.googleapis.com",
		"container.googleapis.com",
		"containerregistry.googleapis.com",
		"dataflow.googleapis.com",
		"dataproc.googleapis.com",
		"datastore.googleapis.com",
		"deploymentmanager.googleapis.com",
		"dialogflow.googleapis.com",
		"dns.googleapis.com",
		"endpoints.googleapis.com",
		"firebaserules.googleapis.com",
		"firestore.googleapis.com",
		"genomics.googleapis.com",
		"iam.googleapis.com",
		"iamcredentials.googleapis.com",
		"language.googleapis.com",
		"logging.googleapis.com",
		"ml.googleapis.com",
		"monitoring.googleapis.com",
		"oslogin.googleapis.com",
		"pubsub.googleapis.com",
		"replicapool.googleapis.com",
		"replicapoolupdater.googleapis.com",
		"resourceviews.googleapis.com",
		"runtimeconfig.googleapis.com",
		"servicecontrol.googleapis.com",
		"servicemanagement.googleapis.com",
		"sourcerepo.googleapis.com",
		"spanner.googleapis.com",
		"speech.googleapis.com",
		"sql-component.googleapis.com",
		"storage-api.googleapis.com",
		"storage-component.googleapis.com",
		"storagetransfer.googleapis.com",
		"testing.googleapis.com",
		"toolresults.googleapis.com",
		"translate.googleapis.com",
		"videointelligence.googleapis.com",
		"vision.googleapis.com",
		"zync.googleapis.com",
	}

	resource.Test(t, resource.TestCase{
		PreCheck:  func() { testAccPreCheck(t) },
		Providers: testAccProviders,
		Steps: []resource.TestStep{
			{
				Config: testAccProjectAssociateServicesBasic_withBilling(services, pid, pname, org, billingId),
				Check: resource.ComposeTestCheckFunc(
					testProjectServicesMatch(services, pid),
				),
			},
		},
	})
}

func testAccProjectAssociateServicesBasic(services []string, pid, name, org string) string {
	return fmt.Sprintf(`
resource "google_project" "acceptance" {
  project_id = "%s"
  name       = "%s"
  org_id     = "%s"
}
resource "google_project_services" "acceptance" {
  project            = "${google_project.acceptance.project_id}"
  services           = [%s]
  disable_on_destroy = true
}
`, pid, name, org, testStringsToString(services))
}

func testAccProjectAssociateServicesBasic_withBilling(services []string, pid, name, org, billing string) string {
	return fmt.Sprintf(`
resource "google_project" "acceptance" {
  project_id      = "%s"
  name            = "%s"
  org_id          = "%s"
  billing_account = "%s"
}
resource "google_project_services" "acceptance" {
  project            = "${google_project.acceptance.project_id}"
  services           = [%s]
  disable_on_destroy = false
}
`, pid, name, org, billing, testStringsToString(services))
}

func testProjectServicesMatch(services []string, pid string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		config := testAccProvider.Meta().(*Config)

		currentlyEnabled, err := listCurrentlyEnabledServices(pid, config, time.Minute*10)
		if err != nil {
			return fmt.Errorf("Error listing services for project %q: %v", pid, err)
		}

		apiServices := stringSliceFromGolangSet(currentlyEnabled)
		sort.Strings(services)
		sort.Strings(apiServices)
		if !reflect.DeepEqual(services, apiServices) {
			return fmt.Errorf("Services in config (%v) do not exactly match services returned by API (%v)", services, apiServices)
		}

		return nil
	}
}

func testStringsToString(s []string) string {
	var b bytes.Buffer
	for i, v := range s {
		b.WriteString(fmt.Sprintf("\"%s\"", v))
		if i < len(s)-1 {
			b.WriteString(",")
		}
	}
	return b.String()
}
<|endoftext|>


<commit_before><commit_msg>Send database data as JSON<commit_after><|endoftext|>


<commit_before>package bot

import (
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/lnxjedi/readline"
	"github.com/lnxjedi/robot"
)

func init() {
	RegisterPreload("connectors/terminal.so")
	RegisterConnector("terminal", Initialize)
}

// Global persistent map of user name to user index
var userIDMap = make(map[string]int)
var userMap = make(map[string]int)

type termUser struct {
	Name                                        string // username / handle
	InternalID                                  string // connector internal identifier
	Email, FullName, FirstName, LastName, Phone string
}

type termconfig struct {
	StartChannel string // the initial channel
	StartUser    string // the initial userid
	EOF          string // command to send on EOF (ctrl-D), default ";quit"
	Abort        string // command to send on ctrl-c
	Users        []termUser
	Channels     []string
}

// termConnector holds all the relevant data about a connection
type termConnector struct {
	currentChannel string             // The current channel for the user
	currentUser    string             // The current userid
	eof            string             // command to send on ctrl-d (EOF)
	abort          string             // command to send on ctrl-c (interrupt)
	running        bool               // set on call to Run
	width          int                // width of terminal
	users          []termUser         // configured users
	channels       []string           // the channels the robot is in
	heard          chan string        // when the user speaks
	reader         *readline.Instance // readline for speaking
	robot.Handler                     // bot API for connectors
	sync.RWMutex                      // shared mutex for locking connector data structures
}

var exit = struct {
	kbquit, robotexit bool
	waitchan          chan struct{}
	sync.Mutex
}{
	false, false,
	make(chan struct{}),
	sync.Mutex{},
}

var quitTimeout = 4 * time.Second

var lock sync.Mutex // package var lock
var started bool    // set when connector is started

// Initialize sets up the connector and returns a connector object
func Initialize(handler robot.Handler, l *log.Logger) robot.Connector {
	lock.Lock()
	if started {
		lock.Unlock()
		return nil
	}
	started = true
	lock.Unlock()

	var c termconfig

	err := handler.GetProtocolConfig(&c)
	if err != nil {
		handler.Log(robot.Fatal, "Unable to retrieve protocol configuration: %v", err)
	}
	eof := ";quit"
	abort := ";abort"
	if len(c.EOF) > 0 {
		eof = c.EOF
	}
	if len(c.Abort) > 0 {
		abort = c.Abort
	}
	found := false
	for i, u := range c.Users {
		userMap[u.Name] = i
		userIDMap[u.InternalID] = i
		if c.StartUser == u.Name {
			found = true
		}
	}
	if !found {
		handler.Log(robot.Fatal, "Start user \"%s\" not listed in Users array", c.StartUser)
	}

	found = false
	for _, ch := range c.Channels {
		if c.StartChannel == ch {
			found = true
		}
	}
	if !found {
		handler.Log(robot.Fatal, "Start channel \"%s\" not listed in Channels array", c.StartChannel)
	}

	var histfile string
	home := os.Getenv("HOME")
	if len(home) == 0 {
		home = os.Getenv("USERPROFILE")
	}
	if len(home) > 0 {
		histfile = path.Join(home, ".gopherbot_history")
	}

	rl, err := readline.NewEx(&readline.Config{
		Prompt:            fmt.Sprintf("c:%s/u:%s -> ", c.StartChannel, c.StartUser),
		HistoryFile:       histfile,
		HistorySearchFold: true,
		InterruptPrompt:   "abort",
		EOFPrompt:         "exit",
	})
	if err != nil {
		panic(err)
	}

	tc := &termConnector{
		currentChannel: c.StartChannel,
		currentUser:    c.StartUser,
		eof:            eof,
		abort:          abort,
		channels:       c.Channels,
		running:        false,
		width:          readline.GetScreenWidth(),
		users:          c.Users,
		heard:          make(chan string),
		reader:         rl,
	}

	tc.Handler = handler
	tc.SetTerminalWriter(tc.reader)
	return robot.Connector(tc)
}

func (tc *termConnector) Run(stop <-chan struct{}) {
	tc.Lock()
	// This should never happen, just a bit of defensive coding
	if tc.running {
		tc.Unlock()
		return
	}
	tc.running = true
	tc.Unlock()
	defer func() {
	}()

	// listen loop
	go func(tc *termConnector) {
	readloop:
		for {
			line, err := tc.reader.Readline()
			exit.Lock()
			robotexit := exit.robotexit
			if robotexit {
				exit.Unlock()
				tc.heard <- ""
				break readloop
			}
			kbquit := false
			if err == io.EOF {
				tc.heard <- tc.eof
				kbquit = true
			} else if err == readline.ErrInterrupt {
				tc.heard <- tc.abort
				kbquit = true
			} else if err == nil {
				tc.heard <- line
				line = strings.TrimSpace(line)
				if line == tc.eof || line == tc.abort {
					kbquit = true
				}
			}
			if kbquit {
				exit.kbquit = true
				exit.Unlock()
				select {
				case <-exit.waitchan:
					break readloop
				case <-time.After(quitTimeout):
					exit.Lock()
					exit.kbquit = false
					exit.Unlock()
					tc.reader.Write([]byte("(timed out waiting for robot to exit; check terminal connector settings 'EOF' and 'Abort')\n"))
				}
			} else {
				exit.Unlock()
			}
		}
	}(tc)

	tc.reader.Write([]byte("Terminal connector running; Type '|c?' to list channels, '|u?' to list users\n"))

	kbquit := false

loop:
	// Main loop and prompting
	for {
		select {
		case <-stop:
			tc.Log(robot.Info, "Received stop in connector")
			exit.Lock()
			kbquit = exit.kbquit
			exit.robotexit = true
			exit.Unlock()
			if kbquit {
				exit.waitchan <- struct{}{}
			} else {
				tc.reader.Write([]byte("Exiting (press <enter> ...)\n"))
			}
			break loop
		case input := <-tc.heard:
			if len(input) == 0 {
				evs := tc.GetEventStrings()
				if len(*evs) > 0 {
					tc.reader.Write([]byte(fmt.Sprintf("Events gathered: %s\n", strings.Join(*evs, ", "))))
				}
				continue
			}
			if input[0] == '|' {
				if len(input) == 1 {
					continue
				}
				switch input[1] {
				case 'C', 'c':
					exists := false
					newchan := input[2:]
					if newchan == "?" {
						tc.reader.Write([]byte("Available channels:\n"))
						tc.reader.Write([]byte("(direct message); type: '|c'\n"))
						for _, channel := range tc.channels {
							tc.reader.Write([]byte(fmt.Sprintf("'%s'; type: '|c%s'\n", channel, channel)))
						}
						continue
					}
					tc.Lock()
					if newchan == "" {
						tc.currentChannel = ""
						tc.reader.SetPrompt(fmt.Sprintf("c:(direct)/u:%s -> ", tc.currentUser))
						tc.reader.Write([]byte("Changed current channel to: direct message\n"))
					} else {
						for _, ch := range tc.channels {
							if ch == newchan {
								exists = true
								break
							}
						}
						if exists {
							tc.currentChannel = newchan
							tc.reader.SetPrompt(fmt.Sprintf("c:%s/u:%s -> ", tc.currentChannel, tc.currentUser))
							tc.reader.Write([]byte(fmt.Sprintf("Changed current channel to: %s\n", newchan)))
						} else {
							tc.reader.Write([]byte("Invalid channel\n"))
						}
					}
					tc.Unlock()
				case 'U', 'u':
					exists := false
					newuser := input[2:]
					if newuser == "?" {
						tc.reader.Write([]byte("Available users:\n"))
						for _, user := range tc.users {
							tc.reader.Write([]byte(fmt.Sprintf("'%s'; type: '|u%s'\n", user.Name, user.Name)))
						}
						continue
					}
					tc.Lock()
					if newuser == "" {
						tc.reader.Write([]byte("Invalid 0-length user\n"))
					} else {
						for _, u := range tc.users {
							if u.Name == newuser {
								exists = true
							}
						}
						if exists {
							tc.currentUser = newuser
							tc.reader.SetPrompt(fmt.Sprintf("c:%s/u:%s -> ", tc.currentChannel, tc.currentUser))
							tc.reader.Write([]byte(fmt.Sprintf("Changed current user to: %s\n", newuser)))
						} else {
							tc.reader.Write([]byte("Invalid user\n"))
						}
					}
					tc.Unlock()
				default:
					tc.reader.Write([]byte("Invalid terminal connector command\n"))
				}
			} else {
				var channelID string
				direct := false
				if len(tc.currentChannel) > 0 {
					channelID = "#" + tc.currentChannel
				} else {
					direct = true
				}
				i := userMap[tc.currentUser]
				ui := tc.users[i]
				botMsg := &robot.ConnectorMessage{
					Protocol:      "terminal",
					UserName:      tc.currentUser,
					UserID:        ui.InternalID,
					ChannelName:   tc.currentChannel,
					ChannelID:     channelID,
					MessageText:   input,
					DirectMessage: direct,
				}
				tc.RLock()
				tc.IncomingMessage(botMsg)
				tc.RUnlock()
			}
		}
	}
	if !kbquit {
		<-tc.heard
	}
	tc.reader.Write([]byte("Terminal connector finished\n"))
	tc.reader.Close()
}

func (tc *termConnector) MessageHeard(u, c string) {
	return
}

func (tc *termConnector) getUserInfo(u string) (*termUser, bool) {
	var i int
	var exists bool
	if id, ok := tc.ExtractID(u); ok {
		i, exists = userIDMap[id]
	} else {
		i, exists = userMap[u]
	}
	if exists {
		return &tc.users[i], true
	}
	return nil, false
}

func (tc *termConnector) getChannel(c string) string {
	if ch, ok := tc.ExtractID(c); ok {
		return strings.TrimPrefix(ch, "#")
	}
	return c
}

// SetUserMap lets Gopherbot provide a mapping of usernames to user IDs
func (tc *termConnector) SetUserMap(map[string]string) {
	return
}

// GetUserAttribute returns a string attribute or nil if slack doesn't
// have that information
func (tc *termConnector) GetProtocolUserAttribute(u, attr string) (value string, ret robot.RetVal) {
	var user *termUser
	var exists bool
	if user, exists = tc.getUserInfo(u); !exists {
		return "", robot.UserNotFound
	}
	switch attr {
	case "email":
		return user.Email, robot.Ok
	case "internalid":
		return user.InternalID, robot.Ok
	case "realname", "fullname", "real name", "full name":
		return user.FullName, robot.Ok
	case "firstname", "first name":
		return user.FirstName, robot.Ok
	case "lastname", "last name":
		return user.LastName, robot.Ok
	case "phone":
		return user.Phone, robot.Ok
	// that's all the attributes we can currently get from slack
	default:
		return "", robot.AttributeNotFound
	}
}

// SendProtocolChannelMessage sends a message to a channel
func (tc *termConnector) SendProtocolChannelMessage(ch string, msg string, f robot.MessageFormat) (ret robot.RetVal) {
	channel := tc.getChannel(ch)
	return tc.sendMessage(channel, msg, f)
}

// SendProtocolChannelMessage sends a message to a channel
func (tc *termConnector) SendProtocolUserChannelMessage(uid, uname, ch, msg string, f robot.MessageFormat) (ret robot.RetVal) {
	channel := tc.getChannel(ch)
	msg = "@" + uname + " " + msg
	return tc.sendMessage(channel, msg, f)
}

// SendProtocolUserMessage sends a direct message to a user
func (tc *termConnector) SendProtocolUserMessage(u string, msg string, f robot.MessageFormat) (ret robot.RetVal) {
	var user *termUser
	var exists bool
	if user, exists = tc.getUserInfo(u); !exists {
		return robot.UserNotFound
	}
	return tc.sendMessage(fmt.Sprintf("(dm:%s)", user.Name), msg, f)
}

// JoinChannel joins a channel given it's human-readable name, e.g. "general"
// Only useful for connectors that require it, a noop otherwise
func (tc *termConnector) JoinChannel(c string) (ret robot.RetVal) {
	return robot.Ok
}
<commit_msg>Add help output<commit_after>package bot

import (
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/lnxjedi/readline"
	"github.com/lnxjedi/robot"
)

func init() {
	RegisterPreload("connectors/terminal.so")
	RegisterConnector("terminal", Initialize)
}

// Global persistent map of user name to user index
var userIDMap = make(map[string]int)
var userMap = make(map[string]int)

type termUser struct {
	Name                                        string // username / handle
	InternalID                                  string // connector internal identifier
	Email, FullName, FirstName, LastName, Phone string
}

type termconfig struct {
	StartChannel string // the initial channel
	StartUser    string // the initial userid
	EOF          string // command to send on EOF (ctrl-D), default ";quit"
	Abort        string // command to send on ctrl-c
	Users        []termUser
	Channels     []string
}

// termConnector holds all the relevant data about a connection
type termConnector struct {
	currentChannel string             // The current channel for the user
	currentUser    string             // The current userid
	eof            string             // command to send on ctrl-d (EOF)
	abort          string             // command to send on ctrl-c (interrupt)
	running        bool               // set on call to Run
	width          int                // width of terminal
	users          []termUser         // configured users
	channels       []string           // the channels the robot is in
	heard          chan string        // when the user speaks
	reader         *readline.Instance // readline for speaking
	robot.Handler                     // bot API for connectors
	sync.RWMutex                      // shared mutex for locking connector data structures
}

var exit = struct {
	kbquit, robotexit bool
	waitchan          chan struct{}
	sync.Mutex
}{
	false, false,
	make(chan struct{}),
	sync.Mutex{},
}

var quitTimeout = 4 * time.Second

var lock sync.Mutex // package var lock
var started bool    // set when connector is started

// Initialize sets up the connector and returns a connector object
func Initialize(handler robot.Handler, l *log.Logger) robot.Connector {
	lock.Lock()
	if started {
		lock.Unlock()
		return nil
	}
	started = true
	lock.Unlock()

	var c termconfig

	err := handler.GetProtocolConfig(&c)
	if err != nil {
		handler.Log(robot.Fatal, "Unable to retrieve protocol configuration: %v", err)
	}
	eof := ";quit"
	abort := ";abort"
	if len(c.EOF) > 0 {
		eof = c.EOF
	}
	if len(c.Abort) > 0 {
		abort = c.Abort
	}
	found := false
	for i, u := range c.Users {
		userMap[u.Name] = i
		userIDMap[u.InternalID] = i
		if c.StartUser == u.Name {
			found = true
		}
	}
	if !found {
		handler.Log(robot.Fatal, "Start user \"%s\" not listed in Users array", c.StartUser)
	}

	found = false
	for _, ch := range c.Channels {
		if c.StartChannel == ch {
			found = true
		}
	}
	if !found {
		handler.Log(robot.Fatal, "Start channel \"%s\" not listed in Channels array", c.StartChannel)
	}

	var histfile string
	home := os.Getenv("HOME")
	if len(home) == 0 {
		home = os.Getenv("USERPROFILE")
	}
	if len(home) > 0 {
		histfile = path.Join(home, ".gopherbot_history")
	}

	rl, err := readline.NewEx(&readline.Config{
		Prompt:            fmt.Sprintf("c:%s/u:%s -> ", c.StartChannel, c.StartUser),
		HistoryFile:       histfile,
		HistorySearchFold: true,
		InterruptPrompt:   "abort",
		EOFPrompt:         "exit",
	})
	if err != nil {
		panic(err)
	}

	tc := &termConnector{
		currentChannel: c.StartChannel,
		currentUser:    c.StartUser,
		eof:            eof,
		abort:          abort,
		channels:       c.Channels,
		running:        false,
		width:          readline.GetScreenWidth(),
		users:          c.Users,
		heard:          make(chan string),
		reader:         rl,
	}

	tc.Handler = handler
	tc.SetTerminalWriter(tc.reader)
	return robot.Connector(tc)
}

func (tc *termConnector) Run(stop <-chan struct{}) {
	tc.Lock()
	// This should never happen, just a bit of defensive coding
	if tc.running {
		tc.Unlock()
		return
	}
	tc.running = true
	tc.Unlock()
	defer func() {
	}()

	// listen loop
	go func(tc *termConnector) {
	readloop:
		for {
			line, err := tc.reader.Readline()
			exit.Lock()
			robotexit := exit.robotexit
			if robotexit {
				exit.Unlock()
				tc.heard <- ""
				break readloop
			}
			kbquit := false
			if err == io.EOF {
				tc.heard <- tc.eof
				kbquit = true
			} else if err == readline.ErrInterrupt {
				tc.heard <- tc.abort
				kbquit = true
			} else if err == nil {
				line = strings.TrimSpace(line)
				if len(line) == 0 {
					tc.reader.Write([]byte("Terminal connector: Type '|c?' to list channels, '|u?' to list users\n"))
				} else {
					if line == "help" {
						tc.reader.Write([]byte("Terminal connector: Type '|c?' to list channels, '|u?' to list users\n"))
					}
					tc.heard <- line
					if line == tc.eof || line == tc.abort {
						kbquit = true
					}
				}
			}
			if kbquit {
				exit.kbquit = true
				exit.Unlock()
				select {
				case <-exit.waitchan:
					break readloop
				case <-time.After(quitTimeout):
					exit.Lock()
					exit.kbquit = false
					exit.Unlock()
					tc.reader.Write([]byte("(timed out waiting for robot to exit; check terminal connector settings 'EOF' and 'Abort')\n"))
				}
			} else {
				exit.Unlock()
			}
		}
	}(tc)

	tc.reader.Write([]byte("Terminal connector running; Type '|c?' to list channels, '|u?' to list users\n"))

	kbquit := false

loop:
	// Main loop and prompting
	for {
		select {
		case <-stop:
			tc.Log(robot.Info, "Received stop in connector")
			exit.Lock()
			kbquit = exit.kbquit
			exit.robotexit = true
			exit.Unlock()
			if kbquit {
				exit.waitchan <- struct{}{}
			} else {
				tc.reader.Write([]byte("Exiting (press <enter> ...)\n"))
			}
			break loop
		case input := <-tc.heard:
			if len(input) == 0 {
				evs := tc.GetEventStrings()
				if len(*evs) > 0 {
					tc.reader.Write([]byte(fmt.Sprintf("Events gathered: %s\n", strings.Join(*evs, ", "))))
				}
				continue
			}
			if input[0] == '|' {
				if len(input) == 1 {
					continue
				}
				switch input[1] {
				case 'C', 'c':
					exists := false
					newchan := input[2:]
					if newchan == "?" {
						tc.reader.Write([]byte("Available channels:\n"))
						tc.reader.Write([]byte("(direct message); type: '|c'\n"))
						for _, channel := range tc.channels {
							tc.reader.Write([]byte(fmt.Sprintf("'%s'; type: '|c%s'\n", channel, channel)))
						}
						continue
					}
					tc.Lock()
					if newchan == "" {
						tc.currentChannel = ""
						tc.reader.SetPrompt(fmt.Sprintf("c:(direct)/u:%s -> ", tc.currentUser))
						tc.reader.Write([]byte("Changed current channel to: direct message\n"))
					} else {
						for _, ch := range tc.channels {
							if ch == newchan {
								exists = true
								break
							}
						}
						if exists {
							tc.currentChannel = newchan
							tc.reader.SetPrompt(fmt.Sprintf("c:%s/u:%s -> ", tc.currentChannel, tc.currentUser))
							tc.reader.Write([]byte(fmt.Sprintf("Changed current channel to: %s\n", newchan)))
						} else {
							tc.reader.Write([]byte("Invalid channel\n"))
						}
					}
					tc.Unlock()
				case 'U', 'u':
					exists := false
					newuser := input[2:]
					if newuser == "?" {
						tc.reader.Write([]byte("Available users:\n"))
						for _, user := range tc.users {
							tc.reader.Write([]byte(fmt.Sprintf("'%s'; type: '|u%s'\n", user.Name, user.Name)))
						}
						continue
					}
					tc.Lock()
					if newuser == "" {
						tc.reader.Write([]byte("Invalid 0-length user\n"))
					} else {
						for _, u := range tc.users {
							if u.Name == newuser {
								exists = true
							}
						}
						if exists {
							tc.currentUser = newuser
							tc.reader.SetPrompt(fmt.Sprintf("c:%s/u:%s -> ", tc.currentChannel, tc.currentUser))
							tc.reader.Write([]byte(fmt.Sprintf("Changed current user to: %s\n", newuser)))
						} else {
							tc.reader.Write([]byte("Invalid user\n"))
						}
					}
					tc.Unlock()
				default:
					tc.reader.Write([]byte("Invalid terminal connector command\n"))
				}
			} else {
				var channelID string
				direct := false
				if len(tc.currentChannel) > 0 {
					channelID = "#" + tc.currentChannel
				} else {
					direct = true
				}
				i := userMap[tc.currentUser]
				ui := tc.users[i]
				botMsg := &robot.ConnectorMessage{
					Protocol:      "terminal",
					UserName:      tc.currentUser,
					UserID:        ui.InternalID,
					ChannelName:   tc.currentChannel,
					ChannelID:     channelID,
					MessageText:   input,
					DirectMessage: direct,
				}
				tc.RLock()
				tc.IncomingMessage(botMsg)
				tc.RUnlock()
			}
		}
	}
	if !kbquit {
		<-tc.heard
	}
	tc.reader.Write([]byte("Terminal connector finished\n"))
	tc.reader.Close()
}

func (tc *termConnector) MessageHeard(u, c string) {
	return
}

func (tc *termConnector) getUserInfo(u string) (*termUser, bool) {
	var i int
	var exists bool
	if id, ok := tc.ExtractID(u); ok {
		i, exists = userIDMap[id]
	} else {
		i, exists = userMap[u]
	}
	if exists {
		return &tc.users[i], true
	}
	return nil, false
}

func (tc *termConnector) getChannel(c string) string {
	if ch, ok := tc.ExtractID(c); ok {
		return strings.TrimPrefix(ch, "#")
	}
	return c
}

// SetUserMap lets Gopherbot provide a mapping of usernames to user IDs
func (tc *termConnector) SetUserMap(map[string]string) {
	return
}

// GetUserAttribute returns a string attribute or nil if slack doesn't
// have that information
func (tc *termConnector) GetProtocolUserAttribute(u, attr string) (value string, ret robot.RetVal) {
	var user *termUser
	var exists bool
	if user, exists = tc.getUserInfo(u); !exists {
		return "", robot.UserNotFound
	}
	switch attr {
	case "email":
		return user.Email, robot.Ok
	case "internalid":
		return user.InternalID, robot.Ok
	case "realname", "fullname", "real name", "full name":
		return user.FullName, robot.Ok
	case "firstname", "first name":
		return user.FirstName, robot.Ok
	case "lastname", "last name":
		return user.LastName, robot.Ok
	case "phone":
		return user.Phone, robot.Ok
	// that's all the attributes we can currently get from slack
	default:
		return "", robot.AttributeNotFound
	}
}

// SendProtocolChannelMessage sends a message to a channel
func (tc *termConnector) SendProtocolChannelMessage(ch string, msg string, f robot.MessageFormat) (ret robot.RetVal) {
	channel := tc.getChannel(ch)
	return tc.sendMessage(channel, msg, f)
}

// SendProtocolChannelMessage sends a message to a channel
func (tc *termConnector) SendProtocolUserChannelMessage(uid, uname, ch, msg string, f robot.MessageFormat) (ret robot.RetVal) {
	channel := tc.getChannel(ch)
	msg = "@" + uname + " " + msg
	return tc.sendMessage(channel, msg, f)
}

// SendProtocolUserMessage sends a direct message to a user
func (tc *termConnector) SendProtocolUserMessage(u string, msg string, f robot.MessageFormat) (ret robot.RetVal) {
	var user *termUser
	var exists bool
	if user, exists = tc.getUserInfo(u); !exists {
		return robot.UserNotFound
	}
	return tc.sendMessage(fmt.Sprintf("(dm:%s)", user.Name), msg, f)
}

// JoinChannel joins a channel given it's human-readable name, e.g. "general"
// Only useful for connectors that require it, a noop otherwise
func (tc *termConnector) JoinChannel(c string) (ret robot.RetVal) {
	return robot.Ok
}
<|endoftext|>


<commit_before>package html

import (
	"bytes"
	"code.google.com/p/go.net/html"
	"code.google.com/p/go.net/html/atom"
	"errors"
	"fmt"
	"log"
	"strings"
)

func BannerRemover(lineSeparator string, skipTop int, skipBottom int) func(string) (string, error) {
	return func(text string) (string, error) {

		lines := strings.Split(text, "\n")

		var pos []int
		for no, line := range lines {
			if strings.TrimRight(line, " ") == lineSeparator {
				pos = append(pos, no)
			}
		}

		if len(pos) == 0 {
			return text, errors.New(fmt.Sprintf("html: cannot detect lineSeparator: %s", lineSeparator))
		}

		// TODO: extract parameter
		if len(pos) != 3 {
			return text, errors.New("html: malformed mail content")
		}

		top, bottom := pos[skipTop], pos[len(pos)-skipBottom-1]
		return strings.Join(lines[top+1:bottom-1], "\n"), nil
	}
}

func ExtractText(htmlText string, remover func(string) (string, error)) (string, error) {
	z := html.NewTokenizer(strings.NewReader(htmlText))

	var buf bytes.Buffer
	bodyBlock := false

loop:
	for {
		tokenType := z.Next()
		switch {
		case tokenType == html.StartTagToken:
			if z.Token().DataAtom == atom.Body {
				bodyBlock = true
			}
		case tokenType == html.EndTagToken:
			if z.Token().DataAtom == atom.Body {
				bodyBlock = false
			}
		case tokenType == html.TextToken:
			if bodyBlock {
				buf.Write(z.Text())
			}
		case tokenType == html.ErrorToken:
			break loop
		case z.Err() != nil:
			log.Printf("html: %v\n", z.Err())
			break loop
		}
	}

	return remover(buf.String())
}
<commit_msg>finetune<commit_after>package html

import (
	"bytes"
	"code.google.com/p/go.net/html"
	"code.google.com/p/go.net/html/atom"
	"errors"
	"fmt"
	"io"
	"log"
	"strings"
)

func BannerRemover(lineSeparator string, skipTop int, skipBottom int) func(string) (string, error) {
	return func(text string) (string, error) {

		lines := strings.Split(text, "\n")

		var pos []int
		for no, line := range lines {
			if strings.TrimRight(line, " ") == lineSeparator {
				pos = append(pos, no)
			}
		}

		if len(pos) == 0 {
			return text, errors.New(fmt.Sprintf("html: cannot detect lineSeparator: %s", lineSeparator))
		}

		// TODO: extract parameter
		if len(pos) != 3 {
			return text, errors.New("html: malformed mail content")
		}

		top, bottom := pos[skipTop], pos[len(pos)-skipBottom-1]
		return strings.Join(lines[top+1:bottom-1], "\n"), nil
	}
}

func ExtractText(htmlText string, remover func(string) (string, error)) (string, error) {
	z := html.NewTokenizer(strings.NewReader(htmlText))

	var buf bytes.Buffer
	bodyBlock := false

loop:
	for {
		tokenType := z.Next()
		switch tokenType {
		case html.StartTagToken:
			if z.Token().DataAtom == atom.Body {
				bodyBlock = true
			}
		case html.EndTagToken:
			if z.Token().DataAtom == atom.Body {
				bodyBlock = false
			}
		case html.TextToken:
			if bodyBlock {
				buf.Write(z.Text())
			}
		case html.ErrorToken:
			if z.Err() != io.EOF {
				log.Printf("html: %v\n", z.Err())
			}
			break loop
		}
	}

	return remover(buf.String())
}
<|endoftext|>