<commit_before>package cli

import (
	"fmt"
	"os"
	"path"
	"runtime/pprof"

	log "github.com/Sirupsen/logrus"
	"github.com/bitrise-io/go-utils/freezable"
	"github.com/codegangsta/cli"
)

var (
	// MachineWorkdir ...
	MachineWorkdir = freezable.String{}
	// MachineParamsAdditionalEnvs ...
	MachineParamsAdditionalEnvs = freezable.StringSlice{}
	// MachineConfigTypeID ...
	MachineConfigTypeID = freezable.String{}
)

func before(c *cli.Context) error {
	// Log level
	if logLevel, err := log.ParseLevel(c.String(LogLevelKey)); err != nil {
		log.Fatal("Failed to parse log level:", err)
	} else {
		log.SetLevel(logLevel)
	}

	if len(c.Args()) != 0 && !c.Bool(HelpKey) && !c.Bool(VersionKey) {
		if err := MachineWorkdir.Set(c.String(WorkdirKey)); err != nil {
			log.Fatalf("Failed to set MachineWorkdir: %s", err)
		}
		if MachineWorkdir.String() == "" {
			log.Fatalln("No Workdir specified!")
		}
	}
	MachineWorkdir.Freeze()

	if err := MachineConfigTypeID.Set(c.String(ConfigTypeIDParamKey)); err != nil {
		log.Fatalf("Failed to set MachineConfigTypeID: %s", err)
	}
	log.Infof("MachineConfigTypeID: %s", MachineConfigTypeID)

	if err := MachineParamsAdditionalEnvs.Set(c.StringSlice(EnvironmentParamKey)); err != nil {
		log.Fatalf("Failed to set MachineParamsAdditionalEnvs: %s", err)
	}
	log.Debugf("MachineParamsAdditionalEnvs: %s", MachineParamsAdditionalEnvs)
	MachineParamsAdditionalEnvs.Freeze()

	return nil
}

func printVersion(c *cli.Context) {
	fmt.Fprintf(c.App.Writer, "%v\n", c.App.Version)
}

