<commit_before>package main

import (
	"fmt"
	"io"
	"math"
	"strings"

	"code.google.com/p/go-sqlite/go1/sqlite3"
	"github.com/rwcarlsen/cyan/query"
)

// The number of sql commands to buffer before dumping to the output database.
const DumpFreq = 100000

var (
	preExecStmts = []string{
		"DROP TABLE IF EXISTS Inventories;",
		"DROP TABLE IF EXISTS TimeList;",
		"CREATE TABLE Inventories (SimID TEXT,ResID INTEGER,AgentID INTEGER,StartTime INTEGER,EndTime INTEGER,StateID INTEGER,Quantity REAL);",
		"CREATE TABLE TimeList AS SELECT DISTINCT Time FROM Transactions;",
		query.Index("TimeList", "Time"),
		query.Index("Resources", "SimID", "ID", "StateID"),
		query.Index("Compositions", "SimID", "ID", "IsoID"),
		query.Index("Transactions", "ID"),
		query.Index("ResCreators", "SimID", "ResID"),
		query.Index("Agents", "SimID", "Prototype"),
	}
	postExecStmts = []string{
		query.Index("Inventories", "SimID", "AgentID"),
		query.Index("Inventories", "SimID", "ResID", "StartTime"),
		"ANALYZE;",
	}
	dumpSql    = "INSERT INTO Inventories VALUES (?,?,?,?,?,?,?);"
	resSqlHead = "SELECT ID,TimeCreated,StateID,Quantity FROM "
	resSqlTail = " WHERE Parent1 = ? OR Parent2 = ?;"

	ownerSql = `SELECT tr.ReceiverID, tr.Time FROM Transactions AS tr
				  WHERE tr.ResourceID = ? AND tr.SimID = ?
				  ORDER BY tr.Time ASC;`
	rootsSql = `SELECT res.ID,res.TimeCreated,rc.ModelID,res.StateID,Quantity FROM Resources AS res
				  INNER JOIN ResCreators AS rc ON res.ID = rc.ResID
				  WHERE res.SimID = ? AND rc.SimID = ?;`
)

// Prepare creates necessary indexes and tables required for efficient
// calculation of cyclus simulation inventory information.  Should be called
// once before walking begins.
func Prepare(conn *sqlite3.Conn) (err error) {
	fmt.Println("Creating indexes and inventory table...")
	for _, sql := range preExecStmts {
		if err := conn.Exec(sql); err != nil {
			fmt.Println("    ", err)
		}
	}
	return nil
}

// Finish should be called for a cyclus database after all walkers have
// completed processing inventory data. It creates final indexes and other
// finishing tasks.
func Finish(conn *sqlite3.Conn) (err error) {
	fmt.Println("Creating inventory indexes...")
	for _, sql := range postExecStmts {
		if err := conn.Exec(sql); err != nil {
			return err
		}
	}
	return nil
}

type Node struct {
	ResId     int
	OwnerId   int
	StartTime int
	EndTime   int
	StateId   int
	Quantity  float64
}

// Context encapsulates the logic for building a fast, queryable inventories
// table for a specific simulation from raw cyclus output database.
type Context struct {
	*sqlite3.Conn
	// Simid is the cyclus simulation id targeted by this context.  Must be
	// set.
	Simid       string
	mappednodes map[int32]struct{}
	tmpResTbl   string
	tmpResStmt  *sqlite3.Stmt
	dumpStmt    *sqlite3.Stmt
	ownerStmt   *sqlite3.Stmt
	resCount    int
	nodes       []*Node
}

func NewContext(conn *sqlite3.Conn, simid string, history chan string) *Context {
	return &Context{
		Conn:  conn,
		Simid: simid,
	}
}

func (c *Context) init() {
	c.nodes = make([]*Node, 0, 10000)
	c.mappednodes = map[int32]struct{}{}

	// create temp res table without simid
	fmt.Println("Creating temporary resource table...")
	c.tmpResTbl = "tmp_restbl_" + strings.Replace(c.Simid, "-", "_", -1)
	err := c.Exec("DROP TABLE IF EXISTS " + c.tmpResTbl)
	panicif(err)

	sql := "CREATE TABLE " + c.tmpResTbl + " AS SELECT ID,TimeCreated,Parent1,Parent2,StateID,Quantity FROM Resources WHERE SimID = ?;"
	err = c.Exec(sql, c.Simid)
	panicif(err)

	fmt.Println("Indexing temporary resource table...")
	err = c.Exec(query.Index(c.tmpResTbl, "Parent1"))
	panicif(err)

	err = c.Exec(query.Index(c.tmpResTbl, "Parent2"))
	panicif(err)

	// create prepared statements
	c.tmpResStmt, err = c.Prepare(resSqlHead + c.tmpResTbl + resSqlTail)
	panicif(err)

	c.dumpStmt, err = c.Prepare(dumpSql)
	panicif(err)

	c.ownerStmt, err = c.Prepare(ownerSql)
	panicif(err)
}

// WalkAll constructs the inventories table in the cyclus database alongside
// other tables. Creates several indexes in the process.  Finish should be
// called on the database connection after all simulation id's have been
// walked.
func (c *Context) WalkAll() (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("%v", r)
		}
	}()

	fmt.Printf("--- Building inventories for simid %v ---\n", c.Simid)
	c.init()

	fmt.Println("Retrieving root resource nodes...")
	roots := c.getRoots()

	fmt.Printf("Found %v root nodes\n", len(roots))
	for i, n := range roots {
		fmt.Printf("    Processing root %d...\n", i)
		c.walkDown(n)
	}

	fmt.Println("Dropping temporary resource table...")
	err = c.Exec("DROP TABLE " + c.tmpResTbl)
	panicif(err)

	c.dumpNodes()

	return nil
}

func (c *Context) getRoots() (roots []*Node) {
	sql := "SELECT COUNT(*) FROM ResCreators WHERE SimID = ?"
	stmt, err := c.Query(sql, c.Simid)
	panicif(err)

	n := 0
	err = stmt.Scan(&n)
	panicif(err)
	stmt.Reset()

	roots = make([]*Node, 0, n)
	for stmt, err = c.Query(rootsSql, c.Simid, c.Simid); err == nil; err = stmt.Next() {
		node := &Node{EndTime: math.MaxInt32}
		err := stmt.Scan(&node.ResId, &node.StartTime, &node.OwnerId, &node.StateId, &node.Quantity)
		panicif(err)

		roots = append(roots, node)
	}
	if err != io.EOF {
		panic(err)
	}
	return roots
}

func (c *Context) walkDown(node *Node) {
	if _, ok := c.mappednodes[int32(node.ResId)]; ok {
		return
	}
	c.mappednodes[int32(node.ResId)] = struct{}{}

	// dump if necessary
	c.resCount++
	if c.resCount%DumpFreq == 0 {
		c.dumpNodes()
	}

	// find resource's children
	kids := make([]*Node, 0, 2)
	err := c.tmpResStmt.Query(node.ResId, node.ResId)
	for ; err == nil; err = c.tmpResStmt.Next() {
		child := &Node{EndTime: math.MaxInt32}
		err := c.tmpResStmt.Scan(&child.ResId, &child.StartTime, &child.StateId, &child.Quantity)
		panicif(err)
		node.EndTime = child.StartTime
		kids = append(kids, child)
	}
	if err != io.EOF {
		panic(err)
	}

	// find resources owner changes (that occurred before children)
	owners, times := c.getNewOwners(node.ResId)

	childOwner := node.OwnerId
	if len(owners) > 0 {
		node.EndTime = times[0]
		childOwner = owners[len(owners)-1]

		lastend := math.MaxInt32
		if len(kids) > 0 {
			lastend = kids[0].StartTime
		}
		times = append(times, lastend)
		for i := range owners {
			n := &Node{ResId: node.ResId,
				OwnerId:   owners[i],
				StartTime: times[i],
				EndTime:   times[i+1],
				StateId:   node.StateId,
				Quantity:  node.Quantity,
			}
			c.nodes = append(c.nodes, n)
		}
	}

	c.nodes = append(c.nodes, node)

	// walk down resource's children
	for _, child := range kids {
		child.OwnerId = childOwner
		c.walkDown(child)
	}
}

func (c *Context) getNewOwners(id int) (owners, times []int) {
	var owner, t int
	err := c.ownerStmt.Query(id, c.Simid)
	for ; err == nil; err = c.ownerStmt.Next() {
		err := c.ownerStmt.Scan(&owner, &t)
		panicif(err)

		if id == owner {
			continue
		}
		owners = append(owners, owner)
		times = append(times, t)
	}
	if err != io.EOF {
		panic(err)
	}
	return owners, times
}

func (c *Context) dumpNodes() {
	fmt.Printf("    Dumping inventories (%d resources done)...\n", c.resCount)
	err := c.Exec("BEGIN TRANSACTION;")
	panicif(err)

	for _, n := range c.nodes {
		err = c.dumpStmt.Exec(c.Simid, n.ResId, n.OwnerId, n.StartTime, n.EndTime, n.StateId, n.Quantity)
		panicif(err)
	}
	err = c.Exec("END TRANSACTION;")
	panicif(err)

	c.nodes = c.nodes[:0]
}
<commit_msg>fixed missing index that caused slow processing<commit_after>package main

import (
	"fmt"
	"io"
	"math"
	"strings"

	"code.google.com/p/go-sqlite/go1/sqlite3"
	"github.com/rwcarlsen/cyan/query"
)

// The number of sql commands to buffer before dumping to the output database.
const DumpFreq = 100000

var (
	preExecStmts = []string{
		"DROP TABLE IF EXISTS Inventories;",
		"DROP TABLE IF EXISTS TimeList;",
		"CREATE TABLE Inventories (SimID TEXT,ResID INTEGER,AgentID INTEGER,StartTime INTEGER,EndTime INTEGER,StateID INTEGER,Quantity REAL);",
		"CREATE TABLE TimeList AS SELECT DISTINCT Time FROM Transactions;",
		query.Index("TimeList", "Time"),
		query.Index("Resources", "SimID", "ID", "StateID"),
		query.Index("Compositions", "SimID", "ID", "IsoID"),
		query.Index("Transactions", "SimID", "ResourceID"),
		query.Index("Transactions", "ID"),
		query.Index("ResCreators", "SimID", "ResID"),
		query.Index("Agents", "SimID", "Prototype"),
	}
	postExecStmts = []string{
		query.Index("Inventories", "SimID", "AgentID"),
		query.Index("Inventories", "SimID", "ResID", "StartTime"),
		"ANALYZE;",
	}
	dumpSql    = "INSERT INTO Inventories VALUES (?,?,?,?,?,?,?);"
	resSqlHead = "SELECT ID,TimeCreated,StateID,Quantity FROM "
	resSqlTail = " WHERE Parent1 = ? OR Parent2 = ?;"

	ownerSql = `SELECT tr.ReceiverID, tr.Time FROM Transactions AS tr
				  WHERE tr.ResourceID = ? AND tr.SimID = ?
				  ORDER BY tr.Time ASC;`
	rootsSql = `SELECT res.ID,res.TimeCreated,rc.ModelID,res.StateID,Quantity FROM Resources AS res
				  INNER JOIN ResCreators AS rc ON res.ID = rc.ResID
				  WHERE res.SimID = ? AND rc.SimID = ?;`
)

// Prepare creates necessary indexes and tables required for efficient
// calculation of cyclus simulation inventory information.  Should be called
// once before walking begins.
func Prepare(conn *sqlite3.Conn) (err error) {
	fmt.Println("Creating indexes and inventory table...")
	for _, sql := range preExecStmts {
		if err := conn.Exec(sql); err != nil {
			fmt.Println("    ", err)
		}
	}
	return nil
}

// Finish should be called for a cyclus database after all walkers have
// completed processing inventory data. It creates final indexes and other
// finishing tasks.
func Finish(conn *sqlite3.Conn) (err error) {
	fmt.Println("Creating inventory indexes...")
	for _, sql := range postExecStmts {
		if err := conn.Exec(sql); err != nil {
			return err
		}
	}
	return nil
}

type Node struct {
	ResId     int
	OwnerId   int
	StartTime int
	EndTime   int
	StateId   int
	Quantity  float64
}

// Context encapsulates the logic for building a fast, queryable inventories
// table for a specific simulation from raw cyclus output database.
type Context struct {
	*sqlite3.Conn
	// Simid is the cyclus simulation id targeted by this context.  Must be
	// set.
	Simid       string
	mappednodes map[int32]struct{}
	tmpResTbl   string
	tmpResStmt  *sqlite3.Stmt
	dumpStmt    *sqlite3.Stmt
	ownerStmt   *sqlite3.Stmt
	resCount    int
	nodes       []*Node
}

