<commit_before>package main

import (
	"bytes"
	"io"
	"io/ioutil"
	"launchpad.net/goyaml"
	"path/filepath"
	"strings"

	"github.com/russross/blackfriday"
)

// A Page represents the key-value pairs in a page or posts front-end YAML as
// well as the markup in the body.
type Page map[string]interface{}

// ParsePage will parse a file with front-end YAML and markup content, and
// return a key-value Page structure.
func ParsePage(fn string) (Page, error) {
	c, err := ioutil.ReadFile(fn)
	if err != nil {
		return nil, err
	}
	return parsePage(fn, c)
}

// Helper function that creates a new Page from a byte array, parsing the
// front-end YAML and the markup, and pre-calculating all page-level variables.
func parsePage(fn string, c []byte) (Page, error) {

	page, err := parseMatter(c) //map[string] interface{} { }
	if err != nil {
		return nil, err
	}

	ext := filepath.Ext(fn)
	ext_output := ext
	markdown := isMarkdown(fn)

	// if markdown, change the output extension to html
	if markdown {
		ext_output = ".html"
	}

	page["ext"] = ext
	page["output_ext"] = ext_output
	page["id"] = removeExt(fn)
	page["url"] = replaceExt(fn, ext_output)

	// if markdown, convert to html
	raw := parseContent(c)
	if markdown {
		page["content"] = string(blackfriday.MarkdownCommon(raw))
	} else {
		page["content"] = string(raw)
	}

	if page["layout"] == "" {
		page["layout"] = "default"
	}

	// according to spec, Jekyll allows user to enter either category or
	// categories. Convert single category to string array to be consistent ...
	if category := page.GetString("category"); category != "" {
		page["categories"] = []string{ category }
		delete(page, "category")
	}

	return page, nil
}

// Helper function to parse the front-end yaml matter.
func parseMatter(content []byte) (Page, error) {
	page := map[string] interface{} { }
	err := goyaml.Unmarshal(content, &page)
	return page, err
}

// Helper function that separates the front-end yaml from the markup, and
// and returns only the markup (content) as a byte array.
func parseContent(content []byte) []byte {
	//now we need to parse out the markdown section create
	//buffered reader
	b := bytes.NewBuffer(content)
	m := new(bytes.Buffer)
	streams := 0

	//read each line of the file and read the markdown section
	//which is the second document stream in the yaml file
	parse : for {
		line, err := b.ReadString('\n')
		switch {
		case err == io.EOF && streams >= 2:
			m.WriteString(line)
			break parse
		case err == io.EOF:
			break parse
		case err != nil :
			return nil
		case streams >= 2:
			m.WriteString(line)
		case strings.HasPrefix(line, "---") :
			streams++
		}
	}

	return m.Bytes()
}

// Sets a parameter value.
func (p Page) Set(key string, val interface{}) {
	p[key] = val
}

// Gets a parameter value.
func (p Page) Get(key string) interface{} {
	return p[key]
}

// Gets a parameter value as a string. If none exists return an empty string.
func (p Page) GetString(key string) (str string) {
	if v, ok := p[key]; ok {
		str = v.(string)
	}
	return
}

// Gets a parameter value as a string array.
func (p Page) GetStrings(key string) (strs []string) {
	if v, ok := p[key]; ok {
		strs = v.([]string)
	}
	return []string{}
}

// Gets a parameter value as a byte array.
func (p Page) GetBytes(key string) (b []byte) {
	if v, ok := p[key]; ok {
		b = v.([]byte)
	}
	return
}

// Gets the layout file to use, without the extension.
// Layout files must be placed in the _layouts directory.
func (p Page) GetLayout() string {
	return p.GetString("layout")
}

// Gets the title of the Page.
func (p Page) GetTitle() string {
	return p.GetString("title")
}

// Gets the URL / relative path of the Page.
// e.g. /2008/12/14/my-post.html
func (p Page) GetUrl() string {
	return p.GetString("url")
}

// Gets the Extension of the File (.html, .md, etc)
func (p Page) GetExt() string {
	return p.GetString("ext")
}

// Gets the un-rendered content of the Page.
func (p Page) GetContent() (c string) {
	if v, ok := p["content"]; ok {
		c = v.(string)
	}
	return
}

// Gets the list of tags to which this Post belongs.
func (p Page) GetTags() []string {
	return p.GetStrings("tags")
}

// Gets the list of categories to which this post belongs.
func (p Page) GetCategories() []string {
	return p.GetStrings("categories")
}

<commit_msg>GetStrings crashed if tags or categories were set, always returned empty.<commit_after>package main

import (
	"bytes"
	"io"
	"io/ioutil"
	"launchpad.net/goyaml"
	"path/filepath"
	"strings"

	"github.com/russross/blackfriday"
)

// A Page represents the key-value pairs in a page or posts front-end YAML as
// well as the markup in the body.
type Page map[string]interface{}

// ParsePage will parse a file with front-end YAML and markup content, and
// return a key-value Page structure.
func ParsePage(fn string) (Page, error) {
	c, err := ioutil.ReadFile(fn)
	if err != nil {
		return nil, err
	}
	return parsePage(fn, c)
}

// Helper function that creates a new Page from a byte array, parsing the
// front-end YAML and the markup, and pre-calculating all page-level variables.
func parsePage(fn string, c []byte) (Page, error) {

	page, err := parseMatter(c) //map[string] interface{} { }
	if err != nil {
		return nil, err
	}

	ext := filepath.Ext(fn)
	ext_output := ext
	markdown := isMarkdown(fn)

	// if markdown, change the output extension to html
	if markdown {
		ext_output = ".html"
	}

	page["ext"] = ext
	page["output_ext"] = ext_output
	page["id"] = removeExt(fn)
	page["url"] = replaceExt(fn, ext_output)

	// if markdown, convert to html
	raw := parseContent(c)
	if markdown {
		page["content"] = string(blackfriday.MarkdownCommon(raw))
	} else {
		page["content"] = string(raw)
	}

	if page["layout"] == "" {
		page["layout"] = "default"
	}

	// according to spec, Jekyll allows user to enter either category or
	// categories. Convert single category to string array to be consistent ...
	if category := page.GetString("category"); category != "" {
		page["categories"] = []string{ category }
		delete(page, "category")
	}

	return page, nil
}

// Helper function to parse the front-end yaml matter.
func parseMatter(content []byte) (Page, error) {
	page := map[string] interface{} { }
	err := goyaml.Unmarshal(content, &page)
	return page, err
}

// Helper function that separates the front-end yaml from the markup, and
// and returns only the markup (content) as a byte array.
func parseContent(content []byte) []byte {
	//now we need to parse out the markdown section create
	//buffered reader
	b := bytes.NewBuffer(content)
	m := new(bytes.Buffer)
	streams := 0

	//read each line of the file and read the markdown section
	//which is the second document stream in the yaml file
	parse : for {
		line, err := b.ReadString('\n')
		switch {
		case err == io.EOF && streams >= 2:
			m.WriteString(line)
			break parse
		case err == io.EOF:
			break parse
		case err != nil :
			return nil
		case streams >= 2:
			m.WriteString(line)
		case strings.HasPrefix(line, "---") :
			streams++
		}
	}

	return m.Bytes()
}

// Sets a parameter value.
func (p Page) Set(key string, val interface{}) {
	p[key] = val
}

// Gets a parameter value.
func (p Page) Get(key string) interface{} {
	return p[key]
}

// Gets a parameter value as a string. If none exists return an empty string.
func (p Page) GetString(key string) (str string) {
	if v, ok := p[key]; ok {
		str = v.(string)
	}
	return
}

// Gets a parameter value as a string array.
func (p Page) GetStrings(key string) (strs []string) {
	if v, ok := p[key]; ok {
		for _, s := range v.([]interface{}) {
			strs = append(strs, s.(string))
		}
	}
	return
}

// Gets a parameter value as a byte array.
func (p Page) GetBytes(key string) (b []byte) {
	if v, ok := p[key]; ok {
		b = v.([]byte)
	}
	return
}

// Gets the layout file to use, without the extension.
// Layout files must be placed in the _layouts directory.
func (p Page) GetLayout() string {
	return p.GetString("layout")
}

// Gets the title of the Page.
func (p Page) GetTitle() string {
	return p.GetString("title")
}

// Gets the URL / relative path of the Page.
// e.g. /2008/12/14/my-post.html
func (p Page) GetUrl() string {
	return p.GetString("url")
}

// Gets the Extension of the File (.html, .md, etc)
func (p Page) GetExt() string {
	return p.GetString("ext")
}

// Gets the un-rendered content of the Page.
func (p Page) GetContent() (c string) {
	if v, ok := p["content"]; ok {
		c = v.(string)
	}
	return
}

// Gets the list of tags to which this Post belongs.
func (p Page) GetTags() []string {
	return p.GetStrings("tags")
}

// Gets the list of categories to which this post belongs.
func (p Page) GetCategories() []string {
	return p.GetStrings("categories")
}

<|endoftext|>


<commit_before>package gopass

import (
	"os"
)

// getPasswd returns the input read from terminal.
// If masked is true, typing will be matched by asterisks on the screen.
// Otherwise, typing will echo nothing.
func getPasswd(masked bool) []byte {
	var pass, bs, mask []byte
	if masked {
		bs = []byte("\b \b")
		mask = []byte("*")
	}

	for {
		if v := getch(); v == 127 || v == 8 {
			if l := len(pass); l > 0 {
				pass = pass[:l-1]
				os.Stdout.Write(bs)
			}
		} else if v == 13 || v == 10 {
			break
		} else if v != 0 {
			pass = append(pass, v)
			os.Stdout.Write(mask)
		}
	}
	os.Stdout.WriteString(lineEnding)
	println()
	return pass
}

// GetPasswd returns the password read from the terminal without echoing input.
// The returned byte array does not include end-of-line characters.
func GetPasswd() []byte {
	return getPasswd(false)
}

// GetPasswdMasked returns the password read from the terminal, echoing asterisks.
// The returned byte array does not include end-of-line characters.
func GetPasswdMasked() []byte {
	return getPasswd(true)
}
<commit_msg>Handle Ctrl+C, Fixes #7<commit_after>package gopass

import (
	"os"
)

var (
	ErrInterrupted = errors.New("Interrupted")
)

// getPasswd returns the input read from terminal.
// If masked is true, typing will be matched by asterisks on the screen.
// Otherwise, typing will echo nothing.
func getPasswd(masked bool) ([]byte, error) {
	var err error
	var pass, bs, mask []byte
	if masked {
		bs = []byte("\b \b")
		mask = []byte("*")
	}

	for {
		if v := getch(); v == 127 || v == 8 {
			if l := len(pass); l > 0 {
				pass = pass[:l-1]
				os.Stdout.Write(bs)
			}
		} else if v == 13 || v == 10 {
			break
		} else if v == 3 {
			err = ErrInterrupted
			break
		} else if v != 0 {
			pass = append(pass, v)
			os.Stdout.Write(mask)
		}
	}
	os.Stdout.WriteString(lineEnding)
	println()
	return pass, err
}

// GetPasswd returns the password read from the terminal without echoing input.
// The returned byte array does not include end-of-line characters.
func GetPasswd() ([]byte, error) {
	return getPasswd(false)
}

// GetPasswdMasked returns the password read from the terminal, echoing asterisks.
// The returned byte array does not include end-of-line characters.
func GetPasswdMasked() ([]byte, error) {
	return getPasswd(true)
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"strings"
)

func init() {
	home := os.Getenv("HOME")
	paths := []string{
		home + "/bin",
		home + "/.vvmn/vim/current/bin",
		"/usr/local/bin",
		"/usr/local/opt/coreutils/libexec/gnubin",
	}
	setPath(paths...)
}

func setPath(args ...string) {
	s := os.Getenv("PATH")
	paths := strings.Split(s, ":")
	var newPaths []string
	for _, arg := range args {
		for _, path := range paths {
			if path != arg {
				newPaths = append(newPaths, path)
			}
		}
	}
	newPaths = append(args, newPaths...)
	os.Setenv("PATH", strings.Join(newPaths, ":"))
}
<commit_msg>Make setPath set $PATH correctly<commit_after>package main

import (
	"os"
	"strings"
)

func init() {
	home := os.Getenv("HOME")
	paths := []string{
		home + "/bin",
		home + "/.vvmn/vim/current/bin",
		"/usr/local/bin",
		"/usr/local/opt/coreutils/libexec/gnubin",
	}
	setPath(paths...)
}

func setPath(args ...string) {
	s := os.Getenv("PATH")
	paths := strings.Split(s, ":")
	var newPaths []string
	for _, path := range paths {
		for _, arg := range args {
			if path != arg {
				newPaths = append(newPaths, path)
			}
		}
	}
	newPaths = append(args, newPaths...)
	os.Setenv("PATH", strings.Join(newPaths, ":"))
}
<|endoftext|>


