<commit_before>package i2c

import (
	"errors"
	"strings"
	"testing"

	"gobot.io/x/gobot"
	"gobot.io/x/gobot/gobottest"
)

var _ gobot.Driver = (*AdafruitMotorHatDriver)(nil)

// --------- HELPERS
func initTestAdafruitMotorHatDriver() (driver *AdafruitMotorHatDriver) {
	driver, _ = initTestAdafruitMotorHatDriverWithStubbedAdaptor()
	return
}

func initTestAdafruitMotorHatDriverWithStubbedAdaptor() (*AdafruitMotorHatDriver, *i2cTestAdaptor) {
	adaptor := newI2cTestAdaptor()
	return NewAdafruitMotorHatDriver(adaptor), adaptor
}

// --------- TESTS
func TestNewAdafruitMotorHatDriver(t *testing.T) {
	var adafruit interface{} = NewAdafruitMotorHatDriver(newI2cTestAdaptor())
	_, ok := adafruit.(*AdafruitMotorHatDriver)
	if !ok {
		t.Errorf("AdafruitMotorHatDriver() should have returned a *AdafruitMotorHatDriver")
	}

	a := NewAdafruitMotorHatDriver(newI2cTestAdaptor())
	gobottest.Assert(t, strings.HasPrefix(a.Name(), "AdafruitMotorHat"), true)
}

// Methods
func TestAdafruitMotorHatDriverStart(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()
	gobottest.Refute(t, ada.Connection(), nil)
	gobottest.Assert(t, ada.Start(), nil)
}

func TestAdafruitMotorHatDriverStartError(t *testing.T) {
	d, adaptor := initTestAdafruitMotorHatDriverWithStubbedAdaptor()
	adaptor.i2cWriteImpl = func([]byte) (int, error) {
		return 0, errors.New("write error")
	}
	gobottest.Assert(t, d.Start(), errors.New("write error"))
}

func TestAdafruitMotorHatDriverHalt(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Halt(), nil)
}

func TestSetHatAddresses(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	motorHatAddr := 0x61
	servoHatAddr := 0x41
	gobottest.Assert(t, ada.SetMotorHatAddress(motorHatAddr), nil)
	gobottest.Assert(t, ada.SetServoHatAddress(servoHatAddr), nil)
}

func TestAdafruitMotorHatDriverSetServoMotorFreq(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)

	freq := 60.0
	err := ada.SetServoMotorFreq(freq)
	gobottest.Assert(t, err, nil)
}

func TestAdafruitMotorHatDriverSetServoMotorPulse(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)

	var channel byte = 7
	var on int32 = 1234
	var off int32 = 4321
	err := ada.SetServoMotorPulse(channel, on, off)
	gobottest.Assert(t, err, nil)
}

func TestAdafruitMotorHatDriverSetDCMotorSpeed(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)

	dcMotor := 1
	var speed int32 = 255
	err := ada.SetDCMotorSpeed(dcMotor, speed)
	gobottest.Assert(t, err, nil)
}

func TestAdafruitMotorHatDriverSetDCMotorSpeedError(t *testing.T) {
	ada, a := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)
	a.i2cWriteImpl = func([]byte) (int, error) {
		return 0, errors.New("write error")
	}

	gobottest.Assert(t, ada.SetDCMotorSpeed(1, 255), errors.New("write error"))
}

func TestAdafruitMotorHatDriverRunDCMotor(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)

	dcMotor := 1
	gobottest.Assert(t, ada.RunDCMotor(dcMotor, AdafruitForward), nil)
	gobottest.Assert(t, ada.RunDCMotor(dcMotor, AdafruitBackward), nil)
	gobottest.Assert(t, ada.RunDCMotor(dcMotor, AdafruitRelease), nil)
}

func TestAdafruitMotorHatDriverRunDCMotorError(t *testing.T) {
	ada, a := initTestAdafruitMotorHatDriverWithStubbedAdaptor()
	gobottest.Assert(t, ada.Start(), nil)
	a.i2cWriteImpl = func([]byte) (int, error) {
		return 0, errors.New("write error")
	}

	dcMotor := 1
	gobottest.Assert(t, ada.RunDCMotor(dcMotor, AdafruitForward), errors.New("write error"))
	gobottest.Assert(t, ada.RunDCMotor(dcMotor, AdafruitBackward), errors.New("write error"))
	gobottest.Assert(t, ada.RunDCMotor(dcMotor, AdafruitRelease), errors.New("write error"))
}

func TestAdafruitMotorHatDriverSetStepperMotorSpeed(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)

	stepperMotor := 1
	rpm := 30
	gobottest.Assert(t, ada.SetStepperMotorSpeed(stepperMotor, rpm), nil)
}

func TestAdafruitMotorHatDriverStepperStep(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)

	// NOTE: not using the direction and style constants to prevent importing
	// the i2c package
	stepperMotor := 0
	steps := 50
	err := ada.Step(stepperMotor, steps, 1, 3)
	gobottest.Assert(t, err, nil)
}

func TestAdafruitMotorHatDriverSetName(t *testing.T) {
	d := initTestAdafruitMotorHatDriver()
	d.SetName("TESTME")
	gobottest.Assert(t, d.Name(), "TESTME")
}

func TestAdafruitMotorHatDriverOptions(t *testing.T) {
	d := NewAdafruitMotorHatDriver(newI2cTestAdaptor(), WithBus(2))
	gobottest.Assert(t, d.GetBusOrDefault(1), 2)
}
<commit_msg>i2c: increase test coverage for adafruit motor hat driver<commit_after>package i2c

import (
	"errors"
	"strings"
	"testing"

	"gobot.io/x/gobot"
	"gobot.io/x/gobot/gobottest"
)

var _ gobot.Driver = (*AdafruitMotorHatDriver)(nil)

// --------- HELPERS
func initTestAdafruitMotorHatDriver() (driver *AdafruitMotorHatDriver) {
	driver, _ = initTestAdafruitMotorHatDriverWithStubbedAdaptor()
	return
}

func initTestAdafruitMotorHatDriverWithStubbedAdaptor() (*AdafruitMotorHatDriver, *i2cTestAdaptor) {
	adaptor := newI2cTestAdaptor()
	return NewAdafruitMotorHatDriver(adaptor), adaptor
}

// --------- TESTS
func TestNewAdafruitMotorHatDriver(t *testing.T) {
	var adafruit interface{} = NewAdafruitMotorHatDriver(newI2cTestAdaptor())
	_, ok := adafruit.(*AdafruitMotorHatDriver)
	if !ok {
		t.Errorf("AdafruitMotorHatDriver() should have returned a *AdafruitMotorHatDriver")
	}

	a := NewAdafruitMotorHatDriver(newI2cTestAdaptor())
	gobottest.Assert(t, strings.HasPrefix(a.Name(), "AdafruitMotorHat"), true)
}

// Methods
func TestAdafruitMotorHatDriverStart(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()
	gobottest.Refute(t, ada.Connection(), nil)
	gobottest.Assert(t, ada.Start(), nil)
}

func TestAdafruitMotorHatDriverStartError(t *testing.T) {
	d, adaptor := initTestAdafruitMotorHatDriverWithStubbedAdaptor()
	adaptor.i2cWriteImpl = func([]byte) (int, error) {
		return 0, errors.New("write error")
	}
	gobottest.Assert(t, d.Start(), errors.New("write error"))
}

func TestAdafruitMotorHatDriverHalt(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Halt(), nil)
}

func TestSetHatAddresses(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	motorHatAddr := 0x61
	servoHatAddr := 0x41
	gobottest.Assert(t, ada.SetMotorHatAddress(motorHatAddr), nil)
	gobottest.Assert(t, ada.SetServoHatAddress(servoHatAddr), nil)
}

func TestAdafruitMotorHatDriverSetServoMotorFreq(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)

	freq := 60.0
	err := ada.SetServoMotorFreq(freq)
	gobottest.Assert(t, err, nil)
}

func TestAdafruitMotorHatDriverSetServoMotorPulse(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)

	var channel byte = 7
	var on int32 = 1234
	var off int32 = 4321
	err := ada.SetServoMotorPulse(channel, on, off)
	gobottest.Assert(t, err, nil)
}

func TestAdafruitMotorHatDriverSetDCMotorSpeed(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)

	dcMotor := 1
	var speed int32 = 255
	err := ada.SetDCMotorSpeed(dcMotor, speed)
	gobottest.Assert(t, err, nil)
}

func TestAdafruitMotorHatDriverSetDCMotorSpeedError(t *testing.T) {
	ada, a := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)
	a.i2cWriteImpl = func([]byte) (int, error) {
		return 0, errors.New("write error")
	}

	gobottest.Assert(t, ada.SetDCMotorSpeed(1, 255), errors.New("write error"))
}

func TestAdafruitMotorHatDriverRunDCMotor(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)

	dcMotor := 1
	gobottest.Assert(t, ada.RunDCMotor(dcMotor, AdafruitForward), nil)
	gobottest.Assert(t, ada.RunDCMotor(dcMotor, AdafruitBackward), nil)
	gobottest.Assert(t, ada.RunDCMotor(dcMotor, AdafruitRelease), nil)
}

func TestAdafruitMotorHatDriverRunDCMotorError(t *testing.T) {
	ada, a := initTestAdafruitMotorHatDriverWithStubbedAdaptor()
	gobottest.Assert(t, ada.Start(), nil)
	a.i2cWriteImpl = func([]byte) (int, error) {
		return 0, errors.New("write error")
	}

	dcMotor := 1
	gobottest.Assert(t, ada.RunDCMotor(dcMotor, AdafruitForward), errors.New("write error"))
	gobottest.Assert(t, ada.RunDCMotor(dcMotor, AdafruitBackward), errors.New("write error"))
	gobottest.Assert(t, ada.RunDCMotor(dcMotor, AdafruitRelease), errors.New("write error"))
}

func TestAdafruitMotorHatDriverSetStepperMotorSpeed(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)

	stepperMotor := 1
	rpm := 30
	gobottest.Assert(t, ada.SetStepperMotorSpeed(stepperMotor, rpm), nil)
}

func TestAdafruitMotorHatDriverStepperMicroStep(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)

	// NOTE: not using the direction and style constants to prevent importing
	// the i2c package
	stepperMotor := 0
	steps := 50
	err := ada.Step(stepperMotor, steps, 1, 3)
	gobottest.Assert(t, err, nil)
}

func TestAdafruitMotorHatDriverStepperSingleStep(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)

	// NOTE: not using the direction and style constants to prevent importing
	// the i2c package
	stepperMotor := 0
	steps := 50
	err := ada.Step(stepperMotor, steps, 1, 0)
	gobottest.Assert(t, err, nil)
}

func TestAdafruitMotorHatDriverStepperDoubleStep(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)

	// NOTE: not using the direction and style constants to prevent importing
	// the i2c package
	stepperMotor := 0
	steps := 50
	err := ada.Step(stepperMotor, steps, 1, 1)
	gobottest.Assert(t, err, nil)
}

