<commit_before>package main

import (
	"encoding/json"
	"time"

	"github.com/flynn/flynn/discoverd/client"
	"github.com/flynn/flynn/host/fixer"
	"github.com/flynn/flynn/pkg/cluster"
	"gopkg.in/inconshreveable/log15.v2"
)

const (
	checkInterval  = 10 * time.Second
	retryInterval  = 5 * time.Second
	deadlineLength = 60 * time.Second
)

var monitorLogger = log15.New("component", "cluster-monitor")

type MonitorMetadata struct {
	Enabled bool `json:"enabled,omitempty"`
	Hosts   int  `json:"hosts,omitempty"`
}

type Monitor struct {
	addr       string
	dm         *DiscoverdManager
	discoverd  *discoverdWrapper
	discClient *discoverd.Client
	monitorSvc discoverd.Service
	isLeader   bool
	c          *cluster.Client
	hostCount  int
	deadline   time.Time
	shutdownCh chan struct{}
	logger     log15.Logger
}

func NewMonitor(dm *DiscoverdManager, addr string, logger log15.Logger) *Monitor {
	return &Monitor{
		dm:         dm,
		discoverd:  nil,
		addr:       addr,
		shutdownCh: make(chan struct{}),
		logger:     logger,
	}
}

func (m *Monitor) waitDiscoverd() {
	for {
		if m.dm.localConnected() {
			m.discClient = discoverd.NewClient()
			break
		}
		time.Sleep(retryInterval)
	}
}

func (m *Monitor) waitRaftLeader() {
	for {
		_, err := m.discClient.RaftLeader()
		if err == nil {
			break
		}
		time.Sleep(retryInterval)
	}
}

func (m *Monitor) waitEnabled() {
	log := m.logger.New("fn", "waitEnabled")
	for {
		monitorMeta, err := m.monitorSvc.GetMeta()
		if err != nil {
			time.Sleep(retryInterval)
			continue
		}
		var decodedMeta MonitorMetadata
		if err := json.Unmarshal(monitorMeta.Data, &decodedMeta); err != nil {
			log.Error("monitor metadata unparsable")
			time.Sleep(retryInterval)
			continue
		}
		if decodedMeta.Enabled {
			m.hostCount = decodedMeta.Hosts
			break
		}
		time.Sleep(retryInterval)
	}

}

func (m *Monitor) waitRegister() {
	for {
		isLeader, err := m.discoverd.Register()
		if err == nil {
			m.isLeader = isLeader
			break
		}
		time.Sleep(retryInterval)
	}
}

func (m *Monitor) Run() {
	log := monitorLogger.New("fn", "Run")
	log.Info("waiting for discoverd")
	m.waitDiscoverd()

	log.Info("waiting for raft leader")
	m.waitRaftLeader()

	// we can connect the leader election wrapper now
	m.discoverd = newDiscoverdWrapper(m.addr+":1113", m.logger)
	// connect cluster client now that discoverd is up.
	m.c = cluster.NewClient()

	m.monitorSvc = discoverd.NewService("cluster-monitor")

	log.Info("waiting for monitor service to be enabled for this cluster")
	m.waitEnabled()

	log.Info("registering cluster-monitor")
	m.waitRegister()

	leaderCh := m.discoverd.LeaderCh()
	ticker := time.NewTicker(checkInterval)

	log.Info("starting monitor loop")
	for {
		var isLeader bool
		select {
		case <-m.shutdownCh:
			log.Info("shutting down monitor")
			return
		case isLeader = <-leaderCh:
			m.isLeader = isLeader
			continue
		default:
		}

		select {
		case <-ticker.C:
			if m.isLeader {
				m.checkCluster()
			}
		}
	}
}

func (m *Monitor) checkCluster() {
	log := monitorLogger.New("fn", "checkCluster")
	var faulted bool
	hosts, err := m.c.Hosts()
	if err != nil || len(hosts) < m.hostCount {
		log.Info("waiting for hosts", "current", len(hosts), "want", m.hostCount)
		return
	}

	controllerInstances, _ := discoverd.NewService("controller").Instances()
	if len(controllerInstances) == 0 {
		log.Error("did not find any controller api instances")
		faulted = true
	}

	if _, err := discoverd.NewService("controller-scheduler").Leader(); err != nil && !discoverd.IsNotFound(err) {
		log.Error("error getting scheduler leader, can't determine health")
	} else if err != nil {
		log.Error("scheduler is not up")
		faulted = true
	}

	if faulted && m.deadline.IsZero() {
		log.Error("cluster is unhealthy, setting fault")
		m.deadline = time.Now().Add(deadlineLength)
	} else if !faulted && !m.deadline.IsZero() {
		log.Info("cluster currently healthy, clearing fault")
		m.deadline = time.Time{}
	}

	if !m.deadline.IsZero() && time.Now().After(m.deadline) {
		log.Error("fault deadline reached")
		if err := m.repairCluster(); err != nil {
			log.Error("error repairing cluster", "err", err)
		}
		return
	}
}

func (m *Monitor) repairCluster() error {
	log := monitorLogger.New("fn", "repairCluster")
	log.Info("initiating cluster repair")
	hosts, err := m.c.Hosts()
	if err != nil {
		return err
	}
	f := fixer.NewClusterFixer(hosts, m.c, log)
	// killing the schedulers to prevent interference
	f.KillSchedulers()

	log.Info("checking status of sirenia databases")
	for _, db := range []string{"postgres", "mariadb"} {
		log.Info("checking for database state", "db", db)
		if _, err := discoverd.NewService(db).GetMeta(); err != nil {
			if discoverd.IsNotFound(err) {
				log.Info("skipping recovery of db, no state in discoverd", "db", db)
				continue
			}
			log.Error("error checking database state", "db", db)
			return err
		}
		if err := f.CheckSirenia(db); err != nil {
			if err := f.FixSirenia(db); err != nil {
				if db == "postgres" {
					return err
				} else {
					log.Error("failed database recovery", "db", db)
				}
			}
		}
	}

	// ensure controller api is working
	controllerService := discoverd.NewService("controller")
	controllerInstances, _ := controllerService.Instances()
	if len(controllerInstances) == 0 {
		controllerInstances, err = f.StartAppJob("controller", "web", "controller")
		if err != nil {
			return err
		}
	}
	// fix any formations and start the scheduler again
	if err := f.FixController(controllerInstances, true); err != nil {
		return err
	}
	// zero out the deadline timer
	m.deadline = time.Time{}
	return nil
}

func (m *Monitor) Shutdown() {
	close(m.shutdownCh)
}
<commit_msg>host: Repair MongoDB cluster before the controller<commit_after>package main

import (
	"encoding/json"
	"time"

	"github.com/flynn/flynn/discoverd/client"
	"github.com/flynn/flynn/host/fixer"
	"github.com/flynn/flynn/pkg/cluster"
	"gopkg.in/inconshreveable/log15.v2"
)

const (
	checkInterval  = 10 * time.Second
	retryInterval  = 5 * time.Second
	deadlineLength = 60 * time.Second
)

var monitorLogger = log15.New("component", "cluster-monitor")

type MonitorMetadata struct {
	Enabled bool `json:"enabled,omitempty"`
	Hosts   int  `json:"hosts,omitempty"`
}

type Monitor struct {
	addr       string
	dm         *DiscoverdManager
	discoverd  *discoverdWrapper
	discClient *discoverd.Client
	monitorSvc discoverd.Service
	isLeader   bool
	c          *cluster.Client
	hostCount  int
	deadline   time.Time
	shutdownCh chan struct{}
	logger     log15.Logger
}

func NewMonitor(dm *DiscoverdManager, addr string, logger log15.Logger) *Monitor {
	return &Monitor{
		dm:         dm,
		discoverd:  nil,
		addr:       addr,
		shutdownCh: make(chan struct{}),
		logger:     logger,
	}
}

func (m *Monitor) waitDiscoverd() {
	for {
		if m.dm.localConnected() {
			m.discClient = discoverd.NewClient()
			break
		}
		time.Sleep(retryInterval)
	}
}

func (m *Monitor) waitRaftLeader() {
	for {
		_, err := m.discClient.RaftLeader()
		if err == nil {
			break
		}
		time.Sleep(retryInterval)
	}
}

func (m *Monitor) waitEnabled() {
	log := m.logger.New("fn", "waitEnabled")
	for {
		monitorMeta, err := m.monitorSvc.GetMeta()
		if err != nil {
			time.Sleep(retryInterval)
			continue
		}
		var decodedMeta MonitorMetadata
		if err := json.Unmarshal(monitorMeta.Data, &decodedMeta); err != nil {
			log.Error("monitor metadata unparsable")
			time.Sleep(retryInterval)
			continue
		}
		if decodedMeta.Enabled {
			m.hostCount = decodedMeta.Hosts
			break
		}
		time.Sleep(retryInterval)
	}

}

func (m *Monitor) waitRegister() {
	for {
		isLeader, err := m.discoverd.Register()
		if err == nil {
			m.isLeader = isLeader
			break
		}
		time.Sleep(retryInterval)
	}
}

func (m *Monitor) Run() {
	log := monitorLogger.New("fn", "Run")
	log.Info("waiting for discoverd")
	m.waitDiscoverd()

	log.Info("waiting for raft leader")
	m.waitRaftLeader()

	// we can connect the leader election wrapper now
	m.discoverd = newDiscoverdWrapper(m.addr+":1113", m.logger)
	// connect cluster client now that discoverd is up.
	m.c = cluster.NewClient()

	m.monitorSvc = discoverd.NewService("cluster-monitor")

	log.Info("waiting for monitor service to be enabled for this cluster")
	m.waitEnabled()

	log.Info("registering cluster-monitor")
	m.waitRegister()

	leaderCh := m.discoverd.LeaderCh()
	ticker := time.NewTicker(checkInterval)

	log.Info("starting monitor loop")
	for {
		var isLeader bool
		select {
		case <-m.shutdownCh:
			log.Info("shutting down monitor")
			return
		case isLeader = <-leaderCh:
			m.isLeader = isLeader
			continue
		default:
		}

		select {
		case <-ticker.C:
			if m.isLeader {
				m.checkCluster()
			}
		}
	}
}

func (m *Monitor) checkCluster() {
	log := monitorLogger.New("fn", "checkCluster")
	var faulted bool
	hosts, err := m.c.Hosts()
	if err != nil || len(hosts) < m.hostCount {
		log.Info("waiting for hosts", "current", len(hosts), "want", m.hostCount)
		return
	}

	controllerInstances, _ := discoverd.NewService("controller").Instances()
	if len(controllerInstances) == 0 {
		log.Error("did not find any controller api instances")
		faulted = true
	}

	if _, err := discoverd.NewService("controller-scheduler").Leader(); err != nil && !discoverd.IsNotFound(err) {
		log.Error("error getting scheduler leader, can't determine health")
	} else if err != nil {
		log.Error("scheduler is not up")
		faulted = true
	}

	if faulted && m.deadline.IsZero() {
		log.Error("cluster is unhealthy, setting fault")
		m.deadline = time.Now().Add(deadlineLength)
	} else if !faulted && !m.deadline.IsZero() {
		log.Info("cluster currently healthy, clearing fault")
		m.deadline = time.Time{}
	}

	if !m.deadline.IsZero() && time.Now().After(m.deadline) {
		log.Error("fault deadline reached")
		if err := m.repairCluster(); err != nil {
			log.Error("error repairing cluster", "err", err)
		}
		return
	}
}

func (m *Monitor) repairCluster() error {
	log := monitorLogger.New("fn", "repairCluster")
	log.Info("initiating cluster repair")
	hosts, err := m.c.Hosts()
	if err != nil {
		return err
	}
	f := fixer.NewClusterFixer(hosts, m.c, log)
	// killing the schedulers to prevent interference
	f.KillSchedulers()

	log.Info("checking status of sirenia databases")
	for _, db := range []string{"postgres", "mariadb", "mongodb"} {
		log.Info("checking for database state", "db", db)
		if _, err := discoverd.NewService(db).GetMeta(); err != nil {
			if discoverd.IsNotFound(err) {
				log.Info("skipping recovery of db, no state in discoverd", "db", db)
				continue
			}
			log.Error("error checking database state", "db", db)
			return err
		}
		if err := f.CheckSirenia(db); err != nil {
			if err := f.FixSirenia(db); err != nil {
				if db == "postgres" {
					return err
				} else {
					log.Error("failed database recovery", "db", db)
				}
			}
		}
	}

	// ensure controller api is working
	controllerService := discoverd.NewService("controller")
	controllerInstances, _ := controllerService.Instances()
	if len(controllerInstances) == 0 {
		controllerInstances, err = f.StartAppJob("controller", "web", "controller")
		if err != nil {
			return err
		}
	}
	// fix any formations and start the scheduler again
	if err := f.FixController(controllerInstances, true); err != nil {
		return err
	}
	// zero out the deadline timer
	m.deadline = time.Time{}
	return nil
}

func (m *Monitor) Shutdown() {
	close(m.shutdownCh)
}
<|endoftext|>


<commit_before>// +build windows

package main

import (
	"os"
)

