<commit_before>// Copyright 2016 The fer 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 fer provides a basic framework to run FairMQ-like processes.
// Clients create fer Devices that exchange data via fer message queue sockets.
//
// A client device might look like so:
//
//   import "github.com/sbinet-alice/fer"
//   import "github.com/sbinet-alice/fer/config"
//   type myDevice struct {
//       cfg  config.Device
//       imsg chan fer.Msg
//       omsg chan fer.Msg
//   }
//
// A device needs to implement the fer.Device interface:
//  func (dev *myDevice) Configure(cfg config.Device) error { ... }
//  func (dev *myDevice) Init(ctl fer.Controler)  error { ... }
//  func (dev *myDevice) Run(ctl fer.Controler)   error { ... }
//  func (dev *myDevice) Pause(ctl fer.Controler) error { ... }
//  func (dev *myDevice) Reset(ctl fer.Controler) error { ... }
//
// Typically, the Configure method is used to retrieve the configuration
// associated with the client's device.
// The Init method is used to retrieve the channels of input/output data messages.
// The Run method is an infinite for-loop, selecting on these input/output data
// messages.
// This infinite for-loop will also NEED to listen for the Controler.Done()
// channel to exit that for-loop.
//
// e.g.:
//
//  func (dev *myDevice) Init(ctl fer.Controler) error {
//      imsg, err := ctl.Chan("data-1", 0)
//      omsg, err := ctl.Chan("data-2", 0)
//      dev.imsg = imsg
//      dev.omsg = omsg
//      return nil
//  }
//
//  func (dev *myDevice) Run(ctl fer.Controler) error {
//      for {
//          select {
//          case data := <-dev.imsg:
//              dev.omsg <- bytes.Repeat(data, 2)
//          case <-ctl.Done():
//              return nil
//          }
//      }
//  }
//
// Then, to create an executable:
//
//  package main
//
//  func main() {
//      err := fer.Main(&myDevice{})
//      if err != nil {
//          log.Fatal(err)
//      }
//  }
//
// And voila.
package fer

import (
	"context"
	"fmt"
	"io"
	"os"

	"github.com/sbinet-alice/fer/config"
)

// Main configures and runs a device's execution, managing its state.
func Main(dev Device) error {
	cfg, err := config.Parse()
	if err != nil {
		return err
	}

	if cfg.Control == "" {
		cfg.Control = "static"
	}
	return runDevice(context.Background(), cfg, dev, os.Stdin, os.Stdout)
}

// Device is a handle to what users get to run via the Fer toolkit.
//
// Devices are configured according to command-line flags and a JSON
// configuration file.
// Clients should implement the Run method to receive and send data via
// the Controler data channels.
type Device interface {
	// Configure hands a device its configuration.
	Configure(cfg config.Device) error
	// Init gives a chance to the device to initialize internal
	// data structures, retrieve channels to input/output data.
	Init(ctl Controler) error
	// Run is where the device's main activity happens.
	// Run should loop forever, until the Controler.Done() channel says
	// otherwise.
	Run(ctl Controler) error
	// Pause pauses the device's execution.
	Pause(ctl Controler) error
	// Reset resets the device's internal state.
	Reset(ctl Controler) error
}

// Controler controls devices execution and gives a device access to input and
// output data channels.
type Controler interface {
	Logger
	Chan(name string, i int) (chan Msg, error)
	Done() chan Cmd

	isControler()
}

// Logger gives access to printf-like facilities
type Logger interface {
	Fatalf(format string, v ...interface{})
	Printf(format string, v ...interface{})
}

// Msg is a quantum of data being exchanged between devices.
type Msg struct {
	Data []byte // Data is the message payload.
	Err  error  // Err indicates whether an error occured.
}

// Cmd describes commands to be sent to a device, via a channel.
type Cmd byte

const (
	// CmdInitDevice is the command sent to initialize a device
	CmdInitDevice Cmd = iota
	// CmdInitTask is the command sent to initialize the tasks of a device
	CmdInitTask
	// CmdRun is the command sent to run a device
	CmdRun
	// CmdPause is the command sent to pause the execution of a device
	CmdPause
	// CmdStop is the command sent to stop the execution of a device
	CmdStop
	// CmdResetTask is the command sent to reset the state of the tasks of a device
	CmdResetTask
	// CmdResetDevice is the command sent to reset the state of a device
	CmdResetDevice
	// CmdEnd is the command sent to end a device
	CmdEnd
	// CmdError is the command sent to notify of an error
	CmdError
)

func (cmd Cmd) String() string {
	switch cmd {
	case CmdInitDevice:
		return "INIT_DEVICE"
	case CmdInitTask:
		return "INIT_TASK"
	case CmdRun:
		return "RUN"
	case CmdPause:
		return "PAUSE"
	case CmdStop:
		return "STOP"
	case CmdResetTask:
		return "RESET_TASK"
	case CmdResetDevice:
		return "RESET_DEVICE"
	case CmdEnd:
		return "END"
	case CmdError:
		return "ERROR_FOUND"
	}
	panic(fmt.Errorf("fer: invalid Cmd value (command=%d)", int(cmd)))
}

func runDevice(ctx context.Context, cfg config.Config, dev Device, r io.Reader, w io.Writer) error {
	sys, err := newDevice(ctx, cfg, dev, r, w)
	if err != nil {
		return err
	}

	sys.cmds <- CmdInitDevice
	if cfg.Control == "static" {
		sys.cmds <- CmdRun
	}

	return sys.run(ctx)
}

func broadcast(cmd Cmd, devs ...*device) {
	for _, dev := range devs {
		dev.cmds <- cmd
	}
}
<commit_msg>fer: more package doc<commit_after>// Copyright 2016 The fer 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 fer provides a basic framework to run FairMQ-like processes.
// Clients create fer Devices that exchange data via fer message queue sockets.
//
// A client device might look like so:
//
//   import "github.com/sbinet-alice/fer"
//   import "github.com/sbinet-alice/fer/config"
//   type myDevice struct {
//       cfg  config.Device
//       imsg chan fer.Msg
//       omsg chan fer.Msg
//   }
//
// A device needs to implement the fer.Device interface:
//  func (dev *myDevice) Configure(cfg config.Device) error { ... }
//  func (dev *myDevice) Init(ctl fer.Controler)  error { ... }
//  func (dev *myDevice) Run(ctl fer.Controler)   error { ... }
//  func (dev *myDevice) Pause(ctl fer.Controler) error { ... }
//  func (dev *myDevice) Reset(ctl fer.Controler) error { ... }
//
// Typically, the Configure method is used to retrieve the configuration
// associated with the client's device.
// The Init method is used to retrieve the channels of input/output data messages.
// The Run method is an infinite for-loop, selecting on these input/output data
// messages.
// This infinite for-loop will also NEED to listen for the Controler.Done()
// channel to exit that for-loop.
//
// e.g.:
//
//  func (dev *myDevice) Init(ctl fer.Controler) error {
//      imsg, err := ctl.Chan("data-1", 0)
//      omsg, err := ctl.Chan("data-2", 0)
//      dev.imsg = imsg
//      dev.omsg = omsg
//      return nil
//  }
//
//  func (dev *myDevice) Run(ctl fer.Controler) error {
//      for {
//          select {
//          case data := <-dev.imsg:
//              dev.omsg <- bytes.Repeat(data, 2)
//          case <-ctl.Done():
//              return nil
//          }
//      }
//  }
//
// Then, to create an executable:
//
//  package main
//
//  func main() {
//      err := fer.Main(&myDevice{})
//      if err != nil {
//          log.Fatal(err)
//      }
//  }
//
// Build it as usual and run like so:
//
//  $> go build -o my-device
//  $> ./my-device --help
//  Usage of my-device:
//    -control string
//      	starts device in interactive/static mode (default "interactive")
//    -id string
//      	device ID
//    -mq-config string
//      	path to JSON file holding device configuration
//    -transport string
//      	transport mechanism to use (zeromq, nanomsg, go-chan, ...) (default "zeromq")
//  $> ./my-device --id my-id --mq-config ./path/to/config.json
package fer

import (
	"context"
	"fmt"
	"io"
	"os"

	"github.com/sbinet-alice/fer/config"
)

// Main configures and runs a device's execution, managing its state.
func Main(dev Device) error {
	cfg, err := config.Parse()
	if err != nil {
		return err
	}

	if cfg.Control == "" {
		cfg.Control = "static"
	}
	return runDevice(context.Background(), cfg, dev, os.Stdin, os.Stdout)
}

// Device is a handle to what users get to run via the Fer toolkit.
//
// Devices are configured according to command-line flags and a JSON
// configuration file.
// Clients should implement the Run method to receive and send data via
// the Controler data channels.
type Device interface {
	// Configure hands a device its configuration.
	Configure(cfg config.Device) error
	// Init gives a chance to the device to initialize internal
	// data structures, retrieve channels to input/output data.
	Init(ctl Controler) error
	// Run is where the device's main activity happens.
	// Run should loop forever, until the Controler.Done() channel says
	// otherwise.
	Run(ctl Controler) error
	// Pause pauses the device's execution.
	Pause(ctl Controler) error
	// Reset resets the device's internal state.
	Reset(ctl Controler) error
}

// Controler controls devices execution and gives a device access to input and
// output data channels.
type Controler interface {
	Logger
	Chan(name string, i int) (chan Msg, error)
	Done() chan Cmd

	isControler()
}

// Logger gives access to printf-like facilities
type Logger interface {
	Fatalf(format string, v ...interface{})
	Printf(format string, v ...interface{})
}

// Msg is a quantum of data being exchanged between devices.
type Msg struct {
	Data []byte // Data is the message payload.
	Err  error  // Err indicates whether an error occured.
}

// Cmd describes commands to be sent to a device, via a channel.
type Cmd byte

const (
	// CmdInitDevice is the command sent to initialize a device
	CmdInitDevice Cmd = iota
	// CmdInitTask is the command sent to initialize the tasks of a device
	CmdInitTask
	// CmdRun is the command sent to run a device
	CmdRun
	// CmdPause is the command sent to pause the execution of a device
	CmdPause
	// CmdStop is the command sent to stop the execution of a device
	CmdStop
	// CmdResetTask is the command sent to reset the state of the tasks of a device
	CmdResetTask
	// CmdResetDevice is the command sent to reset the state of a device
	CmdResetDevice
	// CmdEnd is the command sent to end a device
	CmdEnd
	// CmdError is the command sent to notify of an error
	CmdError
)

func (cmd Cmd) String() string {
	switch cmd {
	case CmdInitDevice:
		return "INIT_DEVICE"
	case CmdInitTask:
		return "INIT_TASK"
	case CmdRun:
		return "RUN"
	case CmdPause:
		return "PAUSE"
	case CmdStop:
		return "STOP"
	case CmdResetTask:
		return "RESET_TASK"
	case CmdResetDevice:
		return "RESET_DEVICE"
	case CmdEnd:
		return "END"
	case CmdError:
		return "ERROR_FOUND"
	}
	panic(fmt.Errorf("fer: invalid Cmd value (command=%d)", int(cmd)))
}

func runDevice(ctx context.Context, cfg config.Config, dev Device, r io.Reader, w io.Writer) error {
	sys, err := newDevice(ctx, cfg, dev, r, w)
	if err != nil {
		return err
	}

	sys.cmds <- CmdInitDevice
	if cfg.Control == "static" {
		sys.cmds <- CmdRun
	}

	return sys.run(ctx)
}