func TestAdafruitMotorHatDriverStepperInterleaveStep(t *testing.T) {
	ada, _ := initTestAdafruitMotorHatDriverWithStubbedAdaptor()

	gobottest.Assert(t, ada.Start(), nil)

	// NOTE: not using the direction and style constants to prevent importing
	// the i2c package
	stepperMotor := 0
	steps := 50
	err := ada.Step(stepperMotor, steps, 1, 2)
	gobottest.Assert(t, err, nil)
}

func TestAdafruitMotorHatDriverSetName(t *testing.T) {
	d := initTestAdafruitMotorHatDriver()
	d.SetName("TESTME")
	gobottest.Assert(t, d.Name(), "TESTME")
}

func TestAdafruitMotorHatDriverOptions(t *testing.T) {
	d := NewAdafruitMotorHatDriver(newI2cTestAdaptor(), WithBus(2))
	gobottest.Assert(t, d.GetBusOrDefault(1), 2)
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"ioutil"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"strings"

	"github.com/skip2/go-qrcode"
)

const (
	NormalBlack   = "\033[38;5;0m  \033[0m"
	NormalRed     = "\033[38;5;1m  \033[0m"
	NormalGreen   = "\033[38;5;2m  \033[0m"
	NormalYellow  = "\033[38;5;3m  \033[0m"
	NormalBlue    = "\033[38;5;4m  \033[0m"
	NormalMagenta = "\033[38;5;5m  \033[0m"
	NormalCyan    = "\033[38;5;6m  \033[0m"
	NormalWhite   = "\033[38;5;7m  \033[0m"

	BrightBlack   = "\033[48;5;0m  \033[0m"
	BrightRed     = "\033[48;5;1m  \033[0m"
	BrightGreen   = "\033[48;5;2m  \033[0m"
	BrightYellow  = "\033[48;5;3m  \033[0m"
	BrightBlue    = "\033[48;5;4m  \033[0m"
	BrightMagenta = "\033[48;5;5m  \033[0m"
	BrightCyan    = "\033[48;5;6m  \033[0m"
	BrightWhite   = "\033[48;5;7m  \033[0m"
)

var (
	frontColor      string
	backgroundColor string
	levelString     string
	codeJustify     string
)

func init() {
	flag.Usage = printHelp
	flag.StringVar(&frontColor, "f", "black", "Front color")
	flag.StringVar(&backgroundColor, "b", "white", "Background color")
	flag.StringVar(&levelString, "l", "m", "Error correction level")
	if runtime.GOOS != "windows" {
		flag.StringVar(&codeJustify, "j", "left", "QR-Code justify")
	}
}

func main() {

	flag.Parse()

	var front, back, content, justify string
	var err error
	var level qrcode.RecoveryLevel

	if front, err = parseColor(frontColor); err != nil {
		fmt.Println(err)
		printHelp()
		return
	}

	if back, err = parseColor(backgroundColor); err != nil {
		fmt.Println(err)
		printHelp()
		return
	}

	if level, err = parseLevel(levelString); err != nil {
		fmt.Println(err)
		printHelp()
		return
	}

	if justify, err = parseJustify(codeJustify); err != nil {
		fmt.Println(err)
		printHelp()
		return
	}

	if content = flag.Arg(0); content == "" {
		data, err = ioutil.ReadAll(os.Stdin)
		if err != nil {
			fmt.Println(err)
			return
		}
		content = string(data)
	}

	qr, err := qrcode.New(content, level)
	if err != nil {
		fmt.Println(err)
		return
	}

	screenCols, _ := getTTYSize()

	for ir, row := range qr.Bitmap() {
		lr := len(row)

		if ir == 0 || ir == 1 || ir == 2 ||
			ir == lr-1 || ir == lr-2 || ir == lr-3 {
			continue
		}

		if justify == "center" {
			for spaces := 0; spaces < (screenCols/2 - lr/2 - 2*(3*2)); spaces++ {
				fmt.Print(" ")
			}
		}

		if justify == "right" {
			for spaces := 0; spaces < (screenCols - 2*(lr-3*2)); spaces++ {
				fmt.Print(" ")
			}
		}

		for ic, col := range row {
			lc := len(qr.Bitmap())
			if ic == 0 || ic == 1 || ic == 2 ||
				ic == lc-1 || ic == lc-2 || ic == lc-3 {
				continue
			}
			if col {
				fmt.Print(front)
			} else {
				fmt.Print(back)
			}
		}
		fmt.Println()
	}
}

func parseColor(str string) (color string, err error) {
	s := strings.ToUpper(str)
	switch s {
	case "BLACK":
		color = BrightBlack
	case "RED":
		color = BrightRed
	case "GREEN":
		color = BrightGreen
	case "YELLOW":
		color = BrightYellow
	case "BLUE":
		color = BrightBlue
	case "MAGENTA":
		color = BrightMagenta
	case "CYAN":
		color = BrightCyan
	case "WHITE":
		color = BrightWhite
	default:
		err = errors.New(fmt.Sprintf("'%s' is not support.", str))
	}

	return
}

func parseLevel(str string) (level qrcode.RecoveryLevel, err error) {
	s := strings.ToUpper(str)
	switch s {
	case "L":
		level = qrcode.Low
	case "M":
		level = qrcode.Medium
	case "Q":
		level = qrcode.High
	case "H":
		level = qrcode.Highest
	default:
		err = errors.New(fmt.Sprintf("'%s' is not support.", str))
	}

	return
}

func parseJustify(str string) (justify string, err error) {
	s := strings.ToUpper(str)
	switch s {
	case "LEFT":
		justify = "left"
	case "RIGHT":
		justify = "right"
	case "CENTER":
		justify = "center"
	default:
		err = errors.New(fmt.Sprintf("'%s' is not support.", str))
	}
	return
}

func printHelp() {

	helpStr := `QRCode generater terminal edition.

Supported background colors: [black, red, green, yellow, blue, magenta, cyan, white]
Supported front colors: [black, red, green, yellow, blue, magenta, cyan, white]
Supported error correction levels: [L, M, Q, H]
`
	if runtime.GOOS != "windows" {
		helpStr = helpStr + "\nSupported qr-code justifies: [left, center, right]"
	}
	fmt.Println(helpStr)
	flag.PrintDefaults()
}

func getTTYSize() (int, int) {
	cmd := exec.Command("stty", "size")
	cmd.Stdin = os.Stdin
	out, err := cmd.Output()
	if err != nil {
		return 0, 0
	}
	outStr := strings.Replace(string(out), "\n", "", -1)
	cols, err := strconv.Atoi(strings.Split(outStr, " ")[1])
	if err != nil {
		return 0, 0
	}
	rows, err := strconv.Atoi(strings.Split(outStr, " ")[0])
	if err != nil {
		return 0, 0
	}
	return cols, rows
}
<commit_msg>Fix syntax errors & trim final newline<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"strings"

	"github.com/skip2/go-qrcode"
)

const (
	NormalBlack   = "\033[38;5;0m  \033[0m"
	NormalRed     = "\033[38;5;1m  \033[0m"
	NormalGreen   = "\033[38;5;2m  \033[0m"
	NormalYellow  = "\033[38;5;3m  \033[0m"
	NormalBlue    = "\033[38;5;4m  \033[0m"
	NormalMagenta = "\033[38;5;5m  \033[0m"
	NormalCyan    = "\033[38;5;6m  \033[0m"
	NormalWhite   = "\033[38;5;7m  \033[0m"

	BrightBlack   = "\033[48;5;0m  \033[0m"
	BrightRed     = "\033[48;5;1m  \033[0m"
	BrightGreen   = "\033[48;5;2m  \033[0m"
	BrightYellow  = "\033[48;5;3m  \033[0m"
	BrightBlue    = "\033[48;5;4m  \033[0m"
	BrightMagenta = "\033[48;5;5m  \033[0m"
	BrightCyan    = "\033[48;5;6m  \033[0m"
	BrightWhite   = "\033[48;5;7m  \033[0m"
)

var (
	frontColor      string
	backgroundColor string
	levelString     string
	codeJustify     string
)

func init() {
	flag.Usage = printHelp
	flag.StringVar(&frontColor, "f", "black", "Front color")
	flag.StringVar(&backgroundColor, "b", "white", "Background color")
	flag.StringVar(&levelString, "l", "m", "Error correction level")
	if runtime.GOOS != "windows" {
		flag.StringVar(&codeJustify, "j", "left", "QR-Code justify")
	}
}

func main() {

	flag.Parse()

	var front, back, content, justify string
	var err error
	var level qrcode.RecoveryLevel

	if front, err = parseColor(frontColor); err != nil {
		fmt.Println(err)
		printHelp()
		return
	}

	if back, err = parseColor(backgroundColor); err != nil {
		fmt.Println(err)
		printHelp()
		return
	}

	if level, err = parseLevel(levelString); err != nil {
		fmt.Println(err)
		printHelp()
		return
	}

	if justify, err = parseJustify(codeJustify); err != nil {
		fmt.Println(err)
		printHelp()
		return
	}

	if content = flag.Arg(0); content == "" {
		data, err := ioutil.ReadAll(os.Stdin)
		if err != nil {
			fmt.Println(err)
			return
		}
		content = strings.TrimSuffix(string(data), "\n")
	}

	qr, err := qrcode.New(content, level)
	if err != nil {
		fmt.Println(err)
		return
	}

	screenCols, _ := getTTYSize()

	for ir, row := range qr.Bitmap() {
		lr := len(row)

		if ir == 0 || ir == 1 || ir == 2 ||
			ir == lr-1 || ir == lr-2 || ir == lr-3 {
			continue
		}

		if justify == "center" {
			for spaces := 0; spaces < (screenCols/2 - lr/2 - 2*(3*2)); spaces++ {
				fmt.Print(" ")
			}
		}

		if justify == "right" {
			for spaces := 0; spaces < (screenCols - 2*(lr-3*2)); spaces++ {
				fmt.Print(" ")
			}
		}

		for ic, col := range row {
			lc := len(qr.Bitmap())
			if ic == 0 || ic == 1 || ic == 2 ||
				ic == lc-1 || ic == lc-2 || ic == lc-3 {
				continue
			}
			if col {
				fmt.Print(front)
			} else {
				fmt.Print(back)
			}
		}
		fmt.Println()
	}
}

func parseColor(str string) (color string, err error) {
	s := strings.ToUpper(str)
	switch s {
	case "BLACK":
		color = BrightBlack
	case "RED":
		color = BrightRed
	case "GREEN":
		color = BrightGreen
	case "YELLOW":
		color = BrightYellow
	case "BLUE":
		color = BrightBlue
	case "MAGENTA":
		color = BrightMagenta
	case "CYAN":
		color = BrightCyan
	case "WHITE":
		color = BrightWhite
	default:
		err = errors.New(fmt.Sprintf("'%s' is not support.", str))
	}

	return
}

func parseLevel(str string) (level qrcode.RecoveryLevel, err error) {
	s := strings.ToUpper(str)
	switch s {
	case "L":
		level = qrcode.Low
	case "M":
		level = qrcode.Medium
	case "Q":
		level = qrcode.High
	case "H":
		level = qrcode.Highest
	default:
		err = errors.New(fmt.Sprintf("'%s' is not support.", str))
	}

	return
}