func NewContext(conn *sqlite3.Conn, simid string, history chan string) *Context {
	return &Context{
		Conn:  conn,
		Simid: simid,
	}
}

func (c *Context) init() {
	c.nodes = make([]*Node, 0, 10000)
	c.mappednodes = map[int32]struct{}{}

	// create temp res table without simid
	fmt.Println("Creating temporary resource table...")
	c.tmpResTbl = "tmp_restbl_" + strings.Replace(c.Simid, "-", "_", -1)
	err := c.Exec("DROP TABLE IF EXISTS " + c.tmpResTbl)
	panicif(err)

	sql := "CREATE TABLE " + c.tmpResTbl + " AS SELECT ID,TimeCreated,Parent1,Parent2,StateID,Quantity FROM Resources WHERE SimID = ?;"
	err = c.Exec(sql, c.Simid)
	panicif(err)

	fmt.Println("Indexing temporary resource table...")
	err = c.Exec(query.Index(c.tmpResTbl, "Parent1"))
	panicif(err)

	err = c.Exec(query.Index(c.tmpResTbl, "Parent2"))
	panicif(err)

	// create prepared statements
	c.tmpResStmt, err = c.Prepare(resSqlHead + c.tmpResTbl + resSqlTail)
	panicif(err)

	c.dumpStmt, err = c.Prepare(dumpSql)
	panicif(err)

	c.ownerStmt, err = c.Prepare(ownerSql)
	panicif(err)
}

// WalkAll constructs the inventories table in the cyclus database alongside
// other tables. Creates several indexes in the process.  Finish should be
// called on the database connection after all simulation id's have been
// walked.
func (c *Context) WalkAll() (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("%v", r)
		}
	}()

	fmt.Printf("--- Building inventories for simid %v ---\n", c.Simid)
	c.init()

	fmt.Println("Retrieving root resource nodes...")
	roots := c.getRoots()

	fmt.Printf("Found %v root nodes\n", len(roots))
	for i, n := range roots {
		fmt.Printf("    Processing root %d...\n", i)
		c.walkDown(n)
	}

	fmt.Println("Dropping temporary resource table...")
	err = c.Exec("DROP TABLE " + c.tmpResTbl)
	panicif(err)

	c.dumpNodes()

	return nil
}

func (c *Context) getRoots() (roots []*Node) {
	sql := "SELECT COUNT(*) FROM ResCreators WHERE SimID = ?"
	stmt, err := c.Query(sql, c.Simid)
	panicif(err)

	n := 0
	err = stmt.Scan(&n)
	panicif(err)
	stmt.Reset()

	roots = make([]*Node, 0, n)
	for stmt, err = c.Query(rootsSql, c.Simid, c.Simid); err == nil; err = stmt.Next() {
		node := &Node{EndTime: math.MaxInt32}
		err := stmt.Scan(&node.ResId, &node.StartTime, &node.OwnerId, &node.StateId, &node.Quantity)
		panicif(err)

		roots = append(roots, node)
	}
	if err != io.EOF {
		panic(err)
	}
	return roots
}

func (c *Context) walkDown(node *Node) {
	if _, ok := c.mappednodes[int32(node.ResId)]; ok {
		return
	}
	c.mappednodes[int32(node.ResId)] = struct{}{}

	// dump if necessary
	c.resCount++
	if c.resCount%DumpFreq == 0 {
		c.dumpNodes()
	}

	// find resource's children
	kids := make([]*Node, 0, 2)
	err := c.tmpResStmt.Query(node.ResId, node.ResId)
	for ; err == nil; err = c.tmpResStmt.Next() {
		child := &Node{EndTime: math.MaxInt32}
		err := c.tmpResStmt.Scan(&child.ResId, &child.StartTime, &child.StateId, &child.Quantity)
		panicif(err)
		node.EndTime = child.StartTime
		kids = append(kids, child)
	}
	if err != io.EOF {
		panic(err)
	}

	// find resources owner changes (that occurred before children)
	owners, times := c.getNewOwners(node.ResId)

	childOwner := node.OwnerId
	if len(owners) > 0 {
		node.EndTime = times[0]
		childOwner = owners[len(owners)-1]

		lastend := math.MaxInt32
		if len(kids) > 0 {
			lastend = kids[0].StartTime
		}
		times = append(times, lastend)
		for i := range owners {
			n := &Node{ResId: node.ResId,
				OwnerId:   owners[i],
				StartTime: times[i],
				EndTime:   times[i+1],
				StateId:   node.StateId,
				Quantity:  node.Quantity,
			}
			c.nodes = append(c.nodes, n)
		}
	}

	c.nodes = append(c.nodes, node)

	// walk down resource's children
	for _, child := range kids {
		child.OwnerId = childOwner
		c.walkDown(child)
	}
}

func (c *Context) getNewOwners(id int) (owners, times []int) {
	var owner, t int
	err := c.ownerStmt.Query(id, c.Simid)
	for ; err == nil; err = c.ownerStmt.Next() {
		err := c.ownerStmt.Scan(&owner, &t)
		panicif(err)

		if id == owner {
			continue
		}
		owners = append(owners, owner)
		times = append(times, t)
	}
	if err != io.EOF {
		panic(err)
	}
	return owners, times
}

func (c *Context) dumpNodes() {
	fmt.Printf("    Dumping inventories (%d resources done)...\n", c.resCount)
	err := c.Exec("BEGIN TRANSACTION;")
	panicif(err)

	for _, n := range c.nodes {
		err = c.dumpStmt.Exec(c.Simid, n.ResId, n.OwnerId, n.StartTime, n.EndTime, n.StateId, n.Quantity)
		panicif(err)
	}
	err = c.Exec("END TRANSACTION;")
	panicif(err)

	c.nodes = c.nodes[:0]
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"go/build"
	"go/parser"
	"go/printer"
	"go/token"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

var fake = flag.Bool("n", false, "If true, don't actually do anything")

func main() {
	flag.Parse()
	pkgName := flag.Arg(0)
	if pkgName == "" {
		log.Fatal("need a package name")
	}
	dest := flag.Arg(1)
	if dest == "" {
		log.Fatal("need a destination path")
	}

	err := vendorize(pkgName, pkgName, dest)
	if err != nil {
		log.Fatal(err)
	}
}

func vendorize(proj, path, dest string) error {
	rootPkg, err := buildPackage(path)
	if err != nil {
		return fmt.Errorf("couldn't import %s: %s", path, err)
	}
	if rootPkg.Goroot {
		return fmt.Errorf("can't vendorize packages from GOROOT")
	}

	allImports := getAllImports(rootPkg)

	var pkgs []*build.Package
	for _, imp := range allImports {
		if strings.HasPrefix(imp, proj) || strings.HasPrefix(imp, dest) {
			// don't process things we presumably don't need to vendor
			continue
		}

		pkg, err := buildPackage(imp)
		if err != nil {
			return fmt.Errorf("%s: couldn't import %s: %s", path, imp, err)
		}
		if !pkg.Goroot {
			pkgs = append(pkgs, pkg)
		}
	}

	rewrites := make(map[string]string)
	for _, pkg := range pkgs {
		err := vendorize(proj, pkg.ImportPath, dest)
		if err != nil {
			return fmt.Errorf("couldn't vendorize %s: %s", pkg.ImportPath, err)
		}
		rewrites[pkg.ImportPath] = dest + "/" + pkg.ImportPath
	}

	pkgDir := rootPkg.Dir

	// Copy all the files to the destination.
	if proj != path {
		destImportPath := dest + "/" + path
		gopaths := filepath.SplitList(os.Getenv("GOPATH"))
		gopath := gopaths[len(gopaths)-1]
		pkgDir = filepath.Join(gopath, "src", destImportPath)

		log.Printf("copying contents of %q to %q", rootPkg.Dir, pkgDir)
		err := copyDir(pkgDir, rootPkg.Dir)
		if err != nil {
			return fmt.Errorf("couldn't copy %s: %s", rootPkg, err)
		}
	}

	// Rewrite any imports
	for _, files := range [][]string{
		rootPkg.GoFiles, rootPkg.CgoFiles, rootPkg.TestGoFiles, rootPkg.XTestGoFiles,
	} {
		for _, file := range files {
			if len(rewrites) > 0 {
				destFile := filepath.Join(pkgDir, file)
				err := rewriteFile(destFile, filepath.Join(rootPkg.Dir, file), rewrites)
				if err != nil {
					return fmt.Errorf("%s: couldn't rewrite file %q: %s", path, file, err)
				}
			}
		}
	}
	return nil
}

// copyFile copies the file given by src to dest, creating dest with the permissions given by perm.
func copyFile(dest, src string, perm os.FileMode) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

	out, err := os.OpenFile(dest, os.O_RDWR|os.O_CREATE|os.O_TRUNC, perm)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = io.Copy(out, in)
	return err
}

// copyDir non-recursively copies the contents of the src directory to dest.
func copyDir(dest, src string) error {
	return filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// We don't recurse.
		if path != dest && info.IsDir() {
			return filepath.SkipDir
		}

		srcFile := filepath.Join(src, path)
		destFile := filepath.Join(dest, path)
		log.Printf("copying %q to %q", srcFile, destFile)
		if *fake {
			return nil
		}
		return copyFile(destFile, srcFile, info.Mode().Perm())
	})
}

// getAllImports returns a list of all import paths in the Go files of pkg.
func getAllImports(pkg *build.Package) []string {
	allImports := make(map[string]bool)
	for _, imports := range [][]string{pkg.Imports, pkg.TestImports, pkg.XTestImports} {
		for _, imp := range imports {
			allImports[imp] = true
		}
	}
	result := make([]string, 0, len(allImports))
	for imp := range allImports {
		result = append(result, imp)
	}
	return result
}

// builtPackages keeps a cache of package builds.
var builtPackages map[string]*build.Package

// buildPackage builds a package given by the path.
func buildPackage(path string) (*build.Package, error) {
	if builtPackages == nil {
		builtPackages = make(map[string]*build.Package)
	}
	if pkg, ok := builtPackages[path]; ok {
		return pkg, nil
	}

	ctx := build.Default
	// TODO(kisielk): support relative imports?
	pkg, err := ctx.Import(path, "", 0)
	if err != nil {
		return nil, err
	}
	builtPackages[path] = pkg
	return pkg, nil
}

func rewriteFile(dest, path string, m map[string]string) error {
	f, err := ioutil.TempFile("", "vendorize")
	if err != nil {
		return err
	}
	defer f.Close()
	err = rewriteImports(path, m, f)
	if err != nil {
		return err
	}
	return os.Rename(f.Name(), dest)
}

func rewriteImports(path string, m map[string]string, w io.Writer) error {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
	if err != nil {
		return err
	}

	for _, s := range f.Imports {
		path, err := strconv.Unquote(s.Path.Value)
		if err != nil {
			panic(err)
		}
		if replacement, ok := m[path]; ok {
			s.Path.Value = strconv.Quote(replacement)
		}
	}

	return printer.Fprint(w, fset, f)
}
<commit_msg>First working version.<commit_after>package main