func broadcast(cmd Cmd, devs ...*device) {
	for _, dev := range devs {
		dev.cmds <- cmd
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2015 Miquel Sabaté Solà <mikisabate@gmail.com>
// This file is licensed under the MIT license.
// See the LICENSE file.

package main

import (
	"bytes"
	"flag"
	"os/exec"
)

type gofmt struct {
	set bool
}

func (f *gofmt) setOptions() {
	flag.BoolVar(&f.set, "fmt", false, "Use the fmt tool.")
}

func (f *gofmt) installed() bool {
	return true
}

func (f *gofmt) isSet() bool {
	return f.set
}

func (f *gofmt) run(pack string) bool {
	printBackendStatus("fmt")

	// Unfortunately, gofmt needs either the relative or the absolute path, the
	// import path is not enough.
	if pack != "." {
		pack = packageAbs(pack)
	}

	cmd := exec.Command("gofmt", "-d", "-e", pack)
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()

	if err != nil {
		printResult("fmt", err.Error(), errored)
		return false
	}

	// And print out the results.
	results := out.String()
	if results == "" {
		printResult("fmt", "", ok)
		return true
	}
	printResult("fmt", out.String(), failed)
	return false
}
<commit_msg>fmt: use "go fmt" instead of "gofmt"<commit_after>// Copyright (C) 2015 Miquel Sabaté Solà <mikisabate@gmail.com>
// This file is licensed under the MIT license.
// See the LICENSE file.

package main

import (
	"bytes"
	"flag"
	"os/exec"
)

type gofmt struct {
	set bool
}

func (f *gofmt) setOptions() {
	flag.BoolVar(&f.set, "fmt", false, "Use the fmt tool.")
}

func (f *gofmt) installed() bool {
	return true
}

func (f *gofmt) isSet() bool {
	return f.set
}

func (f *gofmt) run(pack string) bool {
	printBackendStatus("fmt")

	// Unfortunately, gofmt needs either the relative or the absolute path, the
	// import path is not enough.
	if pack != "." {
		pack = packageAbs(pack)
	}

	cmd := exec.Command("go", "fmt", pack)
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()

	if err != nil {
		printResult("fmt", err.Error(), errored)
		return false
	}

	// And print out the results.
	results := out.String()
	if results == "" {
		printResult("fmt", "", ok)
		return true
	}
	printResult("fmt", out.String(), failed)
	return false
}
<|endoftext|>


<commit_before>package main

import (
  "strings"
  "log"
  // "fmt"
//  "regexp"
  "bytes"
  "encoding/json"
  "os"
  "time"
  "strconv"
  "regexp"
  "github.com/lookify/town/cluster"
  dockerapi "github.com/fsouza/go-dockerclient"
)

var (
  SCALE_NUM_REG, _ = regexp.Compile("\\$\\{SCALE_NUM:(.+)\\}")
//  SCALE_TOTAL_REG, _ = regexp.Compile("\\$\\{SCALE_NUM:(.+)\\}")
  HOSTS_REG, _ = regexp.Compile("\\$\\{(.+)_HOSTS\\}")
)

type Town struct {
  cluster *cluster.Cluster
  docker *dockerapi.Client
}

func NewTown() *Town {
  return &Town{
    cluster: nil,
    docker: nil,
  }
}

func (t *Town) ReadFile(name string) {
  var pathLocs = [...]string{
    name + ".yml",
    "/etc/town/" + name + ".yml",
  }

  for _, path := range pathLocs {
    if _, err := os.Stat(path); err == nil {
      t.cluster = cluster.NewCluster(path)
      t.cluster.ReadFile()
      return
    }
  }

  log.Println("ERROR: Could not find file ", name, ".yml")
}

func (t *Town) Connect() {
  endpoint := t.cluster.Application.Docker.Hosts[0] // at the moment use only first
  log.Println("Using Docker API endpont: ", endpoint)
  docker, err := dockerapi.NewClient( endpoint )
  if err != nil {
    log.Println("Can't connect to the docker")
  }
  t.docker = docker
}

func (t *Town) Provision(checkChanged bool) {
  // update containers
  for _, node := range t.cluster.Nodes {
    var buf bytes.Buffer
    var image = strings.Split(node.Container.Image, ":")

    opts := dockerapi.PullImageOptions{
        Repository: image[0],
        // Registry:     "docker.tsuru.io",
        // Tag:          "latest",
        OutputStream: &buf,
    }

    if len(image) > 1 {
      opts.Tag = image[1]
    }
    err := t.docker.PullImage(opts, dockerapi.AuthConfiguration{});
    if err != nil {
      log.Println("Could not pull image ", image)
    }
  }

  allContainers, err := t.docker.ListContainers(dockerapi.ListContainersOptions{
    All: true,
  })
  if err == nil {
    for _, listing := range allContainers {
      container, err := t.docker.InspectContainer(listing.ID)
      if err == nil {
        name := container.Name[1:]
        node, index := t.cluster.FindNodeByName(name)
        if node != nil && index > 0 {
          if node.Container.Exist == nil {
            node.Container.Exist = []cluster.ExistContainer{}
          }
          runningContainer := cluster.NewExistContainer(listing.ID, name, index, container.State.Running)
          if checkChanged {
            node.Container.Changed = t.isChangedImage(node, container)
          } else {
            node.Container.Changed = true
          }
          node.Container.Exist = append(node.Container.Exist, runningContainer)
        }
      } else {
        log.Println("[ERROR] Unable to inspect container:", listing.ID[:12], err)
      }
    }

    if checkChanged {
      t.cluster.AddChangeDependant()
    }
  } else {
    log.Println("[ERROR] Can't start provision")
  }
}

/**
 * Check node and running container for changes.
 * TODO: add cache to docker call.
 **/
func (t *Town) isChangedImage(node *cluster.Node, container *dockerapi.Container) bool {
  var imageName = container.Image
  image , error := t.docker.InspectImage(imageName)
  if error == nil {
    secondImage , secondError := t.docker.InspectImage(node.Container.Image)
    if secondError == nil {
      return secondImage.Created.After(image.Created)
    } else {
      log.Println("[ERROR] Could not inspect image ", node.Container.Name)
    }
  } else {
    log.Println("[ERROR] Could not inspect image ", imageName)
  }
  return false
}

func (t *Town) StopContainers(checkChanged bool) {
  log.Println("Stop...")
  //for node := range t.cluster.nodes {
  for i := len(t.cluster.Nodes) - 1; i >= 0; i-- {
    node := t.cluster.Nodes[i]
    if (!checkChanged || node.Container.Changed) && len(node.Container.Exist) > 0 {
      for _, container := range node.Container.Exist {
        if container.Running {
          err := t.docker.StopContainer(container.ID, 10)
          if err == nil {
            log.Println("   -  ", container.Name)
          } else {
            log.Println("   -  ", container.Name, " failed ", err)
          }
        }
      }
    }
  }
  log.Println("=============================")
}

func (t *Town) RemoveContainers(checkChanged bool) {
  log.Println("Remove...")
  //for node := range t.cluster.nodes {
  for i := len(t.cluster.Nodes) - 1; i >= 0; i-- {
    node := t.cluster.Nodes[i]
    if (!checkChanged || node.Container.Changed) && len(node.Container.Exist) > 0 {
      for _, container := range node.Container.Exist {
        err := t.docker.RemoveContainer(dockerapi.RemoveContainerOptions{
          ID: container.ID,
          RemoveVolumes: false,
        })
        if err == nil {
          log.Println("   -  ", container.Name)
        } else {
          log.Println("   -  ", container.Name, " failed ", err)
        }
      }
    }
  }
  log.Println("=============================")
}

func (t *Town) CreateContainer(node *cluster.Node, index int) (string, string, string) {
  containerName := node.Container.Name + "-" + strconv.Itoa(index)

  log.Println("   -  ", containerName)

  node.Container.Hostname = containerName // ?? Help !!!!

  env :=  make([]string, 0, cap(node.Container.Environment))
  for _, e := range node.Container.Environment {
    env = append(env, t.exec(e, index))
  }

  volumes := make(map[string]struct{})
  binds := make([]string, 0, cap(node.Container.Volumes))
  if len(node.Container.Volumes) > 0 {
    for _, volume := range node.Container.Volumes {
      volume = t.exec(volume, index)
      vol := strings.Split(volume, ":")
      if len(vol) > 1 {
        volumes[vol[1]] = struct{}{}
      } else {
        volumes[vol[0]] = struct{}{}
      }
      binds = append(binds, volume)
    }
  }

  dockerConfig := dockerapi.Config{
      Image: node.Container.Image,
      Hostname: node.Container.Hostname,
      PortSpecs: node.Container.Ports,
      Env: env,
      Volumes: volumes,

      AttachStdout: false,
      AttachStdin: false,
      AttachStderr: false,

      Tty: false,

      //Cmd: []
  }

  if len(node.Container.Command) > 0 {
    cmd := t.exec(node.Container.Command, index)
    dockerConfig.Cmd = []string{ cmd }
  }

  // just info
  //for _, l := range node.status.links {
  //  log.Println("     * ", l)
  //}

  // create links
  links := t.cluster.GetLinks(node)

  portBindings := map[dockerapi.Port][]dockerapi.PortBinding{}
  // create ports
  for _, ports := range node.Container.Ports {

    port := strings.Split(ports, ":")
    var p dockerapi.Port

    if len(port) > 1 {
      p = dockerapi.Port(port[1] + "/tcp")
    } else {
      p = dockerapi.Port(port[0] + "/tcp")
    }

    if portBindings[p] == nil {
      portBindings[p] = [] dockerapi.PortBinding {}
    }

    portBindings[p] = append(portBindings[p], dockerapi.PortBinding{
      HostIP: "",
      HostPort: port[0],
    })
  }

  hostConfig := dockerapi.HostConfig{
    Binds: binds,
    Links: links, //, [],
    PortBindings: portBindings,
    NetworkMode: "bridge",
    PublishAllPorts: false,
    Privileged: node.Container.Privileged,
  }

  opts := dockerapi.CreateContainerOptions{Name: containerName, Config: &dockerConfig, HostConfig: &hostConfig}
  container, err := t.docker.CreateContainer(opts)
  if err == nil {
    retry := 5
    for retry > 0 {
      error := t.docker.StartContainer(container.ID, &hostConfig)
      if error != nil {
        // log.Println("start error: ", error);

        out, err := json.Marshal(container)
        if err != nil {
            panic (err)
        }
        // fmt.Println(string(out))

        retry--;
        if retry == 0 {
          log.Println(" Start failed after 5 retries: ", string(out))
        }
        // log.Println("retry: ", retry);
      } else {
        inspect, inspectError := t.docker.InspectContainer(container.ID)
        if inspectError == nil {
          //links = append(links, inspect.NetworkSettings.IPAddress + "  " + containerName)
          //ids = append(ids, container.ID)
          return container.ID, inspect.NetworkSettings.IPAddress + "  " + containerName, containerName
        } else {
          log.Println("Inpect ", container.ID, " error ", inspectError)
        }
        //retry = 0
        break;
      }
    }
  } else {
    log.Println("error: ", err);
  }

  return "", "", ""
}

func (t *Town) CreateContainers(checkChanged bool) {
  log.Println("Create...")
  for _, node := range t.cluster.Nodes {

    if !checkChanged || node.Container.Changed {
      ids := make([]string, 0, node.Container.Scale )

      hosts := make([]string, 0, node.Container.Scale)

      log.Println(node.Container.Name, "  image: ", node.Container.Image)
      for i := 1; i <= node.Container.Scale; i++ {

        _, host, containerName := t.CreateContainer(node, i) //id

        if len(node.Container.Validate) > 0 {
          t.bashCommand(containerName, node.Container.Validate)
        }

        ids = append(ids, containerName)
        hosts = append(hosts, host)
      }

      if len(ids) > 1 {
        for index, id := range ids {
          var buffer bytes.Buffer

          buffer.WriteString("echo -e '")
          for i := 0; i < len(hosts); i++ {
            if i != index {
              buffer.WriteString("\n")
              buffer.WriteString(hosts[i])
            }
          }
          buffer.WriteString("' >> /etc/hosts; touch /tmp/host-generated")
          t.bashCommand(id, buffer.String() )
        }
      }

      time.Sleep(1000 * time.Millisecond)
    }
  }
}

func (t *Town) bashCommand(id string, command string)  {
  config := dockerapi.CreateExecOptions{
    Container:    id,
    AttachStdin:  true,
    AttachStdout: true,
    AttachStderr: false,
    Tty:          false,
    Cmd:          []string{"bash", "-c", command},
  }
  execObj, err := t.docker.CreateExec(config)
  if err == nil {
    config := dockerapi.StartExecOptions{
      Detach: true,
    }
    err := t.docker.StartExec(execObj.ID, config)
    if err != nil {
      log.Println("Container ", id, " command failed with error: ", err, "\n", command)
    }
  } else {
    log.Println("Container ", id, " command failed with error: ", err, "\n", command)
  }
}

func (t *Town) exec(text string, scale int) string {
  replace := strings.Replace(text, "${SCALE_NUM}", strconv.Itoa(scale), -1)
  match := SCALE_NUM_REG.FindAllStringSubmatch(replace, -1)
  hostMatch := HOSTS_REG.FindAllStringSubmatch(replace, -1)
  if len(match) > 0 {
    if len(match[0]) > 1 {
      nums := strings.Split(match[0][1], ",")
      if len(nums) > (scale - 1) {
        replace = strings.Replace(replace, match[0][0], nums[scale - 1], -1)
      }
    }
  }
  if len(hostMatch) > 0 {
    if len(hostMatch[0]) > 1 {
      //nums := strings.Split(, ",")
      name := strings.ToLower(hostMatch[0][1])
      node := t.cluster.FindNodeByID(name)

      var buffer bytes.Buffer
      for i := 1; i <= node.Container.Scale; i++ {
        buffer.WriteString(name)
        buffer.WriteString("-")
        buffer.WriteString(strconv.Itoa( i ))
        if i != node.Container.Scale {
          buffer.WriteString(",")
        }
      }
      replace = strings.Replace(replace, hostMatch[0][0], buffer.String(), -1)
    }
  }
  return replace
}
<commit_msg>fix exec command<commit_after>package main

import (
  "strings"
  "log"
  // "fmt"
//  "regexp"
  "bytes"
  "encoding/json"
  "os"
  "time"
  "strconv"
  "regexp"
  "github.com/lookify/town/cluster"
  dockerapi "github.com/fsouza/go-dockerclient"
)

var (
  SCALE_NUM_REG, _ = regexp.Compile("\\$\\{SCALE_NUM:(.+)\\}")
//  SCALE_TOTAL_REG, _ = regexp.Compile("\\$\\{SCALE_NUM:(.+)\\}")
  HOSTS_REG, _ = regexp.Compile("\\$\\{(.+)_HOSTS\\}")
)

type Town struct {
  cluster *cluster.Cluster
  docker *dockerapi.Client
}

func NewTown() *Town {
  return &Town{
    cluster: nil,
    docker: nil,
  }
}

func (t *Town) ReadFile(name string) {
  var pathLocs = [...]string{
    name + ".yml",
    "/etc/town/" + name + ".yml",
  }

  for _, path := range pathLocs {
    if _, err := os.Stat(path); err == nil {
      t.cluster = cluster.NewCluster(path)
      t.cluster.ReadFile()
      return
    }
  }

  log.Println("ERROR: Could not find file ", name, ".yml")
}

func (t *Town) Connect() {
  endpoint := t.cluster.Application.Docker.Hosts[0] // at the moment use only first
  log.Println("Using Docker API endpont: ", endpoint)
  docker, err := dockerapi.NewClient( endpoint )
  if err != nil {
    log.Println("Can't connect to the docker")
  }
  t.docker = docker
}

func (t *Town) Provision(checkChanged bool) {
  // update containers
  for _, node := range t.cluster.Nodes {
    var buf bytes.Buffer
    var image = strings.Split(node.Container.Image, ":")

    opts := dockerapi.PullImageOptions{
        Repository: image[0],
        // Registry:     "docker.tsuru.io",
        // Tag:          "latest",
        OutputStream: &buf,
    }

    if len(image) > 1 {
      opts.Tag = image[1]
    }
    err := t.docker.PullImage(opts, dockerapi.AuthConfiguration{});
    if err != nil {
      log.Println("Could not pull image ", image)
    }
  }

  allContainers, err := t.docker.ListContainers(dockerapi.ListContainersOptions{
    All: true,
  })
  if err == nil {
    for _, listing := range allContainers {
      container, err := t.docker.InspectContainer(listing.ID)
      if err == nil {
        name := container.Name[1:]
        node, index := t.cluster.FindNodeByName(name)
        if node != nil && index > 0 {
          if node.Container.Exist == nil {
            node.Container.Exist = []cluster.ExistContainer{}
          }
          runningContainer := cluster.NewExistContainer(listing.ID, name, index, container.State.Running)
          if checkChanged {
            node.Container.Changed = t.isChangedImage(node, container)
          } else {
            node.Container.Changed = true
          }
          node.Container.Exist = append(node.Container.Exist, runningContainer)
        }
      } else {
        log.Println("[ERROR] Unable to inspect container:", listing.ID[:12], err)
      }
    }

    if checkChanged {
      t.cluster.AddChangeDependant()
    }
  } else {
    log.Println("[ERROR] Can't start provision")
  }
}

/**
 * Check node and running container for changes.
 * TODO: add cache to docker call.
 **/
func (t *Town) isChangedImage(node *cluster.Node, container *dockerapi.Container) bool {
  var imageName = container.Image
  image , error := t.docker.InspectImage(imageName)
  if error == nil {
    secondImage , secondError := t.docker.InspectImage(node.Container.Image)
    if secondError == nil {
      return secondImage.Created.After(image.Created)
    } else {
      log.Println("[ERROR] Could not inspect image ", node.Container.Name)
    }
  } else {
    log.Println("[ERROR] Could not inspect image ", imageName)
  }
  return false
}

func (t *Town) StopContainers(checkChanged bool) {
  log.Println("Stop...")
  //for node := range t.cluster.nodes {
  for i := len(t.cluster.Nodes) - 1; i >= 0; i-- {
    node := t.cluster.Nodes[i]
    if (!checkChanged || node.Container.Changed) && len(node.Container.Exist) > 0 {
      for _, container := range node.Container.Exist {
        if container.Running {
          err := t.docker.StopContainer(container.ID, 10)
          if err == nil {
            log.Println("   -  ", container.Name)
          } else {
            log.Println("   -  ", container.Name, " failed ", err)
          }
        }
      }
    }
  }
  log.Println("=============================")
}

func (t *Town) RemoveContainers(checkChanged bool) {
  log.Println("Remove...")
  //for node := range t.cluster.nodes {
  for i := len(t.cluster.Nodes) - 1; i >= 0; i-- {
    node := t.cluster.Nodes[i]
    if (!checkChanged || node.Container.Changed) && len(node.Container.Exist) > 0 {
      for _, container := range node.Container.Exist {
        err := t.docker.RemoveContainer(dockerapi.RemoveContainerOptions{
          ID: container.ID,
          RemoveVolumes: false,
        })
        if err == nil {
          log.Println("   -  ", container.Name)
        } else {
          log.Println("   -  ", container.Name, " failed ", err)
        }
      }
    }
  }
  log.Println("=============================")
}

func (t *Town) CreateContainer(node *cluster.Node, index int) (string, string, string) {
  containerName := node.Container.Name + "-" + strconv.Itoa(index)

  log.Println("   -  ", containerName)

  node.Container.Hostname = containerName // ?? Help !!!!

  env :=  make([]string, 0, cap(node.Container.Environment))
  for _, e := range node.Container.Environment {
    env = append(env, t.exec(e, index))
  }

  volumes := make(map[string]struct{})
  binds := make([]string, 0, cap(node.Container.Volumes))
  if len(node.Container.Volumes) > 0 {
    for _, volume := range node.Container.Volumes {
      volume = t.exec(volume, index)
      vol := strings.Split(volume, ":")
      if len(vol) > 1 {
        volumes[vol[1]] = struct{}{}
      } else {
        volumes[vol[0]] = struct{}{}
      }
      binds = append(binds, volume)
    }
  }

  dockerConfig := dockerapi.Config{
      Image: node.Container.Image,
      Hostname: node.Container.Hostname,
      PortSpecs: node.Container.Ports,
      Env: env,
      Volumes: volumes,

      AttachStdout: false,
      AttachStdin: false,
      AttachStderr: false,

      Tty: false,

      //Cmd: []
  }

  if len(node.Container.Command) > 0 {
    cmd := t.exec(node.Container.Command, index)
    dockerConfig.Cmd = []string{ cmd }
  }

  // just info
  //for _, l := range node.status.links {
  //  log.Println("     * ", l)
  //}

  // create links
  links := t.cluster.GetLinks(node)

  portBindings := map[dockerapi.Port][]dockerapi.PortBinding{}
  // create ports
  for _, ports := range node.Container.Ports {

    port := strings.Split(ports, ":")
    var p dockerapi.Port

    if len(port) > 1 {
      p = dockerapi.Port(port[1] + "/tcp")
    } else {
      p = dockerapi.Port(port[0] + "/tcp")
    }

    if portBindings[p] == nil {
      portBindings[p] = [] dockerapi.PortBinding {}
    }

    portBindings[p] = append(portBindings[p], dockerapi.PortBinding{
      HostIP: "",
      HostPort: port[0],
    })
  }

  hostConfig := dockerapi.HostConfig{
    Binds: binds,
    Links: links, //, [],
    PortBindings: portBindings,
    NetworkMode: "bridge",
    PublishAllPorts: false,
    Privileged: node.Container.Privileged,
  }

  opts := dockerapi.CreateContainerOptions{Name: containerName, Config: &dockerConfig, HostConfig: &hostConfig}
  container, err := t.docker.CreateContainer(opts)
  if err == nil {
    retry := 5
    for retry > 0 {
      error := t.docker.StartContainer(container.ID, &hostConfig)
      if error != nil {
        // log.Println("start error: ", error);

        out, err := json.Marshal(container)
        if err != nil {
            panic (err)
        }
        // fmt.Println(string(out))

        retry--;
        if retry == 0 {
          log.Println(" Start failed after 5 retries: ", string(out))
        }
        // log.Println("retry: ", retry);
      } else {
        inspect, inspectError := t.docker.InspectContainer(container.ID)
        if inspectError == nil {
          //links = append(links, inspect.NetworkSettings.IPAddress + "  " + containerName)
          //ids = append(ids, container.ID)
          return container.ID, inspect.NetworkSettings.IPAddress + "  " + containerName, containerName
        } else {
          log.Println("Inpect ", container.ID, " error ", inspectError)
        }
        //retry = 0
        break;
      }
    }
  } else {
    log.Println("error: ", err);
  }

  return "", "", ""
}

func (t *Town) CreateContainers(checkChanged bool) {
  log.Println("Create...")
  for _, node := range t.cluster.Nodes {

    if !checkChanged || node.Container.Changed {
      ids := make([]string, 0, node.Container.Scale )

      hosts := make([]string, 0, node.Container.Scale)

      log.Println(node.Container.Name, "  image: ", node.Container.Image)
      for i := 1; i <= node.Container.Scale; i++ {

        _, host, containerName := t.CreateContainer(node, i) //id

        if len(node.Container.Validate) > 0 {
          t.bashCommand(containerName, node.Container.Validate)
        }

        ids = append(ids, containerName)
        hosts = append(hosts, host)
      }

      if len(ids) > 1 {
        for index, id := range ids {
          var buffer bytes.Buffer

          buffer.WriteString("echo -e '")
          for i := 0; i < len(hosts); i++ {
            if i != index {
              buffer.WriteString("\n")
              buffer.WriteString(hosts[i])
            }
          }
          buffer.WriteString("' >> /etc/hosts; touch /tmp/host-generated")
          t.bashCommand(id, buffer.String() )
        }
      }

      time.Sleep(1000 * time.Millisecond)
    }
  }
}

func (t *Town) bashCommand(id string, command string)  {
  config := dockerapi.CreateExecOptions{
    Container:    id,
    AttachStdin:  false,
    AttachStdout: false,
    AttachStderr: false,
    Tty:          false,
    Cmd:          []string{"bash", "-c", command},
  }
  execObj, err := t.docker.CreateExec(config)
  if err == nil {
    startConfig := dockerapi.StartExecOptions{
      Detach: false,
    }
    err := t.docker.StartExec(execObj.ID, startConfig)
    if err != nil {
      log.Println("Container ", id, " command failed with error: ", err, "\n", command)
    }
  } else {
    log.Println("Container ", id, " command failed with error: ", err, "\n", command)
  }
}

func (t *Town) exec(text string, scale int) string {
  replace := strings.Replace(text, "${SCALE_NUM}", strconv.Itoa(scale), -1)
  match := SCALE_NUM_REG.FindAllStringSubmatch(replace, -1)
  hostMatch := HOSTS_REG.FindAllStringSubmatch(replace, -1)
  if len(match) > 0 {
    if len(match[0]) > 1 {
      nums := strings.Split(match[0][1], ",")
      if len(nums) > (scale - 1) {
        replace = strings.Replace(replace, match[0][0], nums[scale - 1], -1)
      }
    }
  }
  if len(hostMatch) > 0 {
    if len(hostMatch[0]) > 1 {
      //nums := strings.Split(, ",")
      name := strings.ToLower(hostMatch[0][1])
      node := t.cluster.FindNodeByID(name)

      var buffer bytes.Buffer
      for i := 1; i <= node.Container.Scale; i++ {
        buffer.WriteString(name)
        buffer.WriteString("-")
        buffer.WriteString(strconv.Itoa( i ))
        if i != node.Container.Scale {
          buffer.WriteString(",")
        }
      }
      replace = strings.Replace(replace, hostMatch[0][0], buffer.String(), -1)
    }
  }
  return replace
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
)

type mpath struct {
	name string
	path string
}

type vdir string

func (vd vdir) trunc() string {
	s := string(vd)
	if len(s) > 0 {
		return s[:7]
	}
	return ""
}

var mpathre = regexp.MustCompile(`^\.(.*)###(.*)$`)
var vdirre = regexp.MustCompile(`^[a-f0-9]{64}$`)

func newmpath(mp string) *mpath {
	res := mpathre.FindAllStringSubmatch(mp, -1)
	var mres *mpath
	if res != nil && len(res) == 1 && len(res[0]) == 3 {
		name := res[0][1]
		path := res[0][2]
		path = strings.Replace(path, ",#,", "/", -1)
		mres = &mpath{name, path}
	}
	return mres
}

func (mp *mpath) String() string {
	return "." + mp.name + "###" + strings.Replace(mp.path, "/", ",#,", -1)
}

func newvdir(vd string) vdir {
	res := vdirre.FindString(vd)
	vres := vdir("")
	if res != "" && res == vd {
		vres = vdir(vd)
	}
	return vres
}

type marker struct {
	mp  *mpath
	dir vdir
}

func (mk *marker) String() string {
	return fmt.Sprintf("marker '%s'<%s->%s>", mk.dir.trunc(), mk.mp.name, mk.mp.path)
}

type markers []*marker

var allmarkers = markers{}

func (marks markers) getMarker(name string, path string, dir vdir) *marker {
	var mark *marker
	var pathdiff *marker
	for _, marker := range marks {
		if marker.mp.name == name {
			if marker.mp.path == path {
				mark = marker
				break
			} else {
				pathdiff = marker
			}
		}
	}
	if mark == nil {
		if pathdiff != nil {
			fmt.Printf("Invalid volume path detected: '%s' (not found in marker '%s')\n", path, pathdiff)
			return nil
		}
		mark = &marker{mp: &mpath{name: name, path: path}, dir: dir}
	}
	ldir := mark.dir
	if mark.dir != dir {
		// move dir and ln -s mark.dir dir
		ldir = dir
		cmd(fmt.Sprintf("sudo move /mnt/sda1/var/lib/docker/vfs/dir/%s /mnt/sda1/var/lib/docker/vfs/dir/_%s", dir, dir))
		fmt.Printf("Mark container named '%s' for path '%s' as link to '%s' (from '%s')\n", name, path, mark.dir, dir)
	}
	_, err := cmd(fmt.Sprintf("sudo ls -1L /mnt/sda1/var/lib/docker/vfs/dir/%s", mark.mp.String()))
	if err != nil {
		mustcmd(fmt.Sprintf("sudo ln -s %s /mnt/sda1/var/lib/docker/vfs/dir/%s", ldir, mark.mp.String()))
	}
	return mark
}

type volume struct {
	dir  vdir
	mark *marker
}

type volumes []*volume

var allvolumes = volumes{}

func (v *volume) String() string {
	return fmt.Sprintf("vol '%s'<%v>", v.dir.trunc(), v.mark)
}

func (vols volumes) getVolume(vd vdir, path string, name string) *volume {
	var vol *volume
	for _, volume := range vols {
		if string(volume.dir) == string(vd) {
			vol = volume
			if vol.mark == nil {
				vol.mark = allmarkers.getMarker(name, path, vol.dir)
			}
			if vol.mark == nil {
				return nil
			}
			break
		}
	}
	if vol == nil {
		vol = &volume{dir: vd}
		vol.mark = allmarkers.getMarker(name, path, vol.dir)
		if vol.mark == nil {
			return nil
		}
	}
	return vol
}

type container struct {
	name    string
	id      string
	stopped bool
	volumes []*volume
}

func (c *container) trunc() string {
	if len(c.id) > 0 {
		return c.id[:7]
	}
	return ""
}

func (c *container) String() string {
	return "cnt '" + c.name + "' (" + c.trunc() + ")" + fmt.Sprintf("[%v] - %d vol", c.stopped, len(c.volumes))
}

var containers = []*container{}

func mustcmd(acmd string) string {
	out, err := cmd(acmd)
	if err != nil {
		log.Fatal(fmt.Sprintf("out='%s', err='%s'", out, err))
	}
	return string(out)
}

type fcmd func(cmd string) (string, error)

var cmd = execcmd

func execcmd(cmd string) (string, error) {
	fmt.Println(cmd)
	out, err := exec.Command("sh", "-c", cmd).Output()
	return strings.TrimSpace(string(out)), err
}

func readVolumes() {
	out := mustcmd("sudo ls -a1F /mnt/sda1/var/lib/docker/vfs/dir")
	vollines := strings.Split(out, "\n")
	for _, volline := range vollines {
		dir := volline
		if dir == "./" || dir == "../" || dir == "" {
			continue
		}
		if strings.HasSuffix(dir, "@") {
			dir = dir[:len(dir)-1]
			fdir := fmt.Sprintf("/mnt/sda1/var/lib/docker/vfs/dir/%s", dir)
			mp := newmpath(dir)
			if mp == nil {
				fmt.Printf("Invalid marker detected: '%s'\n", dir)
				mustcmd("sudo rm " + fdir)
			} else {
				dirlink, err := cmd("sudo readlink " + fdir)
				fmt.Printf("---\ndir: '%s'\ndlk: '%s'\nerr='%v'\n", dir, dirlink, err)
				if err != nil {
					fmt.Printf("Invalid marker (no readlink) detected: '%s'\n", dir)
					mustcmd("sudo rm " + fdir)
				} else {
					_, err := cmd("sudo ls /mnt/sda1/var/lib/docker/vfs/dir/" + dirlink)
					if err != nil {
						fmt.Printf("Invalid marker (readlink no ls) detected: '%s'\n", dir)
						mustcmd("sudo rm " + fdir)
					} else {
						vd := newvdir(dirlink)
						if vd == "" {
							fmt.Printf("Invalid marker (readlink no vdir) detected: '%s'\n", dir)
							mustcmd("sudo rm " + fdir)
						} else {
							allmarkers = append(allmarkers, &marker{mp, vd})
						}
					}
				}
			}
		} else if strings.HasSuffix(dir, "/") {

			dir = dir[:len(dir)-1]
			fdir := fmt.Sprintf("/mnt/sda1/var/lib/docker/vfs/dir/%s", dir)
			vd := newvdir(dir)
			if vd == "" {
				fmt.Printf("Invalid volume folder detected: '%s'\n", dir)
				mustcmd("sudo rm " + fdir)
			} else {
				allvolumes = append(allvolumes, &volume{dir: vd})
			}
		} else {
			fdir := fmt.Sprintf("/mnt/sda1/var/lib/docker/vfs/dir/%s", dir)
			fmt.Printf("Invalid file detected: '%s'\n", dir)
			mustcmd("sudo rm " + fdir)
		}
	}
}

func readContainer() {

	out := mustcmd("docker ps -aq --no-trunc")
	contlines := strings.Split(out, "\n")
	// fmt.Println(contlines)
	for _, contline := range contlines {
		if contline == "" {
			continue
		}
		id := contline
		res := mustcmd("docker inspect -f '{{ .Name }},{{ range $key, $value := .Volumes }}{{ $key }},{{ $value }}##~#{{ end }}' " + id)
		// fmt.Println("res1: '" + res + "'")
		name := res[1:strings.Index(res, ",")]
		cont := &container{name: name, id: id}
		res = res[strings.Index(res, ",")+1:]
		// fmt.Println("res2: '" + res + "'")
		vols := strings.Split(res, "##~#")
		// fmt.Println(vols)
		for _, vol := range vols {
			elts := strings.Split(vol, ",")
			if len(elts) == 2 {
				// fmt.Printf("elts: '%v'\n", elts)
				path := elts[0]
				vfs := elts[1]
				if strings.Contains(vfs, "/var/lib/docker/vfs/dir/") {
					vd := newvdir(filepath.Base(vfs))
					if vd == "" {
						fmt.Printf("Invalid volume folder detected: '%s'\n", vfs)
						break
					}
					newvol := allvolumes.getVolume(vd, path, name)
					cont.volumes = append(cont.volumes, newvol)
				}
			}
		}
		containers = append(containers, cont)
		if len(cont.volumes) == 0 {
			fmt.Printf("Orphan container detected: '%s'\n", cont)
		}
	}
}

// docker run --rm -i -t -v `pwd`:`pwd` -w `pwd` --entrypoint="/bin/bash" go -c 'go build gcl.go'
func main() {
	readVolumes()
	readContainer()
}
<commit_msg>gcl.go: add public functions for Containers, Volumes and Markers<commit_after>package main

import (
	"fmt"
	"log"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
)

type mpath struct {
	name string
	path string
}

type vdir string

func (vd vdir) trunc() string {
	s := string(vd)
	if len(s) > 0 {
		return s[:7]
	}
	return ""
}

var mpathre = regexp.MustCompile(`^\.(.*)###(.*)$`)
var vdirre = regexp.MustCompile(`^[a-f0-9]{64}$`)

func newmpath(mp string) *mpath {
	res := mpathre.FindAllStringSubmatch(mp, -1)
	var mres *mpath
	if res != nil && len(res) == 1 && len(res[0]) == 3 {
		name := res[0][1]
		path := res[0][2]
		path = strings.Replace(path, ",#,", "/", -1)
		mres = &mpath{name, path}
	}
	return mres
}

func (mp *mpath) String() string {
	return "." + mp.name + "###" + strings.Replace(mp.path, "/", ",#,", -1)
}

func newvdir(vd string) vdir {
	res := vdirre.FindString(vd)
	vres := vdir("")
	if res != "" && res == vd {
		vres = vdir(vd)
	}
	return vres
}

type marker struct {
	mp  *mpath
	dir vdir
}

func (mk *marker) String() string {
	return fmt.Sprintf("marker '%s'<%s->%s>", mk.dir.trunc(), mk.mp.name, mk.mp.path)
}

type markers []*marker

var allmarkers = markers{}

func (marks markers) getMarker(name string, path string, dir vdir) *marker {
	var mark *marker
	var pathdiff *marker
	for _, marker := range marks {
		if marker.mp.name == name {
			if marker.mp.path == path {
				mark = marker
				break
			} else {
				pathdiff = marker
			}
		}
	}
	if mark == nil {
		if pathdiff != nil {
			fmt.Printf("Invalid volume path detected: '%s' (not found in marker '%s')\n", path, pathdiff)
			return nil
		}
		mark = &marker{mp: &mpath{name: name, path: path}, dir: dir}
	}
	ldir := mark.dir
	if mark.dir != dir {
		// move dir and ln -s mark.dir dir
		ldir = dir
		cmd(fmt.Sprintf("sudo move /mnt/sda1/var/lib/docker/vfs/dir/%s /mnt/sda1/var/lib/docker/vfs/dir/_%s", dir, dir))
		fmt.Printf("Mark container named '%s' for path '%s' as link to '%s' (from '%s')\n", name, path, mark.dir, dir)
	}
	_, err := cmd(fmt.Sprintf("sudo ls -1L /mnt/sda1/var/lib/docker/vfs/dir/%s", mark.mp.String()))
	if err != nil {
		mustcmd(fmt.Sprintf("sudo ln -s %s /mnt/sda1/var/lib/docker/vfs/dir/%s", ldir, mark.mp.String()))
	}
	return mark
}

type volume struct {
	dir  vdir
	mark *marker
}

type volumes []*volume

var allvolumes = volumes{}

func (v *volume) String() string {
	return fmt.Sprintf("vol '%s'<%v>", v.dir.trunc(), v.mark)
}

func (vols volumes) getVolume(vd vdir, path string, name string) *volume {
	var vol *volume
	for _, volume := range vols {
		if string(volume.dir) == string(vd) {
			vol = volume
			if vol.mark == nil {
				vol.mark = allmarkers.getMarker(name, path, vol.dir)
			}
			if vol.mark == nil {
				return nil
			}
			break
		}
	}
	if vol == nil {
		vol = &volume{dir: vd}
		vol.mark = allmarkers.getMarker(name, path, vol.dir)
		if vol.mark == nil {
			return nil
		}
	}
	return vol
}

type container struct {
	name    string
	id      string
	stopped bool
	volumes []*volume
}

func (c *container) trunc() string {
	if len(c.id) > 0 {
		return c.id[:7]
	}
	return ""
}

func (c *container) String() string {
	return "cnt '" + c.name + "' (" + c.trunc() + ")" + fmt.Sprintf("[%v] - %d vol", c.stopped, len(c.volumes))
}

var containers = []*container{}

func mustcmd(acmd string) string {
	out, err := cmd(acmd)
	if err != nil {
		log.Fatal(fmt.Sprintf("out='%s', err='%s'", out, err))
	}
	return string(out)
}

type fcmd func(cmd string) (string, error)

var cmd = execcmd

func execcmd(cmd string) (string, error) {
	fmt.Println(cmd)
	out, err := exec.Command("sh", "-c", cmd).Output()
	return strings.TrimSpace(string(out)), err
}

func readVolumes() {
	out := mustcmd("sudo ls -a1F /mnt/sda1/var/lib/docker/vfs/dir")
	vollines := strings.Split(out, "\n")
	for _, volline := range vollines {
		dir := volline
		if dir == "./" || dir == "../" || dir == "" {
			continue
		}
		if strings.HasSuffix(dir, "@") {
			dir = dir[:len(dir)-1]
			fdir := fmt.Sprintf("/mnt/sda1/var/lib/docker/vfs/dir/%s", dir)
			mp := newmpath(dir)
			if mp == nil {
				fmt.Printf("Invalid marker detected: '%s'\n", dir)
				mustcmd("sudo rm " + fdir)
			} else {
				dirlink, err := cmd("sudo readlink " + fdir)
				fmt.Printf("---\ndir: '%s'\ndlk: '%s'\nerr='%v'\n", dir, dirlink, err)
				if err != nil {
					fmt.Printf("Invalid marker (no readlink) detected: '%s'\n", dir)
					mustcmd("sudo rm " + fdir)
				} else {
					_, err := cmd("sudo ls /mnt/sda1/var/lib/docker/vfs/dir/" + dirlink)
					if err != nil {
						fmt.Printf("Invalid marker (readlink no ls) detected: '%s'\n", dir)
						mustcmd("sudo rm " + fdir)
					} else {
						vd := newvdir(dirlink)
						if vd == "" {
							fmt.Printf("Invalid marker (readlink no vdir) detected: '%s'\n", dir)
							mustcmd("sudo rm " + fdir)
						} else {
							allmarkers = append(allmarkers, &marker{mp, vd})
						}
					}
				}
			}
		} else if strings.HasSuffix(dir, "/") {

			dir = dir[:len(dir)-1]
			fdir := fmt.Sprintf("/mnt/sda1/var/lib/docker/vfs/dir/%s", dir)
			vd := newvdir(dir)
			if vd == "" {
				fmt.Printf("Invalid volume folder detected: '%s'\n", dir)
				mustcmd("sudo rm " + fdir)
			} else {
				allvolumes = append(allvolumes, &volume{dir: vd})
			}
		} else {
			fdir := fmt.Sprintf("/mnt/sda1/var/lib/docker/vfs/dir/%s", dir)
			fmt.Printf("Invalid file detected: '%s'\n", dir)
			mustcmd("sudo rm " + fdir)
		}
	}
}

func readContainer() {

	out := mustcmd("docker ps -aq --no-trunc")
	contlines := strings.Split(out, "\n")
	// fmt.Println(contlines)
	for _, contline := range contlines {
		if contline == "" {
			continue
		}
		id := contline
		res := mustcmd("docker inspect -f '{{ .Name }},{{ range $key, $value := .Volumes }}{{ $key }},{{ $value }}##~#{{ end }}' " + id)
		// fmt.Println("res1: '" + res + "'")
		name := res[1:strings.Index(res, ",")]
		cont := &container{name: name, id: id}
		res = res[strings.Index(res, ",")+1:]
		// fmt.Println("res2: '" + res + "'")
		vols := strings.Split(res, "##~#")
		// fmt.Println(vols)
		for _, vol := range vols {
			elts := strings.Split(vol, ",")
			if len(elts) == 2 {
				// fmt.Printf("elts: '%v'\n", elts)
				path := elts[0]
				vfs := elts[1]
				if strings.Contains(vfs, "/var/lib/docker/vfs/dir/") {
					vd := newvdir(filepath.Base(vfs))
					if vd == "" {
						fmt.Printf("Invalid volume folder detected: '%s'\n", vfs)
						break
					}
					newvol := allvolumes.getVolume(vd, path, name)
					cont.volumes = append(cont.volumes, newvol)
				}
			}
		}
		containers = append(containers, cont)
		if len(cont.volumes) == 0 {
			fmt.Printf("Orphan container detected: '%s'\n", cont)
		}
	}
}

// docker run --rm -i -t -v `pwd`:`pwd` -w `pwd` --entrypoint="/bin/bash" go -c 'go build gcl.go'
func main() {
	readVolumes()
	readContainer()
}

// Containers returns all detected containers
func Containers() []*container {
	return containers
}

var orphanedContainers = []*container{}

// OrphanedContainers returns all detected containers with missing volumes
// They are not orphaned if they have no volumes
// They are orphaned if they have volumes which don't exist
func OrphanedContainers() []*container {
	return orphanedContainers
}

// Volumes returns all detected volumes
func Volumes() volumes {
	return allvolumes
}

var orphanedVolumes = volumes{}

// OrphanedVolumes returns all detected volumes with missing containers
// They are orphaned because no container reference them directly
// or through a marker
func OrphanedVolumes() volumes {
	return orphanedVolumes
}

// Markers returns all detected or created markers
func Markers() markers {
	return allmarkers
}
<|endoftext|>


<commit_before>package hush

import (
	"crypto/hmac"
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sort"
	"strings"

	"golang.org/x/crypto/pbkdf2"

	"github.com/pkg/errors"

	yaml "gopkg.in/yaml.v2"
)

type Branch struct {
	path Path
	val  *Value
}

type Tree struct {
	branches []Branch
	index    map[Path]int

	encryptionKey []byte
	macKey        []byte
}

const safePerm = 0600 // rw- --- ---

func LoadTree() (*Tree, error) {
	hushPath, err := hushPath()
	if err != nil {
		return nil, err
	}

	stat, err := os.Stat(hushPath)
	if os.IsNotExist(err) {
		warn("hush file does not exist. assuming an empty one")
		return &Tree{}, nil
	}
	if err != nil {
		return nil, errors.Wrap(err, "can't stat hush file")
	}
	if (stat.Mode() & os.ModePerm) != safePerm {
		warn("hush file has loose permissions. fixing.")
		err := os.Chmod(hushPath, safePerm)
		if err != nil {
			die("couldn't fix permissions on hush file")
		}
	}

	file, err := os.Open(hushPath)
	if err != nil {
		return nil, errors.Wrap(err, "opening hush file")
	}
	hushData, err := ioutil.ReadAll(file)
	if err != nil {
		return nil, errors.Wrap(err, "can't read hush file")
	}

	keys := make(yaml.MapSlice, 0)
	err = yaml.Unmarshal(hushData, &keys)
	if err != nil {
		return nil, errors.Wrap(err, "can't parse hush file")
	}
	tree := newT(keys)
	return tree, nil
}

func newT(items yaml.MapSlice) *Tree {
	t := &Tree{
		branches: make([]Branch, 0, 3*len(items)),
		index:    make(map[Path]int),
	}
	newT_(items, []string{}, t)
	return t
}

func newT_(items yaml.MapSlice, crumbs []string, t *Tree) {
	n := len(crumbs)
	for _, item := range items {
		key := item.Key.(string)
		crumbs = append(crumbs, key)

		switch val := item.Value.(type) {
		case string:
			p := NewPath(strings.Join(crumbs, "/"))
			privacy := Private
			if p.IsPublic() {
				privacy = Public
			}
			t.set(p, NewEncoded(val, privacy))
		case yaml.MapSlice:
			newT_(val, crumbs, t)
		default:
			panic(fmt.Sprintf("unexpected type: %#v", val))
		}
		crumbs = crumbs[:n] // remove final crumb
	}
}

// implement sort.Interface interface
var _ sort.Interface = &Tree{}

func (t *Tree) Len() int { return len(t.branches) }
func (t *Tree) Less(i, j int) bool {
	return t.branches[i].path < t.branches[j].path
}
func (t *Tree) Swap(i, j int) {
	t.branches[i], t.branches[j] = t.branches[j], t.branches[i]
	t.index[t.branches[i].path] = i
	t.index[t.branches[j].path] = j
}

func (t *Tree) mapSlice() yaml.MapSlice {
	// sort by key
	sort.Stable(t)

	var slice yaml.MapSlice
	for _, branch := range t.branches {
		if branch.path.IsChecksum() {
			// skip checksum. it's appended by Save()
			continue
		}
		crumbs := branch.path.AsCrumbs()
		slice = mapSlice_(slice, crumbs, branch.val.String())
	}
	return slice
}

func mapSlice_(slice yaml.MapSlice, path []string, value string) yaml.MapSlice {
	if len(path) == 0 {
		panic("path should never have 0 length")
	}
	if len(path) == 1 {
		return append(slice, yaml.MapItem{
			Key:   path[0],
			Value: value,
		})
	}

	var inner yaml.MapSlice
	if len(slice) == 0 {
		slice = append(slice, yaml.MapItem{Key: path[0]})
	} else {
		final := slice[len(slice)-1]
		if final.Key.(string) == path[0] {
			inner = final.Value.(yaml.MapSlice)
		} else {
			slice = append(slice, yaml.MapItem{Key: path[0]})
		}
	}
	slice[len(slice)-1].Value = mapSlice_(inner, path[1:], value)
	return slice
}

// Checksum returns a cryptographic message authentication code for this tree.
func (t *Tree) Checksum() []byte {
	if len(t.macKey) < 32 {
		panic("trying to calculate checksum without a MAC key")
	}
	mac := hmac.New(sha256.New, t.macKey)
	for _, branch := range t.branches {
		if branch.path.IsChecksum() {
			continue // don't checksum the checksum
		}
		mac.Write([]byte(branch.path))
		mac.Write([]byte(branch.val.String()))
	}
	sum := mac.Sum(nil)
	return sum
}

func (t *Tree) filter(pattern string) *Tree {
	keep := t.Empty()
	for _, branch := range t.branches {
		if matches(branch.path, pattern) {
			keep.set(branch.path, branch.val)
		}
	}
	return keep
}

func isLowercase(s string) bool {
	return s == strings.ToLower(s)
}

func matches(p Path, pattern string) bool {
	ps := strings.Split(string(p), "\t")
	patterns := strings.Split(pattern, "/")
	if len(patterns) > len(ps) {
		return false
	}

	ignoreCase := isLowercase(pattern)
	for i, pattern := range patterns {
		haystack := ps[i]
		if ignoreCase {
			haystack = strings.ToLower(haystack)
		}
		if !strings.Contains(haystack, pattern) {
			return false
		}
	}
	return true
}

func (t *Tree) get(p Path) (*Value, bool) {
	i, ok := t.index[p]
	if ok {
		return t.branches[i].val, true
	}
	return nil, false
}

func (t *Tree) set(p Path, val *Value) {
	i, ok := t.index[p]
	if ok {
		t.branches[i].val = val
	} else {
		t.branches = append(t.branches, Branch{p, val})
		t.index[p] = len(t.branches) - 1
	}
}

// Encrypt returns a copy of this tree with all leaves encrypted.
func (tree *Tree) Encrypt() *Tree {
	t := tree.Empty()
	for _, branch := range tree.branches {
		p := branch.path
		v := branch.val
		if p.IsPublic() { // don't encrypt public data
			t.set(p, v)
			continue
		}
		if p.IsEncryptionKey() { // value uses different encryption key
			t.set(p, v)
			continue
		}
		t.set(p, v.Ciphertext(t.encryptionKey))
	}
	return t
}

// Encode returns a copy of this tree with all leaves encoded into base64.
func (tree *Tree) Encode() *Tree {
	t := tree.Empty()
	for _, branch := range tree.branches {
		t.set(branch.path, branch.val.Encode())
	}
	return t
}

// Empty returns a copy of this tree with all the keys and values
// removed.  It retains any other data associated with this tree.
func (t *Tree) Empty() *Tree {
	tree := *t // shallow copy
	tree.branches = make([]Branch, 0, len(t.branches))
	tree.index = make(map[Path]int, len(t.branches))
	return &tree
}

// SetPassphrase sets the password that's used for performing
// encryption and decryption.
func (t *Tree) SetPassphrase(password []byte) error {
	p := NewPath("hush-configuration/salt")
	v, ok := t.get(p)
	if !ok {
		return errors.New("hush file missing salt")
	}
	v, err := v.Decode()
	if err != nil {
		return errors.Wrap(err, "decoding salt")
	}
	salt := v.plaintext
	pwKey := stretchPassword(password, salt)

	p = NewPath("hush-configuration/encryption-key")
	v, ok = t.get(p)
	if !ok {
		return errors.New("hush file missing encryption key")
	}
	v, err = v.Plaintext(pwKey)
	if err != nil {
		return fmt.Errorf("incorrect password or corrupted encryption key")
	}
	t.encryptionKey = v.plaintext

	p = NewPath("hush-configuration/mac-key")
	v, ok = t.get(p)
	if !ok {
		return errors.New("hush file missing MAC key")
	}
	v, err = v.Plaintext(pwKey)
	if err != nil {
		return fmt.Errorf("incorrect password or corrupted mac key")
	}
	t.macKey = v.plaintext

	// now that we have a password, we can verify the checksum
	got, ok := t.get(NewPath("hush-tree-checksum"))
	if !ok {
		return errors.New("hush file has no checksum")
	}
	got, err = got.Decode()
	if err != nil {
		return errors.Wrap(err, "decoding checksum")
	}
	expect := t.Checksum()
	if !hmac.Equal(got.plaintext, expect) {
		return errors.New("checksum doesn't match. file modified without hush command?")
	}

	return nil
}

// Decrypt returns a copy of this tree with all leaves decrypted.
func (tree *Tree) Decrypt() *Tree {
	var err error
	t := tree.Empty()
	for _, branch := range tree.branches {
		p := branch.path
		v := branch.val
		if p.IsPublic() { // don't decrypt public data
			t.set(p, v)
			continue
		}
		if p.IsEncryptionKey() { // value uses different encryption key
			t.set(p, v)
			continue
		}
		if p.IsMacKey() { // value uses different encryption key
			t.set(p, v)
			continue
		}
		if p.IsChecksum() { // not encrypted at all
			t.set(p, v)
			continue
		}
		v, err = v.Plaintext(tree.encryptionKey)
		if err != nil {
			panic(fmt.Sprintf("%s: %s", p, err))
		}
		t.set(p, v)
	}
	return t
}

// Print displays a tree for human consumption.
func (tree *Tree) Print(w io.Writer) error {
	tree = tree.Decrypt()
	slice := tree.mapSlice()
	data, err := yaml.Marshal(slice)
	if err != nil {
		return errors.Wrap(err, "printing tree")
	}

	_, err = w.Write(data)
	return err
}

// Save stores a tree to disk for permanent, private archival.
func (tree *Tree) Save() error {
	tree = tree.Encrypt().Encode()
	slice := tree.mapSlice()

	data, err := yaml.Marshal(slice)
	if err != nil {
		return errors.Wrap(err, "saving tree")
	}

	// save to temporary file
	file, err := ioutil.TempFile("", "hush-")
	if err != nil {
		return errors.Wrap(err, "saving tree")
	}
	err = os.Chmod(file.Name(), safePerm)
	if err != nil {
		return errors.Wrap(err, "saving tree")
	}
	_, err = file.Write(data)
	io.WriteString(file, "hush-tree-checksum: ")
	io.WriteString(file, NewPlaintext(tree.Checksum(), Public).Encode().String())
	io.WriteString(file, "\n")
	file.Close()
	if err != nil {
		return errors.Wrap(err, "saving tree")
	}

	// move temporary file over top of permanent file
	hushPath, err := hushPath()
	if os.IsNotExist(err) {
		err = nil // we can create the file
	}
	if err != nil {
		return errors.Wrap(err, "saving tree")
	}
	err = os.Rename(file.Name(), hushPath)
	return errors.Wrap(err, "saving tree")
}

// stretchPassword converts a password and a salt into a
// cryptographically secure key.
func stretchPassword(password, salt []byte) []byte {
	pwKey := pbkdf2.Key(
		password, salt,
		2<<15, // iteration count (about 80ms on modern server)
		32,    // desired key size in bytes
		sha256.New,
	)
	return pwKey
}
<commit_msg>Factor out tree.Sort()<commit_after>package hush

import (
	"crypto/hmac"
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"sort"
	"strings"

	"golang.org/x/crypto/pbkdf2"

	"github.com/pkg/errors"

	yaml "gopkg.in/yaml.v2"
)

type Branch struct {
	path Path
	val  *Value
}

type Tree struct {
	branches []Branch
	index    map[Path]int

	encryptionKey []byte
	macKey        []byte
}

const safePerm = 0600 // rw- --- ---

func LoadTree() (*Tree, error) {
	hushPath, err := hushPath()
	if err != nil {
		return nil, err
	}

	stat, err := os.Stat(hushPath)
	if os.IsNotExist(err) {
		warn("hush file does not exist. assuming an empty one")
		return &Tree{}, nil
	}
	if err != nil {
		return nil, errors.Wrap(err, "can't stat hush file")
	}
	if (stat.Mode() & os.ModePerm) != safePerm {
		warn("hush file has loose permissions. fixing.")
		err := os.Chmod(hushPath, safePerm)
		if err != nil {
			die("couldn't fix permissions on hush file")
		}
	}

	file, err := os.Open(hushPath)
	if err != nil {
		return nil, errors.Wrap(err, "opening hush file")
	}
	hushData, err := ioutil.ReadAll(file)
	if err != nil {
		return nil, errors.Wrap(err, "can't read hush file")
	}

	keys := make(yaml.MapSlice, 0)
	err = yaml.Unmarshal(hushData, &keys)
	if err != nil {
		return nil, errors.Wrap(err, "can't parse hush file")
	}
	tree := newT(keys)
	return tree, nil
}

func newT(items yaml.MapSlice) *Tree {
	t := &Tree{
		branches: make([]Branch, 0, 3*len(items)),
		index:    make(map[Path]int),
	}
	newT_(items, []string{}, t)
	return t
}

func newT_(items yaml.MapSlice, crumbs []string, t *Tree) {
	n := len(crumbs)
	for _, item := range items {
		key := item.Key.(string)
		crumbs = append(crumbs, key)

		switch val := item.Value.(type) {
		case string:
			p := NewPath(strings.Join(crumbs, "/"))
			privacy := Private
			if p.IsPublic() {
				privacy = Public
			}
			t.set(p, NewEncoded(val, privacy))
		case yaml.MapSlice:
			newT_(val, crumbs, t)
		default:
			panic(fmt.Sprintf("unexpected type: %#v", val))
		}
		crumbs = crumbs[:n] // remove final crumb
	}
}

// implement sort.Interface interface
var _ sort.Interface = &Tree{}

func (t *Tree) Len() int { return len(t.branches) }
func (t *Tree) Less(i, j int) bool {
	return t.branches[i].path < t.branches[j].path
}
func (t *Tree) Swap(i, j int) {
	t.branches[i], t.branches[j] = t.branches[j], t.branches[i]
	t.index[t.branches[i].path] = i
	t.index[t.branches[j].path] = j
}

// Sort sorts the tree in place and defragments any deleted branches.
func (t *Tree) Sort() {
	sort.Stable(t)
}

func (t *Tree) mapSlice() yaml.MapSlice {
	t.Sort()

	var slice yaml.MapSlice
	for _, branch := range t.branches {
		if branch.path.IsChecksum() {
			// skip checksum. it's appended by Save()
			continue
		}
		crumbs := branch.path.AsCrumbs()
		slice = mapSlice_(slice, crumbs, branch.val.String())
	}
	return slice
}

func mapSlice_(slice yaml.MapSlice, path []string, value string) yaml.MapSlice {
	if len(path) == 0 {
		panic("path should never have 0 length")
	}
	if len(path) == 1 {
		return append(slice, yaml.MapItem{
			Key:   path[0],
			Value: value,
		})
	}

	var inner yaml.MapSlice
	if len(slice) == 0 {
		slice = append(slice, yaml.MapItem{Key: path[0]})
	} else {
		final := slice[len(slice)-1]
		if final.Key.(string) == path[0] {
			inner = final.Value.(yaml.MapSlice)
		} else {
			slice = append(slice, yaml.MapItem{Key: path[0]})
		}
	}
	slice[len(slice)-1].Value = mapSlice_(inner, path[1:], value)
	return slice
}

// Checksum returns a cryptographic message authentication code for this tree.
func (t *Tree) Checksum() []byte {
	if len(t.macKey) < 32 {
		panic("trying to calculate checksum without a MAC key")
	}
	mac := hmac.New(sha256.New, t.macKey)
	for _, branch := range t.branches {
		if branch.path.IsChecksum() {
			continue // don't checksum the checksum
		}
		mac.Write([]byte(branch.path))
		mac.Write([]byte(branch.val.String()))
	}
	sum := mac.Sum(nil)
	return sum
}

func (t *Tree) filter(pattern string) *Tree {
	keep := t.Empty()
	for _, branch := range t.branches {
		if matches(branch.path, pattern) {
			keep.set(branch.path, branch.val)
		}
	}
	return keep
}

func isLowercase(s string) bool {
	return s == strings.ToLower(s)
}

func matches(p Path, pattern string) bool {
	ps := strings.Split(string(p), "\t")
	patterns := strings.Split(pattern, "/")
	if len(patterns) > len(ps) {
		return false
	}

	ignoreCase := isLowercase(pattern)
	for i, pattern := range patterns {
		haystack := ps[i]
		if ignoreCase {
			haystack = strings.ToLower(haystack)
		}
		if !strings.Contains(haystack, pattern) {
			return false
		}
	}
	return true
}

func (t *Tree) get(p Path) (*Value, bool) {
	i, ok := t.index[p]
	if ok {
		return t.branches[i].val, true
	}
	return nil, false
}

func (t *Tree) set(p Path, val *Value) {
	i, ok := t.index[p]
	if ok {
		t.branches[i].val = val
	} else {
		t.branches = append(t.branches, Branch{p, val})
		t.index[p] = len(t.branches) - 1
	}
}

// Encrypt returns a copy of this tree with all leaves encrypted.
func (tree *Tree) Encrypt() *Tree {
	t := tree.Empty()
	for _, branch := range tree.branches {
		p := branch.path
		v := branch.val
		if p.IsPublic() { // don't encrypt public data
			t.set(p, v)
			continue
		}
		if p.IsEncryptionKey() { // value uses different encryption key
			t.set(p, v)
			continue
		}
		t.set(p, v.Ciphertext(t.encryptionKey))
	}
	return t
}

// Encode returns a copy of this tree with all leaves encoded into base64.
func (tree *Tree) Encode() *Tree {
	t := tree.Empty()
	for _, branch := range tree.branches {
		t.set(branch.path, branch.val.Encode())
	}
	return t
}

// Empty returns a copy of this tree with all the keys and values
// removed.  It retains any other data associated with this tree.
func (t *Tree) Empty() *Tree {
	tree := *t // shallow copy
	tree.branches = make([]Branch, 0, len(t.branches))
	tree.index = make(map[Path]int, len(t.branches))
	return &tree
}

// SetPassphrase sets the password that's used for performing
// encryption and decryption.
func (t *Tree) SetPassphrase(password []byte) error {
	p := NewPath("hush-configuration/salt")
	v, ok := t.get(p)
	if !ok {
		return errors.New("hush file missing salt")
	}
	v, err := v.Decode()
	if err != nil {
		return errors.Wrap(err, "decoding salt")
	}
	salt := v.plaintext
	pwKey := stretchPassword(password, salt)

	p = NewPath("hush-configuration/encryption-key")
	v, ok = t.get(p)
	if !ok {
		return errors.New("hush file missing encryption key")
	}
	v, err = v.Plaintext(pwKey)
	if err != nil {
		return fmt.Errorf("incorrect password or corrupted encryption key")
	}
	t.encryptionKey = v.plaintext

	p = NewPath("hush-configuration/mac-key")
	v, ok = t.get(p)
	if !ok {
		return errors.New("hush file missing MAC key")
	}
	v, err = v.Plaintext(pwKey)
	if err != nil {
		return fmt.Errorf("incorrect password or corrupted mac key")
	}
	t.macKey = v.plaintext

	// now that we have a password, we can verify the checksum
	got, ok := t.get(NewPath("hush-tree-checksum"))
	if !ok {
		return errors.New("hush file has no checksum")
	}
	got, err = got.Decode()
	if err != nil {
		return errors.Wrap(err, "decoding checksum")
	}
	expect := t.Checksum()
	if !hmac.Equal(got.plaintext, expect) {
		return errors.New("checksum doesn't match. file modified without hush command?")
	}

	return nil
}

// Decrypt returns a copy of this tree with all leaves decrypted.
func (tree *Tree) Decrypt() *Tree {
	var err error
	t := tree.Empty()
	for _, branch := range tree.branches {
		p := branch.path
		v := branch.val
		if p.IsPublic() { // don't decrypt public data
			t.set(p, v)
			continue
		}
		if p.IsEncryptionKey() { // value uses different encryption key
			t.set(p, v)
			continue
		}
		if p.IsMacKey() { // value uses different encryption key
			t.set(p, v)
			continue
		}
		if p.IsChecksum() { // not encrypted at all
			t.set(p, v)
			continue
		}
		v, err = v.Plaintext(tree.encryptionKey)
		if err != nil {
			panic(fmt.Sprintf("%s: %s", p, err))
		}
		t.set(p, v)
	}
	return t
}

// Print displays a tree for human consumption.
func (tree *Tree) Print(w io.Writer) error {
	tree = tree.Decrypt()
	slice := tree.mapSlice()
	data, err := yaml.Marshal(slice)
	if err != nil {
		return errors.Wrap(err, "printing tree")
	}

	_, err = w.Write(data)
	return err
}

// Save stores a tree to disk for permanent, private archival.
func (tree *Tree) Save() error {
	tree = tree.Encrypt().Encode()
	slice := tree.mapSlice()

	data, err := yaml.Marshal(slice)
	if err != nil {
		return errors.Wrap(err, "saving tree")
	}

	// save to temporary file
	file, err := ioutil.TempFile("", "hush-")
	if err != nil {
		return errors.Wrap(err, "saving tree")
	}
	err = os.Chmod(file.Name(), safePerm)
	if err != nil {
		return errors.Wrap(err, "saving tree")
	}
	_, err = file.Write(data)
	io.WriteString(file, "hush-tree-checksum: ")
	io.WriteString(file, NewPlaintext(tree.Checksum(), Public).Encode().String())
	io.WriteString(file, "\n")
	file.Close()
	if err != nil {
		return errors.Wrap(err, "saving tree")
	}

	// move temporary file over top of permanent file
	hushPath, err := hushPath()
	if os.IsNotExist(err) {
		err = nil // we can create the file
	}
	if err != nil {
		return errors.Wrap(err, "saving tree")
	}
	err = os.Rename(file.Name(), hushPath)
	return errors.Wrap(err, "saving tree")
}

// stretchPassword converts a password and a salt into a
// cryptographically secure key.
func stretchPassword(password, salt []byte) []byte {
	pwKey := pbkdf2.Key(
		password, salt,
		2<<15, // iteration count (about 80ms on modern server)
		32,    // desired key size in bytes
		sha256.New,
	)
	return pwKey
}
<|endoftext|>


<commit_before>package reverseproxy

<commit_msg>Removing trie<commit_after><|endoftext|>


<commit_before>package ghg

import (
	"archive/tar"
	"archive/zip"
	"bytes"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/octokit/go-octokit/octokit"
	"github.com/pkg/errors"
)

type ghg struct {
	binDir  string
	target  string
	version string
	client  *octokit.Client
}

func getOctCli(token string) *octokit.Client {
	var auth octokit.AuthMethod
	if token != "" {
		auth = octokit.TokenAuth{AccessToken: token}
	}
	return octokit.NewClient(auth)
}

func (gh *ghg) install() error {
	owner, repo, err := gh.getRepoAndOwner()
	if err != nil {
		return errors.Wrap(err, "failed to resolve target")
	}
	url, err := octokit.ReleasesLatestURL.Expand(octokit.M{"owner": owner, "repo": repo})
	if err != nil {
		return errors.Wrap(err, "failed to build GitHub URL")
	}
	release, r := gh.client.Releases(url).Latest()
	if r.HasError() {
		return errors.Wrap(r.Err, "failed to fetch latest release")
	}
	tag := release.TagName
	goarch := runtime.GOARCH
	goos := runtime.GOOS
	var urls []string
	for _, asset := range release.Assets {
		name := asset.Name
		if strings.Contains(name, goarch) && strings.Contains(name, goos) {
			urls = append(urls, fmt.Sprintf("https://github.com/%s/%s/releases/download/%s/%s", owner, repo, tag, name))
		}
	}

	for _, url := range urls {
		archivePath, err := download(url)
		if err != nil {
			return errors.Wrap(err, "failed to download")
		}
		workDir := filepath.Join(filepath.Dir(archivePath), "work")
		err = extract(archivePath, workDir)
		if err != nil {
			return errors.Wrap(err, "failed to extract")
		}
	}
	return nil
}

func download(url string) (fpath string, err error) {
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		err = errors.Wrap(err, "failed to create request")
		return
	}
	req.Header.Set("User-Agent", fmt.Sprintf("ghg/%s", version))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		err = errors.Wrap(err, "failed to create request")
		return
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		err = errors.Wrap(err, "failed to read response")
		return
	}
	archiveBase := path.Base(url)
	tempdir, err := ioutil.TempDir("", "ghg-")
	if err != nil {
		err = errors.Wrap(err, "failed to create tempdir")
		return
	}
	defer func() {
		if err != nil {
			os.RemoveAll(tempdir)
		}
	}()
	fpath = filepath.Join(tempdir, archiveBase)
	f, err := os.OpenFile(filepath.Join(tempdir, archiveBase), os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		err = errors.Wrap(err, "failed to open file")
		return
	}
	defer f.Close()
	_, err = f.Write(body)
	if err != nil {
		err = errors.Wrap(err, "failed to read response")
		return
	}
	return fpath, nil
}