func (h Host) HostInfo() (HostInfo, error) {
	ret := HostInfo{}
	hostname, err := os.Hostname()
	if err != nil {
		return ret, err
	}

	ret.Hostname = hostname

	return ret, nil
}
<commit_msg>implements Host.Uptime on windows<commit_after>// +build windows

package main

import (
	"os"
	"syscall"
)

func (h Host) HostInfo() (HostInfo, error) {
	ret := HostInfo{}
	hostname, err := os.Hostname()
	if err != nil {
		return ret, err
	}

	ret.Hostname = hostname

	kernel32, err := syscall.LoadLibrary("kernel32.dll")
	if err != nil {
		return ret, err
	}
	defer syscall.FreeLibrary(kernel32)
	GetTickCount, _ := syscall.GetProcAddress(kernel32, "GetTickCount")

	uptimemsec, _, err := syscall.Syscall(uintptr(GetTickCount), 0, 0, 0, 0)

	ret.Uptime = int64(uptimemsec) / 1000
	return ret, nil
}
<|endoftext|>


<commit_before>// Package parth provides path parsing for segment unmarshaling and slicing.
//
// Along with string, all basic non-alias types are supported. An interface is
// available for implementation by user-defined types. When handling an int,
// uint, or float of any size, the first and longest valid value within the
// specified segment will be used.
package parth

import (
	"errors"
)

// Unmarshaler is the interface implemented by types that can unmarshal a path
// segment representation of themselves. It is safe to assume that the segment
// data will not include slashes.
type Unmarshaler interface {
	UnmarshalSegment(string) error
}

// Err{Name} values facilitate error identification.
var (
	ErrUnknownType = errors.New("unknown type provided")

	ErrFirstSegNotFound = errors.New("first segment not found by index")
	ErrLastSegNotFound  = errors.New("last segment not found by index")
	ErrSegOrderReversed = errors.New("first segment must precede last segment")
	ErrKeySegNotFound   = errors.New("segment not found by key")

	ErrDataUnparsable = errors.New("data cannot be parsed")
)

// Segment locates the path segment indicated by the index i and unmarshals it
// into the provided type v. An error is returned if: 1. The type is not a
// pointer to an instance of one of the basic non-alias types and does not
// implement the Unmarshaler interface; 2. The index is out of range of the
// path; 3. The located path segment data cannot be parsed as the type or if an
// error is returned by an Unmarshaler implementation.
func Segment(path string, i int, v interface{}) error { //nolint
	var err error

	switch v := v.(type) {
	case *bool:
		*v, err = segmentToBool(path, i)

	case *float32:
		var f float64
		f, err = segmentToFloatN(path, i, 32)
		*v = float32(f)

	case *float64:
		*v, err = segmentToFloatN(path, i, 64)

	case *int:
		var n int64
		n, err = segmentToIntN(path, i, 0)
		*v = int(n)

	case *int16:
		var n int64
		n, err = segmentToIntN(path, i, 16)
		*v = int16(n)

	case *int32:
		var n int64
		n, err = segmentToIntN(path, i, 32)
		*v = int32(n)

	case *int64:
		*v, err = segmentToIntN(path, i, 64)

	case *int8:
		var n int64
		n, err = segmentToIntN(path, i, 8)
		*v = int8(n)

	case *string:
		*v, err = segmentToString(path, i)

	case *uint:
		var n uint64
		n, err = segmentToUintN(path, i, 0)
		*v = uint(n)

	case *uint16:
		var n uint64
		n, err = segmentToUintN(path, i, 16)
		*v = uint16(n)

	case *uint32:
		var n uint64
		n, err = segmentToUintN(path, i, 32)
		*v = uint32(n)

	case *uint64:
		*v, err = segmentToUintN(path, i, 64)

	case *uint8:
		var n uint64
		n, err = segmentToUintN(path, i, 8)
		*v = uint8(n)

	case Unmarshaler:
		var s string
		s, err = segmentToString(path, i)
		if err == nil {
			err = v.UnmarshalSegment(s)
		}

	default:
		err = ErrUnknownType
	}

	return err
}

// Sequent is similar to Segment, but uses a key to locate a segment and then
// unmarshal the subsequent segment. It is a simple wrapper over SubSeg with an
// index of 0.
func Sequent(path, key string, v interface{}) error {
	return SubSeg(path, key, 0, v)
}

// Span returns the path segments between two segment indexes i and j including
// the first segment. An error is returned if: 1. Either index is out of range
// of the path; 2. The first index i does not precede the last index j.
// Providing a 0 for the last index is a special case which acts as an alias
// for the end of the path.
func Span(path string, i, j int) (string, error) {
	var f, l int
	var ok bool

	if i < 0 {
		f, ok = segStartIndexFromEnd(path, i)
	} else {
		f, ok = segStartIndexFromStart(path, i)
	}
	if !ok {
		return "", ErrFirstSegNotFound
	}

	if j > 0 {
		l, ok = segEndIndexFromStart(path, j)
	} else {
		l, ok = segEndIndexFromEnd(path, j)
	}
	if !ok {
		return "", ErrLastSegNotFound
	}

	if f == l {
		return "", nil
	}

	if f > l {
		return "", ErrSegOrderReversed
	}

	return path[f:l], nil
}

// SubSeg is similar to Segment, but only handles the portion of the path
// subsequent to the provided key. For example, to access the segment
// immediately after a key, an index of 0 should be provided (see Sequent).
func SubSeg(path, key string, i int, v interface{}) error { //nolint
	var err error

	switch v := v.(type) {
	case *bool:
		*v, err = subSegToBool(path, key, i)

	case *float32:
		var f float64
		f, err = subSegToFloatN(path, key, i, 32)
		*v = float32(f)

	case *float64:
		*v, err = subSegToFloatN(path, key, i, 64)

	case *int:
		var n int64
		n, err = subSegToIntN(path, key, i, 0)
		*v = int(n)

	case *int16:
		var n int64
		n, err = subSegToIntN(path, key, i, 16)
		*v = int16(n)

	case *int32:
		var n int64
		n, err = subSegToIntN(path, key, i, 32)
		*v = int32(n)

	case *int64:
		*v, err = subSegToIntN(path, key, i, 64)

	case *int8:
		var n int64
		n, err = subSegToIntN(path, key, i, 8)
		*v = int8(n)

	case *string:
		*v, err = subSegToString(path, key, i)

	case *uint:
		var n uint64
		n, err = subSegToUintN(path, key, i, 0)
		*v = uint(n)

	case *uint16:
		var n uint64
		n, err = subSegToUintN(path, key, i, 16)
		*v = uint16(n)

	case *uint32:
		var n uint64
		n, err = subSegToUintN(path, key, i, 32)
		*v = uint32(n)

	case *uint64:
		*v, err = subSegToUintN(path, key, i, 64)

	case *uint8:
		var n uint64
		n, err = subSegToUintN(path, key, i, 8)
		*v = uint8(n)

	case Unmarshaler:
		var s string
		s, err = subSegToString(path, key, i)
		if err == nil {
			err = v.UnmarshalSegment(s)
		}

	default:
		err = ErrUnknownType
	}

	return err
}

// SubSpan is similar to Span, but only handles the portion of the path
// subsequent to the provided key.
func SubSpan(path, key string, i, j int) (string, error) {
	si, ok := segIndexByKey(path, key)
	if !ok {
		return "", ErrKeySegNotFound
	}

	if i >= 0 {
		i++
	}
	if j > 0 {
		j++
	}

	s, err := Span(path[si:], i, j)
	if err != nil {
		return "", err
	}

	return s, nil
}

// Parth manages path and error data for processing a single path multiple
// times while error checking only once. Only the first encountered error is
// stored as all subsequent calls to Parth methods that can error are elided.
type Parth struct {
	path string
	err  error
}

// New constructs a pointer to an instance of Parth around the provided path.
func New(path string) *Parth {
	return &Parth{path: path}
}

// NewBySpan constructs a pointer to an instance of Parth after preprocessing
// the provided path with Span.
func NewBySpan(path string, i, j int) *Parth {
	s, err := Span(path, i, j)
	return &Parth{s, err}
}

// NewBySubSpan constructs a pointer to an instance of Parth after
// preprocessing the provided path with Span.
func NewBySubSpan(path, key string, i, j int) *Parth {
	s, err := SubSpan(path, key, i, j)
	return &Parth{s, err}
}

// Err returns the first error encountered by the *Parth receiver.
func (p *Parth) Err() error {
	return p.err
}

// Segment operates the same as the package-level function Segment.
func (p *Parth) Segment(i int, v interface{}) {
	if p.err != nil {
		return
	}

	p.err = Segment(p.path, i, v)
}

// Sequent operates the same as the package-level function Sequent.
func (p *Parth) Sequent(key string, v interface{}) {
	p.SubSeg(key, 0, v)
}

// Span operates the same as the package-level function Span.
func (p *Parth) Span(i, j int) string {
	if p.err != nil {
		return ""
	}

	s, err := Span(p.path, i, j)
	p.err = err

	return s
}

// SubSeg operates the same as the package-level function SubSeg.
func (p *Parth) SubSeg(key string, i int, v interface{}) {
	if p.err != nil {
		return
	}

	p.err = SubSeg(p.path, key, i, v)
}

// SubSpan operates the same as the package-level function SubSpan.
func (p *Parth) SubSpan(key string, i, j int) string {
	if p.err != nil {
		return ""
	}

	s, err := SubSpan(p.path, key, i, j)
	p.err = err

	return s
}
<commit_msg>Clarify/correct docs.<commit_after>// Package parth provides path parsing for segment unmarshaling and slicing.
//
// Along with string, all basic non-alias types are supported. An interface is
// available for implementation by user-defined types. When handling an int,
// uint, or float of any size, the first valid value within the specified
// segment will be used.
package parth

import (
	"errors"
)

// Unmarshaler is the interface implemented by types that can unmarshal a path
// segment representation of themselves. It is safe to assume that the segment
// data will not include slashes.
type Unmarshaler interface {
	UnmarshalSegment(string) error
}

// Err{Name} values facilitate error identification.
var (
	ErrUnknownType = errors.New("unknown type provided")

	ErrFirstSegNotFound = errors.New("first segment not found by index")
	ErrLastSegNotFound  = errors.New("last segment not found by index")
	ErrSegOrderReversed = errors.New("first segment must precede last segment")
	ErrKeySegNotFound   = errors.New("segment not found by key")

	ErrDataUnparsable = errors.New("data cannot be parsed")
)

// Segment locates the path segment indicated by the index i and unmarshals it
// into the provided type v. An error is returned if: 1. The type is not a
// pointer to an instance of one of the basic non-alias types and does not
// implement the Unmarshaler interface; 2. The index is out of range of the
// path; 3. The located path segment data cannot be parsed as the provided type
// or if an error is returned when using a provided Unmarshaler implementation.
func Segment(path string, i int, v interface{}) error { //nolint
	var err error

	switch v := v.(type) {
	case *bool:
		*v, err = segmentToBool(path, i)

	case *float32:
		var f float64
		f, err = segmentToFloatN(path, i, 32)
		*v = float32(f)

	case *float64:
		*v, err = segmentToFloatN(path, i, 64)

	case *int:
		var n int64
		n, err = segmentToIntN(path, i, 0)
		*v = int(n)

	case *int16:
		var n int64
		n, err = segmentToIntN(path, i, 16)
		*v = int16(n)

	case *int32:
		var n int64
		n, err = segmentToIntN(path, i, 32)
		*v = int32(n)

	case *int64:
		*v, err = segmentToIntN(path, i, 64)

	case *int8:
		var n int64
		n, err = segmentToIntN(path, i, 8)
		*v = int8(n)

	case *string:
		*v, err = segmentToString(path, i)

	case *uint:
		var n uint64
		n, err = segmentToUintN(path, i, 0)
		*v = uint(n)

	case *uint16:
		var n uint64
		n, err = segmentToUintN(path, i, 16)
		*v = uint16(n)

	case *uint32:
		var n uint64
		n, err = segmentToUintN(path, i, 32)
		*v = uint32(n)

	case *uint64:
		*v, err = segmentToUintN(path, i, 64)

	case *uint8:
		var n uint64
		n, err = segmentToUintN(path, i, 8)
		*v = uint8(n)

	case Unmarshaler:
		var s string
		s, err = segmentToString(path, i)
		if err == nil {
			err = v.UnmarshalSegment(s)
		}

	default:
		err = ErrUnknownType
	}

	return err
}

// Sequent is similar to Segment, but uses a key to locate a segment and then
// unmarshal the subsequent segment. It is a simple wrapper over SubSeg with an
// index of 0.
func Sequent(path, key string, v interface{}) error {
	return SubSeg(path, key, 0, v)
}

// Span returns the path segments between two segment indexes i and j including
// the first segment. An error is returned if: 1. Either index is out of range
// of the path; 2. The first index i does not precede the last index j.
// Providing a 0 for the last index is a special case which acts as an alias
// for the end of the path.
func Span(path string, i, j int) (string, error) {
	var f, l int
	var ok bool

	if i < 0 {
		f, ok = segStartIndexFromEnd(path, i)
	} else {
		f, ok = segStartIndexFromStart(path, i)
	}
	if !ok {
		return "", ErrFirstSegNotFound
	}

	if j > 0 {
		l, ok = segEndIndexFromStart(path, j)
	} else {
		l, ok = segEndIndexFromEnd(path, j)
	}
	if !ok {
		return "", ErrLastSegNotFound
	}

	if f == l {
		return "", nil
	}

	if f > l {
		return "", ErrSegOrderReversed
	}

	return path[f:l], nil
}