import (
	"flag"
	"fmt"
	"go/build"
	"go/parser"
	"go/printer"
	"go/token"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

var fake = flag.Bool("n", false, "If true, don't actually do anything")

// The last component of GOPATH
var gopath string

func main() {
	gopaths := filepath.SplitList(os.Getenv("GOPATH"))
	gopath = gopaths[len(gopaths)-1]
	if gopath == "" {
		log.Fatal("GOPATH must be set")
	}

	flag.Parse()
	pkgName := flag.Arg(0)
	if pkgName == "" {
		log.Fatal("need a package name")
	}
	dest := flag.Arg(1)
	if dest == "" {
		log.Fatal("need a destination path")
	}

	err := vendorize(pkgName, pkgName, dest)
	if err != nil {
		log.Fatal(err)
	}
}

func vendorize(proj, path, dest string) error {
	rootPkg, err := buildPackage(path)
	if err != nil {
		return fmt.Errorf("couldn't import %s: %s", path, err)
	}
	if rootPkg.Goroot {
		return fmt.Errorf("can't vendorize packages from GOROOT")
	}

	allImports := getAllImports(rootPkg)

	var pkgs []*build.Package
	for _, imp := range allImports {
		if strings.HasPrefix(imp, proj) || strings.HasPrefix(imp, dest) {
			// don't process things we presumably don't need to vendor
			continue
		}

		pkg, err := buildPackage(imp)
		if err != nil {
			return fmt.Errorf("%s: couldn't import %s: %s", path, imp, err)
		}
		if !pkg.Goroot {
			pkgs = append(pkgs, pkg)
		}
	}

	rewrites := make(map[string]string)
	for _, pkg := range pkgs {
		err := vendorize(proj, pkg.ImportPath, dest)
		if err != nil {
			return fmt.Errorf("couldn't vendorize %s: %s", pkg.ImportPath, err)
		}
		rewrites[pkg.ImportPath] = dest + "/" + pkg.ImportPath
	}

	pkgDir := rootPkg.Dir

	// Copy all the files to the destination.
	if proj != path {
		pkgDir = filepath.Join(gopath, "src", dest, path)
		err := os.MkdirAll(pkgDir, 0770)
		if err != nil {
			return fmt.Errorf("%s: couldn't make destination directory: %s", path, pkgDir)
		}

		log.Printf("copying contents of %q to %q", rootPkg.Dir, pkgDir)
		err = copyDir(pkgDir, rootPkg.Dir)
		if err != nil {
			return fmt.Errorf("couldn't copy %s: %s", rootPkg.ImportPath, err)
		}
	}

	// Rewrite any imports
	for _, files := range [][]string{
		rootPkg.GoFiles, rootPkg.CgoFiles, rootPkg.TestGoFiles, rootPkg.XTestGoFiles,
	} {
		for _, file := range files {
			if len(rewrites) > 0 {
				destFile := filepath.Join(pkgDir, file)
				log.Printf("rewriting imports in %q", destFile)
				err := rewriteFile(destFile, filepath.Join(rootPkg.Dir, file), rewrites)
				if err != nil {
					return fmt.Errorf("%s: couldn't rewrite file %q: %s", path, file, err)
				}
			}
		}
	}
	return nil
}

// copyFile copies the file given by src to dest, creating dest with the permissions given by perm.
func copyFile(dest, src string, perm os.FileMode) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

	out, err := os.OpenFile(dest, os.O_RDWR|os.O_CREATE|os.O_TRUNC, perm)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = io.Copy(out, in)
	return err
}

// copyDir non-recursively copies the contents of the src directory to dest.
func copyDir(dest, src string) error {
	return filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// We don't recurse.
		if info.IsDir() {
			if path != src {
				return filepath.SkipDir
			}
			return nil
		}

		relPath, err := filepath.Rel(src, path)
		if err != nil {
			return err
		}
		destFile := filepath.Join(dest, relPath)
		log.Printf("copying %q to %q", path, destFile)
		if *fake {
			return nil
		}
		return copyFile(destFile, path, info.Mode().Perm())
	})
}

// getAllImports returns a list of all import paths in the Go files of pkg.
func getAllImports(pkg *build.Package) []string {
	allImports := make(map[string]bool)
	for _, imports := range [][]string{pkg.Imports, pkg.TestImports, pkg.XTestImports} {
		for _, imp := range imports {
			allImports[imp] = true
		}
	}
	result := make([]string, 0, len(allImports))
	for imp := range allImports {
		result = append(result, imp)
	}
	return result
}

// builtPackages keeps a cache of package builds.
var builtPackages map[string]*build.Package

// buildPackage builds a package given by the path.
func buildPackage(path string) (*build.Package, error) {
	if builtPackages == nil {
		builtPackages = make(map[string]*build.Package)
	}
	if pkg, ok := builtPackages[path]; ok {
		return pkg, nil
	}

	ctx := build.Default
	// TODO(kisielk): support relative imports?
	pkg, err := ctx.Import(path, "", 0)
	if err != nil {
		return nil, err
	}
	builtPackages[path] = pkg
	return pkg, nil
}

func rewriteFile(dest, path string, m map[string]string) error {
	f, err := ioutil.TempFile("", "vendorize")
	if err != nil {
		return err
	}
	defer f.Close()
	err = rewriteImports(path, m, f)
	if err != nil {
		return err
	}
	return os.Rename(f.Name(), dest)
}

func rewriteImports(path string, m map[string]string, w io.Writer) error {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
	if err != nil {
		return err
	}

	for _, s := range f.Imports {
		path, err := strconv.Unquote(s.Path.Value)
		if err != nil {
			panic(err)
		}
		if replacement, ok := m[path]; ok {
			s.Path.Value = strconv.Quote(replacement)
		}
	}

	return printer.Fprint(w, fset, f)
}
<|endoftext|>


<commit_before>package epp

import (
	"bytes"
	"encoding/xml"
	"testing"

	"github.com/nbio/st"
)

func TestConnCheck(t *testing.T) {
	c := testLogin(t)

	dcr, err := c.CheckDomain("google.com")
	st.Expect(t, err, nil)
	st.Reject(t, dcr, nil)
	st.Expect(t, len(dcr.Checks), 1)
	st.Expect(t, dcr.Checks[0].Domain, "google.com")
	st.Expect(t, dcr.Checks[0].Available, false)

	dcr, err = c.CheckDomain("dmnr-test-x759824vim-i2.com")
	st.Expect(t, err, nil)
	st.Reject(t, dcr, nil)
	st.Expect(t, len(dcr.Checks), 1)
	st.Expect(t, dcr.Checks[0].Domain, "dmnr-test-x759824vim-i2.com")
	st.Expect(t, dcr.Checks[0].Available, true)

	dcr, err = c.CheckDomain("--dmnr-test--.com")
	st.Reject(t, err, nil)
	st.Expect(t, dcr, (*DomainCheckResponse)(nil))
}

func TestEncodeDomainCheck(t *testing.T) {
	var buf bytes.Buffer
	err := encodeDomainCheck(&buf, []string{"hello.com", "foo.domains", "xn--ninja.net"})
	st.Expect(t, err, nil)
	st.Expect(t, buf.String(), `<?xml version="1.0" encoding="UTF-8"?>
<epp xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="urn:ietf:params:xml:ns:epp-1.0 epp-1.0.xsd" xmlns="urn:ietf:params:xml:ns:epp-1.0"><command><check xmlns:domain="urn:ietf:params:xml:ns:domain-1.0"><domain:check><domain:name>hello.com</domain:name></domain:check><domain:check><domain:name>foo.domains</domain:name></domain:check><domain:check><domain:name>xn--ninja.net</domain:name></domain:check></check></command></epp>`)
	var msg message
	err = xml.Unmarshal(buf.Bytes(), &msg)
	st.Expect(t, err, nil)
}

func TestScanCheckDomainResponse(t *testing.T) {
	var res response_
	dcr := &res.DomainCheckResponse

	d := decoder(testXMLDomainCheckResponse)
	err := IgnoreEOF(scanResponse.Scan(d, &res))
	st.Expect(t, err, nil)
	st.Expect(t, len(dcr.Checks), 1)
	st.Expect(t, dcr.Checks[0].Domain, "good.memorial")
	st.Expect(t, dcr.Checks[0].Available, true)
	st.Expect(t, dcr.Checks[0].Reason, "premium name")
	st.Expect(t, len(dcr.Charges), 1)
	st.Expect(t, dcr.Charges[0].Domain, "good.memorial")
	st.Expect(t, dcr.Charges[0].Category, "premium")
	st.Expect(t, dcr.Charges[0].CategoryName, "BBB+")
}

var testXMLDomainCheckResponse = `<?xml version="1.0" encoding="utf-8"?>
<epp xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="urn:ietf:params:xml:ns:epp-1.0 epp-1.0.xsd" xmlns="urn:ietf:params:xml:ns:epp-1.0">
	<response>
		<result code="1000">
			<msg>Command completed successfully</msg>
		</result>
		<resData>
			<domain:chkData xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="urn:ietf:params:xml:ns:domain-1.0 domain-1.0.xsd" xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
				<domain:cd>
					<domain:name avail="1">good.memorial</domain:name>
					<domain:reason>premium name</domain:reason>
				</domain:cd>
			</domain:chkData>
		</resData>
		<extension>
			<charge:chkData xmlns:charge="http://www.unitedtld.com/epp/charge-1.0">
				<charge:cd>
					<charge:name>good.memorial</charge:name>
					<charge:set>
						<charge:category name="BBB+">premium</charge:category>
						<charge:type>price</charge:type>
						<charge:amount command="create">100.00</charge:amount>
						<charge:amount command="renew">100.00</charge:amount>
						<charge:amount command="transfer">100.00</charge:amount>
						<charge:amount command="update" name="restore">50.00</charge:amount>
					</charge:set>
				</charge:cd>
			</charge:chkData>
		</extension>
		<trID>
			<clTRID>0000000000000002</clTRID>
			<svTRID>83fa5767-5624-4be5-9e54-0b3a52f9de5b:1</svTRID>
		</trID>
	</response>
</epp>`

func BenchmarkEncodeDomainCheck(b *testing.B) {
	var buf bytes.Buffer
	domains := []string{"hello.com"}
	for i := 0; i < b.N; i++ {
		encodeDomainCheck(&buf, domains)
	}
}
<commit_msg>Benchmark encoding and scanning domain:check request/response.<commit_after>package epp

import (
	"bytes"
	"encoding/xml"
	"testing"

	"github.com/nbio/st"
)

func TestConnCheck(t *testing.T) {
	c := testLogin(t)

	dcr, err := c.CheckDomain("google.com")
	st.Expect(t, err, nil)
	st.Reject(t, dcr, nil)
	st.Expect(t, len(dcr.Checks), 1)
	st.Expect(t, dcr.Checks[0].Domain, "google.com")
	st.Expect(t, dcr.Checks[0].Available, false)

	dcr, err = c.CheckDomain("dmnr-test-x759824vim-i2.com")
	st.Expect(t, err, nil)
	st.Reject(t, dcr, nil)
	st.Expect(t, len(dcr.Checks), 1)
	st.Expect(t, dcr.Checks[0].Domain, "dmnr-test-x759824vim-i2.com")
	st.Expect(t, dcr.Checks[0].Available, true)

	dcr, err = c.CheckDomain("--dmnr-test--.com")
	st.Reject(t, err, nil)
	st.Expect(t, dcr, (*DomainCheckResponse)(nil))
}

func TestEncodeDomainCheck(t *testing.T) {
	var buf bytes.Buffer
	err := encodeDomainCheck(&buf, []string{"hello.com", "foo.domains", "xn--ninja.net"})
	st.Expect(t, err, nil)
	st.Expect(t, buf.String(), `<?xml version="1.0" encoding="UTF-8"?>
<epp xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="urn:ietf:params:xml:ns:epp-1.0 epp-1.0.xsd" xmlns="urn:ietf:params:xml:ns:epp-1.0"><command><check xmlns:domain="urn:ietf:params:xml:ns:domain-1.0"><domain:check><domain:name>hello.com</domain:name></domain:check><domain:check><domain:name>foo.domains</domain:name></domain:check><domain:check><domain:name>xn--ninja.net</domain:name></domain:check></check></command></epp>`)
	var msg message
	err = xml.Unmarshal(buf.Bytes(), &msg)
	st.Expect(t, err, nil)
}

func TestScanCheckDomainResponse(t *testing.T) {
	var res response_
	dcr := &res.DomainCheckResponse

	d := decoder(testXMLDomainCheckResponse)
	err := IgnoreEOF(scanResponse.Scan(d, &res))
	st.Expect(t, err, nil)
	st.Expect(t, len(dcr.Checks), 1)
	st.Expect(t, dcr.Checks[0].Domain, "good.memorial")
	st.Expect(t, dcr.Checks[0].Available, true)
	st.Expect(t, dcr.Checks[0].Reason, "premium name")
	st.Expect(t, len(dcr.Charges), 1)
	st.Expect(t, dcr.Charges[0].Domain, "good.memorial")
	st.Expect(t, dcr.Charges[0].Category, "premium")
	st.Expect(t, dcr.Charges[0].CategoryName, "BBB+")
}

func BenchmarkEncodeDomainCheck(b *testing.B) {
	var buf bytes.Buffer
	domains := []string{"hello.com"}
	for i := 0; i < b.N; i++ {
		encodeDomainCheck(&buf, domains)
	}
}

func BenchmarkScanDomainCheckResponse(b *testing.B) {
	for i := 0; i < b.N; i++ {
		b.StopTimer()
		d := decoder(testXMLDomainCheckResponse)
		b.StartTimer()
		var res response_
		scanResponse.Scan(d, &res)
	}
}

