<commit_before>package bridge

import (
	"errors"
	"log"
	"net"
	"net/url"
	"os"
	"path"
	"regexp"
	"strconv"
	"strings"
	"sync"

	dockerapi "github.com/fsouza/go-dockerclient"
)

var serviceIDPattern = regexp.MustCompile(`^(.+?):([a-zA-Z0-9][a-zA-Z0-9_.-]+):[0-9]+(?::udp)?$`)

type Bridge struct {
	sync.Mutex
	registry       RegistryAdapter
	docker         *dockerapi.Client
	services       map[string][]*Service
	deadContainers map[string]*DeadContainer
	config         Config
}

func New(docker *dockerapi.Client, adapterUri string, config Config) (*Bridge, error) {
	uri, err := url.Parse(adapterUri)
	if err != nil {
		return nil, errors.New("bad adapter uri: " + adapterUri)
	}
	factory, found := AdapterFactories.Lookup(uri.Scheme)
	if !found {
		return nil, errors.New("unrecognized adapter: " + adapterUri)
	}

	log.Println("Using", uri.Scheme, "adapter:", uri)
	return &Bridge{
		docker:         docker,
		config:         config,
		registry:       factory.New(uri),
		services:       make(map[string][]*Service),
		deadContainers: make(map[string]*DeadContainer),
	}, nil
}

func (b *Bridge) Ping() error {
	return b.registry.Ping()
}

func (b *Bridge) Add(containerId string) {
	b.Lock()
	defer b.Unlock()
	b.add(containerId, false)
}

func (b *Bridge) Remove(containerId string) {
	b.remove(containerId, true)
}

func (b *Bridge) RemoveOnExit(containerId string) {
	b.remove(containerId, b.shouldRemove(containerId))
}

func (b *Bridge) Refresh() {
	b.Lock()
	defer b.Unlock()

	for containerId, deadContainer := range b.deadContainers {
		deadContainer.TTL -= b.config.RefreshInterval
		if deadContainer.TTL <= 0 {
			delete(b.deadContainers, containerId)
		}
	}

	for containerId, services := range b.services {
		for _, service := range services {
			err := b.registry.Refresh(service)
			if err != nil {
				log.Println("refresh failed:", service.ID, err)
				continue
			}
			log.Println("refreshed:", containerId[:12], service.ID)
		}
	}
}

func (b *Bridge) Sync(quiet bool) {
	b.Lock()
	defer b.Unlock()

	containers, err := b.docker.ListContainers(dockerapi.ListContainersOptions{})
	if err != nil && quiet {
		log.Println("error listing containers, skipping sync")
		return
	} else if err != nil && !quiet {
		log.Fatal(err)
	}

	log.Printf("Syncing services on %d containers", len(containers))

	// NOTE: This assumes reregistering will do the right thing, i.e. nothing..
	for _, listing := range containers {
		services := b.services[listing.ID]
		if services == nil {
			b.add(listing.ID, quiet)
		} else {
			for _, service := range services {
				err := b.registry.Register(service)
				if err != nil {
					log.Println("sync register failed:", service, err)
				}
			}
		}
	}

	// Clean up services that were registered previously, but aren't
	// acknowledged within registrator
	if b.config.Cleanup {
		// Remove services if its corresponding container is not running
		log.Println("Listing non-exited containers")
		filters := map[string][]string{"status": {"created", "restarting", "running", "paused"}}
		nonExitedContainers, err := b.docker.ListContainers(dockerapi.ListContainersOptions{Filters: filters})
		if err != nil {
			log.Println("error listing nonExitedContainers, skipping sync", err)
			return
		}
		for listingId, _ := range b.services {
			found := false
			for _, container := range nonExitedContainers {
				if listingId == container.ID {
					found = true
					break
				}
			}
			// This is a container that does not exist
			if !found {
				log.Printf("stale: Removing service %s because it does not exist", listingId)
				go b.RemoveOnExit(listingId)
			}
		}

		log.Println("Cleaning up dangling services")
		extServices, err := b.registry.Services()
		if err != nil {
			log.Println("cleanup failed:", err)
			return
		}

	Outer:
		for _, extService := range extServices {
			matches := serviceIDPattern.FindStringSubmatch(extService.ID)
			if len(matches) != 3 {
				// There's no way this was registered by us, so leave it
				continue
			}
			serviceHostname := matches[1]
			if serviceHostname != Hostname {
				// ignore because registered on a different host
				continue
			}
			serviceContainerName := matches[2]
			for _, listing := range b.services {
				for _, service := range listing {
					if service.Name == extService.Name && serviceContainerName == service.Origin.container.Name[1:] {
						continue Outer
					}
				}
			}
			log.Println("dangling:", extService.ID)
			err := b.registry.Deregister(extService)
			if err != nil {
				log.Println("deregister failed:", extService.ID, err)
				continue
			}
			log.Println(extService.ID, "removed")
		}
	}
}

func (b *Bridge) add(containerId string, quiet bool) {
	if d := b.deadContainers[containerId]; d != nil {
		b.services[containerId] = d.Services
		delete(b.deadContainers, containerId)
	}

	if b.services[containerId] != nil {
		log.Println("container, ", containerId[:12], ", already exists, ignoring")
		// Alternatively, remove and readd or resubmit.
		return
	}

	container, err := b.docker.InspectContainer(containerId)
	if err != nil {
		log.Println("unable to inspect container:", containerId[:12], err)
		return
	}

	ports := make(map[string]ServicePort)

	// Extract configured host port mappings, relevant when using --net=host
	for port, published := range container.HostConfig.PortBindings {
		ports[string(port)] = servicePort(container, port, published)
	}

	// Extract runtime port mappings, relevant when using --net=bridge
	for port, published := range container.NetworkSettings.Ports {
		ports[string(port)] = servicePort(container, port, published)
	}

	if len(ports) == 0 && !quiet {
		log.Println("ignored:", container.ID[:12], "no published ports")
		return
	}

	for _, port := range ports {
		if b.config.Internal != true && port.HostPort == "" {
			if !quiet {
				log.Println("ignored:", container.ID[:12], "port", port.ExposedPort, "not published on host")
			}
			continue
		}
		service := b.newService(port, len(ports) > 1)
		if service == nil {
			if !quiet {
				log.Println("ignored:", container.ID[:12], "service on port", port.ExposedPort)
			}
			continue
		}
		err := b.registry.Register(service)
		if err != nil {
			log.Println("register failed:", service, err)
			continue
		}
		b.services[container.ID] = append(b.services[container.ID], service)
		log.Println("added:", container.ID[:12], service.ID)
	}
}

func (b *Bridge) newService(port ServicePort, isgroup bool) *Service {
	container := port.container
	defaultName := strings.Split(path.Base(container.Config.Image), ":")[0]

	// not sure about this logic. kind of want to remove it.
	hostname := Hostname
	if hostname == "" {
		hostname = port.HostIP
	}
	if port.HostIP == "0.0.0.0" {
		ip, err := net.ResolveIPAddr("ip", hostname)
		if err == nil {
			port.HostIP = ip.String()
		}
	}

	if b.config.HostIp != "" {
		port.HostIP = b.config.HostIp
	}

	metadata, metadataFromPort := serviceMetaData(container.Config, port.ExposedPort)

	ignore := mapDefault(metadata, "ignore", "")
	if ignore != "" {
		return nil
	}

	service := new(Service)
	service.Origin = port
	service.ID = hostname + ":" + container.Name[1:] + ":" + port.ExposedPort
	service.Name = mapDefault(metadata, "name", defaultName)
	if isgroup && !metadataFromPort["name"] {
		service.Name += "-" + port.ExposedPort
	}
	var p int
	if b.config.Internal == true {
		service.IP = port.ExposedIP
		p, _ = strconv.Atoi(port.ExposedPort)
	} else {
		service.IP = port.HostIP
		p, _ = strconv.Atoi(port.HostPort)
	}
	service.Port = p

	if port.PortType == "udp" {
		service.Tags = combineTags(
			mapDefault(metadata, "tags", ""), b.config.ForceTags, "udp")
		service.ID = service.ID + ":udp"
	} else {
		service.Tags = combineTags(
			mapDefault(metadata, "tags", ""), b.config.ForceTags)
	}

	id := mapDefault(metadata, "id", "")
	if id != "" {
		service.ID = id
	}

	delete(metadata, "id")
	delete(metadata, "tags")
	delete(metadata, "name")
	service.Attrs = metadata
	service.TTL = b.config.RefreshTtl

	return service
}

func (b *Bridge) remove(containerId string, deregister bool) {
	b.Lock()
	defer b.Unlock()

	if deregister {
		deregisterAll := func(services []*Service) {
			for _, service := range services {
				err := b.registry.Deregister(service)
				if err != nil {
					log.Println("deregister failed:", service.ID, err)
					continue
				}
				log.Println("removed:", containerId[:12], service.ID)
			}
		}
		deregisterAll(b.services[containerId])
		if d := b.deadContainers[containerId]; d != nil {
			deregisterAll(d.Services)
			delete(b.deadContainers, containerId)
		}
	} else if b.config.RefreshTtl != 0 && b.services[containerId] != nil {
		// need to stop the refreshing, but can't delete it yet
		b.deadContainers[containerId] = &DeadContainer{b.config.RefreshTtl, b.services[containerId]}
	}
	delete(b.services, containerId)
}

// bit set on ExitCode if it represents an exit via a signal
var dockerSignaledBit = 128

func (b *Bridge) shouldRemove(containerId string) bool {
	if b.config.DeregisterCheck == "always" {
		return true
	}
	container, err := b.docker.InspectContainer(containerId)
	if _, ok := err.(*dockerapi.NoSuchContainer); ok {
		// the container has already been removed from Docker
		// e.g. probabably run with "--rm" to remove immediately
		// so its exit code is not accessible
		log.Printf("registrator: container %v was removed, could not fetch exit code", containerId[:12])
		return true
	}

	switch {
	case err != nil:
		log.Printf("registrator: error fetching status for container %v on \"die\" event: %v\n", containerId[:12], err)
		return false
	case container.State.Running:
		log.Printf("registrator: not removing container %v, still running", containerId[:12])
		return false
	case container.State.ExitCode == 0:
		return true
	case container.State.ExitCode&dockerSignaledBit == dockerSignaledBit:
		return true
	}
	return false
}

var Hostname string

func init() {
	// It's ok for Hostname to ultimately be an empty string
	// An empty string will fall back to trying to make a best guess
	Hostname, _ = os.Hostname()
}
<commit_msg>Add containerid in default service tags<commit_after>package bridge

import (
	"errors"
	"log"
	"net"
	"net/url"
	"os"
	"path"
	"regexp"
	"strconv"
	"strings"
	"sync"

	dockerapi "github.com/fsouza/go-dockerclient"
)

var serviceIDPattern = regexp.MustCompile(`^(.+?):([a-zA-Z0-9][a-zA-Z0-9_.-]+):[0-9]+(?::udp)?$`)

type Bridge struct {
	sync.Mutex
	registry       RegistryAdapter
	docker         *dockerapi.Client
	services       map[string][]*Service
	deadContainers map[string]*DeadContainer
	config         Config
}

func New(docker *dockerapi.Client, adapterUri string, config Config) (*Bridge, error) {
	uri, err := url.Parse(adapterUri)
	if err != nil {
		return nil, errors.New("bad adapter uri: " + adapterUri)
	}
	factory, found := AdapterFactories.Lookup(uri.Scheme)
	if !found {
		return nil, errors.New("unrecognized adapter: " + adapterUri)
	}

	log.Println("Using", uri.Scheme, "adapter:", uri)
	return &Bridge{
		docker:         docker,
		config:         config,
		registry:       factory.New(uri),
		services:       make(map[string][]*Service),
		deadContainers: make(map[string]*DeadContainer),
	}, nil
}

func (b *Bridge) Ping() error {
	return b.registry.Ping()
}

func (b *Bridge) Add(containerId string) {
	b.Lock()
	defer b.Unlock()
	b.add(containerId, false)
}

func (b *Bridge) Remove(containerId string) {
	b.remove(containerId, true)
}

func (b *Bridge) RemoveOnExit(containerId string) {
	b.remove(containerId, b.shouldRemove(containerId))
}

func (b *Bridge) Refresh() {
	b.Lock()
	defer b.Unlock()

	for containerId, deadContainer := range b.deadContainers {
		deadContainer.TTL -= b.config.RefreshInterval
		if deadContainer.TTL <= 0 {
			delete(b.deadContainers, containerId)
		}
	}

	for containerId, services := range b.services {
		for _, service := range services {
			err := b.registry.Refresh(service)
			if err != nil {
				log.Println("refresh failed:", service.ID, err)
				continue
			}
			log.Println("refreshed:", containerId[:12], service.ID)
		}
	}
}