// Run the Envman CLI.
func Run() {
	cli.VersionPrinter = printVersion

	app := cli.NewApp()
	app.Name = path.Base(os.Args[0])
	app.Usage = "bitrise-machine"
	app.Version = "0.9.7"

	app.Author = ""
	app.Email = ""

	app.Before = before

	app.Flags = appFlags
	app.Commands = commands

	if cpuProfileFilePath := os.Getenv(CPUProfileFilePathEnvKey); cpuProfileFilePath != "" {
		log.Infof("Enabling CPU Profiler, writing results into file: %s", cpuProfileFilePath)
		f, err := os.Create(cpuProfileFilePath)
		if err != nil {
			log.Fatalf("CPU Profile: failed to create file: %s", err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatalf("CPU Profile: failed to start CPU profiler: %s", err)
		}
		defer pprof.StopCPUProfile()
	}

	if err := app.Run(os.Args); err != nil {
		log.Fatal("Finished with error:", err)
	}
}
<commit_msg>LOG: MachineConfigTypeID - only debug print<commit_after>package cli

import (
	"fmt"
	"os"
	"path"
	"runtime/pprof"

	log "github.com/Sirupsen/logrus"
	"github.com/bitrise-io/go-utils/freezable"
	"github.com/codegangsta/cli"
)

var (
	// MachineWorkdir ...
	MachineWorkdir = freezable.String{}
	// MachineParamsAdditionalEnvs ...
	MachineParamsAdditionalEnvs = freezable.StringSlice{}
	// MachineConfigTypeID ...
	MachineConfigTypeID = freezable.String{}
)

func before(c *cli.Context) error {
	// Log level
	if logLevel, err := log.ParseLevel(c.String(LogLevelKey)); err != nil {
		log.Fatal("Failed to parse log level:", err)
	} else {
		log.SetLevel(logLevel)
	}

	if len(c.Args()) != 0 && !c.Bool(HelpKey) && !c.Bool(VersionKey) {
		if err := MachineWorkdir.Set(c.String(WorkdirKey)); err != nil {
			log.Fatalf("Failed to set MachineWorkdir: %s", err)
		}
		if MachineWorkdir.String() == "" {
			log.Fatalln("No Workdir specified!")
		}
	}
	MachineWorkdir.Freeze()

	if err := MachineConfigTypeID.Set(c.String(ConfigTypeIDParamKey)); err != nil {
		log.Fatalf("Failed to set MachineConfigTypeID: %s", err)
	}
	log.Debugf("MachineConfigTypeID: %s", MachineConfigTypeID)

	if err := MachineParamsAdditionalEnvs.Set(c.StringSlice(EnvironmentParamKey)); err != nil {
		log.Fatalf("Failed to set MachineParamsAdditionalEnvs: %s", err)
	}
	log.Debugf("MachineParamsAdditionalEnvs: %s", MachineParamsAdditionalEnvs)
	MachineParamsAdditionalEnvs.Freeze()

	return nil
}

func printVersion(c *cli.Context) {
	fmt.Fprintf(c.App.Writer, "%v\n", c.App.Version)
}

// Run the Envman CLI.
func Run() {
	cli.VersionPrinter = printVersion

	app := cli.NewApp()
	app.Name = path.Base(os.Args[0])
	app.Usage = "bitrise-machine"
	app.Version = "0.9.7"

	app.Author = ""
	app.Email = ""

	app.Before = before

	app.Flags = appFlags
	app.Commands = commands

	if cpuProfileFilePath := os.Getenv(CPUProfileFilePathEnvKey); cpuProfileFilePath != "" {
		log.Infof("Enabling CPU Profiler, writing results into file: %s", cpuProfileFilePath)
		f, err := os.Create(cpuProfileFilePath)
		if err != nil {
			log.Fatalf("CPU Profile: failed to create file: %s", err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatalf("CPU Profile: failed to start CPU profiler: %s", err)
		}
		defer pprof.StopCPUProfile()
	}

	if err := app.Run(os.Args); err != nil {
		log.Fatal("Finished with error:", err)
	}
}
<|endoftext|>


<commit_before>package kapacitor

import (
	"fmt"
	"log"
	"sort"
	"sync"
	"time"

	"github.com/influxdata/kapacitor/expvar"
	"github.com/influxdata/kapacitor/models"
	"github.com/influxdata/kapacitor/pipeline"
	"github.com/influxdata/kapacitor/tick/stateful"
)

type CombineNode struct {
	node
	c *pipeline.CombineNode

	expressions        []stateful.Expression
	expressionsByGroup map[models.GroupID][]stateful.Expression
	scopePools         []stateful.ScopePool

	cardinalityMu sync.RWMutex

	combination combination
}

// Create a new CombineNode, which combines a stream with itself dynamically.
func newCombineNode(et *ExecutingTask, n *pipeline.CombineNode, l *log.Logger) (*CombineNode, error) {
	cn := &CombineNode{
		c:                  n,
		node:               node{Node: n, et: et, logger: l},
		expressionsByGroup: make(map[models.GroupID][]stateful.Expression),
		combination:        combination{max: n.Max},
	}

	// Create stateful expressions
	cn.expressions = make([]stateful.Expression, len(n.Lambdas))
	cn.scopePools = make([]stateful.ScopePool, len(n.Lambdas))
	for i, lambda := range n.Lambdas {
		statefulExpr, err := stateful.NewExpression(lambda.Expression)
		if err != nil {
			return nil, fmt.Errorf("Failed to compile %v expression: %v", i, err)
		}
		cn.expressions[i] = statefulExpr
		cn.scopePools[i] = stateful.NewScopePool(stateful.FindReferenceVariables(lambda.Expression))
	}
	cn.node.runF = cn.runCombine
	return cn, nil
}

type buffer struct {
	Time       time.Time
	Name       string
	Group      models.GroupID
	Dimensions models.Dimensions
	Points     []rawPoint
}

type timeList []time.Time

func (t timeList) Len() int           { return len(t) }
func (t timeList) Less(i, j int) bool { return t[i].Before(t[j]) }
func (t timeList) Swap(i, j int)      { t[i], t[j] = t[j], t[i] }

func (n *CombineNode) runCombine([]byte) error {
	valueF := func() int64 {
		n.cardinalityMu.RLock()
		l := len(n.expressionsByGroup)
		n.cardinalityMu.RUnlock()
		return int64(l)
	}
	n.statMap.Set(statCardinalityGauge, expvar.NewIntFuncGauge(valueF))

	switch n.Wants() {
	case pipeline.StreamEdge:
		buffers := make(map[models.GroupID]*buffer)
		for p, ok := n.ins[0].NextPoint(); ok; p, ok = n.ins[0].NextPoint() {
			n.timer.Start()
			t := p.Time.Round(n.c.Tolerance)
			currentBuf, ok := buffers[p.Group]
			if !ok {
				currentBuf = &buffer{
					Time:       t,
					Name:       p.Name,
					Group:      p.Group,
					Dimensions: p.Dimensions,
				}
				buffers[p.Group] = currentBuf
			}
			rp := rawPoint{
				Time:   t,
				Fields: p.Fields,
				Tags:   p.Tags,
			}
			if t.Equal(currentBuf.Time) {
				currentBuf.Points = append(currentBuf.Points, rp)
			} else {
				if err := n.combineBuffer(currentBuf); err != nil {
					return err
				}
				currentBuf.Time = t
				currentBuf.Name = p.Name
				currentBuf.Group = p.Group
				currentBuf.Dimensions = p.Dimensions
				currentBuf.Points = currentBuf.Points[0:1]
				currentBuf.Points[0] = rp
			}
			n.timer.Stop()
		}
	case pipeline.BatchEdge:
		allBuffers := make(map[models.GroupID]map[time.Time]*buffer)
		groupTimes := make(map[models.GroupID]time.Time)
		for b, ok := n.ins[0].NextBatch(); ok; b, ok = n.ins[0].NextBatch() {
			n.timer.Start()
			t := b.TMax.Round(n.c.Tolerance)
			buffers, ok := allBuffers[b.Group]
			if !ok {
				buffers = make(map[time.Time]*buffer)
				allBuffers[b.Group] = buffers
				groupTimes[b.Group] = t
			}
			groupTime := groupTimes[b.Group]
			if !t.Equal(groupTime) {
				// Set new groupTime
				groupTimes[b.Group] = t
				// Combine/Emit all old buffers
				times := make(timeList, 0, len(buffers))
				for t := range buffers {
					times = append(times, t)
				}
				sort.Sort(times)
				for _, t := range times {
					if err := n.combineBuffer(buffers[t]); err != nil {
						return err
					}
					delete(buffers, t)
				}
			}
			for _, p := range b.Points {
				t := p.Time.Round(n.c.Tolerance)
				currentBuf, ok := buffers[t]
				if !ok {
					currentBuf = &buffer{
						Time:       t,
						Name:       b.Name,
						Group:      b.Group,
						Dimensions: b.PointDimensions(),
					}
					buffers[t] = currentBuf
				}
				currentBuf.Points = append(currentBuf.Points, rawPoint{
					Time:   t,
					Fields: p.Fields,
					Tags:   p.Tags,
				})
			}
			n.timer.Stop()
		}
	}
	return nil
}

// Simple container for point data.
type rawPoint struct {
	Time   time.Time
	Fields models.Fields
	Tags   models.Tags
}

// Combine a set of points into all their combinations.
func (n *CombineNode) combineBuffer(buf *buffer) error {
	if len(buf.Points) == 0 {
		return nil
	}
	l := len(n.expressions)
	n.cardinalityMu.RLock()
	expressions, ok := n.expressionsByGroup[buf.Group]
	n.cardinalityMu.RUnlock()
	if !ok {
		expressions = make([]stateful.Expression, l)
		for i, expr := range n.expressions {
			expressions[i] = expr.CopyReset()
		}
		n.cardinalityMu.Lock()
		n.expressionsByGroup[buf.Group] = expressions
		n.cardinalityMu.Unlock()
	}

	// Compute matching result for all points
	matches := make([]map[int]bool, l)
	for i := 0; i < l; i++ {
		matches[i] = make(map[int]bool, len(buf.Points))
	}
	for idx, p := range buf.Points {
		for i := range expressions {
			matched, err := EvalPredicate(expressions[i], n.scopePools[i], p.Time, p.Fields, p.Tags)
			if err != nil {
				n.incrementErrorCount()
				n.logger.Println("E! evaluating lambda expression:", err)
			}
			matches[i][idx] = matched
		}
	}

	p := models.Point{
		Name:       buf.Name,
		Group:      buf.Group,
		Dimensions: buf.Dimensions,
	}
	dimensions := p.Dimensions.ToSet()
	set := make([]rawPoint, l)
	return n.combination.Do(len(buf.Points), l, func(indices []int) error {
		valid := true
		for s := 0; s < l; s++ {
			found := false
			for i := range indices {
				if matches[s][indices[i]] {
					set[s] = buf.Points[indices[i]]
					indices = append(indices[0:i], indices[i+1:]...)
					found = true
					break
				}
			}
			if !found {
				valid = false
				break
			}
		}
		if valid {
			rp := n.merge(set, dimensions)

			p.Time = rp.Time.Round(n.c.Tolerance)
			p.Fields = rp.Fields
			p.Tags = rp.Tags

			n.timer.Pause()
			for _, out := range n.outs {
				err := out.CollectPoint(p)
				if err != nil {
					return err
				}
			}
			n.timer.Resume()
		}
		return nil
	})
}

// Merge a set of points into a single point.
func (n *CombineNode) merge(points []rawPoint, dimensions map[string]bool) rawPoint {
	fields := make(models.Fields, len(points[0].Fields)*len(points))
	tags := make(models.Tags, len(points[0].Tags)*len(points))

	for i, p := range points {
		for field, value := range p.Fields {
			fields[n.c.Names[i]+n.c.Delimiter+field] = value
		}
		for tag, value := range p.Tags {
			if !dimensions[tag] {
				tags[n.c.Names[i]+n.c.Delimiter+tag] = value
			} else {
				tags[tag] = value
			}
		}
	}

	return rawPoint{
		Time:   points[0].Time,
		Fields: fields,
		Tags:   tags,
	}
}

// Type for performing actions on a set of combinations.
type combination struct {
	max int64
}

// Do action for each combination, based on combinatorial logic n choose k.
// If n choose k > max an error is returned
func (c combination) Do(n, k int, f func(indices []int) error) error {
	if count := c.Count(int64(n), int64(k)); count > c.max {
		return fmt.Errorf("refusing to perform combination as total combinations %d exceeds max combinations %d", count, c.max)
	} else if count == -1 {
		// Nothing to do
		return nil
	}

	indices := make([]int, k)
	indicesCopy := make([]int, k)
	for i := 0; i < k; i++ {
		indices[i] = i
	}
	copy(indicesCopy, indices)
	if err := f(indicesCopy); err != nil {
		return err
	}
	for {
		i := k - 1
		for ; i >= 0; i-- {
			if indices[i] != i+n-k {
				break
			}
		}
		if i == -1 {
			return nil
		}
		indices[i]++
		for j := i + 1; j < k; j++ {
			indices[j] = indices[j-1] + 1
		}
		copy(indicesCopy, indices)
		if err := f(indicesCopy); err != nil {
			return err
		}
	}
}

// Count the number of possible combinations of n choose k.
func (c combination) Count(n, k int64) int64 {
	if n < k {
		return -1
	}
	count := int64(1)
	for i := int64(0); i < k; i++ {
		count = (count * (n - i)) / (i + 1)
	}
	return count
}
<commit_msg>Change cardinalityMu to expressionsByGroupMu<commit_after>package kapacitor

import (
	"fmt"
	"log"
	"sort"
	"sync"
	"time"

	"github.com/influxdata/kapacitor/expvar"
	"github.com/influxdata/kapacitor/models"
	"github.com/influxdata/kapacitor/pipeline"
	"github.com/influxdata/kapacitor/tick/stateful"
)

type CombineNode struct {
	node
	c *pipeline.CombineNode

	expressions        []stateful.Expression
	expressionsByGroup map[models.GroupID][]stateful.Expression
	scopePools         []stateful.ScopePool

	expressionsByGroupMu sync.RWMutex

	combination combination
}

// Create a new CombineNode, which combines a stream with itself dynamically.
func newCombineNode(et *ExecutingTask, n *pipeline.CombineNode, l *log.Logger) (*CombineNode, error) {
	cn := &CombineNode{
		c:                  n,
		node:               node{Node: n, et: et, logger: l},
		expressionsByGroup: make(map[models.GroupID][]stateful.Expression),
		combination:        combination{max: n.Max},
	}

	// Create stateful expressions
	cn.expressions = make([]stateful.Expression, len(n.Lambdas))
	cn.scopePools = make([]stateful.ScopePool, len(n.Lambdas))
	for i, lambda := range n.Lambdas {
		statefulExpr, err := stateful.NewExpression(lambda.Expression)
		if err != nil {
			return nil, fmt.Errorf("Failed to compile %v expression: %v", i, err)
		}
		cn.expressions[i] = statefulExpr
		cn.scopePools[i] = stateful.NewScopePool(stateful.FindReferenceVariables(lambda.Expression))
	}
	cn.node.runF = cn.runCombine
	return cn, nil
}

type buffer struct {
	Time       time.Time
	Name       string
	Group      models.GroupID
	Dimensions models.Dimensions
	Points     []rawPoint
}

type timeList []time.Time

func (t timeList) Len() int           { return len(t) }
func (t timeList) Less(i, j int) bool { return t[i].Before(t[j]) }
func (t timeList) Swap(i, j int)      { t[i], t[j] = t[j], t[i] }

func (n *CombineNode) runCombine([]byte) error {
	valueF := func() int64 {
		n.expressionsByGroupMu.RLock()
		l := len(n.expressionsByGroup)
		n.expressionsByGroupMu.RUnlock()
		return int64(l)
	}
	n.statMap.Set(statCardinalityGauge, expvar.NewIntFuncGauge(valueF))

	switch n.Wants() {
	case pipeline.StreamEdge:
		buffers := make(map[models.GroupID]*buffer)
		for p, ok := n.ins[0].NextPoint(); ok; p, ok = n.ins[0].NextPoint() {
			n.timer.Start()
			t := p.Time.Round(n.c.Tolerance)
			currentBuf, ok := buffers[p.Group]
			if !ok {
				currentBuf = &buffer{
					Time:       t,
					Name:       p.Name,
					Group:      p.Group,
					Dimensions: p.Dimensions,
				}
				buffers[p.Group] = currentBuf
			}
			rp := rawPoint{
				Time:   t,
				Fields: p.Fields,
				Tags:   p.Tags,
			}
			if t.Equal(currentBuf.Time) {
				currentBuf.Points = append(currentBuf.Points, rp)
			} else {
				if err := n.combineBuffer(currentBuf); err != nil {
					return err
				}
				currentBuf.Time = t
				currentBuf.Name = p.Name
				currentBuf.Group = p.Group
				currentBuf.Dimensions = p.Dimensions
				currentBuf.Points = currentBuf.Points[0:1]
				currentBuf.Points[0] = rp
			}
			n.timer.Stop()
		}
	case pipeline.BatchEdge:
		allBuffers := make(map[models.GroupID]map[time.Time]*buffer)
		groupTimes := make(map[models.GroupID]time.Time)
		for b, ok := n.ins[0].NextBatch(); ok; b, ok = n.ins[0].NextBatch() {
			n.timer.Start()
			t := b.TMax.Round(n.c.Tolerance)
			buffers, ok := allBuffers[b.Group]
			if !ok {
				buffers = make(map[time.Time]*buffer)
				allBuffers[b.Group] = buffers
				groupTimes[b.Group] = t
			}
			groupTime := groupTimes[b.Group]
			if !t.Equal(groupTime) {
				// Set new groupTime
				groupTimes[b.Group] = t
				// Combine/Emit all old buffers
				times := make(timeList, 0, len(buffers))
				for t := range buffers {
					times = append(times, t)
				}
				sort.Sort(times)
				for _, t := range times {
					if err := n.combineBuffer(buffers[t]); err != nil {
						return err
					}
					delete(buffers, t)
				}
			}
			for _, p := range b.Points {
				t := p.Time.Round(n.c.Tolerance)
				currentBuf, ok := buffers[t]
				if !ok {
					currentBuf = &buffer{
						Time:       t,
						Name:       b.Name,
						Group:      b.Group,
						Dimensions: b.PointDimensions(),
					}
					buffers[t] = currentBuf
				}
				currentBuf.Points = append(currentBuf.Points, rawPoint{
					Time:   t,
					Fields: p.Fields,
					Tags:   p.Tags,
				})
			}
			n.timer.Stop()
		}
	}
	return nil
}

// Simple container for point data.
type rawPoint struct {
	Time   time.Time
	Fields models.Fields
	Tags   models.Tags
}

// Combine a set of points into all their combinations.
func (n *CombineNode) combineBuffer(buf *buffer) error {
	if len(buf.Points) == 0 {
		return nil
	}
	l := len(n.expressions)
	n.expressionsByGroupMu.RLock()
	expressions, ok := n.expressionsByGroup[buf.Group]
	n.expressionsByGroupMu.RUnlock()
	if !ok {
		expressions = make([]stateful.Expression, l)
		for i, expr := range n.expressions {
			expressions[i] = expr.CopyReset()
		}
		n.expressionsByGroupMu.Lock()
		n.expressionsByGroup[buf.Group] = expressions
		n.expressionsByGroupMu.Unlock()
	}

	// Compute matching result for all points
	matches := make([]map[int]bool, l)
	for i := 0; i < l; i++ {
		matches[i] = make(map[int]bool, len(buf.Points))
	}
	for idx, p := range buf.Points {
		for i := range expressions {
			matched, err := EvalPredicate(expressions[i], n.scopePools[i], p.Time, p.Fields, p.Tags)
			if err != nil {
				n.incrementErrorCount()
				n.logger.Println("E! evaluating lambda expression:", err)
			}
			matches[i][idx] = matched
		}
	}

	p := models.Point{
		Name:       buf.Name,
		Group:      buf.Group,
		Dimensions: buf.Dimensions,
	}
	dimensions := p.Dimensions.ToSet()
	set := make([]rawPoint, l)
	return n.combination.Do(len(buf.Points), l, func(indices []int) error {
		valid := true
		for s := 0; s < l; s++ {
			found := false
			for i := range indices {
				if matches[s][indices[i]] {
					set[s] = buf.Points[indices[i]]
					indices = append(indices[0:i], indices[i+1:]...)
					found = true
					break
				}
			}
			if !found {
				valid = false
				break
			}
		}
		if valid {
			rp := n.merge(set, dimensions)

			p.Time = rp.Time.Round(n.c.Tolerance)
			p.Fields = rp.Fields
			p.Tags = rp.Tags

			n.timer.Pause()
			for _, out := range n.outs {
				err := out.CollectPoint(p)
				if err != nil {
					return err
				}
			}
			n.timer.Resume()
		}
		return nil
	})
}

// Merge a set of points into a single point.
func (n *CombineNode) merge(points []rawPoint, dimensions map[string]bool) rawPoint {
	fields := make(models.Fields, len(points[0].Fields)*len(points))
	tags := make(models.Tags, len(points[0].Tags)*len(points))

	for i, p := range points {
		for field, value := range p.Fields {
			fields[n.c.Names[i]+n.c.Delimiter+field] = value
		}
		for tag, value := range p.Tags {
			if !dimensions[tag] {
				tags[n.c.Names[i]+n.c.Delimiter+tag] = value
			} else {
				tags[tag] = value
			}
		}
	}

	return rawPoint{
		Time:   points[0].Time,
		Fields: fields,
		Tags:   tags,
	}
}

// Type for performing actions on a set of combinations.
type combination struct {
	max int64
}

// Do action for each combination, based on combinatorial logic n choose k.
// If n choose k > max an error is returned
func (c combination) Do(n, k int, f func(indices []int) error) error {
	if count := c.Count(int64(n), int64(k)); count > c.max {
		return fmt.Errorf("refusing to perform combination as total combinations %d exceeds max combinations %d", count, c.max)
	} else if count == -1 {
		// Nothing to do
		return nil
	}

	indices := make([]int, k)
	indicesCopy := make([]int, k)
	for i := 0; i < k; i++ {
		indices[i] = i
	}
	copy(indicesCopy, indices)
	if err := f(indicesCopy); err != nil {
		return err
	}
	for {
		i := k - 1
		for ; i >= 0; i-- {
			if indices[i] != i+n-k {
				break
			}
		}
		if i == -1 {
			return nil
		}
		indices[i]++
		for j := i + 1; j < k; j++ {
			indices[j] = indices[j-1] + 1
		}
		copy(indicesCopy, indices)
		if err := f(indicesCopy); err != nil {
			return err
		}
	}
}

// Count the number of possible combinations of n choose k.
func (c combination) Count(n, k int64) int64 {
	if n < k {
		return -1
	}
	count := int64(1)
	for i := int64(0); i < k; i++ {
		count = (count * (n - i)) / (i + 1)
	}
	return count
}
<|endoftext|>


<commit_before>// libasink v0.0.1
//
// (c) Ground Six
//
// @package libasink
// @version 0.0.1
//
// @author Harry Lawrence <http://github.com/hazbo>
//
// License: MIT
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.

package asink

import (
    "os"
    "os/exec"
    "sync"
    "strings"
)

// Command represents the data needed to run a
// command through Asink with its Args and
// number of times for it to be executed.
type Command struct {
    Name       string
    AsyncCount int
    RelCount   int
    Dir        string
    Args       []string
    Env        []string
    Callback   func(command string)
    Dummy      bool
}

// NewCommand creates a new instance of Command with some
// default values. The command string is the
// only initial value that is required.
func NewCommand(name string) Command {
    return Command{
        Name: name,
        AsyncCount: 1,
        RelCount: 1,
        Dir: getWorkingDirectory(),
        Args: []string{},
        Env:  []string{},
        Output: "",
        Callback: func(command string){},
        Dummy: false,
    }
}

// Exec implemented to satisfy the task's Execer
// interface. Loops through the AsyncCount
// to concurrently execute the command.
func (c Command) Exec() bool {
    var wg sync.WaitGroup

    command := make(chan Command)

    c.setenv()
    c.getenv()
    c.chdir()

    for i := 0; i != c.AsyncCount; i++ {
        wg.Add(1)
        go runCommand(command, &wg)
        command <- c
    }

    close(command)
    wg.Wait()
    return true
}

// Sets env vars before executing a command.
func (c Command) setenv() {
    for _, e := range c.Env {
        kv := strings.Split(e, "=")
        os.Setenv(kv[0], kv[1])
    }
}

// Gets env vars from the Name, Args and Env
// parts of the command object.
func (c *Command) getenv() {
    c.Name = os.ExpandEnv(c.Name)
    c.Dir = os.ExpandEnv(c.Dir)
    for ai, a := range c.Args {
        c.Args[ai] = os.ExpandEnv(a)
    }
    for ei, e := range c.Env {
        c.Env[ei] = os.ExpandEnv(e)
    }
    c.setenv()
}

// Changes to a specified dir before executing
// a command.
func (c Command) chdir() {
    os.Chdir(getWorkingDirectory())
    os.Chdir(c.Dir)
}

// Is called within Exec, the actual command
// execution happens in here.
func runCommand(command chan Command, wg *sync.WaitGroup) {
    defer wg.Done()

    c := <- command

    for j := 0; j != c.RelCount; j++ {
        c.Callback(c.Name + " " + strings.Join(c.Args, " "))
        if c.Dummy == false {
            cmd := exec.Command(c.Name, c.Args...)
            cmd.Stdout = os.Stdout
            cmd.Stderr = os.Stderr
            cmd.Run()
        }
    }
}
<commit_msg>Bug, should have ran tests<commit_after>// libasink v0.0.1
//
// (c) Ground Six
//
// @package libasink
// @version 0.0.1
//
// @author Harry Lawrence <http://github.com/hazbo>
//
// License: MIT
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.

package asink

import (
    "os"
    "os/exec"
    "sync"
    "strings"
)

// Command represents the data needed to run a
// command through Asink with its Args and
// number of times for it to be executed.
type Command struct {
    Name       string
    AsyncCount int
    RelCount   int
    Dir        string
    Args       []string
    Env        []string
    Callback   func(command string)
    Dummy      bool
}

// NewCommand creates a new instance of Command with some
// default values. The command string is the
// only initial value that is required.
func NewCommand(name string) Command {
    return Command{
        Name: name,
        AsyncCount: 1,
        RelCount: 1,
        Dir: getWorkingDirectory(),
        Args: []string{},
        Env:  []string{},
        Callback: func(command string){},
        Dummy: false,
    }
}

// Exec implemented to satisfy the task's Execer
// interface. Loops through the AsyncCount
// to concurrently execute the command.
func (c Command) Exec() bool {
    var wg sync.WaitGroup

    command := make(chan Command)

    c.setenv()
    c.getenv()
    c.chdir()

    for i := 0; i != c.AsyncCount; i++ {
        wg.Add(1)
        go runCommand(command, &wg)
        command <- c
    }

    close(command)
    wg.Wait()
    return true
}

// Sets env vars before executing a command.
func (c Command) setenv() {
    for _, e := range c.Env {
        kv := strings.Split(e, "=")
        os.Setenv(kv[0], kv[1])
    }
}

// Gets env vars from the Name, Args and Env
// parts of the command object.
func (c *Command) getenv() {
    c.Name = os.ExpandEnv(c.Name)
    c.Dir = os.ExpandEnv(c.Dir)
    for ai, a := range c.Args {
        c.Args[ai] = os.ExpandEnv(a)
    }
    for ei, e := range c.Env {
        c.Env[ei] = os.ExpandEnv(e)
    }
    c.setenv()
}

// Changes to a specified dir before executing
// a command.
func (c Command) chdir() {
    os.Chdir(getWorkingDirectory())
    os.Chdir(c.Dir)
}

// Is called within Exec, the actual command
// execution happens in here.
func runCommand(command chan Command, wg *sync.WaitGroup) {
    defer wg.Done()

    c := <- command

    for j := 0; j != c.RelCount; j++ {
        c.Callback(c.Name + " " + strings.Join(c.Args, " "))
        if c.Dummy == false {
            cmd := exec.Command(c.Name, c.Args...)
            cmd.Stdout = os.Stdout
            cmd.Stderr = os.Stderr
            cmd.Run()
        }
    }
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/md5"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/brotherlogic/goserver"
	"golang.org/x/net/context"
	"google.golang.org/grpc"

	pbd "github.com/brotherlogic/discovery/proto"
	pb "github.com/brotherlogic/gobuildslave/proto"
	pbs "github.com/brotherlogic/goserver/proto"
)

// Server the main server type
type Server struct {
	*goserver.GoServer
	runner *Runner
	disk   diskChecker
	jobs   map[string]*pb.JobDetails
}

func (s *Server) monitor(job *pb.JobDetails) {
	for true {
		switch job.State {
		case pb.JobDetails_ACKNOWLEDGED:
			job.StartTime = 0
			job.State = pb.JobDetails_BUILDING
			s.runner.Checkout(job.GetSpec().Name)
			job.State = pb.JobDetails_BUILT
		case pb.JobDetails_BUILT:
			s.runner.Run(job)
			for job.StartTime == 0 {
				time.Sleep(waitTime)
			}
			job.State = pb.JobDetails_PENDING
		case pb.JobDetails_KILLING:
			s.runner.kill(job)
			if !isAlive(job.GetSpec()) {
				job.State = pb.JobDetails_DEAD
			}
		case pb.JobDetails_UPDATE_STARTING:
			s.runner.Update(job)
			job.State = pb.JobDetails_RUNNING
		case pb.JobDetails_PENDING:
			time.Sleep(time.Minute)
			if isAlive(job.GetSpec()) {
				job.State = pb.JobDetails_RUNNING
			} else {
				job.State = pb.JobDetails_DEAD
			}
		case pb.JobDetails_RUNNING:
			time.Sleep(waitTime)
			if !isAlive(job.GetSpec()) {
				job.State = pb.JobDetails_DEAD
			}
		case pb.JobDetails_DEAD:
			log.Printf("RERUNNING BECAUSE WERE DEAD (%v)", job)
			job.State = pb.JobDetails_ACKNOWLEDGED
		}
	}
}

func getHash(file string) (string, error) {
	env := os.Environ()
	home := ""
	for _, s := range env {
		if strings.HasPrefix(s, "HOME=") {
			home = s[5:]
		}
	}

	gpath := home + "/gobuild"

	f, err := os.Open(strings.Replace(file, "$GOPATH", gpath, 1))
	if err != nil {
		return "", err
	}
	defer f.Close()

	h := md5.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}

	return string(h.Sum(nil)), nil
}

func getIP(name string, server string) (string, int) {
	conn, _ := grpc.Dial("192.168.86.64:50055", grpc.WithInsecure())
	defer conn.Close()

	registry := pbd.NewDiscoveryServiceClient(conn)
	entry := pbd.RegistryEntry{Name: name, Identifier: server}
	r, err := registry.Discover(context.Background(), &entry)

	if err != nil {
		return "", -1
	}

	return r.Ip, int(r.Port)
}

// updateState of the runner command
func isAlive(spec *pb.JobSpec) bool {
	elems := strings.Split(spec.Name, "/")
	dServer, dPort := getIP(elems[len(elems)-1], spec.Server)

	if dPort > 0 {
		dConn, err := grpc.Dial(dServer+":"+strconv.Itoa(dPort), grpc.WithInsecure())
		if err != nil {
			return false
		}
		defer dConn.Close()

		c := pbs.NewGoserverServiceClient(dConn)
		resp, err := c.IsAlive(context.Background(), &pbs.Alive{})

		if err != nil || resp.Name != elems[len(elems)-1] {
			log.Printf("FOUND DEAD SERVER: (%v) %v -> %v", spec, err, resp)
		}

		return err == nil
	}

	//Mark as false if we can't locate the job
	return false
}

// DoRegister Registers this server
func (s Server) DoRegister(server *grpc.Server) {
	pb.RegisterGoBuildSlaveServer(server, &s)
}

// ReportHealth determines if the server is healthy
func (s Server) ReportHealth() bool {
	return true
}

// Mote promotes/demotes this server
func (s Server) Mote(master bool) error {
	return nil
}

//Init builds the default runner framework
func Init() *Runner {
	r := &Runner{gopath: "goautobuild", m: &sync.Mutex{}}
	r.runner = runCommand
	go r.run()
	return r
}

func runCommand(c *runnerCommand) {
	if c == nil || c.command == nil {
		return
	}

	env := os.Environ()
	home := ""
	for _, s := range env {
		if strings.HasPrefix(s, "HOME=") {
			home = s[5:]
		}
	}

	gpath := home + "/gobuild"
	c.command.Path = strings.Replace(c.command.Path, "$GOPATH", gpath, -1)
	for i := range c.command.Args {
		c.command.Args[i] = strings.Replace(c.command.Args[i], "$GOPATH", gpath, -1)
	}

	path := fmt.Sprintf("GOPATH=" + home + "/gobuild")
	found := false
	envl := os.Environ()
	for i, blah := range envl {
		if strings.HasPrefix(blah, "GOPATH") {
			envl[i] = path
			found = true
		}
	}
	if !found {
		envl = append(envl, path)
	}
	c.command.Env = envl

	out, _ := c.command.StdoutPipe()

	c.command.Start()

	if !c.background {
		str := ""

		if out != nil {
			buf := new(bytes.Buffer)
			buf.ReadFrom(out)
			str = buf.String()
		}

		c.command.Wait()
		c.output = str
		c.complete = true
	} else {
		c.details.StartTime = time.Now().Unix()
	}
}

func (diskChecker prodDiskChecker) diskUsage(path string) int64 {
	return diskUsage(path)
}

func (s *Server) rebuildLoop() {
	for true {
		time.Sleep(time.Minute * 60)

		var rebuildList []*pb.JobDetails
		var hashList []string
		for _, job := range s.runner.backgroundTasks {
			if time.Since(job.started) > time.Hour {
				rebuildList = append(rebuildList, job.details)
				hashList = append(hashList, job.hash)
			}
		}

		for i := range rebuildList {
			s.runner.Rebuild(rebuildList[i], hashList[i])
		}
	}
}

func main() {
	var quiet = flag.Bool("quiet", false, "Show all output")
	flag.Parse()

	if *quiet {
		log.SetFlags(0)
		log.SetOutput(ioutil.Discard)
	}

	s := Server{&goserver.GoServer{}, Init(), prodDiskChecker{}, make(map[string]*pb.JobDetails)}
	s.Register = s
	s.PrepServer()
	s.GoServer.Killme = false
	s.RegisterServingTask(s.rebuildLoop)
	s.RegisterServer("gobuildslave", false)
	s.Serve()
}
<commit_msg>Fixed spec of dead servering<commit_after>package main

import (
	"bytes"
	"crypto/md5"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/brotherlogic/goserver"
	"golang.org/x/net/context"
	"google.golang.org/grpc"

	pbd "github.com/brotherlogic/discovery/proto"
	pb "github.com/brotherlogic/gobuildslave/proto"
	pbs "github.com/brotherlogic/goserver/proto"
)

// Server the main server type
type Server struct {
	*goserver.GoServer
	runner *Runner
	disk   diskChecker
	jobs   map[string]*pb.JobDetails
}

func (s *Server) monitor(job *pb.JobDetails) {
	for true {
		switch job.State {
		case pb.JobDetails_ACKNOWLEDGED:
			job.StartTime = 0
			job.State = pb.JobDetails_BUILDING
			s.runner.Checkout(job.GetSpec().Name)
			job.State = pb.JobDetails_BUILT
		case pb.JobDetails_BUILT:
			s.runner.Run(job)
			for job.StartTime == 0 {
				time.Sleep(waitTime)
			}
			job.State = pb.JobDetails_PENDING
		case pb.JobDetails_KILLING:
			s.runner.kill(job)
			if !isAlive(job.GetSpec()) {
				job.State = pb.JobDetails_DEAD
			}
		case pb.JobDetails_UPDATE_STARTING:
			s.runner.Update(job)
			job.State = pb.JobDetails_RUNNING
		case pb.JobDetails_PENDING:
			time.Sleep(time.Minute)
			if isAlive(job.GetSpec()) {
				job.State = pb.JobDetails_RUNNING
			} else {
				job.State = pb.JobDetails_DEAD
			}
		case pb.JobDetails_RUNNING:
			time.Sleep(waitTime)
			if !isAlive(job.GetSpec()) {
				job.State = pb.JobDetails_DEAD
			}
		case pb.JobDetails_DEAD:
			log.Printf("RERUNNING BECAUSE WERE DEAD (%v)", job)
			job.State = pb.JobDetails_ACKNOWLEDGED
		}
	}
}

func getHash(file string) (string, error) {
	env := os.Environ()
	home := ""
	for _, s := range env {
		if strings.HasPrefix(s, "HOME=") {
			home = s[5:]
		}
	}

	gpath := home + "/gobuild"

	f, err := os.Open(strings.Replace(file, "$GOPATH", gpath, 1))
	if err != nil {
		return "", err
	}
	defer f.Close()

	h := md5.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}

	return string(h.Sum(nil)), nil
}

