<commit_before>package epp

// <epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
//   <command>
//     <check>
//     <domain:check xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
//       <domain:name>dmnr-test-1234.com</domain:name>
//     </domain:check>
//     </check>
//     <clTRID>ABC-12345</clTRID>
//   </command>
// </epp>

// <resData>
//  <domain:chkData xmlns:domain="urn:ietf:params:xml:ns:domain-1.0" xsi:schemaLocation="urn:ietf:params:xml:ns:domain-1.0 domain-1.0.xsd">
//   <domain:cd>
//    <domain:name avail="0">ace.pizza</domain:name>
//    <domain:reason>Premium Domain Name</domain:reason>
//   </domain:cd>
//  </domain:chkData>
// </resData>

// CheckDomain queries the EPP server for the availability status of one or more domains.
func (c *Conn) CheckDomain(domains ...string) (*DomainCheck, error) {
	req := message{
		Command: &command{
			Check: &check{
				DomainCheck: &domainCheck{
					Domains: domains,
				},
			},
			TxnID: c.id(),
		},
	}
	err := c.writeMessage(&req)
	if err != nil {
		return nil, err
	}
	msg := message{}
	err = c.readMessage(&msg)
	if err != nil {
		return nil, err
	}
	res := msg.Response
	if res == nil || res.ResponseData == nil || res.ResponseData.DomainCheckData == nil {
		return nil, ErrResponseMalformed
	}
	dc := DomainCheck(*res.ResponseData.DomainCheckData)
	return &dc, nil
}

// FIXME: backwards compatibility
type DomainCheck domainCheckData
<commit_msg>golint again<commit_after>package epp

// <epp xmlns="urn:ietf:params:xml:ns:epp-1.0">
//   <command>
//     <check>
//     <domain:check xmlns:domain="urn:ietf:params:xml:ns:domain-1.0">
//       <domain:name>dmnr-test-1234.com</domain:name>
//     </domain:check>
//     </check>
//     <clTRID>ABC-12345</clTRID>
//   </command>
// </epp>

// <resData>
//  <domain:chkData xmlns:domain="urn:ietf:params:xml:ns:domain-1.0" xsi:schemaLocation="urn:ietf:params:xml:ns:domain-1.0 domain-1.0.xsd">
//   <domain:cd>
//    <domain:name avail="0">ace.pizza</domain:name>
//    <domain:reason>Premium Domain Name</domain:reason>
//   </domain:cd>
//  </domain:chkData>
// </resData>

// CheckDomain queries the EPP server for the availability status of one or more domains.
func (c *Conn) CheckDomain(domains ...string) (*DomainCheck, error) {
	req := message{
		Command: &command{
			Check: &check{
				DomainCheck: &domainCheck{
					Domains: domains,
				},
			},
			TxnID: c.id(),
		},
	}
	err := c.writeMessage(&req)
	if err != nil {
		return nil, err
	}
	msg := message{}
	err = c.readMessage(&msg)
	if err != nil {
		return nil, err
	}
	res := msg.Response
	if res == nil || res.ResponseData == nil || res.ResponseData.DomainCheckData == nil {
		return nil, ErrResponseMalformed
	}
	dc := DomainCheck(*res.ResponseData.DomainCheckData)
	return &dc, nil
}

// DomainCheck exists for backwards compatibility.
// FIXME: remove/improve this.
type DomainCheck domainCheckData
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"go/build"
	"os"
	"path/filepath"
	"strings"
)

type stringSet map[string]struct{}

func (ss stringSet) add(val string) { ss[val] = struct{}{} }
func (ss stringSet) has(val string) bool {
	_, ok := ss[val]
	return ok
}

// compute the set of dependencies for a list of packages. The packages must
// already be present in toolDirPath.
func dependencies(pkgs []string) stringSet {
	deps := stringSet{}

	buildCtx := build.Default
	buildCtx.GOPATH = toolDirPath

	var resolve func(string, []string)
	resolve = func(parent string, pkgs []string) {
		for _, pkg := range pkgs {
			if !strings.Contains(pkg, ".") {
				continue
			}

			p, err := buildCtx.Import(pkg, filepath.Join(toolDirPath, "src", parent), 0)
			if err != nil {
				fatal(fmt.Sprintf("couldn't import package %q", pkg), err)
			}

			if deps.has(p.ImportPath) {
				continue
			}

			deps.add(p.ImportPath)
			resolve(p.ImportPath, p.Imports)
		}
	}

	resolve("", pkgs)

	return deps
}

// Remove unused files and unused packages from toolDirPath.
func clean(pkgs []string) {
	deps := dependencies(pkgs)
	base := filepath.Join(toolDirPath, "src")

	// Resolve any symlinks in the packages to keep, because we're going
	// to walk through the file system, so we need to trim stuff by
	// _filename_.
	for pkgPath := range deps {
		fullPath := filepath.Join(base, pkgPath)
		resolved, err := filepath.EvalSymlinks(fullPath)
		if err != nil {
			fatal(fmt.Sprintf("couldn't eval symlinks in %q", pkgPath), err)
		}
		// Undo the filepath.Join from above
		pkgPath, err = filepath.Rel(base, resolved)
		if err != nil {
			fatal(fmt.Sprintf("couldn't eval symlinks in %q", pkgPath), err)
		}
		deps.add(pkgPath)
	}

	var toDelete []string
	err := filepath.Walk(base, func(path string, info os.FileInfo, err error) error {
		// Bubble up errors
		if err != nil {
			return err
		}

		// Skip the root directory
		if base == path {
			return nil
		}

		// Get the package directory
		pkg, err := filepath.Rel(base, path)
		if err != nil {
			return err
		}

		// Delete files in packages that aren't in packages we need to build the
		// tools, and any non-go, non-legal files in packages we *do* need.
		if info.Mode().IsRegular() {
			pkg = filepath.Dir(pkg)
			if !(deps.has(pkg) && keepFile(path)) {
				toDelete = append(toDelete, path)
			}
			return nil
		}

		// If the path is a directory that's specially marked for preservation, keep
		// it and all its contents.
		if info.IsDir() && preserveDirectory(path) {
			return filepath.SkipDir
		}

		// If the folder is a kept package or a parent, don't delete it and keep recursing
		for p := range deps {
			if strings.HasPrefix(p, pkg) {
				return nil
			}
		}

		// Otherwise this is a package that isn't imported at all. Delete it and stop recursing
		toDelete = append(toDelete, path)
		return filepath.SkipDir
	})

	if err != nil {
		fatal("unable to clean _tools", err)
	}

	for _, path := range toDelete {
		err = os.RemoveAll(path)
		if err != nil {
			fatal("unable to remove file or directory", err)
		}
	}
}

func keepFile(filename string) bool {
	if strings.HasSuffix(filename, "_test.go") {
		return false
	}
	if strings.HasSuffix(filename, ".go") {
		return true
	}
	if strings.HasSuffix(filename, ".s") {
		return true
	}
	if isLegalFile(filename) {
		return true
	}
	return false
}

var commonLegalFilePrefixes = []string{
	"licence", // UK spelling
	"license", // US spelling
	"copying",
	"unlicense",
	"copyright",
	"copyleft",
	"authors",
	"contributors",
	"readme", // often has a license inline
}

var commonLegalFileSubstrings = []string{
	"legal",
	"notice",
	"disclaimer",
	"patent",
	"third-party",
	"thirdparty",
}

func isLegalFile(filename string) bool {
	base := strings.ToLower(filepath.Base(filename))
	for _, p := range commonLegalFilePrefixes {
		if strings.HasPrefix(base, p) {
			return true
		}
	}
	for _, s := range commonLegalFileSubstrings {
		if strings.Contains(base, s) {
			return true
		}
	}
	return false
}

// List of directories that should be completely preserved if they are present.
var preservedDirectories = []string{
	// gometalinter vendors its own linters and relies on this directory's
	// existence. See issue #7.
	filepath.Join("github.com", "alecthomas", "gometalinter", "_linters"),
}