func (b *Bridge) Sync(quiet bool) {
	b.Lock()
	defer b.Unlock()

	containers, err := b.docker.ListContainers(dockerapi.ListContainersOptions{})
	if err != nil && quiet {
		log.Println("error listing containers, skipping sync")
		return
	} else if err != nil && !quiet {
		log.Fatal(err)
	}

	log.Printf("Syncing services on %d containers", len(containers))

	// NOTE: This assumes reregistering will do the right thing, i.e. nothing..
	for _, listing := range containers {
		services := b.services[listing.ID]
		if services == nil {
			b.add(listing.ID, quiet)
		} else {
			for _, service := range services {
				err := b.registry.Register(service)
				if err != nil {
					log.Println("sync register failed:", service, err)
				}
			}
		}
	}

	// Clean up services that were registered previously, but aren't
	// acknowledged within registrator
	if b.config.Cleanup {
		// Remove services if its corresponding container is not running
		log.Println("Listing non-exited containers")
		filters := map[string][]string{"status": {"created", "restarting", "running", "paused"}}
		nonExitedContainers, err := b.docker.ListContainers(dockerapi.ListContainersOptions{Filters: filters})
		if err != nil {
			log.Println("error listing nonExitedContainers, skipping sync", err)
			return
		}
		for listingId, _ := range b.services {
			found := false
			for _, container := range nonExitedContainers {
				if listingId == container.ID {
					found = true
					break
				}
			}
			// This is a container that does not exist
			if !found {
				log.Printf("stale: Removing service %s because it does not exist", listingId)
				go b.RemoveOnExit(listingId)
			}
		}

		log.Println("Cleaning up dangling services")
		extServices, err := b.registry.Services()
		if err != nil {
			log.Println("cleanup failed:", err)
			return
		}

	Outer:
		for _, extService := range extServices {
			matches := serviceIDPattern.FindStringSubmatch(extService.ID)
			if len(matches) != 3 {
				// There's no way this was registered by us, so leave it
				continue
			}
			serviceHostname := matches[1]
			if serviceHostname != Hostname {
				// ignore because registered on a different host
				continue
			}
			serviceContainerName := matches[2]
			for _, listing := range b.services {
				for _, service := range listing {
					if service.Name == extService.Name && serviceContainerName == service.Origin.container.Name[1:] {
						continue Outer
					}
				}
			}
			log.Println("dangling:", extService.ID)
			err := b.registry.Deregister(extService)
			if err != nil {
				log.Println("deregister failed:", extService.ID, err)
				continue
			}
			log.Println(extService.ID, "removed")
		}
	}
}

func (b *Bridge) add(containerId string, quiet bool) {
	if d := b.deadContainers[containerId]; d != nil {
		b.services[containerId] = d.Services
		delete(b.deadContainers, containerId)
	}

	if b.services[containerId] != nil {
		log.Println("container, ", containerId[:12], ", already exists, ignoring")
		// Alternatively, remove and readd or resubmit.
		return
	}

	container, err := b.docker.InspectContainer(containerId)
	if err != nil {
		log.Println("unable to inspect container:", containerId[:12], err)
		return
	}

	ports := make(map[string]ServicePort)

	// Extract configured host port mappings, relevant when using --net=host
	for port, published := range container.HostConfig.PortBindings {
		ports[string(port)] = servicePort(container, port, published)
	}

	// Extract runtime port mappings, relevant when using --net=bridge
	for port, published := range container.NetworkSettings.Ports {
		ports[string(port)] = servicePort(container, port, published)
	}

	if len(ports) == 0 && !quiet {
		log.Println("ignored:", container.ID[:12], "no published ports")
		return
	}

	for _, port := range ports {
		if b.config.Internal != true && port.HostPort == "" {
			if !quiet {
				log.Println("ignored:", container.ID[:12], "port", port.ExposedPort, "not published on host")
			}
			continue
		}
		service := b.newService(port, len(ports) > 1)
		if service == nil {
			if !quiet {
				log.Println("ignored:", container.ID[:12], "service on port", port.ExposedPort)
			}
			continue
		}
		err := b.registry.Register(service)
		if err != nil {
			log.Println("register failed:", service, err)
			continue
		}
		b.services[container.ID] = append(b.services[container.ID], service)
		log.Println("added:", container.ID[:12], service.ID)
	}
}

func (b *Bridge) newService(port ServicePort, isgroup bool) *Service {
	container := port.container
	defaultName := strings.Split(path.Base(container.Config.Image), ":")[0]

	// not sure about this logic. kind of want to remove it.
	hostname := Hostname
	if hostname == "" {
		hostname = port.HostIP
	}
	if port.HostIP == "0.0.0.0" {
		ip, err := net.ResolveIPAddr("ip", hostname)
		if err == nil {
			port.HostIP = ip.String()
		}
	}

	if b.config.HostIp != "" {
		port.HostIP = b.config.HostIp
	}

	metadata, metadataFromPort := serviceMetaData(container.Config, port.ExposedPort)

	ignore := mapDefault(metadata, "ignore", "")
	if ignore != "" {
		return nil
	}

	service := new(Service)
	service.Origin = port
	service.ID = hostname + ":" + container.Name[1:] + ":" + port.ExposedPort
	service.Name = mapDefault(metadata, "name", defaultName)
	if isgroup && !metadataFromPort["name"] {
		service.Name += "-" + port.ExposedPort
	}
	var p int
	if b.config.Internal == true {
		service.IP = port.ExposedIP
		p, _ = strconv.Atoi(port.ExposedPort)
	} else {
		service.IP = port.HostIP
		p, _ = strconv.Atoi(port.HostPort)
	}
	service.Port = p

	if port.PortType == "udp" {
		service.Tags = combineTags(
			mapDefault(metadata, "tags", ""), b.config.ForceTags, "containerid_"+container.ID[:12], "udp")
		service.ID = service.ID + ":udp"
	} else {
		service.Tags = combineTags(
			mapDefault(metadata, "tags", ""), b.config.ForceTags, "containerid_"+container.ID[:12])
	}

	id := mapDefault(metadata, "id", "")
	if id != "" {
		service.ID = id
	}

	delete(metadata, "id")
	delete(metadata, "tags")
	delete(metadata, "name")
	service.Attrs = metadata
	service.TTL = b.config.RefreshTtl

	return service
}

func (b *Bridge) remove(containerId string, deregister bool) {
	b.Lock()
	defer b.Unlock()

	if deregister {
		deregisterAll := func(services []*Service) {
			for _, service := range services {
				err := b.registry.Deregister(service)
				if err != nil {
					log.Println("deregister failed:", service.ID, err)
					continue
				}
				log.Println("removed:", containerId[:12], service.ID)
			}
		}
		deregisterAll(b.services[containerId])
		if d := b.deadContainers[containerId]; d != nil {
			deregisterAll(d.Services)
			delete(b.deadContainers, containerId)
		}
	} else if b.config.RefreshTtl != 0 && b.services[containerId] != nil {
		// need to stop the refreshing, but can't delete it yet
		b.deadContainers[containerId] = &DeadContainer{b.config.RefreshTtl, b.services[containerId]}
	}
	delete(b.services, containerId)
}

// bit set on ExitCode if it represents an exit via a signal
var dockerSignaledBit = 128

func (b *Bridge) shouldRemove(containerId string) bool {
	if b.config.DeregisterCheck == "always" {
		return true
	}
	container, err := b.docker.InspectContainer(containerId)
	if _, ok := err.(*dockerapi.NoSuchContainer); ok {
		// the container has already been removed from Docker
		// e.g. probabably run with "--rm" to remove immediately
		// so its exit code is not accessible
		log.Printf("registrator: container %v was removed, could not fetch exit code", containerId[:12])
		return true
	}

	switch {
	case err != nil:
		log.Printf("registrator: error fetching status for container %v on \"die\" event: %v\n", containerId[:12], err)
		return false
	case container.State.Running:
		log.Printf("registrator: not removing container %v, still running", containerId[:12])
		return false
	case container.State.ExitCode == 0:
		return true
	case container.State.ExitCode&dockerSignaledBit == dockerSignaledBit:
		return true
	}
	return false
}

var Hostname string

func init() {
	// It's ok for Hostname to ultimately be an empty string
	// An empty string will fall back to trying to make a best guess
	Hostname, _ = os.Hostname()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"time"

	"github.com/codegangsta/cli"
	"github.com/xenolf/lego/acme"
)

func checkFolder(path string) error {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return os.MkdirAll(path, 0700)
	}
	return nil
}

func setup(c *cli.Context) (*Configuration, *Account, *acme.Client) {
	err := checkFolder(c.GlobalString("path"))
	if err != nil {
		logger().Fatalf("Could not check/create path: %s", err.Error())
	}

	conf := NewConfiguration(c)
	if len(c.GlobalString("email")) == 0 {
		logger().Fatal("You have to pass an account (email address) to the program using --email or -m")
	}

	//TODO: move to account struct? Currently MUST pass email.
	acc := NewAccount(c.GlobalString("email"), conf)

	client, err := acme.NewClient(c.GlobalString("server"), acc, conf.RsaBits())
	if err != nil {
		logger().Fatalf("Could not create client: %s", err.Error())
	}

	if len(c.GlobalStringSlice("exclude")) > 0 {
		client.ExcludeChallenges(conf.ExcludedSolvers())
	}

	if c.GlobalIsSet("http") {
		client.SetHTTPAddress(c.GlobalString("http"))
	}

	if c.GlobalIsSet("tls") {
		client.SetTLSAddress(c.GlobalString("tls"))
	}

	if c.GlobalIsSet("dns") {
		var err error
		var provider acme.ChallengeProvider
		switch c.GlobalString("dns") {
		case "cloudflare":
			provider, err = acme.NewDNSProviderCloudFlare("", "")
		case "digitalocean":
			authToken := os.Getenv("DO_AUTH_TOKEN")

			provider, err = acme.NewDNSProviderDigitalOcean(authToken)
		case "dnsimple":
			provider, err = acme.NewDNSProviderDNSimple("", "")
		case "route53":
			awsRegion := os.Getenv("AWS_REGION")
			provider, err = acme.NewDNSProviderRoute53("", "", awsRegion)
		case "rfc2136":
			nameserver := os.Getenv("RFC2136_NAMESERVER")
			zone := os.Getenv("RFC2136_ZONE")
			tsigAlgorithm := os.Getenv("RFC2136_TSIG_ALGORITHM")
			tsigKey := os.Getenv("RFC2136_TSIG_KEY")
			tsigSecret := os.Getenv("RFC2136_TSIG_SECRET")

			provider, err = acme.NewDNSProviderRFC2136(nameserver, zone, tsigAlgorithm, tsigKey, tsigSecret)
		case "manual":
			provider, err = acme.NewDNSProviderManual()
		}

		if err != nil {
			logger().Fatal(err)
		}

		client.SetChallengeProvider(acme.DNS01, provider)

		// --dns=foo indicates that the user specifically want to do a DNS challenge
		// infer that the user also wants to exclude all other challenges
		client.ExcludeChallenges([]acme.Challenge{acme.HTTP01, acme.TLSSNI01})
	}

	return conf, acc, client
}

func saveCertRes(certRes acme.CertificateResource, conf *Configuration) {
	// We store the certificate, private key and metadata in different files
	// as web servers would not be able to work with a combined file.
	certOut := path.Join(conf.CertPath(), certRes.Domain+".crt")
	privOut := path.Join(conf.CertPath(), certRes.Domain+".key")
	metaOut := path.Join(conf.CertPath(), certRes.Domain+".json")

	err := ioutil.WriteFile(certOut, certRes.Certificate, 0600)
	if err != nil {
		logger().Fatalf("Unable to save Certificate for domain %s\n\t%s", certRes.Domain, err.Error())
	}

	err = ioutil.WriteFile(privOut, certRes.PrivateKey, 0600)
	if err != nil {
		logger().Fatalf("Unable to save PrivateKey for domain %s\n\t%s", certRes.Domain, err.Error())
	}

	jsonBytes, err := json.MarshalIndent(certRes, "", "\t")
	if err != nil {
		logger().Fatalf("Unable to marshal CertResource for domain %s\n\t%s", certRes.Domain, err.Error())
	}

	err = ioutil.WriteFile(metaOut, jsonBytes, 0600)
	if err != nil {
		logger().Fatalf("Unable to save CertResource for domain %s\n\t%s", certRes.Domain, err.Error())
	}
}

func run(c *cli.Context) {
	conf, acc, client := setup(c)
	if acc.Registration == nil {
		reg, err := client.Register()
		if err != nil {
			logger().Fatalf("Could not complete registration\n\t%s", err.Error())
		}

		acc.Registration = reg
		acc.Save()

		logger().Print("!!!! HEADS UP !!!!")
		logger().Printf(`
		Your account credentials have been saved in your Let's Encrypt
		configuration directory at "%s".
		You should make a secure backup	of this folder now. This
		configuration directory will also contain certificates and
		private keys obtained from Let's Encrypt so making regular
		backups of this folder is ideal.`, conf.AccountPath(c.GlobalString("email")))

	}

	if acc.Registration.Body.Agreement == "" {
		reader := bufio.NewReader(os.Stdin)
		logger().Printf("Please review the TOS at %s", acc.Registration.TosURL)

		for {
			logger().Println("Do you accept the TOS? Y/n")
			text, err := reader.ReadString('\n')
			if err != nil {
				logger().Fatalf("Could not read from console -> %s", err.Error())
			}

			text = strings.Trim(text, "\r\n")

			if text == "n" {
				logger().Fatal("You did not accept the TOS. Unable to proceed.")
			}

			if text == "Y" || text == "y" || text == "" {
				err = client.AgreeToTOS()
				if err != nil {
					logger().Fatalf("Could not agree to tos -> %s", err)
				}
				acc.Save()
				break
			}

			logger().Println("Your input was invalid. Please answer with one of Y/y, n or by pressing enter.")
		}
	}

	if len(c.GlobalStringSlice("domains")) == 0 {
		logger().Fatal("Please specify --domains or -d")
	}

	cert, failures := client.ObtainCertificate(c.GlobalStringSlice("domains"), true, nil)
	if len(failures) > 0 {
		for k, v := range failures {
			logger().Printf("[%s] Could not obtain certificates\n\t%s", k, v.Error())
		}

		// Make sure to return a non-zero exit code if ObtainSANCertificate
		// returned at least one error. Due to us not returning partial
		// certificate we can just exit here instead of at the end.
		os.Exit(1)
	}

	err := checkFolder(conf.CertPath())
	if err != nil {
		logger().Fatalf("Could not check/create path: %s", err.Error())
	}

	saveCertRes(cert, conf)
}