<commit_before>// Package pcap is a wrapper around the pcap library.
package pcap

/*
#cgo LDFLAGS: -lpcap
#include <stdlib.h>
#include <pcap.h>
*/
import "C"

import (
	"errors"
	"fmt"
	"net"
	"syscall"
	"time"
	"unsafe"
)

// Pcap wraps a pcap_t struct.
type Pcap struct {
	cptr     *C.pcap_t
	sampling int // 1:N; a bit of a hack
}

type pcapError struct{ string }

// Stat records statistics about packets received and dropped.
type Stat struct {
	PacketsReceived  uint32
	PacketsDropped   uint32
	PacketsIfDropped uint32
}

// Interface describes a single network interface on a host.
type Interface struct {
	Name        string
	Description string
	Addresses   []IFAddress
	// TODO: add more elements
}

// IFAddress is the interface address.
type IFAddress struct {
	IP      net.IP
	Netmask net.IPMask
	// TODO: add broadcast + PtP dst ?
}

// Version returns the current pcap library version.
func Version() string { return C.GoString(C.pcap_lib_version()) }

// Datalink TODO
func (p *Pcap) Datalink() int { return int(C.pcap_datalink(p.cptr)) }

func (e *pcapError) Error() string { return e.string }

// Geterror converts the last pcap error to a Go error.
func (p *Pcap) Geterror() error { return &pcapError{C.GoString(C.pcap_geterr(p.cptr))} }

// Next wraps libpcap NextEx
func (p *Pcap) Next() (pkt *Packet) { rv, _ := p.NextEx(); return rv }

// Create TODO
func Create(device string) (*Pcap, error) {
	dev := C.CString(device)
	defer C.free(unsafe.Pointer(dev))

	buf := (*C.char)(C.calloc(errbufSize, 1))
	defer C.free(unsafe.Pointer(buf))

	cptr := C.pcap_create(dev, buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr:     cptr,
		sampling: 1,
	}, nil
}

// SetBufferSize sets buffer size (in bytes) on the activated handle.
func (p *Pcap) SetBufferSize(sz int32) error {
	if C.pcap_set_buffer_size(p.cptr, C.int(sz)) != 0 {
		return p.Geterror()
	}
	return nil
}

// SetPromisc sets promiscuous mode on the handle. It should be called before
// activation.
func (p *Pcap) SetPromisc(promisc bool) error {
	pro := int32(0)
	if promisc {
		pro = 1
	}

	if C.pcap_set_promisc(p.cptr, C.int(pro)) != 0 {
		return p.Geterror()
	}
	return nil
}

// SetSnapLen TODO
func (p *Pcap) SetSnapLen(s int32) error {
	if C.pcap_set_snaplen(p.cptr, C.int(s)) != 0 {
		return p.Geterror()
	}
	return nil
}

// SetReadTimeout (milliseconds) that will be used on a capture handle when it
// is activated.
func (p *Pcap) SetReadTimeout(toMs int32) error {
	if C.pcap_set_timeout(p.cptr, C.int(toMs)) != 0 {
		return p.Geterror()
	}
	return nil
}

// Activate a packet capture handle to look at packets on the network, with the
// options that were set on the handle being in effect.
func (p *Pcap) Activate() error {
	if C.pcap_activate(p.cptr) != 0 {
		return p.Geterror()
	}
	return nil
}

// OpenLive opens a device and returns a handler.
func OpenLive(device string, snaplen int32, promisc bool, timeoutMS int32) (*Pcap, error) {
	buf := (*C.char)(C.calloc(errbufSize, 1))
	defer C.free(unsafe.Pointer(buf))

	dev := C.CString(device)
	defer C.free(unsafe.Pointer(dev))

	pro := int32(0)
	if promisc {
		pro = 1
	}

	cptr := C.pcap_open_live(dev, C.int(snaplen), C.int(pro), C.int(timeoutMS), buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr:     cptr,
		sampling: 1,
	}, nil
}

// OpenOffline provides a Pcap over a .pcap file.
func OpenOffline(file string) (*Pcap, error) {
	buf := (*C.char)(C.calloc(errbufSize, 1))
	defer C.free(unsafe.Pointer(buf))

	cf := C.CString(file)
	defer C.free(unsafe.Pointer(cf))

	cptr := C.pcap_open_offline(cf, buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr:     cptr,
		sampling: 1,
	}, nil
}

// Close calls pcap_close on the underlying pcap_t.
func (p *Pcap) Close() {
	C.pcap_close(p.cptr)
}

// SetSampling sets the sample rate of the handle. We perform sampling in the
// Go wrapper library, as efficiently as we can at that layer.
func (p *Pcap) SetSampling(rate float64) {
	// Take 1 packet every N (approximate).
	//  rate=0.50 sampling=1.0/0.50=2
	//  rate=0.33 sampling=1.0/0.33=3
	//  rate=0.15 sampling=1.0/0.15=6.66=6 (alas)
	p.sampling = int(1.0 / rate)
}

// NextEx gets the next packet on the handle.
func (p *Pcap) NextEx() (*Packet, int32) {
	var pkthdr *C.struct_pcap_pkthdr
	var bufPtr *C.u_char
	var result int32
	for i := 0; i < p.sampling; i++ {
		// "The struct pcap_pkthdr and the packet data are not to be freed by
		// the caller, and are not guaranteed to be valid after the next call
		// ... if the code needs them to remain valid, it must make a copy of
		// them." --pcap_next_ex(3)
		//
		// I believe this is not a terrible way to do sampling. We still incur
		// the cost of pcap_next_ex, which does copy from kernel to userspace,
		// but we skip any further allocations and a lot of analysis in our
		// client code.
		result = int32(C.pcap_next_ex(p.cptr, &pkthdr, &bufPtr))
	}

	buf := unsafe.Pointer(bufPtr)
	if buf == nil {
		return nil, result
	}

	pkt := &Packet{
		Time:   time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec)),
		Caplen: uint32(pkthdr.caplen),
		Len:    uint32(pkthdr.len),
		Data:   C.GoBytes(buf, C.int(pkthdr.caplen)),
	}
	return pkt, result
}

// Getstats TODO
func (p *Pcap) Getstats() (*Stat, error) {
	var cstats _Ctype_struct_pcap_stat
	if C.pcap_stats(p.cptr, &cstats) == -1 {
		return nil, p.Geterror()
	}

	return &Stat{
		PacketsReceived:  uint32(cstats.ps_recv),
		PacketsDropped:   uint32(cstats.ps_drop),
		PacketsIfDropped: uint32(cstats.ps_ifdrop),
	}, nil
}

// SetFilter TODO
func (p *Pcap) SetFilter(expr string) error {
	var bpf _Ctype_struct_bpf_program
	cexpr := C.CString(expr)
	defer C.free(unsafe.Pointer(cexpr))

	if C.pcap_compile(p.cptr, &bpf, cexpr, 1, 0) == -1 {
		return p.Geterror()
	}
	defer C.pcap_freecode(&bpf)

	if C.pcap_setfilter(p.cptr, &bpf) == -1 {
		return p.Geterror()
	}
	return nil
}

// SetDatalink TODO
func (p *Pcap) SetDatalink(dlt int) error {
	if C.pcap_set_datalink(p.cptr, C.int(dlt)) == -1 {
		return p.Geterror()
	}
	return nil
}

// DatalinkValueToName string
func DatalinkValueToName(dlt int) string {
	if name := C.pcap_datalink_val_to_name(C.int(dlt)); name != nil {
		return C.GoString(name)
	}
	return ""
}

// DatalinkValueToDescription TODO
func DatalinkValueToDescription(dlt int) string {
	if desc := C.pcap_datalink_val_to_description(C.int(dlt)); desc != nil {
		return C.GoString(desc)
	}
	return ""
}

// FindAllDevs TODO
func FindAllDevs() ([]Interface, error) {
	buf := (*C.char)(C.calloc(errbufSize, 1))
	defer C.free(unsafe.Pointer(buf))

	var alldevsp *C.pcap_if_t
	if C.pcap_findalldevs((**C.pcap_if_t)(&alldevsp), buf) == -1 {
		return nil, errors.New(C.GoString(buf))
	}
	defer C.pcap_freealldevs((*C.pcap_if_t)(alldevsp))

	ifs := []Interface{}
	for dev := alldevsp; dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		ifs = append(ifs, Interface{
			Name:        C.GoString(dev.name),
			Description: C.GoString(dev.description),
			Addresses:   findAllAddresses(dev.addresses),
			// TODO: add more elements
		})
	}
	return ifs, nil
}

func findAllAddresses(addresses *_Ctype_struct_pcap_addr) []IFAddress {
	// TODO - make it support more than IPv4 and IPv6?
	a := []IFAddress{}
	for curaddr := addresses; curaddr != nil; curaddr = (*_Ctype_struct_pcap_addr)(curaddr.next) {
		if curaddr.addr == nil {
			continue
		}

		ip, err := sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr)))
		if err != nil {
			continue
		}
		netmask, err := sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr)))
		if err != nil {
			continue
		}

		a = append(a, IFAddress{
			IP:      ip,
			Netmask: netmask,
		})
	}
	return a
}

func sockaddrToIP(rsa *syscall.RawSockaddr) ([]byte, error) {
	switch rsa.Family {
	case syscall.AF_INET:
		pp := (*syscall.RawSockaddrInet4)(unsafe.Pointer(rsa))
		ip := make([]byte, 4)
		for i := 0; i < len(ip); i++ {
			ip[i] = pp.Addr[i]
		}
		return ip, nil

	case syscall.AF_INET6:
		pp := (*syscall.RawSockaddrInet6)(unsafe.Pointer(rsa))
		ip := make([]byte, 16)
		for i := 0; i < len(ip); i++ {
			ip[i] = pp.Addr[i]
		}
		return ip, nil

	default:
		return nil, fmt.Errorf("unsupported address type %d", rsa.Family)
	}
}

// Inject TODO
func (p *Pcap) Inject(data []byte) error {
	buf := (*C.char)(C.malloc((C.size_t)(len(data))))
	defer C.free(unsafe.Pointer(buf))

	for i := 0; i < len(data); i++ {
		*(*byte)(unsafe.Pointer(uintptr(unsafe.Pointer(buf)) + uintptr(i))) = data[i]
	}

	if C.pcap_inject(p.cptr, unsafe.Pointer(buf), (C.size_t)(len(data))) == -1 {
		return p.Geterror()
	}
	return nil
}
<commit_msg>Link pcap static<commit_after>// Package pcap is a wrapper around the pcap library.
package pcap

/*
#cgo LDFLAGS: -Wl,--start-group,-Bstatic -lpcap -Wl,--end-group,-Bdynamic
#include <stdlib.h>
#include <pcap.h>
*/
import "C"

import (
	"errors"
	"fmt"
	"net"
	"syscall"
	"time"
	"unsafe"
)

// Pcap wraps a pcap_t struct.
type Pcap struct {
	cptr     *C.pcap_t
	sampling int // 1:N; a bit of a hack
}

type pcapError struct{ string }

// Stat records statistics about packets received and dropped.
type Stat struct {
	PacketsReceived  uint32
	PacketsDropped   uint32
	PacketsIfDropped uint32
}

// Interface describes a single network interface on a host.
type Interface struct {
	Name        string
	Description string
	Addresses   []IFAddress
	// TODO: add more elements
}

// IFAddress is the interface address.
type IFAddress struct {
	IP      net.IP
	Netmask net.IPMask
	// TODO: add broadcast + PtP dst ?
}

// Version returns the current pcap library version.
func Version() string { return C.GoString(C.pcap_lib_version()) }

// Datalink TODO
func (p *Pcap) Datalink() int { return int(C.pcap_datalink(p.cptr)) }

func (e *pcapError) Error() string { return e.string }

// Geterror converts the last pcap error to a Go error.
func (p *Pcap) Geterror() error { return &pcapError{C.GoString(C.pcap_geterr(p.cptr))} }

// Next wraps libpcap NextEx
func (p *Pcap) Next() (pkt *Packet) { rv, _ := p.NextEx(); return rv }

// Create TODO
func Create(device string) (*Pcap, error) {
	dev := C.CString(device)
	defer C.free(unsafe.Pointer(dev))

	buf := (*C.char)(C.calloc(errbufSize, 1))
	defer C.free(unsafe.Pointer(buf))

	cptr := C.pcap_create(dev, buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr:     cptr,
		sampling: 1,
	}, nil
}

// SetBufferSize sets buffer size (in bytes) on the activated handle.
func (p *Pcap) SetBufferSize(sz int32) error {
	if C.pcap_set_buffer_size(p.cptr, C.int(sz)) != 0 {
		return p.Geterror()
	}
	return nil
}

// SetPromisc sets promiscuous mode on the handle. It should be called before
// activation.
func (p *Pcap) SetPromisc(promisc bool) error {
	pro := int32(0)
	if promisc {
		pro = 1
	}

	if C.pcap_set_promisc(p.cptr, C.int(pro)) != 0 {
		return p.Geterror()
	}
	return nil
}