// checks whether path is in the list of preserved directories.
func preserveDirectory(path string) bool {
	for _, d := range preservedDirectories {
		if strings.HasSuffix(path, d) {
			return true
		}
	}
	return false
}
<commit_msg>Fix sqlboiler (#25)<commit_after>package main

import (
	"fmt"
	"go/build"
	"os"
	"path/filepath"
	"strings"
)

type stringSet map[string]struct{}

func (ss stringSet) add(val string) { ss[val] = struct{}{} }
func (ss stringSet) has(val string) bool {
	_, ok := ss[val]
	return ok
}

// compute the set of dependencies for a list of packages. The packages must
// already be present in toolDirPath.
func dependencies(pkgs []string) stringSet {
	deps := stringSet{}

	buildCtx := build.Default
	buildCtx.GOPATH = toolDirPath

	var resolve func(string, []string)
	resolve = func(parent string, pkgs []string) {
		for _, pkg := range pkgs {
			if !strings.Contains(pkg, ".") {
				continue
			}

			p, err := buildCtx.Import(pkg, filepath.Join(toolDirPath, "src", parent), 0)
			if err != nil {
				fatal(fmt.Sprintf("couldn't import package %q", pkg), err)
			}

			if deps.has(p.ImportPath) {
				continue
			}

			deps.add(p.ImportPath)
			resolve(p.ImportPath, p.Imports)
		}
	}

	resolve("", pkgs)

	return deps
}

// Remove unused files and unused packages from toolDirPath.
func clean(pkgs []string) {
	deps := dependencies(pkgs)
	base := filepath.Join(toolDirPath, "src")

	// Resolve any symlinks in the packages to keep, because we're going
	// to walk through the file system, so we need to trim stuff by
	// _filename_.
	for pkgPath := range deps {
		fullPath := filepath.Join(base, pkgPath)
		resolved, err := filepath.EvalSymlinks(fullPath)
		if err != nil {
			fatal(fmt.Sprintf("couldn't eval symlinks in %q", pkgPath), err)
		}
		// Undo the filepath.Join from above
		pkgPath, err = filepath.Rel(base, resolved)
		if err != nil {
			fatal(fmt.Sprintf("couldn't eval symlinks in %q", pkgPath), err)
		}
		deps.add(pkgPath)
	}

	var toDelete []string
	err := filepath.Walk(base, func(path string, info os.FileInfo, err error) error {
		// Bubble up errors
		if err != nil {
			return err
		}

		// Skip the root directory
		if base == path {
			return nil
		}

		// Get the package directory
		pkg, err := filepath.Rel(base, path)
		if err != nil {
			return err
		}

		// Delete files in packages that aren't in packages we need to build the
		// tools, and any non-go, non-legal files in packages we *do* need.
		if info.Mode().IsRegular() {
			pkg = filepath.Dir(pkg)
			if !(deps.has(pkg) && keepFile(path)) {
				toDelete = append(toDelete, path)
			}
			return nil
		}

		// If the path is a directory that's specially marked for preservation, keep
		// it and all its contents.
		if info.IsDir() && preserveDirectory(path) {
			return filepath.SkipDir
		}

		// If the folder is a kept package or a parent, don't delete it and keep recursing
		for p := range deps {
			if strings.HasPrefix(p, pkg) {
				return nil
			}
		}

		// Otherwise this is a package that isn't imported at all. Delete it and stop recursing
		toDelete = append(toDelete, path)
		return filepath.SkipDir
	})

	if err != nil {
		fatal("unable to clean _tools", err)
	}

	for _, path := range toDelete {
		err = os.RemoveAll(path)
		if err != nil {
			fatal("unable to remove file or directory", err)
		}
	}
}

func keepFile(filename string) bool {
	if strings.HasSuffix(filename, "_test.go") {
		return false
	}
	if strings.HasSuffix(filename, ".go") {
		return true
	}
	if strings.HasSuffix(filename, ".s") {
		return true
	}
	if isLegalFile(filename) {
		return true
	}
	return false
}

var commonLegalFilePrefixes = []string{
	"licence", // UK spelling
	"license", // US spelling
	"copying",
	"unlicense",
	"copyright",
	"copyleft",
	"authors",
	"contributors",
	"readme", // often has a license inline
}

var commonLegalFileSubstrings = []string{
	"legal",
	"notice",
	"disclaimer",
	"patent",
	"third-party",
	"thirdparty",
}

func isLegalFile(filename string) bool {
	base := strings.ToLower(filepath.Base(filename))
	for _, p := range commonLegalFilePrefixes {
		if strings.HasPrefix(base, p) {
			return true
		}
	}
	for _, s := range commonLegalFileSubstrings {
		if strings.Contains(base, s) {
			return true
		}
	}
	return false
}

// List of directories that should be completely preserved if they are present.
var preservedDirectories = []string{
	// gometalinter vendors its own linters and relies on this directory's
	// existence. See issue #7.
	filepath.Join("github.com", "alecthomas", "gometalinter", "_linters"),

	// sqlboiler requires template files to exist at runtime. See pull request #25.
	filepath.Join("github.com", "vattle", "sqlboiler", "templates"),
	filepath.Join("github.com", "vattle", "sqlboiler", "templates_test"),
}

// checks whether path is in the list of preserved directories.
func preserveDirectory(path string) bool {
	for _, d := range preservedDirectories {
		if strings.HasSuffix(path, d) {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>// clock counts down to or up from a target time.
package main

import (
	"fmt"
	"os"
	"time"
)

func main() {
	target := time.Date(2016, 6, 16, 0, 0, 0, 0, time.UTC)
	motto := "Simply Go"
	printTargetTime(target, motto)
	exitOnEnterKey()

	var previous time.Time
	for {
		now := time.Now()
		now = now.Add(time.Duration(-now.Nanosecond())) // truncate to second
		if now != previous {
			previous = now
			remaining := target.Sub(now)
			printTimeRemaining(now, remaining)
		}
		time.Sleep(50 * time.Millisecond)
	}
}

func exitOnEnterKey() {
	go func() {
		buf := make([]byte, 1)
		_, _ = os.Stdin.Read(buf)
		os.Exit(0)
	}()
}

const (
	indent          = "\t"
	highlight_start = "\x1b[1;35m"
	highlight_end   = "\x1b[0m"
)

func printTargetTime(target time.Time, motto string) {
	fmt.Print(indent, highlight_start, motto, highlight_end, "\n")
	fmt.Print(indent, target.Format(time.UnixDate), "\n")
}

func printTimeRemaining(now time.Time, remaining time.Duration) {
	var sign string
	if remaining > 0 {
		sign = "-" // countdown is "T minus..."
	} else {
		sign = "+" // count up is "T plus..."
		remaining = -remaining
	}

	var days int
	if remaining >= 24*time.Hour {
		days = int(remaining / (24 * time.Hour))
		remaining = remaining % (24 * time.Hour)
	}

	fmt.Print(indent, now.Format(time.UnixDate), "  ", sign)
	if days > 0 {
		fmt.Print(days, "d")
	}
	fmt.Print(remaining, "          \r")
}
<commit_msg>Solstice<commit_after>// clock counts down to or up from a target time.
package main

import (
	"fmt"
	"os"
	"time"
)

func main() {
	target := time.Date(2016, 6, 20, 18, 34, 9, 0, time.Local)
	motto := "Simply Go"
	printTargetTime(target, motto)
	exitOnEnterKey()

	var previous time.Time
	for {
		now := time.Now()
		now = now.Add(time.Duration(-now.Nanosecond())) // truncate to second
		if now != previous {
			previous = now
			remaining := target.Sub(now)
			printTimeRemaining(now, remaining)
		}
		time.Sleep(50 * time.Millisecond)
	}
}

func exitOnEnterKey() {
	go func() {
		buf := make([]byte, 1)
		_, _ = os.Stdin.Read(buf)
		os.Exit(0)
	}()
}

const (
	indent          = "\t"
	highlight_start = "\x1b[1;35m"
	highlight_end   = "\x1b[0m"
)

func printTargetTime(target time.Time, motto string) {
	fmt.Print(indent, highlight_start, motto, highlight_end, "\n")
	fmt.Print(indent, target.Format(time.UnixDate), "\n")
}

func printTimeRemaining(now time.Time, remaining time.Duration) {
	var sign string
	if remaining > 0 {
		sign = "-" // countdown is "T minus..."
	} else {
		sign = "+" // count up is "T plus..."
		remaining = -remaining
	}

	var days int
	if remaining >= 24*time.Hour {
		days = int(remaining / (24 * time.Hour))
		remaining = remaining % (24 * time.Hour)
	}

	fmt.Print(indent, now.Format(time.UnixDate), "  ", sign)
	if days > 0 {
		fmt.Print(days, "d")
	}
	fmt.Print(remaining, "          \r")
}
<|endoftext|>


<commit_before>package gst

/*
#include <gst/gst.h>
static gboolean clock_time_is_valid(gint64 t) {
	return GST_CLOCK_TIME_IS_VALID(t);
}
*/
import "C"

type Clock struct {
	GstObj
}

func (c *Clock) g() *C.GstClock {
	return (*C.GstClock)(c.GetPtr())
}

func (c *Clock) AsClock() *Clock {
	return c
}

func ClockTimeIsValid(t int64) bool {
	return C.clock_time_is_valid((C.gint64)(t)) == 1
}
<commit_msg>add ClockTime definition & DurationOutput for ClockTime<commit_after>package gst

/*
#include <gst/gst.h>
static gboolean clock_time_is_valid(gint64 t) {
	return GST_CLOCK_TIME_IS_VALID(t);
}
static gchar* clock_duration_output(GstClockTime t) {
	return g_strdup_printf("%" GST_TIME_FORMAT, GST_TIME_ARGS(t));
}
*/
import "C"

type Clock struct {
	GstObj
}

func (c *Clock) g() *C.GstClock {
	return (*C.GstClock)(c.GetPtr())
}

func (c *Clock) AsClock() *Clock {
	return c
}

type ClockTime C.GstClockTime

func (t *ClockTime) g() *C.GstClockTime {
	return (*C.GstClockTime)(t)
}

func (t ClockTime) AsUint64() uint64 {
	return uint64(t)
}

func (t ClockTime) DurationOutput() string {
	return C.GoString((*C.char)(C.clock_duration_output(C.GstClockTime(t))))
}

func ClockTimeIsValid(t int64) bool {
	return C.clock_time_is_valid((C.gint64)(t)) == 1
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"os"
	"os/user"
	"path/filepath"
	"regexp"
	"strconv"

	"github.com/codegangsta/cli"
	"github.com/shenwei356/cnote/cnotedb"
)

var (
	funcs  map[string]func(c *cli.Context)
	DBFILE string
	notedb *cnotedb.NoteDB
)

func init() {
	// DBFILE
	usr, err := user.Current()
	if err != nil {
		os.Exit(1)
	}
	DBFILE = filepath.Join(usr.HomeDir, ".cnote")

	funcs = make(map[string]func(c *cli.Context))
	funcs["new"] = funNew
	funcs["del"] = funDel
	funcs["use"] = funUse
	funcs["list"] = funLs

	funcs["add"] = funAdd
	funcs["rm"] = funRm

	funcs["tag"] = funTag
	funcs["search"] = funSearch

	funcs["dump"] = funDump
	funcs["wipe"] = funWipe
	funcs["restore"] = funRestore
	funcs["import"] = funImport

}

func request_reply(message, reply string) (bool, error) {
	fmt.Printf(message, reply)

	reader := bufio.NewReader(os.Stdin)
	str, err := reader.ReadString('\n')
	if err != nil {
		return false, err
	}

	str = regexp.MustCompile(`[\r\n]`).ReplaceAllString(str, "")
	str = regexp.MustCompile(`^\s+|\s+$`).ReplaceAllString(str, "")
	if str != "yes" {
		fmt.Println("\ngiven up.")
		return false, nil
	}

	return true, nil
}

func getFunc(funcs map[string]func(c *cli.Context), name string) func(c *cli.Context) {
	if f, ok := funcs[name]; ok {
		return f
	} else {
		return func(c *cli.Context) {
			fmt.Printf("command %s not implemented\n", name)
		}
	}
}

func funLs(c *cli.Context) {
	if len(c.Args()) > 0 {
		fmt.Println("no arguments should be given.")
		return
	}

	for _, notename := range notedb.NotesList {

		// read note
		note, err := notedb.ReadNote(notename)
		if err != nil {
			fmt.Println(err)
			os.Exit(0)
		}

		fmt.Printf("note: %s\t(#. of items: %d, last update: %s).",
			notename, note.Sum, note.LastUpdate)
		if notedb.CurrentNote != nil &&
			notename == notedb.CurrentNote.NoteID {

			fmt.Printf(" (current note)")
		}
		fmt.Println()
	}
}

func funNew(c *cli.Context) {
	if len(c.Args()) == 0 {
		fmt.Println("note name needed.")
		return
	}
	if len(c.Args()) > 1 {
		fmt.Println("only one note name allowed.")
		return
	}

	notename := c.Args().First()

	err := notedb.NewNote(notename)
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Printf("note \"%s\" created.\n", notename)
	fmt.Printf("current note: \"%s\".\n", notename)
}

func funDel(c *cli.Context) {
	if len(c.Args()) == 0 {
		fmt.Println("note name needed.")
		return
	}
	notename := c.Args().First()

	note, err := notedb.ReadNote(notename)
	if err != nil {
		fmt.Println(err)
		os.Exit(0)
	}

	reply, err := request_reply(
		fmt.Sprintf("==============================================================\n"+
			" Attention, it will delete all the %d items of note \"%s\".\n"+
			"==============================================================\n",
			note.Sum, notename)+
			" Type \"%s\" to continue:",
		"yes")
	if err != nil {
		fmt.Println(err)
		os.Exit(0)
	}

	if reply == false {
		return
	}

	err = notedb.DeleteNote(notename)
	if err != nil {
		fmt.Println(err)
		os.Exit(0)
	}
}

func funUse(c *cli.Context) {
	if len(c.Args()) == 0 {
		fmt.Println("note name needed.")
		return
	}
	if len(c.Args()) > 1 {
		fmt.Println("only one note name allowed.")
		return
	}

	notename := c.Args().First()
	err := notedb.UseNote(notename)
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Printf("current note: \"%s\" (last update: %s).\n",
		notename, notedb.CurrentNote.LastUpdate)
}

func funAdd(c *cli.Context) {
	if len(c.Args()) != 2 {
		fmt.Println("tag and content needed.")
		return
	}

	item, err := notedb.AddNoteItem(c.Args()[0], c.Args()[1])
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Println(item)
}

func funRm(c *cli.Context) {
	if len(c.Args()) == 0 {
		fmt.Println("item ID needed.")
		return
	}

	for _, itemid := range c.Args() {

		itemid, err := strconv.Atoi(itemid)
		if err != nil {
			fmt.Println("item ID should be positive integer.")
			continue
		}

		// read item and print it, in case of misdeleteing
		item, err = notedb.ReadNoteItem(note, itemid)
		if err != nil {
			return err
		}

		err = notedb.RemoveNoteItem(notedb.CurrentNote, itemid)
		if err != nil {
			fmt.Println(err)
			continue
		}

		fmt.Println(item)
	}
}

func funTag(c *cli.Context) {
	// list all tags
	note, err := notedb.GetCurrentNote()
	if err != nil {
		fmt.Println(err)
		os.Exit(0)
	}

	if len(c.Args()) == 0 {
		for tag, taginfo := range note.Tags {
			fmt.Printf("tag: %s\t(#. of items: %d).\n", tag, len(taginfo))
		}
		return
	}

	items, err := notedb.ItemByTag(c.Args())
	if err != nil {
		fmt.Println(err)
		os.Exit(0)
	}

	for _, item := range items {
		fmt.Println(item)
	}
}

func funSearch(c *cli.Context) {
	if len(c.Args()) == 0 {
		fmt.Println("search keyword needed.")
		return
	}

	items, err := notedb.ItemByRegexp(c.Args())
	if err != nil {
		fmt.Println(err)
		os.Exit(0)
	}

	for _, item := range items {
		fmt.Println(item)
	}
}

func funDump(c *cli.Context) {
	if len(c.Args()) > 0 {
		fmt.Println("no arguments should be given.")
		return
	}

	err := notedb.Dump()
	if err != nil {
		fmt.Println(err)
		os.Exit(0)
	}
}

func funWipe(c *cli.Context) {
	if len(c.Args()) > 0 {
		fmt.Println("no arguments should be given.")
		return
	}

	reply, err := request_reply(
		"========================================\n"+
			" Attention, it will clear all the data.\n"+
			"========================================\n"+
			" Type \"%s\" to continue:",
		"yes")
	if err != nil {
		fmt.Println(err)
		os.Exit(0)
	}

	if reply == false {
		return
	}

	err = notedb.Wipe()
	if err != nil {
		fmt.Println(err)
		os.Exit(0)
	}
}

func funRestore(c *cli.Context) {
	if len(c.Args()) != 1 {
		fmt.Println("dumpped filename needed.")
		return
	}

	reply, err := request_reply(
		"========================================\n"+
			" Attention, it will clear all the data.\n"+
			"========================================\n"+
			" Type \"%s\" to continue:",
		"yes")
	if err != nil {
		fmt.Println(err)
		os.Exit(0)
	}

	if reply == false {
		return
	}

	err = notedb.Restore(c.Args().First())
	if err != nil {
		fmt.Println(err)
		os.Exit(0)
	}
}

func funImport(c *cli.Context) {
	if len(c.Args()) != 3 {
		fmt.Println("three arguments needed: <notename in your cnote>" +
			" <notename in dumpped note> <dumpped filename>.")
		return
	}
	notename, othernotename, filename := c.Args()[0], c.Args()[1], c.Args()[2]
	n, err := notedb.Import(notename, othernotename, filename)
	if err != nil {
		fmt.Println(err)
		os.Exit(0)
	}
	fmt.Printf("%d items imported into note \"%s\".\n", n, notename)
}

func main() {
	notedb = cnotedb.NewNoteDB(DBFILE)
	defer notedb.Close()

	app := cli.NewApp()
	app.Name = "cnote"
	app.Usage = "A command line note app. https://github.com/shenwei356/cnote"
	app.Version = "1.1 (2014-07-20)"
	app.Author = "Wei Shen"
	app.Email = "shenwei356@gmail.com"

	app.Commands = []cli.Command{
		{
			Name:   "new",
			Usage:  "Create a new note",
			Action: getFunc(funcs, "new"),
		},
		{
			Name:   "del",
			Usage:  "Delete a note",
			Action: getFunc(funcs, "del"),
		},
		{
			Name:   "use",
			Usage:  "Select a note",
			Action: getFunc(funcs, "use"),
		},
		{
			Name:      "list",
			ShortName: "ls",
			Usage:     "List all notes",
			Action:    getFunc(funcs, "list"),
		},
		{
			Name:   "add",
			Usage:  "Add a note item",
			Action: getFunc(funcs, "add"),
		},
		{
			Name:   "rm",
			Usage:  "Remove a note item",
			Action: getFunc(funcs, "rm"),
		},
		{
			Name:      "tag",
			ShortName: "t",
			Usage:     "List items by tags. List all tags if no arguments given",
			Action:    getFunc(funcs, "tag"),
		},
		{
			Name:      "search",
			ShortName: "s",
			Usage:     "Search items with regular expression",
			Action:    getFunc(funcs, "search"),
		},
		{
			Name:   "dump",
			Usage:  "Dump whole database, for backup or transfer",
			Action: getFunc(funcs, "dump"),
		},
		{
			Name:   "wipe",
			Usage:  "Attention! Wipe whole database",
			Action: getFunc(funcs, "wipe"),
		},
		{
			Name:   "restore",
			Usage:  "Wipe whole database, and restore from dumpped file",
			Action: getFunc(funcs, "restore"),
		},
		{
			Name:   "import",
			Usage:  "Import note items from dumpped data",
			Action: getFunc(funcs, "import"),
		},
	}

	app.Run(os.Args)
}
<commit_msg>fix bug that using os.Exit() would exit too early<commit_after>package main

import (
	"bufio"
	"fmt"
	"os"
	"os/user"
	"path/filepath"
	"regexp"
	"strconv"

	"github.com/codegangsta/cli"
	"github.com/shenwei356/cnote/cnotedb"
)

var (
	funcs  map[string]func(c *cli.Context)
	DBFILE string
	notedb *cnotedb.NoteDB
)

func init() {
	// DBFILE
	usr, err := user.Current()
	if err != nil {
		os.Exit(1)
	}
	DBFILE = filepath.Join(usr.HomeDir, ".cnote")

	funcs = make(map[string]func(c *cli.Context))
	funcs["new"] = funNew
	funcs["del"] = funDel
	funcs["use"] = funUse
	funcs["list"] = funLs

	funcs["add"] = funAdd
	funcs["rm"] = funRm

	funcs["tag"] = funTag
	funcs["search"] = funSearch

	funcs["dump"] = funDump
	funcs["wipe"] = funWipe
	funcs["restore"] = funRestore
	funcs["import"] = funImport

}

func request_reply(message, reply string) (bool, error) {
	fmt.Printf(message, reply)

	reader := bufio.NewReader(os.Stdin)
	str, err := reader.ReadString('\n')
	if err != nil {
		return false, err
	}

	str = regexp.MustCompile(`[\r\n]`).ReplaceAllString(str, "")
	str = regexp.MustCompile(`^\s+|\s+$`).ReplaceAllString(str, "")
	if str != "yes" {
		fmt.Println("\ngiven up.")
		return false, nil
	}

	return true, nil
}

func getFunc(funcs map[string]func(c *cli.Context), name string) func(c *cli.Context) {
	if f, ok := funcs[name]; ok {
		return f
	} else {
		return func(c *cli.Context) {
			fmt.Printf("command %s not implemented\n", name)
		}
	}
}

func funLs(c *cli.Context) {
	if len(c.Args()) > 0 {
		fmt.Println("no arguments should be given.")
		return
	}

	for _, notename := range notedb.NotesList {

		// read note
		note, err := notedb.ReadNote(notename)
		if err != nil {
			fmt.Println(err)
			return
		}

		fmt.Printf("note: %s\t(#. of items: %d, last update: %s).",
			notename, note.Sum, note.LastUpdate)
		if notedb.CurrentNote != nil &&
			notename == notedb.CurrentNote.NoteID {

			fmt.Printf(" (current note)")
		}
		fmt.Println()
	}
}

func funNew(c *cli.Context) {
	if len(c.Args()) == 0 {
		fmt.Println("note name needed.")
		return
	}
	if len(c.Args()) > 1 {
		fmt.Println("only one note name allowed.")
		return
	}

	notename := c.Args().First()

	err := notedb.NewNote(notename)
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Printf("note \"%s\" created.\n", notename)
	fmt.Printf("current note: \"%s\" (last update: %s).\n",
		notename, notedb.CurrentNote.LastUpdate)
}

func funDel(c *cli.Context) {
	if len(c.Args()) == 0 {
		fmt.Println("note name needed.")
		return
	}
	notename := c.Args().First()

	note, err := notedb.ReadNote(notename)
	if err != nil {
		fmt.Println(err)
		return
	}

	reply, err := request_reply(
		fmt.Sprintf("==============================================================\n"+
			" Attention, it will delete all the %d items of note \"%s\".\n"+
			"==============================================================\n",
			note.Sum, notename)+
			" Type \"%s\" to continue:",
		"yes")
	if err != nil {
		fmt.Println(err)
		return
	}

	if reply == false {
		return
	}

	err = notedb.DeleteNote(notename)
	if err != nil {
		fmt.Println(err)
		return
	}
}

func funUse(c *cli.Context) {
	if len(c.Args()) == 0 {
		fmt.Println("note name needed.")
		return
	}
	if len(c.Args()) > 1 {
		fmt.Println("only one note name allowed.")
		return
	}

	notename := c.Args().First()
	err := notedb.UseNote(notename)
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Printf("current note: \"%s\" (last update: %s).\n",
		notename, notedb.CurrentNote.LastUpdate)
}

func funAdd(c *cli.Context) {
	if len(c.Args()) != 2 {
		fmt.Println("tag and content needed.")
		return
	}

	item, err := notedb.AddNoteItem(c.Args()[0], c.Args()[1])
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Println(item)
}

func funRm(c *cli.Context) {
	if len(c.Args()) == 0 {
		fmt.Println("item ID needed.")
		return
	}

	for _, itemid := range c.Args() {

		itemid, err := strconv.Atoi(itemid)
		if err != nil {
			fmt.Println("item ID should be positive integer.")
			continue
		}

		// read item and print it, in case of misdeleteing
		item, err := notedb.ReadNoteItem(notedb.CurrentNote, itemid)
		if err != nil {
			return
		}

		err = notedb.RemoveNoteItem(notedb.CurrentNote, itemid)
		if err != nil {
			fmt.Println(err)
			continue
		}

		fmt.Println(item)
	}
}

func funTag(c *cli.Context) {
	// list all tags
	note, err := notedb.GetCurrentNote()
	if err != nil {
		fmt.Println(err)
		return
	}

	if len(c.Args()) == 0 {
		for tag, taginfo := range note.Tags {
			fmt.Printf("tag: %s\t(#. of items: %d).\n", tag, len(taginfo))
		}
		return
	}

	items, err := notedb.ItemByTag(c.Args())
	if err != nil {
		fmt.Println(err)
		return
	}

	for _, item := range items {
		fmt.Println(item)
	}
}

func funSearch(c *cli.Context) {
	if len(c.Args()) == 0 {
		fmt.Println("search keyword needed.")
		return
	}

	items, err := notedb.ItemByRegexp(c.Args())
	if err != nil {
		fmt.Println(err)
		return
	}

	for _, item := range items {
		fmt.Println(item)
	}
}

func funDump(c *cli.Context) {
	if len(c.Args()) > 0 {
		fmt.Println("no arguments should be given.")
		return
	}

	err := notedb.Dump()
	if err != nil {
		fmt.Println(err)
		return
	}
}

func funWipe(c *cli.Context) {
	if len(c.Args()) > 0 {
		fmt.Println("no arguments should be given.")
		return
	}

	reply, err := request_reply(
		"========================================\n"+
			" Attention, it will clear all the data.\n"+
			"========================================\n"+
			" Type \"%s\" to continue:",
		"yes")
	if err != nil {
		fmt.Println(err)
		return
	}

	if reply == false {
		return
	}

	err = notedb.Wipe()
	if err != nil {
		fmt.Println(err)
		return
	}
}

func funRestore(c *cli.Context) {
	if len(c.Args()) != 1 {
		fmt.Println("dumpped filename needed.")
		return
	}

	reply, err := request_reply(
		"========================================\n"+
			" Attention, it will clear all the data.\n"+
			"========================================\n"+
			" Type \"%s\" to continue:",
		"yes")
	if err != nil {
		fmt.Println(err)
		return
	}

	if reply == false {
		return
	}

	err = notedb.Restore(c.Args().First())
	if err != nil {
		fmt.Println(err)
		return
	}
}

func funImport(c *cli.Context) {
	if len(c.Args()) != 3 {
		fmt.Println("three arguments needed: <notename in your cnote>" +
			" <notename in dumpped note> <dumpped filename>.")
		return
	}
	notename, othernotename, filename := c.Args()[0], c.Args()[1], c.Args()[2]
	n, err := notedb.Import(notename, othernotename, filename)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Printf("%d items imported into note \"%s\".\n", n, notename)
}

func main() {
	notedb = cnotedb.NewNoteDB(DBFILE)
	defer notedb.Close()

	app := cli.NewApp()
	app.Name = "cnote"
	app.Usage = "A command line note app. https://github.com/shenwei356/cnote"
	app.Version = "1.1 (2014-07-20)"
	app.Author = "Wei Shen"
	app.Email = "shenwei356@gmail.com"

	app.Commands = []cli.Command{
		{
			Name:   "new",
			Usage:  "Create a new note",
			Action: getFunc(funcs, "new"),
		},
		{
			Name:   "del",
			Usage:  "Delete a note",
			Action: getFunc(funcs, "del"),
		},
		{
			Name:   "use",
			Usage:  "Select a note",
			Action: getFunc(funcs, "use"),
		},
		{
			Name:      "list",
			ShortName: "ls",
			Usage:     "List all notes",
			Action:    getFunc(funcs, "list"),
		},
		{
			Name:   "add",
			Usage:  "Add a note item",
			Action: getFunc(funcs, "add"),
		},
		{
			Name:   "rm",
			Usage:  "Remove a note item",
			Action: getFunc(funcs, "rm"),
		},
		{
			Name:      "tag",
			ShortName: "t",
			Usage:     "List items by tags. List all tags if no arguments given",
			Action:    getFunc(funcs, "tag"),
		},
		{
			Name:      "search",
			ShortName: "s",
			Usage:     "Search items with regular expression",
			Action:    getFunc(funcs, "search"),
		},
		{
			Name:   "dump",
			Usage:  "Dump whole database, for backup or transfer",
			Action: getFunc(funcs, "dump"),
		},
		{
			Name:   "wipe",
			Usage:  "Attention! Wipe whole database",
			Action: getFunc(funcs, "wipe"),
		},
		{
			Name:   "restore",
			Usage:  "Wipe whole database, and restore from dumpped file",
			Action: getFunc(funcs, "restore"),
		},
		{
			Name:   "import",
			Usage:  "Import note items from dumpped data",
			Action: getFunc(funcs, "import"),
		},
	}

	app.Run(os.Args)
}
<|endoftext|>


<commit_before>// Copyright © 2013 Steve Francia <spf@spf13.com>.
//
// 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.

// Commands similar to git, go tools and other modern CLI tools
// inspired by go, go-Commander, gh and subcommand

package cobra

import (
	"bytes"
	"fmt"
	flag "github.com/spf13/pflag"
	"os"
	"strings"
)

var _ = flag.ContinueOnError

type Flag interface {
	Args() []string
}

type Flags interface {
	Lookup(string) *Flag
	VisitAll(fn func(*Flag))
	Parse(arguments []string) error
}

// A Commander holds the configuration for the command line tool.
type Commander struct {
	// A Commander is also a Command for top level and global help & flags
	Command
	//ExitOnError, ContinueOnError or PanicOnError
	behavior flag.ErrorHandling

	args []string
}

func NewCommander() (c *Commander) {
	c = new(Commander)
	return
}

func (c *Commander) setFlagBehavior(b flag.ErrorHandling) error {
	if b == flag.ExitOnError || b == flag.ContinueOnError || b == flag.PanicOnError {
		c.behavior = b
		return nil
	}
	return fmt.Errorf("%v is not a valid behavior", b)
}

func (c *Commander) SetName(name string) {
	c.name = name
}

func (c *Commander) SetArgs(a []string) {
	c.args = a
}

func (c *Commander) Execute() (err error) {
	if len(c.args) == 0 {
		err = c.execute(os.Args[1:])
	} else {
		err = c.execute(c.args)
	}
	return
}

// Command is just that, a command for your application.
// eg.  'go run' ... 'run' is the command. Cobra requires
// you to define the usage and description as part of your command
// definition to ensure usability.
type Command struct {
	// Name is the command name, usually the executable's name.
	name string
	// The one-line usage message.
	Use string
	// The short description shown in the 'help' output.
	Short string
	// The long message shown in the 'help <this-command>' output.
	Long string
	// Set of flags specific to this command.
	flags *flag.FlagSet
	// Set of flags children commands will inherit
	pflags *flag.FlagSet
	// Run runs the command.
	// The args are the arguments after the command name.
	Run func(cmd *Command, args []string)
	// Commands is the list of commands supported by this Commander program.
	commands []*Command
	// Parent Command for this command
	parent *Command
	// Commander
	//cmdr *Commander
	flagErrorBuf *bytes.Buffer
}

// find the target command given the args and command tree
func (c *Command) Find(args []string) (cmd *Command, a []string, err error) {
	if c == nil {
		return nil, nil, fmt.Errorf("Called find() on a nil Command")
	}

	validSubCommand := false
	if len(args) > 1 && c.HasSubCommands() {
		for _, cmd := range c.commands {
			if cmd.Name() == args[0] {
				validSubCommand = true
				return cmd.Find(args[1:])
			}
		}
	}
	if !validSubCommand && c.Runnable() {
		return c, args, nil
	}

	return nil, nil, nil
}

// execute the command determined by args and the command tree
func (c *Command) execute(args []string) (err error) {
	err = fmt.Errorf("unknown subcommand %q\nRun 'help' for usage.\n", args[0])

	if c == nil {
		return fmt.Errorf("Called Execute() on a nil Command")
	}

	cmd, a, e := c.Find(args)
	if e == nil {
		err = cmd.ParseFlags(a)
		if err != nil {
			return err
		} else {
			argWoFlags := cmd.Flags().Args()
			cmd.Run(cmd, argWoFlags)
			return nil
		}
	}
	err = e
	return err
}

// Used for testing
func (c *Command) ResetCommands() {
	c.commands = nil
}

// Add one or many commands as children of this
func (c *Command) AddCommand(cmds ...*Command) {
	for i, x := range cmds {
		if cmds[i] == c {
			panic("Command can't be a child of itself")
		}
		cmds[i].parent = c
		c.commands = append(c.commands, x)
	}
}

// The full usage for a given command (including parents)
func (c *Command) Usage(depth ...int) string {
	i := 0
	if len(depth) > 0 {
		i = depth[0]
	}

	if c.HasParent() {
		return c.parent.Usage(i+1) + " " + c.Use
	} else if i > 0 {
		return c.Name()
	} else {
		return c.Use
	}
}

func (c *Command) DebugFlags() {
	fmt.Println("DebugFlags called on", c.Name())
	var debugflags func(*Command)

	debugflags = func(x *Command) {
		if x.HasFlags() || x.HasPersistentFlags() {
			fmt.Println(x.Name())
		}
		if x.HasFlags() {
			x.flags.VisitAll(func(f *flag.Flag) {
				if x.HasPersistentFlags() {
					if x.persistentFlag(f.Name) == nil {
						fmt.Println("  -"+f.Shorthand+",", "--"+f.Name, "["+f.DefValue+"]", "", f.Value, "  [L]")
					} else {
						fmt.Println("  -"+f.Shorthand+",", "--"+f.Name, "["+f.DefValue+"]", "", f.Value, "  [LP]")
					}
				} else {
					fmt.Println("  -"+f.Shorthand+",", "--"+f.Name, "["+f.DefValue+"]", "", f.Value, "  [L]")
				}
			})
		}
		if x.HasPersistentFlags() {
			x.pflags.VisitAll(func(f *flag.Flag) {
				if x.HasFlags() {
					if x.flags.Lookup(f.Name) == nil {
						fmt.Println("  -"+f.Shorthand+",", "--"+f.Name, "["+f.DefValue+"]", "", f.Value, "  [P]")
					}
				} else {
					fmt.Println("  -"+f.Shorthand+",", "--"+f.Name, "["+f.DefValue+"]", "", f.Value, "  [P]")
				}
			})
		}
		fmt.Println(x.flagErrorBuf)
		if x.HasSubCommands() {
			for _, y := range x.commands {
				debugflags(y)
			}
		}
	}

	debugflags(c)
}

// Usage prints the usage details to the standard output.
func (c *Command) PrintUsage() {
	if c.Runnable() {
		fmt.Printf("usage: %s\n\n", c.Usage())
	}

	fmt.Println(strings.Trim(c.Long, "\n"))
}

// Name returns the command's name: the first word in the use line.
func (c *Command) Name() string {
	if c.name != "" {
		return c.name
	}
	name := c.Use
	i := strings.Index(name, " ")
	if i >= 0 {
		name = name[:i]
	}
	return name
}

// Determine if the command is itself runnable
func (c *Command) Runnable() bool {
	return c.Run != nil
}

// Determine if the command has children commands
func (c *Command) HasSubCommands() bool {
	return len(c.commands) > 0
}

// Determine if the command has children commands
func (c *Command) HasParent() bool {
	return c.parent != nil
}

// Get the Commands FlagSet
func (c *Command) Flags() *flag.FlagSet {
	if c.flags == nil {
		c.flags = flag.NewFlagSet(c.Name(), flag.ContinueOnError)
		if c.flagErrorBuf == nil {
			c.flagErrorBuf = new(bytes.Buffer)
		}
		c.flags.SetOutput(c.flagErrorBuf)
	}
	return c.flags
}

// Get the Commands Persistent FlagSet
func (c *Command) PersistentFlags() *flag.FlagSet {
	if c.pflags == nil {
		c.pflags = flag.NewFlagSet(c.Name(), flag.ContinueOnError)
		if c.flagErrorBuf == nil {
			c.flagErrorBuf = new(bytes.Buffer)
		}
		c.pflags.SetOutput(c.flagErrorBuf)
	}
	return c.pflags
}

// Intended for use in testing
func (c *Command) ResetFlags() {
	c.flagErrorBuf = new(bytes.Buffer)
	c.flags = flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	c.flags.SetOutput(c.flagErrorBuf)
	c.pflags = flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	c.pflags.SetOutput(c.flagErrorBuf)
}

func (c *Command) HasFlags() bool {
	return c.Flags().HasFlags()
}

func (c *Command) HasPersistentFlags() bool {
	return c.PersistentFlags().HasFlags()
}

// Climbs up the command tree looking for matching flag
func (c *Command) Flag(name string) (flag *flag.Flag) {
	flag = c.Flags().Lookup(name)

	if flag == nil {
		flag = c.persistentFlag(name)
	}

	return
}

// recursively find matching persistent flag
func (c *Command) persistentFlag(name string) (flag *flag.Flag) {
	if c.HasPersistentFlags() {
		flag = c.PersistentFlags().Lookup(name)
	}

	if flag == nil && c.HasParent() {
		flag = c.parent.persistentFlag(name)
	}
	return
}

// Parses persistent flag tree & local flags
func (c *Command) ParseFlags(args []string) (err error) {
	c.mergePersistentFlags()

	err = c.Flags().Parse(args)
	if err != nil {
		return err
	}
	return nil
}

func (c *Command) mergePersistentFlags() {
	var rmerge func(x *Command)

	rmerge = func(x *Command) {
		if x.HasPersistentFlags() {
			x.PersistentFlags().VisitAll(func(f *flag.Flag) {
				if c.Flags().Lookup(f.Name) == nil {
					c.Flags().AddFlag(f)
				}
			})
		}
		if x.HasParent() {
			rmerge(x.parent)
		}
	}

	rmerge(c)
}

// Climbs up the command tree parsing flags from top to bottom
func (c *Command) ParsePersistentFlags(args []string) (err error) {
	if !c.HasParent() || (c.parent.HasPersistentFlags() && c.parent.PersistentFlags().Parsed()) {
		if c.HasPersistentFlags() {
			err = c.PersistentFlags().Parse(args)
			if err != nil {
				return err
			}
		}
	} else {
		if c.HasParent() && c.parent.HasPersistentFlags() {
			err = c.parent.ParsePersistentFlags(args)
			if err != nil {
				return err
			}
		}
	}
	return
}
<commit_msg>Documentation & housecleaning<commit_after>// Copyright © 2013 Steve Francia <spf@spf13.com>.
//
// 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.

// Commands similar to git, go tools and other modern CLI tools
// inspired by go, go-Commander, gh and subcommand

package cobra

import (
	"bytes"
	"fmt"
	flag "github.com/spf13/pflag"
	"os"
	"strings"
)

var _ = flag.ContinueOnError

// A Commander holds the configuration for the command line tool.
type Commander struct {
	// A Commander is also a Command for top level and global help & flags
	Command

	args []string
}

// Provide the user with a new commander.
// Not of a lot of value today, was intended to do more than just
// create a new commander.
func NewCommander() (c *Commander) {
	c = new(Commander)
	return
}

// Name for commander, should match application name
func (c *Commander) SetName(name string) {
	c.name = name
}

// os.Args[1:] by default, if desired, can be overridden
// particularly useful when testing.
func (c *Commander) SetArgs(a []string) {
	c.args = a
}

// Call execute to use the args (os.Args[1:] by default)
// and run through the command tree finding appropriate matches
// for commands and then corresponding flags.
func (c *Commander) Execute() (err error) {
	if len(c.args) == 0 {
		err = c.execute(os.Args[1:])
	} else {
		err = c.execute(c.args)
	}
	return
}

// Command is just that, a command for your application.
// eg.  'go run' ... 'run' is the command. Cobra requires
// you to define the usage and description as part of your command
// definition to ensure usability.
type Command struct {
	// Name is the command name, usually the executable's name.
	name string
	// The one-line usage message.
	Use string
	// The short description shown in the 'help' output.
	Short string
	// The long message shown in the 'help <this-command>' output.
	Long string
	// Set of flags specific to this command.
	flags *flag.FlagSet
	// Set of flags children commands will inherit
	pflags *flag.FlagSet
	// Run runs the command.
	// The args are the arguments after the command name.
	Run func(cmd *Command, args []string)
	// Commands is the list of commands supported by this Commander program.
	commands []*Command
	// Parent Command for this command
	parent *Command
	// Commander
	//cmdr *Commander
	flagErrorBuf *bytes.Buffer
}

// find the target command given the args and command tree
func (c *Command) Find(args []string) (cmd *Command, a []string, err error) {
	if c == nil {
		return nil, nil, fmt.Errorf("Called find() on a nil Command")
	}

	validSubCommand := false
	if len(args) > 1 && c.HasSubCommands() {
		for _, cmd := range c.commands {
			if cmd.Name() == args[0] {
				validSubCommand = true
				return cmd.Find(args[1:])
			}
		}
	}
	if !validSubCommand && c.Runnable() {
		return c, args, nil
	}

	return nil, nil, nil
}

// execute the command determined by args and the command tree
func (c *Command) execute(args []string) (err error) {
	err = fmt.Errorf("unknown subcommand %q\nRun 'help' for usage.\n", args[0])

	if c == nil {
		return fmt.Errorf("Called Execute() on a nil Command")
	}

	cmd, a, e := c.Find(args)
	if e == nil {
		err = cmd.ParseFlags(a)
		if err != nil {
			return err
		} else {
			argWoFlags := cmd.Flags().Args()
			cmd.Run(cmd, argWoFlags)
			return nil
		}
	}
	err = e
	return err
}

// Used for testing
func (c *Command) ResetCommands() {
	c.commands = nil
}

// Add one or many commands as children of this
func (c *Command) AddCommand(cmds ...*Command) {
	for i, x := range cmds {
		if cmds[i] == c {
			panic("Command can't be a child of itself")
		}
		cmds[i].parent = c
		c.commands = append(c.commands, x)
	}
}

// The full usage for a given command (including parents)
func (c *Command) Usage(depth ...int) string {
	i := 0
	if len(depth) > 0 {
		i = depth[0]
	}

	if c.HasParent() {
		return c.parent.Usage(i+1) + " " + c.Use
	} else if i > 0 {
		return c.Name()
	} else {
		return c.Use
	}
}

// For use in determining which flags have been assigned to which commands
// and which persist
func (c *Command) DebugFlags() {
	fmt.Println("DebugFlags called on", c.Name())
	var debugflags func(*Command)

	debugflags = func(x *Command) {
		if x.HasFlags() || x.HasPersistentFlags() {
			fmt.Println(x.Name())
		}
		if x.HasFlags() {
			x.flags.VisitAll(func(f *flag.Flag) {
				if x.HasPersistentFlags() {
					if x.persistentFlag(f.Name) == nil {
						fmt.Println("  -"+f.Shorthand+",", "--"+f.Name, "["+f.DefValue+"]", "", f.Value, "  [L]")
					} else {
						fmt.Println("  -"+f.Shorthand+",", "--"+f.Name, "["+f.DefValue+"]", "", f.Value, "  [LP]")
					}
				} else {
					fmt.Println("  -"+f.Shorthand+",", "--"+f.Name, "["+f.DefValue+"]", "", f.Value, "  [L]")
				}
			})
		}
		if x.HasPersistentFlags() {
			x.pflags.VisitAll(func(f *flag.Flag) {
				if x.HasFlags() {
					if x.flags.Lookup(f.Name) == nil {
						fmt.Println("  -"+f.Shorthand+",", "--"+f.Name, "["+f.DefValue+"]", "", f.Value, "  [P]")
					}
				} else {
					fmt.Println("  -"+f.Shorthand+",", "--"+f.Name, "["+f.DefValue+"]", "", f.Value, "  [P]")
				}
			})
		}
		fmt.Println(x.flagErrorBuf)
		if x.HasSubCommands() {
			for _, y := range x.commands {
				debugflags(y)
			}
		}
	}

	debugflags(c)
}

// Usage prints the usage details to the standard output.
func (c *Command) PrintUsage() {
	if c.Runnable() {
		fmt.Printf("usage: %s\n\n", c.Usage())
	}

	fmt.Println(strings.Trim(c.Long, "\n"))
}

// Name returns the command's name: the first word in the use line.
func (c *Command) Name() string {
	if c.name != "" {
		return c.name
	}
	name := c.Use
	i := strings.Index(name, " ")
	if i >= 0 {
		name = name[:i]
	}
	return name
}

// Determine if the command is itself runnable
func (c *Command) Runnable() bool {
	return c.Run != nil
}

// Determine if the command has children commands
func (c *Command) HasSubCommands() bool {
	return len(c.commands) > 0
}

// Determine if the command is a child command
func (c *Command) HasParent() bool {
	return c.parent != nil
}

// Get the Commands FlagSet
func (c *Command) Flags() *flag.FlagSet {
	if c.flags == nil {
		c.flags = flag.NewFlagSet(c.Name(), flag.ContinueOnError)
		if c.flagErrorBuf == nil {
			c.flagErrorBuf = new(bytes.Buffer)
		}
		c.flags.SetOutput(c.flagErrorBuf)
	}
	return c.flags
}

// Get the Commands Persistent FlagSet
func (c *Command) PersistentFlags() *flag.FlagSet {
	if c.pflags == nil {
		c.pflags = flag.NewFlagSet(c.Name(), flag.ContinueOnError)
		if c.flagErrorBuf == nil {
			c.flagErrorBuf = new(bytes.Buffer)
		}
		c.pflags.SetOutput(c.flagErrorBuf)
	}
	return c.pflags
}

// Intended for use in testing
func (c *Command) ResetFlags() {
	c.flagErrorBuf = new(bytes.Buffer)
	c.flags = flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	c.flags.SetOutput(c.flagErrorBuf)
	c.pflags = flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	c.pflags.SetOutput(c.flagErrorBuf)
}

// Does the command contain flags (local not persistent)
func (c *Command) HasFlags() bool {
	return c.Flags().HasFlags()
}

// Does the command contain persistent flags
func (c *Command) HasPersistentFlags() bool {
	return c.PersistentFlags().HasFlags()
}

// Climbs up the command tree looking for matching flag
func (c *Command) Flag(name string) (flag *flag.Flag) {
	flag = c.Flags().Lookup(name)

	if flag == nil {
		flag = c.persistentFlag(name)
	}

	return
}

// recursively find matching persistent flag
func (c *Command) persistentFlag(name string) (flag *flag.Flag) {
	if c.HasPersistentFlags() {
		flag = c.PersistentFlags().Lookup(name)
	}

	if flag == nil && c.HasParent() {
		flag = c.parent.persistentFlag(name)
	}
	return
}

// Parses persistent flag tree & local flags
func (c *Command) ParseFlags(args []string) (err error) {
	c.mergePersistentFlags()

	err = c.Flags().Parse(args)
	if err != nil {
		return err
	}
	return nil
}

func (c *Command) mergePersistentFlags() {
	var rmerge func(x *Command)

	rmerge = func(x *Command) {
		if x.HasPersistentFlags() {
			x.PersistentFlags().VisitAll(func(f *flag.Flag) {
				if c.Flags().Lookup(f.Name) == nil {
					c.Flags().AddFlag(f)
				}
			})
		}
		if x.HasParent() {
			rmerge(x.parent)
		}
	}

	rmerge(c)
}
<|endoftext|>


<commit_before>package goridge

import (
	"bytes"
	"io"
	"net/rpc"
	"sync"

	j "github.com/json-iterator/go"
	"github.com/spiral/errors"
)

var json = j.ConfigCompatibleWithStandardLibrary

// Codec represent net/rpc bridge over Goridge socket relay.
type Codec struct {
	relay  Relay
	closed bool
	frame  *Frame
	codec  sync.Map
}

// NewCodec initiates new server rpc codec over socket connection.
func NewCodec(rwc io.ReadWriteCloser) *Codec {
	return &Codec{
		relay: NewSocketRelay(rwc),
		codec: sync.Map{},
	}
}

// NewCodecWithRelay initiates new server rpc codec with a relay of choice.
func NewCodecWithRelay(relay Relay) *Codec {
	return &Codec{relay: relay}
}

// WriteResponse marshals response, byte slice or error to remote party.
func (c *Codec) WriteResponse(r *rpc.Response, body interface{}) error {
	const op = errors.Op("codec WriteResponse")
	frame := NewFrame()

	// SEQ_ID + METHOD_NAME_LEN
	frame.WriteOptions(uint32(r.Seq), uint32(len(r.ServiceMethod)))
	// Write protocol version
	frame.WriteVersion(VERSION_1)

	// Store response flags
	codec, ok := c.codec.Load(r.Seq)
	if ok == false {
		// fallback codec
		frame.WriteFlags(CODEC_GOB)
	} else {
		frame.WriteFlags(codec.(FrameFlag))
	}

	// initialize buffer
	buf := new(bytes.Buffer)
	// writeServiceMethod to the buffer
	buf.WriteString(r.ServiceMethod)

	// if error returned, we sending it via relay and return error from WriteResponse
	if r.Error != "" {
		// Append error flag
		return c.handleError(r, frame, buf)
	}

	flags := frame.ReadFlags()

	if flags&byte(CODEC_JSON) != 0 {
		err := encodeJSON(buf, body)
		if err != nil {
			return c.handleError(r, frame, buf)
		}
	}

	if flags&byte(CODEC_GOB) != 0 {
		err := encodeGob(buf, body)
		if err != nil {
			return c.handleError(r, frame, buf)
		}
	}

	if flags&byte(CODEC_RAW) != 0 {
		err := encodeRaw(buf, frame)
		if err != nil {
			return c.handleError(r, frame, buf)
		}
	}

	if flags&byte(CODEC_MSGPACK) != 0 {
		err := encodeMsgPack(buf, frame)
		if err != nil {
			return c.handleError(r, frame, buf)
		}
	}

	// everything written correct
	frame.WritePayloadLen(uint32(buf.Len()))
	frame.WritePayload(buf.Bytes())

	frame.WriteCRC()

	return c.relay.Send(frame)
}

func (c *Codec) handleError(r *rpc.Response, frame *Frame, buf *bytes.Buffer) error {
	// just to be sure, remove all data from buffer and write new
	buf.Truncate(0)
	buf.WriteString(r.ServiceMethod)

	const op = errors.Op("handle codec error")
	frame.WriteFlags(ERROR)
	// write data to the gob
	buf.WriteString(r.Error)

	frame.WritePayloadLen(uint32(buf.Len()))
	frame.WritePayload(buf.Bytes())

	frame.WriteCRC()
	_ = c.relay.Send(frame)
	return errors.E(op, errors.Str(r.Error))
}

// ReadRequestHeader receives
func (c *Codec) ReadRequestHeader(r *rpc.Request) error {
	const op = errors.Op("codec: read request header")
	frame := NewFrame()
	err := c.relay.Receive(frame)
	if err != nil {
		return err
	}

	// opts[0] sequence ID
	// opts[1] service method name offset from payload in bytes
	opts := frame.ReadOptions()
	if len(opts) != 2 {
		return errors.E(op, errors.Str("should be 2 options. SEQ_ID and METHOD_LEN"))
	}

	r.Seq = uint64(opts[0])
	r.ServiceMethod = string(frame.Payload()[:opts[1]])
	c.frame = frame
	return c.storeCodec(r, frame.ReadFlags())
}

func (c *Codec) storeCodec(r *rpc.Request, flag byte) error {
	if flag&byte(CODEC_JSON) != 0 {
		c.codec.Store(r.Seq, CODEC_JSON)
		return nil
	}

	if flag&byte(CODEC_GOB) != 0 {
		c.codec.Store(r.Seq, CODEC_GOB)
		return nil
	}

	if flag&byte(CODEC_RAW) != 0 {
		c.codec.Store(r.Seq, CODEC_RAW)
		return nil
	}

	if flag&byte(CODEC_MSGPACK) != 0 {
		c.codec.Store(r.Seq, CODEC_MSGPACK)
		return nil
	}

	// If no flags provided, store GOB
	c.codec.Store(r.Seq, CODEC_GOB)
	return nil
}

// ReadRequestBody fetches prefixed body data and automatically unmarshal it as json. RawBody flag will populate
// []byte lice argument for rpc method.
func (c *Codec) ReadRequestBody(out interface{}) error {
	const op = errors.Op("codec read request body")
	if out == nil {
		return nil
	}

	defer func() {
		c.frame = nil
	}()

	flags := c.frame.ReadFlags()

	if flags&byte(CODEC_JSON) != byte(0) {
		return decodeJSON(out, c.frame)
	}

	if flags&byte(CODEC_GOB) != byte(0) {
		return decodeGob(out, c.frame)
	}

	if flags&byte(CODEC_RAW) != byte(0) {
		return decodeRaw(out, c.frame)
	}

	if flags&byte(CODEC_MSGPACK) != byte(0) {
		return decodeMsgPack(out, c.frame)
	}

	return errors.E(op, errors.Str("unknown decoder used in frame"))
}

// Close underlying socket.
func (c *Codec) Close() error {
	if c.closed {
		return nil
	}

	c.closed = true
	return c.relay.Close()
}
<commit_msg>Fix body in msgpack<commit_after>package goridge

import (
	"bytes"
	"io"
	"net/rpc"
	"sync"

	j "github.com/json-iterator/go"
	"github.com/spiral/errors"
)

var json = j.ConfigCompatibleWithStandardLibrary

// Codec represent net/rpc bridge over Goridge socket relay.
type Codec struct {
	relay  Relay
	closed bool
	frame  *Frame
	codec  sync.Map
}

// NewCodec initiates new server rpc codec over socket connection.
func NewCodec(rwc io.ReadWriteCloser) *Codec {
	return &Codec{
		relay: NewSocketRelay(rwc),
		codec: sync.Map{},
	}
}

// NewCodecWithRelay initiates new server rpc codec with a relay of choice.
func NewCodecWithRelay(relay Relay) *Codec {
	return &Codec{relay: relay}
}

// WriteResponse marshals response, byte slice or error to remote party.
func (c *Codec) WriteResponse(r *rpc.Response, body interface{}) error {
	const op = errors.Op("codec WriteResponse")
	frame := NewFrame()

	// SEQ_ID + METHOD_NAME_LEN
	frame.WriteOptions(uint32(r.Seq), uint32(len(r.ServiceMethod)))
	// Write protocol version
	frame.WriteVersion(VERSION_1)

	// Store response flags
	codec, ok := c.codec.Load(r.Seq)
	if ok == false {
		// fallback codec
		frame.WriteFlags(CODEC_GOB)
	} else {
		frame.WriteFlags(codec.(FrameFlag))
	}

	// initialize buffer
	buf := new(bytes.Buffer)
	// writeServiceMethod to the buffer
	buf.WriteString(r.ServiceMethod)

	// if error returned, we sending it via relay and return error from WriteResponse
	if r.Error != "" {
		// Append error flag
		return c.handleError(r, frame, buf)
	}

	flags := frame.ReadFlags()

	if flags&byte(CODEC_JSON) != 0 {
		err := encodeJSON(buf, body)
		if err != nil {
			return c.handleError(r, frame, buf)
		}
	}

	if flags&byte(CODEC_GOB) != 0 {
		err := encodeGob(buf, body)
		if err != nil {
			return c.handleError(r, frame, buf)
		}
	}

	if flags&byte(CODEC_RAW) != 0 {
		err := encodeRaw(buf, body)
		if err != nil {
			return c.handleError(r, frame, buf)
		}
	}

	if flags&byte(CODEC_MSGPACK) != 0 {
		err := encodeMsgPack(buf, body)
		if err != nil {
			return c.handleError(r, frame, buf)
		}
	}

	// everything written correct
	frame.WritePayloadLen(uint32(buf.Len()))
	frame.WritePayload(buf.Bytes())

	frame.WriteCRC()

	return c.relay.Send(frame)
}

func (c *Codec) handleError(r *rpc.Response, frame *Frame, buf *bytes.Buffer) error {
	// just to be sure, remove all data from buffer and write new
	buf.Truncate(0)
	buf.WriteString(r.ServiceMethod)

	const op = errors.Op("handle codec error")
	frame.WriteFlags(ERROR)
	// write data to the gob
	buf.WriteString(r.Error)

	frame.WritePayloadLen(uint32(buf.Len()))
	frame.WritePayload(buf.Bytes())

	frame.WriteCRC()
	_ = c.relay.Send(frame)
	return errors.E(op, errors.Str(r.Error))
}

// ReadRequestHeader receives
func (c *Codec) ReadRequestHeader(r *rpc.Request) error {
	const op = errors.Op("codec: read request header")
	frame := NewFrame()
	err := c.relay.Receive(frame)
	if err != nil {
		return err
	}

	// opts[0] sequence ID
	// opts[1] service method name offset from payload in bytes
	opts := frame.ReadOptions()
	if len(opts) != 2 {
		return errors.E(op, errors.Str("should be 2 options. SEQ_ID and METHOD_LEN"))
	}

	r.Seq = uint64(opts[0])
	r.ServiceMethod = string(frame.Payload()[:opts[1]])
	c.frame = frame
	return c.storeCodec(r, frame.ReadFlags())
}

func (c *Codec) storeCodec(r *rpc.Request, flag byte) error {
	if flag&byte(CODEC_JSON) != 0 {
		c.codec.Store(r.Seq, CODEC_JSON)
		return nil
	}

	if flag&byte(CODEC_GOB) != 0 {
		c.codec.Store(r.Seq, CODEC_GOB)
		return nil
	}

	if flag&byte(CODEC_RAW) != 0 {
		c.codec.Store(r.Seq, CODEC_RAW)
		return nil
	}

	if flag&byte(CODEC_MSGPACK) != 0 {
		c.codec.Store(r.Seq, CODEC_MSGPACK)
		return nil
	}

	// If no flags provided, store GOB
	c.codec.Store(r.Seq, CODEC_GOB)
	return nil
}

// ReadRequestBody fetches prefixed body data and automatically unmarshal it as json. RawBody flag will populate
// []byte lice argument for rpc method.
func (c *Codec) ReadRequestBody(out interface{}) error {
	const op = errors.Op("codec read request body")
	if out == nil {
		return nil
	}

	defer func() {
		c.frame = nil
	}()

	flags := c.frame.ReadFlags()

	if flags&byte(CODEC_JSON) != byte(0) {
		return decodeJSON(out, c.frame)
	}

	if flags&byte(CODEC_GOB) != byte(0) {
		return decodeGob(out, c.frame)
	}

	if flags&byte(CODEC_RAW) != byte(0) {
		return decodeRaw(out, c.frame)
	}

	if flags&byte(CODEC_MSGPACK) != byte(0) {
		return decodeMsgPack(out, c.frame)
	}

	return errors.E(op, errors.Str("unknown decoder used in frame"))
}

// Close underlying socket.
func (c *Codec) Close() error {
	if c.closed {
		return nil
	}

	c.closed = true
	return c.relay.Close()
}
<|endoftext|>


<commit_before>package multiaddr

import (
	"bytes"
	"fmt"
	"strings"

	"github.com/multiformats/go-varint"
)

func stringToBytes(s string) ([]byte, error) {
	// consume trailing slashes
	s = strings.TrimRight(s, "/")

	var b bytes.Buffer
	sp := strings.Split(s, "/")

	if sp[0] != "" {
		return nil, fmt.Errorf("failed to parse multiaddr %q: must begin with /", s)
	}

	// consume first empty elem
	sp = sp[1:]

	if len(sp) == 0 {
		return nil, fmt.Errorf("failed to parse multiaddr %q: empty multiaddr", s)
	}

	for len(sp) > 0 {
		name := sp[0]
		p := ProtocolWithName(name)
		if p.Code == 0 {
			return nil, fmt.Errorf("failed to parse multiaddr %q: unknown protocol %s", s, sp[0])
		}
		_, _ = b.Write(p.VCode)
		sp = sp[1:]

		if p.Size == 0 { // no length.
			continue
		}

		if len(sp) < 1 {
			return nil, fmt.Errorf("failed to parse multiaddr %q: unexpected end of multiaddr", s)
		}

		if p.Path {
			// it's a path protocol (terminal).
			// consume the rest of the address as the next component.
			sp = []string{"/" + strings.Join(sp, "/")}
		}

		a, err := p.Transcoder.StringToBytes(sp[0])
		if err != nil {
			return nil, fmt.Errorf("failed to parse multiaddr %q: invalid value %q for protocol %s: %s", s, sp[0], p.Name, err)
		}
		if p.Size < 0 { // varint size.
			_, _ = b.Write(varint.ToUvarint(uint64(len(a))))
		}
		b.Write(a)
		sp = sp[1:]
	}

	return b.Bytes(), nil
}

func validateBytes(b []byte) (err error) {
	if len(b) == 0 {
		return fmt.Errorf("empty multiaddr")
	}
	for len(b) > 0 {
		code, n, err := ReadVarintCode(b)
		if err != nil {
			return err
		}

		b = b[n:]
		p := ProtocolWithCode(code)
		if p.Code == 0 {
			return fmt.Errorf("no protocol with code %d", code)
		}

		if p.Size == 0 {
			continue
		}

		n, size, err := sizeForAddr(p, b)
		if err != nil {
			return err
		}

		b = b[n:]

		if len(b) < size || size < 0 {
			return fmt.Errorf("invalid value for size %d", len(b))
		}

		err = p.Transcoder.ValidateBytes(b[:size])
		if err != nil {
			return err
		}

		b = b[size:]
	}

	return nil
}

func readComponent(b []byte) (int, Component, error) {
	var offset int
	code, n, err := ReadVarintCode(b)
	if err != nil {
		return 0, Component{}, err
	}
	offset += n

	p := ProtocolWithCode(code)
	if p.Code == 0 {
		return 0, Component{}, fmt.Errorf("no protocol with code %d", code)
	}

	if p.Size == 0 {
		return offset, Component{
			bytes:    b[:offset],
			offset:   offset,
			protocol: p,
		}, nil
	}

	n, size, err := sizeForAddr(p, b[offset:])
	if err != nil {
		return 0, Component{}, err
	}

	offset += n

	if len(b[offset:]) < size || size < 0 {
		return 0, Component{}, fmt.Errorf("invalid value for size %d", len(b[offset:]))
	}

	return offset + size, Component{
		bytes:    b[:offset+size],
		protocol: p,
		offset:   offset,
	}, nil
}

func bytesToString(b []byte) (ret string, err error) {
	if len(b) == 0 {
		return "", fmt.Errorf("empty multiaddr")
	}
	var buf strings.Builder

	for len(b) > 0 {
		n, c, err := readComponent(b)
		if err != nil {
			return "", err
		}
		b = b[n:]
		c.writeTo(&buf)
	}

	return buf.String(), nil
}

func sizeForAddr(p Protocol, b []byte) (skip, size int, err error) {
	switch {
	case p.Size > 0:
		return 0, (p.Size / 8), nil
	case p.Size == 0:
		return 0, 0, nil
	default:
		size, n, err := ReadVarintCode(b)
		if err != nil {
			return 0, 0, err
		}
		return n, size, nil
	}
}

func bytesSplit(b []byte) ([][]byte, error) {
	var ret [][]byte
	for len(b) > 0 {
		code, n, err := ReadVarintCode(b)
		if err != nil {
			return nil, err
		}

		p := ProtocolWithCode(code)
		if p.Code == 0 {
			return nil, fmt.Errorf("no protocol with code %d", b[0])
		}

		n2, size, err := sizeForAddr(p, b[n:])
		if err != nil {
			return nil, err
		}

		length := n + n2 + size
		ret = append(ret, b[:length])
		b = b[length:]
	}

	return ret, nil
}
<commit_msg>remove unused bytesSplit function<commit_after>package multiaddr

import (
	"bytes"
	"fmt"
	"strings"

	"github.com/multiformats/go-varint"
)

func stringToBytes(s string) ([]byte, error) {
	// consume trailing slashes
	s = strings.TrimRight(s, "/")

	var b bytes.Buffer
	sp := strings.Split(s, "/")

	if sp[0] != "" {
		return nil, fmt.Errorf("failed to parse multiaddr %q: must begin with /", s)
	}

	// consume first empty elem
	sp = sp[1:]

	if len(sp) == 0 {
		return nil, fmt.Errorf("failed to parse multiaddr %q: empty multiaddr", s)
	}

	for len(sp) > 0 {
		name := sp[0]
		p := ProtocolWithName(name)
		if p.Code == 0 {
			return nil, fmt.Errorf("failed to parse multiaddr %q: unknown protocol %s", s, sp[0])
		}
		_, _ = b.Write(p.VCode)
		sp = sp[1:]

		if p.Size == 0 { // no length.
			continue
		}

		if len(sp) < 1 {
			return nil, fmt.Errorf("failed to parse multiaddr %q: unexpected end of multiaddr", s)
		}

		if p.Path {
			// it's a path protocol (terminal).
			// consume the rest of the address as the next component.
			sp = []string{"/" + strings.Join(sp, "/")}
		}

		a, err := p.Transcoder.StringToBytes(sp[0])
		if err != nil {
			return nil, fmt.Errorf("failed to parse multiaddr %q: invalid value %q for protocol %s: %s", s, sp[0], p.Name, err)
		}
		if p.Size < 0 { // varint size.
			_, _ = b.Write(varint.ToUvarint(uint64(len(a))))
		}
		b.Write(a)
		sp = sp[1:]
	}

	return b.Bytes(), nil
}

func validateBytes(b []byte) (err error) {
	if len(b) == 0 {
		return fmt.Errorf("empty multiaddr")
	}
	for len(b) > 0 {
		code, n, err := ReadVarintCode(b)
		if err != nil {
			return err
		}

		b = b[n:]
		p := ProtocolWithCode(code)
		if p.Code == 0 {
			return fmt.Errorf("no protocol with code %d", code)
		}

		if p.Size == 0 {
			continue
		}

		n, size, err := sizeForAddr(p, b)
		if err != nil {
			return err
		}

		b = b[n:]

		if len(b) < size || size < 0 {
			return fmt.Errorf("invalid value for size %d", len(b))
		}

		err = p.Transcoder.ValidateBytes(b[:size])
		if err != nil {
			return err
		}

		b = b[size:]
	}

	return nil
}

func readComponent(b []byte) (int, Component, error) {
	var offset int
	code, n, err := ReadVarintCode(b)
	if err != nil {
		return 0, Component{}, err
	}
	offset += n

	p := ProtocolWithCode(code)
	if p.Code == 0 {
		return 0, Component{}, fmt.Errorf("no protocol with code %d", code)
	}

	if p.Size == 0 {
		return offset, Component{
			bytes:    b[:offset],
			offset:   offset,
			protocol: p,
		}, nil
	}

	n, size, err := sizeForAddr(p, b[offset:])
	if err != nil {
		return 0, Component{}, err
	}

	offset += n

	if len(b[offset:]) < size || size < 0 {
		return 0, Component{}, fmt.Errorf("invalid value for size %d", len(b[offset:]))
	}

	return offset + size, Component{
		bytes:    b[:offset+size],
		protocol: p,
		offset:   offset,
	}, nil
}

func bytesToString(b []byte) (ret string, err error) {
	if len(b) == 0 {
		return "", fmt.Errorf("empty multiaddr")
	}
	var buf strings.Builder

	for len(b) > 0 {
		n, c, err := readComponent(b)
		if err != nil {
			return "", err
		}
		b = b[n:]
		c.writeTo(&buf)
	}

	return buf.String(), nil
}

func sizeForAddr(p Protocol, b []byte) (skip, size int, err error) {
	switch {
	case p.Size > 0:
		return 0, (p.Size / 8), nil
	case p.Size == 0:
		return 0, 0, nil
	default:
		size, n, err := ReadVarintCode(b)
		if err != nil {
			return 0, 0, err
		}
		return n, size, nil
	}
}
<|endoftext|>


<commit_before>// Package colly implements a HTTP scraping framework
package colly

import (
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/html"
	"golang.org/x/net/html/charset"

	"github.com/PuerkitoBio/goquery"
)

// Collector provides the scraper instance for a scraping job
type Collector struct {
	// UserAgent is the User-Agent string used by HTTP requests
	UserAgent string
	// MaxDepth limits the recursion depth of visited URLs.
	// Set it to 0 for infinite recursion (default).
	MaxDepth int
	// AllowedDomains is a domain whitelist.
	// Leave it blank to allow any domains to be visited
	AllowedDomains []string
	// AllowURLRevisit allows multiple downloads of the same URL
	AllowURLRevisit bool
	// MaxBodySize is the limit of the retrieved response body in bytes.
	// `0` means unlimited.
	// The default value for MaxBodySize is 10MB (10 * 1024 * 1024 bytes).
	MaxBodySize int
	// CacheDir specifies a location where GET requests are cached as files.
	// When it's not defined, caching is disabled.
	CacheDir          string
	visitedURLs       []string
	htmlCallbacks     map[string]HTMLCallback
	requestCallbacks  []RequestCallback
	responseCallbacks []ResponseCallback
	backend           *httpBackend
	wg                *sync.WaitGroup
	lock              *sync.Mutex
}

// Request is the representation of a HTTP request made by a Collector
type Request struct {
	// URL is the parsed URL of the HTTP request
	URL *url.URL
	// Headers contains the Request's HTTP headers
	Headers *http.Header
	// Ctx is a context between a Request and a Response
	Ctx *Context
	// Depth is the number of the parents of this request
	Depth     int
	collector *Collector
}

// Response is the representation of a HTTP response made by a Collector
type Response struct {
	// StatusCode is the status code of the Response
	StatusCode int
	// Body is the content of the Response
	Body []byte
	// Ctx is a context between a Request and a Response
	Ctx *Context
	// Request is the Request object of the response
	Request *Request
	// Headers contains the Response's HTTP headers
	Headers *http.Header
}

// HTMLElement is the representation of a HTML tag.
type HTMLElement struct {
	// Name is the name of the tag
	Name       string
	Text       string
	attributes []html.Attribute
	// Request is the request object of the element's HTML document
	Request *Request
	// Response is the Response object of the element's HTML document
	Response *Response
	// DOM is the goquery parsed DOM object of the page. DOM is relative
	// to the current HTMLElement
	DOM *goquery.Selection
}

// Context provides a tiny layer for passing data between callbacks
type Context struct {
	contextMap map[string]string
	lock       *sync.RWMutex
}

// RequestCallback is a type alias for OnRequest callback functions
type RequestCallback func(*Request)

// ResponseCallback is a type alias for OnResponse callback functions
type ResponseCallback func(*Response)

// HTMLCallback is a type alias for OnHTML callback functions
type HTMLCallback func(*HTMLElement)

// NewCollector creates a new Collector instance with default configuration
func NewCollector() *Collector {
	c := &Collector{}
	c.Init()
	return c
}

// NewContext initializes a new Context instance
func NewContext() *Context {
	return &Context{
		contextMap: make(map[string]string),
		lock:       &sync.RWMutex{},
	}
}

// Init initializes the Collector's private variables and sets default
// configuration for the Collector
func (c *Collector) Init() {
	c.UserAgent = "colly - https://github.com/asciimoo/colly"
	c.MaxDepth = 0
	c.visitedURLs = make([]string, 0, 8)
	c.htmlCallbacks = make(map[string]HTMLCallback, 0)
	c.requestCallbacks = make([]RequestCallback, 0, 8)
	c.responseCallbacks = make([]ResponseCallback, 0, 8)
	c.MaxBodySize = 10 * 1024 * 1024
	c.backend = &httpBackend{}
	c.backend.Init()
	c.wg = &sync.WaitGroup{}
	c.lock = &sync.Mutex{}
}

// Visit starts Collector's collecting job by creating a
// request to the URL specified in parameter.
// Visit also calls the previously provided OnRequest,
// OnResponse, OnHTML callbacks
func (c *Collector) Visit(URL string) error {
	return c.scrape(URL, "GET", 1, nil, nil)
}

// Post starts a collector job by creating a POST request.
// Post also calls the previously provided OnRequest,
// OnResponse, OnHTML callbacks
func (c *Collector) Post(URL string, requestData map[string]string) error {
	return c.scrape(URL, "POST", 1, createFormReader(requestData), nil)
}

// PostRaw starts a collector job by creating a POST request with raw binary data.
// Post also calls the previously provided callbacks
func (c *Collector) PostRaw(URL string, requestData []byte) error {
	return c.scrape(URL, "POST", 1, bytes.NewReader(requestData), nil)
}

func (c *Collector) scrape(u, method string, depth int, requestData io.Reader, ctx *Context) error {
	c.wg.Add(1)
	defer c.wg.Done()
	if err := c.requestCheck(u, depth); err != nil {
		return err
	}
	parsedURL, err := url.Parse(u)
	if err != nil {
		return err
	}
	if !c.isDomainAllowed(parsedURL.Host) {
		return errors.New("Forbidden domain")
	}
	req, err := http.NewRequest(method, u, requestData)
	if err != nil {
		return err
	}
	req.Header.Set("User-Agent", c.UserAgent)
	if ctx == nil {
		ctx = NewContext()
	}
	request := &Request{
		URL:       parsedURL,
		Headers:   &req.Header,
		Ctx:       ctx,
		Depth:     depth,
		collector: c,
	}

	c.handleOnRequest(request)

	if method == "POST" && req.Header.Get("Content-Type") == "" {
		req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	}
	response, err := c.backend.Cache(req, c.MaxBodySize, c.CacheDir)
	// TODO add OnError callback to handle these cases
	if err != nil {
		return err
	}
	response.Ctx = ctx
	response.Request = request
	response.fixCharset()

	c.handleOnResponse(response)

	c.handleOnHTML(request, response)

	return nil
}

func (c *Collector) requestCheck(u string, depth int) error {
	if u == "" {
		return errors.New("Missing URL")
	}
	if c.MaxDepth > 0 && c.MaxDepth < depth {
		return errors.New("Max depth limit reached")
	}
	if !c.AllowURLRevisit {
		for _, u2 := range c.visitedURLs {
			if u2 == u {
				return errors.New("URL already visited")
			}
		}
		c.lock.Lock()
		c.visitedURLs = append(c.visitedURLs, u)
		c.lock.Unlock()
	}
	return nil
}

func (c *Collector) isDomainAllowed(domain string) bool {
	if c.AllowedDomains == nil || len(c.AllowedDomains) == 0 {
		return true
	}
	for _, d2 := range c.AllowedDomains {
		if d2 == domain {
			return true
		}
	}
	return false
}

// Wait returns when the collector jobs are finished
func (c *Collector) Wait() {
	c.wg.Wait()
}

// OnRequest registers a function. Function will be executed on every
// request made by the Collector
func (c *Collector) OnRequest(f RequestCallback) {
	c.lock.Lock()
	c.requestCallbacks = append(c.requestCallbacks, f)
	c.lock.Unlock()
}

// OnResponse registers a function. Function will be executed on every response
func (c *Collector) OnResponse(f ResponseCallback) {
	c.lock.Lock()
	c.responseCallbacks = append(c.responseCallbacks, f)
	c.lock.Unlock()
}

// OnHTML registers a function. Function will be executed on every HTML
// element matched by the `goquerySelector` parameter.
// `goquerySelector` is a selector used by https://github.com/PuerkitoBio/goquery
func (c *Collector) OnHTML(goquerySelector string, f HTMLCallback) {
	c.lock.Lock()
	c.htmlCallbacks[goquerySelector] = f
	c.lock.Unlock()
}

// WithTransport allows you to set a custom http.Transport for this collector.
func (c *Collector) WithTransport(transport *http.Transport) {
	c.backend.Client.Transport = transport
}

// DisableCookies turns off cookie handling for this collector
func (c *Collector) DisableCookies() {
	c.backend.Client.Jar = nil
}

// SetRequestTimeout overrides the default timeout (10 seconds) for this collector
func (c *Collector) SetRequestTimeout(timeout time.Duration) {
	c.backend.Client.Timeout = timeout
}

func (c *Collector) handleOnRequest(r *Request) {
	for _, f := range c.requestCallbacks {
		f(r)
	}
}

func (c *Collector) handleOnResponse(r *Response) {
	for _, f := range c.responseCallbacks {
		f(r)
	}
}

func (c *Collector) handleOnHTML(req *Request, resp *Response) {
	if strings.Index(strings.ToLower(resp.Headers.Get("Content-Type")), "html") == -1 {
		return
	}
	doc, err := goquery.NewDocumentFromReader(bytes.NewBuffer(resp.Body))
	if err != nil {
		return
	}
	for expr, f := range c.htmlCallbacks {
		doc.Find(expr).Each(func(i int, s *goquery.Selection) {
			for _, n := range s.Nodes {
				e := &HTMLElement{
					Name:       n.Data,
					Request:    req,
					Response:   resp,
					Text:       goquery.NewDocumentFromNode(n).Text(),
					DOM:        s,
					attributes: n.Attr,
				}
				f(e)
			}
		})
	}
}

// Limit adds a new `LimitRule` to the collector
func (c *Collector) Limit(rule *LimitRule) error {
	return c.backend.Limit(rule)
}

// Limits adds new `LimitRule`s to the collector
func (c *Collector) Limits(rules []*LimitRule) error {
	return c.backend.Limits(rules)
}

// SetCookies handles the receipt of the cookies in a reply for the given URL
func (c *Collector) SetCookies(URL string, cookies []*http.Cookie) error {
	if c.backend.Client.Jar == nil {
		return errors.New("Cookie jar is not available")
	}
	u, err := url.Parse(URL)
	if err != nil {
		return err
	}
	c.backend.Client.Jar.SetCookies(u, cookies)
	return nil
}

// Cookies returns the cookies to send in a request for the given URL.
func (c *Collector) Cookies(URL string) []*http.Cookie {
	if c.backend.Client.Jar == nil {
		return nil
	}
	u, err := url.Parse(URL)
	if err != nil {
		return nil
	}
	return c.backend.Client.Jar.Cookies(u)
}

// Attr returns the selected attribute of a HTMLElement or empty string
// if no attribute found
func (h *HTMLElement) Attr(k string) string {
	for _, a := range h.attributes {
		if a.Key == k {
			return a.Val
		}
	}
	return ""
}

// AbsoluteURL returns with the resolved absolute URL of an URL chunk.
// AbsoluteURL returns empty string if the URL chunk is a fragment or
// could not be parsed
func (r *Request) AbsoluteURL(u string) string {
	if strings.HasPrefix(u, "#") {
		return ""
	}
	absURL, err := r.URL.Parse(u)
	if err != nil {
		return ""
	}
	absURL.Fragment = ""
	if absURL.Scheme == "//" {
		absURL.Scheme = r.URL.Scheme
	}
	return absURL.String()
}

// Visit continues Collector's collecting job by creating a
// request and preserves the Context of the previous request.
// Visit also calls the previously provided OnRequest,
// OnResponse, OnHTML callbacks
func (r *Request) Visit(URL string) error {
	return r.collector.scrape(r.AbsoluteURL(URL), "GET", r.Depth+1, nil, r.Ctx)
}

// Post continues a collector job by creating a POST request and preserves the Context
// of the previous request.
// Post also calls the previously provided OnRequest, OnResponse, OnHTML callbacks
func (r *Request) Post(URL string, requestData map[string]string) error {
	return r.collector.scrape(r.AbsoluteURL(URL), "POST", r.Depth+1, createFormReader(requestData), r.Ctx)
}

// PostRaw starts a collector job by creating a POST request with raw binary data.
// PostRaw preserves the Context of the previous request
// and calls the previously provided callbacks
func (r *Request) PostRaw(URL string, requestData []byte) error {
	return r.collector.scrape(r.AbsoluteURL(URL), "POST", r.Depth+1, bytes.NewReader(requestData), r.Ctx)
}

// UnmarshalBinary decodes Context value to nil
// This function is used by request caching
func (c *Context) UnmarshalBinary(_ []byte) error {
	return nil
}

// MarshalBinary encodes Context value
// This function is used by request caching
func (c *Context) MarshalBinary() (_ []byte, _ error) {
	return nil, nil
}

// Put stores a value in Context
func (c *Context) Put(key, value string) {
	c.lock.Lock()
	c.contextMap[key] = value
	c.lock.Unlock()
}

// Get retrieves a value from Context. If no value found for `k`
// Get returns an empty string if key not found
func (c *Context) Get(key string) string {
	c.lock.RLock()
	if v, ok := c.contextMap[key]; ok {
		c.lock.RUnlock()
		return v
	}
	c.lock.RUnlock()
	return ""
}

func createFormReader(data map[string]string) io.Reader {
	form := url.Values{}
	for k, v := range data {
		form.Add(k, v)
	}
	return strings.NewReader(form.Encode())
}

func (r *Response) fixCharset() {
	contentType := strings.ToLower(r.Headers.Get("Content-Type"))
	if strings.Index(contentType, "charset") == -1 {
		return
	}
	if strings.Index(contentType, "utf-8") != -1 || strings.Index(contentType, "utf8") != -1 {
		return
	}
	encodedBodyReader, err := charset.NewReader(bytes.NewReader(r.Body), contentType)
	if err != nil {
		return
	}
	tmpBody, err := ioutil.ReadAll(encodedBodyReader)
	if err != nil {
		return
	}
	r.Body = tmpBody
}
<commit_msg>[enh] provide default URL scheme<commit_after>// Package colly implements a HTTP scraping framework
package colly

import (
	"bytes"
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/html"
	"golang.org/x/net/html/charset"

	"github.com/PuerkitoBio/goquery"
)

// Collector provides the scraper instance for a scraping job
type Collector struct {
	// UserAgent is the User-Agent string used by HTTP requests
	UserAgent string
	// MaxDepth limits the recursion depth of visited URLs.
	// Set it to 0 for infinite recursion (default).
	MaxDepth int
	// AllowedDomains is a domain whitelist.
	// Leave it blank to allow any domains to be visited
	AllowedDomains []string
	// AllowURLRevisit allows multiple downloads of the same URL
	AllowURLRevisit bool
	// MaxBodySize is the limit of the retrieved response body in bytes.
	// `0` means unlimited.
	// The default value for MaxBodySize is 10MB (10 * 1024 * 1024 bytes).
	MaxBodySize int
	// CacheDir specifies a location where GET requests are cached as files.
	// When it's not defined, caching is disabled.
	CacheDir          string
	visitedURLs       []string
	htmlCallbacks     map[string]HTMLCallback
	requestCallbacks  []RequestCallback
	responseCallbacks []ResponseCallback
	backend           *httpBackend
	wg                *sync.WaitGroup
	lock              *sync.Mutex
}

// Request is the representation of a HTTP request made by a Collector
type Request struct {
	// URL is the parsed URL of the HTTP request
	URL *url.URL
	// Headers contains the Request's HTTP headers
	Headers *http.Header
	// Ctx is a context between a Request and a Response
	Ctx *Context
	// Depth is the number of the parents of this request
	Depth     int
	collector *Collector
}

// Response is the representation of a HTTP response made by a Collector
type Response struct {
	// StatusCode is the status code of the Response
	StatusCode int
	// Body is the content of the Response
	Body []byte
	// Ctx is a context between a Request and a Response
	Ctx *Context
	// Request is the Request object of the response
	Request *Request
	// Headers contains the Response's HTTP headers
	Headers *http.Header
}

// HTMLElement is the representation of a HTML tag.
type HTMLElement struct {
	// Name is the name of the tag
	Name       string
	Text       string
	attributes []html.Attribute
	// Request is the request object of the element's HTML document
	Request *Request
	// Response is the Response object of the element's HTML document
	Response *Response
	// DOM is the goquery parsed DOM object of the page. DOM is relative
	// to the current HTMLElement
	DOM *goquery.Selection
}

// Context provides a tiny layer for passing data between callbacks
type Context struct {
	contextMap map[string]string
	lock       *sync.RWMutex
}

// RequestCallback is a type alias for OnRequest callback functions
type RequestCallback func(*Request)

// ResponseCallback is a type alias for OnResponse callback functions
type ResponseCallback func(*Response)

// HTMLCallback is a type alias for OnHTML callback functions
type HTMLCallback func(*HTMLElement)

// NewCollector creates a new Collector instance with default configuration
func NewCollector() *Collector {
	c := &Collector{}
	c.Init()
	return c
}

// NewContext initializes a new Context instance
func NewContext() *Context {
	return &Context{
		contextMap: make(map[string]string),
		lock:       &sync.RWMutex{},
	}
}

// Init initializes the Collector's private variables and sets default
// configuration for the Collector
func (c *Collector) Init() {
	c.UserAgent = "colly - https://github.com/asciimoo/colly"
	c.MaxDepth = 0
	c.visitedURLs = make([]string, 0, 8)
	c.htmlCallbacks = make(map[string]HTMLCallback, 0)
	c.requestCallbacks = make([]RequestCallback, 0, 8)
	c.responseCallbacks = make([]ResponseCallback, 0, 8)
	c.MaxBodySize = 10 * 1024 * 1024
	c.backend = &httpBackend{}
	c.backend.Init()
	c.wg = &sync.WaitGroup{}
	c.lock = &sync.Mutex{}
}

// Visit starts Collector's collecting job by creating a
// request to the URL specified in parameter.
// Visit also calls the previously provided OnRequest,
// OnResponse, OnHTML callbacks
func (c *Collector) Visit(URL string) error {
	return c.scrape(URL, "GET", 1, nil, nil)
}

// Post starts a collector job by creating a POST request.
// Post also calls the previously provided OnRequest,
// OnResponse, OnHTML callbacks
func (c *Collector) Post(URL string, requestData map[string]string) error {
	return c.scrape(URL, "POST", 1, createFormReader(requestData), nil)
}

// PostRaw starts a collector job by creating a POST request with raw binary data.
// Post also calls the previously provided callbacks
func (c *Collector) PostRaw(URL string, requestData []byte) error {
	return c.scrape(URL, "POST", 1, bytes.NewReader(requestData), nil)
}

func (c *Collector) scrape(u, method string, depth int, requestData io.Reader, ctx *Context) error {
	c.wg.Add(1)
	defer c.wg.Done()
	if err := c.requestCheck(u, depth); err != nil {
		return err
	}
	parsedURL, err := url.Parse(u)
	if err != nil {
		return err
	}
	if parsedURL.Scheme == "" {
		parsedURL.Scheme = "http"
	}
	if !c.isDomainAllowed(parsedURL.Host) {
		return errors.New("Forbidden domain")
	}
	req, err := http.NewRequest(method, parsedURL.String(), requestData)
	if err != nil {
		return err
	}
	req.Header.Set("User-Agent", c.UserAgent)
	if ctx == nil {
		ctx = NewContext()
	}
	request := &Request{
		URL:       parsedURL,
		Headers:   &req.Header,
		Ctx:       ctx,
		Depth:     depth,
		collector: c,
	}

	c.handleOnRequest(request)

	if method == "POST" && req.Header.Get("Content-Type") == "" {
		req.Header.Add("Content-Type", "application/x-www-form-urlencoded")
	}
	response, err := c.backend.Cache(req, c.MaxBodySize, c.CacheDir)
	// TODO add OnError callback to handle these cases
	if err != nil {
		return err
	}
	response.Ctx = ctx
	response.Request = request
	response.fixCharset()

	c.handleOnResponse(response)

	c.handleOnHTML(request, response)

	return nil
}

func (c *Collector) requestCheck(u string, depth int) error {
	if u == "" {
		return errors.New("Missing URL")
	}
	if c.MaxDepth > 0 && c.MaxDepth < depth {
		return errors.New("Max depth limit reached")
	}
	if !c.AllowURLRevisit {
		for _, u2 := range c.visitedURLs {
			if u2 == u {
				return errors.New("URL already visited")
			}
		}
		c.lock.Lock()
		c.visitedURLs = append(c.visitedURLs, u)
		c.lock.Unlock()
	}
	return nil
}

func (c *Collector) isDomainAllowed(domain string) bool {
	if c.AllowedDomains == nil || len(c.AllowedDomains) == 0 {
		return true
	}
	for _, d2 := range c.AllowedDomains {
		if d2 == domain {
			return true
		}
	}
	return false
}

// Wait returns when the collector jobs are finished
func (c *Collector) Wait() {
	c.wg.Wait()
}

// OnRequest registers a function. Function will be executed on every
// request made by the Collector
func (c *Collector) OnRequest(f RequestCallback) {
	c.lock.Lock()
	c.requestCallbacks = append(c.requestCallbacks, f)
	c.lock.Unlock()
}

// OnResponse registers a function. Function will be executed on every response
func (c *Collector) OnResponse(f ResponseCallback) {
	c.lock.Lock()
	c.responseCallbacks = append(c.responseCallbacks, f)
	c.lock.Unlock()
}

// OnHTML registers a function. Function will be executed on every HTML
// element matched by the `goquerySelector` parameter.
// `goquerySelector` is a selector used by https://github.com/PuerkitoBio/goquery
func (c *Collector) OnHTML(goquerySelector string, f HTMLCallback) {
	c.lock.Lock()
	c.htmlCallbacks[goquerySelector] = f
	c.lock.Unlock()
}

// WithTransport allows you to set a custom http.Transport for this collector.
func (c *Collector) WithTransport(transport *http.Transport) {
	c.backend.Client.Transport = transport
}

// DisableCookies turns off cookie handling for this collector
func (c *Collector) DisableCookies() {
	c.backend.Client.Jar = nil
}

// SetRequestTimeout overrides the default timeout (10 seconds) for this collector
func (c *Collector) SetRequestTimeout(timeout time.Duration) {
	c.backend.Client.Timeout = timeout
}

func (c *Collector) handleOnRequest(r *Request) {
	for _, f := range c.requestCallbacks {
		f(r)
	}
}

func (c *Collector) handleOnResponse(r *Response) {
	for _, f := range c.responseCallbacks {
		f(r)
	}
}

func (c *Collector) handleOnHTML(req *Request, resp *Response) {
	if strings.Index(strings.ToLower(resp.Headers.Get("Content-Type")), "html") == -1 {
		return
	}
	doc, err := goquery.NewDocumentFromReader(bytes.NewBuffer(resp.Body))
	if err != nil {
		return
	}
	for expr, f := range c.htmlCallbacks {
		doc.Find(expr).Each(func(i int, s *goquery.Selection) {
			for _, n := range s.Nodes {
				e := &HTMLElement{
					Name:       n.Data,
					Request:    req,
					Response:   resp,
					Text:       goquery.NewDocumentFromNode(n).Text(),
					DOM:        s,
					attributes: n.Attr,
				}
				f(e)
			}
		})
	}
}

// Limit adds a new `LimitRule` to the collector
func (c *Collector) Limit(rule *LimitRule) error {
	return c.backend.Limit(rule)
}

// Limits adds new `LimitRule`s to the collector
func (c *Collector) Limits(rules []*LimitRule) error {
	return c.backend.Limits(rules)
}

// SetCookies handles the receipt of the cookies in a reply for the given URL
func (c *Collector) SetCookies(URL string, cookies []*http.Cookie) error {
	if c.backend.Client.Jar == nil {
		return errors.New("Cookie jar is not available")
	}
	u, err := url.Parse(URL)
	if err != nil {
		return err
	}
	c.backend.Client.Jar.SetCookies(u, cookies)
	return nil
}

// Cookies returns the cookies to send in a request for the given URL.
func (c *Collector) Cookies(URL string) []*http.Cookie {
	if c.backend.Client.Jar == nil {
		return nil
	}
	u, err := url.Parse(URL)
	if err != nil {
		return nil
	}
	return c.backend.Client.Jar.Cookies(u)
}

// Attr returns the selected attribute of a HTMLElement or empty string
// if no attribute found
func (h *HTMLElement) Attr(k string) string {
	for _, a := range h.attributes {
		if a.Key == k {
			return a.Val
		}
	}
	return ""
}

// AbsoluteURL returns with the resolved absolute URL of an URL chunk.
// AbsoluteURL returns empty string if the URL chunk is a fragment or
// could not be parsed
func (r *Request) AbsoluteURL(u string) string {
	if strings.HasPrefix(u, "#") {
		return ""
	}
	absURL, err := r.URL.Parse(u)
	if err != nil {
		return ""
	}
	absURL.Fragment = ""
	if absURL.Scheme == "//" {
		absURL.Scheme = r.URL.Scheme
	}
	return absURL.String()
}

// Visit continues Collector's collecting job by creating a
// request and preserves the Context of the previous request.
// Visit also calls the previously provided OnRequest,
// OnResponse, OnHTML callbacks
func (r *Request) Visit(URL string) error {
	return r.collector.scrape(r.AbsoluteURL(URL), "GET", r.Depth+1, nil, r.Ctx)
}

// Post continues a collector job by creating a POST request and preserves the Context
// of the previous request.
// Post also calls the previously provided OnRequest, OnResponse, OnHTML callbacks
func (r *Request) Post(URL string, requestData map[string]string) error {
	return r.collector.scrape(r.AbsoluteURL(URL), "POST", r.Depth+1, createFormReader(requestData), r.Ctx)
}

// PostRaw starts a collector job by creating a POST request with raw binary data.
// PostRaw preserves the Context of the previous request
// and calls the previously provided callbacks
func (r *Request) PostRaw(URL string, requestData []byte) error {
	return r.collector.scrape(r.AbsoluteURL(URL), "POST", r.Depth+1, bytes.NewReader(requestData), r.Ctx)
}

// UnmarshalBinary decodes Context value to nil
// This function is used by request caching
func (c *Context) UnmarshalBinary(_ []byte) error {
	return nil
}

// MarshalBinary encodes Context value
// This function is used by request caching
func (c *Context) MarshalBinary() (_ []byte, _ error) {
	return nil, nil
}

// Put stores a value in Context
func (c *Context) Put(key, value string) {
	c.lock.Lock()
	c.contextMap[key] = value
	c.lock.Unlock()
}

// Get retrieves a value from Context. If no value found for `k`
// Get returns an empty string if key not found
func (c *Context) Get(key string) string {
	c.lock.RLock()
	if v, ok := c.contextMap[key]; ok {
		c.lock.RUnlock()
		return v
	}
	c.lock.RUnlock()
	return ""
}

func createFormReader(data map[string]string) io.Reader {
	form := url.Values{}
	for k, v := range data {
		form.Add(k, v)
	}
	return strings.NewReader(form.Encode())
}

func (r *Response) fixCharset() {
	contentType := strings.ToLower(r.Headers.Get("Content-Type"))
	if strings.Index(contentType, "charset") == -1 {
		return
	}
	if strings.Index(contentType, "utf-8") != -1 || strings.Index(contentType, "utf8") != -1 {
		return
	}
	encodedBodyReader, err := charset.NewReader(bytes.NewReader(r.Body), contentType)
	if err != nil {
		return
	}
	tmpBody, err := ioutil.ReadAll(encodedBodyReader)
	if err != nil {
		return
	}
	r.Body = tmpBody
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"path/filepath"
	"sync"
	"time"
)

// conch holds a file info group, and done, out, and err channels.
type conch struct {
	fig  *fileInfoGroup
	done chan struct{}
	out  chan *fileOutput
	err  chan error
}

// newConch returns a pointer to a new conch with channels setup.
func newConch(fig *fileInfoGroup) *conch {
	return &conch{
		fig:  fig,
		done: make(chan struct{}),
		out:  make(chan *fileOutput),
		err:  make(chan error, 1),
	}
}

// doneChan returns the conch done channel.
func (c *conch) doneChan() chan struct{} {
	return c.done
}

// feedPaths is a generator that sends paths out for processing. If canceled,
// an error is returned.
func (c *conch) feedPaths(paths chan string) {
	for _, v := range c.fig.fsi {
		select {
		case paths <- filepath.Join(c.fig.dir, v.Name()):
		case <-c.done:
			c.err <- errors.New("canceled")
		}
	}

	close(paths)
}

// digest processes the file located at the currently provided path, and
// sends out a new result.
func (c *conch) digest(paths <-chan string) {
	for p := range paths {
		r := newFileOutput(p)

		if slow {
			select {
			case <-time.After(time.Second):
			case <-c.done:
				return
			}
		}

		select {
		case c.out <- r:
		case <-c.done:
			return
		}
	}
}

// fanout sets-up digest goroutines according to width. Each digest goroutine
// waits for data from the paths channel and the entire function collapses when
// completed.
func (c *conch) fanout(paths chan string) {
	var wg sync.WaitGroup
	wg.Add(width)

	// setup digesters by width
	for i := 0; i < width; i++ {
		go func() {
			c.digest(paths)
			wg.Done()
		}()
	}

	wg.Wait()
	close(c.out)
}

// run wires up the paths chan, calls the path generator, calls the fanout, and
// returns the out and err channels.
func (c *conch) run() (<-chan *fileOutput, <-chan error) {
	paths := make(chan string)

	go c.feedPaths(paths)
	go c.fanout(paths)

	return c.out, c.err
}
<commit_msg>Ensure that paths are closed in all cases<commit_after>package main

import (
	"errors"
	"path/filepath"
	"sync"
	"time"
)

// conch holds a file info group, and done, out, and err channels.
type conch struct {
	fig  *fileInfoGroup
	done chan struct{}
	out  chan *fileOutput
	err  chan error
}

// newConch returns a pointer to a new conch with channels setup.
func newConch(fig *fileInfoGroup) *conch {
	return &conch{
		fig:  fig,
		done: make(chan struct{}),
		out:  make(chan *fileOutput),
		err:  make(chan error, 1),
	}
}

// doneChan returns the conch done channel.
func (c *conch) doneChan() chan struct{} {
	return c.done
}

// feedPaths is a generator that sends paths out for processing. If canceled,
// an error is returned.
func (c *conch) feedPaths(paths chan string) {
	defer close(paths)
	for _, v := range c.fig.fsi {
		select {
		case paths <- filepath.Join(c.fig.dir, v.Name()):
		case <-c.done:
			c.err <- errors.New("canceled")
			return
		}
	}
}

// digest processes the file located at the currently provided path, and
// sends out a new result.
func (c *conch) digest(paths <-chan string) {
	for p := range paths {
		r := newFileOutput(p)

		if slow {
			select {
			case <-time.After(time.Second):
			case <-c.done:
				return
			}
		}

		select {
		case c.out <- r:
		case <-c.done:
			return
		}
	}
}

// fanout sets-up digest goroutines according to width. Each digest goroutine
// waits for data from the paths channel and the entire function collapses when
// completed.
func (c *conch) fanout(paths chan string) {
	var wg sync.WaitGroup
	wg.Add(width)

	// setup digesters by width
	for i := 0; i < width; i++ {
		go func() {
			c.digest(paths)
			wg.Done()
		}()
	}

	wg.Wait()
	close(c.out)
}

// run wires up the paths chan, calls the path generator, calls the fanout, and
// returns the out and err channels.
func (c *conch) run() (<-chan *fileOutput, <-chan error) {
	paths := make(chan string)

	go c.feedPaths(paths)
	go c.fanout(paths)

	return c.out, c.err
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"syscall"

	"github.com/kelseyhightower/confd/backends"
	"github.com/kelseyhightower/confd/log"
	"github.com/kelseyhightower/confd/resource/template"
)

func main() {
	flag.Parse()
	if printVersion {
		fmt.Printf("confd %s\n", Version)
		os.Exit(0)
	}
	if err := initConfig(); err != nil {
		log.Fatal(err.Error())
	}
	log.Notice("Starting confd")
	storeClient, err := backends.New(backendsConfig)
	if err != nil {
		log.Fatal(err.Error())
	}
	templateConfig.StoreClient = storeClient
	if onetime {
		if err := template.Process(templateConfig); err != nil {
			os.Exit(1)
		}
		os.Exit(0)
	}
	stopChan := make(chan bool)
	doneChan := make(chan bool)
	errChan := make(chan error, 10)
	var processor template.Processor
	switch {
	case watch:
		processor = template.WatchProcessor(templateConfig, stopChan, doneChan, errChan)
	default:
		processor = template.IntervalProcessor(templateConfig, stopChan, doneChan, errChan, config.Interval)
	}
	go processor.Process()
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
	for {
		select {
		case err := <-errChan:
			log.Error(err.Error())
		case s := <-signalChan:
			log.Info(fmt.Sprintf("Captured %v. Exiting...", s))
			close(doneChan)
		case <-doneChan:
			os.Exit(0)
		}
	}
}
<commit_msg>Enable watch from the config file<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"syscall"

	"github.com/kelseyhightower/confd/backends"
	"github.com/kelseyhightower/confd/log"
	"github.com/kelseyhightower/confd/resource/template"
)

func main() {
	flag.Parse()
	if printVersion {
		fmt.Printf("confd %s\n", Version)
		os.Exit(0)
	}
	if err := initConfig(); err != nil {
		log.Fatal(err.Error())
	}
	log.Notice("Starting confd")
	storeClient, err := backends.New(backendsConfig)
	if err != nil {
		log.Fatal(err.Error())
	}
	templateConfig.StoreClient = storeClient
	if onetime {
		if err := template.Process(templateConfig); err != nil {
			os.Exit(1)
		}
		os.Exit(0)
	}
	stopChan := make(chan bool)
	doneChan := make(chan bool)
	errChan := make(chan error, 10)
	var processor template.Processor
	switch {
	case config.Watch:
		processor = template.WatchProcessor(templateConfig, stopChan, doneChan, errChan)
	default:
		processor = template.IntervalProcessor(templateConfig, stopChan, doneChan, errChan, config.Interval)
	}
	go processor.Process()
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
	for {
		select {
		case err := <-errChan:
			log.Error(err.Error())
		case s := <-signalChan:
			log.Info(fmt.Sprintf("Captured %v. Exiting...", s))
			close(doneChan)
		case <-doneChan:
			os.Exit(0)
		}
	}
}
<|endoftext|>