<commit_before>package execdriver

import (
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/docker/docker/daemon/execdriver/native/template"
	"github.com/docker/docker/pkg/ulimit"
	"github.com/docker/libcontainer"
	"github.com/docker/libcontainer/cgroups/fs"
	"github.com/docker/libcontainer/configs"
)

// Context is a generic key value pair that allows
// arbatrary data to be sent
type Context map[string]string

var (
	ErrNotRunning              = errors.New("Container is not running")
	ErrWaitTimeoutReached      = errors.New("Wait timeout reached")
	ErrDriverAlreadyRegistered = errors.New("A driver already registered this docker init function")
	ErrDriverNotFound          = errors.New("The requested docker init has not been found")
)

type StartCallback func(*ProcessConfig, int)

// Driver specific information based on
// processes registered with the driver
type Info interface {
	IsRunning() bool
}

// Terminal in an interface for drivers to implement
// if they want to support Close and Resize calls from
// the core
type Terminal interface {
	io.Closer
	Resize(height, width int) error
}

type TtyTerminal interface {
	Master() libcontainer.Console
}

// ExitStatus provides exit reasons for a container.
type ExitStatus struct {
	// The exit code with which the container exited.
	ExitCode int

	// Whether the container encountered an OOM.
	OOMKilled bool
}

type Driver interface {
	Run(c *Command, pipes *Pipes, startCallback StartCallback) (ExitStatus, error) // Run executes the process and blocks until the process exits and returns the exit code
	// Exec executes the process in an existing container, blocks until the process exits and returns the exit code
	Exec(c *Command, processConfig *ProcessConfig, pipes *Pipes, startCallback StartCallback) (int, error)
	Kill(c *Command, sig int) error
	Pause(c *Command) error
	Unpause(c *Command) error
	Name() string                                 // Driver name
	Info(id string) Info                          // "temporary" hack (until we move state from core to plugins)
	GetPidsForContainer(id string) ([]int, error) // Returns a list of pids for the given container.
	Terminate(c *Command) error                   // kill it with fire
	Clean(id string) error                        // clean all traces of container exec
	Stats(id string) (*ResourceStats, error)      // Get resource stats for a running container
}

// Network settings of the container
type Network struct {
	Interface      *NetworkInterface `json:"interface"` // if interface is nil then networking is disabled
	Mtu            int               `json:"mtu"`
	ContainerID    string            `json:"container_id"` // id of the container to join network.
	HostNetworking bool              `json:"host_networking"`
}

// IPC settings of the container
type Ipc struct {
	ContainerID string `json:"container_id"` // id of the container to join ipc.
	HostIpc     bool   `json:"host_ipc"`
}

// PID settings of the container
type Pid struct {
	HostPid bool `json:"host_pid"`
}

type NetworkInterface struct {
	Gateway              string `json:"gateway"`
	IPAddress            string `json:"ip"`
	IPPrefixLen          int    `json:"ip_prefix_len"`
	MacAddress           string `json:"mac"`
	Bridge               string `json:"bridge"`
	GlobalIPv6Address    string `json:"global_ipv6"`
	LinkLocalIPv6Address string `json:"link_local_ipv6"`
	GlobalIPv6PrefixLen  int    `json:"global_ipv6_prefix_len"`
	IPv6Gateway          string `json:"ipv6_gateway"`
}

type Resources struct {
	Memory     int64            `json:"memory"`
	MemorySwap int64            `json:"memory_swap"`
	CpuShares  int64            `json:"cpu_shares"`
	CpusetCpus string           `json:"cpuset_cpus"`
	Rlimits    []*ulimit.Rlimit `json:"rlimits"`
}

type ResourceStats struct {
	*libcontainer.Stats
	Read        time.Time `json:"read"`
	MemoryLimit int64     `json:"memory_limit"`
	SystemUsage uint64    `json:"system_usage"`
}

type Mount struct {
	Source      string `json:"source"`
	Destination string `json:"destination"`
	Writable    bool   `json:"writable"`
	Private     bool   `json:"private"`
	Slave       bool   `json:"slave"`
}

// Describes a process that will be run inside a container.
type ProcessConfig struct {
	exec.Cmd `json:"-"`

	Privileged bool     `json:"privileged"`
	User       string   `json:"user"`
	Tty        bool     `json:"tty"`
	Entrypoint string   `json:"entrypoint"`
	Arguments  []string `json:"arguments"`
	Terminal   Terminal `json:"-"` // standard or tty terminal
	Console    string   `json:"-"` // dev/console path
}

// Process wrapps an os/exec.Cmd to add more metadata
type Command struct {
	ID                 string            `json:"id"`
	Rootfs             string            `json:"rootfs"` // root fs of the container
	ReadonlyRootfs     bool              `json:"readonly_rootfs"`
	InitPath           string            `json:"initpath"` // dockerinit
	WorkingDir         string            `json:"working_dir"`
	ConfigPath         string            `json:"config_path"` // this should be able to be removed when the lxc template is moved into the driver
	Network            *Network          `json:"network"`
	Ipc                *Ipc              `json:"ipc"`
	Pid                *Pid              `json:"pid"`
	Resources          *Resources        `json:"resources"`
	Mounts             []Mount           `json:"mounts"`
	AllowedDevices     []*configs.Device `json:"allowed_devices"`
	AutoCreatedDevices []*configs.Device `json:"autocreated_devices"`
	CapAdd             []string          `json:"cap_add"`
	CapDrop            []string          `json:"cap_drop"`
	ContainerPid       int               `json:"container_pid"`  // the pid for the process inside a container
	ProcessConfig      ProcessConfig     `json:"process_config"` // Describes the init process of the container.
	ProcessLabel       string            `json:"process_label"`
	MountLabel         string            `json:"mount_label"`
	LxcConfig          []string          `json:"lxc_config"`
	AppArmorProfile    string            `json:"apparmor_profile"`
	CgroupParent       string            `json:"cgroup_parent"` // The parent cgroup for this command.
}

func InitContainer(c *Command) *configs.Config {
	container := template.New()

	container.Hostname = getEnv("HOSTNAME", c.ProcessConfig.Env)
	container.Cgroups.Name = c.ID
	container.Cgroups.AllowedDevices = c.AllowedDevices
	container.Readonlyfs = c.ReadonlyRootfs
	container.Devices = c.AutoCreatedDevices
	container.Rootfs = c.Rootfs
	container.Readonlyfs = c.ReadonlyRootfs

	// check to see if we are running in ramdisk to disable pivot root
	container.NoPivotRoot = os.Getenv("DOCKER_RAMDISK") != ""

	// Default parent cgroup is "docker". Override if required.
	if c.CgroupParent != "" {
		container.Cgroups.Parent = c.CgroupParent
	}
	return container
}

func getEnv(key string, env []string) string {
	for _, pair := range env {
		parts := strings.Split(pair, "=")
		if parts[0] == key {
			return parts[1]
		}
	}
	return ""
}

func SetupCgroups(container *configs.Config, c *Command) error {
	if c.Resources != nil {
		container.Cgroups.CpuShares = c.Resources.CpuShares
		container.Cgroups.Memory = c.Resources.Memory
		container.Cgroups.MemoryReservation = c.Resources.Memory
		container.Cgroups.MemorySwap = c.Resources.MemorySwap
		container.Cgroups.CpusetCpus = c.Resources.CpusetCpus
	}

	return nil
}

// Returns the network statistics for the network interfaces represented by the NetworkRuntimeInfo.
func getNetworkInterfaceStats(interfaceName string) (*libcontainer.NetworkInterface, error) {
	out := &libcontainer.NetworkInterface{Name: interfaceName}
	// This can happen if the network runtime information is missing - possible if the
	// container was created by an old version of libcontainer.
	if interfaceName == "" {
		return out, nil
	}
	type netStatsPair struct {
		// Where to write the output.
		Out *uint64
		// The network stats file to read.
		File string
	}
	// Ingress for host veth is from the container. Hence tx_bytes stat on the host veth is actually number of bytes received by the container.
	netStats := []netStatsPair{
		{Out: &out.RxBytes, File: "tx_bytes"},
		{Out: &out.RxPackets, File: "tx_packets"},
		{Out: &out.RxErrors, File: "tx_errors"},
		{Out: &out.RxDropped, File: "tx_dropped"},

		{Out: &out.TxBytes, File: "rx_bytes"},
		{Out: &out.TxPackets, File: "rx_packets"},
		{Out: &out.TxErrors, File: "rx_errors"},
		{Out: &out.TxDropped, File: "rx_dropped"},
	}
	for _, netStat := range netStats {
		data, err := readSysfsNetworkStats(interfaceName, netStat.File)
		if err != nil {
			return nil, err
		}
		*(netStat.Out) = data
	}
	return out, nil
}

// Reads the specified statistics available under /sys/class/net/<EthInterface>/statistics
func readSysfsNetworkStats(ethInterface, statsFile string) (uint64, error) {
	data, err := ioutil.ReadFile(filepath.Join("/sys/class/net", ethInterface, "statistics", statsFile))
	if err != nil {
		return 0, err
	}
	return strconv.ParseUint(strings.TrimSpace(string(data)), 10, 64)
}

func Stats(containerDir string, containerMemoryLimit int64, machineMemory int64) (*ResourceStats, error) {
	f, err := os.Open(filepath.Join(containerDir, "state.json"))
	if err != nil {
		return nil, err
	}
	defer f.Close()

	type network struct {
		Type              string
		HostInterfaceName string
	}

	state := struct {
		CgroupPaths map[string]string `json:"cgroup_paths"`
		Networks    []network
	}{}

	if err := json.NewDecoder(f).Decode(&state); err != nil {
		return nil, err
	}
	now := time.Now()

	mgr := fs.Manager{Paths: state.CgroupPaths}
	cstats, err := mgr.GetStats()
	if err != nil {
		return nil, err
	}
	stats := &libcontainer.Stats{CgroupStats: cstats}
	// if the container does not have any memory limit specified set the
	// limit to the machines memory
	memoryLimit := containerMemoryLimit
	if memoryLimit == 0 {
		memoryLimit = machineMemory
	}
	for _, iface := range state.Networks {
		switch iface.Type {
		case "veth":
			istats, err := getNetworkInterfaceStats(iface.HostInterfaceName)
			if err != nil {
				return nil, err
			}
			stats.Interfaces = append(stats.Interfaces, istats)
		}
	}
	return &ResourceStats{
		Stats:       stats,
		Read:        now,
		MemoryLimit: memoryLimit,
	}, nil
}
<commit_msg>Remove duplicate assignment<commit_after>package execdriver

import (
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/docker/docker/daemon/execdriver/native/template"
	"github.com/docker/docker/pkg/ulimit"
	"github.com/docker/libcontainer"
	"github.com/docker/libcontainer/cgroups/fs"
	"github.com/docker/libcontainer/configs"
)

// Context is a generic key value pair that allows
// arbatrary data to be sent
type Context map[string]string

var (
	ErrNotRunning              = errors.New("Container is not running")
	ErrWaitTimeoutReached      = errors.New("Wait timeout reached")
	ErrDriverAlreadyRegistered = errors.New("A driver already registered this docker init function")
	ErrDriverNotFound          = errors.New("The requested docker init has not been found")
)

type StartCallback func(*ProcessConfig, int)

// Driver specific information based on
// processes registered with the driver
type Info interface {
	IsRunning() bool
}

// Terminal in an interface for drivers to implement
// if they want to support Close and Resize calls from
// the core
type Terminal interface {
	io.Closer
	Resize(height, width int) error
}

type TtyTerminal interface {
	Master() libcontainer.Console
}

// ExitStatus provides exit reasons for a container.
type ExitStatus struct {
	// The exit code with which the container exited.
	ExitCode int

	// Whether the container encountered an OOM.
	OOMKilled bool
}

type Driver interface {
	Run(c *Command, pipes *Pipes, startCallback StartCallback) (ExitStatus, error) // Run executes the process and blocks until the process exits and returns the exit code
	// Exec executes the process in an existing container, blocks until the process exits and returns the exit code
	Exec(c *Command, processConfig *ProcessConfig, pipes *Pipes, startCallback StartCallback) (int, error)
	Kill(c *Command, sig int) error
	Pause(c *Command) error
	Unpause(c *Command) error
	Name() string                                 // Driver name
	Info(id string) Info                          // "temporary" hack (until we move state from core to plugins)
	GetPidsForContainer(id string) ([]int, error) // Returns a list of pids for the given container.
	Terminate(c *Command) error                   // kill it with fire
	Clean(id string) error                        // clean all traces of container exec
	Stats(id string) (*ResourceStats, error)      // Get resource stats for a running container
}

// Network settings of the container
type Network struct {
	Interface      *NetworkInterface `json:"interface"` // if interface is nil then networking is disabled
	Mtu            int               `json:"mtu"`
	ContainerID    string            `json:"container_id"` // id of the container to join network.
	HostNetworking bool              `json:"host_networking"`
}

// IPC settings of the container
type Ipc struct {
	ContainerID string `json:"container_id"` // id of the container to join ipc.
	HostIpc     bool   `json:"host_ipc"`
}