func revoke(c *cli.Context) {

	conf, _, client := setup(c)

	err := checkFolder(conf.CertPath())
	if err != nil {
		logger().Fatalf("Could not check/create path: %s", err.Error())
	}

	for _, domain := range c.GlobalStringSlice("domains") {
		logger().Printf("Trying to revoke certificate for domain %s", domain)

		certPath := path.Join(conf.CertPath(), domain+".crt")
		certBytes, err := ioutil.ReadFile(certPath)

		err = client.RevokeCertificate(certBytes)
		if err != nil {
			logger().Fatalf("Error while revoking the certificate for domain %s\n\t%s", domain, err.Error())
		} else {
			logger().Print("Certificate was revoked.")
		}
	}
}

func renew(c *cli.Context) {
	conf, _, client := setup(c)

	if len(c.GlobalStringSlice("domains")) <= 0 {
		logger().Fatal("Please specify at least one domain.")
	}

	domain := c.GlobalStringSlice("domains")[0]

	// load the cert resource from files.
	// We store the certificate, private key and metadata in different files
	// as web servers would not be able to work with a combined file.
	certPath := path.Join(conf.CertPath(), domain+".crt")
	privPath := path.Join(conf.CertPath(), domain+".key")
	metaPath := path.Join(conf.CertPath(), domain+".json")

	certBytes, err := ioutil.ReadFile(certPath)
	if err != nil {
		logger().Fatalf("Error while loading the certificate for domain %s\n\t%s", domain, err.Error())
	}

	if c.IsSet("days") {
		expTime, err := acme.GetPEMCertExpiration(certBytes)
		if err != nil {
			logger().Printf("Could not get Certification expiration for domain %s", domain)
		}

		if int(expTime.Sub(time.Now()).Hours()/24.0) > c.Int("days") {
			return
		}
	}

	metaBytes, err := ioutil.ReadFile(metaPath)
	if err != nil {
		logger().Fatalf("Error while loading the meta data for domain %s\n\t%s", domain, err.Error())
	}

	var certRes acme.CertificateResource
	err = json.Unmarshal(metaBytes, &certRes)
	if err != nil {
		logger().Fatalf("Error while marshalling the meta data for domain %s\n\t%s", domain, err.Error())
	}

	if c.Bool("reuse-key") {
		keyBytes, err := ioutil.ReadFile(privPath)
		if err != nil {
			logger().Fatalf("Error while loading the private key for domain %s\n\t%s", domain, err.Error())
		}
		certRes.PrivateKey = keyBytes
	}

	certRes.Certificate = certBytes

	newCert, err := client.RenewCertificate(certRes, true)
	if err != nil {
		logger().Fatalf("%s", err.Error())
	}

	saveCertRes(newCert, conf)
}
<commit_msg>CLI: Give helpful error message if --http/--tls is given without colon<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"io/ioutil"
	"os"
	"path"
	"strings"
	"time"

	"github.com/codegangsta/cli"
	"github.com/xenolf/lego/acme"
)

func checkFolder(path string) error {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return os.MkdirAll(path, 0700)
	}
	return nil
}

func setup(c *cli.Context) (*Configuration, *Account, *acme.Client) {
	err := checkFolder(c.GlobalString("path"))
	if err != nil {
		logger().Fatalf("Could not check/create path: %s", err.Error())
	}

	conf := NewConfiguration(c)
	if len(c.GlobalString("email")) == 0 {
		logger().Fatal("You have to pass an account (email address) to the program using --email or -m")
	}

	//TODO: move to account struct? Currently MUST pass email.
	acc := NewAccount(c.GlobalString("email"), conf)

	client, err := acme.NewClient(c.GlobalString("server"), acc, conf.RsaBits())
	if err != nil {
		logger().Fatalf("Could not create client: %s", err.Error())
	}

	if len(c.GlobalStringSlice("exclude")) > 0 {
		client.ExcludeChallenges(conf.ExcludedSolvers())
	}

	if c.GlobalIsSet("http") {
		if strings.Index(c.GlobalString("http"), ":") == -1 {
			logger().Fatalf("The --http switch only accepts interface:port or :port for its argument.")
		}
		client.SetHTTPAddress(c.GlobalString("http"))
	}

	if c.GlobalIsSet("tls") {
		if strings.Index(c.GlobalString("tls"), ":") == -1 {
			logger().Fatalf("The --tls switch only accepts interface:port or :port for its argument.")
		}
		client.SetTLSAddress(c.GlobalString("tls"))
	}

	if c.GlobalIsSet("dns") {
		var err error
		var provider acme.ChallengeProvider
		switch c.GlobalString("dns") {
		case "cloudflare":
			provider, err = acme.NewDNSProviderCloudFlare("", "")
		case "digitalocean":
			authToken := os.Getenv("DO_AUTH_TOKEN")

			provider, err = acme.NewDNSProviderDigitalOcean(authToken)
		case "dnsimple":
			provider, err = acme.NewDNSProviderDNSimple("", "")
		case "route53":
			awsRegion := os.Getenv("AWS_REGION")
			provider, err = acme.NewDNSProviderRoute53("", "", awsRegion)
		case "rfc2136":
			nameserver := os.Getenv("RFC2136_NAMESERVER")
			zone := os.Getenv("RFC2136_ZONE")
			tsigAlgorithm := os.Getenv("RFC2136_TSIG_ALGORITHM")
			tsigKey := os.Getenv("RFC2136_TSIG_KEY")
			tsigSecret := os.Getenv("RFC2136_TSIG_SECRET")

			provider, err = acme.NewDNSProviderRFC2136(nameserver, zone, tsigAlgorithm, tsigKey, tsigSecret)
		case "manual":
			provider, err = acme.NewDNSProviderManual()
		}

		if err != nil {
			logger().Fatal(err)
		}

		client.SetChallengeProvider(acme.DNS01, provider)

		// --dns=foo indicates that the user specifically want to do a DNS challenge
		// infer that the user also wants to exclude all other challenges
		client.ExcludeChallenges([]acme.Challenge{acme.HTTP01, acme.TLSSNI01})
	}

	return conf, acc, client
}

func saveCertRes(certRes acme.CertificateResource, conf *Configuration) {
	// We store the certificate, private key and metadata in different files
	// as web servers would not be able to work with a combined file.
	certOut := path.Join(conf.CertPath(), certRes.Domain+".crt")
	privOut := path.Join(conf.CertPath(), certRes.Domain+".key")
	metaOut := path.Join(conf.CertPath(), certRes.Domain+".json")

	err := ioutil.WriteFile(certOut, certRes.Certificate, 0600)
	if err != nil {
		logger().Fatalf("Unable to save Certificate for domain %s\n\t%s", certRes.Domain, err.Error())
	}

	err = ioutil.WriteFile(privOut, certRes.PrivateKey, 0600)
	if err != nil {
		logger().Fatalf("Unable to save PrivateKey for domain %s\n\t%s", certRes.Domain, err.Error())
	}

	jsonBytes, err := json.MarshalIndent(certRes, "", "\t")
	if err != nil {
		logger().Fatalf("Unable to marshal CertResource for domain %s\n\t%s", certRes.Domain, err.Error())
	}

	err = ioutil.WriteFile(metaOut, jsonBytes, 0600)
	if err != nil {
		logger().Fatalf("Unable to save CertResource for domain %s\n\t%s", certRes.Domain, err.Error())
	}
}

func run(c *cli.Context) {
	conf, acc, client := setup(c)
	if acc.Registration == nil {
		reg, err := client.Register()
		if err != nil {
			logger().Fatalf("Could not complete registration\n\t%s", err.Error())
		}

		acc.Registration = reg
		acc.Save()

		logger().Print("!!!! HEADS UP !!!!")
		logger().Printf(`
		Your account credentials have been saved in your Let's Encrypt
		configuration directory at "%s".
		You should make a secure backup	of this folder now. This
		configuration directory will also contain certificates and
		private keys obtained from Let's Encrypt so making regular
		backups of this folder is ideal.`, conf.AccountPath(c.GlobalString("email")))

	}

	if acc.Registration.Body.Agreement == "" {
		reader := bufio.NewReader(os.Stdin)
		logger().Printf("Please review the TOS at %s", acc.Registration.TosURL)

		for {
			logger().Println("Do you accept the TOS? Y/n")
			text, err := reader.ReadString('\n')
			if err != nil {
				logger().Fatalf("Could not read from console -> %s", err.Error())
			}

			text = strings.Trim(text, "\r\n")

			if text == "n" {
				logger().Fatal("You did not accept the TOS. Unable to proceed.")
			}

			if text == "Y" || text == "y" || text == "" {
				err = client.AgreeToTOS()
				if err != nil {
					logger().Fatalf("Could not agree to tos -> %s", err)
				}
				acc.Save()
				break
			}

			logger().Println("Your input was invalid. Please answer with one of Y/y, n or by pressing enter.")
		}
	}

	if len(c.GlobalStringSlice("domains")) == 0 {
		logger().Fatal("Please specify --domains or -d")
	}

	cert, failures := client.ObtainCertificate(c.GlobalStringSlice("domains"), true, nil)
	if len(failures) > 0 {
		for k, v := range failures {
			logger().Printf("[%s] Could not obtain certificates\n\t%s", k, v.Error())
		}

		// Make sure to return a non-zero exit code if ObtainSANCertificate
		// returned at least one error. Due to us not returning partial
		// certificate we can just exit here instead of at the end.
		os.Exit(1)
	}

	err := checkFolder(conf.CertPath())
	if err != nil {
		logger().Fatalf("Could not check/create path: %s", err.Error())
	}

	saveCertRes(cert, conf)
}

func revoke(c *cli.Context) {

	conf, _, client := setup(c)

	err := checkFolder(conf.CertPath())
	if err != nil {
		logger().Fatalf("Could not check/create path: %s", err.Error())
	}

	for _, domain := range c.GlobalStringSlice("domains") {
		logger().Printf("Trying to revoke certificate for domain %s", domain)

		certPath := path.Join(conf.CertPath(), domain+".crt")
		certBytes, err := ioutil.ReadFile(certPath)

		err = client.RevokeCertificate(certBytes)
		if err != nil {
			logger().Fatalf("Error while revoking the certificate for domain %s\n\t%s", domain, err.Error())
		} else {
			logger().Print("Certificate was revoked.")
		}
	}
}

func renew(c *cli.Context) {
	conf, _, client := setup(c)

	if len(c.GlobalStringSlice("domains")) <= 0 {
		logger().Fatal("Please specify at least one domain.")
	}

	domain := c.GlobalStringSlice("domains")[0]

	// load the cert resource from files.
	// We store the certificate, private key and metadata in different files
	// as web servers would not be able to work with a combined file.
	certPath := path.Join(conf.CertPath(), domain+".crt")
	privPath := path.Join(conf.CertPath(), domain+".key")
	metaPath := path.Join(conf.CertPath(), domain+".json")

	certBytes, err := ioutil.ReadFile(certPath)
	if err != nil {
		logger().Fatalf("Error while loading the certificate for domain %s\n\t%s", domain, err.Error())
	}

	if c.IsSet("days") {
		expTime, err := acme.GetPEMCertExpiration(certBytes)
		if err != nil {
			logger().Printf("Could not get Certification expiration for domain %s", domain)
		}

		if int(expTime.Sub(time.Now()).Hours()/24.0) > c.Int("days") {
			return
		}
	}

	metaBytes, err := ioutil.ReadFile(metaPath)
	if err != nil {
		logger().Fatalf("Error while loading the meta data for domain %s\n\t%s", domain, err.Error())
	}

	var certRes acme.CertificateResource
	err = json.Unmarshal(metaBytes, &certRes)
	if err != nil {
		logger().Fatalf("Error while marshalling the meta data for domain %s\n\t%s", domain, err.Error())
	}

	if c.Bool("reuse-key") {
		keyBytes, err := ioutil.ReadFile(privPath)
		if err != nil {
			logger().Fatalf("Error while loading the private key for domain %s\n\t%s", domain, err.Error())
		}
		certRes.PrivateKey = keyBytes
	}

	certRes.Certificate = certBytes

	newCert, err := client.RenewCertificate(certRes, true)
	if err != nil {
		logger().Fatalf("%s", err.Error())
	}

	saveCertRes(newCert, conf)
}
<|endoftext|>


<commit_before>// +build integration

package tests

import (
	"testing"

	"github.com/deis/deis/tests/utils"
)

var (
	configListCmd         = "config:list --app={{.AppName}}"
	configSetCmd          = "config:set FOO=讲台 --app={{.AppName}}"
	configSet2Cmd         = "config:set FOO=10 --app={{.AppName}}"
	configSet3Cmd         = "config:set POWERED_BY=\"the Deis team\" --app={{.AppName}}"
	configSetBuildpackCmd = "config:set BUILDPACK_URL=https://github.com/heroku/heroku-buildpack-go#98f37cc --app={{.AppName}}"
	configUnsetCmd        = "config:unset FOO --app={{.AppName}}"
)

func TestConfig(t *testing.T) {
	params := configSetup(t)
	configSetTest(t, params)
	configListTest(t, params, false)
	appsOpenTest(t, params)
	configUnsetTest(t, params)
	configListTest(t, params, true)
	limitsSetTest(t, params, 4)
	appsOpenTest(t, params)
	limitsUnsetTest(t, params, 6)
	appsOpenTest(t, params)
	//tagsTest(t, params, 8)
	appsOpenTest(t, params)
	utils.AppsDestroyTest(t, params)
}