var testXMLDomainCheckResponse = `<?xml version="1.0" encoding="utf-8"?>
<epp xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="urn:ietf:params:xml:ns:epp-1.0 epp-1.0.xsd" xmlns="urn:ietf:params:xml:ns:epp-1.0">
	<response>
		<result code="1000">
			<msg>Command completed successfully</msg>
		</result>
		<resData>
			<domain:chkData xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="urn:ietf:params:xml:ns:domain-1.0 domain-1.0.xsd" xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
				<domain:cd>
					<domain:name avail="1">good.memorial</domain:name>
					<domain:reason>premium name</domain:reason>
				</domain:cd>
			</domain:chkData>
		</resData>
		<extension>
			<charge:chkData xmlns:charge="http://www.unitedtld.com/epp/charge-1.0">
				<charge:cd>
					<charge:name>good.memorial</charge:name>
					<charge:set>
						<charge:category name="BBB+">premium</charge:category>
						<charge:type>price</charge:type>
						<charge:amount command="create">100.00</charge:amount>
						<charge:amount command="renew">100.00</charge:amount>
						<charge:amount command="transfer">100.00</charge:amount>
						<charge:amount command="update" name="restore">50.00</charge:amount>
					</charge:set>
				</charge:cd>
			</charge:chkData>
		</extension>
		<trID>
			<clTRID>0000000000000002</clTRID>
			<svTRID>83fa5767-5624-4be5-9e54-0b3a52f9de5b:1</svTRID>
		</trID>
	</response>
</epp>`

func BenchmarkEncodeDomainCheck(b *testing.B) {
	var buf bytes.Buffer
	domains := []string{"hello.com"}
	for i := 0; i < b.N; i++ {
		encodeDomainCheck(&buf, domains)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus/promhttp"

	"github.com/grafana/loki/pkg/canary/comparator"
	"github.com/grafana/loki/pkg/canary/reader"
	"github.com/grafana/loki/pkg/canary/writer"
)

func main() {

	lName := flag.String("labelname", "name", "The label name for this instance of loki-canary to use in the log selector")
	lVal := flag.String("labelvalue", "loki-canary", "The unique label value for this instance of loki-canary to use in the log selector")
	port := flag.Int("port", 3500, "Port which loki-canary should expose metrics")
	addr := flag.String("addr", "", "The Loki server URL:Port, e.g. loki:3100")
	tls := flag.Bool("tls", false, "Does the loki connection use TLS?")
	user := flag.String("user", "", "Loki username")
	pass := flag.String("pass", "", "Loki password")

	interval := flag.Duration("interval", 1000*time.Millisecond, "Duration between log entries")
	size := flag.Int("size", 100, "Size in bytes of each log line")
	wait := flag.Duration("wait", 60*time.Second, "Duration to wait for log entries before reporting them lost")
	pruneInterval := flag.Duration("pruneinterval", 60*time.Second, "Frequency to check sent vs received logs, also the frequency which queries for missing logs will be dispatched to loki")
	buckets := flag.Int("buckets", 10, "Number of buckets in the response_latency histogram")
	flag.Parse()

	if *addr == "" {
		_, _ = fmt.Fprintf(os.Stderr, "Must specify a Loki address with -addr\n")
		os.Exit(1)
	}

	sentChan := make(chan time.Time)
	receivedChan := make(chan time.Time)

	w := writer.NewWriter(os.Stdout, sentChan, *interval, *size)
	r := reader.NewReader(os.Stderr, receivedChan, *tls, *addr, *user, *pass, *lName, *lVal)
	c := comparator.NewComparator(os.Stderr, *wait, *pruneInterval, *buckets, sentChan, receivedChan, r)

	http.Handle("/metrics", promhttp.Handler())
	go func() {
		err := http.ListenAndServe(":"+strconv.Itoa(*port), nil)
		if err != nil {
			panic(err)
		}
	}()

	interrupt := make(chan os.Signal, 1)
	terminate := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)
	signal.Notify(terminate, syscall.SIGTERM)

	for {
		select {
		case <-interrupt:
			_, _ = fmt.Fprintf(os.Stderr, "suspending indefinetely\n")
			w.Stop()
			r.Stop()
			c.Stop()
		case <-terminate:
			_, _ = fmt.Fprintf(os.Stderr, "shutting down\n")
			w.Stop()
			r.Stop()
			c.Stop()
			return
		}
	}

}
<commit_msg>style(loki-canary): spelling  (#832)<commit_after>package main

import (
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus/promhttp"

	"github.com/grafana/loki/pkg/canary/comparator"
	"github.com/grafana/loki/pkg/canary/reader"
	"github.com/grafana/loki/pkg/canary/writer"
)

func main() {

	lName := flag.String("labelname", "name", "The label name for this instance of loki-canary to use in the log selector")
	lVal := flag.String("labelvalue", "loki-canary", "The unique label value for this instance of loki-canary to use in the log selector")
	port := flag.Int("port", 3500, "Port which loki-canary should expose metrics")
	addr := flag.String("addr", "", "The Loki server URL:Port, e.g. loki:3100")
	tls := flag.Bool("tls", false, "Does the loki connection use TLS?")
	user := flag.String("user", "", "Loki username")
	pass := flag.String("pass", "", "Loki password")

	interval := flag.Duration("interval", 1000*time.Millisecond, "Duration between log entries")
	size := flag.Int("size", 100, "Size in bytes of each log line")
	wait := flag.Duration("wait", 60*time.Second, "Duration to wait for log entries before reporting them lost")
	pruneInterval := flag.Duration("pruneinterval", 60*time.Second, "Frequency to check sent vs received logs, also the frequency which queries for missing logs will be dispatched to loki")
	buckets := flag.Int("buckets", 10, "Number of buckets in the response_latency histogram")
	flag.Parse()

	if *addr == "" {
		_, _ = fmt.Fprintf(os.Stderr, "Must specify a Loki address with -addr\n")
		os.Exit(1)
	}

	sentChan := make(chan time.Time)
	receivedChan := make(chan time.Time)

	w := writer.NewWriter(os.Stdout, sentChan, *interval, *size)
	r := reader.NewReader(os.Stderr, receivedChan, *tls, *addr, *user, *pass, *lName, *lVal)
	c := comparator.NewComparator(os.Stderr, *wait, *pruneInterval, *buckets, sentChan, receivedChan, r)

	http.Handle("/metrics", promhttp.Handler())
	go func() {
		err := http.ListenAndServe(":"+strconv.Itoa(*port), nil)
		if err != nil {
			panic(err)
		}
	}()

	interrupt := make(chan os.Signal, 1)
	terminate := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)
	signal.Notify(terminate, syscall.SIGTERM)

	for {
		select {
		case <-interrupt:
			_, _ = fmt.Fprintf(os.Stderr, "suspending indefinitely\n")
			w.Stop()
			r.Stop()
			c.Stop()
		case <-terminate:
			_, _ = fmt.Fprintf(os.Stderr, "shutting down\n")
			w.Stop()
			r.Stop()
			c.Stop()
			return
		}
	}

}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"os"

	"github.com/containers/storage/pkg/archive"
	"github.com/containers/storage/pkg/mflag"
	"github.com/containers/storage/storage"
)

func changes(flags *mflag.FlagSet, action string, m storage.Mall, args []string) int {
	if len(args) < 1 {
		return 1
	}
	to := args[0]
	from := ""
	if len(args) >= 2 {
		from = args[1]
	}
	changes, err := m.Changes(to, from)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return 1
	}
	if jsonOutput {
		json.NewEncoder(os.Stdout).Encode(changes)
	} else {
		for _, change := range changes {
			what := "?"
			switch change.Kind {
			case archive.ChangeAdd:
				what = "Add"
			case archive.ChangeModify:
				what = "Modify"
			case archive.ChangeDelete:
				what = "Delete"
			}
			fmt.Printf("%s %q\n", what, change.Path)
		}
	}
	return 0
}

func diff(flags *mflag.FlagSet, action string, m storage.Mall, args []string) int {
	if len(args) < 2 {
		return 1
	}
	to := args[0]
	from := ""
	if len(args) >= 2 {
		from = args[1]
	}
	reader, err := m.Diff(to, from)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return 1
	}
	_, err = io.Copy(os.Stdout, reader)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return 1
	}
	return 0
}

func applyDiff(flags *mflag.FlagSet, action string, m storage.Mall, args []string) int {
	if len(args) < 1 {
		return 1
	}
	_, err := m.ApplyDiff(args[0], os.Stdin)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return 1
	}
	return 0
}

func diffSize(flags *mflag.FlagSet, action string, m storage.Mall, args []string) int {
	if len(args) < 1 {
		return 1
	}
	to := args[0]
	from := ""
	if len(args) >= 2 {
		from = args[1]
	}
	n, err := m.DiffSize(to, from)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return 1
	}
	fmt.Printf("%d\n", n)
	return 0
}

func init() {
	commands = append(commands, command{
		names:       []string{"changes"},
		usage:       "Compare two layers",
		optionsHelp: "[options [...]] layerNameOrID [referenceLayerNameOrID]",
		minArgs:     1,
		maxArgs:     2,
		action:      changes,
		addFlags: func(flags *mflag.FlagSet, cmd *command) {
			flags.BoolVar(&jsonOutput, []string{"-json", "j"}, jsonOutput, "Prefer JSON output")
		},
	})
	commands = append(commands, command{
		names:       []string{"diffsize"},
		usage:       "Compare two layers",
		optionsHelp: "[options [...]] layerNameOrID [referenceLayerNameOrID]",
		minArgs:     1,
		maxArgs:     2,
		action:      diffSize,
	})
	commands = append(commands, command{
		names:       []string{"diff"},
		usage:       "Compare two layers",
		optionsHelp: "[options [...]] layerNameOrID [referenceLayerNameOrID]",
		minArgs:     1,
		maxArgs:     2,
		action:      diff,
	})
	commands = append(commands, command{
		names:       []string{"applydiff"},
		optionsHelp: "[options [...]] layerNameOrID [referenceLayerNameOrID]",
		usage:       "Apply a diff to a layer",
		minArgs:     1,
		maxArgs:     1,
		action:      applyDiff,
	})
}
<commit_msg>Handle compressed diffs in oci-storage applydiff<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"os"

	"github.com/containers/storage/pkg/archive"
	"github.com/containers/storage/pkg/mflag"
	"github.com/containers/storage/storage"
)

var (
	applyDiffFile       = ""
	applyDiffCompressed = false
)

func changes(flags *mflag.FlagSet, action string, m storage.Mall, args []string) int {
	if len(args) < 1 {
		return 1
	}
	to := args[0]
	from := ""
	if len(args) >= 2 {
		from = args[1]
	}
	changes, err := m.Changes(to, from)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return 1
	}
	if jsonOutput {
		json.NewEncoder(os.Stdout).Encode(changes)
	} else {
		for _, change := range changes {
			what := "?"
			switch change.Kind {
			case archive.ChangeAdd:
				what = "Add"
			case archive.ChangeModify:
				what = "Modify"
			case archive.ChangeDelete:
				what = "Delete"
			}
			fmt.Printf("%s %q\n", what, change.Path)
		}
	}
	return 0
}

func diff(flags *mflag.FlagSet, action string, m storage.Mall, args []string) int {
	if len(args) < 2 {
		return 1
	}
	to := args[0]
	from := ""
	if len(args) >= 2 {
		from = args[1]
	}
	reader, err := m.Diff(to, from)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return 1
	}
	_, err = io.Copy(os.Stdout, reader)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return 1
	}
	return 0
}

func applyDiff(flags *mflag.FlagSet, action string, m storage.Mall, args []string) int {
	if len(args) < 1 {
		return 1
	}
	diffStream := io.Reader(os.Stdin)
	if applyDiffFile != "" {
		if f, err := os.Open(applyDiffFile); err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err)
			return 1
		} else {
			diffStream = f
		}
	}
	if applyDiffCompressed {
		if c, err := archive.DecompressStream(diffStream); err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err)
			return 1
		} else {
			diffStream = c
		}
	}
	_, err := m.ApplyDiff(args[0], diffStream)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return 1
	}
	return 0
}

func diffSize(flags *mflag.FlagSet, action string, m storage.Mall, args []string) int {
	if len(args) < 1 {
		return 1
	}
	to := args[0]
	from := ""
	if len(args) >= 2 {
		from = args[1]
	}
	n, err := m.DiffSize(to, from)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%v\n", err)
		return 1
	}
	fmt.Printf("%d\n", n)
	return 0
}