func getIP(name string, server string) (string, int) {
	conn, _ := grpc.Dial("192.168.86.64:50055", grpc.WithInsecure())
	defer conn.Close()

	registry := pbd.NewDiscoveryServiceClient(conn)
	entry := pbd.RegistryEntry{Name: name, Identifier: server}
	r, err := registry.Discover(context.Background(), &entry)

	if err != nil {
		return "", -1
	}

	return r.Ip, int(r.Port)
}

// updateState of the runner command
func isAlive(spec *pb.JobSpec) bool {
	elems := strings.Split(spec.Name, "/")
	dServer, dPort := getIP(elems[len(elems)-1], spec.Server)

	if dPort > 0 {
		dConn, err := grpc.Dial(dServer+":"+strconv.Itoa(dPort), grpc.WithInsecure())
		if err != nil {
			return false
		}
		defer dConn.Close()

		c := pbs.NewGoserverServiceClient(dConn)
		resp, err := c.IsAlive(context.Background(), &pbs.Alive{})

		if err != nil || resp.Name != elems[len(elems)-1] {
			log.Printf("FOUND DEAD SERVER: (%v) %v -> %v", spec, err, resp)
			return false
		}

		return true
	}

	//Mark as false if we can't locate the job
	return false
}

// DoRegister Registers this server
func (s Server) DoRegister(server *grpc.Server) {
	pb.RegisterGoBuildSlaveServer(server, &s)
}