func extract(src, dest string) error {
	base := filepath.Base(src)
	if strings.HasSuffix(base, ".zip") {
		return unzip(src, dest)
	}
	if strings.HasSuffix(base, ".tar.gz") || strings.HasSuffix(base, ".tgz") {
		return extractTar(src, dest)
	}
	return nil
}

func unzip(src, dest string) error {
	r, err := zip.OpenReader(src)
	if err != nil {
		return err
	}
	defer r.Close()

	var move = func(f *zip.File) error {
		rc, err := f.Open()
		if err != nil {
			return err
		}
		defer rc.Close()

		path := filepath.Join(dest, f.Name)
		if f.FileInfo().IsDir() {
			os.MkdirAll(path, f.Mode())
		} else {
			os.MkdirAll(filepath.Dir(path), 0755)
			f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
			if err != nil {
				return errors.Wrap(err, "failed to openfile")
			}
			defer f.Close()
			_, err = io.Copy(f, rc)
			if err != nil {
				return errors.Wrap(err, "failed to copy")
			}
		}
		return nil
	}

	for _, f := range r.File {
		err := move(f)
		if err != nil {
			return err
		}
	}

	return nil
}

func extractTar(src, dest string) error {
	file, err := os.Open(src)
	if err != nil {
		return errors.Wrap(err, fmt.Sprintf("failed to open tgz file: %s", src))
	}
	defer file.Close()

	rdr, err := gzip.NewReader(file)
	if err != nil {
		return errors.Wrap(err, fmt.Sprintf("failed to read tgz file: %s", src))
	}
	defer rdr.Close()

	tr := tar.NewReader(rdr)

	var header *tar.Header
	for {
		header, err = tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return errors.Wrap(err, "failed to read tar")
		}

		var buf *bytes.Buffer
		_, err = io.Copy(buf, tr)
		if err != nil {
			return errors.Wrap(err, "failed io.Copy")
		}

		if err = ioutil.WriteFile(filepath.Join(dest, header.Name), buf.Bytes(), 0755); err != nil {
			return errors.Wrap(err, fmt.Sprintf("failed to extract file: %s", src))
		}
	}
	return nil
}