func init() {
	commands = append(commands, command{
		names:       []string{"changes"},
		usage:       "Compare two layers",
		optionsHelp: "[options [...]] layerNameOrID [referenceLayerNameOrID]",
		minArgs:     1,
		maxArgs:     2,
		action:      changes,
		addFlags: func(flags *mflag.FlagSet, cmd *command) {
			flags.BoolVar(&jsonOutput, []string{"-json", "j"}, jsonOutput, "Prefer JSON output")
		},
	})
	commands = append(commands, command{
		names:       []string{"diffsize"},
		usage:       "Compare two layers",
		optionsHelp: "[options [...]] layerNameOrID [referenceLayerNameOrID]",
		minArgs:     1,
		maxArgs:     2,
		action:      diffSize,
	})
	commands = append(commands, command{
		names:       []string{"diff"},
		usage:       "Compare two layers",
		optionsHelp: "[options [...]] layerNameOrID [referenceLayerNameOrID]",
		minArgs:     1,
		maxArgs:     2,
		action:      diff,
	})
	commands = append(commands, command{
		names:       []string{"applydiff"},
		optionsHelp: "[options [...]] layerNameOrID [referenceLayerNameOrID]",
		usage:       "Apply a diff to a layer",
		minArgs:     1,
		maxArgs:     1,
		action:      applyDiff,
		addFlags: func(flags *mflag.FlagSet, cmd *command) {
			flags.StringVar(&applyDiffFile, []string{"-file", "f"}, "", "Read from file instead of stdin")
			flags.BoolVar(&applyDiffCompressed, []string{"-compressed", "c"}, applyDiffCompressed, "Diff is compressed")
		},
	})
}
<|endoftext|>


<commit_before>package commands

import (
	"fmt"
	"os"
	"strconv"

	"github.com/spf13/cobra"
)

var (
	// Commands
	cmdTasks = &cobra.Command{
		Use:   "tasks",
		Short: "gets all tasks",
		Long: `tasks gets the user's tasks.
        `,
		Run: func(cmd *cobra.Command, args []string) {
			if listID == 0 {
				renderOutput(newClient(cmd).Tasks())
			} else {
				renderOutput(newClient(cmd).TasksForListID(listID))
			}
		},
	}

	cmdTask = &cobra.Command{
		Use:   "task <task-id>",
		Short: "gets the task for the provided task id",
		Long: `task gets a task specified by <task-id>
        `,
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) != 1 {
				fmt.Printf("incorrect number of arguments provided\n\n")
				cmd.Usage()
				os.Exit(2)
			}

			idInt, err := strconv.Atoi(args[0])
			if err != nil {
				fmt.Printf("error parsing taskID: %v\n\n", err)
				cmd.Usage()
				os.Exit(2)
			}
			id := uint(idInt)

			renderOutput(newClient(cmd).Task(
				id,
			))
		},
	}

	cmdDeleteTask = &cobra.Command{
		Use:   "delete-task <task-id>",
		Short: "deletes the task for the provided task id",
		Long: `delete-task deletes the task specified by <task-id>
        `,
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) != 1 {
				fmt.Printf("incorrect number of arguments provided\n\n")
				cmd.Usage()
				os.Exit(2)
			}

			idInt, err := strconv.Atoi(args[0])
			if err != nil {
				fmt.Printf("error parsing taskID: %v\n\n", err)
				cmd.Usage()
				os.Exit(2)
			}
			id := uint(idInt)

			client := newClient(cmd)
			task, err := client.Task(id)
			if err != nil {
				fmt.Printf("error getting task: %v\n\n", err)
				cmd.Usage()
				os.Exit(2)
			}

			err = client.DeleteTask(task)
			if err != nil {
				handleError(err)
			}

			fmt.Printf("task %d deleted successfully\n", id)
		},
	}

	cmdDeleteAllTasks = &cobra.Command{
		Use:   "delete-all-tasks",
		Short: "deletes all tasks",
		Long: `delete-all-tasks deletes all tasks.
        `,
		Run: func(cmd *cobra.Command, args []string) {
			err := newClient(cmd).DeleteAllTasks()
			if err != nil {
				handleError(err)
			}
			fmt.Printf("all tasks deleted successfully\n")
		},
	}
)

func init() {
	cmdTasks.Flags().UintVarP(&listID, listIDLongFlag, listIDShortFlag, 0, "filter by listID")
}
<commit_msg>Add CLI flag --completd to tasks methods.<commit_after>package commands

import (
	"fmt"
	"os"
	"strconv"

	"github.com/spf13/cobra"
)

const (
	completedTasksLongFlag = "completed"
)

var (
	// Flags
	completedTasks bool

	// Commands
	cmdTasks = &cobra.Command{
		Use:   "tasks",
		Short: "gets all tasks",
		Long: `tasks gets the user's tasks.
        `,
		Run: func(cmd *cobra.Command, args []string) {
			// Currently sending completed=false is the same as not sending completed
			// Checking for whether the flag has changed protects us from potentially
			// breaking changes, i.e. if the tasks endpoint decides to run all tasks,
			// not just non-completed ones.

			if listID == 0 {
				if cmd.Flags().Changed(completedTasksLongFlag) {
					renderOutput(newClient(cmd).CompletedTasks(completedTasks))
				} else {
					renderOutput(newClient(cmd).Tasks())
				}
			} else {
				if cmd.Flags().Changed(completedTasksLongFlag) {
					renderOutput(newClient(cmd).CompletedTasksForListID(listID, completedTasks))
				} else {
					renderOutput(newClient(cmd).TasksForListID(listID))
				}
			}
		},
	}

	cmdTask = &cobra.Command{
		Use:   "task <task-id>",
		Short: "gets the task for the provided task id",
		Long: `task gets a task specified by <task-id>
        `,
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) != 1 {
				fmt.Printf("incorrect number of arguments provided\n\n")
				cmd.Usage()
				os.Exit(2)
			}

			idInt, err := strconv.Atoi(args[0])
			if err != nil {
				fmt.Printf("error parsing taskID: %v\n\n", err)
				cmd.Usage()
				os.Exit(2)
			}
			id := uint(idInt)

			renderOutput(newClient(cmd).Task(
				id,
			))
		},
	}

	cmdDeleteTask = &cobra.Command{
		Use:   "delete-task <task-id>",
		Short: "deletes the task for the provided task id",
		Long: `delete-task deletes the task specified by <task-id>
        `,
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) != 1 {
				fmt.Printf("incorrect number of arguments provided\n\n")
				cmd.Usage()
				os.Exit(2)
			}

			idInt, err := strconv.Atoi(args[0])
			if err != nil {
				fmt.Printf("error parsing taskID: %v\n\n", err)
				cmd.Usage()
				os.Exit(2)
			}
			id := uint(idInt)

			client := newClient(cmd)
			task, err := client.Task(id)
			if err != nil {
				fmt.Printf("error getting task: %v\n\n", err)
				cmd.Usage()
				os.Exit(2)
			}

			err = client.DeleteTask(task)
			if err != nil {
				handleError(err)
			}

			fmt.Printf("task %d deleted successfully\n", id)
		},
	}

	cmdDeleteAllTasks = &cobra.Command{
		Use:   "delete-all-tasks",
		Short: "deletes all tasks",
		Long: `delete-all-tasks deletes all tasks.
        `,
		Run: func(cmd *cobra.Command, args []string) {
			err := newClient(cmd).DeleteAllTasks()
			if err != nil {
				handleError(err)
			}
			fmt.Printf("all tasks deleted successfully\n")
		},
	}
)

func init() {
	cmdTasks.Flags().UintVarP(&listID, listIDLongFlag, listIDShortFlag, 0, "filter by listID")
	cmdTasks.Flags().BoolVar(&completedTasks, completedTasksLongFlag, false, "filter for completed tasks")
}
<|endoftext|>


<commit_before>package cli

import (
	"encoding/json"
	"fmt"
	"os"
	"os/user"
	"reflect"
	"strconv"
	"strings"
)

var defaultConfig = `{
	"ping" : {
		"timeout" : "2s",
		"interval": "1s",
		"count":	4
	},
	"hping" : {
		"timeout"  : "2s",
		"interval" : "1s",
		"method"   : "HEAD",
		"data"	   : "mylg",
		"count"	   : 5
	},
	"web" : {
		"port"	   : 8080,
		"address"  : "127.0.0.1"
	},
	"scan" : {
		"port"     : "1-500"
	},
	"trace" : {
		"wait"  : "2s",
		"theme" : "dark"
	},
	"snmp" : {
		"community"     : "public",
		"timeout"       : "1s",
		"version"       : "2c",
		"retries"       : 1,
		"port"          : 161,
		"securitylevel" : "noauthnopriv",
		"authpass"      : "nopass",
		"authproto"     : "sha",
		"Privacypass"   : "nopass",
		"Privacyproto"  : "aes"
	}
}`

// Config represents configuration
type Config struct {
	Ping  Ping  `json:"ping"`
	Hping HPing `json:"hping"`
	Web   Web   `json:"web"`
	Scan  Scan  `json:"scan"`
	Trace Trace `json:"trace"`
	Snmp  SNMP  `json:"snmp"`
}

// Ping represents ping command options
type Ping struct {
	Timeout  string `json:"timeout" tag:"lower"`
	Interval string `json:"interval" tag:"lower"`
	Count    int    `json:"count"`
}

// HPing represents ping command options
type HPing struct {
	Timeout  string `json:"timeout" tag:"lower"`
	Interval string `json:"interval" tag:"lower"`
	Method   string `json:"method" tag:"upper"`
	Data     string `json:"data"`
	Count    int    `json:"count"`
}

// Web represents web command options
type Web struct {
	Port    int    `json:"port"`
	Address string `json:"address"`
}

// Scan represents scan command options
type Scan struct {
	Port string `json:"port"`
}

// Trace represents trace command options
type Trace struct {
	Wait  string `json:"wait" tag:"lower"`
	Theme string `json:"theme" tag:"lower"`
}

// SNMP represents nms command options
type SNMP struct {
	Community     string `json:"community"`
	Timeout       string `json:"timeout" tag:"lower"`
	Version       string `json:"version" tag:"lower"`
	Retries       int    `json:"retries"`
	Port          int    `json:"port"`
	Securitylevel string `json:"securitylevel"`
	Authpass      string `json:"authpass"`
	Authproto     string `json:"authproto" tag:"lower"`
	Privacypass   string `json:"privacypass"`
	Privacyproto  string `json:"privacyproto" tag:"lower"`
}

// WriteConfig write config to disk
func WriteConfig(cfg Config) error {
	f, err := cfgFile()
	if err != nil {
		return err
	}

	h, err := os.Create(f)
	if err != nil {
		return err
	}

	b, err := json.Marshal(cfg)
	if err != nil {
		return err
	}
	_, err = h.Write(b)
	if err != nil {
		return err
	}
	h.Close()

	return nil
}

// GetOptions returns option(s)/value(s) for specific command
func GetOptions(s interface{}, key string) ([]string, []interface{}) {
	var (
		opts []string
		vals []interface{}
	)
	v := reflect.ValueOf(s)
	t := v.Type()

	for i := 0; i < v.NumField(); i++ {
		if t.Field(i).Name == key {
			f := v.Field(i)
			ft := f.Type()
			for j := 0; j < f.NumField(); j++ {
				vals = append(vals, f.Field(j))
				opts = append(opts, ft.Field(j).Name)
			}
			break
		}
	}
	return opts, vals
}

// GetCMDNames returns command line names
func GetCMDNames(s interface{}) []string {
	var fields []string

	v := reflect.ValueOf(s)
	t := v.Type()

	for i := 0; i < t.NumField(); i++ {
		fields = append(fields, t.Field(i).Name)
	}
	return fields
}

// UpgradeConfig adds / removes new command(s)/option(s)
func UpgradeConfig(cfg *Config) error {
	var (
		conf  map[string]interface{}
		cConf Config
	)

	b := make([]byte, 2048)
	f, err := cfgFile()
	if err != nil {
		return err
	}

	h, err := os.Open(f)
	if err != nil {
		return err
	}
	n, _ := h.Read(b)
	b = b[:n]
	// load saved/old config to conf
	json.Unmarshal(b, &conf)
	// load default config to cConf
	json.Unmarshal([]byte(defaultConfig), &cConf)

	for _, cmd := range GetCMDNames(cConf) {
		opts, vals := GetOptions(cConf, cmd)
		for i, opt := range opts {
			if v, ok := conf[strings.ToLower(cmd)].(interface{}); ok {
				if _, ok = v.(map[string]interface{})[strings.ToLower(opt)]; !ok {
					args := fmt.Sprintf("%s %s %v", cmd, opt, vals[i])
					SetConfig(args, cfg)
				}
			} else {
				// there is new command
				args := fmt.Sprintf("%s %s %v", cmd, opt, vals[i])
				SetConfig(args, cfg)
			}
		}
	}
	return nil
}