// SetSnapLen TODO
func (p *Pcap) SetSnapLen(s int32) error {
	if C.pcap_set_snaplen(p.cptr, C.int(s)) != 0 {
		return p.Geterror()
	}
	return nil
}

// SetReadTimeout (milliseconds) that will be used on a capture handle when it
// is activated.
func (p *Pcap) SetReadTimeout(toMs int32) error {
	if C.pcap_set_timeout(p.cptr, C.int(toMs)) != 0 {
		return p.Geterror()
	}
	return nil
}

// Activate a packet capture handle to look at packets on the network, with the
// options that were set on the handle being in effect.
func (p *Pcap) Activate() error {
	if C.pcap_activate(p.cptr) != 0 {
		return p.Geterror()
	}
	return nil
}

// OpenLive opens a device and returns a handler.
func OpenLive(device string, snaplen int32, promisc bool, timeoutMS int32) (*Pcap, error) {
	buf := (*C.char)(C.calloc(errbufSize, 1))
	defer C.free(unsafe.Pointer(buf))

	dev := C.CString(device)
	defer C.free(unsafe.Pointer(dev))

	pro := int32(0)
	if promisc {
		pro = 1
	}

	cptr := C.pcap_open_live(dev, C.int(snaplen), C.int(pro), C.int(timeoutMS), buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr:     cptr,
		sampling: 1,
	}, nil
}

// OpenOffline provides a Pcap over a .pcap file.
func OpenOffline(file string) (*Pcap, error) {
	buf := (*C.char)(C.calloc(errbufSize, 1))
	defer C.free(unsafe.Pointer(buf))

	cf := C.CString(file)
	defer C.free(unsafe.Pointer(cf))

	cptr := C.pcap_open_offline(cf, buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr:     cptr,
		sampling: 1,
	}, nil
}

// Close calls pcap_close on the underlying pcap_t.
func (p *Pcap) Close() {
	C.pcap_close(p.cptr)
}

// SetSampling sets the sample rate of the handle. We perform sampling in the
// Go wrapper library, as efficiently as we can at that layer.
func (p *Pcap) SetSampling(rate float64) {
	// Take 1 packet every N (approximate).
	//  rate=0.50 sampling=1.0/0.50=2
	//  rate=0.33 sampling=1.0/0.33=3
	//  rate=0.15 sampling=1.0/0.15=6.66=6 (alas)
	p.sampling = int(1.0 / rate)
}

// NextEx gets the next packet on the handle.
func (p *Pcap) NextEx() (*Packet, int32) {
	var pkthdr *C.struct_pcap_pkthdr
	var bufPtr *C.u_char
	var result int32
	for i := 0; i < p.sampling; i++ {
		// "The struct pcap_pkthdr and the packet data are not to be freed by
		// the caller, and are not guaranteed to be valid after the next call
		// ... if the code needs them to remain valid, it must make a copy of
		// them." --pcap_next_ex(3)
		//
		// I believe this is not a terrible way to do sampling. We still incur
		// the cost of pcap_next_ex, which does copy from kernel to userspace,
		// but we skip any further allocations and a lot of analysis in our
		// client code.
		result = int32(C.pcap_next_ex(p.cptr, &pkthdr, &bufPtr))
	}

	buf := unsafe.Pointer(bufPtr)
	if buf == nil {
		return nil, result
	}

	pkt := &Packet{
		Time:   time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec)),
		Caplen: uint32(pkthdr.caplen),
		Len:    uint32(pkthdr.len),
		Data:   C.GoBytes(buf, C.int(pkthdr.caplen)),
	}
	return pkt, result
}

// Getstats TODO
func (p *Pcap) Getstats() (*Stat, error) {
	var cstats _Ctype_struct_pcap_stat
	if C.pcap_stats(p.cptr, &cstats) == -1 {
		return nil, p.Geterror()
	}

	return &Stat{
		PacketsReceived:  uint32(cstats.ps_recv),
		PacketsDropped:   uint32(cstats.ps_drop),
		PacketsIfDropped: uint32(cstats.ps_ifdrop),
	}, nil
}

// SetFilter TODO
func (p *Pcap) SetFilter(expr string) error {
	var bpf _Ctype_struct_bpf_program
	cexpr := C.CString(expr)
	defer C.free(unsafe.Pointer(cexpr))

	if C.pcap_compile(p.cptr, &bpf, cexpr, 1, 0) == -1 {
		return p.Geterror()
	}
	defer C.pcap_freecode(&bpf)

	if C.pcap_setfilter(p.cptr, &bpf) == -1 {
		return p.Geterror()
	}
	return nil
}

// SetDatalink TODO
func (p *Pcap) SetDatalink(dlt int) error {
	if C.pcap_set_datalink(p.cptr, C.int(dlt)) == -1 {
		return p.Geterror()
	}
	return nil
}

// DatalinkValueToName string
func DatalinkValueToName(dlt int) string {
	if name := C.pcap_datalink_val_to_name(C.int(dlt)); name != nil {
		return C.GoString(name)
	}
	return ""
}

// DatalinkValueToDescription TODO
func DatalinkValueToDescription(dlt int) string {
	if desc := C.pcap_datalink_val_to_description(C.int(dlt)); desc != nil {
		return C.GoString(desc)
	}
	return ""
}

// FindAllDevs TODO
func FindAllDevs() ([]Interface, error) {
	buf := (*C.char)(C.calloc(errbufSize, 1))
	defer C.free(unsafe.Pointer(buf))

	var alldevsp *C.pcap_if_t
	if C.pcap_findalldevs((**C.pcap_if_t)(&alldevsp), buf) == -1 {
		return nil, errors.New(C.GoString(buf))
	}
	defer C.pcap_freealldevs((*C.pcap_if_t)(alldevsp))

	ifs := []Interface{}
	for dev := alldevsp; dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		ifs = append(ifs, Interface{
			Name:        C.GoString(dev.name),
			Description: C.GoString(dev.description),
			Addresses:   findAllAddresses(dev.addresses),
			// TODO: add more elements
		})
	}
	return ifs, nil
}

func findAllAddresses(addresses *_Ctype_struct_pcap_addr) []IFAddress {
	// TODO - make it support more than IPv4 and IPv6?
	a := []IFAddress{}
	for curaddr := addresses; curaddr != nil; curaddr = (*_Ctype_struct_pcap_addr)(curaddr.next) {
		if curaddr.addr == nil {
			continue
		}

		ip, err := sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr)))
		if err != nil {
			continue
		}
		netmask, err := sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr)))
		if err != nil {
			continue
		}

		a = append(a, IFAddress{
			IP:      ip,
			Netmask: netmask,
		})
	}
	return a
}

func sockaddrToIP(rsa *syscall.RawSockaddr) ([]byte, error) {
	switch rsa.Family {
	case syscall.AF_INET:
		pp := (*syscall.RawSockaddrInet4)(unsafe.Pointer(rsa))
		ip := make([]byte, 4)
		for i := 0; i < len(ip); i++ {
			ip[i] = pp.Addr[i]
		}
		return ip, nil

	case syscall.AF_INET6:
		pp := (*syscall.RawSockaddrInet6)(unsafe.Pointer(rsa))
		ip := make([]byte, 16)
		for i := 0; i < len(ip); i++ {
			ip[i] = pp.Addr[i]
		}
		return ip, nil

	default:
		return nil, fmt.Errorf("unsupported address type %d", rsa.Family)
	}
}

// Inject TODO
func (p *Pcap) Inject(data []byte) error {
	buf := (*C.char)(C.malloc((C.size_t)(len(data))))
	defer C.free(unsafe.Pointer(buf))

	for i := 0; i < len(data); i++ {
		*(*byte)(unsafe.Pointer(uintptr(unsafe.Pointer(buf)) + uintptr(i))) = data[i]
	}

	if C.pcap_inject(p.cptr, unsafe.Pointer(buf), (C.size_t)(len(data))) == -1 {
		return p.Geterror()
	}
	return nil
}
<|endoftext|>


<commit_before>package peco

import (
	"os"
	"reflect"
	"time"

	"golang.org/x/net/context"

	"github.com/google/btree"
	"github.com/lestrrat/go-pdebug"
	"github.com/peco/peco/hub"
	"github.com/peco/peco/internal/util"
	"github.com/peco/peco/pipeline"
	"github.com/peco/peco/sig"
	"github.com/pkg/errors"
)

const version = "v0.3.6"

type errIgnorable struct {
	err error
}

func (e errIgnorable) Ignorable() bool { return true }
func (e errIgnorable) Causer() error {
	return e.err
}
func (e errIgnorable) Error() string {
	return e.err.Error()
}
func makeIgnorable(err error) error {
	return errIgnorable{err: err}
}

// Inputseq is a list of keys that the user typed
type Inputseq []string

func (is *Inputseq) Add(s string) {
	*is = append(*is, s)
}

func (is Inputseq) KeyNames() []string {
	return is
}

func (is Inputseq) Len() int {
	return len(is)
}

func (is *Inputseq) Reset() {
	*is = []string(nil)
}

func New() *Peco {
	return &Peco{
		Argv:              os.Args,
		currentLineBuffer: NewMemoryBuffer(), // XXX revisit this
		queryExecDelay:    50 * time.Millisecond,
		readyCh:           make(chan struct{}),
		screen:            &Termbox{},
		selection:         NewSelection(),
	}
}

func (p *Peco) Ready() <-chan struct{} {
	return p.readyCh
}

func (p *Peco) Screen() Screen {
	return p.screen
}

func (p *Peco) Styles() *StyleSet {
	return &p.styles
}

func (p *Peco) Prompt() string {
	return p.prompt
}

func (p *Peco) Inputseq() *Inputseq {
	return &p.inputseq
}

func (p *Peco) Context() context.Context {
	return p.ctx
}

func (p *Peco) LayoutType() string {
	return p.layoutType
}

func (p *Peco) Location() *Location {
	return &p.location
}

func (p *Peco) ResultCh() chan Line {
	return p.resultCh
}

func (p *Peco) SetResultCh(ch chan Line) {
	p.resultCh = ch
}

func (p *Peco) Selection() *Selection {
	return p.selection
}

func (s RangeStart) Valid() bool {
	return s.valid
}

func (s RangeStart) Value() int {
	return s.val
}

func (s *RangeStart) SetValue(n int) {
	s.val = n
	s.valid = true
}

func (s *RangeStart) Reset() {
	s.valid = false
}

func (p *Peco) SelectionRangeStart() *RangeStart {
	return &p.selectionRangeStart
}

func (p *Peco) SingleKeyJumpShowPrefix() bool {
	return p.singleKeyJumpShowPrefix
}

func (p *Peco) SingleKeyJumpPrefixes() []rune {
	return p.singleKeyJumpPrefixes
}

func (p *Peco) SingleKeyJumpMode() bool {
	return p.singleKeyJumpMode
}

func (p *Peco) SetSingleKeyJumpMode(b bool) {
	p.singleKeyJumpMode = b
}

func (p *Peco) ToggleSingleKeyJumpMode() {
	p.singleKeyJumpMode = !p.singleKeyJumpMode
}

func (p *Peco) SingleKeyJumpIndex(ch rune) (uint, bool) {
	// FIXME: use p.keyjump or something instead of p.config
	n, ok := p.config.SingleKeyJump.PrefixMap[ch]
	return n, ok
}

func (p *Peco) Source() pipeline.Source {
	return p.source
}

func (p *Peco) Filters() *FilterSet {
	return &p.filters
}

func (p *Peco) Query() *Query {
	return &p.query
}

func (p *Peco) QueryExecDelay() time.Duration {
	return p.queryExecDelay
}

func (p *Peco) Caret() *Caret {
	return &p.caret
}

func (p *Peco) Hub() *hub.Hub {
	return p.hub
}

func (p *Peco) Err() error {
	return p.err
}

func (p *Peco) Exit(err error) {
	p.err = err
	if cf := p.cancelFunc; cf != nil {
		cf()
	}
}

func (p *Peco) Keymap() Keymap {
	return p.keymap
}

func (p *Peco) Setup() (err error) {
	if pdebug.Enabled {
		g := pdebug.Marker("Peco.Setup").BindError(&err)
		defer g.End()
	}

	if err := p.config.Init(); err != nil {
		return errors.Wrap(err, "failed to initialize config")
	}

	if err := parseCommandLine(&p.Options, &p.args, p.Argv); err != nil {
		return errors.Wrap(err, "failed to parse command line")
	}

	// Read config
	if err := readConfig(&p.config, p.Options.OptRcfile); err != nil {
		return errors.Wrap(err, "failed to setup configuration")
	}

	// Take Args, Config, Options, and apply the configuration to
	// the peco object
	if err := p.ApplyConfig(); err != nil {
		return errors.Wrap(err, "failed to apply configuration")
	}

	// XXX p.Keymap et al should be initialized around here
	p.hub = hub.New(5)

	// Setup source buffer
	src, err := p.SetupSource()
	if err != nil {
		return errors.Wrap(err, "failed to setup input source")
	}
	p.source = src
	p.ResetCurrentLineBuffer()
	return nil
}