func (gh *ghg) getRepoAndOwner() (owner, repo string, err error) {
	arr := strings.SplitN(gh.target, "/", 2)
	if len(arr) < 1 {
		return "", "", fmt.Errorf("target invalid")
	}
	owner = arr[0]
	repo = arr[0]
	if len(arr) > 1 {
		repo = arr[1]
	}
	return
}
<commit_msg>use github.com/mholt/archiver<commit_after>package ghg

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/mholt/archiver"
	"github.com/octokit/go-octokit/octokit"
	"github.com/pkg/errors"
)

type ghg struct {
	binDir  string
	target  string
	version string
	client  *octokit.Client
}

func getOctCli(token string) *octokit.Client {
	var auth octokit.AuthMethod
	if token != "" {
		auth = octokit.TokenAuth{AccessToken: token}
	}
	return octokit.NewClient(auth)
}

func (gh *ghg) install() error {
	owner, repo, err := gh.getRepoAndOwner()
	if err != nil {
		return errors.Wrap(err, "failed to resolve target")
	}
	url, err := octokit.ReleasesLatestURL.Expand(octokit.M{"owner": owner, "repo": repo})
	if err != nil {
		return errors.Wrap(err, "failed to build GitHub URL")
	}
	release, r := gh.client.Releases(url).Latest()
	if r.HasError() {
		return errors.Wrap(r.Err, "failed to fetch latest release")
	}
	tag := release.TagName
	goarch := runtime.GOARCH
	goos := runtime.GOOS
	var urls []string
	for _, asset := range release.Assets {
		name := asset.Name
		if strings.Contains(name, goarch) && strings.Contains(name, goos) {
			urls = append(urls, fmt.Sprintf("https://github.com/%s/%s/releases/download/%s/%s", owner, repo, tag, name))
		}
	}

	for _, url := range urls {
		archivePath, err := download(url)
		if err != nil {
			return errors.Wrap(err, "failed to download")
		}
		workDir := filepath.Join(filepath.Dir(archivePath), "work")
		os.MkdirAll(workDir, 0755)
		err = extract(archivePath, workDir)
		if err != nil {
			return errors.Wrap(err, "failed to extract")
		}
	}
	return nil
}