// LoadConfig loads configuration
func LoadConfig() Config {
	var cfg Config

	cfg = ReadConfig()
	UpgradeConfig(&cfg)
	return cfg
}

// InitConfig creates new config file
func InitConfig(f string) ([]byte, error) {
	h, err := os.Create(f)
	if err != nil {
		return []byte(""), err
	}

	h.Chmod(os.FileMode(int(0600)))
	h.WriteString(defaultConfig)
	h.Close()

	return []byte(defaultConfig), nil
}

// ReadConfig reads configuration from existing
// or default configuration
func ReadConfig() Config {
	var (
		b    = make([]byte, 2048)
		conf Config
		err  error
	)
	f, err := cfgFile()
	if err != nil {

	}

	h, err := os.Open(f)

	if err != nil {
		switch {
		case os.IsNotExist(err):
			if b, err = InitConfig(f); err != nil {
				println(err.Error())
			}
		case os.IsPermission(err):
			println("cannot read configuration file due to insufficient permissions")
			b = []byte(defaultConfig)
		default:
			println(err.Error())
			b = []byte(defaultConfig)
		}
	} else {
		n, _ := h.Read(b)
		b = b[:n]
	}

	err = json.Unmarshal(b, &conf)
	if err != nil {
		println(err.Error())
		b = []byte(defaultConfig)
		json.Unmarshal(b, &conf)
	}

	return conf
}

// ReadDefaultConfig returns default configuration
func ReadDefaultConfig() (Config, error) {
	var (
		b    = make([]byte, 2048)
		conf Config
	)
	b = []byte(defaultConfig)
	err := json.Unmarshal(b, &conf)
	return conf, err
}

// cfgFile returns config file
func cfgFile() (string, error) {
	user, err := user.Current()
	if err != nil {
		return "", err
	}
	return user.HomeDir + "/.mylg.config", nil
}

//
func optionProp(v reflect.Value, opt string) (string, string) {
	opt = strings.Title(opt)
	for i := 0; i < v.NumField(); i++ {
		if v.Type().Field(i).Name == opt {
			field := v.Type().Field(i)
			return field.Tag.Get("tag"), field.Type.Name()
		}
	}
	return "", "nan"
}

// SetConfig handles update option's value
func SetConfig(args string, s *Config) error {
	var (
		v       reflect.Value
		integer int64
		float   float64
		err     error
	)

	//args = strings.ToLower(args)
	f := strings.Fields(args)
	if len(f) < 3 {
		helpSet()
		return fmt.Errorf("syntax error")
	}

	cmd, opt, val := f[0], f[1], f[2]
	opt = strings.Title(opt)

	v = reflect.ValueOf(s)
	v = reflect.Indirect(v)
	v = v.FieldByName(strings.Title(cmd))

	if !v.IsValid() {
		return fmt.Errorf("invalid command")
	}

	tags, valType := optionProp(v, opt)
	val = applyTag(val, tags)

	switch valType {
	case "string":
		// string
		err = SetValue(v.Addr(), opt, fmt.Sprintf("%v", val))
	case "int":
		// integer
		if integer, err = strconv.ParseInt(val, 10, 64); err == nil {
			err = SetValue(v.Addr(), strings.Title(opt), integer)
		} else {
			err = fmt.Errorf("the value should be integer")
		}
	case "float":
		// float
		if float, err = strconv.ParseFloat(val, 64); err == nil {
			err = SetValue(v.Addr(), opt, float)
		} else {
			err = fmt.Errorf("the value should be float")
		}
	default:
		err = fmt.Errorf("invalid option")
	}

	if err != nil {
		return err
	}

	// save config
	err = WriteConfig(*s)
	return err
}

// SetValue set optioni's value
func SetValue(v reflect.Value, rec string, val interface{}) error {

	if v.Kind() != reflect.Ptr {
		return fmt.Errorf("not a pointer value")
	}

	v = reflect.Indirect(v)
	switch v.Kind() {
	case reflect.Int:
		if value, ok := val.(int64); ok {
			v.SetInt(value)
		} else {
			return fmt.Errorf("the value should be integer")
		}
	case reflect.Float64:
		if value, ok := val.(float64); ok {
			v.SetFloat(value)
		} else {
			return fmt.Errorf("the value should be float")
		}
	case reflect.String:
		if value, ok := val.(string); ok {
			v.SetString(value)
		} else {
			return fmt.Errorf("the value shouldn't be number")
		}
	case reflect.Struct:
		for i := 0; i < v.NumField(); i++ {
			if v.Type().Field(i).Name == rec {
				err := SetValue(v.Field(i).Addr(), rec, val)
				if err != nil {
					return err
				}
			}
		}
	}
	return nil
}

// ShowConfig prints the configuration
func ShowConfig(s *Config) {
	var v reflect.Value

	v = reflect.ValueOf(s)
	v = reflect.Indirect(v)

	for i := 0; i < v.NumField(); i++ {
		cmd := v.Type().Field(i).Name
		cmd = strings.ToLower(cmd)

		vv := v.Field(i).Addr()
		vv = reflect.Indirect(vv)

		for j := 0; j < vv.NumField(); j++ {
			subCmd := vv.Type().Field(j).Name
			subCmd = strings.ToLower(subCmd)
			value := vv.Field(j)
			fmt.Printf("set %-8s %-15s %v\n", cmd, subCmd, value)
		}
	}
}

func applyTag(val string, typ string) string {
	switch {
	case strings.Contains(typ, "lower"):
		return strings.ToLower(val)
	case strings.Contains(typ, "upper"):
		return strings.ToUpper(val)
	default:
		return val
	}

}

// helpSet shows set command
func helpSet() {
	println(`
          usage:
               set command option value
          example:
               set ping timeout 2s
	`)

}
<commit_msg>increase port range<commit_after>package cli

import (
	"encoding/json"
	"fmt"
	"os"
	"os/user"
	"reflect"
	"strconv"
	"strings"
)

var defaultConfig = `{
	"ping" : {
		"timeout" : "2s",
		"interval": "1s",
		"count":	4
	},
	"hping" : {
		"timeout"  : "2s",
		"interval" : "1s",
		"method"   : "HEAD",
		"data"	   : "mylg",
		"count"	   : 5
	},
	"web" : {
		"port"	   : 8080,
		"address"  : "127.0.0.1"
	},
	"scan" : {
		"port"     : "1-1024"
	},
	"trace" : {
		"wait"  : "2s",
		"theme" : "dark"
	},
	"snmp" : {
		"community"     : "public",
		"timeout"       : "1s",
		"version"       : "2c",
		"retries"       : 1,
		"port"          : 161,
		"securitylevel" : "noauthnopriv",
		"authpass"      : "nopass",
		"authproto"     : "sha",
		"Privacypass"   : "nopass",
		"Privacyproto"  : "aes"
	}
}`

// Config represents configuration
type Config struct {
	Ping  Ping  `json:"ping"`
	Hping HPing `json:"hping"`
	Web   Web   `json:"web"`
	Scan  Scan  `json:"scan"`
	Trace Trace `json:"trace"`
	Snmp  SNMP  `json:"snmp"`
}

// Ping represents ping command options
type Ping struct {
	Timeout  string `json:"timeout" tag:"lower"`
	Interval string `json:"interval" tag:"lower"`
	Count    int    `json:"count"`
}

// HPing represents ping command options
type HPing struct {
	Timeout  string `json:"timeout" tag:"lower"`
	Interval string `json:"interval" tag:"lower"`
	Method   string `json:"method" tag:"upper"`
	Data     string `json:"data"`
	Count    int    `json:"count"`
}

// Web represents web command options
type Web struct {
	Port    int    `json:"port"`
	Address string `json:"address"`
}

// Scan represents scan command options
type Scan struct {
	Port string `json:"port"`
}

// Trace represents trace command options
type Trace struct {
	Wait  string `json:"wait" tag:"lower"`
	Theme string `json:"theme" tag:"lower"`
}

// SNMP represents nms command options
type SNMP struct {
	Community     string `json:"community"`
	Timeout       string `json:"timeout" tag:"lower"`
	Version       string `json:"version" tag:"lower"`
	Retries       int    `json:"retries"`
	Port          int    `json:"port"`
	Securitylevel string `json:"securitylevel"`
	Authpass      string `json:"authpass"`
	Authproto     string `json:"authproto" tag:"lower"`
	Privacypass   string `json:"privacypass"`
	Privacyproto  string `json:"privacyproto" tag:"lower"`
}

// WriteConfig write config to disk
func WriteConfig(cfg Config) error {
	f, err := cfgFile()
	if err != nil {
		return err
	}

	h, err := os.Create(f)
	if err != nil {
		return err
	}

	b, err := json.Marshal(cfg)
	if err != nil {
		return err
	}
	_, err = h.Write(b)
	if err != nil {
		return err
	}
	h.Close()

	return nil
}

// GetOptions returns option(s)/value(s) for specific command
func GetOptions(s interface{}, key string) ([]string, []interface{}) {
	var (
		opts []string
		vals []interface{}
	)
	v := reflect.ValueOf(s)
	t := v.Type()

	for i := 0; i < v.NumField(); i++ {
		if t.Field(i).Name == key {
			f := v.Field(i)
			ft := f.Type()
			for j := 0; j < f.NumField(); j++ {
				vals = append(vals, f.Field(j))
				opts = append(opts, ft.Field(j).Name)
			}
			break
		}
	}
	return opts, vals
}

// GetCMDNames returns command line names
func GetCMDNames(s interface{}) []string {
	var fields []string

	v := reflect.ValueOf(s)
	t := v.Type()

	for i := 0; i < t.NumField(); i++ {
		fields = append(fields, t.Field(i).Name)
	}
	return fields
}

// UpgradeConfig adds / removes new command(s)/option(s)
func UpgradeConfig(cfg *Config) error {
	var (
		conf  map[string]interface{}
		cConf Config
	)

	b := make([]byte, 2048)
	f, err := cfgFile()
	if err != nil {
		return err
	}

	h, err := os.Open(f)
	if err != nil {
		return err
	}
	n, _ := h.Read(b)
	b = b[:n]
	// load saved/old config to conf
	json.Unmarshal(b, &conf)
	// load default config to cConf
	json.Unmarshal([]byte(defaultConfig), &cConf)

	for _, cmd := range GetCMDNames(cConf) {
		opts, vals := GetOptions(cConf, cmd)
		for i, opt := range opts {
			if v, ok := conf[strings.ToLower(cmd)].(interface{}); ok {
				if _, ok = v.(map[string]interface{})[strings.ToLower(opt)]; !ok {
					args := fmt.Sprintf("%s %s %v", cmd, opt, vals[i])
					SetConfig(args, cfg)
				}
			} else {
				// there is new command
				args := fmt.Sprintf("%s %s %v", cmd, opt, vals[i])
				SetConfig(args, cfg)
			}
		}
	}
	return nil
}

// LoadConfig loads configuration
func LoadConfig() Config {
	var cfg Config

	cfg = ReadConfig()
	UpgradeConfig(&cfg)
	return cfg
}

// InitConfig creates new config file
func InitConfig(f string) ([]byte, error) {
	h, err := os.Create(f)
	if err != nil {
		return []byte(""), err
	}

	h.Chmod(os.FileMode(int(0600)))
	h.WriteString(defaultConfig)
	h.Close()

	return []byte(defaultConfig), nil
}

// ReadConfig reads configuration from existing
// or default configuration
func ReadConfig() Config {
	var (
		b    = make([]byte, 2048)
		conf Config
		err  error
	)
	f, err := cfgFile()
	if err != nil {

	}

	h, err := os.Open(f)

	if err != nil {
		switch {
		case os.IsNotExist(err):
			if b, err = InitConfig(f); err != nil {
				println(err.Error())
			}
		case os.IsPermission(err):
			println("cannot read configuration file due to insufficient permissions")
			b = []byte(defaultConfig)
		default:
			println(err.Error())
			b = []byte(defaultConfig)
		}
	} else {
		n, _ := h.Read(b)
		b = b[:n]
	}

	err = json.Unmarshal(b, &conf)
	if err != nil {
		println(err.Error())
		b = []byte(defaultConfig)
		json.Unmarshal(b, &conf)
	}

	return conf
}