// ReportHealth determines if the server is healthy
func (s Server) ReportHealth() bool {
	return true
}

// Mote promotes/demotes this server
func (s Server) Mote(master bool) error {
	return nil
}

//Init builds the default runner framework
func Init() *Runner {
	r := &Runner{gopath: "goautobuild", m: &sync.Mutex{}}
	r.runner = runCommand
	go r.run()
	return r
}

func runCommand(c *runnerCommand) {
	if c == nil || c.command == nil {
		return
	}

	env := os.Environ()
	home := ""
	for _, s := range env {
		if strings.HasPrefix(s, "HOME=") {
			home = s[5:]
		}
	}

	gpath := home + "/gobuild"
	c.command.Path = strings.Replace(c.command.Path, "$GOPATH", gpath, -1)
	for i := range c.command.Args {
		c.command.Args[i] = strings.Replace(c.command.Args[i], "$GOPATH", gpath, -1)
	}

	path := fmt.Sprintf("GOPATH=" + home + "/gobuild")
	found := false
	envl := os.Environ()
	for i, blah := range envl {
		if strings.HasPrefix(blah, "GOPATH") {
			envl[i] = path
			found = true
		}
	}
	if !found {
		envl = append(envl, path)
	}
	c.command.Env = envl

	out, _ := c.command.StdoutPipe()

	c.command.Start()

	if !c.background {
		str := ""

		if out != nil {
			buf := new(bytes.Buffer)
			buf.ReadFrom(out)
			str = buf.String()
		}

		c.command.Wait()
		c.output = str
		c.complete = true
	} else {
		c.details.StartTime = time.Now().Unix()
	}
}

func (diskChecker prodDiskChecker) diskUsage(path string) int64 {
	return diskUsage(path)
}

func (s *Server) rebuildLoop() {
	for true {
		time.Sleep(time.Minute * 60)

		var rebuildList []*pb.JobDetails
		var hashList []string
		for _, job := range s.runner.backgroundTasks {
			if time.Since(job.started) > time.Hour {
				rebuildList = append(rebuildList, job.details)
				hashList = append(hashList, job.hash)
			}
		}

		for i := range rebuildList {
			s.runner.Rebuild(rebuildList[i], hashList[i])
		}
	}
}

func main() {
	var quiet = flag.Bool("quiet", false, "Show all output")
	flag.Parse()

	if *quiet {
		log.SetFlags(0)
		log.SetOutput(ioutil.Discard)
	}

	s := Server{&goserver.GoServer{}, Init(), prodDiskChecker{}, make(map[string]*pb.JobDetails)}
	s.Register = s
	s.PrepServer()
	s.GoServer.Killme = false
	s.RegisterServingTask(s.rebuildLoop)
	s.RegisterServer("gobuildslave", false)
	s.Serve()
}
<|endoftext|>


<commit_before>package timeutil

import (
	"time"
)

var (
	TimeFormat = "2006-01-02T15:04:05"
)

func ParseTime(timeStr string) time.Time {
	loc, _ := time.LoadLocation("UTC")
	t, _ := time.ParseInLocation(TimeFormat, timeStr, loc)

	return t
}

func FormatTimestamp(ts int64) string {
	return time.Unix(ts, 0).Format(TimeFormat)
}
<commit_msg>minor refactoring<commit_after>package timeutil

import (
	"time"
)

var (
	/*TimeFormat  yyyy-mm-ddTHH:MM:ss time format */
	TimeFormat = "2006-01-02T15:04:05"
)

func ParseTime(timeStr string) time.Time {
	loc, _ := time.LoadLocation("UTC")
	t, _ := time.ParseInLocation(TimeFormat, timeStr, loc)

	return t
}

func FormatTimestamp(ts int64) string {
	return time.Unix(ts, 0).Format(TimeFormat)
}
<|endoftext|>


<commit_before>package butlerd

import (
	"context"
	"encoding/json"
	"io"
	"io/ioutil"
	"net/http"
	"strconv"

	"github.com/sourcegraph/jsonrpc2"
)

type httpCallStream struct {
	cid    string
	method string

	w  http.ResponseWriter
	r  *http.Request
	hh *httpHandler

	jrh jsonrpc2.Handler

	ctx    context.Context
	cancel context.CancelFunc

	readCh chan []byte
}

var _ jsonrpc2.ObjectStream = (*httpCallStream)(nil)
var _ jsonrpc2.Handler = (*httpCallStream)(nil)

func (s *httpCallStream) ReadObject(v interface{}) error {
	select {
	case msg := <-s.readCh:
		return json.Unmarshal(msg, v)
	case <-s.ctx.Done():
		return io.EOF
	}
}

func (s *httpCallStream) WriteObject(obj interface{}) error {
	marshalled, err := json.Marshal(obj)
	if err != nil {
		return err
	}

	intermediate := make(map[string]interface{})

	err = json.Unmarshal(marshalled, &intermediate)
	if err != nil {
		return err
	}

	_, hasError := intermediate["error"]
	_, hasResult := intermediate["result"]
	if hasError || hasResult {
		// responses are written as http responses
		s.w.Header().Set("content-type", "application/json")
		s.w.Header().Set("cache-control", "no-cache")
		s.w.WriteHeader(200)
		s.w.Write(marshalled)
		return nil
	}

	allowFailures := false
	_, hasMethod := intermediate["method"]
	if !hasMethod {
		// then it must be a notification
		allowFailures = true
	}

	if s.cid == "" {
		return HTTPError(428, "Server tried to send a request, but cid was not specified for this call")
	}

	// notifications or server-side requests are sent to event-stream
	fs, ok := s.hh.getFeedStream(s.cid)
	if !ok {
		if allowFailures {
			return nil
		}
		return HTTPError(428, "Server tried to send a request, but client is not listening to feed")
	}

	fs.requestCh <- marshalled
	return nil
}

func (s *httpCallStream) Handle(ctx context.Context, conn *jsonrpc2.Conn, req *jsonrpc2.Request) {
	// handle asynchronously so we can process server->client requests
	go func() {
		defer s.cancel()
		s.jrh.Handle(ctx, conn, req)
	}()
}

func (s *httpCallStream) Close() error {
	s.cancel()
	return nil
}

func (s *httpCallStream) Wait(parentCtx context.Context) error {
	idString := s.r.Header.Get("x-id")
	if idString == "" {
		return HTTPError(400, "Missing request ID x-id")
	}
	id, err := strconv.ParseInt(idString, 10, 64)
	if err != nil {
		return HTTPError(400, "x-id must be an integer")
	}

	body, err := ioutil.ReadAll(s.r.Body)
	if err != nil {
		return err
	}

	s.readCh = make(chan []byte, 1)
	req := map[string]interface{}{
		"id":     id,
		"method": s.method,
		"params": json.RawMessage(body),
	}
	reqJSON, err := json.Marshal(req)
	if err != nil {
		return err
	}
	s.readCh <- reqJSON

	s.hh.putCallStream(s.cid, s)
	defer s.hh.removeCallStream(s.cid)

	ctx, cancel := context.WithCancel(parentCtx)
	s.ctx = ctx
	s.cancel = cancel
	defer cancel()

	conn := jsonrpc2.NewConn(ctx, s, s)
	<-conn.DisconnectNotify()
	return nil
}

func (s *httpCallStream) cancelWith(status int, msg string) {
	s.w.WriteHeader(status)
	s.w.Write([]byte(msg))
	s.cancel()
}
<commit_msg>Allow some slack when waiting for feedstream<commit_after>package butlerd

import (
	"context"
	"encoding/json"
	"io"
	"io/ioutil"
	"net/http"
	"strconv"
	"time"

	"github.com/sourcegraph/jsonrpc2"
)

type httpCallStream struct {
	cid    string
	method string

	w  http.ResponseWriter
	r  *http.Request
	hh *httpHandler

	jrh jsonrpc2.Handler

	ctx    context.Context
	cancel context.CancelFunc

	readCh chan []byte
}

var _ jsonrpc2.ObjectStream = (*httpCallStream)(nil)
var _ jsonrpc2.Handler = (*httpCallStream)(nil)

func (s *httpCallStream) ReadObject(v interface{}) error {
	select {
	case msg := <-s.readCh:
		return json.Unmarshal(msg, v)
	case <-s.ctx.Done():
		return io.EOF
	}
}