func download(url string) (fpath string, err error) {
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		err = errors.Wrap(err, "failed to create request")
		return
	}
	req.Header.Set("User-Agent", fmt.Sprintf("ghg/%s", version))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		err = errors.Wrap(err, "failed to create request")
		return
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		err = errors.Wrap(err, "failed to read response")
		return
	}
	archiveBase := path.Base(url)
	tempdir, err := ioutil.TempDir("", "ghg-")
	if err != nil {
		err = errors.Wrap(err, "failed to create tempdir")
		return
	}
	defer func() {
		if err != nil {
			os.RemoveAll(tempdir)
		}
	}()
	fpath = filepath.Join(tempdir, archiveBase)
	f, err := os.OpenFile(filepath.Join(tempdir, archiveBase), os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		err = errors.Wrap(err, "failed to open file")
		return
	}
	defer f.Close()
	_, err = f.Write(body)
	if err != nil {
		err = errors.Wrap(err, "failed to read response")
		return
	}
	return fpath, nil
}

func extract(src, dest string) error {
	base := filepath.Base(src)
	if strings.HasSuffix(base, ".zip") {
		return archiver.Unzip(src, dest)
	}
	if strings.HasSuffix(base, ".tar.gz") || strings.HasSuffix(base, ".tgz") {
		return archiver.UntarGz(src, dest)
	}
	return fmt.Errorf("failed to extract file: %s", src)
}

