<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Binary to decimal floating point conversion.
// Algorithm:
//   1) store mantissa in multiprecision decimal
//   2) shift decimal by exponent
//   3) read digits out & format

package strconv

import "math"

// TODO: move elsewhere?
type floatInfo struct {
	mantbits uint
	expbits  uint
	bias     int
}

var float32info = floatInfo{23, 8, -127}
var float64info = floatInfo{52, 11, -1023}

// FormatFloat converts the floating-point number f to a string,
// according to the format fmt and precision prec.  It rounds the
// result assuming that the original was obtained from a floating-point
// value of bitSize bits (32 for float32, 64 for float64).
//
// The format fmt is one of
// 'b' (-ddddp±ddd, a binary exponent),
// 'e' (-d.dddde±dd, a decimal exponent),
// 'E' (-d.ddddE±dd, a decimal exponent),
// 'f' (-ddd.dddd, no exponent),
// 'g' ('e' for large exponents, 'f' otherwise), or
// 'G' ('E' for large exponents, 'f' otherwise).
//
// The precision prec controls the number of digits
// (excluding the exponent) printed by the 'e', 'E', 'f', 'g', and 'G' formats.
// For 'e', 'E', and 'f' it is the number of digits after the decimal point.
// For 'g' and 'G' it is the total number of digits.
// The special precision -1 uses the smallest number of digits
// necessary such that Atof32 will return f exactly.
//
// Ftoa32(f) is not the same as Ftoa64(float32(f)),
// because correct rounding and the number of digits
// needed to identify f depend on the precision of the representation.
func FormatFloat(f float64, fmt byte, prec, bitSize int) string {
	return string(genericFtoa(make([]byte, 0, max(prec+4, 24)), f, fmt, prec, bitSize))
}

// AppendFloat appends the string form of the floating-point number f,
// as generated by FormatFloat, to dst and returns the extended buffer.
func AppendFloat(dst []byte, f float64, fmt byte, prec int, bitSize int) []byte {
	return genericFtoa(dst, f, fmt, prec, bitSize)
}

func genericFtoa(dst []byte, val float64, fmt byte, prec, bitSize int) []byte {
	var bits uint64
	var flt *floatInfo
	switch bitSize {
	case 32:
		bits = uint64(math.Float32bits(float32(val)))
		flt = &float32info
	case 64:
		bits = math.Float64bits(val)
		flt = &float64info
	default:
		panic("strconv: illegal AppendFloat/FormatFloat bitSize")
	}

	neg := bits>>(flt.expbits+flt.mantbits) != 0
	exp := int(bits>>flt.mantbits) & (1<<flt.expbits - 1)
	mant := bits & (uint64(1)<<flt.mantbits - 1)

	switch exp {
	case 1<<flt.expbits - 1:
		// Inf, NaN
		var s string
		switch {
		case mant != 0:
			s = "NaN"
		case neg:
			s = "-Inf"
		default:
			s = "+Inf"
		}
		return append(dst, s...)

	case 0:
		// denormalized
		exp++

	default:
		// add implicit top bit
		mant |= uint64(1) << flt.mantbits
	}
	exp += flt.bias

	// Pick off easy binary format.
	if fmt == 'b' {
		return fmtB(dst, neg, mant, exp, flt)
	}

	// Create exact decimal representation.
	// The shift is exp - flt.mantbits because mant is a 1-bit integer
	// followed by a flt.mantbits fraction, and we are treating it as
	// a 1+flt.mantbits-bit integer.
	d := new(decimal)
	d.Assign(mant)
	d.Shift(exp - int(flt.mantbits))

	// Round appropriately.
	// Negative precision means "only as much as needed to be exact."
	shortest := false
	if prec < 0 {
		shortest = true
		roundShortest(d, mant, exp, flt)
		switch fmt {
		case 'e', 'E':
			prec = d.nd - 1
		case 'f':
			prec = max(d.nd-d.dp, 0)
		case 'g', 'G':
			prec = d.nd
		}
	} else {
		switch fmt {
		case 'e', 'E':
			d.Round(prec + 1)
		case 'f':
			d.Round(d.dp + prec)
		case 'g', 'G':
			if prec == 0 {
				prec = 1
			}
			d.Round(prec)
		}
	}

	switch fmt {
	case 'e', 'E':
		return fmtE(dst, neg, d, prec, fmt)
	case 'f':
		return fmtF(dst, neg, d, prec)
	case 'g', 'G':
		// trailing fractional zeros in 'e' form will be trimmed.
		eprec := prec
		if eprec > d.nd && d.nd >= d.dp {
			eprec = d.nd
		}
		// %e is used if the exponent from the conversion
		// is less than -4 or greater than or equal to the precision.
		// if precision was the shortest possible, use precision 6 for this decision.
		if shortest {
			eprec = 6
		}
		exp := d.dp - 1
		if exp < -4 || exp >= eprec {
			if prec > d.nd {
				prec = d.nd
			}
			return fmtE(dst, neg, d, prec-1, fmt+'e'-'g')
		}
		if prec > d.dp {
			prec = d.nd
		}
		return fmtF(dst, neg, d, max(prec-d.dp, 0))
	}

	// unknown format
	return append(dst, '%', fmt)
}

// Round d (= mant * 2^exp) to the shortest number of digits
// that will let the original floating point value be precisely
// reconstructed.  Size is original floating point size (64 or 32).
func roundShortest(d *decimal, mant uint64, exp int, flt *floatInfo) {
	// If mantissa is zero, the number is zero; stop now.
	if mant == 0 {
		d.nd = 0
		return
	}

	// TODO(rsc): Unless exp == minexp, if the number of digits in d
	// is less than 17, it seems likely that it would be
	// the shortest possible number already.  So maybe we can
	// bail out without doing the extra multiprecision math here.

	// Compute upper and lower such that any decimal number
	// between upper and lower (possibly inclusive)
	// will round to the original floating point number.

	// d = mant << (exp - mantbits)
	// Next highest floating point number is mant+1 << exp-mantbits.
	// Our upper bound is halfway inbetween, mant*2+1 << exp-mantbits-1.
	upper := new(decimal)
	upper.Assign(mant*2 + 1)
	upper.Shift(exp - int(flt.mantbits) - 1)

	// d = mant << (exp - mantbits)
	// Next lowest floating point number is mant-1 << exp-mantbits,
	// unless mant-1 drops the significant bit and exp is not the minimum exp,
	// in which case the next lowest is mant*2-1 << exp-mantbits-1.
	// Either way, call it mantlo << explo-mantbits.
	// Our lower bound is halfway inbetween, mantlo*2+1 << explo-mantbits-1.
	minexp := flt.bias + 1 // minimum possible exponent
	var mantlo uint64
	var explo int
	if mant > 1<<flt.mantbits || exp == minexp {
		mantlo = mant - 1
		explo = exp
	} else {
		mantlo = mant*2 - 1
		explo = exp - 1
	}
	lower := new(decimal)
	lower.Assign(mantlo*2 + 1)
	lower.Shift(explo - int(flt.mantbits) - 1)

	// The upper and lower bounds are possible outputs only if
	// the original mantissa is even, so that IEEE round-to-even
	// would round to the original mantissa and not the neighbors.
	inclusive := mant%2 == 0

	// Now we can figure out the minimum number of digits required.
	// Walk along until d has distinguished itself from upper and lower.
	for i := 0; i < d.nd; i++ {
		var l, m, u byte // lower, middle, upper digits
		if i < lower.nd {
			l = lower.d[i]
		} else {
			l = '0'
		}
		m = d.d[i]
		if i < upper.nd {
			u = upper.d[i]
		} else {
			u = '0'
		}

		// Okay to round down (truncate) if lower has a different digit
		// or if lower is inclusive and is exactly the result of rounding down.
		okdown := l != m || (inclusive && l == m && i+1 == lower.nd)

		// Okay to round up if upper has a different digit and
		// either upper is inclusive or upper is bigger than the result of rounding up.
		okup := m != u && (inclusive || i+1 < upper.nd)

		// If it's okay to do either, then round to the nearest one.
		// If it's okay to do only one, do it.
		switch {
		case okdown && okup:
			d.Round(i + 1)
			return
		case okdown:
			d.RoundDown(i + 1)
			return
		case okup:
			d.RoundUp(i + 1)
			return
		}
	}
}

// %e: -d.ddddde±dd
func fmtE(dst []byte, neg bool, d *decimal, prec int, fmt byte) []byte {
	// sign
	if neg {
		dst = append(dst, '-')
	}

	// first digit
	ch := byte('0')
	if d.nd != 0 {
		ch = d.d[0]
	}
	dst = append(dst, ch)

	// .moredigits
	if prec > 0 {
		dst = append(dst, '.')
		for i := 1; i <= prec; i++ {
			ch = '0'
			if i < d.nd {
				ch = d.d[i]
			}
			dst = append(dst, ch)
		}
	}

	// e±
	dst = append(dst, fmt)
	exp := d.dp - 1
	if d.nd == 0 { // special case: 0 has exponent 0
		exp = 0
	}
	if exp < 0 {
		ch = '-'
		exp = -exp
	} else {
		ch = '+'
	}
	dst = append(dst, ch)

	// dddd
	var buf [3]byte
	i := len(buf)
	for exp >= 10 {
		i--
		buf[i] = byte(exp%10 + '0')
		exp /= 10
	}
	// exp < 10
	i--
	buf[i] = byte(exp + '0')

	// leading zeroes
	if i > len(buf)-2 {
		i--
		buf[i] = '0'
	}

	return append(dst, buf[i:]...)
}

// %f: -ddddddd.ddddd
func fmtF(dst []byte, neg bool, d *decimal, prec int) []byte {
	// sign
	if neg {
		dst = append(dst, '-')
	}

	// integer, padded with zeros as needed.
	if d.dp > 0 {
		var i int
		for i = 0; i < d.dp && i < d.nd; i++ {
			dst = append(dst, d.d[i])
		}
		for ; i < d.dp; i++ {
			dst = append(dst, '0')
		}
	} else {
		dst = append(dst, '0')
	}

	// fraction
	if prec > 0 {
		dst = append(dst, '.')
		for i := 0; i < prec; i++ {
			ch := byte('0')
			if j := d.dp + i; 0 <= j && j < d.nd {
				ch = d.d[j]
			}
			dst = append(dst, ch)
		}
	}

	return dst
}

// %b: -ddddddddp+ddd
func fmtB(dst []byte, neg bool, mant uint64, exp int, flt *floatInfo) []byte {
	var buf [50]byte
	w := len(buf)
	exp -= int(flt.mantbits)
	esign := byte('+')
	if exp < 0 {
		esign = '-'
		exp = -exp
	}
	n := 0
	for exp > 0 || n < 1 {
		n++
		w--
		buf[w] = byte(exp%10 + '0')
		exp /= 10
	}
	w--
	buf[w] = esign
	w--
	buf[w] = 'p'
	n = 0
	for mant > 0 || n < 1 {
		n++
		w--
		buf[w] = byte(mant%10 + '0')
		mant /= 10
	}
	if neg {
		w--
		buf[w] = '-'
	}
	return append(dst, buf[w:]...)
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}
<commit_msg>strconv: remove obsolete comment.<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Binary to decimal floating point conversion.
// Algorithm:
//   1) store mantissa in multiprecision decimal
//   2) shift decimal by exponent
//   3) read digits out & format

package strconv

import "math"

// TODO: move elsewhere?
type floatInfo struct {
	mantbits uint
	expbits  uint
	bias     int
}

var float32info = floatInfo{23, 8, -127}
var float64info = floatInfo{52, 11, -1023}

// FormatFloat converts the floating-point number f to a string,
// according to the format fmt and precision prec.  It rounds the
// result assuming that the original was obtained from a floating-point
// value of bitSize bits (32 for float32, 64 for float64).
//
// The format fmt is one of
// 'b' (-ddddp±ddd, a binary exponent),
// 'e' (-d.dddde±dd, a decimal exponent),
// 'E' (-d.ddddE±dd, a decimal exponent),
// 'f' (-ddd.dddd, no exponent),
// 'g' ('e' for large exponents, 'f' otherwise), or
// 'G' ('E' for large exponents, 'f' otherwise).
//
// The precision prec controls the number of digits
// (excluding the exponent) printed by the 'e', 'E', 'f', 'g', and 'G' formats.
// For 'e', 'E', and 'f' it is the number of digits after the decimal point.
// For 'g' and 'G' it is the total number of digits.
// The special precision -1 uses the smallest number of digits
// necessary such that ParseFloat will return f exactly.
func FormatFloat(f float64, fmt byte, prec, bitSize int) string {
	return string(genericFtoa(make([]byte, 0, max(prec+4, 24)), f, fmt, prec, bitSize))
}

