<commit_before>package main

import (
	"fmt"

	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/juju/cmd"
	"launchpad.net/juju-core/juju/environs"
	"launchpad.net/juju-core/juju/log"
	"launchpad.net/juju-core/juju/state"
	"launchpad.net/tomb"

	// register providers
	_ "launchpad.net/juju-core/juju/environs/ec2"
)

// ProvisioningAgent is a cmd.Command responsible for running a provisioning agent.
type ProvisioningAgent struct {
	Conf AgentConf
}

// Info returns usage information for the command.
func (a *ProvisioningAgent) Info() *cmd.Info {
	return &cmd.Info{"provisioning", "", "run a juju provisioning agent", ""}
}

// Init initializes the command for running.
func (a *ProvisioningAgent) Init(f *gnuflag.FlagSet, args []string) error {
	a.Conf.addFlags(f)
	if err := f.Parse(true, args); err != nil {
		return err
	}
	return a.Conf.checkArgs(f.Args())
}

// Run runs a provisioning agent.
func (a *ProvisioningAgent) Run(_ *cmd.Context) error {
	// TODO(dfc) place the logic in a loop with a suitable delay
	p, err := NewProvisioner(&a.Conf.StateInfo)
	if err != nil {
		return err
	}
	return p.Wait()
}

type Provisioner struct {
	st      *state.State
	info    *state.Info
	environ environs.Environ
	tomb    tomb.Tomb

	environWatcher  *state.ConfigWatcher
	machinesWatcher *state.MachinesWatcher

	// machine.Id => environs.Instance
	instances map[int]environs.Instance
	// instance.Id => *state.Machine
	machines map[string]*state.Machine
}

// NewProvisioner returns a Provisioner.
func NewProvisioner(info *state.Info) (*Provisioner, error) {
	st, err := state.Open(info)
	if err != nil {
		return nil, err
	}
	p := &Provisioner{
		st:        st,
		info:      info,
		instances: make(map[int]environs.Instance),
		machines:  make(map[string]*state.Machine),
	}
	go p.loop()
	return p, nil
}

func (p *Provisioner) loop() {
	defer p.tomb.Done()
	defer p.st.Close()
	p.environWatcher = p.st.WatchEnvironConfig()
	for {
		select {
		case <-p.tomb.Dying():
			return
		case config, ok := <-p.environWatcher.Changes():
			if !ok {
				err := p.environWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			var err error
			p.environ, err = environs.NewEnviron(config.Map())
			if err != nil {
				log.Printf("provisioner loaded invalid environment configuration: %v", err)
				continue
			}
			log.Printf("provisioner loaded new environment configuration")

			p.innerLoop()
		}
	}
}

func (p *Provisioner) innerLoop() {
	// call processMachines to stop any unknown instances before watching machines.
	if err := p.processMachines(nil, nil); err != nil {
		p.tomb.Kill(err)
	}
	p.machinesWatcher = p.st.WatchMachines()
	for {
		select {
		case <-p.tomb.Dying():
			return
		case change, ok := <-p.environWatcher.Changes():
			if !ok {
				err := p.environWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			config, err := environs.NewConfig(change.Map())
			if err != nil {
				log.Printf("provisioner loaded invalid environment configuration: %v", err)
				continue
			}
			p.environ.SetConfig(config)
			log.Printf("provisioner loaded new environment configuration")
		case machines, ok := <-p.machinesWatcher.Changes():
			if !ok {
				err := p.machinesWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			// TODO(dfc) fire process machines periodically to shut down unknown
			// instances.
			if err := p.processMachines(machines.Added, machines.Deleted); err != nil {
				p.tomb.Kill(err)
			}
		}
	}
}

// Wait waits for the Provisioner to exit.
func (p *Provisioner) Wait() error {
	return p.tomb.Wait()
}

// Stop stops the Provisioner and returns any error encountered while
// provisioning.
func (p *Provisioner) Stop() error {
	p.tomb.Kill(nil)
	return p.tomb.Wait()
}

func (p *Provisioner) processMachines(added, removed []*state.Machine) error {
	// step 1. find which of the added machines have not
	// yet been allocated a started instance.
	notstarted, err := p.findNotStarted(added)
	if err != nil {
		return err
	}

	// step 2. start an instance for any machines we found.
	if err := p.startMachines(notstarted); err != nil {
		return err
	}

	// step 3. stop all machines that were removed from the state.
	stopping, err := p.instancesForMachines(removed)
	if err != nil {
		return err
	}

	// step 4. find instances which are running but have no machine
	// associated with them.
	unknown, err := p.findUnknownInstances()
	if err != nil {
		return err
	}

	return p.stopInstances(append(stopping, unknown...))
}

// findUnknownInstances finds instances which are not associated with a machine.
func (p *Provisioner) findUnknownInstances() ([]environs.Instance, error) {
	all, err := p.environ.AllInstances()
	if err != nil {
		return nil, err
	}
	instances := make(map[string]environs.Instance)
	for _, i := range all {
		instances[i.Id()] = i
	}
	// TODO(dfc) this is very inefficient, p.machines cache may help.
	machines, err := p.st.AllMachines()
	if err != nil {
		return nil, err
	}
	for _, m := range machines {
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			// TODO(dfc) InstanceId should return an error if the id isn't set.
			continue
		}
		delete(instances, id)
	}
	var unknown []environs.Instance
	for _, i := range instances {
		unknown = append(unknown, i)
	}
	return unknown, nil
}

// findNotStarted finds machines without an InstanceId set, these are defined as not started.
func (p *Provisioner) findNotStarted(machines []*state.Machine) ([]*state.Machine, error) {
	var notstarted []*state.Machine
	for _, m := range machines {
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			// TODO(dfc) InstanceId should return an error if the id isn't set.
			notstarted = append(notstarted, m)
		} else {
			log.Printf("machine %s already started as instance %q", m, id)
		}
	}
	return notstarted, nil
}

func (p *Provisioner) startMachines(machines []*state.Machine) error {
	for _, m := range machines {
		if err := p.startMachine(m); err != nil {
			return err
		}
	}
	return nil
}

func (p *Provisioner) startMachine(m *state.Machine) error {
	// TODO(dfc) the state.Info passed to environ.StartInstance remains contentious
	// however as the PA only knows one state.Info, and that info is used by MAs and
	// UAs to locate the ZK for this environment, it is logical to use the same
	// state.Info as the PA.
	inst, err := p.environ.StartInstance(m.Id(), p.info)
	if err != nil {
		log.Printf("provisioner can't start machine %s: %v", m, err)
		return err
	}

	// assign the instance id to the machine
	if err := m.SetInstanceId(inst.Id()); err != nil {
		return err
	}

	// populate the local cache
	p.instances[m.Id()] = inst
	p.machines[inst.Id()] = m
	log.Printf("provisioner started machine %s as instance %s", m, inst.Id())
	return nil
}

func (p *Provisioner) stopInstances(instances []environs.Instance) error {
	// Although calling StopInstance with an empty slice should produce no change in the
	// provider, environs like dummy do not consider this a noop.
	if len(instances) == 0 {
		return nil
	}
	if err := p.environ.StopInstances(instances); err != nil {
		return err
	}

	// cleanup cache
	for _, i := range instances {
		if m, ok := p.machines[i.Id()]; ok {
			delete(p.machines, i.Id())
			delete(p.instances, m.Id())
		}
	}
	return nil
}

// instanceForMachine returns the environs.Instance that represents this machine's instance.
func (p *Provisioner) instanceForMachine(m *state.Machine) (environs.Instance, error) {
	inst, ok := p.instances[m.Id()]
	if !ok {
		// not cached locally, ask the environ.
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			// TODO(dfc) InstanceId should return an error if the id isn't set.
			return nil, fmt.Errorf("machine %s not found", m)
		}
		// TODO(dfc) this should be batched, or the cache preloaded at startup to
		// avoid N calls to the envirion.
		insts, err := p.environ.Instances([]string{id})
		if err != nil {
			// the provider doesn't know about this instance, give up.
			return nil, err
		}
		inst = insts[0]
	}
	return inst, nil
}

// instancesForMachines returns a list of environs.Instance that represent the list of machines running
// in the provider.
func (p *Provisioner) instancesForMachines(machines []*state.Machine) ([]environs.Instance, error) {
	var insts []environs.Instance
	for _, m := range machines {
		inst, err := p.instanceForMachine(m)
		if err != nil {
			return nil, err
		}
		insts = append(insts, inst)
	}
	return insts, nil
}
<commit_msg>cmd/jujud: add retry strategy<commit_after>package main

import (
	"fmt"
	"time"

	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/juju/cmd"
	"launchpad.net/juju-core/juju/environs"
	"launchpad.net/juju-core/juju/log"
	"launchpad.net/juju-core/juju/state"
	"launchpad.net/tomb"

	// register providers
	_ "launchpad.net/juju-core/juju/environs/ec2"
)

var retryDuration = 10 * time.Second

// ProvisioningAgent is a cmd.Command responsible for running a provisioning agent.
type ProvisioningAgent struct {
	Conf AgentConf
}

// Info returns usage information for the command.
func (a *ProvisioningAgent) Info() *cmd.Info {
	return &cmd.Info{"provisioning", "", "run a juju provisioning agent", ""}
}

// Init initializes the command for running.
func (a *ProvisioningAgent) Init(f *gnuflag.FlagSet, args []string) error {
	a.Conf.addFlags(f)
	if err := f.Parse(true, args); err != nil {
		return err
	}
	return a.Conf.checkArgs(f.Args())
}

// Run runs a provisioning agent.
func (a *ProvisioningAgent) Run(_ *cmd.Context) error {
	for {
		p, err := NewProvisioner(&a.Conf.StateInfo)
		if err == nil {
			if err = p.Wait(); err == nil {
				// if Wait returns nil then we consider that a signal
				// that the process should exit the retry logic.
				return nil
			}
		}
		log.Printf("restarting provisioner after error: %v", err)
		time.Sleep(retryDuration)
	}
	panic("unreachable")
}

type Provisioner struct {
	st      *state.State
	info    *state.Info
	environ environs.Environ
	tomb    tomb.Tomb

	environWatcher  *state.ConfigWatcher
	machinesWatcher *state.MachinesWatcher

	// machine.Id => environs.Instance
	instances map[int]environs.Instance
	// instance.Id => *state.Machine
	machines map[string]*state.Machine
}

// NewProvisioner returns a Provisioner.
func NewProvisioner(info *state.Info) (*Provisioner, error) {
	st, err := state.Open(info)
	if err != nil {
		return nil, err
	}
	p := &Provisioner{
		st:        st,
		info:      info,
		instances: make(map[int]environs.Instance),
		machines:  make(map[string]*state.Machine),
	}
	go p.loop()
	return p, nil
}