func (gh *ghg) getRepoAndOwner() (owner, repo string, err error) {
	arr := strings.SplitN(gh.target, "/", 2)
	if len(arr) < 1 {
		return "", "", fmt.Errorf("target invalid")
	}
	owner = arr[0]
	repo = arr[0]
	if len(arr) > 1 {
		repo = arr[1]
	}
	return
}
<|endoftext|>


<commit_before>package ghg

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"

	"github.com/mholt/archiver"
	"github.com/octokit/go-octokit/octokit"
	"github.com/pkg/errors"
)

type ghg struct {
	binDir  string
	target  string
	version string
	client  *octokit.Client
}

func getOctCli(token string) *octokit.Client {
	var auth octokit.AuthMethod
	if token != "" {
		auth = octokit.TokenAuth{AccessToken: token}
	}
	return octokit.NewClient(auth)
}

var archiveReg = regexp.MustCompile(`\.(?:zip|tgz|tar\.gz)$`)

func (gh *ghg) install() error {
	owner, repo, err := gh.getRepoAndOwner()
	if err != nil {
		return errors.Wrap(err, "failed to resolve target")
	}
	url, err := octokit.ReleasesLatestURL.Expand(octokit.M{"owner": owner, "repo": repo})
	if err != nil {
		return errors.Wrap(err, "failed to build GitHub URL")
	}
	release, r := gh.client.Releases(url).Latest()
	if r.HasError() {
		return errors.Wrap(r.Err, "failed to fetch latest release")
	}
	tag := release.TagName
	goarch := runtime.GOARCH
	goos := runtime.GOOS
	var urls []string
	for _, asset := range release.Assets {
		name := asset.Name
		if strings.Contains(name, goarch) && strings.Contains(name, goos) && archiveReg.MatchString(name) {
			urls = append(urls, fmt.Sprintf("https://github.com/%s/%s/releases/download/%s/%s", owner, repo, tag, name))
		}
	}

	for _, url := range urls {
		archivePath, err := download(url)
		if err != nil {
			return errors.Wrap(err, "failed to download")
		}
		workDir := filepath.Join(filepath.Dir(archivePath), "work")
		os.MkdirAll(workDir, 0755)
		err = extract(archivePath, workDir)
		if err != nil {
			return errors.Wrap(err, "failed to extract")
		}
		err = pickupExecutable(workDir, ".")
		if err != nil {
			return errors.Wrap(err, "failed to pickup")
		}
	}
	return nil
}