func parseJustify(str string) (justify string, err error) {
	s := strings.ToUpper(str)
	switch s {
	case "LEFT":
		justify = "left"
	case "RIGHT":
		justify = "right"
	case "CENTER":
		justify = "center"
	default:
		err = errors.New(fmt.Sprintf("'%s' is not support.", str))
	}
	return
}

func printHelp() {

	helpStr := `QRCode generater terminal edition.

Supported background colors: [black, red, green, yellow, blue, magenta, cyan, white]
Supported front colors: [black, red, green, yellow, blue, magenta, cyan, white]
Supported error correction levels: [L, M, Q, H]
`
	if runtime.GOOS != "windows" {
		helpStr = helpStr + "\nSupported qr-code justifies: [left, center, right]"
	}
	fmt.Println(helpStr)
	flag.PrintDefaults()
}

func getTTYSize() (int, int) {
	cmd := exec.Command("stty", "size")
	cmd.Stdin = os.Stdin
	out, err := cmd.Output()
	if err != nil {
		return 0, 0
	}
	outStr := strings.Replace(string(out), "\n", "", -1)
	cols, err := strconv.Atoi(strings.Split(outStr, " ")[1])
	if err != nil {
		return 0, 0
	}
	rows, err := strconv.Atoi(strings.Split(outStr, " ")[0])
	if err != nil {
		return 0, 0
	}
	return cols, rows
}
<|endoftext|>


<commit_before>package slinguist

const OtherLanguage = "Other"

var (
	extensionsByLanguage map[string][]string
	ignoredExtensions    = map[string]bool{
		".asc": true, ".cgi": true, ".fcgi": true, ".gml": true, ".fx": true,
		".vhost": true,
	}
	auxiliaryLanguages = map[string]bool{
		"Other": true, "XML": true, "YAML": true, "TOML": true, "INI": true,
		"JSON": true, "TeX": true, "Public Key": true, "AsciiDoc": true,
		"AGS Script": true, "VimL": true, "Diff": true, "CMake": true, "fish": true,
		"Awk": true, "Graphviz (DOT)": true, "Markdown": true, "desktop": true,
		"XSLT": true, "SQL": true, "RMarkdown": true, "IRC log": true,
		"reStructuredText": true, "Twig": true, "CSS": true, "Batchfile": true,
		"Text": true, "HTML+ERB": true, "HTML": true, "Gettext Catalog": true,
		"Smarty": true, "Raw token data": true,
	}
)

func init() {
	for l, _ := range ignoredExtensions {
		languagesByExtension[l] = []string{OtherLanguage}
	}

	extensionsByLanguage = reverseStringListMap(languagesByExtension)
}

// GetLanguageExtensions returns the different extensions being used by the
// language.
func GetLanguageExtensions(language string) []string {
	return extensionsByLanguage[language]
}

// GetLanguage return the Language for a given filename and file content.
func GetLanguage(filename string, content []byte) string {
	if lang, safe := GetLanguageByExtension(filename); safe {
		return lang
	}

	lang, _ := GetLanguageByContent(filename, content)
	return lang
}

func reverseStringListMap(i map[string][]string) (o map[string][]string) {
	o = map[string][]string{}
	for key, set := range i {
		for _, value := range set {
			o[value] = append(o[value], key)
		}
	}

	return
}
<commit_msg>Make ExtensionsByLanguage public<commit_after>package slinguist

const OtherLanguage = "Other"

var (
	ExtensionsByLanguage map[string][]string
	ignoredExtensions    = map[string]bool{
		".asc": true, ".cgi": true, ".fcgi": true, ".gml": true, ".fx": true,
		".vhost": true,
	}
	auxiliaryLanguages = map[string]bool{
		"Other": true, "XML": true, "YAML": true, "TOML": true, "INI": true,
		"JSON": true, "TeX": true, "Public Key": true, "AsciiDoc": true,
		"AGS Script": true, "VimL": true, "Diff": true, "CMake": true, "fish": true,
		"Awk": true, "Graphviz (DOT)": true, "Markdown": true, "desktop": true,
		"XSLT": true, "SQL": true, "RMarkdown": true, "IRC log": true,
		"reStructuredText": true, "Twig": true, "CSS": true, "Batchfile": true,
		"Text": true, "HTML+ERB": true, "HTML": true, "Gettext Catalog": true,
		"Smarty": true, "Raw token data": true,
	}
)

func init() {
	for l, _ := range ignoredExtensions {
		languagesByExtension[l] = []string{OtherLanguage}
	}

	ExtensionsByLanguage = reverseStringListMap(languagesByExtension)
}

// GetLanguageExtensions returns the different extensions being used by the
// language.
func GetLanguageExtensions(language string) []string {
	return ExtensionsByLanguage[language]
}

// GetLanguage return the Language for a given filename and file content.
func GetLanguage(filename string, content []byte) string {
	if lang, safe := GetLanguageByExtension(filename); safe {
		return lang
	}

	lang, _ := GetLanguageByContent(filename, content)
	return lang
}

func reverseStringListMap(i map[string][]string) (o map[string][]string) {
	o = map[string][]string{}
	for key, set := range i {
		for _, value := range set {
			o[value] = append(o[value], key)
		}
	}

	return
}
<|endoftext|>


<commit_before>package lexer

import (
  "github.com/Southern/scanner"
  "regexp"
  "strings"
)

var regex = map[string]map[string]scanner.Definition{
  "comments": map[string]scanner.Definition{
    // Hash comments
    "hash": scanner.Definition{regexp.MustCompile("^#(.+)"), "COMMENT"},

    // Single line comments
    "oneline": scanner.Definition{regexp.MustCompile("^\\/{2,}\\s*(.+)"), "COMMENT"},

    // Multi-line comments
    "multiline": scanner.Definition{regexp.MustCompile("^/\\*([^*]|[\r\n]|(\\*+([^*/]|[\r\n])))*\\*+/"), "COMMENT"},
  },

  "strings": map[string]scanner.Definition{
    // Double quote strings
    "double": scanner.Definition{regexp.MustCompile("^\"([^\"\\n]|\\\")+\""), "STRING"},

    // Single quote strings
    "single": scanner.Definition{regexp.MustCompile("^'([^'\\n]|\\')+'"), "STRING"},
  },

  "operators": map[string]scanner.Definition{
    // Common operators found in almost all languages
    "common": scanner.Definition{regexp.MustCompile(strings.Join([]string{
      "^(",
      strings.Join([]string{
        // ++, --, &&, ||
        "([+\\-&|]{2})",
        // !, !=, <, <<, <=, <<=, >, >>, >>>, >=, >>=, >>>=, ^, ^=, +, +=, -,
        // -=, %, %=, *, *=
        "((<{2}|>{2,3}|[!|&<>^+\\-=%/*])=?)",
      }, "|"),
      ")",
    }, "")), "OPERATOR"},
  },
}
<commit_msg>Strings shouldn't require text inside of them.<commit_after>package lexer

import (
  "github.com/Southern/scanner"
  "regexp"
  "strings"
)

var regex = map[string]map[string]scanner.Definition{
  "comments": map[string]scanner.Definition{
    // Hash comments
    "hash": scanner.Definition{regexp.MustCompile("^#(.+)"), "COMMENT"},

    // Single line comments
    "oneline": scanner.Definition{regexp.MustCompile("^\\/{2,}\\s*(.+)"), "COMMENT"},

    // Multi-line comments
    "multiline": scanner.Definition{regexp.MustCompile("^/\\*([^*]|[\r\n]|(\\*+([^*/]|[\r\n])))*\\*+/"), "COMMENT"},
  },

  "strings": map[string]scanner.Definition{
    // Double quote strings
    "double": scanner.Definition{regexp.MustCompile("^\"([^\"\\n]|\\\")*\""), "STRING"},

    // Single quote strings
    "single": scanner.Definition{regexp.MustCompile("^'([^'\\n]|\\')*'"), "STRING"},
  },

  "operators": map[string]scanner.Definition{
    // Common operators found in almost all languages
    "common": scanner.Definition{regexp.MustCompile(strings.Join([]string{
      "^(",
      strings.Join([]string{
        // ++, --, &&, ||
        "([+\\-&|]{2})",
        // !, !=, <, <<, <=, <<=, >, >>, >>>, >=, >>=, >>>=, ^, ^=, +, +=, -,
        // -=, %, %=, *, *=
        "((<{2}|>{2,3}|[!|&<>^+\\-=%/*])=?)",
      }, "|"),
      ")",
    }, "")), "OPERATOR"},
  },
}
<|endoftext|>


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

package tp

import (
	"context"
	"crypto/tls"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/henrylee2cn/goutil"
	"github.com/henrylee2cn/goutil/pool"
	"github.com/henrylee2cn/teleport/codec"
	"github.com/henrylee2cn/teleport/socket"
	"github.com/henrylee2cn/teleport/utils"
)

// Packet types
const (
	TypeUndefined byte = 0
	TypePull      byte = 1
	TypeReply     byte = 2 // reply to pull
	TypePush      byte = 3
)

// TypeText returns the packet type text.
// If the type is undefined returns 'Undefined'.
func TypeText(typ byte) string {
	switch typ {
	case TypePull:
		return "PULL"
	case TypeReply:
		return "REPLY"
	case TypePush:
		return "PUSH"
	default:
		return "Undefined"
	}
}

// Internal Framework Rerror code.
// Note: Recommended custom code is greater than 1000.
const (
	CodeUnknownError        = -1
	CodeConnClosed          = 102
	CodeWriteFailed         = 104
	CodeDialFailed          = 105
	CodeBadPacket           = 400
	CodeUnauthorized        = 401
	CodeNotFound            = 404
	CodePtypeNotAllowed     = 405
	CodeHandleTimeout       = 408
	CodeInternalServerError = 500
	CodeBadGateway          = 502

	// CodeConflict                      = 409
	// CodeUnsupportedTx                 = 410
	// CodeUnsupportedCodecType          = 415
	// CodeServiceUnavailable            = 503
	// CodeGatewayTimeout                = 504
	// CodeVariantAlsoNegotiates         = 506
	// CodeInsufficientStorage           = 507
	// CodeLoopDetected                  = 508
	// CodeNotExtended                   = 510
	// CodeNetworkAuthenticationRequired = 511
)

// CodeText returns the reply error code text.
// If the type is undefined returns 'Unknown Error'.
func CodeText(rerrCode int32) string {
	switch rerrCode {
	case CodeBadPacket:
		return "Bad Packet"
	case CodeUnauthorized:
		return "Unauthorized"
	case CodeDialFailed:
		return "Dial Failed"
	case CodeConnClosed:
		return "Connection Closed"
	case CodeWriteFailed:
		return "Write Failed"
	case CodeNotFound:
		return "Not Found"
	case CodeHandleTimeout:
		return "Handle Timeout"
	case CodePtypeNotAllowed:
		return "Packet Type Not Allowed"
	case CodeInternalServerError:
		return "Internal Server Error"
	case CodeBadGateway:
		return "Bad Gateway"
	case CodeUnknownError:
		fallthrough
	default:
		return "Unknown Error"
	}
}