// AppendFloat appends the string form of the floating-point number f,
// as generated by FormatFloat, to dst and returns the extended buffer.
func AppendFloat(dst []byte, f float64, fmt byte, prec int, bitSize int) []byte {
	return genericFtoa(dst, f, fmt, prec, bitSize)
}

func genericFtoa(dst []byte, val float64, fmt byte, prec, bitSize int) []byte {
	var bits uint64
	var flt *floatInfo
	switch bitSize {
	case 32:
		bits = uint64(math.Float32bits(float32(val)))
		flt = &float32info
	case 64:
		bits = math.Float64bits(val)
		flt = &float64info
	default:
		panic("strconv: illegal AppendFloat/FormatFloat bitSize")
	}

	neg := bits>>(flt.expbits+flt.mantbits) != 0
	exp := int(bits>>flt.mantbits) & (1<<flt.expbits - 1)
	mant := bits & (uint64(1)<<flt.mantbits - 1)

	switch exp {
	case 1<<flt.expbits - 1:
		// Inf, NaN
		var s string
		switch {
		case mant != 0:
			s = "NaN"
		case neg:
			s = "-Inf"
		default:
			s = "+Inf"
		}
		return append(dst, s...)

	case 0:
		// denormalized
		exp++

	default:
		// add implicit top bit
		mant |= uint64(1) << flt.mantbits
	}
	exp += flt.bias

	// Pick off easy binary format.
	if fmt == 'b' {
		return fmtB(dst, neg, mant, exp, flt)
	}

	// Create exact decimal representation.
	// The shift is exp - flt.mantbits because mant is a 1-bit integer
	// followed by a flt.mantbits fraction, and we are treating it as
	// a 1+flt.mantbits-bit integer.
	d := new(decimal)
	d.Assign(mant)
	d.Shift(exp - int(flt.mantbits))

	// Round appropriately.
	// Negative precision means "only as much as needed to be exact."
	shortest := false
	if prec < 0 {
		shortest = true
		roundShortest(d, mant, exp, flt)
		switch fmt {
		case 'e', 'E':
			prec = d.nd - 1
		case 'f':
			prec = max(d.nd-d.dp, 0)
		case 'g', 'G':
			prec = d.nd
		}
	} else {
		switch fmt {
		case 'e', 'E':
			d.Round(prec + 1)
		case 'f':
			d.Round(d.dp + prec)
		case 'g', 'G':
			if prec == 0 {
				prec = 1
			}
			d.Round(prec)
		}
	}

	switch fmt {
	case 'e', 'E':
		return fmtE(dst, neg, d, prec, fmt)
	case 'f':
		return fmtF(dst, neg, d, prec)
	case 'g', 'G':
		// trailing fractional zeros in 'e' form will be trimmed.
		eprec := prec
		if eprec > d.nd && d.nd >= d.dp {
			eprec = d.nd
		}
		// %e is used if the exponent from the conversion
		// is less than -4 or greater than or equal to the precision.
		// if precision was the shortest possible, use precision 6 for this decision.
		if shortest {
			eprec = 6
		}
		exp := d.dp - 1
		if exp < -4 || exp >= eprec {
			if prec > d.nd {
				prec = d.nd
			}
			return fmtE(dst, neg, d, prec-1, fmt+'e'-'g')
		}
		if prec > d.dp {
			prec = d.nd
		}
		return fmtF(dst, neg, d, max(prec-d.dp, 0))
	}

	// unknown format
	return append(dst, '%', fmt)
}

// Round d (= mant * 2^exp) to the shortest number of digits
// that will let the original floating point value be precisely
// reconstructed.  Size is original floating point size (64 or 32).
func roundShortest(d *decimal, mant uint64, exp int, flt *floatInfo) {
	// If mantissa is zero, the number is zero; stop now.
	if mant == 0 {
		d.nd = 0
		return
	}

	// TODO(rsc): Unless exp == minexp, if the number of digits in d
	// is less than 17, it seems likely that it would be
	// the shortest possible number already.  So maybe we can
	// bail out without doing the extra multiprecision math here.

	// Compute upper and lower such that any decimal number
	// between upper and lower (possibly inclusive)
	// will round to the original floating point number.

	// d = mant << (exp - mantbits)
	// Next highest floating point number is mant+1 << exp-mantbits.
	// Our upper bound is halfway inbetween, mant*2+1 << exp-mantbits-1.
	upper := new(decimal)
	upper.Assign(mant*2 + 1)
	upper.Shift(exp - int(flt.mantbits) - 1)

	// d = mant << (exp - mantbits)
	// Next lowest floating point number is mant-1 << exp-mantbits,
	// unless mant-1 drops the significant bit and exp is not the minimum exp,
	// in which case the next lowest is mant*2-1 << exp-mantbits-1.
	// Either way, call it mantlo << explo-mantbits.
	// Our lower bound is halfway inbetween, mantlo*2+1 << explo-mantbits-1.
	minexp := flt.bias + 1 // minimum possible exponent
	var mantlo uint64
	var explo int
	if mant > 1<<flt.mantbits || exp == minexp {
		mantlo = mant - 1
		explo = exp
	} else {
		mantlo = mant*2 - 1
		explo = exp - 1
	}
	lower := new(decimal)
	lower.Assign(mantlo*2 + 1)
	lower.Shift(explo - int(flt.mantbits) - 1)

	// The upper and lower bounds are possible outputs only if
	// the original mantissa is even, so that IEEE round-to-even
	// would round to the original mantissa and not the neighbors.
	inclusive := mant%2 == 0

	// Now we can figure out the minimum number of digits required.
	// Walk along until d has distinguished itself from upper and lower.
	for i := 0; i < d.nd; i++ {
		var l, m, u byte // lower, middle, upper digits
		if i < lower.nd {
			l = lower.d[i]
		} else {
			l = '0'
		}
		m = d.d[i]
		if i < upper.nd {
			u = upper.d[i]
		} else {
			u = '0'
		}

		// Okay to round down (truncate) if lower has a different digit
		// or if lower is inclusive and is exactly the result of rounding down.
		okdown := l != m || (inclusive && l == m && i+1 == lower.nd)

		// Okay to round up if upper has a different digit and
		// either upper is inclusive or upper is bigger than the result of rounding up.
		okup := m != u && (inclusive || i+1 < upper.nd)

		// If it's okay to do either, then round to the nearest one.
		// If it's okay to do only one, do it.
		switch {
		case okdown && okup:
			d.Round(i + 1)
			return
		case okdown:
			d.RoundDown(i + 1)
			return
		case okup:
			d.RoundUp(i + 1)
			return
		}
	}
}

// %e: -d.ddddde±dd
func fmtE(dst []byte, neg bool, d *decimal, prec int, fmt byte) []byte {
	// sign
	if neg {
		dst = append(dst, '-')
	}

	// first digit
	ch := byte('0')
	if d.nd != 0 {
		ch = d.d[0]
	}
	dst = append(dst, ch)

	// .moredigits
	if prec > 0 {
		dst = append(dst, '.')
		for i := 1; i <= prec; i++ {
			ch = '0'
			if i < d.nd {
				ch = d.d[i]
			}
			dst = append(dst, ch)
		}
	}

	// e±
	dst = append(dst, fmt)
	exp := d.dp - 1
	if d.nd == 0 { // special case: 0 has exponent 0
		exp = 0
	}
	if exp < 0 {
		ch = '-'
		exp = -exp
	} else {
		ch = '+'
	}
	dst = append(dst, ch)

	// dddd
	var buf [3]byte
	i := len(buf)
	for exp >= 10 {
		i--
		buf[i] = byte(exp%10 + '0')
		exp /= 10
	}
	// exp < 10
	i--
	buf[i] = byte(exp + '0')

	// leading zeroes
	if i > len(buf)-2 {
		i--
		buf[i] = '0'
	}

	return append(dst, buf[i:]...)
}

// %f: -ddddddd.ddddd
func fmtF(dst []byte, neg bool, d *decimal, prec int) []byte {
	// sign
	if neg {
		dst = append(dst, '-')
	}

	// integer, padded with zeros as needed.
	if d.dp > 0 {
		var i int
		for i = 0; i < d.dp && i < d.nd; i++ {
			dst = append(dst, d.d[i])
		}
		for ; i < d.dp; i++ {
			dst = append(dst, '0')
		}
	} else {
		dst = append(dst, '0')
	}

	// fraction
	if prec > 0 {
		dst = append(dst, '.')
		for i := 0; i < prec; i++ {
			ch := byte('0')
			if j := d.dp + i; 0 <= j && j < d.nd {
				ch = d.d[j]
			}
			dst = append(dst, ch)
		}
	}

	return dst
}

// %b: -ddddddddp+ddd
func fmtB(dst []byte, neg bool, mant uint64, exp int, flt *floatInfo) []byte {
	var buf [50]byte
	w := len(buf)
	exp -= int(flt.mantbits)
	esign := byte('+')
	if exp < 0 {
		esign = '-'
		exp = -exp
	}
	n := 0
	for exp > 0 || n < 1 {
		n++
		w--
		buf[w] = byte(exp%10 + '0')
		exp /= 10
	}
	w--
	buf[w] = esign
	w--
	buf[w] = 'p'
	n = 0
	for mant > 0 || n < 1 {
		n++
		w--
		buf[w] = byte(mant%10 + '0')
		mant /= 10
	}
	if neg {
		w--
		buf[w] = '-'
	}
	return append(dst, buf[w:]...)
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"strings"

	"github.com/jawher/mow.cli"

	docker "github.com/bywan/go-dockercommand"

	dockerclient "github.com/fsouza/go-dockerclient"
)

var allContainers []dockerclient.APIContainers

func run(cmd *cli.Cmd) {
	cmd.Spec = "[--home|--scm-key|--registry]... [--restart|--update]"

	bzkHome := cmd.String(cli.StringOpt{
		Name:   "home",
		Desc:   "Bazooka's work directory",
		EnvVar: "BZK_HOME",
	})
	scmKey := cmd.String(cli.StringOpt{
		Name:   "scm-key",
		Desc:   "Location of the private SSH Key Bazooka will use for SCM Fetch",
		EnvVar: "BZK_SCM_KEYFILE",
	})
	registry := cmd.String(cli.StringOpt{
		Name:   "registry",
		EnvVar: "BZK_REGISTRY",
	})
	dockerSock := cmd.String(cli.StringOpt{
		Name:   "docker-sock",
		Desc:   "Location of the Docker unix socket, usually /var/run/docker.sock",
		EnvVar: "BZK_DOCKERSOCK",
	})

	forceRestart := cmd.Bool(cli.BoolOpt{
		Name: "r restart",
		Desc: "Restart Bazooka if already running",
	})
	forceUpdate := cmd.Bool(cli.BoolOpt{
		Name: "u update",
		Desc: "Update Bazooka to the latest version by pulling new images from the registry",
	})

	cmd.Action = func() {
		config, err := loadConfig()
		if err != nil {
			log.Fatal(fmt.Errorf("Unable to load Bazooka config, reason is: %v\n", err))
		}
		if len(*bzkHome) == 0 {
			if len(config.Home) == 0 {
				*bzkHome = interactiveInput("Bazooka Home Folder")
				config.Home = *bzkHome
			} else {
				*bzkHome = config.Home
			}
		}

		if len(*dockerSock) == 0 {
			if len(config.DockerSock) == 0 {
				*dockerSock = interactiveInput("Docker Socket path")
				config.DockerSock = *dockerSock
			} else {
				*dockerSock = config.DockerSock
			}
		}

		if len(*scmKey) == 0 {
			if len(config.SCMKey) == 0 {
				*scmKey = interactiveInput("Bazooka Default SCM private key")
				config.SCMKey = *scmKey
			} else {
				*scmKey = config.SCMKey
			}
		}

		err = saveConfig(config)
		if err != nil {
			log.Fatal(fmt.Errorf("Unable to save Bazooka config, reason is: %v\n", err))
		}

		client, err := docker.NewDocker("")
		if err != nil {
			log.Fatal(err)
		}

		allContainers, err = client.Ps(&docker.PsOptions{
			All: true,
		})
		if err != nil {
			log.Fatal(err)
		}

		if *forceUpdate {
			log.Printf("Pulling Bazooka images to check for new versions\n")
			mandatoryImages := []string{"server", "web", "orchestration", "parser"}
			optionalImages := []string{"parser-java", "parser-golang", "scm-git",
				"runner-java", "runner-java:oraclejdk8", "runner-java:oraclejdk7", "runner-java:oraclejdk6", "runner-java:openjdk8", "runner-java:openjdk7", "runner-java:openjdk6",
				"runner-golang", "runner-golang:1.2.2", "runner-golang:1.3", "runner-golang:1.3.1", "runner-golang:1.3.2", "runner-golang:1.3.3", "runner-golang:1.4"}
			for _, image := range mandatoryImages {
				err = client.Pull(&docker.PullOptions{Image: getImageLocation(*registry, fmt.Sprintf("bazooka/%s", image))})
				if err != nil {
					log.Fatal(fmt.Errorf("Unable to pull required image for Bazooka, reason is: %v\n", err))
				}
			}
			for _, image := range optionalImages {
				err = client.Pull(&docker.PullOptions{Image: getImageLocation(*registry, fmt.Sprintf("bazooka/%s", image))})
				if err != nil {
					log.Printf("Unable to pull image for Bazooka, as it is an optional one, let's move on. Reason is: %v\n", err)
				}
			}
		}

		mongoRestarted, err := ensureContainerIsRestarted(client, &docker.RunOptions{
			Name: "bzk_mongodb",
			// Using the official mongo image from dockerhub, this may need a change later
			Image:  "mongo",
			Detach: true,
		}, false)

		serverRestarted, err := ensureContainerIsRestarted(client, &docker.RunOptions{
			Name:   "bzk_server",
			Image:  getImageLocation(*registry, "bazooka/server"),
			Detach: true,
			VolumeBinds: []string{
				fmt.Sprintf("%s:/bazooka", *bzkHome),
				fmt.Sprintf("%s:/var/run/docker.sock", *dockerSock),
			},
			Links: []string{"bzk_mongodb:mongo"},
			Env:   getServerEnv(*bzkHome, *dockerSock, *scmKey),
			PortBindings: map[dockerclient.Port][]dockerclient.PortBinding{
				"3000/tcp": []dockerclient.PortBinding{
					dockerclient.PortBinding{HostPort: "3000"},
				},
			},
		}, mongoRestarted || *forceRestart || *forceUpdate)

		_, err = ensureContainerIsRestarted(client, &docker.RunOptions{
			Name:   "bzk_web",
			Image:  getImageLocation(*registry, "bazooka/web"),
			Detach: true,
			Links:  []string{"bzk_server:server"},
			PortBindings: map[dockerclient.Port][]dockerclient.PortBinding{
				"80/tcp": []dockerclient.PortBinding{
					dockerclient.PortBinding{HostPort: "8000"},
				},
			},
		}, serverRestarted || *forceRestart || *forceUpdate)

		if err != nil {
			log.Fatal(err)
		}
	}

}