// PID settings of the container
type Pid struct {
	HostPid bool `json:"host_pid"`
}

type NetworkInterface struct {
	Gateway              string `json:"gateway"`
	IPAddress            string `json:"ip"`
	IPPrefixLen          int    `json:"ip_prefix_len"`
	MacAddress           string `json:"mac"`
	Bridge               string `json:"bridge"`
	GlobalIPv6Address    string `json:"global_ipv6"`
	LinkLocalIPv6Address string `json:"link_local_ipv6"`
	GlobalIPv6PrefixLen  int    `json:"global_ipv6_prefix_len"`
	IPv6Gateway          string `json:"ipv6_gateway"`
}

type Resources struct {
	Memory     int64            `json:"memory"`
	MemorySwap int64            `json:"memory_swap"`
	CpuShares  int64            `json:"cpu_shares"`
	CpusetCpus string           `json:"cpuset_cpus"`
	Rlimits    []*ulimit.Rlimit `json:"rlimits"`
}

type ResourceStats struct {
	*libcontainer.Stats
	Read        time.Time `json:"read"`
	MemoryLimit int64     `json:"memory_limit"`
	SystemUsage uint64    `json:"system_usage"`
}

type Mount struct {
	Source      string `json:"source"`
	Destination string `json:"destination"`
	Writable    bool   `json:"writable"`
	Private     bool   `json:"private"`
	Slave       bool   `json:"slave"`
}

// Describes a process that will be run inside a container.
type ProcessConfig struct {
	exec.Cmd `json:"-"`

	Privileged bool     `json:"privileged"`
	User       string   `json:"user"`
	Tty        bool     `json:"tty"`
	Entrypoint string   `json:"entrypoint"`
	Arguments  []string `json:"arguments"`
	Terminal   Terminal `json:"-"` // standard or tty terminal
	Console    string   `json:"-"` // dev/console path
}

// Process wrapps an os/exec.Cmd to add more metadata
type Command struct {
	ID                 string            `json:"id"`
	Rootfs             string            `json:"rootfs"` // root fs of the container
	ReadonlyRootfs     bool              `json:"readonly_rootfs"`
	InitPath           string            `json:"initpath"` // dockerinit
	WorkingDir         string            `json:"working_dir"`
	ConfigPath         string            `json:"config_path"` // this should be able to be removed when the lxc template is moved into the driver
	Network            *Network          `json:"network"`
	Ipc                *Ipc              `json:"ipc"`
	Pid                *Pid              `json:"pid"`
	Resources          *Resources        `json:"resources"`
	Mounts             []Mount           `json:"mounts"`
	AllowedDevices     []*configs.Device `json:"allowed_devices"`
	AutoCreatedDevices []*configs.Device `json:"autocreated_devices"`
	CapAdd             []string          `json:"cap_add"`
	CapDrop            []string          `json:"cap_drop"`
	ContainerPid       int               `json:"container_pid"`  // the pid for the process inside a container
	ProcessConfig      ProcessConfig     `json:"process_config"` // Describes the init process of the container.
	ProcessLabel       string            `json:"process_label"`
	MountLabel         string            `json:"mount_label"`
	LxcConfig          []string          `json:"lxc_config"`
	AppArmorProfile    string            `json:"apparmor_profile"`
	CgroupParent       string            `json:"cgroup_parent"` // The parent cgroup for this command.
}

func InitContainer(c *Command) *configs.Config {
	container := template.New()

	container.Hostname = getEnv("HOSTNAME", c.ProcessConfig.Env)
	container.Cgroups.Name = c.ID
	container.Cgroups.AllowedDevices = c.AllowedDevices
	container.Devices = c.AutoCreatedDevices
	container.Rootfs = c.Rootfs
	container.Readonlyfs = c.ReadonlyRootfs

	// check to see if we are running in ramdisk to disable pivot root
	container.NoPivotRoot = os.Getenv("DOCKER_RAMDISK") != ""

	// Default parent cgroup is "docker". Override if required.
	if c.CgroupParent != "" {
		container.Cgroups.Parent = c.CgroupParent
	}
	return container
}

func getEnv(key string, env []string) string {
	for _, pair := range env {
		parts := strings.Split(pair, "=")
		if parts[0] == key {
			return parts[1]
		}
	}
	return ""
}

func SetupCgroups(container *configs.Config, c *Command) error {
	if c.Resources != nil {
		container.Cgroups.CpuShares = c.Resources.CpuShares
		container.Cgroups.Memory = c.Resources.Memory
		container.Cgroups.MemoryReservation = c.Resources.Memory
		container.Cgroups.MemorySwap = c.Resources.MemorySwap
		container.Cgroups.CpusetCpus = c.Resources.CpusetCpus
	}

	return nil
}

// Returns the network statistics for the network interfaces represented by the NetworkRuntimeInfo.
func getNetworkInterfaceStats(interfaceName string) (*libcontainer.NetworkInterface, error) {
	out := &libcontainer.NetworkInterface{Name: interfaceName}
	// This can happen if the network runtime information is missing - possible if the
	// container was created by an old version of libcontainer.
	if interfaceName == "" {
		return out, nil
	}
	type netStatsPair struct {
		// Where to write the output.
		Out *uint64
		// The network stats file to read.
		File string
	}
	// Ingress for host veth is from the container. Hence tx_bytes stat on the host veth is actually number of bytes received by the container.
	netStats := []netStatsPair{
		{Out: &out.RxBytes, File: "tx_bytes"},
		{Out: &out.RxPackets, File: "tx_packets"},
		{Out: &out.RxErrors, File: "tx_errors"},
		{Out: &out.RxDropped, File: "tx_dropped"},

		{Out: &out.TxBytes, File: "rx_bytes"},
		{Out: &out.TxPackets, File: "rx_packets"},
		{Out: &out.TxErrors, File: "rx_errors"},
		{Out: &out.TxDropped, File: "rx_dropped"},
	}
	for _, netStat := range netStats {
		data, err := readSysfsNetworkStats(interfaceName, netStat.File)
		if err != nil {
			return nil, err
		}
		*(netStat.Out) = data
	}
	return out, nil
}

// Reads the specified statistics available under /sys/class/net/<EthInterface>/statistics
func readSysfsNetworkStats(ethInterface, statsFile string) (uint64, error) {
	data, err := ioutil.ReadFile(filepath.Join("/sys/class/net", ethInterface, "statistics", statsFile))
	if err != nil {
		return 0, err
	}
	return strconv.ParseUint(strings.TrimSpace(string(data)), 10, 64)
}

func Stats(containerDir string, containerMemoryLimit int64, machineMemory int64) (*ResourceStats, error) {
	f, err := os.Open(filepath.Join(containerDir, "state.json"))
	if err != nil {
		return nil, err
	}
	defer f.Close()

	type network struct {
		Type              string
		HostInterfaceName string
	}

	state := struct {
		CgroupPaths map[string]string `json:"cgroup_paths"`
		Networks    []network
	}{}

	if err := json.NewDecoder(f).Decode(&state); err != nil {
		return nil, err
	}
	now := time.Now()

	mgr := fs.Manager{Paths: state.CgroupPaths}
	cstats, err := mgr.GetStats()
	if err != nil {
		return nil, err
	}
	stats := &libcontainer.Stats{CgroupStats: cstats}
	// if the container does not have any memory limit specified set the
	// limit to the machines memory
	memoryLimit := containerMemoryLimit
	if memoryLimit == 0 {
		memoryLimit = machineMemory
	}
	for _, iface := range state.Networks {
		switch iface.Type {
		case "veth":
			istats, err := getNetworkInterfaceStats(iface.HostInterfaceName)
			if err != nil {
				return nil, err
			}
			stats.Interfaces = append(stats.Interfaces, istats)
		}
	}
	return &ResourceStats{
		Stats:       stats,
		Read:        now,
		MemoryLimit: memoryLimit,
	}, nil
}
<|endoftext|>


<commit_before>package json

import (
	"testing"

	v2net "github.com/v2ray/v2ray-core/common/net"
	v2nettesting "github.com/v2ray/v2ray-core/common/net/testing"
	"github.com/v2ray/v2ray-core/testing/unit"
)

func TestDomainMatching(t *testing.T) {
	assert := unit.Assert(t)

	rule := &FieldRule{
		Domain: "v2ray.com",
	}
	dest := v2net.NewTCPDestination(v2net.DomainAddress("www.v2ray.com", 80))
	assert.Bool(rule.Apply(dest)).IsTrue()
}

func TestPortMatching(t *testing.T) {
	assert := unit.Assert(t)

	rule := &FieldRule{
		Port: &v2nettesting.PortRange{
			FromValue: 0,
			ToValue:   100,
		},
	}
	dest := v2net.NewTCPDestination(v2net.DomainAddress("www.v2ray.com", 80))
	assert.Bool(rule.Apply(dest)).IsTrue()
}
<commit_msg>test case for ip rule<commit_after>package json

import (
	"testing"

	v2net "github.com/v2ray/v2ray-core/common/net"
	v2nettesting "github.com/v2ray/v2ray-core/common/net/testing"
	"github.com/v2ray/v2ray-core/testing/unit"
)

func TestDomainMatching(t *testing.T) {
	assert := unit.Assert(t)

	rule := &FieldRule{
		Domain: "v2ray.com",
	}
	dest := v2net.NewTCPDestination(v2net.DomainAddress("www.v2ray.com", 80))
	assert.Bool(rule.Apply(dest)).IsTrue()
}

func TestPortMatching(t *testing.T) {
	assert := unit.Assert(t)

	rule := &FieldRule{
		Port: &v2nettesting.PortRange{
			FromValue: 0,
			ToValue:   100,
		},
	}
	dest := v2net.NewTCPDestination(v2net.DomainAddress("www.v2ray.com", 80))
	assert.Bool(rule.Apply(dest)).IsTrue()
}

func TestIPMatching(t *testing.T) {
	assert := unit.Assert(t)

	rawJson := `{
    "type": "field",
    "ip": "10.0.0.0/8",
    "tag": "test"
  }`
	rule := parseRule([]byte(rawJson))
	dest := v2net.NewTCPDestination(v2net.IPAddress([]byte{10, 0, 0, 1}, 80))
	assert.Bool(rule.Apply(dest)).IsTrue()
}
<|endoftext|>


<commit_before>// Copyright 2014 Manu Martinez-Almeida.  All rights reserved.
// Use of this source code is governed by a MIT style
// license that can be found in the LICENSE file.

package gin

import (
	"fmt"
	"io"
	"time"
)

var (
	green   = string([]byte{27, 91, 57, 55, 59, 52, 50, 109})
	white   = string([]byte{27, 91, 57, 48, 59, 52, 55, 109})
	yellow  = string([]byte{27, 91, 57, 55, 59, 52, 51, 109})
	red     = string([]byte{27, 91, 57, 55, 59, 52, 49, 109})
	blue    = string([]byte{27, 91, 57, 55, 59, 52, 52, 109})
	magenta = string([]byte{27, 91, 57, 55, 59, 52, 53, 109})
	cyan    = string([]byte{27, 91, 57, 55, 59, 52, 54, 109})
	reset   = string([]byte{27, 91, 48, 109})
)

func ErrorLogger() HandlerFunc {
	return ErrorLoggerT(ErrorTypeAny)
}

func ErrorLoggerT(typ ErrorType) HandlerFunc {
	return func(c *Context) {
		c.Next()

		if !c.Writer.Written() {
			json := c.Errors.ByType(typ).JSON()
			if json != nil {
				c.JSON(-1, json)
			}
		}
	}
}

// Instances a Logger middleware that will write the logs to gin.DefaultWriter
// By default gin.DefaultWriter = os.Stdout
func Logger() HandlerFunc {
	return LoggerWithWriter(DefaultWriter)
}

// Instance a Logger middleware with the specified writter buffer.
// Example: os.Stdout, a file opened in write mode, a socket...
func LoggerWithWriter(out io.Writer) HandlerFunc {
	return func(c *Context) {
		// Start timer
		start := time.Now()
		path := c.Request.URL.Path

		// Process request
		c.Next()

		// Stop timer
		end := time.Now()
		latency := end.Sub(start)

		clientIP := c.ClientIP()
		method := c.Request.Method
		statusCode := c.Writer.Status()
		statusColor := colorForStatus(statusCode)
		methodColor := colorForMethod(method)
		comment := c.Errors.ByType(ErrorTypePrivate).String()

		fmt.Fprintf(out, "[GIN] %v |%s %3d %s| %13v | %s |%s  %s %-7s %s\n%s",
			end.Format("2006/01/02 - 15:04:05"),
			statusColor, statusCode, reset,
			latency,
			clientIP,
			methodColor, reset, method,
			path,
			comment,
		)
	}
}

func colorForStatus(code int) string {
	switch {
	case code >= 200 && code < 300:
		return green
	case code >= 300 && code < 400:
		return white
	case code >= 400 && code < 500:
		return yellow
	default:
		return red
	}
}