// Internal Framework Rerror string.
var (
	rerrUnknownError        = NewRerror(CodeUnknownError, CodeText(CodeUnknownError), "")
	rerrDialFailed          = NewRerror(CodeDialFailed, CodeText(CodeDialFailed), "")
	rerrConnClosed          = NewRerror(CodeConnClosed, CodeText(CodeConnClosed), "")
	rerrWriteFailed         = NewRerror(CodeWriteFailed, CodeText(CodeWriteFailed), "")
	rerrBadPacket           = NewRerror(CodeBadPacket, CodeText(CodeBadPacket), "")
	rerrNotFound            = NewRerror(CodeNotFound, CodeText(CodeNotFound), "")
	rerrCodePtypeNotAllowed = NewRerror(CodePtypeNotAllowed, CodeText(CodePtypeNotAllowed), "")
	rerrHandleTimeout       = NewRerror(CodeHandleTimeout, CodeText(CodeHandleTimeout), "")
	rerrInternalServerError = NewRerror(CodeInternalServerError, CodeText(CodeInternalServerError), "")
)

// IsConnRerror determines whether the error is a connection error
func IsConnRerror(rerr *Rerror) bool {
	if rerr == nil {
		return false
	}
	if rerr.Code == CodeDialFailed || rerr.Code == CodeConnClosed {
		return true
	}
	return false
}

const (
	// MetaRerror reply error metadata key
	MetaRerror = "X-Reply-Error"
	// MetaRealId real ID metadata key
	MetaRealId = "X-Real-ID"
	// MetaRealIp real IP metadata key
	MetaRealIp = "X-Real-IP"
	// MetaAcceptBodyCodec the key of body codec that the sender wishes to accept
	MetaAcceptBodyCodec = "X-Accept-Body-Codec"
)

// WithRealId sets the real ID to metadata.
func WithRealId(id string) socket.PacketSetting {
	return socket.WithAddMeta(MetaRealId, id)
}

// WithRealIp sets the real IP to metadata.
func WithRealIp(ip string) socket.PacketSetting {
	return socket.WithAddMeta(MetaRealIp, ip)
}

// WithAcceptBodyCodec sets the body codec that the sender wishes to accept.
// Note: If the specified codec is invalid, the receiver will ignore the mate data.
func WithAcceptBodyCodec(bodyCodec byte) socket.PacketSetting {
	if bodyCodec == codec.NilCodecId {
		return func(*socket.Packet) {}
	}
	return socket.WithAddMeta(MetaAcceptBodyCodec, strconv.FormatUint(uint64(bodyCodec), 10))
}

// GetAcceptBodyCodec gets the body codec that the sender wishes to accept.
// Note: If the specified codec is invalid, the receiver will ignore the mate data.
func GetAcceptBodyCodec(meta *utils.Args) (byte, bool) {
	s := meta.Peek(MetaAcceptBodyCodec)
	if len(s) == 0 || len(s) > 3 {
		return 0, false
	}
	b, err := strconv.ParseUint(goutil.BytesToString(s), 10, 8)
	if err != nil {
		return 0, false
	}
	c := byte(b)
	return c, c != codec.NilCodecId
}

// WithContext sets the packet handling context.
//  func WithContext(ctx context.Context) socket.PacketSetting
var WithContext = socket.WithContext

// WithSeq sets the packet sequence.
//  func WithSeq(seq uint64) socket.PacketSetting
var WithSeq = socket.WithSeq

// WithPtype sets the packet type.
//  func WithPtype(ptype byte) socket.PacketSetting
var WithPtype = socket.WithPtype

// WithUri sets the packet URI string.
//  func WithUri(uri string) socket.PacketSetting
var WithUri = socket.WithUri

// WithUriObject sets the packet URI object.
//  func WithUriObject(uriObject *url.URL) socket.PacketSetting
var WithUriObject = socket.WithUriObject

// WithQuery sets the packet URI query parameter.
//  func WithQuery(key, value string) socket.PacketSetting
var WithQuery = socket.WithQuery

// WithAddMeta adds 'key=value' metadata argument.
// Multiple values for the same key may be added.
//  func WithAddMeta(key, value string) socket.PacketSetting
var WithAddMeta = socket.WithAddMeta

// WithSetMeta sets 'key=value' metadata argument.
//  func WithSetMeta(key, value string) socket.PacketSetting
var WithSetMeta = socket.WithSetMeta

// WithBodyCodec sets the body codec.
//  func WithBodyCodec(bodyCodec byte) socket.PacketSetting
var WithBodyCodec = socket.WithBodyCodec

// WithBody sets the body object.
//  func WithBody(body interface{}) socket.PacketSetting
var WithBody = socket.WithBody

// WithNewBody resets the function of geting body.
//  func WithNewBody(newBodyFunc socket.NewBodyFunc) socket.PacketSetting
var WithNewBody = socket.WithNewBody

// WithXferPipe sets transfer filter pipe.
//  func WithXferPipe(filterId ...byte) socket.PacketSetting
var WithXferPipe = socket.WithXferPipe

// GetPacket gets a *Packet form packet stack.
// Note:
//  newBodyFunc is only for reading form connection;
//  settings are only for writing to connection.
//  func GetPacket(settings ...socket.PacketSetting) *socket.Packet
var GetPacket = socket.GetPacket

// PutPacket puts a *socket.Packet to packet stack.
//  func PutPacket(p *socket.Packet)
var PutPacket = socket.PutPacket

var (
	_maxGoroutinesAmount      = (1024 * 1024 * 8) / 8 // max memory 8GB (8KB/goroutine)
	_maxGoroutineIdleDuration time.Duration
	_gopool                   = pool.NewGoPool(_maxGoroutinesAmount, _maxGoroutineIdleDuration)
)

// SetGopool set or reset go pool config.
// Note: Make sure to call it before calling NewPeer() and Go()
func SetGopool(maxGoroutinesAmount int, maxGoroutineIdleDuration time.Duration) {
	_maxGoroutinesAmount, _maxGoroutineIdleDuration := maxGoroutinesAmount, maxGoroutineIdleDuration
	if _gopool != nil {
		_gopool.Stop()
	}
	_gopool = pool.NewGoPool(_maxGoroutinesAmount, _maxGoroutineIdleDuration)
}

// Go similar to go func, but return false if insufficient resources.
func Go(fn func()) bool {
	if err := _gopool.Go(fn); err != nil {
		Warnf("%s", err.Error())
		return false
	}
	return true
}

// AnywayGo similar to go func, but concurrent resources are limited.
func AnywayGo(fn func()) {
TRYGO:
	if !Go(fn) {
		time.Sleep(time.Second)
		goto TRYGO
	}
}

var printPidOnce sync.Once

func doPrintPid() {
	printPidOnce.Do(func() {
		Printf("The current process PID: %d", os.Getpid())
	})
}

type fakePullCmd struct {
	output    *socket.Packet
	reply     interface{}
	rerr      *Rerror
	inputMeta *utils.Args
}

// NewFakePullCmd creates a fake PullCmd.
func NewFakePullCmd(uri string, args, reply interface{}, rerr *Rerror) PullCmd {
	return &fakePullCmd{
		output: socket.NewPacket(
			socket.WithPtype(TypePull),
			socket.WithUri(uri),
			socket.WithBody(args),
		),
		reply:     reply,
		rerr:      rerr,
		inputMeta: utils.AcquireArgs(),
	}
}

var closeChan = func() <-chan struct{} {
	ch := make(chan struct{})
	close(ch)
	return ch
}()

// Done returns the chan that indicates whether it has been completed.
func (f *fakePullCmd) Done() <-chan struct{} {
	return closeChan
}

// Output returns writed packet.
func (f *fakePullCmd) Output() *socket.Packet {
	return f.output
}

// Context carries a deadline, a cancelation signal, and other values across
// API boundaries.
func (f *fakePullCmd) Context() context.Context {
	return f.output.Context()
}

// Result returns the pull result.
func (f *fakePullCmd) Result() (interface{}, *Rerror) {
	return f.reply, f.rerr
}

// Rerror returns the pull error.
func (f *fakePullCmd) Rerror() *Rerror {
	return f.rerr
}

// InputBodyCodec gets the body codec type of the input packet.
func (f *fakePullCmd) InputBodyCodec() byte {
	return codec.NilCodecId
}

// InputMeta returns the header metadata of input packet.
func (f *fakePullCmd) InputMeta() *utils.Args {
	return f.inputMeta
}

// CostTime returns the pulled cost time.
// If PeerConfig.CountTime=false, always returns 0.
func (f *fakePullCmd) CostTime() time.Duration {
	return 0
}

// NewTlsConfigFromFile creates a new TLS config.
func NewTlsConfigFromFile(tlsCertFile, tlsKeyFile string) (*tls.Config, error) {
	cert, err := tls.LoadX509KeyPair(tlsCertFile, tlsKeyFile)
	if err != nil {
		return nil, err
	}
	return &tls.Config{
		Certificates:             []tls.Certificate{cert},
		NextProtos:               []string{"http/1.1", "h2"},
		PreferServerCipherSuites: true,
		CurvePreferences: []tls.CurveID{
			tls.CurveP256,
			tls.X25519,
		},
		MinVersion: tls.VersionTLS12,
		CipherSuites: []uint16{
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305,
			tls.TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_RSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_RSA_WITH_AES_128_GCM_SHA256,
		},
	}, nil
}
<commit_msg>closeChan -> closedChan<commit_after>// Copyright 2015-2017 HenryLee. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package tp

import (
	"context"
	"crypto/tls"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/henrylee2cn/goutil"
	"github.com/henrylee2cn/goutil/pool"
	"github.com/henrylee2cn/teleport/codec"
	"github.com/henrylee2cn/teleport/socket"
	"github.com/henrylee2cn/teleport/utils"
)

// Packet types
const (
	TypeUndefined byte = 0
	TypePull      byte = 1
	TypeReply     byte = 2 // reply to pull
	TypePush      byte = 3
)

// TypeText returns the packet type text.
// If the type is undefined returns 'Undefined'.
func TypeText(typ byte) string {
	switch typ {
	case TypePull:
		return "PULL"
	case TypeReply:
		return "REPLY"
	case TypePush:
		return "PUSH"
	default:
		return "Undefined"
	}
}

// Internal Framework Rerror code.
// Note: Recommended custom code is greater than 1000.
const (
	CodeUnknownError        = -1
	CodeConnClosed          = 102
	CodeWriteFailed         = 104
	CodeDialFailed          = 105
	CodeBadPacket           = 400
	CodeUnauthorized        = 401
	CodeNotFound            = 404
	CodePtypeNotAllowed     = 405
	CodeHandleTimeout       = 408
	CodeInternalServerError = 500
	CodeBadGateway          = 502

	// CodeConflict                      = 409
	// CodeUnsupportedTx                 = 410
	// CodeUnsupportedCodecType          = 415
	// CodeServiceUnavailable            = 503
	// CodeGatewayTimeout                = 504
	// CodeVariantAlsoNegotiates         = 506
	// CodeInsufficientStorage           = 507
	// CodeLoopDetected                  = 508
	// CodeNotExtended                   = 510
	// CodeNetworkAuthenticationRequired = 511
)