func (p *Peco) Run(ctx context.Context) (err error) {
	if pdebug.Enabled {
		g := pdebug.Marker("Peco.Run").BindError(&err)
		defer g.End()
	}

	if err := p.Setup(); err != nil {
		return errors.Wrap(err, "failed to setup peco")
	}
	// screen.Init must be called within Run() because we
	// want to make sure to call screen.Close() after getting
	// out of Run()
	p.screen.Init()
	defer p.screen.Close()

	var _cancel func()
	ctx, _cancel = context.WithCancel(ctx)
	cancel := func() {
		if pdebug.Enabled {
			pdebug.Printf("Peco.Run cancel called")
		}
		_cancel()
	}

	// keep *this* ctx (not the Background one), as calling `cancel`
	// only affects the wrapped context
	p.ctx = ctx

	// remember this cancel func so p.Exit works (XXX requires locking?)
	p.cancelFunc = cancel

	loopers := []interface {
		Loop(ctx context.Context, cancel func()) error
	}{
		NewInput(p, p.Keymap(), p.screen.PollEvent()),
		NewView(p),
		NewFilter(p),
		sig.New(sig.SigReceivedHandlerFunc(func(sig os.Signal) {
			p.Exit(errors.New("received signal: " + sig.String()))
		})),
	}

	for _, l := range loopers {
		go l.Loop(ctx, cancel)
	}

	if pdebug.Enabled {
		pdebug.Printf("peco is now ready, go go go!")
	}

	close(p.readyCh)
	<-ctx.Done()

	return p.Err()
}

func parseCommandLine(opts *CLIOptions, args *[]string, argv []string) error {
	remaining, err := opts.parse(argv)
	if err != nil {
		return errors.Wrap(err, "failed to parse command line options")
	}

	if opts.OptHelp {
		os.Stdout.Write(opts.help())
		return makeIgnorable(errors.New("user asked to show help message"))
	}

	if opts.OptRcfile == "" {
		if file, err := LocateRcfile(locateRcfileIn); err != nil {
			opts.OptRcfile = file
		}
	}

	*args = remaining

	return nil
}

func (p *Peco) SetupSource() (s *Source, err error) {
	if pdebug.Enabled {
		g := pdebug.Marker("Peco.SetupSource").BindError(&err)
		defer g.End()
	}

	var in *os.File
	switch {
	case len(p.args) > 1:
		in, err = os.Open(p.args[1])
		if err != nil {
			return nil, errors.Wrap(err, "failed to open file for input")
		}
	case !util.IsTty(os.Stdin.Fd()):
		in = os.Stdin
	default:
		return nil, errors.New("you must supply something to work with via filename or stdin")
	}
	defer in.Close()

	src := NewSource(in, p.enableSep)

	// Block until we receive something from `in`
	if pdebug.Enabled {
		pdebug.Printf("Blocking until we read something in source...")
	}
	go src.Setup(p)
	<-src.Ready()

	return src, nil
}

func readConfig(cfg *Config, filename string) error {
	if filename != "" {
		if err := cfg.ReadFilename(filename); err != nil {
			return errors.Wrap(err, "failed to read config file")
		}
	}

	return nil
}

func (p *Peco) populateCommandList() error {
	for _, v := range p.config.Command {
		if len(v.Args) == 0 {
			continue
		}
		makeCommandAction(&v).Register("ExecuteCommand." + v.Name)
	}

	return nil
}

func (p *Peco) ApplyConfig() error {
	// If layoutType is not set and is set in the config, set it
	if p.layoutType == "" {
		if v := p.config.Layout; v != "" {
			p.layoutType = v
		} else {
			p.layoutType = DefaultLayoutType
		}
	}

	if err := p.populateCommandList(); err != nil {
		return errors.Wrap(err, "failed to populate command list")
	}

	if err := p.populateFilters(); err != nil {
		return errors.Wrap(err, "failed to populate filters")
	}

	if err := p.populateKeymap(); err != nil {
		return errors.Wrap(err, "failed to populate keymap")
	}

	if err := p.populateStyles(); err != nil {
		return errors.Wrap(err, "failed to populate styles")
	}

	return nil
}

func (p *Peco) populateFilters() error {
	p.filters.Add(NewIgnoreCaseFilter())
	p.filters.Add(NewCaseSensitiveFilter())
	p.filters.Add(NewSmartCaseFilter())
	p.filters.Add(NewRegexpFilter())
	return nil
}

func (p *Peco) populateKeymap() error {
	// Create a new keymap object
	k := NewKeymap(p.config.Keymap, p.config.Action)
	if err := k.ApplyKeybinding(); err != nil {
		return errors.Wrap(err, "failed to apply key bindings")
	}
	p.keymap = k
	return nil
}

func (p *Peco) populateStyles() error {
	p.styles = *(p.config.Style)
	return nil
}

func (p *Peco) CurrentLineBuffer() Buffer {
	return p.currentLineBuffer
}

func (p *Peco) SetCurrentLineBuffer(b Buffer) {
	if pdebug.Enabled {
		g := pdebug.Marker("Peco.SetCurrentLineBuffer %s", reflect.TypeOf(b).String())
		defer g.End()
	}
	p.currentLineBuffer = b
	p.Hub().SendDraw(false)
}

func (p *Peco) ResetCurrentLineBuffer() {
	p.SetCurrentLineBuffer(p.source)
}

func (p *Peco) ExecQuery() bool {
	if pdebug.Enabled {
		g := pdebug.Marker("Peco.ExecQuery")
		defer g.End()
	}

	select {
	case <-p.Ready():
	default:
		if pdebug.Enabled {
			pdebug.Printf("peco is not ready yet, ignoring.")
		}
		return false
	}

	// If this is an empty query, reset the display to show
	// the raw source buffer
	q := p.Query()
	if q.Len() <= 0 {
		if pdebug.Enabled {
			pdebug.Printf("empty query, reset buffer")
		}
		p.ResetCurrentLineBuffer()
		return true
	}

	delay := p.QueryExecDelay()
	if delay <= 0 {
		// No delay, execute immediately
		p.Hub().SendQuery(q.String())
		return true
	}

	p.queryExecMutex.Lock()
	defer p.queryExecMutex.Unlock()

	if p.queryExecTimer != nil {
		return true
	}

	// Wait $delay millisecs before sending the query
	// if a new input comes in, batch them up
	p.queryExecTimer = time.AfterFunc(delay, func() {
		p.Hub().SendQuery(q.String())

		p.queryExecMutex.Lock()
		defer p.queryExecMutex.Unlock()

		p.queryExecTimer = nil
	})
	return true
}

func (p *Peco) collectResults() {
	// In rare cases where the result channel is not setup
	// prior to call to this method, bail out
	if p.resultCh == nil {
		return
	}

	p.selection.Ascend(func(it btree.Item) bool {
		p.resultCh <- it.(Line)
		return true
	})
	close(p.resultCh)
}
<commit_msg>fix --version<commit_after>package peco

import (
	"os"
	"reflect"
	"time"

	"golang.org/x/net/context"

	"github.com/google/btree"
	"github.com/lestrrat/go-pdebug"
	"github.com/peco/peco/hub"
	"github.com/peco/peco/internal/util"
	"github.com/peco/peco/pipeline"
	"github.com/peco/peco/sig"
	"github.com/pkg/errors"
)

const version = "v0.3.6"

type errIgnorable struct {
	err error
}

func (e errIgnorable) Ignorable() bool { return true }
func (e errIgnorable) Causer() error {
	return e.err
}
func (e errIgnorable) Error() string {
	return e.err.Error()
}
func makeIgnorable(err error) error {
	return errIgnorable{err: err}
}

// Inputseq is a list of keys that the user typed
type Inputseq []string

func (is *Inputseq) Add(s string) {
	*is = append(*is, s)
}

func (is Inputseq) KeyNames() []string {
	return is
}

func (is Inputseq) Len() int {
	return len(is)
}

func (is *Inputseq) Reset() {
	*is = []string(nil)
}

func New() *Peco {
	return &Peco{
		Argv:              os.Args,
		currentLineBuffer: NewMemoryBuffer(), // XXX revisit this
		queryExecDelay:    50 * time.Millisecond,
		readyCh:           make(chan struct{}),
		screen:            &Termbox{},
		selection:         NewSelection(),
	}
}

func (p *Peco) Ready() <-chan struct{} {
	return p.readyCh
}

func (p *Peco) Screen() Screen {
	return p.screen
}

func (p *Peco) Styles() *StyleSet {
	return &p.styles
}

func (p *Peco) Prompt() string {
	return p.prompt
}

func (p *Peco) Inputseq() *Inputseq {
	return &p.inputseq
}

func (p *Peco) Context() context.Context {
	return p.ctx
}

func (p *Peco) LayoutType() string {
	return p.layoutType
}

func (p *Peco) Location() *Location {
	return &p.location
}

func (p *Peco) ResultCh() chan Line {
	return p.resultCh
}

func (p *Peco) SetResultCh(ch chan Line) {
	p.resultCh = ch
}

func (p *Peco) Selection() *Selection {
	return p.selection
}

func (s RangeStart) Valid() bool {
	return s.valid
}

func (s RangeStart) Value() int {
	return s.val
}

func (s *RangeStart) SetValue(n int) {
	s.val = n
	s.valid = true
}

func (s *RangeStart) Reset() {
	s.valid = false
}

func (p *Peco) SelectionRangeStart() *RangeStart {
	return &p.selectionRangeStart
}

func (p *Peco) SingleKeyJumpShowPrefix() bool {
	return p.singleKeyJumpShowPrefix
}

func (p *Peco) SingleKeyJumpPrefixes() []rune {
	return p.singleKeyJumpPrefixes
}

func (p *Peco) SingleKeyJumpMode() bool {
	return p.singleKeyJumpMode
}

func (p *Peco) SetSingleKeyJumpMode(b bool) {
	p.singleKeyJumpMode = b
}

func (p *Peco) ToggleSingleKeyJumpMode() {
	p.singleKeyJumpMode = !p.singleKeyJumpMode
}

func (p *Peco) SingleKeyJumpIndex(ch rune) (uint, bool) {
	// FIXME: use p.keyjump or something instead of p.config
	n, ok := p.config.SingleKeyJump.PrefixMap[ch]
	return n, ok
}

func (p *Peco) Source() pipeline.Source {
	return p.source
}

func (p *Peco) Filters() *FilterSet {
	return &p.filters
}

func (p *Peco) Query() *Query {
	return &p.query
}

func (p *Peco) QueryExecDelay() time.Duration {
	return p.queryExecDelay
}

func (p *Peco) Caret() *Caret {
	return &p.caret
}

func (p *Peco) Hub() *hub.Hub {
	return p.hub
}

func (p *Peco) Err() error {
	return p.err
}

func (p *Peco) Exit(err error) {
	p.err = err
	if cf := p.cancelFunc; cf != nil {
		cf()
	}
}

func (p *Peco) Keymap() Keymap {
	return p.keymap
}

func (p *Peco) Setup() (err error) {
	if pdebug.Enabled {
		g := pdebug.Marker("Peco.Setup").BindError(&err)
		defer g.End()
	}

	if err := p.config.Init(); err != nil {
		return errors.Wrap(err, "failed to initialize config")
	}

	if err := parseCommandLine(&p.Options, &p.args, p.Argv); err != nil {
		return errors.Wrap(err, "failed to parse command line")
	}

	// Read config
	if err := readConfig(&p.config, p.Options.OptRcfile); err != nil {
		return errors.Wrap(err, "failed to setup configuration")
	}

	// Take Args, Config, Options, and apply the configuration to
	// the peco object
	if err := p.ApplyConfig(); err != nil {
		return errors.Wrap(err, "failed to apply configuration")
	}

	// XXX p.Keymap et al should be initialized around here
	p.hub = hub.New(5)

	// Setup source buffer
	src, err := p.SetupSource()
	if err != nil {
		return errors.Wrap(err, "failed to setup input source")
	}
	p.source = src
	p.ResetCurrentLineBuffer()
	return nil
}