func (s *httpCallStream) WriteObject(obj interface{}) error {
	marshalled, err := json.Marshal(obj)
	if err != nil {
		return err
	}

	intermediate := make(map[string]interface{})

	err = json.Unmarshal(marshalled, &intermediate)
	if err != nil {
		return err
	}

	_, hasError := intermediate["error"]
	_, hasResult := intermediate["result"]
	if hasError || hasResult {
		// responses are written as http responses
		s.w.Header().Set("content-type", "application/json")
		s.w.Header().Set("cache-control", "no-cache")
		s.w.WriteHeader(200)
		s.w.Write(marshalled)
		return nil
	}

	allowFailures := false
	_, hasMethod := intermediate["method"]
	if !hasMethod {
		// then it must be a notification
		allowFailures = true
	}

	if s.cid == "" {
		return HTTPError(428, "Server tried to send a request, but cid was not specified for this call")
	}

	// notifications or server-side requests are sent to event-stream
	var fs *httpFeedStream
	var ok bool

	for tries := 200; tries > 0; tries-- {
		// allow 2s of slack
		fs, ok = s.hh.getFeedStream(s.cid)
		if ok {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}

	if !ok {
		if allowFailures {
			return nil
		}
		return HTTPError(428, "Server tried to send a request, but client is not listening to feed")
	}

	fs.requestCh <- marshalled
	return nil
}

func (s *httpCallStream) Handle(ctx context.Context, conn *jsonrpc2.Conn, req *jsonrpc2.Request) {
	// handle asynchronously so we can process server->client requests
	go func() {
		defer s.cancel()
		s.jrh.Handle(ctx, conn, req)
	}()
}

func (s *httpCallStream) Close() error {
	s.cancel()
	return nil
}

func (s *httpCallStream) Wait(parentCtx context.Context) error {
	idString := s.r.Header.Get("x-id")
	if idString == "" {
		return HTTPError(400, "Missing request ID x-id")
	}
	id, err := strconv.ParseInt(idString, 10, 64)
	if err != nil {
		return HTTPError(400, "x-id must be an integer")
	}

	body, err := ioutil.ReadAll(s.r.Body)
	if err != nil {
		return err
	}

	s.readCh = make(chan []byte, 1)
	req := map[string]interface{}{
		"id":     id,
		"method": s.method,
		"params": json.RawMessage(body),
	}
	reqJSON, err := json.Marshal(req)
	if err != nil {
		return err
	}
	s.readCh <- reqJSON

	s.hh.putCallStream(s.cid, s)
	defer s.hh.removeCallStream(s.cid)

	ctx, cancel := context.WithCancel(parentCtx)
	s.ctx = ctx
	s.cancel = cancel
	defer cancel()

	conn := jsonrpc2.NewConn(ctx, s, s)
	<-conn.DisconnectNotify()
	return nil
}

func (s *httpCallStream) cancelWith(status int, msg string) {
	s.w.WriteHeader(status)
	s.w.Write([]byte(msg))
	s.cancel()
}
<|endoftext|>


<commit_before>package goat

import (
	"bencode"
	"fmt"
	"strconv"
)

// Tracker announce request
func TrackerAnnounce(user UserRecord, query map[string]string, resChan chan []byte) {
	// Store announce information in struct
	announce := MapToAnnounceLog(query)

	// Request to store announce
	go announce.Save()

	Static.LogChan <- fmt.Sprintf("announce: [ip: %s, port:%d]", announce.Ip, announce.Port)
	Static.LogChan <- fmt.Sprintf("announce: [info_hash: %s]", announce.InfoHash)

	// Only report announce when needed
	if announce.Event != "" {
		Static.LogChan <- fmt.Sprintf("announce: [event: %s]", announce.Event)
	}

	// Check for a matching file via info_hash
	file := new(FileRecord).Load(announce.InfoHash, "info_hash")
	if file == (FileRecord{}) {
		// Torrent is not currently registered
		resChan <- TrackerError("Unregistered torrent")

		// Create an entry in file table for this hash, but mark it as unverified
		file.InfoHash = announce.InfoHash
		file.Verified = false
		file.Leechers = 0
		file.Seeders = 0
		file.Completed = 0
		go file.Save()
		return
	}

	// Ensure file is verified, meaning we will permit tracking of it
	if !file.Verified {
		resChan <- TrackerError("Unverified torrent")
		return
	}

	// Check for existing record for this user with this file
	fileUser := new(FileUserRecord).Load(file.Id, user.Id)
	if fileUser == (FileUserRecord{}) {
		// Create new relationship
		fileUser.FileId = file.Id
		fileUser.UserId = user.Id
		fileUser.Active = true
		fileUser.Completed = false
		fileUser.Announced = 1
		fileUser.Uploaded = announce.Uploaded
		fileUser.Downloaded = announce.Downloaded
		fileUser.Left = announce.Left

		// Add a leecher to this file
		file.Leechers = file.Leechers + 1
	} else {
		// Else, pre-existing record, so update
		// Check for stopped status
		if announce.Event != "stopped" {
			fileUser.Active = true
		} else {
			fileUser.Active = false
		}

		// Check for completion
		if announce.Event == "completed" && announce.Left == 0 {
			fileUser.Completed = true

			// Mark file as completed by another user
			file.Completed = file.Completed + 1

			// Decrement leecher, add seeder
			file.Leechers = file.Leechers - 1
			file.Seeders = file.Seeders + 1
		} else {
			fileUser.Completed = false
		}

		// Add an announce
		fileUser.Announced = fileUser.Announced + 1

		// Do math on statistics
		fileUser.Uploaded = fileUser.Uploaded + announce.Uploaded
		fileUser.Downloaded = fileUser.Downloaded + announce.Downloaded
		fileUser.Left = announce.Left
	}

	// Update File record
	go file.Save()

	// Insert or update the FileUser record
	go fileUser.Save()

	// Check for numwant parameter, return up to that number of peers
	// Default is 50 per protocol
	numwant := 50
	if _, ok := query["numwant"]; ok {
		// Verify numwant is an integer
		num, err := strconv.Atoi(query["numwant"])
		if err == nil {
			numwant = num
		}
	}

	// Tracker announce response, with generated peerlist of length numwant, excluding this user
	resChan <- bencode.EncDictMap(map[string][]byte{
		"interval":     bencode.EncInt(RandRange(3200, 4000)),
		"min interval": bencode.EncInt(1800),
		"peers":        bencode.EncBytes(file.PeerList(query["ip"], numwant)),
	})
}

// Report a bencoded []byte response as specified by input string
func TrackerError(err string) []byte {
	return bencode.EncDictMap(map[string][]byte{
		"failure reason": bencode.EncString(err),
		"interval":       bencode.EncInt(RandRange(3200, 4000)),
		"min interval":   bencode.EncInt(1800),
	})
}
<commit_msg>Fix logic bug in leecher count<commit_after>package goat

import (
	"bencode"
	"fmt"
	"strconv"
)

// Tracker announce request
func TrackerAnnounce(user UserRecord, query map[string]string, resChan chan []byte) {
	// Store announce information in struct
	announce := MapToAnnounceLog(query)

	// Request to store announce
	go announce.Save()

	Static.LogChan <- fmt.Sprintf("announce: [ip: %s, port:%d]", announce.Ip, announce.Port)
	Static.LogChan <- fmt.Sprintf("announce: [info_hash: %s]", announce.InfoHash)

	// Only report announce when needed
	if announce.Event != "" {
		Static.LogChan <- fmt.Sprintf("announce: [event: %s]", announce.Event)
	}

	// Check for a matching file via info_hash
	file := new(FileRecord).Load(announce.InfoHash, "info_hash")
	if file == (FileRecord{}) {
		// Torrent is not currently registered
		resChan <- TrackerError("Unregistered torrent")

		// Create an entry in file table for this hash, but mark it as unverified
		file.InfoHash = announce.InfoHash
		file.Verified = false
		file.Leechers = 0
		file.Seeders = 0
		file.Completed = 0
		go file.Save()
		return
	}

	// Ensure file is verified, meaning we will permit tracking of it
	if !file.Verified {
		resChan <- TrackerError("Unverified torrent")
		return
	}

	// Check for existing record for this user with this file
	fileUser := new(FileUserRecord).Load(file.Id, user.Id)
	if fileUser == (FileUserRecord{}) {
		// Create new relationship
		fileUser.FileId = file.Id
		fileUser.UserId = user.Id
		fileUser.Active = true
		fileUser.Completed = false
		fileUser.Announced = 1
		fileUser.Uploaded = announce.Uploaded
		fileUser.Downloaded = announce.Downloaded
		fileUser.Left = announce.Left

		// Add a leecher to this file, UNLESS they have already completed it
		if announce.Left = 0 || announce.Event = "completed" {
			file.Seeders = file.Seeders + 1
		} else {
			file.Leechers = file.Leechers + 1
		}
	} else {
		// Else, pre-existing record, so update
		// Check for stopped status
		if announce.Event != "stopped" {
			fileUser.Active = true
		} else {
			fileUser.Active = false
		}

		// Check for completion
		if announce.Event == "completed" && announce.Left == 0 {
			fileUser.Completed = true

			// Mark file as completed by another user
			file.Completed = file.Completed + 1

			// Decrement leecher, add seeder
			file.Leechers = file.Leechers - 1
			file.Seeders = file.Seeders + 1
		} else {
			fileUser.Completed = false
		}

		// Add an announce
		fileUser.Announced = fileUser.Announced + 1

		// Do math on statistics
		fileUser.Uploaded = fileUser.Uploaded + announce.Uploaded
		fileUser.Downloaded = fileUser.Downloaded + announce.Downloaded
		fileUser.Left = announce.Left
	}

	// Update File record
	go file.Save()

	// Insert or update the FileUser record
	go fileUser.Save()

	// Check for numwant parameter, return up to that number of peers
	// Default is 50 per protocol
	numwant := 50
	if _, ok := query["numwant"]; ok {
		// Verify numwant is an integer
		num, err := strconv.Atoi(query["numwant"])
		if err == nil {
			numwant = num
		}
	}

	// Tracker announce response, with generated peerlist of length numwant, excluding this user
	resChan <- bencode.EncDictMap(map[string][]byte{
		"interval":     bencode.EncInt(RandRange(3200, 4000)),
		"min interval": bencode.EncInt(1800),
		"peers":        bencode.EncBytes(file.PeerList(query["ip"], numwant)),
	})
}

// Report a bencoded []byte response as specified by input string
func TrackerError(err string) []byte {
	return bencode.EncDictMap(map[string][]byte{
		"failure reason": bencode.EncString(err),
		"interval":       bencode.EncInt(RandRange(3200, 4000)),
		"min interval":   bencode.EncInt(1800),
	})
}
<|endoftext|>


<commit_before>package gopnik

type TileCoord struct {
	X    uint64
	Y    uint64
	Zoom uint64
	Size uint64
	Tags []string
}

func (tc *TileCoord) Equals(coord *TileCoord) bool {
	if tc.Zoom == coord.Zoom && tc.X == coord.X && tc.Y == coord.Y && tc.Size == coord.Size {
		return true
	} else {
		return false
	}
}
<commit_msg>fix TileCoord Equals function<commit_after>package gopnik

type TileCoord struct {
	X    uint64
	Y    uint64
	Zoom uint64
	Size uint64
	Tags []string
}

func (tc *TileCoord) Equals(coord *TileCoord) bool {
	if tc.Zoom == coord.Zoom && tc.X == coord.X && tc.Y == coord.Y && tc.Size == coord.Size && len(tc.Tags) == len(coord.Tags) {
		for _, tag := range tc.Tags {
			for _, tag2 := range coord.Tags {
				if tag2 != tag {
					return false
				}
			}
		}
		return true
	} else {
		return false
	}
}
<|endoftext|>


<commit_before>package jet

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"reflect"
	"sync"
	"text/template"
)

type Set struct {
	Dirs      []string             // directories for look to template files
	templates map[string]*Template // parsed templates
	escapee   SafeWriter
	globals   Scope        // global scope for this template set
	tmx       sync.RWMutex // template parsing mutext
	gmx       sync.RWMutex // global variables map mutext
}

// AddGlobal (add|set)s a global variable to the set
func (s *Set) AddGlobal(key string, i interface{}) (val interface{}, override bool) {
	s.gmx.Lock()
	val, override = s.globals[key]
	if s.globals == nil {
		s.globals = make(Scope)
	}
	s.globals[key] = reflect.ValueOf(i)
	s.gmx.Unlock()
	return
}

// NewSet creates a new set, dir specifies a list of directories entries to search for templates
func NewSet(dir ...string) *Set {
	return &Set{Dirs: dir, templates: make(map[string]*Template)}
}

// NewHTMLSet creates a new set, dir specifies a list of directories entries to search for templates
func NewHTMLSet(dir ...string) *Set {
	return &Set{Dirs: dir, escapee: template.HTMLEscape, templates: make(map[string]*Template)}
}

// NewSafeSet creates a new set, dir specifies a list of directories entries to search for templates
func NewSafeSet(escapee SafeWriter, dir ...string) *Set {
	return &Set{Dirs: dir, escapee: escapee, templates: make(map[string]*Template)}
}

// AddPath adds a path to the directories entries
func (s *Set) AddPath(path string) {
	s.Dirs = append([]string{path}, s.Dirs...)
}

//AddGopathPath adds a path to directories
func (s *Set) AddGopathPath(path string) {
	paths := filepath.SplitList(os.Getenv("GOPATH"))
	for i := 0; i < len(paths); i++ {
		path, err := filepath.Abs(filepath.Join(paths[i], path))

		if err != nil {
			panic(errors.New("Can't add this path err: " + err.Error()))
		}

		if fstats, err := os.Stat(path); os.IsNotExist(err) == false && fstats.IsDir() {
			s.Dirs = append([]string{path}, s.Dirs...)
		}
	}
}

func (s *Set) load(name, content string) (template *Template, err error) {
	if content == "" {
		for i := 0; i < len(s.Dirs); i++ {
			fileName := path.Join(s.Dirs[i], name)
			var bytestring []byte
			bytestring, err = ioutil.ReadFile(fileName)
			if err == nil {
				content = string(bytestring)
				break
			}
		}
		if content == "" && err != nil {
			return
		}
	}

	template, err = s.parse(name, content)
	return
}

func (s *Set) loadTemplate(name, content string) (template *Template, err error) {
	var ok bool
	if template, ok = s.templates[name]; ok {
		return
	}
	template, err = s.load(name, content)
	s.templates[name] = template
	return
}

func (s *Set) GetTemplate(name string) (template *Template, ok bool) {
	s.tmx.RLock()
	template, ok = s.templates[name]
	s.tmx.RUnlock()
	return
}

func (s *Set) LoadTemplate(name, content string) (template *Template, err error) {
	var ok bool

	s.tmx.RLock()
	if template, ok = s.templates[name]; ok {
		s.tmx.RUnlock()
		return
	}

	s.tmx.RUnlock()
	s.tmx.Lock()
	defer s.tmx.Unlock()

	template, ok = s.templates[name]
	if ok && template != nil {
		return
	}

	template, err = s.load(name, content)
	s.templates[name] = template // saves the template
	return
}