// CodeText returns the reply error code text.
// If the type is undefined returns 'Unknown Error'.
func CodeText(rerrCode int32) string {
	switch rerrCode {
	case CodeBadPacket:
		return "Bad Packet"
	case CodeUnauthorized:
		return "Unauthorized"
	case CodeDialFailed:
		return "Dial Failed"
	case CodeConnClosed:
		return "Connection Closed"
	case CodeWriteFailed:
		return "Write Failed"
	case CodeNotFound:
		return "Not Found"
	case CodeHandleTimeout:
		return "Handle Timeout"
	case CodePtypeNotAllowed:
		return "Packet Type Not Allowed"
	case CodeInternalServerError:
		return "Internal Server Error"
	case CodeBadGateway:
		return "Bad Gateway"
	case CodeUnknownError:
		fallthrough
	default:
		return "Unknown Error"
	}
}

// Internal Framework Rerror string.
var (
	rerrUnknownError        = NewRerror(CodeUnknownError, CodeText(CodeUnknownError), "")
	rerrDialFailed          = NewRerror(CodeDialFailed, CodeText(CodeDialFailed), "")
	rerrConnClosed          = NewRerror(CodeConnClosed, CodeText(CodeConnClosed), "")
	rerrWriteFailed         = NewRerror(CodeWriteFailed, CodeText(CodeWriteFailed), "")
	rerrBadPacket           = NewRerror(CodeBadPacket, CodeText(CodeBadPacket), "")
	rerrNotFound            = NewRerror(CodeNotFound, CodeText(CodeNotFound), "")
	rerrCodePtypeNotAllowed = NewRerror(CodePtypeNotAllowed, CodeText(CodePtypeNotAllowed), "")
	rerrHandleTimeout       = NewRerror(CodeHandleTimeout, CodeText(CodeHandleTimeout), "")
	rerrInternalServerError = NewRerror(CodeInternalServerError, CodeText(CodeInternalServerError), "")
)

// IsConnRerror determines whether the error is a connection error
func IsConnRerror(rerr *Rerror) bool {
	if rerr == nil {
		return false
	}
	if rerr.Code == CodeDialFailed || rerr.Code == CodeConnClosed {
		return true
	}
	return false
}

const (
	// MetaRerror reply error metadata key
	MetaRerror = "X-Reply-Error"
	// MetaRealId real ID metadata key
	MetaRealId = "X-Real-ID"
	// MetaRealIp real IP metadata key
	MetaRealIp = "X-Real-IP"
	// MetaAcceptBodyCodec the key of body codec that the sender wishes to accept
	MetaAcceptBodyCodec = "X-Accept-Body-Codec"
)

// WithRealId sets the real ID to metadata.
func WithRealId(id string) socket.PacketSetting {
	return socket.WithAddMeta(MetaRealId, id)
}

// WithRealIp sets the real IP to metadata.
func WithRealIp(ip string) socket.PacketSetting {
	return socket.WithAddMeta(MetaRealIp, ip)
}

// WithAcceptBodyCodec sets the body codec that the sender wishes to accept.
// Note: If the specified codec is invalid, the receiver will ignore the mate data.
func WithAcceptBodyCodec(bodyCodec byte) socket.PacketSetting {
	if bodyCodec == codec.NilCodecId {
		return func(*socket.Packet) {}
	}
	return socket.WithAddMeta(MetaAcceptBodyCodec, strconv.FormatUint(uint64(bodyCodec), 10))
}

// GetAcceptBodyCodec gets the body codec that the sender wishes to accept.
// Note: If the specified codec is invalid, the receiver will ignore the mate data.
func GetAcceptBodyCodec(meta *utils.Args) (byte, bool) {
	s := meta.Peek(MetaAcceptBodyCodec)
	if len(s) == 0 || len(s) > 3 {
		return 0, false
	}
	b, err := strconv.ParseUint(goutil.BytesToString(s), 10, 8)
	if err != nil {
		return 0, false
	}
	c := byte(b)
	return c, c != codec.NilCodecId
}

// WithContext sets the packet handling context.
//  func WithContext(ctx context.Context) socket.PacketSetting
var WithContext = socket.WithContext

// WithSeq sets the packet sequence.
//  func WithSeq(seq uint64) socket.PacketSetting
var WithSeq = socket.WithSeq

// WithPtype sets the packet type.
//  func WithPtype(ptype byte) socket.PacketSetting
var WithPtype = socket.WithPtype

// WithUri sets the packet URI string.
//  func WithUri(uri string) socket.PacketSetting
var WithUri = socket.WithUri

// WithUriObject sets the packet URI object.
//  func WithUriObject(uriObject *url.URL) socket.PacketSetting
var WithUriObject = socket.WithUriObject

// WithQuery sets the packet URI query parameter.
//  func WithQuery(key, value string) socket.PacketSetting
var WithQuery = socket.WithQuery

// WithAddMeta adds 'key=value' metadata argument.
// Multiple values for the same key may be added.
//  func WithAddMeta(key, value string) socket.PacketSetting
var WithAddMeta = socket.WithAddMeta

// WithSetMeta sets 'key=value' metadata argument.
//  func WithSetMeta(key, value string) socket.PacketSetting
var WithSetMeta = socket.WithSetMeta

// WithBodyCodec sets the body codec.
//  func WithBodyCodec(bodyCodec byte) socket.PacketSetting
var WithBodyCodec = socket.WithBodyCodec

// WithBody sets the body object.
//  func WithBody(body interface{}) socket.PacketSetting
var WithBody = socket.WithBody

// WithNewBody resets the function of geting body.
//  func WithNewBody(newBodyFunc socket.NewBodyFunc) socket.PacketSetting
var WithNewBody = socket.WithNewBody

// WithXferPipe sets transfer filter pipe.
//  func WithXferPipe(filterId ...byte) socket.PacketSetting
var WithXferPipe = socket.WithXferPipe

// GetPacket gets a *Packet form packet stack.
// Note:
//  newBodyFunc is only for reading form connection;
//  settings are only for writing to connection.
//  func GetPacket(settings ...socket.PacketSetting) *socket.Packet
var GetPacket = socket.GetPacket

// PutPacket puts a *socket.Packet to packet stack.
//  func PutPacket(p *socket.Packet)
var PutPacket = socket.PutPacket

var (
	_maxGoroutinesAmount      = (1024 * 1024 * 8) / 8 // max memory 8GB (8KB/goroutine)
	_maxGoroutineIdleDuration time.Duration
	_gopool                   = pool.NewGoPool(_maxGoroutinesAmount, _maxGoroutineIdleDuration)
)

// SetGopool set or reset go pool config.
// Note: Make sure to call it before calling NewPeer() and Go()
func SetGopool(maxGoroutinesAmount int, maxGoroutineIdleDuration time.Duration) {
	_maxGoroutinesAmount, _maxGoroutineIdleDuration := maxGoroutinesAmount, maxGoroutineIdleDuration
	if _gopool != nil {
		_gopool.Stop()
	}
	_gopool = pool.NewGoPool(_maxGoroutinesAmount, _maxGoroutineIdleDuration)
}

// Go similar to go func, but return false if insufficient resources.
func Go(fn func()) bool {
	if err := _gopool.Go(fn); err != nil {
		Warnf("%s", err.Error())
		return false
	}
	return true
}

// AnywayGo similar to go func, but concurrent resources are limited.
func AnywayGo(fn func()) {
TRYGO:
	if !Go(fn) {
		time.Sleep(time.Second)
		goto TRYGO
	}
}

var printPidOnce sync.Once

func doPrintPid() {
	printPidOnce.Do(func() {
		Printf("The current process PID: %d", os.Getpid())
	})
}

type fakePullCmd struct {
	output    *socket.Packet
	reply     interface{}
	rerr      *Rerror
	inputMeta *utils.Args
}

// NewFakePullCmd creates a fake PullCmd.
func NewFakePullCmd(uri string, args, reply interface{}, rerr *Rerror) PullCmd {
	return &fakePullCmd{
		output: socket.NewPacket(
			socket.WithPtype(TypePull),
			socket.WithUri(uri),
			socket.WithBody(args),
		),
		reply:     reply,
		rerr:      rerr,
		inputMeta: utils.AcquireArgs(),
	}
}

var closedChan = func() <-chan struct{} {
	ch := make(chan struct{})
	close(ch)
	return ch
}()

// Done returns the chan that indicates whether it has been completed.
func (f *fakePullCmd) Done() <-chan struct{} {
	return closedChan
}

// Output returns writed packet.
func (f *fakePullCmd) Output() *socket.Packet {
	return f.output
}

// Context carries a deadline, a cancelation signal, and other values across
// API boundaries.
func (f *fakePullCmd) Context() context.Context {
	return f.output.Context()
}

// Result returns the pull result.
func (f *fakePullCmd) Result() (interface{}, *Rerror) {
	return f.reply, f.rerr
}

// Rerror returns the pull error.
func (f *fakePullCmd) Rerror() *Rerror {
	return f.rerr
}

// InputBodyCodec gets the body codec type of the input packet.
func (f *fakePullCmd) InputBodyCodec() byte {
	return codec.NilCodecId
}

// InputMeta returns the header metadata of input packet.
func (f *fakePullCmd) InputMeta() *utils.Args {
	return f.inputMeta
}

// CostTime returns the pulled cost time.
// If PeerConfig.CountTime=false, always returns 0.
func (f *fakePullCmd) CostTime() time.Duration {
	return 0
}

// NewTlsConfigFromFile creates a new TLS config.
func NewTlsConfigFromFile(tlsCertFile, tlsKeyFile string) (*tls.Config, error) {
	cert, err := tls.LoadX509KeyPair(tlsCertFile, tlsKeyFile)
	if err != nil {
		return nil, err
	}
	return &tls.Config{
		Certificates:             []tls.Certificate{cert},
		NextProtos:               []string{"http/1.1", "h2"},
		PreferServerCipherSuites: true,
		CurvePreferences: []tls.CurveID{
			tls.CurveP256,
			tls.X25519,
		},
		MinVersion: tls.VersionTLS12,
		CipherSuites: []uint16{
			tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305,
			tls.TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305,
			tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
			tls.TLS_RSA_WITH_AES_256_GCM_SHA384,
			tls.TLS_RSA_WITH_AES_128_GCM_SHA256,
		},
	}, nil
}
<|endoftext|>


<commit_before>package fixtures

import (
	. "github.com/onsi/gingko"
	. "github.com/onsi/gomega"
	"testing"
)

func init() {
	Describe("Using ginkgo", func() {
		It("TestSomethingImportant", func() {
			whatever := &UselessStruct{}
			t.Fail(whatever.ImportantField != "SECRET_PASSWORD")
		})
	})
}
<commit_msg>Updates to goldmaster file<commit_after>package fixtures

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"testing"
)

func init() {
	Describe("Testing with ginkgo", func() {
		It("TestSomethingImportant", func() {

			whatever := &UselessStruct{}
			t.Fail(whatever.ImportantField != "SECRET_PASSWORD")
		})
	})
}
<|endoftext|>


<commit_before>package hypervisors