func configSetup(t *testing.T) *utils.DeisTestConfig {
	cfg := utils.GetGlobalConfig()
	cfg.AppName = "configsample"
	utils.Execute(t, authLoginCmd, cfg, false, "")
	utils.Execute(t, gitCloneCmd, cfg, false, "")
	if err := utils.Chdir(cfg.ExampleApp); err != nil {
		t.Fatal(err)
	}
	utils.Execute(t, appsCreateCmd, cfg, false, "")
	// ensure envvars with spaces work fine on `git push`
	// https://github.com/deis/deis/issues/2477
	utils.Execute(t, configSet3Cmd, cfg, false, "the Deis team")
	// ensure custom buildpack URLS are in order
	utils.Execute(t, configSetBuildpackCmd, cfg, false, "https://github.com/heroku/heroku-buildpack-go#98f37cc")
	utils.Execute(t, gitPushCmd, cfg, false, "")
	utils.CurlWithFail(t, cfg, false, "the Deis team")
	if err := utils.Chdir(".."); err != nil {
		t.Fatal(err)
	}
	return cfg
}

func configListTest(
	t *testing.T, params *utils.DeisTestConfig, notflag bool) {
	utils.CheckList(t, configListCmd, params, "FOO", notflag)
}

func configSetTest(t *testing.T, params *utils.DeisTestConfig) {
	utils.Execute(t, configSetCmd, params, false, "讲台")
	utils.CheckList(t, appsInfoCmd, params, "(v5)", false)
	utils.Execute(t, configSet2Cmd, params, false, "10")
	utils.CheckList(t, appsInfoCmd, params, "(v6)", false)
}

func configUnsetTest(t *testing.T, params *utils.DeisTestConfig) {
	utils.Execute(t, configUnsetCmd, params, false, "")
	utils.CheckList(t, appsInfoCmd, params, "(v7)", false)
	utils.CheckList(t, "run env --app={{.AppName}}", params, "FOO", true)
}
<commit_msg>test(tests/config_test.go): add DEIS_APP and DEIS_RELEASE tests<commit_after>// +build integration

package tests

import (
	"testing"

	"github.com/deis/deis/tests/utils"
)

var (
	configListCmd         = "config:list --app={{.AppName}}"
	configSetCmd          = "config:set FOO=讲台 --app={{.AppName}}"
	configSet2Cmd         = "config:set FOO=10 --app={{.AppName}}"
	configSet3Cmd         = "config:set POWERED_BY=\"the Deis team\" --app={{.AppName}}"
	configSetBuildpackCmd = "config:set BUILDPACK_URL=https://github.com/heroku/heroku-buildpack-go#98f37cc --app={{.AppName}}"
	configUnsetCmd        = "config:unset FOO --app={{.AppName}}"
)

func TestConfig(t *testing.T) {
	params := configSetup(t)
	configSetTest(t, params)
	configListTest(t, params, false)
	appsOpenTest(t, params)
	configUnsetTest(t, params)
	configListTest(t, params, true)
	limitsSetTest(t, params, 4)
	appsOpenTest(t, params)
	limitsUnsetTest(t, params, 6)
	appsOpenTest(t, params)
	//tagsTest(t, params, 8)
	appsOpenTest(t, params)
	utils.AppsDestroyTest(t, params)
}

func configSetup(t *testing.T) *utils.DeisTestConfig {
	cfg := utils.GetGlobalConfig()
	cfg.AppName = "configsample"
	utils.Execute(t, authLoginCmd, cfg, false, "")
	utils.Execute(t, gitCloneCmd, cfg, false, "")
	if err := utils.Chdir(cfg.ExampleApp); err != nil {
		t.Fatal(err)
	}
	utils.Execute(t, appsCreateCmd, cfg, false, "")
	// ensure envvars with spaces work fine on `git push`
	// https://github.com/deis/deis/issues/2477
	utils.Execute(t, configSet3Cmd, cfg, false, "the Deis team")
	// ensure custom buildpack URLS are in order
	utils.Execute(t, configSetBuildpackCmd, cfg, false, "https://github.com/heroku/heroku-buildpack-go#98f37cc")
	utils.Execute(t, gitPushCmd, cfg, false, "")
	utils.CurlWithFail(t, cfg, false, "the Deis team")
	utils.CheckList(t, "run env --app={{.AppName}}", cfg, "DEIS_APP", false)
	utils.CheckList(t, "run env --app={{.AppName}}", cfg, "DEIS_RELEASE", false)
	if err := utils.Chdir(".."); err != nil {
		t.Fatal(err)
	}
	return cfg
}

func configListTest(
	t *testing.T, params *utils.DeisTestConfig, notflag bool) {
	utils.CheckList(t, configListCmd, params, "FOO", notflag)
}

func configSetTest(t *testing.T, params *utils.DeisTestConfig) {
	utils.Execute(t, configSetCmd, params, false, "讲台")
	utils.CheckList(t, appsInfoCmd, params, "(v5)", false)
	utils.Execute(t, configSet2Cmd, params, false, "10")
	utils.CheckList(t, appsInfoCmd, params, "(v6)", false)
}

func configUnsetTest(t *testing.T, params *utils.DeisTestConfig) {
	utils.Execute(t, configUnsetCmd, params, false, "")
	utils.CheckList(t, appsInfoCmd, params, "(v7)", false)
	utils.CheckList(t, "run env --app={{.AppName}}", params, "FOO", true)
}
<|endoftext|>


<commit_before>package restpc

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"time"
)

type Request interface {
	HTTP() *http.Request
	RemoteIP() (string, error)
	URL() *url.URL
	Host() string
	HandlerName() string

	Body() ([]byte, error)
	BodyMap() (map[string]interface{}, error)
	BodyTo(model interface{}) error

	GetHeader(string) string

	GetString(key string, sources ...FromX) (*string, error)
	GetStringList(key string, sources ...FromX) ([]string, error)
	GetInt(key string, sources ...FromX) (*int, error)
	GetFloat(key string, sources ...FromX) (*float64, error)
	GetBool(key string, sources ...FromX) (*bool, error)
	GetTime(key string, sources ...FromX) (*time.Time, error)

	FullMap() map[string]interface{}
}

type FromX interface {
	GetString(req Request, key string) (*string, error)
	GetStringList(req Request, key string) ([]string, error)
	GetInt(req Request, key string) (*int, error)
	GetFloat(req Request, key string) (*float64, error)
	GetBool(req Request, key string) (*bool, error)
	GetTime(req Request, key string) (*time.Time, error)
}

var DefaultParamSources = []FromX{
	FromBody,
	FromForm,
	FromContext,
	FromEmpty,
}

type requestImp struct {
	r           *http.Request // must be set initially
	handlerName string        // must be set initially
	body        []byte
	bodyErr     error
	bodyMap     map[string]interface{}
	bodyMapErr  error
}

func (req *requestImp) HTTP() *http.Request {
	return req.r
}

func (req *requestImp) RemoteIP() (string, error) {
	remoteIp, _, err := net.SplitHostPort(req.r.RemoteAddr)
	if err != nil {
		return "", NewError(
			Internal, "", err,
			"r.RemoteAddr", req.r.RemoteAddr,
		)
	}
	return remoteIp, nil
}

func (req *requestImp) URL() *url.URL {
	return req.r.URL
}

func (req *requestImp) Host() string {
	return req.r.Host
}

func (req *requestImp) HandlerName() string {
	return req.handlerName
}

func (req *requestImp) Body() ([]byte, error) {
	if req.body != nil {
		return req.body, nil
	}
	if req.bodyErr != nil {
		return nil, req.bodyErr
	}
	if req.r.Body == nil {
		return nil, nil
	}
	body, err := ioutil.ReadAll(req.r.Body)
	if err != nil {
		req.bodyErr = err
		log.Println(err)
	}
	req.body = body
	req.r.Body.Close()
	req.r.Body = nil
	return body, nil
}

func (req *requestImp) BodyMap() (map[string]interface{}, error) {
	if req.bodyMap != nil {
		return req.bodyMap, nil
	}
	if req.bodyMapErr != nil {
		return nil, req.bodyMapErr
	}
	data := map[string]interface{}{}
	body, err := req.Body()
	if err != nil {
		req.bodyMapErr = err
		return nil, err
	}
	if len(body) > 0 {
		err = json.Unmarshal(body, &data)
		if err != nil {
			req.bodyMapErr = err
			log.Println(err)
			// return nil, err // FIXME
		}
	}
	req.bodyMap = data
	return data, nil
}

func (req *requestImp) BodyTo(model interface{}) error {
	body, err := req.Body()
	if err != nil {
		return err
	}
	err = json.Unmarshal(body, model)
	if err != nil {
		return NewError(InvalidArgument, "request body is not a valid json", err)
	}
	return nil
}

func (req *requestImp) GetHeader(key string) string {
	return req.r.Header.Get(key)
}

func (req *requestImp) GetString(key string, sources ...FromX) (*string, error) {
	if len(sources) == 0 {
		sources = DefaultParamSources
	}
	for _, source := range sources {
		value, err := source.GetString(req, key)
		if err != nil {
			return nil, err
		}
		if value != nil {
			return value, nil
		}
	}
	return nil, NewError(
		InvalidArgument,
		fmt.Sprintf("missing '%v'", key),
		nil,
	)
}

func (req *requestImp) GetStringList(key string, sources ...FromX) ([]string, error) {
	if len(sources) == 0 {
		sources = DefaultParamSources
	}
	for _, source := range sources {
		value, err := source.GetStringList(req, key)
		if err != nil {
			return nil, err
		}
		if value != nil {
			return value, nil
		}
	}
	return nil, NewError(
		InvalidArgument,
		fmt.Sprintf("missing '%v'", key),
		nil,
	)
}

func (req *requestImp) GetInt(key string, sources ...FromX) (*int, error) {
	if len(sources) == 0 {
		sources = DefaultParamSources
	}
	for _, source := range sources {
		value, err := source.GetInt(req, key)
		if err != nil {
			return nil, err
		}
		if value != nil {
			return value, nil
		}
	}
	return nil, NewError(
		InvalidArgument,
		fmt.Sprintf("missing '%v'", key),
		nil,
	)
}

func (req *requestImp) GetFloat(key string, sources ...FromX) (*float64, error) {
	if len(sources) == 0 {
		sources = DefaultParamSources
	}
	for _, source := range sources {
		value, err := source.GetFloat(req, key)
		if err != nil {
			return nil, err
		}
		if value != nil {
			return value, nil
		}
	}
	return nil, NewError(
		InvalidArgument,
		fmt.Sprintf("missing '%v'", key),
		nil,
	)
}

func (req *requestImp) GetBool(key string, sources ...FromX) (*bool, error) {
	if len(sources) == 0 {
		sources = DefaultParamSources
	}
	for _, source := range sources {
		value, err := source.GetBool(req, key)
		if err != nil {
			return nil, err
		}
		if value != nil {
			return value, nil
		}
	}
	return nil, NewError(
		InvalidArgument,
		fmt.Sprintf("missing '%v'", key),
		nil,
	)
}

func (req *requestImp) GetTime(key string, sources ...FromX) (*time.Time, error) {
	if len(sources) == 0 {
		sources = DefaultParamSources
	}
	for _, source := range sources {
		value, err := source.GetTime(req, key)
		if err != nil {
			return nil, err
		}
		if value != nil {
			return value, nil
		}
	}
	return nil, NewError(
		InvalidArgument,
		fmt.Sprintf("missing '%v'", key),
		nil,
	)
}

func (req *requestImp) HeaderCopy() http.Header {
	header := http.Header{}
	for key, values := range req.r.Header {
		header[key] = values
	}
	return header
}

func (req *requestImp) HeaderStrippedAuth() http.Header {
	header := req.HeaderCopy()
	authHader, ok := header["Authorization"]
	if ok {
		authHaderNew := make([]string, len(authHader))
		for i := 0; i < len(authHader); i++ {
			authHaderNew[i] = "[REMOVED]"
		}
		header["Authorization"] = authHaderNew
	}
	return header
}

func (req *requestImp) FullMap() map[string]interface{} {
	bodyMap, _ := req.BodyMap()
	urlStr := req.URL().String()
	remoteIP, _ := req.RemoteIP()
	return map[string]interface{}{
		"bodyMap":  bodyMap,
		"url":      urlStr,
		"form":     req.r.Form,
		"header":   req.HeaderStrippedAuth(),
		"remoteIP": remoteIP,
		"context":  req.r.Context(),
	}
}
<commit_msg>make use of new code MissingArgument in returning errors<commit_after>package restpc

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"time"
)

type Request interface {
	HTTP() *http.Request
	RemoteIP() (string, error)
	URL() *url.URL
	Host() string
	HandlerName() string

	Body() ([]byte, error)
	BodyMap() (map[string]interface{}, error)
	BodyTo(model interface{}) error

	GetHeader(string) string

	GetString(key string, sources ...FromX) (*string, error)
	GetStringList(key string, sources ...FromX) ([]string, error)
	GetInt(key string, sources ...FromX) (*int, error)
	GetFloat(key string, sources ...FromX) (*float64, error)
	GetBool(key string, sources ...FromX) (*bool, error)
	GetTime(key string, sources ...FromX) (*time.Time, error)

	FullMap() map[string]interface{}
}

type FromX interface {
	GetString(req Request, key string) (*string, error)
	GetStringList(req Request, key string) ([]string, error)
	GetInt(req Request, key string) (*int, error)
	GetFloat(req Request, key string) (*float64, error)
	GetBool(req Request, key string) (*bool, error)
	GetTime(req Request, key string) (*time.Time, error)
}

var DefaultParamSources = []FromX{
	FromBody,
	FromForm,
	FromContext,
	FromEmpty,
}

type requestImp struct {
	r           *http.Request // must be set initially
	handlerName string        // must be set initially
	body        []byte
	bodyErr     error
	bodyMap     map[string]interface{}
	bodyMapErr  error
}

func (req *requestImp) HTTP() *http.Request {
	return req.r
}

func (req *requestImp) RemoteIP() (string, error) {
	remoteIp, _, err := net.SplitHostPort(req.r.RemoteAddr)
	if err != nil {
		return "", NewError(
			Internal, "", err,
			"r.RemoteAddr", req.r.RemoteAddr,
		)
	}
	return remoteIp, nil
}