func (p *Provisioner) loop() {
	defer p.tomb.Done()
	defer p.st.Close()
	p.environWatcher = p.st.WatchEnvironConfig()
	for {
		select {
		case <-p.tomb.Dying():
			return
		case config, ok := <-p.environWatcher.Changes():
			if !ok {
				err := p.environWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			var err error
			p.environ, err = environs.NewEnviron(config.Map())
			if err != nil {
				log.Printf("provisioner loaded invalid environment configuration: %v", err)
				continue
			}
			log.Printf("provisioner loaded new environment configuration")

			p.innerLoop()
		}
	}
}

func (p *Provisioner) innerLoop() {
	// call processMachines to stop any unknown instances before watching machines.
	if err := p.processMachines(nil, nil); err != nil {
		p.tomb.Kill(err)
	}
	p.machinesWatcher = p.st.WatchMachines()
	for {
		select {
		case <-p.tomb.Dying():
			return
		case change, ok := <-p.environWatcher.Changes():
			if !ok {
				err := p.environWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			config, err := environs.NewConfig(change.Map())
			if err != nil {
				log.Printf("provisioner loaded invalid environment configuration: %v", err)
				continue
			}
			p.environ.SetConfig(config)
			log.Printf("provisioner loaded new environment configuration")
		case machines, ok := <-p.machinesWatcher.Changes():
			if !ok {
				err := p.machinesWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			// TODO(dfc) fire process machines periodically to shut down unknown
			// instances.
			if err := p.processMachines(machines.Added, machines.Deleted); err != nil {
				p.tomb.Kill(err)
			}
		}
	}
}

// Wait waits for the Provisioner to exit.
func (p *Provisioner) Wait() error {
	return p.tomb.Wait()
}

// Stop stops the Provisioner and returns any error encountered while
// provisioning.
func (p *Provisioner) Stop() error {
	p.tomb.Kill(nil)
	return p.tomb.Wait()
}

func (p *Provisioner) processMachines(added, removed []*state.Machine) error {
	// step 1. find which of the added machines have not
	// yet been allocated a started instance.
	notstarted, err := p.findNotStarted(added)
	if err != nil {
		return err
	}

	// step 2. start an instance for any machines we found.
	if err := p.startMachines(notstarted); err != nil {
		return err
	}

	// step 3. stop all machines that were removed from the state.
	stopping, err := p.instancesForMachines(removed)
	if err != nil {
		return err
	}

	// step 4. find instances which are running but have no machine
	// associated with them.
	unknown, err := p.findUnknownInstances()
	if err != nil {
		return err
	}

	return p.stopInstances(append(stopping, unknown...))
}

// findUnknownInstances finds instances which are not associated with a machine.
func (p *Provisioner) findUnknownInstances() ([]environs.Instance, error) {
	all, err := p.environ.AllInstances()
	if err != nil {
		return nil, err
	}
	instances := make(map[string]environs.Instance)
	for _, i := range all {
		instances[i.Id()] = i
	}
	// TODO(dfc) this is very inefficient, p.machines cache may help.
	machines, err := p.st.AllMachines()
	if err != nil {
		return nil, err
	}
	for _, m := range machines {
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			// TODO(dfc) InstanceId should return an error if the id isn't set.
			continue
		}
		delete(instances, id)
	}
	var unknown []environs.Instance
	for _, i := range instances {
		unknown = append(unknown, i)
	}
	return unknown, nil
}

// findNotStarted finds machines without an InstanceId set, these are defined as not started.
func (p *Provisioner) findNotStarted(machines []*state.Machine) ([]*state.Machine, error) {
	var notstarted []*state.Machine
	for _, m := range machines {
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			// TODO(dfc) InstanceId should return an error if the id isn't set.
			notstarted = append(notstarted, m)
		} else {
			log.Printf("machine %s already started as instance %q", m, id)
		}
	}
	return notstarted, nil
}

func (p *Provisioner) startMachines(machines []*state.Machine) error {
	for _, m := range machines {
		if err := p.startMachine(m); err != nil {
			return err
		}
	}
	return nil
}

func (p *Provisioner) startMachine(m *state.Machine) error {
	// TODO(dfc) the state.Info passed to environ.StartInstance remains contentious
	// however as the PA only knows one state.Info, and that info is used by MAs and
	// UAs to locate the ZK for this environment, it is logical to use the same
	// state.Info as the PA.
	inst, err := p.environ.StartInstance(m.Id(), p.info)
	if err != nil {
		log.Printf("provisioner can't start machine %s: %v", m, err)
		return err
	}

	// assign the instance id to the machine
	if err := m.SetInstanceId(inst.Id()); err != nil {
		return err
	}

	// populate the local cache
	p.instances[m.Id()] = inst
	p.machines[inst.Id()] = m
	log.Printf("provisioner started machine %s as instance %s", m, inst.Id())
	return nil
}

func (p *Provisioner) stopInstances(instances []environs.Instance) error {
	// Although calling StopInstance with an empty slice should produce no change in the
	// provider, environs like dummy do not consider this a noop.
	if len(instances) == 0 {
		return nil
	}
	if err := p.environ.StopInstances(instances); err != nil {
		return err
	}

	// cleanup cache
	for _, i := range instances {
		if m, ok := p.machines[i.Id()]; ok {
			delete(p.machines, i.Id())
			delete(p.instances, m.Id())
		}
	}
	return nil
}

// instanceForMachine returns the environs.Instance that represents this machine's instance.
func (p *Provisioner) instanceForMachine(m *state.Machine) (environs.Instance, error) {
	inst, ok := p.instances[m.Id()]
	if !ok {
		// not cached locally, ask the environ.
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			// TODO(dfc) InstanceId should return an error if the id isn't set.
			return nil, fmt.Errorf("machine %s not found", m)
		}
		// TODO(dfc) this should be batched, or the cache preloaded at startup to
		// avoid N calls to the envirion.
		insts, err := p.environ.Instances([]string{id})
		if err != nil {
			// the provider doesn't know about this instance, give up.
			return nil, err
		}
		inst = insts[0]
	}
	return inst, nil
}

// instancesForMachines returns a list of environs.Instance that represent the list of machines running
// in the provider.
func (p *Provisioner) instancesForMachines(machines []*state.Machine) ([]environs.Instance, error) {
	var insts []environs.Instance
	for _, m := range machines {
		inst, err := p.instanceForMachine(m)
		if err != nil {
			return nil, err
		}
		insts = append(insts, inst)
	}
	return insts, nil
}
<|endoftext|>


<commit_before>package main

import (
    "encoding/json"
    "fmt"
    "github.com/manterfield/docker-snitch"
    docker "github.com/fsouza/go-dockerclient"
    "net/http"
    "strings"
)

func index() http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        fmt.Fprintf(w, "Hello world!")
    })
}

func containers() http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        status := r.URL.Query().Get("status")
        containers := dockersnitch.Containers(status)

        w.Header().Set("Content-Type", "application/json")
        json.NewEncoder(w).Encode(containers)
    })
}

func containerOpts() http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        status := r.URL.Query().Get("status")
        containers := dockersnitch.Containers(status)

        w.Header().Set("Content-Type", "application/json")
        json.NewEncoder(w).Encode(rundeckContainerOpts(containers))
    })
}

type RundeckContainerOpt struct {
    Name string `json:"name"`
    Value string `json:"value"`
}

func rundeckContainerOpts(ctrs []docker.APIContainers) []RundeckContainerOpt {
    var containerOpts []RundeckContainerOpt
    var ctrOpt RundeckContainerOpt
    var optName string

    for _, ctr := range ctrs {
        if ctr.Image == "tictocstech/docker-snitch" {
            continue
        }
        optName = strings.Join(append(ctr.Names, ctr.Image), " ")
        optName = strings.TrimLeft(optName, "/")
        ctrOpt = RundeckContainerOpt{Name: optName, Value: ctr.ID}
        containerOpts = append(containerOpts, ctrOpt)
    }

    return containerOpts
}

func main() {
    http.Handle("/", index())
    http.Handle("/containers", containers())
    http.Handle("/container-opts", containerOpts())
    http.ListenAndServe(":8080", nil)
}
<commit_msg>Added sorting of container opts endpoint<commit_after>package main

import (
    "encoding/json"
    "fmt"
    "github.com/manterfield/docker-snitch"
    docker "github.com/fsouza/go-dockerclient"
    "net/http"
    "strings"
)

func index() http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        fmt.Fprintf(w, "Hello world!")
    })
}

func containers() http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        status := r.URL.Query().Get("status")
        containers := dockersnitch.Containers(status)

        w.Header().Set("Content-Type", "application/json")
        json.NewEncoder(w).Encode(containers)
    })
}

func containerOpts() http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        status := r.URL.Query().Get("status")
        containers := dockersnitch.Containers(status)

        w.Header().Set("Content-Type", "application/json")
        json.NewEncoder(w).Encode(getContainerOpts(containers))
    })
}

type ContainerOpt struct {
    Name string `json:"name"`
    Value string `json:"value"`
}

func getContainerOpts(ctrs []docker.APIContainers) []ContainerOpt {
    var containerOpts []ContainerOpt
    var ctrOpt ContainerOpt
    var optName string

    for _, ctr := range ctrs {
        if ctr.Image == "tictocstech/docker-snitch" {
            continue
        }
        optName = strings.Join(append(ctr.Names, ctr.Image), " ")
        optName = strings.TrimLeft(optName, "/")
        ctrOpt = ContainerOpt{Name: optName, Value: ctr.ID}
        containerOpts = append(containerOpts, ctrOpt)
    }

    sort.Slice(containerOpts, func(i, j int) bool {
        return containerOpts[i].Name < containerOpts[j].Name
    })
    return containerOpts
}

func main() {
    http.Handle("/", index())
    http.Handle("/containers", containers())
    http.Handle("/container-opts", containerOpts())
    http.ListenAndServe(":8080", nil)
}
<|endoftext|>


<commit_before>package modata

import (
  "fmt"
  "web"
  "sort"
  "time"
)

const REPL_INTERVAL = 10 * time.Second

type ReplicationServer struct {
  MyContact Contact

  server *web.Server
  blockServer *BlockServer
}

func BuildKeyMap(contacts ContactList) map[string][]string {
  keyMapping := make(map[string][]string)
  for _, contact := range contacts {
    status, data, _ := JsonGet(contact.ToHttpAddress() + "/keys", Contact{})
    if (status == OK) {
      for _, key := range data.([]interface{}) {
        skey := key.(string)
        lst, ok := keyMapping[skey]
        if ok {
          keyMapping[skey] = append(lst, contact.ToHttpAddress())
        } else {
          keyMapping[skey] = []string{contact.ToHttpAddress()}
        }
      }
    }
  }
  return keyMapping
}