func download(url string) (fpath string, err error) {
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		err = errors.Wrap(err, "failed to create request")
		return
	}
	req.Header.Set("User-Agent", fmt.Sprintf("ghg/%s", version))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		err = errors.Wrap(err, "failed to create request")
		return
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		err = errors.Wrap(err, "failed to read response")
		return
	}
	archiveBase := path.Base(url)
	tempdir, err := ioutil.TempDir("", "ghg-")
	if err != nil {
		err = errors.Wrap(err, "failed to create tempdir")
		return
	}
	defer func() {
		if err != nil {
			os.RemoveAll(tempdir)
		}
	}()
	fpath = filepath.Join(tempdir, archiveBase)
	f, err := os.OpenFile(filepath.Join(tempdir, archiveBase), os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		err = errors.Wrap(err, "failed to open file")
		return
	}
	defer f.Close()
	_, err = f.Write(body)
	if err != nil {
		err = errors.Wrap(err, "failed to read response")
		return
	}
	return fpath, nil
}

func extract(src, dest string) error {
	base := filepath.Base(src)
	if strings.HasSuffix(base, ".zip") {
		return archiver.Unzip(src, dest)
	}
	if strings.HasSuffix(base, ".tar.gz") || strings.HasSuffix(base, ".tgz") {
		return archiver.UntarGz(src, dest)
	}
	return fmt.Errorf("failed to extract file: %s", src)
}

func (gh *ghg) getRepoAndOwner() (owner, repo string, err error) {
	arr := strings.SplitN(gh.target, "/", 2)
	if len(arr) < 1 {
		return "", "", fmt.Errorf("target invalid")
	}
	owner = arr[0]
	repo = arr[0]
	if len(arr) > 1 {
		repo = arr[1]
	}
	return
}

var executableReg = regexp.MustCompile(`^[a-z][-_a-zA-Z0-9]+(?:\.exe)?$`)

func pickupExecutable(src, dest string) error {
	defer os.RemoveAll(src)
	return filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() {
			return err
		}
		if name := info.Name(); (info.Mode()&0111) != 0 && executableReg.MatchString(name) {
			return os.Rename(path, filepath.Join(dest, name))
		}
		return nil
	})
}
<commit_msg>we can specify version<commit_after>package ghg

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"

	"github.com/mholt/archiver"
	"github.com/octokit/go-octokit/octokit"
	"github.com/pkg/errors"
)

type ghg struct {
	binDir  string
	target  string
	version string
	client  *octokit.Client
}

func getOctCli(token string) *octokit.Client {
	var auth octokit.AuthMethod
	if token != "" {
		auth = octokit.TokenAuth{AccessToken: token}
	}
	return octokit.NewClient(auth)
}

var releaseByTagURL = octokit.Hyperlink("repos/{owner}/{repo}/releases/tags/{tag}")
var archiveReg = regexp.MustCompile(`\.(?:zip|tgz|tar\.gz)$`)

func (gh *ghg) install() error {
	owner, repo, tag, err := gh.getOwnerRepoAndTag()
	if err != nil {
		return errors.Wrap(err, "failed to resolve target")
	}
	var url *url.URL
	if tag == "" {
		url, err = octokit.ReleasesLatestURL.Expand(octokit.M{"owner": owner, "repo": repo})
	} else {
		url, err = releaseByTagURL.Expand(octokit.M{"owner": owner, "repo": repo, "tag": tag})
	}
	if err != nil {
		return errors.Wrap(err, "failed to build GitHub URL")
	}
	release, r := gh.client.Releases(url).Latest()
	if r.HasError() {
		return errors.Wrap(r.Err, "failed to fetch latest release")
	}
	tag = release.TagName
	goarch := runtime.GOARCH
	goos := runtime.GOOS
	var urls []string
	for _, asset := range release.Assets {
		name := asset.Name
		if strings.Contains(name, goarch) && strings.Contains(name, goos) && archiveReg.MatchString(name) {
			urls = append(urls, fmt.Sprintf("https://github.com/%s/%s/releases/download/%s/%s", owner, repo, tag, name))
		}
	}

	for _, url := range urls {
		archivePath, err := download(url)
		if err != nil {
			return errors.Wrap(err, "failed to download")
		}
		workDir := filepath.Join(filepath.Dir(archivePath), "work")
		os.MkdirAll(workDir, 0755)
		err = extract(archivePath, workDir)
		if err != nil {
			return errors.Wrap(err, "failed to extract")
		}
		err = pickupExecutable(workDir, ".")
		if err != nil {
			return errors.Wrap(err, "failed to pickup")
		}
	}
	return nil
}

func download(url string) (fpath string, err error) {
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		err = errors.Wrap(err, "failed to create request")
		return
	}
	req.Header.Set("User-Agent", fmt.Sprintf("ghg/%s", version))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		err = errors.Wrap(err, "failed to create request")
		return
	}
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		err = errors.Wrap(err, "failed to read response")
		return
	}
	archiveBase := path.Base(url)
	tempdir, err := ioutil.TempDir("", "ghg-")
	if err != nil {
		err = errors.Wrap(err, "failed to create tempdir")
		return
	}
	defer func() {
		if err != nil {
			os.RemoveAll(tempdir)
		}
	}()
	fpath = filepath.Join(tempdir, archiveBase)
	f, err := os.OpenFile(filepath.Join(tempdir, archiveBase), os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		err = errors.Wrap(err, "failed to open file")
		return
	}
	defer f.Close()
	_, err = f.Write(body)
	if err != nil {
		err = errors.Wrap(err, "failed to read response")
		return
	}
	return fpath, nil
}

func extract(src, dest string) error {
	base := filepath.Base(src)
	if strings.HasSuffix(base, ".zip") {
		return archiver.Unzip(src, dest)
	}
	if strings.HasSuffix(base, ".tar.gz") || strings.HasSuffix(base, ".tgz") {
		return archiver.UntarGz(src, dest)
	}
	return fmt.Errorf("failed to extract file: %s", src)
}

var targetReg = regexp.MustCompile(`^(?:([^/]+)/)?([^@]+)(?:@(.+))?$`)

func (gh *ghg) getOwnerRepoAndTag() (owner, repo, tag string, err error) {
	if matches := targetReg.FindStringSubmatch(gh.target); len(matches) == 4 {
		owner = matches[1]
		repo = matches[2]
		tag = matches[3]
		if owner == "" {
			owner = repo
		}
	}
	return
}

var executableReg = regexp.MustCompile(`^[a-z][-_a-zA-Z0-9]+(?:\.exe)?$`)

func pickupExecutable(src, dest string) error {
	defer os.RemoveAll(src)
	return filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() {
			return err
		}
		if name := info.Name(); (info.Mode()&0111) != 0 && executableReg.MatchString(name) {
			return os.Rename(path, filepath.Join(dest, name))
		}
		return nil
	})
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
)

func main() {
	flag.Parse()

	commands := map[string]command{
		"ruler": makeRulerCommand(),
	}

	args := flag.Args()

	if len(args) <= 0 {
		flag.Usage()
		os.Exit(1)
	}
	cmd, ok := commands[args[0]]
	if !ok {
		printError("%q is not a valid subcommand\n", args[0])
		os.Exit(1)
	}
	if err := cmd.fn(args[1:]); err != nil {
		printError("error running %q: %s", args[0], err)
		os.Exit(1)
	}
}

type command struct {
	fs *flag.FlagSet
	fn func([]string) error
}

func makeRulerCommand() command {
	fs := flag.NewFlagSet("ruler", flag.ExitOnError)
	spacesToIndent := fs.Int("indent", 0, "the number of spaces to print before the ruler")
	fn := func(args []string) error {
		fs.Parse(args)
		_, err := fmt.Println(ruler(*spacesToIndent))
		return err
	}
	return command{
		fs: fs,
		fn: fn,
	}
}