func ensureContainerIsRestarted(client *docker.Docker, options *docker.RunOptions, needRestart bool) (bool, error) {
	container, err := getContainer(allContainers, options.Name)
	if err != nil {
		log.Printf("Container %s not found, Starting it\n", options.Name)
		_, err := client.Run(options)
		return true, err
	}
	if needRestart {
		log.Printf("Restarting Container %s\n", options.Name)
		err = client.Rm(&docker.RmOptions{
			Container: []string{container.ID},
			Force:     true,
		})
		if err != nil {
			return false, err
		}
		_, err := client.Run(options)
		return true, err
	}
	if strings.HasPrefix(container.Status, "Up") {
		log.Printf("Container %s already Up & Running, keeping on\n", options.Name)
		return false, nil
	}
	log.Printf("Container %s is not `Up`, starting it\n", options.Name)
	return true, client.Start(&docker.StartOptions{
		ID: container.ID,
	})

}

func getServerEnv(home, dockerSock, scmKey string) map[string]string {
	envMap := map[string]string{
		"BZK_HOME":       home,
		"BZK_DOCKERSOCK": dockerSock,
	}
	if len(scmKey) > 0 {
		envMap["BZK_SCM_KEYFILE"] = scmKey
	}
	return envMap
}

func getContainer(containers []dockerclient.APIContainers, name string) (dockerclient.APIContainers, error) {
	for _, container := range containers {
		if contains(container.Names, name) || contains(container.Names, "/"+name) {
			return container, nil
		}
	}
	return dockerclient.APIContainers{}, fmt.Errorf("Container not found")
}

func getImageLocation(registry, image string) string {
	if len(registry) > 0 {
		return fmt.Sprintf("%s/%s", registry, image)
	}
	return image
}

func contains(slice []string, item string) bool {
	set := make(map[string]struct{}, len(slice))
	for _, s := range slice {
		set[s] = struct{}{}
	}
	_, ok := set[item]
	return ok
}
<commit_msg>Remove redundant type<commit_after>package main

import (
	"fmt"
	"log"
	"strings"

	"github.com/jawher/mow.cli"

	docker "github.com/bywan/go-dockercommand"

	dockerclient "github.com/fsouza/go-dockerclient"
)

var allContainers []dockerclient.APIContainers

func run(cmd *cli.Cmd) {
	cmd.Spec = "[--home|--scm-key|--registry]... [--restart|--update]"

	bzkHome := cmd.String(cli.StringOpt{
		Name:   "home",
		Desc:   "Bazooka's work directory",
		EnvVar: "BZK_HOME",
	})
	scmKey := cmd.String(cli.StringOpt{
		Name:   "scm-key",
		Desc:   "Location of the private SSH Key Bazooka will use for SCM Fetch",
		EnvVar: "BZK_SCM_KEYFILE",
	})
	registry := cmd.String(cli.StringOpt{
		Name:   "registry",
		EnvVar: "BZK_REGISTRY",
	})
	dockerSock := cmd.String(cli.StringOpt{
		Name:   "docker-sock",
		Desc:   "Location of the Docker unix socket, usually /var/run/docker.sock",
		EnvVar: "BZK_DOCKERSOCK",
	})

	forceRestart := cmd.Bool(cli.BoolOpt{
		Name: "r restart",
		Desc: "Restart Bazooka if already running",
	})
	forceUpdate := cmd.Bool(cli.BoolOpt{
		Name: "u update",
		Desc: "Update Bazooka to the latest version by pulling new images from the registry",
	})

	cmd.Action = func() {
		config, err := loadConfig()
		if err != nil {
			log.Fatal(fmt.Errorf("Unable to load Bazooka config, reason is: %v\n", err))
		}
		if len(*bzkHome) == 0 {
			if len(config.Home) == 0 {
				*bzkHome = interactiveInput("Bazooka Home Folder")
				config.Home = *bzkHome
			} else {
				*bzkHome = config.Home
			}
		}

		if len(*dockerSock) == 0 {
			if len(config.DockerSock) == 0 {
				*dockerSock = interactiveInput("Docker Socket path")
				config.DockerSock = *dockerSock
			} else {
				*dockerSock = config.DockerSock
			}
		}

		if len(*scmKey) == 0 {
			if len(config.SCMKey) == 0 {
				*scmKey = interactiveInput("Bazooka Default SCM private key")
				config.SCMKey = *scmKey
			} else {
				*scmKey = config.SCMKey
			}
		}

		err = saveConfig(config)
		if err != nil {
			log.Fatal(fmt.Errorf("Unable to save Bazooka config, reason is: %v\n", err))
		}

		client, err := docker.NewDocker("")
		if err != nil {
			log.Fatal(err)
		}

		allContainers, err = client.Ps(&docker.PsOptions{
			All: true,
		})
		if err != nil {
			log.Fatal(err)
		}

		if *forceUpdate {
			log.Printf("Pulling Bazooka images to check for new versions\n")
			mandatoryImages := []string{"server", "web", "orchestration", "parser"}
			optionalImages := []string{"parser-java", "parser-golang", "scm-git",
				"runner-java", "runner-java:oraclejdk8", "runner-java:oraclejdk7", "runner-java:oraclejdk6", "runner-java:openjdk8", "runner-java:openjdk7", "runner-java:openjdk6",
				"runner-golang", "runner-golang:1.2.2", "runner-golang:1.3", "runner-golang:1.3.1", "runner-golang:1.3.2", "runner-golang:1.3.3", "runner-golang:1.4"}
			for _, image := range mandatoryImages {
				err = client.Pull(&docker.PullOptions{Image: getImageLocation(*registry, fmt.Sprintf("bazooka/%s", image))})
				if err != nil {
					log.Fatal(fmt.Errorf("Unable to pull required image for Bazooka, reason is: %v\n", err))
				}
			}
			for _, image := range optionalImages {
				err = client.Pull(&docker.PullOptions{Image: getImageLocation(*registry, fmt.Sprintf("bazooka/%s", image))})
				if err != nil {
					log.Printf("Unable to pull image for Bazooka, as it is an optional one, let's move on. Reason is: %v\n", err)
				}
			}
		}

		mongoRestarted, err := ensureContainerIsRestarted(client, &docker.RunOptions{
			Name: "bzk_mongodb",
			// Using the official mongo image from dockerhub, this may need a change later
			Image:  "mongo",
			Detach: true,
		}, false)

		serverRestarted, err := ensureContainerIsRestarted(client, &docker.RunOptions{
			Name:   "bzk_server",
			Image:  getImageLocation(*registry, "bazooka/server"),
			Detach: true,
			VolumeBinds: []string{
				fmt.Sprintf("%s:/bazooka", *bzkHome),
				fmt.Sprintf("%s:/var/run/docker.sock", *dockerSock),
			},
			Links: []string{"bzk_mongodb:mongo"},
			Env:   getServerEnv(*bzkHome, *dockerSock, *scmKey),
			PortBindings: map[dockerclient.Port][]dockerclient.PortBinding{
				"3000/tcp": {{HostPort: "3000"}},
			},
		}, mongoRestarted || *forceRestart || *forceUpdate)

		_, err = ensureContainerIsRestarted(client, &docker.RunOptions{
			Name:   "bzk_web",
			Image:  getImageLocation(*registry, "bazooka/web"),
			Detach: true,
			Links:  []string{"bzk_server:server"},
			PortBindings: map[dockerclient.Port][]dockerclient.PortBinding{
				"80/tcp": {{HostPort: "8000"}},
			},
		}, serverRestarted || *forceRestart || *forceUpdate)

		if err != nil {
			log.Fatal(err)
		}
	}

}

func ensureContainerIsRestarted(client *docker.Docker, options *docker.RunOptions, needRestart bool) (bool, error) {
	container, err := getContainer(allContainers, options.Name)
	if err != nil {
		log.Printf("Container %s not found, Starting it\n", options.Name)
		_, err := client.Run(options)
		return true, err
	}
	if needRestart {
		log.Printf("Restarting Container %s\n", options.Name)
		err = client.Rm(&docker.RmOptions{
			Container: []string{container.ID},
			Force:     true,
		})
		if err != nil {
			return false, err
		}
		_, err := client.Run(options)
		return true, err
	}
	if strings.HasPrefix(container.Status, "Up") {
		log.Printf("Container %s already Up & Running, keeping on\n", options.Name)
		return false, nil
	}
	log.Printf("Container %s is not `Up`, starting it\n", options.Name)
	return true, client.Start(&docker.StartOptions{
		ID: container.ID,
	})

}

func getServerEnv(home, dockerSock, scmKey string) map[string]string {
	envMap := map[string]string{
		"BZK_HOME":       home,
		"BZK_DOCKERSOCK": dockerSock,
	}
	if len(scmKey) > 0 {
		envMap["BZK_SCM_KEYFILE"] = scmKey
	}
	return envMap
}

func getContainer(containers []dockerclient.APIContainers, name string) (dockerclient.APIContainers, error) {
	for _, container := range containers {
		if contains(container.Names, name) || contains(container.Names, "/"+name) {
			return container, nil
		}
	}
	return dockerclient.APIContainers{}, fmt.Errorf("Container not found")
}

func getImageLocation(registry, image string) string {
	if len(registry) > 0 {
		return fmt.Sprintf("%s/%s", registry, image)
	}
	return image
}

func contains(slice []string, item string) bool {
	set := make(map[string]struct{}, len(slice))
	for _, s := range slice {
		set[s] = struct{}{}
	}
	_, ok := set[item]
	return ok
}
<|endoftext|>


<commit_before>package upcloud

import (
	"fmt"
	"github.com/UpCloudLtd/upcloud-go-sdk/upcloud"
	"github.com/UpCloudLtd/upcloud-go-sdk/upcloud/request"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/helper/communicator"
	"github.com/mitchellh/packer/packer"
	"log"
)

// The unique ID for this builder.
const BuilderId = "upcloudltd.upcloud"

// Builder represents a Packer Builder.
type Builder struct {
	config *Config
	runner multistep.Runner
}

// Prepare processes the build configuration parameters and validates the configuration
func (self *Builder) Prepare(raws ...interface{}) ([]string, []string, err error) {
	// Parse and create the configuration
	self.config, err = NewConfig(raws...)

	if err != nil {
		return nil, nil, err
	}

	// Check that the client/service is usable
	service := self.config.GetService()

	if _, err := service.GetAccount(); err != nil {
		return nil, nil, err
	}

	// Check that the specified storage device is a template
	storageDetails, err := service.GetStorageDetails(&request.GetStorageDetailsRequest{
		UUID: self.config.StorageUUID,
	})

	if err == nil && storageDetails.Type != upcloud.StorageTypeTemplate {
		err = fmt.Errorf("The specified storage UUID is of invalid type \"%s\"", storageDetails.Type)
	}

	if err != nil {
		return nil, nil, err
	}

	return nil, nil, nil
}