import (
	"encoding/gob"
	"flag"
	"fmt"
	"io"
	"net"
	"strings"
	"time"

	"github.com/Symantec/Dominator/fleetmanager/topology"
	"github.com/Symantec/Dominator/lib/constants"
	"github.com/Symantec/Dominator/lib/errors"
	"github.com/Symantec/Dominator/lib/log/prefixlogger"
	"github.com/Symantec/Dominator/lib/srpc"
	proto "github.com/Symantec/Dominator/proto/hypervisor"
)

var (
	desiredAddressPoolSize = flag.Uint("desiredAddressPoolSize", 16,
		"Desired number of free addresses to maintain in Hypervisor")
	maximumAddressPoolSize = flag.Uint("maximumAddressPoolSize", 24,
		"Maximum number of free addresses to maintain in Hypervisor")
	minimumAddressPoolSize = flag.Uint("minimumAddressPoolSize", 8,
		"Minimum number of free addresses to maintain in Hypervisor")
)

func checkPoolLimits() error {
	if *desiredAddressPoolSize < *minimumAddressPoolSize {
		return fmt.Errorf(
			"desiredAddressPoolSize: %d is less than minimumAddressPoolSize: %d",
			*desiredAddressPoolSize, *minimumAddressPoolSize)
	}
	if *desiredAddressPoolSize > *maximumAddressPoolSize {
		return fmt.Errorf(
			"desiredAddressPoolSize: %d is greater than maximumAddressPoolSize: %d",
			*desiredAddressPoolSize, *maximumAddressPoolSize)
	}
	return nil
}

func (m *Manager) updateTopology(t *topology.Topology) {
	machines, err := t.ListMachines("")
	if err != nil {
		m.logger.Println(err)
		return
	}
	deleteList := m.updateTopologyLocked(t, machines)
	for _, hypervisor := range deleteList {
		m.ipStorer.UnregisterHypervisor(hypervisor.machine.HostIpAddress)
		hypervisor.delete()
	}
}

func (m *Manager) updateTopologyLocked(t *topology.Topology,
	machines []*topology.Machine) []*hypervisorType {
	hypervisorsToDelete := make(map[string]struct{}, len(machines))
	m.mutex.Lock()
	defer m.mutex.Unlock()
	m.topology = t
	for hypervisorName := range m.hypervisors {
		hypervisorsToDelete[hypervisorName] = struct{}{}
	}
	for _, machine := range machines {
		delete(hypervisorsToDelete, machine.Hostname)
		if hypervisor, ok := m.hypervisors[machine.Hostname]; ok {
			go hypervisor.update(machine)
		} else {
			hypervisor := &hypervisorType{
				logger:  prefixlogger.New(machine.Hostname+": ", m.logger),
				machine: machine,
				vms:     make(map[string]*vmInfoType),
			}
			m.hypervisors[machine.Hostname] = hypervisor
			go m.manageHypervisorLoop(hypervisor, machine.Hostname)
		}
	}
	deleteList := make([]*hypervisorType, 0, len(hypervisorsToDelete))
	for hypervisorName := range hypervisorsToDelete {
		deleteList = append(deleteList, m.hypervisors[hypervisorName])
		delete(m.hypervisors, hypervisorName)
	}
	subnetsToDelete := make(map[string]struct{}, len(m.subnets))
	for gatewayIp := range m.subnets {
		subnetsToDelete[gatewayIp] = struct{}{}
	}
	t.Walk(func(directory *topology.Directory) error {
		for _, tSubnet := range directory.Subnets {
			gatewayIp := tSubnet.IpGateway.String()
			delete(subnetsToDelete, gatewayIp)
			m.subnets[gatewayIp] = m.makeSubnet(tSubnet)
		}
		return nil
	})
	for gatewayIp := range subnetsToDelete {
		delete(m.subnets, gatewayIp)
	}
	return deleteList
}

func (h *hypervisorType) update(machine *topology.Machine) {
	h.mutex.Lock()
	defer h.mutex.Unlock()
	h.machine = machine
}

func (h *hypervisorType) delete() {
	h.mutex.Lock()
	defer h.mutex.Unlock()
	h.deleteScheduled = true
	if h.conn != nil {
		h.conn.Close()
		h.conn = nil
	}
}

func (h *hypervisorType) isDeleteScheduled() bool {
	h.mutex.RLock()
	defer h.mutex.RUnlock()
	return h.deleteScheduled
}

func (m *Manager) manageHypervisorLoop(h *hypervisorType, hostname string) {
	for !h.isDeleteScheduled() {
		sleepTime := m.manageHypervisor(h, hostname)
		time.Sleep(sleepTime)
	}
}

func (m *Manager) manageHypervisor(h *hypervisorType,
	hostname string) time.Duration {
	failureProbeStatus := probeStatusBad
	defer func() {
		h.mutex.Lock()
		defer h.mutex.Unlock()
		h.probeStatus = failureProbeStatus
		if h.conn != nil {
			h.conn.Close()
			h.conn = nil
		}
	}()
	client, err := srpc.DialHTTP("tcp",
		fmt.Sprintf("%s:%d", hostname, constants.HypervisorPortNumber),
		time.Minute)
	if err != nil {
		h.logger.Debugln(1, err)
		if err == srpc.ErrorNoSrpcEndpoint {
			failureProbeStatus = probeStatusNoSrpc
		}
		return time.Second
	}
	defer client.Close()
	conn, err := client.Call("Hypervisor.GetUpdates")
	if err != nil {
		if strings.HasPrefix(err.Error(), "unknown service") {
			h.logger.Debugln(1, err)
			failureProbeStatus = probeStatusNoService
			return time.Minute
		} else {
			h.logger.Println(err)
		}
		return time.Second
	}
	h.mutex.Lock()
	h.probeStatus = probeStatusGood
	if h.deleteScheduled {
		h.mutex.Unlock()
		conn.Close()
		return 0
	}
	h.conn = conn
	h.mutex.Unlock()
	decoder := gob.NewDecoder(conn)
	h.logger.Debugln(0, "waiting for Update messages")
	firstUpdate := true
	for {
		var update proto.Update
		if err := decoder.Decode(&update); err != nil {
			if err != io.EOF {
				h.logger.Println(err)
			}
			return time.Second
		}
		m.processHypervisorUpdate(h, update, firstUpdate)
		firstUpdate = false
	}
}

func (m *Manager) processAddressPoolUpdates(h *hypervisorType,
	update proto.Update) {
	if update.HaveAddressPool {
		addresses := make([]net.IP, 0, len(update.AddressPool))
		for _, address := range update.AddressPool {
			addresses = append(addresses, address.IpAddress)
		}
		err := m.ipStorer.SetIPsForHypervisor(h.machine.HostIpAddress,
			addresses)
		if err != nil {
			h.logger.Println(err)
		}
	}
	if update.HaveNumFree {
		if update.NumFreeAddresses < *minimumAddressPoolSize {
			hypervisorAddress := fmt.Sprintf("%s:%d",
				h.machine.Hostname, constants.HypervisorPortNumber)
			m.mutex.Lock()
			defer m.mutex.Unlock()
			tSubnets, err := m.topology.GetSubnetsForMachine(h.machine.Hostname)
			if err != nil {
				h.logger.Println(err)
				return
			}
			freeIPs, err := m.findFreeIPs(tSubnets,
				*desiredAddressPoolSize-update.NumFreeAddresses)
			if err != nil {
				h.logger.Println(err)
				return
			}
			addresses := make([]proto.Address, 0, len(freeIPs))
			for _, ip := range freeIPs {
				addresses = append(addresses, proto.Address{
					IpAddress: ip,
					MacAddress: fmt.Sprintf("52:54:%02x:%02x:%02x:%02x",
						ip[0], ip[1], ip[2], ip[3]),
				})
			}
			client, err := srpc.DialHTTP("tcp", hypervisorAddress, time.Minute)
			if err != nil {
				h.logger.Println(err)
				return
			}
			defer client.Close()
			request := proto.AddAddressesToPoolRequest{addresses}
			var reply proto.AddAddressesToPoolResponse
			err = client.RequestReply("Hypervisor.AddAddressesToPool",
				request, &reply)
			if err == nil {
				err = errors.New(reply.Error)
			}
			if err != nil {
				h.logger.Println(err)
				return
			}
			m.ipStorer.AddIPsForHypervisor(h.machine.HostIpAddress, freeIPs)
			h.logger.Debugf(0, "replenished pool with %d addresses\n",
				len(addresses))
		} else if update.NumFreeAddresses > *maximumAddressPoolSize {
			hypervisorAddress := fmt.Sprintf("%s:%d",
				h.machine.Hostname, constants.HypervisorPortNumber)
			client, err := srpc.DialHTTP("tcp", hypervisorAddress, time.Minute)
			if err != nil {
				h.logger.Println(err)
				return
			}
			defer client.Close()
			request := proto.RemoveExcessAddressesFromPoolRequest{
				*desiredAddressPoolSize}
			var reply proto.RemoveExcessAddressesFromPoolResponse
			err = client.RequestReply(
				"Hypervisor.RemoveExcessAddressesFromPool",
				request, &reply)
			if err == nil {
				err = errors.New(reply.Error)
			}
			if err != nil {
				h.logger.Println(err)
				return
			}
			h.logger.Debugf(0, "removed %d excess addresses from pool\n",
				update.NumFreeAddresses-*desiredAddressPoolSize)
		}
	}
}

func (m *Manager) processHypervisorUpdate(h *hypervisorType,
	update proto.Update, firstUpdate bool) {
	if update.HaveSubnets { // Must do subnets first.
		m.processSubnetsUpdates(h, update.Subnets)
	}
	m.processAddressPoolUpdates(h, update)
	if update.HaveVMs {
		if firstUpdate {
			m.processInitialVMs(h, update.VMs)
		} else {
			m.processVmUpdates(h, update.VMs)
		}
	}
}

func (m *Manager) processInitialVMs(h *hypervisorType,
	vms map[string]*proto.VmInfo) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	vmsToDelete := make(map[string]struct{}, len(h.vms))
	for ipAddr := range h.vms {
		vmsToDelete[ipAddr] = struct{}{}
	}
	for ipAddr, protoVm := range vms {
		delete(vmsToDelete, ipAddr)
		if vm, ok := h.vms[ipAddr]; ok {
			vm.VmInfo = *protoVm
		} else {
			vm := &vmInfoType{ipAddr, *protoVm, h}
			h.vms[ipAddr] = vm
			m.vms[ipAddr] = vm
		}
	}
	for ipAddr := range vmsToDelete {
		delete(h.vms, ipAddr)
		delete(m.vms, ipAddr)
	}
}