func (rs *ReplicationServer) IntelligentReplication () {
  contacts := rs.blockServer.routingTable.AllContacts()
  contacts = append(contacts, rs.blockServer.contact)
  keyMapping := BuildKeyMap(contacts)

  // Stores repl_count -> keys
  replicationMap := make(map[int][]string)
  for key, value := range keyMapping {
    replicationCount := len(value)
    lst, ok := replicationMap[replicationCount]
    if ok {
      replicationMap[replicationCount] = append(lst, key)
    } else {
      replicationMap[replicationCount] = []string{key}
    }
  }

  replicationCounts := make([]int,0)
  for key, _ := range replicationMap {
    replicationCounts = append(replicationCounts, key)
  }
  sort.Ints(replicationCounts)

  max := 3
  for _, value := range replicationCounts {
    if value >= rs.blockServer.routingTable.k {
      // Don't replicate fully replicated nodes
      break
    }
    keysToReplicate := replicationMap[value]
    for _, key := range keysToReplicate {
      fmt.Printf("Re-replicating %v with replication %v\n", key, value)
      server := keyMapping[key][0]
      JsonPostUrl(server + "/replicate?key=" + key, Contact{})
    }
    max -= 1
    if (max <= 0) {
      break
    }
  }

}

func (rs *ReplicationServer) KeyMap(c *web.Context) string{
  contacts := rs.blockServer.routingTable.AllContacts()
  contacts = append(contacts, rs.blockServer.contact)
  keyMapping := BuildKeyMap(contacts)
  finalData := make(map[string]interface{})
  finalData["stats"] = map[string]int{"numKeys": len(keyMapping), "contacts": len(contacts),
                        "averageKeysPerContact": len(keyMapping) /
                        len(contacts)}
  finalData["keys"] = make(map[string]interface{})
  for key, value := range keyMapping {
    sort.Strings(value)
    finalData["keys"].(map[string]interface{})[key] = make(map[string]interface{})
    finalData["keys"].(map[string]interface{})[key].(map[string]interface{})["replicationCount"] = len(value)
    finalData["keys"].(map[string]interface{})[key].(map[string]interface{})["nodes"] = value
  }
  return RespondWithData(finalData)
}

func StartReplicationServer(name string, bs *BlockServer) *ReplicationServer{
  rs := new(ReplicationServer)

  rs.server = web.NewServer()
  rs.blockServer = bs

  go func() {
    // Identifier for this node

    rs.server.Get("/contacts", func (c *web.Context) string {
      c.ContentType("json")
      return RespondWithData(bs.routingTable.AllContacts())
    })

    rs.server.Get("/keymap", func (c *web.Context) string {
      c.ContentType("json")
      return rs.KeyMap(c)
    })

    fmt.Printf("Listening on %v\n", name)
    rs.server.Run(name)
  }();

  // create a thread to do intelligent replication every now and again
  go func() {
    time.Sleep(10 * time.Second)
    for {
      rs.IntelligentReplication()
      time.Sleep(REPL_INTERVAL)
    }
  }()

  return rs
}

<commit_msg>Making replication less aggressive<commit_after>package modata

import (
  "fmt"
  "web"
  "sort"
  "time"
)

const REPL_INTERVAL = 10 * time.Second

type ReplicationServer struct {
  MyContact Contact

  server *web.Server
  blockServer *BlockServer
}

func BuildKeyMap(contacts ContactList) map[string][]string {
  keyMapping := make(map[string][]string)
  for _, contact := range contacts {
    status, data, _ := JsonGet(contact.ToHttpAddress() + "/keys", Contact{})
    if (status == OK) {
      for _, key := range data.([]interface{}) {
        skey := key.(string)
        lst, ok := keyMapping[skey]
        if ok {
          keyMapping[skey] = append(lst, contact.ToHttpAddress())
        } else {
          keyMapping[skey] = []string{contact.ToHttpAddress()}
        }
      }
    }
  }
  return keyMapping
}

func (rs *ReplicationServer) IntelligentReplication () {
  contacts := rs.blockServer.routingTable.AllContacts()
  contacts = append(contacts, rs.blockServer.contact)
  keyMapping := BuildKeyMap(contacts)

  // Stores repl_count -> keys
  replicationMap := make(map[int][]string)
  for key, value := range keyMapping {
    replicationCount := len(value)
    lst, ok := replicationMap[replicationCount]
    if ok {
      replicationMap[replicationCount] = append(lst, key)
    } else {
      replicationMap[replicationCount] = []string{key}
    }
  }

  replicationCounts := make([]int,0)
  for key, _ := range replicationMap {
    replicationCounts = append(replicationCounts, key)
  }
  sort.Ints(replicationCounts)

  max := 3
  for _, value := range replicationCounts {
    if value >= rs.blockServer.routingTable.k - 1 {
      // Don't replicate fully replicated nodes
      break
    }
    keysToReplicate := replicationMap[value]
    for _, key := range keysToReplicate {
      fmt.Printf("Re-replicating %v with replication %v\n", key, value)
      server := keyMapping[key][0]
      JsonPostUrl(server + "/replicate?key=" + key, Contact{})
    }
    max -= 1
    if (max <= 0) {
      break
    }
  }

}

func (rs *ReplicationServer) KeyMap(c *web.Context) string{
  contacts := rs.blockServer.routingTable.AllContacts()
  contacts = append(contacts, rs.blockServer.contact)
  keyMapping := BuildKeyMap(contacts)
  finalData := make(map[string]interface{})
  finalData["stats"] = map[string]int{"numKeys": len(keyMapping), "contacts": len(contacts),
                        "averageKeysPerContact": len(keyMapping) /
                        len(contacts)}
  finalData["keys"] = make(map[string]interface{})
  for key, value := range keyMapping {
    sort.Strings(value)
    finalData["keys"].(map[string]interface{})[key] = make(map[string]interface{})
    finalData["keys"].(map[string]interface{})[key].(map[string]interface{})["replicationCount"] = len(value)
    finalData["keys"].(map[string]interface{})[key].(map[string]interface{})["nodes"] = value
  }
  return RespondWithData(finalData)
}

func StartReplicationServer(name string, bs *BlockServer) *ReplicationServer{
  rs := new(ReplicationServer)

  rs.server = web.NewServer()
  rs.blockServer = bs

  go func() {
    // Identifier for this node

    rs.server.Get("/contacts", func (c *web.Context) string {
      c.ContentType("json")
      return RespondWithData(bs.routingTable.AllContacts())
    })

    rs.server.Get("/keymap", func (c *web.Context) string {
      c.ContentType("json")
      return rs.KeyMap(c)
    })

    fmt.Printf("Listening on %v\n", name)
    rs.server.Run(name)
  }();

  // create a thread to do intelligent replication every now and again
  go func() {
    time.Sleep(10 * time.Second)
    for {
      rs.IntelligentReplication()
      time.Sleep(REPL_INTERVAL)
    }
  }()

  return rs
}

<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package saltpack

import "fmt"

type receiverKeys struct {
	_struct       bool   `codec:",toarray"`
	ReceiverKID   []byte `codec:"receiver_key_id"`
	PayloadKeyBox []byte `codec:"payloadkey"`
}

// Version is a major.minor pair that shows the version of the whole file
type Version struct {
	_struct bool `codec:",toarray"`
	Major   int  `codec:"major"`
	Minor   int  `codec:"minor"`
}

// encryptionHeader is the first packet in an encrypted message.
// It contains the encryptions of the session keys, and various
// message metadata.
type encryptionHeader struct {
	_struct         bool           `codec:",toarray"`
	FormatName      string         `codec:"format_name"`
	Version         Version        `codec:"vers"`
	Type            MessageType    `codec:"type"`
	Ephemeral       []byte         `codec:"ephemeral"`
	SenderSecretbox []byte         `codec:"sendersecretbox"`
	Receivers       []receiverKeys `codec:"rcvrs"`
	seqno           packetSeqno
}

// encryptionBlock contains a block of encrypted data. It contains
// the ciphertext, and any necessary authentication Tags.
type encryptionBlock struct {
	_struct            bool     `codec:",toarray"`
	HashAuthenticators [][]byte `codec:"authenticators"`
	PayloadCiphertext  []byte   `codec:"ctext"`
	seqno              packetSeqno
}

func (h *encryptionHeader) validate() error {
	if h.Type != MessageTypeEncryption {
		return ErrWrongMessageType{MessageTypeEncryption, h.Type}
	}
	if h.Version.Major != SaltpackCurrentVersion.Major {
		return ErrBadVersion{h.Version}
	}
	return nil
}

// SignatureHeader is the first packet in a signed message.
type SignatureHeader struct {
	_struct      bool        `codec:",toarray"`
	FormatName   string      `codec:"format_name"`
	Version      Version     `codec:"vers"`
	Type         MessageType `codec:"type"`
	SenderPublic []byte      `codec:"sender_public"`
	Nonce        []byte      `codec:"nonce"`
}

func newSignatureHeader(sender SigningPublicKey, msgType MessageType) (*SignatureHeader, error) {
	if sender == nil {
		return nil, ErrInvalidParameter{message: "no public signing key provided"}
	}
	nonce, err := newSigNonce()
	if err != nil {
		return nil, err
	}

	header := &SignatureHeader{
		FormatName:   SaltpackFormatName,
		Version:      SaltpackCurrentVersion,
		Type:         msgType,
		SenderPublic: sender.ToKID(),
		Nonce:        nonce[:],
	}

	return header, nil
}

func (h *SignatureHeader) validate(msgType MessageType) error {
	if h.Type != msgType {
		return ErrWrongMessageType{
			wanted:   msgType,
			received: h.Type,
		}
	}
	if h.Version.Major != SaltpackCurrentVersion.Major {
		return ErrBadVersion{h.Version}
	}

	if msgType != MessageTypeAttachedSignature && msgType != MessageTypeDetachedSignature {
		return ErrInvalidParameter{message: fmt.Sprintf("signature header must be MessageTypeAttachedSignature or MessageTypeDetachedSignature, not %d", msgType)}
	}

	return nil
}

// signatureBlock contains a block of signed data.
type signatureBlock struct {
	_struct      bool   `codec:",toarray"`
	Signature    []byte `codec:"signature"`
	PayloadChunk []byte `codec:"payload_chunk"`
	seqno        packetSeqno
}
<commit_msg>Add anonymous receiver option to saltpack encrypt<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package saltpack

import "fmt"

type receiverKeys struct {
	_struct       bool   `codec:",toarray"`
	ReceiverKID   []byte `codec:"receiver_key_id"`
	PayloadKeyBox []byte `codec:"payloadkey"`
}

// Version is a major.minor pair that shows the version of the whole file
type Version struct {
	_struct bool `codec:",toarray"`
	Major   int  `codec:"major"`
	Minor   int  `codec:"minor"`
}

// encryptionHeader is the first packet in an encrypted message.
// It contains the encryptions of the session keys, and various
// message metadata.
type encryptionHeader struct {
	_struct         bool           `codec:",toarray"`
	FormatName      string         `codec:"format_name"`
	Version         Version        `codec:"vers"`
	Type            MessageType    `codec:"type"`
	Ephemeral       []byte         `codec:"ephemeral"`
	SenderSecretbox []byte         `codec:"sendersecretbox"`
	Receivers       []receiverKeys `codec:"rcvrs,omitempty"`
	seqno           packetSeqno
}

// encryptionBlock contains a block of encrypted data. It contains
// the ciphertext, and any necessary authentication Tags.
type encryptionBlock struct {
	_struct            bool     `codec:",toarray"`
	HashAuthenticators [][]byte `codec:"authenticators"`
	PayloadCiphertext  []byte   `codec:"ctext"`
	seqno              packetSeqno
}