func (req *requestImp) URL() *url.URL {
	return req.r.URL
}

func (req *requestImp) Host() string {
	return req.r.Host
}

func (req *requestImp) HandlerName() string {
	return req.handlerName
}

func (req *requestImp) Body() ([]byte, error) {
	if req.body != nil {
		return req.body, nil
	}
	if req.bodyErr != nil {
		return nil, req.bodyErr
	}
	if req.r.Body == nil {
		return nil, nil
	}
	body, err := ioutil.ReadAll(req.r.Body)
	if err != nil {
		req.bodyErr = err
		log.Println(err)
	}
	req.body = body
	req.r.Body.Close()
	req.r.Body = nil
	return body, nil
}

func (req *requestImp) BodyMap() (map[string]interface{}, error) {
	if req.bodyMap != nil {
		return req.bodyMap, nil
	}
	if req.bodyMapErr != nil {
		return nil, req.bodyMapErr
	}
	data := map[string]interface{}{}
	body, err := req.Body()
	if err != nil {
		req.bodyMapErr = err
		return nil, err
	}
	if len(body) > 0 {
		err = json.Unmarshal(body, &data)
		if err != nil {
			req.bodyMapErr = err
			log.Println(err)
			// return nil, err // FIXME
		}
	}
	req.bodyMap = data
	return data, nil
}

func (req *requestImp) BodyTo(model interface{}) error {
	body, err := req.Body()
	if err != nil {
		return err
	}
	err = json.Unmarshal(body, model)
	if err != nil {
		return NewError(InvalidArgument, "request body is not a valid json", err)
	}
	return nil
}

func (req *requestImp) GetHeader(key string) string {
	return req.r.Header.Get(key)
}

func (req *requestImp) GetString(key string, sources ...FromX) (*string, error) {
	if len(sources) == 0 {
		sources = DefaultParamSources
	}
	for _, source := range sources {
		value, err := source.GetString(req, key)
		if err != nil {
			return nil, err
		}
		if value != nil {
			return value, nil
		}
	}
	return nil, NewError(
		MissingArgument,
		fmt.Sprintf("missing '%v'", key),
		nil,
	)
}

func (req *requestImp) GetStringList(key string, sources ...FromX) ([]string, error) {
	if len(sources) == 0 {
		sources = DefaultParamSources
	}
	for _, source := range sources {
		value, err := source.GetStringList(req, key)
		if err != nil {
			return nil, err
		}
		if value != nil {
			return value, nil
		}
	}
	return nil, NewError(
		MissingArgument,
		fmt.Sprintf("missing '%v'", key),
		nil,
	)
}

func (req *requestImp) GetInt(key string, sources ...FromX) (*int, error) {
	if len(sources) == 0 {
		sources = DefaultParamSources
	}
	for _, source := range sources {
		value, err := source.GetInt(req, key)
		if err != nil {
			return nil, err
		}
		if value != nil {
			return value, nil
		}
	}
	return nil, NewError(
		MissingArgument,
		fmt.Sprintf("missing '%v'", key),
		nil,
	)
}

func (req *requestImp) GetFloat(key string, sources ...FromX) (*float64, error) {
	if len(sources) == 0 {
		sources = DefaultParamSources
	}
	for _, source := range sources {
		value, err := source.GetFloat(req, key)
		if err != nil {
			return nil, err
		}
		if value != nil {
			return value, nil
		}
	}
	return nil, NewError(
		MissingArgument,
		fmt.Sprintf("missing '%v'", key),
		nil,
	)
}

func (req *requestImp) GetBool(key string, sources ...FromX) (*bool, error) {
	if len(sources) == 0 {
		sources = DefaultParamSources
	}
	for _, source := range sources {
		value, err := source.GetBool(req, key)
		if err != nil {
			return nil, err
		}
		if value != nil {
			return value, nil
		}
	}
	return nil, NewError(
		MissingArgument,
		fmt.Sprintf("missing '%v'", key),
		nil,
	)
}

func (req *requestImp) GetTime(key string, sources ...FromX) (*time.Time, error) {
	if len(sources) == 0 {
		sources = DefaultParamSources
	}
	for _, source := range sources {
		value, err := source.GetTime(req, key)
		if err != nil {
			return nil, err
		}
		if value != nil {
			return value, nil
		}
	}
	return nil, NewError(
		MissingArgument,
		fmt.Sprintf("missing '%v'", key),
		nil,
	)
}

func (req *requestImp) HeaderCopy() http.Header {
	header := http.Header{}
	for key, values := range req.r.Header {
		header[key] = values
	}
	return header
}

func (req *requestImp) HeaderStrippedAuth() http.Header {
	header := req.HeaderCopy()
	authHader, ok := header["Authorization"]
	if ok {
		authHaderNew := make([]string, len(authHader))
		for i := 0; i < len(authHader); i++ {
			authHaderNew[i] = "[REMOVED]"
		}
		header["Authorization"] = authHaderNew
	}
	return header
}

func (req *requestImp) FullMap() map[string]interface{} {
	bodyMap, _ := req.BodyMap()
	urlStr := req.URL().String()
	remoteIP, _ := req.RemoteIP()
	return map[string]interface{}{
		"bodyMap":  bodyMap,
		"url":      urlStr,
		"form":     req.r.Form,
		"header":   req.HeaderStrippedAuth(),
		"remoteIP": remoteIP,
		"context":  req.r.Context(),
	}
}
<|endoftext|>


<commit_before>package archiver

import (
	"bytes"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
)

func TestArchiver(t *testing.T) {
	for name, ar := range SupportedFormats {
		name, ar := name, ar
		t.Run(name, func(t *testing.T) {
			t.Parallel()
			// skip RAR for now
			if _, ok := ar.(rarFormat); ok {
				t.Skip("not supported")
			}
			testWriteRead(t, name, ar)
			testMakeOpen(t, name, ar)
			testMakeOpenWithDestinationEndingInSlash(t, name, ar)
		})
	}
}

// testWriteRead performs a symmetric test by using ar.Write to generate an archive
// from the test corpus, then using ar.Read to extract the archive and comparing
// the contents to ensure they are equal.
func testWriteRead(t *testing.T, name string, ar Archiver) {
	buf := new(bytes.Buffer)
	tmp, err := ioutil.TempDir("", "archiver")
	if err != nil {
		t.Fatalf("[%s] %v", name, err)
	}
	defer os.RemoveAll(tmp)

	// Test creating archive
	err = ar.Write(buf, []string{"testdata"})
	if err != nil {
		t.Fatalf("[%s] writing archive: didn't expect an error, but got: %v", name, err)
	}

	// Test extracting archive
	err = ar.Read(buf, tmp)
	if err != nil {
		t.Fatalf("[%s] reading archive: didn't expect an error, but got: %v", name, err)
	}

	// Check that what was extracted is what was compressed
	symmetricTest(t, name, tmp)
}

// testMakeOpen performs a symmetric test by using ar.Make to make an archive
// from the test corpus, then using ar.Open to open the archive and comparing
// the contents to ensure they are equal.
func testMakeOpen(t *testing.T, name string, ar Archiver) {
	tmp, err := ioutil.TempDir("", "archiver")
	if err != nil {
		t.Fatalf("[%s] %v", name, err)
	}
	defer os.RemoveAll(tmp)

	// Test creating archive
	outfile := filepath.Join(tmp, "test-"+name)
	err = ar.Make(outfile, []string{"testdata"})
	if err != nil {
		t.Fatalf("[%s] making archive: didn't expect an error, but got: %v", name, err)
	}

	if !ar.Match(outfile) {
		t.Fatalf("[%s] identifying format should be 'true', but got 'false'", name)
	}

	// Test extracting archive
	dest := filepath.Join(tmp, "extraction_test")
	os.Mkdir(dest, 0755)
	err = ar.Open(outfile, dest)
	if err != nil {
		t.Fatalf("[%s] extracting archive [%s -> %s]: didn't expect an error, but got: %v", name, outfile, dest, err)
	}

	// Check that what was extracted is what was compressed
	symmetricTest(t, name, dest)
}

// testMakeOpenWithDestinationEndingInSlash is similar to testMakeOpen except that
// it tests the case where destination path has a terminating forward slash especially
// on Windows os.
func testMakeOpenWithDestinationEndingInSlash(t *testing.T, name string, ar Archiver) {
	tmp, err := ioutil.TempDir("", "archiver")
	if err != nil {
		t.Fatalf("[%s] %v", name, err)
	}
	defer os.RemoveAll(tmp)

	// Test creating archive
	outfile := filepath.Join(tmp, "test-"+name)
	err = ar.Make(outfile, []string{"testdata"})
	if err != nil {
		t.Fatalf("[%s] making archive: didn't expect an error, but got: %v", name, err)
	}

	if !ar.Match(outfile) {
		t.Fatalf("[%s] identifying format should be 'true', but got 'false'", name)
	}

	// Test extracting archive with destination that has a slash at the end
	dest := filepath.Join(tmp, "extraction_test")
	os.Mkdir(dest, 0755)
	err = ar.Open(outfile, dest+"/")
	if err != nil {
		t.Fatalf("[%s] extracting archive [%s -> %s]: didn't expect an error, but got: %v", name, outfile, dest, err)
	}

	// Check that what was extracted is what was compressed
	symmetricTest(t, name, dest)
}

// symmetricTest compares the contents of a destination directory to the contents
// of the test corpus and tests that they are equal.
func symmetricTest(t *testing.T, name, dest string) {
	var expectedFileCount int
	filepath.Walk("testdata", func(fpath string, info os.FileInfo, err error) error {
		expectedFileCount++
		return nil
	})

	// If outputs equals inputs, we're good; traverse output files
	// and compare file names, file contents, and file count.
	var actualFileCount int
	filepath.Walk(dest, func(fpath string, info os.FileInfo, err error) error {
		if fpath == dest {
			return nil
		}
		actualFileCount++

		origPath, err := filepath.Rel(dest, fpath)
		if err != nil {
			t.Fatalf("[%s] %s: Error inducing original file path: %v", name, fpath, err)
		}

		if info.IsDir() {
			// stat dir instead of read file
			_, err = os.Stat(origPath)
			if err != nil {
				t.Fatalf("[%s] %s: Couldn't stat original directory (%s): %v", name,
					fpath, origPath, err)
			}
			return nil
		}

		expectedFileInfo, err := os.Stat(origPath)
		if err != nil {
			t.Fatalf("[%s] %s: Error obtaining original file info: %v", name, fpath, err)
		}
		expected, err := ioutil.ReadFile(origPath)
		if err != nil {
			t.Fatalf("[%s] %s: Couldn't open original file (%s) from disk: %v", name,
				fpath, origPath, err)
		}

		actualFileInfo, err := os.Stat(fpath)
		if err != nil {
			t.Fatalf("[%s] %s: Error obtaining actual file info: %v", name, fpath, err)
		}
		actual, err := ioutil.ReadFile(fpath)
		if err != nil {
			t.Fatalf("[%s] %s: Couldn't open new file from disk: %v", name, fpath, err)
		}

		if actualFileInfo.Mode() != expectedFileInfo.Mode() {
			t.Fatalf("[%s] %s: File mode differed between on disk and compressed", name,
				expectedFileInfo.Mode().String()+" : "+actualFileInfo.Mode().String())
		}
		if !bytes.Equal(expected, actual) {
			t.Fatalf("[%s] %s: File contents differed between on disk and compressed", name, origPath)
		}

		return nil
	})

	if got, want := actualFileCount, expectedFileCount; got != want {
		t.Fatalf("[%s] Expected %d resulting files, got %d", name, want, got)
	}
}

func BenchmarkMake(b *testing.B) {
	tmp, err := ioutil.TempDir("", "archiver")
	if err != nil {
		b.Fatal(err)
	}
	defer os.RemoveAll(tmp)

	for name, ar := range SupportedFormats {
		name, ar := name, ar
		b.Run(name, func(b *testing.B) {
			// skip RAR for now
			if _, ok := ar.(rarFormat); ok {
				b.Skip("not supported")
			}
			outfile := filepath.Join(tmp, "benchMake-"+name)
			for i := 0; i < b.N; i++ {
				err = ar.Make(outfile, []string{"testdata"})
				if err != nil {
					b.Fatalf("making archive: didn't expect an error, but got: %v", err)
				}
			}
		})
	}
}

func BenchmarkOpen(b *testing.B) {
	tmp, err := ioutil.TempDir("", "archiver")
	if err != nil {
		b.Fatal(err)
	}
	defer os.RemoveAll(tmp)

	for name, ar := range SupportedFormats {
		name, ar := name, ar
		b.Run(name, func(b *testing.B) {
			// skip RAR for now
			if _, ok := ar.(rarFormat); ok {
				b.Skip("not supported")
			}
			// prepare a archive
			outfile := filepath.Join(tmp, "benchMake-"+name)
			err = ar.Make(outfile, []string{"testdata"})
			if err != nil {
				b.Fatalf("open archive: didn't expect an error, but got: %v", err)
			}
			// prepare extraction destination
			dest := filepath.Join(tmp, "extraction_test")
			os.Mkdir(dest, 0755)

			// let's go
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				err = ar.Open(outfile, dest)
				if err != nil {
					b.Fatalf("open archive: didn't expect an error, but got: %v", err)
				}
			}
		})
	}
}
<commit_msg>Minor format change<commit_after>package archiver

import (
	"bytes"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
)

func TestArchiver(t *testing.T) {
	for name, ar := range SupportedFormats {
		name, ar := name, ar
		t.Run(name, func(t *testing.T) {
			t.Parallel()
			// skip RAR for now
			if _, ok := ar.(rarFormat); ok {
				t.Skip("not supported")
			}
			testWriteRead(t, name, ar)
			testMakeOpen(t, name, ar)
			testMakeOpenWithDestinationEndingInSlash(t, name, ar)
		})
	}
}