func (p *Peco) Run(ctx context.Context) (err error) {
	if pdebug.Enabled {
		g := pdebug.Marker("Peco.Run").BindError(&err)
		defer g.End()
	}

	if err := p.Setup(); err != nil {
		return errors.Wrap(err, "failed to setup peco")
	}
	// screen.Init must be called within Run() because we
	// want to make sure to call screen.Close() after getting
	// out of Run()
	p.screen.Init()
	defer p.screen.Close()

	var _cancel func()
	ctx, _cancel = context.WithCancel(ctx)
	cancel := func() {
		if pdebug.Enabled {
			pdebug.Printf("Peco.Run cancel called")
		}
		_cancel()
	}

	// keep *this* ctx (not the Background one), as calling `cancel`
	// only affects the wrapped context
	p.ctx = ctx

	// remember this cancel func so p.Exit works (XXX requires locking?)
	p.cancelFunc = cancel

	loopers := []interface {
		Loop(ctx context.Context, cancel func()) error
	}{
		NewInput(p, p.Keymap(), p.screen.PollEvent()),
		NewView(p),
		NewFilter(p),
		sig.New(sig.SigReceivedHandlerFunc(func(sig os.Signal) {
			p.Exit(errors.New("received signal: " + sig.String()))
		})),
	}

	for _, l := range loopers {
		go l.Loop(ctx, cancel)
	}

	if pdebug.Enabled {
		pdebug.Printf("peco is now ready, go go go!")
	}

	close(p.readyCh)
	<-ctx.Done()

	return p.Err()
}

func parseCommandLine(opts *CLIOptions, args *[]string, argv []string) error {
	remaining, err := opts.parse(argv)
	if err != nil {
		return errors.Wrap(err, "failed to parse command line options")
	}

	if opts.OptHelp {
		os.Stdout.Write(opts.help())
		return makeIgnorable(errors.New("user asked to show help message"))
	}

	if opts.OptVersion {
		os.Stdout.Write([]byte("peco version " + version))
		return makeIgnorable(errors.New("user asked to show version"))
	}

	if opts.OptRcfile == "" {
		if file, err := LocateRcfile(locateRcfileIn); err != nil {
			opts.OptRcfile = file
		}
	}

	*args = remaining

	return nil
}

func (p *Peco) SetupSource() (s *Source, err error) {
	if pdebug.Enabled {
		g := pdebug.Marker("Peco.SetupSource").BindError(&err)
		defer g.End()
	}

	var in *os.File
	switch {
	case len(p.args) > 1:
		in, err = os.Open(p.args[1])
		if err != nil {
			return nil, errors.Wrap(err, "failed to open file for input")
		}
	case !util.IsTty(os.Stdin.Fd()):
		in = os.Stdin
	default:
		return nil, errors.New("you must supply something to work with via filename or stdin")
	}
	defer in.Close()

	src := NewSource(in, p.enableSep)

	// Block until we receive something from `in`
	if pdebug.Enabled {
		pdebug.Printf("Blocking until we read something in source...")
	}
	go src.Setup(p)
	<-src.Ready()

	return src, nil
}

func readConfig(cfg *Config, filename string) error {
	if filename != "" {
		if err := cfg.ReadFilename(filename); err != nil {
			return errors.Wrap(err, "failed to read config file")
		}
	}

	return nil
}

func (p *Peco) populateCommandList() error {
	for _, v := range p.config.Command {
		if len(v.Args) == 0 {
			continue
		}
		makeCommandAction(&v).Register("ExecuteCommand." + v.Name)
	}

	return nil
}

func (p *Peco) ApplyConfig() error {
	// If layoutType is not set and is set in the config, set it
	if p.layoutType == "" {
		if v := p.config.Layout; v != "" {
			p.layoutType = v
		} else {
			p.layoutType = DefaultLayoutType
		}
	}

	if err := p.populateCommandList(); err != nil {
		return errors.Wrap(err, "failed to populate command list")
	}

	if err := p.populateFilters(); err != nil {
		return errors.Wrap(err, "failed to populate filters")
	}

	if err := p.populateKeymap(); err != nil {
		return errors.Wrap(err, "failed to populate keymap")
	}

	if err := p.populateStyles(); err != nil {
		return errors.Wrap(err, "failed to populate styles")
	}

	return nil
}

func (p *Peco) populateFilters() error {
	p.filters.Add(NewIgnoreCaseFilter())
	p.filters.Add(NewCaseSensitiveFilter())
	p.filters.Add(NewSmartCaseFilter())
	p.filters.Add(NewRegexpFilter())
	return nil
}

func (p *Peco) populateKeymap() error {
	// Create a new keymap object
	k := NewKeymap(p.config.Keymap, p.config.Action)
	if err := k.ApplyKeybinding(); err != nil {
		return errors.Wrap(err, "failed to apply key bindings")
	}
	p.keymap = k
	return nil
}

func (p *Peco) populateStyles() error {
	p.styles = *(p.config.Style)
	return nil
}

func (p *Peco) CurrentLineBuffer() Buffer {
	return p.currentLineBuffer
}

func (p *Peco) SetCurrentLineBuffer(b Buffer) {
	if pdebug.Enabled {
		g := pdebug.Marker("Peco.SetCurrentLineBuffer %s", reflect.TypeOf(b).String())
		defer g.End()
	}
	p.currentLineBuffer = b
	p.Hub().SendDraw(false)
}

func (p *Peco) ResetCurrentLineBuffer() {
	p.SetCurrentLineBuffer(p.source)
}

func (p *Peco) ExecQuery() bool {
	if pdebug.Enabled {
		g := pdebug.Marker("Peco.ExecQuery")
		defer g.End()
	}

	select {
	case <-p.Ready():
	default:
		if pdebug.Enabled {
			pdebug.Printf("peco is not ready yet, ignoring.")
		}
		return false
	}

	// If this is an empty query, reset the display to show
	// the raw source buffer
	q := p.Query()
	if q.Len() <= 0 {
		if pdebug.Enabled {
			pdebug.Printf("empty query, reset buffer")
		}
		p.ResetCurrentLineBuffer()
		return true
	}

	delay := p.QueryExecDelay()
	if delay <= 0 {
		// No delay, execute immediately
		p.Hub().SendQuery(q.String())
		return true
	}

	p.queryExecMutex.Lock()
	defer p.queryExecMutex.Unlock()

	if p.queryExecTimer != nil {
		return true
	}

	// Wait $delay millisecs before sending the query
	// if a new input comes in, batch them up
	p.queryExecTimer = time.AfterFunc(delay, func() {
		p.Hub().SendQuery(q.String())

		p.queryExecMutex.Lock()
		defer p.queryExecMutex.Unlock()

		p.queryExecTimer = nil
	})
	return true
}

func (p *Peco) collectResults() {
	// In rare cases where the result channel is not setup
	// prior to call to this method, bail out
	if p.resultCh == nil {
		return
	}

	p.selection.Ascend(func(it btree.Item) bool {
		p.resultCh <- it.(Line)
		return true
	})
	close(p.resultCh)
}
<|endoftext|>


<commit_before>// Copyright 2013 Jari Takkala and Brian Dignan. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"launchpad.net/tomb"
	"log"
	"net"
	"reflect"
	"sort"
	//"strconv"
	"syscall"
	"time"
)

var Protocol = [19]byte{'B', 'i', 't', 'T', 'o', 'r', 'r', 'e', 'n', 't', ' ', 'p', 'r', 'o', 't', 'o', 'c', 'o', 'l'}

// Message ID values
const (
	MsgChoke int = iota
	MsgUnchoke
	MsgInterested
	MsgNotInterested
	MsgHave
	MsgBitfield
	MsgRequest
	MsgPiece
	MsgCancel
	MsgPort
)

// PeerTuple represents a single IP+port pair of a peer
type PeerTuple struct {
	IP   net.IP
	Port uint16
}

type Peer struct {
	conn           *net.TCPConn
	peerName       string
	amChoking      bool
	amInterested   bool
	peerChoking    bool
	peerInterested bool
	ourBitfield    []bool
	peerBitfield   []bool
	initiator      bool
	peerID         []byte
	keepalive      <-chan time.Time // channel for sending keepalives
	lastTxKeepalive  time.Time
	lastRxKeepalive  time.Time
	read           chan []byte
	infoHash       []byte
	diskIOChans    diskIOPeerChans
	peerManagerChans peerManagerChans
	contRxChans    ControllerPeerChans
	contTxChans    PeerControllerChans
	stats          PeerStats
	t              tomb.Tomb
}

type PeerStats struct {
	read int
	write int
	errors int
}

type PeerManager struct {
	peers        map[string]*Peer
	infoHash     []byte
	peerChans    peerManagerChans
	serverChans  serverPeerChans
	trackerChans trackerPeerChans
	diskIOChans  diskIOPeerChans
	contChans    ControllerPeerManagerChans
	peerContChans PeerControllerChans
	t            tomb.Tomb
}

type peerManagerChans struct {
	deadPeer chan string
}

type PeerComms struct {
	peerName     string
	chans 		 ControllerPeerChans
}

func NewPeerComms(peerName string, cpc ControllerPeerChans) *PeerComms {
	pc := new(PeerComms)
	pc.peerName = peerName
	pc.chans = cpc
	return pc
}

type PeerInfo struct {
	peerName        string
	isChoked        bool // The peer is connected but choked. Defaults to TRUE (choked)
	availablePieces []bool
	activeRequests  map[int]struct{}
	qtyPiecesNeeded int                   // The quantity of pieces that this peer has that we haven't yet downloaded.
	chans 			ControllerPeerChans
}

type Handshake struct {
	Len      uint8
	Protocol [19]byte
	Reserved [8]uint8
	InfoHash [20]byte
	PeerID   [20]byte
}

func NewPeerInfo(quantityOfPieces int, peerComms PeerComms) *PeerInfo {
	pi := new(PeerInfo)

	pi.peerName = peerComms.peerName
	pi.chans = peerComms.chans

	pi.isChoked = true // By default, a peer starts as being choked by the other side.
	pi.availablePieces = make([]bool, quantityOfPieces)
	pi.activeRequests = make(map[int]struct{})

	return pi
}

// Sent by the peer to controller indicating a 'choke' state change. It either went from unchoked to choked,
// or from choked to unchoked.
type PeerChokeStatus struct {
	peerName   string
	isChoked bool
}

type SortedPeers []*PeerInfo

func (sp SortedPeers) Less(i, j int) bool {
	return sp[i].qtyPiecesNeeded <= sp[j].qtyPiecesNeeded
}

func (sp SortedPeers) Swap(i, j int) {
	tmp := sp[i]
	sp[i] = sp[j]
	sp[j] = tmp
}

func (sp SortedPeers) Len() int {
	return len(sp)
}

func sortedPeersByQtyPiecesNeeded(peers map[string]*PeerInfo) SortedPeers {
	peerInfoSlice := make(SortedPeers, 0)

	for _, peerInfo := range peers {
		peerInfoSlice = append(peerInfoSlice, peerInfo)
	}
	sort.Sort(peerInfoSlice)

	return peerInfoSlice
}

func NewPeerManager(infoHash []byte, diskIOChans diskIOPeerChans, serverChans serverPeerChans, trackerChans trackerPeerChans) *PeerManager {
	pm := new(PeerManager)
	pm.infoHash = infoHash
	pm.diskIOChans = diskIOChans
	pm.serverChans = serverChans
	pm.trackerChans = trackerChans
	pm.peerChans.deadPeer = make(chan string)
	pm.peers = make(map[string]*Peer)
	pm.contChans.newPeer = make(chan PeerComms)
	pm.contChans.deadPeer = make(chan string)
	pm.peerContChans.chokeStatus = make(chan PeerChokeStatus)
	pm.peerContChans.havePiece = make(chan chan HavePiece)
	return pm
}

func ConnectToPeer(peerTuple PeerTuple, connCh chan *net.TCPConn) {
	raddr := net.TCPAddr{peerTuple.IP, int(peerTuple.Port), ""}
	log.Println("Connecting to", raddr)
	conn, err := net.DialTCP("tcp4", nil, &raddr)
	if err != nil {
		if e, ok := err.(*net.OpError); ok {
			if e.Err == syscall.ECONNREFUSED {
				log.Println("ConnectToPeer : Connection Refused:", raddr)
				return
			}
		}
		log.Fatal(err)
	}
	log.Println("ConnectToPeer : Connected:", raddr)
	connCh <- conn
}

func NewPeer(
			peerName string,
			infoHash []byte,
			initiator bool,
			diskIOChans diskIOPeerChans,
			contRxChans ControllerPeerChans,
			contTxChans PeerControllerChans) *Peer {
	p := &Peer{
			peerName: peerName,
			infoHash: infoHash,
			amChoking: true,
			amInterested: false,
			peerChoking: true,
			peerInterested: false,
			initiator: initiator,
			diskIOChans: diskIOChans,
			contRxChans: contRxChans,
			contTxChans: contTxChans}
	p.read = make(chan []byte)
	return p
}

func constructMessage(id int, payload []byte) (msg []byte, err error) {
	msg = make([]byte, 4)

	// Store the length of payload + id in network byte order
	binary.BigEndian.PutUint32(msg, uint32(len(payload) + 1))
	msg = append(msg, byte(id))
	msg = append(msg, payload...)

	return
}