// ReadDefaultConfig returns default configuration
func ReadDefaultConfig() (Config, error) {
	var (
		b    = make([]byte, 2048)
		conf Config
	)
	b = []byte(defaultConfig)
	err := json.Unmarshal(b, &conf)
	return conf, err
}

// cfgFile returns config file
func cfgFile() (string, error) {
	user, err := user.Current()
	if err != nil {
		return "", err
	}
	return user.HomeDir + "/.mylg.config", nil
}

//
func optionProp(v reflect.Value, opt string) (string, string) {
	opt = strings.Title(opt)
	for i := 0; i < v.NumField(); i++ {
		if v.Type().Field(i).Name == opt {
			field := v.Type().Field(i)
			return field.Tag.Get("tag"), field.Type.Name()
		}
	}
	return "", "nan"
}

// SetConfig handles update option's value
func SetConfig(args string, s *Config) error {
	var (
		v       reflect.Value
		integer int64
		float   float64
		err     error
	)

	//args = strings.ToLower(args)
	f := strings.Fields(args)
	if len(f) < 3 {
		helpSet()
		return fmt.Errorf("syntax error")
	}

	cmd, opt, val := f[0], f[1], f[2]
	opt = strings.Title(opt)

	v = reflect.ValueOf(s)
	v = reflect.Indirect(v)
	v = v.FieldByName(strings.Title(cmd))

	if !v.IsValid() {
		return fmt.Errorf("invalid command")
	}

	tags, valType := optionProp(v, opt)
	val = applyTag(val, tags)

	switch valType {
	case "string":
		// string
		err = SetValue(v.Addr(), opt, fmt.Sprintf("%v", val))
	case "int":
		// integer
		if integer, err = strconv.ParseInt(val, 10, 64); err == nil {
			err = SetValue(v.Addr(), strings.Title(opt), integer)
		} else {
			err = fmt.Errorf("the value should be integer")
		}
	case "float":
		// float
		if float, err = strconv.ParseFloat(val, 64); err == nil {
			err = SetValue(v.Addr(), opt, float)
		} else {
			err = fmt.Errorf("the value should be float")
		}
	default:
		err = fmt.Errorf("invalid option")
	}

	if err != nil {
		return err
	}

	// save config
	err = WriteConfig(*s)
	return err
}

// SetValue set optioni's value
func SetValue(v reflect.Value, rec string, val interface{}) error {

	if v.Kind() != reflect.Ptr {
		return fmt.Errorf("not a pointer value")
	}

	v = reflect.Indirect(v)
	switch v.Kind() {
	case reflect.Int:
		if value, ok := val.(int64); ok {
			v.SetInt(value)
		} else {
			return fmt.Errorf("the value should be integer")
		}
	case reflect.Float64:
		if value, ok := val.(float64); ok {
			v.SetFloat(value)
		} else {
			return fmt.Errorf("the value should be float")
		}
	case reflect.String:
		if value, ok := val.(string); ok {
			v.SetString(value)
		} else {
			return fmt.Errorf("the value shouldn't be number")
		}
	case reflect.Struct:
		for i := 0; i < v.NumField(); i++ {
			if v.Type().Field(i).Name == rec {
				err := SetValue(v.Field(i).Addr(), rec, val)
				if err != nil {
					return err
				}
			}
		}
	}
	return nil
}

// ShowConfig prints the configuration
func ShowConfig(s *Config) {
	var v reflect.Value

	v = reflect.ValueOf(s)
	v = reflect.Indirect(v)

	for i := 0; i < v.NumField(); i++ {
		cmd := v.Type().Field(i).Name
		cmd = strings.ToLower(cmd)

		vv := v.Field(i).Addr()
		vv = reflect.Indirect(vv)

		for j := 0; j < vv.NumField(); j++ {
			subCmd := vv.Type().Field(j).Name
			subCmd = strings.ToLower(subCmd)
			value := vv.Field(j)
			fmt.Printf("set %-8s %-15s %v\n", cmd, subCmd, value)
		}
	}
}

func applyTag(val string, typ string) string {
	switch {
	case strings.Contains(typ, "lower"):
		return strings.ToLower(val)
	case strings.Contains(typ, "upper"):
		return strings.ToUpper(val)
	default:
		return val
	}

}

// helpSet shows set command
func helpSet() {
	println(`
          usage:
               set command option value
          example:
               set ping timeout 2s
	`)

}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"os"
)

func exitIfError(err error) {
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(2)
	}
}
<commit_msg>move errors<commit_after>package cli

import (
	"fmt"
	"os"
)

func exitIfError(err error) {
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(2)
	}
}

// failf prints to standard error a formatted error and usage message and
// returns the error.
func (cmd *CLI) failf(format string, a ...interface{}) error {
	err := fmt.Errorf(format, a...)
	fmt.Fprintln(cmd.ErrOutput(), err)
	fmt.Fprintln(cmd.ErrOutput(), "")
	cmd.Usage()
	return err
}

func (cmd *CLI) errf(format string, a ...interface{}) {
	cmd.handleErr(cmd.failf(format, a...))
}

func (cmd *CLI) panicf(format string, a ...interface{}) {
	panic(cmd.failf(format, a...))
}

func (cmd *CLI) handleErr(err error) {
	if err != nil {
		switch cmd.ErrorHandling {
		case ExitOnError:
			os.Exit(2)
		case PanicOnError:
			panic(err)
		}
	}
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"os"

	"github.com/vektra/cypress/router"
)

type Router struct {
	ConfigFile string `short:"c" long:"config" description:"path to config file"`
}

func (rt *Router) Execute(args []string) error {
	r := router.NewRouter()

	f, err := os.Open(rt.ConfigFile)
	if err != nil {
		return err
	}

	defer f.Close()

	err = r.LoadConfig(f)
	if err != nil {
		return err
	}

	err = r.Open()
	if err != nil {
		return err
	}

	defer r.Close()

	fmt.Printf("Router loaded and running\n%d routes active\n", len(r.Routes()))

	c := make(chan bool)

	Lifecycle.OnShutdown(func() {
		c <- true
	})

	<-c

	return nil
}

func init() {
	addCommand("router", "Route streams", "Route streams based on a config", &Router{})
}
<commit_msg>Properly cleanup router<commit_after>package cli

import (
	"fmt"
	"os"

	"github.com/vektra/cypress/router"
)

type Router struct {
	ConfigFile string `short:"c" long:"config" description:"path to config file"`
}

func (rt *Router) Execute(args []string) error {
	r := router.NewRouter()

	f, err := os.Open(rt.ConfigFile)
	if err != nil {
		return err
	}

	defer f.Close()

	err = r.LoadConfig(f)
	if err != nil {
		return err
	}

	err = r.Open()
	if err != nil {
		return err
	}

	fmt.Printf("Router loaded and running\n%d routes active\n", len(r.Routes()))

	Lifecycle.OnShutdown(func() {
		r.Close()
	})

	select {}

	return nil
}

func init() {
	addCommand("router", "Route streams", "Route streams based on a config", &Router{})
}
<|endoftext|>


<commit_before>// This is free and unencumbered software released into the public
// domain.  For more information, see <http://unlicense.org> or the
// accompanying UNLICENSE file.

package commander

import (
	"time"

	"github.com/nelsam/gxui"
	"github.com/nelsam/gxui/mixins"
)

const maxStatusAge = 5 * time.Second

var (
	cmdColor = gxui.Color{
		R: 0.3,
		G: 0.3,
		B: 0.6,
		A: 1,
	}
	displayColor = gxui.Color{
		R: 0.3,
		G: 1,
		B: 0.6,
		A: 1,
	}

	ColorErr = gxui.Color{
		R: 1,
		G: 0.2,
		B: 0,
		A: 1,
	}
	ColorWarn = gxui.Color{
		R: 0.8,
		G: 0.7,
		B: 0.1,
		A: 1,
	}
	ColorInfo = gxui.Color{
		R: 0.1,
		G: 0.8,
		B: 0,
		A: 1,
	}
)

// Completer is a type which defines when a gxui.KeyboardEvent
// completes an action.  Types returned from
// "../commands".Command.Next() may implement this if they don't want
// to immediately complete when the "enter" key is pressed.
type Completer interface {
	// Complete returns whether or not the key signals a completion of
	// the input.
	Complete(gxui.KeyboardEvent) bool
}

// A Command is any command that can be run by the commandBox.  Commands
// have many optional interfaces they can implement - see other
// interfaces in this package for what else can be implemented by a
// command.
type Command interface {
	// Name returns the name of the command
	Name() string

	// Menu returns the name of the menu that the command should be
	// displayed under.
	Menu() string
}

// A Starter is a type of command which needs to initialize itself
// whenever the command is started.
type Starter interface {
	// Start starts the command.  The element that the command is
	// targeting will be passed in as target.  If the returned
	// status element is non-nil, it will be displayed as an
	// element to display the current status of the command to
	// the user.
	Start(target gxui.Control) (status gxui.Control)
}

// An InputQueue is a type of command which needs to read user input
// before executing a command.
type InputQueue interface {
	// Next returns the next element for reading user input. By
	// default, this is called every time the commander receives a
	// gxui.KeyEnter event in KeyPress.  If there are situations where
	// this is not the desired behavior, the returned gxui.Focusable
	// can consume the gxui.KeyboardEvent.  If the input element has
	// other keyboard events that would trigger completion, it can
	// implement Completer, which will allow it to define when it
	// is complete.
	//
	// Next will continue to be called until it returns nil, at which
	// point the command is assumed to be done.
	Next() gxui.Focusable
}

// An Executor is a type of command that needs to execute some
// operation on one or more elements.  It will continue to be
// called for every element currently in the UI until it returns
// true for consume.
//
// If executed is never returned as true, an error will be
// displayed to the user stating that the command never
// executed.
type Executor interface {
	Exec(interface{}) (executed, consume bool)
}

// A Statuser is a type that needs to display its status after being
// run.  The commands should use their discretion for status colors,
// but colors for some common message types are exported by this
// package to keep things consistent.
type Statuser interface {
	// Status returns the element to display for the command's status.
	// The element will be removed after some time.
	Status() gxui.Control
}

// ColorSetter is a type that can have its color set.
type ColorSetter interface {
	// SetColor is called when a command element is displayed, so
	// that it matches the color theme of the commander.
	SetColor(gxui.Color)
}

type commandBox struct {
	mixins.LinearLayout

	driver     gxui.Driver
	controller Controller

	label   gxui.Label
	current Command
	display gxui.Control
	input   gxui.Focusable
	status  gxui.Control

	statusTimer *time.Timer
}

func newCommandBox(driver gxui.Driver, theme gxui.Theme, controller Controller) *commandBox {
	box := &commandBox{
		driver:     driver,
		controller: controller,
	}

	box.label = theme.CreateLabel()
	box.label.SetColor(cmdColor)

	box.LinearLayout.Init(box, theme)
	box.SetDirection(gxui.LeftToRight)
	box.AddChild(box.label)
	box.Clear()
	return box
}

func (b *commandBox) Finish() {
	defer b.controller.Editor().Focus()
	statuser, ok := b.current.(Statuser)
	if !ok {
		b.Clear()
		return
	}
	b.status = statuser.Status()
	if b.status == nil {
		b.Clear()
		return
	}
	b.clearDisplay()
	b.clearInput()
	b.AddChild(b.status)
	b.statusTimer = time.AfterFunc(maxStatusAge, func() {
		b.driver.CallSync(func() {
			b.Clear()
		})
	})
}

func (b *commandBox) Clear() {
	b.label.SetText("none")
	b.clearDisplay()
	b.clearInput()
	b.clearStatus()
	b.current = nil
}

func (b *commandBox) Run(command Command) (needsInput bool) {
	if b.statusTimer != nil {
		b.statusTimer.Stop()
	}
	b.current = command

	b.label.SetText(b.current.Name())
	b.startCurrent()
	return b.nextInput()
}

func (b *commandBox) startCurrent() {
	starter, ok := b.current.(Starter)
	if !ok {
		return
	}
	b.display = starter.Start(b.controller)
	if b.display == nil {
		return
	}
	if colorSetter, ok := b.display.(ColorSetter); ok {
		colorSetter.SetColor(displayColor)
	}
	b.AddChild(b.display)
}