func (t *Template) String() (template string) {

	if t.extends != nil {
		template += fmt.Sprintf("{{extends %q}}", t.extends.ParseName)
	}

	for _, _import := range t.imports {
		template += fmt.Sprintf("\n{{import %q}}", _import.ParseName)
	}

	template += t.root.String()
	return
}

func (t *Template) addBlocks(blocks map[string]*BlockNode) {
	if len(blocks) > 0 {
		if t.processedBlocks == nil {
			t.processedBlocks = make(map[string]*BlockNode)
		}
		for key, value := range blocks {
			t.processedBlocks[key] = value
		}
	}
}

type Scope map[string]reflect.Value

func (scope Scope) Set(name string, v interface{}) {
	scope[name] = reflect.ValueOf(v)
}

func (scope Scope) SetPtr(name string, v interface{}) {
	scope[name] = reflect.ValueOf(v).Elem()
}

func (t *Template) Execute(w io.Writer, variables Scope, data interface{}) (err error) {
	st := pool_State.Get().(*State)
	defer st.recover(&err)

	root := t.root
	if t.extends != nil {
		root = t.extends.root
	}
	if data != nil {
		st.context = reflect.ValueOf(data)
	}

	st.blocks = t.processedBlocks

	st.set = t.set
	st.Writer = w
	st.variables = variables

	st.executeList(root)
	return
}
<commit_msg>Remove SetPtr method from jet.Scope<commit_after>package jet

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"reflect"
	"sync"
	"text/template"
)

type Set struct {
	Dirs      []string             // directories for look to template files
	templates map[string]*Template // parsed templates
	escapee   SafeWriter
	globals   Scope        // global scope for this template set
	tmx       sync.RWMutex // template parsing mutext
	gmx       sync.RWMutex // global variables map mutext
}

// AddGlobal (add|set)s a global variable to the set
func (s *Set) AddGlobal(key string, i interface{}) (val interface{}, override bool) {
	s.gmx.Lock()
	val, override = s.globals[key]
	if s.globals == nil {
		s.globals = make(Scope)
	}
	s.globals[key] = reflect.ValueOf(i)
	s.gmx.Unlock()
	return
}

// NewSet creates a new set, dir specifies a list of directories entries to search for templates
func NewSet(dir ...string) *Set {
	return &Set{Dirs: dir, templates: make(map[string]*Template)}
}

// NewHTMLSet creates a new set, dir specifies a list of directories entries to search for templates
func NewHTMLSet(dir ...string) *Set {
	return &Set{Dirs: dir, escapee: template.HTMLEscape, templates: make(map[string]*Template)}
}

// NewSafeSet creates a new set, dir specifies a list of directories entries to search for templates
func NewSafeSet(escapee SafeWriter, dir ...string) *Set {
	return &Set{Dirs: dir, escapee: escapee, templates: make(map[string]*Template)}
}

// AddPath adds a path to the directories entries
func (s *Set) AddPath(path string) {
	s.Dirs = append([]string{path}, s.Dirs...)
}

//AddGopathPath adds a path to directories
func (s *Set) AddGopathPath(path string) {
	paths := filepath.SplitList(os.Getenv("GOPATH"))
	for i := 0; i < len(paths); i++ {
		path, err := filepath.Abs(filepath.Join(paths[i], path))

		if err != nil {
			panic(errors.New("Can't add this path err: " + err.Error()))
		}

		if fstats, err := os.Stat(path); os.IsNotExist(err) == false && fstats.IsDir() {
			s.Dirs = append([]string{path}, s.Dirs...)
		}
	}
}

func (s *Set) load(name, content string) (template *Template, err error) {
	if content == "" {
		for i := 0; i < len(s.Dirs); i++ {
			fileName := path.Join(s.Dirs[i], name)
			var bytestring []byte
			bytestring, err = ioutil.ReadFile(fileName)
			if err == nil {
				content = string(bytestring)
				break
			}
		}
		if content == "" && err != nil {
			return
		}
	}

	template, err = s.parse(name, content)
	return
}

func (s *Set) loadTemplate(name, content string) (template *Template, err error) {
	var ok bool
	if template, ok = s.templates[name]; ok {
		return
	}
	template, err = s.load(name, content)
	s.templates[name] = template
	return
}

func (s *Set) GetTemplate(name string) (template *Template, ok bool) {
	s.tmx.RLock()
	template, ok = s.templates[name]
	s.tmx.RUnlock()
	return
}

func (s *Set) LoadTemplate(name, content string) (template *Template, err error) {
	var ok bool

	s.tmx.RLock()
	if template, ok = s.templates[name]; ok {
		s.tmx.RUnlock()
		return
	}

	s.tmx.RUnlock()
	s.tmx.Lock()
	defer s.tmx.Unlock()

	template, ok = s.templates[name]
	if ok && template != nil {
		return
	}

	template, err = s.load(name, content)
	s.templates[name] = template // saves the template
	return
}

func (t *Template) String() (template string) {

	if t.extends != nil {
		template += fmt.Sprintf("{{extends %q}}", t.extends.ParseName)
	}

	for _, _import := range t.imports {
		template += fmt.Sprintf("\n{{import %q}}", _import.ParseName)
	}

	template += t.root.String()
	return
}

func (t *Template) addBlocks(blocks map[string]*BlockNode) {
	if len(blocks) > 0 {
		if t.processedBlocks == nil {
			t.processedBlocks = make(map[string]*BlockNode)
		}
		for key, value := range blocks {
			t.processedBlocks[key] = value
		}
	}
}

type Scope map[string]reflect.Value

func (scope Scope) Set(name string, v interface{}) {
	scope[name] = reflect.ValueOf(v)
}

func (t *Template) Execute(w io.Writer, variables Scope, data interface{}) (err error) {
	st := pool_State.Get().(*State)
	defer st.recover(&err)

	root := t.root
	if t.extends != nil {
		root = t.extends.root
	}
	if data != nil {
		st.context = reflect.ValueOf(data)
	}

	st.blocks = t.processedBlocks

	st.set = t.set
	st.Writer = w
	st.variables = variables

	st.executeList(root)
	return
}
<|endoftext|>


<commit_before>// Copyright 2014 ISRG.  All rights reserved
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package ca

import (
	"crypto"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"time"

	"github.com/letsencrypt/boulder/core"
	blog "github.com/letsencrypt/boulder/log"
	"github.com/letsencrypt/boulder/policy"

	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cloudflare/cfssl/auth"
	cfsslConfig "github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cloudflare/cfssl/config"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cloudflare/cfssl/helpers"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cloudflare/cfssl/ocsp"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cloudflare/cfssl/signer"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cloudflare/cfssl/signer/remote"
)

type Config struct {
	Server       string
	AuthKey      string
	Profile      string
	TestMode     bool
	DBDriver     string
	DBName       string
	SerialPrefix int
	// A PEM-encoded copy of the issuer certificate.
	IssuerCert string
	// This field is only allowed if TestMode is true, indicating that we are
	// signing with a local key. In production we will use an HSM and this
	// IssuerKey must be empty (and TestMode must be false). PEM-encoded private
	// key used for signing certificates and OCSP responses.
	IssuerKey string
}

// CertificateAuthorityImpl represents a CA that signs certificates, CRLs, and
// OCSP responses.
type CertificateAuthorityImpl struct {
	profile    string
	Signer     signer.Signer
	OCSPSigner ocsp.Signer
	SA         core.StorageAuthority
	PA         core.PolicyAuthority
	DB         core.CertificateAuthorityDatabase
	log        *blog.AuditLogger
	Prefix     int // Prepended to the serial number
}

// NewCertificateAuthorityImpl creates a CA that talks to a remote CFSSL
// instance.  (To use a local signer, simply instantiate CertificateAuthorityImpl
// directly.)  Communications with the CA are authenticated with MACs,
// using CFSSL's authenticated signature scheme.  A CA created in this way
// issues for a single profile on the remote signer, which is indicated
// by name in this constructor.
func NewCertificateAuthorityImpl(cadb core.CertificateAuthorityDatabase, config Config) (*CertificateAuthorityImpl, error) {
	var ca *CertificateAuthorityImpl
	var err error
	logger := blog.GetAuditLogger()
	logger.Notice("Certificate Authority Starting")

	if config.SerialPrefix == 0 || config.SerialPrefix >= 256 {
		err = errors.New("Must have non-zero serial prefix less than 256 for CA.")
		return nil, err
	}

	// Create the remote signer
	localProfile := cfsslConfig.SigningProfile{
		Expiry:       time.Hour,     // BOGUS: Required by CFSSL, but not used
		RemoteName:   config.Server, // BOGUS: Only used as a flag by CFSSL
		RemoteServer: config.Server,
		UseSerialSeq: true,
	}

	localProfile.Provider, err = auth.New(config.AuthKey, nil)
	if err != nil {
		return nil, err
	}

	signer, err := remote.NewSigner(&cfsslConfig.Signing{Default: &localProfile})
	if err != nil {
		return nil, err
	}

	issuer, err := loadIssuer(config.IssuerCert)
	if err != nil {
		return nil, err
	}

	// In test mode, load a private key from a file. In production, use an HSM.
	if !config.TestMode {
		err = errors.New("OCSP signing with a PKCS#11 key not yet implemented.")
		return nil, err
	}
	issuerKey, err := loadIssuerKey(config.IssuerKey)
	if err != nil {
		return nil, err
	}

	// Set up our OCSP signer. Note this calls for both the issuer cert and the
	// OCSP signing cert, which are the same in our case.
	ocspSigner, err := ocsp.NewSigner(issuer, issuer, issuerKey,
		time.Hour*24*4)

	pa := policy.NewPolicyAuthorityImpl()

	ca = &CertificateAuthorityImpl{
		Signer:     signer,
		OCSPSigner: ocspSigner,
		profile:    config.Profile,
		PA:         pa,
		DB:         cadb,
		Prefix:     config.SerialPrefix,
		log:        logger,
	}
	return ca, err
}

func loadIssuer(filename string) (issuerCert *x509.Certificate, err error) {
	if filename == "" {
		err = errors.New("Issuer certificate was not provided in config.")
		return
	}
	issuerCertPEM, err := ioutil.ReadFile(filename)
	if err != nil {
		return
	}
	issuerCert, err = helpers.ParseCertificatePEM(issuerCertPEM)
	return
}

func loadIssuerKey(filename string) (issuerKey crypto.Signer, err error) {
	if filename == "" {
		err = errors.New("IssuerKey must be provided in test mode.")
		return
	}

	pem, err := ioutil.ReadFile(filename)
	if err != nil {
		return
	}
	issuerKey, err = helpers.ParsePrivateKeyPEM(pem)
	return
}

func (ca *CertificateAuthorityImpl) RevokeCertificate(serial string) (err error) {
	certDER, err := ca.SA.GetCertificate(serial)
	if err != nil {
		return err
	}
	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return err
	}

	// Per https://tools.ietf.org/html/rfc5280, CRLReason 0 is "unspecified."
	// TODO: Add support for specifying reason.
	reason := 0

	signRequest := ocsp.SignRequest{
		Certificate: cert,
		Status:      string(core.OCSPStatusRevoked),
		Reason:      reason,
		RevokedAt:   time.Now(),
	}
	ocspResponse, err := ca.OCSPSigner.Sign(signRequest)
	if err != nil {
		return err
	}
	err = ca.SA.MarkCertificateRevoked(serial, ocspResponse, reason)
	return err
}