// Run executes the actual build steps
func (self *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packer.Artifact, error) {
	// Create the service
	service := self.config.GetService()

	// Set up the state which is used to share state between the steps
	state := new(multistep.BasicStateBag)
	state.Put("config", *self.config)
	state.Put("service", *service)
	state.Put("hook", hook)
	state.Put("ui", ui)

	// Build the steps
	steps := []multistep.Step{
		&StepCreateSSHKey{
			Debug:        self.config.PackerDebug,
			DebugKeyPath: fmt.Sprintf("packer-builder-upcloud-%s.pem", self.config.PackerBuildName),
		},
		new(StepCreateServer),
		&communicator.StepConnect{
			Config:    &self.config.Comm,
			Host:      sshHostCallback,
			SSHConfig: sshConfigCallback,
		},
		new(common.StepProvision),
		new(StepTemplatizeStorage),
	}

	// Create the runner which will run the steps we just build
	self.runner = &multistep.BasicRunner{Steps: steps}
	self.runner.Run(state)

	if rawErr, ok := state.GetOk("error"); ok {
		return nil, rawErr.(error)
	}

	// Extract the final storage details from the state
	rawDetails, ok := state.GetOk("storage_details")

	if !ok {
		log.Println("No storage details found in state, the build was probably cancelled")
		return nil, nil
	}

	storageDetails := rawDetails.(*upcloud.StorageDetails)

	// Create an artifact and return it
	artifact := &Artifact{
		UUID:    storageDetails.UUID,
		Zone:    storageDetails.Zone,
		Title:   storageDetails.Title,
		service: service,
	}

	return artifact, nil
}

// Cancel is called when the build is cancelled
func (self *Builder) Cancel() {
	if self.runner != nil {
		log.Println("Cancelling the step runner ...")
		self.runner.Cancel()
	}

	fmt.Println("Cancelling the builder ...")
}
<commit_msg>Update builder.go<commit_after>package upcloud

import (
	"fmt"
	"github.com/UpCloudLtd/upcloud-go-sdk/upcloud"
	"github.com/UpCloudLtd/upcloud-go-sdk/upcloud/request"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/common"
	"github.com/mitchellh/packer/helper/communicator"
	"github.com/mitchellh/packer/packer"
	"log"
)

// The unique ID for this builder.
const BuilderId = "upcloudltd.upcloud"

// Builder represents a Packer Builder.
type Builder struct {
	config *Config
	runner multistep.Runner
}

// Prepare processes the build configuration parameters and validates the configuration
func (self *Builder) Prepare(raws ...interface{}) ([]string, []string, error) {
	// Parse and create the configuration
	self.config, err = NewConfig(raws...)

	if err != nil {
		return nil, nil, err
	}

	// Check that the client/service is usable
	service := self.config.GetService()

	if _, err := service.GetAccount(); err != nil {
		return nil, nil, err
	}

	// Check that the specified storage device is a template
	storageDetails, err := service.GetStorageDetails(&request.GetStorageDetailsRequest{
		UUID: self.config.StorageUUID,
	})

	if err == nil && storageDetails.Type != upcloud.StorageTypeTemplate {
		err = fmt.Errorf("The specified storage UUID is of invalid type \"%s\"", storageDetails.Type)
	}

	if err != nil {
		return nil, nil, err
	}

	return nil, nil, nil
}

// Run executes the actual build steps
func (self *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packer.Artifact, error) {
	// Create the service
	service := self.config.GetService()

	// Set up the state which is used to share state between the steps
	state := new(multistep.BasicStateBag)
	state.Put("config", *self.config)
	state.Put("service", *service)
	state.Put("hook", hook)
	state.Put("ui", ui)

	// Build the steps
	steps := []multistep.Step{
		&StepCreateSSHKey{
			Debug:        self.config.PackerDebug,
			DebugKeyPath: fmt.Sprintf("packer-builder-upcloud-%s.pem", self.config.PackerBuildName),
		},
		new(StepCreateServer),
		&communicator.StepConnect{
			Config:    &self.config.Comm,
			Host:      sshHostCallback,
			SSHConfig: sshConfigCallback,
		},
		new(common.StepProvision),
		new(StepTemplatizeStorage),
	}

	// Create the runner which will run the steps we just build
	self.runner = &multistep.BasicRunner{Steps: steps}
	self.runner.Run(state)

	if rawErr, ok := state.GetOk("error"); ok {
		return nil, rawErr.(error)
	}

	// Extract the final storage details from the state
	rawDetails, ok := state.GetOk("storage_details")

	if !ok {
		log.Println("No storage details found in state, the build was probably cancelled")
		return nil, nil
	}

	storageDetails := rawDetails.(*upcloud.StorageDetails)

	// Create an artifact and return it
	artifact := &Artifact{
		UUID:    storageDetails.UUID,
		Zone:    storageDetails.Zone,
		Title:   storageDetails.Title,
		service: service,
	}

	return artifact, nil
}

// Cancel is called when the build is cancelled
func (self *Builder) Cancel() {
	if self.runner != nil {
		log.Println("Cancelling the step runner ...")
		self.runner.Cancel()
	}

	fmt.Println("Cancelling the builder ...")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net/rpc"

	"pault.ag/go/minion/minion"
	"pault.ag/go/service"
)

var minionCommand = Command{
	Name:  "minion",
	Run:   minionRun,
	Usage: ``,
}

type minionService struct {
	service.Node
}

func (m *minionService) Register() {
	minion := minion.MinionRemote{Arches: []string{"amd64", "all"}}
	rpc.Register(&minion)
}

func minionRun(config MinionConfig, cmd *Command, args []string) {
	log.Printf("Bringing Minion online\n")
	node := minionService{}
	node.Register()
	log.Printf("Diling coordinator\n")
	conn, err := service.DialFromKeys(
		fmt.Sprintf("%s:%d", config.Host, config.Port),
		config.Cert, config.Key, config.CaCert,
	)
	if err != nil {
		log.Fatalf("Error! %s\n", err)
	}
	log.Printf("Doing what they say!\n")
	service.ServeConn(conn)
}
<commit_msg>add things.<commit_after>package main

import (
	"fmt"
	"log"
	"net/rpc"
	"strings"

	"pault.ag/go/minion/minion"
	"pault.ag/go/service"
)

var minionCommand = Command{
	Name:  "minion",
	Run:   minionRun,
	Usage: ``,
}

var archs *string

func init() {
	archs = minionCommand.Flag.String("arch", "", "comma seperated arches")
}

type minionService struct {
	service.Node
}

func (m *minionService) Register() {
	if *archs == "" {
		log.Fatalf("No archs given\n")
	}
	minion := minion.MinionRemote{Arches: strings.Split(*archs, ",")}
	rpc.Register(&minion)
}

func minionRun(config MinionConfig, cmd *Command, args []string) {
	log.Printf("Bringing Minion online\n")
	node := minionService{}
	node.Register()
	log.Printf("Diling coordinator\n")
	conn, err := service.DialFromKeys(
		fmt.Sprintf("%s:%d", config.Host, config.Port),
		config.Cert, config.Key, config.CaCert,
	)
	if err != nil {
		log.Fatalf("Error! %s\n", err)
	}
	log.Printf("Doing what they say!\n")
	service.ServeConn(conn)
}
<|endoftext|>


<commit_before>package client

import (
	"log"
	"encoding/json"
	"path/filepath"

	"github.com/dnaeon/gru/minion"

	"code.google.com/p/go-uuid/uuid"
	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
	etcdclient "github.com/coreos/etcd/client"
)

type EtcdClient struct {
	// KeysAPI client to etcd
	KAPI etcdclient.KeysAPI
}

func NewEtcdClient(cfg etcdclient.Config) *EtcdClient {
	c, err := etcdclient.New(cfg)
	if err != nil {
		log.Fatal(err)
	}

	kapi := etcdclient.NewKeysAPI(c)
	klient := &EtcdClient{
		KAPI: kapi,
	}

	return klient
}