func (h *encryptionHeader) validate() error {
	if h.Type != MessageTypeEncryption {
		return ErrWrongMessageType{MessageTypeEncryption, h.Type}
	}
	if h.Version.Major != SaltpackCurrentVersion.Major {
		return ErrBadVersion{h.Version}
	}
	return nil
}

// SignatureHeader is the first packet in a signed message.
type SignatureHeader struct {
	_struct      bool        `codec:",toarray"`
	FormatName   string      `codec:"format_name"`
	Version      Version     `codec:"vers"`
	Type         MessageType `codec:"type"`
	SenderPublic []byte      `codec:"sender_public"`
	Nonce        []byte      `codec:"nonce"`
}

func newSignatureHeader(sender SigningPublicKey, msgType MessageType) (*SignatureHeader, error) {
	if sender == nil {
		return nil, ErrInvalidParameter{message: "no public signing key provided"}
	}
	nonce, err := newSigNonce()
	if err != nil {
		return nil, err
	}

	header := &SignatureHeader{
		FormatName:   SaltpackFormatName,
		Version:      SaltpackCurrentVersion,
		Type:         msgType,
		SenderPublic: sender.ToKID(),
		Nonce:        nonce[:],
	}

	return header, nil
}

func (h *SignatureHeader) validate(msgType MessageType) error {
	if h.Type != msgType {
		return ErrWrongMessageType{
			wanted:   msgType,
			received: h.Type,
		}
	}
	if h.Version.Major != SaltpackCurrentVersion.Major {
		return ErrBadVersion{h.Version}
	}

	if msgType != MessageTypeAttachedSignature && msgType != MessageTypeDetachedSignature {
		return ErrInvalidParameter{message: fmt.Sprintf("signature header must be MessageTypeAttachedSignature or MessageTypeDetachedSignature, not %d", msgType)}
	}

	return nil
}

// signatureBlock contains a block of signed data.
type signatureBlock struct {
	_struct      bool   `codec:",toarray"`
	Signature    []byte `codec:"signature"`
	PayloadChunk []byte `codec:"payload_chunk"`
	seqno        packetSeqno
}
<|endoftext|>


<commit_before>// json-tidy pretty prints JSON
package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
)

var (
	prefix = flag.String("prefix", "", "Prefix string")
	indent = flag.String("indent", "\t", "Identation string")
)

func die(err error) {
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func main() {
	flag.Usage = func() {
		fmt.Fprint(os.Stderr, `Usage of json-tidy:

json-tidy [opts] [file|url|-]
        Gets input (defaults to stdin) and prints clean json to stdout.
`)
		flag.PrintDefaults()
	}
	flag.Parse()

	if flag.NArg() > 1 {
		flag.Usage()
		die(errors.New("Too many arguments"))
	}

	var src io.Reader = os.Stdin
	if arg := flag.Arg(0); arg != "" && arg != "-" {
		if u, err := url.Parse(arg); err == nil &&
			// It's a URL
			u.Scheme == "http" || u.Scheme == "https" {
			rsp, err := http.Get(arg)
			die(err)
			defer rsp.Body.Close()
			src = rsp.Body
		} else {
			// It's a file
			f, err := os.Open(arg)
			die(err)
			defer f.Close()
			src = f
		}
	}

	dec := json.NewDecoder(src)
	dec.UseNumber() // Preserve number formatting

	var data interface{}

	for dec.More() {
		die(dec.Decode(&data))

		b, err := json.MarshalIndent(&data, *prefix, *indent)
		die(err)

		_, err = os.Stdout.Write(b)
		die(err)
	}
	// Trailing newline
	os.Stdout.Write([]byte{'\n'})
}
<commit_msg>Main: Invert control, remove die function, safely defer close<commit_after>// json-tidy pretty prints JSON
package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
)

var (
	prefix = flag.String("prefix", "", "Prefix string")
	indent = flag.String("indent", "\t", "Identation string")
)

func main() {
	if err := Run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func Run() (err error) {
	flag.Usage = func() {
		fmt.Fprint(os.Stderr, `Usage of json-tidy:

json-tidy [opts] [file|url|-]
        Gets input (defaults to stdin) and prints clean json to stdout.
`)
		flag.PrintDefaults()
	}
	flag.Parse()

	if flag.NArg() > 1 {
		flag.Usage()
		return errors.New("Too many arguments")
	}

	var src io.Reader = os.Stdin
	if arg := flag.Arg(0); arg != "" && arg != "-" {
		if u, err := url.Parse(arg); err == nil &&
			// It's a URL
			u.Scheme == "http" || u.Scheme == "https" {
			rsp, err := http.Get(arg)
			if err != nil {
				return err
			}
			defer DeferClose(&err, rsp.Body.Close)
			src = rsp.Body
		} else {
			// It's a file
			f, err := os.Open(arg)
			if err != nil {
				return err
			}
			defer DeferClose(&err, f.Close)
			src = f
		}
	}

	dec := json.NewDecoder(src)
	dec.UseNumber() // Preserve number formatting

	var data interface{}

	for dec.More() {
		err = dec.Decode(&data)
		if err != nil {
			return err
		}

		b, err := json.MarshalIndent(&data, *prefix, *indent)
		if err != nil {
			return err
		}

		_, err = os.Stdout.Write(b)
		if err != nil {
			return err
		}
	}
	// Trailing newline
	_, err = os.Stdout.Write([]byte{'\n'})
	return err
}

func DeferClose(err *error, f func() error) {
	newErr := f()
	if *err == nil {
		*err = newErr
	}
}
<|endoftext|>


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

// +build host coreos src kvm

package main

import (
	"fmt"
	"os"
	"testing"

	"github.com/coreos/rkt/common"
	"github.com/coreos/rkt/tests/testutils"
)

// TestNonRootReadInfo tests that non-root users that can do rkt list, rkt image list.
func TestNonRootReadInfo(t *testing.T) {
	if !common.SupportsUserNS() {
		t.Skip("User namespaces are not supported on this host.")
	}

	if err := checkUserNS(); err != nil {
		t.Skip("User namespaces don't work on this host.")
	}

	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	uid, _ := ctx.GetUidGidRktBinOwnerNotRoot()

	rktGid, err := common.LookupGid(common.RktGroup)
	if err != nil {
		t.Skipf("Skipping the test because there's no %q group", common.RktGroup)
	}

	if err := ctx.SetupDataDir(); err != nil {
		t.Fatalf("failed to setup data dir: %v", err)
	}

	// Launch some pods, this creates the environment for later testing.
	imgs := []struct {
		name    string
		msg     string
		imgFile string
	}{
		{name: "inspect-1", msg: "foo-1"},
		{name: "inspect-2", msg: "foo-2"},
		{name: "inspect-3", msg: "foo-3"},
	}

	for i, img := range imgs {
		imgName := fmt.Sprintf("rkt-%s.aci", img.name)
		imgs[i].imgFile = patchTestACI(imgName, fmt.Sprintf("--name=%s", img.name), fmt.Sprintf("--exec=/inspect --print-msg=%s --exit-code=0", img.msg))
		defer os.Remove(imgs[i].imgFile)
	}

	runCmds := []string{
		// Run with overlay, without private-users.
		fmt.Sprintf("%s --insecure-options=image run --mds-register=false %s", ctx.Cmd(), imgs[0].imgFile),

		// Run without overlay, without private-users.
		fmt.Sprintf("%s --insecure-options=image run --no-overlay --mds-register=false %s", ctx.Cmd(), imgs[1].imgFile),

		// Run without overlay, with private-users.
		fmt.Sprintf("%s --insecure-options=image run --no-overlay --private-users --mds-register=false %s", ctx.Cmd(), imgs[2].imgFile),
	}

	for i, cmd := range runCmds {
		t.Logf("#%d: Running %s", i, cmd)
		runRktAndCheckOutput(t, cmd, imgs[i].msg, false)
	}

	imgListCmd := fmt.Sprintf("%s image list", ctx.Cmd())
	t.Logf("Running %s", imgListCmd)
	runRktAsUidGidAndCheckOutput(t, imgListCmd, "inspect-", false, uid, rktGid)
}

// TestNonRootFetchRmGCImage tests that non-root users can remove all images but
// cannot gc images.
func TestNonRootFetchRmGCImage(t *testing.T) {
	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	uid, _ := ctx.GetUidGidRktBinOwnerNotRoot()

	rktGid, err := common.LookupGid(common.RktGroup)
	if err != nil {
		t.Skipf("Skipping the test because there's no %q group", common.RktGroup)
	}

	if err := ctx.SetupDataDir(); err != nil {
		t.Fatalf("failed to setup data dir: %v", err)
	}

	rootImg := patchTestACI("rkt-inspect-root-rm.aci", "--exec=/inspect --print-msg=foobar")
	defer os.Remove(rootImg)
	rootImgHash, err := importImageAndFetchHash(t, ctx, "", rootImg)
	if err != nil {
		t.Fatalf("%v", err)
	}

	// Launch/gc a pod so we can test non-root image gc.
	runCmd := fmt.Sprintf("%s --insecure-options=image run --mds-register=false %s", ctx.Cmd(), rootImg)
	runRktAndCheckOutput(t, runCmd, "foobar", false)

	ctx.RunGC()

	// Should not be able to do image gc.
	// We can't touch the treestores even as members of the rkt group.
	imgGCCmd := fmt.Sprintf("%s image gc", ctx.Cmd())
	t.Logf("Running %s", imgGCCmd)
	runRktAsUidGidAndCheckOutput(t, imgGCCmd, "permission denied", false, true, uid, rktGid)

	// Should be able to remove the image fetched by root since we're in the rkt group.
	imgRmCmd := fmt.Sprintf("%s image rm %s", ctx.Cmd(), rootImgHash)
	t.Logf("Running %s", imgRmCmd)
	runRktAsUidGidAndCheckOutput(t, imgRmCmd, "successfully removed", false, false, uid, rktGid)

	// Should be able to remove the image fetched by ourselves.
	nonrootImg := patchTestACI("rkt-inspect-non-root-rm.aci", "--exec=/inspect")
	defer os.Remove(nonrootImg)
	nonrootImgHash, err := importImageAndFetchHashAsUidGid(t, ctx, nonrootImg, "", uid, rktGid)
	if err != nil {
		t.Fatalf("%v", err)
	}

	imgRmCmd = fmt.Sprintf("%s image rm %s", ctx.Cmd(), nonrootImgHash)
	t.Logf("Running %s", imgRmCmd)
	runRktAsUidGidAndCheckOutput(t, imgRmCmd, "successfully removed", false, false, uid, rktGid)
}
<commit_msg>tests: fix runRktAsUidGidAndCheckOutput parameters<commit_after>// Copyright 2015 The rkt Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build host coreos src kvm

package main

import (
	"fmt"
	"os"
	"testing"

	"github.com/coreos/rkt/common"
	"github.com/coreos/rkt/tests/testutils"
)

// TestNonRootReadInfo tests that non-root users that can do rkt list, rkt image list.
func TestNonRootReadInfo(t *testing.T) {
	if !common.SupportsUserNS() {
		t.Skip("User namespaces are not supported on this host.")
	}

	if err := checkUserNS(); err != nil {
		t.Skip("User namespaces don't work on this host.")
	}

	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	uid, _ := ctx.GetUidGidRktBinOwnerNotRoot()

	rktGid, err := common.LookupGid(common.RktGroup)
	if err != nil {
		t.Skipf("Skipping the test because there's no %q group", common.RktGroup)
	}

	if err := ctx.SetupDataDir(); err != nil {
		t.Fatalf("failed to setup data dir: %v", err)
	}

	// Launch some pods, this creates the environment for later testing.
	imgs := []struct {
		name    string
		msg     string
		imgFile string
	}{
		{name: "inspect-1", msg: "foo-1"},
		{name: "inspect-2", msg: "foo-2"},
		{name: "inspect-3", msg: "foo-3"},
	}

	for i, img := range imgs {
		imgName := fmt.Sprintf("rkt-%s.aci", img.name)
		imgs[i].imgFile = patchTestACI(imgName, fmt.Sprintf("--name=%s", img.name), fmt.Sprintf("--exec=/inspect --print-msg=%s --exit-code=0", img.msg))
		defer os.Remove(imgs[i].imgFile)
	}

	runCmds := []string{
		// Run with overlay, without private-users.
		fmt.Sprintf("%s --insecure-options=image run --mds-register=false %s", ctx.Cmd(), imgs[0].imgFile),

		// Run without overlay, without private-users.
		fmt.Sprintf("%s --insecure-options=image run --no-overlay --mds-register=false %s", ctx.Cmd(), imgs[1].imgFile),

		// Run without overlay, with private-users.
		fmt.Sprintf("%s --insecure-options=image run --no-overlay --private-users --mds-register=false %s", ctx.Cmd(), imgs[2].imgFile),
	}

	for i, cmd := range runCmds {
		t.Logf("#%d: Running %s", i, cmd)
		runRktAndCheckOutput(t, cmd, imgs[i].msg, false)
	}

	imgListCmd := fmt.Sprintf("%s image list", ctx.Cmd())
	t.Logf("Running %s", imgListCmd)
	runRktAsUidGidAndCheckOutput(t, imgListCmd, "inspect-", false, false, uid, rktGid)
}

// TestNonRootFetchRmGCImage tests that non-root users can remove all images but
// cannot gc images.
func TestNonRootFetchRmGCImage(t *testing.T) {
	ctx := testutils.NewRktRunCtx()
	defer ctx.Cleanup()

	uid, _ := ctx.GetUidGidRktBinOwnerNotRoot()

	rktGid, err := common.LookupGid(common.RktGroup)
	if err != nil {
		t.Skipf("Skipping the test because there's no %q group", common.RktGroup)
	}

	if err := ctx.SetupDataDir(); err != nil {
		t.Fatalf("failed to setup data dir: %v", err)
	}

	rootImg := patchTestACI("rkt-inspect-root-rm.aci", "--exec=/inspect --print-msg=foobar")
	defer os.Remove(rootImg)
	rootImgHash, err := importImageAndFetchHash(t, ctx, "", rootImg)
	if err != nil {
		t.Fatalf("%v", err)
	}

	// Launch/gc a pod so we can test non-root image gc.
	runCmd := fmt.Sprintf("%s --insecure-options=image run --mds-register=false %s", ctx.Cmd(), rootImg)
	runRktAndCheckOutput(t, runCmd, "foobar", false)

	ctx.RunGC()

	// Should not be able to do image gc.
	// We can't touch the treestores even as members of the rkt group.
	imgGCCmd := fmt.Sprintf("%s image gc", ctx.Cmd())
	t.Logf("Running %s", imgGCCmd)
	runRktAsUidGidAndCheckOutput(t, imgGCCmd, "permission denied", false, true, uid, rktGid)

	// Should be able to remove the image fetched by root since we're in the rkt group.
	imgRmCmd := fmt.Sprintf("%s image rm %s", ctx.Cmd(), rootImgHash)
	t.Logf("Running %s", imgRmCmd)
	runRktAsUidGidAndCheckOutput(t, imgRmCmd, "successfully removed", false, false, uid, rktGid)

	// Should be able to remove the image fetched by ourselves.
	nonrootImg := patchTestACI("rkt-inspect-non-root-rm.aci", "--exec=/inspect")
	defer os.Remove(nonrootImg)
	nonrootImgHash, err := importImageAndFetchHashAsUidGid(t, ctx, nonrootImg, "", uid, rktGid)
	if err != nil {
		t.Fatalf("%v", err)
	}

	imgRmCmd = fmt.Sprintf("%s image rm %s", ctx.Cmd(), nonrootImgHash)
	t.Logf("Running %s", imgRmCmd)
	runRktAsUidGidAndCheckOutput(t, imgRmCmd, "successfully removed", false, false, uid, rktGid)
}
<|endoftext|>


<commit_before>// Copyright 2014 go-dockerclient authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package docker

import (
	"bufio"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func TestEventListeners(t *testing.T) {
	response := `{"status":"create","id":"dfdf82bd3881","from":"base:latest","time":1374067924}
{"status":"start","id":"dfdf82bd3881","from":"base:latest","time":1374067924}
{"status":"stop","id":"dfdf82bd3881","from":"base:latest","time":1374067966}
{"status":"destroy","id":"dfdf82bd3881","from":"base:latest","time":1374067970}
`

	var req http.Request
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		rsc := bufio.NewScanner(strings.NewReader(response))
		for rsc.Scan() {
			w.Write([]byte(rsc.Text()))
			w.(http.Flusher).Flush()
			time.Sleep(10 * time.Millisecond)
		}
		req = *r
	}))
	defer server.Close()

	client, err := NewClient(server.URL)
	if err != nil {
		t.Errorf("Failed to create client: %s", err)
	}
	client.SkipServerVersionCheck = true

	listener := make(chan *APIEvents, 10)
	defer func() { time.Sleep(10 * time.Millisecond); client.RemoveEventListener(listener) }()

	err = client.AddEventListener(listener)
	if err != nil {
		t.Errorf("Failed to add event listener: %s", err)
	}

	timeout := time.After(1 * time.Second)
	var count int

	for {
		select {
		case msg := <-listener:
			t.Logf("Recieved: %s", *msg)
			count++
			err = checkEvent(count, msg)
			if err != nil {
				t.Fatalf("Check event failed: %s", err)
			}
			if count == 4 {
				return
			}
		case <-timeout:
			t.Fatal("TestEventListeners timed out waiting on events")
		}
	}
}