func (m *Manager) processSubnetsUpdates(h *hypervisorType,
	haveSubnets []proto.Subnet) {
	haveSubnetsMap := make(map[string]int, len(haveSubnets))
	for index, subnet := range haveSubnets {
		haveSubnetsMap[subnet.Id] = index
	}
	t, err := m.getTopology()
	if err != nil {
		h.logger.Println(err)
		return
	}
	needSubnets, err := t.GetSubnetsForMachine(h.machine.Hostname)
	if err != nil {
		h.logger.Println(err)
		return
	}
	subnetsToAdd := make([]proto.Subnet, 0)
	for _, needSubnet := range needSubnets {
		if index, ok := haveSubnetsMap[needSubnet.Id]; ok {
			haveSubnet := haveSubnets[index]
			if !needSubnet.IpGateway.Equal(haveSubnet.IpGateway) {
				h.logger.Printf("subnet mismatch: %s: %s!=%s\n",
					needSubnet.Id, needSubnet.IpGateway, haveSubnet.IpGateway)
			}
		} else {
			subnetsToAdd = append(subnetsToAdd, needSubnet.Subnet)
		}
	}
	if len(subnetsToAdd) < 1 {
		return
	}
	client, err := srpc.DialHTTP("tcp",
		fmt.Sprintf("%s:%d",
			h.machine.Hostname, constants.HypervisorPortNumber),
		time.Minute)
	if err != nil {
		h.logger.Println(err)
		return
	}
	defer client.Close()
	request := proto.AddSubnetsRequest{Subnets: subnetsToAdd}
	var reply proto.AddSubnetsResponse
	err = client.RequestReply("Hypervisor.AddSubnets", request, &reply)
	if err == nil {
		err = errors.New(reply.Error)
	}
	if err != nil {
		h.logger.Println(err)
		return
	}
	h.logger.Debugf(0, "Added %d subnets\n", len(subnetsToAdd))
}

func (m *Manager) processVmUpdates(h *hypervisorType,
	updateVMs map[string]*proto.VmInfo) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	for ipAddr, protoVm := range updateVMs {
		if len(protoVm.Volumes) < 1 {
			delete(h.vms, ipAddr)
			delete(m.vms, ipAddr)
		} else {
			if vm, ok := h.vms[ipAddr]; ok {
				vm.VmInfo = *protoVm
			} else {
				vm := &vmInfoType{ipAddr, *protoVm, h}
				h.vms[ipAddr] = vm
				m.vms[ipAddr] = vm
			}
		}
	}
}
<commit_msg>Add -manageHypervisors option to fleet-manager.<commit_after>package hypervisors

import (
	"encoding/gob"
	"flag"
	"fmt"
	"io"
	"net"
	"strings"
	"time"

	"github.com/Symantec/Dominator/fleetmanager/topology"
	"github.com/Symantec/Dominator/lib/constants"
	"github.com/Symantec/Dominator/lib/errors"
	"github.com/Symantec/Dominator/lib/log/prefixlogger"
	"github.com/Symantec/Dominator/lib/srpc"
	proto "github.com/Symantec/Dominator/proto/hypervisor"
)

var (
	desiredAddressPoolSize = flag.Uint("desiredAddressPoolSize", 16,
		"Desired number of free addresses to maintain in Hypervisor")
	manageHypervisors = flag.Bool("manageHypervisors", false,
		"If true, manage hypervisors")
	maximumAddressPoolSize = flag.Uint("maximumAddressPoolSize", 24,
		"Maximum number of free addresses to maintain in Hypervisor")
	minimumAddressPoolSize = flag.Uint("minimumAddressPoolSize", 8,
		"Minimum number of free addresses to maintain in Hypervisor")
)

func checkPoolLimits() error {
	if *desiredAddressPoolSize < *minimumAddressPoolSize {
		return fmt.Errorf(
			"desiredAddressPoolSize: %d is less than minimumAddressPoolSize: %d",
			*desiredAddressPoolSize, *minimumAddressPoolSize)
	}
	if *desiredAddressPoolSize > *maximumAddressPoolSize {
		return fmt.Errorf(
			"desiredAddressPoolSize: %d is greater than maximumAddressPoolSize: %d",
			*desiredAddressPoolSize, *maximumAddressPoolSize)
	}
	return nil
}

func (m *Manager) updateTopology(t *topology.Topology) {
	machines, err := t.ListMachines("")
	if err != nil {
		m.logger.Println(err)
		return
	}
	deleteList := m.updateTopologyLocked(t, machines)
	for _, hypervisor := range deleteList {
		m.ipStorer.UnregisterHypervisor(hypervisor.machine.HostIpAddress)
		hypervisor.delete()
	}
}

func (m *Manager) updateTopologyLocked(t *topology.Topology,
	machines []*topology.Machine) []*hypervisorType {
	hypervisorsToDelete := make(map[string]struct{}, len(machines))
	m.mutex.Lock()
	defer m.mutex.Unlock()
	m.topology = t
	for hypervisorName := range m.hypervisors {
		hypervisorsToDelete[hypervisorName] = struct{}{}
	}
	for _, machine := range machines {
		delete(hypervisorsToDelete, machine.Hostname)
		if hypervisor, ok := m.hypervisors[machine.Hostname]; ok {
			go hypervisor.update(machine)
		} else {
			hypervisor := &hypervisorType{
				logger:  prefixlogger.New(machine.Hostname+": ", m.logger),
				machine: machine,
				vms:     make(map[string]*vmInfoType),
			}
			m.hypervisors[machine.Hostname] = hypervisor
			go m.manageHypervisorLoop(hypervisor, machine.Hostname)
		}
	}
	deleteList := make([]*hypervisorType, 0, len(hypervisorsToDelete))
	for hypervisorName := range hypervisorsToDelete {
		deleteList = append(deleteList, m.hypervisors[hypervisorName])
		delete(m.hypervisors, hypervisorName)
	}
	subnetsToDelete := make(map[string]struct{}, len(m.subnets))
	for gatewayIp := range m.subnets {
		subnetsToDelete[gatewayIp] = struct{}{}
	}
	t.Walk(func(directory *topology.Directory) error {
		for _, tSubnet := range directory.Subnets {
			gatewayIp := tSubnet.IpGateway.String()
			delete(subnetsToDelete, gatewayIp)
			m.subnets[gatewayIp] = m.makeSubnet(tSubnet)
		}
		return nil
	})
	for gatewayIp := range subnetsToDelete {
		delete(m.subnets, gatewayIp)
	}
	return deleteList
}

func (h *hypervisorType) update(machine *topology.Machine) {
	h.mutex.Lock()
	defer h.mutex.Unlock()
	h.machine = machine
}

func (h *hypervisorType) delete() {
	h.mutex.Lock()
	defer h.mutex.Unlock()
	h.deleteScheduled = true
	if h.conn != nil {
		h.conn.Close()
		h.conn = nil
	}
}

func (h *hypervisorType) isDeleteScheduled() bool {
	h.mutex.RLock()
	defer h.mutex.RUnlock()
	return h.deleteScheduled
}

func (m *Manager) manageHypervisorLoop(h *hypervisorType, hostname string) {
	for !h.isDeleteScheduled() {
		sleepTime := m.manageHypervisor(h, hostname)
		time.Sleep(sleepTime)
	}
}

func (m *Manager) manageHypervisor(h *hypervisorType,
	hostname string) time.Duration {
	failureProbeStatus := probeStatusBad
	defer func() {
		h.mutex.Lock()
		defer h.mutex.Unlock()
		h.probeStatus = failureProbeStatus
		if h.conn != nil {
			h.conn.Close()
			h.conn = nil
		}
	}()
	client, err := srpc.DialHTTP("tcp",
		fmt.Sprintf("%s:%d", hostname, constants.HypervisorPortNumber),
		time.Minute)
	if err != nil {
		h.logger.Debugln(1, err)
		if err == srpc.ErrorNoSrpcEndpoint {
			failureProbeStatus = probeStatusNoSrpc
		}
		return time.Second
	}
	defer client.Close()
	conn, err := client.Call("Hypervisor.GetUpdates")
	if err != nil {
		if strings.HasPrefix(err.Error(), "unknown service") {
			h.logger.Debugln(1, err)
			failureProbeStatus = probeStatusNoService
			return time.Minute
		} else {
			h.logger.Println(err)
		}
		return time.Second
	}
	h.mutex.Lock()
	h.probeStatus = probeStatusGood
	if h.deleteScheduled {
		h.mutex.Unlock()
		conn.Close()
		return 0
	}
	h.conn = conn
	h.mutex.Unlock()
	decoder := gob.NewDecoder(conn)
	h.logger.Debugln(0, "waiting for Update messages")
	firstUpdate := true
	for {
		var update proto.Update
		if err := decoder.Decode(&update); err != nil {
			if err != io.EOF {
				h.logger.Println(err)
			}
			return time.Second
		}
		m.processHypervisorUpdate(h, update, firstUpdate)
		firstUpdate = false
	}
}

func (m *Manager) processAddressPoolUpdates(h *hypervisorType,
	update proto.Update) {
	if update.HaveAddressPool {
		addresses := make([]net.IP, 0, len(update.AddressPool))
		for _, address := range update.AddressPool {
			addresses = append(addresses, address.IpAddress)
		}
		err := m.ipStorer.SetIPsForHypervisor(h.machine.HostIpAddress,
			addresses)
		if err != nil {
			h.logger.Println(err)
		}
	}
	if update.HaveNumFree {
		if update.NumFreeAddresses < *minimumAddressPoolSize {
			hypervisorAddress := fmt.Sprintf("%s:%d",
				h.machine.Hostname, constants.HypervisorPortNumber)
			m.mutex.Lock()
			defer m.mutex.Unlock()
			tSubnets, err := m.topology.GetSubnetsForMachine(h.machine.Hostname)
			if err != nil {
				h.logger.Println(err)
				return
			}
			freeIPs, err := m.findFreeIPs(tSubnets,
				*desiredAddressPoolSize-update.NumFreeAddresses)
			if err != nil {
				h.logger.Println(err)
				return
			}
			addresses := make([]proto.Address, 0, len(freeIPs))
			for _, ip := range freeIPs {
				addresses = append(addresses, proto.Address{
					IpAddress: ip,
					MacAddress: fmt.Sprintf("52:54:%02x:%02x:%02x:%02x",
						ip[0], ip[1], ip[2], ip[3]),
				})
			}
			client, err := srpc.DialHTTP("tcp", hypervisorAddress, time.Minute)
			if err != nil {
				h.logger.Println(err)
				return
			}
			defer client.Close()
			request := proto.AddAddressesToPoolRequest{addresses}
			var reply proto.AddAddressesToPoolResponse
			err = client.RequestReply("Hypervisor.AddAddressesToPool",
				request, &reply)
			if err == nil {
				err = errors.New(reply.Error)
			}
			if err != nil {
				h.logger.Println(err)
				return
			}
			m.ipStorer.AddIPsForHypervisor(h.machine.HostIpAddress, freeIPs)
			h.logger.Debugf(0, "replenished pool with %d addresses\n",
				len(addresses))
		} else if update.NumFreeAddresses > *maximumAddressPoolSize {
			hypervisorAddress := fmt.Sprintf("%s:%d",
				h.machine.Hostname, constants.HypervisorPortNumber)
			client, err := srpc.DialHTTP("tcp", hypervisorAddress, time.Minute)
			if err != nil {
				h.logger.Println(err)
				return
			}
			defer client.Close()
			request := proto.RemoveExcessAddressesFromPoolRequest{
				*desiredAddressPoolSize}
			var reply proto.RemoveExcessAddressesFromPoolResponse
			err = client.RequestReply(
				"Hypervisor.RemoveExcessAddressesFromPool",
				request, &reply)
			if err == nil {
				err = errors.New(reply.Error)
			}
			if err != nil {
				h.logger.Println(err)
				return
			}
			h.logger.Debugf(0, "removed %d excess addresses from pool\n",
				update.NumFreeAddresses-*desiredAddressPoolSize)
		}
	}
}