func (c *EtcdClient) SubmitTask(u uuid.UUID, t minion.MinionTask) error {
	minionRootDir := filepath.Join(minion.EtcdMinionSpace, u.String())
	queueDir := filepath.Join(minionRootDir, "queue")

	_, err := c.KAPI.Get(context.Background(), minionRootDir, nil)
	if err != nil {
		return err
	}

	data, err := json.Marshal(t)
	if err != nil {
		return err
	}

	_, err = c.KAPI.CreateInOrder(context.Background(), queueDir, string(data), nil)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>NewEtcdClient returns Client interface<commit_after>package client

import (
	"log"
	"encoding/json"
	"path/filepath"

	"github.com/dnaeon/gru/minion"

	"code.google.com/p/go-uuid/uuid"
	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
	etcdclient "github.com/coreos/etcd/client"
)

type EtcdClient struct {
	// KeysAPI client to etcd
	KAPI etcdclient.KeysAPI
}

func NewEtcdClient(cfg etcdclient.Config) Client {
	c, err := etcdclient.New(cfg)
	if err != nil {
		log.Fatal(err)
	}

	kapi := etcdclient.NewKeysAPI(c)
	klient := &EtcdClient{
		KAPI: kapi,
	}

	return klient
}

func (c *EtcdClient) SubmitTask(u uuid.UUID, t minion.MinionTask) error {
	minionRootDir := filepath.Join(minion.EtcdMinionSpace, u.String())
	queueDir := filepath.Join(minionRootDir, "queue")

	_, err := c.KAPI.Get(context.Background(), minionRootDir, nil)
	if err != nil {
		return err
	}

	data, err := json.Marshal(t)
	if err != nil {
		return err
	}

	_, err = c.KAPI.CreateInOrder(context.Background(), queueDir, string(data), nil)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// +build ignore

// Simple tool that tries to understand simple function definitions as
// used in describing the Windwos Native API.
//
// Comments in Go files that begin with "func:" are parsed as a C
// function prototype.
//
// Comments in Go files that begin with "type:" are parsed as a C
// type definitions.
package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"go/parser"
	"go/token"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/scanner"
)

type State int

const (
	StateInit State = iota
	StateParam
	StateMember
	StateExit
)

type Direction int

const (
	DirectionUnspecified Direction = iota
	DirectionIn
	DirectionOut
	DirectionInOut
)

type FunctionParameterDefinition struct {
	Direction
	Type string
	Name string
}

type FunctionDefinition struct {
	Type   string
	Name   string
	Params []FunctionParameterDefinition
}

type StructMemberDefinition struct {
	Name string
	Type string
}

type StructDefinition struct {
	Name    string
	Members []StructMemberDefinition
}

var translation = map[string]string{
	"NTSTATUS":       "NtStatus",
	"HANDLE":         "Handle",
	"VOID":           "byte",
	"ULONG":          "uint32",
	"LONG":           "int32",
	"USHORT":         "uint16",
	"SHORT":          "int16",
	"WSTR":           "uint16",
	"BOOLEAN":        "bool",
	"LARGE_INTEGER":  "int64",
	"ULARGE_INTEGER": "uint64",
}

func translate(from string) (to string) {
	var ok bool
	if to, ok = translation[from]; ok {
		return
	}
	if from[0] == 'P' {
		return "*" + translate(from[1:])
	}
	if strings.Contains(from, "_") {
		words := strings.Split(from, "_")
		for _, w := range words {
			to = to + strings.Title(strings.ToLower(w))
		}
	}
	return
}

// ParseFunctionDefinition is a half-arsed parser for C function
// definitions as found in MSDN documentation for NTDLL member
// function.
//
// Example (adapted from
// https://msdn.microsoft.com/en-us/library/windows/hardware/ff567029(v=vs.85).aspx):
//
// NTSTATUS NtOpenSection(
//   _Out_ PHANDLE            SectionHandle,
//   _In_  ACCESS_MASK        DesiredAccess,
//   _In_  POBJECT_ATTRIBUTES ObjectAttributes
// );
func ParseFunctionDefinition(rd io.Reader) (*FunctionDefinition, error) {
	s := &scanner.Scanner{Mode: scanner.GoTokens}
	s.Init(rd)
	state := StateInit
	var f FunctionDefinition
	tokens := make([]string, 0)
	for {
		r := s.Scan()
		if r == scanner.EOF {
			break
		}
		switch state {
		case StateInit:
			switch r {
			case '(':
				if len(tokens) < 2 {
					return nil, fmt.Errorf("function definition needs at least a name and a type")
				}
				f.Name = tokens[0]
				for _, t := range tokens[1:] {
					switch t {
					// ignore WINAPI calling convention
					case "WINAPI":
					default:
						f.Type = t
					}
				}
				f.Type = translate(f.Type)
				tokens = tokens[0:0]
				state = StateParam
			case scanner.Ident:
				tokens = append([]string{s.TokenText()}, tokens...)
			default:
				return nil, fmt.Errorf("parse error: got %s", scanner.TokenString(r))
			}
		case StateParam:
			switch r {
			case scanner.Ident:
				tokens = append([]string{s.TokenText()}, tokens...)
			case ',', ')':
				if len(tokens) < 2 {
					return nil, fmt.Errorf("function parameter needs at least a name and a type")
				}
				p := FunctionParameterDefinition{Name: tokens[0]}
				for _, t := range tokens[1:] {
					switch t {
					case "_In_", "_In_opt_":
						p.Direction = DirectionIn
					case "_Out_", "_Out_opt_":
						p.Direction = DirectionOut
					case "_Inout_":
						p.Direction = DirectionInOut
					case "_Reserved_":
					default:
						p.Type = t
					}
				}
				p.Type = translate(p.Type)
				f.Params = append(f.Params, p)
				tokens = tokens[0:0]
				if r == ')' {
					state = StateExit
				}
			}
		case StateExit:
			switch r {
			case ';':
				state = StateInit
				break
			}
		}
	}
	if state != StateInit {
		return nil, fmt.Errorf("parse error: wrong state %d", state)
	}
	return &f, nil
}

// ParseStructDefinition is a half-arsed parser for C struct
// definitions as found in MSDN documentation for NTDLL / Windows
// Driver types.
//
// Example:
//
// typedef struct _OBJECT_ATTRIBUTES {
//   ULONG           Length;
//   HANDLE          RootDirectory;
//   PUNICODE_STRING ObjectName;
//   ULONG           Attributes;
//   PVOID           SecurityDescriptor;
//   PVOID           SecurityQualityOfService;
// } OBJECT_ATTRIBUTES, *POBJECT_ATTRIBUTES;
func ParseStructDefinition(rd io.Reader) (*StructDefinition, error) {
	s := &scanner.Scanner{Mode: scanner.GoTokens}
	s.Init(rd)
	state := StateInit
	var sd StructDefinition
	var name string
	for {
		r := s.Scan()
		if r == scanner.EOF {
			break
		}
		switch state {
		case StateInit:
			e := fmt.Errorf("expecting typedef struct _X {")
			if r != scanner.Ident || s.TokenText() != "typedef" {
				return nil, e
			}
			if r = s.Scan(); r != scanner.Ident || s.TokenText() != "struct" {
				return nil, e
			}
			if r = s.Scan(); r != scanner.Ident || !strings.HasPrefix(s.TokenText(), "_") {
				return nil, e
			}
			name = s.TokenText()[1:]
			sd.Name = translate(name)
			if r = s.Scan(); r != '{' {
				return nil, e
			}
			state = StateMember
		case StateMember:
			if r == '}' {
				state = StateExit
				continue
			}
			if r != scanner.Ident {
				return nil, fmt.Errorf("%s: expecting type, got '%s'", name, s.TokenText())
			}
			m := StructMemberDefinition{Type: translate(s.TokenText())}
			if r = s.Scan(); r != scanner.Ident {
				return nil, fmt.Errorf("%s: expecting type / name pair, got '%s'", name, s.TokenText())
			}
			m.Name = s.TokenText()
			if r = s.Scan(); r != ';' {
				return nil, fmt.Errorf("%s: expecting semicolon after type / name pair, got '%s'", name, s.TokenText())
			}
			sd.Members = append(sd.Members, m)
		case StateExit:
			if r == ';' {
				state = StateInit
				break
			}
			var typ string
			if r == '*' {
				typ = "*"
				r = s.Scan()
			}
			if r != scanner.Ident {
				return nil, fmt.Errorf("expecting type def, got '%s'", r)
			}
			typ += s.TokenText()
			if strings.HasPrefix(typ, "*P") {
				typ = typ[2:]
			}
			if typ != name {
				return nil, fmt.Errorf("expecting type %s, got %s", name, typ)
			}
			if s.Peek() == ',' {
				s.Scan()
			}
		}
	}
	if state != StateInit {
		return nil, fmt.Errorf("parse error: wrong state %d", state)
	}
	return &sd, nil
}

func main() {
	var outfile string
	var functions []FunctionDefinition
	var structs []StructDefinition
	flag.StringVar(&outfile, "output", "", "output file")
	flag.Parse()
	if flag.NArg() != 1 {
		log.Fatal("pass exactly one file")
	}
	infile := filepath.Clean(flag.Arg(0))
	fset := token.NewFileSet()
	if outfile == "" {
		ext := filepath.Ext(infile)
		outfile = infile[:len(infile)-len(ext)] + "_generated" + ext
	}
	f, err := os.Open(infile)
	if err != nil {
		log.Fatal(err)
	}
	file, err := parser.ParseFile(fset, "", f, parser.ParseComments)
	if err != nil {
		log.Fatal(err)
	}
	for _, cgroup := range file.Comments {
		comment := strings.TrimSpace(cgroup.Text())
		if strings.HasPrefix(comment, "func:") {
			comment = comment[5:]
			f, err := ParseFunctionDefinition(strings.NewReader(comment))
			if err != nil {
				log.Fatal(err)
			}
			functions = append(functions, *f)
		} else if strings.HasPrefix(comment, "type:") {
			comment = comment[5:]
			s, err := ParseStructDefinition(strings.NewReader(comment))
			if err != nil {
				log.Fatal(err)
			}
			structs = append(structs, *s)
		}
	}
	f.Close()
	buf := &bytes.Buffer{}
	fmt.Fprintf(buf, `// This file was autogenerated using %s
// DO NOT EDIT.
package %s

`, strings.Join(append([]string{filepath.Base(os.Args[0])}, os.Args[1:]...), " "), file.Name)

funcs:
	for _, function := range functions {
		for _, param := range function.Params {
			if param.Type == "bool" {
				fmt.Fprintln(buf, `import "unsafe"`)
				break funcs
			}
		}
	}

	fmt.Fprintln(buf, "var (")
	for _, function := range functions {
		fmt.Fprintf(buf, `proc%[1]s = modntdll.NewProc("%[1]s")
`, function.Name)
	}
	fmt.Fprintln(buf, ")")

	for _, st := range structs {
		fmt.Fprintf(buf, "type %s struct {\n", st.Name)
		for _, m := range st.Members {
			fmt.Fprintf(buf, "%s %s\n", m.Name, m.Type)
		}
		fmt.Fprint(buf, "}\n\n")
	}

	for _, function := range functions {
		var plist, alist []string
		for _, param := range function.Params {
			alist = append(alist, fmt.Sprintf("%s %s", param.Name, param.Type))
			if param.Type == "bool" {
				plist = append(plist, fmt.Sprintf("fromBool(%s)", param.Name))
			} else if param.Type[0] == '*' {
				plist = append(plist, fmt.Sprintf("uintptr(unsafe.Pointer(%s))", param.Name))
			} else {
				plist = append(plist, fmt.Sprintf("uintptr(%s)", param.Name))
			}
			// alist = append(alist, fmt.Sprintf("p%d %s", n, param.Type))
			// plist = append(plist, "p"+strconv.Itoa(n))
		}
		fmt.Fprintf(buf, `func %[1]s(%[2]s) %[3]s {
	r0, _, _ := proc%[1]s.Call(%[4]s)
	return %[3]s(r0)
}

`, function.Name, strings.Join(alist, ", "), function.Type, strings.Join(plist, ", "))
	}

	f, err = os.Create(outfile)
	if err != nil {
		log.Fatal(err)
	}
	data, err := format.Source(buf.Bytes())
	if err != nil {
		log.Fatal(string(buf.Bytes()), err)
	}
	f.Write(data)
	f.Close()
}
<commit_msg>mkcode: Fix logic for when to import "unsafe" package<commit_after>// +build ignore

// Simple tool that tries to understand simple function definitions as
// used in describing the Windwos Native API.
//
// Comments in Go files that begin with "func:" are parsed as a C
// function prototype.
//
// Comments in Go files that begin with "type:" are parsed as a C
// type definitions.
package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"go/parser"
	"go/token"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/scanner"
)

type State int

const (
	StateInit State = iota
	StateParam
	StateMember
	StateExit
)

type Direction int

const (
	DirectionUnspecified Direction = iota
	DirectionIn
	DirectionOut
	DirectionInOut
)

type FunctionParameterDefinition struct {
	Direction
	Type string
	Name string
}

type FunctionDefinition struct {
	Type   string
	Name   string
	Params []FunctionParameterDefinition
}

type StructMemberDefinition struct {
	Name string
	Type string
}

type StructDefinition struct {
	Name    string
	Members []StructMemberDefinition
}

var translation = map[string]string{
	"NTSTATUS":       "NtStatus",
	"HANDLE":         "Handle",
	"VOID":           "byte",
	"ULONG":          "uint32",
	"LONG":           "int32",
	"USHORT":         "uint16",
	"SHORT":          "int16",
	"WSTR":           "uint16",
	"BOOLEAN":        "bool",
	"LARGE_INTEGER":  "int64",
	"ULARGE_INTEGER": "uint64",
}

func translate(from string) (to string) {
	var ok bool
	if to, ok = translation[from]; ok {
		return
	}
	if from[0] == 'P' {
		return "*" + translate(from[1:])
	}
	if strings.Contains(from, "_") {
		words := strings.Split(from, "_")
		for _, w := range words {
			to = to + strings.Title(strings.ToLower(w))
		}
	}
	return
}

// ParseFunctionDefinition is a half-arsed parser for C function
// definitions as found in MSDN documentation for NTDLL member
// function.
//
// Example (adapted from
// https://msdn.microsoft.com/en-us/library/windows/hardware/ff567029(v=vs.85).aspx):
//
// NTSTATUS NtOpenSection(
//   _Out_ PHANDLE            SectionHandle,
//   _In_  ACCESS_MASK        DesiredAccess,
//   _In_  POBJECT_ATTRIBUTES ObjectAttributes
// );
func ParseFunctionDefinition(rd io.Reader) (*FunctionDefinition, error) {
	s := &scanner.Scanner{Mode: scanner.GoTokens}
	s.Init(rd)
	state := StateInit
	var f FunctionDefinition
	tokens := make([]string, 0)
	for {
		r := s.Scan()
		if r == scanner.EOF {
			break
		}
		switch state {
		case StateInit:
			switch r {
			case '(':
				if len(tokens) < 2 {
					return nil, fmt.Errorf("function definition needs at least a name and a type")
				}
				f.Name = tokens[0]
				for _, t := range tokens[1:] {
					switch t {
					// ignore WINAPI calling convention
					case "WINAPI":
					default:
						f.Type = t
					}
				}
				f.Type = translate(f.Type)
				tokens = tokens[0:0]
				state = StateParam
			case scanner.Ident:
				tokens = append([]string{s.TokenText()}, tokens...)
			default:
				return nil, fmt.Errorf("parse error: got %s", scanner.TokenString(r))
			}
		case StateParam:
			switch r {
			case scanner.Ident:
				tokens = append([]string{s.TokenText()}, tokens...)
			case ',', ')':
				if len(tokens) < 2 {
					return nil, fmt.Errorf("function parameter needs at least a name and a type")
				}
				p := FunctionParameterDefinition{Name: tokens[0]}
				for _, t := range tokens[1:] {
					switch t {
					case "_In_", "_In_opt_":
						p.Direction = DirectionIn
					case "_Out_", "_Out_opt_":
						p.Direction = DirectionOut
					case "_Inout_":
						p.Direction = DirectionInOut
					case "_Reserved_":
					default:
						p.Type = t
					}
				}
				p.Type = translate(p.Type)
				f.Params = append(f.Params, p)
				tokens = tokens[0:0]
				if r == ')' {
					state = StateExit
				}
			}
		case StateExit:
			switch r {
			case ';':
				state = StateInit
				break
			}
		}
	}
	if state != StateInit {
		return nil, fmt.Errorf("parse error: wrong state %d", state)
	}
	return &f, nil
}