func checkEvent(index int, event *APIEvents) error {
	if event.ID != "dfdf82bd3881" {
		return fmt.Errorf("event ID did not match. Expected dfdf82bd3881 got %s", event.ID)
	}
	if event.From != "base:latest" {
		return fmt.Errorf("event from did not match. Expected base:latest got %s", event.From)
	}
	var status string
	switch index {
	case 1:
		status = "create"
	case 2:
		status = "start"
	case 3:
		status = "stop"
	case 4:
		status = "destroy"
	}
	if event.Status != status {
		return fmt.Errorf("event status did not match. Expected %s got %s", status, event.Status)
	}
	return nil
}
<commit_msg>Add a test for TLS events<commit_after>// Copyright 2014 go-dockerclient authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package docker

import (
	"bufio"
	"crypto/tls"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func TestEventListeners(t *testing.T) {
	response := `{"status":"create","id":"dfdf82bd3881","from":"base:latest","time":1374067924}
{"status":"start","id":"dfdf82bd3881","from":"base:latest","time":1374067924}
{"status":"stop","id":"dfdf82bd3881","from":"base:latest","time":1374067966}
{"status":"destroy","id":"dfdf82bd3881","from":"base:latest","time":1374067970}
`

	var req http.Request
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		rsc := bufio.NewScanner(strings.NewReader(response))
		for rsc.Scan() {
			w.Write([]byte(rsc.Text()))
			w.(http.Flusher).Flush()
			time.Sleep(10 * time.Millisecond)
		}
		req = *r
	}))
	defer server.Close()

	client, err := NewClient(server.URL)
	if err != nil {
		t.Errorf("Failed to create client: %s", err)
	}
	client.SkipServerVersionCheck = true

	listener := make(chan *APIEvents, 10)
	defer func() { time.Sleep(10 * time.Millisecond); client.RemoveEventListener(listener) }()

	err = client.AddEventListener(listener)
	if err != nil {
		t.Errorf("Failed to add event listener: %s", err)
	}

	timeout := time.After(1 * time.Second)
	var count int

	for {
		select {
		case msg := <-listener:
			t.Logf("Recieved: %s", *msg)
			count++
			err = checkEvent(count, msg)
			if err != nil {
				t.Fatalf("Check event failed: %s", err)
			}
			if count == 4 {
				return
			}
		case <-timeout:
			t.Fatal("TestEventListeners timed out waiting on events")
		}
	}
}

func TestTLSEventListeners(t *testing.T) {
	response := `{"status":"create","id":"dfdf82bd3881","from":"base:latest","time":1374067924}
{"status":"start","id":"dfdf82bd3881","from":"base:latest","time":1374067924}
{"status":"stop","id":"dfdf82bd3881","from":"base:latest","time":1374067966}
{"status":"destroy","id":"dfdf82bd3881","from":"base:latest","time":1374067970}
`

	var req http.Request
	server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		rsc := bufio.NewScanner(strings.NewReader(response))
		for rsc.Scan() {
			w.Write([]byte(rsc.Text()))
			w.(http.Flusher).Flush()
			time.Sleep(10 * time.Millisecond)
		}
		req = *r
	}))
	server.TLS = &tls.Config{InsecureSkipVerify: true}
	server.StartTLS()
	defer server.Close()

	client, err := NewTLSClient(server.URL, "testing/data/cert.pem", "testing/data/key.pem", "")
	if err != nil {
		t.Errorf("Failed to create client: %s", err)
	}
	client.SkipServerVersionCheck = true

	listener := make(chan *APIEvents, 10)
	defer func() { time.Sleep(10 * time.Millisecond); client.RemoveEventListener(listener) }()

	err = client.AddEventListener(listener)
	if err != nil {
		t.Errorf("Failed to add event listener: %s", err)
	}

	timeout := time.After(1 * time.Second)
	var count int

	for {
		select {
		case msg := <-listener:
			t.Logf("Recieved: %s", *msg)
			count++
			err = checkEvent(count, msg)
			if err != nil {
				t.Fatalf("Check event failed: %s", err)
			}
			if count == 4 {
				return
			}
		case <-timeout:
			t.Fatal("TestTLSEventListeners timed out waiting on events")
		}
	}
}