func verifyHandshake(handshake *Handshake, infoHash []byte) error {
	if int(handshake.Len) != len(Protocol) {
		err := fmt.Sprintf("Unexpected length for pstrlen (wanted %d, got %d)", len(Protocol), int(handshake.Len))
		return errors.New(err)
	}
	if !bytes.Equal(handshake.Protocol[:], Protocol[:]) {
		err := fmt.Sprintf("Protocol mismtach: got %s, expected %s", handshake.Protocol, Protocol)
		return errors.New(err)
	}
	if !bytes.Equal(handshake.InfoHash[:], infoHash) {
		err := fmt.Sprintf("Invalid infoHash: got %x, expected %x", handshake.InfoHash, infoHash)
		return errors.New(err)
	}
	return nil
}

func (p *Peer) sendHaveToController(pieces []HavePiece) {

}

func (p *Peer) readBytesFromConn(numBytes int) []byte {
	result := make([]byte, numBytes)
	_, err := io.ReadFull(p.conn, result)
	if err != nil {
		log.Fatalf("Encountered an error when attempting to read %d bytes from %s", numBytes, p.peerName)
	}
	return result
}

func (p *Peer) decodeMessage(payload []byte) {
	if len(payload) == 0 {
		// keepalive
		return
	}

	messageID := binary.BigEndian.Uint32(p.readBytesFromConn(1))

	switch messageID {
	case 0:
		// Choke Message
		break
	case 1:
		// Unchoke Message
		break
	case 2:
		// Interested Message
		break
	case 3:
		// Not Interested Message
		break
	case 4:
		// Have Message

		// Determine the piece number
		pieceNum := int(binary.BigEndian.Uint32(p.readBytesFromConn(4)))

		// Update the local peer bitfield
		p.peerBitfield[pieceNum] = true

		// Send a single HavePiece struct to the controller
		have := make([]HavePiece, 1)
		have[1] = HavePiece{pieceNum: pieceNum, peerName: p.peerName}
		break
	case 5:
		// Bitfield Message

		// Record the peer bitfield locally

		// Break the bitfield into a slice of HavePiece structs and send them
		// to the controller
		break
	case 6:
		// Request Message
		break
	case 7:
		// Piece Message
		break
	case 8:
		// Cancel Message
		break
	case 9:
		// Port Message
		break
	}
}

func (p *Peer) Reader() {
	log.Println("Peer : Reader : Started")

	var handshake Handshake
	binary.Read(p.conn, binary.BigEndian, &handshake)
	err := verifyHandshake(&handshake, p.infoHash)
	if err != nil {
		p.conn.Close()
		return
	}
	p.peerID = handshake.PeerID[:]

	for {
		length := make([]byte, 4)
		n, err := io.ReadFull(p.conn, length)
		if err != nil {
			// FIXME Are there any cases where we would not read length?
			return
		}
		payload := make([]byte, binary.BigEndian.Uint32(length))
		n, err = io.ReadFull(p.conn, payload)
		if err != nil {
			// FIXME if this is not a keepalive, we should
			// definitely get a payload
			return
		}
		log.Printf("Read %d bytes of %x\n", n, payload)
		p.decodeMessage(payload)
		//p.read <- buf
	}
}

func (p *Peer) sendHandshake() {
	log.Println("Peer : sendHandshake : Started")
	defer log.Println("Peer : sendHandshake : Completed")

	handshake := Handshake{
		Len: uint8(len(Protocol)),
		Protocol: Protocol,
		PeerID: PeerID,
	}
	copy(handshake.InfoHash[:], p.infoHash)
	err := binary.Write(p.conn, binary.BigEndian, &handshake)
	if err != nil {
		// TODO: Handle errors
		log.Fatal(err)
	}
	p.stats.write += int(reflect.TypeOf(handshake).Size())
}

func (p *Peer) Stop() error {
	log.Println("Peer : Stop : Stopping")
	p.t.Kill(nil)
	return p.t.Wait()
}

func (p *Peer) Run() {
	log.Println("Peer : Run : Started")
	defer log.Println("Peer : Run : Completed")

	go p.sendHandshake()
	go p.Reader()

	for {
		select {
		case <-p.keepalive:
		case <-p.read:
			fmt.Println("p.read")
		//case buf := <-p.read:
			//fmt.Println("Read from peer:", buf)
		case <-p.t.Dying():
			p.peerManagerChans.deadPeer <- p.conn.RemoteAddr().String()
			return
		}
	}
}

func (pm *PeerManager) Stop() error {
	log.Println("PeerManager : Stop : Stopping")
	pm.t.Kill(nil)
	return pm.t.Wait()
}

func (pm *PeerManager) Run() {
	log.Println("PeerManager : Run : Started")
	defer pm.t.Done()
	defer log.Println("PeerManager : Run : Completed")

	for {
		select {
		case peer := <-pm.trackerChans.peers:
			peerName := fmt.Sprintf("%s:%d", peer.IP.String(), peer.Port)
			_, ok := pm.peers[peerName]
			if !ok {
				// FIXME Code in this block is duplicated below

				// Create the Controller->Peer chans struct
				contTxChans := *NewControllerPeerChans()

				// Construct the Peer object
				pm.peers[peerName] = NewPeer(
					peerName,
					pm.infoHash,
					true,
					pm.diskIOChans,
					contTxChans,
					pm.peerContChans)

				// Give the controller the channels that it will use to
				// transmit messages to this new peer
				go func() {
					pm.contChans.newPeer <- PeerComms{peerName: peerName, chans: contTxChans}
				}()

				// Have the 'peer' routine create an outbound
				// TCP connection to the remote peer
				go ConnectToPeer(peer, pm.serverChans.conns)
			}
		case conn := <-pm.serverChans.conns:
			_, ok := pm.peers[conn.RemoteAddr().String()]
			if !ok {
				// Create the Controller->Peer chans struct
				contTxChans := *NewControllerPeerChans()

				// Construct the Peer object
				peerName := conn.RemoteAddr().String()
				pm.peers[peerName] = NewPeer(
					peerName,
					pm.infoHash,
					false,
					pm.diskIOChans,
					contTxChans,
					pm.peerContChans)

				// Give the controller the channels that it will use to
				// transmit messages to this new peer
				go func() {
					pm.contChans.newPeer <- PeerComms{peerName: peerName, chans: contTxChans}
				}()
			}
			// Associate the connection with the peer object and start the peer
			pm.peers[conn.RemoteAddr().String()].conn = conn
			go pm.peers[conn.RemoteAddr().String()].Run()
		case peer := <-pm.peerChans.deadPeer:
			log.Printf("PeerManager : Deleting peer %s\n", peer)
			delete(pm.peers, peer)
		case <-pm.t.Dying():
			for _, peer := range pm.peers {
				peer.Stop()
			}
			return
		}
	}
}
<commit_msg>Finished writing the HavePiece case in peer.decodeMessage()<commit_after>// Copyright 2013 Jari Takkala and Brian Dignan. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"launchpad.net/tomb"
	"log"
	"net"
	"reflect"
	"sort"
	//"strconv"
	"syscall"
	"time"
)

var Protocol = [19]byte{'B', 'i', 't', 'T', 'o', 'r', 'r', 'e', 'n', 't', ' ', 'p', 'r', 'o', 't', 'o', 'c', 'o', 'l'}

// Message ID values
const (
	MsgChoke int = iota
	MsgUnchoke
	MsgInterested
	MsgNotInterested
	MsgHave
	MsgBitfield
	MsgRequest
	MsgPiece
	MsgCancel
	MsgPort
)

// PeerTuple represents a single IP+port pair of a peer
type PeerTuple struct {
	IP   net.IP
	Port uint16
}

type Peer struct {
	conn           *net.TCPConn
	peerName       string
	amChoking      bool
	amInterested   bool
	peerChoking    bool
	peerInterested bool
	ourBitfield    []bool
	peerBitfield   []bool
	initiator      bool
	peerID         []byte
	keepalive      <-chan time.Time // channel for sending keepalives
	lastTxKeepalive  time.Time
	lastRxKeepalive  time.Time
	read           chan []byte
	infoHash       []byte
	diskIOChans    diskIOPeerChans
	peerManagerChans peerManagerChans
	contRxChans    ControllerPeerChans
	contTxChans    PeerControllerChans
	stats          PeerStats
	t              tomb.Tomb
}

type PeerStats struct {
	read int
	write int
	errors int
}

type PeerManager struct {
	peers        map[string]*Peer
	infoHash     []byte
	peerChans    peerManagerChans
	serverChans  serverPeerChans
	trackerChans trackerPeerChans
	diskIOChans  diskIOPeerChans
	contChans    ControllerPeerManagerChans
	peerContChans PeerControllerChans
	t            tomb.Tomb
}

type peerManagerChans struct {
	deadPeer chan string
}

type PeerComms struct {
	peerName     string
	chans 		 ControllerPeerChans
}

func NewPeerComms(peerName string, cpc ControllerPeerChans) *PeerComms {
	pc := new(PeerComms)
	pc.peerName = peerName
	pc.chans = cpc
	return pc
}

type PeerInfo struct {
	peerName        string
	isChoked        bool // The peer is connected but choked. Defaults to TRUE (choked)
	availablePieces []bool
	activeRequests  map[int]struct{}
	qtyPiecesNeeded int                   // The quantity of pieces that this peer has that we haven't yet downloaded.
	chans 			ControllerPeerChans
}

type Handshake struct {
	Len      uint8
	Protocol [19]byte
	Reserved [8]uint8
	InfoHash [20]byte
	PeerID   [20]byte
}

func NewPeerInfo(quantityOfPieces int, peerComms PeerComms) *PeerInfo {
	pi := new(PeerInfo)

	pi.peerName = peerComms.peerName
	pi.chans = peerComms.chans

	pi.isChoked = true // By default, a peer starts as being choked by the other side.
	pi.availablePieces = make([]bool, quantityOfPieces)
	pi.activeRequests = make(map[int]struct{})

	return pi
}

// Sent by the peer to controller indicating a 'choke' state change. It either went from unchoked to choked,
// or from choked to unchoked.
type PeerChokeStatus struct {
	peerName   string
	isChoked bool
}

type SortedPeers []*PeerInfo

func (sp SortedPeers) Less(i, j int) bool {
	return sp[i].qtyPiecesNeeded <= sp[j].qtyPiecesNeeded
}

func (sp SortedPeers) Swap(i, j int) {
	tmp := sp[i]
	sp[i] = sp[j]
	sp[j] = tmp
}

func (sp SortedPeers) Len() int {
	return len(sp)
}

func sortedPeersByQtyPiecesNeeded(peers map[string]*PeerInfo) SortedPeers {
	peerInfoSlice := make(SortedPeers, 0)

	for _, peerInfo := range peers {
		peerInfoSlice = append(peerInfoSlice, peerInfo)
	}
	sort.Sort(peerInfoSlice)

	return peerInfoSlice
}

func NewPeerManager(infoHash []byte, diskIOChans diskIOPeerChans, serverChans serverPeerChans, trackerChans trackerPeerChans) *PeerManager {
	pm := new(PeerManager)
	pm.infoHash = infoHash
	pm.diskIOChans = diskIOChans
	pm.serverChans = serverChans
	pm.trackerChans = trackerChans
	pm.peerChans.deadPeer = make(chan string)
	pm.peers = make(map[string]*Peer)
	pm.contChans.newPeer = make(chan PeerComms)
	pm.contChans.deadPeer = make(chan string)
	pm.peerContChans.chokeStatus = make(chan PeerChokeStatus)
	pm.peerContChans.havePiece = make(chan chan HavePiece)
	return pm
}

func ConnectToPeer(peerTuple PeerTuple, connCh chan *net.TCPConn) {
	raddr := net.TCPAddr{peerTuple.IP, int(peerTuple.Port), ""}
	log.Println("Connecting to", raddr)
	conn, err := net.DialTCP("tcp4", nil, &raddr)
	if err != nil {
		if e, ok := err.(*net.OpError); ok {
			if e.Err == syscall.ECONNREFUSED {
				log.Println("ConnectToPeer : Connection Refused:", raddr)
				return
			}
		}
		log.Fatal(err)
	}
	log.Println("ConnectToPeer : Connected:", raddr)
	connCh <- conn
}

func NewPeer(
			peerName string,
			infoHash []byte,
			initiator bool,
			diskIOChans diskIOPeerChans,
			contRxChans ControllerPeerChans,
			contTxChans PeerControllerChans) *Peer {
	p := &Peer{
			peerName: peerName,
			infoHash: infoHash,
			amChoking: true,
			amInterested: false,
			peerChoking: true,
			peerInterested: false,
			initiator: initiator,
			diskIOChans: diskIOChans,
			contRxChans: contRxChans,
			contTxChans: contTxChans}
	p.read = make(chan []byte)
	return p
}

func constructMessage(id int, payload []byte) (msg []byte, err error) {
	msg = make([]byte, 4)

	// Store the length of payload + id in network byte order
	binary.BigEndian.PutUint32(msg, uint32(len(payload) + 1))
	msg = append(msg, byte(id))
	msg = append(msg, payload...)

	return
}