// ParseStructDefinition is a half-arsed parser for C struct
// definitions as found in MSDN documentation for NTDLL / Windows
// Driver types.
//
// Example:
//
// typedef struct _OBJECT_ATTRIBUTES {
//   ULONG           Length;
//   HANDLE          RootDirectory;
//   PUNICODE_STRING ObjectName;
//   ULONG           Attributes;
//   PVOID           SecurityDescriptor;
//   PVOID           SecurityQualityOfService;
// } OBJECT_ATTRIBUTES, *POBJECT_ATTRIBUTES;
func ParseStructDefinition(rd io.Reader) (*StructDefinition, error) {
	s := &scanner.Scanner{Mode: scanner.GoTokens}
	s.Init(rd)
	state := StateInit
	var sd StructDefinition
	var name string
	for {
		r := s.Scan()
		if r == scanner.EOF {
			break
		}
		switch state {
		case StateInit:
			e := fmt.Errorf("expecting typedef struct _X {")
			if r != scanner.Ident || s.TokenText() != "typedef" {
				return nil, e
			}
			if r = s.Scan(); r != scanner.Ident || s.TokenText() != "struct" {
				return nil, e
			}
			if r = s.Scan(); r != scanner.Ident || !strings.HasPrefix(s.TokenText(), "_") {
				return nil, e
			}
			name = s.TokenText()[1:]
			sd.Name = translate(name)
			if r = s.Scan(); r != '{' {
				return nil, e
			}
			state = StateMember
		case StateMember:
			if r == '}' {
				state = StateExit
				continue
			}
			if r != scanner.Ident {
				return nil, fmt.Errorf("%s: expecting type, got '%s'", name, s.TokenText())
			}
			m := StructMemberDefinition{Type: translate(s.TokenText())}
			if r = s.Scan(); r != scanner.Ident {
				return nil, fmt.Errorf("%s: expecting type / name pair, got '%s'", name, s.TokenText())
			}
			m.Name = s.TokenText()
			if r = s.Scan(); r != ';' {
				return nil, fmt.Errorf("%s: expecting semicolon after type / name pair, got '%s'", name, s.TokenText())
			}
			sd.Members = append(sd.Members, m)
		case StateExit:
			if r == ';' {
				state = StateInit
				break
			}
			var typ string
			if r == '*' {
				typ = "*"
				r = s.Scan()
			}
			if r != scanner.Ident {
				return nil, fmt.Errorf("expecting type def, got '%s'", r)
			}
			typ += s.TokenText()
			if strings.HasPrefix(typ, "*P") {
				typ = typ[2:]
			}
			if typ != name {
				return nil, fmt.Errorf("expecting type %s, got %s", name, typ)
			}
			if s.Peek() == ',' {
				s.Scan()
			}
		}
	}
	if state != StateInit {
		return nil, fmt.Errorf("parse error: wrong state %d", state)
	}
	return &sd, nil
}

func main() {
	var outfile string
	var functions []FunctionDefinition
	var structs []StructDefinition
	flag.StringVar(&outfile, "output", "", "output file")
	flag.Parse()
	if flag.NArg() != 1 {
		log.Fatal("pass exactly one file")
	}
	infile := filepath.Clean(flag.Arg(0))
	fset := token.NewFileSet()
	if outfile == "" {
		ext := filepath.Ext(infile)
		outfile = infile[:len(infile)-len(ext)] + "_generated" + ext
	}
	f, err := os.Open(infile)
	if err != nil {
		log.Fatal(err)
	}
	file, err := parser.ParseFile(fset, "", f, parser.ParseComments)
	if err != nil {
		log.Fatal(err)
	}
	for _, cgroup := range file.Comments {
		comment := strings.TrimSpace(cgroup.Text())
		if strings.HasPrefix(comment, "func:") {
			comment = comment[5:]
			f, err := ParseFunctionDefinition(strings.NewReader(comment))
			if err != nil {
				log.Fatal(err)
			}
			functions = append(functions, *f)
		} else if strings.HasPrefix(comment, "type:") {
			comment = comment[5:]
			s, err := ParseStructDefinition(strings.NewReader(comment))
			if err != nil {
				log.Fatal(err)
			}
			structs = append(structs, *s)
		}
	}
	f.Close()
	buf := &bytes.Buffer{}
	fmt.Fprintf(buf, `// This file was autogenerated using %s
// DO NOT EDIT.
package %s

`, strings.Join(append([]string{filepath.Base(os.Args[0])}, os.Args[1:]...), " "), file.Name)

funcs:
	for _, function := range functions {
		for _, param := range function.Params {
			if param.Type[0] == '*' {
				fmt.Fprintln(buf, `import "unsafe"`)
				break funcs
			}
		}
	}

	fmt.Fprintln(buf, "var (")
	for _, function := range functions {
		fmt.Fprintf(buf, `proc%[1]s = modntdll.NewProc("%[1]s")
`, function.Name)
	}
	fmt.Fprintln(buf, ")")

	for _, st := range structs {
		fmt.Fprintf(buf, "type %s struct {\n", st.Name)
		for _, m := range st.Members {
			fmt.Fprintf(buf, "%s %s\n", m.Name, m.Type)
		}
		fmt.Fprint(buf, "}\n\n")
	}

	for _, function := range functions {
		var plist, alist []string
		for _, param := range function.Params {
			alist = append(alist, fmt.Sprintf("%s %s", param.Name, param.Type))
			if param.Type == "bool" {
				plist = append(plist, fmt.Sprintf("fromBool(%s)", param.Name))
			} else if param.Type[0] == '*' {
				plist = append(plist, fmt.Sprintf("uintptr(unsafe.Pointer(%s))", param.Name))
			} else {
				plist = append(plist, fmt.Sprintf("uintptr(%s)", param.Name))
			}
			// alist = append(alist, fmt.Sprintf("p%d %s", n, param.Type))
			// plist = append(plist, "p"+strconv.Itoa(n))
		}
		fmt.Fprintf(buf, `func %[1]s(%[2]s) %[3]s {
	r0, _, _ := proc%[1]s.Call(%[4]s)
	return %[3]s(r0)
}

`, function.Name, strings.Join(alist, ", "), function.Type, strings.Join(plist, ", "))
	}

	f, err = os.Create(outfile)
	if err != nil {
		log.Fatal(err)
	}
	data, err := format.Source(buf.Bytes())
	if err != nil {
		log.Fatal(string(buf.Bytes()), err)
	}
	f.Write(data)
	f.Close()
}
<|endoftext|>


<commit_before>package gitlab

import (
	"github.com/stretchr/gomniauth"
	"github.com/stretchr/gomniauth/common"
	"github.com/stretchr/gomniauth/oauth2"
	"github.com/stretchr/objx"
	"net/http"
)

// Constants
gitlabDefaultScope    string = "user"
gitlabName            string = "gitlab"
gitlabDisplayName     string = "Gitlab"
gitlabAuthURL         string = "https://gitlab.com/oauth/authorize"
gitlabTokenURL        string = "https://github.com/login/oauth/access_token"
gitlabEndpointProfile string = "https://api.github.com/user"

// GitlabProvider implements the Provider interface and provides Gitlab
// OAuth2 communication capabilities.
type GitlabProvider struct {
	config         *common.Config
	tripperFactory common.TripperFactory
}

func New(clientId, clientSecret, redirectUrl string) *GitlabProvider {

	p := new(GithubProvider)
	p.config = &common.Config{Map: objx.MSI(
		oauth2.OAuth2KeyAuthURL, gitlabAuthURL,
		oauth2.OAuth2KeyTokenURL, gitlabTokenURL,
		oauth2.OAuth2KeyClientID, clientId,
		oauth2.OAuth2KeySecret, clientSecret,
		oauth2.OAuth2KeyRedirectUrl, redirectUrl,
		oauth2.OAuth2KeyScope, githubDefaultScope,
		oauth2.OAuth2KeyAccessType, oauth2.OAuth2AccessTypeOnline,
		oauth2.OAuth2KeyApprovalPrompt, oauth2.OAuth2ApprovalPromptAuto,
		oauth2.OAuth2KeyResponseType, oauth2.OAuth2KeyCode)}
	return p
}

// TripperFactory gets an OAuth2TripperFactory
func (provider *GitlabProvider) TripperFactory() common.TripperFactory {

	if provider.tripperFactory == nil {
		provider.tripperFactory = new(oauth2.OAuth2TripperFactory)
	}

	return provider.tripperFactory
}

// PublicData gets a public readable view of this provider.
func (provider *GitlabProvider) PublicData(options map[string]interface{}) (interface{}, error) {
	return gomniauth.ProviderPublicData(provider, options)
}

// Name is the unique name for this provider.
func (provider *GitlabProvider) Name() string {
	return githubName
}

// DisplayName is the human readable name for this provider.
func (provider *GitlabProvider) DisplayName() string {
	return gitlabDisplayName
}

// GetBeginAuthURL gets the URL that the client must visit in order
// to begin the authentication process.
//
// The state argument contains anything you wish to have sent back to your
// callback endpoint.
// The options argument takes any options used to configure the auth request
// sent to the provider. In the case of OAuth2, the options map can contain:
//   1. A "scope" key providing the desired scope(s). It will be merged with the default scope.
func (provider *GitlabProvider) GetBeginAuthURL(state *common.State, options objx.Map) (string, error) {
	if options != nil {
		scope := oauth2.MergeScopes(options.Get(oauth2.OAuth2KeyScope).Str(), gitlabDefaultScope)
		provider.config.Set(oauth2.OAuth2KeyScope, scope)
	}
	return oauth2.GetBeginAuthURLWithBase(provider.config.Get(oauth2.OAuth2KeyAuthURL).Str(), state, provider.config)
}

// Get makes an authenticated request and returns the data in the
// response as a data map.
func (provider *GitlabProvider) Get(creds *common.Credentials, endpoint string) (objx.Map, error) {
	return oauth2.Get(provider, creds, endpoint)
}

// GetUser uses the specified common.Credentials to access the users profile
// from the remote provider, and builds the appropriate User object.
func (provider *GitlabProvider) GetUser(creds *common.Credentials) (common.User, error) {

	profileData, err := provider.Get(creds, gitlabEndpointProfile)

	if err != nil {
		return nil, err
	}

	// build user
	user := NewUser(profileData, creds, provider)

	return user, nil
}

// CompleteAuth takes a map of arguments that are used to
// complete the authorisation process, completes it, and returns
// the appropriate Credentials.
func (provider *GitlabProvider) CompleteAuth(data objx.Map) (*common.Credentials, error) {
	return oauth2.CompleteAuth(provider.TripperFactory(), data, provider.config, provider)
}

// GetClient returns an authenticated http.Client that can be used to make requests to
// protected Github resources
func (provider *GitlabProvider) GetClient(creds *common.Credentials) (*http.Client, error) {
	return oauth2.GetClient(provider.TripperFactory(), creds, provider)
}
<commit_msg>a few fixes<commit_after>package gitlab

import (
	"github.com/stretchr/gomniauth"
	"github.com/stretchr/gomniauth/common"
	"github.com/stretchr/gomniauth/oauth2"
	"github.com/stretchr/objx"
	"net/http"
)

// Constants
gitlabDefaultScope    string = "user"
gitlabName            string = "gitlab"
gitlabDisplayName     string = "Gitlab"
gitlabAuthURL         string = "https://gitlab.com/oauth/authorize"
gitlabTokenURL        string = "https://gitlab.com/login/oauth/token"
gitlabEndpointProfile string = "https://api.gitlab.com/user"

// GitlabProvider implements the Provider interface and provides Gitlab
// OAuth2 communication capabilities.
type GitlabProvider struct {
	config         *common.Config
	tripperFactory common.TripperFactory
}

func New(clientId, clientSecret, redirectUrl string) *GitlabProvider {

	p := new(GithubProvider)
	p.config = &common.Config{Map: objx.MSI(
		oauth2.OAuth2KeyAuthURL, gitlabAuthURL,
		oauth2.OAuth2KeyTokenURL, gitlabTokenURL,
		oauth2.OAuth2KeyClientID, clientId,
		oauth2.OAuth2KeySecret, clientSecret,
		oauth2.OAuth2KeyRedirectUrl, redirectUrl,
		oauth2.OAuth2KeyScope, gitlabDefaultScope,
		oauth2.OAuth2KeyAccessType, oauth2.OAuth2AccessTypeOnline,
		oauth2.OAuth2KeyApprovalPrompt, oauth2.OAuth2ApprovalPromptAuto,
		oauth2.OAuth2KeyResponseType, oauth2.OAuth2KeyCode)}
	return p
}

// TripperFactory gets an OAuth2TripperFactory
func (provider *GitlabProvider) TripperFactory() common.TripperFactory {

	if provider.tripperFactory == nil {
		provider.tripperFactory = new(oauth2.OAuth2TripperFactory)
	}

	return provider.tripperFactory
}

// PublicData gets a public readable view of this provider.
func (provider *GitlabProvider) PublicData(options map[string]interface{}) (interface{}, error) {
	return gomniauth.ProviderPublicData(provider, options)
}

// Name is the unique name for this provider.
func (provider *GitlabProvider) Name() string {
	return githubName
}

// DisplayName is the human readable name for this provider.
func (provider *GitlabProvider) DisplayName() string {
	return gitlabDisplayName
}