func checkEvent(index int, event *APIEvents) error {
	if event.ID != "dfdf82bd3881" {
		return fmt.Errorf("event ID did not match. Expected dfdf82bd3881 got %s", event.ID)
	}
	if event.From != "base:latest" {
		return fmt.Errorf("event from did not match. Expected base:latest got %s", event.From)
	}
	var status string
	switch index {
	case 1:
		status = "create"
	case 2:
		status = "start"
	case 3:
		status = "stop"
	case 4:
		status = "destroy"
	}
	if event.Status != status {
		return fmt.Errorf("event status did not match. Expected %s got %s", status, event.Status)
	}
	return nil
}
<|endoftext|>


<commit_before>package workercmd

import (
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/localip"
	"fmt"
	concourseCmd "github.com/concourse/concourse/cmd"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
)

// This runs Guardian using the gdn binary
func (cmd *WorkerCommand) guardianRunner(logger lager.Logger) (ifrit.Runner, error) {
	depotDir := filepath.Join(cmd.WorkDir.Path(), "depot")

	// must be readable by other users so unprivileged containers can run their
	// own `initc' process
	err := os.MkdirAll(depotDir, 0755)
	if err != nil {
		return nil, err
	}

	members := grouper.Members{}

	gdnFlags := []string{}

	if cmd.Garden.Config.Path() != "" {
		gdnFlags = append(gdnFlags, "--config", cmd.Garden.Config.Path())
	}

	gdnServerFlags := []string{
		"--bind-ip", cmd.BindIP.IP.String(),
		"--bind-port", fmt.Sprintf("%d", cmd.BindPort),

		"--depot", depotDir,
		"--properties-path", filepath.Join(cmd.WorkDir.Path(), "garden-properties.json"),

		"--time-format", "rfc3339",

		// disable graph and grootfs setup; all images passed to Concourse
		// containers are raw://
		"--no-image-plugin",
	}

	for _, dnsServer := range cmd.Garden.DNSServers {
		gdnServerFlags = append(gdnServerFlags, "--dns-server", dnsServer)
	}

	if cmd.ContainerNetworkPool != "" {
		gdnServerFlags = append(gdnServerFlags, "--network-pool", cmd.ContainerNetworkPool)
	}

	gdnServerFlags = append(gdnServerFlags, detectGardenFlags(logger)...)

	if cmd.Garden.DNS.Enable {
		dnsProxyRunner, err := cmd.dnsProxyRunner(logger.Session("dns-proxy"))
		if err != nil {
			return nil, err
		}

		lip, err := localip.LocalIP()
		if err != nil {
			return nil, err
		}

		members = append(members, grouper.Member{
			Name: "dns-proxy",
			Runner: concourseCmd.NewLoggingRunner(
				logger.Session("dns-proxy-runner"),
				dnsProxyRunner,
			),
		})

		gdnServerFlags = append(gdnServerFlags, "--dns-server", lip)

		// must permit access to host network in order for DNS proxy address to be
		// reachable
		gdnServerFlags = append(gdnServerFlags, "--allow-host-access")
	}

	gdnArgs := append(gdnFlags, append([]string{"server"}, gdnServerFlags...)...)

	bin := "gdn"
	if cmd.Garden.Bin != "" {
		bin = cmd.Garden.Bin
	}

	gdnCmd := exec.Command(bin, gdnArgs...)
	gdnCmd.Stdout = os.Stdout
	gdnCmd.Stderr = os.Stderr
	gdnCmd.SysProcAttr = &syscall.SysProcAttr{
		Pdeathsig: syscall.SIGKILL,
	}

	members = append(members, grouper.Member{
		Name: "gdn",
		Runner: concourseCmd.NewLoggingRunner(
			logger.Session("gdn-runner"),
			CmdRunner{gdnCmd},
		),
	})

	return grouper.NewParallel(os.Interrupt, members), nil
}

var gardenEnvPrefix = "CONCOURSE_GARDEN_"

func detectGardenFlags(logger lager.Logger) []string {
	env := os.Environ()

	flags := []string{}
	for _, e := range env {
		spl := strings.SplitN(e, "=", 2)
		if len(spl) != 2 {
			logger.Info("bogus-env", lager.Data{"env": spl})
			continue
		}

		name := spl[0]
		val := spl[1]

		if !strings.HasPrefix(name, gardenEnvPrefix) {
			continue
		}

		strip := strings.Replace(name, gardenEnvPrefix, "", 1)
		flag := flagify(strip)

		logger.Info("forwarding-garden-env-var", lager.Data{
			"env":  name,
			"flag": flag,
		})

		vals := strings.Split(val, ",")

		for _, v := range vals {
			flags = append(flags, "--"+flag, v)
		}

		// clear out env (as twentythousandtonnesofcrudeoil does)
		_ = os.Unsetenv(name)
	}

	return flags
}

func flagify(env string) string {
	return strings.Replace(strings.ToLower(env), "_", "-", -1)
}
<commit_msg>workercmd: add linux build flag for guardian<commit_after>// +build linux

package workercmd

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"

	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/localip"
	concourseCmd "github.com/concourse/concourse/cmd"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
)

// This runs Guardian using the gdn binary
func (cmd *WorkerCommand) guardianRunner(logger lager.Logger) (ifrit.Runner, error) {
	depotDir := filepath.Join(cmd.WorkDir.Path(), "depot")

	// must be readable by other users so unprivileged containers can run their
	// own `initc' process
	err := os.MkdirAll(depotDir, 0755)
	if err != nil {
		return nil, err
	}

	members := grouper.Members{}

	gdnFlags := []string{}

	if cmd.Garden.Config.Path() != "" {
		gdnFlags = append(gdnFlags, "--config", cmd.Garden.Config.Path())
	}

	gdnServerFlags := []string{
		"--bind-ip", cmd.BindIP.IP.String(),
		"--bind-port", fmt.Sprintf("%d", cmd.BindPort),

		"--depot", depotDir,
		"--properties-path", filepath.Join(cmd.WorkDir.Path(), "garden-properties.json"),

		"--time-format", "rfc3339",

		// disable graph and grootfs setup; all images passed to Concourse
		// containers are raw://
		"--no-image-plugin",
	}

	for _, dnsServer := range cmd.Garden.DNSServers {
		gdnServerFlags = append(gdnServerFlags, "--dns-server", dnsServer)
	}

	if cmd.ContainerNetworkPool != "" {
		gdnServerFlags = append(gdnServerFlags, "--network-pool", cmd.ContainerNetworkPool)
	}

	gdnServerFlags = append(gdnServerFlags, detectGardenFlags(logger)...)

	if cmd.Garden.DNS.Enable {
		dnsProxyRunner, err := cmd.dnsProxyRunner(logger.Session("dns-proxy"))
		if err != nil {
			return nil, err
		}

		lip, err := localip.LocalIP()
		if err != nil {
			return nil, err
		}

		members = append(members, grouper.Member{
			Name: "dns-proxy",
			Runner: concourseCmd.NewLoggingRunner(
				logger.Session("dns-proxy-runner"),
				dnsProxyRunner,
			),
		})

		gdnServerFlags = append(gdnServerFlags, "--dns-server", lip)

		// must permit access to host network in order for DNS proxy address to be
		// reachable
		gdnServerFlags = append(gdnServerFlags, "--allow-host-access")
	}

	gdnArgs := append(gdnFlags, append([]string{"server"}, gdnServerFlags...)...)

	bin := "gdn"
	if cmd.Garden.Bin != "" {
		bin = cmd.Garden.Bin
	}

	gdnCmd := exec.Command(bin, gdnArgs...)
	gdnCmd.Stdout = os.Stdout
	gdnCmd.Stderr = os.Stderr
	gdnCmd.SysProcAttr = &syscall.SysProcAttr{
		Pdeathsig: syscall.SIGKILL,
	}

	members = append(members, grouper.Member{
		Name: "gdn",
		Runner: concourseCmd.NewLoggingRunner(
			logger.Session("gdn-runner"),
			CmdRunner{gdnCmd},
		),
	})

	return grouper.NewParallel(os.Interrupt, members), nil
}

var gardenEnvPrefix = "CONCOURSE_GARDEN_"

func detectGardenFlags(logger lager.Logger) []string {
	env := os.Environ()

	flags := []string{}
	for _, e := range env {
		spl := strings.SplitN(e, "=", 2)
		if len(spl) != 2 {
			logger.Info("bogus-env", lager.Data{"env": spl})
			continue
		}

		name := spl[0]
		val := spl[1]

		if !strings.HasPrefix(name, gardenEnvPrefix) {
			continue
		}

		strip := strings.Replace(name, gardenEnvPrefix, "", 1)
		flag := flagify(strip)

		logger.Info("forwarding-garden-env-var", lager.Data{
			"env":  name,
			"flag": flag,
		})

		vals := strings.Split(val, ",")

		for _, v := range vals {
			flags = append(flags, "--"+flag, v)
		}

		// clear out env (as twentythousandtonnesofcrudeoil does)
		_ = os.Unsetenv(name)
	}

	return flags
}