// IssueCertificate attempts to convert a CSR into a signed Certificate, while
// enforcing all policies.
func (ca *CertificateAuthorityImpl) IssueCertificate(csr x509.CertificateRequest) (core.Certificate, error) {
	emptyCert := core.Certificate{}
	var err error
	// XXX Take in authorizations and verify that union covers CSR?
	// Pull hostnames from CSR
	hostNames := csr.DNSNames // DNSNames + CN from CSR
	var commonName string
	if len(csr.Subject.CommonName) > 0 {
		commonName = csr.Subject.CommonName
	} else if len(hostNames) > 0 {
		commonName = hostNames[0]
	} else {
		err = errors.New("Cannot issue a certificate without a hostname.")
		ca.log.WarningErr(err)
		return emptyCert, err
	}

	if len(hostNames) == 0 {
		hostNames = []string{commonName}
	}

	identifier := core.AcmeIdentifier{Type: core.IdentifierDNS, Value: commonName}
	if err = ca.PA.WillingToIssue(identifier); err != nil {
		err = errors.New("Policy forbids issuing for name " + commonName)
		ca.log.AuditErr(err)
		return emptyCert, err
	}
	for _, name := range hostNames {
		identifier = core.AcmeIdentifier{Type: core.IdentifierDNS, Value: name}
		if err = ca.PA.WillingToIssue(identifier); err != nil {
			err = errors.New("Policy forbids issuing for name " + name)
			ca.log.AuditErr(err)
			return emptyCert, err
		}
	}

	// Convert the CSR to PEM
	csrPEM := string(pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE REQUEST",
		Bytes: csr.Raw,
	}))

	// Get the next serial number
	ca.DB.Begin()
	serialDec, err := ca.DB.IncrementAndGetSerial()
	if err != nil {
		return emptyCert, err
	}
	serialHex := fmt.Sprintf("%02X%014X", ca.Prefix, serialDec)

	// Send the cert off for signing
	req := signer.SignRequest{
		Request: csrPEM,
		Profile: ca.profile,
		Hosts:   hostNames,
		Subject: &signer.Subject{
			CN: commonName,
		},
		SerialSeq: serialHex,
	}

	certPEM, err := ca.Signer.Sign(req)
	if err != nil {
		ca.DB.Rollback()
		return emptyCert, err
	}

	if len(certPEM) == 0 {
		err = errors.New("No certificate returned by server")
		ca.log.WarningErr(err)
		ca.DB.Rollback()
		return emptyCert, err
	}

	block, _ := pem.Decode(certPEM)
	if block == nil || block.Type != "CERTIFICATE" {
		err = errors.New("Invalid certificate value returned")
		ca.log.WarningErr(err)
		ca.DB.Rollback()
		return emptyCert, err
	}
	certDER := block.Bytes

	cert := core.Certificate{
		DER:    certDER,
		Status: core.StatusValid,
	}
	if err != nil {
		return emptyCert, err
	}

	// Store the cert with the certificate authority, if provided
	_, err = ca.SA.AddCertificate(certDER)
	if err != nil {
		ca.DB.Rollback()
		return emptyCert, err
	}

	ca.DB.Commit()
	return cert, nil
}
<commit_msg>check serial is > 0<commit_after>// Copyright 2014 ISRG.  All rights reserved
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package ca

import (
	"crypto"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"time"

	"github.com/letsencrypt/boulder/core"
	blog "github.com/letsencrypt/boulder/log"
	"github.com/letsencrypt/boulder/policy"

	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cloudflare/cfssl/auth"
	cfsslConfig "github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cloudflare/cfssl/config"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cloudflare/cfssl/helpers"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cloudflare/cfssl/ocsp"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cloudflare/cfssl/signer"
	"github.com/letsencrypt/boulder/Godeps/_workspace/src/github.com/cloudflare/cfssl/signer/remote"
)

type Config struct {
	Server       string
	AuthKey      string
	Profile      string
	TestMode     bool
	DBDriver     string
	DBName       string
	SerialPrefix int
	// A PEM-encoded copy of the issuer certificate.
	IssuerCert string
	// This field is only allowed if TestMode is true, indicating that we are
	// signing with a local key. In production we will use an HSM and this
	// IssuerKey must be empty (and TestMode must be false). PEM-encoded private
	// key used for signing certificates and OCSP responses.
	IssuerKey string
}

// CertificateAuthorityImpl represents a CA that signs certificates, CRLs, and
// OCSP responses.
type CertificateAuthorityImpl struct {
	profile    string
	Signer     signer.Signer
	OCSPSigner ocsp.Signer
	SA         core.StorageAuthority
	PA         core.PolicyAuthority
	DB         core.CertificateAuthorityDatabase
	log        *blog.AuditLogger
	Prefix     int // Prepended to the serial number
}

// NewCertificateAuthorityImpl creates a CA that talks to a remote CFSSL
// instance.  (To use a local signer, simply instantiate CertificateAuthorityImpl
// directly.)  Communications with the CA are authenticated with MACs,
// using CFSSL's authenticated signature scheme.  A CA created in this way
// issues for a single profile on the remote signer, which is indicated
// by name in this constructor.
func NewCertificateAuthorityImpl(cadb core.CertificateAuthorityDatabase, config Config) (*CertificateAuthorityImpl, error) {
	var ca *CertificateAuthorityImpl
	var err error
	logger := blog.GetAuditLogger()
	logger.Notice("Certificate Authority Starting")

	if config.SerialPrefix <= 0 || config.SerialPrefix >= 256 {
		err = errors.New("Must have a positive non-zero serial prefix less than 256 for CA.")
		return nil, err
	}

	// Create the remote signer
	localProfile := cfsslConfig.SigningProfile{
		Expiry:       time.Hour,     // BOGUS: Required by CFSSL, but not used
		RemoteName:   config.Server, // BOGUS: Only used as a flag by CFSSL
		RemoteServer: config.Server,
		UseSerialSeq: true,
	}

	localProfile.Provider, err = auth.New(config.AuthKey, nil)
	if err != nil {
		return nil, err
	}

	signer, err := remote.NewSigner(&cfsslConfig.Signing{Default: &localProfile})
	if err != nil {
		return nil, err
	}

	issuer, err := loadIssuer(config.IssuerCert)
	if err != nil {
		return nil, err
	}

	// In test mode, load a private key from a file. In production, use an HSM.
	if !config.TestMode {
		err = errors.New("OCSP signing with a PKCS#11 key not yet implemented.")
		return nil, err
	}
	issuerKey, err := loadIssuerKey(config.IssuerKey)
	if err != nil {
		return nil, err
	}

	// Set up our OCSP signer. Note this calls for both the issuer cert and the
	// OCSP signing cert, which are the same in our case.
	ocspSigner, err := ocsp.NewSigner(issuer, issuer, issuerKey,
		time.Hour*24*4)

	pa := policy.NewPolicyAuthorityImpl()

	ca = &CertificateAuthorityImpl{
		Signer:     signer,
		OCSPSigner: ocspSigner,
		profile:    config.Profile,
		PA:         pa,
		DB:         cadb,
		Prefix:     config.SerialPrefix,
		log:        logger,
	}
	return ca, err
}

func loadIssuer(filename string) (issuerCert *x509.Certificate, err error) {
	if filename == "" {
		err = errors.New("Issuer certificate was not provided in config.")
		return
	}
	issuerCertPEM, err := ioutil.ReadFile(filename)
	if err != nil {
		return
	}
	issuerCert, err = helpers.ParseCertificatePEM(issuerCertPEM)
	return
}

func loadIssuerKey(filename string) (issuerKey crypto.Signer, err error) {
	if filename == "" {
		err = errors.New("IssuerKey must be provided in test mode.")
		return
	}

	pem, err := ioutil.ReadFile(filename)
	if err != nil {
		return
	}
	issuerKey, err = helpers.ParsePrivateKeyPEM(pem)
	return
}

func (ca *CertificateAuthorityImpl) RevokeCertificate(serial string) (err error) {
	certDER, err := ca.SA.GetCertificate(serial)
	if err != nil {
		return err
	}
	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return err
	}

	// Per https://tools.ietf.org/html/rfc5280, CRLReason 0 is "unspecified."
	// TODO: Add support for specifying reason.
	reason := 0

	signRequest := ocsp.SignRequest{
		Certificate: cert,
		Status:      string(core.OCSPStatusRevoked),
		Reason:      reason,
		RevokedAt:   time.Now(),
	}
	ocspResponse, err := ca.OCSPSigner.Sign(signRequest)
	if err != nil {
		return err
	}
	err = ca.SA.MarkCertificateRevoked(serial, ocspResponse, reason)
	return err
}

// IssueCertificate attempts to convert a CSR into a signed Certificate, while
// enforcing all policies.
func (ca *CertificateAuthorityImpl) IssueCertificate(csr x509.CertificateRequest) (core.Certificate, error) {
	emptyCert := core.Certificate{}
	var err error
	// XXX Take in authorizations and verify that union covers CSR?
	// Pull hostnames from CSR
	hostNames := csr.DNSNames // DNSNames + CN from CSR
	var commonName string
	if len(csr.Subject.CommonName) > 0 {
		commonName = csr.Subject.CommonName
	} else if len(hostNames) > 0 {
		commonName = hostNames[0]
	} else {
		err = errors.New("Cannot issue a certificate without a hostname.")
		ca.log.WarningErr(err)
		return emptyCert, err
	}

	if len(hostNames) == 0 {
		hostNames = []string{commonName}
	}

	identifier := core.AcmeIdentifier{Type: core.IdentifierDNS, Value: commonName}
	if err = ca.PA.WillingToIssue(identifier); err != nil {
		err = errors.New("Policy forbids issuing for name " + commonName)
		ca.log.AuditErr(err)
		return emptyCert, err
	}
	for _, name := range hostNames {
		identifier = core.AcmeIdentifier{Type: core.IdentifierDNS, Value: name}
		if err = ca.PA.WillingToIssue(identifier); err != nil {
			err = errors.New("Policy forbids issuing for name " + name)
			ca.log.AuditErr(err)
			return emptyCert, err
		}
	}

	// Convert the CSR to PEM
	csrPEM := string(pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE REQUEST",
		Bytes: csr.Raw,
	}))

	// Get the next serial number
	ca.DB.Begin()
	serialDec, err := ca.DB.IncrementAndGetSerial()
	if err != nil {
		return emptyCert, err
	}
	serialHex := fmt.Sprintf("%02X%014X", ca.Prefix, serialDec)

	// Send the cert off for signing
	req := signer.SignRequest{
		Request: csrPEM,
		Profile: ca.profile,
		Hosts:   hostNames,
		Subject: &signer.Subject{
			CN: commonName,
		},
		SerialSeq: serialHex,
	}

	certPEM, err := ca.Signer.Sign(req)
	if err != nil {
		ca.DB.Rollback()
		return emptyCert, err
	}

	if len(certPEM) == 0 {
		err = errors.New("No certificate returned by server")
		ca.log.WarningErr(err)
		ca.DB.Rollback()
		return emptyCert, err
	}

	block, _ := pem.Decode(certPEM)
	if block == nil || block.Type != "CERTIFICATE" {
		err = errors.New("Invalid certificate value returned")
		ca.log.WarningErr(err)
		ca.DB.Rollback()
		return emptyCert, err
	}
	certDER := block.Bytes

	cert := core.Certificate{
		DER:    certDER,
		Status: core.StatusValid,
	}
	if err != nil {
		return emptyCert, err
	}

	// Store the cert with the certificate authority, if provided
	_, err = ca.SA.AddCertificate(certDER)
	if err != nil {
		ca.DB.Rollback()
		return emptyCert, err
	}

	ca.DB.Commit()
	return cert, nil
}
<|endoftext|>


<commit_before>// Package redis implements a Gondola cache driver using redis.
//
// The URL format for this driver is:
//
//  - redis://host[:port][#password={pw}&db={number}]
//
// If no db is provided, it defaults to -1.
package redis

import (
	"fmt"
	"github.com/vmihailenco/redis"
	"gnd.la/cache/driver"
	"gnd.la/config"
	"strconv"
)

type redisDriver struct {
	*redis.Client
}

func (r *redisDriver) Set(key string, b []byte, timeout int) error {
	var req *redis.StatusReq
	if timeout == 0 {
		req = r.Client.Set(key, string(b))
	} else {
		req = r.Client.SetEx(key, int64(timeout), string(b))
	}
	return req.Err()
}

func (r *redisDriver) Get(key string) ([]byte, error) {
	req := r.Client.Get(key)
	if err := req.Err(); err != nil {
		if err == redis.Nil {
			return nil, nil
		}
		return nil, err
	}
	return []byte(req.Val()), nil
}

func (r *redisDriver) GetMulti(keys []string) (map[string][]byte, error) {
	req := r.Client.MGet(keys...)
	if err := req.Err(); err != nil {
		return nil, err
	}
	value := make(map[string][]byte, len(keys))
	for ii, v := range req.Val() {
		if v != nil {
			s := v.(string)
			value[keys[ii]] = []byte(s)
		}
	}
	return value, nil
}

func (r *redisDriver) Delete(key string) error {
	req := r.Client.Del(key)
	return req.Err()
}

func (r *redisDriver) Connection() interface{} {
	return r.Client
}

func redisOpener(url *config.URL) (driver.Driver, error) {
	password := url.Fragment.Get("password")
	db := int64(-1)
	if d := url.Fragment.Get("db"); d != "" {
		val, err := strconv.ParseInt(d, 0, 64)
		if err != nil {
			return nil, fmt.Errorf("invalid db %q, must be an integer", d)
		}
		db = val
	}
	conn := driver.DefaultPort(url.Value, 6379)
	client := redis.NewTCPClient(conn, password, db)
	return &redisDriver{Client: client}, nil
}

func init() {
	driver.Register("redis", redisOpener)
}
<commit_msg>Use redigo to implement the cache Redis driver<commit_after>// Package redis implements a Gondola cache driver using redis.
//
// The URL format for this driver is:
//
//  - redis://host[:port][#password={pw}&db={number}&max_idle={number}&max_active={number}&idle_timeout={seconds}]
//
// If no db is provided, it defaults to -1.
// For the defaults and the explanation for the rest of the parameters,
// see DefaultMaxIdle, DefaultMaxActive and DefaultIdleTimeout.
package redis