func (m *Manager) processHypervisorUpdate(h *hypervisorType,
	update proto.Update, firstUpdate bool) {
	if *manageHypervisors {
		if update.HaveSubnets { // Must do subnets first.
			m.processSubnetsUpdates(h, update.Subnets)
		}
		m.processAddressPoolUpdates(h, update)
	}
	if update.HaveVMs {
		if firstUpdate {
			m.processInitialVMs(h, update.VMs)
		} else {
			m.processVmUpdates(h, update.VMs)
		}
	}
}

func (m *Manager) processInitialVMs(h *hypervisorType,
	vms map[string]*proto.VmInfo) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	vmsToDelete := make(map[string]struct{}, len(h.vms))
	for ipAddr := range h.vms {
		vmsToDelete[ipAddr] = struct{}{}
	}
	for ipAddr, protoVm := range vms {
		delete(vmsToDelete, ipAddr)
		if vm, ok := h.vms[ipAddr]; ok {
			vm.VmInfo = *protoVm
		} else {
			vm := &vmInfoType{ipAddr, *protoVm, h}
			h.vms[ipAddr] = vm
			m.vms[ipAddr] = vm
		}
	}
	for ipAddr := range vmsToDelete {
		delete(h.vms, ipAddr)
		delete(m.vms, ipAddr)
	}
}

func (m *Manager) processSubnetsUpdates(h *hypervisorType,
	haveSubnets []proto.Subnet) {
	haveSubnetsMap := make(map[string]int, len(haveSubnets))
	for index, subnet := range haveSubnets {
		haveSubnetsMap[subnet.Id] = index
	}
	t, err := m.getTopology()
	if err != nil {
		h.logger.Println(err)
		return
	}
	needSubnets, err := t.GetSubnetsForMachine(h.machine.Hostname)
	if err != nil {
		h.logger.Println(err)
		return
	}
	subnetsToAdd := make([]proto.Subnet, 0)
	for _, needSubnet := range needSubnets {
		if index, ok := haveSubnetsMap[needSubnet.Id]; ok {
			haveSubnet := haveSubnets[index]
			if !needSubnet.IpGateway.Equal(haveSubnet.IpGateway) {
				h.logger.Printf("subnet mismatch: %s: %s!=%s\n",
					needSubnet.Id, needSubnet.IpGateway, haveSubnet.IpGateway)
			}
		} else {
			subnetsToAdd = append(subnetsToAdd, needSubnet.Subnet)
		}
	}
	if len(subnetsToAdd) < 1 {
		return
	}
	client, err := srpc.DialHTTP("tcp",
		fmt.Sprintf("%s:%d",
			h.machine.Hostname, constants.HypervisorPortNumber),
		time.Minute)
	if err != nil {
		h.logger.Println(err)
		return
	}
	defer client.Close()
	request := proto.AddSubnetsRequest{Subnets: subnetsToAdd}
	var reply proto.AddSubnetsResponse
	err = client.RequestReply("Hypervisor.AddSubnets", request, &reply)
	if err == nil {
		err = errors.New(reply.Error)
	}
	if err != nil {
		h.logger.Println(err)
		return
	}
	h.logger.Debugf(0, "Added %d subnets\n", len(subnetsToAdd))
}

func (m *Manager) processVmUpdates(h *hypervisorType,
	updateVMs map[string]*proto.VmInfo) {
	m.mutex.Lock()
	defer m.mutex.Unlock()
	for ipAddr, protoVm := range updateVMs {
		if len(protoVm.Volumes) < 1 {
			delete(h.vms, ipAddr)
			delete(m.vms, ipAddr)
		} else {
			if vm, ok := h.vms[ipAddr]; ok {
				vm.VmInfo = *protoVm
			} else {
				vm := &vmInfoType{ipAddr, *protoVm, h}
				h.vms[ipAddr] = vm
				m.vms[ipAddr] = vm
			}
		}
	}
}
<|endoftext|>


<commit_before>package g

import (
	"encoding/json"
	log "github.com/Sirupsen/logrus"
	"sync"

	"github.com/Cepave/open-falcon-backend/modules/query/logger"
	"github.com/toolkits/file"
)

type HttpConfig struct {
	Enabled bool   `json:"enabled"`
	Listen  string `json:"listen"`
}

type GinHttpConfig struct {
	Enabled bool   `json:"enabled"`
	Listen  string `json:"listen"`
}

type GraphConfig struct {
	ConnTimeout int32             `json:"connTimeout"`
	CallTimeout int32             `json:"callTimeout"`
	MaxConns    int32             `json:"maxConns"`
	MaxIdle     int32             `json:"maxIdle"`
	Replicas    int32             `json:"replicas"`
	Cluster     map[string]string `json:"cluster"`
}

type ApiConfig struct {
	Name      string `json:"name"`
	Token     string `json:"token"`
	Contact   string `json:"contact"`
	Event     string `json:"event"`
	Map       string `json:"map"`
	Geo       string `json:"geo"`
	Uplink    string `json:"uplink"`
	Query     string `json:"query"`
	Dashboard string `json:"dashboard"`
	Max       int    `json:"max"`
}

type DbConfig struct {
	Addr string `json:"addr"`
	Idle int    `json:"idle"`
	Max  int    `json:"max"`
}

type NqmLogConfig struct {
	JsonrpcUrl string `json:"jsonrpcUrl"`
}

type NqmConfig struct {
	Addr string `json:"addr"`
	Idle int    `json:"idle"`
	Max  int    `json:"max"`
}

type GrpcConfig struct {
	Enabled bool `json:"enabled"`
	Port    int  `json:"port"`
}
type GraphDB struct {
	Addr  string `json:"addr"`
	Idle  int    `json:"idle"`
	Max   int    `json:"max"`
	Limit int    `json:"limit"`
}

type GlobalConfig struct {
	Debug   bool           `json:"debug"`
	Http    *HttpConfig    `json:"http"`
	Graph   *GraphConfig   `json:"graph"`
	Api     *ApiConfig     `json:"api"`
	Db      *DbConfig      `json:"db"`
	NqmLog  *NqmLogConfig  `json:nqmlog`
	Nqm     *NqmConfig     `json:"nqm"`
	Grpc    *GrpcConfig    `json:"grpc"`
	GinHttp *GinHttpConfig `json:"gin_http"`
	GraphDB *GraphDB       `json:"graphdb"`
}

var (
	ConfigFile string
	config     *GlobalConfig
	configLock = new(sync.RWMutex)
)

// Gets the configuration
func Config() *GlobalConfig {
	configLock.RLock()
	defer configLock.RUnlock()
	return config
}

// Sets the config directly
func SetConfig(newConfig *GlobalConfig) {
	configLock.RLock()
	defer configLock.RUnlock()
	config = newConfig
}

func ParseConfig(cfg string) {
	if cfg == "" {
		log.Fatalln("config file not specified: use -c $filename")
	}

	if !file.IsExist(cfg) {
		log.Fatalln("config file specified not found:", cfg)
	}

	ConfigFile = cfg

	configContent, err := file.ToTrimString(cfg)
	if err != nil {
		log.Fatalln("read config file", cfg, "error:", err.Error())
	}

	var c GlobalConfig
	err = json.Unmarshal([]byte(configContent), &c)
	if err != nil {
		log.Fatalln("parse config file", cfg, "error:", err.Error())
	}

	SetConfig(&c)

	logger.InitLogger(c.Debug)
	log.Println("g.ParseConfig ok, file", cfg)
}
<commit_msg>[OWL-743][query] add Local string<commit_after>package g

import (
	"encoding/json"
	log "github.com/Sirupsen/logrus"
	"sync"

	"github.com/Cepave/open-falcon-backend/modules/query/logger"
	"github.com/toolkits/file"
)

type HttpConfig struct {
	Enabled bool   `json:"enabled"`
	Listen  string `json:"listen"`
}

type GinHttpConfig struct {
	Enabled bool   `json:"enabled"`
	Listen  string `json:"listen"`
}

type GraphConfig struct {
	ConnTimeout int32             `json:"connTimeout"`
	CallTimeout int32             `json:"callTimeout"`
	MaxConns    int32             `json:"maxConns"`
	MaxIdle     int32             `json:"maxIdle"`
	Replicas    int32             `json:"replicas"`
	Cluster     map[string]string `json:"cluster"`
}

type ApiConfig struct {
	Name      string `json:"name"`
	Token     string `json:"token"`
	Contact   string `json:"contact"`
	Event     string `json:"event"`
	Map       string `json:"map"`
	Geo       string `json:"geo"`
	Uplink    string `json:"uplink"`
	Query     string `json:"query"`
	Dashboard string `json:"dashboard"`
	Max       int    `json:"max"`
}

type DbConfig struct {
	Addr string `json:"addr"`
	Idle int    `json:"idle"`
	Max  int    `json:"max"`
}

type NqmLogConfig struct {
	JsonrpcUrl string `json:"jsonrpcUrl"`
}

type NqmConfig struct {
	Addr string `json:"addr"`
	Idle int    `json:"idle"`
	Max  int    `json:"max"`
}

type GrpcConfig struct {
	Enabled bool `json:"enabled"`
	Port    int  `json:"port"`
}
type GraphDB struct {
	Addr  string `json:"addr"`
	Idle  int    `json:"idle"`
	Max   int    `json:"max"`
	Limit int    `json:"limit"`
}

type GlobalConfig struct {
	Debug   bool           `json:"debug"`
	Http    *HttpConfig    `json:"http"`
	Graph   *GraphConfig   `json:"graph"`
	Api     *ApiConfig     `json:"api"`
	Db      *DbConfig      `json:"db"`
	Local   string         `json:"local"`
	NqmLog  *NqmLogConfig  `json:nqmlog`
	Nqm     *NqmConfig     `json:"nqm"`
	Grpc    *GrpcConfig    `json:"grpc"`
	GinHttp *GinHttpConfig `json:"gin_http"`
	GraphDB *GraphDB       `json:"graphdb"`
}

var (
	ConfigFile string
	config     *GlobalConfig
	configLock = new(sync.RWMutex)
)

// Gets the configuration
func Config() *GlobalConfig {
	configLock.RLock()
	defer configLock.RUnlock()
	return config
}

// Sets the config directly
func SetConfig(newConfig *GlobalConfig) {
	configLock.RLock()
	defer configLock.RUnlock()
	config = newConfig
}

func ParseConfig(cfg string) {
	if cfg == "" {
		log.Fatalln("config file not specified: use -c $filename")
	}

	if !file.IsExist(cfg) {
		log.Fatalln("config file specified not found:", cfg)
	}

	ConfigFile = cfg

	configContent, err := file.ToTrimString(cfg)
	if err != nil {
		log.Fatalln("read config file", cfg, "error:", err.Error())
	}

	var c GlobalConfig
	err = json.Unmarshal([]byte(configContent), &c)
	if err != nil {
		log.Fatalln("parse config file", cfg, "error:", err.Error())
	}

	SetConfig(&c)

	logger.InitLogger(c.Debug)
	log.Println("g.ParseConfig ok, file", cfg)
}
<|endoftext|>