func verifyHandshake(handshake *Handshake, infoHash []byte) error {
	if int(handshake.Len) != len(Protocol) {
		err := fmt.Sprintf("Unexpected length for pstrlen (wanted %d, got %d)", len(Protocol), int(handshake.Len))
		return errors.New(err)
	}
	if !bytes.Equal(handshake.Protocol[:], Protocol[:]) {
		err := fmt.Sprintf("Protocol mismtach: got %s, expected %s", handshake.Protocol, Protocol)
		return errors.New(err)
	}
	if !bytes.Equal(handshake.InfoHash[:], infoHash) {
		err := fmt.Sprintf("Invalid infoHash: got %x, expected %x", handshake.InfoHash, infoHash)
		return errors.New(err)
	}
	return nil
}

// Send zero or more HavePiece messages to the controller. To send zero HavePiece
// messages (like when sending an empty bitfield), the pieces argument should
// just be an empty slice.
// NOTE: This function will potentially block and should be run
// as a separate goroutine.
func (p *Peer) sendHaveMessagesToController(pieces []HavePiece) {
	// make an inner channel that will be used to send the individual HavePiece
	// messages to the controller.
	innerChan := make(chan HavePiece)
	p.contTxChans.havePiece <- innerChan

	for _, havePiece := range pieces {
		innerChan <- havePiece
	}

	// close the inner channel to signal to the controller that we're finished
	// sending HavePiece messages.
	close(innerChan)
}

func (p *Peer) readBytesFromConn(numBytes int) []byte {
	result := make([]byte, numBytes)
	_, err := io.ReadFull(p.conn, result)
	if err != nil {
		log.Fatalf("Encountered an error when attempting to read %d bytes from %s", numBytes, p.peerName)
	}
	return result
}

func (p *Peer) decodeMessage(payload []byte) {
	if len(payload) == 0 {
		// keepalive
		return
	}

	messageID := binary.BigEndian.Uint32(p.readBytesFromConn(1))

	switch messageID {
	case 0:
		// Choke Message
		break
	case 1:
		// Unchoke Message
		break
	case 2:
		// Interested Message
		break
	case 3:
		// Not Interested Message
		break
	case 4:
		// Have Message

		// Determine the piece number
		pieceNum := int(binary.BigEndian.Uint32(p.readBytesFromConn(4)))

		// Update the local peer bitfield
		p.peerBitfield[pieceNum] = true

		// Send a single HavePiece struct to the controller
		have := make([]HavePiece, 1)
		have[1] = HavePiece{pieceNum: pieceNum, peerName: p.peerName}
		go p.sendHaveMessagesToController(have)
		break
	case 5:
		// Bitfield Message

		// Record the peer bitfield locally

		// Break the bitfield into a slice of HavePiece structs and send them
		// to the controller
		break
	case 6:
		// Request Message
		break
	case 7:
		// Piece Message
		break
	case 8:
		// Cancel Message
		break
	case 9:
		// Port Message
		break
	}
}

func (p *Peer) Reader() {
	log.Println("Peer : Reader : Started")

	var handshake Handshake
	binary.Read(p.conn, binary.BigEndian, &handshake)
	err := verifyHandshake(&handshake, p.infoHash)
	if err != nil {
		p.conn.Close()
		return
	}
	p.peerID = handshake.PeerID[:]

	for {
		length := make([]byte, 4)
		n, err := io.ReadFull(p.conn, length)
		if err != nil {
			// FIXME Are there any cases where we would not read length?
			return
		}
		payload := make([]byte, binary.BigEndian.Uint32(length))
		n, err = io.ReadFull(p.conn, payload)
		if err != nil {
			// FIXME if this is not a keepalive, we should
			// definitely get a payload
			return
		}
		log.Printf("Read %d bytes of %x\n", n, payload)
		p.decodeMessage(payload)
		//p.read <- buf
	}
}

func (p *Peer) sendHandshake() {
	log.Println("Peer : sendHandshake : Started")
	defer log.Println("Peer : sendHandshake : Completed")

	handshake := Handshake{
		Len: uint8(len(Protocol)),
		Protocol: Protocol,
		PeerID: PeerID,
	}
	copy(handshake.InfoHash[:], p.infoHash)
	err := binary.Write(p.conn, binary.BigEndian, &handshake)
	if err != nil {
		// TODO: Handle errors
		log.Fatal(err)
	}
	p.stats.write += int(reflect.TypeOf(handshake).Size())
}

func (p *Peer) Stop() error {
	log.Println("Peer : Stop : Stopping")
	p.t.Kill(nil)
	return p.t.Wait()
}

func (p *Peer) Run() {
	log.Println("Peer : Run : Started")
	defer log.Println("Peer : Run : Completed")

	go p.sendHandshake()
	go p.Reader()

	for {
		select {
		case <-p.keepalive:
		case <-p.read:
			fmt.Println("p.read")
		//case buf := <-p.read:
			//fmt.Println("Read from peer:", buf)
		case <-p.t.Dying():
			p.peerManagerChans.deadPeer <- p.conn.RemoteAddr().String()
			return
		}
	}
}

func (pm *PeerManager) Stop() error {
	log.Println("PeerManager : Stop : Stopping")
	pm.t.Kill(nil)
	return pm.t.Wait()
}

func (pm *PeerManager) Run() {
	log.Println("PeerManager : Run : Started")
	defer pm.t.Done()
	defer log.Println("PeerManager : Run : Completed")

	for {
		select {
		case peer := <-pm.trackerChans.peers:
			peerName := fmt.Sprintf("%s:%d", peer.IP.String(), peer.Port)
			_, ok := pm.peers[peerName]
			if !ok {
				// FIXME Code in this block is duplicated below

				// Create the Controller->Peer chans struct
				contTxChans := *NewControllerPeerChans()

				// Construct the Peer object
				pm.peers[peerName] = NewPeer(
					peerName,
					pm.infoHash,
					true,
					pm.diskIOChans,
					contTxChans,
					pm.peerContChans)

				// Give the controller the channels that it will use to
				// transmit messages to this new peer
				go func() {
					pm.contChans.newPeer <- PeerComms{peerName: peerName, chans: contTxChans}
				}()

				// Have the 'peer' routine create an outbound
				// TCP connection to the remote peer
				go ConnectToPeer(peer, pm.serverChans.conns)
			}
		case conn := <-pm.serverChans.conns:
			_, ok := pm.peers[conn.RemoteAddr().String()]
			if !ok {
				// Create the Controller->Peer chans struct
				contTxChans := *NewControllerPeerChans()

				// Construct the Peer object
				peerName := conn.RemoteAddr().String()
				pm.peers[peerName] = NewPeer(
					peerName,
					pm.infoHash,
					false,
					pm.diskIOChans,
					contTxChans,
					pm.peerContChans)

				// Give the controller the channels that it will use to
				// transmit messages to this new peer
				go func() {
					pm.contChans.newPeer <- PeerComms{peerName: peerName, chans: contTxChans}
				}()
			}
			// Associate the connection with the peer object and start the peer
			pm.peers[conn.RemoteAddr().String()].conn = conn
			go pm.peers[conn.RemoteAddr().String()].Run()
		case peer := <-pm.peerChans.deadPeer:
			log.Printf("PeerManager : Deleting peer %s\n", peer)
			delete(pm.peers, peer)
		case <-pm.t.Dying():
			for _, peer := range pm.peers {
				peer.Stop()
			}
			return
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/rcrowley/go-metrics"
)

var (
	output    = flag.String("o", "-", "output file")
	workers   = flag.Int("n", 4, "parallel connections")
	chunkSize = flag.Int("c", 100*1024, "chunk size")
	debug     = flag.Bool("d", false, "debug")
)

type chunk struct {
	offset int64
	size   int64
	file   string
}

type download struct {
	wg       sync.WaitGroup
	url      string
	size     int64
	startPos int64
	writer   io.Writer

	succeeded     counter
	downloaded    counter
	downloadMeter metrics.Meter
	progress      counter
}

type counter struct {
	v int64
}

func (c *counter) add(delta int64) {
	atomic.AddInt64(&c.v, delta)
}

func (c *counter) value() int64 {
	return atomic.LoadInt64(&c.v)
}

func newDownload(url string) (*download, error) {
	size, err := contentSize(url)
	if err != nil {
		return nil, err
	}

	return &download{
		url:           url,
		size:          size,
		writer:        os.Stdout,
		downloadMeter: metrics.NewMeter(),
	}, nil
}

// chunkFeeder generates chunk structures and sends them to the out channel.
// The chunk size is determined via commandline params. The last chunk is
// is sized accordingly.
func (d *download) chunkFeeder(out chan<- *chunk) {
	// when resuming a download. the chunk feeder won't emit chunks
	// that are already been written.
	d.succeeded.add(d.startPos)

	if *debug {
		log.Printf("Feeder is running")
	}
	for i := int64(0); ; i++ {
		size := int64(*chunkSize)
		off := i * size

		if off+size <= d.startPos {
			continue
		}
		if d.startPos-off > 0 && d.startPos-off < size {
			size -= d.startPos - off
			off = d.startPos
		}

		if off >= d.size {
			break
		}
		if off+size > d.size {
			size = off + size - d.size
		}
		if *debug {
			log.Printf("Enqueuing chunk %d size %d", off, size)
		}
		out <- &chunk{
			offset: off,
			size:   size}
	}
	if *debug {
		log.Printf("Feeder is closing")
	}
	close(out)
}

// chunkWorker task chunks from the in channel, fetches the chunk from
// the http resource into the buffer in the chunk structure.
// The completed chunk is then sent to the out channel.
// Any error causes to panic.
func (d *download) chunkWorker(n int, in <-chan *chunk, out chan<- *chunk) {
	if *debug {
		log.Println("Started chunk worker", n)
	}

	for ch := range in {
		if *debug {
			log.Printf("Chunk worker %d fetching chunk %d (%d)\n", n, ch.offset, ch.size)
		}

		basename := *output
		if *output == "-" {
			basename = "/tmp/pget.tmp"
		}

		offset := ch.offset
		size := ch.size

		var f *os.File
		var err error

		tmpName := fmt.Sprintf("%s.chunk-%d", basename, ch.offset)
		if st, err := os.Stat(tmpName); os.IsNotExist(err) {
			f, err = os.Create(tmpName)
			if err != nil {
				log.Fatalf("Cannot create temp file for chunk for offset %d, %s", ch.offset, err)
			}
		} else if err == nil {
			f, err = os.OpenFile(tmpName, os.O_APPEND|os.O_WRONLY, 0)
			if err != nil {
				log.Fatalf("Cannot open temp file for appending chunk for offset %d, %s", ch.offset, err)
			}
			offset += st.Size()
			size -= st.Size()
		} else {
			log.Fatal("Cannot stat ", tmpName, err)
		}

		var body io.ReadCloser
		if size < 1 {
			body = ioutil.NopCloser(&bytes.Buffer{})
		} else {
			if *debug {
				log.Printf("Worker %d fetching from %d -> size %d\n", n, offset, size)
			}

			body, err = d.fetch(offset, size)
			if err != nil {
				log.Fatal(err)
			}
			// mark the whole chunk as succeeded
			// even if we continued an interrupted download.
			d.succeeded.add(ch.size)
		}

		if _, err = io.Copy(f, body); err != nil {
			log.Fatal(err)
		}
		body.Close()
		f.Close()

		ch.file = tmpName
		out <- ch
	}

	if *debug {
		log.Println("Finishing chunk worker", n)
	}
}

// chunkWriter reads completed chunks from the in channel,
// and writes them out in the correct order.
// Out of order chunks are kept in memory until the missing pieces
// are delivered.
func (d *download) chunkWriter(in <-chan *chunk) {
	readyChunks := make(map[int64]*chunk)
	outPos := d.startPos
	d.progress.add(d.startPos)

	for ch := range in {
		if *debug {
			log.Printf("Chunk writer processing completed chunk %d (%d)\n", ch.offset, ch.size)
		}
		readyChunks[ch.offset] = ch

		for {
			if c, ok := readyChunks[outPos]; ok {
				f, err := os.Open(c.file)
				if err != nil {
					log.Fatal(err)
				}
				n, err := io.Copy(d.writer, f)
				if err != nil {
					log.Fatal(err)
				}
				d.progress.add(n)

				f.Close()
				//os.Remove(c.file)

				outPos += c.size
			} else {
				break
			}
		}
	}
}

// start
func start() error {
	url := flag.Arg(0)

	d, err := newDownload(url)
	if err != nil {
		return err
	}

	lastPos := int64(0)
	if *output != "-" {
		//f, err := os.Create(*output)
		f, err := os.OpenFile(*output, os.O_RDWR|os.O_CREATE, 0666)
		if err != nil {
			return err
		}
		defer f.Close()
		d.writer = f

		lastPos, err = f.Seek(0, 2)
		if err != nil {
			log.Fatal(err)
		}
	}
	d.startPos = lastPos

	if *debug {
		log.Println("Size", d.size)
		log.Println("Last pos", lastPos)
	}

	ich := make(chan *chunk, 2)
	och := make(chan *chunk, 2)

	// spawn the chunk feeder.
	go d.chunkFeeder(ich)

	// spawn the workers in such a way
	// that we know when they are all done.
	for i := 0; i < *workers; i++ {
		d.wg.Add(1)
		go func(i int) {
			d.chunkWorker(i, ich, och)
			d.wg.Done()
		}(i)
	}

	// when all the workers are gone
	// close the chunk writer input
	go func() {
		d.wg.Wait()
		close(och)
	}()

	done := make(chan struct{})
	go d.showProgress(done)

	// the chunk writer block until
	// all the chunks are received, reordered and written.
	d.chunkWriter(och)

	close(done)
	d.printProgress()

	return nil
}

func (d *download) showProgress(done <-chan struct{}) {
	d.printProgress()

	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			d.printProgress()
		case <-done:
			return
		}
	}
}