import (
	"fmt"
	"time"

	"gnd.la/cache/driver"
	"gnd.la/config"

	"github.com/garyburd/redigo/redis"
)

const (
	// DefaultMaxIdle is the maximum number of idle connections
	// kept in the connection pool.
	DefaultMaxIdle = 8
	// DefaultMaxActive is the maximum number of connections that
	// will be open at any given time. Setting it to zero, the
	// default value, won't limit the number of connections.
	DefaultMaxActive = 0
	// DefaultIdleTimeout is the amount of seconds after an idle
	// connection will be dropped from the pool.
	DefaultIdleTimeout = 300
)

type redisDriver struct {
	pool *redis.Pool
}

func (r *redisDriver) Set(key string, b []byte, timeout int) error {
	conn := r.pool.Get()
	var err error
	if timeout == 0 {
		_, err = conn.Do("SET", key, b)
	} else {
		_, err = conn.Do("SETEX", key, int32(timeout), b)
	}
	conn.Close()
	return err
}

func (r *redisDriver) Get(key string) ([]byte, error) {
	conn := r.pool.Get()
	reply, err := conn.Do("GET", key)
	conn.Close()
	if err != nil {
		return nil, err
	}
	switch reply := reply.(type) {
	case []byte:
		return reply, nil
	case string:
		return []byte(reply), nil
	case redis.Error:
		return nil, reply
	}
	// nil returned, item was not present
	return nil, nil
}

func (r *redisDriver) GetMulti(keys []string) (map[string][]byte, error) {
	args := make([]interface{}, len(keys))
	for ii, v := range keys {
		args[ii] = v
	}
	conn := r.pool.Get()
	reply, err := conn.Do("MGET", args...)
	conn.Close()
	if err != nil {
		return nil, err
	}
	if e, ok := reply.(redis.Error); ok {
		return nil, e
	}
	values := reply.([]interface{})
	ret := make(map[string][]byte, len(keys))
	for ii, v := range values {
		if v != nil {
			if b, ok := v.([]byte); ok {
				ret[keys[ii]] = b
			}
		}
	}
	return ret, nil
}

func (r *redisDriver) Delete(key string) error {
	conn := r.pool.Get()
	_, err := conn.Do("DEL", key)
	conn.Close()
	return err
}

func (r *redisDriver) Connection() interface{} {
	return r.pool
}

func (r *redisDriver) Close() error {
	return r.pool.Close()
}

func redisOpener(url *config.URL) (driver.Driver, error) {
	password := url.Fragment.Get("password")
	db := -1
	maxIdle := DefaultMaxIdle
	maxActive := DefaultMaxActive
	idleTimeout := DefaultIdleTimeout
	if d := url.Fragment.Get("db"); d != "" {
		val, ok := url.Fragment.Int("db")
		if !ok {
			return nil, fmt.Errorf("invalid db %q, must be an integer", d)
		}
		db = val
	}
	if v, ok := url.Fragment.Int("max_idle"); ok {
		maxIdle = v
	}
	if v, ok := url.Fragment.Int("max_active"); ok {
		maxActive = v
	}
	if v, ok := url.Fragment.Int("idle_timeout"); ok {
		idleTimeout = v
	}
	server := driver.DefaultPort(url.Value, 6379)
	pool := &redis.Pool{
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial("tcp", server)
			if err != nil {
				return nil, err
			}
			if password != "" {
				if _, err := c.Do("AUTH", password); err != nil {
					c.Close()
					return nil, err
				}
			}
			if db != -1 {
				if _, err := c.Do("SELECT", db); err != nil {
					c.Close()
					return nil, err
				}
			}
			return c, err
		},
		MaxIdle:     maxIdle,
		MaxActive:   maxActive,
		IdleTimeout: time.Duration(idleTimeout) * time.Second,
	}
	return &redisDriver{pool: pool}, nil
}

func init() {
	driver.Register("redis", redisOpener)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"

	"github.com/gorilla/websocket"

	"github.com/lxc/lxd/shared"
)

func rsyncSendSetup(path string) (*exec.Cmd, net.Conn, io.ReadCloser, error) {
	/*
	 * It's sort of unfortunate, but there's no library call to get a
	 * temporary name, so we get the file and close it and use its name.
	 */
	f, err := ioutil.TempFile("", "lxd_rsync_")
	if err != nil {
		return nil, nil, nil, err
	}
	f.Close()
	os.Remove(f.Name())

	/*
	 * The way rsync works, it invokes a subprocess that does the actual
	 * talking (given to it by a -E argument). Since there isn't an easy
	 * way for us to capture this process' stdin/stdout, we just use netcat
	 * and write to/from a unix socket.
	 *
	 * In principle we don't need this socket. It seems to me that some
	 * clever invocation of rsync --server --sender and usage of that
	 * process' stdin/stdout could work around the need for this socket,
	 * but I couldn't get it to work. Another option would be to look at
	 * the spawned process' first child and read/write from its
	 * stdin/stdout, but that also seemed messy. In any case, this seems to
	 * work just fine.
	 */
	l, err := net.Listen("unix", f.Name())
	if err != nil {
		return nil, nil, nil, err
	}

	/*
	 * Here, the path /tmp/foo is ignored. Since we specify localhost,
	 * rsync thinks we are syncing to a remote host (in this case, the
	 * other end of the lxd websocket), and so the path specified on the
	 * --server instance of rsync takes precedence.
	 *
	 * Additionally, we use sh -c instead of just calling nc directly
	 * because rsync passes a whole bunch of arguments to the wrapper
	 * command (i.e. the command to run on --server). However, we're
	 * hardcoding that at the other end, so we can just ignore it.
	 */
	rsyncCmd := fmt.Sprintf("sh -c \"%s netcat %s\"", execPath, f.Name())
	cmd := exec.Command(
		"rsync",
		"-arvP",
		"--devices",
		"--numeric-ids",
		"--partial",
		path,
		"localhost:/tmp/foo",
		"-e",
		rsyncCmd)

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, nil, nil, err
	}

	if err := cmd.Start(); err != nil {
		return nil, nil, nil, err
	}

	conn, err := l.Accept()
	if err != nil {
		return nil, nil, nil, err
	}
	l.Close()

	return cmd, conn, stderr, nil
}

// RsyncSend sets up the sending half of an rsync, to recursively send the
// directory pointed to by path over the websocket.
func RsyncSend(path string, conn *websocket.Conn, readWrapper func(io.ReadCloser) io.ReadCloser) error {
	cmd, dataSocket, stderr, err := rsyncSendSetup(path)
	if dataSocket != nil {
		defer dataSocket.Close()
	}
	if err != nil {
		return err
	}

	readPipe := io.ReadCloser(dataSocket)
	if readWrapper != nil {
		readPipe = readWrapper(dataSocket)
	}

	readDone, writeDone := shared.WebsocketMirror(conn, dataSocket, readPipe, nil, nil)

	output, err := ioutil.ReadAll(stderr)
	if err != nil {
		shared.LogDebugf("problem reading rsync stderr %s", err)
	}

	err = cmd.Wait()
	if err != nil {
		shared.LogDebugf("problem with rsync send of %s: %s: %s", path, err, string(output))
	}

	<-readDone
	<-writeDone

	return err
}

// RsyncRecv sets up the receiving half of the websocket to rsync (the other
// half set up by RsyncSend), putting the contents in the directory specified
// by path.
func RsyncRecv(path string, conn *websocket.Conn, writeWrapper func(io.WriteCloser) io.WriteCloser) error {
	cmd := exec.Command("rsync",
		"--server",
		"-vlogDtpre.iLsfx",
		"--numeric-ids",
		"--devices",
		"--partial",
		".",
		path)

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return err
	}

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}

	if err := cmd.Start(); err != nil {
		return err
	}

	writePipe := io.WriteCloser(stdin)
	if writeWrapper != nil {
		writePipe = writeWrapper(stdin)
	}

	readDone, writeDone := shared.WebsocketMirror(conn, writePipe, stdout, nil, nil)
	data, err2 := ioutil.ReadAll(stderr)
	if err2 != nil {
		shared.LogDebugf("error reading rsync stderr: %s", err2)
		return err2
	}

	err = cmd.Wait()
	if err != nil {
		shared.LogDebugf("rsync recv error for path %s: %s: %s", path, err, string(data))
	}

	<-readDone
	<-writeDone

	return err
}
<commit_msg>rsync: Tweak logging a bit<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"

	"github.com/gorilla/websocket"

	"github.com/lxc/lxd/shared"
)

func rsyncSendSetup(path string) (*exec.Cmd, net.Conn, io.ReadCloser, error) {
	/*
	 * It's sort of unfortunate, but there's no library call to get a
	 * temporary name, so we get the file and close it and use its name.
	 */
	f, err := ioutil.TempFile("", "lxd_rsync_")
	if err != nil {
		return nil, nil, nil, err
	}
	f.Close()
	os.Remove(f.Name())

	/*
	 * The way rsync works, it invokes a subprocess that does the actual
	 * talking (given to it by a -E argument). Since there isn't an easy
	 * way for us to capture this process' stdin/stdout, we just use netcat
	 * and write to/from a unix socket.
	 *
	 * In principle we don't need this socket. It seems to me that some
	 * clever invocation of rsync --server --sender and usage of that
	 * process' stdin/stdout could work around the need for this socket,
	 * but I couldn't get it to work. Another option would be to look at
	 * the spawned process' first child and read/write from its
	 * stdin/stdout, but that also seemed messy. In any case, this seems to
	 * work just fine.
	 */
	l, err := net.Listen("unix", f.Name())
	if err != nil {
		return nil, nil, nil, err
	}

	/*
	 * Here, the path /tmp/foo is ignored. Since we specify localhost,
	 * rsync thinks we are syncing to a remote host (in this case, the
	 * other end of the lxd websocket), and so the path specified on the
	 * --server instance of rsync takes precedence.
	 *
	 * Additionally, we use sh -c instead of just calling nc directly
	 * because rsync passes a whole bunch of arguments to the wrapper
	 * command (i.e. the command to run on --server). However, we're
	 * hardcoding that at the other end, so we can just ignore it.
	 */
	rsyncCmd := fmt.Sprintf("sh -c \"%s netcat %s\"", execPath, f.Name())
	cmd := exec.Command(
		"rsync",
		"-arvP",
		"--devices",
		"--numeric-ids",
		"--partial",
		path,
		"localhost:/tmp/foo",
		"-e",
		rsyncCmd)

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, nil, nil, err
	}

	if err := cmd.Start(); err != nil {
		return nil, nil, nil, err
	}

	conn, err := l.Accept()
	if err != nil {
		return nil, nil, nil, err
	}
	l.Close()

	return cmd, conn, stderr, nil
}

// RsyncSend sets up the sending half of an rsync, to recursively send the
// directory pointed to by path over the websocket.
func RsyncSend(path string, conn *websocket.Conn, readWrapper func(io.ReadCloser) io.ReadCloser) error {
	cmd, dataSocket, stderr, err := rsyncSendSetup(path)
	if err != nil {
		return err
	}

	if dataSocket != nil {
		defer dataSocket.Close()
	}

	readPipe := io.ReadCloser(dataSocket)
	if readWrapper != nil {
		readPipe = readWrapper(dataSocket)
	}

	readDone, writeDone := shared.WebsocketMirror(conn, dataSocket, readPipe, nil, nil)

	output, err := ioutil.ReadAll(stderr)
	if err != nil {
		return err
	}

	err = cmd.Wait()
	if err != nil {
		shared.LogErrorf("Rsync send failed: %s: %s: %s", path, err, string(output))
	}

	<-readDone
	<-writeDone

	return err
}

// RsyncRecv sets up the receiving half of the websocket to rsync (the other
// half set up by RsyncSend), putting the contents in the directory specified
// by path.
func RsyncRecv(path string, conn *websocket.Conn, writeWrapper func(io.WriteCloser) io.WriteCloser) error {
	cmd := exec.Command("rsync",
		"--server",
		"-vlogDtpre.iLsfx",
		"--numeric-ids",
		"--devices",
		"--partial",
		".",
		path)

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return err
	}

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}

	if err := cmd.Start(); err != nil {
		return err
	}

	writePipe := io.WriteCloser(stdin)
	if writeWrapper != nil {
		writePipe = writeWrapper(stdin)
	}

	readDone, writeDone := shared.WebsocketMirror(conn, writePipe, stdout, nil, nil)
	output, err := ioutil.ReadAll(stderr)
	if err != nil {
		return err
	}

	err = cmd.Wait()
	if err != nil {
		shared.LogErrorf("Rsync receive failed: %s: %s: %s", path, err, string(output))
	}

	<-readDone
	<-writeDone

	return err
}
<|endoftext|>