// SubSeg is similar to Segment, but only handles the portion of the path
// subsequent to the provided key. For example, to access the segment
// immediately after a key, an index of 0 should be provided (see Sequent).
func SubSeg(path, key string, i int, v interface{}) error { //nolint
	var err error

	switch v := v.(type) {
	case *bool:
		*v, err = subSegToBool(path, key, i)

	case *float32:
		var f float64
		f, err = subSegToFloatN(path, key, i, 32)
		*v = float32(f)

	case *float64:
		*v, err = subSegToFloatN(path, key, i, 64)

	case *int:
		var n int64
		n, err = subSegToIntN(path, key, i, 0)
		*v = int(n)

	case *int16:
		var n int64
		n, err = subSegToIntN(path, key, i, 16)
		*v = int16(n)

	case *int32:
		var n int64
		n, err = subSegToIntN(path, key, i, 32)
		*v = int32(n)

	case *int64:
		*v, err = subSegToIntN(path, key, i, 64)

	case *int8:
		var n int64
		n, err = subSegToIntN(path, key, i, 8)
		*v = int8(n)

	case *string:
		*v, err = subSegToString(path, key, i)

	case *uint:
		var n uint64
		n, err = subSegToUintN(path, key, i, 0)
		*v = uint(n)

	case *uint16:
		var n uint64
		n, err = subSegToUintN(path, key, i, 16)
		*v = uint16(n)

	case *uint32:
		var n uint64
		n, err = subSegToUintN(path, key, i, 32)
		*v = uint32(n)

	case *uint64:
		*v, err = subSegToUintN(path, key, i, 64)

	case *uint8:
		var n uint64
		n, err = subSegToUintN(path, key, i, 8)
		*v = uint8(n)

	case Unmarshaler:
		var s string
		s, err = subSegToString(path, key, i)
		if err == nil {
			err = v.UnmarshalSegment(s)
		}

	default:
		err = ErrUnknownType
	}

	return err
}

// SubSpan is similar to Span, but only handles the portion of the path
// subsequent to the provided key.
func SubSpan(path, key string, i, j int) (string, error) {
	si, ok := segIndexByKey(path, key)
	if !ok {
		return "", ErrKeySegNotFound
	}

	if i >= 0 {
		i++
	}
	if j > 0 {
		j++
	}

	s, err := Span(path[si:], i, j)
	if err != nil {
		return "", err
	}

	return s, nil
}

// Parth manages path and error data for processing a single path multiple
// times while error checking only once. Only the first encountered error is
// stored as all subsequent calls to Parth methods that can error are elided.
type Parth struct {
	path string
	err  error
}

// New constructs a pointer to an instance of Parth around the provided path.
func New(path string) *Parth {
	return &Parth{path: path}
}

// NewBySpan constructs a pointer to an instance of Parth after preprocessing
// the provided path with Span.
func NewBySpan(path string, i, j int) *Parth {
	s, err := Span(path, i, j)
	return &Parth{s, err}
}

// NewBySubSpan constructs a pointer to an instance of Parth after
// preprocessing the provided path with SubSpan.
func NewBySubSpan(path, key string, i, j int) *Parth {
	s, err := SubSpan(path, key, i, j)
	return &Parth{s, err}
}

// Err returns the first error encountered by the *Parth receiver.
func (p *Parth) Err() error {
	return p.err
}

// Segment operates the same as the package-level function Segment.
func (p *Parth) Segment(i int, v interface{}) {
	if p.err != nil {
		return
	}

	p.err = Segment(p.path, i, v)
}

// Sequent operates the same as the package-level function Sequent.
func (p *Parth) Sequent(key string, v interface{}) {
	p.SubSeg(key, 0, v)
}

// Span operates the same as the package-level function Span.
func (p *Parth) Span(i, j int) string {
	if p.err != nil {
		return ""
	}

	s, err := Span(p.path, i, j)
	p.err = err

	return s
}

// SubSeg operates the same as the package-level function SubSeg.
func (p *Parth) SubSeg(key string, i int, v interface{}) {
	if p.err != nil {
		return
	}

	p.err = SubSeg(p.path, key, i, v)
}

// SubSpan operates the same as the package-level function SubSpan.
func (p *Parth) SubSpan(key string, i, j int) string {
	if p.err != nil {
		return ""
	}

	s, err := SubSpan(p.path, key, i, j)
	p.err = err

	return s
}
<|endoftext|>


<commit_before>/* Copyright (c) 2014, Daniel Martí <mvdan@mvdan.cc> */
/* See LICENSE for licensing information */

package main