func ruler(spacesToIndent int) string {
	var ruler string
	for i := 0; i < spacesToIndent; i++ {
		ruler += " "
	}
	for i := 1; i <= 140; i++ {
		if i%5 == 0 {
			ruler += "|"
		} else {
			ruler += "-"
		}
	}
	return ruler
}

func printError(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, os.Args[0]+": "+format, args...)
}
<commit_msg>Define default usage<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"path/filepath"
)

var programName string

func init() {
	programName = filepath.Base(os.Args[0])
}

func main() {
	commands := map[string]command{
		"ruler": makeRulerCommand(),
	}

	flag.Usage = func() {
		fmt.Printf("%s is a Twitter command-line interface.\n\n", programName)
		var maxCommandLen int
		for name := range commands {
			if len := len(name); len > maxCommandLen {
				maxCommandLen = len
			}
		}
		for name, subcommand := range commands {
			formatString := fmt.Sprintf("%%%ds: %%s\n", maxCommandLen)
			fmt.Printf(formatString, name, subcommand.desc)
		}
	}

	flag.Parse()

	args := flag.Args()

	if len(args) <= 0 {
		flag.Usage()
		os.Exit(2)
	}
	cmd, ok := commands[args[0]]
	if !ok {
		printError("%q is not a valid subcommand\n", args[0])
		os.Exit(1)
	}
	if err := cmd.fn(args[1:]); err != nil {
		printError("error running %q: %s", args[0], err)
		os.Exit(1)
	}
}

type command struct {
	fs   *flag.FlagSet
	desc string
	fn   func([]string) error
}

func makeRulerCommand() command {
	fs := flag.NewFlagSet("ruler", flag.ExitOnError)
	spacesToIndent := fs.Int("indent", 0, "the number of spaces to print before the ruler")
	fn := func(args []string) error {
		fs.Parse(args)
		_, err := fmt.Println(ruler(*spacesToIndent))
		return err
	}
	return command{
		fs:   fs,
		desc: "prints a 140-character ruler",
		fn:   fn,
	}
}

func ruler(spacesToIndent int) string {
	var ruler string
	for i := 0; i < spacesToIndent; i++ {
		ruler += " "
	}
	for i := 1; i <= 140; i++ {
		if i%5 == 0 {
			ruler += "|"
		} else {
			ruler += "-"
		}
	}
	return ruler
}

func printError(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, programName+": "+format, args...)
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
)

const VERSION = "v0.0.2"

var logger = log.New(os.Stderr, "", 0)

type OriginConfig struct {
	Addr   string
	Direct bool
}

func main() {
	port := flag.Int("p", 0, "")
	ip := flag.String("i", "127.0.0.1", "")
	originPort := flag.Int("P", 0, "")
	originHost := flag.String("H", "", "")
	originDirect := flag.Bool("D", false, "")
	certPath := flag.String("c", "", "")
	keyPath := flag.String("k", "", "")
	outputLogFormat := flag.String("o", "default", "")
	version := flag.Bool("version", false, "")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [OPTIONS]\n\n", os.Args[0])
		fmt.Println("Options:")
		fmt.Println("  -p:        Port (Default: 443)")
		fmt.Println("  -i:        IP Address (Default: 127.0.0.1)")
		fmt.Println("  -P:        Origin port")
		fmt.Println("  -H:        Origin host")
		fmt.Println("  -D:        Use HTTP/2 direct mode to connect origin")
		fmt.Println("  -c:        Certificate file")
		fmt.Println("  -k:        Certificate key file")
		fmt.Println("  -o:        Output log format (default or json, Default: default)")
		fmt.Println("  --version: Display version information and exit.")
		fmt.Println("  --help:    Display this help and exit.")
		os.Exit(1)
	}

	flag.Parse()

	if *version {
		logger.Printf("h2a %s\n", VERSION)
		os.Exit(0)
	}

	if *port == 0 {
		*port = 443
	}
	addr := fmt.Sprintf("%s:%d", *ip, *port)

	if *originPort == 0 {
		logger.Fatalln("Origin port is not specified")
	}
	if *originHost == "" {
		logger.Fatalln("Origin host is not specified")
	}
	originConfig := OriginConfig{
		Addr:   fmt.Sprintf("%s:%d", *originHost, *originPort),
		Direct: *originDirect,
	}

	var formatter Formatter
	if *outputLogFormat == "json" {
		formatter = JSONFormatter
	} else {
		formatter = DefaultFormatter
	}

	cert, err := tls.LoadX509KeyPair(*certPath, *keyPath)
	if err != nil {
		logger.Fatalln("Invalid certificate file")
	}

	config := &tls.Config{}
	config.Certificates = []tls.Certificate{cert}
	config.NextProtos = append(config.NextProtos, "h2", "h2-16", "h2-15", "h2-14", "http/1.1")

	server, err := tls.Listen("tcp", addr, config)
	if err != nil {
		logger.Fatalf("Could not bind address - %s\n", addr)
	}

	defer server.Close()

	for {
		remoteConn, err := server.Accept()
		if err != nil {
			logger.Printf("Unable to accept: %s", err)
			continue
		}

		go handlePeer(remoteConn, originConfig, formatter)
	}
}

func handlePeer(remoteConn net.Conn, originConfig OriginConfig, formatter Formatter) {
	var originConn net.Conn
	var err error
	h2 := true

	defer remoteConn.Close()

	dumper := NewFrameDumper(remoteConn.RemoteAddr(), formatter)
	defer dumper.Close()

	remoteCh, remoteErrCh := handleConnection(remoteConn)

	select {
	case chunk := <-remoteCh:
		connState := remoteConn.(*tls.Conn).ConnectionState()

		if connState.NegotiatedProtocol == "" {
			connState.NegotiatedProtocol = "http/1.1"
		}
		if connState.NegotiatedProtocol == "http/1.1" {
			h2 = true
		}

		dumper.DumpConnectionState(connState)

		config := &tls.Config{}
		config.CipherSuites = []uint16{connState.CipherSuite}
		config.ServerName = connState.ServerName
		config.InsecureSkipVerify = true
		config.NextProtos = append(config.NextProtos, connState.NegotiatedProtocol)

		if originConfig.Direct {
			originConn, err = net.Dial("tcp", originConfig.Addr)
		} else {
			dialer := new(net.Dialer)
			originConn, err = tls.DialWithDialer(dialer, "tcp", originConfig.Addr, config)
		}
		if err != nil {
			logger.Printf("Unable to connect to the origin: %s", err)
			return
		}

		defer originConn.Close()

		_, err = originConn.Write(chunk)
		if err != nil {
			logger.Printf("Unable to write data to the origin: %s", err)
			return
		}

		if h2 {
			dumper.DumpFrame(chunk, true)
		}

	case err := <-remoteErrCh:
		if err != io.EOF {
			logger.Printf("Connection error: %s", err)
		}
		return
	}

	originCh, originErrCh := handleConnection(originConn)

	for {
		select {
		case chunk := <-remoteCh:
			_, err := originConn.Write(chunk)
			if err != nil {
				logger.Printf("Unable to write data to the origin: %s", err)
				return
			}

			if h2 {
				dumper.DumpFrame(chunk, true)
			}

		case err := <-remoteErrCh:
			if err != io.EOF {
				logger.Printf("Connection error: %s", err)
			}
			return

		case chunk := <-originCh:
			_, err := remoteConn.Write(chunk)
			if err != nil {
				logger.Printf("Unable to write data to the connection: %s", err)
				return
			}

			if h2 {
				dumper.DumpFrame(chunk, false)
			}

		case err := <-originErrCh:
			if err != io.EOF {
				logger.Printf("Origin error: %s", err)
			}
			return
		}
	}
}

func handleConnection(conn net.Conn) (<-chan []byte, <-chan error) {
	dataCh := make(chan []byte)
	errCh := make(chan error, 1)

	go func() {
		for {
			buf := make([]byte, 16384)

			n, err := conn.Read(buf)
			if err != nil {
				errCh <- err
				break
			}

			dataCh <- buf[:n]
		}
	}()

	return dataCh, errCh
}
<commit_msg>Refactor code<commit_after>package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
)

const VERSION = "v0.0.2"

var logger = log.New(os.Stderr, "", 0)

type OriginConfig struct {
	Addr   string
	Direct bool
}

func main() {
	port := flag.String("p", "", "")
	ip := flag.String("i", "127.0.0.1", "")
	originPort := flag.String("P", "", "")
	originHost := flag.String("H", "", "")
	originDirect := flag.Bool("D", false, "")
	certPath := flag.String("c", "", "")
	keyPath := flag.String("k", "", "")
	outputLogFormat := flag.String("o", "default", "")
	version := flag.Bool("version", false, "")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [OPTIONS]\n\n", os.Args[0])
		fmt.Println("Options:")
		fmt.Println("  -p:        Port (Default: 443)")
		fmt.Println("  -i:        IP Address (Default: 127.0.0.1)")
		fmt.Println("  -P:        Origin port")
		fmt.Println("  -H:        Origin host")
		fmt.Println("  -D:        Use HTTP/2 direct mode to connect origin")
		fmt.Println("  -c:        Certificate file")
		fmt.Println("  -k:        Certificate key file")
		fmt.Println("  -o:        Output log format (default or json, Default: default)")
		fmt.Println("  --version: Display version information and exit.")
		fmt.Println("  --help:    Display this help and exit.")
		os.Exit(1)
	}

	flag.Parse()

	if *version {
		logger.Printf("h2a %s\n", VERSION)
		os.Exit(0)
	}

	if *port == "" {
		*port = "443"
	}
	addr := net.JoinHostPort(*ip, *port)

	if *originPort == "" {
		logger.Fatalln("Origin port is not specified")
	}
	if *originHost == "" {
		logger.Fatalln("Origin host is not specified")
	}
	originConfig := OriginConfig{
		Addr:   net.JoinHostPort(*originHost, *originPort),
		Direct: *originDirect,
	}

	var formatter Formatter
	if *outputLogFormat == "json" {
		formatter = JSONFormatter
	} else {
		formatter = DefaultFormatter
	}

	cert, err := tls.LoadX509KeyPair(*certPath, *keyPath)
	if err != nil {
		logger.Fatalln("Invalid certificate file")
	}

	config := &tls.Config{}
	config.Certificates = []tls.Certificate{cert}
	config.NextProtos = append(config.NextProtos, "h2", "h2-16", "h2-15", "h2-14", "http/1.1")

	server, err := tls.Listen("tcp", addr, config)
	if err != nil {
		logger.Fatalf("Could not bind address - %s\n", addr)
	}

	defer server.Close()

	for {
		remoteConn, err := server.Accept()
		if err != nil {
			logger.Printf("Unable to accept: %s", err)
			continue
		}

		go handlePeer(remoteConn, originConfig, formatter)
	}
}

func handlePeer(remoteConn net.Conn, originConfig OriginConfig, formatter Formatter) {
	var originConn net.Conn
	var err error
	h2 := true

	defer remoteConn.Close()

	dumper := NewFrameDumper(remoteConn.RemoteAddr(), formatter)
	defer dumper.Close()

	remoteCh, remoteErrCh := handleConnection(remoteConn)

	select {
	case chunk := <-remoteCh:
		connState := remoteConn.(*tls.Conn).ConnectionState()

		if connState.NegotiatedProtocol == "" {
			connState.NegotiatedProtocol = "http/1.1"
		}
		if connState.NegotiatedProtocol == "http/1.1" {
			h2 = true
		}

		dumper.DumpConnectionState(connState)

		config := &tls.Config{}
		config.CipherSuites = []uint16{connState.CipherSuite}
		config.ServerName = connState.ServerName
		config.InsecureSkipVerify = true
		config.NextProtos = append(config.NextProtos, connState.NegotiatedProtocol)

		if originConfig.Direct {
			originConn, err = net.Dial("tcp", originConfig.Addr)
		} else {
			originConn, err = tls.Dial("tcp", originConfig.Addr, config)
		}
		if err != nil {
			logger.Printf("Unable to connect to the origin: %s", err)
			return
		}

		defer originConn.Close()

		_, err = originConn.Write(chunk)
		if err != nil {
			logger.Printf("Unable to write data to the origin: %s", err)
			return
		}

		if h2 {
			dumper.DumpFrame(chunk, true)
		}

	case err := <-remoteErrCh:
		if err != io.EOF {
			logger.Printf("Connection error: %s", err)
		}
		return
	}

	originCh, originErrCh := handleConnection(originConn)

	for {
		select {
		case chunk := <-remoteCh:
			_, err := originConn.Write(chunk)
			if err != nil {
				logger.Printf("Unable to write data to the origin: %s", err)
				return
			}

			if h2 {
				dumper.DumpFrame(chunk, true)
			}

		case err := <-remoteErrCh:
			if err != io.EOF {
				logger.Printf("Connection error: %s", err)
			}
			return

		case chunk := <-originCh:
			_, err := remoteConn.Write(chunk)
			if err != nil {
				logger.Printf("Unable to write data to the connection: %s", err)
				return
			}

			if h2 {
				dumper.DumpFrame(chunk, false)
			}

		case err := <-originErrCh:
			if err != io.EOF {
				logger.Printf("Origin error: %s", err)
			}
			return
		}
	}
}

func handleConnection(conn net.Conn) (<-chan []byte, <-chan error) {
	dataCh := make(chan []byte)
	errCh := make(chan error, 1)

	go func() {
		for {
			buf := make([]byte, 16384)

			n, err := conn.Read(buf)
			if err != nil {
				errCh <- err
				break
			}

			dataCh <- buf[:n]
		}
	}()

	return dataCh, errCh
}
<|endoftext|>