// GetBeginAuthURL gets the URL that the client must visit in order
// to begin the authentication process.
//
// The state argument contains anything you wish to have sent back to your
// callback endpoint.
// The options argument takes any options used to configure the auth request
// sent to the provider. In the case of OAuth2, the options map can contain:
//   1. A "scope" key providing the desired scope(s). It will be merged with the default scope.
func (provider *GitlabProvider) GetBeginAuthURL(state *common.State, options objx.Map) (string, error) {
	if options != nil {
		scope := oauth2.MergeScopes(options.Get(oauth2.OAuth2KeyScope).Str(), gitlabDefaultScope)
		provider.config.Set(oauth2.OAuth2KeyScope, scope)
	}
	return oauth2.GetBeginAuthURLWithBase(provider.config.Get(oauth2.OAuth2KeyAuthURL).Str(), state, provider.config)
}

// Get makes an authenticated request and returns the data in the
// response as a data map.
func (provider *GitlabProvider) Get(creds *common.Credentials, endpoint string) (objx.Map, error) {
	return oauth2.Get(provider, creds, endpoint)
}

// GetUser uses the specified common.Credentials to access the users profile
// from the remote provider, and builds the appropriate User object.
func (provider *GitlabProvider) GetUser(creds *common.Credentials) (common.User, error) {

	profileData, err := provider.Get(creds, gitlabEndpointProfile)

	if err != nil {
		return nil, err
	}

	// build user
	user := NewUser(profileData, creds, provider)

	return user, nil
}

// CompleteAuth takes a map of arguments that are used to
// complete the authorisation process, completes it, and returns
// the appropriate Credentials.
func (provider *GitlabProvider) CompleteAuth(data objx.Map) (*common.Credentials, error) {
	return oauth2.CompleteAuth(provider.TripperFactory(), data, provider.config, provider)
}

// GetClient returns an authenticated http.Client that can be used to make requests to
// protected Github resources
func (provider *GitlabProvider) GetClient(creds *common.Credentials) (*http.Client, error) {
	return oauth2.GetClient(provider.TripperFactory(), creds, provider)
}
<|endoftext|>


<commit_before>package aws_signing_client

import (
	"net/http"
	"testing"

	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/signer/v4"
	"strings"
)

type testRoundTripper struct {
	err error
}

var (
	creds   = credentials.NewStaticCredentials("ID", "SECRET", "TOKEN")
	signer  *v4.Signer
	client  *http.Client
	service string
	region  string

	rt        *testRoundTripper
	newClient *http.Client
	err       error
)

func init() {
	Init()
}

func Init() {
	rt = &testRoundTripper{}
	signer = v4.NewSigner(creds)
	client = http.DefaultClient
	client.Transport = rt
	service = "es"
	region = "us-east-1"
	newClient, _ = nc()
	err = nil
}

func nc() (*http.Client, error) {
	return New(signer, client, service, region)
}

//   _   _                ____ _ _            _
//  | \ | | _____      __/ ___| (_) ___ _ __ | |_
//  |  \| |/ _ \ \ /\ / / |   | | |/ _ \ '_ \| __|
//  | |\  |  __/\ V  V /| |___| | |  __/ | | | |_
//  |_| \_|\___| \_/\_/  \____|_|_|\___|_| |_|\__|
//

// TestNewClientWithoutSigner tests the NewClient() function when it is not passed a *v4.Signer.
func TestNewClientWithoutSigner(t *testing.T) {
	Init()
	signer = nil
	_, err = nc()
	if err != (MissingSignerError{}) {
		t.Error("Error was not of type MissingSignerError")
	}
}

// TestNewClientWithoutService tests the NewClient() function when it is not passed a service string.
func TestNewClientWithoutService(t *testing.T) {
	Init()
	service = ""
	_, err = nc()
	if err != (MissingServiceError{}) {
		t.Error("Error was not of type MissingServiceError")
	}
}

// TestNewClientWithoutRegion tests the NewClient() function when it is not passed a region string.
func TestNewClientWithoutRegion(t *testing.T) {
	Init()
	region = ""
	_, err = nc()
	if err != (MissingRegionError{}) {
		t.Error("Error was not of type MissingRegionError")
	}
}

// TestNewClient tests the NewClient() function when all is right in the World.
func TestNewClient(t *testing.T) {
	Init()
	newClient, err = nc()
	switch {
	case err != nil:
		t.Errorf("An unexpected error occurred while creating a new client with valid parameters: %s", err)
	case newClient == nil:
		t.Error("A nil *http.Client was returned while creating a new client with valid parameters")
	}
}

//   ____                       _ _____     _
//  |  _ \ ___  _   _ _ __   __| |_   _| __(_)_ __
//  | |_) / _ \| | | | '_ \ / _` | | || '__| | '_ \
//  |  _ < (_) | |_| | | | | (_| | | || |  | | |_) |
//  |_| \_\___/ \__,_|_| |_|\__,_| |_||_|  |_| .__/
//                                           |_|

var passedReq *http.Request

// TestRoundTripSignsGetRequest ensures that a GET request is signed before sending.
func TestRoundTripSignsGetRequest(t *testing.T) {
	Init()
	_, err = newClient.Get("https://google.com")
	checkSignatures(t)
}

// TestRoundTripSignsPostRequest ensures that a GET request is signed before sending.
func TestRoundTripSignsPostRequest(t *testing.T) {
	Init()
	_, err = newClient.Post("https://google.com", "application/json", strings.NewReader("{}"))
	checkSignatures(t)
}

func checkSignatures(t *testing.T) {
	switch {
	case err != nil:
		t.Errorf("An unexpected error occurred while making a request: %s", err)
	case passedReq.Header == nil:
		t.Error("nil headers were returned from the signing request")
	case len(passedReq.Header["x-amz-date"]) == 0:
		t.Error("No 'x-amz-date' header was returned from the signing request")
	case len(passedReq.Header["x-amz-security-token"]) == 0:
		t.Error("No 'x-amz-security-token' header was returned from the signing request")
	}
}

func (rt *testRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	passedReq = req
	return &http.Response{}, rt.err
}
<commit_msg>Updating tests to check for every required signed header key<commit_after>package aws_signing_client

import (
	"net/http"
	"testing"

	"strings"

	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/signer/v4"
	"fmt"
)

type testRoundTripper struct {
	err error
}

var (
	creds   = credentials.NewStaticCredentials("ID", "SECRET", "TOKEN")
	v4s     *v4.Signer
	client  *http.Client
	service string
	region  string

	rt        *testRoundTripper
	newClient *http.Client
	err       error
)

func init() {
	Init()
}

func Init() {
	rt = &testRoundTripper{}
	v4s = v4.NewSigner(creds)
	client = http.DefaultClient
	client.Transport = rt
	service = "es"
	region = "us-east-1"
	newClient, _ = nc()
	err = nil
}

func nc() (*http.Client, error) {
	return New(v4s, client, service, region)
}

//   _   _                ____ _ _            _
//  | \ | | _____      __/ ___| (_) ___ _ __ | |_
//  |  \| |/ _ \ \ /\ / / |   | | |/ _ \ '_ \| __|
//  | |\  |  __/\ V  V /| |___| | |  __/ | | | |_
//  |_| \_|\___| \_/\_/  \____|_|_|\___|_| |_|\__|
//

// TestNewClientWithoutSigner tests the NewClient() function when it is not passed a *v4.Signer.
func TestNewClientWithoutSigner(t *testing.T) {
	Init()
	v4s = nil
	_, err = nc()
	if err != (MissingSignerError{}) {
		t.Error("Error was not of type MissingSignerError")
	}
}

// TestNewClientWithoutService tests the NewClient() function when it is not passed a service string.
func TestNewClientWithoutService(t *testing.T) {
	Init()
	service = ""
	_, err = nc()
	if err != (MissingServiceError{}) {
		t.Error("Error was not of type MissingServiceError")
	}
}

// TestNewClientWithoutRegion tests the NewClient() function when it is not passed a region string.
func TestNewClientWithoutRegion(t *testing.T) {
	Init()
	region = ""
	_, err = nc()
	if err != (MissingRegionError{}) {
		t.Error("Error was not of type MissingRegionError")
	}
}

// TestNewClient tests the NewClient() function when all is right in the World.
func TestNewClient(t *testing.T) {
	Init()
	newClient, err = nc()
	switch {
	case err != nil:
		t.Errorf("An unexpected error occurred while creating a new client with valid parameters: %s", err)
	case newClient == nil:
		t.Error("A nil *http.Client was returned while creating a new client with valid parameters")
	}
}

//   ____                       _ _____     _
//  |  _ \ ___  _   _ _ __   __| |_   _| __(_)_ __
//  | |_) / _ \| | | | '_ \ / _` | | || '__| | '_ \
//  |  _ < (_) | |_| | | | | (_| | | || |  | | |_) |
//  |_| \_\___/ \__,_|_| |_|\__,_| |_||_|  |_| .__/
//                                           |_|

var passedReq *http.Request

// TestRoundTripSignsGetRequest ensures that a GET request is signed before sending.
func TestRoundTripSignsGetRequest(t *testing.T) {
	Init()
	_, err = newClient.Get("https://google.com")
	checkSignatures(t)
}

// TestRoundTripSignsPostRequest ensures that a GET request is signed before sending.
func TestRoundTripSignsPostRequest(t *testing.T) {
	Init()
	_, err = newClient.Post("https://google.com", "application/json", strings.NewReader("{}"))
	checkSignatures(t)
}

func checkSignatures(t *testing.T) {
	auth, ok := passedReq.Header["Authorization"]
	switch  {
	case err != nil:
		t.Errorf("An unexpected error occurred while making a request: %s", err)
	case len(passedReq.Header["X-Amz-Date"]) == 0:
		t.Error("No 'X-Amz-Date' header was returned from the signing request")
	case len(passedReq.Header["X-Amz-Security-Token"]) == 0:
		t.Error("No 'X-Amz-Security-Token' header was returned from the signing request")
	case !ok:
		t.Error("No 'Authorization' header was returned from the signing request")
	case !strings.HasPrefix(auth[0], "AWS4-HMAC-SHA256 "):
		t.Error("Authorization header returned does not begin with 'AWS4-HMAC-SHA256'")
	}
}

func (rt *testRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	passedReq = req
	return &http.Response{}, rt.err
}
<|endoftext|>


<commit_before>package statsd

import (
	"bytes"
	"fmt"
	"net"
	"os"
	"reflect"
	"regexp"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/quipo/statsd/event"
)

// MockNetConn is a mock for net.Conn
type MockNetConn struct {
	buf bytes.Buffer
}

func (mock *MockNetConn) Read(b []byte) (n int, err error) {
	return mock.buf.Read(b)
}
func (mock *MockNetConn) Write(b []byte) (n int, err error) {
	return mock.buf.Write(append(b, '\n'))
}
func (mock MockNetConn) Close() error {
	mock.buf.Truncate(0)
	return nil
}
func (mock MockNetConn) LocalAddr() net.Addr {
	return nil
}
func (mock MockNetConn) RemoteAddr() net.Addr {
	return nil
}
func (mock MockNetConn) SetDeadline(t time.Time) error {
	return nil
}
func (mock MockNetConn) SetReadDeadline(t time.Time) error {
	return nil
}
func (mock MockNetConn) SetWriteDeadline(t time.Time) error {
	return nil
}

func newLocalListenerUDP(t *testing.T) (*net.UDPConn, *net.UDPAddr) {
	addr := fmt.Sprintf(":%d", getFreePort())
	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		t.Fatal(err)
	}
	ln, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		t.Fatal(err)
	}
	return ln, udpAddr
}

func TestTotal(t *testing.T) {
	ln, udpAddr := newLocalListenerUDP(t)
	defer ln.Close()

	prefix := "myproject."

	client := NewStatsdClient(udpAddr.String(), prefix)

	ch := make(chan string, 0)

	s := map[string]int64{
		"a:b:c": 5,
		"d:e:f": 2,
		"x:b:c": 5,
		"g.h.i": 1,
	}

	expected := make(map[string]int64)
	for k, v := range s {
		expected[k] = v
	}

	// also test %HOST% replacement
	s["zz.%HOST%"] = 1
	hostname, err := os.Hostname()
	expected["zz."+hostname] = 1

	go doListenUDP(t, ln, ch, len(s))

	err = client.CreateSocket()
	if nil != err {
		t.Fatal(err)
	}
	defer client.Close()

	for k, v := range s {
		client.Total(k, v)
	}

	actual := make(map[string]int64)

	re := regexp.MustCompile(`^(.*)\:(\d+)\|(\w).*$`)

	for i := len(s); i > 0; i-- {
		x := <-ch
		x = strings.TrimSpace(x)
		//fmt.Println(x)
		if !strings.HasPrefix(x, prefix) {
			t.Errorf("Metric without expected prefix: expected '%s', actual '%s'", prefix, x)
			break
		}
		vv := re.FindStringSubmatch(x)
		if vv[3] != "t" {
			t.Errorf("Metric without expected suffix: expected 't', actual '%s'", vv[3])
		}
		v, err := strconv.ParseInt(vv[2], 10, 64)
		if err != nil {
			t.Error(err)
		}
		actual[vv[1][len(prefix):]] = v
	}

	if !reflect.DeepEqual(expected, actual) {
		t.Errorf("did not receive all metrics: Expected: %T %v, Actual: %T %v ", expected, expected, actual, actual)
	}
}