func colorForMethod(method string) string {
	switch method {
	case "GET":
		return blue
	case "POST":
		return cyan
	case "PUT":
		return yellow
	case "DELETE":
		return red
	case "PATCH":
		return green
	case "HEAD":
		return magenta
	case "OPTIONS":
		return white
	default:
		return reset
	}
}
<commit_msg>We do not longer have to call JSON() explicitly<commit_after>// Copyright 2014 Manu Martinez-Almeida.  All rights reserved.
// Use of this source code is governed by a MIT style
// license that can be found in the LICENSE file.

package gin

import (
	"fmt"
	"io"
	"time"
)

var (
	green   = string([]byte{27, 91, 57, 55, 59, 52, 50, 109})
	white   = string([]byte{27, 91, 57, 48, 59, 52, 55, 109})
	yellow  = string([]byte{27, 91, 57, 55, 59, 52, 51, 109})
	red     = string([]byte{27, 91, 57, 55, 59, 52, 49, 109})
	blue    = string([]byte{27, 91, 57, 55, 59, 52, 52, 109})
	magenta = string([]byte{27, 91, 57, 55, 59, 52, 53, 109})
	cyan    = string([]byte{27, 91, 57, 55, 59, 52, 54, 109})
	reset   = string([]byte{27, 91, 48, 109})
)

func ErrorLogger() HandlerFunc {
	return ErrorLoggerT(ErrorTypeAny)
}

func ErrorLoggerT(typ ErrorType) HandlerFunc {
	return func(c *Context) {
		c.Next()
		// avoid writting if we already wrote into the response body
		if !c.Writer.Written() {
			errors := c.Errors.ByType(typ)
			if len(errors) > 0 {
				c.JSON(-1, errors)
			}
		}
	}
}

// Instances a Logger middleware that will write the logs to gin.DefaultWriter
// By default gin.DefaultWriter = os.Stdout
func Logger() HandlerFunc {
	return LoggerWithWriter(DefaultWriter)
}

// Instance a Logger middleware with the specified writter buffer.
// Example: os.Stdout, a file opened in write mode, a socket...
func LoggerWithWriter(out io.Writer) HandlerFunc {
	return func(c *Context) {
		// Start timer
		start := time.Now()
		path := c.Request.URL.Path

		// Process request
		c.Next()

		// Stop timer
		end := time.Now()
		latency := end.Sub(start)

		clientIP := c.ClientIP()
		method := c.Request.Method
		statusCode := c.Writer.Status()
		statusColor := colorForStatus(statusCode)
		methodColor := colorForMethod(method)
		comment := c.Errors.ByType(ErrorTypePrivate).String()

		fmt.Fprintf(out, "[GIN] %v |%s %3d %s| %13v | %s |%s  %s %-7s %s\n%s",
			end.Format("2006/01/02 - 15:04:05"),
			statusColor, statusCode, reset,
			latency,
			clientIP,
			methodColor, reset, method,
			path,
			comment,
		)
	}
}

func colorForStatus(code int) string {
	switch {
	case code >= 200 && code < 300:
		return green
	case code >= 300 && code < 400:
		return white
	case code >= 400 && code < 500:
		return yellow
	default:
		return red
	}
}

func colorForMethod(method string) string {
	switch method {
	case "GET":
		return blue
	case "POST":
		return cyan
	case "PUT":
		return yellow
	case "DELETE":
		return red
	case "PATCH":
		return green
	case "HEAD":
		return magenta
	case "OPTIONS":
		return white
	default:
		return reset
	}
}
<|endoftext|>


<commit_before>package logger

import (
	"fmt"
	"log"

	"github.com/mgutz/ansi"
)

const (
	Silent = iota
	Error  = iota
	Warn   = iota
	Listen = iota
	Init   = iota
	Info   = iota
	Debug  = iota
	max    = iota
)

type Logger struct{ Level int }

var prefixes = [...]string{
	"",
	ansi.Color("[***error***]", "black+B:red"),
	ansi.Color("[**warning**]", "red:yellow+h"),
	ansi.Color("[ listening ]", "black:cyan+h"),
	/*      */ "[initialized]",
	/*      */ "[information]",
	/*      */ "[   debug   ]"}

func (logger *Logger) Log(level int, message string) {
	if logger.Level == Silent {
		return
	}
	// If an invalid Level has been set, output an error and the message as info.
	if logger.Level > Debug {
		log.Printf("%s logger.Level must be < %d", prefixes[Error], max)
		log.Printf("%s %s", prefixes[Info], message)
		return
	}
	if 0 <= level && level <= logger.Level {
		log.Printf("%s %s", prefixes[level], message)
	}
}

func New(level int) (*Logger, error) {
	if level > Debug {
		return nil, fmt.Errorf("logger.Level must be < %d", max)
	}
	return &Logger{Level: level}, nil
}
<commit_msg>update logger<commit_after>package logger

import (
	"fmt"
	"log"

	"github.com/mgutz/ansi"
)

const (
	Silent  = iota
	Error   = iota
	Blocked = iota
	Warn    = iota
	Listen  = iota
	Install = iota
	Init    = iota
	Request = iota
	Info    = iota
	Debug   = iota
	max     = iota
)

type Logger struct{ level int }

var LogLevel = map[string]int{
	"silent":  Silent,
	"error":   Error,
	"blocked": Blocked,
	"warn":    Warn,
	"listen":  Listen,
	"install": Install,
	"init":    Init,
	"request": Request,
	"info":    Info,
	"debug":   Debug}

var prefixes = [...]string{
	"",
	ansi.Color("[***error***]", "black+B:red"),
	ansi.Color("[**blocked**]", "255+Bb:165"),
	ansi.Color("[**warning**]", "red:yellow+h"),
	ansi.Color("[ listening ]", "black:cyan+h"),
	/*      */ "[ installed ]",
	/*      */ "[initialized]",
	/*      */ "[  request  ]",
	/*      */ "[information]",
	/*      */ "[   debug   ]"}

func out(loggerLevel, messageLevel int, format string, v ...interface{}) {
	if loggerLevel == Silent || messageLevel == Silent {
		return
	}
	message := fmt.Sprintf(format, v...)
	if messageLevel > Silent && messageLevel <= loggerLevel {
		log.Printf("%s %s", prefixes[messageLevel], message)
	}
}

func (logger *Logger) Level() int { return logger.level }

func (logger *Logger) SetLevel(level int) error {
	if level > Debug || level < Silent {
		return fmt.Errorf("logger level must be >= %d and < %d", Silent, max)
	}
	logger.level = level
	return nil
}

func (logger *Logger) Error(format string, v ...interface{}) {
	out(logger.level, Error, format, v...)
}

func (logger *Logger) Blocked(format string, v ...interface{}) {
	out(logger.level, Blocked, format, v...)
}

func (logger *Logger) Warn(format string, v ...interface{}) {
	out(logger.level, Warn, format, v...)
}

func (logger *Logger) Listen(format string, v ...interface{}) {
	out(logger.level, Listen, format, v...)
}

func (logger *Logger) Install(format string, v ...interface{}) {
	out(logger.level, Install, format, v...)
}

func (logger *Logger) Init(format string, v ...interface{}) {
	out(logger.level, Init, format, v...)
}

func (logger *Logger) Request(format string, v ...interface{}) {
	out(logger.level, Request, format, v...)
}

func (logger *Logger) Info(format string, v ...interface{}) {
	out(logger.level, Info, format, v...)
}

func (logger *Logger) Debug(format string, v ...interface{}) {
	out(logger.level, Debug, format, v...)
}

func MustError(format string, v ...interface{}) {
	out(max, Error, format, v...)
}

func MustBlocked(format string, v ...interface{}) {
	out(max, Blocked, format, v...)
}

func MustWarn(format string, v ...interface{}) {
	out(max, Warn, format, v...)
}

func MustListen(format string, v ...interface{}) {
	out(max, Listen, format, v...)
}

func MustInstall(format string, v ...interface{}) {
	out(max, Install, format, v...)
}

func MustInit(format string, v ...interface{}) {
	out(max, Init, format, v...)
}

func MustRequest(format string, v ...interface{}) {
	out(max, Request, format, v...)
}

func MustInfo(format string, v ...interface{}) {
	out(max, Info, format, v...)
}

func MustDebug(format string, v ...interface{}) {
	out(max, Debug, format, v...)
}

func New(level int) (*Logger, error) {
	logger := new(Logger)
	if err := logger.SetLevel(level); err != nil {
		return nil, err
	} else {
		return logger, nil
	}
}
<|endoftext|>


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

package subprocess

import (
	"github.com/zenoss/glog"

	"errors"
	"os"
	"os/exec"
	"sync"
	"syscall"
	"time"
)

// Instance manages a subprocess instance.
type Instance struct {
	command        string
	args           []string
	env            []string
	commandExit    chan error // used to send command exit values to the parent controller
	closing        chan chan error
	closeLock      sync.Mutex     // mutex to synchronize Close() calls
	sigtermTimeout time.Duration  // sigterm timeout
	signalChan     chan os.Signal // used to send signals to the command process
}

// New creates a subprocess.Instance
func New(sigtermTimeout time.Duration, env []string, command string, args ...string) (*Instance, chan error, error) {
	s := &Instance{
		command:        command,
		args:           args,
		env:            env,
		commandExit:    make(chan error, 1),
		sigtermTimeout: sigtermTimeout,
		signalChan:     make(chan os.Signal, 1),
	}
	go s.loop()
	return s, s.commandExit, nil
}

// Notify sends the sig to the subprocess instance.
func (s *Instance) Notify(sig os.Signal) {
	select {
	case s.signalChan <- sig:
		glog.V(1).Infof("Notify: sending signal %v", sig)
	default:
		glog.Warningf("Notify: unable to send signal %v, because signalChan is full", sig)
	}
}

// Close signals the subprocess to shutdown via sigterm. If sigterm fails to shutdown
// withing the s.timeout, a sigkill is issued.
func (s *Instance) Close() error {
	if s == nil {
		return nil
	}
	s.closeLock.Lock()
	defer s.closeLock.Unlock()
	if s.closing == nil {
		return errors.New("subprocess instance already closed")
	}
	errc := make(chan error)
	s.closing <- errc
	err := <-errc
	s.closing = nil
	return err
}

func (s *Instance) loop() {

	setUpCmd := func(exitChan chan error) *exec.Cmd {
		glog.Infof("about to execute: %s , %v[%d]", s.command, s.args, len(s.args))
		cmd := exec.Command(s.command, s.args...)
		cmd.Env = s.env
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.Stdin = os.Stdin
		go func() {
			exitChan <- cmd.Run()
		}()
		return cmd
	}

	processExit := make(chan error, 1) //lets us know when the process exits
	cmd := setUpCmd(processExit)
	var returnChan chan error
	sigterm := make(chan error)
	sigkill := make(<-chan time.Time)

	closing := s.closing
	for {

		select {
		case s := <-s.signalChan:
			cmd.Process.Signal(s)
			glog.V(1).Infof("loop: signal sent%v", s)

		case exitError := <-processExit:
			glog.V(1).Infof("loop: process exited with error %v", exitError)
			select {
			case s.commandExit <- exitError: // tell our the parent controller that the command has exited
			default:
				glog.Warningf("Received child exit = %#v but the commandExit channel is full", exitError)
			}
			return

		case returnChan = <-closing:
			cmd.Process.Signal(syscall.SIGQUIT)
			// setup hard killing the task
			sigkill = time.After(s.sigtermTimeout)
			closing = nil // avoid this case again
			go func(rchan chan error) {
				rchan <- cmd.Wait()
			}(sigterm)

		case err := <-sigterm:
			returnChan <- err
			return

		case <-sigkill:
			cmd.Process.Signal(syscall.SIGKILL)
			// our previous goroutine should take care of doing the right thing
			returnChan <- errors.New("subprocess instance sigkilled")
			return
		}

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

package subprocess

import (
	"github.com/zenoss/glog"

	"errors"
	"os"
	"os/exec"
	"sync"
	"syscall"
	"time"
)

// Instance manages a subprocess instance.
type Instance struct {
	command        string
	args           []string
	env            []string
	commandExit    chan error // used to send command exit values to the parent controller
	closing        chan chan error
	closeLock      sync.Mutex     // mutex to synchronize Close() calls
	sigtermTimeout time.Duration  // sigterm timeout
	signalChan     chan os.Signal // used to send signals to the command process
}

// New creates a subprocess.Instance
func New(sigtermTimeout time.Duration, env []string, command string, args ...string) (*Instance, chan error, error) {
	s := &Instance{
		command:        command,
		args:           args,
		env:            env,
		commandExit:    make(chan error, 1),
		sigtermTimeout: sigtermTimeout,
		signalChan:     make(chan os.Signal, 1),
	}
	go s.loop()
	return s, s.commandExit, nil
}

// Notify sends the sig to the subprocess instance.
func (s *Instance) Notify(sig os.Signal) {
	select {
	case s.signalChan <- sig:
		glog.V(1).Infof("Notify: sending signal %v", sig)
	}
}

// Close signals the subprocess to shutdown via sigterm. If sigterm fails to shutdown
// withing the s.timeout, a sigkill is issued.
func (s *Instance) Close() error {
	if s == nil {
		return nil
	}
	s.closeLock.Lock()
	defer s.closeLock.Unlock()
	if s.closing == nil {
		return errors.New("subprocess instance already closed")
	}
	errc := make(chan error)
	s.closing <- errc
	err := <-errc
	s.closing = nil
	return err
}

func (s *Instance) loop() {

	setUpCmd := func(exitChan chan error) *exec.Cmd {
		glog.Infof("about to execute: %s , %v[%d]", s.command, s.args, len(s.args))
		cmd := exec.Command(s.command, s.args...)
		cmd.Env = s.env
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.Stdin = os.Stdin
		go func() {
			exitChan <- cmd.Run()
		}()
		return cmd
	}

	processExit := make(chan error, 1) //lets us know when the process exits
	cmd := setUpCmd(processExit)
	var returnChan chan error
	sigterm := make(chan error)
	sigkill := make(<-chan time.Time)

	closing := s.closing
	for {

		select {
		case s := <-s.signalChan:
			cmd.Process.Signal(s)
			glog.V(1).Infof("loop: signal sent%v", s)

		case exitError := <-processExit:
			glog.V(1).Infof("loop: process exited with error %v", exitError)
			select {
			case s.commandExit <- exitError: // tell our the parent controller that the command has exited
			}
			return

		case returnChan = <-closing:
			cmd.Process.Signal(syscall.SIGQUIT)
			// setup hard killing the task
			sigkill = time.After(s.sigtermTimeout)
			closing = nil // avoid this case again
			go func(rchan chan error) {
				rchan <- cmd.Wait()
			}(sigterm)

		case err := <-sigterm:
			returnChan <- err
			return

		case <-sigkill:
			cmd.Process.Signal(syscall.SIGKILL)
			// our previous goroutine should take care of doing the right thing
			returnChan <- errors.New("subprocess instance sigkilled")
			return
		}

	}
}
<|endoftext|>


<commit_before>package termite

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"net/http/pprof"
	"os"

	"github.com/hanwen/termite/stats"
)

func serveLog(worker *Worker, w http.ResponseWriter, req *http.Request) {
	sz := int64(500 * 1024)
	sizeStr, ok := req.URL.Query()["size"]
	if ok {
		fmt.Scanf(sizeStr[0], "%d", &sz)
	}

	logReq := LogRequest{Whence: os.SEEK_END, Off: -sz, Size: sz}
	logRep := LogResponse{}
	err := worker.Log(&logReq, &logRep)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "<title>Termite worker error</title>")
		fmt.Fprintf(w, "Error: %s", err.Error())
		return
	}

	w.Header().Set("Content-Type", "text/plain")
	w.Write(logRep.Data)
}