// testWriteRead performs a symmetric test by using ar.Write to generate an archive
// from the test corpus, then using ar.Read to extract the archive and comparing
// the contents to ensure they are equal.
func testWriteRead(t *testing.T, name string, ar Archiver) {
	buf := new(bytes.Buffer)
	tmp, err := ioutil.TempDir("", "archiver")
	if err != nil {
		t.Fatalf("[%s] %v", name, err)
	}
	defer os.RemoveAll(tmp)

	// Test creating archive
	err = ar.Write(buf, []string{"testdata"})
	if err != nil {
		t.Fatalf("[%s] writing archive: didn't expect an error, but got: %v", name, err)
	}

	// Test extracting archive
	err = ar.Read(buf, tmp)
	if err != nil {
		t.Fatalf("[%s] reading archive: didn't expect an error, but got: %v", name, err)
	}

	// Check that what was extracted is what was compressed
	symmetricTest(t, name, tmp)
}

// testMakeOpen performs a symmetric test by using ar.Make to make an archive
// from the test corpus, then using ar.Open to open the archive and comparing
// the contents to ensure they are equal.
func testMakeOpen(t *testing.T, name string, ar Archiver) {
	tmp, err := ioutil.TempDir("", "archiver")
	if err != nil {
		t.Fatalf("[%s] %v", name, err)
	}
	defer os.RemoveAll(tmp)

	// Test creating archive
	outfile := filepath.Join(tmp, "test-"+name)
	err = ar.Make(outfile, []string{"testdata"})
	if err != nil {
		t.Fatalf("[%s] making archive: didn't expect an error, but got: %v", name, err)
	}

	if !ar.Match(outfile) {
		t.Fatalf("[%s] identifying format should be 'true', but got 'false'", name)
	}

	// Test extracting archive
	dest := filepath.Join(tmp, "extraction_test")
	os.Mkdir(dest, 0755)
	err = ar.Open(outfile, dest)
	if err != nil {
		t.Fatalf("[%s] extracting archive [%s -> %s]: didn't expect an error, but got: %v", name, outfile, dest, err)
	}

	// Check that what was extracted is what was compressed
	symmetricTest(t, name, dest)
}

// testMakeOpenWithDestinationEndingInSlash is similar to testMakeOpen except that
// it tests the case where destination path has a terminating forward slash especially
// on Windows os.
func testMakeOpenWithDestinationEndingInSlash(t *testing.T, name string, ar Archiver) {
	tmp, err := ioutil.TempDir("", "archiver")
	if err != nil {
		t.Fatalf("[%s] %v", name, err)
	}
	defer os.RemoveAll(tmp)

	// Test creating archive
	outfile := filepath.Join(tmp, "test-"+name)
	err = ar.Make(outfile, []string{"testdata"})
	if err != nil {
		t.Fatalf("[%s] making archive: didn't expect an error, but got: %v", name, err)
	}

	if !ar.Match(outfile) {
		t.Fatalf("[%s] identifying format should be 'true', but got 'false'", name)
	}

	// Test extracting archive with destination that has a slash at the end
	dest := filepath.Join(tmp, "extraction_test")
	os.Mkdir(dest, 0755)
	err = ar.Open(outfile, dest+"/")
	if err != nil {
		t.Fatalf("[%s] extracting archive [%s -> %s]: didn't expect an error, but got: %v", name, outfile, dest, err)
	}

	// Check that what was extracted is what was compressed
	symmetricTest(t, name, dest)
}

// symmetricTest compares the contents of a destination directory to the contents
// of the test corpus and tests that they are equal.
func symmetricTest(t *testing.T, name, dest string) {
	var expectedFileCount int
	filepath.Walk("testdata", func(fpath string, info os.FileInfo, err error) error {
		expectedFileCount++
		return nil
	})

	// If outputs equals inputs, we're good; traverse output files
	// and compare file names, file contents, and file count.
	var actualFileCount int
	filepath.Walk(dest, func(fpath string, info os.FileInfo, err error) error {
		if fpath == dest {
			return nil
		}
		actualFileCount++

		origPath, err := filepath.Rel(dest, fpath)
		if err != nil {
			t.Fatalf("[%s] %s: Error inducing original file path: %v", name, fpath, err)
		}

		if info.IsDir() {
			// stat dir instead of read file
			_, err = os.Stat(origPath)
			if err != nil {
				t.Fatalf("[%s] %s: Couldn't stat original directory (%s): %v", name,
					fpath, origPath, err)
			}
			return nil
		}

		expectedFileInfo, err := os.Stat(origPath)
		if err != nil {
			t.Fatalf("[%s] %s: Error obtaining original file info: %v", name, fpath, err)
		}
		expected, err := ioutil.ReadFile(origPath)
		if err != nil {
			t.Fatalf("[%s] %s: Couldn't open original file (%s) from disk: %v", name,
				fpath, origPath, err)
		}

		actualFileInfo, err := os.Stat(fpath)
		if err != nil {
			t.Fatalf("[%s] %s: Error obtaining actual file info: %v", name, fpath, err)
		}
		actual, err := ioutil.ReadFile(fpath)
		if err != nil {
			t.Fatalf("[%s] %s: Couldn't open new file from disk: %v", name, fpath, err)
		}

		if actualFileInfo.Mode() != expectedFileInfo.Mode() {
			t.Fatalf("[%s] %s: File mode differed between on disk and compressed", name,
				expectedFileInfo.Mode().String()+" : "+actualFileInfo.Mode().String())
		}
		if !bytes.Equal(expected, actual) {
			t.Fatalf("[%s] %s: File contents differed between on disk and compressed", name, origPath)
		}

		return nil
	})

	if got, want := actualFileCount, expectedFileCount; got != want {
		t.Fatalf("[%s] Expected %d resulting files, got %d", name, want, got)
	}
}

func BenchmarkMake(b *testing.B) {
	tmp, err := ioutil.TempDir("", "archiver")
	if err != nil {
		b.Fatal(err)
	}
	defer os.RemoveAll(tmp)

	for name, ar := range SupportedFormats {
		name, ar := name, ar
		b.Run(name, func(b *testing.B) {
			// skip RAR for now
			if _, ok := ar.(rarFormat); ok {
				b.Skip("not supported")
			}
			outfile := filepath.Join(tmp, "benchMake-"+name)
			for i := 0; i < b.N; i++ {
				err = ar.Make(outfile, []string{"testdata"})
				if err != nil {
					b.Fatalf("making archive: didn't expect an error, but got: %v", err)
				}
			}
		})
	}
}

func BenchmarkOpen(b *testing.B) {
	tmp, err := ioutil.TempDir("", "archiver")
	if err != nil {
		b.Fatal(err)
	}
	defer os.RemoveAll(tmp)

	for name, ar := range SupportedFormats {
		name, ar := name, ar
		b.Run(name, func(b *testing.B) {
			// skip RAR for now
			if _, ok := ar.(rarFormat); ok {
				b.Skip("not supported")
			}
			// prepare a archive
			outfile := filepath.Join(tmp, "benchMake-"+name)
			err = ar.Make(outfile, []string{"testdata"})
			if err != nil {
				b.Fatalf("open archive: didn't expect an error, but got: %v", err)
			}
			// prepare extraction destination
			dest := filepath.Join(tmp, "extraction_test")
			os.Mkdir(dest, 0755)

			// let's go
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				err = ar.Open(outfile, dest)
				if err != nil {
					b.Fatalf("open archive: didn't expect an error, but got: %v", err)
				}
			}
		})
	}
}
<|endoftext|>


<commit_before>package grocket

import (
    "time"
    "log"

    "github.com/kiril/btree"
)

/*
 * So, let's figure out how we're going to store this shit.
 * I want an index from {id -> event}
 * and I want an ordered index from {due -> event}
 */

type IndexedEvent struct {
    Event  *Event
    Bucket *TimeBucket
}

var eventsById = make(map[string]*IndexedEvent)
var bucketByTimeIndex = btree.NewBtree()

func FindBucketByTime(due time.Time) *TimeBucket {
    key, error := due.MarshalBinary()
    if error != nil {
        log.Fatal(error)
    }
    binary, searchError := bucketByTimeIndex.Search(key)
    if error != nil {
        log.Fatal(searchError)
    }
    if binary == nil {
        return nil
    }
    bucket := &TimeBucket{}
    bucket.UnmarshalBinary(binary)
    return bucket

}

func FindOrCreateTimeBucket(due time.Time) *TimeBucket {
    bucket := FindBucketByTime(due)
    if bucket == nil {
        bucket = &TimeBucket{Time: due, EventIds: [][]byte{}}
        SaveTimeBucket(bucket)
    }
    return bucket
}

func NextTimeBucket() *TimeBucket {
    first, error := bucketByTimeIndex.Left()
    if error != nil {
        return nil
    }

    if first != nil {
        bucket := &TimeBucket{}
        bucket.UnmarshalBinary(first)
        return bucket
    }

    return nil
}

func SaveTimeBucket(bucket *TimeBucket) error {
    key, timeError := bucket.Time.MarshalBinary()
    if timeError != nil {
        log.Fatal(timeError)
    }

    value, valueError := bucket.MarshalBinary()
    if valueError != nil {
        log.Fatal(valueError)
    }

    existing := FindBucketByTime(bucket.Time)
    if existing != nil {
        return bucketByTimeIndex.Update(key, value)
    } else {
        return bucketByTimeIndex.Insert(key, value)
    }
}

func RemoveTimeBucket(bucket *TimeBucket) error {
    key, timeError := bucket.Time.MarshalBinary()
    if timeError != nil {
        log.Fatal(timeError, "couldn't marshal time")
    }
    return bucketByTimeIndex.Delete(key)
}

func (indexed IndexedEvent) AddToBucket() {
    indexed.Bucket.AddEvent(indexed.Event)
}

func StoreEvent(event *Event) {
    indexed := eventsById[event.Id]
    if indexed != nil {
        if ! indexed.Event.Due.Equal(event.Due) {
            indexed.Bucket.RemoveEvent(indexed.Event)
            indexed = &IndexedEvent{
                Event: event,
                Bucket: FindOrCreateTimeBucket(event.Due),
            }

        } else {
            indexed.Event = event
        }

    } else {
        indexed = &IndexedEvent{
            Event: event,
            Bucket: FindOrCreateTimeBucket(event.Due),
        }
    }

    indexed.AddToBucket()
    SaveTimeBucket(indexed.Bucket)
    eventsById[event.Id] = indexed
}

func RetrieveEventById(id string) *Event {
    return eventsById[id].Event
}

func ProbabilisticSleepDuration() time.Duration {
    return time.Millisecond * 100
}
<commit_msg>remove outdated comment<commit_after>package grocket

import (
    "time"
    "log"

    "github.com/kiril/btree"
)

type IndexedEvent struct {
    Event  *Event
    Bucket *TimeBucket
}

var eventsById = make(map[string]*IndexedEvent)
var bucketByTimeIndex = btree.NewBtree()

func FindBucketByTime(due time.Time) *TimeBucket {
    key, error := due.MarshalBinary()
    if error != nil {
        log.Fatal(error)
    }
    binary, searchError := bucketByTimeIndex.Search(key)
    if error != nil {
        log.Fatal(searchError)
    }
    if binary == nil {
        return nil
    }
    bucket := &TimeBucket{}
    bucket.UnmarshalBinary(binary)
    return bucket

}

func FindOrCreateTimeBucket(due time.Time) *TimeBucket {
    bucket := FindBucketByTime(due)
    if bucket == nil {
        bucket = &TimeBucket{Time: due, EventIds: [][]byte{}}
        SaveTimeBucket(bucket)
    }
    return bucket
}

func NextTimeBucket() *TimeBucket {
    first, error := bucketByTimeIndex.Left()
    if error != nil {
        return nil
    }

    if first != nil {
        bucket := &TimeBucket{}
        bucket.UnmarshalBinary(first)
        return bucket
    }

    return nil
}

func SaveTimeBucket(bucket *TimeBucket) error {
    key, timeError := bucket.Time.MarshalBinary()
    if timeError != nil {
        log.Fatal(timeError)
    }

    value, valueError := bucket.MarshalBinary()
    if valueError != nil {
        log.Fatal(valueError)
    }

    existing := FindBucketByTime(bucket.Time)
    if existing != nil {
        return bucketByTimeIndex.Update(key, value)
    } else {
        return bucketByTimeIndex.Insert(key, value)
    }
}

func RemoveTimeBucket(bucket *TimeBucket) error {
    key, timeError := bucket.Time.MarshalBinary()
    if timeError != nil {
        log.Fatal(timeError, "couldn't marshal time")
    }
    return bucketByTimeIndex.Delete(key)
}

func (indexed IndexedEvent) AddToBucket() {
    indexed.Bucket.AddEvent(indexed.Event)
}

func StoreEvent(event *Event) {
    indexed := eventsById[event.Id]
    if indexed != nil {
        if ! indexed.Event.Due.Equal(event.Due) {
            indexed.Bucket.RemoveEvent(indexed.Event)
            indexed = &IndexedEvent{
                Event: event,
                Bucket: FindOrCreateTimeBucket(event.Due),
            }

        } else {
            indexed.Event = event
        }

    } else {
        indexed = &IndexedEvent{
            Event: event,
            Bucket: FindOrCreateTimeBucket(event.Due),
        }
    }

    indexed.AddToBucket()
    SaveTimeBucket(indexed.Bucket)
    eventsById[event.Id] = indexed
}

func RetrieveEventById(id string) *Event {
    return eventsById[id].Event
}

func ProbabilisticSleepDuration() time.Duration {
    return time.Millisecond * 100
}
<|endoftext|>