import (
	"compress/zlib"
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const (
	idSize    = 8 // should be between 6 and 256
	indexTmpl = "index.html"
	maxSize   = 1 << 20 // whole POST body

	// GET error messages
	invalidId     = "Invalid paste id."
	pasteNotFound = "Paste doesn't exist."
	unknownError  = "Something went terribly wrong."
	// POST error messages
	missingForm = "Paste could not be found inside the posted form."
)

var siteUrl = flag.String("u", "http://localhost:9090", "URL of the site")
var listen = flag.String("l", "localhost:9090", "Host and port to listen to")
var dataDir = flag.String("d", "data", "Directory to store all the pastes in")
var lifeTimeStr = flag.String("t", "12h", "Lifetime of the pastes (units: s,m,h)")
var lifeTime time.Duration

const chars = "abcdefghijklmnopqrstuvwxyz0123456789"

var validId *regexp.Regexp = regexp.MustCompile("^[a-zA-Z0-9]{" + strconv.FormatInt(idSize, 10) + "}$")

var indexTemplate *template.Template

func pathId(id string) string {
	return path.Join(id[0:2], id[2:4], id[4:])
}

const (
	_        = iota
	KB int64 = 1 << (10 * iota)
	MB
)

func readableSize(b int64) string {
	switch {
	case b >= MB:
		return fmt.Sprintf("%.2fMB", float64(b)/float64(MB))
	case b >= KB:
		return fmt.Sprintf("%.2fKB", float64(b)/float64(KB))
	}
	return fmt.Sprintf("%dB", b)
}

func randomId() string {
	s := make([]byte, idSize)
	var offset uint = 0
	for {
		r := rand.Int63()
		for i := 0; i < 8; i++ {
			randbyte := int(r&0xff) % len(chars)
			s[offset] = chars[randbyte]
			offset++
			if offset == idSize {
				return string(s)
			}
			r >>= 8
		}
	}
	return strings.Repeat(chars[0:1], idSize)
}

func endLife(path string) {
	err := os.Remove(path)
	if err == nil {
		log.Printf("Removed paste: %s", path)
	} else {
		log.Printf("Could not end the life of %s: %s", path, err)
		programDeath(path, 2*time.Minute)
	}
}

func programDeath(path string, after time.Duration) {
	timer := time.NewTimer(after)
	go func() {
		<-timer.C
		endLife(path)
	}()
}

func handler(w http.ResponseWriter, r *http.Request) {
	var err error
	switch r.Method {
	case "GET":
		var id, pastePath string
		id = r.URL.Path[1:]
		if len(id) == 0 {
			indexTemplate.Execute(w, *siteUrl)
			return
		}
		if !validId.MatchString(id) {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", invalidId)
			return
		}
		id = strings.ToLower(id)
		pastePath = pathId(id)
		pasteFile, err := os.Open(pastePath)
		if err != nil {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "%s\n", pasteNotFound)
			return
		}
		compReader, err := zlib.NewReader(pasteFile)
		if err != nil {
			log.Printf("Could not open a compression reader for %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		io.Copy(w, compReader)
		compReader.Close()
		pasteFile.Close()

	case "POST":
		r.Body = http.MaxBytesReader(w, r.Body, maxSize)
		var id, pastePath string
		for {
			id = randomId()
			pastePath = pathId(id)
			if _, err := os.Stat(pastePath); os.IsNotExist(err) {
				break
			}
		}
		if err = r.ParseMultipartForm(maxSize); err != nil {
			log.Printf("Could not parse POST multipart form: %s", err)
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", err)
			return
		}
		var content string
		if vs, found := r.Form["paste"]; found {
			content = vs[0]
		} else {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", missingForm)
			return
		}
		dir, _ := path.Split(pastePath)
		if err = os.MkdirAll(dir, 0700); err != nil {
			log.Printf("Could not create directories leading to %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		programDeath(pastePath, lifeTime)
		pasteFile, err := os.OpenFile(pastePath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
		if err != nil {
			log.Printf("Could not create new paste pasteFile %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		compWriter := zlib.NewWriter(pasteFile)
		b, err := io.WriteString(compWriter, content)
		compWriter.Close()
		pasteFile.Close()
		if err != nil {
			log.Printf("Could not write compressed data into %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		log.Printf("Created a new paste: %s (%s)", pastePath, readableSize(int64(b)))
		fmt.Fprintf(w, "%s/%s\n", *siteUrl, id)
	}
}

func walkFunc(path string, info os.FileInfo, err error) error {
	if err != nil {
		return err
	}
	if info.IsDir() {
		return nil
	}
	deathTime := info.ModTime().Add(lifeTime)
	now := time.Now()
	if deathTime.Before(now) {
		go endLife(path)
		return nil
	}
	var lifeLeft time.Duration
	if deathTime.After(now.Add(lifeTime)) {
		lifeLeft = lifeTime
	} else {
		lifeLeft = deathTime.Sub(now)
	}
	log.Printf("Recovered paste %s has %s left", path, lifeLeft)
	programDeath(path, lifeLeft)
	return nil
}

func main() {
	var err error
	flag.Parse()
	if lifeTime, err = time.ParseDuration(*lifeTimeStr); err != nil {
		log.Printf("Invalid lifetime '%s': %s", lifeTimeStr, err)
		return
	}
	if indexTemplate, err = template.ParseFiles(indexTmpl); err != nil {
		log.Printf("Could not load template %s: %s", indexTmpl, err)
		return
	}
	if err = os.MkdirAll(*dataDir, 0700); err != nil {
		log.Printf("Could not create data directory %s: %s", *dataDir, err)
		return
	}
	if err = os.Chdir(*dataDir); err != nil {
		log.Printf("Could not enter data directory %s: %s", *dataDir, err)
		return
	}
	if err = filepath.Walk(".", walkFunc); err != nil {
		log.Printf("Could not recover data directory %s: %s", *dataDir, err)
		return
	}
	log.Printf("idSize   = %d", idSize)
	log.Printf("maxSize  = %s", readableSize(maxSize))
	log.Printf("siteUrl  = %s", *siteUrl)
	log.Printf("listen   = %s", *listen)
	log.Printf("dataDir  = %s", *dataDir)
	log.Printf("lifeTime = %s", lifeTime)
	http.HandleFunc("/", handler)
	log.Fatal(http.ListenAndServe(*listen, nil))
}
<commit_msg>Add -s to specify the maximum size<commit_after>/* Copyright (c) 2014, Daniel Martí <mvdan@mvdan.cc> */
/* See LICENSE for licensing information */

package main

import (
	"compress/zlib"
	"errors"
	"flag"
	"fmt"
	"html/template"
	"io"
	"log"
	"math/rand"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const (
	idSize    = 8 // should be between 6 and 256
	indexTmpl = "index.html"
	chars     = "abcdefghijklmnopqrstuvwxyz0123456789"

	// GET error messages
	invalidId     = "Invalid paste id."
	pasteNotFound = "Paste doesn't exist."
	unknownError  = "Something went terribly wrong."
	// POST error messages
	missingForm = "Paste could not be found inside the posted form."
)

var (
	siteUrl     = flag.String("u", "http://localhost:9090", "URL of the site")
	listen      = flag.String("l", "localhost:9090", "Host and port to listen to")
	dataDir     = flag.String("d", "data", "Directory to store all the pastes in")
	lifeTimeStr = flag.String("t", "12h", "Lifetime of the pastes (units: s,m,h)")
	maxSizeStr  = flag.String("s", "1M", "Maximum size of POSTs in bytes (units: B,K,M)")

	lifeTime time.Duration
	maxSize  ByteSize

	validId       = regexp.MustCompile("^[a-zA-Z0-9]{" + strconv.FormatInt(idSize, 10) + "}$")
	regexByteSize = regexp.MustCompile(`^([\d\.]+)\s*([KM]?B|[BKM])$`)
	indexTemplate *template.Template
)

func pathId(id string) string {
	return path.Join(id[0:2], id[2:4], id[4:])
}

type ByteSize int64

const (
	B ByteSize = 1 << (10 * iota)
	KB
	MB
)

func (b ByteSize) String() string {
	switch {
	case b >= MB:
		return fmt.Sprintf("%.2f MB", float64(b)/float64(MB))
	case b >= KB:
		return fmt.Sprintf("%.2f KB", float64(b)/float64(KB))
	}
	return fmt.Sprintf("%d B", b)
}

func parseByteSize(str string) (ByteSize, error) {
	if !regexByteSize.MatchString(str) {
		return 0, errors.New("Could not parse size in bytes")
	}
	parts := regexByteSize.FindStringSubmatch(str)
	size, _ := strconv.ParseFloat(string(parts[1]), 64)

	switch string(parts[2]) {
	case "KB", "K":
		size *= float64(KB)
	case "MB", "M":
		size *= float64(MB)
	}
	return ByteSize(size), nil
}

func randomId() string {
	s := make([]byte, idSize)
	var offset uint = 0
	for {
		r := rand.Int63()
		for i := 0; i < 8; i++ {
			randbyte := int(r&0xff) % len(chars)
			s[offset] = chars[randbyte]
			offset++
			if offset == idSize {
				return string(s)
			}
			r >>= 8
		}
	}
	return strings.Repeat(chars[0:1], idSize)
}

func endLife(path string) {
	err := os.Remove(path)
	if err == nil {
		log.Printf("Removed paste: %s", path)
	} else {
		log.Printf("Could not end the life of %s: %s", path, err)
		programDeath(path, 2*time.Minute)
	}
}

func programDeath(path string, after time.Duration) {
	timer := time.NewTimer(after)
	go func() {
		<-timer.C
		endLife(path)
	}()
}

func handler(w http.ResponseWriter, r *http.Request) {
	var err error
	switch r.Method {
	case "GET":
		var id, pastePath string
		id = r.URL.Path[1:]
		if len(id) == 0 {
			indexTemplate.Execute(w, *siteUrl)
			return
		}
		if !validId.MatchString(id) {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", invalidId)
			return
		}
		id = strings.ToLower(id)
		pastePath = pathId(id)
		pasteFile, err := os.Open(pastePath)
		if err != nil {
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "%s\n", pasteNotFound)
			return
		}
		compReader, err := zlib.NewReader(pasteFile)
		if err != nil {
			log.Printf("Could not open a compression reader for %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		io.Copy(w, compReader)
		compReader.Close()
		pasteFile.Close()

	case "POST":
		r.Body = http.MaxBytesReader(w, r.Body, int64(maxSize))
		var id, pastePath string
		for {
			id = randomId()
			pastePath = pathId(id)
			if _, err := os.Stat(pastePath); os.IsNotExist(err) {
				break
			}
		}
		if err = r.ParseMultipartForm(int64(maxSize)); err != nil {
			log.Printf("Could not parse POST multipart form: %s", err)
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", err)
			return
		}
		var content string
		if vs, found := r.Form["paste"]; found {
			content = vs[0]
		} else {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s\n", missingForm)
			return
		}
		dir, _ := path.Split(pastePath)
		if err = os.MkdirAll(dir, 0700); err != nil {
			log.Printf("Could not create directories leading to %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		programDeath(pastePath, lifeTime)
		pasteFile, err := os.OpenFile(pastePath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
		if err != nil {
			log.Printf("Could not create new paste pasteFile %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		compWriter := zlib.NewWriter(pasteFile)
		b, err := io.WriteString(compWriter, content)
		compWriter.Close()
		pasteFile.Close()
		if err != nil {
			log.Printf("Could not write compressed data into %s: %s", pastePath, err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "%s\n", unknownError)
			return
		}
		writtenSize := ByteSize(b)
		log.Printf("Created a new paste: %s (%s)", pastePath, writtenSize)
		fmt.Fprintf(w, "%s/%s\n", *siteUrl, id)
	}
}

func walkFunc(path string, info os.FileInfo, err error) error {
	if err != nil {
		return err
	}
	if info.IsDir() {
		return nil
	}
	deathTime := info.ModTime().Add(lifeTime)
	now := time.Now()
	if deathTime.Before(now) {
		go endLife(path)
		return nil
	}
	var lifeLeft time.Duration
	if deathTime.After(now.Add(lifeTime)) {
		lifeLeft = lifeTime
	} else {
		lifeLeft = deathTime.Sub(now)
	}
	log.Printf("Recovered paste %s has %s left", path, lifeLeft)
	programDeath(path, lifeLeft)
	return nil
}

func main() {
	var err error
	flag.Parse()
	if lifeTime, err = time.ParseDuration(*lifeTimeStr); err != nil {
		log.Printf("Invalid lifetime '%s': %s", lifeTimeStr, err)
		return
	}
	if maxSize, err = parseByteSize(*maxSizeStr); err != nil {
		log.Printf("Invalid max size '%s': %s", maxSizeStr, err)
		return
	}
	if indexTemplate, err = template.ParseFiles(indexTmpl); err != nil {
		log.Printf("Could not load template %s: %s", indexTmpl, err)
		return
	}
	if err = os.MkdirAll(*dataDir, 0700); err != nil {
		log.Printf("Could not create data directory %s: %s", *dataDir, err)
		return
	}
	if err = os.Chdir(*dataDir); err != nil {
		log.Printf("Could not enter data directory %s: %s", *dataDir, err)
		return
	}
	if err = filepath.Walk(".", walkFunc); err != nil {
		log.Printf("Could not recover data directory %s: %s", *dataDir, err)
		return
	}
	log.Printf("idSize   = %d", idSize)
	log.Printf("maxSize  = %s", maxSize)
	log.Printf("siteUrl  = %s", *siteUrl)
	log.Printf("listen   = %s", *listen)
	log.Printf("dataDir  = %s", *dataDir)
	log.Printf("lifeTime = %s", lifeTime)
	http.HandleFunc("/", handler)
	log.Fatal(http.ListenAndServe(*listen, nil))
}
<|endoftext|>


<commit_before>package endtoend_test

import (
    "flag"
	. "launchpad.net/gocheck"
    "launchpad.net/goamz/aws"
    "launchpad.net/goamz/ec2"
	"testing"
    "time"
)

func Test(t *testing.T) { TestingT(t) }

type S struct {
    vm *VM
}

var _ = Suite(&S{})

var flagDesc = "enable end-to-end tests that creates a machine in amazon, you'll need a AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY to run this tests."
var enableSuite = flag.Bool("endtoend", false, flagDesc)

type VM struct {
    instanceId string
    ec2 *ec2.EC2
}

func (s *S) stopOnStateChange(toState string, c *C) {
    ticker := time.Tick(time.Minute)
    for _ = range ticker {
        instResp, err := s.vm.ec2.Instances([]string{s.vm.instanceId}, nil)
        c.Check(err, IsNil)
        state := instResp.Reservations[0].Instances[0].State
        if state.Name == toState {
            break
        }
    }
}

func (s *S) newVM(c *C) {
    auth, err := aws.EnvAuth()
    c.Check(err, IsNil)
    e := ec2.New(auth, aws.USEast)
    s.vm = &VM{ec2: e}
    options := ec2.RunInstances{
        ImageId: "ami-ccf405a5", // ubuntu maverik
        InstanceType: "t1.micro",
    }
    resp, err := e.RunInstances(&options)
    c.Check(err, IsNil)
    instanceId := resp.Instances[0].InstanceId
    s.vm.instanceId = instanceId
    // wait until instance is up
    s.stopOnStateChange("running", c)
}

func (s *S) destroyVM(c *C) {
    _, err := s.vm.ec2.TerminateInstances([]string{s.vm.instanceId})
    c.Check(err, IsNil)
    s.stopOnStateChange("terminated", c)
}

func (s *S) SetUpSuite(c *C) {
    if !*enableSuite {
        c.Skip("skipping end-to-end suite, use -endtoend to enable")
    }
    s.newVM(c)
}

func (s *S) TearDown(c *C) {
    s.destroyVM(c)
}

func (s *S) TestTrueIsTrue(c *C) {
    c.Assert(true, Equals, true)
}
<commit_msg>go fmt<commit_after>package endtoend_test

import (
	"flag"
	"launchpad.net/goamz/aws"
	"launchpad.net/goamz/ec2"
	. "launchpad.net/gocheck"
	"testing"
	"time"
)

func Test(t *testing.T) { TestingT(t) }

type S struct {
	vm *VM
}

var _ = Suite(&S{})

var flagDesc = "enable end-to-end tests that creates a machine in amazon, you'll need a AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY to run this tests."
var enableSuite = flag.Bool("endtoend", false, flagDesc)

type VM struct {
	instanceId string
	ec2        *ec2.EC2
}

func (s *S) stopOnStateChange(toState string, c *C) {
	ticker := time.Tick(time.Minute)
	for _ = range ticker {
		instResp, err := s.vm.ec2.Instances([]string{s.vm.instanceId}, nil)
		c.Check(err, IsNil)
		state := instResp.Reservations[0].Instances[0].State
		if state.Name == toState {
			break
		}
	}
}

func (s *S) newVM(c *C) {
	auth, err := aws.EnvAuth()
	c.Check(err, IsNil)
	e := ec2.New(auth, aws.USEast)
	s.vm = &VM{ec2: e}
	options := ec2.RunInstances{
		ImageId:      "ami-ccf405a5", // ubuntu maverik
		InstanceType: "t1.micro",
	}
	resp, err := e.RunInstances(&options)
	c.Check(err, IsNil)
	instanceId := resp.Instances[0].InstanceId
	s.vm.instanceId = instanceId
	// wait until instance is up
	s.stopOnStateChange("running", c)
}

func (s *S) destroyVM(c *C) {
	_, err := s.vm.ec2.TerminateInstances([]string{s.vm.instanceId})
	c.Check(err, IsNil)
	s.stopOnStateChange("terminated", c)
}

func (s *S) SetUpSuite(c *C) {
	if !*enableSuite {
		c.Skip("skipping end-to-end suite, use -endtoend to enable")
	}
	s.newVM(c)
}

func (s *S) TearDown(c *C) {
	s.destroyVM(c)
}

func (s *S) TestTrueIsTrue(c *C) {
	c.Assert(true, Equals, true)
}
<|endoftext|>


<commit_before>package core

import (
	"context"
	"time"

	"chain/core/query"
	"chain/core/query/chql"
	"chain/cos/bc"
	"chain/database/pg"
	"chain/errors"
	"chain/net/http/httpjson"
)

var (
	ErrBadIndexConfig = errors.New("index configuration invalid")
)

// createIndex is an http handler for creating indexes.
//
// POST /create-index
func (a *api) createIndex(ctx context.Context, in struct {
	Alias    string `json:"alias"`
	Type     string `json:"type"`
	Query    string `json:"query"`
	Unspents bool   `json:"unspents"`
}) (*query.Index, error) {
	if !query.IndexTypes[in.Type] {
		return nil, errors.WithDetailf(ErrBadIndexConfig, "unknown index type %q", in.Type)
	}
	if in.Unspents && in.Type != query.IndexTypeBalance {
		return nil, errors.WithDetail(ErrBadIndexConfig, "unspents flag is only valid for balance indexes")
	}
	if in.Alias == "" {
		return nil, errors.WithDetail(httpjson.ErrBadRequest, "missing index alias")
	}

	idx, err := a.indexer.CreateIndex(ctx, in.Alias, in.Type, in.Query, in.Unspents)
	return idx, errors.Wrap(err, "creating the new index")
}

// listIndexes is an http handler for listing ChQL indexes.
//
// POST /list-indexes
func (a *api) listIndexes(ctx context.Context, query requestQuery) (page, error) {
	limit := defGenericPageSize

	indexes, cursor, err := a.indexer.ListIndexes(ctx, query.Cursor, limit)
	if err != nil {
		return page{}, errors.Wrap(err, "listing indexes")
	}

	query.Cursor = cursor
	return page{
		Items:    httpjson.Array(indexes),
		LastPage: len(indexes) < limit,
		Query:    query,
	}, nil
}

var (
	ErrNeitherIndexNorQuery = errors.New("must provide either index or query")
	ErrBothIndexAndQuery    = errors.New("cannot provide both index and query")
)

// listTransactions is an http handler for listing transactions matching
// a ChQL query or index.
//
// POST /list-transactions
func (a *api) listTransactions(ctx context.Context, in requestQuery) (result page, err error) {
	if (in.IndexID != "" || in.IndexAlias != "") && in.ChQL != "" {
		return result, errors.WithDetail(httpjson.ErrBadRequest, "cannot provide both index and query")
	}
	if in.EndTimeMS == 0 {
		in.EndTimeMS = bc.Millis(time.Now())
	}

	var (
		q   chql.Query
		cur query.TxCursor
	)

	// Build the ChQL query
	if in.IndexAlias != "" || in.IndexID != "" {
		idx, err := a.indexer.GetIndex(ctx, in.IndexID, in.IndexAlias, query.IndexTypeTransaction)
		if err != nil {
			return result, err
		}
		if idx == nil {
			return result, errors.WithDetail(pg.ErrUserInputNotFound, "transaction index not found")
		}
		q = idx.Query
	} else {
		q, err = chql.Parse(in.ChQL)
		if err != nil {
			return result, err
		}
	}

	// Either parse the provided cursor or look one up for the time range.
	if in.Cursor != "" {
		cur, err = query.DecodeTxCursor(in.Cursor)
		if err != nil {
			return result, errors.Wrap(err, "decoding cursor")
		}
	} else {
		cur, err = a.indexer.LookupTxCursor(ctx, in.StartTimeMS, in.EndTimeMS)
		if err != nil {
			return result, err
		}
	}

	limit := defGenericPageSize
	txns, nextCur, err := a.indexer.Transactions(ctx, q, in.ChQLParams, cur, limit)
	if err != nil {
		return result, errors.Wrap(err, "running tx query")
	}

	out := in
	out.Cursor = nextCur.String()
	return page{
		Items:    httpjson.Array(txns),
		LastPage: len(txns) < limit,
		Query:    out,
	}, nil
}

// listAccounts is an http handler for listing accounts matching
// a ChQL query or index.
//
// POST /list-accounts
func (a *api) listAccounts(ctx context.Context, in requestQuery) (page, error) {
	limit := defGenericPageSize

	// Build the ChQL query
	q, err := chql.Parse(in.ChQL)
	if err != nil {
		return page{}, errors.Wrap(err, "parsing acc query")
	}
	cur := in.Cursor

	// Use the ChQL query engine for querying account tags.
	accounts, cur, err := a.indexer.Accounts(ctx, q, in.ChQLParams, cur, limit)
	if err != nil {
		return page{}, errors.Wrap(err, "running acc query")
	}

	// Pull in the accounts by the IDs
	out := in
	out.Cursor = cur
	return page{
		Items:    httpjson.Array(accounts),
		LastPage: len(accounts) < limit,
		Query:    out,
	}, nil
}

// POST /list-balances
func (a *api) listBalances(ctx context.Context, in requestQuery) (result page, err error) {
	if (in.IndexID != "" || in.IndexAlias != "") && in.ChQL != "" {
		return result, errors.WithDetail(httpjson.ErrBadRequest, "cannot provide both index and query")
	}
	if in.TimestampMS == 0 {
		in.TimestampMS = bc.Millis(time.Now())
	}

	var q chql.Query
	if in.IndexID != "" || in.IndexAlias != "" {
		idx, err := a.indexer.GetIndex(ctx, in.IndexID, in.IndexAlias, query.IndexTypeBalance)
		if err != nil {
			return result, err
		}
		if idx == nil {
			return result, errors.WithDetail(pg.ErrUserInputNotFound, "balance index not found")
		}
		q = idx.Query
	} else {
		q, err = chql.Parse(in.ChQL)
		if err != nil {
			return result, err
		}
	}

	// TODO(jackson): paginate this endpoint.
	balances, err := a.indexer.Balances(ctx, q, in.ChQLParams, in.TimestampMS)
	if err != nil {
		return result, err
	}

	result.Items = httpjson.Array(balances)
	result.LastPage = true
	result.Query = in
	return result, nil
}

// POST /list-unspent-outputs
func (a *api) listUnspentOutputs(ctx context.Context, in requestQuery) (result page, err error) {
	if (in.IndexID != "" || in.IndexAlias != "") && in.ChQL != "" {
		return result, errors.WithDetail(httpjson.ErrBadRequest, "cannot provide both index and query")
	}

	var q chql.Query
	if in.IndexID != "" || in.IndexAlias != "" {
		idx, err := a.indexer.GetIndex(ctx, in.IndexID, in.IndexAlias, query.IndexTypeBalance)
		if err != nil {
			return result, err
		}
		if idx == nil {
			return result, errors.WithDetail(pg.ErrUserInputNotFound, "balance index not found")
		}
		if !idx.Unspents {
			return result, errors.WithDetail(httpjson.ErrBadRequest, "balance index doesn't support output indexing")
		}
		q = idx.Query
	} else {
		q, err = chql.Parse(in.ChQL)
		if err != nil {
			return result, err
		}
	}

	var cursor *query.OutputsCursor
	if in.Cursor != "" {
		cursor, err = query.DecodeOutputsCursor(in.Cursor)
		if err != nil {
			return result, errors.Wrap(err, "decoding cursor")
		}
	}

	limit := defGenericPageSize
	outputs, newCursor, err := a.indexer.Outputs(ctx, q, in.ChQLParams, in.TimestampMS, cursor, limit)
	if err != nil {
		return result, errors.Wrap(err, "querying outputs")
	}

	outQuery := in
	outQuery.Cursor = newCursor.String()
	return page{
		Items:    outputs,
		LastPage: len(outputs) < limit,
		Query:    outQuery,
	}, nil
}

// listAssets is an http handler for listing assets matching
// a ChQL query or index.
//
// POST /list-assets
func (a *api) listAssets(ctx context.Context, in requestQuery) (page, error) {
	limit := defGenericPageSize

	// Build the ChQL query
	q, err := chql.Parse(in.ChQL)
	if err != nil {
		return page{}, err
	}
	cur := in.Cursor

	// Use the ChQL query engine for querying asset tags.
	var assets []map[string]interface{}
	assets, cur, err = a.indexer.Assets(ctx, q, in.ChQLParams, cur, limit)
	if err != nil {
		return page{}, errors.Wrap(err, "running asset query")
	}

	out := in
	out.Cursor = cur
	return page{
		Items:    httpjson.Array(assets),
		LastPage: len(assets) < limit,
		Query:    out,
	}, nil
}
<commit_msg>core: add list-unspent-outputs default timestamp<commit_after>package core

import (
	"context"
	"time"

	"chain/core/query"
	"chain/core/query/chql"
	"chain/cos/bc"
	"chain/database/pg"
	"chain/errors"
	"chain/net/http/httpjson"
)

var (
	ErrBadIndexConfig = errors.New("index configuration invalid")
)

// createIndex is an http handler for creating indexes.
//
// POST /create-index
func (a *api) createIndex(ctx context.Context, in struct {
	Alias    string `json:"alias"`
	Type     string `json:"type"`
	Query    string `json:"query"`
	Unspents bool   `json:"unspents"`
}) (*query.Index, error) {
	if !query.IndexTypes[in.Type] {
		return nil, errors.WithDetailf(ErrBadIndexConfig, "unknown index type %q", in.Type)
	}
	if in.Unspents && in.Type != query.IndexTypeBalance {
		return nil, errors.WithDetail(ErrBadIndexConfig, "unspents flag is only valid for balance indexes")
	}
	if in.Alias == "" {
		return nil, errors.WithDetail(httpjson.ErrBadRequest, "missing index alias")
	}

	idx, err := a.indexer.CreateIndex(ctx, in.Alias, in.Type, in.Query, in.Unspents)
	return idx, errors.Wrap(err, "creating the new index")
}

// listIndexes is an http handler for listing ChQL indexes.
//
// POST /list-indexes
func (a *api) listIndexes(ctx context.Context, query requestQuery) (page, error) {
	limit := defGenericPageSize

	indexes, cursor, err := a.indexer.ListIndexes(ctx, query.Cursor, limit)
	if err != nil {
		return page{}, errors.Wrap(err, "listing indexes")
	}

	query.Cursor = cursor
	return page{
		Items:    httpjson.Array(indexes),
		LastPage: len(indexes) < limit,
		Query:    query,
	}, nil
}

var (
	ErrNeitherIndexNorQuery = errors.New("must provide either index or query")
	ErrBothIndexAndQuery    = errors.New("cannot provide both index and query")
)

// listTransactions is an http handler for listing transactions matching
// a ChQL query or index.
//
// POST /list-transactions
func (a *api) listTransactions(ctx context.Context, in requestQuery) (result page, err error) {
	if (in.IndexID != "" || in.IndexAlias != "") && in.ChQL != "" {
		return result, errors.WithDetail(httpjson.ErrBadRequest, "cannot provide both index and query")
	}
	if in.EndTimeMS == 0 {
		in.EndTimeMS = bc.Millis(time.Now())
	}

	var (
		q   chql.Query
		cur query.TxCursor
	)

	// Build the ChQL query
	if in.IndexAlias != "" || in.IndexID != "" {
		idx, err := a.indexer.GetIndex(ctx, in.IndexID, in.IndexAlias, query.IndexTypeTransaction)
		if err != nil {
			return result, err
		}
		if idx == nil {
			return result, errors.WithDetail(pg.ErrUserInputNotFound, "transaction index not found")
		}
		q = idx.Query
	} else {
		q, err = chql.Parse(in.ChQL)
		if err != nil {
			return result, err
		}
	}

	// Either parse the provided cursor or look one up for the time range.
	if in.Cursor != "" {
		cur, err = query.DecodeTxCursor(in.Cursor)
		if err != nil {
			return result, errors.Wrap(err, "decoding cursor")
		}
	} else {
		cur, err = a.indexer.LookupTxCursor(ctx, in.StartTimeMS, in.EndTimeMS)
		if err != nil {
			return result, err
		}
	}

	limit := defGenericPageSize
	txns, nextCur, err := a.indexer.Transactions(ctx, q, in.ChQLParams, cur, limit)
	if err != nil {
		return result, errors.Wrap(err, "running tx query")
	}

	out := in
	out.Cursor = nextCur.String()
	return page{
		Items:    httpjson.Array(txns),
		LastPage: len(txns) < limit,
		Query:    out,
	}, nil
}

// listAccounts is an http handler for listing accounts matching
// a ChQL query or index.
//
// POST /list-accounts
func (a *api) listAccounts(ctx context.Context, in requestQuery) (page, error) {
	limit := defGenericPageSize

	// Build the ChQL query
	q, err := chql.Parse(in.ChQL)
	if err != nil {
		return page{}, errors.Wrap(err, "parsing acc query")
	}
	cur := in.Cursor

	// Use the ChQL query engine for querying account tags.
	accounts, cur, err := a.indexer.Accounts(ctx, q, in.ChQLParams, cur, limit)
	if err != nil {
		return page{}, errors.Wrap(err, "running acc query")
	}

	// Pull in the accounts by the IDs
	out := in
	out.Cursor = cur
	return page{
		Items:    httpjson.Array(accounts),
		LastPage: len(accounts) < limit,
		Query:    out,
	}, nil
}

// POST /list-balances
func (a *api) listBalances(ctx context.Context, in requestQuery) (result page, err error) {
	if (in.IndexID != "" || in.IndexAlias != "") && in.ChQL != "" {
		return result, errors.WithDetail(httpjson.ErrBadRequest, "cannot provide both index and query")
	}
	if in.TimestampMS == 0 {
		in.TimestampMS = bc.Millis(time.Now())
	}

	var q chql.Query
	if in.IndexID != "" || in.IndexAlias != "" {
		idx, err := a.indexer.GetIndex(ctx, in.IndexID, in.IndexAlias, query.IndexTypeBalance)
		if err != nil {
			return result, err
		}
		if idx == nil {
			return result, errors.WithDetail(pg.ErrUserInputNotFound, "balance index not found")
		}
		q = idx.Query
	} else {
		q, err = chql.Parse(in.ChQL)
		if err != nil {
			return result, err
		}
	}

	// TODO(jackson): paginate this endpoint.
	balances, err := a.indexer.Balances(ctx, q, in.ChQLParams, in.TimestampMS)
	if err != nil {
		return result, err
	}

	result.Items = httpjson.Array(balances)
	result.LastPage = true
	result.Query = in
	return result, nil
}

// POST /list-unspent-outputs
func (a *api) listUnspentOutputs(ctx context.Context, in requestQuery) (result page, err error) {
	if (in.IndexID != "" || in.IndexAlias != "") && in.ChQL != "" {
		return result, errors.WithDetail(httpjson.ErrBadRequest, "cannot provide both index and query")
	}
	if in.TimestampMS == 0 {
		in.TimestampMS = bc.Millis(time.Now())
	}

	var q chql.Query
	if in.IndexID != "" || in.IndexAlias != "" {
		idx, err := a.indexer.GetIndex(ctx, in.IndexID, in.IndexAlias, query.IndexTypeBalance)
		if err != nil {
			return result, err
		}
		if idx == nil {
			return result, errors.WithDetail(pg.ErrUserInputNotFound, "balance index not found")
		}
		if !idx.Unspents {
			return result, errors.WithDetail(httpjson.ErrBadRequest, "balance index doesn't support output indexing")
		}
		q = idx.Query
	} else {
		q, err = chql.Parse(in.ChQL)
		if err != nil {
			return result, err
		}
	}

	var cursor *query.OutputsCursor
	if in.Cursor != "" {
		cursor, err = query.DecodeOutputsCursor(in.Cursor)
		if err != nil {
			return result, errors.Wrap(err, "decoding cursor")
		}
	}

	limit := defGenericPageSize
	outputs, newCursor, err := a.indexer.Outputs(ctx, q, in.ChQLParams, in.TimestampMS, cursor, limit)
	if err != nil {
		return result, errors.Wrap(err, "querying outputs")
	}

	outQuery := in
	outQuery.Cursor = newCursor.String()
	return page{
		Items:    outputs,
		LastPage: len(outputs) < limit,
		Query:    outQuery,
	}, nil
}

// listAssets is an http handler for listing assets matching
// a ChQL query or index.
//
// POST /list-assets
func (a *api) listAssets(ctx context.Context, in requestQuery) (page, error) {
	limit := defGenericPageSize

	// Build the ChQL query
	q, err := chql.Parse(in.ChQL)
	if err != nil {
		return page{}, err
	}
	cur := in.Cursor

	// Use the ChQL query engine for querying asset tags.
	var assets []map[string]interface{}
	assets, cur, err = a.indexer.Assets(ctx, q, in.ChQLParams, cur, limit)
	if err != nil {
		return page{}, errors.Wrap(err, "running asset query")
	}

	out := in
	out.Cursor = cur
	return page{
		Items:    httpjson.Array(assets),
		LastPage: len(assets) < limit,
		Query:    out,
	}, nil
}
<|endoftext|>


<commit_before>/*
 * Copyright 2017 the original author or 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 http

import (
	dispatcher "github.com/sk8sio/function-sidecar/pkg/dispatcher"
	"net/http"
	"bytes"
	"io/ioutil"
	"log"
)

type httpDispatcher struct {
	Foo string
}

func (httpDispatcher) Dispatch(in interface{}) (interface{}, error) {
	slice := ([]byte)(in.(string))
	resp, err := http.Post("http://localhost:8080", "text/plain", bytes.NewReader(slice))

	if err != nil {
		log.Printf("Error invoking http://localhost:8080: %v", err)
		return nil, err
	}
	defer resp.Body.Close()
	out, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Printf("Error reading response", err)
		return nil, err
	}

	return string(out), nil
}

func NewHttpDispatcher() dispatcher.Dispatcher {
	return httpDispatcher{""}
}
<commit_msg>Add a 10s timeout for http connection<commit_after>/*
 * Copyright 2017 the original author or 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 http

import (
	dispatcher "github.com/sk8sio/function-sidecar/pkg/dispatcher"
	"net/http"
	"bytes"
	"io/ioutil"
	"log"
	"time"
)

type httpDispatcher struct {
	Foo string
}

func (httpDispatcher) Dispatch(in interface{}) (interface{}, error) {
	slice := ([]byte)(in.(string))

	client := http.Client{
		Timeout: time.Duration(10 * time.Second),
	}
	resp, err := client.Post("http://localhost:8080", "text/plain", bytes.NewReader(slice))

	if err != nil {
		log.Printf("Error invoking http://localhost:8080: %v", err)
		return nil, err
	}
	defer resp.Body.Close()
	out, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Printf("Error reading response", err)
		return nil, err
	}

	return string(out), nil
}

func NewHttpDispatcher() dispatcher.Dispatcher {
	return httpDispatcher{""}
}
<|endoftext|>


<commit_before>package http

import (
	"encoding/json"
	"net/http"
)

func Handler() http.Handler {
	mux := http.NewServeMux()
	//the handle functions goes here
	mux.Handle("/v1/user/", handleUser())

	return mux //mux implements http.Handler interface, so we can use it in the ListenAndServe method
}

func handleUser() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case "GET":
			handleUserGet(w, r)
		case "PUT":
			handleUserPut(w, r)
		case "POST":
			handleUserPost(w, r)
		case "DELETE":
			handleUserDelete(w, r)
		default:
			respondError(w, http.StatusOK, nil)
		}
	})
}

func handleUserGet(w http.ResponseWriter, r *http.Request) {
	w.Header().Add("Content-Type", "application/json")

	w.WriteHeader(http.StatusOK)
	w.Write([]byte("Hello World")) //this should be json data
}

func handleUserPut(w http.ResponseWriter, r *http.Request) {
	respondError(w, http.StatusForbidden, nil)
}

func handleUserPost(w http.ResponseWriter, r *http.Request) {
	respondOk(w, nil)
}

func handleUserDelete(w http.ResponseWriter, r *http.Request) {

}

func respondError(w http.ResponseWriter, status int, err error) {
	w.Header().Add("Content-Type", "application/json")
	w.WriteHeader(status)
}

func respondOk(w http.ResponseWriter, body interface{}) {
	w.Header().Add("Content-Type", "application/json")

	if body == nil {
		w.WriteHeader(http.StatusNoContent)
	} else {
		w.WriteHeader(http.StatusOK)
		enc := json.NewEncoder(w)
		enc.Encode(body)
	}
}
<commit_msg>litte docu<commit_after>package http

import (
	"encoding/json"
	"net/http"
)

func Handler() http.Handler {
	mux := http.NewServeMux()
	//the handle functions goes here
	mux.Handle("/v1/user/", handleUser())

	return mux //mux implements http.Handler interface, so we can use it in the ListenAndServe method
}

func handleUser() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case "GET": //handle HTTP.GET
			handleUserGet(w, r)
		case "PUT": //handle HTTP.PUT
			handleUserPut(w, r)
		case "POST": //handle HTTP.POST
			handleUserPost(w, r)
		case "DELETE": //handle HTTP.DELETE
			handleUserDelete(w, r)
		default:
			respondError(w, http.StatusMethodNotAllowed, nil)
		}
	})
}

func handleUserGet(w http.ResponseWriter, r *http.Request) {
	w.Header().Add("Content-Type", "application/json")

	w.WriteHeader(http.StatusOK)
	w.Write([]byte("Hello World")) //this should be json data
}

func handleUserPut(w http.ResponseWriter, r *http.Request) {
	respondError(w, http.StatusForbidden, nil)
}

func handleUserPost(w http.ResponseWriter, r *http.Request) {
	respondOk(w, nil)
}

func handleUserDelete(w http.ResponseWriter, r *http.Request) {

}

func respondError(w http.ResponseWriter, status int, err error) {
	w.Header().Add("Content-Type", "application/json")
	w.WriteHeader(status)
}

func respondOk(w http.ResponseWriter, body interface{}) {
	w.Header().Add("Content-Type", "application/json")

	if body == nil {
		w.WriteHeader(http.StatusNoContent)
	} else {
		w.WriteHeader(http.StatusOK)
		enc := json.NewEncoder(w)
		enc.Encode(body)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"testing"
)

func TestParseArgs_Copy(t *testing.T) {
	good := [][]string{
		[]string{"cp", "one", "two"},
		[]string{"copy", "One", "Two"},
	}

	for _, args := range good {
		cmd, err := ParseArgs(args)
		if err != nil {
			t.Fatalf("Expected %v to parse, error: %v", args, err)
		}

		if cmd == nil {
			t.Fatalf("Expected %v to parse", args)
		}

		if _, ok := cmd.(*Copy); !ok {
			t.Fatalf("Expected %v to produce a Copy command", args)
		}
	}

	bad := [][]string{
		[]string{"cp", "one"},
		[]string{"cp", "one", "two", "three"},
		[]string{"copy", "one"},
		[]string{"copy", "one", "two", "three"},
	}

	for _, args := range bad {
		_, err := ParseArgs(args)
		if err == nil {
			t.Fatalf("Expected %v to fail to parse", args)
		}
	}
}
<commit_msg>More robust and extensible ParseArgs test<commit_after>package main

import (
	"reflect"
	"testing"
)

type parseCase struct {
	Args    []string
	Command Command
}

var (
	goodParseCases = []parseCase{
		{
			Args: []string{"cp", "one", "two"},
			Command: &Copy{
				OldVaultName: "one",
				NewVaultName: "two",
			},
		},
		{
			Args: []string{"copy", "one", "two"},
			Command: &Copy{
				OldVaultName: "one",
				NewVaultName: "two",
			},
		},
	}

	badParseCases = []parseCase{
		{
			Args: []string{"cp", "one"},
		},
		{
			Args: []string{"cp", "one", "two", "three"},
		},
		{
			Args: []string{"copy", "one"},
		},
		{
			Args: []string{"copy", "one", "two", "three"},
		},
	}
)

type parseExpectation struct {
	Args    []string
	Command Command
}

func TestParseArgs(t *testing.T) {
	for _, good := range goodParseCases {
		cmd, err := ParseArgs(good.Args)
		if err != nil {
			t.Fatalf("Failed to parse '%v': %v", good.Args, err)
		}

		if !reflect.DeepEqual(good.Command, cmd) {
			t.Fatalf("Expected command: %#v, got: %#v", good.Command, cmd)
		}
	}

	for _, bad := range badParseCases {
		_, err := ParseArgs(bad.Args)
		if err == nil {
			t.Fatalf("Expected '%v' to fail to parse", bad.Args)
		}
	}
}
<|endoftext|>


<commit_before>package cwl

// Arguments ...
type Arguments []Argument

// New constructs "Arguments" struct.
func (baseCommands Arguments) New(i interface{}) Arguments {
	dest := Arguments{}
	switch x := i.(type) {
	case string:
		argument := Argument{}
		argument.Value = x
		dest = append(dest, argument)
	case []interface{}:
		for _, elm := range x {
			argument := Argument{}
			switch val := elm.(type) {
			case string:
				argument.Value = val
			case map[string]interface{}:
				argument.CommandLineBinding = val
			}
			dest = append(dest, argument)
		}
	}
	return dest
}

// Argument
type Argument struct {
	Value             string
	CommandLineBinding map[string]interface{}
	// TODO support Expression
}
<commit_msg>Exec go fmt<commit_after>package cwl

// Arguments ...
type Arguments []Argument

// New constructs "Arguments" struct.
func (baseCommands Arguments) New(i interface{}) Arguments {
	dest := Arguments{}
	switch x := i.(type) {
	case string:
		argument := Argument{}
		argument.Value = x
		dest = append(dest, argument)
	case []interface{}:
		for _, elm := range x {
			argument := Argument{}
			switch val := elm.(type) {
			case string:
				argument.Value = val
			case map[string]interface{}:
				argument.CommandLineBinding = val
			}
			dest = append(dest, argument)
		}
	}
	return dest
}

// Argument
type Argument struct {
	Value              string
	CommandLineBinding map[string]interface{}
	// TODO support Expression
}
<|endoftext|>


<commit_before>package ngorm

import (
	"fmt"
	"reflect"

	"github.com/ngorm/ngorm/model"
	"github.com/ngorm/ngorm/scope"
	"github.com/ngorm/ngorm/util"
)

// Association provides utility functions for dealing with association queries
type Association struct {
	db     *DB
	column string
	field  *model.Field
}

// Find find out all related associations
func (a *Association) Find(v interface{}) error {
	return a.db.Related(v, a.column)
}

// Append append new associations for many2many, has_many, replace current association for has_one, belongs_to
func (a *Association) Append(values ...interface{}) error {
	return a.Save(values...)
}

// Save save passed values as associations
func (a *Association) Save(values ...interface{}) error {
	if len(values) > 0 {
		e := a.db.e
		field := a.field
		rel := field.Relationship
		var v reflect.Value
		if rel.Kind == "has_one" {
			if len(values) > 1 {
				return fmt.Errorf("relation %s expect one struct value got %d", rel.Kind, len(values))
			}
			v = reflect.New(field.Field.Type())
			err := a.Find(v.Interface())
			if err != nil {
				return err
			}
			v = v.Elem()
			ov := reflect.ValueOf(values[0])
			if ov.Kind() == reflect.Ptr {
				ov = ov.Elem()
			}
			ovTyp := ov.Type()
			for i := 0; i < ovTyp.NumField(); i++ {
				fTyp := ovTyp.Field(i)
				fv := ov.FieldByName(fTyp.Name)
				if !isZero(fv) {
					fEv := v.FieldByName(fTyp.Name)
					fEv.Set(fv)
				}
			}
		} else {
			v = reflect.MakeSlice(field.Struct.Type, 0, 0)
			if v.Kind() == reflect.Ptr {
				v = v.Elem()
			}
			for _, value := range values {
				fv := reflect.ValueOf(value)
				if fv.Kind() == reflect.Ptr {
					fv = fv.Elem()
				}
				v = reflect.Append(v, fv)
			}
		}
		field.Field.Set(v)
		return a.db.Begin().Save(e.Scope.Value)
	}
	return nil
}

func isZero(v reflect.Value) bool {
	return v.Interface() == reflect.Zero(v.Type()).Interface()
}

// Count return the count of current associations
func (a *Association) Count() (int, error) {
	var (
		count      = 0
		rel        = a.field.Relationship
		fieldValue = a.field.Field.Interface()
		query      = a.db.Begin().Model(fieldValue)
	)
	if rel.Kind == "many_to_many" {
		err := scope.JoinWith(rel.JoinTableHandler, query.e, a.db.e.Scope.Value)
		if err != nil {
			return 0, err
		}
	} else if rel.Kind == "has_many" || rel.Kind == "has_one" {
		primaryKeys := util.ColumnAsArray(rel.AssociationForeignFieldNames, a.db.e.Scope.Value)
		query = query.Where(
			fmt.Sprintf("%v IN (%v)",
				scope.ToQueryCondition(a.db.e, rel.ForeignDBNames),
				scope.ToQueryMarks(primaryKeys)),
			util.ToQueryValues(primaryKeys)...,
		)
	} else if rel.Kind == "belongs_to" {
		primaryKeys := util.ColumnAsArray(rel.ForeignFieldNames, a.db.e.Scope.Value)
		query = query.Where(
			fmt.Sprintf("%v IN (%v)",
				scope.ToQueryCondition(a.db.e, rel.AssociationForeignDBNames),
				scope.ToQueryMarks(primaryKeys)),
			util.ToQueryValues(primaryKeys)...,
		)
	}

	if rel.PolymorphicType != "" {
		query = query.Where(
			fmt.Sprintf("%v%v = ?",
				a.db.e.Dialect.QueryFieldName(
					scope.QuotedTableName(a.db.e, fieldValue)),
				scope.Quote(a.db.e, rel.PolymorphicDBName)),
			rel.PolymorphicValue,
		)
	}
	err := query.Count(&count)
	if err != nil {
		return 0, err
	}
	return count, nil
}
<commit_msg>[ngorm] add godoc<commit_after>package ngorm

import (
	"fmt"
	"reflect"

	"github.com/ngorm/ngorm/model"
	"github.com/ngorm/ngorm/scope"
	"github.com/ngorm/ngorm/util"
)

// Association provides utility functions for dealing with association queries
type Association struct {
	db     *DB
	column string
	field  *model.Field
}

// Find find out all related associations
func (a *Association) Find(v interface{}) error {
	return a.db.Related(v, a.column)
}

// Append append new associations for many2many, has_many, replace current
// association for has_one, belongs_to
//
// This wraps around Association.Save, verbatim  meaning you can have the same
// effect with Save method.
func (a *Association) Append(values ...interface{}) error {
	return a.Save(values...)
}

// Save save passed values as associations. This expects to have a single value
// for a has_one, belongs_to relationships. You can pass one or more values for
// many_to_many relationship.
func (a *Association) Save(values ...interface{}) error {
	if len(values) > 0 {
		e := a.db.e
		field := a.field
		rel := field.Relationship
		var v reflect.Value
		if rel.Kind == "has_one" {
			if len(values) > 1 {
				return fmt.Errorf("relation %s expect one struct value got %d", rel.Kind, len(values))
			}
			v = reflect.New(field.Field.Type())
			err := a.Find(v.Interface())
			if err != nil {
				return err
			}
			v = v.Elem()
			ov := reflect.ValueOf(values[0])
			if ov.Kind() == reflect.Ptr {
				ov = ov.Elem()
			}
			ovTyp := ov.Type()
			for i := 0; i < ovTyp.NumField(); i++ {
				fTyp := ovTyp.Field(i)
				fv := ov.FieldByName(fTyp.Name)
				if !isZero(fv) {
					fEv := v.FieldByName(fTyp.Name)
					fEv.Set(fv)
				}
			}
		} else {
			v = reflect.MakeSlice(field.Struct.Type, 0, 0)
			if v.Kind() == reflect.Ptr {
				v = v.Elem()
			}
			for _, value := range values {
				fv := reflect.ValueOf(value)
				if fv.Kind() == reflect.Ptr {
					fv = fv.Elem()
				}
				v = reflect.Append(v, fv)
			}
		}
		field.Field.Set(v)
		return a.db.Begin().Save(e.Scope.Value)
	}
	return nil
}

func isZero(v reflect.Value) bool {
	return v.Interface() == reflect.Zero(v.Type()).Interface()
}

// Count return the count of current associations
func (a *Association) Count() (int, error) {
	var (
		count      = 0
		rel        = a.field.Relationship
		fieldValue = a.field.Field.Interface()
		query      = a.db.Begin().Model(fieldValue)
	)
	if rel.Kind == "many_to_many" {
		err := scope.JoinWith(rel.JoinTableHandler, query.e, a.db.e.Scope.Value)
		if err != nil {
			return 0, err
		}
	} else if rel.Kind == "has_many" || rel.Kind == "has_one" {
		primaryKeys := util.ColumnAsArray(rel.AssociationForeignFieldNames, a.db.e.Scope.Value)
		query = query.Where(
			fmt.Sprintf("%v IN (%v)",
				scope.ToQueryCondition(a.db.e, rel.ForeignDBNames),
				scope.ToQueryMarks(primaryKeys)),
			util.ToQueryValues(primaryKeys)...,
		)
	} else if rel.Kind == "belongs_to" {
		primaryKeys := util.ColumnAsArray(rel.ForeignFieldNames, a.db.e.Scope.Value)
		query = query.Where(
			fmt.Sprintf("%v IN (%v)",
				scope.ToQueryCondition(a.db.e, rel.AssociationForeignDBNames),
				scope.ToQueryMarks(primaryKeys)),
			util.ToQueryValues(primaryKeys)...,
		)
	}

	if rel.PolymorphicType != "" {
		query = query.Where(
			fmt.Sprintf("%v%v = ?",
				a.db.e.Dialect.QueryFieldName(
					scope.QuotedTableName(a.db.e, fieldValue)),
				scope.Quote(a.db.e, rel.PolymorphicDBName)),
			rel.PolymorphicValue,
		)
	}
	err := query.Count(&count)
	if err != nil {
		return 0, err
	}
	return count, nil
}
<|endoftext|>


<commit_before>package db

import (
	"database/sql"
	"database/sql/driver"
	"errors"
	"fmt"
	"io"
	"strings"
	"time"

	"code.cloudfoundry.org/lager"
	"github.com/Masterminds/squirrel"
	"github.com/concourse/concourse/atc/db/encryption"
	"github.com/concourse/concourse/atc/db/lock"
	"github.com/concourse/concourse/atc/db/migration"
	multierror "github.com/hashicorp/go-multierror"
	"github.com/lib/pq"
)

//go:generate counterfeiter . Conn

type Conn interface {
	Bus() NotificationsBus
	EncryptionStrategy() encryption.Strategy

	Ping() error
	Driver() driver.Driver

	Begin() (Tx, error)
	Exec(query string, args ...interface{}) (sql.Result, error)
	Prepare(query string) (*sql.Stmt, error)
	Query(query string, args ...interface{}) (*sql.Rows, error)
	QueryRow(query string, args ...interface{}) squirrel.RowScanner

	SetMaxIdleConns(n int)
	SetMaxOpenConns(n int)
	Stats() sql.DBStats

	Close() error
	Name() string
}

//go:generate counterfeiter . Tx

type Tx interface {
	Commit() error
	Exec(query string, args ...interface{}) (sql.Result, error)
	Prepare(query string) (*sql.Stmt, error)
	Query(query string, args ...interface{}) (*sql.Rows, error)
	QueryRow(query string, args ...interface{}) squirrel.RowScanner
	Rollback() error
	Stmt(stmt *sql.Stmt) *sql.Stmt
}

func Open(logger lager.Logger, sqlDriver string, sqlDataSource string, newKey *encryption.Key, oldKey *encryption.Key, connectionName string, lockFactory lock.LockFactory) (Conn, error) {
	for {
		var strategy encryption.Strategy
		if newKey != nil {
			strategy = newKey
		} else {
			strategy = encryption.NewNoEncryption()
		}

		sqlDb, err := migration.NewOpenHelper(sqlDriver, sqlDataSource, lockFactory, strategy).Open()
		if err != nil {
			if shouldRetry(err) {
				logger.Error("failed-to-open-db-retrying", err)
				time.Sleep(5 * time.Second)
				continue
			}

			return nil, err
		}

		switch {
		case oldKey != nil && newKey == nil:
			err = decryptToPlaintext(logger.Session("decrypt"), sqlDb, oldKey)
		case oldKey != nil && newKey != nil:
			err = encryptWithNewKey(logger.Session("rotate"), sqlDb, newKey, oldKey)
		}
		if err != nil {
			return nil, err
		}

		if newKey != nil {
			err = encryptPlaintext(logger.Session("encrypt"), sqlDb, newKey)
			if err != nil {
				return nil, err
			}
		}

		listener := pq.NewDialListener(keepAliveDialer{}, sqlDataSource, time.Second, time.Minute, nil)

		return &db{
			DB: sqlDb,

			bus:        NewNotificationsBus(listener, sqlDb),
			encryption: strategy,
			name:       connectionName,
		}, nil
	}
}

func shouldRetry(err error) bool {
	if strings.Contains(err.Error(), "dial ") {
		return true
	}

	if pqErr, ok := err.(*pq.Error); ok {
		return pqErr.Code.Name() == "cannot_connect_now"
	}

	return false
}

var encryptedColumns = map[string]string{
	"teams":          "legacy_auth",
	"resources":      "config",
	"jobs":           "config",
	"resource_types": "config",
	"builds":         "private_plan",
	"cert_cache":     "cert",
}

func encryptPlaintext(logger lager.Logger, sqlDB *sql.DB, key *encryption.Key) error {
	for table, col := range encryptedColumns {
		rows, err := sqlDB.Query(`
			SELECT id, ` + col + `
			FROM ` + table + `
			WHERE nonce IS NULL
			AND ` + col + ` IS NOT NULL
		`)
		if err != nil {
			return err
		}

		tLog := logger.Session("table", lager.Data{
			"table": table,
		})

		encryptedRows := 0

		for rows.Next() {
			var (
				id  int
				val sql.NullString
			)

			err := rows.Scan(&id, &val)
			if err != nil {
				tLog.Error("failed-to-scan", err)
				return err
			}

			if !val.Valid {
				continue
			}

			rLog := tLog.Session("row", lager.Data{
				"id": id,
			})

			encrypted, nonce, err := key.Encrypt([]byte(val.String))
			if err != nil {
				rLog.Error("failed-to-encrypt", err)
				return err
			}

			_, err = sqlDB.Exec(`
				UPDATE `+table+`
				SET `+col+` = $1, nonce = $2
				WHERE id = $3
			`, encrypted, nonce, id)
			if err != nil {
				rLog.Error("failed-to-update", err)
				return err
			}

			encryptedRows++
		}

		if encryptedRows > 0 {
			tLog.Info("encrypted-existing-plaintext-data", lager.Data{
				"rows": encryptedRows,
			})
		}
	}

	return nil
}

func decryptToPlaintext(logger lager.Logger, sqlDB *sql.DB, oldKey *encryption.Key) error {
	for table, col := range encryptedColumns {
		rows, err := sqlDB.Query(`
			SELECT id, nonce, ` + col + `
			FROM ` + table + `
			WHERE nonce IS NOT NULL
		`)
		if err != nil {
			return err
		}

		tLog := logger.Session("table", lager.Data{
			"table": table,
		})

		decryptedRows := 0

		for rows.Next() {
			var (
				id         int
				val, nonce string
			)

			err := rows.Scan(&id, &nonce, &val)
			if err != nil {
				tLog.Error("failed-to-scan", err)
				return err
			}

			rLog := tLog.Session("row", lager.Data{
				"id": id,
			})

			decrypted, err := oldKey.Decrypt(val, &nonce)
			if err != nil {
				rLog.Error("failed-to-decrypt", err)
				return err
			}

			_, err = sqlDB.Exec(`
				UPDATE `+table+`
				SET `+col+` = $1, nonce = NULL
				WHERE id = $2
			`, decrypted, id)
			if err != nil {
				rLog.Error("failed-to-update", err)
				return err
			}

			decryptedRows++
		}

		if decryptedRows > 0 {
			tLog.Info("decrypted-existing-encrypted-data", lager.Data{
				"rows": decryptedRows,
			})
		}
	}

	return nil
}

var ErrEncryptedWithUnknownKey = errors.New("row encrypted with neither old nor new key")

func encryptWithNewKey(logger lager.Logger, sqlDB *sql.DB, newKey *encryption.Key, oldKey *encryption.Key) error {
	for table, col := range encryptedColumns {
		rows, err := sqlDB.Query(`
			SELECT id, nonce, ` + col + `
			FROM ` + table + `
			WHERE nonce IS NOT NULL
		`)
		if err != nil {
			return err
		}

		tLog := logger.Session("table", lager.Data{
			"table": table,
		})

		encryptedRows := 0

		for rows.Next() {
			var (
				id         int
				val, nonce string
			)

			err := rows.Scan(&id, &nonce, &val)
			if err != nil {
				tLog.Error("failed-to-scan", err)
				return err
			}

			rLog := tLog.Session("row", lager.Data{
				"id": id,
			})

			decrypted, err := oldKey.Decrypt(val, &nonce)
			if err != nil {
				_, err = newKey.Decrypt(val, &nonce)
				if err == nil {
					rLog.Debug("already-encrypted-with-new-key")
					continue
				}

				logger.Error("failed-to-decrypt-with-either-key", err)
				return ErrEncryptedWithUnknownKey
			}

			encrypted, newNonce, err := newKey.Encrypt(decrypted)
			if err != nil {
				rLog.Error("failed-to-encrypt", err)
				return err
			}

			_, err = sqlDB.Exec(`
				UPDATE `+table+`
				SET `+col+` = $1, nonce = $2
				WHERE id = $3
			`, encrypted, newNonce, id)
			if err != nil {
				rLog.Error("failed-to-update", err)
				return err
			}

			encryptedRows++
		}

		if encryptedRows > 0 {
			tLog.Info("re-encrypted-existing-encrypted-data", lager.Data{
				"rows": encryptedRows,
			})
		}
	}

	return nil
}

type db struct {
	*sql.DB

	bus        NotificationsBus
	encryption encryption.Strategy
	name       string
}

func (db *db) Name() string {
	return db.name
}

func (db *db) Bus() NotificationsBus {
	return db.bus
}

func (db *db) EncryptionStrategy() encryption.Strategy {
	return db.encryption
}

func (db *db) Close() error {
	var errs error
	dbErr := db.DB.Close()
	if dbErr != nil {
		errs = multierror.Append(errs, dbErr)
	}

	busErr := db.bus.Close()
	if busErr != nil {
		errs = multierror.Append(errs, busErr)
	}

	return errs
}

// Close ignores errors, and should used with defer.
// makes errcheck happy that those errs are captured
func Close(c io.Closer) {
	_ = c.Close()
}

func (db *db) Begin() (Tx, error) {
	tx, err := db.DB.Begin()
	if err != nil {
		return nil, err
	}

	return &dbTx{tx, GlobalConnectionTracker.Track()}, nil
}

func (db *db) Exec(query string, args ...interface{}) (sql.Result, error) {
	defer GlobalConnectionTracker.Track().Release()
	return db.DB.Exec(query, args...)
}

func (db *db) Prepare(query string) (*sql.Stmt, error) {
	defer GlobalConnectionTracker.Track().Release()
	return db.DB.Prepare(query)
}

func (db *db) Query(query string, args ...interface{}) (*sql.Rows, error) {
	defer GlobalConnectionTracker.Track().Release()
	return db.DB.Query(query, args...)
}

// to conform to squirrel.Runner interface
func (db *db) QueryRow(query string, args ...interface{}) squirrel.RowScanner {
	defer GlobalConnectionTracker.Track().Release()
	return db.DB.QueryRow(query, args...)
}

type dbTx struct {
	*sql.Tx

	session *ConnectionSession
}

// to conform to squirrel.Runner interface
func (tx *dbTx) QueryRow(query string, args ...interface{}) squirrel.RowScanner {
	return tx.Tx.QueryRow(query, args...)
}

func (tx *dbTx) Commit() error {
	defer tx.session.Release()
	return tx.Tx.Commit()
}

func (tx *dbTx) Rollback() error {
	defer tx.session.Release()
	return tx.Tx.Rollback()
}

// Rollback ignores errors, and should be used with defer.
// makes errcheck happy that those errs are captured
func Rollback(tx Tx) {
	_ = tx.Rollback()
}

type nonOneRowAffectedError struct {
	RowsAffected int64
}

func (err nonOneRowAffectedError) Error() string {
	return fmt.Sprintf("expected 1 row to be updated; got %d", err.RowsAffected)
}
<commit_msg>atc: support encrypting tables with non-'id' pkeys<commit_after>package db

import (
	"database/sql"
	"database/sql/driver"
	"errors"
	"fmt"
	"io"
	"strings"
	"time"

	"code.cloudfoundry.org/lager"
	"github.com/Masterminds/squirrel"
	"github.com/concourse/concourse/atc/db/encryption"
	"github.com/concourse/concourse/atc/db/lock"
	"github.com/concourse/concourse/atc/db/migration"
	multierror "github.com/hashicorp/go-multierror"
	"github.com/lib/pq"
)

//go:generate counterfeiter . Conn

type Conn interface {
	Bus() NotificationsBus
	EncryptionStrategy() encryption.Strategy

	Ping() error
	Driver() driver.Driver

	Begin() (Tx, error)
	Exec(query string, args ...interface{}) (sql.Result, error)
	Prepare(query string) (*sql.Stmt, error)
	Query(query string, args ...interface{}) (*sql.Rows, error)
	QueryRow(query string, args ...interface{}) squirrel.RowScanner

	SetMaxIdleConns(n int)
	SetMaxOpenConns(n int)
	Stats() sql.DBStats

	Close() error
	Name() string
}

//go:generate counterfeiter . Tx

type Tx interface {
	Commit() error
	Exec(query string, args ...interface{}) (sql.Result, error)
	Prepare(query string) (*sql.Stmt, error)
	Query(query string, args ...interface{}) (*sql.Rows, error)
	QueryRow(query string, args ...interface{}) squirrel.RowScanner
	Rollback() error
	Stmt(stmt *sql.Stmt) *sql.Stmt
}

func Open(logger lager.Logger, sqlDriver string, sqlDataSource string, newKey *encryption.Key, oldKey *encryption.Key, connectionName string, lockFactory lock.LockFactory) (Conn, error) {
	for {
		var strategy encryption.Strategy
		if newKey != nil {
			strategy = newKey
		} else {
			strategy = encryption.NewNoEncryption()
		}

		sqlDb, err := migration.NewOpenHelper(sqlDriver, sqlDataSource, lockFactory, strategy).Open()
		if err != nil {
			if shouldRetry(err) {
				logger.Error("failed-to-open-db-retrying", err)
				time.Sleep(5 * time.Second)
				continue
			}

			return nil, err
		}

		switch {
		case oldKey != nil && newKey == nil:
			err = decryptToPlaintext(logger.Session("decrypt"), sqlDb, oldKey)
		case oldKey != nil && newKey != nil:
			err = encryptWithNewKey(logger.Session("rotate"), sqlDb, newKey, oldKey)
		}
		if err != nil {
			return nil, err
		}

		if newKey != nil {
			err = encryptPlaintext(logger.Session("encrypt"), sqlDb, newKey)
			if err != nil {
				return nil, err
			}
		}

		listener := pq.NewDialListener(keepAliveDialer{}, sqlDataSource, time.Second, time.Minute, nil)

		return &db{
			DB: sqlDb,

			bus:        NewNotificationsBus(listener, sqlDb),
			encryption: strategy,
			name:       connectionName,
		}, nil
	}
}

func shouldRetry(err error) bool {
	if strings.Contains(err.Error(), "dial ") {
		return true
	}

	if pqErr, ok := err.(*pq.Error); ok {
		return pqErr.Code.Name() == "cannot_connect_now"
	}

	return false
}

type encryptedColumn struct {
	Table      string
	Column     string
	PrimaryKey string
}

var encryptedColumns = []encryptedColumn{
	{"teams", "legacy_auth", "id"},
	{"resources", "config", "id"},
	{"jobs", "config", "id"},
	{"resource_types", "config", "id"},
	{"builds", "private_plan", "id"},
	{"vert_cache", "cert", "domain"},
}

func encryptPlaintext(logger lager.Logger, sqlDB *sql.DB, key *encryption.Key) error {
	for _, ec := range encryptedColumns {
		rows, err := sqlDB.Query(`
			SELECT ` + ec.PrimaryKey + `, ` + ec.Column + `
			FROM ` + ec.Table + `
			WHERE nonce IS NULL
			AND ` + ec.Column + ` IS NOT NULL
		`)
		if err != nil {
			return err
		}

		tLog := logger.Session("table", lager.Data{
			"table": ec.Table,
		})

		encryptedRows := 0

		for rows.Next() {
			var (
				key interface{}
				val sql.NullString
			)

			err := rows.Scan(&key, &val)
			if err != nil {
				tLog.Error("failed-to-scan", err)
				return err
			}

			if !val.Valid {
				continue
			}

			rLog := tLog.Session("row", lager.Data{
				"key": key,
			})

			encrypted, nonce, err := key.Encrypt([]byte(val.String))
			if err != nil {
				rLog.Error("failed-to-encrypt", err)
				return err
			}

			_, err = sqlDB.Exec(`
				UPDATE `+ec.Table+`
				SET `+ec.Column+` = $1, nonce = $2
				WHERE `+ec.PrimaryKey+` = $3
			`, encrypted, nonce, key)
			if err != nil {
				rLog.Error("failed-to-update", err)
				return err
			}

			encryptedRows++
		}

		if encryptedRows > 0 {
			tLog.Info("encrypted-existing-plaintext-data", lager.Data{
				"rows": encryptedRows,
			})
		}
	}

	return nil
}

func decryptToPlaintext(logger lager.Logger, sqlDB *sql.DB, oldKey *encryption.Key) error {
	for _, ec := range encryptedColumns {
		rows, err := sqlDB.Query(`
			SELECT ` + ec.PrimaryKey + `, nonce, ` + ec.Column + `
			FROM ` + ec.Table + `
			WHERE nonce IS NOT NULL
		`)
		if err != nil {
			return err
		}

		tLog := logger.Session("table", lager.Data{
			"table": ec.Table,
		})

		decryptedRows := 0

		for rows.Next() {
			var (
				key        interface{}
				val, nonce string
			)

			err := rows.Scan(&key, &nonce, &val)
			if err != nil {
				tLog.Error("failed-to-scan", err)
				return err
			}

			rLog := tLog.Session("row", lager.Data{
				"key": key,
			})

			decrypted, err := oldKey.Decrypt(val, &nonce)
			if err != nil {
				rLog.Error("failed-to-decrypt", err)
				return err
			}

			_, err = sqlDB.Exec(`
				UPDATE `+ec.Table+`
				SET `+ec.Column+` = $1, nonce = NULL
				WHERE `+ec.PrimaryKey+` = $2
			`, decrypted, key)
			if err != nil {
				rLog.Error("failed-to-update", err)
				return err
			}

			decryptedRows++
		}

		if decryptedRows > 0 {
			tLog.Info("decrypted-existing-encrypted-data", lager.Data{
				"rows": decryptedRows,
			})
		}
	}

	return nil
}

var ErrEncryptedWithUnknownKey = errors.New("row encrypted with neither old nor new key")

func encryptWithNewKey(logger lager.Logger, sqlDB *sql.DB, newKey *encryption.Key, oldKey *encryption.Key) error {
	for _, ec := range encryptedColumns {
		rows, err := sqlDB.Query(`
			SELECT ` + ec.PrimaryKey + `, nonce, ` + ec.Column + `
			FROM ` + ec.Table + `
			WHERE nonce IS NOT NULL
		`)
		if err != nil {
			return err
		}

		tLog := logger.Session("table", lager.Data{
			"table": ec.Table,
		})

		encryptedRows := 0

		for rows.Next() {
			var (
				key        interface{}
				val, nonce string
			)

			err := rows.Scan(&key, &nonce, &val)
			if err != nil {
				tLog.Error("failed-to-scan", err)
				return err
			}

			rLog := tLog.Session("row", lager.Data{
				"key": key,
			})

			decrypted, err := oldKey.Decrypt(val, &nonce)
			if err != nil {
				_, err = newKey.Decrypt(val, &nonce)
				if err == nil {
					rLog.Debug("already-encrypted-with-new-key")
					continue
				}

				logger.Error("failed-to-decrypt-with-either-key", err)
				return ErrEncryptedWithUnknownKey
			}

			encrypted, newNonce, err := newKey.Encrypt(decrypted)
			if err != nil {
				rLog.Error("failed-to-encrypt", err)
				return err
			}

			_, err = sqlDB.Exec(`
				UPDATE `+ec.Table+`
				SET `+ec.Column+` = $1, nonce = $2
				WHERE `+ec.PrimaryKey+` = $3
			`, encrypted, newNonce, key)
			if err != nil {
				rLog.Error("failed-to-update", err)
				return err
			}

			encryptedRows++
		}

		if encryptedRows > 0 {
			tLog.Info("re-encrypted-existing-encrypted-data", lager.Data{
				"rows": encryptedRows,
			})
		}
	}

	return nil
}

type db struct {
	*sql.DB

	bus        NotificationsBus
	encryption encryption.Strategy
	name       string
}

func (db *db) Name() string {
	return db.name
}

func (db *db) Bus() NotificationsBus {
	return db.bus
}

func (db *db) EncryptionStrategy() encryption.Strategy {
	return db.encryption
}

func (db *db) Close() error {
	var errs error
	dbErr := db.DB.Close()
	if dbErr != nil {
		errs = multierror.Append(errs, dbErr)
	}

	busErr := db.bus.Close()
	if busErr != nil {
		errs = multierror.Append(errs, busErr)
	}

	return errs
}

// Close ignores errors, and should used with defer.
// makes errcheck happy that those errs are captured
func Close(c io.Closer) {
	_ = c.Close()
}

func (db *db) Begin() (Tx, error) {
	tx, err := db.DB.Begin()
	if err != nil {
		return nil, err
	}

	return &dbTx{tx, GlobalConnectionTracker.Track()}, nil
}

func (db *db) Exec(query string, args ...interface{}) (sql.Result, error) {
	defer GlobalConnectionTracker.Track().Release()
	return db.DB.Exec(query, args...)
}

func (db *db) Prepare(query string) (*sql.Stmt, error) {
	defer GlobalConnectionTracker.Track().Release()
	return db.DB.Prepare(query)
}

func (db *db) Query(query string, args ...interface{}) (*sql.Rows, error) {
	defer GlobalConnectionTracker.Track().Release()
	return db.DB.Query(query, args...)
}

// to conform to squirrel.Runner interface
func (db *db) QueryRow(query string, args ...interface{}) squirrel.RowScanner {
	defer GlobalConnectionTracker.Track().Release()
	return db.DB.QueryRow(query, args...)
}

type dbTx struct {
	*sql.Tx

	session *ConnectionSession
}

// to conform to squirrel.Runner interface
func (tx *dbTx) QueryRow(query string, args ...interface{}) squirrel.RowScanner {
	return tx.Tx.QueryRow(query, args...)
}

func (tx *dbTx) Commit() error {
	defer tx.session.Release()
	return tx.Tx.Commit()
}

func (tx *dbTx) Rollback() error {
	defer tx.session.Release()
	return tx.Tx.Rollback()
}

// Rollback ignores errors, and should be used with defer.
// makes errcheck happy that those errs are captured
func Rollback(tx Tx) {
	_ = tx.Rollback()
}

type nonOneRowAffectedError struct {
	RowsAffected int64
}

func (err nonOneRowAffectedError) Error() string {
	return fmt.Sprintf("expected 1 row to be updated; got %d", err.RowsAffected)
}
<|endoftext|>