func (b *commandBox) Current() Command {
	return b.current
}

func (b *commandBox) KeyPress(event gxui.KeyboardEvent) (consume bool) {
	if event.Modifier == 0 && event.Key == gxui.KeyEscape {
		return false
	}
	isEnter := event.Modifier == 0 && event.Key == gxui.KeyEnter
	complete := isEnter
	if completer, ok := b.input.(Completer); ok {
		complete = completer.Complete(event)
	}
	if complete {
		hasMore := b.nextInput()
		complete = !hasMore
	}
	return !(complete && isEnter)
}

func (b *commandBox) HasFocus() bool {
	if b.input == nil {
		return false
	}
	return b.input.HasFocus()
}

func (b *commandBox) clearDisplay() {
	if b.display == nil {
		return
	}
	b.RemoveChild(b.display)
	b.display = nil
}

func (b *commandBox) clearInput() {
	if b.input == nil {
		return
	}
	b.RemoveChild(b.input)
	b.input = nil
}

func (b *commandBox) clearStatus() {
	if b.status == nil {
		return
	}
	b.RemoveChild(b.status)
	b.status = nil
}

func (b *commandBox) nextInput() (more bool) {
	queue, ok := b.current.(InputQueue)
	if !ok {
		return false
	}
	next := queue.Next()
	if next == nil {
		return false
	}
	b.clearInput()
	b.input = next
	b.AddChild(b.input)
	gxui.SetFocus(b.input)
	return true
}
<commit_msg>Clear command box before running a new command<commit_after>// This is free and unencumbered software released into the public
// domain.  For more information, see <http://unlicense.org> or the
// accompanying UNLICENSE file.

package commander

import (
	"time"

	"github.com/nelsam/gxui"
	"github.com/nelsam/gxui/mixins"
)

const maxStatusAge = 5 * time.Second

var (
	cmdColor = gxui.Color{
		R: 0.3,
		G: 0.3,
		B: 0.6,
		A: 1,
	}
	displayColor = gxui.Color{
		R: 0.3,
		G: 1,
		B: 0.6,
		A: 1,
	}

	ColorErr = gxui.Color{
		R: 1,
		G: 0.2,
		B: 0,
		A: 1,
	}
	ColorWarn = gxui.Color{
		R: 0.8,
		G: 0.7,
		B: 0.1,
		A: 1,
	}
	ColorInfo = gxui.Color{
		R: 0.1,
		G: 0.8,
		B: 0,
		A: 1,
	}
)

// Completer is a type which defines when a gxui.KeyboardEvent
// completes an action.  Types returned from
// "../commands".Command.Next() may implement this if they don't want
// to immediately complete when the "enter" key is pressed.
type Completer interface {
	// Complete returns whether or not the key signals a completion of
	// the input.
	Complete(gxui.KeyboardEvent) bool
}

// A Command is any command that can be run by the commandBox.  Commands
// have many optional interfaces they can implement - see other
// interfaces in this package for what else can be implemented by a
// command.
type Command interface {
	// Name returns the name of the command
	Name() string

	// Menu returns the name of the menu that the command should be
	// displayed under.
	Menu() string
}

// A Starter is a type of command which needs to initialize itself
// whenever the command is started.
type Starter interface {
	// Start starts the command.  The element that the command is
	// targeting will be passed in as target.  If the returned
	// status element is non-nil, it will be displayed as an
	// element to display the current status of the command to
	// the user.
	Start(target gxui.Control) (status gxui.Control)
}

// An InputQueue is a type of command which needs to read user input
// before executing a command.
type InputQueue interface {
	// Next returns the next element for reading user input. By
	// default, this is called every time the commander receives a
	// gxui.KeyEnter event in KeyPress.  If there are situations where
	// this is not the desired behavior, the returned gxui.Focusable
	// can consume the gxui.KeyboardEvent.  If the input element has
	// other keyboard events that would trigger completion, it can
	// implement Completer, which will allow it to define when it
	// is complete.
	//
	// Next will continue to be called until it returns nil, at which
	// point the command is assumed to be done.
	Next() gxui.Focusable
}

// An Executor is a type of command that needs to execute some
// operation on one or more elements.  It will continue to be
// called for every element currently in the UI until it returns
// true for consume.
//
// If executed is never returned as true, an error will be
// displayed to the user stating that the command never
// executed.
type Executor interface {
	Exec(interface{}) (executed, consume bool)
}

// A Statuser is a type that needs to display its status after being
// run.  The commands should use their discretion for status colors,
// but colors for some common message types are exported by this
// package to keep things consistent.
type Statuser interface {
	// Status returns the element to display for the command's status.
	// The element will be removed after some time.
	Status() gxui.Control
}

// ColorSetter is a type that can have its color set.
type ColorSetter interface {
	// SetColor is called when a command element is displayed, so
	// that it matches the color theme of the commander.
	SetColor(gxui.Color)
}

type commandBox struct {
	mixins.LinearLayout

	driver     gxui.Driver
	controller Controller

	label   gxui.Label
	current Command
	display gxui.Control
	input   gxui.Focusable
	status  gxui.Control

	statusTimer *time.Timer
}

func newCommandBox(driver gxui.Driver, theme gxui.Theme, controller Controller) *commandBox {
	box := &commandBox{
		driver:     driver,
		controller: controller,
	}

	box.label = theme.CreateLabel()
	box.label.SetColor(cmdColor)

	box.LinearLayout.Init(box, theme)
	box.SetDirection(gxui.LeftToRight)
	box.AddChild(box.label)
	box.Clear()
	return box
}

func (b *commandBox) Finish() {
	defer b.controller.Editor().Focus()
	statuser, ok := b.current.(Statuser)
	if !ok {
		b.Clear()
		return
	}
	b.status = statuser.Status()
	if b.status == nil {
		b.Clear()
		return
	}
	b.clearDisplay()
	b.clearInput()
	b.AddChild(b.status)
	b.statusTimer = time.AfterFunc(maxStatusAge, func() {
		b.driver.CallSync(func() {
			b.Clear()
		})
	})
}

func (b *commandBox) Clear() {
	b.label.SetText("none")
	b.clearDisplay()
	b.clearInput()
	b.clearStatus()
	b.current = nil
}

func (b *commandBox) Run(command Command) (needsInput bool) {
	b.Clear()
	if b.statusTimer != nil {
		b.statusTimer.Stop()
	}
	b.current = command

	b.label.SetText(b.current.Name())
	b.startCurrent()
	return b.nextInput()
}

func (b *commandBox) startCurrent() {
	starter, ok := b.current.(Starter)
	if !ok {
		return
	}
	b.display = starter.Start(b.controller)
	if b.display == nil {
		return
	}
	if colorSetter, ok := b.display.(ColorSetter); ok {
		colorSetter.SetColor(displayColor)
	}
	b.AddChild(b.display)
}

func (b *commandBox) Current() Command {
	return b.current
}

func (b *commandBox) KeyPress(event gxui.KeyboardEvent) (consume bool) {
	if event.Modifier == 0 && event.Key == gxui.KeyEscape {
		return false
	}
	isEnter := event.Modifier == 0 && event.Key == gxui.KeyEnter
	complete := isEnter
	if completer, ok := b.input.(Completer); ok {
		complete = completer.Complete(event)
	}
	if complete {
		hasMore := b.nextInput()
		complete = !hasMore
	}
	return !(complete && isEnter)
}

func (b *commandBox) HasFocus() bool {
	if b.input == nil {
		return false
	}
	return b.input.HasFocus()
}

func (b *commandBox) clearDisplay() {
	if b.display == nil {
		return
	}
	b.RemoveChild(b.display)
	b.display = nil
}

func (b *commandBox) clearInput() {
	if b.input == nil {
		return
	}
	b.RemoveChild(b.input)
	b.input = nil
}

func (b *commandBox) clearStatus() {
	if b.status == nil {
		return
	}
	b.RemoveChild(b.status)
	b.status = nil
}

func (b *commandBox) nextInput() (more bool) {
	queue, ok := b.current.(InputQueue)
	if !ok {
		return false
	}
	next := queue.Next()
	if next == nil {
		return false
	}
	b.clearInput()
	b.input = next
	b.AddChild(b.input)
	gxui.SetFocus(b.input)
	return true
}
<|endoftext|>


<commit_before>package knuckles

import (
	"fmt"
	"net/http"
	"net/url"
	"time"
)

type Backend struct {
	Name          string
	Endpoint      string
	CheckInterval time.Duration
	CheckURL      *url.URL
	Alive         bool
	quitChan      chan bool
	notifyChan    chan BackendStatus
}

type BackendStatus struct {
	Name  string
	Alive bool
}

type BackendSettings struct {
	Endpoint      string
	CheckUrl      string
	CheckInterval time.Duration
	Updates       chan BackendStatus
}

func NewBackend(name string, settings BackendSettings) (*Backend, error) {
	var err error
	be := new(Backend)
	be.Name = name
	be.Alive = false
	be.quitChan = make(chan bool)

	if settings.CheckInterval == 0 ||
		settings.CheckUrl == "" ||
		settings.Endpoint == "" ||
		settings.Updates == nil {
		return nil, fmt.Errorf("Missing parameters")
	}
	be.Endpoint = settings.Endpoint
	be.notifyChan = settings.Updates
	be.CheckInterval = settings.CheckInterval

	be.CheckURL, err = url.Parse(settings.CheckUrl)
	return be, err
}

func (self *Backend) Start() {
	run := true

	self.CheckHealth()
	for run {
		select {
		case <-self.quitChan:
			run = false
		case <-time.After(self.CheckInterval):
			self.CheckHealth()
		}
	}
	self.quitChan <- true
}

func (self *Backend) Stop() {
	self.quitChan <- true
	<-self.quitChan
	close(self.quitChan)
}

func (self *Backend) CheckHealth() {
	current := self.Alive

	client := &http.Client{}

	req, _ := http.NewRequest("GET", self.CheckURL.String(), nil)
	req.Close = true

	resp, err := client.Do(req)
	if err == nil {
		defer resp.Body.Close()
	}

	if err != nil || (resp.StatusCode >= 400) {
		self.Alive = false
	} else {
		self.Alive = true
	}

	if current != self.Alive {
		status := BackendStatus{
			Name:  self.Name,
			Alive: self.Alive,
		}
		self.notifyChan <- status
	}
}
<commit_msg>do not follow redirects on health checks<commit_after>package knuckles

import (
	"fmt"
	"net/http"
	"net/url"
	"time"
)

type Backend struct {
	Name          string
	Endpoint      string
	CheckInterval time.Duration
	CheckURL      *url.URL
	Alive         bool
	quitChan      chan bool
	notifyChan    chan BackendStatus
}

type BackendStatus struct {
	Name  string
	Alive bool
}

type BackendSettings struct {
	Endpoint      string
	CheckUrl      string
	CheckInterval time.Duration
	Updates       chan BackendStatus
}

func NewBackend(name string, settings BackendSettings) (*Backend, error) {
	var err error
	be := new(Backend)
	be.Name = name
	be.Alive = false
	be.quitChan = make(chan bool)

	if settings.CheckInterval == 0 ||
		settings.CheckUrl == "" ||
		settings.Endpoint == "" ||
		settings.Updates == nil {
		return nil, fmt.Errorf("Missing parameters")
	}
	be.Endpoint = settings.Endpoint
	be.notifyChan = settings.Updates
	be.CheckInterval = settings.CheckInterval

	be.CheckURL, err = url.Parse(settings.CheckUrl)
	return be, err
}

func (self *Backend) Start() {
	run := true

	self.CheckHealth()
	for run {
		select {
		case <-self.quitChan:
			run = false
		case <-time.After(self.CheckInterval):
			self.CheckHealth()
		}
	}
	self.quitChan <- true
}

func (self *Backend) Stop() {
	self.quitChan <- true
	<-self.quitChan
	close(self.quitChan)
}

func (self *Backend) CheckHealth() {
	current := self.Alive

	tr := &http.Transport{}

	req, _ := http.NewRequest("GET", self.CheckURL.String(), nil)
	req.Close = true

	resp, err := tr.RoundTrip(req)
	if err == nil {
		defer resp.Body.Close()
	}

	if err != nil {
		self.Alive = false
	} else {
		self.Alive = true
	}

	if current != self.Alive {
		status := BackendStatus{
			Name:  self.Name,
			Alive: self.Alive,
		}
		self.notifyChan <- status
	}
}
<|endoftext|>