<commit_before>package lib

import (
	"github.com/stretchr/testify/assert"
	"io/ioutil"
	"os"
	"testing"
	"github.com/rafecolton/go-fileutils"
	"log"
	"path"
)

var (
	cwd string
	testDirectory string
)

func init() {
	var err error
	cwd, err = os.Getwd()
	if err != nil {
		panic(err)
	}
}

func changeToTestDir() {
	if len(testDirectory) > 0 {
		panic("testDirectory is already set: " + testDirectory)
	}

	path, err := ioutil.TempDir("", "x")
	if err != nil {
		panic(err)
	}

	if err := os.MkdirAll(path, 0655); err != nil {
		panic(err)
	}

	if err := os.Chdir(path); err != nil {
		panic(err)
	}
	testDirectory = path
}

func cleanup() {
	if _cwd, err := os.Getwd(); err != nil {
		panic(err)
	} else {
		if len(testDirectory) > 0 && _cwd != cwd {
			if err := os.Chdir(cwd); err != nil {
				panic(err)
			}
			if err := os.RemoveAll(testDirectory); err != nil {
				panic(err)
			}
			testDirectory = ""
		}
	}
}

func TestDetectionOfMavenWrapper(t *testing.T) {
	changeToTestDir()
	defer cleanup()

	content := "#!/bin/sh\necho -n x"
	ioutil.WriteFile("./mvnw", []byte(content), 0700)
	file, _ := os.Create("maven.log")

	maven, _ := NewMaven(file)

	assert.Equal(t, maven.command, "./mvnw")
}

func TestMavenNotFound(t *testing.T) {
	changeToTestDir()
	defer cleanup()

	os.Setenv("PATH", "")
	file, _ := os.Create("maven.log")

	_, err := NewMaven(file)

	assert.NotNil(t, err, "should raised an error")
}

func TestMavenWrapperFound(t *testing.T) {
	changeToTestDir()
	defer cleanup()

	os.Setenv("PATH", ".")

	content := "#!/bin/sh\necho -n x"
	ioutil.WriteFile("mvn", []byte(content), 0700)
	file, _ := os.Create("maven.log")

	// action
	maven, _ := NewMaven(file)

	logContent, _ := readFile("maven.log")

	assert.Equal(t, maven.command, "mvn")
	assert.Equal(t, logContent, "x")
}

func TestMavenParentPomUpdate(t *testing.T) {
	changeToTestDir()
	defer cleanup()

	log.Print(testDirectory)
	sourcePath := path.Dir(cwd + "/../test-projects/simple-parent-update")
	projectDir := "simple-parent-update"
	if err := fileutils.CpR(sourcePath, "x"); err != nil {
		panic(err)
	}
	if err := os.Chdir(testDirectory + "/x/" + projectDir); err != nil {
		panic(err)
	}
	x, _ := os.Getwd()
	log.Print("cwd: " + projectDir + " " + x)

	file, _ := os.Create("maven.log")

	// action
	maven, _ := NewMaven(file)
	maven.UpdateParent()

	logContent, _ := readFile("maven.log")
	log.Print(logContent)
}
<commit_msg>completes happy path test<commit_after>package lib

import (
	"github.com/stretchr/testify/assert"
	"io/ioutil"
	"os"
	"testing"
	"github.com/rafecolton/go-fileutils"
	"path"
	"strings"
)

var (
	cwd string
	testDirectory string
)

func init() {
	var err error
	cwd, err = os.Getwd()
	if err != nil {
		panic(err)
	}
}

func changeToTestDir() {
	if len(testDirectory) > 0 {
		panic("testDirectory is already set: " + testDirectory)
	}

	path, err := ioutil.TempDir("", "x")
	if err != nil {
		panic(err)
	}

	if err := os.MkdirAll(path, 0655); err != nil {
		panic(err)
	}

	if err := os.Chdir(path); err != nil {
		panic(err)
	}
	testDirectory = path
}

func cleanup() {
	if _cwd, err := os.Getwd(); err != nil {
		panic(err)
	} else {
		if len(testDirectory) > 0 && _cwd != cwd {
			if err := os.Chdir(cwd); err != nil {
				panic(err)
			}
			if err := os.RemoveAll(testDirectory); err != nil {
				panic(err)
			}
			testDirectory = ""
		}
	}
}

func TestDetectionOfMavenWrapper(t *testing.T) {
	changeToTestDir()
	defer cleanup()

	content := "#!/bin/sh\necho -n x"
	ioutil.WriteFile("./mvnw", []byte(content), 0700)
	file, _ := os.Create("maven.log")

	maven, _ := NewMaven(file)

	assert.Equal(t, maven.command, "./mvnw")
}

func TestMavenNotFound(t *testing.T) {
	changeToTestDir()
	defer cleanup()

	os.Setenv("PATH", "")
	file, _ := os.Create("maven.log")

	_, err := NewMaven(file)

	assert.NotNil(t, err, "should raised an error")
}

func TestMavenWrapperFound(t *testing.T) {
	changeToTestDir()
	defer cleanup()

	os.Setenv("PATH", ".")

	content := "#!/bin/sh\necho -n x"
	ioutil.WriteFile("mvn", []byte(content), 0700)
	file, _ := os.Create("maven.log")

	// action
	maven, _ := NewMaven(file)

	logContent, _ := readFile("maven.log")

	assert.Equal(t, maven.command, "mvn")
	assert.Equal(t, logContent, "x")
}

func copyTestProjectToTestDirectory(testProjectName string) {
	sourcePath := path.Dir(cwd + "/../test-projects/" + testProjectName)
	if err := fileutils.CpR(sourcePath, "x"); err != nil {
		panic(err)
	}
	if err := os.Chdir("x/" + testProjectName); err != nil {
		panic(err)
	}
}

func TestMavenParentPomUpdate(t *testing.T) {
	changeToTestDir()
	defer cleanup()

	copyTestProjectToTestDirectory("simple-parent-update")

	file, _ := os.Create("maven.log")

	// action
	maven, _ := NewMaven(file)
	maven.UpdateParent()

	var updateMessage string
	errors := make([]string, 0)
	logContent, _ := readFile("maven.log")
	lines := strings.Split(logContent, "\n")
	for _, line := range lines {
		updateToken := "[INFO] Updating parent from "
		if strings.HasPrefix(line, updateToken) {
			updateMessage = line
		} else {
			warnToken := "[WARNING]"
			errorToken := "[ERROR]"
			if strings.HasPrefix(line, warnToken) || strings.HasPrefix(line, errorToken) {
				errors=append(errors, line)
			}
		}
	}

	assert.Empty(t, errors)
	assert.NotZero(t, updateMessage)
	assert.True(t, strings.HasPrefix(updateMessage, "[INFO] Updating parent from 1.3.7.RELEASE to "))
}
<|endoftext|>


<commit_before>package herd

import (
	"fmt"
	"github.com/Symantec/Dominator/lib/filesystem"
	"github.com/Symantec/Dominator/lib/filter"
	"github.com/Symantec/Dominator/lib/hash"
	subproto "github.com/Symantec/Dominator/proto/sub"
	"path"
	"syscall"
	"time"
)

type state struct {
	subFS                   *filesystem.FileSystem
	requiredFS              *filesystem.FileSystem
	requiredInodeToSubInode map[uint64]uint64
	inodesChanged           map[uint64]bool   // Required inode number.
	inodesCreated           map[uint64]string // Required inode number.
	subFilenameToInode      map[string]uint64
	subObjectCacheUsage     map[hash.Hash]uint64
}

func (sub *Sub) buildUpdateRequest(request *subproto.UpdateRequest) {
	fmt.Println("buildUpdateRequest()") // TODO(rgooch): Delete debugging.
	var state state
	state.subFS = &sub.fileSystem.FileSystem
	requiredImage := sub.herd.getImage(sub.requiredImage)
	state.requiredFS = requiredImage.FileSystem
	filter := requiredImage.Filter
	request.Triggers = requiredImage.Triggers
	state.requiredInodeToSubInode = make(map[uint64]uint64)
	state.inodesChanged = make(map[uint64]bool)
	state.inodesCreated = make(map[uint64]string)
	state.subObjectCacheUsage = make(map[hash.Hash]uint64,
		len(sub.fileSystem.ObjectCache))
	var rusageStart, rusageStop syscall.Rusage
	syscall.Getrusage(syscall.RUSAGE_SELF, &rusageStart)
	// Populate subObjectCacheUsage.
	for _, hash := range sub.fileSystem.ObjectCache {
		state.subObjectCacheUsage[hash] = 0
	}
	compareDirectories(request, &state,
		&state.subFS.DirectoryInode, &state.requiredFS.DirectoryInode,
		"/", filter)
	// Look for multiply used objects and tell the sub.
	for obj, useCount := range state.subObjectCacheUsage {
		if useCount > 1 {
			if request.MultiplyUsedObjects == nil {
				request.MultiplyUsedObjects = make(map[hash.Hash]uint64)
			}
			request.MultiplyUsedObjects[obj] = useCount
		}
	}
	syscall.Getrusage(syscall.RUSAGE_SELF, &rusageStop)
	sub.lastComputeUpdateCpuDuration = time.Duration(
		rusageStop.Utime.Sec)*time.Second +
		time.Duration(rusageStop.Utime.Usec)*time.Microsecond -
		time.Duration(rusageStart.Utime.Sec)*time.Second -
		time.Duration(rusageStart.Utime.Usec)*time.Microsecond
	sub.herd.logger.Printf(
		"buildUpdateRequest(%s) took: %s user CPU time\n",
		sub.hostname, sub.lastComputeUpdateCpuDuration)
}

func compareDirectories(request *subproto.UpdateRequest, state *state,
	subDirectory, requiredDirectory *filesystem.DirectoryInode,
	myPathName string, filter *filter.Filter) {
	// First look for entries that should be deleted.
	if subDirectory != nil {
		for name := range subDirectory.EntriesByName {
			pathname := path.Join(myPathName, name)
			if filter.Match(pathname) {
				continue
			}
			if _, ok := requiredDirectory.EntriesByName[name]; !ok {
				request.PathsToDelete = append(request.PathsToDelete, pathname)
			}
		}
	}
	for name, requiredEntry := range requiredDirectory.EntriesByName {
		pathname := path.Join(myPathName, name)
		if filter.Match(pathname) {
			continue
		}
		var subEntry *filesystem.DirectoryEntry
		if subDirectory != nil {
			if se, ok := subDirectory.EntriesByName[name]; ok {
				subEntry = se
			}
		}
		if subEntry == nil {
			addEntry(request, state, requiredEntry, pathname)
		} else {
			compareEntries(request, state, subEntry, requiredEntry, pathname,
				filter)
		}
		// If a directory: descend (possibly with the directory for the sub).
		requiredInode := requiredEntry.Inode()
		if requiredInode, ok := requiredInode.(*filesystem.DirectoryInode); ok {
			var subInode *filesystem.DirectoryInode
			if subEntry != nil {
				if si, ok := subEntry.Inode().(*filesystem.DirectoryInode); ok {
					subInode = si
				}
			}
			compareDirectories(request, state, subInode, requiredInode,
				pathname, filter)
		}
	}
}

func addEntry(request *subproto.UpdateRequest, state *state,
	requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	requiredInode := requiredEntry.Inode()
	if requiredInode, ok := requiredInode.(*filesystem.DirectoryInode); ok {
		makeDirectory(request, requiredInode, myPathName, true)
	} else {
		addInode(request, state, requiredEntry, myPathName)
	}
}

func compareEntries(request *subproto.UpdateRequest, state *state,
	subEntry, requiredEntry *filesystem.DirectoryEntry,
	myPathName string, filter *filter.Filter) {
	subInode := subEntry.Inode()
	requiredInode := requiredEntry.Inode()
	sameType, sameMetadata, sameData := filesystem.CompareInodes(
		subInode, requiredInode, nil)
	if requiredInode, ok := requiredInode.(*filesystem.DirectoryInode); ok {
		if sameMetadata {
			return
		}
		if sameType {
			makeDirectory(request, requiredInode, myPathName, false)
		} else {
			makeDirectory(request, requiredInode, myPathName, true)
		}
		return
	}
	if sameType && sameData && sameMetadata {
		relink(request, state, subEntry, requiredEntry, myPathName)
		return
	}
	if sameType && sameData {
		updateMetadata(request, state, requiredEntry, myPathName)
		relink(request, state, subEntry, requiredEntry, myPathName)
		return
	}
	addInode(request, state, requiredEntry, myPathName)
}

func relink(request *subproto.UpdateRequest, state *state,
	subEntry, requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	subInum, ok := state.requiredInodeToSubInode[requiredEntry.InodeNumber]
	if !ok {
		state.requiredInodeToSubInode[requiredEntry.InodeNumber] =
			subEntry.InodeNumber
		return
	}
	if subInum == subEntry.InodeNumber {
		return
	}
	makeHardlink(request,
		myPathName, state.subFS.InodeToFilenamesTable[subInum][0])
}

func makeHardlink(request *subproto.UpdateRequest, newLink, target string) {
	var hardlink subproto.Hardlink
	hardlink.NewLink = newLink
	hardlink.Target = target
	request.HardlinksToMake = append(request.HardlinksToMake, hardlink)
}

func updateMetadata(request *subproto.UpdateRequest, state *state,
	requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	if state.inodesChanged[requiredEntry.InodeNumber] {
		return
	}
	var inode subproto.Inode
	inode.Name = myPathName
	inode.GenericInode = requiredEntry.Inode()
	request.InodesToChange = append(request.InodesToChange, inode)
	state.inodesChanged[requiredEntry.InodeNumber] = true
}