// percentRatio returns the ration of two integers as a percentage.
func percentRatio(a, b int64) float64 {
	return float64(a) / float64(b) * 100.0
}

func (d *download) printProgress() {
	fmt.Fprintf(os.Stderr, "succeeded: %.2f%%; progress: %.2f%%; downloaded: %s; rate: %s/s\n", percentRatio(d.succeeded.value(), d.size), percentRatio(d.progress.value(), d.size), humanize.Bytes(uint64(d.downloaded.value())), humanize.Bytes(uint64(d.downloadMeter.Rate1())))
}

// contentSize returns the size of the resource
// at the given url.
func contentSize(url string) (int64, error) {
	resp, err := http.Get(url)
	if err != nil {
		return 0, err
	}
	return resp.ContentLength, nil
}

// followRedirect ensures that all headers are passed through the redirect.
// This is necessary for the range header to work across redirects.
func followRedirect(req *http.Request, via []*http.Request) error {
	req.Header = via[0].Header
	return nil
}

// fetch returns a ReadCloser with the content of a
// section of the resource found at url starting at from
// and len long.
func (d *download) fetch(from, len int64) (io.ReadCloser, error) {
	req, err := http.NewRequest("GET", d.url, nil)
	if err != nil {
		return nil, err
	}

	req.Header.Add("Range", fmt.Sprintf("bytes=%d-%d", from, from+len-1))

	client := &http.Client{
		CheckRedirect: followRedirect,
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusPartialContent {
		defer resp.Body.Close()
		if resp.StatusCode == http.StatusOK {
			return nil, fmt.Errorf("Server doesn't implement partial content")
		}
		return nil, errors.New(resp.Status)
	}
	return &measuringReader{resp.Body, &d.downloaded, d.downloadMeter}, nil
}

type measuringReader struct {
	io.ReadCloser
	*counter
	metrics.Meter
}

func (m *measuringReader) Read(p []byte) (int, error) {
	n, err := m.ReadCloser.Read(p)
	m.add(int64(n))
	m.Mark(int64(n))
	return n, err
}

func main() {
	flag.Parse()
	if err := start(); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Fix last chunk. Fix if temp chunks are too big. Retry N times with exp backoff.<commit_after>package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sync"
	"sync/atomic"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/rcrowley/go-metrics"
)

var (
	output    = flag.String("o", "-", "output file")
	workers   = flag.Int("n", 4, "parallel connections")
	chunkSize = flag.Int("c", 100*1024, "chunk size")
	debug     = flag.Bool("d", false, "debug")
)

const (
	numRetries = 8
)

type chunk struct {
	offset int64
	size   int64
	file   string
}

type download struct {
	wg       sync.WaitGroup
	url      string
	size     int64
	startPos int64
	writer   io.Writer

	succeeded     counter
	downloaded    counter
	downloadMeter metrics.Meter
	progress      counter
}

type counter struct {
	v int64
}

func (c *counter) add(delta int64) {
	atomic.AddInt64(&c.v, delta)
}

func (c *counter) value() int64 {
	return atomic.LoadInt64(&c.v)
}

func newDownload(url string) (*download, error) {
	size, err := contentSize(url)
	if err != nil {
		return nil, err
	}

	return &download{
		url:           url,
		size:          size,
		writer:        os.Stdout,
		downloadMeter: metrics.NewMeter(),
	}, nil
}

// chunkFeeder generates chunk structures and sends them to the out channel.
// The chunk size is determined via commandline params. The last chunk is
// is sized accordingly.
func (d *download) chunkFeeder(out chan<- *chunk) {
	// when resuming a download. the chunk feeder won't emit chunks
	// that are already been written.
	d.succeeded.add(d.startPos)

	if *debug {
		log.Printf("Feeder is running")
	}
	for i := int64(0); ; i++ {
		size := int64(*chunkSize)
		off := i * size

		if off+size <= d.startPos {
			continue
		}
		if d.startPos-off > 0 && d.startPos-off < size {
			size -= d.startPos - off
			off = d.startPos
		}

		if off >= d.size {
			break
		}
		if off+size >= d.size {
			size = d.size - off
		}
		if *debug {
			log.Printf("Enqueuing chunk %d size %d", off, size)
		}
		out <- &chunk{
			offset: off,
			size:   size}
	}
	if *debug {
		log.Printf("Feeder is closing")
	}
	close(out)
}

// chunkWorker task chunks from the in channel, fetches the chunk from
// the http resource into the buffer in the chunk structure.
// The completed chunk is then sent to the out channel.
// Any error causes to panic.
func (d *download) chunkWorker(n int, in <-chan *chunk, out chan<- *chunk) {
	if *debug {
		log.Println("Started chunk worker", n)
	}

	for ch := range in {
		if *debug {
			log.Printf("Chunk worker %d fetching chunk %d (%d)\n", n, ch.offset, ch.size)
		}

		basename := *output
		if *output == "-" {
			basename = "/tmp/pget.tmp"
		}

		offset := ch.offset
		size := ch.size

		var f *os.File
		var err error

		tmpName := fmt.Sprintf("%s.chunk-%d", basename, ch.offset)
		if st, err := os.Stat(tmpName); os.IsNotExist(err) {
			f, err = os.Create(tmpName)
			if err != nil {
				log.Fatalf("Cannot create temp file for chunk for offset %d, %s", ch.offset, err)
			}
		} else if err == nil {
			f, err = os.OpenFile(tmpName, os.O_APPEND|os.O_WRONLY, 0)
			if err != nil {
				log.Fatalf("Cannot open temp file for appending chunk for offset %d, %s", ch.offset, err)
			}
			offset += st.Size()
			size -= st.Size()
		} else {
			log.Fatal("Cannot stat ", tmpName, err)
		}

		var body io.ReadCloser
		if size < 1 {
			if *debug {
				log.Printf("Chunk worker %d NOT fetching chunk %d (%d), because tmp data is complete.\n", n, ch.offset, ch.size)
			}

			body = ioutil.NopCloser(&bytes.Buffer{})
		} else {
			if *debug {
				log.Printf("Worker %d fetching chunk %d (%d) partial from %d -> size %d\n", n, ch.offset, ch.size, offset, size)
			}

			body, err = d.retryFetch(offset, size)
			if err != nil {
				log.Fatal(err)
			}
			// mark the whole chunk as succeeded
			// even if we continued an interrupted download.
			d.succeeded.add(ch.size)
		}

		if bs, err := io.Copy(f, body); err != nil {
			log.Fatalf("Error saving chunk %d (%d) after %d bytes: %s", ch.offset, ch.size, bs, err)
		}
		body.Close()
		f.Close()

		ch.file = tmpName
		out <- ch
	}

	if *debug {
		log.Println("Finishing chunk worker", n)
	}
}

// chunkWriter reads completed chunks from the in channel,
// and writes them out in the correct order.
// Out of order chunks are kept in memory until the missing pieces
// are delivered.
func (d *download) chunkWriter(in <-chan *chunk) {
	readyChunks := make(map[int64]*chunk)
	outPos := d.startPos
	d.progress.add(d.startPos)

	for ch := range in {
		if *debug {
			log.Printf("Chunk writer processing completed chunk %d (%d)\n", ch.offset, ch.size)
		}
		readyChunks[ch.offset] = ch

		for {
			if c, ok := readyChunks[outPos]; ok {
				f, err := os.Open(c.file)
				if err != nil {
					log.Fatal(err)
				}

				n, err := io.CopyN(d.writer, f, ch.size)
				if err != nil {
					log.Fatalf("Error merging chunk %d (%d) after %d bytes: %s", ch.offset, ch.size, n, err)
				}
				d.progress.add(n)

				f.Close()
				os.Remove(c.file)

				outPos += c.size
			} else {
				break
			}
		}
	}
}

// start
func start() error {
	url := flag.Arg(0)

	d, err := newDownload(url)
	if err != nil {
		return err
	}

	lastPos := int64(0)
	if *output != "-" {
		//f, err := os.Create(*output)
		f, err := os.OpenFile(*output, os.O_RDWR|os.O_CREATE, 0666)
		if err != nil {
			return err
		}
		defer f.Close()
		d.writer = f

		lastPos, err = f.Seek(0, 2)
		if err != nil {
			log.Fatal(err)
		}
	}
	d.startPos = lastPos

	if *debug {
		log.Println("Size", d.size)
		log.Println("Last pos", lastPos)
	}

	ich := make(chan *chunk, 2)
	och := make(chan *chunk, 2)

	// spawn the chunk feeder.
	go d.chunkFeeder(ich)

	// spawn the workers in such a way
	// that we know when they are all done.
	for i := 0; i < *workers; i++ {
		d.wg.Add(1)
		go func(i int) {
			d.chunkWorker(i, ich, och)
			d.wg.Done()
		}(i)
	}

	// when all the workers are gone
	// close the chunk writer input
	go func() {
		d.wg.Wait()
		close(och)
	}()

	done := make(chan struct{})
	go d.showProgress(done)

	// the chunk writer block until
	// all the chunks are received, reordered and written.
	d.chunkWriter(och)

	close(done)
	d.printProgress()

	return nil
}

func (d *download) showProgress(done <-chan struct{}) {
	d.printProgress()

	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			d.printProgress()
		case <-done:
			return
		}
	}
}

// percentRatio returns the ration of two integers as a percentage.
func percentRatio(a, b int64) float64 {
	return float64(a) / float64(b) * 100.0
}

func (d *download) printProgress() {
	fmt.Fprintf(os.Stderr, "succeeded: %.2f%%; progress: %.2f%%; downloaded: %s; rate: %s/s\n", percentRatio(d.succeeded.value(), d.size), percentRatio(d.progress.value(), d.size), humanize.Bytes(uint64(d.downloaded.value())), humanize.Bytes(uint64(d.downloadMeter.Rate1())))
}

// contentSize returns the size of the resource
// at the given url.
func contentSize(url string) (int64, error) {
	resp, err := http.Get(url)
	if err != nil {
		return 0, err
	}
	return resp.ContentLength, nil
}

// followRedirect ensures that all headers are passed through the redirect.
// This is necessary for the range header to work across redirects.
func followRedirect(req *http.Request, via []*http.Request) error {
	req.Header = via[0].Header
	return nil
}

// retryFetch will retry a fetch a few times.
func (d *download) retryFetch(from, len int64) (r io.ReadCloser, err error) {
	t := 1 * time.Second
	for i := 0; i < numRetries; i++ {
		r, err = d.fetch(from, len)
		if err == nil {
			return
		}

		if *debug {
			log.Printf("waiting %s, retry %d", t, i+1)
		}
		time.Sleep(t)
		t *= 2
	}
	return
}

// fetch returns a ReadCloser with the content of a
// section of the resource found at url starting at from
// and len long.
func (d *download) fetch(from, len int64) (io.ReadCloser, error) {
	req, err := http.NewRequest("GET", d.url, nil)
	if err != nil {
		return nil, err
	}

	req.Header.Add("Range", fmt.Sprintf("bytes=%d-%d", from, from+len-1))

	client := &http.Client{
		CheckRedirect: followRedirect,
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusPartialContent {
		defer resp.Body.Close()
		if resp.StatusCode == http.StatusOK {
			return nil, fmt.Errorf("Server doesn't implement partial content")
		}
		return nil, errors.New(resp.Status)
	}
	return &measuringReader{resp.Body, &d.downloaded, d.downloadMeter}, nil
}

type measuringReader struct {
	io.ReadCloser
	*counter
	metrics.Meter
}

func (m *measuringReader) Read(p []byte) (int, error) {
	n, err := m.ReadCloser.Read(p)
	m.add(int64(n))
	m.Mark(int64(n))
	return n, err
}

func main() {
	flag.Parse()
	if err := start(); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>