func flagify(env string) string {
	return strings.Replace(strings.ToLower(env), "_", "-", -1)
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package codegen

import (
	"fmt"
	"strings"

	"github.com/pkg/errors"
	"go.uber.org/thriftrw/compile"
)

// TypeConverter can generate a function body that converts two thriftrw
// FieldGroups from one to another. It's assumed that the converted code
// operates on two variables, "in" and "out" and that both are a go struct.
type TypeConverter struct {
	Lines  []string
	Helper PackageNameResolver
}

// PackageNameResolver interface allows for resolving what the
// package name for a thrift file is. This depends on where the
// thrift-based structs are generated.
type PackageNameResolver interface {
	TypePackageName(thriftFile string) (string, error)
}

func (c *TypeConverter) getGoTypeName(
	valueType compile.TypeSpec,
) (string, error) {
	switch s := valueType.(type) {
	case *compile.BoolSpec:
		return "bool", nil
	case *compile.I8Spec:
		return "int8", nil
	case *compile.I16Spec:
		return "int16", nil
	case *compile.I32Spec:
		return "int32", nil
	case *compile.I64Spec:
		return "int64", nil
	case *compile.DoubleSpec:
		return "float64", nil
	case *compile.StringSpec:
		return "string", nil
	case *compile.BinarySpec:
		return "[]byte", nil
	case *compile.MapSpec:
		panic("Not Implemented")
	case *compile.SetSpec:
		panic("Not Implemented")
	case *compile.ListSpec:
		panic("Not Implemented")
	case *compile.EnumSpec, *compile.StructSpec, *compile.TypedefSpec:
		return c.getIdentifierName(s)
	default:
		panic(fmt.Sprintf("Unknown type (%T) %v", valueType, valueType))
	}
}

func (c *TypeConverter) getIdentifierName(
	fieldType compile.TypeSpec,
) (string, error) {
	pkgName, err := c.Helper.TypePackageName(fieldType.ThriftFile())
	if err != nil {
		return "", errors.Wrapf(
			err,
			"could not lookup fieldType when building converter for %s :",
			fieldType.ThriftName(),
		)
	}
	typeName := pkgName + "." + fieldType.ThriftName()
	return typeName, nil
}

func (c *TypeConverter) genConverterForStruct(
	toFieldName string,
	toFieldType *compile.StructSpec,
	fromFieldType compile.TypeSpec,
	fromIdentifier string,
	keyPrefix string,
	indent string,
) error {
	toIdentifier := indent + "out." + keyPrefix

	typeName, err := c.getIdentifierName(toFieldType)
	if err != nil {
		return err
	}
	subToFields := toFieldType.Fields

	fromFieldStruct, ok := fromFieldType.(*compile.StructSpec)
	if !ok {
		return errors.Errorf(
			"could not convert struct fields, "+
				"incompatible type for %s :",
			toFieldName,
		)
	}

	line := "if " + fromIdentifier + " != nil {"
	c.Lines = append(c.Lines, line)

	line = "	" + toIdentifier + " = &" + typeName + "{}"
	c.Lines = append(c.Lines, line)

	subFromFields := fromFieldStruct.Fields
	err = c.genStructConverter(
		keyPrefix+".",
		indent+"	",
		subFromFields,
		subToFields,
	)
	if err != nil {
		return err
	}

	line = "} else {"
	c.Lines = append(c.Lines, line)

	line = "	" + toIdentifier + " = nil"
	c.Lines = append(c.Lines, line)

	line = "}"
	c.Lines = append(c.Lines, line)

	return nil
}

func (c *TypeConverter) genConverterForPrimitive(
	toField *compile.FieldSpec,
	toIdentifier string,
	fromIdentifier string,
) error {
	var line string
	typeName, err := c.getGoTypeName(toField.Type)
	if err != nil {
		return err
	}

	if toField.Required {
		line = toIdentifier + " = " + typeName + "(" + fromIdentifier + ")"
	} else {
		line = toIdentifier + " = (*" + typeName + ")(" + fromIdentifier + ")"
	}
	c.Lines = append(c.Lines, line)
	return nil
}

func (c *TypeConverter) genConverterForList(
	toFieldType *compile.ListSpec,
	toField *compile.FieldSpec,
	fromField *compile.FieldSpec,
	toIdentifier string,
	fromIdentifier string,
	keyPrefix string,
	indent string,
) error {
	typeName, err := c.getGoTypeName(toFieldType.ValueSpec)
	if err != nil {
		return err
	}

	valueStruct, isStruct := toFieldType.ValueSpec.(*compile.StructSpec)
	if isStruct {
		line := toIdentifier + " = make([]*" +
			typeName + ", len(" + fromIdentifier + "))"
		c.Lines = append(c.Lines, line)
	} else {
		line := toIdentifier + " = make([]" +
			typeName + ", len(" + fromIdentifier + "))"
		c.Lines = append(c.Lines, line)
	}

	line := "for index, value := range " + fromIdentifier + " {"
	c.Lines = append(c.Lines, line)

	if isStruct {
		fromFieldType, ok := fromField.Type.(*compile.ListSpec)
		if !ok {
			return errors.Errorf(
				"Could not convert field (%s): type is not list",
				fromField.Name,
			)
		}

		err = c.genConverterForStruct(
			toField.Name,
			valueStruct,
			fromFieldType.ValueSpec,
			"value",
			keyPrefix+strings.Title(toField.Name)+"[index]",
			indent,
		)
		if err != nil {
			return err
		}
	} else {
		line = toIdentifier + "[index] = " + typeName + "(value)"
		c.Lines = append(c.Lines, line)
	}

	line = "}"
	c.Lines = append(c.Lines, line)
	return nil
}

func (c *TypeConverter) genConverterForMap(
	toFieldType *compile.MapSpec,
	toField *compile.FieldSpec,
	fromField *compile.FieldSpec,
	toIdentifier string,
	fromIdentifier string,
	keyPrefix string,
	indent string,
) error {
	typeName, err := c.getGoTypeName(toFieldType.ValueSpec)
	if err != nil {
		return err
	}

	_, isStringKey := toFieldType.KeySpec.(*compile.StringSpec)
	if !isStringKey {
		return errors.Errorf(
			"could not convert key (%s), map is not string-keyed.",
			toField.Name,
		)
	}

	valueStruct, isStruct := toFieldType.ValueSpec.(*compile.StructSpec)
	if isStruct {
		line := toIdentifier + " = make(map[string]*" +
			typeName + ", len(" + fromIdentifier + "))"
		c.Lines = append(c.Lines, line)
	} else {
		line := toIdentifier + " = make(map[string]" +
			typeName + ", len(" + fromIdentifier + "))"
		c.Lines = append(c.Lines, line)
	}

	line := "for key, value := range " + fromIdentifier + " {"
	c.Lines = append(c.Lines, line)

	if isStruct {
		fromFieldType, ok := fromField.Type.(*compile.MapSpec)
		if !ok {
			return errors.Errorf(
				"Could not convert field (%s): type is not map",
				fromField.Name,
			)
		}

		err = c.genConverterForStruct(
			toField.Name,
			valueStruct,
			fromFieldType.ValueSpec,
			"value",
			keyPrefix+strings.Title(toField.Name)+"[key]",
			indent,
		)
		if err != nil {
			return err
		}
	} else {
		line = toIdentifier + "[key] = " + typeName + "(value)"
		c.Lines = append(c.Lines, line)
	}

	line = "}"
	c.Lines = append(c.Lines, line)
	return nil
}

func (c *TypeConverter) genStructConverter(
	keyPrefix string,
	indent string,
	fromFields []*compile.FieldSpec,
	toFields []*compile.FieldSpec,
) error {
	for i := 0; i < len(toFields); i++ {
		toField := toFields[i]

		var fromField *compile.FieldSpec
		for j := 0; j < len(fromFields); j++ {
			if fromFields[j].Name == toField.Name {
				fromField = fromFields[j]
				break
			}
		}

		if fromField == nil {
			return errors.Errorf(
				"cannot map by name for the field %s",
				toField.Name,
			)
		}

		toIdentifier := indent + "out." + keyPrefix + strings.Title(toField.Name)
		fromIdentifier := "in." + keyPrefix + strings.Title(fromField.Name)

		// Override thrift type names to avoid naming collisions between endpoint
		// and client types.
		switch toFieldType := toField.Type.(type) {
		case
			*compile.BoolSpec,
			*compile.I8Spec,
			*compile.I16Spec,
			*compile.I32Spec,
			*compile.EnumSpec,
			*compile.I64Spec,
			*compile.DoubleSpec,
			*compile.StringSpec:

			err := c.genConverterForPrimitive(
				toField, toIdentifier, fromIdentifier,
			)
			if err != nil {
				return err
			}
		case *compile.BinarySpec:
			line := toIdentifier + " = []byte(" + fromIdentifier + ")"
			c.Lines = append(c.Lines, line)
		case *compile.TypedefSpec:
			typeName, err := c.getIdentifierName(toField.Type)
			if err != nil {
				return err
			}

			var line string
			// TODO: typedef for struct is invalid here ...
			if toField.Required {
				line = toIdentifier + " = " + typeName + "(" + fromIdentifier + ")"
			} else {
				line = toIdentifier + " = (*" + typeName + ")(" + fromIdentifier + ")"
			}
			c.Lines = append(c.Lines, line)

		case *compile.StructSpec:
			err := c.genConverterForStruct(
				toField.Name,
				toFieldType,
				fromField.Type,
				fromIdentifier,
				keyPrefix+strings.Title(toField.Name),
				indent,
			)
			if err != nil {
				return err
			}
		case *compile.ListSpec:
			err := c.genConverterForList(
				toFieldType,
				toField,
				fromField,
				toIdentifier,
				fromIdentifier,
				keyPrefix,
				indent,
			)
			if err != nil {
				return err
			}
		case *compile.MapSpec:
			err := c.genConverterForMap(
				toFieldType,
				toField,
				fromField,
				toIdentifier,
				fromIdentifier,
				keyPrefix,
				indent,
			)
			if err != nil {
				return err
			}
		default:
			// fmt.Printf("Unknown type %s for field %s \n",
			// 	toField.Type.TypeCode().String(), toField.Name,
			// )

			// pkgName, err := h.TypePackageName(toField.Type.ThriftFile())
			// if err != nil {
			// 	return nil, err
			// }
			// typeName := pkgName + "." + toField.Type.ThriftName()
			// line := prefix + "(*" + typeName + ")" + postfix
			// c.Lines = append(c.Lines, line)
		}
	}

	return nil
}

// GenStructConverter will add lines to the TypeConverter for mapping
// from one go struct to another based on two thriftrw.FieldGroups
func (c *TypeConverter) GenStructConverter(
	fromFields []*compile.FieldSpec,
	toFields []*compile.FieldSpec,
) error {
	err := c.genStructConverter("", "", fromFields, toFields)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>codegen/type_converter: cleanup comment<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package codegen

import (
	"fmt"
	"strings"

	"github.com/pkg/errors"
	"go.uber.org/thriftrw/compile"
)

// TypeConverter can generate a function body that converts two thriftrw
// FieldGroups from one to another. It's assumed that the converted code
// operates on two variables, "in" and "out" and that both are a go struct.
type TypeConverter struct {
	Lines  []string
	Helper PackageNameResolver
}

// PackageNameResolver interface allows for resolving what the
// package name for a thrift file is. This depends on where the
// thrift-based structs are generated.
type PackageNameResolver interface {
	TypePackageName(thriftFile string) (string, error)
}

func (c *TypeConverter) getGoTypeName(
	valueType compile.TypeSpec,
) (string, error) {
	switch s := valueType.(type) {
	case *compile.BoolSpec:
		return "bool", nil
	case *compile.I8Spec:
		return "int8", nil
	case *compile.I16Spec:
		return "int16", nil
	case *compile.I32Spec:
		return "int32", nil
	case *compile.I64Spec:
		return "int64", nil
	case *compile.DoubleSpec:
		return "float64", nil
	case *compile.StringSpec:
		return "string", nil
	case *compile.BinarySpec:
		return "[]byte", nil
	case *compile.MapSpec:
		panic("Not Implemented")
	case *compile.SetSpec:
		panic("Not Implemented")
	case *compile.ListSpec:
		panic("Not Implemented")
	case *compile.EnumSpec, *compile.StructSpec, *compile.TypedefSpec:
		return c.getIdentifierName(s)
	default:
		panic(fmt.Sprintf("Unknown type (%T) %v", valueType, valueType))
	}
}

func (c *TypeConverter) getIdentifierName(
	fieldType compile.TypeSpec,
) (string, error) {
	pkgName, err := c.Helper.TypePackageName(fieldType.ThriftFile())
	if err != nil {
		return "", errors.Wrapf(
			err,
			"could not lookup fieldType when building converter for %s :",
			fieldType.ThriftName(),
		)
	}
	typeName := pkgName + "." + fieldType.ThriftName()
	return typeName, nil
}

func (c *TypeConverter) genConverterForStruct(
	toFieldName string,
	toFieldType *compile.StructSpec,
	fromFieldType compile.TypeSpec,
	fromIdentifier string,
	keyPrefix string,
	indent string,
) error {
	toIdentifier := indent + "out." + keyPrefix

	typeName, err := c.getIdentifierName(toFieldType)
	if err != nil {
		return err
	}
	subToFields := toFieldType.Fields

	fromFieldStruct, ok := fromFieldType.(*compile.StructSpec)
	if !ok {
		return errors.Errorf(
			"could not convert struct fields, "+
				"incompatible type for %s :",
			toFieldName,
		)
	}

	line := "if " + fromIdentifier + " != nil {"
	c.Lines = append(c.Lines, line)

	line = "	" + toIdentifier + " = &" + typeName + "{}"
	c.Lines = append(c.Lines, line)

	subFromFields := fromFieldStruct.Fields
	err = c.genStructConverter(
		keyPrefix+".",
		indent+"	",
		subFromFields,
		subToFields,
	)
	if err != nil {
		return err
	}

	line = "} else {"
	c.Lines = append(c.Lines, line)

	line = "	" + toIdentifier + " = nil"
	c.Lines = append(c.Lines, line)

	line = "}"
	c.Lines = append(c.Lines, line)

	return nil
}

func (c *TypeConverter) genConverterForPrimitive(
	toField *compile.FieldSpec,
	toIdentifier string,
	fromIdentifier string,
) error {
	var line string
	typeName, err := c.getGoTypeName(toField.Type)
	if err != nil {
		return err
	}

	if toField.Required {
		line = toIdentifier + " = " + typeName + "(" + fromIdentifier + ")"
	} else {
		line = toIdentifier + " = (*" + typeName + ")(" + fromIdentifier + ")"
	}
	c.Lines = append(c.Lines, line)
	return nil
}

func (c *TypeConverter) genConverterForList(
	toFieldType *compile.ListSpec,
	toField *compile.FieldSpec,
	fromField *compile.FieldSpec,
	toIdentifier string,
	fromIdentifier string,
	keyPrefix string,
	indent string,
) error {
	typeName, err := c.getGoTypeName(toFieldType.ValueSpec)
	if err != nil {
		return err
	}

	valueStruct, isStruct := toFieldType.ValueSpec.(*compile.StructSpec)
	if isStruct {
		line := toIdentifier + " = make([]*" +
			typeName + ", len(" + fromIdentifier + "))"
		c.Lines = append(c.Lines, line)
	} else {
		line := toIdentifier + " = make([]" +
			typeName + ", len(" + fromIdentifier + "))"
		c.Lines = append(c.Lines, line)
	}

	line := "for index, value := range " + fromIdentifier + " {"
	c.Lines = append(c.Lines, line)

	if isStruct {
		fromFieldType, ok := fromField.Type.(*compile.ListSpec)
		if !ok {
			return errors.Errorf(
				"Could not convert field (%s): type is not list",
				fromField.Name,
			)
		}

		err = c.genConverterForStruct(
			toField.Name,
			valueStruct,
			fromFieldType.ValueSpec,
			"value",
			keyPrefix+strings.Title(toField.Name)+"[index]",
			indent,
		)
		if err != nil {
			return err
		}
	} else {
		line = toIdentifier + "[index] = " + typeName + "(value)"
		c.Lines = append(c.Lines, line)
	}

	line = "}"
	c.Lines = append(c.Lines, line)
	return nil
}

func (c *TypeConverter) genConverterForMap(
	toFieldType *compile.MapSpec,
	toField *compile.FieldSpec,
	fromField *compile.FieldSpec,
	toIdentifier string,
	fromIdentifier string,
	keyPrefix string,
	indent string,
) error {
	typeName, err := c.getGoTypeName(toFieldType.ValueSpec)
	if err != nil {
		return err
	}

	_, isStringKey := toFieldType.KeySpec.(*compile.StringSpec)
	if !isStringKey {
		return errors.Errorf(
			"could not convert key (%s), map is not string-keyed.",
			toField.Name,
		)
	}

	valueStruct, isStruct := toFieldType.ValueSpec.(*compile.StructSpec)
	if isStruct {
		line := toIdentifier + " = make(map[string]*" +
			typeName + ", len(" + fromIdentifier + "))"
		c.Lines = append(c.Lines, line)
	} else {
		line := toIdentifier + " = make(map[string]" +
			typeName + ", len(" + fromIdentifier + "))"
		c.Lines = append(c.Lines, line)
	}

	line := "for key, value := range " + fromIdentifier + " {"
	c.Lines = append(c.Lines, line)

	if isStruct {
		fromFieldType, ok := fromField.Type.(*compile.MapSpec)
		if !ok {
			return errors.Errorf(
				"Could not convert field (%s): type is not map",
				fromField.Name,
			)
		}

		err = c.genConverterForStruct(
			toField.Name,
			valueStruct,
			fromFieldType.ValueSpec,
			"value",
			keyPrefix+strings.Title(toField.Name)+"[key]",
			indent,
		)
		if err != nil {
			return err
		}
	} else {
		line = toIdentifier + "[key] = " + typeName + "(value)"
		c.Lines = append(c.Lines, line)
	}

	line = "}"
	c.Lines = append(c.Lines, line)
	return nil
}

func (c *TypeConverter) genStructConverter(
	keyPrefix string,
	indent string,
	fromFields []*compile.FieldSpec,
	toFields []*compile.FieldSpec,
) error {
	for i := 0; i < len(toFields); i++ {
		toField := toFields[i]

		var fromField *compile.FieldSpec
		for j := 0; j < len(fromFields); j++ {
			if fromFields[j].Name == toField.Name {
				fromField = fromFields[j]
				break
			}
		}

		if fromField == nil {
			return errors.Errorf(
				"cannot map by name for the field %s",
				toField.Name,
			)
		}

		toIdentifier := indent + "out." + keyPrefix + strings.Title(toField.Name)
		fromIdentifier := "in." + keyPrefix + strings.Title(fromField.Name)

		// Override thrift type names to avoid naming collisions between endpoint
		// and client types.
		switch toFieldType := toField.Type.(type) {
		case
			*compile.BoolSpec,
			*compile.I8Spec,
			*compile.I16Spec,
			*compile.I32Spec,
			*compile.EnumSpec,
			*compile.I64Spec,
			*compile.DoubleSpec,
			*compile.StringSpec:

			err := c.genConverterForPrimitive(
				toField, toIdentifier, fromIdentifier,
			)
			if err != nil {
				return err
			}
		case *compile.BinarySpec:
			line := toIdentifier + " = []byte(" + fromIdentifier + ")"
			c.Lines = append(c.Lines, line)
		case *compile.TypedefSpec:
			typeName, err := c.getIdentifierName(toField.Type)
			if err != nil {
				return err
			}

			var line string
			// TODO: typedef for struct is invalid here ...
			if toField.Required {
				line = toIdentifier + " = " + typeName + "(" + fromIdentifier + ")"
			} else {
				line = toIdentifier + " = (*" + typeName + ")(" + fromIdentifier + ")"
			}
			c.Lines = append(c.Lines, line)

		case *compile.StructSpec:
			err := c.genConverterForStruct(
				toField.Name,
				toFieldType,
				fromField.Type,
				fromIdentifier,
				keyPrefix+strings.Title(toField.Name),
				indent,
			)
			if err != nil {
				return err
			}
		case *compile.ListSpec:
			err := c.genConverterForList(
				toFieldType,
				toField,
				fromField,
				toIdentifier,
				fromIdentifier,
				keyPrefix,
				indent,
			)
			if err != nil {
				return err
			}
		case *compile.MapSpec:
			err := c.genConverterForMap(
				toFieldType,
				toField,
				fromField,
				toIdentifier,
				fromIdentifier,
				keyPrefix,
				indent,
			)
			if err != nil {
				return err
			}
		default:
			// fmt.Printf("Unknown type %s for field %s \n",
			// 	toField.Type.TypeCode().String(), toField.Name,
			// )

			// pkgName, err := h.TypePackageName(toField.Type.ThriftFile())
			// if err != nil {
			// 	return nil, err
			// }
			// typeName := pkgName + "." + toField.Type.ThriftName()
			// line := toIdentifier + "(*" + typeName + ")" + postfix
			// c.Lines = append(c.Lines, line)
		}
	}

	return nil
}

// GenStructConverter will add lines to the TypeConverter for mapping
// from one go struct to another based on two thriftrw.FieldGroups
func (c *TypeConverter) GenStructConverter(
	fromFields []*compile.FieldSpec,
	toFields []*compile.FieldSpec,
) error {
	err := c.genStructConverter("", "", fromFields, toFields)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package doggy

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"net/http"
	"os"
	"strings"
	"sync/atomic"
)

// Key to use when setting the request ID.
type ctxKeyRequestID int

// RequestIDKey is the key that holds th unique request ID in a request context.
const RequestIDKey ctxKeyRequestID = 0

var prefix string
var reqid uint64

func init() {
	hostname, err := os.Hostname()
	if hostname == "" || err != nil {
		hostname = "localhost"
	}
	var buf [12]byte
	var b64 string
	for len(b64) < 10 {
		rand.Read(buf[:])
		b64 = base64.StdEncoding.EncodeToString(buf[:])
		b64 = strings.NewReplacer("+", "", "/", "").Replace(b64)
	}

	prefix = fmt.Sprintf("%s/%s", hostname, b64[0:10])
}

func TraceID(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	myid := atomic.AddUint64(&reqid, 1)
	traceID := fmt.Sprintf("%s-%06d", prefix, myid)
	ctx := context.WithValue(r.Context(), RequestIDKey, traceID)
	next(rw, r.WithContext(ctx))
	rw.Header().Set("TraceID", traceID)
}

func GetReqID(ctx context.Context) string {
	if ctx == nil {
		return ""
	}
	if reqID, ok := ctx.Value(RequestIDKey).(string); ok {
		return reqID
	}
	return ""
}
<commit_msg>unexport RequestIDKey<commit_after>package doggy

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"net/http"
	"os"
	"strings"
	"sync/atomic"
)

// Key to use when setting the request ID.
type ctxKeyRequestID int

// RequestIDKey is the key that holds th unique request ID in a request context.
const requestIDKey ctxKeyRequestID = 0

var prefix string
var reqid uint64

func init() {
	hostname, err := os.Hostname()
	if hostname == "" || err != nil {
		hostname = "localhost"
	}
	var buf [12]byte
	var b64 string
	for len(b64) < 10 {
		rand.Read(buf[:])
		b64 = base64.StdEncoding.EncodeToString(buf[:])
		b64 = strings.NewReplacer("+", "", "/", "").Replace(b64)
	}

	prefix = fmt.Sprintf("%s/%s", hostname, b64[0:10])
}

func TraceID(rw http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	myid := atomic.AddUint64(&reqid, 1)
	traceID := fmt.Sprintf("%s-%06d", prefix, myid)
	ctx := context.WithValue(r.Context(), requestIDKey, traceID)
	next(rw, r.WithContext(ctx))
	rw.Header().Set("TraceID", traceID)
}

func GetReqID(ctx context.Context) string {
	if ctx == nil {
		return ""
	}
	if reqID, ok := ctx.Value(requestIDKey).(string); ok {
		return reqID
	}
	return ""
}
<|endoftext|>