func makeDirectory(request *subproto.UpdateRequest,
	requiredInode *filesystem.DirectoryInode, pathName string, create bool) {
	var newInode subproto.Inode
	newInode.Name = pathName
	var newDirectoryInode filesystem.DirectoryInode
	newDirectoryInode.Mode = requiredInode.Mode
	newDirectoryInode.Uid = requiredInode.Uid
	newDirectoryInode.Gid = requiredInode.Gid
	newInode.GenericInode = &newDirectoryInode
	if create {
		request.DirectoriesToMake = append(request.DirectoriesToMake, newInode)
	} else {
		request.InodesToChange = append(request.InodesToChange, newInode)
	}
}

func addInode(request *subproto.UpdateRequest, state *state,
	requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	requiredInode := requiredEntry.Inode()
	if name, ok := state.inodesCreated[requiredEntry.InodeNumber]; ok {
		makeHardlink(request, myPathName, name)
		return
	}
	// Try to find a sibling inode.
	names := state.requiredFS.InodeToFilenamesTable[requiredEntry.InodeNumber]
	if len(names) > 1 {
		var sameDataInode filesystem.GenericInode
		var sameDataName string
		for _, name := range names {
			if inum, found := state.getSubInodeFromFilename(name); found {
				subInode := state.subFS.InodeTable[inum]
				_, sameMetadata, sameData := filesystem.CompareInodes(
					subInode, requiredInode, nil)
				if sameMetadata && sameData {
					makeHardlink(request, myPathName, name)
					return
				}
				if sameData {
					sameDataInode = subInode
					sameDataName = name
				}
			}
		}
		if sameDataInode != nil {
			updateMetadata(request, state, requiredEntry, sameDataName)
			makeHardlink(request, myPathName, sameDataName)
			return
		}
	}
	if inode, ok := requiredEntry.Inode().(*filesystem.RegularInode); ok {
		if inode.Size > 0 {
			if _, ok := state.subObjectCacheUsage[inode.Hash]; ok {
				state.subObjectCacheUsage[inode.Hash]++
			} else {
				// Not in object cache: grab it from file-system.
				if state.subFS.HashToInodesTable == nil {
					state.subFS.BuildHashToInodesTable()
				}
				if ilist, ok := state.subFS.HashToInodesTable[inode.Hash]; ok {
					var fileToCopy subproto.FileToCopyToCache
					fileToCopy.Name =
						state.subFS.InodeToFilenamesTable[ilist[0]][0]
					fileToCopy.Hash = inode.Hash
					request.FilesToCopyToCache = append(
						request.FilesToCopyToCache, fileToCopy)
					state.subObjectCacheUsage[inode.Hash] = 1
				} else {
					panic("No object in cache for: " + myPathName)
				}
			}
		}
	}
	var inode subproto.Inode
	inode.Name = myPathName
	inode.GenericInode = requiredEntry.Inode()
	request.InodesToMake = append(request.InodesToMake, inode)
	state.inodesCreated[requiredEntry.InodeNumber] = myPathName
}

func (state *state) getSubInodeFromFilename(name string) (uint64, bool) {
	if state.subFilenameToInode == nil {
		fmt.Println("Making subFilenameToInode map...") // HACK
		state.subFilenameToInode = make(map[string]uint64)
		for inum, names := range state.subFS.InodeToFilenamesTable {
			for _, n := range names {
				state.subFilenameToInode[n] = inum
			}
		}
		fmt.Println("Made subFilenameToInode map") // HACK
	}
	inum, ok := state.subFilenameToInode[name]
	return inum, ok
}
<commit_msg>Remove debugging message from buildUpdateRequest().<commit_after>package herd

import (
	"fmt"
	"github.com/Symantec/Dominator/lib/filesystem"
	"github.com/Symantec/Dominator/lib/filter"
	"github.com/Symantec/Dominator/lib/hash"
	subproto "github.com/Symantec/Dominator/proto/sub"
	"path"
	"syscall"
	"time"
)

type state struct {
	subFS                   *filesystem.FileSystem
	requiredFS              *filesystem.FileSystem
	requiredInodeToSubInode map[uint64]uint64
	inodesChanged           map[uint64]bool   // Required inode number.
	inodesCreated           map[uint64]string // Required inode number.
	subFilenameToInode      map[string]uint64
	subObjectCacheUsage     map[hash.Hash]uint64
}

func (sub *Sub) buildUpdateRequest(request *subproto.UpdateRequest) {
	var state state
	state.subFS = &sub.fileSystem.FileSystem
	requiredImage := sub.herd.getImage(sub.requiredImage)
	state.requiredFS = requiredImage.FileSystem
	filter := requiredImage.Filter
	request.Triggers = requiredImage.Triggers
	state.requiredInodeToSubInode = make(map[uint64]uint64)
	state.inodesChanged = make(map[uint64]bool)
	state.inodesCreated = make(map[uint64]string)
	state.subObjectCacheUsage = make(map[hash.Hash]uint64,
		len(sub.fileSystem.ObjectCache))
	var rusageStart, rusageStop syscall.Rusage
	syscall.Getrusage(syscall.RUSAGE_SELF, &rusageStart)
	// Populate subObjectCacheUsage.
	for _, hash := range sub.fileSystem.ObjectCache {
		state.subObjectCacheUsage[hash] = 0
	}
	compareDirectories(request, &state,
		&state.subFS.DirectoryInode, &state.requiredFS.DirectoryInode,
		"/", filter)
	// Look for multiply used objects and tell the sub.
	for obj, useCount := range state.subObjectCacheUsage {
		if useCount > 1 {
			if request.MultiplyUsedObjects == nil {
				request.MultiplyUsedObjects = make(map[hash.Hash]uint64)
			}
			request.MultiplyUsedObjects[obj] = useCount
		}
	}
	syscall.Getrusage(syscall.RUSAGE_SELF, &rusageStop)
	sub.lastComputeUpdateCpuDuration = time.Duration(
		rusageStop.Utime.Sec)*time.Second +
		time.Duration(rusageStop.Utime.Usec)*time.Microsecond -
		time.Duration(rusageStart.Utime.Sec)*time.Second -
		time.Duration(rusageStart.Utime.Usec)*time.Microsecond
	sub.herd.logger.Printf(
		"buildUpdateRequest(%s) took: %s user CPU time\n",
		sub.hostname, sub.lastComputeUpdateCpuDuration)
}

func compareDirectories(request *subproto.UpdateRequest, state *state,
	subDirectory, requiredDirectory *filesystem.DirectoryInode,
	myPathName string, filter *filter.Filter) {
	// First look for entries that should be deleted.
	if subDirectory != nil {
		for name := range subDirectory.EntriesByName {
			pathname := path.Join(myPathName, name)
			if filter.Match(pathname) {
				continue
			}
			if _, ok := requiredDirectory.EntriesByName[name]; !ok {
				request.PathsToDelete = append(request.PathsToDelete, pathname)
			}
		}
	}
	for name, requiredEntry := range requiredDirectory.EntriesByName {
		pathname := path.Join(myPathName, name)
		if filter.Match(pathname) {
			continue
		}
		var subEntry *filesystem.DirectoryEntry
		if subDirectory != nil {
			if se, ok := subDirectory.EntriesByName[name]; ok {
				subEntry = se
			}
		}
		if subEntry == nil {
			addEntry(request, state, requiredEntry, pathname)
		} else {
			compareEntries(request, state, subEntry, requiredEntry, pathname,
				filter)
		}
		// If a directory: descend (possibly with the directory for the sub).
		requiredInode := requiredEntry.Inode()
		if requiredInode, ok := requiredInode.(*filesystem.DirectoryInode); ok {
			var subInode *filesystem.DirectoryInode
			if subEntry != nil {
				if si, ok := subEntry.Inode().(*filesystem.DirectoryInode); ok {
					subInode = si
				}
			}
			compareDirectories(request, state, subInode, requiredInode,
				pathname, filter)
		}
	}
}

func addEntry(request *subproto.UpdateRequest, state *state,
	requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	requiredInode := requiredEntry.Inode()
	if requiredInode, ok := requiredInode.(*filesystem.DirectoryInode); ok {
		makeDirectory(request, requiredInode, myPathName, true)
	} else {
		addInode(request, state, requiredEntry, myPathName)
	}
}

func compareEntries(request *subproto.UpdateRequest, state *state,
	subEntry, requiredEntry *filesystem.DirectoryEntry,
	myPathName string, filter *filter.Filter) {
	subInode := subEntry.Inode()
	requiredInode := requiredEntry.Inode()
	sameType, sameMetadata, sameData := filesystem.CompareInodes(
		subInode, requiredInode, nil)
	if requiredInode, ok := requiredInode.(*filesystem.DirectoryInode); ok {
		if sameMetadata {
			return
		}
		if sameType {
			makeDirectory(request, requiredInode, myPathName, false)
		} else {
			makeDirectory(request, requiredInode, myPathName, true)
		}
		return
	}
	if sameType && sameData && sameMetadata {
		relink(request, state, subEntry, requiredEntry, myPathName)
		return
	}
	if sameType && sameData {
		updateMetadata(request, state, requiredEntry, myPathName)
		relink(request, state, subEntry, requiredEntry, myPathName)
		return
	}
	addInode(request, state, requiredEntry, myPathName)
}

func relink(request *subproto.UpdateRequest, state *state,
	subEntry, requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	subInum, ok := state.requiredInodeToSubInode[requiredEntry.InodeNumber]
	if !ok {
		state.requiredInodeToSubInode[requiredEntry.InodeNumber] =
			subEntry.InodeNumber
		return
	}
	if subInum == subEntry.InodeNumber {
		return
	}
	makeHardlink(request,
		myPathName, state.subFS.InodeToFilenamesTable[subInum][0])
}

func makeHardlink(request *subproto.UpdateRequest, newLink, target string) {
	var hardlink subproto.Hardlink
	hardlink.NewLink = newLink
	hardlink.Target = target
	request.HardlinksToMake = append(request.HardlinksToMake, hardlink)
}

func updateMetadata(request *subproto.UpdateRequest, state *state,
	requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	if state.inodesChanged[requiredEntry.InodeNumber] {
		return
	}
	var inode subproto.Inode
	inode.Name = myPathName
	inode.GenericInode = requiredEntry.Inode()
	request.InodesToChange = append(request.InodesToChange, inode)
	state.inodesChanged[requiredEntry.InodeNumber] = true
}

func makeDirectory(request *subproto.UpdateRequest,
	requiredInode *filesystem.DirectoryInode, pathName string, create bool) {
	var newInode subproto.Inode
	newInode.Name = pathName
	var newDirectoryInode filesystem.DirectoryInode
	newDirectoryInode.Mode = requiredInode.Mode
	newDirectoryInode.Uid = requiredInode.Uid
	newDirectoryInode.Gid = requiredInode.Gid
	newInode.GenericInode = &newDirectoryInode
	if create {
		request.DirectoriesToMake = append(request.DirectoriesToMake, newInode)
	} else {
		request.InodesToChange = append(request.InodesToChange, newInode)
	}
}

func addInode(request *subproto.UpdateRequest, state *state,
	requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	requiredInode := requiredEntry.Inode()
	if name, ok := state.inodesCreated[requiredEntry.InodeNumber]; ok {
		makeHardlink(request, myPathName, name)
		return
	}
	// Try to find a sibling inode.
	names := state.requiredFS.InodeToFilenamesTable[requiredEntry.InodeNumber]
	if len(names) > 1 {
		var sameDataInode filesystem.GenericInode
		var sameDataName string
		for _, name := range names {
			if inum, found := state.getSubInodeFromFilename(name); found {
				subInode := state.subFS.InodeTable[inum]
				_, sameMetadata, sameData := filesystem.CompareInodes(
					subInode, requiredInode, nil)
				if sameMetadata && sameData {
					makeHardlink(request, myPathName, name)
					return
				}
				if sameData {
					sameDataInode = subInode
					sameDataName = name
				}
			}
		}
		if sameDataInode != nil {
			updateMetadata(request, state, requiredEntry, sameDataName)
			makeHardlink(request, myPathName, sameDataName)
			return
		}
	}
	if inode, ok := requiredEntry.Inode().(*filesystem.RegularInode); ok {
		if inode.Size > 0 {
			if _, ok := state.subObjectCacheUsage[inode.Hash]; ok {
				state.subObjectCacheUsage[inode.Hash]++
			} else {
				// Not in object cache: grab it from file-system.
				if state.subFS.HashToInodesTable == nil {
					state.subFS.BuildHashToInodesTable()
				}
				if ilist, ok := state.subFS.HashToInodesTable[inode.Hash]; ok {
					var fileToCopy subproto.FileToCopyToCache
					fileToCopy.Name =
						state.subFS.InodeToFilenamesTable[ilist[0]][0]
					fileToCopy.Hash = inode.Hash
					request.FilesToCopyToCache = append(
						request.FilesToCopyToCache, fileToCopy)
					state.subObjectCacheUsage[inode.Hash] = 1
				} else {
					panic("No object in cache for: " + myPathName)
				}
			}
		}
	}
	var inode subproto.Inode
	inode.Name = myPathName
	inode.GenericInode = requiredEntry.Inode()
	request.InodesToMake = append(request.InodesToMake, inode)
	state.inodesCreated[requiredEntry.InodeNumber] = myPathName
}

func (state *state) getSubInodeFromFilename(name string) (uint64, bool) {
	if state.subFilenameToInode == nil {
		fmt.Println("Making subFilenameToInode map...") // HACK
		state.subFilenameToInode = make(map[string]uint64)
		for inum, names := range state.subFS.InodeToFilenamesTable {
			for _, n := range names {
				state.subFilenameToInode[n] = inum
			}
		}
		fmt.Println("Made subFilenameToInode map") // HACK
	}
	inum, ok := state.subFilenameToInode[name]
	return inum, ok
}
<|endoftext|>