func doListenUDP(t *testing.T, conn *net.UDPConn, ch chan string, n int) {
	for n > 0 {
		// Handle the connection in a new goroutine.
		// The loop then returns to accepting, so that
		// multiple connections may be served concurrently.
		go func(c *net.UDPConn, ch chan string) {
			buffer := make([]byte, 1024)
			size, err := c.Read(buffer)
			// size, address, err := sock.ReadFrom(buffer) <- This starts printing empty and nil values below immediatly
			if err != nil {
				fmt.Println(string(buffer), size, err)
				t.Fatal(err)
			}
			ch <- string(buffer)
		}(conn, ch)
		n--
	}
}

func doListenTCP(t *testing.T, conn net.Listener, ch chan string, n int) {
	client, err := conn.Accept()
	for {
		if err != nil {
			t.Fatal(err)
		}

		buf := make([]byte, 1024)
		c, err := client.Read(buf)
		if err != nil {
			if err.Error() == "EOF" {
				return
			}
			t.Fatal(err)
		}

		for _, s := range bytes.Split(buf[:c], []byte{'\n'}) {
			ch <- string(s)
		}
	}
}

func newLocalListenerTCP(t *testing.T) (string, net.Listener) {
	addr := fmt.Sprintf("127.0.0.1:%d", getFreePort())
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		t.Fatal(err)
	}
	return addr, ln
}

func TestTCP(t *testing.T) {
	addr, ln := newLocalListenerTCP(t)
	defer ln.Close()

	prefix := "myproject."
	client := NewStatsdClient(addr, prefix)

	ch := make(chan string, 0)

	s := map[string]int64{
		"a:b:c": 5,
		"d:e:f": 2,
		"x:b:c": 5,
		"g.h.i": 1,
	}

	expected := make(map[string]int64)
	for k, v := range s {
		expected[k] = v
	}

	// also test %HOST% replacement
	s["zz.%HOST%"] = 1
	hostname, err := os.Hostname()
	expected["zz."+hostname] = 1

	go doListenTCP(t, ln, ch, len(s))

	err = client.CreateTCPSocket()
	if nil != err {
		t.Fatal(err)
	}
	defer client.Close()

	for k, v := range s {
		client.Total(k, v)
	}

	actual := make(map[string]int64)

	re := regexp.MustCompile(`^(.*)\:(\d+)\|(\w).*$`)

	for i := len(s); i > 0; i-- {
		x := <-ch
		x = strings.TrimSpace(x)
		//fmt.Println(x)
		if !strings.HasPrefix(x, prefix) {
			t.Errorf("Metric without expected prefix: expected '%s', actual '%s'", prefix, x)
			break
		}
		vv := re.FindStringSubmatch(x)
		if vv[3] != "t" {
			t.Errorf("Metric without expected suffix: expected 't', actual '%s'", vv[3])
		}
		v, err := strconv.ParseInt(vv[2], 10, 64)
		if err != nil {
			t.Error(err)
		}
		actual[vv[1][len(prefix):]] = v
	}

	if !reflect.DeepEqual(expected, actual) {
		t.Errorf("did not receive all metrics: Expected: %T %v, Actual: %T %v \n", expected, expected, actual, actual)
	}
}

func TestSendEvents(t *testing.T) {
	c := NewStatsdClient("127.0.0.1:1201", "test")
	c.conn = &MockNetConn{} // mock connection

	// override with a small size
	UDPPayloadSize = 40

	e1 := &event.Increment{Name: "test1", Value: 123}
	e2 := &event.Increment{Name: "test2", Value: 432}
	e3 := &event.Increment{Name: "test3", Value: 111}
	e4 := &event.Gauge{Name: "test4", Value: 12435}

	events := map[string]event.Event{
		"test1": e1,
		"test2": e2,
		"test3": e3,
		"test4": e4,
	}

	err := c.SendEvents(events)
	if nil != err {
		t.Error(err)
	}

	b1 := make([]byte, UDPPayloadSize*3)
	n, err2 := c.conn.Read(b1)
	if nil != err2 {
		t.Error(err2)
	}
	nStats := len(strings.Split(strings.TrimSpace(string(b1[:n])), "\n"))
	if nStats != len(events) {
		t.Errorf("Was expecting %d events, got %d:  %s", len(events), nStats, string(b1))
	}
}

// getFreePort Ask the kernel for a free open port that is ready to use
func getFreePort() int {
	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
	if err != nil {
		panic(err)
	}

	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		panic(err)
	}
	defer l.Close()
	return l.Addr().(*net.TCPAddr).Port
}
<commit_msg>client_test: do not receive empty leftover after newline<commit_after>package statsd

import (
	"bytes"
	"fmt"
	"net"
	"os"
	"reflect"
	"regexp"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/quipo/statsd/event"
)

// MockNetConn is a mock for net.Conn
type MockNetConn struct {
	buf bytes.Buffer
}

func (mock *MockNetConn) Read(b []byte) (n int, err error) {
	return mock.buf.Read(b)
}
func (mock *MockNetConn) Write(b []byte) (n int, err error) {
	return mock.buf.Write(append(b, '\n'))
}
func (mock MockNetConn) Close() error {
	mock.buf.Truncate(0)
	return nil
}
func (mock MockNetConn) LocalAddr() net.Addr {
	return nil
}
func (mock MockNetConn) RemoteAddr() net.Addr {
	return nil
}
func (mock MockNetConn) SetDeadline(t time.Time) error {
	return nil
}
func (mock MockNetConn) SetReadDeadline(t time.Time) error {
	return nil
}
func (mock MockNetConn) SetWriteDeadline(t time.Time) error {
	return nil
}

func newLocalListenerUDP(t *testing.T) (*net.UDPConn, *net.UDPAddr) {
	addr := fmt.Sprintf(":%d", getFreePort())
	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		t.Fatal(err)
	}
	ln, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		t.Fatal(err)
	}
	return ln, udpAddr
}

func TestTotal(t *testing.T) {
	ln, udpAddr := newLocalListenerUDP(t)
	defer ln.Close()

	prefix := "myproject."

	client := NewStatsdClient(udpAddr.String(), prefix)

	ch := make(chan string, 0)

	s := map[string]int64{
		"a:b:c": 5,
		"d:e:f": 2,
		"x:b:c": 5,
		"g.h.i": 1,
	}

	expected := make(map[string]int64)
	for k, v := range s {
		expected[k] = v
	}

	// also test %HOST% replacement
	s["zz.%HOST%"] = 1
	hostname, err := os.Hostname()
	expected["zz."+hostname] = 1

	go doListenUDP(t, ln, ch, len(s))

	err = client.CreateSocket()
	if nil != err {
		t.Fatal(err)
	}
	defer client.Close()

	for k, v := range s {
		client.Total(k, v)
	}

	actual := make(map[string]int64)

	re := regexp.MustCompile(`^(.*)\:(\d+)\|(\w).*$`)

	for i := len(s); i > 0; i-- {
		x := <-ch
		x = strings.TrimSpace(x)
		//fmt.Println(x)
		if !strings.HasPrefix(x, prefix) {
			t.Errorf("Metric without expected prefix: expected '%s', actual '%s'", prefix, x)
			break
		}
		vv := re.FindStringSubmatch(x)
		if vv[3] != "t" {
			t.Errorf("Metric without expected suffix: expected 't', actual '%s'", vv[3])
		}
		v, err := strconv.ParseInt(vv[2], 10, 64)
		if err != nil {
			t.Error(err)
		}
		actual[vv[1][len(prefix):]] = v
	}

	if !reflect.DeepEqual(expected, actual) {
		t.Errorf("did not receive all metrics: Expected: %T %v, Actual: %T %v ", expected, expected, actual, actual)
	}
}

func doListenUDP(t *testing.T, conn *net.UDPConn, ch chan string, n int) {
	for n > 0 {
		// Handle the connection in a new goroutine.
		// The loop then returns to accepting, so that
		// multiple connections may be served concurrently.
		go func(c *net.UDPConn, ch chan string) {
			buffer := make([]byte, 1024)
			size, err := c.Read(buffer)
			// size, address, err := sock.ReadFrom(buffer) <- This starts printing empty and nil values below immediatly
			if err != nil {
				fmt.Println(string(buffer), size, err)
				t.Fatal(err)
			}
			ch <- string(buffer)
		}(conn, ch)
		n--
	}
}

func doListenTCP(t *testing.T, conn net.Listener, ch chan string, n int) {
	client, err := conn.Accept()
	for {
		if err != nil {
			t.Fatal(err)
		}

		buf := make([]byte, 1024)
		c, err := client.Read(buf)
		if err != nil {
			if err.Error() == "EOF" {
				return
			}
			t.Fatal(err)
		}

		for _, s := range bytes.Split(buf[:c], []byte{'\n'}) {
			if len(s) > 0 {
				ch <- string(s)
			}
		}
	}
}

func newLocalListenerTCP(t *testing.T) (string, net.Listener) {
	addr := fmt.Sprintf("127.0.0.1:%d", getFreePort())
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		t.Fatal(err)
	}
	return addr, ln
}

func TestTCP(t *testing.T) {
	addr, ln := newLocalListenerTCP(t)
	defer ln.Close()

	prefix := "myproject."
	client := NewStatsdClient(addr, prefix)

	ch := make(chan string, 0)

	s := map[string]int64{
		"a:b:c": 5,
		"d:e:f": 2,
		"x:b:c": 5,
		"g.h.i": 1,
	}

	expected := make(map[string]int64)
	for k, v := range s {
		expected[k] = v
	}

	// also test %HOST% replacement
	s["zz.%HOST%"] = 1
	hostname, err := os.Hostname()
	expected["zz."+hostname] = 1

	go doListenTCP(t, ln, ch, len(s))

	err = client.CreateTCPSocket()
	if nil != err {
		t.Fatal(err)
	}
	defer client.Close()

	for k, v := range s {
		client.Total(k, v)
	}

	actual := make(map[string]int64)

	re := regexp.MustCompile(`^(.*)\:(\d+)\|(\w).*$`)

	for i := len(s); i > 0; i-- {
		x := <-ch
		x = strings.TrimSpace(x)
		//fmt.Println(x)
		if !strings.HasPrefix(x, prefix) {
			t.Errorf("Metric without expected prefix: expected '%s', actual '%s'", prefix, x)
			break
		}
		vv := re.FindStringSubmatch(x)
		if vv[3] != "t" {
			t.Errorf("Metric without expected suffix: expected 't', actual '%s'", vv[3])
		}
		v, err := strconv.ParseInt(vv[2], 10, 64)
		if err != nil {
			t.Error(err)
		}
		actual[vv[1][len(prefix):]] = v
	}

	if !reflect.DeepEqual(expected, actual) {
		t.Errorf("did not receive all metrics: Expected: %T %v, Actual: %T %v \n", expected, expected, actual, actual)
	}
}

func TestSendEvents(t *testing.T) {
	c := NewStatsdClient("127.0.0.1:1201", "test")
	c.conn = &MockNetConn{} // mock connection

	// override with a small size
	UDPPayloadSize = 40

	e1 := &event.Increment{Name: "test1", Value: 123}
	e2 := &event.Increment{Name: "test2", Value: 432}
	e3 := &event.Increment{Name: "test3", Value: 111}
	e4 := &event.Gauge{Name: "test4", Value: 12435}

	events := map[string]event.Event{
		"test1": e1,
		"test2": e2,
		"test3": e3,
		"test4": e4,
	}

	err := c.SendEvents(events)
	if nil != err {
		t.Error(err)
	}

	b1 := make([]byte, UDPPayloadSize*3)
	n, err2 := c.conn.Read(b1)
	if nil != err2 {
		t.Error(err2)
	}
	nStats := len(strings.Split(strings.TrimSpace(string(b1[:n])), "\n"))
	if nStats != len(events) {
		t.Errorf("Was expecting %d events, got %d:  %s", len(events), nStats, string(b1))
	}
}

// getFreePort Ask the kernel for a free open port that is ready to use
func getFreePort() int {
	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
	if err != nil {
		panic(err)
	}

	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		panic(err)
	}
	defer l.Close()
	return l.Addr().(*net.TCPAddr).Port
}
<|endoftext|>