func serveStatus(worker *Worker, w http.ResponseWriter, r *http.Request) {
	statusReq := WorkerStatusRequest{}
	status := WorkerStatusResponse{}

	err := worker.Status(&statusReq, &status)
	if err != nil {
		// TODO - escaping.
		fmt.Fprintf(w, "
Status error: %v", err)
	}

	addr := fmt.Sprintf("%s:%d", cname, worker.listener.Addr().(*net.TCPAddr).Port)
	fmt.Fprintf(w, "
Worker %s
Version %s
Jobs %d\n",
		addr, status.Version, status.MaxJobCount)
	fmt.Fprintf(w, "
Worker log %s\n", addr, addr)

	if !status.Accepting {
		fmt.Fprintf(w, "shutting down")
	}
	stats.CpuStatsWriteHttp(w, status.CpuStats, status.DiskStats)

	fmt.Fprintf(w, "
Total CPU: %s", status.TotalCpu.Percent())
	fmt.Fprintf(w, "
Content cache hit rate: %.0f %%, Age %d",
		100.0*status.ContentCacheHitRate,
		status.ContentCacheHitAge)

	m := status.MemStat
	fmt.Fprintf(w, "
HeapIdle: %v, HeapInUse: %v",
		m.HeapIdle, m.HeapInuse)

	stats.CountStatsWriteHttp(w, status.PhaseNames, status.PhaseCounts)

	for _, mirrorStatus := range status.MirrorStatus {
		mirrorStatusHtml(w, mirrorStatus)
	}
}

func mirrorStatusHtml(w http.ResponseWriter, s MirrorStatusResponse) {
	fmt.Fprintf(w, "
Mirror %s
", s.Root)
	for _, s := range s.RpcTimings {
		fmt.Fprintf(w, "
%s", s)
	}

	running := 0
	for _, fs := range s.Fses {
		running += len(fs.Tasks)
	}

	fmt.Fprintf(w, "
%d maximum jobs, %d running, %d waiting tasks, %d unused filesystems.\n",
		s.Granted, running, s.WaitingTasks, s.IdleFses)
	if !s.Accepting {
		fmt.Fprintf(w, "
shutting down\n")
	}

	fmt.Fprintf(w, "
\n")
	for _, v := range s.Fses {
		fmt.Fprintf(w, "
id %s: %s
\n", v.Id, v.Mem)
		for _, t := range v.Tasks {
			fmt.Fprintf(w, "
%s\n", t)
		}
		fmt.Fprintf(w, "
\n")
	}
	fmt.Fprintf(w, "
\n")
}

func (w *Worker) serveStatus() {
	l, err := net.Listen("tcp", ":0")
	if err != nil {
		log.Println("status serve:", err)
		return
	}

	w.httpStatusPort = l.Addr().(*net.TCPAddr).Port
	log.Printf("Serving status on port %d", w.httpStatusPort)
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(wr http.ResponseWriter, r *http.Request) {
		serveStatus(w, wr, r)
	})
	mux.HandleFunc("/log", func(wr http.ResponseWriter, r *http.Request) {
		serveLog(w, wr, r)
	})

	mux.Handle("/debug/pprof/", http.HandlerFunc(pprof.Index))
	mux.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
    	mux.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
	mux.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))

	err = http.Serve(l, mux)
	if err != nil {
		log.Println("status serve:", err)
		return
	}
}
<commit_msg>Add direct status link to workerstatus page.<commit_after>package termite

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"net/http/pprof"
	"os"

	"github.com/hanwen/termite/stats"
)

func serveLog(worker *Worker, w http.ResponseWriter, req *http.Request) {
	sz := int64(500 * 1024)
	sizeStr, ok := req.URL.Query()["size"]
	if ok {
		fmt.Scanf(sizeStr[0], "%d", &sz)
	}

	logReq := LogRequest{Whence: os.SEEK_END, Off: -sz, Size: sz}
	logRep := LogResponse{}
	err := worker.Log(&logReq, &logRep)
	if err != nil {
		w.WriteHeader(http.StatusBadRequest)
		fmt.Fprintf(w, "<title>Termite worker error</title>")
		fmt.Fprintf(w, "Error: %s", err.Error())
		return
	}

	w.Header().Set("Content-Type", "text/plain")
	w.Write(logRep.Data)
}

func serveStatus(worker *Worker, w http.ResponseWriter, r *http.Request) {
	statusReq := WorkerStatusRequest{}
	status := WorkerStatusResponse{}

	err := worker.Status(&statusReq, &status)
	if err != nil {
		// TODO - escaping.
		fmt.Fprintf(w, "
Status error: %v", err)
	}

	addr := fmt.Sprintf("%s:%d", cname, worker.listener.Addr().(*net.TCPAddr).Port)
	fmt.Fprintf(w, "
Worker %s (status)
Version %s
Jobs %d\n",
		addr, cname, worker.httpStatusPort, status.Version, status.MaxJobCount)
	fmt.Fprintf(w, "
Worker log %s\n", addr, addr)

	if !status.Accepting {
		fmt.Fprintf(w, "shutting down")
	}
	stats.CpuStatsWriteHttp(w, status.CpuStats, status.DiskStats)

	fmt.Fprintf(w, "
Total CPU: %s", status.TotalCpu.Percent())
	fmt.Fprintf(w, "
Content cache hit rate: %.0f %%, Age %d",
		100.0*status.ContentCacheHitRate,
		status.ContentCacheHitAge)

	m := status.MemStat
	fmt.Fprintf(w, "
HeapIdle: %v, HeapInUse: %v",
		m.HeapIdle, m.HeapInuse)

	stats.CountStatsWriteHttp(w, status.PhaseNames, status.PhaseCounts)

	for _, mirrorStatus := range status.MirrorStatus {
		mirrorStatusHtml(w, mirrorStatus)
	}
}

func mirrorStatusHtml(w http.ResponseWriter, s MirrorStatusResponse) {
	fmt.Fprintf(w, "
Mirror %s
", s.Root)
	for _, s := range s.RpcTimings {
		fmt.Fprintf(w, "
%s", s)
	}

	running := 0
	for _, fs := range s.Fses {
		running += len(fs.Tasks)
	}

	fmt.Fprintf(w, "
%d maximum jobs, %d running, %d waiting tasks, %d unused filesystems.\n",
		s.Granted, running, s.WaitingTasks, s.IdleFses)
	if !s.Accepting {
		fmt.Fprintf(w, "
shutting down\n")
	}

	fmt.Fprintf(w, "
\n")
	for _, v := range s.Fses {
		fmt.Fprintf(w, "
id %s: %s
\n", v.Id, v.Mem)
		for _, t := range v.Tasks {
			fmt.Fprintf(w, "
%s\n", t)
		}
		fmt.Fprintf(w, "
\n")
	}
	fmt.Fprintf(w, "
\n")
}

func (w *Worker) serveStatus() {
	l, err := net.Listen("tcp", ":0")
	if err != nil {
		log.Println("status serve:", err)
		return
	}

	w.httpStatusPort = l.Addr().(*net.TCPAddr).Port
	log.Printf("Serving status on port %d", w.httpStatusPort)
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(wr http.ResponseWriter, r *http.Request) {
		serveStatus(w, wr, r)
	})
	mux.HandleFunc("/log", func(wr http.ResponseWriter, r *http.Request) {
		serveLog(w, wr, r)
	})

	mux.Handle("/debug/pprof/", http.HandlerFunc(pprof.Index))
	mux.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
    	mux.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
	mux.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))

	err = http.Serve(l, mux)
	if err != nil {
		log.Println("status serve:", err)
		return
	}
}
<|endoftext|>


<commit_before>package errors

// Auth 权限错误
type Auth struct {
	CodeError
}

func (e Auth) Error() string {
	return e.CodeError.Error()
}

// NewAuth 新建
func NewAuth(code int, msgs ...interface{}) error {
	return Auth{NewCode(code, msgs...)}
}
<commit_msg>修改auth，增加forbidden<commit_after>package errors

// Unauthorized 权限错误
type Unauthorized struct {
	CodeError
}

func (e Unauthorized) Error() string {
	return e.CodeError.Error()
}

// NewUnauthorized 新建
func NewUnauthorized(code int, msgs ...interface{}) error {
	return Unauthorized{NewCode(code, msgs...)}
}

// Forbidden 权限错误
type Forbidden struct {
	CodeError
}

func (e Forbidden) Error() string {
	return e.CodeError.Error()
}

// NewForbidden 新建
func NewForbidden(code int, msgs ...interface{}) error {
	return Forbidden{NewCode(code, msgs...)}
}
<|endoftext|>


<commit_before>package memberlist

import (
	"os"
	"testing"
	"time"
)

// CheckInteg will skip a test if integration testing is not enabled.
func CheckInteg(t *testing.T) {
	if !IsInteg() {
		t.SkipNow()
	}
}

// IsInteg returns a boolean telling you if we're in integ testing mode.
func IsInteg() bool {
	return os.Getenv("INTEG_TESTS") != ""
}

// Tests the memberlist by creating a cluster of 100 nodes
// and checking that we get strong convergence of changes.
func TestMemberlist_Integ(t *testing.T) {
	CheckInteg(t)

	num := 16
	var members []*Memberlist

	ch := make(chan *Node, num)
	leaveCh := make(chan *Node, num)

	for i := 0; i < num; i++ {
		addr, _ := GetBindAddr()
		c := DefaultConfig()
		c.Name = addr
		c.BindAddr = addr
		c.RTT = 100 * time.Microsecond
		c.ProbeInterval = 5 * time.Millisecond
		c.GossipInterval = 3 * time.Millisecond
		c.PushPullInterval = 20 * time.Millisecond

		if i == 0 {
			m, err := Create(c)
			if err != nil {
				t.Fatalf("unexpected err: %s", err)
			}
			members = append(members, m)
			defer m.Shutdown()
			m.config.JoinCh = ch
			m.config.LeaveCh = leaveCh
		} else {
			last := members[i-1]
			m, err := Join(c, []string{last.config.Name})
			if err != nil {
				t.Fatalf("unexpected err: %s", err)
			}
			members = append(members, m)
			defer m.Shutdown()
		}
	}

	// Wait for node 1 to see all the others
	for i := 1; i < num; i++ {
		select {
		case <-ch:
		case l := <-leaveCh:
			t.Fatalf("unexpected leave %s", l)
		case <-time.After(10 * time.Millisecond):
			t.Fatalf("timeout")
		}
	}
	time.Sleep(10 * time.Millisecond)

	for idx, m := range members {
		if len(m.Members()) != num {
			t.Fatalf("bad num %d at idx %d", len(m.Members()), idx)
		}
	}
}
<commit_msg>Scale up the integ test to 32 nodes<commit_after>package memberlist

import (
	"os"
	"testing"
	"time"
)

// CheckInteg will skip a test if integration testing is not enabled.
func CheckInteg(t *testing.T) {
	if !IsInteg() {
		t.SkipNow()
	}
}

// IsInteg returns a boolean telling you if we're in integ testing mode.
func IsInteg() bool {
	return os.Getenv("INTEG_TESTS") != ""
}

// Tests the memberlist by creating a cluster of 100 nodes
// and checking that we get strong convergence of changes.
func TestMemberlist_Integ(t *testing.T) {
	CheckInteg(t)

	num := 32
	var members []*Memberlist

	joinCh := make(chan *Node, num)
	leaveCh := make(chan *Node, num)

	for i := 0; i < num; i++ {
		addr, _ := GetBindAddr()
		c := DefaultConfig()
		c.Name = addr
		c.BindAddr = addr
		c.RTT = 200 * time.Microsecond
		c.ProbeInterval = 5 * time.Millisecond
		c.GossipInterval = 5 * time.Millisecond
		c.PushPullInterval = 100 * time.Millisecond

		if i == 0 {
			m, err := Create(c)
			if err != nil {
				t.Fatalf("unexpected err: %s", err)
			}
			members = append(members, m)
			defer m.Shutdown()
			m.config.JoinCh = joinCh
			m.config.LeaveCh = leaveCh
		} else {
			last := members[i-1]
			m, err := Join(c, []string{last.config.Name})
			if err != nil {
				t.Fatalf("unexpected err: %s", err)
			}
			members = append(members, m)
			defer m.Shutdown()
		}
	}

	// Wait for node 1 to see all the others
	time.Sleep(250 * time.Millisecond)

	for idx, m := range members {
		if m.NumMembers() != num {
			t.Fatalf("bad num %d at idx %d", len(m.Members()), idx)
		}
	}
}
<|endoftext|>


<commit_before>package raft

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"testing"
	"time"
)

// CheckInteg will skip a test if integration testing is not enabled.
func CheckInteg(t *testing.T) {
	if !IsInteg() {
		t.SkipNow()
	}
}

// IsInteg returns a boolean telling you if we're in integ testing mode.
func IsInteg() bool {
	return os.Getenv("INTEG_TESTS") != ""
}

type RaftEnv struct {
	dir      string
	conf     *Config
	fsm      *MockFSM
	store    *InmemStore
	snapshot *FileSnapshotStore
	peers    *JSONPeers
	trans    *NetworkTransport
	raft     *Raft
}

func (r *RaftEnv) Release() {
	log.Printf("[WARN] Release node at %v", r.raft.localAddr)
	f := r.raft.Shutdown()
	if err := f.Error(); err != nil {
		panic(err)
	}
	r.trans.Close()
	os.RemoveAll(r.dir)
}

func MakeRaft(t *testing.T, conf *Config) *RaftEnv {
	env := &RaftEnv{}

	// Set the config
	if conf == nil {
		conf = inmemConfig(t)
	}
	env.conf = conf

	dir, err := ioutil.TempDir("", "raft")
	if err != nil {
		t.Fatalf("err: %v ", err)
	}
	env.dir = dir

	stable := NewInmemStore()
	env.store = stable

	snap, err := NewFileSnapshotStore(dir, 3, nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	env.snapshot = snap

	env.fsm = &MockFSM{}

	trans, err := NewTCPTransport("127.0.0.1:0", nil, 2, time.Second, nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	env.trans = trans

	env.peers = NewJSONPeers(dir, trans)

	log.Printf("[INFO] Starting node at %v", trans.LocalAddr())
	raft, err := NewRaft(conf, env.fsm, stable, stable, snap, env.peers, trans)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	env.raft = raft
	return env
}

func WaitFor(env *RaftEnv, state RaftState) error {
	limit := time.Now().Add(200 * time.Millisecond)
	for env.raft.State() != state {
		if time.Now().Before(limit) {
			time.Sleep(10 * time.Millisecond)
		} else {
			return fmt.Errorf("failed to transition to state %v", state)
		}
	}
	return nil
}

func WaitForAny(state RaftState, envs []*RaftEnv) (*RaftEnv, error) {
	limit := time.Now().Add(200 * time.Millisecond)
CHECK:
	for _, env := range envs {
		if env.raft.State() == state {
			return env, nil
		}
	}
	if time.Now().Before(limit) {
		goto WAIT
	}
	return nil, fmt.Errorf("failed to find node in %v state", state)
WAIT:
	time.Sleep(10 * time.Millisecond)
	goto CHECK
}

func WaitFuture(f Future, t *testing.T) error {
	timer := time.AfterFunc(200*time.Millisecond, func() {
		panic(fmt.Errorf("timeout waiting for future %v", f))
	})
	defer timer.Stop()
	return f.Error()
}

func NoErr(err error, t *testing.T) {
	if err != nil {
		t.Fatalf("err: %v", err)
	}
}

func CheckConsistent(envs []*RaftEnv, t *testing.T) {
	limit := time.Now().Add(400 * time.Millisecond)
	first := envs[0]
	var err error
CHECK:
	l1 := len(first.fsm.logs)
	for i := 1; i < len(envs); i++ {
		env := envs[i]
		l2 := len(env.fsm.logs)
		if l1 != l2 {
			err = fmt.Errorf("log length mismatch %d %d", l1, l2)
			goto ERR
		}
		for idx, log := range first.fsm.logs {
			other := env.fsm.logs[idx]
			if bytes.Compare(log, other) != 0 {
				err = fmt.Errorf("log %d mismatch %v %v", idx, log, other)
				goto ERR
			}
		}
	}
	return
ERR:
	if time.Now().After(limit) {
		t.Fatalf("%v", err)
	}
	time.Sleep(20 * time.Millisecond)
	goto CHECK
}

// Tests Raft by creating a cluster, growing it to 5 nodes while
// causing various stressful conditions
func TestRaft_Integ(t *testing.T) {
	CheckInteg(t)
	conf := DefaultConfig()
	conf.HeartbeatTimeout = 50 * time.Millisecond
	conf.ElectionTimeout = 50 * time.Millisecond
	conf.LeaderLeaseTimeout = 50 * time.Millisecond
	conf.CommitTimeout = 5 * time.Millisecond
	conf.SnapshotThreshold = 100
	conf.TrailingLogs = 10
	conf.EnableSingleNode = true

	// Create a single node
	env1 := MakeRaft(t, conf)
	NoErr(WaitFor(env1, Leader), t)

	// Do some commits
	var futures []Future
	for i := 0; i < 100; i++ {
		futures = append(futures, env1.raft.Apply([]byte(fmt.Sprintf("test%d", i)), 0))
	}
	for _, f := range futures {
		NoErr(WaitFuture(f, t), t)
		log.Printf("[DEBUG] Applied %v", f)
	}

	// Do a snapshot
	NoErr(WaitFuture(env1.raft.Snapshot(), t), t)

	// Join a few nodes!
	var envs []*RaftEnv
	for i := 0; i < 4; i++ {
		env := MakeRaft(t, conf)
		addr := env.trans.LocalAddr()
		NoErr(WaitFuture(env1.raft.AddPeer(addr), t), t)
		envs = append(envs, env)
	}

	// Wait for a leader
	leader, err := WaitForAny(Leader, append([]*RaftEnv{env1}, envs...))
	NoErr(err, t)

	// Do some more commits
	futures = nil
	for i := 0; i < 100; i++ {
		futures = append(futures, leader.raft.Apply([]byte(fmt.Sprintf("test%d", i)), 0))
	}
	for _, f := range futures {
		NoErr(WaitFuture(f, t), t)
		log.Printf("[DEBUG] Applied %v", f)
	}

	// Shoot two nodes in the head!
	rm1, rm2 := envs[0], envs[1]
	rm1.Release()
	rm2.Release()
	envs = envs[2:]
	time.Sleep(10 * time.Millisecond)

	// Wait for a leader
	leader, err = WaitForAny(Leader, append([]*RaftEnv{env1}, envs...))
	NoErr(err, t)

	// Do some more commits
	futures = nil
	for i := 0; i < 100; i++ {
		futures = append(futures, leader.raft.Apply([]byte(fmt.Sprintf("test%d", i)), 0))
	}
	for _, f := range futures {
		NoErr(WaitFuture(f, t), t)
		log.Printf("[DEBUG] Applied %v", f)
	}

	// Join a few new nodes!
	for i := 0; i < 2; i++ {
		env := MakeRaft(t, conf)
		addr := env.trans.LocalAddr()
		NoErr(WaitFuture(leader.raft.AddPeer(addr), t), t)
		envs = append(envs, env)
	}

	// Remove the old nodes
	NoErr(WaitFuture(leader.raft.RemovePeer(rm1.raft.localAddr), t), t)
	NoErr(WaitFuture(leader.raft.RemovePeer(rm2.raft.localAddr), t), t)

	// Shoot the leader
	env1.Release()
	time.Sleep(3 * conf.HeartbeatTimeout)

	// Wait for a leader
	leader, err = WaitForAny(Leader, envs)
	NoErr(err, t)

	allEnvs := append([]*RaftEnv{env1}, envs...)
	CheckConsistent(allEnvs, t)

	if len(env1.fsm.logs) != 300 {
		t.Fatalf("should apply 300 logs! %d", len(env1.fsm.logs))
	}

	for _, e := range envs {
		e.Release()
	}
}
<commit_msg>Use a more idiomatic initialization style<commit_after>package raft

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"testing"
	"time"
)

// CheckInteg will skip a test if integration testing is not enabled.
func CheckInteg(t *testing.T) {
	if !IsInteg() {
		t.SkipNow()
	}
}

// IsInteg returns a boolean telling you if we're in integ testing mode.
func IsInteg() bool {
	return os.Getenv("INTEG_TESTS") != ""
}

type RaftEnv struct {
	dir      string
	conf     *Config
	fsm      *MockFSM
	store    *InmemStore
	snapshot *FileSnapshotStore
	peers    *JSONPeers
	trans    *NetworkTransport
	raft     *Raft
}

func (r *RaftEnv) Release() {
	log.Printf("[WARN] Release node at %v", r.raft.localAddr)
	f := r.raft.Shutdown()
	if err := f.Error(); err != nil {
		panic(err)
	}
	r.trans.Close()
	os.RemoveAll(r.dir)
}

func MakeRaft(t *testing.T, conf *Config) *RaftEnv {
	// Set the config
	if conf == nil {
		conf = inmemConfig(t)
	}

	dir, err := ioutil.TempDir("", "raft")
	if err != nil {
		t.Fatalf("err: %v ", err)
	}

	stable := NewInmemStore()

	snap, err := NewFileSnapshotStore(dir, 3, nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	env := &RaftEnv{
		conf:     conf,
		dir:      dir,
		store:    stable,
		snapshot: snap,
		fsm:      &MockFSM{},
	}

	trans, err := NewTCPTransport("127.0.0.1:0", nil, 2, time.Second, nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	env.trans = trans

	env.peers = NewJSONPeers(dir, trans)

	log.Printf("[INFO] Starting node at %v", trans.LocalAddr())
	raft, err := NewRaft(conf, env.fsm, stable, stable, snap, env.peers, trans)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	env.raft = raft
	return env
}

func WaitFor(env *RaftEnv, state RaftState) error {
	limit := time.Now().Add(200 * time.Millisecond)
	for env.raft.State() != state {
		if time.Now().Before(limit) {
			time.Sleep(10 * time.Millisecond)
		} else {
			return fmt.Errorf("failed to transition to state %v", state)
		}
	}
	return nil
}

func WaitForAny(state RaftState, envs []*RaftEnv) (*RaftEnv, error) {
	limit := time.Now().Add(200 * time.Millisecond)
CHECK:
	for _, env := range envs {
		if env.raft.State() == state {
			return env, nil
		}
	}
	if time.Now().Before(limit) {
		goto WAIT
	}
	return nil, fmt.Errorf("failed to find node in %v state", state)
WAIT:
	time.Sleep(10 * time.Millisecond)
	goto CHECK
}

func WaitFuture(f Future, t *testing.T) error {
	timer := time.AfterFunc(200*time.Millisecond, func() {
		panic(fmt.Errorf("timeout waiting for future %v", f))
	})
	defer timer.Stop()
	return f.Error()
}

func NoErr(err error, t *testing.T) {
	if err != nil {
		t.Fatalf("err: %v", err)
	}
}

func CheckConsistent(envs []*RaftEnv, t *testing.T) {
	limit := time.Now().Add(400 * time.Millisecond)
	first := envs[0]
	var err error
CHECK:
	l1 := len(first.fsm.logs)
	for i := 1; i < len(envs); i++ {
		env := envs[i]
		l2 := len(env.fsm.logs)
		if l1 != l2 {
			err = fmt.Errorf("log length mismatch %d %d", l1, l2)
			goto ERR
		}
		for idx, log := range first.fsm.logs {
			other := env.fsm.logs[idx]
			if bytes.Compare(log, other) != 0 {
				err = fmt.Errorf("log %d mismatch %v %v", idx, log, other)
				goto ERR
			}
		}
	}
	return
ERR:
	if time.Now().After(limit) {
		t.Fatalf("%v", err)
	}
	time.Sleep(20 * time.Millisecond)
	goto CHECK
}

// Tests Raft by creating a cluster, growing it to 5 nodes while
// causing various stressful conditions
func TestRaft_Integ(t *testing.T) {
	CheckInteg(t)
	conf := DefaultConfig()
	conf.HeartbeatTimeout = 50 * time.Millisecond
	conf.ElectionTimeout = 50 * time.Millisecond
	conf.LeaderLeaseTimeout = 50 * time.Millisecond
	conf.CommitTimeout = 5 * time.Millisecond
	conf.SnapshotThreshold = 100
	conf.TrailingLogs = 10
	conf.EnableSingleNode = true

	// Create a single node
	env1 := MakeRaft(t, conf)
	NoErr(WaitFor(env1, Leader), t)

	// Do some commits
	var futures []Future
	for i := 0; i < 100; i++ {
		futures = append(futures, env1.raft.Apply([]byte(fmt.Sprintf("test%d", i)), 0))
	}
	for _, f := range futures {
		NoErr(WaitFuture(f, t), t)
		log.Printf("[DEBUG] Applied %v", f)
	}

	// Do a snapshot
	NoErr(WaitFuture(env1.raft.Snapshot(), t), t)

	// Join a few nodes!
	var envs []*RaftEnv
	for i := 0; i < 4; i++ {
		env := MakeRaft(t, conf)
		addr := env.trans.LocalAddr()
		NoErr(WaitFuture(env1.raft.AddPeer(addr), t), t)
		envs = append(envs, env)
	}

	// Wait for a leader
	leader, err := WaitForAny(Leader, append([]*RaftEnv{env1}, envs...))
	NoErr(err, t)

	// Do some more commits
	futures = nil
	for i := 0; i < 100; i++ {
		futures = append(futures, leader.raft.Apply([]byte(fmt.Sprintf("test%d", i)), 0))
	}
	for _, f := range futures {
		NoErr(WaitFuture(f, t), t)
		log.Printf("[DEBUG] Applied %v", f)
	}

	// Shoot two nodes in the head!
	rm1, rm2 := envs[0], envs[1]
	rm1.Release()
	rm2.Release()
	envs = envs[2:]
	time.Sleep(10 * time.Millisecond)

	// Wait for a leader
	leader, err = WaitForAny(Leader, append([]*RaftEnv{env1}, envs...))
	NoErr(err, t)

	// Do some more commits
	futures = nil
	for i := 0; i < 100; i++ {
		futures = append(futures, leader.raft.Apply([]byte(fmt.Sprintf("test%d", i)), 0))
	}
	for _, f := range futures {
		NoErr(WaitFuture(f, t), t)
		log.Printf("[DEBUG] Applied %v", f)
	}

	// Join a few new nodes!
	for i := 0; i < 2; i++ {
		env := MakeRaft(t, conf)
		addr := env.trans.LocalAddr()
		NoErr(WaitFuture(leader.raft.AddPeer(addr), t), t)
		envs = append(envs, env)
	}

	// Remove the old nodes
	NoErr(WaitFuture(leader.raft.RemovePeer(rm1.raft.localAddr), t), t)
	NoErr(WaitFuture(leader.raft.RemovePeer(rm2.raft.localAddr), t), t)

	// Shoot the leader
	env1.Release()
	time.Sleep(3 * conf.HeartbeatTimeout)

	// Wait for a leader
	leader, err = WaitForAny(Leader, envs)
	NoErr(err, t)

	allEnvs := append([]*RaftEnv{env1}, envs...)
	CheckConsistent(allEnvs, t)

	if len(env1.fsm.logs) != 300 {
		t.Fatalf("should apply 300 logs! %d", len(env1.fsm.logs))
	}

	for _, e := range envs {
		e.Release()
	}
}
<|endoftext|>


<commit_before>package etcd

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"
	"time"
)

// See SetConsistency for how to use these constants.
const (
	// Using strings rather than iota because the consistency level
	// could be persisted to disk, so it'd be better to use
	// human-readable values.
	STRONG_CONSISTENCY = "STRONG"
	WEAK_CONSISTENCY   = "WEAK"
)

const (
	defaultBufferSize = 10
)

func init() {
	rand.Seed(int64(time.Now().Nanosecond()))
}

type Config struct {
	CertFile    string        `json:"certFile"`
	KeyFile     string        `json:"keyFile"`
	CaCertFile  []string      `json:"caCertFiles"`
	DialTimeout time.Duration `json:"timeout"`
	Consistency string        `json:"consistency"`
}

type credentials struct {
	username string
	password string
}

type Client struct {
	config      Config   `json:"config"`
	cluster     *Cluster `json:"cluster"`
	httpClient  *http.Client
	credentials *credentials
	transport   *http.Transport
	persistence io.Writer
	cURLch      chan string
	// CheckRetry can be used to control the policy for failed requests
	// and modify the cluster if needed.
	// The client calls it before sending requests again, and
	// stops retrying if CheckRetry returns some error. The cases that
	// this function needs to handle include no response and unexpected
	// http status code of response.
	// If CheckRetry is nil, client will call the default one
	// `DefaultCheckRetry`.
	// Argument cluster is the etcd.Cluster object that these requests have been made on.
	// Argument numReqs is the number of http.Requests that have been made so far.
	// Argument lastResp is the http.Responses from the last request.
	// Argument err is the reason of the failure.
	CheckRetry func(cluster *Cluster, numReqs int,
		lastResp http.Response, err error) error
}

// NewClient create a basic client that is configured to be used
// with the given machine list.
func NewClient(machines []string) *Client {
	config := Config{
		// default timeout is one second
		DialTimeout: time.Second,
		Consistency: WEAK_CONSISTENCY,
	}

	client := &Client{
		cluster: NewCluster(machines),
		config:  config,
	}

	client.initHTTPClient()
	client.saveConfig()

	return client
}

// NewTLSClient create a basic client with TLS configuration
func NewTLSClient(machines []string, cert, key, caCert string) (*Client, error) {
	// overwrite the default machine to use https
	if len(machines) == 0 {
		machines = []string{"https://127.0.0.1:4001"}
	}

	config := Config{
		// default timeout is one second
		DialTimeout: time.Second,
		Consistency: WEAK_CONSISTENCY,
		CertFile:    cert,
		KeyFile:     key,
		CaCertFile:  make([]string, 0),
	}

	client := &Client{
		cluster: NewCluster(machines),
		config:  config,
	}

	err := client.initHTTPSClient(cert, key)
	if err != nil {
		return nil, err
	}

	err = client.AddRootCA(caCert)

	client.saveConfig()

	return client, nil
}

// NewClientFromFile creates a client from a given file path.
// The given file is expected to use the JSON format.
func NewClientFromFile(fpath string) (*Client, error) {
	fi, err := os.Open(fpath)
	if err != nil {
		return nil, err
	}

	defer func() {
		if err := fi.Close(); err != nil {
			panic(err)
		}
	}()

	return NewClientFromReader(fi)
}

// NewClientFromReader creates a Client configured from a given reader.
// The configuration is expected to use the JSON format.
func NewClientFromReader(reader io.Reader) (*Client, error) {
	c := new(Client)

	b, err := ioutil.ReadAll(reader)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(b, c)
	if err != nil {
		return nil, err
	}
	if c.config.CertFile == "" {
		c.initHTTPClient()
	} else {
		err = c.initHTTPSClient(c.config.CertFile, c.config.KeyFile)
	}

	if err != nil {
		return nil, err
	}

	for _, caCert := range c.config.CaCertFile {
		if err := c.AddRootCA(caCert); err != nil {
			return nil, err
		}
	}

	return c, nil
}

// Override the Client's HTTP Transport object
func (c *Client) SetTransport(tr *http.Transport) {
	c.httpClient.Transport = tr
	c.transport = tr
}

func (c *Client) SetCredentials(username, password string) {
	c.credentials = &credentials{username, password}
}

func (c *Client) Close() {
	c.transport.DisableKeepAlives = true
	c.transport.CloseIdleConnections()
}

// initHTTPClient initializes a HTTP client for etcd client
func (c *Client) initHTTPClient() {
	c.transport = &http.Transport{
		Dial: c.DefaultDial,
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true,
		},
	}
	c.httpClient = &http.Client{Transport: c.transport}
}

// initHTTPClient initializes a HTTPS client for etcd client
func (c *Client) initHTTPSClient(cert, key string) error {
	if cert == "" || key == "" {
		return errors.New("Require both cert and key path")
	}

	tlsCert, err := tls.LoadX509KeyPair(cert, key)
	if err != nil {
		return err
	}

	tlsConfig := &tls.Config{
		Certificates:       []tls.Certificate{tlsCert},
		InsecureSkipVerify: true,
	}

	tr := &http.Transport{
		TLSClientConfig: tlsConfig,
		Dial:            c.DefaultDial,
	}

	c.httpClient = &http.Client{Transport: tr}
	return nil
}

// SetPersistence sets a writer to which the config will be
// written every time it's changed.
func (c *Client) SetPersistence(writer io.Writer) {
	c.persistence = writer
}

// SetConsistency changes the consistency level of the client.
//
// When consistency is set to STRONG_CONSISTENCY, all requests,
// including GET, are sent to the leader.  This means that, assuming
// the absence of leader failures, GET requests are guaranteed to see
// the changes made by previous requests.
//
// When consistency is set to WEAK_CONSISTENCY, other requests
// are still sent to the leader, but GET requests are sent to a
// random server from the server pool.  This reduces the read
// load on the leader, but it's not guaranteed that the GET requests
// will see changes made by previous requests (they might have not
// yet been committed on non-leader servers).
func (c *Client) SetConsistency(consistency string) error {
	if !(consistency == STRONG_CONSISTENCY || consistency == WEAK_CONSISTENCY) {
		return errors.New("The argument must be either STRONG_CONSISTENCY or WEAK_CONSISTENCY.")
	}
	c.config.Consistency = consistency
	return nil
}

// Sets the DialTimeout value
func (c *Client) SetDialTimeout(d time.Duration) {
	c.config.DialTimeout = d
}

// AddRootCA adds a root CA cert for the etcd client
func (c *Client) AddRootCA(caCert string) error {
	if c.httpClient == nil {
		return errors.New("Client has not been initialized yet!")
	}

	certBytes, err := ioutil.ReadFile(caCert)
	if err != nil {
		return err
	}

	tr, ok := c.httpClient.Transport.(*http.Transport)

	if !ok {
		panic("AddRootCA(): Transport type assert should not fail")
	}

	if tr.TLSClientConfig.RootCAs == nil {
		caCertPool := x509.NewCertPool()
		ok = caCertPool.AppendCertsFromPEM(certBytes)
		if ok {
			tr.TLSClientConfig.RootCAs = caCertPool
		}
		tr.TLSClientConfig.InsecureSkipVerify = false
	} else {
		ok = tr.TLSClientConfig.RootCAs.AppendCertsFromPEM(certBytes)
	}

	if !ok {
		err = errors.New("Unable to load caCert")
	}

	c.config.CaCertFile = append(c.config.CaCertFile, caCert)
	c.saveConfig()

	return err
}

// SetCluster updates cluster information using the given machine list.
func (c *Client) SetCluster(machines []string) bool {
	success := c.internalSyncCluster(machines)
	return success
}

func (c *Client) GetCluster() []string {
	return c.cluster.Machines
}

// SyncCluster updates the cluster information using the internal machine list.
// If no members are found, the intenral machine list is left untouched.
func (c *Client) SyncCluster() bool {
	return c.internalSyncCluster(c.cluster.Machines)
}

// internalSyncCluster syncs cluster information using the given machine list.
func (c *Client) internalSyncCluster(machines []string) bool {
	// comma-separated list of machines in the cluster.
	members := ""

	for _, machine := range machines {
		httpPath := c.createHttpPath(machine, path.Join(version, "members"))
		resp, err := c.httpClient.Get(httpPath)
		if err != nil {
			// try another machine in the cluster
			continue
		}

		if resp.StatusCode != http.StatusOK { // fall-back to old endpoint
			httpPath := c.createHttpPath(machine, path.Join(version, "machines"))
			resp, err := c.httpClient.Get(httpPath)
			if err != nil {
				// try another machine in the cluster
				continue
			}
			b, err := ioutil.ReadAll(resp.Body)
			resp.Body.Close()
			if err != nil {
				// try another machine in the cluster
				continue
			}
			members = string(b)
		} else {
			b, err := ioutil.ReadAll(resp.Body)
			resp.Body.Close()
			if err != nil {
				// try another machine in the cluster
				continue
			}

			var mCollection memberCollection
			if err := json.Unmarshal(b, &mCollection); err != nil {
				// try another machine
				continue
			}

			urls := make([]string, 0)
			for _, m := range mCollection {
				urls = append(urls, m.ClientURLs...)
			}

			members = strings.Join(urls, ",")
		}

		// We should never do an empty cluster update.
		if members == "" {
			continue
		}

		// update Machines List
		c.cluster.updateFromStr(members)
		logger.Debug("sync.machines ", c.cluster.Machines)
		c.saveConfig()
		return true
	}

	return false
}

// createHttpPath creates a complete HTTP URL.
// serverName should contain both the host name and a port number, if any.
func (c *Client) createHttpPath(serverName string, _path string) string {
	u, err := url.Parse(serverName)
	if err != nil {
		panic(err)
	}

	u.Path = path.Join(u.Path, _path)

	if u.Scheme == "" {
		u.Scheme = "http"
	}
	return u.String()
}

// DefaultDial attempts to open a TCP connection to the provided address, explicitly
// enabling keep-alives with a one-second interval.
func (c *Client) DefaultDial(network, addr string) (net.Conn, error) {
	conn, err := net.DialTimeout(network, addr, c.config.DialTimeout)
	if err != nil {
		return nil, err
	}

	tcpConn, ok := conn.(*net.TCPConn)
	if !ok {
		return nil, errors.New("Failed type-assertion of net.Conn as *net.TCPConn")
	}

	// Keep TCP alive to check whether or not the remote machine is down
	if err = tcpConn.SetKeepAlive(true); err != nil {
		return nil, err
	}

	if err = tcpConn.SetKeepAlivePeriod(time.Second); err != nil {
		return nil, err
	}

	return tcpConn, nil
}

func (c *Client) OpenCURL() {
	c.cURLch = make(chan string, defaultBufferSize)
}

func (c *Client) CloseCURL() {
	c.cURLch = nil
}

func (c *Client) sendCURL(command string) {
	go func() {
		select {
		case c.cURLch <- command:
		default:
		}
	}()
}

func (c *Client) RecvCURL() string {
	return <-c.cURLch
}

// saveConfig saves the current config using c.persistence.
func (c *Client) saveConfig() error {
	if c.persistence != nil {
		b, err := json.Marshal(c)
		if err != nil {
			return err
		}

		_, err = c.persistence.Write(b)
		if err != nil {
			return err
		}
	}

	return nil
}

// MarshalJSON implements the Marshaller interface
// as defined by the standard JSON package.
func (c *Client) MarshalJSON() ([]byte, error) {
	b, err := json.Marshal(struct {
		Config  Config   `json:"config"`
		Cluster *Cluster `json:"cluster"`
	}{
		Config:  c.config,
		Cluster: c.cluster,
	})

	if err != nil {
		return nil, err
	}

	return b, nil
}

// UnmarshalJSON implements the Unmarshaller interface
// as defined by the standard JSON package.
func (c *Client) UnmarshalJSON(b []byte) error {
	temp := struct {
		Config  Config   `json:"config"`
		Cluster *Cluster `json:"cluster"`
	}{}
	err := json.Unmarshal(b, &temp)
	if err != nil {
		return err
	}

	c.cluster = temp.Cluster
	c.config = temp.Config
	return nil
}
<commit_msg>rely on net.Dialer to set keepalives and ignore the errors<commit_after>package etcd

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"
	"time"
)

// See SetConsistency for how to use these constants.
const (
	// Using strings rather than iota because the consistency level
	// could be persisted to disk, so it'd be better to use
	// human-readable values.
	STRONG_CONSISTENCY = "STRONG"
	WEAK_CONSISTENCY   = "WEAK"
)

const (
	defaultBufferSize = 10
)

func init() {
	rand.Seed(int64(time.Now().Nanosecond()))
}

type Config struct {
	CertFile    string        `json:"certFile"`
	KeyFile     string        `json:"keyFile"`
	CaCertFile  []string      `json:"caCertFiles"`
	DialTimeout time.Duration `json:"timeout"`
	Consistency string        `json:"consistency"`
}

type credentials struct {
	username string
	password string
}

type Client struct {
	config      Config   `json:"config"`
	cluster     *Cluster `json:"cluster"`
	httpClient  *http.Client
	credentials *credentials
	transport   *http.Transport
	persistence io.Writer
	cURLch      chan string
	// CheckRetry can be used to control the policy for failed requests
	// and modify the cluster if needed.
	// The client calls it before sending requests again, and
	// stops retrying if CheckRetry returns some error. The cases that
	// this function needs to handle include no response and unexpected
	// http status code of response.
	// If CheckRetry is nil, client will call the default one
	// `DefaultCheckRetry`.
	// Argument cluster is the etcd.Cluster object that these requests have been made on.
	// Argument numReqs is the number of http.Requests that have been made so far.
	// Argument lastResp is the http.Responses from the last request.
	// Argument err is the reason of the failure.
	CheckRetry func(cluster *Cluster, numReqs int,
		lastResp http.Response, err error) error
}

// NewClient create a basic client that is configured to be used
// with the given machine list.
func NewClient(machines []string) *Client {
	config := Config{
		// default timeout is one second
		DialTimeout: time.Second,
		Consistency: WEAK_CONSISTENCY,
	}

	client := &Client{
		cluster: NewCluster(machines),
		config:  config,
	}

	client.initHTTPClient()
	client.saveConfig()

	return client
}

// NewTLSClient create a basic client with TLS configuration
func NewTLSClient(machines []string, cert, key, caCert string) (*Client, error) {
	// overwrite the default machine to use https
	if len(machines) == 0 {
		machines = []string{"https://127.0.0.1:4001"}
	}

	config := Config{
		// default timeout is one second
		DialTimeout: time.Second,
		Consistency: WEAK_CONSISTENCY,
		CertFile:    cert,
		KeyFile:     key,
		CaCertFile:  make([]string, 0),
	}

	client := &Client{
		cluster: NewCluster(machines),
		config:  config,
	}

	err := client.initHTTPSClient(cert, key)
	if err != nil {
		return nil, err
	}

	err = client.AddRootCA(caCert)

	client.saveConfig()

	return client, nil
}

// NewClientFromFile creates a client from a given file path.
// The given file is expected to use the JSON format.
func NewClientFromFile(fpath string) (*Client, error) {
	fi, err := os.Open(fpath)
	if err != nil {
		return nil, err
	}

	defer func() {
		if err := fi.Close(); err != nil {
			panic(err)
		}
	}()

	return NewClientFromReader(fi)
}

// NewClientFromReader creates a Client configured from a given reader.
// The configuration is expected to use the JSON format.
func NewClientFromReader(reader io.Reader) (*Client, error) {
	c := new(Client)

	b, err := ioutil.ReadAll(reader)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(b, c)
	if err != nil {
		return nil, err
	}
	if c.config.CertFile == "" {
		c.initHTTPClient()
	} else {
		err = c.initHTTPSClient(c.config.CertFile, c.config.KeyFile)
	}

	if err != nil {
		return nil, err
	}

	for _, caCert := range c.config.CaCertFile {
		if err := c.AddRootCA(caCert); err != nil {
			return nil, err
		}
	}

	return c, nil
}

// Override the Client's HTTP Transport object
func (c *Client) SetTransport(tr *http.Transport) {
	c.httpClient.Transport = tr
	c.transport = tr
}

func (c *Client) SetCredentials(username, password string) {
	c.credentials = &credentials{username, password}
}

func (c *Client) Close() {
	c.transport.DisableKeepAlives = true
	c.transport.CloseIdleConnections()
}

// initHTTPClient initializes a HTTP client for etcd client
func (c *Client) initHTTPClient() {
	c.transport = &http.Transport{
		Dial: c.DefaultDial,
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true,
		},
	}
	c.httpClient = &http.Client{Transport: c.transport}
}

// initHTTPClient initializes a HTTPS client for etcd client
func (c *Client) initHTTPSClient(cert, key string) error {
	if cert == "" || key == "" {
		return errors.New("Require both cert and key path")
	}

	tlsCert, err := tls.LoadX509KeyPair(cert, key)
	if err != nil {
		return err
	}

	tlsConfig := &tls.Config{
		Certificates:       []tls.Certificate{tlsCert},
		InsecureSkipVerify: true,
	}

	tr := &http.Transport{
		TLSClientConfig: tlsConfig,
		Dial:            c.DefaultDial,
	}

	c.httpClient = &http.Client{Transport: tr}
	return nil
}

// SetPersistence sets a writer to which the config will be
// written every time it's changed.
func (c *Client) SetPersistence(writer io.Writer) {
	c.persistence = writer
}

// SetConsistency changes the consistency level of the client.
//
// When consistency is set to STRONG_CONSISTENCY, all requests,
// including GET, are sent to the leader.  This means that, assuming
// the absence of leader failures, GET requests are guaranteed to see
// the changes made by previous requests.
//
// When consistency is set to WEAK_CONSISTENCY, other requests
// are still sent to the leader, but GET requests are sent to a
// random server from the server pool.  This reduces the read
// load on the leader, but it's not guaranteed that the GET requests
// will see changes made by previous requests (they might have not
// yet been committed on non-leader servers).
func (c *Client) SetConsistency(consistency string) error {
	if !(consistency == STRONG_CONSISTENCY || consistency == WEAK_CONSISTENCY) {
		return errors.New("The argument must be either STRONG_CONSISTENCY or WEAK_CONSISTENCY.")
	}
	c.config.Consistency = consistency
	return nil
}

// Sets the DialTimeout value
func (c *Client) SetDialTimeout(d time.Duration) {
	c.config.DialTimeout = d
}

// AddRootCA adds a root CA cert for the etcd client
func (c *Client) AddRootCA(caCert string) error {
	if c.httpClient == nil {
		return errors.New("Client has not been initialized yet!")
	}

	certBytes, err := ioutil.ReadFile(caCert)
	if err != nil {
		return err
	}

	tr, ok := c.httpClient.Transport.(*http.Transport)

	if !ok {
		panic("AddRootCA(): Transport type assert should not fail")
	}

	if tr.TLSClientConfig.RootCAs == nil {
		caCertPool := x509.NewCertPool()
		ok = caCertPool.AppendCertsFromPEM(certBytes)
		if ok {
			tr.TLSClientConfig.RootCAs = caCertPool
		}
		tr.TLSClientConfig.InsecureSkipVerify = false
	} else {
		ok = tr.TLSClientConfig.RootCAs.AppendCertsFromPEM(certBytes)
	}

	if !ok {
		err = errors.New("Unable to load caCert")
	}

	c.config.CaCertFile = append(c.config.CaCertFile, caCert)
	c.saveConfig()

	return err
}

// SetCluster updates cluster information using the given machine list.
func (c *Client) SetCluster(machines []string) bool {
	success := c.internalSyncCluster(machines)
	return success
}

func (c *Client) GetCluster() []string {
	return c.cluster.Machines
}

// SyncCluster updates the cluster information using the internal machine list.
// If no members are found, the intenral machine list is left untouched.
func (c *Client) SyncCluster() bool {
	return c.internalSyncCluster(c.cluster.Machines)
}

// internalSyncCluster syncs cluster information using the given machine list.
func (c *Client) internalSyncCluster(machines []string) bool {
	// comma-separated list of machines in the cluster.
	members := ""

	for _, machine := range machines {
		httpPath := c.createHttpPath(machine, path.Join(version, "members"))
		resp, err := c.httpClient.Get(httpPath)
		if err != nil {
			// try another machine in the cluster
			continue
		}

		if resp.StatusCode != http.StatusOK { // fall-back to old endpoint
			httpPath := c.createHttpPath(machine, path.Join(version, "machines"))
			resp, err := c.httpClient.Get(httpPath)
			if err != nil {
				// try another machine in the cluster
				continue
			}
			b, err := ioutil.ReadAll(resp.Body)
			resp.Body.Close()
			if err != nil {
				// try another machine in the cluster
				continue
			}
			members = string(b)
		} else {
			b, err := ioutil.ReadAll(resp.Body)
			resp.Body.Close()
			if err != nil {
				// try another machine in the cluster
				continue
			}

			var mCollection memberCollection
			if err := json.Unmarshal(b, &mCollection); err != nil {
				// try another machine
				continue
			}

			urls := make([]string, 0)
			for _, m := range mCollection {
				urls = append(urls, m.ClientURLs...)
			}

			members = strings.Join(urls, ",")
		}

		// We should never do an empty cluster update.
		if members == "" {
			continue
		}

		// update Machines List
		c.cluster.updateFromStr(members)
		logger.Debug("sync.machines ", c.cluster.Machines)
		c.saveConfig()
		return true
	}

	return false
}

// createHttpPath creates a complete HTTP URL.
// serverName should contain both the host name and a port number, if any.
func (c *Client) createHttpPath(serverName string, _path string) string {
	u, err := url.Parse(serverName)
	if err != nil {
		panic(err)
	}

	u.Path = path.Join(u.Path, _path)

	if u.Scheme == "" {
		u.Scheme = "http"
	}
	return u.String()
}

// DefaultDial attempts to open a TCP connection to the provided address, explicitly
// enabling keep-alives with a one-second interval.
func (c *Client) DefaultDial(network, addr string) (net.Conn, error) {
	dialer := net.Dialer{
		Timeout:   c.config.DialTimeout,
		KeepAlive: time.Second,
	}

	return dialer.Dial(network, addr)
}

func (c *Client) OpenCURL() {
	c.cURLch = make(chan string, defaultBufferSize)
}

func (c *Client) CloseCURL() {
	c.cURLch = nil
}

func (c *Client) sendCURL(command string) {
	go func() {
		select {
		case c.cURLch <- command:
		default:
		}
	}()
}

func (c *Client) RecvCURL() string {
	return <-c.cURLch
}

// saveConfig saves the current config using c.persistence.
func (c *Client) saveConfig() error {
	if c.persistence != nil {
		b, err := json.Marshal(c)
		if err != nil {
			return err
		}

		_, err = c.persistence.Write(b)
		if err != nil {
			return err
		}
	}

	return nil
}

// MarshalJSON implements the Marshaller interface
// as defined by the standard JSON package.
func (c *Client) MarshalJSON() ([]byte, error) {
	b, err := json.Marshal(struct {
		Config  Config   `json:"config"`
		Cluster *Cluster `json:"cluster"`
	}{
		Config:  c.config,
		Cluster: c.cluster,
	})

	if err != nil {
		return nil, err
	}

	return b, nil
}

// UnmarshalJSON implements the Unmarshaller interface
// as defined by the standard JSON package.
func (c *Client) UnmarshalJSON(b []byte) error {
	temp := struct {
		Config  Config   `json:"config"`
		Cluster *Cluster `json:"cluster"`
	}{}
	err := json.Unmarshal(b, &temp)
	if err != nil {
		return err
	}

	c.cluster = temp.Cluster
	c.config = temp.Config
	return nil
}
<|endoftext|>


<commit_before>package kube

import (
	"bytes"
	"fmt"
	"io"
	"strings"
	"time"

	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"github.com/Sirupsen/logrus"
	"github.com/ghodss/yaml"
	"io/ioutil"
	"net/http"
)

// KubeClient is the interface to access the kubernetes job API
type KubeClient interface {
	CreateJob(job *Job) error
	GetSecret(namespace string, secretName string) (map[string]string, error)
	DeployResourceFile(resourceFile []byte) error
}

// concrete implementation of a job client
type realKubeClient struct {
	*http.Client
	address string
	token   string
}

// NewClient returns a new KubeClient connecting to the address. This uses the service
// account credentials
func NewClient(address string) (KubeClient, error) {
	// create tls client
	cacertFile := "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
	capem, err := ioutil.ReadFile(cacertFile)
	if err != nil {
		return nil, err
	}
	cacert := x509.NewCertPool()
	if !cacert.AppendCertsFromPEM(capem) {
		return nil, fmt.Errorf("unable to load certificate authority")
	}
	config := &tls.Config{RootCAs: cacert}
	transport := &http.Transport{TLSClientConfig: config}

	// read token
	client := &http.Client{Transport: transport}
	tokenFile := "/var/run/secrets/kubernetes.io/serviceaccount/token"
	token, err := ioutil.ReadFile(tokenFile)
	if err != nil {
		return nil, err
	}
	return &realKubeClient{client, address, string(token)}, nil
}

// Create a new kubernetes Job with the given job
func (r *realKubeClient) CreateJob(job *Job) error {
	url := "/apis/extensions/v1beta1/namespaces/" + job.Metadata["namespace"].(string) + "/jobs"
	data, err := json.Marshal(job)
	if err != nil {
		return err
	}
	byteData := bytes.NewReader(data)
	return r.doPost(url, byteData)

}

// DeployResourceFile deploys a Kubernbetes YAML spec file as Kubernetes resources
func (r *realKubeClient) DeployResourceFile(resourceFile []byte) error {

	resources := strings.Split(string(resourceFile), "---")

	for _, resource := range resources {

		// if it's empty, skip
		if strings.TrimSpace(resource) == "" {
			continue
		}

		logrus.Info("deploying to kubernetes: ", resource)

		data, err := yaml.YAMLToJSON([]byte(resource))
		if err != nil {
			logrus.WithError(err).Error("unable to convert yaml to json")
			return err
		}

		var out map[string]interface{}
		err = json.Unmarshal(data, &out)
		if err != nil {
			logrus.WithError(err).Error("unable to unmarshal json to map")
			return err
		}

		// if unmarshalled data is nil, skip
		if out == nil {
			continue
		}

		kind := strings.ToLower(out["kind"].(string)) + "s"
		metadata := out["metadata"]
		namespace := "default"
		name := ""
		if metadata != nil {
			namespace = metadata.(map[string]interface{})["namespace"].(string)
			name = metadata.(map[string]interface{})["name"].(string)
		}

		// endpoint is /api/v1/namespaces/{namespace}/{resourceType}
		uri := fmt.Sprintf("/api/v1/namespaces/%s/%s/%s", namespace, kind, name)

		err = r.doGet(uri, &out)
		if out != nil {
			err = r.doDelete(uri)
			time.Sleep(30 * time.Second)
			if err != nil {
				logrus.WithError(err).Error("unable to DELETE resource")
			}
		}

		postUri := fmt.Sprintf("/api/v1/namespaces/%s/%s", namespace, kind)
		err = r.doPost(postUri, bytes.NewReader(data))
		if err != nil {
			logrus.WithError(err).Error("unable to POST data")
			return err
		}
	}

	return nil
}

//Get secret with a given namespace and secret name
func (r *realKubeClient) GetSecret(namespace string, secretName string) (map[string]string, error) {
	secret := &Secret{}
	secrets := make(map[string]string)

	uri := "/api/v1/namespaces/" + namespace + "/secrets/" + secretName

	err := r.doGet(uri, secret)
	if err != nil {
		return nil, err
	}

	for key, value := range secret.Data {
		decodedValue, err := base64.StdEncoding.DecodeString(value)

		if err != nil {
			logrus.WithError(err).Println("Unable to decode secret", key, value)
			continue
		}
		secrets[key] = string(decodedValue)
	}
	return secrets, nil
}

func (r *realKubeClient) doGet(uri string, response interface{}) error {
	req, err := r.createRequest("GET", uri, nil)
	if err != nil {
		return err
	}
	res, err := r.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return err
	}
	if res.StatusCode != 200 {
		return fmt.Errorf("%d: %s", res.StatusCode, string(body))
	}
	err = json.Unmarshal(body, response)
	if err != nil {
		return err
	}
	return nil
}

func (r *realKubeClient) doDelete(uri string) error {

	req, err := r.createRequest("DELETE", uri, nil)
	if err != nil {
		return err
	}
	res, err := r.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return err
	}
	if res.StatusCode != 200 {
		return fmt.Errorf("%d: %s", res.StatusCode, string(body))
	}

	return nil
}

func (r *realKubeClient) doPost(uri string, data io.Reader) error {
	req, err := r.createRequest("POST", uri, data)
	if err != nil {
		return err
	}
	res, err := r.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode == http.StatusCreated || res.StatusCode == http.StatusOK {
		return nil
	}
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return err
	}
	return fmt.Errorf("%d: %s", res.StatusCode, string(body))
}

func (r *realKubeClient) createRequest(method, uri string, data io.Reader) (*http.Request, error) {
	req, err := http.NewRequest(method, r.address+uri, data)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Authorization", "Bearer "+r.token)
	return req, nil
}
<commit_msg>wait till resource has been removed<commit_after>package kube

import (
	"bytes"
	"fmt"
	"io"
	"strings"
	"time"

	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"github.com/Sirupsen/logrus"
	"github.com/ghodss/yaml"
	"io/ioutil"
	"net/http"
)

// KubeClient is the interface to access the kubernetes job API
type KubeClient interface {
	CreateJob(job *Job) error
	GetSecret(namespace string, secretName string) (map[string]string, error)
	DeployResourceFile(resourceFile []byte) error
}

// concrete implementation of a job client
type realKubeClient struct {
	*http.Client
	address string
	token   string
}

// NewClient returns a new KubeClient connecting to the address. This uses the service
// account credentials
func NewClient(address string) (KubeClient, error) {
	// create tls client
	cacertFile := "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
	capem, err := ioutil.ReadFile(cacertFile)
	if err != nil {
		return nil, err
	}
	cacert := x509.NewCertPool()
	if !cacert.AppendCertsFromPEM(capem) {
		return nil, fmt.Errorf("unable to load certificate authority")
	}
	config := &tls.Config{RootCAs: cacert}
	transport := &http.Transport{TLSClientConfig: config}

	// read token
	client := &http.Client{Transport: transport}
	tokenFile := "/var/run/secrets/kubernetes.io/serviceaccount/token"
	token, err := ioutil.ReadFile(tokenFile)
	if err != nil {
		return nil, err
	}
	return &realKubeClient{client, address, string(token)}, nil
}

// Create a new kubernetes Job with the given job
func (r *realKubeClient) CreateJob(job *Job) error {
	url := "/apis/extensions/v1beta1/namespaces/" + job.Metadata["namespace"].(string) + "/jobs"
	data, err := json.Marshal(job)
	if err != nil {
		return err
	}
	byteData := bytes.NewReader(data)
	return r.doPost(url, byteData)

}

// DeployResourceFile deploys a Kubernbetes YAML spec file as Kubernetes resources
func (r *realKubeClient) DeployResourceFile(resourceFile []byte) error {

	resources := strings.Split(string(resourceFile), "---")

	for _, resource := range resources {

		// if it's empty, skip
		if strings.TrimSpace(resource) == "" {
			continue
		}

		logrus.Info("deploying to kubernetes: ", resource)

		data, err := yaml.YAMLToJSON([]byte(resource))
		if err != nil {
			logrus.WithError(err).Error("unable to convert yaml to json")
			return err
		}

		var out map[string]interface{}
		err = json.Unmarshal(data, &out)
		if err != nil {
			logrus.WithError(err).Error("unable to unmarshal json to map")
			return err
		}

		// if unmarshalled data is nil, skip
		if out == nil {
			continue
		}

		kind := strings.ToLower(out["kind"].(string)) + "s"
		metadata := out["metadata"]
		namespace := "default"
		name := ""
		if metadata != nil {
			namespace = metadata.(map[string]interface{})["namespace"].(string)
			name = metadata.(map[string]interface{})["name"].(string)
		}

		// endpoint is /api/v1/namespaces/{namespace}/{resourceType}
		uri := fmt.Sprintf("/api/v1/namespaces/%s/%s/%s", namespace, kind, name)

		err = r.doGet(uri, &out)
		if err == nil {
			err = r.doDelete(uri)
			time.Sleep(45 * time.Second)
			if err != nil {
				logrus.WithError(err).Error("unable to DELETE resource")
			}
		}

		postUri := fmt.Sprintf("/api/v1/namespaces/%s/%s", namespace, kind)
		err = r.doPost(postUri, bytes.NewReader(data))
		if err != nil {
			logrus.WithError(err).Error("unable to POST data")
			return err
		}

	}

	return nil
}

//Get secret with a given namespace and secret name
func (r *realKubeClient) GetSecret(namespace string, secretName string) (map[string]string, error) {
	secret := &Secret{}
	secrets := make(map[string]string)

	uri := "/api/v1/namespaces/" + namespace + "/secrets/" + secretName

	err := r.doGet(uri, secret)
	if err != nil {
		return nil, err
	}

	for key, value := range secret.Data {
		decodedValue, err := base64.StdEncoding.DecodeString(value)

		if err != nil {
			logrus.WithError(err).Println("Unable to decode secret", key, value)
			continue
		}
		secrets[key] = string(decodedValue)
	}
	return secrets, nil
}

func (r *realKubeClient) doGet(uri string, response interface{}) error {
	req, err := r.createRequest("GET", uri, nil)
	if err != nil {
		return err
	}
	res, err := r.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return err
	}
	if res.StatusCode != 200 {
		return fmt.Errorf("%d: %s", res.StatusCode, string(body))
	}
	err = json.Unmarshal(body, response)
	if err != nil {
		return err
	}
	return nil
}

func (r *realKubeClient) doDelete(uri string) error {

	req, err := r.createRequest("DELETE", uri, nil)
	if err != nil {
		return err
	}
	res, err := r.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return err
	}
	if res.StatusCode != 200 {
		return fmt.Errorf("%d: %s", res.StatusCode, string(body))
	}

	return nil
}

func (r *realKubeClient) doPost(uri string, data io.Reader) error {
	req, err := r.createRequest("POST", uri, data)
	if err != nil {
		return err
	}
	res, err := r.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	if res.StatusCode == http.StatusCreated || res.StatusCode == http.StatusOK {
		return nil
	}
	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return err
	}
	return fmt.Errorf("%d: %s", res.StatusCode, string(body))
}

func (r *realKubeClient) createRequest(method, uri string, data io.Reader) (*http.Request, error) {
	req, err := http.NewRequest(method, r.address+uri, data)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Authorization", "Bearer "+r.token)
	return req, nil
}
<|endoftext|>